← back to Exo
Use time-weighted power sampling (#2038)
414132ae9cdf0f5a9a63b4ee95d0c3d465b2c448 · 2026-05-07 03:42:14 -0700 · Alex Cheema
## Why
The power sampler currently averages sampled wattage values
arithmetically. That can be materially wrong when sample intervals are
uneven: a short high-power spike gets the same weight as a long steady
interval. Energy should be computed by integrating power over time, and
average power should be derived from energy / elapsed time.
## How
- Store each power sample with its relative timestamp.
- Anchor the first sample at `t=0` and take a final sample at `elapsed`
when producing results.
- Integrate per-node power using the trapezoidal rule.
- Sum node energy for total cluster energy, then derive total average
system power from total energy / elapsed.
- Add focused unit tests for uneven sample intervals and the
single-sample fallback.
## Tests
- `uv run pytest src/exo/utils/tests/test_power_sampler.py`
- `uv run basedpyright`
- `uv run ruff check src/exo/utils/power_sampler.py
src/exo/utils/tests/test_power_sampler.py`
- `nix fmt`
Files touched
M src/exo/utils/power_sampler.pyM src/exo/utils/tests/test_power_sampler.py
Diff
commit 414132ae9cdf0f5a9a63b4ee95d0c3d465b2c448
Author: Alex Cheema <41707476+AlexCheema@users.noreply.github.com>
Date: Thu May 7 03:42:14 2026 -0700
Use time-weighted power sampling (#2038)
## Why
The power sampler currently averages sampled wattage values
arithmetically. That can be materially wrong when sample intervals are
uneven: a short high-power spike gets the same weight as a long steady
interval. Energy should be computed by integrating power over time, and
average power should be derived from energy / elapsed time.
## How
- Store each power sample with its relative timestamp.
- Anchor the first sample at `t=0` and take a final sample at `elapsed`
when producing results.
- Integrate per-node power using the trapezoidal rule.
- Sum node energy for total cluster energy, then derive total average
system power from total energy / elapsed.
- Add focused unit tests for uneven sample intervals and the
single-sample fallback.
## Tests
- `uv run pytest src/exo/utils/tests/test_power_sampler.py`
- `uv run basedpyright`
- `uv run ruff check src/exo/utils/power_sampler.py
src/exo/utils/tests/test_power_sampler.py`
- `nix fmt`
---
src/exo/utils/power_sampler.py | 53 +++++++++++++++++++++++--------
src/exo/utils/tests/test_power_sampler.py | 30 +++++++++++++++++
2 files changed, 70 insertions(+), 13 deletions(-)
diff --git a/src/exo/utils/power_sampler.py b/src/exo/utils/power_sampler.py
index b8e985a1..c6e61b41 100644
--- a/src/exo/utils/power_sampler.py
+++ b/src/exo/utils/power_sampler.py
@@ -19,19 +19,21 @@ class PowerSampler:
):
self._get_node_system = get_node_system
self._interval = interval
- self._samples: defaultdict[NodeId, list[SystemPerformanceProfile]] = (
- defaultdict(list)
- )
+ self._samples: defaultdict[
+ NodeId, list[tuple[float, SystemPerformanceProfile]]
+ ] = defaultdict(list)
self._start_time: float | None = None
self._stopped = False
- def _take_sample(self) -> None:
+ def _take_sample(self, t_rel: float | None = None) -> None:
+ assert self._start_time is not None
+ ts = t_rel if t_rel is not None else time.perf_counter() - self._start_time
for node_id, profile in self._get_node_system().items():
- self._samples[node_id].append(profile)
+ self._samples[node_id].append((ts, profile))
async def run(self) -> None:
self._start_time = time.perf_counter()
- self._take_sample()
+ self._take_sample(t_rel=0.0)
while not self._stopped:
await anyio.sleep(self._interval)
self._take_sample()
@@ -39,26 +41,51 @@ class PowerSampler:
def result(self) -> PowerUsage:
self._stopped = True
assert self._start_time is not None, "result() called before run()"
- self._take_sample()
elapsed = time.perf_counter() - self._start_time
+ self._take_sample(t_rel=elapsed)
node_stats: list[NodePowerStats] = []
- for node_id, profiles in self._samples.items():
- n = len(profiles)
+ total_energy_j = 0.0
+ for node_id, ts_profiles in self._samples.items():
+ n = len(ts_profiles)
if n == 0:
continue
+ node_energy_j = trapezoidal_energy(ts_profiles, elapsed)
+ avg_power_w = node_energy_j / elapsed if elapsed > 0 else 0.0
+ total_energy_j += node_energy_j
node_stats.append(
NodePowerStats(
node_id=node_id,
samples=n,
- avg_sys_power=sum(p.sys_power for p in profiles) / n,
+ avg_sys_power=avg_power_w,
)
)
- total_avg_sys = sum(ns.avg_sys_power for ns in node_stats)
+ total_avg_sys_w = total_energy_j / elapsed if elapsed > 0 else 0.0
return PowerUsage(
elapsed_seconds=elapsed,
nodes=node_stats,
- total_avg_sys_power_watts=total_avg_sys,
- total_energy_joules=total_avg_sys * elapsed,
+ total_avg_sys_power_watts=total_avg_sys_w,
+ total_energy_joules=total_energy_j,
)
+
+
+def trapezoidal_energy(
+ ts_profiles: list[tuple[float, SystemPerformanceProfile]],
+ elapsed: float,
+) -> float:
+ """Integrate sys_power(t) over the sample window using the trapezoidal rule.
+ First sample is anchored at t=0 and last at t=elapsed (set by `run` /
+ `result`), so the integral spans the full request interval. Falls back to
+ power * elapsed when only one sample exists (constant-power assumption)."""
+ if len(ts_profiles) == 1:
+ return ts_profiles[0][1].sys_power * elapsed
+ energy_j = 0.0
+ for i in range(1, len(ts_profiles)):
+ t_prev, p_prev = ts_profiles[i - 1]
+ t_cur, p_cur = ts_profiles[i]
+ dt = t_cur - t_prev
+ if dt <= 0:
+ continue
+ energy_j += (p_prev.sys_power + p_cur.sys_power) / 2.0 * dt
+ return energy_j
diff --git a/src/exo/utils/tests/test_power_sampler.py b/src/exo/utils/tests/test_power_sampler.py
index 69f4ccee..7880936c 100644
--- a/src/exo/utils/tests/test_power_sampler.py
+++ b/src/exo/utils/tests/test_power_sampler.py
@@ -111,6 +111,36 @@ async def test_empty_state() -> None:
assert result.total_energy_joules == 0.0
+def test_trapezoidal_unit_dt_weighting() -> None:
+ """Pure unit test on the integration helper. Crafted samples where the
+ arithmetic mean is wildly wrong vs the time-weighted result."""
+ from exo.utils.power_sampler import trapezoidal_energy
+
+ # 5 s window. Power = 10 W for the first 4.9 s, then 100 W for the last 0.1 s.
+ # Three samples: t=0 W=10, t=4.9 W=10, t=5.0 W=100.
+ samples = [
+ (0.0, _make_profile(10.0)),
+ (4.9, _make_profile(10.0)),
+ (5.0, _make_profile(100.0)),
+ ]
+ energy = trapezoidal_energy(samples, elapsed=5.0)
+ # (10+10)/2 * 4.9 + (10+100)/2 * 0.1 = 49 + 5.5 = 54.5 J
+ assert abs(energy - 54.5) < 1e-9
+ avg = energy / 5.0 # 10.9 W
+ # Arithmetic mean of the three samples would be (10+10+100)/3 ≈ 40 W.
+ # Trapezoidal correctly weights each segment by its dt.
+ assert abs(avg - 10.9) < 1e-9
+
+
+def test_trapezoidal_unit_single_sample() -> None:
+ """One sample: no window to integrate over, so fall back to constant power
+ over the elapsed duration."""
+ from exo.utils.power_sampler import trapezoidal_energy
+
+ samples = [(0.0, _make_profile(42.0))]
+ assert trapezoidal_energy(samples, elapsed=3.0) == 42.0 * 3.0
+
+
async def test_result_stops_sampling() -> None:
"""Calling result() should stop the sampler's run loop."""
state: dict[NodeId, SystemPerformanceProfile] = {
← edef8004 Store custom model cards in State (#2024)
·
back to Exo
·
Integration tests infra (#1995) fa571313 →