[malibio]

Problem

The process-wide atexit handler (shared by embedding and chat) only drained the embedding-only LLAMA_STATES registry before dropping LLAMA_BACKEND. ChatEngine::load_model registered the handler but never its own state, so a loaded chat model's Metal residency sets outlived the backend on exit() — the same SIGABRT-on-quit crash class previously fixed for embeddings.

Fix

• embedding.rs: replace the LlamaState-typed registry with GPU_STATES: Vec<Arc<dyn GpuResidentState>>, a trait implemented for any Mutex<Option<T>> model slot. The atexit callback releases every registered state before dropping the backend. register_state_for_cleanup is now generic and pub(crate), and it deduplicates re-registration of the same slot when a model reloads.
• chat/mod.rs: load_model registers self.state right after register_atexit_handler(), the same way embedding does.

Tests

• New registry_releases_states_of_distinct_types: registers two different state types (plus a duplicate registration) and checks each is registered once, released (dropped), and idempotent. It works on the registered entries directly rather than draining the global registry, which would unload models that parallel tests have loaded.
• nlp-engine lib: 125 passed; clippy -D warnings clean; --no-default-features and chat-only feature builds compile.
• Pre-push gate: full pyramid passed.

Out of scope

Graceful shutdown (drain_and_release_gpu, LocalAgentService::shutdown) still doesn't unload the chat engine. The atexit path now covers the crash either way, because it runs on every exit().

🤖 Generated with Claude Code

[malibio]

Pragmatic Code Review — PR #3043

Automated pragmatic-code-review: REQUEST CHANGES (posted as a comment — GitHub blocks self-approval).

Review Type

Initial review. The PR had no earlier reviews or comments.

Requirements Check (#2952)

• 🟢 Root cause addressed: the shared atexit cleanup() now drains one GPU_STATES registry that holds both embedding and chat slots, before LLAMA_BACKEND is dropped. ChatEngine::load_model registers self.state after every successful load. This matches the issue's first suggested fix, generalizing the registry through a trait object.
• 🔴 Regression introduced: the registry keeps a strong Arc to the chat slot. Production code never calls ChatEngine::unload_model. It releases a chat model only by dropping the engine. So from this PR on, unloading or switching a model no longer frees GPU memory. See finding 1.

Code Review Findings

🔴 Critical / Blocker

1. The registry now keeps every chat model resident after unload or switch (multi-GB Metal leak)
📁 packages/nlp-engine/src/embedding.rs:134-141 (register_state_for_cleanup) and packages/nlp-engine/src/chat/mod.rs:291

GPU_STATES stores Arc<dyn GpuResidentState>, which is a strong reference to the service's Arc<Mutex<Option<ChatLlamaState>>>. Entries are pruned only when is_loaded() returns false, meaning the slot is None. Here is how production releases a chat model:

• LocalAgentService::unload_model calls shared.reset_to_noop_engine(), which calls set_engine(Arc::new(NoOpInferenceEngine)) (packages/daemon/src/services/local_agent_service.rs:377, :1637)
• A model switch calls replace_engine / set_engine_if_changed

Both paths drop the LlamaChatInferenceEngine, and with it the Arc<ChatEngine>. Neither calls ChatEngine::unload_model(). A grep shows no production caller, and neither LlamaChatInferenceEngine nor ChatEngine implements Drop. Before this PR, dropping the engine dropped the only Arc to the state, which freed the model and context. After this PR, the registry's Arc keeps the Some(ChatLlamaState) alive. It is never None, so it is never pruned. Consequences:

• Unload does nothing to GPU memory. The model stays resident until the process exits.
• Model switch A to B keeps both models resident. On the 16 GB Mac target, the second load will probably OOM or fall back to CPU. Every further switch leaks one more model.

The same latent problem already existed for embedding: an EmbeddingService dropped without shutdown() stayed pinned. It stayed hidden only because embedding services are long-lived.

Fix (small): have the registry observe slots, not own them. Store Weak<dyn GpuResidentState>. Weak<Mutex<Option<T>>> coerces to Weak<dyn Trait>.

static GPU_STATES: Mutex<Vec<Weak<dyn GpuResidentState>>> = Mutex::new(Vec::new());

pub(crate) fn register_state_for_cleanup<T: Send + 'static>(state: &Arc<Mutex<Option<T>>>) {
    let state: Arc<dyn GpuResidentState> = Arc::clone(state) as _;
    let weak = Arc::downgrade(&state);
    let mut states = GPU_STATES.lock().unwrap_or_else(|p| p.into_inner());
    states.retain(|s| s.strong_count() > 0 && !Weak::ptr_eq(s, &weak));
    states.push(weak);
}
// atexit: states.iter().filter_map(Weak::upgrade).filter(|s| s.release()).count()

This also removes the need to lock slots during pruning (see finding 2). A slot that is unloaded but still owned has almost no cost to keep, and the ptr_eq dedupe covers reloads.

