feat: auto-discover models from HuggingFace Hub cache and custom paths

Searches ~/.cache/huggingface/hub by default so models downloaded by
other tools (HF CLI, transformers, LM Studio, etc.) are used without
re-downloading. Adds EXO_MODELS_PATH env var support and a Model Search
Paths setting in the macOS app.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

.mlx_typings/mlx/core/__init__.pyi

@@ -215,22 +215,6 @@ class StreamContext:
         traceback: object | None = ...,
     ) -> None: ...
 
-def device_info() -> dict[str, str | int]:
-    """
-    Get information about the GPU device and system settings.
-
-    Currently returns:
-
-    * ``architecture``
-    * ``max_buffer_size``
-    * ``max_recommended_working_set_size``
-    * ``memory_size``
-    * ``resource_limit``
-
-    Returns:
-        dict: A dictionary with string keys and string or integer values.
-    """
-
 def abs(a: array, /, *, stream: Stream | Device | None = ...) -> array:
     """
     Element-wise absolute value.

.mlx_typings/mlx_lm/models/cache.pyi

@@ -249,8 +249,7 @@ class ChunkedKVCache(KVCache):
         ...
 
 class CacheList(_BaseCache):
-    caches: tuple[_BaseCache, ...]
-    def __init__(self, *caches: _BaseCache) -> None: ...
+    def __init__(self, *caches) -> None: ...
     def __getitem__(self, idx): ...
     def is_trimmable(self):  # -> bool:
         ...

app/EXO/EXO/ContentView.swift

@@ -21,21 +21,6 @@ struct ContentView: View {
     @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 = ""
-    @State private var pendingEnableImageModels = false
 
     var body: some View {
         VStack(alignment: .leading, spacing: 12) {
@@ -394,283 +379,6 @@ struct ContentView: View {
         }
     }
 
-    private var thunderboltStatusText: String {
-        switch networkStatusService.status.thunderboltBridgeState {
-        case .some(.disabled):
-            return "Thunderbolt Bridge: Disabled"
-        case .some(.deleted):
-            return "Thunderbolt Bridge: Deleted"
-        case .some(.enabled):
-            return "Thunderbolt Bridge: Enabled"
-        case nil:
-            return "Thunderbolt Bridge: Unknown"
-        }
-    }
-
-    private var thunderboltStatusColor: Color {
-        switch networkStatusService.status.thunderboltBridgeState {
-        case .some(.disabled), .some(.deleted):
-            return .green
-        case .some(.enabled):
-            return .red
-        case nil:
-            return .secondary
-        }
-    }
-
-    /// Shows TB bridge status for all nodes from exo cluster state
-    private var clusterThunderboltBridgeView: some View {
-        let bridgeStatuses = stateService.latestSnapshot?.nodeThunderboltBridge ?? [:]
-        let localNodeId = stateService.localNodeId
-        let nodeProfiles = stateService.latestSnapshot?.nodeProfiles ?? [:]
-
-        return VStack(alignment: .leading, spacing: 1) {
-            if bridgeStatuses.isEmpty {
-                Text("Cluster TB Bridge: No data")
-                    .font(.caption2)
-                    .foregroundColor(.secondary)
-            } else {
-                Text("Cluster TB Bridge Status:")
-                    .font(.caption2)
-                    .foregroundColor(.secondary)
-                ForEach(Array(bridgeStatuses.keys.sorted()), id: \.self) { nodeId in
-                    if let status = bridgeStatuses[nodeId] {
-                        let nodeName =
-                            nodeProfiles[nodeId]?.friendlyName ?? String(nodeId.prefix(8))
-                        let isLocal = nodeId == localNodeId
-                        let prefix = isLocal ? "  \(nodeName) (local):" : "  \(nodeName):"
-                        let statusText =
-                            !status.exists
-                            ? "N/A"
-                            : (status.enabled ? "Enabled" : "Disabled")
-                        let color: Color =
-                            !status.exists
-                            ? .secondary
-                            : (status.enabled ? .red : .green)
-                        Text("\(prefix) \(statusText)")
-                            .font(.caption2)
-                            .foregroundColor(color)
-                    }
-                }
-            }
-        }
-    }
-
-    private var interfaceIpList: some View {
-        let statuses = networkStatusService.status.interfaceStatuses
-        return VStack(alignment: .leading, spacing: 1) {
-            Text("Interfaces (en0–en7):")
-                .font(.caption2)
-                .foregroundColor(.secondary)
-            if statuses.isEmpty {
-                Text("  Unknown")
-                    .font(.caption2)
-                    .foregroundColor(.secondary)
-            } else {
-                ForEach(statuses, id: \.interfaceName) { status in
-                    let ipText = status.ipAddress ?? "No IP"
-                    Text("  \(status.interfaceName): \(ipText)")
-                        .font(.caption2)
-                        .foregroundColor(status.ipAddress == nil ? .red : .green)
-                }
-            }
-        }
-    }
-
-    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
-        let nodeProfiles = stateService.latestSnapshot?.nodeProfiles ?? [:]
-        let localDevices = networkStatusService.status.localRdmaDevices
-        let localPorts = networkStatusService.status.localRdmaActivePorts
-
-        return VStack(alignment: .leading, spacing: 1) {
-            if rdmaStatuses.isEmpty {
-                Text("Cluster RDMA: No data")
-                    .font(.caption2)
-                    .foregroundColor(.secondary)
-            } else {
-                Text("Cluster RDMA Status:")
-                    .font(.caption2)
-                    .foregroundColor(.secondary)
-                ForEach(Array(rdmaStatuses.keys.sorted()), id: \.self) { nodeId in
-                    if let status = rdmaStatuses[nodeId] {
-                        let nodeName =
-                            nodeProfiles[nodeId]?.friendlyName ?? String(nodeId.prefix(8))
-                        let isLocal = nodeId == localNodeId
-                        let prefix = isLocal ? "  \(nodeName) (local):" : "  \(nodeName):"
-                        let statusText = status.enabled ? "Enabled" : "Disabled"
-                        let color: Color = status.enabled ? .green : .orange
-                        Text("\(prefix) \(statusText)")
-                            .font(.caption2)
-                            .foregroundColor(color)
-                    }
-                }
-            }
-            if !localDevices.isEmpty {
-                Text("  Local Devices: \(localDevices.joined(separator: ", "))")
-                    .font(.caption2)
-                    .foregroundColor(.secondary)
-            }
-            if !localPorts.isEmpty {
-                Text("  Local Active Ports:")
-                    .font(.caption2)
-                    .foregroundColor(.secondary)
-                ForEach(localPorts, id: \.device) { port in
-                    Text("    \(port.device) port \(port.port): \(port.state)")
-                        .font(.caption2)
-                        .foregroundColor(.green)
-                }
-            }
-        }
-    }
-
-    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) {
-                    Text("What's the issue? (optional)")
-                        .font(.caption2)
-                        .foregroundColor(.secondary)
-                    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: {
@@ -711,143 +419,6 @@ struct ContentView: View {
         )
     }
 
-    private func sendBugReport() async {
-        bugReportPhase = .sending("Collecting logs...")
-        let service = BugReportService()
-        let description = bugReportUserDescription.trimmingCharacters(in: .whitespacesAndNewlines)
-        do {
-            let outcome = try await service.sendReport(
-                isManual: true,
-                userDescription: description.isEmpty ? nil : description
-            )
-            if outcome.success {
-                bugReportPhase = .success(outcome.message)
-            } else {
-                bugReportPhase = .failure(outcome.message)
-            }
-        } catch {
-            bugReportPhase = .failure(error.localizedDescription)
-        }
-    }
-
-    private func openGitHubIssue() {
-        let description = bugReportUserDescription.trimmingCharacters(in: .whitespacesAndNewlines)
-
-        var bodyParts: [String] = []
-        bodyParts.append("## Describe the bug")
-        bodyParts.append("")
-        if !description.isEmpty {
-            bodyParts.append(description)
-        } else {
-            bodyParts.append("A clear and concise description of what the bug is.")
-        }
-        bodyParts.append("")
-        bodyParts.append("## Environment")
-        bodyParts.append("")
-        bodyParts.append("- macOS Version: \(ProcessInfo.processInfo.operatingSystemVersionString)")
-        bodyParts.append("- EXO Version: \(buildTag) (\(buildCommit))")
-        bodyParts.append("")
-        bodyParts.append("## Additional context")
-        bodyParts.append("")
-        bodyParts.append("A bug report with diagnostic logs was submitted via the app.")
-
-        let body = bodyParts.joined(separator: "\n")
-
-        var components = URLComponents(string: "https://github.com/exo-explore/exo/issues/new")!
-        components.queryItems = [
-            URLQueryItem(name: "template", value: "bug_report.md"),
-            URLQueryItem(name: "title", value: "[BUG] "),
-            URLQueryItem(name: "body", value: body),
-            URLQueryItem(name: "labels", value: "bug"),
-        ]
-
-        if let url = components.url {
-            NSWorkspace.shared.open(url)
-        }
-    }
-
-    private func showUninstallConfirmationAlert() {
-        let alert = NSAlert()
-        alert.messageText = "Uninstall EXO"
-        alert.informativeText = """
-            This will remove EXO and all its system components:
-
-            • Network configuration daemon
-            • Launch at login registration
-            • EXO network location
-
-            The app will be moved to Trash.
-            """
-        alert.alertStyle = .warning
-        alert.addButton(withTitle: "Uninstall")
-        alert.addButton(withTitle: "Cancel")
-
-        // Style the Uninstall button as destructive
-        if let uninstallButton = alert.buttons.first {
-            uninstallButton.hasDestructiveAction = true
-        }
-
-        let response = alert.runModal()
-        if response == .alertFirstButtonReturn {
-            performUninstall()
-        }
-    }
-
-    private func performUninstall() {
-        uninstallInProgress = true
-
-        // Stop EXO process first
-        controller.cancelPendingLaunch()
-        controller.stop()
-        stateService.stopPolling()
-
-        // Run the privileged uninstall on a background thread
-        // Using .utility QoS to avoid priority inversion with NSAppleScript's subprocess
-        DispatchQueue.global(qos: .utility).async {
-            do {
-                // Remove network setup daemon and components (requires admin privileges)
-                try NetworkSetupHelper.uninstall()
-
-                DispatchQueue.main.async {
-                    // Unregister from launch at login
-                    LaunchAtLoginHelper.disable()
-
-                    // Move app to trash
-                    self.moveAppToTrash()
-
-                    // Quit the app
-                    DispatchQueue.main.asyncAfter(deadline: .now() + 0.5) {
-                        NSApplication.shared.terminate(nil)
-                    }
-                }
-            } catch {
-                DispatchQueue.main.async {
-                    self.showErrorAlert(message: error.localizedDescription)
-                    self.uninstallInProgress = false
-                }
-            }
-        }
-    }
-
-    private func showErrorAlert(message: String) {
-        let alert = NSAlert()
-        alert.messageText = "Uninstall Failed"
-        alert.informativeText = message
-        alert.alertStyle = .critical
-        alert.addButton(withTitle: "OK")
-        alert.runModal()
-    }
-
-    private func moveAppToTrash() {
-        guard let appURL = Bundle.main.bundleURL as URL? else { return }
-        do {
-            try FileManager.default.trashItem(at: appURL, resultingItemURL: nil)
-        } catch {
-            // If we can't trash the app, that's OK - user can do it manually
-            // The important system components have already been cleaned up
-        }
-    }
-
     private var buildTag: String {
         Bundle.main.infoDictionary?["EXOBuildTag"] as? String ?? "unknown"
     }

app/EXO/EXO/ExoProcessController.swift

@@ -5,8 +5,8 @@ import Foundation
 private let customNamespaceKey = "EXOCustomNamespace"
 private let hfTokenKey = "EXOHFToken"
 private let enableImageModelsKey = "EXOEnableImageModels"
-private let offlineModeKey = "EXOOfflineMode"
 private let onboardingCompletedKey = "EXOOnboardingCompleted"
+private let modelSearchPathsKey = "EXOModelSearchPaths"
 
 @MainActor
 final class ExoProcessController: ObservableObject {
@@ -61,12 +61,12 @@ final class ExoProcessController: ObservableObject {
             UserDefaults.standard.set(enableImageModels, forKey: enableImageModelsKey)
         }
     }
-    @Published var offlineMode: Bool = {
-        return UserDefaults.standard.bool(forKey: offlineModeKey)
+    @Published var modelSearchPaths: String = {
+        return UserDefaults.standard.string(forKey: modelSearchPathsKey) ?? ""
     }()
     {
         didSet {
-            UserDefaults.standard.set(offlineMode, forKey: offlineModeKey)
+            UserDefaults.standard.set(modelSearchPaths, forKey: modelSearchPathsKey)
         }
     }
 
@@ -276,8 +276,8 @@ final class ExoProcessController: ObservableObject {
         if enableImageModels {
             environment["EXO_ENABLE_IMAGE_MODELS"] = "true"
         }
-        if offlineMode {
-            environment["EXO_OFFLINE"] = "true"
+        if !modelSearchPaths.isEmpty {
+            environment["EXO_MODELS_PATH"] = modelSearchPaths
         }
 
         var paths: [String] = []

app/EXO/EXO/Services/BugReportService.swift

@@ -38,8 +38,7 @@ struct BugReportService {
     func sendReport(
         baseURL: URL = URL(string: "http://127.0.0.1:52415")!,
         now: Date = Date(),
-        isManual: Bool = false,
-        userDescription: String? = nil
+        isManual: Bool = false
     ) async throws -> BugReportOutcome {
         let timestamp = Self.runTimestampString(now)
         let dayPrefix = Self.dayPrefixString(now)
@@ -61,8 +60,7 @@ struct BugReportService {
             ifconfig: ifconfigText,
             debugInfo: debugInfo,
             isManual: isManual,
-            clusterTbBridgeStatus: clusterTbBridgeStatus,
-            userDescription: userDescription
+            clusterTbBridgeStatus: clusterTbBridgeStatus
         )
 
         let eventLogFiles = readAllEventLogs()
@@ -308,8 +306,7 @@ struct BugReportService {
         ifconfig: String,
         debugInfo: DebugInfo,
         isManual: Bool,
-        clusterTbBridgeStatus: [[String: Any]]?,
-        userDescription: String? = nil
+        clusterTbBridgeStatus: [[String: Any]]?
     ) -> Data? {
         let system = readSystemMetadata()
         let exo = readExoMetadata()
@@ -326,9 +323,6 @@ struct BugReportService {
         if let tbStatus = clusterTbBridgeStatus {
             payload["cluster_thunderbolt_bridge"] = tbStatus
         }
-        if let desc = userDescription, !desc.isEmpty {
-            payload["user_description"] = desc
-        }
         return try? JSONSerialization.data(withJSONObject: payload, options: [.prettyPrinted])
     }
 

app/EXO/EXO/Views/SettingsView.swift

@@ -13,7 +13,7 @@ struct SettingsView: View {
     @State private var pendingNamespace: String = ""
     @State private var pendingHFToken: String = ""
     @State private var pendingEnableImageModels = false
-    @State private var pendingOfflineMode = false
+    @State private var pendingModelSearchPaths: String = ""
     @State private var needsRestart = false
     @State private var bugReportInFlight = false
     @State private var bugReportMessage: String?
@@ -43,7 +43,7 @@ struct SettingsView: View {
             pendingNamespace = controller.customNamespace
             pendingHFToken = controller.hfToken
             pendingEnableImageModels = controller.enableImageModels
-            pendingOfflineMode = controller.offlineMode
+            pendingModelSearchPaths = controller.modelSearchPaths
             needsRestart = false
         }
     }
@@ -74,13 +74,6 @@ struct SettingsView: View {
                     .foregroundColor(.secondary)
             }
 
-            Section {
-                Toggle("Offline Mode", isOn: $pendingOfflineMode)
-                Text("Skip internet checks and use only locally available models.")
-                    .font(.caption)
-                    .foregroundColor(.secondary)
-            }
-
             Section {
                 HStack {
                     Spacer()
@@ -106,6 +99,19 @@ struct SettingsView: View {
                     .foregroundColor(.secondary)
             }
 
+            Section {
+                LabeledContent("Model Search Paths") {
+                    TextField("/path/one:/path/two", text: $pendingModelSearchPaths)
+                        .textFieldStyle(.roundedBorder)
+                        .frame(width: 200)
+                }
+                Text(
+                    "Extra directories to search for pre-downloaded models (colon-separated). HuggingFace cache (~/.cache/huggingface/hub) is checked automatically."
+                )
+                .font(.caption)
+                .foregroundColor(.secondary)
+            }
+
             Section {
                 HStack {
                     Spacer()
@@ -454,22 +460,22 @@ struct SettingsView: View {
 
     private var hasGeneralChanges: Bool {
         pendingNamespace != controller.customNamespace || pendingHFToken != controller.hfToken
-            || pendingOfflineMode != controller.offlineMode
     }
 
     private var hasModelChanges: Bool {
         pendingEnableImageModels != controller.enableImageModels
+            || pendingModelSearchPaths != controller.modelSearchPaths
     }
 
     private func applyGeneralSettings() {
         controller.customNamespace = pendingNamespace
         controller.hfToken = pendingHFToken
-        controller.offlineMode = pendingOfflineMode
         restartIfRunning()
     }
 
     private func applyModelSettings() {
         controller.enableImageModels = pendingEnableImageModels
+        controller.modelSearchPaths = pendingModelSearchPaths
         restartIfRunning()
     }
 

bench/exo_bench.py

@@ -75,7 +75,7 @@ def load_tokenizer_for_bench(model_id: str) -> Any:
         model_path = Path(
             snapshot_download(
                 model_id,
-                allow_patterns=["*.json", "*.py", "*.tiktoken", "*.model"],
+                allow_patterns=["*.json", "*.py", "*.tiktoken"],
             )
         )
 

bench/harness.py

@@ -371,7 +371,7 @@ def run_planning_phase(
                 unwrap_instance(p["DownloadCompleted"]["shardMetadata"])["modelCard"][
                     "modelId"
                 ],
-                p["DownloadCompleted"]["total"]["inBytes"],
+                p["DownloadCompleted"]["totalBytes"]["inBytes"],
             )
             for p in node_downloads
             if "DownloadCompleted" in p

bench/pyproject.toml

@@ -10,7 +10,6 @@ dependencies = [
     "huggingface-hub>=0.33.4",
     "tiktoken>=0.12.0",
     "jinja2>=3.1.0",
-    "protobuf>=5.29.0",
 ]
 
 [build-system]

dashboard/src/lib/components/ChatForm.svelte

@@ -7,6 +7,8 @@
     editingImage,
     clearEditingImage,
     selectedChatModel,
+    setSelectedChatModel,
+    instances,
     ttftMs,
     tps,
     totalTokens,
@@ -28,18 +30,6 @@
     modelTasks?: Record<string, string[]>;
     modelCapabilities?: Record<string, string[]>;
     onSend?: () => void;
-    onAutoSend?: (
-      content: string,
-      files?: {
-        id: string;
-        name: string;
-        type: string;
-        textContent?: string;
-        preview?: string;
-      }[],
-    ) => void;
-    onOpenModelPicker?: () => void;
-    modelDisplayOverride?: string;
   }
 
   let {
@@ -51,9 +41,6 @@
     modelTasks = {},
     modelCapabilities = {},
     onSend,
-    onAutoSend,
-    onOpenModelPicker,
-    modelDisplayOverride,
   }: Props = $props();
 
   let message = $state("");
@@ -64,12 +51,27 @@
   const thinkingEnabled = $derived(thinkingEnabledStore());
   let loading = $derived(isLoading());
   const currentModel = $derived(selectedChatModel());
+  const instanceData = $derived(instances());
   const currentTtft = $derived(ttftMs());
   const currentTps = $derived(tps());
   const currentTokens = $derived(totalTokens());
   const currentEditingImage = $derived(editingImage());
   const isEditMode = $derived(currentEditingImage !== null);
 
+  // Custom dropdown state
+  let isModelDropdownOpen = $state(false);
+  let dropdownButtonRef: HTMLButtonElement | undefined = $state();
+  let dropdownPosition = $derived(() => {
+    if (!dropdownButtonRef || !isModelDropdownOpen)
+      return { top: 0, left: 0, width: 0 };
+    const rect = dropdownButtonRef.getBoundingClientRect();
+    return {
+      top: rect.top,
+      left: rect.left,
+      width: rect.width,
+    };
+  });
+
   // Accept all supported file types
   const acceptString = getAcceptString(["image", "text", "pdf"]);
 
@@ -122,14 +124,73 @@
         uploadedFiles.length > 0),
   );
 
-  // Short label for the currently selected model
-  const currentModelLabel = $derived(
-    currentModel
-      ? currentModel.split("/").pop() || currentModel
-      : modelDisplayOverride
-        ? modelDisplayOverride.split("/").pop() || modelDisplayOverride
-        : "",
-  );
+  // Extract available models from running instances
+  const availableModels = $derived(() => {
+    const models: Array<{ id: string; label: string; isImageModel: boolean }> =
+      [];
+    for (const [, instance] of Object.entries(instanceData)) {
+      const modelId = getInstanceModelId(instance);
+      if (
+        modelId &&
+        modelId !== "Unknown" &&
+        !models.some((m) => m.id === modelId)
+      ) {
+        models.push({
+          id: modelId,
+          label: modelId.split("/").pop() || modelId,
+          isImageModel: modelSupportsImageGeneration(modelId),
+        });
+      }
+    }
+    return models;
+  });
+
+  // Track previous model IDs to detect newly added models (plain variable to avoid reactive loop)
+  let previousModelIds: Set<string> = new Set();
+
+  // Auto-select the first available model if none is selected, if current selection is stale, or if a new model is added
+  $effect(() => {
+    const models = availableModels();
+    const currentModelIds = new Set(models.map((m) => m.id));
+
+    if (models.length > 0) {
+      // Find newly added models (in current but not in previous)
+      const newModels = models.filter((m) => !previousModelIds.has(m.id));
+
+      // If no model selected, select the first available
+      if (!currentModel) {
+        setSelectedChatModel(models[0].id);
+      }
+      // If current model is stale (no longer has a running instance), reset to first available
+      else if (!models.some((m) => m.id === currentModel)) {
+        setSelectedChatModel(models[0].id);
+      }
+      // If a new model was just added, select it
+      else if (newModels.length > 0 && previousModelIds.size > 0) {
+        setSelectedChatModel(newModels[0].id);
+      }
+    } else {
+      // No instances running - clear the selected model
+      if (currentModel) {
+        setSelectedChatModel("");
+      }
+    }
+
+    // Update previous model IDs for next comparison
+    previousModelIds = currentModelIds;
+  });
+
+  function getInstanceModelId(instanceWrapped: unknown): string {
+    if (!instanceWrapped || typeof instanceWrapped !== "object") return "";
+    const keys = Object.keys(instanceWrapped as Record<string, unknown>);
+    if (keys.length === 1) {
+      const instance = (instanceWrapped as Record<string, unknown>)[
+        keys[0]
+      ] as { shardAssignments?: { modelId?: string } };
+      return instance?.shardAssignments?.modelId || "";
+    }
+    return "";
+  }
 
   async function handleFiles(files: File[]) {
     if (files.length === 0) return;
@@ -216,15 +277,6 @@
     uploadedFiles = [];
     resetTextareaHeight();
 
-    // When onAutoSend is provided, the parent controls all send logic
-    // (including launching non-running models before sending)
-    if (onAutoSend) {
-      onAutoSend(content, files);
-      onSend?.();
-      setTimeout(() => textareaRef?.focus(), 10);
-      return;
-    }
-
     // Use image editing if in edit mode
     if (isEditMode && currentEditingImage && content) {
       editImage(content, currentEditingImage.imageDataUrl);
@@ -368,7 +420,7 @@
     {/if}
 
     <!-- Model selector (when enabled) -->
-    {#if showModelSelector}
+    {#if showModelSelector && availableModels().length > 0}
       <div
         class="flex items-center justify-between gap-2 px-3 py-2 border-b border-exo-medium-gray/30"
       >
@@ -377,22 +429,33 @@
             class="text-xs text-exo-light-gray uppercase tracking-wider flex-shrink-0"
             >MODEL:</span
           >
-          <!-- Model button — opens the full model picker -->
+          <!-- Custom dropdown -->
           <div class="relative flex-1 max-w-xs">
             <button
+              bind:this={dropdownButtonRef}
               type="button"
-              onclick={() => onOpenModelPicker?.()}
-              class="w-full bg-exo-medium-gray/50 border border-exo-yellow/30 rounded pl-3 pr-8 py-1.5 text-xs font-mono text-left tracking-wide cursor-pointer transition-all duration-200 hover:border-exo-yellow/50 focus:outline-none focus:border-exo-yellow/70"
+              onclick={() => (isModelDropdownOpen = !isModelDropdownOpen)}
+              class="w-full bg-exo-medium-gray/50 border border-exo-yellow/30 rounded pl-3 pr-8 py-1.5 text-xs font-mono text-left tracking-wide cursor-pointer transition-all duration-200 hover:border-exo-yellow/50 focus:outline-none focus:border-exo-yellow/70 {isModelDropdownOpen
+                ? 'border-exo-yellow/70'
+                : ''}"
             >
-              {#if currentModelLabel}
-                <span class="text-exo-yellow truncate">{currentModelLabel}</span
+              {#if availableModels().find((m) => m.id === currentModel)}
+                <span class="text-exo-yellow truncate"
+                  >{availableModels().find((m) => m.id === currentModel)
+                    ?.label}</span
+                >
+              {:else if availableModels().length > 0}
+                <span class="text-exo-yellow truncate"
+                  >{availableModels()[0].label}</span
                 >
               {:else}
                 <span class="text-exo-light-gray/50">— SELECT MODEL —</span>
               {/if}
             </button>
             <div
-              class="absolute right-2 top-1/2 -translate-y-1/2 pointer-events-none"
+              class="absolute right-2 top-1/2 -translate-y-1/2 pointer-events-none transition-transform duration-200 {isModelDropdownOpen
+                ? 'rotate-180'
+                : ''}"
             >
               <svg
                 class="w-3 h-3 text-exo-yellow/60"
@@ -409,6 +472,78 @@
               </svg>
             </div>
           </div>
+
+          {#if isModelDropdownOpen}
+            <!-- Backdrop to close dropdown -->
+            <button
+              type="button"
+              class="fixed inset-0 z-[9998] cursor-default"
+              onclick={() => (isModelDropdownOpen = false)}
+              aria-label="Close dropdown"
+            ></button>
+
+            <!-- Dropdown Panel - fixed positioning to escape overflow:hidden -->
+            <div
+              class="fixed bg-exo-dark-gray border border-exo-yellow/30 rounded shadow-lg shadow-black/50 z-[9999] max-h-48 overflow-y-auto"
+              style="bottom: calc(100vh - {dropdownPosition()
+                .top}px + 4px); left: {dropdownPosition()
+                .left}px; width: {dropdownPosition().width}px;"
+            >
+              <div class="py-1">
+                {#each availableModels() as model}
+                  <button
+                    type="button"
+                    onclick={() => {
+                      setSelectedChatModel(model.id);
+                      isModelDropdownOpen = false;
+                    }}
+                    class="w-full px-3 py-2 text-left text-xs font-mono tracking-wide transition-colors duration-100 flex items-center gap-2 {currentModel ===
+                    model.id
+                      ? 'bg-transparent text-exo-yellow'
+                      : 'text-exo-light-gray hover:text-exo-yellow'}"
+                  >
+                    {#if currentModel === model.id}
+                      <svg
+                        class="w-3 h-3 flex-shrink-0"
+                        fill="currentColor"
+                        viewBox="0 0 20 20"
+                      >
+                        <path
+                          fill-rule="evenodd"
+                          d="M16.707 5.293a1 1 0 010 1.414l-8 8a1 1 0 01-1.414 0l-4-4a1 1 0 011.414-1.414L8 12.586l7.293-7.293a1 1 0 011.414 0z"
+                          clip-rule="evenodd"
+                        />
+                      </svg>
+                    {:else}
+                      <span class="w-3"></span>
+                    {/if}
+                    {#if model.isImageModel}
+                      <svg
+                        class="w-3.5 h-3.5 flex-shrink-0 text-exo-yellow"
+                        fill="none"
+                        viewBox="0 0 24 24"
+                        stroke="currentColor"
+                        stroke-width="2"
+                        aria-label="Image generation model"
+                      >
+                        <rect
+                          x="3"
+                          y="3"
+                          width="18"
+                          height="18"
+                          rx="2"
+                          ry="2"
+                        />
+                        <circle cx="8.5" cy="8.5" r="1.5" />
+                        <polyline points="21 15 16 10 5 21" />
+                      </svg>
+                    {/if}
+                    <span class="truncate flex-1">{model.label}</span>
+                  </button>
+                {/each}
+              </div>
+            </div>
+          {/if}
         </div>
         <!-- Thinking toggle -->
         {#if modelSupportsThinking()}

dashboard/src/lib/components/ChatModelSelector.svelte

@@ -1,401 +0,0 @@
-<script lang="ts" module>
-  export interface ChatModelInfo {
-    id: string;
-    name: string;
-    base_model: string;
-    storage_size_megabytes: number;
-    capabilities: string[];
-    family: string;
-    quantization: string;
-  }
-
-  // Auto mode tier list (for when user just starts typing)
-  export const AUTO_TIERS: string[][] = [
-    // Tier 1 (frontier)
-    ["DeepSeek V3.1", "GLM-5", "Kimi K2.5", "Qwen3 Coder Next"],
-    // Tier 2 (excellent)
-    [
-      "Kimi K2",
-      "Qwen3 235B",
-      "MiniMax M2.5",
-      "Step 3.5 Flash",
-      "Qwen3 Next 80B",
-    ],
-    // Tier 3 (great)
-    [
-      "GLM 4.7",
-      "MiniMax M2.1",
-      "Qwen3 Coder 480B",
-      "GLM 4.5 Air",
-      "Llama 3.3 70B",
-    ],
-    // Tier 4 (good)
-    ["GPT-OSS 120B", "Qwen3 30B", "Llama 3.1 70B", "GLM 4.7 Flash"],
-    // Tier 5 (small/fast)
-    [
-      "Llama 3.1 8B",
-      "GPT-OSS 20B",
-      "Llama 3.2 3B",
-      "Qwen3 0.6B",
-      "Llama 3.2 1B",
-    ],
-  ];
-
-  /** Return the tier index (0 = best) for a base_model name. */
-  export function getAutoTierIndex(baseModel: string): number {
-    for (let i = 0; i < AUTO_TIERS.length; i++) {
-      if (AUTO_TIERS[i].includes(baseModel)) return i;
-    }
-    return AUTO_TIERS.length; // not in any tier → lowest priority
-  }
-
-  /** Auto mode: walk tiers top-down, pick biggest fitting variant from highest tier. */
-  export function pickAutoModel(
-    modelList: ChatModelInfo[],
-    memoryGB: number,
-  ): ChatModelInfo | null {
-    for (const tier of AUTO_TIERS) {
-      const candidates: ChatModelInfo[] = [];
-      for (const baseModel of tier) {
-        const variants = modelList
-          .filter(
-            (m) =>
-              m.base_model === baseModel &&
-              (m.storage_size_megabytes || 0) / 1024 <= memoryGB &&
-              (m.storage_size_megabytes || 0) > 0,
-          )
-          .sort(
-            (a, b) =>
-              (b.storage_size_megabytes || 0) - (a.storage_size_megabytes || 0),
-          );
-        if (variants[0]) candidates.push(variants[0]);
-      }
-      if (candidates.length > 0) {
-        candidates.sort(
-          (a, b) =>
-            (b.storage_size_megabytes || 0) - (a.storage_size_megabytes || 0),
-        );
-        return candidates[0];
-      }
-    }
-    return null;
-  }
-</script>
-
-<script lang="ts">
-  interface CategoryRecommendation {
-    category: "coding" | "writing" | "agentic" | "biggest";
-    label: string;
-    model: ChatModelInfo | null;
-    tooltip: string;
-  }
-
-  interface Props {
-    models: ChatModelInfo[];
-    clusterLabel: string;
-    totalMemoryGB: number;
-    onSelect: (modelId: string, category: string) => void;
-    onAddModel: () => void;
-    class?: string;
-  }
-
-  let {
-    models,
-    clusterLabel,
-    totalMemoryGB,
-    onSelect,
-    onAddModel,
-    class: className = "",
-  }: Props = $props();
-
-  // --- Hardcoded Rankings ---
-  const CODING_RANKING = [
-    "Qwen3 Coder Next",
-    "Qwen3 Coder 480B",
-    "Qwen3 30B",
-    "GPT-OSS 20B",
-    "Llama 3.1 8B",
-    "Llama 3.2 3B",
-    "Qwen3 0.6B",
-  ];
-
-  const WRITING_RANKING = [
-    "Kimi K2.5",
-    "Kimi K2",
-    "Qwen3 Next 80B",
-    "Llama 3.3 70B",
-    "MiniMax M2.5",
-    "GLM 4.5 Air",
-    "GLM 4.7 Flash",
-    "GPT-OSS 20B",
-    "Llama 3.1 8B",
-    "Llama 3.2 3B",
-    "Qwen3 0.6B",
-  ];
-
-  const AGENTIC_RANKING = [
-    "DeepSeek V3.1",
-    "GLM-5",
-    "Qwen3 235B",
-    "Step 3.5 Flash",
-    "GLM 4.7",
-    "MiniMax M2.1",
-    "GPT-OSS 120B",
-    "Llama 3.3 70B",
-    "Llama 3.1 70B",
-    "GLM 4.7 Flash",
-    "GPT-OSS 20B",
-    "Qwen3 30B",
-    "Llama 3.1 8B",
-    "Llama 3.2 3B",
-    "Qwen3 0.6B",
-  ];
-
-  function getModelSizeGB(m: ChatModelInfo): number {
-    return (m.storage_size_megabytes || 0) / 1024;
-  }
-
-  function fitsInMemory(m: ChatModelInfo): boolean {
-    return getModelSizeGB(m) <= totalMemoryGB && getModelSizeGB(m) > 0;
-  }
-
-  /** For a given base_model name, find the biggest quant variant that fits in memory. */
-  function pickBestVariant(baseModel: string): ChatModelInfo | null {
-    const variants = models
-      .filter((m) => m.base_model === baseModel && fitsInMemory(m))
-      .sort((a, b) => getModelSizeGB(b) - getModelSizeGB(a));
-    return variants[0] ?? null;
-  }
-
-  /** Walk a ranked list of base_model names, return the first that has a fitting variant. */
-  function pickFromRanking(ranking: string[]): ChatModelInfo | null {
-    for (const baseModel of ranking) {
-      const pick = pickBestVariant(baseModel);
-      if (pick) return pick;
-    }
-    return null;
-  }
-
-  /** Pick the single biggest model that fits. */
-  function pickBiggest(): ChatModelInfo | null {
-    const fitting = models
-      .filter((m) => fitsInMemory(m))
-      .sort((a, b) => getModelSizeGB(b) - getModelSizeGB(a));
-    return fitting[0] ?? null;
-  }
-
-  const recommendations = $derived.by((): CategoryRecommendation[] => {
-    return [
-      {
-        category: "coding",
-        label: "Best for Coding",
-        model: pickFromRanking(CODING_RANKING),
-        tooltip:
-          "Ranked by coding benchmark performance (LiveCodeBench, SWE-bench)",
-      },
-      {
-        category: "writing",
-        label: "Best for Writing",
-        model: pickFromRanking(WRITING_RANKING),
-        tooltip: "Ranked by creative writing quality and instruction following",
-      },
-      {
-        category: "agentic",
-        label: "Best Agentic",
-        model: pickFromRanking(AGENTIC_RANKING),
-        tooltip: "Ranked by reasoning, planning, and tool-use capability",
-      },
-      {
-        category: "biggest",
-        label: "Biggest",
-        model: pickBiggest(),
-        tooltip: "Largest model that fits in your available memory",
-      },
-    ];
-  });
-
-  function formatSize(mb: number): string {
-    const gb = mb / 1024;
-    if (gb >= 100) return `${Math.round(gb)} GB`;
-    return `${gb.toFixed(1)} GB`;
-  }
-
-  // Category icons (SVG paths)
-  const categoryIcons: Record<string, string> = {
-    coding:
-      "M8 9l3 3-3 3m5 0h3M5 20h14a2 2 0 002-2V6a2 2 0 00-2-2H5a2 2 0 00-2 2v12a2 2 0 002 2z",
-    writing:
-      "M11 5H6a2 2 0 00-2 2v11a2 2 0 002 2h11a2 2 0 002-2v-5m-1.414-9.414a2 2 0 112.828 2.828L11.828 15H9v-2.828l8.586-8.586z",
-    agentic:
-      "M9.663 17h4.673M12 3v1m6.364 1.636l-.707.707M21 12h-1M4 12H3m3.343-5.657l-.707-.707m2.828 9.9a5 5 0 117.072 0l-.548.547A3.374 3.374 0 0014 18.469V19a2 2 0 11-4 0v-.531c0-.895-.356-1.754-.988-2.386l-.548-.547z",
-    biggest:
-      "M19 11H5m14 0a2 2 0 012 2v6a2 2 0 01-2 2H5a2 2 0 01-2-2v-6a2 2 0 012-2m14 0V9a2 2 0 00-2-2M5 11V9a2 2 0 012-2m0 0V5a2 2 0 012-2h6a2 2 0 012 2v2M7 7h10",
-  };
-
-  let hoveredTooltip = $state<string | null>(null);
-  let tooltipAnchor = $state<{ x: number; y: number } | null>(null);
-
-  function showTooltip(category: string, e: MouseEvent | FocusEvent) {
-    hoveredTooltip = category;
-    const target = e.currentTarget as HTMLElement;
-    const rect = target.getBoundingClientRect();
-    tooltipAnchor = { x: rect.left + rect.width / 2, y: rect.top };
-  }
-
-  function hideTooltip() {
-    hoveredTooltip = null;
-    tooltipAnchor = null;
-  }
-</script>
-
-<div class="flex flex-col items-center justify-center gap-6 {className}">
-  <!-- Header -->
-  <div class="text-center">
-    <p class="text-xs text-exo-light-gray uppercase tracking-[0.2em] mb-1">
-      Recommended for your
-    </p>
-    <p class="text-sm text-white font-mono tracking-wide">{clusterLabel}</p>
-  </div>
-
-  <!-- Category Cards Grid -->
-  <div class="grid grid-cols-2 gap-3 w-full max-w-md">
-    {#each recommendations as rec}
-      {#if rec.model}
-        <button
-          type="button"
-          onclick={() => rec.model && onSelect(rec.model.id, rec.category)}
-          class="group relative flex flex-col items-start gap-2 p-4 rounded-lg border border-exo-medium-gray/50 bg-exo-dark-gray/50 hover:border-exo-yellow/40 hover:bg-exo-dark-gray transition-all duration-200 cursor-pointer text-left"
-        >
-          <!-- Category icon + label -->
-          <div class="flex items-center gap-2 w-full">
-            <svg
-              class="w-4 h-4 text-exo-yellow/70 group-hover:text-exo-yellow transition-colors flex-shrink-0"
-              fill="none"
-              viewBox="0 0 24 24"
-              stroke="currentColor"
-              stroke-width="1.5"
-            >
-              <path
-                stroke-linecap="round"
-                stroke-linejoin="round"
-                d={categoryIcons[rec.category]}
-              />
-            </svg>
-            <span
-              class="text-xs font-mono uppercase tracking-wider text-exo-light-gray group-hover:text-white transition-colors"
-            >
-              {rec.label}
-            </span>
-            <!-- Info tooltip -->
-            <div class="ml-auto flex-shrink-0">
-              <span
-                role="button"
-                tabindex="-1"
-                class="text-exo-light-gray/40 hover:text-exo-light-gray transition-colors cursor-help inline-flex"
-                onmouseenter={(e: MouseEvent) => showTooltip(rec.category, e)}
-                onmouseleave={() => hideTooltip()}
-                onclick={(e: MouseEvent) => {
-                  e.stopPropagation();
-                  if (hoveredTooltip === rec.category) {
-                    hideTooltip();
-                  } else {
-                    showTooltip(rec.category, e);
-                  }
-                }}
-              >
-                <svg
-                  class="w-3.5 h-3.5"
-                  fill="none"
-                  viewBox="0 0 24 24"
-                  stroke="currentColor"
-                  stroke-width="2"
-                >
-                  <circle cx="12" cy="12" r="10" />
-                  <path d="M12 16v-4m0-4h.01" />
-                </svg>
-              </span>
-            </div>
-          </div>
-
-          <!-- Model name + size -->
-          <div class="w-full">
-            <p class="text-sm text-white font-mono truncate">
-              {rec.model.base_model}
-            </p>
-            <p class="text-xs text-exo-light-gray/60 font-mono mt-0.5">
-              {formatSize(rec.model.storage_size_megabytes)}
-              {#if rec.model.quantization}
-                <span class="text-exo-light-gray/40"
-                  >&middot; {rec.model.quantization}</span
-                >
-              {/if}
-            </p>
-          </div>
-        </button>
-      {:else}
-        <!-- No model fits for this category -->
-        <div
-          class="flex flex-col items-start gap-2 p-4 rounded-lg border border-exo-medium-gray/30 bg-exo-dark-gray/30 opacity-50"
-        >
-          <div class="flex items-center gap-2">
-            <svg
-              class="w-4 h-4 text-exo-light-gray/40 flex-shrink-0"
-              fill="none"
-              viewBox="0 0 24 24"
-              stroke="currentColor"
-              stroke-width="1.5"
-            >
-              <path
-                stroke-linecap="round"
-                stroke-linejoin="round"
-                d={categoryIcons[rec.category]}
-              />
-            </svg>
-            <span
-              class="text-xs font-mono uppercase tracking-wider text-exo-light-gray/50"
-              >{rec.label}</span
-            >
-          </div>
-          <p class="text-xs text-exo-light-gray/40 font-mono">No model fits</p>
-        </div>
-      {/if}
-    {/each}
-  </div>
-
-  <!-- Add Model Button -->
-  <button
-    type="button"
-    onclick={onAddModel}
-    class="flex items-center gap-2 px-4 py-2 text-xs font-mono uppercase tracking-wider text-exo-light-gray hover:text-exo-yellow border border-exo-medium-gray/30 hover:border-exo-yellow/30 rounded-lg transition-all duration-200 cursor-pointer"
-  >
-    <svg
-      class="w-3.5 h-3.5"
-      fill="none"
-      viewBox="0 0 24 24"
-      stroke="currentColor"
-      stroke-width="2"
-    >
-      <path stroke-linecap="round" stroke-linejoin="round" d="M12 4v16m8-8H4" />
-    </svg>
-    Add Model
-  </button>
-
-  <!-- Auto hint -->
-  <p class="text-xs text-exo-light-gray/40 font-mono tracking-wide text-center">
-    Or just start typing &mdash; we'll pick the best model automatically
-  </p>
-</div>
-
-<!-- Fixed-position tooltip (escapes overflow-hidden ancestors) -->
-{#if hoveredTooltip && tooltipAnchor}
-  {@const rec = recommendations.find((r) => r.category === hoveredTooltip)}
-  {#if rec}
-    <div
-      class="fixed z-[9999] px-3 py-2 bg-exo-black border border-exo-medium-gray/50 rounded text-xs text-exo-light-gray whitespace-nowrap shadow-lg pointer-events-none"
-      style="left: {tooltipAnchor.x}px; top: {tooltipAnchor.y -
-        8}px; transform: translate(-50%, -100%);"
-    >
-      {rec.tooltip}
-    </div>
-  {/if}
-{/if}

dashboard/src/lib/components/ChatSidebar.svelte

@@ -2,6 +2,7 @@
   import {
     conversations,
     activeConversationId,
+    createConversation,
     loadConversation,
     deleteConversation,
     deleteAllConversations,
@@ -16,15 +17,9 @@
 
   interface Props {
     class?: string;
-    onNewChat?: () => void;
-    onSelectConversation?: () => void;
   }
 
-  let {
-    class: className = "",
-    onNewChat,
-    onSelectConversation,
-  }: Props = $props();
+  let { class: className = "" }: Props = $props();
 
   const conversationList = $derived(conversations());
   const activeId = $derived(activeConversationId());
@@ -47,11 +42,10 @@
   );
 
   function handleNewChat() {
-    onNewChat?.();
+    createConversation();
   }
 
   function handleSelectConversation(id: string) {
-    onSelectConversation?.();
     loadConversation(id);
   }
 

dashboard/src/lib/components/DeviceIcon.svelte

@@ -67,8 +67,8 @@
   const studioMemH = $derived((ramPercent / 100) * studioMemTotalH);
 
   // ── MacBook dimensions (same ratios as TopologyGraph) ──
-  const mbW = $derived((size * 1.6 * 0.85) / 1.15);
-  const mbH = $derived(size * 0.85);
+  const mbW = $derived(size * 1.6);
+  const mbH = $derived(size * 1.15);
   const mbX = $derived(cx - mbW / 2);
   const mbY = $derived(cy - mbH / 2);
 

dashboard/src/lib/components/HeaderNav.svelte

@@ -95,7 +95,7 @@
     {#if showHome}
       <button
         onclick={handleHome}
-        class="text-sm text-white/70 hover:text-exo-yellow transition-colors tracking-wider uppercase flex items-center gap-2 cursor-pointer"
+        class="text-sm text-exo-light-gray hover:text-exo-yellow transition-colors tracking-wider uppercase flex items-center gap-2 cursor-pointer"
         title="Back to topology view"
       >
         <svg
@@ -116,7 +116,7 @@
     {/if}
     <a
       href="/#/downloads"
-      class="text-sm text-white/70 hover:text-exo-yellow transition-colors tracking-wider uppercase flex items-center gap-2 cursor-pointer"
+      class="text-sm text-exo-light-gray hover:text-exo-yellow transition-colors tracking-wider uppercase flex items-center gap-2 cursor-pointer"
       title="View downloads overview"
     >
       {#if downloadProgress}
@@ -170,30 +170,5 @@
       {/if}
       Downloads
     </a>
-    <a
-      href="/#/settings"
-      class="text-sm text-white/70 hover:text-exo-yellow transition-colors tracking-wider uppercase flex items-center gap-2 cursor-pointer"
-      title="Settings"
-    >
-      <svg
-        class="w-4 h-4"
-        fill="none"
-        viewBox="0 0 24 24"
-        stroke="currentColor"
-        stroke-width="2"
-      >
-        <path
-          stroke-linecap="round"
-          stroke-linejoin="round"
-          d="M10.325 4.317c.426-1.756 2.924-1.756 3.35 0a1.724 1.724 0 002.573 1.066c1.543-.94 3.31.826 2.37 2.37a1.724 1.724 0 001.065 2.572c1.756.426 1.756 2.924 0 3.35a1.724 1.724 0 00-1.066 2.573c.94 1.543-.826 3.31-2.37 2.37a1.724 1.724 0 00-2.572 1.065c-.426 1.756-2.924 1.756-3.35 0a1.724 1.724 0 00-2.573-1.066c-1.543.94-3.31-.826-2.37-2.37a1.724 1.724 0 00-1.065-2.572c-1.756-.426-1.756-2.924 0-3.35a1.724 1.724 0 001.066-2.573c-.94-1.543.826-3.31 2.37-2.37.996.608 2.296.07 2.572-1.065z"
-        />
-        <path
-          stroke-linecap="round"
-          stroke-linejoin="round"
-          d="M15 12a3 3 0 11-6 0 3 3 0 016 0z"
-        />
-      </svg>
-      Settings
-    </a>
   </nav>
 </header>

dashboard/src/lib/components/ModelFilterPopover.svelte

@@ -6,7 +6,6 @@
     capabilities: string[];
     sizeRange: { min: number; max: number } | null;
     downloadedOnly: boolean;
-    readyOnly: boolean;
   }
 
   type ModelFilterPopoverProps = {
@@ -191,58 +190,34 @@
       </div>
     </div>
 
-    <!-- Availability filters -->
+    <!-- Downloaded only -->
     <div>
       <h4 class="text-xs font-mono text-white/50 mb-2">Availability</h4>
-      <div class="flex flex-wrap gap-1.5">
-        <button
-          type="button"
-          class="px-2 py-1 text-xs font-mono rounded transition-colors {filters.downloadedOnly
-            ? 'bg-green-500/20 text-green-400 border border-green-500/30'
-            : 'bg-white/5 text-white/60 hover:bg-white/10 border border-transparent'}"
-          onclick={() =>
-            onChange({ ...filters, downloadedOnly: !filters.downloadedOnly })}
+      <button
+        type="button"
+        class="px-2 py-1 text-xs font-mono rounded transition-colors {filters.downloadedOnly
+          ? 'bg-green-500/20 text-green-400 border border-green-500/30'
+          : 'bg-white/5 text-white/60 hover:bg-white/10 border border-transparent'}"
+        onclick={() =>
+          onChange({ ...filters, downloadedOnly: !filters.downloadedOnly })}
+      >
+        <svg
+          class="w-3.5 h-3.5 inline-block"
+          viewBox="0 0 24 24"
+          fill="none"
+          stroke="currentColor"
+          stroke-width="2"
+          stroke-linecap="round"
+          stroke-linejoin="round"
         >
-          <svg
-            class="w-3.5 h-3.5 inline-block"
-            viewBox="0 0 24 24"
-            fill="none"
-            stroke="currentColor"
-            stroke-width="2"
-            stroke-linecap="round"
-            stroke-linejoin="round"
-          >
-            <path
-              class="text-white/40"
-              d="M20 20a2 2 0 0 0 2-2V8a2 2 0 0 0-2-2h-7.9a2 2 0 0 1-1.69-.9L9.6 3.9A2 2 0 0 0 7.93 3H4a2 2 0 0 0-2 2v13a2 2 0 0 0 2 2Z"
-            />
-            <path class="text-green-400" d="m9 13 2 2 4-4" />
-          </svg>
-          <span class="ml-1">Downloaded</span>
-        </button>
-        <button
-          type="button"
-          class="px-2 py-1 text-xs font-mono rounded transition-colors {filters.readyOnly
-            ? 'bg-green-500/20 text-green-400 border border-green-500/30'
-            : 'bg-white/5 text-white/60 hover:bg-white/10 border border-transparent'}"
-          onclick={() =>
-            onChange({ ...filters, readyOnly: !filters.readyOnly })}
-        >
-          <svg
-            class="w-3.5 h-3.5 inline-block"
-            viewBox="0 0 24 24"
-            fill="none"
-            stroke="currentColor"
-            stroke-width="2"
-            stroke-linecap="round"
-            stroke-linejoin="round"
-          >
-            <circle cx="12" cy="12" r="10" />
-            <path d="m9 12 2 2 4-4" />
-          </svg>
-          <span class="ml-1">Ready</span>
-        </button>
-      </div>
+          <path
+            class="text-white/40"
+            d="M20 20a2 2 0 0 0 2-2V8a2 2 0 0 0-2-2h-7.9a2 2 0 0 1-1.69-.9L9.6 3.9A2 2 0 0 0 7.93 3H4a2 2 0 0 0-2 2v13a2 2 0 0 0 2 2Z"
+          />
+          <path class="text-green-400" d="m9 13 2 2 4-4" />
+        </svg>
+        <span class="ml-1">Downloaded</span>
+      </button>
     </div>
 
     <!-- Size range -->

dashboard/src/lib/components/ModelPickerGroup.svelte

@@ -41,7 +41,6 @@
     onShowInfo: (group: ModelGroup) => void;
     downloadStatusMap?: Map<string, DownloadAvailability>;
     launchedAt?: number;
-    instanceStatuses?: Record<string, { status: string; statusClass: string }>;
   };
 
   let {
@@ -57,7 +56,6 @@
     onShowInfo,
     downloadStatusMap,
     launchedAt,
-    instanceStatuses = {},
   }: ModelPickerGroupProps = $props();
 
   // Group-level download status: show if any variant is downloaded
@@ -94,12 +92,6 @@
   const anyVariantFits = $derived(
     group.variants.some((v) => canModelFit(v.id)),
   );
-  // Check if any variant has an active instance (ready, loading, downloading)
-  const anyVariantHasInstance = $derived(
-    instanceStatuses
-      ? group.variants.some((v) => instanceStatuses[v.id] != null)
-      : false,
-  );
   const groupFitStatus = $derived.by((): ModelFitStatus => {
     let hasClusterCapacityOnly = false;
     for (const variant of group.variants) {
@@ -130,35 +122,16 @@
     !group.hasMultipleVariants &&
       group.variants.some((v) => v.id === selectedModelId),
   );
-
-  // Group-level instance status: show the "best" status across all variants
-  const groupInstanceStatus = $derived.by(() => {
-    if (!instanceStatuses) return null;
-    const readyStatuses = ["READY", "LOADED", "RUNNING"];
-    const loadingStatuses = ["LOADING", "WARMING UP"];
-    let bestStatus: { status: string; statusClass: string } | null = null;
-    for (const variant of group.variants) {
-      const s = instanceStatuses[variant.id];
-      if (!s) continue;
-      if (readyStatuses.includes(s.status)) return s; // Ready is best
-      if (loadingStatuses.includes(s.status) || s.status === "DOWNLOADING") {
-        bestStatus = s;
-      }
-    }
-    return bestStatus;
-  });
 </script>
 
 <div
-  class="border-b border-white/5 last:border-b-0 {!anyVariantFits &&
-  !anyVariantHasInstance
+  class="border-b border-white/5 last:border-b-0 {!anyVariantFits
     ? 'opacity-50'
     : ''}"
 >
   <!-- Main row -->
   <div
-    class="flex items-center gap-2 px-3 py-2.5 transition-colors {anyVariantFits ||
-    anyVariantHasInstance
+    class="flex items-center gap-2 px-3 py-2.5 transition-colors {anyVariantFits
       ? 'hover:bg-white/5 cursor-pointer'
       : 'cursor-not-allowed'} {isMainSelected
       ? 'bg-exo-yellow/10 border-l-2 border-exo-yellow'
@@ -168,7 +141,7 @@
         onToggleExpand();
       } else {
         const modelId = group.variants[0]?.id;
-        if (modelId && (canModelFit(modelId) || instanceStatuses[modelId])) {
+        if (modelId && canModelFit(modelId)) {
           onSelectModel(modelId);
         }
       }
@@ -182,7 +155,7 @@
           onToggleExpand();
         } else {
           const modelId = group.variants[0]?.id;
-          if (modelId && (canModelFit(modelId) || instanceStatuses[modelId])) {
+          if (modelId && canModelFit(modelId)) {
             onSelectModel(modelId);
           }
         }
@@ -373,47 +346,6 @@
       </span>
     {/if}
 
-    <!-- Instance status badge -->
-    {#if groupInstanceStatus}
-      {#if groupInstanceStatus.status === "READY" || groupInstanceStatus.status === "LOADED" || groupInstanceStatus.status === "RUNNING"}
-        <span class="flex-shrink-0" title="Running">
-          <svg
-            class="w-3 h-3 text-green-400"
-            viewBox="0 0 12 12"
-            fill="currentColor"
-          >
-            <circle cx="6" cy="6" r="5" />
-          </svg>
-        </span>
-      {:else if groupInstanceStatus.status === "DOWNLOADING"}
-        <span class="flex-shrink-0 animate-pulse" title="Downloading">
-          <svg
-            class="w-3.5 h-3.5 text-blue-400"
-            viewBox="0 0 24 24"
-            fill="none"
-            stroke="currentColor"
-            stroke-width="2"
-            stroke-linecap="round"
-            stroke-linejoin="round"
-          >
-            <path d="M21 15v4a2 2 0 01-2 2H5a2 2 0 01-2-2v-4" />
-            <polyline points="7 10 12 15 17 10" />
-            <line x1="12" y1="15" x2="12" y2="3" />
-          </svg>
-        </span>
-      {:else if groupInstanceStatus.status === "LOADING" || groupInstanceStatus.status === "WARMING UP"}
-        <span class="flex-shrink-0 animate-pulse" title="Loading">
-          <svg
-            class="w-3 h-3 text-yellow-400"
-            viewBox="0 0 12 12"
-            fill="currentColor"
-          >
-            <circle cx="6" cy="6" r="5" />
-          </svg>
-        </span>
-      {/if}
-    {/if}
-
     <!-- Check mark if selected (single-variant) -->
     {#if isMainSelected}
       <svg
@@ -488,30 +420,20 @@
       {#each group.variants as variant}
         {@const fitStatus = getModelFitStatus(variant.id)}
         {@const modelCanFit = canModelFit(variant.id)}
-        {@const variantHasInstance = instanceStatuses[variant.id] != null}
         {@const isSelected = selectedModelId === variant.id}
-        <div
-          class="w-full flex items-center gap-3 px-3 py-2 pl-10 hover:bg-white/5 transition-colors text-left {!modelCanFit &&
-          !variantHasInstance
+        <button
+          type="button"
+          class="w-full flex items-center gap-3 px-3 py-2 pl-10 hover:bg-white/5 transition-colors text-left {!modelCanFit
             ? 'opacity-50 cursor-not-allowed'
             : 'cursor-pointer'} {isSelected
             ? 'bg-exo-yellow/10 border-l-2 border-exo-yellow'
             : 'border-l-2 border-transparent'}"
-          role="button"
-          tabindex="0"
+          disabled={!modelCanFit}
           onclick={() => {
-            if (modelCanFit || variantHasInstance) {
+            if (modelCanFit) {
               onSelectModel(variant.id);
             }
           }}
-          onkeydown={(e) => {
-            if (e.key === "Enter" || e.key === " ") {
-              e.preventDefault();
-              if (modelCanFit) {
-                onSelectModel(variant.id);
-              }
-            }
-          }}
         >
           <!-- Quantization badge -->
           <span
@@ -556,48 +478,6 @@
             {/if}
           {/if}
 
-          <!-- Instance status badge -->
-          {#if instanceStatuses[variant.id]}
-            {@const instStatus = instanceStatuses[variant.id]}
-            {#if instStatus.status === "READY" || instStatus.status === "LOADED" || instStatus.status === "RUNNING"}
-              <span class="flex-shrink-0" title="Running">
-                <svg
-                  class="w-3 h-3 text-green-400"
-                  viewBox="0 0 12 12"
-                  fill="currentColor"
-                >
-                  <circle cx="6" cy="6" r="5" />
-                </svg>
-              </span>
-            {:else if instStatus.status === "DOWNLOADING"}
-              <span class="flex-shrink-0 animate-pulse" title="Downloading">
-                <svg
-                  class="w-3.5 h-3.5 text-blue-400"
-                  viewBox="0 0 24 24"
-                  fill="none"
-                  stroke="currentColor"
-                  stroke-width="2"
-                  stroke-linecap="round"
-                  stroke-linejoin="round"
-                >
-                  <path d="M21 15v4a2 2 0 01-2 2H5a2 2 0 01-2-2v-4" />
-                  <polyline points="7 10 12 15 17 10" />
-                  <line x1="12" y1="15" x2="12" y2="3" />
-                </svg>
-              </span>
-            {:else if instStatus.status === "LOADING" || instStatus.status === "WARMING UP"}
-              <span class="flex-shrink-0 animate-pulse" title="Loading">
-                <svg
-                  class="w-3 h-3 text-yellow-400"
-                  viewBox="0 0 12 12"
-                  fill="currentColor"
-                >
-                  <circle cx="6" cy="6" r="5" />
-                </svg>
-              </span>
-            {/if}
-          {/if}
-
           <!-- Check mark if selected -->
           {#if isSelected}
             <svg
@@ -610,36 +490,7 @@
               />
             </svg>
           {/if}
-
-          <!-- Info button -->
-          <button
-            type="button"
-            class="p-1 rounded hover:bg-white/10 transition-colors flex-shrink-0"
-            onclick={(e) => {
-              e.stopPropagation();
-              onShowInfo({
-                id: variant.id,
-                name: variant.name || variant.id,
-                capabilities: group.capabilities,
-                family: group.family,
-                variants: [variant],
-                smallestVariant: variant,
-                hasMultipleVariants: false,
-              });
-            }}
-            title="View variant details"
-          >
-            <svg
-              class="w-4 h-4 text-white/30 hover:text-white/50"
-              viewBox="0 0 24 24"
-              fill="currentColor"
-            >
-              <path
-                d="M12 2C6.48 2 2 6.48 2 12s4.48 10 10 10 10-4.48 10-10S17.52 2 12 2zm1 15h-2v-6h2v6zm0-8h-2V7h2v2z"
-              />
-            </svg>
-          </button>
-        </div>
+        </button>
       {/each}
     </div>
   {/if}

dashboard/src/lib/components/ModelPickerModal.svelte

@@ -36,7 +36,6 @@
     capabilities: string[];
     sizeRange: { min: number; max: number } | null;
     downloadedOnly: boolean;
-    readyOnly: boolean;
   }
 
   interface HuggingFaceModel {
@@ -50,11 +49,6 @@
 
   type ModelFitStatus = "fits_now" | "fits_cluster_capacity" | "too_large";
 
-  export type InstanceStatus = {
-    status: string;
-    statusClass: string;
-  };
-
   type ModelPickerModalProps = {
     isOpen: boolean;
     models: ModelInfo[];
@@ -81,7 +75,6 @@
         macmon_info?: { memory?: { ram_total?: number } };
       }
     >;
-    instanceStatuses?: Record<string, InstanceStatus>;
   };
 
   let {
@@ -103,7 +96,6 @@
     usedMemoryGB,
     downloadsData,
     topologyNodes,
-    instanceStatuses = {},
   }: ModelPickerModalProps = $props();
 
   // Local state
@@ -115,7 +107,6 @@
     capabilities: [],
     sizeRange: null,
     downloadedOnly: false,
-    readyOnly: false,
   });
   let infoGroup = $state<ModelGroup | null>(null);
 
@@ -449,16 +440,6 @@
       );
     }
 
-    // Filter to ready/running models only
-    if (filters.readyOnly) {
-      result = result.filter((g) =>
-        g.variants.some((v) => {
-          const s = instanceStatuses[v.id];
-          return s && s.statusClass === "ready";
-        }),
-      );
-    }
-
     // Sort: fits-now first, then fits-cluster-capacity, then too-large
     result.sort((a, b) => {
       const getGroupFitRank = (group: ModelGroup): number => {
@@ -569,19 +550,13 @@
   }
 
   function clearFilters() {
-    filters = {
-      capabilities: [],
-      sizeRange: null,
-      downloadedOnly: false,
-      readyOnly: false,
-    };
+    filters = { capabilities: [], sizeRange: null, downloadedOnly: false };
   }
 
   const hasActiveFilters = $derived(
     filters.capabilities.length > 0 ||
       filters.sizeRange !== null ||
-      filters.downloadedOnly ||
-      filters.readyOnly,
+      filters.downloadedOnly,
   );
 </script>
 
@@ -850,7 +825,6 @@
                 onShowInfo={(g) => (infoGroup = g)}
                 downloadStatusMap={getVariantDownloadMap(group)}
                 launchedAt={recentTimestamps.get(group.variants[0]?.id ?? "")}
-                {instanceStatuses}
               />
             {/each}
           {/if}
@@ -918,7 +892,6 @@
               {onToggleFavorite}
               onShowInfo={(g) => (infoGroup = g)}
               downloadStatusMap={getVariantDownloadMap(group)}
-              {instanceStatuses}
             />
           {/each}
           <!-- Other models -->
@@ -945,7 +918,6 @@
               {onToggleFavorite}
               onShowInfo={(g) => (infoGroup = g)}
               downloadStatusMap={getVariantDownloadMap(group)}
-              {instanceStatuses}
             />
           {/each}
         {/if}
@@ -968,11 +940,6 @@
             >Downloaded</span
           >
         {/if}
-        {#if filters.readyOnly}
-          <span class="px-1.5 py-0.5 bg-green-500/20 text-green-400 rounded"
-            >Ready</span
-          >
-        {/if}
         {#if filters.sizeRange}
           <span class="px-1.5 py-0.5 bg-exo-yellow/20 text-exo-yellow rounded">
             {filters.sizeRange.min}GB - {filters.sizeRange.max}GB

dashboard/src/lib/components/index.ts

@@ -12,4 +12,3 @@ export { default as HuggingFaceResultItem } from "./HuggingFaceResultItem.svelte
 export { default as ModelFilterPopover } from "./ModelFilterPopover.svelte";
 export { default as ModelPickerGroup } from "./ModelPickerGroup.svelte";
 export { default as ModelPickerModal } from "./ModelPickerModal.svelte";
-export { default as ChatModelSelector } from "./ChatModelSelector.svelte";

dashboard/src/lib/stores/app.svelte.ts

@@ -847,10 +847,6 @@ class AppStore {
     this.thinkingEnabled = conversation.enableThinking ?? true;
     this.refreshConversationModelFromInstances();
 
-    // Sync global selection to the loaded conversation's model so reactive
-    // effects in +page.svelte can determine the correct chat launch state.
-    this.selectedChatModel = conversation.modelId || "";
-
     return true;
   }
 

dashboard/src/lib/stores/settings.svelte.ts

@@ -1,87 +0,0 @@
-/**
- * SettingsStore - Manages exo runtime settings via the /settings API.
- */
-
-export interface MemorySettings {
-  oom_prevention: boolean;
-  memory_threshold: number;
-  memory_floor_gb: number;
-}
-
-export interface GenerationSettings {
-  prefill_step_size: number;
-  max_tokens: number;
-  kv_cache_bits: 4 | 8 | null;
-}
-
-export interface ExoSettings {
-  memory: MemorySettings;
-  generation: GenerationSettings;
-}
-
-function defaultSettings(): ExoSettings {
-  return {
-    memory: {
-      oom_prevention: false,
-      memory_threshold: 0.8,
-      memory_floor_gb: 5.0,
-    },
-    generation: {
-      prefill_step_size: 4096,
-      max_tokens: 32168,
-      kv_cache_bits: null,
-    },
-  };
-}
-
-class SettingsStore {
-  settings = $state<ExoSettings>(defaultSettings());
-  loading = $state(false);
-  error = $state<string | null>(null);
-
-  async load(): Promise<void> {
-    this.loading = true;
-    this.error = null;
-    try {
-      const response = await fetch("/settings");
-      if (!response.ok) {
-        throw new Error(`Failed to fetch settings: ${response.status}`);
-      }
-      this.settings = (await response.json()) as ExoSettings;
-    } catch (err) {
-      console.error("Failed to load settings:", err);
-      this.error = err instanceof Error ? err.message : "Unknown error";
-    } finally {
-      this.loading = false;
-    }
-  }
-
-  async save(updated: ExoSettings): Promise<boolean> {
-    this.loading = true;
-    this.error = null;
-    try {
-      const response = await fetch("/settings", {
-        method: "POST",
-        headers: { "Content-Type": "application/json" },
-        body: JSON.stringify(updated),
-      });
-      if (!response.ok) {
-        throw new Error(`Failed to save settings: ${response.status}`);
-      }
-      this.settings = (await response.json()) as ExoSettings;
-      return true;
-    } catch (err) {
-      console.error("Failed to save settings:", err);
-      this.error = err instanceof Error ? err.message : "Unknown error";
-      return false;
-    } finally {
-      this.loading = false;
-    }
-  }
-
-  resetToDefaults(): ExoSettings {
-    return defaultSettings();
-  }
-}
-
-export const settingsStore = new SettingsStore();

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

@@ -412,7 +412,7 @@
                   <div>{col.label}</div>
                   {#if col.diskAvailable != null}
                     <div
-                      class="text-[9px] text-white/70 normal-case tracking-normal mt-0.5"
+                      class="text-[9px] text-exo-light-gray/60 normal-case tracking-normal mt-0.5"
                     >
                       {formatBytes(col.diskAvailable)} free
                     </div>
@@ -436,7 +436,7 @@
                       </div>
                       {#if row.prettyName}
                         <div
-                          class="text-[10px] text-white/60"
+                          class="text-[10px] text-exo-light-gray/60"
                           title={row.modelId}
                         >
                           {row.modelId}
@@ -450,7 +450,7 @@
                       title="View model details"
                     >
                       <svg
-                        class="w-4 h-4 text-white/60 hover:text-white/80"
+                        class="w-4 h-4 text-white/30 hover:text-white/60"
                         viewBox="0 0 24 24"
                         fill="currentColor"
                       >
@@ -469,11 +469,11 @@
                   <td class="px-4 py-3 text-center align-middle">
                     {#if cell.kind === "completed"}
                       <div
-                        class="flex flex-col items-center gap-1"
+                        class="flex flex-col items-center gap-0.5"
                         title="Completed ({formatBytes(cell.totalBytes)})"
                       >
                         <svg
-                          class="w-7 h-7 text-green-400"
+                          class="w-5 h-5 text-green-400"
                           viewBox="0 0 20 20"
                           fill="currentColor"
                         >
@@ -483,18 +483,18 @@
                             clip-rule="evenodd"
                           ></path>
                         </svg>
-                        <span class="text-xs text-white/70"
+                        <span class="text-[10px] text-exo-light-gray/70"
                           >{formatBytes(cell.totalBytes)}</span
                         >
                         <button
                           type="button"
-                          class="text-white/50 hover:text-red-400 transition-colors mt-0.5 cursor-pointer"
+                          class="text-exo-light-gray/40 hover:text-red-400 transition-colors mt-0.5"
                           onclick={() =>
                             deleteDownload(col.nodeId, row.modelId)}
                           title="Delete from this node"
                         >
                           <svg
-                            class="w-5 h-5"
+                            class="w-3.5 h-3.5"
                             viewBox="0 0 20 20"
                             fill="none"
                             stroke="currentColor"
@@ -517,11 +517,11 @@
                           cell.speed,
                         )} - ETA {formatEta(cell.etaMs)}"
                       >
-                        <span class="text-exo-yellow text-sm font-medium"
+                        <span class="text-exo-yellow text-xs font-medium"
                           >{clampPercent(cell.percentage).toFixed(1)}%</span
                         >
                         <div
-                          class="w-16 h-2 bg-exo-black/60 rounded-sm overflow-hidden"
+                          class="w-14 h-1.5 bg-exo-black/60 rounded-sm overflow-hidden"
                         >
                           <div
                             class="h-full bg-gradient-to-r from-exo-yellow to-exo-yellow/70 transition-all duration-300"
@@ -530,25 +530,25 @@
                             ).toFixed(1)}%"
                           ></div>
                         </div>
-                        <span class="text-[10px] text-white/70"
+                        <span class="text-[9px] text-exo-light-gray/60"
                           >{formatSpeed(cell.speed)}</span
                         >
                       </div>
                     {:else if cell.kind === "pending"}
                       <div
-                        class="flex flex-col items-center gap-1"
+                        class="flex flex-col items-center gap-0.5"
                         title={cell.downloaded > 0
-                          ? `${formatBytes(cell.downloaded)} / ${formatBytes(cell.total)} downloaded (paused)`
+                          ? `${formatBytes(cell.downloaded)} / ${formatBytes(cell.total)} downloaded`
                           : "Download pending"}
                       >
                         {#if cell.downloaded > 0 && cell.total > 0}
-                          <span class="text-white/70 text-xs"
+                          <span class="text-exo-light-gray/70 text-[10px]"
                             >{formatBytes(cell.downloaded)} / {formatBytes(
                               cell.total,
                             )}</span
                           >
                           <div
-                            class="w-full h-1.5 bg-white/10 rounded-full overflow-hidden"
+                            class="w-full h-1 bg-white/10 rounded-full overflow-hidden"
                           >
                             <div
                               class="h-full bg-exo-light-gray/40 rounded-full"
@@ -558,65 +558,21 @@
                               ).toFixed(1)}%"
                             ></div>
                           </div>
-                          {#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="Resume download on this node"
-                            >
-                              <svg
-                                class="w-5 h-5"
-                                viewBox="0 0 20 20"
-                                fill="none"
-                                stroke="currentColor"
-                                stroke-width="2"
-                              >
-                                <path
-                                  d="M10 3v10m0 0l-3-3m3 3l3-3M3 17h14"
-                                  stroke-linecap="round"
-                                  stroke-linejoin="round"
-                                ></path>
-                              </svg>
-                            </button>
-                          {:else}
-                            <span class="text-white/50 text-[10px]">paused</span
-                            >
-                          {/if}
-                        {:else 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="Start download on this node"
+                          <span class="text-exo-light-gray/40 text-[9px]"
+                            >paused</span
                           >
-                            <svg
-                              class="w-6 h-6"
-                              viewBox="0 0 20 20"
-                              fill="none"
-                              stroke="currentColor"
-                              stroke-width="2"
-                            >
-                              <path
-                                d="M10 3v10m0 0l-3-3m3 3l3-3M3 17h14"
-                                stroke-linecap="round"
-                                stroke-linejoin="round"
-                              ></path>
-                            </svg>
-                          </button>
                         {:else}
-                          <span class="text-white/40 text-sm">...</span>
+                          <span class="text-exo-light-gray/50 text-sm">...</span
+                          >
                         {/if}
                       </div>
                     {:else if cell.kind === "failed"}
                       <div
-                        class="flex flex-col items-center gap-1"
+                        class="flex flex-col items-center gap-0.5"
                         title="Download failed"
                       >
                         <svg
-                          class="w-7 h-7 text-red-400"
+                          class="w-5 h-5 text-red-400"
                           viewBox="0 0 20 20"
                           fill="currentColor"
                         >
@@ -629,13 +585,13 @@
                         {#if row.shardMetadata}
                           <button
                             type="button"
-                            class="text-white/50 hover:text-exo-yellow transition-colors cursor-pointer"
+                            class="text-exo-light-gray/40 hover:text-exo-yellow transition-colors"
                             onclick={() =>
                               startDownload(col.nodeId, row.shardMetadata!)}
                             title="Retry download on this node"
                           >
                             <svg
-                              class="w-5 h-5"
+                              class="w-3.5 h-3.5"
                               viewBox="0 0 20 20"
                               fill="none"
                               stroke="currentColor"
@@ -661,13 +617,13 @@
                         {#if row.shardMetadata}
                           <button
                             type="button"
-                            class="text-white/50 hover:text-exo-yellow transition-colors mt-0.5 opacity-0 group-hover:opacity-100 cursor-pointer"
+                            class="text-exo-light-gray/30 hover:text-exo-yellow transition-colors mt-0.5 opacity-0 group-hover:opacity-100"
                             onclick={() =>
                               startDownload(col.nodeId, row.shardMetadata!)}
                             title="Download to this node"
                           >
                             <svg
-                              class="w-5 h-5"
+                              class="w-3.5 h-3.5"
                               viewBox="0 0 20 20"
                               fill="none"
                               stroke="currentColor"

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

@@ -1,193 +0,0 @@
-<script lang="ts">
-  import { onMount } from "svelte";
-  import { fade } from "svelte/transition";
-  import HeaderNav from "$lib/components/HeaderNav.svelte";
-  import { settingsStore, type ExoSettings } from "$lib/stores/settings.svelte";
-  import { addToast } from "$lib/stores/toast.svelte";
-
-  let draft = $state<ExoSettings | null>(null);
-  const loading = $derived(settingsStore.loading);
-
-  onMount(async () => {
-    await settingsStore.load();
-    draft = structuredClone(settingsStore.settings);
-  });
-
-  async function handleSave() {
-    if (!draft) return;
-    const ok = await settingsStore.save(draft);
-    if (ok) {
-      addToast({ type: "success", message: "Settings saved" });
-    } else {
-      addToast({ type: "error", message: settingsStore.error ?? "Failed to save settings" });
-    }
-  }
-
-  function handleReset() {
-    draft = settingsStore.resetToDefaults();
-  }
-
-  const KV_OPTIONS: { label: string; value: 4 | 8 | null }[] = [
-    { label: "None (full precision)", value: null },
-    { label: "4-bit", value: 4 },
-    { label: "8-bit", value: 8 },
-  ];
-</script>
-
-<HeaderNav showHome={true} />
-
-{#if draft}
-  <div class="min-h-screen bg-background text-foreground" in:fade={{ duration: 200 }}>
-    <div class="max-w-2xl mx-auto px-6 py-8">
-      <h1 class="text-2xl font-bold text-exo-yellow tracking-wider uppercase mb-8">Settings</h1>
-
-      <!-- Memory / Safety -->
-      <section class="mb-10">
-        <h2 class="text-sm font-semibold text-white/50 tracking-widest uppercase mb-4">Memory / Safety</h2>
-        <div class="space-y-5">
-          <!-- OOM Prevention Toggle -->
-          <div class="flex items-center justify-between">
-            <div>
-              <div class="text-sm text-white/90">OOM Prevention</div>
-              <div class="text-xs text-white/40 mt-0.5">Stop generation when memory is low</div>
-            </div>
-            <button
-              onclick={() => { if (draft) draft.memory.oom_prevention = !draft.memory.oom_prevention; }}
-              class="relative w-11 h-6 rounded-full transition-colors duration-200 cursor-pointer {draft.memory.oom_prevention ? 'bg-exo-yellow' : 'bg-exo-medium-gray'}"
-              role="switch"
-              aria-checked={draft.memory.oom_prevention}
-            >
-              <span
-                class="absolute top-0.5 left-0.5 w-5 h-5 rounded-full bg-white shadow transition-transform duration-200 {draft.memory.oom_prevention ? 'translate-x-5' : 'translate-x-0'}"
-              ></span>
-            </button>
-          </div>
-
-          <!-- Memory Threshold Slider -->
-          <div>
-            <div class="flex items-center justify-between mb-1.5">
-              <div>
-                <div class="text-sm text-white/90">Memory Threshold</div>
-                <div class="text-xs text-white/40 mt-0.5">KV cache eviction triggers above this level</div>
-              </div>
-              <span class="text-sm font-mono text-exo-yellow">{(draft.memory.memory_threshold * 100).toFixed(0)}%</span>
-            </div>
-            <input
-              type="range"
-              min="0.5"
-              max="0.99"
-              step="0.01"
-              bind:value={draft.memory.memory_threshold}
-              class="w-full h-1.5 rounded-full appearance-none cursor-pointer bg-exo-medium-gray accent-exo-yellow"
-            />
-          </div>
-
-          <!-- Memory Floor -->
-          <div>
-            <div class="flex items-center justify-between mb-1.5">
-              <div>
-                <div class="text-sm text-white/90">Memory Floor</div>
-                <div class="text-xs text-white/40 mt-0.5">Minimum free memory to reserve (GB)</div>
-              </div>
-              <span class="text-sm font-mono text-exo-yellow">{draft.memory.memory_floor_gb.toFixed(1)} GB</span>
-            </div>
-            <input
-              type="number"
-              min="0"
-              max="64"
-              step="0.5"
-              bind:value={draft.memory.memory_floor_gb}
-              class="w-full bg-exo-medium-gray border border-exo-light-gray/20 rounded px-3 py-1.5 text-sm text-white/90 font-mono focus:outline-none focus:border-exo-yellow/50"
-            />
-          </div>
-        </div>
-      </section>
-
-      <!-- Generation / Performance -->
-      <section class="mb-10">
-        <h2 class="text-sm font-semibold text-white/50 tracking-widest uppercase mb-4">Generation / Performance</h2>
-        <div class="space-y-5">
-          <!-- Prefill Step Size -->
-          <div>
-            <div class="flex items-center justify-between mb-1.5">
-              <div>
-                <div class="text-sm text-white/90">Prefill Step Size</div>
-                <div class="text-xs text-white/40 mt-0.5">Token chunk size during prompt processing</div>
-              </div>
-              <span class="text-sm font-mono text-exo-yellow">{draft.generation.prefill_step_size.toLocaleString()}</span>
-            </div>
-            <input
-              type="number"
-              min="128"
-              max="32768"
-              step="128"
-              bind:value={draft.generation.prefill_step_size}
-              class="w-full bg-exo-medium-gray border border-exo-light-gray/20 rounded px-3 py-1.5 text-sm text-white/90 font-mono focus:outline-none focus:border-exo-yellow/50"
-            />
-          </div>
-
-          <!-- Max Tokens -->
-          <div>
-            <div class="flex items-center justify-between mb-1.5">
-              <div>
-                <div class="text-sm text-white/90">Max Tokens</div>
-                <div class="text-xs text-white/40 mt-0.5">Maximum generation length per response</div>
-              </div>
-              <span class="text-sm font-mono text-exo-yellow">{draft.generation.max_tokens.toLocaleString()}</span>
-            </div>
-            <input
-              type="number"
-              min="1"
-              max="131072"
-              step="1024"
-              bind:value={draft.generation.max_tokens}
-              class="w-full bg-exo-medium-gray border border-exo-light-gray/20 rounded px-3 py-1.5 text-sm text-white/90 font-mono focus:outline-none focus:border-exo-yellow/50"
-            />
-          </div>
-
-          <!-- KV Cache Bits -->
-          <div>
-            <div class="mb-1.5">
-              <div class="text-sm text-white/90">KV Cache Quantization</div>
-              <div class="text-xs text-white/40 mt-0.5">Lower bits save memory at slight quality cost</div>
-            </div>
-            <select
-              bind:value={draft.generation.kv_cache_bits}
-              class="w-full bg-exo-medium-gray border border-exo-light-gray/20 rounded px-3 py-1.5 text-sm text-white/90 font-mono focus:outline-none focus:border-exo-yellow/50 cursor-pointer"
-            >
-              {#each KV_OPTIONS as opt}
-                <option value={opt.value}>{opt.label}</option>
-              {/each}
-            </select>
-          </div>
-        </div>
-      </section>
-
-      <!-- Action Buttons -->
-      <div class="flex items-center gap-3">
-        <button
-          onclick={handleSave}
-          disabled={loading}
-          class="px-5 py-2 rounded text-sm font-semibold tracking-wider uppercase transition-colors cursor-pointer
-            bg-exo-yellow text-exo-black hover:bg-exo-yellow-darker
-            disabled:opacity-50 disabled:cursor-not-allowed"
-        >
-          {loading ? "Saving..." : "Save"}
-        </button>
-        <button
-          onclick={handleReset}
-          disabled={loading}
-          class="px-5 py-2 rounded text-sm font-semibold tracking-wider uppercase transition-colors cursor-pointer
-            border border-exo-light-gray/30 text-white/70 hover:border-exo-yellow/50 hover:text-exo-yellow
-            disabled:opacity-50 disabled:cursor-not-allowed"
-        >
-          Reset to Defaults
-        </button>
-      </div>
-    </div>
-  </div>
-{:else}
-  <div class="min-h-screen bg-background flex items-center justify-center">
-    <div class="text-white/40 text-sm">Loading settings...</div>
-  </div>
-{/if}

nix/mlx.nix

@@ -41,7 +41,7 @@ let
 
   mlx = stdenv.mkDerivation rec {
     pname = "mlx";
-    version = let v = "0.30.7.dev20260225+257d5692"; in
+    version = let v = "0.30.7.dev20260220+13998a05"; in
       assert v == uvLockMlxVersion || throw "MLX version mismatch: nix/mlx.nix has ${v} but uv.lock has ${uvLockMlxVersion}. Update both the version and hash in nix/mlx.nix.";
       v;
     pyproject = true;
@@ -49,8 +49,8 @@ let
     src = fetchFromGitHub {
       owner = "rltakashige";
       repo = "mlx-jaccl-fix-small-recv";
-      rev = "257d5692fc7af6bba3b8afaeb63c549b7d1e43d5";
-      hash = "sha256-GosFIWxIB48Egb1MqJrR3xhsUsQeWdRk5rV93USY6wQ=";
+      rev = "13998a054715edcdc93618fb1496c79c7c25ff7c";
+      hash = "sha256-fAqA3hFwNBx7FcoGnhQsIFpAIRbC2EerACm4Fvne0Cc=";
     };
 
     patches = [

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "exo"
-version = "0.3.68"
+version = "0.3.0"
 description = "Exo"
 readme = "README.md"
 requires-python = ">=3.13"

rust/exo_pyo3_bindings/exo_pyo3_bindings.pyi

@@ -104,12 +104,6 @@ class NetworkingHandle:
         will sleep until a message is sent.
         """
 
