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

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",
 ]

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/Models/ClusterState.swift

@@ -264,7 +264,7 @@ struct NodeDownloadStatus {
     init?(statusKey: String, payload: NodeDownloadPayload) {
         guard let nodeId = payload.nodeId else { return nil }
         self.nodeId = nodeId
-        self.progress = statusKey == "ModelDownloading" ? payload.downloadProgress : nil
+        self.progress = statusKey == "DownloadOngoing" ? payload.downloadProgress : nil
     }
 }
 

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

@@ -30,8 +30,9 @@ 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,

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())}")
@@ -331,8 +421,8 @@ def run_planning_phase(
         node_downloads = client.get_node_downloads(node_id) or []
 
         already_downloaded = any(
-            "ModelReady" in p
-            and unwrap_instance(p["ModelReady"]["shardMetadata"])["modelCard"][
+            "DownloadCompleted" in p
+            and unwrap_instance(p["DownloadCompleted"]["shardMetadata"])["modelCard"][
                 "modelId"
             ]
             == full_model_id
@@ -370,13 +460,14 @@ def run_planning_phase(
 
         completed = [
             (
-                unwrap_instance(p["ModelReady"]["shardMetadata"])["modelCard"][
+                unwrap_instance(p["DownloadCompleted"]["shardMetadata"])["modelCard"][
                     "modelId"
                 ],
-                p["ModelReady"]["total"]["inBytes"],
+                p["DownloadCompleted"]["total"]["inBytes"],
             )
             for p in node_downloads
-            if "ModelReady" in p and not p["ModelReady"].get("readOnly", False)
+            if "DownloadCompleted" in p
+            and not p["DownloadCompleted"].get("readOnly", False)
         ]
         for del_model, size in sorted(completed, key=lambda x: x[1]):
             logger.info(f"Deleting {del_model} from {node_id} ({size // (1024**2)}MB)")
@@ -409,20 +500,20 @@ def run_planning_phase(
         for node_id in node_ids:
             node_downloads = client.get_node_downloads(node_id) or []
             done = any(
-                "ModelReady" in p
-                and unwrap_instance(p["ModelReady"]["shardMetadata"])["modelCard"][
-                    "modelId"
-                ]
+                "DownloadCompleted" in p
+                and unwrap_instance(p["DownloadCompleted"]["shardMetadata"])[
+                    "modelCard"
+                ]["modelId"]
                 == full_model_id
                 for p in node_downloads
             )
             failed = [
-                p["ModelDownloadFailed"]["errorMessage"]
+                p["DownloadFailed"]["errorMessage"]
                 for p in node_downloads
-                if "ModelDownloadFailed" in p
-                and unwrap_instance(p["ModelDownloadFailed"]["shardMetadata"])[
-                    "modelCard"
-                ]["modelId"]
+                if "DownloadFailed" in p
+                and unwrap_instance(p["DownloadFailed"]["shardMetadata"])["modelCard"][
+                    "modelId"
+                ]
                 == full_model_id
             ]
             if failed:
@@ -464,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
@@ -497,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"
@@ -531,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/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)

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/lib/stores/app.svelte.ts

@@ -261,13 +261,6 @@ interface RawStateResponse {
     string,
     { total: { inBytes: number }; available: { inBytes: number } }
   >;
-  nodeStorageConfig?: Record<
-    string,
-    {
-      maxStorage: { inBytes: number } | null;
-      storagePolicy: "manual" | "auto-evict";
-    }
-  >;
 }
 
 export interface MessageAttachment {
@@ -584,15 +577,6 @@ class AppStore {
     >
   >({});
   nodeRdmaCtl = $state<Record<string, { enabled: boolean }>>({});
-  nodeStorageConfig = $state<
-    Record<
-      string,
-      {
-        maxStorage: { inBytes: number } | null;
-        storagePolicy: "manual" | "auto-evict";
-      }
-    >
-  >({});
   nodeThunderboltBridge = $state<
     Record<
       string,
@@ -1367,7 +1351,6 @@ class AppStore {
       this.thunderboltBridgeCycles = data.thunderboltBridgeCycles ?? [];
       // Thunderbolt bridge status per node
       this.nodeThunderboltBridge = data.nodeThunderboltBridge ?? {};
-      this.nodeStorageConfig = data.nodeStorageConfig ?? {};
       this.lastUpdate = Date.now();
       // Connection recovered
       if (!this.isConnected) {
@@ -3426,29 +3409,6 @@ class AppStore {
     }
   }
 
-  async setStorageConfig(
-    nodeIds: string[] | null,
-    maxStorageGb: number | null,
-    storagePolicy: "manual" | "auto-evict",
-  ): Promise<void> {
-    try {
-      const response = await fetch("/storage", {
-        method: "PUT",
-        headers: { "Content-Type": "application/json" },
-        body: JSON.stringify({ nodeIds, maxStorageGb, storagePolicy }),
-      });
-      if (!response.ok) {
-        const errorText = await response.text();
-        throw new Error(
-          `Failed to set storage config: ${response.status} - ${errorText}`,
-        );
-      }
-    } catch (error) {
-      console.error("Error setting storage config:", error);
-      throw error;
-    }
-  }
-
   /**
    * List all available traces
    */
@@ -3539,7 +3499,6 @@ export const deleteInstanceLink = (linkId: string) =>
   appStore.deleteInstanceLink(linkId);
 export const downloads = () => appStore.downloads;
 export const nodeDisk = () => appStore.nodeDisk;
-export const nodeStorageConfig = () => appStore.nodeStorageConfig;
 export const placementPreviews = () => appStore.placementPreviews;
 export const selectedPreviewModelId = () => appStore.selectedPreviewModelId;
 export const isLoadingPreviews = () => appStore.isLoadingPreviews;
@@ -3665,11 +3624,6 @@ export const cancelDownload = (nodeId: string, modelId: string) =>
   appStore.cancelDownload(nodeId, modelId);
 export const deleteDownload = (nodeId: string, modelId: string) =>
   appStore.deleteDownload(nodeId, modelId);
-export const setStorageConfig = (
-  nodeIds: string[] | null,
-  maxStorageGb: number | null,
-  storagePolicy: "manual" | "auto-evict",
-) => appStore.setStorageConfig(nodeIds, maxStorageGb, storagePolicy);
 
 // Trace actions
 export const listTraces = () => appStore.listTraces();

dashboard/src/lib/utils/downloads.ts

@@ -5,7 +5,7 @@
  *   Record<NodeId, Array<TaggedDownloadEntry>>
  *
  * Each entry is a tagged union object like:
- *   { "ModelReady": { shard_metadata: { "PipelineShardMetadata": { model_card: { model_id: "..." }, ... } }, ... } }
+ *   { "DownloadCompleted": { shard_metadata: { "PipelineShardMetadata": { model_card: { model_id: "..." }, ... } }, ... } }
  */
 
 /** Unwrap one level of tagged-union envelope, returning [tag, payload]. */
@@ -49,7 +49,7 @@ export function extractShardMetadata(
   return shardMetadata as Record<string, unknown>;
 }
 
-/** Get the download tag (ModelReady, ModelDownloading, etc.) from a wrapped entry. */
+/** Get the download tag (DownloadCompleted, DownloadOngoing, etc.) from a wrapped entry. */
 export function getDownloadTag(
   entry: unknown,
 ): [string, Record<string, unknown>] | null {
@@ -73,7 +73,7 @@ function* iterNodeDownloads(
   }
 }
 
-/** Check if a specific model is fully downloaded (ModelReady) on a specific node. */
+/** Check if a specific model is fully downloaded (DownloadCompleted) on a specific node. */
 export function isModelDownloadedOnNode(
   downloadsData: Record<string, unknown[]>,
   nodeId: string,
@@ -83,12 +83,12 @@ export function isModelDownloadedOnNode(
   if (!Array.isArray(nodeDownloads)) return false;
 
   for (const [tag, , entryModelId] of iterNodeDownloads(nodeDownloads)) {
-    if (tag === "ModelReady" && entryModelId === modelId) return true;
+    if (tag === "DownloadCompleted" && entryModelId === modelId) return true;
   }
   return false;
 }
 
-/** Get all node IDs where a model is fully downloaded (ModelReady). */
+/** Get all node IDs where a model is fully downloaded (DownloadCompleted). */
 export function getNodesWithModelDownloaded(
   downloadsData: Record<string, unknown[]>,
   modelId: string,
@@ -122,7 +122,7 @@ export function getShardMetadataForModel(
       const shard = extractShardMetadata(payload);
       if (!shard) continue;
 
-      if (tag === "ModelReady") return shard;
+      if (tag === "DownloadCompleted") return shard;
       if (!fallback) fallback = shard;
     }
   }
@@ -131,7 +131,7 @@ export function getShardMetadataForModel(
 
 /**
  * Get the download status tag for a specific model on a specific node.
- * Returns the "best" status: ModelReady > ModelDownloading > others.
+ * Returns the "best" status: DownloadCompleted > DownloadOngoing > others.
  */
 export function getModelDownloadStatus(
   downloadsData: Record<string, unknown[]>,
@@ -144,8 +144,8 @@ export function getModelDownloadStatus(
   let best: string | null = null;
   for (const [tag, , entryModelId] of iterNodeDownloads(nodeDownloads)) {
     if (entryModelId !== modelId) continue;
-    if (tag === "ModelReady") return tag;
-    if (tag === "ModelDownloading") best = tag;
+    if (tag === "DownloadCompleted") return tag;
+    if (tag === "DownloadOngoing") best = tag;
     else if (!best) best = tag;
   }
   return best;

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 {
@@ -1581,14 +1606,12 @@
     progress: DownloadProgress | null;
     perNode: NodeDownloadStatus[];
     failedError: string | null;
-    rejectedError: string | null;
   } {
     const empty = {
       isDownloading: false,
       progress: null,
       perNode: [] as NodeDownloadStatus[],
       failedError: null,
-      rejectedError: null,
     };
 
     if (!downloadsData || Object.keys(downloadsData).length === 0) {
@@ -1620,8 +1643,8 @@
         const downloadModelId = extractModelIdFromDownload(downloadPayload);
         if (!downloadModelId || downloadModelId !== modelId) continue;
 
-        // ModelDownloadFailed — return with any data collected so far
-        if (downloadKind === "ModelDownloadFailed") {
+        // DownloadFailed — return with any data collected so far
+        if (downloadKind === "DownloadFailed") {
           return {
             isDownloading: false,
             progress: null,
@@ -1630,33 +1653,20 @@
               (downloadPayload.errorMessage as string) ||
               (downloadPayload.error_message as string) ||
               "Download failed",
-            rejectedError: null,
-          };
-        }
-
-        // ModelRejected — storage limit exceeded
-        if (downloadKind === "ModelRejected") {
-          return {
-            isDownloading: false,
-            progress: null,
-            perNode: Array.from(perNodeMap.values()),
-            failedError: null,
-            rejectedError:
-              (downloadPayload.reason as string) || "Storage limit exceeded",
           };
         }
 
         if (
-          downloadKind !== "ModelDownloading" &&
-          downloadKind !== "ModelNotDownloading" &&
-          downloadKind !== "ModelReady"
+          downloadKind !== "DownloadOngoing" &&
+          downloadKind !== "DownloadPending" &&
+          downloadKind !== "DownloadCompleted"
         )
           continue;
 
         const nodeName =
           data?.nodes?.[nodeId]?.friendly_name ?? nodeId.slice(0, 8);
 
-        if (downloadKind === "ModelReady") {
+        if (downloadKind === "DownloadCompleted") {
           perNodeMap.set(nodeId, {
             nodeId,
             nodeName,
@@ -1667,7 +1677,7 @@
           continue;
         }
 
-        if (downloadKind === "ModelNotDownloading") {
+        if (downloadKind === "DownloadPending") {
           const pendingDownloaded = getBytes(
             downloadPayload.downloaded ??
               downloadPayload.downloaded_bytes ??
@@ -1691,7 +1701,7 @@
           continue;
         }
 
-        // ModelDownloading
+        // DownloadOngoing
         const progress = parseDownloadProgress(downloadPayload);
         if (
           !progress ||
@@ -1737,7 +1747,6 @@
         progress: null,
         perNode,
         failedError: null,
-        rejectedError: null,
       };
     }
 
@@ -1758,7 +1767,6 @@
       },
       perNode,
       failedError: null,
-      rejectedError: null,
     };
   }
 
@@ -1843,17 +1851,6 @@
       };
     }
 
-    if (result.rejectedError) {
-      return {
-        isDownloading: false,
-        isFailed: true,
-        errorMessage: result.rejectedError,
-        progress: null,
-        statusText: "REJECTED",
-        perNode: [],
-      };
-    }
-
     if (!result.isDownloading) {
       const statusInfo = deriveInstanceStatus(instanceWrapped);
       return {
@@ -2091,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?: {
@@ -2503,13 +2501,7 @@
     if (Object.keys(prev).length > 0) {
       for (const [id, currentStatus] of Object.entries(currentStatuses)) {
         const prevStatus = prev[id];
-        if (prevStatus === currentStatus) continue;
-        if (
-          !prevStatus &&
-          currentStatus !== "REJECTED" &&
-          currentStatus !== "FAILED"
-        )
-          continue;
+        if (!prevStatus || prevStatus === currentStatus) continue;
 
         const modelId = getInstanceModelId(instanceData[id]);
         const shortName = modelId
@@ -2543,14 +2535,6 @@
           addToast({ type: "error", message: `Model failed: ${shortName}` });
         }
 
-        if (prevStatus !== "REJECTED" && currentStatus === "REJECTED") {
-          addToast({
-            type: "warning",
-            message: `Storage limit exceeded: ${shortName}`,
-            duration: 8000,
-          });
-        }
-
         // Any -> Shutdown
         if (prevStatus !== "SHUTDOWN" && currentStatus === "SHUTDOWN") {
           addToast({ type: "info", message: `Model shut down: ${shortName}` });
@@ -2561,46 +2545,6 @@
     previousInstanceStatuses = currentStatuses;
   });
 
-  // ── Download rejection toasts (independent of instances) ──
-  // Instances are deleted immediately after rejection, so the instance-based
-  // toast logic above never sees them. Watch downloads directly instead.
-  let previousRejectedModels = new Set<string>();
-
-  $effect(() => {
-    const currentRejected = new Set<string>();
-    if (downloadsData && typeof downloadsData === "object") {
-      for (const nodeDownloads of Object.values(downloadsData)) {
-        if (!Array.isArray(nodeDownloads)) continue;
-        for (const entry of nodeDownloads) {
-          if (!entry || typeof entry !== "object") continue;
-          const keys = Object.keys(entry as Record<string, unknown>);
-          if (keys.length !== 1) continue;
-          if (keys[0] === "ModelRejected") {
-            const payload = (entry as Record<string, unknown>)[
-              keys[0]
-            ] as Record<string, unknown>;
-            const modelId = extractModelIdFromDownload(payload);
-            if (modelId) currentRejected.add(modelId);
-          }
-        }
-      }
-    }
-
-    if (previousRejectedModels.size > 0 || currentRejected.size > 0) {
-      for (const modelId of currentRejected) {
-        if (!previousRejectedModels.has(modelId)) {
-          const shortName = modelId.split("/").pop() ?? modelId;
-          addToast({
-            type: "warning",
-            message: `Storage limit exceeded: ${shortName}`,
-            duration: 8000,
-          });
-        }
-      }
-    }
-    previousRejectedModels = currentRejected;
-  });
-
   // ── Connection status toasts ──
   let previousConnectionStatus: boolean | null = null;
 
@@ -3517,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"
@@ -4905,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"
@@ -5052,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
@@ -5854,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 ===
@@ -5877,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 ===
@@ -5899,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 -->

dashboard/src/routes/downloads/+page.svelte

@@ -6,20 +6,17 @@
     topologyData,
     downloads,
     nodeDisk,
-    nodeStorageConfig,
     refreshState,
     lastUpdate as lastUpdateStore,
     startDownload,
     cancelDownload,
     deleteDownload,
-    setStorageConfig,
   } from "$lib/stores/app.svelte";
   import {
     getDownloadTag,
     extractModelIdFromDownload,
     extractShardMetadata,
   } from "$lib/utils/downloads";
-  import { addToast } from "$lib/stores/toast.svelte";
   import HeaderNav from "$lib/components/HeaderNav.svelte";
 
   type CellStatus =
@@ -39,15 +36,7 @@
         total: number;
         modelDirectory?: string;
       }
-    | { kind: "failed"; errorMessage?: string; modelDirectory?: string }
-    | {
-        kind: "rejected";
-        reason: string;
-        requiredBytes: number;
-        availableBytes: number;
-        limitBytes?: number;
-        modelDirectory?: string;
-      }
+    | { kind: "failed"; modelDirectory?: string }
     | { kind: "not_present" };
 
   type ModelCardInfo = {
@@ -73,14 +62,11 @@
     label: string;
     diskAvailable?: number;
     diskTotal?: number;
-    storageLimit?: number;
-    storagePolicy?: "manual" | "auto-evict";
   };
 
   const data = $derived(topologyData());
   const downloadsData = $derived(downloads());
   const nodeDiskData = $derived(nodeDisk());
-  const storageConfigData = $derived(nodeStorageConfig());
 
   function getNodeLabel(nodeId: string): string {
     const node = data?.nodes?.[nodeId];
@@ -137,37 +123,10 @@
     return Math.min(100, Math.max(0, value as number));
   }
 
-  function getNodeUsedStorage(nodeId: string): number {
-    const nodeDownloads = downloadsData?.[nodeId];
-    if (!nodeDownloads || !Array.isArray(nodeDownloads)) return 0;
-    let total = 0;
-    for (const entry of nodeDownloads) {
-      const tagged = getDownloadTag(entry);
-      if (!tagged) continue;
-      const [tag, payload] = tagged;
-      if (tag === "ModelReady") {
-        total += getBytes(payload.total);
-      } else if (tag === "ModelDownloading") {
-        const prog = (payload.download_progress ?? payload.downloadProgress) as
-          | Record<string, unknown>
-          | undefined;
-        if (prog) total += getBytes(prog.downloaded);
-      }
-    }
-    return total;
-  }
-
-  function storageBarColor(percent: number): string {
-    if (percent >= 90) return "bg-red-500";
-    if (percent >= 70) return "bg-yellow-500";
-    return "bg-green-500";
-  }
-
   const CELL_PRIORITY: Record<CellStatus["kind"], number> = {
-    completed: 5,
-    downloading: 4,
-    pending: 3,
-    rejected: 2,
+    completed: 4,
+    downloading: 3,
+    pending: 2,
     failed: 1,
     not_present: 0,
   };
@@ -220,80 +179,6 @@
   let nodeColumns = $state<NodeColumn[]>([]);
   let infoRow = $state<ModelRow | null>(null);
 
-  let storageConfigNode = $state<NodeColumn | null>(null);
-  let configMaxGb = $state<number | null>(null);
-  let configNoLimit = $state(true);
-  let configPolicy = $state<"manual" | "auto-evict">("manual");
-  let configSaving = $state(false);
-  let configApplyAll = $state(false);
-  let configDiskTotalGb = $derived(
-    storageConfigNode
-      ? Math.round((storageConfigNode.diskTotal ?? 0) / 1024 ** 3)
-      : 0,
-  );
-  let configEffectiveCapacityGb = $derived.by(() => {
-    if (!storageConfigNode) return 0;
-    const diskAvail = storageConfigNode.diskAvailable ?? 0;
-    const exoUsed = getNodeUsedStorage(storageConfigNode.nodeId);
-    return Math.round((diskAvail + exoUsed) / 1024 ** 3);
-  });
-  let configLimitExceedsDisk = $derived(
-    !configNoLimit &&
-      configMaxGb != null &&
-      configMaxGb > configEffectiveCapacityGb &&
-      configEffectiveCapacityGb > 0,
-  );
-
-  function openStorageConfig(col: NodeColumn) {
-    storageConfigNode = col;
-    if (col.storageLimit != null) {
-      configNoLimit = false;
-      configMaxGb = Math.round(col.storageLimit / 1024 ** 3);
-    } else {
-      configNoLimit = true;
-      configMaxGb = null;
-    }
-    configPolicy = col.storagePolicy ?? "manual";
-    configApplyAll = false;
-  }
-
-  async function freeSpaceAndRetry(
-    nodeId: string,
-    shardMetadata: Record<string, unknown>,
-  ) {
-    try {
-      const col = nodeColumns.find((c) => c.nodeId === nodeId);
-      const limitGb = col?.storageLimit ? col.storageLimit / 1024 ** 3 : null;
-      await setStorageConfig([nodeId], limitGb, "auto-evict");
-      await startDownload(nodeId, shardMetadata);
-      refreshState();
-    } catch (error) {
-      addToast({
-        type: "error",
-        message: `Failed: ${error instanceof Error ? error.message : String(error)}`,
-      });
-    }
-  }
-
-  async function saveStorageConfig() {
-    if (!storageConfigNode) return;
-    configSaving = true;
-    try {
-      const maxGb = configNoLimit ? null : configMaxGb;
-      const nodeIds = configApplyAll ? null : [storageConfigNode.nodeId];
-      await setStorageConfig(nodeIds, maxGb, configPolicy);
-      storageConfigNode = null;
-      refreshState();
-    } catch (error) {
-      addToast({
-        type: "error",
-        message: `Failed to save storage config: ${error instanceof Error ? error.message : String(error)}`,
-      });
-    } finally {
-      configSaving = false;
-    }
-  }
-
   $effect(() => {
     try {
       if (!downloadsData || Object.keys(downloadsData).length === 0) {
@@ -305,14 +190,11 @@
       const allNodeIds = Object.keys(downloadsData);
       const columns: NodeColumn[] = allNodeIds.map((nodeId) => {
         const diskInfo = nodeDiskData?.[nodeId];
-        const storageConfig = storageConfigData?.[nodeId];
         return {
           nodeId,
           label: getNodeLabel(nodeId),
           diskAvailable: diskInfo?.available?.inBytes,
           diskTotal: diskInfo?.total?.inBytes,
-          storageLimit: storageConfig?.maxStorage?.inBytes ?? undefined,
-          storagePolicy: storageConfig?.storagePolicy,
         };
       });
 
@@ -353,14 +235,10 @@
             ((payload.model_directory ?? payload.modelDirectory) as string) ||
             undefined;
           let cell: CellStatus;
-          if (tag === "ModelReady") {
+          if (tag === "DownloadCompleted") {
             const totalBytes = getBytes(payload.total);
-            cell = {
-              kind: "completed",
-              totalBytes,
-              modelDirectory,
-            };
-          } else if (tag === "ModelDownloading") {
+            cell = { kind: "completed", totalBytes, modelDirectory };
+          } else if (tag === "DownloadOngoing") {
             const rawProgress =
               payload.download_progress ?? payload.downloadProgress ?? {};
             const prog = rawProgress as Record<string, unknown>;
@@ -380,21 +258,8 @@
               etaMs,
               modelDirectory,
             };
-          } else if (tag === "ModelRejected") {
-            cell = {
-              kind: "rejected",
-              reason: (payload.reason as string) ?? "Storage limit exceeded",
-              requiredBytes: getBytes(payload.required),
-              availableBytes: getBytes(payload.available),
-              limitBytes: getBytes(payload.limit),
-              modelDirectory,
-            };
-          } else if (tag === "ModelDownloadFailed") {
-            const errorMessage =
-              (payload.error_message as string) ??
-              (payload.errorMessage as string) ??
-              undefined;
-            cell = { kind: "failed", errorMessage, modelDirectory };
+          } else if (tag === "DownloadFailed") {
+            cell = { kind: "failed", modelDirectory };
           } else {
             const downloaded = getBytes(
               payload.downloaded ??
@@ -420,13 +285,12 @@
       }
 
       function rowSortKey(row: ModelRow): number {
-        // in progress (4) -> completed (3) -> rejected/paused (2) -> not started (1) -> not present (0)
+        // in progress (4) -> completed (3) -> paused (2) -> not started (1) -> not present (0)
         let best = 0;
         for (const cell of Object.values(row.cells)) {
           let score = 0;
           if (cell.kind === "downloading") score = 4;
           else if (cell.kind === "completed") score = 3;
-          else if (cell.kind === "rejected") score = 2;
           else if (cell.kind === "pending" && cell.downloaded > 0)
             score = 2; // paused
           else if (cell.kind === "pending" || cell.kind === "failed") score = 1; // not started
@@ -596,57 +460,15 @@
                 Model
               </th>
               {#each nodeColumns as col}
-                {@const usedStorage = getNodeUsedStorage(col.nodeId)}
-                {@const quotaLimit = col.storageLimit}
-                {@const diskAvail = col.diskAvailable ?? 0}
-                {@const storageMax =
-                  quotaLimit != null
-                    ? Math.min(quotaLimit, diskAvail + usedStorage)
-                    : diskAvail + usedStorage}
-                {@const storagePercent =
-                  storageMax > 0
-                    ? Math.min(100, (usedStorage / storageMax) * 100)
-                    : 0}
                 <th
                   class="px-4 py-3 text-[11px] uppercase tracking-wider text-exo-light-gray font-medium text-center whitespace-nowrap min-w-[120px]"
                 >
-                  <div class="flex items-center justify-center gap-1">
-                    <span>{col.label}</span>
-                    <button
-                      type="button"
-                      class="p-0.5 rounded hover:bg-white/10 transition-colors"
-                      onclick={() => openStorageConfig(col)}
-                      title="Storage settings"
-                      aria-label="Storage settings for {col.label}"
+                  <div>{col.label}</div>
+                  {#if col.diskAvailable != null}
+                    <div
+                      class="text-[9px] text-white/70 normal-case tracking-normal mt-0.5"
                     >
-                      <svg
-                        class="w-3.5 h-3.5 text-white/40 hover:text-exo-yellow transition-colors"
-                        viewBox="0 0 20 20"
-                        fill="currentColor"
-                      >
-                        <path
-                          fill-rule="evenodd"
-                          d="M11.49 3.17c-.38-1.56-2.6-1.56-2.98 0a1.532 1.532 0 01-2.286.948c-1.372-.836-2.942.734-2.106 2.106.54.886.061 2.042-.947 2.287-1.561.379-1.561 2.6 0 2.978a1.532 1.532 0 01.947 2.287c-.836 1.372.734 2.942 2.106 2.106a1.532 1.532 0 012.287.947c.379 1.561 2.6 1.561 2.978 0a1.533 1.533 0 012.287-.947c1.372.836 2.942-.734 2.106-2.106a1.533 1.533 0 01.947-2.287c1.561-.379 1.561-2.6 0-2.978a1.532 1.532 0 01-.947-2.287c.836-1.372-.734-2.942-2.106-2.106a1.532 1.532 0 01-2.287-.947zM10 13a3 3 0 100-6 3 3 0 000 6z"
-                          clip-rule="evenodd"
-                        />
-                      </svg>
-                    </button>
-                  </div>
-                  {#if storageMax > 0}
-                    <div class="text-[9px] normal-case tracking-normal mt-1">
-                      <div
-                        class="w-full h-1.5 bg-white/10 rounded-full overflow-hidden"
-                      >
-                        <div
-                          class="h-full rounded-full transition-all duration-300 {storageBarColor(
-                            storagePercent,
-                          )}"
-                          style="width: {storagePercent.toFixed(1)}%"
-                        ></div>
-                      </div>
-                      <div class="text-white/60 mt-0.5">
-                        {formatBytes(usedStorage)} / {formatBytes(storageMax)}
-                      </div>
+                      {formatBytes(col.diskAvailable)} free
                     </div>
                   {/if}
                 </th>
@@ -814,85 +636,36 @@
                           <span class="text-white/40 text-sm">...</span>
                         {/if}
                       </div>
-                    {:else if cell.kind === "rejected"}
+                    {:else if cell.kind === "failed"}
                       <div
                         class="flex flex-col items-center gap-1"
-                        title={cell.reason}
+                        title="Download failed"
                       >
                         <svg
-                          class="w-7 h-7 text-orange-400"
+                          class="w-7 h-7 text-red-400"
                           viewBox="0 0 20 20"
                           fill="currentColor"
                         >
                           <path
                             fill-rule="evenodd"
-                            d="M8.257 3.099c.765-1.36 2.722-1.36 3.486 0l5.58 9.92c.75 1.334-.213 2.98-1.742 2.98H4.42c-1.53 0-2.493-1.646-1.743-2.98l5.58-9.92zM11 13a1 1 0 11-2 0 1 1 0 012 0zm-1-8a1 1 0 00-1 1v3a1 1 0 002 0V6a1 1 0 00-1-1z"
+                            d="M4.293 4.293a1 1 0 011.414 0L10 8.586l4.293-4.293a1 1 0 111.414 1.414L11.414 10l4.293 4.293a1 1 0 01-1.414 1.414L10 11.414l-4.293 4.293a1 1 0 01-1.414-1.414L8.586 10 4.293 5.707a1 1 0 010-1.414z"
                             clip-rule="evenodd"
                           ></path>
                         </svg>
-                        <span class="text-[10px] text-orange-400/80"
-                          >Need {formatBytes(cell.requiredBytes)}</span
-                        >
-                        <span class="text-[10px] text-white/50"
-                          >{formatBytes(cell.availableBytes)} avail</span
-                        >
-                        {#if row.shardMetadata}
-                          <div class="flex items-center gap-2 mt-0.5">
-                            <button
-                              type="button"
-                              class="text-[9px] text-white/50 hover:text-orange-300 transition-colors cursor-pointer border border-white/10 hover:border-orange-400/40 rounded px-1.5 py-0.5"
-                              onclick={() =>
-                                freeSpaceAndRetry(
-                                  col.nodeId,
-                                  row.shardMetadata!,
-                                )}
-                              title="Switch to auto-evict, remove least-recently-used models, and retry download"
-                            >
-                              Free space & retry
-                            </button>
+                        <div class="flex gap-1">
+                          {#if row.shardMetadata}
                             <button
                               type="button"
                               class="text-white/50 hover:text-exo-yellow transition-colors cursor-pointer"
                               onclick={() =>
                                 startDownload(col.nodeId, row.shardMetadata!)}
-                              title="Retry download (without freeing space)"
+                              title="Retry download on this node"
                             >
                               {@render downloadIcon()}
                             </button>
-                          </div>
-                        {/if}
-                      </div>
-                    {:else if cell.kind === "failed"}
-                      <div class="flex flex-col items-center gap-1">
-                        <!-- Error icon with tooltip -->
-                        <div class="relative group">
-                          <svg
-                            class="w-7 h-7 text-red-400 cursor-help"
-                            viewBox="0 0 20 20"
-                            fill="currentColor"
-                          >
-                            <path
-                              fill-rule="evenodd"
-                              d="M18 10a8 8 0 11-16 0 8 8 0 0116 0zm-7 4a1 1 0 11-2 0 1 1 0 012 0zm-1-9a1 1 0 00-1 1v4a1 1 0 102 0V6a1 1 0 00-1-1z"
-                              clip-rule="evenodd"
-                            ></path>
-                          </svg>
-                          <div
-                            class="absolute top-full left-1/2 -translate-x-1/2 mt-2 px-3 py-2 bg-black/95 border border-red-500/30 rounded-lg text-[10px] text-red-300 opacity-0 group-hover:opacity-100 transition-opacity pointer-events-none z-50 max-w-[300px] text-center"
-                          >
-                            {cell.errorMessage ?? "Download failed"}
-                          </div>
+                          {/if}
+                          {@render deleteButton(col.nodeId, row.modelId)}
                         </div>
-                        {#if row.shardMetadata}
-                          <button
-                            type="button"
-                            class="text-[9px] text-white/50 hover:text-exo-yellow transition-colors cursor-pointer border border-white/10 hover:border-exo-yellow/40 rounded px-1.5 py-0.5"
-                            onclick={() =>
-                              startDownload(col.nodeId, row.shardMetadata!)}
-                          >
-                            Retry
-                          </button>
-                        {/if}
                       </div>
                     {:else}
                       <div
@@ -1044,11 +817,9 @@
                       ? 'bg-green-500/10 text-green-400/80 border border-green-500/20'
                       : cellStatus.kind === 'downloading'
                         ? 'bg-exo-yellow/10 text-exo-yellow/80 border border-exo-yellow/20'
-                        : cellStatus.kind === 'rejected'
-                          ? 'bg-orange-500/10 text-orange-400/80 border border-orange-500/20'
-                          : cellStatus.kind === 'failed'
-                            ? 'bg-red-500/10 text-red-400/80 border border-red-500/20'
-                            : 'bg-white/5 text-white/50 border border-white/10'}"
+                        : cellStatus.kind === 'failed'
+                          ? 'bg-red-500/10 text-red-400/80 border border-red-500/20'
+                          : 'bg-white/5 text-white/50 border border-white/10'}"
                   >
                     {col.label}
                     {#if cellStatus.kind === "downloading" && "percentage" in cellStatus}
@@ -1073,221 +844,8 @@
   </div>
 {/if}
 
-<!-- Storage config modal -->
-{#if storageConfigNode}
-  <div
-    class="fixed inset-0 z-[60] bg-black/60"
-    transition:fade={{ duration: 150 }}
-    onclick={() => (storageConfigNode = null)}
-    role="presentation"
-  ></div>
-  <div
-    class="fixed z-[60] top-1/2 left-1/2 -translate-x-1/2 -translate-y-1/2 w-[min(80vw,360px)] bg-exo-dark-gray border border-exo-yellow/10 rounded-lg shadow-2xl p-4"
-    transition:fly={{ y: 10, duration: 200, easing: cubicOut }}
-    role="dialog"
-    aria-modal="true"
-    onkeydown={(e) => {
-      if (e.key === "Escape") storageConfigNode = null;
-    }}
-  >
-    <div class="flex items-start justify-between mb-4">
-      <h3 class="font-mono text-sm text-white">
-        Storage — {configApplyAll ? "All nodes" : storageConfigNode.label}
-      </h3>
-      <button
-        type="button"
-        class="p-1 rounded hover:bg-white/10 transition-colors text-white/50"
-        onclick={() => (storageConfigNode = null)}
-        aria-label="Close storage settings"
-      >
-        <svg class="w-4 h-4" viewBox="0 0 24 24" fill="currentColor">
-          <path
-            d="M19 6.41L17.59 5 12 10.59 6.41 5 5 6.41 10.59 12 5 17.59 6.41 19 12 13.41 17.59 19 19 17.59 13.41 12 19 6.41z"
-          />
-        </svg>
-      </button>
-    </div>
-
-    <div class="space-y-4">
-      <!-- Apply to all nodes -->
-      {#if nodeColumns.length > 1}
-        <label class="flex items-center gap-2 cursor-pointer">
-          <input
-            type="checkbox"
-            bind:checked={configApplyAll}
-            class="accent-exo-yellow w-4 h-4"
-          />
-          <span class="text-xs font-mono text-white/80">Apply to all nodes</span
-          >
-        </label>
-      {/if}
-
-      <!-- No limit checkbox -->
-      <label class="flex items-center gap-2 cursor-pointer">
-        <input
-          type="checkbox"
-          bind:checked={configNoLimit}
-          onchange={() => {
-            if (!configNoLimit && configMaxGb == null) {
-              configMaxGb = configDiskTotalGb || 50;
-            }
-          }}
-          class="accent-exo-yellow w-4 h-4"
-        />
-        <span class="text-xs font-mono text-white/80">Unlimited storage</span>
-      </label>
-
-      <!-- Max storage slider -->
-      <div class="space-y-1.5">
-        <div class="flex items-baseline justify-between">
-          <label
-            class="text-[11px] font-mono text-white/50 uppercase tracking-wider"
-            for="storage-max-gb"
-          >
-            Max storage
-          </label>
-          <span
-            class="text-xs font-mono tabular-nums transition-opacity {configNoLimit
-              ? 'opacity-30'
-              : 'text-white'}"
-          >
-            {configMaxGb ?? 0} GB
-          </span>
-        </div>
-        <input
-          id="storage-max-gb"
-          type="range"
-          min="1"
-          max={Math.max(configDiskTotalGb, configMaxGb ?? 1)}
-          step="1"
-          bind:value={configMaxGb}
-          disabled={configNoLimit}
-          class="slider w-full h-1.5 rounded-full appearance-none cursor-pointer
-            disabled:opacity-30 disabled:cursor-not-allowed"
-        />
-        <div
-          class="flex justify-between text-[10px] font-mono text-white/30 transition-opacity {configNoLimit
-            ? 'opacity-30'
-            : ''}"
-        >
-          <span>1 GB</span>
-          <span>{Math.max(configDiskTotalGb, configMaxGb ?? 1)} GB</span>
-        </div>
-      </div>
-
-      <!-- Disk capacity warning -->
-      {#if configLimitExceedsDisk}
-        <div
-          class="flex items-start gap-2 px-3 py-2 rounded bg-orange-500/10 border border-orange-500/20"
-        >
-          <svg
-            class="w-4 h-4 text-orange-400 shrink-0 mt-0.5"
-            viewBox="0 0 20 20"
-            fill="currentColor"
-          >
-            <path
-              fill-rule="evenodd"
-              d="M8.257 3.099c.765-1.36 2.722-1.36 3.486 0l5.58 9.92c.75 1.334-.213 2.98-1.742 2.98H4.42c-1.53 0-2.493-1.646-1.743-2.98l5.58-9.92zM11 13a1 1 0 11-2 0 1 1 0 012 0zm-1-8a1 1 0 00-1 1v3a1 1 0 002 0V6a1 1 0 00-1-1z"
-              clip-rule="evenodd"
-            />
-          </svg>
-          <p class="text-[10px] text-orange-300/80 font-mono">
-            Disk only has {configEffectiveCapacityGb} GB available for models. The
-            {configMaxGb} GB limit has no effect.
-          </p>
-        </div>
-      {/if}
-
-      <!-- Policy selector -->
-      <div class="space-y-1.5">
-        <div
-          class="text-[11px] font-mono text-white/50 uppercase tracking-wider"
-        >
-          Eviction policy
-        </div>
-        <div class="flex gap-1">
-          <button
-            type="button"
-            class="flex-1 px-3 py-1.5 rounded text-xs font-mono transition-colors
-              {configPolicy === 'manual'
-              ? 'bg-exo-yellow/20 text-exo-yellow border border-exo-yellow/40'
-              : 'bg-exo-black/40 text-white/50 border border-exo-medium-gray/30 hover:text-white/70'}"
-            onclick={() => (configPolicy = "manual")}
-          >
-            Manual
-          </button>
-          <button
-            type="button"
-            class="flex-1 px-3 py-1.5 rounded text-xs font-mono transition-colors
-              {configPolicy === 'auto-evict'
-              ? 'bg-exo-yellow/20 text-exo-yellow border border-exo-yellow/40'
-              : 'bg-exo-black/40 text-white/50 border border-exo-medium-gray/30 hover:text-white/70'}"
-            onclick={() => (configPolicy = "auto-evict")}
-          >
-            Auto-evict
-          </button>
-        </div>
-        <p class="text-[10px] text-white/40 font-mono">
-          {#if configPolicy === "manual"}
-            Downloads that exceed the limit are rejected. Delete models
-            manually.
-          {:else}
-            Oldest unused models are automatically removed to make room.
-          {/if}
-        </p>
-      </div>
-    </div>
-
-    <!-- Actions -->
-    <div class="flex justify-end gap-2 mt-5">
-      <button
-        type="button"
-        class="px-3 py-1.5 rounded text-xs font-mono text-white/50 hover:text-white/70 transition-colors"
-        onclick={() => (storageConfigNode = null)}
-      >
-        Cancel
-      </button>
-      <button
-        type="button"
-        class="px-3 py-1.5 rounded text-xs font-mono bg-exo-yellow/20 text-exo-yellow border border-exo-yellow/40 hover:bg-exo-yellow/30 transition-colors disabled:opacity-50"
-        onclick={saveStorageConfig}
-        disabled={configSaving ||
-          (!configNoLimit && (configMaxGb == null || configMaxGb <= 0))}
-      >
-        {configSaving ? "Saving..." : "Save"}
-      </button>
-    </div>
-  </div>
-{/if}
-
 <style>
   table {
     min-width: max-content;
   }
-
-  .slider {
-    background: rgba(255, 255, 255, 0.1);
-  }
-  .slider::-webkit-slider-thumb {
-    -webkit-appearance: none;
-    width: 14px;
-    height: 14px;
-    border-radius: 50%;
-    background: #f5c518;
-    cursor: pointer;
-  }
-  .slider::-moz-range-thumb {
-    width: 14px;
-    height: 14px;
-    border-radius: 50%;
-    border: none;
-    background: #f5c518;
-    cursor: pointer;
-  }
-  .slider:disabled::-webkit-slider-thumb {
-    cursor: not-allowed;
-  }
-  .slider:disabled::-moz-range-thumb {
-    cursor: not-allowed;
-  }
 </style>

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

@@ -26,6 +26,7 @@ dependencies = [
   "msgspec>=0.19.0",
   "zstandard>=0.23.0",
   "transformers>=5.6.2",
+  "nvidia-ml-py>=13.595.45",
 ]
 
 [project.scripts]
@@ -36,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",
@@ -47,27 +47,28 @@ dev = [
 build = ["nanobind"]
 mlx-none = ["anyio"]
 mlx = [
-  "mlx==0.32.0",
+  "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'",
+
 ]
 mlx-cpu = ["exo[mlx]", "mlx-cpu==0.31.2; sys_platform == 'linux'"]
-mlx-cuda12 = [
-  "exo[mlx]",
-  "mlx-cuda-12==0.32.0; sys_platform == 'linux'",
-  "nvidia-ml-py>=13.595.45",
-]
-mlx-cuda13 = [
-  "exo[mlx]",
-  "mlx-cuda-13==0.32.0; sys_platform == 'linux'",
-  "nvidia-ml-py>=13.595.45",
+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'",
 ]
 
 ###
@@ -75,40 +76,28 @@ mlx-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'" },
-  { url = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv/releases/download/mlx_cuda/mlx-0.32.0-cp313-cp313-manylinux_2_35_aarch64.whl", marker = "sys_platform == 'linux' and platform_machine == 'aarch64'" },
-  { url = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv/releases/download/mlx_cuda/mlx-0.32.0-cp313-cp313-manylinux_2_35_x86_64.whl", marker = "sys_platform == 'linux' and platform_machine != 'aarch64'" },
-
-]
+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 = "https://github.com/evanev7/mflux", branch = "exo2" }
+mflux = { git = "http://github.com/evanev7/mflux", branch = "exo" }
+vllm = { git = "http://github.com/evanev7/vllm", branch = "exo2" }
 torch = [
-  { index = "pytorch-cpu", marker = "sys_platform == 'linux' and extra == 'mlx-cpu' and extra != 'mlx-cuda13' and extra != 'mlx-cuda12'" },
-  { index = "pytorch-cu128", marker = "sys_platform == 'linux' and extra == 'mlx-cuda12' and extra != 'mlx-cuda13' " },
-  { index = "pytorch-cu130", marker = "sys_platform == 'linux' and extra == 'mlx-cuda13'" },
-]
-mlx-cuda-12 = [
-  { url = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv/releases/download/mlx_cuda/mlx_cuda_12-0.32.0-py3-none-manylinux_2_35_aarch64.whl", marker = "sys_platform == 'linux' and platform_machine == 'aarch64'" },
-  { url = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv/releases/download/mlx_cuda/mlx_cuda_12-0.32.0-py3-none-manylinux_2_35_x86_64.whl", marker = "sys_platform == 'linux' and platform_machine != 'aarch64'" },
-]
-mlx-cuda-13 = [
-  { url = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv/releases/download/mlx_cuda/mlx_cuda_13-0.32.0-py3-none-manylinux_2_35_aarch64.whl", marker = "sys_platform == 'linux' and platform_machine == 'aarch64'" },
-  { url = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv/releases/download/mlx_cuda/mlx_cuda_13-0.32.0-py3-none-manylinux_2_35_x86_64.whl", marker = "sys_platform == 'linux' and platform_machine != 'aarch64'" },
+  { 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 != 'mlx-cuda13' and extra != 'mlx-cuda12'" },
-  { index = "pytorch-cu128", marker = "sys_platform == 'linux' and extra == 'mlx-cuda12' and extra != 'mlx-cuda13'" },
-  { index = "pytorch-cu130", marker = "sys_platform == 'linux' and extra == 'mlx-cuda13'" },
+  { 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 != 'mlx-cuda13' and extra != 'mlx-cuda12'" },
-  { index = "pytorch-cu128", marker = "sys_platform == 'linux' and extra == 'mlx-cuda12' and extra != 'mlx-cuda13'" },
-  { index = "pytorch-cu130", marker = "sys_platform == 'linux' and extra == 'mlx-cuda13'" },
+  { 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')" },
 ]
 
 [[tool.uv.index]]
@@ -135,7 +124,7 @@ build-backend = "uv_build"
 ###
 
 [tool.basedpyright]
-include = ["src", "bench", "tools"]
+include = ["src", "bench"]
 typeCheckingMode = "strict"
 failOnWarnings = true
 
@@ -169,13 +158,6 @@ reportMissingModuleSource = false
 [[tool.basedpyright.executionEnvironments]]
 root = "src"
 
-[[tool.basedpyright.executionEnvironments]]
-root = "bench"
-extraPaths = ["tools/src"]
-
-[[tool.basedpyright.executionEnvironments]]
-root = "tools/src"
-
 
 ###
 # uv configuration
@@ -194,6 +176,11 @@ conflicts = [
     { extra = "mlx-cpu" },
     { extra = "mlx-none" },
   ],
+  [
+    { extra = "vllm-cuda13" },
+    { extra = "mlx-cuda12" },
+    { extra = "vllm-none" },
+  ],
 ]
 
 [tool.uv.extra-build-dependencies]
@@ -254,5 +241,5 @@ pythonpath = "."
 asyncio_mode = "auto"
 markers = ["slow: marks tests as slow (deselected by default)"]
 env = ["EXO_TESTS=1"]
-addopts = "-m 'not slow' --ignore=tests"
+addopts = "-m 'not slow' --ignore=tests/start_distributed_test.py"
 filterwarnings = ["ignore:builtin type Swig:DeprecationWarning"]

python/parts.nix

@@ -15,13 +15,7 @@ let
         else if (builtins.elem "mlx-cuda12" members.exo or [ ]) then "mlx-cuda-12"
         else "mlx-cpu";
       python = pkgs.python313;
-
-      cuda_cccl_compat = pkgs.runCommand "cuda-cccl-compat" { } ''
-        mkdir -p $out/include
-        ln -s ${cudaPackages.cuda_cccl}/include $out/include/cccl
-      '';
       cudaLibs = with cudaPackages; [
-        cuda_crt
         cuda_cudart
         cuda_cccl
         cuda_cupti
@@ -39,10 +33,6 @@ let
         libnvshmem
         nccl
       ];
-      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 ]);
-      };
       exoOverlay = final: prev: {
         # Replace workspace exo_pyo3_bindings with Nix-built wheel.
         # Preserve passthru so mkVirtualEnv can resolve dependency groups.
@@ -174,6 +164,159 @@ let
             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: {
@@ -201,41 +344,28 @@ let
       );
       # 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 =
-        let
-          libPath = lib.makeLibraryPath (
-            [ pkgs.stdenv.cc.cc.lib ] ++ lib.optionals cudaSupport [ cudaRoot ]
-          );
-        in
-        text: name: pkgs.writeShellApplication {
-          inherit name;
-          text = ''
-            LD_LIBRARY_PATH="${libPath}''${LD_LIBRARY_PATH:+:}''${LD_LIBRARY_PATH:-}" exec \
-              ${lib.optionalString cudaSupport "nixglhost "} ${text}
-          '';
-          runtimeEnv = {
-            EXO_DASHBOARD_DIR = self'.packages.dashboard;
-            EXO_RESOURCES_DIR = inputs.self + /resources;
-          };
-          runtimeInputs = [
-            (venv name)
-          ] ++ lib.optionals cudaSupport [ pkgs.nix-gl-host ]
-          ++ lib.optionals isDarwin [ pkgs.macmon ];
-          passthru = {
-            venv = venv name;
-            evenv = ((pythonSet.overrideScope editableOverlay).mkVirtualEnv "${name}-evenv" (members // { exo = (members.exo or [ ]) ++ [ "dev" ]; })).overrideAttrs (_: {
-              venvSkip = [ "lib/python${python.pythonVersion}/site-packages/mlx/share/cmake/*" "lib/python${python.pythonVersion}/site-packages/build_backend.py" ];
-            });
-          } // lib.optionalAttrs cudaSupport {
-            inherit cudaRoot;
-          };
+      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 = [
+          (venv name)
+          pkgs.nix-gl-host
+        ]
+        ++ lib.optionals isDarwin [ pkgs.macmon ];
+        passthru = {
+          venv = venv name;
+          evenv = ((pythonSet.overrideScope editableOverlay).mkVirtualEnv "${name}-evenv" (members // { exo = (members.exo or [ ]) ++ [ "dev" ]; })).overrideAttrs (_: { venvSkip = [ "lib/python${python.pythonVersion}/site-packages/mlx/share/cmake/*" "lib/python${python.pythonVersion}/site-packages/build_backend.py" ]; });
+        };
+      };
     in
     {
       inherit venv;
       mkPythonScript = path: mkApp ''python ${path} "$@"'';
-      exo = mkApp ''exo "$@"'' "exo";
+      mkExo = mkApp ''exo "$@"'';
     };
 in
 {
@@ -243,18 +373,18 @@ in
     { self', pkgs, unfreePkgs, lib, ... }:
     let
       inherit (pkgs.stdenv.hostPlatform) isLinux;
-      inherit (mkPythonSet { inherit self' pkgs lib; members = { exo = [ "mlx-cpu" ]; }; }) exo;
+      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" "mlx-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 = [ "mlx-cpu" ];
+        exo = [ "mlx-cpu" "vllm-none" ];
         exo-bench = [ ]; # Include pytest, pytest-asyncio, pytest-env
       };
       }).mkPythonScript;
@@ -264,14 +394,12 @@ in
         runtimeInputs = [ pkgs.python313 ];
         text = ''exec python ${path} "$@"'';
       };
-      # if someone is particularly interested in cuda12 support in nix, please open an issue.
-      # until then, it's more hassle than its worth
-      #cuda12Set = mkPythonSet { inherit self' lib; inherit (unfreePkgs.pkgsCuda.cudaPackages_12) pkgs; members = { exo = [ "mlx-cuda12" ]; }; };
-      cuda13Set = mkPythonSet { inherit self' lib; inherit (unfreePkgs.pkgsCuda.cudaPackages_13) pkgs; members = { exo = [ "mlx-cuda13" ]; }; };
+      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 = {
-        inherit exo;
+        exo = mkExo "exo";
         # for running tests in ci
         exo-test-env = testVenv;
         exo-bench = mkBenchScript "exo-bench" (inputs.self + /bench/exo_bench.py);
@@ -280,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 = cuda12Set.exo;
-        exo-cuda-13 = cuda13Set.exo;
+        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/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "maturin"
 
 [project]
 name = "exo_pyo3_bindings"
-version = "0.2.2"
+version = "0.2.1"
 description = "Add your description here"
 readme = "README.md"
 authors = [

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

@@ -20,9 +20,7 @@ from fastapi.staticfiles import StaticFiles
 from hypercorn.asyncio import serve  # pyright: ignore[reportUnknownVariableType]
 from hypercorn.config import Config
 from hypercorn.typing import ASGIFramework
-from hypercorn.utils import LifespanTimeoutError
 from loguru import logger
-from pydantic import Field
 
 from exo.api.adapters.chat_completions import (
     chat_request_to_text_generation,
@@ -135,12 +133,13 @@ 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.storage import calculate_used_storage
 from exo.shared.tracing import TraceEvent, compute_stats, export_trace, load_trace_file
 from exo.shared.types.chunks import (
     ErrorChunk,
@@ -167,7 +166,6 @@ from exo.shared.types.commands import (
     PlaceInstance,
     SendInputChunk,
     SetInstanceLink,
-    SetStorageConfig,
     StartDownload,
     TaskCancelled,
     TaskFinished,
@@ -184,7 +182,6 @@ from exo.shared.types.events import (
 from exo.shared.types.instance_link import InstanceLink, InstanceLinkId
 from exo.shared.types.memory import Memory
 from exo.shared.types.state import State
-from exo.shared.types.storage import StorageConfig, StoragePolicy
 from exo.shared.types.tasks import (
     ImageEdits as ImageEditsTask,
 )
@@ -198,28 +195,15 @@ from exo.shared.types.text_generation import (
     Base64ImageHash,
     TextGenerationTaskParams,
 )
-from exo.shared.types.worker.downloads import ModelReady
+from exo.shared.types.worker.downloads import DownloadCompleted
 from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
 from exo.shared.types.worker.shards import Sharding
 from exo.utils.banner import print_startup_banner
 from exo.utils.channels import Receiver, Sender, channel
 from exo.utils.disk_event_log import DiskEventLog
 from exo.utils.power_sampler import PowerSampler
-from exo.utils.pydantic_ext import FrozenModel
 from exo.utils.task_group import TaskGroup
 
-
-class SetStorageConfigRequest(FrozenModel):
-    node_ids: list[NodeId] | None = None
-    max_storage_gb: Annotated[float, Field(ge=0)] | None = None
-    storage_policy: StoragePolicy = "manual"
-
-
-class NodeStorageInfo(FrozenModel):
-    config: StorageConfig
-    used: Memory
-
-
 _API_EVENT_LOG_DIR = EXO_EVENT_LOG_DIR / "api"
 ONBOARDING_COMPLETE_FILE = EXO_CACHE_HOME / "onboarding_complete"
 
@@ -409,9 +393,6 @@ class API:
         self.app.post("/download/start")(self.start_download)
         self.app.delete("/download/{node_id}/{model_id:path}")(self.delete_download)
         self.app.post("/download/cancel")(self.cancel_download)
-        self.app.get("/storage")(self.get_storage)
-        self.app.get("/storage/{node_id}")(self.get_storage_node)
-        self.app.put("/storage")(self.set_storage_config)
         self.app.get("/v1/traces")(self.list_traces)
         self.app.post("/v1/traces/delete")(self.delete_traces)
         self.app.get("/v1/traces/{task_id}")(self.get_trace)
@@ -500,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
@@ -546,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:
@@ -564,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:
@@ -661,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:
@@ -1654,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, ModelReady):
+                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())
@@ -1676,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
@@ -1740,13 +1723,13 @@ 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()
             for node_downloads in self.state.downloads.values():
                 for dl in node_downloads:
-                    if isinstance(dl, ModelReady):
+                    if isinstance(dl, DownloadCompleted):
                         downloaded_model_ids.add(dl.shard_metadata.model_card.model_id)
             cards = [c for c in cards if c.model_id in downloaded_model_ids]
 
@@ -1768,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
             ]
@@ -1791,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,
@@ -1807,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")
 
@@ -1877,21 +1861,12 @@ class API:
                     await anyio.sleep_forever()
                 finally:
                     with anyio.CancelScope(shield=True):
-                        # IMPORTANT: when new queues are added, update this (for proper shutdown semantics)
-                        self._shutdown_queues(self._text_generation_queues)
-                        self._shutdown_queues(self._image_generation_queues)
-
                         shutdown_ev.set()
         finally:
             self._event_log.close()
             self.command_sender.close()
             self.event_receiver.close()
 
-    @staticmethod
-    def _shutdown_queues[K, V](queues: dict[K, Sender[V]]):
-        for v in queues.values():
-            v.close()
-
     async def run_api(self, ev: anyio.Event):
         cfg = Config()
         cfg.bind = [f"0.0.0.0:{self.port}"]
@@ -1899,23 +1874,12 @@ class API:
         cfg.accesslog = None
         cfg.errorlog = "-"
         cfg.logger_class = InterceptLogger
-
-        # prevents hangs when mid-request and connection refuses to close
-        cfg.graceful_timeout = 2  # seconds
-        cfg.shutdown_timeout = 3  # seconds
-
         with anyio.CancelScope(shield=True):
-            try:
-                await serve(
-                    cast(ASGIFramework, self.app),
-                    cfg,
-                    shutdown_trigger=ev.wait,
-                )
-            except LifespanTimeoutError as e:
-                logger.warning(
-                    "Graceful server shutdown timed out, some connections forcebly closed"
-                )
-                logger.opt(exception=e).debug("")
+            await serve(
+                cast(ASGIFramework, self.app),
+                cfg,
+                shutdown_trigger=ev.wait,
+            )
 
     async def _apply_state(self):
         with self.event_receiver as events:
@@ -2033,48 +1997,6 @@ class API:
         await self._send_download(command)
         return CancelDownloadResponse(command_id=command.command_id)
 
-    async def get_storage(self) -> dict[str, NodeStorageInfo]:
-        result: dict[str, NodeStorageInfo] = {}
-        for node_id, config in self.state.node_storage_config.items():
-            downloads = list(self.state.downloads.get(node_id, ()))
-            used = calculate_used_storage(downloads)
-            result[node_id] = NodeStorageInfo(config=config, used=used)
-        return result
-
-    async def get_storage_node(self, node_id: NodeId) -> NodeStorageInfo:
-        config = self.state.node_storage_config.get(node_id, StorageConfig())
-        downloads = list(self.state.downloads.get(node_id, ()))
-        used = calculate_used_storage(downloads)
-        return NodeStorageInfo(config=config, used=used)
-
-    async def set_storage_config(
-        self, request: SetStorageConfigRequest
-    ) -> dict[str, str | list[str]]:
-        max_storage = (
-            Memory.from_gb(request.max_storage_gb)
-            if request.max_storage_gb is not None
-            else None
-        )
-
-        target_node_ids = (
-            request.node_ids
-            if request.node_ids is not None
-            else list(self.state.node_storage_config.keys())
-        )
-
-        command_ids: list[str] = []
-        for node_id in target_node_ids:
-            command = SetStorageConfig(
-                target_node_id=node_id,
-                max_storage=max_storage,
-                storage_policy=request.storage_policy,
-            )
-            await self.command_sender.send(
-                ForwarderCommand(origin=self._system_id, command=command)
-            )
-            command_ids.append(str(command.command_id))
-        return {"status": "ok", "commandIds": command_ids}
-
     @staticmethod
     def _get_trace_path(task_id: str) -> Path:
         trace_path = EXO_TRACING_CACHE_DIR / f"trace_{task_id}.json"

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

@@ -1,12 +1,8 @@
 from __future__ import annotations
 
-import json
 from dataclasses import dataclass, field
-from datetime import UTC, datetime
 from pathlib import Path
 
-import aiofiles
-import aiofiles.os as aios
 import anyio
 from anyio import BrokenResourceError, ClosedResourceError, current_time, to_thread
 from loguru import logger
@@ -19,19 +15,8 @@ from exo.download.download_utils import (
     resolve_existing_model,
 )
 from exo.download.shard_downloader import ShardDownloader
-from exo.shared.constants import (
-    EXO_DEFAULT_MODELS_DIR,
-    EXO_MODEL_USAGE_FILE,
-    EXO_MODELS_DIRS,
-    EXO_MODELS_READ_ONLY_DIRS,
-)
-from exo.shared.models import model_cards
-from exo.shared.models.model_cards import ModelId
-from exo.shared.storage import (
-    calculate_used_storage,
-    decide_storage_action,
-    persist_storage_config,
-)
+from exo.shared.constants import EXO_DEFAULT_MODELS_DIR, EXO_MODELS_READ_ONLY_DIRS
+from exo.shared.models.model_cards import ModelId, get_model_cards
 from exo.shared.types.commands import (
     CancelDownload,
     DeleteDownload,
@@ -41,26 +26,15 @@ from exo.shared.types.commands import (
 from exo.shared.types.common import NodeId
 from exo.shared.types.events import (
     Event,
-    IndexedEvent,
-    InstanceCreated,
-    InstanceDeleted,
     NodeDownloadProgress,
-    StorageConfigUpdated,
 )
 from exo.shared.types.memory import Memory
-from exo.shared.types.storage import (
-    StorageAllow,
-    StorageConfig,
-    StorageEvict,
-    StorageReject,
-)
 from exo.shared.types.worker.downloads import (
-    ModelDownloadFailed,
-    ModelDownloading,
-    ModelNotDownloading,
-    ModelReady,
-    ModelRejected,
-    ModelStatus,
+    DownloadCompleted,
+    DownloadFailed,
+    DownloadOngoing,
+    DownloadPending,
+    DownloadProgress,
 )
 from exo.shared.types.worker.shards import PipelineShardMetadata, ShardMetadata
 from exo.utils.channels import Receiver, Sender
@@ -72,18 +46,12 @@ class DownloadCoordinator:
     node_id: NodeId
     shard_downloader: ShardDownloader
     download_command_receiver: Receiver[ForwarderDownloadCommand]
-    event_receiver: Receiver[IndexedEvent]
     event_sender: Sender[Event]
     offline: bool = False
-    storage_config: StorageConfig = field(default_factory=StorageConfig)
 
     # Local state
-    download_status: dict[ModelId, ModelStatus] = field(default_factory=dict)
+    download_status: dict[ModelId, DownloadProgress] = field(default_factory=dict)
     active_downloads: dict[ModelId, anyio.CancelScope] = field(default_factory=dict)
-    _deleting: set[ModelId] = field(default_factory=set)
-
-    _model_last_used: dict[ModelId, datetime] = field(default_factory=dict)
-    _active_model_ids: set[ModelId] = field(default_factory=set)
 
     _tg: TaskGroup = field(init=False, default_factory=TaskGroup)
     _stopped: anyio.Event = field(init=False, default_factory=anyio.Event)
@@ -98,28 +66,13 @@ class DownloadCoordinator:
     def _default_model_dir(model_id: ModelId) -> str:
         return str(EXO_DEFAULT_MODELS_DIR / model_id.normalize())
 
-    @staticmethod
-    def _get_disk_free() -> Memory | None:
-        """Get free disk space for the first available models directory."""
-        import shutil
-
-        for candidate_dir in EXO_MODELS_DIRS:
-            if not candidate_dir.exists():
-                continue
-            try:
-                usage = shutil.disk_usage(candidate_dir)
-                return Memory.from_bytes(usage.free)
-            except OSError:
-                continue
-        return None
-
     def _completed_from_path(
         self,
         shard: ShardMetadata,
         found: Path,
         total: Memory,
-    ) -> ModelReady:
-        return ModelReady(
+    ) -> DownloadCompleted:
+        return DownloadCompleted(
             shard_metadata=shard,
             node_id=self.node_id,
             total=total,
@@ -143,7 +96,7 @@ class DownloadCoordinator:
                         callback_shard, found, progress.total
                     )
                 else:
-                    completed = ModelReady(
+                    completed = DownloadCompleted(
                         shard_metadata=callback_shard,
                         node_id=self.node_id,
                         total=progress.total,
@@ -159,7 +112,7 @@ class DownloadCoordinator:
                 and current_time() - self._last_progress_time.get(model_id, 0.0)
                 > throttle_interval_secs
             ):
-                ongoing = ModelDownloading(
+                ongoing = DownloadOngoing(
                     node_id=self.node_id,
                     shard_metadata=callback_shard,
                     download_progress=map_repo_download_progress_to_download_progress_data(
@@ -181,40 +134,13 @@ class DownloadCoordinator:
         logger.info(
             f"Starting DownloadCoordinator{' (offline mode)' if self.offline else ''}"
         )
-        await self._load_model_usage()
         try:
             async with self._tg as tg:
                 tg.start_soon(self._command_processor)
                 tg.start_soon(self._emit_existing_download_progress)
-                tg.start_soon(self._event_watcher)
         finally:
             self._stopped.set()
 
-    async def _event_watcher(self) -> None:
-        active_instances: dict[str, ModelId] = {}
-        with self.event_receiver as events:
-            async for indexed_event in events:
-                match indexed_event.event:
-                    case StorageConfigUpdated(node_id=node_id) if (
-                        node_id == self.node_id
-                    ):
-                        self.storage_config = indexed_event.event.storage_config
-                        await self.clear_rejections()
-                        await persist_storage_config(indexed_event.event.storage_config)
-                    case InstanceCreated(instance=instance):
-                        active_instances[instance.instance_id] = (
-                            instance.shard_assignments.model_id
-                        )
-                        await self.update_active_models(set(active_instances.values()))
-                    case InstanceDeleted(instance_id=instance_id):
-                        if instance_id in active_instances:
-                            del active_instances[instance_id]
-                            await self.update_active_models(
-                                set(active_instances.values())
-                            )
-                    case _:
-                        pass
-
     async def shutdown(self) -> None:
         self._tg.cancel_tasks()
         await self._stopped.wait()
@@ -241,10 +167,10 @@ class DownloadCoordinator:
             current_status = self.download_status[model_id]
             downloaded = Memory()
             total = Memory()
-            if isinstance(current_status, ModelDownloading):
+            if isinstance(current_status, DownloadOngoing):
                 downloaded = current_status.download_progress.downloaded
                 total = current_status.download_progress.total
-            pending = ModelNotDownloading(
+            pending = DownloadPending(
                 shard_metadata=current_status.shard_metadata,
                 node_id=self.node_id,
                 model_directory=self._default_model_dir(model_id),
@@ -259,16 +185,14 @@ class DownloadCoordinator:
     async def _start_download(self, shard: ShardMetadata) -> None:
         model_id = shard.model_card.model_id
 
-        # Check if already downloading or complete
+        # Check if already downloading, complete, or recently failed
         if model_id in self.download_status:
             status = self.download_status[model_id]
-            if isinstance(status, (ModelDownloading, ModelReady)):
+            if isinstance(status, (DownloadOngoing, DownloadCompleted, DownloadFailed)):
                 logger.debug(
-                    f"Download for {model_id} skipped: current status is {type(status).__name__}"
+                    f"Download for {model_id} already in progress, complete, or failed, skipping"
                 )
                 return
-            if isinstance(status, (ModelRejected, ModelDownloadFailed)):
-                del self.download_status[model_id]
 
         # Check all model directories for pre-existing complete models
         found_path = await to_thread.run_sync(
@@ -285,27 +209,8 @@ class DownloadCoordinator:
             )
             return
 
-        disk_free = await to_thread.run_sync(self._get_disk_free)
-        action = decide_storage_action(
-            shard.model_card.storage_size,
-            self.storage_config,
-            list(self.download_status.values()),
-            self._model_last_used,
-            frozenset(self._active_model_ids),
-            disk_free=disk_free,
-        )
-        match action:
-            case StorageReject(reason=reason, available=available):
-                await self._reject_download(shard, reason, available)
-                return
-            case StorageEvict(model_ids=model_ids):
-                if not await self._execute_evictions(model_ids, shard):
-                    return
-            case StorageAllow():
-                pass
-
         # Emit pending status
-        progress = ModelNotDownloading(
+        progress = DownloadPending(
             shard_metadata=shard,
             node_id=self.node_id,
             model_directory=self._default_model_dir(model_id),
@@ -327,7 +232,7 @@ class DownloadCoordinator:
                     shard, found, initial_progress.total
                 )
             else:
-                completed = ModelReady(
+                completed = DownloadCompleted(
                     shard_metadata=shard,
                     node_id=self.node_id,
                     total=initial_progress.total,
@@ -343,7 +248,7 @@ class DownloadCoordinator:
             logger.warning(
                 f"Offline mode: model {model_id} is not fully available locally, cannot download"
             )
-            failed = ModelDownloadFailed(
+            failed = DownloadFailed(
                 shard_metadata=shard,
                 node_id=self.node_id,
                 error_message=f"Model files not found locally in offline mode: {model_id}",
@@ -362,7 +267,7 @@ class DownloadCoordinator:
         model_id = shard.model_card.model_id
 
         # Emit ongoing status
-        status = ModelDownloading(
+        status = DownloadOngoing(
             node_id=self.node_id,
             shard_metadata=shard,
             download_progress=map_repo_download_progress_to_download_progress_data(
@@ -379,7 +284,7 @@ class DownloadCoordinator:
                     await self.shard_downloader.ensure_shard(shard)
             except Exception as e:
                 logger.error(f"Download failed for {model_id}: {e}")
-                failed = ModelDownloadFailed(
+                failed = DownloadFailed(
                     shard_metadata=shard,
                     node_id=self.node_id,
                     error_message=str(e),
@@ -389,6 +294,9 @@ class DownloadCoordinator:
                 await self.event_sender.send(
                     NodeDownloadProgress(download_progress=failed)
                 )
+            except anyio.get_cancelled_exc_class():
+                # ignore cancellation - let cleanup do its thing
+                pass
             finally:
                 self.active_downloads.pop(model_id, None)
 
@@ -396,15 +304,13 @@ class DownloadCoordinator:
         self._tg.start_soon(download_wrapper, scope)
         self.active_downloads[model_id] = scope
 
-    async def _remove_model_from_disk(self, model_id: ModelId) -> bool:
+    async def _delete_download(self, model_id: ModelId) -> None:
         # Protect read-only models from deletion
         if model_id in self.download_status:
             current = self.download_status[model_id]
-            if isinstance(current, ModelReady) and current.read_only:
-                logger.warning(
-                    f"Refusing to delete read-only model {model_id} (from EXO_MODELS_READ_ONLY_DIRS)"
-                )
-                return False
+            if isinstance(current, DownloadCompleted) and current.read_only:
+                logger.warning(f"Refusing to delete read-only model {model_id}")
+                return
 
         # Cancel if active
         if model_id in self.active_downloads:
@@ -415,22 +321,15 @@ class DownloadCoordinator:
         logger.info(f"Deleting model files for {model_id}")
         deleted = await delete_model(model_id)
 
-        if not deleted:
-            logger.warning(f"Failed to delete model {model_id} from disk")
-            return False
-
-        logger.info(f"Successfully deleted model {model_id}")
-        return True
-
-    async def _delete_download(self, model_id: ModelId) -> bool:
-        success = await self._remove_model_from_disk(model_id)
-        if not success:
-            return False
+        if deleted:
+            logger.info(f"Successfully deleted model {model_id}")
+        else:
+            logger.warning(f"Model {model_id} was not found on disk")
 
         # Emit pending status to reset UI state, then remove from local tracking
         if model_id in self.download_status:
             current_status = self.download_status[model_id]
-            pending = ModelNotDownloading(
+            pending = DownloadPending(
                 shard_metadata=current_status.shard_metadata,
                 node_id=self.node_id,
                 model_directory=self._default_model_dir(model_id),
@@ -440,35 +339,22 @@ class DownloadCoordinator:
             )
             del self.download_status[model_id]
 
-        return True
-
     async def _emit_existing_download_progress(self) -> None:
         while True:
             try:
                 logger.debug(
                     "DownloadCoordinator: Fetching and emitting existing download progress..."
                 )
-
                 async for (
                     _,
                     progress,
                 ) in self.shard_downloader.get_shard_download_status():
                     model_id = progress.shard.model_card.model_id
 
-                    # Don't overwrite status while deletion is in progress
-                    if model_id in self._deleting:
-                        continue
-
                     # Active downloads emit progress via the callback — don't overwrite
                     if model_id in self.active_downloads:
                         continue
 
-                    if isinstance(
-                        self.download_status.get(model_id),
-                        (ModelRejected, ModelDownloadFailed),
-                    ):
-                        continue
-
                     if progress.status == "complete":
                         found = await to_thread.run_sync(
                             resolve_existing_model,
@@ -476,11 +362,11 @@ class DownloadCoordinator:
                             progress.shard.model_card,
                         )
                         if found is not None:
-                            status: ModelStatus = self._completed_from_path(
+                            status: DownloadProgress = self._completed_from_path(
                                 progress.shard, found, progress.total
                             )
                         else:
-                            status = ModelReady(
+                            status = DownloadCompleted(
                                 node_id=self.node_id,
                                 shard_metadata=progress.shard,
                                 total=progress.total,
@@ -489,7 +375,9 @@ class DownloadCoordinator:
                     elif progress.status in ["in_progress", "not_started"]:
                         # TODO(ciaran): temporary solution
                         # Don't downgrade a model that is already confirmed complete.
-                        if isinstance(self.download_status.get(model_id), ModelReady):
+                        if isinstance(
+                            self.download_status.get(model_id), DownloadCompleted
+                        ):
                             continue
                         # The per-file size check compares local files against
                         # the latest HF "main" revision, which is a moving
@@ -509,7 +397,7 @@ class DownloadCoordinator:
                                 progress.shard, found, progress.total
                             )
                         elif progress.downloaded_this_session.in_bytes == 0:
-                            status = ModelNotDownloading(
+                            status = DownloadPending(
                                 node_id=self.node_id,
                                 shard_metadata=progress.shard,
                                 model_directory=self._default_model_dir(model_id),
@@ -517,7 +405,7 @@ class DownloadCoordinator:
                                 total=progress.total,
                             )
                         else:
-                            status = ModelDownloading(
+                            status = DownloadOngoing(
                                 node_id=self.node_id,
                                 shard_metadata=progress.shard,
                                 download_progress=map_repo_download_progress_to_download_progress_data(
@@ -528,19 +416,19 @@ class DownloadCoordinator:
                     else:
                         continue
 
-                    self.download_status[model_id] = status
+                    self.download_status[progress.shard.model_card.model_id] = status
                     await self.event_sender.send(
                         NodeDownloadProgress(download_progress=status)
                     )
                 # 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
                         if isinstance(
                             self.download_status.get(mid),
-                            (ModelReady, ModelDownloading, ModelDownloadFailed),
+                            (DownloadCompleted, DownloadOngoing, DownloadFailed),
                         ):
                             continue
                         found = await to_thread.run_sync(
@@ -555,8 +443,10 @@ class DownloadCoordinator:
                                 end_layer=card.n_layers,
                                 n_layers=card.n_layers,
                             )
-                            path_completed: ModelStatus = self._completed_from_path(
-                                path_shard, found, card.storage_size
+                            path_completed: DownloadProgress = (
+                                self._completed_from_path(
+                                    path_shard, found, card.storage_size
+                                )
                             )
                             self.download_status[mid] = path_completed
                             await self.event_sender.send(
@@ -571,116 +461,3 @@ class DownloadCoordinator:
                     f"DownloadCoordinator: Error emitting existing download progress: {e}"
                 )
             await anyio.sleep(60)
-
-    async def _reject_download(
-        self, shard: ShardMetadata, reason: str, available: Memory
-    ) -> None:
-        model_id = shard.model_card.model_id
-        rejected = ModelRejected(
-            shard_metadata=shard,
-            node_id=self.node_id,
-            model_directory=self._default_model_dir(model_id),
-            reason=reason,
-            required=shard.model_card.storage_size,
-            available=available if available.in_bytes > 0 else Memory(),
-            limit=self.storage_config.max_storage,
-        )
-        self.download_status[model_id] = rejected
-        await self.event_sender.send(NodeDownloadProgress(download_progress=rejected))
-
-    async def _execute_evictions(
-        self, model_ids: list[ModelId], shard: ShardMetadata
-    ) -> bool:
-        """Execute disk deletions for the given model IDs. Returns False on failure."""
-        target_model_id = shard.model_card.model_id
-
-        for evict_model_id in model_ids:
-            logger.info(
-                f"Auto-evicting model {evict_model_id} to free space for {target_model_id}"
-            )
-            evicted_status = self.download_status.get(evict_model_id)
-            self._deleting.add(evict_model_id)
-            try:
-                success = await self._remove_model_from_disk(evict_model_id)
-            finally:
-                self._deleting.discard(evict_model_id)
-
-            if not success:
-                current_used = calculate_used_storage(
-                    list(self.download_status.values())
-                )
-                if self.storage_config.max_storage is not None:
-                    current_available = self.storage_config.max_storage - current_used
-                else:
-                    disk_free = self._get_disk_free()
-                    current_available = disk_free if disk_free is not None else Memory()
-                await self._reject_download(
-                    shard,
-                    f"Failed to delete model {evict_model_id} from disk",
-                    current_available,
-                )
-                return False
-
-            if evicted_status is not None:
-                not_downloading = ModelNotDownloading(
-                    shard_metadata=evicted_status.shard_metadata,
-                    node_id=self.node_id,
-                    model_directory=self._default_model_dir(evict_model_id),
-                )
-                await self.event_sender.send(
-                    NodeDownloadProgress(download_progress=not_downloading)
-                )
-                del self.download_status[evict_model_id]
-
-        return True
-
-    async def clear_rejections(self) -> None:
-        rejected = [
-            (model_id, status)
-            for model_id, status in self.download_status.items()
-            if isinstance(status, ModelRejected)
-        ]
-        for model_id, status in rejected:
-            logger.info(
-                f"Clearing ModelRejected for {model_id} after storage config change"
-            )
-            pending = ModelNotDownloading(
-                shard_metadata=status.shard_metadata,
-                node_id=self.node_id,
-                model_directory=self._default_model_dir(model_id),
-            )
-            self.download_status[model_id] = pending
-            await self.event_sender.send(
-                NodeDownloadProgress(download_progress=pending)
-            )
-
-    async def update_active_models(self, active_model_ids: set[ModelId]) -> None:
-        new_models = active_model_ids - self._active_model_ids
-        for mid in new_models:
-            self._model_last_used[mid] = datetime.now(UTC)
-        self._active_model_ids = active_model_ids.copy()
-        if new_models:
-            await self._persist_model_usage()
-
-    async def _persist_model_usage(self) -> None:
-        try:
-            await aios.makedirs(EXO_MODEL_USAGE_FILE.parent, exist_ok=True)
-            data = {mid: ts.isoformat() for mid, ts in self._model_last_used.items()}
-            async with aiofiles.open(EXO_MODEL_USAGE_FILE, "w") as f:
-                await f.write(json.dumps(data))
-        except Exception as e:
-            logger.warning(f"Failed to persist model usage: {e}")
-
-    async def _load_model_usage(self) -> None:
-        try:
-            if await aios.path.exists(EXO_MODEL_USAGE_FILE):
-                async with aiofiles.open(EXO_MODEL_USAGE_FILE, "r") as f:
-                    raw: dict[str, str] = json.loads(await f.read())  # pyright: ignore[reportAny]
-                self._model_last_used = {
-                    ModelId(k): datetime.fromisoformat(v) for k, v in raw.items()
-                }
-                logger.debug(
-                    f"Loaded model usage for {len(self._model_last_used)} models"
-                )
-        except Exception as e:
-            logger.warning(f"Failed to load model usage: {e}")

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

@@ -1,5 +1,6 @@
 from abc import ABC, abstractmethod
 from collections.abc import Awaitable
+from copy import copy
 from datetime import timedelta
 from pathlib import Path
 from typing import AsyncIterator, Callable
@@ -76,7 +77,9 @@ class NoopShardDownloader(ShardDownloader):
     async def get_shard_download_status_for_shard(
         self, shard: ShardMetadata
     ) -> RepoDownloadProgress:
-        return NOOP_DOWNLOAD_PROGRESS.model_copy(update={"shard": shard})
+        dp = copy(NOOP_DOWNLOAD_PROGRESS)
+        dp.shard = shard
+        return dp
 
 
 NOOP_DOWNLOAD_PROGRESS = RepoDownloadProgress(

src/exo/download/tests/test_auto_eviction.py

@@ -1,546 +0,0 @@
-"""Tests for auto-eviction in the DownloadCoordinator.
-
-Tests exercise _start_download (the production entry point) to verify that
-storage quota checks and LRU eviction work end-to-end through the coordinator.
-"""
-
-from datetime import UTC, datetime
-from pathlib import Path
-from unittest.mock import AsyncMock, patch
-
-from anyio.streams.memory import MemoryObjectStreamState
-
-from exo.download.coordinator import DownloadCoordinator
-from exo.download.shard_downloader import NoopShardDownloader
-from exo.shared.models.model_cards import ModelCard, ModelId, ModelTask
-from exo.shared.types.commands import ForwarderDownloadCommand
-from exo.shared.types.common import NodeId
-from exo.shared.types.events import Event, IndexedEvent, NodeDownloadProgress
-from exo.shared.types.memory import Memory
-from exo.shared.types.storage import StorageConfig
-from exo.shared.types.worker.downloads import (
-    ModelNotDownloading,
-    ModelReady,
-    ModelRejected,
-)
-from exo.shared.types.worker.shards import PipelineShardMetadata, ShardMetadata
-from exo.utils.channels import Receiver, Sender
-
-MODEL_A = ModelId("org/model-a")
-MODEL_B = ModelId("org/model-b")
-MODEL_C = ModelId("org/model-c")
-MODEL_NEW = ModelId("org/model-new")
-NODE_ID = NodeId("test-node")
-
-
-def _shard(model_id: ModelId, size_gb: float) -> ShardMetadata:
-    return PipelineShardMetadata(
-        model_card=ModelCard(
-            model_id=model_id,
-            storage_size=Memory.from_gb(size_gb),
-            n_layers=32,
-            hidden_size=1000,
-            supports_tensor=True,
-            tasks=[ModelTask.TextGeneration],
-        ),
-        device_rank=0,
-        world_size=1,
-        start_layer=0,
-        end_layer=32,
-        n_layers=32,
-    )
-
-
-def _completed(model_id: ModelId, size_gb: float) -> ModelReady:
-    return ModelReady(
-        node_id=NODE_ID,
-        shard_metadata=_shard(model_id, size_gb),
-        total=Memory.from_gb(size_gb),
-    )
-
-
-def _make_coordinator(
-    storage_config: StorageConfig,
-    download_status: dict[ModelId, ModelReady | ModelRejected],
-    model_last_used: dict[ModelId, datetime] | None = None,
-) -> tuple[DownloadCoordinator, Receiver[Event]]:
-    state = MemoryObjectStreamState[Event](max_buffer_size=100)
-    event_sender = Sender[Event](_state=state)
-    event_receiver = Receiver[Event](_state=state)
-
-    cmd_state: MemoryObjectStreamState[ForwarderDownloadCommand] = (
-        MemoryObjectStreamState(max_buffer_size=100)
-    )
-    cmd_receiver: Receiver[ForwarderDownloadCommand] = Receiver(_state=cmd_state)
-
-    idx_state: MemoryObjectStreamState[IndexedEvent] = MemoryObjectStreamState(
-        max_buffer_size=100
-    )
-    idx_receiver: Receiver[IndexedEvent] = Receiver(_state=idx_state)
-
-    coordinator = DownloadCoordinator(
-        node_id=NODE_ID,
-        shard_downloader=NoopShardDownloader(),
-        download_command_receiver=cmd_receiver,
-        event_receiver=idx_receiver,
-        event_sender=event_sender,
-        storage_config=storage_config,
-    )
-    coordinator.download_status = dict(download_status)
-    if model_last_used is not None:
-        coordinator._model_last_used = model_last_used  # pyright: ignore[reportPrivateUsage]
-
-    return coordinator, event_receiver
-
-
-async def _start_download(
-    coordinator: DownloadCoordinator, shard: ShardMetadata
-) -> None:
-    await coordinator._start_download(shard)  # pyright: ignore[reportPrivateUsage]
-
-
-class TestStartDownloadAutoEviction:
-    """Tests that go through _start_download — the production entry point."""
-
-    @patch(
-        "exo.download.coordinator.delete_model",
-        new_callable=AsyncMock,
-        return_value=True,
-    )
-    @patch("exo.download.coordinator.resolve_existing_model", return_value=None)
-    async def test_evicts_oldest_model_to_fit_new_download(
-        self, _mock_resolve: AsyncMock, mock_delete: AsyncMock
-    ) -> None:
-        """_start_download should trigger auto-eviction of the oldest model."""
-        config = StorageConfig(
-            max_storage=Memory.from_gb(10), storage_policy="auto-evict"
-        )
-        coordinator, _ = _make_coordinator(
-            config,
-            {MODEL_A: _completed(MODEL_A, 4), MODEL_B: _completed(MODEL_B, 4)},
-            model_last_used={
-                MODEL_A: datetime(2024, 1, 1, tzinfo=UTC),
-                MODEL_B: datetime(2024, 6, 1, tzinfo=UTC),
-            },
-        )
-
-        await _start_download(coordinator, _shard(MODEL_NEW, 5))
-
-        # MODEL_A (oldest) should have been evicted
-        mock_delete.assert_called_once_with(MODEL_A)
-        assert MODEL_A not in coordinator.download_status
-
-    @patch(
-        "exo.download.coordinator.delete_model",
-        new_callable=AsyncMock,
-        return_value=True,
-    )
-    @patch("exo.download.coordinator.resolve_existing_model", return_value=None)
-    async def test_evicts_multiple_in_lru_order(
-        self, _mock_resolve: AsyncMock, mock_delete: AsyncMock
-    ) -> None:
-        """_start_download evicts multiple models oldest-first until space is freed."""
-        config = StorageConfig(
-            max_storage=Memory.from_gb(10), storage_policy="auto-evict"
-        )
-        coordinator, _ = _make_coordinator(
-            config,
-            {
-                MODEL_A: _completed(MODEL_A, 3),
-                MODEL_B: _completed(MODEL_B, 3),
-                MODEL_C: _completed(MODEL_C, 3),
-            },
-            model_last_used={
-                MODEL_A: datetime(2024, 1, 1, tzinfo=UTC),
-                MODEL_B: datetime(2024, 6, 1, tzinfo=UTC),
-                MODEL_C: datetime(2024, 12, 1, tzinfo=UTC),
-            },
-        )
-
-        # Need 8 GiB, have 1 GiB free — need to free 7 GiB
-        await _start_download(coordinator, _shard(MODEL_NEW, 8))
-
-        evicted = [call.args[0] for call in mock_delete.call_args_list]
-        assert evicted == [MODEL_A, MODEL_B, MODEL_C]
-
-    @patch(
-        "exo.download.coordinator.delete_model",
-        new_callable=AsyncMock,
-        return_value=True,
-    )
-    @patch("exo.download.coordinator.resolve_existing_model", return_value=None)
-    async def test_rejects_when_cannot_free_enough_space(
-        self, _mock_resolve: AsyncMock, mock_delete: AsyncMock
-    ) -> None:
-        """_start_download emits DownloadRejected when eviction can't free enough."""
-        config = StorageConfig(
-            max_storage=Memory.from_gb(10), storage_policy="auto-evict"
-        )
-        coordinator, _ = _make_coordinator(
-            config,
-            {MODEL_A: _completed(MODEL_A, 2)},
-        )
-
-        await _start_download(coordinator, _shard(MODEL_NEW, 20))
-
-        mock_delete.assert_not_called()
-        assert isinstance(coordinator.download_status[MODEL_NEW], ModelRejected)
-
-    @patch(
-        "exo.download.coordinator.delete_model",
-        new_callable=AsyncMock,
-        return_value=True,
-    )
-    @patch("exo.download.coordinator.resolve_existing_model", return_value=None)
-    async def test_no_eviction_when_space_available(
-        self, _mock_resolve: AsyncMock, mock_delete: AsyncMock
-    ) -> None:
-        """_start_download proceeds without evicting when enough space exists."""
-        config = StorageConfig(
-            max_storage=Memory.from_gb(10), storage_policy="auto-evict"
-        )
-        coordinator, _ = _make_coordinator(
-            config,
-            {MODEL_A: _completed(MODEL_A, 2)},
-        )
-
-        await _start_download(coordinator, _shard(MODEL_NEW, 5))
-
-        mock_delete.assert_not_called()
-        # Download should have started (not rejected)
-        assert MODEL_NEW in coordinator.download_status
-        assert not isinstance(coordinator.download_status[MODEL_NEW], ModelRejected)
-
-    @patch(
-        "exo.download.coordinator.delete_model",
-        new_callable=AsyncMock,
-        return_value=True,
-    )
-    @patch("exo.download.coordinator.resolve_existing_model", return_value=None)
-    async def test_manual_policy_rejects_instead_of_evicting(
-        self, _mock_resolve: AsyncMock, mock_delete: AsyncMock
-    ) -> None:
-        """With manual policy, _start_download rejects instead of auto-evicting."""
-        config = StorageConfig(max_storage=Memory.from_gb(10), storage_policy="manual")
-        coordinator, _ = _make_coordinator(
-            config,
-            {MODEL_A: _completed(MODEL_A, 4), MODEL_B: _completed(MODEL_B, 4)},
-        )
-
-        await _start_download(coordinator, _shard(MODEL_NEW, 5))
-
-        mock_delete.assert_not_called()
-        assert isinstance(coordinator.download_status[MODEL_NEW], ModelRejected)
-
-    @patch(
-        "exo.download.coordinator.delete_model",
-        new_callable=AsyncMock,
-        return_value=True,
-    )
-    @patch("exo.download.coordinator.resolve_existing_model", return_value=None)
-    async def test_eviction_emits_not_downloading_event_for_evicted_model(
-        self, _mock_resolve: AsyncMock, mock_delete: AsyncMock
-    ) -> None:
-        """Evicted models emit ModelNotDownloading events and are removed from status."""
-        config = StorageConfig(
-            max_storage=Memory.from_gb(10), storage_policy="auto-evict"
-        )
-        coordinator, event_receiver = _make_coordinator(
-            config,
-            {MODEL_A: _completed(MODEL_A, 4), MODEL_B: _completed(MODEL_B, 4)},
-            model_last_used={
-                MODEL_A: datetime(2024, 1, 1, tzinfo=UTC),
-                MODEL_B: datetime(2024, 6, 1, tzinfo=UTC),
-            },
-        )
-
-        await _start_download(coordinator, _shard(MODEL_NEW, 5))
-
-        events = event_receiver.collect()
-        eviction_events = [
-            e
-            for e in events
-            if isinstance(e, NodeDownloadProgress)
-            and isinstance(e.download_progress, ModelNotDownloading)
-            and e.download_progress.shard_metadata.model_card.model_id == MODEL_A
-        ]
-        assert len(eviction_events) == 1
-        assert MODEL_A not in coordinator.download_status
-
-
-class TestActiveModelProtection:
-    """Tests that update_active_models protects models from eviction.
-
-    These tests use the public API (update_active_models) to mark models as
-    active, then trigger eviction through _start_download.
-    """
-
-    @patch(
-        "exo.download.coordinator.delete_model",
-        new_callable=AsyncMock,
-        return_value=True,
-    )
-    @patch("exo.download.coordinator.resolve_existing_model", return_value=None)
-    async def test_active_model_not_evicted(
-        self, _mock_resolve: AsyncMock, mock_delete: AsyncMock
-    ) -> None:
-        """A model marked active via update_active_models must not be evicted."""
-        config = StorageConfig(
-            max_storage=Memory.from_gb(10), storage_policy="auto-evict"
-        )
-        coordinator, _ = _make_coordinator(
-            config,
-            {MODEL_A: _completed(MODEL_A, 4), MODEL_B: _completed(MODEL_B, 4)},
-            model_last_used={
-                MODEL_A: datetime(2024, 1, 1, tzinfo=UTC),  # oldest
-                MODEL_B: datetime(2024, 6, 1, tzinfo=UTC),
-            },
-        )
-
-        # Mark MODEL_A as active through the public API
-        await coordinator.update_active_models({MODEL_A})
-
-        await _start_download(coordinator, _shard(MODEL_NEW, 5))
-
-        # MODEL_A is active — MODEL_B should be evicted instead
-        mock_delete.assert_called_once_with(MODEL_B)
-        assert MODEL_A in coordinator.download_status
-
-    @patch(
-        "exo.download.coordinator.delete_model",
-        new_callable=AsyncMock,
-        return_value=True,
-    )
-    @patch("exo.download.coordinator.resolve_existing_model", return_value=None)
-    async def test_all_active_models_rejected(
-        self, _mock_resolve: AsyncMock, mock_delete: AsyncMock
-    ) -> None:
-        """When all models are active, eviction is impossible — download is rejected."""
-        config = StorageConfig(
-            max_storage=Memory.from_gb(10), storage_policy="auto-evict"
-        )
-        coordinator, _ = _make_coordinator(
-            config,
-            {MODEL_A: _completed(MODEL_A, 4), MODEL_B: _completed(MODEL_B, 4)},
-        )
-
-        await coordinator.update_active_models({MODEL_A, MODEL_B})
-
-        await _start_download(coordinator, _shard(MODEL_NEW, 5))
-
-        mock_delete.assert_not_called()
-        assert isinstance(coordinator.download_status[MODEL_NEW], ModelRejected)
-
-
-class TestDiskDeleteFailure:
-    """Tests that eviction fails properly when disk delete fails."""
-
-    @patch(
-        "exo.download.coordinator.delete_model",
-        new_callable=AsyncMock,
-        return_value=False,
-    )
-    @patch("exo.download.coordinator.resolve_existing_model", return_value=None)
-    async def test_eviction_rejected_on_disk_delete_failure(
-        self, _mock_resolve: AsyncMock, mock_delete: AsyncMock
-    ) -> None:
-        """When disk delete fails, auto-eviction emits DownloadRejected."""
-        config = StorageConfig(
-            max_storage=Memory.from_gb(10), storage_policy="auto-evict"
-        )
-        coordinator, _ = _make_coordinator(
-            config,
-            {MODEL_A: _completed(MODEL_A, 4), MODEL_B: _completed(MODEL_B, 4)},
-            model_last_used={
-                MODEL_A: datetime(2024, 1, 1, tzinfo=UTC),
-                MODEL_B: datetime(2024, 6, 1, tzinfo=UTC),
-            },
-        )
-
-        await _start_download(coordinator, _shard(MODEL_NEW, 5))
-
-        # Should have tried to delete oldest model and failed
-        mock_delete.assert_called_once_with(MODEL_A)
-        # New model should be rejected
-        assert isinstance(coordinator.download_status[MODEL_NEW], ModelRejected)
-        # Eviction target should still be in download_status (not removed)
-        assert MODEL_A in coordinator.download_status
-
-
-class TestLruPersistence:
-    """Tests for _persist_model_usage and _load_model_usage round-trip."""
-
-    async def test_persist_then_load_round_trip(self, tmp_path: Path) -> None:
-        """Persisting then loading recovers the same data."""
-        usage_file = tmp_path / "model_usage.json"
-        coordinator, _ = _make_coordinator(StorageConfig(), {})
-        coordinator._model_last_used = {  # pyright: ignore[reportPrivateUsage]
-            MODEL_A: datetime(2024, 1, 15, 12, 30, 0, tzinfo=UTC),
-            MODEL_B: datetime(2024, 6, 1, 0, 0, 0, tzinfo=UTC),
-        }
-
-        with patch("exo.download.coordinator.EXO_MODEL_USAGE_FILE", usage_file):
-            await coordinator._persist_model_usage()  # pyright: ignore[reportPrivateUsage]
-
-            # Create a fresh coordinator and load
-            coordinator2, _ = _make_coordinator(StorageConfig(), {})
-            await coordinator2._load_model_usage()  # pyright: ignore[reportPrivateUsage]
-
-        assert coordinator2._model_last_used == coordinator._model_last_used  # pyright: ignore[reportPrivateUsage]
-
-    async def test_load_missing_file_returns_empty(self, tmp_path: Path) -> None:
-        """Loading when file doesn't exist returns empty dict."""
-        usage_file = tmp_path / "nonexistent" / "model_usage.json"
-        coordinator, _ = _make_coordinator(StorageConfig(), {})
-
-        with patch("exo.download.coordinator.EXO_MODEL_USAGE_FILE", usage_file):
-            await coordinator._load_model_usage()  # pyright: ignore[reportPrivateUsage]
-
-        assert coordinator._model_last_used == {}  # pyright: ignore[reportPrivateUsage]
-
-    async def test_load_corrupt_json_returns_empty(self, tmp_path: Path) -> None:
-        """Loading corrupt JSON logs warning and returns empty dict."""
-        usage_file = tmp_path / "model_usage.json"
-        usage_file.write_text("not valid json {{{")
-        coordinator, _ = _make_coordinator(StorageConfig(), {})
-
-        with patch("exo.download.coordinator.EXO_MODEL_USAGE_FILE", usage_file):
-            await coordinator._load_model_usage()  # pyright: ignore[reportPrivateUsage]
-
-        assert coordinator._model_last_used == {}  # pyright: ignore[reportPrivateUsage]
-
-
-class TestEvictionEvents:
-    """Tests that eviction emits the correct sequence of events."""
-
-    @patch(
-        "exo.download.coordinator.delete_model",
-        new_callable=AsyncMock,
-        return_value=True,
-    )
-    @patch("exo.download.coordinator.resolve_existing_model", return_value=None)
-    async def test_eviction_emits_not_downloading_event(
-        self, _mock_resolve: AsyncMock, mock_delete: AsyncMock
-    ) -> None:
-        """Eviction emits a ModelNotDownloading event and removes from status."""
-        config = StorageConfig(
-            max_storage=Memory.from_gb(10), storage_policy="auto-evict"
-        )
-        coordinator, event_receiver = _make_coordinator(
-            config,
-            {MODEL_A: _completed(MODEL_A, 4), MODEL_B: _completed(MODEL_B, 4)},
-            model_last_used={
-                MODEL_A: datetime(2024, 1, 1, tzinfo=UTC),
-                MODEL_B: datetime(2024, 6, 1, tzinfo=UTC),
-            },
-        )
-
-        await _start_download(coordinator, _shard(MODEL_NEW, 5))
-
-        events = event_receiver.collect()
-        eviction_events = [
-            e
-            for e in events
-            if isinstance(e, NodeDownloadProgress)
-            and isinstance(e.download_progress, ModelNotDownloading)
-            and e.download_progress.shard_metadata.model_card.model_id == MODEL_A
-        ]
-        assert len(eviction_events) == 1
-        assert MODEL_A not in coordinator.download_status
-
-    @patch(
-        "exo.download.coordinator.delete_model",
-        new_callable=AsyncMock,
-        return_value=True,
-    )
-    @patch("exo.download.coordinator.resolve_existing_model", return_value=None)
-    async def test_multi_eviction_emits_event_per_model(
-        self, _mock_resolve: AsyncMock, _mock_delete: AsyncMock
-    ) -> None:
-        """Each evicted model gets its own ModelNotDownloading event."""
-        config = StorageConfig(
-            max_storage=Memory.from_gb(10), storage_policy="auto-evict"
-        )
-        coordinator, event_receiver = _make_coordinator(
-            config,
-            {
-                MODEL_A: _completed(MODEL_A, 3),
-                MODEL_B: _completed(MODEL_B, 3),
-                MODEL_C: _completed(MODEL_C, 3),
-            },
-            model_last_used={
-                MODEL_A: datetime(2024, 1, 1, tzinfo=UTC),
-                MODEL_B: datetime(2024, 6, 1, tzinfo=UTC),
-                MODEL_C: datetime(2024, 12, 1, tzinfo=UTC),
-            },
-        )
-
-        await _start_download(coordinator, _shard(MODEL_NEW, 8))
-
-        events = event_receiver.collect()
-        eviction_events = [
-            e
-            for e in events
-            if isinstance(e, NodeDownloadProgress)
-            and isinstance(e.download_progress, ModelNotDownloading)
-            and e.download_progress.shard_metadata.model_card.model_id != MODEL_NEW
-        ]
-        evicted_model_ids = [
-            e.download_progress.shard_metadata.model_card.model_id
-            for e in eviction_events
-        ]
-        assert evicted_model_ids == [MODEL_A, MODEL_B, MODEL_C]
-        for mid in [MODEL_A, MODEL_B, MODEL_C]:
-            assert mid not in coordinator.download_status
-
-
-class TestClearRejections:
-    """Tests for clear_rejections behavior."""
-
-    async def test_clear_rejections_resets_rejected(self) -> None:
-        """clear_rejections resets ModelRejected to ModelNotDownloading."""
-        config = StorageConfig(
-            max_storage=Memory.from_gb(10), storage_policy="auto-evict"
-        )
-        rejected = ModelRejected(
-            node_id=NODE_ID,
-            shard_metadata=_shard(MODEL_B, 4),
-            reason="Not enough space",
-            required=Memory.from_gb(4),
-            available=Memory.from_gb(1),
-            limit=Memory.from_gb(10),
-        )
-        coordinator, _ = _make_coordinator(
-            config,
-            {MODEL_A: _completed(MODEL_A, 4), MODEL_B: rejected},
-        )
-
-        await coordinator.clear_rejections()
-
-        # Completed should remain unchanged
-        assert isinstance(coordinator.download_status[MODEL_A], ModelReady)
-        # Rejected should be cleared
-        assert isinstance(coordinator.download_status[MODEL_B], ModelNotDownloading)
-
-    async def test_clear_rejections_on_policy_only_change(self) -> None:
-        """clear_rejections fires even when only the policy changes (no limit change)."""
-        config = StorageConfig(max_storage=Memory.from_gb(10), storage_policy="manual")
-        rejected = ModelRejected(
-            node_id=NODE_ID,
-            shard_metadata=_shard(MODEL_A, 4),
-            reason="Manual policy",
-            required=Memory.from_gb(4),
-            available=Memory.from_gb(1),
-            limit=Memory.from_gb(10),
-        )
-        coordinator, _ = _make_coordinator(
-            config,
-            {MODEL_A: rejected, MODEL_B: _completed(MODEL_B, 3)},
-        )
-
-        await coordinator.clear_rejections()
-
-        # Rejected should be cleared to Pending
-        assert isinstance(coordinator.download_status[MODEL_A], ModelNotDownloading)
-        # Completed should be unchanged
-        assert isinstance(coordinator.download_status[MODEL_B], ModelReady)

src/exo/download/tests/test_cancel_download.py

@@ -18,9 +18,9 @@ from exo.shared.types.commands import (
     StartDownload,
 )
 from exo.shared.types.common import NodeId, SystemId
-from exo.shared.types.events import Event, IndexedEvent, NodeDownloadProgress
+from exo.shared.types.events import Event, NodeDownloadProgress
 from exo.shared.types.memory import Memory
-from exo.shared.types.worker.downloads import ModelNotDownloading
+from exo.shared.types.worker.downloads import DownloadPending
 from exo.shared.types.worker.shards import PipelineShardMetadata, ShardMetadata
 from exo.utils.channels import Receiver, Sender, channel
 
@@ -139,13 +139,11 @@ def _setup_coordinator(
 ]:
     cmd_send, cmd_recv = channel[ForwarderDownloadCommand]()
     event_send, event_recv = channel[Event]()
-    _idx_send, idx_recv = channel[IndexedEvent]()
     wrapped = SingletonShardDownloader(downloader)
     coordinator = DownloadCoordinator(
         node_id=NODE_ID,
         shard_downloader=wrapped,
         download_command_receiver=cmd_recv,
-        event_receiver=idx_recv,
         event_sender=event_send,
     )
     return coordinator, cmd_send, event_recv
@@ -153,15 +151,15 @@ def _setup_coordinator(
 
 async def _wait_for_pending(
     event_recv: Receiver[Event], model_id: ModelId, timeout: float = 2.0
-) -> ModelNotDownloading | None:
-    """Drain events until we see a ModelNotDownloading for the given model, or timeout."""
+) -> DownloadPending | None:
+    """Drain events until we see a DownloadPending for the given model, or timeout."""
     try:
         async with asyncio.timeout(timeout):
             while True:
                 event = await event_recv.receive()
                 if (
                     isinstance(event, NodeDownloadProgress)
-                    and isinstance(event.download_progress, ModelNotDownloading)
+                    and isinstance(event.download_progress, DownloadPending)
                     and event.download_progress.shard_metadata.model_card.model_id
                     == model_id
                 ):
@@ -171,7 +169,7 @@ async def _wait_for_pending(
 
 
 async def test_cancel_active_download_transitions_to_pending() -> None:
-    """Cancelling an in-progress download should emit a ModelNotDownloading event
+    """Cancelling an in-progress download should emit a DownloadPending event
     and remove the model from active_downloads."""
     slow_downloader = SlowShardDownloader()
     coordinator, cmd_send, event_recv = _setup_coordinator(slow_downloader)
@@ -191,7 +189,7 @@ async def test_cancel_active_download_transitions_to_pending() -> None:
         # Wait for the download to actually start (blocking in ensure_shard)
         await asyncio.wait_for(slow_downloader.download_started.wait(), timeout=2.0)
 
-        # Drain any events emitted before the cancel (initial ModelNotDownloading, DownloadOngoing)
+        # Drain any events emitted before the cancel (initial DownloadPending, DownloadOngoing)
         while True:
             try:
                 async with asyncio.timeout(0.1):
@@ -207,9 +205,9 @@ async def test_cancel_active_download_transitions_to_pending() -> None:
             )
         )
 
-        # Should receive a ModelNotDownloading event with preserved progress
+        # Should receive a DownloadPending event with preserved progress
         pending = await _wait_for_pending(event_recv, MODEL_ID)
-        assert pending is not None, "Cancel should emit ModelNotDownloading"
+        assert pending is not None, "Cancel should emit DownloadPending"
         assert pending.shard_metadata.model_card.model_id == MODEL_ID
         assert pending.total == Memory.from_mb(100), "Should preserve total bytes"
 
@@ -220,7 +218,7 @@ async def test_cancel_active_download_transitions_to_pending() -> None:
         assert MODEL_ID not in coordinator.active_downloads
         # But should still be in download_status as pending
         assert MODEL_ID in coordinator.download_status
-        assert isinstance(coordinator.download_status[MODEL_ID], ModelNotDownloading)
+        assert isinstance(coordinator.download_status[MODEL_ID], DownloadPending)
     finally:
         await coordinator.shutdown()
         coordinator_task.cancel()
@@ -244,7 +242,7 @@ async def test_cancel_nonexistent_download_is_noop() -> None:
             )
         )
 
-        # Should NOT receive any ModelNotDownloading event
+        # Should NOT receive any DownloadPending event
         pending = await _wait_for_pending(event_recv, MODEL_ID, timeout=0.5)
         assert pending is None, "Cancel of non-existent download should not emit events"
 
@@ -284,7 +282,7 @@ async def test_cancel_then_resume_download() -> None:
             )
         )
         pending = await _wait_for_pending(event_recv, MODEL_ID)
-        assert pending is not None, "Cancel should emit ModelNotDownloading"
+        assert pending is not None, "Cancel should emit DownloadPending"
 
         await asyncio.sleep(0.05)
 

src/exo/download/tests/test_download_status_not_lost.py

@@ -21,11 +21,11 @@ from exo.download.shard_downloader import ShardDownloader
 from exo.shared.models.model_cards import ModelCard, ModelId, ModelTask
 from exo.shared.types.commands import ForwarderDownloadCommand
 from exo.shared.types.common import NodeId
-from exo.shared.types.events import Event, IndexedEvent, NodeDownloadProgress
+from exo.shared.types.events import Event, NodeDownloadProgress
 from exo.shared.types.memory import Memory
 from exo.shared.types.worker.downloads import (
-    ModelNotDownloading,
-    ModelReady,
+    DownloadCompleted,
+    DownloadPending,
 )
 from exo.shared.types.worker.shards import PipelineShardMetadata, ShardMetadata
 from exo.utils.channels import Receiver, Sender, channel
@@ -128,13 +128,11 @@ def _setup_coordinator(
 ]:
     cmd_send, cmd_recv = channel[ForwarderDownloadCommand]()
     event_send, event_recv = channel[Event]()
-    _indexed_send, indexed_recv = channel[IndexedEvent]()
     wrapped = SingletonShardDownloader(downloader)
     coordinator = DownloadCoordinator(
         node_id=NODE_ID,
         shard_downloader=wrapped,
         download_command_receiver=cmd_recv,
-        event_receiver=indexed_recv,
         event_sender=event_send,
     )
     return coordinator, cmd_send, event_recv
@@ -155,14 +153,14 @@ async def _collect_events(
 
 
 async def test_completed_status_not_downgraded_by_rescan() -> None:
-    """A model already marked ModelReady must not revert to
-    ModelNotDownloading when the periodic rescan reports a non-complete
+    """A model already marked DownloadCompleted must not revert to
+    DownloadPending when the periodic rescan reports a non-complete
     file-size status (regression test for #1918)."""
     downloader = FakeShardDownloader(status="not_started")
     coordinator, _cmd_send, event_recv = _setup_coordinator(downloader)
 
     # Pre-seed the coordinator with a completed status for the model
-    completed = ModelReady(
+    completed = DownloadCompleted(
         node_id=NODE_ID,
         shard_metadata=SHARD,
         total=Memory.from_mb(100),
@@ -176,21 +174,21 @@ async def test_completed_status_not_downgraded_by_rescan() -> None:
         # Wait for the rescan to process (it should skip the completed model)
         events = await _collect_events(event_recv, timeout=1.5)
 
-        # The model must still be ModelReady — not downgraded
-        assert isinstance(coordinator.download_status[MODEL_ID], ModelReady), (
-            f"Expected ModelReady but got {type(coordinator.download_status[MODEL_ID]).__name__}"
+        # The model must still be DownloadCompleted — not downgraded
+        assert isinstance(coordinator.download_status[MODEL_ID], DownloadCompleted), (
+            f"Expected DownloadCompleted but got {type(coordinator.download_status[MODEL_ID]).__name__}"
         )
 
-        # No ModelNotDownloading event should have been emitted for this model
+        # No DownloadPending event should have been emitted for this model
         pending_events = [
             e
             for e in events
             if isinstance(e, NodeDownloadProgress)
-            and isinstance(e.download_progress, ModelNotDownloading)
+            and isinstance(e.download_progress, DownloadPending)
             and e.download_progress.shard_metadata.model_card.model_id == MODEL_ID
         ]
         assert len(pending_events) == 0, (
-            f"Expected no ModelNotDownloading events for completed model, got {len(pending_events)}"
+            f"Expected no DownloadPending events for completed model, got {len(pending_events)}"
         )
     finally:
         await coordinator.shutdown()
@@ -202,7 +200,7 @@ async def test_completed_status_not_downgraded_by_rescan() -> None:
 async def test_incomplete_model_with_files_present_detected_as_complete() -> None:
     """When the per-file size check says not_started but resolve_existing_model
     confirms the model directory is complete, the model should be marked
-    ModelReady (regression test for #1918 — initial scan case)."""
+    DownloadCompleted (regression test for #1918 — initial scan case)."""
     downloader = FakeShardDownloader(status="not_started")
     coordinator, _cmd_send, event_recv = _setup_coordinator(downloader)
 
@@ -215,21 +213,25 @@ async def test_incomplete_model_with_files_present_detected_as_complete() -> Non
         try:
             events = await _collect_events(event_recv, timeout=1.5)
 
-            # The model should be ModelReady (resolve_existing_model confirmed it)
-            assert isinstance(coordinator.download_status.get(MODEL_ID), ModelReady), (
-                f"Expected ModelReady but got "
+            # The model should be DownloadCompleted (resolve_existing_model confirmed it)
+            assert isinstance(
+                coordinator.download_status.get(MODEL_ID), DownloadCompleted
+            ), (
+                f"Expected DownloadCompleted but got "
                 f"{type(coordinator.download_status.get(MODEL_ID)).__name__}"
             )
 
-            # Should have emitted a ModelReady event
+            # Should have emitted a DownloadCompleted event
             completed_events = [
                 e
                 for e in events
                 if isinstance(e, NodeDownloadProgress)
-                and isinstance(e.download_progress, ModelReady)
+                and isinstance(e.download_progress, DownloadCompleted)
                 and e.download_progress.shard_metadata.model_card.model_id == MODEL_ID
             ]
-            assert len(completed_events) > 0, "Expected at least one ModelReady event"
+            assert len(completed_events) > 0, (
+                "Expected at least one DownloadCompleted event"
+            )
         finally:
             await coordinator.shutdown()
             coordinator_task.cancel()
@@ -240,7 +242,7 @@ async def test_incomplete_model_with_files_present_detected_as_complete() -> Non
 async def test_genuinely_incomplete_model_stays_pending() -> None:
     """When the per-file size check says not_started and resolve_existing_model
     returns None (model truly not complete), the model should correctly be
-    ModelNotDownloading."""
+    DownloadPending."""
     downloader = FakeShardDownloader(status="not_started")
     coordinator, _cmd_send, event_recv = _setup_coordinator(downloader)
 
@@ -253,24 +255,24 @@ async def test_genuinely_incomplete_model_stays_pending() -> None:
         try:
             events = await _collect_events(event_recv, timeout=1.5)
 
-            # The model should be ModelNotDownloading
+            # The model should be DownloadPending
             assert isinstance(
-                coordinator.download_status.get(MODEL_ID), ModelNotDownloading
+                coordinator.download_status.get(MODEL_ID), DownloadPending
             ), (
-                f"Expected ModelNotDownloading but got "
+                f"Expected DownloadPending but got "
                 f"{type(coordinator.download_status.get(MODEL_ID)).__name__}"
             )
 
-            # Should have emitted a ModelNotDownloading event
+            # Should have emitted a DownloadPending event
             pending_events = [
                 e
                 for e in events
                 if isinstance(e, NodeDownloadProgress)
-                and isinstance(e.download_progress, ModelNotDownloading)
+                and isinstance(e.download_progress, DownloadPending)
                 and e.download_progress.shard_metadata.model_card.model_id == MODEL_ID
             ]
             assert len(pending_events) > 0, (
-                "Expected at least one ModelNotDownloading event"
+                "Expected at least one DownloadPending event"
             )
         finally:
             await coordinator.shutdown()

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/download/tests/test_re_download.py

@@ -19,9 +19,9 @@ from exo.shared.types.commands import (
     StartDownload,
 )
 from exo.shared.types.common import NodeId, SystemId
-from exo.shared.types.events import Event, IndexedEvent, NodeDownloadProgress
+from exo.shared.types.events import Event, NodeDownloadProgress
 from exo.shared.types.memory import Memory
-from exo.shared.types.worker.downloads import ModelReady
+from exo.shared.types.worker.downloads import DownloadCompleted
 from exo.shared.types.worker.shards import PipelineShardMetadata, ShardMetadata
 from exo.utils.channels import Receiver, Sender, channel
 
@@ -132,7 +132,6 @@ async def test_re_download_after_delete_completes() -> None:
     cmd_send: Sender[ForwarderDownloadCommand]
     cmd_send, cmd_recv = channel[ForwarderDownloadCommand]()
     event_send, event_recv = channel[Event]()
-    _idx_send, idx_recv = channel[IndexedEvent]()
 
     fake_downloader = FakeShardDownloader()
     wrapped_downloader = SingletonShardDownloader(fake_downloader)
@@ -140,7 +139,6 @@ async def test_re_download_after_delete_completes() -> None:
         node_id=NODE_ID,
         shard_downloader=wrapped_downloader,
         download_command_receiver=cmd_recv,
-        event_receiver=idx_recv,
         event_sender=event_send,
     )
 
@@ -196,7 +194,7 @@ async def test_re_download_after_delete_completes() -> None:
 
 async def _wait_for_download_completed(
     event_recv: Receiver[Event], model_id: ModelId, timeout: float = 2.0
-) -> ModelReady | None:
+) -> DownloadCompleted | None:
     """Drain events until we see a DownloadCompleted for the given model, or timeout."""
     try:
         async with asyncio.timeout(timeout):
@@ -204,7 +202,7 @@ async def _wait_for_download_completed(
                 event = await event_recv.receive()
                 if (
                     isinstance(event, NodeDownloadProgress)
-                    and isinstance(event.download_progress, ModelReady)
+                    and isinstance(event.download_progress, DownloadCompleted)
                     and event.download_progress.shard_metadata.model_card.model_id
                     == model_id
                 ):

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
 
@@ -21,12 +20,8 @@ from exo.routing.router import Router, get_node_id_keypair
 from exo.shared.constants import EXO_LOG
 from exo.shared.election import Election, ElectionResult
 from exo.shared.logging import logger_cleanup, logger_setup
-from exo.shared.storage import load_storage_config
 from exo.shared.types.common import NodeId, SessionId
-from exo.shared.types.storage import StoragePolicy
 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
@@ -73,21 +68,14 @@ class Node:
 
         logger.info(f"Starting node {node_id}")
 
-        storage_config = await load_storage_config(
-            max_storage_gb=args.max_storage_gb,
-            storage_policy=args.storage_policy,
-        )
-
         # Create DownloadCoordinator (unless --no-downloads)
         if not args.no_downloads:
             download_coordinator = DownloadCoordinator(
                 node_id,
                 exo_shard_downloader(offline=args.offline),
-                download_command_receiver=router.receiver(topics.DOWNLOAD_COMMANDS),
-                event_receiver=event_router.receiver(),
                 event_sender=event_router.sender(),
+                download_command_receiver=router.receiver(topics.DOWNLOAD_COMMANDS),
                 offline=args.offline,
-                storage_config=storage_config,
             )
         else:
             download_coordinator = None
@@ -243,22 +231,15 @@ class Node:
                 if result.is_new_master:
                     if self.download_coordinator:
                         await self.download_coordinator.shutdown()
-                        storage_config = self.download_coordinator.storage_config
-                        active_model_ids = self.download_coordinator._active_model_ids  # pyright: ignore[reportPrivateUsage]
-                        model_last_used = self.download_coordinator._model_last_used  # pyright: ignore[reportPrivateUsage]
                         self.download_coordinator = DownloadCoordinator(
                             self.node_id,
                             exo_shard_downloader(offline=self.offline),
+                            event_sender=self.event_router.sender(),
                             download_command_receiver=self.router.receiver(
                                 topics.DOWNLOAD_COMMANDS
                             ),
-                            event_receiver=self.event_router.receiver(),
-                            event_sender=self.event_router.sender(),
                             offline=self.offline,
-                            storage_config=storage_config,
                         )
-                        self.download_coordinator._active_model_ids = active_model_ids  # pyright: ignore[reportPrivateUsage]
-                        self.download_coordinator._model_last_used = model_last_used  # pyright: ignore[reportPrivateUsage]
                         self._tg.start_soon(self.download_coordinator.run)
                     if self.worker:
                         await self.worker.shutdown()
@@ -283,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}")
@@ -337,7 +306,6 @@ def main():
     finally:
         logger.info("EXO Shutdown complete")
         logger_cleanup()
-        del pidfile
 
 
 class Args(FrozenModel):
@@ -351,11 +319,8 @@ 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
-    max_storage_gb: float | None = None
-    storage_policy: StoragePolicy | None = None
 
     @classmethod
     def parse(cls) -> Self:
@@ -413,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],
@@ -448,20 +408,6 @@ class Args(FrozenModel):
             dest="fast_synch",
             help="Force MLX FAST_SYNCH off",
         )
-        parser.add_argument(
-            "--max-storage-gb",
-            type=float,
-            dest="max_storage_gb",
-            default=None,
-            help="Maximum storage for downloaded models in GB (default: unlimited)",
-        )
-        parser.add_argument(
-            "--storage-policy",
-            choices=["manual", "auto-evict"],
-            dest="storage_policy",
-            default=None,
-            help="Storage policy: 'manual' rejects on exceed, 'auto-evict' removes LRU models (default: manual)",
-        )
 
         args = parser.parse_args()
         return cls(**vars(args))  # pyright: ignore[reportAny] - We are intentionally validating here, we can't do it statically

src/exo/master/main.py

@@ -13,7 +13,6 @@ from exo.master.placement import (
 from exo.master.placement_utils import find_ip_prioritised
 from exo.shared.apply import apply
 from exo.shared.constants import EXO_EVENT_LOG_DIR, EXO_TRACING_ENABLED
-from exo.shared.storage import get_download_rejected_events
 from exo.shared.types.commands import (
     AddCustomModelCard,
     CreateInstance,
@@ -28,7 +27,6 @@ from exo.shared.types.commands import (
     RequestEventLog,
     SendInputChunk,
     SetInstanceLink,
-    SetStorageConfig,
     TaskCancelled,
     TaskFinished,
     TestCommand,
@@ -46,10 +44,8 @@ from exo.shared.types.events import (
     InstanceLinkCreated,
     InstanceLinkDeleted,
     LocalForwarderEvent,
-    NodeDownloadProgress,
     NodeGatheredInfo,
     NodeTimedOut,
-    StorageConfigUpdated,
     TaskCreated,
     TaskDeleted,
     TaskStatusUpdated,
@@ -59,7 +55,6 @@ from exo.shared.types.events import (
 )
 from exo.shared.types.instance_link import InstanceLink
 from exo.shared.types.state import State
-from exo.shared.types.storage import StorageConfig
 from exo.shared.types.tasks import (
     ImageEdits as ImageEditsTask,
 )
@@ -73,7 +68,6 @@ from exo.shared.types.tasks import (
 from exo.shared.types.tasks import (
     TextGeneration as TextGenerationTask,
 )
-from exo.shared.types.worker.downloads import ModelDownloadFailed, ModelRejected
 from exo.shared.types.worker.instances import InstanceId
 from exo.utils.channels import Receiver, Sender
 from exo.utils.disk_event_log import DiskEventLog
@@ -171,9 +165,7 @@ class Master:
         with self.command_receiver as commands:
             async for forwarder_command in commands:
                 try:
-                    logger.info(
-                        f"Executing command from {forwarder_command.origin}: {forwarder_command.command}"
-                    )
+                    logger.info(f"Executing command: {forwarder_command.command}")
 
                     generated_events: list[Event] = []
                     command = forwarder_command.command
@@ -188,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}
@@ -373,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
@@ -447,16 +459,6 @@ class Master:
                             generated_events.append(
                                 InstanceLinkDeleted(link_id=command.link_id)
                             )
-                        case SetStorageConfig():
-                            generated_events.append(
-                                StorageConfigUpdated(
-                                    node_id=command.target_node_id,
-                                    storage_config=StorageConfig(
-                                        max_storage=command.max_storage,
-                                        storage_policy=command.storage_policy,
-                                    ),
-                                )
-                            )
                         case RequestEventLog():
                             # We should just be able to send everything, since other buffers will ignore old messages
                             # rate limit to 1000 at a time
@@ -522,20 +524,6 @@ class Master:
                     indexed = IndexedEvent(event=event, idx=len(self._event_log))
                     self.state = apply(self.state, indexed)
 
-                    if isinstance(event, NodeDownloadProgress) and isinstance(
-                        event.download_progress, (ModelRejected, ModelDownloadFailed)
-                    ):
-                        dp = event.download_progress
-                        cleanup_events = get_download_rejected_events(
-                            dp.shard_metadata.model_card.model_id,
-                            dp.node_id,
-                            self.state.instances,
-                            self.state.tasks,
-                        )
-                        for cleanup_event in cleanup_events:
-                            logger.info(f"Download failure cleanup: {cleanup_event}")
-                            await self.event_sender.send(cleanup_event)
-
                     self._event_log.append(event)
                     await self._send_event(indexed)
 

src/exo/master/placement.py

@@ -28,15 +28,14 @@ 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 (
-    ModelDownloadFailed,
-    ModelDownloading,
-    ModelNotDownloading,
-    ModelReady,
-    ModelRejected,
-    ModelStatus,
+    DownloadCompleted,
+    DownloadFailed,
+    DownloadOngoing,
+    DownloadPending,
+    DownloadProgress,
 )
 from exo.shared.types.worker.instances import (
     Instance,
@@ -44,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
@@ -62,28 +62,26 @@ def add_instance_to_placements(
 def _get_node_download_fraction(
     node_id: NodeId,
     model_id: ModelId,
-    download_status: Mapping[NodeId, Sequence[ModelStatus]],
+    download_status: Mapping[NodeId, Sequence[DownloadProgress]],
 ) -> float:
     """Return the download fraction (0.0–1.0) for a model on a given node."""
     for progress in download_status.get(node_id, []):
         if progress.shard_metadata.model_card.model_id != model_id:
             continue
         match progress:
-            case ModelReady():
+            case DownloadCompleted():
                 return 1.0
-            case ModelDownloading():
+            case DownloadOngoing():
                 total = progress.download_progress.total.in_bytes
                 return (
                     progress.download_progress.downloaded.in_bytes / total
                     if total > 0
                     else 0.0
                 )
-            case ModelNotDownloading():
+            case DownloadPending():
                 total = progress.total.in_bytes
                 return progress.downloaded.in_bytes / total if total > 0 else 0.0
-            case ModelDownloadFailed():
-                return 0.0
-            case ModelRejected():
+            case DownloadFailed():
                 return 0.0
     return 0.0
 
@@ -91,7 +89,7 @@ def _get_node_download_fraction(
 def _cycle_download_score(
     cycle: Cycle,
     model_id: ModelId,
-    download_status: Mapping[NodeId, Sequence[ModelStatus]],
+    download_status: Mapping[NodeId, Sequence[DownloadProgress]],
 ) -> float:
     """Sum of download fractions across all nodes in a cycle."""
     return sum(
@@ -107,8 +105,7 @@ def place_instance(
     node_memory: Mapping[NodeId, MemoryUsage],
     node_network: Mapping[NodeId, NodeNetworkInfo],
     required_nodes: set[NodeId] | None = None,
-    download_status: Mapping[NodeId, Sequence[ModelStatus]] | None = None,
-    node_rdma_ctl: Mapping[NodeId, NodeRdmaCtlStatus] | None = None,
+    download_status: Mapping[NodeId, Sequence[DownloadProgress]] | None = None,
 ) -> dict[InstanceId, Instance]:
     cycles = topology.get_cycles()
     candidate_cycles = list(filter(lambda it: len(it) >= command.min_nodes, cycles))
@@ -170,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:
@@ -216,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,
@@ -280,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
 
@@ -337,14 +329,14 @@ def get_transition_events(
 
 def cancel_unnecessary_downloads(
     instances: Mapping[InstanceId, Instance],
-    download_status: Mapping[NodeId, Sequence[ModelStatus]],
+    download_status: Mapping[NodeId, Sequence[DownloadProgress]],
 ) -> Sequence[DownloadCommand]:
     commands: list[DownloadCommand] = []
     currently_downloading = [
         (k, v.shard_metadata.model_card.model_id)
         for k, vs in download_status.items()
         for v in vs
-        if isinstance(v, (ModelDownloading))
+        if isinstance(v, (DownloadOngoing))
     ]
     active_models = set(
         (

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,
@@ -34,10 +30,10 @@ from exo.shared.types.text_generation import (
 )
 from exo.shared.types.topology import Connection, SocketConnection
 from exo.shared.types.worker.downloads import (
+    DownloadCompleted,
+    DownloadFailed,
+    DownloadOngoing,
     DownloadProgressData,
-    ModelDownloadFailed,
-    ModelDownloading,
-    ModelReady,
 )
 from exo.shared.types.worker.instances import (
     Instance,
@@ -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,
@@ -778,7 +636,7 @@ def test_placement_prefers_cycle_with_downloaded_model(
     # node_b has the model fully downloaded, node_a does not
     download_status = {
         node_b: [
-            ModelReady(
+            DownloadCompleted(
                 node_id=node_b,
                 shard_metadata=shard_meta,
                 total=model_card.storage_size,
@@ -825,7 +683,7 @@ def test_placement_prefers_cycle_with_higher_download_progress(
     # node_a: 30% downloaded, node_b: 80% downloaded
     download_status = {
         node_a: [
-            ModelDownloading(
+            DownloadOngoing(
                 node_id=node_a,
                 shard_metadata=shard_meta,
                 download_progress=DownloadProgressData(
@@ -841,7 +699,7 @@ def test_placement_prefers_cycle_with_higher_download_progress(
             ),
         ],
         node_b: [
-            ModelDownloading(
+            DownloadOngoing(
                 node_id=node_b,
                 shard_metadata=shard_meta,
                 download_progress=DownloadProgressData(
@@ -898,7 +756,7 @@ def test_placement_does_not_prefer_cycle_with_failed_download(
     # node_b has a failed download — should not be preferred
     download_status = {
         node_b: [
-            ModelDownloadFailed(
+            DownloadFailed(
                 node_id=node_b,
                 shard_metadata=shard_meta,
                 error_message="connection reset",

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,
@@ -21,7 +20,6 @@ from exo.shared.types.events import (
     NodeGatheredInfo,
     NodeTimedOut,
     RunnerStatusUpdated,
-    StorageConfigUpdated,
     TaskAcknowledged,
     TaskCreated,
     TaskDeleted,
@@ -44,7 +42,7 @@ from exo.shared.types.profiling import (
 from exo.shared.types.state import State
 from exo.shared.types.tasks import Task, TaskId, TaskStatus
 from exo.shared.types.topology import Connection, RDMAConnection
-from exo.shared.types.worker.downloads import ModelStatus
+from exo.shared.types.worker.downloads import DownloadProgress
 from exo.shared.types.worker.instances import Instance, InstanceId
 from exo.shared.types.worker.runners import (
     RunnerId,
@@ -61,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:
@@ -89,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():
@@ -123,8 +109,6 @@ def event_apply(event: Event, state: State) -> State:
             return apply_instance_link_created(event, state)
         case InstanceLinkDeleted():
             return apply_instance_link_deleted(event, state)
-        case StorageConfigUpdated():
-            return apply_storage_config_updated(event, state)
 
 
 def apply(state: State, event: IndexedEvent) -> State:
@@ -144,13 +128,18 @@ def apply_node_download_progress(event: NodeDownloadProgress, state: State) -> S
     """
     dp = event.download_progress
     node_id = dp.node_id
-    model_id = dp.shard_metadata.model_card.model_id
 
     current = list(state.downloads.get(node_id, ()))
 
     replaced = False
     for i, existing_dp in enumerate(current):
-        if existing_dp.shard_metadata.model_card.model_id == model_id:
+        # TODO(ciaran): deduplicate by model_id for now. Will need to use
+        # shard_metadata again when pipeline and tensor downloads differ.
+        # For now this is fine
+        if (
+            existing_dp.shard_metadata.model_card.model_id
+            == dp.shard_metadata.model_card.model_id
+        ):
             current[i] = dp
             replaced = True
             break
@@ -158,7 +147,7 @@ def apply_node_download_progress(event: NodeDownloadProgress, state: State) -> S
     if not replaced:
         current.append(dp)
 
-    new_downloads: Mapping[NodeId, Sequence[ModelStatus]] = {
+    new_downloads: Mapping[NodeId, Sequence[DownloadProgress]] = {
         **state.downloads,
         node_id: current,
     }
@@ -317,10 +306,8 @@ 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_storage_config = {
-        key: value
-        for key, value in state.node_storage_config.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)
@@ -344,7 +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_storage_config": node_storage_config,
+            "node_vllm": node_vllm,
             "thunderbolt_bridge_cycles": thunderbolt_bridge_cycles,
         }
     )
@@ -371,15 +358,17 @@ 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():
             update["node_disk"] = {**state.node_disk, event.node_id: info.disk_usage}
         case NodeConfig():
-            update["node_storage_config"] = {
-                **state.node_storage_config,
-                event.node_id: info.storage_config,
-            }
+            pass
         case MiscData():
             current_identity = state.node_identities.get(event.node_id, NodeIdentity())
             new_identity = current_identity.model_copy(
@@ -419,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,
@@ -434,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():
@@ -461,24 +443,15 @@ 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)
 
 
-def apply_storage_config_updated(event: StorageConfigUpdated, state: State) -> State:
-    new_node_storage_config = {
-        **state.node_storage_config,
-        event.node_id: event.storage_config,
-    }
-    return state.model_copy(update={"node_storage_config": new_node_storage_config})
-
-
 def apply_topology_edge_created(event: TopologyEdgeCreated, state: State) -> State:
     topology = copy.deepcopy(state.topology)
     topology.add_connection(event.conn)
@@ -490,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

@@ -68,12 +68,7 @@ DASHBOARD_DIR = (
 # Log files (data/logs or cache)
 EXO_LOG_DIR = EXO_CACHE_HOME / "exo_log"
 EXO_LOG = EXO_LOG_DIR / "exo.log"
-EXO_RUNNER_LOG_DIR = EXO_LOG_DIR / "runner_log"
-EXO_RUNNER_STDOUT_LOG = EXO_RUNNER_LOG_DIR / "stdout.log"
-EXO_RUNNER_STDERR_LOG = EXO_RUNNER_LOG_DIR / "stderr.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"
@@ -99,8 +94,6 @@ EXO_ENABLE_IMAGE_MODELS = (
 
 EXO_OFFLINE = os.getenv("EXO_OFFLINE", "false").lower() == "true"
 
-EXO_MODEL_USAGE_FILE = EXO_DATA_HOME / "model_usage.json"
-
 EXO_TRACING_ENABLED = os.getenv("EXO_TRACING_ENABLED", "false").lower() == "true"
 
 ENABLE_DISAGGREGATION = os.getenv("ENABLE_DISAGGREGATION", "false").lower() == "true"

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

@@ -1,253 +0,0 @@
-import tomllib
-from collections.abc import Mapping, Sequence
-from datetime import UTC, datetime
-
-import anyio
-import tomlkit
-from loguru import logger
-from tomlkit.exceptions import TOMLKitError
-
-from exo.shared.constants import EXO_CONFIG_FILE
-from exo.shared.models.model_cards import ModelId
-from exo.shared.types.common import NodeId
-from exo.shared.types.events import InstanceDeleted, TaskStatusUpdated
-from exo.shared.types.memory import Memory
-from exo.shared.types.storage import (
-    StorageAllow,
-    StorageConfig,
-    StorageDecision,
-    StorageEvict,
-    StoragePolicy,
-    StorageReject,
-)
-from exo.shared.types.tasks import Task, TaskId, TaskStatus
-from exo.shared.types.worker.downloads import (
-    ModelDownloading,
-    ModelReady,
-    ModelStatus,
-)
-from exo.shared.types.worker.instances import Instance, InstanceId
-
-
-def calculate_used_storage(downloads: Sequence[ModelStatus]) -> Memory:
-    total = Memory()
-    for dp in downloads:
-        if isinstance(dp, ModelReady):
-            total = total + dp.total
-        elif isinstance(dp, ModelDownloading):
-            total = total + dp.download_progress.total
-    return total
-
-
-def check_storage_quota(
-    model_size: Memory,
-    config: StorageConfig,
-    downloads: Sequence[ModelStatus],
-) -> tuple[bool, str]:
-    if config.max_storage is None:
-        return True, ""
-
-    used = calculate_used_storage(downloads)
-    available = config.max_storage - used
-
-    if model_size <= available:
-        return True, ""
-
-    return (
-        False,
-        f"Need {model_size.in_gb:.1f} GiB, only {max(0, available.in_gb):.1f} GiB available within {config.max_storage.in_gb:.1f} GiB limit",
-    )
-
-
-def get_lru_eviction_candidates(
-    downloads: Sequence[ModelStatus],
-    model_last_used: Mapping[ModelId, datetime],
-    active_model_ids: frozenset[ModelId],
-) -> list[tuple[ModelId, ModelReady]]:
-    candidates: list[tuple[ModelId, ModelReady]] = []
-    for dp in downloads:
-        if not isinstance(dp, ModelReady):
-            continue
-        if dp.read_only:
-            continue
-        model_id = dp.shard_metadata.model_card.model_id
-        if model_id in active_model_ids:
-            continue
-        candidates.append((model_id, dp))
-
-    candidates.sort(
-        key=lambda item: model_last_used.get(item[0], datetime.min.replace(tzinfo=UTC))
-    )
-    return candidates
-
-
-def compute_evictions_needed(
-    model_size: Memory,
-    available: Memory,
-    candidates: list[tuple[ModelId, ModelReady]],
-) -> list[ModelId] | None:
-    if model_size <= available:
-        return []
-
-    space_needed = model_size - available
-    freed = Memory()
-    to_evict: list[ModelId] = []
-
-    for model_id, completed in candidates:
-        to_evict.append(model_id)
-        freed = freed + completed.total
-        if freed >= space_needed:
-            return to_evict
-
-    return None
-
-
-def decide_storage_action(
-    model_size: Memory,
-    config: StorageConfig,
-    downloads: Sequence[ModelStatus],
-    model_last_used: Mapping[ModelId, datetime],
-    active_model_ids: frozenset[ModelId],
-    disk_free: Memory | None = None,
-) -> StorageDecision:
-    """Pure decision function: given storage state, decide whether to allow, evict, or reject.
-
-    If ``disk_free`` is provided, it is used alongside the quota to determine
-    the effective available space. This ensures evictions are triggered when the
-    physical disk is full, even if the quota accounting says there is room.
-    """
-    if config.max_storage is None:
-        if disk_free is not None and model_size > disk_free:
-            # No quota set, but disk is physically full — try auto-evict if enabled
-            if config.storage_policy == "auto-evict":
-                candidates = get_lru_eviction_candidates(
-                    downloads, model_last_used, active_model_ids
-                )
-                to_evict = compute_evictions_needed(model_size, disk_free, candidates)
-                if to_evict is not None:
-                    return StorageEvict(model_ids=to_evict)
-            return StorageReject(
-                reason=f"Need {model_size.in_gb:.1f} GiB but only {disk_free.in_gb:.1f} GiB free on disk",
-                available=disk_free,
-            )
-        return StorageAllow()
-
-    used = calculate_used_storage(downloads)
-    raw_quota_available = config.max_storage - used
-    quota_available = (
-        raw_quota_available if raw_quota_available.in_bytes >= 0 else Memory()
-    )
-
-    # Effective available is the minimum of quota headroom and physical disk free space
-    available = quota_available
-    if disk_free is not None and disk_free < available:
-        available = disk_free
-
-    if model_size <= available:
-        return StorageAllow()
-
-    reason = (
-        f"Need {model_size.in_gb:.1f} GiB, only {max(0, available.in_gb):.1f} GiB available"
-        f" (quota: {quota_available.in_gb:.1f} GiB, disk: {disk_free.in_gb:.1f} GiB)"
-        if disk_free is not None
-        else f"Need {model_size.in_gb:.1f} GiB, only {max(0, available.in_gb):.1f} GiB available within {config.max_storage.in_gb:.1f} GiB limit"
-    )
-
-    if config.storage_policy == "auto-evict":
-        candidates = get_lru_eviction_candidates(
-            downloads, model_last_used, active_model_ids
-        )
-        to_evict = compute_evictions_needed(model_size, available, candidates)
-
-        if to_evict is not None:
-            return StorageEvict(model_ids=to_evict)
-
-        return StorageReject(
-            reason="Cannot free enough space even after evicting all eligible models",
-            available=available,
-        )
-
-    return StorageReject(
-        reason=reason,
-        available=available,
-    )
-
-
-def get_download_rejected_events(
-    rejected_model_id: ModelId,
-    rejected_node_id: NodeId,
-    instances: Mapping[InstanceId, Instance],
-    tasks: Mapping[TaskId, Task],
-) -> list[TaskStatusUpdated | InstanceDeleted]:
-    """Pure function: compute events needed to clean up after a download rejection."""
-    events: list[TaskStatusUpdated | InstanceDeleted] = []
-    for instance_id, instance in instances.items():
-        if (
-            instance.shard_assignments.model_id == rejected_model_id
-            and rejected_node_id in instance.shard_assignments.node_to_runner
-        ):
-            for task in tasks.values():
-                if task.instance_id == instance_id and task.task_status in (
-                    TaskStatus.Pending,
-                    TaskStatus.Running,
-                ):
-                    events.append(
-                        TaskStatusUpdated(
-                            task_id=task.task_id,
-                            task_status=TaskStatus.Failed,
-                        )
-                    )
-            events.append(InstanceDeleted(instance_id=instance_id))
-    return events
-
-
-async def load_storage_config(
-    *,
-    max_storage_gb: float | None = None,
-    storage_policy: StoragePolicy | None = None,
-) -> StorageConfig:
-    """Load StorageConfig from config.toml, overlaying any CLI arg overrides."""
-    base = StorageConfig()
-    cfg_file = anyio.Path(EXO_CONFIG_FILE)
-    try:
-        await cfg_file.parent.mkdir(parents=True, exist_ok=True)
-        await cfg_file.touch(exist_ok=True)
-        raw = (await cfg_file.read_bytes()).decode("utf-8")
-        if raw.strip():
-            data = tomllib.loads(raw)
-            base = StorageConfig.from_disk(data)
-    except (OSError, tomllib.TOMLDecodeError, ValueError, KeyError):
-        logger.warning("Failed to read storage config from config file, using defaults")
-
-    resolved_max_storage = (
-        Memory.from_gb(max_storage_gb)
-        if max_storage_gb is not None
-        else base.max_storage
-    )
-    resolved_policy = (
-        storage_policy if storage_policy is not None else base.storage_policy
-    )
-    return StorageConfig(
-        max_storage=resolved_max_storage, storage_policy=resolved_policy
-    )
-
-
-async def persist_storage_config(config: StorageConfig) -> None:
-    """Persist StorageConfig to config.toml, preserving other config keys."""
-    cfg_path = anyio.Path(EXO_CONFIG_FILE)
-    await cfg_path.parent.mkdir(parents=True, exist_ok=True)
-
-    doc = tomlkit.document()
-    try:
-        raw = (await cfg_path.read_bytes()).decode("utf-8")
-        if raw.strip():
-            doc = tomlkit.parse(raw)
-    except (FileNotFoundError, TOMLKitError, UnicodeDecodeError):
-        pass
-
-    # Clear max_storage_gb so it doesn't linger when max_storage is None
-    doc.pop("max_storage_gb", None)  # pyright: ignore[reportUnknownMemberType]
-    doc.update(config.to_disk())  # pyright: ignore[reportUnknownMemberType]
-
-    await cfg_path.write_text(tomlkit.dumps(doc))  # pyright: ignore[reportUnknownMemberType]
-    logger.debug(f"Persisted storage config to {cfg_path}")

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_node_download.py

@@ -4,14 +4,14 @@ from exo.shared.types.common import NodeId
 from exo.shared.types.events import NodeDownloadProgress
 from exo.shared.types.memory import Memory
 from exo.shared.types.state import State
-from exo.shared.types.worker.downloads import ModelReady
+from exo.shared.types.worker.downloads import DownloadCompleted
 from exo.worker.tests.constants import MODEL_A_ID, MODEL_B_ID
 
 
 def test_apply_node_download_progress():
     state = State()
     shard1 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=0, world_size=2)
-    event = ModelReady(
+    event = DownloadCompleted(
         node_id=NodeId("node-1"),
         shard_metadata=shard1,
         total=Memory(),
@@ -27,12 +27,12 @@ def test_apply_node_download_progress():
 def test_apply_two_node_download_progress():
     shard1 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=0, world_size=2)
     shard2 = get_pipeline_shard_metadata(MODEL_B_ID, device_rank=0, world_size=2)
-    event1 = ModelReady(
+    event1 = DownloadCompleted(
         node_id=NodeId("node-1"),
         shard_metadata=shard1,
         total=Memory(),
     )
-    event2 = ModelReady(
+    event2 = DownloadCompleted(
         node_id=NodeId("node-1"),
         shard_metadata=shard2,
         total=Memory(),

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/tests/test_apply/test_apply_storage.py

@@ -1,45 +0,0 @@
-from exo.shared.apply import apply_node_timed_out, apply_storage_config_updated
-from exo.shared.types.common import NodeId
-from exo.shared.types.events import NodeTimedOut, StorageConfigUpdated
-from exo.shared.types.memory import Memory
-from exo.shared.types.state import State
-from exo.shared.types.storage import StorageConfig
-
-NODE_A = NodeId("node-a")
-NODE_B = NodeId("node-b")
-
-
-def test_storage_config_updated_adds_config() -> None:
-    state = State()
-    config = StorageConfig(max_storage=Memory.from_gb(10), storage_policy="manual")
-    event = StorageConfigUpdated(node_id=NODE_A, storage_config=config)
-
-    new_state = apply_storage_config_updated(event, state)
-
-    assert NODE_A in new_state.node_storage_config
-    assert new_state.node_storage_config[NODE_A].max_storage == Memory.from_gb(10)
-    assert new_state.node_storage_config[NODE_A].storage_policy == "manual"
-
-
-def test_storage_config_updated_overwrites_existing() -> None:
-    config1 = StorageConfig(max_storage=Memory.from_gb(10), storage_policy="manual")
-    state = State(node_storage_config={NODE_A: config1})
-
-    config2 = StorageConfig(max_storage=Memory.from_gb(20), storage_policy="auto-evict")
-    event = StorageConfigUpdated(node_id=NODE_A, storage_config=config2)
-
-    new_state = apply_storage_config_updated(event, state)
-
-    assert new_state.node_storage_config[NODE_A].max_storage == Memory.from_gb(20)
-    assert new_state.node_storage_config[NODE_A].storage_policy == "auto-evict"
-
-
-def test_node_timed_out_cleans_up_storage_config() -> None:
-    config = StorageConfig(max_storage=Memory.from_gb(10))
-    state = State(node_storage_config={NODE_A: config, NODE_B: config})
-
-    event = NodeTimedOut(node_id=NODE_A)
-    new_state = apply_node_timed_out(event, state)
-
-    assert NODE_A not in new_state.node_storage_config
-    assert NODE_B in new_state.node_storage_config

src/exo/shared/tests/test_storage.py

@@ -1,634 +0,0 @@
-from datetime import UTC, datetime
-from pathlib import Path
-from unittest.mock import patch
-
-from exo.shared.models.model_cards import ModelId
-from exo.shared.storage import (
-    calculate_used_storage,
-    check_storage_quota,
-    compute_evictions_needed,
-    decide_storage_action,
-    get_download_rejected_events,
-    get_lru_eviction_candidates,
-    load_storage_config,
-    persist_storage_config,
-)
-from exo.shared.tests.conftest import get_pipeline_shard_metadata
-from exo.shared.types.common import NodeId
-from exo.shared.types.events import InstanceDeleted, TaskStatusUpdated
-from exo.shared.types.memory import Memory
-from exo.shared.types.storage import (
-    StorageAllow,
-    StorageConfig,
-    StorageEvict,
-    StorageReject,
-)
-from exo.shared.types.tasks import LoadModel, TaskId, TaskStatus
-from exo.shared.types.worker.downloads import (
-    DownloadProgressData,
-    ModelDownloading,
-    ModelNotDownloading,
-    ModelReady,
-)
-from exo.shared.types.worker.instances import InstanceId, MlxRingInstance
-from exo.shared.types.worker.runners import RunnerId, ShardAssignments
-
-MODEL_A = ModelId("org/model-a")
-MODEL_B = ModelId("org/model-b")
-MODEL_C = ModelId("org/model-c")
-MODEL_D = ModelId("org/model-d")
-NODE_ID = "node-1"
-
-
-def _completed(
-    model_id: ModelId, size_gb: float, read_only: bool = False
-) -> ModelReady:
-    shard = get_pipeline_shard_metadata(model_id, device_rank=0)
-    return ModelReady(
-        node_id=NODE_ID,  # type: ignore[arg-type]
-        shard_metadata=shard,
-        total=Memory.from_gb(size_gb),
-        read_only=read_only,
-    )
-
-
-class TestCheckStorageQuota:
-    def test_unlimited_allows(self) -> None:
-        config = StorageConfig(max_storage=None)
-        allowed, _reason = check_storage_quota(Memory.from_gb(10), config, [])
-        assert allowed is True
-        assert _reason == ""
-
-    def test_under_limit_allows(self) -> None:
-        config = StorageConfig(max_storage=Memory.from_gb(20))
-        downloads = [_completed(MODEL_A, 5)]
-        allowed, _reason = check_storage_quota(Memory.from_gb(10), config, downloads)
-        assert allowed is True
-
-    def test_over_limit_rejects(self) -> None:
-        config = StorageConfig(max_storage=Memory.from_gb(10))
-        downloads = [_completed(MODEL_A, 5)]
-        allowed, reason = check_storage_quota(Memory.from_gb(8), config, downloads)
-        assert allowed is False
-        assert "Need" in reason
-        assert "available" in reason
-
-    def test_exact_fit_allows(self) -> None:
-        config = StorageConfig(max_storage=Memory.from_gb(10))
-        downloads = [_completed(MODEL_A, 5)]
-        allowed, _ = check_storage_quota(Memory.from_gb(5), config, downloads)
-        assert allowed is True
-
-
-class TestGetLruEvictionCandidates:
-    def test_excludes_active_models(self) -> None:
-        downloads = [_completed(MODEL_A, 5), _completed(MODEL_B, 3)]
-        last_used = {
-            MODEL_A: datetime(2024, 1, 1, tzinfo=UTC),
-            MODEL_B: datetime(2024, 1, 2, tzinfo=UTC),
-        }
-        candidates = get_lru_eviction_candidates(
-            downloads, last_used, frozenset({MODEL_A})
-        )
-        assert len(candidates) == 1
-        assert candidates[0][0] == MODEL_B
-
-    def test_excludes_read_only(self) -> None:
-        downloads = [_completed(MODEL_A, 5, read_only=True), _completed(MODEL_B, 3)]
-        candidates = get_lru_eviction_candidates(downloads, {}, frozenset())
-        assert len(candidates) == 1
-        assert candidates[0][0] == MODEL_B
-
-    def test_sorts_oldest_first(self) -> None:
-        downloads = [_completed(MODEL_A, 5), _completed(MODEL_B, 3)]
-        last_used = {
-            MODEL_A: datetime(2024, 6, 1, tzinfo=UTC),
-            MODEL_B: datetime(2024, 1, 1, tzinfo=UTC),
-        }
-        candidates = get_lru_eviction_candidates(downloads, last_used, frozenset())
-        assert candidates[0][0] == MODEL_B
-        assert candidates[1][0] == MODEL_A
-
-    def test_models_without_usage_get_min(self) -> None:
-        downloads = [_completed(MODEL_A, 5), _completed(MODEL_B, 3)]
-        last_used = {MODEL_A: datetime(2024, 6, 1, tzinfo=UTC)}
-        candidates = get_lru_eviction_candidates(downloads, last_used, frozenset())
-        assert candidates[0][0] == MODEL_B  # no usage -> datetime.min
-
-
-class TestComputeEvictionsNeeded:
-    def test_sufficient_candidates(self) -> None:
-        candidates = [
-            (MODEL_A, _completed(MODEL_A, 5)),
-            (MODEL_B, _completed(MODEL_B, 3)),
-        ]
-        result = compute_evictions_needed(
-            Memory.from_gb(6), Memory.from_gb(2), candidates
-        )
-        assert result is not None
-        assert MODEL_A in result
-
-    def test_insufficient_candidates(self) -> None:
-        candidates = [
-            (MODEL_A, _completed(MODEL_A, 2)),
-        ]
-        result = compute_evictions_needed(
-            Memory.from_gb(10), Memory.from_gb(2), candidates
-        )
-        assert result is None
-
-    def test_no_eviction_needed(self) -> None:
-        candidates = [(MODEL_A, _completed(MODEL_A, 5))]
-        result = compute_evictions_needed(
-            Memory.from_gb(3), Memory.from_gb(5), candidates
-        )
-        assert result == []
-
-    def test_evicts_in_lru_order(self) -> None:
-        """Eviction picks candidates in the order given (oldest first from LRU sort)."""
-        candidates = [
-            (MODEL_A, _completed(MODEL_A, 2)),  # oldest
-            (MODEL_B, _completed(MODEL_B, 2)),  # newer
-            (MODEL_C, _completed(MODEL_C, 2)),  # newest
-        ]
-        result = compute_evictions_needed(
-            Memory.from_gb(3), Memory.from_gb(1), candidates
-        )
-        assert result == [MODEL_A]
-
-    def test_evicts_multiple_until_enough_space(self) -> None:
-        """When one model isn't enough, evicts multiple in LRU order."""
-        candidates = [
-            (MODEL_A, _completed(MODEL_A, 1)),
-            (MODEL_B, _completed(MODEL_B, 1)),
-            (MODEL_C, _completed(MODEL_C, 1)),
-        ]
-        # Need 4 GiB, have 1 GiB available — need 3 GiB freed
-        result = compute_evictions_needed(
-            Memory.from_gb(4), Memory.from_gb(1), candidates
-        )
-        assert result == [MODEL_A, MODEL_B, MODEL_C]
-
-    def test_evicts_minimum_needed(self) -> None:
-        """Stops evicting as soon as enough space is freed."""
-        candidates = [
-            (MODEL_A, _completed(MODEL_A, 3)),
-            (MODEL_B, _completed(MODEL_B, 3)),
-        ]
-        # Need 5 GiB, have 2 GiB — need 3 GiB freed. Model A alone suffices.
-        result = compute_evictions_needed(
-            Memory.from_gb(5), Memory.from_gb(2), candidates
-        )
-        assert result == [MODEL_A]
-
-
-class TestCalculateUsedStorage:
-    def test_only_counts_completed_and_ongoing(self) -> None:
-        """Pending and rejected downloads should not count toward used storage."""
-        shard_a = get_pipeline_shard_metadata(MODEL_A, device_rank=0)
-        shard_b = get_pipeline_shard_metadata(MODEL_B, device_rank=0)
-        downloads = [
-            _completed(MODEL_A, 5),
-            ModelNotDownloading(
-                node_id=NODE_ID,  # type: ignore[arg-type]
-                shard_metadata=shard_a,
-            ),
-            ModelDownloading(
-                node_id=NODE_ID,  # type: ignore[arg-type]
-                shard_metadata=shard_b,
-                download_progress=DownloadProgressData(
-                    total=Memory.from_gb(10),
-                    downloaded=Memory.from_gb(3),
-                    downloaded_this_session=Memory.from_gb(3),
-                    completed_files=1,
-                    total_files=5,
-                    speed=0,
-                    eta_ms=0,
-                    files={},
-                ),
-            ),
-        ]
-        used = calculate_used_storage(downloads)
-        # 5 GiB completed + 10 GiB ongoing total = 15 GiB
-        assert abs(used.in_gb - 15.0) < 0.01
-
-    def test_empty_downloads(self) -> None:
-        assert calculate_used_storage([]).in_bytes == 0
-
-
-class TestGetLruEvictionCandidatesExtended:
-    def test_excludes_non_completed_downloads(self) -> None:
-        """Only DownloadCompleted entries are eviction candidates."""
-        shard_a = get_pipeline_shard_metadata(MODEL_A, device_rank=0)
-        downloads = [
-            _completed(MODEL_B, 3),
-            ModelNotDownloading(
-                node_id=NODE_ID,  # type: ignore[arg-type]
-                shard_metadata=shard_a,
-            ),
-        ]
-        candidates = get_lru_eviction_candidates(downloads, {}, frozenset())
-        assert len(candidates) == 1
-        assert candidates[0][0] == MODEL_B
-
-    def test_all_active_returns_empty(self) -> None:
-        """When all completed models are active, no candidates available."""
-        downloads = [_completed(MODEL_A, 5), _completed(MODEL_B, 3)]
-        candidates = get_lru_eviction_candidates(
-            downloads, {}, frozenset({MODEL_A, MODEL_B})
-        )
-        assert candidates == []
-
-    def test_three_models_lru_order(self) -> None:
-        """Three models sorted correctly: oldest used first."""
-        downloads = [
-            _completed(MODEL_A, 2),
-            _completed(MODEL_B, 3),
-            _completed(MODEL_C, 1),
-        ]
-        last_used = {
-            MODEL_A: datetime(2024, 3, 1, tzinfo=UTC),
-            MODEL_B: datetime(2024, 1, 1, tzinfo=UTC),
-            MODEL_C: datetime(2024, 6, 1, tzinfo=UTC),
-        }
-        candidates = get_lru_eviction_candidates(downloads, last_used, frozenset())
-        assert [c[0] for c in candidates] == [MODEL_B, MODEL_A, MODEL_C]
-
-    def test_mixed_active_readonly_and_regular(self) -> None:
-        """Only non-active, non-read-only completed models are candidates."""
-        downloads = [
-            _completed(MODEL_A, 5, read_only=True),  # excluded: read-only
-            _completed(MODEL_B, 3),  # excluded: active
-            _completed(MODEL_C, 2),  # candidate
-            _completed(MODEL_D, 1),  # candidate
-        ]
-        last_used = {
-            MODEL_C: datetime(2024, 6, 1, tzinfo=UTC),
-            MODEL_D: datetime(2024, 1, 1, tzinfo=UTC),
-        }
-        candidates = get_lru_eviction_candidates(
-            downloads, last_used, frozenset({MODEL_B})
-        )
-        assert [c[0] for c in candidates] == [MODEL_D, MODEL_C]
-
-
-class TestEndToEndEvictionScenario:
-    """Tests that combine LRU candidate selection with eviction computation."""
-
-    def test_evicts_oldest_model_to_fit_new_one(self) -> None:
-        """10 GiB limit, 3 completed models totaling 9 GiB,
-        need 3 GiB for new model — should evict the oldest."""
-        downloads = [
-            _completed(MODEL_A, 3),  # oldest used
-            _completed(MODEL_B, 3),
-            _completed(MODEL_C, 3),  # newest used
-        ]
-        last_used = {
-            MODEL_A: datetime(2024, 1, 1, tzinfo=UTC),
-            MODEL_B: datetime(2024, 6, 1, tzinfo=UTC),
-            MODEL_C: datetime(2024, 12, 1, tzinfo=UTC),
-        }
-        config = StorageConfig(max_storage=Memory.from_gb(10))
-        new_model_size = Memory.from_gb(3)
-
-        # Step 1: quota check fails
-        allowed, _ = check_storage_quota(new_model_size, config, downloads)
-        assert not allowed
-
-        # Step 2: get candidates in LRU order
-        candidates = get_lru_eviction_candidates(downloads, last_used, frozenset())
-        assert candidates[0][0] == MODEL_A  # oldest
-
-        # Step 3: compute what to evict
-        used = calculate_used_storage(downloads)
-        assert config.max_storage is not None
-        available = config.max_storage - used
-        to_evict = compute_evictions_needed(new_model_size, available, candidates)
-        assert to_evict == [MODEL_A]
-
-    def test_protects_currently_active_model(self) -> None:
-        """Active model should not be evicted even if it's the oldest."""
-        downloads = [
-            _completed(MODEL_A, 4),  # oldest but active
-            _completed(MODEL_B, 4),  # next oldest, evictable
-        ]
-        last_used = {
-            MODEL_A: datetime(2024, 1, 1, tzinfo=UTC),
-            MODEL_B: datetime(2024, 6, 1, tzinfo=UTC),
-        }
-        config = StorageConfig(max_storage=Memory.from_gb(10))
-        new_model_size = Memory.from_gb(5)
-
-        allowed, _ = check_storage_quota(new_model_size, config, downloads)
-        assert not allowed
-
-        # MODEL_A is active — should be excluded
-        candidates = get_lru_eviction_candidates(
-            downloads, last_used, frozenset({MODEL_A})
-        )
-        assert len(candidates) == 1
-        assert candidates[0][0] == MODEL_B
-
-        used = calculate_used_storage(downloads)
-        assert config.max_storage is not None
-        available = config.max_storage - used
-        to_evict = compute_evictions_needed(new_model_size, available, candidates)
-        assert to_evict == [MODEL_B]
-
-    def test_cannot_evict_enough_returns_none(self) -> None:
-        """When all evictable space isn't enough, returns None."""
-        downloads = [
-            _completed(MODEL_A, 4, read_only=True),  # can't evict
-            _completed(MODEL_B, 2),  # can evict but only 2 GiB
-        ]
-        config = StorageConfig(max_storage=Memory.from_gb(10))
-        new_model_size = Memory.from_gb(8)
-
-        candidates = get_lru_eviction_candidates(downloads, {}, frozenset())
-        used = calculate_used_storage(downloads)
-        assert config.max_storage is not None
-        available = config.max_storage - used
-        to_evict = compute_evictions_needed(new_model_size, available, candidates)
-        assert to_evict is None
-
-
-class TestDecideStorageAction:
-    """Tests for the decide_storage_action pure function."""
-
-    def test_unlimited_allows(self) -> None:
-        config = StorageConfig(max_storage=None)
-        action = decide_storage_action(Memory.from_gb(10), config, [], {}, frozenset())
-        assert isinstance(action, StorageAllow)
-
-    def test_under_limit_allows(self) -> None:
-        config = StorageConfig(max_storage=Memory.from_gb(20))
-        downloads = [_completed(MODEL_A, 5)]
-        action = decide_storage_action(
-            Memory.from_gb(10), config, downloads, {}, frozenset()
-        )
-        assert isinstance(action, StorageAllow)
-
-    def test_manual_policy_rejects(self) -> None:
-        config = StorageConfig(max_storage=Memory.from_gb(10), storage_policy="manual")
-        downloads = [_completed(MODEL_A, 5)]
-        action = decide_storage_action(
-            Memory.from_gb(8), config, downloads, {}, frozenset()
-        )
-        assert isinstance(action, StorageReject)
-        assert "Need" in action.reason
-
-    def test_auto_evict_returns_evict(self) -> None:
-        config = StorageConfig(
-            max_storage=Memory.from_gb(10), storage_policy="auto-evict"
-        )
-        downloads = [_completed(MODEL_A, 4), _completed(MODEL_B, 4)]
-        last_used = {
-            MODEL_A: datetime(2024, 1, 1, tzinfo=UTC),
-            MODEL_B: datetime(2024, 6, 1, tzinfo=UTC),
-        }
-        action = decide_storage_action(
-            Memory.from_gb(5), config, downloads, last_used, frozenset()
-        )
-        assert isinstance(action, StorageEvict)
-        assert MODEL_A in action.model_ids
-
-    def test_auto_evict_rejects_when_impossible(self) -> None:
-        config = StorageConfig(
-            max_storage=Memory.from_gb(10), storage_policy="auto-evict"
-        )
-        downloads = [_completed(MODEL_A, 2)]
-        action = decide_storage_action(
-            Memory.from_gb(20), config, downloads, {}, frozenset()
-        )
-        assert isinstance(action, StorageReject)
-        assert "Cannot free enough" in action.reason
-
-    def test_auto_evict_protects_active_models(self) -> None:
-        config = StorageConfig(
-            max_storage=Memory.from_gb(10), storage_policy="auto-evict"
-        )
-        downloads = [_completed(MODEL_A, 4), _completed(MODEL_B, 4)]
-        last_used = {
-            MODEL_A: datetime(2024, 1, 1, tzinfo=UTC),
-            MODEL_B: datetime(2024, 6, 1, tzinfo=UTC),
-        }
-        # MODEL_A is oldest but active — should evict MODEL_B instead
-        action = decide_storage_action(
-            Memory.from_gb(5), config, downloads, last_used, frozenset({MODEL_A})
-        )
-        assert isinstance(action, StorageEvict)
-        assert action.model_ids == [MODEL_B]
-
-    def test_auto_evict_all_active_rejects(self) -> None:
-        config = StorageConfig(
-            max_storage=Memory.from_gb(10), storage_policy="auto-evict"
-        )
-        downloads = [_completed(MODEL_A, 4), _completed(MODEL_B, 4)]
-        action = decide_storage_action(
-            Memory.from_gb(5),
-            config,
-            downloads,
-            {},
-            frozenset({MODEL_A, MODEL_B}),
-        )
-        assert isinstance(action, StorageReject)
-
-
-def _make_instance(
-    instance_id: InstanceId,
-    model_id: ModelId,
-    node_id: NodeId,
-) -> MlxRingInstance:
-    shard = get_pipeline_shard_metadata(model_id, device_rank=0)
-    runner_id = RunnerId()
-    return MlxRingInstance(
-        instance_id=instance_id,
-        shard_assignments=ShardAssignments(
-            model_id=model_id,
-            runner_to_shard={runner_id: shard},
-            node_to_runner={node_id: runner_id},
-        ),
-        hosts_by_node={node_id: []},
-        ephemeral_port=0,
-    )
-
-
-class TestGetDownloadRejectedEvents:
-    """Tests for the get_download_rejected_events pure function."""
-
-    def test_deletes_instance_for_rejected_model(self) -> None:
-        node_id = NodeId("node-1")
-        instance_id = InstanceId()
-        instance = _make_instance(instance_id, MODEL_A, node_id)
-        events = get_download_rejected_events(
-            MODEL_A, node_id, {instance_id: instance}, {}
-        )
-        assert len(events) == 1
-        assert isinstance(events[0], InstanceDeleted)
-        assert events[0].instance_id == instance_id
-
-    def test_fails_pending_tasks_before_deleting_instance(self) -> None:
-        node_id = NodeId("node-1")
-        instance_id = InstanceId()
-        instance = _make_instance(instance_id, MODEL_A, node_id)
-        task_id = TaskId()
-        task = LoadModel(
-            task_id=task_id,
-            instance_id=instance_id,
-            task_status=TaskStatus.Pending,
-        )
-        events = get_download_rejected_events(
-            MODEL_A, node_id, {instance_id: instance}, {task_id: task}
-        )
-        assert len(events) == 2
-        assert isinstance(events[0], TaskStatusUpdated)
-        assert events[0].task_status == TaskStatus.Failed
-        assert isinstance(events[1], InstanceDeleted)
-
-    def test_ignores_different_model(self) -> None:
-        node_id = NodeId("node-1")
-        instance_id = InstanceId()
-        instance = _make_instance(instance_id, MODEL_B, node_id)
-        events = get_download_rejected_events(
-            MODEL_A, node_id, {instance_id: instance}, {}
-        )
-        assert events == []
-
-    def test_ignores_different_node(self) -> None:
-        node_id = NodeId("node-1")
-        other_node = NodeId("node-2")
-        instance_id = InstanceId()
-        instance = _make_instance(instance_id, MODEL_A, other_node)
-        events = get_download_rejected_events(
-            MODEL_A, node_id, {instance_id: instance}, {}
-        )
-        assert events == []
-
-    def test_skips_completed_tasks(self) -> None:
-        node_id = NodeId("node-1")
-        instance_id = InstanceId()
-        instance = _make_instance(instance_id, MODEL_A, node_id)
-        task_id = TaskId()
-        task = LoadModel(
-            task_id=task_id,
-            instance_id=instance_id,
-            task_status=TaskStatus.Complete,
-        )
-        events = get_download_rejected_events(
-            MODEL_A, node_id, {instance_id: instance}, {task_id: task}
-        )
-        # Only InstanceDeleted, no TaskStatusUpdated for completed task
-        assert len(events) == 1
-        assert isinstance(events[0], InstanceDeleted)
-
-
-class TestPersistStorageConfig:
-    """Tests for persist_storage_config I/O."""
-
-    async def test_round_trip(self, tmp_path: Path) -> None:
-        cfg_file = tmp_path / "config.toml"
-        config = StorageConfig(
-            max_storage=Memory.from_gb(50), storage_policy="auto-evict"
-        )
-        with patch("exo.shared.storage.EXO_CONFIG_FILE", cfg_file):
-            await persist_storage_config(config)
-            loaded = await load_storage_config()
-
-        assert loaded.storage_policy == "auto-evict"
-        assert loaded.max_storage is not None
-        assert abs(loaded.max_storage.in_gb - 50.0) < 0.1
-
-    async def test_preserves_other_keys(self, tmp_path: Path) -> None:
-        cfg_file = tmp_path / "config.toml"
-        cfg_file.write_text('some_other_key = "hello"\n')
-        config = StorageConfig(max_storage=Memory.from_gb(10))
-        with patch("exo.shared.storage.EXO_CONFIG_FILE", cfg_file):
-            await persist_storage_config(config)
-
-        contents = cfg_file.read_text()
-        assert "some_other_key" in contents
-        assert "hello" in contents
-        assert "max_storage_gb" in contents
-
-    async def test_clears_max_storage_gb_when_unlimited(self, tmp_path: Path) -> None:
-        cfg_file = tmp_path / "config.toml"
-        cfg_file.write_text('max_storage_gb = 50\nstorage_policy = "auto-evict"\n')
-        config = StorageConfig(max_storage=None, storage_policy="manual")
-        with patch("exo.shared.storage.EXO_CONFIG_FILE", cfg_file):
-            await persist_storage_config(config)
-
-        contents = cfg_file.read_text()
-        assert "max_storage_gb" not in contents
-        assert "manual" in contents
-
-    async def test_creates_file_if_missing(self, tmp_path: Path) -> None:
-        cfg_file = tmp_path / "subdir" / "config.toml"
-        config = StorageConfig(max_storage=Memory.from_gb(25))
-        with patch("exo.shared.storage.EXO_CONFIG_FILE", cfg_file):
-            await persist_storage_config(config)
-
-        assert cfg_file.exists()
-        assert "max_storage_gb" in cfg_file.read_text()
-
-
-class TestLoadStorageConfig:
-    """Tests for load_storage_config I/O."""
-
-    async def test_defaults_when_empty_file(self, tmp_path: Path) -> None:
-        cfg_file = tmp_path / "config.toml"
-        cfg_file.write_text("")
-        with patch("exo.shared.storage.EXO_CONFIG_FILE", cfg_file):
-            config = await load_storage_config()
-
-        assert config.max_storage is None
-        assert config.storage_policy == "manual"
-
-    async def test_reads_from_file(self, tmp_path: Path) -> None:
-        cfg_file = tmp_path / "config.toml"
-        cfg_file.write_text('max_storage_gb = 30.0\nstorage_policy = "auto-evict"\n')
-        with patch("exo.shared.storage.EXO_CONFIG_FILE", cfg_file):
-            config = await load_storage_config()
-
-        assert config.storage_policy == "auto-evict"
-        assert config.max_storage is not None
-        assert abs(config.max_storage.in_gb - 30.0) < 0.1
-
-    async def test_cli_overrides_file(self, tmp_path: Path) -> None:
-        cfg_file = tmp_path / "config.toml"
-        cfg_file.write_text('max_storage_gb = 30.0\nstorage_policy = "manual"\n')
-        with patch("exo.shared.storage.EXO_CONFIG_FILE", cfg_file):
-            config = await load_storage_config(
-                max_storage_gb=100.0, storage_policy="auto-evict"
-            )
-
-        assert config.storage_policy == "auto-evict"
-        assert config.max_storage is not None
-        assert abs(config.max_storage.in_gb - 100.0) < 0.1
-
-    async def test_partial_cli_override(self, tmp_path: Path) -> None:
-        """CLI overrides only the fields provided, file values used for the rest."""
-        cfg_file = tmp_path / "config.toml"
-        cfg_file.write_text('max_storage_gb = 30.0\nstorage_policy = "auto-evict"\n')
-        with patch("exo.shared.storage.EXO_CONFIG_FILE", cfg_file):
-            config = await load_storage_config(max_storage_gb=50.0)
-
-        assert config.storage_policy == "auto-evict"  # from file
-        assert config.max_storage is not None
-        assert abs(config.max_storage.in_gb - 50.0) < 0.1  # from CLI
-
-    async def test_defaults_on_corrupt_file(self, tmp_path: Path) -> None:
-        cfg_file = tmp_path / "config.toml"
-        cfg_file.write_text("not valid toml {{{")
-        with patch("exo.shared.storage.EXO_CONFIG_FILE", cfg_file):
-            config = await load_storage_config()
-
-        assert config.max_storage is None
-        assert config.storage_policy == "manual"
-
-    async def test_creates_file_if_missing(self, tmp_path: Path) -> None:
-        cfg_file = tmp_path / "subdir" / "config.toml"
-        with patch("exo.shared.storage.EXO_CONFIG_FILE", cfg_file):
-            config = await load_storage_config()
-
-        assert config.max_storage is None
-        assert cfg_file.exists()

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

@@ -8,8 +8,6 @@ from exo.shared.models.model_cards import ModelCard, ModelId
 from exo.shared.types.chunks import InputImageChunk
 from exo.shared.types.common import CommandId, NodeId, SystemId
 from exo.shared.types.instance_link import InstanceLinkId
-from exo.shared.types.memory import Memory
-from exo.shared.types.storage import StoragePolicy
 from exo.shared.types.text_generation import TextGenerationTaskParams
 from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
 from exo.shared.types.worker.shards import Sharding, ShardMetadata
@@ -105,12 +103,6 @@ class DeleteInstanceLink(BaseCommand):
 DownloadCommand = StartDownload | DeleteDownload | CancelDownload
 
 
-class SetStorageConfig(BaseCommand):
-    target_node_id: NodeId
-    max_storage: Memory | None
-    storage_policy: StoragePolicy
-
-
 Command = (
     TestCommand
     | RequestEventLog
@@ -127,7 +119,6 @@ Command = (
     | DeleteCustomModelCard
     | SetInstanceLink
     | DeleteInstanceLink
-    | SetStorageConfig
 )
 
 

src/exo/shared/types/events.py

@@ -8,9 +8,8 @@ from exo.shared.topology import Connection
 from exo.shared.types.chunks import Chunk, InputImageChunk
 from exo.shared.types.common import CommandId, Id, ModelId, NodeId, SessionId, SystemId
 from exo.shared.types.instance_link import InstanceLink, InstanceLinkId
-from exo.shared.types.storage import StorageConfig
 from exo.shared.types.tasks import Task, TaskId, TaskStatus
-from exo.shared.types.worker.downloads import ModelStatus
+from exo.shared.types.worker.downloads import DownloadProgress
 from exo.shared.types.worker.instances import Instance, InstanceId
 from exo.shared.types.worker.runners import RunnerId, RunnerStatus
 from exo.utils.info_gatherer.info_gatherer import GatheredInfo
@@ -88,7 +87,7 @@ class NodeGatheredInfo(BaseEvent):
 
 
 class NodeDownloadProgress(BaseEvent):
-    download_progress: ModelStatus
+    download_progress: DownloadProgress
 
 
 class ChunkGenerated(BaseEvent):
@@ -147,12 +146,6 @@ class InstanceLinkDeleted(BaseEvent):
     link_id: InstanceLinkId
 
 
-@final
-class StorageConfigUpdated(BaseEvent):
-    node_id: NodeId
-    storage_config: StorageConfig
-
-
 Event = (
     TestEvent
     | TaskCreated
@@ -176,7 +169,6 @@ Event = (
     | CustomModelCardDeleted
     | InstanceLinkCreated
     | InstanceLinkDeleted
-    | StorageConfigUpdated
 )
 
 

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,
@@ -19,9 +18,8 @@ from exo.shared.types.profiling import (
     SystemPerformanceProfile,
     ThunderboltBridgeStatus,
 )
-from exo.shared.types.storage import StorageConfig
 from exo.shared.types.tasks import Task, TaskId
-from exo.shared.types.worker.downloads import ModelStatus
+from exo.shared.types.worker.downloads import DownloadProgress
 from exo.shared.types.worker.instances import Instance, InstanceId
 from exo.shared.types.worker.runners import RunnerId, RunnerStatus
 from exo.utils.pydantic_ext import FrozenModel
@@ -45,7 +43,7 @@ class State(FrozenModel):
     )
     instances: Mapping[InstanceId, Instance] = {}
     runners: Mapping[RunnerId, RunnerStatus] = {}
-    downloads: Mapping[NodeId, Sequence[ModelStatus]] = {}
+    downloads: Mapping[NodeId, Sequence[DownloadProgress]] = {}
     tasks: Mapping[TaskId, Task] = {}
     last_seen: Mapping[NodeId, datetime] = {}
     topology: Topology = Field(default_factory=Topology)
@@ -60,7 +58,7 @@ class State(FrozenModel):
     node_thunderbolt: Mapping[NodeId, NodeThunderboltInfo] = {}
     node_thunderbolt_bridge: Mapping[NodeId, ThunderboltBridgeStatus] = {}
     node_rdma_ctl: Mapping[NodeId, NodeRdmaCtlStatus] = {}
-    node_storage_config: Mapping[NodeId, StorageConfig] = {}
+    node_vllm: Mapping[NodeId, bool] = {}
 
     # Detected cycles where all nodes have Thunderbolt bridge enabled (>2 nodes)
     thunderbolt_bridge_cycles: Sequence[Sequence[NodeId]] = []
@@ -68,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/storage.py

@@ -1,52 +0,0 @@
-from typing import Any, Literal, Self, final
-
-from exo.shared.models.model_cards import ModelId
-from exo.shared.types.memory import Memory
-from exo.utils.pydantic_ext import FrozenModel
-
-StoragePolicy = Literal["manual", "auto-evict"]
-
-
-@final
-class StorageConfig(FrozenModel):
-    max_storage: Memory | None = None
-    storage_policy: StoragePolicy = "manual"
-
-    @classmethod
-    def from_disk(cls, data: dict[str, Any]) -> Self:
-        """Parse from a TOML config dict (e.g. from tomllib)."""
-        max_storage: Memory | None = None
-        if "max_storage_gb" in data:
-            gb = float(data["max_storage_gb"])  # pyright: ignore[reportAny]
-            if gb < 0:
-                raise ValueError(f"max_storage_gb must be non-negative, got {gb}")
-            max_storage = Memory.from_gb(gb)
-        policy: StoragePolicy = data.get("storage_policy", "manual")  # pyright: ignore[reportAny]
-        return cls(max_storage=max_storage, storage_policy=policy)
-
-    def to_disk(self) -> dict[str, Any]:
-        """Serialize to a dict suitable for writing to TOML."""
-        result: dict[str, Any] = {}
-        if self.max_storage is not None:
-            result["max_storage_gb"] = round(self.max_storage.in_gb, 2)
-        result["storage_policy"] = self.storage_policy
-        return result
-
-
-@final
-class StorageAllow(FrozenModel):
-    pass
-
-
-@final
-class StorageEvict(FrozenModel):
-    model_ids: list[ModelId]
-
-
-@final
-class StorageReject(FrozenModel):
-    reason: str
-    available: Memory
-
-
-StorageDecision = StorageAllow | StorageEvict | StorageReject

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

@@ -23,43 +23,32 @@ class DownloadProgressData(FrozenModel):
     files: dict[str, "DownloadProgressData"]
 
 
-class BaseModelStatus(TaggedModel):
+class BaseDownloadProgress(TaggedModel):
     node_id: NodeId
     shard_metadata: ShardMetadata
     model_directory: str = ""
 
 
-class ModelNotDownloading(BaseModelStatus):
+class DownloadPending(BaseDownloadProgress):
     downloaded: Memory = Memory()
     total: Memory = Memory()
 
 
-class ModelReady(BaseModelStatus):
+class DownloadCompleted(BaseDownloadProgress):
     total: Memory
     read_only: bool = False
 
 
-class ModelDownloadFailed(BaseModelStatus):
+class DownloadFailed(BaseDownloadProgress):
     error_message: str
 
 
-class ModelDownloading(BaseModelStatus):
+class DownloadOngoing(BaseDownloadProgress):
     download_progress: DownloadProgressData
 
 
-class ModelRejected(BaseModelStatus):
-    reason: str
-    required: Memory
-    available: Memory
-    limit: Memory | None = None
-
-
-ModelStatus = (
-    ModelNotDownloading
-    | ModelReady
-    | ModelDownloadFailed
-    | ModelDownloading
-    | ModelRejected
+DownloadProgress = (
+    DownloadPending | DownloadCompleted | DownloadFailed | DownloadOngoing
 )
 
 

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

@@ -21,7 +21,6 @@ from exo.shared.types.profiling import (
     NetworkInterfaceInfo,
     ThunderboltBridgeStatus,
 )
-from exo.shared.types.storage import StorageConfig
 from exo.shared.types.thunderbolt import (
     ThunderboltConnection,
     ThunderboltConnectivity,
@@ -32,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,
@@ -295,8 +295,6 @@ class ThunderboltBridgeInfo(TaggedModel):
 class NodeConfig(TaggedModel):
     """Node configuration from EXO_CONFIG_FILE, reloaded from the file only at startup. Other changes should come in through the API and propagate from there"""
 
-    storage_config: StorageConfig = StorageConfig()
-
     @classmethod
     async def gather(cls) -> Self | None:
         cfg_file = anyio.Path(EXO_CONFIG_FILE)
@@ -306,7 +304,7 @@ class NodeConfig(TaggedModel):
             try:
                 contents = (await f.read()).decode("utf-8")
                 data = tomllib.loads(contents)
-                return cls(storage_config=StorageConfig.from_disk(data))
+                return cls.model_validate(data)
             except (tomllib.TOMLDecodeError, UnicodeDecodeError, ValidationError):
                 logger.warning("Invalid config file, skipping...")
                 return None
@@ -356,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
@@ -364,6 +380,8 @@ GatheredInfo = (
     | MacThunderboltConnections
     | RdmaCtlStatus
     | ThunderboltBridgeInfo
+    | NvmlMetrics
+    | VllmCapability
     | NodeConfig
     | MiscData
     | StaticNodeInformation
@@ -422,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)
@@ -430,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()
@@ -478,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 == ""

src/exo/utils/tests/test_power_sampler.py

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

src/exo/worker/disaggregated/protocol.py

@@ -1,3 +1,4 @@
+from dataclasses import dataclass
 from typing import BinaryIO, Literal
 
 import msgspec
@@ -23,7 +24,28 @@ class TensorBlob(msgspec.Struct):
     data: bytes
 
 
-class KVChunk(msgspec.Struct, tag="kv_chunk"):
+class _KVChunkHeader(msgspec.Struct, tag="kv_chunk"):
+    """Wire-side KV chunk metadata. Raw `keys` then `values` bytes follow on
+    the stream, lengths given by `keys_len` / `values_len`. Splitting them out
+    of the msgpack frame lets the producer pass tensor buffers via the buffer
+    protocol straight into the socket (one host-side memcpy total).
+    """
+
+    layer_idx: int
+    num_tokens: int
+    n_heads: int
+    head_dim: int
+    dtype: DType
+    keys_len: int
+    values_len: int
+
+
+@dataclass(frozen=True)
+class KVChunk:
+    """In-memory KV chunk reconstructed by `read_message` from
+    `_KVChunkHeader` + the raw bytes that follow on the wire.
+    """
+
     layer_idx: int
     num_tokens: int
     n_heads: int
@@ -51,10 +73,13 @@ class ErrorMessage(msgspec.Struct, tag="error"):
     message: str
 
 
+_WireMessage = _KVChunkHeader | ArraysState | Done | ErrorMessage
 Message = KVChunk | ArraysState | Done | ErrorMessage
 
 _msg_encoder = msgspec.msgpack.Encoder()
-_msg_decoder: msgspec.msgpack.Decoder[Message] = msgspec.msgpack.Decoder(Message)
+_msg_decoder: msgspec.msgpack.Decoder[_WireMessage] = msgspec.msgpack.Decoder(
+    _WireMessage
+)
 _header_encoder = msgspec.msgpack.Encoder()
 _header_decoder: msgspec.msgpack.Decoder[Header] = msgspec.msgpack.Decoder(Header)
 
@@ -99,7 +124,7 @@ def read_header(stream: BinaryIO) -> Header:
         raise ProtocolError(f"Bad header: {exc}") from exc
 
 
-def write_message(stream: BinaryIO, msg: Message) -> None:
+def write_message(stream: BinaryIO, msg: _WireMessage) -> None:
     write_frame(stream, _msg_encoder.encode(msg))
 
 
@@ -108,9 +133,22 @@ def read_message(stream: BinaryIO) -> Message | None:
     if not payload:
         return None
     try:
-        return _msg_decoder.decode(payload)
+        msg = _msg_decoder.decode(payload)
     except msgspec.DecodeError as exc:
         raise ProtocolError(f"Bad message: {exc}") from exc
+    if isinstance(msg, _KVChunkHeader):
+        keys = _read_exactly(stream, msg.keys_len)
+        values = _read_exactly(stream, msg.values_len)
+        return KVChunk(
+            layer_idx=msg.layer_idx,
+            num_tokens=msg.num_tokens,
+            n_heads=msg.n_heads,
+            head_dim=msg.head_dim,
+            dtype=msg.dtype,
+            keys=keys,
+            values=values,
+        )
+    return msg
 
 
 def write_kv_chunk(
@@ -121,21 +159,35 @@ def write_kv_chunk(
     n_heads: int,
     head_dim: int,
     dtype: DType,
-    keys: bytes,
-    values: bytes,
+    keys: "bytes | memoryview",
+    values: "bytes | memoryview",
 ) -> None:
-    write_message(
-        stream,
-        KVChunk(
+    """Stream KV chunk metadata + raw key/value bytes to the wire.
+
+    `keys` / `values` may be bytes-like (bytes, bytearray, memoryview) — the
+    raw payload is written directly to the buffered stream after the
+    msgpack-framed header, avoiding a memcpy through the msgpack encoder.
+    """
+    keys_len = len(keys)
+    values_len = len(values)
+    header_payload = _msg_encoder.encode(
+        _KVChunkHeader(
             layer_idx=layer_idx,
             num_tokens=num_tokens,
             n_heads=n_heads,
             head_dim=head_dim,
             dtype=dtype,
-            keys=keys,
-            values=values,
-        ),
+            keys_len=keys_len,
+            values_len=values_len,
+        )
     )
+    stream.write(len(header_payload).to_bytes(4, "big"))
+    stream.write(header_payload)
+    stream.write(keys)
+    stream.write(values)
+    # No per-chunk flush: the K/V payload is far larger than the
+    # BufferedWriter's internal buffer so it bypasses to the socket directly.
+    # The trailing `Done` frame's `write_frame` flushes once at the end.
 
 
 def write_arrays_state(

src/exo/worker/disaggregated/server.py

@@ -21,6 +21,7 @@ class PrefillRequest(msgspec.Struct):
     model_id: str = ""
     token_ids: list[int] = msgspec.field(default_factory=list)
     start_pos: int = 0
+    use_prefix_cache: bool = True
 
 
 _request_encoder = msgspec.msgpack.Encoder()
@@ -56,7 +57,10 @@ class _PrefillHandler(socketserver.StreamRequestHandler):
         super().setup()
         sock = cast(socket.socket, self.request)
         sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
-        sock.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF, 4 * 1024 * 1024)
+        # 64MB send buffer: K/V chunks are ~33MB each; a small SNDBUF
+        # back-pressures the writer thread between chunks and serializes
+        # network with compute.
+        sock.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF, 64 * 1024 * 1024)
 
     def handle(self) -> None:
         server = cast(PrefillServer, self.server)

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

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

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

@@ -34,7 +34,7 @@ class MlxBuilder(Builder):
     model_id: ModelId
     event_sender: MpSender[Event]
     cancel_receiver: MpReceiver[TaskId]
-    inference_model: Model | None = None
+    model: Model | None = None
     tokenizer: TokenizerWrapper | None = None
     group: mx.distributed.Group | None = None
     vision_processor: VisionProcessor | None = None
@@ -44,14 +44,14 @@ class MlxBuilder(Builder):
 
     def load(self, bound_instance: BoundInstance) -> Generator[ModelLoadingResponse]:
         (
-            self.inference_model,
+            self.model,
             self.tokenizer,
             self.vision_processor,
         ) = yield from load_mlx_items(bound_instance, self.group)
 
     def close(self) -> None:
         with contextlib.suppress(NameError, AttributeError):
-            del self.inference_model
+            del self.model
         with contextlib.suppress(NameError, AttributeError):
             del self.tokenizer
         with contextlib.suppress(NameError, AttributeError):
@@ -60,7 +60,7 @@ class MlxBuilder(Builder):
     def build(
         self,
     ) -> Engine:
-        assert self.inference_model
+        assert self.model
         assert self.tokenizer
 
         vision_processor = self.vision_processor
@@ -86,7 +86,7 @@ class MlxBuilder(Builder):
         if os.environ.get("EXO_NO_BATCH"):
             logger.info("using SequentialGenerator (batching disabled)")
             return SequentialGenerator(
-                model=self.inference_model,
+                model=self.model,
                 tokenizer=self.tokenizer,
                 group=self.group,
                 tool_parser=tool_parser,
@@ -100,7 +100,7 @@ class MlxBuilder(Builder):
         else:
             logger.info("using BatchGenerator")
             return BatchGenerator(
-                model=self.inference_model,
+                model=self.model,
                 tokenizer=self.tokenizer,
                 group=self.group,
                 tool_parser=tool_parser,

src/exo/worker/engines/mlx/disaggregated/adapter.py

@@ -22,7 +22,6 @@ from exo.worker.disaggregated.protocol import (
     write_kv_chunk,
 )
 from exo.worker.engines.mlx.types import KVCacheType
-from exo.worker.runner.bootstrap import logger
 
 _STR_TO_MX: dict[DType, mx.Dtype] = {
     "bfloat16": mx.bfloat16,
@@ -90,6 +89,18 @@ def nhd_to_bhsd(t: mx.array) -> mx.array:
     return mx.expand_dims(mx.transpose(t, (1, 0, 2)), 0)
 
 
+def _rotating_to_temporal(buf: mx.array, idx: int, offset: int, keep: int) -> mx.array:
+    seq = int(buf.shape[2])
+    if idx == seq:
+        return buf
+    if idx < offset:
+        return mx.concatenate(
+            [buf[..., :keep, :], buf[..., idx:, :], buf[..., keep:idx, :]],
+            axis=2,
+        )
+    return buf[..., :idx, :]
+
+
 def send_mlx_kv_cache(
     stream: BinaryIO,
     caches: KVCacheType,
@@ -103,7 +114,7 @@ def send_mlx_kv_cache(
         match c:
             case QuantizedKVCache() | CacheList() | DeepseekV4Cache():
                 raise NotImplementedError
-            case KVCache() | RotatingKVCache():
+            case KVCache():
                 keys = c.keys
                 values = c.values
                 if keys is None or values is None:
@@ -132,11 +143,39 @@ def send_mlx_kv_cache(
                     keys=array_to_bytes(k_nhd),
                     values=array_to_bytes(v_nhd),
                 )
-                if tokens_sent != 0 and num_tokens != tokens_sent:
-                    logger.critical(
-                        f"Unexpected number of tokens sent {num_tokens} != {tokens_sent}"
-                    )
-                tokens_sent = num_tokens
+                tokens_sent = max(tokens_sent, num_tokens)
+            case RotatingKVCache():
+                keys = c.keys
+                values = c.values
+                if keys is None or values is None:
+                    continue
+                offset = int(c.offset)
+                if offset <= 0:
+                    continue
+                idx = int(c._idx)
+                keep = int(c.keep)
+                with mx.stream(mx.Device(mx.cpu)):
+                    k_temporal = _rotating_to_temporal(keys, idx, offset, keep)
+                    v_temporal = _rotating_to_temporal(values, idx, offset, keep)
+                    k = mx.array(k_temporal)
+                    v = mx.array(v_temporal)
+                    k_nhd = bhsd_to_nhd(k)
+                    v_nhd = bhsd_to_nhd(v)
+                    mx.eval(k_nhd, v_nhd)
+                num_tokens = int(k_nhd.shape[0])
+                n_heads = int(k_nhd.shape[1])
+                head_dim = int(k_nhd.shape[2])
+                write_kv_chunk(
+                    stream,
+                    layer_idx=layer_idx,
+                    num_tokens=num_tokens,
+                    n_heads=n_heads,
+                    head_dim=head_dim,
+                    dtype=dtype,
+                    keys=array_to_bytes(k_nhd),
+                    values=array_to_bytes(v_nhd),
+                )
+                tokens_sent = max(tokens_sent, offset)
             case ArraysCache():
                 blobs: list[TensorBlob] = []
                 for a in c.state:

src/exo/worker/engines/mlx/disaggregated/client.py

@@ -79,6 +79,7 @@ def remote_prefill_fetch(
         result = PrefillResult(header=header)
         kv_by_layer: dict[int, list[KVChunk]] = defaultdict(list)
         chunks_received = 0
+        done_seen = False
 
         while True:
             msg = read_message(stream)
@@ -93,10 +94,17 @@ def remote_prefill_fetch(
                 result.arrays[msg.layer_idx] = msg.arrays
             elif isinstance(msg, Done):
                 result.total_tokens = msg.total_tokens
+                done_seen = True
                 break
             else:
                 raise RuntimeError(f"Prefill server error [{msg.code}]: {msg.message}")
 
+        if not done_seen:
+            raise ConnectionError(
+                "Prefill server closed before Done frame "
+                f"(received {chunks_received} kv chunks, {len(result.arrays)} arrays)"
+            )
+
         result.kv_chunks = dict(kv_by_layer)
         return result
     finally:

src/exo/worker/engines/mlx/generator/batch_generate.py

@@ -215,14 +215,18 @@ class ExoBatchGenerator:
         with vision_ctx:
             if use_remote and task_params.prefill_endpoint is not None:
                 try:
+                    # Send full prompt; producer's vLLM APC handles the prefix
+                    # match. `start_pos` aligns the writer's skip_tokens with
+                    # the consumer's locally-cached prefix.
                     _prefill_tps, _prefill_tokens, cache_snapshots = remote_prefill(
-                        prompt_tokens[:-1],
+                        all_prompt_tokens[:-1],
                         cache,
                         on_prefill_progress,
                         endpoint=task_params.prefill_endpoint,
                         request_id=str(uuid.uuid4()),
                         model_id=str(task_params.model),
                         start_pos=prefix_hit_length,
+                        use_prefix_cache=not is_bench or task_params.use_prefix_cache,
                     )
                     remote_prefilled = True
                 except Exception: