fix(model cache): release shared weights when a cache goes away

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

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

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

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

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,20 +105,23 @@ 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
 
     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
 
     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,17 +227,20 @@ 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
             )
 
     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())
 
 

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