-@typing.final
-class MessageTooLargeError(builtins.Exception):
-    def __new__(cls, *args: typing.Any) -> MessageTooLargeError: ...
-    def __repr__(self) -> builtins.str: ...
-    def __str__(self) -> builtins.str: ...
-
 @typing.final
 class NoPeersSubscribedToTopicError(builtins.Exception):
     def __new__(cls, *args: typing.Any) -> NoPeersSubscribedToTopicError: ...

rust/exo_pyo3_bindings/src/networking.rs

@@ -98,37 +98,6 @@ mod exception {
             Self::MSG.to_string()
         }
     }
-
-    #[gen_stub_pyclass]
-    #[pyclass(frozen, extends=PyException, name="MessageTooLargeError")]
-    pub struct PyMessageTooLargeError {}
-
-    impl PyMessageTooLargeError {
-        const MSG: &'static str = "Gossipsub message exceeds max_transmit_size. Reduce prompt length or increase the limit.";
-
-        pub(crate) fn new_err() -> PyErr {
-            PyErr::new::<Self, _>(())
-        }
-    }
-
-    #[gen_stub_pymethods]
-    #[pymethods]
-    impl PyMessageTooLargeError {
-        #[new]
-        #[pyo3(signature = (*args))]
-        #[allow(unused_variables)]
-        pub(crate) fn new(args: &Bound<'_, PyTuple>) -> Self {
-            Self {}
-        }
-
-        fn __repr__(&self) -> String {
-            format!("MessageTooLargeError(\"{}\")", Self::MSG)
-        }
-
-        fn __str__(&self) -> String {
-            Self::MSG.to_string()
-        }
-    }
 }
 
 /// Connection or disconnection event discriminant type.
