diff --git a/.lil/changes/lmcache-checkpoint-on-evict-ownership.json b/.lil/changes/lmcache-checkpoint-on-evict-ownership.json new file mode 100644 index 00000000000..8a1a22687fd --- /dev/null +++ b/.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"] +} 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/docs/source/mp/l2_storage/index.rst b/docs/source/mp/l2_storage/index.rst index a89c8c3a188..d4d6b8ca5d1 100644 --- a/docs/source/mp/l2_storage/index.rst +++ b/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. diff --git a/lmcache/integration/vllm/checkpoint_scheduler.py b/lmcache/integration/vllm/checkpoint_scheduler.py index 1038455515d..f07e9ecedcf 100644 --- a/lmcache/integration/vllm/checkpoint_scheduler.py +++ b/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) diff --git a/lmcache/v1/distributed/admission.py b/lmcache/v1/distributed/admission.py index 4638f37f0cc..0d211fd0afb 100644 --- a/lmcache/v1/distributed/admission.py +++ b/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() diff --git a/lmcache/v1/distributed/checkpoint_retention.py b/lmcache/v1/distributed/checkpoint_retention.py index 1cb29a77af3..6d3dc04bc49 100644 --- a/lmcache/v1/distributed/checkpoint_retention.py +++ b/lmcache/v1/distributed/checkpoint_retention.py @@ -1,29 +1,19 @@ # SPDX-License-Identifier: Apache-2.0 -""" -Decide which recurrent checkpoint pages deserve L1 memory and L2 storage. - -A conversation publishes new request-boundary checkpoints every turn, and -each one supersedes the previous turn's endpoint state. Pages shared with the -newer checkpoint stay alive through it, so only the pages unique to the older -checkpoint (recurrent endpoint state, a partial attention page, auxiliary -state) become dead weight. This module keeps that knowledge: - -* Superseded pages are dropped first from L1 and L2 and are never written to - L2 when they leave L1. -* With write-on-evict storage, a current checkpoint page is written to L2 - once, when L1 is about to evict it, instead of on every request. -* L2 residency per adapter is tracked so a page that already reached L2 is - not written again. - -Everything is bounded in memory and safe to call from the store, eviction and -request-handler threads. Forgetting an entry only costs an extra write or a -later eviction; correctness never depends on this state, because restores -validate every page and fall back to a shorter checkpoint when one is gone. +"""Coordinate checkpoint persistence, replica accounting and coherent retirement. + +Supersession is an eviction hint, never proof that an older branch is dead. +Deferred writes start under RAM pressure or shutdown. Directory lifecycle +callbacks retire affected checkpoints before their final known copy is reclaimed. + +Supersession hints are bounded independently. Replica accounting is exact for +inventoried physical payloads and shrinks when adapters report deletion; it must +not forget a live replica because a telemetry history limit was reached. """ # Standard from collections import OrderedDict from collections.abc import Callable, Iterable +from typing import TypeVar import threading import time @@ -33,6 +23,7 @@ from lmcache.v1.distributed.internal_api import L2AdapterListener logger = init_logger(__name__) +_DeleteResult = TypeVar("_DeleteResult") class _BoundedSet: @@ -89,16 +80,17 @@ class CheckpointRetention: persist: Callback that asynchronously stores keys to L2; installed by the storage manager once its store controller exists. max_superseded: Bound on remembered superseded pages. - max_tracked: Bound on remembered L2-resident pages per adapter. - persist_timeout: Seconds after which a page whose write never - completed may be evicted from L1 without reaching L2. + max_tracked: Legacy compatibility argument; physical replica inventory + is exact and bounded by stored payloads, not a telemetry history cap. + persist_timeout: Legacy compatibility argument. Elapsed write time alone + no longer permits losing a retained checkpoint page. """ def __init__( self, *, write_on_evict: bool = False, - persist: Callable[[list[ObjectKey]], None] | None = None, + persist: Callable[[list[ObjectKey]], list[ObjectKey] | None] | None = None, max_superseded: int = 262144, max_tracked: int = 1048576, max_generations: int = 65536, @@ -106,13 +98,12 @@ def __init__( ) -> None: self._write_on_evict = write_on_evict self._persist = persist - self._max_tracked = max_tracked - self._persist_timeout = persist_timeout self._lock = threading.Lock() self._superseded = _BoundedSet(max_superseded) self._superseded_generations = _BoundedSet(max_generations) # adapter id -> key -> size in bytes self._resident: dict[int, OrderedDict[ObjectKey, int]] = {} + self._unknown_inventories: set[int] = set() # key -> monotonic time its write was requested self._pending: dict[ObjectKey, float] = {} self._stats = { @@ -124,13 +115,93 @@ def __init__( "write_on_evict_requests": 0, "write_on_evict_persisted": 0, "write_on_evict_timeouts": 0, + "retired_checkpoints": 0, } + self._evict: Callable | None = None + self._pressure: Callable[[int], int] | None = None + self._retained: Callable[[ObjectKey], bool] | None = None + self._order: Callable[[list[ObjectKey]], list[ObjectKey]] | None = None + self._reconcile: Callable[[], None] | None = None + + def set_lifecycle( + self, + *, + evict: Callable | None, + pressure: Callable[[int], int] | None, + retained: Callable[[ObjectKey], bool] | None, + order: Callable[[list[ObjectKey]], list[ObjectKey]] | None = None, + reconcile: Callable[[], None] | None = None, + ) -> None: + """Attach the directory authority; callbacks run outside retention locks.""" + self._evict, self._pressure, self._retained = evict, pressure, retained + self._order = order + self._reconcile = reconcile + + def evict( + self, + keys: list[ObjectKey], + tier: str | int, + delete: Callable[[list[ObjectKey]], _DeleteResult], + ) -> _DeleteResult: + """Retire last-copy owners before deletion; return the callback result. + + ``tier`` identifies L1 or an adapter. The callback receives only keys + eligible for deletion, and its exceptions propagate to the caller. + """ + if self._evict is not None: + return self._evict(keys, tier, delete) + return delete(keys) + + def retire_under_pressure(self, target_bytes: int) -> int: + """Retire cold generations to free capacity after bounded admission waits.""" + return self._pressure(target_bytes) if self._pressure is not None else 0 + + def is_retained(self, key: ObjectKey) -> bool: + """Whether a checkpoint page still has a generation owner.""" + return self._retained(key) if self._retained is not None else True + + def resident_adapters( + self, key: ObjectKey, expected_size: int | None = None + ) -> set[int]: + """Return known replicas, optionally requiring the manifest's byte size. + + Args: + key: Checkpoint page to inspect. + expected_size: Required payload bytes, or None for presence only. + + Returns: + Adapter IDs with a matching recorded replica. + """ + with self._lock: + return { + adapter + for adapter, resident in self._resident.items() + if key in resident + and (expected_size is None or resident[key] == expected_size) + } + + def record_retirement(self, count: int) -> None: + """Count explicitly retired generations separately from failed restores.""" + with self._lock: + self._stats["retired_checkpoints"] += count + + def set_inventory_complete(self, adapter_id: int) -> None: + """Mark a successfully enumerated adapter inventory as authoritative.""" + with self._lock: + self._unknown_inventories.discard(adapter_id) + + def inventories_complete(self) -> bool: + """Whether missing residency is proof of absence across all adapters.""" + with self._lock: + return not self._unknown_inventories @property def write_on_evict(self) -> bool: return self._write_on_evict - def set_persist(self, persist: Callable[[list[ObjectKey]], None]) -> None: + def set_persist( + self, persist: Callable[[list[ObjectKey]], list[ObjectKey] | None] + ) -> None: """Install the asynchronous L2 store callback.""" self._persist = persist @@ -138,11 +209,22 @@ def listener_for(self, adapter_id: int) -> L2AdapterListener: """Return a listener that records one adapter's checkpoint pages.""" with self._lock: self._resident.setdefault(adapter_id, OrderedDict()) + self._unknown_inventories.add(adapter_id) return _AdapterResidencyListener(self, adapter_id) def forget_adapter(self, adapter_id: int) -> None: + with self._lock: + keys = list(self._resident.get(adapter_id, {})) + self.evict( + keys, + adapter_id, + lambda selected: self.record_l2_absent(adapter_id, selected), + ) with self._lock: self._resident.pop(adapter_id, None) + self._unknown_inventories.discard(adapter_id) + if self._reconcile is not None: + self._reconcile() # ----- supersession --------------------------------------------------- @@ -196,12 +278,13 @@ def record_l2_present( for key, size in zip(keys, sizes, strict=False): if not is_recurrent_checkpoint_key(key): continue - resident[key] = size + if size or key not in resident: + resident[key] = size resident.move_to_end(key) if self._pending.pop(key, None) is not None: self._stats["write_on_evict_persisted"] += 1 - while len(resident) > self._max_tracked: - resident.popitem(last=False) + # Residency is bounded by physical tier contents, not an LRU sample. + # Forgetting a live replica would make retirement decisions unsound. def record_l2_absent(self, adapter_id: int, keys: list[ObjectKey]) -> None: with self._lock: @@ -211,9 +294,32 @@ def record_l2_absent(self, adapter_id: int, keys: list[ObjectKey]) -> None: for key in keys: resident.pop(key, None) - def is_l2_resident(self, key: ObjectKey) -> bool: + def is_l2_resident(self, key: ObjectKey, expected_size: int | None = None) -> bool: + """Whether a replica exists with the optional required payload byte size.""" with self._lock: - return any(key in resident for resident in self._resident.values()) + return any( + key in resident + and (expected_size is None or resident[key] == expected_size) + for resident in self._resident.values() + ) + + def replica_may_match( + self, adapter_id: int, key: ObjectKey, expected_size: int + ) -> bool: + """Allow lookup unless inventory positively proves a byte-size mismatch. + + Args: + adapter_id: Backend being considered for the load. + key: Requested object key. + expected_size: Required logical payload bytes. + + Returns: + True for unknown or matching sizes; False for a known mismatch. + """ + with self._lock: + resident = self._resident.get(adapter_id) + size = resident.get(key) if resident is not None else None + return size is None or size == expected_size def superseded_in_adapter( self, adapter_id: int, max_bytes: int @@ -247,42 +353,52 @@ def record_l1_superseded_drops(self, count: int) -> None: def needs_persist_before_evict(self, key: ObjectKey) -> bool: """Whether L1 must keep ``key`` until its L2 write completes. - True for a current checkpoint page that is not in L2 yet, while its - write is still expected to finish. A superseded page, an ordinary KV - chunk or a page whose write timed out may be evicted. + True for a checkpoint page without a known L2 copy. Ordinary KV + chunks follow their existing policy. Capacity-driven loss first retires + the owning generations; age and supersession alone never allow loss. """ if not self._write_on_evict or not is_recurrent_checkpoint_key(key): return False - now = time.monotonic() with self._lock: - if key in self._superseded: - return False if any(key in resident for resident in self._resident.values()): self._pending.pop(key, None) return False - requested = self._pending.get(key) - if requested is not None and now - requested > self._persist_timeout: - del self._pending[key] - self._stats["write_on_evict_timeouts"] += 1 - logger.warning( - "Checkpoint page write to L2 did not complete within %.0f s; " - "evicting it from L1 without an L2 copy", - self._persist_timeout, - ) - return False return True def request_persist(self, keys: list[ObjectKey]) -> int: """Ask the store controller to write pages not already requested.""" + if self._order is not None: + keys = self._order(keys) now = time.monotonic() with self._lock: - fresh = [key for key in keys if key not in self._pending] + fresh = [ + key + for key in keys + if key not in self._pending + and not any(key in resident for resident in self._resident.values()) + ] for key in fresh: self._pending[key] = now self._stats["write_on_evict_requests"] += len(fresh) if fresh and self._persist is not None: - self._persist(fresh) - return len(fresh) + accepted = self._persist(fresh) + if accepted is not None: + accepted_keys = set(accepted) + with self._lock: + for key in fresh: + if key not in accepted_keys: + self._pending.pop(key, None) + return ( + len(accepted_keys) + if fresh and self._persist is not None and accepted is not None + else len(fresh) + ) + + def forget_pending(self, keys: Iterable[ObjectKey]) -> None: + """Forget requests for pages no retained generation needs.""" + with self._lock: + for key in keys: + self._pending.pop(key, None) def pending_count(self) -> int: with self._lock: diff --git a/lmcache/v1/distributed/config.py b/lmcache/v1/distributed/config.py index add647874cc..3d92481e804 100644 --- a/lmcache/v1/distributed/config.py +++ b/lmcache/v1/distributed/config.py @@ -332,12 +332,12 @@ class StorageManagerConfig: """ Total deadline for capacity-only L1 store admission retries. """ checkpoint_shutdown_flush_seconds: float = 30.0 - """ With write-on-evict checkpoint storage, time budget for writing - current checkpoint pages still in L1 to L2 during a clean shutdown. """ + """ Time budget to drain checkpoint persistence during clean shutdown. + On-evict also initiates writes of retained RAM-only checkpoint pages. """ checkpoint_write_timeout_seconds: float = 120.0 - """ With write-on-evict checkpoint storage, how long L1 keeps a page - whose L2 write has not completed before evicting it without a copy. """ + """ Legacy compatibility setting; elapsed write time no longer permits + reclaiming an active buffer. Capacity pressure uses the admission deadline. """ def __post_init__(self) -> None: if self.store_admission_timeout_seconds < 0: @@ -601,17 +601,17 @@ def add_storage_manager_args( "--checkpoint-shutdown-flush-seconds", type=float, default=30.0, - help="With --l2-store-policy checkpoint_on_evict, time budget for " - "writing current recurrent checkpoints still in L1 to L2 on a clean " - "shutdown. 0 skips the flush.", + help="Time budget to drain recurrent checkpoint persistence on clean " + "shutdown. on-evict also starts writes for retained RAM-only pages. " + "0 skips the flush; incomplete checkpoints may be lost.", ) parser.add_argument( "--checkpoint-write-timeout-seconds", type=float, default=120.0, - help="With --l2-store-policy checkpoint_on_evict, how long L1 keeps a " - "checkpoint page whose L2 write has not completed before evicting it " - "without an L2 copy.", + help="Legacy compatibility option; no longer releases pages by write age. " + "Capacity pressure may retire checkpoints within the store admission " + "deadline, but active I/O retains its buffers until completion.", ) # L2 Policies diff --git a/lmcache/v1/distributed/internal_api.py b/lmcache/v1/distributed/internal_api.py index fb8c7009a0e..5995ebe3de3 100644 --- a/lmcache/v1/distributed/internal_api.py +++ b/lmcache/v1/distributed/internal_api.py @@ -46,6 +46,18 @@ class L1ManagerListener(EventListener): Listener for L1 manager events """ + def on_l1_committed_objects( + self, keys: list[ObjectKey], sizes: list[int] + ) -> list[ObjectKey]: + """Accept ownership before committed pages become evictable. + + Runs under the L1 lock; never call L1 from this callback. Returned + keys acquire one non-expiring internal pin, which the listener must + eventually release after its work drains. Default listeners own none. + """ + self.on_l1_keys_write_finished(keys) + return [] + @abstractmethod def on_l1_keys_reserved_read(self, keys: list[ObjectKey]): """ diff --git a/lmcache/v1/distributed/l1_manager.py b/lmcache/v1/distributed/l1_manager.py index 0f4848235fd..ce3ea45f90e 100644 --- a/lmcache/v1/distributed/l1_manager.py +++ b/lmcache/v1/distributed/l1_manager.py @@ -4,6 +4,7 @@ """ # Standard +from collections.abc import Callable from dataclasses import dataclass from itertools import chain, islice from typing import Literal @@ -56,6 +57,10 @@ class L1ObjectState: """ Whether the last write finished. An expired write lock releases the object but never makes an unfinished write readable. """ + internal_readers: int = 0 + """Non-expiring ownership for publication and local backend operations.""" + internal_writer: bool = False + def available_for_read(self) -> bool: """Check if the object is available for read. @@ -76,6 +81,8 @@ def available_for_write(self) -> bool: return ( not self.write_lock.is_locked() and not self.read_lock.is_locked() + and not self.internal_readers + and not self.internal_writer and not self.is_temporary ) @@ -260,6 +267,9 @@ def reserve_read( self, keys: list[ObjectKey], read_locks: int = 1, + *, + internal: bool = False, + retained_only: bool = False, ) -> dict[ObjectKey, L1OperationResult]: """Reserve read access for the given keys. @@ -270,6 +280,8 @@ def reserve_read( one per worker that consumes a read lock for the same key (e.g. MLA models with TP > 1). + internal: Use non-expiring ownership until the consumer releases it. + retained_only: Reject temporary pages that disappear after reads. Returns: A dictionary mapping each object key to a tuple @@ -289,15 +301,18 @@ def reserve_read( ret[key] = (L1Error.KEY_NOT_EXIST, None) continue - if not entry.available_for_read(): + if not entry.available_for_read() or (retained_only and entry.is_temporary): ret[key] = (L1Error.KEY_NOT_READABLE, None) continue # TODO(perf): support a count argument in # TTLLock.lock() to avoid Python for-loop # overhead (TTLLock is C++ std::atomic). - for _ in range(total): - entry.read_lock.lock() + if internal: + entry.internal_readers += total + else: + for _ in range(total): + entry.read_lock.lock() ret[key] = (L1Error.SUCCESS, entry.memory_obj) successful_keys.append(key) @@ -311,6 +326,62 @@ def reserve_read( ) return ret + @l1_mgr_synchronized + def readable_keys( + self, keys: list[ObjectKey], *, retained_only: bool = False + ) -> list[ObjectKey]: + """Snapshot readable keys without acquiring or releasing ownership. + + With retained_only, exclude temporary pages that disappear when their + readers finish. The snapshot does not pin returned objects. + + Args: + keys: Object keys to inspect, in result order. + retained_only: Whether to exclude temporary pages. + + Returns: + Readable keys at the instant of the snapshot. + """ + return [ + key + for key in keys + if (entry := self._objects.get(key)) is not None + and entry.available_for_read() + and not (retained_only and entry.is_temporary) + ] + + @l1_mgr_synchronized + def release_internal_reads(self, keys: list[ObjectKey]) -> None: + """Release one non-expiring ownership pin per key after work drains. + + Raises ValueError on an unmatched release. Internal ownership must + never be released merely because a client TTL or I/O deadline expires. + """ + for key in keys: + entry = self._objects.get(key) + if entry is None or entry.internal_readers <= 0: + raise ValueError("unmatched internal L1 ownership release") + entry.internal_readers -= 1 + if ( + entry.is_temporary + and not entry.internal_readers + and not entry.read_lock.is_locked() + ): + meta = self._object_meta(entry.memory_obj) + self._memory_manager.free([entry.memory_obj]) + del self._objects[key] + self._uncommitted.discard(key) + for listener in self._registered_listeners: + listener.on_l1_keys_deleted_by_manager([key]) + self._event_bus.publish( + Event( + event_type=EventType.L1_KEYS_EVICTED, + metadata={"keys": [key], "meta": [meta]}, + ) + ) + if keys: + self._notify_capacity_change() + @l1_mgr_synchronized def unsafe_read( self, @@ -341,7 +412,7 @@ def unsafe_read( ret[key] = (L1Error.KEY_NOT_EXIST, None) continue - if not entry.read_lock.is_locked(): + if not entry.read_lock.is_locked() and not entry.internal_readers: ret[key] = (L1Error.KEY_NOT_READABLE, None) continue @@ -419,7 +490,11 @@ def finish_read( # overhead (TTLLock is C++ std::atomic). for _ in range(total): entry.read_lock.unlock() - if entry.is_temporary and not entry.read_lock.is_locked(): + if ( + entry.is_temporary + and not entry.read_lock.is_locked() + and not entry.internal_readers + ): # NOTE: temporary objects shouldn't have write-locks need_to_free.append(entry.memory_obj) need_to_free_keys.append(key) @@ -459,6 +534,8 @@ def reserve_write( is_temporary: list[bool], layout_desc: MemoryLayoutDesc, mode: Literal["new", "update", "all"] = "all", + *, + internal: bool = False, ) -> dict[ObjectKey, L1OperationResult]: """Reserve write access for the given keys. @@ -566,6 +643,9 @@ def reserve_write( for listener in self._registered_listeners: listener.on_l1_keys_reserved_write(successful_keys) + if internal: + for key in successful_keys: + self._objects[key].internal_writer = True self._event_bus.publish( Event( event_type=EventType.L1_WRITE_RESERVED, @@ -639,7 +719,7 @@ def finish_write( ret[key] = L1Error.KEY_NOT_EXIST continue - if not entry.write_lock.is_locked(): + if not entry.write_lock.is_locked() and not entry.internal_writer: logger.warning( "L1Manager: finish write on non-write-locked key %s, " "potential inconsistent data might be written", @@ -658,6 +738,7 @@ def finish_write( continue entry.write_lock.unlock() + entry.internal_writer = False entry.committed = True self._uncommitted.discard(key) ret[key] = L1Error.SUCCESS @@ -665,7 +746,15 @@ def finish_write( successful_keys_meta.append(self._object_meta(entry.memory_obj)) for listener in self._registered_listeners: - listener.on_l1_keys_write_finished(successful_keys) + owned = listener.on_l1_committed_objects( + successful_keys, + [ + self._objects[key].memory_obj.get_physical_size() + for key in successful_keys + ], + ) + for key in owned: + self._objects[key].internal_readers += 1 self._event_bus.publish( Event( event_type=EventType.L1_WRITE_FINISHED, @@ -690,7 +779,11 @@ def abort_write(self, keys: list[ObjectKey]) -> dict[ObjectKey, L1Error]: if entry is None: ret[key] = L1Error.KEY_NOT_EXIST continue - if not entry.write_lock.is_locked() or entry.read_lock.is_locked(): + if ( + (not entry.write_lock.is_locked() and not entry.internal_writer) + or entry.read_lock.is_locked() + or entry.internal_readers + ): ret[key] = L1Error.KEY_IN_WRONG_STATE continue entry.write_lock.unlock() @@ -719,6 +812,9 @@ def finish_write_and_reserve_read( self, keys: list[ObjectKey], read_locks: int = 1, + *, + notify_store: bool = False, + internal: bool = False, ) -> dict[ObjectKey, L1OperationResult]: """Atomically finish write and acquire read lock for the given keys. @@ -732,6 +828,8 @@ def finish_write_and_reserve_read( read_locks: Total read locks acquired per key -- one per TP worker that consumes a read lock for the same key (e.g. MLA models with TP > 1). + notify_store: Announce newly produced pages to the store policy, + rather than treating them as pages prefetched from L2. Returns: A dictionary mapping each object key to a tuple of @@ -753,7 +851,7 @@ def finish_write_and_reserve_read( ret[key] = (L1Error.KEY_NOT_EXIST, None) continue - if not entry.write_lock.is_locked(): + if not entry.write_lock.is_locked() and not entry.internal_writer: logger.warning( "L1Manager: finish_write_and_reserve_read on " "non-write-locked key %s", @@ -771,16 +869,31 @@ def finish_write_and_reserve_read( continue entry.write_lock.unlock() + entry.internal_writer = False entry.committed = True self._uncommitted.discard(key) - for _ in range(total): - entry.read_lock.lock() + if internal: + entry.internal_readers += total + else: + for _ in range(total): + entry.read_lock.lock() ret[key] = (L1Error.SUCCESS, entry.memory_obj) successful_keys.append(key) successful_keys_meta.append(self._object_meta(entry.memory_obj)) for listener in self._registered_listeners: - listener.on_l1_keys_finish_write_and_reserve_read(successful_keys) + if notify_store: + owned = listener.on_l1_committed_objects( + successful_keys, + [ + self._objects[key].memory_obj.get_physical_size() + for key in successful_keys + ], + ) + for key in owned: + self._objects[key].internal_readers += 1 + else: + listener.on_l1_keys_finish_write_and_reserve_read(successful_keys) self._event_bus.publish( Event( event_type=EventType.L1_WRITE_FINISHED_AND_READ_RESERVED, @@ -789,6 +902,30 @@ def finish_write_and_reserve_read( ) return ret + @l1_mgr_synchronized + def handoff_internal_reads( + self, + keys: list[ObjectKey], + accept: Callable[[list[ObjectKey], list[int]], list[ObjectKey]], + ) -> list[ObjectKey]: + """Atomically offer readable objects to a bounded local work queue. + + accept(keys, sizes) returns accepted keys without calling L1. Those + keys acquire one internal pin before L1 unlocks. + """ + readable = [ + key + for key in dict.fromkeys(keys) + if key in self._objects and self._objects[key].available_for_read() + ] + owned = accept( + readable, + [self._objects[key].memory_obj.get_physical_size() for key in readable], + ) + for key in owned: + self._objects[key].internal_readers += 1 + return owned + @l1_mgr_synchronized def delete( self, keys: list[ObjectKey], force: bool = False @@ -819,7 +956,12 @@ def delete( ret[key] = L1Error.KEY_NOT_EXIST continue - locked = entry.read_lock.is_locked() or entry.write_lock.is_locked() + locked = ( + entry.read_lock.is_locked() + or entry.write_lock.is_locked() + or entry.internal_readers > 0 + or entry.internal_writer + ) if locked and not force: ret[key] = L1Error.KEY_IS_LOCKED continue @@ -863,7 +1005,10 @@ def reclaim_abandoned_writes(self) -> int: for key in list(self._uncommitted): entry = self._objects.get(key) if entry is not None and ( - entry.write_lock.is_locked() or entry.read_lock.is_locked() + entry.write_lock.is_locked() + or entry.read_lock.is_locked() + or entry.internal_readers > 0 + or entry.internal_writer ): continue self._uncommitted.discard(key) @@ -954,7 +1099,12 @@ def clear(self, force: bool = False) -> None: locked_count = 0 for key, entry in list(self._objects.items()): - if entry.write_lock.is_locked() or entry.read_lock.is_locked(): + if ( + entry.write_lock.is_locked() + or entry.read_lock.is_locked() + or entry.internal_readers + or entry.internal_writer + ): locked_count += 1 continue keys_to_clear.append(key) @@ -1032,6 +1182,20 @@ def num_objects(self) -> int: """Return the number of objects currently tracked in L1.""" return len(self._objects) + @l1_mgr_synchronized + def readable_key_sizes(self) -> dict[ObjectKey, int]: + """Snapshot committed RAM contents, including pinned readable pages.""" + return { + key: entry.memory_obj.get_size() + for key, entry in self._objects.items() + if entry.available_for_read() + } + + @l1_mgr_synchronized + def keys(self) -> list[ObjectKey]: + """Snapshot allocated keys for coordinated administrative reclamation.""" + return list(self._objects) + def is_key_evictable(self, key: ObjectKey) -> bool: """Check if a key is eligible for eviction (not locked). @@ -1049,7 +1213,12 @@ def is_key_evictable(self, key: ObjectKey) -> bool: entry = self._objects.get(key, None) if entry is None: return False - return not entry.read_lock.is_locked() and not entry.write_lock.is_locked() + return ( + not entry.read_lock.is_locked() + and not entry.write_lock.is_locked() + and not entry.internal_readers + and not entry.internal_writer + ) def get_memory_usage(self) -> tuple[int, int]: """Get the current memory usage of L1 cache. diff --git a/lmcache/v1/distributed/l2_adapters/base.py b/lmcache/v1/distributed/l2_adapters/base.py index 1e118821611..14ebadf2aea 100644 --- a/lmcache/v1/distributed/l2_adapters/base.py +++ b/lmcache/v1/distributed/l2_adapters/base.py @@ -668,6 +668,13 @@ def get_existing_key_sizes(self) -> Mapping[ObjectKey, int]: """ return _EMPTY_KEY_SIZES + def has_complete_inventory(self) -> bool: + """Whether a successful inventory snapshot proves absent keys missing. + + Remote backends without a complete inventory require lookup validation. + """ + return False + def _initialize_usage(self, key_sizes: Mapping[ObjectKey, int]) -> None: """Seed byte accounting before the adapter accepts operations. diff --git a/lmcache/v1/distributed/l2_adapters/fs_l2_adapter.py b/lmcache/v1/distributed/l2_adapters/fs_l2_adapter.py index 4d98535e30b..9235f0a48aa 100644 --- a/lmcache/v1/distributed/l2_adapters/fs_l2_adapter.py +++ b/lmcache/v1/distributed/l2_adapters/fs_l2_adapter.py @@ -12,6 +12,8 @@ from __future__ import annotations # Standard +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager from pathlib import Path from typing import TYPE_CHECKING, Optional, Union import asyncio @@ -486,6 +488,10 @@ def __init__(self, config: FSL2AdapterConfig): self._completed_lookup_tasks: dict[L2TaskId, Bitmap] = {} self._completed_load_tasks: dict[L2TaskId, Bitmap] = {} self._lock = threading.Lock() + self._close_lock = threading.Lock() + self._closing = False + self._closed = False + self._key_operations: dict[ObjectKey, tuple[asyncio.Lock, int]] = {} # Background asyncio event loop self._loop = asyncio.new_event_loop() @@ -525,12 +531,13 @@ def submit_store_task( objects: list[MemoryObj], ) -> L2TaskId: with self._lock: + if self._closing: + raise RuntimeError("filesystem adapter is closing") task_id = self._get_next_task_id() - - asyncio.run_coroutine_threadsafe( - self._execute_store(keys, objects, task_id), - self._loop, - ) + asyncio.run_coroutine_threadsafe( + self._execute_store(keys, objects, task_id), + self._loop, + ) return task_id def pop_completed_store_tasks( @@ -556,12 +563,13 @@ def submit_lookup_and_lock_task( self, keys: list[ObjectKey], group_layout_descs: dict[int, MemoryLayoutDesc] ) -> L2TaskId: with self._lock: + if self._closing: + raise RuntimeError("filesystem adapter is closing") task_id = self._get_next_task_id() - - asyncio.run_coroutine_threadsafe( - self._execute_lookup(keys, task_id), - self._loop, - ) + asyncio.run_coroutine_threadsafe( + self._execute_lookup(keys, task_id), + self._loop, + ) return task_id def query_lookup_and_lock_result(self, task_id: L2TaskId) -> Bitmap | None: @@ -583,12 +591,13 @@ def submit_load_task( objects: list[MemoryObj], ) -> L2TaskId: with self._lock: + if self._closing: + raise RuntimeError("filesystem adapter is closing") task_id = self._get_next_task_id() - - asyncio.run_coroutine_threadsafe( - self._execute_load(keys, objects, task_id), - self._loop, - ) + asyncio.run_coroutine_threadsafe( + self._execute_load(keys, objects, task_id), + self._loop, + ) return task_id def query_load_result(self, task_id: L2TaskId) -> Bitmap | None: @@ -609,6 +618,37 @@ def report_status(self) -> dict: "event_loop_alive": self._loop_thread.is_alive(), } + def has_complete_inventory(self) -> bool: + """The local filesystem can enumerate all persisted payloads.""" + return True + + def get_existing_key_sizes(self) -> dict[ObjectKey, int]: + """Inventory atomically published objects for checkpoint recovery. + + Called before the adapter is exposed to serving or eviction. Temporary + files and unrecognized layouts are excluded; external file corruption + still requires ordinary restore validation. + """ + inventory = {} + for path in self._base_path.rglob("*.data"): + if path.is_symlink() or not path.is_file(): + continue + relative = path.relative_to(self._base_path) + key = ( + _filename_to_object_key(path.name) + if len(relative.parts) == 1 + else _bounded_relative_path_to_object_key(relative) + ) + if key is None: + continue + try: + size = path.stat().st_size + except FileNotFoundError: + continue + if size > 0: + inventory[key] = size + return inventory + # ------------------------------------------------------------------ # Eviction Interface # ------------------------------------------------------------------ @@ -620,23 +660,20 @@ def delete(self, keys: list[ObjectKey]) -> None: keys: The object keys to delete. Note: - No per-key locks: a delete racing a load turns that load - into a miss; racing a store of the same key may leave the - key re-stored. + Same-key stores and deletes include their residency notification + in one ordered operation. Concurrent loads may observe a miss. """ if not keys: return - try: + with self._lock: + if self._closing: + raise RuntimeError("filesystem adapter is closing") fut = asyncio.run_coroutine_threadsafe( - self._execute_delete(keys), - self._loop, + self._execute_delete(keys), self._loop ) - deleted_keys, deleted_sizes = fut.result(timeout=30.0) - except Exception as e: - logger.warning("FSL2Adapter delete failed: %s", e) - return - if deleted_keys: - self._notify_keys_deleted(deleted_keys, deleted_sizes) + # The coroutine publishes deletion before releasing its per-key lock. + # Closure drains accepted operations rather than cancelling active I/O. + fut.result() # ``get_usage()`` is inherited from ``L2AdapterInterface``. The base # class maintains byte totals via ``_notify_keys_stored`` / @@ -650,32 +687,33 @@ def delete(self, keys: list[ObjectKey]) -> None: # ------------------------------------------------------------------ def close(self) -> None: - async def _stop_tasks(): + """Fence submissions and drain accepted I/O before releasing resources.""" + + async def drain() -> None: tasks = [ - t - for t in asyncio.all_tasks(self._loop) - if t is not asyncio.current_task() + task + for task in asyncio.all_tasks(self._loop) + if task is not asyncio.current_task() ] - for task in tasks: - task.cancel() if tasks: await asyncio.gather(*tasks, return_exceptions=True) - - if self._loop.is_running(): - fut = asyncio.run_coroutine_threadsafe(_stop_tasks(), self._loop) - try: - fut.result(timeout=5) - except Exception: - pass - self._loop.call_soon_threadsafe(self._loop.stop) - - self._loop_thread.join() - self._loop.close() - - self._store_efd.close() - self._lookup_efd.close() - self._load_efd.close() - logger.info("FSL2Adapter closed") + await self._loop.shutdown_default_executor() + + with self._close_lock: + if self._closed: + return + with self._lock: + self._closing = True + if self._loop.is_running(): + asyncio.run_coroutine_threadsafe(drain(), self._loop).result() + self._loop.call_soon_threadsafe(self._loop.stop) + self._loop_thread.join() + self._loop.close() + self._store_efd.close() + self._lookup_efd.close() + self._load_efd.close() + self._closed = True + logger.info("FSL2Adapter closed") # ------------------------------------------------------------------ # Internal helpers @@ -842,6 +880,21 @@ def _write_with_odirect( except OSError: pass + @asynccontextmanager + async def _key_operation(self, key: ObjectKey) -> AsyncIterator[None]: + """Order I/O and replica notifications; retain locks only while in use.""" + lock, users = self._key_operations.get(key, (asyncio.Lock(), 0)) + self._key_operations[key] = (lock, users + 1) + try: + async with lock: + yield + finally: + _, users = self._key_operations[key] + if users == 1: + del self._key_operations[key] + else: + self._key_operations[key] = (lock, users - 1) + # ---- store ---------------------------------------------------------- async def _execute_store( @@ -852,74 +905,74 @@ async def _execute_store( ) -> None: success = True bytes_written = 0 - stored_keys: list[ObjectKey] = [] - stored_sizes: list[int] = [] try: for key, obj in zip(keys, objects, strict=True): - # Skip if already stored on disk - if await self._existing_key_path(key) is not None: - continue - file_path, tmp_template = self._key_to_file_and_tmp_path(key) - buf = obj.byte_array - size = len(buf) - tmp_path = tmp_template.with_name( - f"{tmp_template.name}.{os.getpid()}.{uuid.uuid4().hex}" - ) - self._require_representable_path(tmp_path) - file_path.parent.mkdir(parents=True, exist_ok=True) - - try: - # Decide whether O_DIRECT is usable - do_odirect = self._use_odirect - if do_odirect: - aligned = self._os_disk_bs > 0 and size % self._os_disk_bs == 0 - if not aligned: - logger.warning( - "Cannot use O_DIRECT for " - "writing size %d, not " - "aligned to block size " - "%d.", - size, - self._os_disk_bs, + async with self._key_operation(key): + # Skip if already stored on disk + if await self._existing_key_path(key) is not None: + continue + file_path, tmp_template = self._key_to_file_and_tmp_path(key) + buf = obj.byte_array + size = len(buf) + tmp_path = tmp_template.with_name( + f"{tmp_template.name}.{os.getpid()}.{uuid.uuid4().hex}" + ) + self._require_representable_path(tmp_path) + file_path.parent.mkdir(parents=True, exist_ok=True) + + try: + # Decide whether O_DIRECT is usable + do_odirect = self._use_odirect + if do_odirect: + aligned = ( + self._os_disk_bs > 0 and size % self._os_disk_bs == 0 ) - do_odirect = False - - if do_odirect: - await self._loop.run_in_executor( - None, - self._write_with_odirect, - tmp_path, - buf, - ) - else: - async with aiofiles.open(tmp_path, "xb") as f: - await f.write(buf) + if not aligned: + logger.warning( + "Cannot use O_DIRECT for " + "writing size %d, not " + "aligned to block size " + "%d.", + size, + self._os_disk_bs, + ) + do_odirect = False + + if do_odirect: + await self._loop.run_in_executor( + None, + self._write_with_odirect, + tmp_path, + buf, + ) + else: + async with aiofiles.open(tmp_path, "xb") as f: + await f.write(buf) - published = await self._loop.run_in_executor( - None, _publish_temp_file, tmp_path, file_path - ) - if published: - bytes_written += size - stored_keys.append(key) - stored_sizes.append(size) - logger.debug( - "FSL2Adapter stored key %s (%d bytes)", - file_path.name, - size, + published = await self._loop.run_in_executor( + None, _publish_temp_file, tmp_path, file_path ) - else: - logger.debug( - "FSL2Adapter duplicate store kept existing key %s", - file_path.name, + if published: + bytes_written += size + self._notify_keys_stored([key], [size]) + logger.debug( + "FSL2Adapter stored key %s (%d bytes)", + file_path.name, + size, + ) + else: + logger.debug( + "FSL2Adapter duplicate store kept existing key %s", + file_path.name, + ) + except Exception: + logger.exception( + "FSL2Adapter failed to store %s", + file_path, ) - except Exception: - logger.exception( - "FSL2Adapter failed to store %s", - file_path, - ) - if await aiofiles.os.path.exists(tmp_path): - await aiofiles.os.unlink(tmp_path) - success = False + if await aiofiles.os.path.exists(tmp_path): + await aiofiles.os.unlink(tmp_path) + success = False except Exception: logger.exception( "FSL2Adapter store task %s failed", @@ -927,9 +980,6 @@ async def _execute_store( ) success = False - if stored_keys: - self._notify_keys_stored(stored_keys, stored_sizes) - with self._lock: self._completed_store_tasks[task_id] = L2StoreResult(success, bytes_written) self._store_efd.notify() @@ -1056,26 +1106,28 @@ async def _execute_delete( sem = asyncio.Semaphore(self._DELETE_CONCURRENCY) async def _delete_one(key: ObjectKey) -> tuple[ObjectKey, int] | None: - try: - file_path = await self._existing_key_path(key) - except ValueError: - logger.exception("FSL2Adapter rejected unrepresentable key %s", key) - return None - if file_path is None: - return None - async with sem: + async with self._key_operation(key): try: - size = (await aiofiles.os.stat(file_path)).st_size - await aiofiles.os.unlink(file_path) - except FileNotFoundError: + file_path = await self._existing_key_path(key) + except ValueError: + logger.exception("FSL2Adapter rejected unrepresentable key %s", key) return None - except Exception: - logger.exception( - "FSL2Adapter failed to delete %s", - file_path, - ) + if file_path is None: return None - return key, size + async with sem: + try: + size = (await aiofiles.os.stat(file_path)).st_size + await aiofiles.os.unlink(file_path) + except FileNotFoundError: + return None + except Exception: + logger.exception( + "FSL2Adapter failed to delete %s", + file_path, + ) + return None + self._notify_keys_deleted([key], [size]) + return key, size results = await asyncio.gather(*(_delete_one(k) for k in keys)) diff --git a/lmcache/v1/distributed/l2_adapters/native_connector_l2_adapter.py b/lmcache/v1/distributed/l2_adapters/native_connector_l2_adapter.py index 2ed9fb7d106..dc1da875331 100644 --- a/lmcache/v1/distributed/l2_adapters/native_connector_l2_adapter.py +++ b/lmcache/v1/distributed/l2_adapters/native_connector_l2_adapter.py @@ -136,6 +136,7 @@ def __init__( """ super().__init__(max_capacity_bytes=int(max_capacity_gb * (1024**3))) existing_key_sizes = dict(initial_key_sizes or {}) + self._inventory_complete = initial_key_sizes is not None self._initialize_usage(existing_key_sizes) self._client = native_client self._client_fd: int = int(native_client.event_fd()) @@ -170,6 +171,7 @@ def __init__( # Pending delete events for synchronous delete() calls self._pending_delete_events: dict[L2TaskId, threading.Event] = {} + self._delete_errors: dict[L2TaskId, str] = {} # Synchronous stores use private completions so StoreController cannot # consume another thread's result from ``_completed_stores``. @@ -191,6 +193,9 @@ def __init__( # Lock for all shared state above self._lock = threading.Lock() + self._close_lock = threading.Lock() + self._closing = False + self._closed = False # Background demux thread self._stop = threading.Event() @@ -201,6 +206,10 @@ def __init__( ) self._demux_thread.start() + def has_complete_inventory(self) -> bool: + """Only connectors supplied with a startup inventory prove absence.""" + return self._inventory_complete + def get_existing_key_sizes(self) -> Mapping[ObjectKey, int]: """Return a detached, immutable native-resident key snapshot.""" with self._lock: @@ -253,6 +262,8 @@ def submit_store_task( # with demux thread. The native submit is # non-blocking so holding the lock is brief. with self._lock: + if self._closing: + raise RuntimeError("native adapter is closing") task_id = self._get_next_task_id() if not self._try_reserve_store_bytes(reserved_bytes): self._completed_stores[task_id] = L2StoreResult(False, 0) @@ -325,6 +336,8 @@ def store_objects_sync( } with self._lock: + if self._closing: + raise RuntimeError("native adapter is closing") task_id = self._get_next_task_id() if not self._try_reserve_store_bytes(reserved_bytes): logger.warning( @@ -384,6 +397,8 @@ def submit_lookup_and_lock_task( key_strings = [_object_key_to_string(k) for k in keys] with self._lock: + if self._closing: + raise RuntimeError("native adapter is closing") task_id = self._get_next_task_id() future_id = int(self._client.submit_batch_exists(key_strings)) self._pending_ops[future_id] = ( @@ -422,6 +437,8 @@ def submit_load_task( memviews = [_obj_to_memoryview(obj) for obj in objects] with self._lock: + if self._closing: + raise RuntimeError("native adapter is closing") task_id = self._get_next_task_id() future_id = int(self._client.submit_batch_get(key_strings, memviews)) self._pending_ops[future_id] = ( @@ -445,7 +462,8 @@ def delete(self, keys: list[ObjectKey]) -> None: """Delete a batch of keys from the remote backend. Submits a batch delete to the native connector and blocks - until the demux thread signals completion (up to 30s timeout). + until the demux thread signals completion. A slow unlink retains its + lifecycle fence; a timeout cannot cancel an executing backend operation. Fires ``_notify_keys_deleted`` on success so eviction policy tracking stays in sync. @@ -459,6 +477,8 @@ def delete(self, keys: list[ObjectKey]) -> None: done_event = threading.Event() with self._lock: + if self._closing: + raise RuntimeError("native adapter is closing") # A key becomes lookup-locked only when the native EXISTS # completion is demultiplexed. Protect pending lookup keys as well, # otherwise eviction can delete a key between lookup submission and @@ -494,20 +514,16 @@ def delete(self, keys: list[ObjectKey]) -> None: self._pending_delete_events[task_id] = done_event # Block until demux thread signals completion - if not done_event.wait(timeout=30.0): - with self._lock: - self._pending_delete_events.pop(task_id, None) - # Note: _pending_ops entry may already be consumed - # by the demux thread; pop is safe either way. - for fid, entry in list(self._pending_ops.items()): - if entry[1] == task_id: - self._pending_ops.pop(fid, None) - break + while not done_event.wait(timeout=30.0): logger.warning( - "delete() timed out after 30s for %d keys", + "delete() still draining for %d keys; retaining deletion fence", len(delete_keys), ) - return + + with self._lock: + error = self._delete_errors.pop(task_id, None) + if error is not None: + raise RuntimeError(error) # ``_notify_keys_deleted`` is fired by the demux thread (with # accurate per-key sizes drawn from ``_key_sizes``) when the @@ -547,31 +563,47 @@ def report_status(self) -> dict[str, Any]: # --------------------------------------------------------------- def close(self) -> None: - self._stop.set() - self._demux_thread.join(timeout=5) - - with self._lock: - sync_events = [ - event for event, _result in self._pending_sync_store_events.values() - ] - self._pending_sync_store_events.clear() - abandoned_stores = list(self._pending_store_sizes.values()) - self._pending_store_sizes.clear() - self._pending_ops.clear() - # Keep ``abandoned_stores`` (and therefore every buffer owner) alive - # until the native client has stopped its worker threads. - self._client.close() - abandoned_reservations = sum( - pending.reserved_bytes for pending in abandoned_stores - ) - if abandoned_reservations: - self._release_store_reservation(abandoned_reservations) - for event in sync_events: - event.set() + """Drain native I/O before settling waiters and releasing buffer owners.""" + with self._close_lock: + if self._closed: + return + with self._lock: + self._closing = True + self._stop.set() + self._demux_thread.join() - self._store_efd.close() - self._lookup_efd.close() - self._load_efd.close() + with self._lock: + sync_events = [ + event for event, _result in self._pending_sync_store_events.values() + ] + abandoned_stores = list(self._pending_store_sizes.values()) + pending_deletes = dict(self._pending_delete_events) + # Keep buffer owners and deletion fences until every syscall drains. + # If native close fails, do not claim that releasing them is safe. + self._client.close() + with self._lock: + self._pending_sync_store_events.clear() + self._pending_store_sizes.clear() + self._pending_ops.clear() + self._pending_delete_events.clear() + for task_id in pending_deletes: + self._delete_errors[task_id] = ( + "native adapter closed after I/O drain; " + "delete outcome unavailable" + ) + abandoned_reservations = sum( + pending.reserved_bytes for pending in abandoned_stores + ) + if abandoned_reservations: + self._release_store_reservation(abandoned_reservations) + for event in sync_events: + event.set() + for event in pending_deletes.values(): + event.set() + self._store_efd.close() + self._lookup_efd.close() + self._load_efd.close() + self._closed = True # --------------------------------------------------------------- # Internal helpers @@ -617,6 +649,7 @@ def _demux_loop(self) -> None: sizes_deleted: list[int] = [] settled_store_reservations = 0 async_store_completed = False + async_store_results: dict[int, L2StoreResult] = {} # Events for synchronous ``delete()`` callers. We set these # AFTER firing ``_notify_keys_deleted`` below so that when # the caller unblocks and calls ``get_usage()``, the base @@ -689,7 +722,7 @@ def _demux_loop(self) -> None: keys_stored.append(key) sizes_stored.append(0) if op_type == self._OP_STORE: - self._completed_stores[task_id] = L2StoreResult( + async_store_results[task_id] = L2StoreResult( completion_ok, task_bytes ) async_store_completed = True @@ -736,6 +769,10 @@ def _demux_loop(self) -> None: self._load_efd.notify() elif op_type == self._OP_DELETE: + if not ok: + self._delete_errors[task_id] = str( + error or "native delete failed" + ) if result_bools is not None and lookup_keys is not None: for i, deleted in enumerate(result_bools): if not deleted: @@ -758,8 +795,6 @@ def _demux_loop(self) -> None: keys_stored, sizes_stored, ) - if async_store_completed: - self._store_efd.notify() for event in sync_store_done_events: event.set() @@ -772,6 +807,12 @@ def _demux_loop(self) -> None: self._notify_keys_accessed(keys_accessed) if keys_deleted: self._notify_keys_deleted(keys_deleted, sizes_deleted) + # Async consumers may release source ownership on completion. Make + # replica events visible first, including partial store successes. + if async_store_completed: + with self._lock: + self._completed_stores.update(async_store_results) + self._store_efd.notify() # Unblock any synchronous ``delete()`` callers only AFTER # ``_notify_keys_deleted`` has updated the base class byte # accounting, so ``get_usage()`` never briefly reports stale diff --git a/lmcache/v1/distributed/storage_controllers/checkpoint_evict_store_policy.py b/lmcache/v1/distributed/storage_controllers/checkpoint_evict_store_policy.py index 6f8a6d47aca..87b20d61a73 100644 --- a/lmcache/v1/distributed/storage_controllers/checkpoint_evict_store_policy.py +++ b/lmcache/v1/distributed/storage_controllers/checkpoint_evict_store_policy.py @@ -7,14 +7,15 @@ states that the next turn of the same conversation supersedes within seconds, and pushes still-useful checkpoints out of L2 by capacity. -This policy keeps new checkpoint pages in L1 only. When L1 is about to evict -a page that is still current (not superseded by a newer checkpoint of the -same conversation) and not yet in L2, the page is written to L2 first and -evicted after the write completes. A conversation that stays active never -writes its checkpoints; one that goes idle writes its latest checkpoint once. -Superseded pages are evicted without a write. Ordinary KV chunks keep -write-through behavior. On a clean shutdown, current checkpoint pages still -in L1 are written to L2 within the configured time budget. +This policy keeps new checkpoint pages in L1 until memory pressure or shutdown. +Retained pages without an L2 copy are offered for bounded asynchronous +persistence before normal eviction. Supersession guides victim selection but +does not prove an older branch is dead. Sustained capacity pressure may retire +cold checkpoints coherently before reclaiming their last pages. + +Ordinary KV chunks keep write-through behavior. Clean shutdown persists retained +pages within its time budget; unfinished checkpoints may be lost. Write savings +depend on retirement before persistence, not merely conversation activity. Select it with ``--l2-store-policy checkpoint_on_evict``. """ diff --git a/lmcache/v1/distributed/storage_controllers/eviction_controller.py b/lmcache/v1/distributed/storage_controllers/eviction_controller.py index 61befdda3c3..c54eb3a453e 100644 --- a/lmcache/v1/distributed/storage_controllers/eviction_controller.py +++ b/lmcache/v1/distributed/storage_controllers/eviction_controller.py @@ -138,6 +138,7 @@ def __init__( self._event_bus = get_event_bus() self._last_extra_log = time.monotonic() self._immediate_request = threading.Event() + self._pressure_since: float | None = None self._write_back_enabled = eviction_config.write_back_on_evict self._periodic_backup_enabled = eviction_config.periodic_flush_interval > 0 @@ -209,7 +210,23 @@ def eligible(key: ObjectKey) -> bool: def _request_persist(self, to_persist: list[ObjectKey]) -> None: if self._retention is not None and to_persist: - self._retention.request_persist(list(dict.fromkeys(to_persist))) + keys = list(dict.fromkeys(to_persist)) + retained = [key for key in keys if self._retention.is_retained(key)] + orphaned = [key for key in keys if key not in retained] + if orphaned: + self._delete_l1(orphaned) + self._retention.request_persist(retained) + + def _delete_l1(self, keys: list[ObjectKey]) -> dict[ObjectKey, L1Error]: + if self._retention is not None: + result = self._retention.evict(keys, "l1", self._l1_manager.delete) + else: + result = self._l1_manager.delete(keys) + if any(error == L1Error.SUCCESS for error in result.values()): + self._reclamation_generation = ( + getattr(self, "_reclamation_generation", 0) + 1 + ) + return result def _drop_superseded(self) -> int: """Evict superseded checkpoint pages from L1 before any LRU victim.""" @@ -223,15 +240,36 @@ def _drop_superseded(self) -> int: ][:_SUPERSEDED_DROP_BATCH] if not victims: return 0 - result = self._l1_manager.delete(victims) + to_persist = [ + key for key in victims if retention.needs_persist_before_evict(key) + ] + self._request_persist(to_persist) + result = self._delete_l1([key for key in victims if key not in to_persist]) dropped = sum(1 for error in result.values() if error == L1Error.SUCCESS) retention.record_l1_superseded_drops(dropped) return dropped def request_immediate_eviction(self) -> None: """Wake the eviction loop for a capacity-blocked store.""" + now = time.monotonic() + if ( + getattr(self, "_pressure_since", None) is None + or now - getattr(self, "_last_pressure_request", 0.0) > 1.0 + ): + self._pressure_since = now + self._last_pressure_request = now self._immediate_request.set() + def set_pressure_grace(self, seconds: float) -> None: + """Allocate part of the existing admission deadline to persistence.""" + self._pressure_grace = max(0.0, seconds) + self._pressure_since = None + + @property + def trigger_watermark(self) -> float: + """L1 pressure threshold used to reserve publication headroom.""" + return self._eviction_config.trigger_watermark + def stop(self) -> None: """Stop promptly even when the eviction loop is waiting.""" self._stop_flag.set() @@ -428,6 +466,7 @@ def eviction_loop(self): logger.exception("L1 eviction pass failed; retrying next pass") def _run_eviction_pass(self, immediate: bool) -> None: + reclamation = getattr(self, "_reclamation_generation", 0) self._l1_manager.reclaim_abandoned_writes() watermark = self._eviction_config.trigger_watermark eviction_ratio = self._eviction_config.eviction_ratio @@ -436,7 +475,8 @@ def _run_eviction_pass(self, immediate: bool) -> None: if self._eviction_config.extra_logging_enabled: self._maybe_log_memory_usage(used_bytes, total_bytes) usage = 0 if total_bytes == 0 else used_bytes / total_bytes - if usage < watermark: + if usage < watermark and not immediate: + self._pressure_since = None now = time.monotonic() if ( backup_interval > 0 @@ -469,6 +509,7 @@ def _run_eviction_pass(self, immediate: bool) -> None: if self._drop_superseded(): used_bytes, total_bytes = self._l1_manager.get_memory_usage() if total_bytes and used_bytes / total_bytes < watermark: + self._pressure_since = None self._publish_triggered(usage, watermark) return to_persist: list[ObjectKey] = [] @@ -479,6 +520,21 @@ def _run_eviction_pass(self, immediate: bool) -> None: for action in actions: self.execute_eviction_action(action) self._request_persist(to_persist) + if getattr(self, "_reclamation_generation", 0) != reclamation: + self._pressure_since = None + pressure_since = getattr(self, "_pressure_since", None) + if ( + immediate + and self._retention is not None + and pressure_since is not None + and time.monotonic() - pressure_since + >= getattr(self, "_pressure_grace", 1.0) + ): + retired_bytes = self._retention.retire_under_pressure( + max(1, int(total_bytes * eviction_ratio)) + ) + if retired_bytes: + self._pressure_since = None self._publish_triggered(usage, watermark) def execute_eviction_action(self, action: EvictionAction): @@ -488,13 +544,13 @@ def execute_eviction_action(self, action: EvictionAction): else: logger.error("L2 eviction destination requires writeback") logger.error("Treating it as DISCARD.") - self._l1_manager.delete(action.keys) + self._delete_l1(action.keys) elif action.destination == EvictionDestination.DISCARD: - self._l1_manager.delete(action.keys) + self._delete_l1(action.keys) else: logger.error("Unsupported eviction destination: %s", action.destination) logger.error("Treating it as DISCARD.") - self._l1_manager.delete(action.keys) + self._delete_l1(action.keys) def emergency_evict_bytes( self, @@ -595,6 +651,10 @@ def emergency_evict_bytes( break self._request_persist(to_persist) + if free < target_free_bytes: + # Continue through the same bounded pressure policy used by + # store admission while the prefetch caller retries for room. + self.request_immediate_eviction() evicted_keys = max(0, initial_objects - self._l1_manager.num_objects()) logger.info( @@ -750,7 +810,7 @@ def _flush_one_l2_batch_then_delete( ) return _SyncFlushResult.FAILURE - delete_result = self._l1_manager.delete(readable_keys) + delete_result = self._delete_l1(readable_keys) not_deleted = [ key for key, error in delete_result.items() if error != L1Error.SUCCESS ] @@ -1099,12 +1159,23 @@ def _check_and_evict_by_cache_salt(self, state: L2AdapterEvictionState): def _execute_eviction_action( self, adapter: L2AdapterInterface, action: EvictionAction ): + def delete(keys: list[ObjectKey]) -> None: + if self._retention is None: + adapter.delete(keys) + return + adapter_id = next( + state.adapter_id + for state in self._adapter_states + if state.adapter is adapter + ) + self._retention.evict(keys, adapter_id, adapter.delete) + if action.destination == EvictionDestination.DISCARD: - adapter.delete(action.keys) + delete(action.keys) else: logger.error("Unsupported eviction destination: %s", action.destination) logger.error("Treating it as DISCARD.") - adapter.delete(action.keys) + delete(action.keys) if action.keys: get_event_bus().publish( diff --git a/lmcache/v1/distributed/storage_controllers/prefetch_controller.py b/lmcache/v1/distributed/storage_controllers/prefetch_controller.py index 54451854ca2..c489365a54a 100644 --- a/lmcache/v1/distributed/storage_controllers/prefetch_controller.py +++ b/lmcache/v1/distributed/storage_controllers/prefetch_controller.py @@ -53,7 +53,7 @@ from dataclasses import dataclass, field from itertools import groupby from operator import attrgetter -from typing import TYPE_CHECKING, Iterable +from typing import TYPE_CHECKING, Callable, Iterable import enum import select import threading @@ -193,6 +193,19 @@ class PrefetchPhase(enum.Enum): PLAN_AND_LOAD = enum.auto() +@dataclass(frozen=True) +class PrefetchResult: + """Consumed prefetch outcome; reservation failure is not disk absence. + + found indexes the requested keys. reservation_failed records allocation + failure or contention observed while reserving the actual load buffers. + A zero found bit alone does not distinguish absence from backend failure. + """ + + found: Bitmap + reservation_failed: bool = False + + @dataclass class InFlightPrefetchRequest: """Tracks a single prefetch request across its lifecycle phases.""" @@ -225,6 +238,7 @@ class InFlightPrefetchRequest: l2_adapter2readlocks: dict[int, Bitmap] = field(default_factory=dict) # True once the prefix hit was stored/published. hit_reported: bool = False + reservation_failed: bool = False # Load phase: adapter_idx -> bitmap of key indices to load load_plan: dict[int, Bitmap] = field(default_factory=dict) @@ -273,6 +287,8 @@ class PrefetchController(StorageControllerInterface): adapter_descriptors: Descriptors for each L2 adapter (same order). policy: The prefetch policy for load plan decisions. max_in_flight: Maximum number of concurrent prefetch requests. + replica_may_match: Optional predicate rejecting known payload-size + mismatches by adapter ID, key and expected logical bytes. """ # Singleton dispatch for the in-flight load gauges: tests may construct @@ -289,6 +305,7 @@ def __init__( adapter_descriptors: list[AdapterDescriptor], policy: PrefetchPolicy, max_in_flight: int = 8, + replica_may_match: Callable[[int, ObjectKey, int], bool] | None = None, ) -> None: self._l1_manager = l1_manager self._l2_adapters: dict[int, L2AdapterInterface] = { @@ -300,6 +317,7 @@ def __init__( } self._policy = policy self._max_in_flight = max_in_flight + self._replica_may_match = replica_may_match self._l1_eviction_controller: L1EvictionController | None = None # Adapters that are being drained and will be removed after all @@ -338,6 +356,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 +507,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 bitmap and allocation evidence atomically; None means pending. + + This shares consumption with query_prefetch_result; both discard lookup + bookkeeping exactly once. No missing bitmap bit proves disk absence. + """ 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 @@ -608,16 +641,26 @@ def stop(self) -> None: """ Signal the loop to stop and wait for the thread to join. - Cleans up any in-flight requests (releases L1 write locks, - L2 locks) before returning. + Releases lookup locks, but active load destinations stay owned until + adapters drain and release_stopped_ownership is called. """ self._stop_flag.set() self._submission_efd.notify() self._thread.join() + self._stopped_write_keys = [ + key + for request in self._in_flight_requests.values() + for key in request.write_reserved_keys + ] self._cleanup_in_flight_requests() self._submission_efd.close() self._adapter_ctrl_efd.close() + def release_stopped_ownership(self) -> None: + """Release unpublished load destinations after all adapters have drained.""" + self._discard_unloaded_buffers(self._stopped_write_keys) + self._stopped_write_keys = [] + def add_adapter( self, adapter_id: int, @@ -1102,6 +1145,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: @@ -1215,6 +1259,7 @@ def _reserve_grouped_buffers( is_temporary=[not retention_map[key] for key in group_keys], layout_desc=request.group_layout_descs[group_id], mode="new", + internal=True, ) ) return results @@ -1385,8 +1430,22 @@ def _poll_lookup_results( ) if result is None: continue - request.lookup_results[adapter_idx] = result request.l2_adapter2readlocks[adapter_idx] = result + if self._replica_may_match is None: + request.lookup_results[adapter_idx] = result + else: + eligible = Bitmap(len(request.keys)) + sizes = { + group: get_size_bytes(layout.shapes, layout.dtypes) + for group, layout in request.group_layout_descs.items() + } + for i, key in enumerate(request.keys): + if result.test(i) and self._replica_may_match( + adapter_idx, key, sizes[key.object_group_id] + ): + eligible.set(i) + # Preserve the original locks: rejected hits still need unlock. + request.lookup_results[adapter_idx] = eligible del request.pending_lookup_tasks[adapter_idx] def _poll_load_results( @@ -1600,6 +1659,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() @@ -1638,9 +1700,8 @@ def _cleanup_in_flight_requests(self) -> None: """Release resources for any in-flight requests during shutdown.""" l1_mgr = self._l1_manager for request in self._in_flight_requests.values(): - if request.phase == PrefetchPhase.PLAN_AND_LOAD: - if request.write_reserved_keys: - self._discard_unloaded_buffers(request.write_reserved_keys) + # An executing backend load can still write these reservations. + # release_stopped_ownership discards them after adapter teardown. self._release_l2_locks(request, keep={}) if request.l1_readlocks.popcount() > 0: l1_mgr.finish_read( diff --git a/lmcache/v1/distributed/storage_controllers/store_controller.py b/lmcache/v1/distributed/storage_controllers/store_controller.py index 2b51517313f..e32e6886050 100644 --- a/lmcache/v1/distributed/storage_controllers/store_controller.py +++ b/lmcache/v1/distributed/storage_controllers/store_controller.py @@ -11,14 +11,17 @@ # Standard from collections import Counter, OrderedDict, defaultdict +from collections.abc import Callable from dataclasses import dataclass import enum import select import threading +import time # First Party from lmcache.logging import init_logger -from lmcache.v1.distributed.api import ObjectKey +from lmcache.v1.distributed.api import ObjectKey, is_recurrent_checkpoint_key +from lmcache.v1.distributed.checkpoint_retention import CheckpointRetention from lmcache.v1.distributed.error import L1Error from lmcache.v1.distributed.internal_api import L1ManagerListener, L2AdapterListener from lmcache.v1.distributed.l1_manager import L1Manager @@ -88,13 +91,121 @@ class StoreListener(L1ManagerListener, L2AdapterListener): so the controller stores a reused key once rather than on every reuse. """ - def __init__(self) -> None: + def __init__( + self, + *, + budget_bytes: int = 0, + select_new: Callable[[list[ObjectKey]], set[ObjectKey]] | None = None, + ) -> None: self._pending_keys: list[ObjectKey] = [] self._reused_keys: list[ObjectKey] = [] self._residency_events: list[tuple[L2Residency, list[ObjectKey]]] = [] self._tracks_l2_residency = False self._lock = threading.Lock() self._event_fd = create_event_notifier() + self._budget_bytes = budget_bytes + self._select_new = select_new + self._owned: dict[ObjectKey, tuple[int, bool]] = {} + self._owned_bytes = 0 + self._rejected = 0 + self._active: set[ObjectKey] = set() + + def on_l1_committed_objects( + self, keys: list[ObjectKey], sizes: list[int] + ) -> list[ObjectKey]: + if self._select_new is None: + self.on_l1_keys_write_finished(keys) + return [] + selected = self._select_new(keys) + pairs = [ + (key, size) + for key, size in zip(keys, sizes, strict=True) + if key in selected + ] + return self.claim( + [key for key, _ in pairs], [size for _, size in pairs], reuse=False + ) + + def claim( + self, keys: list[ObjectKey], sizes: list[int], *, reuse: bool = True + ) -> list[ObjectKey]: + """Accept a byte-bounded batch during L1's atomic ownership handoff.""" + accepted = [] + with self._lock: + for key, size in zip(keys, sizes, strict=True): + if key in self._owned: + continue + if self._owned_bytes + size > self._budget_bytes: + self._rejected += 1 + continue + self._owned[key] = (size, reuse) + self._owned_bytes += size + accepted.append(key) + (self._reused_keys if reuse else self._pending_keys).extend(accepted) + if accepted: + self._event_fd.notify() + return accepted + + def release_owned(self, keys: list[ObjectKey]) -> list[ObjectKey]: + """Forget completed/cancelled queue ownership; caller releases L1 pins.""" + released = [] + with self._lock: + for key in keys: + item = self._owned.pop(key, None) + if item is not None: + self._owned_bytes -= item[0] + released.append(key) + return released + + def ownership(self) -> dict[ObjectKey, tuple[int, bool]]: + """Snapshot of bounded work, including its unique source byte sizes.""" + with self._lock: + return dict(self._owned) + + def owns(self, key: ObjectKey) -> bool: + """Return whether queued or active work owns this source page.""" + with self._lock: + return key in self._owned + + def queued(self, key: ObjectKey) -> bool: + """Whether ownership can still be cancelled before backend submission.""" + with self._lock: + return key in self._owned and key not in self._active + + def begin_submission(self, keys: list[ObjectKey]) -> list[ObjectKey]: + """Atomically claim queued pages against cancellation.""" + with self._lock: + claimed = [ + key + for key in dict.fromkeys(keys) + if key in self._owned and key not in self._active + ] + self._active.update(claimed) + return claimed + + def end_submission(self, keys: list[ObjectKey]) -> None: + """Return drained pages to cancelable queue ownership.""" + with self._lock: + self._active.difference_update(keys) + + def cancel_queued(self, keys: list[ObjectKey]) -> list[ObjectKey]: + """Cancel unsubmitted work; caller releases only the returned pins.""" + with self._lock: + cancelled = [ + key for key in keys if key in self._owned and key not in self._active + ] + for key in cancelled: + size, _ = self._owned.pop(key) + self._owned_bytes -= size + return cancelled + + def report_ownership(self) -> dict[str, int]: + """Return byte admission accounting without copying the work queue.""" + with self._lock: + return { + "owned_source_bytes": self._owned_bytes, + "rejected_queue_pages": self._rejected, + } def get_event_fd(self) -> int: """ @@ -313,6 +424,8 @@ def __init__( l2_adapters: list[L2AdapterInterface], adapter_descriptors: list[AdapterDescriptor], policy: StorePolicy, + retention: CheckpointRetention | None = None, + max_pending_bytes: int | None = None, ) -> None: self._l1_manager = l1_manager self._l2_adapters: dict[int, L2AdapterInterface] = { @@ -323,6 +436,18 @@ def __init__( desc.index: desc for desc in adapter_descriptors } self._policy = policy + self._retention = retention + _, capacity = l1_manager.get_memory_usage() + self._max_pending_bytes = ( + max_pending_bytes + if max_pending_bytes is not None + else min(capacity, max(4096, capacity // 4)) + ) + self._retry: dict[ObjectKey, tuple[float, int]] = {} + self._failures: dict[ObjectKey, int] = {} + self._completed_destinations: dict[ObjectKey, set[int]] = {} + self._destination_attempts: dict[ObjectKey, dict[int, int]] = {} + self._accepting = True # Reuse admission state, touched only by the store loop thread. self._tracks_l2_residency = policy.uses_reuse_admission() @@ -339,7 +464,9 @@ def __init__( self._pending_adapter_ops: list[AddAdapterOp | RemoveAdapterOp] = [] self._adapter_ctrl_efd = create_event_notifier() - self._listener = StoreListener() + self._listener = StoreListener( + budget_bytes=self._max_pending_bytes, select_new=self._select_new_keys + ) self._l1_manager.register_listener(self._listener) if self._tracks_l2_residency: self._listener.enable_l2_residency_tracking() @@ -351,6 +478,7 @@ def __init__( # Composite key is needed because task IDs are only unique # within a single adapter, not across adapters. self._in_flight_tasks: dict[tuple[int, L2TaskId], InFlightStoreTask] = {} + self._active_keys: Counter[ObjectKey] = Counter() # Shadow counter for status reporting (updated in background loop) self._status_in_flight_count: int = 0 @@ -402,15 +530,19 @@ def stop(self) -> None: """ Signal the loop to stop, wait for the thread to join. - Releases all in-flight read locks on shutdown so that - L1 objects are not permanently locked. + Queued work releases its ownership; active backend buffers remain + pinned until adapters drain and release_stopped_ownership is called. """ + self._accepting = False self._stop_flag.set() # Wake up the poll loop so it can exit promptly self._listener.notify() self._thread.join() - self._cleanup_in_flight_tasks() - self._listener.close() + # Active backend operations still own their source buffers. The storage + # manager closes/drains adapters before release_stopped_ownership(). + self._discard_owned( + [key for key in self._listener.ownership() if not self._key_in_flight(key)] + ) self._adapter_ctrl_efd.close() def report_status(self) -> dict: @@ -426,9 +558,11 @@ def report_status(self) -> dict: "num_l2_adapters": len(self._l2_adapters), "num_active_adapters": len(self._l2_adapters) - num_draining, "num_draining_adapters": num_draining, + **self._listener.report_ownership(), + "max_pending_bytes": self._max_pending_bytes, } - def submit_reused_keys(self, keys: list[ObjectKey]) -> None: + def submit_reused_keys(self, keys: list[ObjectKey]) -> list[ObjectKey]: """ Report keys that a completed restore read from L1. @@ -440,8 +574,43 @@ def submit_reused_keys(self, keys: list[ObjectKey]) -> None: keys: Keys whose restore, including the consumer's copy, has completed. """ - if keys: - self._listener.enqueue_reused_keys(keys) + if keys and self._accepting and self._tracks_l2_residency: + return self._l1_manager.handoff_internal_reads(keys, self._listener.claim) + return [] + + def release_stopped_ownership(self) -> None: + """Release pins only after every backend has drained or closed.""" + self._cleanup_in_flight_tasks() + self._discard_owned(list(self._listener.ownership())) + self._listener.close() + + def submit_checkpoint_keys(self, keys: list[ObjectKey]) -> list[ObjectKey]: + """Queue explicit persistence through the policy's existing write path.""" + if not self._accepting: + return [] + if self._tracks_l2_residency: + return self.submit_reused_keys(keys) + return self._l1_manager.handoff_internal_reads( + keys, + lambda available, sizes: self._listener.claim( + available, sizes, reuse=False + ), + ) + + def has_pending_work(self) -> bool: + """Whether queued, retrying or active work still owns source pages.""" + return bool(self._listener.ownership() or self._status_in_flight_count) + + def owns_queued(self, key: ObjectKey) -> bool: + """Report cancelable source ownership without releasing active I/O.""" + return self._listener.queued(key) + + def cancel_queued(self, keys: list[ObjectKey]) -> None: + """Cancel retired work atomically against backend submission.""" + released = self._listener.cancel_queued(keys) + if released: + self._l1_manager.release_internal_reads(released) + self._listener.notify() def add_adapter( self, @@ -557,6 +726,7 @@ def _store_loop(self) -> None: while not self._stop_flag.is_set(): # First, apply runtime add/remove of the L2 adapters. self._apply_pending_adapter_ops(poller) + self._retry_ready() ready = poller.poll(STORE_LOOP_POLL_TIMEOUT_MS) @@ -679,7 +849,10 @@ def _process_new_keys(self, keys: list[ObjectKey]) -> None: keys (list[ObjectKey]): Keys that finished writing to L1. """ - for group in _group_keys_by_shape(keys).values(): + owned = self._listener.ownership() + for group in _group_keys_by_shape( + [key for key in keys if key in owned] + ).values(): self._submit_store_for_single_shape(group) def _process_reuse_events(self) -> None: @@ -691,14 +864,33 @@ def _process_reuse_events(self) -> None: else: for key in keys: self._l2_resident.pop(key, None) + owned = self._listener.ownership() fresh = [ key for key in dict.fromkeys(reused) - if key not in self._l2_resident and key not in self._reuse_in_flight + if (self._retention is not None or key not in self._l2_resident) + and key not in self._reuse_in_flight + and key in owned ] for group in _group_keys_by_shape(fresh).values(): plan = self._policy.select_reuse_targets(group, self._routing_descriptors()) - self._reuse_in_flight.update(self._submit_plan(plan)) + submitted = set(self._submit_plan(plan)) + self._reuse_in_flight.update(submitted) + self._discard_owned( + [ + key + for key in group + if key not in submitted and key not in self._retry + ] + ) + fresh_keys = set(fresh) + self._discard_owned( + [ + key + for key in reused + if key not in fresh_keys and not self._key_in_flight(key) + ] + ) def _routing_descriptors(self) -> list[AdapterDescriptor]: """Return descriptors of live adapters that may receive new stores.""" @@ -713,14 +905,48 @@ def _routing_descriptors(self) -> list[AdapterDescriptor]: def _submit_store_for_single_shape(self, keys: list[ObjectKey]) -> None: """Submit ``keys`` (all same shape) to their target adapters.""" plan = self._policy.select_store_targets(keys, self._routing_descriptors()) - self._submit_plan(plan) + submitted = set(self._submit_plan(plan)) + self._discard_owned( + [key for key in keys if key not in submitted and key not in self._retry] + ) def _submit_plan(self, plan: dict[int, list[ObjectKey]]) -> list[ObjectKey]: """Submit one store task per adapter in ``plan``; return submitted keys.""" l1_mgr = self._l1_manager submitted: list[ObjectKey] = [] + plan = {adapter: list(dict.fromkeys(keys)) for adapter, keys in plan.items()} + + # Fairness is per page: a failing first adapter must not starve a + # healthy replica when active destination bytes are serialized. + choices: dict[ObjectKey, int] = {} + for adapter_index, keys in plan.items(): + for key in keys: + if adapter_index in self._completed_destinations.get(key, ()): + continue + if ( + self._retention is not None + and adapter_index in self._retention.resident_adapters(key) + ): + continue + attempts = self._destination_attempts.get(key, {}) + previous = choices.get(key) + if previous is None or attempts.get(adapter_index, 0) < attempts.get( + previous, 0 + ): + choices[key] = adapter_index for adapter_index, target_keys in plan.items(): + if not target_keys: + continue + deferred = [ + key + for key in target_keys + if key in choices and choices[key] != adapter_index + ] + self._schedule_retry(deferred) + target_keys = [ + key for key in target_keys if choices.get(key) == adapter_index + ] if not target_keys: continue @@ -733,7 +959,28 @@ def _submit_plan(self, plan: dict[int, list[ObjectKey]]) -> list[ObjectKey]: continue # Reserve read to get MemoryObj references and hold read locks - read_results = l1_mgr.reserve_read(target_keys) + if self._retention is not None: + target_keys = [ + key + for key in target_keys + if adapter_index not in self._retention.resident_adapters(key) + ] + target_keys = [ + key + for key in target_keys + if adapter_index not in self._completed_destinations.get(key, ()) + ] + if not target_keys: + continue + # One destination at a time keeps active I/O bytes within the + # same source-byte budget even with multiple replication targets. + active = [key for key in target_keys if self._key_in_flight(key)] + self._schedule_retry(active) + target_keys = [key for key in target_keys if not self._key_in_flight(key)] + if not target_keys: + continue + target_keys = self._listener.begin_submission(target_keys) + read_results = l1_mgr.reserve_read(target_keys, internal=True) successful_keys = [] successful_objs = [] @@ -786,10 +1033,25 @@ def _submit_plan(self, plan: dict[int, list[ObjectKey]]) -> list[ObjectKey]: ) if not successful_keys: + self._listener.end_submission(target_keys) continue + self._listener.end_submission(list(set(target_keys) - set(successful_keys))) + adapter = self._l2_adapters[adapter_index] - task_id = adapter.submit_store_task(successful_keys, successful_objs) + for key in successful_keys: + attempts = self._destination_attempts.setdefault(key, {}) + attempts[adapter_index] = attempts.get(adapter_index, 0) + 1 + try: + task_id = adapter.submit_store_task(successful_keys, successful_objs) + except Exception: + l1_mgr.release_internal_reads(successful_keys) + self._listener.end_submission(successful_keys) + self._schedule_retry(successful_keys) + logger.exception( + "L2 store submission failed for adapter %d", adapter_index + ) + continue submitted.extend(successful_keys) self._in_flight_tasks[(adapter_index, task_id)] = InFlightStoreTask( @@ -798,6 +1060,7 @@ def _submit_plan(self, plan: dict[int, list[ObjectKey]]) -> list[ObjectKey]: read_locked_keys=list(successful_keys), ) self._status_in_flight_count += 1 + self._active_keys.update(successful_keys) # All objects for a single store task share one layout (L1 # allocates uniform MemoryObjs per chunk), so total bytes is @@ -868,9 +1131,14 @@ def _finalize_store( l1_mgr = self._l1_manager success = task.l2_store_result - l1_mgr.finish_read(task.read_locked_keys) + l1_mgr.release_internal_reads(task.read_locked_keys) + self._listener.end_submission(task.read_locked_keys) del self._in_flight_tasks[task_key] self._status_in_flight_count -= 1 + for key in task.keys: + self._active_keys[key] -= 1 + if self._active_keys[key] == 0: + del self._active_keys[key] self._reuse_in_flight.difference_update(task.keys) l2_name = self._adapter_descriptors[adapter_index].type_name @@ -880,7 +1148,15 @@ def _finalize_store( "l2_name": l2_name, "bytes_transferred": task.l2_bytes_transferred, } + delete_keys = [] if success: + for key in task.keys: + # Checkpoint replicas have exact deletion-aware accounting. + # Ordinary chunks only need progress across this bounded batch. + if self._retention is None or not is_recurrent_checkpoint_key(key): + self._completed_destinations.setdefault(key, set()).add( + adapter_index + ) self._event_bus.publish( Event( event_type=EventType.L2_STORE_COMPLETED, @@ -898,11 +1174,9 @@ def _finalize_store( adapter_index, len(task.keys), ) - if self._tracks_l2_residency: + if self._tracks_l2_residency and self._retention is None: self._remember_l2_resident(task.keys) delete_keys = self._policy.select_l1_deletions(task.keys) - if delete_keys: - l1_mgr.delete(delete_keys) else: self._event_bus.publish( Event( @@ -920,6 +1194,19 @@ def _finalize_store( adapter_index, task.keys, ) + self._schedule_retry(task.keys) + self._discard_owned( + [ + key + for key in task.keys + if not self._key_in_flight(key) and key not in self._retry + ] + ) + if delete_keys: + if self._retention is not None: + self._retention.evict(delete_keys, "l1", l1_mgr.delete) + else: + l1_mgr.delete(delete_keys) def _remember_l2_resident(self, keys: list[ObjectKey]) -> None: """Record keys as present in L2, forgetting the oldest beyond the bound.""" @@ -942,5 +1229,84 @@ def _cleanup_in_flight_tasks(self) -> None: adapter_index, len(task.read_locked_keys), ) - l1_mgr.finish_read(task.read_locked_keys) + l1_mgr.release_internal_reads(task.read_locked_keys) + self._listener.end_submission(task.read_locked_keys) self._in_flight_tasks.clear() + self._active_keys.clear() + self._status_in_flight_count = 0 + + def _select_new_keys(self, keys: list[ObjectKey]) -> set[ObjectKey]: + if not self._accepting: + return set() + plan = self._policy.select_store_targets( + keys, list(self._adapter_descriptors.values()) + ) + return {key for group in plan.values() for key in group} + + def _key_in_flight(self, key: ObjectKey) -> bool: + return key in self._active_keys + + def _discard_owned(self, keys: list[ObjectKey]) -> None: + released = self._listener.release_owned(keys) + if released: + self._l1_manager.release_internal_reads(released) + if self._retention is not None: + self._retention.forget_pending(released) + for key in keys: + self._retry.pop(key, None) + self._failures.pop(key, None) + self._completed_destinations.pop(key, None) + self._destination_attempts.pop(key, None) + + def _schedule_retry(self, keys: list[ObjectKey]) -> None: + owned = self._listener.ownership() + for key in keys: + if key not in owned: + continue + attempt = self._failures.get(key, 0) + 1 + self._failures[key] = attempt + self._retry[key] = ( + time.monotonic() + min(5.0, 0.1 * 2 ** min(attempt, 6)), + attempt, + ) + + def _retry_ready(self) -> None: + owned = self._listener.ownership() + for key in list(self._retry): + if key not in owned: + self._discard_owned([key]) + cancelled = [ + key + for key in owned + if is_recurrent_checkpoint_key(key) + and self._retention is not None + and not self._retention.is_retained(key) + and not self._key_in_flight(key) + ] + self._discard_owned(cancelled) + now = time.monotonic() + ready = [ + key + for key, (when, _) in self._retry.items() + if when <= now and not self._key_in_flight(key) + ] + for key in ready: + self._retry.pop(key, None) + for group in _group_keys_by_shape(ready).values(): + reuse = [key for key in group if owned.get(key, (0, False))[1]] + reuse_keys = set(reuse) + new = [key for key in group if key not in reuse_keys] + plans = [ + self._policy.select_reuse_targets(reuse, self._routing_descriptors()), + self._policy.select_store_targets(new, self._routing_descriptors()), + ] + submitted = set() + for plan in plans: + submitted.update(self._submit_plan(plan)) + self._discard_owned( + [ + key + for key in group + if key not in submitted and key not in self._retry + ] + ) diff --git a/lmcache/v1/distributed/storage_manager.py b/lmcache/v1/distributed/storage_manager.py index 3d6350d8fae..071d14db8f6 100644 --- a/lmcache/v1/distributed/storage_manager.py +++ b/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, ) @@ -118,6 +121,7 @@ def __init__(self, config: StorageManagerConfig): self._checkpoint_shutdown_flush_seconds = ( config.checkpoint_shutdown_flush_seconds ) + self._checkpoint_shutdown_flushed = False self._checkpoint_retention = CheckpointRetention( write_on_evict=store_policy.writes_checkpoints_on_evict(), persist_timeout=config.checkpoint_write_timeout_seconds, @@ -129,6 +133,9 @@ def __init__(self, config: StorageManagerConfig): eviction_config=config.eviction_config, ) self._eviction_controller.set_checkpoint_retention(self._checkpoint_retention) + self._eviction_controller.set_pressure_grace( + self._store_admission_timeout_seconds / 4 + ) self._eviction_controller.start() # L2 adapters and store controller. When an adapter config carries @@ -200,10 +207,11 @@ def __init__(self, config: StorageManagerConfig): l2_adapters=list(self._l2_adapters.values()), adapter_descriptors=list(self._adapter_descriptors.values()), policy=store_policy, + retention=self._checkpoint_retention, ) self._store_controller.start() self._checkpoint_retention.set_persist( - self._store_controller.submit_reused_keys + self._store_controller.submit_checkpoint_keys ) # Prefetch controller @@ -213,6 +221,7 @@ def __init__(self, config: StorageManagerConfig): adapter_descriptors=list(self._adapter_descriptors.values()), policy=create_prefetch_policy(config.prefetch_policy), max_in_flight=config.prefetch_max_in_flight, + replica_may_match=self._checkpoint_retention.replica_may_match, ) self._prefetch_controller.start() @@ -274,6 +283,8 @@ def reserve_write_detailed( keys: list[ObjectKey], layout_desc: MemoryLayoutDesc, mode: Literal["new", "update", "all"], + *, + internal: bool = False, ) -> dict[ObjectKey, tuple[L1Error, MemoryObj | None]]: """Reserve objects while preserving each L1 failure reason. @@ -294,6 +305,7 @@ def reserve_write_detailed( is_temporary=[False] * len(keys), layout_desc=layout_desc, mode=mode, + internal=internal, ) result = {k: m for k, (e, m) in reserve_result.items() if m is not None} @@ -337,13 +349,77 @@ def reserve_write( if memory_obj is not None } - def get_readable_keys(self, keys: list[ObjectKey]) -> list[ObjectKey]: - """Return complete readable keys without retaining their read locks.""" - results = self._l1_manager.reserve_read(keys) - readable = [key for key, (_error, obj) in results.items() if obj is not None] - if readable: - self._l1_manager.finish_read(readable) - return readable + def get_readable_keys( + self, keys: list[ObjectKey], *, retained_only: bool = False + ) -> list[ObjectKey]: + """Snapshot readable keys; optionally exclude automatically freed pages.""" + return self._l1_manager.readable_keys(keys, retained_only=retained_only) + + def pin_readable_keys( + self, keys: list[ObjectKey], *, retained_only: bool = False + ) -> list[ObjectKey]: + """Acquire ownership of existing pages; return the keys actually pinned. + + Release these ownership pins with release_read_pins, not the restore + completion API: publication ownership is not consumer reuse. + """ + return [ + key + for key, (_, obj) in self._l1_manager.reserve_read( + keys, internal=True, retained_only=retained_only + ).items() + if obj is not None + ] + + def release_read_pins(self, keys: list[ObjectKey]) -> None: + """Release publication ownership without triggering reuse writes.""" + if keys: + self._l1_manager.release_internal_reads(keys) + + def finish_write_pinned(self, keys: list[ObjectKey]) -> list[ObjectKey]: + """Commit new payloads and atomically keep them pinned for publication. + + Return successfully committed keys. The caller owns one read pin per + returned key, including on partial success, and must release it. + """ + result = self._l1_manager.finish_write_and_reserve_read( + keys, notify_store=True, internal=True + ) + succeeded = [ + key for key, (error, _) in result.items() if error == L1Error.SUCCESS + ] + self._event_bus.publish( + Event( + event_type=EventType.SM_WRITE_FINISHED, + metadata={ + "succeeded_keys": succeeded, + "failed_keys": [key for key in keys if key not in succeeded], + }, + ) + ) + return succeeded + + def checkpoint_publication_budget(self) -> int: + """Bound aggregate pending publication bytes below the L1 watermark.""" + _, total = self._l1_manager.get_memory_usage() + return int(total * self._eviction_controller.trigger_watermark) + + def checkpoint_allocation_bytes(self, page_bytes: int) -> int: + """Return the aligned RAM charge for one checkpoint page.""" + alignment = self._l1_memory_desc.align_bytes + return ((page_bytes + alignment - 1) // alignment) * alignment + + def is_l1_evictable(self, key: ObjectKey) -> bool: + """Check eligibility; actual deletion rechecks ownership atomically.""" + return self._l1_manager.is_key_evictable(key) + + def can_cancel_persistence(self, key: ObjectKey) -> bool: + """Whether only queued backend work can be cancelled for this page.""" + return self._store_controller.owns_queued(key) + + def cancel_queued_persistence(self, keys: list[ObjectKey]) -> None: + """Release unsubmitted persistence for coherently retired pages.""" + self._store_controller.cancel_queued(keys) @property def store_admission_timeout_seconds(self) -> float: @@ -611,23 +687,36 @@ def _track_checkpoint_residency( retention.record_l2_present( adapter_id, list(inventory), list(inventory.values()) ) + if adapter.has_complete_inventory(): + retention.set_inventory_complete(adapter_id) + + def flush_checkpoints(self) -> None: + """Perform the configured bounded flush while the directory is open.""" + if self._checkpoint_shutdown_flushed: + return + try: + self._flush_current_checkpoints() + finally: + # No payload may be retired after its directory has closed. + self._eviction_controller.stop() + self._l2_eviction_controller.stop() + self._checkpoint_shutdown_flushed = True def _flush_current_checkpoints(self) -> None: """Write current checkpoint pages still only in L1 to L2 on shutdown.""" retention = self._checkpoint_retention budget = self._checkpoint_shutdown_flush_seconds - if not retention.write_on_evict or budget <= 0 or not self._has_l2_adapters(): + if budget <= 0 or not self._has_l2_adapters(): return - count = self._l1_manager.num_objects() - keys, _ = self._l1_manager.get_evictable_keys(limit=count, scan_limit=count) + keys = self._l1_manager.readable_key_sizes() pending = [ key for key in keys if is_recurrent_checkpoint_key(key) - and not retention.is_superseded(key) and not retention.is_l2_resident(key) + and retention.is_retained(key) ] - if not pending: + if not pending and not self._store_controller.has_pending_work(): return logger.info( "Writing %d current checkpoint pages to L2 before shutdown (budget %.0f s)", @@ -635,10 +724,13 @@ def _flush_current_checkpoints(self) -> None: budget, ) start = time.monotonic() - retention.request_persist(pending) while time.monotonic() - start < budget: pending = [key for key in pending if not retention.is_l2_resident(key)] - if not pending: + if retention.write_on_evict: + retention.request_persist(pending) + if not self._store_controller.has_pending_work() and ( + not pending or not retention.write_on_evict + ): break time.sleep(0.2) if pending: @@ -953,13 +1045,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 original-position hits and actual load-reservation evidence. + + Legacy bitmap queries consume the same result. Backend misses alone do + not prove missing payload; no post-hoc free-memory inference is made. + """ 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 @@ -983,7 +1089,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]): """ @@ -1024,9 +1130,15 @@ def delete_l1_keys( and the number refused because they were locked (non-force only). Missing keys are a no-op, so the operation is idempotent. """ - results = self._l1_manager.delete(keys, force=force) + results = self._checkpoint_retention.evict( + keys, + "l1-force" if force else "l1", + lambda selected: self._l1_manager.delete(selected, force=force), + ) deleted = sum(1 for err in results.values() if err == L1Error.SUCCESS) - skipped = sum(1 for err in results.values() if err == L1Error.KEY_IS_LOCKED) + skipped = sum( + 1 for err in results.values() if err == L1Error.KEY_IS_LOCKED + ) + len(set(keys) - set(results)) return deleted, skipped def unsafe_read( @@ -1379,14 +1491,23 @@ def clear(self, force: bool = False): If False (default), only clear unlocked objects, keeping write-locked and read-locked objects intact. """ - self._l1_manager.clear(force=force) + if force: + # Forced administrative deletion retains its documented unsafe-copy + # semantics; retire directory entries before discarding buffers. + self._checkpoint_retention.evict( + self._l1_manager.keys(), + "l1-force", + lambda keys: self._l1_manager.delete(keys, force=True), + ) + else: + self.delete_l1_keys(self._l1_manager.keys()) def close(self): """ Close the storage manager and release all resources. """ try: - self._flush_current_checkpoints() + self.flush_checkpoints() except Exception: logger.exception("Shutdown checkpoint flush failed") self._l1_manager.begin_shutdown() @@ -1400,6 +1521,9 @@ def close(self): for adapter in self._l2_adapters.values(): adapter.close() + self._store_controller.release_stopped_ownership() + self._prefetch_controller.release_stopped_ownership() + self._l1_manager.close() def report_status(self) -> dict: diff --git a/lmcache/v1/multiprocess/cache_control/object_service.py b/lmcache/v1/multiprocess/cache_control/object_service.py index dc836577cf4..ae9522a8487 100644 --- a/lmcache/v1/multiprocess/cache_control/object_service.py +++ b/lmcache/v1/multiprocess/cache_control/object_service.py @@ -13,7 +13,7 @@ import asyncio # First Party -from lmcache.v1.distributed.api import EncodedObjectKey, Tier +from lmcache.v1.distributed.api import EncodedObjectKey, ObjectKey, Tier from lmcache.v1.multiprocess.cache_control.errors import ( InvalidRequest, NotFound, @@ -150,12 +150,14 @@ async def delete_objects( L2 deletion is idempotent at the adapter level (absent / locked keys are skipped). The L2 adapter is only resolved when the tier includes L2, so a pure ``l1`` delete works on an L1-only server. + Duplicate keys are processed and counted once per tier. Returns: ``{"deleted", "skipped", "ok"[, "error"]}``: ``deleted`` is the total - keys removed across the requested tiers (L1 removals plus the L2 batch - size); ``skipped`` is the L1 keys refused because they were locked - (non-force only); ``ok`` is ``False`` with ``error`` set when the L2 + L1 keys removed plus L2 keys submitted after checkpoint protection; + adapters may additionally skip busy or absent keys. ``skipped`` counts + L1 lock refusals and L2 checkpoint protection refusals. + ``ok`` is ``False`` with ``error`` set when the L2 adapter raised (a structured failure, not a crash). Raises: @@ -175,6 +177,7 @@ async def delete_objects( parsed.append(cache_key.to_object_key()) except ValueError as exc: raise InvalidRequest(f"keys[{i}]: {exc}") from None + parsed = list(dict.fromkeys(parsed)) deleted = 0 skipped = 0 @@ -189,10 +192,21 @@ async def delete_objects( skipped += l1_skipped if tier in (Tier.L2, Tier.ALL): - _, adapter = self._resolve_adapter(adapter_selector) + descriptor, adapter = self._resolve_adapter(adapter_selector) + + def delete_selected(selected: list[ObjectKey]) -> int: + adapter.delete(selected) + return len(selected) + try: - await asyncio.to_thread(adapter.delete, parsed) - deleted += len(parsed) + selected_count = await asyncio.to_thread( + self._engine.storage_manager.checkpoint_retention.evict, + parsed, + descriptor.index, + delete_selected, + ) + deleted += selected_count + skipped += len(parsed) - selected_count except Exception as exc: # noqa: BLE001 - structured result, not a crash ok = False error = str(exc) diff --git a/lmcache/v1/multiprocess/checkpoint_index.py b/lmcache/v1/multiprocess/checkpoint_index.py index 0f3d09d54dc..96ecf06d924 100644 --- a/lmcache/v1/multiprocess/checkpoint_index.py +++ b/lmcache/v1/multiprocess/checkpoint_index.py @@ -8,6 +8,7 @@ """ # Standard +from collections.abc import Callable from dataclasses import dataclass, field from pathlib import Path import json @@ -172,7 +173,11 @@ def __init__( self._max_entries = max_entries self._max_pending = max_pending self._pending: dict[str, _PendingManifest] = {} - self._lock = threading.Lock() + self._lock = threading.RLock() + self._listeners: list[Callable[[str, CheckpointManifest | str], None]] = [] + self._close_callbacks: list[Callable[[], None]] = [] + self._closed = False + self._closing = False self._db = sqlite3.connect( str(path) if path is not None else ":memory:", check_same_thread=False ) @@ -234,6 +239,42 @@ def __init__( "SELECT COALESCE(MAX(access_order), 0) FROM checkpoints" ).fetchone()[0] + @property + def lifecycle_lock(self) -> threading.RLock: + """Serialize directory mutations with payload ownership transitions. + + Callers must not wait for capacity or backend I/O while holding this + reentrant lock. L1 callbacks must never acquire it. + """ + return self._lock + + def add_lifecycle_listener( + self, + listener: Callable[[str, CheckpointManifest | str], None], + before_close: Callable[[], None], + ) -> None: + """Subscribe to committed changes under the lifecycle lock. + + Events are publish/touch (manifest) and remove/abort (generation). + The close callback runs outside the lock while SQLite is still open. + Listeners must be nonblocking and must not raise. + """ + with self._lock: + self._listeners.append(listener) + self._close_callbacks.append(before_close) + + def manifests(self) -> list[CheckpointManifest]: + """Return published manifests ordered from least to most recently used.""" + with self._lock: + generations = self._db.execute( + "SELECT generation FROM checkpoints ORDER BY access_order" + ).fetchall() + return [ + manifest + for (generation,) in generations + if (manifest := self.get(generation)) is not None + ] + def begin(self, manifest: CheckpointManifest) -> bool: """Stage a complete manifest before rank acknowledgements arrive. @@ -252,6 +293,8 @@ def begin(self, manifest: CheckpointManifest) -> bool: """ with self._lock: pending = self._pending.get(manifest.generation) + if self._closing: + return False if pending is not None: if pending.manifest != manifest: raise ValueError("checkpoint generation changed during store") @@ -309,6 +352,19 @@ def acknowledge(self, generation: str, rank: int) -> bool: self._access_order += 1 tail = json.dumps(prefix.tail_tokens).encode() with self._db: + removed = [ + row[0] + for row in self._db.execute( + "SELECT generation FROM checkpoints WHERE namespace=? " + "AND start_tokens=? AND prefix_hash=? AND tail=?", + ( + prefix.namespace, + prefix.start_tokens, + prefix.prefix_hash, + tail, + ), + ) + ] # The generation replaced for this exact prefix, if any. self._db.execute( "DELETE FROM checkpoint_tails WHERE generation IN (" @@ -340,6 +396,14 @@ def acknowledge(self, generation: str, rank: int) -> bool: _prefix_hashes(prefix.tail_tokens)[-1], ), ) + removed.extend( + row[0] + for row in self._db.execute( + "SELECT generation FROM checkpoints ORDER BY access_order DESC " + "LIMIT -1 OFFSET ?", + (self._max_entries,), + ) + ) for table in ("checkpoint_tails", "checkpoints"): self._db.execute( f"DELETE FROM {table} WHERE generation IN (" @@ -348,6 +412,9 @@ def acknowledge(self, generation: str, rank: int) -> bool: (self._max_entries,), ) del self._pending[generation] + for retired in removed: + self._notify("remove", retired) + self._notify("publish", manifest) return True def is_pending(self, manifest: CheckpointManifest) -> bool: @@ -415,6 +482,7 @@ def find(self, prefixes: tuple[CheckpointPrefix, ...]) -> CheckpointManifest | N "UPDATE checkpoints SET access_order=? WHERE generation=?", (self._access_order, best.generation), ) + self._notify("touch", best) return best def get(self, generation: str) -> CheckpointManifest | None: @@ -493,6 +561,7 @@ def abort_stale(self, older_than_seconds: float) -> list[str]: ] for generation in stale: del self._pending[generation] + self._notify("abort", generation) return stale def abort(self, generation: str) -> None: @@ -503,6 +572,7 @@ def abort(self, generation: str) -> None: """ with self._lock: self._pending.pop(generation, None) + self._notify("abort", generation) def invalidate(self, generation: str) -> None: """Remove only the generation whose payload retrieval failed. @@ -511,11 +581,13 @@ def invalidate(self, generation: str) -> None: generation: Failed generation, not merely its shared token prefix. Invalidating a replaced generation cannot remove its replacement. """ - with self._lock, self._db: - for table in ("checkpoint_tails", "checkpoints"): - self._db.execute( - f"DELETE FROM {table} WHERE generation=?", (generation,) - ) + with self._lock: + with self._db: + for table in ("checkpoint_tails", "checkpoints"): + self._db.execute( + f"DELETE FROM {table} WHERE generation=?", (generation,) + ) + self._notify("remove", generation) def report_status(self) -> dict[str, int]: """Return directory counts without asserting that payload pages are resident. @@ -535,5 +607,18 @@ def report_status(self) -> dict[str, int]: def close(self) -> None: """Discard pending generations and close the database connection.""" with self._lock: + if self._closed: + return + self._closing = True + for callback in self._close_callbacks: + callback() + with self._lock: + for generation in list(self._pending): + self._notify("abort", generation) self._pending.clear() self._db.close() + self._closed = True + + def _notify(self, event: str, value: CheckpointManifest | str) -> None: + for listener in self._listeners: + listener(event, value) diff --git a/lmcache/v1/multiprocess/checkpoint_storage.py b/lmcache/v1/multiprocess/checkpoint_storage.py index 9918aeb8562..876ae21a7df 100644 --- a/lmcache/v1/multiprocess/checkpoint_storage.py +++ b/lmcache/v1/multiprocess/checkpoint_storage.py @@ -7,11 +7,12 @@ """ # Standard +from collections import OrderedDict +from collections.abc import Callable from dataclasses import dataclass, field -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, TypeVar import hashlib import json -import threading import time import uuid @@ -43,6 +44,7 @@ from lmcache.v1.memory_management import MemoryObj logger = init_logger(__name__) +_DeleteResult = TypeVar("_DeleteResult") # A retrieve waiting for RAM asks the L1 eviction loop again after this many # seconds, as store admission does: a pass frees nothing while its victims are @@ -214,6 +216,7 @@ class _StoreLease: manifest: CheckpointManifest rank: int keys: list[ObjectKey] + reused: list[ObjectKey] = field(default_factory=list) started: float = field(default_factory=time.monotonic) @@ -243,13 +246,10 @@ class CheckpointPayloadStore: index: Directory used for all-rank publication and stale invalidation. max_leases: Shared bound for pending stores and retrieves. Exhaustion rejects admission without recycling a worker's live SHM buffers. - abandoned_after_seconds: Age after which a lease or pending generation - is treated as abandoned by a dead worker and released; 0 disables - reclaiming. It must be far longer than any live transfer. A worker - whose copy could not drain (UnsafeCheckpointCopyError) keeps its - lease on purpose because its DMA may still touch those pages; - reclaiming releases that lease too, which is safe only because a - copy still running after this long means a hung CUDA context. + abandoned_after_seconds: Age after which unexposed lookups and pending + generations are cancelled; 0 disables automatic cancellation. + Exposed copy buffers stay owned until the worker reports completion. + Elapsed time cannot prove that a CUDA transfer has stopped. The caller stages a manifest with ``index.begin`` before rank stores. Each successful store acknowledgement follows a drained worker D2H transfer. @@ -277,7 +277,38 @@ def __init__( self._store_ranks: set[tuple[str, int]] = set() self._retrieves: dict[str, _RetrieveLease] = {} self._retrieve_admissions: set[str] = set() - self._lock = threading.Lock() + self._lock = index.lifecycle_lock + self._completed: dict[str, list[ObjectKey]] = {} + self._completed_ranks: dict[str, set[int]] = {} + self._publication_bytes: dict[str, int] = {} + self._generations: OrderedDict[str, dict[ObjectKey, int]] = OrderedDict() + self._owners: dict[ObjectKey, set[str]] = {} + self._published: set[str] = set() + self._deleting: dict[ObjectKey, set[str | int]] = {} + self._closed = False + self._accepting = True + index.add_lifecycle_listener(self._index_changed, self._before_index_close) + storage.checkpoint_retention.set_lifecycle( + evict=self._evict_pages, + pressure=self._retire_under_pressure, + retained=self._is_retained, + order=self._persistence_order, + reconcile=self._reconcile_replicas, + ) + with self._lock: + for manifest in index.manifests(): + self._register_manifest(manifest) + self._published.add(manifest.generation) + keys = list(self._generations[manifest.generation]) + readable = set(storage.get_readable_keys(keys, retained_only=True)) + if storage.checkpoint_retention.inventories_complete() and any( + key not in readable + and not storage.checkpoint_retention.is_l2_resident( + key, self._generations[manifest.generation][key] + ) + for key in keys + ): + index.invalidate(manifest.generation) def prepare_store( self, manifest: CheckpointManifest, rank: int @@ -304,8 +335,28 @@ def prepare_store( return None identity = (manifest.generation, rank) with self._lock: - if identity in self._store_ranks: + if not self._accepting: + return None + if identity in self._store_ranks or rank in self._completed_ranks.get( + manifest.generation, set() + ): raise ValueError("checkpoint rank already has a store lease") + if any(key in self._deleting for group in key_groups for key in group): + return AdmissionFailure.BUSY + self._register_manifest(manifest) + if manifest.generation not in self._publication_bytes: + footprint = manifest.world_size * sum( + self._storage.checkpoint_allocation_bytes(group.page_bytes) + * len(group.positions) + for group in groups + ) + if ( + footprint + sum(self._publication_bytes.values()) + > self._storage.checkpoint_publication_budget() + ): + self._index.abort(manifest.generation) + return None + self._publication_bytes[manifest.generation] = footprint if ( len(self._store_ranks) + len(self._retrieves) @@ -316,18 +367,23 @@ def prepare_store( return None self._store_ranks.add(identity) reserved: list[ObjectKey] = [] + reused: list[ObjectKey] = [] def attempt() -> AdmissionAttempt[ tuple[tuple[ShmSlotDescriptor | None, ...], ...] ]: - if not self._index.is_pending(manifest): - return AdmissionAttempt.failure(AdmissionFailure.CONFLICT) + with self._lock: + if not self._index.is_pending(manifest): + return AdmissionAttempt.failure(AdmissionFailure.CONFLICT) + if any(key in self._deleting for group in key_groups for key in group): + return AdmissionAttempt.failure(AdmissionFailure.BUSY) slots: list[tuple[ShmSlotDescriptor | None, ...]] = [] for group, keys in zip(groups, key_groups, strict=True): detailed = self._storage.reserve_write_detailed( list(keys), MemoryLayoutDesc([torch.Size([group.page_bytes])], [torch.uint8]), "new", + internal=True, ) objects = { key: obj @@ -341,10 +397,17 @@ def attempt() -> AdmissionAttempt[ if detailed.get(key, (L1Error.KEY_NOT_EXIST, None))[0] is L1Error.KEY_NOT_WRITABLE ] - readable = set(self._storage.get_readable_keys(existing_candidates)) + readable = set( + self._storage.pin_readable_keys( + existing_candidates, retained_only=True + ) + ) + reused.extend(readable) if len(objects) + len(readable) != len(keys): self._storage.abort_write(reserved) reserved.clear() + self._storage.release_read_pins(reused) + reused.clear() missing = [ key for key in keys @@ -400,13 +463,14 @@ def attempt() -> AdmissionAttempt[ return AdmissionFailure.BUSY if retryable else None lease_id = uuid.uuid4().hex with self._lock: - self._stores[lease_id] = _StoreLease(manifest, rank, reserved) + self._stores[lease_id] = _StoreLease(manifest, rank, reserved, reused) admitted = True return CheckpointSlots(lease_id, outcome.value) finally: if not admitted: try: self._storage.abort_write(reserved) + self._storage.release_read_pins(reused) finally: try: if not retryable: @@ -430,28 +494,32 @@ def finish_store(self, lease_id: str, success: bool) -> bool: lease = self._stores.pop(lease_id, None) if lease is None: return False - try: - if not success: - self._storage.abort_write(lease.keys) - self._index.abort(lease.manifest.generation) - return False - self._storage.finish_write(lease.keys) - if len(self._storage.get_readable_keys(lease.keys)) != len(lease.keys): - self._index.abort(lease.manifest.generation) - return False - return self._index.acknowledge(lease.manifest.generation, lease.rank) - except Exception: - # Commit can fail before releasing every exclusive write lock. - # abort_write leaves already-readable objects intact; none become - # discoverable through this generation after publication is aborted. + with self._lock: + pinned = list(lease.reused) try: + if not success or not self._index.is_pending(lease.manifest): + self._storage.abort_write(lease.keys) + self._index.abort(lease.manifest.generation) + return False + committed = self._storage.finish_write_pinned(lease.keys) + pinned.extend(committed) + if len(committed) != len(lease.keys): + self._storage.abort_write(lease.keys) + self._index.abort(lease.manifest.generation) + return False + self._completed.setdefault(lease.manifest.generation, []).extend(pinned) + self._completed_ranks.setdefault(lease.manifest.generation, set()).add( + lease.rank + ) + pinned = [] + return self._index.acknowledge(lease.manifest.generation, lease.rank) + except Exception: self._storage.abort_write(lease.keys) - finally: self._index.abort(lease.manifest.generation) - raise - finally: - with self._lock: - self._store_ranks.remove((lease.manifest.generation, lease.rank)) + raise + finally: + self._storage.release_read_pins(pinned) + self._store_ranks.discard((lease.manifest.generation, lease.rank)) def begin_retrieve(self, manifest: CheckpointManifest, rank: int) -> str | None: """Start asynchronous RAM/filesystem lookup for every page of one rank. @@ -473,6 +541,8 @@ def begin_retrieve(self, manifest: CheckpointManifest, rank: int) -> str | None: ] lease_id = uuid.uuid4().hex with self._lock: + if not self._accepting: + return None if ( len(self._store_ranks) + len(self._retrieves) @@ -515,36 +585,17 @@ def _maybe_reclaim(self) -> None: self.reclaim_abandoned() def reclaim_abandoned(self) -> int: - """Release leases and pending generations that a dead worker left behind. + """Cancel stale publication and unexposed lookups without recycling DMA. - A worker that dies between PREPARE_STORE and FINISH_STORE, or between - a ready retrieve and its release, never returns its lease. The lease - keeps its L1 locks and counts against ``max_leases``, and its pending - generation can never be staged again. Once older than - ``abandoned_after_seconds``, store leases are aborted, ready retrieve - leases release their read locks, and pending lookups are cancelled and - drained as their prefetch completes. + Completed ranks release their publication pins when their generation + aborts. Live writers and exposed retrieves require explicit completion; + their buffers remain bounded by the pool and lease admission limits. Returns: - Number of leases and pending generations released. + Number of cancelled lookups and aborted pending generations. """ cutoff = time.monotonic() - self._abandoned_after with self._lock: - stores = [ - (lease_id, lease) - for lease_id, lease in self._stores.items() - if lease.started < cutoff - ] - for lease_id, lease in stores: - del self._stores[lease_id] - self._store_ranks.discard((lease.manifest.generation, lease.rank)) - ready = [ - (lease_id, lease) - for lease_id, lease in self._retrieves.items() - if lease.slots is not None and lease.started < cutoff - ] - for lease_id, _lease in ready: - del self._retrieves[lease_id] lookups = [ lease_id for lease_id, lease in self._retrieves.items() @@ -552,11 +603,6 @@ def reclaim_abandoned(self) -> int: ] for lease_id in lookups: self._retrieves[lease_id].cancelled = True - for _lease_id, store_lease in stores: - self._storage.abort_write(store_lease.keys) - self._index.abort(store_lease.manifest.generation) - for _lease_id, read_lease in ready: - self._storage.finish_read_prefetched(read_lease.keys) for lease_id in lookups: try: # A cancelled lookup releases its locks once its prefetch ends. @@ -564,13 +610,11 @@ def reclaim_abandoned(self) -> int: except KeyError: pass stale = self._index.abort_stale(self._abandoned_after) - reclaimed = len(stores) + len(ready) + len(lookups) + len(stale) + reclaimed = len(lookups) + len(stale) if reclaimed: logger.warning( - "Released %d store leases, %d retrieve leases, %d lookups and %d " + "Cancelled %d lookups and %d " "pending checkpoint generations abandoned for over %.0f s", - len(stores), - len(ready), len(lookups), len(stale), self._abandoned_after, @@ -592,6 +636,31 @@ def report_status(self) -> dict[str, int]: "max_leases": self._max_leases, } + def quiesce(self) -> None: + """Close admission atomically with the idle check. + + Raises RuntimeError if a worker still owns copy buffers or an admission + is in progress. A failed attempt leaves the service usable so workers + can finish and the caller can retry closure. + """ + with self._lock: + if self._store_ranks or self._retrieves or self._retrieve_admissions: + raise RuntimeError( + "Checkpoint worker copy leases must drain before close" + ) + self._accepting = False + + def prepare_terminal_shutdown(self) -> None: + """Stop admission and flush the directory without releasing live copies. + + Terminal process shutdown may leave leases from workers already killed + by the supervisor. Those buffers remain owned until process termination; + closing their metadata never authorizes recycling their memory. + """ + with self._lock: + self._accepting = False + self._index.close() + def poll_retrieve(self, lease_id: str) -> CheckpointSlots | bool | None: """Return slots only when every required page is readable and pinned. @@ -600,16 +669,12 @@ def poll_retrieve(self, lease_id: str) -> CheckpointSlots | bool | None: Returns: None while prefetch is pending, False on a miss/cancellation, or - the pinned slots. A miss releases all acquired locks and invalidates - only the failed generation. It must not advance computed tokens. - - 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 - at most the storage admission timeout; a lease that never gets - room misses without invalidating its generation. + the pinned slots. A miss releases all acquired locks and must not + advance computed tokens. Reservation failures retry within the + original admission deadline, using aligned allocation sizes. + Unknown backend misses preserve the generation: failed I/O is not + proof of missing stored pages. Coordinated retirement and recovery + remove manifests whose payloads are known to be lost. Raises: KeyError: If the lease is unknown or already finished. @@ -621,13 +686,21 @@ def poll_retrieve(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 + ) + pinned = self._storage.pin_readable_keys(lease.keys) + if pinned != lease.keys: + self._storage.release_read_pins(pinned) + self._storage.finish_read_prefetched(lease.keys) + return self._repeat_or_miss(lease_id, lease, [], True) try: keys, objects = self._storage.unsafe_read(lease.keys) if keys != lease.keys or len(objects) != len(keys): @@ -646,41 +719,54 @@ def poll_retrieve(self, lease_id: str) -> CheckpointSlots | bool | None: lease.slots = CheckpointSlots(lease_id, tuple(slots)) return lease.slots except Exception: + self._storage.release_read_pins(pinned) self._storage.finish_read_prefetched(lease.keys) self._index.invalidate(lease.manifest.generation) del self._retrieves[lease_id] 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 = False, ) -> 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. An unavailable page does not prove that its stored copy is lost. """ lease.readable = len(readable_keys) if lease.cancelled: return self._miss(lease_id, lease, "was cancelled by the engine") + if self._index.get(lease.manifest.generation) is None: + return self._miss( + lease_id, lease, "missed; its checkpoint is no longer listed" + ) groups = checkpoint_page_groups(lease.manifest) readable = set(readable_keys) lease.unread_bytes = sum( - groups[key.object_group_id].page_bytes + self._storage.checkpoint_allocation_bytes( + groups[key.object_group_id].page_bytes + ) for key in lease.keys if key not in readable ) - used, total = self._storage.get_l1_usage() - # Without L2, or with more pages than RAM can hold, no repeat loads them. + _, total = self._storage.get_l1_usage() 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) ): lease.handle = None return self._repeat_with_room(lease_id, lease) - self._index.invalidate(lease.manifest.generation) - return self._miss(lease_id, lease, "missed; its checkpoint is no longer listed") + # Adapters can report a miss for failed I/O as well as absence. Only + # coordinated retirement/recovery provides authoritative loss evidence. + return self._miss( + lease_id, lease, "was unavailable; its checkpoint stays listed" + ) def _repeat_with_room(self, lease_id: str, lease: _RetrieveLease) -> bool | None: """Repeat a waiting lookup once RAM can hold its unread pages. @@ -692,16 +778,16 @@ 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") - used, total = self._storage.get_l1_usage() - if total - used >= lease.unread_bytes: - 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" ) + used, total = self._storage.get_l1_usage() + if total - used >= lease.unread_bytes: + lease.handle = self._lookup(lease_id, lease.manifest, lease.keys) + lease.lookups += 1 + return None if now - lease.eviction_requested >= _ROOM_REQUEST_INTERVAL_SECONDS: lease.eviction_requested = now self._storage.request_immediate_eviction() @@ -737,11 +823,13 @@ def finish_retrieve(self, lease_id: str) -> None: return if lease.slots is None: raise ValueError("checkpoint retrieval has not produced a read lease") - del self._retrieves[lease_id] - # A completed restore proves reuse; reuse-admission store policies - # store checkpoint pages to L2 only after such a report. - self._storage.notify_keys_reused(lease.keys) - self._storage.finish_read_prefetched(lease.keys) + # Keep admission/shutdown accounting until ownership cleanup ends. + try: + self._storage.notify_keys_reused(lease.keys) + self._storage.finish_read_prefetched(lease.keys) + finally: + self._storage.release_read_pins(lease.keys) + del self._retrieves[lease_id] def cancel_retrieve(self, lease_id: str) -> None: """Mark a pending lookup for draining without exposing its SHM slots. @@ -761,3 +849,233 @@ def cancel_retrieve(self, lease_id: str) -> None: "prepared checkpoint retrieval requires copy completion" ) lease.cancelled = True + + def _index_changed(self, event: str, value: CheckpointManifest | str) -> None: + if event in {"publish", "touch"}: + assert isinstance(value, CheckpointManifest) + self._register_manifest(value) + self._generations.move_to_end(value.generation) + if event == "publish": + self._published.add(value.generation) + if event not in {"publish", "abort", "remove"}: + return + generation = ( + value.generation if isinstance(value, CheckpointManifest) else value + ) + if event in {"abort", "remove"}: + # Aborting an already published generation is a no-op. + if event != "abort" or generation not in self._published: + self._published.discard(generation) + for key in self._generations.pop(generation, {}): + owners = self._owners[key] + owners.discard(generation) + if not owners: + del self._owners[key] + self._storage.checkpoint_retention.forget_pending([key]) + self._storage.release_read_pins(self._completed.pop(generation, [])) + self._completed_ranks.pop(generation, None) + self._publication_bytes.pop(generation, None) + + def _before_index_close(self) -> None: + if self._closed: + return + with self._lock: + self._accepting = False + for generation in list(self._generations): + if generation in self._published: + continue + self._index.abort(generation) + try: + self._storage.flush_checkpoints() + except Exception: + logger.exception( + "Checkpoint shutdown flush failed; retiring RAM-only entries" + ) + with self._lock: + retention = self._storage.checkpoint_retention + lost = [ + generation + for generation in self._published + if any( + not retention.is_l2_resident(key, size) + for key, size in self._generations[generation].items() + ) + ] + for generation in lost: + self._index.invalidate(generation) + if lost: + logger.warning("Shutdown retired %d RAM-only checkpoints", len(lost)) + retention.record_retirement(len(lost)) + self._closed = True + retention.set_lifecycle(evict=None, pressure=None, retained=None) + + def _register_manifest(self, manifest: CheckpointManifest) -> None: + if manifest.generation in self._generations: + return + sizes: dict[ObjectKey, int] = {} + groups = checkpoint_page_groups(manifest) + for rank in range(manifest.world_size): + for group, keys in zip( + groups, checkpoint_object_keys(manifest, rank), strict=True + ): + sizes.update((key, group.page_bytes) for key in keys) + self._generations[manifest.generation] = sizes + for key in sizes: + self._owners.setdefault(key, set()).add(manifest.generation) + + def _is_retained(self, key: ObjectKey) -> bool: + with self._lock: + return key in self._owners + + def _reconcile_replicas(self) -> None: + """Adapter removal can invalidate unenumerated recovered dependencies.""" + with self._lock: + retention = self._storage.checkpoint_retention + for generation in list(self._published): + keys = list(self._generations[generation]) + readable = set( + self._storage.get_readable_keys(keys, retained_only=True) + ) + if any( + key not in readable + and not retention.is_l2_resident( + key, self._generations[generation][key] + ) + for key in keys + ): + self._index.invalidate(generation) + retention.record_retirement(1) + + def _persistence_order(self, keys: list[ObjectKey]) -> list[ObjectKey]: + """Prefer candidates that complete a checkpoint with fewer missing bytes.""" + with self._lock: + retention = self._storage.checkpoint_retention + candidates = {owner for key in keys for owner in self._owners.get(key, ())} + missing = { + generation: [ + key + for key in self._generations[generation] + if not retention.is_l2_resident( + key, self._generations[generation][key] + ) + ] + for generation in candidates + } + ranked = sorted( + [ + generation + for generation in reversed(self._generations) + if generation in candidates + ], + key=lambda generation: sum( + self._generations[generation][key] for key in missing[generation] + ), + ) + ordered = dict.fromkeys( + key for generation in ranked for key in missing[generation] + ) + return list(ordered) + + def _evict_pages( + self, + keys: list[ObjectKey], + tier: str | int, + delete: Callable[[list[ObjectKey]], _DeleteResult], + ) -> _DeleteResult: + """Fence payload deletion and retire affected last-copy generations.""" + keys = list(dict.fromkeys(keys)) + force = tier == "l1-force" + if force: + tier = "l1" + with self._lock: + if self._closed: + return delete(keys) + retention = self._storage.checkpoint_retention + readable = ( + set(self._storage.get_readable_keys(keys, retained_only=True)) + if tier != "l1" + else set() + ) + selected = [] + retired: set[str] = set() + for key in keys: + if tier in self._deleting.get(key, set()): + continue + if ( + tier == "l1" + and not force + and not self._storage.is_l1_evictable(key) + ): + continue + deleting = self._deleting.get(key, set()) + owners = self._owners.get(key, set()) + retained_l1 = tier != "l1" and key in readable and "l1" not in deleting + at_risk = { + generation + for generation in owners + if not retained_l1 + and not ( + retention.resident_adapters( + key, self._generations[generation][key] + ) + - deleting + - {tier} + ) + } + if at_risk - self._published and not force: + continue + if force: + for generation in list(at_risk - self._published): + self._index.abort(generation) + retired.update(at_risk & self._published) + selected.append(key) + for generation in retired: + self._index.invalidate(generation) + retention.record_retirement(len(retired)) + for key in selected: + self._deleting.setdefault(key, set()).add(tier) + try: + return delete(selected) + finally: + with self._lock: + for key in selected: + self._deleting[key].discard(tier) + if not self._deleting[key]: + del self._deleting[key] + + def _retire_under_pressure(self, target_bytes: int) -> int: + """Free unique pages of cold, inactive generations, preserving sharing.""" + freed = 0 + with self._lock: + candidates = list(self._generations) + for generation in candidates: + if freed >= target_bytes: + break + with self._lock: + if generation not in self._published: + continue + sizes = self._generations[generation] + unique = [ + key + for key in sizes + if self._owners.get(key) == {generation} + and ( + self._storage.is_l1_evictable(key) + or self._storage.can_cancel_persistence(key) + ) + ] + if not unique: + continue + # Shared dependencies survive retirement and need no extra write. + self._index.invalidate(generation) + self._storage.checkpoint_retention.record_retirement(1) + self._storage.cancel_queued_persistence(unique) + self._storage.delete_l1_keys(unique) + # Bound selection by bytes retired, including queued pins that + # the store loop will release. Admission still checks actual RAM. + freed += sum( + self._storage.checkpoint_allocation_bytes(sizes[key]) + for key in unique + ) + logger.info("Retired checkpoint under RAM pressure: %s", generation) + return freed diff --git a/lmcache/v1/multiprocess/modules/checkpoint.py b/lmcache/v1/multiprocess/modules/checkpoint.py index ec3990049d7..4c4cdcb1271 100644 --- a/lmcache/v1/multiprocess/modules/checkpoint.py +++ b/lmcache/v1/multiprocess/modules/checkpoint.py @@ -209,9 +209,9 @@ def supersede( checkpoints, which other conversations share, are kept. Pages the new checkpoint references are current again. - Superseded pages are evicted first and are never written to L2 on - eviction. Their manifests stay listed, so a request that branches - from an older turn can still restore one while its pages last. + Supersession is an eviction preference, not permission to lose a + retained branch. Persist its last copy or retire its manifest before + reclaiming pages that no remaining generation needs. Returns: Number of pages newly marked superseded. @@ -363,7 +363,9 @@ def close(self) -> None: This module never frees buffers solely because a copy took too long. The owning process shutdown must coordinate GPU worker termination. """ - status = self._payloads.report_status() - if status["store_leases"] or status["retrieve_leases"]: - raise RuntimeError("Checkpoint worker copy leases must drain before close") + self._payloads.quiesce() self._index.close() + + def prepare_terminal_shutdown(self) -> None: + """Repair persistent metadata while retaining orphaned copy buffers.""" + self._payloads.prepare_terminal_shutdown() diff --git a/lmcache/v1/multiprocess/server.py b/lmcache/v1/multiprocess/server.py index 15e7f661df8..4b582a94028 100644 --- a/lmcache/v1/multiprocess/server.py +++ b/lmcache/v1/multiprocess/server.py @@ -165,7 +165,10 @@ def report_status(self) -> dict: return status def close(self) -> None: - """Close all modules and release shared resources.""" + """Terminal shutdown: flush metadata before refusing unsafe buffer release.""" + for module in self._modules: + if isinstance(module, CheckpointModule): + module.prepare_terminal_shutdown() for module in self._modules: module.close() self._context.close() diff --git a/tests/v1/distributed/test_checkpoint_retention.py b/tests/v1/distributed/test_checkpoint_retention.py index 8459e902b77..2b9c4351a4a 100644 --- a/tests/v1/distributed/test_checkpoint_retention.py +++ b/tests/v1/distributed/test_checkpoint_retention.py @@ -81,12 +81,12 @@ def test_ordinary_kv_chunks_are_never_held_or_tracked(): assert not retention.is_l2_resident(key) -def test_superseded_page_is_evictable_without_a_write(): +def test_superseded_page_still_needs_persistence(): retention = CheckpointRetention(write_on_evict=True) old, shared = make_key(1), make_key(2) assert retention.mark_superseded("gen-old", [old, shared]) == 2 assert retention.is_superseded_generation("gen-old") - assert not retention.needs_persist_before_evict(old) + assert retention.needs_persist_before_evict(old) # A newer checkpoint references the shared page again. retention.mark_current([shared]) assert not retention.is_superseded(shared) @@ -95,7 +95,7 @@ def test_superseded_page_is_evictable_without_a_write(): assert retention.mark_superseded("gen-old-2", [old]) == 0 -def test_timed_out_write_allows_eviction(): +def test_write_age_alone_does_not_allow_last_copy_eviction(): retention = CheckpointRetention(write_on_evict=True, persist_timeout=5.0) key = make_key(1) with patch("time.monotonic", return_value=100.0): @@ -103,9 +103,9 @@ def test_timed_out_write_allows_eviction(): with patch("time.monotonic", return_value=104.0): assert retention.needs_persist_before_evict(key) with patch("time.monotonic", return_value=106.0): - assert not retention.needs_persist_before_evict(key) - assert retention.report_status()["write_on_evict_timeouts"] == 1 - assert retention.pending_count() == 0 + assert retention.needs_persist_before_evict(key) + assert retention.report_status()["write_on_evict_timeouts"] == 0 + assert retention.pending_count() == 1 def test_superseded_in_adapter_respects_residency_and_budget(): @@ -139,7 +139,8 @@ def test_bounds_forget_oldest_entries(): retention.mark_superseded("gen", keys) assert retention.superseded_keys() == keys[1:] retention.record_l2_present(0, keys, [1, 1, 1]) - assert not retention.is_l2_resident(keys[0]) + # Supersession hints are bounded; replica presence cannot silently expire. + assert retention.is_l2_resident(keys[0]) assert retention.is_l2_resident(keys[2]) diff --git a/tests/v1/distributed/test_fs_l2_adapter_persistence.py b/tests/v1/distributed/test_fs_l2_adapter_persistence.py index 9d2710ef35b..de139464d87 100644 --- a/tests/v1/distributed/test_fs_l2_adapter_persistence.py +++ b/tests/v1/distributed/test_fs_l2_adapter_persistence.py @@ -5,7 +5,9 @@ from pathlib import Path from typing import cast from unittest.mock import patch +import asyncio import os +import threading import time # Third Party @@ -39,6 +41,91 @@ def _memory_obj(data: bytes) -> MemoryObj: return cast(MemoryObj, _Buffer(data)) +def test_close_waits_for_running_executor_io(tmp_path: Path) -> None: + """Coroutine cancellation cannot release memory still read by a syscall.""" + adapter = FSL2Adapter(FSL2AdapterConfig(base_path=str(tmp_path))) + entered, release, closed = threading.Event(), threading.Event(), threading.Event() + + def blocked_write(path: Path, data: memoryview) -> bool: + entered.set() + assert release.wait(5) + path.write_bytes(data) + return True + + def close() -> None: + adapter.close() + closed.set() + + # Force the actual executor-backed O_DIRECT dispatch without host O_DIRECT. + with ( + patch.object(adapter, "_use_odirect", True), + patch.object(adapter, "_os_disk_bs", 4096), + patch.object(adapter, "_write_with_odirect", blocked_write), + ): + adapter.submit_store_task([_long_key()], [_memory_obj(bytes(4096))]) + assert entered.wait(5) + closer = threading.Thread(target=close) + closer.start() + try: + assert not closed.wait(0.2) + finally: + release.set() + closer.join(5) + assert closed.is_set() + + +def test_close_keeps_delete_pending_until_unlink_drains(tmp_path: Path) -> None: + adapter = FSL2Adapter(FSL2AdapterConfig(base_path=str(tmp_path))) + key = _long_key() + _wait_for_store(adapter, key, b"retained payload") + entered, release = threading.Event(), threading.Event() + deleted, closed = threading.Event(), threading.Event() + errors: list[Exception] = [] + + def unlink(path: Path) -> None: + entered.set() + assert release.wait(10) + os.unlink(path) + + async def delayed_unlink(path: Path) -> None: + await asyncio.to_thread(unlink, path) + + def delete() -> None: + try: + adapter.delete([key]) + except Exception as exc: + errors.append(exc) + finally: + deleted.set() + + def close() -> None: + adapter.close() + closed.set() + + # Hold the real executor syscall so cancellation cannot mimic completion. + with patch( + "lmcache.v1.distributed.l2_adapters.fs_l2_adapter.aiofiles.os.unlink", + delayed_unlink, + ): + deletion = threading.Thread(target=delete) + closer = threading.Thread(target=close) + deletion.start() + try: + assert entered.wait(5) + closer.start() + assert not deleted.wait(0.2) + assert not closed.is_set() + finally: + release.set() + deletion.join(5) + if closer.ident is not None: + closer.join(5) + assert deleted.is_set() and closed.is_set() + assert not errors + assert key not in adapter.get_existing_key_sizes() + adapter.close() + + def _long_key(*, model_suffix: str = "", salt_suffix: str = "") -> ObjectKey: return ObjectKey( chunk_hash=b"\xde\xad\xbe\xef" * 8, diff --git a/tests/v1/distributed/test_store_ownership_cpu.py b/tests/v1/distributed/test_store_ownership_cpu.py new file mode 100644 index 00000000000..c1f9619b7be --- /dev/null +++ b/tests/v1/distributed/test_store_ownership_cpu.py @@ -0,0 +1,284 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Persistence ownership, failure and byte bounds without a GPU runtime.""" + +# Standard +from collections.abc import Iterator +from contextlib import contextmanager +from pathlib import Path +import threading +import time +import uuid + +# Third Party +import torch + +# First Party +from lmcache.v1.distributed.api import ( + RECURRENT_CHECKPOINT_MODEL_PREFIX, + MemoryLayoutDesc, + ObjectKey, + PrefetchRequestSpec, +) +from lmcache.v1.distributed.checkpoint_retention import CheckpointRetention +from lmcache.v1.distributed.config import L1ManagerConfig, L1MemoryManagerConfig +from lmcache.v1.distributed.error import L1Error +from lmcache.v1.distributed.l1_manager import L1Manager +from lmcache.v1.distributed.l2_adapters.mock_l2_adapter import MockL2AdapterConfig +from lmcache.v1.distributed.l2_adapters.fs_l2_adapter import ( + FSL2Adapter, + FSL2AdapterConfig, +) +from lmcache.v1.distributed.l2_adapters.native_connector_l2_adapter import ( + NativeConnectorL2Adapter, + _object_key_to_string, +) +from lmcache.v1.distributed.storage_controllers.store_controller import StoreController +from lmcache.v1.distributed.storage_controllers.prefetch_controller import ( + PrefetchController, +) +from lmcache.v1.distributed.storage_controllers.prefetch_policy import ( + DefaultPrefetchPolicy, +) +from lmcache.v1.distributed.storage_controllers.store_policy import ( + AdapterDescriptor, + DefaultStorePolicy, +) +from tests.v1.distributed.test_native_connector_l2_adapter import MockNativeConnector +from tests.v1.distributed.test_fs_l2_adapter_persistence import _memory_obj +from tests.v1.multiprocess.test_checkpoint_storage import wait_for + + +class RecordingConnector(MockNativeConnector): + def __init__(self) -> None: + super().__init__() + self.submissions: list[list[str]] = [] + + def submit_batch_set(self, keys: list[str], memoryviews: list) -> int: + self.submissions.append(list(keys)) + return super().submit_batch_set(keys, memoryviews) + + +@contextmanager +def pipeline( + replicas: int = 1, +) -> Iterator[tuple[L1Manager, StoreController, list[RecordingConnector]]]: + l1 = L1Manager( + L1ManagerConfig( + L1MemoryManagerConfig( + 65536, False, shm_name=f"lmcache_test_{uuid.uuid4().hex}" + ), + read_ttl_seconds=1, + ) + ) + clients = [RecordingConnector() for _ in range(replicas)] + adapters = [ + NativeConnectorL2Adapter(client, initial_key_sizes={}) for client in clients + ] + retention = CheckpointRetention() + for i, adapter in enumerate(adapters): + adapter.register_listener(retention.listener_for(i)) + controller = StoreController( + l1, + adapters, + [ + AdapterDescriptor( + i, MockL2AdapterConfig(max_size_gb=1, mock_bandwidth_gb=1) + ) + for i in range(replicas) + ], + DefaultStorePolicy(), + retention=retention, + max_pending_bytes=8192, + ) + try: + yield l1, controller, clients + finally: + # Tests start the controller after inspecting the queue handoff. + controller.stop() + for adapter in adapters: + adapter.close() + controller.release_stopped_ownership() + l1.close() + + +def write( + l1: L1Manager, count: int = 2, *, checkpoint: bool = False +) -> list[ObjectKey]: + model = ( + RECURRENT_CHECKPOINT_MODEL_PREFIX + "test" if checkpoint else "ordinary-test" + ) + keys = [ObjectKey(ObjectKey.IntHash2Bytes(i), model, 0) for i in range(count)] + layout = MemoryLayoutDesc([torch.Size([4096])], [torch.uint8]) + results = l1.reserve_write(keys, [False] * count, layout) + for error, obj in results.values(): + assert error == L1Error.SUCCESS and obj is not None + obj.tensor.fill_(33) + l1.finish_write(keys) + return keys + + +def test_queue_owns_pages_before_store_loop_and_bounds_bytes() -> None: + with pipeline() as (l1, controller, _): + keys = write(l1, 3) + status = controller.report_status() + assert status["owned_source_bytes"] == status["max_pending_bytes"] == 8192 + assert status["rejected_queue_pages"] == 1 + result = l1.delete(keys) + assert result[keys[0]] == result[keys[1]] == L1Error.KEY_IS_LOCKED + assert result[keys[2]] == L1Error.SUCCESS + controller.start() + wait_for(lambda: not controller.has_pending_work(), "store did not drain") + assert all(value == L1Error.SUCCESS for value in l1.delete(keys[:2]).values()) + + +def test_active_store_ownership_survives_ttl_and_stop() -> None: + with pipeline() as (l1, controller, clients): + clients[0].suppress_set_completion = True + keys = write(l1) + controller.start() + wait_for( + lambda: controller.report_status()["in_flight_task_count"] == 1, + "no submission", + ) + time.sleep(1.1) + controller.cancel_queued(keys) + assert all(value == L1Error.KEY_IS_LOCKED for value in l1.delete(keys).values()) + clients[0].complete_suppressed_sets() + wait_for(lambda: not controller.has_pending_work(), "completion did not drain") + assert all(value == L1Error.SUCCESS for value in l1.delete(keys).values()) + + +def test_cancelled_queue_never_submits_recycled_buffers() -> None: + with pipeline() as (l1, controller, clients): + keys = write(l1) + controller.cancel_queued(keys) + assert all(value == L1Error.SUCCESS for value in l1.delete(keys).values()) + controller.start() + assert not controller.has_pending_work() + assert clients[0].submissions == [] + + +def test_partial_failure_retries_only_missing_pages() -> None: + with pipeline() as (l1, controller, clients): + keys = write(l1, checkpoint=True) + failed = _object_key_to_string(keys[1]) + clients[0].fail_set_keys.add(failed) + clients[0].report_set_per_key_results = True + controller.start() + wait_for(lambda: len(clients[0].submissions) >= 2, "partial store not retried") + assert clients[0].submissions[1] == [failed] + clients[0].fail_set_keys.clear() + wait_for(lambda: not controller.has_pending_work(), "retry did not drain") + assert all(value == L1Error.SUCCESS for value in l1.delete(keys).values()) + + +def test_cancel_and_rewrite_same_keys_does_not_duplicate_io() -> None: + with pipeline() as (l1, controller, clients): + keys = write(l1) + controller.cancel_queued(keys) + assert all(value == L1Error.SUCCESS for value in l1.delete(keys).values()) + assert write(l1) == keys + controller.start() + wait_for( + lambda: not controller.has_pending_work(), "rewritten work did not drain" + ) + assert sum(len(batch) for batch in clients[0].submissions) == len(keys) + + +def test_replication_serializes_destinations_within_byte_budget() -> None: + with pipeline(2) as (l1, controller, clients): + clients[0].suppress_set_completion = True + write(l1) + controller.start() + wait_for(lambda: len(clients[0].submissions) == 1, "first replica missing") + assert not clients[1].submissions + assert controller.report_status()["owned_source_bytes"] == 8192 + clients[0].complete_suppressed_sets() + wait_for(lambda: not controller.has_pending_work(), "replication did not drain") + assert len(clients[1].submissions) == 1 + + +def test_failed_first_destination_does_not_starve_healthy_replica() -> None: + with pipeline(2) as (l1, controller, clients): + clients[0].raise_set_submit = True + write(l1, checkpoint=True) + controller.start() + wait_for(lambda: bool(clients[1].submissions), "healthy replica was starved") + clients[0].raise_set_submit = False + wait_for( + lambda: not controller.has_pending_work(), "recovered replica did not drain" + ) + + +def test_prefetch_owns_destination_through_ttl_and_backend_shutdown( + tmp_path: Path, monkeypatch +) -> None: + l1 = L1Manager( + L1ManagerConfig( + L1MemoryManagerConfig( + 65536, False, shm_name=f"lmcache_test_{uuid.uuid4().hex}" + ), + write_ttl_seconds=1, + ) + ) + config = FSL2AdapterConfig(str(tmp_path)) + adapter = FSL2Adapter(config) + key = ObjectKey(ObjectKey.IntHash2Bytes(1), "prefetch-ownership", 0) + task = adapter.submit_store_task([key], [_memory_obj(bytes(4096))]) + wait_for( + lambda: task in adapter.pop_completed_store_tasks(), + "seed store did not complete", + ) + entered, release, closed = threading.Event(), threading.Event(), threading.Event() + + def blocked_read(path: Path, destination: memoryview) -> int: + entered.set() + assert release.wait(5) + data = path.read_bytes() + destination[:] = data + return len(data) + + monkeypatch.setattr(adapter, "_use_odirect", True) + monkeypatch.setattr(adapter, "_os_disk_bs", 4096) + monkeypatch.setattr(adapter, "_read_with_odirect", blocked_read) + controller = PrefetchController( + l1, [adapter], [AdapterDescriptor(0, config)], DefaultPrefetchPolicy() + ) + controller.start() + controller.submit_prefetch_request( + PrefetchRequestSpec( + [key], {0: MemoryLayoutDesc([torch.Size([4096])], [torch.uint8])} + ) + ) + stopped = False + try: + assert entered.wait(5) + time.sleep(1.1) + assert l1.reclaim_abandoned_writes() == 0 + assert l1.delete([key])[key] == L1Error.KEY_IS_LOCKED + controller.stop() + stopped = True + assert l1.get_memory_usage()[0] == 4096 + + def close_backend() -> None: + adapter.close() + closed.set() + + closer = threading.Thread(target=close_backend) + closer.start() + try: + assert not closed.wait(0.2) + finally: + release.set() + closer.join(5) + assert closed.is_set() + controller.release_stopped_ownership() + assert l1.get_memory_usage()[0] == 0 + finally: + release.set() + if not stopped: + controller.stop() + if not closed.is_set(): + adapter.close() + controller.release_stopped_ownership() + l1.close() diff --git a/tests/v1/multiprocess/test_checkpoint_abandoned_leases.py b/tests/v1/multiprocess/test_checkpoint_abandoned_leases.py index b5d25d84fd8..4d3e86b0e67 100644 --- a/tests/v1/multiprocess/test_checkpoint_abandoned_leases.py +++ b/tests/v1/multiprocess/test_checkpoint_abandoned_leases.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -"""Leases and pending generations of a dead worker are eventually released.""" +"""Stale metadata may expire; exposed copy buffers require drained completion.""" # Standard from dataclasses import replace @@ -21,7 +21,7 @@ ABANDONED_AFTER = 0.2 -def test_abandoned_store_lease_and_generation_are_released() -> None: +def test_abandoned_store_metadata_expires_but_copy_buffers_stay_owned() -> None: with open_store() as (_service, index, _storage, _mapping): service = CheckpointPayloadStore( _storage, index, abandoned_after_seconds=ABANDONED_AFTER @@ -35,7 +35,9 @@ def test_abandoned_store_lease_and_generation_are_released() -> None: time.sleep(ABANDONED_AFTER + 0.1) assert service.reclaim_abandoned() == 1 status = service.report_status() - assert status["store_leases"] == 0 + assert status["store_leases"] == 1 + keys = [key for group in checkpoint_object_keys(entry, 0) for key in group] + assert _storage.delete_l1_keys(keys)[0] == 0 assert index.report_status()["pending_generations"] == 0 assert not service.finish_store(lease.lease_id, True) # The generation can be staged and stored again. @@ -58,7 +60,7 @@ def test_pending_generation_without_a_lease_is_released() -> None: assert index.begin(entry) -def test_abandoned_ready_retrieve_releases_its_read_locks() -> None: +def test_abandoned_ready_retrieve_keeps_copy_ownership() -> None: with open_store() as (_service, index, storage, mapping): service = CheckpointPayloadStore( storage, index, abandoned_after_seconds=ABANDONED_AFTER @@ -76,8 +78,10 @@ def test_abandoned_ready_retrieve_releases_its_read_locks() -> None: # The reader holds the pages until it releases them. assert storage.delete_l1_keys(keys)[0] == 0 time.sleep(ABANDONED_AFTER + 0.1) - assert service.reclaim_abandoned() == 1 - assert service.report_status()["retrieve_leases"] == 0 + assert service.reclaim_abandoned() == 0 + assert service.report_status()["retrieve_leases"] == 1 + assert storage.delete_l1_keys(keys)[0] == 0 + service.finish_retrieve(lease_id) assert storage.delete_l1_keys(keys)[0] == len(keys) diff --git a/tests/v1/multiprocess/test_checkpoint_eviction_lifecycle.py b/tests/v1/multiprocess/test_checkpoint_eviction_lifecycle.py new file mode 100644 index 00000000000..9360b08cdaf --- /dev/null +++ b/tests/v1/multiprocess/test_checkpoint_eviction_lifecycle.py @@ -0,0 +1,224 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Loss under storage pressure must preserve coherent checkpoint lookup.""" + +# Standard +from dataclasses import replace +import hashlib +import time + +# Third Party +import pytest + +# First Party +from tests.v1.multiprocess.test_checkpoint_storage import ( + CheckpointSlots, + drain_l2_stores, + entry_keys, + large_manifest, + make_content_manifest, + make_manifest, + open_store, + publish_all, + restores, + wait_for, +) +from lmcache.v1.distributed.l2_adapters.fs_l2_adapter import ( + FSL2Adapter, + FSL2AdapterConfig, +) + + +def test_last_copy_deletion_retires_only_affected_branch() -> None: + key = lambda value: hashlib.sha256(value.encode()).hexdigest() + older = make_content_manifest( + "older", prefix_token=3, attention_keys=(key("a"), key("b"), key("old")) + ) + newer = make_content_manifest( + "newer", prefix_token=4, attention_keys=(key("a"), key("b"), key("new")) + ) + with open_store() as (service, index, storage, mapping): + publish_all(service, index, mapping, older) + publish_all(service, index, mapping, newer) + unique = list(set(entry_keys(older)) - set(entry_keys(newer))) + assert storage.delete_l1_keys(unique)[0] == len(unique) + assert index.find((older.prefix,)) is None + assert index.find((newer.prefix,)) == newer + assert restores(service, mapping, newer) + + +def test_clear_retires_ram_only_manifests() -> None: + with open_store() as (service, index, storage, mapping): + entry = make_manifest() + publish_all(service, index, mapping, entry) + storage.clear() + assert index.find((entry.prefix,)) is None + + +@pytest.mark.parametrize("native", [False, True]) +def test_disk_failure_retires_cold_checkpoints_and_admits_new_work( + tmp_path, monkeypatch, native +) -> None: + with open_store( + tmp_path, + native, + store_policy="checkpoint_on_evict", + shutdown_flush_seconds=0, + admission_timeout_seconds=2, + ) as resources: + service, index, storage, mapping = resources + adapter = storage.l2_adapters()[0][1] + + def unavailable(*_args, **_kwargs): + raise OSError("injected disk unavailable") + + monkeypatch.setattr(adapter, "submit_store_task", unavailable) + entries = [large_manifest(500 + i) for i in range(7)] + started = time.monotonic() + for entry in entries: + publish_all(service, index, mapping, entry) + assert time.monotonic() - started < 15 + assert storage.get_admission_stats()["exhausted_timeouts"] == 0 + assert storage.checkpoint_retention.report_status()["retired_checkpoints"] > 0 + assert index.find((entries[-1].prefix,)) == entries[-1] + assert restores(service, mapping, entries[-1]) + for entry in entries: + if index.find((entry.prefix,)) is not None: + assert set(storage.get_readable_keys(entry_keys(entry))) == set( + entry_keys(entry) + ) + status = storage.report_status() + # All work is bounded independently of how long failures persist. + store_status = status["store_controller"] + assert store_status["owned_source_bytes"] <= store_status["max_pending_bytes"] + + +@pytest.mark.parametrize("native", [False, True]) +def test_transient_submission_failure_retries_and_restores( + tmp_path, monkeypatch, native +) -> None: + with open_store( + tmp_path, native, store_policy="checkpoint_on_evict", shutdown_flush_seconds=0 + ) as (service, index, storage, mapping): + entry = replace(make_manifest(), world_size=1) + publish_all(service, index, mapping, entry) + adapter = storage.l2_adapters()[0][1] + submit = adapter.submit_store_task + calls = [] + + def fail_once(keys, objects): + calls.append(list(keys)) + if len(calls) == 1: + raise OSError("injected transient submission failure") + return submit(keys, objects) + + monkeypatch.setattr(adapter, "submit_store_task", fail_once) + storage.checkpoint_retention.request_persist(entry_keys(entry)) + wait_for( + lambda: all( + storage.checkpoint_retention.is_l2_resident(key) + for key in entry_keys(entry) + ), + "retry did not persist checkpoint", + ) + drain_l2_stores(storage) + assert len(calls) > 1 + storage.clear() + assert index.find((entry.prefix,)) == entry + assert restores(service, mapping, entry) + + +def test_zero_shutdown_budget_removes_ram_only_directory_entries(tmp_path) -> None: + entry = make_manifest() + with open_store( + tmp_path, True, store_policy="checkpoint_on_evict", shutdown_flush_seconds=0 + ) as (service, index, storage, mapping): + publish_all(service, index, mapping, entry) + assert index.find((entry.prefix,)) == entry + with open_store(tmp_path, True) as (_, index, _, _): + assert index.find((entry.prefix,)) is None + + +def test_duplicate_completed_rank_cannot_accumulate_pins() -> None: + entry = make_manifest() + with open_store() as (service, index, _, mapping): + assert index.begin(entry) + lease = service.prepare_store(entry, 0) + assert isinstance(lease, CheckpointSlots) + for group in lease.groups: + for slot in group: + assert slot is not None + mapping[slot.offset : slot.offset + slot.length] = b"\0" * slot.length + assert not service.finish_store(lease.lease_id, True) + with pytest.raises(ValueError, match="already has a store lease"): + service.prepare_store(entry, 0) + index.abort(entry.generation) + + +def test_unknown_inventory_preserves_lookup_validation(tmp_path, monkeypatch) -> None: + entry = replace(make_manifest(), world_size=1) + with open_store(tmp_path) as (service, index, storage, mapping): + publish_all(service, index, mapping, entry) + drain_l2_stores(storage) + monkeypatch.setattr(FSL2Adapter, "has_complete_inventory", lambda _self: False) + monkeypatch.setattr(FSL2Adapter, "get_existing_key_sizes", lambda _self: {}) + with open_store(tmp_path) as (service, index, _, mapping): + assert index.find((entry.prefix,)) == entry + assert restores(service, mapping, entry) + + +def test_terminal_shutdown_flushes_complete_state_and_keeps_orphan_buffers( + tmp_path, +) -> None: + complete = replace(make_manifest(), world_size=1) + orphan = replace(make_manifest(), world_size=1) + with open_store(tmp_path, True, store_policy="checkpoint_on_evict") as resources: + service, index, storage, mapping = resources + publish_all(service, index, mapping, complete) + assert index.begin(orphan) + lease = service.prepare_store(orphan, 0) + assert isinstance(lease, CheckpointSlots) + service.prepare_terminal_shutdown() + keys = entry_keys(orphan) + assert storage.delete_l1_keys(keys) == (0, len(keys)) + assert service.prepare_store(orphan, 0) is None + assert not service.finish_store(lease.lease_id, False) + with open_store(tmp_path, True) as (service, index, _, mapping): + assert index.find((complete.prefix,)) == complete + assert restores(service, mapping, complete) + + +def test_persistence_retries_after_last_adapter_is_replaced(tmp_path) -> None: + entry = replace(make_manifest(), world_size=1) + with open_store(tmp_path, store_policy="checkpoint_on_evict") as resources: + service, index, storage, mapping = resources + publish_all(service, index, mapping, entry) + storage.delete_l2_adapter(storage.l2_adapters()[0][0].index) + retention = storage.checkpoint_retention + assert retention.request_persist(entry_keys(entry)) > 0 + wait_for( + lambda: retention.pending_count() == 0, + "unroutable persistence stayed pending", + ) + storage.add_l2_adapter(FSL2AdapterConfig(str(tmp_path / "replacement"))) + assert retention.request_persist(entry_keys(entry)) > 0 + wait_for( + lambda: all(retention.is_l2_resident(key) for key in entry_keys(entry)), + "replacement adapter never received retained pages", + ) + + +def test_terminal_shutdown_keeps_exposed_retrieve_ownership(tmp_path) -> None: + from tests.v1.multiprocess.test_checkpoint_storage import poll + + entry = replace(make_manifest(), world_size=1) + with open_store(tmp_path, True, store_policy="checkpoint_on_evict") as resources: + service, index, storage, mapping = resources + publish_all(service, index, mapping, entry) + lease = service.begin_retrieve(entry, 0) + assert lease is not None and isinstance(poll(service, lease), CheckpointSlots) + service.prepare_terminal_shutdown() + assert storage.delete_l1_keys(entry_keys(entry))[0] == 0 + service.finish_retrieve(lease) + with open_store(tmp_path, True) as (service, index, _, mapping): + assert index.find((entry.prefix,)) == entry + assert restores(service, mapping, entry) diff --git a/tests/v1/multiprocess/test_checkpoint_publication_ownership.py b/tests/v1/multiprocess/test_checkpoint_publication_ownership.py new file mode 100644 index 00000000000..d375ee8fe3a --- /dev/null +++ b/tests/v1/multiprocess/test_checkpoint_publication_ownership.py @@ -0,0 +1,126 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Publication must own reused pages and every completed rank until commit.""" + +# Standard +from dataclasses import replace +import hashlib +import time + +# First Party +from tests.v1.multiprocess.test_checkpoint_storage import ( + CheckpointSlots, + checkpoint_object_keys, + make_content_manifest, + make_manifest, + open_store, + poll, + publish_all, +) + + +def fill(mapping, lease: CheckpointSlots) -> None: + for group in lease.groups: + for slot in group: + if slot is not None: + mapping[slot.offset : slot.offset + slot.length] = b"\x33" * slot.length + + +def test_reused_pages_remain_owned_until_publication() -> None: + key = lambda value: hashlib.sha256(value.encode()).hexdigest() + older = make_content_manifest( + "older", prefix_token=3, attention_keys=(key("a"), key("b"), key("old")) + ) + newer = make_content_manifest( + "newer", prefix_token=4, attention_keys=(key("a"), key("b"), key("new")) + ) + with open_store() as (service, index, storage, mapping): + publish_all(service, index, mapping, older) + assert index.begin(newer) + lease = service.prepare_store(newer, 0) + assert isinstance(lease, CheckpointSlots) + shared = list(checkpoint_object_keys(newer, 0)[0][:2]) + assert lease.groups[0][:2] == (None, None) + assert storage.delete_l1_keys(shared, force=False) == (0, 2) + fill(mapping, lease) + assert service.finish_store(lease.lease_id, True) + retrieve = service.begin_retrieve(newer, 0) + assert retrieve is not None + restored = poll(service, retrieve) + assert isinstance(restored, CheckpointSlots) + service.finish_retrieve(restored.lease_id) + + +def test_completed_rank_stays_owned_until_all_ranks_publish() -> None: + manifest = replace(make_manifest(), world_size=2) + with open_store() as (service, index, storage, mapping): + assert index.begin(manifest) + first = service.prepare_store(manifest, 0) + assert isinstance(first, CheckpointSlots) + fill(mapping, first) + assert not service.finish_store(first.lease_id, True) + keys = [key for group in checkpoint_object_keys(manifest, 0) for key in group] + assert storage.delete_l1_keys(keys, force=False) == (0, len(keys)) + second = service.prepare_store(manifest, 1) + assert isinstance(second, CheckpointSlots) + fill(mapping, second) + assert service.finish_store(second.lease_id, True) + assert storage.delete_l1_keys(keys, force=False) == (len(keys), 0) + + +def test_abort_releases_completed_rank_but_not_live_writer() -> None: + manifest = replace(make_manifest(), world_size=2) + with open_store() as (service, index, storage, mapping): + assert index.begin(manifest) + first = service.prepare_store(manifest, 0) + second = service.prepare_store(manifest, 1) + assert isinstance(first, CheckpointSlots) + assert isinstance(second, CheckpointSlots) + fill(mapping, first) + assert not service.finish_store(first.lease_id, True) + index.abort(manifest.generation) + keys = [key for group in checkpoint_object_keys(manifest, 0) for key in group] + live = [key for group in checkpoint_object_keys(manifest, 1) for key in group] + assert storage.delete_l1_keys(keys, force=False) == (len(keys), 0) + assert storage.delete_l1_keys(live, force=False) == (0, len(live)) + assert not service.finish_store(second.lease_id, False) + assert not service.finish_store(second.lease_id, False) + + +def test_copy_ownership_survives_client_lock_expiry() -> None: + manifest = replace(make_manifest(), world_size=1) + with open_store(lock_ttl_seconds=1) as (service, index, storage, mapping): + assert index.begin(manifest) + lease = service.prepare_store(manifest, 0) + assert isinstance(lease, CheckpointSlots) + keys = [key for group in checkpoint_object_keys(manifest, 0) for key in group] + time.sleep(1.1) + assert storage.delete_l1_keys(keys) == (0, len(keys)) + fill(mapping, lease) + assert service.finish_store(lease.lease_id, True) + retrieve = service.begin_retrieve(manifest, 0) + assert retrieve is not None + assert isinstance(poll(service, retrieve), CheckpointSlots) + time.sleep(1.1) + assert storage.delete_l1_keys(keys) == (0, len(keys)) + service.finish_retrieve(retrieve) + assert storage.delete_l1_keys(keys) == (len(keys), 0) + + +def test_quiesce_prevents_admission_after_idle_check() -> None: + with open_store() as (service, index, _, _): + manifest = make_manifest() + assert index.begin(manifest) + service.quiesce() + assert service.prepare_store(manifest, 0) is None + assert service.begin_retrieve(manifest, 0) is None + + +def test_failed_copy_releases_internal_writer_after_client_ttl() -> None: + manifest = replace(make_manifest(), world_size=1) + with open_store(lock_ttl_seconds=1) as (service, index, storage, _): + assert index.begin(manifest) + lease = service.prepare_store(manifest, 0) + assert isinstance(lease, CheckpointSlots) + time.sleep(1.1) + assert not service.finish_store(lease.lease_id, False) + assert storage.get_l1_usage()[0] == 0 diff --git a/tests/v1/multiprocess/test_checkpoint_review_regressions.py b/tests/v1/multiprocess/test_checkpoint_review_regressions.py new file mode 100644 index 00000000000..a09dfd2968f --- /dev/null +++ b/tests/v1/multiprocess/test_checkpoint_review_regressions.py @@ -0,0 +1,536 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Counterexamples from the independent cache-integrity review.""" + +# Standard +from dataclasses import replace +from pathlib import Path +from types import SimpleNamespace +from typing import Any +from unittest.mock import patch +from unittest.mock import Mock +import asyncio +import hashlib +import threading + +# Third Party +import pytest +import torch + +# First Party +from lmcache.v1.distributed.api import ( + EncodedObjectKey, + MemoryLayoutDesc, + ObjectKey, + Tier, +) +from lmcache.v1.distributed.admission import AdmissionFailure +from lmcache.v1.multiprocess.cache_control.object_service import ObjectService +from lmcache.integration.vllm.checkpoint_scheduler import ( + CheckpointSchedulerBridge, + _Lookup, + _PendingTask, +) +from lmcache.v1.multiprocess.checkpoint_identity import CheckpointTokenRoots +from lmcache.v1.distributed.l2_adapters.fs_l2_adapter import ( + FSL2AdapterConfig, + _object_key_to_relative_path, +) +from lmcache.v1.distributed.l2_adapters.native_connector_l2_adapter import ( + NativeConnectorL2Adapter, +) +from tests.v1.distributed.test_native_connector_l2_adapter import ( + MockNativeConnector, + create_object_key, +) +from tests.v1.multiprocess.test_checkpoint_storage import ( + CheckpointSlots, + drain_l2_stores, + entry_keys, + large_manifest, + make_content_manifest, + make_manifest, + open_store, + poll, + publish_all, + publish_before_restart, + restores, + wait_for, +) + + +@pytest.mark.parametrize("native", [False, True]) +def test_capacity_return_after_failed_load_preserves_checkpoint( + tmp_path: Path, native: bool +) -> None: + entry = large_manifest(550, page_bytes=128 * 1024) + publish_before_restart(tmp_path, native, [entry]) + with open_store(tmp_path, native, admission_timeout_seconds=0.15) as ( + service, + index, + storage, + mapping, + ): + filler = ObjectKey(hashlib.sha256(b"active-writer").digest(), "filler", 0, 0) + _, total = storage.get_l1_usage() + lookup = storage.submit_prefetch_task + query = storage.query_prefetch_status_detailed + attempts = 0 + released = set() + + def reserve_then_lookup(*args: Any, **kwargs: Any) -> Any: + nonlocal attempts + attempts += 1 + result = storage.reserve_write_detailed( + [filler], + MemoryLayoutDesc([torch.Size([total])], [torch.uint8]), + "new", + internal=True, + ) + assert result[filler][1] is not None + return lookup(*args, **kwargs) + + def release_after_lookup(handle: Any) -> Any: + result = query(handle) + if result is not None and attempts not in released: + assert result.reservation_failed + storage.abort_write([filler]) + released.add(attempts) + return result + + with ( + patch.object( + storage, "submit_prefetch_task", side_effect=reserve_then_lookup + ), + patch.object( + storage, + "query_prefetch_status_detailed", + side_effect=release_after_lookup, + ), + ): + assert poll(service, service.begin_retrieve(entry, 0)) is False + assert attempts >= 2 + assert index.get(entry.generation) == entry + assert restores(service, mapping, entry) + + +@pytest.mark.parametrize("native", [False, True]) +def test_public_l2_delete_retires_last_copy(tmp_path: Path, native: bool) -> None: + with open_store(tmp_path, native) as (service, index, storage, mapping): + entry = replace(make_manifest(), world_size=1) + publish_all(service, index, mapping, entry) + drain_l2_stores(storage) + storage.clear() + keys = entry_keys(entry) + encoded = [ + EncodedObjectKey( + k.chunk_hash.hex(), + k.model_name, + k.kv_rank, + k.object_group_id, + k.cache_salt, + ) + for k in keys + ] + result = asyncio.run( + ObjectService(SimpleNamespace(storage_manager=storage)).delete_objects( + Tier.L2, None, encoded + ) + ) + assert result == {"deleted": len(keys), "skipped": 0, "ok": True} + assert index.find((entry.prefix,)) is None + + +@pytest.mark.parametrize("native", [False, True]) +@pytest.mark.parametrize("policy", ["default", "retain"]) +def test_temporary_restore_is_not_a_retained_replica( + tmp_path: Path, native: bool, policy: str +) -> None: + with open_store( + tmp_path, + native, + store_policy="checkpoint_on_evict", + prefetch_policy=policy, + shutdown_flush_seconds=0, + ) as (service, index, storage, mapping): + entry = make_content_manifest( + "temporary-source", + prefix_token=1, + attention_keys=tuple( + hashlib.sha256(bytes([i])).hexdigest() for i in range(3) + ), + ) + keys = entry_keys(entry) + publish_all(service, index, mapping, entry) + retention = storage.checkpoint_retention + retention.request_persist(keys) + wait_for(lambda: all(retention.is_l2_resident(k) for k in keys), "persist") + drain_l2_stores(storage) + storage.clear() + lease = poll(service, service.begin_retrieve(entry, 0)) + assert isinstance(lease, CheckpointSlots) + descriptor, adapter = storage.l2_adapters()[0] + retention.evict(keys, descriptor.index, adapter.delete) + assert not any(retention.is_l2_resident(k) for k in keys) + if policy == "default": + assert index.find((entry.prefix,)) is None + # A new publication may not inherit pages freed with this reader. + replacement = replace(entry, generation=entry.generation + "-new") + assert index.begin(replacement) + assert service.prepare_store(replacement, 0) is AdmissionFailure.BUSY + index.abort(replacement.generation) + else: + assert index.find((entry.prefix,)) == entry + service.finish_retrieve(lease.lease_id) + assert index.find((entry.prefix,)) is None or len( + storage.get_readable_keys(keys) + ) == len(keys) + + +@pytest.mark.parametrize("via_api", [False, True]) +def test_store_delete_notifications_follow_physical_order( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, via_api: bool +) -> None: + with open_store( + tmp_path, store_policy="checkpoint_on_evict", shutdown_flush_seconds=0 + ) as (service, index, storage, mapping): + entry = replace(make_manifest(), world_size=1) + publish_all(service, index, mapping, entry) + descriptor, adapter = storage.l2_adapters()[0] + keys = entry_keys(entry) + first, second = keys[:2] + blocked, resume = threading.Event(), threading.Event() + original = adapter._existing_key_path + + async def barrier(key: ObjectKey) -> Any: + if key == second and not blocked.is_set(): + blocked.set() + await asyncio.to_thread(resume.wait, 10) + return await original(key) + + monkeypatch.setattr(adapter, "_existing_key_path", barrier) + storage.checkpoint_retention.request_persist(keys) + try: + assert blocked.wait(5) + if via_api: + encoded = EncodedObjectKey( + first.chunk_hash.hex(), + first.model_name, + first.kv_rank, + first.object_group_id, + first.cache_salt, + ) + result = asyncio.run( + ObjectService( + SimpleNamespace(storage_manager=storage) + ).delete_objects(Tier.L2, None, [encoded]) + ) + assert result["ok"] + else: + storage.checkpoint_retention.evict( + [first], descriptor.index, adapter.delete + ) + assert first not in adapter.get_existing_key_sizes() + finally: + resume.set() + drain_l2_stores(storage) + assert storage.checkpoint_retention.is_l2_resident(first) == ( + first in adapter.get_existing_key_sizes() + ) + storage.clear() + assert index.find((entry.prefix,)) is None or restores(service, mapping, entry) + + +@pytest.mark.parametrize("native", [False, True]) +def test_aligned_capacity_shortage_preserves_disk_checkpoint( + tmp_path: Path, native: bool +) -> None: + entry = replace(make_manifest(), world_size=1) + publish_before_restart(tmp_path, native, [entry]) + with open_store(tmp_path, native, admission_timeout_seconds=0.15) as ( + service, + index, + storage, + mapping, + ): + filler = ObjectKey(hashlib.sha256(b"alignment-filler").digest(), "filler", 0, 0) + _, total = storage.get_l1_usage() + reserved = storage.reserve_write_detailed( + [filler], + MemoryLayoutDesc([torch.Size([total - 4096])], [torch.uint8]), + "new", + internal=True, + ) + assert reserved[filler][1] is not None + try: + assert poll(service, service.begin_retrieve(entry, 0)) is False + assert index.get(entry.generation) == entry + finally: + storage.abort_write([filler]) + assert restores(service, mapping, entry) + + +def test_expired_prefetch_ownership_does_not_invalidate_disk_checkpoint( + tmp_path: Path, +) -> None: + entry = replace(make_manifest(), world_size=1) + publish_before_restart(tmp_path, False, [entry]) + with open_store(tmp_path) as (service, index, storage, mapping): + pin = storage.pin_readable_keys + first = True + + def expired_once(keys: list[ObjectKey], **kwargs: Any) -> list[ObjectKey]: + nonlocal first + if first: + first = False + return [] + return pin(keys, **kwargs) + + with patch.object(storage, "pin_readable_keys", side_effect=expired_once): + assert restores(service, mapping, entry) + assert index.get(entry.generation) == entry + + +class _DelayedDelete(MockNativeConnector): + """Model native close draining a syscall and discarding its completion.""" + + def __init__(self) -> None: + super().__init__() + self.submitted = threading.Event() + self.draining = threading.Event() + self.release = threading.Event() + + def submit_batch_delete(self, keys: list[str]) -> int: + self.submitted.set() + return 1 + + def close(self) -> None: + self.draining.set() + assert self.release.wait(10) + super().close() + + +@pytest.mark.parametrize("remove_adapter", [False, True]) +def test_native_close_settles_delete_only_after_backend_drain( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, remove_adapter: bool +) -> None: + client = _DelayedDelete() + key = create_object_key(987) + adapter = NativeConnectorL2Adapter(client, initial_key_sizes={key: 4096}) + monkeypatch.setattr( + "lmcache.v1.distributed.storage_manager.create_l2_adapter", lambda *_: adapter + ) + errors: list[Exception] = [] + deleted = threading.Event() + closed = threading.Event() + + with open_store(tmp_path, shutdown_flush_seconds=0) as (_, _, storage, _): + tier = storage.l2_adapters()[0][0].index + + def delete() -> None: + try: + storage.checkpoint_retention.evict([key], tier, adapter.delete) + except Exception as exc: + errors.append(exc) + finally: + deleted.set() + + def close() -> None: + if remove_adapter: + storage.delete_l2_adapter(tier) + else: + adapter.close() + closed.set() + + deletion = threading.Thread(target=delete) + closer = threading.Thread(target=close) + deletion.start() + try: + assert client.submitted.wait(5) + closer.start() + assert client.draining.wait(5) + assert not deleted.wait(0.1) + assert not closed.is_set() + with pytest.raises(RuntimeError, match="closing"): + adapter.delete([key]) + finally: + client.release.set() + deletion.join(5) + if closer.ident is not None: + closer.join(5) + assert deleted.is_set() and closed.is_set() + assert len(errors) == 1 + assert "outcome unavailable" in str(errors[0]) + adapter.close() + + +@pytest.mark.parametrize("native", [False, True]) +@pytest.mark.parametrize("via_api", [False, True]) +def test_duplicate_delete_releases_every_fence( + tmp_path: Path, native: bool, via_api: bool +) -> None: + with open_store(tmp_path, native, shutdown_flush_seconds=0) as ( + service, + index, + storage, + mapping, + ): + entry = make_content_manifest( + "duplicate-source", + prefix_token=1, + attention_keys=tuple( + hashlib.sha256(bytes([i])).hexdigest() for i in range(3) + ), + ) + publish_all(service, index, mapping, entry) + drain_l2_stores(storage) + keys = entry_keys(entry) + selected = [keys[0], keys[0], keys[1]] + if via_api: + encoded = [ + EncodedObjectKey( + k.chunk_hash.hex(), + k.model_name, + k.kv_rank, + k.object_group_id, + k.cache_salt, + ) + for k in selected + ] + result = asyncio.run( + ObjectService(SimpleNamespace(storage_manager=storage)).delete_objects( + Tier.L2, None, encoded + ) + ) + assert result == {"deleted": 2, "skipped": 0, "ok": True} + else: + descriptor, adapter = storage.l2_adapters()[0] + storage.checkpoint_retention.evict( + selected, descriptor.index, adapter.delete + ) + new = replace(entry, generation="duplicate-next") + assert index.begin(new) + lease = service.prepare_store(new, 0) + assert isinstance(lease, CheckpointSlots) + service.finish_store(lease.lease_id, False) + + +def _retry_roots( + roots: CheckpointTokenRoots, manifest: Any, *, copy_succeeded: bool = False +) -> tuple[Any, ...]: + """Exercise the actual all-rank failure transition with an in-memory RPC sink.""" + client, manager = Mock(), Mock() + manager.acknowledge_external_boundary_checkpoint.return_value = False + bridge = CheckpointSchedulerBridge(manager, client, {}, 1) + # Seed completed transfer state without constructing a GPU allocator. + bridge._lookups["consumer"] = _Lookup(roots, Mock()) + bridge._tasks["failed"] = _PendingTask( + SimpleNamespace(direction="RETRIEVE", manifest=manifest), + SimpleNamespace(checkpoint_id=1, num_tokens=manifest.prefix.num_tokens), + "consumer", + ) + bridge.complete({"failed": {0: copy_succeeded}}) + assert bridge._lookups["consumer"].roots is roots + return client.submit_request.call_args.args[1][0] + + +@pytest.mark.parametrize("length", [1, 4095, 4096, 4097, 8193]) +def test_failed_restore_bounds_query_without_changing_identity(length: int) -> None: + roots = CheckpointTokenRoots.build("retry", list(range(length + 3))) + manifest = replace(make_manifest(), prefix=roots.prefix(length)) + query = _retry_roots(roots, manifest) + expected = CheckpointTokenRoots.build("retry", list(range(length - 1))).roots + assert query == expected + + +def test_publication_race_retries_same_checkpoint_without_longer_candidates() -> None: + roots = CheckpointTokenRoots.build("retry", list(range(12))) + manifest = replace(make_manifest(), prefix=roots.prefix(11)) + expected = CheckpointTokenRoots.build("retry", list(range(11))).roots + assert _retry_roots(roots, manifest, copy_succeeded=True) == expected + + +@pytest.mark.parametrize("native", [False, True]) +def test_failed_restore_uses_intact_shorter_without_forgetting_longest( + tmp_path: Path, native: bool +) -> None: + roots = CheckpointTokenRoots.build("retry", list(range(12))) + longer = replace(make_manifest(), world_size=1, prefix=roots.prefix(11)) + shorter = replace(make_manifest(), world_size=1, prefix=roots.prefix(8)) + with open_store(tmp_path, native, admission_timeout_seconds=0.1) as ( + service, + index, + storage, + mapping, + ): + publish_all(service, index, mapping, shorter) + publish_all(service, index, mapping, longer) + drain_l2_stores(storage) + storage.clear() + # Corrupt only this test's payload after the startup inventory was read. + path = ( + tmp_path / "payloads" / _object_key_to_relative_path(entry_keys(longer)[-1]) + ) + path.write_bytes(b"") + assert index.find(roots.roots) == longer + assert not restores(service, mapping, longer) + assert index.get(longer.generation) == longer + candidate = index.find(_retry_roots(roots, longer)) + assert candidate == shorter + assert restores(service, mapping, candidate) + + +@pytest.mark.parametrize("native", [False, True]) +def test_recovery_retires_known_truncated_payload(tmp_path: Path, native: bool) -> None: + entry = replace(make_manifest(), world_size=1) + publish_before_restart(tmp_path, native, [entry]) + key = entry_keys(entry)[-1] + path = tmp_path / "payloads" / _object_key_to_relative_path(key) + path.write_bytes(b"x") + with open_store(tmp_path, native) as (_, index, _, _): + assert index.get(entry.generation) is None + + +@pytest.mark.parametrize("another_complete_copy", [False, True]) +def test_malformed_replica_cannot_certify_last_copy( + tmp_path: Path, another_complete_copy: bool +) -> None: + with open_store( + tmp_path, store_policy="checkpoint_on_evict", shutdown_flush_seconds=0 + ) as (service, index, storage, mapping): + entry = replace(make_manifest(), world_size=1) + publish_all(service, index, mapping, entry) + key = entry_keys(entry)[-1] + retention = storage.checkpoint_retention + retention.record_l2_present(100, [key], [1]) + if another_complete_copy: + retention.record_l2_present(101, [key], [96]) + assert storage.delete_l1_keys([key]) == (1, 0) + assert (index.get(entry.generation) is not None) == another_complete_copy + + +@pytest.mark.parametrize("native", [False, True]) +@pytest.mark.parametrize("first_copy", ["known_short", "unknown_size"]) +def test_prefetch_selects_usable_replica( + tmp_path: Path, native: bool, first_copy: str +) -> None: + entry = replace(make_manifest(), world_size=1) + with open_store(tmp_path, native, admission_timeout_seconds=0.1) as ( + service, + index, + storage, + mapping, + ): + storage.add_l2_adapter(FSL2AdapterConfig(str(tmp_path / "backup"))) + publish_all(service, index, mapping, entry) + drain_l2_stores(storage) + key = entry_keys(entry)[-1] + relative = _object_key_to_relative_path(key) + assert (tmp_path / "backup" / relative).stat().st_size == 96 + if first_copy == "known_short": + (tmp_path / "payloads" / relative).write_bytes(b"x") + storage.checkpoint_retention.record_l2_present(0, [key], [1]) + else: + storage.checkpoint_retention.record_l2_absent(0, [key]) + storage.clear() + assert index.get(entry.generation) == entry + assert restores(service, mapping, entry) diff --git a/tests/v1/multiprocess/test_checkpoint_storage.py b/tests/v1/multiprocess/test_checkpoint_storage.py index 2f472f0a8ad..4df853dcd11 100644 --- a/tests/v1/multiprocess/test_checkpoint_storage.py +++ b/tests/v1/multiprocess/test_checkpoint_storage.py @@ -84,22 +84,29 @@ def open_store( shutdown_flush_seconds: float = 30.0, admission_timeout_seconds: float = 8.0, prefetch_policy: str = "default", + lock_ttl_seconds: int = 300, + disk_capacity_gb: float = 0, ) -> Iterator[tuple[CheckpointPayloadStore, CheckpointIndex, StorageManager, mmap]]: name = shm_name or f"lmcache_l1_pool_checkpoint_test_{uuid.uuid4().hex}" size = 4 * 1024 * 1024 + adapters = [] + if path is not None: + adapter = (FSNativeL2AdapterConfig if native else FSL2AdapterConfig)( + str(path / "payloads"), + **({"max_capacity_gb": disk_capacity_gb} if native else {}), + ) + if disk_capacity_gb: + adapter.eviction_config = EvictionConfig(eviction_policy="LRU") + adapters.append(adapter) storage = StorageManager( StorageManagerConfig( - L1ManagerConfig(L1MemoryManagerConfig(size, False, shm_name=name)), - EvictionConfig(eviction_policy="LRU"), - l2_adapter_config=L2AdaptersConfig( - [ - (FSNativeL2AdapterConfig if native else FSL2AdapterConfig)( - str(path / "payloads") - ) - ] - if path is not None - else [] + L1ManagerConfig( + L1MemoryManagerConfig(size, False, shm_name=name), + read_ttl_seconds=lock_ttl_seconds, + write_ttl_seconds=lock_ttl_seconds, ), + EvictionConfig(eviction_policy="LRU"), + l2_adapter_config=L2AdaptersConfig(adapters), store_policy=store_policy, prefetch_policy=prefetch_policy, checkpoint_shutdown_flush_seconds=shutdown_flush_seconds, @@ -692,7 +699,7 @@ def fail_commit(_keys: list[Any]) -> None: raise RuntimeError("injected storage commit failure") with monkeypatch.context() as patch: - patch.setattr(storage, "finish_write", fail_commit) + patch.setattr(storage, "finish_write_pinned", fail_commit) with pytest.raises(RuntimeError, match="injected storage commit failure"): service.finish_store(lease.lease_id, True) assert index.find((entry.prefix,)) is None @@ -785,7 +792,8 @@ def test_payload_capacity_miss_preserves_a_pinned_checkpoint(store: Any) -> None def fill(mapping: mmap, lease: CheckpointSlots, rank: int) -> None: for group_id, group in enumerate(lease.groups): for page_id, slot in enumerate(group): - assert slot is not None + if slot is None: + continue mapping[slot.offset : slot.offset + slot.length] = ( bytes([rank * 16 + group_id * 4 + page_id]) * slot.length ) @@ -881,6 +889,8 @@ def test_http_nonforced_clear_preserves_checkpoint_read_and_write_leases( key for group in checkpoint_object_keys(pending, 0) for key in group ] assert storage.get_readable_keys(committed) == committed + # A completed rank still owns its pages until its peers publish or abort. + index.abort(pending.generation) response = client.post("/cache/clear", json={"force": False}) assert response.status_code == 200 assert storage.get_readable_keys(retained + committed) == [] @@ -1180,6 +1190,7 @@ def drain_l2_stores(storage: StorageManager) -> None: == status["pending_reused_keys_count"] == status["in_flight_task_count"] == 0 + and status.get("owned_source_bytes", 0) == 0 ): return time.sleep(0.01) @@ -1233,7 +1244,7 @@ def test_reuse_admission_persists_only_restored_checkpoints( ) service.finish_retrieve(lease.lease_id) # Never restored before the restart, so its pages never reached L2. - assert index.find((unused.prefix,)) == unused + assert index.find((unused.prefix,)) is None lease_id = service.begin_retrieve(unused, 0) assert lease_id is not None assert poll(service, lease_id) is False @@ -1337,10 +1348,10 @@ def test_evict_admission_writes_checkpoint_only_when_l1_evicts_it( @pytest.mark.parametrize("native", [False, True]) -def test_superseded_checkpoint_is_dropped_first_and_never_written( +def test_superseded_checkpoint_is_persisted_before_ordinary_eviction( tmp_path: Path, native: bool ) -> None: - """An older turn's unique pages leave L1 first and never reach L2.""" + """Supersession prioritizes eviction without destroying a retained branch.""" older, newer = large_manifest(200), large_manifest(201) with open_store( tmp_path, native, store_policy="checkpoint_on_evict", shutdown_flush_seconds=0 @@ -1356,7 +1367,7 @@ def test_superseded_checkpoint_is_dropped_first_and_never_written( "current checkpoint pages never reached L2", ) drain_l2_stores(storage) - assert not any(retention.is_l2_resident(key) for key in entry_keys(older)) + assert all(retention.is_l2_resident(key) for key in entry_keys(older)) assert retention.report_status()["l1_superseded_drops"] >= len( entry_keys(older) ) @@ -1364,7 +1375,7 @@ def test_superseded_checkpoint_is_dropped_first_and_never_written( tmp_path, native, store_policy="checkpoint_on_evict", shutdown_flush_seconds=0 ) as (service, index, storage, mapping): assert restores(service, mapping, newer) - assert not restores(service, mapping, older) + assert restores(service, mapping, older) @pytest.mark.parametrize("flush_seconds", [0.0, 30.0]) @@ -1389,7 +1400,7 @@ def test_shutdown_flush_writes_current_checkpoints_for_the_next_start( tmp_path, True, store_policy="checkpoint_on_evict", shutdown_flush_seconds=0 ) as (service, index, storage, mapping): assert restores(service, mapping, current) == (flush_seconds > 0) - assert not restores(service, mapping, superseded) + assert restores(service, mapping, superseded) == (flush_seconds > 0) def publish_before_restart( diff --git a/tests/v1/multiprocess/test_checkpoint_write_volume.py b/tests/v1/multiprocess/test_checkpoint_write_volume.py new file mode 100644 index 00000000000..f98b643c8fe --- /dev/null +++ b/tests/v1/multiprocess/test_checkpoint_write_volume.py @@ -0,0 +1,97 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Measure payload writes and restore outcomes before and after final drain.""" + +# Standard +from dataclasses import replace +from pathlib import Path +import json + +# Third Party +import pytest + +# First Party +from lmcache.v1.distributed.api import ObjectKey +from lmcache.v1.distributed.internal_api import L2AdapterListener +from tests.v1.multiprocess.test_checkpoint_storage import ( + drain_l2_stores, + large_manifest, + make_manifest, + open_store, + publish_all, + restores, +) + + +class WriteCounter(L2AdapterListener): + def __init__(self) -> None: + self.payload_bytes = 0 + + def on_l2_keys_stored(self, keys: list[ObjectKey], sizes: list[int]) -> None: + self.payload_bytes += sum(sizes) + + def on_l2_keys_accessed(self, keys: list[ObjectKey]) -> None: + pass + + def on_l2_keys_deleted(self, keys: list[ObjectKey]) -> None: + pass + + +@pytest.mark.parametrize("case", ["replacement", "retained_branches", "disk_pressure"]) +def test_payload_write_volume(tmp_path: Path, case: str) -> None: + entries = ( + [large_manifest(900 + i) for i in range(16)] + if case == "disk_pressure" + else [replace(make_manifest(), world_size=1) for _ in range(12)] + ) + entries = [ + replace(entry, prefix=replace(entry.prefix, tail_tokens=(900 + i,))) + for i, entry in enumerate(entries) + ] + if case == "replacement": + entries = [replace(entry, prefix=entries[0].prefix) for entry in entries] + totals = {} + for policy in ("default", "checkpoint_on_evict"): + path = tmp_path / policy + path.mkdir() + counter = WriteCounter() + with open_store( + path, + True, + store_policy=policy, + disk_capacity_gb=(6 / 1024 if case == "disk_pressure" else 0), + ) as resources: + service, index, storage, mapping = resources + storage.l2_adapters()[0][1].register_listener(counter) + for entry in entries: + publish_all(service, index, mapping, entry) + # A completed turn separates write bursts equally for both policies. + drain_l2_stores(storage) + before = counter.payload_bytes + assert restores(service, mapping, entries[-1]) + after = counter.payload_bytes + with open_store(path, True, store_policy=policy) as resources: + service, index, _, mapping = resources + retained = [ + entry for entry in entries if index.find((entry.prefix,)) == entry + ] + restored = sum(restores(service, mapping, entry) for entry in retained) + assert restored == len(retained) + totals[policy] = { + "before_drain": before, + "after_drain": after, + "restored_after_restart": restored, + } + print( + "WRITE_VOLUME " + json.dumps({"case": case, "policies": totals}, sort_keys=True) + ) + if case == "replacement": + assert ( + totals["checkpoint_on_evict"]["after_drain"] + < totals["default"]["after_drain"] + ) + if case == "retained_branches": + assert totals["checkpoint_on_evict"]["before_drain"] == 0 + assert ( + totals["checkpoint_on_evict"]["after_drain"] + == totals["default"]["after_drain"] + ) diff --git a/tests/v1/test_vllm_semantic_checkpoint_transfer.py b/tests/v1/test_vllm_semantic_checkpoint_transfer.py index 44c063222a5..efc92c32f6a 100644 --- a/tests/v1/test_vllm_semantic_checkpoint_transfer.py +++ b/tests/v1/test_vllm_semantic_checkpoint_transfer.py @@ -7,7 +7,6 @@ from multiprocessing import shared_memory from types import SimpleNamespace from typing import Any, cast -from unittest.mock import patch import json import os import time @@ -26,6 +25,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, @@ -40,7 +40,6 @@ from vllm.v1.request import Request # noqa: E402 # First Party -from lmcache.integration.vllm import checkpoint_scheduler # noqa: E402 from lmcache.integration.vllm.checkpoint_copy import CheckpointPageCopier # noqa: E402 from lmcache.integration.vllm.checkpoint_scheduler import ( # noqa: E402 CheckpointEngineTask, @@ -485,6 +484,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 +507,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 +597,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 +679,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 +719,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 @@ -711,8 +755,8 @@ def drain_predecessor() -> None: worker.close() -def test_failed_restore_falls_back_to_the_longest_remaining_checkpoint() -> None: - """A restore whose pages are gone retries the directory, not the prompt.""" +def test_retired_restore_selects_the_longest_remaining_checkpoint() -> None: + """Coordinated last-copy loss removes the unusable candidate before lookup.""" with open_checkpoint_rpc() as (client, module, mapping, _name): manager = make_manager() bridge = CheckpointSchedulerBridge( @@ -803,22 +847,16 @@ def run(task: CheckpointEngineTask) -> dict[int, bool]: consumer = make_request("consumer") attempts: list[tuple[int, bool]] = [] deadline = time.monotonic() + 10 - with patch.object(checkpoint_scheduler.logger, "info") as info: - while not bridge.poll_prefix(consumer): - assert time.monotonic() < deadline - for task in bridge.take_tasks(): - results = run(task) - attempts.append( - (task.manifest.prefix.num_tokens, all(results.values())) - ) - time.sleep(0.001) + while not bridge.poll_prefix(consumer): + assert time.monotonic() < deadline + for task in bridge.take_tasks(): + results = run(task) + attempts.append( + (task.manifest.prefix.num_tokens, all(results.values())) + ) + time.sleep(0.001) - assert attempts == [(11, False), (8, True)] - failure = next( - call.args for call in info.call_args_list if "failed" in call.args[0] - ) - assert failure[1:4] == (11, consumer.request_id, [0, 1, 2, 3]) - assert 0 <= failure[4] < 10 + assert attempts == [(8, True)] assert manager.get_computed_blocks(consumer)[1] == 8 assert bridge.external_tokens(consumer) == 8 assert not bridge.has_pending @@ -846,7 +884,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 +998,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 +1011,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