@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 45779ef99f

[Bugfix][NIXL] Fix multiple handles xfer race (#56104)

Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai> Co-authored-by: Nick Hill <nickhill123@gmail.com>

Nicolò Lucchesi committed Sep 15, 2026 at 06:27 UTC 45779ef99f8c7f4bd59f86a48acccc860fc3c97c

5 files changed +233 -151

tests/v1/kv_connector/unit/test_nixl_connector.py

+77 -84

@@ -4,6 +4,7 @@

4 import contextlib

5 import inspect

6 import os

7 +import queue

8 import tempfile

9 import textwrap

10 import time

@@ -1952,17 +1953,43 @@ def test_mixed_memory_local_descriptors_split_by_memory_type():

1953 assert memory_types == ["DRAM", "VRAM"]

1954

1955

1955 -def test_mixed_memory_read_notifies_after_both_transfers_finish():

1956 +@pytest.fixture

1957 +def recv_worker():

1958 + """Receive lifecycle state without distributed or device initialization."""

1959 worker = object.__new__(NixlConnectorWorker)

1957 - worker._desc_is_dram_by_block_size = {16: np.array([True, True, False, False])}

1958 - worker._desc_pos_by_block_size = {16: np.array([0, 1, 0, 1])}

1959 - worker._dram_src_handles_by_block_size = {16: 10}

1960 - worker._recving_metadata = {"request": MagicMock()}

1960 + worker.transfer_topo = MagicMock()

1961 + worker.transfer_topo.block_size_ratio.return_value = 1

1962 + worker.transfer_topo.get_engine_info.return_value = SimpleNamespace(

1963 + remote_block_size=16, remote_physical_blocks_per_logical=1

1964 + )

1965 + worker._physical_blocks_per_logical_kv_block = 1

1966 + worker._recving_metadata = {"request": MagicMock(local_block_ids=([1, 2, 3],))}

1967 worker._recving_transfers = defaultdict(list)

1968 + worker._failed_recv_reqs = queue.Queue()

1969 + worker._recv_failures = set()

1970 + worker._replicated_pcp_done_sending = set()

1971 + worker._invalid_block_ids = queue.Queue()

1972 worker._pending_recv_notifs = {}

1963 - worker._failed_inflight_recvs = set()

1964 - worker.xfer_stats = MagicMock()

1973 + worker._reqs_to_send = {}

1974 + worker._replicated_pcp_done_sending = set()

1975 + worker._is_hma_required = False

1976 + worker._has_mamba = False

1977 + worker.use_host_buffer = False

1978 + worker.enable_permute_local_kv = False

1979 + worker.enable_heterogeneous_attn_post_process = False

1980 + worker.tp_rank = 0

1981 + worker._log_failure = MagicMock() # type: ignore[method-assign]

1982 + worker.xfer_stats = NixlKVConnectorStats()

1983 worker.nixl_wrapper = MagicMock()

1984 + worker.nixl_wrapper.get_xfer_telemetry.return_value = get_default_xfer_telemetry()

1985 + return worker

1986 +

1987 +

1988 +def test_mixed_memory_read_notifies_after_both_transfers_finish(recv_worker):

1989 + worker = recv_worker

1990 + worker._desc_is_dram_by_block_size = {16: np.array([True, True, False, False])}

1991 + worker._desc_pos_by_block_size = {16: np.array([0, 1, 0, 1])}

1992 + worker._dram_src_handles_by_block_size = {16: 10}

1993 worker.nixl_wrapper.make_prepped_xfer.side_effect = [101, 102]

1994 worker.nixl_wrapper.check_xfer_state.return_value = "DONE"

1995

@@ -1989,27 +2016,20 @@ def test_mixed_memory_read_notifies_after_both_transfers_finish():

2016 np.testing.assert_array_equal(device_read.args[4], [7])

2017 worker.nixl_wrapper.send_notif.assert_not_called()

2018

1992 - assert worker._pop_done_transfers(worker._recving_transfers) == {"request"}

2019 + assert worker.get_finished() == (set(), {"request"})

2020 worker.nixl_wrapper.send_notif.assert_called_once_with(

2021 "prefill", notif_msg=b"request:1"

2022 )

2023

2024

1998 -def test_mixed_memory_read_failure_does_not_notify_producer():

2025 +def test_mixed_memory_read_failure_does_not_notify_producer(recv_worker):

2026 """A failed half of a split READ must suppress its success notification."""

2000 - worker = object.__new__(NixlConnectorWorker)

2001 - worker._recving_metadata = {"request": MagicMock()}

2027 + worker = recv_worker

2028 worker._recving_transfers = {"request": [101, 102]}

2029 worker._pending_recv_notifs = {"request": [("prefill", b"request:1")]}

2004 - worker._failed_inflight_recvs = set()

2005 - worker._is_hma_required = True

2006 - worker._failed_recv_reqs = MagicMock()

2007 - worker._log_failure = MagicMock() # type: ignore[method-assign]

2008 - worker.xfer_stats = MagicMock()

2009 - worker.nixl_wrapper = MagicMock()

2030 worker.nixl_wrapper.check_xfer_state.side_effect = ["ERR", "DONE"]

2031

2012 - worker._pop_done_transfers(worker._recving_transfers)

2032 + assert worker.get_finished() == (set(), {"request"})

2033

2034 assert "request" not in worker._pending_recv_notifs

2035 worker.nixl_wrapper.send_notif.assert_not_called()

@@ -2969,13 +2989,10 @@ def test_handshake_failure_returns_finished(default_vllm_config, dist_init):

2989 # Wait for handshake to fail

2990 time.sleep(0.3)

2991

2972 - # Check that blocks were marked invalid

2973 - invalid_blocks = connector.get_block_ids_with_load_errors()

2974 - assert invalid_blocks == {1, 2, 3}

2975 -

2976 - # Check that request appears in get_finished

2992 + assert connector.get_block_ids_with_load_errors() == set()

2993 _, done_recving = connector.get_finished(finished_req_ids=set())

2994 assert request_id in done_recving

2995 + assert connector.get_block_ids_with_load_errors() == {1, 2, 3}

2996

2997

2998 @patch(

@@ -3026,6 +3043,7 @@ def test_transfer_setup_failure_returns_finished(

3043 time.sleep(0.1)

3044 connector.start_load_kv(dummy_ctx)

3045

3046 + assert connector.get_block_ids_with_load_errors() == set()

3047 results = connector.get_transfer_results(finished_req_ids=set())

3048 assert request_id in results.finished_recving

3049 assert results.failed_recving == {request_id}

@@ -3248,74 +3266,49 @@ def test_failed_request_skips_kv_postprocessing(

3266 assert invalid_blocks == {1, 2, 3}

3267

3268

3251 -@patch(

3252 - "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper",

3253 - FailingNixlWrapper,

3254 -)

3255 -def test_handles_failing_in_separate_polls_do_not_kill_the_engine(

3256 - default_vllm_config, dist_init

3257 -):

3258 - """Failures across polls report once, after the last writer is terminal."""

3259 - vllm_config = create_vllm_config()

3260 - connector = NixlConnector(

3261 - vllm_config, KVConnectorRole.WORKER, make_kv_cache_config(block_size=16)

3262 - )

3263 - connector.connector_worker = FakeNixlConnectorWorker(

3264 - vllm_config, connector.engine_id, hand_shake_latency=0

3265 - )

3266 - worker = connector.connector_worker

3269 +@pytest.mark.parametrize("sibling_state", ["ERR", "DONE"])

3270 +def test_recv_failure_waits_for_sibling_transfer(recv_worker, sibling_state):

3271 + """Failure must not release blocks while a sibling can still write to them."""

3272 + worker = recv_worker

3273 + worker._recving_transfers["request"] = [101, 102]

3274 + worker._pending_recv_notifs = {"request": [("prefill", b"request:1")]}

3275 + worker.nixl_wrapper.check_xfer_state.side_effect = ["ERR", "PROC", sibling_state]

3276

3268 - request_id = "test_two_handles_failing_in_separate_polls"

3269 - metadata = NixlConnectorMetadata()

3270 - metadata.add_new_req_to_recv(

3271 - request_id=request_id,

3272 - local_block_ids=([1, 2, 3],),

3273 - kv_transfer_params={

3274 - "remote_block_ids": ([4, 5, 6],),

3275 - "remote_engine_id": FakeNixlConnectorWorker.REMOTE_ENGINE_ID,

3276 - "remote_request_id": f"prefill-{request_id}",

3277 - "remote_host": "localhost",

3278 - "remote_port": 1234,

3279 - "remote_tp_size": 1,

3280 - },

3281 - )

3282 - connector.bind_connector_metadata(metadata)

3283 - dummy_ctx = ForwardContext(no_compile_layers={}, attn_metadata={}, slot_mapping={})

3284 - connector.start_load_kv(dummy_ctx)

3285 - connector.bind_connector_metadata(NixlConnectorMetadata())

3286 - time.sleep(0.1)

3287 - connector.start_load_kv(dummy_ctx)

3277 + results = worker.get_transfer_results()

3278 + assert results.finished_recving == results.failed_recving == set()

3279 + assert worker.get_block_ids_with_load_errors() == set()

3280 + assert "request" in worker._recving_metadata

3281 + worker.nixl_wrapper.release_xfer_handle.assert_called_once_with(101)

3282

3289 - # Give the request a second handle, as a second remote rank would.

3290 - handles = worker._recving_transfers[request_id]

3291 - assert len(handles) == 1

3292 - first, second = handles[0], object()

3293 - worker._recving_transfers[request_id] = [first, second]

3283 + results = worker.get_transfer_results()

3284 + assert results.finished_recving == results.failed_recving == {"request"}

3285 + assert worker.get_block_ids_with_load_errors() == {1, 2, 3}

3286 + worker.nixl_wrapper.send_notif.assert_not_called()

3287

3295 - # Poll 1: the first handle has failed, the second is still running.

3296 - def first_failed(handle):

3297 - return "ERR" if handle is first else "PROC"

3288 + results = worker.get_transfer_results()

3289 + assert results.finished_recving == results.failed_recving == set()

3290 + assert worker.get_block_ids_with_load_errors() == set()

3291

3299 - with patch.object(

3300 - worker.nixl_wrapper, "check_xfer_state", side_effect=first_failed

3301 - ):

3302 - _, done_recving = connector.get_finished(finished_req_ids=set())

3292

3304 - assert request_id not in done_recving

3305 - assert request_id in worker._recving_metadata

3306 - assert worker._recving_transfers[request_id] == [second]

3307 - assert connector.get_block_ids_with_load_errors() == set()

3293 +def test_recv_failure_waits_for_unpollable_handle_to_be_released(recv_worker):

3294 + """A polling exception and failed release do not prove that DMA has stopped."""

3295 + worker = recv_worker

3296 + worker._recving_transfers["request"] = [101]

3297 + worker.nixl_wrapper.check_xfer_state.side_effect = [RuntimeError("poll"), "DONE"]

3298 + worker.nixl_wrapper.release_xfer_handle.side_effect = [

3299 + RuntimeError("release"),

3300 + None,

3301 + ]

3302

3309 - # Only the last terminal handle permits the scheduler to reuse the blocks.

3310 - with patch.object(worker.nixl_wrapper, "check_xfer_state", return_value="ERR"):

3311 - _, done_recving = connector.get_finished(finished_req_ids=set())

3303 + results = worker.get_transfer_results()

3304 + assert results.finished_recving == results.failed_recving == set()

3305 + assert worker.get_block_ids_with_load_errors() == set()

3306 + assert "request" in worker._recving_metadata

3307

3313 - assert done_recving == {request_id}

3314 - assert request_id not in worker._recving_transfers

3315 - assert request_id not in worker._recving_metadata

3316 - assert connector.get_block_ids_with_load_errors() == {1, 2, 3}

3317 - _, done_recving = connector.get_finished(finished_req_ids=set())

3318 - assert done_recving == set()

3308 + results = worker.get_transfer_results()

3309 + assert results.finished_recving == results.failed_recving == {"request"}

3310 + assert worker.get_block_ids_with_load_errors() == {1, 2, 3}

3311 + assert worker.nixl_wrapper.release_xfer_handle.call_count == 2

3312

3313

3314 def _set_test_speculative_config(

tests/v1/kv_connector/unit/test_nixl_push_connector.py

+68 -5

@@ -351,7 +351,7 @@ class _StubWriterWorker(NixlPushConnectorWorker):

351 # Base worker fields touched by start_load_kv / _get_new_notifs.

352 w._recving_metadata = {}

353 w._pending_recv_notifs = {}

354 - w._failed_inflight_recvs = set()

354 + w._recv_failures = set()

355 w._recving_transfers = defaultdict(list)

356 w._reqs_to_process = set()

357 w._reqs_to_send = {}

@@ -877,6 +877,68 @@ class TestPushWriterNotifs:

877 assert evicted == ["req-done"]

878 assert w._push_writer_wake.is_set()

879

880 + @staticmethod

881 + def _pollable_worker() -> _StubWriterWorker:

882 + """Stub worker with the extra base-worker state the real

883 + ``get_transfer_results`` needs to poll ``_sending_transfers``."""

884 + w = _StubWriterWorker.fresh()

885 + w.transfer_topo = MagicMock()

886 + w.nixl_wrapper = MagicMock()

887 + w.xfer_stats = MagicMock()

888 + w._log_failure = MagicMock() # type: ignore[method-assign]

889 + w._failed_recv_reqs = queue.Queue()

890 + w._invalid_block_ids = queue.Queue()

891 + w._pending_recv_notifs = {}

892 + w._replicated_pcp_done_sending = set()

893 + w.use_host_buffer = False

894 + return w

895 +

896 + def _make_sending_req(self, w: _StubWriterWorker) -> str:

897 + request_id = "req-send-1"

898 + w._sending_transfers[request_id] = [101, 102]

899 + w._reqs_to_send[request_id] = time.perf_counter() + 60

900 + w._reqs_to_process.add(request_id)

901 + return request_id

902 +

903 + def test_failed_push_send_is_not_reported_done_sending(self):

904 + """A failed P-side WRITE must not surface as finished sending.

905 +

906 + Reporting it would free P-side blocks as if the KV had been

907 + delivered while D still waits for the data; the request must be

908 + left to the lease / watchdog instead.

909 + """

910 + w = self._pollable_worker()

911 + request_id = self._make_sending_req(w)

912 + w.nixl_wrapper.check_xfer_state.side_effect = ["ERR", "DONE"]

913 +

914 + results = w.get_transfer_results()

915 +

916 + assert request_id not in results.finished_sending

917 + # Lease tracking is preserved so the watchdog can reschedule.

918 + assert request_id in w._reqs_to_send

919 + assert request_id in w._reqs_to_process

920 + # Both handles were still cleaned up.

921 + assert request_id not in w._sending_transfers

922 +

923 + def test_completed_push_send_waits_for_consumer_notif(self):

924 + """P-side sends are completed by consumer notifs / lease expiry,

925 + not by local WRITE-handle completion."""

926 + w = self._pollable_worker()

927 + request_id = self._make_sending_req(w)

928 + w.nixl_wrapper.check_xfer_state.return_value = "DONE"

929 +

930 + results = w.get_transfer_results()

931 +

932 + assert request_id not in results.finished_sending

933 + assert request_id in w._reqs_to_send

934 + assert request_id not in w._sending_transfers

935 +

936 + # The consumer notif is what completes the send.

937 + w._pending_completion_notifs.put(f"{request_id}:1".encode())

938 + results = w.get_transfer_results()

939 + assert request_id in results.finished_sending

940 + assert request_id not in w._reqs_to_send

941 +

942

943 # ----------------------------------------------------------------- #

944 # Negative / error-path tests #

@@ -1079,7 +1141,7 @@ class TestPushWriterNegative:

1141 # Wake set so the writer drains NIXL notifs even when idle.

1142 assert w._push_writer_wake.is_set()

1143

1082 - def test_failed_push_transfer_finishes_sending(self):

1144 + def test_failed_push_transfer_does_not_evict_writer_state(self):

1145 w = _StubWriterWorker.fresh()

1146 w.nixl_wrapper = MagicMock()

1147 w.nixl_wrapper.check_xfer_state.return_value = "ERR"

@@ -1096,11 +1158,12 @@ class TestPushWriterNegative:

1158 ):

1159 results = w.get_transfer_results()

1160

1099 - assert results.finished_sending == {"req-failed"}

1161 + assert results.finished_sending == set()

1162 assert results.finished_recving == set()

1163 assert "req-failed" not in w._sending_transfers

1102 - assert "req-failed" not in w._reqs_to_send

1103 - assert "req-failed" not in w._reqs_to_process

1164 + assert "req-failed" in w._reqs_to_send

1165 + assert "req-failed" in w._reqs_to_process

1166 + assert w._evict_finished_inbox.empty()

1167 w.nixl_wrapper.release_xfer_handle.assert_called_once_with(17)

1168

1169 def test_get_new_notifs_unknown_request_is_logged_and_skipped(self, caplog):

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

+63 -50

@@ -692,7 +692,7 @@ class NixlBaseConnectorWorker:

692 # Uses Queue for thread-safe cross-thread coordination with the

693 # background handshake thread, matching the _ready_requests pattern.

694 self._failed_recv_reqs: queue.Queue[ReqId] = queue.Queue()

695 - self._failed_inflight_recvs: set[ReqId] = set()

695 + self._recv_failures: set[ReqId] = set()

696 self._pending_recv_notifs: dict[ReqId, list[tuple[str, bytes]]] = {}

697

698 # Handshake metadata of this worker for NIXL transfers.

@@ -1212,7 +1212,7 @@ class NixlBaseConnectorWorker:

1212 error=e,

1213 meta=meta,

1214 )

1215 - self._handle_failed_transfer(req_id, None)

1215 + self._failed_recv_reqs.put(req_id)

1216

1217 fut.add_done_callback(request_ready)

1218

@@ -2618,22 +2618,32 @@ class NixlBaseConnectorWorker:

2618 """

2619 assert self.transfer_topo is not None

2620 done_sending = self._get_new_notifs()

2621 - done_recving = self._pop_done_transfers(self._recving_transfers)

2621 + done_recving, newly_failed = self._pop_done_transfers(self._recving_transfers)

2622 + if newly_failed:

2623 + self._recv_failures.update(newly_failed)

2624

2625 done_sending.update(self._replicated_pcp_done_sending)

2626 self._replicated_pcp_done_sending.clear()

2627

2626 - # Drain queue of requests where handshake or transfer setup failed.

2627 - failed_recv_reqs = set[ReqId]()

2628 + # Process receive failures reported by background threads.

2629 while not self._failed_recv_reqs.empty():

2630 try:

2630 - failed_recv_reqs.add(self._failed_recv_reqs.get_nowait())

2631 + req_id = self._failed_recv_reqs.get_nowait()

2632 except queue.Empty:

2633 break

2634 + self._handle_failed_transfer(req_id, None, self._recv_failures)

2635 +

2636 + failed_recv_reqs: set[ReqId] = set()

2637 + if self._recv_failures:

2638 + # Keep failure bookkeeping independent of healthy request count.

2639 + for req_id in tuple(self._recv_failures):

2640 + if req_id not in self._recving_metadata:

2641 + self._recv_failures.remove(req_id)

2642 + elif req_id not in self._recving_transfers:

2643 + self._recv_failures.remove(req_id)

2644 + failed_recv_reqs.add(req_id)

2645 + done_recving.update(failed_recv_reqs)

2646

2634 - # Add failed requests to done_recving for scheduler tracking

2635 - # (blocks are already marked invalid, scheduler will handle recompute)

2636 - done_recving.update(failed_recv_reqs)

2647 if len(done_sending) > 0 or len(done_recving) > 0:

2648 logger.debug(

2649 "Rank %s, get_finished: %s requests done sending "

@@ -2654,12 +2664,17 @@ class NixlBaseConnectorWorker:

2664

2665 # Skip KV sync and post-processing for failed requests

2666 if req_id in failed_recv_reqs:

2667 + self._pending_recv_notifs.pop(req_id, None)

2668 + # TODO (NickLucche) handle failed transfer for HMA.

2669 + if not self._is_hma_required:

2670 + self._invalid_block_ids.put(set(meta.local_block_ids[0]))

2671 logger.warning(

2672 "Skipping KV post-processing for failed request %s",

2673 req_id,

2674 )

2675 continue

2676

2677 + self._send_pending_recv_notifs(req_id)

2678 assert meta.remote is not None

2679 if self.use_host_buffer:

2680 self.sync_recved_kv_to_device(req_id, meta)

@@ -2783,22 +2798,26 @@ class NixlBaseConnectorWorker:

2798 new_expiry,

2799 )

2800

2786 - def _pop_done_transfers(self, transfers: dict[str, list[int]]) -> set[str]:

2787 - """

2788 - Pop completed xfers by checking for DONE state.

2801 + def _pop_done_transfers(

2802 + self, transfers: dict[str, list[int]]

2803 + ) -> tuple[set[str], set[str]]:

2804 + """Poll transfers, retaining handles until they can be released.

2805 +

2806 Args:

2790 - transfers: dict of req_id -> list[running_xfer]

2807 + transfers: Outstanding handles by request, updated in place.

2808 +

2809 Returns:

2792 - set of req_ids that have all done xfers

2810 + Requests with no outstanding handles and requests with observed

2811 + failures. Failed requests may still have outstanding handles.

2812 """

2813 done_req_ids: set[str] = set()

2814 + failed_req_ids: set[str] = set()

2815 for req_id, handles in list(transfers.items()):

2816 in_progress = []

2817 for handle in handles:

2818 try:

2819 xfer_state = self.nixl_wrapper.check_xfer_state(handle)

2820 if xfer_state == "DONE":

2801 - # Get telemetry from NIXL

2821 res = self.nixl_wrapper.get_xfer_telemetry(handle)

2822 self.xfer_stats.record_transfer(res)

2823 self.nixl_wrapper.release_xfer_handle(handle)

@@ -2807,58 +2826,52 @@ class NixlBaseConnectorWorker:

2826 else:

2827 self._log_failure(

2828 failure_type="transfer_failed",

2810 - msg="Deferring request completion until its last "

2811 - "xfer is terminal",

2829 req_id=req_id,

2830 xfer_state=xfer_state,

2831 )

2815 - self._handle_failed_transfer(req_id, handle)

2832 + if not self._handle_failed_transfer(

2833 + req_id, handle, failed_req_ids

2834 + ):

2835 + in_progress.append(handle)

2836 except Exception as e:

2837 self._log_failure(

2838 failure_type="transfer_exception",

2819 - msg="Handle is unpollable; treating it as terminal",

2839 req_id=req_id,

2840 error=e,

2841 )

2823 - self._handle_failed_transfer(req_id, handle)

2842 + if not self._handle_failed_transfer(req_id, handle, failed_req_ids):

2843 + in_progress.append(handle)

2844

2845 if not in_progress:

2826 - del transfers[req_id]

2846 done_req_ids.add(req_id)

2828 - if req_id in self._failed_inflight_recvs:

2829 - self._failed_inflight_recvs.remove(req_id)

2830 - self._report_failed_recv(req_id)

2831 - else:

2832 - self._send_pending_recv_notifs(req_id)

2847 + del transfers[req_id]

2848 else:

2849 transfers[req_id] = in_progress

2835 - return done_req_ids

2850 + return done_req_ids, failed_req_ids

2851

2837 - def _handle_failed_transfer(self, req_id: str, handle: int | None):

2838 - """

2839 - Handle a failed transfer by marking all (logical) blocks as invalid and

2840 - recording the failure.

2841 -

2842 - Args:

2843 - req_id: The request ID.

2844 - handle: The transfer handle.

2845 - """

2846 - # A sibling READ may still be writing these blocks. Retain metadata

2847 - # and defer invalidation until every handle is terminal.

2848 - self._pending_recv_notifs.pop(req_id, None)

2849 - if req_id in self._recving_transfers:

2850 - self._failed_inflight_recvs.add(req_id)

2851 - else:

2852 - self._report_failed_recv(req_id)

2853 - if handle is not None:

2852 + def _try_release_xfer_handle(self, req_id: str, handle: int) -> bool:

2853 + """Release a handle, returning False if the caller must retain it."""

2854 + try:

2855 self.nixl_wrapper.release_xfer_handle(handle)

2856 + except Exception as e:

2857 + # A status error does not guarantee that the backend stopped DMA.

2858 + self._log_failure(

2859 + failure_type="transfer_release_failed",

2860 + msg="Retaining handle and blocks until release succeeds",

2861 + req_id=req_id,

2862 + error=e,

2863 + )

2864 + return False

2865 + return True

2866 +

2867 + def _handle_failed_transfer(

2868 + self, req_id: str, handle: int | None, failed_req_ids: set[str] | None = None

2869 + ) -> bool:

2870 + """Record a failure and release its handle, returning False to retain it."""

2871 self.xfer_stats.record_failed_transfer()

2856 -

2857 - def _report_failed_recv(self, req_id: str) -> None:

2858 - if (meta := self._recving_metadata.get(req_id)) is not None:

2859 - if not self._is_hma_required:

2860 - self._invalid_block_ids.put(set(meta.local_block_ids[0]))

2861 - self._failed_recv_reqs.put(req_id)

2872 + if failed_req_ids is not None:

2873 + failed_req_ids.add(req_id)

2874 + return handle is None or self._try_release_xfer_handle(req_id, handle)

2875

2876 def _send_pending_recv_notifs(self, req_id: str) -> None:

2877 """Send notifications deferred by split DRAM/VRAM reads."""

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

+11 -6

@@ -149,7 +149,7 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):

149 engine_id,

150 )

151 self.xfer_stats.record_kv_expired_req()

152 - self._handle_failed_transfer(req_id, None)

152 + self._handle_failed_transfer(req_id, None, self._recv_failures)

153 return

154

155 if any(len(group) > 0 for group in meta.local_block_ids):

@@ -493,16 +493,19 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):

