metrale_bench/hardware/
energy_sampler.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-26: The I/O half of in-window energy: one long-lived `nvidia-smi -lms`
4//! child whose lines a background task parses, plus the per-run [`EnergyMeter`].
5//!
6//! Owner: bench hardware.
7//! Invariants:
8//! - [`EnergyMeter`] samples only a loopback target: the sampler reads the local
9//!   GPU, which is the serving GPU only then.
10//! - When the sampler cannot start, the meter holds no sampler, baseline or cost,
11//!   so the record carries no `gpu_rail_*` keys: absent, never zero.
12//!
13//! One child rather than a spawn per sample: measured 2026-09-20 on GB10, a
14//! single `nvidia-smi` spawn cost 10–20 ms of CPU, while the loop-mode child used
15//! 0.02 s of CPU over 10 s at 250 ms. The child's CPU time is read when it is
16//! stopped and goes on the record ([`SamplerCost`]).
17
18use std::sync::Arc;
19use std::sync::atomic::{AtomicU64, Ordering};
20use std::time::{Duration, Instant};
21
22use anyhow::{Context, Result, bail};
23use parking_lot::Mutex;
24use tokio::io::{AsyncBufReadExt, BufReader};
25
26use super::energy::{
27    EnergyWindow, IDLE_BASELINE_SECS, PowerSample, RAIL_NOTE, SAMPLE_PERIOD_MS, SamplerCost,
28    integrate, parse_line,
29};
30use crate::plugin::TargetEndpoint;
31use crate::result::LogLine;
32
33/// 2026-09-26: The fields the child streams, in the column order [`parse_line`] reads.
34const QUERY: &str = "--query-gpu=power.draw.average,\
35                     clocks_event_reasons.sw_power_cap,\
36                     clocks_event_reasons.hw_power_brake_slowdown";
37
38/// 2026-09-26: How long [`EnergySampler::spawn`] waits for the first reading
39/// before it declares the rail unreadable.
40const FIRST_SAMPLE_TIMEOUT: Duration = Duration::from_secs(3);
41
42pub struct EnergySampler {
43    child: tokio::process::Child,
44    pid: u32,
45    started: Instant,
46    samples: Arc<Mutex<Vec<PowerSample>>>,
47    rejected: Arc<AtomicU64>,
48}
49
50impl EnergySampler {
51    /// 2026-09-26: Start the child and wait for its first reading.
52    ///
53    /// `Err` when `nvidia-smi` cannot be spawned, exits before it is polled, or
54    /// produces no parseable reading within `FIRST_SAMPLE_TIMEOUT`.
55    /// [`EnergyMeter::start`] logs the error and the run goes on without energy.
56    pub async fn spawn() -> Result<Self> {
57        let mut child = tokio::process::Command::new("nvidia-smi")
58            .arg(QUERY)
59            .arg("--format=csv,noheader,nounits")
60            .arg("-lms")
61            .arg(SAMPLE_PERIOD_MS.to_string())
62            .stdin(std::process::Stdio::null())
63            .stdout(std::process::Stdio::piped())
64            .stderr(std::process::Stdio::null())
65            .kill_on_drop(true)
66            .spawn()
67            .context("spawning nvidia-smi for power sampling")?;
68        let pid = child
69            .id()
70            .context("nvidia-smi exited before it was polled")?;
71        let stdout = child.stdout.take().context("nvidia-smi stdout not piped")?;
72        let samples = Arc::new(Mutex::new(Vec::new()));
73        let rejected = Arc::new(AtomicU64::new(0));
74        let (sink, counter) = (samples.clone(), rejected.clone());
75        tokio::spawn(async move {
76            let mut lines = BufReader::new(stdout).lines();
77            while let Ok(Some(line)) = lines.next_line().await {
78                let at = Instant::now();
79                match parse_line(&line, at) {
80                    Some(s) => sink.lock().push(s),
81                    None => {
82                        counter.fetch_add(1, Ordering::Relaxed);
83                    }
84                }
85            }
86        });
87        let this = Self {
88            child,
89            pid,
90            started: Instant::now(),
91            samples,
92            rejected,
93        };
94        let deadline = Instant::now() + FIRST_SAMPLE_TIMEOUT;
95        while this.samples.lock().is_empty() {
96            if Instant::now() >= deadline {
97                let rejected = this.rejected.load(Ordering::Relaxed);
98                bail!(
99                    "nvidia-smi produced no power reading within {FIRST_SAMPLE_TIMEOUT:?} \
100                     ({rejected} unparseable line(s)) — the GPU rail is not readable here"
101                );
102            }
103            tokio::time::sleep(Duration::from_millis(20)).await;
104        }
105        Ok(this)
106    }
107
108    /// 2026-09-26: The integral over `[start, end]` of the readings so far.
109    pub fn window(&self, start: Instant, end: Instant) -> Option<EnergyWindow> {
110        integrate(&self.samples.lock(), start, end)
111    }
112
113    /// 2026-09-26: Sleep [`IDLE_BASELINE_SECS`] and integrate whatever the box
114    /// drew meanwhile. [`EnergyMeter::start`] calls it; its caller must have the
115    /// model loaded and no request in flight.
116    pub async fn idle_baseline(&self) -> Option<EnergyWindow> {
117        let start = Instant::now();
118        tokio::time::sleep(Duration::from_secs(IDLE_BASELINE_SECS)).await;
119        self.window(start, Instant::now())
120    }
121
122    /// 2026-09-26: CPU seconds the child has consumed: the first field of
123    /// `/proc/<pid>/schedstat`, read as nanoseconds.
124    fn cpu_seconds(&self) -> Option<f64> {
125        std::fs::read_to_string(format!("/proc/{}/schedstat", self.pid))
126            .ok()?
127            .split_whitespace()
128            .next()?
129            .parse::<u64>()
130            .ok()
131            .map(|ns| ns as f64 / 1e9)
132    }
133
134    /// 2026-09-26: Measure the child's cost, then kill it.
135    pub async fn stop(mut self) -> SamplerCost {
136        let cost = SamplerCost {
137            cpu_s: self.cpu_seconds(),
138            wall_s: self.started.elapsed().as_secs_f64(),
139            samples: self.samples.lock().len(),
140            rejected_lines: self.rejected.load(Ordering::Relaxed),
141        };
142        if let Err(e) = self.child.kill().await {
143            tracing::debug!("energy sampler already exited: {e}");
144        }
145        cost
146    }
147}
148
149/// 2026-09-26: The energy lifecycle of one benchmark run: start the sampler
150/// (loopback targets only), take the idle baseline, hand out per-window
151/// integrals, stop and report the cost. The concurrency sweep and decode-floor
152/// drivers hold one as a field and log the lines it returns.
153///
154/// When the rail is unreadable the run continues: the returned lines say why,
155/// and the record carries no `gpu_rail_*` keys.
156#[derive(Default)]
157pub struct EnergyMeter {
158    sampler: Option<EnergySampler>,
159    idle: Option<EnergyWindow>,
160    cost: Option<SamplerCost>,
161}
162
163impl EnergyMeter {
164    /// 2026-09-26: Start sampling for `target`, then sample the idle baseline.
165    /// Call with the model loaded and nothing in flight, before the first
166    /// measured window.
167    pub async fn start(&mut self, target: &TargetEndpoint) -> Vec<LogLine> {
168        let mut lines = Vec::new();
169        if !target.is_loopback() {
170            lines.push(LogLine::info(format!(
171                "energy: {} is not loopback — the local GPU rail is not the serving GPU, so \
172                 no energy is recorded for this run",
173                target.base_url
174            )));
175            return lines;
176        }
177        match EnergySampler::spawn().await {
178            Ok(sampler) => {
179                lines.push(LogLine::info(RAIL_NOTE));
180                self.idle = sampler.idle_baseline().await;
181                match &self.idle {
182                    Some(idle) => lines.push(LogLine::info(format!(
183                        "energy: idle baseline (model loaded, nothing in flight) — {}",
184                        idle.one_line(None)
185                    ))),
186                    None => lines.push(LogLine::warn(
187                        "energy: the idle baseline window produced no samples; totals will be \
188                         recorded without an above-idle figure",
189                    )),
190                }
191                self.sampler = Some(sampler);
192            }
193            Err(e) => lines.push(LogLine::warn(format!("energy: not recorded — {e:#}"))),
194        }
195        lines
196    }
197
198    /// 2026-09-26: The integral over one measured window, when sampling is live.
199    pub fn window(&self, start: Instant, end: Instant) -> Option<EnergyWindow> {
200        self.sampler.as_ref()?.window(start, end)
201    }
202
203    pub fn idle(&self) -> Option<&EnergyWindow> {
204        self.idle.as_ref()
205    }
206
207    /// 2026-09-26: Stop the sampler; its cost becomes part of this meter's
208    /// record keys. Returns the line to log, or `None` when nothing was running.
209    pub async fn stop(&mut self) -> Option<LogLine> {
210        let cost = self.sampler.take()?.stop().await;
211        let line = LogLine::info(cost.one_line());
212        self.cost = Some(cost);
213        Some(line)
214    }
215
216    /// 2026-09-26: The run-level keys: the idle baseline and the sampler's cost.
217    /// Per-window keys come from [`EnergyWindow::metrics`].
218    pub fn metrics(&self, m: &mut std::collections::BTreeMap<String, f64>) {
219        if let Some(idle) = &self.idle {
220            m.insert("gpu_rail_idle_power_w".to_string(), idle.mean_power_w);
221            m.insert(
222                "gpu_rail_idle_power_samples".to_string(),
223                idle.samples as f64,
224            );
225        }
226        if let Some(cost) = &self.cost {
227            cost.metrics(m);
228        }
229    }
230}
metrale_bench/hardware/energy_sampler.rs

metrale_bench/hardware/
energy_sampler.rs
