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

.envrc

@@ -1 +1,8 @@
 use flake
+
+# creates .venv if doesn't exist and loads its environment
+export VIRTUAL_ENV=".venv"
+if ! [ -d "./$VIRTUAL_ENV" ]; then
+  uv venv
+fi
+layout python

.gitignore

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

.mlx_typings/mlx_lm/models/cache.pyi

@@ -191,10 +191,13 @@ class RotatingKVCache(_BaseCache):
     def state(self, v):  # -> None:
         ...
     @property
-    def meta_state(self) -> tuple[str, ...]: ...
+    def meta_state(self):  # -> tuple[str, ...]:
+        ...
     @meta_state.setter
-    def meta_state(self, v: tuple[str, ...]) -> None: ...
-    def is_trimmable(self) -> bool: ...
+    def meta_state(self, v):  # -> None:
+        ...
+    def is_trimmable(self):  # -> bool:
+        ...
     def trim(self, n: int) -> int: ...
     def to_quantized(
         self, group_size: int = ..., bits: int = ...

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.toml

@@ -12,24 +12,24 @@ timeout = 7200.0
 settle_timeout = 60.0
 
 # Workload
-pp = [4096, 8192]
-tg = [128]
+pp = [4096]
+tg = [512]
 repeat = 1
 warmup = 0
 
 json_out = "bench/prefill_decode_results.json"
 
 [prefill]
-model = "sakamakismile/Qwen3.6-27B-NVFP4"
-node = "gx10-de89"
-instance_meta = "vllm"
+model = "mlx-community/gpt-oss-20b-MXFP4-Q8"
+node = "mike"
+instance_meta = "ring"
 sharding = "pipeline"
 min_nodes = 1
 max_nodes = 1
 
 [decode]
-model = "mlx-community/Qwen3.6-27B-4bit"
-node = "Ryuichi’s MacBook Pro"
+model = "mlx-community/gpt-oss-20b-MXFP4-Q8"
+node = "james"
 instance_meta = "ring"
 sharding = "pipeline"
 min_nodes = 1

bench/prefill_decode_bench.py

@@ -31,14 +31,12 @@ from typing import Any
 
 from exo_bench import (
     PromptSizer,
-    SystemMetricsSampler,
     format_peak_memory,
     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,
@@ -279,7 +277,6 @@ def _run_phase(
     warmup: int,
     repeat: int,
     common_meta: dict[str, Any],
-    sampler: SystemMetricsSampler | None = None,
 ) -> list[dict[str, Any]]:
     logger.info(f"=== phase: {label} (model={model_id}) ===")
     rows: list[dict[str, Any]] = []
@@ -290,13 +287,10 @@ def _run_phase(
     for pp, tg in pp_tg_pairs:
         logger.info(f"--- {label}: pp={pp} tg={tg} ---")
         runs: list[dict[str, Any]] = []
-        inference_windows: list[tuple[float, float]] = []
         for r in range(repeat):
             time.sleep(2)
             try:
-                inf_t0 = time.monotonic()
                 row, actual_pp_tokens = run_one(client, model_id, pp, tg, prompt_sizer)
-                inference_windows.append((inf_t0, time.monotonic()))
             except Exception as e:
                 logger.error(e)
                 continue
@@ -320,26 +314,11 @@ def _run_phase(
             gtok = mean(x["stats"]["generation_tokens"] for x in runs)
             peak = mean(x["stats"]["peak_memory_usage"]["inBytes"] for x in runs)
             avg_elapsed = mean(x["elapsed_s"] for x in runs)
-            energy_str = ""
-            if sampler is not None and inference_windows:
-                joules = sum(
-                    sampler.energy_between(t0, t1) for t0, t1 in inference_windows
-                )
-                inf_seconds = sum(t1 - t0 for t0, t1 in inference_windows)
-                avg_watts = joules / inf_seconds if inf_seconds > 0 else 0.0
-                energy_per_run = joules / len(runs) if runs else 0.0
-                energy_str = (
-                    f"    energy={joules:.1f}J ({avg_watts:.1f}W avg over "
-                    f"{inf_seconds:.1f}s inference, {energy_per_run:.1f}J/run)"
-                )
-                for run_row, (t0, t1) in zip(runs, inference_windows, strict=False):
-                    run_row["energy_joules"] = sampler.energy_between(t0, t1)
-                    run_row["inference_window_s"] = t1 - t0
             logger.info(
                 f"[{label}] prompt_tps={prompt_tps:.2f} gen_tps={gen_tps:.2f}    "
                 f"prompt_tokens={ptok} gen_tokens={gtok}    "
                 f"peak_memory={format_peak_memory(peak)}    "
-                f"avg_elapsed={avg_elapsed:.2f}s{energy_str}"
+                f"avg_elapsed={avg_elapsed:.2f}s"
             )
         time.sleep(2)
     return rows
@@ -352,36 +331,14 @@ def _summarise(rows: list[dict[str, Any]]) -> dict[tuple[int, int], dict[str, fl
         grouped.setdefault(key, []).append(r)
     out: dict[tuple[int, int], dict[str, float]] = {}
     for key, runs in grouped.items():
-        energy_runs = [x.get("energy_joules") for x in runs if "energy_joules" in x]
-        window_runs = [
-            x.get("inference_window_s") for x in runs if "inference_window_s" in x
-        ]
         out[key] = {
             "prompt_tps": mean(x["stats"]["prompt_tps"] for x in runs),
             "gen_tps": mean(x["stats"]["generation_tps"] for x in runs),
             "elapsed_s": mean(x["elapsed_s"] for x in runs),
-            "prompt_tokens": mean(x["stats"]["prompt_tokens"] for x in runs),
-            "gen_tokens": mean(x["stats"]["generation_tokens"] for x in runs),
-            "energy_j": mean(energy_runs) if energy_runs else 0.0,
-            "inference_window_s": mean(window_runs) if window_runs else 0.0,
         }
     return out
 
 
-def _normalised_seconds(summary: dict[str, float], pp: int, tg: int) -> float | None:
-    """Wall-clock time implied by reported tps for the *configured* pp/tg.
-
-    elapsed_s is not comparable across phases when models EOS at different
-    lengths. This formula reconstructs "what would this phase take to do
-    pp prompt tokens + tg generation tokens" using its own reported rates.
-    """
-    p_tps = summary.get("prompt_tps", 0.0)
-    g_tps = summary.get("gen_tps", 0.0)
-    if p_tps <= 0 or g_tps <= 0:
-        return None
-    return pp / p_tps + tg / g_tps
-
-
 def _print_diff(
     disagg_rows: list[dict[str, Any]],
     decode_alone_rows: list[dict[str, Any]],
@@ -392,17 +349,14 @@ def _print_diff(
     prefill_alone = _summarise(prefill_alone_rows)
     keys = set(disagg.keys()) | set(decode_alone.keys()) | set(prefill_alone.keys())
 
-    width = 110
+    width = 64
     for key in sorted(keys):
         pp, tg = key
         logger.info("─" * width)
         logger.info(f"  pp={pp}  tg={tg}")
         logger.info("─" * width)
         logger.info(
-            f"  {'phase':<16} {'elapsed':>9}  {'norm':>9}  "
-            f"{'prompt_tps':>11}  {'gen_tps':>8}  "
-            f"{'p_tok':>6}  {'g_tok':>6}  "
-            f"{'energy':>9}  {'avg_W':>7}"
+            f"  {'phase':<16} {'elapsed':>10}  {'prompt_tps':>11}  {'gen_tps':>9}"
         )
         for label, summary in (
             ("disaggregated", disagg.get(key)),
@@ -410,51 +364,26 @@ def _print_diff(
             ("prefill_alone", prefill_alone.get(key)),
         ):
             if summary is None:
-                logger.info(
-                    f"  {label:<16} {'—':>9}  {'—':>9}  "
-                    f"{'—':>11}  {'—':>8}  {'—':>6}  {'—':>6}  "
-                    f"{'—':>9}  {'—':>7}"
-                )
+                logger.info(f"  {label:<16} {'—':>10}  {'—':>11}  {'—':>9}")
                 continue
-            norm = _normalised_seconds(summary, pp, tg)
-            norm_str = f"{norm:>8.2f}s" if norm is not None else f"{'—':>9}"
-            energy = summary.get("energy_j", 0.0)
-            window = summary.get("inference_window_s", 0.0)
-            energy_str = f"{energy:>8.1f}J" if energy > 0 else f"{'—':>9}"
-            avg_w = energy / window if window > 0 else 0.0
-            avg_w_str = f"{avg_w:>6.1f}W" if avg_w > 0 else f"{'—':>7}"
             logger.info(
                 f"  {label:<16} "
-                f"{summary['elapsed_s']:>8.2f}s  "
-                f"{norm_str}  "
+                f"{summary['elapsed_s']:>9.2f}s  "
                 f"{summary['prompt_tps']:>11.1f}  "
-                f"{summary['gen_tps']:>8.2f}  "
-                f"{summary['prompt_tokens']:>6.0f}  "
-                f"{summary['gen_tokens']:>6.0f}  "
-                f"{energy_str}  "
-                f"{avg_w_str}"
+                f"{summary['gen_tps']:>9.2f}"
             )
 
         d = disagg.get(key)
         da = decode_alone.get(key)
         pa = prefill_alone.get(key)
-        d_norm = _normalised_seconds(d, pp, tg) if d else None
-        if d_norm and da:
-            da_norm = _normalised_seconds(da, pp, tg)
-            if da_norm:
-                logger.info(
-                    f"  norm speedup vs decode_alone:  {da_norm / d_norm:.2f}x  "
-                    f"(prefill {d['prompt_tps'] / da['prompt_tps']:.2f}x, "
-                    f"decode {d['gen_tps'] / da['gen_tps']:.2f}x)"
-                )
-        if d_norm and pa:
-            pa_norm = _normalised_seconds(pa, pp, tg)
-            if pa_norm:
-                logger.info(
-                    f"  norm speedup vs prefill_alone: {pa_norm / d_norm:.2f}x  "
-                    f"(prefill {d['prompt_tps'] / pa['prompt_tps']:.2f}x, "
-                    f"decode {d['gen_tps'] / pa['gen_tps']:.2f}x)"
-                )
+        if d and da and d["elapsed_s"] > 0:
+            logger.info(
+                f"  speedup vs decode_alone:  {da['elapsed_s'] / d['elapsed_s']:.2f}x"
+            )
+        if d and pa and d["elapsed_s"] > 0:
+            logger.info(
+                f"  speedup vs prefill_alone: {pa['elapsed_s'] / d['elapsed_s']:.2f}x"
+            )
     logger.info("─" * width)
 
 
@@ -752,16 +681,6 @@ def main() -> int:
     link_id = ""
     prefill_alive = False
     decode_alive = False
-    sampler_nodes = sorted(
-        {
-            *node_ids_from_instance(prefill_instance),
-            *node_ids_from_instance(decode_instance),
-        }
-    )
-    sampler = SystemMetricsSampler(
-        ExoClient(args.host, args.port, timeout_s=30), sampler_nodes
-    )
-    sampler.start()
     try:
         logger.info("Creating prefill instance...")
         client.request_json("POST", "/instance", body={"instance": prefill_instance})
@@ -780,7 +699,6 @@ def main() -> int:
                 warmup=args.warmup,
                 repeat=args.repeat,
                 common_meta=common_meta,
-                sampler=sampler,
             )
             all_rows.extend(prefill_alone_rows)
 
@@ -810,7 +728,6 @@ def main() -> int:
             warmup=args.warmup,
             repeat=args.repeat,
             common_meta=common_meta,
-            sampler=sampler,
         )
         all_rows.extend(disagg_rows)
 
@@ -835,13 +752,11 @@ def main() -> int:
                 warmup=args.warmup,
                 repeat=args.repeat,
                 common_meta=common_meta,
-                sampler=sampler,
             )
             all_rows.extend(decode_alone_rows)
 
             _print_diff(disagg_rows, decode_alone_rows, prefill_alone_rows)
     finally:
-        sampler.stop()
         with contextlib.suppress(ExoHttpError):
             if link_id:
                 _delete_instance_link(client, link_id)

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/lib/components/ChatSidebar.svelte

@@ -202,7 +202,6 @@
     let instanceType: string | null = null;
     if (instanceTag === "MlxRingInstance") instanceType = "MLX Ring";
     else if (instanceTag === "MlxJacclInstance") instanceType = "MLX RDMA";
-    else if (instanceTag === "VllmInstance") instanceType = "vLLM";
 
     let sharding: string | null = null;
     const inst = instance as {

dashboard/src/lib/components/DeviceIcon.svelte

@@ -9,7 +9,7 @@
    */
 
   interface Props {
-    /** "macbook pro" | "mac studio" | "mac mini" | "dgx spark" | "linux" etc. */
+    /** "macbook pro" | "mac studio" | "mac mini" etc. */
     deviceType: string;
     /** Center X coordinate in SVG space */
     cx: number;
@@ -38,43 +38,10 @@
   const LOGO_NATIVE_WIDTH = 814;
   const LOGO_NATIVE_HEIGHT = 1000;
 
-  // NVIDIA logo SVG path
-  const NVIDIA_LOGO_PATH =
-    "M0.81 0.429V0.299c0.013 -0.001 0.026 -0.002 0.038 -0.002 0.355 -0.011 0.588 0.306 0.588 0.306S1.186 0.952 0.916 0.952c-0.036 0 -0.071 -0.006 -0.105 -0.017V0.542c0.138 0.017 0.166 0.078 0.249 0.216l0.185 -0.155s-0.135 -0.177 -0.362 -0.177c-0.024 -0.001 -0.048 0.001 -0.072 0.003m0 -0.429v0.194l0.038 -0.002c0.494 -0.017 0.816 0.405 0.816 0.405s-0.37 0.45 -0.754 0.45c-0.034 0 -0.066 -0.003 -0.099 -0.009v0.12c0.027 0.003 0.055 0.006 0.082 0.006 0.358 0 0.618 -0.183 0.869 -0.399 0.042 0.034 0.212 0.114 0.247 0.15 -0.238 0.2 -0.794 0.361 -1.11 0.361 -0.03 0 -0.059 -0.002 -0.088 -0.005v0.169h1.362V0zm0 0.935v0.102c-0.331 -0.059 -0.423 -0.404 -0.423 -0.404s0.159 -0.176 0.423 -0.205v0.112h-0.001C0.671 0.524 0.562 0.654 0.562 0.654s0.062 0.218 0.248 0.282m-0.588 -0.316s0.196 -0.29 0.589 -0.32V0.194C0.376 0.229 0 0.597 0 0.597s0.213 0.616 0.81 0.672v-0.112c-0.438 -0.054 -0.588 -0.538 -0.588 -0.538";
-
   const wireColor = "rgba(179,179,179,0.8)";
   const strokeWidth = 1.5;
 
   const modelLower = $derived(deviceType.toLowerCase());
-  const isSpark = $derived(
-    modelLower.includes("dgx") || modelLower.includes("gx10"),
-  );
-  const isLinux = $derived(!isSpark && modelLower.startsWith("linux"));
-  const isLinuxLaptop = $derived(isLinux && modelLower.includes("laptop"));
-
-  // ── DGX Spark dimensions ──
-  const dgxW = $derived(size * 1.55);
-  const dgxH = $derived(size * 0.58);
-  const dgxX = $derived(cx - dgxW / 2);
-  const dgxY = $derived(cy - dgxH / 2);
-  const dgxChassisX = $derived(dgxX - dgxW * 0.03);
-  const dgxChassisW = $derived(dgxW * 1.05);
-  const dgxHandleW = $derived(dgxW * 0.27);
-  const dgxHandleGap = $derived(dgxH * 0.05);
-  const dgxHandleH = $derived(dgxH - dgxHandleGap * 2);
-  const dgxHandleY = $derived(dgxY + dgxHandleGap);
-  const dgxInnerHandleW = $derived(dgxW * 0.12);
-  const dgxInnerHandleH = $derived(dgxHandleH - dgxH * 0.06);
-  const dgxLeftHandleX = $derived(dgxX + 4);
-  const dgxRightHandleX = $derived(dgxX + dgxW - dgxHandleW - 4);
-  const dgxClipId = $derived(`di-dgx-${uid}`);
-  const dgxTextureId = $derived(`di-dgx-tex-${uid}`);
-
-  // ── Linux Desktop dimensions (reuses Mac Studio proportions) ──
-  const linuxDesktopClipId = $derived(`di-linux-desktop-${uid}`);
-
-  // ── Linux Laptop dimensions (reuses MacBook proportions) ──
-  const linuxScreenClipId = $derived(`di-linux-screen-${uid}`);
 
   // ── Mac Studio dimensions (same ratios as TopologyGraph) ──
   const studioW = $derived(size * 1.25);
@@ -147,264 +114,7 @@
   const studioClipId = $derived(`di-studio-${uid}`);
 </script>
 
-{#if isSpark}
-  <!-- DGX Spark -->
-  <defs>
-    <clipPath id={dgxClipId}>
-      <rect x={dgxX} y={dgxY} width={dgxW} height={dgxH} rx="3" />
-    </clipPath>
-    <pattern
-      id={dgxTextureId}
-      patternUnits="userSpaceOnUse"
-      width="8"
-      height="8"
-    >
-      <rect width="8" height="8" fill="#6f6248" />
-      <circle cx="2" cy="2" r="1" fill="#5a4f3b" opacity="0.5" />
-      <circle cx="6" cy="6" r="1" fill="#4a4232" opacity="0.45" />
-    </pattern>
-  </defs>
-
-  <!-- Main body -->
-  <rect
-    x={dgxChassisX}
-    y={dgxY}
-    width={dgxChassisW}
-    height={dgxH}
-    rx="3"
-    fill="url(#{dgxTextureId})"
-    stroke={wireColor}
-    stroke-width={strokeWidth}
-  />
-
-  <!-- Side border accents -->
-  <rect
-    x={dgxChassisX}
-    y={dgxY}
-    width={dgxW * 0.02}
-    height={dgxH}
-    fill="#8a7a56"
-  />
-  <rect
-    x={dgxChassisX + dgxChassisW - dgxW * 0.02}
-    y={dgxY}
-    width={dgxW * 0.02}
-    height={dgxH}
-    fill="#8a7a56"
-  />
-
-  <!-- Memory fill -->
-  {#if ramPercent > 0}
-    <rect
-      x={dgxX}
-      y={dgxY + dgxH - (ramPercent / 100) * dgxH}
-      width={dgxW}
-      height={(ramPercent / 100) * dgxH}
-      fill="rgba(255,215,0,0.45)"
-      clip-path="url(#{dgxClipId})"
-    />
-  {/if}
-
-  <!-- Left handle -->
-  <rect
-    x={dgxLeftHandleX}
-    y={dgxHandleY}
-    width={dgxHandleW}
-    height={dgxHandleH}
-    rx="2.4"
-    fill="#b3a170"
-    stroke="#403723"
-    stroke-width="0.7"
-  />
-  <rect
-    x={dgxLeftHandleX + dgxHandleW * 0.06}
-    y={dgxHandleY + dgxH * 0.03}
-    width={dgxInnerHandleW}
-    height={dgxInnerHandleH}
-    rx="1.6"
-    fill="#8a7a56"
-  />
-
-  <!-- Right handle -->
-  <rect
-    x={dgxRightHandleX}
-    y={dgxHandleY}
-    width={dgxHandleW}
-    height={dgxHandleH}
-    rx="2.4"
-    fill="#b3a170"
-    stroke="#403723"
-    stroke-width="0.7"
-  />
-  <rect
-    x={dgxRightHandleX + dgxHandleW - dgxInnerHandleW - dgxHandleW * 0.08}
-    y={dgxHandleY + dgxH * 0.03}
-    width={dgxInnerHandleW}
-    height={dgxInnerHandleH}
-    rx="1.6"
-    fill="#8a7a56"
-  />
-
-  <!-- NVIDIA logo (rotated 90deg on left handle) -->
-  {@const badgeW = dgxW * 0.09}
-  {@const badgeH = dgxHandleH * 0.5}
-  {@const badgeX = dgxLeftHandleX + dgxHandleW - badgeW - dgxHandleW * 0.06}
-  {@const badgeYPos = dgxHandleY + (dgxHandleH - badgeH) / 2}
-  {@const textSz = badgeW * 0.58}
-  {@const logoW = textSz * 1.2}
-  {@const logoH = logoW * (1.438 / 2.174)}
-  {@const ctrX = badgeX + badgeW / 2 - badgeW * 0.03}
-  {@const ctrY = badgeYPos + badgeH / 2}
-  {@const labelGap = badgeW * 0.15}
-  {@const totalW = logoW + labelGap + textSz * 3.6}
-  <g transform="rotate(90 {ctrX} {ctrY})">
-    <svg
-      x={ctrX - totalW / 2}
-      y={ctrY - logoH / 2}
-      width={logoW}
-      height={logoH}
-      viewBox="0 0 2.174 1.438"
-    >
-      <path d={NVIDIA_LOGO_PATH} fill="#76b900" />
-    </svg>
-    <text
-      x={ctrX - totalW / 2 + logoW + labelGap}
-      y={ctrY}
-      text-anchor="start"
-      dominant-baseline="middle"
-      fill="#8a7a56"
-      font-size={textSz}
-      font-family="monospace"
-      font-weight="700">NVIDIA</text
-    >
-  </g>
-{:else if isLinuxLaptop}
-  <!-- Linux Laptop — MacBook shape with Tux logo -->
-  <defs>
-    <clipPath id={linuxScreenClipId}>
-      <rect
-        x={mbScreenX + mbBezel}
-        y={mbY + mbBezel}
-        width={mbScreenW - mbBezel * 2}
-        height={mbScreenH - mbBezel * 2}
-        rx="2"
-      />
-    </clipPath>
-  </defs>
-
-  <rect
-    x={mbScreenX}
-    y={mbY}
-    width={mbScreenW}
-    height={mbScreenH}
-    rx="3"
-    fill="#1a1a1a"
-    stroke={wireColor}
-    stroke-width={strokeWidth}
-  />
-  <rect
-    x={mbScreenX + mbBezel}
-    y={mbY + mbBezel}
-    width={mbScreenW - mbBezel * 2}
-    height={mbScreenH - mbBezel * 2}
-    rx="2"
-    fill="#0a0a12"
-  />
-  {#if ramPercent > 0}
-    <rect
-      x={mbScreenX + mbBezel}
-      y={mbY + mbBezel + (mbMemTotalH - mbMemH)}
-      width={mbScreenW - mbBezel * 2}
-      height={mbMemH}
-      fill="rgba(255,215,0,0.85)"
-      clip-path="url(#{linuxScreenClipId})"
-    />
-  {/if}
-
-  <!-- Terminal prompt on screen -->
-  <text
-    x={cx}
-    y={mbY + mbScreenH / 2}
-    text-anchor="middle"
-    dominant-baseline="middle"
-    fill="#FFFFFF"
-    opacity="0.9"
-    font-size={mbScreenH * 0.25}
-    font-family="SF Mono, Monaco, monospace"
-    font-weight="700">{">_"}</text
-  >
-
-  <path
-    d="M {mbBaseTopX} {mbBaseY} L {mbBaseTopX +
-      mbBaseTopW} {mbBaseY} L {mbBaseBottomX + mbBaseBottomW} {mbBaseY +
-      mbBaseH} L {mbBaseBottomX} {mbBaseY + mbBaseH} Z"
-    fill="#2c2c2c"
-    stroke={wireColor}
-    stroke-width="1"
-  />
-  <rect
-    x={mbKbX}
-    y={mbKbY}
-    width={mbKbW}
-    height={mbKbH}
-    fill="rgba(0,0,0,0.2)"
-    rx="2"
-  />
-  <rect
-    x={mbTpX}
-    y={mbTpY}
-    width={mbTpW}
-    height={mbTpH}
-    fill="rgba(255,255,255,0.08)"
-    rx="2"
-  />
-{:else if isLinux}
-  <!-- Linux Desktop — Mac Studio shape with Tux logo -->
-  <defs>
-    <clipPath id={linuxDesktopClipId}>
-      <rect
-        x={studioX}
-        y={studioY + studioTopH}
-        width={studioW}
-        height={studioH - studioTopH}
-        rx={studioCorner - 1}
-      />
-    </clipPath>
-  </defs>
-
-  <rect
-    x={studioX}
-    y={studioY}
-    width={studioW}
-    height={studioH}
-    rx={studioCorner}
-    fill="#1a1a1a"
-    stroke={wireColor}
-    stroke-width={strokeWidth}
-  />
-  {#if ramPercent > 0}
-    <rect
-      x={studioX}
-      y={studioY + studioTopH + (studioMemTotalH - studioMemH)}
-      width={studioW}
-      height={studioMemH}
-      fill="rgba(255,215,0,0.75)"
-      clip-path="url(#{linuxDesktopClipId})"
-    />
-  {/if}
-
-  <!-- Terminal prompt on front face -->
-  <text
-    x={cx}
-    y={studioY + studioTopH + (studioH - studioTopH) / 2}
-    text-anchor="middle"
-    dominant-baseline="middle"
-    fill="rgba(255,255,255,0.5)"
-    font-size={(studioH - studioTopH) * 0.4}
-    font-family="SF Mono, Monaco, monospace"
-    font-weight="700">{">_"}</text
-  >
-{:else if modelLower === "mac studio" || modelLower === "mac mini"}
+{#if modelLower === "mac studio" || modelLower === "mac mini"}
   <!-- Mac Studio / Mac Mini -->
   <defs>
     <clipPath id={studioClipId}>

dashboard/src/lib/components/ModelCard.svelte

@@ -23,7 +23,7 @@
     } | null;
     nodes?: Record<string, NodeInfo>;
     sharding?: "Pipeline" | "Tensor";
-    runtime?: "MlxRing" | "MlxJaccl" | "Vllm";
+    runtime?: "MlxRing" | "MlxJaccl";
     onLaunch?: () => void;
     tags?: string[];
     apiPreview?: PlacementPreview | null;
@@ -168,10 +168,8 @@
 
   function getDeviceType(
     name: string,
-  ): "macbook" | "studio" | "mini" | "dgx" | "linux" | "unknown" {
+  ): "macbook" | "studio" | "mini" | "unknown" {
     const lower = name.toLowerCase();
-    if (lower.includes("dgx") || lower.includes("gx10")) return "dgx";
-    if (lower.includes("linux")) return "linux";
     if (lower.includes("macbook")) return "macbook";
     if (lower.includes("studio")) return "studio";
     if (lower.includes("mini")) return "mini";
@@ -578,17 +576,13 @@
         class="px-1.5 py-0.5 text-xs font-mono tracking-wider uppercase bg-exo-medium-gray/30 text-exo-light-gray border border-exo-medium-gray/40"
         title={runtime === "MlxRing"
           ? "Ring: standard networking. Works over any connection (Wi-Fi, Ethernet, Thunderbolt)."
-          : runtime === "MlxJaccl"
-            ? "RDMA: direct memory access over Thunderbolt. Significantly faster for multi-device inference."
-            : "vLLM: NVIDIA CUDA inference engine."}
+          : "RDMA: direct memory access over Thunderbolt. Significantly faster for multi-device inference."}
       >
         {runtime === "MlxRing"
           ? "MLX Ring"
           : runtime === "MlxJaccl"
             ? "MLX RDMA"
-            : runtime === "Vllm"
-              ? "vLLM"
-              : runtime}
+            : runtime}
       </span>
     </div>
 
@@ -996,81 +990,6 @@
                     />
                   {/if}
                 </g>
-              {:else if node.deviceType === "dgx"}
-                <!-- DGX Spark icon -->
-                {@const s = node.iconSize}
-                {@const dgxW = s * 1.4}
-                {@const dgxH = s * 0.52}
-                <g transform="translate({-dgxW / 2}, {-dgxH / 2})">
-                  <!-- Chassis -->
-                  <rect
-                    x="0"
-                    y="0"
-                    width={dgxW}
-                    height={dgxH}
-                    rx="2"
-                    fill="#6f6248"
-                    stroke={node.isUsed ? "#FFD700" : "#4B5563"}
-                    stroke-width="1.5"
-                  />
-                  <!-- Side accents -->
-                  <rect
-                    x="0"
-                    y="0"
-                    width={dgxW * 0.02}
-                    height={dgxH}
-                    fill="#8a7a56"
-                  />
-                  <rect
-                    x={dgxW - dgxW * 0.02}
-                    y="0"
-                    width={dgxW * 0.02}
-                    height={dgxH}
-                    fill="#8a7a56"
-                  />
-                  <!-- Left handle -->
-                  <rect
-                    x={dgxW * 0.04}
-                    y={dgxH * 0.08}
-                    width={dgxW * 0.22}
-                    height={dgxH * 0.84}
-                    rx="2"
-                    fill="#b3a170"
-                    stroke="#403723"
-                    stroke-width="0.5"
-                  />
-                  <!-- Right handle -->
-                  <rect
-                    x={dgxW - dgxW * 0.04 - dgxW * 0.22}
-                    y={dgxH * 0.08}
-                    width={dgxW * 0.22}
-                    height={dgxH * 0.84}
-                    rx="2"
-                    fill="#b3a170"
-                    stroke="#403723"
-                    stroke-width="0.5"
-                  />
-                  <!-- Memory fill -->
-                  <rect
-                    x="2"
-                    y={dgxH - dgxH * (node.currentPercent / 100)}
-                    width={dgxW - 4}
-                    height={dgxH * (node.currentPercent / 100)}
-                    fill="rgba(255,215,0,0.35)"
-                  />
-                  {#if node.modelUsageGB > 0 && node.isUsed}
-                    <rect
-                      x="2"
-                      y={dgxH - dgxH * (node.newPercent / 100)}
-                      width={dgxW - 4}
-                      height={dgxH *
-                        ((node.newPercent - node.currentPercent) / 100)}
-                      fill="#FFD700"
-                      filter="url(#memGlow-{filterId})"
-                      class="animate-pulse-slow"
-                    />
-                  {/if}
-                </g>
               {:else}
                 <!-- Unknown device - hexagon -->
                 <g

dashboard/src/lib/components/ModelPickerGroup.svelte

@@ -9,7 +9,6 @@
     capabilities?: string[];
     family?: string;
     is_custom?: boolean;
-    requires_vllm?: boolean;
   }
 
   interface ModelGroup {
@@ -20,7 +19,6 @@
     variants: ModelInfo[];
     smallestVariant: ModelInfo;
     hasMultipleVariants: boolean;
-    requiresVllm: boolean;
   }
 
   type DownloadAvailability = {
@@ -215,14 +213,6 @@
         <span class="font-mono text-sm text-white truncate">
           {group.name}
         </span>
-        {#if group.requiresVllm}
-          <span
-            class="text-[10px] font-mono px-1.5 py-0.5 rounded bg-orange-500/15 text-orange-300 border border-orange-400/30 flex-shrink-0 tracking-wider uppercase"
-            title="Requires vLLM runtime"
-          >
-            vLLM
-          </span>
-        {/if}
         <!-- Capability icons -->
         {#each group.capabilities.filter((c) => c !== "text") as cap}
           {#if cap === "thinking"}
@@ -533,15 +523,6 @@
             {variant.quantization || "default"}
           </span>
 
-          {#if variant.requires_vllm}
-            <span
-              class="text-[10px] font-mono px-1.5 py-0.5 rounded bg-orange-500/15 text-orange-300 border border-orange-400/30 flex-shrink-0 tracking-wider uppercase"
-              title="Requires vLLM runtime"
-            >
-              vLLM
-            </span>
-          {/if}
-
           <!-- Size -->
           <span
             class="text-xs font-mono flex-1 {getSizeClassForFitStatus(
@@ -647,7 +628,6 @@
                 variants: [variant],
                 smallestVariant: variant,
                 hasMultipleVariants: false,
-                requiresVllm: variant.requires_vllm === true,
               });
             }}
             title="View variant details"

dashboard/src/lib/components/ModelPickerModal.svelte

@@ -22,7 +22,6 @@
     is_custom?: boolean;
     tasks?: string[];
     hugging_face_id?: string;
-    requires_vllm?: boolean;
   }
 
   interface ModelGroup {
@@ -33,7 +32,6 @@
     variants: ModelInfo[];
     smallestVariant: ModelInfo;
     hasMultipleVariants: boolean;
-    requiresVllm: boolean;
   }
 
   interface FilterState {
@@ -398,7 +396,6 @@
           variants: [],
           smallestVariant: model,
           hasMultipleVariants: false,
-          requiresVllm: true,
         });
       }
 
@@ -433,7 +430,6 @@
           (a.storage_size_megabytes || 0) - (b.storage_size_megabytes || 0),
       );
       group.hasMultipleVariants = group.variants.length > 1;
-      group.requiresVllm = group.variants.every((v) => v.requires_vllm);
     }
 
     // Convert to array and sort by smallest variant size (biggest first)
@@ -591,7 +587,6 @@
           variants: [model],
           smallestVariant: model,
           hasMultipleVariants: false,
-          requiresVllm: model.requires_vllm === true,
         });
       }
     }
@@ -1170,17 +1165,6 @@
           <span class="text-white/40">Variants:</span>
           <span class="text-white/70">{infoGroup.variants.length}</span>
         </div>
-        {#if infoGroup.requiresVllm}
-          <div class="flex items-center gap-2">
-            <span class="text-white/40">Runtime:</span>
-            <span
-              class="text-[10px] font-mono px-1.5 py-0.5 rounded bg-orange-500/15 text-orange-300 border border-orange-400/30 tracking-wider uppercase"
-            >
-              vLLM
-            </span>
-            <span class="text-white/40 text-[11px]">required</span>
-          </div>
-        {/if}
         {#if infoGroup.variants.length > 0}
           <div class="mt-3 pt-3 border-t border-exo-yellow/10">
             <span class="text-white/40">Available quantizations:</span>

dashboard/src/lib/components/PrefillDecodeDisaggregation.svelte

@@ -219,7 +219,7 @@
         Prefill vs Decode
       </summary>
       <div class="mt-2 text-white/80 text-sm leading-relaxed">
-        Prefill is the compute-bound pass that consumes the entire prompt and
+        Prefill is the compute-heavy pass that consumes the entire prompt and
         builds a KV cache. Decode is the memory-bandwidth-bound loop that emits
         tokens sequentially from that cache. The two phases have very different
         bottlenecks, so running them on different hardware can be substantially

dashboard/src/lib/components/TopologyGraph.svelte

@@ -117,10 +117,6 @@
   const LOGO_NATIVE_WIDTH = 814;
   const LOGO_NATIVE_HEIGHT = 1000;
 
-  // NVIDIA logo SVG path (from exo-nvidia)
-  const NVIDIA_LOGO_PATH =
-    "M0.81 0.429V0.299c0.013 -0.001 0.026 -0.002 0.038 -0.002 0.355 -0.011 0.588 0.306 0.588 0.306S1.186 0.952 0.916 0.952c-0.036 0 -0.071 -0.006 -0.105 -0.017V0.542c0.138 0.017 0.166 0.078 0.249 0.216l0.185 -0.155s-0.135 -0.177 -0.362 -0.177c-0.024 -0.001 -0.048 0.001 -0.072 0.003m0 -0.429v0.194l0.038 -0.002c0.494 -0.017 0.816 0.405 0.816 0.405s-0.37 0.45 -0.754 0.45c-0.034 0 -0.066 -0.003 -0.099 -0.009v0.12c0.027 0.003 0.055 0.006 0.082 0.006 0.358 0 0.618 -0.183 0.869 -0.399 0.042 0.034 0.212 0.114 0.247 0.15 -0.238 0.2 -0.794 0.361 -1.11 0.361 -0.03 0 -0.059 -0.002 -0.088 -0.005v0.169h1.362V0zm0 0.935v0.102c-0.331 -0.059 -0.423 -0.404 -0.423 -0.404s0.159 -0.176 0.423 -0.205v0.112h-0.001C0.671 0.524 0.562 0.654 0.562 0.654s0.062 0.218 0.248 0.282m-0.588 -0.316s0.196 -0.29 0.589 -0.32V0.194C0.376 0.229 0 0.597 0 0.597s0.213 0.616 0.81 0.672v-0.112c-0.438 -0.054 -0.588 -0.538 -0.588 -0.538";
-
   function formatBytes(bytes: number, decimals = 1): string {
     if (!bytes || bytes === 0) return "0B";
     const k = 1024;
@@ -558,13 +554,6 @@
       const clipPathId = `clip-${nodeInfo.id.replace(/[^a-zA-Z0-9]/g, "-")}`;
 
       const modelLower = modelId.toLowerCase();
-      const identity = identitiesData[nodeInfo.id];
-      const nameLower = (friendlyName || "").toLowerCase();
-      const isSpark = modelLower.includes("dgx") || modelLower.includes("gx10");
-      const isLinux =
-        !isSpark &&
-        (modelLower.startsWith("linux") || identity?.osVersion === "Linux");
-      const isLinuxLaptop = isLinux && modelLower.includes("laptop");
 
       // Check node states for styling
       const isHighlighted = highlightedNodes.has(nodeInfo.id);
@@ -634,382 +623,7 @@
           `${friendlyName}\nID: ${nodeInfo.id.slice(-8)}\nMemory: ${formatBytes(ramUsed)}/${formatBytes(ramTotal)}`,
         );
 
-      if (isSpark) {
-        // NVIDIA DGX Spark — gold chassis with textured front, side handles, and NVIDIA badge
-        iconBaseWidth = nodeRadius * 1.55;
-        iconBaseHeight = nodeRadius * 0.58;
-        const x = nodeInfo.x - iconBaseWidth / 2;
-        const y = nodeInfo.y - iconBaseHeight / 2;
-        const chassisX = x - iconBaseWidth * 0.03;
-        const chassisWidth = iconBaseWidth * 1.05;
-        const cornerRadius = 3;
-
-        const dgxClipId = `dgx-clip-${nodeInfo.id.replace(/[^a-zA-Z0-9]/g, "-")}`;
-        defs
-          .append("clipPath")
-          .attr("id", dgxClipId)
-          .append("rect")
-          .attr("x", x)
-          .attr("y", y)
-          .attr("width", iconBaseWidth)
-          .attr("height", iconBaseHeight)
-          .attr("rx", cornerRadius);
-
-        // Chassis texture pattern
-        const textureId = `chassis-texture-${nodeInfo.id.replace(/[^a-zA-Z0-9]/g, "-")}`;
-        defs
-          .append("pattern")
-          .attr("id", textureId)
-          .attr("patternUnits", "userSpaceOnUse")
-          .attr("width", 8)
-          .attr("height", 8);
-        const texturePattern = defs.select(`#${textureId}`);
-        texturePattern
-          .append("rect")
-          .attr("width", 8)
-          .attr("height", 8)
-          .attr("fill", "#6f6248");
-        texturePattern
-          .append("circle")
-          .attr("cx", 2)
-          .attr("cy", 2)
-          .attr("r", 1)
-          .attr("fill", "#5a4f3b")
-          .attr("opacity", 0.5);
-        texturePattern
-          .append("circle")
-          .attr("cx", 6)
-          .attr("cy", 6)
-          .attr("r", 1)
-          .attr("fill", "#4a4232")
-          .attr("opacity", 0.45);
-
-        // Main body
-        nodeG
-          .append("rect")
-          .attr("class", "node-outline")
-          .attr("x", chassisX)
-          .attr("y", y)
-          .attr("width", chassisWidth)
-          .attr("height", iconBaseHeight)
-          .attr("rx", cornerRadius)
-          .attr("fill", `url(#${textureId})`)
-          .attr("stroke", wireColor)
-          .attr("stroke-width", strokeWidth);
-
-        // Side border accents
-        const sideThickness = iconBaseWidth * 0.02;
-        nodeG
-          .append("rect")
-          .attr("x", chassisX)
-          .attr("y", y)
-          .attr("width", sideThickness)
-          .attr("height", iconBaseHeight)
-          .attr("fill", "#8a7a56");
-        nodeG
-          .append("rect")
-          .attr("x", chassisX + chassisWidth - sideThickness)
-          .attr("y", y)
-          .attr("width", sideThickness)
-          .attr("height", iconBaseHeight)
-          .attr("fill", "#8a7a56");
-
-        // Memory fill (bottom up)
-        if (ramUsagePercent > 0) {
-          const memFillHeight = (ramUsagePercent / 100) * iconBaseHeight;
-          nodeG
-            .append("rect")
-            .attr("x", x)
-            .attr("y", y + iconBaseHeight - memFillHeight)
-            .attr("width", iconBaseWidth)
-            .attr("height", memFillHeight)
-            .attr("fill", "rgba(255,215,0,0.45)")
-            .attr("clip-path", `url(#${dgxClipId})`);
-        }
-
-        // Side handles with inner recess
-        const handleWidth = iconBaseWidth * 0.27;
-        const handleGap = iconBaseHeight * 0.05;
-        const handleHeight = iconBaseHeight - handleGap * 2;
-        const handleY = y + handleGap;
-        const innerHandleWidth = iconBaseWidth * 0.12;
-        const innerHandleHeight = handleHeight - iconBaseHeight * 0.06;
-        const leftHandleX = x + 4;
-        const rightHandleX = x + iconBaseWidth - handleWidth - 4;
-
-        // Left handle
-        nodeG
-          .append("rect")
-          .attr("x", leftHandleX)
-          .attr("y", handleY)
-          .attr("width", handleWidth)
-          .attr("height", handleHeight)
-          .attr("rx", 2.4)
-          .attr("fill", "#b3a170")
-          .attr("stroke", "#403723")
-          .attr("stroke-width", 0.7);
-        nodeG
-          .append("rect")
-          .attr("x", leftHandleX + handleWidth * 0.06)
-          .attr("y", handleY + iconBaseHeight * 0.03)
-          .attr("width", innerHandleWidth)
-          .attr("height", innerHandleHeight)
-          .attr("rx", 1.6)
-          .attr("fill", "#8a7a56");
-
-        // Right handle
-        nodeG
-          .append("rect")
-          .attr("x", rightHandleX)
-          .attr("y", handleY)
-          .attr("width", handleWidth)
-          .attr("height", handleHeight)
-          .attr("rx", 2.4)
-          .attr("fill", "#b3a170")
-          .attr("stroke", "#403723")
-          .attr("stroke-width", 0.7);
-        nodeG
-          .append("rect")
-          .attr(
-            "x",
-            rightHandleX + handleWidth - innerHandleWidth - handleWidth * 0.08,
-          )
-          .attr("y", handleY + iconBaseHeight * 0.03)
-          .attr("width", innerHandleWidth)
-          .attr("height", innerHandleHeight)
-          .attr("rx", 1.6)
-          .attr("fill", "#8a7a56");
-
-        // NVIDIA logo + text label (rotated 90 deg on left handle)
-        const badgeWidth = iconBaseWidth * 0.09;
-        const badgeHeight = handleHeight * 0.5;
-        const badgeX =
-          leftHandleX + handleWidth - badgeWidth - handleWidth * 0.06;
-        const badgeY = handleY + (handleHeight - badgeHeight) / 2;
-        const textSize = badgeWidth * 0.58;
-        const logoWidth = textSize * 1.2;
-        const logoHeight = logoWidth * (1.438 / 2.174);
-        const centerX = badgeX + badgeWidth / 2 - badgeWidth * 0.03;
-        const centerY = badgeY + badgeHeight / 2;
-        const gap = badgeWidth * 0.15;
-        const totalWidth = logoWidth + gap + textSize * 3.6;
-
-        const labelGroup = nodeG
-          .append("g")
-          .attr("transform", `rotate(90 ${centerX} ${centerY})`);
-
-        labelGroup
-          .append("svg")
-          .attr("x", centerX - totalWidth / 2)
-          .attr("y", centerY - logoHeight / 2)
-          .attr("width", logoWidth)
-          .attr("height", logoHeight)
-          .attr("viewBox", "0 0 2.174 1.438")
-          .append("path")
-          .attr("d", NVIDIA_LOGO_PATH)
-          .attr("fill", "#76b900");
-
-        labelGroup
-          .append("text")
-          .attr("x", centerX - totalWidth / 2 + logoWidth + gap)
-          .attr("y", centerY)
-          .attr("text-anchor", "start")
-          .attr("dominant-baseline", "middle")
-          .attr("fill", "#8a7a56")
-          .attr("font-size", textSize)
-          .attr("font-family", "monospace")
-          .attr("font-weight", "700")
-          .text("NVIDIA");
-      } else if (isLinuxLaptop) {
-        // Linux Laptop — same shape as MacBook but with Tux logo
-        iconBaseWidth = nodeRadius * 1.6;
-        iconBaseHeight = nodeRadius * 1.15;
-        const x = nodeInfo.x - iconBaseWidth / 2;
-        const y = nodeInfo.y - iconBaseHeight / 2;
-
-        const screenHeight = iconBaseHeight * 0.7;
-        const baseHeight = iconBaseHeight * 0.3;
-        const screenWidth = iconBaseWidth * 0.85;
-        const screenX = nodeInfo.x - screenWidth / 2;
-        const screenBezel = 3;
-
-        const linuxScreenClipId = `linux-screen-${nodeInfo.id.replace(/[^a-zA-Z0-9]/g, "-")}`;
-        defs
-          .append("clipPath")
-          .attr("id", linuxScreenClipId)
-          .append("rect")
-          .attr("x", screenX + screenBezel)
-          .attr("y", y + screenBezel)
-          .attr("width", screenWidth - screenBezel * 2)
-          .attr("height", screenHeight - screenBezel * 2)
-          .attr("rx", 2);
-
-        // Screen outer frame
-        nodeG
-          .append("rect")
-          .attr("class", "node-outline")
-          .attr("x", screenX)
-          .attr("y", y)
-          .attr("width", screenWidth)
-          .attr("height", screenHeight)
-          .attr("rx", 3)
-          .attr("fill", "#1a1a1a")
-          .attr("stroke", wireColor)
-          .attr("stroke-width", strokeWidth);
-
-        // Screen inner
-        nodeG
-          .append("rect")
-          .attr("x", screenX + screenBezel)
-          .attr("y", y + screenBezel)
-          .attr("width", screenWidth - screenBezel * 2)
-          .attr("height", screenHeight - screenBezel * 2)
-          .attr("rx", 2)
-          .attr("fill", "#0a0a12");
-
-        // Memory fill on screen
-        if (ramUsagePercent > 0) {
-          const memFillTotalHeight = screenHeight - screenBezel * 2;
-          const memFillActualHeight =
-            (ramUsagePercent / 100) * memFillTotalHeight;
-          nodeG
-            .append("rect")
-            .attr("x", screenX + screenBezel)
-            .attr(
-              "y",
-              y + screenBezel + (memFillTotalHeight - memFillActualHeight),
-            )
-            .attr("width", screenWidth - screenBezel * 2)
-            .attr("height", memFillActualHeight)
-            .attr("fill", "rgba(255,215,0,0.85)")
-            .attr("clip-path", `url(#${linuxScreenClipId})`);
-        }
-
-        // Terminal prompt on screen
-        nodeG
-          .append("text")
-          .attr("x", nodeInfo.x)
-          .attr("y", y + screenHeight / 2)
-          .attr("text-anchor", "middle")
-          .attr("dominant-baseline", "middle")
-          .attr("fill", "#FFFFFF")
-          .attr("opacity", 0.9)
-          .attr("font-size", screenHeight * 0.25)
-          .attr("font-family", "SF Mono, Monaco, monospace")
-          .attr("font-weight", "700")
-          .text(">_");
-
-        // Keyboard base (trapezoidal)
-        const baseY = y + screenHeight;
-        const baseTopWidth = screenWidth;
-        const baseBottomWidth = iconBaseWidth;
-        const baseTopX = nodeInfo.x - baseTopWidth / 2;
-        const baseBottomX = nodeInfo.x - baseBottomWidth / 2;
-
-        nodeG
-          .append("path")
-          .attr(
-            "d",
-            `M ${baseTopX} ${baseY} L ${baseTopX + baseTopWidth} ${baseY} L ${baseBottomX + baseBottomWidth} ${baseY + baseHeight} L ${baseBottomX} ${baseY + baseHeight} Z`,
-          )
-          .attr("fill", "#2c2c2c")
-          .attr("stroke", wireColor)
-          .attr("stroke-width", 1);
-
-        // Keyboard area
-        const keyboardX = baseTopX + 6;
-        const keyboardY = baseY + 3;
-        const keyboardWidth = baseTopWidth - 12;
-        const keyboardHeight = baseHeight * 0.55;
-        nodeG
-          .append("rect")
-          .attr("x", keyboardX)
-          .attr("y", keyboardY)
-          .attr("width", keyboardWidth)
-          .attr("height", keyboardHeight)
-          .attr("fill", "rgba(0,0,0,0.2)")
-          .attr("rx", 2);
-
-        // Trackpad
-        const trackpadWidth = baseTopWidth * 0.4;
-        const trackpadX = nodeInfo.x - trackpadWidth / 2;
-        const trackpadY = baseY + keyboardHeight + 5;
-        const trackpadHeight = baseHeight * 0.3;
-        nodeG
-          .append("rect")
-          .attr("x", trackpadX)
-          .attr("y", trackpadY)
-          .attr("width", trackpadWidth)
-          .attr("height", trackpadHeight)
-          .attr("fill", "rgba(255,255,255,0.08)")
-          .attr("rx", 2);
-      } else if (isLinux) {
-        // Linux Desktop — same shape as Mac Studio but with Tux logo
-        iconBaseWidth = nodeRadius * 1.25;
-        iconBaseHeight = nodeRadius * 0.85;
-        const x = nodeInfo.x - iconBaseWidth / 2;
-        const y = nodeInfo.y - iconBaseHeight / 2;
-        const cornerRadius = 4;
-        const topSurfaceHeight = iconBaseHeight * 0.15;
-
-        const linuxDesktopClipId = `linux-desktop-${nodeInfo.id.replace(/[^a-zA-Z0-9]/g, "-")}`;
-        defs
-          .append("clipPath")
-          .attr("id", linuxDesktopClipId)
-          .append("rect")
-          .attr("x", x)
-          .attr("y", y + topSurfaceHeight)
-          .attr("width", iconBaseWidth)
-          .attr("height", iconBaseHeight - topSurfaceHeight)
-          .attr("rx", cornerRadius - 1);
-
-        // Main body
-        nodeG
-          .append("rect")
-          .attr("class", "node-outline")
-          .attr("x", x)
-          .attr("y", y)
-          .attr("width", iconBaseWidth)
-          .attr("height", iconBaseHeight)
-          .attr("rx", cornerRadius)
-          .attr("fill", "#1a1a1a")
-          .attr("stroke", wireColor)
-          .attr("stroke-width", strokeWidth);
-
-        // Memory fill
-        if (ramUsagePercent > 0) {
-          const memFillTotalHeight = iconBaseHeight - topSurfaceHeight;
-          const memFillActualHeight =
-            (ramUsagePercent / 100) * memFillTotalHeight;
-          nodeG
-            .append("rect")
-            .attr("x", x)
-            .attr(
-              "y",
-              y + topSurfaceHeight + (memFillTotalHeight - memFillActualHeight),
-            )
-            .attr("width", iconBaseWidth)
-            .attr("height", memFillActualHeight)
-            .attr("fill", "rgba(255,215,0,0.75)")
-            .attr("clip-path", `url(#${linuxDesktopClipId})`);
-        }
-
-        // Terminal prompt on front face
-        nodeG
-          .append("text")
-          .attr("x", nodeInfo.x)
-          .attr(
-            "y",
-            y + topSurfaceHeight + (iconBaseHeight - topSurfaceHeight) / 2,
-          )
-          .attr("text-anchor", "middle")
-          .attr("dominant-baseline", "middle")
-          .attr("fill", "rgba(255,255,255,0.5)")
-          .attr("font-size", (iconBaseHeight - topSurfaceHeight) * 0.4)
-          .attr("font-family", "SF Mono, Monaco, monospace")
-          .attr("font-weight", "700")
-          .text(">_");
-      } else if (modelLower === "mac studio") {
+      if (modelLower === "mac studio") {
         // Mac Studio - classic cube with memory fill
         iconBaseWidth = nodeRadius * 1.25;
         iconBaseHeight = nodeRadius * 0.85;
@@ -1568,12 +1182,8 @@
           debugLabelY += debugLineHeight;
         }
 
-        const dbgIdentity = identitiesData[nodeInfo.id];
-        if (dbgIdentity?.osVersion) {
-          const osLabel =
-            dbgIdentity.osVersion === "Linux"
-              ? "Linux"
-              : `macOS ${dbgIdentity.osVersion}${dbgIdentity.osBuildVersion ? ` (${dbgIdentity.osBuildVersion})` : ""}`;
+        const identity = identitiesData[nodeInfo.id];
+        if (identity?.osVersion) {
           nodeG
             .append("text")
             .attr("x", nodeInfo.x)
@@ -1582,7 +1192,9 @@
             .attr("fill", "rgba(179,179,179,0.7)")
             .attr("font-size", debugFontSize)
             .attr("font-family", "SF Mono, Monaco, monospace")
-            .text(osLabel);
+            .text(
+              `macOS ${identity.osVersion}${identity.osBuildVersion ? ` (${identity.osBuildVersion})` : ""}`,
+            );
         }
       }
     });

dashboard/src/routes/+page.svelte

@@ -65,7 +65,6 @@
     nodeThunderboltBridge,
     nodeIdentities,
     isConnected,
-    featureFlags,
     type DownloadProgress,
     type PlacementPreview,
   } from "$lib/stores/app.svelte";
@@ -703,10 +702,7 @@
       ? Object.keys(topologyData()!.nodes).length
       : 1;
     const sharding = nodeCount <= 1 ? "Pipeline" : selectedSharding;
-    const instanceType =
-      nodeCount <= 1 && selectedInstanceType === "MlxJaccl"
-        ? "MlxRing"
-        : selectedInstanceType;
+    const instanceType = nodeCount <= 1 ? "MlxRing" : selectedInstanceType;
     try {
       const placementResponse = await fetch(
         `/instance/placement?model_id=${encodeURIComponent(modelId)}&sharding=${sharding}&instance_meta=${instanceType}&min_nodes=1`,
@@ -787,7 +783,6 @@
       quantization?: string;
       base_model?: string;
       capabilities?: string[];
-      requires_vllm?: boolean;
     }>
   >([]);
   type ModelMemoryFitStatus =
@@ -891,7 +886,7 @@
   }
 
   let selectedSharding = $state<"Pipeline" | "Tensor">("Pipeline");
-  type InstanceMeta = "MlxRing" | "MlxJaccl" | "Vllm";
+  type InstanceMeta = "MlxRing" | "MlxJaccl";
 
   // Launch defaults persistence
   const LAUNCH_DEFAULTS_KEY = "exo-launch-defaults-v2";
@@ -937,12 +932,7 @@
     // Apply sharding and instance type unconditionally
     selectedSharding = defaults.sharding;
     selectedInstanceType =
-      defaults.instanceType === "MlxRing"
-        ? "MlxRing"
-        : defaults.instanceType === "Vllm"
-          ? "Vllm"
-          : "MlxJaccl";
-    userPickedInstanceType = true;
+      defaults.instanceType === "MlxRing" ? "MlxRing" : "MlxJaccl";
 
     // Apply minNodes if valid (between 1 and maxNodes)
     if (
@@ -964,23 +954,6 @@
   }
 
   let selectedInstanceType = $state<InstanceMeta>("MlxRing");
-  let userPickedInstanceType = $state(false);
-  $effect(() => {
-    if (!userPickedInstanceType && featureFlags()["vllm_available"]) {
-      selectedInstanceType = "Vllm";
-    }
-  });
-  const selectedModelRequiresVllm = $derived.by((): boolean => {
-    const id = selectedPreviewModelId();
-    if (!id) return false;
-    const model = models.find((m) => m.id === id);
-    return model?.requires_vllm === true;
-  });
-  $effect(() => {
-    if (selectedModelRequiresVllm) {
-      selectedInstanceType = "Vllm";
-    }
-  });
   let selectedMinNodes = $state<number>(1);
   let minNodesInitialized = $state(false);
   let launchingModelId = $state<string | null>(null);
@@ -1173,7 +1146,9 @@
   }
 
   const matchesSelectedRuntime = (runtime: InstanceMeta): boolean =>
-    runtime === selectedInstanceType;
+    selectedInstanceType === "MlxRing"
+      ? runtime === "MlxRing"
+      : runtime === "MlxJaccl";
 
   // Helper to check if a model can be launched (has valid placement with >= minNodes)
   function canModelFit(modelId: string): boolean {
@@ -2088,7 +2063,6 @@
     let instanceType = "Unknown";
     if (instanceTag === "MlxRingInstance") instanceType = "MLX Ring";
     else if (instanceTag === "MlxJacclInstance") instanceType = "MLX RDMA";
-    else if (instanceTag === "VllmInstance") instanceType = "vLLM";
 
     const inst = instance as {
       shardAssignments?: {
@@ -3461,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"
@@ -4848,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"
@@ -4994,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
@@ -5795,18 +5772,14 @@
                     </div>
                     <div class="flex gap-2">
                       <button
-                        disabled={selectedModelRequiresVllm}
                         onclick={() => {
-                          if (selectedModelRequiresVllm) return;
                           selectedInstanceType = "MlxRing";
-                          userPickedInstanceType = true;
                           saveLaunchDefaults();
                         }}
-                        class="flex items-center gap-2 py-1.5 px-3 text-xs font-mono border rounded transition-all duration-200 {selectedModelRequiresVllm
-                          ? 'opacity-40 cursor-not-allowed bg-transparent text-white/40 border-exo-medium-gray/30'
-                          : selectedInstanceType === 'MlxRing'
-                            ? 'cursor-pointer bg-transparent text-exo-yellow border-exo-yellow'
-                            : 'cursor-pointer bg-transparent text-white/70 border-exo-medium-gray/50 hover:border-exo-yellow/50'}"
+                        class="flex items-center gap-2 py-1.5 px-3 text-xs font-mono border rounded transition-all duration-200 cursor-pointer {selectedInstanceType ===
+                        'MlxRing'
+                          ? 'bg-transparent text-exo-yellow border-exo-yellow'
+                          : 'bg-transparent text-white/70 border-exo-medium-gray/50 hover:border-exo-yellow/50'}"
                       >
                         <span
                           class="w-3 h-3 rounded-full border-2 flex items-center justify-center {selectedInstanceType ===
@@ -5822,18 +5795,14 @@
                         TCP/IP
                       </button>
                       <button
-                        disabled={selectedModelRequiresVllm}
                         onclick={() => {
-                          if (selectedModelRequiresVllm) return;
                           selectedInstanceType = "MlxJaccl";
-                          userPickedInstanceType = true;
                           saveLaunchDefaults();
                         }}
-                        class="flex items-center gap-2 py-1.5 px-3 text-xs font-mono border rounded transition-all duration-200 {selectedModelRequiresVllm
-                          ? 'opacity-40 cursor-not-allowed bg-transparent text-white/40 border-exo-medium-gray/30'
-                          : selectedInstanceType === 'MlxJaccl'
-                            ? 'cursor-pointer bg-transparent text-exo-yellow border-exo-yellow'
-                            : 'cursor-pointer bg-transparent text-white/70 border-exo-medium-gray/50 hover:border-exo-yellow/50'}"
+                        class="flex items-center gap-2 py-1.5 px-3 text-xs font-mono border rounded transition-all duration-200 cursor-pointer {selectedInstanceType ===
+                        'MlxJaccl'
+                          ? 'bg-transparent text-exo-yellow border-exo-yellow'
+                          : 'bg-transparent text-white/70 border-exo-medium-gray/50 hover:border-exo-yellow/50'}"
                       >
                         <span
                           class="w-3 h-3 rounded-full border-2 flex items-center justify-center {selectedInstanceType ===
