Skip to main content

ruvector_profiler/
power.rs

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