metrale_model_engine/model/types/
teardown.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-26: `TransformerModel` teardown: `adopt_weight_store` and `release_pools`.
4//!
5//! Owner: model-engine.
6//! Invariants:
7//! - `release_pools` attempts every release even after an earlier one fails, and returns the first error.
8
9use metrale_gpu_runtime::gpu::GpuBackend;
10
11use super::TransformerModel;
12
13impl TransformerModel {
14    /// 2026-09-25: Hand the model the ledger of its own weights, for teardown.
15    pub fn adopt_weight_store(&mut self, store: metrale_model_weights::weights::WeightStore) {
16        self.weight_store = Some(store);
17    }
18
19    /// 2026-09-25: Release every GPU resource this model owns, in this order:
20    /// derived weights, SSM snapshots, the SSM state pool, the KV cache, the
21    /// buffer arena, the weight store, then whatever the backend still tracks
22    /// (`GpuBackend::sweep_unreleased`). Every step is attempted even after an
23    /// earlier one fails; the first error is returned, with its label as
24    /// context.
25    pub(in crate::model) fn release_pools(&mut self) -> anyhow::Result<()> {
26        use metrale_core::scope::ModelResource;
27
28        let gpu: &dyn GpuBackend = self.gpu.as_ref();
29        let mut first_error: Option<anyhow::Error> = None;
30        let mut attempt = |label: &'static str, r: anyhow::Result<()>| {
31            if let Err(e) = r
32                && first_error.is_none()
33            {
34                first_error = Some(e.context(label));
35            }
36        };
37
38        attempt("derived weights", self.derived.release(gpu));
39        attempt("ssm snapshots", self.ssm_snapshots.release(gpu));
40        // 2026-09-25: The pool is Arc'd because sequences hold slots. Another
41        // live handle means a sequence was not released before teardown; that
42        // is reported as an error rather than forced.
43        match std::sync::Arc::get_mut(&mut self.ssm_pool) {
44            Some(pool) => attempt("ssm state pool", pool.release(gpu)),
45            None => attempt(
46                "ssm state pool",
47                Err(anyhow::anyhow!(
48                    "{} handle(s) still hold the SSM pool — a sequence was not \
49                     released before teardown",
50                    std::sync::Arc::strong_count(&self.ssm_pool) - 1
51                )),
52            ),
53        }
54        attempt("kv cache", self.kv_cache.lock().release(gpu));
55        attempt("buffer arena", self.buffers.release(gpu));
56        // 2026-09-25: The weight store goes after every owner above, because
57        // the layers hold pointers into it.
58        if let Some(mut store) = self.weight_store.take() {
59            attempt("weight store", store.release(gpu));
60        }
61        // 2026-09-25: Last: allocations no owner above released, such as
62        // weights the loaders fused into layer-owned buffers.
63        let swept = gpu.sweep_unreleased();
64        if swept > 0 {
65            tracing::warn!(target: "metrale_model_engine::model::types", "teardown swept {swept} allocation(s) that no ModelResource \
66                 released — they are reclaimed, but each one is memory whose \
67                 owner is unaccounted for"
68            );
69        }
70
71        match first_error {
72            Some(e) => Err(e),
73            None => Ok(()),
74        }
75    }
76}
metrale_model_engine/model/types/teardown.rs

metrale_model_engine/model/types/
teardown.rs
