diff --git a/tests/v1/kv_offload/tiering/p2p/test_zmq_transport.py b/tests/v1/kv_offload/tiering/p2p/test_zmq_transport.py index 101d4291cfd..f814a04887c 100644 --- a/tests/v1/kv_offload/tiering/p2p/test_zmq_transport.py +++ b/tests/v1/kv_offload/tiering/p2p/test_zmq_transport.py @@ -219,6 +219,7 @@ class TestZmqTransportConnectivity: # Pruning is synchronous within poll(). transport_a.poll() assert len(transport_a._connections) == 0 + assert new_conns[0]._sockets.dealer.closed finally: transport_a.close() transport_b.close() @@ -228,3 +229,110 @@ class TestZmqTransportConnectivity: transport, _ = _make_transport() transport.close() transport.close() # should not raise + + +class TestZmqReconnect: + """Reconnecting to a peer whose connection died (real ZMQ sockets). + + A session marks its connection dead while handling messages, which happens + after the transport's own sweep has run for that tick — so a dead + connection stays registered until the next poll(). Reconnecting in that + window must succeed, and the retired connection must release its sockets. + Real sockets are required: a mock reports every attribute as closed. + """ + + def test_close_after_mark_dead_releases_sockets(self): + """close() releases sockets even when mark_dead() ran first. + + mark_dead() must not set the flag close() guards on, or every peer + disconnect leaks a DEALER and a monitor socket. + """ + transport, _ = _make_transport() + try: + conn = transport.connect(f"127.0.0.1:{_free_port()}") + dealer, monitor = conn._sockets.dealer, conn._sockets.monitor + + conn.mark_dead() + assert not conn.alive + assert not dealer.closed + + conn.close() + assert dealer.closed + assert monitor.closed + finally: + transport.close() + + def test_connect_retires_dead_connection(self): + """connect() replaces a registered-but-dead connection.""" + transport, _ = _make_transport() + try: + # The peer port is never bound — only the peer id matters here. + peer_id = f"127.0.0.1:{_free_port()}" + dead = transport.connect(peer_id) + dead.mark_dead() + + conn = transport.connect(peer_id) + + assert conn is not dead + assert conn.alive + assert transport._connections[peer_id] is conn + assert dead._sockets.dealer.closed + finally: + transport.close() + + def test_repeated_reconnect_to_same_peer(self): + """A flapping peer stays reconnectable. + + Monitor endpoints are inproc addresses that libzmq releases + asynchronously, so deriving one from peer_id alone makes each + reconnect race the previous teardown and fail with EADDRINUSE. + """ + transport, _ = _make_transport() + try: + peer_id = f"127.0.0.1:{_free_port()}" + for _ in range(10): + conn = transport.connect(peer_id) + conn.mark_dead() + transport.poll() + assert conn._sockets.dealer.closed + + assert not transport._connections + finally: + transport.close() + + def test_inbound_message_survives_dead_registration(self): + """A reconnecting peer's first message is not dropped. + + poll() must retire connections killed by their session before routing + traffic, otherwise the message is enqueued into the dead connection and + discarded when it is swept — and a session announces itself only once. + Covers only that between-polls window: a peer dying while poll() runs, + or dying silently until the heartbeat expires, is out of scope. + """ + transport_a, port_a = _make_transport() + transport_b, _ = _make_transport() + + try: + conn_b = transport_b.connect(f"127.0.0.1:{port_a}") + conn_b.send({"type": "connect", "seq": 1}) + + inbound = _wait_for_inbound(transport_a)[0] + _wait_for_messages(transport_a, inbound, 1) + + inbound.mark_dead() + conn_b.send({"type": "connect", "seq": 2}) + + # Wait until the frame is readable on the ROUTER, so the message is + # known to have arrived rather than merely being slow. + poller = zmq.Poller() + poller.register(transport_a._router, zmq.POLLIN) + assert poller.poll(2000), "message never reached the ROUTER" + + new_conns = _wait_for_inbound(transport_a) + assert len(new_conns) == 1 + assert new_conns[0] is not inbound + msgs = _wait_for_messages(transport_a, new_conns[0], 1) + assert msgs == [{"type": "connect", "seq": 2}] + finally: + transport_a.close() + transport_b.close() diff --git a/vllm/v1/kv_offload/tiering/p2p/control/base.py b/vllm/v1/kv_offload/tiering/p2p/control/base.py index 6b4d4cfcb59..7016cfb162a 100644 --- a/vllm/v1/kv_offload/tiering/p2p/control/base.py +++ b/vllm/v1/kv_offload/tiering/p2p/control/base.py @@ -97,6 +97,7 @@ class ControlConnection(ABC): After this call, alive returns False. The session should stop using this connection and the transport will clean it up. + Resources are not released here — close() must still run. """ ... @@ -134,6 +135,10 @@ class ControlTransport(ABC): The connection's send queue is live immediately — messages sent before the remote peer's poll() will be buffered. + + A connection to peer_id that is registered but no longer alive is + retired and replaced; a live one is a duplicate and must not be + replaced. """ ... diff --git a/vllm/v1/kv_offload/tiering/p2p/control/zmq.py b/vllm/v1/kv_offload/tiering/p2p/control/zmq.py index 4e9069a471c..42dc2c69846 100644 --- a/vllm/v1/kv_offload/tiering/p2p/control/zmq.py +++ b/vllm/v1/kv_offload/tiering/p2p/control/zmq.py @@ -55,12 +55,16 @@ class ZmqConnection(ControlConnection): def __init__(self, peer_id: str, sockets: _Sockets) -> None: super().__init__(peer_id) self._sockets = sockets + # _dead: the peer is gone. _closed: the sockets have been released. + # Distinct, so mark_dead() cannot turn close() into a no-op and leak + # the DEALER and its monitor socket. + self._dead = False self._closed = False self._inbox: list[dict] = [] def send(self, msg: dict) -> None: """Send a msgpack-encoded message to this peer.""" - if self._closed: + if not self.alive: raise RuntimeError( f"ZmqConnection: send on closed connection to {self.peer_id}" ) @@ -77,12 +81,13 @@ class ZmqConnection(ControlConnection): @property def alive(self) -> bool: - return not self._closed + return not (self._dead or self._closed) def close(self) -> None: if self._closed: return self._closed = True + self._dead = True logger.info("ZmqConnection: closing connection to %s", self.peer_id) self._sockets.monitor.close() self._sockets.dealer.close() @@ -92,8 +97,12 @@ class ZmqConnection(ControlConnection): self._inbox.append(msg) def mark_dead(self) -> None: - """Mark connection as disconnected.""" - self._closed = True + """Mark connection as disconnected. + + Only flips liveness: the sockets stay open until close() releases + them, so the owner can still drain recv() before tearing down. + """ + self._dead = True @property def monitor_socket(self) -> zmq.Socket: @@ -114,6 +123,9 @@ class ZmqTransport(ControlTransport): self._connections: dict[str, ZmqConnection] = {} self._pending_inbound: list[tuple[str, dict]] = [] + # Monotonic suffix for inproc monitor endpoints — see + # _open_connection() for why peer_id alone is not enough. + self._monitor_seq = 0 self._zmq_ctx = zmq.Context() self._router: zmq.Socket = self._zmq_ctx.socket(zmq.ROUTER) @@ -127,10 +139,23 @@ class ZmqTransport(ControlTransport): # ------------------------------------------------------------------ def connect(self, peer_id: str) -> ZmqConnection: - """Open an outbound connection to a remote peer.""" - assert peer_id not in self._connections, ( - f"ZmqConnection to {peer_id} already exists" - ) + """Open an outbound connection to a remote peer. + + A dead connection can still be registered: its owning session may mark + it dead after this tick's sweep already ran, and poll() only + unregisters it on the next pass. Retire such an entry instead of + asserting, so a reconnect landing in that window succeeds. A live + entry is still a genuine duplicate. + """ + existing = self._connections.get(peer_id) + if existing is not None: + assert not existing.alive, f"ZmqConnection to {peer_id} already exists" + logger.info( + "ZmqTransport %s: retiring dead connection to %s before reconnect", + self._local_id, + peer_id, + ) + self._connections.pop(peer_id).close() logger.info( "ZmqTransport %s: opening OUTBOUND connection to %s", self._local_id, @@ -146,8 +171,18 @@ class ZmqTransport(ControlTransport): - Checks monitors for disconnections - Removes and closes dead connections """ + # Retire connections a session killed since the last poll() before + # routing traffic: otherwise a reconnecting peer's first message is + # enqueued into the dead connection and discarded along with it. + # This closes only that between-polls window. A peer that dies while + # poll() is running is not seen until the next _check_monitors(), and + # one that dies silently not until the ZMQ heartbeat expires; both + # still take a message into a doomed connection and are out of scope. + self._sweep_dead_connections() + self._recv_router() self._check_monitors() + self._sweep_dead_connections() # Create connections for new inbound peers new_connections: list[ControlConnection] | None = None @@ -166,10 +201,6 @@ class ZmqTransport(ControlTransport): conn.enqueue(msg) self._pending_inbound.clear() - # Remove dead connections - for pid in [p for p, c in self._connections.items() if not c.alive]: - self._connections.pop(pid).close() - return ( new_connections if new_connections is not None else _EMPTY_NEW_CONNECTIONS ) @@ -210,8 +241,13 @@ class ZmqTransport(ControlTransport): _apply_heartbeat(dealer) dealer.identity = self._local_id.encode() + # Unique per connection, not per peer: libzmq releases an inproc + # endpoint on its reaper thread after the DEALER's close() has already + # returned, so reusing the peer-derived address on a reconnect races + # that teardown and fails with EADDRINUSE. safe_id = peer_id.replace(":", "-").replace("/", "-") - monitor_addr = f"inproc://p2p-monitor-{safe_id}" + monitor_addr = f"inproc://p2p-monitor-{safe_id}-{self._monitor_seq}" + self._monitor_seq += 1 dealer.monitor(monitor_addr, zmq.EVENT_DISCONNECTED) monitor_sock = self._zmq_ctx.socket(zmq.PAIR) @@ -231,6 +267,11 @@ class ZmqTransport(ControlTransport): ) return conn + def _sweep_dead_connections(self) -> None: + """Unregister and release every connection that is no longer alive.""" + for pid in [p for p, c in self._connections.items() if not c.alive]: + self._connections.pop(pid).close() + def _recv_router(self) -> None: """Non-blocking: receive all pending messages from ROUTER.""" while True: