Optimizations 6

.envrc

@@ -1,8 +1 @@
 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,4 +40,3 @@ bench/**/*.json
 tmp/models
 /build/exo
 /.claude/skills
-/.claude

.mlx_typings/mlx_lm/models/cache.pyi

@@ -191,13 +191,10 @@ 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):  # -> None:
-        ...
-    def is_trimmable(self):  # -> bool:
-        ...
+    def meta_state(self, v: tuple[str, ...]) -> None: ...
+    def is_trimmable(self) -> bool: ...
     def trim(self, n: int) -> int: ...
     def to_quantized(
         self, group_size: int = ..., bits: int = ...

AGENTS.md

@@ -120,81 +120,6 @@ 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,13 +916,11 @@ dependencies = [
  "libp2p",
  "log",
  "networking",
- "pidfile-rs",
  "pin-project",
  "pyo3",
  "pyo3-async-runtimes",
  "pyo3-log",
  "pyo3-stub-gen",
- "thiserror 2.0.17",
  "tokio",
  "util",
 ]
@@ -966,16 +964,6 @@ 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"
@@ -1801,9 +1789,9 @@ checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe"
 
 [[package]]
 name = "libc"
-version = "0.2.186"
+version = "0.2.178"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66"
+checksum = "37c93d8daa9d8a012fd8ab92f088405fb202ea0b6ab73ee2482ae66af4f42091"
 
 [[package]]
 name = "libp2p"
@@ -2819,18 +2807,6 @@ 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"
@@ -3084,18 +3060,6 @@ 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"
@@ -3704,7 +3668,7 @@ dependencies = [
  "netlink-packet-utils",
  "netlink-proto",
  "netlink-sys",
- "nix 0.26.4",
+ "nix",
  "thiserror 1.0.69",
  "tokio",
 ]

README.md

@@ -550,120 +550,6 @@ 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,13 +16,22 @@ 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 = ""
@@ -285,13 +294,6 @@ 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()
@@ -475,6 +477,40 @@ 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
@@ -523,6 +559,127 @@ 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: {
@@ -563,6 +720,61 @@ 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"
@@ -645,6 +857,13 @@ 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,7 +22,6 @@ 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)
@@ -47,7 +46,6 @@ 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()
@@ -68,7 +66,6 @@ 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 = ClusterStateService.makeNonCachingSession()
+        session: URLSession = .shared
     ) {
         self.baseURL = baseURL
         self.endpoint = baseURL.appendingPathComponent("state")
@@ -27,23 +27,6 @@ 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

@@ -1,242 +0,0 @@
-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,6 +21,8 @@ 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 {
@@ -200,6 +202,8 @@ struct SettingsView: View {
                 VStack(alignment: .leading, spacing: 2) {
                     rdmaStatusView
                 }
+
+                sendBugReportButton
             }
 
             Section("Danger Zone") {
@@ -500,8 +504,50 @@ 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

@@ -1,14 +0,0 @@
-"""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

@@ -1,56 +0,0 @@
-"""``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

@@ -1,345 +0,0 @@
-"""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

@@ -1,220 +0,0 @@
-"""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

@@ -1,270 +0,0 @@
-"""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

@@ -1,125 +0,0 @@
-"""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

@@ -1,109 +0,0 @@
-"""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

@@ -1,275 +0,0 @@
-"""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

@@ -1,70 +0,0 @@
-# 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

@@ -1,34 +0,0 @@
-# 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,8 +15,9 @@ from pathlib import Path
 from typing import Any, Literal
 
 import httpx
-from exo_tools.client import ExoClient, ExoHttpError
-from exo_tools.harness import (
+from harness import (
+    ExoClient,
+    ExoHttpError,
     add_common_instance_args,
     capture_cluster_snapshot,
     instance_id_from_instance,

bench/exo_bench.py

@@ -24,12 +24,15 @@ 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 exo_tools.client import ExoClient, ExoHttpError
-from exo_tools.harness import (
+from harness import (
+    ExoClient,
+    ExoHttpError,
     add_common_instance_args,
     capture_cluster_snapshot,
     find_existing_instance,
@@ -43,44 +46,125 @@ from exo_tools.harness import (
     wait_for_instance_ready,
 )
 from loguru import logger
+from transformers import AutoTokenizer
 
-# 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,
-)
+# 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 = _LibPromptSizer
-load_tokenizer_for_bench = _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
 
 
-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 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 format_peak_memory(b: float) -> str:
@@ -186,6 +270,159 @@ 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,8 +42,9 @@ from pathlib import Path
 from typing import Any
 
 import httpx
-from exo_tools.client import ExoClient, ExoHttpError
-from exo_tools.harness import (
+from harness import (
+    ExoClient,
+    ExoHttpError,
     add_common_instance_args,
     capture_cluster_snapshot,
     find_existing_instance,

bench/harness.py

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

bench/lib/__init__.py

@@ -1,18 +0,0 @@
-"""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

@@ -1,215 +0,0 @@
-"""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

@@ -1,194 +0,0 @@
-"""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

@@ -1,428 +0,0 @@
-"""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

@@ -1,183 +0,0 @@
-"""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

@@ -1,308 +0,0 @@
-"""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

@@ -1,269 +0,0 @@
-"""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

@@ -1,139 +0,0 @@
-"""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

@@ -1,65 +0,0 @@
-"""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

@@ -1,186 +0,0 @@
-"""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

@@ -1,171 +0,0 @@
-"""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

@@ -1,189 +0,0 @@
-"""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]
-tg = [512]
+pp = [4096, 8192]
+tg = [128]
 repeat = 1
 warmup = 0
 
 json_out = "bench/prefill_decode_results.json"
 
 [prefill]
-model = "mlx-community/gpt-oss-20b-MXFP4-Q8"
-node = "mike"
-instance_meta = "ring"
+model = "sakamakismile/Qwen3.6-27B-NVFP4"
+node = "gx10-de89"
+instance_meta = "vllm"
 sharding = "pipeline"
 min_nodes = 1
 max_nodes = 1
 
 [decode]
-model = "mlx-community/gpt-oss-20b-MXFP4-Q8"
-node = "james"
+model = "mlx-community/Qwen3.6-27B-4bit"
+node = "Ryuichi’s MacBook Pro"
 instance_meta = "ring"
 sharding = "pipeline"
 min_nodes = 1

bench/prefill_decode_bench.py

@@ -31,12 +31,14 @@ from typing import Any
 
 from exo_bench import (
     PromptSizer,
+    SystemMetricsSampler,
     format_peak_memory,
     load_tokenizer_for_bench,
     parse_int_list,
 )
-from exo_tools.client import ExoClient, ExoHttpError
-from exo_tools.harness import (
+from harness import (
+    ExoClient,
+    ExoHttpError,
     add_common_instance_args,
     instance_id_from_instance,
     node_ids_from_instance,
@@ -277,6 +279,7 @@ 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]] = []
@@ -287,10 +290,13 @@ 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
@@ -314,11 +320,26 @@ 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"
+                f"avg_elapsed={avg_elapsed:.2f}s{energy_str}"
             )
         time.sleep(2)
     return rows
@@ -331,14 +352,36 @@ 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]],
@@ -349,14 +392,17 @@ def _print_diff(
     prefill_alone = _summarise(prefill_alone_rows)
     keys = set(disagg.keys()) | set(decode_alone.keys()) | set(prefill_alone.keys())
 
-    width = 64
+    width = 110
     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':>10}  {'prompt_tps':>11}  {'gen_tps':>9}"
+            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}"
         )
         for label, summary in (
             ("disaggregated", disagg.get(key)),
@@ -364,26 +410,51 @@ def _print_diff(
             ("prefill_alone", prefill_alone.get(key)),
         ):
             if summary is None:
-                logger.info(f"  {label:<16} {'—':>10}  {'—':>11}  {'—':>9}")
+                logger.info(
+                    f"  {label:<16} {'—':>9}  {'—':>9}  "
+                    f"{'—':>11}  {'—':>8}  {'—':>6}  {'—':>6}  "
+                    f"{'—':>9}  {'—':>7}"
+                )
                 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']:>9.2f}s  "
+                f"{summary['elapsed_s']:>8.2f}s  "
+                f"{norm_str}  "
                 f"{summary['prompt_tps']:>11.1f}  "
-                f"{summary['gen_tps']:>9.2f}"
+                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}"
             )
 
         d = disagg.get(key)
         da = decode_alone.get(key)
         pa = prefill_alone.get(key)
-        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"
-            )
+        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)"
+                )
     logger.info("─" * width)
 
 
