[Bugfix] Handle P2P NIXL submission failures

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()

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):

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

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.