| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Unit tests for vllm.v1.kv_offload.tiering.p2p.control.zmq.""" |
| 4 | |
| 5 | from __future__ import annotations |
| 6 | |
| 7 | import socket |
| 8 | import time |
| 9 | |
| 10 | import pytest |
| 11 | import zmq |
| 12 | |
| 13 | from vllm.v1.kv_offload.tiering.p2p.control.zmq import ( |
| 14 | ZmqConnection, |
| 15 | ZmqTransport, |
| 16 | _Sockets, |
| 17 | ) |
| 18 | |
| 19 | |
| 20 | def _free_port() -> int: |
| 21 | """Find a free TCP port.""" |
| 22 | with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: |
| 23 | s.bind(("127.0.0.1", 0)) |
| 24 | return s.getsockname()[1] |
| 25 | |
| 26 | |
| 27 | def _make_transport(host: str = "127.0.0.1", attempts: int = 8): |
| 28 | """Construct a ZmqTransport on a fresh port, retrying on bind collisions. |
| 29 | |
| 30 | Why: _free_port() releases the probe socket before ZmqTransport binds the |
| 31 | same port — a parallel test run can steal it in between. Retrying on |
| 32 | ZMQError/OSError closes that race without a production change. |
| 33 | """ |
| 34 | last_err: Exception | None = None |
| 35 | for _ in range(attempts): |
| 36 | port = _free_port() |
| 37 | try: |
| 38 | return ZmqTransport(f"{host}:{port}", host, port), port |
| 39 | except (zmq.ZMQError, OSError) as e: |
| 40 | last_err = e |
| 41 | assert last_err is not None |
| 42 | raise last_err |
| 43 | |
| 44 | |
| 45 | def _wait_for_inbound(transport: ZmqTransport, deadline: float = 2.0): |
| 46 | """Poll until at least one new inbound connection is accepted, or fail.""" |
| 47 | end = time.monotonic() + deadline |
| 48 | while time.monotonic() < end: |
| 49 | new = transport.poll() |
| 50 | if new: |
| 51 | return new |
| 52 | time.sleep(0.005) |
| 53 | raise AssertionError(f"no inbound connection within {deadline}s") |
| 54 | |
| 55 | |
| 56 | def _wait_for_messages( |
| 57 | transport: ZmqTransport, |
| 58 | conn: ZmqConnection, |
| 59 | n: int, |
| 60 | deadline: float = 2.0, |
| 61 | ) -> list[dict]: |
| 62 | """Poll until `conn` has received at least `n` messages, then return them.""" |
| 63 | end = time.monotonic() + deadline |
| 64 | msgs: list[dict] = [] |
| 65 | while time.monotonic() < end: |
| 66 | transport.poll() |
| 67 | msgs.extend(conn.recv()) |
| 68 | if len(msgs) >= n: |
| 69 | return msgs |
| 70 | time.sleep(0.005) |
| 71 | raise AssertionError(f"got {len(msgs)}/{n} messages within {deadline}s") |
| 72 | |
| 73 | |
| 74 | def _make_mock_connection(peer_id: str = "test:1234") -> ZmqConnection: |
| 75 | """Create a ZmqConnection with mock sockets for unit testing.""" |
| 76 | from unittest.mock import MagicMock |
| 77 | |
| 78 | sockets = _Sockets(dealer=MagicMock(), monitor=MagicMock()) |
| 79 | return ZmqConnection(peer_id, sockets) |
| 80 | |
| 81 | |
| 82 | class TestZmqConnection: |
| 83 | """Tests for ZmqConnection in isolation (no real sockets).""" |
| 84 | |
| 85 | def test_enqueue_and_recv(self): |
| 86 | """Messages enqueued are returned by recv() in order.""" |
| 87 | conn = _make_mock_connection() |
| 88 | |
| 89 | conn.enqueue({"type": "a"}) |
| 90 | conn.enqueue({"type": "b"}) |
| 91 | |
| 92 | msgs = conn.recv() |
| 93 | assert list(msgs) == [{"type": "a"}, {"type": "b"}] |
| 94 | # Second recv is empty |
| 95 | assert not conn.recv() |
| 96 | |
| 97 | def test_recv_returns_empty_initially(self): |
| 98 | conn = _make_mock_connection() |
| 99 | assert not conn.recv() |
| 100 | |
| 101 | def test_alive_initially_true(self): |
| 102 | conn = _make_mock_connection() |
| 103 | assert conn.alive is True |
| 104 | |
| 105 | def test_mark_dead(self): |
| 106 | conn = _make_mock_connection() |
| 107 | conn.mark_dead() |
| 108 | assert conn.alive is False |
| 109 | |
| 110 | def test_send_raises_when_closed(self): |
| 111 | conn = _make_mock_connection() |
| 112 | conn.mark_dead() |
| 113 | |
| 114 | with pytest.raises(RuntimeError, match="closed connection"): |
| 115 | conn.send({"type": "test"}) |
| 116 | |
| 117 | |
| 118 | class TestZmqTransportConnectivity: |
| 119 | """Integration tests for ZmqTransport with real ZMQ sockets.""" |
| 120 | |
| 121 | def test_connect_and_send_message(self): |
| 122 | """Two transports can connect and exchange messages.""" |
| 123 | transport_a, port_a = _make_transport() |
| 124 | transport_b, port_b = _make_transport() |
| 125 | |
| 126 | try: |
| 127 | peer_a_id = f"127.0.0.1:{port_a}" |
| 128 | conn_b_to_a = transport_b.connect(peer_a_id) |
| 129 | conn_b_to_a.send({"type": "hello", "data": 42}) |
| 130 | |
| 131 | new_conns = _wait_for_inbound(transport_a) |
| 132 | assert len(new_conns) == 1 |
| 133 | |
| 134 | conn_a_from_b = new_conns[0] |
| 135 | assert conn_a_from_b.peer_id == f"127.0.0.1:{port_b}" |
| 136 | |
| 137 | msgs = _wait_for_messages(transport_a, conn_a_from_b, 1) |
| 138 | assert msgs == [{"type": "hello", "data": 42}] |
| 139 | finally: |
| 140 | transport_a.close() |
| 141 | transport_b.close() |
| 142 | |
| 143 | def test_bidirectional_messaging(self): |
| 144 | """Both sides can send and receive after connection.""" |
| 145 | transport_a, port_a = _make_transport() |
| 146 | transport_b, _ = _make_transport() |
| 147 | |
| 148 | try: |
| 149 | conn_b = transport_b.connect(f"127.0.0.1:{port_a}") |
| 150 | conn_b.send({"type": "connect", "from": "b"}) |
| 151 | |
| 152 | new_conns = _wait_for_inbound(transport_a) |
| 153 | assert len(new_conns) == 1 |
| 154 | conn_a = new_conns[0] |
| 155 | |
| 156 | conn_a.send({"type": "reply", "from": "a"}) |
| 157 | |
| 158 | msgs = _wait_for_messages(transport_b, conn_b, 1) |
| 159 | assert msgs == [{"type": "reply", "from": "a"}] |
| 160 | finally: |
| 161 | transport_a.close() |
| 162 | transport_b.close() |
| 163 | |
| 164 | def test_poll_returns_empty_when_no_connections(self): |
| 165 | transport, _ = _make_transport() |
| 166 | try: |
| 167 | assert not transport.poll() |
| 168 | finally: |
| 169 | transport.close() |
| 170 | |
| 171 | def test_multiple_messages(self): |
| 172 | """Multiple messages are buffered and returned together.""" |
| 173 | transport_a, port_a = _make_transport() |
| 174 | transport_b, _ = _make_transport() |
| 175 | |
| 176 | try: |
| 177 | conn_b = transport_b.connect(f"127.0.0.1:{port_a}") |
| 178 | conn_b.send({"seq": 1}) |
| 179 | conn_b.send({"seq": 2}) |
| 180 | conn_b.send({"seq": 3}) |
| 181 | |
| 182 | new_conns = _wait_for_inbound(transport_a) |
| 183 | assert len(new_conns) == 1 |
| 184 | conn_a = new_conns[0] |
| 185 | |
| 186 | msgs = _wait_for_messages(transport_a, conn_a, 3) |
| 187 | assert [m["seq"] for m in msgs] == [1, 2, 3] |
| 188 | finally: |
| 189 | transport_a.close() |
| 190 | transport_b.close() |
| 191 | |
| 192 | def test_duplicate_connect_asserts(self): |
| 193 | """Connecting to the same peer twice raises AssertionError.""" |
| 194 | # port_a is never bound — we just need a syntactically-valid peer id. |
| 195 | port_a = _free_port() |
| 196 | transport_b, _ = _make_transport() |
| 197 | try: |
| 198 | transport_b.connect(f"127.0.0.1:{port_a}") |
| 199 | with pytest.raises(AssertionError, match="already exists"): |
| 200 | transport_b.connect(f"127.0.0.1:{port_a}") |
| 201 | finally: |
| 202 | transport_b.close() |
| 203 | |
| 204 | def test_dead_connection_removed_on_poll(self): |
| 205 | """Dead connections are cleaned up during poll.""" |
| 206 | transport_a, port_a = _make_transport() |
| 207 | transport_b, _ = _make_transport() |
| 208 | |
| 209 | try: |
| 210 | conn_b = transport_b.connect(f"127.0.0.1:{port_a}") |
| 211 | conn_b.send({"type": "hello"}) |
| 212 | |
| 213 | new_conns = _wait_for_inbound(transport_a) |
| 214 | assert len(new_conns) == 1 |
| 215 | |
| 216 | # Mark the inbound connection dead manually. |
| 217 | new_conns[0].mark_dead() |
| 218 | |
| 219 | # Pruning is synchronous within poll(). |
| 220 | transport_a.poll() |
| 221 | assert len(transport_a._connections) == 0 |
| 222 | assert new_conns[0]._sockets.dealer.closed |
| 223 | finally: |
| 224 | transport_a.close() |
| 225 | transport_b.close() |
| 226 | |
| 227 | def test_close_is_idempotent(self): |
| 228 | """Calling close() twice doesn't raise.""" |
| 229 | transport, _ = _make_transport() |
| 230 | transport.close() |
| 231 | transport.close() # should not raise |
| 232 | |
| 233 | |
| 234 | class TestZmqReconnect: |
| 235 | """Reconnecting to a peer whose connection died (real ZMQ sockets). |
| 236 | |
| 237 | A session marks its connection dead while handling messages, which happens |
| 238 | after the transport's own sweep has run for that tick — so a dead |
| 239 | connection stays registered until the next poll(). Reconnecting in that |
| 240 | window must succeed, and the retired connection must release its sockets. |
| 241 | Real sockets are required: a mock reports every attribute as closed. |
| 242 | """ |
| 243 | |
| 244 | def test_close_after_mark_dead_releases_sockets(self): |
| 245 | """close() releases sockets even when mark_dead() ran first. |
| 246 | |
| 247 | mark_dead() must not set the flag close() guards on, or every peer |
| 248 | disconnect leaks a DEALER and a monitor socket. |
| 249 | """ |
| 250 | transport, _ = _make_transport() |
| 251 | try: |
| 252 | conn = transport.connect(f"127.0.0.1:{_free_port()}") |
| 253 | dealer, monitor = conn._sockets.dealer, conn._sockets.monitor |
| 254 | |
| 255 | conn.mark_dead() |
| 256 | assert not conn.alive |
| 257 | assert not dealer.closed |
| 258 | |
| 259 | conn.close() |
| 260 | assert dealer.closed |
| 261 | assert monitor.closed |
| 262 | finally: |
| 263 | transport.close() |
| 264 | |
| 265 | def test_connect_retires_dead_connection(self): |
| 266 | """connect() replaces a registered-but-dead connection.""" |
| 267 | transport, _ = _make_transport() |
| 268 | try: |
| 269 | # The peer port is never bound — only the peer id matters here. |
| 270 | peer_id = f"127.0.0.1:{_free_port()}" |
| 271 | dead = transport.connect(peer_id) |
| 272 | dead.mark_dead() |
| 273 | |
| 274 | conn = transport.connect(peer_id) |
| 275 | |
| 276 | assert conn is not dead |
| 277 | assert conn.alive |
| 278 | assert transport._connections[peer_id] is conn |
| 279 | assert dead._sockets.dealer.closed |
| 280 | finally: |
| 281 | transport.close() |
| 282 | |
| 283 | def test_repeated_reconnect_to_same_peer(self): |
| 284 | """A flapping peer stays reconnectable. |
| 285 | |
| 286 | Monitor endpoints are inproc addresses that libzmq releases |
| 287 | asynchronously, so deriving one from peer_id alone makes each |
| 288 | reconnect race the previous teardown and fail with EADDRINUSE. |
| 289 | """ |
| 290 | transport, _ = _make_transport() |
| 291 | try: |
| 292 | peer_id = f"127.0.0.1:{_free_port()}" |
| 293 | for _ in range(10): |
| 294 | conn = transport.connect(peer_id) |
| 295 | conn.mark_dead() |
| 296 | transport.poll() |
| 297 | assert conn._sockets.dealer.closed |
| 298 | |
| 299 | assert not transport._connections |
| 300 | finally: |
| 301 | transport.close() |
| 302 | |
| 303 | def test_inbound_message_survives_dead_registration(self): |
| 304 | """A reconnecting peer's first message is not dropped. |
| 305 | |
| 306 | poll() must retire connections killed by their session before routing |
| 307 | traffic, otherwise the message is enqueued into the dead connection and |
| 308 | discarded when it is swept — and a session announces itself only once. |
| 309 | Covers only that between-polls window: a peer dying while poll() runs, |
| 310 | or dying silently until the heartbeat expires, is out of scope. |
| 311 | """ |
| 312 | transport_a, port_a = _make_transport() |
| 313 | transport_b, _ = _make_transport() |
| 314 | |
| 315 | try: |
| 316 | conn_b = transport_b.connect(f"127.0.0.1:{port_a}") |
| 317 | conn_b.send({"type": "connect", "seq": 1}) |
| 318 | |
| 319 | inbound = _wait_for_inbound(transport_a)[0] |
| 320 | _wait_for_messages(transport_a, inbound, 1) |
| 321 | |
| 322 | inbound.mark_dead() |
| 323 | conn_b.send({"type": "connect", "seq": 2}) |
| 324 | |
| 325 | # Wait until the frame is readable on the ROUTER, so the message is |
| 326 | # known to have arrived rather than merely being slow. |
| 327 | poller = zmq.Poller() |
| 328 | poller.register(transport_a._router, zmq.POLLIN) |
| 329 | assert poller.poll(2000), "message never reached the ROUTER" |
| 330 | |
| 331 | new_conns = _wait_for_inbound(transport_a) |
| 332 | assert len(new_conns) == 1 |
| 333 | assert new_conns[0] is not inbound |
| 334 | msgs = _wait_for_messages(transport_a, new_conns[0], 1) |
| 335 | assert msgs == [{"type": "connect", "seq": 2}] |
| 336 | finally: |
| 337 | transport_a.close() |
| 338 | transport_b.close() |