@@ -681,6 +752,16 @@ 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})
@@ -699,6 +780,7 @@ def main() -> int:
                 warmup=args.warmup,
                 repeat=args.repeat,
                 common_meta=common_meta,
+                sampler=sampler,
             )
             all_rows.extend(prefill_alone_rows)
 
@@ -728,6 +810,7 @@ def main() -> int:
             warmup=args.warmup,
             repeat=args.repeat,
             common_meta=common_meta,
+            sampler=sampler,
         )
         all_rows.extend(disagg_rows)
 
@@ -752,11 +835,13 @@ 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,7 +16,6 @@ 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,6 +202,7 @@
     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" etc. */
+    /** "macbook pro" | "mac studio" | "mac mini" | "dgx spark" | "linux" etc. */
     deviceType: string;
     /** Center X coordinate in SVG space */
     cx: number;
@@ -38,10 +38,43 @@
   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);
@@ -114,7 +147,264 @@
   const studioClipId = $derived(`di-studio-${uid}`);
 </script>
 
-{#if modelLower === "mac studio" || modelLower === "mac mini"}
+{#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"}
   <!-- 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";
+    runtime?: "MlxRing" | "MlxJaccl" | "Vllm";
     onLaunch?: () => void;
     tags?: string[];
     apiPreview?: PlacementPreview | null;
@@ -168,8 +168,10 @@
 
   function getDeviceType(
     name: string,
-  ): "macbook" | "studio" | "mini" | "unknown" {
+  ): "macbook" | "studio" | "mini" | "dgx" | "linux" | "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";
@@ -576,13 +578,17 @@
         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)."
-          : "RDMA: direct memory access over Thunderbolt. Significantly faster for multi-device inference."}
+          : runtime === "MlxJaccl"
+            ? "RDMA: direct memory access over Thunderbolt. Significantly faster for multi-device inference."
+            : "vLLM: NVIDIA CUDA inference engine."}
       >
         {runtime === "MlxRing"
           ? "MLX Ring"
           : runtime === "MlxJaccl"
             ? "MLX RDMA"
-            : runtime}
+            : runtime === "Vllm"
+              ? "vLLM"
+              : runtime}
       </span>
     </div>
 