@@ -5848,41 +5817,7 @@
                         </span>
                         RDMA (Fast)
                       </button>
-                      {#if featureFlags()["vllm_available"] || selectedModelRequiresVllm}
-                        <button
-                          onclick={() => {
-                            selectedInstanceType = "Vllm";
-                            userPickedInstanceType = true;
-                            saveLaunchDefaults();
-                          }}
-                          class="flex items-center gap-2 py-1.5 px-3 text-xs font-mono border rounded transition-all duration-200 cursor-pointer {selectedInstanceType ===
-                          'Vllm'
-                            ? 'bg-transparent text-exo-yellow border-exo-yellow'
-                            : 'bg-transparent text-white/70 border-exo-medium-gray/50 hover:border-exo-yellow/50'}"
-                        >
-                          <span
-                            class="w-3 h-3 rounded-full border-2 flex items-center justify-center {selectedInstanceType ===
-                            'Vllm'
-                              ? 'border-exo-yellow'
-                              : 'border-exo-medium-gray'}"
-                          >
-                            {#if selectedInstanceType === "Vllm"}
-                              <span
-                                class="w-1.5 h-1.5 rounded-full bg-exo-yellow"
-                              ></span>
-                            {/if}
-                          </span>
-                          vLLM (CUDA)
-                        </button>
-                      {/if}
                     </div>
-                    {#if selectedModelRequiresVllm}
-                      <div
-                        class="mt-2 text-[11px] font-mono text-orange-300/80"
-                      >
-                        This model requires vLLM.
-                      </div>
-                    {/if}
                   </div>
 
                   <!-- Minimum Devices -->

flake.nix

@@ -146,7 +146,7 @@
                 config.treefmt.build.wrapper
 
                 # PYTHON
-                self'.packages.exo.passthru.evenv
+                self'.packages.editableVenv
                 uv
 
                 # RUST

justfile

@@ -40,19 +40,6 @@ build-app: rust-rebuild sync-clean package
     xcodebuild build -project app/EXO/EXO.xcodeproj -scheme EXO -configuration Debug -derivedDataPath app/EXO/build
     @echo "\nBuild complete. Run with:\n  open {{justfile_directory()}}/app/EXO/build/Build/Products/Debug/EXO.app"
 
-sync-cuda:
-    #!/usr/bin/env bash
-    set -euo pipefail
-    uv sync --extra vllm-cuda13 --extra mlx-cpu --no-install-package vllm
-    dest=".venv/lib/python3.13/site-packages"
-    [[ -d $dest/vllm ]] || {
-        nix build .#exo-cuda-13.passthru.evenv
-        # will also grab vllm-0.19.1-distinfo
-        cp -aL result/lib/python3.13/site-packages/vllm* .venv/lib/python3.13/site-packages
-        chmod -R u+rwX .venv/lib/python3.13/site-packages/vllm*
-        rm result
-    }
-
 clean:
     rm -rf **/__pycache__
     rm -rf target/

nix/vllm-setuppy-cmake.patch

@@ -1,26 +0,0 @@
-diff --git a/setup.py b/setup.py
-index 6dc2ed028..bdcc6354a 100644
---- a/setup.py
-+++ b/setup.py
-@@ -18,6 +18,13 @@ from setuptools import Extension, setup
- from setuptools.command.build_ext import build_ext
- 
- 
-+if "NIX_ATTRS_JSON_FILE" in os.environ:
-+    with open(os.environ["NIX_ATTRS_JSON_FILE"], "r") as f:
-+        NIX_ATTRS = json.load(f)
-+else:
-+    NIX_ATTRS = { "cmakeFlags": os.environ.get("cmakeFlags", "").split() }
-+
-+
- def load_module_from_path(module_name, path):
-     spec = importlib.util.spec_from_file_location(module_name, path)
-     module = importlib.util.module_from_spec(spec)
-@@ -184,6 +191,7 @@ class cmake_build_ext(build_ext):
-         cmake_args = [
-             "-DCMAKE_BUILD_TYPE={}".format(cfg),
-             "-DVLLM_TARGET_DEVICE={}".format(VLLM_TARGET_DEVICE),
-+            *NIX_ATTRS["cmakeFlags"],
-         ]
- 
-         verbose = envs.VERBOSE

pyproject.toml

@@ -15,18 +15,21 @@ dependencies = [
   "huggingface-hub>=1.8.0",
   "psutil>=7.0.0",
   "loguru>=0.7.3",
-  "exo-pyo3-bindings",               # rust bindings
+  "exo-pyo3-bindings",                         # rust bindings
   "anyio==4.11.0",
-  "tiktoken>=0.12.0",                # required for kimi k2 tokenizer
+  "mlx==0.31.2; sys_platform == 'darwin'",
+  "mlx-lm; sys_platform=='darwin'",
+  "tiktoken>=0.12.0",                          # required for kimi k2 tokenizer
   "hypercorn>=0.18.0",
   "openai-harmony>=0.0.8",
   "httpx>=0.28.1",
   "tomlkit>=0.14.0",
+  "mflux==0.17.2; sys_platform == 'darwin'",
   "python-multipart>=0.0.21",
   "msgspec>=0.19.0",
   "zstandard>=0.23.0",
+  "mlx-vlm>=0.3.11; sys_platform == 'darwin'",
   "transformers>=5.6.2",
-  "nvidia-ml-py>=13.595.45",
 ]
 
 [project.scripts]
@@ -37,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",
@@ -45,30 +49,26 @@ dev = [
 
 [project.optional-dependencies]
 build = ["nanobind"]
-mlx-none = ["anyio"]
-mlx = [
-  "mlx==0.31.2",
-  "mlx-lm",
-  "mlx-vlm>=0.3.11",
-  "mflux==0.17.5",
-  # pinning vllms versions for consistency.
-  "torch==2.10.0; sys_platform == 'darwin'",
-  "torch==2.10.0; sys_platform == 'linux'",
-  "torchaudio==2.10.0; sys_platform == 'darwin'",
-  "torchaudio==2.10.0; sys_platform == 'linux'",
-  "torchvision==0.25.0; sys_platform == 'darwin'",
-  "torchvision==0.25.0; sys_platform == 'linux'",
-
+cpu = [
+  "mlx==0.31.1; sys_platform == 'linux'",
+  "mlx-cpu==0.31.1; sys_platform == 'linux'",
+  "mlx-lm; sys_platform == 'linux'",
+  "mlx-vlm>=0.3.11; sys_platform== 'linux'",
+  "torch>=2.10.0; sys_platform == 'linux'",
 ]
-mlx-cpu = ["exo[mlx]", "mlx-cpu==0.31.2; sys_platform == 'linux'"]
-mlx-cuda12 = ["exo[mlx]", "mlx-cuda-12==0.31.1; sys_platform == 'linux'"]
-mlx-cuda13 = ["exo[mlx]", "mlx-cuda-13==0.31.1; sys_platform == 'linux'"]
-vllm-none = ["anyio"]
-vllm-cuda13 = [
-  "vllm[cuda13, fastsafetensors]; sys_platform == 'linux'",
-  "torch==2.10.0; sys_platform == 'linux'",
-  "torchaudio==2.10.0; sys_platform == 'linux'",
-  "torchvision==0.25.0; sys_platform == 'linux'",
+cuda12 = [
+  "mlx==0.31.1; sys_platform == 'linux'",
+  "mlx-cuda-12==0.31.1; sys_platform == 'linux'",
+  "mlx-lm; sys_platform == 'linux'",
+  "mlx-vlm>=0.3.11; sys_platform== 'linux'",
+  "torch>=2.10.0; sys_platform == 'linux'",
+]
+cuda13 = [
+  "mlx==0.31.1; sys_platform == 'linux'",
+  "mlx-cuda-13==0.31.1; sys_platform == 'linux'",
+  "mlx-lm; sys_platform == 'linux'",
+  "mlx-vlm>=0.3.11; sys_platform== 'linux'",
+  "torch>=2.10.0; sys_platform == 'linux'",
 ]
 
 ###
@@ -76,29 +76,18 @@ vllm-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 }
 mlx = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git", branch = "address-rdma-gpu-locks", marker = "sys_platform == 'darwin'" }
 mlx-lm = { git = "https://github.com/rltakashige/mlx-lm", branch = "leo/deepseek-v4" }
-mflux = { git = "http://github.com/evanev7/mflux", branch = "exo" }
-vllm = { git = "http://github.com/evanev7/vllm", branch = "exo2" }
 torch = [
-  { index = "pytorch-cpu", marker = "sys_platform == 'linux' and extra == 'mlx-cpu' and extra != 'vllm-cuda13' and extra != 'mlx-cuda13' and extra != 'mlx-cuda12'" },
-  { index = "pytorch-cu128", marker = "sys_platform == 'linux' and extra == 'mlx-cuda12' and extra != 'mlx-cuda13' and extra != 'vllm-cuda13'" },
-  { index = "pytorch-cu130", marker = "sys_platform == 'linux' and (extra == 'mlx-cuda13' or extra == 'vllm-cuda13')" },
-]
-torchvision = [
-  { index = "pytorch-cpu", marker = "sys_platform == 'linux' and extra == 'mlx-cpu' and extra != 'vllm-cuda13' and extra != 'mlx-cuda13' and extra != 'mlx-cuda12'" },
-  { index = "pytorch-cu128", marker = "sys_platform == 'linux' and extra == 'mlx-cuda12' and extra != 'mlx-cuda13' and extra != 'vllm-cuda13'" },
-  { index = "pytorch-cu130", marker = "sys_platform == 'linux' and (extra == 'mlx-cuda13' or extra == 'vllm-cuda13')" },
-]
-torchaudio = [
-  { index = "pytorch-cpu", marker = "sys_platform == 'linux' and extra == 'mlx-cpu' and extra != 'vllm-cuda13' and extra != 'mlx-cuda13' and extra != 'mlx-cuda12'" },
-  { index = "pytorch-cu128", marker = "sys_platform == 'linux' and extra == 'mlx-cuda12' and extra != 'mlx-cuda13' and extra != 'vllm-cuda13'" },
-  { index = "pytorch-cu130", marker = "sys_platform == 'linux' and (extra == 'mlx-cuda13' or extra == 'vllm-cuda13')" },
+  { index = "pytorch-cu130", marker = "sys_platform == 'linux' and extra == 'cuda13' and extra != 'cpu' and extra != 'cuda12'" },
+  { index = "pytorch-cu120", marker = "sys_platform == 'linux' and extra == 'cuda12' and extra != 'cpu' and extra != 'cuda13'" },
+  { index = "pytorch-cpu", marker = "(extra != 'cuda12' and extra != 'cuda13' and sys_platform == 'linux') or sys_platform == 'darwin'" },
 ]
+vllm = { git = "https://github.com/hmellor/vllm.git", branch = "transformers-v5" }
 
 [[tool.uv.index]]
 name = "pytorch-cu130"
@@ -106,8 +95,8 @@ url = "https://download.pytorch.org/whl/cu130"
 explicit = true
 
 [[tool.uv.index]]
-name = "pytorch-cu128"
-url = "https://download.pytorch.org/whl/cu128"
+name = "pytorch-cu120"
+url = "https://download.pytorch.org/whl/cu120"
 explicit = true
 
 [[tool.uv.index]]
@@ -124,7 +113,7 @@ build-backend = "uv_build"
 ###
 
 [tool.basedpyright]
-include = ["src", "bench"]
+include = ["src", "bench", "tools"]
 typeCheckingMode = "strict"
 failOnWarnings = true
 
@@ -158,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
@@ -168,19 +166,11 @@ root = "src"
 required-version = ">=0.8.6"
 prerelease = "allow"
 environments = ["sys_platform == 'darwin'", "sys_platform == 'linux'"]
-override-dependencies = ["opencv-python; python_version < '0'"]
-conflicts = [
-  [
-    { extra = "mlx-cuda13" },
-    { extra = "mlx-cuda12" },
-    { extra = "mlx-cpu" },
-    { extra = "mlx-none" },
-  ],
-  [
-    { extra = "vllm-cuda13" },
-    { extra = "mlx-cuda12" },
-    { extra = "vllm-none" },
-  ],
+conflicts = [[{ extra = "cuda12" }, { extra = "cuda13" }, { extra = "cpu" }]]
+constraint-dependencies = ["transformers>=5.6.2"]
+override-dependencies = [
+  "mlx==0.31.1; sys_platform=='linux'",
+  "mlx; sys_platform=='darwin'",
 ]
 
 [tool.uv.extra-build-dependencies]
@@ -195,7 +185,6 @@ mlx = [
   "ninja",
 ]
 mlx-lm = ["setuptools"]
-mflux = ["uv_build"]
 xgrammar = [
   "nanobind",
   "setuptools",
@@ -237,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"]

python/parts.nix

@@ -10,10 +10,8 @@ let
       inherit (pkgs.stdenv.hostPlatform) isLinux isDarwin isx86_64;
       inherit (pkgs.config) cudaSupport;
       inherit (pkgs) cudaPackages;
-      libmlx_source =
-        if (builtins.elem "mlx-cuda13" members.exo or [ ]) then "mlx-cuda-13"
-        else if (builtins.elem "mlx-cuda12" members.exo or [ ]) then "mlx-cuda-12"
-        else "mlx-cpu";
+      cuda13Support = cudaSupport && cudaPackages.cudaMajorVersion == "13";
+      libmlx_source = if cuda13Support then "mlx-cuda-13" else if cudaSupport then "mlx-cuda-12" else "mlx-cpu";
       python = pkgs.python313;
       cudaLibs = with cudaPackages; [
         cuda_cudart
@@ -115,213 +113,37 @@ let
               });
           } // lib.optionalAttrs isLinux {
           mlx = prev.mlx.overrideAttrs (old: {
-            nativeBuildInputs = old.nativeBuildInputs ++ lib.optionals cudaSupport [ pkgs.autoAddDriverRunpath ];
             buildInputs = old.buildInputs ++ lib.optionals cudaSupport cudaLibs;
+            autoPatchelfIgnoreMissingDeps = lib.optionals cudaSupport [ "libcuda.so.1" ];
             postInstall = ''
               cp -r "${final.${libmlx_source}}/${final.python.sitePackages}/mlx" "$out/${final.python.sitePackages}/mlx/"
             '';
-            autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
           });
         } // lib.optionalAttrs cudaSupport {
           "${libmlx_source}" = prev."${libmlx_source}".overrideAttrs (old: {
-            nativeBuildInputs = old.nativeBuildInputs ++ [ pkgs.autoAddDriverRunpath ];
             buildInputs = old.buildInputs ++ cudaLibs;
             autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
           });
           nvidia-cufile = prev.nvidia-cufile.overrideAttrs (old: {
-            nativeBuildInputs = old.nativeBuildInputs ++ [ pkgs.autoAddDriverRunpath ];
             buildInputs = old.buildInputs ++ [ pkgs.rdma-core ];
+            autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
           });
           nvidia-cusolver = prev.nvidia-cusolver.overrideAttrs (old: {
-            nativeBuildInputs = old.nativeBuildInputs ++ [ pkgs.autoAddDriverRunpath ];
             buildInputs = old.buildInputs ++ cudaLibs;
+            autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
           });
           nvidia-nvshmem-cu13 = prev.nvidia-nvshmem-cu13.overrideAttrs (old: {
-            nativeBuildInputs = old.nativeBuildInputs ++ [ pkgs.autoAddDriverRunpath ];
             buildInputs = old.buildInputs ++ [ pkgs.rdma-core pkgs.pmix pkgs.libfabric pkgs.ucx pkgs.openmpi ];
+            autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
           });
           nvidia-cusparse = prev.nvidia-cusparse.overrideAttrs (old: {
-            nativeBuildInputs = old.nativeBuildInputs ++ [ pkgs.autoAddDriverRunpath ];
-            buildInputs = old.buildInputs ++ cudaLibs;
+            buildInputs = old.buildInputs ++ [ cudaLibs ];
+            autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
           });
           torch = prev.torch.overrideAttrs (old: {
-            nativeBuildInputs = old.nativeBuildInputs ++ [ pkgs.autoAddDriverRunpath ];
             buildInputs = old.buildInputs ++ cudaLibs;
             autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
           });
-          torchaudio = prev.torchaudio.overrideAttrs (old: {
-            nativeBuildInputs = old.nativeBuildInputs ++ [ pkgs.autoAddDriverRunpath ];
-            buildInputs = old.buildInputs ++ [ cudaPackages.cuda_cudart ];
-            preFixup = "addAutoPatchelfSearchPath '${final.torch}'";
-          });
-          torchvision = prev.torchvision.overrideAttrs (old: {
-            nativeBuildInputs = old.nativeBuildInputs ++ [ pkgs.autoAddDriverRunpath ];
-            preFixup = "addAutoPatchelfSearchPath '${final.torch}'";
-          });
-
-          torch-c-dlpack-ext = prev.torch-c-dlpack-ext.overrideAttrs (old: {
-            buildInputs = old.buildInputs ++ cudaLibs;
-            autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
-            preFixup = "addAutoPatchelfSearchPath '${final.torch}'";
-          });
-          # Currently treating vllm as a cuda dep. it obviously exists as a non cuda dep
-          vllm = prev.vllm.overrideAttrs (old:
-            let
-              cuda_cccl_compat = pkgs.runCommand "cuda-cccl-compat" { } ''
-                mkdir -p $out/include
-                ln -s ${cudaPackages.cuda_cccl}/include $out/include/cccl
-              '';
-
-              cudaRoot = pkgs.symlinkJoin {
-                name = "cuda-merged-exo";
-                paths = builtins.concatMap (p: [ (lib.getBin p) (lib.getLib p) (lib.getDev p) ]) (cudaLibs ++ [ cudaPackages.cuda_nvcc cuda_cccl_compat ]);
-              };
-
-              cutlass = pkgs.fetchFromGitHub {
-                name = "cutlass-source";
-                owner = "NVIDIA";
-                repo = "cutlass";
-                tag = "v4.2.1";
-                hash = "sha256-iP560D5Vwuj6wX1otJhwbvqe/X4mYVeKTpK533Wr5gY=";
-              };
-              triton-kernels = pkgs.fetchFromGitHub {
-                owner = "triton-lang";
-                repo = "triton";
-                tag = "v3.6.0";
-                hash = "sha256-JFSpQn+WsNnh7CAPlcpOcUp0nyKXNbJEANdXqmkt4Tc=";
-              };
-
-              cutlass-flashmla = pkgs.fetchFromGitHub {
-                owner = "NVIDIA";
-                repo = "cutlass";
-                rev = "147f5673d0c1c3dcf66f78d677fd647e4a020219";
-                hash = "sha256-dHQto08IwTDOIuFUp9jwm1MWkFi8v2YJ/UESrLuG71g=";
-              };
-
-              flashmla = pkgs.stdenv.mkDerivation {
-                pname = "flashmla";
-                version = "1.0.0";
-
-                src = pkgs.fetchFromGitHub {
-                  name = "FlashMLA-source";
-                  owner = "vllm-project";
-                  repo = "FlashMLA";
-                  rev = "c2afa9cb93e674d5a9120a170a6da57b89267208";
-                  hash = "sha256-pKlwxV6G9iHag/jbu3bAyvYvnu5TbrQwUMFV0AlGC3s=";
-                };
-
-                dontConfigure = true;
-
-                buildPhase = ''
-                  rm -rf csrc/cutlass
-                  ln -sf ${cutlass-flashmla} csrc/cutlass
-                '';
-
-                installPhase = ''
-                  cp -rva . $out
-                '';
-              };
-              qutlass = pkgs.fetchFromGitHub {
-                name = "qutlass-source";
-                owner = "IST-DASLab";
-                repo = "qutlass";
-                rev = "830d2c4537c7396e14a02a46fbddd18b5d107c65";
-                hash = "sha256-aG4qd0vlwP+8gudfvHwhtXCFmBOJKQQTvcwahpEqC84=";
-              };
-              vllm-flash-attn = pkgs.stdenv.mkDerivation {
-                pname = "vllm-flash-attn";
-                version = "2.7.2.post1";
-
-                src = pkgs.fetchFromGitHub {
-                  name = "flash-attention-source";
-                  owner = "vllm-project";
-                  repo = "flash-attention";
-                  rev = "188be16520ceefdc625fdf71365585d2ee348fe2";
-                  hash = "sha256-Osec+/IF3+UDtbIhDMBXzUeWJ7hDJNb5FpaVaziPSgM=";
-                };
-
-                patches = [
-                  (pkgs.fetchpatch {
-                    url = "https://github.com/Dao-AILab/flash-attention/commit/dad67c88d4b6122c69d0bed1cebded0cded71cea.patch";
-                    hash = "sha256-JSgXWItOp5KRpFbTQj/cZk+Tqez+4mEz5kmH5EUeQN4=";
-                  })
-                  (pkgs.fetchpatch {
-                    url = "https://github.com/Dao-AILab/flash-attention/commit/e26dd28e487117ee3e6bc4908682f41f31e6f83a.patch";
-                    hash = "sha256-NkCEowXSi+tiWu74Qt+VPKKavx0H9JeteovSJKToK9A=";
-                  })
-                ];
-
-                dontConfigure = true;
-
-                buildPhase = ''
-                  rm -rf csrc/cutlass
-                  ln -sf ${cutlass} csrc/cutlass
-                '';
-
-                installPhase = ''
-                  cp -rva . $out
-                '';
-              };
-            in
-            {
-              patches = (old.patches or [ ]) ++ [ ../nix/vllm-setuppy-cmake.patch ];
-              nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [
-                pkgs.cmake
-                pkgs.ninja
-                pkgs.autoAddDriverRunpath
-              ] ++ lib.optionals cudaSupport [
-                cudaPackages.cuda_nvcc
-              ];
-              # TODO: vllm rocm/cpu
-              VLLM_TARGET_DEVICE = "empty";
-              preConfigure = ''
-                export MAX_JOBS="$NIX_BUILD_CORES"
-              '';
-
-              # TODO: vllm non cuda13 support, more arch's, etc.
-            } // lib.optionalAttrs cudaSupport {
-              buildInputs = cudaLibs ++ [ cudaRoot ];
-
-              VLLM_CUDA_VERSION = cudaPackages.cudaMajorMinorVersion;
-              CUDA_HOME = "${cudaRoot}";
-              CUDAToolkit_ROOT = "${cudaRoot}";
-              CUDACXX = "${cudaRoot}/bin/nvcc";
-              VLLM_CUTLASS_SRC_DIR = "${lib.getDev cutlass}";
-              VLLM_TARGET_DEVICE = "cuda";
-              TORCH_CUDA_ARCH_LIST = "12.0;12.1";
-              TRITON_KERNELS_SRC_DIR = "${lib.getDev triton-kernels}/python/triton_kernels/triton_kernels";
-              FLASH_MLA_SRC_DIR = "${lib.getDev flashmla}";
-              QUTLASS_SRC_DIR = "${lib.getDev qutlass}";
-              VLLM_FLASH_ATTN_SRC_DIR = "${lib.getDev vllm-flash-attn}";
-              CAFFE2_USE_CUDNN = "ON";
-              CAFFE2_USE_CUFILE = "ON";
-              CUTLASS_ENABLE_CUBLAS = "ON";
-              CUTLASS_NVCC_ARCHS_ENABLED = "12.0;12.1";
-
-              cmakeFlags = [
-                (lib.cmakeBool "CMAKE_SKIP_INSTALL_RPATH" true)
-                (lib.cmakeBool "CMAKE_BUILD_WITH_INSTALL_RPATH" true)
-                (lib.cmakeFeature "CUDA_HOME" "${cudaRoot}")
-                (lib.cmakeFeature "CUDAToolkit_ROOT" "${cudaRoot}")
-                (lib.cmakeFeature "CMAKE_CUDA_COMPILER" "${cudaRoot}/bin/nvcc")
-                (lib.cmakeFeature "CMAKE_PREFIX_PATH" "${cudaRoot}")
-                (lib.cmakeFeature "FETCHCONTENT_SOURCE_DIR_CUTLASS" "${lib.getDev cutlass}")
-                (lib.cmakeFeature "FLASH_MLA_SRC_DIR" "${lib.getDev flashmla}")
-                (lib.cmakeFeature "VLLM_FLASH_ATTN_SRC_DIR" "${lib.getDev vllm-flash-attn}")
-                (lib.cmakeFeature "QUTLASS_SRC_DIR" "${lib.getDev qutlass}")
-                (lib.cmakeFeature "TORCH_CUDA_ARCH_LIST" "12.0;12.1")
-                (lib.cmakeFeature "CUTLASS_NVCC_ARCHS_ENABLED" "${cudaPackages.flags.cmakeCudaArchitecturesString}")
-                (lib.cmakeFeature "CUDA_TOOLKIT_ROOT_DIR" "${cudaRoot}")
-                (lib.cmakeFeature "CAFFE2_USE_CUDNN" "ON")
-                (lib.cmakeFeature "CAFFE2_USE_CUFILE" "ON")
-                (lib.cmakeFeature "CUTLASS_ENABLE_CUBLAS" "ON")
-              ];
-            });
-
-        } // lib.optionalAttrs (cudaSupport && isx86_64) {
-          numba = prev.numba.overrideAttrs (old: {
-            buildInputs = (old.buildInputs or [ ]) ++ [ pkgs.tbb ];
-          });
         };
       pyprojectOverlay = workspace.mkPyprojectOverlay {
         sourcePreference = "wheel";
@@ -342,28 +164,24 @@ let
           buildSystemsOverlay
         ]
       );
