diff --git a/docs/features/kv_offloading_usage.md b/docs/features/kv_offloading_usage.md index 72f838c7d2ca..503ed8d3de56 100644 --- a/docs/features/kv_offloading_usage.md +++ b/docs/features/kv_offloading_usage.md @@ -76,6 +76,7 @@ vllm serve \ | `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). | | `offload_prompt_only` | no | `true` | both | If `true`, only prompt (prefill) blocks are offloaded; decode blocks are skipped. | +| `hit_pending_deadline_s` | no | `60` | both | Seconds a request may stay deferred because a candidate block is still write-in-flight in the primary tier. Past this, that block is treated as a miss for the rest of the request and the candidate prefix is truncated there, so the request recomputes locally rather than deferring until the client times out. The deadline is timed per blocking block: when the block holding a request up changes, the new one starts a deadline of its own instead of inheriting the elapsed wait. The default sits above the P2P worst case for a live promotion (`_LOAD_TIMEOUT_S` 30s + `_ABORT_ACK_TIMEOUT_S` 10s), so on that tier it fires only for writes that have genuinely leaked. **The CPU, filesystem and object-store backends have no equivalent bound on transfer time**, so there a healthy but slow write can cross the deadline too; the only consequence is a cache miss and local recomputation, never incorrect output. Expirations are counted by `vllm:kv_offload_hit_pending_deadline_expired`; a non-zero value is worth investigating on P2P but may simply reflect slow promotions on the other backends. Set `0` to disable and defer indefinitely. | | `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. Full-attention groups only; sliding-window/SSM groups 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_connector/unit/offloading_connector/test_scheduler.py b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py index 661078041f10..91207a168141 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py @@ -12,6 +12,9 @@ ) from tests.v1.kv_connector.unit.utils import EOS_TOKEN_ID from vllm.distributed.kv_events import MEDIUM_CPU, BlockRemoved, BlockStored +from vllm.distributed.kv_transfer.kv_connector.v1.offloading import ( + scheduler as sched_module, +) from vllm.distributed.kv_transfer.kv_connector.v1.offloading.common import ( OffloadingConnectorMetadata, OffloadingWorkerMetadata, @@ -21,8 +24,11 @@ _ConnectorMetricName, ) from vllm.distributed.kv_transfer.kv_connector.v1.offloading.scheduler import ( + GroupOffloadConfig, OffloadingConnectorScheduler, RequestOffloadState, + SchedulerOffloadConfig, + _LookupScan, is_store_reachable_swa_chunk, ) from vllm.v1.core.kv_cache_utils import BlockHash @@ -59,6 +65,11 @@ def _reduce_kv_connector_stats(runner): return reduced +def _stub_prefix_scan(num_hit_chunks: int): + """Stub `_maximal_prefix_lookup` with a fixed hit count, no HIT_PENDING.""" + return lambda keys, ctx, *_: _LookupScan(num_hit_chunks, None) + + def test_scheduler_reports_allocation_failure(request_runner): runner = request_runner( block_size=4, @@ -316,7 +327,7 @@ def take_raw_events(): req_status.req, group_config, 0, - ) + ).num_hit_chunks == 1 ) assert key in tracker._pending_event_metadata @@ -487,7 +498,7 @@ def test_offloading_connector(request_runner, async_scheduling: bool): runner.manager.prepare_store.side_effect = lambda keys, req_context: ( generate_store_output([]) ) - runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 1 + runner.connector_scheduler._maximal_prefix_lookup = _stub_prefix_scan(1) runner.run(decoded_tokens=[EOS_TOKEN_ID], expected_loaded=(0, 1, 2)) # single block lookup with a hit in a middle block @@ -495,7 +506,7 @@ def test_offloading_connector(request_runner, async_scheduling: bool): runner.manager.prepare_store.side_effect = lambda keys, req_context: ( generate_store_output([]) ) - runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 1 + runner.connector_scheduler._maximal_prefix_lookup = _stub_prefix_scan(1) runner.run(decoded_tokens=[EOS_TOKEN_ID], expected_loaded=(3, 4, 5)) @@ -553,7 +564,7 @@ def test_request_preemption(request_runner, async_scheduling: bool): # request should now return from preemption # re-load [0, ..., 8] from the CPU and store [9, 10, 11] - runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 3 + runner.connector_scheduler._maximal_prefix_lookup = _stub_prefix_scan(3) runner.manager.prepare_store.side_effect = lambda keys, req_context: ( generate_store_output(keys) ) @@ -712,7 +723,7 @@ def test_concurrent_lookups_of_the_same_prefix(request_runner, async_scheduling: # start a request to load the first block, but don't complete runner.scheduler.reset_prefix_cache() runner.new_request(token_ids=[0] * tokens_per_chunk) - runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 1 + runner.connector_scheduler._maximal_prefix_lookup = _stub_prefix_scan(1) runner.run( decoded_tokens=[], complete_transfers=False, @@ -724,7 +735,7 @@ def test_concurrent_lookups_of_the_same_prefix(request_runner, async_scheduling: # start a new request to load the same first block runner.new_request(token_ids=[0] * tokens_per_chunk) - runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 1 + runner.connector_scheduler._maximal_prefix_lookup = _stub_prefix_scan(1) runner.run( decoded_tokens=[], complete_transfers=False, @@ -776,7 +787,7 @@ def test_abort_loading_requests(request_runner, async_scheduling: bool): # start a request to load the first block, but don't complete runner.scheduler.reset_prefix_cache() runner.new_request(token_ids=[0] * tokens_per_chunk) - runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 1 + runner.connector_scheduler._maximal_prefix_lookup = _stub_prefix_scan(1) runner.run( decoded_tokens=[], complete_transfers=False, @@ -1083,16 +1094,39 @@ def _make_scheduler_with_lookup( _LOOKUP_GROUP_CONFIG = MagicMock() -def _maximal_lookup(sched, keys, start_chunk_idx: int = 0): +def _maximal_scan( + sched, keys, start_chunk_idx: int = 0, downgrade_hit_pending: bool = False +): + """Run a prefix scan and return the full _LookupScan.""" return sched._maximal_prefix_lookup( keys, _EMPTY_REQ_CTX, _LOOKUP_REQ, _LOOKUP_GROUP_CONFIG, start_chunk_idx, + downgrade_hit_pending, ) +def _maximal_lookup( + sched, keys, start_chunk_idx: int = 0, downgrade_hit_pending: bool = False +): + return _maximal_scan( + sched, keys, start_chunk_idx, downgrade_hit_pending + ).num_hit_chunks + + +def _sliding_scan(sched, keys, window: int, downgrade_hit_pending: bool = False): + """Run a sliding-window scan and return the full _LookupScan.""" + return sched._sliding_window_lookup( + keys, window, _EMPTY_REQ_CTX, downgrade_hit_pending + ) + + +def _sliding_lookup(sched, keys, window: int, downgrade_hit_pending: bool = False): + return _sliding_scan(sched, keys, window, downgrade_hit_pending).num_hit_chunks + + class TestMaximalPrefixLookup: def test_all_hit(self): sched = _make_scheduler_with_lookup({1: LookupResult.HIT, 2: LookupResult.HIT}) @@ -1212,43 +1246,518 @@ def test_retry_stops_at_miss(self): assert sched.manager.lookup.call_count == 2 sched._events_tracker.record_lookup.assert_not_called() + def test_hit_pending_downgrade_truncates_prefix(self): + """Once downgraded, the prefix is cut at the first HIT_PENDING. + + The pending chunk must not be counted: prepare_load asserts + block.is_ready, so including it would turn the hang into a crash. + """ + sched = _make_scheduler_with_lookup( + { + 1: LookupResult.HIT, + 2: LookupResult.HIT_PENDING, + 3: LookupResult.HIT, + } + ) + # Resolves to 1 — not None (no longer deferred), and not 2 or 3. + assert ( + _maximal_lookup(sched, to_keys([1, 2, 3]), downgrade_hit_pending=True) == 1 + ) + # Scanning stops at the pending key, so key 3 is never looked up. + assert sched.manager.lookup.call_count == 2 + + def test_hit_pending_downgrade_leaves_pure_hits_alone(self): + """Downgrading changes nothing when no block is pending.""" + sched = _make_scheduler_with_lookup({1: LookupResult.HIT, 2: LookupResult.HIT}) + assert _maximal_lookup(sched, to_keys([1, 2]), downgrade_hit_pending=True) == 2 + + def test_hit_pending_downgrade_does_not_rescue_retry(self): + """RETRY still defers after a downgrade — only HIT_PENDING is bounded.""" + sched = _make_scheduler_with_lookup({1: LookupResult.RETRY}) + scan = _maximal_scan(sched, to_keys([1]), downgrade_hit_pending=True) + assert scan.num_hit_chunks is None + assert scan.first_pending_key is None + + +class TestLookupScanSignal: + """The scan reports *which* block deferred it, so the deadline can key on it. + + An identity rather than a bare flag is what stops a newly pending key from + inheriting an older key's elapsed time. + """ + + def test_prefix_scan_reports_hit_pending(self): + sched = _make_scheduler_with_lookup({1: LookupResult.HIT_PENDING}) + keys = to_keys([1]) + assert _maximal_scan(sched, keys).first_pending_key == keys[0] + + def test_prefix_scan_reports_the_lowest_index_pending_key(self): + """The prefix cannot resolve past the earliest pending block.""" + sched = _make_scheduler_with_lookup( + { + 1: LookupResult.HIT, + 2: LookupResult.HIT_PENDING, + 3: LookupResult.HIT_PENDING, + } + ) + keys = to_keys([1, 2, 3]) + assert _maximal_scan(sched, keys).first_pending_key == keys[1] + + def test_prefix_scan_does_not_report_retry(self): + sched = _make_scheduler_with_lookup({1: LookupResult.RETRY}) + assert _maximal_scan(sched, to_keys([1])).first_pending_key is None + + def test_prefix_scan_reports_nothing_when_resolved(self): + sched = _make_scheduler_with_lookup({1: LookupResult.HIT}) + assert _maximal_scan(sched, to_keys([1])).first_pending_key is None + + def test_sliding_scan_reports_hit_pending(self): + sched = _make_scheduler_with_lookup({1: LookupResult.HIT_PENDING}) + keys = to_keys([1]) + assert _sliding_scan(sched, keys, 1).first_pending_key == keys[0] + + def test_sliding_scan_reports_the_highest_index_pending_key(self): + """This scan walks backwards, so its first blocker is the last key.""" + sched = _make_scheduler_with_lookup( + { + 1: LookupResult.HIT_PENDING, + 2: LookupResult.HIT_PENDING, + 3: LookupResult.HIT, + } + ) + keys = to_keys([1, 2, 3]) + assert _sliding_scan(sched, keys, 3).first_pending_key == keys[1] + + def test_sliding_scan_does_not_report_retry(self): + sched = _make_scheduler_with_lookup({1: LookupResult.RETRY}) + assert _sliding_scan(sched, to_keys([1]), 1).first_pending_key is None + + def test_downgraded_hit_pending_is_not_reported(self): + """A downgraded pending block is a miss, so it must not re-arm.""" + sched = _make_scheduler_with_lookup({1: LookupResult.HIT_PENDING}) + scan = _maximal_scan(sched, to_keys([1]), downgrade_hit_pending=True) + assert scan.first_pending_key is None + assert scan.num_hit_chunks == 0 + + +# The shipped default for hit_pending_deadline_s, hardcoded rather than +# imported so this module still imports against a tree without the fix — the +# regression test below has to fail on behavior, not on a missing symbol. +# tests/v1/kv_offload/test_spec_config.py pins the real default to this value. +_SHIPPED_DEADLINE_S = 60.0 + + +class _FakeClock: + """Monotonic clock the test advances by hand. + + Lets a deferral span the deadline without sleeping, and keeps the tests + honest: they never set the deadline knob, so the scheduler runs on the + shipped default. That is what makes them fail on an unfixed tree by + behavior rather than by a missing keyword argument. + """ + + def __init__(self, start: float = 10_000.0): + self.now = start + + def __call__(self) -> float: + return self.now + + def advance(self, seconds: float) -> None: + self.now += seconds + + +def _install_fake_clock(monkeypatch, sched_module) -> _FakeClock: + """Swap the scheduler module's `time` for one the test drives.""" + clock = _FakeClock() + monkeypatch.setattr(sched_module, "time", SimpleNamespace(monotonic=clock)) + return clock + + +def _make_deadline_scheduler(lookup_results, **config_overrides): + """Build a real scheduler wired for get_num_new_matched_tokens(). + + Only the attributes that code path actually touches are constructed, so + this needs no engine, no KV cache and no accelerator. `lookup_results` is + the same live dict `_make_scheduler_with_lookup` reads, so a test can + change what the manager reports between scheduling steps by mutating it. + + The deadline is deliberately left unset so it takes the shipped default. + """ + sched = _make_scheduler_with_lookup(lookup_results) + sched.config = SchedulerOffloadConfig( + kv_group_configs=( + GroupOffloadConfig( + group_idx=0, + tokens_per_block=4, + tokens_per_chunk=4, + hashes_per_chunk=1, + kv_event_group_spec=MagicMock(), + sliding_window_size_in_chunks=None, + ), + ), + blocks_per_chunk=1, + num_workers=1, + offload_prompt_only=False, + **config_overrides, + ) + sched._lookup_groups = (0,) + sched._sliding_window_groups = () + sched._mamba_align_size = None + sched._chunks_being_loaded = {} + sched._connector_stats = OffloadingConnectorStats() + sched._req_status = {} + return sched + + +def _make_deadline_request(req_id: str = "req-hit-pending"): + """A 3-chunk request whose block hashes decode to keys 1, 2 and 3.""" + return SimpleNamespace( + request_id=req_id, + block_hashes=[b"1", b"2", b"3"], + num_tokens=12, + num_prompt_tokens=12, + kv_transfer_params=None, + skip_reading_prefix_cache=False, + is_finished=lambda: False, + ) + + +def _expiry_count(sched) -> int: + return sched._connector_stats.reduce().get( + _ConnectorMetricName.HIT_PENDING_DEADLINE_EXPIRED, 0 + ) + + +class TestHitPendingDeadline: + """Request-level behavior of the HIT_PENDING deadline (issue #49829).""" + + def test_deadline_releases_wedged_request(self, monkeypatch): + """A permanently-HIT_PENDING key must not defer a request forever. + + A leaked primary-tier write leaves a key HIT_PENDING, and every + request whose candidate prefix reaches it is deferred on every + scheduling step until the client gives up. + + This is the regression guard for #49829, and it deliberately touches + no API added by the fix: it only advances the clock and reads the + return value of get_num_new_matched_tokens(). On an unfixed tree every + step keeps returning (None, False) and the last assertion fails. + """ + clock = _install_fake_clock(monkeypatch, sched_module) + sched = _make_deadline_scheduler( + { + 1: LookupResult.HIT_PENDING, + 2: LookupResult.HIT_PENDING, + 3: LookupResult.HIT_PENDING, + } + ) + request = _make_deadline_request() + sched.on_new_request(request) + + # The request defers, as it does today. + assert sched.get_num_new_matched_tokens(request, 0) == (None, False) + + # Still inside the deadline: the scheduler keeps re-polling and the + # request keeps deferring, burning a lookup per candidate key per step. + clock.advance(_SHIPPED_DEADLINE_S / 2) + assert sched.get_num_new_matched_tokens(request, 0) == (None, False) + + # Past the deadline. One more step observes the expiry... + clock.advance(_SHIPPED_DEADLINE_S) + assert sched.get_num_new_matched_tokens(request, 0) == (None, False) + + # ...and the request is released: 0 external tokens, so it recomputes + # locally instead of waiting on a write that is not coming. + assert sched.get_num_new_matched_tokens(request, 0) == (0, False) + + # It stays released no matter how long the key remains pending. + for _ in range(100): + assert sched.get_num_new_matched_tokens(request, 0) == (0, False) + + def test_partial_prefix_survives_expiry(self, monkeypatch): + """Ready chunks before the pending one are still served after expiry. + + Behavioral like the test above: no fix-added API is referenced. + """ + clock = _install_fake_clock(monkeypatch, sched_module) + sched = _make_deadline_scheduler( + { + 1: LookupResult.HIT, + 2: LookupResult.HIT, + 3: LookupResult.HIT_PENDING, + } + ) + request = _make_deadline_request() + sched.on_new_request(request) + + assert sched.get_num_new_matched_tokens(request, 0) == (None, False) + clock.advance(_SHIPPED_DEADLINE_S + 1.0) + sched.get_num_new_matched_tokens(request, 0) + + # Two ready chunks of 4 tokens each are loaded; the pending third is + # dropped rather than counted, so it never reaches prepare_load. + assert sched.get_num_new_matched_tokens(request, 0) == (8, True) + + def test_expiry_is_recorded_once_per_request(self, monkeypatch): + """The expiry is sticky, so the counter must not tick every step.""" + clock = _install_fake_clock(monkeypatch, sched_module) + sched = _make_deadline_scheduler({1: LookupResult.HIT_PENDING}) + request = _make_deadline_request() + sched.on_new_request(request) + req_status = sched._req_status[request.request_id] + + sched.get_num_new_matched_tokens(request, 0) + clock.advance(_SHIPPED_DEADLINE_S + 1.0) + for _ in range(5): + sched.get_num_new_matched_tokens(request, 0) + + assert req_status.hit_pending_expired is True + assert req_status.hit_pending_start_time is None + assert _expiry_count(sched) == 1 + + def test_disarms_on_retry_only_pass(self, monkeypatch): + """A pass that deferred without HIT_PENDING must clear the timer. + + Otherwise a request that saw one transient HIT_PENDING and then + deferred for unrelated reasons would carry a stale armed timer and + eventually expire without any stalled write. + """ + _install_fake_clock(monkeypatch, sched_module) + results = {1: LookupResult.HIT_PENDING} + sched = _make_deadline_scheduler(results) + request = _make_deadline_request() + sched.on_new_request(request) + req_status = sched._req_status[request.request_id] + + sched.get_num_new_matched_tokens(request, 0) + assert req_status.hit_pending_start_time is not None + + # The pending write landed elsewhere; now the backend only asks for a + # retry. Nothing is pending any more, so the timer must be dropped. + results[1] = LookupResult.RETRY + assert sched.get_num_new_matched_tokens(request, 0) == (None, False) + assert req_status.hit_pending_start_time is None + assert req_status.hit_pending_expired is False + + def test_new_blocker_does_not_inherit_the_previous_clock(self, monkeypatch): + """A newly pending key gets its own deadline, not the old key's clock. + + Without the identity check, a request 59s into waiting on key 1 would + carry that elapsed time over to key 2 the moment key 1 resolved and + expire a second later, despite key 2's own write being healthy. Any + key inheriting an almost-spent clock leaves the 60s default no margin. + """ + clock = _install_fake_clock(monkeypatch, sched_module) + results = { + 1: LookupResult.HIT_PENDING, + 2: LookupResult.HIT_PENDING, + 3: LookupResult.HIT_PENDING, + } + sched = _make_deadline_scheduler(results) + request = _make_deadline_request() + sched.on_new_request(request) + req_status = sched._req_status[request.request_id] + + # Key 1 blocks the prefix and is nearly out of time. + assert sched.get_num_new_matched_tokens(request, 0) == (None, False) + clock.advance(_SHIPPED_DEADLINE_S - 1.0) + assert sched.get_num_new_matched_tokens(request, 0) == (None, False) + assert req_status.hit_pending_expired is False + + # Key 1's write lands, so key 2 becomes the blocker. Its own write has + # only just started; it must not be charged for key 1's wait. + results[1] = LookupResult.HIT + assert sched.get_num_new_matched_tokens(request, 0) == (None, False) + clock.advance(2.0) + assert sched.get_num_new_matched_tokens(request, 0) == (None, False) + assert req_status.hit_pending_expired is False + assert _expiry_count(sched) == 0 + + # Key 2 does expire, but only after a full deadline of its own. + clock.advance(_SHIPPED_DEADLINE_S) + sched.get_num_new_matched_tokens(request, 0) + assert req_status.hit_pending_expired is True + assert _expiry_count(sched) == 1 + + def test_unchanged_blocker_keeps_accruing(self, monkeypatch): + """Re-arming is for a *changed* blocker only. + + Counterpart to the test above: if every pass re-armed, a permanently + stalled key would never expire and #49829 would be back. + """ + clock = _install_fake_clock(monkeypatch, sched_module) + sched = _make_deadline_scheduler({1: LookupResult.HIT_PENDING}) + request = _make_deadline_request() + sched.on_new_request(request) + req_status = sched._req_status[request.request_id] + + # Split the wait across two passes that together reach the deadline. + armed_at = None + for _ in range(2): + assert sched.get_num_new_matched_tokens(request, 0) == (None, False) + if armed_at is None: + armed_at = req_status.hit_pending_start_time + assert armed_at is not None + # Same blocker, so the arming instant never moves. + assert req_status.hit_pending_start_time == armed_at + assert req_status.hit_pending_expired is False + clock.advance(_SHIPPED_DEADLINE_S / 2) + + # The elapsed time accrued across passes rather than restarting. + sched.get_num_new_matched_tokens(request, 0) + assert req_status.hit_pending_expired is True + assert _expiry_count(sched) == 1 + + def test_disarm_clears_the_tracked_key(self, monkeypatch): + """Disarming drops the key with the start time, never just one. + + A key left behind with no start time would read as a *different* + blocker on the next pass and silently re-arm instead of arming. + """ + _install_fake_clock(monkeypatch, sched_module) + results = {1: LookupResult.HIT_PENDING} + sched = _make_deadline_scheduler(results) + request = _make_deadline_request() + sched.on_new_request(request) + req_status = sched._req_status[request.request_id] + + sched.get_num_new_matched_tokens(request, 0) + assert req_status.hit_pending_start_time is not None + assert req_status.hit_pending_key is not None + + # Nothing is pending any more, so both halves of the timer go. + results[1] = LookupResult.MISS + sched.get_num_new_matched_tokens(request, 0) + assert req_status.hit_pending_start_time is None + assert req_status.hit_pending_key is None + + def test_retry_deferral_never_expires(self, monkeypatch): + """RETRY is out of scope: even a long-armed timer cannot fire on it. + + The RETRY/backoff path is owned separately; this deadline must leave + it exactly as it was. + """ + clock = _install_fake_clock(monkeypatch, sched_module) + sched = _make_deadline_scheduler({1: LookupResult.RETRY}) + request = _make_deadline_request() + sched.on_new_request(request) + req_status = sched._req_status[request.request_id] + + # Pre-arm the timer, then let far more than the deadline elapse on + # RETRY-only passes. + req_status.hit_pending_start_time = clock() + for _ in range(5): + clock.advance(_SHIPPED_DEADLINE_S) + assert sched.get_num_new_matched_tokens(request, 0) == (None, False) + + assert req_status.hit_pending_expired is False + assert _expiry_count(sched) == 0 + + def test_resolved_lookup_disarms(self, monkeypatch): + """A lookup that resolves clears the timer, so runs must be unbroken.""" + _install_fake_clock(monkeypatch, sched_module) + results = {1: LookupResult.HIT_PENDING} + sched = _make_deadline_scheduler(results) + request = _make_deadline_request() + sched.on_new_request(request) + req_status = sched._req_status[request.request_id] + + sched.get_num_new_matched_tokens(request, 0) + assert req_status.hit_pending_start_time is not None + + results[1] = LookupResult.MISS + assert sched.get_num_new_matched_tokens(request, 0) == (0, False) + assert req_status.hit_pending_start_time is None + + def test_in_flight_transfers_do_not_accrue_deadline(self, monkeypatch): + """Waiting on this request's own transfers is not HIT_PENDING waiting.""" + _install_fake_clock(monkeypatch, sched_module) + sched = _make_deadline_scheduler({1: LookupResult.HIT_PENDING}) + request = _make_deadline_request() + sched.on_new_request(request) + req_status = sched._req_status[request.request_id] + + sched.get_num_new_matched_tokens(request, 0) + assert req_status.hit_pending_start_time is not None + + req_status.transfer_jobs.add(1) + assert sched.get_num_new_matched_tokens(request, 0) == (None, False) + assert req_status.hit_pending_start_time is None + + def test_zero_deadline_disables(self, monkeypatch): + """0 restores the pre-fix behavior: defer on HIT_PENDING forever.""" + clock = _install_fake_clock(monkeypatch, sched_module) + sched = _make_deadline_scheduler( + {1: LookupResult.HIT_PENDING}, hit_pending_deadline_s=0.0 + ) + request = _make_deadline_request() + sched.on_new_request(request) + req_status = sched._req_status[request.request_id] + + for _ in range(5): + clock.advance(_SHIPPED_DEADLINE_S) + assert sched.get_num_new_matched_tokens(request, 0) == (None, False) + + assert req_status.hit_pending_start_time is None + assert req_status.hit_pending_expired is False + assert _expiry_count(sched) == 0 + + def test_reset_cache_clears_expiry(self, monkeypatch): + """A cache wipe removes the pending writes, so the expiry must go too.""" + clock = _install_fake_clock(monkeypatch, sched_module) + sched = _make_deadline_scheduler({1: LookupResult.HIT_PENDING}) + request = _make_deadline_request() + sched.on_new_request(request) + req_status = sched._req_status[request.request_id] + + sched.get_num_new_matched_tokens(request, 0) + clock.advance(_SHIPPED_DEADLINE_S + 1.0) + sched.get_num_new_matched_tokens(request, 0) + assert req_status.hit_pending_expired is True + + sched._current_batch_load_jobs = [] + sched._current_batch_jobs_to_flush = set() + sched._current_batch_allocated_block_ids = set() + sched._job_counter = 0 + sched._jobs = {} + sched._block_id_to_pending_jobs = {} + sched.reset_cache() + + assert req_status.hit_pending_expired is False + assert req_status.hit_pending_start_time is None + class TestSlidingWindowLookup: def test_all_hit_exact_window(self): sched = _make_scheduler_with_lookup({1: LookupResult.HIT, 2: LookupResult.HIT}) - assert sched._sliding_window_lookup(to_keys([1, 2]), 2, _EMPTY_REQ_CTX) == 2 + assert _sliding_lookup(sched, to_keys([1, 2]), 2) == 2 def test_all_miss(self): sched = _make_scheduler_with_lookup({}) - assert sched._sliding_window_lookup(to_keys([1, 2, 3]), 1, _EMPTY_REQ_CTX) == 0 + assert _sliding_lookup(sched, to_keys([1, 2, 3]), 1) == 0 def test_window_at_end(self): sched = _make_scheduler_with_lookup({2: LookupResult.HIT, 3: LookupResult.HIT}) - assert sched._sliding_window_lookup(to_keys([1, 2, 3]), 2, _EMPTY_REQ_CTX) == 3 + assert _sliding_lookup(sched, to_keys([1, 2, 3]), 2) == 3 def test_window_in_middle(self): sched = _make_scheduler_with_lookup({2: LookupResult.HIT, 3: LookupResult.HIT}) - assert ( - sched._sliding_window_lookup(to_keys([1, 2, 3, 4]), 2, _EMPTY_REQ_CTX) == 3 - ) + assert _sliding_lookup(sched, to_keys([1, 2, 3, 4]), 2) == 3 def test_no_full_window_falls_back_to_prefix(self): sched = _make_scheduler_with_lookup({1: LookupResult.HIT, 2: LookupResult.HIT}) - assert sched._sliding_window_lookup(to_keys([1, 2, 3]), 3, _EMPTY_REQ_CTX) == 2 + assert _sliding_lookup(sched, to_keys([1, 2, 3]), 3) == 2 def test_single_block_window(self): sched = _make_scheduler_with_lookup({2: LookupResult.HIT, 3: LookupResult.HIT}) - assert sched._sliding_window_lookup(to_keys([1, 2, 3]), 1, _EMPTY_REQ_CTX) == 3 + assert _sliding_lookup(sched, to_keys([1, 2, 3]), 1) == 3 def test_gap_resets_consecutive(self): sched = _make_scheduler_with_lookup( {2: LookupResult.HIT, 3: LookupResult.HIT, 4: LookupResult.HIT} ) # [1, 2, 3, 0, 4] — gap at 0 resets, window of 2 found at [2,3] - assert ( - sched._sliding_window_lookup(to_keys([1, 2, 3, 0, 4]), 2, _EMPTY_REQ_CTX) - == 3 - ) + assert _sliding_lookup(sched, to_keys([1, 2, 3, 0, 4]), 2) == 3 def test_window_prefers_rightmost(self): sched = _make_scheduler_with_lookup( @@ -1261,10 +1770,7 @@ def test_window_prefers_rightmost(self): ) # two valid windows: [1,2] at positions 0-1 and [4,5] at positions 3-4 # scans right-to-left, finds [4,5] first - assert ( - sched._sliding_window_lookup(to_keys([1, 2, 3, 4, 5]), 2, _EMPTY_REQ_CTX) - == 5 - ) + assert _sliding_lookup(sched, to_keys([1, 2, 3, 4, 5]), 2) == 5 def test_prefix_fallback_with_gap(self): sched = _make_scheduler_with_lookup( @@ -1276,20 +1782,17 @@ def test_prefix_fallback_with_gap(self): } ) # window of 4 not found contiguously (gap at 1) - assert ( - sched._sliding_window_lookup(to_keys([2, 1, 3, 4, 5]), 4, _EMPTY_REQ_CTX) - == 1 - ) + assert _sliding_lookup(sched, to_keys([2, 1, 3, 4, 5]), 4) == 1 def test_empty(self): sched = _make_scheduler_with_lookup({}) - assert sched._sliding_window_lookup([], 1, _EMPTY_REQ_CTX) == 0 + assert _sliding_lookup(sched, [], 1) == 0 def test_retry_defers(self): sched = _make_scheduler_with_lookup( {1: LookupResult.HIT, 2: LookupResult.RETRY} ) - assert sched._sliding_window_lookup(to_keys([1, 2]), 2, _EMPTY_REQ_CTX) is None + assert _sliding_lookup(sched, to_keys([1, 2]), 2) is None def test_retry_with_full_window_still_defers(self): """Even if a real window is found after a RETRY, result is deferred.""" @@ -1303,10 +1806,7 @@ def test_retry_with_full_window_still_defers(self): 4: LookupResult.HIT, } ) - assert ( - sched._sliding_window_lookup(to_keys([1, 2, 3, 4]), 2, _EMPTY_REQ_CTX) - is None - ) + assert _sliding_lookup(sched, to_keys([1, 2, 3, 4]), 2) is None def test_hit_pending_counts_as_hit(self): """HIT_PENDING counts toward the consecutive-hit streak.""" @@ -1314,7 +1814,7 @@ def test_hit_pending_counts_as_hit(self): {1: LookupResult.HIT, 2: LookupResult.HIT_PENDING} ) # window=2: both count as hits, but defer_lookup is set - assert sched._sliding_window_lookup(to_keys([1, 2]), 2, _EMPTY_REQ_CTX) is None + assert _sliding_lookup(sched, to_keys([1, 2]), 2) is None def test_hit_pending_does_not_break_streak(self): """HIT_PENDING in the middle of a window doesn't reset the streak.""" @@ -1322,9 +1822,32 @@ def test_hit_pending_does_not_break_streak(self): {1: LookupResult.HIT, 2: LookupResult.HIT_PENDING, 3: LookupResult.HIT} ) # window=3: right-to-left finds 3(HIT),2(HIT_PENDING),1(HIT) = 3 consecutive - assert ( - sched._sliding_window_lookup(to_keys([1, 2, 3]), 3, _EMPTY_REQ_CTX) is None + assert _sliding_lookup(sched, to_keys([1, 2, 3]), 3) is None + + def test_hit_pending_downgrade_resets_streak(self): + """Once downgraded, HIT_PENDING behaves exactly like MISS. + + The returned window must never span the pending block — prepare_load + asserts block.is_ready, so counting it would trade the hang for a crash. + """ + sched = _make_scheduler_with_lookup( + {1: LookupResult.HIT, 2: LookupResult.HIT_PENDING, 3: LookupResult.HIT} ) + # window=2: right-to-left 3(HIT), 2(downgraded to miss) resets the + # streak, 1(HIT) — no window of 2 exists, so this resolves to 1 + # instead of deferring. + assert _sliding_lookup(sched, to_keys([1, 2, 3]), 2, True) == 1 + + def test_hit_pending_downgrade_allows_window_clear_of_pending(self): + """A window entirely past the pending block still resolves.""" + sched = _make_scheduler_with_lookup( + { + 1: LookupResult.HIT_PENDING, + 2: LookupResult.HIT, + 3: LookupResult.HIT, + } + ) + assert _sliding_lookup(sched, to_keys([1, 2, 3]), 2, True) == 3 @pytest.mark.parametrize("async_scheduling", [True, False]) @@ -1362,7 +1885,7 @@ def test_request_level_policy_stores_all_blocks(request_runner, async_scheduling # New request with 2 offloaded chunks; first matches what's in CPU. runner.new_request(token_ids=[0] * tokens_per_chunk * 2) - runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 1 + runner.connector_scheduler._maximal_prefix_lookup = _stub_prefix_scan(1) runner.manager.prepare_store.side_effect = lambda keys, req_context: ( generate_store_output(keys) ) @@ -1393,7 +1916,7 @@ def test_loads_do_not_populate_fence_index(request_runner): async_scheduling=False, ) runner.new_request(token_ids=[0] * 12) - runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 1 + runner.connector_scheduler._maximal_prefix_lookup = _stub_prefix_scan(1) runner.run(decoded_tokens=[], complete_transfers=False) assert runner.connector_scheduler._block_id_to_pending_jobs == {} @@ -1439,7 +1962,7 @@ def capture_fence(): runner.scheduler.reset_prefix_cache() runner.new_request(token_ids=[0] * 4) - runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 1 + runner.connector_scheduler._maximal_prefix_lookup = _stub_prefix_scan(1) runner.manager.prepare_store.side_effect = lambda keys, req_context: ( generate_store_output([]) ) @@ -1490,7 +2013,7 @@ def capture_fence(): runner.scheduler.reset_prefix_cache() runner.new_request(token_ids=[1] * 4) - runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 0 + runner.connector_scheduler._maximal_prefix_lookup = _stub_prefix_scan(0) runner.manager.prepare_store.side_effect = lambda keys, req_context: ( generate_store_output([]) ) @@ -1761,7 +2284,7 @@ def test_reset_cache(request_runner, async_scheduling: bool): # Leave the load in-flight so that reset_cache must flush it. runner.scheduler.reset_prefix_cache() runner.new_request(token_ids=[0] * tokens_per_chunk) - runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 1 + runner.connector_scheduler._maximal_prefix_lookup = _stub_prefix_scan(1) runner.manager.prepare_store.side_effect = lambda keys, req_context: ( generate_store_output([]) ) @@ -1948,7 +2471,9 @@ def test_async_preempt_readmit_before_transfer_output_is_deferred(request_runner # preemption batch's ModelRunnerOutput is consumed by update_from_output(). free_block_queue.num_free_blocks = num_free_blocks_empty assert runner.scheduler.reset_prefix_cache() - runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: len(keys) + runner.connector_scheduler._maximal_prefix_lookup = lambda keys, *_: _LookupScan( + len(keys), None + ) readmit_output = runner.scheduler.schedule() @@ -2062,7 +2587,7 @@ def test_swa_alignment_skip(request_runner, async_scheduling: bool): runner.scheduler.reset_prefix_cache() runner.new_request(token_ids=[0] * num_tokens + [1]) runner.manager.lookup.return_value = LookupResult.HIT - runner.connector_scheduler._maximal_prefix_lookup = lambda keys, ctx, *_: 2 + runner.connector_scheduler._maximal_prefix_lookup = _stub_prefix_scan(2) runner.run( decoded_tokens=[EOS_TOKEN_ID], # Group 0: full prefix lookup hits 2 offloaded chunks @@ -2399,12 +2924,12 @@ def test_sw_lookup_inflates_query_max(self, request_runner): captured_keys: list = [] orig_sw_lookup = type(sched)._sliding_window_lookup - def capturing_sw_lookup(self_arg, keys, window, req_context): + def capturing_sw_lookup(self_arg, keys, window, req_context, *rest): captured_keys.append(list(keys)) - return orig_sw_lookup(self_arg, keys, window, req_context) + return orig_sw_lookup(self_arg, keys, window, req_context, *rest) - sched._sliding_window_lookup = lambda keys, window, req_ctx: ( - capturing_sw_lookup(sched, keys, window, req_ctx) + sched._sliding_window_lookup = lambda keys, window, req_ctx, *rest: ( + capturing_sw_lookup(sched, keys, window, req_ctx, *rest) ) req_status = self._make_req_status( diff --git a/tests/v1/kv_offload/test_spec_config.py b/tests/v1/kv_offload/test_spec_config.py new file mode 100644 index 000000000000..c4f30fb0966c --- /dev/null +++ b/tests/v1/kv_offload/test_spec_config.py @@ -0,0 +1,82 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for tunables parsed by the shared OffloadingSpec base class.""" + +from collections.abc import Mapping +from typing import Any + +import pytest + +from vllm.v1.kv_offload.config import ( + OffloadingCacheConfig, + OffloadingConfig, + OffloadingGroupConfig, + OffloadingModelConfig, + OffloadingParallelConfig, +) +from vllm.v1.kv_offload.cpu.spec import CPUOffloadingSpec +from vllm.v1.kv_offload.tiering.spec import TieringOffloadingSpec + +# Both concrete specs must pick these up: TieringOffloadingSpec derives from +# CPUOffloadingSpec, which derives from OffloadingSpec. +SPEC_CLASSES = [CPUOffloadingSpec, TieringOffloadingSpec] + + +def make_config(extra_config: Mapping[str, Any]) -> OffloadingConfig: + return OffloadingConfig( + groups=(OffloadingGroupConfig(tokens_per_block=16, layer_names=("layer0",)),), + worker_kv_bytes_per_block=1024, + enable_kv_cache_events=False, + extra_config={"cpu_bytes_to_use": 1 << 20, **extra_config}, + engine_id="test-engine", + model=OffloadingModelConfig(name="test-model", dtype="float16"), + cache=OffloadingCacheConfig(tokens_per_hash=16, blocks_per_chunk=1), + parallel=OffloadingParallelConfig( + rank=0, + world_size=1, + tp_size=1, + pp_size=1, + pcp_size=1, + dcp_size=1, + data_parallel_index=0, + is_parallelism_agnostic=True, + ), + ) + + +@pytest.mark.parametrize("spec_cls", SPEC_CLASSES) +def test_hit_pending_deadline_default(spec_cls): + """The default clears the P2P load+abort-ack ceiling (30s + 10s).""" + assert spec_cls(make_config({})).hit_pending_deadline_s == 60.0 + + +@pytest.mark.parametrize("spec_cls", SPEC_CLASSES) +def test_hit_pending_deadline_override(spec_cls): + spec = spec_cls(make_config({"hit_pending_deadline_s": 12.5})) + assert spec.hit_pending_deadline_s == 12.5 + + +@pytest.mark.parametrize("spec_cls", SPEC_CLASSES) +def test_hit_pending_deadline_zero_allowed(spec_cls): + """0 is the documented opt-out, not an error.""" + assert ( + spec_cls(make_config({"hit_pending_deadline_s": 0})).hit_pending_deadline_s == 0 + ) + + +@pytest.mark.parametrize("spec_cls", SPEC_CLASSES) +def test_hit_pending_deadline_rejects_negative(spec_cls): + with pytest.raises(ValueError, match="hit_pending_deadline_s"): + spec_cls(make_config({"hit_pending_deadline_s": -1})) + + +@pytest.mark.parametrize("spec_cls", SPEC_CLASSES) +def test_hit_pending_deadline_rejects_nan(spec_cls): + """NaN must not slip past validation. + + `nan < 0` is False, so a bare `< 0` guard admits it; every later + `now - start < nan` is then also False, expiring the request on its + second deferred pass instead of honouring any deadline. + """ + with pytest.raises(ValueError, match="hit_pending_deadline_s"): + spec_cls(make_config({"hit_pending_deadline_s": float("nan")})) diff --git a/tests/v1/kv_offload/tiering/test_tiering_offloading.py b/tests/v1/kv_offload/tiering/test_tiering_offloading.py index b19a270efb46..868a5493dd4a 100644 --- a/tests/v1/kv_offload/tiering/test_tiering_offloading.py +++ b/tests/v1/kv_offload/tiering/test_tiering_offloading.py @@ -639,6 +639,38 @@ def test_lookup_shared_block_no_duplicate_promotion(self, manager_setup): assert list(job_metadata.keys) == [shared_block] assert job_metadata.req_context is ctx_a + def test_hit_pending_is_stable_and_skips_secondary_tiers(self, manager_setup): + """A primary-pending key stays HIT_PENDING and never re-enters promotion. + + The HIT_PENDING deadline for issue #49829 lives in the connector + scheduler, above this manager, so that one fix covers every spec. + This pins the two manager-side properties that design relies on: + repeated lookups keep reporting HIT_PENDING (the manager has no + deadline of its own to trip over), and a pending key short-circuits + before the secondary tiers, so downgrading it upstream can never + collide with the in-flight write by starting a second promotion. + """ + shared_block = to_keys([0])[0] + self.secondary_tier1.blocks[shared_block] = True + + ctx_a = ReqContext(req_id="req_a") + ctx_b = ReqContext(req_id="req_b") + + # req_a initiates the promotion, leaving the primary slot write-pending. + assert self.manager.lookup(shared_block, ctx_a) is LookupResult.RETRY + + self.secondary_tier1.lookup = MagicMock(wraps=self.secondary_tier1.lookup) + self.secondary_tier2.lookup = MagicMock(wraps=self.secondary_tier2.lookup) + + # req_b hot-loops on the same key while that write is in flight, as + # the scheduler's per-step scan does. + results = {self.manager.lookup(shared_block, ctx_b) for _ in range(100)} + assert results == {LookupResult.HIT_PENDING} + + # The primary short-circuit means no secondary tier is ever consulted. + self.secondary_tier1.lookup.assert_not_called() + self.secondary_tier2.lookup.assert_not_called() + def test_complete_store_forwards_req_context_to_submit_store(self, manager_setup): """complete_store cascades to secondary tiers with the correct req_context.""" blocks = to_keys(range(2)) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/metrics.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/metrics.py index 58023c393561..26ad4f70f53d 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/metrics.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/metrics.py @@ -36,6 +36,7 @@ class _ConnectorMetricName: LOOKUP_SYNC_DELAY = "vllm:kv_offload_lookup_sync_delay_seconds" LOOKUP_ASYNC_DELAY = "vllm:kv_offload_lookup_async_delay_seconds" ALLOCATION_FAILURE = "vllm:kv_offload_allocation_failure" + HIT_PENDING_DEADLINE_EXPIRED = "vllm:kv_offload_hit_pending_deadline_expired" class _TransferType: @@ -126,6 +127,17 @@ def get_connector_metric_definitions() -> dict[str, OffloadingMetricMetadata]: "Number of KV offload store allocation attempts that failed." ), ), + _ConnectorMetricName.HIT_PENDING_DEADLINE_EXPIRED: OffloadingCounterMetadata( + documentation=( + "Number of requests whose HIT_PENDING deadline expired, causing " + "pending blocks to be treated as misses for the rest of the " + "request. A non-zero value means a write stayed pending past " + "the deadline, which is not proof of a leak: only the P2P " + "secondary tier bounds transfer time, so on the CPU, " + "filesystem and object-store backends a healthy but slow " + "write can cross the deadline too." + ), + ), } diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py index fecce3d28fec..c58583dcddb4 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -38,6 +38,7 @@ SlidingWindowSpec, ) from vllm.v1.kv_offload.base import ( + DEFAULT_HIT_PENDING_DEADLINE_S, GPULoadStoreSpec, Locality, LookupResult, @@ -164,6 +165,10 @@ class SchedulerOffloadConfig(NamedTuple): blocks_per_chunk: int num_workers: int offload_prompt_only: bool + # Seconds a request may stay deferred on HIT_PENDING before pending blocks + # are treated as misses for the rest of that request. 0 disables the + # deadline. See OffloadingSpec.__init__ for how the default is derived. + hit_pending_deadline_s: float = DEFAULT_HIT_PENDING_DEADLINE_S @classmethod def from_spec( @@ -255,9 +260,23 @@ def _alignment_chunk_count( ), blocks_per_chunk=spec.blocks_per_chunk, offload_prompt_only=spec.offload_prompt_only, + hit_pending_deadline_s=spec.hit_pending_deadline_s, ) +class _LookupScan(NamedTuple): + """Result of scanning one KV group's keys through the offloading manager.""" + + # Number of hit chunks (prefix scan) or the window end index (sliding + # window scan); None if the backend deferred this scan. + num_hit_chunks: int | None + # The first key this scan saw return a live (non-downgraded) HIT_PENDING, + # in scan order; None if it saw none. Reported out of the scan itself, so + # the deadline can be keyed on the blocking key without a second + # manager.lookup() sweep over the same keys. + first_pending_key: OffloadKey | None + + @dataclass class RequestGroupState: offload_keys: list[OffloadKey] = field(default_factory=list) @@ -288,6 +307,24 @@ class RequestOffloadState: deferred_lookup_start_time: float | None = None # True once on_request_finished has been signaled to the manager. finished_signaled: bool = False + # time.monotonic() at which the timer below was armed for + # hit_pending_key; None while not armed. Distinct from + # deferred_lookup_start_time, which is armed by any deferral (including + # RETRY) and only feeds the LOOKUP_ASYNC_DELAY histogram. + hit_pending_start_time: float | None = None + # The key hit_pending_start_time is counting for. A later pass reporting a + # different key means a different stalled write, which gets its own full + # deadline, so the timer re-arms rather than letting it inherit the + # elapsed time already accrued. + hit_pending_key: OffloadKey | None = None + # Set once the deadline fires, and sticky for the rest of the request: + # subsequent lookups downgrade HIT_PENDING to a miss rather than + # re-arming the timer. + hit_pending_expired: bool = False + # Written by _lookup(): the live (non-downgraded) HIT_PENDING key the pass + # just performed observed, or None if it observed none. Drives arming, + # re-arming and disarming. + observed_pending_key: OffloadKey | None = None def __post_init__(self) -> None: self.group_states = tuple( @@ -519,6 +556,71 @@ def _maybe_observe_lookup_async_delay( time.monotonic() - start_time, ) + def _update_hit_pending_deadline( + self, req_status: RequestOffloadState, now: float + ) -> None: + """Arm, disarm, or expire the HIT_PENDING deadline of a deferred request. + + Only deferrals caused by HIT_PENDING are bounded. The timer disarms on + a pass where *no* group observed a live HIT_PENDING; if any group did, + it stays armed even when another group deferred for an unrelated + reason (RETRY, or hit chunks already being loaded), since the stalled + write that group is waiting on is still the thing being bounded. + Without the disarm, a request that saw one transient HIT_PENDING and + then deferred only for unrelated reasons would keep a stale armed + timer and expire with no stalled write anywhere. + + The timer is keyed on *which* block is blocking, not merely on the + fact that one is: without the identity check, a key that turns pending + while an older one is still unresolved would inherit the older key's + elapsed time and expire despite a perfectly healthy write. The cost is + that the bound is per-blocker, so a request that keeps acquiring new + blockers can wait longer than one deadline in total. That is the + conservative direction — each change of identity means the previous + write landed — and a leaked write pins one key forever, so the case + this deadline exists for still expires on time. + """ + observed_key = req_status.observed_pending_key + if observed_key is None: + # Start time and key are one piece of state, cleared together. + req_status.hit_pending_start_time = None + req_status.hit_pending_key = None + return + + deadline = self.config.hit_pending_deadline_s + if deadline <= 0 or req_status.hit_pending_expired: + return + + start_time = req_status.hit_pending_start_time + if start_time is None or observed_key != req_status.hit_pending_key: + # Not armed yet, or a different block is blocking now. Either way + # this key starts its own deadline from scratch. + req_status.hit_pending_start_time = now + req_status.hit_pending_key = observed_key + return + + if now - start_time < deadline: + return + + # The write backing this block is not landing. Stop counting pending + # blocks as hits for this request so it recomputes locally instead of + # deferring until the client gives up. Sticky: later lookups must not + # re-arm the timer. + req_status.hit_pending_expired = True + req_status.hit_pending_start_time = None + req_status.hit_pending_key = None + self._connector_stats.increase_counter( + _ConnectorMetricName.HIT_PENDING_DEADLINE_EXPIRED + ) + logger.warning( + "Request %s was deferred on HIT_PENDING for %.1fs (deadline %.1fs); " + "treating pending blocks as misses for the rest of this request. " + "An offload write is likely stalled or leaked.", + req_status.req.request_id, + now - start_time, + deadline, + ) + def _generate_job_id(self) -> int: job_id = self._job_counter self._job_counter += 1 @@ -549,11 +651,22 @@ def _maximal_prefix_lookup( req: Request, group_config: GroupOffloadConfig, start_chunk_idx: int, - ) -> int | None: + downgrade_hit_pending: bool = False, + ) -> _LookupScan: """Return the number of consecutive offloaded chunks from the start, - or None if the backend deferred a lookup.""" + or None if the backend deferred a lookup. + + When `downgrade_hit_pending` is set, the request's HIT_PENDING deadline + has expired: pending blocks are treated as misses so the scan can + resolve instead of deferring again. + + Also reports the first key seen pending. This scan runs forward, so + that is the lowest-index pending key: the block that bounds how far + the prefix can resolve, and so the one the deadline should time. + """ hit_count = 0 defer_lookup = False + first_pending_key: OffloadKey | None = None for local_idx, key in enumerate(keys): result = self.manager.lookup(key, req_context) match result: @@ -566,6 +679,14 @@ def _maximal_prefix_lookup( ) hit_count += 1 case LookupResult.HIT_PENDING: + if downgrade_hit_pending: + # Deadline expired: treat as a miss and truncate the + # prefix here. Must break *before* incrementing — + # counting a pending block would later hand it to + # prepare_load, which asserts block.is_ready. + break + if first_pending_key is None: + first_pending_key = key defer_lookup = True hit_count += 1 case LookupResult.RETRY: @@ -574,29 +695,52 @@ def _maximal_prefix_lookup( defer_lookup = True case LookupResult.MISS: break - return hit_count if not defer_lookup else None + return _LookupScan( + hit_count if not defer_lookup else None, + first_pending_key, + ) def _sliding_window_lookup( self, keys: Sequence[OffloadKey], sliding_window_size: int, req_context: ReqContext, - ) -> int | None: + downgrade_hit_pending: bool = False, + ) -> _LookupScan: """Return the end index (in `keys`) of the last run of `sliding_window_size` consecutive hits, scanning from the end. - Returns 0 on miss, None if the backend deferred a lookup.""" + Returns 0 on miss, None if the backend deferred a lookup. + + When `downgrade_hit_pending` is set, the request's HIT_PENDING deadline + has expired: pending blocks are treated as misses so the scan can + resolve instead of deferring again. + + Also reports the first key seen pending. This scan runs backwards, so + unlike the prefix case that is the *highest*-index pending key. Either + end works: the choice only has to be deterministic, so an unchanged + blocker keeps accruing rather than re-arming every pass. + """ defer_lookup = False + first_pending_key: OffloadKey | None = None consecutive_hits = 0 for idx in range(len(keys) - 1, -1, -1): match self.manager.lookup(keys[idx], req_context): case LookupResult.HIT: consecutive_hits += 1 case LookupResult.HIT_PENDING: - # Block is in cache, just not readable yet — counts - # as hit for the consecutive streak. Don't break: - # keep scanning to let manager kick off async lookups. - defer_lookup = True - consecutive_hits += 1 + if downgrade_hit_pending: + # Deadline expired: treat exactly as a miss, so no + # returned window can span a block that prepare_load + # would reject for not being ready. + consecutive_hits = 0 + else: + # Block is in cache, just not readable yet — counts + # as hit for the consecutive streak. Don't break: + # keep scanning to let manager kick off async lookups. + defer_lookup = True + if first_pending_key is None: + first_pending_key = keys[idx] + consecutive_hits += 1 case LookupResult.RETRY: # Block location uncertain — does not count as hit. # Don't break: keep scanning to let manager kick off @@ -606,8 +750,14 @@ def _sliding_window_lookup( case LookupResult.MISS: consecutive_hits = 0 if consecutive_hits == sliding_window_size: - return idx + sliding_window_size if not defer_lookup else None - return consecutive_hits if not defer_lookup else None + return _LookupScan( + idx + sliding_window_size if not defer_lookup else None, + first_pending_key, + ) + return _LookupScan( + consecutive_hits if not defer_lookup else None, + first_pending_key, + ) def _touch(self, req_status: RequestOffloadState): for group_config, group_state in zip( @@ -654,6 +804,11 @@ def _lookup(self, req_status: RequestOffloadState) -> int | None: defer_lookup = False lookup_groups = self._lookup_groups + # Re-observed from scratch on every pass; the caller uses it to arm, + # re-arm or disarm the HIT_PENDING deadline. + req_status.observed_pending_key = None + downgrade_hit_pending = req_status.hit_pending_expired + # Tracks which eagle groups have already popped their volatile trailing chunk # in the current convergence iteration. Reset when a non-eagle group # tightens the hit boundary, requiring a fresh pop. @@ -707,22 +862,34 @@ def _lookup(self, req_status: RequestOffloadState) -> int | None: # have backend-confirmed hits num_hit_chunks: int | None if sliding_window_size_in_chunks is None: - num_hit_chunks = self._maximal_prefix_lookup( + scan = self._maximal_prefix_lookup( offload_keys, req_status.req_context, req_status.req, group_config, start_chunk_idx, + downgrade_hit_pending, ) else: required_window = sliding_window_size_in_chunks if is_eagle_unverified: required_window += 1 - num_hit_chunks = self._sliding_window_lookup( + scan = self._sliding_window_lookup( offload_keys, required_window, req_status.req_context, + downgrade_hit_pending, ) + num_hit_chunks = scan.num_hit_chunks + if ( + scan.first_pending_key is not None + and req_status.observed_pending_key is None + ): + # First blocker reported this pass wins. Groups are walked + # in fixed self._lookup_groups order, so an unchanged set + # of stalled writes picks the same key every pass and the + # timer accrues instead of resetting. + req_status.observed_pending_key = scan.first_pending_key if num_hit_chunks == 0: return 0 @@ -844,6 +1011,10 @@ def get_num_new_matched_tokens( "Delaying request %s since it still has in-flight transfers", request.request_id, ) + # No lookup runs on this path, so no HIT_PENDING is observed and + # the deadline must not keep accruing. + req_status.hit_pending_start_time = None + req_status.hit_pending_key = None return None, False req_status.update_offload_keys() @@ -852,6 +1023,10 @@ def get_num_new_matched_tokens( num_hit_tokens: int | None if request.skip_reading_prefix_cache: num_hit_tokens = 0 + # The prefix cache is not consulted at all, so nothing can be + # observed as HIT_PENDING on this path. + req_status.hit_pending_start_time = None + req_status.hit_pending_key = None else: lookup_start = time.monotonic() num_hit_tokens = self._lookup(req_status) @@ -862,8 +1037,11 @@ def get_num_new_matched_tokens( if num_hit_tokens is None: if req_status.deferred_lookup_start_time is None: req_status.deferred_lookup_start_time = lookup_start + self._update_hit_pending_deadline(req_status, lookup_start) else: self._maybe_observe_lookup_async_delay(req_status) + req_status.hit_pending_start_time = None + req_status.hit_pending_key = None req_status.update_num_hit_chunks(num_computed_tokens + (num_hit_tokens or 0)) self._touch(req_status) @@ -1451,6 +1629,12 @@ def reset_cache(self) -> None: for group_state in status.group_states: group_state.next_stored_chunk_idx = 0 status.transfer_jobs.clear() + # The cache is gone, so the writes that were pending no longer + # exist. Carrying an expiry over would suppress hits that are now + # legitimately re-fetchable. + status.hit_pending_start_time = None + status.hit_pending_key = None + status.hit_pending_expired = False # Discard jobs and save job_counter to be able to discard worker responses self._stale_job_threshold = self._job_counter diff --git a/vllm/v1/kv_offload/base.py b/vllm/v1/kv_offload/base.py index b0db86cf9bff..51c7e6e3ead6 100644 --- a/vllm/v1/kv_offload/base.py +++ b/vllm/v1/kv_offload/base.py @@ -25,6 +25,11 @@ # Use the helper functions below to construct / decompose keys. OffloadKey = NewType("OffloadKey", bytes) +# Default for the `hit_pending_deadline_s` tunable; see OffloadingSpec.__init__ +# for how the value is derived. Shared with SchedulerOffloadConfig so the +# default is stated in exactly one place. +DEFAULT_HIT_PENDING_DEADLINE_S = 60.0 + def make_offload_key(block_hash: bytes, group_idx: int) -> OffloadKey: """Pack a block hash and group index into an `OffloadKey`.""" @@ -554,6 +559,47 @@ def __init__(self, config: OffloadingConfig): self.extra_config.get("offload_prompt_only", True) ) + # Maximum time, in seconds, that a request may stay deferred because a + # candidate block keeps looking up as HIT_PENDING (present in the + # primary tier, but with a write still in flight). Past this, the block + # is treated as a miss for the remainder of the request and the + # candidate prefix is truncated there, so the request recomputes + # locally instead of deferring until the client times out. + # + # The check cannot distinguish two cases: + # - The request is waiting on a promotion it initiated itself. That + # wait tracks a live transfer, so it is the only case a too-short + # deadline can harm, by discarding a hit about to land. How long it + # may legitimately run is backend-specific: the P2P secondary tier + # bounds it at _LOAD_TIMEOUT_S (30s) + _ABORT_ACK_TIMEOUT_S (10s) = + # 40s, after which the finished-job poll resolves the block either + # way. The CPU, filesystem and object-store backends have no + # equivalent ceiling, so there a healthy but slow write can cross + # any fixed deadline, costing a cache miss and local recomputation + # but never correctness. + # - The request touched a key left write-pending by an earlier, + # unrelated request whose write leaked. No live transfer exists and + # nothing bounds the wait. This is the case that hangs, and it is + # resolved at any deadline value. + # The default clears P2P's 40s ceiling with margin, since the second + # case is insensitive to the value. It is therefore derived from the + # one backend that publishes a bound; deployments on a slow filesystem + # or object store may need to raise it. Set 0 to disable the deadline + # and defer indefinitely on HIT_PENDING. + self.hit_pending_deadline_s: float = float( + self.extra_config.get( + "hit_pending_deadline_s", DEFAULT_HIT_PENDING_DEADLINE_S + ) + ) + # Negated comparison so NaN is rejected too: `nan < 0` is False, which + # would let NaN through and make every deadline comparison downstream + # False, expiring the request on its second deferred pass. + if not (self.hit_pending_deadline_s >= 0): + raise ValueError( + "'hit_pending_deadline_s' must be non-negative, got " + f"{self.hit_pending_deadline_s}" + ) + self.tokens_per_block = tuple(group.tokens_per_block for group in config.groups) self.tokens_per_hash = config.cache.tokens_per_hash self.blocks_per_chunk = config.cache.blocks_per_chunk