diff --git a/invokeai/app/services/model_load/model_load_default.py b/invokeai/app/services/model_load/model_load_default.py index 47d385cb564..d5ea060b4e7 100644 --- a/invokeai/app/services/model_load/model_load_default.py +++ b/invokeai/app/services/model_load/model_load_default.py @@ -111,7 +111,10 @@ def load_model_from_path( ram_cache = self.ram_cache cache_key = str(model_path) try: - return LoadedModelWithoutConfig(cache_record=ram_cache.get(key=cache_key), cache=ram_cache) + # Retrieve with a first-use claim so the record cannot be evicted between the lookup + # and the wrapper's first lock (see ModelCache.get_with_first_use_claim). + cache_record, first_use_claim = ram_cache.get_with_first_use_claim(key=cache_key) + return LoadedModelWithoutConfig(cache_record=cache_record, cache=ram_cache, first_use_claim=first_use_claim) except IndexError: pass @@ -160,9 +163,17 @@ def diffusers_load_directory(directory: Path) -> AnyModel: # a worker sharing this cache built it while we waited. with MODEL_LOAD_LOCK.write_lock(): try: - return LoadedModelWithoutConfig(cache_record=ram_cache.get(key=cache_key), cache=ram_cache) + cache_record, first_use_claim = ram_cache.get_with_first_use_claim(key=cache_key) + return LoadedModelWithoutConfig( + cache_record=cache_record, cache=ram_cache, first_use_claim=first_use_claim + ) except IndexError: pass raw_model = loader(model_path) - ram_cache.put(key=cache_key, model=raw_model) - return LoadedModelWithoutConfig(cache_record=ram_cache.get(key=cache_key), cache=ram_cache) + # The admission claim shields the record until this retrieval's own claim takes over + # (see ModelCache.put), so nothing can evict the model still held in raw_model here. + admission_claim = ram_cache.put(key=cache_key, model=raw_model, claim_admission=True) + cache_record, first_use_claim = ram_cache.get_with_first_use_claim(key=cache_key) + if admission_claim is not None: + admission_claim.release() + return LoadedModelWithoutConfig(cache_record=cache_record, cache=ram_cache, first_use_claim=first_use_claim) diff --git a/invokeai/backend/model_manager/load/load_base.py b/invokeai/backend/model_manager/load/load_base.py index 9df213f90a8..c8bac9b370b 100644 --- a/invokeai/backend/model_manager/load/load_base.py +++ b/invokeai/backend/model_manager/load/load_base.py @@ -18,7 +18,11 @@ from invokeai.backend.model_manager.load.model_cache.cached_model.cached_model_with_partial_load import ( CachedModelWithPartialLoad, ) -from invokeai.backend.model_manager.load.model_cache.model_cache import MODEL_LOAD_LOCK, ModelCache +from invokeai.backend.model_manager.load.model_cache.model_cache import ( + MODEL_LOAD_LOCK, + FirstUseClaim, + ModelCache, +) from invokeai.backend.model_manager.taxonomy import AnyModel, SubModelType @@ -53,17 +57,85 @@ class LoadedModelWithoutConfig: do not have a state_dict, in which case this value will be None. """ - def __init__(self, cache_record: CacheRecord, cache: ModelCache): + def __init__(self, cache_record: CacheRecord, cache: ModelCache, first_use_claim: Optional[FirstUseClaim] = None): self._cache_record = cache_record self._cache = cache - release_first_use_grace = getattr(cache, "release_first_use_grace", None) - self._first_use_finalizer = ( - finalize(self, release_first_use_grace, cache_record) - if cache_record.awaiting_first_use and release_first_use_grace is not None + # The record must stay shielded from the eviction sweeps until this wrapper's first lock: + # without that, a sweep racing the window — a peer's budget reconcile, another model's + # make-room, or the cache's shutdown() — evicts the record out from under the wrapper, + # detaching it from the cache's RAM accounting and (for shared weights) from store + # ownership while its tensors live on, so a peer's reload of the key mints a duplicate + # canonical copy the budget counts once. + # + # `first_use_claim` is that shield, and callers should always supply one: it was armed by + # ModelCache.get_with_first_use_claim() under the same lock acquisition as the lookup, so + # the window is covered from its very first instruction. Adopting it is just holding the + # reference — its hold is released on the first lock (_end_first_use_window) or, if this + # wrapper is dropped un-entered, by the claim's own finalizer when it dies with us. + self._first_use_claim = first_use_claim + # Fallback for a record obtained through plain get() (or by a caller that could not get a + # claim): arm the hold here instead. The instructions between that get() returning and + # this constructor are then unshielded — an eviction landing exactly there is the + # pre-existing, tolerated issue-7513 detached path, and register_first_use_hold declines + # to arm on a record that already lost that race. The finalizer below also covers the + # put()-set admission grace for a record whose hold could not be armed (no deferred worker + # running). Both release routes quote the epoch the hold was armed under, so a hold the + # cache's dead-worker recovery already zeroed is never re-released against a successor. + release_grace = getattr(cache, "release_first_use_grace", None) + register_hold = getattr(cache, "register_first_use_hold", None) + self._first_use_hold_epoch: Optional[int] = ( + register_hold(cache_record) + if first_use_claim is None and register_hold is not None and release_grace is not None else None ) + self._first_use_finalizer = None + try: + if ( + first_use_claim is None + and release_grace is not None + and (self._first_use_hold_epoch is not None or cache_record.awaiting_first_use) + ): + self._first_use_finalizer = finalize( + self, + release_grace, + cache_record, + self._first_use_hold_epoch is not None, + self._first_use_hold_epoch if self._first_use_hold_epoch is not None else 0, + ) + self._first_use_finalizer.atexit = False + except BaseException: + # finalize() allocates, and the cache runs at the RAM ceiling by design. A hold armed + # above with no finalizer to carry its release would shield its record from every + # eviction path — shutdown()'s sweep included — for the life of the process. Detach + # first: a failure AFTER the finalizer was registered would otherwise leave it live, + # and its later release would double-decrement, consuming a hold that by then may + # belong to a different holder. + if self._first_use_finalizer is not None: + self._first_use_finalizer.detach() + self._first_use_finalizer = None + if self._first_use_hold_epoch is not None: + release_hold = getattr(cache, "release_first_use_hold", None) + if release_hold is not None: + release_hold(cache_record, self._first_use_hold_epoch) + self._first_use_hold_epoch = None + raise + + def _end_first_use_window(self) -> None: + """This wrapper's first lock ended its get()->lock() window: the record is now pinned by + its lock count, so drop the abandonment finalizer and release the first-use hold. Runs at + most once — later re-entries of the context manager find nothing to release.""" + if self._first_use_claim is not None: + claim, self._first_use_claim = self._first_use_claim, None + claim.release() if self._first_use_finalizer is not None: - self._first_use_finalizer.atexit = False + self._first_use_finalizer.detach() + self._first_use_finalizer = None + if self._first_use_hold_epoch is not None: + hold_epoch = self._first_use_hold_epoch + self._first_use_hold_epoch = None + release_hold = getattr(self._cache, "release_first_use_hold", None) + if release_hold is not None: + release_hold(self._cache_record, hold_epoch) def __enter__(self) -> AnyModel: # Hold the MODEL_LOAD_LOCK read lock across the VRAM load (lock() runs @@ -72,8 +144,7 @@ def __enter__(self) -> AnyModel: # Acquired before the cache's own lock to keep a consistent lock order (see MODEL_LOAD_LOCK). with MODEL_LOAD_LOCK.read_lock(): self._cache.lock(self._cache_record, None) - if self._first_use_finalizer is not None: - self._first_use_finalizer.detach() + self._end_first_use_window() try: self.repair_required_tensors_on_device() return self.model @@ -96,8 +167,7 @@ def model_on_device( # See __enter__ for why the VRAM load is wrapped in the read lock. with MODEL_LOAD_LOCK.read_lock(): self._cache.lock(self._cache_record, working_mem_bytes) - if self._first_use_finalizer is not None: - self._first_use_finalizer.detach() + self._end_first_use_window() try: self.repair_required_tensors_on_device() yield (self._cache_record.cached_model.get_cpu_state_dict(), self._cache_record.cached_model.model) @@ -113,8 +183,7 @@ def model(self) -> AnyModel: def model_in_ram(self) -> Generator[AnyModel, None, None]: """Pin the model's cache record in RAM without moving the model to its execution device.""" self._cache.lock_in_ram(self._cache_record) - if self._first_use_finalizer is not None: - self._first_use_finalizer.detach() + self._end_first_use_window() try: yield self.model finally: @@ -166,8 +235,14 @@ def unload_from_vram(self, vram_bytes_to_free: int, keep_required_weights_in_vra class LoadedModel(LoadedModelWithoutConfig): """Context manager object that mediates transfer from RAM<->VRAM.""" - def __init__(self, config: Optional[AnyModelConfig], cache_record: CacheRecord, cache: ModelCache): - super().__init__(cache_record=cache_record, cache=cache) + def __init__( + self, + config: Optional[AnyModelConfig], + cache_record: CacheRecord, + cache: ModelCache, + first_use_claim: Optional[FirstUseClaim] = None, + ): + super().__init__(cache_record=cache_record, cache=cache, first_use_claim=first_use_claim) self.config = config diff --git a/invokeai/backend/model_manager/load/load_default.py b/invokeai/backend/model_manager/load/load_default.py index 6491fc3f9e5..a8727b476b0 100644 --- a/invokeai/backend/model_manager/load/load_default.py +++ b/invokeai/backend/model_manager/load/load_default.py @@ -18,6 +18,7 @@ from invokeai.backend.model_manager.load.model_cache.cache_record import CacheRecord from invokeai.backend.model_manager.load.model_cache.model_cache import ( MODEL_LOAD_LOCK, + FirstUseClaim, ModelCache, get_model_cache_key, ) @@ -221,8 +222,13 @@ def load_model(self, model_config: AnyModelConfig, submodel_type: Optional[SubMo if not model_path.exists(): raise FileNotFoundError(f"Files for model '{model_config.name}' not found at {model_path}") - cache_record = self._load_and_cache(model_config, submodel_type) - return LoadedModel(config=model_config, cache_record=cache_record, cache=self._ram_cache) + cache_record, first_use_claim = self._load_and_cache(model_config, submodel_type) + return LoadedModel( + config=model_config, + cache_record=cache_record, + cache=self._ram_cache, + first_use_claim=first_use_claim, + ) @property def ram_cache(self) -> ModelCache: @@ -257,11 +263,23 @@ def _get_execution_device( return None - def _load_and_cache(self, config: AnyModelConfig, submodel_type: Optional[SubModelType] = None) -> CacheRecord: + def _load_and_cache( + self, config: AnyModelConfig, submodel_type: Optional[SubModelType] = None + ) -> tuple[CacheRecord, Optional[FirstUseClaim]]: + """Return the model's cache record together with the first-use claim shielding it. + + The claim is armed inside the cache lookup itself (see ModelCache.get_with_first_use_claim) + and belongs to the LoadedModel this record is about to be wrapped in, which releases it at + its first lock. It is carried out of here rather than being armed at the wrapper's + construction so that the record cannot be evicted anywhere along the way — this method's + two returns, the shell registration below, and load_model()'s own frame. If the load + raises after the claim is armed, the claim dies with the frame that holds it and releases + the hold itself. + """ stats_name = ":".join([config.base, config.type, config.name, (submodel_type or "")]) cache_key = get_model_cache_key(config.key, submodel_type) try: - return self._ram_cache.get(key=cache_key, stats_name=stats_name) + return self._ram_cache.get_with_first_use_claim(key=cache_key, stats_name=stats_name) except IndexError: pass @@ -280,7 +298,7 @@ def _load_and_cache(self, config: AnyModelConfig, submodel_type: Optional[SubMod # entry while we waited for the mutex. (Workers on other devices use a different cache, # so they will still miss here and construct their own copy — which is intended.) try: - return self._ram_cache.get(key=cache_key, stats_name=stats_name) + return self._ram_cache.get_with_first_use_claim(key=cache_key, stats_name=stats_name) except IndexError: pass @@ -326,15 +344,22 @@ def _load_and_cache(self, config: AnyModelConfig, submodel_type: Optional[SubMod # Determine execution device from model config, considering submodel type execution_device = self._get_execution_device(config, submodel_type) - self._ram_cache.put( + admission_claim = self._ram_cache.put( cache_key, model=loaded_model, execution_device=execution_device, + claim_admission=True, + ) + # Retrieve immediately, and hold the admission claim across the retrieval: the claim + # shields the new record until this frame's own claim takes over, so nothing — a peer's + # reconcile, another model's make-room, the cache's shutdown() — can evict the model + # this loader is still holding, and a load that dies in between releases the shield by + # dropping the claim rather than leaving a flag standing that nothing can clear. + cache_record, first_use_claim = self._ram_cache.get_with_first_use_claim( + key=cache_key, stats_name=stats_name ) - # Retrieve immediately: the new record carries the cache's post-admission grace until - # it is locked, and keeping put() and get() adjacent means no failure in between can - # leave a graced record whose loader never comes back for it. - cache_record = self._ram_cache.get(key=cache_key, stats_name=stats_name) + if admission_claim is not None: + admission_claim.release() # Register the shell only after put() has created the shared entry (via the wrapper's # acquire); it is dropped automatically when that entry's last reference is released. @@ -343,7 +368,7 @@ def _load_and_cache(self, config: AnyModelConfig, submodel_type: Optional[SubMod if shared_store is not None: shared_store.set_shell(cache_key, shell_to_register) - return cache_record + return cache_record, first_use_claim def get_size_fs( self, config: AnyModelConfig, model_path: Path, submodel_type: Optional[SubModelType] = None diff --git a/invokeai/backend/model_manager/load/model_cache/cache_record.py b/invokeai/backend/model_manager/load/model_cache/cache_record.py index a266cdb79e5..61243997e7d 100644 --- a/invokeai/backend/model_manager/load/model_cache/cache_record.py +++ b/invokeai/backend/model_manager/load/model_cache/cache_record.py @@ -1,4 +1,6 @@ +import weakref from dataclasses import dataclass +from typing import Optional from invokeai.backend.model_manager.load.model_cache.cached_model.cached_model_only_full_load import ( CachedModelOnlyFullLoad, @@ -30,13 +32,64 @@ class CacheRecord: # under the in-flight load, breaking the loader's get() or detaching a live model from the # cache's RAM accounting. The grace deliberately survives get(): get() is synchronized, and # its own lock-release hook may run a pending reconcile before the caller can lock the record - # it was just handed. The flag cannot shield a record forever: the cache's local make_room - # path ignores it (cold loads are serialized under MODEL_LOAD_LOCK, so make_room can never - # see another loader's entry inside the put()->lock() window), and the next admission on the - # same cache clears any flag still standing (see the sweep in ModelCache.put()), so a load - # that errors out — or a LoadedModel dropped without ever locking — cannot dodge budget - # reconciles indefinitely. + # it was just handed. The flag cannot shield a record forever: the synchronous eviction + # paths (make_room, drop_model) ignore it, and the next admission on the same cache clears + # any flag still standing (see the sweep in ModelCache.put()), so a load that errors out + # between put() and the LoadedModel's construction cannot dodge budget reconciles + # indefinitely. That backstop is why the flag is withheld whenever it could outlive its + # releasers: an admission made with no deferred worker running, and one made after + # shutdown() — after which no further put() is guaranteed to run the sweep. For the same + # reason shutdown() retires every flag still standing rather than letting its sweep + # stale-retain the record it shields: a load cancelled before its retrieval leaves no wrapper, + # so nothing but that sweep could ever have released it. From the retrieval on, the window is + # tracked by first_use_holds below, whose release is guaranteed by a finalizer rather than by + # the sweep. awaiting_first_use: bool = False + # Count of live holders of this record that have retrieved it but not yet locked it. Armed by + # ModelCache.register_first_use_hold() — from the cache lookup itself + # (ModelCache.get_with_first_use_claim, so that the window has no unshielded head), or from + # LoadedModelWithoutConfig.__init__ for a record obtained through plain get() — and released + # exactly once per holder: on the holder's first lock, or by the weakref finalizer of the + # FirstUseClaim (or the wrapper) that owns it, if it is dropped without ever locking. Unlike + # awaiting_first_use, these holds are NOT swept by the next admission: a warm get()'s holder + # can legitimately sit un-entered across another model's cold load (a node retrieves several + # models before entering their contexts), and its finalizer guarantees the release the sweep + # exists to backstop. The only recovery sweep is + # ModelCache zeroing the counts once the deferred worker — the thread that carries + # finalizer-initiated releases — is gone (from inside the dying worker itself, and as a + # backstop at the next worker start and at shutdown), since a release dispatched toward a + # dead worker may be dropped and would otherwise shield the record forever. + first_use_holds: int = 0 + # Bumped whenever stranded holds are zeroed (dead-worker recovery). Every hold release + # carries the epoch it was armed under and is ignored across a bump: without this, a + # surviving wrapper's late release — or a release enqueued before the old worker died and + # drained after the restart — would decrement a FRESH hold armed by a different wrapper + # under the healthy replacement worker, silently unshielding that wrapper's window. + first_use_holds_epoch: int = 0 + # Weak reference to the FirstUseClaim handed to the loader that admitted this record + # (ModelCache.put(claim_admission=True)), or None if the admission was not claimed. Resolving + # it answers, synchronously and with no release path of any kind, the one question the + # asynchronous eviction sweeps actually need to ask about a just-admitted record: is the load + # that put it here still running? While it resolves, that loader's own local still holds these + # tensors, so evicting the record would release shared-store ownership (and debit the budget) + # for bytes that are still resident — a peer's reload would then mint a duplicate canonical + # the budget counts once — and would fail the load with an IndexError from a retrieval that + # can no longer find its own model. + # + # This is why the shield the asynchronous sweeps consult is the reference and not the two + # flags above. awaiting_first_use is unowned: any holder's abandonment + # (_release_abandoned_holder) and the next admission's sweep clear it on behalf of whoever + # ran, so one load can strip another's shield. first_use_holds is owned but its release rides + # on the deferred worker, so dead-worker recovery must zero it — deliberately trading a live + # holder's shield for the certainty that no record stays shielded forever. A weak reference + # has neither problem: nothing has to release it, so nothing can release it early or fail to + # release it at all. + # + # Only the asynchronous sweeps consult it, through in_first_use_window below. The synchronous + # paths (make_room, drop_model, unlock's stale eviction) gate on first_use_holds alone and + # have never honoured this window's flags — the trade documented on awaiting_first_use — so + # they are unaffected either way. + admission_claim_ref: Optional["weakref.ref"] = None def lock(self) -> None: """Lock this record.""" @@ -51,3 +104,23 @@ def unlock(self) -> None: def is_locked(self) -> bool: """Return true if record is locked.""" return self._locks > 0 + + @property + def admission_in_flight(self) -> bool: + """True while the load that admitted this record is still between its put() and its + retrieval — see admission_claim_ref.""" + claim_ref = self.admission_claim_ref + return claim_ref is not None and claim_ref() is not None + + @property + def in_first_use_window(self) -> bool: + """True while a load, a live first-use claim, or a live LoadedModel wrapper is between + obtaining this record and locking it — including the load that admitted it, for as long as + it has not come back to retrieve it (admission_in_flight). The asynchronous eviction sweeps (shutdown, budget + reconcile, peer-requested eviction) treat such a record like a locked one: evicting it would detach a record whose + holder is about to lock it, splitting the model from the cache's RAM accounting and — for + shared weights — releasing store ownership while the tensors live on, so a peer's reload + would mint a duplicate canonical copy. The synchronous paths (make_room, drop_model, + unlock's stale eviction) honor only the first_use_holds half — see awaiting_first_use for + why an orphaned grace must stay reachable there.""" + return self.awaiting_first_use or self.first_use_holds > 0 or self.admission_in_flight diff --git a/invokeai/backend/model_manager/load/model_cache/model_cache.py b/invokeai/backend/model_manager/load/model_cache/model_cache.py index 4904ff5d1e3..8a8a42935bf 100644 --- a/invokeai/backend/model_manager/load/model_cache/model_cache.py +++ b/invokeai/backend/model_manager/load/model_cache/model_cache.py @@ -8,7 +8,7 @@ from dataclasses import dataclass from functools import wraps from logging import Logger -from typing import Any, Callable, Dict, Generator, List, Optional, Protocol +from typing import Any, Callable, Dict, Generator, List, NamedTuple, Optional, Protocol import psutil import torch @@ -58,6 +58,100 @@ _DEFERRED_STOP = object() +class _AbandonedHolderRelease(NamedTuple): + """Deferred-work item: a LoadedModel wrapper was dropped without ever locking its record. + + `held_first_use` records whether that wrapper had armed a first-use hold (see + CacheRecord.first_use_holds), so the release decrements only what its own wrapper armed — a + grace-only wrapper (one constructed while no worker was running to arm a hold) must not + consume a hold that belongs to a different, still-live wrapper of the same record. + `hold_epoch` is the CacheRecord.first_use_holds_epoch the hold was armed under; a release + from before a dead-worker zeroing sweep must not decrement a hold armed after it. + + The record is referenced WEAKLY. Only a worker drains this queue, and an item can outlive + every chance to be drained: _dispatch_deferred's liveness gate is unsynchronized, so a + finalizer that read the worker slot just before it was retired still enqueues against a thread + that is alive only because it is unwinding, and on a shut-down cache no replacement worker is + guaranteed to start and drain it. A strong reference would pin that record — and through it + the model's whole CPU state dict — for the life of the process, while the store and the budget + both report the bytes as released. + + Resolving weakly costs nothing. While the release still has work to do, the record is the live + occupant of its key in _cached_models, which holds it strongly; and a record already dropped + from there is one _release_abandoned_holder can do nothing useful for — everything past its + identity check is gated on live occupancy, and the hold decrement above that check is inert on + a detached record, since a detached record is never re-attached (put() mints a fresh + CacheRecord) and every reader of first_use_holds is itself gated on occupancy. + """ + + cache_entry_ref: "weakref.ReferenceType[CacheRecord]" + held_first_use: bool + hold_epoch: int + + +class FirstUseClaim: + """Ownership of a first-use hold armed by `ModelCache.get_with_first_use_claim()`. + + The hold is armed under the SAME lock acquisition as the lookup that produced the record, + which is what makes the shield gapless. Arming it from the LoadedModel constructor instead + (its original home) leaves the caller's whole get()->construct stretch unshielded, and that + stretch is not always a few instructions: the configured loader retrieves its record deep + inside `_load_and_cache` and does the shared-store shell registration and two returns before + `load_model` wraps it. An eviction landing in that gap — a shutdown sweep, a peer's reconcile + — detaches the record its holder is about to lock, releasing shared-store ownership while the + tensors live on, so a peer's reload of the key mints a duplicate canonical copy that the + budget counts once (JPPhoto review, 2026-08-30). + + The claim releases its hold exactly once, whichever comes first: + + - `release()`, called by the LoadedModel wrapper that adopted the claim when it reaches its + first lock (see `LoadedModelWithoutConfig._end_first_use_window`); + - its weakref finalizer, when the claim is dropped without that ever happening — with the + wrapper that adopted it (dropped un-entered), or with the frame that requested it, if the + load raised before any wrapper could be built. The finalizer route goes through + `ModelCache.release_first_use_grace`, so it never takes a lock on the collecting thread + (see that method), and it quotes the epoch the hold was armed under, so a hold that + dead-worker recovery already zeroed is never re-released against a successor hold. + + An unadopted claim is therefore not a leak: dropping it is a complete release. That is what + lets the record be shielded from the moment it is looked up, before any wrapper exists to own + the shield. + """ + + def __init__(self, cache: "ModelCache", cache_entry: CacheRecord, hold_epoch: int) -> None: + self._cache = cache + self._cache_entry = cache_entry + self._hold_epoch = hold_epoch + # Seeded first so release() is well defined on a claim whose construction failed partway. + # It cannot make the construction atomic — weakref.finalize registers itself before it + # returns, so an async exception landing between that call and the store below leaves a + # live finalizer for a claim _claim_first_use is about to hand the hold back for, and that + # finalizer's later release is one this class cannot intercept. A one-bytecode window, + # unreachable without an asynchronous exception, and it costs at worst one holder's shield + # rather than a permanently shielded record. + self._finalizer: Optional[weakref.finalize] = None + # The finalizer's arguments keep the cache and the record alive for as long as the claim + # itself is — exactly the lifetime over which the release still has work to do. + self._finalizer = weakref.finalize(self, cache.release_first_use_grace, cache_entry, True, hold_epoch) + self._finalizer.atexit = False + + def release(self) -> None: + """Release the hold now, synchronously. Idempotent. + + Called from the adopting wrapper's first lock, which runs in an ordinary thread context + (never a finalizer), so the direct — locking — release is safe here. + """ + finalizer, self._finalizer = self._finalizer, None + if finalizer is None: + return + if finalizer.detach() is None: + # Already dead or detached: a racing release() on this same claim has the hold. + # Releasing it again would decrement a hold armed by a different holder and unshield + # their window. + return + self._cache.release_first_use_hold(self._cache_entry, self._hold_epoch, admission_claim=self) + + def _run_deferred_work(cache_ref: "weakref.ReferenceType[ModelCache]", work_queue: "queue.SimpleQueue[object]") -> None: """Drain one ModelCache's deferred-work queue until it is stopped or the cache is collected. @@ -68,35 +162,67 @@ def _run_deferred_work(cache_ref: "weakref.ReferenceType[ModelCache]", work_queu `weakref.finalize` registered alongside this thread (see ModelCache._ensure_deferred_worker) pushes _DEFERRED_STOP when the cache is collected, so a parked worker wakes and exits rather than leaking a thread per abandoned cache. + + The worker outlives shutdown() on purpose: a record retained by the shutdown sweep because a + live LoadedModel wrapper was still inside its get()->lock() window has no future unlock() if + that wrapper is dropped un-entered — the wrapper's finalizer, carried by this worker, is the + only thing left that can evict the record and release its shared weights and budget bytes. + + An exit that is NOT one of those two orderly endings runs ModelCache._recover_from_dead_worker + before the thread unwinds. Every shield this queue exists to release is granted only while a + worker is alive, so a worker that stops without a successor leaves those shields with nothing + to lift them; recovering from inside the dying thread makes that recovery independent of a + later admission, which after shutdown() may never come. """ - while True: - work = work_queue.get() - cache = None - try: - if work is _DEFERRED_STOP: - return + stopped_on_purpose = False + try: + while True: + work = work_queue.get() + cache = None + cache_entry = None + try: + if work is _DEFERRED_STOP: + stopped_on_purpose = True + return + cache = cache_ref() + if cache is None: + # The cache was collected; nothing can ever need doing again. + stopped_on_purpose = True + return + if work is _DEFERRED_RECONCILE: + cache._reconcile_budget_if_pending() + else: + assert isinstance(work, _AbandonedHolderRelease) + cache_entry = work.cache_entry_ref() + if cache_entry is not None: + cache._release_abandoned_holder(cache_entry, work.held_first_use, work.hold_epoch) + except Exception: + if cache is not None: + cache._logger.exception("Error processing deferred model-cache work") + finally: + # Drop both references before blocking on the next get(): locals stay bound for as + # long as this frame lives. `work` may carry a CacheRecord, which transitively holds + # its model's CPU weights — and _release_abandoned_holder can evict that very + # record, removing it from the cache AND subtracting its bytes from the RamBudget, + # so holding it would leave the budget under-reporting a model that is still + # resident. `cache` must go for the same reason this function takes a weakref at + # all, and `cache_entry` is the strong reference the queue deliberately does not + # keep — parking on the next get() while still holding it would defeat the point. + work = None + cache = None + cache_entry = None + finally: + # Reached only for an abnormal end (a BaseException — including one raised asynchronously + # into this thread — or a failure of work_queue.get() itself). The two orderly exits above + # both mean no shield can outlive the worker: _DEFERRED_STOP is pushed by the + # cache-collection finalizer, and a collected cache has no records left to shield. + if not stopped_on_purpose: cache = cache_ref() - if cache is None: - # The cache was collected; nothing can ever need doing again. - return - if cache._shutdown_event.is_set(): - continue - if work is _DEFERRED_RECONCILE: - cache._reconcile_budget_if_pending() - else: - assert isinstance(work, CacheRecord) - cache._release_first_use_grace(work) - except Exception: if cache is not None: - cache._logger.exception("Error processing deferred model-cache work") - finally: - # Drop both references before blocking on the next get(): locals stay bound for as long - # as this frame lives. `work` may be a CacheRecord, which transitively holds its model's - # CPU weights — and _release_first_use_grace's release hook can evict that very record, - # removing it from the cache AND subtracting its bytes from the RamBudget, so holding it - # would leave the budget under-reporting a model that is still resident. `cache` must go - # for the same reason this function takes a weakref at all. - work = None + try: + cache._recover_from_dead_worker(threading.current_thread()) + except Exception: + cache._logger.exception("Error recovering from a dead model-cache deferred-work thread") cache = None @@ -562,27 +688,96 @@ def _on_timeout(self) -> None: @synchronized def shutdown(self) -> None: - """Shutdown the model cache, cancelling any pending timers.""" + """Shutdown the model cache: cancel any pending timers and evict the resident records.""" if self._shutdown_event.is_set(): return self._shutdown_event.set() - self._deferred_work_queue.put(_DEFERRED_STOP) + # The deferred worker is deliberately NOT stopped here. The sweep below can only mark a + # record stale when something still references it, and one of those referents — a live + # LoadedModel wrapper that has not yet locked its record — may simply be dropped instead + # of used. Its finalizer-initiated release, carried by the worker, is then the only event + # left that can evict the record; stopping the worker at shutdown would strand such + # records (and their shared-store references and budget bytes) for the life of the + # process. The worker parks on its queue and is stopped by the cache-collection finalizer + # registered in _ensure_deferred_worker when the ModelCache itself is finally dropped. if self._timeout_timer is not None: self._timeout_timer.cancel() self._timeout_timer = None - # Release the resident records' shared-weights references now, synchronously. A shut-down - # cache serves no more loads, and waiting for collection would leave the store's refcounts - # (and canonical tensors) to the wrappers' finalizers — which only ENQUEUE, and at teardown - # there may be no later store operation to drain the queue. shutdown() runs in a normal - # thread context, so the direct (locking) release is safe here. - for cache_entry in self._cached_models.values(): - cache_entry.cached_model.release_shared_weights() + # If no live worker remains, the shields it was carrying releases for have nothing left to + # lift them and would make the sweep below stale-retain their records forever: the + # wrappers' finalizer releases were (or will be) dropped by the dispatch check, and no + # unlock() is coming for a never-locked holder. Lift them now so the sweep can evict those + # records (see _recover_stranded_shields for the trade this takes). + # + # The condition is "no live worker", not "a dead thread still in the slot": a worker that + # ran its own dying recovery retires the slot, and an empty slot means just as surely that + # nothing is left to carry a release. Keying on the dead thread would skip this lift in + # exactly the case the dying recovery deliberately leaves for it — a grace it declined to + # touch because the cache was still live at the time. + if self._deferred_work_thread is None or not self._deferred_work_thread.is_alive(): + self._recover_stranded_shields() + # Evict the resident records now rather than merely releasing their shared-store + # references. Releasing while retaining the records would make the accounting lie two + # ways: the store stops counting bytes whose tensors the retained wrappers still hold (so + # a post-shutdown load of the same key on a peer cache registers a duplicate canonical + # alongside the still-resident released copy), and a later eviction of such a record — + # put() after shutdown() is reachable, see the note in put() — reads uses_shared_weights + # as already-False and debits the non-shared budget for bytes that were admitted as + # shared. Routing through _delete_cache_entry() keeps store ownership until the record + # itself goes away, so the accounting stays truthful at every point. The release must be + # synchronous regardless: waiting for collection would leave the refcounts to the + # wrappers' finalizers — which only ENQUEUE, and at teardown there may be no later store + # operation to drain the queue. shutdown() runs in a normal thread context, so the direct + # (locking) release inside _delete_cache_entry() is safe here. + # + # Records still in use keep their references: entries locked by an in-flight generation + # (Invoker.stop() stops the model manager before the session processor, whose workers are + # cancelled but not joined) and entries a live LoadedModel wrapper has retrieved but not + # yet locked (first_use_holds) are marked stale instead. The eventual release evicts them + # through this same path: unlock() once the generation lets go, or the wrapper's + # abandonment finalizer (via the deferred worker, see _release_abandoned_holder) if the + # wrapper is dropped without ever locking. Without the window check, shutdown() racing the + # gap between the retrieval and the wrapper's __enter__() would evict the very record its + # holder is about to lock: the holder would proceed on a detached record (the tolerated + # issue-7513 path) whose shared-store ownership was just released, so a peer's reload of + # the same key would mint a duplicate canonical copy while the budget counted only one. A + # record never released keeps its bytes — and its accounting — until process exit, which + # is the truthful description of a model that really is still resident. + + # An admission whose load never came back for it is NOT retained here, and needs no + # special case to avoid it: the loaders claim their admissions (put(claim_admission=True)), + # so a load cancelled or errored between its put() and its retrieval drops that claim by + # dying, and the claim's release — the same abandonment path a dropped LoadedModel takes — + # evicts the record with its shared-store reference and its budget bytes. Without it, the + # put()-set grace on such a record had no releaser left after shutdown (its other two are + # the loader's own get() -> lock() and the sweep at the top of the NEXT put(), which is no + # longer guaranteed to run), and the sweep below stale-RETAINED the record for the life of + # the cache object (JPPhoto review, 2026-08-30). + # + # Retiring the flag here instead — the obvious shortcut — is not safe: it is unowned, so + # a standing grace does not mean nobody is still working on the record. It is equally the + # state of a load still between its put() and its retrieval, and of a live un-entered + # LoadedModel wrapper whose hold a worker death zeroed. Evicting either would be the very + # accounting lie this method exists to avoid — the holder's local still holds the tensors + # whose shared-store ownership the eviction releases, so a peer's reload mints a duplicate + # canonical the budget counts once — on top of failing an in-flight load with an + # IndexError from a retrieval that no longer finds its own model. + for cache_entry in list(self._cached_models.values()): + if cache_entry.is_locked or cache_entry.in_first_use_window: + cache_entry.is_stale = True + else: + self._delete_cache_entry(cache_entry) @synchronized @record_activity def put( - self, key: str, model: AnyModel, execution_device: Optional[torch.device] = None, prefetch: bool = False - ) -> None: + self, + key: str, + model: AnyModel, + execution_device: Optional[torch.device] = None, + prefetch: bool = False, + claim_admission: bool = False, + ) -> Optional[FirstUseClaim]: """Add a model to the cache. Args: @@ -594,12 +789,20 @@ def put( single-file pipeline load) and no loader will retrieve it after this call. It is admitted without the post-admission grace, so budget reconciles may evict it immediately. + claim_admission: Shield the new record with an owned first-use hold, alongside the + swept post-admission grace, and return the claim that owns it (None if no hold + could be armed, or if this call was a no-op because the key was already resident). + Every loader that is going to retrieve the record it just admitted should ask for + one and hold it until its retrieval has armed a claim of its own: an owned shield + cannot outlive its load, so a load that is cancelled or errors before its + retrieval releases the admission by dying, instead of leaving a grace standing + that nothing can clear (see shutdown()). """ if key in self._cached_models: self._logger.debug( f"Attempted to add model {key} ({model.__class__.__name__}), but it already exists in the cache. No action necessary." ) - return + return None # Start (or revive) the worker before mutating cache state. Every deferred task concerns an # admitted record, so this makes later dispatch queue-only while avoiding a thread for an @@ -611,8 +814,12 @@ def put( # its final one, so a flag that survives to the next admission is stale — its loader # either errored out before retrieving the model or dropped the LoadedModel without ever # locking it. Clear such flags so an orphaned record cannot dodge budget reconciles - # indefinitely. (An entry retrieved but not yet locked loses its shield here; if a - # reconcile then evicts it, lock() falls back to the tolerated issue-7513 path and + # indefinitely. Only the put()-set grace is swept: an entry whose LoadedModel wrapper is + # already constructed is tracked by first_use_holds instead, which a concurrent cold load + # must NOT clear — a node may retrieve several models before entering any of their + # contexts — and whose release is guaranteed by the wrapper's finalizer rather than by + # this sweep. (An entry retrieved but with no wrapper hold yet loses its shield here; if + # a reconcile then evicts it, lock() falls back to the tolerated issue-7513 path and # proceeds on the detached record.) for stale_entry in self._cached_models.values(): stale_entry.awaiting_first_use = False @@ -684,17 +891,50 @@ def put( # ordinarily evictable. # # Neither does an admission made while no deferred worker is running to release the grace - # if the loader abandons the model. Both states are reachable — put() after shutdown() - # (Invoker.stop() stops model_manager before session_processor, so an in-flight generation - # can land here), and a worker that could not be started under thread exhaustion — and in - # both the flag would never be cleared, leaving the record permanently invisible to every - # asynchronous eviction path while its bytes stay charged to the shared budget. Without the - # grace the record is merely ordinarily evictable; lock() still clears the flag on the - # normal path, so nothing changes when the worker is healthy. + # if the loader abandons the model (a worker that could not be started under thread + # exhaustion, or one lost to an unexpected error before this put()'s restart attempt + # could succeed): the flag would never be cleared, leaving the record permanently + # invisible to every asynchronous eviction path while its bytes stay charged to the + # shared budget. Without the grace the record is merely ordinarily evictable; lock() + # still clears the flag on the normal path, so nothing changes when the worker is + # healthy. + # + # Nor does an admission after shutdown() (Invoker.stop() stops model_manager before + # session_processor, so an in-flight generation can land here). The grace has exactly two + # releasers: the loader's own forward progress (get() -> lock(), or the wrapper's + # abandonment finalizer once one exists), and the sweep at the top of the NEXT put(). A + # load cancelled between this put() and the LoadedModel's construction — no wrapper, so no + # finalizer — depends entirely on that sweep, and after shutdown no further put() is + # guaranteed to come: the flag would stand for the life of the process, hiding the record + # from every asynchronous eviction path while its bytes stay charged to the shared budget. + # Withholding it costs only the shield: the record is ordinarily evictable, is_stale below + # makes its eventual release evict it, and a loader that does come back gets the same + # first_use_holds shield as any other wrapper. If an eviction wins the race to a record + # whose loader is still between put() and get(), that get() raises rather than falling + # back — the same trade the synchronous paths (make_room, drop_model) have always made + # against this flag, taken here only for admissions into an already-shut-down cache. + # + # A claimed admission gets the grace too, but as a courtesy rather than a guarantee: what + # actually shields that window is the record's weak reference to the claim (see the + # arming at the end of this method, and CacheRecord.admission_claim_ref). Neither flag is + # dependable on its own here — _clear_stranded_first_use_holds zeroes every hold the + # moment the deferred worker dies, and the grace is unowned, so any other holder's + # abandonment or the next admission's sweep clears it on behalf of whoever ran. Shielding + # the window with either alone turns one of those ordinary events into a failed load: the + # record is evicted by the next reconcile while its loader still holds the tensors, and + # the loader's own retrieval then raises. worker_running = self._deferred_work_thread is not None and self._deferred_work_thread.is_alive() + shutting_down = self._shutdown_event.is_set() cache_record = CacheRecord( - key=key, cached_model=wrapped_model, awaiting_first_use=not prefetch and worker_running + key=key, + cached_model=wrapped_model, + awaiting_first_use=not prefetch and worker_running and not shutting_down, ) + # An admission after shutdown() (reachable, see above) missed the shutdown sweep, so + # nothing would ever evict it: mark it stale at birth so its final release — unlock(), or + # the abandonment path — evicts it instead of leaving it resident until process exit. + if shutting_down: + cache_record.is_stale = True self._cached_models[key] = cache_record self._cache_stack.append(key) # Account this model's RAM in the global budget. Shared weights are tracked once by the @@ -724,6 +964,33 @@ def put( # everything unlocked, so the only entry a self-reconcile could ever claim is the # model just admitted — evicting it out from under its own loader.) + # Shield the put() -> retrieval window with an owned hold on top of the grace above when + # the caller asked for one. The claim lives in the loader's frame, so a load cancelled or + # errored before its retrieval releases it by dying — which is what lets shutdown() retain + # a claimed admission for a load that is really still in flight (whose local still holds + # these very tensors) while retiring the unowned graces that nothing can release. The + # loader hands the shield over to its retrieval's own claim and drops this one. + # + # Armed last, once the admission is fully committed. _claim_first_use can raise + # (weakref.finalize allocates, and this cache runs at the RAM ceiling by design), and a + # raise between the record's insertion above and the budget accounting would leave a + # resident, store-owning record that the budget never counted — whose eventual eviction + # then debits bytes it never added, permanently stealing another record's charge (the + # debit is clamped to what this cache has tracked). Everything above this point is + # committed, so the worst a failure here costs is the shield: the record is an ordinary + # graced admission, exactly as if the caller had not asked to claim it. + if not claim_admission or prefetch: + return None + admission_claim = self._claim_first_use(cache_record) + if admission_claim is not None: + # The record's weak reference to the claim — not the hold, and not the grace — is what + # actually shields this window (see CacheRecord.admission_claim_ref): it expires with + # the loader's frame, so no event has to release it, and neither a worker death + # (which zeroes holds) nor another holder's abandonment (which clears the grace) can + # strip it while the load is still running. + cache_record.admission_claim_ref = weakref.ref(admission_claim) + return admission_claim + def _warn_once(self, topic: str, device: torch.device, message: str) -> None: """Log `message` the first time `topic` arises for `device`. @@ -770,8 +1037,10 @@ def cached_model_keys(self) -> set[str]: if self._ram_budget is not None and self._budget_reconcile_pending.is_set() and not self._lock._is_owned(): self._dispatch_deferred(_DEFERRED_RECONCILE) - def release_first_use_grace(self, cache_entry: CacheRecord) -> None: - """Make an abandoned, never-locked record available for budget eviction. + def release_first_use_grace( + self, cache_entry: CacheRecord, held_first_use: bool = False, hold_epoch: int = 0 + ) -> None: + """Make an abandoned, never-locked record available for eviction again. Called from a `weakref.finalize` callback (see LoadedModelWithoutConfig), which runs at an arbitrary decref/garbage-collection point in an arbitrary thread. That thread may already @@ -790,13 +1059,19 @@ def release_first_use_grace(self, cache_entry: CacheRecord) -> None: The work is therefore handed to the cache's background worker, exactly as cached_model_keys() does with its own reconcile. SimpleQueue.put() is reentrant and never waits on the cache, store or budget locks, so it is safe from a finalizer. + + `held_first_use` says whether the dropped wrapper had armed a first-use hold (see + register_first_use_hold); the deferred release decrements only what that wrapper armed, + and only while `hold_epoch` still matches the record's — a hold zeroed by dead-worker + recovery must not be re-released against a successor hold. """ - # Unsynchronized read: the flag is monotonic (put() is the only writer that sets it, and - # only on a brand-new record), so a False reading is always final and there is nothing to - # release. Losing a race here at worst queues work that no-ops under the lock. - if not cache_entry.awaiting_first_use: + # Unsynchronized reads: awaiting_first_use is monotonic (put() is the only writer that + # sets it, and only on a brand-new record) and a caller passing held_first_use owns the + # hold it is releasing, so a nothing-to-release reading is final. Losing a race here at + # worst queues work that no-ops under the lock. + if not held_first_use and not cache_entry.awaiting_first_use: return - self._dispatch_deferred(cache_entry) + self._dispatch_deferred(_AbandonedHolderRelease(weakref.ref(cache_entry), held_first_use, hold_epoch)) def _ensure_deferred_worker(self) -> None: """Start the background worker if it is not currently running. Caller must hold the lock. @@ -808,14 +1083,19 @@ def _ensure_deferred_worker(self) -> None: record would keep shielding an idle cache from eviction, which is exactly the failure this mechanism exists to prevent. - Never revives the worker after shutdown(): that call's `_DEFERRED_STOP` is still queued, so - a new thread would consume it and exit immediately, and a shut-down cache has no deferred - work worth doing. + Revival applies after shutdown() too: the worker is what carries abandonment releases for + the records the shutdown sweep retained, so a post-shutdown admission must restore it the + same as any other. """ - if self._shutdown_event.is_set(): - return if self._deferred_work_thread is not None and self._deferred_work_thread.is_alive(): return + # No worker is running. Recover unconditionally rather than only when a dead thread is + # still in the slot: the previous worker may have died without running its own recovery (a + # fork, or a death this call raced), or its recovery may have failed partway through and + # retired the slot on the way out. The sweep is idempotent and costs nothing on a cache + # that has no shields standing, so re-running it at every worker start is what makes a + # failed recovery retryable. + self._recover_stranded_shields() thread = threading.Thread( target=_run_deferred_work, args=(weakref.ref(self), self._deferred_work_queue), @@ -849,32 +1129,228 @@ def _dispatch_deferred(self, work: object) -> None: Neither caller may block: release_first_use_grace() runs inside a weakref finalizer (see its docstring) and cached_model_keys() has a no-stall contract. `SimpleQueue.put()` satisfies both, but only the worker ever drains the queue, so enqueueing while no worker is running - would grow it without bound — and, for a CacheRecord, pin that model's CPU weights for the - life of the process. Drop the item instead. - - Dropping loses nothing real, because put() only grants the first-use grace when a worker is - running to release it (see put()). So a dropped CacheRecord is one that was never shielded - from eviction in the first place, and a dropped reconcile is re-run by the synchronized - release hook of the next cache operation — put() additionally clears any stale grace flags - itself. What must never happen is a record that is shielded with nothing left to unshield - it; that is what the pairing of these two rules prevents. - - A shutdown() landing between the check and the put() can still strand a single item in the - queue. The cache is being torn down at that point and the cost is bounded by one record, so - that race is tolerated rather than paid for with a lock this method cannot take. + would grow it without bound. Drop the item instead. + + Dropping loses nothing irrecoverable, because the shields this queue releases are only + granted while a worker is running (put()'s grace and register_first_use_hold's holds are + both gated on worker liveness). A dropped release can therefore only belong to a shield + granted under a worker that has since died — and every dead-worker recovery zeros exactly + those holds (from the dying worker itself, and as a backstop at the next worker start and + at shutdown), while put() sweeps stale grace flags itself. A dropped reconcile is re-run by + the synchronized release hook of the next cache operation. What must never happen is a + record that is shielded with nothing left to unshield it; that is what the pairing of these + rules prevents. + + The liveness gate is deliberately unsynchronized (a finalizer must not take the cache + lock), so it is only advisory: a caller that read the slot just before a dying worker + retired it still enqueues against a thread that is alive only because it is unwinding, and + and that item may never be drained — no replacement worker is guaranteed to start on a + shut-down cache. _AbandonedHolderRelease therefore holds its record weakly, so a stranded + item pins nothing. (One that IS drained later, by a replacement worker, is harmless: the + epoch bump from the recovery blocks its hold decrement, and its record is by then either + gone or past the point where clearing an orphaned grace matters.) """ - if self._shutdown_event.is_set(): - return thread = self._deferred_work_thread if thread is None or not thread.is_alive(): return self._deferred_work_queue.put(work) @synchronized - def _release_first_use_grace(self, cache_entry: CacheRecord) -> None: - """Clear an abandoned record's grace, then let the release hook reconcile the budget.""" - if self._cached_models.get(cache_entry.key) is cache_entry and not cache_entry.is_locked: + def register_first_use_hold(self, cache_entry: CacheRecord) -> Optional[int]: + """Shield a record while a holder is between retrieving it and its first lock. Returns the + hold's epoch when armed, None when it could not be. + + Normally reached through get_with_first_use_claim(), which arms the hold in the same lock + acquisition as the lookup; called directly only by a LoadedModel wrapper built from a + record that was retrieved with plain get(). + + The eviction sweeps treat a held record like a locked one (see + CacheRecord.in_first_use_window) — in particular, shutdown() retains it with its + shared-store ownership and budget accounting intact instead of evicting it out from under + the holder. The caller (LoadedModelWithoutConfig) releases the hold exactly once: via + release_first_use_hold() on its first lock, or via its weakref finalizer if it is dropped + without ever locking — both quoting the returned epoch, so a hold that dead-worker + recovery already zeroed is never re-released against a successor (see + CacheRecord.first_use_holds_epoch). Because the finalizer route travels through the + deferred worker, the hold is only granted while a worker is running to carry it — the + same liveness gate as put()'s admission grace — after first attempting to revive a dead + worker, which also clears any holds stranded by the death (see _ensure_deferred_worker). + + Not armed for a record that is no longer the occupant under its key: an eviction already + won the race against this wrapper's construction, the tolerated issue-7513 detached path + is already in effect, and a hold on a detached record shields nothing. + """ + self._ensure_deferred_worker() + if self._deferred_work_thread is None or not self._deferred_work_thread.is_alive(): + return None + if self._cached_models.get(cache_entry.key) is not cache_entry: + return None + cache_entry.first_use_holds += 1 + return cache_entry.first_use_holds_epoch + + @synchronized + def release_first_use_hold( + self, cache_entry: CacheRecord, hold_epoch: int, admission_claim: Optional["FirstUseClaim"] = None + ) -> None: + """Release a register_first_use_hold() hold whose holder reached its first lock. + + `admission_claim`, when given, is the claim doing the releasing: if the record still + points at it as its admission shield (CacheRecord.admission_claim_ref), that shield is + retired here too. The shield is object liveness, not claim validity, so without this a + spent claim would go on shielding the record for as long as the loader's frame — or the + traceback that captured it — kept the object alive, well past the retrieval it was meant + to cover. + """ + if cache_entry.first_use_holds > 0 and cache_entry.first_use_holds_epoch == hold_epoch: + cache_entry.first_use_holds -= 1 + if admission_claim is not None: + claim_ref = cache_entry.admission_claim_ref + if claim_ref is not None and claim_ref() is admission_claim: + cache_entry.admission_claim_ref = None + + def _clear_stranded_first_use_holds(self) -> None: + """Zero every record's holds after the deferred worker is found dead. Caller must hold + the cache lock. + + A hold's finalizer-initiated release may have been dispatched toward the dead thread and + dropped — finalizers fire once, so a dropped release is never retried — and such a hold + would shield its record from every eviction path forever. Bumping the epoch makes the + surviving wrappers' own releases (and any release still sitting in the queue from before + the death) no-ops, so they cannot consume holds armed afresh under a later worker. A + still-live wrapper unshielded here merely falls back to the tolerated issue-7513 detached + path if an eviction actually races its lock, which is recoverable; a permanently shielded + record is not. + """ + dropped: list[tuple[str, int]] = [] + for entry in self._cached_models.values(): + if entry.first_use_holds > 0: + dropped.append((entry.key, entry.first_use_holds)) + entry.first_use_holds = 0 + entry.first_use_holds_epoch += 1 + # Report only once every record is actually unshielded. A logging handler that raises is + # one of the ways the worker dies in the first place (see _ensure_deferred_worker), and + # logging inside the loop would let that same handler abort this sweep partway through, + # leaving the records it had not reached shielded with nothing left to unshield them. + for key, count in dropped: + self._logger.warning( + f"Dropping {count} first-use hold(s) on cache entry {key}: the deferred-work thread died, " + "so their releases may have been lost." + ) + + def _recover_stranded_shields(self) -> None: + """Lift the first-use shields whose release depended on a deferred worker that is gone. + Caller must hold the cache lock. Does not evict — see _evict_stale_unshielded_entries. + + Holds are always lifted: a hold's only releases are its wrapper's first lock and, if that + never comes, the finalizer-initiated release this worker carried. With the worker gone the + latter is dropped by _dispatch_deferred, and finalizers never fire twice, so a hold left + standing shields its record from every eviction path for the life of the process. A + wrapper unshielded here merely falls back to the tolerated issue-7513 detached path if an + eviction really does race its lock, which is recoverable; a permanently shielded record is + not. + + The put()-set admission grace is lifted only once the cache is shut down, because only + then has it actually lost a releaser. On a live cache the grace's backstop is the sweep at + the top of the next put() — not the worker — and that sweep still runs; clearing it here + would instead unshield a load that is only midway between its put() and its get(), whose + record a reconcile could then evict out from under it, turning a worker death into a + failed load. After shutdown that backstop is gone (no further put() is guaranteed) and + put() no longer arms the grace at all. + + The same rule, for the same reason, applies to the admission shield a claimed put() + published (CacheRecord.admission_claim_ref). On a live cache it needs no recovery at all: + nothing has to release it, so a worker death cannot strand it — it expires by itself when + the loader's frame does, and the record becomes ordinarily evictable again. Once the cache + is shut down that is no longer enough: the record is stale, and the eviction its expiry + should trigger travels through the dead worker, so the shield would keep it (and its + shared-store reference, and its budget bytes) resident for the life of the process. + Retiring it lets _evict_stale_unshielded_entries reclaim it — the same trade this method + already makes for holds, and it costs the same thing: if the loader really is still + running, it falls back to the tolerated issue-7513 detached path. + """ + self._clear_stranded_first_use_holds() + if not self._shutdown_event.is_set(): + return + for entry in self._cached_models.values(): + entry.awaiting_first_use = False + entry.admission_claim_ref = None + + def _evict_stale_unshielded_entries(self) -> None: + """Evict records left stale with no lock and no first-use shield. Caller holds the lock. + + Only for a shut-down cache whose deferred worker has died. A record the shutdown sweep + marked stale and retained *because* of a shield now has nothing left to evict it: the + shield is lifted, no unlock() is coming for a holder that never locked, no abandonment + release can be carried, and no further admission is guaranteed to run make_room. Before + shutdown this is deliberately not done — the ordinary eviction paths are all still + running, and evicting here would release a record's shared-store ownership while live + wrappers still hold its tensors, which is the accounting lie shutdown() itself refuses to + make (a peer's reload would mint a duplicate canonical while the budget counted one). + """ + evicted = 0 + for entry in list(self._cached_models.values()): + if entry.is_stale and not entry.is_locked and not entry.in_first_use_window: + self._logger.debug(f"Evicting stale cache entry {entry.key} stranded by the deferred worker's death.") + self._delete_cache_entry(entry) + evicted += 1 + if evicted: + try: + gc.collect() + TorchDevice.empty_cache() + except Exception: + # Deferrable housekeeping: empty_cache() can raise from a sick CUDA context after + # an eviction, and this runs in a dying worker's last frame. The eviction itself + # is already done and must not be undone by a failure to hand memory back. + self._logger.exception("Error releasing device memory after a stranded-record eviction") + + @synchronized + def _recover_from_dead_worker(self, worker: threading.Thread) -> None: + """Run stranded-shield recovery from inside the deferred worker as it dies abnormally. + + The other two recovery sites both depend on something else happening first: the next + admission (_ensure_deferred_worker) or shutdown(). Neither is guaranteed — a cache that is + already shut down takes no more admissions, and shutdown()'s own liveness check passes if + the worker is still running at that moment and dies a moment later. That leaves the + records the shutdown sweep retained for a live holder shielded by holds nothing can + release. Recovering here makes the recovery a property of the death itself. + """ + if self._deferred_work_thread is not worker: + # A replacement worker has already taken this slot, so the shields standing now were + # granted under it and are its to release. _ensure_deferred_worker cleared whatever + # this thread stranded when it started that replacement. + return + # Retire the slot first. This thread is still `is_alive()` while it unwinds its own frame, + # so a concurrent admission would otherwise read the worker as healthy and arm a shield + # that this recovery is about to zero. With the slot empty, put() withholds the grace and + # register_first_use_hold starts a replacement worker instead. (An unsynchronized + # _dispatch_deferred that read the slot just before this can still enqueue against the + # dying thread; that item is simply never drained, and _AbandonedHolderRelease holds its + # record weakly so it pins nothing.) + self._deferred_work_thread = None + self._recover_stranded_shields() + if self._shutdown_event.is_set(): + self._evict_stale_unshielded_entries() + + @synchronized + def _release_abandoned_holder(self, cache_entry: CacheRecord, held_first_use: bool, hold_epoch: int) -> None: + """Deferred-worker handler for a LoadedModel wrapper dropped without ever locking. + + Releases whatever shield the wrapper held, then — if the abandoned record is stale + (shutdown() or drop_model() marked it while the wrapper kept it retained) and nothing else + holds it — evicts it here, because no unlock() is ever coming to run the usual + stale-eviction path. The synchronized release hook then reconciles the budget as usual. + """ + if held_first_use and cache_entry.first_use_holds > 0 and cache_entry.first_use_holds_epoch == hold_epoch: + cache_entry.first_use_holds -= 1 + if self._cached_models.get(cache_entry.key) is not cache_entry: + return + if not cache_entry.is_locked: cache_entry.awaiting_first_use = False + if cache_entry.is_stale and not cache_entry.is_locked and not cache_entry.in_first_use_window: + self._delete_cache_entry(cache_entry) + gc.collect() + TorchDevice.empty_cache() + self._logger.debug(f"Evicted stale cache entry {cache_entry.key} after its holder was abandoned.") @synchronized def _get_cache_snapshot(self) -> dict[str, CacheEntrySnapshot]: @@ -937,6 +1413,50 @@ def get(self, key: str, stats_name: Optional[str] = None) -> CacheRecord: return cache_entry + @synchronized + def get_with_first_use_claim( + self, key: str, stats_name: Optional[str] = None + ) -> tuple[CacheRecord, Optional[FirstUseClaim]]: + """get(), with the retrieved record's first-use hold armed atomically with the lookup. + + Every caller that is going to wrap the record in a LoadedModel should retrieve it through + this method rather than get(): the returned claim shields the record from the asynchronous + eviction sweeps for the whole stretch between the lookup and the wrapper's first lock, + with no unshielded gap at the front of it (see FirstUseClaim). Both calls below run under + this frame's lock acquisition, and the synchronized decorator's reconcile hook fires only + on the outermost frame, so no eviction can run between the lookup and the arming. + + The claim must be handed to the wrapper + (`LoadedModelWithoutConfig(..., first_use_claim=claim)`) or simply dropped; either way its + hold is released exactly once and the caller has nothing to clean up. + + Returns `(record, None)` when no hold could be armed — no deferred worker is running to + carry the claim's finalizer release, or an eviction already detached the record. The + caller is then exactly where it was before this method existed: the record may be evicted + under it, and lock() falls back to the tolerated issue-7513 detached path. + + Raises IndexError if the model is not in the cache, exactly as get() does. + """ + cache_entry = self.get(key, stats_name) + return cache_entry, self._claim_first_use(cache_entry) + + def _claim_first_use(self, cache_entry: CacheRecord) -> Optional[FirstUseClaim]: + """Arm a first-use hold and wrap it in the claim that owns its release. Caller holds the + lock. Returns None when no hold could be armed.""" + hold_epoch = self.register_first_use_hold(cache_entry) + if hold_epoch is None: + return None + try: + return FirstUseClaim(self, cache_entry, hold_epoch) + except BaseException: + # The hold is armed, but the object that was to carry its release does not exist — + # weakref.finalize() allocates, and this cache runs at the RAM ceiling by design. An + # orphaned hold is the one failure this whole mechanism must not produce: it shields + # its record from every eviction path, shutdown()'s sweep included, for the life of + # the process. Hand the hold back before letting the failure out. + self.release_first_use_hold(cache_entry, hold_epoch) + raise + @synchronized @record_activity def lock_in_ram(self, cache_entry: CacheRecord) -> None: @@ -1030,7 +1550,22 @@ def unlock(self, cache_entry: CacheRecord) -> None: # If `drop_model()` marked this entry stale (e.g. settings changed while a generation # was using it), evict now so the next load rebuilds with the new settings rather than # silently reusing the pre-change cached module. - if cache_entry.is_stale and not cache_entry.is_locked and cache_entry.key in self._cached_models: + # Identity check, not key membership: if this record was already detached (error-path + # delete) and the key re-admitted, the occupant is a different, non-stale record that must + # not be evicted — and no cleared-callback should fire for a no-op. + # A first-use hold defers the eviction the same way a lock does: another wrapper may + # already hold this record for its own upcoming lock, and evicting here would detach it + # mid-window. Whatever releases the hold — that holder's own unlock() after use, or the + # abandonment path (_release_abandoned_holder) — performs the stale eviction instead. + # (Only the hold half of the first-use window can be live here: awaiting_first_use is + # cleared by every lock entry point, and a brand-new record has no lockers before that, + # so no unlock() can observe it set.) + if ( + cache_entry.is_stale + and not cache_entry.is_locked + and cache_entry.first_use_holds == 0 + and self._cached_models.get(cache_entry.key) is cache_entry + ): bytes_freed = cache_entry.cached_model.total_bytes() self._delete_cache_entry(cache_entry) if self.stats: @@ -1546,7 +2081,13 @@ def _make_room_internal(self, bytes_needed: int) -> None: model_key = self._cache_stack[pos] cache_entry = self._cached_models[model_key] - if not cache_entry.is_locked: + # A first-use hold shields here too: a wrapper obtained warm can sit un-entered while + # another model's cold load makes room, and evicting its record would detach it from + # the holder about to lock it (see CacheRecord.first_use_holds). The put()-set grace + # deliberately does NOT shield from this path (matching its long-standing semantics): + # its releaser is the loader's own forward progress, not a finalizer, so an orphaned + # grace must stay reachable by the synchronous eviction paths. + if not cache_entry.is_locked and cache_entry.first_use_holds == 0: ram_bytes_freed += cache_entry.cached_model.total_bytes() self._logger.debug( f"Dropping {model_key} from RAM cache to free {(cache_entry.cached_model.total_bytes() / MB):.2f}MB." @@ -1632,7 +2173,7 @@ def evict_unlocked_for_peer(self, is_satisfied: Callable[[], bool]) -> Optional[ pos = 0 while pos < len(self._cache_stack) and not is_satisfied(): cache_entry = self._cached_models[self._cache_stack[pos]] - if cache_entry.is_locked or cache_entry.awaiting_first_use: + if cache_entry.is_locked or cache_entry.in_first_use_window: pos += 1 continue self._logger.debug( @@ -1696,16 +2237,21 @@ def _reconcile_budget_if_pending(self, blocking: bool = True) -> None: if self._ram_budget.available() >= 0: return self._budget_reconcile_pending.set() + models_cleared = 0 + # Acquire immediately before the try: a BaseException delivered between the two leaks + # the RLock, and a cache lock owned by a thread that is unwinding blocks every other + # thread for the life of the process. This narrows that window to the interpreter's + # own check between acquire() returning and SETUP_FINALLY rather than closing it — + # inherent to acquiring a lock without `with`, which the blocking flag rules out here. if not self._lock.acquire(blocking=blocking): # Contended (non-blocking caller only): the holder releases through a reconcile # hook, which re-runs this reconcile with the flag still set. return - models_cleared = 0 try: pos = 0 while pos < len(self._cache_stack) and self._ram_budget.available() < 0: cache_entry = self._cached_models[self._cache_stack[pos]] - if cache_entry.is_locked or cache_entry.awaiting_first_use: + if cache_entry.is_locked or cache_entry.in_first_use_window: pos += 1 continue self._logger.debug( @@ -1727,31 +2273,39 @@ def _reconcile_budget_if_pending(self, blocking: bool = True) -> None: # Satisfied: loop back to retire the flag via the guarded clear above. def _delete_cache_entry(self, cache_entry: CacheRecord) -> None: - """Delete cache_entry from the cache if it exists. No exception is thrown if it doesn't exist.""" - was_present = cache_entry.key in self._cached_models + """Delete cache_entry from the cache if it is the record currently held under its key. + No exception is thrown if it is absent (or the key is now held by a different record).""" + # Identity, not key membership: a record can be deleted while still locked (the VRAM-move + # error paths) and the key re-admitted before the record's last unlock() runs the + # stale-eviction path. A key-only check would pop the NEW record from the cache — detaching + # it from all accounting — and, the old record's shared release having already happened, + # read uses_shared_weights as False and debit the non-shared budget for bytes that were + # admitted as shared. The identity guard makes a delete of a detached record a full no-op, + # which also keeps the release exactly-once for double-deletes (release_shared_weights is + # itself idempotent, but the budget debit is not). + if self._cached_models.get(cache_entry.key) is not cache_entry: + return self._cache_stack = [key for key in self._cache_stack if key != cache_entry.key] - self._cached_models.pop(cache_entry.key, None) + del self._cached_models[cache_entry.key] # Drop this device's reference to the shared canonical CPU weights so they can be freed once - # the last device releases them. Guard on was_present so a double-delete doesn't - # double-release (release_shared_weights is itself idempotent, but a re-added entry under the - # same key must not be released by a stale delete). - if was_present: - uses_shared = cache_entry.cached_model.uses_shared_weights - total_bytes = cache_entry.cached_model.total_bytes() - cache_entry.cached_model.release_shared_weights() - # Drop the matching non-shared contribution from the global budget (shared weights are - # released via the store above). Captured before release_shared_weights() flips the flag. - if self._ram_budget is not None and not uses_shared: - self._ram_budget.remove_non_shared(total_bytes, cache=self) + # the last device releases them. + uses_shared = cache_entry.cached_model.uses_shared_weights + total_bytes = cache_entry.cached_model.total_bytes() + cache_entry.cached_model.release_shared_weights() + # Drop the matching non-shared contribution from the global budget (shared weights are + # released via the store above). Captured before release_shared_weights() flips the flag. + if self._ram_budget is not None and not uses_shared: + self._ram_budget.remove_non_shared(total_bytes, cache=self) @synchronized def drop_model(self, model_key: str) -> int: """Drop all cache entries belonging to a model so the next load rebuilds them. Cache keys are `` or `:` (see `get_model_cache_key`), - so a single model may have multiple entries. Locked entries are marked `is_stale` and - evicted by `unlock()` as soon as the last lock releases — without that, a setting - toggled during an in-flight generation would survive on the locked entry and quietly + so a single model may have multiple entries. Locked entries — and entries inside a + first-use window (a LoadedModel wrapper obtained but not yet locked) — are marked + `is_stale` and evicted as soon as the last lock (or the window) releases — without that, + a setting toggled during an in-flight generation would survive on the locked entry and quietly get reused by the next generation. Returns the number of entries immediately dropped (locked entries that are only marked @@ -1765,7 +2319,12 @@ def drop_model(self, model_key: str) -> int: dropped: list[CacheRecord] = [] bytes_freed = 0 for entry in matching: - if entry.is_locked: + # A record with a first-use hold is deferred exactly like a locked one: a live + # LoadedModel wrapper is about to lock it, and evicting now would detach the record + # mid-window. The stale mark makes the hold's release — unlock() after use, or the + # abandonment path — perform the eviction. (An orphaned put()-grace, by contrast, + # stays evictable here, as it always has been on the synchronous paths.) + if entry.is_locked or entry.first_use_holds > 0: entry.is_stale = True continue bytes_freed += entry.cached_model.total_bytes() diff --git a/tests/backend/model_manager/load/model_cache/test_model_cache_ram_budget.py b/tests/backend/model_manager/load/model_cache/test_model_cache_ram_budget.py index c734e96def1..273f18ab43c 100644 --- a/tests/backend/model_manager/load/model_cache/test_model_cache_ram_budget.py +++ b/tests/backend/model_manager/load/model_cache/test_model_cache_ram_budget.py @@ -118,6 +118,1283 @@ def test_shutdown_releases_shared_weights_synchronously(mock_logger): assert store._deferred_releases.qsize() == 0 +def test_shutdown_evicts_unlocked_records(mock_logger): + """shutdown() must route resident records through eviction, not merely release their + shared-store references: a released-but-retained record keeps its tensors alive while the + store (and budget) report zero — accounting that no longer describes reality.""" + store = SharedCpuWeightsStore() + budget = RamBudget(max_bytes=10**12, shared_store=store) + cache = _make_cache(store, budget, mock_logger) + cache.put("m", DummyModule()) + record = _use_and_release(cache, "m") + wrapper_ref = weakref.ref(record.cached_model) + del record + + cache.shutdown() + assert "m" not in cache._cached_models + assert store.refcount("m") == 0 + assert budget.total_in_use() == 0 + # The record is gone, so the tensors really are released: the zero accounting is true. + assert _collect_until(lambda: wrapper_ref() is None) + + +def test_shutdown_retains_locked_records_with_their_accounting(mock_logger): + """A record locked by an in-flight generation at shutdown() keeps its shared-store reference: + its tensors really are resident, so the store and budget must keep saying so. unlock() then + evicts it through the ordinary stale path, releasing exactly once.""" + store = SharedCpuWeightsStore() + budget = RamBudget(max_bytes=10**12, shared_store=store) + cache = _make_cache(store, budget, mock_logger) + cache.put("m", DummyModule()) + record = cache.get("m") + cache.lock(record, None) + + cache.shutdown() + # Still locked: ownership and accounting are retained. + assert "m" in cache._cached_models + assert store.refcount("m") == 1 + assert budget.total_in_use() == S + + cache.unlock(record) + # The last unlock evicts the stale-marked record and returns the accounting to zero. + assert "m" not in cache._cached_models + assert store.refcount("m") == 0 + assert budget.total_in_use() == 0 + + +def test_shutdown_retains_admission_window_records(mock_logger): + """A record inside the put()->lock() admission window (awaiting_first_use) at shutdown() is + retained like a locked one: its loader is about to lock it, and evicting it would release + shared ownership while the loader still holds the tensors. The post-use unlock evicts it.""" + store = SharedCpuWeightsStore() + budget = RamBudget(max_bytes=10**12, shared_store=store) + cache = _make_cache(store, budget, mock_logger) + cache.put("m", DummyModule()) # not yet locked: awaiting_first_use is set + assert cache._cached_models["m"].awaiting_first_use + + cache.shutdown() + assert "m" in cache._cached_models + assert store.refcount("m") == 1 + assert budget.total_in_use() == S + + _use_and_release(cache, "m") + assert "m" not in cache._cached_models + assert store.refcount("m") == 0 + assert budget.total_in_use() == 0 + + +def test_stale_eviction_ignores_a_readmitted_record_under_the_same_key(mock_logger): + """A stale-marked record can be detached while still locked (the VRAM-move error paths call + _delete_cache_entry on a locked record) and the key re-admitted before its last unlock(). + The stale eviction must match the record by IDENTITY: a key-only match would pop the new + record — detaching it from all accounting — and debit the budget for the old record's bytes.""" + store = SharedCpuWeightsStore() + budget = RamBudget(max_bytes=10**12, shared_store=store) + cache = _make_cache(store, budget, mock_logger) + cache.put("m", DummyModule()) + record_1 = cache.get("m") + cache.lock(record_1, None) + + cache.shutdown() # marks the locked record stale + # Simulate the error-path delete of the locked record (see _move_model_to_vram/_ram), then a + # post-shutdown re-admission of the same key (reachable: see the put()-after-shutdown() note). + cache._delete_cache_entry(record_1) + cache.put("m", DummyModule()) + record_2 = cache._cached_models["m"] + assert record_2 is not record_1 + in_use_after_readmission = budget.total_in_use() + assert in_use_after_readmission == S + + # The detached record's last unlock must not evict the re-admitted record or touch the budget. + cache.unlock(record_1) + assert cache._cached_models.get("m") is record_2 + assert store.refcount("m") == 1 + assert budget.total_in_use() == in_use_after_readmission + + +def test_no_duplicate_canonical_when_peer_reloads_after_shutdown(mock_logger): + """The canonical entry must survive while a locked holder retains it, so a peer cache + reloading the key after this cache's shutdown() adopts the SAME canonical tensors instead of + registering a second copy alongside the still-resident one.""" + store = SharedCpuWeightsStore() + budget = RamBudget(max_bytes=10**12, shared_store=store) + cache_a = _make_cache(store, budget, mock_logger) + cache_b = _make_cache(store, budget, mock_logger) + try: + cache_a.put("m", DummyModule()) + record_a = cache_a.get("m") + cache_a.lock(record_a, None) + canonical_before = store.peek("m") + + cache_a.shutdown() + cache_b.put("m", DummyModule()) + _use_and_release(cache_b, "m") + # cache_b adopted the existing canonical: one copy in RAM, referenced by both holders. + assert store.peek("m") is canonical_before + assert store.refcount("m") == 2 + assert budget.total_in_use() == S + + cache_a.unlock(record_a) + assert store.refcount("m") == 1 + assert budget.total_in_use() == S + finally: + cache_b.shutdown() + + +def test_shutdown_retains_record_inside_get_to_lock_window(mock_logger): + """shutdown() racing the gap between get() and the LoadedModel's first lock must retain the + record (JPPhoto review, 2026-08-13): a warm record is past its admission grace, so without + the wrapper's first-use hold the sweep would evict it, releasing shared-store ownership while + the holder proceeds to lock the detached record — and a peer's reload of the same key would + then mint a duplicate canonical copy that the budget counts only once.""" + store = SharedCpuWeightsStore() + budget = RamBudget(max_bytes=10**12, shared_store=store) + cache_a = _make_cache(store, budget, mock_logger) + cache_b = _make_cache(store, budget, mock_logger) + try: + cache_a.put("m", DummyModule()) + _use_and_release(cache_a, "m") # warm: past the admission grace, unlocked + + # A generation retrieves the model; shutdown() lands before it enters the context. + loaded_model = LoadedModelWithoutConfig(cache_record=cache_a.get("m"), cache=cache_a) + canonical_before = store.peek("m") + cache_a.shutdown() + + record = cache_a._cached_models.get("m") + assert record is loaded_model._cache_record, "shutdown() evicted the record mid-window" + assert record.is_stale + assert store.refcount("m") == 1 + assert budget.total_in_use() == S + + # A peer reloading the key while the holder is still using it adopts the SAME canonical. + with loaded_model as _model: + assert cache_a._cached_models.get("m") is loaded_model._cache_record, "locked a detached record" + cache_b.put("m", DummyModule()) + _use_and_release(cache_b, "m") + assert store.peek("m") is canonical_before, "peer reload minted a duplicate canonical" + assert store.refcount("m") == 2 + assert budget.total_in_use() == S + + # Exiting the context is the record's last release: the stale mark set at shutdown() + # evicts it with its accounting. + assert "m" not in cache_a._cached_models + assert store.refcount("m") == 1 + assert budget.total_in_use() == S + finally: + cache_b.shutdown() + + +def test_shutdown_retains_a_claimed_admission_for_its_in_flight_load(mock_logger): + """A loader between its put() and its retrieval still holds the canonical tensors it just + admitted, so shutdown() must retain that record. Evicting it would release shared-store + ownership (and debit the budget) for bytes that are still resident — a peer's reload would + then mint a duplicate canonical the budget counts once — and would fail the in-flight load + with an IndexError from a retrieval that no longer finds its own model. Invoker.stop() stops + the model manager BEFORE the session processor, whose workers are cancelled but not joined, + so that load is genuinely still running.""" + store = SharedCpuWeightsStore() + budget = RamBudget(max_bytes=10**12, shared_store=store) + cache_a = _make_cache(store, budget, mock_logger) + cache_b = _make_cache(store, budget, mock_logger) + try: + # The loader admits its model and still holds it; shutdown() lands before the retrieval. + admission_claim = cache_a.put("m", DummyModule(), claim_admission=True) + assert admission_claim is not None, "the admission was not claimable" + record = cache_a._cached_models["m"] + canonical_before = store.peek("m") + cache_a.shutdown() + + assert cache_a._cached_models.get("m") is record, "shutdown() evicted a record mid-admission" + assert record.is_stale + assert store.refcount("m") == 1 + assert budget.total_in_use() == S + + # A peer loading the key now adopts the same canonical rather than duplicating it. + cache_b.put("m", DummyModule()) + _use_and_release(cache_b, "m") + assert store.peek("m") is canonical_before, "peer reload minted a duplicate canonical" + assert store.refcount("m") == 2 + + # The loader's retrieval still finds its model, and hands the shield to its own claim. + retrieved, first_use_claim = cache_a.get_with_first_use_claim("m") + assert retrieved is record + admission_claim.release() + + loaded_model = LoadedModelWithoutConfig(cache_record=record, cache=cache_a, first_use_claim=first_use_claim) + with loaded_model as _model: + assert cache_a._cached_models.get("m") is record, "locked a detached record" + assert "m" not in cache_a._cached_models + assert store.refcount("m") == 1 + assert budget.total_in_use() == S + finally: + cache_b.shutdown() + + +def test_a_cancelled_claimed_admission_is_evicted_at_shutdown(mock_logger): + """JPPhoto's report in its own order (2026-08-30): a load that admits a model and is then + cancelled before its retrieval leaves no wrapper, so no finalizer and no unlock() is ever + coming, and after shutdown no further put() is guaranteed to run the sweep that was the + put()-set grace's last backstop — the record, its shared-store reference and its budget charge + used to stand until the cache object was collected. The admission claim is what ends it: + dropping it releases the admission, so the shutdown sweep finds an ordinary idle record.""" + store = SharedCpuWeightsStore() + budget = RamBudget(max_bytes=10**12, shared_store=store) + cache = _make_cache(store, budget, mock_logger) + admission_claim = cache.put("m", DummyModule(), claim_admission=True) + assert admission_claim is not None + record = cache._cached_models["m"] + + del admission_claim # the load is cancelled between its put() and its retrieval + gc.collect() + assert _wait_until(lambda: not record.in_first_use_window), "the cancelled admission stayed shielded" + + cache.shutdown() + assert "m" not in cache._cached_models, "shutdown() retained an admission nothing can release" + assert store.refcount("m") == 0 + assert budget.total_in_use() == 0 + + +def test_a_dropped_admission_claim_lets_shutdown_evict_its_record(mock_logger): + """The other half of the claimed admission: a load cancelled before its retrieval drops the + claim, and that release is what evicts the record the shutdown sweep retained for it — with + its shared-store reference and its budget bytes.""" + store = SharedCpuWeightsStore() + budget = RamBudget(max_bytes=10**12, shared_store=store) + cache = _make_cache(store, budget, mock_logger) + admission_claim = cache.put("m", DummyModule(), claim_admission=True) + assert admission_claim is not None + record = cache._cached_models["m"] + assert record.first_use_holds == 1 + assert record.awaiting_first_use, "a claimed admission keeps the grace as its dead-worker fallback" + + cache.shutdown() + assert cache._cached_models.get("m") is record + + del admission_claim # the load was cancelled before it could retrieve the model + gc.collect() + assert _wait_until(lambda: "m" not in cache._cached_models), "a cancelled admission stayed resident" + assert store.refcount("m") == 0 + assert budget.total_in_use() == 0 + + +def test_a_dead_worker_falls_back_to_the_admission_grace(mock_logger): + """A claimed admission must keep the swept grace beside its hold. Dead-worker recovery zeroes + every hold — a hold whose finalizer-carried release was dropped would shield its record + forever — while it deliberately leaves a live cache's grace alone, whose backstop is the next + put()'s sweep rather than the worker. Shielding the admission window with the hold alone would + therefore turn a worker death into a failed load: the next reconcile evicts the record while + the loader still holds its tensors, and the loader's own retrieval raises.""" + store = SharedCpuWeightsStore() + budget = RamBudget(max_bytes=10**12, shared_store=store) + cache = _make_cache(store, budget, mock_logger) + try: + admission_claim = cache.put("m", DummyModule(), claim_admission=True) + assert admission_claim is not None + record = cache._cached_models["m"] + assert record.first_use_holds == 1 and record.awaiting_first_use + + _kill_worker_abnormally(cache, record) + assert record.first_use_holds == 0, "premise broken: the recovery did not zero the hold" + assert record.awaiting_first_use, "the live cache's admission grace was lifted with the hold" + + # An asynchronous eviction sweep must still skip the record, and the loader's retrieval + # must still find the model it admitted. + cache.evict_unlocked_for_peer(lambda: False) + assert "m" in cache._cached_models, "the admission was evicted out from under its loader" + assert store.refcount("m") == 1 + retrieved, _claim = cache.get_with_first_use_claim("m") + assert retrieved is record + finally: + cache.shutdown() + + +def test_a_failed_admission_claim_leaves_the_admission_fully_accounted(mock_logger, monkeypatch): + """The claim is armed only once put() has committed. A raise while arming must not leave a + resident, store-owning record that the budget never counted: the record's eventual eviction + would then debit bytes that were never added, and the debit is clamped to what this cache has + tracked — so it lands on a different, still-resident model's charge, permanently.""" + store = SharedCpuWeightsStore() + budget = RamBudget(max_bytes=10**12, shared_store=store) + # keep_ram_copy=False -> not deduplicated, so each model is charged to the non-shared total. + cache = _make_cache(store, budget, mock_logger, keep_ram_copy=False) + + def boom(*args, **kwargs): + raise MemoryError("no room for the claim") + + try: + cache.put("resident", DummyModule()) + one_model = budget.total_in_use() + assert one_model > 0 + + monkeypatch.setattr(model_cache_module, "FirstUseClaim", boom) + with pytest.raises(MemoryError): + cache.put("claimed", DummyModule(), claim_admission=True) + monkeypatch.undo() + + assert "claimed" in cache._cached_models, "premise: the admission itself completed" + assert budget.total_in_use() == 2 * one_model, "a resident admission went uncounted" + + # Evicting it debits only its own bytes; the other model's charge survives. + cache._delete_cache_entry(cache._cached_models["claimed"]) + assert budget.total_in_use() == one_model + finally: + monkeypatch.undo() + cache.shutdown() + + +def test_shutdown_retains_a_mid_admission_record_after_a_worker_death_and_restart(mock_logger): + """Neither flag survives an ordinary sequence of events, so neither can be what shutdown() + keys its retention on. A worker death zeroes the admission's hold; a replacement worker + (started by any later retrieval or admission) then hides the death from shutdown()'s + dead-worker recovery, and that later admission's sweep clears the unowned grace. With both + gone, shutdown() must still retain the record on the strength of the live claim: its loader + has not come back yet and still holds these tensors.""" + store = SharedCpuWeightsStore() + budget = RamBudget(max_bytes=10**12, shared_store=store) + cache = _make_cache(store, budget, mock_logger) + try: + admission_claim = cache.put("m", DummyModule(), claim_admission=True) + assert admission_claim is not None + record = cache._cached_models["m"] + assert record.first_use_holds == 1 and record.awaiting_first_use + + _kill_worker_abnormally(cache, record) + assert record.first_use_holds == 0, "premise broken: the recovery did not zero the hold" + + # An unrelated admission restarts the worker — so shutdown() takes its live-worker path, + # skipping the dead-worker recovery — and its sweep clears the unowned grace. + cache.put("other", DummyModule()) + assert cache._deferred_work_thread is not None and cache._deferred_work_thread.is_alive() + assert not record.awaiting_first_use, "premise broken: the sweep left the grace standing" + + cache.shutdown() + assert cache._cached_models.get("m") is record, "shutdown() evicted a record mid-admission" + assert record.is_stale + assert store.refcount("m") == 1 + + # The loader comes back and finds its own model, as it would have before any of this. + retrieved, _claim = cache.get_with_first_use_claim("m") + assert retrieved is record + finally: + cache.shutdown() + + +def test_a_live_admission_survives_both_unowned_shields_being_stripped(mock_logger): + """The same property against the asynchronous sweeps, isolated: with the hold zeroed and the + grace cleared — each by an event outside the loader's control — the record is still shielded + while its admission claim lives, and evicting it is only correct once that claim is gone.""" + store = SharedCpuWeightsStore() + budget = RamBudget(max_bytes=10**12, shared_store=store) + cache = _make_cache(store, budget, mock_logger) + admission_claim = cache.put("m", DummyModule(), claim_admission=True) + assert admission_claim is not None + record = cache._cached_models["m"] + + record.first_use_holds = 0 # as dead-worker recovery would leave it + record.awaiting_first_use = False # as another holder's abandonment, or the next put(), would + assert record.in_first_use_window, "the live admission claim stopped shielding its record" + + cache.evict_unlocked_for_peer(lambda: False) + assert "m" in cache._cached_models, "a peer eviction took the record out from under its loader" + cache.shutdown() + assert cache._cached_models.get("m") is record, "shutdown() evicted a record mid-admission" + + del admission_claim # the load ends; only now is the record nobody's + gc.collect() + assert _wait_until(lambda: "m" not in cache._cached_models), "the finished admission stayed resident" + assert store.refcount("m") == 0 + assert budget.total_in_use() == 0 + + +def test_shutdown_retains_an_un_entered_wrapper_whose_hold_a_worker_death_zeroed(mock_logger): + """The grace is unowned, so a standing one does not mean nobody is working on the record: + here it is the last shield of a live LoadedModel wrapper whose hold the dead-worker recovery + zeroed. shutdown() must not treat it as an orphan — evicting it releases shared-store + ownership while the wrapper still holds the tensors (so a peer's reload mints a duplicate + canonical the budget counts once) and leaves the wrapper locking a detached record.""" + store = SharedCpuWeightsStore() + budget = RamBudget(max_bytes=10**12, shared_store=store) + cache_a = _make_cache(store, budget, mock_logger) + cache_b = _make_cache(store, budget, mock_logger) + try: + # An earlier, unrelated resident: admitted first, because a later put() would sweep the + # grace this test needs standing. + cache_a.put("other", DummyModule()) + _use_and_release(cache_a, "other") + + cache_a.put("m", DummyModule()) + record, claim = cache_a.get_with_first_use_claim("m") + loaded_model = LoadedModelWithoutConfig(cache_record=record, cache=cache_a, first_use_claim=claim) + assert record.awaiting_first_use and record.first_use_holds == 1 + canonical_before = store.peek("m") + + _kill_worker_abnormally(cache_a, record) + assert record.first_use_holds == 0, "premise broken: the recovery did not zero the hold" + assert record.awaiting_first_use, "premise broken: a live cache's grace was lifted" + + # A retrieval of the OTHER model restarts the worker (register_first_use_hold revives it), + # so shutdown() takes its live-worker path and never runs the dead-worker recovery — while + # our record keeps the grace as its only shield. + other_holder = cache_a.get_with_first_use_claim("other") + assert cache_a._deferred_work_thread is not None and cache_a._deferred_work_thread.is_alive() + assert record.first_use_holds == 0 and record.awaiting_first_use + + cache_a.shutdown() + assert cache_a._cached_models.get("m") is record, "shutdown() evicted a live wrapper's record" + assert store.refcount("m") == 1 + + with loaded_model as _model: + assert cache_a._cached_models.get("m") is record, "locked a detached record" + cache_b.put("m", DummyModule()) + _use_and_release(cache_b, "m") + assert store.peek("m") is canonical_before, "peer reload minted a duplicate canonical" + del other_holder + finally: + cache_b.shutdown() + + +def test_a_dead_worker_reclaims_a_record_left_shielded_by_an_admission_claim(mock_logger): + """The admission shield needs no recovery on a live cache — nothing has to release it, so a + worker death cannot strand it. Once the cache is shut down it does: the record is stale, and + the eviction its expiry should trigger travels through the dead worker. The recovery must + therefore retire it, or the model, its shared-store reference and its budget bytes stay + resident for the life of the process.""" + store = SharedCpuWeightsStore() + budget = RamBudget(max_bytes=10**12, shared_store=store) + cache = _make_cache(store, budget, mock_logger) + admission_claim = cache.put("m", DummyModule(), claim_admission=True) + assert admission_claim is not None + record = cache._cached_models["m"] + + cache.shutdown() + assert cache._cached_models.get("m") is record and record.is_stale + + # The worker dies after shutdown, with the claim still alive: its own dying recovery is the + # last event this record will ever see. + _kill_worker_abnormally(cache, record) + assert "m" not in cache._cached_models, "a record shielded with nothing left to unshield it" + assert store.refcount("m") == 0 + assert budget.total_in_use() == 0 + + +def test_releasing_an_admission_claim_retires_its_shield(mock_logger): + """The shield is the claim's liveness, so a spent claim must retire it explicitly: the loader + releases at its retrieval, but the object itself lives on to the end of its frame — or for as + long as a traceback holds that frame — and until then the record would stay invisible to every + asynchronous eviction path.""" + store = SharedCpuWeightsStore() + budget = RamBudget(max_bytes=10**12, shared_store=store) + cache = _make_cache(store, budget, mock_logger) + try: + admission_claim = cache.put("m", DummyModule(), claim_admission=True) + assert admission_claim is not None + record = cache._cached_models["m"] + assert record.admission_in_flight + + # The loader retrieves, uses and releases the model; the claim object is still alive here. + _use_and_release(cache, "m") + admission_claim.release() + assert not record.admission_in_flight, "a spent claim went on shielding its record" + assert record.in_first_use_window is False + + cache.evict_unlocked_for_peer(lambda: False) + assert "m" not in cache._cached_models, "a finished load's record stayed unreclaimable" + assert store.refcount("m") == 0 + assert budget.total_in_use() == 0 + finally: + cache.shutdown() + + +def test_a_failed_claim_construction_does_not_strand_its_hold(mock_logger, monkeypatch): + """The hold is armed before the claim that owns its release exists. If constructing the claim + raises — it allocates a weakref.finalize, and this cache runs at the RAM ceiling by design — + the hold must be handed back: an orphaned hold shields its record from every eviction path, + make_room, drop_model and shutdown's sweep included, for the life of the process.""" + store = SharedCpuWeightsStore() + budget = RamBudget(max_bytes=10**12, shared_store=store) + cache = _make_cache(store, budget, mock_logger) + cache.put("m", DummyModule()) + _use_and_release(cache, "m") + record = cache._cached_models["m"] + + def boom(*args, **kwargs): + raise MemoryError("no room for the claim") + + monkeypatch.setattr(model_cache_module, "FirstUseClaim", boom) + with pytest.raises(MemoryError): + cache.get_with_first_use_claim("m") + monkeypatch.undo() + + assert record.first_use_holds == 0, "a failed claim left its hold armed with no owner" + cache.shutdown() + assert "m" not in cache._cached_models + assert store.refcount("m") == 0 + assert budget.total_in_use() == 0 + + +def test_claim_shields_the_record_from_the_lookup_onward(mock_logger): + """The first-use hold must be armed by the lookup itself, not by the wrapper's constructor + (JPPhoto review, 2026-08-30). The stretch in between is not a few instructions — the + configured loader retrieves its record inside _load_and_cache and then does the shared-store + shell registration and two returns — and a shutdown landing there would evict the record its + holder is about to lock, releasing shared-store ownership while the tensors live on, so a + peer's reload of the key would mint a duplicate canonical copy the budget counts once.""" + store = SharedCpuWeightsStore() + budget = RamBudget(max_bytes=10**12, shared_store=store) + cache_a = _make_cache(store, budget, mock_logger) + cache_b = _make_cache(store, budget, mock_logger) + try: + cache_a.put("m", DummyModule()) + _use_and_release(cache_a, "m") # warm: past the admission grace, unlocked + + # The retrieval a load is about to wrap; shutdown() lands before the wrapper exists. + record, claim = cache_a.get_with_first_use_claim("m") + assert claim is not None + assert record.first_use_holds == 1, "the lookup did not arm the shield" + canonical_before = store.peek("m") + cache_a.shutdown() + + assert cache_a._cached_models.get("m") is record, "shutdown() evicted the record before its wrapper existed" + assert record.is_stale + assert store.refcount("m") == 1 + assert budget.total_in_use() == S + + # The wrapper is built late, adopts the claim, and locks a still-attached record. + loaded_model = LoadedModelWithoutConfig(cache_record=record, cache=cache_a, first_use_claim=claim) + with loaded_model as _model: + assert cache_a._cached_models.get("m") is record, "locked a detached record" + cache_b.put("m", DummyModule()) + _use_and_release(cache_b, "m") + assert store.peek("m") is canonical_before, "peer reload minted a duplicate canonical" + assert store.refcount("m") == 2 + assert budget.total_in_use() == S + + # Exiting the context is the record's last release: the shutdown's stale mark evicts it. + assert "m" not in cache_a._cached_models + assert store.refcount("m") == 1 + assert budget.total_in_use() == S + finally: + cache_b.shutdown() + + +def test_an_unadopted_claim_releases_its_hold_when_dropped(mock_logger): + """A claim is armed before any wrapper exists, so a load that raises in between must not leave + the hold standing: dropping the claim is itself a complete release, and on a shut-down cache + it evicts the record the sweep retained for it, with its accounting.""" + store = SharedCpuWeightsStore() + budget = RamBudget(max_bytes=10**12, shared_store=store) + cache = _make_cache(store, budget, mock_logger) + cache.put("m", DummyModule()) + _use_and_release(cache, "m") + + record, claim = cache.get_with_first_use_claim("m") + assert claim is not None and record.first_use_holds == 1 + + cache.shutdown() + assert cache._cached_models.get("m") is record, "shutdown() evicted the record out from under its claim" + + del claim # the load raised before it could build a wrapper + gc.collect() + assert _wait_until(lambda: "m" not in cache._cached_models), "a dropped claim stranded its record" + assert record.first_use_holds == 0 + assert store.refcount("m") == 0 + assert budget.total_in_use() == 0 + + +def test_a_spent_claim_cannot_consume_a_later_holders_hold(mock_logger): + """A claim releases exactly once. Once its wrapper has locked, the claim is spent, so its + later collection must not decrement a hold a different, still-live holder armed in the + meantime — that would unshield the second holder's window and let shutdown() evict the record + out from under it.""" + store = SharedCpuWeightsStore() + budget = RamBudget(max_bytes=10**12, shared_store=store) + cache = _make_cache(store, budget, mock_logger) + cache.put("m", DummyModule()) + _use_and_release(cache, "m") + + record, claim_a = cache.get_with_first_use_claim("m") + wrapper_a = LoadedModelWithoutConfig(cache_record=record, cache=cache, first_use_claim=claim_a) + with wrapper_a as _model: + pass + assert record.first_use_holds == 0, "the first lock did not release the adopted claim" + + # Holder B arms a fresh hold on the same record. + record_b, claim_b = cache.get_with_first_use_claim("m") + assert record_b is record and record.first_use_holds == 1 + + del wrapper_a, claim_a + gc.collect() + assert not _wait_until(lambda: record.first_use_holds == 0, timeout=0.5), "a spent claim consumed a later hold" + + cache.shutdown() + assert cache._cached_models.get("m") is record, "shutdown() evicted the record out from under holder B" + + del claim_b + gc.collect() + assert _wait_until(lambda: "m" not in cache._cached_models), "holder B's claim did not release" + assert store.refcount("m") == 0 + assert budget.total_in_use() == 0 + + +def test_abandoned_holder_reaches_zero_after_shutdown(mock_logger): + """A record retained by the shutdown sweep for a wrapper that is then dropped un-entered must + still reach zero (JPPhoto review, 2026-08-13): no unlock() is ever coming, so the wrapper's + abandonment finalizer — carried by the deferred worker, which therefore must outlive + shutdown() — is the only event left that can evict the record and release its shared-store + reference and budget bytes.""" + store = SharedCpuWeightsStore() + budget = RamBudget(max_bytes=10**12, shared_store=store) + cache = _make_cache(store, budget, mock_logger) + cache.put("m", DummyModule()) + loaded_model = LoadedModelWithoutConfig(cache_record=cache.get("m"), cache=cache) + + cache.shutdown() + assert "m" in cache._cached_models, "shutdown() evicted the record out from under its holder" + + del loaded_model + gc.collect() + assert _wait_until(lambda: "m" not in cache._cached_models), "the abandoned record was never evicted" + assert store.refcount("m") == 0 + assert budget.total_in_use() == 0 + + +def test_cold_admission_sweep_does_not_clear_wrapper_holds(mock_logger): + """put()'s stale-grace sweep must not unshield a live wrapper: a node may retrieve several + models and only then enter their contexts, so another model's cold admission (and its + make-room) can land inside a warm wrapper's get()->lock() window. The hold — unlike the + put()-set grace — survives the sweep, so a shutdown() after that admission still retains the + record for its holder.""" + store = SharedCpuWeightsStore() + budget = RamBudget(max_bytes=10**12, shared_store=store) + cache = _make_cache(store, budget, mock_logger) + cache.put("first", DummyModule()) + _use_and_release(cache, "first") + loaded_model = LoadedModelWithoutConfig(cache_record=cache.get("first"), cache=cache) + record = loaded_model._cache_record + assert record.first_use_holds == 1 + + cache.put("second", DummyModule()) # sweeps stale grace flags; must leave holds alone + assert record.first_use_holds == 1, "the cold admission cleared a live wrapper's hold" + + cache.shutdown() + assert cache._cached_models.get("first") is record, "shutdown() evicted the held record" + + with loaded_model as _model: + assert cache._cached_models.get("first") is record, "locked a detached record" + assert "first" not in cache._cached_models + assert store.refcount("first") == 0 + + +def test_post_shutdown_admission_is_evicted_after_its_use(mock_logger): + """put() after shutdown() (reachable: Invoker.stop() stops model_manager before + session_processor) missed the shutdown sweep, so the record is marked stale at admission — + its final release evicts it instead of leaving it resident until process exit.""" + store = SharedCpuWeightsStore() + budget = RamBudget(max_bytes=10**12, shared_store=store) + cache = _make_cache(store, budget, mock_logger) + cache.put("early", DummyModule()) + _use_and_release(cache, "early") + cache.shutdown() + + cache.put("late", DummyModule()) + assert cache._cached_models["late"].is_stale, "a post-shutdown admission must be stale at birth" + _use_and_release(cache, "late") + assert "late" not in cache._cached_models + assert store.refcount("late") == 0 + assert budget.total_in_use() == 0 + + +def test_make_room_skips_held_records(mock_logger): + """make_room (another model's cold load, or the keep-alive timeout's clear) must treat a + record with a live first-use hold like a locked one: its wrapper is about to lock it, and + evicting it would detach the record mid-window. Once the hold is consumed by the first lock, + the record is ordinary evictable content again.""" + store = SharedCpuWeightsStore() + budget = RamBudget(max_bytes=10**12, shared_store=store) + cache = _make_cache(store, budget, mock_logger) + try: + cache.put("held", DummyModule()) + _use_and_release(cache, "held") + loaded_model = LoadedModelWithoutConfig(cache_record=cache.get("held"), cache=cache) + + cache.make_room(10**12) + assert "held" in cache._cached_models, "make_room evicted a held record mid-window" + + with loaded_model as _model: + pass + cache.make_room(10**12) + assert "held" not in cache._cached_models + finally: + cache.shutdown() + + +def test_drop_model_defers_eviction_for_held_record(mock_logger): + """drop_model() must defer a held record exactly as it defers a locked one: mark it stale and + let the hold's release perform the eviction, instead of detaching the record from the holder + about to lock it.""" + store = SharedCpuWeightsStore() + budget = RamBudget(max_bytes=10**12, shared_store=store) + cache = _make_cache(store, budget, mock_logger) + try: + cache.put("m", DummyModule()) + _use_and_release(cache, "m") + loaded_model = LoadedModelWithoutConfig(cache_record=cache.get("m"), cache=cache) + record = loaded_model._cache_record + + dropped = cache.drop_model("m") + assert dropped == 0, "drop_model evicted a held record instead of deferring" + assert cache._cached_models.get("m") is record + assert record.is_stale + + with loaded_model as _model: + assert cache._cached_models.get("m") is record, "locked a detached record" + assert "m" not in cache._cached_models + assert budget.total_in_use() == 0 + finally: + cache.shutdown() + + +def test_unlock_stale_eviction_defers_to_live_holder(mock_logger): + """The last unlock() of a stale record must not evict it while another wrapper still holds it + for its own upcoming lock — that wrapper's own release performs the eviction instead.""" + store = SharedCpuWeightsStore() + budget = RamBudget(max_bytes=10**12, shared_store=store) + cache = _make_cache(store, budget, mock_logger) + cache.put("m", DummyModule()) + _use_and_release(cache, "m") + record = cache.get("m") + cache.lock(record, None) # generation A is using the model + loaded_model = LoadedModelWithoutConfig(cache_record=cache.get("m"), cache=cache) # holder B + + cache.shutdown() # marks the locked record stale + cache.unlock(record) # A finishes; B's hold defers the stale eviction + assert cache._cached_models.get("m") is record, "unlock evicted a record another wrapper holds" + + with loaded_model as _model: + assert cache._cached_models.get("m") is record, "locked a detached record" + assert "m" not in cache._cached_models + assert store.refcount("m") == 0 + assert budget.total_in_use() == 0 + + +def test_worker_death_zeroes_stranded_first_use_holds(mock_logger): + """A hold whose abandonment release was dispatched into a dead worker is dropped for good — + finalizers fire once — so the next worker start must zero the surviving holds: the + alternative is a record shielded from every eviction path for the life of the process.""" + store = SharedCpuWeightsStore() + budget = RamBudget(max_bytes=10**12, shared_store=store) + cache = _make_cache(store, budget, mock_logger) + try: + cache.put("m", DummyModule()) + loaded_model = LoadedModelWithoutConfig(cache_record=cache.get("m"), cache=cache) + record = loaded_model._cache_record + assert record.first_use_holds == 1 + + # The worker dies (the way a raising log handler would kill it) ... + cache._deferred_work_queue.put(model_cache_module._DEFERRED_STOP) + assert cache._deferred_work_thread is not None + cache._deferred_work_thread.join(timeout=10) + # ... and the wrapper's abandonment release is dispatched into the void. + del loaded_model + gc.collect() + assert record.first_use_holds == 1, "the release was not dropped — dead-worker premise broken" + + # The next admission starts a replacement worker, which clears the stranded hold. + cache.put("next", DummyModule()) + assert record.first_use_holds == 0, "a stranded hold survived the worker restart" + finally: + cache.shutdown() + + +def test_stale_hold_release_cannot_steal_a_fresh_hold(mock_logger): + """A release from before a dead-worker zeroing sweep must not decrement a hold armed after + it: the zeroing bumps the record's hold epoch, and releases quote the epoch they were armed + under. Without that, a surviving wrapper's late first-lock release would silently consume a + different wrapper's fresh shield, and a shutdown() in that wrapper's window would evict the + record out from under it — the exact defect the holds exist to prevent.""" + store = SharedCpuWeightsStore() + budget = RamBudget(max_bytes=10**12, shared_store=store) + cache = _make_cache(store, budget, mock_logger) + try: + cache.put("m", DummyModule()) + _use_and_release(cache, "m") + wrapper_a = LoadedModelWithoutConfig(cache_record=cache.get("m"), cache=cache) + record = wrapper_a._cache_record + assert record.first_use_holds == 1 + + # The worker dies; the next admission zeroes the stranded hold and bumps the epoch. + cache._deferred_work_queue.put(model_cache_module._DEFERRED_STOP) + assert cache._deferred_work_thread is not None + cache._deferred_work_thread.join(timeout=10) + cache.put("x", DummyModule()) + assert record.first_use_holds == 0 + + # A fresh wrapper arms a new-epoch hold under the healthy replacement worker. + wrapper_b = LoadedModelWithoutConfig(cache_record=cache.get("m"), cache=cache) + assert record.first_use_holds == 1 + + # Wrapper A's late first-lock release quotes the old epoch: it must be a no-op. + with wrapper_a as _model: + pass + assert record.first_use_holds == 1, "a stale release consumed the fresh wrapper's hold" + + # And shutdown() in wrapper B's window therefore still retains the record for it. + cache.shutdown() + assert cache._cached_models.get("m") is record, "shutdown() evicted the record out from under its holder" + with wrapper_b as _model: + assert cache._cached_models.get("m") is record, "locked a detached record" + assert "m" not in cache._cached_models + assert store.refcount("m") == 0 + finally: + cache.shutdown() + + +def test_stale_release_drained_after_worker_restart_is_rejected_by_epoch(mock_logger): + """A release enqueued while the old worker was alive survives its death in the SimpleQueue + (the queue is never cleared on restart) and is drained by the replacement worker AFTER + dead-worker recovery zeroed the holds and bumped the epoch. The abandonment handler must + reject it: the hold it quotes was already accounted for by the zeroing, so honoring it would + consume a FRESH hold armed by a different wrapper under the healthy worker.""" + store = SharedCpuWeightsStore() + budget = RamBudget(max_bytes=10**12, shared_store=store) + cache = _make_cache(store, budget, mock_logger) + try: + cache.put("m", DummyModule()) + _use_and_release(cache, "m") + record = cache.get("m") + epoch_before = cache.register_first_use_hold(record) # wrapper A's hold + assert epoch_before is not None and record.first_use_holds == 1 + + # The worker dies; A's abandonment release is already sitting in the queue, undrained. + cache._deferred_work_queue.put(model_cache_module._DEFERRED_STOP) + assert cache._deferred_work_thread is not None + cache._deferred_work_thread.join(timeout=10) + + # The next admission runs recovery (zero + epoch bump) and starts the replacement worker. + cache.put("x", DummyModule()) + assert record.first_use_holds == 0 + assert record.first_use_holds_epoch == epoch_before + 1 + + # Wrapper B arms a fresh hold under the healthy worker. + epoch_after = cache.register_first_use_hold(record) + assert epoch_after == epoch_before + 1 and record.first_use_holds == 1 + + # The replacement worker drains A's stale release (invoked directly here — it is exactly + # what _run_deferred_work does with the surviving queue item): it must be a no-op. + cache._release_abandoned_holder(record, True, epoch_before) + assert record.first_use_holds == 1, "a stale queued release consumed the fresh wrapper's hold" + + # B's own release, quoting the current epoch, works normally. + cache.release_first_use_hold(record, epoch_after) + assert record.first_use_holds == 0 + finally: + cache.shutdown() + + +def test_shutdown_clears_holds_stranded_by_a_dead_worker(mock_logger): + """shutdown() must run the dead-worker hold recovery itself: a hold whose abandonment + release was dropped by the dead-thread dispatch check has no other releaser — no unlock() is + coming for a never-locked holder, and after shutdown no put() is guaranteed to run the usual + next-start recovery — so without this the sweep would stale-retain the record, its + shared-store refcount and its budget bytes for the life of the process.""" + store = SharedCpuWeightsStore() + budget = RamBudget(max_bytes=10**12, shared_store=store) + cache = _make_cache(store, budget, mock_logger) + cache.put("m", DummyModule()) + _use_and_release(cache, "m") + loaded_model = LoadedModelWithoutConfig(cache_record=cache.get("m"), cache=cache) + record = loaded_model._cache_record + assert record.first_use_holds == 1 + + # The worker dies; the wrapper is dropped and its release is dispatched into the void. + cache._deferred_work_queue.put(model_cache_module._DEFERRED_STOP) + assert cache._deferred_work_thread is not None + cache._deferred_work_thread.join(timeout=10) + del loaded_model + gc.collect() + assert record.first_use_holds == 1, "the release was not dropped — dead-worker premise broken" + + cache.shutdown() + assert "m" not in cache._cached_models, "shutdown() stale-retained a record nothing can ever release" + assert store.refcount("m") == 0 + assert budget.total_in_use() == 0 + + +class _WorkerKill(BaseException): + """Deliberately not an ``Exception``: the deferred worker catches those per work item, so this + escapes its loop the way an asynchronously raised BaseException — or a failure of the queue's + own get() — would, i.e. the abnormal death the recovery path exists for.""" + + +def _kill_worker_abnormally(cache: ModelCache, record) -> threading.Thread: + """Kill the deferred worker without the orderly _DEFERRED_STOP, and wait for it to unwind.""" + worker = cache._deferred_work_thread + assert worker is not None and worker.is_alive() + with patch.object(cache, "_release_abandoned_holder", side_effect=_WorkerKill): + cache._deferred_work_queue.put(model_cache_module._AbandonedHolderRelease(weakref.ref(record), False, 0)) + worker.join(timeout=10) + assert not worker.is_alive(), "the worker did not die" + return worker + + +def test_post_shutdown_admission_gets_no_first_use_grace(mock_logger): + """put() after shutdown() must not arm the post-admission grace. + + The grace's backstop releaser is the sweep at the top of the NEXT put(), and after shutdown no + further put() is guaranteed. A load cancelled between put() and the LoadedModel's construction + leaves no wrapper, hence no finalizer either — so an armed flag would stand for the life of the + process, hiding the record from every asynchronous eviction path while its bytes stay charged + to the shared budget.""" + store = SharedCpuWeightsStore() + budget = RamBudget(max_bytes=10**12, shared_store=store) + cache = _make_cache(store, budget, mock_logger) + cache.put("early", DummyModule()) + _use_and_release(cache, "early") + cache.shutdown() + # The worker deliberately outlives shutdown(), so worker liveness is not what withholds the + # grace here. + assert cache._deferred_work_thread is not None and cache._deferred_work_thread.is_alive() + + cache.put("late", DummyModule()) # ... and this load is cancelled before it ever calls get() + record = cache._cached_models["late"] + assert not record.awaiting_first_use, "a post-shutdown admission armed a grace nothing can release" + assert not record.in_first_use_window + + # Unshielded, so the asynchronous paths can take the abandoned record back. + assert cache.evict_unlocked_for_peer(lambda: False) == 1 + assert "late" not in cache._cached_models + assert store.refcount("late") == 0 + assert budget.total_in_use() == 0 + + +@pytest.mark.filterwarnings("ignore::pytest.PytestUnhandledThreadExceptionWarning") +def test_worker_death_after_shutdown_recovers_its_stranded_holds(mock_logger): + """A worker that dies *after* shutdown()'s liveness check must run the recovery itself. + + shutdown() retains a record whose wrapper is still inside its get()->lock() window, counting on + the worker to carry the wrapper's abandonment release. If the worker then dies, that release is + dropped by the dead-thread dispatch check, no unlock() is coming for a never-locked holder, and + a shut-down cache takes no further admission to run the next-start recovery — so the record, + its shared-store reference and its budget bytes would be stranded for the life of the + process.""" + store = SharedCpuWeightsStore() + budget = RamBudget(max_bytes=10**12, shared_store=store) + cache = _make_cache(store, budget, mock_logger) + cache.put("m", DummyModule()) + _use_and_release(cache, "m") + loaded_model = LoadedModelWithoutConfig(cache_record=cache.get("m"), cache=cache) + record = loaded_model._cache_record + assert record.first_use_holds == 1 + + # shutdown() sees a live worker, so it retains the held record for its holder. + cache.shutdown() + assert cache._cached_models.get("m") is record + assert store.refcount("m") == 1 + + _kill_worker_abnormally(cache, record) + + assert record.first_use_holds == 0, "the dying worker left its stranded hold standing" + assert "m" not in cache._cached_models, "a record nothing can ever release stayed resident" + assert store.refcount("m") == 0 + assert budget.total_in_use() == 0 + assert cache._deferred_work_thread is None, "the dead worker was left in the slot as if healthy" + + # The surviving wrapper is still live; dropping it must not disturb the settled accounting. + del loaded_model + gc.collect() + assert budget.total_in_use() == 0 + + +def test_dying_worker_recovery_leaves_a_replacement_workers_shields_alone(mock_logger): + """The dying worker's recovery is scoped by identity: once a replacement has taken the slot, + the shields standing are the replacement's, and _ensure_deferred_worker already cleared + whatever the dead thread stranded. A late unwind must not zero them.""" + store = SharedCpuWeightsStore() + budget = RamBudget(max_bytes=10**12, shared_store=store) + cache = _make_cache(store, budget, mock_logger) + try: + cache.put("m", DummyModule()) + _use_and_release(cache, "m") + loaded_model = LoadedModelWithoutConfig(cache_record=cache.get("m"), cache=cache) + record = loaded_model._cache_record + assert record.first_use_holds == 1 + + # A thread that is not the current worker reports its death. + cache._recover_from_dead_worker(threading.current_thread()) + assert record.first_use_holds == 1, "an impostor's recovery zeroed the live worker's holds" + assert cache._deferred_work_thread is not None and cache._deferred_work_thread.is_alive() + + # And the hold still works: shutdown() retains the record for its holder. + cache.shutdown() + assert cache._cached_models.get("m") is record + with loaded_model as _model: + pass + assert "m" not in cache._cached_models + assert budget.total_in_use() == 0 + finally: + cache.shutdown() + + +@pytest.mark.filterwarnings("ignore::pytest.PytestUnhandledThreadExceptionWarning") +def test_dying_worker_does_not_unshield_a_live_caches_fresh_admission(mock_logger): + """A worker death must not turn into a failed load on a cache that is still running. + + The dying worker is `is_alive()` for as long as it unwinds, so a cold load landing in that + window starts no replacement worker and is admitted with the ordinary post-admission grace. + That grace's backstop is the sweep at the top of the NEXT put(), not the worker, and the + loader has not even called get() yet — clearing it from the dying frame would leave the record + exposed to a reconcile and the loader's get() raising IndexError.""" + store = SharedCpuWeightsStore() + budget = RamBudget(max_bytes=10**12, shared_store=store) + cache = _make_cache(store, budget, mock_logger) + try: + cache.put("resident", DummyModule()) + record = _use_and_release(cache, "resident") + worker = cache._deferred_work_thread + assert worker is not None + + # Hold the worker inside its dying recovery so the admission below lands mid-unwind. + entered = threading.Event() + released = threading.Event() + real_recover = cache._recover_from_dead_worker + + def gated(dying_worker): + entered.set() + assert released.wait(timeout=10) + real_recover(dying_worker) + + with ( + patch.object(cache, "_release_abandoned_holder", side_effect=_WorkerKill), + patch.object(cache, "_recover_from_dead_worker", side_effect=gated), + ): + cache._deferred_work_queue.put(model_cache_module._AbandonedHolderRelease(weakref.ref(record), False, 0)) + assert entered.wait(timeout=10), "the worker never reached its recovery" + + # A cold load lands while the worker is still alive-but-unwinding. + assert worker.is_alive() + cache.put("loading", DummyModule()) + admitted = cache._cached_models["loading"] + assert admitted.awaiting_first_use, "premise broken: the admission was not graced" + + released.set() + worker.join(timeout=10) + + assert admitted.awaiting_first_use, "the dying worker unshielded a load still between put() and get()" + # The loader's get() therefore still finds its model. + assert cache.get("loading") is admitted + finally: + cache.shutdown() + + +def test_dead_worker_recovery_at_admission_keeps_a_retained_records_accounting(mock_logger): + """The recovery run when an admission finds the worker dead must not evict a record whose + tensors live wrappers still hold. + + Releasing store ownership while the weights live on is the accounting lie shutdown() itself + refuses to make: the store stops counting bytes that are still resident, so a peer's reload of + the key mints a duplicate canonical while the budget counts one. Lifting the stranded shield + is enough — a replacement worker is started two lines later, and the ordinary eviction paths + are all still reachable.""" + store = SharedCpuWeightsStore() + budget = RamBudget(max_bytes=10**12, shared_store=store) + cache = _make_cache(store, budget, mock_logger) + cache.put("m", DummyModule()) + _use_and_release(cache, "m") + wrapper_a = LoadedModelWithoutConfig(cache_record=cache.get("m"), cache=cache) + record = wrapper_a._cache_record + + # shutdown() retains the record for its holder ... + cache.shutdown() + assert cache._cached_models.get("m") is record + + # ... and the worker is then lost without running its own recovery (a fork, or a death this + # admission races). + cache._deferred_work_queue.put(model_cache_module._DEFERRED_STOP) + assert cache._deferred_work_thread is not None + cache._deferred_work_thread.join(timeout=10) + + # A second wrapper's construction runs the dead-worker recovery from _ensure_deferred_worker. + wrapper_b = LoadedModelWithoutConfig(cache_record=cache.get("m"), cache=cache) + assert cache._cached_models.get("m") is record, "the recovery detached a record wrappers still hold" + assert store.refcount("m") == 1, "shared-store ownership was released while the tensors live on" + assert budget.total_in_use() == S + assert record.first_use_holds == 1, "the replacement worker's hold was not armed for wrapper B" + + # And the record still settles to zero once that holder is done with it. + with wrapper_b as _model: + pass + assert "m" not in cache._cached_models + assert store.refcount("m") == 0 + assert budget.total_in_use() == 0 + del wrapper_a + + +def test_first_use_grace_is_lifted_only_once_the_cache_is_shut_down(mock_logger): + """The two halves of the recovery's grace rule, at the seam itself: on a live cache the grace + survives (the next put()'s sweep is still its backstop); once shut down it is lifted, because + no further put() is guaranteed and a grace stale-retained by the shutdown sweep would + otherwise shield its record for the life of the process.""" + store = SharedCpuWeightsStore() + budget = RamBudget(max_bytes=10**12, shared_store=store) + cache = _make_cache(store, budget, mock_logger) + cache.put("m", DummyModule()) + record = cache._cached_models["m"] + assert record.awaiting_first_use + + cache._recover_stranded_shields() + assert record.awaiting_first_use, "a live cache's grace was lifted before it lost its backstop" + + # shutdown() stale-retains the graced record; recovery must then lift it. + cache.shutdown() + assert cache._cached_models.get("m") is record + assert record.is_stale + cache._recover_stranded_shields() + assert not record.awaiting_first_use, "a shut-down cache's orphaned grace was left standing" + cache._evict_stale_unshielded_entries() + assert "m" not in cache._cached_models + assert store.refcount("m") == 0 + assert budget.total_in_use() == 0 + + +def test_queued_abandonment_release_does_not_pin_its_record(mock_logger): + """A release sitting in the deferred queue must pin nothing. + + Such an item can outlive every chance to drain it: _dispatch_deferred's liveness gate is + unsynchronized, so a finalizer that read the worker slot just before it was retired still + enqueues against a thread that is alive only because it is unwinding, and on a shut-down cache + nothing drains the queue again. A strong reference to the CacheRecord would hold that model's + whole CPU state dict for the life of the process while the store and the budget both reported + the bytes released.""" + store = SharedCpuWeightsStore() + budget = RamBudget(max_bytes=10**12, shared_store=store) + cache = _make_cache(store, budget, mock_logger) + cache.put("parked", DummyModule()) + parked = _use_and_release(cache, "parked") + cache.put("m", DummyModule()) + record = _use_and_release(cache, "m") + record_ref = weakref.ref(record) + module_ref = weakref.ref(record.cached_model.model) + + # Park the worker on an earlier item so the release below stays queued behind it. + entered = threading.Event() + unblock = threading.Event() + + def park(*_args): + entered.set() + assert unblock.wait(timeout=10) + + with patch.object(cache, "_release_abandoned_holder", side_effect=park): + cache._deferred_work_queue.put(model_cache_module._AbandonedHolderRelease(weakref.ref(parked), False, 0)) + assert entered.wait(timeout=10), "the worker never picked up the parking item" + + # The real dispatch path, with a worker that is alive but will never get to this item. + cache.release_first_use_grace(record, held_first_use=True, hold_epoch=record.first_use_holds_epoch) + assert cache._deferred_work_queue.qsize() == 1, "premise broken: the release was not queued" + + cache.shutdown() + assert "m" not in cache._cached_models + assert store.refcount("m") == 0 + + # The accounting says those bytes are gone, so they must really be gone. + del record + assert _collect_until(lambda: record_ref() is None), "the queued release pinned the CacheRecord" + assert module_ref() is None, "the evicted model is still resident in RAM" + assert cache._deferred_work_queue.qsize() == 1, "premise broken: the item was drained after all" + + unblock.set() + + +@pytest.mark.filterwarnings("ignore::pytest.PytestUnhandledThreadExceptionWarning") +def test_shutdown_lifts_a_grace_the_dying_worker_left_standing(mock_logger): + """The dying worker deliberately leaves a live cache's admission grace alone — the sweep at the + top of the next put() is still its backstop. shutdown() is where that backstop runs out, and it + must lift the grace even though the dying worker retired the worker slot on its way out. Keying + that lift on a dead thread still sitting in the slot would skip exactly the case the dying + recovery handed to it, stranding the record, its shared-store reference and its budget bytes + for the life of the process.""" + store = SharedCpuWeightsStore() + budget = RamBudget(max_bytes=10**12, shared_store=store) + cache = _make_cache(store, budget, mock_logger) + cache.put("resident", DummyModule()) + resident = _use_and_release(cache, "resident") + + # A load admitted and then cancelled before get(): no wrapper is ever built, so no finalizer + # will ever exist to release this grace. + cache.put("orphan", DummyModule()) + orphan = cache._cached_models["orphan"] + assert orphan.awaiting_first_use + + _kill_worker_abnormally(cache, resident) + assert cache._deferred_work_thread is None, "premise broken: the dying worker kept the slot" + assert orphan.awaiting_first_use, "the grace was lifted while the next put() was still its backstop" + + cache.shutdown() + assert "orphan" not in cache._cached_models, "shutdown() stale-retained a grace nothing can release" + assert store.refcount("orphan") == 0 + assert budget.total_in_use() == 0 + + +@pytest.mark.filterwarnings("ignore::pytest.PytestUnhandledThreadExceptionWarning") +def test_worker_start_retries_a_recovery_that_failed_partway(mock_logger): + """A dying worker's recovery can itself fail — a logging handler that raises is one of the ways + the worker dies in the first place — and it has already retired the worker slot by then. The + next worker start must re-run the recovery rather than key on a dead thread still sitting in + the slot, or the hold it never reached stays shielded for good.""" + store = SharedCpuWeightsStore() + budget = RamBudget(max_bytes=10**12, shared_store=store) + cache = _make_cache(store, budget, mock_logger) + cache.put("m", DummyModule()) + _use_and_release(cache, "m") + loaded_model = LoadedModelWithoutConfig(cache_record=cache.get("m"), cache=cache) + record = loaded_model._cache_record + assert record.first_use_holds == 1 + + with patch.object(cache, "_clear_stranded_first_use_holds", side_effect=RuntimeError("logging handler blew up")): + _kill_worker_abnormally(cache, record) + assert record.first_use_holds == 1, "premise broken: the recovery did not fail" + assert cache._deferred_work_thread is None + + # The next worker start re-runs the recovery. + cache.put("next", DummyModule()) + assert record.first_use_holds == 0, "a recovery that failed partway was never retried" + + # And with the hold gone, the record settles instead of being shielded for good. + del loaded_model + cache.shutdown() + assert "m" not in cache._cached_models + assert store.refcount("m") == 0 + + +def test_hold_recovery_unshields_every_record_before_reporting_any(mock_logger): + """A logging handler that raises is one of the ways the deferred worker dies in the first + place, so the recovery must unshield every record before it reports any of them. Logging inside + the loop would let that same handler abort the sweep partway, leaving the records it never + reached shielded with nothing left to unshield them.""" + store = SharedCpuWeightsStore() + budget = RamBudget(max_bytes=10**12, shared_store=store) + cache = _make_cache(store, budget, mock_logger) + try: + wrappers = [] + for key in ("a", "b"): + cache.put(key, DummyModule()) + _use_and_release(cache, key) + wrappers.append(LoadedModelWithoutConfig(cache_record=cache.get(key), cache=cache)) + assert all(w._cache_record.first_use_holds == 1 for w in wrappers) + + raising_logger = MagicMock() + raising_logger.getEffectiveLevel.return_value = logging.INFO + raising_logger.warning.side_effect = RuntimeError("logging handler blew up") + with patch.object(cache, "_logger", raising_logger): + with pytest.raises(RuntimeError): + cache._clear_stranded_first_use_holds() + + assert [w._cache_record.first_use_holds for w in wrappers] == [0, 0], ( + "a raising log handler aborted the sweep partway, leaving a record shielded for good" + ) + finally: + cache.shutdown() + + def test_dropped_cache_releases_shared_weights_on_collection(mock_logger): """A cache dropped without shutdown() must not strand its shared-weights references: the store's refcount and bytes — and therefore the budget total — must return to zero once the @@ -294,9 +1571,11 @@ def test_model_in_ram_on_a_cold_record_ends_the_grace_and_detaches_the_finalizer with loaded_model.model_in_ram(): assert record.is_locked assert not record.awaiting_first_use - # The grace has been consumed by the pin; the finalizer must be detached so a later GC - # of the handle does not queue a redundant grace release. - assert not loaded_model._first_use_finalizer.alive + # The grace has been consumed by the pin; the finalizer must be dropped (and the + # wrapper's first-use hold released) so a later GC of the handle does not queue a + # redundant release. + assert loaded_model._first_use_finalizer is None + assert record.first_use_holds == 0 assert not record.is_locked # Post-pin, the record is ordinary evictable cache content. @@ -1139,16 +2418,30 @@ def fail_first_worker_start(thread: threading.Thread) -> None: cache.shutdown() -def test_shutdown_stops_deferred_worker(mock_logger): +def test_deferred_worker_survives_shutdown_and_exits_on_collection(mock_logger): + """shutdown() must NOT stop the deferred worker: a record the shutdown sweep retained because + a live LoadedModel wrapper had not locked it yet has no future unlock() if that wrapper is + simply dropped — the wrapper's abandonment release, carried by the worker, is the only event + left that can evict it. The worker exits when the cache itself is collected (via the + finalizer registered at worker start), not before.""" store = SharedCpuWeightsStore() budget = RamBudget(max_bytes=int(S * 1.4), shared_store=store) cache = _make_cache(store, budget, mock_logger) cache.put("model", DummyModule()) cache.shutdown() - cache._deferred_work_thread.join(timeout=10) + worker = cache._deferred_work_thread + assert worker is not None + # join() with a timeout rather than a bare is_alive() check: a worker that shutdown() told to + # stop may not have consumed the sentinel yet, and a racy alive-reading would pass anyway. + worker.join(timeout=2) + assert worker.is_alive(), "shutdown() stopped the deferred worker" - assert not cache._deferred_work_thread.is_alive() + cache_ref = weakref.ref(cache) + del cache + assert _wait_until(lambda: (gc.collect(), cache_ref() is None)[1]), "the cache was not collected" + worker.join(timeout=10) + assert not worker.is_alive(), "the worker outlived its collected cache" @pytest.mark.parametrize("finalizer_order", ["before_cache_release", "after_cache_release"]) @@ -1540,8 +2833,9 @@ def test_deferred_dispatch_is_dropped_when_no_worker_is_running(mock_logger): assert idle_cache.cached_model_keys() == set() assert idle_cache._deferred_work_queue.qsize() == 0 - # After shutdown the same must hold for a cache that *does* have a (now stopped) worker. - busy_cache.shutdown() + # The same must hold for a cache whose worker has died (shutdown() no longer stops the + # worker, so simulate a death the way a raising log handler would cause one). + busy_cache._deferred_work_queue.put(model_cache_module._DEFERRED_STOP) assert busy_cache._deferred_work_thread is not None busy_cache._deferred_work_thread.join(timeout=10) assert not busy_cache._deferred_work_thread.is_alive() @@ -1641,33 +2935,55 @@ def test_dropped_non_shared_cache_releases_only_its_budget_charge(mock_logger): assert budget.total_in_use() == S -def test_admission_without_a_worker_gets_no_first_use_grace(mock_logger): +def test_admission_without_a_worker_gets_no_first_use_grace_or_holds(mock_logger, monkeypatch): """A record must never be shielded from eviction with nothing left able to unshield it. - put() after shutdown() is reachable in production — Invoker.stop() stops model_manager before - session_processor, so an in-flight generation can admit a model after every cache has been shut - down — and _ensure_deferred_worker deliberately does not revive the worker there. Granting the - grace anyway would leave the record permanently invisible to both asynchronous eviction paths - while its bytes stayed charged to the shared budget. + put() (and register_first_use_hold) first try to revive a dead worker, so the no-worker state + only persists when the thread cannot be started at all (RLIMIT_NPROC, a container's pids.max). + Granting the grace — or arming a wrapper hold — there would leave the record permanently + invisible to the asynchronous eviction paths while its bytes stayed charged to the shared + budget, because the releases both travel through the worker. """ store = SharedCpuWeightsStore() budget = RamBudget(max_bytes=int(S * 8), shared_store=store) cache = _make_cache(store, budget, mock_logger) + real_thread_start = threading.Thread.start + + def fail_worker_start(thread: threading.Thread) -> None: + if thread.name == "model-cache-deferred-work": + raise RuntimeError("forced thread start failure") + real_thread_start(thread) + try: cache.put("normal", DummyModule()) assert cache._cached_models["normal"].awaiting_first_use, "a healthy admission keeps its grace" + _use_and_release(cache, "normal") - cache.shutdown() + # Kill the worker, then fail every restart so the cache truly has none. + cache._deferred_work_queue.put(model_cache_module._DEFERRED_STOP) assert cache._deferred_work_thread is not None cache._deferred_work_thread.join(timeout=10) + monkeypatch.setattr(threading.Thread, "start", fail_worker_start) cache.put("late", DummyModule()) record = cache._cached_models["late"] - assert cache._deferred_work_thread is not None and not cache._deferred_work_thread.is_alive() assert not record.awaiting_first_use, "admitted with a grace no worker can ever release" - # Being unshielded, it is reachable by the synchronous eviction path. + # A wrapper constructed now must not arm a hold (its finalizer's release would be + # dropped), and must therefore not register the finalizer either. + loaded_model = LoadedModelWithoutConfig(cache_record=cache.get("late"), cache=cache) + assert record.first_use_holds == 0, "armed a hold no worker can ever release" + assert loaded_model._first_use_finalizer is None + + # Same for a retrieval that asks for a claim: there is nothing to carry its release. + record_again, claim = cache.get_with_first_use_claim("late") + assert record_again is record + assert claim is None, "armed a claim no worker can ever release" + assert record.first_use_holds == 0 + + # Being unshielded, the record is reachable by the synchronous eviction path. cache._delete_cache_entry(record) assert "late" not in cache._cached_models finally: + monkeypatch.undo() cache.shutdown() diff --git a/tests/backend/model_manager/load/model_cache/test_model_cache_timeout.py b/tests/backend/model_manager/load/model_cache/test_model_cache_timeout.py index 5b27b3d8ef6..74a1337dcfe 100644 --- a/tests/backend/model_manager/load/model_cache/test_model_cache_timeout.py +++ b/tests/backend/model_manager/load/model_cache/test_model_cache_timeout.py @@ -109,18 +109,26 @@ def test_no_timeout_keeps_models(model_cache_no_timeout): def test_shutdown_cancels_timer(model_cache_with_timeout): - """Test that shutdown properly cancels the timeout timer.""" + """Test that shutdown properly cancels the timeout timer and evicts resident records.""" cache = model_cache_with_timeout - # Add a model to start the timer + # Add a model to start the timer, and complete its load-use cycle so the record is an + # ordinary idle resident (a record still inside the put()->lock() admission window is + # deliberately retained by shutdown()). test_tensor = torch.randn(10, 10) cache.put("test_model", test_tensor) + record = cache.get("test_model") + cache.lock(record, None) + cache.unlock(record) + assert cache._timeout_timer is not None # Shutdown the cache cache.shutdown() - # Wait for what would be the timeout - time.sleep(1.0) + # The timer is cancelled and the idle record is evicted with its accounting. + assert cache._timeout_timer is None + assert "test_model" not in cache._cached_models - # The model should still be in the cache since shutdown was called - assert "test_model" in cache._cached_models + # Wait for what would be the timeout; the cancelled timer must not fire or arm a new one. + time.sleep(1.0) + assert cache._timeout_timer is None