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, 09:45 UTC
[V1] Update zmq socket creation in nixl connector (#18148)
Signed-off-by: Russell Bryant <rbryant@redhat.com>
Russell Bryant committed
May 15, 2025 at 02:17 UTC
a8f5aec20ad685851f972847c0567db270d9845f
3 files changed
+34
-15
tests/test_utils.py
+6
-1
@@ -17,7 +17,7 @@ from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config
17
from vllm.utils import (CacheInfo, FlexibleArgumentParser, LRUCache,
18
MemorySnapshot, PlaceholderModule, StoreBoolean,
19
bind_kv_cache, deprecate_kwargs, get_open_port,
20
- make_zmq_socket, memory_profiling,
20
+ make_zmq_path, make_zmq_socket, memory_profiling,
21
merge_async_iterators, sha256, split_zmq_path,
22
supports_kw, swap_dict_values)
23
@@ -714,3 +714,8 @@ def test_make_zmq_socket_ipv6():
714
# Clean up
715
zsock.close()
716
ctx.term()
717
+
718
+
719
+def test_make_zmq_path():
720
+ assert make_zmq_path("tcp", "127.0.0.1", "5555") == "tcp://127.0.0.1:5555"
721
+ assert make_zmq_path("tcp", "::1", "5555") == "tcp://[::1]:5555"
vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+10
-14
@@ -21,7 +21,7 @@ from vllm.distributed.parallel_state import (
21
get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size,
22
get_tp_group)
23
from vllm.logger import init_logger
24
-from vllm.utils import round_down
24
+from vllm.utils import make_zmq_path, make_zmq_socket, round_down
25
from vllm.v1.core.sched.output import SchedulerOutput
26
from vllm.v1.request import RequestStatus
27
@@ -379,7 +379,7 @@ class NixlConnectorWorker:
379
# hack to keeps us moving. We will switch when moving to etcd
380
# or where we have a single ZMQ socket in the scheduler.
381
port = envs.VLLM_NIXL_SIDE_CHANNEL_PORT + rank
382
- path = f"tcp://{host}:{port}"
382
+ path = make_zmq_path("tcp", host, port)
383
logger.debug("Starting listening on path: %s", path)
384
with zmq_ctx(zmq.ROUTER, path) as sock:
385
ready_event.set()
@@ -397,7 +397,7 @@ class NixlConnectorWorker:
397
# NOTE(rob): we need each rank to have a unique port. This is
398
# a hack to keep us moving. We will switch when moving to etcd
399
# or where we have a single ZMQ socket in the scheduler.
400
- path = f"tcp://{host}:{port + self.rank}"
400
+ path = make_zmq_path("tcp", host, port + self.rank)
401
logger.debug("Querying metadata on path: %s", path)
402
with zmq_ctx(zmq.REQ, path) as sock:
403
# Send query for the request.
@@ -741,20 +741,16 @@ class NixlConnectorWorker:
741
def zmq_ctx(socket_type: Any, addr: str) -> Iterator[zmq.Socket]:
742
"""Context manager for a ZMQ socket"""
743
744
+ if socket_type not in (zmq.ROUTER, zmq.REQ):
745
+ raise ValueError(f"Unexpected socket type: {socket_type}")
746
+
747
ctx: Optional[zmq.Context] = None
748
try:
749
ctx = zmq.Context() # type: ignore[attr-defined]
747
-
748
- if socket_type == zmq.ROUTER:
749
- socket = ctx.socket(zmq.ROUTER)
750
- socket.bind(addr)
751
- elif socket_type == zmq.REQ:
752
- socket = ctx.socket(zmq.REQ)
753
- socket.connect(addr)
754
- else:
755
- raise ValueError(f"Unexpected socket type: {socket_type}")
756
-
757
- yield socket
750
+ yield make_zmq_socket(ctx=ctx,
751
+ path=addr,
752
+ socket_type=socket_type,
753
+ bind=socket_type == zmq.ROUTER)
754
finally:
755
if ctx is not None:
756
ctx.destroy(linger=0)
vllm/utils.py
+18
@@ -2350,6 +2350,24 @@ def split_zmq_path(path: str) -> Tuple[str, str, str]:
2350
return scheme, host, port
2351
2352
2353
+def make_zmq_path(scheme: str, host: str, port: Optional[int] = None) -> str:
2354
+ """Make a ZMQ path from its parts.
2355
+
2356
+ Args:
2357
+ scheme: The ZMQ transport scheme (e.g. tcp, ipc, inproc).
2358
+ host: The host - can be an IPv4 address, IPv6 address, or hostname.
2359
+ port: Optional port number, only used for TCP sockets.
2360
+
2361
+ Returns:
2362
+ A properly formatted ZMQ path string.
2363
+ """
2364
+ if not port:
2365
+ return f"{scheme}://{host}"
2366
+ if is_valid_ipv6_address(host):
2367
+ return f"{scheme}://[{host}]:{port}"
2368
+ return f"{scheme}://{host}:{port}"
2369
+
2370
+
2371
# Adapted from: https://github.com/sgl-project/sglang/blob/v0.4.1/python/sglang/srt/utils.py#L783 # noqa: E501
2372
def make_zmq_socket(
2373
ctx: Union[zmq.asyncio.Context, zmq.Context], # type: ignore[name-defined]