From f959cefb6a2ddf53489ae93f68fec673abc66761 Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Wed, 29 Jul 2026 18:03:56 -0400 Subject: [PATCH 1/9] fix(model cache): release shared weights when a cache goes away MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Nothing released a cache's SharedCpuWeightsStore references except _delete_cache_entry(): shutdown() left every resident record's refcount held, and a cache dropped without shutdown() (test teardown; any future wiring that rebuilds caches at runtime) stranded the canonical tensors and their accounting forever. Today's production wiring tears the store down together with its caches, so the live exposure is cross-test pollution of the process-global store and RAM pinned past ModelManagerService.stop() — but the refcount invariant ('every acquire is paired with exactly one release') was simply not upheld, and this makes it self-healing before any wiring change turns it into a real peer-accounting bug. Two mechanisms, for the two ways a cache goes away: - shutdown() now releases its resident records' shared references synchronously — it runs in a normal thread context, so the direct (locking) release is safe there, and teardown does not depend on a later store operation happening. - Each wrapper registers a weakref.finalize fallback for the dropped-without-shutdown case. The finalizer runs in GC context, where taking the store's non-reentrant lock could self-deadlock (a collection can fire inside acquire()'s critical section on the same thread — the rule ModelCache.release_first_use_grace documents), so it only ENQUEUES into a SimpleQueue; every public store method drains the queue under the lock. The finalizer is registered inside the acquire's try (a registration failure must release too), its args carry the key and canonical dict rather than the wrapper (finalize holds args strongly — referencing self would make the wrapper immortal), and release_shared_weights() detaches it before releasing synchronously so eviction-then-collection releases exactly once. The state-dict identity keeps releases correct across invalidate()'s retired entries. RamBudget.total_in_use() now documents why its store read must stay outside the budget lock: the drain allocates under the store lock, so GC can run _on_cache_collected (store→budget) there, and a budget→store order anywhere would complete the deadlock cycle. Six regression tests, verified to fail before the fix, covering: shutdown releases synchronously with an empty queue; collection returns refcount/bytes/budget to zero; the collection-time release is enqueue-only (never applied inline by GC); eviction + collection release exactly once across two caches; a retired (invalidated) entry is freed by a collected holder; and the partial-load wrapper behaves like the full-load one. One existing test relied on an abandoned wrapper leaking its reference and now binds it. Co-Authored-By: Claude Fable 5 --- .../cached_model_only_full_load.py | 26 +++- .../cached_model_with_partial_load.py | 26 +++- .../load/model_cache/model_cache.py | 7 + .../load/model_cache/ram_budget.py | 5 + .../load/model_cache/shared_cpu_weights.py | 73 +++++++-- .../test_cached_model_shared_weights.py | 6 +- .../test_model_cache_ram_budget.py | 145 ++++++++++++++++++ 7 files changed, 268 insertions(+), 20 deletions(-) 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 3950ed22eab..62be14cc023 100644 --- a/invokeai/backend/model_manager/load/model_cache/model_cache.py +++ b/invokeai/backend/model_manager/load/model_cache/model_cache.py @@ -542,6 +542,13 @@ def shutdown(self) -> None: 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() @synchronized @record_activity 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 656a10be534..b4f465edaa8 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,150 @@ 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_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) From 3c33a92f3e544f4b21315ba5702ac136321e1b48 Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Wed, 12 Aug 2026 21:41:27 -0400 Subject: [PATCH 2/9] fix(model cache): evict records at shutdown() instead of only releasing shared weights MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit shutdown() released the resident records' shared-store references while retaining the records themselves, so the accounting stopped describing reality: - The store (and RamBudget) reported zero for bytes whose tensors the retained wrappers still held. - A post-shutdown load of the same key on a peer cache registered a duplicate canonical alongside the still-resident released copy. - A post-shutdown eviction of a released record (put() after shutdown() is reachable: Invoker.stop() stops the model manager before the session processor) read uses_shared_weights as already-False and debited the non-shared budget for bytes that were admitted as shared. shutdown() now routes idle records through _delete_cache_entry(), which releases shared ownership and budget accounting together, exactly once. Records still in use — locked by an in-flight generation or inside the put()->lock() admission window — keep their references and are marked stale; unlock() evicts them through the existing stale path when the generation lets go, so the accounting stays truthful at every point. All five regression tests verified to fail against the previous shutdown() behavior. Follow-on to #9403, addressing JPPhoto's review comment there. Co-Authored-By: Claude Fable 5 --- .../load/model_cache/model_cache.py | 35 +++++-- .../test_model_cache_ram_budget.py | 94 +++++++++++++++++++ .../model_cache/test_model_cache_timeout.py | 20 ++-- 3 files changed, 135 insertions(+), 14 deletions(-) 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 45757ca9366..d71c5b8fb31 100644 --- a/invokeai/backend/model_manager/load/model_cache/model_cache.py +++ b/invokeai/backend/model_manager/load/model_cache/model_cache.py @@ -534,7 +534,7 @@ 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() @@ -542,13 +542,32 @@ def shutdown(self) -> None: 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() + # 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 the put()->lock() admission window + # (awaiting_first_use) are marked stale instead, and unlock() evicts them through this + # same path once the generation lets go. A record never unlocked 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.awaiting_first_use: + cache_entry.is_stale = True + else: + self._delete_cache_entry(cache_entry) @synchronized @record_activity 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 3efe2d2fe12..27d7fac04ae 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,100 @@ 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_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_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 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 From ad6c23268fc36af3bcf3aeafca45a0bb20a882e7 Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Wed, 12 Aug 2026 21:52:42 -0400 Subject: [PATCH 3/9] fix(model cache): match records by identity in stale eviction and _delete_cache_entry MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Surfaced by adversarial review of the shutdown() change: a stale-marked record can be detached while still locked (the VRAM-move error paths call _delete_cache_entry on a locked record) and its key re-admitted before the record's last unlock(). The stale-eviction path matched by key only, so it popped the NEW record — detaching it from the cache and all accounting — and, the old record's shared release having already happened, read uses_shared_weights as False and debited the non-shared budget for bytes that were admitted as shared. The hazard predates the shutdown() change (drop_model() sets the same flag), but shutdown() now arms stale marks at every server stop that overlaps in-flight work, so close it here: _delete_cache_entry() and unlock()'s stale eviction act only when the record passed in IS the record currently held under its key; a delete of a detached record is a full no-op. Regression test verified to fail against the key-only matching. Co-Authored-By: Claude Fable 5 --- .../load/model_cache/model_cache.py | 44 ++++++++++++------- .../test_model_cache_ram_budget.py | 29 ++++++++++++ 2 files changed, 58 insertions(+), 15 deletions(-) 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 d71c5b8fb31..f6823a45e6c 100644 --- a/invokeai/backend/model_manager/load/model_cache/model_cache.py +++ b/invokeai/backend/model_manager/load/model_cache/model_cache.py @@ -983,7 +983,14 @@ 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. + if ( + cache_entry.is_stale + and not cache_entry.is_locked + 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: @@ -1597,22 +1604,29 @@ 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: 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 27d7fac04ae..8a5d6f1d2f8 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 @@ -183,6 +183,35 @@ def test_shutdown_retains_admission_window_records(mock_logger): 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 From 5fda854c4780232ae94d02b35db4bab2eecab670 Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Sun, 16 Aug 2026 22:13:03 -0400 Subject: [PATCH 4/9] fix(model cache): track get()->lock() holders through shutdown and abandonment MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two defects found in review of the shutdown eviction change (JPPhoto, 2026-08-13): 1. shutdown() racing the gap between get() and the LoadedModel's first lock evicted the warm record out from under its holder: the holder locked a detached record whose shared-store ownership had just been released, so a peer's reload of the same key minted a duplicate canonical copy while the budget counted one. 2. A record retained by the shutdown sweep for a never-locked holder could never be evicted if that holder was simply dropped: the abandonment finalizer's deferred work was discarded post-shutdown (and the worker was stopped), pinning the record, its shared-store refcount and its budget bytes for the life of the process. The fix tracks every wrapper's get()->lock() window with a per-record hold count (CacheRecord.first_use_holds), armed in LoadedModelWithoutConfig's constructor and released exactly once per wrapper — on its first lock, or by its weakref finalizer if it is dropped un-entered. Held records are treated like locked ones by every eviction path (shutdown, budget reconcile, peer-requested eviction, make_room, drop_model, unlock's stale eviction); stale-marked records whose last holder is abandoned are evicted by the deferred worker, which now outlives shutdown() for exactly that purpose (it already exits via the cache-collection finalizer). Holds are only granted while a worker is alive to carry the finalizer's release, and a worker death zeroes surviving holds at the next start so no record can stay shielded with nothing left to unshield it. Admissions landing after shutdown() are marked stale at birth so their final release evicts them too. Co-Authored-By: Claude Fable 5 --- .../backend/model_manager/load/load_base.py | 43 ++- .../load/model_cache/cache_record.py | 35 ++- .../load/model_cache/model_cache.py | 230 ++++++++++---- .../test_model_cache_ram_budget.py | 280 ++++++++++++++++-- 4 files changed, 496 insertions(+), 92 deletions(-) diff --git a/invokeai/backend/model_manager/load/load_base.py b/invokeai/backend/model_manager/load/load_base.py index 7225fd1402f..0fd986769f1 100644 --- a/invokeai/backend/model_manager/load/load_base.py +++ b/invokeai/backend/model_manager/load/load_base.py @@ -56,15 +56,37 @@ 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 get() and this wrapper's 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 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). + release_grace = getattr(cache, "release_first_use_grace", None) + register_hold = getattr(cache, "register_first_use_hold", None) + self._holds_first_use = ( + bool(register_hold(cache_record)) if register_hold is not None and release_grace is not None else False ) - if self._first_use_finalizer is not None: + self._first_use_finalizer = None + if release_grace is not None and (self._holds_first_use or cache_record.awaiting_first_use): + self._first_use_finalizer = finalize(self, release_grace, cache_record, self._holds_first_use) 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._holds_first_use: + self._holds_first_use = False + release_hold = getattr(self._cache, "release_first_use_hold", None) + if release_hold is not None: + release_hold(self._cache_record) + 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 +94,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 +117,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 +133,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..6bf6532b2bd 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,24 @@ 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._ensure_deferred_worker() zeroing the counts after the deferred worker — the + # thread that carries finalizer-initiated releases — is found dead, since holds released into + # a dead worker are dropped and would otherwise shield the record forever. + first_use_holds: int = 0 def lock(self) -> None: """Lock this record.""" @@ -51,3 +62,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/model_cache.py b/invokeai/backend/model_manager/load/model_cache/model_cache.py index f6823a45e6c..498a4e246f6 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,19 @@ _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. + """ + + cache_entry: CacheRecord + held_first_use: bool + + 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 +80,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 +96,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) 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. @@ -538,7 +554,14 @@ def shutdown(self) -> None: 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 @@ -558,13 +581,20 @@ def shutdown(self) -> None: # # 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 the put()->lock() admission window - # (awaiting_first_use) are marked stale instead, and unlock() evicts them through this - # same path once the generation lets go. A record never unlocked keeps its bytes — and - # its accounting — until process exit, which is the truthful description of a model that - # really is still resident. + # 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.awaiting_first_use: + if cache_entry.is_locked or cache_entry.in_first_use_window: cache_entry.is_stale = True else: self._delete_cache_entry(cache_entry) @@ -602,8 +632,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 @@ -649,17 +683,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 @@ -723,8 +764,8 @@ 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) -> 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 @@ -743,13 +784,17 @@ 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. """ - # 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)) def _ensure_deferred_worker(self) -> None: """Start the background worker if it is not currently running. Caller must hold the lock. @@ -761,14 +806,27 @@ 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. Any first-use hold armed while it was alive + # may have had its finalizer-initiated release dispatched into the dead thread and + # dropped (finalizers fire once, so a dropped release is never retried) — such a hold + # would shield its record from every eviction path forever. Zero the counts: 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 thread = threading.Thread( target=_run_deferred_work, args=(weakref.ref(self), self._deferred_work_queue), @@ -805,29 +863,67 @@ 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) -> bool: + """Shield a record while a just-constructed LoadedModel wrapper is between get() and its + first lock. Returns whether the hold was armed. + + 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. Because that finalizer 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). + """ + self._ensure_deferred_worker() + if self._deferred_work_thread is None or not self._deferred_work_thread.is_alive(): + return False + cache_entry.first_use_holds += 1 + return True + + @synchronized + def release_first_use_hold(self, cache_entry: CacheRecord) -> None: + """Release a register_first_use_hold() hold whose wrapper reached its first lock.""" + if cache_entry.first_use_holds > 0: + cache_entry.first_use_holds -= 1 + + @synchronized + def _release_abandoned_holder(self, cache_entry: CacheRecord, held_first_use: bool) -> 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: + 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]: @@ -986,9 +1082,17 @@ def unlock(self, cache_entry: CacheRecord) -> None: # 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() @@ -1423,7 +1527,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." @@ -1509,7 +1619,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( @@ -1582,7 +1692,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( @@ -1633,9 +1743,10 @@ 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 @@ -1649,7 +1760,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 8a5d6f1d2f8..63bd4541183 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 @@ -241,6 +241,219 @@ def test_no_duplicate_canonical_when_peer_reloads_after_shutdown(mock_logger): 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_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 @@ -417,9 +630,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. @@ -1223,16 +1438,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"]) @@ -1624,8 +1853,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() @@ -1725,33 +1955,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() From d000e0aef1e586d313b8e704f7abbec8455093bd Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Sun, 16 Aug 2026 22:41:56 -0400 Subject: [PATCH 5/9] fix(model cache): epoch-guard hold releases and recover stranded holds at shutdown MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Hardening from adversarial review of the first-use-hold mechanism: - Hold releases (the wrapper's first-lock release and the abandonment finalizer's deferred release) now quote the epoch the hold was armed under, and dead-worker recovery bumps the record's epoch when it zeroes stranded holds. Without this, a surviving wrapper's late release — or a release enqueued before the 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. - shutdown() now runs the dead-worker hold recovery itself (and clears the put()-grace flags in the same situation): a hold whose abandonment release was dropped by the dead-thread dispatch check has no other releaser, and after shutdown no put() is guaranteed to run the usual next-start recovery — the sweep would stale-retain the record, its shared-store refcount and its budget bytes for the life of the process. - register_first_use_hold() declines to arm on a record that is no longer the occupant under its key: an eviction already won the race against the wrapper's construction and a hold on a detached record shields nothing. Co-Authored-By: Claude Fable 5 --- .../backend/model_manager/load/load_base.py | 39 ++++-- .../load/model_cache/cache_record.py | 13 +- .../load/model_cache/model_cache.py | 102 +++++++++++----- .../test_model_cache_ram_budget.py | 114 ++++++++++++++++++ 4 files changed, 222 insertions(+), 46 deletions(-) diff --git a/invokeai/backend/model_manager/load/load_base.py b/invokeai/backend/model_manager/load/load_base.py index 0fd986769f1..d4094b5557c 100644 --- a/invokeai/backend/model_manager/load/load_base.py +++ b/invokeai/backend/model_manager/load/load_base.py @@ -56,22 +56,34 @@ class LoadedModelWithoutConfig: def __init__(self, cache_record: CacheRecord, cache: ModelCache): self._cache_record = cache_record self._cache = cache - # Shield the record for the window between get() and this wrapper's 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 hold is released exactly once: on the + # 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). + # 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._holds_first_use = ( - bool(register_hold(cache_record)) if register_hold is not None and release_grace is not None else False + 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 ) self._first_use_finalizer = None - if release_grace is not None and (self._holds_first_use or cache_record.awaiting_first_use): - self._first_use_finalizer = finalize(self, release_grace, cache_record, self._holds_first_use) + 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: @@ -81,11 +93,12 @@ def _end_first_use_window(self) -> None: if self._first_use_finalizer is not None: self._first_use_finalizer.detach() self._first_use_finalizer = None - if self._holds_first_use: - self._holds_first_use = False + 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) + release_hold(self._cache_record, hold_epoch) def __enter__(self) -> AnyModel: # Hold the MODEL_LOAD_LOCK read lock across the VRAM load (lock() runs 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 6bf6532b2bd..a921bca96d2 100644 --- a/invokeai/backend/model_manager/load/model_cache/cache_record.py +++ b/invokeai/backend/model_manager/load/model_cache/cache_record.py @@ -44,10 +44,17 @@ class CacheRecord: # 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._ensure_deferred_worker() zeroing the counts after the deferred worker — the - # thread that carries finalizer-initiated releases — is found dead, since holds released into - # a dead worker are dropped and would otherwise shield the record forever. + # 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.""" 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 498a4e246f6..4f984ae4109 100644 --- a/invokeai/backend/model_manager/load/model_cache/model_cache.py +++ b/invokeai/backend/model_manager/load/model_cache/model_cache.py @@ -64,10 +64,13 @@ class _AbandonedHolderRelease(NamedTuple): 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: @@ -100,7 +103,7 @@ def _run_deferred_work(cache_ref: "weakref.ReferenceType[ModelCache]", work_queu cache._reconcile_budget_if_pending() else: assert isinstance(work, _AbandonedHolderRelease) - cache._release_abandoned_holder(work.cache_entry, work.held_first_use) + 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") @@ -565,6 +568,23 @@ def shutdown(self) -> None: 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 @@ -764,7 +784,9 @@ 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, held_first_use: bool = False) -> None: + 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 @@ -786,7 +808,9 @@ def release_first_use_grace(self, cache_entry: CacheRecord, held_first_use: bool 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. + 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 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 @@ -794,7 +818,7 @@ def release_first_use_grace(self, cache_entry: CacheRecord, held_first_use: bool # 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(_AbandonedHolderRelease(cache_entry, held_first_use)) + 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. @@ -813,20 +837,8 @@ def _ensure_deferred_worker(self) -> None: 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. Any first-use hold armed while it was alive - # may have had its finalizer-initiated release dispatched into the dead thread and - # dropped (finalizers fire once, so a dropped release is never retried) — such a hold - # would shield its record from every eviction path forever. Zero the counts: 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 + # 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), @@ -878,34 +890,64 @@ def _dispatch_deferred(self, work: object) -> None: self._deferred_work_queue.put(work) @synchronized - def register_first_use_hold(self, cache_entry: CacheRecord) -> bool: + 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 whether the hold was armed. + 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. Because that finalizer 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). + 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 False + return None + if self._cached_models.get(cache_entry.key) is not cache_entry: + return None cache_entry.first_use_holds += 1 - return True + return cache_entry.first_use_holds_epoch @synchronized - def release_first_use_hold(self, cache_entry: CacheRecord) -> None: + 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: + 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) -> None: + 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 @@ -913,7 +955,7 @@ def _release_abandoned_holder(self, cache_entry: CacheRecord, held_first_use: bo 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: + 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 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 63bd4541183..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 @@ -454,6 +454,120 @@ def test_worker_death_zeroes_stranded_first_use_holds(mock_logger): 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 From 05ca84ecfe1e405f3cedcd3ecc41683b72ec2d34 Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Sat, 29 Aug 2026 12:36:32 -0400 Subject: [PATCH 6/9] fix(model cache): withhold the post-shutdown grace and recover from the worker's own death Two follow-ups from review. put() after shutdown() no longer arms the post-admission grace. That flag'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 stayed charged to the shared budget. Withholding it costs only the shield -- the record stays stale at birth, so its eventual release still evicts it, and a loader that does come back gets the ordinary first_use_holds shield. The deferred worker now runs stranded-shield recovery from inside its own dying frame. Previously recovery depended on something else happening first -- the next admission, or shutdown() -- and neither is guaranteed when the worker dies *after* shutdown()'s liveness check: the records the shutdown sweep retained for a live holder were left shielded by holds nothing could release. The recovery is scoped by thread identity (a replacement worker's shields are its own) and retires the worker slot before sweeping, so a concurrent admission cannot arm a shield the recovery is about to zero. It also drains the queue the dead worker left behind, whose _AbandonedHolderRelease items pin their models' CPU weights. _ensure_deferred_worker and shutdown() now share the same recovery, which also lifts orphaned admission graces and evicts whatever that leaves unshielded. Three tests, each reverted-and-confirmed-failing against the code it guards. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01GjAw9hpJe8d1GyvpdKJFGw --- .../load/model_cache/cache_record.py | 15 +- .../load/model_cache/model_cache.py | 188 +++++++++++++----- .../test_model_cache_ram_budget.py | 115 +++++++++++ 3 files changed, 265 insertions(+), 53 deletions(-) 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 a921bca96d2..24057359056 100644 --- a/invokeai/backend/model_manager/load/model_cache/cache_record.py +++ b/invokeai/backend/model_manager/load/model_cache/cache_record.py @@ -34,8 +34,11 @@ class CacheRecord: # 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. + # 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. 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 @@ -44,10 +47,10 @@ class CacheRecord: # 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. + # 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 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 16029b69270..a7771194089 100644 --- a/invokeai/backend/model_manager/load/model_cache/model_cache.py +++ b/invokeai/backend/model_manager/load/model_cache/model_cache.py @@ -89,33 +89,57 @@ def _run_deferred_work(cache_ref: "weakref.ReferenceType[ModelCache]", work_queu 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 + 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._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 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. + work = None + cache = 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 work is _DEFERRED_RECONCILE: - cache._reconcile_budget_if_pending() - else: - 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 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. - 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 @@ -596,23 +620,14 @@ def shutdown(self) -> None: 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 the worker died before this shutdown and neither its own dying recovery nor a later + # admission cleared what it stranded, the surviving shields 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, 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). 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 + 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 @@ -763,17 +778,33 @@ def put( # 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. + # 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. 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 self._shutdown_event.is_set(): + if shutting_down: cache_record.is_stale = True self._cached_models[key] = cache_record self._cache_stack.append(key) @@ -903,8 +934,9 @@ def _ensure_deferred_worker(self) -> None: 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() + # The previous worker died unexpectedly (and without running its own recovery — a fork, + # or a death raced by this call); recover the shields it stranded. + self._recover_stranded_shields() thread = threading.Thread( target=_run_deferred_work, args=(weakref.ref(self), self._deferred_work_queue), @@ -1012,6 +1044,68 @@ def _clear_stranded_first_use_holds(self) -> None: entry.first_use_holds = 0 entry.first_use_holds_epoch += 1 + def _recover_stranded_shields(self) -> None: + """Lift every shield whose release depended on a deferred worker that is no longer there, + and evict whatever that leaves unshielded. Caller must hold the cache lock. + + Both first-use shields are granted only while a worker is alive, because both can end up + needing a finalizer-carried release: a hold always does, and an admission grace does + whenever the loader abandons the model after constructing its wrapper. Once the worker is + gone those releases are dropped by _dispatch_deferred and finalizers never fire twice, so + the shields would stand for the life of the process — invisible to every asynchronous + eviction path while their bytes stay charged to the shared budget. A wrapper unshielded + here merely falls back to the tolerated issue-7513 detached path if an eviction really + does race its lock; a loader still inside the put()->get() gap can instead see its get() + raise. Both are recoverable, and a permanently shielded record is not — the same trade the + synchronous paths (make_room, drop_model) have always made against the grace. + + A record the shutdown sweep (or drop_model) marked stale and retained *because* of one of + those shields is then evicted here: with the shield gone and no lock outstanding, no + unlock() and no abandonment release is ever coming to run the usual stale eviction. + """ + self._clear_stranded_first_use_holds() + for entry in self._cached_models.values(): + entry.awaiting_first_use = False + # Drop what the dead worker left queued. The sweeps above already cover its semantics — + # the epoch bump makes every queued hold release a no-op, and a dropped reconcile is + # re-run by the next lock-release hook — but each queued _AbandonedHolderRelease holds a + # CacheRecord, and through it that model's CPU weights. Nothing else drains this queue + # unless a later admission starts a replacement worker, which after shutdown() may never + # come. (No _DEFERRED_STOP can be in flight: that is pushed by the cache-collection + # finalizer, and this cache is alive.) + while True: + try: + self._deferred_work_queue.get_nowait() + except queue.Empty: + break + 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) + + @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. + self._deferred_work_thread = None + self._recover_stranded_shields() + @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. 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 06abae555af..0130d7b348e 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 @@ -568,6 +568,121 @@ def test_shutdown_clears_holds_stranded_by_a_dead_worker(mock_logger): 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(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() + + 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 From 5ec7937bd766a650d61941abd0d3196735050348 Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Sat, 29 Aug 2026 13:06:38 -0400 Subject: [PATCH 7/9] fix(model cache): narrow the dead-worker recovery and stop it pinning records Adversarial review of the previous commit found three problems with the shared recovery it introduced. The recovery cleared the put()-set admission grace unconditionally. On a live cache that is a new failure mode, not a fix: 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 grace, which the recovery then zeroed while the loader was still between put() and get() -- a reconcile could evict the record and the loader's get() would raise IndexError. The grace only actually loses a releaser once the cache is shut down (its backstop is the next put()'s sweep, not the worker), so it is now lifted only then. The recovery also evicted stale-unshielded records from _ensure_deferred_worker, which register_first_use_hold calls before arming -- so a second wrapper's construction could detach the very record it was about to shield, releasing shared-store ownership while live wrappers still held the tensors. That is the accounting lie shutdown() itself refuses to make. The eviction moved to _evict_stale_unshielded_entries, called only from the dying worker and only on a shut-down cache, where nothing else can ever run it; it now also collects and empties the device cache the way the other abandonment path does. Keeping shutdown()'s call to pure field assignments restores its old property that the branch cannot raise before the resident-record sweep. The queue drain the previous commit added did not close the pin it targeted: _dispatch_deferred's liveness gate is unsynchronized, so a finalizer that read the worker slot just before it was retired still enqueues after the drain. The drain is gone; _AbandonedHolderRelease now holds its record weakly instead, so a stranded item pins nothing, and the worker clears the strong reference it resolves before parking on the next get(). Also moves _reconcile_budget_if_pending's lock acquisition adjacent to its try: a BaseException between the two leaked the cache RLock to an unwinding thread, blocking every other thread for the life of the process. Five tests, each reverted-and-confirmed-failing against the code it guards. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01GjAw9hpJe8d1GyvpdKJFGw --- .../load/model_cache/model_cache.py | 123 ++++++++----- .../test_model_cache_ram_budget.py | 171 +++++++++++++++++- 2 files changed, 251 insertions(+), 43 deletions(-) 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 a7771194089..28bae84ea78 100644 --- a/invokeai/backend/model_manager/load/model_cache/model_cache.py +++ b/invokeai/backend/model_manager/load/model_cache/model_cache.py @@ -67,9 +67,19 @@ class _AbandonedHolderRelease(NamedTuple): 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 the worker drains this queue, and an item enqueued as + the worker is dying is never drained at all: _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. 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 matters the record is the live occupant of its key in + _cached_models, which holds it strongly; a record that has already been dropped from there is + exactly the case _release_abandoned_holder's identity check no-ops anyway. """ - cache_entry: CacheRecord + cache_entry_ref: "weakref.ReferenceType[CacheRecord]" held_first_use: bool hold_epoch: int @@ -101,6 +111,7 @@ def _run_deferred_work(cache_ref: "weakref.ReferenceType[ModelCache]", work_queu while True: work = work_queue.get() cache = None + cache_entry = None try: if work is _DEFERRED_STOP: stopped_on_purpose = True @@ -114,7 +125,9 @@ def _run_deferred_work(cache_ref: "weakref.ReferenceType[ModelCache]", work_queu cache._reconcile_budget_if_pending() else: assert isinstance(work, _AbandonedHolderRelease) - cache._release_abandoned_holder(work.cache_entry, work.held_first_use, work.hold_epoch) + 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") @@ -125,9 +138,11 @@ def _run_deferred_work(cache_ref: "weakref.ReferenceType[ModelCache]", work_queu # 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. + # 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 @@ -915,7 +930,7 @@ def release_first_use_grace( # 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(_AbandonedHolderRelease(cache_entry, held_first_use, hold_epoch)) + 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. @@ -970,17 +985,23 @@ 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. + 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 _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. + 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 + that item is then never drained. _AbandonedHolderRelease therefore holds its record weakly, + so a stranded item pins nothing. """ thread = self._deferred_work_thread if thread is None or not thread.is_alive(): @@ -1045,43 +1066,53 @@ def _clear_stranded_first_use_holds(self) -> None: entry.first_use_holds_epoch += 1 def _recover_stranded_shields(self) -> None: - """Lift every shield whose release depended on a deferred worker that is no longer there, - and evict whatever that leaves unshielded. Caller must hold the cache lock. - - Both first-use shields are granted only while a worker is alive, because both can end up - needing a finalizer-carried release: a hold always does, and an admission grace does - whenever the loader abandons the model after constructing its wrapper. Once the worker is - gone those releases are dropped by _dispatch_deferred and finalizers never fire twice, so - the shields would stand for the life of the process — invisible to every asynchronous - eviction path while their bytes stay charged to the shared budget. A wrapper unshielded - here merely falls back to the tolerated issue-7513 detached path if an eviction really - does race its lock; a loader still inside the put()->get() gap can instead see its get() - raise. Both are recoverable, and a permanently shielded record is not — the same trade the - synchronous paths (make_room, drop_model) have always made against the grace. - - A record the shutdown sweep (or drop_model) marked stale and retained *because* of one of - those shields is then evicted here: with the shield gone and no lock outstanding, no - unlock() and no abandonment release is ever coming to run the usual stale eviction. + """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, so what is lifted here is only a grace armed before + the shutdown and stale-retained by its sweep. """ 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 - # Drop what the dead worker left queued. The sweeps above already cover its semantics — - # the epoch bump makes every queued hold release a no-op, and a dropped reconcile is - # re-run by the next lock-release hook — but each queued _AbandonedHolderRelease holds a - # CacheRecord, and through it that model's CPU weights. Nothing else drains this queue - # unless a later admission starts a replacement worker, which after shutdown() may never - # come. (No _DEFERRED_STOP can be in flight: that is pushed by the cache-collection - # finalizer, and this cache is alive.) - while True: - try: - self._deferred_work_queue.get_nowait() - except queue.Empty: - break + + 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: + gc.collect() + TorchDevice.empty_cache() @synchronized def _recover_from_dead_worker(self, worker: threading.Thread) -> None: @@ -1102,9 +1133,14 @@ def _recover_from_dead_worker(self, worker: threading.Thread) -> None: # 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. + # 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: @@ -1968,11 +2004,14 @@ 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 would + # leak the RLock, and a cache lock owned by a thread that is unwinding blocks every + # other thread for the life of the process. 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: 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 0130d7b348e..0d119ad4372 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 @@ -579,7 +579,7 @@ def _kill_worker_abnormally(cache: ModelCache, record) -> threading.Thread: 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(record, False, 0)) + 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 @@ -683,6 +683,175 @@ def test_dying_worker_recovery_leaves_a_replacement_workers_shields_alone(mock_l 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() + + 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 From 4445d2dbc611f9e33118b3cfdc54681475f4ce0a Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Sat, 29 Aug 2026 13:34:30 -0400 Subject: [PATCH 8/9] fix(model cache): key the dead-worker backstops on worker liveness, not the slot A second adversarial pass found that retiring the worker slot from inside the dying worker silently disabled both remaining recovery sites, which gated on "a dead thread still occupying the slot". The dying recovery deliberately leaves a live cache's admission grace standing -- the next put()'s sweep is still its backstop -- and hands the lift to shutdown(); with the slot already empty, shutdown() skipped it and stale-retained the record, its shared-store reference and its budget bytes for the life of the process. Both gates now key on "no live worker": shutdown() lifts when the slot is empty or dead, and the worker start recovers unconditionally (it has already returned if a worker is alive). That also makes a failed recovery retryable, which matters because the recovery was not exception-safe and had already retired the slot by the time it could raise. _clear_stranded_first_use_holds now unshields every record before reporting any of them -- a logging handler that raises is one of the ways the worker dies in the first place, and logging inline let that same handler abort the sweep partway -- and the post-eviction gc/empty_cache housekeeping, which the codebase already documents can raise from a sick CUDA context, no longer takes the eviction down with it. Also corrects two overstated claims in the weakref rationale: a stranded queue item can be drained later by a replacement worker (the queue is per-cache, not per-worker), and the hold decrement in _release_abandoned_holder runs before the identity check -- it is inert on a detached record for a different reason, which the docstring now gives. Moving _reconcile_budget_if_pending's acquire adjacent to its try narrows the RLock-leak window rather than closing it; said so. Three tests, each reverted-and-confirmed-failing against the code it guards. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01GjAw9hpJe8d1GyvpdKJFGw --- .../load/model_cache/model_cache.py | 93 ++++++++++++------- .../test_model_cache_ram_budget.py | 91 ++++++++++++++++++ 2 files changed, 153 insertions(+), 31 deletions(-) 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 28bae84ea78..765e4b4423e 100644 --- a/invokeai/backend/model_manager/load/model_cache/model_cache.py +++ b/invokeai/backend/model_manager/load/model_cache/model_cache.py @@ -68,15 +68,20 @@ class _AbandonedHolderRelease(NamedTuple): `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 the worker drains this queue, and an item enqueued as - the worker is dying is never drained at all: _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. 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 matters the record is the live occupant of its key in - _cached_models, which holds it strongly; a record that has already been dropped from there is - exactly the case _release_abandoned_holder's identity check no-ops anyway. + 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]" @@ -635,13 +640,18 @@ def shutdown(self) -> None: if self._timeout_timer is not None: self._timeout_timer.cancel() self._timeout_timer = None - # If the worker died before this shutdown and neither its own dying recovery nor a later - # admission cleared what it stranded, the surviving shields 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, 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). - if self._deferred_work_thread is not None and not self._deferred_work_thread.is_alive(): + # 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 @@ -948,10 +958,13 @@ def _ensure_deferred_worker(self) -> None: """ 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 (and without running its own recovery — a fork, - # or a death raced by this call); recover the shields it stranded. - self._recover_stranded_shields() + # 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), @@ -1000,8 +1013,11 @@ def _dispatch_deferred(self, work: object) -> None: 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 - that item is then never drained. _AbandonedHolderRelease therefore holds its record weakly, - so a stranded item pins nothing. + 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.) """ thread = self._deferred_work_thread if thread is None or not thread.is_alive(): @@ -1056,14 +1072,21 @@ def _clear_stranded_first_use_holds(self) -> None: 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: - 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." - ) + 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. @@ -1111,8 +1134,14 @@ def _evict_stale_unshielded_entries(self) -> None: self._delete_cache_entry(entry) evicted += 1 if evicted: - gc.collect() - TorchDevice.empty_cache() + 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: @@ -2005,9 +2034,11 @@ def _reconcile_budget_if_pending(self, blocking: bool = True) -> None: return self._budget_reconcile_pending.set() models_cleared = 0 - # Acquire immediately before the try: a BaseException delivered between the two would - # leak the RLock, and a cache lock owned by a thread that is unwinding blocks every - # other thread for the life of the process. + # 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. 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 0d119ad4372..03efc0ffeab 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 @@ -852,6 +852,97 @@ def park(*_args): 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 From 2269534fb4d928f7ef12a5836449b958421936bd Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Mon, 31 Aug 2026 00:32:48 -0400 Subject: [PATCH 9/9] fix(model cache): claim the first-use window at the lookup, and let a cancelled admission release itself Two findings from review. The first-use shield was armed by LoadedModel's constructor, leaving the whole stretch between the cache lookup and that constructor unshielded -- and that stretch is not a couple of instructions: the configured loader retrieves its record inside _load_and_cache and then does the shared-store shell registration and two returns before load_model wraps it. A shutdown sweep or a peer's reconcile landing there detached the record its holder was about to lock, releasing shared-store ownership while the tensors lived on, so a peer's reload minted a duplicate canonical the budget counted once. ModelCache.get_with_first_use_claim() now arms the hold in the same lock acquisition as the lookup and hands back a FirstUseClaim that owns it: the wrapper adopts the claim and releases it at its first lock, and a claim dropped without ever being adopted -- the load raised before a wrapper existed -- releases the hold by dying. shutdown() stale-retained a record carrying only the put()-set admission grace. That grace's three releasers are the loader's own get()->lock(), the abandonment finalizer of a wrapper built from the record, and the sweep at the top of the next put(); a load cancelled between its put() and its retrieval has neither of the first two, and after shutdown no further put() is guaranteed to run the third, so the record, its shared-store reference and its budget charge stood until the cache object was collected. put(claim_admission=True) now hands the loader a claim over that window too, so such a load releases its admission by dying and the shutdown sweep finds an ordinary idle record. Retiring the grace at shutdown instead -- the obvious shortcut, and what the first two drafts of this commit did -- is not safe. The flag is unowned, so a standing grace does not mean nobody is 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 wrapper whose hold a worker death zeroed. Both were evicted out from under their holder, with the duplicate-canonical accounting lie and an IndexError from a retrieval that no longer found its own model. For the same reason the admission window is not shielded by either flag but by a weak reference to the claim (CacheRecord.admission_claim_ref): nothing has to release it, so neither a worker death (which zeroes holds) nor another holder's abandonment (which clears the grace) can make a running load look finished. _recover_stranded_shields retires it once the cache is shut down, where the eviction its expiry should trigger would otherwise travel through a dead worker -- the same trade that method already makes for holds. The claim is armed only after put() has committed its accounting, so a failure to allocate it cannot leave a resident, store-owning record the budget never counted, and both hand-back guards release the hold when the object that was to carry its release cannot be built. Nine tests, each reverted-and-confirmed-failing against the code it guards. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01FCitU6EFcp76AaauNfaWTz --- .../services/model_load/model_load_default.py | 19 +- .../backend/model_manager/load/load_base.py | 99 ++-- .../model_manager/load/load_default.py | 47 +- .../load/model_cache/cache_record.py | 68 ++- .../load/model_cache/model_cache.py | 246 +++++++++- .../test_model_cache_ram_budget.py | 458 ++++++++++++++++++ 6 files changed, 859 insertions(+), 78 deletions(-) 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 da1cb305f3d..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,43 +57,76 @@ 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 - # 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. + # 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 register_hold is not None and release_grace is not None else None + 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 - 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 + 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.detach() self._first_use_finalizer = None @@ -198,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 24057359056..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, @@ -36,17 +38,23 @@ class CacheRecord: # 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. 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. + # 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 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 + # 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 @@ -58,6 +66,30 @@ class CacheRecord: # 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.""" @@ -73,14 +105,22 @@ 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 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 + """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 + 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 765e4b4423e..8a8a42935bf 100644 --- a/invokeai/backend/model_manager/load/model_cache/model_cache.py +++ b/invokeai/backend/model_manager/load/model_cache/model_cache.py @@ -89,6 +89,69 @@ class _AbandonedHolderRelease(NamedTuple): 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. @@ -669,18 +732,36 @@ def shutdown(self) -> None: # # 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. + # 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 @@ -690,8 +771,13 @@ def shutdown(self) -> None: @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: @@ -703,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 @@ -819,6 +913,16 @@ def put( # 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( @@ -860,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`. @@ -1026,8 +1157,12 @@ def _dispatch_deferred(self, work: object) -> None: @synchronized 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. + """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 @@ -1054,10 +1189,24 @@ def register_first_use_hold(self, cache_entry: CacheRecord) -> Optional[int]: 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.""" + 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 @@ -1106,14 +1255,25 @@ def _recover_stranded_shields(self) -> None: 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, so what is lifted here is only a grace armed before - the shutdown and stale-retained by its sweep. + 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. @@ -1253,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: 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 03efc0ffeab..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 @@ -284,6 +284,458 @@ def test_shutdown_retains_record_inside_get_to_lock_window(mock_logger): 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 @@ -2523,6 +2975,12 @@ def fail_worker_start(thread: threading.Thread) -> None: 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