@tduncan / Nvidia-Optimization-Stack-VLLM / commits / c4f6ce4874

[Bugfix][NIXL] Restore successful push completion reporting (#58188)

Signed-off-by: Dao Le <Dao007forever@gmail.com> Co-authored-by: Codex <noreply@openai.com>

Dao007forever committed Sep 23, 2026 at 09:39 UTC c4f6ce48749a092b02525af16947350305650573

2 files changed +111 -20

tests/v1/kv_connector/unit/test_nixl_push_connector.py

+97 -13

@@ -50,7 +50,10 @@ from vllm.distributed.kv_transfer.kv_connector.v1.nixl.push_worker import (

50 from vllm.distributed.kv_transfer.kv_connector.v1.nixl.stats import (

51 NixlKVConnectorStats,

52 )

53 -from vllm.distributed.kv_transfer.kv_connector.v1.nixl.tp_mapping import TPMapping

53 +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.tp_mapping import (

54 + ReadSpec,

55 + TPMapping,

56 +)

57 from vllm.distributed.kv_transfer.kv_connector.v1.nixl.utils import (

58 get_base_request_id,

59 )

@@ -344,6 +347,7 @@ class _StubWriterWorker(NixlPushConnectorWorker):

347 )

348

349 w._sending_transfers = defaultdict[ReqId, list[TransferHandle]](list)

350 + w._send_failures = set()

351 w._sending_transfers_lock = threading.Lock()

352 w._push_finished_blocks = {}

353 w._pending_d_registrations = {}

@@ -941,30 +945,25 @@ class TestPushWriterNotifs:

945 assert request_id not in w._sending_transfers

946

947 @pytest.mark.parametrize(("pp_size", "is_hma"), [(1, False), (2, True)])

944 - def test_completed_push_send_waits_for_consumer_notif(self, pp_size, is_hma):

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

946 - not by local WRITE-handle completion."""

948 + def test_completed_push_send_releases_blocks_without_consumer_notif(

949 + self, pp_size, is_hma

950 + ):

951 + """All WRITEs completing lets P reclaim its blocks without a D ACK."""

952 w = self._pollable_worker()

953 w.pp_size = pp_size

954 w._is_hma_required = is_hma

955 request_id = self._make_sending_req(w)

956 w.nixl_wrapper.check_xfer_state.side_effect = ["DONE", "PROC", "DONE"]

957

953 - results = w.get_transfer_results()

954 - assert request_id not in results.finished_sending

955 - assert w._sending_transfers[request_id] == [102]

956 - assert request_id in w._reqs_to_send

957 -

958 results = w.get_transfer_results()

959

960 assert request_id not in results.finished_sending

961 assert request_id in w._reqs_to_send

962 - assert request_id not in w._sending_transfers

962 + assert w._sending_transfers[request_id] == [102]

963 + assert w._evict_finished_inbox.empty()

964

964 - # The consumer notif is what completes the send.

965 - w._pending_completion_notifs.put(f"{request_id}:1".encode())

965 results = w.get_transfer_results()

967 - assert request_id in results.finished_sending

966 + assert results.finished_sending == {request_id}

967 assert request_id not in w._reqs_to_send

968 assert request_id not in w._reqs_to_process

969 assert w._evict_finished_inbox.get_nowait() == request_id

@@ -973,6 +972,91 @@ class TestPushWriterNotifs:

972 assert w.nixl_wrapper.release_xfer_handle.call_count == 2

973 w.xfer_stats.record_kv_expired_req.assert_not_called()

974

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

976 + def test_push_failure_survives_until_lease_expiry(self, sibling_state):

977 + """A later sibling completion must not turn a failed send into success."""

978 + w = self._pollable_worker()

979 + request_id = self._make_sending_req(w)

980 + w.nixl_wrapper.check_xfer_state.side_effect = ["ERR", "PROC", sibling_state]

981 +

982 + assert w.get_transfer_results().finished_sending == set()

983 + assert w._sending_transfers[request_id] == [102]

984 + assert w.get_transfer_results().finished_sending == set()

985 + assert request_id not in w._sending_transfers

986 + assert request_id in w._reqs_to_send

987 + assert w._evict_finished_inbox.empty()

988 +

989 + w.expected_consumer_notifications_by_req = {}

990 + w._reqs_to_send[request_id] = time.perf_counter() - 1

991 + assert w.get_transfer_results().finished_sending == {request_id}

992 + assert request_id not in w._send_failures

993 + assert w._evict_finished_inbox.get_nowait() == request_id

994 +

995 + @pytest.mark.parametrize("release_fails", [False, True])

996 + def test_push_submission_failure_is_remembered_after_sibling_success(

997 + self, release_fails

998 + ):

999 + """A submission error remains a failure even if later polls succeed."""

1000 + w = self._pollable_worker()

1001 + request_id = self._make_sending_req(w)

1002 + w.transfer_topo.block_size_ratio.return_value = 1

1003 + w._apply_prefix_caching = MagicMock(return_value=([[1]], [[1]]))

1004 + w._compute_desc_ids = MagicMock(return_value=[0])

1005 + w.dst_num_blocks = {w.engine_id: 8}

1006 + w.dst_region_num_blocks[w.engine_id] = [8]

1007 + w.dst_region_group_ids[w.engine_id] = [0]

1008 + w.dst_uses_region_group_mapping[w.engine_id] = False

1009 + w.nixl_wrapper.make_prepped_xfer.return_value = 101

1010 + w.nixl_wrapper.transfer.side_effect = RuntimeError("submit")

1011 + if release_fails:

1012 + w.nixl_wrapper.release_xfer_handle.side_effect = [

1013 + RuntimeError("release"),

1014 + None,

1015 + None,

1016 + ]

1017 +

1018 + handle = w._xfer_blocks(

1019 + read_spec=ReadSpec(0, [[1]], [[1]]),

1020 + dst_engine_id=w.engine_id,

1021 + request_id=request_id,

1022 + remote_request_id="decode-request",

1023 + local_xfer_side_handle=1,

1024 + remote_xfer_side_handle=2,

1025 + )

1026 + assert handle == (101 if release_fails else None)

1027 + w._sending_transfers[request_id] = [102]

1028 + if handle is not None:

1029 + w._sending_transfers[request_id].append(handle)

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

1031 +

1032 + assert w.get_transfer_results().finished_sending == set()

1033 + assert request_id in w._reqs_to_send

1034 + assert request_id not in w._sending_transfers

1035 + assert w._evict_finished_inbox.empty()

1036 +

1037 + @pytest.mark.parametrize("first_state", ["ERR", "DONE"])

1038 + def test_push_completion_after_expiry_is_not_reported_twice(self, first_state):

1039 + """A WRITE finishing after lease expiry must not complete the request again."""

1040 + w = self._pollable_worker()

1041 + request_id = self._make_sending_req(w)

1042 + w.nixl_wrapper.check_xfer_state.side_effect = [

1043 + first_state,

1044 + "PROC",

1045 + "PROC",

1046 + "DONE",

1047 + ]

1048 + assert w.get_transfer_results().finished_sending == set()

1049 +

1050 + w.expected_consumer_notifications_by_req = {}

1051 + w._reqs_to_send[request_id] = time.perf_counter() - 1

1052 + assert w.get_transfer_results().finished_sending == {request_id}

1053 + assert w._evict_finished_inbox.get_nowait() == request_id

1054 +

1055 + assert w.get_transfer_results().finished_sending == set()

1056 + assert request_id not in w._sending_transfers

1057 + assert request_id not in w._send_failures

1058 + assert w._evict_finished_inbox.empty()

1059 +

1060

1061 # ----------------------------------------------------------------- #

1062 # Negative / error-path tests #

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

+14 -7

@@ -103,6 +103,7 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker):

103 # (``_pop_done_transfers``); guarded by

104 # ``_sending_transfers_lock``.

105 self._sending_transfers = defaultdict[ReqId, list[TransferHandle]](list)

106 + self._send_failures: set[ReqId] = set()

107 self._sending_transfers_lock = threading.Lock()

108

109 # Writer-thread owned matching state.

@@ -162,6 +163,7 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker):

163 for handle in handles:

164 self.nixl_wrapper.release_xfer_handle(handle)

165 self._sending_transfers.clear()

166 + self._send_failures.clear()

167 super().shutdown()

168

169 # --- Engine-main-thread entry point -------------------------------- #

@@ -728,6 +730,8 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker):

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

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

732 # via the lease / watchdog.

733 + with self._sending_transfers_lock:

734 + self._send_failures.add(request_id)

735 if not self._handle_failed_transfer(request_id, handle):

736 return handle

737 return None

@@ -811,13 +815,16 @@ class NixlPushConnectorWorker(NixlBaseConnectorWorker):

815 done_pushing, failed_pushing = self._pop_done_transfers(

816 self._sending_transfers

817 )

814 - # A failed send must never be reported as done: its blocks

815 - # are freed via the lease / watchdog instead.

816 - done_pushing = {

817 - req_id

818 - for req_id in done_pushing - failed_pushing

819 - if req_id in self._recving_metadata

820 - }

818 + # Remember failures until the final sibling WRITE completes.

819 + self._send_failures.update(failed_pushing)

820 + successful = {

821 + req_id

822 + for req_id in done_pushing - self._send_failures

823 + if req_id in self._reqs_to_send or req_id in self._reqs_to_process

824 + }

825 + self._send_failures.difference_update(done_pushing | done_sending)

826 + # Expired requests were already reported, even if their WRITEs finish later.

827 + done_pushing = successful

828 for req_id in done_pushing:

829 self._reqs_to_send.pop(req_id, None)

830 self._reqs_to_process.discard(req_id)