@@ -990,6 +996,81 @@
                     />
                   {/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,6 +9,7 @@
     capabilities?: string[];
     family?: string;
     is_custom?: boolean;
+    requires_vllm?: boolean;
   }
 
   interface ModelGroup {
@@ -19,6 +20,7 @@
     variants: ModelInfo[];
     smallestVariant: ModelInfo;
     hasMultipleVariants: boolean;
+    requiresVllm: boolean;
   }
 
   type DownloadAvailability = {
@@ -213,6 +215,14 @@
         <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"}
@@ -523,6 +533,15 @@
             {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(
@@ -628,6 +647,7 @@
                 variants: [variant],
                 smallestVariant: variant,
                 hasMultipleVariants: false,
+                requiresVllm: variant.requires_vllm === true,
               });
             }}
             title="View variant details"

dashboard/src/lib/components/ModelPickerModal.svelte

@@ -22,6 +22,7 @@
     is_custom?: boolean;
     tasks?: string[];
     hugging_face_id?: string;
+    requires_vllm?: boolean;
   }
 
   interface ModelGroup {
@@ -32,6 +33,7 @@
     variants: ModelInfo[];
     smallestVariant: ModelInfo;
     hasMultipleVariants: boolean;
+    requiresVllm: boolean;
   }
 
   interface FilterState {
@@ -396,6 +398,7 @@
           variants: [],
           smallestVariant: model,
           hasMultipleVariants: false,
+          requiresVllm: true,
         });
       }
 
@@ -430,6 +433,7 @@
           (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)
@@ -587,6 +591,7 @@
           variants: [model],
           smallestVariant: model,
           hasMultipleVariants: false,
+          requiresVllm: model.requires_vllm === true,
         });
       }
     }
@@ -1165,6 +1170,17 @@
           <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-heavy pass that consumes the entire prompt and
+        Prefill is the compute-bound 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,6 +117,10 @@
   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;
@@ -554,6 +558,13 @@
       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);
@@ -623,7 +634,382 @@
           `${friendlyName}\nID: ${nodeInfo.id.slice(-8)}\nMemory: ${formatBytes(ramUsed)}/${formatBytes(ramTotal)}`,
         );
 
-      if (modelLower === "mac studio") {
+      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") {
         // Mac Studio - classic cube with memory fill
         iconBaseWidth = nodeRadius * 1.25;
         iconBaseHeight = nodeRadius * 0.85;
@@ -1182,8 +1568,12 @@
           debugLabelY += debugLineHeight;
         }
 
-        const identity = identitiesData[nodeInfo.id];
-        if (identity?.osVersion) {
+        const dbgIdentity = identitiesData[nodeInfo.id];
+        if (dbgIdentity?.osVersion) {
+          const osLabel =
+            dbgIdentity.osVersion === "Linux"
+              ? "Linux"
+              : `macOS ${dbgIdentity.osVersion}${dbgIdentity.osBuildVersion ? ` (${dbgIdentity.osBuildVersion})` : ""}`;
           nodeG
             .append("text")
             .attr("x", nodeInfo.x)
@@ -1192,9 +1582,7 @@
             .attr("fill", "rgba(179,179,179,0.7)")
             .attr("font-size", debugFontSize)
             .attr("font-family", "SF Mono, Monaco, monospace")
-            .text(
-              `macOS ${identity.osVersion}${identity.osBuildVersion ? ` (${identity.osBuildVersion})` : ""}`,
-            );
+            .text(osLabel);
         }
       }
     });

dashboard/src/routes/+page.svelte

@@ -65,6 +65,7 @@
     nodeThunderboltBridge,
     nodeIdentities,
     isConnected,
+    featureFlags,
     type DownloadProgress,
     type PlacementPreview,
   } from "$lib/stores/app.svelte";
@@ -702,7 +703,10 @@
       ? Object.keys(topologyData()!.nodes).length
       : 1;
     const sharding = nodeCount <= 1 ? "Pipeline" : selectedSharding;
-    const instanceType = nodeCount <= 1 ? "MlxRing" : selectedInstanceType;
+    const instanceType =
+      nodeCount <= 1 && selectedInstanceType === "MlxJaccl"
+        ? "MlxRing"
+        : selectedInstanceType;
     try {
       const placementResponse = await fetch(
         `/instance/placement?model_id=${encodeURIComponent(modelId)}&sharding=${sharding}&instance_meta=${instanceType}&min_nodes=1`,
@@ -783,6 +787,7 @@
       quantization?: string;
       base_model?: string;
       capabilities?: string[];
+      requires_vllm?: boolean;
     }>
   >([]);
   type ModelMemoryFitStatus =
@@ -886,7 +891,7 @@
   }
 
   let selectedSharding = $state<"Pipeline" | "Tensor">("Pipeline");