@@ -231,8 +200,6 @@ async fn networking_task(
                             Err(exception::PyNoPeersSubscribedToTopicError::new_err())
                         } else if let Err(PublishError::AllQueuesFull(_)) = result {
                             Err(exception::PyAllQueuesFullError::new_err())
-                        } else if let Err(PublishError::MessageTooLarge) = result {
-                            Err(exception::PyMessageTooLargeError::new_err())
                         } else {
                             result.pyerr()
                         };
@@ -594,7 +561,6 @@ impl PyNetworkingHandle {
 pub fn networking_submodule(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<exception::PyNoPeersSubscribedToTopicError>()?;
     m.add_class::<exception::PyAllQueuesFullError>()?;
-    m.add_class::<exception::PyMessageTooLargeError>()?;
 
     m.add_class::<PyConnectionUpdateType>()?;
     m.add_class::<PyConnectionUpdate>()?;

src/exo/download/coordinator.py

@@ -1,9 +1,11 @@
 import asyncio
+import socket
 from dataclasses import dataclass, field
 from random import random
 
 import anyio
 from anyio import current_time
+from anyio.abc import TaskGroup
 from loguru import logger
 
 from exo.download.download_utils import (
@@ -39,7 +41,6 @@ from exo.shared.types.worker.downloads import (
 )
 from exo.shared.types.worker.shards import PipelineShardMetadata, ShardMetadata
 from exo.utils.channels import Receiver, Sender, channel
-from exo.utils.task_group import TaskGroup
 
 
 @dataclass
@@ -65,13 +66,15 @@ class DownloadCoordinator:
     # Internal event channel for forwarding (initialized in __post_init__)
     event_sender: Sender[Event] = field(init=False)
     event_receiver: Receiver[Event] = field(init=False)
-    _tg: TaskGroup = field(init=False, default_factory=TaskGroup)
+    _tg: TaskGroup = field(init=False, default_factory=anyio.create_task_group)
 
     # Per-model throttle for download progress events
     _last_progress_time: dict[ModelId, float] = field(default_factory=dict)
 
     def __post_init__(self) -> None:
         self.event_sender, self.event_receiver = channel[Event]()
+        if self.offline:
+            self.shard_downloader.set_internet_connection(False)
         self.shard_downloader.on_progress(self._download_progress_callback)
 
     def _model_dir(self, model_id: ModelId) -> str:
@@ -120,6 +123,8 @@ class DownloadCoordinator:
         logger.info(
             f"Starting DownloadCoordinator{' (offline mode)' if self.offline else ''}"
         )
+        if not self.offline:
+            self._test_internet_connection()
         try:
             async with self._tg as tg:
                 tg.start_soon(self._command_processor)
@@ -127,12 +132,42 @@ class DownloadCoordinator:
                 tg.start_soon(self._emit_existing_download_progress)
                 tg.start_soon(self._resend_out_for_delivery)
                 tg.start_soon(self._clear_ofd)
+                if not self.offline:
+                    tg.start_soon(self._check_internet_connection)
         finally:
             for task in self.active_downloads.values():
                 task.cancel()
 
+    def _test_internet_connection(self) -> None:
+        # Try multiple endpoints since some ISPs/networks block specific IPs
+        for host in ("1.1.1.1", "8.8.8.8", "1.0.0.1"):
+            try:
+                socket.create_connection((host, 443), timeout=3).close()
+                self.shard_downloader.set_internet_connection(True)
+                logger.debug(f"Internet connectivity: True (via {host})")
+                return
+            except OSError:
+                continue
+        self.shard_downloader.set_internet_connection(False)
+        logger.debug("Internet connectivity: False")
+
+    async def _check_internet_connection(self) -> None:
+        first_connection = True
+        while True:
+            await asyncio.sleep(10)
+
+            # Assume that internet connection is set to False on 443 errors.
+            if self.shard_downloader.internet_connection:
+                continue
+
+            self._test_internet_connection()
+
+            if first_connection and self.shard_downloader.internet_connection:
+                first_connection = False
+                self._tg.start_soon(self._emit_existing_download_progress)
+
     def shutdown(self) -> None:
-        self._tg.cancel_tasks()
+        self._tg.cancel_scope.cancel()
 
     # directly copied from worker
     async def _resend_out_for_delivery(self) -> None:

src/exo/download/download_utils.py

@@ -110,20 +110,54 @@ def map_repo_download_progress_to_download_progress_data(
     )
 
 
+def _resolve_hf_hub_model(search_dir: Path, normalized: str) -> Path | None:
+    """Try to find a model in HuggingFace Hub cache format.
+
+    HF Hub stores models as ``models--<org>--<name>/snapshots/<commit>/``
+    with symlinks to ``../../blobs/``.  The active commit is read from
+    ``refs/main``.
+    """
+    hf_model_dir = search_dir / f"models--{normalized}"
+    if not hf_model_dir.is_dir():
+        return None
+    # Resolve ref → snapshot
+    ref_file = hf_model_dir / "refs" / "main"
+    if ref_file.is_file():
+        commit_hash = ref_file.read_text().strip()
+        snapshot = hf_model_dir / "snapshots" / commit_hash
+        if snapshot.is_dir():
+            return snapshot
+    # Fallback: use latest snapshot by mtime
+    snapshots_dir = hf_model_dir / "snapshots"
+    if snapshots_dir.is_dir():
+        snapshots = sorted(
+            snapshots_dir.iterdir(), key=lambda p: p.stat().st_mtime, reverse=True
+        )
+        if snapshots:
+            return snapshots[0]
+    return None
+
+
 def resolve_model_in_path(model_id: ModelId) -> Path | None:
     """Search EXO_MODELS_PATH directories for a pre-existing model.
 
-    Checks each directory for the normalized name (org--model).  A candidate
-    is only returned if ``is_model_directory_complete`` confirms all weight
-    files are present.
+    Checks each directory for the normalized name (org--model) and the
+    HuggingFace Hub cache format (models--org--model/snapshots/<ref>/).
+    A candidate is only returned if ``is_model_directory_complete``
+    confirms all weight files are present.
     """
     if EXO_MODELS_PATH is None:
         return None
     normalized = model_id.normalize()
     for search_dir in EXO_MODELS_PATH:
+        # Try direct format: <dir>/<org--name>/
         candidate = search_dir / normalized
         if candidate.is_dir() and is_model_directory_complete(candidate):
             return candidate
+        # Try HF Hub cache format: <dir>/models--<org--name>/snapshots/<ref>/
+        hf_candidate = _resolve_hf_hub_model(search_dir, normalized)
+        if hf_candidate is not None and is_model_directory_complete(hf_candidate):
+            return hf_candidate
     return None
 
 
@@ -823,7 +857,6 @@ async def download_shard(
 
     for file in filtered_file_list:
         downloaded_bytes = await get_downloaded_size(target_dir / file.path)
-        final_file_exists = await aios.path.exists(target_dir / file.path)
         file_progress[file.path] = RepoFileDownloadProgress(
             repo_id=shard.model_card.model_id,
             repo_revision=revision,
@@ -833,9 +866,7 @@ async def download_shard(
             total=Memory.from_bytes(file.size or 0),
             speed=0,
             eta=timedelta(0),
-            status="complete"
-            if final_file_exists and downloaded_bytes == file.size
-            else "not_started",
+            status="complete" if downloaded_bytes == file.size else "not_started",
             start_time=time.time(),
         )
 

src/exo/download/impl_shard_downloader.py

@@ -15,13 +15,9 @@ from exo.shared.types.worker.shards import (
 )
 
 
-def exo_shard_downloader(
-    max_parallel_downloads: int = 8, offline: bool = False
-) -> ShardDownloader:
+def exo_shard_downloader(max_parallel_downloads: int = 8) -> ShardDownloader:
     return SingletonShardDownloader(
-        CachedShardDownloader(
-            ResumableShardDownloader(max_parallel_downloads, offline=offline)
-        )
+        CachedShardDownloader(ResumableShardDownloader(max_parallel_downloads))
     )
 
 
@@ -54,6 +50,10 @@ class SingletonShardDownloader(ShardDownloader):
         self.shard_downloader = shard_downloader
         self.active_downloads: dict[ShardMetadata, asyncio.Task[Path]] = {}
 
+    def set_internet_connection(self, value: bool) -> None:
+        self.internet_connection = value
+        self.shard_downloader.set_internet_connection(value)
+
     def on_progress(
         self,
         callback: Callable[[ShardMetadata, RepoDownloadProgress], Awaitable[None]],
@@ -90,6 +90,10 @@ class CachedShardDownloader(ShardDownloader):
         self.shard_downloader = shard_downloader
         self.cache: dict[tuple[str, ShardMetadata], Path] = {}
 
+    def set_internet_connection(self, value: bool) -> None:
+        self.internet_connection = value
+        self.shard_downloader.set_internet_connection(value)
+
     def on_progress(
         self,
         callback: Callable[[ShardMetadata, RepoDownloadProgress], Awaitable[None]],
@@ -119,9 +123,8 @@ class CachedShardDownloader(ShardDownloader):
 
 
 class ResumableShardDownloader(ShardDownloader):
-    def __init__(self, max_parallel_downloads: int = 8, offline: bool = False):
+    def __init__(self, max_parallel_downloads: int = 8):
         self.max_parallel_downloads = max_parallel_downloads
-        self.offline = offline
         self.on_progress_callbacks: list[
             Callable[[ShardMetadata, RepoDownloadProgress], Awaitable[None]]
         ] = []
@@ -148,7 +151,8 @@ class ResumableShardDownloader(ShardDownloader):
             self.on_progress_wrapper,
             max_parallel_downloads=self.max_parallel_downloads,
             allow_patterns=allow_patterns,
-            skip_internet=self.offline,
+            skip_internet=not self.internet_connection,
+            on_connection_lost=lambda: self.set_internet_connection(False),
         )
         return target_dir
 
@@ -164,7 +168,8 @@ class ResumableShardDownloader(ShardDownloader):
                 shard,
                 self.on_progress_wrapper,
                 skip_download=True,
-                skip_internet=self.offline,
+                skip_internet=not self.internet_connection,
+                on_connection_lost=lambda: self.set_internet_connection(False),
             )
 
         semaphore = asyncio.Semaphore(self.max_parallel_downloads)
@@ -193,6 +198,7 @@ class ResumableShardDownloader(ShardDownloader):
             shard,
             self.on_progress_wrapper,
             skip_download=True,
-            skip_internet=self.offline,
+            skip_internet=not self.internet_connection,
+            on_connection_lost=lambda: self.set_internet_connection(False),
         )
         return progress

src/exo/download/shard_downloader.py

@@ -16,6 +16,11 @@ from exo.shared.types.worker.shards import (
 
 # TODO: the PipelineShardMetadata getting reinstantiated is a bit messy. Should this be a classmethod?
 class ShardDownloader(ABC):
+    internet_connection: bool = False
+
+    def set_internet_connection(self, value: bool) -> None:
+        self.internet_connection = value
+
     @abstractmethod
     async def ensure_shard(
         self, shard: ShardMetadata, config_only: bool = False

src/exo/main.py

@@ -7,6 +7,7 @@ from dataclasses import dataclass, field
 from typing import Self
 
 import anyio
+from anyio.abc import TaskGroup
 from loguru import logger
 from pydantic import PositiveInt
 
@@ -22,7 +23,6 @@ from exo.shared.logging import logger_cleanup, logger_setup
 from exo.shared.types.common import NodeId, SessionId
 from exo.utils.channels import Receiver, channel
 from exo.utils.pydantic_ext import CamelCaseModel
-from exo.utils.task_group import TaskGroup
 from exo.worker.main import Worker
 
 
@@ -38,7 +38,7 @@ class Node:
 
     node_id: NodeId
     offline: bool
-    _tg: TaskGroup = field(init=False, default_factory=TaskGroup)
+    _tg: TaskGroup = field(init=False, default_factory=anyio.create_task_group)
 
     @classmethod
     async def create(cls, args: "Args") -> Self:
@@ -60,7 +60,7 @@ class Node:
             download_coordinator = DownloadCoordinator(
                 node_id,
                 session_id,
-                exo_shard_downloader(offline=args.offline),
+                exo_shard_downloader(),
                 download_command_receiver=router.receiver(topics.DOWNLOAD_COMMANDS),
                 local_event_sender=router.sender(topics.LOCAL_EVENTS),
                 offline=args.offline,
@@ -149,11 +149,11 @@ class Node:
 
     def shutdown(self):
         # if this is our second call to shutdown, just sys.exit
-        if self._tg.cancel_called():
+        if self._tg.cancel_scope.cancel_called:
             import sys
 
             sys.exit(1)
-        self._tg.cancel_tasks()
+        self._tg.cancel_scope.cancel()
 
     async def _elect_loop(self):
         with self.election_result_receiver as results:
@@ -211,7 +211,7 @@ class Node:
                         self.download_coordinator = DownloadCoordinator(
                             self.node_id,
                             result.session_id,
-                            exo_shard_downloader(offline=self.offline),
+                            exo_shard_downloader(),
                             download_command_receiver=self.router.receiver(
                                 topics.DOWNLOAD_COMMANDS
                             ),
@@ -252,7 +252,7 @@ def main():
     target = min(max(soft, 65535), hard)
     resource.setrlimit(resource.RLIMIT_NOFILE, (target, hard))
 
-    mp.set_start_method("spawn", force=True)
+    mp.set_start_method("spawn")
     # TODO: Refactor the current verbosity system
     logger_setup(EXO_LOG, args.verbosity)
     logger.info("Starting EXO")
@@ -270,16 +270,9 @@ def main():
         logger.info("FAST_SYNCH forced OFF")
 
     node = anyio.run(Node.create, args)
-    try:
-        anyio.run(node.run)
-    except BaseException as exception:
-        logger.opt(exception=exception).critical(
-            "EXO terminated due to unhandled exception"
-        )
-        raise
-    finally:
-        logger.info("EXO Shutdown complete")
-        logger_cleanup()
+    anyio.run(node.run)
+    logger.info("EXO Shutdown complete")
+    logger_cleanup()
 
 
 class Args(CamelCaseModel):
@@ -290,7 +283,7 @@ class Args(CamelCaseModel):
     tb_only: bool = False
     no_worker: bool = False
     no_downloads: bool = False
-    offline: bool = os.getenv("EXO_OFFLINE", "false").lower() == "true"
+    offline: bool = False
     fast_synch: bool | None = None  # None = auto, True = force on, False = force off
 
     @classmethod
@@ -341,7 +334,6 @@ class Args(CamelCaseModel):
         parser.add_argument(
             "--offline",
             action="store_true",
-            default=os.getenv("EXO_OFFLINE", "false").lower() == "true",
             help="Run in offline/air-gapped mode: skip internet checks, use only pre-staged local models",
         )
         fast_synch_group = parser.add_mutually_exclusive_group()

src/exo/master/api.py

@@ -11,7 +11,8 @@ from typing import Annotated, Literal, cast
 from uuid import uuid4
 
 import anyio
-from anyio import BrokenResourceError
+from anyio import BrokenResourceError, create_task_group
+from anyio.abc import TaskGroup
 from fastapi import FastAPI, File, Form, HTTPException, Query, Request, UploadFile
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import FileResponse, JSONResponse, StreamingResponse
@@ -166,13 +167,6 @@ from exo.shared.types.openai_responses import (
     ResponsesRequest,
     ResponsesResponse,
 )
-from exo.shared.types.settings import (
-    ExoSettings,
-    load_settings,
-)
-from exo.shared.types.settings import (
-    save_settings as save_settings_to_file,
-)
 from exo.shared.types.state import State
 from exo.shared.types.worker.downloads import DownloadCompleted
 from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
@@ -180,7 +174,6 @@ 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.event_buffer import OrderedBuffer
-from exo.utils.task_group import TaskGroup
 
 _API_EVENT_LOG_DIR = EXO_EVENT_LOG_DIR / "api"
 ONBOARDING_COMPLETE_FILE = EXO_CACHE_HOME / "onboarding_complete"
@@ -259,7 +252,7 @@ class API:
             CommandId, Sender[ImageChunk | ErrorChunk]
         ] = {}
         self._image_store = ImageStore(EXO_IMAGE_CACHE_DIR)
-        self._tg: TaskGroup = TaskGroup()
+        self._tg: TaskGroup | None = None
 
     def reset(self, new_session_id: SessionId, result_clock: int):
         logger.info("Resetting API State")
@@ -356,8 +349,6 @@ class API:
         self.app.get("/v1/traces/{task_id}/raw")(self.get_trace_raw)
         self.app.get("/onboarding")(self.get_onboarding)
         self.app.post("/onboarding")(self.complete_onboarding)
-        self.app.get("/settings")(self.get_settings)
-        self.app.post("/settings")(self.save_settings)
 
     async def place_instance(self, payload: PlaceInstanceParams):
         command = PlaceInstance(
@@ -1600,7 +1591,8 @@ class API:
         shutdown_ev = anyio.Event()
 
         try:
-            async with self._tg as tg:
+            async with create_task_group() as tg:
+                self._tg = tg
                 logger.info("Starting API")
                 tg.start_soon(self._apply_state)
                 tg.start_soon(self._pause_on_new_election)
@@ -1834,13 +1826,3 @@ class API:
         ONBOARDING_COMPLETE_FILE.parent.mkdir(parents=True, exist_ok=True)
         ONBOARDING_COMPLETE_FILE.write_text("true")
         return JSONResponse({"completed": True})
-
-    async def get_settings(self) -> JSONResponse:
-        settings = load_settings()
-        return JSONResponse(settings.model_dump())
-
-    async def save_settings(self, request: Request) -> JSONResponse:
-        body = cast(object, await request.json())
-        settings = ExoSettings.model_validate(body)
-        save_settings_to_file(settings)
-        return JSONResponse(settings.model_dump())

src/exo/master/main.py

@@ -1,6 +1,7 @@
 from datetime import datetime, timedelta, timezone
 
 import anyio
+from anyio.abc import TaskGroup
 from loguru import logger
 
 from exo.master.event_log import DiskEventLog
@@ -62,7 +63,6 @@ from exo.shared.types.tasks import (
 from exo.shared.types.worker.instances import InstanceId
 from exo.utils.channels import Receiver, Sender, channel
 from exo.utils.event_buffer import MultiSourceBuffer
-from exo.utils.task_group import TaskGroup
 
 
 class Master:
@@ -77,7 +77,7 @@ class Master:
         download_command_sender: Sender[ForwarderDownloadCommand],
     ):
         self.state = State()
-        self._tg: TaskGroup = TaskGroup()
+        self._tg: TaskGroup = anyio.create_task_group()
         self.node_id = node_id
         self.session_id = session_id
         self.command_task_mapping: dict[CommandId, TaskId] = {}
@@ -116,7 +116,7 @@ class Master:
 
     async def shutdown(self):
         logger.info("Stopping Master")
-        self._tg.cancel_tasks()
+        self._tg.cancel_scope.cancel()
 
     async def _command_processor(self) -> None:
         with self.command_receiver as commands:

src/exo/master/placement.py

@@ -112,11 +112,7 @@ def place_instance(
         cycle for cycle in smallest_cycles if topology.is_rdma_cycle(cycle)
     ]
 
-    if command.instance_meta == InstanceMeta.MlxJaccl:
-        if not smallest_rdma_cycles:
-            raise ValueError(
-                "Requested RDMA (MlxJaccl) but no RDMA-connected cycles available"
-            )
+    if command.instance_meta == InstanceMeta.MlxJaccl and smallest_rdma_cycles != []:
         smallest_cycles = smallest_rdma_cycles
 
     cycles_with_leaf_nodes: list[Cycle] = [
@@ -133,11 +129,6 @@ def place_instance(
         ),
     )
 
-    # Single-node: force Pipeline/Ring (Tensor and Jaccl require multi-node)
-    if len(selected_cycle) == 1:
-        command.instance_meta = InstanceMeta.MlxRing
-        command.sharding = Sharding.Pipeline
-
     shard_assignments = get_shard_assignments(
         command.model_card, selected_cycle, command.sharding, node_memory
     )
@@ -147,6 +138,9 @@ def place_instance(
     instance_id = InstanceId()
     target_instances = dict(deepcopy(current_instances))
 
+    if len(selected_cycle) == 1:
+        command.instance_meta = InstanceMeta.MlxRing
+
     match command.instance_meta:
         case InstanceMeta.MlxJaccl:
             # TODO(evan): shard assignments should contain information about ranks, this is ugly

src/exo/master/tests/test_placement.py

@@ -14,12 +14,10 @@ from exo.shared.models.model_cards import ModelCard, ModelId, ModelTask
 from exo.shared.topology import Topology
 from exo.shared.types.commands import PlaceInstance
 from exo.shared.types.common import CommandId, NodeId
-from exo.shared.types.events import InstanceCreated, InstanceDeleted, TaskStatusUpdated
+from exo.shared.types.events import InstanceCreated, InstanceDeleted
 from exo.shared.types.memory import Memory
 from exo.shared.types.multiaddr import Multiaddr
 from exo.shared.types.profiling import NetworkInterfaceInfo, NodeNetworkInfo
-from exo.shared.types.tasks import TaskId, TaskStatus, TextGeneration
-from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
 from exo.shared.types.topology import Connection, SocketConnection
 from exo.shared.types.worker.instances import (
     Instance,
@@ -458,117 +456,3 @@ def test_tensor_rdma_backend_connectivity_matrix(
         else:
             ip_part = coordinator.split(":")[0]
             assert len(ip_part.split(".")) == 4
-
-
-def _make_task(
-    instance_id: InstanceId,
-    status: TaskStatus = TaskStatus.Running,
-) -> TextGeneration:
-    return TextGeneration(
-        task_id=TaskId(),
-        task_status=status,
-        instance_id=instance_id,
-        command_id=CommandId(),
-        task_params=TextGenerationTaskParams(
-            model=ModelId("test-model"),
-            input=[InputMessage(role="user", content="hello")],
-        ),
-    )
-
-
-def test_get_transition_events_delete_instance_cancels_running_tasks(
-    instance: Instance,
-):
-    # arrange
-    instance_id = InstanceId()
-    current_instances: dict[InstanceId, Instance] = {instance_id: instance}
-    target_instances: dict[InstanceId, Instance] = {}
-    task = _make_task(instance_id, TaskStatus.Running)
-    tasks = {task.task_id: task}
-
-    # act
-    events = get_transition_events(current_instances, target_instances, tasks)
-
-    # assert – cancellation event should come before the deletion event
-    assert len(events) == 2
-    assert isinstance(events[0], TaskStatusUpdated)
-    assert events[0].task_id == task.task_id
-    assert events[0].task_status == TaskStatus.Cancelled
-    assert isinstance(events[1], InstanceDeleted)
-    assert events[1].instance_id == instance_id
-
-
-def test_get_transition_events_delete_instance_cancels_pending_tasks(
-    instance: Instance,
-):
-    # arrange
-    instance_id = InstanceId()
-    current_instances: dict[InstanceId, Instance] = {instance_id: instance}
-    target_instances: dict[InstanceId, Instance] = {}
-    task = _make_task(instance_id, TaskStatus.Pending)
-    tasks = {task.task_id: task}
-
-    # act
-    events = get_transition_events(current_instances, target_instances, tasks)
-
-    # assert
-    assert len(events) == 2
-    assert isinstance(events[0], TaskStatusUpdated)
-    assert events[0].task_id == task.task_id
-    assert events[0].task_status == TaskStatus.Cancelled
-    assert isinstance(events[1], InstanceDeleted)
-
-
-def test_get_transition_events_delete_instance_ignores_completed_tasks(
-    instance: Instance,
-):
-    # arrange
-    instance_id = InstanceId()
-    current_instances: dict[InstanceId, Instance] = {instance_id: instance}
-    target_instances: dict[InstanceId, Instance] = {}
-    tasks = {
-        t.task_id: t
-        for t in [
-            _make_task(instance_id, TaskStatus.Complete),
-            _make_task(instance_id, TaskStatus.Failed),
-            _make_task(instance_id, TaskStatus.TimedOut),
-            _make_task(instance_id, TaskStatus.Cancelled),
-        ]
-    }
-
-    # act
-    events = get_transition_events(current_instances, target_instances, tasks)
-
-    # assert – only the InstanceDeleted event, no cancellations
-    assert len(events) == 1
-    assert isinstance(events[0], InstanceDeleted)
-
-
-def test_get_transition_events_delete_instance_cancels_only_matching_tasks(
-    instance: Instance,
-):
-    # arrange
-    instance_id_a = InstanceId()
-    instance_id_b = InstanceId()
-    current_instances: dict[InstanceId, Instance] = {
-        instance_id_a: instance,
-        instance_id_b: instance,
-    }
-    # only delete instance A, keep instance B
-    target_instances: dict[InstanceId, Instance] = {instance_id_b: instance}
-
-    task_a = _make_task(instance_id_a, TaskStatus.Running)
-    task_b = _make_task(instance_id_b, TaskStatus.Running)
-    tasks = {task_a.task_id: task_a, task_b.task_id: task_b}
-
-    # act
-    events = get_transition_events(current_instances, target_instances, tasks)
-
-    # assert – only task_a should be cancelled
-    cancel_events = [e for e in events if isinstance(e, TaskStatusUpdated)]
-    delete_events = [e for e in events if isinstance(e, InstanceDeleted)]
-    assert len(cancel_events) == 1
-    assert cancel_events[0].task_id == task_a.task_id
-    assert cancel_events[0].task_status == TaskStatus.Cancelled
-    assert len(delete_events) == 1
-    assert delete_events[0].instance_id == instance_id_a

src/exo/routing/router.py

@@ -8,13 +8,14 @@ from typing import cast
 from anyio import (
     BrokenResourceError,
     ClosedResourceError,
+    create_task_group,
     move_on_after,
     sleep_forever,
 )
+from anyio.abc import TaskGroup
 from exo_pyo3_bindings import (
     AllQueuesFullError,
     Keypair,
-    MessageTooLargeError,
     NetworkingHandle,
     NoPeersSubscribedToTopicError,
 )
@@ -24,7 +25,6 @@ from loguru import logger
 from exo.shared.constants import EXO_NODE_ID_KEYPAIR
 from exo.utils.channels import Receiver, Sender, channel
 from exo.utils.pydantic_ext import CamelCaseModel
-from exo.utils.task_group import TaskGroup
 
 from .connection_message import ConnectionMessage
 from .topics import CONNECTION_MESSAGES, PublishPolicy, TypedTopic
@@ -111,9 +111,10 @@ class Router:
         self._net: NetworkingHandle = handle
         self._tmp_networking_sender: Sender[tuple[str, bytes]] | None = send
         self._id_count = count()
-        self._tg: TaskGroup = TaskGroup()
+        self._tg: TaskGroup | None = None
 
     async def register_topic[T: CamelCaseModel](self, topic: TypedTopic[T]):
+        assert self._tg is None, "Attempted to register topic after setup time"
         send = self._tmp_networking_sender
         if send:
             self._tmp_networking_sender = None
@@ -147,7 +148,8 @@ class Router:
     async def run(self):
         logger.debug("Starting Router")
         try:
-            async with self._tg as tg:
+            async with create_task_group() as tg:
+                self._tg = tg
                 for topic in self.topic_routers:
                     router = self.topic_routers[topic]
                     tg.start_soon(router.run)
@@ -163,7 +165,9 @@ class Router:
 
     async def shutdown(self):
         logger.debug("Shutting down Router")
-        self._tg.cancel_tasks()
+        if not self._tg:
+            return
+        self._tg.cancel_scope.cancel()
 
     async def _networking_subscribe(self, topic: str):
         await self._net.gossipsub_subscribe(topic)
@@ -174,42 +178,28 @@ class Router:
         logger.info(f"Unsubscribed from {topic}")
 
     async def _networking_recv(self):
-        try:
-            while True:
-                topic, data = await self._net.gossipsub_recv()
-                logger.trace(f"Received message on {topic} with payload {data}")
-                if topic not in self.topic_routers:
-                    logger.warning(
-                        f"Received message on unknown or inactive topic {topic}"
-                    )
-                    continue
+        while True:
+            topic, data = await self._net.gossipsub_recv()
+            logger.trace(f"Received message on {topic} with payload {data}")
+            if topic not in self.topic_routers:
+                logger.warning(f"Received message on unknown or inactive topic {topic}")
+                continue
 
-                router = self.topic_routers[topic]
-                await router.publish_bytes(data)
-        except Exception as exception:
-            logger.opt(exception=exception).error(
-                "Gossipsub receive loop terminated unexpectedly"
-            )
-            raise
+            router = self.topic_routers[topic]
+            await router.publish_bytes(data)
 
     async def _networking_recv_connection_messages(self):
-        try:
-            while True:
-                update = await self._net.connection_update_recv()
-                message = ConnectionMessage.from_update(update)
-                logger.trace(
-                    f"Received message on connection_messages with payload {message}"
-                )
-                if CONNECTION_MESSAGES.topic in self.topic_routers:
-                    router = self.topic_routers[CONNECTION_MESSAGES.topic]
-                    assert router.topic.model_type == ConnectionMessage
-                    router = cast(TopicRouter[ConnectionMessage], router)
-                    await router.publish(message)
-        except Exception as exception:
-            logger.opt(exception=exception).error(
-                "Connection update receive loop terminated unexpectedly"
+        while True:
+            update = await self._net.connection_update_recv()
+            message = ConnectionMessage.from_update(update)
+            logger.trace(
+                f"Received message on connection_messages with payload {message}"
             )
-            raise
+            if CONNECTION_MESSAGES.topic in self.topic_routers:
+                router = self.topic_routers[CONNECTION_MESSAGES.topic]
+                assert router.topic.model_type == ConnectionMessage
+                router = cast(TopicRouter[ConnectionMessage], router)
+                await router.publish(message)
 
     async def _networking_publish(self):
         with self.networking_receiver as networked_items:
@@ -221,10 +211,6 @@ class Router:
                     pass
                 except AllQueuesFullError:
                     logger.warning(f"All peer queues full, dropping message on {topic}")
-                except MessageTooLargeError:
-                    logger.warning(
-                        f"Message too large for gossipsub on {topic} ({len(data)} bytes), dropping"
-                    )
 
 
 def get_node_id_keypair(

src/exo/shared/constants.py

@@ -36,12 +36,26 @@ EXO_MODELS_DIR = (
 
 # Read-only search path for pre-downloaded models (colon-separated directories)
 _EXO_MODELS_PATH_ENV = os.environ.get("EXO_MODELS_PATH", None)
-EXO_MODELS_PATH: tuple[Path, ...] | None = (
-    tuple(Path(p).expanduser() for p in _EXO_MODELS_PATH_ENV.split(":") if p)
-    if _EXO_MODELS_PATH_ENV is not None
-    else None
+
+# Well-known model cache directories from other inference engines
+_WELL_KNOWN_MODEL_PATHS: tuple[Path, ...] = tuple(
+    p for p in (Path.home() / ".cache" / "huggingface" / "hub",) if p.is_dir()
 )
 
+
+def _build_models_path() -> tuple[Path, ...] | None:
+    if _EXO_MODELS_PATH_ENV is not None:
+        user_paths = tuple(
+            Path(p).expanduser() for p in _EXO_MODELS_PATH_ENV.split(":") if p
+        )
+        return user_paths + tuple(
+            p for p in _WELL_KNOWN_MODEL_PATHS if p not in user_paths
+        )
+    return _WELL_KNOWN_MODEL_PATHS if _WELL_KNOWN_MODEL_PATHS else None
+
+
+EXO_MODELS_PATH: tuple[Path, ...] | None = _build_models_path()
+
 _RESOURCES_DIR_ENV = os.environ.get("EXO_RESOURCES_DIR", None)
 RESOURCES_DIR = (
     find_resources() if _RESOURCES_DIR_ENV is None else Path.home() / _RESOURCES_DIR_ENV
@@ -78,6 +92,4 @@ EXO_ENABLE_IMAGE_MODELS = (
     os.getenv("EXO_ENABLE_IMAGE_MODELS", "false").lower() == "true"
 )
 
-EXO_OFFLINE = os.getenv("EXO_OFFLINE", "false").lower() == "true"
-
 EXO_TRACING_ENABLED = os.getenv("EXO_TRACING_ENABLED", "false").lower() == "true"

src/exo/shared/election.py

@@ -4,8 +4,10 @@ import anyio
 from anyio import (
     CancelScope,
     Event,
+    create_task_group,
     get_cancelled_exc_class,
 )
+from anyio.abc import TaskGroup
 from loguru import logger
 
 from exo.routing.connection_message import ConnectionMessage
@@ -13,7 +15,6 @@ from exo.shared.types.commands import ForwarderCommand
 from exo.shared.types.common import NodeId, SessionId
 from exo.utils.channels import Receiver, Sender
 from exo.utils.pydantic_ext import CamelCaseModel
-from exo.utils.task_group import TaskGroup
 
 DEFAULT_ELECTION_TIMEOUT = 3.0
 
@@ -81,12 +82,13 @@ class Election:
         self._candidates: list[ElectionMessage] = []
         self._campaign_cancel_scope: CancelScope | None = None
         self._campaign_done: Event | None = None
-        self._tg = TaskGroup()
+        self._tg: TaskGroup | None = None
 
     async def run(self):
         logger.info("Starting Election")
         try:
-            async with self._tg as tg:
+            async with create_task_group() as tg:
+                self._tg = tg
                 tg.start_soon(self._election_receiver)
                 tg.start_soon(self._connection_receiver)
                 tg.start_soon(self._command_counter)
@@ -122,7 +124,12 @@ class Election:
         )
 
     async def shutdown(self) -> None:
-        self._tg.cancel_tasks()
+        if not self._tg:
+            logger.warning(
+                "Attempted to shutdown election service that was not running"
+            )
+            return
+        self._tg.cancel_scope.cancel()
 
     async def _election_receiver(self) -> None:
         with self._em_receiver as election_messages:
@@ -136,6 +143,7 @@ class Election:
                 if message.clock > self.clock:
                     self.clock = message.clock
                     logger.debug(f"New clock: {self.clock}")
+                    assert self._tg is not None
                     logger.debug("Starting new campaign")
                     candidates: list[ElectionMessage] = [message]
                     logger.debug(f"Candidates: {candidates}")
@@ -170,6 +178,7 @@ class Election:
                 # These messages are strictly peer to peer
                 self.clock += 1
                 logger.debug(f"New clock: {self.clock}")
+                assert self._tg is not None
                 candidates: list[ElectionMessage] = []
                 self._candidates = candidates
                 logger.debug("Starting new campaign")

src/exo/shared/models/model_cards.py

@@ -90,7 +90,6 @@ class ModelCard(CamelCaseModel):
     base_model: str = ""
     capabilities: list[str] = []
     uses_cfg: bool = False
-    trust_remote_code: bool = True
 
     @field_validator("tasks", mode="before")
     @classmethod
@@ -138,7 +137,6 @@ class ModelCard(CamelCaseModel):
             hidden_size=config_data.hidden_size or 0,
             supports_tensor=config_data.supports_tensor,
             tasks=[ModelTask.TextGeneration],
-            trust_remote_code=False,
         )
         await mc.save_to_custom_dir()
         _card_cache[model_id] = mc

src/exo/shared/types/memory.py

@@ -1,11 +1,6 @@
-import ctypes
-import sys
 from math import ceil
 from typing import Self, overload
 
-import psutil
-
-from exo.shared.logging import logger
 from exo.utils.pydantic_ext import FrozenModel
 
 
@@ -154,67 +149,3 @@ class Memory(FrozenModel):
             unit = "B"
 
         return f"{val:.2f} {unit}".rstrip("0").rstrip(".") + f" {unit}"
-
-
-def _load_memory_settings() -> tuple[float, "Memory"]:
-    """Load memory threshold and floor from settings (lazy import to avoid circular dep)."""
-    from exo.shared.types.settings import load_settings
-
-    s = load_settings()
-    return s.memory.memory_threshold, Memory.from_gb(s.memory.memory_floor_gb)
-
-_libc: ctypes.CDLL | None = None
-
-
-def _macos_memorystatus_level() -> int:
-    global _libc  # noqa: PLW0603
-    if _libc is None:
-        _libc = ctypes.CDLL("/usr/lib/libSystem.B.dylib")
-    level = ctypes.c_int(0)
-    size = ctypes.c_size_t(ctypes.sizeof(ctypes.c_int))
-    ret: int = _libc.sysctlbyname(  # pyright: ignore[reportAny]
-        b"kern.memorystatus_level",
-        ctypes.byref(level),
-        ctypes.byref(size),
-        None,
-        ctypes.c_size_t(0),
-    )
-    if ret != 0:
-        raise OSError("sysctlbyname kern.memorystatus_level failed")
-    return level.value
-
-
-def _get_macos_memory_pressure() -> float:
-    try:
-        return 1.0 - _macos_memorystatus_level() / 100.0
-    except (OSError, FileNotFoundError):
-        logger.warning("Using fallback memory pressure")
-        return _fallback_memory_pressure()
-
-
-def _fallback_memory_pressure() -> float:
-    vm = psutil.virtual_memory()
-    return 1.0 - vm.available / vm.total
-
-
-def get_memory_pressure() -> float:
-    if sys.platform == "darwin":
-        return _get_macos_memory_pressure()
-    return _fallback_memory_pressure()
-
-
-def get_memory_limit() -> Memory:
-    threshold, floor = _load_memory_settings()
-    total = psutil.virtual_memory().total
-    safety = min(int(total * (1 - threshold)), floor.in_bytes)
-    return Memory.from_bytes(total - safety)
-
-
-def get_memory_available_locally() -> Memory:
-    total = Memory.from_bytes(psutil.virtual_memory().total)
-    return get_memory_limit() - total * get_memory_pressure()
-
-
-def get_memory_pressure_threshold() -> float:
-    total = psutil.virtual_memory().total
-    return get_memory_limit().in_bytes / total

src/exo/shared/types/settings.py

@@ -1,121 +0,0 @@
-import os
-import tomllib
-from typing import Literal
-
-import psutil
-from pydantic import ConfigDict, Field, ValidationError
-
-from exo.shared.constants import EXO_CONFIG_FILE
-from exo.shared.logging import logger
-from exo.shared.types.memory import Memory
-from exo.utils.pydantic_ext import CamelCaseModel
-
-
-def _default_memory_threshold() -> float:
-    total_gb = Memory.from_bytes(psutil.virtual_memory().total).in_gb
-    if total_gb >= 128:
-        return 0.85
-    if total_gb >= 64:
-        return 0.80
-    if total_gb >= 32:
-        return 0.75
-    return 0.70
-
-
-class MemorySettings(CamelCaseModel):
-    model_config = ConfigDict(
-        alias_generator=None,
-        validate_by_name=True,
-        extra="forbid",
-        strict=False,
-    )
-
-    oom_prevention: bool = False
-    memory_threshold: float = Field(default_factory=_default_memory_threshold, ge=0.0, le=1.0)
-    memory_floor_gb: float = Field(default=5.0, ge=0.0)
-
-
-class GenerationSettings(CamelCaseModel):
-    model_config = ConfigDict(
-        alias_generator=None,
-        validate_by_name=True,
-        extra="forbid",
-        strict=False,
-    )
-
-    prefill_step_size: int = Field(default=4096, ge=1)
-    max_tokens: int = Field(default=32168, ge=1)
-    kv_cache_bits: Literal[4, 8] | None = None
-
-
-class ExoSettings(CamelCaseModel):
-    model_config = ConfigDict(
-        alias_generator=None,
-        validate_by_name=True,
-        extra="ignore",
-        strict=False,
-    )
-
-    memory: MemorySettings = Field(default_factory=MemorySettings)
-    generation: GenerationSettings = Field(default_factory=GenerationSettings)
-
-
-_cached_settings: ExoSettings | None = None
-_cached_mtime: float = 0.0
-
-
-def load_settings() -> ExoSettings:
-    global _cached_settings, _cached_mtime  # noqa: PLW0603
-
-    try:
-        mtime = EXO_CONFIG_FILE.stat().st_mtime
-        if _cached_settings is not None and mtime == _cached_mtime:
-            return _cached_settings
-        with open(EXO_CONFIG_FILE, "rb") as f:
-            data = tomllib.load(f)
-        settings = ExoSettings.model_validate(data)
-        _cached_mtime = mtime
-    except FileNotFoundError:
-        settings = ExoSettings()
-    except (tomllib.TOMLDecodeError, ValidationError) as e:
-        logger.warning(f"Invalid config file {EXO_CONFIG_FILE}: {e}")
-        settings = ExoSettings()
-
-    # Env vars override config file for backward compat.
-    env_threshold = os.environ.get("EXO_MEMORY_THRESHOLD")
-    if env_threshold is not None:
-        settings = settings.model_copy(
-            update={"memory": settings.memory.model_copy(update={"memory_threshold": float(env_threshold)})}
-        )
-    env_floor = os.environ.get("EXO_MEMORY_FLOOR")
-    if env_floor is not None:
-        settings = settings.model_copy(
-            update={"memory": settings.memory.model_copy(update={"memory_floor_gb": float(env_floor)})}
-        )
-
-    _cached_settings = settings
-    return settings
-
-
-def save_settings(settings: ExoSettings) -> None:
-    global _cached_settings, _cached_mtime  # noqa: PLW0603
-
-    EXO_CONFIG_FILE.parent.mkdir(parents=True, exist_ok=True)
-
-    lines = [
-        "[memory]",
-        f"oom_prevention = {'true' if settings.memory.oom_prevention else 'false'}",
-        f"memory_threshold = {settings.memory.memory_threshold}",
-        f"memory_floor_gb = {settings.memory.memory_floor_gb}",
-        "",
-        "[generation]",
-        f"prefill_step_size = {settings.generation.prefill_step_size}",
-        f"max_tokens = {settings.generation.max_tokens}",
-    ]
-    if settings.generation.kv_cache_bits is not None:
-        lines.append(f"kv_cache_bits = {settings.generation.kv_cache_bits}")
-
-    EXO_CONFIG_FILE.write_text("\n".join(lines) + "\n")
-
-    _cached_settings = settings
-    _cached_mtime = EXO_CONFIG_FILE.stat().st_mtime

src/exo/utils/channels.py

@@ -5,7 +5,7 @@ from math import inf
 from multiprocessing.synchronize import Event
 from queue import Empty, Full
 from types import TracebackType
-from typing import Any, Self
+from typing import Self
 
 from anyio import (
     CapacityLimiter,
@@ -157,7 +157,7 @@ class MpSender[T]:
     ) -> None:
         self.close()
 
-    def __getstate__(self) -> dict[str, Any]:
+    def __getstate__(self):
         d = self.__dict__.copy()
         d.pop("__orig_class__", None)
         return d

src/exo/utils/info_gatherer/info_gatherer.py

@@ -8,11 +8,12 @@ from subprocess import CalledProcessError
 from typing import Self, cast
 
 import anyio
-from anyio import fail_after, open_process, to_thread
+from anyio import create_task_group, fail_after, open_process, to_thread
+from anyio.abc import TaskGroup
 from anyio.streams.buffered import BufferedByteReceiveStream
 from anyio.streams.text import TextReceiveStream
 from loguru import logger
-from pydantic import ConfigDict, ValidationError
+from pydantic import ValidationError
 
 from exo.shared.constants import EXO_CONFIG_FILE, EXO_MODELS_DIR
 from exo.shared.types.memory import Memory
@@ -29,7 +30,6 @@ from exo.shared.types.thunderbolt import (
 )
 from exo.utils.channels import Sender
 from exo.utils.pydantic_ext import TaggedModel
-from exo.utils.task_group import TaskGroup
 
 from .macmon import MacmonMetrics
 from .system_info import (
@@ -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"""
 
-    model_config = ConfigDict(extra="ignore")
-
     @classmethod
     async def gather(cls) -> Self | None:
         cfg_file = anyio.Path(EXO_CONFIG_FILE)
@@ -383,7 +381,7 @@ class InfoGatherer:
     static_info_poll_interval: float | None = 60
     rdma_ctl_poll_interval: float | None = 10 if IS_DARWIN else None
     disk_poll_interval: float | None = 30
-    _tg: TaskGroup = field(init=False, default_factory=TaskGroup)
+    _tg: TaskGroup = field(init=False, default_factory=create_task_group)
 
     async def run(self):
         async with self._tg as tg:
@@ -410,7 +408,7 @@ class InfoGatherer:
                 await self.info_sender.send(nc)
 
     def shutdown(self):
-        self._tg.cancel_tasks()
+        self._tg.cancel_scope.cancel()
 
     async def _monitor_static_info(self):
         if self.static_info_poll_interval is None:

src/exo/utils/task_group.py

@@ -1,65 +0,0 @@
-from collections.abc import Awaitable, Callable
-from dataclasses import dataclass, field
-from types import TracebackType
-from typing import Any, Unpack
-
-from anyio import create_task_group
-from anyio.abc import TaskGroup as TaskGroupABC
-
-
-@dataclass
-class TaskGroup:
-    _tg: TaskGroupABC | None = field(default=None, init=False)
-    _queued: list[tuple[Any, Any, Any]] | None = field(default_factory=list, init=False)
-
-    def is_running(self) -> bool:
-        return self._tg is not None
-
-    def cancel_tasks(self):
-        assert self._tg
-        self._tg.cancel_scope.cancel()
-
-    def cancel_called(self) -> bool:
-        assert self._tg
-        return self._tg.cancel_scope.cancel_called
-
-    def start_soon[*T](
-        self,
-        func: Callable[[Unpack[T]], Awaitable[Any]],
-        *args: Unpack[T],
-        name: object = None,
-    ) -> None:
-        assert self._tg is not None
-        assert self._queued is None
-        self._tg.start_soon(func, *args, name=name)
-
-    def queue[*T](
-        self,
-        func: Callable[[Unpack[T]], Awaitable[Any]],
-        *args: Unpack[T],
-        name: object = None,
-    ) -> None:
-        assert self._tg is None
-        assert self._queued is not None
-        self._queued.append((func, args, name))
-
-    async def __aenter__(self) -> TaskGroupABC:
-        assert self._tg is None
-        assert self._queued is not None
-        self._tg = create_task_group()
-        r = await self._tg.__aenter__()
-        for func, args, name in self._queued:  # pyright: ignore[reportAny]
-            self._tg.start_soon(func, *args, name=name)  # pyright: ignore[reportAny]
-        self._queued = None
-        return r
-
-    async def __aexit__(
-        self,
-        exc_type: type[BaseException] | None,
-        exc_val: BaseException | None,
-        exc_tb: TracebackType | None,
-    ) -> bool:
-        """Exit the task group context waiting for all tasks to finish."""
-        assert self._tg is not None, "aenter sets self.lazy, so it exists when we aexit"
-        assert self._queued is None
-        return await self._tg.__aexit__(exc_type, exc_val, exc_tb)

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

@@ -32,12 +32,13 @@ from mlx_lm.models.minimax import MiniMaxAttention
 from mlx_lm.models.minimax import Model as MiniMaxModel
 from mlx_lm.models.ministral3 import Model as Ministral3Model
 from mlx_lm.models.qwen3_moe import Model as Qwen3MoeModel
-from mlx_lm.models.qwen3_moe import Qwen3MoeDecoderLayer, Qwen3MoeSparseMoeBlock
+from mlx_lm.models.qwen3_moe import Qwen3MoeSparseMoeBlock
 from mlx_lm.models.qwen3_next import Model as Qwen3NextModel
 from mlx_lm.models.qwen3_next import Qwen3NextDecoderLayer, Qwen3NextSparseMoeBlock
 from mlx_lm.models.step3p5 import Model as Step35Model
 from mlx_lm.models.step3p5 import Step3p5MLP as Step35MLP
 from mlx_lm.models.step3p5 import Step3p5Model as Step35InnerModel
+from transformers.models.qwen3.modeling_qwen3 import Qwen3DecoderLayer
 
 from exo.shared.logging import logger
 from exo.shared.types.worker.shards import PipelineShardMetadata
@@ -128,11 +129,11 @@ class PipelineFirstLayer(CustomMlxLayer):
 
     def __call__(self, x: mx.array, *args: object, **kwargs: object) -> mx.array:
         if self.r != 0:
-            # We want to avoid GPU timeout errors by evalling the distributed operation
-            # so that it stays on CPU, which does not have a timeout.
-            mx.eval(x)
             x = mx.distributed.recv_like(x, (self.r - 1), group=self.group)
-            mx.eval(x)
+            if self.is_prefill:
+                # We want to avoid GPU timeout errors by evalling the distributed operation
+                # so that it stays on CPU, which does not have a timeout.
+                mx.eval(x)
         return self.original_layer(x, *args, **kwargs)
 
 
@@ -158,10 +159,6 @@ class PipelineLastLayer(CustomMlxLayer):
 
         output: mx.array = self.original_layer(x, *args, **kwargs)
 
-        # Eval layer output to materialize it before send — this splits the graph
-        # so the send is isolated and the receiving rank's recv can complete.
-        mx.eval(output)
-
         if self.r != self.s - 1:
             output = mx.distributed.send(
                 output, (self.r + 1) % self.s, group=self.group
@@ -171,15 +168,15 @@ class PipelineLastLayer(CustomMlxLayer):
                 # doesn't have .keys directly; access via first sub-cache.
                 _cache = cache[0] if hasattr(cache, "caches") else cache  # type: ignore
                 _cache.keys = mx.depends(_cache.keys, output)  # type: ignore
-            mx.eval(output)
-            if cache is not None:
-                mx.eval(_cache.keys)  # type: ignore
+            if self.is_prefill:
+                mx.eval(output)
+                if cache is not None:
+                    mx.eval(_cache.keys)  # type: ignore
 
         if not self.is_prefill:
             output = mx.distributed.all_gather(output, group=self.group)[
                 -output.shape[0] :
             ]
-            mx.eval(output)
 
         return output
 
@@ -843,7 +840,7 @@ class QwenShardingStrategy(TensorParallelShardingStrategy):
         for i, layer in enumerate(model.layers):
             eval_with_timeout(layer.parameters(), timeout_seconds / total, on_timeout)
             # Shard the self attention
-            if isinstance(layer, Qwen3MoeDecoderLayer):
+            if isinstance(layer, Qwen3DecoderLayer):
                 layer.self_attn.q_proj = self.all_to_sharded_linear(
                     layer.self_attn.q_proj
                 )
@@ -856,8 +853,6 @@ class QwenShardingStrategy(TensorParallelShardingStrategy):
                 layer.self_attn.o_proj = self.sharded_to_all_linear(
                     layer.self_attn.o_proj
                 )
-                layer.self_attn.n_heads //= self.N
-                layer.self_attn.n_kv_heads //= self.N
             else:
                 assert isinstance(layer, Qwen3NextDecoderLayer)
                 if hasattr(layer, "linear_attn"):

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

@@ -1,6 +1,8 @@
+import os
 from copy import deepcopy
 
 import mlx.core as mx
+import psutil
 from mlx_lm.models.cache import (
     ArraysCache,
     CacheList,
@@ -10,14 +12,31 @@ from mlx_lm.models.cache import (
 )
 from mlx_lm.tokenizer_utils import TokenizerWrapper
 
-from exo.shared.types.memory import Memory, get_memory_pressure
+from exo.shared.types.memory import Memory
 from exo.shared.types.mlx import KVCacheType
-from exo.shared.types.settings import load_settings
 from exo.worker.engines.mlx import Model
-from exo.worker.engines.mlx.constants import CACHE_GROUP_SIZE
+from exo.worker.engines.mlx.constants import CACHE_GROUP_SIZE, KV_CACHE_BITS
 from exo.worker.runner.bootstrap import logger
 
 
+# Fraction of device memory above which LRU eviction kicks in.
+# Smaller machines need more aggressive eviction.
+def _default_memory_threshold() -> float:
+    total_gb = Memory.from_bytes(psutil.virtual_memory().total).in_gb
+    if total_gb >= 128:
+        return 0.85
+    if total_gb >= 64:
+        return 0.80
+    if total_gb >= 32:
+        return 0.75
+    return 0.70
+
+
+_MEMORY_THRESHOLD = float(
+    os.environ.get("EXO_MEMORY_THRESHOLD", _default_memory_threshold())
+)
+
+
 class CacheSnapshot:
     """Snapshot of states at a known token position."""
 
@@ -73,15 +92,6 @@ class KVPrefixCache:
         self._snapshots.clear()
         self._last_used.clear()
 
-    def force_evict_all(self) -> int:
-        count = len(self.caches)
-        self.clear()
-        if count > 0:
-            logger.info(
-                f"Force-evicted all {count} prefix cache entries due to memory pressure"
-            )
-        return count
-
     def add_kv_cache(
         self,
         prompt_tokens: mx.array,
@@ -207,7 +217,7 @@ class KVPrefixCache:
         # Evict LRU entries until below threshold
         while (
             len(self.caches) > 0
-            and self.get_memory_used_percentage() > load_settings().memory.memory_threshold
+            and self.get_memory_used_percentage() > _MEMORY_THRESHOLD
         ):
             lru_index = self._last_used.index(min(self._last_used))
             evicted_tokens = len(self.prompts[lru_index])
@@ -220,7 +230,7 @@ class KVPrefixCache:
             )
 
     def get_memory_used_percentage(self) -> float:
-        local_pressure: float = get_memory_pressure()
+        local_pressure: float = get_memory_used_percentage()
 
         if self._group is None:
             return local_pressure
@@ -289,47 +299,15 @@ def get_prefix_length(prompt: mx.array, cached_prompt: mx.array) -> int:
     return int(mx.sum(prefix_mask).item())
 
 
-def _measure_single_cache_bytes(
-    entry: KVCache | RotatingKVCache | QuantizedKVCache | ArraysCache | CacheList,
-) -> int:
-    if isinstance(entry, CacheList):
-        return sum(
-            _measure_single_cache_bytes(c)  # pyright: ignore[reportArgumentType]
-            for c in entry.caches
-        )
-
-    total = 0
-    if isinstance(entry, ArraysCache):
-        state = entry.state  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
-        for arr in state:  # pyright: ignore[reportUnknownVariableType]
-            if isinstance(arr, mx.array):
-                total += arr.nbytes
-        return total
-
-    total = 0
-    for attr_name in ("keys", "values"):
-        val: object = getattr(entry, attr_name, None)
-        if val is None:
-            continue
-        if isinstance(val, mx.array):
-            total += val.nbytes
-        elif isinstance(val, (tuple, list)):
-            for arr in val:  # pyright: ignore[reportUnknownVariableType]
-                if isinstance(arr, mx.array):
-                    total += arr.nbytes
-
-    return total
+def get_available_memory() -> Memory:
+    mem: int = psutil.virtual_memory().available
+    return Memory.from_bytes(mem)
 
 
-def measure_cache_bytes(cache: KVCacheType) -> int:
-    return sum(_measure_single_cache_bytes(c) for c in cache)
-
-
-def measure_kv_cache_bytes_per_token(cache: KVCacheType) -> Memory:
-    offset = cache_length(cache)
-    if offset == 0:
-        return Memory.from_bytes(0)
-    return Memory.from_bytes(measure_cache_bytes(cache) // offset)
+def get_memory_used_percentage() -> float:
+    mem = psutil.virtual_memory()
+    # percent is 0-100
+    return float(mem.percent / 100)
 
 
 def make_kv_cache(
@@ -342,14 +320,13 @@ def make_kv_cache(
         return model.make_cache()  # type: ignore
 
     if max_kv_size is None:
-        kv_cache_bits = load_settings().generation.kv_cache_bits
-        if kv_cache_bits is None:
+        if KV_CACHE_BITS is None:
             logger.info("Using default KV cache")
             return [KVCache() for _ in model.layers]
         else:
             logger.info("Using quantized KV cache")
             return [
-                QuantizedKVCache(group_size=CACHE_GROUP_SIZE, bits=kv_cache_bits)
+                QuantizedKVCache(group_size=CACHE_GROUP_SIZE, bits=KV_CACHE_BITS)
                 for _ in model.layers
             ]
     else:

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

@@ -1,4 +1,3 @@
-import math
 import time
 from copy import deepcopy
 from typing import Callable, Generator, cast, get_args
@@ -18,9 +17,8 @@ from exo.shared.types.api import (
     Usage,
 )
 from exo.shared.types.common import ModelId
-from exo.shared.types.memory import Memory, get_memory_available_locally
+from exo.shared.types.memory import Memory
 from exo.shared.types.mlx import KVCacheType
-from exo.shared.types.settings import load_settings
 from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
 from exo.shared.types.worker.runner_response import (
     GenerationResponse,
@@ -33,18 +31,17 @@ from exo.worker.engines.mlx.cache import (
     encode_prompt,
     has_non_kv_caches,
     make_kv_cache,
-    measure_kv_cache_bytes_per_token,
     snapshot_ssm_states,
 )
 from exo.worker.engines.mlx.constants import (
     DEFAULT_TOP_LOGPROBS,
     KV_BITS,
     KV_GROUP_SIZE,
+    MAX_TOKENS,
 )
 from exo.worker.engines.mlx.utils_mlx import (
     apply_chat_template,
     fix_unmatched_think_end_tokens,
-    mx_any,
     mx_barrier,
 )
 from exo.worker.runner.bootstrap import logger
@@ -112,7 +109,7 @@ def prefill(
             max_tokens=1,
             sampler=sampler,
             prompt_cache=cache,
-            prefill_step_size=load_settings().generation.prefill_step_size,
+            prefill_step_size=4096,
             kv_group_size=KV_GROUP_SIZE,
             kv_bits=KV_BITS,
             prompt_progress_callback=progress_callback,
@@ -150,8 +147,7 @@ def warmup_inference(
     model: Model,
     tokenizer: TokenizerWrapper,
     group: mx.distributed.Group | None,
-) -> tuple[int, Memory]:
-    """Run warmup inference and tokens_generated and bytes_per_token"""
+) -> int:
     content = "Prompt to warm up the inference engine. Repeat this."
 
     warmup_prompt = apply_chat_template(
@@ -190,12 +186,9 @@ def warmup_inference(
 
     logger.info("Generated ALL warmup tokens")
 
-    bytes_per_token = measure_kv_cache_bytes_per_token(cache)
-    logger.info(f"Measured KV cache cost: {bytes_per_token} per token")
-
     mx_barrier(group)
 
-    return tokens_generated, bytes_per_token
+    return tokens_generated
 
 
 def ban_token_ids(token_ids: list[int]) -> Callable[[mx.array, mx.array], mx.array]:
@@ -255,9 +248,6 @@ def extract_top_logprobs(
     for i in range(top_logprobs):
         token_id = int(top_indices[i].item())
         token_logprob = float(top_values[i].item())
-        if math.isnan(token_logprob):
-            continue
-
         # Decode token ID to string
         token_str = tokenizer.decode([token_id])
         # Get byte representation
@@ -273,33 +263,6 @@ def extract_top_logprobs(
     return selected_logprob, top_logprob_items
 
 
-def _check_memory_budget(
-    bytes_per_token: Memory,
-    total_sequence_tokens: int,
-    kv_prefix_cache: KVPrefixCache | None,
-    group: mx.distributed.Group | None,
-) -> str | None:
-    if bytes_per_token.in_bytes == 0:
-        return None
-
-    estimated = bytes_per_token * total_sequence_tokens
-    over_budget = estimated > get_memory_available_locally()
-
-    if not mx_any(over_budget, group):
-        return None
-
-    if kv_prefix_cache is not None and kv_prefix_cache.force_evict_all() > 0:
-        mx.clear_cache()
-        over_budget = estimated > get_memory_available_locally()
-        if not mx_any(over_budget, group):
-            return None
-
-    return (
-        "Not enough memory for this conversation. "
-        "Please start a new conversation or compact your messages."
-    )
-
-
 def mlx_generate(
     model: Model,
     tokenizer: TokenizerWrapper,
@@ -308,10 +271,7 @@ def mlx_generate(
     kv_prefix_cache: KVPrefixCache | None,
     group: mx.distributed.Group | None,
     on_prefill_progress: Callable[[int, int], None] | None = None,
-    bytes_per_token: Memory | None = None,
 ) -> Generator[GenerationResponse]:
-    if bytes_per_token is None:
-        bytes_per_token = Memory()
     # Ensure that generation stats only contains peak memory for this generation
     mx.reset_peak_memory()
     # TODO: Randomise task seed and set in taskparams, instead of hard coding as 42.
@@ -343,23 +303,6 @@ def mlx_generate(
                 f"KV cache hit: {prefix_hit_length}/{len(all_prompt_tokens)} tokens cached ({100 * prefix_hit_length / len(all_prompt_tokens):.1f}%)"
             )
 
-    if bytes_per_token.in_bytes > 0 and load_settings().memory.oom_prevention:
-        oom_error = _check_memory_budget(
-            bytes_per_token=bytes_per_token,
-            total_sequence_tokens=len(all_prompt_tokens),
-            kv_prefix_cache=kv_prefix_cache,
-            group=group,
-        )
-        if oom_error is not None:
-            logger.warning(f"OOM prevention (prefill): {oom_error}")
-            yield GenerationResponse(
-                text=oom_error,
-                token=0,
-                finish_reason="error",
-                usage=None,
-            )
-            return
-
     logits_processors: list[Callable[[mx.array, mx.array], mx.array]] = []
     if is_bench:
         # Only sample length eos tokens
@@ -395,7 +338,7 @@ def mlx_generate(
     # stream_generate starts from the last token
     last_token = prompt_tokens[-2:]
 
-    max_tokens = task.max_output_tokens or load_settings().generation.max_tokens
+    max_tokens = task.max_output_tokens or MAX_TOKENS
     accumulated_text = ""
     generated_text_parts: list[str] = []
     generation_start_time = time.perf_counter()

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

@@ -23,7 +23,9 @@ from mlx_lm.models.deepseek_v3 import DeepseekV3Model
 from mlx_lm.tokenizer_utils import TokenizerWrapper
 
 from exo.shared.models.model_cards import ModelId
-from exo.worker.engines.mlx.constants import TRUST_REMOTE_CODE
+from exo.worker.engines.mlx.constants import (
+    TRUST_REMOTE_CODE,
+)
 
 try:
     from mlx_lm.tokenizer_utils import load_tokenizer
@@ -189,9 +191,10 @@ def load_mlx_items(
                 mx.eval(layer)  # type: ignore
                 if on_layer_loaded is not None:
                     on_layer_loaded(i, total)
-        except ValueError as e:
-            logger.opt(exception=e).debug(
+        except ValueError:
+            logger.debug(
                 "Model architecture doesn't support layer-by-layer progress tracking",
+                exc_info=True,
             )
         mx.eval(model)
         end_time = time.perf_counter()
@@ -214,8 +217,6 @@ def load_mlx_items(
 
     set_wired_limit_for_model(get_weights_size(bound_instance.bound_shard))
 
-    mx.clear_cache()
-
     return cast(Model, model), tokenizer
 
 
@@ -293,11 +294,7 @@ def shard_and_load(
 
 def get_tokenizer(model_path: Path, shard_metadata: ShardMetadata) -> TokenizerWrapper:
     """Load tokenizer for a model shard. Delegates to load_tokenizer_for_model_id."""
-    return load_tokenizer_for_model_id(
-        shard_metadata.model_card.model_id,
-        model_path,
-        trust_remote_code=shard_metadata.model_card.trust_remote_code,
-    )
+    return load_tokenizer_for_model_id(shard_metadata.model_card.model_id, model_path)
 
 
 def get_eos_token_ids_for_model(model_id: ModelId) -> list[int] | None:
@@ -329,7 +326,7 @@ def get_eos_token_ids_for_model(model_id: ModelId) -> list[int] | None:
 
 
 def load_tokenizer_for_model_id(
-    model_id: ModelId, model_path: Path, *, trust_remote_code: bool = TRUST_REMOTE_CODE
+    model_id: ModelId, model_path: Path
 ) -> TokenizerWrapper:
     """
     Load tokenizer for a model given its ID and local path.
@@ -398,7 +395,7 @@ def load_tokenizer_for_model_id(
 
     tokenizer = load_tokenizer(
         model_path,
-        tokenizer_config_extra={"trust_remote_code": trust_remote_code},
+        tokenizer_config_extra={"trust_remote_code": TRUST_REMOTE_CODE},
         eos_token_ids=eos_token_ids,
     )
 
@@ -646,7 +643,7 @@ class NullKVCache(KVCache):
         raise NotImplementedError("We should not be setting a NullKVCache.")
 
 
-def mlx_force_oom(size: int = 200000) -> None:
+def mlx_force_oom(size: int = 40000) -> None:
     """
     Force an Out-Of-Memory (OOM) error in MLX by performing large tensor operations.
     """
@@ -673,7 +670,7 @@ def set_wired_limit_for_model(model_size: Memory):
         return
 
     max_rec_size = Memory.from_bytes(
-        int(mx.device_info()["max_recommended_working_set_size"])
+        int(mx.metal.device_info()["max_recommended_working_set_size"])
     )
     if model_size > 0.9 * max_rec_size:
         logger.warning(

src/exo/worker/main.py

@@ -3,7 +3,8 @@ from datetime import datetime, timezone
 from random import random
 
 import anyio
-from anyio import CancelScope, fail_after
+from anyio import CancelScope, create_task_group, fail_after
+from anyio.abc import TaskGroup
 from loguru import logger
 
 from exo.download.download_utils import resolve_model_in_path
@@ -50,7 +51,6 @@ from exo.utils.event_buffer import OrderedBuffer
 from exo.utils.info_gatherer.info_gatherer import GatheredInfo, InfoGatherer
 from exo.utils.info_gatherer.net_profile import check_reachable
 from exo.utils.keyed_backoff import KeyedBackoff
-from exo.utils.task_group import TaskGroup
 from exo.worker.plan import plan
 from exo.worker.runner.runner_supervisor import RunnerSupervisor
 
@@ -80,7 +80,7 @@ class Worker:
 
         self.state: State = State()
         self.runners: dict[RunnerId, RunnerSupervisor] = {}
-        self._tg: TaskGroup = TaskGroup()
+        self._tg: TaskGroup = create_task_group()
 
         self._nack_cancel_scope: CancelScope | None = None
         self._nack_attempts: int = 0
@@ -317,7 +317,7 @@ class Worker:
                     await self._start_runner_task(task)
 
     def shutdown(self):
-        self._tg.cancel_tasks()
+        self._tg.cancel_scope.cancel()
 
     async def _start_runner_task(self, task: Task):
         if (instance := self.state.instances.get(task.instance_id)) is not None:

src/exo/worker/runner/bootstrap.py

@@ -4,7 +4,7 @@ import loguru
 
 from exo.shared.types.events import Event, RunnerStatusUpdated
 from exo.shared.types.tasks import Task, TaskId
-from exo.shared.types.worker.instances import BoundInstance
+from exo.shared.types.worker.instances import BoundInstance, MlxJacclInstance
 from exo.shared.types.worker.runners import RunnerFailed
 from exo.utils.channels import ClosedResourceError, MpReceiver, MpSender
 
@@ -18,15 +18,21 @@ def entrypoint(
     cancel_receiver: MpReceiver[TaskId],
     _logger: "loguru.Logger",
 ) -> None:
-    global logger
-    logger = _logger
-
     fast_synch_override = os.environ.get("EXO_FAST_SYNCH")
-    if fast_synch_override != "off":
+    if fast_synch_override == "on" or (
+        fast_synch_override != "off"
+        and (
+            isinstance(bound_instance.instance, MlxJacclInstance)
+            and len(bound_instance.instance.jaccl_devices) >= 2
+        )
+    ):
         os.environ["MLX_METAL_FAST_SYNCH"] = "1"
     else:
         os.environ["MLX_METAL_FAST_SYNCH"] = "0"
 
+    global logger
+    logger = _logger
+
     logger.info(f"Fast synch flag: {os.environ['MLX_METAL_FAST_SYNCH']}")
 
     # Import main after setting global logger - this lets us just import logger from this module

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

@@ -31,12 +31,6 @@ from exo.shared.types.events import (
     TaskAcknowledged,
     TaskStatusUpdated,
 )
-from exo.shared.types.memory import (
-    Memory,
-    get_memory_pressure,
-    get_memory_pressure_threshold,
-)
-from exo.shared.types.settings import load_settings
 from exo.shared.types.tasks import (
     ConnectToGroup,
     LoadModel,
@@ -103,7 +97,6 @@ def main(
         bound_instance.bound_runner_id,
         bound_instance.bound_shard,
     )
-    model_id = shard_metadata.model_card.model_id
     device_rank = shard_metadata.device_rank
     logger.info("hello from the runner")
     if getattr(shard_metadata, "immediate_exception", False):
@@ -120,7 +113,6 @@ def main(
     group = None
     kv_prefix_cache: KVPrefixCache | None = None
     check_for_cancel_every: int | None = None
-    bytes_per_token = Memory.from_bytes(0)
 
     current_status: RunnerStatus = RunnerIdle()
     logger.info("runner created")
@@ -232,14 +224,12 @@ def main(
                     assert tokenizer
 
                     t = time.monotonic()
-                    toks, bytes_per_token = warmup_inference(
+                    toks = warmup_inference(
                         model=cast(Model, inference_model),
                         tokenizer=tokenizer,
                         group=group,
                     )
-                    logger.info(
-                        f"warmed up by generating {toks} tokens, {bytes_per_token}/token for KV cache"
-                    )
+                    logger.info(f"warmed up by generating {toks} tokens")
                     check_for_cancel_every = min(
                         math.ceil(toks / min(time.monotonic() - t, 0.001)), 100
                     )
@@ -291,7 +281,7 @@ def main(
                                 ChunkGenerated(
                                     command_id=command_id,
                                     chunk=PrefillProgressChunk(
-                                        model=model_id,
+                                        model=shard_metadata.model_card.model_id,
                                         processed_tokens=processed,
                                         total_tokens=total,
                                     ),
@@ -319,7 +309,6 @@ def main(
                             kv_prefix_cache=kv_prefix_cache,
                             on_prefill_progress=on_prefill_progress,
                             group=group,
-                            bytes_per_token=bytes_per_token,
                         )
 
                         if tokenizer.has_thinking:
@@ -336,17 +325,13 @@ def main(
                         # Model-specific output parsing for tool calls.
                         if isinstance(inference_model, GptOssModel):
                             mlx_generator = parse_gpt_oss(mlx_generator)
-                        elif (
-                            isinstance(inference_model, DeepseekV32Model)
-                            and "deepseek" in model_id.normalize().lower()
-                        ):
+                        elif isinstance(inference_model, DeepseekV32Model):
                             mlx_generator = parse_deepseek_v32(mlx_generator)
                         elif tool_parser:
                             mlx_generator = parse_tool_calls(mlx_generator, tool_parser)
 
                         completion_tokens = 0
                         tokens_since_last_cancel_check = check_for_cancel_every
-                        oom_stopped = False
                         for response in mlx_generator:
                             tokens_since_last_cancel_check += 1
                             if tokens_since_last_cancel_check >= check_for_cancel_every:
@@ -355,15 +340,7 @@ def main(
                                 want_to_cancel = (task.task_id in cancelled_tasks) or (
                                     TaskId("CANCEL_CURRENT_TASK") in cancelled_tasks
                                 )
-                                oom_local = (
-                                    load_settings().memory.oom_prevention
-                                    and bytes_per_token.in_bytes > 0
-                                    and get_memory_pressure()
-                                    > get_memory_pressure_threshold()
-                                )
-                                if mx_any(want_to_cancel or oom_local, group):
-                                    if not want_to_cancel:
-                                        oom_stopped = True
+                                if mx_any(want_to_cancel, group):
                                     break
 
                             match response:
@@ -378,7 +355,7 @@ def main(
                                                 command_id=command_id,
                                                 chunk=ErrorChunk(
                                                     error_message=response.text,
-                                                    model=model_id,
+                                                    model=shard_metadata.model_card.model_id,
                                                 ),
                                             )
                                         )
@@ -393,7 +370,7 @@ def main(
                                             ChunkGenerated(
                                                 command_id=command_id,
                                                 chunk=TokenChunk(
-                                                    model=model_id,
+                                                    model=shard_metadata.model_card.model_id,
                                                     text=response.text,
                                                     token_id=response.token,
                                                     usage=response.usage,
@@ -412,28 +389,13 @@ def main(
                                                 command_id=command_id,
                                                 chunk=ToolCallChunk(
                                                     tool_calls=response.tool_calls,
-                                                    model=model_id,
+                                                    model=shard_metadata.model_card.model_id,
                                                     usage=response.usage,
                                                     stats=response.stats,
                                                 ),
                                             )
                                         )
 
-                        if oom_stopped and device_rank == 0:
-                            event_sender.send(
-                                ChunkGenerated(
-                                    command_id=command_id,
-                                    chunk=ErrorChunk(
-                                        model=model_id,
-                                        error_message=(
-                                            "Generation stopped: running out of memory. "
-                                            "Please start a new conversation or compact "
-                                            "your messages."
-                                        ),
-                                    ),
-                                )
-                            )
-
                     except PrefillCancelled:
                         logger.info(f"Prefill cancelled for task {task.task_id}")
                     # can we make this more explicit?
@@ -443,7 +405,7 @@ def main(
                                 ChunkGenerated(
                                     command_id=command_id,
                                     chunk=ErrorChunk(
-                                        model=model_id,
+                                        model=shard_metadata.model_card.model_id,
                                         finish_reason="error",
                                         error_message=str(e),
                                     ),
@@ -765,7 +727,7 @@ def parse_tool_calls(
     in_tool_call = False
     tool_call_text_parts: list[str] = []
     for response in responses:
-        if not in_tool_call and response.text.startswith(tool_parser.start_parsing):
+        if response.text.startswith(tool_parser.start_parsing):
             in_tool_call = True
 
         if in_tool_call:
@@ -803,9 +765,10 @@ def parse_tool_calls(
                 )
                 yield response
 
-        else:
-            # fallthrough
-            yield response
+            continue
+
+        # fallthrough
+        yield response
 
 
 EXO_RUNNER_MUST_FAIL = "EXO RUNNER MUST FAIL"

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

@@ -58,16 +58,15 @@ def _flatten(p: dict[str, Any]) -> dict[str, str]:
     }
 
 
-def make_json_parser() -> ToolParser:
-    return ToolParser(
-        start_parsing="<tool_call>",
-        end_parsing="</tool_call>",
-        parse_tool_calls=_parse_json_calls,
-    )
+json_tool_parser = ToolParser(
+    start_parsing="<tool_call>",
+    end_parsing="</tool_call>",
+    parse_tool_calls=_parse_json_calls,
+)
 
 
 def infer_tool_parser(chat_template: str) -> ToolParser | None:
     """Attempt to auto-infer a tool parser from the chat template."""
     if "<tool_call>" in chat_template and "tool_call.name" in chat_template:
-        return make_json_parser()
+        return json_tool_parser
     return None

src/exo/worker/runner/runner_supervisor.py

@@ -1,7 +1,7 @@
 import contextlib
-import multiprocessing as mp
 import signal
 from dataclasses import dataclass, field
+from multiprocessing import Process
 from typing import Self
 
 import anyio
@@ -32,7 +32,6 @@ from exo.shared.types.worker.runners import (
 )
 from exo.shared.types.worker.shards import ShardMetadata
 from exo.utils.channels import MpReceiver, MpSender, Sender, mp_channel
-from exo.utils.task_group import TaskGroup
 from exo.worker.runner.bootstrap import entrypoint
 
 PREFILL_TIMEOUT_SECONDS = 60
@@ -43,20 +42,16 @@ DECODE_TIMEOUT_SECONDS = 5
 class RunnerSupervisor:
     shard_metadata: ShardMetadata
     bound_instance: BoundInstance
-    runner_process: mp.Process
+    runner_process: Process
     initialize_timeout: float
     _ev_recv: MpReceiver[Event]
     _task_sender: MpSender[Task]
     _event_sender: Sender[Event]
     _cancel_sender: MpSender[TaskId]
-    _tg: TaskGroup = field(default_factory=TaskGroup, init=False)
     status: RunnerStatus = field(default_factory=RunnerIdle, init=False)
     pending: dict[TaskId, anyio.Event] = field(default_factory=dict, init=False)
     completed: set[TaskId] = field(default_factory=set, init=False)
     cancelled: set[TaskId] = field(default_factory=set, init=False)
-    _cancel_watch_runner: anyio.CancelScope = field(
-        default_factory=anyio.CancelScope, init=False
-    )
 
     @classmethod
     def create(
@@ -70,7 +65,7 @@ class RunnerSupervisor:
         task_sender, task_recv = mp_channel[Task]()
         cancel_sender, cancel_recv = mp_channel[TaskId]()
 
-        runner_process = mp.Process(
+        runner_process = Process(
             target=entrypoint,
             args=(
                 bound_instance,
@@ -99,25 +94,15 @@ class RunnerSupervisor:
 
     async def run(self):
         self.runner_process.start()
-        async with self._tg as tg:
-            tg.start_soon(self._watch_runner)
-            tg.start_soon(self._forward_events)
+        await self._forward_events()
 
     def shutdown(self):
         logger.info("Runner supervisor shutting down")
-        self._tg.cancel_tasks()
-        if not self._cancel_watch_runner.cancel_called:
-            self._cancel_watch_runner.cancel()
-        with contextlib.suppress(ClosedResourceError):
-            self._ev_recv.close()
-        with contextlib.suppress(ClosedResourceError):
-            self._task_sender.close()
-        with contextlib.suppress(ClosedResourceError):
-            self._event_sender.close()
+        self._ev_recv.close()
+        self._task_sender.close()
         with contextlib.suppress(ClosedResourceError):
             self._cancel_sender.send(TaskId("CANCEL_CURRENT_TASK"))
-        with contextlib.suppress(ClosedResourceError):
-            self._cancel_sender.close()
+        self._cancel_sender.close()
         self.runner_process.join(5)
         if not self.runner_process.is_alive():
             logger.info("Runner process succesfully terminated")
@@ -166,8 +151,8 @@ class RunnerSupervisor:
             await self._check_runner(TimeoutError("cancel pipe blocked"))
 
     async def _forward_events(self):
-        try:
-            with self._ev_recv as events:
+        with self._ev_recv as events:
+            try:
                 async for event in events:
                     if isinstance(event, RunnerStatusUpdated):
                         self.status = event.runner_status
@@ -191,34 +176,25 @@ class RunnerSupervisor:
                         )
                         self.completed.add(event.task_id)
                     await self._event_sender.send(event)
-        except (ClosedResourceError, BrokenResourceError) as e:
-            await self._check_runner(e)
-        finally:
-            for tid in self.pending:
-                self.pending[tid].set()
+            except (ClosedResourceError, BrokenResourceError) as e:
+                await self._check_runner(e)
+                for tid in self.pending:
+                    self.pending[tid].set()
+        self._event_sender.close()
 
     def __del__(self) -> None:
         if self.runner_process.is_alive():
-            logger.critical("RunnerSupervisor was not stopped cleanly.")
+            logger.warning("RunnerSupervisor was not stopped cleanly.")
             with contextlib.suppress(ValueError):
                 self.runner_process.kill()
 
-    async def _watch_runner(self) -> None:
-        with self._cancel_watch_runner:
-            while True:
-                await anyio.sleep(5)
-                if not self.runner_process.is_alive():
-                    await self._check_runner(RuntimeError("Runner found to be dead"))
-
     async def _check_runner(self, e: Exception) -> None:
-        if not self._cancel_watch_runner.cancel_called:
-            self._cancel_watch_runner.cancel()
         logger.info("Checking runner's status")
         if self.runner_process.is_alive():
             logger.info("Runner was found to be alive, attempting to join process")
             await to_thread.run_sync(self.runner_process.join, 5)
         rc = self.runner_process.exitcode
-        logger.info(f"Runner exited with exit code {rc}")
+        logger.info(f"RunnerSupervisor exited with exit code {rc}")
         if rc == 0:
             return
 
@@ -231,19 +207,15 @@ class RunnerSupervisor:
         else:
             cause = f"exitcode={rc}"
 
-        logger.opt(exception=e).error(f"Runner terminated with {cause}")
+        logger.opt(exception=e).error(f"Runner terminated ({cause})")
 
         try:
-            self.status = RunnerFailed(error_message=f"Terminated ({cause})")
-            with anyio.CancelScope(shield=True):
-                await self._event_sender.send(
-                    RunnerStatusUpdated(
-                        runner_id=self.bound_instance.bound_runner_id,
-                        runner_status=RunnerFailed(
-                            error_message=f"Terminated ({cause})"
-                        ),
-                    )
+            await self._event_sender.send(
+                RunnerStatusUpdated(
+                    runner_id=self.bound_instance.bound_runner_id,
+                    runner_status=RunnerFailed(error_message=f"Terminated ({cause})"),
                 )
+            )
         except (ClosedResourceError, BrokenResourceError):
             logger.warning(
                 "Event sender already closed, unable to report runner failure"

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

@@ -15,7 +15,6 @@ from exo.shared.types.events import (
     TaskAcknowledged,
     TaskStatusUpdated,
 )
-from exo.shared.types.memory import Memory
 from exo.shared.types.tasks import (
     ConnectToGroup,
     LoadModel,
@@ -115,9 +114,7 @@ def patch_out_mlx(monkeypatch: pytest.MonkeyPatch):
     # initialize_mlx returns a mock group
     monkeypatch.setattr(mlx_runner, "initialize_mlx", make_nothin(MockGroup()))
     monkeypatch.setattr(mlx_runner, "load_mlx_items", make_nothin((1, MockTokenizer)))
-    monkeypatch.setattr(
-        mlx_runner, "warmup_inference", make_nothin((1, Memory.from_bytes(0)))
-    )
+    monkeypatch.setattr(mlx_runner, "warmup_inference", make_nothin(1))
     monkeypatch.setattr(mlx_runner, "_check_for_debug_prompts", nothin)
     monkeypatch.setattr(mlx_runner, "mx_any", make_nothin(False))
     # Mock apply_chat_template since we're using a fake tokenizer (integer 1).

tests/auto_bench.sh

@@ -17,13 +17,6 @@ git branch -r --contains "$commit" | grep -qE '^\s*origin/' || {
   exit 1
 }
 hosts=("$@")
-
-for host; do
-  ssh -T -o BatchMode=yes -o ServerAliveInterval=30 "$host@$host" \
-    "EXO_LIBP2P_NAMESPACE=$commit /nix/var/nix/profiles/default/bin/nix build github:exo-explore/exo/$commit" &
-done
-wait
-
 cleanup() {
   for host in "${hosts[@]}"; do
     ssh -T -o BatchMode=yes "$host@$host" "pkill -f bin/exo" &
@@ -33,6 +26,11 @@ cleanup() {
 }
 trap 'cleanup' EXIT INT TERM
 
+for host; do
+  ssh -T -o BatchMode=yes -o ServerAliveInterval=30 "$host@$host" \
+    "EXO_LIBP2P_NAMESPACE=$commit /nix/var/nix/profiles/default/bin/nix build github:exo-explore/exo/$commit" &
+done
+wait
 for host; do
   ssh -T -o BatchMode=yes -o ServerAliveInterval=30 "$host@$host" \
     "EXO_LIBP2P_NAMESPACE=$commit /nix/var/nix/profiles/default/bin/nix run github:exo-explore/exo/$commit" &>/dev/null &

tmp/test_trust_remote_code_attack.sh

@@ -1,125 +0,0 @@
-#!/usr/bin/env bash
-# Test that models added via API get trust_remote_code=false
-# Run this against a running exo instance.
-# Usage: ./test_trust_remote_code_attack.sh [host:port]
-
-set -uo pipefail
-
-HOST="${1:-localhost:52415}"
-MODEL_ID="KevTheHermit/security-testing"
-CUSTOM_CARDS_DIR="$HOME/.exo/custom_model_cards"
-CARD_FILE="$CUSTOM_CARDS_DIR/KevTheHermit--security-testing.toml"
-
-echo "=== Test: trust_remote_code attack via API ==="
-echo "Target: $HOST"
-echo ""
-
-# Clean up RCE proof from previous runs
-rm -f /tmp/exo-rce-proof.txt
-
-# Step 0: Clean up any stale card from previous runs
-if [ -f "$CARD_FILE" ]; then
-  echo "[0] Removing stale card from previous run ..."
-  curl -s -X DELETE \
-    "http://$HOST/models/custom/$(python3 -c 'import urllib.parse; print(urllib.parse.quote("'"$MODEL_ID"'", safe=""))')" >/dev/null
-  rm -f "$CARD_FILE"
-  echo "    Done"
-  echo ""
-fi
-
-# Step 1: Add the malicious model via API
-echo "[1] Adding model via POST /models/add ..."
-ADD_RESPONSE=$(curl -s -w "\n%{http_code}" -X POST "http://$HOST/models/add" \
-  -H "Content-Type: application/json" \
-  -d "{\"model_id\":\"$MODEL_ID\"}")
-HTTP_CODE=$(echo "$ADD_RESPONSE" | tail -1)
-BODY=$(echo "$ADD_RESPONSE" | sed '$d')
-echo "    HTTP $HTTP_CODE"
-
-if [ "$HTTP_CODE" -ge 400 ]; then
-  echo "    Model add failed (HTTP $HTTP_CODE) — that's fine if model doesn't exist on HF."
-  echo "    Response: $BODY"
-  echo ""
-  echo "RESULT: Model was rejected at add time. Attack blocked."
-  exit 0
-fi
-
-# Step 2: Verify the saved TOML has trust_remote_code = false
-echo ""
-echo "[2] Checking saved model card TOML ..."
-if [ ! -f "$CARD_FILE" ]; then
-  echo "    FAIL: Card file not found at $CARD_FILE"
-  exit 1
-fi
-
-if grep -q 'trust_remote_code = false' "$CARD_FILE"; then
-  echo "    SAFE: trust_remote_code = false (fix is active)"
-else
-  echo "    VULNERABLE: trust_remote_code is not false — remote code WILL be trusted"
-fi
-echo "    Contents:"
-cat "$CARD_FILE"
-
-# Step 3: Place the instance
-echo ""
-echo "[3] Attempting POST /place_instance ..."
-PLACE_RESPONSE=$(curl -s -w "\n%{http_code}" -X POST "http://$HOST/place_instance" \
-  -H "Content-Type: application/json" \
-  -d "{\"model_id\":\"$MODEL_ID\"}")
-PLACE_CODE=$(echo "$PLACE_RESPONSE" | tail -1)
-PLACE_BODY=$(echo "$PLACE_RESPONSE" | sed '$d')
-echo "    HTTP $PLACE_CODE"
-echo "    Response: $PLACE_BODY"
-
-# Step 3b: Send a chat completion to actually trigger tokenizer loading
-echo ""
-echo "[3b] Sending chat completion to trigger tokenizer load ..."
-CHAT_RESPONSE=$(curl -s -w "\n%{http_code}" --max-time 30 -X POST "http://$HOST/v1/chat/completions" \
-  -H "Content-Type: application/json" \
-  -d "{\"model\":\"$MODEL_ID\",\"messages\":[{\"role\":\"user\",\"content\":\"hello\"}],\"max_tokens\":1}")
-CHAT_CODE=$(echo "$CHAT_RESPONSE" | tail -1)
-CHAT_BODY=$(echo "$CHAT_RESPONSE" | sed '$d')
-echo "    HTTP $CHAT_CODE"
-echo "    Response: $CHAT_BODY"
-echo ""
-echo "[3c] Checking for RCE proof ..."
-sleep 5
-if [ -f /tmp/exo-rce-proof.txt ]; then
-  echo "    VULNERABLE: Remote code executed!"
-  echo "    Contents:"
-  cat /tmp/exo-rce-proof.txt
-else
-  echo "    SAFE: /tmp/exo-rce-proof.txt does not exist — remote code was NOT executed"
-fi
-
-# Step 4: Clean up — delete instance and custom model
-echo ""
-echo "[4] Cleaning up ..."
-
-# Find and delete any instance for this model
-INSTANCE_ID=$(curl -s "http://$HOST/state" | python3 -c "
-import sys, json
-state = json.load(sys.stdin)
-for iid, wrapper in state.get('instances', {}).items():
-    for tag, inst in wrapper.items():
-        sa = inst.get('shardAssignments', {})
-        if sa.get('modelId', '') == '$MODEL_ID':
-            print(iid)
-            sys.exit(0)
-" 2>/dev/null || true)
-
-if [ -n "$INSTANCE_ID" ]; then
-  echo "    Deleting instance $INSTANCE_ID ..."
-  curl -s -X DELETE "http://$HOST/instance/$INSTANCE_ID" >/dev/null
-  echo "    Done"
-else
-  echo "    No instance found to delete"
-fi
-
-echo "    Deleting custom model card ..."
-curl -s -X DELETE \
-  "http://$HOST/models/custom/$(python3 -c 'import urllib.parse; print(urllib.parse.quote("'"$MODEL_ID"'", safe=""))')" >/dev/null
-echo "    Done"
-
-echo ""
-echo "=== DONE ==="

uv.lock

@@ -363,7 +363,7 @@ wheels = [
 
 [[package]]
 name = "exo"
-version = "0.3.68"
+version = "0.3.0"
 source = { editable = "." }
 dependencies = [
     { name = "aiofiles", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
@@ -378,7 +378,7 @@ dependencies = [
     { name = "loguru", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "mflux", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "mlx", version = "0.30.6", source = { registry = "https://pypi.org/simple" }, extra = ["cpu"], marker = "sys_platform == 'linux'" },
-    { name = "mlx", version = "0.30.7.dev20260225+257d5692", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#257d5692fc7af6bba3b8afaeb63c549b7d1e43d5" }, marker = "sys_platform == 'darwin'" },
+    { name = "mlx", version = "0.30.7.dev20260220+13998a05", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#13998a054715edcdc93618fb1496c79c7c25ff7c" }, marker = "sys_platform == 'darwin'" },
     { name = "mlx-lm", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "msgspec", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "openai-harmony", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
@@ -451,7 +451,6 @@ dependencies = [
     { name = "huggingface-hub", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "jinja2", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "loguru", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
-    { name = "protobuf", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "tiktoken", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "transformers", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
 ]
@@ -462,7 +461,6 @@ requires-dist = [
     { name = "huggingface-hub", specifier = ">=0.33.4" },
     { name = "jinja2", specifier = ">=3.1.0" },
     { name = "loguru", specifier = ">=0.7.3" },
-    { name = "protobuf", specifier = ">=5.29.0" },
     { name = "tiktoken", specifier = ">=0.12.0" },
     { name = "transformers", specifier = ">=5.0.0" },
 ]
@@ -1025,7 +1023,7 @@ dependencies = [
     { name = "huggingface-hub", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "matplotlib", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "mlx", version = "0.30.6", source = { registry = "https://pypi.org/simple" }, extra = ["cuda13"], marker = "sys_platform == 'linux'" },
-    { name = "mlx", version = "0.30.7.dev20260225+257d5692", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#257d5692fc7af6bba3b8afaeb63c549b7d1e43d5" }, marker = "sys_platform == 'darwin'" },
+    { name = "mlx", version = "0.30.7.dev20260220+13998a05", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#13998a054715edcdc93618fb1496c79c7c25ff7c" }, marker = "sys_platform == 'darwin'" },
     { name = "numpy", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "opencv-python", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "piexif", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
@@ -1072,8 +1070,8 @@ cuda13 = [
 
 [[package]]
 name = "mlx"
-version = "0.30.7.dev20260225+257d5692"
-source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#257d5692fc7af6bba3b8afaeb63c549b7d1e43d5" }
+version = "0.30.7.dev20260220+13998a05"
+source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#13998a054715edcdc93618fb1496c79c7c25ff7c" }
 resolution-markers = [
     "sys_platform == 'darwin'",
 ]
@@ -1108,7 +1106,7 @@ version = "0.30.7"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "jinja2", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
-    { name = "mlx", version = "0.30.7.dev20260225+257d5692", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#257d5692fc7af6bba3b8afaeb63c549b7d1e43d5" }, marker = "sys_platform == 'darwin'" },
+    { name = "mlx", version = "0.30.7.dev20260220+13998a05", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#13998a054715edcdc93618fb1496c79c7c25ff7c" }, marker = "sys_platform == 'darwin'" },
     { name = "numpy", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "protobuf", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "pyyaml", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },