diff --git a/invokeai/backend/model_manager/load/load_base.py b/invokeai/backend/model_manager/load/load_base.py index 7225fd1402f..d4094b5557c 100644 --- a/invokeai/backend/model_manager/load/load_base.py +++ b/invokeai/backend/model_manager/load/load_base.py @@ -56,15 +56,50 @@ class LoadedModelWithoutConfig: def __init__(self, cache_record: CacheRecord, cache: ModelCache): 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 - else None + # Shield the record for the window between this wrapper's construction and its first + # lock: without it, an eviction sweep racing that gap — a peer's budget reconcile, + # another model's make-room, or the cache's shutdown() — would evict the record out from + # under this wrapper, detaching it from the cache's RAM accounting and (for shared + # weights) from store ownership while its tensors live on. The few instructions between + # get() returning and this constructor arming the hold remain 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 + # hold is released exactly once: on the + # first lock (_end_first_use_window), or by the finalizer below if this wrapper is + # dropped without ever locking. The finalizer 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 hold. + 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 register_hold is not None and release_grace is not None else None ) - if self._first_use_finalizer is not None: + self._first_use_finalizer = None + if 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 + 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_finalizer is not None: + 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 # load_state_dict(assign=True), which calls register_parameter) so it can't overlap a @@ -72,8 +107,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 +130,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 +146,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: 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..a921bca96d2 100644 --- a/invokeai/backend/model_manager/load/model_cache/cache_record.py +++ b/invokeai/backend/model_manager/load/model_cache/cache_record.py @@ -30,13 +30,31 @@ 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. From the wrapper's construction on, the window is tracked by first_use_holds + # below, whose release is guaranteed by the wrapper's finalizer rather than by the sweep. awaiting_first_use: bool = False + # Count of live LoadedModel wrappers holding this record that have not yet locked it. Armed by + # ModelCache.register_first_use_hold() (called from LoadedModelWithoutConfig.__init__) and + # released exactly once per wrapper — on the wrapper's first lock, or by its weakref finalizer + # if it is dropped without ever locking. Unlike awaiting_first_use, these holds are NOT swept + # by the next admission: a warm get()'s wrapper 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 after the deferred worker — the thread that carries + # finalizer-initiated releases — is found dead (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 def lock(self) -> None: """Lock this record.""" @@ -51,3 +69,15 @@ def unlock(self) -> None: def is_locked(self) -> bool: """Return true if record is locked.""" return self._locks > 0 + + @property + def in_first_use_window(self) -> bool: + """True while a load or a live LoadedModel wrapper is between obtaining this record and + locking it. 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 diff --git a/invokeai/backend/model_manager/load/model_cache/cached_model/cached_model_only_full_load.py b/invokeai/backend/model_manager/load/model_cache/cached_model/cached_model_only_full_load.py index ccf40654575..226f65f6998 100644 --- a/invokeai/backend/model_manager/load/model_cache/cached_model/cached_model_only_full_load.py +++ b/invokeai/backend/model_manager/load/model_cache/cached_model/cached_model_only_full_load.py @@ -1,3 +1,4 @@ +import weakref from typing import Any import torch @@ -41,6 +42,7 @@ def __init__( # under `cache_key`; `release_shared_weights()` must be called exactly once on eviction. self._shared_store: SharedCpuWeightsStore | None = None self._shared_key: str | None = None + self._shared_release_finalizer: weakref.finalize | None = None # A CPU read-only copy of the model's state dict. self._cpu_state_dict: dict[str, torch.Tensor] | None = None @@ -57,9 +59,24 @@ def __init__( if canonical is not cpu_state_dict: model.load_state_dict(canonical, assign=True) cpu_state_dict = canonical + # A cache dropped without a shutdown() never routes its records through + # _delete_cache_entry, so nothing would call release_shared_weights() and the + # canonical tensors would stay resident (and counted by the RAM budget) + # forever. The finalizer must not reference `self` (its args are held strongly + # — that would make the wrapper immortal) and must not take the store's + # non-reentrant lock (it runs in GC context): release_deferred only enqueues; + # the store applies it on its next operation. release_shared_weights() detaches + # this on the normal eviction path, so the release happens exactly once either + # way. Registered inside this try so a failure here (e.g. MemoryError) + # releases the just-acquired reference too. + self._shared_release_finalizer = weakref.finalize( + self, shared_store.release_deferred, cache_key, canonical + ) + self._shared_release_finalizer.atexit = False except Exception: - # The re-point failed after acquiring a reference; release it so the shared - # entry's refcount isn't leaked (this wrapper will never enter the cache). + # The re-point or finalizer registration failed after acquiring a reference; + # release it so the shared entry's refcount isn't leaked (this wrapper will + # never enter the cache). self.release_shared_weights() raise self._cpu_state_dict = cpu_state_dict @@ -92,6 +109,11 @@ def release_shared_weights(self) -> None: no-op. After release, the shared store frees the canonical tensors once the last device that held this key releases it. """ + if self._shared_release_finalizer is not None: + # The eviction path is releasing synchronously; the collection-time fallback must not + # release the same reference a second time. + self._shared_release_finalizer.detach() + self._shared_release_finalizer = None if self._shared_store is not None and self._shared_key is not None: self._shared_store.release(self._shared_key, self._cpu_state_dict) self._shared_store = None diff --git a/invokeai/backend/model_manager/load/model_cache/cached_model/cached_model_with_partial_load.py b/invokeai/backend/model_manager/load/model_cache/cached_model/cached_model_with_partial_load.py index cae68e24331..00be7d3d950 100644 --- a/invokeai/backend/model_manager/load/model_cache/cached_model/cached_model_with_partial_load.py +++ b/invokeai/backend/model_manager/load/model_cache/cached_model/cached_model_with_partial_load.py @@ -1,3 +1,5 @@ +import weakref + import torch from invokeai.backend.model_manager.load.model_cache.shared_cpu_weights import SharedCpuWeightsStore @@ -29,6 +31,7 @@ def __init__( # under `cache_key`; `release_shared_weights()` must be called exactly once on eviction. self._shared_store: SharedCpuWeightsStore | None = None self._shared_key: str | None = None + self._shared_release_finalizer: weakref.finalize | None = None # Assigned for real at the end of __init__; initialized here so the acquire-failure path # below can call release_shared_weights(), which reads it, before that assignment runs. self._cpu_state_dict: dict[str, torch.Tensor] | None = None @@ -69,9 +72,23 @@ def __init__( if canonical is not cpu_state_dict: self._model.load_state_dict(canonical, assign=True) cpu_state_dict = canonical + # A cache dropped without a shutdown() never routes its records through + # _delete_cache_entry, so nothing would call release_shared_weights() and the + # canonical tensors would stay resident (and counted by the RAM budget) forever. + # The finalizer must not reference `self` (its args are held strongly — that would + # make the wrapper immortal) and must not take the store's non-reentrant lock (it + # runs in GC context): release_deferred only enqueues; the store applies it on its + # next operation. release_shared_weights() detaches this on the normal eviction + # path, so the release happens exactly once either way. Registered inside this try + # so a failure here (e.g. MemoryError) releases the just-acquired reference too. + self._shared_release_finalizer = weakref.finalize( + self, shared_store.release_deferred, cache_key, canonical + ) + self._shared_release_finalizer.atexit = False except Exception: - # The re-point failed after acquiring a reference; release it so the shared entry's - # refcount isn't leaked (this wrapper will never be inserted into the cache). + # The re-point or finalizer registration failed after acquiring a reference; + # release it so the shared entry's refcount isn't leaked (this wrapper will never + # be inserted into the cache). self.release_shared_weights() raise @@ -175,6 +192,11 @@ def release_shared_weights(self) -> None: no-op. After release, the shared store frees the canonical tensors once the last device that held this key releases it. """ + if self._shared_release_finalizer is not None: + # The eviction path is releasing synchronously; the collection-time fallback must not + # release the same reference a second time. + self._shared_release_finalizer.detach() + self._shared_release_finalizer = None if self._shared_store is not None and self._shared_key is not None: self._shared_store.release(self._shared_key, self._cpu_state_dict) self._shared_store = None 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 8f8296b8674..4f984ae4109 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 @@ -57,6 +57,22 @@ _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. + """ + + cache_entry: CacheRecord + held_first_use: bool + hold_epoch: int + + 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. @@ -67,6 +83,11 @@ 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. """ while True: work = work_queue.get() @@ -78,20 +99,18 @@ def _run_deferred_work(cache_ref: "weakref.ReferenceType[ModelCache]", work_queu 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) + assert isinstance(work, _AbandonedHolderRelease) + cache._release_abandoned_holder(work.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 be a CacheRecord, which transitively holds its model's - # CPU weights — and _release_first_use_grace's release hook can evict that very record, + # 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. @@ -534,14 +553,71 @@ 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 + # If the worker died before this shutdown and no admission has revived it, the holds it + # stranded would make the sweep below stale-retain their records forever: the wrappers' + # finalizer releases were (or will be) dropped by the dead-thread dispatch check, no + # unlock() is coming for a never-locked holder, and after shutdown no put() is guaranteed + # to run the usual dead-worker recovery. Clear them now so the sweep can evict the + # records; a holder that does still lock falls back to the tolerated issue-7513 path. Its + # abandoned put()-grace counterpart is cleared for the same reason: with the worker dead, + # a wrapper's grace release can no longer arrive either. (For a loader still inside the + # put()->get() gap this can turn the grace into a failed load — its get() raises rather + # than falling back — but that requires the worker's abnormal death AND shutdown() inside + # that gap, and the alternative is retaining the record forever if the loader instead + # abandoned it. The synchronous paths have always accepted the same trade: make_room and + # drop_model ignore the grace outright.) + if self._deferred_work_thread is not None and not self._deferred_work_thread.is_alive(): + self._clear_stranded_first_use_holds() + for cache_entry in self._cached_models.values(): + cache_entry.awaiting_first_use = False + # 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 inside a first-use window — the put()->lock() + # admission grace, or a LoadedModel wrapper obtained from get() and not yet entered — 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 get() 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. + 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 @@ -576,8 +652,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 @@ -623,17 +703,24 @@ 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. put() after shutdown() (Invoker.stop() stops model_manager before + # session_processor, so an in-flight generation can land here) is NOT such a state: the + # worker outlives shutdown() precisely so these releases keep flowing. worker_running = self._deferred_work_thread is not None and self._deferred_work_thread.is_alive() cache_record = CacheRecord( key=key, cached_model=wrapped_model, awaiting_first_use=not prefetch and worker_running ) + # 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 self._shutdown_event.is_set(): + 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 @@ -697,8 +784,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 @@ -717,13 +806,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(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. @@ -735,14 +830,15 @@ 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 + if self._deferred_work_thread is not None: + # The previous worker died unexpectedly; recover the holds it stranded. + self._clear_stranded_first_use_holds() thread = threading.Thread( target=_run_deferred_work, args=(weakref.ref(self), self._deferred_work_queue), @@ -779,29 +875,97 @@ def _dispatch_deferred(self, work: object) -> None: 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. + 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 _ensure_deferred_worker zeros exactly + those holds when it starts the replacement, 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. """ - 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 just-constructed LoadedModel wrapper is between get() and its + first lock. Returns the hold's epoch when armed, None when it could not be. + + 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) -> None: + """Release a register_first_use_hold() hold whose wrapper reached its first lock.""" + if cache_entry.first_use_holds > 0 and cache_entry.first_use_holds_epoch == hold_epoch: + cache_entry.first_use_holds -= 1 + + 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. + """ + for entry in self._cached_models.values(): + if entry.first_use_holds > 0: + self._logger.warning( + f"Dropping {entry.first_use_holds} first-use hold(s) on cache entry {entry.key}: the " + "deferred-work thread died, so their releases may have been lost." + ) + entry.first_use_holds = 0 + entry.first_use_holds_epoch += 1 + + @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]: @@ -957,7 +1121,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: @@ -1390,7 +1569,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." @@ -1476,7 +1661,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( @@ -1549,7 +1734,7 @@ def _reconcile_budget_if_pending(self, blocking: bool = True) -> None: 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( @@ -1571,31 +1756,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 @@ -1609,7 +1802,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/invokeai/backend/model_manager/load/model_cache/ram_budget.py b/invokeai/backend/model_manager/load/model_cache/ram_budget.py index fdc194ebe2a..55b3c3d2603 100644 --- a/invokeai/backend/model_manager/load/model_cache/ram_budget.py +++ b/invokeai/backend/model_manager/load/model_cache/ram_budget.py @@ -111,6 +111,11 @@ def remove_non_shared(self, nbytes: int, cache: Optional["ModelCache"] = None) - def total_in_use(self) -> int: """The true total RAM used by the model caches: shared weights (counted once) + non-shared.""" + # The store read MUST stay outside self._lock. The store's deferred-release drain runs + # under the store lock and allocates, so a cyclic GC can fire there and run + # _on_cache_collected, which takes THIS lock (store → budget on one thread). If any thread + # held the budget lock while calling into the store (budget → store), the two orders would + # deadlock against each other. shared = self._store.total_bytes_in_use() if self._store is not None else 0 with self._lock: non_shared = self._non_shared_bytes diff --git a/invokeai/backend/model_manager/load/model_cache/shared_cpu_weights.py b/invokeai/backend/model_manager/load/model_cache/shared_cpu_weights.py index 3e1fc9ad512..2992e1adfb5 100644 --- a/invokeai/backend/model_manager/load/model_cache/shared_cpu_weights.py +++ b/invokeai/backend/model_manager/load/model_cache/shared_cpu_weights.py @@ -1,3 +1,4 @@ +import queue import threading from dataclasses import dataclass, field @@ -50,6 +51,14 @@ class SharedCpuWeightsStore: def __init__(self) -> None: self._lock = threading.Lock() + # Releases posted from GC context (a cached-model wrapper's weakref.finalize when its cache + # was dropped without shutdown()/clear()). A finalizer must not take self._lock: it can run + # at any allocation point on any thread — including inside acquire()'s critical section, + # where taking this non-reentrant lock again would self-deadlock the process (see + # ModelCache.release_first_use_grace for the same constraint). SimpleQueue.put is + # lock-free/reentrant, so finalizers only enqueue; every public method drains the queue + # under the lock before doing its own work. + self._deferred_releases: queue.SimpleQueue[tuple[str, dict[str, torch.Tensor] | None]] = queue.SimpleQueue() self._entries: dict[str, _SharedWeightsEntry] = {} # Entries forgotten by `invalidate()` while still referenced by live cached models (e.g. a # locked, stale-marked cache entry mid-generation). They can no longer be acquired or peeked, @@ -76,6 +85,7 @@ def acquire(self, key: str, state_dict: dict[str, torch.Tensor]) -> dict[str, to re-pointing its module at these tensors and dropping the `state_dict` it passed in. """ with self._lock: + self._drain_deferred_locked() entry = self._entries.get(key) if entry is None: entry = _SharedWeightsEntry( @@ -95,6 +105,7 @@ def peek(self, key: str) -> dict[str, torch.Tensor] | None: itself increment the count. """ with self._lock: + self._drain_deferred_locked() entry = self._entries.get(key) return entry.state_dict if entry is not None else None @@ -102,6 +113,7 @@ def set_shell(self, key: str, shell: object) -> None: """Register the empty (meta-weight) structural clone for `key`, if an entry exists and none is set yet. A no-op when the key has no canonical entry (e.g. keep_ram_copy disabled).""" with self._lock: + self._drain_deferred_locked() entry = self._entries.get(key) if entry is not None and entry.shell is None: entry.shell = shell @@ -109,6 +121,7 @@ def set_shell(self, key: str, shell: object) -> None: def get_shell(self, key: str) -> object | None: """Return the registered meta-weight shell for `key`, or None if absent.""" with self._lock: + self._drain_deferred_locked() entry = self._entries.get(key) return entry.shell if entry is not None else None @@ -127,21 +140,47 @@ def release(self, key: str, state_dict: dict[str, torch.Tensor] | None = None) - lets go. """ with self._lock: - entry = self._entries.get(key) - if entry is not None and (state_dict is None or entry.state_dict is state_dict): - entry.refcount -= 1 - if entry.refcount <= 0: - del self._entries[key] + self._drain_deferred_locked() + self._release_locked(key, state_dict) + + def release_deferred(self, key: str, state_dict: dict[str, torch.Tensor] | None = None) -> None: + """Post a release to be applied by the next store operation, WITHOUT taking the store lock. + + This is the only release entry point that is safe from GC context (weakref.finalize + callbacks, __del__): it only enqueues. A finalizer can fire at any allocation point on any + thread — including while that same thread holds self._lock inside acquire() — so taking the + non-reentrant lock here could self-deadlock the process. Every public method drains the + queue under the lock, so the released bytes disappear from the accounting no later than the + next store operation (in particular, the next `total_bytes_in_use()` / budget query). + """ + self._deferred_releases.put((key, state_dict)) + + def _drain_deferred_locked(self) -> None: + """Apply all pending deferred releases. Caller must hold self._lock.""" + while True: + try: + key, state_dict = self._deferred_releases.get_nowait() + except queue.Empty: return - # Not the live canonical for `key` — it may be a retired (invalidated) entry whose - # tensors are still being counted against the RAM budget. - if state_dict is not None: - for i, retired in enumerate(self._retired): - if retired.state_dict is state_dict: - retired.refcount -= 1 - if retired.refcount <= 0: - del self._retired[i] - return + self._release_locked(key, state_dict) + + def _release_locked(self, key: str, state_dict: dict[str, torch.Tensor] | None) -> None: + """The body of release(). Caller must hold self._lock.""" + entry = self._entries.get(key) + if entry is not None and (state_dict is None or entry.state_dict is state_dict): + entry.refcount -= 1 + if entry.refcount <= 0: + del self._entries[key] + return + # Not the live canonical for `key` — it may be a retired (invalidated) entry whose + # tensors are still being counted against the RAM budget. + if state_dict is not None: + for i, retired in enumerate(self._retired): + if retired.state_dict is state_dict: + retired.refcount -= 1 + if retired.refcount <= 0: + del self._retired[i] + return def invalidate(self, model_key: str) -> int: """Forget the canonical entries (and shells) for `model_key` and all of its submodels, so no @@ -158,6 +197,7 @@ def invalidate(self, model_key: str) -> int: """ prefix = f"{model_key}:" with self._lock: + self._drain_deferred_locked() doomed = [key for key in self._entries if key == model_key or key.startswith(prefix)] for key in doomed: entry = self._entries.pop(key) @@ -169,11 +209,13 @@ def invalidate(self, model_key: str) -> int: def __contains__(self, key: str) -> bool: with self._lock: + self._drain_deferred_locked() return key in self._entries def refcount(self, key: str) -> int: """Return the current refcount for `key`, or 0 if not present.""" with self._lock: + self._drain_deferred_locked() entry = self._entries.get(key) return entry.refcount if entry is not None else 0 @@ -185,6 +227,7 @@ def total_bytes_in_use(self) -> int: it — i.e. the true RAM footprint of cached weights, not the per-device double-count. """ with self._lock: + self._drain_deferred_locked() return sum(entry.total_bytes for entry in self._entries.values()) + sum( entry.total_bytes for entry in self._retired ) @@ -192,10 +235,12 @@ def total_bytes_in_use(self) -> int: def retired_bytes(self) -> int: """Return the total size (in bytes) of retired (invalidated but still referenced) entries.""" with self._lock: + self._drain_deferred_locked() return sum(entry.total_bytes for entry in self._retired) def keys(self) -> list[str]: with self._lock: + self._drain_deferred_locked() return list(self._entries.keys()) diff --git a/tests/backend/model_manager/load/model_cache/cached_model/test_cached_model_shared_weights.py b/tests/backend/model_manager/load/model_cache/cached_model/test_cached_model_shared_weights.py index 79a34ff0d96..cfd9a8da501 100644 --- a/tests/backend/model_manager/load/model_cache/cached_model/test_cached_model_shared_weights.py +++ b/tests/backend/model_manager/load/model_cache/cached_model/test_cached_model_shared_weights.py @@ -113,9 +113,10 @@ def load_state_dict(self, *args, **kwargs): # type: ignore[override] def test_acquire_is_released_if_repoint_fails(): - # First device registers the canonical weights (refcount 1). + # First device registers the canonical weights (refcount 1). The wrapper must stay bound: an + # abandoned wrapper's collection-time finalizer releases its reference (by design). store = SharedCpuWeightsStore() - CachedModelWithPartialLoad(DummyModule(), CPU, keep_ram_copy=True, shared_store=store, cache_key="m") + first = CachedModelWithPartialLoad(DummyModule(), CPU, keep_ram_copy=True, shared_store=store, cache_key="m") assert store.refcount("m") == 1 # Second device adopts the canonical copy, but its re-point throws. The just-acquired reference @@ -124,3 +125,4 @@ def test_acquire_is_released_if_repoint_fails(): CachedModelWithPartialLoad(_RepointFailsModule(), CPU, keep_ram_copy=True, shared_store=store, cache_key="m") assert store.refcount("m") == 1 # back to just the first device, not leaked at 2 + assert first.uses_shared_weights # keep the first wrapper alive through the assertions above 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 d57e1970d0a..4b737b4c5f9 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 @@ -13,6 +13,7 @@ from unittest.mock import MagicMock, patch import pytest +import torch from invokeai.backend.model_manager.load.load_base import LoadedModelWithoutConfig from invokeai.backend.model_manager.load.model_cache import model_cache as model_cache_module @@ -91,6 +92,600 @@ def test_shared_model_counts_once_in_global_budget(mock_logger): cache_b.shutdown() +def _collect_until(predicate, attempts: int = 5) -> bool: + """gc.collect() until predicate() holds — finalizer chains can need more than one pass.""" + for _ in range(attempts): + gc.collect() + if predicate(): + return True + return predicate() + + +def test_shutdown_releases_shared_weights_synchronously(mock_logger): + """shutdown() must release its resident records' shared references itself: the finalizer + fallback only enqueues, and at teardown there may be no later store operation to drain the + queue — so relying on collection would leave the canonical tensors pinned indefinitely.""" + 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") + assert store.refcount("m") == 1 + + cache.shutdown() + # No gc, no further store activity needed: the release was synchronous. + assert store._entries.get("m") is None + 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_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 + + +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 + cache is collected.""" + store = SharedCpuWeightsStore() + budget = RamBudget(max_bytes=10**12, shared_store=store) + cache = _make_cache(store, budget, mock_logger) # keep_ram_copy=True -> shared weights + cache.put("m", DummyModule()) + _use_and_release(cache, "m") + assert store.refcount("m") == 1 + assert budget.total_in_use() == S + + cache_ref = weakref.ref(cache) + del cache + assert _collect_until(lambda: cache_ref() is None) + assert store.refcount("m") == 0 + assert store.total_bytes_in_use() == 0 + assert budget.total_in_use() == 0 + + +def test_collection_release_is_deferred_not_taken_under_the_store_lock(mock_logger): + """The collection-time release must only ENQUEUE: it runs in GC context, where taking the + store's non-reentrant lock (e.g. while another frame on the same thread is inside acquire()) + would self-deadlock the process. The queue is drained by the next store operation.""" + 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") + + cache_ref = weakref.ref(cache) + del cache + assert _collect_until(lambda: cache_ref() is None) + # The finalizer has fired, but it must not have touched the entries directly: the refcount is + # still 1 when read without the public (draining) API, and the release sits in the queue. + assert store._entries["m"].refcount == 1 + assert store._deferred_releases.qsize() == 1 + # The next public operation applies it. + assert store.refcount("m") == 0 + assert store._deferred_releases.qsize() == 0 + + +def test_normal_eviction_and_collection_release_exactly_once(mock_logger): + """An entry evicted through _delete_cache_entry (which calls release_shared_weights) must not + be released AGAIN when its wrapper is later collected — the second device's reference would be + freed out from under it.""" + 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) + cache_a.put("m", DummyModule()) + cache_b.put("m", DummyModule()) + _use_and_release(cache_a, "m") + _use_and_release(cache_b, "m") + assert store.refcount("m") == 2 + + # Normal eviction on cache_a releases its reference synchronously (and detaches the fallback). + assert cache_a.evict_unlocked_for_peer(lambda: False) == 1 + assert store.refcount("m") == 1 + + # Collecting cache_a afterwards must not decrement again on cache_b's behalf. + ref_a = weakref.ref(cache_a) + del cache_a + assert _collect_until(lambda: ref_a() is None) + assert store.refcount("m") == 1 + assert budget.total_in_use() == S + + ref_b = weakref.ref(cache_b) + del cache_b + assert _collect_until(lambda: ref_b() is None) + assert store.refcount("m") == 0 + assert budget.total_in_use() == 0 + + +def test_dropped_cache_releases_a_retired_shared_entry(mock_logger): + """invalidate() moves a still-referenced entry to the retired list, matched later by state-dict + identity. A holder that is collected (rather than evicted) must still free the retired entry's + accounting via the deferred release.""" + 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") + assert store.invalidate("m") == 1 + assert store.retired_bytes() == S + + cache_ref = weakref.ref(cache) + del cache + assert _collect_until(lambda: cache_ref() is None) + assert store.retired_bytes() == 0 + assert store.total_bytes_in_use() == 0 + assert budget.total_in_use() == 0 + + +def test_collected_partial_load_wrapper_releases_shared_weights(mock_logger): + """CachedModelWithPartialLoad (the partial-loading wrapper) has the same collection-time + release as CachedModelOnlyFullLoad.""" + from invokeai.backend.model_manager.load.model_cache.cached_model.cached_model_with_partial_load import ( + CachedModelWithPartialLoad, + ) + + store = SharedCpuWeightsStore() + wrapped = CachedModelWithPartialLoad( + model=DummyModule(), + compute_device=torch.device("cpu"), + keep_ram_copy=True, + shared_store=store, + cache_key="m", + ) + assert store.refcount("m") == 1 + + wrapped_ref = weakref.ref(wrapped) + del wrapped + assert _collect_until(lambda: wrapped_ref() is None) + assert store.refcount("m") == 0 + assert store.total_bytes_in_use() == 0 + + def test_non_shared_model_counts_per_device(mock_logger): store = SharedCpuWeightsStore() budget = RamBudget(max_bytes=10**12, shared_store=store) @@ -149,9 +744,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. @@ -955,16 +1552,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"]) @@ -1356,8 +1967,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() @@ -1457,33 +2069,49 @@ 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 + + # 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