diff --git a/.lil/changes/lmcache-restore-admission.json b/.lil/changes/lmcache-restore-admission.json new file mode 100644 index 00000000000..2a460446ba8 --- /dev/null +++ b/.lil/changes/lmcache-restore-admission.json @@ -0,0 +1,21 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "lmcache-restore-admission", + "category": "fix", + "summary": "Recurrent checkpoint restores wait for GPU and copy capacity instead of admitting requests to recompute cached prefixes.", + "models": [ + "all" + ], + "compatibility": "Requires the paired vLLM restore-admission change; deploy vLLM before enabling this LMCache change.", + "details": [ + "Capacity waits retain the answered checkpoint lookup and do not consume the timeout for an unanswered directory request.", + "A local checkpoint evicted before admission triggers a fresh external lookup, while genuine misses and exhausted transfer retries retain recomputation fallback.", + "This change does not repair missing checkpoint pages in cache storage." + ], + "requires": [ + "vllm-restore-admission" + ], + "pull_requests": [ + 100 + ] +} diff --git a/lmcache/integration/vllm/checkpoint_scheduler.py b/lmcache/integration/vllm/checkpoint_scheduler.py index 1038455515d..fa05ca574c4 100644 --- a/lmcache/integration/vllm/checkpoint_scheduler.py +++ b/lmcache/integration/vllm/checkpoint_scheduler.py @@ -76,9 +76,11 @@ class _Lookup: done: bool = False task_id: str | None = None checkpoint_id: int | None = None + local_checkpoint_id: int | None = None attempts: int = 1 started: float = field(default_factory=time.monotonic) capacity_refusal_logged: bool = False + task_refusal_logged: bool = False class CheckpointSchedulerBridge: @@ -91,10 +93,10 @@ class CheckpointSchedulerBridge: Worker layout is added after all ranks report identical descriptors. world_size: Number of engine ranks contributing to one atomic generation. max_tasks: Admission limit for collective stores and restores. - lookup_timeout: Seconds a request may wait for directory replies, - including shorter-checkpoint retries, and a store for its begin - reply. A request whose lookup has not been answered by then is - admitted to recompute its prompt; an unanswered store is aborted. + lookup_timeout: Seconds allowed for each directory reply (including + retries) and for a store's begin reply. Capacity waiting does not + consume this deadline. An unanswered lookup permits prompt + recomputation; an unanswered store is aborted. A started worker copy is never abandoned. The scheduler must keep issuing connector-only steps while has_pending is @@ -175,9 +177,8 @@ def poll_prefix(self, request: "Request") -> bool: Returns: False while layout negotiation, lookup or H2D is pending. True when - ordinary GPU admission may proceed, including when an external - import cannot reserve capacity. An unadmitted request may poll - again to retry that reservation without another directory lookup. + ordinary GPU admission may proceed. Capacity and copy-task pressure + keep an available checkpoint waiting without another lookup. Reusing a finished request's ID waits for its admitted copies to drain, so their completion cannot cancel or erase another lookup. """ @@ -188,11 +189,15 @@ def poll_prefix(self, request: "Request") -> bool: ): return True local = self._cache.find(request, request.num_tokens) - if local is not None and local.num_tokens == request.num_tokens: + state = self._lookups.get(request.request_id) + if ( + state is None + and local is not None + and local.num_tokens == request.num_tokens + ): return True if self._layout is None: return False - state = self._lookups.get(request.request_id) if state is None: roots = self._roots(request) state = _Lookup( @@ -202,6 +207,27 @@ def poll_prefix(self, request: "Request") -> bool: self._lookups[request.request_id] = state return False if state.done: + # Local reuse can also disappear while ordinary admission waits. + selected_id = ( + state.local_checkpoint_id + if state.local_checkpoint_id is not None + else state.checkpoint_id + ) + if ( + selected_id is not None + and not ( + self._manager.external_boundary_admission_ready(request.request_id) + ) + and (local is None or local.checkpoint_id != selected_id) + ): + logger.info( + "Recurrent checkpoint invalidated before admission for " + "request %s; looking up its prefix again", + request.request_id, + ) + self._manager.release_external_boundary_admission(request.request_id) + del self._lookups[request.request_id] + return self.poll_prefix(request) return True if state.task_id is not None: return False @@ -227,21 +253,25 @@ def poll_prefix(self, request: "Request") -> bool: ) state.done = True return True - if manifest is None or ( - local is not None and manifest.prefix.num_tokens <= local.num_tokens - ): + if manifest is None: state.done = True return True - if len(self._tasks) >= self._max_tasks: - logger.info( - "Recurrent checkpoint restore of %d tokens skipped for request %s: " - "%d checkpoint copies in flight; recomputing its prompt", - manifest.prefix.num_tokens, - request.request_id, - len(self._tasks), - ) + if local is not None and manifest.prefix.num_tokens <= local.num_tokens: + # Keep local selection separate from external cache-hit accounting. + state.local_checkpoint_id = local.checkpoint_id state.done = True return True + if len(self._tasks) >= self._max_tasks: + if not state.task_refusal_logged: + state.task_refusal_logged = True + logger.info( + "Recurrent checkpoint restore of %d tokens waiting for " + "copy capacity for request %s (%d copies in flight)", + manifest.prefix.num_tokens, + request.request_id, + len(self._tasks), + ) + return False try: payload, positions = self._validate_manifest(manifest, state.roots) checkpoint = self._manager.reserve_external_boundary_checkpoint( @@ -251,6 +281,7 @@ def poll_prefix(self, request: "Request") -> bool: draft_prefix_len=payload["draft_prefix_len"], kind=payload["kind"], num_ranks=self._world_size, + reserve_admission=True, ) except (ValueError, KeyError, TypeError): logger.warning( @@ -263,21 +294,25 @@ def poll_prefix(self, request: "Request") -> bool: state.done = True return True if checkpoint is None: - # Insufficient GPU capacity is not an external-cache miss. Permit - # ordinary admission, but retain the manifest so an unadmitted - # request can retry its import after other owners release pages. + # Resource pressure is not a cache miss. Keep the answered manifest + # while runnable owners release capacity. if not state.capacity_refusal_logged: state.capacity_refusal_logged = True logger.info( "Recurrent checkpoint restore of %d tokens deferred for " - "request %s: not enough free GPU blocks (%d free); it is " - "admitted without the restore unless capacity returns first", + "request %s: waiting for restore and execution capacity " + "(%d free GPU blocks)", manifest.prefix.num_tokens, request.request_id, self._manager.block_pool.get_num_free_blocks(), ) - return True + return False task = self._make_task(manifest, checkpoint, "RETRIEVE") + logger.debug( + "Recurrent checkpoint restore of %d tokens copying for request %s", + manifest.prefix.num_tokens, + request.request_id, + ) self._tasks[task.task_id] = _PendingTask(task, checkpoint, request.request_id) state.task_id = task.task_id return False @@ -523,6 +558,22 @@ def complete(self, results: dict[str, dict[int, bool]]) -> None: ) if published and state is not None: state.checkpoint_id = pending.checkpoint.checkpoint_id + logger.debug( + "Recurrent checkpoint restore of %d tokens ready for " + "request %s; ownership retained until admission", + pending.checkpoint.num_tokens, + pending.request_id, + ) + elif state is not None: + retry = state.attempts < _MAX_LOOKUP_ATTEMPTS + logger.info( + "Recurrent checkpoint publication invalidated for " + "request %s%s", + pending.request_id, + "; retrying lookup" + if retry + else "; recomputing its prompt", + ) elapsed = time.monotonic() - pending.created if elapsed > _SLOW_RESTORE_SECONDS: logger.info( @@ -540,7 +591,6 @@ def complete(self, results: dict[str, dict[int, bool]]) -> None: state is not None and pending.request_id not in self._cancelled and state.attempts < _MAX_LOOKUP_ATTEMPTS - and time.monotonic() - state.started < self._lookup_timeout ) logger.info( "Recurrent checkpoint restore of %d tokens failed for " @@ -567,6 +617,7 @@ def complete(self, results: dict[str, dict[int, bool]]) -> None: ) state.task_id = None state.attempts += 1 + state.started = time.monotonic() state.future = self._client.submit_request( RequestType.CHECKPOINT_FIND, [state.roots.roots] ) @@ -578,6 +629,7 @@ def complete(self, results: dict[str, dict[int, bool]]) -> None: def finish_request(self, request_id: str) -> None: """Forget lookup state, retaining copy pins until every admitted task drains.""" + self._manager.release_external_boundary_admission(request_id) self._lookups.pop(request_id, None) if any(task.request_id == request_id for task in self._tasks.values()): self._cancelled.add(request_id) diff --git a/tests/v1/test_vllm_semantic_checkpoint_transfer.py b/tests/v1/test_vllm_semantic_checkpoint_transfer.py index 44c063222a5..ffa64844a84 100644 --- a/tests/v1/test_vllm_semantic_checkpoint_transfer.py +++ b/tests/v1/test_vllm_semantic_checkpoint_transfer.py @@ -26,6 +26,7 @@ from vllm.lora.request import LoRARequest # noqa: E402 from vllm.sampling_params import SamplingParams # noqa: E402 from vllm.utils.hashing import sha256 # noqa: E402 +from vllm.v1.core.boundary_checkpoint import BoundaryCheckpoint # noqa: E402 from vllm.v1.core.kv_cache_manager import KVCacheManager # noqa: E402 from vllm.v1.core.kv_cache_utils import ( # noqa: E402 get_request_block_hasher, @@ -485,6 +486,11 @@ def test_worker_metadata_requires_bound_rank() -> None: "rejected-store-reused-request-id", "inflight-store-reused-request-id", "capacity-retry", + "capacity-timeout", + "task-capacity", + "local-hit-during-copy", + "invalidated-before-use", + "publication-invalidated", ], ) def test_semantic_roundtrip_collective_visibility_and_cancellation( @@ -503,6 +509,8 @@ def test_semantic_roundtrip_collective_visibility_and_cancellation( "parallel": {"tp": 4, "dcp": 1}, }, 4, + max_tasks=1 if outcome == "task-capacity" else 32, + lookup_timeout=0.05 if outcome == "capacity-timeout" else 60.0, ) layout = {"schema_version": 1, "page_bytes": 128} bridge.accept_layouts({rank: layout for rank in range(4)}) @@ -591,30 +599,28 @@ def record(kind: RequestType, args: list[Any]) -> Any: assert roots and all(root.namespace == prefix.namespace for root in roots) assert sum(len(root.tail_tokens) for root in roots) >= prefix.num_tokens consumer = make_request("consumer") - if outcome == "capacity-retry": - # Ordinary admission may also fail while another request owns - # the pool. An available manifest must remain retryable until - # this consumer is admitted or cancelled. + if outcome in ("capacity-retry", "capacity-timeout"): + # Resource pressure must not authorize cold admission. pressure = manager.block_pool.get_new_blocks( manager.block_pool.get_num_free_blocks() ) try: - deadline = time.monotonic() + 5 - while not bridge.poll_prefix(consumer): - assert time.monotonic() < deadline + for _ in range(100): + assert not bridge.poll_prefix(consumer) time.sleep(0.001) assert bridge.take_tasks() == [] - # Persistent pressure still permits ordinary admission; - # retaining the manifest must not introduce a wait loop. - assert bridge.poll_prefix(consumer) + assert not bridge.poll_prefix(consumer) finally: manager.block_pool.free_blocks(pressure) inflight_store = None - if outcome == "inflight-store-reused-request-id": + if outcome in ("inflight-store-reused-request-id", "task-capacity"): # This cancelled producer has different tokens but the same # public ID as the waiting consumer. Its admitted copy is still # live while the consumer looks up the first producer's data. - predecessor = make_request("consumer", first_token=20) + predecessor = make_request( + "blocker" if outcome == "task-capacity" else "consumer", + first_token=20, + ) predecessor_checkpoint = manager.reserve_external_boundary_checkpoint( predecessor, 11, @@ -675,8 +681,28 @@ def drain_predecessor() -> None: assert len(tasks) == 1 restore_task = tasks[0] assert manager.get_computed_blocks(consumer)[1] == 0 + if outcome == "local-hit-during-copy": + local = make_request("local-producer") + local_checkpoint = manager.reserve_external_boundary_checkpoint( + local, + 11, + manager.boundary_checkpoint_page_positions(11), + draft_prefix_len=11, + kind="prompt", + num_ranks=1, + ) + assert local_checkpoint is not None + assert manager.acknowledge_external_boundary_checkpoint( + local_checkpoint.checkpoint_id, 0 + ) + assert not bridge.poll_prefix(consumer) if outcome == "cancelled": bridge.finish_request(consumer.request_id) + if outcome == "publication-invalidated": + assert manager.boundary_checkpoints is not None + manager.boundary_checkpoints.invalidate_block( + restore_task.block_ids[-1][0] + ) for rank in range(4): future = worker.submit( CheckpointTransferJob( @@ -695,12 +721,32 @@ def drain_predecessor() -> None: } ) if rank < 3: - assert manager.get_computed_blocks(consumer)[1] == 0 + if outcome != "local-hit-during-copy": + assert manager.get_computed_blocks(consumer)[1] == 0 assert manager.block_pool.get_num_free_blocks() < before - assert manager.block_pool.get_num_free_blocks() == before - expected_tokens = 0 if outcome in ("rank-miss", "cancelled") else 11 + expected_tokens = ( + 0 + if outcome in ("rank-miss", "cancelled", "publication-invalidated") + else 11 + ) + if expected_tokens: + assert manager.block_pool.get_num_free_blocks() < before + assert not manager.reset_prefix_cache() + else: + assert manager.block_pool.get_num_free_blocks() == before assert manager.get_computed_blocks(consumer)[1] == expected_tokens assert bridge.external_tokens(consumer) == expected_tokens + if outcome == "invalidated-before-use": + # Preemption releases request ownership while its connector + # lookup state can still remember the previous successful import. + checkpoint = consumer.boundary_checkpoint + assert checkpoint is not None + manager.free(consumer) + assert manager.boundary_checkpoints is not None + manager.boundary_checkpoints.invalidate(checkpoint.checkpoint_id) + assert not bridge.poll_prefix(consumer) + bridge.finish_request(consumer.request_id) + assert manager.block_pool.get_num_free_blocks() == before if inflight_store is not None: drain_predecessor() assert not bridge.has_pending @@ -846,7 +892,108 @@ def make_bridge(client, manager, lookup_timeout: float) -> CheckpointSchedulerBr return bridge -@pytest.mark.parametrize("reply", ["never", "error"]) +@pytest.mark.parametrize("prefix, already_local", [(11, False), (8, True)]) +def test_local_reuse_retries_external_restore_after_capacity_eviction( + prefix: int, already_local: bool +) -> None: + """Losing a local selection must not turn an external hit into cold prefill.""" + manager = make_manager() + manifest = None + lookups = 0 + + def submit(kind: RequestType, args: list[Any]) -> SimpleNamespace: + nonlocal manifest, lookups + if kind == RequestType.CHECKPOINT_BEGIN: + manifest = args[0] + if kind == RequestType.CHECKPOINT_FIND: + lookups += 1 + result = manifest if kind == RequestType.CHECKPOINT_FIND else True + return SimpleNamespace(query=lambda: True, result=lambda: result) + + bridge = make_bridge(SimpleNamespace(submit_request=submit), manager, 60) + producer = make_request("producer") + + def publish() -> BoundaryCheckpoint: + checkpoint = manager.reserve_external_boundary_checkpoint( + producer, + prefix, + manager.boundary_checkpoint_page_positions(prefix), + draft_prefix_len=prefix, + kind="prompt", + num_ranks=1, + ) + assert checkpoint is not None + assert manager.acknowledge_external_boundary_checkpoint( + checkpoint.checkpoint_id, 0 + ) + return checkpoint + + checkpoint = publish() + bridge.store(producer, checkpoint) + (store,) = bridge.take_tasks() + bridge.complete({store.task_id: {rank: True for rank in range(4)}}) + bridge.finish_request(producer.request_id) + assert manager.reset_prefix_cache() + consumer = make_request("consumer") + if already_local: + checkpoint = publish() + assert not bridge.poll_prefix(consumer) + if not already_local: + checkpoint = publish() + # Other requests occupy everything except the unpinned cached bundle. + pressure = manager.block_pool.get_new_blocks( + manager.block_pool.get_num_free_blocks() - len(checkpoint.dependencies) + ) + assert bridge.poll_prefix(consumer) + blocks, tokens, _ = manager.get_computed_blocks(consumer) + assert tokens == prefix + assert bridge.external_tokens(consumer) == 0 + assert bridge.poll_prefix(consumer) + assert lookups == 1 + assert bridge.take_tasks() == [] + assert ( + manager.allocate_slots( + consumer, + 1, + num_new_computed_tokens=tokens, + new_computed_blocks=blocks, + full_sequence_must_fit=True, + ) + is None + ) + # A running request grows into the cached bundle before capacity returns. + growth = manager.block_pool.get_new_blocks(1) + manager.block_pool.free_blocks(pressure + growth) + assert manager.get_computed_blocks(consumer)[1] == 0 + assert not bridge.poll_prefix(consumer) + assert lookups == 2 + assert not bridge.poll_prefix(consumer) + (restore,) = bridge.take_tasks() + bridge.complete({restore.task_id: {rank: True for rank in range(4)}}) + assert bridge.poll_prefix(consumer) + blocks, tokens, _ = manager.get_computed_blocks(consumer) + assert tokens == prefix + assert bridge.external_tokens(consumer) == prefix + assert ( + manager.allocate_slots( + consumer, + 1, + num_new_computed_tokens=tokens, + new_computed_blocks=blocks, + full_sequence_must_fit=True, + ) + is not None + ) + _, copies = manager.take_kv_cache_block_copies() + manager.block_pool.free_blocks(copies) + bridge.finish_request(consumer.request_id) + manager.free(consumer) + assert not bridge.has_pending + assert manager.external_boundary_reserved_blocks() == 0 + assert manager.block_pool.get_num_free_blocks() == 63 + + +@pytest.mark.parametrize("reply", ["never", "error", "miss"]) def test_lookup_without_a_usable_reply_admits_the_request( reply: str, monkeypatch: pytest.MonkeyPatch ) -> None: @@ -859,7 +1006,10 @@ def test_lookup_without_a_usable_reply_admits_the_request( future = ( SimpleNamespace(query=lambda: False, result=failed.result) if reply == "never" - else SimpleNamespace(query=lambda: True, result=failed.result) + else SimpleNamespace( + query=lambda: True, + result=(lambda: None) if reply == "miss" else failed.result, + ) ) consumer = make_request("consumer") with monkeypatch.context() as patch: @@ -869,6 +1019,7 @@ def test_lookup_without_a_usable_reply_admits_the_request( assert not bridge.poll_prefix(consumer) time.sleep(0.3) assert bridge.poll_prefix(consumer) + assert bridge.poll_prefix(consumer) assert bridge.take_tasks() == [] assert bridge.external_tokens(consumer) == 0 assert not bridge.has_pending