Improve exo harness with path state (#1815)

<img width="3224" height="1476" alt="image"
src="https://github.com/user-attachments/assets/d90a7d8a-9fe5-43a1-a715-1ef7ecc15422"
/>

bench/eval_tool_calls.py

@@ -17,6 +17,7 @@ from harness import (
     ExoClient,
     ExoHttpError,
     add_common_instance_args,
+    capture_cluster_snapshot,
     instance_id_from_instance,
     nodes_used_in_instance,
     resolve_model_short_id,
@@ -1006,6 +1007,7 @@ Examples:
         sys.exit(1)
 
     time.sleep(1)
+    cluster_snapshot = capture_cluster_snapshot(exo)
     all_results: list[ScenarioResult] = []
 
     try:
@@ -1084,16 +1086,19 @@ Examples:
                 print(f"  - {r.name} [{r.api}/{r.phase}]: {r.error}", file=log)
 
     json_results = [result_to_dict(r) for r in all_results]
+    output: dict[str, Any] = {"results": json_results}
+    if cluster_snapshot:
+        output["cluster"] = cluster_snapshot
 
     if args.stdout:
-        print(json.dumps(json_results, indent=2))
+        print(json.dumps(output, indent=2))
     else:
         json_path = args.json_out
         parent = os.path.dirname(json_path)
         if parent:
             os.makedirs(parent, exist_ok=True)
         with open(json_path, "w") as f:
-            json.dump(json_results, f, indent=2)
+            json.dump(output, f, indent=2)
             f.write("\n")
         print(f"\nJSON results written to {json_path}", file=log)
 

bench/exo_bench.py

@@ -22,6 +22,7 @@ import contextlib
 import itertools
 import json
 import sys
+import threading
 import time
 from collections.abc import Callable
 from concurrent.futures import ThreadPoolExecutor, as_completed
@@ -33,7 +34,9 @@ from harness import (
     ExoClient,
     ExoHttpError,
     add_common_instance_args,
+    capture_cluster_snapshot,
     instance_id_from_instance,
+    node_ids_from_instance,
     nodes_used_in_instance,
     resolve_model_short_id,
     run_planning_phase,
@@ -131,6 +134,91 @@ def format_peak_memory(b: float) -> str:
     raise ValueError("You're using petabytes of memory. Something went wrong...")
 
 
+_SAMPLER_METRICS = ("gpuUsage", "temp", "sysPower", "pcpuUsage", "ecpuUsage")
+
+
+class SystemMetricsSampler:
+    def __init__(self, client: ExoClient, node_ids: list[str], interval_s: float = 1.0):
+        self._client = client
+        self._node_ids = node_ids
+        self._interval_s = interval_s
+        self._samples: dict[str, list[tuple[float, dict[str, float]]]] = {
+            nid: [] for nid in node_ids
+        }
+        self._stop = threading.Event()
+        self._thread: threading.Thread | None = None
+
+    def start(self) -> None:
+        self._stop.clear()
+        self._thread = threading.Thread(target=self._poll_loop, daemon=True)
+        self._thread.start()
+
+    def stop(self) -> None:
+        self._stop.set()
+        if self._thread:
+            self._thread.join(timeout=5)
+
+    def _poll_loop(self) -> None:
+        while not self._stop.is_set():
+            t = time.monotonic()
+            for nid in self._node_ids:
+                try:
+                    data = self._client.get_node_system(nid)
+                    if data:
+                        self._samples[nid].append(
+                            (t, {k: data.get(k, 0.0) for k in _SAMPLER_METRICS})
+                        )
+                except Exception:
+                    pass
+            self._stop.wait(self._interval_s)
+
+    def energy_between(self, t0: float, t1: float) -> float:
+        total_joules = 0.0
+        for _nid, samples in self._samples.items():
+            window = [(t, s["sysPower"]) for t, s in samples if t0 <= t <= t1]
+            if len(window) >= 2:
+                for i in range(1, len(window)):
+                    dt = window[i][0] - window[i - 1][0]
+                    avg_power = (window[i][1] + window[i - 1][1]) / 2
+                    total_joules += avg_power * dt
+            elif len(window) == 1:
+                total_joules += window[0][1] * (t1 - t0)
+        return total_joules
+
+    def summarize(self) -> dict[str, dict[str, dict[str, float]]]:
+        result: dict[str, dict[str, dict[str, float]]] = {}
+        for nid, samples in self._samples.items():
+            if not samples:
+                continue
+            metrics: dict[str, dict[str, float]] = {}
+            for key in _SAMPLER_METRICS:
+                values = [s[key] for t, s in samples]
+                metrics[key] = {
+                    "min": round(min(values), 2),
+                    "max": round(max(values), 2),
+                    "mean": round(mean(values), 2),
+                    "samples": len(values),
+                }
+            result[nid] = metrics
+        return result
+
+    def print_summary(self, placement_label: str) -> None:
+        summary = self.summarize()
+        if not summary:
+            return
+        logger.info(f"--- System Metrics ({placement_label}) ---")
+        for nid, metrics in summary.items():
+            gpu = metrics.get("gpuUsage", {})
+            temp = metrics.get("temp", {})
+            power = metrics.get("sysPower", {})
+            logger.info(
+                f"  {nid}:  "
+                f"GPU {gpu.get('mean', 0) * 100:.0f}% avg ({gpu.get('min', 0) * 100:.0f}–{gpu.get('max', 0) * 100:.0f}%)  |  "
+                f"{temp.get('mean', 0):.1f}°C avg  |  "
+                f"{power.get('mean', 0):.1f}W avg"
+            )
+
+
 def parse_int_list(values: list[str]) -> list[int]:
     items: list[int] = []
     for v in values:
@@ -279,6 +367,17 @@ def main() -> int:
         action="store_true",
         help="Force all pp×tg combinations (cartesian product) even when lists have equal length.",
     )
+    ap.add_argument(
+        "--no-system-metrics",
+        action="store_true",
+        help="Disable GPU utilization, temperature, and power collection during inference.",
+    )
+    ap.add_argument(
+        "--metrics-interval",
+        type=float,
+        default=1.0,
+        help="System metrics polling interval in seconds (default: 1.0).",
+    )
     args = ap.parse_args()
 
     pp_list = parse_int_list(args.pp)
@@ -364,7 +463,9 @@ def main() -> int:
     else:
         logger.info("Download: model already cached")
 
+    cluster_snapshot = capture_cluster_snapshot(client)
     all_rows: list[dict[str, Any]] = []
+    all_system_metrics: dict[str, dict[str, dict[str, float]]] = {}
 
     for preview in selected:
         instance = preview["instance"]
@@ -390,6 +491,16 @@ def main() -> int:
 
         time.sleep(1)
 
+        sampler: SystemMetricsSampler | None = None
+        if not args.no_system_metrics:
+            nids = node_ids_from_instance(instance)
+            sampler = SystemMetricsSampler(
+                ExoClient(args.host, args.port, timeout_s=30),
+                nids,
+                interval_s=args.metrics_interval,
+            )
+            sampler.start()
+
         try:
             for i in range(args.warmup):
                 run_one_completion(
@@ -408,15 +519,18 @@ def main() -> int:
                 for concurrency in concurrency_list:
                     logger.info(f"--- pp={pp} tg={tg} concurrency={concurrency} ---")
                     runs: list[dict[str, Any]] = []
+                    inference_windows: list[tuple[float, float]] = []
                     for r in range(args.repeat):
                         time.sleep(3)
 
                         if concurrency <= 1:
                             # Sequential: single request
                             try:
+                                inf_t0 = time.monotonic()
                                 row, actual_pp_tokens = run_one_completion(
                                     client, full_model_id, pp, tg, prompt_sizer
                                 )
+                                inference_windows.append((inf_t0, time.monotonic()))
                             except Exception as e:
                                 logger.error(e)
                                 continue
@@ -457,6 +571,7 @@ def main() -> int:
                                     c, full_model_id, _pp, _tg, prompt_sizer
                                 )
 
+                            inf_t0 = time.monotonic()
                             with ThreadPoolExecutor(max_workers=concurrency) as pool:
                                 futures = {
                                     pool.submit(_run_concurrent, i): i
@@ -468,6 +583,7 @@ def main() -> int:
                                     except Exception as e:
                                         logger.error(f"Concurrent request failed: {e}")
                                         batch_errors += 1
+                            inference_windows.append((inf_t0, time.monotonic()))
 
                             for idx, (row, actual_pp_tokens) in enumerate(
                                 batch_results
@@ -522,13 +638,30 @@ def main() -> int:
                             x["stats"]["peak_memory_usage"]["inBytes"] for x in runs
                         )
 
-                        logger.info(
+                        summary = (
                             f"prompt_tps={prompt_tps:.2f} gen_tps={gen_tps:.2f}    "
                             f"prompt_tokens={ptok} gen_tokens={gtok}    "
-                            f"peak_memory={format_peak_memory(peak)}\n"
+                            f"peak_memory={format_peak_memory(peak)}"
                         )
+                        if sampler and inference_windows:
+                            joules = sum(
+                                sampler.energy_between(t0, t1)
+                                for t0, t1 in inference_windows
+                            )
+                            inf_seconds = sum(t1 - t0 for t0, t1 in inference_windows)
+                            avg_watts = joules / inf_seconds if inf_seconds > 0 else 0
+                            summary += f"    energy={joules:.1f}J ({avg_watts:.1f}W avg over {inf_seconds:.1f}s inference)"
+                        logger.info(f"{summary}\n")
                     time.sleep(2)
         finally:
+            if sampler:
+                sampler.stop()
+                placement_label = f"{sharding}/{instance_meta}/{n_nodes} nodes"
+                sampler.print_summary(placement_label)
+                placement_metrics = sampler.summarize()
+                if placement_metrics:
+                    all_system_metrics.update(placement_metrics)
+
             try:
                 client.request_json("DELETE", f"/instance/{instance_id}")
             except ExoHttpError as e:
@@ -539,11 +672,17 @@ def main() -> int:
 
             time.sleep(5)
 
+    output: dict[str, Any] = {"runs": all_rows}
+    if cluster_snapshot:
+        output["cluster"] = cluster_snapshot
+    if all_system_metrics:
+        output["system_metrics"] = all_system_metrics
+
     if args.stdout:
-        json.dump(all_rows, sys.stdout, indent=2, ensure_ascii=False)
+        json.dump(output, sys.stdout, indent=2, ensure_ascii=False)
     elif args.json_out:
         with open(args.json_out, "w", encoding="utf-8") as f:
-            json.dump(all_rows, f, indent=2, ensure_ascii=False)
+            json.dump(output, f, indent=2, ensure_ascii=False)
         logger.debug(f"\nWrote results JSON: {args.json_out}")
 
     return 0

bench/exo_eval.py

@@ -46,6 +46,7 @@ from harness import (
     ExoClient,
     ExoHttpError,
     add_common_instance_args,
+    capture_cluster_snapshot,
     instance_id_from_instance,
     nodes_used_in_instance,
     resolve_model_short_id,
@@ -1027,16 +1028,18 @@ def save_results(
     concurrency: int,
     results: list[QuestionResult],
     scores: dict[str, Any],
+    cluster: dict[str, Any] | None = None,
 ) -> Path:
     out_dir = Path(results_dir) / model.replace("/", "_") / benchmark_name
     out_dir.mkdir(parents=True, exist_ok=True)
     ts = time.strftime("%Y%m%d_%H%M%S")
     path = out_dir / f"c{concurrency}_{ts}.json"
 
-    data = {
+    data: dict[str, Any] = {
         "benchmark": benchmark_name,
         "model": model,
         "concurrency": concurrency,
+        **({"cluster": cluster} if cluster else {}),
         "scores": scores,
         "results": [
             {
@@ -1231,8 +1234,10 @@ def main() -> int:
                 client.request_json("DELETE", f"/instance/{instance_id}")
             return 1
         time.sleep(1)
+        cluster_snapshot = capture_cluster_snapshot(client)
     else:
         full_model_id = args.model
+        cluster_snapshot = None
 
     # Auto-detect reasoning from model config
     model_config = load_model_config(full_model_id)
@@ -1328,6 +1333,7 @@ def main() -> int:
                             c,
                             results,
                             scores,
+                            cluster=cluster_snapshot,
                         )
                         results_by_c[c] = results
                 if len(results_by_c) >= 2:
@@ -1358,6 +1364,7 @@ def main() -> int:
                         args.num_concurrent,
                         results,
                         scores,
+                        cluster=cluster_snapshot,
                     )
     finally:
         if instance_id is not None:

bench/harness.py

@@ -69,6 +69,35 @@ class ExoClient:
     def post_bench_chat_completions(self, payload: dict[str, Any]) -> dict[str, Any]:
         return self.request_json("POST", "/bench/chat/completions", body=payload)
 
+    def get_state_path(self, path: str) -> Any:
+        try:
+            return self.request_json("GET", f"/state/{path}")
+        except ExoHttpError as e:
+            if e.status == 404:
+                return None
+            raise
+
+    def get_instance(self, instance_id: str) -> dict[str, Any] | None:
+        return self.get_state_path(f"instances/{instance_id}")
+
+    def get_runner(self, runner_id: str) -> dict[str, Any] | None:
+        return self.get_state_path(f"runners/{runner_id}")
+
+    def get_node_downloads(self, node_id: str) -> list[dict[str, Any]] | None:
+        return self.get_state_path(f"downloads/{node_id}")
+
+    def get_node_disk(self, node_id: str) -> dict[str, Any] | None:
+        return self.get_state_path(f"nodeDisk/{node_id}")
+
+    def get_node_system(self, node_id: str) -> dict[str, Any] | None:
+        return self.get_state_path(f"nodeSystem/{node_id}")
+
+    def get_node_identities(self) -> dict[str, Any] | None:
+        return self.get_state_path("nodeIdentities")
+
+    def get_topology(self) -> dict[str, Any] | None:
+        return self.get_state_path("topology")
+
 
 def unwrap_instance(instance: dict[str, Any]) -> dict[str, Any]:
     if len(instance) != 1:
@@ -97,6 +126,11 @@ def runner_ids_from_instance(instance: dict[str, Any]) -> list[str]:
     return list(runner_to_shard.keys())
 
 
+def node_ids_from_instance(instance: dict[str, Any]) -> list[str]:
+    inner = unwrap_instance(instance)
+    return list(inner["shardAssignments"]["nodeToRunner"].keys())
+
+
 def runner_ready(runner: dict[str, Any]) -> bool:
     return "RunnerReady" in runner
 
@@ -116,13 +150,12 @@ def wait_for_instance_ready(
 ) -> None:
     start_time = time.time()
     instance_existed = False
+    last_loaded: dict[str, int] = {}
     while time.time() - start_time < timeout:
-        state = client.request_json("GET", "/state")
-        instances = state.get("instances", {})
+        instance = client.get_instance(instance_id)
 
-        if instance_id not in instances:
+        if instance is None:
             if instance_existed:
-                # Instance was deleted after being created - likely due to runner failure
                 raise RuntimeError(
                     f"Instance {instance_id} was deleted (runner may have failed)"
                 )
@@ -130,18 +163,25 @@ def wait_for_instance_ready(
             continue
 
         instance_existed = True
-        instance = instances[instance_id]
-        runner_ids = runner_ids_from_instance(instance)
-        runners = state.get("runners", {})
+        rids = runner_ids_from_instance(instance)
 
-        # Check for failed runners first
-        for rid in runner_ids:
-            runner = runners.get(rid, {})
+        all_ready = True
+        for rid in rids:
+            runner = client.get_runner(rid) or {}
             if runner_failed(runner):
                 error_msg = get_runner_failed_message(runner) or "Unknown error"
                 raise RuntimeError(f"Runner {rid} failed: {error_msg}")
+            if "RunnerLoading" in runner:
+                loading = runner["RunnerLoading"]
+                loaded = loading.get("layersLoaded", 0)
+                total = loading.get("totalLayers", 0)
+                if total > 0 and last_loaded.get(rid) != loaded:
+                    last_loaded[rid] = loaded
+                    logger.debug(f"Runner {rid}: loading layers {loaded}/{total}")
+            if not runner_ready(runner):
+                all_ready = False
 
-        if all(runner_ready(runners.get(rid, {})) for rid in runner_ids):
+        if all_ready:
             return
 
         time.sleep(0.1)
@@ -165,6 +205,23 @@ def wait_for_instance_gone(
     raise TimeoutError(f"Instance {instance_id} did not get deleted within {timeout=}")
 
 
+def capture_cluster_snapshot(client: ExoClient) -> dict[str, Any]:
+    snapshot: dict[str, Any] = {}
+    identities = client.get_node_identities()
+    if identities:
+        snapshot["nodeIdentities"] = identities
+    topology = client.get_topology()
+    if topology:
+        snapshot["topology"] = topology
+    node_memory = client.get_state_path("nodeMemory")
+    if node_memory:
+        snapshot["nodeMemory"] = node_memory
+    node_system = client.get_state_path("nodeSystem")
+    if node_system:
+        snapshot["nodeSystem"] = node_system
+    return snapshot
+
+
 def resolve_model_short_id(
     client: ExoClient, model_arg: str, *, force_download: bool = False
 ) -> tuple[str, str]:
@@ -326,16 +383,11 @@ def run_planning_phase(
     node_ids = list(inner["shardAssignments"]["nodeToRunner"].keys())
     runner_to_shard = inner["shardAssignments"]["runnerToShard"]
 
-    state = client.request_json("GET", "/state")
-    downloads = state.get("downloads", {})
-    node_disk = state.get("nodeDisk", {})
-
     needs_download = False
 
     for node_id in node_ids:
-        node_downloads = downloads.get(node_id, [])
+        node_downloads = client.get_node_downloads(node_id) or []
 
-        # Check if model already downloaded on this node
         already_downloaded = any(
             "DownloadCompleted" in p
             and unwrap_instance(p["DownloadCompleted"]["shardMetadata"])["modelCard"][
@@ -349,8 +401,7 @@ def run_planning_phase(
 
         needs_download = True
 
-        # Wait for disk info if settle_deadline is set
-        disk_info = node_disk.get(node_id, {})
+        disk_info = client.get_node_disk(node_id) or {}
         backoff = _SETTLE_INITIAL_BACKOFF_S
         while not disk_info and settle_deadline and time.monotonic() < settle_deadline:
             remaining = settle_deadline - time.monotonic()
@@ -359,9 +410,7 @@ def run_planning_phase(
             )
             time.sleep(min(backoff, remaining))
             backoff = min(backoff * _SETTLE_BACKOFF_MULTIPLIER, _SETTLE_MAX_BACKOFF_S)
-            state = client.request_json("GET", "/state")
-            node_disk = state.get("nodeDisk", {})
-            disk_info = node_disk.get(node_id, {})
+            disk_info = client.get_node_disk(node_id) or {}
 
         if not disk_info:
             logger.warning(f"No disk info for {node_id}, skipping space check")
@@ -377,7 +426,6 @@ def run_planning_phase(
                 f"have {avail // (1024**3)}GB. Use --danger-delete-downloads to free space."
             )
 
-        # Delete from smallest to largest (skip read-only models)
         completed = [
             (
                 unwrap_instance(p["DownloadCompleted"]["shardMetadata"])["modelCard"][
@@ -417,21 +465,20 @@ def run_planning_phase(
     # Wait for downloads
     start = time.time()
     while time.time() - start < timeout:
-        state = client.request_json("GET", "/state")
-        downloads = state.get("downloads", {})
         all_done = True
         for node_id in node_ids:
+            node_downloads = client.get_node_downloads(node_id) or []
             done = any(
                 "DownloadCompleted" in p
                 and unwrap_instance(p["DownloadCompleted"]["shardMetadata"])[
                     "modelCard"
                 ]["modelId"]
                 == full_model_id
-                for p in downloads.get(node_id, [])
+                for p in node_downloads
             )
             failed = [
                 p["DownloadFailed"]["errorMessage"]
-                for p in downloads.get(node_id, [])
+                for p in node_downloads
                 if "DownloadFailed" in p
                 and unwrap_instance(p["DownloadFailed"]["shardMetadata"])["modelCard"][
                     "modelId"
@@ -442,6 +489,27 @@ def run_planning_phase(
                 raise RuntimeError(f"Download failed on {node_id}: {failed[0]}")
             if not done:
                 all_done = False
+                ongoing = [
+                    p
+                    for p in node_downloads
+                    if "DownloadOngoing" in p
+                    and unwrap_instance(p["DownloadOngoing"]["shardMetadata"])[
+                        "modelCard"
+                    ]["modelId"]
+                    == full_model_id
+                ]
+                if ongoing:
+                    prog = ongoing[0]["DownloadOngoing"]["downloadProgress"]
+                    speed_mb = prog.get("speed", 0) / (1024 * 1024)
+                    eta_s = prog.get("etaMs", 0) / 1000
+                    dl_bytes = prog.get("downloaded", {}).get("inBytes", 0)
+                    total_bytes = prog.get("total", {}).get("inBytes", 0)
+                    pct = (dl_bytes / total_bytes * 100) if total_bytes else 0
+                    logger.info(
+                        f"Downloading on {node_id}: {pct:.1f}% @ {speed_mb:.1f} MB/s, "
+                        f"ETA {eta_s:.0f}s "
+                        f"({prog.get('completedFiles', 0)}/{prog.get('totalFiles', 0)} files)"
+                    )
         if all_done:
             if download_t0 is not None:
                 return time.perf_counter() - download_t0