[PD][PushConnector] Record last activity of remotes to allow clean up of stale ones (#50234)
Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai>
Nicolò Lucchesi committed
Aug 7, 2026 at 11:43 UTC
5ec47f3e48e7f6da9b6caa1c804b3887f832a788
3 files changed
+59
-2
tests/v1/kv_connector/unit/test_nixl_push_connector.py
+55
@@ -341,6 +341,7 @@ class _StubWriterWorker(NixlPushConnectorWorker):
341
# Single non-hybrid attention group, matching the stub block id lists.
342
w._has_mamba = False
343
w._group_spec_types = (FullAttentionSpec,)
344
+ w._engine_last_active = {}
345
346
# Track _do_start_push_kv invocations.
347
calls: list[tuple[str, Any, dict[str, Any]]] = []
@@ -619,6 +620,60 @@ def test_writer_loop_drains_deferred_push_inbox():
620
assert w.start_push_calls[0][0] == "req-retry"
621
622
623
+def _eviction_worker(engine_ttl: float) -> NixlPushConnectorWorker:
624
+ """A push worker wired to drive the base ``_ensure_handshake`` eviction
625
+ path (``_evict_stale_engines`` + ``_cleanup_remote_engine``)."""
626
+ w = _StubWriterWorker.fresh()
627
+ w._engine_ttl = engine_ttl
628
+ w._engine_last_active = {}
629
+ w._engine_clock_offset = {}
630
+ w._handshake_lock = threading.RLock()
631
+ # _cleanup_remote_engine touches these when reaping an engine.
632
+ w.nixl_wrapper = MagicMock()
633
+ w.dst_xfer_side_handles = {}
634
+ w.kv_caches_base_addr = {}
635
+ w.dst_num_blocks = {}
636
+ w.tp_mappings = {}
637
+ w.transfer_topo = None
638
+ w._logical_to_kernel_block_ids = lambda blocks, ratio: blocks
639
+ w.writes = []
640
+ w._xfer_blocks_for_req = lambda req_id, meta: w.writes.append(req_id)
641
+ return w
642
+
643
+
644
+def test_active_push_refreshes_engine_last_active():
645
+ """The base ``_ensure_handshake`` refreshes liveness only on a new
646
+ handshake, so an active push to an already-connected engine must refresh
647
+ ``_engine_last_active`` itself or it is reaped mid-stream (S1)."""
648
+ w = _eviction_worker(engine_ttl=3600.0)
649
+ # Already-connected D -> _ensure_handshake returns None, WRITE runs inline.
650
+ w._remote_agents["decode-engine"] = {(0, 0): "agent-decode"}
651
+ stale = time.perf_counter() - 5.0
652
+ w._engine_last_active["decode-engine"] = stale
653
+
654
+ _real_do_start_push_kv(w, "req", ([1, 2, 3],), _registration_data("req"))
655
+
656
+ assert w._engine_last_active["decode-engine"] > stale
657
+ assert w.writes == ["req"]
658
+
659
+
660
+def test_stale_engine_evicted_on_push():
661
+ """A push drives ``_ensure_handshake`` -> ``_evict_stale_engines``, so a D
662
+ engine silent past its TTL is reaped -- no _remote_agents / NIXL agent leak
663
+ across D scale up/down (S1)."""
664
+ w = _eviction_worker(engine_ttl=30.0)
665
+ w._remote_agents["D-old"] = {(0, 0): "agent-D-old"}
666
+ w._engine_last_active["D-old"] = time.perf_counter() - 10_000.0
667
+ # The engine being pushed to is already connected.
668
+ w._remote_agents["decode-engine"] = {(0, 0): "agent-decode"}
669
+
670
+ _real_do_start_push_kv(w, "req", ([1, 2, 3],), _registration_data("req"))
671
+
672
+ assert "D-old" not in w._remote_agents
673
+ assert "D-old" not in w._engine_last_active
674
+ w.nixl_wrapper.remove_remote_agent.assert_called_once_with("agent-D-old")
675
+
676
+
677
class TestPushWriterNotifs:
678
def test_get_new_notifs_processes_forwarded_completion_notif(self):
679
"""Non-PUSH_REG notifs forwarded by the writer thread are drained
vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py
+1
-2
@@ -2560,8 +2560,7 @@ class NixlBaseConnectorWorker:
2560
2561
# Drop the cached clock offset; it is re-measured on the next handshake.
2562
self._engine_clock_offset.pop(engine_id, None)
2563
- # Push P-side engines are tracked in _remote_agents but not in
2564
- # _engine_last_active (they don't participate in stale eviction), so
2563
+ # A just-completed handshake may not have recorded activity yet, so
2564
# tolerate a missing entry.
2565
last_active = self._engine_last_active.pop(engine_id, None)
2566
if log_eviction and last_active is not None:
vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py
+3
@@ -450,6 +450,9 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker):
450
451
fut.add_done_callback(_on_handshake)
452
return
453
+ # Keep the engine alive while it is actively receiving pushes, mirroring
454
+ # how pull-mode transfers touch _engine_last_active in start_load_kv.
455
+ self._engine_last_active[decode_engine_id] = time.perf_counter()
456
457
# Both sides stay logical here; ``_xfer_blocks_for_req`` converts each
458
# to physical with its own physical-blocks-per-logical ratio -- P uses