[Bugfix][KV Offload][P2P] Keep store pins and AbortAck until NIXL transfers stop (#59790)

Signed-off-by: apex-mochen <2756823972@qq.com>
Co-authored-by: Codex <noreply@openai.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Co-authored-by: Kevin H. Luu <khluu000@gmail.com>

tests/v1/kv_offload/tiering/p2p/test_data_transport.py

@@ -14,6 +14,15 @@ import pytest
 from vllm.v1.kv_offload.tiering.p2p.data.base import PollResult
 from vllm.v1.kv_offload.tiering.p2p.data.nixl import NixlTransport
 
+
+class nixlRemoteDisconnectError(Exception):  # noqa: N801
+    """Stand-in for NIXL's binding exception of the same name."""
+
+
+def _remote_disconnect() -> Exception:
+    return nixlRemoteDisconnectError("NIXL_ERR_REMOTE_DISCONNECT")
+
+
 # ---------------------------------------------------------------------------
 # DataTransport base class tests
 # ---------------------------------------------------------------------------
@@ -286,25 +295,101 @@ class TestNixlTransportWithMockedAgent:
         assert transport.cancel([999, 1000]) == []
         assert transport.cancel([999, 1000], mode="wait") == []
 
-    def test_cancel_wait_release_succeeds(self):
-        """wait-mode cancel that succeeds pops the entry and returns []."""
+    @pytest.mark.parametrize("state", ["DONE", "ERR"])
+    def test_cancel_wait_release_succeeds(self, state):
+        """wait-mode cancel of a stopped transfer pops it and returns []."""
         transport = self._make_transport()
         transport.add_remote_peer("peer:1", b"meta", 0x1000, 8, 1024)
 
         tid = transport.write_blocks("peer:1", [0], [1])
         assert tid in transport._inflight
 
+        transport._agent.check_xfer_state.return_value = state
         result = transport.cancel([tid], mode="wait")
         assert result == []
         assert tid not in transport._inflight
         transport._agent.release_xfer_handle.assert_called_once()
 
+    @pytest.mark.parametrize("state", ["PROC", "PEND"])
+    def test_cancel_wait_keeps_running_transfer(self, state):
+        """Releasing a running NIXL WRITE does not stop it, so wait-mode
+        cancel must not release it or report it as cancelled."""
+        transport = self._make_transport()
+        transport.add_remote_peer("peer:1", b"meta", 0x1000, 8, 1024)
+
+        tid = transport.write_blocks("peer:1", [0], [1])
+        transport._agent.check_xfer_state.return_value = state
+
+        assert transport.cancel([tid], mode="wait") == [tid]
+        assert tid in transport._inflight
+        transport._agent.release_xfer_handle.assert_not_called()
+
+    def test_cancel_wait_state_check_raises(self):
+        """A failed state check is not proof the transfer stopped."""
+        transport = self._make_transport()
+        transport.add_remote_peer("peer:1", b"meta", 0x1000, 8, 1024)
+
+        tid = transport.write_blocks("peer:1", [0], [1])
+        transport._agent.check_xfer_state.side_effect = RuntimeError("NOT_FOUND")
+
+        assert transport.cancel([tid], mode="wait") == [tid]
+        assert tid in transport._inflight
+        transport._agent.release_xfer_handle.assert_not_called()
+
+    def test_cancel_wait_remote_disconnect_is_terminal(self):
+        """A dead peer can no longer be written to: wait-mode cancel drops
+        the transfer even if its handle can no longer be released."""
+        transport = self._make_transport()
+        transport.add_remote_peer("peer:1", b"meta", 0x1000, 8, 1024)
+
+        tid = transport.write_blocks("peer:1", [0], [1])
+        transport._agent.check_xfer_state.side_effect = _remote_disconnect()
+        transport._agent.release_xfer_handle.side_effect = RuntimeError("gone")
+
+        assert transport.cancel([tid], mode="wait") == []
+        assert tid not in transport._inflight
+        assert transport.poll() == PollResult(done=(), failed=())
+
+    @pytest.mark.parametrize(
+        "exc",
+        [
+            _remote_disconnect(),
+            # Recognized by status code even if the binding class is renamed.
+            RuntimeError("getXferStatus: NIXL_ERR_REMOTE_DISCONNECT"),
+        ],
+    )
+    def test_poll_reports_remote_disconnect_as_failed(self, exc):
+        """check_xfer_state raises instead of returning ERR once the peer
+        disconnected; poll must surface that as a failure, not keep the
+        transfer inflight forever."""
+        transport = self._make_transport()
+        transport.add_remote_peer("peer:1", b"meta", 0x1000, 8, 1024)
+
+        tid = transport.write_blocks("peer:1", [0], [1])
+        transport._agent.check_xfer_state.side_effect = exc
+
+        result = transport.poll()
+        assert result.done == ()
+        assert tid in result.failed
+        assert tid not in transport._inflight
+
+    def test_poll_keeps_transfer_on_other_state_check_errors(self):
+        transport = self._make_transport()
+        transport.add_remote_peer("peer:1", b"meta", 0x1000, 8, 1024)
+
+        tid = transport.write_blocks("peer:1", [0], [1])
+        transport._agent.check_xfer_state.side_effect = RuntimeError("NOT_FOUND")
+
+        assert transport.poll() == PollResult(done=(), failed=())
+        assert tid in transport._inflight
+
     def test_cancel_wait_release_raises(self):
         """wait-mode cancel keeps the entry and returns the tid on raise."""
         transport = self._make_transport()
         transport.add_remote_peer("peer:1", b"meta", 0x1000, 8, 1024)
 
         tid = transport.write_blocks("peer:1", [0], [1])
+        transport._agent.check_xfer_state.return_value = "DONE"
         transport._agent.release_xfer_handle.side_effect = RuntimeError(
             "NIXL_ERR_REPOST_ACTIVE"
         )

tests/v1/kv_offload/tiering/p2p/test_sessions.py

@@ -15,10 +15,12 @@ from __future__ import annotations
 import time
 from collections.abc import Callable, Sequence
 from concurrent.futures import Future
+from types import SimpleNamespace
 
 import numpy as np
 import pytest
 
+import vllm.v1.kv_offload.tiering.p2p.session.server as server_mod
 from vllm.v1.kv_offload.base import (
     LookupResult,
     OffloadKey,
@@ -1632,9 +1634,11 @@ class TestServerFlows:
         assert _srv_abort_started(session, "req-1") is None
         assert tid not in session._server._inflight
 
-    def test_abort_fetch_force_cancels_after_timeout(self):
-        """If wait-mode never drains, the deadline forces immediate
-        cancel and an ack is still sent."""
+    def test_abort_fetch_never_force_cancels_after_timeout(self):
+        """Past the drain deadline a transfer that is still running is not
+        force-released and no ack is sent: releasing the handle or removing
+        the peer does not stop a NIXL WRITE, so the ack would let the peer
+        reuse blocks that are still being written."""
         session, conn, transport = _make_session()
         _activate(session, conn)
         tid = 42
@@ -1658,21 +1662,28 @@ class TestServerFlows:
         )
         session.poll()
         assert _srv_abort_started(session, "req-1") is not None
+        peers = dict(transport._remote_peers)
         # Backdate past the drain deadline.
         session._server._pending_aborts[("req-1", 0)] = (
             time.monotonic() - _CANCEL_DRAIN_TIMEOUT_S - 1.0
         )
-        # Even if the transport still claims it can't cancel, the
-        # session must force-pop and ack.
         transport._cancel_calls.clear()
 
-        session.poll()
+        for _ in range(3):
+            session.poll()
+        assert not any(m[TYPE_KEY] == AbortAckMsg.TYPE for m in conn._sent)
+        assert tid in session._server._inflight
+        assert all(mode == "wait" for _, mode in transport._cancel_calls)
+        assert transport._remote_peers == peers
 
-        ack = next(m for m in conn._sent if m[TYPE_KEY] == AbortAckMsg.TYPE)
-        assert ack[AbortAckMsg.KV_REQUEST_ID] == "req-1"
+        # Once NIXL reports the transfer terminal, the ack follows.
+        transport._cancel_still_inflight.discard(tid)
+        transport._poll_done.append(tid)
+        session.poll()
+        acks = [m for m in conn._sent if m[TYPE_KEY] == AbortAckMsg.TYPE]
+        assert len(acks) == 1
         assert _srv_abort_started(session, "req-1") is None
         assert tid not in session._server._inflight
-        assert ([tid], "immediate") in transport._cancel_calls
 
     def test_abort_fetch_idempotent_while_draining(self):
         """Receiving AbortFetchMsg twice for the same kv_request_id
@@ -1731,10 +1742,12 @@ class TestServerFlows:
         stores = session.poll().stores
         assert StoreResult(job_id=1, success=False) in stores
 
-    def test_store_timeout_then_late_completion_no_duplicate(self):
-        """A job timed out by _timeout_pending_store_jobs must not also
-        emit a contradictory StoreResult(success=True) when the transport
-        later reports the same transfer as done."""
+    @pytest.mark.parametrize("cancel_pending", [False, True])
+    def test_store_timeout_cancels_stuck_transfer_after_confirmation(
+        self, cancel_pending, monkeypatch
+    ):
+        # Retry on every poll; pacing is covered by the *_paced tests.
+        monkeypatch.setattr(server_mod, "_CANCEL_RETRY_INTERVAL_S", 0.0)
         session, conn, transport = _make_session()
         _activate(session, conn)
         session.add_stored_blocks("req-1", [b"k1"], [0], job_id=1)
@@ -1749,25 +1762,143 @@ class TestServerFlows:
         )
         session.poll()
         tid = next(iter(transport._transfers))
-
-        # Backdate the store job so the next poll times it out.
         session._server._store_jobs[1] = time.monotonic() - 60.0
-        stores = session.poll().stores
-        assert StoreResult(job_id=1, success=False) in stores
-        assert StoreResult(job_id=1, success=True) not in stores
+        if cancel_pending:
+            transport._cancel_still_inflight.add(tid)
+            for _ in range(3):
+                assert session.poll().stores == []
+                assert 1 in session._server._store_jobs
+                assert tid in session._server._inflight
+                assert not any(m[TYPE_KEY] == TransferDoneMsg.TYPE for m in conn._sent)
+            transport._cancel_still_inflight.clear()
+        assert session.poll().stores == [StoreResult(job_id=1, success=False)]
+        assert transport._cancel_calls
+        assert all(mode == "wait" for _, mode in transport._cancel_calls)
+        assert tid not in session._server._inflight
+        assert tid not in transport._transfers
+        terminal = [m for m in conn._sent if m[TYPE_KEY] == TransferDoneMsg.TYPE]
+        assert len(terminal) == 1
+        assert terminal[0][TransferDoneMsg.SUCCESS] is False
+        assert session.poll().stores == []
 
-        # Transport later reports the same transfer as done — must not
-        # emit a second (contradictory) StoreResult for job_id=1.
-        transport._poll_done.append(tid)
-        stores = session.poll().stores
-        assert all(s.job_id != 1 for s in stores), (
-            f"unexpected duplicate StoreResult after timeout: {stores}"
+    @pytest.mark.parametrize("success", [True, False])
+    def test_store_timeout_keeps_pins_until_cancel_or_late_terminal(self, success):
+        session, conn, transport = _make_session()
+        _activate(session, conn)
+        session.add_stored_blocks("req-1", [b"k1"], [0], job_id=1)
+        session.add_stored_blocks("req-2", [b"k2"], [1], job_id=2)
+        conn.enqueue(
+            {
+                TYPE_KEY: FetchMsg.TYPE,
+                FetchMsg.ROUND_SEQ: 0,
+                FetchMsg.KV_REQUEST_ID: "req-1",
+                FetchMsg.KEYS: [b"k1"],
+                FetchMsg.BLOCK_INDEXES: [5],
+            }
         )
+        session.poll()
+        tid = next(iter(transport._transfers))
+        transport._cancel_still_inflight.add(tid)
 
-    def test_store_timeout_then_late_failure_no_duplicate(self):
-        """Symmetric guard: a timed-out job must not also emit a second
-        StoreResult(success=False) when the transport later reports the
-        same transfer as failed."""
+        for jid in (1, 2):
+            session._server._store_jobs[jid] = time.monotonic() - 60.0
+        # A StoreResult releases the primary slot. Only the job still
+        # waiting for a fetch may time out while the transfer is active.
+        assert session.poll().stores == [StoreResult(job_id=2, success=False)]
+        assert session.poll().stores == []
+        assert not any(m[TYPE_KEY] == TransferDoneMsg.TYPE for m in conn._sent)
+
+        completed = transport._poll_done if success else transport._poll_failed
+        completed.append(tid)
+        assert session.poll().stores == [StoreResult(job_id=1, success=False)]
+        terminal = [m for m in conn._sent if m[TYPE_KEY] == TransferDoneMsg.TYPE]
+        assert len(terminal) == 1
+        assert terminal[0][TransferDoneMsg.SUCCESS] is False
+        assert session.poll().stores == []
+
+    @pytest.mark.parametrize("success", [True, False, None])
+    def test_store_timeout_waits_for_other_transfers_in_round(
+        self, success, monkeypatch
+    ):
+        monkeypatch.setattr(server_mod, "_CANCEL_RETRY_INTERVAL_S", 0.0)
+        session, conn, transport = _make_session()
+        _activate(session, conn)
+        conn.enqueue(
+            {
+                TYPE_KEY: FetchMsg.TYPE,
+                FetchMsg.ROUND_SEQ: 0,
+                FetchMsg.KV_REQUEST_ID: "req-1",
+                FetchMsg.KEYS: [b"k1", b"k2"],
+                FetchMsg.BLOCK_INDEXES: [5, 6],
+            }
+        )
+        session.poll()
+        session.add_stored_blocks("req-1", [b"k1"], [0], job_id=1)
+        session.add_stored_blocks("req-1", [b"k2"], [1], job_id=2)
+        first, second = list(transport._transfers)
+        session._server._store_jobs[1] = time.monotonic() - 60.0
+        transport._cancel_still_inflight.add(second)
+        assert session.poll().stores == [StoreResult(job_id=1, success=False)]
+        assert first not in session._server._inflight
+        assert second in session._server._inflight
+        assert 2 in session._server._store_jobs
+        assert not any(m[TYPE_KEY] == TransferDoneMsg.TYPE for m in conn._sent)
+        # Additional supply must not restart the timed-out round.
+        session.add_stored_blocks("req-1", [b"k3"], [2], job_id=3)
+        assert session.poll().stores == [StoreResult(job_id=3, success=False)]
+        assert set(session._server._inflight) == {second}
+        if success is None:
+            transport._cancel_still_inflight.clear()
+        else:
+            completed = transport._poll_done if success else transport._poll_failed
+            completed.append(second)
+        assert session.poll().stores == [StoreResult(job_id=2, success=False)]
+        terminal = [m for m in conn._sent if m[TYPE_KEY] == TransferDoneMsg.TYPE]
+        assert len(terminal) == 1
+        assert terminal[0][TransferDoneMsg.SUCCESS] is False
+        assert session.poll().stores == []
+
+    def test_failed_transfer_keeps_sibling_pins_until_terminal(self, monkeypatch):
+        """A failed transfer fails its round, but a sibling transfer that may
+        still be writing is not force-cancelled: its pin and the round's
+        TransferDoneMsg wait until it is terminal."""
+        monkeypatch.setattr(server_mod, "_CANCEL_RETRY_INTERVAL_S", 0.0)
+        session, conn, transport = _make_session()
+        _activate(session, conn)
+        conn.enqueue(
+            {
+                TYPE_KEY: FetchMsg.TYPE,
+                FetchMsg.ROUND_SEQ: 0,
+                FetchMsg.KV_REQUEST_ID: "req-1",
+                FetchMsg.KEYS: [b"k1", b"k2"],
+                FetchMsg.BLOCK_INDEXES: [5, 6],
+            }
+        )
+        session.poll()
+        session.add_stored_blocks("req-1", [b"k1"], [0], job_id=1)
+        session.add_stored_blocks("req-1", [b"k2"], [1], job_id=2)
+        first, second = list(transport._transfers)
+        transport._cancel_still_inflight.add(second)
+
+        transport._poll_failed.append(first)
+        assert session.poll().stores == [StoreResult(job_id=1, success=False)]
+        for _ in range(3):
+            assert session.poll().stores == []
+        assert second in session._server._inflight
+        assert 2 in session._server._store_jobs
+        assert transport._cancel_calls
+        assert all(mode == "wait" for _, mode in transport._cancel_calls)
+        assert not any(m[TYPE_KEY] == TransferDoneMsg.TYPE for m in conn._sent)
+
+        transport._cancel_still_inflight.clear()
+        assert session.poll().stores == [StoreResult(job_id=2, success=False)]
+        terminal = [m for m in conn._sent if m[TYPE_KEY] == TransferDoneMsg.TYPE]
+        assert len(terminal) == 1
+        assert terminal[0][TransferDoneMsg.SUCCESS] is False
+        assert session.poll().stores == []
+
+    def test_abort_during_store_timeout_keeps_unconfirmed_pins(self, monkeypatch):
+        monkeypatch.setattr(server_mod, "_CANCEL_RETRY_INTERVAL_S", 0.0)
         session, conn, transport = _make_session()
         _activate(session, conn)
         session.add_stored_blocks("req-1", [b"k1"], [0], job_id=1)
@@ -1782,18 +1913,219 @@ class TestServerFlows:
         )
         session.poll()
         tid = next(iter(transport._transfers))
-
+        transport._cancel_still_inflight.add(tid)
         session._server._store_jobs[1] = time.monotonic() - 60.0
-        stores = session.poll().stores
-        assert [s for s in stores if s.job_id == 1] == [
-            StoreResult(job_id=1, success=False)
-        ]
-
-        transport._poll_failed.append(tid)
-        stores = session.poll().stores
-        assert all(s.job_id != 1 for s in stores), (
-            f"unexpected duplicate StoreResult after timeout: {stores}"
+        assert session.poll().stores == []
+        conn.enqueue(
+            {
+                TYPE_KEY: AbortFetchMsg.TYPE,
+                AbortFetchMsg.KV_REQUEST_ID: "req-1",
+                AbortFetchMsg.ROUND_SEQ: 0,
+            }
         )
+        assert session.poll().stores == []
+        session._server._pending_aborts[("req-1", 0)] = time.monotonic() - 60.0
+        assert session.poll().stores == []
+        assert 1 in session._server._store_jobs
+        assert tid in session._server._inflight
+        assert all(mode == "wait" for _, mode in transport._cancel_calls)
+        assert not any(m[TYPE_KEY] == AbortAckMsg.TYPE for m in conn._sent)
+        transport._cancel_still_inflight.clear()
+        assert session.poll().stores == [StoreResult(job_id=1, success=False)]
+        assert sum(m[TYPE_KEY] == AbortAckMsg.TYPE for m in conn._sent) == 1
+        assert session.poll().stores == []
+
+    @staticmethod
+    def _fake_clock(monkeypatch) -> list[float]:
+        clock = [1000.0]
+        monkeypatch.setattr(
+            server_mod, "time", SimpleNamespace(monotonic=lambda: clock[0])
+        )
+        return clock
+
+    @staticmethod
+    def _submit_unconfirmed_transfer(session, conn, transport) -> int:
+        """Start one transfer for job 1 whose wait-cancel never confirms."""
+        session.add_stored_blocks("req-1", [b"k1"], [0], job_id=1)
+        conn.enqueue(
+            {
+                TYPE_KEY: FetchMsg.TYPE,
+                FetchMsg.ROUND_SEQ: 0,
+                FetchMsg.KV_REQUEST_ID: "req-1",
+                FetchMsg.KEYS: [b"k1"],
+                FetchMsg.BLOCK_INDEXES: [5],
+            }
+        )
+        session.poll()
+        tid = next(iter(transport._transfers))
+        transport._cancel_still_inflight.add(tid)
+        return tid
+
+    @staticmethod
+    def _wait_cancels(transport) -> int:
+        return sum(1 for _, mode in transport._cancel_calls if mode == "wait")
+
+    def test_store_timeout_cancel_retry_is_paced(self, monkeypatch):
+        clock = self._fake_clock(monkeypatch)
+        session, conn, transport = _make_session()
+        _activate(session, conn)
+        tid = self._submit_unconfirmed_transfer(session, conn, transport)
+        session._server._store_jobs[1] = clock[0] - 60.0
+
+        assert session.poll().stores == []
+        assert self._wait_cancels(transport) == 1
+        start = clock[0]
+        for dt in (0.0, 0.05, 0.099):
+            clock[0] = start + dt
+            for _ in range(5):
+                assert session.poll().stores == []
+        assert self._wait_cancels(transport) == 1
+
+        clock[0] = start + server_mod._CANCEL_RETRY_INTERVAL_S + 1e-3
+        assert session.poll().stores == []
+        assert self._wait_cancels(transport) == 2
+
+        transport._cancel_still_inflight.clear()
+        clock[0] += server_mod._CANCEL_RETRY_INTERVAL_S + 1e-3
+        assert session.poll().stores == [StoreResult(job_id=1, success=False)]
+        assert self._wait_cancels(transport) == 3
+        assert tid not in session._server._inflight
+
+    def test_abort_drain_cancel_retry_is_paced(self, monkeypatch):
+        clock = self._fake_clock(monkeypatch)
+        session, conn, transport = _make_session()
+        _activate(session, conn)
+        self._submit_unconfirmed_transfer(session, conn, transport)
+        conn.enqueue(
+            {
+                TYPE_KEY: AbortFetchMsg.TYPE,
+                AbortFetchMsg.KV_REQUEST_ID: "req-1",
+                AbortFetchMsg.ROUND_SEQ: 0,
+            }
+        )
+        # The first drain attempt is immediate.
+        session.poll()
+        assert self._wait_cancels(transport) == 1
+        for _ in range(5):
+            session.poll()
+        assert self._wait_cancels(transport) == 1
+
+        clock[0] += server_mod._CANCEL_RETRY_INTERVAL_S + 1e-3
+        session.poll()
+        assert self._wait_cancels(transport) == 2
+        assert not any(m[TYPE_KEY] == AbortAckMsg.TYPE for m in conn._sent)
+
+        transport._cancel_still_inflight.clear()
+        clock[0] += server_mod._CANCEL_RETRY_INTERVAL_S + 1e-3
+        session.poll()
+        assert sum(m[TYPE_KEY] == AbortAckMsg.TYPE for m in conn._sent) == 1
+
+    def test_unconfirmed_store_cancel_never_acks_or_tears_down(self, monkeypatch):
+        """A cancellation that is never confirmed must not release the pin,
+        send AbortAck, force-cancel, or remove the peer, however long it
+        stays pending. AbortAck follows the release once it is confirmed."""
+        clock = self._fake_clock(monkeypatch)
+        session, conn, transport = _make_session()
+        _activate(session, conn)
+        peers = dict(transport._remote_peers)
+        tid = self._submit_unconfirmed_transfer(session, conn, transport)
+        session._server._store_jobs[1] = clock[0] - 60.0
+        assert session.poll().stores == []
+
+        abort = {
+            TYPE_KEY: AbortFetchMsg.TYPE,
+            AbortFetchMsg.KV_REQUEST_ID: "req-1",
+            AbortFetchMsg.ROUND_SEQ: 0,
+        }
+        started_at = clock[0]
+        conn.enqueue(dict(abort))
+        assert session.poll().stores == []
+        # A duplicate abort neither resets the drain start nor adds a cancel.
+        conn.enqueue(dict(abort))
+        assert session.poll().stores == []
+        assert session._server._pending_aborts[("req-1", 0)] == started_at
+        # The store sweep and the abort drain share the round's pacing.
+        assert self._wait_cancels(transport) == 1
+
+        ticks = 50
+        for step in range(1, ticks + 1):
+            clock[0] = started_at + step * 10.0
+            assert session.poll().stores == []
+        assert clock[0] - started_at > _CANCEL_DRAIN_TIMEOUT_S
+        assert self._wait_cancels(transport) == 1 + ticks
+        assert all(mode == "wait" for _, mode in transport._cancel_calls)
+        assert 1 in session._server._store_jobs
+        assert tid in session._server._inflight
+        assert transport._remote_peers == peers
+        assert not any(
+            m[TYPE_KEY] in (AbortAckMsg.TYPE, TransferDoneMsg.TYPE) for m in conn._sent
+        )
+
+        transport._cancel_still_inflight.clear()
+        clock[0] += server_mod._CANCEL_RETRY_INTERVAL_S + 1e-3
+        assert session.poll().stores == [StoreResult(job_id=1, success=False)]
+        assert 1 not in session._server._store_jobs
+        assert sum(m[TYPE_KEY] == AbortAckMsg.TYPE for m in conn._sent) == 1
+        assert session.poll().stores == []
+        assert sum(m[TYPE_KEY] == AbortAckMsg.TYPE for m in conn._sent) == 1
+
+    def test_abort_keeps_pins_only_of_jobs_being_written(self, monkeypatch):
+        """An abort fails the round's jobs at once, except those an inflight
+        transfer still reads; their pins and the ack wait for the transfer."""
+        monkeypatch.setattr(server_mod, "_CANCEL_RETRY_INTERVAL_S", 0.0)
+        session, conn, transport = _make_session()
+        _activate(session, conn)
+        tid = self._submit_unconfirmed_transfer(session, conn, transport)
+        # Job 2 supplies a block nobody demanded: no transfer reads it.
+        session.add_stored_blocks("req-1", [b"k2"], [1], job_id=2)
+        conn.enqueue(
+            {
+                TYPE_KEY: AbortFetchMsg.TYPE,
+                AbortFetchMsg.KV_REQUEST_ID: "req-1",
+                AbortFetchMsg.ROUND_SEQ: 0,
+            }
+        )
+        assert session.poll().stores == [StoreResult(job_id=2, success=False)]
+        for _ in range(3):
+            assert session.poll().stores == []
+        assert 1 in session._server._store_jobs
+        assert tid in session._server._inflight
+        assert not any(m[TYPE_KEY] == AbortAckMsg.TYPE for m in conn._sent)
+
+        transport._cancel_still_inflight.clear()
+        assert session.poll().stores == [StoreResult(job_id=1, success=False)]
+        assert sum(m[TYPE_KEY] == AbortAckMsg.TYPE for m in conn._sent) == 1
+        assert not any(m[TYPE_KEY] == TransferDoneMsg.TYPE for m in conn._sent)
+        assert session.poll().stores == []
+
+    @pytest.mark.parametrize("success", [True, False])
+    def test_abort_acks_after_late_terminal_without_transfer_done(
+        self, success, monkeypatch
+    ):
+        """A transfer that finishes on its own after the abort releases its
+        pin and lets the parked abort ack; the aborted round sends no
+        TransferDoneMsg."""
+        monkeypatch.setattr(server_mod, "_CANCEL_RETRY_INTERVAL_S", 0.0)
+        session, conn, transport = _make_session()
+        _activate(session, conn)
+        tid = self._submit_unconfirmed_transfer(session, conn, transport)
+        conn.enqueue(
+            {
+                TYPE_KEY: AbortFetchMsg.TYPE,
+                AbortFetchMsg.KV_REQUEST_ID: "req-1",
+                AbortFetchMsg.ROUND_SEQ: 0,
+            }
+        )
+        assert session.poll().stores == []
+        assert 1 in session._server._store_jobs
+
+        completed = transport._poll_done if success else transport._poll_failed
+        completed.append(tid)
+        assert session.poll().stores == [StoreResult(job_id=1, success=False)]
+        assert sum(m[TYPE_KEY] == AbortAckMsg.TYPE for m in conn._sent) == 1
+        assert not any(m[TYPE_KEY] == TransferDoneMsg.TYPE for m in conn._sent)
+        assert session.poll().stores == []
+        assert sum(m[TYPE_KEY] == AbortAckMsg.TYPE for m in conn._sent) == 1
 
     def test_fetch_after_symmetric_supply_timed_out_reads_nothing(self):
         """A lookup pin whose store job timed out has been reported failed,

vllm/v1/kv_offload/tiering/p2p/data/base.py

@@ -296,12 +296,12 @@ class DataTransport(ABC):
                     return []. Matches the legacy fire-and-forget
                     behavior — the caller does not wait for the
                     underlying transfer to drain.
-                "wait": attempt to release each handle. If the release
-                    cannot complete because the transfer is still
-                    PROC/PEND, the entry stays in the inflight set and
-                    its id is included in the returned list. The
-                    caller is expected to keep calling poll() until
-                    every returned id surfaces in done/failed.
+                "wait": release only transfers that have stopped. A
+                    transfer still PROC/PEND (or whose release fails)
+                    stays in the inflight set and its id is included in
+                    the returned list. The caller is expected to keep
+                    calling poll() until every returned id surfaces in
+                    done/failed.
 
         Returns:
             For mode="wait", the subset of *transfer_ids* still

vllm/v1/kv_offload/tiering/p2p/data/nixl.py

@@ -29,6 +29,22 @@ logger = init_logger(__name__)
 _EMPTY_POLL_RESULT: PollResult = PollResult(done=(), failed=())
 
 
+def _is_remote_disconnect(exc: Exception) -> bool:
+    """True if a NIXL state check failed because the peer disconnected.
+
+    NIXL raises ``nixlRemoteDisconnectError`` (NIXL_ERR_REMOTE_DISCONNECT)
+    from ``check_xfer_state`` once the target endpoint is gone, e.g. after
+    the peer process died, instead of reporting an ``"ERR"`` state. Matched
+    by name and status code so the bindings need not be imported here; an
+    unrecognized error keeps the transfer inflight, the conservative side.
+    """
+    return type(
+        exc
+    ).__name__ == "nixlRemoteDisconnectError" or "NIXL_ERR_REMOTE_DISCONNECT" in str(
+        exc
+    )
+
+
 class _Inflight(NamedTuple):
     """A submitted-but-not-yet-drained transfer.
 
@@ -325,7 +341,9 @@ class NixlTransport(DataTransport):
                     transfer_id,
                     exc,
                 )
-                continue
+                if not _is_remote_disconnect(exc):
+                    continue
+                state = "ERR"
             if state == "DONE":
                 if done_ids is None:
                     done_ids = []
@@ -357,10 +375,14 @@ class NixlTransport(DataTransport):
     ) -> list[int]:
         """Cancel inflight transfers by their IDs.
 
-        See ``DataTransport.cancel`` for the contract. In "wait" mode,
-        transfers whose ``release_xfer_handle`` raises (NIXL could not
-        complete the abort because the backend is still draining) stay
-        in ``self._inflight`` so a later ``poll()`` will observe them.
+        See ``DataTransport.cancel`` for the contract. In "wait" mode, a
+        transfer is released only once NIXL reports it terminal. Releasing
+        a PROC/PEND handle does not stop the WRITE: the UCX backend frees
+        the request without cancelling RMA, and the bytes keep landing. So
+        PROC/PEND transfers, and those whose state check or
+        ``release_xfer_handle`` raises, stay in ``self._inflight`` for a
+        later ``poll()`` to observe. A state check that fails because the
+        peer disconnected counts as terminal.
         """
         if mode == "immediate":
             handles = [
@@ -376,6 +398,27 @@ class NixlTransport(DataTransport):
             entry = self._inflight.get(tid)
             if entry is None:
                 continue
+            try:
+                state = self._agent.check_xfer_state(entry.handle)
+            except Exception as exc:
+                logger.debug(
+                    "NixlTransport %s: cancel state check failed for "
+                    "transfer_id=%d: %s",
+                    self._agent_name,
+                    tid,
+                    exc,
+                )
+                if not _is_remote_disconnect(exc):
+                    still_inflight.append(tid)
+                    continue
+                # The handle may no longer be releasable; it is terminal
+                # either way, so drop it instead of retrying forever.
+                self._release_handles([entry.handle])
+                del self._inflight[tid]
+                continue
+            if state in ("PROC", "PEND"):
+                still_inflight.append(tid)
+                continue
             try:
                 self._agent.release_xfer_handle(entry.handle)
             except Exception as exc:

vllm/v1/kv_offload/tiering/p2p/manager.py

@@ -1018,11 +1018,12 @@ class P2PSecondaryTierManager(SecondaryTierManager):
     def _drain_inflight_for_shutdown(self) -> None:
         """Best-effort drain of inflight transfers before closing _data.
 
-        Mirrors session._drain_abort but as a single bounded loop. Collects
-        inflight transfer_ids from each session, repeatedly calls
+        Collects inflight transfer_ids from each session, repeatedly calls
         _data.cancel(..., mode="wait") and _data.poll() so handles can
-        surface as done/failed, and falls back to mode="immediate" once
-        _SHUTDOWN_DRAIN_TIMEOUT_S elapses so a wedged peer can't hang us.
+        surface as done/failed. Unlike session._drain_abort, which never
+        force-releases, it falls back to mode="immediate" once
+        _SHUTDOWN_DRAIN_TIMEOUT_S elapses: the engine is exiting, and a
+        wedged peer must not hang shutdown.
         """
         ids = [tid for s in self._sessions.values() for tid in s._server._inflight]
         if not ids:

vllm/v1/kv_offload/tiering/p2p/session/server.py

@@ -39,6 +39,10 @@ logger = init_logger(__name__)
 
 _STORE_TIMEOUT_S = 30.0
 _CANCEL_DRAIN_TIMEOUT_S = 10.0
+# Minimum spacing between wait-mode cancel attempts for one round. Each
+# attempt calls into NIXL on the scheduler thread, so an unconfirmed
+# cancellation is retried at this rate instead of on every poll tick.
+_CANCEL_RETRY_INTERVAL_S = 0.1
 # Cap on time the server holds a HIT_PENDING / RETRY key from an inbound
 # LookupMsg before falling back to MISS. Long enough that an in-flight
 # primary write or a just-started promotion typically completes; short
@@ -85,6 +89,13 @@ class _OutboundRequestState:
     remaining: int = 0  # blocks that need to be transferred to client
     finishing: bool = False  # Signal finish request ASAP
     inflight: int = 0  # transfers submitted for this round, not yet polled
+    # The round was cancelled (a store job timed out, or the peer aborted
+    # it) while transfers may still be writing to the peer. Its jobs fail,
+    # but a job keeps its pin until every transfer carrying it is terminal:
+    # NIXL cannot stop a WRITE that is already in flight.
+    cancelled: bool = False
+    # Last wait-mode cancel attempt for this round's transfers (monotonic).
+    last_cancel_at: float | None = None
     # Job IDs that submit_store'd blocks for this round and have not
     # yet emitted a StoreResult. The terminal-finalize helper drains
     # this set; poll-done and poll-failed discard entries as their
@@ -272,6 +283,9 @@ class ServerRole:
         # Parked aborts awaiting drain, keyed by (kv_request_id, round)
         # with the abort start time.
         self._pending_aborts: dict[tuple[str, int], float] = {}
+        # Parked aborts already warned about for exceeding
+        # _CANCEL_DRAIN_TIMEOUT_S.
+        self._slow_aborts: set[tuple[str, int]] = set()
 
     # ------------------------------------------------------------------
     # State helpers
@@ -322,6 +336,10 @@ class ServerRole:
         rnd = st.outbound.get(round_seq)
         if rnd is None:
             rnd = st.outbound[round_seq] = _OutboundRequestState()
+        if rnd.cancelled:
+            self._store_jobs.pop(job_id)
+            self._pending_store_results.append(StoreResult(job_id, success=False))
+            return
         if from_lookup:
             rnd.lookup_supplied = True
         result = rnd.add_stored_blocks(keys, block_ids, job_id)
@@ -765,7 +783,9 @@ class ServerRole:
                     tid,
                 )
                 continue
-            results.extend(self._settle_xfer_jobs(xfer, success=True))
+            results.extend(
+                self._settle_xfer_jobs(xfer, success=not xfer.round.cancelled)
+            )
             rnd = xfer.round
             st = self._requests.get(xfer.kv_request_id)
             if st is not None and st.outbound.get(xfer.round_key) is rnd:
@@ -773,7 +793,10 @@ class ServerRole:
                 assert rnd.remaining >= 0, (
                     f"remaining went negative for kv_request_id={xfer.kv_request_id}"
                 )
-                if rnd.remaining == 0:
+                if rnd.cancelled:
+                    if rnd.inflight == 0:
+                        self._finalize_cancelled(xfer.kv_request_id, xfer.round_key)
+                elif rnd.remaining == 0:
                     self._finalize_outbound(
                         xfer.kv_request_id, xfer.round_key, success=True
                     )
@@ -783,7 +806,6 @@ class ServerRole:
                     )
             self._maybe_prune(xfer.kv_request_id)
 
-        failed_rounds: list[tuple[str, _OutboundRequestState]] | None = None
         for tid in poll_result.failed:
             xfer = self._inflight_pop(tid)
             if xfer is None:
@@ -796,38 +818,21 @@ class ServerRole:
                     tid,
                 )
                 continue
-            results.extend(self._settle_xfer_jobs(xfer, success=False))
             rnd = xfer.round
             st = self._requests.get(xfer.kv_request_id)
-            if st is not None and st.outbound.get(xfer.round_key) is rnd:
-                del st.outbound[xfer.round_key]
-                if failed_rounds is None:
-                    failed_rounds = []
-                failed_rounds.append((xfer.kv_request_id, rnd))
-                self._send(
-                    {
-                        TYPE_KEY: TransferDoneMsg.TYPE,
-                        TransferDoneMsg.KV_REQUEST_ID: xfer.kv_request_id,
-                        TransferDoneMsg.SUCCESS: False,
-                        TransferDoneMsg.ROUND_SEQ: xfer.round_key,
-                    }
-                )
+            in_round = st is not None and st.outbound.get(xfer.round_key) is rnd
+            if in_round and not rnd.cancelled:
+                # One failed transfer fails the round, but its other
+                # transfers may still be writing. Cancel them like a
+                # timed-out round: their pins and the TransferDoneMsg wait
+                # until they are terminal.
+                rnd.cancelled = True
+                results.extend(self._fail_unwritten_round_jobs(rnd))
+            results.extend(self._settle_xfer_jobs(xfer, success=False))
+            if in_round and rnd.inflight == 0:
+                self._finalize_cancelled(xfer.kv_request_id, xfer.round_key)
             self._maybe_prune(xfer.kv_request_id)
 
-        # Cancel each failed round's other inflight and fail its
-        # remaining store jobs — nothing else will settle them.
-        if failed_rounds:
-            for kv_request_id, rnd in failed_rounds:
-                ids_to_cancel = [
-                    tid for tid, x in self._inflight.items() if x.round is rnd
-                ]
-                for tid in ids_to_cancel:
-                    self._inflight_pop(tid)
-                if ids_to_cancel:
-                    self._transport.cancel(ids_to_cancel)
-                results.extend(self._fail_round_jobs(rnd))
-                self._maybe_prune(kv_request_id)
-
         return results
 
     def collect_idle_timeouts(self) -> list[StoreResult]:
@@ -870,6 +875,7 @@ class ServerRole:
         self._requests.clear()
         self._serve_pending.clear()
         self._pending_aborts.clear()
+        self._slow_aborts.clear()
         self._finished_lookup_ctxs.clear()
         return failed_stores, failed_serves
 
@@ -918,7 +924,10 @@ class ServerRole:
         result.
         """
         results: list[StoreResult] = []
+        still_pinned = self._read_job_ids() if xfer.round.cancelled else set()
         for job_id in xfer.job_ids:
+            if job_id in still_pinned:
+                continue
             if self._store_jobs.pop(job_id, None) is None:
                 continue
             results.append(StoreResult(job_id=job_id, success=success))
@@ -939,6 +948,48 @@ class ServerRole:
         rnd.pending_job_ids.clear()
         return results
 
+    def _finalize_cancelled(self, kv_request_id: str, round_key: int) -> None:
+        """Finalize a cancelled round once its last transfer is terminal.
+
+        An aborted round gets no TransferDoneMsg: its parked abort answers
+        the peer with AbortAckMsg on the next drain.
+        """
+        self._finalize_outbound(
+            kv_request_id,
+            round_key,
+            success=False,
+            send_done=(kv_request_id, round_key) not in self._pending_aborts,
+        )
+
+    def _read_job_ids(self, rnd: _OutboundRequestState | None = None) -> set[int]:
+        """Jobs whose blocks an inflight transfer (of ``rnd``, if given)
+        still reads; their pins must be held."""
+        return {
+            jid
+            for xfer in self._inflight.values()
+            if rnd is None or xfer.round is rnd
+            for jid in xfer.job_ids
+        }
+
+    def _fail_unwritten_round_jobs(
+        self, rnd: _OutboundRequestState
+    ) -> list[StoreResult]:
+        """Fail a cancelled round's jobs that no inflight transfer carries.
+
+        Jobs that an inflight transfer carries keep their pins; they settle
+        when that transfer turns terminal (``_settle_xfer_jobs``).
+        """
+        written = self._read_job_ids(rnd)
+        results: list[StoreResult] = []
+        for job_id in list(rnd.pending_job_ids):
+            if job_id in written:
+                continue
+            rnd.pending_job_ids.discard(job_id)
+            if self._store_jobs.pop(job_id, None) is None:
+                continue
+            results.append(StoreResult(job_id=job_id, success=False))
+        return results
+
     def _finalize_outbound(
         self,
         kv_request_id: str,
@@ -953,7 +1004,12 @@ class ServerRole:
         fetch). Other rounds of the id are untouched.
         """
         st = self._requests[kv_request_id]
-        req = st.outbound.pop(round_key)
+        req = st.outbound[round_key]
+        if req.cancelled and req.inflight:
+            return
+        del st.outbound[round_key]
+        if req.cancelled:
+            success = False
         if success is None:
             success = req.demand_received and req.remaining == 0
         settled = self._fail_round_jobs(req) if not success else None
@@ -992,41 +1048,56 @@ class ServerRole:
     def _drain_abort(self, kv_request_id: str, round_seq: int) -> None:
         """One drain attempt for a pending abort.
 
-        Detaches the aborted round, then asks the transport to cancel its
-        inflight transfers in ``mode="wait"``. Sends ``AbortAckMsg`` once
-        nothing remains inflight, or after ``_CANCEL_DRAIN_TIMEOUT_S``
-        falls back to ``mode="immediate"`` and acks anyway.
+        Cancels the aborted round, then asks the transport to cancel its
+        inflight transfers in ``mode="wait"``. Sends ``AbortAckMsg`` only
+        once none of them can still write to the peer: the ack lets the
+        peer reuse the destination blocks. A transfer that NIXL does not
+        report terminal is never force-released, however long the drain
+        takes: neither releasing its handle nor removing the remote agent
+        stops a WRITE that is already in flight. Past
+        ``_CANCEL_DRAIN_TIMEOUT_S`` the drain only logs a warning.
         """
         st = self._requests[kv_request_id]
-        rnd = st.outbound.pop(round_seq, None)
-        if rnd is not None:
-            # Its transfers are being cancelled; fail its jobs now
-            # instead of leaking them to the store timeout.
-            self._pending_store_results.extend(self._fail_round_jobs(rnd))
+        rnd = st.outbound.get(round_seq)
+        if rnd is not None and not rnd.cancelled:
+            # Fail the round's jobs now instead of leaking them to the
+            # store timeout, but keep the pins of those still being written.
+            rnd.cancelled = True
+            self._pending_store_results.extend(self._fail_unwritten_round_jobs(rnd))
         ids = [
             tid
             for tid, x in self._inflight.items()
             if x.kv_request_id == kv_request_id and x.round_key == round_seq
         ]
         if not ids:
+            if rnd is not None and st.outbound.get(round_seq) is rnd:
+                self._finalize_outbound(
+                    kv_request_id, round_seq, success=False, send_done=False
+                )
             self._finalize_abort(kv_request_id, round_seq)
             return
 
-        started_at = self._pending_aborts[(kv_request_id, round_seq)]
-        if time.monotonic() - started_at >= _CANCEL_DRAIN_TIMEOUT_S:
-            for tid in ids:
-                self._inflight_pop(tid)
-            self._transport.cancel(ids, mode="immediate")
+        now = time.monotonic()
+        key = (kv_request_id, round_seq)
+        if (
+            now - self._pending_aborts[key] >= _CANCEL_DRAIN_TIMEOUT_S
+            and key not in self._slow_aborts
+        ):
+            self._slow_aborts.add(key)
             logger.warning(
-                "P2PSession %s: cancel drain timed out for kv_request_id=%s,"
-                " force-canceled %d transfers",
+                "P2PSession %s: abort of kv_request_id=%s round=%s still "
+                "waiting after %.0fs for %d transfers to stop; holding their "
+                "pins and the AbortAck until NIXL reports them terminal",
                 self._peer_id,
                 kv_request_id,
+                round_seq,
+                _CANCEL_DRAIN_TIMEOUT_S,
                 len(ids),
             )
-            self._finalize_abort(kv_request_id, round_seq)
-            return
 
+        # All of the aborted round's transfers share one round state.
+        if not self._cancel_due(self._inflight[ids[0]].round, now):
+            return
         still = self._transport.cancel(ids, mode="wait")
         # Tids the transport successfully released are gone from its
         # _inflight; mirror that in session bookkeeping so they don't
@@ -1037,10 +1108,30 @@ class ServerRole:
             if tid not in still_set:
                 self._inflight_pop(tid)
         if not still:
+            if rnd is not None and st.outbound.get(round_seq) is rnd:
+                self._finalize_outbound(
+                    kv_request_id, round_seq, success=False, send_done=False
+                )
             self._finalize_abort(kv_request_id, round_seq)
 
+    @staticmethod
+    def _cancel_due(rnd: _OutboundRequestState, now: float) -> bool:
+        """Rate-limit wait-mode cancels of ``rnd``'s transfers.
+
+        Shared by the store-timeout sweep and the abort drain, so a round
+        gets at most one attempt per ``_CANCEL_RETRY_INTERVAL_S``.
+        """
+        if (
+            rnd.last_cancel_at is not None
+            and now - rnd.last_cancel_at < _CANCEL_RETRY_INTERVAL_S
+        ):
+            return False
+        rnd.last_cancel_at = now
+        return True
+
     def _finalize_abort(self, kv_request_id: str, round_seq: int) -> None:
         self._pending_aborts.pop((kv_request_id, round_seq), None)
+        self._slow_aborts.discard((kv_request_id, round_seq))
         self._send(
             {
                 TYPE_KEY: AbortAckMsg.TYPE,
@@ -1121,7 +1212,8 @@ class ServerRole:
     def _timeout_pending_store_jobs(self) -> list[StoreResult]:
         if not self._store_jobs:
             return []
-        deadline = time.monotonic() - _STORE_TIMEOUT_S
+        now = time.monotonic()
+        deadline = now - _STORE_TIMEOUT_S
         timed_out: list[int] | None = None
         for jid, submitted_at in self._store_jobs.items():
             if submitted_at <= deadline:
@@ -1129,10 +1221,44 @@ class ServerRole:
                     timed_out = []
                 timed_out.append(jid)
         if timed_out is None:
-            return []
+            if not any(xfer.round.cancelled for xfer in self._inflight.values()):
+                return []
+            timed_out = []
         results: list[StoreResult] = []
+        expired = set(timed_out)
+        rounds = {
+            (xfer.kv_request_id, xfer.round_key): xfer.round
+            for xfer in self._inflight.values()
+            if xfer.round.cancelled or expired.intersection(xfer.job_ids)
+        }
+        for (kv_request_id, round_key), rnd in rounds.items():
+            rnd.cancelled = True
+            if not self._cancel_due(rnd, now):
+                continue
+            tids = [tid for tid, xfer in self._inflight.items() if xfer.round is rnd]
+            still = set(self._transport.cancel(tids, mode="wait"))
+            for tid in tids:
+                if tid in still:
+                    continue
+                xfer = self._inflight_pop(tid)
+                assert xfer is not None
+                results.extend(self._settle_xfer_jobs(xfer, success=False))
+            st = self._requests.get(kv_request_id)
+            if (
+                rnd.inflight == 0
+                and st is not None
+                and st.outbound.get(round_key) is rnd
+            ):
+                self._finalize_cancelled(kv_request_id, round_key)
+
+        # Cancellation may still be draining. Keep every affected pin until
+        # the transport confirms that no transfer reads its primary slot.
+        inflight_job_ids = self._read_job_ids()
         for jid in timed_out:
-            del self._store_jobs[jid]
+            if jid in inflight_job_ids:
+                continue
+            if self._store_jobs.pop(jid, None) is None:
+                continue
             results.append(StoreResult(job_id=jid, success=False))
             logger.warning("P2PSession %s: store job %d timed out", self._peer_id, jid)
         # A failed store job's primary slots are unpinned by the engine, so its