[MP] Preserve checkpoint integrity during deferred-write eviction

.lil/changes/lmcache-checkpoint-on-evict-ownership.json

@@ -0,0 +1,15 @@
+{
+  "schema": "local-inference-release-change/v1",
+  "id": "lmcache-checkpoint-on-evict-ownership",
+  "category": "fix",
+  "summary": "Keep deferred checkpoint storage coherent under memory and disk pressure",
+  "models": ["Qwen3.8-Flash-Next", "GLM-5.3-Flash"],
+  "compatibility": "No user action required.",
+  "details": [
+    "Checkpoint publication and active transfers retain their pages until completion; supersession alone no longer discards an older branch's unwritten pages.",
+    "Persistence queues are bounded by allocated bytes, shared pages are deduplicated, and transient failures retry missing copies.",
+    "Sustained memory pressure can retire cold checkpoints before reclaiming their last copy. Shutdown drains writes within its budget and removes incomplete entries; restart validates recovered payload availability."
+  ],
+  "authors": ["ktsaou"],
+  "requires": ["b12x-checkpoint-shutdown-order"]
+}

.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
+  ]
+}

docs/source/mp/l2_storage/index.rst

@@ -117,9 +117,9 @@ Select policies via CLI:
    * - ``--l2-store-policy``
      - ``checkpoint_on_evict``
      - Store ordinary keys like ``default``.  Keep new recurrent checkpoint
-       pages in L1 and store a page to L2 once, when L1 is about to evict it
-       and no newer checkpoint of the same conversation has superseded it.
-       On shutdown, current pages still only in L1 are written within
+       pages in L1 and persist retained pages under memory pressure.
+       Supersession is an eviction hint, not proof an older branch is dead.
+       On shutdown, retained pages still only in L1 are written within
        ``--checkpoint-shutdown-flush-seconds``.  See
        :ref:`checkpoint-retention`.
    * - ``--l2-prefetch-policy``
@@ -458,32 +458,42 @@ not extend.  Checkpoints of the same request and ``instruction``
 checkpoints, which other conversations share, are never marked.  With every
 store policy:
 
-* L1 eviction drops superseded pages before any LRU victim.
-* L2 eviction deletes superseded pages before any LRU victim.
+* Supersession and recency guide eviction; shared dependencies remain protected.
+* Planned loss of the final known copy retires affected directory entries first.
 
 With ``--l2-store-policy checkpoint_on_evict`` in addition:
 
-* A current checkpoint page is written to L2 once, when L1 is about to
-  evict it, instead of on every request.  L1 keeps the page until the
-  write completes (at most ``--checkpoint-write-timeout-seconds``).
-* A superseded page is never written to L2.
-* Shutdown writes current pages still only in L1, so they can be restored
-  after a restart.
-
-Superseded manifests stay listed, so a request that branches from an older
-turn can still restore it while its pages last; when a page is gone, the
-restore falls back to the longest remaining checkpoint.
+* Retained checkpoint pages are offered to L2 under RAM pressure or shutdown.
+  Existing disk copies are reused, including pages shared between checkpoints.
+* Queued and active persistence have an allocated-byte budget. Transient failures
+  retry missing replicas; a failed replica does not starve a healthy destination.
+* If capacity cannot be reclaimed within bounded admission waits, cold checkpoints
+  may be retired coherently. Unneeded queued writes can be cancelled, while active
+  operations retain their buffers until completion. The legacy
+  ``--checkpoint-write-timeout-seconds`` option no longer releases pages by age.
+* Shutdown attempts persistence within ``--checkpoint-shutdown-flush-seconds``;
+  remaining RAM-only generations may be lost and are removed from the directory.
+  Other policies drain already eligible writes without making unused on-reuse
+  checkpoints eligible.
+
+An older branch remains eligible until retirement. Recovery prunes manifests
+whose payloads are provably missing; adapters without a complete inventory keep
+lookup candidates for validation. Unknown inventory never certifies another copy
+during planned deletion. Restores validate all required pages and fall back to
+the longest complete compatible checkpoint still retained.
 
 **Sizing.**  With write-through (``default``) the L2 retention time is about::
 
     L2 capacity / (requests per minute x checkpoint bytes per request)
 
