← back to Exo
bench: restore --danger-delete-downloads planning phase (#1542)
42e1e7322bc0fef65df441b11668ff6b0391c903 · 2026-02-19 15:42:02 +0000 · Jake Hillion
c2f2111b extracted shared utilities from exo_bench.py into harness.py
but accidentally dropped the run_planning_phase function and
--danger-delete-downloads CLI argument in the process.
Restored run_planning_phase in harness.py (where its dependencies now
live) and re-added the --danger-delete-downloads argument to
add_common_instance_args. Re-wired the planning phase call in
exo_bench.py's main() before the benchmark loop.
Files touched
M bench/eval_tool_calls.pyM bench/exo_bench.pyM bench/harness.py
Diff
commit 42e1e7322bc0fef65df441b11668ff6b0391c903
Author: Jake Hillion <jake@hillion.co.uk>
Date: Thu Feb 19 15:42:02 2026 +0000
bench: restore --danger-delete-downloads planning phase (#1542)
c2f2111b extracted shared utilities from exo_bench.py into harness.py
but accidentally dropped the run_planning_phase function and
--danger-delete-downloads CLI argument in the process.
Restored run_planning_phase in harness.py (where its dependencies now
live) and re-added the --danger-delete-downloads argument to
add_common_instance_args. Re-wired the planning phase call in
exo_bench.py's main() before the benchmark loop.
---
bench/eval_tool_calls.py | 16 +++++
bench/exo_bench.py | 15 +++++
bench/harness.py | 150 +++++++++++++++++++++++++++++++++++++++++++++++
3 files changed, 181 insertions(+)
diff --git a/bench/eval_tool_calls.py b/bench/eval_tool_calls.py
index 2d55d6ca..13cb2537 100644
--- a/bench/eval_tool_calls.py
+++ b/bench/eval_tool_calls.py
@@ -20,6 +20,7 @@ from harness import (
instance_id_from_instance,
nodes_used_in_instance,
resolve_model_short_id,
+ run_planning_phase,
settle_and_fetch_placements,
wait_for_instance_gone,
wait_for_instance_ready,
@@ -962,6 +963,21 @@ Examples:
selected.sort(key=_placement_sort_key)
preview = selected[0]
+
+ settle_deadline = (
+ time.monotonic() + args.settle_timeout if args.settle_timeout > 0 else None
+ )
+
+ print("Planning phase: checking downloads...", file=log)
+ run_planning_phase(
+ exo,
+ full_model_id,
+ preview,
+ args.danger_delete_downloads,
+ args.timeout,
+ settle_deadline,
+ )
+
instance = preview["instance"]
instance_id = instance_id_from_instance(instance)
sharding = str(preview["sharding"])
diff --git a/bench/exo_bench.py b/bench/exo_bench.py
index 9f6f5f02..f6fcd342 100644
--- a/bench/exo_bench.py
+++ b/bench/exo_bench.py
@@ -35,6 +35,7 @@ from harness import (
instance_id_from_instance,
nodes_used_in_instance,
resolve_model_short_id,
+ run_planning_phase,
settle_and_fetch_placements,
wait_for_instance_gone,
wait_for_instance_ready,
@@ -332,6 +333,20 @@ def main() -> int:
if args.dry_run:
return 0
+ settle_deadline = (
+ time.monotonic() + args.settle_timeout if args.settle_timeout > 0 else None
+ )
+
+ 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:
diff --git a/bench/harness.py b/bench/harness.py
index c8ae9318..58aa8435 100644
--- a/bench/harness.py
+++ b/bench/harness.py
@@ -282,6 +282,151 @@ def settle_and_fetch_placements(
return selected
+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 add_common_instance_args(ap: argparse.ArgumentParser) -> None:
ap.add_argument("--host", default=os.environ.get("EXO_HOST", "localhost"))
ap.add_argument(
@@ -325,3 +470,8 @@ def add_common_instance_args(ap: argparse.ArgumentParser) -> None:
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.",
+ )
← aa3f106f fix: import ResponsesStreamEvent and DRY up SSE formatting (
·
back to Exo
·
simplify rust ident module 4c4c6ce9 →