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")