[Bugfix][KV Offload][P2P] Retain unconfirmed source WRITEs across session teardown

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

@@ -909,6 +909,77 @@ def test_raising_state_check_warns_once(self, monkeypatch):
             transport.poll(peer_id="peer:1")
         assert sum("check_xfer_state failed" in w for w in warnings) == 1
 
+    def test_remove_remote_peer_retires_running_transfer(self):
+        """A WRITE still running at peer removal keeps its handle and the
+        peer's registration until it stops, and leaves every poll() so a
+        reconnected peer under the same id never sees it."""
+        transport = self._make_transport()
+        agent = transport._agent
+        transport.add_remote_peer("peer:1", b"meta", 0x1000, 8, 1024)
+        tid = transport.write_blocks("peer:1", [0], [1])
+
+        transport.remove_remote_peer("peer:1")
+        agent.release_dlist_handle.assert_not_called()
+        agent.remove_remote_agent.assert_not_called()
+        transport.add_remote_peer("peer:1", b"meta", 0x1000, 8, 1024)
+        agent.check_xfer_state.return_value = "DONE"
+        assert not transport.poll(peer_id="peer:1").done
+        assert not transport.poll().done
+        assert tid in transport._inflight
+
+        assert transport.poll_retired() == [tid]
+        assert tid not in transport._inflight
+        agent.release_xfer_handle.assert_called_once()
+        # The old registration only; the new one stays registered.
+        agent.release_dlist_handle.assert_called_once()
+        agent.remove_remote_agent.assert_called_once()
+        assert "peer:1" in transport._remote_dlists
+
+    @pytest.mark.parametrize("state", ["PROC", "PEND"])
+    def test_poll_retired_keeps_running_transfer(self, state):
+        transport = self._make_transport()
+        transport.add_remote_peer("peer:1", b"meta", 0x1000, 8, 1024)
+        tid = transport.write_blocks("peer:1", [0], [1])
+        transport.remove_remote_peer("peer:1")
+        transport._agent.check_xfer_state.return_value = state
+
+        assert transport.poll_retired() == []
+        assert tid in transport._inflight
+        transport._agent.release_xfer_handle.assert_not_called()
+
+    def test_poll_retired_state_check_raise_is_not_terminal(self):
+        """Only a remote disconnect proves a retired WRITE 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.remove_remote_peer("peer:1")
+        transport._agent.check_xfer_state.side_effect = nixlNotFoundError("gone")
+        assert transport.poll_retired() == []
+        transport._agent.check_xfer_state.side_effect = _remote_disconnect()
+        assert transport.poll_retired() == [tid]
+        assert "peer:1" not in transport._peer_faults
+
+    def test_force_release_retired_releases_registration(self):
+        transport = self._make_transport()
+        agent = transport._agent
+        transport.add_remote_peer("peer:1", b"meta", 0x1000, 8, 1024)
+        tid = transport.write_blocks("peer:1", [0], [1])
+        assert tid is not None
+        transport.remove_remote_peer("peer:1")
+
+        transport.force_release_retired([tid])
+
+        assert tid not in transport._inflight
+        agent.release_xfer_handle.assert_called_once()
+        agent.remove_remote_agent.assert_called_once()
+
+    def test_remove_remote_peer_without_inflight_releases_at_once(self):
+        transport = self._make_transport()
+        transport.add_remote_peer("peer:1", b"meta", 0x1000, 8, 1024)
+        transport.remove_remote_peer("peer:1")
+        transport._agent.release_dlist_handle.assert_called_once()
+        transport._agent.remove_remote_agent.assert_called_once()
+
     def test_remove_remote_peer_clears_the_fault(self):
         transport = self._make_transport()
         for _ in range(3):

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

@@ -39,6 +39,7 @@
 from vllm.v1.kv_offload.tiering.p2p import manager as manager_module
 from vllm.v1.kv_offload.tiering.p2p.manager import (
     _RECONNECT_MAX_BACKOFF_S,
+    _RETAINED_HOLD_CAP_S,
     _UNBOUND_STORE_TIMEOUT_S,
     P2PSecondaryTierManager,
     _annotate_req_context,
@@ -139,6 +140,8 @@ def _make_manager() -> P2PSecondaryTierManager:
     mgr._reaped_stores = {}
     mgr._unbound_store_timeout_s = _UNBOUND_STORE_TIMEOUT_S
     mgr._failed_serve_ctxs = []
+    mgr._retained = {}
+    mgr._retained_force_released = 0
     return mgr
 
 
@@ -727,6 +730,7 @@ def __init__(
         close_stores: list[int] | None = None,
         close_failed_serves: list[ReqContext] | None = None,
         last_activity_at: float | None = None,
+        close_retained: dict[int, frozenset[int]] | None = None,
     ) -> None:
         self.peer_id = peer_id
         self.alive = alive
@@ -741,6 +745,7 @@ def __init__(
         self._close_req_ids = close_req_ids or []
         self._close_stores = close_stores or []
         self._close_failed_serves = close_failed_serves or []
+        self._close_retained = close_retained or {}
         self.requests: list[tuple[int, str]] = []
         self.stores_added: list[tuple[str, list, object, int]] = []
         self.attached: list[object] = []
@@ -788,6 +793,7 @@ def close(self):
             failed_req_ids=self._close_req_ids,
             failed_stores=self._close_stores,
             failed_serves=self._close_failed_serves,
+            retained=self._close_retained,
         )
 
 
@@ -2325,6 +2331,171 @@ def test_disabled_when_zero(self, monkeypatch):
 # ---------------------------------------------------------------------------
 
 
+class _RetiringFakeData:
+    """Fake DataTransport whose retired transfers stop on demand."""
+
+    def __init__(self) -> None:
+        self.stopped: list[int] = []
+        self.force_released: list[int] = []
+        self.removed: list[str] = []
+
+    def remove_remote_peer(self, pid) -> None:
+        self.removed.append(pid)
+
+    def poll_retired(self) -> list[int]:
+        stopped, self.stopped = self.stopped, []
+        return stopped
+
+    def force_release_retired(self, transfer_ids) -> None:
+        self.force_released.extend(transfer_ids)
+
+
+class TestRetainedTransfers:
+    """A reaped session's still-running transfers keep their source pins
+    until the transport reports them stopped, or the hold cap expires."""
+
+    PEER = "10.0.0.1:8000"
+
+    def _reap(self, retained, close_stores=None):
+        mgr = _make_manager()
+        data = _RetiringFakeData()
+        mgr._data = data  # type: ignore[assignment]
+        mgr._sessions[self.PEER] = _FakeSession(  # type: ignore[assignment]
+            peer_id=self.PEER,
+            alive=False,
+            connected=True,
+            close_stores=close_stores,
+            close_retained=retained,
+        )
+        mgr._reap_dead_sessions()
+        return mgr, data
+
+    def test_reap_holds_pins_until_transfer_stops(self):
+        mgr, data = self._reap({7: frozenset({1})}, close_stores=[2])
+        assert mgr._finished_jobs == [JobResult(job_id=2, success=False)]
+        assert data.removed == [self.PEER]
+
+        mgr._poll_retained()
+        assert mgr._finished_jobs == [JobResult(job_id=2, success=False)]
+
+        data.stopped = [7]
+        mgr._poll_retained()
+        assert mgr._finished_jobs[-1] == JobResult(job_id=1, success=False)
+        assert mgr._retained == {}
+
+    def test_shared_job_held_until_last_transfer_stops(self):
+        mgr, data = self._reap({7: frozenset({1}), 8: frozenset({1})})
+        data.stopped = [7]
+        mgr._poll_retained()
+        assert mgr._finished_jobs == []
+        data.stopped = [8]
+        mgr._poll_retained()
+        assert mgr._finished_jobs == [JobResult(job_id=1, success=False)]
+
+    def test_hold_cap_force_releases(self, monkeypatch):
+        now = [1000.0]
+        monkeypatch.setattr(manager_module.time, "monotonic", lambda: now[0])
+        mgr, data = self._reap({7: frozenset({1})})
+
+        now[0] += _RETAINED_HOLD_CAP_S - 1
+        mgr._poll_retained()
+        assert data.force_released == []
+
+        now[0] += 1
+        mgr._poll_retained()
+        assert data.force_released == [7]
+        assert mgr._finished_jobs == [JobResult(job_id=1, success=False)]
+        assert mgr._retained_force_released == 1
+        assert mgr._retained == {}
+
+    def test_idle_eviction_retains_too(self, monkeypatch):
+        monkeypatch.setenv("VLLM_P2P_IDLE_TIMEOUT_S", "10")
+        monkeypatch.setenv("VLLM_P2P_MAX_PEERS", "1")
+        import vllm.envs
+
+        vllm.envs.__dict__.pop("VLLM_P2P_IDLE_TIMEOUT_S", None)
+        vllm.envs.__dict__.pop("VLLM_P2P_MAX_PEERS", None)
+        mgr = _make_manager()
+        mgr._data = _RetiringFakeData()  # type: ignore[assignment]
+        mgr._sessions[self.PEER] = _FakeSession(  # type: ignore[assignment]
+            peer_id=self.PEER,
+            last_activity_at=time.monotonic() - 20,
+            close_retained={7: frozenset({1})},
+        )
+        mgr._reap_idle_sessions()
+        assert self.PEER not in mgr._sessions
+        assert mgr._retained[7][:2] == (self.PEER, frozenset({1}))
+        assert mgr._finished_jobs == []
+
+    def test_shutdown_fails_retained_jobs(self):
+        mgr, _ = self._reap({7: frozenset({1})})
+        mgr._data = _ShutdownFakeData()  # type: ignore[assignment]
+        mgr._control = _ShutdownFakeControl()  # type: ignore[assignment]
+        mgr.shutdown()
+        assert JobResult(job_id=1, success=False) in mgr._finished_jobs
+        assert mgr._retained == {}
+
+    def test_outbound_waits_for_retained_transfers(self):
+        """A new session would load the peer's NIXL agent under the same
+        name while the reaped one's registration still holds it."""
+        mgr, data = self._reap({7: frozenset({1})})
+
+        class NoConnect:
+            def connect(self, peer_id):
+                raise AssertionError("connect() while transfers are retained")
+
+        mgr._control = NoConnect()  # type: ignore[assignment]
+        assert mgr._get_or_create_session(self.PEER) is None
+
+        data.stopped = [7]
+        mgr._poll_retained()
+        assert not mgr._retained_peer(self.PEER)
+
+    def test_inbound_is_told_to_disconnect_while_retained(self):
+        mgr, data = self._reap({7: frozenset({1})})
+
+        class Conn(_RecordingConn):
+            def __init__(self, peer_id: str) -> None:
+                super().__init__(peer_id)
+                self.sent: list[dict] = []
+
+            def send(self, msg: dict) -> None:
+                self.sent.append(msg)
+
+        conn = Conn(self.PEER)
+        mgr._accept_new_peers([conn])  # type: ignore[list-item]
+        assert conn.sent == [{TYPE_KEY: DisconnectMsg.TYPE}]
+        assert conn.close_calls == 1
+        assert self.PEER not in mgr._sessions
+
+        data.stopped = [7]
+        mgr._poll_retained()
+        assert not mgr._retained_peer(self.PEER)
+
+    def test_other_peers_are_not_held(self):
+        mgr, _ = self._reap({7: frozenset({1})})
+        assert not mgr._retained_peer("10.0.0.2:8000")
+
+    def test_drain_jobs_waits_for_retained_transfers(self, monkeypatch):
+        """reset_cache() asserts every job finished once drain_jobs()
+        returns, so a retained transfer's store jobs must be settled."""
+        mgr, data = self._reap({7: frozenset({1})})
+        mgr._control = _DrainCtrl()  # type: ignore[assignment]
+        polls = [0]
+
+        def poll_retired():
+            polls[0] += 1
+            return [7] if polls[0] == 3 else []
+
+        data.poll_retired = poll_retired  # type: ignore[method-assign]
+        monkeypatch.setattr(manager_module.time, "sleep", lambda s: None)
+        mgr.drain_jobs()
+
+        assert polls[0] == 3
+        assert mgr._retained == {}
+        assert JobResult(job_id=1, success=False) in mgr._finished_jobs
+
+
 class TestReconnectBackoff:
     """A peer whose session ended on a data-plane fault reconnects at once the
     first time; repeats back off from 1 s up to _RECONNECT_MAX_BACKOFF_S."""

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