-  type InstanceMeta = "MlxRing" | "MlxJaccl";
+  type InstanceMeta = "MlxRing" | "MlxJaccl" | "Vllm";
 
   // Launch defaults persistence
   const LAUNCH_DEFAULTS_KEY = "exo-launch-defaults-v2";
@@ -932,7 +937,12 @@
     // Apply sharding and instance type unconditionally
     selectedSharding = defaults.sharding;
     selectedInstanceType =
-      defaults.instanceType === "MlxRing" ? "MlxRing" : "MlxJaccl";
+      defaults.instanceType === "MlxRing"
+        ? "MlxRing"
+        : defaults.instanceType === "Vllm"
+          ? "Vllm"
+          : "MlxJaccl";
+    userPickedInstanceType = true;
 
     // Apply minNodes if valid (between 1 and maxNodes)
     if (
@@ -954,6 +964,23 @@
   }
 
   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);
@@ -1146,9 +1173,7 @@
   }
 
   const matchesSelectedRuntime = (runtime: InstanceMeta): boolean =>
-    selectedInstanceType === "MlxRing"
-      ? runtime === "MlxRing"
-      : runtime === "MlxJaccl";
+    runtime === selectedInstanceType;
 
   // Helper to check if a model can be launched (has valid placement with >= minNodes)
   function canModelFit(modelId: string): boolean {
@@ -2063,6 +2088,7 @@
     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?: {
@@ -3435,7 +3461,6 @@
             >
               <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"
@@ -4823,7 +4848,6 @@
                 >
                   <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"
@@ -4970,7 +4994,6 @@
                   >
                     <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
@@ -5772,14 +5795,18 @@
                     </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 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'}"
+                        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'}"
                       >
                         <span
                           class="w-3 h-3 rounded-full border-2 flex items-center justify-center {selectedInstanceType ===
@@ -5795,14 +5822,18 @@
                         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 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'}"
+                        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'}"
                       >
                         <span
                           class="w-3 h-3 rounded-full border-2 flex items-center justify-center {selectedInstanceType ===
@@ -5817,7 +5848,41 @@
                         </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.editableVenv
+                self'.packages.exo.passthru.evenv
                 uv
 
                 # RUST

justfile

@@ -40,6 +40,19 @@ 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

@@ -0,0 +1,26 @@
+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,21 +15,18 @@ 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",
-  "mlx==0.31.2; sys_platform == 'darwin'",
-  "mlx-lm; sys_platform=='darwin'",
-  "tiktoken>=0.12.0",                          # required for kimi k2 tokenizer
+  "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]
@@ -40,7 +37,6 @@ 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",
@@ -49,26 +45,30 @@ dev = [
 
 [project.optional-dependencies]
 build = ["nanobind"]
-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-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'",
+
 ]
-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'",
+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'",
 ]
 
 ###
@@ -76,18 +76,29 @@ cuda13 = [
 ###
 
 [tool.uv.workspace]
-members = ["rust/exo_pyo3_bindings", "bench", "tools"]
+members = ["rust/exo_pyo3_bindings", "bench"]
 
 [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-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'" },
+  { 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')" },
 ]
-vllm = { git = "https://github.com/hmellor/vllm.git", branch = "transformers-v5" }
 
 [[tool.uv.index]]
 name = "pytorch-cu130"
@@ -95,8 +106,8 @@ url = "https://download.pytorch.org/whl/cu130"
 explicit = true
 
 [[tool.uv.index]]
-name = "pytorch-cu120"
-url = "https://download.pytorch.org/whl/cu120"
+name = "pytorch-cu128"
+url = "https://download.pytorch.org/whl/cu128"
 explicit = true
 
 [[tool.uv.index]]
@@ -113,7 +124,7 @@ build-backend = "uv_build"
 ###
 
 [tool.basedpyright]
-include = ["src", "bench", "tools"]
+include = ["src", "bench"]
 typeCheckingMode = "strict"
 failOnWarnings = true
 
@@ -147,15 +158,6 @@ 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
@@ -166,11 +168,19 @@ root = "tools/src"
 required-version = ">=0.8.6"
 prerelease = "allow"
 environments = ["sys_platform == 'darwin'", "sys_platform == 'linux'"]
-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'",
+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" },
+  ],
 ]
 
 [tool.uv.extra-build-dependencies]
@@ -185,6 +195,7 @@ mlx = [
   "ninja",
 ]
 mlx-lm = ["setuptools"]
+mflux = ["uv_build"]
 xgrammar = [
   "nanobind",
   "setuptools",
@@ -226,12 +237,9 @@ 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"]
-# `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"
+addopts = "-m 'not slow' --ignore=tests/start_distributed_test.py"
 filterwarnings = ["ignore:builtin type Swig:DeprecationWarning"]

python/parts.nix

@@ -10,8 +10,10 @@ let
       inherit (pkgs.stdenv.hostPlatform) isLinux isDarwin isx86_64;
       inherit (pkgs.config) cudaSupport;
       inherit (pkgs) cudaPackages;
-      cuda13Support = cudaSupport && cudaPackages.cudaMajorVersion == "13";
-      libmlx_source = if cuda13Support then "mlx-cuda-13" else if cudaSupport then "mlx-cuda-12" else "mlx-cpu";
+      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";
       python = pkgs.python313;
       cudaLibs = with cudaPackages; [
         cuda_cudart
@@ -113,37 +115,213 @@ 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: {
-            buildInputs = old.buildInputs ++ [ cudaLibs ];
-            autoPatchelfIgnoreMissingDeps = [ "libcuda.so.1" ];
+            nativeBuildInputs = old.nativeBuildInputs ++ [ pkgs.autoAddDriverRunpath ];
+            buildInputs = old.buildInputs ++ cudaLibs;
           });
           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";
@@ -164,24 +342,28 @@ let
           buildSystemsOverlay
         ]
       );
