metrale_telemetry/
energy.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-26: Energy attribution: the NVML counter's deltas, paired with the
4//! tokens produced over the same interval.
5//!
6//! Two cumulative totals are kept, GPU-rail millijoules and tokens, and
7//! sampled together into [`EnergyRing`]. Every ratio is formed from
8//! differences of those totals: J/token over the last
9//! [`crate::hub::LIVE_WINDOW`] samples, J/token since the first sample, and a
10//! request's energy (its token share of its lifetime window). No ratio is
11//! accumulated.
12//!
13//! The sampler thread reads the counter at the device cadence, not per
14//! scheduler step, so a window's ΔE/Δtokens is exact up to one sample period
15//! at each edge. A backwards step is a wrap or a reset;
16//! [`CounterTracker::advance`] tells them apart.
17//!
18//! Owner: telemetry.
19//! Invariants: none beyond the types.
20
21use crate::seqcell::SeqCell;
22use std::sync::atomic::{AtomicU64, Ordering};
23
24/// 2026-09-26: What one counter reading did to the running total.
25#[derive(Clone, Copy, Debug, PartialEq, Eq)]
26pub enum Advance {
27    /// 2026-09-26: The first reading: it anchors the series and contributes nothing.
28    First,
29    /// 2026-09-26: The counter moved forward by this many mJ.
30    Delta(u64),
31    /// 2026-09-26: The counter wrapped past `u64::MAX`; this many mJ elapsed.
32    Wrapped(u64),
33    /// 2026-09-26: The counter restarted; the new reading, the mJ since the
34    /// restart, is charged.
35    Reset(u64),
36}
37
38impl Advance {
39    pub fn millijoules(self) -> u64 {
40        match self {
41            Self::First => 0,
42            Self::Delta(d) | Self::Wrapped(d) | Self::Reset(d) => d,
43        }
44    }
45}
46
47/// 2026-09-26: Turns raw counter readings into non-negative deltas.
48#[derive(Clone, Copy, Debug, Default)]
49pub struct CounterTracker {
50    prev: Option<u64>,
51}
52
53/// 2026-09-26: A backwards step is a wrap only when the previous reading sat
54/// in the top quarter of the range and the new one in the bottom quarter;
55/// any other backwards step is a restart.
56const QUARTER: u64 = u64::MAX / 4;
57
58impl CounterTracker {
59    pub fn advance(&mut self, raw: u64) -> Advance {
60        let prev = self.prev.replace(raw);
61        match prev {
62            None => Advance::First,
63            Some(p) if raw >= p => Advance::Delta(raw - p),
64            Some(p) if p > u64::MAX - QUARTER && raw < QUARTER => {
65                Advance::Wrapped(raw.wrapping_sub(p))
66            }
67            Some(_) => Advance::Reset(raw),
68        }
69    }
70}
71
72/// 2026-09-26: GPU-rail joules per token for a (ΔmJ, Δtokens) pair. `None`
73/// when either is 0, so the result is never 0 or ∞.
74pub fn joules_per_token(delta_mj: u64, tokens: u64) -> Option<f64> {
75    (tokens > 0 && delta_mj > 0).then(|| delta_mj as f64 / 1000.0 / tokens as f64)
76}
77
78/// 2026-09-26: One ring point: cumulative energy and tokens at one sample instant.
79#[derive(Clone, Copy, Debug, PartialEq, Eq)]
80pub struct EnergyPoint {
81    pub at_ns: u64,
82    pub energy_mj: u64,
83    pub tokens: u64,
84}
85
86/// 2026-09-26: Ring capacity: 8192 samples, ~13.6 minutes at the 10 Hz
87/// serving cadence. A request older than the ring is attributed from the
88/// oldest point it still holds.
89pub const RING_SLOTS: usize = 8192;
90
91/// 2026-09-26: The sampled (time, energy, tokens) series. `push` assumes one
92/// writer (the sampler thread); any number of readers read without a lock.
93pub struct EnergyRing {
94    slots: Box<[SeqCell<3>]>,
95    written: AtomicU64,
96}
97
98impl std::fmt::Debug for EnergyRing {
99    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
100        f.debug_struct("EnergyRing")
101            .field("written", &self.written.load(Ordering::Relaxed))
102            .finish_non_exhaustive()
103    }
104}
105
106impl EnergyRing {
107    pub fn with_capacity(slots: usize) -> Self {
108        assert!(slots >= 2, "a ring needs two points to form a window");
109        Self {
110            slots: (0..slots).map(|_| SeqCell::new()).collect(),
111            written: AtomicU64::new(0),
112        }
113    }
114
115    pub fn push(&self, p: EnergyPoint) {
116        let n = self.written.load(Ordering::Relaxed);
117        let i = (n % self.slots.len() as u64) as usize;
118        self.slots[i].store([p.at_ns, p.energy_mj, p.tokens]);
119        self.written.store(n + 1, Ordering::Release);
120    }
121
122    fn point(&self, logical: u64) -> EnergyPoint {
123        let (_, [at_ns, energy_mj, tokens]) =
124            self.slots[(logical % self.slots.len() as u64) as usize].load();
125        EnergyPoint {
126            at_ns,
127            energy_mj,
128            tokens,
129        }
130    }
131
132    /// 2026-09-26: The live logical index range `[lo, hi)`.
133    fn live(&self) -> (u64, u64) {
134        let hi = self.written.load(Ordering::Acquire);
135        // 2026-09-26: One slot of slack: the writer may be overwriting the oldest.
136        let lo = hi.saturating_sub(self.slots.len() as u64 - 1);
137        (lo, hi)
138    }
139
140    /// 2026-09-26: The point `back` samples before the newest (`0` = the newest).
141    pub fn back(&self, back: u64) -> Option<EnergyPoint> {
142        let (lo, hi) = self.live();
143        let i = hi.checked_sub(1 + back)?;
144        (i >= lo).then(|| self.point(i))
145    }
146
147    /// 2026-09-26: The newest point at or before `t_ns`; the oldest point when
148    /// every live point is later; `None` for an empty ring.
149    pub fn at_or_before(&self, t_ns: u64) -> Option<EnergyPoint> {
150        let (lo, hi) = self.live();
151        if lo == hi {
152            return None;
153        }
154        // 2026-09-26: Binary search for the first point after `t_ns`, in [lo, hi).
155        let (mut a, mut b) = (lo, hi);
156        while a < b {
157            let mid = a + (b - a) / 2;
158            if self.point(mid).at_ns <= t_ns {
159                a = mid + 1;
160            } else {
161                b = mid;
162            }
163        }
164        Some(self.point(if a == lo { lo } else { a - 1 }))
165    }
166}
167
168/// 2026-09-26: Millijoules attributed to a request that emitted `tokens` over
169/// `[start_ns, end_ns]`: its share, by tokens, of the energy counted over that
170/// window. `None` when the ring is empty or the window saw no tokens.
171pub fn request_millijoules(
172    ring: &EnergyRing,
173    start_ns: u64,
174    end_ns: u64,
175    tokens: u64,
176) -> Option<u64> {
177    let p0 = ring.at_or_before(start_ns)?;
178    let p1 = ring.at_or_before(end_ns)?;
179    let de = p1.energy_mj.checked_sub(p0.energy_mj)?;
180    let dt = p1.tokens.checked_sub(p0.tokens).filter(|t| *t > 0)?;
181    Some((de as u128 * tokens as u128 / dt as u128) as u64)
182}
183
184#[cfg(test)]
185#[path = "energy_tests.rs"]
186mod tests;
metrale_telemetry/energy.rs

metrale_telemetry/
energy.rs