493 self._recving_transfers[request_id].append(handle)

494 return True

495 except Exception as e:

496 - # mark all (logical) blocks for this request as invalid

496 self._log_failure(

497 failure_type="transfer_setup_failed",

498 req_id=request_id,

500 - msg="Marking blocks as invalid",

499 + msg="Deferring failure reporting until outstanding transfers finish",

500 error=e,

501 dst_engine_id=dst_engine_id,

502 remote_rank=remote_rank,

503 )

505 - self._handle_failed_transfer(request_id, handle)

504 + if not self._handle_failed_transfer(

505 + request_id, handle, self._recv_failures

506 + ):

507 + assert handle is not None

508 + self._recving_transfers[request_id].append(handle)

509 return False

510

511 def _read_blocks_mixed(

@@ -544,7 +547,8 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):

547 )

548 except Exception:

549 for handle in handles:

547 - self.nixl_wrapper.release_xfer_handle(handle)

550 + if not self._try_release_xfer_handle(request_id, handle):

551 + self._recving_transfers[request_id].append(handle)

552 raise

553

554 self._pending_recv_notifs.setdefault(request_id, []).append(

@@ -555,7 +559,8 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):

559 self.nixl_wrapper.transfer(handle)

560 except Exception:

561 for unstarted in handles[i:]:

