[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>

tests/v1/kv_connector/unit/test_nixl_connector.py

@@ -2874,6 +2874,87 @@ def test_failed_request_skips_kv_postprocessing(
     assert invalid_blocks == {1, 2, 3}
 
 
+@patch(
+    "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",
+    FailingNixlWrapper,
+)
+def test_handles_failing_in_separate_polls_do_not_kill_the_engine(
+    default_vllm_config, dist_init
+):
+    """A request whose handles fail in different polls must not crash the engine.
+
+    One transfer handle is created per remote rank, and when a peer goes away
+    they do not all fail in the same poll: one errors while another is still
+    PROC. The first failure reports the request via _failed_recv_reqs and
+    get_finished() pops its metadata; when the remaining handles fail in a
+    later poll, the request must be cleaned up without being reported again.
+
+    Reporting twice kills the EngineCore: the scheduler's assert in
+    _update_from_kv_xfer_finished only expects a finished recv for a request
+    still waiting for KVs, having moved this one out of
+    WAITING_FOR_REMOTE_KVS to recompute locally after the first report.
+    """
+    vllm_config = create_vllm_config()
+    connector = NixlConnector(
+        vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16)
+    )
+    connector.connector_worker = FakeNixlConnectorWorker(
+        vllm_config, connector.engine_id, hand_shake_latency=0
+    )
+    worker = connector.connector_worker
+
+    request_id = "test_two_handles_failing_in_separate_polls"
+    metadata = NixlConnectorMetadata()
+    metadata.add_new_req_to_recv(
+        request_id=request_id,
+        local_block_ids=([1, 2, 3],),
+        kv_transfer_params={
+            "remote_block_ids": ([4, 5, 6],),
+            "remote_engine_id": FakeNixlConnectorWorker.REMOTE_ENGINE_ID,
+            "remote_request_id": f"prefill-{request_id}",
+            "remote_host": "localhost",
+            "remote_port": 1234,
+            "remote_tp_size": 1,
+        },
+    )
+    connector.bind_connector_metadata(metadata)
+    dummy_ctx = ForwardContext(no_compile_layers={}, attn_metadata={}, slot_mapping={})
+    connector.start_load_kv(dummy_ctx)
+    connector.bind_connector_metadata(NixlConnectorMetadata())
+    time.sleep(0.1)
+    connector.start_load_kv(dummy_ctx)
+
+    # Give the request a second handle, as a second remote rank would.
+    handles = worker._recving_transfers[request_id]
+    assert len(handles) == 1
+    first, second = handles[0], object()
+    worker._recving_transfers[request_id] = [first, second]
+
+    # Poll 1: the first handle has failed, the second is still running.
+    def first_failed(handle):
+        return "ERR" if handle is first else "PROC"
+
+    with patch.object(
+        worker.nixl_wrapper, "check_xfer_state", side_effect=first_failed
+    ):
+        _, done_recving = connector.get_finished(finished_req_ids=set())
+
+    assert request_id in done_recving
+    assert request_id not in worker._recving_metadata
+    assert worker._recving_transfers[request_id] == [second]
+
+    # Poll 2: the second handle fails too. The request was already reported,
+    # so it must not be reported a second time: the scheduler has since moved
+    # it out of WAITING_FOR_REMOTE_KVS to recompute locally, and asserts on a
+    # finished recv for a request in that state. Its remaining handles are
+    # still released and the transfer entry removed.
+    with patch.object(worker.nixl_wrapper, "check_xfer_state", return_value="ERR"):
+        _, done_recving = connector.get_finished(finished_req_ids=set())
+
+    assert request_id not in done_recving
+    assert request_id not in worker._recving_transfers
+
+
 def _set_test_speculative_config(
     vllm_config,
     *,

vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py

@@ -2521,7 +2521,11 @@ class NixlBaseConnectorWorker:
 
             if not in_progress:
                 # Only report request as completed when all transfers are done.
-                done_req_ids.add(req_id)
+                # A request failed in an earlier poll was already reported via
+                # _failed_recv_reqs and its metadata popped by get_finished();
+                # don't report it again, just drop the remaining handles.
+                if req_id in self._recving_metadata:
+                    done_req_ids.add(req_id)
                 del transfers[req_id]
             else:
                 transfers[req_id] = in_progress
@@ -2536,11 +2540,15 @@ class NixlBaseConnectorWorker:
             req_id: The request ID.
             handle: The transfer handle.
         """
-        # Use .get() here as the metadata cleanup is handled by get_finished()
+        # (multi-read) One handle is created per remote rank, and they do not
+        # all fail in the same _pop_done_transfers poll. The request is
+        # reported failed on the first one, which pops its metadata in
+        # get_finished(); on the later failures only the handle cleanup is left.
         # TODO (NickLucche) handle failed transfer for HMA.
-        if (meta := self._recving_metadata.get(req_id)) and not self._is_hma_required:
-            self._invalid_block_ids.put(set(meta.local_block_ids[0]))
-        self._failed_recv_reqs.put(req_id)
+        if (meta := self._recving_metadata.get(req_id)) is not None:
+            if not self._is_hma_required:
+                self._invalid_block_ids.put(set(meta.local_block_ids[0]))
+            self._failed_recv_reqs.put(req_id)
         if handle is not None:
             self.nixl_wrapper.release_xfer_handle(handle)
         self.xfer_stats.record_failed_transfer()