-      # mlx and mlx-cuda ship clashing cmake files - we dont need them at runtime anyway
-      venv = name: (pythonSet.mkVirtualEnv "${name}-venv" members).overrideAttrs (_: { venvSkip = [ "lib/python${python.pythonVersion}/site-packages/mlx/share/cmake/*" "lib/python${python.pythonVersion}/site-packages/build_backend.py" ]; });
-      mkApp = text: name: pkgs.writeShellApplication {
+      venv = name: (pythonSet.mkVirtualEnv "${name}-env" members).overrideAttrs (_: { venvSkip = [ "lib/python${python.pythonVersion}/site-packages/mlx/share/cmake/*" ]; });
+      mkApp = cmd: name: pkgs.writeShellApplication {
         inherit name;
-        text = "exec " + lib.optionalString cudaSupport "nixglhost " + text;
         runtimeEnv = {
           EXO_DASHBOARD_DIR = self'.packages.dashboard;
           EXO_RESOURCES_DIR = inputs.self + /resources;
         };
         runtimeInputs = [
+          # mlx and mlx-cuda ship clashing cmake files - we dont need them at runtime anyway
           (venv name)
-          pkgs.nix-gl-host
         ]
         ++ lib.optionals isDarwin [ pkgs.macmon ];
-        passthru = {
-          venv = venv name;
-          evenv = ((pythonSet.overrideScope editableOverlay).mkVirtualEnv "${name}-evenv" (members // { exo = (members.exo or [ ]) ++ [ "dev" ]; })).overrideAttrs (_: { venvSkip = [ "lib/python${python.pythonVersion}/site-packages/mlx/share/cmake/*" "lib/python${python.pythonVersion}/site-packages/build_backend.py" ]; });
-        };
+        text = "exec " + lib.optionalString cudaSupport "${lib.getExe pkgs.nix-gl-host} " + cmd;
       };
     in
     {
       inherit venv;
+      editablePythonSet = pythonSet.overrideScope editableOverlay;
       mkPythonScript = path: mkApp ''python ${path} "$@"'';
       mkExo = mkApp ''exo "$@"'';
     };
@@ -373,18 +191,18 @@ in
     { self', pkgs, unfreePkgs, lib, ... }:
     let
       inherit (pkgs.stdenv.hostPlatform) isLinux;
-      inherit (mkPythonSet { inherit self' pkgs lib; members = { exo = [ "mlx-cpu" "vllm-none" ]; }; }) mkExo;
+      inherit (mkPythonSet { inherit self' pkgs lib; members = { exo = [ "cpu" ]; }; }) editablePythonSet mkExo;
 
       # Virtual environment with dev dependencies for testing
       testVenv = (mkPythonSet {
         inherit self' pkgs lib; members = {
-        exo = [ "dev" "mlx-cpu" "vllm-none" ]; # Include pytest, pytest-asyncio, pytest-env
+        exo = [ "dev" "cpu" ]; # Include pytest, pytest-asyncio, pytest-env
       };
       }).venv "exo-test";
 
       mkBenchScript = (mkPythonSet {
         inherit self' pkgs lib; members = {
-        exo = [ "mlx-cpu" "vllm-none" ];
+        exo = [ "cpu" ];
         exo-bench = [ ]; # Include pytest, pytest-asyncio, pytest-env
       };
       }).mkPythonScript;
@@ -394,12 +212,12 @@ in
         runtimeInputs = [ pkgs.python313 ];
         text = ''exec python ${path} "$@"'';
       };
-      cuda12Set = mkPythonSet { inherit self' lib; inherit (unfreePkgs.pkgsCuda.cudaPackages_12) pkgs; members = { exo = [ "mlx-cuda12" "vllm-none" ]; }; };
-      cuda13Set = mkPythonSet { inherit self' lib; inherit (unfreePkgs.pkgsCuda.cudaPackages_13) pkgs; members = { exo = [ "mlx-cpu" "vllm-cuda13" ]; }; };
+
     in
     {
       packages = {
         exo = mkExo "exo";
+        editableVenv = editablePythonSet.mkVirtualEnv "exo-dev-env" { exo = [ "dev" ]; };
         # for running tests in ci
         exo-test-env = testVenv;
         exo-bench = mkBenchScript "exo-bench" (inputs.self + /bench/exo_bench.py);
@@ -408,8 +226,8 @@ in
         # used by ./tests/run_exo_on.sh
         exo-get-all-models-on-cluster = mkSimplePythonScript "exo-get-all-models-on-cluster" (inputs.self + /tests/get_all_models_on_cluster.py);
       } // lib.optionalAttrs isLinux {
-        exo-cuda-12 = cuda12Set.mkExo "exo-cuda-12";
-        exo-cuda-13 = cuda13Set.mkExo "exo-cuda-13";
+        exo-cuda-12 = (mkPythonSet { inherit self' lib; inherit (unfreePkgs.pkgsCuda.cudaPackages_12) pkgs; members = { exo = [ "cuda12" ]; }; }).mkExo "exo-cuda-12";
+        exo-cuda-13 = (mkPythonSet { inherit self' lib; inherit (unfreePkgs.pkgsCuda.cudaPackages_13) pkgs; members = { exo = [ "cuda13" ]; }; }).mkExo "exo-cuda-13";
       };
 
       checks = {

resources/inference_model_cards/2imi9--gpt-oss-20B-NVFP4A16-BF16.toml

@@ -1,21 +0,0 @@
-model_id = "2imi9/gpt-oss-20B-NVFP4A16-BF16"
-n_layers = 24
-hidden_size = 2880
-num_key_value_heads = 8
-supports_tensor = false
-tasks = ["TextGeneration"]
-family = "gpt-oss"
-quantization = "nvfp4"
-base_model = "GPT-OSS 20B"
-capabilities = ["text", "thinking"]
-reasoning_dialect = "channel"
-context_length = 131072
-requires_vllm = true
-
-[storage_size]
-in_bytes = 41829514752
-
-[sampling_defaults]
-temperature = 1.0
-top_p = 1.0
-top_k = 0

resources/inference_model_cards/nvidia--Qwen3-30B-A3B-NVFP4.toml

@@ -1,27 +0,0 @@
-model_id = "nvidia/Qwen3-30B-A3B-NVFP4"
-n_layers = 48
-hidden_size = 2048
-num_key_value_heads = 4
-supports_tensor = false
-tasks = ["TextGeneration"]
-family = "qwen"
-quantization = "nvfp4"
-base_model = "Qwen3 30B"
-capabilities = ["text", "thinking", "thinking_toggle"]
-context_length = 32768
-requires_vllm = true
-
-[storage_size]
-in_bytes = 18087458688
-
-[sampling_defaults]
-temperature = 0.6
-top_p = 0.95
-top_k = 20
-min_p = 0.0
-
-[sampling_defaults.non_thinking]
-temperature = 0.7
-top_p = 0.8
-top_k = 20
-min_p = 0.0

resources/inference_model_cards/openai--gpt-oss-120b.toml

@@ -1,20 +0,0 @@
-model_id = "openai/gpt-oss-120b"
-n_layers = 36
-hidden_size = 2880
-num_key_value_heads = 8
-supports_tensor = false
-tasks = ["TextGeneration"]
-family = "gpt-oss"
-quantization = "mxfp4"
-base_model = "GPT-OSS 120B"
-capabilities = ["text", "thinking"]
-reasoning_dialect = "channel"
-context_length = 131072
-
-[storage_size]
-in_bytes = 65248815744
-
-[sampling_defaults]
-temperature = 1.0
-top_p = 1.0
-top_k = 0

resources/inference_model_cards/sakamakismile--Qwen3.6-27B-NVFP4.toml

@@ -1,32 +0,0 @@
-model_id = "sakamakismile/Qwen3.6-27B-NVFP4"
-n_layers = 64
-hidden_size = 5120
-num_key_value_heads = 4
-supports_tensor = false
-tasks = ["TextGeneration"]
-family = "qwen"
-quantization = "nvfp4"
-base_model = "Qwen3.6 27B"
-capabilities = ["text", "thinking", "thinking_toggle", "vision"]
-reasoning_dialect = "post_last_user"
-context_length = 262144
-requires_vllm = true
-
-[storage_size]
-in_bytes = 16703361232
-
-[sampling_defaults]
-temperature = 1.0
-top_p = 0.95
-top_k = 20
-min_p = 0.0
-repetition_penalty = 1.0
-presence_penalty = 1.5
-
-[sampling_defaults.non_thinking]
-temperature = 0.7
-top_p = 0.8
-top_k = 20
-min_p = 0.0
-repetition_penalty = 1.0
-presence_penalty = 1.5

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)

scripts/check_serve_prefill.py

@@ -1,222 +0,0 @@
-#!/usr/bin/env python
-"""Standalone smoke test for VllmEngine.serve_prefill.
-
-Loads a real vLLM engine, runs serve_prefill against an in-memory buffer
-twice in a row with the same prompt, and verifies both runs produce a
-well-formed wire stream (header -> KV chunks -> Done).
-
-The second run is the regression guard: with vLLM APC enabled this would
-trip the chunked-prefill + APC + custom kv-connector CUDA assert
-(`vectorized_gather_kernel: ind >= ind_dim_size`) and the server would
-close the socket before the Done frame.
-
-Usage on the Spark (gx10-de89):
-
-    cd /home/larry/exo
-    /nix/store/2b82iz9ac0pxqafrgxmgdkq8sr2hwlx6-exo-cuda-13-venv/bin/python \\
-        scripts/check_serve_prefill.py Qwen/Qwen3-0.6B
-
-Exits 0 on success, non-zero with a diagnostic on failure.
-"""
-
-from __future__ import annotations
-
-import contextlib
-import io
-import os
-import sys
-import traceback
-from pathlib import Path
-from typing import cast
-
-
-def _ensure_repo_on_path() -> None:
-    repo = Path(__file__).resolve().parent.parent
-    src = repo / "src"
-    if str(src) not in sys.path:
-        sys.path.insert(0, str(src))
-
-
-_ensure_repo_on_path()
-
-from exo.shared.types.common import ModelId  # noqa: E402
-from exo.worker.disaggregated.protocol import (  # noqa: E402
-    ArraysState,
-    Done,
-    ErrorMessage,
-    KVChunk,
-    read_header,
-    read_message,
-)
-from exo.worker.disaggregated.server import PrefillRequest  # noqa: E402
-
-
-def _make_token_ids(n: int) -> list[int]:
-    return [(i * 1009 + 17) % 30000 + 100 for i in range(n)]
-
-
-def _decode(
-    payload: bytes,
-) -> tuple[list[KVChunk], list[ArraysState], Done | None, ErrorMessage | None]:
-    buf = io.BytesIO(payload)
-    _ = read_header(buf)
-    chunks: list[KVChunk] = []
-    arrays: list[ArraysState] = []
-    done: Done | None = None
-    error: ErrorMessage | None = None
-    while True:
-        msg = read_message(buf)
-        if msg is None:
-            break
-        if isinstance(msg, KVChunk):
-            chunks.append(msg)
-        elif isinstance(msg, ArraysState):
-            arrays.append(msg)
-        elif isinstance(msg, Done):
-            done = msg
-            break
-        elif isinstance(msg, ErrorMessage):
-            error = msg
-            break
-    return chunks, arrays, done, error
-
-
-def _build_engine(model_id: ModelId) -> object:
-    from exo.worker.engines.vllm.engine import VllmEngine
-    from exo.worker.engines.vllm.generator import VllmBatchEngine, load_vllm_engine
-    from exo.worker.engines.vllm.kv_connector import (
-        ExoKVProducerConnector,
-        _patch_gdn_capture,
-        _patch_vllm_for_connector,
-    )
-
-    _patch_vllm_for_connector(ExoKVProducerConnector)
-    _patch_gdn_capture()
-
-    llm_engine, tool_parser = load_vllm_engine(
-        model_id=model_id,
-        trust_remote_code=False,
-        n_layers=1,
-        kv_connector_cls=ExoKVProducerConnector,
-    )
-    gen = VllmBatchEngine(engine=llm_engine, model_id=model_id)
-
-    class _S:
-        def send(self, _: object) -> None: ...
-
-    class _R:
-        def collect(self) -> list[object]:
-            return []
-
-    return VllmEngine(
-        tool_parser=tool_parser,
-        model_id=model_id,
-        cancel_receiver=cast("object", _R()),  # pyright: ignore[reportArgumentType]
-        event_sender=cast("object", _S()),  # pyright: ignore[reportArgumentType]
-        _gen=gen,
-        max_concurrent_requests=1,
-    )
-
-
-def _run_one(engine: object, n_tokens: int, label: str) -> int:
-    request = PrefillRequest(
-        request_id=f"check-{label}-{os.getpid()}",
-        model_id="ignored",
-        token_ids=_make_token_ids(n_tokens),
-        start_pos=0,
-        use_prefix_cache=True,
-    )
-    buf = io.BytesIO()
-    engine.serve_prefill(request, buf)  # pyright: ignore[reportAttributeAccessIssue]
-    payload = buf.getvalue()
-    if not payload:
-        raise AssertionError(f"{label}: server wrote nothing")
-
-    chunks, arrays, done, error = _decode(payload)
-    if error is not None:
-        raise AssertionError(
-            f"{label}: server returned ErrorMessage [{error.code}]: {error.message}"
-        )
-    if done is None:
-        raise AssertionError(
-            f"{label}: stream did not end with Done "
-            f"({len(chunks)} kv chunks, {len(arrays)} arrays)"
-        )
-    if done.total_tokens <= 0:
-        raise AssertionError(f"{label}: Done reported {done.total_tokens} tokens")
-    if not chunks:
-        raise AssertionError(f"{label}: no KV chunks shipped")
-
-    expected = max(0, n_tokens - 2)
-    if done.total_tokens < expected - 64:
-        raise AssertionError(
-            f"{label}: got {done.total_tokens} tokens, expected ~{expected}"
-        )
-    print(
-        f"  [{label}] OK: tokens={done.total_tokens} "
-        f"kv_chunks={len(chunks)} arrays={len(arrays)}"
-    )
-    return done.total_tokens
-
-
-def main(argv: list[str]) -> int:
-    if len(argv) < 2:
-        print(__doc__)
-        return 2
-    model_id = ModelId(argv[1])
-
-    from exo.download.download_utils import build_model_path
-
-    model_path = build_model_path(model_id)
-    if not model_path.exists():
-        print(f"FAIL: model {model_id} not found at {model_path}")
-        return 1
-    print(f"Loading vLLM engine for {model_id} ({model_path}) ...")
-
-    engine = _build_engine(model_id)
-    failures: list[str] = []
-    try:
-        try:
-            t1 = _run_one(engine, n_tokens=512, label="run1-fresh")
-        except AssertionError as e:
-            failures.append(f"run1: {e}")
-            t1 = 0
-        try:
-            t2 = _run_one(engine, n_tokens=512, label="run2-same-prompt")
-        except AssertionError as e:
-            failures.append(f"run2: {e}")
-            t2 = 0
-        if t1 and t2 and t1 != t2:
-            failures.append(
-                f"run1 returned {t1} tokens but run2 returned {t2} (should match)"
-            )
-        try:
-            ta = _run_one(engine, n_tokens=256, label="run3-shorter")
-            tb = _run_one(engine, n_tokens=768, label="run4-longer")
-            if ta and tb and tb <= ta:
-                failures.append(
-                    f"longer prompt should produce more tokens: 256->{ta} 768->{tb}"
-                )
-        except AssertionError as e:
-            failures.append(f"length-variation: {e}")
-    finally:
-        with contextlib.suppress(Exception):
-            engine.close()  # pyright: ignore[reportAttributeAccessIssue]
-
-    if failures:
-        print()
-        print("FAIL")
-        for f in failures:
-            print(f"  - {f}")
-        return 1
-    print()
-    print("PASS")
-    return 0
-
-
-if __name__ == "__main__":
-    try:
-        sys.exit(main(sys.argv))
-    except Exception:
-        traceback.print_exc()
-        sys.exit(1)

setup_linklocal_dgx.sh

@@ -1,124 +0,0 @@
-#!/usr/bin/env bash
-set -Eeuo pipefail
-
-SELF_IP="169.254.100.1"
-PEER_IP="169.254.100.2"
-PREFIX="16"
-IFACE="enP7s7"
-USE_NM="auto"
-DRY_RUN=0
-
-usage() {
-  cat <<EOF
-Usage: sudo $(basename "$0") [options]
-
-Configure a Linux Ethernet interface with a static IPv4 for a host-to-host
-link to a Mac peer.
-
-Defaults: this host = ${SELF_IP}/${PREFIX}, peer = ${PEER_IP}, iface = ${IFACE}.
-
-Options:
-  --iface IFACE     Default: ${IFACE}
-  --self-ip IP      Default: ${SELF_IP}
-  --peer-ip IP      For verification ping. Default: ${PEER_IP}
-  --prefix N        Default: ${PREFIX}
-  --no-nm           Use 'ip addr' directly (transient, no NetworkManager).
-  --dry-run         Print actions without applying.
-  -h, --help        Show this help.
-EOF
-}
-
-while (($#)); do
-  case "$1" in
-  --iface)
-    shift
-    IFACE="${1:?}"
-    ;;
-  --self-ip)
-    shift
-    SELF_IP="${1:?}"
-    ;;
-  --peer-ip)
-    shift
-    PEER_IP="${1:?}"
-    ;;
-  --prefix)
-    shift
-    PREFIX="${1:?}"
-    ;;
-  --no-nm) USE_NM=no ;;
-  --dry-run) DRY_RUN=1 ;;
-  -h | --help)
-    usage
-    exit 0
-    ;;
-  *)
-    echo "Unknown arg: $1" >&2
-    usage >&2
-    exit 1
-    ;;
-  esac
-  shift
-done
-
-[[ $EUID -eq 0 ]] || {
-  echo "Run as root." >&2
-  exit 1
-}
-
-run() {
-  printf '+'
-  printf ' %q' "$@"
-  printf '\n'
-  ((DRY_RUN)) || "$@"
-}
-
-ip link show "$IFACE" >/dev/null 2>&1 || {
-  echo "Interface $IFACE does not exist." >&2
-  exit 1
-}
-
-if [[ $USE_NM == "auto" ]]; then
-  if command -v nmcli >/dev/null 2>&1 && systemctl is-active --quiet NetworkManager 2>/dev/null; then
-    USE_NM=yes
-  else
-    USE_NM=no
-  fi
-fi
-
-if [[ $USE_NM == "yes" ]]; then
-  CONN="$(nmcli -g GENERAL.CONNECTION device show "$IFACE" 2>/dev/null | head -n1 || true)"
-  if [[ -z $CONN || $CONN == "--" ]]; then
-    CONN="static-${IFACE}"
-    run nmcli connection add type ethernet ifname "$IFACE" con-name "$CONN"
-  fi
-  run nmcli connection modify "$CONN" \
-    connection.interface-name "$IFACE" \
-    connection.autoconnect yes \
-    connection.autoconnect-priority 100 \
-    ipv4.method manual \
-    ipv4.addresses "${SELF_IP}/${PREFIX}" \
-    ipv4.gateway "" \
-    ipv4.dns "" \
-    ipv4.never-default yes \
-    ipv6.method link-local \
-    ipv6.addr-gen-mode stable-privacy
-  run nmcli connection up "$CONN"
-else
-  run ip link set "$IFACE" up
-  run ip addr flush dev "$IFACE"
-  run ip addr add "${SELF_IP}/${PREFIX}" dev "$IFACE"
-fi
-
-if ((!DRY_RUN)); then
-  printf '\n'
-  ip -br addr show "$IFACE"
-  printf '\n'
-  if ping -c2 -W2 "$PEER_IP" >/dev/null 2>&1; then
-    echo "OK: $PEER_IP reachable on $IFACE."
-  else
-    echo "WARN: $PEER_IP not reachable yet."
-    echo "      Verify the peer is configured (run setup_linklocal_mac.sh on the Mac)."
-    echo "      ip neigh show dev $IFACE  # check for the peer MAC"
-  fi
-fi

setup_linklocal_mac.sh

@@ -1,170 +0,0 @@
-#!/usr/bin/env bash
-set -Eeuo pipefail
-
-SELF_IP="169.254.100.2"
-PEER_IP="169.254.100.1"
-NETMASK="255.255.0.0"
-IFACE=""
-DRY_RUN=0
-
-usage() {
-  cat <<EOF
-Usage: sudo $(basename "$0") [options]
-
-Configure a Mac Ethernet interface with a static IPv4 for a host-to-host link
-to the DGX/GX10 peer.
-
-Defaults: this Mac = ${SELF_IP}, peer = ${PEER_IP}, mask = ${NETMASK}.
-
-Options:
-  --iface IFACE     Interface (e.g. en12). Default: auto-detect.
-  --self-ip IP      This Mac's address. Default: ${SELF_IP}.
-  --peer-ip IP      Peer for verification ping. Default: ${PEER_IP}.
-  --netmask MASK    Default: ${NETMASK}.
-  --dry-run         Print actions without applying.
-  -h, --help        Show this help.
-EOF
-}
-
-while (($#)); do
-  case "$1" in
-  --iface)
-    shift
-    IFACE="${1:?}"
-    ;;
-  --self-ip)
-    shift
-    SELF_IP="${1:?}"
-    ;;
-  --peer-ip)
-    shift
-    PEER_IP="${1:?}"
-    ;;
-  --netmask)
-    shift
-    NETMASK="${1:?}"
-    ;;
-  --dry-run) DRY_RUN=1 ;;
-  -h | --help)
-    usage
-    exit 0
-    ;;
-  *)
-    echo "Unknown arg: $1" >&2
-    usage >&2
-    exit 1
-    ;;
-  esac
-  shift
-done
-
-[[ $EUID -eq 0 ]] || {
-  echo "Run with sudo." >&2
-  exit 1
-}
-
-run() {
-  printf '+'
-  printf ' %q' "$@"
-  printf '\n'
-  ((DRY_RUN)) || "$@"
-}
-
-target_subnet_prefix() {
-  local ip="$1"
-  printf '%s.' "${ip%.*}"
-}
-
-iface_score() {
-  local iface="$1" info subnet
-  info="$(ifconfig "$iface" 2>/dev/null || true)"
-  [[ -n $info ]] || {
-    echo 0
-    return
-  }
-  grep -q 'status: active' <<<"$info" || {
-    echo 0
-    return
-  }
-  subnet="$(target_subnet_prefix "$SELF_IP")"
-  if grep -qE "inet ${subnet//./\\.}" <<<"$info"; then
-    echo 100
-    return
-  fi
-  if grep -qE 'inet 169\.254\.' <<<"$info"; then
-    echo 80
-    return
-  fi
-  if ! grep -qE '^[[:space:]]*inet ' <<<"$info"; then
-    echo 60
-    return
-  fi
-  echo 10
-}
-
-detect_iface() {
-  local best="" best_score=0 iface score
-  for iface in $(ifconfig -l); do
-    [[ $iface =~ ^en[0-9]+$ ]] || continue
-    score="$(iface_score "$iface")"
-    if ((score > best_score)); then
-      best="$iface"
-      best_score="$score"
-    fi
-  done
-  ((best_score >= 60)) || return 1
-  printf '%s\n' "$best"
-}
-
-iface_to_service() {
-  local iface="$1" line port=""
-  while IFS= read -r line; do
-    if [[ $line == "Hardware Port: "* ]]; then
-      port="${line#Hardware Port: }"
-    elif [[ $line == "Device: $iface" ]]; then
-      printf '%s\n' "$port"
-      return 0
-    fi
-  done < <(networksetup -listallhardwareports)
-  return 1
-}
-
-if [[ -z $IFACE ]]; then
-  IFACE="$(detect_iface || true)"
-  [[ -n $IFACE ]] || {
-    echo "Could not auto-detect a wired interface. Pass --iface enX." >&2
-    echo "Active interfaces:" >&2
-    ifconfig -l | tr ' ' '\n' | grep -E '^en[0-9]+$' | while read -r i; do
-      printf '  %-6s %s\n' "$i" "$(ifconfig "$i" | grep -E 'status:|inet ' | tr '\n' ' ')" >&2
-    done
-    exit 1
-  }
-  echo "Auto-detected interface: $IFACE"
-fi
-
-ifconfig "$IFACE" >/dev/null 2>&1 || {
-  echo "Interface $IFACE does not exist." >&2
-  exit 1
-}
-
-SERVICE="$(iface_to_service "$IFACE" || true)"
-[[ -n $SERVICE ]] || {
-  echo "No network service maps to $IFACE. Check System Settings -> Network." >&2
-  exit 1
-}
-echo "Network service: $SERVICE"
-
-run networksetup -setmanual "$SERVICE" "$SELF_IP" "$NETMASK" ""
-
-if ((!DRY_RUN)); then
-  printf '\n'
-  ifconfig "$IFACE" | grep -E 'inet |status:'
-  printf '\n'
-  if ping -c2 -t3 "$PEER_IP" >/dev/null 2>&1; then
-    echo "OK: $PEER_IP reachable on $IFACE."
-  else
-    echo "WARN: $PEER_IP not reachable yet."
-    echo "      Verify the peer is configured (run setup_linklocal_dgx.sh on the GX10)."
-    echo "      arp -an -i $IFACE  # check for the peer MAC"
-  fi
-fi

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
@@ -526,8 +525,8 @@ class API:
                         )
                     ]
                 )
