@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 251 lines 7.98 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import socket
4 import threading
5
6 import pytest
7 import zmq
8
9 from vllm.utils import network_utils
10 from vllm.utils.network_utils import (
11 get_file_store_init_method,
12 get_open_port,
13 get_open_ports_list,
14 get_tcp_uri,
15 join_host_port,
16 make_zmq_path,
17 make_zmq_socket,
18 split_host_port,
19 split_zmq_path,
20 )
21
22
23 def test_get_file_store_init_method_is_unique():
24 init_methods = {get_file_store_init_method() for _ in range(2)}
25
26 assert len(init_methods) == 2
27 assert all(method.startswith("file://") for method in init_methods)
28
29
30 def _call_with_timeout(func, timeout: float = 10.0):
31 """Run func in a daemon thread so a livelock regression fails the test
32 quickly instead of hanging the suite."""
33 result: dict = {}
34
35 def target():
36 try:
37 result["value"] = func()
38 except BaseException as e:
39 result["error"] = e
40
41 thread = threading.Thread(target=target, daemon=True)
42 thread.start()
43 thread.join(timeout)
44 assert not thread.is_alive(), f"call did not finish within {timeout}s"
45 if "error" in result:
46 raise result["error"]
47 return result["value"]
48
49
50 def test_get_open_port(monkeypatch: pytest.MonkeyPatch):
51 with monkeypatch.context() as m:
52 m.setenv("VLLM_PORT", "5678")
53 # make sure we can get multiple ports, even if the env var is set
54 with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s1:
55 s1.bind(("localhost", get_open_port()))
56 with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s2:
57 s2.bind(("localhost", get_open_port()))
58 with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s3:
59 s3.bind(("localhost", get_open_port()))
60
61
62 def test_get_open_port_vllm_port_in_dp_reserved_range(
63 monkeypatch: pytest.MonkeyPatch,
64 ):
65 # VLLM_PORT falling inside the data-parallel reserved window used to make
66 # get_open_port() loop forever (issue #50024). It must instead return a
67 # port outside the reserved range.
68 with monkeypatch.context() as m:
69 m.setenv("VLLM_DP_MASTER_PORT", "5680")
70 # 5682 is inside [5680, 5690).
71 m.setenv("VLLM_PORT", "5682")
72 port = _call_with_timeout(get_open_port)
73 assert port not in range(5680, 5690)
74
75
76 def test_get_open_ports_list_with_vllm_port(monkeypatch: pytest.MonkeyPatch):
77 with monkeypatch.context() as m:
78 m.setenv("VLLM_PORT", "5678")
79 ports = get_open_ports_list(5)
80 assert len(ports) == 5
81 assert len(set(ports)) == 5, "ports must be unique"
82
83 # verify every port is actually bindable
84 sockets = []
85 try:
86 for p in ports:
87 s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
88 s.bind(("localhost", p))
89 sockets.append(s)
90 finally:
91 for s in sockets:
92 s.close()
93
94
95 def test_get_open_port_skips_reserved_dp_master_ports(
96 monkeypatch: pytest.MonkeyPatch,
97 ):
98 """VLLM_PORT == VLLM_DP_MASTER_PORT must not livelock by returning the
99 same free-but-reserved port forever (upstream issue #50024)."""
100 base = 42123
101 with monkeypatch.context() as m:
102 m.setenv("VLLM_PORT", str(base))
103 m.setenv("VLLM_DP_MASTER_PORT", str(base))
104 port = _call_with_timeout(get_open_port)
105 assert port >= base + 10
106
107
108 def test_get_open_ports_list_skips_reserved_dp_master_ports(
109 monkeypatch: pytest.MonkeyPatch,
110 ):
111 """The VLLM_PORT band scan must not hand out ports inside the window
112 reserved for the data parallel master process."""
113 base = 42223
114 reserved = range(base, base + 10)
115 with monkeypatch.context() as m:
116 m.setenv("VLLM_PORT", str(base))
117 m.setenv("VLLM_DP_MASTER_PORT", str(base))
118 ports = _call_with_timeout(lambda: get_open_ports_list(5))
119 assert len(ports) == 5
120 assert len(set(ports)) == 5, "ports must be unique"
121 assert all(p not in reserved for p in ports)
122
123
124 def test_get_open_port_ephemeral_skips_reserved_range(
125 monkeypatch: pytest.MonkeyPatch,
126 ):
127 """An ephemeral port that happens to fall in the reserved range is
128 replaced by a rescan starting past the range."""
129 base = 42323
130 with monkeypatch.context() as m:
131 m.delenv("VLLM_PORT", raising=False)
132 m.setenv("VLLM_DP_MASTER_PORT", str(base))
133 m.setattr(
134 network_utils,
135 "_get_open_port",
136 lambda start_port=None, max_attempts=None: (
137 base if start_port is None else start_port
138 ),
139 )
140 assert _call_with_timeout(get_open_port) == base + 10
141
142
143 def test_get_open_port_ephemeral_without_dp_master_port(
144 monkeypatch: pytest.MonkeyPatch,
145 ):
146 """Without VLLM_PORT and VLLM_DP_MASTER_PORT, an ephemeral port is
147 returned as before."""
148 with monkeypatch.context() as m:
149 m.delenv("VLLM_PORT", raising=False)
150 m.delenv("VLLM_DP_MASTER_PORT", raising=False)
151 port = _call_with_timeout(get_open_port)
152 with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
153 s.bind(("localhost", port))
154
155
156 @pytest.mark.parametrize(
157 "path,expected",
158 [
159 ("ipc://some_path", ("ipc", "some_path", "")),
160 ("tcp://127.0.0.1:5555", ("tcp", "127.0.0.1", "5555")),
161 ("tcp://[::1]:5555", ("tcp", "::1", "5555")), # IPv6 address
162 ("inproc://some_identifier", ("inproc", "some_identifier", "")),
163 ],
164 )
165 def test_split_zmq_path(path, expected):
166 assert split_zmq_path(path) == expected
167
168
169 @pytest.mark.parametrize(
170 "invalid_path",
171 [
172 "invalid_path", # Missing scheme
173 "tcp://127.0.0.1", # Missing port
174 "tcp://[::1]", # Missing port for IPv6
175 "tcp://:5555", # Missing host
176 ],
177 )
178 def test_split_zmq_path_invalid(invalid_path):
179 with pytest.raises(ValueError):
180 split_zmq_path(invalid_path)
181
182
183 def test_make_zmq_socket_ipv6():
184 # Check if IPv6 is supported by trying to create an IPv6 socket
185 try:
186 sock = socket.socket(socket.AF_INET6, socket.SOCK_STREAM)
187 sock.close()
188 except OSError:
189 pytest.skip("IPv6 is not supported on this system")
190
191 ctx = zmq.Context()
192 ipv6_path = "tcp://[::]:5555" # IPv6 loopback address
193 socket_type = zmq.REP # Example socket type
194
195 # Create the socket
196 zsock: zmq.Socket = make_zmq_socket(ctx, ipv6_path, socket_type)
197
198 # Verify that the IPV6 option is set
199 assert zsock.getsockopt(zmq.IPV6) == 1, (
200 "IPV6 option should be enabled for IPv6 addresses"
201 )
202
203 # Clean up
204 zsock.close()
205 ctx.term()
206
207
208 def test_make_zmq_path():
209 assert make_zmq_path("tcp", "127.0.0.1", 5555) == "tcp://127.0.0.1:5555"
210 assert make_zmq_path("tcp", "::1", 5555) == "tcp://[::1]:5555"
211
212
213 def test_get_tcp_uri():
214 assert get_tcp_uri("127.0.0.1", 5555) == "tcp://127.0.0.1:5555"
215 assert get_tcp_uri("::1", 5555) == "tcp://[::1]:5555"
216
217
218 def test_split_host_port():
219 # valid ipv4
220 assert split_host_port("127.0.0.1:5555") == ("127.0.0.1", 5555)
221 # invalid ipv4
222 with pytest.raises(ValueError):
223 # multi colon
224 assert split_host_port("127.0.0.1::5555")
225 with pytest.raises(ValueError):
226 # tailing colon
227 assert split_host_port("127.0.0.1:5555:")
228 with pytest.raises(ValueError):
229 # no colon
230 assert split_host_port("127.0.0.15555")
231 with pytest.raises(ValueError):
232 # none int port
233 assert split_host_port("127.0.0.1:5555a")
234
235 # valid ipv6
236 assert split_host_port("[::1]:5555") == ("::1", 5555)
237 # invalid ipv6
238 with pytest.raises(ValueError):
239 # multi colon
240 assert split_host_port("[::1]::5555")
241 with pytest.raises(IndexError):
242 # no colon
243 assert split_host_port("[::1]5555")
244 with pytest.raises(ValueError):
245 # none int port
246 assert split_host_port("[::1]:5555a")
247
248
249 def test_join_host_port():
250 assert join_host_port("127.0.0.1", 5555) == "127.0.0.1:5555"
251 assert join_host_port("::1", 5555) == "[::1]:5555"