metrale_speculative/
snapshot.rs1use 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 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 pub delivered_tps: f32,
39 pub steps_total: u64,
40 pub published_at: Instant,
42}
43
44#[derive(Default)]
48pub struct SnapshotCell(Mutex<Option<SchedulerSnapshot>>);
49
50impl SnapshotCell {
51 pub fn publish(&self, s: SchedulerSnapshot) {
53 *self.0.lock() = Some(s);
54 }
55
56 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}