spark_model/layers/ops/
model_stats.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! [`ModelStats`] — diagnostic counters and one-shot latches owned by the model.
4//!
5//! Sibling to [`ModelLevers`](super::ModelLevers): the levers say what a model's
6//! kernels *do*, this records what they *did*. Both are owned by
7//! `TransformerModel` and lent to every `ForwardContext`.
8//!
9//! Telemetry is not exempt from scoping. A counter that spans a model swap
10//! averages two models together and describes neither, and a one-shot dump
11//! latch that has already fired suppresses the *next* model's dump — the exact
12//! artifact someone asked for by setting the flag. Nothing here changes
13//! generation, so the failure is a wrong measurement rather than a wrong
14//! answer; for a diagnostic those are the same kind of defect.
15//!
16//! Counters are atomics because they are mutated through the shared `&` that
17//! `ForwardContext` hands out. The change from the statics they replace is
18//! where they live, not how they are written.
19
20use std::sync::atomic::AtomicU64;
21
22/// Per-model diagnostic state.
23#[derive(Debug, Default)]
24pub struct ModelStats {
25    /// MoE expert-union sampling (`ModelLevers::moe_union_stats`): calls seen,
26    /// calls sampled, and the running unique-expert / slot totals behind the
27    /// periodic aggregate line.
28    pub moe_union: MoeUnionStats,
29    /// One-shot latches for the `ATLAS_*_DUMP` diagnostics. A latch is per
30    /// model so a swap re-arms the dump instead of silently swallowing it.
31    pub dumped: DumpLatches,
32}
33
34/// Expert-union sampling counters for one model.
35#[derive(Debug, Default)]
36pub struct MoeUnionStats {
37    pub calls: AtomicU64,
38    pub samples: AtomicU64,
39    pub unique_sum: AtomicU64,
40    pub slots_sum: AtomicU64,
41}
42
43/// One-shot diagnostic latches for one model.
44///
45/// The named fields are the latches with a caller that already holds the
46/// struct; [`keyed`](DumpLatches::keyed) covers the long tail of
47/// `ATLAS_*_DUMP` gates, which are numerous, scattered, and identical in
48/// shape — a field each would be noise, and a static each is the bug.
49#[derive(Debug, Default)]
50pub struct DumpLatches {
51    /// Ad-hoc latches, keyed by a `&'static str` naming the dump.
52    fired: std::sync::Mutex<std::collections::BTreeSet<&'static str>>,
53}
54
55impl ModelStats {
56    pub fn new() -> Self {
57        Self::default()
58    }
59
60    /// `true` the first time THIS model reaches `key`, `false` after.
61    ///
62    /// The general per-model latch. Log-dedup gates used a `static Once` each,
63    /// which is correct for "print this line once" and wrong for "print this
64    /// line once per model": after a swap the new model's kernel-route and
65    /// fallback lines — the ones that say which path a model actually took —
66    /// were suppressed by the previous model's shot, and those lines exist to
67    /// be read when a model behaves unexpectedly.
68    ///
69    /// Namespace keys by purpose (`"log:..."`, `"dump:..."`) so two unrelated
70    /// sites cannot collide.
71    pub fn once(&self, key: &'static str) -> bool {
72        self.dumped.keyed(key)
73    }
74}
75
76impl DumpLatches {
77    /// `true` exactly once per model for `key`. Use for the `ATLAS_*_DUMP`
78    /// gates that would otherwise each grow a `static AtomicBool`.
79    ///
80    /// Call it only when the dump is actually wanted — it consumes the shot,
81    /// so gating on the env var FIRST keeps a disabled dump from burning it.
82    pub fn keyed(&self, key: &'static str) -> bool {
83        self.fired
84            .lock()
85            .expect("dump latches poisoned")
86            .insert(key)
87    }
88}
89
90#[cfg(test)]
91mod tests {
92    use super::*;
93    use std::sync::atomic::Ordering;
94
95    #[test]
96    fn a_keyed_latch_fires_once_per_key_per_model() {
97        let a = ModelStats::new();
98        assert!(a.dumped.keyed("dflash_block"));
99        assert!(!a.dumped.keyed("dflash_block"), "and only once");
100        assert!(a.dumped.keyed("dflash_ctx"), "a different dump is separate");
101        assert!(
102            ModelStats::new().dumped.keyed("dflash_block"),
103            "and a new model re-arms every key"
104        );
105    }
106
107    #[test]
108    fn two_models_count_expert_unions_independently() {
109        let a = ModelStats::new();
110        let b = ModelStats::new();
111        a.moe_union.calls.fetch_add(9, Ordering::Relaxed);
112        assert_eq!(a.moe_union.calls.load(Ordering::Relaxed), 9);
113        assert_eq!(
114            b.moe_union.calls.load(Ordering::Relaxed),
115            0,
116            "a second model starts clean rather than inheriting a mean"
117        );
118    }
119}
spark_model/layers/ops/model_stats.rs

spark_model/layers/ops/
model_stats.rs
