@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 285 lines 9.64 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Tests for control-layer ABCs, ZmqClientConnection, and ECSession."""
4
5 from unittest.mock import MagicMock, patch
6
7 import pytest
8 import zmq
9
10 from vllm.distributed.ec_transfer.ec_connector.cpu.control.base import (
11 ControlConnection,
12 )
13
14
15 class _ConcreteConn(ControlConnection):
16 """Minimal concrete subclass for ABC-instantiation tests."""
17
18 def __init__(self):
19 self._alive = True
20 self._sent = []
21 self._inbox = []
22
23 @property
24 def alive(self) -> bool:
25 return self._alive
26
27 def send(self, msg: bytes) -> None:
28 self._sent.append(msg)
29
30 def recv(self) -> list[bytes]:
31 msgs, self._inbox = self._inbox, []
32 return msgs
33
34 def mark_dead(self) -> None:
35 self._alive = False
36
37 def close(self) -> None:
38 self._alive = False
39
40
41 def test_control_connection_abc_cannot_be_instantiated_directly():
42 with pytest.raises(TypeError):
43 ControlConnection() # type: ignore[abstract]
44
45
46 def test_concrete_connection_send_recv():
47 conn = _ConcreteConn()
48 assert conn.alive
49 conn.send(b"hello")
50 conn._inbox.append(b"world")
51 assert conn.recv() == [b"world"]
52 assert conn._sent == [b"hello"]
53
54
55 def test_mark_dead_sets_alive_false():
56 conn = _ConcreteConn()
57 conn.mark_dead()
58 assert not conn.alive
59
60
61 def test_close_sets_alive_false():
62 conn = _ConcreteConn()
63 conn.close()
64 assert not conn.alive
65
66
67 # ── ZmqClientConnection ───────────────────────────────────────────────────────
68
69
70 from vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq import ( # noqa: E402
71 ZmqClientConnection,
72 ZmqClientTransport,
73 )
74 from vllm.distributed.ec_transfer.ec_connector.cpu.utils import PeerAddr # noqa: E402
75
76
77 def _conn(messages: list[list[bytes]], monitor=None) -> ZmqClientConnection:
78 """Build a ZmqClientConnection whose dealer yields the given frame-lists."""
79 dealer = MagicMock()
80 dealer.recv_multipart.side_effect = [*messages, zmq.Again()]
81 return ZmqClientConnection(dealer=dealer, monitor=monitor)
82
83
84 def test_zmq_client_connection_recv_extracts_payloads_from_envelopes():
85 conn = _conn([[b"", b"msg1"], [b"", b"msg2"]])
86 assert conn.recv() == [b"msg1", b"msg2"]
87
88
89 def test_zmq_client_connection_recv_skips_frames_without_empty_delimiter():
90 # A frame that does not start with b"" is malformed and must be dropped.
91 conn = _conn([[b"", b"ok"], [b"no-delimiter", b"bad"]])
92 assert conn.recv() == [b"ok"]
93
94
95 def test_zmq_client_connection_recv_breaks_on_non_again_exception():
96 # A socket error should stop the drain and return whatever was collected.
97 dealer = MagicMock()
98 dealer.recv_multipart.side_effect = [[b"", b"first"], OSError("broken")]
99 conn = ZmqClientConnection(dealer=dealer, monitor=None)
100 assert conn.recv() == [b"first"]
101
102
103 def test_zmq_client_connection_send_prefixes_empty_delimiter():
104 dealer = MagicMock()
105 conn = ZmqClientConnection(dealer=dealer, monitor=None)
106 conn.send(b"payload")
107 dealer.send_multipart.assert_called_once_with([b"", b"payload"])
108
109
110 def test_zmq_client_connection_mark_dead_does_not_close_socket():
111 # mark_dead signals liveness only; resource teardown is close()'s job.
112 dealer = MagicMock()
113 conn = ZmqClientConnection(dealer=dealer, monitor=None)
114 conn.mark_dead()
115 assert not conn.alive
116 dealer.close.assert_not_called()
117
118
119 def test_zmq_client_connection_close_without_monitor():
120 dealer = MagicMock()
121 conn = ZmqClientConnection(dealer=dealer, monitor=None)
122 conn.close()
123 dealer.close.assert_called_once_with(linger=0)
124 assert not conn.alive
125
126
127 def test_zmq_client_connection_close_with_monitor_disables_then_closes_both():
128 # The disable_monitor → monitor.close → dealer.close sequence must be
129 # respected: closing the dealer before disabling the monitor can cause
130 # a zmq assertion in some versions.
131 dealer = MagicMock()
132 monitor = MagicMock()
133 conn = ZmqClientConnection(dealer=dealer, monitor=monitor)
134 conn.close()
135 dealer.disable_monitor.assert_called_once()
136 monitor.close.assert_called_once_with(linger=0)
137 dealer.close.assert_called_once_with(linger=0)
138 assert not conn.alive
139
140
141 def test_zmq_client_connection_close_is_idempotent():
142 dealer = MagicMock()
143 conn = ZmqClientConnection(dealer=dealer, monitor=None)
144 conn.close()
145 conn.close() # must not raise
146 assert not conn.alive
147
148
149 # ── ZmqClientTransport ────────────────────────────────────────────────────────
150
151
152 def _make_transport() -> tuple[ZmqClientTransport, MagicMock]:
153 """Return a ZmqClientTransport with a mocked ZMQ context."""
154 ctx = MagicMock()
155 ctx.socket.return_value = MagicMock()
156 with (
157 patch(
158 "vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq.zmq.Context",
159 return_value=ctx,
160 ),
161 patch(
162 "vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq.make_zmq_socket",
163 return_value=MagicMock(),
164 ),
165 patch(
166 "vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq.make_zmq_path",
167 return_value="tcp://host:0",
168 ),
169 ):
170 transport = ZmqClientTransport()
171 transport._ctx = ctx
172 return transport, ctx
173
174
175 def _inject_conn(transport: ZmqClientTransport, addr: PeerAddr) -> ZmqClientConnection:
176 """Inject a ZmqClientConnection directly into the transport pool."""
177 conn = ZmqClientConnection(dealer=MagicMock(), monitor=MagicMock())
178 transport._connections[addr] = conn
179 return conn
180
181
182 @patch("vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq.make_zmq_socket")
183 @patch("vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq.make_zmq_path")
184 @patch("vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq.zmq.Context")
185 def test_zmq_client_transport_connect_caches_connection(
186 mock_ctx_cls, mock_path, mock_sock
187 ):
188 mock_ctx_cls.return_value = MagicMock()
189 mock_ctx_cls.return_value.socket.return_value = MagicMock()
190 mock_sock.return_value = MagicMock()
191 transport = ZmqClientTransport()
192 addr: PeerAddr = ("host", 1234)
193 assert transport.connect(addr) is transport.connect(addr)
194
195
196 @patch("vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq.make_zmq_socket")
197 @patch("vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq.make_zmq_path")
198 @patch("vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq.zmq.Context")
199 def test_zmq_client_transport_connect_creates_independent_conns_per_addr(
200 mock_ctx_cls, mock_path, mock_sock
201 ):
202 mock_ctx_cls.return_value = MagicMock()
203 mock_ctx_cls.return_value.socket.return_value = MagicMock()
204 mock_sock.return_value = MagicMock()
205 transport = ZmqClientTransport()
206 conn_a = transport.connect(("host-a", 1111))
207 conn_b = transport.connect(("host-b", 2222))
208 assert conn_a is not conn_b
209
210
211 @patch("vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq.recv_monitor_message")
212 def test_zmq_client_transport_poll_dead_detects_disconnected_peer(mock_recv_mon):
213 transport, _ = _make_transport()
214 addr: PeerAddr = ("host", 1234)
215 conn = _inject_conn(transport, addr)
216 mock_recv_mon.side_effect = [{"event": zmq.EVENT_DISCONNECTED}, zmq.Again()]
217
218 dead = transport.poll_dead()
219
220 assert addr in dead
221 assert not conn.alive
222
223
224 @patch("vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq.recv_monitor_message")
225 def test_zmq_client_transport_poll_dead_skips_conn_without_monitor(mock_recv_mon):
226 transport, _ = _make_transport()
227 addr: PeerAddr = ("host", 9999)
228 conn = ZmqClientConnection(dealer=MagicMock(), monitor=None)
229 transport._connections[addr] = conn
230
231 dead = transport.poll_dead()
232
233 assert dead == []
234 mock_recv_mon.assert_not_called()
235
236
237 @patch("vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq.recv_monitor_message")
238 def test_zmq_client_transport_poll_dead_ignores_non_disconnect_events(mock_recv_mon):
239 transport, _ = _make_transport()
240 addr: PeerAddr = ("host", 1234)
241 conn = _inject_conn(transport, addr)
242 # A CONNECTED event should not trigger a dead notification.
243 mock_recv_mon.side_effect = [{"event": zmq.EVENT_CONNECTED}, zmq.Again()]
244
245 dead = transport.poll_dead()
246
247 assert dead == []
248 assert conn.alive
249
250
251 def test_zmq_client_transport_remove_closes_and_evicts_conn():
252 transport, _ = _make_transport()
253 addr: PeerAddr = ("host", 5678)
254 conn = _inject_conn(transport, addr)
255
256 removed = transport.remove(addr)
257
258 assert removed is conn
259 assert addr not in transport._connections
260 conn.dealer.close.assert_called()
261
262
263 def test_zmq_client_transport_remove_returns_none_for_unknown_addr():
264 transport, _ = _make_transport()
265 assert transport.remove(("ghost", 0)) is None
266
267
268 @patch("vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq.make_zmq_socket")
269 @patch("vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq.make_zmq_path")
270 @patch("vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq.zmq.Context")
271 def test_zmq_client_transport_close_closes_all_connections(
272 mock_ctx_cls, mock_path, mock_sock
273 ):
274 mock_ctx = MagicMock()
275 mock_ctx_cls.return_value = mock_ctx
276 mock_ctx.socket.return_value = MagicMock()
277 dealer_a, dealer_b = MagicMock(), MagicMock()
278 mock_sock.side_effect = [dealer_a, dealer_b]
279 transport = ZmqClientTransport()
280 transport.connect(("a", 1))
281 transport.connect(("b", 2))
282 transport.close()
283 dealer_a.close.assert_called()
284 dealer_b.close.assert_called()
285 assert transport._connections == {}