fix(checkpoint): keep intact disk checkpoints listed when a restore cannot get RAM (ktsaou's #105 + follow-ups)

.lil/changes/lmcache-105.json

@@ -0,0 +1,23 @@
+{
+  "schema": "local-inference-release-change/v1",
+  "id": "lmcache-105",
+  "category": "fix",
+  "summary": "A checkpoint restore that cannot get RAM no longer removes an intact disk checkpoint from the directory",
+  "models": [
+    "GLM-5.3-Flash",
+    "Qwen3.8-Flash-Next"
+  ],
+  "compatibility": "No user action required.",
+  "details": [
+    "A restore from disk loads a checkpoint's pages only into RAM reserved for all of them. When that reservation failed, the restore was repeated once, and if the repeat failed too while RAM looked free by the time the cache checked (RAM held briefly by other restores or writes, or free RAM split into pieces smaller than a page), the checkpoint was removed from the directory although every page was intact on disk. Later turns of that conversation then recomputed their prompts instead of restoring.",
+    "The cache now decides from the allocation result itself: a restore whose pages could not be reserved is repeated, 20 ms apart, until the store admission timeout (default 8 s), and then misses with its checkpoint still listed, so a later request restores it. Only a restore that got its RAM and still found pages unreadable retires the checkpoint, also when the first lookup answers after the timeout. A checkpoint larger than the whole RAM tier also stays listed, and the room a waiting restore asks eviction for now includes each page's 4 KiB alignment."
+  ],
+  "pull_requests": [
+    105
+  ],
+  "requires": [],
+  "authors": [
+    "ktsaou",
+    "voipmonitor"
+  ]
+}

lmcache/v1/distributed/storage_controllers/prefetch_controller.py

@@ -193,6 +193,18 @@ class PrefetchPhase(enum.Enum):
     PLAN_AND_LOAD = enum.auto()
 
 
+@dataclass(frozen=True)
+class PrefetchResult:
+    """Retained keys and whether reserving their load buffers failed.
+
+    Allocation failure or contention does not prove stored pages are absent.
+    ``found`` indexes the keys submitted to the prefetch request.
+    """
+
+    found: Bitmap
+    reservation_failed: bool = False
+
+
 @dataclass
 class InFlightPrefetchRequest:
     """Tracks a single prefetch request across its lifecycle phases."""
@@ -245,6 +257,8 @@ class InFlightPrefetchRequest:
     """Maps object_group_id to that group's layout (one ``MemoryLayoutDesc``
     describes a single group's MemoryObj). Covers every object group."""
 
+    reservation_failed: bool = False
+
     def all_lookups_done(self) -> bool:
         return len(self.pending_lookup_tasks) == 0
 
@@ -338,6 +352,7 @@ def __init__(
         self._prefetch_results_lock = threading.Lock()
         self._prefetch_results_cv = threading.Condition(self._prefetch_results_lock)
         self._completed_results: dict[PrefetchRequestId, Bitmap] = {}
+        self._completed_reservation_failures: set[PrefetchRequestId] = set()
 
         # Map eventfds to adapter indices for quick lookup in poll.
         # Relies on the L2AdapterInterface contract that every adapter
@@ -488,12 +503,26 @@ def query_prefetch_result(self, request_id: PrefetchRequestId) -> Bitmap | None:
             query_lookup_result after calling this function, otherwise it will
             get None forever.
         """
+        result = self.query_prefetch_result_detailed(request_id)
+        return result.found if result is not None else None
+
+    def query_prefetch_result_detailed(
+        self, request_id: PrefetchRequestId
+    ) -> PrefetchResult | None:
+        """Consume hits and allocation evidence for ``request_id`` atomically.
+
+        Returns None while pending or already consumed. Bitmap-only queries
+        consume the same result and release the same lookup bookkeeping.
+        """
         with self._prefetch_results_lock:
-            result = self._completed_results.pop(request_id, None)
-        if result is not None:
-            with self._lookup_results_lock:
-                self._completed_lookups.pop(request_id, None)
-        return result
+            found = self._completed_results.pop(request_id, None)
+            if found is None:
+                return None
+            failed = request_id in self._completed_reservation_failures
+            self._completed_reservation_failures.discard(request_id)
+        with self._lookup_results_lock:
+            self._completed_lookups.pop(request_id, None)
+        return PrefetchResult(found, failed)
 
     def wait_prefetch_result(
         self, request_id: PrefetchRequestId, timeout: float
@@ -1102,6 +1131,7 @@ def _reserve_load_buffers(
                 request.write_reserved_objs[key] = mem_obj
                 reserved.add(key)
                 continue
+            request.reservation_failed = True
             if err == L1Error.OUT_OF_MEMORY:
                 oom_keys.append(key)
             elif err == L1Error.KEY_NOT_WRITABLE:
@@ -1600,6 +1630,9 @@ def _release_l2_locks(
     def _complete_request(self, request_id: PrefetchRequestId, result: Bitmap) -> None:
         """Store the retained-key bitmap and remove from in-flight tracking."""
         with self._prefetch_results_lock:
+            request = self._in_flight_requests.get(request_id)
+            if request is not None and request.reservation_failed:
+                self._completed_reservation_failures.add(request_id)
             self._completed_results[request_id] = result
             # Wake any WAIT_PREFETCH_STATUS handler blocked on this result.
             self._prefetch_results_cv.notify_all()

lmcache/v1/distributed/storage_manager.py

@@ -56,6 +56,9 @@
     PrefetchController,
     StoreController,
 )
+from lmcache.v1.distributed.storage_controllers.prefetch_controller import (
+    PrefetchResult,
+)
 from lmcache.v1.distributed.storage_controllers.prefetch_policy import (
     create_prefetch_policy,
 )
@@ -1030,13 +1033,27 @@ def query_prefetch_status(
             done, None if it's still in progress. Derive the prefix hit count
             via ``count_leading_ones``.
         """
+        result = self.query_prefetch_status_detailed(handle)
+        return result.found if result is not None else None
+
+    def query_prefetch_status_detailed(
+        self, handle: PrefetchHandle
+    ) -> PrefetchResult | None:
+        """Consume hits and load-allocation evidence for ``handle``.
+
+        Returns None while pending or already consumed. Found bits index the
+        original requested keys. Bitmap-only queries consume the same result.
+        """
         l2_r: Bitmap | None = None
+        reservation_failed = False
         if handle.prefetch_request_id != -1:
-            l2_r = self._prefetch_controller.query_prefetch_result(
+            result = self._prefetch_controller.query_prefetch_result_detailed(
                 handle.prefetch_request_id
             )
-            if l2_r is None:
+            if result is None:
                 return None
+            l2_r = result.found
+            reservation_failed = result.reservation_failed
 
         found = self._combine_found(handle, l2_r)
         # popcount (not count_leading_ones) so the log is accurate for
@@ -1060,7 +1077,7 @@ def query_prefetch_status(
                 handle.external_request_id,
                 handle.prefetch_request_id,
             )
-        return found
+        return PrefetchResult(found, reservation_failed)
 
     def touch_l1_keys(self, keys: list[ObjectKey]):
         """

lmcache/v1/multiprocess/checkpoint_storage.py

@@ -49,6 +49,8 @@
 # seconds, as store admission does: a pass frees nothing while its victims are
 # still pinned by other restores or L2 writes.
 _ROOM_REQUEST_INTERVAL_SECONDS = 0.5
+# Pause before repeating a lookup whose load buffers could not be reserved.
+_RESERVATION_RETRY_SECONDS = 0.02
 # Released cancelled lookups remembered so that a later poll from their worker
 # still gets the miss; a poll of an older one raises KeyError.
 _MAX_RELEASED_CANCELLED_LOOKUPS = 4096
@@ -237,6 +239,7 @@ class _RetrieveLease:
     readable: int = 0
     unread_bytes: int = 0
     eviction_requested: float = 0.0
+    resubmit_after: float = 0.0
 
 
 class CheckpointPayloadStore:
@@ -624,8 +627,8 @@ def poll_retrieve(self, lease_id: str) -> CheckpointSlots | bool | None:
             The storage manager loads pages from L2 only into RAM it reserved
             for all of them, so a full RAM leaves stored pages unread. A lookup
             with unread pages is therefore repeated once RAM has room for them,
-            and only a repeated lookup that had that room invalidates the
-            generation. The lease stays pending while eviction makes room, for
+            and only a repeated lookup without allocation failure invalidates
+            the generation. The lease stays pending while eviction makes room, for
             at most the storage admission timeout; a lease that never gets
             room misses without invalidating its generation.
 
@@ -680,13 +683,16 @@ def _poll(self, lease_id: str) -> CheckpointSlots | bool | None:
                 return lease.slots
             if lease.handle is None:
                 return self._repeat_with_room(lease_id, lease)
-            found = self._storage.query_prefetch_status(lease.handle)
-            if found is None:
+            result = self._storage.query_prefetch_status_detailed(lease.handle)
+            if result is None:
                 return None
+            found = result.found
             readable_keys = [key for i, key in enumerate(lease.keys) if found.test(i)]
             if len(readable_keys) != len(lease.keys) or lease.cancelled:
                 self._storage.finish_read_prefetched(readable_keys)
-                return self._repeat_or_miss(lease_id, lease, readable_keys)
+                return self._repeat_or_miss(
+                    lease_id, lease, readable_keys, result.reservation_failed
+                )
             try:
                 keys, objects = self._storage.unsafe_read(lease.keys)
                 if keys != lease.keys or len(objects) != len(keys):
@@ -711,30 +717,51 @@ def _poll(self, lease_id: str) -> CheckpointSlots | bool | None:
                 raise
 
     def _repeat_or_miss(
-        self, lease_id: str, lease: _RetrieveLease, readable_keys: list[ObjectKey]
+        self,
+        lease_id: str,
+        lease: _RetrieveLease,
+        readable_keys: list[ObjectKey],
+        reservation_failed: bool,
     ) -> bool | None:
         """Repeat or end a lookup whose pages were not all pinned.
 
         Called with the lookup's read locks already released. Returns None
         when the lookup will be repeated, otherwise False after forgetting the
-        lease. Only a repeated lookup that had room in RAM invalidates.
+        lease. Only a repeated lookup without allocation failure invalidates.
         """
         lease.readable = len(readable_keys)
         if lease.cancelled:
             return self._miss(lease_id, lease, "was cancelled by the engine")
         groups = checkpoint_page_groups(lease.manifest)
         readable = set(readable_keys)
+        alignment = self._storage.l1_memory_desc.align_bytes
         lease.unread_bytes = sum(
-            groups[key.object_group_id].page_bytes
+            ((groups[key.object_group_id].page_bytes + alignment - 1) // alignment)
+            * alignment
             for key in lease.keys
             if key not in readable
         )
         used, total = self._storage.get_l1_usage()
+        if reservation_failed and lease.unread_bytes > total:
+            return self._miss(
+                lease_id, lease, "cannot fit in RAM; its checkpoint stays listed"
+            )
+        if reservation_failed:
+            now = time.monotonic()
+            if now - lease.started >= self._storage.store_admission_timeout_seconds:
+                return self._miss(
+                    lease_id, lease, "found no room in RAM; its checkpoint stays listed"
+                )
+            lease.resubmit_after = now + _RESERVATION_RETRY_SECONDS
         # Without L2, or with more pages than RAM can hold, no repeat loads them.
         if (
             self._storage.l2_adapters()
             and lease.unread_bytes <= total
-            and (lease.lookups == 1 or total - used < lease.unread_bytes)
+            and (
+                reservation_failed
+                or lease.lookups == 1
+                or total - used < lease.unread_bytes
+            )
         ):
             lease.handle = None
             return self._repeat_with_room(lease_id, lease)
@@ -751,12 +778,12 @@ def _repeat_with_room(self, lease_id: str, lease: _RetrieveLease) -> bool | None
         """
         if lease.cancelled:
             return self._miss(lease_id, lease, "was cancelled by the engine")
+        now = time.monotonic()
         used, total = self._storage.get_l1_usage()
-        if total - used >= lease.unread_bytes:
+        if total - used >= lease.unread_bytes and now >= lease.resubmit_after:
             lease.handle = self._lookup(lease_id, lease.manifest, lease.keys)
             lease.lookups += 1
             return None
-        now = time.monotonic()
         if now - lease.started >= self._storage.store_admission_timeout_seconds:
             return self._miss(
                 lease_id, lease, "found no room in RAM; its checkpoint stays listed"

tests/v1/multiprocess/test_checkpoint_cancel_drain.py

@@ -141,10 +141,10 @@ def test_failed_poll_of_a_pending_lookup_is_not_fatal() -> None:
 def _stall_lookups(monkeypatch: pytest.MonkeyPatch, storage: Any) -> threading.Event:
     """Keep every prefetch pending until the returned event is set."""
     answer = threading.Event()
-    query = storage.query_prefetch_status
+    query = storage.query_prefetch_status_detailed
     monkeypatch.setattr(
         storage,
-        "query_prefetch_status",
+        "query_prefetch_status_detailed",
         lambda handle: query(handle) if answer.is_set() else None,
     )
     return answer