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

[lstein]

Summary

Follow-on to #9263, addressing the first of the non-blocking issues deferred from review there: a ModelCache dropped without shutdown()/clear() never routes its resident records through _delete_cache_entry() — the only caller of release_shared_weights() — so the process-global SharedCpuWeightsStore kept the entry's refcount and canonical tensors forever, and every surviving peer cache saw phantom RAM in the shared budget (evicting or refusing capacity for bytes no live cache held).

Design

Each cached-model wrapper now registers a weakref.finalize fallback when it acquires shared weights. Two constraints shape the implementation:

• The finalizer runs in GC context, where taking the store's non-reentrant lock could self-deadlock: a collection can fire inside acquire()'s critical section on the same thread — the same rule ModelCache.release_first_use_grace documents. So the finalizer only enqueues the release into a SimpleQueue (lock-free, reentrant-safe); every public store method drains the queue under the lock, so the bytes disappear from the accounting no later than the next store operation — in particular the next budget query.
• The eviction path stays exactly-once: release_shared_weights() detaches the finalizer before releasing synchronously, so a wrapper that was evicted and later collected cannot decrement a peer's reference. The finalizer's args carry the key and the canonical dict — not the wrapper (finalize holds args strongly; referencing self would make the wrapper immortal) — and the state-dict identity check keeps the release correct across invalidate()'s retired entries.

Tests

Five regression tests, each verified to fail before the fix:

• dropping a cache returns the store's refcount, bytes, and RamBudget.total_in_use() to zero (the test JPPhoto specified);
• the collection-time release is enqueue-only — the refcount is untouched until the next store operation drains it (this is the GC-deadlock-safety property, asserted deterministically);
• eviction + later collection release exactly once across two caches sharing a key;
• a retired (invalidate()d while referenced) entry is freed by a collected holder via state-dict identity;
• CachedModelWithPartialLoad behaves identically to CachedModelOnlyFullLoad.

One existing test constructed a wrapper without binding it and relied on the abandoned wrapper leaking its reference; it now binds the wrapper.

Status

Stacked on #9263 (lstein/feat/multi-gpu); the diff shows that branch's commits until it merges. Marked draft until then — rebase onto main and un-draft after #9263 lands.

🤖 Generated with Claude Code

[JPPhoto]

Looks good! Just one thing to address in a different PR:

• invokeai/backend/model_manager/load/model_cache/model_cache.py:550: shutdown() releases shared-store ownership but retains cache records/models. RAM accounting becomes zero while tensors remain live; later same-key acquire creates duplicate canonical weights. Test: put shared model, call shutdown(), verify record/tensors remain while store reports zero, then reacquire same key and compare state-dict identity. Consider clearing records during shutdown, or retain ownership until record eviction; preserve accounting truth.

  Why does this matter?

  • shutdown() is not a hard barrier:
    • It does not prevent later put(), get(), or lock() calls.
    • ModelManagerService.stop() shuts caches down before stopping the session processor.
    • Session workers are cancelled but not joined; the code explicitly says put() after shutdown can occur.

  Evidence: model_cache.py:633, model_manager_default.py:64, session_processor_default.py:532.

  The issue is real, though likely non-blocking unless shutdown races active work.

[lstein]

Confirmed and addressed in #9494 (draft until this PR merges). It routes shutdown()'s idle records through _delete_cache_entry() and marks in-use records stale for unlock() to evict, per your "retain ownership until record eviction" suggestion — plus an identity guard in the stale-eviction path that an adversarial review of the fix surfaced. Details in the PR body.