← back to Exo
Improve exo harness with path state (#1815)
2efbb8ab4fc806971f9221b3d9ea80e490cf7670 · 2026-03-30 11:20:46 +0100 · rltakashige
<img width="3224" height="1476" alt="image"
src="https://github.com/user-attachments/assets/d90a7d8a-9fe5-43a1-a715-1ef7ecc15422"
/>
Files touched
M bench/eval_tool_calls.pyM bench/exo_bench.pyM bench/exo_eval.pyM bench/harness.py
Diff
commit 2efbb8ab4fc806971f9221b3d9ea80e490cf7670
Author: rltakashige <rl.takashige@gmail.com>
Date: Mon Mar 30 11:20:46 2026 +0100
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 | 9 ++-
bench/exo_bench.py | 147 +++++++++++++++++++++++++++++++++++++++++++++--
bench/exo_eval.py | 9 ++-
bench/harness.py | 124 ++++++++++++++++++++++++++++++---------
4 files changed, 254 insertions(+), 35 deletions(-)
diff --git a/bench/eval_tool_calls.py b/bench/eval_tool_calls.py
index 13cb2537..18071473 100644
--- a/bench/eval_tool_calls.py
+++ b/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)
diff --git a/bench/exo_bench.py b/bench/exo_bench.py
index 2daa4927..97b6b3ac 100644
--- a/bench/exo_bench.py
+++ b/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
diff --git a/bench/exo_eval.py b/bench/exo_eval.py
index e0770e75..fb4d55f3 100644
--- a/bench/exo_eval.py
+++ b/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:
diff --git a/bench/harness.py b/bench/harness.py
index 3d771de7..440a97cf 100644
--- a/bench/harness.py
+++ b/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 all(runner_ready(runners.get(rid, {})) for rid in runner_ids):
+ 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_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
← c6c5a3e7 feat: /state/paths (#1796)
·
back to Exo
·
Add multimodality! (#1802) 635801d5 →