From e39adf457425bc85e49c69684b40d000aaa98a52 Mon Sep 17 00:00:00 2001 From: thunguo Date: Thu, 20 Aug 2026 17:08:09 +0800 Subject: [PATCH 1/2] [Bugfix][KV Offload] Bound primary HIT_PENDING waits Track HIT_PENDING deadlines per request and key, then fall back to a stable MISS when the configurable timeout expires. Add cleanup, observability, documentation, and regression tests. Fixes #49829 Signed-off-by: thunguo --- docs/features/kv_offloading_usage.md | 1 + .../tiering/test_tiering_offloading.py | 121 ++++++++++++++++++ vllm/v1/kv_offload/tiering/base.py | 1 + vllm/v1/kv_offload/tiering/manager.py | 89 ++++++++++++- vllm/v1/kv_offload/tiering/metrics.py | 3 + vllm/v1/kv_offload/tiering/spec.py | 14 ++ 6 files changed, 223 insertions(+), 6 deletions(-) diff --git a/docs/features/kv_offloading_usage.md b/docs/features/kv_offloading_usage.md index e2450c662549..da91acd86742 100644 --- a/docs/features/kv_offloading_usage.md +++ b/docs/features/kv_offloading_usage.md @@ -75,6 +75,7 @@ vllm serve \ | `store_threshold` | no | `0` | single-tier | Min lookups before a block is offloaded. Values ≥ 2 are rejected by `TieringOffloadingSpec`. | | `max_tracker_size` | no | `64000` | single-tier | Max entries in the lookup tracker. | | `secondary_tiers` | no | `[]` | multi-tier | List of secondary tier configs (see below). | +| `hit_pending_timeout_s` | no | `5.0` | multi-tier | Maximum seconds a request waits for an in-flight CPU-primary write before treating that block as a miss and recomputing it locally. Must be positive. | | `offload_prompt_only` | no | `true` | both | If `true`, only prompt (prefill) blocks are offloaded; decode blocks are skipped. | | `self_describing_kv_events` | no | `false` | both | Opt-in. When `true` *and* KV cache events are enabled (`--kv-events-config` with `enable_kv_cache_events`), the connector emits self-describing block-granular `BlockStored`/`BlockRemoved` payloads (constituent block hashes, whole-chunk `token_ids`, per-block `block_size`, parent hash, LoRA + group/cache-spec metadata) instead of the placeholder fallback, so external KV-event consumers can index offloaded blocks. Inert unless events are enabled. With `TieringOffloadingSpec`, a CPU promotion is self-describing when a local request observes its primary-tier `HIT` before event translation; otherwise its stored event may retain the placeholder, while a later `HIT` can backfill metadata for removal. Pending-removal/re-promotion races and externally initiated promotions may also produce placeholders, and consumers must ignore removals for unknown hashes. Partial recurrent tails emit the hash-aligned portion from the physical block start through the tail boundary. Other sliding-window/SSM chunks keep the placeholder fallback. In chunk mode (`block_size` > GPU block size, or `blocks_per_chunk` > 1), overlapping chunks re-announce shared per-block hashes, so consumers must reference-count (deduplicate) repeated store/remove announcements. | | `spec_module_path` | no | — | both | Python import path for a custom `OffloadingSpec` not in the built-in registry. Required only when `spec_name` is not built-in (advanced). | diff --git a/tests/v1/kv_offload/tiering/test_tiering_offloading.py b/tests/v1/kv_offload/tiering/test_tiering_offloading.py index 9af714ac9720..b2d537645097 100644 --- a/tests/v1/kv_offload/tiering/test_tiering_offloading.py +++ b/tests/v1/kv_offload/tiering/test_tiering_offloading.py @@ -17,6 +17,7 @@ import pytest import torch +import vllm.v1.kv_offload.tiering.manager as tiering_manager_module from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import ( OffloadingConnectorStats, ) @@ -170,6 +171,10 @@ def test_tiering_spec_collects_secondary_metric_definitions(monkeypatch): metrics[TieringOffloadingMetrics.PROMOTION_ALLOCATION_FAILURES], OffloadingCounterMetadata, ) + assert isinstance( + metrics[TieringOffloadingMetrics.HIT_PENDING_TIMEOUTS], + OffloadingCounterMetadata, + ) for metric_name in ( TieringOffloadingMetrics.PRIMARY_WRITE_USAGE_PERC, TieringOffloadingMetrics.PRIMARY_READ_USAGE_PERC, @@ -218,6 +223,14 @@ def test_tiering_manager_aggregates_secondary_stats(): assert MetricsSecondaryTierManager.MY_TIER_METRIC not in second_stats.data["data"] +@pytest.mark.parametrize("timeout", [0.0, -1.0, float("inf"), float("nan")]) +def test_tiering_manager_rejects_invalid_hit_pending_timeout(timeout): + with pytest.raises(ValueError, match="must be a positive number"): + TieringOffloadingManager( + primary_tier=MagicMock(), hit_pending_timeout_s=timeout + ) + + class TestExampleSecondaryTierManager: """Tests for ExampleSecondaryTierManager implementation.""" @@ -288,6 +301,14 @@ def _start_request(self, req_context: ReqContext = _CTX): if req_context.req_id not in self.manager._req_state: self.manager.on_new_request(req_context) + @staticmethod + def _install_clock(monkeypatch) -> list[float]: + now = [100.0] + clock = MagicMock() + clock.monotonic.side_effect = lambda: now[0] + monkeypatch.setattr(tiering_manager_module, "time", clock) + return now + def test_failed_promotion_finalizes_primary_with_failure(self, manager_setup): """A failed promotion still finalizes the primary slots with success=False; the tier corrects its own verdict.""" @@ -470,6 +491,104 @@ def test_lookup_from_primary(self, manager_setup): # Lookup should find all blocks in primary assert count_hits(self.manager, blocks) == 3 + def test_primary_hit_pending_times_out_and_stays_miss( + self, manager_setup, monkeypatch + ): + """A stalled primary write cannot defer one request indefinitely.""" + now = self._install_clock(monkeypatch) + + ctx = ReqContext(req_id="stalled-primary") + self._start_request(ctx) + block = to_keys([10])[0] + self.primary_tier.lookup = MagicMock(return_value=LookupResult.HIT_PENDING) + + assert self.manager.lookup(block, ctx) is LookupResult.HIT_PENDING + now[0] += 4.999 + assert self.manager.lookup(block, ctx) is LookupResult.HIT_PENDING + now[0] += 0.001 + assert self.manager.lookup(block, ctx) is LookupResult.MISS + + stats = self.manager.get_stats() + assert stats is not None + assert ( + stats.data["data"][TieringOffloadingMetrics.HIT_PENDING_TIMEOUTS][()] == 1 + ) + + self.primary_tier.lookup.return_value = LookupResult.HIT + assert self.manager.lookup(block, ctx) is LookupResult.MISS + + self.manager.on_request_finished(ctx) + assert ctx.req_id not in self.manager._primary_pending_lookups + + def test_primary_hit_pending_deadline_is_per_request( + self, manager_setup, monkeypatch + ): + """A later request gets its own full wait window for the same key.""" + now = self._install_clock(monkeypatch) + + block = to_keys([10])[0] + first_ctx = ReqContext(req_id="first-waiter") + second_ctx = ReqContext(req_id="second-waiter") + self._start_request(first_ctx) + self._start_request(second_ctx) + self.primary_tier.lookup = MagicMock(return_value=LookupResult.HIT_PENDING) + + assert self.manager.lookup(block, first_ctx) is LookupResult.HIT_PENDING + now[0] += 4.0 + assert self.manager.lookup(block, second_ctx) is LookupResult.HIT_PENDING + now[0] += 1.0 + assert self.manager.lookup(block, first_ctx) is LookupResult.MISS + assert self.manager.lookup(block, second_ctx) is LookupResult.HIT_PENDING + + def test_primary_hit_pending_deadline_is_per_key(self, manager_setup, monkeypatch): + """A later key gets its own full wait window within one request.""" + now = self._install_clock(monkeypatch) + + ctx = ReqContext(req_id="multi-key-waiter") + self._start_request(ctx) + first_block, second_block = to_keys([10, 11]) + self.primary_tier.lookup = MagicMock(return_value=LookupResult.HIT_PENDING) + + assert self.manager.lookup(first_block, ctx) is LookupResult.HIT_PENDING + now[0] += 4.0 + assert self.manager.lookup(second_block, ctx) is LookupResult.HIT_PENDING + now[0] += 1.0 + assert self.manager.lookup(first_block, ctx) is LookupResult.MISS + assert self.manager.lookup(second_block, ctx) is LookupResult.HIT_PENDING + + def test_primary_hit_pending_deadline_clears_on_resolution( + self, manager_setup, monkeypatch + ): + """A resolved primary lookup does not retain its old deadline.""" + now = self._install_clock(monkeypatch) + + ctx = ReqContext(req_id="resolving-primary") + self._start_request(ctx) + block = to_keys([10])[0] + self.primary_tier.lookup = MagicMock(return_value=LookupResult.HIT_PENDING) + + assert self.manager.lookup(block, ctx) is LookupResult.HIT_PENDING + now[0] += 4.0 + self.primary_tier.lookup.return_value = LookupResult.HIT + assert self.manager.lookup(block, ctx) is LookupResult.HIT + + now[0] += 2.0 + self.primary_tier.lookup.return_value = LookupResult.HIT_PENDING + assert self.manager.lookup(block, ctx) is LookupResult.HIT_PENDING + + def test_promotion_starts_hit_pending_deadline(self, manager_setup, monkeypatch): + """The lookup that initiates a promotion starts the wait window.""" + now = self._install_clock(monkeypatch) + + ctx = ReqContext(req_id="stalled-promotion") + self._start_request(ctx) + block = to_keys([10])[0] + self.secondary_tier1.blocks[block] = True + + assert self.manager.lookup(block, ctx) is LookupResult.HIT_PENDING + now[0] += 5.0 + assert self.manager.lookup(block, ctx) is LookupResult.MISS + def test_promotion_from_secondary(self, manager_setup): """Test promotion of blocks from secondary to primary tier.""" blocks = to_keys(range(3)) @@ -1101,6 +1220,7 @@ def test_reset_cache_clears_orchestrator_state(self, manager_setup): is LookupResult.HIT_PENDING ) assert self.manager._pending_load_submissions + assert self.manager._primary_pending_lookups # Request-level tier registration. self.secondary_tier1.on_new_request = lambda req_context: ( @@ -1123,6 +1243,7 @@ def test_reset_cache_clears_orchestrator_state(self, manager_setup): # Orchestrator state cleared. assert self.manager._jobs == {} assert self.manager._pending_load_submissions == {} + assert self.manager._primary_pending_lookups == {} assert set(self.manager._req_state) == {_CTX.req_id, rl_ctx.req_id} assert self.manager._processed_jobs_this_step is False diff --git a/vllm/v1/kv_offload/tiering/base.py b/vllm/v1/kv_offload/tiering/base.py index c1e39811e341..5e1cc65a651b 100644 --- a/vllm/v1/kv_offload/tiering/base.py +++ b/vllm/v1/kv_offload/tiering/base.py @@ -52,6 +52,7 @@ class TieringOffloadingMetrics: PROMOTION_ALLOCATION_FAILURES = ( "vllm:kv_offload_tiering_promotion_allocation_failures" ) + HIT_PENDING_TIMEOUTS = "vllm:kv_offload_tiering_hit_pending_timeouts" ACTIVE_PROMOTION_JOBS = "vllm:kv_offload_tiering_active_promotion_jobs" ACTIVE_CASCADE_JOBS = "vllm:kv_offload_tiering_active_cascade_jobs" diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index eb5d5784d432..497b21f81402 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -20,6 +20,7 @@ protecting blocks from eviction until complete_read() is called """ +import math import time from collections.abc import Collection, Iterable, Sequence from dataclasses import dataclass, field @@ -58,6 +59,8 @@ logger = init_logger(__name__) +DEFAULT_HIT_PENDING_TIMEOUT_S = 5.0 + @dataclass class PendingPromotion: @@ -183,6 +186,7 @@ def __init__( self, primary_tier: CPUPrimaryTierOffloadingManager, secondary_tiers: list[SecondaryTierManager] | None = None, + hit_pending_timeout_s: float = DEFAULT_HIT_PENDING_TIMEOUT_S, ): """ Initialize the TieringOffloadingManager. @@ -191,9 +195,20 @@ def __init__( primary_tier: The primary tier manager (CPU-based). secondary_tiers: List of secondary tier managers (e.g., Storage, Network). Can be None or empty list. + hit_pending_timeout_s: Maximum time a request waits for a primary + write before treating that key as a miss. """ + if ( + isinstance(hit_pending_timeout_s, bool) + or not isinstance(hit_pending_timeout_s, (int, float)) + or not math.isfinite(hit_pending_timeout_s) + or hit_pending_timeout_s <= 0 + ): + raise ValueError("hit_pending_timeout_s must be a positive number") + self.primary_tier: CPUPrimaryTierOffloadingManager = primary_tier self.secondary_tiers = secondary_tiers or [] + self._hit_pending_timeout_s = float(hit_pending_timeout_s) self._job_id_counter: int = 0 # Job tracking: maps job_id to metadata for all in-flight transfers. @@ -224,6 +239,11 @@ def __init__( # complete_store(), since complete_store() can still submit cascades. self._req_state: dict[str, RequestState] = {} + # Per-request deadlines for primary-tier writes observed in flight. + # A None deadline marks a key that timed out and must remain a miss for + # the rest of that request. + self._primary_pending_lookups: dict[str, dict[OffloadKey, float | None]] = {} + # Cached ParentManager wrappers for each secondary tier. self._tier_parents: dict[SecondaryTierManager, _SecondaryTierFacingParent] = { tier: _SecondaryTierFacingParent(self, tier_idx) @@ -234,6 +254,44 @@ def __init__( def _transfer_jobs(self) -> dict[JobId, JobMetadata]: return self._jobs + def _primary_lookup_timed_out( + self, key: OffloadKey, req_context: ReqContext + ) -> bool: + pending = self._primary_pending_lookups.get(req_context.req_id) + return pending is not None and key in pending and pending[key] is None + + def _clear_primary_pending_lookup( + self, key: OffloadKey, req_context: ReqContext + ) -> None: + req_id = req_context.req_id + pending = self._primary_pending_lookups.get(req_id) + if pending is None: + return + pending.pop(key, None) + if not pending: + del self._primary_pending_lookups[req_id] + + def _track_primary_pending_lookup( + self, + key: OffloadKey, + req_context: ReqContext, + now: float, + ) -> LookupResult: + pending = self._primary_pending_lookups.setdefault(req_context.req_id, {}) + if key not in pending: + pending[key] = now + self._hit_pending_timeout_s + return LookupResult.HIT_PENDING + + deadline = pending[key] + if deadline is None: + return LookupResult.MISS + if now < deadline: + return LookupResult.HIT_PENDING + + pending[key] = None + self._metrics.on_hit_pending_timeout() + return LookupResult.MISS + def _next_job_id(self) -> JobId: """Generate a unique job ID for async transfer tracking.""" job_id = self._job_id_counter @@ -367,9 +425,23 @@ def lookup( # in time for a promotion this lookup may initiate. self._maybe_process_finished_jobs() + # Once this request has fallen back to local recomputation for a key, + # keep that verdict stable even if the stalled primary write later + # completes. + if self._primary_lookup_timed_out(key, req_context): + return LookupResult.MISS + start_time = time.monotonic() primary_hit = self.primary_tier.lookup(key, req_context) - lookup_duration = time.monotonic() - start_time + lookup_end = time.monotonic() + lookup_duration = lookup_end - start_time + primary_was_pending = primary_hit is LookupResult.HIT_PENDING + if primary_was_pending: + primary_hit = self._track_primary_pending_lookup( + key, req_context, lookup_end + ) + else: + self._clear_primary_pending_lookup(key, req_context) self._metrics.on_lookup( req_context, key, @@ -377,10 +449,8 @@ def lookup( primary_hit, lookup_duration, ) - if primary_hit is LookupResult.HIT: - return LookupResult.HIT - if primary_hit is LookupResult.HIT_PENDING: - return LookupResult.HIT_PENDING + if primary_was_pending or primary_hit is LookupResult.HIT: + return primary_hit any_retry = False for i, tier in enumerate(self.secondary_tiers): @@ -401,7 +471,11 @@ def lookup( lookup_duration, ) promoted = self._initiate_promotion(i, key, req_context) - return LookupResult.MISS if not promoted else LookupResult.HIT_PENDING + if not promoted: + return LookupResult.MISS + return self._track_primary_pending_lookup( + key, req_context, time.monotonic() + ) if result is LookupResult.RETRY: any_retry = True self._metrics.on_lookup( @@ -712,6 +786,7 @@ def on_new_request( Returns REQUEST_LEVEL if ANY secondary tier wants request-level. Only stores REQUEST_LEVEL tier decisions for use in prepare_store. """ + self._primary_pending_lookups.pop(req_context.req_id, None) state = RequestState(req_context=req_context) self._metrics.on_new_request(req_context) for tier_idx, tier in enumerate(self.secondary_tiers): @@ -738,6 +813,7 @@ def on_request_finished( *, exclude_tier_idx: int | None = None, ) -> None: + self._primary_pending_lookups.pop(req_context.req_id, None) self.primary_tier.on_request_finished(req_context) state = self._req_state[req_context.req_id] state.is_finished = True @@ -843,6 +919,7 @@ def reset_cache(self) -> None: # reset below invalidates; their submit_load() has not yet been # called so no tier I/O is touching that memory. self._pending_load_submissions.clear() + self._primary_pending_lookups.clear() self._metrics.assert_idle() finished_req_ids = [] diff --git a/vllm/v1/kv_offload/tiering/metrics.py b/vllm/v1/kv_offload/tiering/metrics.py index 60907142ac38..0d43e3e53f3f 100644 --- a/vllm/v1/kv_offload/tiering/metrics.py +++ b/vllm/v1/kv_offload/tiering/metrics.py @@ -130,6 +130,9 @@ def on_promotion_allocation_failure(self) -> None: TieringOffloadingMetrics.PROMOTION_ALLOCATION_FAILURES ) + def on_hit_pending_timeout(self) -> None: + self._stats.increase_counter(TieringOffloadingMetrics.HIT_PENDING_TIMEOUTS) + def take_stats(self) -> OffloadingConnectorStats | None: active_transfer_stats = OffloadingConnectorStats() self._observe_active_transfer_stats(active_transfer_stats) diff --git a/vllm/v1/kv_offload/tiering/spec.py b/vllm/v1/kv_offload/tiering/spec.py index bb0325ca2842..a3bbf51d8e22 100644 --- a/vllm/v1/kv_offload/tiering/spec.py +++ b/vllm/v1/kv_offload/tiering/spec.py @@ -16,6 +16,8 @@ - cache_policy_module_path: (optional) Python import path to load eviction_policy from when it names an out-of-tree CachePolicy not registered via CachePolicyFactory + - hit_pending_timeout_s: (optional) Maximum time a request waits for an + in-flight primary-tier write before recomputing locally (default: 5 seconds) - secondary_tiers: (optional) List of secondary tier configurations Each secondary tier config is a dict with: - type: (required) Type of secondary tier (e.g., "example", "fs", @@ -74,6 +76,7 @@ from vllm.v1.kv_offload.tiering.base import TieringOffloadingMetrics from vllm.v1.kv_offload.tiering.factory import SecondaryTierFactory from vllm.v1.kv_offload.tiering.manager import ( + DEFAULT_HIT_PENDING_TIMEOUT_S, CPUPrimaryTierOffloadingManager, TieringOffloadingManager, ) @@ -209,6 +212,14 @@ def build_metric_definitions( ), ) ) + metrics[TieringOffloadingMetrics.HIT_PENDING_TIMEOUTS] = ( + OffloadingCounterMetadata( + documentation=( + "Number of primary-tier writes that exceeded the per-request " + "HIT_PENDING deadline." + ), + ) + ) metrics[TieringOffloadingMetrics.PRIMARY_WRITE_USAGE_PERC] = ( OffloadingGaugeMetadata( documentation=( @@ -335,6 +346,9 @@ def get_manager(self) -> OffloadingManager: tiering_manager = TieringOffloadingManager( primary_tier=primary_tier, secondary_tiers=secondary_tiers, + hit_pending_timeout_s=self.extra_config.get( + "hit_pending_timeout_s", DEFAULT_HIT_PENDING_TIMEOUT_S + ), ) self._manager = tiering_manager except Exception: From dfe0fe212db1981cb166771103cf5a5dc6237b4d Mon Sep 17 00:00:00 2001 From: thunguo Date: Thu, 20 Aug 2026 19:32:03 +0800 Subject: [PATCH 2/2] [Bugfix][KV Offload] Preserve timeout across failed promotions Keep each request-key deadline across failed promotion rounds so retries cannot renew the deferral window. Preserve primary lookup metric semantics, add rate-limited diagnostics, and extend regression coverage. Signed-off-by: thunguo --- .../tiering/test_tiering_offloading.py | 61 ++++++++++++++- vllm/v1/kv_offload/tiering/manager.py | 76 ++++++++++++------- vllm/v1/kv_offload/tiering/spec.py | 6 +- 3 files changed, 108 insertions(+), 35 deletions(-) diff --git a/tests/v1/kv_offload/tiering/test_tiering_offloading.py b/tests/v1/kv_offload/tiering/test_tiering_offloading.py index b2d537645097..9f54f10c4c6d 100644 --- a/tests/v1/kv_offload/tiering/test_tiering_offloading.py +++ b/tests/v1/kv_offload/tiering/test_tiering_offloading.py @@ -11,6 +11,7 @@ 5. Eviction coordination between tiers """ +import logging from collections.abc import Iterable from unittest.mock import MagicMock @@ -225,10 +226,11 @@ def test_tiering_manager_aggregates_secondary_stats(): @pytest.mark.parametrize("timeout", [0.0, -1.0, float("inf"), float("nan")]) def test_tiering_manager_rejects_invalid_hit_pending_timeout(timeout): - with pytest.raises(ValueError, match="must be a positive number"): + with pytest.raises(ValueError, match="must be a positive number") as exc_info: TieringOffloadingManager( primary_tier=MagicMock(), hit_pending_timeout_s=timeout ) + assert repr(timeout) in str(exc_info.value) class TestExampleSecondaryTierManager: @@ -304,9 +306,7 @@ def _start_request(self, req_context: ReqContext = _CTX): @staticmethod def _install_clock(monkeypatch) -> list[float]: now = [100.0] - clock = MagicMock() - clock.monotonic.side_effect = lambda: now[0] - monkeypatch.setattr(tiering_manager_module, "time", clock) + monkeypatch.setattr(tiering_manager_module.time, "monotonic", lambda: now[0]) return now def test_failed_promotion_finalizes_primary_with_failure(self, manager_setup): @@ -513,6 +513,7 @@ def test_primary_hit_pending_times_out_and_stays_miss( assert ( stats.data["data"][TieringOffloadingMetrics.HIT_PENDING_TIMEOUTS][()] == 1 ) + assert TieringOffloadingMetrics.BLOCK_QUERIES not in stats.data["data"] self.primary_tier.lookup.return_value = LookupResult.HIT assert self.manager.lookup(block, ctx) is LookupResult.MISS @@ -520,6 +521,32 @@ def test_primary_hit_pending_times_out_and_stays_miss( self.manager.on_request_finished(ctx) assert ctx.req_id not in self.manager._primary_pending_lookups + def test_primary_hit_pending_timeout_warning_is_rate_limited( + self, manager_setup, monkeypatch, caplog + ): + now = self._install_clock(monkeypatch) + + ctx = ReqContext(req_id="warning-rate-limit") + self._start_request(ctx) + first_block, second_block = to_keys([10, 11]) + self.primary_tier.lookup = MagicMock(return_value=LookupResult.HIT_PENDING) + + assert self.manager.lookup(first_block, ctx) is LookupResult.HIT_PENDING + assert self.manager.lookup(second_block, ctx) is LookupResult.HIT_PENDING + + with caplog.at_level(logging.WARNING, logger=tiering_manager_module.__name__): + now[0] += 5.0 + assert self.manager.lookup(first_block, ctx) is LookupResult.MISS + assert self.manager.lookup(second_block, ctx) is LookupResult.MISS + + messages = [ + record.getMessage() + for record in caplog.records + if "HIT_PENDING timeout" in record.getMessage() + ] + assert len(messages) == 1 + assert repr(first_block) in messages[0] + def test_primary_hit_pending_deadline_is_per_request( self, manager_setup, monkeypatch ): @@ -589,6 +616,32 @@ def test_promotion_starts_hit_pending_deadline(self, manager_setup, monkeypatch) now[0] += 5.0 assert self.manager.lookup(block, ctx) is LookupResult.MISS + def test_failed_promotions_do_not_reset_hit_pending_deadline( + self, manager_setup, monkeypatch + ): + """Repeated failed promotions share one total deferral budget.""" + now = self._install_clock(monkeypatch) + + ctx = ReqContext(req_id="repeated-promotion-failures") + self._start_request(ctx) + block = to_keys([10])[0] + self.secondary_tier1.blocks[block] = True + + def fail_promotion(job: TransferJob) -> None: + self.secondary_tier1.completed_jobs.append( + JobResult(job_id=job.job_id, success=False) + ) + + self.secondary_tier1.submit_load = fail_promotion + + for _ in range(2): + assert self.manager.lookup(block, ctx) is LookupResult.HIT_PENDING + self._simulate_on_schedule_end() + now[0] += 3.0 + self._simulate_on_schedule_end() + + assert self.manager.lookup(block, ctx) is LookupResult.MISS + def test_promotion_from_secondary(self, manager_setup): """Test promotion of blocks from secondary to primary tier.""" blocks = to_keys(range(3)) diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index 497b21f81402..8df080cb31f7 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -60,6 +60,7 @@ logger = init_logger(__name__) DEFAULT_HIT_PENDING_TIMEOUT_S = 5.0 +HIT_PENDING_WARNING_INTERVAL_S = 60.0 @dataclass @@ -79,6 +80,12 @@ class RequestState: request_level_tiers: set[int] | None = None +@dataclass(slots=True) +class PrimaryPendingLookupState: + deadlines: dict[OffloadKey, float] = field(default_factory=dict) + downgraded_keys: set[OffloadKey] = field(default_factory=set) + + class JobMetadata(NamedTuple): transfer_job: TransferJob tier_idx: int @@ -198,17 +205,16 @@ def __init__( hit_pending_timeout_s: Maximum time a request waits for a primary write before treating that key as a miss. """ - if ( - isinstance(hit_pending_timeout_s, bool) - or not isinstance(hit_pending_timeout_s, (int, float)) - or not math.isfinite(hit_pending_timeout_s) - or hit_pending_timeout_s <= 0 - ): - raise ValueError("hit_pending_timeout_s must be a positive number") + if not math.isfinite(hit_pending_timeout_s) or hit_pending_timeout_s <= 0: + raise ValueError( + "hit_pending_timeout_s must be a positive number, " + f"got {hit_pending_timeout_s!r}" + ) self.primary_tier: CPUPrimaryTierOffloadingManager = primary_tier self.secondary_tiers = secondary_tiers or [] - self._hit_pending_timeout_s = float(hit_pending_timeout_s) + self._hit_pending_timeout_s = hit_pending_timeout_s + self._next_hit_pending_warning_at: float = 0.0 self._job_id_counter: int = 0 # Job tracking: maps job_id to metadata for all in-flight transfers. @@ -239,10 +245,8 @@ def __init__( # complete_store(), since complete_store() can still submit cascades. self._req_state: dict[str, RequestState] = {} - # Per-request deadlines for primary-tier writes observed in flight. - # A None deadline marks a key that timed out and must remain a miss for - # the rest of that request. - self._primary_pending_lookups: dict[str, dict[OffloadKey, float | None]] = {} + # Per-request primary-tier writes observed in flight. + self._primary_pending_lookups: dict[str, PrimaryPendingLookupState] = {} # Cached ParentManager wrappers for each secondary tier. self._tier_parents: dict[SecondaryTierManager, _SecondaryTierFacingParent] = { @@ -257,18 +261,19 @@ def _transfer_jobs(self) -> dict[JobId, JobMetadata]: def _primary_lookup_timed_out( self, key: OffloadKey, req_context: ReqContext ) -> bool: - pending = self._primary_pending_lookups.get(req_context.req_id) - return pending is not None and key in pending and pending[key] is None + state = self._primary_pending_lookups.get(req_context.req_id) + return state is not None and key in state.downgraded_keys def _clear_primary_pending_lookup( self, key: OffloadKey, req_context: ReqContext ) -> None: req_id = req_context.req_id - pending = self._primary_pending_lookups.get(req_id) - if pending is None: + state = self._primary_pending_lookups.get(req_id) + if state is None: return - pending.pop(key, None) - if not pending: + state.deadlines.pop(key, None) + state.downgraded_keys.discard(key) + if not state.deadlines and not state.downgraded_keys: del self._primary_pending_lookups[req_id] def _track_primary_pending_lookup( @@ -277,19 +282,31 @@ def _track_primary_pending_lookup( req_context: ReqContext, now: float, ) -> LookupResult: - pending = self._primary_pending_lookups.setdefault(req_context.req_id, {}) - if key not in pending: - pending[key] = now + self._hit_pending_timeout_s - return LookupResult.HIT_PENDING + state = self._primary_pending_lookups.setdefault( + req_context.req_id, PrimaryPendingLookupState() + ) + if key in state.downgraded_keys: + return LookupResult.MISS - deadline = pending[key] + deadline = state.deadlines.get(key) if deadline is None: - return LookupResult.MISS + state.deadlines[key] = now + self._hit_pending_timeout_s + return LookupResult.HIT_PENDING if now < deadline: return LookupResult.HIT_PENDING - pending[key] = None + del state.deadlines[key] + state.downgraded_keys.add(key) self._metrics.on_hit_pending_timeout() + if now >= self._next_hit_pending_warning_at: + logger.warning( + "Primary-tier write for key %r exceeded the %.1f-second " + "HIT_PENDING timeout; treating it as a miss for request %s", + key, + self._hit_pending_timeout_s, + req_context.req_id, + ) + self._next_hit_pending_warning_at = now + HIT_PENDING_WARNING_INTERVAL_S return LookupResult.MISS def _next_job_id(self) -> JobId: @@ -432,21 +449,22 @@ def lookup( return LookupResult.MISS start_time = time.monotonic() - primary_hit = self.primary_tier.lookup(key, req_context) + observed_primary_hit = self.primary_tier.lookup(key, req_context) lookup_end = time.monotonic() lookup_duration = lookup_end - start_time - primary_was_pending = primary_hit is LookupResult.HIT_PENDING + primary_hit = observed_primary_hit + primary_was_pending = observed_primary_hit is LookupResult.HIT_PENDING if primary_was_pending: primary_hit = self._track_primary_pending_lookup( key, req_context, lookup_end ) - else: + elif observed_primary_hit is LookupResult.HIT: self._clear_primary_pending_lookup(key, req_context) self._metrics.on_lookup( req_context, key, self._metrics.primary_tier_label, - primary_hit, + observed_primary_hit, lookup_duration, ) if primary_was_pending or primary_hit is LookupResult.HIT: diff --git a/vllm/v1/kv_offload/tiering/spec.py b/vllm/v1/kv_offload/tiering/spec.py index a3bbf51d8e22..55ad58b6049c 100644 --- a/vllm/v1/kv_offload/tiering/spec.py +++ b/vllm/v1/kv_offload/tiering/spec.py @@ -346,8 +346,10 @@ def get_manager(self) -> OffloadingManager: tiering_manager = TieringOffloadingManager( primary_tier=primary_tier, secondary_tiers=secondary_tiers, - hit_pending_timeout_s=self.extra_config.get( - "hit_pending_timeout_s", DEFAULT_HIT_PENDING_TIMEOUT_S + hit_pending_timeout_s=float( + self.extra_config.get( + "hit_pending_timeout_s", DEFAULT_HIT_PENDING_TIMEOUT_S + ) ), ) self._manager = tiering_manager