diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py index 839d36a9a0..52bd4c65ea 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector.py +++ b/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, *, 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 04c50df0f5..4516c851cf 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 @@ -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()