From d484b062e8c351ef590042e8f840c74898e1258d Mon Sep 17 00:00:00 2001 From: Jake Hillion Date: Fri, 20 Feb 2026 15:33:08 +0000 Subject: [PATCH] bench: add download timing to bench output (#1566) The bench script downloads models during the planning phase but doesn't record how long the download took, making it difficult to track download performance for a given model over time. Modified `run_planning_phase` to return download metadata: whether a fresh download occurred, the wall-clock duration, and the model size in bytes. These fields are included in every JSON output row alongside the existing per-run metrics, and a summary line is logged to the console. This allows filtering bench results by `download_occurred` and grouping by `model_id` to compute average download times across runs. Test plan: ``` # existing model jake@maverick:/data/users/jake/repos/exo/ > nix run .#exo-bench -- --host s1 --model mlx-community/gpt-oss-120b-MXFP4-Q8 --pp 128 --tg 128 ... 2026-02-20 15:23:49.081 | INFO | __main__:main:340 - Planning phase: checking downloads... 2026-02-20 15:23:49.152 | INFO | harness:run_planning_phase:402 - Started download on 12D3KooWKx41iikn188ozrxSdoG26g88jFCfie9wEA1eQR8csbPm 2026-02-20 15:23:49.184 | INFO | __main__:main:352 - Download: model already cached ... Wrote results JSON: bench/results.json jake@maverick:/data/users/jake/repos/exo/ > cat bench/results.json [ { "elapsed_s": 2.9446684420108795, "output_text_preview": "The user just typed a long series of \"a\". Possibly they are testing. There's no explicit question. Could be they want a response? Might be a test of handling long input. We can respond politely, ask i", "stats": { "prompt_tps": 117.7872141515621, "generation_tps": 85.49598231498028, "prompt_tokens": 129, "generation_tokens": 128, "peak_memory_usage": { "inBytes": 68215145744 } }, "model_short_id": "gpt-oss-120b-MXFP4-Q8", "model_id": "mlx-community/gpt-oss-120b-MXFP4-Q8", "placement_sharding": "Pipeline", "placement_instance_meta": "MlxRing", "placement_nodes": 1, "instance_id": "68babc2a-6e94-4c70-aa07-7ec681f7c856", "pp_tokens": 128, "tg": 128, "repeat_index": 0 } ]% # no change to output ``` ``` # missing model jake@maverick:/data/users/jake/repos/exo/ > nix run .#exo-bench -- --host s1 --model mlx-community/Meta-Llama-3.1-8B-Instruct-4bit --pp 128 --tg 128 ... 2026-02-20 15:24:42.553 | INFO | __main__:main:340 - Planning phase: checking downloads... 2026-02-20 15:24:42.625 | INFO | harness:run_planning_phase:402 - Started download on 12D3KooWKx41iikn188ozrxSdoG26g88jFCfie9wEA1eQR8csbPm 2026-02-20 15:25:37.494 | INFO | __main__:main:350 - Download: 54.9s (freshly downloaded) ... Wrote results JSON: bench/results.json jake@maverick:/data/users/jake/repos/exo/ > cat bench/results.json [ { "elapsed_s": 1.500349276990164, "output_text_preview": "It seems like you've entered a large number of 'a's. If you'd like to discuss something or ask a question, I'm here to help. If not, is there anything else I can assist you with? \n\nIf you're intereste", "stats": { "prompt_tps": 395.43264952543666, "generation_tps": 128.03520443181478, "prompt_tokens": 129, "generation_tokens": 128, "peak_memory_usage": { "inBytes": 5116952079 } }, "model_short_id": "Meta-Llama-3.1-8B-Instruct-4bit", "model_id": "mlx-community/Meta-Llama-3.1-8B-Instruct-4bit", "placement_sharding": "Pipeline", "placement_instance_meta": "MlxRing", "placement_nodes": 1, "instance_id": "ccd9bd71-d4cc-4b75-a37f-98090544626a", "pp_tokens": 128, "tg": 128, "repeat_index": 0, "download_duration_s": 54.88322358299047 } ]% # one new field ``` --- bench/exo_bench.py | 11 ++++++++++- bench/harness.py | 19 +++++++++++++++---- 2 files changed, 25 insertions(+), 5 deletions(-) diff --git a/bench/exo_bench.py b/bench/exo_bench.py index f6fcd342..e60d0872 100644 --- a/bench/exo_bench.py +++ b/bench/exo_bench.py @@ -338,7 +338,7 @@ def main() -> int: ) logger.info("Planning phase: checking downloads...") - run_planning_phase( + download_duration_s = run_planning_phase( client, full_model_id, selected[0], @@ -346,6 +346,10 @@ def main() -> int: args.timeout, settle_deadline, ) + if download_duration_s is not None: + logger.info(f"Download: {download_duration_s:.1f}s (freshly downloaded)") + else: + logger.info("Download: model already cached") all_rows: list[dict[str, Any]] = [] @@ -409,6 +413,11 @@ def main() -> int: "pp_tokens": actual_pp_tokens, "tg": tg, "repeat_index": r, + **( + {"download_duration_s": download_duration_s} + if download_duration_s is not None + else {} + ), } ) runs.append(row) diff --git a/bench/harness.py b/bench/harness.py index 58aa8435..263c8e91 100644 --- a/bench/harness.py +++ b/bench/harness.py @@ -289,8 +289,12 @@ def run_planning_phase( danger_delete: bool, timeout: float, settle_deadline: float | None, -) -> None: - """Check disk space and ensure model is downloaded before benchmarking.""" +) -> float | None: + """Check disk space and ensure model is downloaded before benchmarking. + + Returns the wall-clock download duration in seconds if a fresh download + was needed, or None if the model was already cached on all nodes. + """ # Get model size from /models models = client.request_json("GET", "/models") or {} model_bytes = 0 @@ -303,7 +307,7 @@ def run_planning_phase( logger.warning( f"Could not determine size for {full_model_id}, skipping disk check" ) - return + return None # Get nodes from preview inner = unwrap_instance(preview["instance"]) @@ -314,6 +318,8 @@ def run_planning_phase( downloads = state.get("downloads", {}) node_disk = state.get("nodeDisk", {}) + needs_download = False + for node_id in node_ids: node_downloads = downloads.get(node_id, []) @@ -329,6 +335,8 @@ def run_planning_phase( if already_downloaded: continue + needs_download = True + # Wait for disk info if settle_deadline is set disk_info = node_disk.get(node_id, {}) backoff = _SETTLE_INITIAL_BACKOFF_S @@ -379,6 +387,7 @@ def run_planning_phase( raise RuntimeError(f"Could not free enough space on {node_id}") # Start downloads (idempotent) + download_t0 = time.perf_counter() if needs_download else None for node_id in node_ids: runner_id = inner["shardAssignments"]["nodeToRunner"][node_id] shard = runner_to_shard[runner_id] @@ -421,7 +430,9 @@ def run_planning_phase( if not done: all_done = False if all_done: - return + if download_t0 is not None: + return time.perf_counter() - download_t0 + return None time.sleep(1) raise TimeoutError("Downloads did not complete in time")