A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 22:15 UTC
[Core][Distributed] Use IPC (domain socket) ZMQ socket for local comms (#13688)
Nick Hill committed
Feb 23, 2025 at 02:54 UTC
5a2ba16f5c6cf508d6a8e8a82cdfd4741190166a
1 file changed
+21
-21
vllm/distributed/device_communicators/shm_broadcast.py
+21
-21
@@ -19,7 +19,8 @@ from zmq import SUB, SUBSCRIBE, XPUB, XPUB_VERBOSE, Context # type: ignore
19
import vllm.envs as envs
20
from vllm.distributed.utils import StatelessProcessGroup
21
from vllm.logger import init_logger
22
-from vllm.utils import get_ip, get_open_port, is_valid_ipv6_address
22
+from vllm.utils import (get_ip, get_open_port, get_open_zmq_ipc_path,
23
+ is_valid_ipv6_address)
24
25
VLLM_RINGBUFFER_WARNING_INTERVAL = envs.VLLM_RINGBUFFER_WARNING_INTERVAL
26
@@ -165,12 +166,12 @@ class ShmRingBuffer:
166
167
@dataclass
168
class Handle:
168
- connect_ip: str
169
local_reader_ranks: List[int] = field(default_factory=list)
170
171
buffer_handle: Optional[Tuple[int, int, int, str]] = None
172
- local_subscribe_port: Optional[int] = None
173
- remote_subscribe_port: Optional[int] = None
172
+ local_subscribe_addr: Optional[str] = None
173
+ remote_subscribe_addr: Optional[str] = None
174
+ remote_addr_ipv6: bool = False
175
176
177
class MessageQueue:
@@ -192,9 +193,6 @@ class MessageQueue:
193
n_remote_reader = n_reader - n_local_reader
194
self.n_remote_reader = n_remote_reader
195
195
- if connect_ip is None:
196
- connect_ip = get_ip() if n_remote_reader > 0 else "127.0.0.1"
197
-
196
context = Context()
197
198
if n_local_reader > 0:
@@ -212,32 +210,34 @@ class MessageQueue:
210
# message. otherwise, we will only receive the first subscription
211
# see http://api.zeromq.org/3-3:zmq-setsockopt for more details
212
self.local_socket.setsockopt(XPUB_VERBOSE, True)
215
- local_subscribe_port = get_open_port()
216
- socket_addr = f"tcp://127.0.0.1:{local_subscribe_port}"
217
- logger.debug("Binding to %s", socket_addr)
218
- self.local_socket.bind(socket_addr)
213
+ local_subscribe_addr = get_open_zmq_ipc_path()
214
+ logger.debug("Binding to %s", local_subscribe_addr)
215
+ self.local_socket.bind(local_subscribe_addr)
216
217
self.current_idx = 0
221
-
218
else:
219
self.buffer = None # type: ignore
224
- local_subscribe_port = None
220
+ local_subscribe_addr = None
221
self.local_socket = None
222
self.current_idx = -1
223
224
+ remote_addr_ipv6 = False
225
if n_remote_reader > 0:
226
# for remote readers, we will:
227
# create a publish-subscribe socket to communicate large data
228
+ if not connect_ip:
229
+ connect_ip = get_ip()
230
self.remote_socket = context.socket(XPUB)
231
self.remote_socket.setsockopt(XPUB_VERBOSE, True)
232
remote_subscribe_port = get_open_port()
233
if is_valid_ipv6_address(connect_ip):
234
self.remote_socket.setsockopt(IPV6, 1)
235
+ remote_addr_ipv6 = True
236
socket_addr = f"tcp://*:{remote_subscribe_port}"
237
self.remote_socket.bind(socket_addr)
238
-
238
+ remote_subscribe_addr = f"tcp://{connect_ip}:{remote_subscribe_port}"
239
else:
240
- remote_subscribe_port = None
240
+ remote_subscribe_addr = None
241
self.remote_socket = None
242
243
self._is_writer = True
@@ -247,12 +247,12 @@ class MessageQueue:
247
self._is_remote_reader = False
248
249
self.handle = Handle(
250
- connect_ip=connect_ip,
250
local_reader_ranks=local_reader_ranks,
251
buffer_handle=self.buffer.handle()
252
if self.buffer is not None else None,
254
- local_subscribe_port=local_subscribe_port,
255
- remote_subscribe_port=remote_subscribe_port,
253
+ local_subscribe_addr=local_subscribe_addr,
254
+ remote_subscribe_addr=remote_subscribe_addr,
255
+ remote_addr_ipv6=remote_addr_ipv6,
256
)
257
258
logger.info("vLLM message queue communication handle: %s", self.handle)
@@ -278,7 +278,7 @@ class MessageQueue:
278
279
self.local_socket = context.socket(SUB)
280
self.local_socket.setsockopt_string(SUBSCRIBE, "")
281
- socket_addr = f"tcp://127.0.0.1:{handle.local_subscribe_port}"
281
+ socket_addr = handle.local_subscribe_addr
282
logger.debug("Connecting to %s", socket_addr)
283
self.local_socket.connect(socket_addr)
284
@@ -294,9 +294,9 @@ class MessageQueue:
294
295
self.remote_socket = context.socket(SUB)
296
self.remote_socket.setsockopt_string(SUBSCRIBE, "")
297
- if is_valid_ipv6_address(handle.connect_ip):
297
+ if handle.remote_addr_ipv6:
298
self.remote_socket.setsockopt(IPV6, 1)
299
- socket_addr = f"tcp://{handle.connect_ip}:{handle.remote_subscribe_port}"
299
+ socket_addr = handle.remote_subscribe_addr
300
logger.debug("Connecting to %s", socket_addr)
301
self.remote_socket.connect(socket_addr)
302