diff --git a/docs/features/kv_offloading_usage.md b/docs/features/kv_offloading_usage.md index 54a54d8988..a5c8684b42 100644 --- a/docs/features/kv_offloading_usage.md +++ b/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 \ + --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. diff --git a/tests/v1/kv_offload/tiering/p2p/test_manager.py b/tests/v1/kv_offload/tiering/p2p/test_manager.py index dc0e384d7c..aceec1cd8f 100644 --- a/tests/v1/kv_offload/tiering/p2p/test_manager.py +++ b/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 @@ class TestGetFinished: 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 @@ class TestGetFinished: # 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 @@ class TestBidirectionalManager: 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 diff --git a/vllm/v1/kv_offload/tiering/p2p/manager.py b/vllm/v1/kv_offload/tiering/p2p/manager.py index 2cd4901710..ba1973b767 100644 --- a/vllm/v1/kv_offload/tiering/p2p/manager.py +++ b/vllm/v1/kv_offload/tiering/p2p/manager.py @@ -47,10 +47,18 @@ logger = init_logger(__name__) # 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 @@ class P2PSecondaryTierManager(SecondaryTierManager): 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 @@ class P2PSecondaryTierManager(SecondaryTierManager): 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 @@ class P2PSecondaryTierManager(SecondaryTierManager): # 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 @@ class P2PSecondaryTierManager(SecondaryTierManager): 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 @@ class P2PSecondaryTierManager(SecondaryTierManager): 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 @@ class P2PSecondaryTierManager(SecondaryTierManager): """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 @@ class P2PSecondaryTierManager(SecondaryTierManager): 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 @@ class P2PSecondaryTierManager(SecondaryTierManager): "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 @@ class P2PSecondaryTierManager(SecondaryTierManager): # 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 @@ class P2PSecondaryTierManager(SecondaryTierManager): JobResult(job_id=batch.job_id, success=False) ) self._unbound_stores.clear() + self._reaped_stores.clear() self._control.close() self._data.close()