Warmup = 2

## Motivation

Process-isolated runner crashes and C-extension failures can write
directly to fd-level stdout/stderr, bypassing Python/loguru. We need to
capture that output per runner process without polluting the main
process or other workers, and without breaking operation when the parent
stdio is detached.

## Changes

- Added `AsyncProcess`, a spawn-only multiprocessing wrapper that
redirects child stdout/stderr to pipes and exposes them as in-memory
`Receiver[bytes]`s
- Replaced runner-supervisor's raw `multiprocessing.Process` usage with
`AsyncProcess`
- Added `--no-stdio`, redirecting stdin/stdout/stderr to `/dev/null`
after logging is configured
- Disabled verbose MLX
- Added tests covering stdio capture, child crashes, repeated bad
children, SIGTERM/SIGKILL shutdown escalation, stdio detachment, and
spawning captured children from a stdio-detached parent

## Why It Works

The parent can redirect its own stdio fds to `/dev/null`, while
`AsyncProcess` installs fresh pipe fds over fd 1 and 2 inside each
spawned child. That keeps stdio-detached parents quiet while preserving
per-runner stdout/stderr capture. Runner shutdown is still bounded:
SIGTERM grace first, then SIGKILL escalation if needed.

Next direction: the runner supervisor currently drains captured output
and logs it as stdout/debug and stderr/warning. This should be split
into more useful process-isolated error reporting instead of just log
forwarding (regex match on errors to obtain "reason" string, best
effort).

## Test Plan

### Manual Testing

Ran on 4 Mac Minis in a Thunderbolt 4 ring, can see that runner's
stdout/stderr contents are being captured.

### Automated Testing

- Added async-process tests for fd-level stdout/stderr capture, Python
traceback capture, bounded-buffer output, child `exit`/abort, parent
stdio preservation, fd leak checks, spawn-context mp channels, and
SIGTERM/SIGKILL shutdown behavior
- Added stdio-detach tests proving stdio detaches to `/dev/null`, a
stdio-detached parent can still spawn and capture a child, and the same
stdio-detached parent can spawn/capture multiple children sequentially
- Updated runner-supervisor tests for the new `AsyncProcess.exitcode`
path

.gitignore

@@ -40,3 +40,4 @@ bench/**/*.json
 tmp/models
 /build/exo
 /.claude/skills
+/.claude

AGENTS.md

@@ -120,6 +120,81 @@ From .cursorrules:
 
 Tests use pytest-asyncio with `asyncio_mode = "auto"`. Tests are in `tests/` subdirectories alongside the code they test. The `EXO_TESTS=1` env var is set during tests.
 
+Integration tests live in `tests/` (root) and are opt-in via `--ignore=tests` in the default pytest addopts. They require an `eco`-managed cluster:
+
+```bash
+uv run pytest tests/ -v                    # constraint-driven host pick
+uv run pytest tests/ -v --hosts s4         # explicit host override
+```
+
+## Benchmarking
+
+Benchmarks live in `bench/`. The framework is a CLI with subcommands; each benchmark is a small library module under `bench/lib/<name>.py` plus a CLI front-end under `bench/cli/<name>.py`.
+
+```
+bench/
+├── lib/                       # composable, typed building blocks
+│   ├── prompt.py              # PromptSizer, load_tokenizer_for_bench
+│   ├── completion.py          # run_one_completion + typed payloads
+│   ├── session.py             # BenchSession (cluster + client + instance)
+│   ├── results.py             # RunMetadata, ResultsBundle, JSON schema
+│   ├── model_meta.py          # HF API: total weights size, max context, layers
+│   ├── cluster.py             # managed_cluster + managed_instance ctx-managers
+│   └── context_scaling.py     # prompt-TPS / decode-TPS vs context-size sweep
+├── cli/                       # CLI subcommands
+│   ├── _common.py             # shared argparse args + SharedOptions
+│   ├── context_scaling.py     # `python -m bench.cli context-scaling …`
+│   └── __main__.py            # subcommand dispatcher
+└── exo_bench.py, prefill_decode_bench.py
+                                # legacy CLI scripts; PromptSizer / run_one_completion
+                                # / load_tokenizer_for_bench are re-exports of bench.lib.
+```
+
+Run a benchmark:
+
+```bash
+# Defaults assume a multi-node, Thunderbolt-connected cluster with tensor
+# parallelism + JACCL: --sharding Tensor --comm MlxJaccl --thunderbolt a2a.
+# Memory + disk minimums are auto-derived from HF metadata.
+uv run python -m bench.cli context-scaling \
+  --model mlx-community/Qwen3-30B-A3B-4bit --nodes 2 --num-steps 32
+
+# Single-node smoke: opt out of TB / tensor / jaccl
+uv run python -m bench.cli context-scaling --hosts s4 \
+  --model mlx-community/Llama-3.2-1B-Instruct-4bit --num-steps 4 \
+  --sharding Pipeline --comm MlxRing --thunderbolt none
+
+# From a TOML config (CLI flags override config values)
+uv run python -m bench.cli context-scaling \
+  --config bench/configs/context_scaling.example.toml --hosts s4,s9
+```
+
+Shared CLI flags (every subcommand inherits these via `bench/cli/_common.py`):
+- `--config <path>.toml` — load run parameters from a TOML file
+- `--model`, `--sharding {Pipeline,Tensor}` (default Tensor), `--comm {MlxRing,MlxJaccl}` (default MlxJaccl), `--min-nodes` — placement
+- `--hosts`, `--nodes` (number of cluster hosts; distinct from `--min-nodes`), `--thunderbolt {a2a,ring,none}` (default a2a), `--chip` — host pool
+- `--min-memory-gb`, `--max-memory-gb`, `--min-disk-gb`, `--max-disk-gb` (minimums auto-derived from HF model size when not supplied)
+- `--evict-downloads` (default on; auto-evicts smallest-first when disk is short)
+- `--cleanup-instance` (default on; deletes the instance on exit)
+- `--output-dir`, `--tag key=value` (repeatable)
+
+Run a multi-run campaign from a single TOML file (each `[[runs]]` = its own cluster deploy + bench + teardown; `[defaults]` is shared, per-run keys override; `[plot]` triggers a comparison PNG):
+
+```bash
+uv run python -m bench.cli campaign bench/configs/llama-family-smoke.toml
+```
+
+Plot any results JSON to a PNG (auto-detects benchmark type from `metadata.benchmark`):
+
+```bash
+uv run python -m bench.cli plot bench/results/context_scaling/latest.json
+uv run python -m bench.cli plot a.json b.json --label-tag operator   # multi-run comparison
+```
+
+Adding a new benchmark = (1) write a `bench/lib/<name>.py` exposing a typed `run(session, params, bundle)` callable; (2) add a `bench/cli/<name>.py` with `add_subparser(...)` + `run(args) -> Path`; (3) register the imports in `bench/cli/__main__.py`. To enable plotting for the new benchmark, add a `render_<name>(inputs)` function in `bench/lib/plotting.py` and a dispatch entry in `bench/cli/plot.py::run`.
+
+Results land at `bench/results/<benchmark>/<run_id>.json` (with a `latest.json` symlink alongside) containing metadata (exo SHA, hostname, platform, ISO timestamps, methodology version, user tags), full cluster snapshot, the resolved + derived params, per-step rows, optional cold-control rows, and any derived summaries (e.g. `t_cum_seconds[]` for context-scaling).
+
 ## Dashboard UI Testing & Screenshots
 
 ### Building and Running the Dashboard

Cargo.lock

@@ -916,11 +916,13 @@ dependencies = [
  "libp2p",
  "log",
  "networking",
+ "pidfile-rs",
  "pin-project",
  "pyo3",
  "pyo3-async-runtimes",
  "pyo3-log",
  "pyo3-stub-gen",
+ "thiserror 2.0.17",
  "tokio",
  "util",
 ]
@@ -964,6 +966,16 @@ version = "0.1.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "3a3076410a55c90011c298b04d0cfa770b00fa04e1e3c97d3f6c9de105a03844"
 
+[[package]]
+name = "flopen"
+version = "0.1.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "fbfb8b5fbd1f27929f216650081a07b6ceb0741f0542c8c43ff7ef8e93a35a5d"
+dependencies = [
+ "libc",
+ "nix 0.31.2",
+]
+
 [[package]]
 name = "fnv"
 version = "1.0.7"
@@ -1789,9 +1801,9 @@ checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe"
 
 [[package]]
 name = "libc"
-version = "0.2.178"
+version = "0.2.186"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "37c93d8daa9d8a012fd8ab92f088405fb202ea0b6ab73ee2482ae66af4f42091"
+checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66"
 
 [[package]]
 name = "libp2p"
@@ -2807,6 +2819,18 @@ dependencies = [
  "libc",
 ]
 
+[[package]]
+name = "nix"
+version = "0.31.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5d6d0705320c1e6ba1d912b5e37cf18071b6c2e9b7fa8215a1e8a7651966f5d3"
+dependencies = [
+ "bitflags 2.10.0",
+ "cfg-if",
+ "cfg_aliases",
+ "libc",
+]
+
 [[package]]
 name = "nohash-hasher"
 version = "0.2.0"
@@ -3060,6 +3084,18 @@ dependencies = [
  "siphasher",
 ]
 
+[[package]]
+name = "pidfile-rs"
+version = "0.3.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d1a8aa9a30b1b65ef48b333931b80f2324a14e00208eb2b8f5788f1180791bcc"
+dependencies = [
+ "flopen",
+ "libc",
+ "log",
+ "thiserror 1.0.69",
+]
+
 [[package]]
 name = "pin-project"
 version = "1.1.10"
@@ -3668,7 +3704,7 @@ dependencies = [
  "netlink-packet-utils",
  "netlink-proto",
  "netlink-sys",
- "nix",
+ "nix 0.26.4",
  "thiserror 1.0.69",
  "tokio",
 ]

README.md

@@ -550,6 +550,120 @@ uv run bench/exo_bench.py \
 
 The tool outputs performance metrics including prompt tokens per second (prompt_tps), generation tokens per second (generation_tps), and peak memory usage for each configuration.
 
+### Composable benchmarks (CLI)
+
+For benchmarks that need an `eco`-managed cluster and a stable JSON result format, exo ships a CLI under `bench/cli/`. The CLI handles cluster lifecycle, instance placement, model-metadata resolution (HuggingFace), and result capture; benchmark logic lives in `bench/lib/` so each new benchmark is a small library module + a CLI subcommand.
+
+**Run the prompt-TPS / decode-TPS vs context-size sweep:**
+
+The defaults assume a multi-node, Thunderbolt-connected cluster with tensor parallelism + JACCL — the typical exo benchmarking setup:
+
+```bash
+# Defaults: --sharding Tensor --comm MlxJaccl --thunderbolt a2a, with
+# memory/disk minimums auto-derived from the HF model size. eco picks
+# `--nodes` hosts from its inventory that form a TB clique and satisfy
+# those constraints.
+uv run python -m bench.cli context-scaling \
+  --model mlx-community/Qwen3-30B-A3B-4bit --nodes 2 --num-steps 32
+
+# Pin to specific hosts (defaults still apply for sharding/comm/topology)
+uv run python -m bench.cli context-scaling --hosts s4,s9 \
+  --model mlx-community/Qwen3-30B-A3B-4bit --num-steps 16
+
+# Single-node smoke test: explicit single-node placement overrides
+uv run python -m bench.cli context-scaling --hosts s4 \
+  --model mlx-community/Llama-3.2-1B-Instruct-4bit --num-steps 4 \
+  --sharding Pipeline --comm MlxRing --thunderbolt none
+
+# Override the auto-derived ramp / cold controls
+uv run python -m bench.cli context-scaling --hosts s4,s9 --model X \
+  --pp-step 4096 --num-steps 32 --cold-controls 8192,32768,65536,131072
+
+# Custom output dir + tags
+uv run python -m bench.cli context-scaling --hosts s4,s9 --model X \
+  --output-dir bench/results/2026-05-10/ --tag operator=$USER --tag run=full
+
+# Run from a TOML config (CLI flags override values from the file)
+uv run python -m bench.cli context-scaling \
+  --config bench/configs/context_scaling.example.toml
+```
+
+**Shared flags (every benchmark subcommand has these):**
+
+- `--model` — HuggingFace model id (required)
+- `--config <path>.toml` — load run parameters from a TOML file
+- `--sharding {Pipeline,Tensor}` (default **Tensor**) — sharding mode
+- `--comm {MlxRing,MlxJaccl}` (default **MlxJaccl**) — inter-node comm mode
+- `--min-nodes N` (default 1) — minimum nodes for the placement
+- `--hosts s4,s9` — pin to specific hosts; bypasses constraint search
+- `--nodes N` (default 1) — number of cluster hosts to reserve when `--hosts` is unset (distinct from `--min-nodes`, which controls the model's instance placement)
+- `--thunderbolt {a2a,ring,none}` (default **a2a**) — required Thunderbolt topology
+- `--chip "M3 Ultra"` — required chip (substring match; comment to allow any)
+- `--min-memory-gb`, `--max-memory-gb`, `--min-disk-gb`, `--max-disk-gb` — host RAM / disk constraints. The minimums are auto-derived from the HF model size (×1.30 + 1 GiB for memory, ×1.10 + 1 GiB for disk) when not supplied; explicit values always win.
+- `--evict-downloads` (default **on**) — auto-evict existing models smallest-first on disk-full; pass `--no-evict-downloads` to keep
+- `--cleanup-instance` (default **on**) — delete the placed instance after exit; pass `--no-cleanup-instance` to leave it running for debugging
+- `--output-dir bench/results` — base directory for JSON results (subcommands add their own subfolder)
+- `--tag key=value` — append to `metadata.tags` (repeatable)
+
+**Context-scaling-specific flags:**
+
+- `--num-steps N` — number of equally-spaced ramp points (default 32)
+- `--pp-step Δ` — explicit Δ in tokens (overrides auto-derivation from `max_position_embeddings`)
+- `--fraction-of-max F` — when Δ is auto-derived, use `F × max_context` as the upper bound
+- `--tg` — tokens generated per step (default 64)
+- `--warmup` — warmup requests at `pp=Δ` (default 1)
+- `--cold-controls auto` (4 evenly-spaced points across the ramp) **or** `--cold-controls 8192,32768,…` (explicit pp values). Default: no cold controls.
+
+**Output:** each run writes `bench/results/<benchmark>/<run_id>.json` plus a `latest.json` symlink. The JSON contains metadata (exo SHA, hostname, platform, user tags), the full cluster snapshot at run start, the resolved + derived params, per-step rows, cold-control rows, and derived summaries (`t_cum_seconds`, `control_gaps`).
+
+**Multi-run campaigns** — `bench campaign` runs a list of bench invocations from a single TOML file. Each `[[runs]]` entry is its own cluster deploy + bench + teardown, with a shared `[defaults]` table for DRY config:
+
+```toml
+# bench/configs/llama-family-smoke.toml
+[defaults]
+nodes = 4
+num_steps = 8
+fraction_of_max = 0.5
+
+[[runs]]
+subcommand = "context-scaling"
+model = "mlx-community/Llama-3.2-3B-Instruct-4bit"
+[runs.tags]
+model_short = "llama-3.2-3b-4bit"
+
+[[runs]]
+subcommand = "context-scaling"
+model = "mlx-community/Meta-Llama-3.1-8B-Instruct-4bit"
+[runs.tags]
+model_short = "llama-3.1-8b-4bit"
+
+[plot]
+label_tag = "model_short"
+```
+
+```bash
+uv run python -m bench.cli campaign bench/configs/llama-family-smoke.toml
+```
+
+After all runs finish, an optional `[plot]` table triggers a comparison plot per benchmark group (one PNG per benchmark type with ≥2 runs).
+
+**Plotting** — `bench plot` renders any results JSON to a 2-panel PNG (prompt_tps + generation_tps vs pp_tokens, cold controls overlaid as 'x' markers):
+
+```bash
+# Plot the most recent run next to its JSON
+uv run python -m bench.cli plot bench/results/context_scaling/latest.json
+
+# Compare multiple runs (one line per run; legend label = the chosen tag)
+uv run python -m bench.cli plot run_a.json run_b.json --label-tag operator
+
+# Custom output path + title
+uv run python -m bench.cli plot run.json --output /tmp/scaling.png --title "30B 4-node sweep"
+```
+
+The benchmark type is detected from each JSON's `metadata.benchmark`, so the same `plot` command will work for future benchmarks once their renderer is registered in `bench/lib/plotting.py`.
+
+Methodology for the context-scaling benchmark is documented in detail in `bench/lib/context_scaling.py`'s module docstring and in `bench/METHODOLOGY.md`.
+
 ---
 
 ## Hardware Accelerator Support

app/EXO/EXO/ContentView.swift

@@ -16,22 +16,13 @@ struct ContentView: View {
     @EnvironmentObject private var updater: SparkleUpdater
     @EnvironmentObject private var thunderboltBridgeService: ThunderboltBridgeService
     @EnvironmentObject private var settingsWindowController: SettingsWindowController
+    @EnvironmentObject private var bugReportWindowController: BugReportWindowController
     @State private var focusedNode: NodeViewModel?
     @State private var deletingInstanceIDs: Set<String> = []
     @State private var showAllNodes = false
     @State private var showAllInstances = false
     @State private var baseURLCopied = false
     @State private var showAdvanced = false
-    @State private var showDebugInfo = false
-    private enum BugReportPhase: Equatable {
-        case idle
-        case prompting
-        case sending(String)
-        case success(String)
-        case failure(String)
-    }
-    @State private var bugReportPhase: BugReportPhase = .idle
-    @State private var bugReportUserDescription: String = ""
     @State private var uninstallInProgress = false
     @State private var pendingNamespace: String = ""
     @State private var pendingHFToken: String = ""
@@ -294,6 +285,13 @@ struct ContentView: View {
             ) {
                 updater.checkForUpdates()
             }
+            HoverButton(
+                title: "Share Bug Report…",
+                tint: .primary,
+                trailingSystemImage: "ladybug"
+            ) {
+                bugReportWindowController.open()
+            }
             .padding(.bottom, 8)
             HoverButton(title: "Quit", tint: .secondary) {
                 controller.stop()
@@ -477,40 +475,6 @@ struct ContentView: View {
         }
     }
 
-    private var debugSection: some View {
-        VStack(alignment: .leading, spacing: 4) {
-            HoverButton(
-                title: "Debug Info",
-                tint: .primary,
-                trailingSystemImage: showDebugInfo ? "chevron.up" : "chevron.down",
-                small: true
-            ) {
-                showDebugInfo.toggle()
-            }
-            if showDebugInfo {
-                VStack(alignment: .leading, spacing: 4) {
-                    Text("Version: \(buildTag)")
-                        .font(.caption2)
-                        .foregroundColor(.secondary)
-                    Text("Commit: \(buildCommit)")
-                        .font(.caption2)
-                        .foregroundColor(.secondary)
-                    Text(thunderboltStatusText)
-                        .font(.caption2)
-                        .foregroundColor(thunderboltStatusColor)
-                    clusterThunderboltBridgeView
-                    interfaceIpList
-                    rdmaStatusView
-                    sendBugReportButton
-                        .padding(.top, 6)
-                }
-                .padding(.leading, 8)
-                .transition(.opacity)
-            }
-        }
-        .animation(.easeInOut(duration: 0.25), value: showDebugInfo)
-    }
-
     private var rdmaStatusView: some View {
         let rdmaStatuses = stateService.latestSnapshot?.nodeRdmaCtl ?? [:]
         let localNodeId = stateService.localNodeId
@@ -559,127 +523,6 @@ struct ContentView: View {
         }
     }
 
-    private var sendBugReportButton: some View {
-        VStack(alignment: .leading, spacing: 6) {
-            switch bugReportPhase {
-            case .idle:
-                Button {
-                    bugReportPhase = .prompting
-                    bugReportUserDescription = ""
-                } label: {
-                    HStack {
-                        Text("Send Bug Report")
-                            .font(.caption)
-                            .fontWeight(.semibold)
-                        Spacer()
-                    }
-                    .padding(.vertical, 6)
-                    .padding(.horizontal, 8)
-                    .background(
-                        RoundedRectangle(cornerRadius: 6)
-                            .fill(Color.accentColor.opacity(0.12))
-                    )
-                }
-                .buttonStyle(.plain)
-
-            case .prompting:
-                VStack(alignment: .leading, spacing: 6) {
-                    VStack(alignment: .leading, spacing: 2) {
-                        Text("Tell us what went wrong (optional)")
-                            .font(.caption2)
-                            .foregroundColor(.secondary)
-                        Text(
-                            "A quick description of what you were doing and what happened helps us track down the bug for you."
-                        )
-                        .font(.caption2)
-                        .foregroundColor(.secondary)
-                        .opacity(0.8)
-                        .fixedSize(horizontal: false, vertical: true)
-                    }
-                    TextEditor(text: $bugReportUserDescription)
-                        .font(.caption2)
-                        .frame(height: 60)
-                        .overlay(
-                            RoundedRectangle(cornerRadius: 4)
-                                .stroke(Color.secondary.opacity(0.3), lineWidth: 1)
-                        )
-                    HStack(spacing: 8) {
-                        Button("Send") {
-                            Task {
-                                await sendBugReport()
-                            }
-                        }
-                        .font(.caption2)
-                        .buttonStyle(.borderedProminent)
-                        .controlSize(.small)
-                        Button("Cancel") {
-                            bugReportPhase = .idle
-                        }
-                        .font(.caption2)
-                        .buttonStyle(.bordered)
-                        .controlSize(.small)
-                    }
-                }
-                .padding(8)
-                .background(
-                    RoundedRectangle(cornerRadius: 6)
-                        .fill(Color.accentColor.opacity(0.06))
-                )
-
-            case .sending(let message):
-                HStack(spacing: 6) {
-                    ProgressView()
-                        .scaleEffect(0.6)
-                    Text(message)
-                        .font(.caption2)
-                        .foregroundColor(.secondary)
-                }
-
-            case .success(let message):
-                VStack(alignment: .leading, spacing: 6) {
-                    Text(message)
-                        .font(.caption2)
-                        .foregroundColor(.secondary)
-                        .fixedSize(horizontal: false, vertical: true)
-                    Button {
-                        openGitHubIssue()
-                    } label: {
-                        HStack(spacing: 4) {
-                            Image(systemName: "arrow.up.right.square")
-                                .imageScale(.small)
-                            Text("Create GitHub Issue")
-                                .font(.caption2)
-                        }
-                    }
-                    .buttonStyle(.bordered)
-                    .controlSize(.small)
-                    Button("Done") {
-                        bugReportPhase = .idle
-                        bugReportUserDescription = ""
-                    }
-                    .font(.caption2)
-                    .buttonStyle(.plain)
-                    .foregroundColor(.secondary)
-                }
-
-            case .failure(let message):
-                VStack(alignment: .leading, spacing: 4) {
-                    Text(message)
-                        .font(.caption2)
-                        .foregroundColor(.red)
-                        .fixedSize(horizontal: false, vertical: true)
-                    Button("Dismiss") {
-                        bugReportPhase = .idle
-                    }
-                    .font(.caption2)
-                    .buttonStyle(.plain)
-                    .foregroundColor(.secondary)
-                }
-            }
-        }
-        .animation(.easeInOut(duration: 0.2), value: bugReportPhase)
-    }
-
     private var processToggleBinding: Binding<Bool> {
         Binding(
             get: {
@@ -720,61 +563,6 @@ struct ContentView: View {
         )
     }
 
-    private func sendBugReport() async {
-        bugReportPhase = .sending("Collecting logs...")
-        let service = BugReportService()
-        let description = bugReportUserDescription.trimmingCharacters(in: .whitespacesAndNewlines)
-        do {
-            let outcome = try await service.sendReport(
-                isManual: true,
-                userDescription: description.isEmpty ? nil : description
-            )
-            if outcome.success {
-                bugReportPhase = .success(outcome.message)
-            } else {
-                bugReportPhase = .failure(outcome.message)
-            }
-        } catch {
-            bugReportPhase = .failure(error.localizedDescription)
-        }
-    }
-
-    private func openGitHubIssue() {
-        let description = bugReportUserDescription.trimmingCharacters(in: .whitespacesAndNewlines)
-
-        var bodyParts: [String] = []
-        bodyParts.append("## Describe the bug")
-        bodyParts.append("")
-        if !description.isEmpty {
-            bodyParts.append(description)
-        } else {
-            bodyParts.append("A clear and concise description of what the bug is.")
-        }
-        bodyParts.append("")
-        bodyParts.append("## Environment")
-        bodyParts.append("")
-        bodyParts.append("- macOS Version: \(ProcessInfo.processInfo.operatingSystemVersionString)")
-        bodyParts.append("- EXO Version: \(buildTag) (\(buildCommit))")
-        bodyParts.append("")
-        bodyParts.append("## Additional context")
-        bodyParts.append("")
-        bodyParts.append("A bug report with diagnostic logs was submitted via the app.")
-
-        let body = bodyParts.joined(separator: "\n")
-
-        var components = URLComponents(string: "https://github.com/exo-explore/exo/issues/new")!
-        components.queryItems = [
-            URLQueryItem(name: "template", value: "bug_report.md"),
-            URLQueryItem(name: "title", value: "[BUG] "),
-            URLQueryItem(name: "body", value: body),
-            URLQueryItem(name: "labels", value: "bug"),
-        ]
-
-        if let url = components.url {
-            NSWorkspace.shared.open(url)
-        }
-    }
-
     private func showUninstallConfirmationAlert() {
         let alert = NSAlert()
         alert.messageText = "Uninstall EXO"
@@ -857,13 +645,6 @@ struct ContentView: View {
         }
     }
 
-    private var buildTag: String {
-        Bundle.main.infoDictionary?["EXOBuildTag"] as? String ?? "unknown"
-    }
-
-    private var buildCommit: String {
-        Bundle.main.infoDictionary?["EXOBuildCommit"] as? String ?? "unknown"
-    }
 }
 
 private struct HoverButton: View {

app/EXO/EXO/EXOApp.swift

@@ -22,6 +22,7 @@ struct EXOApp: App {
     @StateObject private var updater: SparkleUpdater
     @StateObject private var thunderboltBridgeService: ThunderboltBridgeService
     @StateObject private var settingsWindowController: SettingsWindowController
+    @StateObject private var bugReportWindowController: BugReportWindowController
     private let terminationObserver: TerminationObserver
     private let firstLaunchPopout = FirstLaunchPopout()
     private let ciContext = CIContext(options: nil)
@@ -46,6 +47,7 @@ struct EXOApp: App {
         let thunderboltBridge = ThunderboltBridgeService(clusterStateService: service)
         _thunderboltBridgeService = StateObject(wrappedValue: thunderboltBridge)
         _settingsWindowController = StateObject(wrappedValue: SettingsWindowController())
+        _bugReportWindowController = StateObject(wrappedValue: BugReportWindowController())
         enableLaunchAtLoginIfNeeded()
         // Install LaunchDaemon to disable Thunderbolt Bridge on startup (prevents network loops)
         NetworkSetupHelper.promptAndInstallIfNeeded()
@@ -66,6 +68,7 @@ struct EXOApp: App {
                 .environmentObject(updater)
                 .environmentObject(thunderboltBridgeService)
                 .environmentObject(settingsWindowController)
+                .environmentObject(bugReportWindowController)
         } label: {
             menuBarIcon
                 .onReceive(controller.$isFirstLaunchReady) { ready in

app/EXO/EXO/Services/ClusterStateService.swift

@@ -17,7 +17,7 @@ final class ClusterStateService: ObservableObject {
 
     init(
         baseURL: URL = URL(string: "http://127.0.0.1:52415")!,
-        session: URLSession = .shared
+        session: URLSession = ClusterStateService.makeNonCachingSession()
     ) {
         self.baseURL = baseURL
         self.endpoint = baseURL.appendingPathComponent("state")
@@ -27,6 +27,23 @@ final class ClusterStateService: ObservableObject {
         self.decoder = decoder
     }
 
+    /// `URLSession.shared` carries an on-disk `URLCache` that persists every
+    /// response body under `~/Library/Caches/exolabs.EXO/`. We poll `/state`
+    /// at 2 Hz from `startPolling`, so leaving the shared cache attached
+    /// dirties ~500–620 KB/sec of file-backed memory and trips macOS's
+    /// per-process `disk writes` resource limit (microstackshot reports
+    /// observed on M3 Ultra producing GBs of cached responses per hour).
+    /// Cluster-state polling responses are time-sensitive and small; they
+    /// gain nothing from being cached on disk. Use an ephemeral session
+    /// with `urlCache = nil` so neither response bodies nor metadata
+    /// touch disk.
+    private static func makeNonCachingSession() -> URLSession {
+        let config = URLSessionConfiguration.ephemeral
+        config.urlCache = nil
+        config.requestCachePolicy = .reloadIgnoringLocalCacheData
+        return URLSession(configuration: config)
+    }
+
     func startPolling(interval: TimeInterval = 0.5) {
         stopPolling()
         Task {

app/EXO/EXO/Views/BugReportWindowController.swift

@@ -0,0 +1,242 @@
+import AppKit
+import SwiftUI
+
+/// Manages a standalone window for the bug-report flow.
+/// Ensures only one instance exists and brings it to front on repeated opens.
+@MainActor
+final class BugReportWindowController: ObservableObject {
+    private var window: NSWindow?
+
+    func open() {
+        if let existing = window, existing.isVisible {
+            existing.makeKeyAndOrderFront(nil)
+            NSApp.activate()
+            return
+        }
+
+        let view = BugReportView(onDismiss: { [weak self] in
+            self?.window?.close()
+        })
+
+        let hostingController = NSHostingController(rootView: view)
+        hostingController.sizingOptions = [.preferredContentSize, .minSize]
+
+        let newWindow = NSWindow(contentViewController: hostingController)
+        newWindow.styleMask = [.titled, .closable, .resizable]
+        newWindow.title = "Send a Bug Report"
+        newWindow.center()
+        newWindow.setFrameAutosaveName("ExoBugReportWindow")
+        newWindow.isReleasedWhenClosed = false
+        newWindow.makeKeyAndOrderFront(nil)
+        NSApp.activate()
+
+        window = newWindow
+    }
+}
+
+private struct BugReportView: View {
+    fileprivate enum Phase: Equatable {
+        case prompting
+        case sending(String)
+        case success(String)
+        case failure(String)
+    }
+
+    let onDismiss: () -> Void
+
+    @State private var phase: Phase = .prompting
+    @State private var userDescription: String = ""
+    @FocusState private var descriptionFocused: Bool
+
+    var body: some View {
+        VStack(alignment: .leading, spacing: 12) {
+            switch phase {
+            case .prompting:
+                promptingView
+            case .sending(let message):
+                sendingView(message: message)
+            case .success(let message):
+                successView(message: message)
+            case .failure(let message):
+                failureView(message: message)
+            }
+        }
+        .padding(16)
+        .frame(minWidth: 380)
+        .animation(.easeInOut(duration: 0.2), value: phase)
+        .onAppear { descriptionFocused = true }
+    }
+
+    private var promptingView: some View {
+        VStack(alignment: .leading, spacing: 8) {
+            Text("Description (optional)")
+                .font(.subheadline)
+                .foregroundColor(.secondary)
+            ZStack(alignment: .topLeading) {
+                if userDescription.isEmpty {
+                    Text("What were you doing when it broke?")
+                        .font(.body)
+                        .foregroundColor(Color(nsColor: .placeholderTextColor))
+                        .padding(.horizontal, 10)
+                        .padding(.vertical, 8)
+                        .allowsHitTesting(false)
+                }
+                TextEditor(text: $userDescription)
+                    .font(.body)
+                    .scrollContentBackground(.hidden)
+                    .padding(4)
+                    .frame(height: 72)
+                    .focused($descriptionFocused)
+            }
+            .background(
+                RoundedRectangle(cornerRadius: 6)
+                    .fill(Color(nsColor: .textBackgroundColor))
+            )
+            .overlay(
+                RoundedRectangle(cornerRadius: 6)
+                    .strokeBorder(Color(nsColor: .separatorColor), lineWidth: 1)
+            )
+
+            Text("Diagnostic logs will be uploaded with your report.")
+                .font(.caption)
+                .foregroundColor(.secondary)
+
+            HStack {
+                Spacer()
+                Button("Cancel") { onDismiss() }
+                    .keyboardShortcut(.cancelAction)
+                Button("Send") {
+                    Task { await send() }
+                }
+                .keyboardShortcut(.defaultAction)
+            }
+            .padding(.top, 4)
+        }
+    }
+
+    private func sendingView(message: String) -> some View {
+        VStack(alignment: .leading, spacing: 12) {
+            HStack(spacing: 10) {
+                ProgressView().controlSize(.small)
+                Text(message)
+                    .foregroundColor(.secondary)
+            }
+            HStack {
+                Spacer()
+                Button("Cancel") { onDismiss() }
+                    .keyboardShortcut(.cancelAction)
+                    .disabled(true)
+                Button("Send") {}
+                    .disabled(true)
+            }
+        }
+    }
+
+    private func successView(message: String) -> some View {
+        VStack(alignment: .leading, spacing: 12) {
+            HStack(alignment: .top, spacing: 10) {
+                Image(systemName: "checkmark.circle.fill")
+                    .foregroundColor(.green)
+                    .font(.title2)
+                Text(message)
+                    .fixedSize(horizontal: false, vertical: true)
+            }
+            HStack {
+                Button {
+                    openGitHubIssue()
+                } label: {
+                    HStack(spacing: 4) {
+                        Image(systemName: "arrow.up.right.square")
+                        Text("Open GitHub Issue")
+                    }
+                }
+                Spacer()
+                Button("Done") { onDismiss() }
+                    .keyboardShortcut(.defaultAction)
+            }
+        }
+    }
+
+    private func failureView(message: String) -> some View {
+        VStack(alignment: .leading, spacing: 12) {
+            HStack(alignment: .top, spacing: 10) {
+                Image(systemName: "exclamationmark.triangle.fill")
+                    .foregroundColor(.orange)
+                    .font(.title2)
+                Text(message)
+                    .fixedSize(horizontal: false, vertical: true)
+            }
+            HStack {
+                Spacer()
+                Button("Try Again") {
+                    phase = .prompting
+                }
+                Button("Close") { onDismiss() }
+                    .keyboardShortcut(.defaultAction)
+            }
+        }
+    }
+
+    private func send() async {
+        phase = .sending("Collecting logs and uploading…")
+        let service = BugReportService()
+        let description = userDescription.trimmingCharacters(in: .whitespacesAndNewlines)
+        do {
+            let outcome = try await service.sendReport(
+                isManual: true,
+                userDescription: description.isEmpty ? nil : description
+            )
+            if outcome.success {
+                phase = .success(outcome.message)
+            } else {
+                phase = .failure(outcome.message)
+            }
+        } catch {
+            phase = .failure(error.localizedDescription)
+        }
+    }
+
+    private func openGitHubIssue() {
+        let description = userDescription.trimmingCharacters(in: .whitespacesAndNewlines)
+
+        var bodyParts: [String] = []
+        bodyParts.append("## Describe the bug")
+        bodyParts.append("")
+        if !description.isEmpty {
+            bodyParts.append(description)
+        } else {
+            bodyParts.append("A clear and concise description of what the bug is.")
+        }
+        bodyParts.append("")
+        bodyParts.append("## Environment")
+        bodyParts.append("")
+        bodyParts.append("- macOS Version: \(ProcessInfo.processInfo.operatingSystemVersionString)")
+        bodyParts.append("- EXO Version: \(buildTag) (\(buildCommit))")
+        bodyParts.append("")
+        bodyParts.append("## Additional context")
+        bodyParts.append("")
+        bodyParts.append("A bug report with diagnostic logs was submitted via the app.")
+
+        let body = bodyParts.joined(separator: "\n")
+
+        var components = URLComponents(string: "https://github.com/exo-explore/exo/issues/new")!
+        components.queryItems = [
+            URLQueryItem(name: "template", value: "bug_report.md"),
+            URLQueryItem(name: "title", value: "[BUG] "),
+            URLQueryItem(name: "body", value: body),
+            URLQueryItem(name: "labels", value: "bug"),
+        ]
+
+        if let url = components.url {
+            NSWorkspace.shared.open(url)
+        }
+    }
+
+    private var buildTag: String {
+        Bundle.main.infoDictionary?["EXOBuildTag"] as? String ?? "unknown"
+    }
+
+    private var buildCommit: String {
+        Bundle.main.infoDictionary?["EXOBuildCommit"] as? String ?? "unknown"
+    }
+}

app/EXO/EXO/Views/SettingsView.swift

@@ -21,8 +21,6 @@ struct SettingsView: View {
     @State private var pendingReadOnlyModelsDirs: String = ""
     @State private var pendingCustomEnvironmentVariables: [CustomEnvironmentVariable] = []
     @State private var needsRestart = false
-    @State private var bugReportInFlight = false
-    @State private var bugReportMessage: String?
     @State private var uninstallInProgress = false
 
     var body: some View {
@@ -202,8 +200,6 @@ struct SettingsView: View {
                 VStack(alignment: .leading, spacing: 2) {
                     rdmaStatusView
                 }
-
-                sendBugReportButton
             }
 
             Section("Danger Zone") {
@@ -504,50 +500,8 @@ struct SettingsView: View {
         }
     }
 
-    private var sendBugReportButton: some View {
-        VStack(alignment: .leading, spacing: 4) {
-            Button {
-                Task {
-                    await sendBugReport()
-                }
-            } label: {
-                HStack {
-                    if bugReportInFlight {
-                        ProgressView()
-                            .scaleEffect(0.6)
-                    }
-                    Text("Send Bug Report")
-                        .font(.caption)
-                        .fontWeight(.semibold)
-                    Spacer()
-                }
-            }
-            .disabled(bugReportInFlight)
-
-            if let message = bugReportMessage {
-                Text(message)
-                    .font(.caption2)
-                    .foregroundColor(.secondary)
-                    .fixedSize(horizontal: false, vertical: true)
-            }
-        }
-    }
-
     // MARK: - Actions
 
-    private func sendBugReport() async {
-        bugReportInFlight = true
-        bugReportMessage = "Collecting logs..."
-        let service = BugReportService()
-        do {
-            let outcome = try await service.sendReport(isManual: true)
-            bugReportMessage = outcome.message
-        } catch {
-            bugReportMessage = error.localizedDescription
-        }
-        bugReportInFlight = false
-    }
-
     private func showUninstallConfirmationAlert() {
         let alert = NSAlert()
         alert.messageText = "Uninstall EXO"

bench/cli/__init__.py

@@ -0,0 +1,14 @@
+"""CLI front-ends for bench library benchmarks.
+
+Each benchmark is a sub-package / module with two pieces:
+
+  - a ``run(...)`` callable in ``bench.lib.<name>`` that does the actual
+    measurement (no argparse, no eco, no I/O)
+  - an ``add_subparser(subparsers)`` helper here that wires CLI args to a
+    handler invoking the lib
+
+The main entry point dispatches to the requested subcommand:
+
+    uv run python -m bench.cli context-scaling --hosts s4 \\
+        --model mlx-community/Qwen3-30B-A3B-4bit
+"""

bench/cli/__main__.py

@@ -0,0 +1,56 @@
+"""``python -m bench.cli`` — dispatcher for benchmark subcommands.
+
+To add a new benchmark:
+
+  1. Implement the methodology in ``bench.lib.<name>`` exposing a typed
+     ``run(session, params, bundle)`` callable (no argparse, no eco I/O).
+  2. Implement a ``bench.cli.<name>`` module with an ``add_subparser`` and
+     a ``run(args) -> Path`` handler.
+  3. Add an ``import + add_subparser(subparsers)`` line below.
+"""
+
+from __future__ import annotations
+
+import argparse
+import sys
+from collections.abc import Callable
+from pathlib import Path
+from typing import cast
+
+from bench.cli import campaign, context_scaling, plot
+from bench.cli._common import expand_config_in_argv
+
+
+def _build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        prog="python -m bench.cli",
+        description=(
+            "Composable, eco-managed benchmarks for exo. "
+            "Pick a subcommand and pass its model / cluster options."
+        ),
+    )
+    subparsers = parser.add_subparsers(
+        dest="subcommand",
+        required=True,
+        metavar="SUBCOMMAND",
+    )
+    context_scaling.add_subparser(subparsers)
+    plot.add_subparser(subparsers)
+    campaign.add_subparser(subparsers)
+    return parser
+
+
+def main(argv: list[str] | None = None) -> int:
+    raw_argv = list(argv if argv is not None else sys.argv[1:])
+    expanded = expand_config_in_argv(raw_argv)
+    args = _build_parser().parse_args(expanded)
+    handler = getattr(args, "handler", None)
+    if not callable(handler):
+        subcommand = getattr(args, "subcommand", "<unknown>")
+        raise SystemExit(f"subcommand {subcommand!r} did not register a handler")
+    cast("Callable[[argparse.Namespace], Path]", handler)(args)
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

bench/cli/_common.py

@@ -0,0 +1,345 @@
+"""Shared CLI argument parsing for the bench command-line interface.
+
+Every benchmark subcommand inherits the same model / cluster / output
+arguments via :func:`add_shared_args` and consumes them through
+:class:`SharedOptions`. argparse's ``Namespace.<attr>`` is fundamentally
+typed ``Any``; the :func:`get_arg` / :func:`get_arg_optional` helpers are
+the single boundary where we coerce to typed values.
+
+A ``--config <path>.toml`` flag lets the caller capture a run definition
+in a TOML file. :func:`expand_config_in_argv` rewrites argv in place,
+substituting the config's keys as CLI flags placed *before* any explicit
+user args so that explicit CLI flags always win.
+"""
+
+from __future__ import annotations
+
+import argparse
+import tomllib
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, TypeVar
+
+from exo_tools.cluster import Chip, Thunderbolt
+from exo_tools.harness import Comm, Sharding
+
+_T = TypeVar("_T")
+
+
+def get_arg(args: argparse.Namespace, name: str, type_: type[_T]) -> _T:
+    """Return ``args.<name>``, asserting it's an instance of ``type_``.
+
+    For ``int`` and ``float`` we additionally accept inputs that ``int(.)`` /
+    ``float(.)`` would parse, since argparse's ``type=int`` already coerces
+    cleanly on input but post-`set_defaults` callers may pass raw values.
+    """
+    raw: Any = getattr(args, name)  # type: ignore[reportAny]
+    if isinstance(raw, type_):
+        return raw
+    if type_ is int and isinstance(raw, (int, str)):
+        return int(raw)  # type: ignore[return-value]
+    if type_ is float and isinstance(raw, (int, float, str)):
+        return float(raw)  # type: ignore[return-value]
+    raise TypeError(
+        f"argparse field {name!r} expected {type_.__name__}, got {type(raw).__name__}"  # type: ignore[reportUnknownArgumentType]
+    )
+
+
+def get_arg_optional(args: argparse.Namespace, name: str, type_: type[_T]) -> _T | None:
+    """Like :func:`get_arg` but allows the field to be missing or None."""
+    raw = getattr(args, name, None)
+    if raw is None:
+        return None
+    if isinstance(raw, type_):
+        return raw
+    if type_ is int and isinstance(raw, (int, str)):
+        return int(raw)  # type: ignore[return-value]
+    if type_ is float and isinstance(raw, (int, float, str)):
+        return float(raw)  # type: ignore[return-value]
+    raise TypeError(
+        f"argparse field {name!r} expected {type_.__name__} or None, "
+        f"got {type(raw).__name__}"  # type: ignore[reportUnknownArgumentType]
+    )
+
+
+@dataclass(frozen=True)
+class SharedOptions:
+    """Parsed shared CLI options for any benchmark."""
+
+    model: str
+    hosts: tuple[str, ...]
+    nodes: int
+    thunderbolt: Thunderbolt | None
+    chip: Chip | None
+    min_memory_gb: float | None
+    max_memory_gb: float | None
+    min_disk_gb: float | None
+    max_disk_gb: float | None
+    evict_downloads: bool
+    sharding: Sharding
+    comm: Comm
+    min_nodes: int
+    output_dir: Path
+    tags: dict[str, str]
+    cleanup_instance: bool
+    user_prefix: str
+
+    @classmethod
+    def from_namespace(cls, args: argparse.Namespace) -> SharedOptions:
+        hosts_raw = get_arg_optional(args, "hosts", str)
+        thunderbolt_raw = get_arg_optional(args, "thunderbolt", str)
+        chip_raw = get_arg_optional(args, "chip", str)
+        tag_list_raw: object = getattr(args, "tag", None) or []
+        if isinstance(tag_list_raw, list):
+            tag_list: list[str] = [
+                str(t)  # type: ignore[reportUnknownArgumentType]
+                for t in tag_list_raw  # type: ignore[reportUnknownVariableType]
+            ]
+        else:
+            tag_list = []
+        return cls(
+            model=get_arg(args, "model", str),
+            hosts=tuple(_parse_csv(hosts_raw)) if hosts_raw else (),
+            nodes=get_arg(args, "nodes", int),
+            thunderbolt=Thunderbolt(thunderbolt_raw) if thunderbolt_raw else None,
+            chip=Chip(chip_raw) if chip_raw else None,
+            min_memory_gb=get_arg_optional(args, "min_memory_gb", float),
+            max_memory_gb=get_arg_optional(args, "max_memory_gb", float),
+            min_disk_gb=get_arg_optional(args, "min_disk_gb", float),
+            max_disk_gb=get_arg_optional(args, "max_disk_gb", float),
+            evict_downloads=get_arg(args, "evict_downloads", bool),
+            sharding=Sharding(get_arg(args, "sharding", str)),
+            comm=Comm(get_arg(args, "comm", str)),
+            min_nodes=get_arg(args, "min_nodes", int),
+            output_dir=Path(get_arg(args, "output_dir", str)),
+            tags=_parse_tags(tag_list),
+            cleanup_instance=get_arg(args, "cleanup_instance", bool),
+            user_prefix=get_arg(args, "eco_user_prefix", str),
+        )
+
+
+def add_shared_args(parser: argparse.ArgumentParser) -> None:
+    """Register the shared-arg group on ``parser``.
+
+    The bool flags (``--auto-constrain``, ``--evict-downloads``,
+    ``--cleanup-instance``) all default to True and use
+    :class:`argparse.BooleanOptionalAction` so callers opt out via the
+    ``--no-X`` form (or set ``X = false`` in a TOML config).
+    """
+    g_config = parser.add_argument_group("config file")
+    g_config.add_argument(
+        "--config",
+        default=None,
+        help="TOML file with run parameters. CLI flags placed after --config "
+        "override values from the file.",
+    )
+
+    g_model = parser.add_argument_group("model")
+    g_model.add_argument(
+        "--model",
+        required=True,
+        help="HuggingFace model id. To run multiple models in one go, use "
+        "the 'campaign' subcommand with a TOML file listing each as a "
+        "separate [[runs]] entry.",
+    )
+    g_model.add_argument(
+        "--sharding",
+        default=Sharding.TENSOR.value,
+        choices=[s.value for s in Sharding],
+        help="Sharding mode for the placed instance. Default 'Tensor' (splits "
+        "layers within nodes; pairs with --comm MlxJaccl for high throughput "
+        "on TB-connected clusters). Use 'Pipeline' for layer-per-node sharding "
+        "(typical for single-node smoke tests).",
+    )
+    g_model.add_argument(
+        "--comm",
+        default=Comm.JACCL.value,
+        choices=[c.value for c in Comm],
+        help="Inter-node communication mode. Default 'MlxJaccl' (RDMA over "
+        "Thunderbolt; pairs with --sharding Tensor and --thunderbolt a2a). "
+        "Use 'MlxRing' for ring all-reduce over the regular network.",
+    )
+    g_model.add_argument("--min-nodes", type=int, default=1)
+
+    g_cluster = parser.add_argument_group("cluster")
+    g_cluster.add_argument(
+        "--hosts",
+        default=None,
+        help="Comma-separated host list (e.g. s4,s9). Bypasses constraint search.",
+    )
+    g_cluster.add_argument(
+        "--nodes",
+        type=int,
+        default=1,
+        help="Number of cluster nodes (hosts) to deploy on. "
+        "Distinct from --min-nodes which controls the model's instance placement.",
+    )
+    g_cluster.add_argument(
+        "--thunderbolt",
+        default=Thunderbolt.A2A.value,
+        choices=[t.value for t in Thunderbolt],
+        help="Thunderbolt topology required: 'a2a' (clique, default; needed "
+        "for tensor parallelism + JACCL), 'ring' (cycle; for pipeline + JACCL), "
+        "or 'none' (exclude TB-connected hosts; pair with --sharding Pipeline "
+        "--comm MlxRing).",
+    )
+    g_cluster.add_argument(
+        "--chip",
+        default=None,
+        choices=[c.value for c in Chip],
+        help="Chip required (e.g. 'M3 Ultra')",
+    )
+    g_cluster.add_argument(
+        "--min-memory-gb",
+        type=float,
+        default=None,
+        help="Min RAM (GB) on each host. If unset, auto-derived from the HF "
+        "model size (×1.30 + 1 GiB).",
+    )
+    g_cluster.add_argument(
+        "--max-memory-gb",
+        type=float,
+        default=None,
+        help="Max RAM (GB) on each host. Useful to leave bigger machines "
+        "free for other workloads.",
+    )
+    g_cluster.add_argument(
+        "--min-disk-gb",
+        type=float,
+        default=None,
+        help="Min free disk (GB) on each host. If unset, auto-derived from "
+        "the HF model size (×1.10 + 1 GiB).",
+    )
+    g_cluster.add_argument(
+        "--max-disk-gb",
+        type=float,
+        default=None,
+        help="Max disk (GB) on each host.",
+    )
+
+    g_runtime = parser.add_argument_group("runtime")
+    g_runtime.add_argument(
+        "--evict-downloads",
+        action=argparse.BooleanOptionalAction,
+        default=True,
+        help="Auto-evict existing models (smallest first) when disk is short to "
+        "make room for the bench model. Default on; pass --no-evict-downloads "
+        "to keep existing downloads.",
+    )
+    g_runtime.add_argument(
+        "--cleanup-instance",
+        action=argparse.BooleanOptionalAction,
+        default=True,
+        help="Clean up the placed instance after the benchmark exits. "
+        "Default on; pass --no-cleanup-instance to leave it running for debugging.",
+    )
+    g_runtime.add_argument(
+        "--eco-user-prefix",
+        default="bench",
+        help="USER prefix for the eco session (default: 'bench').",
+    )
+
+    g_output = parser.add_argument_group("output")
+    g_output.add_argument(
+        "--output-dir",
+        default="bench/results",
+        help="Base directory for JSON results. Subcommands may add a sub-folder.",
+    )
+    g_output.add_argument(
+        "--tag",
+        action="append",
+        default=[],
+        help="Add a 'key=value' tag to metadata.tags (repeatable).",
+    )
+
+
+# ---------------------------------------------------------------------------
+# TOML config expansion
+# ---------------------------------------------------------------------------
+
+
+def expand_config_in_argv(argv: list[str]) -> list[str]:
+    """If ``--config <path>`` appears in ``argv``, splice the TOML's contents in.
+
+    The TOML file's keys are converted to CLI flags (``foo_bar`` →
+    ``--foo-bar``) and inserted *before* the user's other args, so explicit
+    CLI flags always override the config. The ``--config <path>`` pair
+    itself is removed from argv. The first arg (the subcommand name) is
+    preserved at index 0.
+
+    Special handling:
+      - ``[tags]`` table → repeated ``--tag key=value`` occurrences
+      - lists → joined as a comma-separated value (matches the parser's
+        CSV handling for ``--hosts``)
+      - bool true/false → ``--key`` / ``--no-key`` (assumes the underlying
+        flag uses :class:`argparse.BooleanOptionalAction`)
+    """
+    if "--config" not in argv:
+        return list(argv)
+
+    idx = argv.index("--config")
+    if idx + 1 >= len(argv):
+        raise ValueError("--config requires a path argument")
+    config_path = Path(argv[idx + 1])
+    if not config_path.is_file():
+        raise FileNotFoundError(f"Config file not found: {config_path}")
+
+    with config_path.open("rb") as f:
+        config_data: dict[str, Any] = tomllib.load(f)
+
+    expanded = _config_to_argv(config_data)
+    stripped = list(argv[:idx]) + list(argv[idx + 2 :])
+    if not stripped:
+        return expanded
+    # The subcommand name must come first; insert config-derived args
+    # right after it so that the user's later explicit args override.
+    return [stripped[0]] + expanded + stripped[1:]
+
+
+def _config_to_argv(data: dict[str, Any]) -> list[str]:
+    """Convert a TOML-loaded dict to a list of argv-style CLI flags."""
+    out: list[str] = []
+    for key in data:
+        value: Any = data[key]  # type: ignore[reportAny]
+        if key == "tags" and isinstance(value, dict):
+            for tag_key, tag_value in value.items():  # type: ignore[reportUnknownVariableType]
+                out.extend(["--tag", f"{tag_key}={tag_value}"])
+            continue
+        if value is None:
+            continue
+        flag = "--" + key.replace("_", "-")
+        if isinstance(value, bool):
+            out.append(flag if value else f"--no-{key.replace('_', '-')}")
+        elif isinstance(value, list):
+            joined = ",".join(
+                str(x)  # type: ignore[reportUnknownArgumentType]
+                for x in value  # type: ignore[reportUnknownVariableType]
+            )
+            out.extend([flag, joined])
+        else:
+            out.extend([flag, str(value)])  # type: ignore[reportAny]
+    return out
+
+
+def _parse_csv(raw: str) -> list[str]:
+    return [s.strip() for s in raw.split(",") if s.strip()]
+
+
+def _parse_tags(raw: list[str]) -> dict[str, str]:
+    out: dict[str, str] = {}
+    for entry in raw:
+        if "=" not in entry:
+            raise argparse.ArgumentTypeError(
+                f"--tag must be 'key=value', got {entry!r}"
+            )
+        k, v = entry.split("=", 1)
+        out[k.strip()] = v.strip()
+    return out
+
+
+@dataclass
+class CommandResult:
+    """Return value from a benchmark CLI handler."""
+
+    output_path: Path | None = None
+    extra: dict[str, str] = field(default_factory=dict)

bench/cli/campaign.py

@@ -0,0 +1,220 @@
+"""Run a campaign of bench invocations from a single TOML file.
+
+A campaign config has a ``[defaults]`` table (applied to every run) and a
+list of ``[[runs]]`` entries (each a fully-formed invocation with its own
+``subcommand``). The campaign runner merges defaults with each run's
+overrides, dispatches to the matching subcommand handler, and collects
+the output JSON paths.
+
+Each run gets its own cluster — the deploy / teardown happens per-run.
+After all runs finish, an optional ``[plot]`` table triggers a comparison
+plot per benchmark group.
+
+Schema::
+
+    [defaults]
+    nodes = 4
+    num_steps = 8
+
+    [[runs]]
+    subcommand = "context-scaling"
+    model = "mlx-community/Llama-3.2-3B-Instruct-4bit"
+    [runs.tags]
+    model_short = "llama-3.2-3b"
+
+    [[runs]]
+    subcommand = "context-scaling"
+    model = "mlx-community/Meta-Llama-3.1-8B-Instruct-4bit"
+    [runs.tags]
+    model_short = "llama-3.1-8b"
+
+    [plot]
+    label_tag = "model_short"
+"""
+
+from __future__ import annotations
+
+import argparse
+import tomllib
+from collections.abc import Callable
+from pathlib import Path
+from typing import Any, cast
+
+from loguru import logger
+
+from bench.cli import context_scaling
+from bench.cli._common import (
+    _config_to_argv,  # type: ignore[reportPrivateUsage]
+    get_arg,
+)
+from bench.lib.plotting import PlotInputs, render_context_scaling
+
+# Each subcommand exposes its argparse via add_subparser. The campaign
+# runner builds a one-off parser per run with only the chosen subcommand
+# registered, parses the run-derived argv, and invokes the handler.
+_SUBCOMMAND_PARSERS: dict[
+    str,
+    Callable[
+        [Any], None
+    ],  # subparsers action — argparse private; Any-typed at boundary
+] = {
+    "context-scaling": context_scaling.add_subparser,
+}
+
+
+def add_subparser(
+    subparsers: argparse._SubParsersAction[argparse.ArgumentParser],  # type: ignore[type-arg]
+) -> None:
+    parser = subparsers.add_parser(
+        "campaign",
+        help="Run a list of bench invocations from a single TOML config.",
+        description=__doc__,
+    )
+    parser.add_argument(
+        "config",
+        type=str,
+        help="TOML campaign file (with [defaults] + [[runs]] tables).",
+    )
+    parser.add_argument(
+        "--no-plot",
+        action="store_true",
+        help="Skip the optional comparison plot at the end of the campaign.",
+    )
+    parser.set_defaults(handler=run)
+
+
+# ---------------------------------------------------------------------------
+
+
+def run(args: argparse.Namespace) -> Path | None:
+    config_path = Path(get_arg(args, "config", str))
+    if not config_path.is_file():
+        raise SystemExit(f"campaign: file not found: {config_path}")
+
+    with config_path.open("rb") as f:
+        raw = tomllib.load(f)
+
+    defaults = _table(raw, "defaults")
+    runs_obj = raw.get("runs")
+    if not isinstance(runs_obj, list) or not runs_obj:
+        raise SystemExit(f"campaign: {config_path}: missing or empty [[runs]] list")
+    runs_raw: list[Any] = cast("list[Any]", runs_obj)
+    plot_cfg = _table(raw, "plot")
+
+    n_runs = len(runs_raw)
+    output_paths: dict[str, list[Path]] = {}
+    for i, run_obj in enumerate(runs_raw):  # type: ignore[reportAny]
+        if not isinstance(run_obj, dict):
+            raise SystemExit(
+                f"campaign: run #{i + 1}: expected a TOML table, "
+                f"got {type(run_obj).__name__}"  # type: ignore[reportUnknownArgumentType]
+            )
+        run_cfg = cast("dict[str, Any]", run_obj)
+
+        merged = _merge(defaults, run_cfg)
+        subcommand_obj: Any = merged.pop("subcommand", None)  # type: ignore[reportAny]
+        if not isinstance(subcommand_obj, str):
+            raise SystemExit(
+                f"campaign: run #{i + 1}: 'subcommand' field is required (str)"
+            )
+        if subcommand_obj not in _SUBCOMMAND_PARSERS:
+            raise SystemExit(
+                f"campaign: run #{i + 1}: unknown subcommand "
+                f"{subcommand_obj!r} (have {sorted(_SUBCOMMAND_PARSERS)})"
+            )
+
+        argv_for_run = _config_to_argv(merged)
+        sub_args = _parse_for_subcommand(subcommand_obj, argv_for_run)
+        handler = getattr(sub_args, "handler", None)
+        if not callable(handler):
+            raise SystemExit(f"campaign: subcommand {subcommand_obj!r} has no handler")
+
+        logger.info(
+            f"campaign: starting run {i + 1}/{n_runs} "
+            f"({subcommand_obj}; {len(merged)} flags)"
+        )
+        out = cast("Callable[[argparse.Namespace], Path]", handler)(sub_args)
+        output_paths.setdefault(subcommand_obj, []).append(Path(out))
+        logger.info(f"campaign: finished run {i + 1}/{n_runs} → {out}")
+
+    last_path: Path | None = None
+    for paths in output_paths.values():
+        if paths:
+            last_path = paths[-1]
+
+    if get_arg(args, "no_plot", bool):
+        return last_path
+
+    comparison = _render_comparisons(output_paths, plot_cfg)
+    return comparison or last_path
+
+
+# ---------------------------------------------------------------------------
+
+
+def _table(data: dict[str, Any], key: str) -> dict[str, Any]:
+    """Return ``data[key]`` if it's a table, else an empty dict."""
+    val: Any = data.get(key)
+    return cast("dict[str, Any]", val) if isinstance(val, dict) else {}
+
+
+def _merge(defaults: dict[str, Any], run: dict[str, Any]) -> dict[str, Any]:
+    """Shallow-merge ``defaults`` with ``run``; ``run`` wins on conflict.
+
+    The ``tags`` table is deep-merged (defaults' tags + run's tags) so a
+    campaign-level operator tag and a per-run model_short tag both survive.
+    """
+    merged: dict[str, Any] = {**defaults, **run}
+    default_tags = _table(defaults, "tags")
+    run_tags = _table(run, "tags")
+    if default_tags or run_tags:
+        merged["tags"] = {**default_tags, **run_tags}
+    return merged
+
+
+def _parse_for_subcommand(
+    subcommand: str, argv_for_run: list[str]
+) -> argparse.Namespace:
+    """Build a one-off parser with ``subcommand`` registered + parse argv."""
+    parser = argparse.ArgumentParser(prog=f"bench campaign:{subcommand}")
+    subparsers = parser.add_subparsers(dest="subcommand", required=True)
+    _SUBCOMMAND_PARSERS[subcommand](subparsers)
+    return parser.parse_args([subcommand] + argv_for_run)
+
+
+def _render_comparisons(
+    output_paths: dict[str, list[Path]],
+    plot_cfg: dict[str, Any],
+) -> Path | None:
+    """Render one comparison plot per benchmark group with ≥2 outputs."""
+    label_tag = _str_or_none(plot_cfg.get("label_tag"))
+    title = _str_or_none(plot_cfg.get("title"))
+
+    last: Path | None = None
+    for subcommand, paths in output_paths.items():
+        if len(paths) < 2:
+            continue
+        if subcommand != "context-scaling":
+            logger.warning(
+                f"campaign: no comparison renderer registered for {subcommand!r}; "
+                "skipping comparison plot"
+            )
+            continue
+        out = paths[0].with_name(f"campaign_{subcommand}_compare.png")
+        last = render_context_scaling(
+            PlotInputs(
+                results=paths,
+                output=out,
+                label_tag=label_tag,
+                title=title,
+            )
+        )
+        logger.info(f"campaign: wrote comparison plot {last}")
+    return last
+
+
+def _str_or_none(value: Any) -> str | None:  # type: ignore[reportAny]
+    return value if isinstance(value, str) else None
+
+
+__all__ = ["add_subparser", "run"]

bench/cli/context_scaling.py

@@ -0,0 +1,270 @@
+"""Context-scaling benchmark — CLI subcommand.
+
+Wraps :func:`bench.lib.context_scaling.run` with:
+
+  - HF model-metadata resolution
+  - Auto-derived constraints (memory, disk) and context ramp (Δ, K)
+  - eco cluster + instance lifecycle (managed_cluster + managed_instance)
+  - Cold-control isolation (delete sweep instance before controls)
+  - JSON results + ``latest.json`` symlink under ``<output-dir>/context_scaling/``
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+from pathlib import Path
+
+from exo_tools.cluster import EcoSession
+from loguru import logger
+
+from bench.cli._common import (
+    SharedOptions,
+    add_shared_args,
+    get_arg,
+    get_arg_optional,
+)
+from bench.lib import context_scaling
+from bench.lib.cluster import managed_cluster, managed_instance
+from bench.lib.context_scaling import (
+    ContextScalingParams,
+    make_cold_control_factory,
+)
+from bench.lib.model_meta import (
+    ModelMeta,
+    derive_cold_controls,
+    derive_context_ramp,
+    fetch_model_meta,
+)
+from bench.lib.results import ResultsBundle, RunMetadata, find_repo_root
+
+
+def add_subparser(
+    subparsers: argparse._SubParsersAction[argparse.ArgumentParser],  # type: ignore[type-arg]
+) -> None:
+    parser = subparsers.add_parser(
+        "context-scaling",
+        help="Prompt-TPS / decode-TPS vs context-size sweep",
+        description=__doc__,
+    )
+    add_shared_args(parser)
+    g = parser.add_argument_group("context-scaling")
+    g.add_argument(
+        "--num-steps",
+        type=int,
+        default=32,
+        help="Number of equally-spaced PP points in the ramp (K).",
+    )
+    g.add_argument(
+        "--pp-step",
+        type=int,
+        default=None,
+        help="Δ (token step). If unset, derived from the model's max context.",
+    )
+    g.add_argument(
+        "--fraction-of-max",
+        type=float,
+        default=1.0,
+        help="When Δ is auto-derived, use this fraction of the model's "
+        "max_position_embeddings as the ramp's upper bound (0 < f ≤ 1).",
+    )
+    g.add_argument(
+        "--tg",
+        type=int,
+        default=64,
+        help="Tokens to generate per step (decode duration; constant across ramp).",
+    )
+    g.add_argument(
+        "--warmup",
+        type=int,
+        default=2,
+        help="Warmup requests at pp=Δ before the measured ramp. "
+        "First warmup is cache-disabled (kernel JIT only); subsequent "
+        "warmups are cache-enabled (the second is the one that primes "
+        "the cache entry with a hot-kernel rate). Default 2 is the "
+        "sweet spot: warmup=0 leaves JIT cost in step 0; warmup=1 has "
+        "step 0 as a 'none' hit (still hot-kernel cold prefill, just "
+        "classified differently).",
+    )
+    g.add_argument(
+        "--cold-controls",
+        type=str,
+        default=None,
+        help="Cold-control pp values to take after the cached sweep. Either "
+        "'auto' (4 evenly-spaced points across the ramp) or a comma-separated "
+        "list of explicit pp values (e.g. '8192,32768,65536'). "
+        "Default: no cold controls.",
+    )
+    g.add_argument(
+        "--sleep-between-s",
+        type=float,
+        default=1.0,
+        help="Seconds to sleep between consecutive sweep requests.",
+    )
+    parser.set_defaults(handler=run)
+
+
+# ---------------------------------------------------------------------------
+
+
+def run(args: argparse.Namespace) -> Path:
+    """Execute the context-scaling benchmark per the parsed args.
+
+    Returns the path of the JSON results file.
+    """
+    shared = SharedOptions.from_namespace(args)
+    repo_root = find_repo_root()
+
+    # 1. Fetch HF metadata up-front; everything else can be derived from it.
+    logger.info(f"fetching HuggingFace metadata for {shared.model}")
+    meta = fetch_model_meta(shared.model)
+    logger.info(
+        f"  weights: {meta.total_weight_gb:.1f}GB; "
+        f"max context: {meta.max_position_embeddings} tokens; "
+        f"layers: {meta.num_hidden_layers}"
+    )
+
+    # 2. Derive constraints (user values always win; otherwise fall back to
+    # ModelMeta heuristics for the *minimums*).
+    min_memory_gb = (
+        shared.min_memory_gb
+        if shared.min_memory_gb is not None
+        else meta.memory_constraint_gb
+    )
+    min_disk_gb = (
+        shared.min_disk_gb
+        if shared.min_disk_gb is not None
+        else meta.disk_constraint_gb
+    )
+    logger.info(f"  cluster constraint: min memory {min_memory_gb:.1f}GB")
+    logger.info(f"  cluster constraint: min disk {min_disk_gb:.1f}GB")
+    if shared.max_memory_gb is not None:
+        logger.info(f"  cluster constraint: max memory {shared.max_memory_gb:.1f}GB")
+    if shared.max_disk_gb is not None:
+        logger.info(f"  cluster constraint: max disk {shared.max_disk_gb:.1f}GB")
+
+    explicit_pp_step = get_arg_optional(args, "pp_step", int)
+    num_steps = get_arg(args, "num_steps", int)
+    if explicit_pp_step is not None:
+        pp_step = explicit_pp_step
+    else:
+        pp_step, num_steps = derive_context_ramp(
+            meta,
+            num_steps=num_steps,
+            fraction_of_max=get_arg(args, "fraction_of_max", float),
+        )
+        logger.info(
+            f"  derived ramp: Δ={pp_step} × K={num_steps} "
+            f"= {pp_step * num_steps} tokens (max {meta.max_position_embeddings})"
+        )
+
+    cold_controls = _resolve_cold_controls(
+        args, meta, pp_step=pp_step, num_steps=num_steps
+    )
+    if cold_controls:
+        logger.info(f"  cold controls: {list(cold_controls)}")
+
+    # 3. Spin up cluster + instance + run.
+    eco = EcoSession(user_prefix=shared.user_prefix)
+    output_dir = (shared.output_dir / "context_scaling").resolve()
+    metadata = RunMetadata.new(
+        benchmark="context_scaling",
+        repo_root=repo_root,
+        tags={**shared.tags, "host_pool": ",".join(shared.hosts) or "<auto>"},
+    )
+    bundle = ResultsBundle(metadata=metadata)
+
+    with (
+        managed_cluster(
+            eco,
+            hosts=list(shared.hosts) or None,
+            count=shared.nodes,
+            thunderbolt=shared.thunderbolt,
+            chip=shared.chip,
+            min_memory_gb=min_memory_gb,
+            max_memory_gb=shared.max_memory_gb,
+            min_disk_gb=min_disk_gb,
+            max_disk_gb=shared.max_disk_gb,
+        ) as cluster,
+        managed_instance(
+            cluster,
+            eco,
+            shared.model,
+            sharding=shared.sharding,
+            comm=shared.comm,
+            min_nodes=shared.min_nodes,
+            evict_downloads=shared.evict_downloads,
+            cleanup_on_exit=shared.cleanup_instance,
+        ) as session,
+    ):
+        params = ContextScalingParams(
+            pp_step=pp_step,
+            num_steps=num_steps,
+            tg=get_arg(args, "tg", int),
+            warmup=get_arg(args, "warmup", int),
+            cold_controls=cold_controls,
+            sleep_between_s=get_arg(args, "sleep_between_s", float),
+        )
+        factory = (
+            make_cold_control_factory(
+                session, shared.sharding, shared.comm, shared.min_nodes
+            )
+            if cold_controls
+            else None
+        )
+        context_scaling.run(session, params, bundle, cold_control_factory=factory)
+
+    out_path = bundle.write_json(output_dir)
+    _update_latest_symlink(out_path)
+    logger.info(f"wrote results → {out_path}")
+
+    _validate_partial_hits(bundle)
+    return out_path
+
+
+# ---------------------------------------------------------------------------
+
+
+def _resolve_cold_controls(
+    args: argparse.Namespace,
+    meta: ModelMeta,
+    *,
+    pp_step: int,
+    num_steps: int,
+) -> tuple[int, ...]:
+    raw = get_arg_optional(args, "cold_controls", str)
+    if raw is None or not raw.strip():
+        return ()
+    if raw.strip().lower() == "auto":
+        return derive_cold_controls(meta, pp_step=pp_step, num_steps=num_steps, count=4)
+    return tuple(int(s.strip()) for s in raw.split(",") if s.strip())
+
+
+def _update_latest_symlink(out_path: Path) -> None:
+    """Update ``<dir>/latest.json`` to point at the newly-written file."""
+    link = out_path.parent / "latest.json"
+    try:
+        if link.is_symlink() or link.exists():
+            link.unlink()
+        os.symlink(out_path.name, link)
+    except OSError as e:
+        logger.warning(f"could not update latest.json symlink: {e}")
+
+
+def _validate_partial_hits(bundle: ResultsBundle) -> None:
+    """Hard-fail if the cached sweep didn't see ``partial`` on every step ≥ 1.
+
+    Step 0 is allowed to be ``exact`` (warmup primed the cache at pp=Δ); a
+    later ``exact`` means Δ was effectively absorbed into the cache and the
+    cold-rate measurement is meaningless. ``none`` means the cache was
+    discarded mid-sweep and ``T_cum`` is unreliable.
+    """
+    cached = [r for r in bundle.runs if r.get("phase") == "cached_sweep"]
+    bad = [r for r in cached[1:] if r.get("prefix_cache_hit") != "partial"]
+    if bad:
+        bad_summary = [(r["step_index"], r["prefix_cache_hit"]) for r in bad]
+        raise RuntimeError(
+            f"{len(bad)} cached-sweep step(s) reported "
+            f"prefix_cache_hit != 'partial': {bad_summary!r}; "
+            "T_cum is unreliable."
+        )

bench/cli/plot.py

@@ -0,0 +1,125 @@
+"""Plot benchmark results — CLI subcommand.
+
+    uv run python -m bench.cli plot bench/results/context_scaling/latest.json
+    uv run python -m bench.cli plot run_a.json run_b.json --label-tag operator
+    uv run python -m bench.cli plot latest.json --output /tmp/scaling.png
+
+The benchmark type is detected from each JSON's ``metadata.benchmark`` —
+all input files must share the same benchmark.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+from typing import Any, cast
+
+from loguru import logger
+
+from bench.cli._common import get_arg_optional
+from bench.lib.plotting import PlotInputs, render_context_scaling
+
+
+def add_subparser(
+    subparsers: argparse._SubParsersAction[argparse.ArgumentParser],  # type: ignore[type-arg]
+) -> None:
+    parser = subparsers.add_parser(
+        "plot",
+        help="Render benchmark JSON result(s) as a PNG.",
+        description=__doc__,
+    )
+    parser.add_argument(
+        "paths",
+        nargs="+",
+        type=str,
+        help="One or more bench results JSON files. Multiple files are "
+        "rendered as a comparison plot (one line per file).",
+    )
+    parser.add_argument(
+        "--output",
+        type=str,
+        default=None,
+        help="PNG output path. Default: replace the first JSON's '.json' "
+        "suffix with '.png' (or '.compare.png' when multiple inputs).",
+    )
+    parser.add_argument(
+        "--label-tag",
+        type=str,
+        default=None,
+        help="Use metadata.tags[<KEY>] as the legend label for each run "
+        "(falls back to run_id if unset or missing).",
+    )
+    parser.add_argument(
+        "--title",
+        type=str,
+        default=None,
+        help="Override the auto-generated figure title.",
+    )
+    parser.set_defaults(handler=run)
+
+
+def run(args: argparse.Namespace) -> Path:
+    paths_raw = getattr(args, "paths", None)
+    if not isinstance(paths_raw, list) or not paths_raw:
+        raise SystemExit("plot: at least one JSON path is required")
+    paths = [
+        Path(str(p))  # type: ignore[reportUnknownArgumentType]
+        for p in cast("list[Any]", paths_raw)  # type: ignore[reportAny]
+    ]
+    for p in paths:
+        if not p.is_file():
+            raise SystemExit(f"plot: file not found: {p}")
+
+    benchmarks = {_benchmark_for(p) for p in paths}
+    if len(benchmarks) != 1:
+        raise SystemExit(
+            f"plot: all input JSONs must share the same benchmark, got {benchmarks!r}"
+        )
+    benchmark = next(iter(benchmarks))
+
+    output_arg = get_arg_optional(args, "output", str)
+    output = Path(output_arg) if output_arg is not None else _default_output(paths)
+    inputs = PlotInputs(
+        results=paths,
+        output=output,
+        label_tag=get_arg_optional(args, "label_tag", str),
+        title=get_arg_optional(args, "title", str),
+    )
+
+    if benchmark == "context_scaling":
+        out_path = render_context_scaling(inputs)
+    else:
+        raise SystemExit(f"plot: no renderer registered for benchmark {benchmark!r}")
+
+    logger.info(f"plot: wrote {out_path}")
+    return out_path
+
+
+# ---------------------------------------------------------------------------
+
+
+def _benchmark_for(path: Path) -> str:
+    with path.open() as f:
+        loaded: Any = json.load(f)  # type: ignore[reportAny]
+    if not isinstance(loaded, dict):
+        raise SystemExit(f"plot: {path}: expected top-level JSON object")
+    metadata: Any = loaded.get("metadata", {})  # type: ignore[reportAny]
+    if not isinstance(metadata, dict):
+        raise SystemExit(f"plot: {path}: metadata is not an object")
+    benchmark: Any = metadata.get("benchmark")  # type: ignore[reportAny]
+    if not isinstance(benchmark, str):
+        raise SystemExit(f"plot: {path}: metadata.benchmark missing or not a string")
+    return benchmark
+
+
+def _default_output(paths: list[Path]) -> Path:
+    """Auto-derive a PNG path next to the first JSON.
+
+    Single input → ``<path>.png`` (replaces ``.json``).
+    Multiple inputs → ``<path>.compare.png`` next to the first JSON.
+    """
+    first = paths[0]
+    if len(paths) == 1:
+        return first.with_suffix(".png")
+    return first.with_name(first.stem + ".compare.png")

bench/cli/tests/test_campaign.py

@@ -0,0 +1,109 @@
+"""Unit tests for ``bench.cli.campaign``.
+
+The pure helpers (defaults+run merge, table-lookup, str-or-none) are
+tested here. End-to-end campaign execution requires a real eco cluster
+and is exercised manually via ``bench campaign <toml>``.
+"""
+
+from __future__ import annotations
+
+from bench.cli.campaign import (
+    _merge,  # type: ignore[reportPrivateUsage]
+    _str_or_none,  # type: ignore[reportPrivateUsage]
+    _table,  # type: ignore[reportPrivateUsage]
+)
+
+# ---------------------------------------------------------------------------
+# _table
+# ---------------------------------------------------------------------------
+
+
+class TestTable:
+    def test_present_table(self) -> None:
+        data = {"defaults": {"nodes": 4}}
+        assert _table(data, "defaults") == {"nodes": 4}
+
+    def test_missing_key_returns_empty(self) -> None:
+        assert _table({}, "absent") == {}
+
+    def test_non_table_value_returns_empty(self) -> None:
+        # `nodes = 4` is an int at top level, not a table; treat as empty.
+        assert _table({"nodes": 4}, "nodes") == {}
+
+    def test_list_value_returns_empty(self) -> None:
+        assert _table({"runs": [{"a": 1}]}, "runs") == {}
+
+
+# ---------------------------------------------------------------------------
+# _merge
+# ---------------------------------------------------------------------------
+
+
+class TestMerge:
+    def test_run_wins_on_conflict(self) -> None:
+        defaults = {"nodes": 4, "tg": 64}
+        run = {"nodes": 2}
+        assert _merge(defaults, run) == {"nodes": 2, "tg": 64}
+
+    def test_disjoint_keys(self) -> None:
+        defaults = {"nodes": 4}
+        run = {"model": "test/foo"}
+        assert _merge(defaults, run) == {"nodes": 4, "model": "test/foo"}
+
+    def test_run_only(self) -> None:
+        assert _merge({}, {"a": 1, "b": 2}) == {"a": 1, "b": 2}
+
+    def test_defaults_only(self) -> None:
+        assert _merge({"a": 1}, {}) == {"a": 1}
+
+    def test_tags_deep_merged_defaults_only(self) -> None:
+        defaults = {"tags": {"operator": "ciaranbor"}}
+        run = {"model": "test/foo"}
+        merged = _merge(defaults, run)
+        assert merged["tags"] == {"operator": "ciaranbor"}
+
+    def test_tags_deep_merged_run_only(self) -> None:
+        defaults = {"nodes": 4}
+        run = {"tags": {"model_short": "llama-3b"}}
+        merged = _merge(defaults, run)
+        assert merged["tags"] == {"model_short": "llama-3b"}
+
+    def test_tags_deep_merged_both(self) -> None:
+        defaults = {"tags": {"operator": "ciaranbor", "campaign": "smoke"}}
+        run = {"tags": {"model_short": "llama-3b"}}
+        merged = _merge(defaults, run)
+        assert merged["tags"] == {
+            "operator": "ciaranbor",
+            "campaign": "smoke",
+            "model_short": "llama-3b",
+        }
+
+    def test_run_tags_override_defaults_tags(self) -> None:
+        defaults = {"tags": {"operator": "ciaranbor"}}
+        run = {"tags": {"operator": "alice"}}
+        merged = _merge(defaults, run)
+        assert merged["tags"] == {"operator": "alice"}
+
+    def test_no_tags_table_means_no_tags_key(self) -> None:
+        # When neither side has tags, we don't synthesise an empty dict.
+        merged = _merge({"nodes": 4}, {"model": "test/foo"})
+        assert "tags" not in merged
+
+
+# ---------------------------------------------------------------------------
+# _str_or_none
+# ---------------------------------------------------------------------------
+
+
+class TestStrOrNone:
+    def test_str_passes_through(self) -> None:
+        assert _str_or_none("hello") == "hello"
+
+    def test_none_returns_none(self) -> None:
+        assert _str_or_none(None) is None
+
+    def test_int_returns_none(self) -> None:
+        assert _str_or_none(42) is None
+
+    def test_list_returns_none(self) -> None:
+        assert _str_or_none(["a", "b"]) is None

bench/cli/tests/test_common.py

@@ -0,0 +1,275 @@
+"""Unit tests for the argparse boundary helpers in ``bench.cli._common``."""
+
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+import pytest
+
+from bench.cli._common import (
+    _config_to_argv,  # type: ignore[reportPrivateUsage]
+    _parse_csv,  # type: ignore[reportPrivateUsage]
+    _parse_tags,  # type: ignore[reportPrivateUsage]
+    expand_config_in_argv,
+    get_arg,
+    get_arg_optional,
+)
+
+# ---------------------------------------------------------------------------
+# _parse_csv
+# ---------------------------------------------------------------------------
+
+
+class TestParseCsv:
+    def test_simple_list(self) -> None:
+        assert _parse_csv("a,b,c") == ["a", "b", "c"]
+
+    def test_strips_whitespace(self) -> None:
+        assert _parse_csv("  a , b , c  ") == ["a", "b", "c"]
+
+    def test_skips_empty_entries(self) -> None:
+        assert _parse_csv("a,,b,") == ["a", "b"]
+        assert _parse_csv(",,,") == []
+
+    def test_empty_string_returns_empty(self) -> None:
+        assert _parse_csv("") == []
+
+    def test_single_value(self) -> None:
+        assert _parse_csv("only") == ["only"]
+
+
+# ---------------------------------------------------------------------------
+# _parse_tags
+# ---------------------------------------------------------------------------
+
+
+class TestParseTags:
+    def test_empty_input_returns_empty_dict(self) -> None:
+        assert _parse_tags([]) == {}
+
+    def test_single_tag(self) -> None:
+        assert _parse_tags(["operator=ciaranbor"]) == {"operator": "ciaranbor"}
+
+    def test_multiple_tags(self) -> None:
+        assert _parse_tags(["a=1", "b=2", "c=3"]) == {"a": "1", "b": "2", "c": "3"}
+
+    def test_strips_whitespace_around_key_and_value(self) -> None:
+        assert _parse_tags(["  key  =  value  "]) == {"key": "value"}
+
+    def test_value_can_contain_equals(self) -> None:
+        assert _parse_tags(["url=http://example.com/?a=b"]) == {
+            "url": "http://example.com/?a=b"
+        }
+
+    def test_later_duplicate_key_wins(self) -> None:
+        # Standard dict behaviour; explicit so we notice if it changes.
+        assert _parse_tags(["k=v1", "k=v2"]) == {"k": "v2"}
+
+    def test_missing_equals_raises(self) -> None:
+        with pytest.raises(argparse.ArgumentTypeError, match="key=value"):
+            _ = _parse_tags(["malformed"])
+
+    def test_one_malformed_in_list_raises(self) -> None:
+        with pytest.raises(argparse.ArgumentTypeError):
+            _ = _parse_tags(["good=1", "bad", "alsogood=2"])
+
+
+# ---------------------------------------------------------------------------
+# get_arg / get_arg_optional
+# ---------------------------------------------------------------------------
+
+
+class TestGetArg:
+    def test_str_passes_through(self) -> None:
+        ns = argparse.Namespace(name="hello")
+        assert get_arg(ns, "name", str) == "hello"
+
+    def test_int_passes_through(self) -> None:
+        ns = argparse.Namespace(count=42)
+        assert get_arg(ns, "count", int) == 42
+
+    def test_int_coerces_from_string(self) -> None:
+        ns = argparse.Namespace(count="42")
+        assert get_arg(ns, "count", int) == 42
+
+    def test_float_passes_through(self) -> None:
+        ns = argparse.Namespace(rate=3.14)
+        assert get_arg(ns, "rate", float) == 3.14
+
+    def test_float_coerces_from_int(self) -> None:
+        ns = argparse.Namespace(rate=3)
+        assert get_arg(ns, "rate", float) == 3.0
+
+    def test_float_coerces_from_string(self) -> None:
+        ns = argparse.Namespace(rate="3.14")
+        assert get_arg(ns, "rate", float) == 3.14
+
+    def test_bool_passes_through(self) -> None:
+        ns = argparse.Namespace(flag=True)
+        assert get_arg(ns, "flag", bool) is True
+
+    def test_wrong_type_raises(self) -> None:
+        ns = argparse.Namespace(name=42)
+        with pytest.raises(TypeError, match="expected str"):
+            _ = get_arg(ns, "name", str)
+
+    def test_missing_attribute_raises(self) -> None:
+        ns = argparse.Namespace()
+        with pytest.raises(AttributeError):
+            _ = get_arg(ns, "missing", str)
+
+
+class TestGetArgOptional:
+    def test_missing_returns_none(self) -> None:
+        ns = argparse.Namespace()
+        assert get_arg_optional(ns, "missing", str) is None
+
+    def test_explicit_none_returns_none(self) -> None:
+        ns = argparse.Namespace(value=None)
+        assert get_arg_optional(ns, "value", str) is None
+
+    def test_present_value_returns_typed(self) -> None:
+        ns = argparse.Namespace(value="present")
+        assert get_arg_optional(ns, "value", str) == "present"
+
+    def test_int_coerces_from_string(self) -> None:
+        ns = argparse.Namespace(value="42")
+        assert get_arg_optional(ns, "value", int) == 42
+
+    def test_float_coerces_from_int(self) -> None:
+        ns = argparse.Namespace(value=42)
+        assert get_arg_optional(ns, "value", float) == 42.0
+
+    def test_wrong_type_raises(self) -> None:
+        ns = argparse.Namespace(value=[1, 2, 3])
+        with pytest.raises(TypeError, match="expected str or None"):
+            _ = get_arg_optional(ns, "value", str)
+
+
+# ---------------------------------------------------------------------------
+# _config_to_argv
+# ---------------------------------------------------------------------------
+
+
+class TestConfigToArgv:
+    def test_empty(self) -> None:
+        assert _config_to_argv({}) == []
+
+    def test_string_value(self) -> None:
+        assert _config_to_argv({"model": "mlx/foo"}) == ["--model", "mlx/foo"]
+
+    def test_int_and_float_values(self) -> None:
+        out = _config_to_argv({"num_steps": 32, "fraction_of_max": 0.5})
+        assert out == ["--num-steps", "32", "--fraction-of-max", "0.5"]
+
+    def test_underscore_keys_become_hyphenated_flags(self) -> None:
+        out = _config_to_argv({"min_memory_gb": 21.0})
+        assert out == ["--min-memory-gb", "21.0"]
+
+    def test_bool_true_emits_flag(self) -> None:
+        assert _config_to_argv({"auto_constrain": True}) == ["--auto-constrain"]
+
+    def test_bool_false_emits_no_form(self) -> None:
+        assert _config_to_argv({"auto_constrain": False}) == ["--no-auto-constrain"]
+
+    def test_none_value_skipped(self) -> None:
+        assert _config_to_argv({"chip": None, "model": "foo"}) == [
+            "--model",
+            "foo",
+        ]
+
+    def test_list_joined_as_csv(self) -> None:
+        out = _config_to_argv({"hosts": ["s4", "s9"], "cold_controls": [1024, 2048]})
+        assert out == [
+            "--hosts",
+            "s4,s9",
+            "--cold-controls",
+            "1024,2048",
+        ]
+
+    def test_tags_table_expands_to_repeated_tag_args(self) -> None:
+        out = _config_to_argv({"tags": {"operator": "ciaranbor", "run": "full"}})
+        # Order within a TOML table is preserved by tomllib
+        assert out == [
+            "--tag",
+            "operator=ciaranbor",
+            "--tag",
+            "run=full",
+        ]
+
+
+# ---------------------------------------------------------------------------
+# expand_config_in_argv
+# ---------------------------------------------------------------------------
+
+
+class TestExpandConfigInArgv:
+    def test_no_config_flag_passthrough(self) -> None:
+        argv = ["context-scaling", "--model", "foo"]
+        assert expand_config_in_argv(argv) == argv
+
+    def test_config_at_end(self, tmp_path: Path) -> None:
+        cfg = tmp_path / "run.toml"
+        _ = cfg.write_text('model = "from_config"\nnum_steps = 16\n')
+        argv = ["context-scaling", "--config", str(cfg)]
+        # Config flags are inserted right after the subcommand
+        assert expand_config_in_argv(argv) == [
+            "context-scaling",
+            "--model",
+            "from_config",
+            "--num-steps",
+            "16",
+        ]
+
+    def test_explicit_cli_overrides_config(self, tmp_path: Path) -> None:
+        cfg = tmp_path / "run.toml"
+        _ = cfg.write_text('model = "from_config"\nnum_steps = 16\n')
+        # User overrides --num-steps explicitly. Argparse takes the last
+        # occurrence for non-append actions, so the user's 32 wins.
+        argv = ["context-scaling", "--config", str(cfg), "--num-steps", "32"]
+        out = expand_config_in_argv(argv)
+        assert out == [
+            "context-scaling",
+            "--model",
+            "from_config",
+            "--num-steps",
+            "16",
+            "--num-steps",
+            "32",
+        ]
+
+    def test_missing_path_arg_raises(self) -> None:
+        with pytest.raises(ValueError, match="--config requires a path"):
+            _ = expand_config_in_argv(["context-scaling", "--config"])
+
+    def test_nonexistent_file_raises(self, tmp_path: Path) -> None:
+        with pytest.raises(FileNotFoundError, match="Config file not found"):
+            _ = expand_config_in_argv(
+                ["context-scaling", "--config", str(tmp_path / "missing.toml")]
+            )
+
+    def test_bool_false_in_config(self, tmp_path: Path) -> None:
+        cfg = tmp_path / "run.toml"
+        _ = cfg.write_text("auto_constrain = false\n")
+        argv = ["context-scaling", "--config", str(cfg)]
+        assert expand_config_in_argv(argv) == [
+            "context-scaling",
+            "--no-auto-constrain",
+        ]
+
+    def test_tags_table(self, tmp_path: Path) -> None:
+        cfg = tmp_path / "run.toml"
+        _ = cfg.write_text(
+            'model = "foo"\n[tags]\noperator = "ciaranbor"\nrun = "full"\n'
+        )
+        argv = ["context-scaling", "--config", str(cfg)]
+        assert expand_config_in_argv(argv) == [
+            "context-scaling",
+            "--model",
+            "foo",
+            "--tag",
+            "operator=ciaranbor",
+            "--tag",
+            "run=full",
+        ]

bench/configs/context_scaling.example.toml

@@ -0,0 +1,70 @@
+# Example context-scaling run configuration.
+#
+# Use it like this:
+#
+#   uv run python -m bench.cli context-scaling --config bench/configs/context_scaling.example.toml
+#
+# CLI flags placed after `--config` override individual values.
+#
+# All shared and subcommand-specific flags can appear here. Keys mirror the
+# CLI flag names with hyphens replaced by underscores. Boolean keys map to
+# `--key` / `--no-key`; lists are joined as CSV; the `[tags]` table maps to
+# repeated `--tag key=value` flags.
+#
+# NOTE: TOML scoping — once a `[table]` header is opened, all subsequent
+# top-level-looking assignments belong to that table until the next header.
+# Keep tables (like `[tags]`) at the END of the file.
+
+# ---- Model + placement ----
+model = "mlx-community/Qwen3-30B-A3B-4bit"
+# sharding = "Tensor"      # default: "Tensor"; pairs with --comm MlxJaccl + --thunderbolt a2a
+# comm = "MlxJaccl"        # default: "MlxJaccl" (RDMA over Thunderbolt)
+# min_nodes = 1
+
+# ---- Cluster ----
+# Either pin to specific hosts...
+# hosts = ["s4"]
+# ...or let eco pick hosts that satisfy the constraints below.
+# nodes = 1
+# chip = "M3 Ultra"        # eco chip name (case-insensitive substring); comment to allow any
+# thunderbolt = "a2a"      # default: "a2a" (clique, for Tensor+JACCL)
+# "ring" (cycle; for Pipeline+JACCL)
+# "none" (exclude TB; pair with sharding=Pipeline + comm=MlxRing for non-TB hosts)
+
+# Memory + disk minimums are auto-derived from the HF model size
+# (×1.30 + 1 GiB for memory, ×1.10 + 1 GiB for disk). Set any of these
+# explicitly to override the auto-derived value.
+# min_memory_gb = 96.0
+# max_memory_gb = 256.0    # leave bigger machines free for other workloads
+# min_disk_gb = 24.0
+# max_disk_gb = 4000.0
+
+# ---- Runtime ----
+# evict_downloads is true by default — frees disk smallest-first to fit
+# the bench model. Set to false to keep existing downloads.
+# evict_downloads = false
+
+# cleanup_instance is true by default — deletes the placed instance on exit.
+# Set to false to leave it running for debugging.
+# cleanup_instance = false
+
+# ---- Output ----
+output_dir = "bench/results"
+
+# ---- Context-scaling sweep ----
+num_steps = 32 # K — number of equally-spaced ramp points
+# pp_step = 1024              # Δ — explicit override; otherwise auto-derived
+# fraction_of_max = 1.0       # use this fraction of max_position_embeddings
+tg = 64        # tokens generated per step
+# warmup = 2   # default: 2 (1 cache-disabled JIT warmup + 1 cache-priming warmup)
+# cold_controls = "auto"      # 4 evenly-spaced controls across the ramp, or:
+# cold_controls = "8192,16384,32768,40960"   # explicit pp values
+sleep_between_s = 1.0
+
+# ---- Tags ----
+# Survive into metadata.tags in the output JSON; useful for filtering or
+# grouping runs across SHAs / hosts / configs. `$USER` is NOT expanded
+# (TOML is literal); pass `--tag operator=$USER` on the CLI for shell expansion.
+# Must be the LAST table in the file (see TOML scoping note above).
+[tags]
+run = "full"

bench/configs/llama-family-smoke.toml

@@ -0,0 +1,34 @@
+# 4-node smoke campaign: two small/medium Llama models, abbreviated ramps,
+# auto-everything else (TB a2a + tensor + JACCL + auto-derived constraints).
+#
+# Run with:
+#   uv run python -m bench.cli campaign bench/configs/llama-family-smoke.toml
+#
+# Each [[runs]] gets its own cluster (deploy + bench + teardown). After
+# both runs finish, a side-by-side comparison plot is written next to the
+# JSONs.
+
+[defaults]
+nodes = 4
+num_steps = 8
+fraction_of_max = 0.5
+
+[defaults.tags]
+campaign = "llama-family-smoke"
+
+[[runs]]
+subcommand = "context-scaling"
+model = "mlx-community/Llama-3.2-3B-Instruct-4bit"
+[runs.tags]
+model_short = "llama-3.2-3b-4bit"
+
+[[runs]]
+subcommand = "context-scaling"
+model = "mlx-community/Meta-Llama-3.1-8B-Instruct-4bit"
+[runs.tags]
+model_short = "llama-3.1-8b-4bit"
+
+# Final comparison plot (one PNG per benchmark group with ≥2 runs).
+[plot]
+label_tag = "model_short"
+title = "Llama 3 family — 4-node tensor + JACCL smoke"

bench/eval_tool_calls.py

@@ -15,9 +15,8 @@ from pathlib import Path
 from typing import Any, Literal
 
 import httpx
-from harness import (
-    ExoClient,
-    ExoHttpError,
+from exo_tools.client import ExoClient, ExoHttpError
+from exo_tools.harness import (
     add_common_instance_args,
     capture_cluster_snapshot,
     instance_id_from_instance,

bench/exo_bench.py

@@ -24,15 +24,12 @@ import json
 import sys
 import threading
 import time
-from collections.abc import Callable
 from concurrent.futures import ThreadPoolExecutor, as_completed
-from pathlib import Path
 from statistics import mean
 from typing import Any
 
-from harness import (
-    ExoClient,
-    ExoHttpError,
+from exo_tools.client import ExoClient, ExoHttpError
+from exo_tools.harness import (
     add_common_instance_args,
     capture_cluster_snapshot,
     find_existing_instance,
@@ -46,125 +43,44 @@ from harness import (
     wait_for_instance_ready,
 )
 from loguru import logger
-from transformers import AutoTokenizer
 
-# Monkey-patch for transformers 5.x compatibility
-# Kimi's tokenization_kimi.py imports bytes_to_unicode from the old location
-# which was moved in transformers 5.0.0rc2
-try:
-    import transformers.models.gpt2.tokenization_gpt2 as gpt2_tokenization
-    from transformers.convert_slow_tokenizer import bytes_to_unicode
+# PromptSizer / run_one_completion / load_tokenizer_for_bench are the
+# canonical, fully-typed implementations under bench/lib/. They are
+# re-exported here for backwards compatibility with prefill_decode_bench.py
+# and any other consumers of `from exo_bench import …`.
+from bench.lib.completion import run_one_completion as _lib_run_one_completion
+from bench.lib.prompt import (
+    PromptSizer as _LibPromptSizer,
+)
+from bench.lib.prompt import (
+    load_tokenizer_for_bench as _lib_load_tokenizer_for_bench,
+)
 
-    if not hasattr(gpt2_tokenization, "bytes_to_unicode"):
-        gpt2_tokenization.bytes_to_unicode = bytes_to_unicode  # type: ignore[attr-defined]
-except ImportError:
-    pass  # transformers < 5.0 or bytes_to_unicode not available
+PromptSizer = _LibPromptSizer
+load_tokenizer_for_bench = _lib_load_tokenizer_for_bench
 
 
-def load_tokenizer_for_bench(model_id: str) -> Any:
-    """
-    Load tokenizer for benchmarking, with special handling for Kimi models.
-
-    Kimi uses a custom TikTokenTokenizer that transformers 5.x can't load via AutoTokenizer.
-    This function replicates the logic from utils_mlx.py for bench compatibility.
-    """
-    model_id_lower = model_id.lower()
-
-    if "kimi-k2" in model_id_lower:
-        import importlib.util
-        import types
-
-        from huggingface_hub import snapshot_download
-
-        # Download/get the model path
-        model_path = Path(
-            snapshot_download(
-                model_id,
-                allow_patterns=["*.json", "*.py", "*.tiktoken", "*.model", "*.jinja"],
-            )
-        )
-
-        sys.path.insert(0, str(model_path))
-
-        # Load tool_declaration_ts first (tokenization_kimi imports it with relative import)
-        tool_decl_path = model_path / "tool_declaration_ts.py"
-        if tool_decl_path.exists():
-            spec = importlib.util.spec_from_file_location(
-                "tool_declaration_ts", tool_decl_path
-            )
-            if spec and spec.loader:
-                tool_decl_module = importlib.util.module_from_spec(spec)
-                sys.modules["tool_declaration_ts"] = tool_decl_module
-                spec.loader.exec_module(tool_decl_module)
-
-        # Load tokenization_kimi with patched source (convert relative to absolute import)
-        tok_path = model_path / "tokenization_kimi.py"
-        source = tok_path.read_text()
-        source = source.replace("from .tool_declaration_ts", "from tool_declaration_ts")
-        spec = importlib.util.spec_from_file_location("tokenization_kimi", tok_path)
-        if spec:
-            tok_module = types.ModuleType("tokenization_kimi")
-            tok_module.__file__ = str(tok_path)
-            sys.modules["tokenization_kimi"] = tok_module
-            exec(compile(source, tok_path, "exec"), tok_module.__dict__)  # noqa: S102
-            TikTokenTokenizer = tok_module.TikTokenTokenizer  # noqa: N806
-        else:
-            from tokenization_kimi import TikTokenTokenizer  # type: ignore[import-not-found]  # noqa: I001
-
-        hf_tokenizer: Any = TikTokenTokenizer.from_pretrained(model_path)
-
-        # Patch encode to use internal tiktoken model directly
-        # transformers 5.x has a bug in the encode->pad path for slow tokenizers
-        def _patched_encode(text: str, **kwargs: object) -> list[int]:
-            # Pass allowed_special="all" to handle special tokens like <|im_user|>
-            return list(hf_tokenizer.model.encode(text, allowed_special="all"))
-
-        hf_tokenizer.encode = _patched_encode
-
-        return hf_tokenizer
-
-    # TODO: Change back to using only transformers
-    try:
-        return AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
-    except (AttributeError, ValueError):
-        from huggingface_hub import snapshot_download
-        from transformers import PretrainedConfig
-
-        model_path = Path(
-            snapshot_download(
-                model_id,
-                allow_patterns=[
-                    "*.json",
-                    "*.py",
-                    "tokenizer.model",
-                    "*.tiktoken",
-                    "tiktoken.model",
-                    "*.txt",
-                    "*.jsonl",
-                    "*.jinja",
-                ],
-            )
-        )
-        stub_kwargs: dict[str, Any] = {}
-        config_file = model_path / "config.json"
-        if config_file.exists():
-            with open(config_file) as f:
-                raw = json.load(f)
-            for key in (
-                "model_type",
-                "max_position_embeddings",
-                "vocab_size",
-                "bos_token_id",
-                "eos_token_id",
-                "pad_token_id",
-            ):
-                if key in raw:
-                    stub_kwargs[key] = raw[key]
-        return AutoTokenizer.from_pretrained(
-            str(model_path),
-            config=PretrainedConfig(**stub_kwargs),
-            trust_remote_code=True,
-        )
+def run_one_completion(
+    client: ExoClient,
+    model_id: str,
+    pp_hint: int,
+    tg: int,
+    prompt_sizer: PromptSizer,
+    *,
+    use_prefix_cache: bool = False,
+    stream: bool = False,
+) -> tuple[dict[str, Any], int]:
+    """Backwards-compatible shim returning a plain ``dict`` row."""
+    row, pp_tokens = _lib_run_one_completion(
+        client,
+        model_id,
+        pp_hint,
+        tg,
+        prompt_sizer,
+        use_prefix_cache=use_prefix_cache,
+        stream=stream,
+    )
+    return dict(row), pp_tokens
 
 
 def format_peak_memory(b: float) -> str:
@@ -270,159 +186,6 @@ def parse_int_list(values: list[str]) -> list[int]:
     return items
 
 
-def run_one_completion(
-    client: ExoClient,
-    model_id: str,
-    pp_hint: int,
-    tg: int,
-    prompt_sizer: PromptSizer,
-    *,
-    use_prefix_cache: bool = False,
-    stream: bool = False,
-) -> tuple[dict[str, Any], int]:
-    content, pp_tokens = prompt_sizer.build(pp_hint)
-    payload: dict[str, Any] = {
-        "model": model_id,
-        "messages": [{"role": "user", "content": content}],
-        "max_tokens": tg,
-        "logprobs": False,
-        "use_prefix_cache": use_prefix_cache,
-    }
-
-    if not stream:
-        payload["stream"] = False
-        t0 = time.perf_counter()
-        out = client.post_bench_chat_completions(payload)
-        elapsed = time.perf_counter() - t0
-
-        stats = out.get("generation_stats")
-        choices = out.get("choices") or [{}]
-        message = choices[0].get("message", {}) if choices else {}
-        content = message.get("content") or ""
-        preview = content[:200] if content else ""
-    else:
-        tokens = 0
-        first_token_time = None
-        t0 = time.perf_counter()
-        text_parts: list[str] = []
-        stats = None
-
-        for raw_line in client.stream_bench_chat_completions(payload):
-            line = raw_line.strip()
-            if line.startswith(": generation_stats "):
-                with contextlib.suppress(json.JSONDecodeError):
-                    stats = json.loads(line[len(": generation_stats ") :])
-                continue
-            if not line.startswith("data: "):
-                continue
-            data = line[6:]
-            if data == "[DONE]":
-                break
-            try:
-                chunk = json.loads(data)
-                delta = chunk.get("choices", [{}])[0].get("delta", {})
-                if delta.get("content"):
-                    if first_token_time is None:
-                        first_token_time = time.perf_counter()
-                    tokens += 1
-                    text_parts.append(delta["content"])
-            except json.JSONDecodeError:
-                pass
-
-        elapsed = time.perf_counter() - t0
-        preview = "".join(text_parts)[:200]
-
-        if not stats:
-            ttft = (first_token_time - t0) if first_token_time else elapsed
-            gen_time = elapsed - ttft if tokens > 1 else elapsed
-            gen_tps = (tokens - 1) / gen_time if tokens > 1 and gen_time > 0 else 0.0
-            prompt_tps = pp_tokens / ttft if ttft > 0 else 0.0
-            stats = {
-                "prompt_tokens": pp_tokens,
-                "generation_tokens": tokens,
-                "prompt_tps": round(prompt_tps, 2),
-                "generation_tps": round(gen_tps, 2),
-                "peak_memory_usage": {"inBytes": 0},
-            }
-
-    return {
-        "elapsed_s": elapsed,
-        "output_text_preview": preview,
-        "stats": stats,
-    }, pp_tokens
-
-
-class PromptSizer:
-    def __init__(self, tokenizer: Any, atom: str = "a "):
-        self.tokenizer = tokenizer
-        self.atom = atom
-        self.count_fn = PromptSizer._make_counter(tokenizer)
-        self.base_tokens = self.count_fn("")
-
-    @staticmethod
-    def _make_counter(tokenizer: Any) -> Callable[[str], int]:
-        def count_fn(user_content: str) -> int:
-            messages = [{"role": "user", "content": user_content}]
-            try:
-                ids = tokenizer.apply_chat_template(
-                    messages, tokenize=True, add_generation_prompt=True
-                )
-            except ValueError:
-                # Models without a Jinja chat template (e.g. DeepSeek V4 which
-                # ships its own Python encoder). Use the exo-side V4 encoder.
-                from exo.worker.engines.mlx.deepseek_v4_encoding import (
-                    encode_messages as encode_v4,
-                )
-
-                prompt = encode_v4(messages, thinking_mode="thinking")
-                ids = tokenizer.encode(prompt, add_special_tokens=False)
-            # Fix for transformers 5.x
-            if hasattr(ids, "input_ids"):
-                ids = ids.input_ids
-            return int(len(ids))
-
-        return count_fn
-
-    def build(self, target_prompt_tokens: int) -> tuple[str, int]:
-        target = int(target_prompt_tokens)
-        if target < self.base_tokens:
-            raise RuntimeError(
-                f"Target ({target}) is smaller than template overhead ({self.base_tokens})."
-            )
-
-        # Estimate tokens per atom using a sample
-        sample_count = 100
-        sample_content = self.atom * sample_count
-        sample_tokens = self.count_fn(sample_content) - self.base_tokens
-        tokens_per_atom = sample_tokens / sample_count
-
-        # Estimate starting point
-        needed_tokens = target - self.base_tokens
-        estimated_atoms = int(needed_tokens / tokens_per_atom)
-
-        # Binary search to find exact atom count
-        low, high = 0, estimated_atoms * 2 + 100
-        while low < high:
-            mid = (low + high) // 2
-            tok = self.count_fn(self.atom * mid)
-            if tok < target:
-                low = mid + 1
-            else:
-                high = mid
-
-        content = self.atom * low
-        tok = self.count_fn(content)
-        logger.info(f"{tok=}")
-
-        if tok != target:
-            raise RuntimeError(
-                f"Overshot: got {tok} tokens (target {target}). "
-                f"Pick a different atom (try ' a' or '\\n' or '0 ')."
-            )
-
-        return content, tok
-
-
 def main() -> int:
     ap = argparse.ArgumentParser(
         prog="exo-bench",

bench/exo_eval.py

@@ -42,9 +42,8 @@ from pathlib import Path
 from typing import Any
 
 import httpx
-from harness import (
-    ExoClient,
-    ExoHttpError,
+from exo_tools.client import ExoClient, ExoHttpError
+from exo_tools.harness import (
     add_common_instance_args,
     capture_cluster_snapshot,
     find_existing_instance,

bench/lib/__init__.py

@@ -0,0 +1,18 @@
+"""Composable bench library for exo.
+
+Provides reusable building blocks for benchmarks:
+
+- :class:`bench.lib.session.BenchSession` — cluster + instance + client wrapper
+- :class:`bench.lib.results.ResultsBundle` — structured results + JSON writer
+- :func:`bench.lib.cluster.managed_cluster` /
+  :func:`bench.lib.cluster.managed_instance` — eco-managed lifecycle ctx-managers
+- :func:`bench.lib.model_meta.fetch_model_meta` — HF metadata fetcher driving
+  cluster constraints + auto-derived context ramps
+- :mod:`bench.lib.context_scaling` — prompt-TPS / decode-TPS vs context-size sweep
+
+CLI entrypoints under ``bench/cli/`` consume this library via
+``python -m bench.cli <subcommand>``. Adding a new benchmark = (i) write
+``bench/lib/<name>.py`` exposing a typed ``run(session, params, bundle)``
+callable, (ii) write ``bench/cli/<name>.py`` with an ``add_subparser`` and
+a handler, (iii) register it in ``_REGISTRY`` in ``bench/cli/__main__.py``.
+"""

bench/lib/cluster.py

@@ -0,0 +1,215 @@
+"""Eco-managed cluster + instance lifecycle helpers for the bench CLI.
+
+Two context managers:
+
+- :func:`managed_cluster` deploys exo on the requested hosts (or via
+  constraint-based reservation) and tears it down on exit.
+- :func:`managed_instance` resolves the model on the cluster, optionally
+  frees disk via ``--danger-delete-downloads`` (default on for benches),
+  places the instance, and deletes it on exit.
+
+The library never reaches for global state — every call takes an
+explicit :class:`EcoSession`. Callers are expected to instantiate one
+session per CLI invocation and use it across both context managers.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import time
+from collections.abc import Iterator
+from contextlib import contextmanager
+from typing import Any, cast
+
+from exo_tools.client import ExoClient
+from exo_tools.cluster import Chip, ClusterInfo, EcoSession, Thunderbolt
+from exo_tools.harness import (
+    Comm,
+    Sharding,
+    cleanup_all_instances,
+    place_instance,
+    resolve_model_short_id,
+    run_planning_phase,
+)
+from loguru import logger
+
+from .session import BenchSession
+
+
+@contextmanager
+def managed_cluster(
+    eco: EcoSession,
+    *,
+    hosts: list[str] | None = None,
+    count: int = 1,
+    thunderbolt: Thunderbolt | None = None,
+    chip: Chip | None = None,
+    min_memory_gb: float | None = None,
+    max_memory_gb: float | None = None,
+    min_disk_gb: float | None = None,
+    max_disk_gb: float | None = None,
+    deploy_timeout_s: int = 600,
+) -> Iterator[ClusterInfo]:
+    """Deploy exo for the duration of the ``with`` block, then ``eco stop``.
+
+    If ``hosts`` is given, deploys on exactly those hosts (constraint flags
+    are ignored — eco doesn't re-validate the explicit list). Otherwise eco
+    reserves any matching hosts that satisfy all of:
+
+      - ``count`` (number of hosts)
+      - ``thunderbolt`` topology (``A2A``, ``RING``, or ``NONE`` to
+        exclude TB-connected hosts)
+      - ``chip`` (substring match against eco's chip names)
+      - memory bounds (``min_memory_gb`` / ``max_memory_gb``)
+      - disk bounds (``min_disk_gb`` / ``max_disk_gb``)
+    """
+    if hosts:
+        cluster = eco.start_deploy(
+            hosts=hosts[:count],
+            wait=True,
+            timeout=deploy_timeout_s,
+        )
+    else:
+        cluster = eco.start_deploy(
+            count=count,
+            thunderbolt=thunderbolt,
+            chip=chip,
+            min_memory_gb=min_memory_gb,
+            max_memory_gb=max_memory_gb,
+            min_disk_gb=min_disk_gb,
+            max_disk_gb=max_disk_gb,
+            wait=True,
+            timeout=deploy_timeout_s,
+        )
+    logger.info(
+        f"cluster deployed: {len(cluster.hosts)} host(s) "
+        f"({', '.join(cluster.hosts)}); namespace={cluster.namespace}"
+    )
+    try:
+        yield cluster
+    finally:
+        with contextlib.suppress(Exception):
+            eco.stop(cluster.hosts)
+        logger.info("cluster stopped")
+
+
+@contextmanager
+def managed_instance(
+    cluster: ClusterInfo,
+    eco: EcoSession,
+    model_id: str,
+    *,
+    sharding: Sharding = Sharding.PIPELINE,
+    comm: Comm = Comm.RING,
+    min_nodes: int = 1,
+    evict_downloads: bool = True,
+    cleanup_on_exit: bool = True,
+    instance_timeout_s: float = 7200.0,
+    settle_timeout_s: float = 60.0,
+) -> Iterator[BenchSession]:
+    """Resolve the model on the cluster, place an instance, yield a session.
+
+    Steps on entry:
+      1. Resolve ``model_id`` to ``(short_id, full_id)`` against the cluster's
+         ``/models`` endpoint (auto-adds from HuggingFace if missing).
+      2. Run the harness's planning phase: validates each node has enough
+         disk for the model and starts the download (or reuses an existing
+         download). When ``evict_downloads=True`` (the default for benches),
+         this also evicts smaller existing models if disk is short.
+      3. Place the instance, wait for it to be ``RunnerReady``.
+      4. Yield a :class:`BenchSession` pointing at the cluster's primary API.
+
+    On exit: deletes the placed instance (and any other lingering
+    instances) so the cluster is clean for the next benchmark.
+    """
+    client = cluster.make_client(timeout_s=instance_timeout_s)
+
+    short_id, full_id = resolve_model_short_id(client, model_id, force_download=True)
+    logger.info(f"resolved model: short_id={short_id} full_id={full_id}")
+
+    # The planning phase needs a concrete preview (instance + runner-to-shard
+    # mapping) to know which nodes to download to. Pull the placements API
+    # directly and take the first valid one — bench cares about disk +
+    # download, not the specific shard mapping.
+    preview = _first_valid_preview(client, full_id, settle_timeout_s)
+    if preview is None:
+        raise RuntimeError(
+            f"No placement available for {full_id} on cluster {cluster.hosts}"
+        )
+
+    duration = run_planning_phase(
+        client,
+        full_id,
+        preview,
+        danger_delete=evict_downloads,
+        timeout=instance_timeout_s,
+        settle_deadline=None,
+    )
+    if duration is not None:
+        logger.info(f"download: {duration:.1f}s (freshly downloaded)")
+    else:
+        logger.info("download: model already cached on all nodes")
+
+    instance_id = place_instance(
+        client,
+        model_id,
+        sharding=sharding,
+        comm=comm,
+        min_nodes=min_nodes,
+        timeout=instance_timeout_s,
+    )
+    logger.info(f"placed instance {instance_id} ({sharding.value}/{comm.value})")
+
+    sess = BenchSession(
+        cluster=cluster,
+        eco=eco,
+        instance_id=instance_id,
+        model_id=short_id,
+        full_model_id=full_id,
+    )
+    try:
+        yield sess
+    finally:
+        if cleanup_on_exit:
+            with contextlib.suppress(Exception):
+                cleanup_all_instances(sess.client)
+        else:
+            logger.info(
+                f"cleanup_on_exit=False: leaving instance(s) on {cluster.hosts}"
+            )
+
+
+def _first_valid_preview(
+    client: ExoClient, full_model_id: str, settle_timeout_s: float
+) -> dict[str, Any] | None:
+    """Poll ``/instance/previews`` until at least one valid preview comes back."""
+    deadline = time.monotonic() + settle_timeout_s
+    backoff_s = 1.0
+    while True:
+        resp_obj: Any = client.request_json(  # type: ignore[reportAny]
+            "GET", "/instance/previews", params={"model_id": full_model_id}
+        )
+        resp: dict[str, Any] = (
+            cast("dict[str, Any]", resp_obj) if isinstance(resp_obj, dict) else {}
+        )
+        previews_raw: object = resp.get("previews") or []
+        previews: list[Any] = (
+            cast("list[Any]", previews_raw) if isinstance(previews_raw, list) else []
+        )
+        for raw in previews:  # type: ignore[reportAny]
+            if not isinstance(raw, dict):
+                continue
+            entry = cast("dict[str, Any]", raw)
+            if entry.get("error") is not None:
+                continue
+            instance = entry.get("instance")
+            if isinstance(instance, dict):
+                return entry
+        if time.monotonic() >= deadline:
+            return None
+        logger.info(
+            f"waiting for placement to appear for {full_model_id} "
+            f"({deadline - time.monotonic():.0f}s remaining)..."
+        )
+        time.sleep(min(backoff_s, max(0.0, deadline - time.monotonic())))
+        backoff_s = min(backoff_s * 2, 30.0)

bench/lib/completion.py

@@ -0,0 +1,194 @@
+"""Typed wrapper around ``/bench/chat/completions`` for benchmarks.
+
+The bench endpoint disables EOS suppression and KV prefix caching by
+default (see ``bench/METHODOLOGY.md``). This module exposes a single
+function :func:`run_one_completion` that:
+
+  1. Builds an exact-token-length prompt via :class:`PromptSizer`.
+  2. POSTs to ``/bench/chat/completions``.
+  3. Returns a ``(BenchRow, prompt_tokens)`` pair where ``BenchRow`` is a
+     :class:`typing.TypedDict` with the fields the caller needs.
+
+Streaming is supported but rarely needed for context-scaling — the
+non-streaming path is the default.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import json
+import time
+from typing import Any, Literal, NotRequired, TypedDict, cast
+
+from exo_tools.client import ExoClient
+
+from .prompt import PromptSizer
+
+PrefixCacheHit = Literal["none", "partial", "exact"]
+
+
+class GenerationStats(TypedDict, total=False):
+    """Server-reported per-task timing stats."""
+
+    prompt_tps: float
+    generation_tps: float
+    prompt_tokens: int
+    generation_tokens: int
+    peak_memory_usage: dict[str, int]
+    prefix_cache_hit: PrefixCacheHit
+
+
+class BenchRow(TypedDict):
+    """Per-request result row returned to callers."""
+
+    elapsed_s: float
+    output_text_preview: str
+    stats: GenerationStats
+    error: NotRequired[str]
+
+
+def _as_dict(value: Any) -> dict[str, Any]:  # type: ignore[reportAny]
+    """Narrow an arbitrary JSON value to a typed ``dict[str, Any]``."""
+    if isinstance(value, dict):
+        return cast("dict[str, Any]", value)
+    return {}
+
+
+def _as_list(value: Any) -> list[Any]:  # type: ignore[reportAny]
+    if isinstance(value, list):
+        return cast("list[Any]", value)
+    return []
+
+
+def _extract_stats(raw_response: dict[str, Any]) -> GenerationStats:
+    stats_obj = raw_response.get("generation_stats")
+    if not isinstance(stats_obj, dict):
+        return {}
+    return cast("GenerationStats", cast("object", stats_obj))
+
+
+def _extract_preview(raw_response: dict[str, Any], limit: int = 200) -> str:
+    choices = _as_list(raw_response.get("choices"))
+    if not choices:
+        return ""
+    first = _as_dict(choices[0])
+    message = _as_dict(first.get("message"))
+    content_obj = message.get("content")
+    if isinstance(content_obj, str):
+        return content_obj[:limit]
+    return ""
+
+
+def run_one_completion(
+    client: ExoClient,
+    model_id: str,
+    pp_hint: int,
+    tg: int,
+    prompt_sizer: PromptSizer,
+    *,
+    use_prefix_cache: bool = False,
+    stream: bool = False,
+) -> tuple[BenchRow, int]:
+    """Send one request to ``/bench/chat/completions`` and return its row.
+
+    ``pp_hint`` is the *target* prompt-token count; the actual prompt is
+    sized via :class:`PromptSizer` and the verified value is returned as
+    the second element of the tuple.
+    """
+    content, pp_tokens = prompt_sizer.build(pp_hint)
+    payload: dict[str, Any] = {
+        "model": model_id,
+        "messages": [{"role": "user", "content": content}],
+        "max_tokens": tg,
+        "logprobs": False,
+        "use_prefix_cache": use_prefix_cache,
+    }
+
+    if not stream:
+        payload["stream"] = False
+        t0 = time.perf_counter()
+        raw_obj = client.post_bench_chat_completions(payload)
+        elapsed = time.perf_counter() - t0
+        raw = _as_dict(raw_obj)
+        return (
+            BenchRow(
+                elapsed_s=elapsed,
+                output_text_preview=_extract_preview(raw),
+                stats=_extract_stats(raw),
+            ),
+            pp_tokens,
+        )
+
+    return _run_streaming(client, payload, pp_tokens)
+
+
+def _run_streaming(
+    client: ExoClient,
+    payload: dict[str, Any],
+    pp_tokens: int,
+) -> tuple[BenchRow, int]:
+    """Streaming variant: parse SSE lines, recover ``GenerationStats``."""
+    payload = {**payload, "stream": True}
+
+    tokens = 0
+    first_token_time: float | None = None
+    t0 = time.perf_counter()
+    text_parts: list[str] = []
+    stats: GenerationStats = {}
+
+    for raw_line in client.stream_bench_chat_completions(payload):
+        line = raw_line.strip()
+        if line.startswith(": generation_stats "):
+            with contextlib.suppress(json.JSONDecodeError):
+                parsed_obj: Any = json.loads(  # type: ignore[reportAny]
+                    line[len(": generation_stats ") :]
+                )
+                if isinstance(parsed_obj, dict):
+                    stats = cast("GenerationStats", cast("object", parsed_obj))
+            continue
+        if not line.startswith("data: "):
+            continue
+        data = line[6:]
+        if data == "[DONE]":
+            break
+        try:
+            chunk_obj: Any = json.loads(data)  # type: ignore[reportAny]
+        except json.JSONDecodeError:
+            continue
+        chunk = _as_dict(chunk_obj)
+        choices = _as_list(chunk.get("choices"))
+        if not choices:
+            continue
+        first = _as_dict(choices[0])
+        delta = _as_dict(first.get("delta"))
+        delta_content_obj = delta.get("content")
+        if isinstance(delta_content_obj, str) and delta_content_obj:
+            if first_token_time is None:
+                first_token_time = time.perf_counter()
+            tokens += 1
+            text_parts.append(delta_content_obj)
+
+    elapsed = time.perf_counter() - t0
+    preview = "".join(text_parts)[:200]
+
+    if not stats:
+        ttft = (first_token_time - t0) if first_token_time is not None else elapsed
+        gen_time = elapsed - ttft if tokens > 1 else elapsed
+        gen_tps = (tokens - 1) / gen_time if tokens > 1 and gen_time > 0 else 0.0
+        prompt_tps = pp_tokens / ttft if ttft > 0 else 0.0
+        stats = GenerationStats(
+            prompt_tokens=pp_tokens,
+            generation_tokens=tokens,
+            prompt_tps=round(prompt_tps, 2),
+            generation_tps=round(gen_tps, 2),
+            peak_memory_usage={"inBytes": 0},
+        )
+
+    return (
+        BenchRow(
+            elapsed_s=elapsed,
+            output_text_preview=preview,
+            stats=stats,
+        ),
+        pp_tokens,
+    )

bench/lib/context_scaling.py

@@ -0,0 +1,428 @@
+"""Prompt-TPS / decode-TPS vs context-size sweep.
+
+Methodology (see also ``bench/METHODOLOGY.md``):
+
+  Run a single ascending ramp of equally-spaced prompt lengths
+  ``pp ∈ {Δ, 2Δ, …, K·Δ}`` with ``prefix_cache=enabled``, ``repeat=1``,
+  ``concurrency=1`` and one warmup at ``pp=Δ``.
+
+  Because each step's prefix is exactly what the previous step left in
+  the cache, every step beyond the first is a *partial* hit and the
+  server-reported ``prompt_tps`` reflects the true cold rate over the
+  fresh ``Δ``-token suffix. We accept the warmup's reported rate as the
+  cold equivalent for ``pp=Δ`` (the warmup itself is the cold prefill).
+
+  ``decode TPS`` is independent of prefill mechanics — every step's
+  ``generation_tps`` is a real decode-rate-at-N data point.
+
+  Cumulative cold-prefill upper bound:
+      ``T_cum(pp_k) = Σ_{i=1..k} (Δ_i / prompt_tps_i)``
+
+  Optional cold-control points (``prefix_cache=disabled``) validate the
+  approximation; the gap quantifies per-task overhead. To preserve the
+  ``none`` cache-hit classification AND ensure the request actually
+  hits a freshly-placed runner (the master picks the instance with the
+  lowest in-flight task count, which is non-deterministic when multiple
+  same-model instances exist), :func:`run` deletes the sweep instance
+  *before* invoking the cold-control factory. The factory itself places
+  a fresh instance per control and deletes it on exit; the
+  :func:`bench.lib.cluster.managed_instance` ctx-manager calls
+  ``cleanup_all_instances`` on exit as a final safety net.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import time
+from collections.abc import Callable, Iterator
+from contextlib import AbstractContextManager, contextmanager
+from dataclasses import asdict, dataclass
+from typing import Any
+
+from exo_tools.client import ExoClient, ExoHttpError
+from exo_tools.harness import (
+    Comm,
+    Sharding,
+    place_instance,
+    wait_for_instance_gone,
+)
+from loguru import logger
+
+from .completion import GenerationStats, PrefixCacheHit, run_one_completion
+from .prompt import PromptSizer
+from .results import ResultsBundle
+from .session import BenchSession
+
+
+@dataclass(frozen=True)
+class ContextScalingParams:
+    """Inputs for a single context-scaling sweep."""
+
+    pp_step: int
+    num_steps: int
+    tg: int
+    warmup: int = 1
+    cold_controls: tuple[int, ...] = ()
+    sleep_between_s: float = 1.0
+
+
+@dataclass
+class StepResult:
+    pp_tokens: int
+    delta_tokens: int
+    prompt_tps: float
+    generation_tps: float
+    prefix_cache_hit: PrefixCacheHit | str
+    prompt_tokens: int
+    generation_tokens: int
+    elapsed_s: float
+    peak_memory_bytes: int = 0
+    output_text_preview: str = ""
+
+
+def _peak_bytes(stats: GenerationStats) -> int:
+    pm = stats.get("peak_memory_usage") or {}
+    return int(pm.get("inBytes") or pm.get("in_bytes") or 0)
+
+
+def _build_step_result(
+    pp_tokens: int,
+    delta_tokens: int,
+    elapsed_s: float,
+    output_text_preview: str,
+    stats: GenerationStats,
+) -> StepResult:
+    return StepResult(
+        pp_tokens=pp_tokens,
+        delta_tokens=delta_tokens,
+        prompt_tps=float(stats.get("prompt_tps") or 0.0),
+        generation_tps=float(stats.get("generation_tps") or 0.0),
+        prefix_cache_hit=stats.get("prefix_cache_hit") or "unknown",
+        prompt_tokens=int(stats.get("prompt_tokens") or pp_tokens),
+        generation_tokens=int(stats.get("generation_tokens") or 0),
+        elapsed_s=elapsed_s,
+        peak_memory_bytes=_peak_bytes(stats),
+        output_text_preview=output_text_preview[:200],
+    )
+
+
+def _run_request(
+    client: ExoClient,
+    full_model_id: str,
+    pp: int,
+    tg: int,
+    sizer: PromptSizer,
+    *,
+    use_prefix_cache: bool,
+) -> tuple[StepResult, int]:
+    """Send one request and return ``(StepResult, actual_pp_tokens)``."""
+    row, actual_pp = run_one_completion(
+        client,
+        full_model_id,
+        pp,
+        tg,
+        sizer,
+        use_prefix_cache=use_prefix_cache,
+        stream=False,
+    )
+    step = _build_step_result(
+        pp_tokens=actual_pp,
+        delta_tokens=actual_pp,  # caller overrides for cached sweep
+        elapsed_s=row["elapsed_s"],
+        output_text_preview=row["output_text_preview"],
+        stats=row["stats"],
+    )
+    return step, actual_pp
+
+
+def _compute_t_cum(steps: list[StepResult]) -> list[float]:
+    t_cum = 0.0
+    out: list[float] = []
+    for s in steps:
+        if s.prompt_tps > 0 and s.delta_tokens > 0:
+            t_cum += s.delta_tokens / s.prompt_tps
+        out.append(round(t_cum, 6))
+    return out
+
+
+def run_cached_sweep(
+    session: BenchSession,
+    params: ContextScalingParams,
+    bundle: ResultsBundle,
+) -> list[StepResult]:
+    """Run the ascending PP sweep with ``prefix_cache=enabled``.
+
+    Mutates ``bundle.runs`` in place and returns the typed step list.
+    """
+    if session.full_model_id is None:
+        raise RuntimeError(
+            "BenchSession.full_model_id must be set for context-scaling."
+        )
+
+    sizer = session.get_prompt_sizer()
+    client = session.client
+    pp_targets = [params.pp_step * i for i in range(1, params.num_steps + 1)]
+    logger.info(
+        f"context-scaling: K={params.num_steps} steps, Δ={params.pp_step} tokens, "
+        f"tg={params.tg}, warmup={params.warmup}, cached"
+    )
+
+    # Warmup discipline:
+    #   - First warmup runs with the prefix cache DISABLED. This triggers
+    #     the MLX kernel JIT compile + KV-buffer alloc for this exact
+    #     (Δ, dtype, batch) shape, but does NOT write a cache entry — so
+    #     the cold-with-JIT rate isn't fossilised.
+    #   - Subsequent warmups run with the prefix cache ENABLED. The
+    #     second one finds an empty cache, does a real cold prefill with
+    #     a HOT kernel, and writes the resulting rate into the cache
+    #     entry at pp=Δ.
+    #   - Step 0 (also cache-enabled) is then an exact hit on that entry
+    #     and reports the hot rate.
+    # Default warmup=2 gives both effects; warmup=1 still does the JIT
+    # warmup but leaves step 0 as a "none" hit (cold prefill at the hot
+    # kernel, creates the cache entry on the way through).
+    for w in range(params.warmup):
+        is_jit_warmup = w == 0
+        kind = "JIT warmup" if is_jit_warmup else "cache-prime warmup"
+        logger.info(
+            f"  warmup {w + 1}/{params.warmup} ({kind}, pp={params.pp_step})"
+        )
+        _run_request(
+            client,
+            session.full_model_id,
+            params.pp_step,
+            params.tg,
+            sizer,
+            use_prefix_cache=not is_jit_warmup,
+        )
+
+    steps: list[StepResult] = []
+    prev_pp = 0
+    for i, pp in enumerate(pp_targets):
+        time.sleep(params.sleep_between_s)
+        try:
+            step, actual_pp = _run_request(
+                client,
+                session.full_model_id,
+                pp,
+                params.tg,
+                sizer,
+                use_prefix_cache=True,
+            )
+        except Exception as e:
+            logger.error(f"step {i + 1}/{params.num_steps} (pp={pp}) failed: {e}")
+            raise
+
+        step.delta_tokens = actual_pp - prev_pp
+        steps.append(step)
+        bundle.runs.append({"step_index": i, "phase": "cached_sweep", **asdict(step)})
+        logger.info(
+            f"  step {i + 1}/{params.num_steps} pp={actual_pp} Δ={step.delta_tokens} "
+            f"prompt_tps={step.prompt_tps:.1f} gen_tps={step.generation_tps:.2f} "
+            f"hit={step.prefix_cache_hit}"
+        )
+        prev_pp = actual_pp
+
+    return steps
+
+
+def run_cold_controls(
+    factory: Callable[[], AbstractContextManager[ExoClient]],
+    session: BenchSession,
+    params: ContextScalingParams,
+    bundle: ResultsBundle,
+) -> list[StepResult]:
+    """Run cold-control points on a fresh instance to preserve ``none`` hits.
+
+    A cold control is a single request at ``pp=N`` with
+    ``prefix_cache=disabled``, executed against a freshly-placed instance
+    (and with no other same-model instance live, so the master's task
+    routing is deterministic). The caller is expected to delete the
+    sweep instance before invoking this — see :func:`run`.
+    """
+    if not params.cold_controls:
+        return []
+    if session.full_model_id is None:
+        raise RuntimeError("BenchSession.full_model_id must be set for cold controls.")
+
+    sizer = session.get_prompt_sizer()
+    out: list[StepResult] = []
+    for control_pp in params.cold_controls:
+        logger.info(f"cold control: pp={control_pp} (fresh instance, cache disabled)")
+        with factory() as fresh_client:
+            step, actual_pp = _run_request(
+                fresh_client,
+                session.full_model_id,
+                control_pp,
+                params.tg,
+                sizer,
+                use_prefix_cache=False,
+            )
+        step.delta_tokens = actual_pp
+        out.append(step)
+        bundle.cold_controls.append({"phase": "cold_control", **asdict(step)})
+        logger.info(
+            f"  cold pp={actual_pp} prompt_tps={step.prompt_tps:.1f} "
+            f"gen_tps={step.generation_tps:.2f} hit={step.prefix_cache_hit}"
+        )
+        if step.prefix_cache_hit != "none":
+            logger.warning(
+                f"cold control at pp={actual_pp} reported "
+                f"prefix_cache_hit={step.prefix_cache_hit!r}; "
+                f"control may not be cold."
+            )
+    return out
+
+
+def derive_summary(
+    steps: list[StepResult],
+    cold_controls: list[StepResult],
+) -> dict[str, Any]:
+    """Compute the cumulative cold-prefill upper bound + control gaps."""
+    t_cum = _compute_t_cum(steps)
+    bracketed = sorted(
+        ((s.pp_tokens, t) for s, t in zip(steps, t_cum, strict=True)),
+        key=lambda x: x[0],
+    )
+    control_gaps: list[dict[str, float]] = []
+    for ctrl in cold_controls:
+        cold_t = ctrl.pp_tokens / ctrl.prompt_tps if ctrl.prompt_tps > 0 else 0.0
+        cum_t = _interp(bracketed, ctrl.pp_tokens)
+        gap = cum_t - cold_t
+        control_gaps.append(
+            {
+                "pp_tokens": ctrl.pp_tokens,
+                "cold_t_seconds": round(cold_t, 4),
+                "t_cum_seconds_at_pp": round(cum_t, 4),
+                "gap_seconds": round(gap, 4),
+                "gap_fraction": round(gap / cold_t, 4) if cold_t > 0 else 0.0,
+            }
+        )
+    return {
+        "t_cum_seconds": t_cum,
+        "control_gaps": control_gaps,
+    }
+
+
+def _interp(points: list[tuple[int, float]], x: int) -> float:
+    """Linear interpolate y at x, given sorted ``(x, y)`` points."""
+    if not points:
+        return 0.0
+    if x <= points[0][0]:
+        return points[0][1]
+    if x >= points[-1][0]:
+        return points[-1][1]
+    for i in range(1, len(points)):
+        x0, y0 = points[i - 1]
+        x1, y1 = points[i]
+        if x0 <= x <= x1 and x1 != x0:
+            return y0 + (y1 - y0) * (x - x0) / (x1 - x0)
+    return points[-1][1]
+
+
+# ---------------------------------------------------------------------------
+# Cold-control instance factory
+# ---------------------------------------------------------------------------
+
+
+def make_cold_control_factory(
+    session: BenchSession,
+    sharding: Sharding,
+    comm: Comm,
+    min_nodes: int,
+    instance_timeout_s: float = 1800.0,
+) -> Callable[[], AbstractContextManager[ExoClient]]:
+    """Return a callable yielding a context manager that places a fresh instance.
+
+    Each ``with factory() as client:`` block places a brand-new instance,
+    yields its client, then deletes the instance on exit. Used to isolate
+    cold-control runs.
+
+    The caller is responsible for ensuring no other same-model instance is
+    live during the ``with`` block — otherwise master routing is
+    non-deterministic and the cold control may be served by a stale runner.
+    See :func:`run` for the orchestration.
+    """
+
+    @contextmanager
+    def factory() -> Iterator[ExoClient]:
+        if session.full_model_id is None:
+            raise RuntimeError("session.full_model_id is unset")
+        client = session.client
+        instance_id = place_instance(
+            client,
+            session.full_model_id,
+            sharding=sharding,
+            comm=comm,
+            min_nodes=min_nodes,
+            timeout=instance_timeout_s,
+        )
+        try:
+            yield client
+        finally:
+            with contextlib.suppress(ExoHttpError):
+                client.request_json("DELETE", f"/instance/{instance_id}")
+            with contextlib.suppress(Exception):
+                wait_for_instance_gone(client, instance_id, timeout=60.0)
+
+    return factory
+
+
+def _delete_instance(client: ExoClient, instance_id: str) -> None:
+    """Best-effort delete of a placed instance."""
+    with contextlib.suppress(ExoHttpError):
+        client.request_json("DELETE", f"/instance/{instance_id}")
+    with contextlib.suppress(Exception):
+        wait_for_instance_gone(client, instance_id, timeout=60.0)
+
+
+def run(
+    session: BenchSession,
+    params: ContextScalingParams,
+    bundle: ResultsBundle,
+    *,
+    cold_control_factory: Callable[[], AbstractContextManager[ExoClient]] | None = None,
+) -> ResultsBundle:
+    """End-to-end: cached sweep + optional cold controls + derived summary.
+
+    To make cold controls truly isolated from the sweep instance, we
+    delete the sweep instance *before* running the controls (otherwise
+    the master might route a control's request to the stale sweep
+    instance, since both match the same ``model_id``). The controls then
+    each place their own fresh instance via the factory.
+    """
+    bundle.params.update(
+        {
+            "pp_step": params.pp_step,
+            "num_steps": params.num_steps,
+            "tg": params.tg,
+            "warmup": params.warmup,
+            "cold_controls": list(params.cold_controls),
+            "sleep_between_s": params.sleep_between_s,
+            "model_id": session.model_id,
+            "full_model_id": session.full_model_id,
+        }
+    )
+    bundle.capture_cluster(session.client)
+
+    cached_steps = run_cached_sweep(session, params, bundle)
+
+    cold_steps: list[StepResult] = []
+    if params.cold_controls and cold_control_factory is not None:
+        # Delete the sweep instance so the cold-control fresh instance is
+        # the only same-model instance live for the duration of the controls.
+        if session.instance_id is not None:
+            logger.info(
+                f"cold controls: deleting sweep instance {session.instance_id} "
+                "to isolate fresh instance routing"
+            )
+            _delete_instance(session.client, session.instance_id)
+            session.instance_id = None
+        cold_steps = run_cold_controls(cold_control_factory, session, params, bundle)
+    elif params.cold_controls and cold_control_factory is None:
+        logger.warning(
+            "Cold controls requested but no cold_control_factory supplied; skipping."
+        )
+
+    bundle.derived.update(derive_summary(cached_steps, cold_steps))
+    return bundle

bench/lib/model_meta.py

@@ -0,0 +1,183 @@
+"""Fetch HuggingFace model metadata for benchmark planning.
+
+Two pieces of metadata drive every benchmark we run:
+
+  1. **Total weight size** — used to derive ``min-memory`` and ``min-disk``
+     constraints when picking a host. We sum the sizes of all
+     ``.safetensors`` (or ``.bin``) shards from the repo's file listing.
+  2. **Max position embeddings** — the model's training context length.
+     Used to bound a context-scaling sweep at the model's max context, and
+     to derive a sensible Δ given a target step count.
+
+The fetcher uses the ``huggingface_hub`` python API, which talks to the
+public HF Hub HTTPS endpoints — no exo cluster required, no download
+of weights.
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass, field
+from typing import Any, cast
+
+# Files that count toward the on-disk weight footprint.
+_WEIGHT_SUFFIXES = (".safetensors", ".bin", ".gguf", ".pt", ".npz")
+
+
+@dataclass(frozen=True)
+class ModelMeta:
+    """Subset of HF metadata that a benchmark needs."""
+
+    model_id: str
+    total_weight_bytes: int
+    max_position_embeddings: int
+    num_hidden_layers: int
+    raw_config: dict[str, Any] = field(default_factory=dict)
+
+    @property
+    def total_weight_gb(self) -> float:
+        return self.total_weight_bytes / (1024**3)
+
+    @property
+    def memory_constraint_gb(self) -> float:
+        """Estimated minimum host memory to hold weights + overhead.
+
+        Picks the model size + 30 % headroom (KV cache, activations,
+        framework bookkeeping). Rounded up to the next whole GiB.
+        """
+        return float(int(self.total_weight_gb * 1.30) + 1)
+
+    @property
+    def disk_constraint_gb(self) -> float:
+        """Disk space the host must have free for the download."""
+        return float(int(self.total_weight_gb * 1.10) + 1)
+
+
+def _read_config_json(model_id: str) -> dict[str, Any]:
+    from huggingface_hub import (
+        hf_hub_download,  # type: ignore[reportUnknownVariableType]
+    )
+
+    raw_path = hf_hub_download(repo_id=model_id, filename="config.json", dry_run=False)
+    with open(raw_path) as f:
+        loaded: Any = json.load(f)  # type: ignore[reportAny]
+    return cast("dict[str, Any]", loaded) if isinstance(loaded, dict) else {}
+
+
+def _sum_weight_sizes(model_id: str) -> int:
+    """Sum sizes of all weight-shard files in the repo's file listing."""
+    from huggingface_hub import HfApi
+
+    api = HfApi()
+    info = api.model_info(repo_id=model_id, files_metadata=True)
+    siblings = info.siblings or []
+    total = 0
+    for sib in siblings:
+        rfilename = getattr(sib, "rfilename", None)
+        size = getattr(sib, "size", None)
+        if not isinstance(rfilename, str) or not isinstance(size, int):
+            continue
+        if any(rfilename.endswith(suf) for suf in _WEIGHT_SUFFIXES):
+            total += size
+    return total
+
+
+def _first_int(config: dict[str, Any], *keys: str) -> int:
+    """Return the first key from ``config`` that holds a usable positive int."""
+    for key in keys:
+        value = config.get(key)
+        if isinstance(value, int) and value > 0:
+            return value
+        if isinstance(value, str):
+            try:
+                parsed = int(value)
+            except ValueError:
+                continue
+            if parsed > 0:
+                return parsed
+    return 0
+
+
+def fetch_model_meta(model_id: str) -> ModelMeta:
+    """Fetch the metadata our benchmarks care about for ``model_id``.
+
+    Args:
+        model_id: HuggingFace repo id, e.g. ``mlx-community/Qwen3-30B-A3B-4bit``.
+
+    Returns:
+        Populated :class:`ModelMeta`.
+
+    Raises:
+        Exception: any HTTP / parse error from ``huggingface_hub`` propagates.
+    """
+    config = _read_config_json(model_id)
+    return ModelMeta(
+        model_id=model_id,
+        total_weight_bytes=_sum_weight_sizes(model_id),
+        max_position_embeddings=_first_int(
+            config,
+            "max_position_embeddings",
+            "max_seq_len",
+            "model_max_length",
+            "n_positions",
+        ),
+        num_hidden_layers=_first_int(
+            config,
+            "num_hidden_layers",
+            "num_layers",
+            "n_layer",
+            "n_layers",
+            "num_decoder_layers",
+        ),
+        raw_config=config,
+    )
+
+
+def derive_context_ramp(
+    meta: ModelMeta,
+    *,
+    num_steps: int,
+    fraction_of_max: float = 1.0,
+    min_pp_step: int = 256,
+    round_to: int = 256,
+) -> tuple[int, int]:
+    """Pick ``(pp_step, num_steps)`` covering ``fraction_of_max`` of the context.
+
+    Δ is rounded down to the nearest ``round_to`` so the per-step prompt is a
+    clean number, and clamped to ``min_pp_step`` for tiny-context models.
+    """
+    if meta.max_position_embeddings <= 0:
+        raise ValueError(
+            f"{meta.model_id} reports max_position_embeddings=0 in config.json"
+        )
+    if not (0.0 < fraction_of_max <= 1.0):
+        raise ValueError(f"fraction_of_max must be in (0, 1], got {fraction_of_max}")
+    if num_steps <= 0:
+        raise ValueError(f"num_steps must be >0, got {num_steps}")
+
+    target_max = int(meta.max_position_embeddings * fraction_of_max)
+    raw_step = max(min_pp_step, target_max // num_steps)
+    pp_step = (raw_step // round_to) * round_to or round_to
+    return pp_step, num_steps
+
+
+def derive_cold_controls(
+    meta: ModelMeta,
+    *,
+    pp_step: int,
+    num_steps: int,
+    count: int = 4,
+) -> tuple[int, ...]:
+    """Pick ``count`` evenly-spaced cold-control points across the ramp.
+
+    Always includes the largest ramp point (``pp_step * num_steps``).
+    Returns control pp values in ascending order, deduped.
+    """
+    if count <= 0:
+        return ()
+    max_pp = pp_step * num_steps
+    if count == 1:
+        return (max_pp,)
+    spaced = sorted({(max_pp * (i + 1)) // count for i in range(count)})
+    # Filter out anything below pp_step (a control at <Δ is meaningless).
+    return tuple(p for p in spaced if p >= pp_step)

bench/lib/plotting.py

@@ -0,0 +1,308 @@
+"""Typed matplotlib renderers for benchmark JSON results.
+
+This module owns the *visualisation* of bench results, mirroring how
+``bench/lib/<name>.py`` owns the methodology and ``bench/cli/<name>.py``
+owns the orchestration. Adding plotting for a new benchmark = a new
+``render_<name>`` function here + a dispatch entry in ``bench/cli/plot.py``.
+
+Functions take typed inputs (``Path`` lists, options) and write a PNG.
+They never touch argparse or stdout — that's the CLI's job.
+
+matplotlib's type stubs are thin (most return values are ``Any``), so all
+calls into ``pyplot`` are concentrated at the bottom of this file with
+targeted ``# type: ignore[reportUnknownMemberType, reportAny]`` per line.
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, cast
+
+# Tab10 cycle from matplotlib's default; we pick colours by index ourselves
+# instead of fishing them out of `Line2D.get_color()` so the strict-type
+# fallout stays small and predictable.
+_COLOR_CYCLE: tuple[str, ...] = (
+    "C0",
+    "C1",
+    "C2",
+    "C3",
+    "C4",
+    "C5",
+    "C6",
+    "C7",
+    "C8",
+    "C9",
+)
+
+
+@dataclass(frozen=True)
+class PlotInputs:
+    """Inputs for any benchmark renderer.
+
+    Attributes:
+        results: One or more bench JSON files. The first is used to
+            auto-derive the title when ``title`` is unset.
+        output: Path to write the PNG to.
+        label_tag: When set, use ``metadata.tags[label_tag]`` as the
+            legend label for each run; otherwise use the run id.
+        title: Override for the figure title.
+    """
+
+    results: list[Path]
+    output: Path
+    label_tag: str | None = None
+    title: str | None = None
+
+
+@dataclass(frozen=True)
+class _RunSeries:
+    """Pre-extracted plot data for one results JSON.
+
+    ``cached_prefill_seconds`` is the cumulative cold-prefill estimate
+    (``T_cum`` from the methodology — read from ``derived.t_cum_seconds``).
+    ``control_prefill_seconds`` is the actual cold prefill time per
+    control (``pp_tokens / prompt_tps`` from the cold-control row).
+    """
+
+    label: str
+    cached_pp: list[int]
+    cached_prefill_seconds: list[float]
+    cached_gen_tps: list[float]
+    control_pp: list[int]
+    control_prefill_seconds: list[float]
+
+
+# ---------------------------------------------------------------------------
+# Pure data extraction (strict-typed, no matplotlib)
+# ---------------------------------------------------------------------------
+
+
+def _load(path: Path) -> dict[str, Any]:
+    """Read a bench JSON file and assert top-level shape."""
+    with path.open() as f:
+        loaded: Any = json.load(f)  # type: ignore[reportAny]
+    if not isinstance(loaded, dict):
+        raise ValueError(f"{path}: expected top-level JSON object")
+    return cast("dict[str, Any]", loaded)
+
+
+# dict[str, Any].get(...) returns Any. The five _get_* helpers below
+# concentrate the Any boundary so the rest of the module can be strict.
+
+
+def _get_dict(d: dict[str, Any], key: str) -> dict[str, Any]:
+    val: Any = d.get(key)
+    return cast("dict[str, Any]", val) if isinstance(val, dict) else {}
+
+
+def _get_list(d: dict[str, Any], key: str) -> list[Any]:
+    val: Any = d.get(key)
+    return cast("list[Any]", val) if isinstance(val, list) else []
+
+
+def _get_str(d: dict[str, Any], key: str, default: str = "") -> str:
+    val: Any = d.get(key, default)  # type: ignore[reportAny]
+    return val if isinstance(val, str) else default
+
+
+def _get_int(row: dict[str, Any], key: str) -> int:
+    val: Any = row.get(key, 0)  # type: ignore[reportAny]
+    if isinstance(val, bool):  # bool is int; reject explicitly
+        return 0
+    if isinstance(val, (int, float)):
+        return int(val)
+    if isinstance(val, str):
+        try:
+            return int(float(val))
+        except ValueError:
+            return 0
+    return 0
+
+
+def _get_float(row: dict[str, Any], key: str) -> float:
+    val: Any = row.get(key, 0.0)  # type: ignore[reportAny]
+    if isinstance(val, bool):
+        return 0.0
+    if isinstance(val, (int, float)):
+        return float(val)
+    if isinstance(val, str):
+        try:
+            return float(val)
+        except ValueError:
+            return 0.0
+    return 0.0
+
+
+def _label_for(data: dict[str, Any], label_tag: str | None) -> str:
+    if label_tag is not None:
+        tags = _get_dict(_get_dict(data, "metadata"), "tags")
+        if label_tag in tags:
+            return _get_str(tags, label_tag, "(unnamed)")
+    return _get_str(_get_dict(data, "metadata"), "run_id", "(unnamed)")
+
+
+def _extract_series(data: dict[str, Any], label: str) -> _RunSeries:
+    """Pre-extract typed lists from a context-scaling bench JSON."""
+    cached_pp: list[int] = []
+    cached_gen_tps: list[float] = []
+    for raw in _get_list(data, "runs"):  # type: ignore[reportAny]
+        if not isinstance(raw, dict):
+            continue
+        row = cast("dict[str, Any]", raw)
+        if _get_str(row, "phase") != "cached_sweep":
+            continue
+        cached_pp.append(_get_int(row, "pp_tokens"))
+        cached_gen_tps.append(_get_float(row, "generation_tps"))
+
+    # Cumulative cold-prefill estimate is computed in derive_summary and
+    # written to derived.t_cum_seconds (parallel to the cached steps).
+    derived = _get_dict(data, "derived")
+    t_cum_raw = _get_list(derived, "t_cum_seconds")
+    cached_prefill_seconds: list[float] = []
+    for raw in t_cum_raw:  # type: ignore[reportAny]
+        if isinstance(raw, (int, float)) and not isinstance(raw, bool):
+            cached_prefill_seconds.append(float(raw))
+
+    # Cold controls give us the actual cold prefill time directly:
+    # pp_tokens / prompt_tps. Skip rows with zero/missing prompt_tps.
+    control_pp: list[int] = []
+    control_prefill_seconds: list[float] = []
+    for raw in _get_list(data, "cold_controls"):  # type: ignore[reportAny]
+        if not isinstance(raw, dict):
+            continue
+        row = cast("dict[str, Any]", raw)
+        pp = _get_int(row, "pp_tokens")
+        tps = _get_float(row, "prompt_tps")
+        if pp > 0 and tps > 0:
+            control_pp.append(pp)
+            control_prefill_seconds.append(pp / tps)
+
+    return _RunSeries(
+        label=label,
+        cached_pp=cached_pp,
+        cached_prefill_seconds=cached_prefill_seconds,
+        cached_gen_tps=cached_gen_tps,
+        control_pp=control_pp,
+        control_prefill_seconds=control_prefill_seconds,
+    )
+
+
+def _auto_title(data: dict[str, Any]) -> str:
+    metadata = _get_dict(data, "metadata")
+    params = _get_dict(data, "params")
+    model = (
+        _get_str(params, "full_model_id") or _get_str(params, "model_id") or "(unknown)"
+    )
+    sha = _get_str(metadata, "exo_sha") or "(no-sha)"
+    host = _get_str(metadata, "hostname") or "(no-host)"
+    return f"{model}\n{sha} on {host}"
+
+
+# ---------------------------------------------------------------------------
+# Matplotlib boundary — each call site has a narrow, justified ignore.
+# ---------------------------------------------------------------------------
+
+
+def render_context_scaling(inputs: PlotInputs) -> Path:
+    """Render a 2-panel context-scaling plot.
+
+    Top:    pp_tokens vs prompt_tps (line per run; cold controls as 'x' scatter)
+    Bottom: pp_tokens vs generation_tps (line per run)
+
+    Each line is a separate result file. Multi-file mode is for comparing
+    runs across exo SHAs / hosts / configs; the title is taken from the
+    first file's metadata unless ``inputs.title`` is set.
+    """
+    if not inputs.results:
+        raise ValueError("at least one results JSON path is required")
+
+    # Validate + extract first so any data-shape error surfaces before we
+    # even import matplotlib.
+    first_data: dict[str, Any] | None = None
+    series: list[_RunSeries] = []
+    for path in inputs.results:
+        data = _load(path)
+        if first_data is None:
+            first_data = data
+        benchmark = _get_str(_get_dict(data, "metadata"), "benchmark")
+        if benchmark != "context_scaling":
+            raise ValueError(
+                f"{path}: expected benchmark=='context_scaling', got {benchmark!r}"
+            )
+        series.append(_extract_series(data, _label_for(data, inputs.label_tag)))
+
+    title = inputs.title
+    if title is None and first_data is not None:
+        title = _auto_title(first_data)
+        if len(inputs.results) > 1:
+            title = f"{title}\n(comparison of {len(inputs.results)} runs)"
+
+    inputs.output.parent.mkdir(parents=True, exist_ok=True)
+    _draw(series, inputs.output, title=title)
+    return inputs.output
+
+
+def _draw(series: list[_RunSeries], output: Path, *, title: str | None) -> None:
+    """Concentrated matplotlib boundary."""
+    import matplotlib
+
+    matplotlib.use("Agg")
+    import matplotlib.pyplot as plt
+
+    fig, axes = plt.subplots(  # type: ignore[reportUnknownMemberType]
+        2, 1, figsize=(10, 8), sharex=True
+    )
+    top: Any = axes[0]  # type: ignore[reportAny]
+    bottom: Any = axes[1]  # type: ignore[reportAny]
+
+    for i, run in enumerate(series):
+        color = _COLOR_CYCLE[i % len(_COLOR_CYCLE)]
+        # Cumulative cold-prefill estimate (T_cum). Only plot points where
+        # we have a t_cum value — skip if derived was empty for this run.
+        n = min(len(run.cached_pp), len(run.cached_prefill_seconds))
+        if n > 0:
+            top.plot(  # type: ignore[reportAny, reportUnknownMemberType]
+                run.cached_pp[:n],
+                run.cached_prefill_seconds[:n],
+                "-o",
+                color=color,
+                label=run.label,
+            )
+        if run.control_pp:
+            top.scatter(  # type: ignore[reportAny, reportUnknownMemberType]
+                run.control_pp,
+                run.control_prefill_seconds,
+                marker="x",
+                s=80,
+                color=color,
+                label=f"{run.label} (cold one-shot)",
+            )
+        bottom.plot(  # type: ignore[reportAny, reportUnknownMemberType]
+            run.cached_pp,
+            run.cached_gen_tps,
+            "-o",
+            color=color,
+            label=run.label,
+        )
+
+    top.set_ylabel("prefill time (s)")  # type: ignore[reportAny, reportUnknownMemberType]
+    top.set_title(  # type: ignore[reportAny, reportUnknownMemberType]
+        "cumulative cold-prefill time vs context size "
+        "(line: T_cum estimate; ✕: cold one-shot control)"
+    )
+    top.grid(True, alpha=0.3)  # type: ignore[reportAny, reportUnknownMemberType]
+    top.legend(loc="best", fontsize=8)  # type: ignore[reportAny, reportUnknownMemberType]
+
+    bottom.set_xlabel("pp_tokens")  # type: ignore[reportAny, reportUnknownMemberType]
+    bottom.set_ylabel("generation_tps (tok/s)")  # type: ignore[reportAny, reportUnknownMemberType]
+    bottom.set_title("decode throughput vs context size")  # type: ignore[reportAny, reportUnknownMemberType]
+    bottom.grid(True, alpha=0.3)  # type: ignore[reportAny, reportUnknownMemberType]
+
+    if title is not None:
+        fig.suptitle(title, fontsize=10)  # type: ignore[reportUnknownMemberType]
+
+    fig.tight_layout()
+    fig.savefig(output, dpi=120, bbox_inches="tight")  # type: ignore[reportUnknownMemberType]
+    plt.close(fig)

bench/lib/prompt.py

@@ -0,0 +1,269 @@
+"""Typed prompt-sizing utilities for benchmarks.
+
+Wraps the HuggingFace ``transformers`` tokenizer (a fundamentally dynamic
+object — different models return different types from
+``apply_chat_template``) behind a small typed API so the rest of the bench
+library can stay strict-typed.
+
+``PromptSizer.build(target)`` returns a ``(content, exact_token_count)``
+pair. Internally it:
+  1. Tokenises the empty user message to learn the chat-template overhead
+     (``base_tokens``).
+  2. Estimates tokens-per-atom from a 100-atom sample.
+  3. Binary-searches over the atom count so the resulting message
+     tokenises to *exactly* ``target`` tokens.
+
+Callers downstream (``run_one_completion`` etc.) receive the verified
+token count, so analysis can confirm the prompt hit its target.
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import json
+import sys
+import types
+from collections.abc import Callable
+from pathlib import Path
+from typing import Any, Final, cast
+
+
+def _coerce_token_ids(raw: object) -> list[int]:
+    """Normalise ``apply_chat_template`` output to a flat list of token ids.
+
+    transformers' ``apply_chat_template`` may return:
+      - ``list[int]`` (slow tokenizers, ``tokenize=True``)
+      - a ``BatchEncoding`` with ``.input_ids`` (fast tokenizers)
+      - a tensor wrapped object (some models)
+
+    We only need ``len(.)`` of the result, so we just need to flatten to a
+    list and return it.
+    """
+    if isinstance(raw, list):
+        return cast("list[int]", raw)
+    input_ids = getattr(raw, "input_ids", None)
+    if isinstance(input_ids, list):
+        return cast("list[int]", input_ids)
+    raise TypeError(
+        f"Unsupported tokenizer output type {type(raw).__name__}; "
+        "expected list[int] or BatchEncoding-like with .input_ids."
+    )
+
+
+def _build_token_counter(tokenizer: object) -> Callable[[str], int]:
+    """Return a closure that counts tokens for a user message.
+
+    Tries ``apply_chat_template`` first; falls back to the DeepSeek-V4
+    Python encoder for models that don't ship a Jinja chat template.
+    """
+    apply_chat_template = cast(
+        Callable[..., object],
+        tokenizer.apply_chat_template,  # type: ignore[reportAttributeAccessIssue, reportUnknownMemberType]
+    )
+    encode = cast(
+        Callable[..., list[int]],
+        tokenizer.encode,  # type: ignore[reportAttributeAccessIssue, reportUnknownMemberType]
+    )
+
+    def count_fn(user_content: str) -> int:
+        messages = [{"role": "user", "content": user_content}]
+        try:
+            raw = apply_chat_template(
+                messages, tokenize=True, add_generation_prompt=True
+            )
+        except ValueError:
+            # Models without a Jinja chat template (e.g. DeepSeek V4 which
+            # ships its own Python encoder). Use the exo-side V4 encoder.
+            from exo.worker.engines.mlx.vendor.deepseek_v4_encoding import (  # type: ignore[reportMissingTypeStubs]
+                encode_messages as encode_v4,
+            )
+
+            prompt = cast(str, encode_v4(messages, thinking_mode="thinking"))  # type: ignore[reportUnknownArgumentType]
+            raw = encode(prompt, add_special_tokens=False)
+        return len(_coerce_token_ids(raw))
+
+    return count_fn
+
+
+class PromptSizer:
+    """Build a chat-completion content string of an exact token length."""
+
+    DEFAULT_ATOM: Final[str] = "a "
+
+    def __init__(self, tokenizer: object, atom: str = DEFAULT_ATOM):
+        self._tokenizer = tokenizer
+        self.atom = atom
+        self._count_fn = _build_token_counter(tokenizer)
+        self.base_tokens = self._count_fn("")
+
+    def count(self, content: str) -> int:
+        """Return the token count for ``content`` after chat-template expansion."""
+        return self._count_fn(content)
+
+    def build(self, target_prompt_tokens: int) -> tuple[str, int]:
+        """Return ``(content, exact_token_count)`` summing to ``target``.
+
+        Raises ``RuntimeError`` if the chosen ``atom`` overshoots the target
+        (try a different atom — see ``DEFAULT_ATOM``).
+        """
+        target = int(target_prompt_tokens)
+        if target < self.base_tokens:
+            raise RuntimeError(
+                f"Target ({target}) is smaller than template overhead "
+                f"({self.base_tokens})."
+            )
+
+        # Estimate tokens per atom using a sample.
+        sample_count = 100
+        sample_tokens = self._count_fn(self.atom * sample_count) - self.base_tokens
+        tokens_per_atom = sample_tokens / sample_count
+        needed_tokens = target - self.base_tokens
+        estimated_atoms = int(needed_tokens / tokens_per_atom)
+
+        # Binary search to find exact atom count.
+        low, high = 0, estimated_atoms * 2 + 100
+        while low < high:
+            mid = (low + high) // 2
+            if self._count_fn(self.atom * mid) < target:
+                low = mid + 1
+            else:
+                high = mid
+
+        content = self.atom * low
+        actual = self._count_fn(content)
+        if actual != target:
+            raise RuntimeError(
+                f"Overshot: got {actual} tokens (target {target}). "
+                f"Pick a different atom (try ' a' or '\\n' or '0 ')."
+            )
+        return content, actual
+
+
+def _load_kimi_tokenizer(model_id: str) -> object:
+    """Special-case Kimi K2's custom TikTokenTokenizer (transformers 5.x quirk)."""
+    from huggingface_hub import (
+        snapshot_download,  # type: ignore[reportUnknownVariableType]
+    )
+
+    raw_path = snapshot_download(
+        model_id,
+        allow_patterns=[
+            "*.json",
+            "*.py",
+            "*.tiktoken",
+            "*.model",
+            "*.jinja",
+        ],
+        dry_run=False,
+    )
+    model_path = Path(raw_path)
+    sys.path.insert(0, str(model_path))
+
+    tool_decl_path = model_path / "tool_declaration_ts.py"
+    if tool_decl_path.exists():
+        spec = importlib.util.spec_from_file_location(
+            "tool_declaration_ts", tool_decl_path
+        )
+        if spec is not None and spec.loader is not None:
+            tool_decl_module = importlib.util.module_from_spec(spec)
+            sys.modules["tool_declaration_ts"] = tool_decl_module
+            spec.loader.exec_module(tool_decl_module)
+
+    tok_path = model_path / "tokenization_kimi.py"
+    source = tok_path.read_text().replace(
+        "from .tool_declaration_ts", "from tool_declaration_ts"
+    )
+    tok_module = types.ModuleType("tokenization_kimi")
+    tok_module.__file__ = str(tok_path)
+    sys.modules["tokenization_kimi"] = tok_module
+    exec(compile(source, str(tok_path), "exec"), tok_module.__dict__)  # noqa: S102
+
+    tik_token_cls = cast(Any, tok_module).TikTokenTokenizer  # type: ignore[reportAny]
+    hf_tokenizer = cast(Any, tik_token_cls.from_pretrained(model_path))  # type: ignore[reportAny]
+
+    # Patch encode to use internal tiktoken model directly (transformers 5.x
+    # bug in the encode→pad path for slow tokenizers).
+    def _patched_encode(text: str, **_kwargs: object) -> list[int]:
+        return list(
+            hf_tokenizer.model.encode(text, allowed_special="all")  # type: ignore[reportAny, reportUnknownMemberType]
+        )
+
+    hf_tokenizer.encode = _patched_encode
+    return cast(object, hf_tokenizer)
+
+
+def load_tokenizer_for_bench(model_id: str) -> object:
+    """Load a HuggingFace tokenizer with bench-specific compatibility shims.
+
+    Returns the tokenizer as ``object`` because transformers' types are
+    fundamentally dynamic (concrete class depends on the model). Callers
+    should pass the result straight to :class:`PromptSizer`.
+    """
+    # Monkey-patch for transformers 5.x: Kimi's tokenization_kimi.py imports
+    # bytes_to_unicode from gpt2_tokenization which moved.
+    try:
+        import transformers.models.gpt2.tokenization_gpt2 as gpt2_tokenization
+        from transformers.convert_slow_tokenizer import bytes_to_unicode
+
+        if not hasattr(gpt2_tokenization, "bytes_to_unicode"):
+            gpt2_tokenization.bytes_to_unicode = bytes_to_unicode  # type: ignore[reportAttributeAccessIssue]
+    except ImportError:
+        pass
+
+    if "kimi-k2" in model_id.lower():
+        return _load_kimi_tokenizer(model_id)
+
+    from transformers import AutoTokenizer
+
+    try:
+        return cast(
+            object,
+            AutoTokenizer.from_pretrained(model_id, trust_remote_code=True),  # type: ignore[reportUnknownMemberType]
+        )
+    except (AttributeError, ValueError):
+        # Some models ship a Jinja template / encoder that AutoTokenizer
+        # can't introspect from HF directly — download artefacts and load
+        # from the local snapshot path.
+        from huggingface_hub import (
+            snapshot_download,  # type: ignore[reportUnknownVariableType]
+        )
+        from transformers import PretrainedConfig
+
+        raw_full_path = snapshot_download(
+            model_id,
+            allow_patterns=[
+                "*.json",
+                "*.py",
+                "tokenizer.model",
+                "*.tiktoken",
+                "tiktoken.model",
+                "*.txt",
+                "*.jsonl",
+                "*.jinja",
+            ],
+            dry_run=False,
+        )
+        model_path = Path(raw_full_path)
+        stub_kwargs: dict[str, Any] = {}
+        config_file = model_path / "config.json"
+        if config_file.exists():
+            with config_file.open() as f:
+                raw_config: dict[str, Any] = json.load(f)  # type: ignore[reportAny]
+            for key in (
+                "model_type",
+                "max_position_embeddings",
+                "vocab_size",
+                "bos_token_id",
+                "eos_token_id",
+                "pad_token_id",
+            ):
+                if key in raw_config:
+                    stub_kwargs[key] = raw_config[key]
+        return cast(
+            object,
+            AutoTokenizer.from_pretrained(  # type: ignore[reportUnknownMemberType]
+                str(model_path),
+                config=PretrainedConfig(**stub_kwargs),  # type: ignore[reportArgumentType, reportAny]
+                trust_remote_code=True,
+            ),
+        )

bench/lib/results.py

@@ -0,0 +1,139 @@
+"""Structured benchmark results — metadata capture + JSON output.
+
+Every benchmark run produces a single JSON file with a stable schema:
+
+- ``metadata``: exo SHA, ISO timestamps, hostnames, and any user-supplied
+  tags identifying the run.
+- ``cluster``: snapshot from the API (node identities, topology, memory).
+- ``params``: the benchmark's input parameters (sweep config, etc).
+- ``runs``: per-request result rows.
+- ``derived``: any computed summaries (``t_cum_seconds`` for context scaling).
+
+The format is intentionally additive so downstream tooling (plot scripts,
+dashboards) can rely on optional fields being absent rather than malformed.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import platform
+import socket
+import subprocess
+from dataclasses import asdict, dataclass, field
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any
+
+from exo_tools.client import ExoClient
+from exo_tools.harness import capture_cluster_snapshot
+
+
+def _git_describe(repo_root: Path) -> str | None:
+    """Return ``<short-sha>[-dirty]`` for the repo at ``repo_root`` or None."""
+    try:
+        sha = subprocess.run(
+            ["git", "rev-parse", "--short=12", "HEAD"],
+            cwd=str(repo_root),
+            capture_output=True,
+            text=True,
+            timeout=5,
+            check=True,
+        ).stdout.strip()
+    except (
+        subprocess.CalledProcessError,
+        FileNotFoundError,
+        subprocess.TimeoutExpired,
+    ):
+        return None
+    try:
+        dirty = subprocess.run(
+            ["git", "status", "--porcelain"],
+            cwd=str(repo_root),
+            capture_output=True,
+            text=True,
+            timeout=5,
+            check=True,
+        ).stdout.strip()
+        return f"{sha}-dirty" if dirty else sha
+    except (
+        subprocess.CalledProcessError,
+        FileNotFoundError,
+        subprocess.TimeoutExpired,
+    ):
+        return sha
+
+
+@dataclass
+class RunMetadata:
+    """Identifies a single bench run."""
+
+    run_id: str
+    benchmark: str
+    started_at: str
+    finished_at: str | None = None
+    exo_sha: str | None = None
+    hostname: str = ""
+    platform: str = ""
+    tags: dict[str, str] = field(default_factory=dict)
+
+    @classmethod
+    def new(
+        cls,
+        benchmark: str,
+        repo_root: Path,
+        *,
+        tags: dict[str, str] | None = None,
+    ) -> RunMetadata:
+        now = datetime.now(timezone.utc)
+        run_id = f"{benchmark}_{now.strftime('%Y%m%dT%H%M%SZ')}_{os.getpid()}"
+        return cls(
+            run_id=run_id,
+            benchmark=benchmark,
+            started_at=now.isoformat(),
+            exo_sha=_git_describe(repo_root),
+            hostname=socket.gethostname(),
+            platform=f"{platform.system()} {platform.release()} ({platform.machine()})",
+            tags=dict(tags or {}),
+        )
+
+
+@dataclass
+class ResultsBundle:
+    """Container for a single benchmark's results, before being written."""
+
+    metadata: RunMetadata
+    params: dict[str, Any] = field(default_factory=dict)
+    cluster: dict[str, Any] = field(default_factory=dict)
+    runs: list[dict[str, Any]] = field(default_factory=list)
+    cold_controls: list[dict[str, Any]] = field(default_factory=list)
+    derived: dict[str, Any] = field(default_factory=dict)
+
+    def capture_cluster(self, client: ExoClient) -> None:
+        """Snapshot the cluster state into ``self.cluster``."""
+        try:
+            snapshot = capture_cluster_snapshot(client)
+            if snapshot:
+                self.cluster.update(snapshot)
+        except Exception:
+            # Non-fatal: a benchmark without cluster snapshot is still valid
+            pass
+
+    def write_json(self, output_dir: Path) -> Path:
+        """Write the bundle as ``<output_dir>/<run_id>.json`` and return the path."""
+        if self.metadata.finished_at is None:
+            self.metadata.finished_at = datetime.now(timezone.utc).isoformat()
+        output_dir.mkdir(parents=True, exist_ok=True)
+        path = output_dir / f"{self.metadata.run_id}.json"
+        with path.open("w", encoding="utf-8") as f:
+            json.dump(asdict(self), f, indent=2, ensure_ascii=False)
+        return path
+
+
+def find_repo_root(start: Path | None = None) -> Path:
+    """Walk upwards from ``start`` (or this file) until a ``.git`` dir is found."""
+    cur = (start or Path(__file__)).resolve()
+    for parent in (cur, *cur.parents):
+        if (parent / ".git").is_dir() or (parent / ".git").is_file():
+            return parent
+    raise RuntimeError(f"Could not locate repo root above {cur}")

bench/lib/session.py

@@ -0,0 +1,65 @@
+"""BenchSession — wires together cluster + client + instance + tokenizer.
+
+Holds the ``EcoSession``, a deployed ``ClusterInfo``, an ``ExoClient`` for
+the cluster's primary endpoint, and (for benchmarks that need exact-token
+prompts) a lazily-constructed :class:`PromptSizer`.
+
+Benchmarks consume this via :func:`bench.lib.cluster.managed_instance`,
+which yields a populated ``BenchSession``. Library helpers (e.g.
+``context_scaling.run``) take a ``BenchSession`` and never reach for
+global state.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any, cast
+
+from exo_tools.client import ExoClient
+from exo_tools.cluster import ClusterInfo, EcoSession, make_client_from_url
+
+from .prompt import PromptSizer, load_tokenizer_for_bench
+
+
+@dataclass
+class BenchSession:
+    """Bundle of cluster + client + (optional) instance for benchmarks."""
+
+    cluster: ClusterInfo
+    eco: EcoSession
+    instance_id: str | None = None
+    model_id: str | None = None
+    full_model_id: str | None = None
+    _prompt_sizer: PromptSizer | None = field(default=None, repr=False)
+
+    @property
+    def client(self) -> ExoClient:
+        return make_client_from_url(self.cluster.api_url)
+
+    def state(self) -> dict[str, Any]:
+        raw: Any = self.client.request_json("GET", "/state")  # type: ignore[reportAny]
+        if isinstance(raw, dict):
+            return cast("dict[str, Any]", raw)
+        return {}
+
+    def instances(self) -> dict[str, Any]:
+        result: Any = self.state().get("instances", {})  # type: ignore[reportAny]
+        if isinstance(result, dict):
+            return cast("dict[str, Any]", result)
+        return {}
+
+    def get_prompt_sizer(self) -> PromptSizer:
+        """Return a cached :class:`PromptSizer` for ``self.full_model_id``.
+
+        Loaded lazily because tokenizer load is expensive and not every
+        benchmark needs prompt sizing.
+        """
+        if self._prompt_sizer is not None:
+            return self._prompt_sizer
+        if self.full_model_id is None:
+            raise RuntimeError(
+                "BenchSession.full_model_id is not set; cannot build a PromptSizer."
+            )
+        tokenizer = load_tokenizer_for_bench(self.full_model_id)
+        self._prompt_sizer = PromptSizer(tokenizer)
+        return self._prompt_sizer

bench/lib/tests/test_context_scaling.py

@@ -0,0 +1,186 @@
+"""Unit tests for the pure helpers in ``bench.lib.context_scaling``.
+
+The orchestration entry points (``run``, ``run_cached_sweep``,
+``run_cold_controls``, ``make_cold_control_factory``) need a real
+``BenchSession`` and exo cluster, so they're exercised end-to-end via
+``python -m bench.cli context-scaling``. This module covers the
+underscore-prefixed pure helpers via direct private-symbol access (the
+private prefix discourages library users; tests for those helpers are
+the explicit exception).
+"""
+
+from __future__ import annotations
+
+import math
+from typing import cast
+
+from bench.lib.context_scaling import (
+    StepResult,
+    _compute_t_cum,  # type: ignore[reportPrivateUsage]
+    _interp,  # type: ignore[reportPrivateUsage]
+    derive_summary,
+)
+
+
+def _step(
+    *,
+    pp: int,
+    delta: int,
+    prompt_tps: float,
+    generation_tps: float = 100.0,
+    hit: str = "partial",
+) -> StepResult:
+    return StepResult(
+        pp_tokens=pp,
+        delta_tokens=delta,
+        prompt_tps=prompt_tps,
+        generation_tps=generation_tps,
+        prefix_cache_hit=hit,
+        prompt_tokens=pp,
+        generation_tokens=32,
+        elapsed_s=delta / prompt_tps if prompt_tps else 0.0,
+    )
+
+
+def _close(actual: float, expected: float, abs_tol: float = 1e-3) -> bool:
+    return math.isclose(actual, expected, abs_tol=abs_tol)
+
+
+# ---------------------------------------------------------------------------
+# _compute_t_cum
+# ---------------------------------------------------------------------------
+
+
+class TestComputeTCum:
+    def test_empty_returns_empty(self) -> None:
+        assert _compute_t_cum([]) == []
+
+    def test_single_step(self) -> None:
+        # 256 tokens at 1024 tps -> 0.25s
+        out = _compute_t_cum([_step(pp=256, delta=256, prompt_tps=1024.0)])
+        assert len(out) == 1
+        assert _close(out[0], 0.25)
+
+    def test_cumulative_sum_across_three_steps(self) -> None:
+        steps = [
+            _step(pp=256, delta=256, prompt_tps=1000.0),  # 0.256s
+            _step(pp=512, delta=256, prompt_tps=2000.0),  # +0.128s = 0.384s
+            _step(pp=768, delta=256, prompt_tps=512.0),  # +0.500s = 0.884s
+        ]
+        out = _compute_t_cum(steps)
+        assert _close(out[0], 0.256)
+        assert _close(out[1], 0.384)
+        assert _close(out[2], 0.884)
+        # Monotonically non-decreasing
+        assert out == sorted(out)
+
+    def test_zero_tps_step_skipped(self) -> None:
+        # A row with prompt_tps == 0 contributes nothing to the cumulative sum
+        steps = [
+            _step(pp=256, delta=256, prompt_tps=1024.0),  # +0.25s
+            _step(pp=512, delta=256, prompt_tps=0.0),  # +0
+            _step(pp=768, delta=256, prompt_tps=512.0),  # +0.5s
+        ]
+        out = _compute_t_cum(steps)
+        assert _close(out[0], 0.25)
+        assert _close(out[1], 0.25)  # unchanged
+        assert _close(out[2], 0.75)
+
+    def test_zero_delta_step_skipped(self) -> None:
+        # Defensive: a Δ=0 row would otherwise add zero anyway, but we
+        # explicitly guard against negative delta + 0/0.
+        steps = [
+            _step(pp=256, delta=256, prompt_tps=1000.0),
+            _step(pp=256, delta=0, prompt_tps=1000.0),  # explicit Δ=0
+        ]
+        out = _compute_t_cum(steps)
+        assert _close(out[0], 0.256)
+        assert _close(out[1], 0.256)
+
+
+# ---------------------------------------------------------------------------
+# _interp
+# ---------------------------------------------------------------------------
+
+
+class TestInterp:
+    def test_empty_points_returns_zero(self) -> None:
+        assert _interp([], 100) == 0.0
+
+    def test_single_point_returns_y(self) -> None:
+        assert _interp([(100, 1.5)], 50) == 1.5
+        assert _interp([(100, 1.5)], 100) == 1.5
+        assert _interp([(100, 1.5)], 200) == 1.5
+
+    def test_clamps_below_first(self) -> None:
+        points = [(100, 0.1), (200, 0.3), (300, 0.6)]
+        assert _interp(points, 0) == 0.1
+        assert _interp(points, 50) == 0.1
+        assert _interp(points, 100) == 0.1
+
+    def test_clamps_above_last(self) -> None:
+        points = [(100, 0.1), (200, 0.3), (300, 0.6)]
+        assert _interp(points, 300) == 0.6
+        assert _interp(points, 500) == 0.6
+        assert _interp(points, 1_000_000) == 0.6
+
+    def test_mid_bracket_linear_interpolation(self) -> None:
+        points = [(100, 0.0), (200, 1.0)]
+        assert _close(_interp(points, 150), 0.5)
+        assert _close(_interp(points, 175), 0.75)
+
+    def test_multi_segment_linear_interpolation(self) -> None:
+        # Two adjacent segments, x=250 falls in the second one
+        points = [(100, 0.1), (200, 0.3), (300, 0.6)]
+        # 200..300: 0.3 + (0.6-0.3) * (250-200)/(300-200) = 0.3 + 0.15 = 0.45
+        assert _close(_interp(points, 250), 0.45)
+
+
+# ---------------------------------------------------------------------------
+# derive_summary
+# ---------------------------------------------------------------------------
+
+
+def _gap_at(summary: dict[str, object], index: int) -> dict[str, float]:
+    """Cast ``summary['control_gaps'][index]`` into the typed shape we expect."""
+    raw = summary["control_gaps"]
+    assert isinstance(raw, list)
+    entry = cast("dict[str, float]", raw[index])
+    return entry
+
+
+class TestDeriveSummary:
+    def test_no_controls_only_t_cum(self) -> None:
+        steps = [
+            _step(pp=256, delta=256, prompt_tps=1024.0),
+            _step(pp=512, delta=256, prompt_tps=1024.0),
+        ]
+        summary = derive_summary(steps, [])
+        t_cum = cast("list[float]", summary["t_cum_seconds"])
+        assert _close(t_cum[0], 0.25)
+        assert _close(t_cum[1], 0.5)
+        assert summary["control_gaps"] == []
+
+    def test_control_gap_at_known_pp(self) -> None:
+        # Sweep: 0.25s @ pp=256, 0.5s @ pp=512
+        steps = [
+            _step(pp=256, delta=256, prompt_tps=1024.0),
+            _step(pp=512, delta=256, prompt_tps=1024.0),
+        ]
+        # Cold control at pp=512, 2x faster than the per-step rate -> 0.25s
+        controls = [_step(pp=512, delta=512, prompt_tps=2048.0, hit="none")]
+        summary = derive_summary(steps, controls)
+        gap = _gap_at(summary, 0)
+        assert gap["pp_tokens"] == 512
+        assert _close(gap["cold_t_seconds"], 0.25, abs_tol=0.01)
+        assert _close(gap["t_cum_seconds_at_pp"], 0.5, abs_tol=0.01)
+        assert _close(gap["gap_seconds"], 0.25, abs_tol=0.01)
+        # gap_fraction = 0.25 / 0.25 = 1.0
+        assert _close(gap["gap_fraction"], 1.0, abs_tol=0.01)
+
+    def test_control_gap_zero_cold_tps_yields_zero_fraction(self) -> None:
+        steps = [_step(pp=256, delta=256, prompt_tps=1000.0)]
+        controls = [_step(pp=256, delta=256, prompt_tps=0.0, hit="none")]
+        gap = _gap_at(derive_summary(steps, controls), 0)
+        assert gap["cold_t_seconds"] == 0.0
+        assert gap["gap_fraction"] == 0.0

bench/lib/tests/test_model_meta.py

@@ -0,0 +1,171 @@
+"""Unit tests for ``bench.lib.model_meta``.
+
+These exercise the pure derivation helpers (no HF round-trip). The HTTP
+fetchers (``fetch_model_meta``, ``_read_config_json``, ``_sum_weight_sizes``)
+hit the public hub and aren't covered here.
+"""
+
+from __future__ import annotations
+
+import math
+
+import pytest
+
+from bench.lib.model_meta import (
+    ModelMeta,
+    derive_cold_controls,
+    derive_context_ramp,
+)
+
+
+def _meta(
+    *,
+    weight_bytes: int = 0,
+    max_pos: int = 4096,
+    layers: int = 32,
+) -> ModelMeta:
+    return ModelMeta(
+        model_id="test/model",
+        total_weight_bytes=weight_bytes,
+        max_position_embeddings=max_pos,
+        num_hidden_layers=layers,
+    )
+
+
+# ---------------------------------------------------------------------------
+# ModelMeta properties
+# ---------------------------------------------------------------------------
+
+
+class TestModelMetaConstraints:
+    def test_zero_weight_yields_one_gib_floor(self) -> None:
+        meta = _meta(weight_bytes=0)
+        # int(0 * 1.30) + 1 == 1; int(0 * 1.10) + 1 == 1
+        assert meta.memory_constraint_gb == 1.0
+        assert meta.disk_constraint_gb == 1.0
+
+    def test_one_gib_weight_rounds_up(self) -> None:
+        meta = _meta(weight_bytes=1 * (1024**3))
+        # int(1.0 * 1.30) + 1 = 2; int(1.0 * 1.10) + 1 = 2
+        assert meta.memory_constraint_gb == 2.0
+        assert meta.disk_constraint_gb == 2.0
+
+    def test_sixteen_gib_weight_uses_30pct_memory_10pct_disk(self) -> None:
+        meta = _meta(weight_bytes=16 * (1024**3))
+        # memory: int(16 * 1.30) + 1 = 21; disk: int(16 * 1.10) + 1 = 18
+        assert meta.memory_constraint_gb == 21.0
+        assert meta.disk_constraint_gb == 18.0
+
+    def test_total_weight_gb_property(self) -> None:
+        meta = _meta(weight_bytes=2_147_483_648)  # 2 GiB exactly
+        assert math.isclose(meta.total_weight_gb, 2.0)
+
+
+# ---------------------------------------------------------------------------
+# derive_context_ramp
+# ---------------------------------------------------------------------------
+
+
+class TestDeriveContextRamp:
+    def test_full_max_evenly_divides_round_to(self) -> None:
+        meta = _meta(max_pos=131072)  # 128k
+        pp_step, num_steps = derive_context_ramp(meta, num_steps=32)
+        # 131072 // 32 = 4096; rounded down to multiple of 256 = 4096
+        assert pp_step == 4096
+        assert num_steps == 32
+        # Top of ramp == max
+        assert pp_step * num_steps == 131072
+
+    def test_qwen30b_a3b_ramp(self) -> None:
+        meta = _meta(max_pos=40960)  # Qwen3-30B-A3B
+        pp_step, num_steps = derive_context_ramp(meta, num_steps=32)
+        # 40960 // 32 = 1280; multiple of 256
+        assert pp_step == 1280
+        assert pp_step * num_steps == 40960
+
+    def test_fraction_of_max_half(self) -> None:
+        meta = _meta(max_pos=131072)
+        pp_step, num_steps = derive_context_ramp(meta, num_steps=8, fraction_of_max=0.5)
+        # half = 65536; 65536 // 8 = 8192
+        assert pp_step == 8192
+        assert num_steps == 8
+
+    def test_min_pp_step_floor(self) -> None:
+        meta = _meta(max_pos=512)
+        # 512 // 32 = 16, but min_pp_step=256 floors it; rounded to 256
+        pp_step, num_steps = derive_context_ramp(meta, num_steps=32)
+        assert pp_step == 256
+        assert num_steps == 32
+
+    def test_round_to_truncates_down(self) -> None:
+        meta = _meta(max_pos=10000)
+        pp_step, _ = derive_context_ramp(meta, num_steps=32, round_to=256)
+        # 10000 // 32 = 312; (312 // 256) * 256 = 256
+        assert pp_step == 256
+
+    def test_round_to_zero_step_falls_back_to_round_to(self) -> None:
+        # Pathological: huge round_to relative to per-step size
+        meta = _meta(max_pos=1024)
+        pp_step, _ = derive_context_ramp(meta, num_steps=8, round_to=1024)
+        # 1024 // 8 = 128, but min_pp_step=256 → 256; (256 // 1024) * 1024 = 0;
+        # `or round_to` rescues to 1024.
+        assert pp_step == 1024
+
+    def test_max_pos_zero_raises(self) -> None:
+        meta = _meta(max_pos=0)
+        with pytest.raises(ValueError, match="max_position_embeddings=0"):
+            _ = derive_context_ramp(meta, num_steps=32)
+
+    @pytest.mark.parametrize("fraction", [0.0, -0.1, 1.5, 2.0])
+    def test_fraction_outside_unit_interval_raises(self, fraction: float) -> None:
+        meta = _meta(max_pos=4096)
+        with pytest.raises(ValueError, match="fraction_of_max"):
+            _ = derive_context_ramp(meta, num_steps=4, fraction_of_max=fraction)
+
+    @pytest.mark.parametrize("steps", [0, -1, -100])
+    def test_num_steps_must_be_positive(self, steps: int) -> None:
+        meta = _meta(max_pos=4096)
+        with pytest.raises(ValueError, match="num_steps"):
+            _ = derive_context_ramp(meta, num_steps=steps)
+
+
+# ---------------------------------------------------------------------------
+# derive_cold_controls
+# ---------------------------------------------------------------------------
+
+
+class TestDeriveColdControls:
+    def test_count_zero_returns_empty_tuple(self) -> None:
+        meta = _meta()
+        assert derive_cold_controls(meta, pp_step=4096, num_steps=32, count=0) == ()
+
+    def test_count_one_returns_top_only(self) -> None:
+        meta = _meta()
+        assert derive_cold_controls(meta, pp_step=4096, num_steps=32, count=1) == (
+            131072,
+        )
+
+    def test_evenly_spaced_four(self) -> None:
+        meta = _meta()
+        out = derive_cold_controls(meta, pp_step=4096, num_steps=32, count=4)
+        # max_pp = 131072; (131072 * (i+1)) // 4 for i in {0,1,2,3}
+        # = {32768, 65536, 98304, 131072}
+        assert out == (32768, 65536, 98304, 131072)
+
+    def test_filters_below_pp_step(self) -> None:
+        meta = _meta()
+        out = derive_cold_controls(meta, pp_step=8192, num_steps=2, count=4)
+        # max_pp = 16384; spaced points = {4096, 8192, 12288, 16384};
+        # 4096 < pp_step=8192 → dropped.
+        assert out == (8192, 12288, 16384)
+
+    def test_dedups_at_low_count_high_step(self) -> None:
+        meta = _meta()
+        # max_pp = 1024; count=2 → spaced = {512, 1024}; 512 < pp_step? No (=).
+        out = derive_cold_controls(meta, pp_step=512, num_steps=2, count=2)
+        assert out == (512, 1024)
+
+    def test_returned_in_ascending_order(self) -> None:
+        meta = _meta()
+        out = derive_cold_controls(meta, pp_step=1024, num_steps=8, count=4)
+        assert list(out) == sorted(out)

bench/lib/tests/test_plotting.py

@@ -0,0 +1,189 @@
+"""Smoke tests for ``bench.lib.plotting``.
+
+Renders a synthetic benchmark JSON to a tmp PNG and verifies the file is
+non-empty. We deliberately don't assert on pixel values — matplotlib
+output isn't byte-stable across versions — but a non-empty PNG with a
+valid header is a strong signal the renderer didn't throw.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Any, cast
+
+import pytest
+
+from bench.lib.plotting import PlotInputs, render_context_scaling
+
+
+def _write_synthetic_run(path: Path, *, run_id: str, model: str = "test/model") -> None:
+    """Write a minimal context-scaling-shaped JSON for plotting tests."""
+    payload = {
+        "metadata": {
+            "run_id": run_id,
+            "benchmark": "context_scaling",
+            "started_at": "2026-05-10T00:00:00Z",
+            "exo_sha": "deadbeef",
+            "hostname": "test-host",
+            "platform": "Linux 6.0 (x86_64)",
+            "tags": {"operator": "tester"},
+        },
+        "params": {
+            "pp_step": 256,
+            "num_steps": 4,
+            "tg": 32,
+            "warmup": 1,
+            "full_model_id": model,
+        },
+        "cluster": {},
+        "runs": [
+            {
+                "step_index": 0,
+                "phase": "cached_sweep",
+                "pp_tokens": 256,
+                "delta_tokens": 256,
+                "prompt_tps": 1800.0,
+                "generation_tps": 410.0,
+                "prefix_cache_hit": "exact",
+                "prompt_tokens": 256,
+                "generation_tokens": 32,
+                "elapsed_s": 0.14,
+                "peak_memory_bytes": 1_000_000_000,
+                "output_text_preview": "",
+            },
+            {
+                "step_index": 1,
+                "phase": "cached_sweep",
+                "pp_tokens": 512,
+                "delta_tokens": 256,
+                "prompt_tps": 2000.0,
+                "generation_tps": 395.0,
+                "prefix_cache_hit": "partial",
+                "prompt_tokens": 512,
+                "generation_tokens": 32,
+                "elapsed_s": 0.13,
+                "peak_memory_bytes": 1_100_000_000,
+                "output_text_preview": "",
+            },
+            {
+                "step_index": 2,
+                "phase": "cached_sweep",
+                "pp_tokens": 768,
+                "delta_tokens": 256,
+                "prompt_tps": 2200.0,
+                "generation_tps": 378.0,
+                "prefix_cache_hit": "partial",
+                "prompt_tokens": 768,
+                "generation_tokens": 32,
+                "elapsed_s": 0.12,
+                "peak_memory_bytes": 1_200_000_000,
+                "output_text_preview": "",
+            },
+        ],
+        "cold_controls": [
+            {
+                "phase": "cold_control",
+                "pp_tokens": 512,
+                "delta_tokens": 512,
+                "prompt_tps": 3200.0,
+                "generation_tps": 400.0,
+                "prefix_cache_hit": "none",
+                "prompt_tokens": 512,
+                "generation_tokens": 32,
+                "elapsed_s": 0.16,
+                "peak_memory_bytes": 1_500_000_000,
+                "output_text_preview": "",
+            },
+        ],
+        "derived": {
+            "t_cum_seconds": [0.14, 0.27, 0.39],
+            "control_gaps": [],
+        },
+    }
+    _ = path.write_text(json.dumps(payload))
+
+
+def _png_is_valid(path: Path) -> bool:
+    """A PNG file starts with the 8-byte magic ``\\x89PNG\\r\\n\\x1a\\n``."""
+    if not path.is_file():
+        return False
+    if path.stat().st_size < 100:
+        return False
+    head = path.read_bytes()[:8]
+    return head == b"\x89PNG\r\n\x1a\n"
+
+
+# ---------------------------------------------------------------------------
+
+
+class TestRenderContextScaling:
+    def test_single_run(self, tmp_path: Path) -> None:
+        json_path = tmp_path / "run.json"
+        _write_synthetic_run(json_path, run_id="r1")
+        out = tmp_path / "out.png"
+
+        returned = render_context_scaling(PlotInputs(results=[json_path], output=out))
+        assert returned == out
+        assert _png_is_valid(out)
+
+    def test_creates_output_parent_dir(self, tmp_path: Path) -> None:
+        json_path = tmp_path / "run.json"
+        _write_synthetic_run(json_path, run_id="r1")
+        out = tmp_path / "nested" / "deep" / "out.png"
+
+        _ = render_context_scaling(PlotInputs(results=[json_path], output=out))
+        assert _png_is_valid(out)
+
+    def test_comparison_two_runs(self, tmp_path: Path) -> None:
+        a = tmp_path / "a.json"
+        b = tmp_path / "b.json"
+        _write_synthetic_run(a, run_id="run-a", model="test/model-a")
+        _write_synthetic_run(b, run_id="run-b", model="test/model-b")
+        out = tmp_path / "compare.png"
+
+        _ = render_context_scaling(PlotInputs(results=[a, b], output=out))
+        assert _png_is_valid(out)
+
+    def test_label_tag_uses_metadata_tag(self, tmp_path: Path) -> None:
+        # Smoke test: just confirm passing label_tag doesn't throw and the
+        # PNG renders. Label content is too matplotlib-internal to inspect.
+        json_path = tmp_path / "run.json"
+        _write_synthetic_run(json_path, run_id="r1")
+        out = tmp_path / "out.png"
+
+        _ = render_context_scaling(
+            PlotInputs(results=[json_path], output=out, label_tag="operator")
+        )
+        assert _png_is_valid(out)
+
+    def test_explicit_title(self, tmp_path: Path) -> None:
+        json_path = tmp_path / "run.json"
+        _write_synthetic_run(json_path, run_id="r1")
+        out = tmp_path / "out.png"
+
+        _ = render_context_scaling(
+            PlotInputs(results=[json_path], output=out, title="Custom Title")
+        )
+        assert _png_is_valid(out)
+
+    def test_empty_results_raises(self, tmp_path: Path) -> None:
+        with pytest.raises(ValueError, match="at least one"):
+            _ = render_context_scaling(
+                PlotInputs(results=[], output=tmp_path / "out.png")
+            )
+
+    def test_wrong_benchmark_raises(self, tmp_path: Path) -> None:
+        # Same shape but with the wrong metadata.benchmark
+        json_path = tmp_path / "run.json"
+        _write_synthetic_run(json_path, run_id="r1")
+        raw_loaded: Any = json.loads(json_path.read_text())  # type: ignore[reportAny]
+        assert isinstance(raw_loaded, dict)
+        data = cast("dict[str, dict[str, str]]", raw_loaded)
+        data["metadata"]["benchmark"] = "something_else"
+        _ = json_path.write_text(json.dumps(data))
+
+        with pytest.raises(ValueError, match="context_scaling"):
+            _ = render_context_scaling(
+                PlotInputs(results=[json_path], output=tmp_path / "out.png")
+            )

bench/prefill_decode_bench.py

@@ -35,9 +35,8 @@ from exo_bench import (
     load_tokenizer_for_bench,
     parse_int_list,
 )
-from harness import (
-    ExoClient,
-    ExoHttpError,
+from exo_tools.client import ExoClient, ExoHttpError
+from exo_tools.harness import (
     add_common_instance_args,
     instance_id_from_instance,
     node_ids_from_instance,

bench/pyproject.toml

@@ -16,6 +16,7 @@ dependencies = [
   "lm-eval[api,math]>=0.4.0",
   "human-eval>=1.0.3",
   "numpy>=1.24.0",
+  "matplotlib>=3.8",
 ]
 
 [build-system]

dashboard/src/routes/+page.svelte

@@ -3435,6 +3435,7 @@
             >
               <li>Connect nodes with TB5 cables</li>
               <li>Boot to Recovery (hold power 10s → Options)</li>
+              <li>Open Terminal from the Utilities menu</li>
               <li>
                 Run
                 <code class="text-yellow-300 bg-yellow-400/10 px-1 rounded"
@@ -4822,6 +4823,7 @@
                 >
                   <li>Connect nodes with TB5 cables</li>
                   <li>Boot to Recovery (hold power 10s → Options)</li>
+                  <li>Open Terminal from the Utilities menu</li>
                   <li>
                     Run
                     <code class="text-yellow-300 bg-yellow-400/10 px-1 rounded"
@@ -4968,6 +4970,7 @@
                   >
                     <li>Connect nodes with TB5 cables</li>
                     <li>Boot to Recovery (hold power 10s → Options)</li>
+                    <li>Open Terminal from the Utilities menu</li>
                     <li>
                       Run
                       <code

pyproject.toml

@@ -40,6 +40,7 @@ exo = "exo.main:main"
 dev = [
   "basedpyright>=1.29.0",
   "pyinstaller>=6.17.0",
+  "playwright>=1.52.0",
   "pytest>=8.4.0",
   "pytest-asyncio>=1.0.0",
   "pytest-env",
@@ -75,7 +76,7 @@ cuda13 = [
 ###
 
 [tool.uv.workspace]
-members = ["rust/exo_pyo3_bindings", "bench"]
+members = ["rust/exo_pyo3_bindings", "bench", "tools"]
 
 [tool.uv.sources]
 exo-pyo3-bindings = { workspace = true }
@@ -112,7 +113,7 @@ build-backend = "uv_build"
 ###
 
 [tool.basedpyright]
-include = ["src", "bench"]
+include = ["src", "bench", "tools"]
 typeCheckingMode = "strict"
 failOnWarnings = true
 
@@ -146,6 +147,15 @@ reportMissingModuleSource = false
 [[tool.basedpyright.executionEnvironments]]
 root = "src"
 
+[[tool.basedpyright.executionEnvironments]]
+root = "bench"
+# `.` keeps `from bench.lib.X import …` resolvable (pytest adds the project
+# root to sys.path; we want type-checking to agree with runtime).
+extraPaths = ["tools/src", "."]
+
+[[tool.basedpyright.executionEnvironments]]
+root = "tools/src"
+
 
 ###
 # uv configuration
@@ -216,9 +226,12 @@ extend-exclude = [
 extend-select = ["I", "N", "B", "A", "PIE", "SIM"]
 
 [tool.pytest.ini_options]
-pythonpath = "."
+pythonpath = ["."]
 asyncio_mode = "auto"
 markers = ["slow: marks tests as slow (deselected by default)"]
 env = ["EXO_TESTS=1"]
-addopts = "-m 'not slow' --ignore=tests/start_distributed_test.py"
+# `tests/` requires an eco cluster (opt-in). `tmp/` holds throwaway scripts
+# that run a top-level `sys.exit(...)` at import time, which otherwise blows
+# up the default pytest collection.
+addopts = "-m 'not slow' --ignore=tests --ignore=tmp"
 filterwarnings = ["ignore:builtin type Swig:DeprecationWarning"]

rust/exo_pyo3_bindings/Cargo.toml

@@ -46,9 +46,12 @@ pyo3-async-runtimes = { version = "0.27.0", features = [
 ] }
 pyo3-log = "0.13.2"
 
+pidfile-rs = "0.3"
+
 # macro dependencies
 extend = { workspace = true }
 delegate = { workspace = true }
+thiserror = "2.0"
 
 # async runtime
 tokio = { workspace = true, features = ["full", "tracing"] }

rust/exo_pyo3_bindings/exo_pyo3_bindings.pyi

@@ -2,6 +2,8 @@
 # ruff: noqa: E501, F401
 
 import builtins
+import os
+import pathlib
 import typing
 
 @typing.final
@@ -69,6 +71,48 @@ class NoPeersSubscribedToTopicError(builtins.Exception):
     def __repr__(self) -> builtins.str: ...
     def __str__(self) -> builtins.str: ...
 
+@typing.final
+class Pidfile:
+    r"""
+    A PID file protected with a lock.
+    
+    An instance of `Pidfile` can be used to manage a PID file: create it,
+    lock it, detect already running daemons. It is backed by [`pidfile`][]
+    functions of `libbsd`/`libutil` which use `flopen` to lock the PID
+    file.
+    
+    When a PID file is created, the process ID of the current process is
+    *not* written there, making it possible to lock the PID file before
+    forking and only write the ID of the forked process when it is ready.
+    
+    The PID file is deleted automatically when the `Pidfile` comes out of
+    the scope. To close the PID file without deleting it, for example, in
+    the parent process of a forked daemon, call `close()`.
+    
+    [`exit`]: https://doc.rust-lang.org/std/process/fn.exit.html
+    [`pidfile`]: https://linux.die.net/man/3/pidfile
+    [`daemon`(3)]: https://linux.die.net/man/3/daemon
+    """
+    def __new__(cls, path: builtins.str | os.PathLike | pathlib.Path, mode: builtins.int) -> Pidfile:
+        r"""
+        Creates a new PID file and locks it.
+        
+        If the PID file cannot be locked, returns `PidfileError::AlreadyRunning` with
+        a PID of the already running process, or `None` if no PID has been written to
+        the PID file yet.
+        """
+    def write(self) -> None:
+        r"""
+        Writes the current process ID to the PID file.
+        
+        The file is truncated before writing.
+        """
+
+@typing.final
+class PidfileError(builtins.Exception):
+    def __repr__(self) -> builtins.str: ...
+    def __str__(self) -> builtins.str: ...
+
 class PyFromSwarm:
     @typing.final
     class Connection(PyFromSwarm):

rust/exo_pyo3_bindings/src/lib.rs

@@ -7,9 +7,11 @@
 mod allow_threading;
 mod ident;
 mod networking;
+mod pidfile;
 
 use crate::ident::PyKeypair;
 use crate::networking::networking_submodule;
+use crate::pidfile::pidfile_submodule;
 use pyo3::prelude::PyModule;
 use pyo3::types::PyModuleMethods;
 use pyo3::{Bound, PyResult, pyclass, pymodule};
@@ -164,6 +166,7 @@ fn main_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
     //       too many importing issues...
     m.add_class::<PyKeypair>()?;
     networking_submodule(m)?;
+    pidfile_submodule(m)?;
 
     // top-level constructs
     // TODO: ...

rust/exo_pyo3_bindings/src/pidfile.rs

@@ -0,0 +1,87 @@
+use pidfile_rs::{Pidfile, PidfileError};
+use pyo3::exceptions::PyException;
+use pyo3::prelude::{PyModule, PyModuleMethods};
+use pyo3::{Bound, PyErr, PyResult, Python, pyclass, pymethods};
+use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pymethods};
+use std::fs::Permissions;
+use std::os::unix::prelude::PermissionsExt;
+use std::path::PathBuf;
+
+#[gen_stub_pyclass]
+#[pyclass(frozen, extends=PyException, name="PidfileError")]
+pub struct PyPidfileError(PidfileError);
+
+impl PyPidfileError {
+    // TODO: I actually like this pattern a LOT more but how to abstract??
+    fn into_pyerr(self, py: Python) -> PyErr {
+        match Bound::new(py, self) {
+            Ok(err) => PyErr::from_value(err.into_any()),
+            Err(err) => err,
+        }
+    }
+}
+
+#[gen_stub_pymethods]
+#[pymethods]
+impl PyPidfileError {
+    fn __repr__(&self) -> String {
+        format!("PidfileError(\"{}\")", self.0)
+    }
+
+    fn __str__(&self) -> String {
+        self.0.to_string()
+    }
+}
+
+/// A PID file protected with a lock.
+///
+/// An instance of `Pidfile` can be used to manage a PID file: create it,
+/// lock it, detect already running daemons. It is backed by [`pidfile`][]
+/// functions of `libbsd`/`libutil` which use `flopen` to lock the PID
+/// file.
+///
+/// When a PID file is created, the process ID of the current process is
+/// *not* written there, making it possible to lock the PID file before
+/// forking and only write the ID of the forked process when it is ready.
+///
+/// The PID file is deleted automatically when the `Pidfile` comes out of
+/// the scope. To close the PID file without deleting it, for example, in
+/// the parent process of a forked daemon, call `close()`.
+///
+/// [`exit`]: https://doc.rust-lang.org/std/process/fn.exit.html
+/// [`pidfile`]: https://linux.die.net/man/3/pidfile
+/// [`daemon`(3)]: https://linux.die.net/man/3/daemon
+#[gen_stub_pyclass]
+#[pyclass(name = "Pidfile")]
+pub struct PyPidfile(Pidfile);
+
+#[gen_stub_pymethods]
+#[pymethods]
+impl PyPidfile {
+    /// Creates a new PID file and locks it.
+    ///
+    /// If the PID file cannot be locked, returns `PidfileError::AlreadyRunning` with
+    /// a PID of the already running process, or `None` if no PID has been written to
+    /// the PID file yet.
+    #[new]
+    fn py_new(py: Python, path: PathBuf, mode: u32) -> PyResult<Self> {
+        Ok(Self(
+            Pidfile::new(&path, Permissions::from_mode(mode))
+                .map_err(|e| PyPidfileError(e).into_pyerr(py))?,
+        ))
+    }
+
+    /// Writes the current process ID to the PID file.
+    ///
+    /// The file is truncated before writing.
+    fn write<'py>(&mut self, py: Python<'py>) -> PyResult<()> {
+        self.0.write().map_err(|e| PyPidfileError(e).into_pyerr(py))
+    }
+}
+
+pub fn pidfile_submodule(m: &Bound<PyModule>) -> PyResult<()> {
+    m.add_class::<PyPidfileError>()?;
+    m.add_class::<PyPidfile>()?;
+
+    Ok(())
+}

rust/exo_pyo3_bindings/tests/test_python.py

@@ -1,10 +1,12 @@
 import asyncio
 
 import pytest
+from _pytest.capture import CaptureFixture
 from exo_pyo3_bindings import (
     Keypair,
     NetworkingHandle,
     NoPeersSubscribedToTopicError,
+    Pidfile,
     PyFromSwarm,
 )
 
@@ -26,6 +28,13 @@ async def test_sleep_on_multiple_items() -> None:
             print("caught it", e)
 
 
+def test_pidfile(capsys: CaptureFixture[str]):
+    with capsys.disabled():
+        print("\nbefore python")
+        scoped_lock_file()
+        print("after python")
+
+
 async def _await_recv(h: NetworkingHandle):
     while True:
         event = await h.recv()
@@ -34,3 +43,7 @@ async def _await_recv(h: NetworkingHandle):
                 print(f"PYTHON: connection update: {c}")
             case PyFromSwarm.Message() as m:
                 print(f"PYTHON: message: {m}")
+
+
+def scoped_lock_file():
+    a = Pidfile("/tmp/lock.pid", 0o0600)

src/exo/api/main.py

@@ -133,12 +133,10 @@ from exo.shared.constants import (
 )
 from exo.shared.election import ElectionMessage
 from exo.shared.logging import InterceptLogger
+from exo.shared.models import model_cards
 from exo.shared.models.model_cards import (
     ModelCard,
     ModelId,
-    add_to_card_cache,
-    get_card,
-    get_model_cards,
 )
 from exo.shared.tracing import TraceEvent, compute_stats, export_trace, load_trace_file
 from exo.shared.types.chunks import (
@@ -481,6 +479,7 @@ class API:
                 topology=self.state.topology,
                 current_instances=self.state.instances,
                 download_status=self.state.downloads,
+                node_rdma_ctl=self.state.node_rdma_ctl,
             )
         except ValueError as exc:
             raise HTTPException(status_code=400, detail=str(exc)) from exc
@@ -544,6 +543,7 @@ class API:
                     current_instances=self.state.instances,
                     required_nodes=required_nodes,
                     download_status=self.state.downloads,
+                    node_rdma_ctl=self.state.node_rdma_ctl,
                 )
             except ValueError as exc:
                 if (model_card.model_id, sharding, instance_meta, 0) not in seen:
@@ -1633,17 +1633,16 @@ class API:
     async def ollama_tags(self) -> OllamaTagsResponse:
         """Returns list of models in Ollama tags format. We return the downloaded ones only."""
 
-        def none_if_empty(value: str) -> str | None:
-            return value or None
-
-        downloaded_model_ids: set[str] = set()
+        downloaded_model_ids: set[ModelId] = set()
         for node_downloads in self.state.downloads.values():
             for dl in node_downloads:
                 if isinstance(dl, DownloadCompleted):
                     downloaded_model_ids.add(dl.shard_metadata.model_card.model_id)
 
         cards = [
-            c for c in await get_model_cards() if c.model_id in downloaded_model_ids
+            c
+            for c in await model_cards.card_cache.list_all()
+            if c.model_id in downloaded_model_ids
         ]
 
         now = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
@@ -1656,8 +1655,8 @@ class API:
                     size=card.storage_size.in_bytes,
                     digest="sha256:000000000000",
                     details=OllamaModelDetails(
-                        family=none_if_empty(card.family),
-                        quantization_level=none_if_empty(card.quantization),
+                        family=card.family or None,
+                        quantization_level=card.quantization or None,
                     ),
                 )
                 for card in cards
@@ -1720,7 +1719,7 @@ class API:
 
     async def get_models(self, status: str | None = Query(default=None)) -> ModelList:
         """Returns list of available models, optionally filtered by being downloaded."""
-        cards = await get_model_cards()
+        cards = await model_cards.card_cache.list_all()
 
         if status == "downloaded":
             downloaded_model_ids: set[str] = set()
@@ -1771,7 +1770,7 @@ class API:
 
         # Immediately update the local cache so the subsequent GET /models
         # returns the new model without waiting for the event round-trip.
-        add_to_card_cache(card)
+        model_cards.card_cache.cc[card.model_id] = card
 
         return ModelListModel(
             id=card.model_id,
@@ -1787,7 +1786,7 @@ class API:
 
     async def delete_custom_model(self, model_id: ModelId) -> JSONResponse:
         """Delete a user-added custom model card and sync deletion across the cluster."""
-        card = get_card(model_id)
+        card = model_cards.card_cache.get(model_id)
         if card is None or not card.is_custom:
             raise HTTPException(status_code=404, detail="Custom model card not found")
 

src/exo/download/coordinator.py

@@ -16,7 +16,8 @@ from exo.download.download_utils import (
 )
 from exo.download.shard_downloader import ShardDownloader
 from exo.shared.constants import EXO_DEFAULT_MODELS_DIR, EXO_MODELS_READ_ONLY_DIRS
-from exo.shared.models.model_cards import ModelId, get_model_cards
+from exo.shared.models import model_cards
+from exo.shared.models.model_cards import ModelId
 from exo.shared.types.commands import (
     CancelDownload,
     DeleteDownload,
@@ -422,7 +423,7 @@ class DownloadCoordinator:
                     )
                 # Scan read-only directories for pre-downloaded models
                 if EXO_MODELS_READ_ONLY_DIRS:
-                    for card in await get_model_cards():
+                    for card in await model_cards.card_cache.list_all():
                         mid = card.model_id
                         if mid in self.active_downloads:
                             continue

src/exo/download/download_utils.py

@@ -1,11 +1,12 @@
 import asyncio
 import hashlib
 import os
+import random
 import shutil
 import ssl
 import time
 import traceback
-from collections.abc import Awaitable
+from collections.abc import Awaitable, Mapping
 from datetime import timedelta
 from pathlib import Path
 from typing import Callable, Literal
@@ -55,6 +56,36 @@ class HuggingFaceAuthenticationError(Exception):
 class HuggingFaceRateLimitError(Exception):
     """429 Huggingface code"""
 
+    def __init__(self, msg: str, retry_after: float | None = None) -> None:
+        super().__init__(msg)
+        self.retry_after = retry_after
+
+
+def _parse_retry_after(headers: Mapping[str, str]) -> float | None:
+    """Parse seconds-to-reset from HF's RateLimit header.
+
+    HF sends e.g. ``ratelimit: "api";r=0;t=52`` on 429s; ``t`` is the wait.
+    Returns ``None`` if the header is missing or has no ``t`` field.
+    """
+    raw = headers.get("RateLimit") or headers.get("ratelimit")
+    if raw is None:
+        return None
+    for part in raw.split(";"):
+        key, _, val = part.strip().partition("=")
+        if key == "t":
+            try:
+                return float(val)
+            except ValueError:
+                return None
+    return None
+
+
+# reset window is 5 min
+_RATE_LIMIT_MAX_SLEEP_SECS = 300.0
+
+# 24h. Manually clear the cache (or `delete_model`) to force a refresh.
+_FILE_LIST_CACHE_TTL_SECS = 24 * 60 * 60
+
 
 async def _build_auth_error_message(status_code: int, model_id: ModelId) -> str:
     token = await get_hf_token()
@@ -348,9 +379,6 @@ async def _build_file_list_from_local_directory(
     return None
 
 
-_fetched_file_lists_this_session: set[str] = set()
-
-
 async def fetch_file_list_with_cache(
     model_id: ModelId,
     revision: str = "main",
@@ -360,13 +388,16 @@ async def fetch_file_list_with_cache(
 ) -> list[FileListEntry]:
     target_dir = await ensure_cache_dir(model_id)
     cache_file = target_dir / f"{model_id.normalize()}--{revision}--file_list.json"
-    cache_key = f"{model_id.normalize()}--{revision}"
 
-    if cache_key in _fetched_file_lists_this_session and await aios.path.exists(
-        cache_file
-    ):
-        async with aiofiles.open(cache_file, "r") as f:
-            return TypeAdapter(list[FileListEntry]).validate_json(await f.read())
+    # cache survives process restarts so cold starts don't re-burst HF
+    if await aios.path.exists(cache_file):
+        try:
+            cache_age = time.time() - (await aios.stat(cache_file)).st_mtime
+        except OSError:
+            cache_age = float("inf")
+        if cache_age < _FILE_LIST_CACHE_TTL_SECS:
+            async with aiofiles.open(cache_file, "r") as f:
+                return TypeAdapter(list[FileListEntry]).validate_json(await f.read())
 
     if skip_internet:
         if await aios.path.exists(cache_file):
@@ -395,7 +426,6 @@ async def fetch_file_list_with_cache(
             await f.write(
                 TypeAdapter(list[FileListEntry]).dump_json(file_list).decode()
             )
-        _fetched_file_lists_this_session.add(cache_key)
         return file_list
     except Exception as e:
         logger.opt(exception=e).warning(
@@ -426,17 +456,29 @@ async def fetch_file_list_with_retry(
     recursive: bool = False,
     on_connection_lost: Callable[[], None] = lambda: None,
 ) -> list[FileListEntry]:
-    n_attempts = 3
+    n_attempts = 5
     for attempt in range(n_attempts):
         try:
             return await _fetch_file_list(model_id, revision, path, recursive)
         except HuggingFaceAuthenticationError:
             raise
+        except HuggingFaceRateLimitError as e:
+            if attempt == n_attempts - 1:
+                raise
+            sleep_for = e.retry_after if e.retry_after is not None else 2.0**attempt
+            sleep_for = min(sleep_for, _RATE_LIMIT_MAX_SLEEP_SECS) + random.uniform(
+                0, 1
+            )
+            logger.warning(
+                f"Rate limited by HuggingFace fetching file list for {model_id}; "
+                f"sleeping {sleep_for:.1f}s before retry {attempt + 2}/{n_attempts}"
+            )
+            await asyncio.sleep(sleep_for)
         except Exception as e:
             on_connection_lost()
             if attempt == n_attempts - 1:
                 raise e
-            await asyncio.sleep(2.0**attempt)
+            await asyncio.sleep(2.0**attempt + random.uniform(0, 1))
     raise Exception(
         f"Failed to fetch file list for {model_id=} {revision=} {path=} {recursive=}"
     )
@@ -447,6 +489,9 @@ async def _fetch_file_list(
 ) -> list[FileListEntry]:
     api_url = f"{get_hf_endpoint()}/api/models/{model_id}/tree/{revision}"
     url = f"{api_url}/{path}" if path else api_url
+    # ?recursive=true returns the whole subtree in one request
+    if recursive:
+        url = f"{url}?recursive=true"
 
     headers = await get_download_headers()
     async with (
@@ -458,7 +503,8 @@ async def _fetch_file_list(
             raise HuggingFaceAuthenticationError(msg)
         elif response.status == 429:
             raise HuggingFaceRateLimitError(
-                f"Couldn't download {model_id} because of HuggingFace rate limit."
+                f"HuggingFace rate limit hit fetching file list for {model_id}",
+                retry_after=_parse_retry_after(response.headers),
             )
         elif response.status == 200:
             data_json = await response.text()
@@ -468,10 +514,14 @@ async def _fetch_file_list(
                 if item.type == "file":
                     files.append(FileListEntry.model_validate(item))
                 elif item.type == "directory" and recursive:
-                    subfiles = await _fetch_file_list(
-                        model_id, revision, item.path, recursive
-                    )
-                    files.extend(subfiles)
+                    # already inlined by ?recursive=true
+                    continue
+            if recursive and len(data) >= 1000:
+                # HF tree endpoint paginates at 1000; we don't follow cursors
+                logger.warning(
+                    f"File list for {model_id} hit the 1000-entry page cap "
+                    "and may be truncated; cursor pagination is not implemented"
+                )
             return files
         else:
             raise Exception(f"Failed to fetch file list: {response.status}")
@@ -552,6 +602,11 @@ async def file_meta(
         if r.status in [401, 403]:
             msg = await _build_auth_error_message(r.status, model_id)
             raise HuggingFaceAuthenticationError(msg)
+        if r.status == 429:
+            raise HuggingFaceRateLimitError(
+                f"HuggingFace rate limit hit fetching metadata for {model_id}/{path}",
+                retry_after=_parse_retry_after(r.headers),
+            )
         content_length = int(
             r.headers.get("x-linked-size") or r.headers.get("content-length") or 0
         )
@@ -571,7 +626,7 @@ async def download_file_with_retry(
     on_connection_lost: Callable[[], None] = lambda: None,
     skip_internet: bool = False,
 ) -> Path:
-    n_attempts = 3
+    n_attempts = 5
     for attempt in range(n_attempts):
         try:
             return await _download_file(
@@ -583,12 +638,16 @@ async def download_file_with_retry(
             raise
         except HuggingFaceRateLimitError as e:
             if attempt == n_attempts - 1:
-                raise e
-            logger.error(
-                f"Download error on attempt {attempt}/{n_attempts} for {model_id=} {revision=} {path=} {target_dir=}"
+                raise
+            sleep_for = e.retry_after if e.retry_after is not None else 2.0**attempt
+            sleep_for = min(sleep_for, _RATE_LIMIT_MAX_SLEEP_SECS) + random.uniform(
+                0, 1
             )
-            logger.error(traceback.format_exc())
-            await asyncio.sleep(2.0**attempt)
+            logger.warning(
+                f"Rate limited by HuggingFace downloading {model_id}/{path}; "
+                f"sleeping {sleep_for:.1f}s before retry {attempt + 2}/{n_attempts}"
+            )
+            await asyncio.sleep(sleep_for)
         except Exception as e:
             if attempt == n_attempts - 1:
                 on_connection_lost()
@@ -597,7 +656,7 @@ async def download_file_with_retry(
                 f"Download error on attempt {attempt + 1}/{n_attempts} for {model_id=} {revision=} {path=} {target_dir=}"
             )
             logger.error(traceback.format_exc())
-            await asyncio.sleep(2.0**attempt)
+            await asyncio.sleep(2.0**attempt + random.uniform(0, 1))
     raise Exception(
         f"Failed to download file {model_id=} {revision=} {path=} {target_dir=}"
     )
@@ -665,6 +724,11 @@ async def _download_file(
             if r.status in [401, 403]:
                 msg = await _build_auth_error_message(r.status, model_id)
                 raise HuggingFaceAuthenticationError(msg)
+            if r.status == 429:
+                raise HuggingFaceRateLimitError(
+                    f"HuggingFace rate limit hit downloading {model_id}/{path}",
+                    retry_after=_parse_retry_after(r.headers),
+                )
             assert r.status in [200, 206], (
                 f"Failed to download {path} from {url}: {r.status}"
             )

src/exo/download/impl_shard_downloader.py

@@ -11,11 +11,11 @@ from exo.download.download_utils import (
     download_shard,
 )
 from exo.download.shard_downloader import ShardDownloader
+from exo.shared.models import model_cards
 from exo.shared.models.model_cards import (
     ModelCard,
     ModelId,
     ModelTask,
-    get_model_cards,
 )
 from exo.shared.types.memory import Memory
 from exo.shared.types.worker.shards import (
@@ -258,7 +258,7 @@ class ResumableShardDownloader(ShardDownloader):
 
         tasks = [
             create_task(download_with_semaphore(model_card))
-            for model_card in await get_model_cards()
+            for model_card in await model_cards.card_cache.list_all()
         ]
 
         for task in asyncio.as_completed(tasks):

src/exo/download/tests/test_offline_mode.py

@@ -1,5 +1,7 @@
 """Tests for offline/air-gapped mode."""
 
+import os
+import time
 from collections.abc import AsyncIterator
 from pathlib import Path
 from unittest.mock import AsyncMock, patch
@@ -231,3 +233,64 @@ class TestFetchFileListOffline:
         raise FileNotFoundError."""
         with pytest.raises(FileNotFoundError, match="No internet"):
             await fetch_file_list_with_cache(model_id, "main", skip_internet=True)
+
+
+class TestFileListCacheTTL:
+    async def test_uses_fresh_cache_without_fetching(
+        self, model_id: ModelId, temp_models_dir: Path
+    ) -> None:
+        from pydantic import TypeAdapter
+
+        cache_dir = temp_models_dir / "caches" / model_id.normalize()
+        await aios.makedirs(cache_dir, exist_ok=True)
+
+        cached_list = [
+            FileListEntry(type="file", path="model.safetensors", size=1000),
+        ]
+        cache_file = cache_dir / f"{model_id.normalize()}--main--file_list.json"
+        async with aiofiles.open(cache_file, "w") as f:
+            await f.write(
+                TypeAdapter(list[FileListEntry]).dump_json(cached_list).decode()
+            )
+
+        with patch(
+            "exo.download.download_utils.fetch_file_list_with_retry",
+            new_callable=AsyncMock,
+        ) as mock_fetch:
+            result = await fetch_file_list_with_cache(model_id, "main")
+
+        assert result == cached_list
+        mock_fetch.assert_not_called()
+
+    async def test_refetches_when_cache_older_than_ttl(
+        self, model_id: ModelId, temp_models_dir: Path
+    ) -> None:
+        from pydantic import TypeAdapter
+
+        from exo.download.download_utils import (
+            _FILE_LIST_CACHE_TTL_SECS,  # pyright: ignore[reportPrivateUsage]
+        )
+
+        cache_dir = temp_models_dir / "caches" / model_id.normalize()
+        await aios.makedirs(cache_dir, exist_ok=True)
+
+        stale_list = [FileListEntry(type="file", path="stale.bin", size=1)]
+        cache_file = cache_dir / f"{model_id.normalize()}--main--file_list.json"
+        async with aiofiles.open(cache_file, "w") as f:
+            await f.write(
+                TypeAdapter(list[FileListEntry]).dump_json(stale_list).decode()
+            )
+
+        old_mtime = time.time() - _FILE_LIST_CACHE_TTL_SECS - 60
+        os.utime(cache_file, (old_mtime, old_mtime))
+
+        fresh_list = [FileListEntry(type="file", path="fresh.bin", size=2)]
+        with patch(
+            "exo.download.download_utils.fetch_file_list_with_retry",
+            new_callable=AsyncMock,
+            return_value=fresh_list,
+        ) as mock_fetch:
+            result = await fetch_file_list_with_cache(model_id, "main")
+
+        assert result == fresh_list
+        mock_fetch.assert_called_once()

src/exo/download/tests/test_rate_limit_handling.py

@@ -0,0 +1,355 @@
+"""Tests for HuggingFace 429 rate-limit handling in download_utils."""
+
+from collections.abc import AsyncIterator
+from pathlib import Path
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import aiofiles.os as aios
+import pytest
+
+from exo.download.download_utils import (
+    HuggingFaceRateLimitError,
+    _download_file,  # pyright: ignore[reportPrivateUsage]
+    _fetch_file_list,  # pyright: ignore[reportPrivateUsage]
+    _parse_retry_after,  # pyright: ignore[reportPrivateUsage]
+    download_file_with_retry,
+    fetch_file_list_with_retry,
+    file_meta,
+)
+from exo.shared.types.common import ModelId
+
+# captured from a real HF 429 on 2026-04-30 (header is lowercased by Cloudfront)
+REAL_HF_429_HEADERS_2026_04_30 = {
+    "ratelimit": '"api";r=0;t=52',
+    "ratelimit-policy": '"fixed window";"api";q=500;w=300',
+}
+
+
+class TestParseRetryAfter:
+    def test_parses_documented_format(self) -> None:
+        assert _parse_retry_after({"RateLimit": '"api";r=0;t=243'}) == 243.0
+
+    def test_parses_real_hf_response(self) -> None:
+        assert _parse_retry_after(REAL_HF_429_HEADERS_2026_04_30) == 52.0
+
+    def test_parses_resolvers_bucket(self) -> None:
+        assert _parse_retry_after({"ratelimit": '"resolvers";r=0;t=120'}) == 120.0
+
+    def test_parses_pages_bucket(self) -> None:
+        assert _parse_retry_after({"ratelimit": '"pages";r=0;t=10'}) == 10.0
+
+    def test_returns_none_when_header_missing(self) -> None:
+        assert _parse_retry_after({}) is None
+
+    def test_returns_none_when_only_retry_after_present(self) -> None:
+        assert _parse_retry_after({"Retry-After": "60"}) is None
+
+    def test_returns_none_when_format_unrecognised(self) -> None:
+        assert _parse_retry_after({"ratelimit": "garbage"}) is None
+
+    def test_handles_extra_whitespace(self) -> None:
+        assert _parse_retry_after({"ratelimit": '"api"; r=0; t=42'}) == 42.0
+
+
+class TestFetchFileListRetry:
+    async def test_uses_retry_after_from_error(self) -> None:
+        sleeps: list[float] = []
+
+        async def fake_sleep(seconds: float) -> None:
+            sleeps.append(seconds)
+
+        async def fake_fetch(*args: object, **kwargs: object) -> list[object]:
+            if not sleeps:
+                raise HuggingFaceRateLimitError("rate limited", retry_after=2.0)
+            return []
+
+        with (
+            patch(
+                "exo.download.download_utils._fetch_file_list", side_effect=fake_fetch
+            ),
+            patch("exo.download.download_utils.asyncio.sleep", side_effect=fake_sleep),
+        ):
+            result = await fetch_file_list_with_retry(ModelId("test/model"))
+
+        assert result == []
+        assert len(sleeps) == 1
+        assert 2.0 <= sleeps[0] < 3.0  # retry_after + jitter[0,1)
+
+    async def test_falls_back_to_exp_backoff_when_no_retry_after(self) -> None:
+        sleeps: list[float] = []
+
+        async def fake_sleep(seconds: float) -> None:
+            sleeps.append(seconds)
+
+        async def fake_fetch(*args: object, **kwargs: object) -> list[object]:
+            if not sleeps:
+                raise HuggingFaceRateLimitError("rate limited", retry_after=None)
+            return []
+
+        with (
+            patch(
+                "exo.download.download_utils._fetch_file_list", side_effect=fake_fetch
+            ),
+            patch("exo.download.download_utils.asyncio.sleep", side_effect=fake_sleep),
+        ):
+            await fetch_file_list_with_retry(ModelId("test/model"))
+
+        assert len(sleeps) == 1
+        assert 1.0 <= sleeps[0] < 2.0  # 2**0 + jitter[0,1)
+
+    async def test_caps_sleep_at_max_window(self) -> None:
+        sleeps: list[float] = []
+
+        async def fake_sleep(seconds: float) -> None:
+            sleeps.append(seconds)
+
+        async def fake_fetch(*args: object, **kwargs: object) -> list[object]:
+            if not sleeps:
+                raise HuggingFaceRateLimitError("rate limited", retry_after=10_000.0)
+            return []
+
+        with (
+            patch(
+                "exo.download.download_utils._fetch_file_list", side_effect=fake_fetch
+            ),
+            patch("exo.download.download_utils.asyncio.sleep", side_effect=fake_sleep),
+        ):
+            await fetch_file_list_with_retry(ModelId("test/model"))
+
+        assert len(sleeps) == 1
+        assert 300.0 <= sleeps[0] < 301.0  # cap + jitter[0,1)
+
+    async def test_retries_up_to_five_times(self) -> None:
+        sleeps: list[float] = []
+
+        async def fake_sleep(seconds: float) -> None:
+            sleeps.append(seconds)
+
+        async def fake_fetch(*args: object, **kwargs: object) -> list[object]:
+            raise HuggingFaceRateLimitError("rate limited", retry_after=1.0)
+
+        with (
+            patch(
+                "exo.download.download_utils._fetch_file_list", side_effect=fake_fetch
+            ),
+            patch("exo.download.download_utils.asyncio.sleep", side_effect=fake_sleep),
+            pytest.raises(HuggingFaceRateLimitError),
+        ):
+            await fetch_file_list_with_retry(ModelId("test/model"))
+
+        assert len(sleeps) == 4  # 5 attempts -> 4 sleeps before giving up
+
+
+class TestDownloadFileRetry:
+    @pytest.fixture
+    async def target_dir(self, tmp_path: Path) -> AsyncIterator[Path]:
+        target = tmp_path / "downloads"
+        await aios.makedirs(target, exist_ok=True)
+        yield target
+
+    async def test_uses_retry_after_from_error(self, target_dir: Path) -> None:
+        sleeps: list[float] = []
+        results: list[Path] = [target_dir / "file.bin"]
+
+        async def fake_sleep(seconds: float) -> None:
+            sleeps.append(seconds)
+
+        async def fake_download(*args: object, **kwargs: object) -> Path:
+            if not sleeps:
+                raise HuggingFaceRateLimitError("rate limited", retry_after=5.0)
+            return results[0]
+
+        with (
+            patch(
+                "exo.download.download_utils._download_file",
+                side_effect=fake_download,
+            ),
+            patch("exo.download.download_utils.asyncio.sleep", side_effect=fake_sleep),
+        ):
+            result = await download_file_with_retry(
+                ModelId("test/model"), "main", "file.bin", target_dir
+            )
+
+        assert result == results[0]
+        assert len(sleeps) == 1
+        assert 5.0 <= sleeps[0] < 6.0
+
+    async def test_caps_sleep_at_max_window(self, target_dir: Path) -> None:
+        sleeps: list[float] = []
+        results: list[Path] = [target_dir / "file.bin"]
+
+        async def fake_sleep(seconds: float) -> None:
+            sleeps.append(seconds)
+
+        async def fake_download(*args: object, **kwargs: object) -> Path:
+            if not sleeps:
+                raise HuggingFaceRateLimitError("rate limited", retry_after=99_999.0)
+            return results[0]
+
+        with (
+            patch(
+                "exo.download.download_utils._download_file",
+                side_effect=fake_download,
+            ),
+            patch("exo.download.download_utils.asyncio.sleep", side_effect=fake_sleep),
+        ):
+            await download_file_with_retry(
+                ModelId("test/model"), "main", "file.bin", target_dir
+            )
+
+        assert len(sleeps) == 1
+        assert 300.0 <= sleeps[0] < 301.0
+
+    async def test_retries_up_to_five_times(self, target_dir: Path) -> None:
+        sleeps: list[float] = []
+
+        async def fake_sleep(seconds: float) -> None:
+            sleeps.append(seconds)
+
+        with (
+            patch(
+                "exo.download.download_utils._download_file",
+                new_callable=AsyncMock,
+                side_effect=HuggingFaceRateLimitError("rate limited", retry_after=1.0),
+            ),
+            patch("exo.download.download_utils.asyncio.sleep", side_effect=fake_sleep),
+            pytest.raises(HuggingFaceRateLimitError),
+        ):
+            await download_file_with_retry(
+                ModelId("test/model"), "main", "file.bin", target_dir
+            )
+
+        assert len(sleeps) == 4
+
+
+def _make_mock_session_returning(
+    response_attrs: dict[str, object], method: str = "get"
+) -> MagicMock:
+    """Build a MagicMock that mimics ``create_http_session`` returning a
+    response whose ``status`` / ``headers`` are set from ``response_attrs``.
+
+    Mocks the chain ``create_http_session().__aenter__() -> session``, and
+    ``session.<method>().__aenter__() -> response``.
+    """
+    mock_response = MagicMock()
+    for k, v in response_attrs.items():
+        setattr(mock_response, k, v)
+
+    mock_session = MagicMock()
+    method_mock = getattr(mock_session, method)  # pyright: ignore[reportAny]
+    method_mock.return_value.__aenter__ = AsyncMock(  # pyright: ignore[reportAny]
+        return_value=mock_response
+    )
+    method_mock.return_value.__aexit__ = AsyncMock(  # pyright: ignore[reportAny]
+        return_value=None
+    )
+
+    mock_factory = MagicMock()
+    mock_factory.return_value.__aenter__ = AsyncMock(  # pyright: ignore[reportAny]
+        return_value=mock_session
+    )
+    mock_factory.return_value.__aexit__ = AsyncMock(  # pyright: ignore[reportAny]
+        return_value=None
+    )
+    return mock_factory
+
+
+REAL_HF_429_HEADER_DICT = {"ratelimit": '"api";r=0;t=52'}
+
+
+class TestRateLimitAtHttpCallSites:
+    """Verify each HF call site translates an HTTP 429 into a
+    ``HuggingFaceRateLimitError`` carrying the parsed ``retry_after``.
+
+    These tests would catch regressions where (a) the 429 branch is
+    deleted, (b) ``_parse_retry_after`` stops being called, or
+    (c) the wrong header object is passed to it.
+    """
+
+    async def test_fetch_file_list_maps_429_to_rate_limit_error(self) -> None:
+        mock_factory = _make_mock_session_returning(
+            {"status": 429, "headers": REAL_HF_429_HEADER_DICT}
+        )
+        with (
+            patch("exo.download.download_utils.create_http_session", mock_factory),
+            pytest.raises(HuggingFaceRateLimitError) as exc_info,
+        ):
+            await _fetch_file_list(ModelId("test/model"), "main")
+        assert exc_info.value.retry_after == 52.0
+
+    async def test_file_meta_maps_429_to_rate_limit_error(self) -> None:
+        mock_factory = _make_mock_session_returning(
+            {"status": 429, "headers": REAL_HF_429_HEADER_DICT}, method="head"
+        )
+        with (
+            patch("exo.download.download_utils.create_http_session", mock_factory),
+            pytest.raises(HuggingFaceRateLimitError) as exc_info,
+        ):
+            await file_meta(ModelId("test/model"), "main", "weights.safetensors")
+        assert exc_info.value.retry_after == 52.0
+
+    async def test_file_meta_maps_429_after_307_redirect(self) -> None:
+        """When the initial HEAD 307s and the redirected HEAD then 429s,
+        the 429 must still surface as ``HuggingFaceRateLimitError``."""
+        # First HEAD -> 307 with a Location header pointing somewhere new.
+        first_response = MagicMock()
+        first_response.status = 307
+        first_response.headers = {"location": "/redirected/url"}
+
+        # Second HEAD (the recursive call) -> 429 with the real-HF header.
+        second_response = MagicMock()
+        second_response.status = 429
+        second_response.headers = REAL_HF_429_HEADER_DICT
+
+        responses = iter([first_response, second_response])
+
+        mock_session = MagicMock()
+        mock_session.head.return_value.__aenter__ = AsyncMock(  # pyright: ignore[reportAny]
+            side_effect=lambda: next(responses)
+        )
+        mock_session.head.return_value.__aexit__ = AsyncMock(  # pyright: ignore[reportAny]
+            return_value=None
+        )
+
+        mock_factory = MagicMock()
+        mock_factory.return_value.__aenter__ = AsyncMock(  # pyright: ignore[reportAny]
+            return_value=mock_session
+        )
+        mock_factory.return_value.__aexit__ = AsyncMock(  # pyright: ignore[reportAny]
+            return_value=None
+        )
+
+        with (
+            patch("exo.download.download_utils.create_http_session", mock_factory),
+            pytest.raises(HuggingFaceRateLimitError) as exc_info,
+        ):
+            await file_meta(ModelId("test/model"), "main", "weights.safetensors")
+        assert exc_info.value.retry_after == 52.0
+
+    async def test_download_file_maps_429_to_rate_limit_error(
+        self, tmp_path: Path
+    ) -> None:
+        target_dir = tmp_path / "downloads"
+        await aios.makedirs(target_dir, exist_ok=True)
+        # No local file -> _download_file goes straight to file_meta then GET.
+        # We need both calls to succeed enough to reach the GET branch:
+        #   - file_meta returns a non-429 (size, etag) so we proceed.
+        #   - the GET then 429s.
+        with (
+            patch(
+                "exo.download.download_utils.file_meta",
+                new_callable=AsyncMock,
+                return_value=(100, "abc123"),
+            ),
+            patch(
+                "exo.download.download_utils.create_http_session",
+                _make_mock_session_returning(
+                    {"status": 429, "headers": REAL_HF_429_HEADER_DICT}
+                ),
+            ),
+            pytest.raises(HuggingFaceRateLimitError) as exc_info,
+        ):
+            await _download_file(
+                ModelId("test/model"), "main", "weights.safetensors", target_dir
+            )
+        assert exc_info.value.retry_after == 52.0

src/exo/main.py

@@ -3,6 +3,7 @@ import multiprocessing as mp
 import os
 import resource
 import signal
+import sys
 from dataclasses import dataclass, field
 from typing import Self
 
@@ -22,6 +23,8 @@ from exo.shared.election import Election, ElectionResult
 from exo.shared.logging import logger_cleanup, logger_setup
 from exo.shared.types.common import NodeId, SessionId
 from exo.utils.channels import Receiver, channel
+from exo.utils.daemon import detach_stdio_to_devnull
+from exo.utils.pidfile import PidfileLockError, acquire_exo_pidfile
 from exo.utils.pydantic_ext import FrozenModel
 from exo.utils.task_group import TaskGroup
 from exo.worker.main import Worker
@@ -264,14 +267,26 @@ class Node:
 
 
 def main():
+    # Exit early if no PID file (not compatible with double-for daemonization yet)
+    try:
+        pidfile = acquire_exo_pidfile()
+    except PidfileLockError as exception:
+        print(exception, file=sys.stderr)
+        raise SystemExit(1) from exception
+
     args = Args.parse()
     soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE)
     target = min(max(soft, 65535), hard)
     resource.setrlimit(resource.RLIMIT_NOFILE, (target, hard))
 
     mp.set_start_method("spawn", force=True)
+
     # TODO: Refactor the current verbosity system
     logger_setup(EXO_LOG, args.verbosity)
+    if args.no_stdio:
+        detach_stdio_to_devnull()
+        logger.info("Detached stdio to /dev/null")
+
     logger.info(f"{'=' * 40}")
     logger.info(f"Starting EXO | pid={os.getpid()}")
     logger.info(f"{'=' * 40}")
@@ -306,6 +321,7 @@ def main():
     finally:
         logger.info("EXO Shutdown complete")
         logger_cleanup()
+        del pidfile
 
 
 class Args(FrozenModel):
@@ -319,6 +335,7 @@ class Args(FrozenModel):
     offline: bool = os.getenv("EXO_OFFLINE", "false").lower() == "true"
     no_batch: bool = False
     fast_synch: bool | None = None  # None = auto, True = force on, False = force off
+    no_stdio: bool = False
     bootstrap_peers: list[str] = []
     libp2p_port: int
 
@@ -378,6 +395,11 @@ class Args(FrozenModel):
             action="store_true",
             help="Disable continuous batching, use sequential generation",
         )
+        parser.add_argument(
+            "--no-stdio",
+            action="store_true",
+            help="Detach stdin/stdout/stderr to /dev/null after logging is configured",
+        )
         parser.add_argument(
             "--bootstrap-peers",
             type=lambda s: [p for p in s.split(",") if p],

src/exo/master/main.py

@@ -365,6 +365,7 @@ class Master:
                                 self.state.node_memory,
                                 self.state.node_network,
                                 download_status=self.state.downloads,
+                                node_rdma_ctl=self.state.node_rdma_ctl,
                             )
                             transition_events = get_transition_events(
                                 self.state.instances, placement, self.state.tasks

src/exo/master/placement.py

@@ -28,7 +28,7 @@ from exo.shared.types.events import (
     TaskStatusUpdated,
 )
 from exo.shared.types.memory import Memory
-from exo.shared.types.profiling import MemoryUsage, NodeNetworkInfo
+from exo.shared.types.profiling import MemoryUsage, NodeNetworkInfo, NodeRdmaCtlStatus
 from exo.shared.types.tasks import Task, TaskId, TaskStatus
 from exo.shared.types.worker.downloads import (
     DownloadCompleted,
@@ -105,6 +105,7 @@ def place_instance(
     node_network: Mapping[NodeId, NodeNetworkInfo],
     required_nodes: set[NodeId] | None = None,
     download_status: Mapping[NodeId, Sequence[DownloadProgress]] | None = None,
+    node_rdma_ctl: Mapping[NodeId, NodeRdmaCtlStatus] | None = None,
 ) -> dict[InstanceId, Instance]:
     cycles = topology.get_cycles()
     candidate_cycles = list(filter(lambda it: len(it) >= command.min_nodes, cycles))
@@ -166,8 +167,18 @@ def place_instance(
 
     smallest_cycles = get_smallest_cycles(cycles_with_sufficient_memory)
 
+    rdma_ctl_status = node_rdma_ctl or {}
+
+    def _all_rdma_ctl_enabled(cycle: Cycle) -> bool:
+        return all(
+            ((status := rdma_ctl_status.get(node_id)) is not None and status.enabled)
+            for node_id in cycle
+        )
+
     smallest_rdma_cycles = [
-        cycle for cycle in smallest_cycles if topology.is_rdma_cycle(cycle)
+        cycle
+        for cycle in smallest_cycles
+        if topology.is_rdma_cycle(cycle) and _all_rdma_ctl_enabled(cycle)
     ]
 
     if command.instance_meta == InstanceMeta.MlxJaccl:

src/exo/master/tests/test_placement.py

@@ -21,7 +21,11 @@ from exo.shared.types.events import (
 )
 from exo.shared.types.memory import Memory
 from exo.shared.types.multiaddr import Multiaddr
-from exo.shared.types.profiling import NetworkInterfaceInfo, NodeNetworkInfo
+from exo.shared.types.profiling import (
+    NetworkInterfaceInfo,
+    NodeNetworkInfo,
+    NodeRdmaCtlStatus,
+)
 from exo.shared.types.tasks import TaskId, TaskStatus, TextGeneration
 from exo.shared.types.text_generation import (
     InputMessage,
@@ -439,8 +443,21 @@ def test_tensor_rdma_backend_connectivity_matrix(
         min_nodes=1,
     )
 
+    node_rdma_ctl = {
+        node_a: NodeRdmaCtlStatus(enabled=True),
+        node_b: NodeRdmaCtlStatus(enabled=True),
+        node_c: NodeRdmaCtlStatus(enabled=True),
+    }
+
     # act
-    placements = place_instance(cic, topology, {}, node_memory, node_network)
+    placements = place_instance(
+        cic,
+        topology,
+        {},
+        node_memory,
+        node_network,
+        node_rdma_ctl=node_rdma_ctl,
+    )
 
     # assert
     assert len(placements) == 1
@@ -482,6 +499,131 @@ def test_tensor_rdma_backend_connectivity_matrix(
             assert len(ip_part.split(".")) == 4
 
 
+def _build_three_node_rdma_topology() -> tuple[
+    Topology, NodeId, NodeId, NodeId, dict[NodeId, NodeNetworkInfo]
+]:
+    topology = Topology()
+    node_a = NodeId()
+    node_b = NodeId()
+    node_c = NodeId()
+
+    ethernet_interface = NetworkInterfaceInfo(name="en0", ip_address="10.0.0.1")
+    ethernet_conn = SocketConnection(
+        sink_multiaddr=Multiaddr(address="/ip4/10.0.0.1/tcp/8000")
+    )
+    node_network = {
+        node_a: NodeNetworkInfo(interfaces=[ethernet_interface]),
+        node_b: NodeNetworkInfo(interfaces=[ethernet_interface]),
+        node_c: NodeNetworkInfo(interfaces=[ethernet_interface]),
+    }
+
+    for n in (node_a, node_b, node_c):
+        topology.add_node(n)
+
+    rdma_pairs = [
+        (node_a, node_b, 3),
+        (node_b, node_a, 3),
+        (node_b, node_c, 4),
+        (node_c, node_b, 4),
+        (node_a, node_c, 5),
+        (node_c, node_a, 5),
+    ]
+    for src, sink, iface in rdma_pairs:
+        topology.add_connection(
+            Connection(source=src, sink=sink, edge=create_rdma_connection(iface))
+        )
+
+    socket_pairs = [
+        (node_a, node_b),
+        (node_b, node_c),
+        (node_c, node_a),
+        (node_a, node_c),
+        (node_b, node_a),
+        (node_c, node_b),
+    ]
+    for src, sink in socket_pairs:
+        topology.add_connection(Connection(source=src, sink=sink, edge=ethernet_conn))
+
+    return topology, node_a, node_b, node_c, node_network
+
+
+def test_place_mlx_jaccl_rejects_when_a_node_has_rdma_ctl_disabled(
+    model_card: ModelCard,
+):
+    # arrange
+    model_card = model_card.model_copy(
+        update={"n_layers": 12, "storage_size": Memory.from_bytes(1500)}
+    )
+    topology, node_a, node_b, node_c, node_network = _build_three_node_rdma_topology()
+    node_memory = {
+        node_a: create_node_memory(500),
+        node_b: create_node_memory(500),
+        node_c: create_node_memory(500),
+    }
+    node_rdma_ctl = {
+        node_a: NodeRdmaCtlStatus(enabled=True),
+        node_b: NodeRdmaCtlStatus(enabled=True),
+        node_c: NodeRdmaCtlStatus(enabled=False),
+    }
+    cic = PlaceInstance(
+        sharding=Sharding.Tensor,
+        instance_meta=InstanceMeta.MlxJaccl,
+        command_id=CommandId(),
+        model_card=model_card,
+        min_nodes=3,
+    )
+
+    # act / assert
+    with pytest.raises(
+        ValueError, match="Requested RDMA \\(MlxJaccl\\) but no RDMA-connected cycles"
+    ):
+        place_instance(
+            cic,
+            topology,
+            {},
+            node_memory,
+            node_network,
+            node_rdma_ctl=node_rdma_ctl,
+        )
+
+
+def test_place_mlx_jaccl_rejects_when_node_rdma_ctl_missing(model_card: ModelCard):
+    """A node with no observed rdma_ctl status must not participate in RDMA placement."""
+    # arrange
+    model_card = model_card.model_copy(
+        update={"n_layers": 12, "storage_size": Memory.from_bytes(1500)}
+    )
+    topology, node_a, node_b, node_c, node_network = _build_three_node_rdma_topology()
+    node_memory = {
+        node_a: create_node_memory(500),
+        node_b: create_node_memory(500),
+        node_c: create_node_memory(500),
+    }
+    # node_c has no rdma_ctl entry at all
+    node_rdma_ctl = {
+        node_a: NodeRdmaCtlStatus(enabled=True),
+        node_b: NodeRdmaCtlStatus(enabled=True),
+    }
+    cic = PlaceInstance(
+        sharding=Sharding.Tensor,
+        instance_meta=InstanceMeta.MlxJaccl,
+        command_id=CommandId(),
+        model_card=model_card,
+        min_nodes=3,
+    )
+
+    # act / assert
+    with pytest.raises(ValueError):
+        place_instance(
+            cic,
+            topology,
+            {},
+            node_memory,
+            node_network,
+            node_rdma_ctl=node_rdma_ctl,
+        )
+
+
 def _make_task(
     instance_id: InstanceId,
     status: TaskStatus = TaskStatus.Running,

src/exo/shared/apply.py

@@ -4,7 +4,8 @@ from datetime import datetime
 
 from loguru import logger
 
-from exo.shared.types.common import NodeId
+from exo.shared.models.model_cards import ModelCard
+from exo.shared.types.common import ModelId, NodeId
 from exo.shared.types.events import (
     ChunkGenerated,
     CustomModelCardAdded,
@@ -65,6 +66,18 @@ from exo.utils.info_gatherer.info_gatherer import (
 )
 
 
+def _is_rdma_ctl_enabled(
+    node_id: NodeId, node_rdma_ctl: Mapping[NodeId, NodeRdmaCtlStatus]
+) -> bool:
+    """A node is RDMA-capable only if rdma_ctl status has been observed as enabled.
+
+    Missing entries default to ``False`` — if we have not yet observed (or the node
+    cannot run) ``rdma_ctl``, it must not participate in an RDMA-backed instance.
+    """
+    status = node_rdma_ctl.get(node_id)
+    return status is not None and status.enabled
+
+
 def event_apply(event: Event, state: State) -> State:
     """Apply an event to state."""
     match event:
@@ -75,10 +88,12 @@ def event_apply(event: Event, state: State) -> State:
             | InputChunkReceived()
             | TracesCollected()
             | TracesMerged()
-            | CustomModelCardAdded()
-            | CustomModelCardDeleted()
         ):  # Pass-through events that don't modify state
             return state
+        case CustomModelCardAdded():
+            return apply_custom_model_card_added(event, state)
+        case CustomModelCardDeleted():
+            return apply_custom_model_card_deleted(event, state)
         case InstanceCreated():
             return apply_instance_created(event, state)
         case InstanceDeleted():
@@ -397,6 +412,9 @@ def apply_node_gathered_info(event: NodeGatheredInfo, state: State) -> State:
                 for nid in state.node_thunderbolt
                 for tb_ident in state.node_thunderbolt[nid].interfaces
             }
+            source_is_rdma_enabled = _is_rdma_ctl_enabled(
+                event.node_id, state.node_rdma_ctl
+            )
             as_rdma_conns = [
                 Connection(
                     source=event.node_id,
@@ -409,6 +427,10 @@ def apply_node_gathered_info(event: NodeGatheredInfo, state: State) -> State:
                 for tb_conn in info.conns
                 if tb_conn.source_uuid in conn_map
                 if tb_conn.sink_uuid in conn_map
+                if source_is_rdma_enabled
+                and _is_rdma_ctl_enabled(
+                    conn_map[tb_conn.sink_uuid][0], state.node_rdma_ctl
+                )
             ]
             topology.replace_all_out_rdma_connections(event.node_id, as_rdma_conns)
         case ThunderboltBridgeInfo():
@@ -432,6 +454,12 @@ def apply_node_gathered_info(event: NodeGatheredInfo, state: State) -> State:
                 **state.node_rdma_ctl,
                 event.node_id: NodeRdmaCtlStatus(enabled=info.enabled),
             }
+            # If RDMA just got disabled on this node, drop any RDMA edges touching it
+            # so placement / topology consumers cannot pick a disabled node for an
+            # RDMA-backed instance. (Edges will repopulate on the next
+            # MacThunderboltConnections poll once both endpoints are enabled again.)
+            if not info.enabled:
+                topology.remove_all_rdma_connections_touching(event.node_id)
 
     return state.model_copy(update=update)
 
@@ -447,3 +475,22 @@ def apply_topology_edge_deleted(event: TopologyEdgeDeleted, state: State) -> Sta
     topology.remove_connection(event.conn)
     # TODO: Clean up removing the reverse connection
     return state.model_copy(update={"topology": topology})
+
+
+def apply_custom_model_card_added(event: CustomModelCardAdded, state: State) -> State:
+    new_cards: Mapping[ModelId, ModelCard] = {
+        **state.custom_model_cards,
+        event.model_card.model_id: event.model_card,
+    }
+    return state.model_copy(update={"custom_model_cards": new_cards})
+
+
+def apply_custom_model_card_deleted(
+    event: CustomModelCardDeleted, state: State
+) -> State:
+    new_cards: Mapping[ModelId, ModelCard] = {
+        model_id: card
+        for model_id, card in state.custom_model_cards.items()
+        if model_id != event.model_id
+    }
+    return state.model_copy(update={"custom_model_cards": new_cards})

src/exo/shared/constants.py

@@ -69,6 +69,7 @@ DASHBOARD_DIR = (
 EXO_LOG_DIR = EXO_CACHE_HOME / "exo_log"
 EXO_LOG = EXO_LOG_DIR / "exo.log"
 EXO_TEST_LOG = EXO_CACHE_HOME / "exo_test.log"
+EXO_PID_FILE = EXO_CACHE_HOME / "exo.pid"
 
 # Identity (config)
 EXO_NODE_ID_KEYPAIR = EXO_CONFIG_HOME / "node_id.keypair"

src/exo/shared/models/model_cards.py

@@ -39,7 +39,57 @@ _BUILTIN_CARD_DIRS = [
     Path(RESOURCES_DIR) / "image_model_cards",
 ]
 
-_card_cache: dict[ModelId, "ModelCard"] = {}
+
+class _CardCache:
+    def __init__(self):
+        self.cc: dict[ModelId, "ModelCard"] = {}
+
+    def get(self, model_id: ModelId) -> "ModelCard | None":
+        return self.cc.get(model_id)
+
+    async def save(self, card: "ModelCard"):
+        self.cc[card.model_id] = card
+        try:
+            await card.save_to_custom_dir()
+        except OSError as e:
+            logger.warning(f"failed to save custom model card ({e.strerror})")
+
+    async def pop(self, model_id: ModelId) -> "ModelCard | None":
+        """Delete a user-added custom model card. Returns True if deleted."""
+        card_path = _custom_cards_dir / (ModelId(model_id).normalize() + ".toml")
+        try:
+            if await card_path.exists():
+                await card_path.unlink()
+                return self.cc.pop(model_id, None)
+        except OSError as e:
+            logger.warning(f"failed to delete custom model card ({e.strerror})")
+
+    async def list_all(self) -> list["ModelCard"]:
+        if len(self.cc) == 0:
+            await self.refresh()
+        if EXO_ENABLE_IMAGE_MODELS:
+            return list(self.cc.values())
+        return [c for c in self.cc.values() if not _is_image_card(c)]
+
+    async def _load_cards_from_dir(self, directory: Path, *, is_custom: bool) -> None:
+        """Load all TOML model cards from a directory into the cache."""
+        async for toml_file in directory.rglob("*.toml"):
+            try:
+                card = await ModelCard.load_from_path(toml_file)
+                if is_custom:
+                    card = card.model_copy(update={"is_custom": True})
+                if self.get(card.model_id) is None:
+                    self.cc[card.model_id] = card
+            except (ValidationError, TOMLKitError):
+                pass
+
+    async def refresh(self) -> None:
+        for path in _BUILTIN_CARD_DIRS:
+            await self._load_cards_from_dir(path, is_custom=False)
+        await self._load_cards_from_dir(_custom_cards_dir, is_custom=True)
+
+
+card_cache = _CardCache()
 
 
 def detect_vision_from_config(model_id: ModelId) -> "VisionCardConfig | None":
@@ -59,42 +109,10 @@ def detect_vision_from_config(model_id: ModelId) -> "VisionCardConfig | None":
     return None
 
 
-async def _load_cards_from_dir(directory: Path, *, is_custom: bool) -> None:
-    """Load all TOML model cards from a directory into the cache."""
-    async for toml_file in directory.rglob("*.toml"):
-        try:
-            card = await ModelCard.load_from_path(toml_file)
-            if is_custom:
-                card = card.model_copy(update={"is_custom": True})
-            if card.model_id not in _card_cache:
-                _card_cache[card.model_id] = card
-        except (ValidationError, TOMLKitError):
-            pass
-
-
-async def _refresh_card_cache() -> None:
-    for path in _BUILTIN_CARD_DIRS:
-        await _load_cards_from_dir(path, is_custom=False)
-    await _load_cards_from_dir(_custom_cards_dir, is_custom=True)
-
-
 def _is_image_card(card: "ModelCard") -> bool:
     return any(t in (ModelTask.TextToImage, ModelTask.ImageToImage) for t in card.tasks)
 
 
-def get_card(model_id: ModelId) -> "ModelCard | None":
-    """Look up a single model card from the cache by ID."""
-    return _card_cache.get(model_id)
-
-
-async def get_model_cards() -> list["ModelCard"]:
-    if len(_card_cache) == 0:
-        await _refresh_card_cache()
-    if EXO_ENABLE_IMAGE_MODELS:
-        return list(_card_cache.values())
-    return [c for c in _card_cache.values() if not _is_image_card(c)]
-
-
 class ModelTask(str, Enum):
     TextGeneration = "TextGeneration"
     TextToImage = "TextToImage"
@@ -196,14 +214,13 @@ class ModelCard(FrozenModel):
     # Is it okay that model card.load defaults to network access if the card doesn't exist? do we want to be more explicit here?
     @staticmethod
     async def load(model_id: ModelId) -> "ModelCard":
-        if model_id not in _card_cache:
-            await _refresh_card_cache()
-        if (mc := _card_cache.get(model_id)) is not None:
+        if card_cache.get(model_id) is None:
+            await card_cache.refresh()
+        if (mc := card_cache.get(model_id)) is not None:
             return mc
 
         mc = await ModelCard.fetch_from_hf(model_id)
         await mc.save_to_custom_dir()
-        _card_cache[model_id] = mc
         return mc
 
     @staticmethod
@@ -233,21 +250,6 @@ class ModelCard(FrozenModel):
         )
 
 
-def add_to_card_cache(card: "ModelCard") -> None:
-    """Add or update a model card in the in-memory cache."""
-    _card_cache[card.model_id] = card
-
-
-async def delete_custom_card(model_id: ModelId) -> bool:
-    """Delete a user-added custom model card. Returns True if deleted."""
-    card_path = _custom_cards_dir / (ModelId(model_id).normalize() + ".toml")
-    if await card_path.exists():
-        await card_path.unlink()
-        _card_cache.pop(model_id, None)
-        return True
-    return False
-
-
 class ConfigData(BaseModel):
     model_config = {"extra": "ignore"}  # Allow unknown fields
 

src/exo/shared/tests/test_apply/test_apply_custom_model_cards.py

@@ -0,0 +1,44 @@
+from exo.shared.apply import apply
+from exo.shared.models.model_cards import ModelCard, ModelTask
+from exo.shared.types.common import ModelId
+from exo.shared.types.events import (
+    CustomModelCardAdded,
+    CustomModelCardDeleted,
+    IndexedEvent,
+)
+from exo.shared.types.memory import Memory
+from exo.shared.types.state import State
+
+
+def _model_card(model_id: ModelId) -> ModelCard:
+    return ModelCard(
+        model_id=model_id,
+        n_layers=1,
+        storage_size=Memory.from_bytes(1),
+        hidden_size=1,
+        supports_tensor=True,
+        tasks=[ModelTask.TextGeneration],
+    )
+
+
+def test_custom_model_card_added_is_reduced_into_state() -> None:
+    card = _model_card(ModelId("custom/model"))
+
+    state = apply(
+        State(),
+        IndexedEvent(idx=0, event=CustomModelCardAdded(model_card=card)),
+    )
+
+    assert state.custom_model_cards == {card.model_id: card}
+
+
+def test_custom_model_card_deleted_removes_card_from_state() -> None:
+    card = _model_card(ModelId("custom/model"))
+    state = State(custom_model_cards={card.model_id: card}, last_event_applied_idx=0)
+
+    state = apply(
+        state,
+        IndexedEvent(idx=1, event=CustomModelCardDeleted(model_id=card.model_id)),
+    )
+
+    assert state.custom_model_cards == {}

src/exo/shared/tests/test_apply/test_apply_rdma_gating.py

@@ -0,0 +1,231 @@
+from datetime import datetime, timezone
+
+from exo.shared.apply import apply_node_gathered_info
+from exo.shared.topology import Topology
+from exo.shared.types.common import NodeId
+from exo.shared.types.events import NodeGatheredInfo
+from exo.shared.types.profiling import (
+    NodeRdmaCtlStatus,
+    NodeThunderboltInfo,
+)
+from exo.shared.types.state import State
+from exo.shared.types.thunderbolt import ThunderboltConnection, ThunderboltIdentifier
+from exo.shared.types.topology import RDMAConnection
+from exo.utils.info_gatherer.info_gatherer import (
+    MacThunderboltConnections,
+    RdmaCtlStatus,
+)
+
+
+def _now() -> str:
+    return datetime.now(timezone.utc).isoformat()
+
+
+def _make_state_with_thunderbolt_idents(
+    *node_ids_and_uuids: tuple[NodeId, str, str],
+    rdma_ctl: dict[NodeId, NodeRdmaCtlStatus] | None = None,
+) -> State:
+    """Build a State with Thunderbolt identifiers per node so the apply MacThunderboltConnections
+    case can resolve uuid -> (node, iface)."""
+    node_thunderbolt = {
+        nid: NodeThunderboltInfo(
+            interfaces=[ThunderboltIdentifier(rdma_interface=iface, domain_uuid=uuid)]
+        )
+        for nid, uuid, iface in node_ids_and_uuids
+    }
+    return State(
+        node_thunderbolt=node_thunderbolt,
+        node_rdma_ctl=rdma_ctl or {},
+    )
+
+
+def _has_rdma_edge(topology: Topology, source: NodeId, sink: NodeId) -> bool:
+    return any(
+        isinstance(edge, RDMAConnection)
+        for edge in topology.get_all_connections_between(source, sink)
+    )
+
+
+def test_mac_thunderbolt_connections_emits_rdma_when_both_endpoints_enabled():
+    node_a = NodeId()
+    node_b = NodeId()
+    state = _make_state_with_thunderbolt_idents(
+        (node_a, "uuid-a", "rdma_en1"),
+        (node_b, "uuid-b", "rdma_en1"),
+        rdma_ctl={
+            node_a: NodeRdmaCtlStatus(enabled=True),
+            node_b: NodeRdmaCtlStatus(enabled=True),
+        },
+    )
+
+    event = NodeGatheredInfo(
+        node_id=node_a,
+        when=_now(),
+        info=MacThunderboltConnections(
+            conns=[ThunderboltConnection(source_uuid="uuid-a", sink_uuid="uuid-b")]
+        ),
+    )
+
+    new_state = apply_node_gathered_info(event, state)
+
+    assert _has_rdma_edge(new_state.topology, node_a, node_b)
+
+
+def test_mac_thunderbolt_connections_skips_rdma_when_source_rdma_ctl_disabled():
+    node_a = NodeId()
+    node_b = NodeId()
+    state = _make_state_with_thunderbolt_idents(
+        (node_a, "uuid-a", "rdma_en1"),
+        (node_b, "uuid-b", "rdma_en1"),
+        rdma_ctl={
+            node_a: NodeRdmaCtlStatus(enabled=False),
+            node_b: NodeRdmaCtlStatus(enabled=True),
+        },
+    )
+
+    event = NodeGatheredInfo(
+        node_id=node_a,
+        when=_now(),
+        info=MacThunderboltConnections(
+            conns=[ThunderboltConnection(source_uuid="uuid-a", sink_uuid="uuid-b")]
+        ),
+    )
+
+    new_state = apply_node_gathered_info(event, state)
+
+    assert not _has_rdma_edge(new_state.topology, node_a, node_b)
+
+
+def test_mac_thunderbolt_connections_skips_rdma_when_sink_rdma_ctl_disabled():
+    node_a = NodeId()
+    node_b = NodeId()
+    state = _make_state_with_thunderbolt_idents(
+        (node_a, "uuid-a", "rdma_en1"),
+        (node_b, "uuid-b", "rdma_en1"),
+        rdma_ctl={
+            node_a: NodeRdmaCtlStatus(enabled=True),
+            node_b: NodeRdmaCtlStatus(enabled=False),
+        },
+    )
+
+    event = NodeGatheredInfo(
+        node_id=node_a,
+        when=_now(),
+        info=MacThunderboltConnections(
+            conns=[ThunderboltConnection(source_uuid="uuid-a", sink_uuid="uuid-b")]
+        ),
+    )
+
+    new_state = apply_node_gathered_info(event, state)
+
+    assert not _has_rdma_edge(new_state.topology, node_a, node_b)
+
+
+def test_mac_thunderbolt_connections_skips_rdma_when_rdma_ctl_status_missing():
+    """Missing rdma_ctl status defaults to not-enabled — node is RDMA-incapable."""
+    node_a = NodeId()
+    node_b = NodeId()
+    state = _make_state_with_thunderbolt_idents(
+        (node_a, "uuid-a", "rdma_en1"),
+        (node_b, "uuid-b", "rdma_en1"),
+        rdma_ctl={
+            node_a: NodeRdmaCtlStatus(enabled=True),
+            # node_b intentionally absent
+        },
+    )
+
+    event = NodeGatheredInfo(
+        node_id=node_a,
+        when=_now(),
+        info=MacThunderboltConnections(
+            conns=[ThunderboltConnection(source_uuid="uuid-a", sink_uuid="uuid-b")]
+        ),
+    )
+
+    new_state = apply_node_gathered_info(event, state)
+
+    assert not _has_rdma_edge(new_state.topology, node_a, node_b)
+
+
+def test_rdma_ctl_status_disabled_purges_existing_rdma_edges():
+    """When a node reports rdma_ctl disabled, all RDMA edges touching it must be removed."""
+    node_a = NodeId()
+    node_b = NodeId()
+
+    # Start with both nodes RDMA-enabled and existing RDMA edges in the topology.
+    state = _make_state_with_thunderbolt_idents(
+        (node_a, "uuid-a", "rdma_en1"),
+        (node_b, "uuid-b", "rdma_en1"),
+        rdma_ctl={
+            node_a: NodeRdmaCtlStatus(enabled=True),
+            node_b: NodeRdmaCtlStatus(enabled=True),
+        },
+    )
+    state = apply_node_gathered_info(
+        NodeGatheredInfo(
+            node_id=node_a,
+            when=_now(),
+            info=MacThunderboltConnections(
+                conns=[ThunderboltConnection(source_uuid="uuid-a", sink_uuid="uuid-b")]
+            ),
+        ),
+        state,
+    )
+    state = apply_node_gathered_info(
+        NodeGatheredInfo(
+            node_id=node_b,
+            when=_now(),
+            info=MacThunderboltConnections(
+                conns=[ThunderboltConnection(source_uuid="uuid-b", sink_uuid="uuid-a")]
+            ),
+        ),
+        state,
+    )
+    assert _has_rdma_edge(state.topology, node_a, node_b)
+    assert _has_rdma_edge(state.topology, node_b, node_a)
+
+    # Now node_a flips to rdma_ctl disabled — both directions of RDMA edge must drop.
+    state = apply_node_gathered_info(
+        NodeGatheredInfo(
+            node_id=node_a, when=_now(), info=RdmaCtlStatus(enabled=False)
+        ),
+        state,
+    )
+
+    assert not _has_rdma_edge(state.topology, node_a, node_b)
+    assert not _has_rdma_edge(state.topology, node_b, node_a)
+    assert state.node_rdma_ctl[node_a].enabled is False
+
+
+def test_topology_remove_all_rdma_connections_touching_keeps_socket_edges():
+    """Purging RDMA edges for a disabled node must not affect non-RDMA edges."""
+    from exo.shared.types.multiaddr import Multiaddr
+    from exo.shared.types.topology import Connection, SocketConnection
+
+    topology = Topology()
+    node_a = NodeId()
+    node_b = NodeId()
+    topology.add_node(node_a)
+    topology.add_node(node_b)
+    topology.add_connection(
+        Connection(
+            source=node_a,
+            sink=node_b,
+            edge=RDMAConnection(
+                source_rdma_iface="rdma_en1", sink_rdma_iface="rdma_en1"
+            ),
+        )
+    )
+    socket_edge = SocketConnection(
+        sink_multiaddr=Multiaddr(address="/ip4/10.0.0.1/tcp/8000")
+    )
+    topology.add_connection(Connection(source=node_a, sink=node_b, edge=socket_edge))
+
+    topology.remove_all_rdma_connections_touching(node_a)
+
+    assert not _has_rdma_edge(topology, node_a, node_b)
+    # Socket edge survives.
+    assert any(
+        isinstance(edge, SocketConnection)
+        for edge in topology.get_all_connections_between(node_a, node_b)
+    )

src/exo/shared/topology.py

@@ -169,6 +169,22 @@ class Topology:
         for conn in new_connections:
             self.add_connection(conn)
 
+    def remove_all_rdma_connections_touching(self, node_id: NodeId) -> None:
+        """Remove every RDMA edge incident to ``node_id`` (incoming or outgoing)."""
+        if node_id not in self._vertex_indices:
+            return
+        rx_idx = self._vertex_indices[node_id]
+        rdma_edge_idxs = [
+            edge_idx
+            for edge_idx in (
+                *self._graph.out_edge_indices(rx_idx),
+                *self._graph.in_edge_indices(rx_idx),
+            )
+            if isinstance(self._graph.get_edge_data_by_index(edge_idx), RDMAConnection)
+        ]
+        for edge_idx in rdma_edge_idxs:
+            self._graph.remove_edge_from_index(edge_idx)
+
     def remove_connection(self, conn: Connection) -> None:
         if (
             conn.source not in self._vertex_indices

src/exo/shared/types/state.py

@@ -5,8 +5,9 @@ from typing import Any, cast
 from pydantic import ConfigDict, Field, field_serializer, field_validator
 from pydantic.alias_generators import to_camel
 
+from exo.shared.models.model_cards import ModelCard
 from exo.shared.topology import Topology, TopologySnapshot
-from exo.shared.types.common import NodeId
+from exo.shared.types.common import ModelId, NodeId
 from exo.shared.types.instance_link import InstanceLink, InstanceLinkId
 from exo.shared.types.profiling import (
     DiskUsage,
@@ -65,6 +66,9 @@ class State(FrozenModel):
     instance_links: Mapping[InstanceLinkId, InstanceLink] = {}
     prefill_server_ports: Mapping[RunnerId, int] = {}
 
+    # User-added model cards. Workers can reconcile their on-disk custom card cache
+    custom_model_cards: Mapping[ModelId, ModelCard] = {}
+
     @field_serializer("topology", mode="plain")
     def _encode_topology(self, value: Topology) -> TopologySnapshot:
         return value.to_snapshot()

src/exo/shared/types/text_generation.py

@@ -135,9 +135,9 @@ class TextGenerationTaskParams(BaseModel, frozen=True):
     prefill_endpoint: str | None = None
 
     def with_card_sampling_defaults(self) -> "TextGenerationTaskParams":
-        from exo.shared.models.model_cards import get_card
+        from exo.shared.models import model_cards
 
-        card = get_card(self.model)
+        card = model_cards.card_cache.get(self.model)
         if card is None:
             return self
 

src/exo/utils/async_process.py

@@ -0,0 +1,290 @@
+from __future__ import annotations
+
+import contextlib
+import faulthandler
+import multiprocessing as mp
+import os
+import sys
+from collections.abc import Callable, Iterable, Mapping
+from multiprocessing.process import BaseProcess
+from multiprocessing.resource_sharer import DupFd
+from typing import final
+
+from anyio import (
+    TASK_STATUS_IGNORED,
+    BrokenResourceError,
+    CancelScope,
+    ClosedResourceError,
+    Event,
+    create_task_group,
+    move_on_after,
+    sleep,
+    wait_readable,
+)
+from anyio.abc import TaskStatus
+from loguru import logger
+
+from exo.utils.channels import Receiver, Sender, channel
+
+_STDOUT_FD = 1
+_STDERR_FD = 2
+_READ_CHUNK_SIZE = 64 * 1024
+_TERMINATE_GRACE_SECONDS = 10.0
+_TERMINATE_RETRY_GRACE_SECONDS = 2.0
+_TERMINATE_ATTEMPTS = 10
+_KILL_GRACE_SECONDS = 5.0
+
+
+@final
+class AsyncProcess:
+    def __init__(
+        self,
+        target: Callable[..., object] | None = None,
+        name: str | None = None,
+        args: Iterable[object] = (),
+        kwargs: Mapping[str, object] | None = None,
+        *,
+        daemon: bool | None = None,
+    ) -> None:
+        # setup state
+        self._target = target
+        self._name = name
+        self._args = args
+        self._kwargs = kwargs
+        self._daemon = daemon
+
+        # lifecycle state
+        self._process: BaseProcess | None = None
+        self._pid: int | None = None
+        self._stdout_tx, self._stdout_rx = channel[bytes]()
+        self._stderr_tx, self._stderr_rx = channel[bytes]()
+        self._started = Event()
+        self._done = Event()
+        self._run_cancel_scope: CancelScope | None = None
+        self._start_error: BaseException | None = None
+        self._exitcode: int | None = None
+
+    async def run(self, *, task_status: TaskStatus[None] = TASK_STATUS_IGNORED) -> None:
+        if self._run_cancel_scope is not None or self._done.is_set():
+            raise RuntimeError("process has already been started")
+
+        stdout_read_fd: int | None = None
+        stdout_write_fd: int | None = None
+        stderr_read_fd: int | None = None
+        stderr_write_fd: int | None = None
+
+        def cleanup_stdio_fd() -> None:
+            nonlocal stdout_read_fd, stdout_write_fd, stderr_read_fd, stderr_write_fd
+            stdout_read_fd = _close_fd(stdout_read_fd)
+            stdout_write_fd = _close_fd(stdout_write_fd)
+            stderr_read_fd = _close_fd(stderr_read_fd)
+            stderr_write_fd = _close_fd(stderr_write_fd)
+
+        try:
+            with CancelScope() as run_cancel_scope:
+                self._run_cancel_scope = run_cancel_scope
+                stdout_read_fd, stdout_write_fd = os.pipe()
+                stderr_read_fd, stderr_write_fd = os.pipe()
+
+                process = mp.Process(
+                    target=_run_with_captured_stdio,
+                    name=self._name,
+                    args=(
+                        DupFd(stdout_write_fd),
+                        DupFd(stderr_write_fd),
+                        self._target,
+                        *self._args,
+                    ),
+                    kwargs={} if self._kwargs is None else self._kwargs,
+                    daemon=self._daemon,
+                )
+                process.start()
+                pid = process.pid
+                if pid is None:
+                    raise RuntimeError("started process has no pid")
+
+                # important to close parent write-side FD to prevent hangs
+                stdout_write_fd = _close_fd(stdout_write_fd)
+                stderr_write_fd = _close_fd(stderr_write_fd)
+
+                self._process = process
+                self._pid = pid
+                self._started.set()
+
+                async with create_task_group() as tg:
+                    tg.start_soon(_drain_fd, stdout_read_fd, self._stdout_tx)
+                    stdout_read_fd = None
+                    tg.start_soon(_drain_fd, stderr_read_fd, self._stderr_tx)
+                    stderr_read_fd = None
+                    task_status.started()
+                    await self.wait()
+        except BaseException as exc:
+            if not self._started.is_set():
+                self._start_error = exc
+                self._started.set()
+            raise
+        finally:
+            try:
+                with CancelScope(shield=True):
+                    await self._terminate_if_still_alive()
+            finally:
+                cleanup_stdio_fd()
+                for tx in (self._stdout_tx, self._stderr_tx):
+                    with contextlib.suppress(Exception):
+                        await tx.aclose()
+                if self._process is not None:
+                    with contextlib.suppress(ValueError):
+                        self._process.close()
+                self._run_cancel_scope = None
+                self._done.set()
+
+    async def stop(self) -> None:
+        if self._run_cancel_scope is None and not self._done.is_set():
+            raise RuntimeError("process has not been started")
+        if self._run_cancel_scope is not None:
+            self._run_cancel_scope.cancel()
+        await self._done.wait()
+
+    async def aclose(self) -> None:
+        await self.stop()
+
+    async def wait(self) -> int:
+        if self._exitcode is not None:
+            return self._exitcode
+
+        await self._started.wait()
+        if self._start_error is not None:
+            raise self._start_error
+        assert self._process is not None
+
+        while True:
+            exitcode = self.exitcode
+            if exitcode is not None:
+                return exitcode
+            await sleep(0.01)
+
+    @property
+    def pid(self) -> int:
+        if self._pid is None:
+            raise RuntimeError("process has not been started")
+        return self._pid
+
+    @property
+    def exitcode(self) -> int | None:
+        if self._exitcode is not None:
+            return self._exitcode
+        if self._process is None:
+            return None
+
+        with contextlib.suppress(ValueError):
+            exitcode = self._process.exitcode
+            if exitcode is not None:
+                self._exitcode = exitcode
+            return exitcode
+        return None
+
+    def is_alive(self) -> bool:
+        if self._process is None:
+            return False
+
+        with contextlib.suppress(ValueError):
+            return self._process.is_alive()
+        return False
+
+    # TODO: maybe in the future if needed, create stdin that is also installed,
+    #       and a ByteSendStream handle is provided for it :)
+
+    @property
+    def stdout(self) -> Receiver[bytes]:
+        return self._stdout_rx
+
+    @property
+    def stderr(self) -> Receiver[bytes]:
+        return self._stderr_rx
+
+    async def _terminate_if_still_alive(self) -> None:
+        process = self._process
+        if process is None:
+            return
+
+        if self.exitcode is not None:
+            return
+
+        with contextlib.suppress(ValueError):
+            if not process.is_alive():
+                return
+
+            logger.warning("Child process didn't shut down successfully, terminating")
+            process.terminate()
+            with move_on_after(_TERMINATE_GRACE_SECONDS):
+                await self.wait()
+
+            if self.exitcode is not None or not process.is_alive():
+                logger.warning("Terminated nicely in the first attempt!")
+                return
+
+            for attempt in range(2, _TERMINATE_ATTEMPTS + 1):
+                process.terminate()
+                with move_on_after(_TERMINATE_RETRY_GRACE_SECONDS):
+                    await self.wait()
+
+                if self.exitcode is not None or not process.is_alive():
+                    logger.warning(f"That took {attempt} attempts :)")
+                    return
+
+            logger.critical("Child process didn't respond to SIGTERM, killing")
+            j = 0
+            while True:
+                process.kill()
+                with move_on_after(_KILL_GRACE_SECONDS):
+                    await self.wait()
+                j += 1
+                if self.exitcode is not None or not process.is_alive():
+                    break
+            logger.warning(f"That took {j} attempts :(")
+
+
+# Spawn-mode multiprocessing requires a module-level target that can be pickled.
+def _run_with_captured_stdio(
+    stdout: DupFd,
+    stderr: DupFd,
+    target: Callable[..., object] | None,
+    *target_args: object,
+    **target_kwargs: object,
+) -> None:
+    stdout_fd = stdout.detach()
+    stderr_fd = stderr.detach()
+
+    try:
+        os.dup2(stdout_fd, _STDOUT_FD)
+        os.dup2(stderr_fd, _STDERR_FD)
+    finally:
+        for fd in (stdout_fd, stderr_fd):
+            if fd not in (_STDOUT_FD, _STDERR_FD):
+                _close_fd(fd)
+
+    faulthandler.enable(file=sys.stderr, all_threads=True)
+    if target is not None:
+        target(*target_args, **target_kwargs)
+
+
+async def _drain_fd(fd: int, tx: Sender[bytes]) -> None:
+    try:
+        while True:
+            await wait_readable(fd)
+            chunk = os.read(fd, _READ_CHUNK_SIZE)
+            if not chunk:
+                return
+            await tx.send(chunk)
+    except (BrokenPipeError, BrokenResourceError, ClosedResourceError):
+        pass
+    finally:
+        _close_fd(fd)
+        await tx.aclose()
+
+
+def _close_fd(fd: int | None) -> None:
+    if fd is None:
+        return
+    with contextlib.suppress(OSError):
+        os.close(fd)

src/exo/utils/daemon.py

@@ -0,0 +1,28 @@
+import os
+import sys
+
+_STDIN_FD = 0
+_STDOUT_FD = 1
+_STDERR_FD = 2
+
+
+def detach_stdio_to_devnull() -> None:
+    """Redirect process stdio file descriptors to /dev/null."""
+
+    for stream in (sys.stdout, sys.stderr, sys.__stdout__, sys.__stderr__):
+        if stream is not None:
+            stream.flush()
+
+    stdin_fd = os.open(os.devnull, os.O_RDONLY)
+    stdout_fd = os.open(os.devnull, os.O_WRONLY)
+    stderr_fd = os.open(os.devnull, os.O_WRONLY)
+
+    try:
+        # dup2 closes the target fd first, but leaves the source fd open.
+        os.dup2(stdin_fd, _STDIN_FD)
+        os.dup2(stdout_fd, _STDOUT_FD)
+        os.dup2(stderr_fd, _STDERR_FD)
+    finally:
+        for fd in (stdin_fd, stdout_fd, stderr_fd):
+            if fd not in (_STDIN_FD, _STDOUT_FD, _STDERR_FD):
+                os.close(fd)

src/exo/utils/pidfile.py

@@ -0,0 +1,28 @@
+from __future__ import annotations
+
+import os
+from typing import Final
+
+from exo_pyo3_bindings import Pidfile, PidfileError
+
+from exo.shared.constants import EXO_PID_FILE
+
+_PIDFILE_MODE: Final = 0o600
+
+
+class PidfileLockError(RuntimeError):
+    pass
+
+
+def acquire_exo_pidfile() -> Pidfile:
+    path = EXO_PID_FILE
+    os.makedirs(os.path.dirname(path), exist_ok=True)
+    try:
+        pidfile = Pidfile(path, _PIDFILE_MODE)
+        pidfile.write()
+    except (OSError, PidfileError) as exception:
+        raise PidfileLockError(
+            f"Failed to acquire EXO pidfile at {path}: {exception}"
+        ) from exception
+
+    return pidfile

src/exo/utils/power_sampler.py

@@ -19,19 +19,21 @@ class PowerSampler:
     ):
         self._get_node_system = get_node_system
         self._interval = interval
-        self._samples: defaultdict[NodeId, list[SystemPerformanceProfile]] = (
-            defaultdict(list)
-        )
+        self._samples: defaultdict[
+            NodeId, list[tuple[float, SystemPerformanceProfile]]
+        ] = defaultdict(list)
         self._start_time: float | None = None
         self._stopped = False
 
-    def _take_sample(self) -> None:
+    def _take_sample(self, t_rel: float | None = None) -> None:
+        assert self._start_time is not None
+        ts = t_rel if t_rel is not None else time.perf_counter() - self._start_time
         for node_id, profile in self._get_node_system().items():
-            self._samples[node_id].append(profile)
+            self._samples[node_id].append((ts, profile))
 
     async def run(self) -> None:
         self._start_time = time.perf_counter()
-        self._take_sample()
+        self._take_sample(t_rel=0.0)
         while not self._stopped:
             await anyio.sleep(self._interval)
             self._take_sample()
@@ -39,26 +41,51 @@ class PowerSampler:
     def result(self) -> PowerUsage:
         self._stopped = True
         assert self._start_time is not None, "result() called before run()"
-        self._take_sample()
         elapsed = time.perf_counter() - self._start_time
+        self._take_sample(t_rel=elapsed)
 
         node_stats: list[NodePowerStats] = []
-        for node_id, profiles in self._samples.items():
-            n = len(profiles)
+        total_energy_j = 0.0
+        for node_id, ts_profiles in self._samples.items():
+            n = len(ts_profiles)
             if n == 0:
                 continue
+            node_energy_j = trapezoidal_energy(ts_profiles, elapsed)
+            avg_power_w = node_energy_j / elapsed if elapsed > 0 else 0.0
+            total_energy_j += node_energy_j
             node_stats.append(
                 NodePowerStats(
                     node_id=node_id,
                     samples=n,
-                    avg_sys_power=sum(p.sys_power for p in profiles) / n,
+                    avg_sys_power=avg_power_w,
                 )
             )
 
-        total_avg_sys = sum(ns.avg_sys_power for ns in node_stats)
+        total_avg_sys_w = total_energy_j / elapsed if elapsed > 0 else 0.0
         return PowerUsage(
             elapsed_seconds=elapsed,
             nodes=node_stats,
-            total_avg_sys_power_watts=total_avg_sys,
-            total_energy_joules=total_avg_sys * elapsed,
+            total_avg_sys_power_watts=total_avg_sys_w,
+            total_energy_joules=total_energy_j,
         )
+
+
+def trapezoidal_energy(
+    ts_profiles: list[tuple[float, SystemPerformanceProfile]],
+    elapsed: float,
+) -> float:
+    """Integrate sys_power(t) over the sample window using the trapezoidal rule.
+    First sample is anchored at t=0 and last at t=elapsed (set by `run` /
+    `result`), so the integral spans the full request interval. Falls back to
+    power * elapsed when only one sample exists (constant-power assumption)."""
+    if len(ts_profiles) == 1:
+        return ts_profiles[0][1].sys_power * elapsed
+    energy_j = 0.0
+    for i in range(1, len(ts_profiles)):
+        t_prev, p_prev = ts_profiles[i - 1]
+        t_cur, p_cur = ts_profiles[i]
+        dt = t_cur - t_prev
+        if dt <= 0:
+            continue
+        energy_j += (p_prev.sys_power + p_cur.sys_power) / 2.0 * dt
+    return energy_j

src/exo/utils/tests/conftest.py

@@ -0,0 +1,8 @@
+import multiprocessing as mp
+
+import pytest
+
+
+@pytest.fixture(scope="session", autouse=True)
+def mp_force_spawn():
+    mp.set_start_method("spawn", force=True)

src/exo/utils/tests/test_async_process.py

@@ -0,0 +1,515 @@
+import contextlib
+import os
+import signal
+import sys
+import time
+from collections.abc import AsyncIterator, Callable
+from types import FrameType
+
+import mlx.core as mx
+import pytest
+from _pytest.capture import CaptureFixture
+from anyio import EndOfStream, create_task_group, fail_after
+from pytest import MonkeyPatch
+
+import exo.utils.async_process as async_process
+from exo.utils.async_process import (
+    AsyncProcess,
+)
+from exo.utils.channels import MpSender, Receiver, mp_channel
+
+
+def _write_to_stdio(prefix: str, *, stderr_suffix: str) -> None:
+    print(f"{prefix}: python stdout")
+    print(f"{prefix}: python stderr {stderr_suffix}", file=sys.stderr)
+    os.write(1, f"{prefix}: fd stdout\n".encode())
+    os.write(2, f"{prefix}: fd stderr {stderr_suffix}\n".encode())
+
+
+def _write_large_output() -> None:
+    os.write(1, b"stdout-0123456789")
+    os.write(2, b"stderr-0123456789")
+
+
+def _write_all(fd: int, data: bytes) -> None:
+    remaining = memoryview(data)
+    while remaining:
+        written = os.write(fd, remaining)
+        remaining = remaining[written:]
+
+
+def _write_large_exact_output(size: int) -> None:
+    _write_all(1, b"stdout:" + (b"x" * size))
+    _write_all(2, b"stderr:" + (b"y" * size))
+
+
+def _raise_after_stderr_write() -> None:
+    os.write(2, b"stderr before exception\n")
+    raise RuntimeError("child boom")
+
+
+def _exit_after_stdio_write(prefix: str, exitcode: int) -> None:
+    os.write(1, f"{prefix}: stdout before _exit\n".encode())
+    os.write(2, f"{prefix}: stderr before _exit\n".encode())
+    os._exit(exitcode)
+
+
+def _abort_after_stdio_write(prefix: str) -> None:
+    os.write(1, f"{prefix}: stdout before abort\n".encode())
+    os.write(2, f"{prefix}: stderr before abort\n".encode())
+    os.abort()
+
+
+def _close_stdio_and_exit() -> None:
+    os.close(1)
+    os.close(2)
+    os._exit(0)
+
+
+def _exit_on_sigterm(exitcode: int) -> None:
+    def handle_sigterm(_signum: int, _frame: FrameType | None) -> None:
+        os._exit(exitcode)
+
+    signal.signal(signal.SIGTERM, handle_sigterm)
+    os.write(1, b"sigterm-ready\n")
+    while True:
+        time.sleep(0.1)
+
+
+def _exit_after_repeated_sigterm(required_count: int, exitcode: int) -> None:
+    sigterm_count = 0
+
+    def handle_sigterm(_signum: int, _frame: FrameType | None) -> None:
+        nonlocal sigterm_count
+        sigterm_count += 1
+        if sigterm_count >= required_count:
+            os._exit(exitcode)
+
+    signal.signal(signal.SIGTERM, handle_sigterm)
+    os.write(1, b"sigterm-ready\n")
+    while True:
+        time.sleep(0.1)
+
+
+def _ignore_sigterm_forever() -> None:
+    signal.signal(signal.SIGTERM, signal.SIG_IGN)
+    os.write(1, b"sigterm-ready\n")
+    while True:
+        time.sleep(0.1)
+
+
+def _sleep_forever() -> None:
+    while True:
+        time.sleep(0.1)
+
+
+def _send_over_mp_channel(send: MpSender[str]) -> None:
+    send.send("hello from child")
+    send.close()
+
+
+def _mlx_force_oom(size: int = 40_000) -> None:
+    """
+    Force an Out-Of-Memory (OOM) error in MLX by performing large tensor operations.
+    """
+    print("CHILD: start")
+
+    mx.set_default_device(mx.gpu)
+    a = mx.random.uniform(shape=(size, size), dtype=mx.float32)
+    b = mx.random.uniform(shape=(size, size), dtype=mx.float32)
+    mx.eval(a, b)
+    c = mx.matmul(a, b)
+    d = mx.matmul(a, c)
+    e = mx.matmul(b, c)
+    f = mx.sigmoid(d + e)
+    mx.eval(f)
+
+    print("CHILD: end")
+
+
+async def _collect_stream(
+    stream: Receiver[bytes],
+    output: bytearray,
+) -> None:
+    while True:
+        try:
+            output.extend(await stream.receive())
+        except EndOfStream:
+            return
+
+
+async def _collect_process_output(
+    process: AsyncProcess,
+) -> tuple[int, bytes, bytes]:
+    stdout = bytearray()
+    stderr = bytearray()
+    exitcodes: list[int] = []
+
+    async with create_task_group() as task_group:
+        task_group.start_soon(_collect_stream, process.stdout, stdout)
+        task_group.start_soon(_collect_stream, process.stderr, stderr)
+        exitcodes.append(await process.wait())
+
+    if not exitcodes:
+        raise RuntimeError("process exited without a return code")
+    return exitcodes[0], bytes(stdout), bytes(stderr)
+
+
+def _fd_identity(fd: int) -> tuple[int, int]:
+    fd_stat = os.fstat(fd)
+    return fd_stat.st_dev, fd_stat.st_ino
+
+
+def _fd_count() -> int | None:
+    for fd_dir in ("/proc/self/fd", "/dev/fd"):
+        with contextlib.suppress(OSError):
+            return len(os.listdir(fd_dir))
+    return None
+
+
+@contextlib.asynccontextmanager
+async def _started_process(process: AsyncProcess) -> AsyncIterator[None]:
+    async with create_task_group() as task_group:
+        await task_group.start(process.run)
+        try:
+            yield
+        finally:
+            await process.stop()
+
+
+async def _run_and_collect(
+    target: Callable[..., object] | None,
+    *,
+    args: tuple[object, ...] = (),
+    kwargs: dict[str, object] | None = None,
+) -> tuple[int, bytes, bytes]:
+    process = AsyncProcess(
+        target,
+        args=args,
+        kwargs=kwargs,
+    )
+    async with _started_process(process):
+        return await _collect_process_output(process)
+
+
+@pytest.mark.anyio
+async def test_spawn_process_captures_stdout_and_stderr_separately(
+    capfd: CaptureFixture[str],
+) -> None:
+    process = AsyncProcess(
+        _write_to_stdio,
+        args=("child",),
+        kwargs={"stderr_suffix": "error"},
+    )
+    async with _started_process(process):
+        exitcode, stdout_bytes, stderr_bytes = await _collect_process_output(process)
+
+    parent_output = capfd.readouterr()
+    stdout = stdout_bytes.decode("utf-8", errors="replace")
+    stderr = stderr_bytes.decode("utf-8", errors="replace")
+
+    assert exitcode == 0
+    assert "child: python stdout" in stdout
+    assert "child: fd stdout" in stdout
+    assert "child: python stderr error" in stderr
+    assert "child: fd stderr error" in stderr
+    assert "child:" not in parent_output.out
+    assert "child:" not in parent_output.err
+
+
+@pytest.mark.anyio
+async def test_process_with_no_target_exits_successfully() -> None:
+    exitcode, stdout, stderr = await _run_and_collect(None)
+
+    assert exitcode == 0
+    assert stdout == b""
+    assert stderr == b""
+
+
+@pytest.mark.anyio
+async def test_output_receivers_and_wait_are_safe_immediately_after_run_starts() -> (
+    None
+):
+    process = AsyncProcess(
+        _write_to_stdio,
+        args=("immediate",),
+        kwargs={"stderr_suffix": "error"},
+    )
+    result: tuple[int, bytes, bytes] | None = None
+
+    async with create_task_group() as task_group:
+        await task_group.start(process.run)
+        try:
+            result = await _collect_process_output(process)
+        finally:
+            await process.stop()
+
+    assert result is not None
+    exitcode, stdout, stderr = result
+    assert exitcode == 0
+    assert b"immediate: fd stdout\n" in stdout
+    assert b"immediate: fd stderr error\n" in stderr
+
+
+@pytest.mark.anyio
+async def test_stop_before_run_raises() -> None:
+    process = AsyncProcess(
+        _write_to_stdio,
+        args=("never",),
+        kwargs={"stderr_suffix": "run"},
+    )
+
+    assert not process.is_alive()
+    with pytest.raises(RuntimeError, match="process has not been started"):
+        await process.stop()
+
+
+@pytest.mark.anyio
+async def test_process_run_is_one_shot() -> None:
+    process = AsyncProcess(None)
+
+    await process.run()
+
+    with pytest.raises(RuntimeError, match="process has already been started"):
+        await process.run()
+
+
+@pytest.mark.anyio
+async def test_process_started_with_task_group_start_can_stop_immediately() -> None:
+    process = AsyncProcess(_sleep_forever)
+
+    async with create_task_group() as task_group:
+        await task_group.start(process.run)
+        assert process.is_alive()
+        with fail_after(2):
+            await process.stop()
+
+    assert not process.is_alive()
+
+
+@pytest.mark.anyio
+async def test_stdout_receiver_yields_bytes_chunks() -> None:
+    process = AsyncProcess(_write_large_output)
+
+    async with _started_process(process):
+        first_stdout = await process.stdout.receive()
+        exitcode, remaining_stdout, stderr = await _collect_process_output(process)
+
+    assert exitcode == 0
+    assert first_stdout + remaining_stdout == b"stdout-0123456789"
+    assert stderr == b"stderr-0123456789"
+
+
+@pytest.mark.anyio
+async def test_output_can_be_read_after_process_exits() -> None:
+    process = AsyncProcess(_write_large_output)
+
+    async with create_task_group() as task_group:
+        await task_group.start(process.run)
+        assert await process.wait() == 0
+
+    assert await process.stdout.receive() == b"stdout-0123456789"
+    assert await process.stderr.receive() == b"stderr-0123456789"
+    with pytest.raises(EndOfStream):
+        await process.stdout.receive()
+    with pytest.raises(EndOfStream):
+        await process.stderr.receive()
+
+
+@pytest.mark.anyio
+async def test_large_stdout_and_stderr_are_not_lost() -> None:
+    size = 1024 * 1024
+    exitcode, stdout, stderr = await _run_and_collect(
+        _write_large_exact_output,
+        args=(size,),
+    )
+
+    assert exitcode == 0
+    assert stdout == b"stdout:" + (b"x" * size)
+    assert stderr == b"stderr:" + (b"y" * size)
+
+
+@pytest.mark.anyio
+async def test_child_exception_traceback_is_captured_from_stderr() -> None:
+    process = AsyncProcess(_raise_after_stderr_write)
+
+    async with _started_process(process):
+        exitcode, _, stderr_bytes = await _collect_process_output(process)
+
+    assert exitcode == 1
+    stderr = stderr_bytes.decode("utf-8", errors="replace")
+    assert "stderr before exception" in stderr
+    assert "RuntimeError: child boom" in stderr
+
+
+@pytest.mark.anyio
+async def test_repeated_bad_children_do_not_pollute_or_replace_parent_stdio(
+    capfd: CaptureFixture[str],
+) -> None:
+    stdout_object = sys.stdout
+    stderr_object = sys.stderr
+    stdout_identity = _fd_identity(1)
+    stderr_identity = _fd_identity(2)
+
+    cases: tuple[tuple[Callable[..., object], tuple[object, ...]], ...] = (
+        (_raise_after_stderr_write, ()),
+        (_exit_after_stdio_write, ("exit-child", 17)),
+        (_abort_after_stdio_write, ("abort-child",)),
+    )
+
+    for iteration in range(3):
+        for target, args in cases:
+            exitcode, stdout, stderr = await _run_and_collect(
+                target,
+                args=args,
+            )
+
+            assert exitcode != 0
+            if target is _exit_after_stdio_write:
+                assert stdout == b"exit-child: stdout before _exit\n"
+                assert stderr == b"exit-child: stderr before _exit\n"
+            elif target is _abort_after_stdio_write:
+                assert b"abort-child: stdout before abort\n" in stdout
+                assert b"abort-child: stderr before abort\n" in stderr
+                assert exitcode == -signal.SIGABRT
+            else:
+                assert stdout == b""
+                assert b"stderr before exception\n" in stderr
+                assert b"RuntimeError: child boom" in stderr
+
+        print(f"parent stdout still works {iteration}")
+        print(f"parent stderr still works {iteration}", file=sys.stderr)
+
+    parent_output = capfd.readouterr()
+
+    assert sys.stdout is stdout_object
+    assert sys.stderr is stderr_object
+    assert _fd_identity(1) == stdout_identity
+    assert _fd_identity(2) == stderr_identity
+    assert "parent stdout still works 0" in parent_output.out
+    assert "parent stdout still works 2" in parent_output.out
+    assert "parent stderr still works 0" in parent_output.err
+    assert "parent stderr still works 2" in parent_output.err
+    assert "exit-child:" not in parent_output.out
+    assert "exit-child:" not in parent_output.err
+    assert "abort-child:" not in parent_output.out
+    assert "abort-child:" not in parent_output.err
+    assert "child boom" not in parent_output.err
+
+
+@pytest.mark.anyio
+async def test_child_can_close_stdio_without_corrupting_parent_stdio(
+    capfd: CaptureFixture[str],
+) -> None:
+    stdout_identity = _fd_identity(1)
+    stderr_identity = _fd_identity(2)
+
+    exitcode, stdout, stderr = await _run_and_collect(_close_stdio_and_exit)
+    os.write(1, b"parent stdout after child closed stdio\n")
+    os.write(2, b"parent stderr after child closed stdio\n")
+    parent_output = capfd.readouterr()
+
+    assert exitcode == 0
+    assert stdout == b""
+    assert stderr == b""
+    assert _fd_identity(1) == stdout_identity
+    assert _fd_identity(2) == stderr_identity
+    assert "parent stdout after child closed stdio" in parent_output.out
+    assert "parent stderr after child closed stdio" in parent_output.err
+
+
+@pytest.mark.anyio
+async def test_repeated_crashing_children_do_not_grow_parent_fd_table() -> None:
+    await _run_and_collect(_exit_after_stdio_write, args=("warmup", 23))
+    before = _fd_count()
+    if before is None:
+        pytest.skip("fd table count is not available on this platform")
+
+    for iteration in range(20):
+        exitcode, stdout, stderr = await _run_and_collect(
+            _exit_after_stdio_write,
+            args=(f"fd-child-{iteration}", 31),
+        )
+
+        assert exitcode == 31
+        assert stdout == f"fd-child-{iteration}: stdout before _exit\n".encode()
+        assert stderr == f"fd-child-{iteration}: stderr before _exit\n".encode()
+
+    after = _fd_count()
+    assert after is not None
+    assert after <= before + 2
+
+
+@pytest.mark.anyio
+async def test_stop_allows_child_to_exit_after_sigterm() -> None:
+    process = AsyncProcess(_exit_on_sigterm, args=(43,))
+
+    async with _started_process(process):
+        assert await process.stdout.receive() == b"sigterm-ready\n"
+
+        with fail_after(2):
+            await process.stop()
+
+    assert process.exitcode == 43
+
+
+@pytest.mark.anyio
+async def test_stop_retries_sigterm_before_sigkill(monkeypatch: MonkeyPatch) -> None:
+    monkeypatch.setattr(async_process, "_TERMINATE_GRACE_SECONDS", 0.01)
+    monkeypatch.setattr(async_process, "_TERMINATE_RETRY_GRACE_SECONDS", 0.01)
+    process = AsyncProcess(_exit_after_repeated_sigterm, args=(3, 44))
+
+    async with _started_process(process):
+        assert await process.stdout.receive() == b"sigterm-ready\n"
+
+        with fail_after(2):
+            await process.stop()
+
+    assert process.exitcode == 44
+
+
+@pytest.mark.anyio
+async def test_stop_escalates_to_sigkill_when_child_ignores_sigterm(
+    monkeypatch: MonkeyPatch,
+) -> None:
+    monkeypatch.setattr(async_process, "_TERMINATE_GRACE_SECONDS", 0.1)
+    monkeypatch.setattr(async_process, "_TERMINATE_RETRY_GRACE_SECONDS", 0.01)
+    process = AsyncProcess(_ignore_sigterm_forever)
+
+    async with _started_process(process):
+        assert await process.stdout.receive() == b"sigterm-ready\n"
+
+        with fail_after(3):
+            await process.stop()
+
+    assert process.exitcode == -signal.SIGKILL
+
+
+@pytest.mark.anyio
+async def test_process_can_use_mp_channel_with_global_spawn_context() -> None:
+    send, recv = mp_channel[str]()
+    process = AsyncProcess(_send_over_mp_channel, args=(send,))
+
+    async with _started_process(process):
+        with fail_after(2):
+            assert await recv.receive_async() == "hello from child"
+            assert await process.wait() == 0
+
+    with contextlib.suppress(Exception):
+        recv.close()
+
+
+@pytest.mark.anyio
+@pytest.mark.skip(reason="manual MLX OOM isolation check")
+async def test_death(capsys: CaptureFixture[str]) -> None:
+    with capsys.disabled():
+        process = AsyncProcess(_mlx_force_oom)
+        stdout = b""
+        stderr = b""
+        async with _started_process(process):
+            _, stdout, stderr = await _collect_process_output(process)
+
+        print("PARENT: done")
+
+        print("CHILD out:", stdout.decode("utf-8", errors="replace"))
+        print("CHILD err:", stderr.decode("utf-8", errors="replace"), "hello :)")

src/exo/utils/tests/test_daemon.py

@@ -0,0 +1,168 @@
+import contextlib
+import os
+from collections.abc import AsyncIterator
+
+import anyio
+import pytest
+from anyio import EndOfStream, create_task_group, fail_after
+
+from exo.utils.async_process import AsyncProcess
+from exo.utils.channels import MpReceiver, MpSender, Receiver, mp_channel
+from exo.utils.daemon import detach_stdio_to_devnull
+
+
+def _write_before_and_after_detach() -> None:
+    os.write(1, b"before stdout\n")
+    os.write(2, b"before stderr\n")
+    detach_stdio_to_devnull()
+    os.write(1, b"after stdout\n")
+    os.write(2, b"after stderr\n")
+
+
+def _write_grandchild_stdio(label: str) -> None:
+    os.write(1, f"{label} stdout\n".encode())
+    os.write(2, f"{label} stderr\n".encode())
+
+
+async def _spawn_grandchild_and_report(
+    result_sender: MpSender[tuple[int, bytes, bytes]],
+    label: str,
+) -> None:
+    result_sender.send(await _collect_spawned_child(label))
+    result_sender.close()
+
+
+async def _collect_spawned_child(label: str) -> tuple[int, bytes, bytes]:
+    process = AsyncProcess(_write_grandchild_stdio, args=(label,))
+    async with _started_process(process):
+        return await _collect_process_output(process)
+
+
+def _detach_stdio_then_spawn_captured_child(
+    result_sender: MpSender[tuple[int, bytes, bytes]],
+) -> None:
+    detach_stdio_to_devnull()
+    anyio.run(_spawn_grandchild_and_report, result_sender, "grandchild")
+
+
+def _detach_stdio_then_spawn_captured_children_sequentially(
+    result_sender: MpSender[list[tuple[int, bytes, bytes]]],
+) -> None:
+    async def run_children() -> list[tuple[int, bytes, bytes]]:
+        results: list[tuple[int, bytes, bytes]] = []
+        for index in range(5):
+            results.append(await _collect_spawned_child(f"grandchild-{index}"))
+        return results
+
+    detach_stdio_to_devnull()
+    result_sender.send(anyio.run(run_children))
+    result_sender.close()
+
+
+async def _collect_stream(stream: Receiver[bytes], output: bytearray) -> None:
+    while True:
+        try:
+            output.extend(await stream.receive())
+        except EndOfStream:
+            return
+
+
+async def _collect_process_output(
+    process: AsyncProcess,
+) -> tuple[int, bytes, bytes]:
+    stdout = bytearray()
+    stderr = bytearray()
+    exitcodes: list[int] = []
+
+    async with create_task_group() as collect_group:
+        collect_group.start_soon(_collect_stream, process.stdout, stdout)
+        collect_group.start_soon(_collect_stream, process.stderr, stderr)
+        exitcodes.append(await process.wait())
+
+    if not exitcodes:
+        raise RuntimeError("process exited without a return code")
+    return exitcodes[0], bytes(stdout), bytes(stderr)
+
+
+@contextlib.asynccontextmanager
+async def _started_process(process: AsyncProcess) -> AsyncIterator[None]:
+    async with create_task_group() as task_group:
+        await task_group.start(process.run)
+        try:
+            yield
+        finally:
+            await process.stop()
+
+
+async def _run_process_and_receive[T](
+    process: AsyncProcess,
+    recv: MpReceiver[T],
+    *,
+    timeout: float,
+) -> tuple[int, T]:
+    async with _started_process(process):
+        with fail_after(timeout):
+            result = await recv.receive_async()
+            exitcode = await process.wait()
+
+    return exitcode, result
+
+
+@pytest.mark.anyio
+async def test_detach_stdio_to_devnull_redirects_stdio_away_from_capture() -> None:
+    process = AsyncProcess(_write_before_and_after_detach)
+
+    async with _started_process(process):
+        exitcode, stdout, stderr = await _collect_process_output(process)
+
+    assert exitcode == 0
+    assert stdout == b"before stdout\n"
+    assert stderr == b"before stderr\n"
+
+
+@pytest.mark.anyio
+async def test_detached_stdio_process_can_spawn_and_capture_child_stdio() -> None:
+    send, recv = mp_channel[tuple[int, bytes, bytes]]()
+    process = AsyncProcess(_detach_stdio_then_spawn_captured_child, args=(send,))
+
+    try:
+        daemonized_parent_exitcode, result = await _run_process_and_receive(
+            process, recv, timeout=5
+        )
+    finally:
+        recv.close()
+
+    child_exitcode, child_stdout, child_stderr = result
+
+    assert daemonized_parent_exitcode == 0
+    assert child_exitcode == 0
+    assert child_stdout == b"grandchild stdout\n"
+    assert child_stderr == b"grandchild stderr\n"
+
+
+@pytest.mark.anyio
+async def test_detached_stdio_process_can_spawn_captured_children_sequentially() -> (
+    None
+):
+    send, recv = mp_channel[list[tuple[int, bytes, bytes]]]()
+    process = AsyncProcess(
+        _detach_stdio_then_spawn_captured_children_sequentially,
+        args=(send,),
+    )
+
+    try:
+        daemonized_parent_exitcode, results = await _run_process_and_receive(
+            process, recv, timeout=10
+        )
+    finally:
+        recv.close()
+
+    assert daemonized_parent_exitcode == 0
+    assert results == [
+        (
+            0,
+            f"grandchild-{index} stdout\n".encode(),
+            f"grandchild-{index} stderr\n".encode(),
+        )
+        for index in range(5)
+    ]

src/exo/utils/tests/test_pidfile.py

@@ -0,0 +1,84 @@
+from __future__ import annotations
+
+import gc
+import os
+import subprocess
+import sys
+import textwrap
+from pathlib import Path
+from typing import Final
+
+import pytest
+
+import exo.utils.pidfile as pidfile
+from exo.utils.pidfile import acquire_exo_pidfile
+
+_CHILD_ACQUIRE_PIDFILE_SCRIPT: Final = textwrap.dedent(
+    """
+    import sys
+    from pathlib import Path
+    from unittest.mock import patch
+
+    import exo.utils.pidfile as pidfile
+    from exo.utils.pidfile import PidfileLockError, acquire_exo_pidfile
+
+    with patch.object(pidfile, "EXO_PID_FILE", Path(sys.argv[1])):
+        try:
+            handle = acquire_exo_pidfile()
+        except PidfileLockError as exception:
+            print(str(exception))
+            raise SystemExit(73) from exception
+
+        del handle
+    """
+)
+
+
+def _use_pidfile_path(monkeypatch: pytest.MonkeyPatch, path: Path) -> None:
+    monkeypatch.setattr(pidfile, "EXO_PID_FILE", path)
+
+
+def _run_child_acquire_pidfile(path: Path) -> subprocess.CompletedProcess[str]:
+    return subprocess.run(
+        [sys.executable, "-c", _CHILD_ACQUIRE_PIDFILE_SCRIPT, str(path)],
+        check=False,
+        capture_output=True,
+        text=True,
+    )
+
+
+def test_acquire_exo_pidfile_writes_current_pid_and_removes_on_drop(
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    path = tmp_path / "exo.pid"
+    _use_pidfile_path(monkeypatch, path)
+
+    handle = acquire_exo_pidfile()
+    assert path.read_text() == str(os.getpid())
+
+    del handle
+    gc.collect()
+
+    assert not path.exists()
+
+
+def test_acquire_exo_pidfile_rejects_second_process(
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    path = tmp_path / "exo.pid"
+    _use_pidfile_path(monkeypatch, path)
+
+    handle = acquire_exo_pidfile()
+    try:
+        blocked_child = _run_child_acquire_pidfile(path)
+        assert blocked_child.returncode == 73
+        assert "Failed to acquire EXO pidfile" in blocked_child.stdout
+    finally:
+        del handle
+        gc.collect()
+
+    unblocked_child = _run_child_acquire_pidfile(path)
+    assert unblocked_child.returncode == 0
+    assert unblocked_child.stdout == ""

src/exo/utils/tests/test_power_sampler.py

@@ -111,6 +111,36 @@ async def test_empty_state() -> None:
     assert result.total_energy_joules == 0.0
 
 
+def test_trapezoidal_unit_dt_weighting() -> None:
+    """Pure unit test on the integration helper. Crafted samples where the
+    arithmetic mean is wildly wrong vs the time-weighted result."""
+    from exo.utils.power_sampler import trapezoidal_energy
+
+    # 5 s window. Power = 10 W for the first 4.9 s, then 100 W for the last 0.1 s.
+    # Three samples: t=0 W=10, t=4.9 W=10, t=5.0 W=100.
+    samples = [
+        (0.0, _make_profile(10.0)),
+        (4.9, _make_profile(10.0)),
+        (5.0, _make_profile(100.0)),
+    ]
+    energy = trapezoidal_energy(samples, elapsed=5.0)
+    # (10+10)/2 * 4.9 + (10+100)/2 * 0.1 = 49 + 5.5 = 54.5 J
+    assert abs(energy - 54.5) < 1e-9
+    avg = energy / 5.0  # 10.9 W
+    # Arithmetic mean of the three samples would be (10+10+100)/3 ≈ 40 W.
+    # Trapezoidal correctly weights each segment by its dt.
+    assert abs(avg - 10.9) < 1e-9
+
+
+def test_trapezoidal_unit_single_sample() -> None:
+    """One sample: no window to integrate over, so fall back to constant power
+    over the elapsed duration."""
+    from exo.utils.power_sampler import trapezoidal_energy
+
+    samples = [(0.0, _make_profile(42.0))]
+    assert trapezoidal_energy(samples, elapsed=3.0) == 42.0 * 3.0
+
+
 async def test_result_stops_sampling() -> None:
     """Calling result() should stop the sampler's run loop."""
     state: dict[NodeId, SystemPerformanceProfile] = {

src/exo/worker/engines/image/builder.py

@@ -143,6 +143,7 @@ class ImageEngine(Engine):
         Generator[tuple[TaskId, Chunk | FinishedResponse | CancelledResponse]] | None
     ) = field(init=False, default=None)
     queue: deque[ImageTask] = field(init=False, default_factory=deque)
+    _cancelled_tasks: set[TaskId] = field(init=False, default_factory=set)
 
     def warmup(self) -> None:
         image = warmup_image_generator(model=self.image_model)
@@ -168,7 +169,11 @@ class ImageEngine(Engine):
             task = self.queue.popleft()
             self.current_gen = self._run_image_task(task.task_id, task.task_params)
             resp = next(self.current_gen, None)
-        return (resp,) if resp is not None else ()
+        return (
+            (resp,)
+            if resp is not None and _is_primary_output_node(self.shard_metadata)
+            else ()
+        )
 
     def close(self) -> None:
         with contextlib.suppress(NameError, AttributeError):

src/exo/worker/engines/mlx/utils_mlx.py

@@ -115,7 +115,8 @@ def mlx_distributed_init(
 
                 os.environ["MLX_HOSTFILE"] = coordination_file
                 os.environ["MLX_RANK"] = str(rank)
-                os.environ["MLX_RING_VERBOSE"] = "1"
+                # os.environ["MLX_RING_VERBOSE"] = "1"  # NOTE: we don't use it enough to care (turn on again if need to)
+
                 group = mx.distributed.init(backend="ring", strict=True)
 
             case MlxJacclInstance(

src/exo/worker/main.py

@@ -10,7 +10,7 @@ from exo.api.types import ImageEditsTaskParams
 from exo.download.download_utils import is_read_only_model_dir, resolve_existing_model
 from exo.shared.apply import apply
 from exo.shared.constants import EXO_MAX_INSTANCE_RETRIES
-from exo.shared.models.model_cards import ModelId, add_to_card_cache, delete_custom_card
+from exo.shared.models.model_cards import ModelId, card_cache
 from exo.shared.types.chunks import InputImageChunk
 from exo.shared.types.commands import (
     DeleteInstance,
@@ -20,8 +20,6 @@ from exo.shared.types.commands import (
 )
 from exo.shared.types.common import CommandId, NodeId, SystemId
 from exo.shared.types.events import (
-    CustomModelCardAdded,
-    CustomModelCardDeleted,
     Event,
     IndexedEvent,
     InputChunkReceived,
@@ -110,6 +108,8 @@ class Worker:
                 tg.start_soon(self.plan_step)
                 tg.start_soon(self._event_applier)
                 tg.start_soon(self._poll_connection_updates)
+                tg.start_soon(self._reconcile_custom_cards)
+
         finally:
             # Actual shutdown code - waits for all tasks to complete before executing.
             logger.info("Stopping Worker")
@@ -151,7 +151,6 @@ class Worker:
                     self.input_chunk_buffer[cmd_id][event.chunk.chunk_index] = (
                         event.chunk
                     )
-
                     if (
                         len(self.input_chunk_buffer[cmd_id])
                         == self.input_chunk_counts[cmd_id]
@@ -172,12 +171,18 @@ class Worker:
                                 )
                             ] = img
 
-                if isinstance(event, CustomModelCardAdded):
-                    await event.model_card.save_to_custom_dir()
-                    add_to_card_cache(event.model_card)
+    async def _reconcile_custom_cards(self) -> None:
+        while True:
+            await anyio.sleep(1)
+            target = dict(self.state.custom_model_cards)
+            for model_id, card in target.items():
+                if card_cache.get(model_id) == card:
+                    continue
+                await card_cache.save(card)
 
-                if isinstance(event, CustomModelCardDeleted):
-                    await delete_custom_card(event.model_id)
+            for card in await card_cache.list_all():
+                if card.model_id not in target:
+                    await card_cache.pop(card.model_id)
 
     async def plan_step(self):
         while True:

src/exo/worker/runner/llm_inference/batch_generator.py

@@ -138,8 +138,10 @@ class SequentialGenerator(Engine):
     def agree_on_tasks(self) -> None:
         """Agree between all ranks about the task ordering (some may have received in different order or not at all)."""
         agreed, different = mx_all_gather_tasks(self._maybe_queue, self.group)
-        self._queue.extend(task for task in self._maybe_queue if task in agreed)
-        self._maybe_queue = [task for task in self._maybe_queue if task in different]
+        # Extend from `agreed` (sorted by task_id on all ranks) to guarantee every
+        # rank enqueues tasks in the same order, preventing TP collective deadlocks.
+        self._queue.extend(agreed)
+        self._maybe_queue = list(different)
 
     def agree_on_cancellations(self) -> None:
         """Agree between all ranks about which tasks to cancel."""
@@ -197,9 +199,14 @@ class SequentialGenerator(Engine):
             self._active = None
             raise
 
-        return itertools.chain(
-            output,
-            map(lambda task: (task, CancelledResponse()), self._cancelled_tasks),
+        return filter(
+            lambda chunk: (
+                not isinstance(chunk[1], GenerationChunk) or self.device_rank == 0
+            ),
+            itertools.chain(
+                output,
+                map(lambda task: (task, CancelledResponse()), self._cancelled_tasks),
+            ),
         )
 
     def _start_next(self) -> None:
@@ -368,8 +375,10 @@ class BatchGenerator(Engine):
     def agree_on_tasks(self) -> None:
         """Agree between all ranks about the task ordering (some may have received in different order or not at all)."""
         agreed, different = mx_all_gather_tasks(self._maybe_queue, self.group)
-        self._queue.extend(task for task in self._maybe_queue if task in agreed)
-        self._maybe_queue = [task for task in self._maybe_queue if task in different]
+        # Extend from `agreed` (sorted by task_id on all ranks) to guarantee every
+        # rank enqueues tasks in the same order, preventing TP collective deadlocks.
+        self._queue.extend(agreed)
+        self._maybe_queue = list(different)
 
     def agree_on_cancellations(self) -> None:
         """Agree between all ranks about which tasks to cancel."""
@@ -449,7 +458,12 @@ class BatchGenerator(Engine):
                 output.append((task.task_id, FinishedResponse()))
                 del self._active_tasks[uid]
 
-        return itertools.chain(output, self._apply_cancellations())
+        return filter(
+            lambda chunk: (
+                not isinstance(chunk[1], GenerationChunk) or self.device_rank == 0
+            ),
+            itertools.chain(output, self._apply_cancellations()),
+        )
 
     def _apply_cancellations(
         self,

src/exo/worker/runner/runner.py

@@ -390,5 +390,5 @@ class Runner:
         chunk: Chunk,
         command_id: CommandId,
     ):
-        if self.device_rank == 0:
-            self.event_sender.send(ChunkGenerated(command_id=command_id, chunk=chunk))
+        assert isinstance(self.generator, Engine)
+        self.event_sender.send(ChunkGenerated(command_id=command_id, chunk=chunk))

src/exo/worker/runner/supervisor.py

@@ -1,5 +1,4 @@
 import contextlib
-import multiprocessing as mp
 import signal
 from dataclasses import dataclass, field
 from typing import Self
@@ -8,7 +7,7 @@ import anyio
 from anyio import (
     BrokenResourceError,
     ClosedResourceError,
-    to_thread,
+    EndOfStream,
 )
 from loguru import logger
 
@@ -41,7 +40,8 @@ from exo.shared.types.worker.runners import (
     RunnerWarmingUp,
 )
 from exo.shared.types.worker.shards import ShardMetadata
-from exo.utils.channels import MpReceiver, MpSender, Sender, mp_channel
+from exo.utils.async_process import AsyncProcess
+from exo.utils.channels import MpReceiver, MpSender, Receiver, Sender, mp_channel
 from exo.utils.task_group import TaskGroup
 from exo.worker.runner.bootstrap import entrypoint
 
@@ -53,7 +53,7 @@ DECODE_TIMEOUT_SECONDS = 5
 class RunnerSupervisor:
     shard_metadata: ShardMetadata
     bound_instance: BoundInstance
-    runner_process: mp.Process
+    runner_process: AsyncProcess
     initialize_timeout: float
     _ev_recv: MpReceiver[Event]
     _task_sender: MpSender[Task]
@@ -81,7 +81,7 @@ class RunnerSupervisor:
         task_sender, task_recv = mp_channel[Task]()
         cancel_sender, cancel_recv = mp_channel[TaskId]()
 
-        runner_process = mp.Process(
+        runner_process = AsyncProcess(
             target=entrypoint,
             args=(
                 bound_instance,
@@ -109,9 +109,25 @@ class RunnerSupervisor:
         return self
 
     async def run(self):
-        self.runner_process.start()
         try:
             async with self._tg as tg:
+                # start the process itself
+                await tg.start(self.runner_process.run)
+
+                # start tasks to drain/collect stdout/stderr into usable errors
+                #
+                # TODO: right now it logs them as warnings, but in the future they should be split
+                #       into being logged AND a seperate task which tries to best-effort figure out cause
+                #       of error and package into error enum, which then is used by rest of app to act on it;
+                #       inferring what the error is would be done by pattern-matching in the text for things
+                #       e.g. certain VLLM error codes and so on
+                tg.start_soon(
+                    self._forward_runner_output, "stdout", self.runner_process.stdout
+                )
+                tg.start_soon(
+                    self._forward_runner_output, "stderr", self.runner_process.stderr
+                )
+
                 tg.start_soon(self._watch_runner)
                 tg.start_soon(self._forward_events)
         finally:
@@ -129,41 +145,11 @@ class RunnerSupervisor:
             with contextlib.suppress(ClosedResourceError):
                 self._cancel_sender.close()
 
-            await to_thread.run_sync(self.runner_process.join, 5)
-
-            if self.runner_process.is_alive():
-                logger.warning(
-                    "Runner process didn't shutdown succesfully, terminating"
+            with anyio.CancelScope(shield=True):
+                await self.runner_process.stop()
+                logger.info(
+                    f"Runner process successfully terminated: {self.runner_process.exitcode}"
                 )
-                self.runner_process.terminate()
-                self.runner_process.join(timeout=10)
-
-                if not self.runner_process.is_alive():
-                    logger.warning("Terminated nicely in the first attempt!")
-
-                else:
-                    # Try really hard to terminate
-                    for i in range(2, 11):
-                        self.runner_process.terminate()
-                        self.runner_process.join(timeout=2)
-                        if not self.runner_process.is_alive():
-                            logger.warning(f"That took {i} attempts :)")
-                            break
-                    # Try even harder to kill
-                    else:
-                        logger.critical(
-                            "Runner process didn't respond to SIGTERM, killing"
-                        )
-                        j = 0
-                        while self.runner_process.is_alive():
-                            j += 1
-                            self.runner_process.kill()
-                            self.runner_process.join(timeout=5)
-                            logger.warning(f"That took {j} attempts :(")
-            else:
-                logger.info("Runner process succesfully terminated")
-
-            self.runner_process.close()
 
     def shutdown(self):
         self._tg.cancel_tasks()
@@ -249,13 +235,33 @@ class RunnerSupervisor:
                 if not self.runner_process.is_alive():
                     await self._check_runner(RuntimeError("Runner found to be dead"))
 
+    async def _forward_runner_output(
+        self,
+        stream_name: str,
+        stream: Receiver[bytes],
+    ) -> None:
+        while True:
+            try:
+                chunk = await stream.receive()
+            except (EndOfStream, ClosedResourceError, BrokenResourceError):
+                return
+
+            message = chunk.decode("utf-8", errors="replace").rstrip()
+            if not message:
+                continue
+            if stream_name == "stderr":
+                logger.warning(f"Runner stderr: {message}")
+            else:
+                logger.debug(f"Runner stdout: {message}")
+
     async def _check_runner(self, e: Exception) -> None:
         if not self._cancel_watch_runner.cancel_called:
             self._cancel_watch_runner.cancel()
         logger.info("Checking runner's status")
         if self.runner_process.is_alive():
-            logger.info("Runner was found to be alive, attempting to join process")
-            await to_thread.run_sync(self.runner_process.join, 5)
+            logger.info("Runner was found to be alive, stopping process")
+            with anyio.CancelScope(shield=True):
+                await self.runner_process.stop()
         rc = self.runner_process.exitcode
         logger.info(f"Runner exited with exit code {rc}")
         if rc == 0:

src/exo/worker/tests/unittests/test_mlx/test_tokenizers.py

@@ -16,7 +16,7 @@ from exo.download.download_utils import (
     fetch_file_list_with_cache,
     resolve_model_dir,
 )
-from exo.shared.models.model_cards import ModelCard, ModelId, get_model_cards
+from exo.shared.models.model_cards import ModelCard, ModelId, card_cache
 from exo.worker.engines.mlx.utils_mlx import (
     get_eos_token_ids_for_model,
     load_tokenizer_for_model_id,
@@ -76,7 +76,7 @@ def get_test_models() -> list[ModelCard]:
     """Get a representative sample of models to test."""
     # Pick one model from each family to test
     families: dict[str, ModelCard] = {}
-    for card in asyncio.run(get_model_cards()):
+    for card in asyncio.run(card_cache.list_all()):
         # Extract family name (e.g., "llama-3.1" from "llama-3.1-8b")
         parts = card.model_id.short().split("-")
         family = "-".join(parts[:2]) if len(parts) >= 2 else parts[0]
@@ -298,7 +298,7 @@ async def test_tokenizer_special_tokens(model_card: ModelCard) -> None:
 async def test_kimi_tokenizer_specifically():
     """Test Kimi tokenizer with its specific patches and quirks."""
     kimi_models = [
-        card for card in await get_model_cards() if "kimi" in card.model_id.lower()
+        card for card in await card_cache.list_all() if "kimi" in card.model_id.lower()
     ]
 
     if not kimi_models:
@@ -350,7 +350,7 @@ async def test_glm_tokenizer_specifically():
 
     glm_model_cards = [
         card
-        for card in await get_model_cards()
+        for card in await card_cache.list_all()
         if contains(card, "glm")
         and not contains(card, "-5")
         and not contains(card, "4.7")

src/exo/worker/tests/unittests/test_runner/test_runner_supervisor.py

@@ -1,4 +1,3 @@
-import multiprocessing as mp
 from typing import cast
 
 import anyio
@@ -16,6 +15,7 @@ from exo.shared.types.text_generation import (
 )
 from exo.shared.types.worker.instances import BoundInstance, InstanceId
 from exo.shared.types.worker.runners import RunnerFailed, RunnerId
+from exo.utils.async_process import AsyncProcess
 from exo.utils.channels import channel, mp_channel
 from exo.worker.runner.supervisor import RunnerSupervisor
 from exo.worker.tests.unittests.conftest import get_bound_mlx_ring_instance
@@ -24,23 +24,11 @@ from exo.worker.tests.unittests.conftest import get_bound_mlx_ring_instance
 class _DeadProcess:
     exitcode = -6
 
-    def start(self) -> None:
-        return None
-
     def is_alive(self) -> bool:
         return False
 
-    def join(self, _timeout: float | None = None) -> None:
-        return None
 
-    def terminate(self) -> None:
-        return None
-
-    def kill(self) -> None:
-        return None
-
-
-@pytest.mark.asyncio
+@pytest.mark.anyio
 async def test_check_runner_emits_error_chunk_for_inflight_text_generation() -> None:
     event_sender, event_receiver = channel[Event]()
     task_sender, _ = mp_channel[Task]()
@@ -57,7 +45,7 @@ async def test_check_runner_emits_error_chunk_for_inflight_text_generation() ->
     supervisor = RunnerSupervisor(
         shard_metadata=bound_instance.bound_shard,
         bound_instance=bound_instance,
-        runner_process=cast("mp.Process", cast(object, _DeadProcess())),
+        runner_process=cast(AsyncProcess, cast(object, _DeadProcess())),
         initialize_timeout=400,
         _ev_recv=ev_recv,
         _task_sender=task_sender,

tests/conftest.py

@@ -0,0 +1,181 @@
+# type: ignore
+"""Pytest configuration for marker-driven exo integration tests.
+
+Test authors declare requirements via markers:
+
+    @pytest.mark.cluster(count=2, thunderbolt='a2a')
+    @pytest.mark.instance('mlx-community/Llama-3.2-1B-Instruct-4bit',
+                          sharding='tensor', comm='jaccl')
+    def test_jaccl_inference(session):
+        resp = session.chat('What is 2+2?')
+        assert '4' in resp
+
+Clusters are cached by `ClusterSpec`; tests with the same cluster_spec
+share a deployment. Each test places its own instance (matching its
+`@pytest.mark.instance`), and instances are cleaned up after the test.
+
+Run with:
+    uv run pytest tests/ -v
+    uv run pytest tests/ -v --hosts s2,s4,s9,s10
+"""
+
+from __future__ import annotations
+
+import contextlib
+import json
+
+import pytest
+from exo_tools.cluster import ClusterInfo, EcoSession
+from exo_tools.harness import cleanup_all_instances, place_instance
+
+from .framework import (
+    ClusterSpec,
+    Session,
+    parse_cluster_marker,
+    parse_instance_marker,
+)
+
+# Single eco session for the entire test process.
+eco = EcoSession(user_prefix="test")
+
+# Cluster cache keyed by ClusterSpec — tests with the same spec share a deployment.
+# Cleared at session teardown.
+_cluster_cache: dict[ClusterSpec, ClusterInfo] = {}
+
+
+def pytest_addoption(parser):
+    parser.addoption(
+        "--hosts",
+        default=None,
+        help="Comma-separated list of hosts (e.g. s2,s4,s9,s10). "
+        "Overrides constraint-based reservation.",
+    )
+
+
+def pytest_configure(config):
+    """Register custom markers."""
+    config.addinivalue_line(
+        "markers",
+        "cluster(count=N, thunderbolt=Thunderbolt|None, min_memory=GB, chip=PATTERN): "
+        "declare cluster requirements for a test",
+    )
+    config.addinivalue_line(
+        "markers",
+        "instance(model_id, sharding=Sharding, comm=Comm, min_nodes=N): "
+        "declare instance placement for a test",
+    )
+
+
+def pytest_report_header(config):
+    """Show the eco user and hosts for this test session."""
+    hosts = config.getoption("--hosts")
+    lines = [f"eco user: {eco.user}"]
+    if hosts:
+        lines.append(f"hosts override: {hosts}")
+    return lines
+
+
+@pytest.fixture(scope="session")
+def _host_pool(request) -> list[str] | None:
+    raw = request.config.getoption("--hosts")
+    if raw:
+        return [h.strip() for h in raw.split(",") if h.strip()]
+    return None
+
+
+@pytest.fixture
+def session(request, _host_pool) -> Session:
+    """Per-test fixture providing a Session matching the test's markers.
+
+    Reads @pytest.mark.cluster and @pytest.mark.instance from the test, deploys
+    a matching cluster (cached across tests with the same spec), places the
+    model, and yields a Session for the test to interact with. Cleans up the
+    instance after the test, and invalidates the cluster cache if the test
+    left nodes disconnected.
+    """
+    cluster_marker = request.node.get_closest_marker("cluster")
+    instance_marker = request.node.get_closest_marker("instance")
+
+    cluster_spec = parse_cluster_marker(cluster_marker)
+    instance_spec = parse_instance_marker(instance_marker)
+
+    # Deploy or reuse a cluster matching the spec
+    cluster = _cluster_cache.get(cluster_spec)
+    if cluster is None:
+        if _host_pool:
+            cluster = eco.start_deploy(
+                hosts=_host_pool[: cluster_spec.count], wait=True
+            )
+        else:
+            cluster = eco.start_deploy(
+                count=cluster_spec.count,
+                thunderbolt=cluster_spec.thunderbolt,
+                chip=cluster_spec.chip,
+                min_memory_gb=cluster_spec.min_memory_gb,
+                wait=True,
+            )
+        _cluster_cache[cluster_spec] = cluster
+
+    # Place an instance for this test if the test specified one
+    instance_id = None
+    if instance_spec is not None:
+        client = cluster.make_client()
+        instance_id = place_instance(
+            client,
+            instance_spec.model_id,
+            sharding=instance_spec.sharding,
+            comm=instance_spec.comm,
+            min_nodes=instance_spec.min_nodes,
+        )
+
+    sess = Session(
+        cluster=cluster,
+        eco=eco,
+        instance_spec=instance_spec,
+        instance_id=instance_id,
+    )
+
+    yield sess
+
+    # ---- Teardown ----
+
+    # If the test left nodes disconnected, invalidate the cluster cache and
+    # stop the cluster so the next test deploys fresh.
+    if sess._stopped_hosts:
+        _cluster_cache.pop(cluster_spec, None)
+        with contextlib.suppress(Exception):
+            eco.stop(sess.cluster.hosts)
+        return
+
+    # Otherwise, clean up any instances created during the test
+    with contextlib.suppress(Exception):
+        cleanup_all_instances(sess.client)
+
+
+# ---------------------------------------------------------------------------
+# Session-level teardown — stop all cached clusters
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture(scope="session", autouse=True)
+def _teardown_clusters():
+    yield
+    for cluster in _cluster_cache.values():
+        with contextlib.suppress(Exception):
+            eco.stop(cluster.hosts)
+    _cluster_cache.clear()
+
+
+def pytest_runtest_makereport(item, call):
+    """Attach cluster logs to the test report when a test fails."""
+    if call.when != "call" or call.excinfo is None:
+        return
+
+    sess = item.funcargs.get("session")
+    if sess is None:
+        return
+    try:
+        logs = eco.logs(sess.cluster.hosts, lines=200)
+        item.add_report_section("call", "Cluster Logs", json.dumps(logs, indent=2))
+    except Exception:
+        pass

tests/framework.py

@@ -0,0 +1,199 @@
+"""Marker-driven test framework for exo integration tests.
+
+Test authors declare requirements via markers:
+
+    @pytest.mark.cluster(count=2, thunderbolt='a2a')
+    @pytest.mark.instance('mlx-community/Llama-3.2-1B-Instruct-4bit',
+                          sharding='tensor', comm='jaccl')
+    def test_jaccl_inference(session):
+        resp = session.chat('What is 2+2?')
+        assert '4' in resp
+
+The `session` fixture reads the markers, deploys the cluster, places the
+instance, and provides a `Session` object. All cluster/instance orchestration
+lives in `exo_tools.harness`; this module is purely the pytest-facing layer.
+"""
+
+from __future__ import annotations
+
+import time
+from dataclasses import dataclass, field
+from typing import Any
+
+from exo_tools.client import ExoClient
+from exo_tools.cluster import (
+    Chip,
+    ClusterInfo,
+    EcoSession,
+    Thunderbolt,
+    make_client_from_url,
+)
+from exo_tools.harness import Comm, Sharding
+
+from exo.api.types.api import (
+    ChatCompletionChoice,
+    ChatCompletionRequest,
+    ChatCompletionResponse,
+)
+
+DEFAULT_MODEL = "mlx-community/Llama-3.2-1B-Instruct-4bit"
+
+
+def _extract_content(resp: ChatCompletionResponse) -> str:
+    """Extract plain-text content from a non-streaming chat completion."""
+    choice = resp.choices[0]
+    if not isinstance(choice, ChatCompletionChoice):
+        raise RuntimeError(
+            f"Expected non-streaming choice, got {type(choice).__name__}"
+        )
+    content = choice.message.content
+    if not isinstance(content, str):
+        raise RuntimeError(f"Expected string content, got {type(content).__name__}")
+    return content
+
+
+@dataclass(frozen=True)
+class ClusterSpec:
+    count: int = 1
+    thunderbolt: Thunderbolt | None = None
+    min_memory_gb: float | None = None
+    chip: Chip | None = None
+
+
+@dataclass(frozen=True)
+class InstanceSpec:
+    model_id: str
+    sharding: Sharding = Sharding.PIPELINE
+    comm: Comm = Comm.RING
+    min_nodes: int = 1
+
+
+def parse_cluster_marker(marker) -> ClusterSpec:
+    if marker is None:
+        return ClusterSpec()
+    return ClusterSpec(
+        count=marker.kwargs.get("count", 1),
+        thunderbolt=marker.kwargs.get("thunderbolt"),
+        min_memory_gb=marker.kwargs.get("min_memory"),
+        chip=marker.kwargs.get("chip"),
+    )
+
+
+def parse_instance_marker(marker) -> InstanceSpec | None:
+    if marker is None:
+        return None
+    if not marker.args:
+        raise ValueError(
+            "@pytest.mark.instance requires a positional model_id argument"
+        )
+    return InstanceSpec(
+        model_id=marker.args[0],
+        sharding=marker.kwargs.get("sharding", Sharding.PIPELINE),
+        comm=marker.kwargs.get("comm", Comm.RING),
+        min_nodes=marker.kwargs.get("min_nodes", 1),
+    )
+
+
+@dataclass
+class Session:
+    cluster: ClusterInfo
+    eco: EcoSession
+    instance_spec: InstanceSpec | None = None
+    instance_id: str | None = None
+    _stopped_hosts: set[str] = field(default_factory=set)
+
+    @property
+    def client(self) -> ExoClient:
+        for host in self.cluster.hosts:
+            if host not in self._stopped_hosts:
+                return make_client_from_url(self.cluster.api_endpoints[host])
+        return self.cluster.make_client()
+
+    @property
+    def state(self) -> dict[str, Any]:
+        return self.client.request_json("GET", "/state") or {}
+
+    @property
+    def instances(self) -> dict[str, Any]:
+        return self.state.get("instances", {})
+
+    # ---- Inference ----
+
+    def chat(self, prompt: str, max_tokens: int = 100) -> str:
+        resp = self.chat_raw(prompt, max_tokens=max_tokens)
+        return _extract_content(resp)
+
+    def chat_raw(self, prompt: str, **kwargs: Any) -> ChatCompletionResponse:
+        if not self.instance_spec:
+            raise RuntimeError(
+                "No instance placed; add @pytest.mark.instance to the test"
+            )
+        max_tokens = kwargs.pop("max_tokens", 100)
+        request = ChatCompletionRequest.model_validate(
+            {
+                "model": self.instance_spec.model_id,
+                "messages": [{"role": "user", "content": prompt}],
+                "max_tokens": max_tokens,
+                **kwargs,
+            }
+        )
+        return self._post_chat(request)
+
+    def multi_turn(self, messages: list[dict[str, str]], max_tokens: int = 100) -> str:
+        if not self.instance_spec:
+            raise RuntimeError(
+                "No instance placed; add @pytest.mark.instance to the test"
+            )
+        request = ChatCompletionRequest.model_validate(
+            {
+                "model": self.instance_spec.model_id,
+                "messages": messages,
+                "max_tokens": max_tokens,
+            }
+        )
+        return _extract_content(self._post_chat(request))
+
+    def _post_chat(self, request: ChatCompletionRequest) -> ChatCompletionResponse:
+        raw = self.client.request_json(
+            "POST",
+            "/v1/chat/completions",
+            body=request.model_dump(exclude_none=True),
+        )
+        return ChatCompletionResponse.model_validate(raw)
+
+    def disconnect_node(self, index: int) -> None:
+        """Stop exo on a node and wait for the cluster to observe the disconnect."""
+        host = self.cluster.hosts[index]
+        self.eco.stop([host], keep=True)
+        self._stopped_hosts.add(host)
+
+    def reconnect_node(self, index: int) -> None:
+        """Restart a previously disconnected node into the existing namespace."""
+        host = self.cluster.hosts[index]
+        self.eco.start_hosts([host], namespace=self.cluster.namespace)
+        self._stopped_hosts.discard(host)
+
+    def wait_ready(
+        self, expected_nodes: int | None = None, timeout: float = 60
+    ) -> None:
+        """Wait until the cluster has exactly `expected_nodes` visible and reporting memory.
+
+        Defaults to the count of non-stopped hosts. Use this after
+        `disconnect_node` / `reconnect_node` to wait for the cluster to settle.
+        """
+        if expected_nodes is None:
+            expected_nodes = len(self.cluster.hosts) - len(self._stopped_hosts)
+        start = time.time()
+        while time.time() - start < timeout:
+            try:
+                state = self.state
+                identities = len(state.get("nodeIdentities", {}))
+                memory = len(state.get("nodeMemory", {}))
+                if identities == expected_nodes and memory == expected_nodes:
+                    return
+            except Exception:
+                pass
+            time.sleep(2.0)
+        raise TimeoutError(
+            f"Cluster did not reach exactly {expected_nodes} ready nodes within {timeout}s"
+        )

tests/test_1node.py

@@ -0,0 +1,75 @@
+# type: ignore
+"""Single-node integration tests.
+
+Run with:
+    uv run pytest tests/test_1node.py -v
+"""
+
+from __future__ import annotations
+
+import time
+
+import pytest
+from exo_tools.harness import is_model_downloaded, place_instance
+
+from .framework import DEFAULT_MODEL, InstanceSpec
+
+
+@pytest.mark.cluster(count=1)
+@pytest.mark.instance(DEFAULT_MODEL)
+def test_place_instance_and_chat(session):
+    resp = session.chat("Say hello in one sentence.")
+    assert len(resp) > 0
+
+
+@pytest.mark.cluster(count=1)
+@pytest.mark.instance(DEFAULT_MODEL)
+def test_chat_multiple_turns(session):
+    first_reply = session.chat("What is 2 + 2?")
+    assert len(first_reply) > 0
+
+    second_reply = session.multi_turn(
+        [
+            {"role": "user", "content": "What is 2 + 2?"},
+            {"role": "assistant", "content": first_reply},
+            {"role": "user", "content": "Now multiply that by 3."},
+        ]
+    )
+    assert len(second_reply) > 0
+
+
+@pytest.mark.cluster(count=1)
+@pytest.mark.instance(DEFAULT_MODEL)
+def test_delete_instance(session):
+    from exo_tools.harness import wait_for_instance_gone
+
+    session.client.request_json("DELETE", f"/instance/{session.instance_id}")
+    wait_for_instance_gone(session.client, session.instance_id, timeout=30.0)
+    assert len(session.instances) == 0, (
+        f"Expected no instances, found {len(session.instances)}"
+    )
+
+
+@pytest.mark.cluster(count=1)
+def test_download_from_scratch(session):
+    """Ensure the model is not on the cluster, then place an instance to
+    trigger a fresh download and verify inference.
+    """
+    node_id = next(iter(session.state.get("nodeIdentities", {})))
+
+    # Delete any existing download — the API call is idempotent
+    session.client.request_json("DELETE", f"/download/{node_id}/{DEFAULT_MODEL}")
+
+    # Poll until the model is gone (it may already be gone)
+    deadline = time.time() + 60.0
+    while time.time() < deadline:
+        if not is_model_downloaded(session.client, DEFAULT_MODEL):
+            break
+        time.sleep(2.0)
+    else:
+        raise AssertionError(f"Expected {DEFAULT_MODEL} to be deleted from cluster")
+
+    place_instance(session.client, DEFAULT_MODEL, timeout=900.0)
+    session.instance_spec = InstanceSpec(model_id=DEFAULT_MODEL)
+    resp = session.chat("Say hello in one sentence.")
+    assert len(resp) > 0

tests/test_2node.py

@@ -0,0 +1,49 @@
+# type: ignore
+"""Two-node integration tests (ring + jaccl parallelism).
+
+Run with:
+    uv run pytest tests/test_2node.py -v
+"""
+
+from __future__ import annotations
+
+import pytest
+from exo_tools.cluster import Thunderbolt
+from exo_tools.harness import Comm, Sharding
+
+from .framework import DEFAULT_MODEL
+
+
+@pytest.mark.cluster(count=2, thunderbolt=Thunderbolt.A2A)
+@pytest.mark.instance(
+    DEFAULT_MODEL, sharding=Sharding.TENSOR, comm=Comm.JACCL, min_nodes=2
+)
+def test_2node_jaccl(session):
+    resp = session.chat("Say hello in one sentence.")
+    assert len(resp) > 0
+
+
+@pytest.mark.cluster(count=2, thunderbolt=Thunderbolt.A2A)
+@pytest.mark.instance(
+    DEFAULT_MODEL, sharding=Sharding.PIPELINE, comm=Comm.RING, min_nodes=2
+)
+def test_2node_ring(session):
+    resp = session.chat("Say hello in one sentence.")
+    assert len(resp) > 0
+
+
+@pytest.mark.cluster(count=2, thunderbolt=Thunderbolt.A2A)
+@pytest.mark.instance(
+    DEFAULT_MODEL, sharding=Sharding.TENSOR, comm=Comm.JACCL, min_nodes=2
+)
+def test_2node_jaccl_multi_turn(session):
+    first = session.chat("What is the capital of France?")
+    assert len(first) > 0
+    second = session.multi_turn(
+        [
+            {"role": "user", "content": "What is the capital of France?"},
+            {"role": "assistant", "content": first},
+            {"role": "user", "content": "What country is it in?"},
+        ]
+    )
+    assert len(second) > 0

tests/test_4node.py

@@ -0,0 +1,32 @@
+# type: ignore
+"""Four-node integration tests.
+
+Run with:
+    uv run pytest tests/test_4node.py -v
+"""
+
+from __future__ import annotations
+
+import pytest
+from exo_tools.cluster import Thunderbolt
+from exo_tools.harness import Comm, Sharding
+
+from .framework import DEFAULT_MODEL
+
+
+@pytest.mark.cluster(count=4, thunderbolt=Thunderbolt.A2A)
+@pytest.mark.instance(
+    DEFAULT_MODEL, sharding=Sharding.PIPELINE, comm=Comm.RING, min_nodes=4
+)
+def test_4node_pipeline_ring(session):
+    resp = session.chat("Say hello in one sentence.")
+    assert len(resp) > 0
+
+
+@pytest.mark.cluster(count=4, thunderbolt=Thunderbolt.A2A)
+@pytest.mark.instance(
+    DEFAULT_MODEL, sharding=Sharding.TENSOR, comm=Comm.JACCL, min_nodes=4
+)
+def test_4node_tensor_jaccl(session):
+    resp = session.chat("Say hello in one sentence.")
+    assert len(resp) > 0

tests/test_dashboard.py

@@ -0,0 +1,102 @@
+# type: ignore
+"""Dashboard end-to-end tests using Playwright (headless Chromium).
+
+Prerequisites:
+    uv run playwright install chromium
+
+Run with:
+    uv run pytest tests/test_dashboard.py -v
+"""
+
+from __future__ import annotations
+
+import contextlib
+
+import pytest
+
+try:
+    from playwright.sync_api import sync_playwright
+
+    _HAS_PLAYWRIGHT = True
+except ImportError:
+    _HAS_PLAYWRIGHT = False
+
+# Check if Chromium is installed by attempting a quick launch
+_HAS_CHROMIUM = False
+if _HAS_PLAYWRIGHT:
+    try:
+        with sync_playwright() as p:
+            browser = p.chromium.launch(headless=True)
+            browser.close()
+        _HAS_CHROMIUM = True
+    except Exception:
+        pass
+
+pytestmark = pytest.mark.skipif(
+    not _HAS_PLAYWRIGHT or not _HAS_CHROMIUM,
+    reason="playwright or chromium not installed (run: uv run playwright install chromium)",
+)
+
+
+def _mark_onboarding_complete(session) -> None:
+    """Mark onboarding complete on the server so the wizard doesn't auto-launch a model."""
+    with contextlib.suppress(Exception):
+        session.client.request_json("POST", "/onboarding")
+
+
+@pytest.mark.cluster(count=1)
+def test_dashboard_chat_inference(session):
+    """Full UI flow: open dashboard, pick a model, send a chat, verify response.
+
+    The instance is created via the dashboard UI (model picker → chat send
+    triggers the dashboard's auto-launch flow), not via @pytest.mark.instance.
+    """
+    _mark_onboarding_complete(session)
+
+    with sync_playwright() as p:
+        browser = p.chromium.launch(headless=True)
+        page = browser.new_page(viewport={"width": 1280, "height": 800})
+        page.goto(session.cluster.api_url, wait_until="networkidle")
+        page.wait_for_timeout(3000)
+        page.screenshot(path="/tmp/dashboard_initial.png")
+
+        # Open the model picker by clicking the "SELECT MODEL" button
+        page.get_by_text("SELECT MODEL", exact=False).first.click()
+        page.wait_for_timeout(1000)
+        page.screenshot(path="/tmp/dashboard_picker_open.png")
+
+        # Search for the model — uses the model id substring; the picker
+        # matches against name/id so "Llama-3.2-1B" filters to the small Llama.
+        search_input = page.locator('input[placeholder*="Search models"]').first
+        search_input.fill("Llama-3.2-1B")
+        page.wait_for_timeout(1500)
+        page.screenshot(path="/tmp/dashboard_picker_search.png")
+
+        # Click the only matching result. The picker shows the model's
+        # display name (e.g. "Llama 3.2 1B") which differs from the model_id.
+        # We click the first visible button-like row in the result list.
+        page.get_by_text("Llama 3.2 1B", exact=False).first.click()
+        page.wait_for_timeout(1500)
+        page.screenshot(path="/tmp/dashboard_model_selected.png")
+
+        # Type a chat message — sending triggers the dashboard's auto-launch
+        # flow: it picks an optimal placement for the selected model and POSTs
+        # to /instance, then sends the chat once the runner is ready.
+        chat_input = page.locator("textarea").first
+        chat_input.fill("Say hello")
+        chat_input.press("Enter")
+        page.screenshot(path="/tmp/dashboard_chat_sent.png")
+
+        # Wait for the instance to launch and respond. Generous timeout
+        # because this includes model placement + load + generation.
+        page.wait_for_timeout(60000)
+        page.screenshot(path="/tmp/dashboard_after_chat.png")
+
+        # Verify an instance was created and the chat got a response
+        instances = session.client.request_json("GET", "/state").get("instances", {})
+        assert len(instances) > 0, "Expected the dashboard to have created an instance"
+
+        body_text = page.text_content("body") or ""
+        assert len(body_text) > 0
+
+        browser.close()

tests/test_resilience.py

@@ -0,0 +1,56 @@
+# type: ignore
+"""Resilience tests: disconnect/reconnect nodes and verify cluster recovery.
+
+Run with:
+    uv run pytest tests/test_resilience.py -v
+"""
+
+from __future__ import annotations
+
+import pytest
+from exo_tools.cluster import Thunderbolt
+from exo_tools.harness import Comm, Sharding, cleanup_all_instances, place_instance
+
+from .framework import DEFAULT_MODEL, InstanceSpec
+
+
+@pytest.mark.cluster(count=2, thunderbolt=Thunderbolt.A2A)
+@pytest.mark.instance(
+    DEFAULT_MODEL, sharding=Sharding.PIPELINE, comm=Comm.RING, min_nodes=2
+)
+def test_node_recovery(session):
+    """Full disconnect/reconnect cycle.
+
+    1. Place a 2-node instance, verify inference
+    2. Disconnect one node
+    3. Place a 1-node instance on remaining node, verify inference
+    4. Reconnect the stopped node, wait for the cluster to reform
+    5. Place a 2-node instance again, verify inference
+    """
+    # --- Phase 1: 2-node inference ---
+    resp = session.chat("Hello")
+    assert len(resp) > 0
+
+    # --- Phase 2: disconnect one node ---
+    session.disconnect_node(1)
+    session.wait_ready(60)
+
+    # Clean up the now-broken 2-node instance
+    cleanup_all_instances(session.client)
+
+    # --- Phase 3: 1-node inference on the remaining node ---
+    place_instance(session.client, DEFAULT_MODEL, min_nodes=1)
+    session.instance_spec = InstanceSpec(model_id=DEFAULT_MODEL, min_nodes=1)
+    resp = session.chat("Hello")
+    assert len(resp) > 0
+
+    # --- Phase 4: reconnect and restore 2-node cluster ---
+    cleanup_all_instances(session.client)
+    session.reconnect_node(1)
+    session.wait_ready(60)
+
+    # --- Phase 5: 2-node inference again ---
+    place_instance(session.client, DEFAULT_MODEL, min_nodes=2)
+    session.instance_spec = InstanceSpec(model_id=DEFAULT_MODEL, min_nodes=2)
+    resp = session.chat("Hello again")
+    assert len(resp) > 0

tools/pyproject.toml

@@ -0,0 +1,10 @@
+[project]
+name = "exo-tools"
+version = "0.1.0"
+description = "Shared tooling for interacting with exo clusters"
+requires-python = ">=3.13"
+dependencies = ["loguru>=0.7.3"]
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"

tools/src/exo_tools/client.py

@@ -0,0 +1,117 @@
+# type: ignore
+"""HTTP client for the exo API."""
+
+from __future__ import annotations
+
+import http.client
+import json
+from collections.abc import Iterator
+from typing import Any
+from urllib.parse import urlencode
+
+
+class ExoHttpError(RuntimeError):
+    def __init__(self, status: int, reason: str, body_preview: str):
+        super().__init__(f"HTTP {status} {reason}: {body_preview}")
+        self.status = status
+
+
+class ExoClient:
+    def __init__(self, host: str, port: int, timeout_s: float = 7200.0):
+        self.host = host
+        self.port = port
+        self.timeout_s = timeout_s
+
+    def request_json(
+        self,
+        method: str,
+        path: str,
+        params: dict[str, Any] | None = None,
+        body: dict[str, Any] | None = None,
+        headers: dict[str, str] | None = None,
+    ) -> Any:
+        if not path.startswith("/"):
+            path = "/" + path
+        if params:
+            path = path + "?" + urlencode(params)
+
+        conn = http.client.HTTPConnection(self.host, self.port, timeout=self.timeout_s)
+        try:
+            payload: bytes | None = None
+            hdrs: dict[str, str] = {"Accept": "application/json"}
+
+            if body is not None:
+                payload = json.dumps(body).encode("utf-8")
+                hdrs["Content-Type"] = "application/json"
+            if headers:
+                hdrs.update(headers)
+
+            conn.request(method.upper(), path, body=payload, headers=hdrs)
+            resp = conn.getresponse()
+            raw = resp.read()
+            text = raw.decode("utf-8", errors="replace") if raw else ""
+
+            if resp.status >= 400:
+                raise ExoHttpError(resp.status, resp.reason, text[:300])
+
+            if not text:
+                return None
+            return json.loads(text)
+        finally:
+            conn.close()
+
+    def post_bench_chat_completions(self, payload: dict[str, Any]) -> dict[str, Any]:
+        return self.request_json("POST", "/bench/chat/completions", body=payload)
+
+    def stream_bench_chat_completions(self, payload: dict[str, Any]) -> Iterator[str]:
+        """POST /bench/chat/completions with stream=True, yielding raw SSE lines."""
+        payload = {**payload, "stream": True}
+        data = json.dumps(payload).encode("utf-8")
+        conn = http.client.HTTPConnection(self.host, self.port, timeout=self.timeout_s)
+        try:
+            conn.request(
+                "POST",
+                "/bench/chat/completions",
+                body=data,
+                headers={
+                    "Content-Type": "application/json",
+                    "Accept": "text/event-stream",
+                },
+            )
+            resp = conn.getresponse()
+            if resp.status >= 400:
+                raw = resp.read().decode("utf-8", errors="replace")
+                raise ExoHttpError(resp.status, resp.reason, raw[:300])
+            for line in resp:
+                yield line.decode("utf-8", errors="replace")
+        finally:
+            conn.close()
+
+    def get_state_path(self, path: str) -> Any:
+        try:
+            return self.request_json("GET", f"/state/{path}")
+        except ExoHttpError as e:
+            if e.status == 404:
+                return None
+            raise
+
+    def get_instance(self, instance_id: str) -> dict[str, Any] | None:
+        return self.get_state_path(f"instances/{instance_id}")
+
+    def get_runner(self, runner_id: str) -> dict[str, Any] | None:
+        return self.get_state_path(f"runners/{runner_id}")
+
+    def get_node_downloads(self, node_id: str) -> list[dict[str, Any]] | None:
+        return self.get_state_path(f"downloads/{node_id}")
+
+    def get_node_disk(self, node_id: str) -> dict[str, Any] | None:
+        return self.get_state_path(f"nodeDisk/{node_id}")
+
+    def get_node_system(self, node_id: str) -> dict[str, Any] | None:
+        return self.get_state_path(f"nodeSystem/{node_id}")
+
+    def get_node_identities(self) -> dict[str, Any] | None:
+        return self.get_state_path("nodeIdentities")
+
+    def get_topology(self) -> dict[str, Any] | None:
+        return self.get_state_path("topology")

tools/src/exo_tools/cluster.py

@@ -0,0 +1,262 @@
+# type: ignore
+"""Cluster lifecycle management via eco.
+
+Provides subprocess wrappers for eco commands (deploy, stop, start, release,
+logs, exec) and a ClusterInfo dataclass. Reusable by integration tests,
+bench, eval, and CI workflows.
+"""
+
+from __future__ import annotations
+
+import atexit
+import contextlib
+import json
+import logging
+import math
+import os
+import signal
+import subprocess
+import uuid
+from dataclasses import dataclass, field
+from enum import Enum
+
+from .client import ExoClient
+
+
+class Thunderbolt(str, Enum):
+    A2A = "a2a"  # all-to-all (eco --tb-a2a)
+    RING = "ring"  # ring topology (eco --tb-ring)
+    NONE = "none"  # exclude Thunderbolt-connected hosts (eco --no-thunderbolt)
+
+
+class Chip(str, Enum):
+    M1 = "M1"
+    M1_PRO = "M1 Pro"
+    M1_MAX = "M1 Max"
+    M1_ULTRA = "M1 Ultra"
+    M2 = "M2"
+    M2_PRO = "M2 Pro"
+    M2_MAX = "M2 Max"
+    M2_ULTRA = "M2 Ultra"
+    M3 = "M3"
+    M3_PRO = "M3 Pro"
+    M3_MAX = "M3 Max"
+    M3_ULTRA = "M3 Ultra"
+    M4 = "M4"
+    M4_PRO = "M4 Pro"
+    M4_MAX = "M4 Max"
+    M4_ULTRA = "M4 Ultra"
+
+
+logger = logging.getLogger("exo_tools.cluster")
+
+# When set, deploy from a GitHub branch/tag instead of local source (rsync).
+_EXO_REF = os.environ.get("EXO_REF")
+
+
+@dataclass
+class ClusterInfo:
+    """Holds the result of an `eco start --deploy` invocation."""
+
+    hosts: list[str]
+    namespace: str
+    api_endpoints: dict[str, str]  # host -> url
+    api_url: str  # primary endpoint for ExoClient
+
+    primary_host: str = ""
+    _host: str = field(init=False, repr=False, default="")
+    _port: int = field(init=False, repr=False, default=52415)
+
+    def __post_init__(self) -> None:
+        if not self.primary_host:
+            self.primary_host = self.hosts[0]
+        url = self.api_url.replace("http://", "").replace("https://", "")
+        parts = url.split(":")
+        self._host = parts[0]
+        self._port = int(parts[1]) if len(parts) > 1 else 52415
+
+    def make_client(self, timeout_s: float = 7200.0) -> ExoClient:
+        return ExoClient(self._host, self._port, timeout_s=timeout_s)
+
+
+class EcoSession:
+    """Manages an eco session with a unique user and automatic cleanup.
+
+    Usage:
+        session = EcoSession(user_prefix="test")
+        cluster = session.start_deploy(count=2, thunderbolt=True)
+        ...
+        session.stop_all()  # or let atexit handle it
+
+    The session registers atexit and signal handlers to ensure cleanup
+    on normal exit, uncaught exceptions, SIGTERM, and SIGHUP. SIGINT
+    is left unhandled so KeyboardInterrupt propagates normally.
+    """
+
+    def __init__(self, user_prefix: str = "test") -> None:
+        self._session_id = uuid.uuid4().hex[:8]
+        self.user = f"{user_prefix}-{self._session_id}"
+        self._env = {**os.environ, "USER": self.user}
+
+        # Register cleanup handlers
+        atexit.register(self.stop_all)
+        for sig in (signal.SIGTERM, signal.SIGHUP):
+            signal.signal(sig, self._signal_handler)
+
+    def _signal_handler(self, signum: int, _frame: object) -> None:
+        self.stop_all()
+        raise SystemExit(128 + signum)
+
+    def stop_all(self) -> None:
+        """Stop all clusters and release all reservations for this session."""
+        with contextlib.suppress(Exception):
+            subprocess.run(
+                ["eco", "stop"],
+                capture_output=True,
+                text=True,
+                timeout=30,
+                env=self._env,
+            )
+
+    def _run(
+        self, args: list[str], *, check: bool = True, timeout: int = 120
+    ) -> subprocess.CompletedProcess[str]:
+        """Run an eco command as this session's user.
+
+        stdout is captured (JSON output), stderr is passed through to the
+        console so eco's progress messages are visible.
+        """
+        logger.info(f"eco: {' '.join(args)}")
+        return subprocess.run(
+            args,
+            stdout=subprocess.PIPE,
+            stderr=None,
+            text=True,
+            check=check,
+            timeout=timeout,
+            env=self._env,
+        )
+
+    def start_deploy(
+        self,
+        hosts: list[str] | None = None,
+        *,
+        count: int | None = None,
+        thunderbolt: Thunderbolt | None = None,
+        chip: Chip | None = None,
+        min_memory_gb: float | None = None,
+        max_memory_gb: float | None = None,
+        min_disk_gb: float | None = None,
+        max_disk_gb: float | None = None,
+        wait: bool = True,
+        ref: str | None = _EXO_REF,
+        timeout: int = 600,
+    ) -> ClusterInfo:
+        """Start and deploy exo on a set of hosts via eco.
+
+        By default, deploys from local source via rsync. Set EXO_REF
+        or pass ref= to deploy from a GitHub branch/tag instead (for CI).
+
+        Selection constraints (memory/disk in GiB, chip substring,
+        Thunderbolt topology) are forwarded as eco CLI flags. Pass
+        ``thunderbolt=Thunderbolt.NONE`` to exclude TB-connected hosts.
+        """
+        cmd: list[str] = ["eco", "--json", "start", "--deploy"]
+        if hosts:
+            cmd.extend(hosts)
+        if count is not None:
+            cmd.extend(["--count", str(count)])
+        if thunderbolt is Thunderbolt.NONE:
+            cmd.append("--no-thunderbolt")
+        elif thunderbolt is not None:
+            cmd.append(f"--tb-{thunderbolt.value}")
+        if chip is not None:
+            cmd.extend(["--chip", chip.value])
+        # eco's GB args are integer-typed. Round mins up + maxes down so
+        # we never relax the user's constraint.
+        if min_memory_gb is not None:
+            cmd.extend(["--min-memory", str(math.ceil(min_memory_gb))])
+        if max_memory_gb is not None:
+            cmd.extend(["--max-memory", str(math.floor(max_memory_gb))])
+        if min_disk_gb is not None:
+            cmd.extend(["--min-disk", str(math.ceil(min_disk_gb))])
+        if max_disk_gb is not None:
+            cmd.extend(["--max-disk", str(math.floor(max_disk_gb))])
+        if wait:
+            cmd.append("--wait")
+        if ref:
+            cmd.extend(["--ref", ref])
+
+        result = self._run(cmd, timeout=timeout)
+        data = json.loads(result.stdout)["data"]
+        endpoints: dict[str, str] = data["api_endpoints"]
+        primary_host = data["hosts"][0]
+
+        return ClusterInfo(
+            hosts=data["hosts"],
+            namespace=data["namespace"],
+            api_endpoints=endpoints,
+            api_url=endpoints[primary_host],
+            primary_host=primary_host,
+        )
+
+    def stop(self, hosts: list[str], *, keep: bool = False, timeout: int = 120) -> None:
+        """Stop exo on the given hosts. If keep=True, keep the reservation."""
+        cmd: list[str] = ["eco", "stop"]
+        cmd.extend(hosts)
+        if keep:
+            cmd.append("--keep")
+        self._run(cmd, timeout=timeout)
+
+    def start_hosts(
+        self, hosts: list[str], *, namespace: str, timeout: int = 300
+    ) -> None:
+        """Start (previously stopped) hosts back into an existing namespace."""
+        cmd: list[str] = ["eco", "--json", "start"]
+        cmd.extend(hosts)
+        cmd.extend(["--namespace", namespace])
+        self._run(cmd, timeout=timeout)
+
+    def release(self, hosts: list[str], timeout: int = 120) -> None:
+        """Release hosts from the reservation."""
+        cmd: list[str] = ["eco", "release"]
+        cmd.extend(hosts)
+        self._run(cmd, timeout=timeout)
+
+    def logs(
+        self, hosts: list[str], lines: int = 500, timeout: int = 60
+    ) -> dict[str, list[str]]:
+        """Fetch recent logs from cluster hosts."""
+        cmd: list[str] = ["eco", "--json", "logs"]
+        cmd.extend(hosts)
+        cmd.extend(["-n", str(lines), "--raw"])
+        result = self._run(cmd, check=False, timeout=timeout)
+        if result.returncode != 0:
+            return {"_error": [result.stderr]}
+        try:
+            return json.loads(result.stdout)
+        except json.JSONDecodeError:
+            return {"_raw": result.stdout.splitlines()}
+
+    def exec(self, hosts: list[str], command: str, timeout: int = 120) -> str:
+        """Run an arbitrary command on the given hosts via eco."""
+        cmd: list[str] = ["eco", "exec"]
+        cmd.extend(hosts)
+        cmd.append("--")
+        cmd.extend(command.split())
+        result = self._run(cmd, check=False, timeout=timeout)
+        return result.stdout
+
+
+def make_client(cluster: ClusterInfo, timeout_s: float = 7200.0) -> ExoClient:
+    """Create an ExoClient from a ClusterInfo."""
+    return cluster.make_client(timeout_s=timeout_s)
+
+
+def make_client_from_url(url: str, timeout_s: float = 7200.0) -> ExoClient:
+    """Create an ExoClient from a URL string like 'http://host:port'."""
+    url_clean = url.replace("http://", "").replace("https://", "")
+    parts = url_clean.split(":")
+    host = parts[0]
+    port = int(parts[1]) if len(parts) > 1 else 52415
+    return ExoClient(host, port, timeout_s=timeout_s)

tools/src/exo_tools/harness.py

@@ -1,129 +1,39 @@
 # type: ignore
+"""Instance lifecycle helpers for exo clusters.
+
+Provides utilities for placing instances, waiting for readiness,
+managing downloads, filtering placements, and common CLI arguments.
+"""
+
 from __future__ import annotations
 
 import argparse
-import http.client
-import json
+import contextlib
 import os
 import time
-from collections.abc import Iterator
+from enum import Enum
 from typing import Any
-from urllib.parse import urlencode
 
 from loguru import logger
 
+from .client import ExoClient, ExoHttpError
+
+
+class Sharding(str, Enum):
+    PIPELINE = "Pipeline"  # layers split across nodes
+    TENSOR = "Tensor"  # layers split within (across nodes)
+
+
+class Comm(str, Enum):
+    RING = "MlxRing"  # ring all-reduce over network
+    JACCL = "MlxJaccl"  # RDMA over Thunderbolt
+
+
 _SETTLE_INITIAL_BACKOFF_S = 1.0
 _SETTLE_MAX_BACKOFF_S = 60.0
 _SETTLE_BACKOFF_MULTIPLIER = 2.0
 
 
-class ExoHttpError(RuntimeError):
-    def __init__(self, status: int, reason: str, body_preview: str):
-        super().__init__(f"HTTP {status} {reason}: {body_preview}")
-        self.status = status
-
-
-class ExoClient:
-    def __init__(self, host: str, port: int, timeout_s: float = 7200.0):
-        self.host = host
-        self.port = port
-        self.timeout_s = timeout_s
-
-    def request_json(
-        self,
-        method: str,
-        path: str,
-        params: dict[str, Any] | None = None,
-        body: dict[str, Any] | None = None,
-        headers: dict[str, str] | None = None,
-    ) -> Any:
-        if not path.startswith("/"):
-            path = "/" + path
-        if params:
-            path = path + "?" + urlencode(params)
-
-        conn = http.client.HTTPConnection(self.host, self.port, timeout=self.timeout_s)
-        try:
-            payload: bytes | None = None
-            hdrs: dict[str, str] = {"Accept": "application/json"}
-
-            if body is not None:
-                payload = json.dumps(body).encode("utf-8")
-                hdrs["Content-Type"] = "application/json"
-            if headers:
-                hdrs.update(headers)
-
-            conn.request(method.upper(), path, body=payload, headers=hdrs)
-            resp = conn.getresponse()
-            raw = resp.read()
-            text = raw.decode("utf-8", errors="replace") if raw else ""
-
-            if resp.status >= 400:
-                raise ExoHttpError(resp.status, resp.reason, text[:300])
-
-            if not text:
-                return None
-            return json.loads(text)
-        finally:
-            conn.close()
-
-    def post_bench_chat_completions(self, payload: dict[str, Any]) -> dict[str, Any]:
-        return self.request_json("POST", "/bench/chat/completions", body=payload)
-
-    def stream_bench_chat_completions(self, payload: dict[str, Any]) -> Iterator[str]:
-        """POST /bench/chat/completions with stream=True, yielding raw SSE lines."""
-        payload = {**payload, "stream": True}
-        data = json.dumps(payload).encode("utf-8")
-        conn = http.client.HTTPConnection(self.host, self.port, timeout=self.timeout_s)
-        try:
-            conn.request(
-                "POST",
-                "/bench/chat/completions",
-                body=data,
-                headers={
-                    "Content-Type": "application/json",
-                    "Accept": "text/event-stream",
-                },
-            )
-            resp = conn.getresponse()
-            if resp.status >= 400:
-                raw = resp.read().decode("utf-8", errors="replace")
-                raise ExoHttpError(resp.status, resp.reason, raw[:300])
-            for line in resp:
-                yield line.decode("utf-8", errors="replace")
-        finally:
-            conn.close()
-
-    def get_state_path(self, path: str) -> Any:
-        try:
-            return self.request_json("GET", f"/state/{path}")
-        except ExoHttpError as e:
-            if e.status == 404:
-                return None
-            raise
-
-    def get_instance(self, instance_id: str) -> dict[str, Any] | None:
-        return self.get_state_path(f"instances/{instance_id}")
-
-    def get_runner(self, runner_id: str) -> dict[str, Any] | None:
-        return self.get_state_path(f"runners/{runner_id}")
-
-    def get_node_downloads(self, node_id: str) -> list[dict[str, Any]] | None:
-        return self.get_state_path(f"downloads/{node_id}")
-
-    def get_node_disk(self, node_id: str) -> dict[str, Any] | None:
-        return self.get_state_path(f"nodeDisk/{node_id}")
-
-    def get_node_system(self, node_id: str) -> dict[str, Any] | None:
-        return self.get_state_path(f"nodeSystem/{node_id}")
-
-    def get_node_identities(self) -> dict[str, Any] | None:
-        return self.get_state_path("nodeIdentities")
-
-    def get_topology(self) -> dict[str, Any] | None:
-        return self.get_state_path("topology")
-
-
 def unwrap_instance(instance: dict[str, Any]) -> dict[str, Any]:
     if len(instance) != 1:
         raise KeyError(f"Expected 1 key, got keys={list(instance.keys())}")
@@ -555,7 +465,6 @@ def find_existing_instance(client: ExoClient, model_id: str) -> str | None:
     except Exception:
         return None
     for inst_id, inst in state.get("instances", {}).items():
-        # Instance structure is nested: {"MlxJacclInstance": {"shardAssignments": {"modelId": ...}}}
         for _inst_type, inner in inst.items():
             if not isinstance(inner, dict):
                 continue
@@ -623,3 +532,112 @@ def add_common_instance_args(ap: argparse.ArgumentParser) -> None:
         action="store_true",
         help="Reuse an existing running instance for this model instead of creating a new one.",
     )
+
+
+# ---------------------------------------------------------------------------
+# Cluster/instance orchestration helpers (used by tests, bench, eval)
+# ---------------------------------------------------------------------------
+
+
+def get_instance_ids(client: ExoClient) -> set[str]:
+    """Return the set of current instance IDs from cluster state."""
+    state = client.request_json("GET", "/state") or {}
+    result: set[str] = set()
+    for instance in state.get("instances", {}).values():
+        with contextlib.suppress(Exception):
+            result.add(instance_id_from_instance(instance))
+    return result
+
+
+def wait_for_cluster_ready(
+    client: ExoClient, expected_nodes: int = 1, timeout: float = 120.0
+) -> None:
+    """Wait until the cluster has all expected nodes visible and reporting memory.
+
+    Placement requires nodeMemory for all nodes in a cycle. This polls until
+    both nodeIdentities and nodeMemory have at least `expected_nodes` entries.
+    """
+    start = time.time()
+    while time.time() - start < timeout:
+        try:
+            state = client.request_json("GET", "/state") or {}
+            if (
+                len(state.get("nodeIdentities", {})) >= expected_nodes
+                and len(state.get("nodeMemory", {})) >= expected_nodes
+            ):
+                return
+        except Exception:
+            pass
+        time.sleep(1.0)
+    raise TimeoutError(f"Cluster not ready: expected {expected_nodes} nodes")
+
+
+def place_instance(
+    client: ExoClient,
+    model_id: str,
+    *,
+    sharding: Sharding = Sharding.PIPELINE,
+    comm: Comm = Comm.RING,
+    min_nodes: int = 1,
+    timeout: float = 600.0,
+    placement_retries: int = 10,
+    placement_retry_delay: float = 10.0,
+) -> str:
+    """Place an instance and wait for it to be ready. Returns the instance_id.
+
+    The /place_instance API returns a command_id, but instances are stored
+    under a separately-generated instance_id. This polls cluster state for the
+    new instance, retrying placement if the cluster is still settling.
+    """
+    wait_for_cluster_ready(client, expected_nodes=min_nodes)
+
+    body = {
+        "model_id": model_id,
+        "sharding": sharding.value,
+        "instance_meta": comm.value,
+        "min_nodes": min_nodes,
+    }
+
+    instance_id: str | None = None
+    for attempt in range(placement_retries):
+        before_ids = get_instance_ids(client)
+        client.request_json("POST", "/place_instance", body=body)
+
+        poll_deadline = time.time() + 30.0
+        while time.time() < poll_deadline:
+            new_ids = get_instance_ids(client) - before_ids
+            if new_ids:
+                instance_id = next(iter(new_ids))
+                break
+            time.sleep(1.0)
+
+        if instance_id is not None:
+            break
+
+        if attempt < placement_retries - 1:
+            time.sleep(placement_retry_delay)
+
+    if instance_id is None:
+        raise TimeoutError(
+            f"Placement failed after {placement_retries} attempts "
+            f"({sharding.value}/{comm.value} for {model_id})"
+        )
+
+    wait_for_instance_ready(client, instance_id, timeout=timeout)
+    return instance_id
+
+
+def cleanup_all_instances(client: ExoClient) -> None:
+    """Remove all running instances from the cluster."""
+    state = client.request_json("GET", "/state") or {}
+    for instance in state.get("instances", {}).values():
+        with contextlib.suppress(Exception):
+            iid = instance_id_from_instance(instance)
+            client.request_json("DELETE", f"/instance/{iid}")
+            wait_for_instance_gone(client, iid, timeout=30.0)
+
+
+def is_model_downloaded(client: ExoClient, model_id: str) -> bool:
+    response = client.request_json("GET", "/models", params={"status": "downloaded"})
+    data = (response or {}).get("data", [])
+    return all(model.get("id") == model_id for model in data)