← back to Exo
auto bench (#1405)
9b5cae3db67418f6914897b54d1db0a9d9b713b6 · 2026-02-06 15:35:46 +0000 · Evan Quiney
runs exo_bench remotely with some nice git QoL
## usage
run tests/auto_bench.sh host1 [host2]
exo bench will be run on those hosts and its output saved to
bench/commit_hash/*.json on all models currently downloaded
Files touched
M .gitignoreM bench/exo_bench.pyM python/parts.nixA tests/auto_bench.shA tests/get_all_models_on_cluster.py
Diff
commit 9b5cae3db67418f6914897b54d1db0a9d9b713b6
Author: Evan Quiney <evanev7@gmail.com>
Date: Fri Feb 6 15:35:46 2026 +0000
auto bench (#1405)
runs exo_bench remotely with some nice git QoL
## usage
run tests/auto_bench.sh host1 [host2]
exo bench will be run on those hosts and its output saved to
bench/commit_hash/*.json on all models currently downloaded
---
.gitignore | 3 +++
bench/exo_bench.py | 5 +++-
python/parts.nix | 22 +++++++++++-----
tests/auto_bench.sh | 53 ++++++++++++++++++++++++++++++++++++++
tests/get_all_models_on_cluster.py | 36 ++++++++++++++++++++++++++
5 files changed, 111 insertions(+), 8 deletions(-)
diff --git a/.gitignore b/.gitignore
index 139fc326..15979672 100644
--- a/.gitignore
+++ b/.gitignore
@@ -32,3 +32,6 @@ dashboard/.svelte-kit/
# host config snapshots
hosts_*.json
.swp
+
+# bench files
+bench/**/*.json
diff --git a/bench/exo_bench.py b/bench/exo_bench.py
index 15cc3a58..eaed6cf5 100644
--- a/bench/exo_bench.py
+++ b/bench/exo_bench.py
@@ -455,6 +455,7 @@ def main() -> int:
default="bench/results.json",
help="Write raw per-run results JSON to this path.",
)
+ ap.add_argument("--stdout", action="store_true", help="Write results to stdout")
ap.add_argument(
"--dry-run", action="store_true", help="List selected placements and exit."
)
@@ -667,7 +668,9 @@ def main() -> int:
time.sleep(5)
- if args.json_out:
+ if args.stdout:
+ json.dump(all_rows, 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)
logger.debug(f"\nWrote results JSON: {args.json_out}")
diff --git a/python/parts.nix b/python/parts.nix
index d3eea784..f9e8f9e1 100644
--- a/python/parts.nix
+++ b/python/parts.nix
@@ -66,7 +66,13 @@
EXO_DASHBOARD_DIR = self'.packages.dashboard;
EXO_RESOURCES_DIR = inputs.self + /resources;
};
- text = ''python ${path}'';
+ text = ''exec python ${path} "$@"'';
+ };
+
+ mkSimplePythonScript = name: path: pkgs.writeShellApplication {
+ inherit name;
+ runtimeInputs = [ pkgs.python313 ];
+ text = ''exec python ${path} "$@"'';
};
exoPackage = pkgs.runCommand "exo"
@@ -85,12 +91,14 @@
in
{
# Python package only available on macOS (requires MLX/Metal)
- packages = lib.optionalAttrs pkgs.stdenv.hostPlatform.isDarwin {
- exo = exoPackage;
- # Test environment for running pytest outside of Nix sandbox (needs GPU access)
- exo-test-env = testVenv;
- exo-bench = mkPythonScript "exo-bench" (inputs.self + /bench/exo_bench.py);
- exo-distributed-test = mkPythonScript "exo-distributed-test" (inputs.self + /tests/headless_runner.py);
+ packages = lib.optionalAttrs pkgs.stdenv.hostPlatform.isDarwin
+ {
+ exo = exoPackage;
+ # Test environment for running pytest outside of Nix sandbox (needs GPU access)
+ exo-test-env = testVenv;
+ exo-bench = mkPythonScript "exo-bench" (inputs.self + /bench/exo_bench.py);
+ } // {
+ exo-get-all-models-on-cluster = mkSimplePythonScript "exo-get-all-models-on-cluster" (inputs.self + /tests/get_all_models_on_cluster.py);
};
checks = {
diff --git a/tests/auto_bench.sh b/tests/auto_bench.sh
new file mode 100755
index 00000000..ac6e44d9
--- /dev/null
+++ b/tests/auto_bench.sh
@@ -0,0 +1,53 @@
+#!/usr/bin/env bash
+
+[ $# -lt 1 ] && {
+ echo "Usage: $0 host1 [host2 ...]"
+ exit 1
+}
+
+[ -z "$(git status --porcelain)" ] || {
+ echo "Uncommitted changes"
+ exit 1
+}
+
+commit=$(git rev-parse HEAD)
+git fetch -q origin
+git branch -r --contains "$commit" | grep -qE '^\s*origin/' || {
+ echo "Not pushed to origin"
+ exit 1
+}
+hosts=("$@")
+cleanup() {
+ for host in "${hosts[@]}"; do
+ ssh -T -o BatchMode=yes "$host@$host" "pkill -f bin/exo" &
+ done
+ sleep 1
+ jobs -pr | xargs -r kill 2>/dev/null || true
+}
+trap 'cleanup' EXIT INT TERM
+
+for host; do
+ ssh -T -o BatchMode=yes -o ServerAliveInterval=30 "$host@$host" \
+ "EXO_LIBP2P_NAMESPACE=$commit /nix/var/nix/profiles/default/bin/nix build github:exo-explore/exo/$commit" &
+done
+wait
+for host; do
+ ssh -T -o BatchMode=yes -o ServerAliveInterval=30 "$host@$host" \
+ "EXO_LIBP2P_NAMESPACE=$commit /nix/var/nix/profiles/default/bin/nix run github:exo-explore/exo/$commit" &>/dev/null &
+done
+
+for host; do
+ echo "Waiting for $host..." 1>&2
+ until curl -sf "http://$host:52415/models" &>/dev/null; do sleep 1; done
+done
+
+echo "Waiting 30s for cluster setup" 1>&2
+sleep 30
+echo "EXO loaded" 1>&2
+bench_runner="${hosts[0]}"
+mkdir -p "./bench/$commit"
+nix run .#exo-get-all-models-on-cluster -- "$bench_runner" | while IFS= read -r model; do
+ echo "running bench for $model" 1>&2
+ ssh -Tn -o BatchMode=yes -o ServerAliveInterval=30 "$bench_runner@$bench_runner" "/nix/var/nix/profiles/default/bin/nix run github:exo-explore/exo/$commit#exo-bench -- --model $model --pp 128 4096 --tg 128 --stdout --skip-tensor-ring" >>"./bench/$commit/${model//\//--}.json"
+ echo
+done
diff --git a/tests/get_all_models_on_cluster.py b/tests/get_all_models_on_cluster.py
new file mode 100755
index 00000000..d150e9ea
--- /dev/null
+++ b/tests/get_all_models_on_cluster.py
@@ -0,0 +1,36 @@
+#!/usr/bin/env python3
+# pyright: reportAny=false
+import json
+import subprocess
+import sys
+from typing import Any, cast
+from urllib.request import urlopen
+
+h = sys.argv[1] if len(sys.argv) > 1 else sys.exit(f"USAGE: {sys.argv[0]} host")
+ts = subprocess.run(
+ ["tailscale", "status"], check=True, text=True, capture_output=True
+).stdout.splitlines()
+ip = next(
+ (sl[0] for line in ts if len(sl := line.split()) >= 2 if sl[1] == h), None
+) or sys.exit(f"{h} not found in tailscale")
+with urlopen(f"http://{ip}:52415/state", timeout=5) as r:
+ data = json.loads(r.read()).get("downloads", {})
+
+
+def mid(x: dict[str, Any]) -> str | None:
+ for k in (
+ "DownloadCompleted",
+ "shardMetadata",
+ "PipelineShardMetadata",
+ "modelCard",
+ "modelId",
+ ):
+ x = x.get(k, {})
+ return cast(str | None, x if x != {} else None)
+
+
+common = set[str].intersection(
+ *[{m for d in nid if (m := mid(d))} for nid in data.values()]
+)
+for c in common:
+ print(c)
← cf7201f9 pyproject: set minimum uv version
·
back to Exo
·
Ciaran/num sync steps (#1396) 63e9cc4f →