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,34 @@ 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. That backstop is why the flag is withheld whenever it could outlive its
+    # releasers: an admission made with no deferred worker running, and one made after
+    # shutdown() — after which no further put() is guaranteed to run the sweep. From the
+    # wrapper's construction on, the window is tracked by first_use_holds below, whose release is
+    # guaranteed by the wrapper's finalizer rather than by the sweep.
     awaiting_first_use: bool = False
+    # Count of live LoadedModel wrappers holding this record that have not yet locked it. Armed by
+    # ModelCache.register_first_use_hold() (called from LoadedModelWithoutConfig.__init__) and
+    # 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 once the deferred worker — the thread that carries
+    # finalizer-initiated releases — is gone (from inside the dying worker itself, and as a
+    # backstop at the next worker start and at shutdown), since a release dispatched toward a
+    # dead worker may be dropped and would otherwise shield the record forever.
+    first_use_holds: int = 0
+    # Bumped whenever stranded holds are zeroed (dead-worker recovery). Every hold release
+    # carries the epoch it was armed under and is ignored across a bump: without this, a
+    # 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 +72,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