Test to add: register a slot, drop the owning Arc, then assert that the inner value's Drop ran (DROPS == 1). That test fails on the current code and would have caught this regression.

🟡 Improvements

2. Registration can now stall behind a chat generation running on a different engine
📁 packages/nlp-engine/src/embedding.rs:139 (retain(|s| s.is_loaded() && …))

register_state_for_cleanup holds the GPU_STATES lock and locks every registered slot through is_loaded(). ChatEngine::generate_blocking (chat/mod.rs:367) holds the chat slot mutex for the entire generation, up to 2048 tokens. So EmbeddingService::initialize() or another chat load_model now blocks until an unrelated in-flight chat turn finishes. The embedding-only registry never had this cross-service coupling, because embedding holds its lock for milliseconds.

This is not a deadlock. Lock order is always GPU_STATES → slot, and no path holds a slot lock while taking GPU_STATES: load_model and initialize both release the slot guard before registering. It is still an unnecessary stall. The Weak/strong_count pruning from finding 1 removes it entirely. If you keep the strong registry, use try_lock and treat WouldBlock as loaded.

3. Quitting during a chat generation now waits for that generation to finish
📁 packages/nlp-engine/src/embedding.rs:122-126 (release_registered_states)

exit() does not stop other threads. The atexit callback blocks on the chat slot mutex until generate_blocking returns. Waiting is the correct choice, because freeing the model while another thread is mid-decode would be worse than the original crash. It is also strictly better than a SIGABRT. The cost is that Quit can hang for the rest of a turn, which can be many seconds on E4B. generate_blocking has no cancel flag today. A follow-up could add a process-wide "shutting down" AtomicBool that the token loop checks and that cleanup() sets before it takes locks. That is not blocking for this PR, but it deserves an issue.

4. The test proves the dedupe and the release on distinct types, but not the lifetime contract
📁 packages/nlp-engine/src/embedding.rs:896-949

The test is well built. It scopes itself to its own entries with Arc::ptr_eq instead of draining the global registry, so other tests' models are untouched, and the function-local DROPS static keeps it isolated. What it does not check is that registration leaves ownership unchanged. That is exactly the invariant finding 1 breaks. Add the drop-the-owner case described above.

🟢 Looks Good

• Trait-object design (KISS): a blanket impl<T: Send> GpuResidentState for Mutex<Option<T>> is the smallest abstraction that unifies both state types with no enum and no per-type registry. Mutex<Option<T>>: Sync requires only T: Send, so the Send + Sync supertrait bound holds. ChatLlamaState's unsafe impl Send already exists, so nothing new is unsafe.
• Arc::ptr_eq on dyn: since Rust 1.72 it ignores vtable metadata (the toolchain here is 1.94.1), so the dedupe cannot give a false negative from a duplicated vtable.
• Pruning and reload: a reload re-registers after Some is written, and retain removes the old entry for the same slot before pushing. This is correct, and it becomes simpler with Weak.
• cfg gating: I checked --features chat-service only and --features embedding-service only (both with --tests) and the default build. All compile cleanly and add no new warnings. --no-default-features --tests fails on llama_cpp_2 paths in chat/mod.rs:45, :1270 and :2718+. Those lines are outside this diff and the failure also happens on main.
• Atexit ordering: it drains states first, then the backend, and uses poison-tolerant locks throughout, the same as before.
• CLAUDE.md: no compatibility shims, no lint suppression, no issue numbers in comments. LLAMA_STATES is fully removed.
• The new test passes locally (cargo test -p nodespace-nlp-engine --lib registry_releases).

Nits

• Nit: embedding.rs:125: release_registered_states returns a count that only the test would use, and the test does not call it. Either have the test assert on the count, or return ().
• Nit: the registry, trait and atexit handler now serve both services but still live in embedding.rs. A small gpu_lifecycle.rs module (backend + registry + atexit) would put the shared ownership where a reader expects it. This is optional and fine as a follow-up.

Recommendation

REQUEST CHANGES. The fix to the atexit path is correct, but it turns every chat unload and model switch into a permanent multi-GB GPU leak, because the registry owns the slot. Switching to a Weak registry and adding a drop-the-owner test is a small change that fixes findings 1 and 2 together. Once that lands, this is an APPROVE.

[malibio]

Automated pragmatic-code-review: REQUEST CHANGES (posted as a comment — GitHub blocks self-approval). The strong-Arc GPU_STATES registry keeps dropped chat engines' models resident on unload/switch; see the review comment for details and the Weak-based fix.

[malibio]

Review addressed (23398e9)

