[Bugfix][NixlPush] Guard _remote_agents read in _do_send_reg_notif (X1) (#56317)
Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai>
Nicolò Lucchesi committed
Sep 14, 2026 at 12:37 UTC
0ca7aef4e341cb4f4d7233291c304a168f4384aa
3 files changed
+111
-2
tests/v1/kv_connector/unit/test_nixl_push_connector.py
+104
@@ -362,6 +362,7 @@ class _StubWriterWorker(NixlPushConnectorWorker):
362
w.world_size = 1
363
w.engine_id = "test-decode-engine"
364
w._remote_agents = {}
365
+ w._handshake_lock = threading.RLock()
366
w._physical_blocks_per_logical_kv_block = 1
367
w._uses_region_group_mapping = False
368
w.region_group_ids = [0]
@@ -418,6 +419,84 @@ def _registration_data(
419
}
420
421
422
+class TestPushWriterRegSend:
423
+ """``_do_send_reg_notif``: the PUSH_REG goes to every handshaked agent of
424
+ the P engine, and the agent map is read under ``_handshake_lock`` so the
425
+ send can't race the handshake done-callback's insert or a concurrent
426
+ stale-engine eviction (X1)."""
427
+
428
+ def _reg_send_worker(self) -> tuple[_StubWriterWorker, list]:
429
+ w = _StubWriterWorker.fresh()
430
+ w.nixl_wrapper = MagicMock()
431
+ failures: list = []
432
+ w._handle_failed_transfer = lambda *a: failures.append(a)
433
+ return w, failures
434
+
435
+ def test_sends_to_all_agents_with_prefix_and_payload(self):
436
+ w, failures = self._reg_send_worker()
437
+ w._remote_agents["prefill-engine"] = {(0, 0): "agent-0", (0, 1): "agent-1"}
438
+ rd = _registration_data("req-1")
439
+
440
+ w._do_send_reg_notif("req-1", rd)
441
+
442
+ assert failures == []
443
+ calls = w.nixl_wrapper.send_notif.call_args_list
444
+ assert len(calls) == 2
445
+ assert {c.args[0] for c in calls} == {"agent-0", "agent-1"}
446
+ for call in calls:
447
+ notif = call.kwargs["notif_msg"]
448
+ assert notif.startswith(PUSH_REG_NOTIF_PREFIX)
449
+ # (Tuples decode back as lists, so compare identity fields.)
450
+ decoded = msgspec.msgpack.decode(notif[len(PUSH_REG_NOTIF_PREFIX) :])
451
+ assert decoded["request_id"] == "req-1"
452
+ assert decoded["decode_engine_id"] == rd["decode_engine_id"]
453
+ assert decoded["remote_engine_id"] == rd["remote_engine_id"]
454
+ assert decoded["local_block_ids"] == [
455
+ list(group) for group in rd["local_block_ids"]
456
+ ]
457
+
458
+ def test_fails_request_when_engine_not_handshaked(self):
459
+ w, failures = self._reg_send_worker()
460
+
461
+ w._do_send_reg_notif("req-1", _registration_data("req-1"))
462
+
463
+ assert failures == [("req-1", None)]
464
+ w.nixl_wrapper.send_notif.assert_not_called()
465
+
466
+ def test_agent_read_is_serialized_with_handshake_lock(self):
467
+ """While another thread holds ``_handshake_lock`` (as the handshake
468
+ done-callback does while writing ``_remote_agents``), the writer's
469
+ check-then-send must not run ahead of the critical section."""
470
+ w, failures = self._reg_send_worker()
471
+ w._remote_agents["prefill-engine"] = {(0, 0): "agent-0"}
472
+ rd = _registration_data("req-1")
473
+
474
+ w._handshake_lock.acquire()
475
+ done = threading.Event()
476
+ errors: list[BaseException] = []
477
+
478
+ def _send():
479
+ try:
480
+ w._do_send_reg_notif("req-1", rd)
481
+ except BaseException as e: # pragma: no cover
482
+ errors.append(e)
483
+ finally:
484
+ done.set()
485
+
486
+ t = threading.Thread(target=_send)
487
+ t.start()
488
+ assert not done.wait(timeout=0.2)
489
+ w.nixl_wrapper.send_notif.assert_not_called()
490
+ assert failures == []
491
+
492
+ w._handshake_lock.release()
493
+ assert done.wait(timeout=2.0)
494
+ t.join(timeout=2.0)
495
+ assert errors == []
496
+ assert failures == []
497
+ assert w.nixl_wrapper.send_notif.call_count == 1
498
+
499
+
500
class TestPushWriterMatching:
501
def test_handle_push_reg_matches_existing_finished_blocks(self):
502
"""PUSH_REG arrives second (P finished first): match + fire."""
@@ -727,6 +806,31 @@ def test_stale_engine_evicted_on_push():
806
w.nixl_wrapper.remove_remote_agent.assert_called_once_with("agent-D-old")
807
808
809
+def test_cleanup_remote_engine_pops_under_handshake_lock():
810
+ """Regression (X1): eviction must not remove ``_remote_agents`` entries
811
+ while another thread is inside the lock's critical section (e.g. the
812
+ push writer reading the map to send a registration)."""
813
+ w = _eviction_worker(engine_ttl=30.0)
814
+ w._remote_agents["D-old"] = {(0, 0): "agent-D-old"}
815
+
816
+ w._handshake_lock.acquire()
817
+ done = threading.Event()
818
+
819
+ def _cleanup():
820
+ w._cleanup_remote_engine("D-old")
821
+ done.set()
822
+
823
+ t = threading.Thread(target=_cleanup)
824
+ t.start()
825
+ assert "D-old" in w._remote_agents
826
+
827
+ w._handshake_lock.release()
828
+ assert done.wait(timeout=2.0)
829
+ t.join(timeout=2.0)
830
+ assert "D-old" not in w._remote_agents
831
+ w.nixl_wrapper.remove_remote_agent.assert_called_once_with("agent-D-old")
832
+
833
+
834
class TestPushWriterNotifs:
835
def test_get_new_notifs_processes_forwarded_completion_notif(self):
836
"""Non-PUSH_REG notifs forwarded by the writer thread are drained
vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py
+4
-1
@@ -3264,7 +3264,10 @@ class NixlBaseConnectorWorker:
3264
# Notif-only engines (push-mode D side) have no descriptor state.
3265
for handle in self.dst_xfer_side_handles.pop(engine_id, {}).values():
3266
self.nixl_wrapper.release_dlist_handle(handle)
3267
- for agent_name in self._remote_agents.pop(engine_id).values():
3267
+ # Pop under the handshake lock; NIXL teardown stays outside it.
3268
+ with self._handshake_lock:
3269
+ agents = self._remote_agents.pop(engine_id)
3270
+ for agent_name in agents.values():
3271
self.nixl_wrapper.remove_remote_agent(agent_name)
3272
3273
self.kv_caches_base_addr.pop(engine_id, None)
vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py
+3
-1
@@ -354,7 +354,9 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker):
354
def _do_send_reg_notif(self, req_id: str, reg_data: dict[str, Any]) -> None:
355
engine_id = reg_data["remote_engine_id"]
356
notif_msg = PUSH_REG_NOTIF_PREFIX + msgspec.msgpack.encode(reg_data)
357
- agents = self._remote_agents.get(engine_id)
357
+ # _remote_agents is mutated on other threads; snapshot under the lock.
358
+ with self._handshake_lock:
359
+ agents = dict(self._remote_agents.get(engine_id) or {})
360
if not agents:
361
logger.error(
362
"No remote agents for engine %s; cannot send registration for %s",