Skip to main content

ruvector_profiler/
power.rs

1#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
2pub struct PowerSample { pub watts: f64, pub timestamp_us: u64 }
3
4#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
5pub struct EnergyResult {
6    pub total_joules: f64,
7    pub mean_watts: f64,
8    pub peak_watts: f64,
9    pub duration_s: f64,
10    pub samples: usize,
11}
12
13/// Trait for reading instantaneous power (NVML, RAPL, etc.).
14pub trait PowerSource { fn read_watts(&self) -> f64; }
15
16/// Fixed-wattage mock for deterministic tests.
17pub struct MockPowerSource { pub watts: f64 }
18impl PowerSource for MockPowerSource { fn read_watts(&self) -> f64 { self.watts } }
19
20/// Trapezoidal integration of power samples (must be sorted by timestamp).
21pub fn estimate_energy(samples: &[PowerSample]) -> EnergyResult {
22    let n = samples.len();
23    if n < 2 {
24        return EnergyResult {
25            total_joules: 0.0, samples: n, duration_s: 0.0,
26            mean_watts: samples.first().map_or(0.0, |s| s.watts),
27            peak_watts: samples.first().map_or(0.0, |s| s.watts),
28        };
29    }
30    let (mut joules, mut peak, mut sum) = (0.0f64, f64::NEG_INFINITY, 0.0f64);
31    for i in 0..n {
32        let w = samples[i].watts;
33        sum += w;
34        if w > peak { peak = w; }
35        if i > 0 {
36            let dt = samples[i].timestamp_us.saturating_sub(samples[i - 1].timestamp_us) as f64 / 1e6;
37            joules += (samples[i - 1].watts + w) / 2.0 * dt;
38        }
39    }
40    let dur = samples.last().unwrap().timestamp_us.saturating_sub(samples[0].timestamp_us) as f64 / 1e6;
41    EnergyResult { total_joules: joules, mean_watts: sum / n as f64, peak_watts: peak, duration_s: dur, samples: n }
42}
43
44pub struct PowerTracker { pub samples: Vec<PowerSample>, pub label: String }
45
46impl PowerTracker {
47    pub fn new(label: &str) -> Self { Self { samples: Vec::new(), label: label.to_string() } }
48
49    pub fn sample(&mut self, source: &dyn PowerSource) {
50        let ts = std::time::SystemTime::now()
51            .duration_since(std::time::UNIX_EPOCH).unwrap_or_default().as_micros() as u64;
52        self.samples.push(PowerSample { watts: source.read_watts(), timestamp_us: ts });
53    }
54
55    pub fn energy(&self) -> EnergyResult { estimate_energy(&self.samples) }
56}
57
58#[cfg(test)]
59mod tests {
60    use super::*;
61    fn ps(w: f64, t: u64) -> PowerSample { PowerSample { watts: w, timestamp_us: t } }
62
63    #[test]
64    fn energy_empty() { let r = estimate_energy(&[]); assert_eq!(r.samples, 0); }
65
66    #[test]
67    fn energy_single() {
68        let r = estimate_energy(&[ps(42.0, 0)]);
69        assert_eq!((r.total_joules, r.mean_watts), (0.0, 42.0));
70    }
71
72    #[test]
73    fn energy_constant_100w_1s() {
74        let r = estimate_energy(&[ps(100.0, 0), ps(100.0, 1_000_000)]);
75        assert!((r.total_joules - 100.0).abs() < 1e-9);
76    }
77
78    #[test]
79    fn energy_ramp() {
80        let r = estimate_energy(&[ps(0.0, 0), ps(200.0, 1_000_000)]);
81        assert!((r.total_joules - 100.0).abs() < 1e-9);
82    }
83
84    #[test]
85    fn mock_source() { assert_eq!(MockPowerSource { watts: 75.0 }.read_watts(), 75.0); }
86
87    #[test]
88    fn tracker_collects() {
89        let src = MockPowerSource { watts: 50.0 };
90        let mut t = PowerTracker::new("gpu");
91        t.sample(&src); t.sample(&src);
92        assert_eq!(t.samples.len(), 2);
93    }
94}