diff --git a/bench/watt_bench.py b/bench/watt_bench.py index e84105d..57ae9c7 100644 --- a/bench/watt_bench.py +++ b/bench/watt_bench.py @@ -222,6 +222,115 @@ class CpuSampler: "window_s": round(window, 1), "n_packages": len(self._pkgs)} +# ----------------------------------------------------------- probe abstraction + +class LocalProbe: + """GPU+CPU power sampled on THIS machine (watt_bench runs on the + GPU box). Wraps PowerSampler + CpuSampler behind the start/stop/ + gpu_stats/cpu_stats interface RemoteProbe also implements.""" + + def __init__(self, gpu_index: int, interval: float, cpu_cmd: str | None): + self._g = PowerSampler(gpu_index, interval) + self._c = CpuSampler(cpu_cmd) + self.gpu_name = self._g.gpu_name + self.available = self._g.available + self.cpu_available = self._c.available + + def start(self): + self._g.samples = [] + self._g.start() + self._c.start() + + def stop(self): + self._g.stop() + self._gs = self._g.stats() + self._cs = self._c.stats() + + def gpu_stats(self) -> dict: + return self._gs + + def cpu_stats(self) -> dict: + return self._cs + + +class RemoteProbe: + """GPU+CPU power sampled on a REMOTE worker box over SSH — the + laptop-driver / worker-reporter mode (fox 2026-05-20). The laptop + runs the workload (retrieval + judge + driving the worker's LLM + endpoint); this orchestrates ``bench/watt_probe.py`` on the worker + for the workload window so the worker reports its own power without + needing shards / arborist / a venv (the probe is stdlib-only). + + start() scp's the probe (once) and launches it in --until-file mode + detached; stop() touches the stop-file, waits for flush, fetches the + JSON, and exposes it in the LocalProbe stats shape.""" + + def __init__(self, host: str, gpu_index: int, gpu_label: str, + interval: float): + self.host = host + self.gpu_index = gpu_index + self.gpu_label = gpu_label or "gpu" + self.interval = interval + tag = time.strftime("%H%M%S", time.gmtime()) + self.stop_file = f"/tmp/wattprobe_{tag}.stop" + self.out_file = f"/tmp/wattprobe_{tag}.json" + self._report: dict | None = None + self.gpu_name = f"remote:{host}" + self.available = self._ensure_probe() + self.cpu_available = self.available # confirmed after stop() + + def _ssh(self, cmd: str, timeout: int = 30): + return subprocess.run( + ["ssh", "-o", "BatchMode=yes", "-o", "ConnectTimeout=10", + self.host, cmd], + capture_output=True, text=True, timeout=timeout) + + def _ensure_probe(self) -> bool: + probe = str(Path(__file__).resolve().parent / "watt_probe.py") + try: + subprocess.run( + ["scp", "-q", "-o", "BatchMode=yes", probe, + f"{self.host}:/tmp/watt_probe.py"], + check=True, timeout=30) + return True + except Exception: # noqa: BLE001 + return False + + def start(self): + self._ssh(f"rm -f {self.stop_file} {self.out_file}") + # Detach so ssh returns while the probe keeps sampling. + self._ssh( + f"nohup python3 /tmp/watt_probe.py --until-file {self.stop_file} " + f"--out {self.out_file} --gpu-index {self.gpu_index} " + f"--gpu-label {self.gpu_label} --interval {self.interval} " + f">/dev/null 2>&1 & echo started") + + def stop(self): + self._ssh(f"touch {self.stop_file}") + time.sleep(self.interval * 3 + 1.0) # let the probe flush its JSON + r = self._ssh(f"cat {self.out_file}") + try: + self._report = json.loads(r.stdout) + self.gpu_name = self._report.get("gpu_name", self.gpu_name) + self.cpu_available = bool(self._report.get("cpu_rapl_available")) + except Exception: # noqa: BLE001 + self._report = None + + def gpu_stats(self) -> dict: + rp = self._report or {} + return {"available": rp.get("gpu_available", False), + "mean_w": rp.get("gpu_mean_w"), "peak_w": rp.get("gpu_peak_w"), + "joules": rp.get("gpu_joules"), "window_s": rp.get("window_s"), + "util_mean_pct": rp.get("gpu_util_mean_pct"), + "n_samples": rp.get("n_samples")} + + def cpu_stats(self) -> dict: + rp = self._report or {} + return {"available": rp.get("cpu_rapl_available", False), + "mean_w": rp.get("cpu_mean_w"), "joules": rp.get("cpu_joules"), + "window_s": rp.get("window_s")} + + # ----------------------------------------------------------- driver def _make_client(endpoint: str): @@ -264,9 +373,17 @@ def main() -> int: ap.add_argument("--gpu-label", default="", help="human label for the GPU tier (e.g. 3090, 4090) " "— stamped into the output for cross-GPU compare") + ap.add_argument("--remote-gpu-host", default="", + help="laptop-driver mode: run the workload HERE " + "(local shards + judge, driving --endpoint) but " + "sample power on this REMOTE worker box over SSH " + "via bench/watt_probe.py. The worker needs no " + "shards / arborist / venv. e.g. ai.foxhop.net") ap.add_argument("--endpoint", default="", - help="override the model endpoint (e.g. localhost on " - "the GPU box). Default = MODELS entry's endpoint.") + help="override the model endpoint. In --remote-gpu-host " + "mode point this at the worker's endpoint (its " + "public URL or an SSH tunnel). Default = MODELS " + "entry's endpoint.") ap.add_argument("--sample-interval", type=float, default=0.5) ap.add_argument("--idle-baseline-s", type=float, default=5.0, help="seconds to sample idle power before driving load") @@ -289,25 +406,39 @@ def main() -> int: outp = Path(a.out_dir) / f"watt_{label}_{ts}.jsonl" outp.parent.mkdir(parents=True, exist_ok=True) - sampler = PowerSampler(a.gpu_index, a.sample_interval) - cpu_probe = CpuSampler(a.cpu_energy_cmd or None) - print(f"watt_bench — gpu_label={label!r} detected={sampler.gpu_name!r} " - f"nvidia-smi={'yes' if sampler.available else 'NO (energy null)'}") - print(f" cpu RAPL: {'yes' if cpu_probe.available else 'NO (root-only / locked — pass --cpu-energy-cmd)'}" - f" ({len(cpu_probe._pkgs)} package(s))") + remote = a.remote_gpu_host.strip() + + def _new_probe(): + if remote: + return RemoteProbe(remote, a.gpu_index, label, a.sample_interval) + return LocalProbe(a.gpu_index, a.sample_interval, + a.cpu_energy_cmd or None) + + probe0 = _new_probe() + mode = f"REMOTE via ssh {remote}" if remote else "LOCAL" + print(f"watt_bench [{mode}] — gpu_label={label!r} " + f"detected={probe0.gpu_name!r} " + f"gpu={'yes' if probe0.available else 'NO'}") + print(f" cpu energy: {'yes' if probe0.cpu_available else 'NO (RAPL root-only / unavailable)'}") print(f" fixture={a.fixture} n={len(items)} models={models} " f"variants={variants} arborist_ref={a.arborist_ref or 'none'}") # Idle baseline. idle = None - if sampler.available and a.idle_baseline_s > 0: - print(f" sampling idle baseline for {a.idle_baseline_s}s …") - sampler.start() + if probe0.available and a.idle_baseline_s > 0: + print(f" sampling idle baseline for {a.idle_baseline_s}s " + f"({'remote' if remote else 'local'}) …") + probe0.start() time.sleep(a.idle_baseline_s) - sampler.stop() - idle = sampler.stats() - print(f" idle: mean={idle['mean_w']}W peak={idle['peak_w']}W") - sampler.samples = [] # reset for the load window + probe0.stop() + ig = probe0.gpu_stats() + idle = {"mean_w": ig.get("mean_w"), "peak_w": ig.get("peak_w"), + "gpu_util_mean_pct": ig.get("util_mean_pct"), + "cpu_mean_w": probe0.cpu_stats().get("mean_w")} + print(f" idle: gpu mean={idle['mean_w']}W peak={idle['peak_w']}W " + f"util={idle['gpu_util_mean_pct']}% cpu={idle['cpu_mean_w']}W" + + (" [util>5% ⇒ contaminated by live traffic]" + if (idle.get('gpu_util_mean_pct') or 0) > 5 else "")) cells = [] # one entry per (arm, model, variant) measured @@ -315,10 +446,8 @@ def main() -> int: endpoint = a.endpoint or MODELS[mkey]["endpoint"] client = _make_client(endpoint) rows = [] - s = PowerSampler(a.gpu_index, a.sample_interval) - cpu = CpuSampler(a.cpu_energy_cmd or None) - s.start() - cpu.start() + probe = _new_probe() + probe.start() t0 = time.time() try: for variant in variants: @@ -354,9 +483,9 @@ def main() -> int: "est_completion_tokens": ctoks}) finally: client.close() - s.stop() - st = s.stats() - cpu_st = cpu.stats() + probe.stop() + st = probe.gpu_stats() + cpu_st = probe.cpu_stats() elapsed = time.time() - t0 nq = len(rows) tot_tok = sum(r["est_completion_tokens"] for r in rows) @@ -368,7 +497,7 @@ def main() -> int: "n_questions": nq, "elapsed_s": round(elapsed, 1), "gpu_power": st, "cpu_power": cpu_st, - "gpu_label": label, "gpu_name": sampler.gpu_name, + "gpu_label": label, "gpu_name": probe.gpu_name, "idle_baseline": idle, "gpu_joules_per_question": (round(gpu_j / nq, 1) if gpu_j and nq else None), @@ -399,19 +528,20 @@ def main() -> int: if a.arborist_ref and a.arborist_ref in MODELS: run_cell("arborist", a.arborist_ref) - report = {"ts": ts, "gpu_label": label, "gpu_name": sampler.gpu_name, - "nvidia_smi_available": sampler.available, - "cpu_rapl_available": cpu_probe.available, - "cpu_rapl_packages": len(cpu_probe._pkgs), + report = {"ts": ts, "gpu_label": label, "gpu_name": probe0.gpu_name, + "mode": "remote" if remote else "local", + "remote_gpu_host": remote or None, + "gpu_available": probe0.available, + "cpu_energy_available": probe0.cpu_available, "idle_baseline": idle, "fixture": a.fixture, "n_items": len(items), "cells": cells} rp = outp.with_suffix(".json").with_name(f"watt_{label}_{ts}.json") rp.write_text(json.dumps(report, indent=2)) print(f"\n per-question rows: {outp}") print(f" energy report: {rp}") - if not sampler.available: - print(" NOTE: nvidia-smi unavailable — energy fields are null. " - "Run this ON the GPU box for real wattage.") + if not probe0.available: + print(" NOTE: GPU power unavailable — energy fields null. Run ON " + "the GPU box, or use --remote-gpu-host to sample a worker.") return 0