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