From f959cefb6a2ddf53489ae93f68fec673abc66761 Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Wed, 29 Jul 2026 18:03:56 -0400 Subject: [PATCH 1/5] 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/5] 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/5] 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/5] 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/5] 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