[object Object]

← back to Exo

bench: add download timing to bench output (#1566)

d484b062e8c351ef590042e8f840c74898e1258d · 2026-02-20 15:33:08 +0000 · Jake Hillion

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
```

Files touched

Diff

commit d484b062e8c351ef590042e8f840c74898e1258d
Author: Jake Hillion <jake@hillion.co.uk>
Date:   Fri Feb 20 15:33:08 2026 +0000

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

← e32b649d fix: enable psutil fallback for memory monitoring when macmo  ·  back to Exo  ·  add system ids (#1536) dc8d42b4 →