-        if any(self.state.node_vllm.values()):
-            instance_combinations.append((Sharding.Pipeline, InstanceMeta.Vllm, 1))
+        # TODO: PDD
+        # instance_combinations.append((Sharding.PrefillDecodeDisaggregation, InstanceMeta.MlxRing, 1))
 
         for sharding, instance_meta, min_nodes in instance_combinations:
             try:
@@ -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:
@@ -640,10 +640,7 @@ class API:
         )
 
     async def get_feature_flags(self) -> dict[str, bool]:
-        return {
-            "disaggregation": ENABLE_DISAGGREGATION,
-            "vllm_available": any(self.state.node_vllm.values()),
-        }
+        return {"disaggregation": ENABLE_DISAGGREGATION}
 
     async def list_instance_links(self) -> list[InstanceLink]:
         if not ENABLE_DISAGGREGATION:
@@ -1636,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())
@@ -1659,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
@@ -1723,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()
@@ -1751,7 +1747,6 @@ class API:
                     capabilities=card.capabilities,
                     reasoning_dialect=card.reasoning_dialect,
                     context_length=card.context_length,
-                    requires_vllm=card.requires_vllm,
                 )
                 for card in cards
             ]
@@ -1775,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,
@@ -1791,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/api/types/api.py

@@ -49,7 +49,6 @@ class ModelListModel(BaseModel):
     base_model: str = Field(default="")
     capabilities: list[str] = Field(default_factory=list)
     reasoning_dialect: ReasoningDialect = "none"
