← back to Exo
Include power usage in bench responses (#1692)
f36fd56c383232ccffdc0c05395aefece5a52972 · 2026-03-10 16:02:55 +0000 · ciaranbor
## Motivation
Add power/energy usage tracking to /bench API responses.
## Changes
- New PowerSampler class that periodically samples sys_power from each
node during inference
- New PowerUsage / NodePowerStats API types
- Integrated into both bench chat completion and bench image generation
endpoints
## Why It Works
Runs concurrently via anyio.create_task_group, reads from existing
state.node_system heartbeat data — no new plumbing needed. Energy =
avg_power × elapsed_time.
## Test Plan
### Manual Testing
Run a /bench request and verify power_usage appears in the response.
### Automated Testing
7 async tests covering single/multi-node, energy math, dynamic power
changes, empty state, and lifecycle.
Files touched
M src/exo/master/api.pyM src/exo/shared/types/api.pyA src/exo/utils/power_sampler.pyA src/exo/utils/tests/test_power_sampler.py
Diff
commit f36fd56c383232ccffdc0c05395aefece5a52972
Author: ciaranbor <81697641+ciaranbor@users.noreply.github.com>
Date: Tue Mar 10 16:02:55 2026 +0000
Include power usage in bench responses (#1692)
## Motivation
Add power/energy usage tracking to /bench API responses.
## Changes
- New PowerSampler class that periodically samples sys_power from each
node during inference
- New PowerUsage / NodePowerStats API types
- Integrated into both bench chat completion and bench image generation
endpoints
## Why It Works
Runs concurrently via anyio.create_task_group, reads from existing
state.node_system heartbeat data — no new plumbing needed. Energy =
avg_power × elapsed_time.
## Test Plan
### Manual Testing
Run a /bench request and verify power_usage appears in the response.
### Automated Testing
7 async tests covering single/multi-node, energy math, dynamic power
changes, empty state, and lifecycle.
---
src/exo/master/api.py | 71 +++++++++-------
src/exo/shared/types/api.py | 15 ++++
src/exo/utils/power_sampler.py | 64 +++++++++++++++
src/exo/utils/tests/test_power_sampler.py | 131 ++++++++++++++++++++++++++++++
4 files changed, 253 insertions(+), 28 deletions(-)
diff --git a/src/exo/master/api.py b/src/exo/master/api.py
index db38c877..8c92fc21 100644
--- a/src/exo/master/api.py
+++ b/src/exo/master/api.py
@@ -174,6 +174,7 @@ from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
from exo.shared.types.worker.shards import Sharding
from exo.utils.banner import print_startup_banner
from exo.utils.channels import Receiver, Sender, channel
+from exo.utils.power_sampler import PowerSampler
from exo.utils.task_group import TaskGroup
_API_EVENT_LOG_DIR = EXO_EVENT_LOG_DIR / "api"
@@ -625,6 +626,7 @@ class API:
async def _collect_text_generation_with_stats(
self, command_id: CommandId
) -> BenchChatCompletionResponse:
+ sampler = PowerSampler(get_node_system=lambda: self.state.node_system)
text_parts: list[str] = []
tool_calls: list[ToolCall] = []
model: ModelId | None = None
@@ -632,41 +634,46 @@ class API:
stats: GenerationStats | None = None
- async for chunk in self._token_chunk_stream(command_id):
- if isinstance(chunk, PrefillProgressChunk):
- continue
+ async with anyio.create_task_group() as tg:
+ tg.start_soon(sampler.run)
- if chunk.finish_reason == "error":
- raise HTTPException(
- status_code=500,
- detail=chunk.error_message or "Internal server error",
- )
+ async for chunk in self._token_chunk_stream(command_id):
+ if isinstance(chunk, PrefillProgressChunk):
+ continue
- if model is None:
- model = chunk.model
+ if chunk.finish_reason == "error":
+ raise HTTPException(
+ status_code=500,
+ detail=chunk.error_message or "Internal server error",
+ )
- if isinstance(chunk, TokenChunk):
- text_parts.append(chunk.text)
+ if model is None:
+ model = chunk.model
- if isinstance(chunk, ToolCallChunk):
- tool_calls.extend(
- ToolCall(
- id=str(uuid4()),
- index=i,
- function=tool,
+ if isinstance(chunk, TokenChunk):
+ text_parts.append(chunk.text)
+
+ if isinstance(chunk, ToolCallChunk):
+ tool_calls.extend(
+ ToolCall(
+ id=str(uuid4()),
+ index=i,
+ function=tool,
+ )
+ for i, tool in enumerate(chunk.tool_calls)
)
- for i, tool in enumerate(chunk.tool_calls)
- )
- stats = chunk.stats or stats
+ stats = chunk.stats or stats
+
+ if chunk.finish_reason is not None:
+ finish_reason = chunk.finish_reason
- if chunk.finish_reason is not None:
- finish_reason = chunk.finish_reason
+ tg.cancel_scope.cancel()
combined_text = "".join(text_parts)
assert model is not None
- resp = BenchChatCompletionResponse(
+ return BenchChatCompletionResponse(
id=command_id,
created=int(time.time()),
model=model,
@@ -682,8 +689,8 @@ class API:
)
],
generation_stats=stats,
+ power_usage=sampler.result(),
)
- return resp
async def _trigger_notify_user_to_download_model(self, model_id: ModelId) -> None:
logger.warning(
@@ -1076,10 +1083,18 @@ class API:
num_images: int,
response_format: str,
) -> BenchImageGenerationResponse:
- images, stats = await self._collect_image_chunks(
- request, command_id, num_images, response_format, capture_stats=True
+ sampler = PowerSampler(get_node_system=lambda: self.state.node_system)
+ images: list[ImageData] = []
+ stats: ImageGenerationStats | None = None
+ async with anyio.create_task_group() as tg:
+ tg.start_soon(sampler.run)
+ images, stats = await self._collect_image_chunks(
+ request, command_id, num_images, response_format, capture_stats=True
+ )
+ tg.cancel_scope.cancel()
+ return BenchImageGenerationResponse(
+ data=images, generation_stats=stats, power_usage=sampler.result()
)
- return BenchImageGenerationResponse(data=images, generation_stats=stats)
async def bench_image_generations(
self, request: Request, payload: BenchImageGenerationTaskParams
diff --git a/src/exo/shared/types/api.py b/src/exo/shared/types/api.py
index 2c0736e0..828ee23b 100644
--- a/src/exo/shared/types/api.py
+++ b/src/exo/shared/types/api.py
@@ -172,8 +172,22 @@ class ImageGenerationStats(BaseModel):
peak_memory_usage: Memory
+class NodePowerStats(BaseModel, frozen=True):
+ node_id: NodeId
+ samples: int
+ avg_sys_power: float
+
+
+class PowerUsage(BaseModel, frozen=True):
+ elapsed_seconds: float
+ nodes: list[NodePowerStats]
+ total_avg_sys_power_watts: float
+ total_energy_joules: float
+
+
class BenchChatCompletionResponse(ChatCompletionResponse):
generation_stats: GenerationStats | None = None
+ power_usage: PowerUsage | None = None
class StreamOptions(BaseModel):
@@ -381,6 +395,7 @@ class ImageGenerationResponse(BaseModel):
class BenchImageGenerationResponse(ImageGenerationResponse):
generation_stats: ImageGenerationStats | None = None
+ power_usage: PowerUsage | None = None
class ImageListItem(BaseModel, frozen=True):
diff --git a/src/exo/utils/power_sampler.py b/src/exo/utils/power_sampler.py
new file mode 100644
index 00000000..b336ec9c
--- /dev/null
+++ b/src/exo/utils/power_sampler.py
@@ -0,0 +1,64 @@
+import time
+from collections import defaultdict
+from collections.abc import Callable, Mapping
+from typing import final
+
+import anyio
+
+from exo.shared.types.api import NodePowerStats, PowerUsage
+from exo.shared.types.common import NodeId
+from exo.shared.types.profiling import SystemPerformanceProfile
+
+
+@final
+class PowerSampler:
+ def __init__(
+ self,
+ get_node_system: Callable[[], Mapping[NodeId, SystemPerformanceProfile]],
+ interval: float = 1.0,
+ ):
+ self._get_node_system = get_node_system
+ self._interval = interval
+ self._samples: defaultdict[NodeId, list[SystemPerformanceProfile]] = (
+ defaultdict(list)
+ )
+ self._start_time: float | None = None
+ self._stopped = False
+
+ def _take_sample(self) -> None:
+ for node_id, profile in self._get_node_system().items():
+ self._samples[node_id].append(profile)
+
+ async def run(self) -> None:
+ self._start_time = time.perf_counter()
+ self._take_sample()
+ while not self._stopped:
+ await anyio.sleep(self._interval)
+ self._take_sample()
+
+ 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
+
+ node_stats: list[NodePowerStats] = []
+ for node_id, profiles in self._samples.items():
+ n = len(profiles)
+ if n == 0:
+ continue
+ node_stats.append(
+ NodePowerStats(
+ node_id=node_id,
+ samples=n,
+ avg_sys_power=sum(p.sys_power for p in profiles) / n,
+ )
+ )
+
+ total_avg_sys = sum(ns.avg_sys_power for ns in node_stats)
+ return PowerUsage(
+ elapsed_seconds=elapsed,
+ nodes=node_stats,
+ total_avg_sys_power_watts=total_avg_sys,
+ total_energy_joules=total_avg_sys * elapsed,
+ )
diff --git a/src/exo/utils/tests/test_power_sampler.py b/src/exo/utils/tests/test_power_sampler.py
new file mode 100644
index 00000000..e9c223ce
--- /dev/null
+++ b/src/exo/utils/tests/test_power_sampler.py
@@ -0,0 +1,131 @@
+from collections.abc import Mapping
+
+import anyio
+import pytest
+
+from exo.shared.types.api import PowerUsage
+from exo.shared.types.common import NodeId
+from exo.shared.types.profiling import SystemPerformanceProfile
+from exo.utils.power_sampler import PowerSampler
+
+
+def _make_profile(sys_power: float) -> SystemPerformanceProfile:
+ return SystemPerformanceProfile(sys_power=sys_power)
+
+
+NODE_A = NodeId("node-a")
+NODE_B = NodeId("node-b")
+
+
+@pytest.fixture
+def single_node_sampler() -> PowerSampler:
+ state: dict[NodeId, SystemPerformanceProfile] = {
+ NODE_A: _make_profile(10.0),
+ }
+ return PowerSampler(get_node_system=lambda: state)
+
+
+@pytest.fixture
+def multi_node_state() -> dict[NodeId, SystemPerformanceProfile]:
+ return {
+ NODE_A: _make_profile(10.0),
+ NODE_B: _make_profile(20.0),
+ }
+
+
+async def test_single_sample(single_node_sampler: PowerSampler) -> None:
+ """A sampler that runs briefly should capture at least the initial sample."""
+ async with anyio.create_task_group() as tg:
+ tg.start_soon(single_node_sampler.run)
+ await anyio.sleep(0.05)
+ tg.cancel_scope.cancel()
+
+ result = single_node_sampler.result()
+ assert len(result.nodes) == 1
+ assert result.nodes[0].node_id == NODE_A
+ assert result.nodes[0].avg_sys_power == 10.0
+ assert result.nodes[0].samples >= 1
+ assert result.elapsed_seconds > 0
+
+
+async def test_multi_node_averaging(
+ multi_node_state: dict[NodeId, SystemPerformanceProfile],
+) -> None:
+ """Power from multiple nodes should be summed for total cluster power."""
+ sampler = PowerSampler(get_node_system=lambda: multi_node_state)
+ async with anyio.create_task_group() as tg:
+ tg.start_soon(sampler.run)
+ await anyio.sleep(0.05)
+ tg.cancel_scope.cancel()
+
+ result = sampler.result()
+ assert len(result.nodes) == 2
+ assert result.total_avg_sys_power_watts == 30.0
+
+
+async def test_energy_calculation(single_node_sampler: PowerSampler) -> None:
+ """Energy (joules) should be avg_power * elapsed_seconds."""
+ async with anyio.create_task_group() as tg:
+ tg.start_soon(single_node_sampler.run)
+ await anyio.sleep(0.3)
+ tg.cancel_scope.cancel()
+
+ result = single_node_sampler.result()
+ expected_energy = result.total_avg_sys_power_watts * result.elapsed_seconds
+ assert result.total_energy_joules == expected_energy
+
+
+async def test_changing_power_is_averaged() -> None:
+ """When power changes mid-sampling, the result should be the average."""
+ state: dict[NodeId, SystemPerformanceProfile] = {
+ NODE_A: _make_profile(10.0),
+ }
+ sampler = PowerSampler(get_node_system=lambda: state, interval=0.05)
+
+ async with anyio.create_task_group() as tg:
+ tg.start_soon(sampler.run)
+ await anyio.sleep(0.15)
+ state[NODE_A] = _make_profile(20.0)
+ await anyio.sleep(0.15)
+ tg.cancel_scope.cancel()
+
+ result = sampler.result()
+ avg = result.nodes[0].avg_sys_power
+ # Should be between 10 and 20, not exactly either
+ assert 10.0 < avg < 20.0
+
+
+async def test_empty_state() -> None:
+ """A sampler with no nodes should return an empty result."""
+ empty: Mapping[NodeId, SystemPerformanceProfile] = {}
+ sampler = PowerSampler(get_node_system=lambda: empty)
+
+ async with anyio.create_task_group() as tg:
+ tg.start_soon(sampler.run)
+ await anyio.sleep(0.05)
+ tg.cancel_scope.cancel()
+
+ result = sampler.result()
+ assert len(result.nodes) == 0
+ assert result.total_avg_sys_power_watts == 0.0
+ assert result.total_energy_joules == 0.0
+
+
+async def test_result_stops_sampling() -> None:
+ """Calling result() should stop the sampler's run loop."""
+ state: dict[NodeId, SystemPerformanceProfile] = {
+ NODE_A: _make_profile(10.0),
+ }
+ sampler = PowerSampler(get_node_system=lambda: state, interval=0.02)
+
+ result: PowerUsage | None = None
+ async with anyio.create_task_group() as tg:
+ tg.start_soon(sampler.run)
+ await anyio.sleep(0.1)
+ result = sampler.result()
+ # run() should exit on its own since _stopped is True
+ await anyio.sleep(0.1)
+ tg.cancel_scope.cancel()
+
+ assert result is not None
+ assert result.nodes[0].samples >= 2
← 82c54dd6 Add support for Nemotron sharding (#1693)
·
back to Exo
·
fix: use StreamingDetokenizer in batch generator to fix emoj 38f0c091 →