-      venv = name: (pythonSet.mkVirtualEnv "${name}-env" members).overrideAttrs (_: { venvSkip = [ "lib/python${python.pythonVersion}/site-packages/mlx/share/cmake/*" ]; });
-      mkApp = cmd: name: pkgs.writeShellApplication {
+      # 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 {
         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 ];
-        text = "exec " + lib.optionalString cudaSupport "${lib.getExe pkgs.nix-gl-host} " + cmd;
+        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" ]; });
+        };
       };
     in
     {
       inherit venv;
-      editablePythonSet = pythonSet.overrideScope editableOverlay;
       mkPythonScript = path: mkApp ''python ${path} "$@"'';
       mkExo = mkApp ''exo "$@"'';
     };
@@ -191,18 +373,18 @@ in
     { self', pkgs, unfreePkgs, lib, ... }:
     let
       inherit (pkgs.stdenv.hostPlatform) isLinux;
-      inherit (mkPythonSet { inherit self' pkgs lib; members = { exo = [ "cpu" ]; }; }) editablePythonSet mkExo;
+      inherit (mkPythonSet { inherit self' pkgs lib; members = { exo = [ "mlx-cpu" "vllm-none" ]; }; }) mkExo;
 
       # Virtual environment with dev dependencies for testing
       testVenv = (mkPythonSet {
         inherit self' pkgs lib; members = {
-        exo = [ "dev" "cpu" ]; # Include pytest, pytest-asyncio, pytest-env
+        exo = [ "dev" "mlx-cpu" "vllm-none" ]; # Include pytest, pytest-asyncio, pytest-env
       };
       }).venv "exo-test";
 
       mkBenchScript = (mkPythonSet {
         inherit self' pkgs lib; members = {
-        exo = [ "cpu" ];
+        exo = [ "mlx-cpu" "vllm-none" ];
         exo-bench = [ ]; # Include pytest, pytest-asyncio, pytest-env
       };
       }).mkPythonScript;
@@ -212,12 +394,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);
@@ -226,8 +408,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 = (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";
+        exo-cuda-12 = cuda12Set.mkExo "exo-cuda-12";
+        exo-cuda-13 = cuda13Set.mkExo "exo-cuda-13";
       };
 
       checks = {

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

@@ -0,0 +1,21 @@
+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

@@ -0,0 +1,27 @@
+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

@@ -0,0 +1,20 @@
+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

@@ -0,0 +1,32 @@
+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,12 +46,9 @@ 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,8 +2,6 @@
 # ruff: noqa: E501, F401
 
 import builtins
-import os
-import pathlib
 import typing
 
 @typing.final
@@ -71,48 +69,6 @@ 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,11 +7,9 @@
 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};
@@ -166,7 +164,6 @@ 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

@@ -1,87 +0,0 @@
-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,12 +1,10 @@
 import asyncio
 
 import pytest
-from _pytest.capture import CaptureFixture
 from exo_pyo3_bindings import (
     Keypair,
     NetworkingHandle,
     NoPeersSubscribedToTopicError,
-    Pidfile,
     PyFromSwarm,
 )
 
@@ -28,13 +26,6 @@ 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()
@@ -43,7 +34,3 @@ 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

@@ -0,0 +1,222 @@
+#!/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

@@ -0,0 +1,124 @@
+#!/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

@@ -0,0 +1,170 @@
+#!/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,10 +133,12 @@ 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 (
@@ -479,7 +481,6 @@ 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
@@ -525,8 +526,8 @@ class API:
                         )
                     ]
                 )
-        # TODO: PDD
-        # instance_combinations.append((Sharding.PrefillDecodeDisaggregation, InstanceMeta.MlxRing, 1))
+        if any(self.state.node_vllm.values()):
+            instance_combinations.append((Sharding.Pipeline, InstanceMeta.Vllm, 1))
 
         for sharding, instance_meta, min_nodes in instance_combinations:
             try:
@@ -543,7 +544,6 @@ 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,7 +640,10 @@ class API:
         )
 
     async def get_feature_flags(self) -> dict[str, bool]:
-        return {"disaggregation": ENABLE_DISAGGREGATION}
+        return {
+            "disaggregation": ENABLE_DISAGGREGATION,
+            "vllm_available": any(self.state.node_vllm.values()),
+        }
 
     async def list_instance_links(self) -> list[InstanceLink]:
         if not ENABLE_DISAGGREGATION:
@@ -1633,16 +1636,17 @@ class API:
     async def ollama_tags(self) -> OllamaTagsResponse:
         """Returns list of models in Ollama tags format. We return the downloaded ones only."""
 