-With ``checkpoint_on_evict`` the L2 receives roughly one current checkpoint
-per conversation that leaves L1, so retention is about::
+With ``checkpoint_on_evict``, write savings depend on checkpoints retiring before
+they need persistence. Retaining every branch can eventually write as much as
+write-through. Estimate retention from the measured workload::
 
-    L2 capacity / (conversations leaving L1 per minute x checkpoint bytes)
+    L2 capacity / checkpoint payload bytes written per minute
 
-and superseded pages are reclaimed first.  The
+The
 ``lmcache_mp_checkpoint_retention`` gauge reports, per ``stat``, superseded
 pages, L1 drops and L2 evictions of superseded pages, write-on-evict
-requests, completions and timeouts, and the checkpoint bytes held in L2.
+requests, completions, explicit retirements, and the checkpoint bytes held in L2.
+The legacy write-timeout counter remains zero; source ownership does not expire.

lmcache/integration/vllm/checkpoint_scheduler.py

@@ -2,7 +2,7 @@
 """Scheduler ownership and all-rank publication of external recurrent checkpoints."""
 
 # Standard
-from dataclasses import dataclass, field
+from dataclasses import dataclass, field, replace
 from typing import TYPE_CHECKING, Any, Literal
 import json
 import math
@@ -25,9 +25,8 @@
 
 logger = init_logger(__name__)
 
-# Directory lookups per request, including the first. A failed restore makes
-# the server invalidate the missing generation, so each retry receives the
-# longest candidate that is still listed rather than recomputing the prompt.
+# Directory lookups per request, including the first. Failed restores retry
+# strictly shorter candidates without requiring global manifest invalidation.
 _MAX_LOOKUP_ATTEMPTS = 4
 
 # Restores slower than this many seconds are logged even when they succeed;
@@ -76,9 +75,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 +92,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 +176,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 +188,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 +206,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 +252,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 +280,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 +293,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 +557,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 +590,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,8 +616,24 @@ def complete(self, results: dict[str, dict[int, bool]]) -> None:
                     )
                     state.task_id = None
                     state.attempts += 1
+                    state.started = time.monotonic()
+                    # A publication race does not make the stored payload
+                    # unusable; only failed worker copies need a shorter prefix.
+                    limit = (
+                        pending.task.manifest.prefix.num_tokens
+                        if all(pending.acknowledgements.values())
+                        else pending.task.manifest.prefix.num_tokens - 1
+                    )
+                    shorter_roots = tuple(
+                        replace(
+                            root,
+                            tail_tokens=root.tail_tokens[: limit - root.start_tokens],
+                        )
+                        for root in state.roots.roots
+                        if root.start_tokens < limit
+                    )
                     state.future = self._client.submit_request(
-                        RequestType.CHECKPOINT_FIND, [state.roots.roots]
+                        RequestType.CHECKPOINT_FIND, [shorter_roots]
                     )
                 elif state is not None:
                     state.done = True
@@ -578,6 +643,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)

lmcache/v1/distributed/admission.py

@@ -103,12 +103,14 @@ def reserve_with_eviction_backpressure(
     if retry_wait_seconds <= 0:
         raise ValueError("retry_wait_seconds must be > 0")
 
-    generation = get_generation()
     first = attempt()
     if first.failure_reason is None:
         return AdmissionOutcome(first.value, None, retries=0, waited=False)
     if first.failure_reason is not AdmissionFailure.CAPACITY:
         return AdmissionOutcome(None, first.failure_reason, retries=0, waited=False)
+    # A failed atomic attempt rolls back its own partial allocations. Observe
+    # capacity afterwards so that rollback cannot wake an immediate busy loop.
+    generation = get_generation()
 
     on_wait()
     deadline = time.monotonic() + timeout_seconds
@@ -143,3 +145,4 @@ def reserve_with_eviction_backpressure(
             return AdmissionOutcome(
                 None, retried.failure_reason, retries=retries, waited=True
             )
+        generation = get_generation()