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

@@ -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 isinstance(chunk, TokenChunk):
-                text_parts.append(chunk.text)
-
-            if isinstance(chunk, ToolCallChunk):
-                tool_calls.extend(
-                    ToolCall(
-                        id=str(uuid4()),
-                        index=i,
-                        function=tool,
+                if chunk.finish_reason == "error":
+                    raise HTTPException(
+                        status_code=500,
+                        detail=chunk.error_message or "Internal server error",
                     )
-                    for i, tool in enumerate(chunk.tool_calls)
-                )
 
-            stats = chunk.stats or stats
+                if model is None:
+                    model = chunk.model
 
-            if chunk.finish_reason is not None:
-                finish_reason = chunk.finish_reason
+                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)
+                    )
+
+                stats = chunk.stats or stats
+
+                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

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):

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,
+        )

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