-    requires_vllm: bool = Field(default=False)
 
 
 class ModelList(BaseModel):

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

@@ -180,31 +180,10 @@ class Master:
                             for link in self.state.instance_links.values():
                                 prefill_only.difference_update(link.decode_instances)
 
-                            # If the user typed a prefill-only model id (e.g.
-                            # the vLLM-side producer of a P/D pair), the
-                            # candidate decode side is whatever it's linked
-                            # to. Expand the requested model id to also
-                            # include those linked decode instances.
-                            requested_model = command.task_params.model
-                            linked_decode_ids: set[InstanceId] = set()
-                            for link in self.state.instance_links.values():
-                                if any(
-                                    self.state.instances.get(pid) is not None
-                                    and self.state.instances[
-                                        pid
-                                    ].shard_assignments.model_id
-                                    == requested_model
-                                    for pid in link.prefill_instances
-                                ):
-                                    linked_decode_ids.update(link.decode_instances)
-
                             for instance in self.state.instances.values():
-                                model_match = (
-                                    instance.shard_assignments.model_id
-                                    == requested_model
-                                ) or (instance.instance_id in linked_decode_ids)
                                 if (
-                                    model_match
+                                    instance.shard_assignments.model_id
+                                    == command.task_params.model
                                     and instance.instance_id not in prefill_only
                                 ):
                                     in_flight = {TaskStatus.Pending, TaskStatus.Running}
