arborist/bench/watt_probe.py

#!/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 classify_power_bands(watt_samples) -> dict:
    """Split power samples into a low band (model resident, between
    requests — 'middle idle') and a high band (active token generation)
    by the LARGEST GAP in the sorted samples. Data-derived — never
    hardcoded; every card/model/inference-server has its own profile
    (fox 2026-05-21). Reports measured POWER STATES, not the duty-cycle
    blend that `mean_w` (joules/window) collapses them into.
    """
    ws = sorted(w for w in watt_samples if w is not None)
    n = len(ws)
    if n < 4:
        return {"available": False, "n_samples": n}
    gi = max(range(n - 1), key=lambda i: ws[i + 1] - ws[i])
    split = (ws[gi] + ws[gi + 1]) / 2.0
    low, high = ws[: gi + 1], ws[gi + 1:]
    low_mean = sum(low) / len(low)
    high_mean = sum(high) / len(high)
    # Distinct generation state only if the high band is clearly above
    # the low band; otherwise the window was effectively unimodal.
    bimodal = high_mean > 1.5 * max(low_mean, 1e-6)
    return {
        "available": True,
        "n_samples": n,
        "bimodal": bimodal,
        "split_w": round(split, 1),
        "low_band_w": round(low_mean, 1),    # serving floor / middle-idle
        "high_band_w": round(high_mean, 1),  # active generation draw
        "duty_cycle": round(len(high) / n, 3),
        "min_w": round(ws[0], 1),
        "peak_w": round(ws[-1], 1),
    }


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

    power_bands = classify_power_bands([w for _, w, _ in samples])

    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,
        # Measured power-state decomposition + raw samples for audit /
        # offline re-analysis (n is small; transmits fine over ssh cat).
        "power_bands": power_bands,
        "watt_samples": [round(w, 1) for _, w, _ in samples],
    }
    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())