Skip to main content

batuta/experiment/
metrics.rs

1//! Energy and cost metrics for experiment tracking.
2//!
3//! This module provides types for tracking energy consumption, cost, and
4//! efficiency metrics during ML experiments.
5
6use serde::{Deserialize, Serialize};
7
8/// Energy consumption metrics
9#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
10pub struct EnergyMetrics {
11    /// Total energy consumed in joules
12    pub total_joules: f64,
13    /// Average power draw in watts
14    pub average_power_watts: f64,
15    /// Peak power draw in watts
16    pub peak_power_watts: f64,
17    /// Duration of measurement in seconds
18    pub duration_seconds: f64,
19    /// CO2 equivalent emissions in grams (based on grid carbon intensity)
20    pub co2_grams: Option<f64>,
21    /// Power Usage Effectiveness (datacenter overhead)
22    pub pue: f64,
23}
24
25impl EnergyMetrics {
26    /// Create new energy metrics
27    pub fn new(
28        total_joules: f64,
29        average_power_watts: f64,
30        peak_power_watts: f64,
31        duration_seconds: f64,
32    ) -> Self {
33        Self {
34            total_joules,
35            average_power_watts,
36            peak_power_watts,
37            duration_seconds,
38            co2_grams: None,
39            pue: 1.0,
40        }
41    }
42
43    /// Calculate CO2 emissions based on carbon intensity (g CO2/kWh)
44    pub fn with_carbon_intensity(mut self, carbon_intensity_g_per_kwh: f64) -> Self {
45        let kwh = self.total_joules / 3_600_000.0;
46        self.co2_grams = Some(kwh * carbon_intensity_g_per_kwh * self.pue);
47        self
48    }
49
50    /// Set the Power Usage Effectiveness factor
51    pub fn with_pue(mut self, pue: f64) -> Self {
52        let old_pue = self.pue;
53        self.pue = pue;
54        // Recalculate CO2 if already set (scale by new PUE / old PUE)
55        if let Some(co2) = self.co2_grams {
56            self.co2_grams = Some(co2 / old_pue * pue);
57        }
58        self
59    }
60
61    /// Calculate energy efficiency in FLOPS per watt
62    pub fn flops_per_watt(&self, total_flops: f64) -> f64 {
63        if self.average_power_watts > 0.0 {
64            total_flops / self.average_power_watts
65        } else {
66            0.0
67        }
68    }
69}
70
71/// Cost metrics for experiments
72#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
73pub struct CostMetrics {
74    /// Compute cost in USD
75    pub compute_cost_usd: f64,
76    /// Storage cost in USD
77    pub storage_cost_usd: f64,
78    /// Network transfer cost in USD
79    pub network_cost_usd: f64,
80    /// Total cost in USD
81    pub total_cost_usd: f64,
82    /// Cost per FLOP in USD
83    pub cost_per_flop: Option<f64>,
84    /// Cost per sample processed
85    pub cost_per_sample: Option<f64>,
86    /// Currency (default USD)
87    pub currency: String,
88}
89
90impl CostMetrics {
91    /// Create new cost metrics
92    pub fn new(compute_cost: f64, storage_cost: f64, network_cost: f64) -> Self {
93        Self {
94            compute_cost_usd: compute_cost,
95            storage_cost_usd: storage_cost,
96            network_cost_usd: network_cost,
97            total_cost_usd: compute_cost + storage_cost + network_cost,
98            cost_per_flop: None,
99            cost_per_sample: None,
100            currency: "USD".to_string(),
101        }
102    }
103
104    /// Add FLOP-based cost calculation
105    pub fn with_flops(mut self, total_flops: f64) -> Self {
106        if total_flops > 0.0 {
107            self.cost_per_flop = Some(self.total_cost_usd / total_flops);
108        }
109        self
110    }
111
112    /// Add sample-based cost calculation
113    pub fn with_samples(mut self, total_samples: u64) -> Self {
114        if total_samples > 0 {
115            self.cost_per_sample = Some(self.total_cost_usd / total_samples as f64);
116        }
117        self
118    }
119}