✅ 🔴 Blocker 1: registry pinned chat models: confirmed. ChatEngine::unload_model has no production caller, and unload/switch drops the engine. GPU_STATES now holds Weak<dyn GpuResidentState>. Registration prunes dead refs (strong_count() == 0) and dedupes with Weak::ptr_eq; atexit upgrades before releasing. New test registry_does_not_keep_a_dropped_owners_state_alive.
✅ 🟡 2: registration stalling behind an in-flight generation: pruning no longer locks slots.
✅ 🟡 4: the test missed the ownership rule: covered by the new test.
✅ Nit: unused release count / is_loaded: removed.
⏭️ 🟡 3: quit waits for an in-flight chat turn: skipped. Waiting is the correct behavior (freeing mid-decode is worse), and interrupting generation on shutdown is a separate feature.
⏭️ Nit: move the registry to gpu_lifecycle.rs: skipped as scope creep; it stays with LLAMA_BACKEND, which it tears down.

Pre-push gate: all tests passed.

Re-review needed: the blocker fix changes the registry's ownership model.

[malibio]

Pragmatic Code Review (Re-review): PR #3043

Automated pragmatic-code-review: APPROVE (posted as a comment — GitHub blocks self-approval).

Review Type

Re-review of b1f94331..23398e93, which addresses the REQUEST CHANGES review.

Previous Review Summary

• Previous recommendations: 6 (1 blocker, 3 improvements, 2 nits)
• Addressed: 4 (Blocker 1, Improvement 2, Improvement 4, and the unused-count nit)
• Not addressed, with a sound reason: 2 (Improvement 3, and the nit about moving the module)
• New issues: 0 blocking

Verification

Blocker 1: the registry pinned models. RESOLVED.

• GPU_STATES now stores Weak<dyn GpuResidentState>. The unsized coercion happens on the Arc before Arc::downgrade, which is correct. The registry now holds no strong reference, so the only way the chat model is released in production also frees it. That path drops LlamaChatInferenceEngine, then Arc<ChatEngine>, then the state Arc. ChatEngine has no Drop impl, and none is needed now.
• Atexit path: filter_map(Weak::upgrade), then release(). The upgraded Arc is a temporary inside the iterator chain, and release() .take()s the inner value. The model and context therefore drop inside release_registered_states, which runs strictly before LLAMA_BACKEND.take(). The ordering is correct.
• Coverage window: load_model and initialize both write Some(state) and then register, with no await in between. A live slot is unregistered only for the few instructions between those two statements. That window existed before this PR, and it matters only if exit() lands in those microseconds. Not a concern.
• Embedding path: EmbeddingService owns its state Arc. The daemon keeps the service as Arc<EmbeddingService> in the watch channel (assembly.rs). On graceful shutdown, release_shared_gpu calls release_gpu_context(), which clears the slot, so atexit finds None. If the service is still alive at exit() without a graceful shutdown, upgrade succeeds and the slot is released. If the service has already been dropped, its model was freed with it. All three cases are correct.

Improvement 2: registration locked slots. RESOLVED. retain now uses only Weak::strong_count and Weak::ptr_eq. Neither takes a slot mutex, so loading can no longer stall behind an in-flight generate_blocking.

New test proves the ownership rule. CONFIRMED. registry_does_not_keep_a_dropped_owners_state_alive registers the slot, drops the only owner, and asserts that Model::drop ran. With the previous strong Arc registry, the registry's clone keeps Some(Model) alive, DROPPED stays false, and the assertion fails. It is a real regression guard. It uses a function-local static, so it is isolated from the other registry test. Both pass locally (cargo test -p nodespace-nlp-engine --lib registry_).

Skip rationales. SOUND.

• Improvement 3 (quit waits for an in-flight turn): waiting is correct, and freeing the model mid-decode would reintroduce a crash. A cooperative cancel on shutdown is a separate feature. It is worth a backlog issue if Quit latency during a turn turns out to be noticeable.
• Nit (move to gpu_lifecycle.rs): keeping the registry beside LLAMA_BACKEND, which it has to precede in teardown, is a defensible grouping. This is optional.

Findings

Critical Issues

None.

Suggested Improvements

None.

Nitpicks

• Nit: embedding.rs register_state_for_cleanup: Weak entries for dropped owners are pruned only on the next registration. Each costs one allocation header, and the list is bounded by the number of loads, so this is negligible. No change needed.

Recommendation

APPROVE. The Weak registry resolves the GPU leak on unload and model switch, removes the cross-service lock coupling, and keeps the atexit ordering intact. The new test pins down the ownership rule. The change is small, focused, and follows CLAUDE.md. Merging is the author's call.

[malibio]

Automated pragmatic-code-review: APPROVE (posted as a comment — GitHub blocks self-approval). Weak registry resolves the blocker and the lock-coupling finding; new test guards the ownership rule. See the re-review comment.

[malibio]

Re-review addressed

Verdict was APPROVE with one non-blocking nit.

⏭️ Nit: entries for dropped owners are cleared only at the next registration. Skipped. Each stale entry is a tiny Weak, the list grows by at most one per model load, and atexit skips dead entries via upgrade(). Pruning elsewhere would add code for no measurable gain.

No code changes; no re-review needed.