-        downloaded_model_ids: set[ModelId] = set()
+        def none_if_empty(value: str) -> str | None:
+            return value or None
+
+        downloaded_model_ids: set[str] = 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 model_cards.card_cache.list_all()
-            if c.model_id in downloaded_model_ids
+            c for c in await get_model_cards() if c.model_id in downloaded_model_ids
         ]
 
         now = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
@@ -1655,8 +1659,8 @@ class API:
                     size=card.storage_size.in_bytes,
                     digest="sha256:000000000000",
                     details=OllamaModelDetails(
-                        family=card.family or None,
-                        quantization_level=card.quantization or None,
+                        family=none_if_empty(card.family),
+                        quantization_level=none_if_empty(card.quantization),
                     ),
                 )
                 for card in cards
@@ -1719,7 +1723,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 model_cards.card_cache.list_all()
+        cards = await get_model_cards()
 
         if status == "downloaded":
             downloaded_model_ids: set[str] = set()
@@ -1747,6 +1751,7 @@ class API:
                     capabilities=card.capabilities,
                     reasoning_dialect=card.reasoning_dialect,
                     context_length=card.context_length,
+                    requires_vllm=card.requires_vllm,
                 )
                 for card in cards
             ]
@@ -1770,7 +1775,7 @@ class API:
 
         # Immediately update the local cache so the subsequent GET /models
         # returns the new model without waiting for the event round-trip.
-        model_cards.card_cache.cc[card.model_id] = card
+        add_to_card_cache(card)
 
         return ModelListModel(
             id=card.model_id,
@@ -1786,7 +1791,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 = model_cards.card_cache.get(model_id)
+        card = get_card(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,6 +49,7 @@ 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,8 +16,7 @@ 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 import model_cards
-from exo.shared.models.model_cards import ModelId
+from exo.shared.models.model_cards import ModelId, get_model_cards
 from exo.shared.types.commands import (
     CancelDownload,
     DeleteDownload,
@@ -423,7 +422,7 @@ class DownloadCoordinator:
                     )
                 # Scan read-only directories for pre-downloaded models
                 if EXO_MODELS_READ_ONLY_DIRS:
-                    for card in await model_cards.card_cache.list_all():
+                    for card in await get_model_cards():
                         mid = card.model_id
                         if mid in self.active_downloads:
                             continue

src/exo/download/download_utils.py

@@ -1,12 +1,11 @@
 import asyncio
 import hashlib
 import os
-import random
 import shutil
 import ssl
 import time
 import traceback
-from collections.abc import Awaitable, Mapping
+from collections.abc import Awaitable
 from datetime import timedelta
 from pathlib import Path
 from typing import Callable, Literal
@@ -56,36 +55,6 @@ 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()
@@ -379,6 +348,9 @@ 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",
@@ -388,16 +360,13 @@ 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}"
 
-    # 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 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())
 
     if skip_internet:
         if await aios.path.exists(cache_file):
@@ -426,6 +395,7 @@ 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(
@@ -456,29 +426,17 @@ async def fetch_file_list_with_retry(
     recursive: bool = False,
     on_connection_lost: Callable[[], None] = lambda: None,
 ) -> list[FileListEntry]:
-    n_attempts = 5
+    n_attempts = 3
     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 + random.uniform(0, 1))
+            await asyncio.sleep(2.0**attempt)
     raise Exception(
         f"Failed to fetch file list for {model_id=} {revision=} {path=} {recursive=}"
     )
@@ -489,9 +447,6 @@ 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 (
@@ -503,8 +458,7 @@ async def _fetch_file_list(
             raise HuggingFaceAuthenticationError(msg)
         elif response.status == 429:
             raise HuggingFaceRateLimitError(
-                f"HuggingFace rate limit hit fetching file list for {model_id}",
-                retry_after=_parse_retry_after(response.headers),
+                f"Couldn't download {model_id} because of HuggingFace rate limit."
             )
         elif response.status == 200:
             data_json = await response.text()
@@ -514,14 +468,10 @@ async def _fetch_file_list(
                 if item.type == "file":
                     files.append(FileListEntry.model_validate(item))
                 elif item.type == "directory" and recursive:
-                    # 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"
-                )
+                    subfiles = await _fetch_file_list(
+                        model_id, revision, item.path, recursive
+                    )
+                    files.extend(subfiles)
             return files
         else:
             raise Exception(f"Failed to fetch file list: {response.status}")
@@ -602,11 +552,6 @@ 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
         )
@@ -626,7 +571,7 @@ async def download_file_with_retry(
     on_connection_lost: Callable[[], None] = lambda: None,
     skip_internet: bool = False,
 ) -> Path:
-    n_attempts = 5
+    n_attempts = 3
     for attempt in range(n_attempts):
         try:
             return await _download_file(
@@ -638,16 +583,12 @@ async def download_file_with_retry(
             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
+                raise e
+            logger.error(
+                f"Download error on attempt {attempt}/{n_attempts} for {model_id=} {revision=} {path=} {target_dir=}"
             )
-            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)
+            logger.error(traceback.format_exc())
+            await asyncio.sleep(2.0**attempt)
         except Exception as e:
             if attempt == n_attempts - 1:
                 on_connection_lost()
@@ -656,7 +597,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 + random.uniform(0, 1))
+            await asyncio.sleep(2.0**attempt)
     raise Exception(
         f"Failed to download file {model_id=} {revision=} {path=} {target_dir=}"
     )
@@ -724,11 +665,6 @@ 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 model_cards.card_cache.list_all()
+            for model_card in await get_model_cards()
         ]
 
         for task in asyncio.as_completed(tasks):

src/exo/download/tests/test_offline_mode.py

@@ -1,7 +1,5 @@
 """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
@@ -233,64 +231,3 @@ 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

@@ -1,355 +0,0 @@
-"""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,7 +3,6 @@ import multiprocessing as mp
 import os
 import resource
 import signal
-import sys
 from dataclasses import dataclass, field
 from typing import Self
 
@@ -23,8 +22,6 @@ 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
@@ -267,26 +264,14 @@ 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}")
@@ -321,7 +306,6 @@ def main():
     finally:
         logger.info("EXO Shutdown complete")
         logger_cleanup()
