diff --git a/bench/watt_probe.py b/bench/watt_probe.py new file mode 100644 index 0000000..ae1737b --- /dev/null +++ b/bench/watt_probe.py @@ -0,0 +1,187 @@ +#!/usr/bin/env python3 +"""Zero-dependency GPU+CPU power probe — runs ON a worker (GPU) box. + +Split out of `bench/watt_bench.py` for the laptop-driver / worker-reporter +architecture (fox 2026-05-20): the laptop holds the shards, runs +retrieval + the judge + the workload loop, and drives the worker's LLM +endpoint over the network; the worker boxes (3090 / 4090) only serve the +model and **report their own power**. Power sampling must run on the box +(nvidia-smi + RAPL are host-local), but everything else stays on the +laptop — so the workers need NO shards, NO judge, NO arborist install, +NO venv. This script is **stdlib-only** (python3 + nvidia-smi + +readable RAPL); copy it over and run it, nothing to install. + +It samples GPU power (`nvidia-smi power.draw`) and CPU package energy +(Intel RAPL `energy_uj`) over a window and emits a JSON energy report. +The laptop orchestrates the window two ways: + + * ``--duration SECONDS`` — sample for a fixed window (laptop sizes it + to the workload), or + * ``--until-file PATH`` — sample until PATH appears (laptop touches it + to stop), the tightly-correlated mode: start probe, run workload, + touch stop-file, collect JSON. + +Output (stdout, or ``--out PATH``): the same energy schema watt_bench's +samplers produce — gpu mean/peak W + joules, cpu mean W + joules, +window, gpu_name — so the laptop merges it into the cell record exactly +as if sampled locally. + +Usage on the worker (driven by the laptop over SSH):: + + # fixed window + python3 watt_probe.py --duration 60 --gpu-label 4090 + # signalled window (laptop: start, run workload, then `touch /tmp/stop`) + python3 watt_probe.py --until-file /tmp/wattprobe.stop --out /tmp/p.json & +""" +from __future__ import annotations + +import argparse +import json +import subprocess +import sys +import time +from pathlib import Path + + +def _nvsmi(query: str, gpu_index: int) -> str | None: + try: + out = subprocess.run( + ["nvidia-smi", f"--id={gpu_index}", f"--query-gpu={query}", + "--format=csv,noheader,nounits"], + capture_output=True, text=True, timeout=5).stdout.strip() + return out.splitlines()[0] if out else None + except Exception: # noqa: BLE001 + return None + + +def _gpu_name(gpu_index: int) -> str: + return _nvsmi("name", gpu_index) or "unknown" + + +def _gpu_watts(gpu_index: int) -> float | None: + v = _nvsmi("power.draw", gpu_index) + try: + return float(v) if v is not None else None + except ValueError: + return None + + +def _gpu_util(gpu_index: int) -> float | None: + v = _nvsmi("utilization.gpu", gpu_index) + try: + return float(v) if v is not None else None + except (ValueError, TypeError): + return None + + +def _rapl_packages() -> list[Path]: + base = Path("/sys/class/powercap") + if not base.exists(): + return [] + # top-level packages only (intel-rapl:N) — NOT the :N:M core/uncore + # subdomains, which would double-count. + return sorted(p / "energy_uj" for p in base.glob("intel-rapl:*") + if p.name.count(":") == 1 and (p / "energy_uj").exists()) + + +def _rapl_total_uj(pkgs: list[Path]) -> int | None: + if not pkgs: + return None + vals = [] + for p in pkgs: + try: + vals.append(int(p.read_text().strip())) + except Exception: # noqa: BLE001 + return None + return sum(vals) + + +def main() -> int: + ap = argparse.ArgumentParser() + ap.add_argument("--gpu-index", type=int, default=0) + ap.add_argument("--gpu-label", default="") + ap.add_argument("--interval", type=float, default=0.5) + g = ap.add_mutually_exclusive_group(required=True) + g.add_argument("--duration", type=float, + help="sample for this many seconds then emit") + g.add_argument("--until-file", default=None, + help="sample until this path exists (laptop touches " + "it to stop) — the workload-correlated mode") + ap.add_argument("--out", default="", + help="write JSON here (default stdout)") + a = ap.parse_args() + + pkgs = _rapl_packages() + gpu_name = _gpu_name(a.gpu_index) + gpu_avail = _gpu_watts(a.gpu_index) is not None + cpu_e0 = _rapl_total_uj(pkgs) + cpu_avail = cpu_e0 is not None + + samples: list[tuple[float, float, float | None]] = [] # (t, watts, util) + t_start = time.time() + stop_path = Path(a.until_file) if a.until_file else None + + def _should_stop() -> bool: + if a.duration is not None: + return (time.time() - t_start) >= a.duration + return stop_path.exists() if stop_path else True + + # If signalled mode and the stop-file already exists, clear our view + # by ignoring a pre-existing file would be wrong — caller is expected + # to remove it before starting. We just begin sampling. + while not _should_stop(): + w = _gpu_watts(a.gpu_index) + if w is not None: + samples.append((time.time(), w, _gpu_util(a.gpu_index))) + time.sleep(a.interval) + + t_end = time.time() + cpu_e1 = _rapl_total_uj(pkgs) + + # GPU: trapezoid-integrate sampled instantaneous power → joules. + gpu_joules = None + gpu_mean = gpu_peak = None + util_mean = None + if samples: + ws = [w for _, w, _ in samples] + gpu_mean = round(sum(ws) / len(ws), 1) + gpu_peak = round(max(ws), 1) + utils = [u for _, _, u in samples if u is not None] + util_mean = round(sum(utils) / len(utils), 1) if utils else None + j = 0.0 + for (t0, w0, _), (t1, w1, _) in zip(samples, samples[1:]): + j += (w0 + w1) / 2.0 * (t1 - t0) + gpu_joules = round(j, 1) + + # CPU: RAPL energy is a direct cumulative counter — end minus start. + cpu_joules = cpu_mean = None + if cpu_avail and cpu_e1 is not None and cpu_e1 >= cpu_e0: + cpu_joules = round((cpu_e1 - cpu_e0) / 1_000_000.0, 1) + win = t_end - t_start + cpu_mean = round(cpu_joules / win, 1) if win else None + + report = { + "gpu_label": a.gpu_label or "gpu", + "gpu_name": gpu_name, + "gpu_available": gpu_avail, + "cpu_rapl_available": cpu_avail, + "cpu_rapl_packages": len(pkgs), + "window_s": round(t_end - t_start, 1), + "n_samples": len(samples), + "gpu_mean_w": gpu_mean, + "gpu_peak_w": gpu_peak, + "gpu_util_mean_pct": util_mean, + "gpu_joules": gpu_joules, + "cpu_mean_w": cpu_mean, + "cpu_joules": cpu_joules, + } + blob = json.dumps(report, indent=2) + if a.out: + Path(a.out).write_text(blob + "\n") + else: + print(blob) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main())