From c34a267661edea6f9bb713b433331b0bd5c371f1 Mon Sep 17 00:00:00 2001 From: Itay Etelis <92247226+Etelis@users.noreply.github.com> Date: Sun, 6 Sep 2026 10:29:31 +0300 Subject: [PATCH] Handle P2P NIXL submission failures Signed-off-by: Itay Etelis <92247226+Etelis@users.noreply.github.com> --- .../tiering/p2p/test_data_transport.py | 31 +++++++++++++ .../kv_offload/tiering/p2p/test_sessions.py | 43 +++++++++++-------- vllm/v1/kv_offload/tiering/p2p/data/nixl.py | 32 ++++++++++---- .../kv_offload/tiering/p2p/session/server.py | 1 + 4 files changed, 79 insertions(+), 28 deletions(-) diff --git a/tests/v1/kv_offload/tiering/p2p/test_data_transport.py b/tests/v1/kv_offload/tiering/p2p/test_data_transport.py index d3bacc8f326f..70959313bbb9 100644 --- a/tests/v1/kv_offload/tiering/p2p/test_data_transport.py +++ b/tests/v1/kv_offload/tiering/p2p/test_data_transport.py @@ -8,6 +8,7 @@ from unittest.mock import MagicMock, patch import numpy as np +import pytest from vllm.v1.kv_offload.tiering.p2p.data.base import PollResult from vllm.v1.kv_offload.tiering.p2p.data.nixl import NixlTransport @@ -122,6 +123,36 @@ def test_write_blocks_returns_transfer_id(self): assert tid is not None assert isinstance(tid, int) + @pytest.mark.parametrize( + "failure_method,release_fails", + [ + ("make_prepped_xfer", False), + ("transfer", False), + ("transfer", True), + ], + ) + def test_write_blocks_submission_exception_returns_failure( + self, failure_method, release_fails + ): + """Setup failures surface even if releasing a created handle also fails.""" + transport = self._make_transport() + transport.add_remote_peer("peer:1", b"meta", 0x1000, 8, 1024) + agent = transport._agent + getattr(agent, failure_method).side_effect = RuntimeError("submission failed") + if release_fails: + agent.release_xfer_handle.side_effect = RuntimeError("release failed") + + assert transport.write_blocks("peer:1", [0], [1]) is None + assert not transport._inflight + assert transport.poll() == PollResult(done=(), failed=()) + if failure_method == "make_prepped_xfer": + agent.transfer.assert_not_called() + agent.release_xfer_handle.assert_not_called() + else: + agent.release_xfer_handle.assert_called_once_with( + agent.make_prepped_xfer.return_value + ) + def test_write_blocks_increments_transfer_id(self): """Each write_blocks call gets a unique transfer_id.""" transport = self._make_transport() diff --git a/tests/v1/kv_offload/tiering/p2p/test_sessions.py b/tests/v1/kv_offload/tiering/p2p/test_sessions.py index 357f2b203bf1..debdd33c5659 100644 --- a/tests/v1/kv_offload/tiering/p2p/test_sessions.py +++ b/tests/v1/kv_offload/tiering/p2p/test_sessions.py @@ -1928,18 +1928,22 @@ def test_finish_request_remaining_zero_emits_success_via_inflight(self): assert StoreResult(job_id=7, success=True) in stores assert _srv_outbound(session, "req-1") is None - def test_write_blocks_failure_finalizes_with_failure(self): - """write_blocks returning None must not leave the request hanging. - - The matched blocks are gone from req.demanded but no inflight - will satisfy them, so remaining > 0 forever. Setting finishing - and calling _finalize_outbound(success=False) immediately (no - other inflight) tells the peer + emits StoreResult(success=False) - without waiting on _STORE_TIMEOUT_S or _LOAD_TIMEOUT_S. - """ + @pytest.mark.parametrize( + "store_first,finish_before_fetch", + [(False, False), (True, False), (True, True)], + ids=["fetch-first", "store-first", "finish-before-fetch"], + ) + def test_write_blocks_failure_finalizes_with_failure( + self, store_first, finish_before_fetch + ): + """Submission failure finalizes once in either store/fetch ordering.""" session, conn, transport = _make_session() _activate(session, conn) - # Decoder demands b"k1". + transport.write_blocks = lambda *a, **kw: None # type: ignore[assignment] + if store_first: + session.add_stored_blocks("req-1", [b"k1"], [0], job_id=42) + if finish_before_fetch: + session.finish_request("req-1") conn.enqueue( { TYPE_KEY: FetchMsg.TYPE, @@ -1949,21 +1953,22 @@ def test_write_blocks_failure_finalizes_with_failure(self): FetchMsg.BLOCK_INDEXES: [10], } ) - session.poll() - # Force write_blocks to fail on the next call. - transport.write_blocks = lambda *a, **kw: None # type: ignore[assignment] - - session.add_stored_blocks("req-1", [b"k1"], [0], job_id=42) + stores = session.poll().stores + if not store_first: + session.add_stored_blocks("req-1", [b"k1"], [0], job_id=42) + stores.extend(session.poll().stores) # Outbound was finalized immediately (no other inflight). assert _srv_outbound(session, "req-1") is None + assert session._dispatch_error_count == 0 + assert stores == [StoreResult(job_id=42, success=False)] + assert session.poll().stores == [] # Peer notified with success=False. - msg = next(m for m in conn._sent if m[TYPE_KEY] == TransferDoneMsg.TYPE) + done = [m for m in conn._sent if m[TYPE_KEY] == TransferDoneMsg.TYPE] + assert len(done) == 1 + msg = done[0] assert msg[TransferDoneMsg.KV_REQUEST_ID] == "req-1" assert msg[TransferDoneMsg.SUCCESS] is False - # Local store job surfaces on the next poll. - stores = session.poll().stores - assert StoreResult(job_id=42, success=False) in stores assert 42 not in session._server._store_jobs def test_partial_match_completes_in_two_rounds(self): diff --git a/vllm/v1/kv_offload/tiering/p2p/data/nixl.py b/vllm/v1/kv_offload/tiering/p2p/data/nixl.py index bf3adf4b604b..1cfd8b20f0b3 100644 --- a/vllm/v1/kv_offload/tiering/p2p/data/nixl.py +++ b/vllm/v1/kv_offload/tiering/p2p/data/nixl.py @@ -158,7 +158,8 @@ def write_blocks( ) -> int | None: """Submit a WRITE transfer to *peer_id*. - Returns a transfer ID, or None if the peer is not registered. + Returns a transfer ID, or None if the peer is not registered or + submission fails. The ID is returned via poll() when the transfer completes or fails. """ remote_dlist = self._remote_dlists.get(peer_id) @@ -177,14 +178,27 @@ def write_blocks( peer_id, len(local_idxs), ) - handle = self._agent.make_prepped_xfer( - "WRITE", - self._local_dlist, - np.asarray(local_idxs, dtype=np.int32), - remote_dlist, - np.asarray(remote_idxs, dtype=np.int32), - ) - self._agent.transfer(handle) + handle = None + try: + handle = self._agent.make_prepped_xfer( + "WRITE", + self._local_dlist, + np.asarray(local_idxs, dtype=np.int32), + remote_dlist, + np.asarray(remote_idxs, dtype=np.int32), + ) + self._agent.transfer(handle) + except Exception as exc: + logger.warning( + "NixlTransport %s: write_blocks failed for peer=%s: %s", + self._agent_name, + peer_id, + exc, + exc_info=True, + ) + if handle is not None: + self._release_handles([handle]) + return None transfer_id = next(self._next_id) self._inflight[transfer_id] = _Inflight(peer_id, handle) return transfer_id diff --git a/vllm/v1/kv_offload/tiering/p2p/session/server.py b/vllm/v1/kv_offload/tiering/p2p/session/server.py index 7d24f3773fb3..9d4ac76f3a2d 100644 --- a/vllm/v1/kv_offload/tiering/p2p/session/server.py +++ b/vllm/v1/kv_offload/tiering/p2p/session/server.py @@ -388,6 +388,7 @@ def on_fetch( return if result.local_idxs: self._submit_transfer(kv_request_id, result, req, round_seq) + return # Prefiller-first mode: finish_request may have run before # fetch arrived. If so, finalize once we know what was # demanded — fully satisfied → success, else early-fail.