-        del pidfile
 
 
 class Args(FrozenModel):
@@ -335,7 +319,6 @@ 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
 
@@ -395,11 +378,6 @@ 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,10 +180,31 @@ 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():
-                                if (
+                                model_match = (
                                     instance.shard_assignments.model_id
-                                    == command.task_params.model
+                                    == requested_model
+                                ) or (instance.instance_id in linked_decode_ids)
+                                if (
+                                    model_match
                                     and instance.instance_id not in prefill_only
                                 ):
                                     in_flight = {TaskStatus.Pending, TaskStatus.Running}
@@ -365,7 +386,6 @@ 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, NodeRdmaCtlStatus
+from exo.shared.types.profiling import MemoryUsage, NodeNetworkInfo
 from exo.shared.types.tasks import Task, TaskId, TaskStatus
 from exo.shared.types.worker.downloads import (
     DownloadCompleted,
@@ -43,6 +43,7 @@ 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
@@ -105,7 +106,6 @@ 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,18 +167,8 @@ 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) and _all_rdma_ctl_enabled(cycle)
+        cycle for cycle in smallest_cycles if topology.is_rdma_cycle(cycle)
     ]
 
     if command.instance_meta == InstanceMeta.MlxJaccl:
@@ -213,7 +203,7 @@ def place_instance(
     )
 
     # Single-node: force Pipeline/Ring (Tensor and Jaccl require multi-node)
-    if len(selected_cycle) == 1:
+    if len(selected_cycle) == 1 and command.instance_meta != InstanceMeta.Vllm:
         command = command.model_copy(
             update={
                 "instance_meta": InstanceMeta.MlxRing,
@@ -277,6 +267,11 @@ 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,7 +375,13 @@ def find_ip_prioritised(
             "maybe_ethernet": 3,
             "thunderbolt": 4,
         }
-    return min(ips, key=lambda ip: priority.get(ip_to_type.get(ip, "unknown"), 2))
+
+    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)
 
 
 def get_mlx_ring_hosts_by_node(

src/exo/master/tests/test_placement.py

@@ -21,11 +21,7 @@ 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,
-    NodeRdmaCtlStatus,
-)
+from exo.shared.types.profiling import NetworkInterfaceInfo, NodeNetworkInfo
 from exo.shared.types.tasks import TaskId, TaskStatus, TextGeneration
 from exo.shared.types.text_generation import (
     InputMessage,
@@ -443,21 +439,8 @@ 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,
-        node_rdma_ctl=node_rdma_ctl,
-    )
+    placements = place_instance(cic, topology, {}, node_memory, node_network)
 
     # assert
     assert len(placements) == 1
@@ -499,131 +482,6 @@ 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,8 +4,7 @@ from datetime import datetime
 
 from loguru import logger
 
-from exo.shared.models.model_cards import ModelCard
-from exo.shared.types.common import ModelId, NodeId
+from exo.shared.types.common import NodeId
 from exo.shared.types.events import (
     ChunkGenerated,
     CustomModelCardAdded,
@@ -60,24 +59,14 @@ 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:
@@ -88,12 +77,10 @@ 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():
@@ -319,6 +306,9 @@ 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 = (
@@ -341,6 +331,7 @@ 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,
         }
     )
@@ -367,6 +358,11 @@ 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():
@@ -412,9 +408,6 @@ 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,
@@ -427,10 +420,6 @@ 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():
@@ -454,12 +443,11 @@ def apply_node_gathered_info(event: NodeGatheredInfo, state: State) -> State:
                 **state.node_rdma_ctl,
                 event.node_id: NodeRdmaCtlStatus(enabled=info.enabled),
             }
-            # If RDMA just got disabled on this node, drop any RDMA edges touching it
-            # so placement / topology consumers cannot pick a disabled node for an
-            # RDMA-backed instance. (Edges will repopulate on the next
-            # MacThunderboltConnections poll once both endpoints are enabled again.)
-            if not info.enabled:
-                topology.remove_all_rdma_connections_touching(event.node_id)
+        case VllmCapability():
+            update["node_vllm"] = {
+                **state.node_vllm,
+                event.node_id: info.available,
+            }
 
     return state.model_copy(update=update)
 
@@ -475,22 +463,3 @@ 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,7 +69,6 @@ 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,57 +39,7 @@ _BUILTIN_CARD_DIRS = [
     Path(RESOURCES_DIR) / "image_model_cards",
 ]
 
