From ee6b35e5ed143c28ecfd57b29ffd268c5781ce80 Mon Sep 17 00:00:00 2001
From: Liran Schour <lirans@il.ibm.com>
Date: Wed, 28 Oct 2026 10:44:57 +0200
Subject: [PATCH 1/2] [Bugfix][KV Offload][P2P] Retain unconfirmed source
 WRITEs across session teardown

Part of #63602 (gap 3).

ServerRole.close() cancelled inflight transfers in "immediate" mode and
returned every store job as failed. That released NIXL handles still in
PROC/PEND and let the engine unpin source blocks a WRITE may still be
reading; remove_remote_peer() then freed the descriptors those handles
were prepped against. The abort drain and store timeout already refuse
this; teardown was the one path that didn't.

- ServerRole.close() cancels in "wait" mode. Transfers still running are
  returned as SessionCloseResult.retained with the store jobs they read;
  those jobs stay pinned instead of being failed.
- NixlTransport.remove_remote_peer() retires the peer's still-inflight
  transfers: they leave poll(), so a reconnect under the same peer_id
  never sees them, and the old registration's release is deferred until
  they stop. poll_retired() reports stopped ones (same terminal rule as
  wait-mode cancel, without faulting the reconnected peer).
- The manager holds retained transfers and fails their store jobs once
  the transport reports them stopped, or at _RETAINED_HOLD_CAP_S (120 s),
  after which force_release_retired() releases them with an error log.

Signed-off-by: Liran Schour <lirans@il.ibm.com>
---
 .../tiering/p2p/test_data_transport.py        |  71 +++++++++++
 .../v1/kv_offload/tiering/p2p/test_manager.py | 111 ++++++++++++++++++
 .../kv_offload/tiering/p2p/test_sessions.py   |  37 ++++++
 vllm/v1/kv_offload/tiering/p2p/data/base.py   |  27 ++++-
 vllm/v1/kv_offload/tiering/p2p/data/nixl.py   |  66 ++++++++++-
 vllm/v1/kv_offload/tiering/p2p/manager.py     |  71 +++++++++++
 .../kv_offload/tiering/p2p/session/server.py  |  25 ++--
 .../kv_offload/tiering/p2p/session/session.py |  11 +-
 8 files changed, 407 insertions(+), 12 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 46822f369030..f3d60e8079db 100644
--- a/tests/v1/kv_offload/tiering/p2p/test_data_transport.py
+++ b/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):
diff --git a/tests/v1/kv_offload/tiering/p2p/test_manager.py b/tests/v1/kv_offload/tiering/p2p/test_manager.py
index c47010b70790..4de4178928c4 100644
--- a/tests/v1/kv_offload/tiering/p2p/test_manager.py
+++ b/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,111 @@ 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 == {}
+
+
 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."""
diff --git a/tests/v1/kv_offload/tiering/p2p/test_sessions.py b/tests/v1/kv_offload/tiering/p2p/test_sessions.py
index 7856867d6a96..368039003ce6 100644
--- a/tests/v1/kv_offload/tiering/p2p/test_sessions.py
+++ b/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,
diff --git a/vllm/v1/kv_offload/tiering/p2p/data/base.py b/vllm/v1/kv_offload/tiering/p2p/data/base.py
index 30456ca2e22d..4a8ca7c10101 100644
--- a/vllm/v1/kv_offload/tiering/p2p/data/base.py
+++ b/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,
diff --git a/vllm/v1/kv_offload/tiering/p2p/data/nixl.py b/vllm/v1/kv_offload/tiering/p2p/data/nixl.py
index 382ed7576527..c62dab63eb22 100644
--- a/vllm/v1/kv_offload/tiering/p2p/data/nixl.py
+++ b/vllm/v1/kv_offload/tiering/p2p/data/nixl.py
@@ -105,11 +105,18 @@ class _Inflight(NamedTuple):
     transport is shared across all peer sessions of the engine.
     ``struck`` is set once the transfer has counted against its peer, so a
     state check that keeps raising charges it once, not once per poll.
+    ``retired`` is set once its peer is removed while it still runs; only
+    poll_retired() sees it from then on.
     """
 
     peer_id: str
     handle: object
     struck: bool = False
+    retired: bool = False
+
+
+# A removed peer's (peer_id, dlist, nixl_name, retired transfer_ids).
+_DeferredRemote = tuple[str, object | None, str | None, set[int]]
 
 
 class NixlTransport(DataTransport):