@@ -386,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,
@@ -43,7 +43,6 @@ from exo.shared.types.worker.instances import (
     InstanceMeta,
     MlxJacclInstance,
     MlxRingInstance,
-    VllmInstance,
 )
 from exo.shared.types.worker.shards import Sharding
 from exo.utils.ports import random_ephemeral_port
@@ -106,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))
@@ -167,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:
@@ -203,7 +213,7 @@ def place_instance(
     )
 
     # Single-node: force Pipeline/Ring (Tensor and Jaccl require multi-node)
-    if len(selected_cycle) == 1 and command.instance_meta != InstanceMeta.Vllm:
+    if len(selected_cycle) == 1:
         command = command.model_copy(
             update={
                 "instance_meta": InstanceMeta.MlxRing,
@@ -267,11 +277,6 @@ def place_instance(
                 hosts_by_node=hosts_by_node,
                 ephemeral_port=ephemeral_port,
             )
-        case InstanceMeta.Vllm:
-            target_instances[instance_id] = VllmInstance(
-                instance_id=instance_id,
-                shard_assignments=shard_assignments,
-            )
 
     return target_instances
 

src/exo/master/placement_utils.py

@@ -375,13 +375,7 @@ def find_ip_prioritised(
             "maybe_ethernet": 3,
             "thunderbolt": 4,
         }
-
-    def _key(ip: str) -> tuple[int, int]:
-        link_local = 0 if ip.startswith("169.254.") else 1
-        type_pri = priority.get(ip_to_type.get(ip, "unknown"), 2)
-        return (link_local, type_pri)
-
-    return min(ips, key=_key)
+    return min(ips, key=lambda ip: priority.get(ip_to_type.get(ip, "unknown"), 2))
 
 
 def get_mlx_ring_hosts_by_node(

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,
@@ -59,14 +60,24 @@ from exo.utils.info_gatherer.info_gatherer import (
     NodeConfig,
     NodeDiskUsage,
     NodeNetworkInterfaces,
-    NvmlMetrics,
     RdmaCtlStatus,
     StaticNodeInformation,
     ThunderboltBridgeInfo,
-    VllmCapability,
 )
 
 
+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:
@@ -77,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():
@@ -306,9 +319,6 @@ def apply_node_timed_out(event: NodeTimedOut, state: State) -> State:
     node_rdma_ctl = {
         key: value for key, value in state.node_rdma_ctl.items() if key != event.node_id
     }
-    node_vllm = {
-        key: value for key, value in state.node_vllm.items() if key != event.node_id
-    }
     # Only recompute cycles if the leaving node had TB bridge enabled
     leaving_node_status = state.node_thunderbolt_bridge.get(event.node_id)
     leaving_node_had_tb_enabled = (
@@ -331,7 +341,6 @@ def apply_node_timed_out(event: NodeTimedOut, state: State) -> State:
             "node_thunderbolt": node_thunderbolt,
             "node_thunderbolt_bridge": node_thunderbolt_bridge,
             "node_rdma_ctl": node_rdma_ctl,
-            "node_vllm": node_vllm,
             "thunderbolt_bridge_cycles": thunderbolt_bridge_cycles,
         }
     )
@@ -358,11 +367,6 @@ def apply_node_gathered_info(event: NodeGatheredInfo, state: State) -> State:
                 event.node_id: info.system_profile,
             }
             update["node_memory"] = {**state.node_memory, event.node_id: info.memory}
-        case NvmlMetrics():
-            update["node_system"] = {
-                **state.node_system,
-                event.node_id: info.system_profile,
-            }
         case MemoryUsage():
             update["node_memory"] = {**state.node_memory, event.node_id: info}
         case NodeDiskUsage():
@@ -408,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,
@@ -420,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():
@@ -443,11 +454,12 @@ def apply_node_gathered_info(event: NodeGatheredInfo, state: State) -> State:
                 **state.node_rdma_ctl,
                 event.node_id: NodeRdmaCtlStatus(enabled=info.enabled),
             }