-
-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()
+_card_cache: dict[ModelId, "ModelCard"] = {}
 
 
 def detect_vision_from_config(model_id: ModelId) -> "VisionCardConfig | None":
@@ -109,10 +59,42 @@ 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"
@@ -168,6 +150,7 @@ 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)
@@ -214,13 +197,14 @@ 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 card_cache.get(model_id) is None:
-            await card_cache.refresh()
-        if (mc := card_cache.get(model_id)) is not None:
+        if model_id not in _card_cache:
+            await _refresh_card_cache()
+        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
@@ -250,6 +234,21 @@ 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
 
@@ -351,7 +350,11 @@ 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."""
+    """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.
+    """
     from exo.download.download_utils import (
         download_file_with_retry,
         resolve_model_dir,
@@ -359,21 +362,25 @@ 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)
-    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())
+    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
 
-    metadata = index_data.metadata
-    if metadata is not None and metadata.total_size is not None:
-        return Memory.from_bytes(metadata.total_size)
+    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)
 
     info = model_info(model_id)
     if info.safetensors is None:

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

@@ -1,44 +0,0 @@
-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

@@ -1,231 +0,0 @@
-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,22 +169,6 @@ 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,9 +5,8 @@ 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 ModelId, NodeId
+from exo.shared.types.common import NodeId
 from exo.shared.types.instance_link import InstanceLink, InstanceLinkId
 from exo.shared.types.profiling import (
     DiskUsage,
@@ -59,6 +58,7 @@ 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,9 +66,6 @@ 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 import model_cards
+        from exo.shared.models.model_cards import get_card
 
-        card = model_cards.card_cache.get(self.model)
+        card = get_card(self.model)
         if card is None:
             return self
 

src/exo/shared/types/worker/instances.py

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

src/exo/utils/async_process.py

@@ -1,290 +0,0 @@
-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

@@ -1,28 +0,0 @@
-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,6 +31,7 @@ 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,
@@ -353,6 +354,24 @@ 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
@@ -361,6 +380,8 @@ GatheredInfo = (
     | MacThunderboltConnections
     | RdmaCtlStatus
     | ThunderboltBridgeInfo
+    | NvmlMetrics
+    | VllmCapability
     | NodeConfig
     | MiscData
     | StaticNodeInformation
@@ -419,6 +440,8 @@ 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)
@@ -427,6 +450,10 @@ 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()
@@ -475,6 +502,16 @@ 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

@@ -0,0 +1,70 @@
+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,6 +1,7 @@
 import platform
 import socket
 import sys
+from pathlib import Path
 from subprocess import CalledProcessError
 
 import psutil
@@ -117,12 +118,90 @@ 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 Mac system information using system_profiler."""
+    """Get system model and chip information.
+
+    On macOS, uses ``system_profiler``.  On Linux, reads DMI data from
+    sysfs and CPU info from ``/proc/cpuinfo``.
+    """
     model = "Unknown Model"
     chip = "Unknown Chip"
 
-    # TODO: better non mac support
+    if sys.platform == "linux":
+        return await _get_linux_model_and_chip()
+
     if sys.platform != "darwin":
         return (model, chip)
 

src/exo/utils/pidfile.py

@@ -1,28 +0,0 @@
-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,21 +19,19 @@ class PowerSampler:
     ):
         self._get_node_system = get_node_system
         self._interval = interval
-        self._samples: defaultdict[
-            NodeId, list[tuple[float, SystemPerformanceProfile]]
-        ] = defaultdict(list)
+        self._samples: defaultdict[NodeId, list[SystemPerformanceProfile]] = (
+            defaultdict(list)
+        )
         self._start_time: float | None = None
         self._stopped = False
 
-    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
+    def _take_sample(self) -> None:
         for node_id, profile in self._get_node_system().items():
-            self._samples[node_id].append((ts, profile))
+            self._samples[node_id].append(profile)
 
     async def run(self) -> None:
         self._start_time = time.perf_counter()
-        self._take_sample(t_rel=0.0)
+        self._take_sample()
         while not self._stopped:
             await anyio.sleep(self._interval)
             self._take_sample()
@@ -41,51 +39,26 @@ 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] = []
-        total_energy_j = 0.0
-        for node_id, ts_profiles in self._samples.items():
-            n = len(ts_profiles)
+        for node_id, profiles in self._samples.items():
+            n = len(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=avg_power_w,
+                    avg_sys_power=sum(p.sys_power for p in profiles) / n,
                 )
             )
 
-        total_avg_sys_w = total_energy_j / elapsed if elapsed > 0 else 0.0
+        total_avg_sys = sum(ns.avg_sys_power for ns in node_stats)
         return PowerUsage(
             elapsed_seconds=elapsed,
             nodes=node_stats,
-            total_avg_sys_power_watts=total_avg_sys_w,
-            total_energy_joules=total_energy_j,
+            total_avg_sys_power_watts=total_avg_sys,
+            total_energy_joules=total_avg_sys * elapsed,
         )
-
-
-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

@@ -1,8 +0,0 @@
-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

@@ -1,515 +0,0 @@
-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

@@ -1,168 +0,0 @@
-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

@@ -1,84 +0,0 @@
-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 == ""