@@ -144,6 +151,9 @@ def __init__(
         self._peer_strikes: dict[str, int] = {}
         # peer_id → fault cause; reported by poll(peer_id) until removal.
         self._peer_faults: dict[str, str] = {}
+        # Remote handles of removed peers that retired transfers still
+        # reference. Released by poll_retired() once those transfers are gone.
+        self._deferred_remote: list[_DeferredRemote] = []
         # Registration is O(num_blocks) NIXL work — add_remote_agent plus a
         # prep_xfer_dlist over every block of the peer's region — so it runs
         # off the scheduler thread, mirroring the NIXL connector's
@@ -263,6 +273,17 @@ def remove_remote_peer(self, peer_id: str) -> None:
             dlist = self._remote_dlists.pop(peer_id, None)
         self._peer_strikes.pop(peer_id, None)
         self._peer_faults.pop(peer_id, None)
+        # Transfers still inflight here may still be writing. Retire them so
+        # a reconnected session under the same peer_id never sees them, and
+        # keep the remote handles they were prepped against until they end.
+        retired: set[int] = set()
+        for tid, entry in self._inflight.items():
+            if entry.peer_id == peer_id and not entry.retired:
+                self._inflight[tid] = entry._replace(retired=True)
+                retired.add(tid)
+        if retired:
+            self._deferred_remote.append((peer_id, dlist, nixl_name, retired))
+            return
         if self._agent is not None:
             self._release_remote_handles(peer_id, dlist, nixl_name)
 
@@ -446,7 +467,7 @@ def poll(self, peer_id: str | None = None) -> PollResult:
         failures: list[tuple[str, str]] | None = None
 
         for transfer_id, entry in self._inflight.items():
-            if peer_id is not None and entry.peer_id != peer_id:
+            if entry.retired or (peer_id is not None and entry.peer_id != peer_id):
                 continue
             try:
                 state = self._agent.check_xfer_state(entry.handle)
@@ -515,6 +536,48 @@ def _poll_result(
             peer_fault=fault,
         )
 
+    def poll_retired(self) -> list[int]:
+        """Release retired transfers that have stopped; return their ids.
+
+        A transfer stopped once NIXL reports it terminal, by the same rule
+        as ``cancel(mode="wait")``, but without charging its peer: the
+        peer is gone, and a reconnected one under the same id is not at
+        fault. A removed peer's remote handles are released once none of
+        its retired transfers remain.
+        """
+        stopped: list[int] = []
+        for tid, entry in self._inflight.items():
+            if not entry.retired:
+                continue
+            try:
+                state = self._agent.check_xfer_state(entry.handle)
+            except Exception as exc:
+                if not _is_terminal_state_error(exc):
+                    continue
+                state = "ERR"
+            if state not in ("PROC", "PEND"):
+                stopped.append(tid)
+        if not stopped:
+            return []
+        self._release_handles([self._inflight.pop(tid).handle for tid in stopped])
+        self._release_deferred_remote()
+        return stopped
+
+    def force_release_retired(self, transfer_ids: Iterable[int]) -> None:
+        """Release retired transfers that never stopped (bounded hold)."""
+        self.cancel(transfer_ids, mode="immediate")
+        self._release_deferred_remote()
+
+    def _release_deferred_remote(self) -> None:
+        pending: list[_DeferredRemote] = []
+        for peer_id, dlist, nixl_name, tids in self._deferred_remote:
+            tids &= self._inflight.keys()
+            if tids:
+                pending.append((peer_id, dlist, nixl_name, tids))
+            elif self._agent is not None:
+                self._release_remote_handles(peer_id, dlist, nixl_name)
+        self._deferred_remote = pending
+
     def cancel(
         self,
         transfer_ids: Iterable[int],
@@ -594,6 +657,7 @@ def close(self) -> None:
             return
         self._release_handles([entry.handle for entry in self._inflight.values()])
         self._inflight.clear()
+        self._release_deferred_remote()
         for peer_id in list(self._remote_dlists):
             self.remove_remote_peer(peer_id)
         if self._local_dlist is not None:
diff --git a/vllm/v1/kv_offload/tiering/p2p/manager.py b/vllm/v1/kv_offload/tiering/p2p/manager.py
index 52aab4bfd266..7c374f834b35 100644
--- a/vllm/v1/kv_offload/tiering/p2p/manager.py
+++ b/vllm/v1/kv_offload/tiering/p2p/manager.py
@@ -71,6 +71,11 @@
 # enough to avoid busy-spinning the scheduler thread.
 _DRAIN_SLEEP_S = 0.001
 
+# Longest a reaped session's still-running transfer keeps its source pins.
+# NIXL may never report a wedged transfer terminal, and a hold without a
+# bound would look like a full cache; past this the handle is released.
+_RETAINED_HOLD_CAP_S = 120.0
+
 # Idle sessions are only reaped when remaining capacity is at most this
 # many slots. Session setup is expensive; below this headroom, keep idle
 # peers so a later burst does not pay handshake again.
@@ -400,6 +405,10 @@ def __init__(
         # at the top of the next ``serve_external_requests`` where the
         # handle is valid.
         self._failed_serve_ctxs: list[ReqContext] = []
+        # transfer_id → (peer_id, store job_ids, retained_at) for transfers
+        # of reaped sessions that may still be writing; see _poll_retained.
+        self._retained: dict[int, tuple[str, frozenset[int], float]] = {}
+        self._retained_force_released = 0
 
     # ------------------------------------------------------------------
     # SecondaryTierManager interface
@@ -852,6 +861,7 @@ def _reap_dead_sessions(self) -> None:
                 self._finished_jobs.append(JobResult(job_id=job_id, success=False))
             for job_id in close_result.failed_stores:
                 self._finished_jobs.append(JobResult(job_id=job_id, success=False))
+            self._retain(pid, close_result.retained)
             # Fail every client-side request (in-flight loads plus unresolved
             # symmetric-P2P probes) toward the dead peer so lookup() returns
             # MISS (local prefill) instead of RETRY forever — even if a fresh
@@ -877,6 +887,59 @@ def _reap_dead_sessions(self) -> None:
                 fault_info,
             )
 
+    def _retain(self, peer_id: str, retained: dict[int, frozenset[int]]) -> None:
+        """Hold a reaped session's still-running transfers and their pins."""
+        if not retained:
+            return
+        now = time.monotonic()
+        for tid, job_ids in retained.items():
+            self._retained[tid] = (peer_id, job_ids, now)
+        logger.warning(
+            "P2P %s: peer %s reaped with %d transfers still running; holding "
+            "their source pins until they stop (at most %.0fs)",
+            self._local_id,
+            peer_id,
+            len(retained),
+            _RETAINED_HOLD_CAP_S,
+        )
+
+    def _poll_retained(self) -> None:
+        """Unpin retained transfers once they stop, or at the hold cap.
+
+        A store job is failed only once no retained transfer reads it.
+        """
+        if not self._retained:
+            return
+        released = list(self._data.poll_retired())
+        deadline = time.monotonic() - _RETAINED_HOLD_CAP_S
+        expired = [
+            tid
+            for tid, (_, _, retained_at) in self._retained.items()
+            if retained_at <= deadline and tid not in released
+        ]
+        if expired:
+            self._data.force_release_retired(expired)
+            self._retained_force_released += len(expired)
+            logger.error(
+                "P2P %s: %d transfers of reaped peers still running after "
+                "%.0fs; releasing their handles and source pins anyway "
+                "(%d so far)",
+                self._local_id,
+                len(expired),
+                _RETAINED_HOLD_CAP_S,
+                self._retained_force_released,
+            )
+        jobs: set[int] = set()
+        for tid in (*released, *expired):
+            entry = self._retained.pop(tid, None)
+            if entry is not None:
+                jobs |= entry[1]
+        for job_id in jobs - self._retained_job_ids():
+            self._finished_jobs.append(JobResult(job_id=job_id, success=False))
+
+    def _retained_job_ids(self) -> set[int]:
+        return {jid for _, job_ids, _ in self._retained.values() for jid in job_ids}
+
     def _schedule_reconnect_backoff(self, peer_id: str) -> float:
         """Record a data-plane fault for *peer_id*; return its reconnect wait.
 
@@ -988,6 +1051,7 @@ def _reap_idle_sessions(self) -> None:
                 self._finished_jobs.append(JobResult(job_id=job_id, success=False))
             for job_id in close_result.failed_stores:
                 self._finished_jobs.append(JobResult(job_id=job_id, success=False))
+            self._retain(pid, close_result.retained)
             self._failed_req_ids.update(close_result.failed_req_ids)
             self._failed_serve_ctxs.extend(close_result.failed_serves)
             self._data.remove_remote_peer(pid)
@@ -1066,6 +1130,7 @@ def _poll_once(self) -> None:
         self._reap_dead_sessions()
         self._reap_unbound_stores()
         self._reap_idle_sessions()
+        self._poll_retained()
 
     # ------------------------------------------------------------------
     # Lifecycle
@@ -1090,6 +1155,11 @@ def shutdown(self) -> None:
                 )
         self._unbound_stores.clear()
         self._reaped_stores.clear()
+        # The drain above already waited on retained transfers too, and
+        # _data.close() releases whatever is left.
+        for job_id in self._retained_job_ids():
+            self._finished_jobs.append(JobResult(job_id=job_id, success=False))
+        self._retained.clear()
         self._control.close()
         self._data.close()
 
@@ -1104,6 +1174,7 @@ def _drain_inflight_for_shutdown(self) -> None:
         wedged peer must not hang shutdown.
         """
         ids = [tid for s in self._sessions.values() for tid in s._server._inflight]
+        ids.extend(self._retained)
         if not ids:
             return
         deadline = time.monotonic() + _SHUTDOWN_DRAIN_TIMEOUT_S
diff --git a/vllm/v1/kv_offload/tiering/p2p/session/server.py b/vllm/v1/kv_offload/tiering/p2p/session/server.py
index eff19222ea29..0f057f080f11 100644
--- a/vllm/v1/kv_offload/tiering/p2p/session/server.py
+++ b/vllm/v1/kv_offload/tiering/p2p/session/server.py
@@ -860,19 +860,28 @@ def drain_pending_aborts(self) -> None:
         for kv_request_id, round_seq in list(self._pending_aborts):
             self._drain_abort(kv_request_id, round_seq)
 
-    def close(self) -> tuple[list[int], list[ReqContext]]:
-        """Tear down. Cancels inflight.
+    def close(
+        self,
+    ) -> tuple[list[int], list[ReqContext], dict[int, frozenset[int]]]:
+        """Tear down. Cancels inflight, retaining what may still write.
 
-        Returns ``(failed_store_job_ids, failed_serves)`` where
+        Returns ``(failed_store_job_ids, failed_serves, retained)``.
         ``failed_serves`` are synthetic lookup ctxs still owing a
         ``parent.on_request_finished``. The session is going away with no
         parent handle in hand, so the manager flushes these in its next
-        ``serve_external_requests``.
+        ``serve_external_requests``. ``retained`` maps each transfer NIXL
+        does not report terminal to the store jobs it still reads: like the
+        abort drain, close never releases such a transfer, so those jobs
+        stay pinned and are left out of ``failed_store_job_ids`` until the
+        transport reports the transfer stopped.
         """
-        failed_stores = list(self._store_jobs.keys())
-        self._store_jobs.clear()
+        retained: dict[int, frozenset[int]] = {}
         if self._inflight:
-            self._transport.cancel(list(self._inflight.keys()))
+            still = self._transport.cancel(list(self._inflight.keys()), mode="wait")
+            retained = {tid: frozenset(self._inflight[tid].job_ids) for tid in still}
+        held = set().union(*retained.values())
+        failed_stores = [jid for jid in self._store_jobs if jid not in held]
+        self._store_jobs.clear()
         self._inflight.clear()
         self._pending_store_results.clear()
         # Surface every synthetic ctx still owing on_request_finished so
@@ -888,7 +897,7 @@ def close(self) -> tuple[list[int], list[ReqContext]]:
         self._pending_aborts.clear()
         self._slow_aborts.clear()
         self._finished_lookup_ctxs.clear()
-        return failed_stores, failed_serves
+        return failed_stores, failed_serves, retained
 
     @property
     def has_inflight_transfers(self) -> bool:
diff --git a/vllm/v1/kv_offload/tiering/p2p/session/session.py b/vllm/v1/kv_offload/tiering/p2p/session/session.py
index 3c435bb7d334..578bd6a96111 100644
--- a/vllm/v1/kv_offload/tiering/p2p/session/session.py
+++ b/vllm/v1/kv_offload/tiering/p2p/session/session.py
@@ -91,13 +91,17 @@ class SessionCloseResult(NamedTuple):
     server stores) the manager must mark failed. `failed_req_ids` is every
     client-side kv_request_id whose lookup() must fail (in-flight loads plus
     unresolved probes); `failed_serves` is the server-side lookup state the
-    dead peer can no longer resolve.
+    dead peer can no longer resolve. `retained` maps each server transfer
+    that may still be writing to the store jobs it reads: they are not in
+    `failed_stores`, and stay pinned until the transport reports it stopped.
     """
 
     failed_jobs: list[int]  # client load job_ids
     failed_req_ids: list[str]  # client kv_request_ids (loads + probes)
     failed_stores: list[int]  # server store job_ids
     failed_serves: list[ReqContext]  # server-side lookup ctxs needing release
+    # transfer_id → store job_ids it still reads
+    retained: dict[int, frozenset[int]] = {}
 
 
 class P2PSession:
@@ -384,6 +388,8 @@ def close(self) -> SessionCloseResult:
         failed_serves: synthetic lookup ctxs still owing
             ``parent.on_request_finished`` (the manager flushes these on
             its next ``serve_external_requests``).
+        retained: server transfers that may still be writing, with the
+            store jobs they read (held until they stop).
         """
         # Abandon any in-flight registration: nothing will ack this peer
         # now. The manager's remove_remote_peer bumps the transport's
@@ -394,7 +400,7 @@ def close(self) -> SessionCloseResult:
             self._pending_registration = None
 
         client_result = self._client.close()
-        failed_stores, failed_serves = self._server.close()
+        failed_stores, failed_serves, retained = self._server.close()
 
         if self._conn is not None:
             with contextlib.suppress(Exception):
@@ -407,6 +413,7 @@ def close(self) -> SessionCloseResult:
             failed_req_ids=client_result.failed_req_ids,
             failed_stores=failed_stores,
             failed_serves=failed_serves,
+            retained=retained,
         )
 
     # ------------------------------------------------------------------

From 2c516742aa0281bae7acc3708c3af6369b211cdd Mon Sep 17 00:00:00 2001
From: Liran Schour <lirans@il.ibm.com>
Date: Wed, 28 Oct 2026 14:40:00 +0200
Subject: [PATCH 2/2] [KV Offload][P2P] Hold a peer out while its retained
 transfers run

A reaped peer's registration stays until its retired transfers stop. A
reconnect meanwhile loads the peer's NIXL agent again under the same
name: NIXL keeps the old UCX endpoints, which may have failed, and the
deferred release's remove_remote_agent then tears down the new
session's agent. Reject outbound and inbound sessions for such a peer
until its retained transfers stop, and lower the hold cap to 30s since
it now also bounds how long the peer is unreachable.

drain_jobs() also waits for retained transfers: reset_cache() asserts
every job finished once it returns, and their store jobs are reported
only when they stop.

Signed-off-by: Liran Schour <lirans@il.ibm.com>
---
 .../v1/kv_offload/tiering/p2p/test_manager.py | 60 +++++++++++++++++++
 vllm/v1/kv_offload/tiering/p2p/manager.py     | 40 +++++++++++--
 .../kv_offload/tiering/p2p/session/session.py |  2 +-
 3 files changed, 95 insertions(+), 7 deletions(-)

diff --git a/tests/v1/kv_offload/tiering/p2p/test_manager.py b/tests/v1/kv_offload/tiering/p2p/test_manager.py
index 4de4178928c4..56611497bd4b 100644
--- a/tests/v1/kv_offload/tiering/p2p/test_manager.py
+++ b/tests/v1/kv_offload/tiering/p2p/test_manager.py
@@ -2435,6 +2435,66 @@ def test_shutdown_fails_retained_jobs(self):
         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
diff --git a/vllm/v1/kv_offload/tiering/p2p/manager.py b/vllm/v1/kv_offload/tiering/p2p/manager.py
index 7c374f834b35..696dde3bf2ee 100644
--- a/vllm/v1/kv_offload/tiering/p2p/manager.py
+++ b/vllm/v1/kv_offload/tiering/p2p/manager.py
@@ -74,7 +74,9 @@
 # Longest a reaped session's still-running transfer keeps its source pins.
 # NIXL may never report a wedged transfer terminal, and a hold without a
 # bound would look like a full cache; past this the handle is released.
-_RETAINED_HOLD_CAP_S = 120.0
+# The peer cannot reconnect while it holds retained transfers (see
+# _retained_peer), so this also bounds how long it stays unreachable.
+_RETAINED_HOLD_CAP_S = 30.0
 
 # Idle sessions are only reaped when remaining capacity is at most this
 # many slots. Session setup is expensive; below this headroom, keep idle
@@ -666,9 +668,11 @@ def drain_jobs(self) -> None:
         """Block until every submitted load/store job has completed or failed.
 
         Loops calling ``_poll_once()`` until no session has outstanding
-        inbound loads or in-flight outbound stores. Mid-flight transfers
-        are NOT cancelled — the caller (``TieringOffloadingManager.reset_cache``)
-        needs the primary memoryview to be quiescent, not aborted. Results
+        inbound loads or in-flight outbound stores, and no reaped session's
+        transfer is still retained (bounded by its hold cap). Mid-flight
+        transfers are NOT cancelled — the caller
+        (``TieringOffloadingManager.reset_cache``) needs the primary
+        memoryview to be quiescent, not aborted. Results
         accumulate in ``_finished_jobs`` and are surfaced by the next
         ``get_finished_jobs()`` call.
         """
@@ -676,7 +680,9 @@ def drain_jobs(self) -> None:
         warned = False
         while True:
             self._poll_once()
-            pending = any(s.has_pending_work for s in self._sessions.values())
+            pending = bool(self._retained) or any(
+                s.has_pending_work for s in self._sessions.values()
+            )
             if not pending:
                 return
             if not warned and time.monotonic() - start > 5.0:
@@ -739,7 +745,7 @@ def _get_or_create_session(self, peer_id: str) -> P2PSession | None:
         session = self._sessions.get(peer_id)
         if session is not None:
             return session
-        if self._reconnect_wait(peer_id) > 0:
+        if self._reconnect_wait(peer_id) > 0 or self._retained_peer(peer_id):
             return None
 
         max_peers = envs.VLLM_P2P_MAX_PEERS
@@ -812,6 +818,17 @@ def _accept_new_peers(self, new_connections: Sequence[ControlConnection]) -> Non
                     conn.send({TYPE_KEY: DisconnectMsg.TYPE})
                 conn.close()
                 continue
+            if self._retained_peer(conn.peer_id):
+                logger.info(
+                    "P2P %s: rejecting inbound peer %s — its reaped "
+                    "session's transfers are still running",
+                    self._local_id,
+                    conn.peer_id,
+                )
+                with contextlib.suppress(Exception):
+                    conn.send({TYPE_KEY: DisconnectMsg.TYPE})
+                conn.close()
+                continue
             try:
                 existing = self._sessions.get(conn.peer_id)
                 if existing is not None:
@@ -940,6 +957,17 @@ def _poll_retained(self) -> None:
     def _retained_job_ids(self) -> set[int]:
         return {jid for _, job_ids, _ in self._retained.values() for jid in job_ids}
 
+    def _retained_peer(self, peer_id: str) -> bool:
+        """Whether a reaped session of *peer_id* still has running transfers.
+
+        Its registration stays until they stop. A new session would load
+        the peer's NIXL agent again under the same name: NIXL keeps the old
+        UCX endpoints, possibly failed, and the old registration's release
+        removes the agent under the new one. So the peer stays out until
+        the transfers stop, at most _RETAINED_HOLD_CAP_S.
+        """
+        return any(pid == peer_id for pid, _, _ in self._retained.values())
+
     def _schedule_reconnect_backoff(self, peer_id: str) -> float:
         """Record a data-plane fault for *peer_id*; return its reconnect wait.
 
diff --git a/vllm/v1/kv_offload/tiering/p2p/session/session.py b/vllm/v1/kv_offload/tiering/p2p/session/session.py
index 578bd6a96111..d23c306d3b26 100644
--- a/vllm/v1/kv_offload/tiering/p2p/session/session.py
+++ b/vllm/v1/kv_offload/tiering/p2p/session/session.py
@@ -101,7 +101,7 @@ class SessionCloseResult(NamedTuple):
     failed_stores: list[int]  # server store job_ids
     failed_serves: list[ReqContext]  # server-side lookup ctxs needing release
     # transfer_id → store job_ids it still reads
-    retained: dict[int, frozenset[int]] = {}
+    retained: dict[int, frozenset[int]]
 
 
 class P2PSession: