[Bugfix][NIXL] Don't assert when a failed transfer is cleaned up twice (#54518)
Signed-off-by: Yizheng Jiao <jyizheng@gmail.com> Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai> Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: NickLucche <nicolo.lucchesi@mistral.ai> Co-authored-by: Nick Hill <nickhill123@gmail.com>
Yizheng Jiao committed
Sep 4, 2026 at 16:20 UTC
eb74fbb3e7f0a917e93fff90ef68da23dbd2b34b
2 files changed
+94
-5
tests/v1/kv_connector/unit/test_nixl_connector.py
+81
@@ -2874,6 +2874,87 @@ def test_failed_request_skips_kv_postprocessing(
2874
assert invalid_blocks == {1, 2, 3}
2875
2876
2877
+@patch(
2878
+ "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
2879
+ FailingNixlWrapper,
2880
+)
2881
+def test_handles_failing_in_separate_polls_do_not_kill_the_engine(
2882
+ default_vllm_config, dist_init
2883
+):
2884
+ """A request whose handles fail in different polls must not crash the engine.
2885
+
2886
+ One transfer handle is created per remote rank, and when a peer goes away
2887
+ they do not all fail in the same poll: one errors while another is still
2888
+ PROC. The first failure reports the request via _failed_recv_reqs and
2889
+ get_finished() pops its metadata; when the remaining handles fail in a
2890
+ later poll, the request must be cleaned up without being reported again.
2891
+
2892
+ Reporting twice kills the EngineCore: the scheduler's assert in
2893
+ _update_from_kv_xfer_finished only expects a finished recv for a request
2894
+ still waiting for KVs, having moved this one out of
2895
+ WAITING_FOR_REMOTE_KVS to recompute locally after the first report.
2896
+ """
2897
+ vllm_config = create_vllm_config()
2898
+ connector = NixlConnector(
2899
+ vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16)
2900
+ )
2901
+ connector.connector_worker = FakeNixlConnectorWorker(
2902
+ vllm_config, connector.engine_id, hand_shake_latency=0
2903
+ )
2904
+ worker = connector.connector_worker
2905
+
2906
+ request_id = "test_two_handles_failing_in_separate_polls"
2907
+ metadata = NixlConnectorMetadata()
2908
+ metadata.add_new_req_to_recv(
2909
+ request_id=request_id,
2910
+ local_block_ids=([1, 2, 3],),
2911
+ kv_transfer_params={
2912
+ "remote_block_ids": ([4, 5, 6],),
2913
+ "remote_engine_id": FakeNixlConnectorWorker.REMOTE_ENGINE_ID,
2914
+ "remote_request_id": f"prefill-{request_id}",
2915
+ "remote_host": "localhost",
2916
+ "remote_port": 1234,
2917
+ "remote_tp_size": 1,
2918
+ },
2919
+ )
2920
+ connector.bind_connector_metadata(metadata)
2921
+ dummy_ctx = ForwardContext(no_compile_layers={}, attn_metadata={}, slot_mapping={})
2922
+ connector.start_load_kv(dummy_ctx)
2923
+ connector.bind_connector_metadata(NixlConnectorMetadata())
2924
+ time.sleep(0.1)
2925
+ connector.start_load_kv(dummy_ctx)
2926
+
2927
+ # Give the request a second handle, as a second remote rank would.
2928
+ handles = worker._recving_transfers[request_id]
2929
+ assert len(handles) == 1
2930
+ first, second = handles[0], object()
2931
+ worker._recving_transfers[request_id] = [first, second]
2932
+
2933
+ # Poll 1: the first handle has failed, the second is still running.
2934
+ def first_failed(handle):
2935
+ return "ERR" if handle is first else "PROC"
2936
+
2937
+ with patch.object(
2938
+ worker.nixl_wrapper, "check_xfer_state", side_effect=first_failed
2939
+ ):
2940
+ _, done_recving = connector.get_finished(finished_req_ids=set())
2941
+
2942
+ assert request_id in done_recving
2943
+ assert request_id not in worker._recving_metadata
2944
+ assert worker._recving_transfers[request_id] == [second]
2945
+
2946
+ # Poll 2: the second handle fails too. The request was already reported,
2947
+ # so it must not be reported a second time: the scheduler has since moved
2948
+ # it out of WAITING_FOR_REMOTE_KVS to recompute locally, and asserts on a
2949
+ # finished recv for a request in that state. Its remaining handles are
2950
+ # still released and the transfer entry removed.
2951
+ with patch.object(worker.nixl_wrapper, "check_xfer_state", return_value="ERR"):
2952
+ _, done_recving = connector.get_finished(finished_req_ids=set())
2953
+
2954
+ assert request_id not in done_recving
2955
+ assert request_id not in worker._recving_transfers
2956
+
2957
+
2958
def _set_test_speculative_config(
2959
vllm_config,
2960
*,
vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py
+13
-5
@@ -2521,7 +2521,11 @@ class NixlBaseConnectorWorker:
2521
2522
if not in_progress:
2523
# Only report request as completed when all transfers are done.
2524
- done_req_ids.add(req_id)
2524
+ # A request failed in an earlier poll was already reported via
2525
+ # _failed_recv_reqs and its metadata popped by get_finished();
2526
+ # don't report it again, just drop the remaining handles.
2527
+ if req_id in self._recving_metadata:
2528
+ done_req_ids.add(req_id)
2529
del transfers[req_id]
2530
else:
2531
transfers[req_id] = in_progress
@@ -2536,11 +2540,15 @@ class NixlBaseConnectorWorker:
2540
req_id: The request ID.
2541
handle: The transfer handle.
2542
"""
2539
- # Use .get() here as the metadata cleanup is handled by get_finished()
2543
+ # (multi-read) One handle is created per remote rank, and they do not
2544
+ # all fail in the same _pop_done_transfers poll. The request is
2545
+ # reported failed on the first one, which pops its metadata in
2546
+ # get_finished(); on the later failures only the handle cleanup is left.
2547
# TODO (NickLucche) handle failed transfer for HMA.
2541
- if (meta := self._recving_metadata.get(req_id)) and not self._is_hma_required:
2542
- self._invalid_block_ids.put(set(meta.local_block_ids[0]))
2543
- self._failed_recv_reqs.put(req_id)
2548
+ if (meta := self._recving_metadata.get(req_id)) is not None:
2549
+ if not self._is_hma_required:
2550
+ self._invalid_block_ids.put(set(meta.local_block_ids[0]))
2551
+ self._failed_recv_reqs.put(req_id)
2552
if handle is not None:
2553
self.nixl_wrapper.release_xfer_handle(handle)
2554
self.xfer_stats.record_failed_transfer()