[PD][PushConnector] Record last activity of remotes on the D side (#52245)
Signed-off-by: Sunita Nadampalli <nadampal@amazon.com> Co-authored-by: Nicolò Lucchesi <nicolo.lucchesi@mistral.ai>
snadampal committed
Sep 24, 2026 at 18:15 UTC
6a86bf626e32555a21018050d57befca1a704344
2 files changed
+61
tests/v1/kv_connector/unit/test_nixl_push_connector.py
+55
@@ -856,6 +856,61 @@ def test_cleanup_remote_engine_pops_under_handshake_lock():
856
w.nixl_wrapper.remove_remote_agent.assert_called_once_with("agent-D-old")
857
858
859
+def _recv_metadata(remote_engine_id: str, req_id: str = "req") -> NixlConnectorMetadata:
860
+ """Build metadata with a single ``reqs_to_recv`` entry pointing at
861
+ ``remote_engine_id`` (the prefill remote D receives a push from)."""
862
+ params = {
863
+ "remote_block_ids": ([0],),
864
+ "remote_engine_id": remote_engine_id,
865
+ "remote_request_id": "p-req",
866
+ "remote_host": "localhost",
867
+ "remote_port": 1234,
868
+ "tp_size": 1,
869
+ }
870
+ meta = NixlConnectorMetadata()
871
+ meta.add_new_req_to_recv(req_id, ([0],), params)
872
+ return meta
873
+
874
+
875
+def test_start_load_kv_refreshes_engine_last_active_on_d_side():
876
+ """D receiving a push refreshes ``_engine_last_active`` for the prefill
877
+ remote, but only once that engine is connected (handshaked into
878
+ ``_remote_agents`` via heartbeats), mirroring the P-side refresh which runs
879
+ after ``_ensure_handshake``. Without it the prefill engine is reaped
880
+ mid-stream once it passes its TTL."""
881
+ w = _StubWriterWorker.fresh()
882
+ w._send_heartbeats = lambda metadata: None
883
+ w._logical_to_kernel_block_ids = lambda x, ratio: x
884
+ # Connected prefill engine (already handshaked in).
885
+ w._remote_agents["prefill-engine"] = {(0, 0): "agent-p"}
886
+
887
+ stale = time.perf_counter() - 5.0
888
+ w._engine_last_active["prefill-engine"] = stale
889
+
890
+ w.start_load_kv(_recv_metadata("prefill-engine"))
891
+
892
+ assert w._engine_last_active["prefill-engine"] > stale
893
+ assert "req" in w._recving_metadata
894
+
895
+
896
+def test_start_load_kv_skips_liveness_for_unconnected_engine():
897
+ """If the prefill engine was never handshaked into ``_remote_agents`` (e.g.
898
+ its heartbeat handshake never succeeded), the D-side refresh must NOT record
899
+ an ``_engine_last_active`` entry -- otherwise a later eviction pass hits the
900
+ ``_remote_agents`` invariant in ``_cleanup_remote_engine`` and crashes on
901
+ the D node."""
902
+ w = _eviction_worker(engine_ttl=30.0)
903
+ w._send_heartbeats = lambda metadata: None
904
+
905
+ w.start_load_kv(_recv_metadata("prefill-engine"))
906
+
907
+ # Unconnected -> no orphaned liveness entry, so eviction has nothing to trip on.
908
+ assert "prefill-engine" not in w._engine_last_active
909
+ assert "prefill-engine" not in w._remote_agents
910
+
911
+ w._evict_stale_engines() # nothing to reap, must not raise
912
+
913
+
914
class TestPushWriterNotifs:
915
def test_get_new_notifs_processes_forwarded_completion_notif(self):
916
"""Non-PUSH_REG notifs forwarded by the writer thread are drained
vllm/distributed/kv_transfer/kv_connector/v1/nixl/push_worker.py
+6
@@ -180,6 +180,12 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker):
180
)
181
assert meta.remote is not None
182
remote_engine_id = meta.remote.engine_id
183
+ # Update last activity from this remote (same as pull), but only for
184
+ # an already-connected engine so we never leave an _engine_last_active
185
+ # entry without a _remote_agents entry (breaks _cleanup_remote_engine).
186
+ if remote_engine_id in self._remote_agents:
187
+ self._engine_last_active[remote_engine_id] = time.perf_counter()
188
+
189
logger.debug(
190
"start_load_kv (push) for request %s from remote engine %s. "
191
"Num local_block_ids: %s. Num remote_block_ids: %s. ",