fix(model cache): evict records at shutdown() instead of only releasing shared weights

invokeai/backend/model_manager/load/load_base.py

@@ -56,24 +56,58 @@ class LoadedModelWithoutConfig:
     def __init__(self, cache_record: CacheRecord, cache: ModelCache):
         self._cache_record = cache_record
         self._cache = cache
-        release_first_use_grace = getattr(cache, "release_first_use_grace", None)
-        self._first_use_finalizer = (
-            finalize(self, release_first_use_grace, cache_record)
-            if cache_record.awaiting_first_use and release_first_use_grace is not None
-            else None
+        # Shield the record for the window between this wrapper's construction and its first
+        # lock: without it, an eviction sweep racing that gap — a peer's budget reconcile,
+        # another model's make-room, or the cache's shutdown() — would evict the record out from
+        # under this wrapper, detaching it from the cache's RAM accounting and (for shared
+        # weights) from store ownership while its tensors live on. The few instructions between
+        # get() returning and this constructor arming the hold remain unshielded — an eviction
+        # landing exactly there is the pre-existing, tolerated issue-7513 detached path, and
+        # register_first_use_hold declines to arm on a record that already lost that race. The
+        # hold is released exactly once: on the
+        # first lock (_end_first_use_window), or by the finalizer below if this wrapper is
+        # dropped without ever locking. The finalizer also covers the put()-set admission grace
+        # for a record whose hold could not be armed (no deferred worker running). Both release
+        # routes quote the epoch the hold was armed under, so a hold the cache's dead-worker
+        # recovery already zeroed is never re-released against a successor hold.
+        release_grace = getattr(cache, "release_first_use_grace", None)
+        register_hold = getattr(cache, "register_first_use_hold", None)
+        self._first_use_hold_epoch: Optional[int] = (
+            register_hold(cache_record) if register_hold is not None and release_grace is not None else None
         )
-        if self._first_use_finalizer is not None:
+        self._first_use_finalizer = None
+        if release_grace is not None and (self._first_use_hold_epoch is not None or cache_record.awaiting_first_use):
+            self._first_use_finalizer = finalize(
+                self,
+                release_grace,
+                cache_record,
+                self._first_use_hold_epoch is not None,
+                self._first_use_hold_epoch if self._first_use_hold_epoch is not None else 0,
+            )
             self._first_use_finalizer.atexit = False
 
+    def _end_first_use_window(self) -> None:
+        """This wrapper's first lock ended its get()->lock() window: the record is now pinned by
+        its lock count, so drop the abandonment finalizer and release the first-use hold. Runs at
+        most once — later re-entries of the context manager find nothing to release."""
+        if self._first_use_finalizer is not None:
+            self._first_use_finalizer.detach()
+            self._first_use_finalizer = None
+        if self._first_use_hold_epoch is not None:
+            hold_epoch = self._first_use_hold_epoch
+            self._first_use_hold_epoch = None
+            release_hold = getattr(self._cache, "release_first_use_hold", None)
+            if release_hold is not None:
+                release_hold(self._cache_record, hold_epoch)
+
     def __enter__(self) -> AnyModel:
         # Hold the MODEL_LOAD_LOCK read lock across the VRAM load (lock() runs
         # load_state_dict(assign=True), which calls register_parameter) so it can't overlap a
         # concurrent model construction that has the global register_parameter -> meta patch active.
         # Acquired before the cache's own lock to keep a consistent lock order (see MODEL_LOAD_LOCK).
         with MODEL_LOAD_LOCK.read_lock():
             self._cache.lock(self._cache_record, None)
-        if self._first_use_finalizer is not None:
-            self._first_use_finalizer.detach()
+        self._end_first_use_window()
         try:
             self.repair_required_tensors_on_device()
             return self.model
@@ -96,8 +130,7 @@ def model_on_device(
         # See __enter__ for why the VRAM load is wrapped in the read lock.
         with MODEL_LOAD_LOCK.read_lock():
             self._cache.lock(self._cache_record, working_mem_bytes)
-        if self._first_use_finalizer is not None:
-            self._first_use_finalizer.detach()
+        self._end_first_use_window()
         try:
             self.repair_required_tensors_on_device()
             yield (self._cache_record.cached_model.get_cpu_state_dict(), self._cache_record.cached_model.model)
@@ -113,8 +146,7 @@ def model(self) -> AnyModel:
     def model_in_ram(self) -> Generator[AnyModel, None, None]:
         """Pin the model's cache record in RAM without moving the model to its execution device."""
         self._cache.lock_in_ram(self._cache_record)
-        if self._first_use_finalizer is not None:
-            self._first_use_finalizer.detach()
+        self._end_first_use_window()
         try:
             yield self.model
         finally:

invokeai/backend/model_manager/load/model_cache/cache_record.py

@@ -30,13 +30,31 @@ class CacheRecord:
     # under the in-flight load, breaking the loader's get() or detaching a live model from the
     # cache's RAM accounting. The grace deliberately survives get(): get() is synchronized, and
     # its own lock-release hook may run a pending reconcile before the caller can lock the record
-    # it was just handed. The flag cannot shield a record forever: the cache's local make_room
-    # path ignores it (cold loads are serialized under MODEL_LOAD_LOCK, so make_room can never
-    # see another loader's entry inside the put()->lock() window), and the next admission on the
-    # same cache clears any flag still standing (see the sweep in ModelCache.put()), so a load
-    # that errors out — or a LoadedModel dropped without ever locking — cannot dodge budget
-    # reconciles indefinitely.
+    # it was just handed. The flag cannot shield a record forever: the synchronous eviction
+    # paths (make_room, drop_model) ignore it, and the next admission on the same cache clears
+    # any flag still standing (see the sweep in ModelCache.put()), so a load that errors out
+    # between put() and the LoadedModel's construction cannot dodge budget reconciles
+    # indefinitely. From the wrapper's construction on, the window is tracked by first_use_holds
+    # below, whose release is guaranteed by the wrapper's finalizer rather than by the sweep.
     awaiting_first_use: bool = False
+    # Count of live LoadedModel wrappers holding this record that have not yet locked it. Armed by
+    # ModelCache.register_first_use_hold() (called from LoadedModelWithoutConfig.__init__) and
+    # released exactly once per wrapper — on the wrapper's first lock, or by its weakref finalizer
+    # if it is dropped without ever locking. Unlike awaiting_first_use, these holds are NOT swept
+    # by the next admission: a warm get()'s wrapper can legitimately sit un-entered across another
+    # model's cold load (a node retrieves several models before entering their contexts), and its
+    # finalizer guarantees the release the sweep exists to backstop. The only recovery sweep is
+    # ModelCache zeroing the counts after the deferred worker — the thread that carries
+    # finalizer-initiated releases — is found dead (at the next worker start, and at shutdown),
+    # since a release dispatched toward a dead worker may be dropped and would otherwise shield
+    # the record forever.
+    first_use_holds: int = 0
+    # Bumped whenever stranded holds are zeroed (dead-worker recovery). Every hold release
+    # carries the epoch it was armed under and is ignored across a bump: without this, a
+    # surviving wrapper's late release — or a release enqueued before the old worker died and
+    # drained after the restart — would decrement a FRESH hold armed by a different wrapper
+    # under the healthy replacement worker, silently unshielding that wrapper's window.
+    first_use_holds_epoch: int = 0
 
     def lock(self) -> None:
         """Lock this record."""
@@ -51,3 +69,15 @@ def unlock(self) -> None:
     def is_locked(self) -> bool:
         """Return true if record is locked."""
         return self._locks > 0
+
+    @property
+    def in_first_use_window(self) -> bool:
+        """True while a load or a live LoadedModel wrapper is between obtaining this record and
+        locking it. The asynchronous eviction sweeps (shutdown, budget reconcile, peer-requested
+        eviction) treat such a record like a locked one: evicting it would detach a record whose
+        holder is about to lock it, splitting the model from the cache's RAM accounting and — for
+        shared weights — releasing store ownership while the tensors live on, so a peer's reload
+        would mint a duplicate canonical copy. The synchronous paths (make_room, drop_model,
+        unlock's stale eviction) honor only the first_use_holds half — see awaiting_first_use for
+        why an orphaned grace must stay reachable there."""
+        return self.awaiting_first_use or self.first_use_holds > 0

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