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.

[lstein]

Both confirmed and fixed in 4445d2d, and the merge conflicts with main are resolved (two hunks, both from #9403's shutdown() shared-weights release loop landing upstream after this branch's merge-base — this PR replaces that loop, so our side won both; all eight tests main added to test_model_cache_ram_budget.py in the meantime survive the resolution).

To fix — put() re-arming the grace after shutdown(). Confirmed exactly as you described. The grace has three releasers — the loader's own lock(), the wrapper's abandonment finalizer, and the sweep at the top of the next put() — and a load cancelled before the LoadedModel is constructed has no wrapper, hence no finalizer, so it depends entirely on that sweep. After shutdown no further put() is guaranteed, and the flag stands for the life of the process. put() now withholds it once _shutdown_event is set; the record is still stale at birth, so a loader that does come back still evicts on release, and it gets the ordinary first_use_holds shield from its wrapper. test_post_shutdown_admission_gets_no_first_use_grace follows your recipe.

Corner case — the worker dying after shutdown()'s liveness check. Confirmed, and taking your second suggestion rather than the first: recovery is now a property of the death itself. _run_deferred_work has an outer finally that distinguishes its two orderly exits (_DEFERRED_STOP, collected cache) from an abnormal one and calls _recover_from_dead_worker(current_thread()), which is scoped by thread identity — a replacement worker's shields are its own — and retires the worker slot before sweeping, so a concurrent admission cannot arm a shield the recovery is about to zero. test_worker_death_after_shutdown_recovers_its_stranded_holds is your recipe with the death made abnormal (a BaseException out of the worker's per-item handler); the record, its store refcount and its budget bytes all reach zero off the death alone.

Adversarially reviewing my own mechanism afterwards was worth more than the original fix, so in the interest of not hiding it: the first two rounds of that fix were wrong in three ways, all now corrected and each pinned by a test that fails against the broken version.

• The recovery cleared the admission grace unconditionally. On a live cache that is a new bug, not a fix: the dying worker is is_alive() for as long as it unwinds, so a cold load landing in that window starts no replacement worker and is admitted with the ordinary grace — which the recovery then zeroed while the loader was still between put() and get(), so a reconcile could evict the record and the loader's get() would raise. The grace only actually loses a releaser once the cache is shut down.
• It also evicted stale-unshielded records from _ensure_deferred_worker, which register_first_use_hold calls before arming — so a second wrapper's construction could detach the very record it was about to shield, releasing shared-store ownership while live wrappers still held the tensors. That is precisely the accounting lie shutdown() itself refuses to make. The eviction now lives in _evict_stale_unshielded_entries, called only from the dying worker and only on a shut-down cache, where nothing else can ever run it.
• Retiring the worker slot silently disabled both remaining recovery sites, which gated on "a dead thread still occupying the slot". The dying recovery deliberately leaves a live cache's grace standing and hands the lift to shutdown() — which, with the slot already empty, skipped it and stale-retained the record forever. Both gates now key on "no live worker", which also makes a failed recovery retryable at the next worker start.

Smaller hardening in the same pass:

• _AbandonedHolderRelease holds its CacheRecord weakly. _dispatch_deferred's liveness gate is unsynchronized by necessity (a finalizer must not take the cache lock), so a caller that read the worker slot just before it was retired still enqueues against a thread that is alive only because it is unwinding — and on a shut-down cache that item may never be drained. A strong reference pinned the model's whole CPU state dict for the life of the process while the store and the budget both reported the bytes released. The worker also clears the strong reference it resolves before parking on the next get(); the existing test_deferred_worker_does_not_pin_the_record_it_just_released caught that one mid-flight.
• _clear_stranded_first_use_holds unshields every record before reporting any of them. A logging handler that raises is one of the ways the worker dies in the first place, and logging inline let that same handler abort the sweep partway, leaving the records it never reached shielded for good. The post-eviction gc/empty_cache housekeeping no longer takes the eviction down with it either.
• _reconcile_budget_if_pending acquires the cache lock adjacent to its try. A BaseException between the two leaked the RLock to an unwinding thread, blocking every other thread for the life of the process. This narrows the window rather than closing it — that residue is inherent to acquiring without with, which the non-blocking flag rules out here — and the comment says so.

Two residuals disclosed rather than papered over. A put()-set grace whose loader is cancelled before get() is still stale-retained at shutdown() while the worker is healthy: the alternative is evicting a record whose loader may still hold its tensors, which is the duplicate-canonical lie, and test_shutdown_retains_admission_window_records pins that choice deliberately. And _evict_stale_unshielded_entries does not fire the on_cache_models_cleared callbacks the other eviction sites do; it only runs on a shut-down cache, so nobody is listening.

Eleven new tests across the three commits. Every production line the round-1 and round-2 fixes touch was individually reverted and confirmed to fail exactly the test that guards it. Full tests/backend/model_manager/ green (1229 passed), ruff clean.

[JPPhoto]

A few reachable and possible issues:

• invokeai/backend/model_manager/load/model_cache/model_cache.py:654-688, 1091-1116: An admission can fail between put() and get(). With a live worker, shutdown() marks it stale but leaves awaiting_first_use; no finalizer or later sweep removes it. Effect: model, shared-store reference, and budget charge persist until cache/process destruction. Likelihood: plausible shutdown cancellation/error. Recovery: manual eviction or restart. Test: put() without get(), call shutdown() while worker lives, then verify record and budget remain.

• invokeai/backend/model_manager/load/load_base.py:59-87; invokeai/backend/model_manager/load/model_cache/model_cache.py:684-688, 1027-1052: Hold registration occurs after get(). If shutdown evicts during that gap, registration returns None, then the wrapper locks a detached record. Effect: peer reload can create duplicate canonical weights while budget counts one copy, risking OOM. Likelihood: rare but reachable during shutdown racing an in-flight load. Recovery: no automatic reattachment; restart/cache rebuild. Test: pause register_first_use_hold(), run shutdown(), resume, then reload the key on a peer and compare store/budget accounting.

Suggestions:

• Consider making wrapper claiming atomic with cache lookup, and explicitly retire unclaimed admissions during shutdown.