[KVConnector][P2P] Configurable unbound-store timeout and one-RTT rejection of a late fetch

docs/features/kv_offloading_usage.md

@@ -203,9 +203,31 @@ Block content hashes must match across instances for peers to exchange blocks (s
 | `port` | no | `$VLLM_P2P_SIDE_CHANNEL_PORT` (`5710`) | Base port for the control socket. Must be reachable from peers. The bound port is `base + data_parallel_index` (one socket per DP replica). When omitted, the base resolves from the env var below. |
 | `backends` | no | `["UCX"]` | NIXL transport backends. See [NixlConnector Usage Guide](nixl_connector_usage.md#selecting-a-nixl-transport-backend-plugin) for available backends and selection guidance. |
 | `num_threads` | no | `4` | NIXL agent worker threads. Only used when `backends` is UCX-only; ignored when any non-UCX backend is requested. |
+| `unbound_store_timeout_s` | no | `60` | Seconds a producer holds stored blocks for a consumer that has not fetched them yet. Raise it for deployments whose prefills outlast the default; the blocks keep primary-tier CPU slots pinned for the whole window. Once it expires a late fetch is rejected in a single round trip, so the consumer falls back to local prefill immediately. |
 
 The `backends` and `num_threads` options mirror the conditional logic used by [`NixlConnector`](nixl_connector_usage.md#selecting-a-nixl-transport-backend-plugin): when any non-UCX backend is configured, NIXL is initialised with `backends=...`; otherwise it falls back to a UCX-only agent with the configured `num_threads`. This lets the P2P tier use a different transport (e.g. `MOONCAKE`, `GDS_MT`, `LIBFABRIC`) than the main `NixlConnector` running in the same process.
 
+A producer parks a request's blocks until the consumer's `FetchMsg` arrives; if none arrives
+within `unbound_store_timeout_s` the blocks are released so they stop pinning primary-tier
+slots. Raise it when prefills legitimately take longer than the default:
+
+```bash
+vllm serve <model> \
+  --kv-transfer-config '{
+    "kv_connector": "OffloadingConnector",
+    "kv_role": "kv_both",
+    "kv_connector_extra_config": {
+      "spec_name": "TieringOffloadingSpec",
+      "secondary_tiers": [
+        {
+          "type": "p2p",
+          "unbound_store_timeout_s": 180
+        }
+      ]
+    }
+  }'
+```
+
 #### Environment Variables
 
 Rather than embedding `host`/`port` in each `secondary_tiers` entry, set them once at deploy time via environment variables (mirroring `VLLM_NIXL_SIDE_CHANNEL_HOST`/`VLLM_NIXL_SIDE_CHANNEL_PORT`). Explicit `host`/`port` config keys, when present, take precedence.

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

@@ -117,6 +117,8 @@ def _make_manager() -> P2PSecondaryTierManager:
     mgr._sessions = {}
     mgr._kv_to_session = {}
     mgr._unbound_stores = {}
+    mgr._reaped_stores = {}
+    mgr._unbound_store_timeout_s = _UNBOUND_STORE_TIMEOUT_S
     mgr._failed_serve_ctxs = []
     return mgr
 
@@ -759,9 +761,9 @@ def test_unbound_store_kept_within_timeout(self):
         assert "req-fresh" in mgr._unbound_stores
 
     def test_unbound_store_reaped_after_timeout(self):
-        """Unbound stores past _UNBOUND_STORE_TIMEOUT_S surface as failed
-        and their kv_request_id lands in _failed_req_ids so a late
-        FetchMsg/lookup doesn't try to satisfy them."""
+        """Unbound stores past the reap deadline surface as failed and their
+        kv_request_id lands in _reaped_stores so a late FetchMsg is rejected
+        instead of parking demand nothing can satisfy."""
         from vllm.v1.kv_offload.tiering.p2p.manager import _UnboundStoreBatch
 
         mgr = self._make()
@@ -779,7 +781,7 @@ def test_unbound_store_reaped_after_timeout(self):
         # 2 baseline + 2 buffered stores
         assert JobResult(job_id=10, success=False) in results
         assert JobResult(job_id=11, success=False) in results
-        assert "req-stale" in mgr._failed_req_ids
+        assert "req-stale" in mgr._reaped_stores
 
     def test_submit_store_parks_unbound_batch(self):
         """submit_store on an unbound id appends a batch with a fresh
@@ -1204,6 +1206,75 @@ def test_both_loads_succeed(self):
         assert 201 in b_ok, f"B loads succeeded: {b_ok}"
         assert 200 in b_ok, f"B stores succeeded: {b_ok}"
 
+    def test_late_fetch_after_reap_fails_immediately(self):
+        """A fetch for a reaped kv_request_id fails in one round trip.
+
+        The producer drops parked blocks once unbound_store_timeout_s
+        expires. A consumer whose fetch arrives after that must learn on the
+        next tick via TransferDoneMsg(success=False) — not stall for
+        _LOAD_TIMEOUT_S and then take the abort path.
+        """
+        from vllm.v1.kv_offload.tiering.p2p.session.client import _LOAD_TIMEOUT_S
+
+        mgr_a, mgr_b = _build_paired_managers()
+        kv_id = "req-late"
+
+        # Record every message type leaving the consumer, so the assertion
+        # below can prove no abort was ever needed.
+        sent_from_a: list[str] = []
+        drain = mgr_a._control._drain_outbound_to
+
+        def recording_drain(peer_local_id: str):
+            out = drain(peer_local_id)
+            sent_from_a.extend(msg.get("type") for _, msg in out)
+            return out
+
+        mgr_a._control._drain_outbound_to = recording_drain  # type: ignore[method-assign]
+
+        # Producer parks the blocks, then reaps them before any fetch lands.
+        mgr_b.submit_store(
+            _job_metadata(
+                job_id=200,
+                keys=[b"b-block"],
+                chunk_ids=[0],
+                kv_params={"remote_decoder": {"kv_request_id": kv_id}},
+            )
+        )
+        mgr_b._unbound_store_timeout_s = 0.0
+        reap_results = list(mgr_b.get_finished_jobs())
+        assert JobResult(job_id=200, success=False) in reap_results
+        assert kv_id in mgr_b._reaped_stores
+
+        # Consumer now asks for the blocks that no longer exist.
+        consumer_params = {
+            "remote_prefiller": {
+                "kv_request_id": kv_id,
+                "remote_host": "B",
+                "remote_port": 2,
+            },
+        }
+        mgr_a.on_new_request(_req_context(consumer_params))
+        mgr_a.submit_load(
+            _job_metadata(
+                job_id=101,
+                keys=[b"b-block"],
+                chunk_ids=[0],
+                kv_params=consumer_params,
+            )
+        )
+
+        started = time.monotonic()
+        all_a: list[JobResult] = []
+        for _ in range(8):
+            all_a.extend(list(mgr_a.get_finished_jobs()))
+            list(mgr_b.get_finished_jobs())
+        elapsed = time.monotonic() - started
+
+        assert JobResult(job_id=101, success=False) in all_a, all_a
+        assert kv_id in mgr_a._failed_req_ids
+        assert "abort_fetch" not in sent_from_a, sent_from_a
+        assert elapsed < _LOAD_TIMEOUT_S
+
 
 # ---------------------------------------------------------------------------
 # _accept_new_peers — duplicate connection rejection

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

@@ -47,10 +47,18 @@
 # Reap unbound store batches that have been parked without a FetchMsg
 # binding them to a session for longer than this. Protects against the
 # prefiller buffering blocks for a decoder that never asks (decoder died,
-# network partition, lost kv_request_id). Must be longer than the per-store
-# deadline so the store-timeout path fires first for individual jobs.
+# network partition, lost kv_request_id). Applies only while a batch is
+# unbound: binding hands it to the session, which times the job out under
+# _STORE_TIMEOUT_S instead.
+# Overridable per tier via the ``unbound_store_timeout_s`` config key.
 _UNBOUND_STORE_TIMEOUT_S = 60.0
 
+# How long a reaped kv_request_id is remembered so a fetch that arrives
+# after the reap can be rejected immediately. Must exceed the consumer's
+# load timeout (session.client._LOAD_TIMEOUT_S): past that the consumer has
+# already given up on its own, so there is nobody left to reject.
+_REAPED_ID_RETENTION_S = 60.0
+
 # Time we wait during shutdown for inflight transfers to drain via
 # cancel(mode="wait") before falling back to mode="immediate". Bounded
 # so a wedged peer can't hang shutdown.
@@ -208,6 +216,7 @@ def __init__(
         port: int | None = None,
         backends: list[str] | None = None,
         num_threads: int = 4,
+        unbound_store_timeout_s: float = _UNBOUND_STORE_TIMEOUT_S,
         **kwargs: Any,
     ) -> None:
         """Initialize the P2P secondary tier manager.
@@ -245,9 +254,29 @@ def __init__(
             num_threads: NIXL agent worker threads for the UCX-only
                 branch. Ignored when ``backends`` contains a non-UCX
                 entry.
+            unbound_store_timeout_s: Seconds a producer holds stored blocks
+                for a consumer that has not fetched them yet, before
+                ``_reap_unbound_stores`` drops them. Raise it for
+                deployments whose prefills outlast the default; the blocks
+                keep primary-tier slots pinned for the whole window.
             **kwargs: Reserved for future tier-specific options.
+
+        Raises:
+            ValueError: If ``unbound_store_timeout_s`` is not a positive
+                number, or anything convertible to one.
         """
         super().__init__(offloading_spec, primary_kv_view, tier_type)
+        try:
+            timeout_s = float(unbound_store_timeout_s)
+        except (TypeError, ValueError):
+            timeout_s = float("nan")
+        if not timeout_s > 0:
+            raise ValueError(
+                f"unbound_store_timeout_s must be a positive number, got "
+                f"{unbound_store_timeout_s!r}"
+            )
+        self._unbound_store_timeout_s = timeout_s
+
         # Block hashes chain from NONE_HASH (see v1/core/kv_cache_utils.py).
         # Peers whose seeds differ compute different hashes for identical
         # content, so lookups silently miss and no KV crosses the wire. The
@@ -303,8 +332,16 @@ def __init__(
         # kv_request_id → list of batches submit_store'd before any peer
         # asked for that id. Drained into a session by _on_session_fetch
         # when the corresponding FetchMsg arrives, or surfaced as failures
-        # by _reap_unbound_stores after _UNBOUND_STORE_TIMEOUT_S.
+        # by _reap_unbound_stores after _unbound_store_timeout_s.
         self._unbound_stores: dict[str, list[_UnboundStoreBatch]] = {}
+        # kv_request_id → time its parked batches were reaped. Nothing under
+        # one of these ids can be served any more, so a FetchMsg for it is
+        # rejected on the spot by _poll_once instead of parking demand until
+        # the consumer's own load timeout, and submit_store fails rather than
+        # re-parking under it — a prefill that outran the timeout would
+        # otherwise keep pinning fresh primary-tier slots, one timeout at a
+        # time. Entries are pruned after _REAPED_ID_RETENTION_S.
+        self._reaped_stores: dict[str, float] = {}
 
         self._finished_jobs: list[JobResult] = []
         # kv_request_ids that hit a transport/session failure; On load lookup()
@@ -390,7 +427,7 @@ def on_request_finished(self, req_context: ReqContext) -> None:
         kv_transfer_params; if a session has bound the id, finish it. If
         no session has bound the id yet, this is a no-op: parked batches
         in `_unbound_stores` are left in place and cleaned up only by
-        `_reap_unbound_stores` after `_UNBOUND_STORE_TIMEOUT_S`.
+        `_reap_unbound_stores` after `_unbound_store_timeout_s`.
         """
         source = req_context.get_state(P2PSourceInfo)
         dest = req_context.get_state(P2PDestInfo)
@@ -447,6 +484,22 @@ def submit_store(self, job_metadata: TransferJob) -> None:
             self._finished_jobs.append(JobResult(job_id=job_id, success=False))
             return
 
+        # This id's earlier batches were already reaped, so the consumer has
+        # either given up or will be rejected by the fetch path. Parking more
+        # would pin primary-tier slots for another full timeout on blocks
+        # nobody can fetch, so fail the job now — a prefill that outran the
+        # timeout falls back to local recompute on the decoder either way.
+        if kv_request_id in self._reaped_stores:
+            logger.warning(
+                "P2P %s: submit_store for reaped kv_request_id=%s job_id=%d "
+                "— failing without parking",
+                self._local_id,
+                kv_request_id,
+                job_id,
+            )
+            self._finished_jobs.append(JobResult(job_id=job_id, success=False))
+            return
+
         # Fast path: a session has already received FetchMsg for this id,
         # so we can route the batch straight into its ServerRole.
         session = self._kv_to_session.get(kv_request_id)
@@ -707,13 +760,27 @@ def _reap_unbound_stores(self) -> None:
         """Time out submit_store batches that no peer has ever fetched.
 
         Walks `_unbound_stores` for entries whose oldest batch is older
-        than `_UNBOUND_STORE_TIMEOUT_S`. Drops the kv_request_id, surfaces
-        every batched job as failed, and adds the id to `_failed_req_ids`
-        so a late inbound FetchMsg short-circuits to a clean rejection.
+        than `_unbound_store_timeout_s`. Drops the kv_request_id, surfaces
+        every batched job as failed, and records the id in `_reaped_stores`
+        so a late inbound FetchMsg is rejected in one round trip and a late
+        `submit_store` fails instead of re-parking under the reaped id.
+
+        Also expires `_reaped_stores` entries whose consumer can no longer
+        be waiting, which is why this runs before the empty check below.
         """
+        now = time.monotonic()
+        if self._reaped_stores:
+            retention_deadline = now - _REAPED_ID_RETENTION_S
+            for kid in [
+                kid
+                for kid, reaped_at in self._reaped_stores.items()
+                if reaped_at <= retention_deadline
+            ]:
+                del self._reaped_stores[kid]
+
         if not self._unbound_stores:
             return
-        deadline = time.monotonic() - _UNBOUND_STORE_TIMEOUT_S
+        deadline = now - self._unbound_store_timeout_s
         expired: list[str] | None = None
         for kid, batches in self._unbound_stores.items():
             # Batches are appended in arrival order, so the head is oldest.
@@ -725,7 +792,7 @@ def _reap_unbound_stores(self) -> None:
             return
         for kid in expired:
             batches = self._unbound_stores.pop(kid)
-            self._failed_req_ids.add(kid)
+            self._reaped_stores[kid] = now
             for batch in batches:
                 self._finished_jobs.append(
                     JobResult(job_id=batch.job_id, success=False)
@@ -735,7 +802,7 @@ def _reap_unbound_stores(self) -> None:
                 "without a fetch — failing %d job(s)",
                 self._local_id,
                 kid,
-                _UNBOUND_STORE_TIMEOUT_S,
+                self._unbound_store_timeout_s,
                 len(batches),
             )
 
@@ -779,6 +846,24 @@ def _poll_once(self) -> None:
             # inline in dispatch, so the replayed add_stored_blocks calls
             # match that demand and submit transfers immediately.
             for kv_request_id in result.new_fetch_ids:
+                if kv_request_id in self._reaped_stores:
+                    # This id's blocks were already reaped, so no
+                    # submit_store will ever satisfy the demand on_fetch
+                    # just recorded. Finalize the round now: the peer gets
+                    # TransferDoneMsg(success=False) on this tick and falls
+                    # back to local prefill, instead of parking until its
+                    # own load timeout expires. The id is deliberately left
+                    # unbound so any later submit_store keeps taking the
+                    # reject path. Batches re-parked between the reap and the
+                    # retention prune are failed here rather than left to a
+                    # second reap, which would pin their slots for another
+                    # full timeout.
+                    session.finish_request(kv_request_id)
+                    for batch in self._unbound_stores.pop(kv_request_id, ()):
+                        self._finished_jobs.append(
+                            JobResult(job_id=batch.job_id, success=False)
+                        )
+                    continue
                 self._kv_to_session[kv_request_id] = session
                 for batch in self._unbound_stores.pop(kv_request_id, ()):
                     session.add_stored_blocks(
@@ -810,6 +895,7 @@ def shutdown(self) -> None:
                     JobResult(job_id=batch.job_id, success=False)
                 )
         self._unbound_stores.clear()
+        self._reaped_stores.clear()
         self._control.close()
         self._data.close()