@@ -2984,6 +2984,43 @@ def test_close_returns_pending_loads_and_stores(self):
         assert set(result.failed_stores) == {10}
         assert result.failed_serves == []
 
+    def test_close_retains_running_transfer(self):
+        """close() must not release a WRITE that may still be landing.
+
+        It cancels in wait mode; a transfer still running is reported as
+        retained with its store job, which stays pinned, and only the jobs
+        no transfer reads are returned as failed.
+        """
+        session, conn, data = _make_session()
+        _activate(session, conn)
+        session.add_stored_blocks("req", [OffloadKey(b"key")], [0], job_id=1)
+        session.add_stored_blocks("idle", [OffloadKey(b"other")], [1], job_id=2)
+        _fetch(conn, "req", [b"key"], [5])
+        session.poll()
+        (tid,) = data._transfers
+        data._cancel_still_inflight.add(tid)
+
+        result = session.close()
+
+        assert data._cancel_calls == [([tid], "wait")]
+        assert tid in data._transfers
+        assert result.retained == {tid: frozenset({1})}
+        assert result.failed_stores == [2]
+
+    def test_close_releases_stopped_transfer(self):
+        """A transfer that has already stopped is released, not retained."""
+        session, conn, data = _make_session()
+        _activate(session, conn)
+        session.add_stored_blocks("req", [OffloadKey(b"key")], [0], job_id=1)
+        _fetch(conn, "req", [b"key"], [5])
+        session.poll()
+
+        result = session.close()
+
+        assert data._transfers == {}
+        assert result.retained == {}
+        assert result.failed_stores == [1]
+
     def test_send_failure_marks_connection_dead(self):
         """A raising send must mark the connection dead, not silently drop
         the message — otherwise the session lingers alive, is never reaped,

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

@@ -65,6 +65,14 @@
   state check failed stays inflight unless the failure proves it can no
   longer write, so the peer's teardown is what bounds it.
 
+remove_remote_peer() does not stop a transfer that is still running: the
+transport retires it instead. A retired transfer leaves every poll(), so
+a reconnected peer under the same id never sees it; poll_retired()
+reports it once it has stopped, and only then are its handle and the
+removed peer's registration released. The caller keeps the transfer's
+source pins until then, or until it gives up and calls
+force_release_retired().
+
 The transport never recovers a peer itself. The session turns a peer
 fault into the same connection death as a missed heartbeat, so the
 manager's reap is the only teardown: DisconnectMsg makes the peer tear
@@ -261,10 +269,27 @@ def add_remote_peer_async(
     def remove_remote_peer(self, peer_id: str) -> None:
         """Unregister a remote peer and release associated resources.
 
-        Inflight transfers to this peer should be cancelled first.
+        Transfers to this peer still inflight are retired, not released
+        (see "Failure model"); the peer's registration is released once
+        they have stopped.
         """
         ...
 
+    def poll_retired(self) -> list[int]:
+        """Release retired transfers that have stopped; return their ids.
+
+        Must be called periodically while any transfer is retired. The
+        default transport retires nothing.
+        """
+        return []
+
+    def force_release_retired(self, transfer_ids: Iterable[int]) -> None:
+        """Release retired transfers that have not stopped.
+
+        Unsafe while one may still be writing; for a bounded hold only.
+        """
+        self.cancel(transfer_ids, mode="immediate")
+
     @abstractmethod
     def write_blocks(
         self,