-        case VllmCapability():
-            update["node_vllm"] = {
-                **state.node_vllm,
-                event.node_id: info.available,
-            }
+            # 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)
 
@@ -463,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"
@@ -150,7 +168,6 @@ class ModelCard(FrozenModel):
     context_length: int = 0
     uses_cfg: bool = False
     trust_remote_code: bool = True
-    requires_vllm: bool = False
     is_custom: bool = False
     vision: VisionCardConfig | None = None
     sampling_defaults: SamplingDefaults = Field(default_factory=SamplingDefaults)
@@ -197,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
@@ -234,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
 
@@ -350,11 +351,7 @@ async def fetch_config_data(model_id: ModelId) -> ConfigData:
 
 
 async def fetch_safetensors_size(model_id: ModelId) -> Memory:
-    """Gets model size from safetensors index or falls back to HF API.
-
-    Single-shard repos don't have a `model.safetensors.index.json`; fall back
-    to the HF API for those.
-    """
+    """Gets model size from safetensors index or falls back to HF API."""
     from exo.download.download_utils import (
         download_file_with_retry,
         resolve_model_dir,
@@ -362,25 +359,21 @@ async def fetch_safetensors_size(model_id: ModelId) -> Memory:
     from exo.shared.types.worker.downloads import ModelSafetensorsIndex
 
     target_dir = await resolve_model_dir(model_id)
-    try:
-        index_path = await download_file_with_retry(
-            model_id,
-            "main",
-            "model.safetensors.index.json",
-            target_dir,
-            lambda curr_bytes, total_bytes, is_renamed: logger.debug(
-                f"Downloading model.safetensors.index.json for {model_id}: {curr_bytes}/{total_bytes} ({is_renamed=})"
-            ),
-        )
-    except FileNotFoundError:
-        index_path = None
+    index_path = await download_file_with_retry(
+        model_id,
+        "main",
+        "model.safetensors.index.json",
+        target_dir,
+        lambda curr_bytes, total_bytes, is_renamed: logger.debug(
+            f"Downloading model.safetensors.index.json for {model_id}: {curr_bytes}/{total_bytes} ({is_renamed=})"
+        ),
+    )
+    async with aiofiles.open(index_path, "r") as f:
+        index_data = ModelSafetensorsIndex.model_validate_json(await f.read())
 
-    if index_path is not None:
-        async with aiofiles.open(index_path, "r") as f:
-            index_data = ModelSafetensorsIndex.model_validate_json(await f.read())
-        metadata = index_data.metadata
-        if metadata is not None and metadata.total_size is not None:
-            return Memory.from_bytes(metadata.total_size)
+    metadata = index_data.metadata
+    if metadata is not None and metadata.total_size is not None:
+        return Memory.from_bytes(metadata.total_size)
 
     info = model_info(model_id)
     if info.safetensors is None:

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,
@@ -58,7 +59,6 @@ class State(FrozenModel):
     node_thunderbolt: Mapping[NodeId, NodeThunderboltInfo] = {}
     node_thunderbolt_bridge: Mapping[NodeId, ThunderboltBridgeStatus] = {}
     node_rdma_ctl: Mapping[NodeId, NodeRdmaCtlStatus] = {}
-    node_vllm: Mapping[NodeId, bool] = {}
 
     # Detected cycles where all nodes have Thunderbolt bridge enabled (>2 nodes)
     thunderbolt_bridge_cycles: Sequence[Sequence[NodeId]] = []
@@ -66,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/shared/types/worker/instances.py

@@ -15,7 +15,6 @@ class InstanceId(Id):
 class InstanceMeta(str, Enum):
     MlxRing = "MlxRing"
     MlxJaccl = "MlxJaccl"
-    Vllm = "Vllm"
 
 
 class BaseInstance(TaggedModel):
@@ -36,12 +35,8 @@ class MlxJacclInstance(BaseInstance):
     jaccl_coordinators: dict[NodeId, str]
 
 
-class VllmInstance(BaseInstance):
-    pass
-
-
 # TODO: Single node instance
-Instance = MlxRingInstance | MlxJacclInstance | VllmInstance
+Instance = MlxRingInstance | MlxJacclInstance
 
 
 class BoundInstance(FrozenModel):

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/banner.py

@@ -25,12 +25,12 @@ def print_startup_banner(port: int) -> None:
     banner = f"""
 ╔═══════════════════════════════════════════════════════════════════════╗
 ║                                                                       ║
-║   ███████╗██╗  ██╗ ██████╗               ██████╗  ██████╗ ██╗  ██╗    ║
-║   ██╔════╝╚██╗██╔╝██╔═══██╗      ██      ██╔══██╗██╔════╝ ╚██╗██╔╝    ║
-║   █████╗   ╚███╔╝ ██║   ██║    ██████╗   ██║  ██║██║  ███╗ ╚███╔╝     ║
-║   ██╔══╝   ██╔██╗ ██║   ██║    ╚═██╔═╝   ██║  ██║██║   ██║ ██╔██╗     ║
-║   ███████╗██╔╝ ██╗╚██████╔╝      ╚═╝     ██████╔╝╚██████╔╝██╔╝ ██╗    ║
-║   ╚══════╝╚═╝  ╚═╝ ╚═════╝               ╚═════╝  ╚═════╝ ╚═╝  ╚═╝    ║
+║   ███████╗██╗  ██╗ ██████╗                                            ║
+║   ██╔════╝╚██╗██╔╝██╔═══██╗                                           ║
+║   █████╗   ╚███╔╝ ██║   ██║                                           ║
+║   ██╔══╝   ██╔██╗ ██║   ██║                                           ║
+║   ███████╗██╔╝ ██╗╚██████╔╝                                           ║
+║   ╚══════╝╚═╝  ╚═╝ ╚═════╝                                            ║
 ║                                                                       ║
 ║   Distributed AI Inference Cluster                                    ║
 ║                                                                       ║

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/info_gatherer/info_gatherer.py

@@ -31,7 +31,6 @@ from exo.utils.pydantic_ext import TaggedModel
 from exo.utils.task_group import TaskGroup
 
 from .macmon import MacmonMetrics
-from .nvml import NvmlMetrics, gather_nvidia_metrics, has_nvml
 from .system_info import (
     get_friendly_name,
     get_model_and_chip,
@@ -354,24 +353,6 @@ async def _gather_iface_map() -> dict[str, str] | None:
     return ports
 
 
-class VllmCapability(TaggedModel):
-    available: bool
-    version: str | None = None
-
-    @classmethod
-    async def gather(cls) -> Self:
-        try:
-            import importlib
-
-            vllm = importlib.import_module("vllm")
-            return cls(
-                available=True,
-                version=cast(str | None, getattr(vllm, "__version__", None)),
-            )
-        except ImportError:
-            return cls(available=False)
-
-
 GatheredInfo = (
     MacmonMetrics
     | MemoryUsage
@@ -380,8 +361,6 @@ GatheredInfo = (
     | MacThunderboltConnections
     | RdmaCtlStatus
     | ThunderboltBridgeInfo
-    | NvmlMetrics
-    | VllmCapability
     | NodeConfig
     | MiscData
     | StaticNodeInformation
@@ -440,8 +419,6 @@ class InfoGatherer:
                 tg.start_soon(self._monitor_rdma_ctl_status, 10)
             if not IS_DARWIN:
                 tg.start_soon(self._monitor_memory_usage, 1)
-                if has_nvml():
-                    tg.start_soon(self._monitor_nvml_metrics, 1)
             tg.start_soon(self._watch_system_info, 10)
             tg.start_soon(self._monitor_misc, 60)
             tg.start_soon(self._monitor_static_info, 60)
@@ -450,10 +427,6 @@ class InfoGatherer:
             nc = await NodeConfig.gather()
             if nc is not None:
                 await self.info_sender.send(nc)
-            try:
-                await self.info_sender.send(await VllmCapability.gather())
-            except Exception as e:
-                logger.warning(f"Error gathering vLLM capability: {e}")
 
     def shutdown(self):
         self._tg.cancel_tasks()
@@ -502,16 +475,6 @@ class InfoGatherer:
                 logger.opt(exception=e).warning("Error gathering Thunderbolt data")
             await anyio.sleep(system_profiler_interval)
 
-    async def _monitor_nvml_metrics(self, nvml_poll_rate: float):
-        while True:
-            try:
-                metrics = gather_nvidia_metrics()
-                if metrics is not None:
-                    await self.info_sender.send(metrics)
-            except Exception as e:
-                logger.opt(exception=e).warning("Error gathering NVML metrics")
-            await anyio.sleep(nvml_poll_rate)
-
     async def _monitor_memory_usage(self, memory_poll_rate: float):
         if self._psutil_enabled:
             return

src/exo/utils/info_gatherer/nvml.py

@@ -1,70 +0,0 @@
-from exo.shared.types.profiling import SystemPerformanceProfile
-from exo.utils.pydantic_ext import TaggedModel
-
-try:
-    import pynvml as nvml
-except ImportError:
-    nvml = None
-
-_CPU_POWER_IDLE = 20.0
-_CPU_POWER_MAX = 100.0
-_GPU_POWER_MAX = 120.0
-
-
-class NvmlMetrics(TaggedModel):
-    system_profile: SystemPerformanceProfile
-
-
-def has_nvml() -> bool:
-    if nvml is None:
-        return False
-    try:
-        nvml.nvmlInit()
-        count = nvml.nvmlDeviceGetCount()
-        nvml.nvmlShutdown()
-        return count > 0
-    except Exception:
-        return False
-
-
-def gather_nvidia_metrics() -> NvmlMetrics | None:
-    if nvml is None:
-        return None
-
-    is_init = False
-    try:
-        nvml.nvmlInit()
-        is_init = True
-        count = nvml.nvmlDeviceGetCount()
-        if count == 0:
-            return None
-
-        total_gpu_util = 0.0
-        total_temp = 0.0
-        total_gpu_power = 0.0
-        for i in range(count):
-            handle = nvml.nvmlDeviceGetHandleByIndex(i)
-            util = nvml.nvmlDeviceGetUtilizationRates(handle)
-            total_gpu_util += float(util.gpu)
-            total_temp += float(
-                nvml.nvmlDeviceGetTemperatureV(handle, nvml.NVML_TEMPERATURE_GPU)
-            )
-            total_gpu_power += float(nvml.nvmlDeviceGetPowerUsage(handle)) / 1000.0
-
-        gpu_load_fraction = min(total_gpu_power / _GPU_POWER_MAX, 1.0)
-        estimated_cpu_power = (
-            _CPU_POWER_IDLE + (_CPU_POWER_MAX - _CPU_POWER_IDLE) * gpu_load_fraction
-        )
-
-        return NvmlMetrics(
-            system_profile=SystemPerformanceProfile(
-                gpu_usage=total_gpu_util / count / 100.0,
-                temp=total_temp / count,
-                sys_power=total_gpu_power + estimated_cpu_power,
-            ),
-        )
-    except Exception:
-        return None
-    finally:
-        if is_init:
-            nvml.nvmlShutdown()

src/exo/utils/info_gatherer/system_info.py

@@ -1,7 +1,6 @@
 import platform
 import socket
 import sys
-from pathlib import Path
 from subprocess import CalledProcessError
 
 import psutil
@@ -118,90 +117,12 @@ async def get_network_interfaces() -> list[NetworkInterfaceInfo]:
     return interfaces_info
 
 
-def _read_dmi_field(name: str) -> str | None:
-    try:
-        path = Path(f"/sys/class/dmi/id/{name}")
-        if path.exists():
-            return path.read_text().strip()
-    except (OSError, PermissionError):
-        pass
-    return None
-
-
-async def _get_linux_model_and_chip() -> tuple[str, str]:
-    model = "Linux"
-    chip = "Unknown Chip"
-
-    product_name = _read_dmi_field("product_name")
-    sys_vendor = _read_dmi_field("sys_vendor")
-
-    # DGX Spark: DMI product_name may be "DGX_Spark" or "gx10" variant
-    product_lower = (product_name or "").lower()
-    if product_name and ("dgx" in product_lower or "gx10" in product_lower):
-        model = "DGX Spark"
-        try:
-            process = await run_process(
-                ["nvidia-smi", "--query-gpu=name", "--format=csv,noheader"]
-            )
-            gpu_name = process.stdout.decode().strip().split("\n")[0]
-            chip = gpu_name if gpu_name and gpu_name != "[N/A]" else "NVIDIA GB10"
-        except (CalledProcessError, FileNotFoundError):
-            chip = "NVIDIA GB10"
-        return (model, chip)
-
-    # Other NVIDIA systems (sys_vendor contains "NVIDIA")
-    if sys_vendor and "NVIDIA" in sys_vendor:
-        model = product_name.replace("_", " ") if product_name else "NVIDIA System"
-        try:
-            process = await run_process(
-                ["nvidia-smi", "--query-gpu=name", "--format=csv,noheader"]
-            )
-            gpu_name = process.stdout.decode().strip().split("\n")[0]
-            if gpu_name and gpu_name != "[N/A]":
-                chip = gpu_name
-        except (CalledProcessError, FileNotFoundError):
-            pass
-        return (model, chip)
-
-    # Generic Linux — detect laptop vs desktop via chassis_type
-    # SMBIOS chassis types: 8,9,10,14,31,32 = portable/laptop
-    chassis_type = _read_dmi_field("chassis_type")
-    laptop_chassis_types = {"8", "9", "10", "14", "31", "32"}
-    if chassis_type in laptop_chassis_types:
-        model = "Linux Laptop"
-    elif chassis_type is not None:
-        model = "Linux Desktop"
-
-    # Also check for battery as a fallback laptop indicator
-    if model == "Linux" and Path("/sys/class/power_supply/BAT0").exists():
-        model = "Linux Laptop"
-
-    # Use /proc/cpuinfo for chip
-    cpuinfo_path = Path("/proc/cpuinfo")
-    if cpuinfo_path.exists():
-        try:
-            for line in cpuinfo_path.read_text().splitlines():
-                if line.startswith("model name"):
-                    chip = line.split(":", 1)[1].strip()
-                    break
-        except OSError:
-            pass
-
-    return (model, chip)
-
-
 async def get_model_and_chip() -> tuple[str, str]:
-    """Get system model and chip information.
-
-    On macOS, uses ``system_profiler``.  On Linux, reads DMI data from
-    sysfs and CPU info from ``/proc/cpuinfo``.
-    """
+    """Get Mac system information using system_profiler."""
     model = "Unknown Model"
     chip = "Unknown Chip"
 
-    if sys.platform == "linux":
-        return await _get_linux_model_and_chip()
-
+    # TODO: better non mac support
     if sys.platform != "darwin":
         return (model, chip)
 

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