From 5ec47f3e48e7f6da9b6caa1c804b3887f832a788 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Lucchesi?= Date: Fri, 7 Aug 2026 11:43:29 +0200 Subject: [PATCH] [PD][PushConnector] Record last activity of remotes to allow clean up of stale ones (#50234) Signed-off-by: NickLucche --- .../unit/test_nixl_push_connector.py | 55 +++++++++++++++++++ .../kv_connector/v1/nixl/base_worker.py | 3 +- .../kv_connector/v1/nixl/push_worker.py | 3 + 3 files changed, 59 insertions(+), 2 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_nixl_push_connector.py b/tests/v1/kv_connector/unit/test_nixl_push_connector.py index 69887de4a0..a5e1eabad8 100644 --- a/tests/v1/kv_connector/unit/test_nixl_push_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_push_connector.py @@ -341,6 +341,7 @@ class _StubWriterWorker(NixlPushConnectorWorker): # Single non-hybrid attention group, matching the stub block id lists. w._has_mamba = False w._group_spec_types = (FullAttentionSpec,) + w._engine_last_active = {} # Track _do_start_push_kv invocations. calls: list[tuple[str, Any, dict[str, Any]]] = [] @@ -619,6 +620,60 @@ def test_writer_loop_drains_deferred_push_inbox(): assert w.start_push_calls[0][0] == "req-retry" +def _eviction_worker(engine_ttl: float) -> NixlPushConnectorWorker: + """A push worker wired to drive the base ``_ensure_handshake`` eviction + path (``_evict_stale_engines`` + ``_cleanup_remote_engine``).""" + w = _StubWriterWorker.fresh() + w._engine_ttl = engine_ttl + w._engine_last_active = {} + w._engine_clock_offset = {} + w._handshake_lock = threading.RLock() + # _cleanup_remote_engine touches these when reaping an engine. + w.nixl_wrapper = MagicMock() + w.dst_xfer_side_handles = {} + w.kv_caches_base_addr = {} + w.dst_num_blocks = {} + w.tp_mappings = {} + w.transfer_topo = None + w._logical_to_kernel_block_ids = lambda blocks, ratio: blocks + w.writes = [] + w._xfer_blocks_for_req = lambda req_id, meta: w.writes.append(req_id) + return w + + +def test_active_push_refreshes_engine_last_active(): + """The base ``_ensure_handshake`` refreshes liveness only on a new + handshake, so an active push to an already-connected engine must refresh + ``_engine_last_active`` itself or it is reaped mid-stream (S1).""" + w = _eviction_worker(engine_ttl=3600.0) + # Already-connected D -> _ensure_handshake returns None, WRITE runs inline. + w._remote_agents["decode-engine"] = {(0, 0): "agent-decode"} + stale = time.perf_counter() - 5.0 + w._engine_last_active["decode-engine"] = stale + + _real_do_start_push_kv(w, "req", ([1, 2, 3],), _registration_data("req")) + + assert w._engine_last_active["decode-engine"] > stale + assert w.writes == ["req"] + + +def test_stale_engine_evicted_on_push(): + """A push drives ``_ensure_handshake`` -> ``_evict_stale_engines``, so a D + engine silent past its TTL is reaped -- no _remote_agents / NIXL agent leak + across D scale up/down (S1).""" + w = _eviction_worker(engine_ttl=30.0) + w._remote_agents["D-old"] = {(0, 0): "agent-D-old"} + w._engine_last_active["D-old"] = time.perf_counter() - 10_000.0 + # The engine being pushed to is already connected. + w._remote_agents["decode-engine"] = {(0, 0): "agent-decode"} + + _real_do_start_push_kv(w, "req", ([1, 2, 3],), _registration_data("req")) + + assert "D-old" not in w._remote_agents + assert "D-old" not in w._engine_last_active + w.nixl_wrapper.remove_remote_agent.assert_called_once_with("agent-D-old") + + class TestPushWriterNotifs: def test_get_new_notifs_processes_forwarded_completion_notif(self): """Non-PUSH_REG notifs forwarded by the writer thread are drained diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py index 13f45dfcc4..e61b2a3921 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py @@ -2560,8 +2560,7 @@ class NixlBaseConnectorWorker: # Drop the cached clock offset; it is re-measured on the next handshake. self._engine_clock_offset.pop(engine_id, None) - # Push P-side engines are tracked in _remote_agents but not in - # _engine_last_active (they don't participate in stale eviction), so + # A just-completed handshake may not have recorded activity yet, so # tolerate a missing entry. last_active = self._engine_last_active.pop(engine_id, None) if log_eviction and last_active is not None: diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py index 13af0314f0..1050e35b7a 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py @@ -450,6 +450,9 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker): fut.add_done_callback(_on_handshake) return + # Keep the engine alive while it is actively receiving pushes, mirroring + # how pull-mode transfers touch _engine_last_active in start_load_kv. + self._engine_last_active[decode_engine_id] = time.perf_counter() # Both sides stay logical here; ``_xfer_blocks_for_req`` converts each # to physical with its own physical-blocks-per-logical ratio -- P uses