[Bugfix][NIXL] Don't evict a remote engine a transfer is still reading from (#54689)
Signed-off-by: Yizheng Jiao <jyizheng@gmail.com>
Yizheng Jiao committed
Sep 11, 2026 at 18:06 UTC
56d72701380bc68955cc0aa9367636df8356e3c8
2 files changed
+61
-2
tests/v1/kv_connector/unit/test_nixl_connector.py
+42
@@ -2444,6 +2444,48 @@ def test_engine_ttl_eviction(default_vllm_config, dist_init):
2444
mock_rem.assert_any_call("agent_1")
2445
2446
2447
+@patch(
2448
+ "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
2449
+ FakeNixlWrapper,
2450
+)
2451
+def test_engine_with_inflight_transfer_is_not_evicted(default_vllm_config, dist_init):
2452
+ """A transfer that outlives the TTL must keep its engine registered.
2453
+
2454
+ _engine_last_active is stamped when a read is issued and not refreshed
2455
+ while it runs, so a stalled transfer -- a peer that has lost its NIC holds
2456
+ one indefinitely -- leaves its engine looking idle. Evicting it releases
2457
+ the dlist handle and remote agent the transfer is still reading through.
2458
+ """
2459
+ worker, engine_id = _setup_worker_with_remote_engine(engine_ttl=10.0)
2460
+ nixl_wrapper = worker.nixl_wrapper
2461
+
2462
+ request_id = "req-still-reading"
2463
+ worker._recving_transfers[request_id] = [MagicMock()]
2464
+ worker._recving_metadata[request_id] = MagicMock(
2465
+ remote=MagicMock(engine_id=engine_id)
2466
+ )
2467
+
2468
+ with (
2469
+ patch.object(nixl_wrapper, "release_dlist_handle") as mock_rel,
2470
+ patch.object(nixl_wrapper, "remove_remote_agent") as mock_rem,
2471
+ ):
2472
+ worker._engine_last_active[engine_id] = time.perf_counter() - 20.0
2473
+
2474
+ worker._evict_stale_engines()
2475
+
2476
+ assert engine_id in worker._remote_agents
2477
+ assert engine_id in worker.dst_xfer_side_handles
2478
+ mock_rel.assert_not_called()
2479
+ mock_rem.assert_not_called()
2480
+
2481
+ # Once the transfer is done the engine is stale like any other.
2482
+ del worker._recving_transfers[request_id]
2483
+ worker._evict_stale_engines()
2484
+
2485
+ assert engine_id not in worker._remote_agents
2486
+ assert mock_rem.call_count == 2
2487
+
2488
+
2489
@patch(
2490
"vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
2491
FakeNixlWrapper,
vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py
+19
-2
@@ -3143,7 +3143,9 @@ class NixlBaseConnectorWorker:
3143
to be "optimal" until a new handshake is performed.
3144
3145
Engines with active transfers or pending handshakes cannot be stale:
3146
- - Active transfers touch _engine_last_active in start_load_kv.
3146
+ - Reads stamp _engine_last_active when they are issued, and engines a
3147
+ transfer is still reading from are held back explicitly, since the
3148
+ stamp is not refreshed while the read runs.
3149
- Pending handshakes don't have an _engine_last_active entry yet
3150
"""
3151
# NOTE (NickLucche): This does NOT currently prevent OOMing if a huge number
@@ -3154,10 +3156,25 @@ class NixlBaseConnectorWorker:
3156
return
3157
3158
now = time.perf_counter()
3159
+ busy = self._engines_with_inflight_transfers()
3160
for eid, last_active in list(self._engine_last_active.items()):
3158
- if now - last_active > self._engine_ttl:
3161
+ if now - last_active > self._engine_ttl and eid not in busy:
3162
self._cleanup_remote_engine(eid)
3163
3164
+ def _engines_with_inflight_transfers(self) -> set[EngineId]:
3165
+ """Remote engines a transfer is still reading from.
3166
+
3167
+ The timestamp is stamped when a read is issued and not refreshed while
3168
+ it runs, so a transfer that outlives the TTL leaves its engine looking
3169
+ idle. A peer that has lost its NIC holds one indefinitely.
3170
+ """
3171
+ return {
3172
+ meta.remote.engine_id
3173
+ for req_id in self._recving_transfers
3174
+ if (meta := self._recving_metadata.get(req_id)) is not None
3175
+ and meta.remote is not None
3176
+ }
3177
+
3178
def _cleanup_remote_engine(
3179
self, engine_id: EngineId, *, log_eviction: bool = True
3180
) -> None: