diff --git a/invokeai/backend/model_manager/load/model_cache/cached_model/cached_model_only_full_load.py b/invokeai/backend/model_manager/load/model_cache/cached_model/cached_model_only_full_load.py index ccf40654575..226f65f6998 100644 --- a/invokeai/backend/model_manager/load/model_cache/cached_model/cached_model_only_full_load.py +++ b/invokeai/backend/model_manager/load/model_cache/cached_model/cached_model_only_full_load.py @@ -1,3 +1,4 @@ +import weakref from typing import Any import torch @@ -41,6 +42,7 @@ def __init__( # under `cache_key`; `release_shared_weights()` must be called exactly once on eviction. self._shared_store: SharedCpuWeightsStore | None = None self._shared_key: str | None = None + self._shared_release_finalizer: weakref.finalize | None = None # A CPU read-only copy of the model's state dict. self._cpu_state_dict: dict[str, torch.Tensor] | None = None @@ -57,9 +59,24 @@ def __init__( if canonical is not cpu_state_dict: model.load_state_dict(canonical, assign=True) cpu_state_dict = canonical + # A cache dropped without a shutdown() never routes its records through + # _delete_cache_entry, so nothing would call release_shared_weights() and the + # canonical tensors would stay resident (and counted by the RAM budget) + # forever. The finalizer must not reference `self` (its args are held strongly + # — that would make the wrapper immortal) and must not take the store's + # non-reentrant lock (it runs in GC context): release_deferred only enqueues; + # the store applies it on its next operation. release_shared_weights() detaches + # this on the normal eviction path, so the release happens exactly once either + # way. Registered inside this try so a failure here (e.g. MemoryError) + # releases the just-acquired reference too. + self._shared_release_finalizer = weakref.finalize( + self, shared_store.release_deferred, cache_key, canonical + ) + self._shared_release_finalizer.atexit = False except Exception: - # The re-point failed after acquiring a reference; release it so the shared - # entry's refcount isn't leaked (this wrapper will never enter the cache). + # The re-point or finalizer registration failed after acquiring a reference; + # release it so the shared entry's refcount isn't leaked (this wrapper will + # never enter the cache). self.release_shared_weights() raise self._cpu_state_dict = cpu_state_dict @@ -92,6 +109,11 @@ def release_shared_weights(self) -> None: no-op. After release, the shared store frees the canonical tensors once the last device that held this key releases it. """ + if self._shared_release_finalizer is not None: + # The eviction path is releasing synchronously; the collection-time fallback must not + # release the same reference a second time. + self._shared_release_finalizer.detach() + self._shared_release_finalizer = None if self._shared_store is not None and self._shared_key is not None: self._shared_store.release(self._shared_key, self._cpu_state_dict) self._shared_store = None diff --git a/invokeai/backend/model_manager/load/model_cache/cached_model/cached_model_with_partial_load.py b/invokeai/backend/model_manager/load/model_cache/cached_model/cached_model_with_partial_load.py index cae68e24331..00be7d3d950 100644 --- a/invokeai/backend/model_manager/load/model_cache/cached_model/cached_model_with_partial_load.py +++ b/invokeai/backend/model_manager/load/model_cache/cached_model/cached_model_with_partial_load.py @@ -1,3 +1,5 @@ +import weakref + import torch from invokeai.backend.model_manager.load.model_cache.shared_cpu_weights import SharedCpuWeightsStore @@ -29,6 +31,7 @@ def __init__( # under `cache_key`; `release_shared_weights()` must be called exactly once on eviction. self._shared_store: SharedCpuWeightsStore | None = None self._shared_key: str | None = None + self._shared_release_finalizer: weakref.finalize | None = None # Assigned for real at the end of __init__; initialized here so the acquire-failure path # below can call release_shared_weights(), which reads it, before that assignment runs. self._cpu_state_dict: dict[str, torch.Tensor] | None = None @@ -69,9 +72,23 @@ def __init__( if canonical is not cpu_state_dict: self._model.load_state_dict(canonical, assign=True) cpu_state_dict = canonical + # A cache dropped without a shutdown() never routes its records through + # _delete_cache_entry, so nothing would call release_shared_weights() and the + # canonical tensors would stay resident (and counted by the RAM budget) forever. + # The finalizer must not reference `self` (its args are held strongly — that would + # make the wrapper immortal) and must not take the store's non-reentrant lock (it + # runs in GC context): release_deferred only enqueues; the store applies it on its + # next operation. release_shared_weights() detaches this on the normal eviction + # path, so the release happens exactly once either way. Registered inside this try + # so a failure here (e.g. MemoryError) releases the just-acquired reference too. + self._shared_release_finalizer = weakref.finalize( + self, shared_store.release_deferred, cache_key, canonical + ) + self._shared_release_finalizer.atexit = False except Exception: - # The re-point failed after acquiring a reference; release it so the shared entry's - # refcount isn't leaked (this wrapper will never be inserted into the cache). + # The re-point or finalizer registration failed after acquiring a reference; + # release it so the shared entry's refcount isn't leaked (this wrapper will never + # be inserted into the cache). self.release_shared_weights() raise @@ -175,6 +192,11 @@ def release_shared_weights(self) -> None: no-op. After release, the shared store frees the canonical tensors once the last device that held this key releases it. """ + if self._shared_release_finalizer is not None: + # The eviction path is releasing synchronously; the collection-time fallback must not + # release the same reference a second time. + self._shared_release_finalizer.detach() + self._shared_release_finalizer = None if self._shared_store is not None and self._shared_key is not None: self._shared_store.release(self._shared_key, self._cpu_state_dict) self._shared_store = None diff --git a/invokeai/backend/model_manager/load/model_cache/model_cache.py b/invokeai/backend/model_manager/load/model_cache/model_cache.py index 8f8296b8674..45757ca9366 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 d57e1970d0a..3efe2d2fe12 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)