From 131fb141a61fc0fcddb5737022b6d54a7da65eaf Mon Sep 17 00:00:00 2001 From: Jake Hillion Date: Wed, 11 Feb 2026 21:58:45 +0000 Subject: [PATCH] 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 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: