From f36fd56c383232ccffdc0c05395aefece5a52972 Mon Sep 17 00:00:00 2001 From: ciaranbor <81697641+ciaranbor@users.noreply.github.com> Date: Tue, 10 Mar 2026 16:02:55 +0000 Subject: [PATCH] Include power usage in bench responses (#1692) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## 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 | 75 ++++++++----- 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, 255 insertions(+), 30 deletions(-) create mode 100644 src/exo/utils/power_sampler.py create mode 100644 src/exo/utils/tests/test_power_sampler.py 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 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 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