[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>
Jason committed
Jul 27, 2026 at 01:07 UTC
53397fbfac4dcd421b9d467199068518b135bc34
3 files changed
+167
-13
tests/v1/kv_offload/tiering/p2p/test_zmq_transport.py
+108
@@ -219,6 +219,7 @@ class TestZmqTransportConnectivity:
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()
@@ -228,3 +229,110 @@ class TestZmqTransportConnectivity:
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()
vllm/v1/kv_offload/tiering/p2p/control/base.py
+5
@@ -97,6 +97,7 @@ class ControlConnection(ABC):
97
98
After this call, alive returns False. The session should
99
stop using this connection and the transport will clean it up.
100
+ Resources are not released here — close() must still run.
101
"""
102
...
103
@@ -134,6 +135,10 @@ class ControlTransport(ABC):
135
136
The connection's send queue is live immediately — messages
137
sent before the remote peer's poll() will be buffered.
138
+
139
+ A connection to peer_id that is registered but no longer alive is
140
+ retired and replaced; a live one is a duplicate and must not be
141
+ replaced.
142
"""
143
...
144
vllm/v1/kv_offload/tiering/p2p/control/zmq.py
+54
-13
@@ -55,12 +55,16 @@ class ZmqConnection(ControlConnection):
55
def __init__(self, peer_id: str, sockets: _Sockets) -> None:
56
super().__init__(peer_id)
57
self._sockets = sockets
58
+ # _dead: the peer is gone. _closed: the sockets have been released.
59
+ # Distinct, so mark_dead() cannot turn close() into a no-op and leak
60
+ # the DEALER and its monitor socket.
61
+ self._dead = False
62
self._closed = False
63
self._inbox: list[dict] = []
64
65
def send(self, msg: dict) -> None:
66
"""Send a msgpack-encoded message to this peer."""
63
- if self._closed:
67
+ if not self.alive:
68
raise RuntimeError(
69
f"ZmqConnection: send on closed connection to {self.peer_id}"
70
)
@@ -77,12 +81,13 @@ class ZmqConnection(ControlConnection):
81
82
@property
83
def alive(self) -> bool:
80
- return not self._closed
84
+ return not (self._dead or self._closed)
85
86
def close(self) -> None:
87
if self._closed:
88
return
89
self._closed = True
90
+ self._dead = True
91
logger.info("ZmqConnection: closing connection to %s", self.peer_id)
92
self._sockets.monitor.close()
93
self._sockets.dealer.close()
@@ -92,8 +97,12 @@ class ZmqConnection(ControlConnection):
97
self._inbox.append(msg)
98
99
def mark_dead(self) -> None:
95
- """Mark connection as disconnected."""
96
- self._closed = True
100
+ """Mark connection as disconnected.
101
+
102
+ Only flips liveness: the sockets stay open until close() releases
103
+ them, so the owner can still drain recv() before tearing down.
104
+ """
105
+ self._dead = True
106
107
@property
108
def monitor_socket(self) -> zmq.Socket:
@@ -114,6 +123,9 @@ class ZmqTransport(ControlTransport):
123
124
self._connections: dict[str, ZmqConnection] = {}
125
self._pending_inbound: list[tuple[str, dict]] = []
126
+ # Monotonic suffix for inproc monitor endpoints — see
127
+ # _open_connection() for why peer_id alone is not enough.
128
+ self._monitor_seq = 0
129
130
self._zmq_ctx = zmq.Context()
131
self._router: zmq.Socket = self._zmq_ctx.socket(zmq.ROUTER)
@@ -127,10 +139,23 @@ class ZmqTransport(ControlTransport):
139
# ------------------------------------------------------------------
140
141
def connect(self, peer_id: str) -> ZmqConnection:
130
- """Open an outbound connection to a remote peer."""
131
- assert peer_id not in self._connections, (
132
- f"ZmqConnection to {peer_id} already exists"
133
- )
142
+ """Open an outbound connection to a remote peer.
143
+
144
+ A dead connection can still be registered: its owning session may mark
145
+ it dead after this tick's sweep already ran, and poll() only
146
+ unregisters it on the next pass. Retire such an entry instead of
147
+ asserting, so a reconnect landing in that window succeeds. A live
148
+ entry is still a genuine duplicate.
149
+ """
150
+ existing = self._connections.get(peer_id)
151
+ if existing is not None:
152
+ assert not existing.alive, f"ZmqConnection to {peer_id} already exists"
153
+ logger.info(
154
+ "ZmqTransport %s: retiring dead connection to %s before reconnect",
155
+ self._local_id,
156
+ peer_id,
157
+ )
158
+ self._connections.pop(peer_id).close()
159
logger.info(
160
"ZmqTransport %s: opening OUTBOUND connection to %s",
161
self._local_id,
@@ -146,8 +171,18 @@ class ZmqTransport(ControlTransport):
171
- Checks monitors for disconnections
172
- Removes and closes dead connections
173
"""
174
+ # Retire connections a session killed since the last poll() before
175
+ # routing traffic: otherwise a reconnecting peer's first message is
176
+ # enqueued into the dead connection and discarded along with it.
177
+ # This closes only that between-polls window. A peer that dies while
178
+ # poll() is running is not seen until the next _check_monitors(), and
179
+ # one that dies silently not until the ZMQ heartbeat expires; both
180
+ # still take a message into a doomed connection and are out of scope.
181
+ self._sweep_dead_connections()
182
+
183
self._recv_router()
184
self._check_monitors()
185
+ self._sweep_dead_connections()
186
187
# Create connections for new inbound peers
188
new_connections: list[ControlConnection] | None = None
@@ -166,10 +201,6 @@ class ZmqTransport(ControlTransport):
201
conn.enqueue(msg)
202
self._pending_inbound.clear()
203
169
- # Remove dead connections
170
- for pid in [p for p, c in self._connections.items() if not c.alive]:
171
- self._connections.pop(pid).close()
172
-
204
return (
205
new_connections if new_connections is not None else _EMPTY_NEW_CONNECTIONS
206
)
@@ -210,8 +241,13 @@ class ZmqTransport(ControlTransport):
241
_apply_heartbeat(dealer)
242
dealer.identity = self._local_id.encode()
243
244
+ # Unique per connection, not per peer: libzmq releases an inproc
245
+ # endpoint on its reaper thread after the DEALER's close() has already
246
+ # returned, so reusing the peer-derived address on a reconnect races
247
+ # that teardown and fails with EADDRINUSE.
248
safe_id = peer_id.replace(":", "-").replace("/", "-")
214
- monitor_addr = f"inproc://p2p-monitor-{safe_id}"
249
+ monitor_addr = f"inproc://p2p-monitor-{safe_id}-{self._monitor_seq}"
250
+ self._monitor_seq += 1
251
dealer.monitor(monitor_addr, zmq.EVENT_DISCONNECTED)
252
253
monitor_sock = self._zmq_ctx.socket(zmq.PAIR)
@@ -231,6 +267,11 @@ class ZmqTransport(ControlTransport):
267
)
268
return conn
269
270
+ def _sweep_dead_connections(self) -> None:
271
+ """Unregister and release every connection that is no longer alive."""
272
+ for pid in [p for p, c in self._connections.items() if not c.alive]:
273
+ self._connections.pop(pid).close()
274
+
275
def _recv_router(self) -> None:
276
"""Non-blocking: receive all pending messages from ROUTER."""
277
while True: