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

[lstein]

Summary

Follow-on to #9403, addressing the issue @JPPhoto flagged in his approving review there: shutdown() released the resident records' shared-store references but kept the records (and their models) in _cached_models, so the accounting stopped describing reality. Concretely:

• The store — and therefore RamBudget.total_in_use() — reported zero for bytes whose tensors the retained wrappers still held.
• shutdown() is not a hard barrier (Invoker.stop() stops the model manager before the session processor, whose workers are cancelled but not joined), so a post-shutdown load of the same key on a peer cache registered a duplicate canonical alongside the still-resident released copy — two copies in RAM, one counted.
• A third defect surfaced while confirming the report: release_shared_weights() flips uses_shared_weights to False, so a post-shutdown eviction of such a record (put() after shutdown() triggering _make_room_internal) debited the non-shared budget for bytes that were admitted as shared — uncounting another still-resident non-shared model's contribution when the cache held a mix.

Design

Per the review suggestion, shutdown() retains ownership until record eviction:

• Idle records are routed through _delete_cache_entry(), which releases shared-store ownership and budget accounting together, exactly once, at the moment the record actually goes away. The release stays synchronous (the original motivation for the shutdown-time release: finalizers only enqueue, and at teardown nothing may drain the queue).
• In-use records — locked by an in-flight generation, or inside the put()→lock() admission window (awaiting_first_use) — keep their references and are marked is_stale; the existing stale path in unlock() evicts them when 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.

Second commit: identity guard in stale eviction

An adversarial review of the first commit surfaced a related pre-existing hazard that the shutdown change arms at every server stop overlapping in-flight work: unlock()'s stale eviction and _delete_cache_entry() matched records by key, not identity. A stale-marked record can be detached while still locked (the VRAM-move error paths delete locked records) and the key re-admitted before the record's last unlock(); the key-only match then popped the new record — detaching it from the cache and all accounting — and debited the non-shared budget for the old record's shared-admitted bytes. Both sites now act only when the record passed in is the current occupant of its key, making a delete of a detached record a full no-op.

Tests

Six regression tests, each verified to fail against the code it guards:

• shutdown evicts idle records and zeroes store refcount + budget, with the wrapper actually collectable (the zero accounting is true);
• a locked record retains its refcount and budget bytes across shutdown, and the last unlock() evicts it exactly once;
• same for a record inside the admission window;
• a peer cache reloading the key after this cache's shutdown adopts the same canonical state dict (identity-checked) instead of registering a duplicate — the reacquire test the review specified;
• a detached stale record's last unlock() leaves a re-admitted same-key record (and the budget) untouched;
• the timeout test now asserts timer cancellation directly and that the idle record is evicted.

Status

Stacked on #9403 (lstein/fix/multigpu-shared-weights-collect); the diff shows that branch's commit until it merges. Marked draft until then — rebase onto main and un-draft after #9403 lands.

🤖 Generated with Claude Code

[JPPhoto]

Merge blockers:

• invokeai/backend/model_manager/load/model_cache/model_cache.py:567-570: shutdown() evicts warm records between get() and LoadedModelWithoutConfig.__enter__(). Delayed lock uses detached record; peer load creates duplicate canonical weights while budget counts only one. Test: warm-load/unlock m, pause after get(), call shutdown(), then enter and load m on peer; assert distinct state dicts and both copies counted.

Other findings/issues:

• invokeai/backend/model_manager/load/model_cache/model_cache.py:541,567-568: Abandoned first-use records are marked stale, but post-shutdown finalizer work is dropped by _dispatch_deferred() at :819-824; no unlock() means record and shared RAM remain pinned. Test: put() normally, create LoadedModelWithoutConfig without entering, shut down, delete it, collect, and assert record/store refcount/budget reach zero.

Suggestions:

• Consider tracking every get()-to-lock() holder through shutdown, with explicit cleanup for abandoned admissions.

[lstein]

Both findings confirmed — thank you, they were both real. Fixed in 5fda854 + d000e0a, following your suggestion of tracking every get()-to-lock() holder through shutdown.

Blocker — shutdown() between get() and __enter__(). Reproduced exactly as described: the warm record is past its admission grace, so the sweep evicted it, the holder locked the detached record via the tolerated issue-7513 path, and a peer reload minted a duplicate canonical while the budget counted one. The fix adds CacheRecord.first_use_holds: armed in LoadedModelWithoutConfig.__init__, released exactly once per wrapper — on its first lock, or by its weakref finalizer if the wrapper is dropped without ever locking. Every eviction path (shutdown sweep, budget reconcile, peer-requested eviction, make_room, drop_model, and unlock's stale eviction) now treats a held record like a locked one: shutdown marks it stale and retains its store ownership and accounting, the holder locks the attached record, and whatever ends the window — the post-use unlock, or the abandonment path — performs the eviction. test_shutdown_retains_record_inside_get_to_lock_window follows your recipe and asserts the peer adopts the same canonical (store.peek identity, refcount 2, budget unchanged).

One residual, disclosed rather than papered over: the few instructions between get() returning and the wrapper's constructor arming the hold remain unshielded. Closing that would mean arming inside get(), whose failure mode is a permanently shielded record whenever the wrapper is never constructed (an exception in between) — and the cache's design treats "shielded with nothing left to unshield it" as strictly worse than the pre-existing tolerated detached-lock fallback. An eviction landing in that gap settles the accounting exactly once (the identity guards from the previous commit), and register_first_use_hold declines to arm on a record that already lost that race.

Non-blocker — abandoned first-use record post-shutdown. Confirmed: the finalizer's deferred work was dropped twice over (_dispatch_deferred's shutdown check plus the worker's own, and the worker had consumed _DEFERRED_STOP anyway). The deferred worker now outlives shutdown() — it is stopped by the existing cache-collection finalizer instead — post-shutdown deferred work is processed, and the abandonment handler itself evicts a stale record whose last holder is gone, because no unlock() is ever coming for it. Your exact test recipe is test_abandoned_holder_reaches_zero_after_shutdown: put() normally, create the LoadedModel without entering, shut down, delete it, collect, and record / store refcount / budget all reach zero.

Hardening that fell out of adversarially reviewing the mechanism, in the same push:

• put() after shutdown() (reachable, per the earlier note) marks the record stale at birth, so its final release evicts it instead of leaving it resident until process exit.
• Holds are granted only while a worker is alive to carry the finalizer's release. A worker death zeroes surviving holds at the next worker start and at shutdown(), and every release quotes the epoch it was armed under — so a stale release (a surviving wrapper's late first lock, or a release enqueued before the death and drained after the restart) can never consume a hold armed afresh by a different wrapper.

Twelve new tests; each was reverted-and-confirmed-failing against the code it guards, including both of your scenarios. Full model_manager suite green, ruff clean.

[JPPhoto]

To fix:

• invokeai/backend/model_manager/load/model_cache/model_cache.py:716-723 (ModelCache.put) re-arms awaiting_first_use for put() calls after shutdown. If loading is canceled before get() or wrapper construction, the stale record remains retained and keeps shared RAM accounted indefinitely. Test: call cache.shutdown(); cache.put("m", DummyModule()); observe cache._cached_models["m"].awaiting_first_use and nonzero budget usage.

Corner/impossible cases:

• invokeai/backend/model_manager/load/model_cache/model_cache.py:584-620, 869-890 (shutdown/_dispatch_deferred) can strand a wrapper hold if the deferred worker dies after shutdown's liveness check. The finalizer then drops its release, while no later admission runs hold recovery. Test: shut down with a live held wrapper, terminate the worker immediately afterward, delete the wrapper, and run gc.collect(); the stale record and first_use_holds remain.

Suggestions:

• Consider disabling admission grace after shutdown, or synchronously evicting unwrapped post-shutdown records.

• Consider making worker termination recovery independent of a later put(), such as a terminal shutdown sweep or a synchronous fallback release.