diff --git a/docs/POWER_MONITORING.md b/docs/POWER_MONITORING.md new file mode 100644 index 000000000..bd1e44a25 --- /dev/null +++ b/docs/POWER_MONITORING.md @@ -0,0 +1,132 @@ +# Power Monitoring + +First-class, **vendor-neutral** power/energy telemetry collected _during_ a +benchmark run, windowed to the performance phase, and written to a sibling +`power.json` (plus a section in `report.txt`). The `Report` / +`result_summary.json` are never mutated — power is a separate artifact, exactly +like profiling. + +## Design: a general framework + pluggable sources + +The core (`settings.power`) knows nothing about GPUs, Prometheus, or any vendor. +A **source** is a small registered builder that turns the config into a sidecar +command + a parse spec. `nvidia_smi` is just the **first reference source**; +everything vendor-specific lives inside its builder, not in the schema. + +``` +setup → [start power sidecar] → warmup → PERFORMANCE phase → teardown → [stop sidecar] + └── trace sliced to this window ──┘ +``` + +The sidecar is best-effort: a broken collector writes `status:"failed"` and the +benchmark completes normally. It can never fail or perturb the run. + +## Config + +The vendor-neutral core is tiny; per-source settings go in `options` (each +source documents its keys): + +```yaml +settings: + power: + source: nvidia_smi # nvidia_smi | prometheus | command | + interval_s: 1.0 # sampling interval (default 1.0) + options: {} # source-specific settings + # advanced (hidden, sane defaults): value_kind, env +``` + +When `source` is unset the whole block is omitted from the serialized config — +disabling power leaves zero footprint. + +### Built-in sources + +```yaml +# NVIDIA GPUs on the load-gen host (the reference source) +power: + source: nvidia_smi + options: { gpu_indices: [0, 1, 2, 3] } # optional; default = all GPUs + +# Remote GPUs / cluster: scrape a server-side exporter (DCGM, etc.) +power: + source: prometheus + options: + url: "http://gpu-node:9400" + query: "DCGM_FI_DEV_POWER_USAGE" # watts; one series per GPU/instance + +# Anything else, zero Python: a program that prints canonical JSONL +power: + source: command + options: + argv: ["ssh", "gpu-node", "my-power-reader", "--interval", "1"] +``` + +The `command` argv must print **one JSONL sample per line**: +`{"ts": , "value": , "label": ""}`. + +## Adding your own source (plugin) + +No core edits, no entry-points — register a builder and pass it settings via +`options`: + +```python +from inference_endpoint.power import power_source, ResolvedSource + +@power_source("redfish") +def _build(cfg): + return ResolvedSource( + argv=["redfish-power", cfg.options["bmc"], "--interval", str(cfg.interval_s)], + fmt="jsonl", value_kind="power_w", + ts_field="ts", value_field="value", label_field="label", csv_header=False, + ) +``` + +```yaml +power: { source: redfish, options: { bmc: "10.0.0.5" } } +``` + +(For non-Python sources, prefer the built-in `command` source — the process +boundary is the plugin API.) + +## Output: `power.json` + +```json +{ + "schema_version": "1.0", + "status": "ok", + "window": {"start_epoch_s": ..., "end_epoch_s": ..., "duration_s": 305.0, "basis": "performance_phase"}, + "totals": { + "energy_j": 1045212.2, + "mean_power_w": 3426.9, + "output_tokens": 4800915, + "energy_per_output_token_j": 0.2177, + "consistent_with_window": true + }, + "sources": [ + {"label": "gpu0", "value_kind": "power_w", "energy_j": 130651.5, + "power_w": {"mean": 428.4, "p50": 410.0, "p95": 658.0, "max": 879.1}, "sample_count": 305} + ], + "provenance": {"command": [...], "interval_s": 1.0, "samples_in_window": 305, "samples_dropped": 0} +} +``` + +### Energy-per-output-token consistency + +`energy_per_output_token_j` is emitted **only** when the token count and the +energy share the same window (single tracked performance phase). Otherwise it is +`null` with an explanatory note — preventing a windowed-energy ÷ global-tokens +mismatch. + +## Remote GPUs / multi-node + +The sidecar runs **where its command runs** (the load-gen host). For remote +GPUs, either point `prometheus` at a server-side exporter (recommended) or use a +`command` source. For an ssh-wrapped command, use `ssh -tt` so the remote +sampler receives the stop signal. Clocks are assumed NTP-synced. + +## Caveats / scope + +- `power_w` energy uses a trapezoid integral with no edge interpolation + (≤ one sample-interval bias at each window edge — negligible at 1 s / 300 s+). +- The window is the full performance phase (includes drain). +- v1 collects a single source; multi-source totals / scope-aware double-count + guards are future work. diff --git a/examples/11_Power_Monitoring/online_with_power.yaml b/examples/11_Power_Monitoring/online_with_power.yaml new file mode 100644 index 000000000..ded51cc1c --- /dev/null +++ b/examples/11_Power_Monitoring/online_with_power.yaml @@ -0,0 +1,47 @@ +# Online benchmark with vendor-agnostic power monitoring. +# See docs/POWER_MONITORING.md. Power telemetry is collected during the +# performance phase and written to /power.json + a section in +# report.txt. A broken collector never fails the benchmark. +name: "online-power-example" +version: "1.0" +type: "online" + +model_params: + name: "meta-llama/Llama-3.1-8B-Instruct" + temperature: 0.0 + max_new_tokens: 128 + +datasets: + - name: cnn_dailymail::llama3_8b + type: "performance" + samples: 2000 + parser: + input: prompt + +settings: + runtime: + min_duration_ms: 300000 + max_duration_ms: 1800000 + + load_pattern: + type: "concurrency" + target_concurrency: 64 + + client: + num_workers: 4 + + # --- Power monitoring ----------------------------------------------------- + # Built-in nvidia-smi sidecar on the load-generator host. For a remote GPU + # node, switch to the `prometheus` source (DCGM exporter) or an ssh-wrapped + # `command` source — see docs/POWER_MONITORING.md. + power: + source: nvidia_smi + interval_s: 1.0 + # options: { gpu_indices: [0, 1, 2, 3] } # default: all visible GPUs + +endpoint_config: + endpoints: + - "http://localhost:8000" + api_key: null + +report_dir: logs/power_example diff --git a/src/inference_endpoint/commands/benchmark/execute.py b/src/inference_endpoint/commands/benchmark/execute.py index 6df97594a..8b725bb79 100644 --- a/src/inference_endpoint/commands/benchmark/execute.py +++ b/src/inference_endpoint/commands/benchmark/execute.py @@ -100,6 +100,13 @@ SessionResult, ) from inference_endpoint.metrics.report import Report +from inference_endpoint.power import ( + PowerCollector, + build_power_report, + write_power_section, +) +from inference_endpoint.power.window import disabled_report +from inference_endpoint.utils import monotime_to_datetime transformers_logging.set_verbosity_error() @@ -148,6 +155,9 @@ class BenchmarkResult: # settings.profiling.engine is set; None otherwise. Rendered into # report.txt and a sibling profiling.json by finalize_benchmark. profiling: dict[str, Any] | None = None + # Power-telemetry sidecar when settings.power.source is set; None otherwise. + # finalize_benchmark windows its trace and writes a sibling power.json. + power: PowerCollector | None = None @dataclass @@ -896,6 +906,15 @@ def _on_phase_start(phase: PhaseConfig) -> None: ) profile_starts.append(rec) + # Power telemetry sidecar (best-effort; never fails the run). Started + # before session.run so it is warm; finalize_benchmark windows its + # trace to the performance phase. + power_cfg = config.settings.power + power_collector: PowerCollector | None = None + if power_cfg.source is not None: + power_collector = PowerCollector(power_cfg, ctx.report_dir / "power") + power_collector.start() + loop.add_signal_handler(signal.SIGINT, session.stop) try: result = await session.run(phases, on_phase_start=_on_phase_start) @@ -926,6 +945,9 @@ def _on_phase_start(phase: PhaseConfig) -> None: rec["error"] or rec["status"], ) profile_stops.append(rec) + # Stop the power sidecar (best-effort; bounded SIGTERM->SIGKILL). + if power_collector is not None: + power_collector.stop() logger.info("Cleaning up...") try: if http_client: @@ -1014,6 +1036,7 @@ def _on_phase_start(phase: PhaseConfig) -> None: report=report, tmpfs_dir=tmpfs_dir, profiling=profiling_payload, + power=power_collector, ) @@ -1066,6 +1089,51 @@ def _salvage_tmpfs(report_dir: Path, tmpfs_dir: Path) -> None: logger.debug(f"Copied {src_events} -> {dst_events}") +def _build_power_payload( + ctx: BenchmarkContext, bench: BenchmarkResult +) -> dict[str, Any] | None: + """Window the power trace to the performance phase and build the JSON payload. + + Best-effort: any failure yields a ``status:"failed"`` payload, never an + exception that would abort finalization. + """ + pc = bench.power + if pc is None: + return None + try: + if pc.resolved is None: + return disabled_report(pc.error or "power collector failed to start") + perf = bench.session.perf_results + if not perf: + return disabled_report("no performance phase to window power against") + # Phase windows are monotonic_ns; convert to wall-clock epoch (the clock + # power samples are stamped in) via the shared monotonic->wall anchor. + start_epoch = monotime_to_datetime(perf[0].start_time_ns).timestamp() + end_epoch = monotime_to_datetime(perf[-1].end_time_ns).timestamp() + output_tokens: int | None = None + if bench.report is not None: + osl = bench.report.output_sequence_lengths or {} + total = osl.get("total") + output_tokens = int(total) if total else None + return build_power_report( + resolved=pc.resolved, + trace_path=pc.trace_path, + window_start_epoch_s=start_epoch, + window_end_epoch_s=end_epoch, + output_tokens=output_tokens, + token_window_basis="performance_phase_tracked", + # Report tokens are scoped to the (single) tracked perf phase, so + # they share the energy window. Multiple perf phases would span gaps. + consistent_with_window=len(perf) == 1, + collector_status=pc.status, + collector_error=pc.error, + interval_s=pc.cfg.interval_s, + ) + except Exception as e: # noqa: BLE001 — power must never break finalization. + logger.warning("power: failed to build report: %s", e) + return disabled_report(f"power report build failed: {e}") + + def finalize_benchmark(ctx: BenchmarkContext, bench: BenchmarkResult) -> None: """Score accuracy, aggregate results, write JSON.""" config = ctx.config @@ -1073,6 +1141,8 @@ def finalize_benchmark(ctx: BenchmarkContext, bench: BenchmarkResult) -> None: collector = bench.collector report = bench.report + power_payload = _build_power_payload(ctx, bench) + # Display report if available (from MetricsAggregator pub/sub snapshot). # result_summary.json is the self-complete machine-readable report (carries # qps/tps/seeds via Report.to_json); report.txt is the full human-readable @@ -1086,6 +1156,8 @@ def finalize_benchmark(ctx: BenchmarkContext, bench: BenchmarkResult) -> None: report.display(fn=lambda s: print(s, file=f)) if bench.profiling is not None: _write_profiling_section(f, bench.profiling) + if power_payload is not None: + write_power_section(f, power_payload) logger.info("Report written to %s", report_txt) # Sibling profiling.json — kept separate so Report stays a pure @@ -1095,6 +1167,10 @@ def finalize_benchmark(ctx: BenchmarkContext, bench: BenchmarkResult) -> None: json.dumps(bench.profiling, indent=2) ) + # Sibling power.json — same rationale: Report stays snapshot-pure. + if power_payload is not None: + (ctx.report_dir / "power.json").write_text(json.dumps(power_payload, indent=2)) + # Write scoring artifacts + copy event log from tmpfs to disk _write_scoring_artifacts(ctx, result, bench.tmpfs_dir) diff --git a/src/inference_endpoint/config/schema.py b/src/inference_endpoint/config/schema.py index 5ad7b9c52..2125fe282 100644 --- a/src/inference_endpoint/config/schema.py +++ b/src/inference_endpoint/config/schema.py @@ -702,6 +702,70 @@ def _validate_url_scheme(cls, v: list[str] | None) -> list[str] | None: return v +@cyclopts.Parameter(name="*") +class PowerConfig(BaseModel): + """Vendor-agnostic power/energy monitoring. + + When ``source`` is set, a sidecar process streams power samples to a trace + file during the performance phase; at finalization the trace is sliced to + the measurement window, integrated into energy, and written to a sibling + ``power.json``. A broken collector never fails or skews the benchmark. + + The sidecar runs *where the command runs* — i.e. the load-generator host. + For remote GPUs either point ``prometheus`` at a server-side DCGM exporter + or supply an ssh-wrapped ``command``; clocks are assumed NTP-synced. + """ + + model_config = ConfigDict(extra="forbid", frozen=True) + + # Vendor-neutral core. Source-specific settings live in `options`, owned and + # validated by the registered source builder (see power/sources.py) — the + # schema knows nothing about GPUs, Prometheus, etc. + source: Annotated[ + str | None, + cyclopts.Parameter( + alias="--power", + help="Power source: nvidia_smi | prometheus | command | ", + ), + ] = Field( + None, + description="Power telemetry source name; None disables monitoring. " + "Built-ins: nvidia_smi, prometheus, command. Register custom sources " + "with @power_source('name').", + ) + interval_s: Annotated[ + float, + cyclopts.Parameter(alias="--power-interval"), + ] = Field(1.0, gt=0, description="Sampling interval in seconds") + options: dict[str, Any] = Field( + default_factory=dict, + description="Source-specific settings (each source documents its keys). " + "e.g. nvidia_smi: {gpu_indices}; prometheus: {url, query}; " + "command: {argv}. command argv must print one JSONL sample per line: " + '{"ts": , "value": , "label": ""}.', + ) + + # --- Advanced (hidden from CLI; rarely needed, sane defaults) ------------ + value_kind: Annotated[ + Literal["power_w", "energy_j"], cyclopts.Parameter(show=False) + ] = Field( + "power_w", + description="Sample semantics: instantaneous watts or a cumulative joule counter", + ) + env: Annotated[dict[str, str], cyclopts.Parameter(show=False)] = Field( + default_factory=dict, description="Extra environment for the sidecar process" + ) + + @model_serializer(mode="wrap") + def _serialize(self, handler: SerializerFunctionWrapHandler) -> dict[str, Any]: + # When disabled (source=None), collapse to an empty mapping so a + # turned-off feature leaves no footprint in the serialized config. + data = handler(self) + if self.source is None: + return {} + return data + + @cyclopts.Parameter(name="*") class Settings(BaseModel): """Test settings.""" @@ -717,6 +781,7 @@ class Settings(BaseModel): ) warmup: WarmupConfig = Field(default_factory=WarmupConfig) profiling: ProfilingConfig = Field(default_factory=ProfilingConfig) + power: PowerConfig = Field(default_factory=PowerConfig) class OfflineSettings(Settings): diff --git a/src/inference_endpoint/config/templates/concurrency_template_full.yaml b/src/inference_endpoint/config/templates/concurrency_template_full.yaml index 382e3cc00..98e09ca6f 100644 --- a/src/inference_endpoint/config/templates/concurrency_template_full.yaml +++ b/src/inference_endpoint/config/templates/concurrency_template_full.yaml @@ -90,6 +90,12 @@ settings: profiling: engine: null # Profile the named inference engine around the performance phase | options: vllm urls: null # URL(s) the profiler start/stop triggers are derived from. When None, derived from endpoint_config.endpoints instead. Use when the profiler admin endpoint differs from the inference endpoint. + power: + source: null # Power telemetry source name; None disables monitoring. Built-ins: nvidia_smi, prometheus, command. Register custom sources with @power_source('name'). + interval_s: 1.0 # Sampling interval in seconds + options: {} # Source-specific settings (each source documents its keys). e.g. nvidia_smi: {gpu_indices}; prometheus: {url, query}; command: {argv}. command argv must print one JSONL sample per line: {"ts": , "value": , "label": ""}. + value_kind: power_w # Sample semantics: instantaneous watts or a cumulative joule counter | options: power_w, energy_j + env: {} # Extra environment for the sidecar process endpoint_config: endpoints: # Endpoint URL(s). Must include scheme, e.g. 'http://host:port'. - http://localhost:8000 diff --git a/src/inference_endpoint/config/templates/offline_template_full.yaml b/src/inference_endpoint/config/templates/offline_template_full.yaml index 6a2998812..e1d88b9cb 100644 --- a/src/inference_endpoint/config/templates/offline_template_full.yaml +++ b/src/inference_endpoint/config/templates/offline_template_full.yaml @@ -90,6 +90,12 @@ settings: profiling: engine: null # Profile the named inference engine around the performance phase | options: vllm urls: null # URL(s) the profiler start/stop triggers are derived from. When None, derived from endpoint_config.endpoints instead. Use when the profiler admin endpoint differs from the inference endpoint. + power: + source: null # Power telemetry source name; None disables monitoring. Built-ins: nvidia_smi, prometheus, command. Register custom sources with @power_source('name'). + interval_s: 1.0 # Sampling interval in seconds + options: {} # Source-specific settings (each source documents its keys). e.g. nvidia_smi: {gpu_indices}; prometheus: {url, query}; command: {argv}. command argv must print one JSONL sample per line: {"ts": , "value": , "label": ""}. + value_kind: power_w # Sample semantics: instantaneous watts or a cumulative joule counter | options: power_w, energy_j + env: {} # Extra environment for the sidecar process endpoint_config: endpoints: # Endpoint URL(s). Must include scheme, e.g. 'http://host:port'. - http://localhost:8000 diff --git a/src/inference_endpoint/config/templates/online_template_full.yaml b/src/inference_endpoint/config/templates/online_template_full.yaml index 6c74b7a6a..b3257806b 100644 --- a/src/inference_endpoint/config/templates/online_template_full.yaml +++ b/src/inference_endpoint/config/templates/online_template_full.yaml @@ -91,6 +91,12 @@ settings: profiling: engine: null # Profile the named inference engine around the performance phase | options: vllm urls: null # URL(s) the profiler start/stop triggers are derived from. When None, derived from endpoint_config.endpoints instead. Use when the profiler admin endpoint differs from the inference endpoint. + power: + source: null # Power telemetry source name; None disables monitoring. Built-ins: nvidia_smi, prometheus, command. Register custom sources with @power_source('name'). + interval_s: 1.0 # Sampling interval in seconds + options: {} # Source-specific settings (each source documents its keys). e.g. nvidia_smi: {gpu_indices}; prometheus: {url, query}; command: {argv}. command argv must print one JSONL sample per line: {"ts": , "value": , "label": ""}. + value_kind: power_w # Sample semantics: instantaneous watts or a cumulative joule counter | options: power_w, energy_j + env: {} # Extra environment for the sidecar process endpoint_config: endpoints: # Endpoint URL(s). Must include scheme, e.g. 'http://host:port'. - http://localhost:8000 diff --git a/src/inference_endpoint/power/__init__.py b/src/inference_endpoint/power/__init__.py new file mode 100644 index 000000000..3f9073b6e --- /dev/null +++ b/src/inference_endpoint/power/__init__.py @@ -0,0 +1,42 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Vendor-agnostic power/energy monitoring. + +A benchmark run can spawn a *sidecar* that streams power telemetry to a trace +file while the performance phase runs; at finalization the trace is sliced to +the measurement window, integrated into energy, and written to a sibling +``power.json`` (the ``Report`` is never mutated, mirroring the profiling +precedent). + +The agnosticism boundary is the *process boundary*: any source that can print +one sample per line (``nvidia-smi``, a Prometheus/DCGM exporter, an +IPMI/redfish/PDU scraper, a cloud metrics CLI, or a user script) plugs in via +the ``command`` source with a field mapping — no Python, no core edits. +""" + +from inference_endpoint.power.collector import PowerCollector +from inference_endpoint.power.render import write_power_section +from inference_endpoint.power.sources import ResolvedSource, power_source, resolve +from inference_endpoint.power.window import build_power_report + +__all__ = [ + "PowerCollector", + "ResolvedSource", + "build_power_report", + "power_source", # decorator: register a custom source + "resolve", + "write_power_section", +] diff --git a/src/inference_endpoint/power/collector.py b/src/inference_endpoint/power/collector.py new file mode 100644 index 000000000..55bb947af --- /dev/null +++ b/src/inference_endpoint/power/collector.py @@ -0,0 +1,153 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Power-telemetry sidecar lifecycle. + +A plain ``subprocess.Popen`` (NOT the ServiceLauncher — we must never abort the +benchmark if the collector misbehaves) that streams a source's stdout into a +trace file. It runs in its own process group so a hung child and any +grandchildren are torn down together. Every method swallows its errors: power +monitoring is strictly best-effort. +""" + +from __future__ import annotations + +import ctypes +import logging +import os +import signal +import subprocess +import sys +from pathlib import Path +from typing import IO + +from inference_endpoint.config.schema import PowerConfig +from inference_endpoint.power.sources import ResolvedSource, resolve + +logger = logging.getLogger(__name__) + +# Grace period after SIGTERM before the sidecar is force-killed. Not a config +# knob — sampling sidecars die fast; tuning this has no practical value. +_STOP_GRACE_S = 5.0 +_PR_SET_PDEATHSIG = 1 + + +def _die_with_parent() -> None: + """preexec_fn: on Linux, request SIGKILL if the parent process dies. + + Backstop for the case where the parent is SIGKILL'd/OOM-killed before + stop() can run — otherwise a detached sidecar would poll forever. + """ + if sys.platform == "linux": + try: + ctypes.CDLL("libc.so.6", use_errno=True).prctl( + _PR_SET_PDEATHSIG, signal.SIGKILL + ) + except Exception: # noqa: BLE001 — best-effort; never block the child. + pass + + +class PowerCollector: + """Start/stop a sidecar that writes a power trace to ``trace_path``.""" + + def __init__(self, cfg: PowerConfig, out_dir: Path) -> None: + self.cfg = cfg + self.out_dir = out_dir + self.trace_path = out_dir / "power_trace.log" + self.stderr_path = out_dir / "power_collect.stderr.log" + self.resolved: ResolvedSource | None = None + self.status = "disabled" + self.error: str | None = None + self._proc: subprocess.Popen[bytes] | None = None + self._trace_fh: IO[bytes] | None = None + self._stderr_fh: IO[bytes] | None = None + + def start(self) -> None: + """Launch the sidecar. Records failure in ``status`` instead of raising.""" + try: + self.out_dir.mkdir(parents=True, exist_ok=True) + self.resolved = resolve(self.cfg) + env = {**os.environ, **self.cfg.env} + self._trace_fh = self.trace_path.open("wb") + self._stderr_fh = self.stderr_path.open("wb") + self._proc = subprocess.Popen( + self.resolved.argv, + stdout=self._trace_fh, + stderr=self._stderr_fh, + env=env, + # Own session/process-group so stop() can signal the whole tree. + start_new_session=True, + # Backstop: die if the parent is SIGKILL'd before stop() runs. + preexec_fn=_die_with_parent, # noqa: PLW1509 (intentional, Linux) + ) + self.status = "ok" + logger.info( + "power: started sidecar pid=%d -> %s (%s)", + self._proc.pid, + self.trace_path, + " ".join(self.resolved.argv), + ) + except Exception as e: # noqa: BLE001 — collector startup must never abort the run. + self.status = "failed" + self.error = f"start failed: {e}" + logger.warning("power: %s", self.error) + self._close_files() + + def stop(self) -> None: + """Terminate the sidecar (SIGTERM → grace → SIGKILL). Never raises.""" + proc = self._proc + if proc is None: + return + try: + if proc.poll() is None: + self._signal_group(proc, signal.SIGTERM) + try: + proc.wait(timeout=_STOP_GRACE_S) + except subprocess.TimeoutExpired: + logger.warning("power: sidecar did not exit; sending SIGKILL") + self._signal_group(proc, signal.SIGKILL) + try: + proc.wait(timeout=_STOP_GRACE_S) + except subprocess.TimeoutExpired: + self.error = "sidecar would not die" + rc = proc.returncode + # nvidia-smi/poll loops are killed by us, so a signal exit is expected. + if rc not in (0, None, -signal.SIGTERM, -signal.SIGKILL) and not self.error: + self.error = f"sidecar exited rc={rc}" + except Exception as e: # noqa: BLE001 — stop must never abort finalization. + self.error = f"stop failed: {e}" + logger.warning("power: %s", self.error) + finally: + self._close_files() + + @staticmethod + def _signal_group(proc: subprocess.Popen[bytes], sig: int) -> None: + try: + if hasattr(os, "killpg") and hasattr(os, "getpgid"): + os.killpg(os.getpgid(proc.pid), sig) + else: # Windows: no process groups + proc.send_signal(sig) + except (ProcessLookupError, PermissionError): + proc.send_signal(sig) + + def _close_files(self) -> None: + for fh in (self._trace_fh, self._stderr_fh): + try: + if fh is not None: + fh.close() + except Exception: # noqa: BLE001 + pass + self._trace_fh = None + self._stderr_fh = None diff --git a/src/inference_endpoint/power/parse.py b/src/inference_endpoint/power/parse.py new file mode 100644 index 000000000..54abb9499 --- /dev/null +++ b/src/inference_endpoint/power/parse.py @@ -0,0 +1,133 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Parse a power trace file into canonical samples. + +Tolerant by design: malformed lines are dropped and counted, never raised, so a +flaky collector degrades gracefully instead of failing the benchmark. +""" + +from __future__ import annotations + +import csv +import json +from dataclasses import dataclass +from datetime import datetime +from pathlib import Path + +from inference_endpoint.power.sources import ResolvedSource + + +@dataclass(frozen=True) +class PowerSample: + ts_epoch_s: float + label: str + value: float + + +@dataclass(frozen=True) +class ParseResult: + samples: list[PowerSample] + dropped: int + + +def _parse_ts(raw: str) -> float: + """Accept epoch seconds (float) or the nvidia-smi ``YYYY/MM/DD HH:MM:SS.mmm``.""" + raw = raw.strip() + try: + return float(raw) + except ValueError: + pass + for fmt in ("%Y/%m/%d %H:%M:%S.%f", "%Y/%m/%d %H:%M:%S"): + try: + return datetime.strptime(raw, fmt).timestamp() + except ValueError: + continue + raise ValueError(f"unparseable timestamp: {raw!r}") + + +def parse_trace(path: Path, src: ResolvedSource) -> ParseResult: + """Read ``path`` per ``src``'s format/field-mapping into canonical samples.""" + if not path.exists(): + return ParseResult(samples=[], dropped=0) + if src.fmt == "csv": + return _parse_csv(path, src) + return _parse_jsonl(path, src) + + +def _parse_jsonl(path: Path, src: ResolvedSource) -> ParseResult: + samples: list[PowerSample] = [] + dropped = 0 + with path.open() as f: + for line in f: + line = line.strip() + if not line: + continue + try: + obj = json.loads(line) + if not isinstance(obj, dict): + dropped += 1 + continue + ts = _parse_ts(str(obj[src.ts_field])) + value = float(obj[src.value_field]) + label = ( + str(obj.get(src.label_field, "default")) + if src.label_field + else "default" + ) + samples.append(PowerSample(ts, label, value)) + except (ValueError, KeyError, TypeError): + dropped += 1 + return ParseResult(samples=samples, dropped=dropped) + + +def _parse_csv(path: Path, src: ResolvedSource) -> ParseResult: + samples: list[PowerSample] = [] + dropped = 0 + # Field mapping is by integer index for headerless CSV (nvidia-smi), else by + # column name. Split paths so the row type stays concrete for each. + with path.open() as f: + if src.csv_header: + for named in csv.DictReader(f): + try: + ts = _parse_ts(named[src.ts_field]) + value = float(named[src.value_field]) + label = ( + str(named[src.label_field]) if src.label_field else "default" + ) + samples.append(PowerSample(ts, label.strip(), value)) + except (ValueError, KeyError, TypeError): + dropped += 1 + else: + try: + ts_i = int(src.ts_field) + val_i = int(src.value_field) + lbl_i = int(src.label_field) if src.label_field else None + except (ValueError, TypeError): + return ParseResult(samples=[], dropped=0) + for row in csv.reader(f): + try: + ts = _parse_ts(row[ts_i]) + value = float(row[val_i]) + label = row[lbl_i].strip() if lbl_i is not None else "default" + # Prefix bare numeric labels only when the source asked for it + # (e.g. nvidia-smi GPU index 0 -> "gpu0"); generic sources + # keep their labels verbatim. + if src.label_prefix and label.isdigit(): + label = f"{src.label_prefix}{label}" + samples.append(PowerSample(ts, label.strip(), value)) + except (ValueError, IndexError, TypeError): + dropped += 1 + return ParseResult(samples=samples, dropped=dropped) diff --git a/src/inference_endpoint/power/prom_poll.py b/src/inference_endpoint/power/prom_poll.py new file mode 100644 index 000000000..1c5d0a840 --- /dev/null +++ b/src/inference_endpoint/power/prom_poll.py @@ -0,0 +1,92 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tiny Prometheus instant-query poller (stdlib only). + +Run as a sidecar: every ``--interval`` seconds it issues ``query`` and prints one +canonical JSONL sample per returned series to stdout: + + {"ts": , "value": , "label": ""} + +One transient HTTP failure must not kill the loop — failures are skipped. The +series ``label`` is the metric's full label set so multi-GPU/exporter results +stay distinct (and energy is not double-counted across series). +""" + +from __future__ import annotations + +import argparse +import json +import sys +import time +import urllib.parse +import urllib.request + + +def _series_label(metric: dict[str, str] | None) -> str: + """Stable label from a Prometheus metric's labels (gpu/instance preferred).""" + if not metric: + return "default" + for key in ("gpu", "device", "instance", "__name__"): + if key in metric: + return str(metric[key]) + return ",".join(f"{k}={v}" for k, v in sorted(metric.items())) + + +def _poll_once(url: str, query: str, timeout: float) -> None: + full = f"{url.rstrip('/')}/api/v1/query?" + urllib.parse.urlencode({"query": query}) + with urllib.request.urlopen(full, timeout=timeout) as resp: # noqa: S310 (user URL) + body = json.loads(resp.read().decode()) + if body.get("status") != "success": + return + now = time.time() + for series in body.get("data", {}).get("result", []): + value_pair = series.get("value") + if not value_pair or len(value_pair) != 2: + continue + try: + value = float(value_pair[1]) + except (ValueError, TypeError): + continue + sample = { + "ts": now, + "value": value, + "label": _series_label(series.get("metric", {})), + } + sys.stdout.write(json.dumps(sample) + "\n") + sys.stdout.flush() + + +def main() -> None: + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--url", required=True) + p.add_argument("--query", required=True) + p.add_argument("--interval", type=float, default=1.0) + p.add_argument("--timeout", type=float, default=None) + args = p.parse_args() + timeout = args.timeout if args.timeout is not None else max(args.interval, 1.0) + while True: + start = time.monotonic() + try: + _poll_once(args.url, args.query, timeout) + except Exception: # noqa: BLE001 — a poll failure must never kill the loop. + pass + # Keep cadence steady regardless of request latency. + elapsed = time.monotonic() - start + time.sleep(max(0.0, args.interval - elapsed)) + + +if __name__ == "__main__": + main() diff --git a/src/inference_endpoint/power/render.py b/src/inference_endpoint/power/render.py new file mode 100644 index 000000000..0aa4f79ad --- /dev/null +++ b/src/inference_endpoint/power/render.py @@ -0,0 +1,68 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Render the Power section into report.txt (mirrors _write_profiling_section).""" + +from __future__ import annotations + +from typing import Any, TextIO + + +def _fmt(v: Any, suffix: str = "") -> str: + if v is None: + return "n/a" + if isinstance(v, float): + return f"{v:,.3f}{suffix}" + return f"{v}{suffix}" + + +def write_power_section(f: TextIO, power: dict[str, Any]) -> None: + """Append a human-readable Power section to report.txt.""" + f.write("\n------------------- Power -------------------\n") + f.write(f"Status: {power.get('status', 'unknown')}\n") + if power.get("error"): + f.write(f"Error: {power['error']}\n") + return + + win = power.get("window", {}) + f.write( + f"Window: {_fmt(win.get('duration_s'), ' s')} " + f"(basis: {win.get('basis', '?')})\n" + ) + + totals = power.get("totals", {}) + f.write(f"Total energy: {_fmt(totals.get('energy_j'), ' J')}\n") + f.write(f"Mean power: {_fmt(totals.get('mean_power_w'), ' W')}\n") + epot = totals.get("energy_per_output_token_j") + if epot is not None: + f.write(f"Energy/output token: {_fmt(epot, ' J')}\n") + elif totals.get("energy_per_output_token_note"): + f.write( + f"Energy/output token: n/a ({totals['energy_per_output_token_note']})\n" + ) + + for src in power.get("sources", []): + pw = src.get("power_w") or {} + f.write( + f" [{src.get('label')}] energy={_fmt(src.get('energy_j'), ' J')} " + f"mean={_fmt(pw.get('mean'), ' W')} " + f"p95={_fmt(pw.get('p95'), ' W')} " + f"max={_fmt(pw.get('max'), ' W')} " + f"(n={src.get('sample_count')})\n" + ) + + prov = power.get("provenance", {}) + if prov.get("samples_dropped"): + f.write(f" (dropped {prov['samples_dropped']} malformed samples)\n") diff --git a/src/inference_endpoint/power/sources.py b/src/inference_endpoint/power/sources.py new file mode 100644 index 000000000..a23906073 --- /dev/null +++ b/src/inference_endpoint/power/sources.py @@ -0,0 +1,163 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Power source registry. + +A source builder turns a :class:`PowerConfig` into a :class:`ResolvedSource` +(a sidecar argv + how to read what it prints). Built-ins register themselves +below; users add their own with the ``@power_source`` decorator and feed it +arbitrary settings via ``cfg.options`` — no core edits, no entry-points: + + from inference_endpoint.power import power_source, ResolvedSource + + @power_source("my_pdu") + def _build(cfg): + return ResolvedSource( + argv=["my-pdu-reader", cfg.options["rack"]], + fmt="jsonl", value_kind="power_w", + ts_field="ts", value_field="value", label_field="label", + csv_header=False, + ) + +For non-Python sources, the built-in ``command`` source runs any program that +prints the canonical JSONL contract — the process boundary is the plugin API. +""" + +from __future__ import annotations + +import sys +from collections.abc import Callable +from dataclasses import dataclass +from typing import Any + +from inference_endpoint.config.schema import PowerConfig + + +@dataclass(frozen=True) +class ResolvedSource: + """A runnable sidecar plus how to read what it prints (all internal).""" + + argv: list[str] + fmt: str # "csv" (nvidia-smi only) | "jsonl" + value_kind: str # "power_w" | "energy_j" + # Field/column mapping: integer-index strings for headerless CSV, key names + # for JSONL. + ts_field: str + value_field: str + label_field: str | None + csv_header: bool + # If set, bare numeric labels are prefixed with it (nvidia GPU 0 -> "gpu0"). + label_prefix: str | None = None + + +PowerSourceBuilder = Callable[[PowerConfig], ResolvedSource] +_REGISTRY: dict[str, PowerSourceBuilder] = {} + + +def power_source(name: str) -> Callable[[PowerSourceBuilder], PowerSourceBuilder]: + """Register a builder under ``name`` (usable as ``power.source``).""" + + def decorator(fn: PowerSourceBuilder) -> PowerSourceBuilder: + _REGISTRY[name] = fn + return fn + + return decorator + + +def resolve(cfg: PowerConfig) -> ResolvedSource: + """Look up and run the builder for ``cfg.source``. Caller guarantees enabled.""" + builder = _REGISTRY.get(cfg.source or "") + if builder is None: + raise ValueError( + f"unknown power source {cfg.source!r}; available: {sorted(_REGISTRY)}" + ) + return builder(cfg) + + +def _jsonl_source(argv: list[str], value_kind: str) -> ResolvedSource: + """A source emitting the canonical JSONL contract {ts, value, label}.""" + return ResolvedSource( + argv=argv, + fmt="jsonl", + value_kind=value_kind, + ts_field="ts", + value_field="value", + label_field="label", + csv_header=False, + ) + + +def _require(cfg: PowerConfig, key: str, source: str) -> Any: + """Fetch a required ``options`` key or raise a clear error.""" + if key not in cfg.options: + raise ValueError(f"power source {source!r} requires options.{key}") + return cfg.options[key] + + +@power_source("nvidia_smi") +def _nvidia_smi(cfg: PowerConfig) -> ResolvedSource: + """First reference source. NVIDIA-specific bits live here, not in the schema. + + options: gpu_indices (list[int], optional; default all visible GPUs). + """ + argv = [ + # stdbuf line-buffers so tail samples aren't lost to block buffering + # when stdout is a file and the sidecar is killed at phase end. + "stdbuf", + "-oL", + "nvidia-smi", + "--query-gpu=timestamp,index,power.draw", + "--format=csv,noheader,nounits", + "-lms", + str(int(cfg.interval_s * 1000)), + ] + gpu_indices = cfg.options.get("gpu_indices") + if gpu_indices: + argv += ["-i", ",".join(str(i) for i in gpu_indices)] + # Headerless CSV columns: timestamp(nvidia fmt), index, power.draw(W). + return ResolvedSource( + argv=argv, + fmt="csv", + value_kind="power_w", + ts_field="0", + value_field="2", + label_field="1", + csv_header=False, + label_prefix="gpu", + ) + + +@power_source("prometheus") +def _prometheus(cfg: PowerConfig) -> ResolvedSource: + """options: url (str), query (PromQL returning watts/joules).""" + argv = [ + sys.executable, + "-m", + "inference_endpoint.power.prom_poll", + "--url", + str(_require(cfg, "url", "prometheus")), + "--query", + str(_require(cfg, "query", "prometheus")), + "--interval", + str(cfg.interval_s), + ] + return _jsonl_source(argv, cfg.value_kind) + + +@power_source("command") +def _command(cfg: PowerConfig) -> ResolvedSource: + """options: argv (list[str]) that prints one canonical JSONL sample per line.""" + argv = _require(cfg, "argv", "command") + return _jsonl_source(list(argv), cfg.value_kind) diff --git a/src/inference_endpoint/power/window.py b/src/inference_endpoint/power/window.py new file mode 100644 index 000000000..330ca8760 --- /dev/null +++ b/src/inference_endpoint/power/window.py @@ -0,0 +1,185 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Slice a power trace to the measurement window and integrate it into energy. + +Produces the ``power.json`` payload: per-series watt statistics, energy over the +window, and an energy-per-output-token figure that is emitted **only** when the +token count and the energy share the same window (otherwise the denominators +would silently mix — the single biggest correctness trap here). +""" + +from __future__ import annotations + +from typing import Any + +from inference_endpoint.power.parse import PowerSample, parse_trace +from inference_endpoint.power.sources import ResolvedSource + + +def _stats(values: list[float]) -> dict[str, float | int]: + if not values: + return {"n": 0} + s = sorted(values) + n = len(s) + return { + "n": n, + "mean": sum(s) / n, + "p50": s[n // 2], + "p95": s[min(n - 1, int(n * 0.95))], + "max": s[-1], + "min": s[0], + } + + +def _energy_j(samples: list[PowerSample], value_kind: str) -> float | None: + """Energy in joules over the samples (already window-filtered, time-sorted). + + ``power_w``: trapezoid integral of watts over seconds. + ``energy_j``: cumulative counter delta (last - first). + Returns None when there is too little data to bracket the window. + """ + if len(samples) < 2: + return None + if value_kind == "energy_j": + # Sum non-negative adjacent deltas; a mid-run counter reset (negative + # delta) makes the whole total unusable, even if last-first is positive. + total = 0.0 + for a, b in zip(samples, samples[1:], strict=False): + delta = b.value - a.value + if delta < 0: + return None # counter reset → unusable + total += delta + return total + total = 0.0 + for a, b in zip(samples, samples[1:], strict=False): + dt = b.ts_epoch_s - a.ts_epoch_s + if dt <= 0: + continue + total += 0.5 * (a.value + b.value) * dt # ponytail: no edge interpolation + return total + + +def build_power_report( + *, + resolved: ResolvedSource, + trace_path: Any, + window_start_epoch_s: float, + window_end_epoch_s: float, + output_tokens: int | None, + token_window_basis: str, + consistent_with_window: bool, + collector_status: str, + collector_error: str | None, + interval_s: float, +) -> dict[str, Any]: + """Assemble the ``power.json`` dict. Pure + total-failure tolerant.""" + parsed = parse_trace(trace_path, resolved) + in_window = sorted( + ( + s + for s in parsed.samples + if window_start_epoch_s <= s.ts_epoch_s <= window_end_epoch_s + ), + key=lambda s: s.ts_epoch_s, + ) + + by_label: dict[str, list[PowerSample]] = {} + for s in in_window: + by_label.setdefault(s.label, []).append(s) + + sources: list[dict[str, Any]] = [] + total_energy_j = 0.0 + have_energy = False + for label, series in sorted(by_label.items()): + energy = _energy_j(series, resolved.value_kind) + watts = [s.value for s in series] if resolved.value_kind == "power_w" else None + entry: dict[str, Any] = { + "label": label, + "value_kind": resolved.value_kind, + "sample_count": len(series), + "energy_j": energy, + } + if watts is not None: + entry["power_w"] = _stats(watts) + sources.append(entry) + if energy is not None: + total_energy_j += energy + have_energy = True + + duration_s = max(0.0, window_end_epoch_s - window_start_epoch_s) + # Status precedence: collector failure > no data > partial (dropped rows or + # a nonzero sidecar exit that still produced some samples) > ok. + if collector_status == "failed": + status = "failed" + elif not in_window: + status = "no_data" + elif parsed.dropped > 0 or collector_error: + status = "partial" + else: + status = "ok" + + energy_total: float | None = total_energy_j if have_energy else None + epot: float | None = None + epot_note: str | None = None + if energy_total is None: + epot_note = "no energy integrated over window" + elif not output_tokens: + epot_note = "output token count unavailable" + elif not consistent_with_window: + epot_note = ( + "token count spans a different window than energy " + f"(token basis: {token_window_basis}); J/token suppressed" + ) + else: + epot = energy_total / output_tokens + + return { + "schema_version": "1.0", + "status": status, + "window": { + "start_epoch_s": window_start_epoch_s, + "end_epoch_s": window_end_epoch_s, + "duration_s": duration_s, + "basis": "performance_phase", + }, + "totals": { + "energy_j": energy_total, + "mean_power_w": (energy_total / duration_s) + if (energy_total is not None and duration_s > 0) + else None, + "output_tokens": output_tokens, + "token_window_basis": token_window_basis, + "consistent_with_window": consistent_with_window, + "energy_per_output_token_j": epot, + "energy_per_output_token_note": epot_note, + }, + "sources": sources, + "provenance": { + "command": resolved.argv, + "interval_s": interval_s, + "value_kind": resolved.value_kind, + "samples_parsed": len(parsed.samples), + "samples_in_window": len(in_window), + "samples_dropped": parsed.dropped, + "collector_status": collector_status, + "collector_error": collector_error, + }, + } + + +def disabled_report(reason: str) -> dict[str, Any]: + """Minimal payload when monitoring could not even start.""" + return {"schema_version": "1.0", "status": "failed", "error": reason} diff --git a/tests/unit/power/__init__.py b/tests/unit/power/__init__.py new file mode 100644 index 000000000..467079831 --- /dev/null +++ b/tests/unit/power/__init__.py @@ -0,0 +1,14 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/tests/unit/power/test_power.py b/tests/unit/power/test_power.py new file mode 100644 index 000000000..c09966675 --- /dev/null +++ b/tests/unit/power/test_power.py @@ -0,0 +1,307 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for vendor-agnostic power monitoring.""" + +from __future__ import annotations + +import sys +import time + +import pytest +from inference_endpoint.config.schema import PowerConfig +from inference_endpoint.power.collector import PowerCollector +from inference_endpoint.power.parse import parse_trace +from inference_endpoint.power.sources import ResolvedSource, power_source, resolve +from inference_endpoint.power.window import build_power_report + +pytestmark = pytest.mark.unit + + +def _jsonl_source(value_kind: str = "power_w") -> ResolvedSource: + return ResolvedSource( + argv=["true"], + fmt="jsonl", + value_kind=value_kind, + ts_field="ts", + value_field="value", + label_field="label", + csv_header=False, + ) + + +# --------------------------------------------------------------------------- # +# Config + sources +# --------------------------------------------------------------------------- # +def test_source_requires_its_options(): + # Each source validates its own options at build time (vendor-neutral core). + with pytest.raises(ValueError, match="options.url"): + resolve(PowerConfig(source="prometheus")) + with pytest.raises(ValueError, match="options.argv"): + resolve(PowerConfig(source="command")) + # nvidia_smi needs nothing extra + resolve(PowerConfig(source="nvidia_smi")) + + +def test_disabled_power_has_no_config_footprint(): + # source=None must serialize to an empty mapping (frictionless when off). + assert PowerConfig().model_dump() == {} + assert PowerConfig(source="nvidia_smi").model_dump()["source"] == "nvidia_smi" + + +def test_resolve_nvidia_smi_argv(): + r = resolve( + PowerConfig( + source="nvidia_smi", options={"gpu_indices": [0, 3]}, interval_s=0.5 + ) + ) + assert "nvidia-smi" in r.argv + assert "-lms" in r.argv and "500" in r.argv + assert r.argv[-2:] == ["-i", "0,3"] + assert r.value_kind == "power_w" and r.fmt == "csv" + + +def test_custom_source_plugin_registration(): + # A user registers their own source with the decorator + cfg.options. + @power_source("test_custom_src") + def _build(cfg: PowerConfig) -> ResolvedSource: + return ResolvedSource( + argv=["echo", cfg.options["tag"]], + fmt="jsonl", + value_kind="power_w", + ts_field="ts", + value_field="value", + label_field="label", + csv_header=False, + ) + + r = resolve(PowerConfig(source="test_custom_src", options={"tag": "hi"})) + assert r.argv == ["echo", "hi"] + + +def test_resolve_unknown_source_raises(): + with pytest.raises(ValueError, match="unknown power source"): + resolve(PowerConfig(source="nope")) + + +# --------------------------------------------------------------------------- # +# Parsing +# --------------------------------------------------------------------------- # +def test_parse_jsonl_drops_malformed(tmp_path): + f = tmp_path / "t.log" + f.write_text( + '{"ts": 1000.0, "value": 100.0, "label": "gpu0"}\n' + "not json\n" + '{"ts": 1001.0, "value": 110.0, "label": "gpu0"}\n' + '{"ts": 1002.0}\n' # missing value + ) + res = parse_trace(f, _jsonl_source()) + assert len(res.samples) == 2 + assert res.dropped == 2 + assert res.samples[0].label == "gpu0" + + +def test_parse_csv_nvidia_format(tmp_path): + f = tmp_path / "smi.csv" + f.write_text( + "2026/05/28 05:51:22.131, 0, 259.16\n" "2026/05/28 05:51:22.131, 1, 260.00\n" + ) + res = parse_trace(f, resolve(PowerConfig(source="nvidia_smi"))) + assert res.dropped == 0 + labels = {s.label for s in res.samples} + assert labels == {"gpu0", "gpu1"} + assert res.samples[0].value == pytest.approx(259.16) + + +def test_parse_missing_file_is_empty(tmp_path): + res = parse_trace(tmp_path / "nope.log", _jsonl_source()) + assert res.samples == [] and res.dropped == 0 + + +# --------------------------------------------------------------------------- # +# Windowing + integration +# --------------------------------------------------------------------------- # +def test_trapezoid_energy_constant_power(tmp_path): + # Constant 100 W for 10 s -> 1000 J. + f = tmp_path / "p.log" + lines = [ + f'{{"ts": {1000.0 + i}, "value": 100.0, "label": "gpu0"}}' for i in range(11) + ] + f.write_text("\n".join(lines) + "\n") + rep = build_power_report( + resolved=_jsonl_source(), + trace_path=f, + window_start_epoch_s=1000.0, + window_end_epoch_s=1010.0, + output_tokens=2000, + token_window_basis="performance_phase_tracked", + consistent_with_window=True, + collector_status="ok", + collector_error=None, + interval_s=1.0, + ) + assert rep["status"] == "ok" + assert rep["totals"]["energy_j"] == pytest.approx(1000.0) + assert rep["totals"]["mean_power_w"] == pytest.approx(100.0) + # 1000 J / 2000 tokens = 0.5 J/token + assert rep["totals"]["energy_per_output_token_j"] == pytest.approx(0.5) + + +def test_energy_counter_kind_uses_delta(tmp_path): + # Cumulative joule counter: 5000 -> 9000 over window => 4000 J. + f = tmp_path / "c.log" + f.write_text( + '{"ts": 1000.0, "value": 5000.0, "label": "node"}\n' + '{"ts": 1005.0, "value": 7000.0, "label": "node"}\n' + '{"ts": 1010.0, "value": 9000.0, "label": "node"}\n' + ) + rep = build_power_report( + resolved=_jsonl_source(value_kind="energy_j"), + trace_path=f, + window_start_epoch_s=1000.0, + window_end_epoch_s=1010.0, + output_tokens=None, + token_window_basis="performance_phase_tracked", + consistent_with_window=True, + collector_status="ok", + collector_error=None, + interval_s=5.0, + ) + assert rep["totals"]["energy_j"] == pytest.approx(4000.0) + + +def test_energy_counter_kind_detects_midrun_reset(tmp_path): + # A counter reset mid-window (100 -> 150 -> 20 -> 110) has a positive + # last-minus-first delta (10) but is unusable; must return None, not 10. + f = tmp_path / "c.log" + f.write_text( + '{"ts": 1000.0, "value": 100.0, "label": "node"}\n' + '{"ts": 1005.0, "value": 150.0, "label": "node"}\n' + '{"ts": 1010.0, "value": 20.0, "label": "node"}\n' + '{"ts": 1015.0, "value": 110.0, "label": "node"}\n' + ) + rep = build_power_report( + resolved=_jsonl_source(value_kind="energy_j"), + trace_path=f, + window_start_epoch_s=1000.0, + window_end_epoch_s=1015.0, + output_tokens=None, + token_window_basis="performance_phase_tracked", + consistent_with_window=True, + collector_status="ok", + collector_error=None, + interval_s=5.0, + ) + assert rep["totals"]["energy_j"] is None + + +def test_epot_suppressed_when_inconsistent(tmp_path): + f = tmp_path / "p.log" + f.write_text( + '{"ts": 1000.0, "value": 100.0, "label": "gpu0"}\n' + '{"ts": 1010.0, "value": 100.0, "label": "gpu0"}\n' + ) + rep = build_power_report( + resolved=_jsonl_source(), + trace_path=f, + window_start_epoch_s=1000.0, + window_end_epoch_s=1010.0, + output_tokens=2000, + token_window_basis="global_run", + consistent_with_window=False, # denominators would mix + collector_status="ok", + collector_error=None, + interval_s=10.0, + ) + assert rep["totals"]["energy_per_output_token_j"] is None + assert "suppressed" in rep["totals"]["energy_per_output_token_note"] + + +def test_status_no_data_when_window_empty(tmp_path): + f = tmp_path / "p.log" + f.write_text('{"ts": 5.0, "value": 100.0, "label": "gpu0"}\n') # outside window + rep = build_power_report( + resolved=_jsonl_source(), + trace_path=f, + window_start_epoch_s=1000.0, + window_end_epoch_s=1010.0, + output_tokens=10, + token_window_basis="performance_phase_tracked", + consistent_with_window=True, + collector_status="ok", + collector_error=None, + interval_s=1.0, + ) + assert rep["status"] == "no_data" + assert rep["totals"]["energy_j"] is None + + +def test_collector_failed_status_propagates(tmp_path): + rep = build_power_report( + resolved=_jsonl_source(), + trace_path=tmp_path / "missing.log", + window_start_epoch_s=1000.0, + window_end_epoch_s=1010.0, + output_tokens=10, + token_window_basis="performance_phase_tracked", + consistent_with_window=True, + collector_status="failed", + collector_error="boom", + interval_s=1.0, + ) + assert rep["status"] == "failed" + + +# --------------------------------------------------------------------------- # +# End-to-end collector with a fake command source +# --------------------------------------------------------------------------- # +def test_collector_end_to_end(tmp_path): + # A fake source: emit 5 JSONL samples at ~50 ms spacing, then exit. + script = ( + "import json,sys,time\n" + "for i in range(5):\n" + " print(json.dumps({'ts': time.time(), 'value': 200.0, 'label': 'gpu0'}), flush=True)\n" + " time.sleep(0.05)\n" + ) + cfg = PowerConfig( + source="command", + options={"argv": [sys.executable, "-c", script]}, + interval_s=0.05, + ) + t0 = time.time() + collector = PowerCollector(cfg, tmp_path / "power") + collector.start() + assert collector.status == "ok" + time.sleep(0.6) # let it finish + collector.stop() + t1 = time.time() + + assert collector.trace_path.exists() + rep = build_power_report( + resolved=collector.resolved, + trace_path=collector.trace_path, + window_start_epoch_s=t0, + window_end_epoch_s=t1, + output_tokens=1000, + token_window_basis="performance_phase_tracked", + consistent_with_window=True, + collector_status=collector.status, + collector_error=collector.error, + interval_s=cfg.interval_s, + ) + assert rep["status"] == "ok" + assert rep["sources"][0]["sample_count"] >= 3 + assert rep["sources"][0]["power_w"]["mean"] == pytest.approx(200.0) + assert rep["totals"]["energy_j"] is not None