Robustify activation caching: rename navigation, hook lifecycle, membership semantics

[Butanium]

Summary

Started as a fix for #535 (cache attribute access broken under collapsing/re-mount renames) and grew into a robustness pass over the activation-cache subsystem. Three headline changes:

1. Cache navigation is structural and shares rename semantics with the envoy tree — fixes Unused renaming scheme make cache access via renamed name fail #535 and the whole class of string-matching bugs around it.
2. Cache hooks are actually removed at trace exit — previously they leaked, and the next trace on the same model silently corrupted the old cache.
3. Membership (in) and lookup ([]) share one resolution semantics — renamed spellings included.

1. Rename navigation (fixes #535)

Cache storage keys are always real module paths (dict-style access, keys(), iteration unchanged). Navigation resolves against a weights-free skeleton of the envoy tree:

• Envoy._skeleton() snapshots the tree as PathNodes: real path, children by name, aliases resolved to their target nodes via Aliaser.resolve(envoy, alias) — the same function Envoy.__getattr__ uses, so cache and envoy cannot disagree about what an alias means. Nodes hold no modules: the skeleton pickles, travels with remote results, and keeps no weights alive.
• CacheDict navigation is per-segment tree descent (PathNode.resolve, real children authoritative over aliases). Navigation refuses to descend into subtrees where nothing was cached, and error messages now distinguish "exists on the model but was never cached" from "no such module".
• Every alias of a module navigates — cache.model.transformer.h[x], cache.model.model.layers[x], cache.model.layers[x], and secondary aliases of collapsing renames ({"model.layers": ["layers", "blocks"]}) all reach the same entry. test_cache_navigation_matches_envoy sweeps the spellings and asserts envoy/cache equivalence.
• List-valued renames (documented Aliaser syntax) no longer crash tracer.cache(); dot-prefixed alias values ({".transformer.h": ".layers"}, also documented) now work on the envoy side too (Aliaser.build normalizes them).
• Out-of-bounds indices on (renamed) module lists raise IndexError; index 1 can no longer string-prefix-match h.10. Old pickles (no skeleton fields) fall back to a storage-derived trie, as do CacheDicts constructed without a model.
• Also fixed along the way: shared-mutable-default state across all CacheDicts in a process; the fixed cache.model root being corrupted by renames like {"model": "foo"}.

2. Hook lifecycle

Persistent cache hooks were never removed: dead mediators are pruned from interleaver.mediators (in __exit__ and check_dangling_mediators) before cancel() runs remove_hooks(), so the handles leaked onto the modules. Consequence: run a second model.trace(...) on the same model and the first trace's cache hooks fire again — entries silently become [Entry, Entry] lists and previously-working .output calls raise AttributeError: 'list' object has no attribute 'output'.

Hooks can't be removed the moment a mediator dies (vLLM re-enters the interleaver context per engine step and cache hooks must keep recording across steps), so pruned mediators' handles are retired to the interleaver and drained in cancel() — the documented single per-trace cleanup path. test_cache_hooks_removed_after_trace asserts hook counts return to baseline and a prior cache survives subsequent traces intact.

3. Membership = lookup

"model.layers.0" in cache was False while cache["model.layers.0"] succeeded. __contains__ now delegates to __getitem__, so in covers renamed spellings and sub-view-relative keys exactly like access does.

Tests

Nine new/extended tests in tests/test_lm.py::TestCache plus one in TestRename (dotted alias values): collapsing renames (incl. list-valued and secondary aliases), root protection (issue #535's original example), envoy/cache navigation equivalence sweep, membership/lookup agreement, state isolation, alias-never-shadows-real-key, old-pickle compatibility, and hook removal across traces.

tests/test_lm.py: 90 passed · test_tiny.py: 18 · test_envoys/test_source/test_iter_edge_cases/test_memory_cleanup/test_multiple_wrappers/test_transform/test_0516_features: 115 passed, 1 xpassed (all on dev base). All three repro snippets from the #535 thread verified. The branch was also adversarially reviewed by an independent agent (A/B behavioral comparison against dev on seven access patterns, pickle/deepcopy round-trips, multi-fire generation, swap renames); its findings are addressed in the later commits.

https://claude.ai/code/session_01GPKVtGmUQWFsS7dMEUsNNe