558 - self.nixl_wrapper.release_xfer_handle(unstarted)

562 + if not self._try_release_xfer_handle(request_id, unstarted):

563 + self._recving_transfers[request_id].append(unstarted)

564 raise

565 self._recving_transfers[request_id].append(handle)

566

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

+14 -6

@@ -341,7 +341,7 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker):

341 self._log_failure(

342 failure_type="push_reg_handshake_failed", req_id=rid, error=e

343 )

344 - self._handle_failed_transfer(rid, None)

344 + self._failed_recv_reqs.put(rid)

345 return

346 # Re-queue for the writer to send now that the handshake is done.

347 self._reg_send_inbox.put((rid, rd))

@@ -363,7 +363,7 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker):

363 engine_id,

364 req_id,

365 )

366 - self._handle_failed_transfer(req_id, None)

366 + self._failed_recv_reqs.put(req_id)

367 return

368 for rank, agent_name in agents.items():

369 try:

@@ -722,9 +722,8 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker):

722 # don't have a ``_recving_metadata`` entry to invalidate, so

723 # we just release the handle and let the engine reschedule

724 # via the lease / watchdog.

725 - if handle is not None:

726 - self.nixl_wrapper.release_xfer_handle(handle)

727 - self.xfer_stats.record_failed_transfer()

725 + if not self._handle_failed_transfer(request_id, handle):

726 + return handle

727 return None

728

729 # --- Notification handling on engine main thread ------------------ #

@@ -803,7 +802,16 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker):

802 # ``_pop_done_transfers`` mutates ``_sending_transfers``; the

803 # writer thread also appends to it, so guard the pop.

804 with self._sending_transfers_lock:

806 - done_pushing = self._pop_done_transfers(self._sending_transfers)

805 + done_pushing, failed_pushing = self._pop_done_transfers(

806 + self._sending_transfers

807 + )

808 + # A failed send must never be reported as done: its blocks

809 + # are freed via the lease / watchdog instead.

810 + done_pushing = {

811 + req_id

812 + for req_id in done_pushing - failed_pushing

813 + if req_id in self._recving_metadata

814 + }

815 for req_id in done_pushing:

816 self._reqs_to_send.pop(req_id, None)

817 self._reqs_to_process.discard(req_id)