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, 12:15 UTC
[Core] Enable IPv6 with vllm.utils.make_zmq_socket() (#16506)
Signed-off-by: Russell Bryant <rbryant@redhat.com>
Russell Bryant committed
May 1, 2025 at 05:38 UTC
fbefc8a78d22b20eac042c586805c7dcbfc66b1c
2 files changed
+81
-1
tests/test_utils.py
+53
-1
@@ -10,13 +10,15 @@ from unittest.mock import patch
10
11
import pytest
12
import torch
13
+import zmq
14
from vllm_test_utils.monitor import monitor
15
16
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,
19
- memory_profiling, merge_async_iterators, sha256,
20
+ make_zmq_socket, memory_profiling,
21
+ merge_async_iterators, sha256, split_zmq_path,
22
supports_kw, swap_dict_values)
23
24
from .utils import create_new_process_for_each_test, error_on_warning
@@ -662,3 +664,53 @@ def test_sha256(input: tuple, output: int):
664
665
# hashing different input, returns different value
666
assert hash != sha256(input + (1, ))
667
+
668
+
669
+@pytest.mark.parametrize(
670
+ "path,expected",
671
+ [
672
+ ("ipc://some_path", ("ipc", "some_path", "")),
673
+ ("tcp://127.0.0.1:5555", ("tcp", "127.0.0.1", "5555")),
674
+ ("tcp://[::1]:5555", ("tcp", "::1", "5555")), # IPv6 address
675
+ ("inproc://some_identifier", ("inproc", "some_identifier", "")),
676
+ ]
677
+)
678
+def test_split_zmq_path(path, expected):
679
+ assert split_zmq_path(path) == expected
680
+
681
+
682
+@pytest.mark.parametrize(
683
+ "invalid_path",
684
+ [
685
+ "invalid_path", # Missing scheme
686
+ "tcp://127.0.0.1", # Missing port
687
+ "tcp://[::1]", # Missing port for IPv6
688
+ "tcp://:5555", # Missing host
689
+ ]
690
+)
691
+def test_split_zmq_path_invalid(invalid_path):
692
+ with pytest.raises(ValueError):
693
+ split_zmq_path(invalid_path)
694
+
695
+
696
+def test_make_zmq_socket_ipv6():
697
+ # Check if IPv6 is supported by trying to create an IPv6 socket
698
+ try:
699
+ sock = socket.socket(socket.AF_INET6, socket.SOCK_STREAM)
700
+ sock.close()
701
+ except socket.error:
702
+ pytest.skip("IPv6 is not supported on this system")
703
+
704
+ ctx = zmq.Context()
705
+ ipv6_path = "tcp://[::]:5555" # IPv6 loopback address
706
+ socket_type = zmq.REP # Example socket type
707
+
708
+ # Create the socket
709
+ zsock: zmq.Socket = make_zmq_socket(ctx, ipv6_path, socket_type)
710
+
711
+ # Verify that the IPV6 option is set
712
+ assert zsock.getsockopt(zmq.IPV6) == 1, "IPV6 option should be enabled for IPv6 addresses"
713
+
714
+ # Clean up
715
+ zsock.close()
716
+ ctx.term()
vllm/utils.py
+28
@@ -45,6 +45,7 @@ from types import MappingProxyType
45
from typing import (TYPE_CHECKING, Any, Callable, Generic, Literal, NamedTuple,
46
Optional, Sequence, Tuple, Type, TypeVar, Union, cast,
47
overload)
48
+from urllib.parse import urlparse
49
from uuid import uuid4
50
51
import cachetools
@@ -2278,6 +2279,27 @@ def get_exception_traceback():
2279
return err_str
2280
2281
2282
+def split_zmq_path(path: str) -> Tuple[str, str, str]:
2283
+ """Split a zmq path into its parts."""
2284
+ parsed = urlparse(path)
2285
+ if not parsed.scheme:
2286
+ raise ValueError(f"Invalid zmq path: {path}")
2287
+
2288
+ scheme = parsed.scheme
2289
+ host = parsed.hostname or ""
2290
+ port = str(parsed.port or "")
2291
+
2292
+ if scheme == "tcp" and not all((host, port)):
2293
+ # The host and port fields are required for tcp
2294
+ raise ValueError(f"Invalid zmq path: {path}")
2295
+
2296
+ if scheme != "tcp" and port:
2297
+ # port only makes sense with tcp
2298
+ raise ValueError(f"Invalid zmq path: {path}")
2299
+
2300
+ return scheme, host, port
2301
+
2302
+
2303
# Adapted from: https://github.com/sgl-project/sglang/blob/v0.4.1/python/sglang/srt/utils.py#L783 # noqa: E501
2304
def make_zmq_socket(
2305
ctx: Union[zmq.asyncio.Context, zmq.Context], # type: ignore[name-defined]
@@ -2317,6 +2339,12 @@ def make_zmq_socket(
2339
if identity is not None:
2340
socket.setsockopt(zmq.IDENTITY, identity)
2341
2342
+ # Determine if the path is a TCP socket with an IPv6 address.
2343
+ # Enable IPv6 on the zmq socket if so.
2344
+ scheme, host, _ = split_zmq_path(path)
2345
+ if scheme == "tcp" and is_valid_ipv6_address(host):
2346
+ socket.setsockopt(zmq.IPV6, 1)
2347
+
2348
if bind:
2349
socket.bind(path)
2350
else: