← back to Exo
bench: add --danger-delete-downloads flag with planning phase
131fb141a61fc0fcddb5737022b6d54a7da65eaf · 2026-02-11 21:58:45 +0000 · Jake Hillion
exo bench previously relied on the worker's plan loop to download
models, which could fail silently or run into disk space issues during
benchmarking. This made it difficult to diagnose download failures.
Added a planning phase that runs before benchmarking to explicitly
handle downloads. It checks available disk space on each node via the
/state endpoint and starts downloads via POST /download/start. When
the --danger-delete-downloads flag is set and there's insufficient
space, it deletes existing models from smallest to largest until
there's room for the benchmark model.
Test plan:
- CI
```
jake@maverick:/data/users/jake/repos/exo/ > nix run .#exo-bench -- --pp 128,2048,4096 --tg 128 --stdout --settle-timeout 10 --host s1 --model mlx-community/gpt-oss-120b-MXFP4-Q8
PyTorch was not found. Models won't be available and only tokenizers, configuration and file/data utilities can be used.
2026-02-16 12:12:11.807 | INFO | __main__:main:710 - pp/tg mode: combinations (product) - 3 pairs
Warning: You are sending unauthenticated requests to the HF Hub. Please set a HF_TOKEN to enable higher rate limits and faster downloads.
2026-02-16 12:12:13.455 | DEBUG | __main__:main:725 - [exo-bench] loaded tokenizer: mlx-community/gpt-oss-120b-MXFP4-Q8 for prompt sizer
2026-02-16 12:12:13.473 | DEBUG | __main__:main:761 - exo-bench model: short_id=gpt-oss-120b-MXFP4-Q8 full_id=mlx-community/gpt-oss-120b-MXFP4-Q8
2026-02-16 12:12:13.473 | INFO | __main__:main:762 - placements: 1
2026-02-16 12:12:13.474 | INFO | __main__:main:764 - - Pipeline / MlxRing / nodes=1
2026-02-16 12:12:13.474 | INFO | __main__:main:771 - Planning phase: checking downloads...
Traceback (most recent call last):
File "/nix/store/q31kmbcfr5bf97290bvbnhrvpc3fv824-source/bench/exo_bench.py", line 885, in <module>
raise SystemExit(main())
~~~~^^
File "/nix/store/q31kmbcfr5bf97290bvbnhrvpc3fv824-source/bench/exo_bench.py", line 772, in main
run_planning_phase(
~~~~~~~~~~~~~~~~~~^
client,
^^^^^^^
...<4 lines>...
settle_deadline,
^^^^^^^^^^^^^^^^
)
^
File "/nix/store/q31kmbcfr5bf97290bvbnhrvpc3fv824-source/bench/exo_bench.py", line 367, in run_planning_phase
raise RuntimeError(
...<2 lines>...
)
RuntimeError: Insufficient disk on 12D3KooWE2C7dzC9d9YJMEfWK3g8og7JdZj3HHXZ8VmGrXYAEnEj: need 65GB, have 55GB. Use --danger-delete-downloads to free space.
jake@maverick:/data/users/jake/repos/exo/ > nix run .#exo-bench -- --pp 128,2048,4096 --tg 128 --stdout --settle-timeout 10 --host s1 --model mlx-community/gpt-oss-120b-MXFP4-Q8 --danger-delete-downloads
PyTorch was not found. Models won't be available and only tokenizers, configuration and file/data utilities can be used.
2026-02-16 12:12:19.626 | INFO | __main__:main:710 - pp/tg mode: combinations (product) - 3 pairs
2026-02-16 12:12:21.262 | DEBUG | __main__:main:725 - [exo-bench] loaded tokenizer: mlx-community/gpt-oss-120b-MXFP4-Q8 for prompt sizer
2026-02-16 12:12:21.280 | DEBUG | __main__:main:761 - exo-bench model: short_id=gpt-oss-120b-MXFP4-Q8 full_id=mlx-community/gpt-oss-120b-MXFP4-Q8
2026-02-16 12:12:21.280 | INFO | __main__:main:762 - placements: 1
2026-02-16 12:12:21.280 | INFO | __main__:main:764 - - Pipeline / MlxRing / nodes=1
2026-02-16 12:12:21.280 | INFO | __main__:main:771 - Planning phase: checking downloads...
2026-02-16 12:12:21.336 | INFO | __main__:run_planning_phase:386 - Deleting mlx-community/Qwen3-0.6B-4bit from 12D3KooWE2C7dzC9d9YJMEfWK3g8og7JdZj3HHXZ8VmGrXYAEnEj (335MB)
2026-02-16 12:12:21.350 | INFO | __main__:run_planning_phase:386 - Deleting mlx-community/Llama-3.2-1B-Instruct-4bit from 12D3KooWE2C7dzC9d9YJMEfWK3g8og7JdZj3HHXZ8VmGrXYAEnEj (679MB)
2026-02-16 12:12:21.363 | INFO | __main__:run_planning_phase:386 - Deleting mlx-community/Llama-3.2-3B-Instruct-4bit from 12D3KooWE2C7dzC9d9YJMEfWK3g8og7JdZj3HHXZ8VmGrXYAEnEj (1740MB)
2026-02-16 12:12:21.373 | INFO | __main__:run_planning_phase:386 - Deleting mlx-community/Llama-3.2-3B-Instruct-8bit from 12D3KooWE2C7dzC9d9YJMEfWK3g8og7JdZj3HHXZ8VmGrXYAEnEj (3264MB)
2026-02-16 12:12:21.384 | INFO | __main__:run_planning_phase:386 - Deleting mlx-community/GLM-4.7-Flash-8bit from 12D3KooWE2C7dzC9d9YJMEfWK3g8og7JdZj3HHXZ8VmGrXYAEnEj (30366MB)
2026-02-16 12:12:21.413 | INFO | __main__:run_planning_phase:407 - Started download on 12D3KooWE2C7dzC9d9YJMEfWK3g8og7JdZj3HHXZ8VmGrXYAEnEj
```
It's not pretty but it works!
Files touched
Diff
commit 131fb141a61fc0fcddb5737022b6d54a7da65eaf
Author: Jake Hillion <jake@hillion.co.uk>
Date: Wed Feb 11 21:58:45 2026 +0000
bench: add --danger-delete-downloads flag with planning phase
exo bench previously relied on the worker's plan loop to download
models, which could fail silently or run into disk space issues during
benchmarking. This made it difficult to diagnose download failures.
Added a planning phase that runs before benchmarking to explicitly
handle downloads. It checks available disk space on each node via the
/state endpoint and starts downloads via POST /download/start. When
the --danger-delete-downloads flag is set and there's insufficient
space, it deletes existing models from smallest to largest until
there's room for the benchmark model.
Test plan:
- CI
```
jake@maverick:/data/users/jake/repos/exo/ > nix run .#exo-bench -- --pp 128,2048,4096 --tg 128 --stdout --settle-timeout 10 --host s1 --model mlx-community/gpt-oss-120b-MXFP4-Q8
PyTorch was not found. Models won't be available and only tokenizers, configuration and file/data utilities can be used.
2026-02-16 12:12:11.807 | INFO | __main__:main:710 - pp/tg mode: combinations (product) - 3 pairs
Warning: You are sending unauthenticated requests to the HF Hub. Please set a HF_TOKEN to enable higher rate limits and faster downloads.
2026-02-16 12:12:13.455 | DEBUG | __main__:main:725 - [exo-bench] loaded tokenizer: mlx-community/gpt-oss-120b-MXFP4-Q8 for prompt sizer
2026-02-16 12:12:13.473 | DEBUG | __main__:main:761 - exo-bench model: short_id=gpt-oss-120b-MXFP4-Q8 full_id=mlx-community/gpt-oss-120b-MXFP4-Q8
2026-02-16 12:12:13.473 | INFO | __main__:main:762 - placements: 1
2026-02-16 12:12:13.474 | INFO | __main__:main:764 - - Pipeline / MlxRing / nodes=1
2026-02-16 12:12:13.474 | INFO | __main__:main:771 - Planning phase: checking downloads...
Traceback (most recent call last):
File "/nix/store/q31kmbcfr5bf97290bvbnhrvpc3fv824-source/bench/exo_bench.py", line 885, in <module>
raise SystemExit(main())
~~~~^^
File "/nix/store/q31kmbcfr5bf97290bvbnhrvpc3fv824-source/bench/exo_bench.py", line 772, in main
run_planning_phase(
~~~~~~~~~~~~~~~~~~^
client,
^^^^^^^
...<4 lines>...
settle_deadline,
^^^^^^^^^^^^^^^^
)
^
File "/nix/store/q31kmbcfr5bf97290bvbnhrvpc3fv824-source/bench/exo_bench.py", line 367, in run_planning_phase
raise RuntimeError(
...<2 lines>...
)
RuntimeError: Insufficient disk on 12D3KooWE2C7dzC9d9YJMEfWK3g8og7JdZj3HHXZ8VmGrXYAEnEj: need 65GB, have 55GB. Use --danger-delete-downloads to free space.
jake@maverick:/data/users/jake/repos/exo/ > nix run .#exo-bench -- --pp 128,2048,4096 --tg 128 --stdout --settle-timeout 10 --host s1 --model mlx-community/gpt-oss-120b-MXFP4-Q8 --danger-delete-downloads
PyTorch was not found. Models won't be available and only tokenizers, configuration and file/data utilities can be used.
2026-02-16 12:12:19.626 | INFO | __main__:main:710 - pp/tg mode: combinations (product) - 3 pairs
2026-02-16 12:12:21.262 | DEBUG | __main__:main:725 - [exo-bench] loaded tokenizer: mlx-community/gpt-oss-120b-MXFP4-Q8 for prompt sizer
2026-02-16 12:12:21.280 | DEBUG | __main__:main:761 - exo-bench model: short_id=gpt-oss-120b-MXFP4-Q8 full_id=mlx-community/gpt-oss-120b-MXFP4-Q8
2026-02-16 12:12:21.280 | INFO | __main__:main:762 - placements: 1
2026-02-16 12:12:21.280 | INFO | __main__:main:764 - - Pipeline / MlxRing / nodes=1
2026-02-16 12:12:21.280 | INFO | __main__:main:771 - Planning phase: checking downloads...
2026-02-16 12:12:21.336 | INFO | __main__:run_planning_phase:386 - Deleting mlx-community/Qwen3-0.6B-4bit from 12D3KooWE2C7dzC9d9YJMEfWK3g8og7JdZj3HHXZ8VmGrXYAEnEj (335MB)
2026-02-16 12:12:21.350 | INFO | __main__:run_planning_phase:386 - Deleting mlx-community/Llama-3.2-1B-Instruct-4bit from 12D3KooWE2C7dzC9d9YJMEfWK3g8og7JdZj3HHXZ8VmGrXYAEnEj (679MB)
2026-02-16 12:12:21.363 | INFO | __main__:run_planning_phase:386 - Deleting mlx-community/Llama-3.2-3B-Instruct-4bit from 12D3KooWE2C7dzC9d9YJMEfWK3g8og7JdZj3HHXZ8VmGrXYAEnEj (1740MB)
2026-02-16 12:12:21.373 | INFO | __main__:run_planning_phase:386 - Deleting mlx-community/Llama-3.2-3B-Instruct-8bit from 12D3KooWE2C7dzC9d9YJMEfWK3g8og7JdZj3HHXZ8VmGrXYAEnEj (3264MB)
2026-02-16 12:12:21.384 | INFO | __main__:run_planning_phase:386 - Deleting mlx-community/GLM-4.7-Flash-8bit from 12D3KooWE2C7dzC9d9YJMEfWK3g8og7JdZj3HHXZ8VmGrXYAEnEj (30366MB)
2026-02-16 12:12:21.413 | INFO | __main__:run_planning_phase:407 - Started download on 12D3KooWE2C7dzC9d9YJMEfWK3g8og7JdZj3HHXZ8VmGrXYAEnEj
```
It's not pretty but it works!
---
bench/exo_bench.py | 171 +++++++++++++++++++++++++++++++++++++++++++++++++++--
1 file changed, 167 insertions(+), 4 deletions(-)
diff --git a/bench/exo_bench.py b/bench/exo_bench.py
index 8358f000..56bb0421 100644
--- a/bench/exo_bench.py
+++ b/bench/exo_bench.py
@@ -288,6 +288,151 @@ def resolve_model_short_id(client: ExoClient, model_arg: str) -> tuple[str, str]
raise ValueError(f"Model not found in /models: {model_arg}")
+def run_planning_phase(
+ client: ExoClient,
+ full_model_id: str,
+ preview: dict[str, Any],
+ danger_delete: bool,
+ timeout: float,
+ settle_deadline: float | None,
+) -> None:
+ """Check disk space and ensure model is downloaded before benchmarking."""
+ # Get model size from /models
+ models = client.request_json("GET", "/models") or {}
+ model_bytes = 0
+ for m in models.get("data", []):
+ if m.get("hugging_face_id") == full_model_id:
+ model_bytes = m.get("storage_size_megabytes", 0) * 1024 * 1024
+ break
+
+ if not model_bytes:
+ logger.warning(
+ f"Could not determine size for {full_model_id}, skipping disk check"
+ )
+ return
+
+ # Get nodes from preview
+ inner = unwrap_instance(preview["instance"])
+ 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", {})
+
+ for node_id in node_ids:
+ node_downloads = downloads.get(node_id, [])
+
+ # Check if model already downloaded on this node
+ already_downloaded = any(
+ "DownloadCompleted" in p
+ and unwrap_instance(p["DownloadCompleted"]["shardMetadata"])["modelCard"][
+ "modelId"
+ ]
+ == full_model_id
+ for p in node_downloads
+ )
+ if already_downloaded:
+ continue
+
+ # Wait for disk info if settle_deadline is set
+ disk_info = node_disk.get(node_id, {})
+ backoff = _SETTLE_INITIAL_BACKOFF_S
+ while not disk_info and settle_deadline and time.monotonic() < settle_deadline:
+ remaining = settle_deadline - time.monotonic()
+ logger.info(
+ f"Waiting for disk info on {node_id} ({remaining:.0f}s remaining)..."
+ )
+ 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, {})
+
+ if not disk_info:
+ logger.warning(f"No disk info for {node_id}, skipping space check")
+ continue
+
+ avail = disk_info.get("available", {}).get("inBytes", 0)
+ if avail >= model_bytes:
+ continue
+
+ if not danger_delete:
+ raise RuntimeError(
+ f"Insufficient disk on {node_id}: need {model_bytes // (1024**3)}GB, "
+ f"have {avail // (1024**3)}GB. Use --danger-delete-downloads to free space."
+ )
+
+ # Delete from smallest to largest
+ completed = [
+ (
+ unwrap_instance(p["DownloadCompleted"]["shardMetadata"])["modelCard"][
+ "modelId"
+ ],
+ p["DownloadCompleted"]["totalBytes"]["inBytes"],
+ )
+ for p in node_downloads
+ if "DownloadCompleted" in p
+ ]
+ for del_model, size in sorted(completed, key=lambda x: x[1]):
+ logger.info(f"Deleting {del_model} from {node_id} ({size // (1024**2)}MB)")
+ client.request_json("DELETE", f"/download/{node_id}/{del_model}")
+ avail += size
+ if avail >= model_bytes:
+ break
+
+ if avail < model_bytes:
+ raise RuntimeError(f"Could not free enough space on {node_id}")
+
+ # Start downloads (idempotent)
+ for node_id in node_ids:
+ runner_id = inner["shardAssignments"]["nodeToRunner"][node_id]
+ shard = runner_to_shard[runner_id]
+ client.request_json(
+ "POST",
+ "/download/start",
+ body={
+ "targetNodeId": node_id,
+ "shardMetadata": shard,
+ },
+ )
+ logger.info(f"Started download on {node_id}")
+
+ # 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:
+ done = any(
+ "DownloadCompleted" in p
+ and unwrap_instance(p["DownloadCompleted"]["shardMetadata"])[
+ "modelCard"
+ ]["modelId"]
+ == full_model_id
+ for p in downloads.get(node_id, [])
+ )
+ failed = [
+ p["DownloadFailed"]["errorMessage"]
+ for p in downloads.get(node_id, [])
+ if "DownloadFailed" in p
+ and unwrap_instance(p["DownloadFailed"]["shardMetadata"])["modelCard"][
+ "modelId"
+ ]
+ == full_model_id
+ ]
+ if failed:
+ raise RuntimeError(f"Download failed on {node_id}: {failed[0]}")
+ if not done:
+ all_done = False
+ if all_done:
+ return
+ time.sleep(1)
+
+ raise TimeoutError("Downloads did not complete in time")
+
+
def placement_filter(instance_meta: str, wanted: str) -> bool:
s = (instance_meta or "").lower()
if wanted == "both":
@@ -535,6 +680,11 @@ def main() -> int:
default=0,
help="Max seconds to wait for the cluster to produce valid placements (0 = try once).",
)
+ ap.add_argument(
+ "--danger-delete-downloads",
+ action="store_true",
+ help="Delete existing models from smallest to largest to make room for benchmark model.",
+ )
args = ap.parse_args()
pp_list = parse_int_list(args.pp)
@@ -569,13 +719,16 @@ def main() -> int:
logger.error("[exo-bench] tokenizer usable but prompt sizing failed")
raise
+ settle_deadline = (
+ time.monotonic() + args.settle_timeout if args.settle_timeout > 0 else None
+ )
+
selected = fetch_and_filter_placements(client, full_model_id, args)
- if not selected and args.settle_timeout > 0:
+ if not selected and settle_deadline:
backoff = _SETTLE_INITIAL_BACKOFF_S
- deadline = time.monotonic() + args.settle_timeout
- while not selected and time.monotonic() < deadline:
- remaining = deadline - time.monotonic()
+ while not selected and time.monotonic() < settle_deadline:
+ remaining = settle_deadline - time.monotonic()
logger.warning(
f"No valid placements yet (cluster may still be settling). "
f"Retrying in {backoff:.1f}s ({remaining:.0f}s remaining)..."
@@ -607,6 +760,16 @@ def main() -> int:
if args.dry_run:
return 0
+ logger.info("Planning phase: checking downloads...")
+ run_planning_phase(
+ client,
+ full_model_id,
+ selected[0],
+ args.danger_delete_downloads,
+ args.timeout,
+ settle_deadline,
+ )
+
all_rows: list[dict[str, Any]] = []
for preview in selected:
← 2d8bfc2e fix: PlaceInstanceParams broken field validator
·
back to Exo
·
api cancellation (#1276) 2759e923 →