[Bugfix][KV Offload][P2P] Fix EngineCore crash reconnecting to a reaped peer (#49823)

Signed-off-by: Jason Yao <wsyjh8@gmail.com>
Co-authored-by: Or Ozeri <oro@il.ibm.com>

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()

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.
         """
         ...
 

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: