metrale_speculative/
snapshot.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: The scheduler's observability snapshot: one `Copy` struct the
4//! scheduler publishes each loop tick into the run's `SnapshotCell`, which the
5//! TUI reads.
6//!
7//! Owner: speculative.
8//! Invariants:
9//! - A `SnapshotCell` holds at most the latest snapshot; `publish` replaces
10//!   it whole.
11
12use std::time::Instant;
13
14use parking_lot::Mutex;
15
16#[derive(Clone, Copy, Debug, PartialEq)]
17pub enum MtpModeSnap {
18    Mtp,
19    Serial,
20    Probing,
21    /// 2026-09-25: No MTP gate in this run.
22    Off,
23}
24
25#[derive(Clone, Copy, Debug)]
26pub struct SchedulerSnapshot {
27    pub active_seqs: u32,
28    pub prefilling_seqs: u32,
29    pub swapped_seqs: u32,
30    pub pending_len: u32,
31    pub kv_blocks_free: u32,
32    pub kv_blocks_total: u32,
33    pub ssm_slots_used: u32,
34    pub ssm_slots_total: u32,
35    pub mtp_mode: MtpModeSnap,
36    /// 2026-09-25: The MTP gate's delivered-throughput EWMA for the current
37    /// mode (tok/s); 0.0 when unmeasured or without a gate.
38    pub delivered_tps: f32,
39    pub steps_total: u64,
40    /// 2026-09-25: When the scheduler published this snapshot, on its clock.
41    pub published_at: Instant,
42}
43
44/// 2026-09-25: The run's snapshot cell: the scheduler writes it, the TUI
45/// reads it. `serve` creates one per run and shares it as an `Arc`, by the
46/// same route as `SchedLevers`.
47#[derive(Default)]
48pub struct SnapshotCell(Mutex<Option<SchedulerSnapshot>>);
49
50impl SnapshotCell {
51    /// 2026-09-25: Replaces the held snapshot.
52    pub fn publish(&self, s: SchedulerSnapshot) {
53        *self.0.lock() = Some(s);
54    }
55
56    /// 2026-09-25: The latest snapshot, if one has been published.
57    pub fn read(&self) -> Option<SchedulerSnapshot> {
58        *self.0.lock()
59    }
60}
61
62#[cfg(test)]
63mod tests {
64    use super::*;
65
66    #[test]
67    fn publish_read_roundtrip() {
68        let cell = SnapshotCell::default();
69        let s = SchedulerSnapshot {
70            active_seqs: 3,
71            prefilling_seqs: 1,
72            swapped_seqs: 0,
73            pending_len: 2,
74            kv_blocks_free: 100,
75            kv_blocks_total: 200,
76            ssm_slots_used: 4,
77            ssm_slots_total: 128,
78            mtp_mode: MtpModeSnap::Mtp,
79            delivered_tps: 42.5,
80            steps_total: 7,
81            published_at: Instant::now(),
82        };
83        cell.publish(s);
84        let r = cell.read().expect("published");
85        assert_eq!(r.active_seqs, 3);
86        assert_eq!(r.mtp_mode, MtpModeSnap::Mtp);
87    }
88}