1use serde::Serialize;
11
12use crate::device::DeviceReading;
13use crate::hub::{DeviceState, Telemetry};
14use crate::instrument::HistogramSnapshot;
15use crate::sched::SchedShape;
16
17#[derive(Clone, Debug, PartialEq, Serialize)]
18pub struct DeviceSnapshot {
19 #[serde(flatten)]
20 pub reading: DeviceReading,
21 pub samples: u64,
23 pub age_ns: u64,
25 pub read_errors: u64,
26}
27
28#[derive(Clone, Debug, PartialEq, Serialize)]
29pub struct EnergySnapshot {
30 pub gpu_mj_total: u64,
33 pub gpu_counter_mj: u64,
35 pub tokens_total: u64,
37 pub joules_per_token_live: Option<f64>,
39 pub joules_per_token_cumulative: Option<f64>,
40 pub counter_resets: u64,
41 pub counter_wraps: u64,
42}
43
44#[derive(Clone, Debug, PartialEq, Serialize)]
45pub struct SchedSnapshot {
46 #[serde(flatten)]
47 pub shape: SchedShape,
48 pub stream_syncs: u64,
49 pub blocking_d2h: u64,
50 pub graph_replays: u64,
51 pub graph_captures: u64,
52}
53
54#[derive(Clone, Debug, PartialEq, Serialize)]
55pub struct SpecRow {
56 pub drafts: usize,
57 pub steps: u64,
58 pub acceptance_by_depth: Vec<f64>,
60}
61
62#[derive(Clone, Debug, PartialEq, Serialize)]
63pub struct RequestSnapshot {
64 pub finished: u64,
65 pub ttft_mean_ms: Option<f64>,
66 pub tpot_mean_ms: Option<f64>,
67 pub e2e_mean_ms: Option<f64>,
68 pub energy_mean_j: Option<f64>,
69 pub energy_unattributed: u64,
70}
71
72#[derive(Clone, Debug, PartialEq, Serialize)]
73pub struct TelemetrySnapshot {
74 pub level: &'static str,
75 pub at_ns: u64,
76 pub device_state: &'static str,
77 pub device: Option<DeviceSnapshot>,
78 pub energy: EnergySnapshot,
79 pub sched: SchedSnapshot,
80 pub spec: Vec<SpecRow>,
81 pub requests: RequestSnapshot,
82}
83
84fn mean(h: &HistogramSnapshot, scale: f64) -> Option<f64> {
85 (h.count > 0).then(|| h.sum as f64 / h.count as f64 / scale)
86}
87
88impl DeviceState {
89 pub fn as_str(self) -> &'static str {
90 match self {
91 Self::NotStarted => "not_started",
92 Self::Available => "available",
93 Self::Unavailable => "unavailable",
94 }
95 }
96}
97
98impl Telemetry {
99 pub fn snapshot(&self) -> TelemetrySnapshot {
102 let at_ns = self.now_ns();
103 let device = self
104 .device_reading()
105 .map(|(reading, t, samples)| DeviceSnapshot {
106 reading,
107 samples,
108 age_ns: at_ns.saturating_sub(t),
109 read_errors: self.device_read_errors.get(),
110 });
111 let s = &self.sched;
112 let r = &self.requests;
113 let spec = self
114 .spec
115 .widths()
116 .into_iter()
117 .map(|d| SpecRow {
118 drafts: d,
119 steps: (0..=d).map(|a| self.spec.steps(d, a)).sum(),
120 acceptance_by_depth: self.spec.acceptance_by_depth(d),
121 })
122 .collect();
123 TelemetrySnapshot {
124 level: self.level().as_str(),
125 at_ns,
126 device_state: self.device_state().as_str(),
127 device,
128 energy: EnergySnapshot {
129 gpu_mj_total: self.energy_mj.get(),
130 gpu_counter_mj: self.energy_counter_raw.get(),
131 tokens_total: self.tokens_since_first_sample(),
132 joules_per_token_live: self.joules_per_token_live.get(),
133 joules_per_token_cumulative: self.joules_per_token_cumulative(),
134 counter_resets: self.energy_resets.get(),
135 counter_wraps: self.energy_wraps.get(),
136 },
137 sched: SchedSnapshot {
138 shape: s.shape(),
139 stream_syncs: s.stream_syncs.get(),
140 blocking_d2h: s.blocking_d2h.get(),
141 graph_replays: s.graph_replays.get(),
142 graph_captures: s.graph_captures.get(),
143 },
144 spec,
145 requests: RequestSnapshot {
146 finished: r.finished.get(),
147 ttft_mean_ms: mean(&r.ttft.snapshot(), 1e6),
148 tpot_mean_ms: mean(&r.tpot.snapshot(), 1e6),
149 e2e_mean_ms: mean(&r.e2e.snapshot(), 1e6),
150 energy_mean_j: mean(&r.energy.snapshot(), 1e3),
151 energy_unattributed: r.energy_unattributed.get(),
152 },
153 }
154 }
155}