Point mlx-lm to davidmcc73 fork with context parallelism support

Store pre-shard head count and distributed group on each attention
layer during sharding, enabling automatic TP→CP switching at runtime
when context length exceeds a threshold.

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

app/EXO/EXO/ContentView.swift

@@ -14,6 +14,7 @@ struct ContentView: View {
     @EnvironmentObject private var networkStatusService: NetworkStatusService
     @EnvironmentObject private var localNetworkChecker: LocalNetworkChecker
     @EnvironmentObject private var updater: SparkleUpdater
+    @EnvironmentObject private var thunderboltBridgeService: ThunderboltBridgeService
     @State private var focusedNode: NodeViewModel?
     @State private var deletingInstanceIDs: Set<String> = []
     @State private var showAllNodes = false
@@ -24,6 +25,8 @@ struct ContentView: View {
     @State private var bugReportMessage: 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) {
@@ -303,6 +306,49 @@ struct ContentView: View {
                             .disabled(pendingNamespace == controller.customNamespace)
                         }
                     }
+                    VStack(alignment: .leading, spacing: 4) {
+                        Text("HuggingFace Token")
+                            .font(.caption2)
+                            .foregroundColor(.secondary)
+                        HStack {
+                            SecureField("optional", text: $pendingHFToken)
+                                .textFieldStyle(.roundedBorder)
+                                .font(.caption2)
+                                .onAppear {
+                                    pendingHFToken = controller.hfToken
+                                }
+                            Button("Save & Restart") {
+                                controller.hfToken = pendingHFToken
+                                if controller.status == .running || controller.status == .starting {
+                                    controller.restart()
+                                }
+                            }
+                            .font(.caption2)
+                            .disabled(pendingHFToken == controller.hfToken)
+                        }
+                    }
+                    Divider()
+                    HStack {
+                        Toggle(
+                            "Enable Image Models (experimental)", isOn: $pendingEnableImageModels
+                        )
+                        .toggleStyle(.switch)
+                        .font(.caption2)
+                        .onAppear {
+                            pendingEnableImageModels = controller.enableImageModels
+                        }
+
+                        Spacer()
+
+                        Button("Save & Restart") {
+                            controller.enableImageModels = pendingEnableImageModels
+                            if controller.status == .running || controller.status == .starting {
+                                controller.restart()
+                            }
+                        }
+                        .font(.caption2)
+                        .disabled(pendingEnableImageModels == controller.enableImageModels)
+                    }
                     HoverButton(title: "Check for Updates", small: true) {
                         updater.checkForUpdates()
                     }
@@ -423,6 +469,44 @@ struct ContentView: View {
         }
     }
 
+    /// 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) {
@@ -465,6 +549,7 @@ struct ContentView: View {
                     Text(thunderboltStatusText)
                         .font(.caption2)
                         .foregroundColor(thunderboltStatusColor)
+                    clusterThunderboltBridgeView
                     interfaceIpList
                     rdmaStatusView
                     sendBugReportButton

app/EXO/EXO/EXOApp.swift

@@ -21,6 +21,7 @@ struct EXOApp: App {
     @StateObject private var networkStatusService: NetworkStatusService
     @StateObject private var localNetworkChecker: LocalNetworkChecker
     @StateObject private var updater: SparkleUpdater
+    @StateObject private var thunderboltBridgeService: ThunderboltBridgeService
     private let terminationObserver: TerminationObserver
     private let ciContext = CIContext(options: nil)
 
@@ -41,10 +42,13 @@ struct EXOApp: App {
         let localNetwork = LocalNetworkChecker()
         _localNetworkChecker = StateObject(wrappedValue: localNetwork)
         _updater = StateObject(wrappedValue: updater)
+        let thunderboltBridge = ThunderboltBridgeService(clusterStateService: service)
+        _thunderboltBridgeService = StateObject(wrappedValue: thunderboltBridge)
         enableLaunchAtLoginIfNeeded()
-        NetworkSetupHelper.ensureLaunchDaemonInstalled()
-        // Check local network access BEFORE launching exo
-        localNetwork.check()
+        // Install LaunchDaemon to disable Thunderbolt Bridge on startup (prevents network loops)
+        NetworkSetupHelper.promptAndInstallIfNeeded()
+        // Check local network access periodically (warning disappears when user grants permission)
+        localNetwork.startPeriodicChecking(interval: 10)
         controller.scheduleLaunch(after: 15)
         service.startPolling()
         networkStatus.startPolling()
@@ -58,6 +62,7 @@ struct EXOApp: App {
                 .environmentObject(networkStatusService)
                 .environmentObject(localNetworkChecker)
                 .environmentObject(updater)
+                .environmentObject(thunderboltBridgeService)
         } label: {
             menuBarIcon
         }
@@ -130,6 +135,7 @@ struct EXOApp: App {
                 "Failed to register EXO for launch at login: \(error.localizedDescription)")
         }
     }
+
 }
 
 /// Helper for managing EXO's launch-at-login registration

app/EXO/EXO/ExoProcessController.swift

@@ -3,6 +3,8 @@ import Combine
 import Foundation
 
 private let customNamespaceKey = "EXOCustomNamespace"
+private let hfTokenKey = "EXOHFToken"
+private let enableImageModelsKey = "EXOEnableImageModels"
 
 @MainActor
 final class ExoProcessController: ObservableObject {
@@ -37,6 +39,22 @@ final class ExoProcessController: ObservableObject {
             UserDefaults.standard.set(customNamespace, forKey: customNamespaceKey)
         }
     }
+    @Published var hfToken: String = {
+        return UserDefaults.standard.string(forKey: hfTokenKey) ?? ""
+    }()
+    {
+        didSet {
+            UserDefaults.standard.set(hfToken, forKey: hfTokenKey)
+        }
+    }
+    @Published var enableImageModels: Bool = {
+        return UserDefaults.standard.bool(forKey: enableImageModelsKey)
+    }()
+    {
+        didSet {
+            UserDefaults.standard.set(enableImageModels, forKey: enableImageModelsKey)
+        }
+    }
 
     private var process: Process?
     private var runtimeDirectoryURL: URL?
@@ -191,6 +209,12 @@ final class ExoProcessController: ObservableObject {
         var environment = ProcessInfo.processInfo.environment
         environment["EXO_RUNTIME_DIR"] = runtimeURL.path
         environment["EXO_LIBP2P_NAMESPACE"] = computeNamespace()
+        if !hfToken.isEmpty {
+            environment["HF_TOKEN"] = hfToken
+        }
+        if enableImageModels {
+            environment["EXO_ENABLE_IMAGE_MODELS"] = "true"
+        }
 
         var paths: [String] = []
         if let existing = environment["PATH"], !existing.isEmpty {

app/EXO/EXO/Models/ClusterState.swift

@@ -5,17 +5,43 @@ import Foundation
 struct ClusterState: Decodable {
     let instances: [String: ClusterInstance]
     let runners: [String: RunnerStatusSummary]
-    let nodeProfiles: [String: NodeProfile]
     let tasks: [String: ClusterTask]
     let topology: Topology?
     let downloads: [String: [NodeDownloadStatus]]
+    let thunderboltBridgeCycles: [[String]]
+
+    // Granular node state (split from the old nodeProfiles)
+    let nodeIdentities: [String: NodeIdentity]
+    let nodeMemory: [String: MemoryInfo]
+    let nodeSystem: [String: SystemInfo]
+    let nodeThunderboltBridge: [String: ThunderboltBridgeStatus]
+
+    /// Computed property for backwards compatibility - merges granular state into NodeProfile
+    var nodeProfiles: [String: NodeProfile] {
+        var profiles: [String: NodeProfile] = [:]
+        let allNodeIds = Set(nodeIdentities.keys)
+            .union(nodeMemory.keys)
+            .union(nodeSystem.keys)
+        for nodeId in allNodeIds {
+            let identity = nodeIdentities[nodeId]
+            let memory = nodeMemory[nodeId]
+            let system = nodeSystem[nodeId]
+            profiles[nodeId] = NodeProfile(
+                modelId: identity?.modelId,
+                chipId: identity?.chipId,
+                friendlyName: identity?.friendlyName,
+                memory: memory,
+                system: system
+            )
+        }
+        return profiles
+    }
 
     init(from decoder: Decoder) throws {
         let container = try decoder.container(keyedBy: CodingKeys.self)
         let rawInstances = try container.decode([String: TaggedInstance].self, forKey: .instances)
         self.instances = rawInstances.mapValues(\.instance)
         self.runners = try container.decode([String: RunnerStatusSummary].self, forKey: .runners)
-        self.nodeProfiles = try container.decode([String: NodeProfile].self, forKey: .nodeProfiles)
         let rawTasks =
             try container.decodeIfPresent([String: TaggedTask].self, forKey: .tasks) ?? [:]
         self.tasks = rawTasks.compactMapValues(\.task)
@@ -24,15 +50,34 @@ struct ClusterState: Decodable {
             try container.decodeIfPresent([String: [TaggedNodeDownload]].self, forKey: .downloads)
             ?? [:]
         self.downloads = rawDownloads.mapValues { $0.compactMap(\.status) }
+        self.thunderboltBridgeCycles =
+            try container.decodeIfPresent([[String]].self, forKey: .thunderboltBridgeCycles) ?? []
+
+        // Granular node state
+        self.nodeIdentities =
+            try container.decodeIfPresent([String: NodeIdentity].self, forKey: .nodeIdentities)
+            ?? [:]
+        self.nodeMemory =
+            try container.decodeIfPresent([String: MemoryInfo].self, forKey: .nodeMemory) ?? [:]
+        self.nodeSystem =
+            try container.decodeIfPresent([String: SystemInfo].self, forKey: .nodeSystem) ?? [:]
+        self.nodeThunderboltBridge =
+            try container.decodeIfPresent(
+                [String: ThunderboltBridgeStatus].self, forKey: .nodeThunderboltBridge
+            ) ?? [:]
     }
 
     private enum CodingKeys: String, CodingKey {
         case instances
         case runners
-        case nodeProfiles
         case topology
         case tasks
         case downloads
+        case thunderboltBridgeCycles
+        case nodeIdentities
+        case nodeMemory
+        case nodeSystem
+        case nodeThunderboltBridge
     }
 }
 
@@ -102,6 +147,18 @@ struct NodeProfile: Decodable {
     let system: SystemInfo?
 }
 
+struct NodeIdentity: Decodable {
+    let modelId: String?
+    let chipId: String?
+    let friendlyName: String?
+}
+
+struct ThunderboltBridgeStatus: Decodable {
+    let enabled: Bool
+    let exists: Bool
+    let serviceName: String?
+}
+
 struct MemoryInfo: Decodable {
     let ramTotal: MemoryValue?
     let ramAvailable: MemoryValue?
@@ -120,16 +177,51 @@ struct SystemInfo: Decodable {
 }
 
 struct Topology: Decodable {
-    let nodes: [TopologyNode]
-    let connections: [TopologyConnection]?
+    /// Node IDs in the topology
+    let nodes: [String]
+    /// Flattened list of connections (source -> sink pairs)
+    let connections: [TopologyConnection]
+
+    init(from decoder: Decoder) throws {
+        let container = try decoder.container(keyedBy: CodingKeys.self)
+        self.nodes = try container.decodeIfPresent([String].self, forKey: .nodes) ?? []
+
+        // Connections come as nested map: { source: { sink: [edges] } }
+        // We flatten to array of (source, sink) pairs
+        var flatConnections: [TopologyConnection] = []
+        if let nested = try container.decodeIfPresent(
+            [String: [String: [AnyCodable]]].self, forKey: .connections
+        ) {
+            for (source, sinks) in nested {
+                for sink in sinks.keys {
+                    flatConnections.append(
+                        TopologyConnection(localNodeId: source, sendBackNodeId: sink))
+                }
+            }
+        }
+        self.connections = flatConnections
+    }
+
+    private enum CodingKeys: String, CodingKey {
+        case nodes
+        case connections
+    }
 }
 
-struct TopologyNode: Decodable {
-    let nodeId: String
-    let nodeProfile: NodeProfile
+/// Placeholder for decoding arbitrary JSON values we don't need to inspect
+private struct AnyCodable: Decodable {
+    init(from decoder: Decoder) throws {
+        // Just consume the value without storing it
+        _ = try? decoder.singleValueContainer().decode(Bool.self)
+        _ = try? decoder.singleValueContainer().decode(Int.self)
+        _ = try? decoder.singleValueContainer().decode(Double.self)
+        _ = try? decoder.singleValueContainer().decode(String.self)
+        _ = try? decoder.singleValueContainer().decode([AnyCodable].self)
+        _ = try? decoder.singleValueContainer().decode([String: AnyCodable].self)
+    }
 }
 
-struct TopologyConnection: Decodable {
+struct TopologyConnection {
     let localNodeId: String
     let sendBackNodeId: String
 }

app/EXO/EXO/Services/BugReportService.swift

@@ -55,12 +55,16 @@ struct BugReportService {
         let stateData = try await stateResult
         let eventsData = try await eventsResult
 
+        // Extract cluster TB bridge status from exo state
+        let clusterTbBridgeStatus = extractClusterTbBridgeStatus(from: stateData)
+
         let reportJSON = makeReportJson(
             timestamp: timestamp,
             hostName: hostName,
             ifconfig: ifconfigText,
             debugInfo: debugInfo,
-            isManual: isManual
+            isManual: isManual,
+            clusterTbBridgeStatus: clusterTbBridgeStatus
         )
 
         let uploads: [(path: String, data: Data?)] = [
@@ -178,18 +182,19 @@ struct BugReportService {
     }
 
     private func readThunderboltBridgeDisabled() -> Bool? {
-        let result = runCommand([
-            "/usr/sbin/networksetup", "-getnetworkserviceenabled", "Thunderbolt Bridge",
-        ])
-        guard result.exitCode == 0 else { return nil }
-        let output = result.output.lowercased()
-        if output.contains("enabled") {
-            return false
+        // Dynamically find the Thunderbolt Bridge service (don't assume the name)
+        guard let serviceName = ThunderboltBridgeDetector.findThunderboltBridgeServiceName() else {
+            // No bridge containing Thunderbolt interfaces exists
+            return nil
         }
-        if output.contains("disabled") {
-            return true
+
+        guard let isEnabled = ThunderboltBridgeDetector.isServiceEnabled(serviceName: serviceName)
+        else {
+            return nil
         }
-        return nil
+
+        // Return true if disabled, false if enabled
+        return !isEnabled
     }
 
     private func readInterfaces() -> [DebugInfo.InterfaceStatus] {
@@ -268,11 +273,12 @@ struct BugReportService {
         hostName: String,
         ifconfig: String,
         debugInfo: DebugInfo,
-        isManual: Bool
+        isManual: Bool,
+        clusterTbBridgeStatus: [[String: Any]]?
     ) -> Data? {
         let system = readSystemMetadata()
         let exo = readExoMetadata()
-        let payload: [String: Any] = [
+        var payload: [String: Any] = [
             "timestamp": timestamp,
             "host": hostName,
             "ifconfig": ifconfig,
@@ -282,9 +288,38 @@ struct BugReportService {
             "exo_commit": exo.commit as Any,
             "report_type": isManual ? "manual" : "automated",
         ]
+        if let tbStatus = clusterTbBridgeStatus {
+            payload["cluster_thunderbolt_bridge"] = tbStatus
+        }
         return try? JSONSerialization.data(withJSONObject: payload, options: [.prettyPrinted])
     }
 
+    /// Extracts cluster-wide Thunderbolt Bridge status from exo state JSON
+    private func extractClusterTbBridgeStatus(from stateData: Data?) -> [[String: Any]]? {
+        guard let data = stateData,
+            let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any],
+            let nodeThunderboltBridge = json["node_thunderbolt_bridge"] as? [String: [String: Any]]
+        else {
+            return nil
+        }
+
+        var result: [[String: Any]] = []
+        for (nodeId, status) in nodeThunderboltBridge {
+            var entry: [String: Any] = ["node_id": nodeId]
+            if let enabled = status["enabled"] as? Bool {
+                entry["enabled"] = enabled
+            }
+            if let exists = status["exists"] as? Bool {
+                entry["exists"] = exists
+            }
+            if let serviceName = status["service_name"] as? String {
+                entry["service_name"] = serviceName
+            }
+            result.append(entry)
+        }
+        return result.isEmpty ? nil : result
+    }
+
     private func readSystemMetadata() -> [String: Any] {
         let hostname = safeRunCommand(["/bin/hostname"])
         let computerName = safeRunCommand(["/usr/sbin/scutil", "--get", "ComputerName"])

app/EXO/EXO/Services/LocalNetworkChecker.swift

@@ -41,6 +41,7 @@ final class LocalNetworkChecker: ObservableObject {
 
     private var connection: NWConnection?
     private var checkTask: Task<Void, Never>?
+    private var periodicTask: Task<Void, Never>?
 
     /// Whether we've completed at least one check (stored in UserDefaults)
     private var hasCompletedInitialCheck: Bool {
@@ -48,10 +49,39 @@ final class LocalNetworkChecker: ObservableObject {
         set { UserDefaults.standard.set(newValue, forKey: Self.hasCompletedInitialCheckKey) }
     }
 
-    /// Checks if local network access is working.
+    /// Checks if local network access is working (one-time check).
     func check() {
+        performCheck()
+    }
+
+    /// Starts periodic checking of local network access.
+    /// Re-checks every `interval` seconds so the warning disappears when user grants permission.
+    func startPeriodicChecking(interval: TimeInterval = 10) {
+        stopPeriodicChecking()
+        // Do an immediate check first
+        performCheck()
+        // Then schedule periodic checks
+        periodicTask = Task { [weak self] in
+            while !Task.isCancelled {
+                try? await Task.sleep(nanoseconds: UInt64(interval * 1_000_000_000))
+                guard !Task.isCancelled else { break }
+                self?.performCheck()
+            }
+        }
+    }
+
+    /// Stops periodic checking.
+    func stopPeriodicChecking() {
+        periodicTask?.cancel()
+        periodicTask = nil
+    }
+
+    private func performCheck() {
         checkTask?.cancel()
-        status = .checking
+        // Only show "checking" status on first check to avoid UI flicker
+        if status == .unknown {
+            status = .checking
+        }
 
         // Use longer timeout on first launch to allow time for permission prompt
         let isFirstCheck = !hasCompletedInitialCheck
@@ -60,12 +90,15 @@ final class LocalNetworkChecker: ObservableObject {
         checkTask = Task { [weak self] in
             guard let self else { return }
 
-            Self.logger.info("Checking local network connectivity (first check: \(isFirstCheck))")
+            Self.logger.debug("Checking local network connectivity (first check: \(isFirstCheck))")
             let result = await self.checkConnectivity(timeout: timeout)
             self.status = result
             self.hasCompletedInitialCheck = true
 
-            Self.logger.info("Local network check complete: \(result.displayText)")
+            // Only log on state changes or first check to reduce noise
+            if isFirstCheck || result != self.status {
+                Self.logger.info("Local network check: \(result.displayText)")
+            }
         }
     }
 
@@ -141,6 +174,7 @@ final class LocalNetworkChecker: ObservableObject {
     }
 
     func stop() {
+        stopPeriodicChecking()
         checkTask?.cancel()
         checkTask = nil
         connection?.cancel()

app/EXO/EXO/Services/NetworkSetupHelper.swift

@@ -7,8 +7,11 @@ enum NetworkSetupHelper {
     private static let daemonLabel = "io.exo.networksetup"
     private static let scriptDestination =
         "/Library/Application Support/EXO/disable_bridge.sh"
+    // Legacy script path from older versions
+    private static let legacyScriptDestination =
+        "/Library/Application Support/EXO/disable_bridge_enable_dhcp.sh"
     private static let plistDestination = "/Library/LaunchDaemons/io.exo.networksetup.plist"
-    private static let requiredStartInterval: Int = 1791
+    private static let requiredStartInterval: Int = 1786
 
     private static let setupScript = """
         #!/usr/bin/env bash
@@ -28,19 +31,69 @@ enum NetworkSetupHelper {
         # Remove Thunderbolt Bridge from VirtualNetworkInterfaces in preferences.plist
         /usr/libexec/PlistBuddy -c "Delete :VirtualNetworkInterfaces:Bridge:bridge0" "$PREFS" 2>/dev/null || true
 
+        networksetup -listlocations | grep -q exo || {
+          networksetup -createlocation exo
+        }
+
+        networksetup -switchtolocation exo
+        networksetup -listallhardwareports \\
+          | awk -F': ' '/Hardware Port: / {print $2}' \\
+          | while IFS=":" read -r name; do
+              case "$name" in
+                "Ethernet Adapter"*)
+                        ;;
+                "Thunderbolt Bridge")
+                        ;;
+                "Thunderbolt "*)
+                  networksetup -listallnetworkservices \\
+                    | grep -q "EXO $name" \\
+                      || networksetup -createnetworkservice "EXO $name" "$name" 2>/dev/null \\
+                      || continue
+                  networksetup -setdhcp "EXO $name"
+                        ;;
+                *)
+                  networksetup -listallnetworkservices \\
+                    | grep -q "$name" \\
+                      || networksetup -createnetworkservice "$name" "$name" 2>/dev/null \\
+                      || continue
+                        ;;
+              esac
+            done
+
         networksetup -listnetworkservices | grep -q "Thunderbolt Bridge" && {
           networksetup -setnetworkserviceenabled "Thunderbolt Bridge" off
         } || true
         """
 
-    static func ensureLaunchDaemonInstalled() {
+    /// Prompts user and installs the LaunchDaemon if not already installed.
+    /// Shows an alert explaining what will be installed before requesting admin privileges.
+    static func promptAndInstallIfNeeded() {
         // Use .utility priority to match NSAppleScript's internal QoS and avoid priority inversion
         Task.detached(priority: .utility) {
+            // If already correctly installed, skip
+            if daemonAlreadyInstalled() {
+                return
+            }
+
+            // Show alert on main thread
+            let shouldInstall = await MainActor.run {
+                let alert = NSAlert()
+                alert.messageText = "EXO Network Configuration"
+                alert.informativeText =
+                    "EXO needs to install a system service to automatically disable Thunderbolt Bridge on startup. This prevents network loops when connecting multiple Macs via Thunderbolt.\n\nYou will be prompted for your administrator password."
+                alert.alertStyle = .informational
+                alert.addButton(withTitle: "Install")
+                alert.addButton(withTitle: "Not Now")
+                return alert.runModal() == .alertFirstButtonReturn
+            }
+
+            guard shouldInstall else {
+                logger.info("User deferred network setup daemon installation")
+                return
+            }
+
             do {
-                if daemonAlreadyInstalled() {
-                    return
-                }
-                try await installLaunchDaemon()
+                try installLaunchDaemon()
                 logger.info("Network setup launch daemon installed and started")
             } catch {
                 logger.error(
@@ -63,48 +116,9 @@ enum NetworkSetupHelper {
     static func hasInstalledComponents() -> Bool {
         let manager = FileManager.default
         let scriptExists = manager.fileExists(atPath: scriptDestination)
+        let legacyScriptExists = manager.fileExists(atPath: legacyScriptDestination)
         let plistExists = manager.fileExists(atPath: plistDestination)
-        return scriptExists || plistExists
-    }
-
-    private static func makeUninstallScript() -> String {
-        """
-        set -euo pipefail
-
-        LABEL="\(daemonLabel)"
-        SCRIPT_DEST="\(scriptDestination)"
-        PLIST_DEST="\(plistDestination)"
-        LOG_OUT="/var/log/\(daemonLabel).log"
-        LOG_ERR="/var/log/\(daemonLabel).err.log"
-
-        # Unload the LaunchDaemon if running
-        launchctl bootout system/"$LABEL" 2>/dev/null || true
-
-        # Remove LaunchDaemon plist
-        rm -f "$PLIST_DEST"
-
-        # Remove the script and parent directory if empty
-        rm -f "$SCRIPT_DEST"
-        rmdir "$(dirname "$SCRIPT_DEST")" 2>/dev/null || true
-
-        # Remove log files
-        rm -f "$LOG_OUT" "$LOG_ERR"
-
-        # Switch back to Automatic network location
-        networksetup -switchtolocation Automatic 2>/dev/null || true
-
-        # Delete the exo network location if it exists
-        networksetup -listlocations | grep -q '^exo$' && {
-          networksetup -deletelocation exo 2>/dev/null || true
-        } || true
-
-        # Re-enable Thunderbolt Bridge if it exists
-        networksetup -listnetworkservices | grep -q "Thunderbolt Bridge" && {
-          networksetup -setnetworkserviceenabled "Thunderbolt Bridge" on 2>/dev/null || true
-        } || true
-
-        echo "EXO network components removed successfully"
-        """
+        return scriptExists || legacyScriptExists || plistExists
     }
 
     private static func daemonAlreadyInstalled() -> Bool {
@@ -140,7 +154,7 @@ enum NetworkSetupHelper {
         return true
     }
 
-    private static func installLaunchDaemon() async throws {
+    private static func installLaunchDaemon() throws {
         let installerScript = makeInstallerScript()
         try runShellAsAdmin(installerScript)
     }
@@ -151,8 +165,19 @@ enum NetworkSetupHelper {
 
         LABEL="\(daemonLabel)"
         SCRIPT_DEST="\(scriptDestination)"
+        LEGACY_SCRIPT_DEST="\(legacyScriptDestination)"
         PLIST_DEST="\(plistDestination)"
+        LOG_OUT="/var/log/\(daemonLabel).log"
+        LOG_ERR="/var/log/\(daemonLabel).err.log"
 
+        # First, completely remove any existing installation
+        launchctl bootout system/"$LABEL" 2>/dev/null || true
+        rm -f "$PLIST_DEST"
+        rm -f "$SCRIPT_DEST"
+        rm -f "$LEGACY_SCRIPT_DEST"
+        rm -f "$LOG_OUT" "$LOG_ERR"
+
+        # Install fresh
         mkdir -p "$(dirname "$SCRIPT_DEST")"
 
         cat > "$SCRIPT_DEST" <<'EOF_SCRIPT'
@@ -184,13 +209,81 @@ enum NetworkSetupHelper {
         </plist>
         EOF_PLIST
 
-        launchctl bootout system/"$LABEL" >/dev/null 2>&1 || true
         launchctl bootstrap system "$PLIST_DEST"
         launchctl enable system/"$LABEL"
         launchctl kickstart -k system/"$LABEL"
         """
     }
 
+    private static func makeUninstallScript() -> String {
+        """
+        set -euo pipefail
+
+        LABEL="\(daemonLabel)"
+        SCRIPT_DEST="\(scriptDestination)"
+        LEGACY_SCRIPT_DEST="\(legacyScriptDestination)"
+        PLIST_DEST="\(plistDestination)"
+        LOG_OUT="/var/log/\(daemonLabel).log"
+        LOG_ERR="/var/log/\(daemonLabel).err.log"
+
+        # Unload the LaunchDaemon if running
+        launchctl bootout system/"$LABEL" 2>/dev/null || true
+
+        # Remove LaunchDaemon plist
+        rm -f "$PLIST_DEST"
+
+        # Remove the script (current and legacy paths) and parent directory if empty
+        rm -f "$SCRIPT_DEST"
+        rm -f "$LEGACY_SCRIPT_DEST"
+        rmdir "$(dirname "$SCRIPT_DEST")" 2>/dev/null || true
+
+        # Remove log files
+        rm -f "$LOG_OUT" "$LOG_ERR"
+
+        # Switch back to Automatic network location
+        networksetup -switchtolocation Automatic 2>/dev/null || true
+
+        # Delete the exo network location if it exists
+        networksetup -listlocations | grep -q '^exo$' && {
+          networksetup -deletelocation exo 2>/dev/null || true
+        } || true
+
+        # Re-enable any Thunderbolt Bridge service if it exists
+        # We find it dynamically by looking for bridges containing Thunderbolt interfaces
+        find_and_enable_thunderbolt_bridge() {
+          # Get Thunderbolt interface devices from hardware ports
+          tb_devices=$(networksetup -listallhardwareports 2>/dev/null | awk '
+            /^Hardware Port:/ { port = tolower(substr($0, 16)) }
+            /^Device:/ { if (port ~ /thunderbolt/) print substr($0, 9) }
+          ')
+          [ -z "$tb_devices" ] && return 0
+
+          # For each bridge device, check if it contains Thunderbolt interfaces
+          for bridge in bridge0 bridge1 bridge2; do
+            members=$(ifconfig "$bridge" 2>/dev/null | awk '/member:/ {print $2}')
+            [ -z "$members" ] && continue
+
+            for tb_dev in $tb_devices; do
+              if echo "$members" | grep -qx "$tb_dev"; then
+                # Find the service name for this bridge device
+                service_name=$(networksetup -listnetworkserviceorder 2>/dev/null | awk -v dev="$bridge" '
+                  /^\\([0-9*]/ { gsub(/^\\([0-9*]+\\) /, ""); svc = $0 }
+                  /Device:/ && $0 ~ dev { print svc; exit }
+                ')
+                if [ -n "$service_name" ]; then
+                  networksetup -setnetworkserviceenabled "$service_name" on 2>/dev/null || true
+                  return 0
+                fi
+              fi
+            done
+          done
+        }
+        find_and_enable_thunderbolt_bridge
+
+        echo "EXO network components removed successfully"
+        """
+    }
+
     private static func runShellAsAdmin(_ script: String) throws {
         let escapedScript =
             script

app/EXO/EXO/Services/NetworkStatusService.swift

@@ -153,22 +153,18 @@ private struct NetworkStatusFetcher {
     }
 
     private func readThunderboltBridgeState() -> ThunderboltState? {
-        let result = runCommand(["networksetup", "-getnetworkserviceenabled", "Thunderbolt Bridge"])
-        guard result.exitCode == 0 else {
-            let lower = result.output.lowercased() + result.error.lowercased()
-            if lower.contains("not a recognized network service") {
-                return .deleted
-            }
+        // Dynamically find the Thunderbolt Bridge service (don't assume the name)
+        guard let serviceName = ThunderboltBridgeDetector.findThunderboltBridgeServiceName() else {
+            // No bridge containing Thunderbolt interfaces exists
+            return .deleted
+        }
+
+        guard let isEnabled = ThunderboltBridgeDetector.isServiceEnabled(serviceName: serviceName)
+        else {
             return nil
         }
-        let output = result.output.lowercased()
-        if output.contains("enabled") {
-            return .enabled
-        }
-        if output.contains("disabled") {
-            return .disabled
-        }
-        return nil
+
+        return isEnabled ? .enabled : .disabled
     }
 
     private func readBridgeInactive() -> Bool? {

app/EXO/EXO/Services/ThunderboltBridgeDetector.swift

@@ -0,0 +1,194 @@
+import Foundation
+import os.log
+
+/// Utility for dynamically detecting Thunderbolt Bridge network services.
+/// This mirrors the Python logic in info_gatherer.py - we never assume the service
+/// is named "Thunderbolt Bridge", instead we find bridges containing Thunderbolt interfaces.
+enum ThunderboltBridgeDetector {
+    private static let logger = Logger(
+        subsystem: "io.exo.EXO", category: "ThunderboltBridgeDetector")
+
+    struct CommandResult {
+        let exitCode: Int32
+        let output: String
+        let error: String
+    }
+
+    /// Find the network service name of a bridge containing Thunderbolt interfaces.
+    /// Returns nil if no such bridge exists.
+    static func findThunderboltBridgeServiceName() -> String? {
+        // 1. Get all Thunderbolt interface devices (e.g., en2, en3)
+        guard let thunderboltDevices = getThunderboltDevices(), !thunderboltDevices.isEmpty else {
+            logger.debug("No Thunderbolt devices found")
+            return nil
+        }
+        logger.debug("Found Thunderbolt devices: \(thunderboltDevices.joined(separator: ", "))")
+
+        // 2. Get bridge services from network service order
+        guard let bridgeServices = getBridgeServices(), !bridgeServices.isEmpty else {
+            logger.debug("No bridge services found")
+            return nil
+        }
+        logger.debug("Found bridge services: \(bridgeServices.keys.joined(separator: ", "))")
+
+        // 3. Find a bridge that contains Thunderbolt interfaces
+        for (bridgeDevice, serviceName) in bridgeServices {
+            let members = getBridgeMembers(bridgeDevice: bridgeDevice)
+            logger.debug(
+                "Bridge \(bridgeDevice) (\(serviceName)) has members: \(members.joined(separator: ", "))"
+            )
+
+            // Check if any Thunderbolt device is a member of this bridge
+            if !members.isDisjoint(with: thunderboltDevices) {
+                logger.info(
+                    "Found Thunderbolt Bridge service: '\(serviceName)' (device: \(bridgeDevice))")
+                return serviceName
+            }
+        }
+
+        logger.debug("No bridge found containing Thunderbolt interfaces")
+        return nil
+    }
+
+    /// Get Thunderbolt interface device names (e.g., en2, en3) from hardware ports.
+    private static func getThunderboltDevices() -> Set<String>? {
+        let result = runCommand(["networksetup", "-listallhardwareports"])
+        guard result.exitCode == 0 else {
+            logger.warning("networksetup -listallhardwareports failed: \(result.error)")
+            return nil
+        }
+
+        var thunderboltDevices: Set<String> = []
+        var currentPort: String?
+
+        for line in result.output.components(separatedBy: .newlines) {
+            let trimmed = line.trimmingCharacters(in: .whitespaces)
+            if trimmed.hasPrefix("Hardware Port:") {
+                currentPort = String(trimmed.dropFirst("Hardware Port:".count)).trimmingCharacters(
+                    in: .whitespaces)
+            } else if trimmed.hasPrefix("Device:"), let port = currentPort {
+                let device = String(trimmed.dropFirst("Device:".count)).trimmingCharacters(
+                    in: .whitespaces)
+                if port.lowercased().contains("thunderbolt") {
+                    thunderboltDevices.insert(device)
+                }
+                currentPort = nil
+            }
+        }
+
+        return thunderboltDevices
+    }
+
+    /// Get mapping of bridge device -> service name from network service order.
+    private static func getBridgeServices() -> [String: String]? {
+        let result = runCommand(["networksetup", "-listnetworkserviceorder"])
+        guard result.exitCode == 0 else {
+            logger.warning("networksetup -listnetworkserviceorder failed: \(result.error)")
+            return nil
+        }
+
+        // Parse service order to find bridge devices and their service names
+        // Format: "(1) Service Name\n(Hardware Port: ..., Device: bridge0)\n"
+        var bridgeServices: [String: String] = [:]
+        var currentService: String?
+
+        for line in result.output.components(separatedBy: .newlines) {
+            let trimmed = line.trimmingCharacters(in: .whitespaces)
+
+            // Match "(N) Service Name" or "(*) Service Name" (disabled)
+            // but NOT "(Hardware Port: ...)" lines
+            if trimmed.hasPrefix("("), trimmed.contains(")"),
+                !trimmed.hasPrefix("(Hardware Port:")
+            {
+                if let parenEnd = trimmed.firstIndex(of: ")") {
+                    let afterParen = trimmed.index(after: parenEnd)
+                    if afterParen < trimmed.endIndex {
+                        currentService =
+                            String(trimmed[afterParen...])
+                            .trimmingCharacters(in: .whitespaces)
+                    }
+                }
+            }
+            // Match "(Hardware Port: ..., Device: bridgeX)"
+            else if let service = currentService, trimmed.contains("Device: bridge") {
+                // Extract device name from "..., Device: bridge0)"
+                if let deviceRange = trimmed.range(of: "Device: ") {
+                    let afterDevice = trimmed[deviceRange.upperBound...]
+                    if let parenIndex = afterDevice.firstIndex(of: ")") {
+                        let device = String(afterDevice[..<parenIndex])
+                        bridgeServices[device] = service
+                    }
+                }
+            }
+        }
+
+        return bridgeServices
+    }
+
+    /// Get member interfaces of a bridge device via ifconfig.
+    private static func getBridgeMembers(bridgeDevice: String) -> Set<String> {
+        let result = runCommand(["ifconfig", bridgeDevice])
+        guard result.exitCode == 0 else {
+            logger.debug("ifconfig \(bridgeDevice) failed")
+            return []
+        }
+
+        var members: Set<String> = []
+        for line in result.output.components(separatedBy: .newlines) {
+            let trimmed = line.trimmingCharacters(in: .whitespaces)
+            if trimmed.hasPrefix("member:") {
+                let parts = trimmed.split(separator: " ")
+                if parts.count > 1 {
+                    members.insert(String(parts[1]))
+                }
+            }
+        }
+
+        return members
+    }
+
+    /// Check if a network service is enabled.
+    static func isServiceEnabled(serviceName: String) -> Bool? {
+        let result = runCommand(["networksetup", "-getnetworkserviceenabled", serviceName])
+        guard result.exitCode == 0 else {
+            logger.warning("Failed to check if '\(serviceName)' is enabled: \(result.error)")
+            return nil
+        }
+
+        let output = result.output.lowercased().trimmingCharacters(in: .whitespacesAndNewlines)
+        if output.contains("enabled") {
+            return true
+        }
+        if output.contains("disabled") {
+            return false
+        }
+        return nil
+    }
+
+    private static func runCommand(_ arguments: [String]) -> CommandResult {
+        let process = Process()
+        process.launchPath = "/usr/bin/env"
+        process.arguments = arguments
+
+        let stdout = Pipe()
+        let stderr = Pipe()
+        process.standardOutput = stdout
+        process.standardError = stderr
+
+        do {
+            try process.run()
+        } catch {
+            return CommandResult(exitCode: -1, output: "", error: error.localizedDescription)
+        }
+        process.waitUntilExit()
+
+        let outputData = stdout.fileHandleForReading.readDataToEndOfFile()
+        let errorData = stderr.fileHandleForReading.readDataToEndOfFile()
+
+        return CommandResult(
+            exitCode: process.terminationStatus,
+            output: String(decoding: outputData, as: UTF8.self),
+            error: String(decoding: errorData, as: UTF8.self)
+        )
+    }
+}

app/EXO/EXO/Services/ThunderboltBridgeService.swift

@@ -0,0 +1,258 @@
+import AppKit
+import Combine
+import Foundation
+import Security
+import SystemConfiguration
+import os.log
+
+@MainActor
+final class ThunderboltBridgeService: ObservableObject {
+    private static let logger = Logger(subsystem: "io.exo.EXO", category: "ThunderboltBridge")
+
+    @Published private(set) var detectedCycle: [String]?
+    @Published private(set) var hasPromptedForCurrentCycle = false
+    @Published private(set) var lastError: String?
+
+    private weak var clusterStateService: ClusterStateService?
+    private var cancellables = Set<AnyCancellable>()
+    private var previousCycleSignature: String?
+
+    init(clusterStateService: ClusterStateService) {
+        self.clusterStateService = clusterStateService
+        setupObserver()
+    }
+
+    private func setupObserver() {
+        guard let service = clusterStateService else { return }
+
+        service.$latestSnapshot
+            .compactMap { $0 }
+            .sink { [weak self] snapshot in
+                self?.checkForCycles(snapshot: snapshot)
+            }
+            .store(in: &cancellables)
+    }
+
+    private func checkForCycles(snapshot: ClusterState) {
+        let cycles = snapshot.thunderboltBridgeCycles
+
+        // Only consider cycles with more than 2 nodes
+        guard let firstCycle = cycles.first, firstCycle.count > 2 else {
+            // No problematic cycles detected, reset state
+            if detectedCycle != nil {
+                detectedCycle = nil
+                previousCycleSignature = nil
+                hasPromptedForCurrentCycle = false
+            }
+            return
+        }
+
+        // Create a signature for this cycle to detect if it changed
+        let cycleSignature = firstCycle.sorted().joined(separator: ",")
+
+        // If this is a new/different cycle, reset the prompt state
+        if cycleSignature != previousCycleSignature {
+            previousCycleSignature = cycleSignature
+            hasPromptedForCurrentCycle = false
+        }
+
+        detectedCycle = firstCycle
+
+        // Only prompt once per cycle
+        if !hasPromptedForCurrentCycle {
+            showDisableBridgePrompt(nodeIds: firstCycle)
+        }
+    }
+
+    private func showDisableBridgePrompt(nodeIds: [String]) {
+        hasPromptedForCurrentCycle = true
+
+        // Get friendly names for the nodes if available
+        let nodeNames = nodeIds.map { nodeId -> String in
+            if let snapshot = clusterStateService?.latestSnapshot,
+                let profile = snapshot.nodeProfiles[nodeId],
+                let friendlyName = profile.friendlyName, !friendlyName.isEmpty
+            {
+                return friendlyName
+            }
+            return String(nodeId.prefix(8))  // Use first 8 chars of node ID as fallback
+        }
+        let machineNames = nodeNames.joined(separator: ", ")
+
+        let alert = NSAlert()
+        alert.messageText = "Thunderbolt Bridge Loop Detected"
+        alert.informativeText = """
+            A Thunderbolt Bridge loop has been detected between \(nodeNames.count) machines: \(machineNames).
+
+            This can cause network packet storms and connectivity issues. Would you like to disable Thunderbolt Bridge on this machine to break the loop?
+            """
+        alert.alertStyle = .warning
+        alert.addButton(withTitle: "Disable Bridge")
+        alert.addButton(withTitle: "Not Now")
+
+        let response = alert.runModal()
+
+        if response == .alertFirstButtonReturn {
+            Task {
+                await disableThunderboltBridge()
+            }
+        }
+    }
+
+    func disableThunderboltBridge() async {
+        Self.logger.info("Attempting to disable Thunderbolt Bridge via SCPreferences")
+        lastError = nil
+
+        do {
+            try await disableThunderboltBridgeWithSCPreferences()
+            Self.logger.info("Successfully disabled Thunderbolt Bridge")
+        } catch {
+            Self.logger.error(
+                "Failed to disable Thunderbolt Bridge: \(error.localizedDescription, privacy: .public)"
+            )
+            lastError = error.localizedDescription
+            showErrorAlert(message: error.localizedDescription)
+        }
+    }
+
+    private func disableThunderboltBridgeWithSCPreferences() async throws {
+        // 1. Create authorization reference
+        var authRef: AuthorizationRef?
+        var status = AuthorizationCreate(nil, nil, [], &authRef)
+        guard status == errAuthorizationSuccess, let authRef = authRef else {
+            throw ThunderboltBridgeError.authorizationFailed
+        }
+
+        defer { AuthorizationFree(authRef, [.destroyRights]) }
+
+        // 2. Request specific network configuration rights
+        let rightName = "system.services.systemconfiguration.network"
+        var item = AuthorizationItem(
+            name: rightName,
+            valueLength: 0,
+            value: nil,
+            flags: 0
+        )
+        var rights = AuthorizationRights(count: 1, items: &item)
+
+        status = AuthorizationCopyRights(
+            authRef,
+            &rights,
+            nil,
+            [.extendRights, .interactionAllowed],
+            nil
+        )
+        guard status == errAuthorizationSuccess else {
+            if status == errAuthorizationCanceled {
+                throw ThunderboltBridgeError.authorizationCanceled
+            }
+            throw ThunderboltBridgeError.authorizationDenied
+        }
+
+        // 3. Create SCPreferences with authorization
+        guard
+            let prefs = SCPreferencesCreateWithAuthorization(
+                kCFAllocatorDefault,
+                "EXO" as CFString,
+                nil,
+                authRef
+            )
+        else {
+            throw ThunderboltBridgeError.preferencesCreationFailed
+        }
+
+        // 4. Lock, modify, commit
+        guard SCPreferencesLock(prefs, true) else {
+            throw ThunderboltBridgeError.lockFailed
+        }
+
+        defer {
+            SCPreferencesUnlock(prefs)
+        }
+
+        // 5. Find the Thunderbolt Bridge service dynamically (don't assume the name)
+        guard let targetServiceName = ThunderboltBridgeDetector.findThunderboltBridgeServiceName()
+        else {
+            throw ThunderboltBridgeError.serviceNotFound
+        }
+
+        guard let allServices = SCNetworkServiceCopyAll(prefs) as? [SCNetworkService] else {
+            throw ThunderboltBridgeError.servicesNotFound
+        }
+
+        var found = false
+        for service in allServices {
+            if let name = SCNetworkServiceGetName(service) as String?,
+                name == targetServiceName
+            {
+                guard SCNetworkServiceSetEnabled(service, false) else {
+                    throw ThunderboltBridgeError.disableFailed
+                }
+                found = true
+                Self.logger.info(
+                    "Found and disabled Thunderbolt Bridge service: '\(targetServiceName)'")
+                break
+            }
+        }
+
+        if !found {
+            throw ThunderboltBridgeError.serviceNotFound
+        }
+
+        // 6. Commit and apply
+        guard SCPreferencesCommitChanges(prefs) else {
+            throw ThunderboltBridgeError.commitFailed
+        }
+
+        guard SCPreferencesApplyChanges(prefs) else {
+            throw ThunderboltBridgeError.applyFailed
+        }
+    }
+
+    private func showErrorAlert(message: String) {
+        let alert = NSAlert()
+        alert.messageText = "Failed to Disable Thunderbolt Bridge"
+        alert.informativeText = message
+        alert.alertStyle = .critical
+        alert.addButton(withTitle: "OK")
+        alert.runModal()
+    }
+}
+
+enum ThunderboltBridgeError: LocalizedError {
+    case authorizationFailed
+    case authorizationCanceled
+    case authorizationDenied
+    case preferencesCreationFailed
+    case lockFailed
+    case servicesNotFound
+    case serviceNotFound
+    case disableFailed
+    case commitFailed
+    case applyFailed
+
+    var errorDescription: String? {
+        switch self {
+        case .authorizationFailed:
+            return "Failed to create authorization"
+        case .authorizationCanceled:
+            return "Authorization was canceled by user"
+        case .authorizationDenied:
+            return "Authorization was denied"
+        case .preferencesCreationFailed:
+            return "Failed to access network preferences"
+        case .lockFailed:
+            return "Failed to lock network preferences for modification"
+        case .servicesNotFound:
+            return "Could not retrieve network services"
+        case .serviceNotFound:
+            return "Thunderbolt Bridge service not found"
+        case .disableFailed:
+            return "Failed to disable Thunderbolt Bridge service"
+        case .commitFailed:
+            return "Failed to save network configuration changes"
+        case .applyFailed:
+            return "Failed to apply network configuration changes"
+        }
+    }
+}

app/EXO/EXO/ViewModels/NodeViewModel.swift

@@ -86,7 +86,7 @@ struct TopologyViewModel {
 
 extension ClusterState {
     func topologyViewModel(localNodeId: String?) -> TopologyViewModel? {
-        let topologyNodeIds = Set(topology?.nodes.map(\.nodeId) ?? [])
+        let topologyNodeIds = Set(topology?.nodes ?? [])
         let allNodes = nodeViewModels().filter {
             topologyNodeIds.isEmpty || topologyNodeIds.contains($0.id)
         }
@@ -95,8 +95,8 @@ extension ClusterState {
         let nodesById = Dictionary(uniqueKeysWithValues: allNodes.map { ($0.id, $0) })
         var orderedNodes: [NodeViewModel] = []
         if let topologyNodes = topology?.nodes {
-            for topoNode in topologyNodes {
-                if let viewModel = nodesById[topoNode.nodeId] {
+            for nodeId in topologyNodes {
+                if let viewModel = nodesById[nodeId] {
                     orderedNodes.append(viewModel)
                 }
             }
@@ -116,7 +116,7 @@ extension ClusterState {
 
         let nodeIds = Set(orderedNodes.map(\.id))
         let edgesArray: [TopologyEdgeViewModel] =
-            topology?.connections?.compactMap { connection in
+            topology?.connections.compactMap { connection in
                 guard nodeIds.contains(connection.localNodeId),
                     nodeIds.contains(connection.sendBackNodeId)
                 else { return nil }

bench/eval_config.toml

@@ -0,0 +1,66 @@
+# exo-eval configuration file
+# See bench/exo_eval.py for usage
+
+[eval]
+# Eval framework type: "lm_eval" | "swe_bench" | "custom"
+type = "lm_eval"
+# Require HuggingFace token (default: true)
+# Set to false if using only public datasets
+require_hf_token = true
+
+# Instance/placement configuration
+# Controls how exo sets up the model instance before running evals
+[instance]
+# Placement strategy: "ring" | "jaccl" | "both"
+instance_meta = "jaccl"
+# Sharding strategy: "pipeline" | "tensor" | "both"
+sharding = "tensor"
+# Node constraints
+min_nodes = 2
+max_nodes = 2
+
+# lm_eval configuration (EleutherAI's lm-evaluation-harness)
+[lm_eval]
+# Tasks to run (list of task names)
+# NOTE: Chat completions API only supports generation-based tasks.
+# Loglikelihood tasks (mmlu, hellaswag, arc) require /v1/completions endpoint.
+#
+# Generation-based tasks (work with chat completions):
+#   - mmlu_pro, mmlu_generative, mmlu_flan_cot_fewshot, mmlu_flan_cot_zeroshot
+#   - gsm8k, gsm8k_cot, gsm8k_cot_zeroshot
+#   - truthfulqa (uses generate_until for some subtasks)
+#   - humaneval, mbpp (code generation)
+#
+# Run `lm_eval --tasks list` to see all available tasks
+tasks = ["mmlu_pro"]
+# Number of few-shot examples (5 is standard for mmlu_pro CoT)
+num_fewshot = 5
+# Batch size (use 1 for API models, "auto" doesn't work)
+batch_size = 1
+# Number of concurrent requests (set > 1 to enable parallelism)
+# Higher values enable better batching throughput
+num_concurrent = 64
+# Apply chat template for instruct/chat models (default: true)
+apply_chat_template = true
+# Use fewshot examples as conversation turns (better for chat models)
+fewshot_as_multiturn = true
+# Optional: limit samples per task (omit or comment out for no limit)
+# limit = 100
+# Output path for results
+output_path = "bench/eval_results"
+
+# SWE-bench configuration (placeholder)
+[swe_bench]
+# SWE-bench dataset
+dataset = "princeton-nlp/SWE-bench_Lite"
+# Maximum workers for parallel execution
+max_workers = 8
+# Path for prediction outputs
+predictions_path = "bench/predictions"
+
+# Custom evaluation script configuration
+[custom]
+# Path to custom evaluation script
+script = "path/to/eval_script.py"
+# Arguments to pass to the script
+args = ["--arg1", "value1"]

bench/exo_eval.py

@@ -0,0 +1,679 @@
+#!/usr/bin/env python3
+# pyright: reportAny=false, reportUnknownMemberType=false, reportUnknownVariableType=false, reportUnknownArgumentType=false
+"""
+exo-eval: Evaluation harness for exo inference system.
+
+Supports multiple evaluation frameworks via TOML configuration:
+- lm_eval: Language model evaluation using EleutherAI's lm-evaluation-harness
+- swe_bench: SWE-bench evaluation (placeholder for future implementation)
+- custom: Custom evaluation scripts
+
+Usage:
+    uv run python -m bench.exo_eval --config bench/eval_config.toml --model Llama-3.2-1b-Instruct-4bit
+    uv run python -m bench.exo_eval --config bench/eval_config.toml --model Llama-3.2-1b-Instruct-4bit --dry-run
+"""
+
+from __future__ import annotations
+
+import argparse
+import contextlib
+import json
+import os
+import subprocess
+import sys
+import time
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any, Literal
+
+# Add parent directory to path for direct script execution
+if __name__ == "__main__" and __package__ is None:
+    sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
+
+import tomlkit
+from huggingface_hub import get_token as get_hf_token
+from loguru import logger
+from tomlkit.exceptions import TOMLKitError
+
+from bench.exo_bench import (
+    ExoClient,
+    ExoHttpError,
+    instance_id_from_instance,
+    nodes_used_in_instance,
+    placement_filter,
+    resolve_model_short_id,
+    sharding_filter,
+    wait_for_instance_gone,
+    wait_for_instance_ready,
+)
+
+EvalType = Literal["lm_eval", "swe_bench", "custom"]
+
+
+def load_config(config_path: str) -> dict[str, Any]:
+    """Load and parse TOML configuration file."""
+    path = Path(config_path)
+    if not path.exists():
+        raise FileNotFoundError(f"Config file not found: {config_path}")
+
+    with open(path, encoding="utf-8") as f:
+        return dict(tomlkit.load(f))
+
+
+def get_eval_type(config: dict[str, Any]) -> EvalType:
+    """Extract evaluation type from config."""
+    eval_section = config.get("eval", {})
+    eval_type = eval_section.get("type", "lm_eval")
+    if eval_type not in ("lm_eval", "swe_bench", "custom"):
+        raise ValueError(f"Unknown eval type: {eval_type}")
+    return eval_type
+
+
+def check_hf_token(config: dict[str, Any]) -> bool:
+    """Check if HuggingFace token is available when required.
+
+    Returns True if token is available or not required, False otherwise.
+    """
+    eval_section = config.get("eval", {})
+    require_hf_token = eval_section.get("require_hf_token", True)
+
+    if not require_hf_token:
+        return True
+
+    token = get_hf_token()
+    if token is None:
+        logger.error(
+            "HuggingFace token not found. "
+            "Set HF_TOKEN environment variable or run 'huggingface-cli login'. "
+            "To disable this check, set require_hf_token = false in [eval] config."
+        )
+        return False
+
+    logger.info("HuggingFace token found")
+    return True
+
+
+def select_placement(
+    client: ExoClient,
+    full_model_id: str,
+    config: dict[str, Any],
+) -> dict[str, Any] | None:
+    """Select a placement based on config preferences."""
+    instance_config = config.get("instance", {})
+
+    # If explicit instance is provided, use it directly
+    if "instance" in instance_config:
+        return instance_config["instance"]
+
+    # Otherwise, select from previews based on preferences
+    instance_meta_pref = instance_config.get("instance_meta", "ring")
+    sharding_pref = instance_config.get("sharding", "pipeline")
+    max_nodes = instance_config.get("max_nodes", 4)
+    min_nodes = instance_config.get("min_nodes", 1)
+
+    previews_resp = client.request_json(
+        "GET", "/instance/previews", params={"model_id": full_model_id}
+    )
+    previews = previews_resp.get("previews") or []
+
+    selected: list[dict[str, Any]] = []
+    for p in previews:
+        if p.get("error") is not None:
+            continue
+        if not placement_filter(str(p.get("instance_meta", "")), instance_meta_pref):
+            continue
+        if not sharding_filter(str(p.get("sharding", "")), sharding_pref):
+            continue
+
+        instance = p.get("instance")
+        if not isinstance(instance, dict):
+            continue
+
+        n = nodes_used_in_instance(instance)
+        if min_nodes <= n <= max_nodes:
+            selected.append(p)
+
+    if not selected:
+        return None
+
+    # Sort by preference: exact match on sharding/meta, then by node count (descending)
+    def sort_key(p: dict[str, Any]) -> tuple[int, int, int]:
+        meta_match = (
+            1 if instance_meta_pref in str(p.get("instance_meta", "")).lower() else 0
+        )
+        sharding_match = 1 if sharding_pref in str(p.get("sharding", "")).lower() else 0
+        n_nodes = nodes_used_in_instance(p["instance"])
+        return (meta_match, sharding_match, n_nodes)
+
+    selected.sort(key=sort_key, reverse=True)
+    return selected[0]
+
+
+def setup_instance(
+    client: ExoClient,
+    full_model_id: str,
+    config: dict[str, Any],
+    dry_run: bool,
+) -> tuple[str | None, dict[str, Any] | None]:
+    """Create and wait for an instance to be ready. Returns (instance_id, preview)."""
+    preview = select_placement(client, full_model_id, config)
+
+    if preview is None:
+        logger.error("No valid placement found matching config preferences")
+        return None, None
+
+    instance_data = preview.get("instance")
+    instance: dict[str, Any] = (
+        instance_data if isinstance(instance_data, dict) else preview
+    )
+    instance_id = instance_id_from_instance(instance)
+
+    sharding = str(preview.get("sharding", "unknown"))
+    instance_meta = str(preview.get("instance_meta", "unknown"))
+    n_nodes = nodes_used_in_instance(instance)
+
+    logger.info(f"Selected placement: {sharding} / {instance_meta} / nodes={n_nodes}")
+    logger.info(f"Instance ID: {instance_id}")
+
+    if dry_run:
+        logger.info("[dry-run] Would create instance and wait for ready")
+        return instance_id, preview
+
+    # Create instance
+    client.request_json("POST", "/instance", body={"instance": instance})
+
+    try:
+        wait_for_instance_ready(client, instance_id)
+        logger.info("Instance is ready")
+        time.sleep(1)  # Brief pause after ready
+        return instance_id, preview
+    except (RuntimeError, TimeoutError) as e:
+        logger.error(f"Failed to initialize instance: {e}")
+        with contextlib.suppress(ExoHttpError):
+            client.request_json("DELETE", f"/instance/{instance_id}")
+        return None, None
+
+
+def teardown_instance(client: ExoClient, instance_id: str) -> None:
+    """Delete an instance and wait for it to be gone."""
+    try:
+        client.request_json("DELETE", f"/instance/{instance_id}")
+    except ExoHttpError as e:
+        if e.status != 404:
+            raise
+    except (ConnectionRefusedError, OSError):
+        logger.warning(
+            f"Could not connect to exo to delete instance {instance_id} (server may be down)"
+        )
+        return
+    try:
+        wait_for_instance_gone(client, instance_id)
+    except (ConnectionRefusedError, OSError, TimeoutError):
+        logger.warning("Could not verify instance deletion (server may be down)")
+        return
+    logger.info(f"Instance {instance_id} deleted")
+
+
+def build_lm_eval_args(
+    config: dict[str, Any],
+    base_url: str,
+    model: str,
+    output_path: str | None,
+    limit: int | None,
+    use_completions: bool,
+) -> list[str]:
+    """Build command-line arguments for lm_eval."""
+    lm_eval_config = config.get("lm_eval", {})
+
+    # Choose model type based on whether tasks need completions API
+    if use_completions:
+        model_type = "local-completions"
+        endpoint_url = f"{base_url}/v1/completions"
+    else:
+        model_type = "local-chat-completions"
+        endpoint_url = f"{base_url}/v1/chat/completions"
+
+    # Build model_args string with num_concurrent and timeout
+    model_args_parts = [f"model={model}", f"base_url={endpoint_url}"]
+    num_concurrent = lm_eval_config.get("num_concurrent")
+    if num_concurrent is not None and num_concurrent > 1:
+        model_args_parts.append(f"num_concurrent={num_concurrent}")
+    # Use a very long timeout (1 week) to handle large request queues
+    timeout = lm_eval_config.get("timeout", 604800)
+    model_args_parts.append(f"timeout={timeout}")
+    model_args = ",".join(model_args_parts)
+
+    args = [
+        sys.executable,
+        "-m",
+        "bench.lm_eval_patched",
+        "--model",
+        model_type,
+        "--model_args",
+        model_args,
+        "--verbosity",
+        "WARNING",
+    ]
+
+    # Tasks
+    tasks = lm_eval_config.get("tasks", ["mmlu"])
+    tasks_str = ",".join(tasks) if isinstance(tasks, list) else str(tasks)
+    args.extend(["--tasks", tasks_str])
+
+    # Few-shot
+    num_fewshot = lm_eval_config.get("num_fewshot")
+    if num_fewshot is not None:
+        args.extend(["--num_fewshot", str(num_fewshot)])
+
+    # Batch size (default to 1 for API models, "auto" doesn't work)
+    batch_size = lm_eval_config.get("batch_size", 1)
+    args.extend(["--batch_size", str(batch_size)])
+
+    # Apply chat template for instruct/chat models (default: true)
+    # Only applies to chat completions, but doesn't hurt to include
+    apply_chat_template = lm_eval_config.get("apply_chat_template", True)
+    if apply_chat_template and not use_completions:
+        args.append("--apply_chat_template")
+
+    # Fewshot as multiturn (optional, works with chat template)
+    fewshot_as_multiturn = lm_eval_config.get("fewshot_as_multiturn", False)
+    if fewshot_as_multiturn and not use_completions:
+        args.append("--fewshot_as_multiturn")
+
+    # Limit (command line overrides config)
+    effective_limit = limit if limit is not None else lm_eval_config.get("limit")
+    if effective_limit is not None:
+        args.extend(["--limit", str(effective_limit)])
+
+    # Output path
+    effective_output = output_path or lm_eval_config.get("output_path")
+    if effective_output:
+        args.extend(["--output_path", effective_output])
+        # Log model responses for post-hoc analysis when output is saved
+        args.append("--log_samples")
+
+    return args
+
+
+def run_lm_eval(
+    config: dict[str, Any],
+    host: str,
+    port: int,
+    model: str,
+    output_path: str | None,
+    limit: int | None,
+    dry_run: bool,
+) -> int:
+    """Run lm_eval evaluation."""
+    lm_eval_config = config.get("lm_eval", {})
+    tasks = lm_eval_config.get("tasks", ["mmlu"])
+    if isinstance(tasks, str):
+        tasks = [tasks]
+
+    exo_base_url = f"http://{host}:{port}"
+
+    # Build args - use native completions or chat completions endpoint directly
+    args = build_lm_eval_args(
+        config, exo_base_url, model, output_path, limit, use_completions=False
+    )
+    logger.info(f"lm_eval command: {' '.join(args)}")
+
+    if dry_run:
+        logger.info("[dry-run] Would execute the above command")
+        return 0
+
+    try:
+        result = subprocess.run(args, check=False)
+
+        # Print token usage summary from exo
+        try:
+            import httpx
+
+            usage_resp = httpx.get(f"{exo_base_url}/v1/usage", timeout=5)
+            if usage_resp.status_code == 200:
+                usage = usage_resp.json()
+                logger.info("--- Token Usage (Total) ---")
+                logger.info(f"  Requests:          {usage.get('total_requests', 0)}")
+                logger.info(
+                    f"  Prompt tokens:     {usage.get('total_prompt_tokens', 0)}"
+                )
+                logger.info(
+                    f"  Completion tokens: {usage.get('total_completion_tokens', 0)}"
+                )
+                logger.info(
+                    f"  Reasoning tokens:  {usage.get('total_reasoning_tokens', 0)}"
+                )
+                logger.info(f"  Total tokens:      {usage.get('total_tokens', 0)}")
+                by_model = usage.get("by_model", {})
+                if by_model:
+                    for model_name, counters in by_model.items():
+                        logger.info(f"--- Token Usage ({model_name}) ---")
+                        logger.info(
+                            f"  Requests:          {counters.get('requests', 0)}"
+                        )
+                        logger.info(
+                            f"  Prompt tokens:     {counters.get('prompt_tokens', 0)}"
+                        )
+                        logger.info(
+                            f"  Completion tokens: {counters.get('completion_tokens', 0)}"
+                        )
+                        logger.info(
+                            f"  Reasoning tokens:  {counters.get('reasoning_tokens', 0)}"
+                        )
+        except Exception:
+            pass  # Usage endpoint not available
+
+        return result.returncode
+    except FileNotFoundError:
+        logger.error("lm_eval not found. Install with: uv sync --extra eval")
+        return 1
+
+
+def run_swe_bench(
+    config: dict[str, Any],
+    host: str,
+    port: int,
+    model: str,
+    output_path: str | None,
+    dry_run: bool,
+) -> int:
+    """Run SWE-bench evaluation (placeholder)."""
+    swe_config = config.get("swe_bench", {})
+
+    dataset = swe_config.get("dataset", "princeton-nlp/SWE-bench_Lite")
+    max_workers = swe_config.get("max_workers", 8)
+    predictions_path = output_path or swe_config.get(
+        "predictions_path", "bench/predictions"
+    )
+
+    logger.info("SWE-bench evaluation configuration:")
+    logger.info(f"  Dataset: {dataset}")
+    logger.info(f"  Model: {model}")
+    logger.info(f"  API endpoint: http://{host}:{port}/v1")
+    logger.info(f"  Max workers: {max_workers}")
+    logger.info(f"  Predictions path: {predictions_path}")
+
+    if dry_run:
+        logger.info("[dry-run] SWE-bench evaluation would be executed")
+        return 0
+
+    logger.warning(
+        "SWE-bench integration is a placeholder. "
+        "Implement swebench inference and evaluation logic as needed."
+    )
+    return 0
+
+
+def run_custom_eval(
+    config: dict[str, Any],
+    host: str,
+    port: int,
+    model: str,
+    output_path: str | None,
+    dry_run: bool,
+) -> int:
+    """Run custom evaluation script."""
+    custom_config = config.get("custom", {})
+
+    script = custom_config.get("script")
+    if not script:
+        logger.error("No script specified in [custom] config section")
+        return 1
+
+    script_path = Path(script)
+    if not script_path.exists():
+        logger.error(f"Custom script not found: {script}")
+        return 1
+
+    script_args = custom_config.get("args", [])
+    if not isinstance(script_args, list):
+        script_args = [str(script_args)]
+
+    # Build environment with exo connection info
+    env = os.environ.copy()
+    env["EXO_HOST"] = host
+    env["EXO_PORT"] = str(port)
+    env["EXO_MODEL"] = model
+    if output_path:
+        env["EXO_OUTPUT_PATH"] = output_path
+
+    cmd = [sys.executable, str(script_path), *script_args]
+    logger.info(f"Custom eval command: {' '.join(cmd)}")
+
+    if dry_run:
+        logger.info("[dry-run] Would execute the above command")
+        return 0
+
+    result = subprocess.run(cmd, env=env, check=False)
+    return result.returncode
+
+
+def write_results_metadata(
+    output_path: str,
+    config: dict[str, Any],
+    host: str,
+    port: int,
+    model: str,
+    eval_type: EvalType,
+    return_code: int,
+    preview: dict[str, Any] | None,
+) -> None:
+    """Write evaluation metadata to a JSON file."""
+    metadata: dict[str, Any] = {
+        "timestamp": datetime.now(timezone.utc).isoformat(),
+        "eval_type": eval_type,
+        "model": model,
+        "api_endpoint": f"http://{host}:{port}/v1",
+        "config": config,
+        "return_code": return_code,
+    }
+
+    if preview:
+        metadata["placement"] = {
+            "sharding": preview.get("sharding"),
+            "instance_meta": preview.get("instance_meta"),
+            "instance_id": instance_id_from_instance(preview["instance"])
+            if "instance" in preview
+            else None,
+        }
+
+    output_dir = Path(output_path)
+    output_dir.mkdir(parents=True, exist_ok=True)
+    metadata_path = output_dir / "eval_metadata.json"
+
+    with open(metadata_path, "w", encoding="utf-8") as f:
+        json.dump(metadata, f, indent=2, ensure_ascii=False, default=str)
+
+    logger.info(f"Wrote evaluation metadata to: {metadata_path}")
+
+
+def main() -> int:
+    """Main entry point for exo-eval."""
+    ap = argparse.ArgumentParser(
+        prog="exo-eval",
+        description="Evaluation harness for exo inference system.",
+    )
+    ap.add_argument(
+        "--config",
+        required=True,
+        help="Path to TOML configuration file",
+    )
+    ap.add_argument(
+        "--host",
+        default=os.environ.get("EXO_HOST", "localhost"),
+        help="exo API host (default: localhost or EXO_HOST env var)",
+    )
+    ap.add_argument(
+        "--port",
+        type=int,
+        default=int(os.environ.get("EXO_PORT", "52415")),
+        help="exo API port (default: 52415 or EXO_PORT env var)",
+    )
+    ap.add_argument(
+        "--model",
+        required=True,
+        help="Model name/ID to evaluate",
+    )
+    ap.add_argument(
+        "--output",
+        default=None,
+        help="Output path for results (overrides config)",
+    )
+    ap.add_argument(
+        "--limit",
+        type=int,
+        default=None,
+        help="Limit samples per task (overrides config, lm_eval only)",
+    )
+    ap.add_argument(
+        "--timeout",
+        type=float,
+        default=604800.0,
+        help="HTTP timeout in seconds (default: 604800 = 1 week)",
+    )
+    ap.add_argument(
+        "--skip-instance-setup",
+        action="store_true",
+        help="Skip instance creation (assume instance already running)",
+    )
+    ap.add_argument(
+        "--pipeline",
+        type=int,
+        default=None,
+        metavar="N",
+        help="Use pipeline sharding with exactly N nodes (overrides config)",
+    )
+    ap.add_argument(
+        "--instance-meta",
+        choices=["ring", "jaccl", "both"],
+        default=None,
+        help="Instance meta preference (overrides config)",
+    )
+    ap.add_argument(
+        "--dry-run",
+        action="store_true",
+        help="Print commands without executing",
+    )
+    args = ap.parse_args()
+
+    logger.info(f"exo-eval starting with config: {args.config}")
+
+    try:
+        config = load_config(args.config)
+    except FileNotFoundError as e:
+        logger.error(str(e))
+        return 1
+    except TOMLKitError as e:
+        logger.error(f"Failed to parse config: {e}")
+        return 1
+
+    eval_type = get_eval_type(config)
+    logger.info(f"Evaluation type: {eval_type}")
+    logger.info(f"Model: {args.model}")
+    logger.info(f"API endpoint: http://{args.host}:{args.port}/v1")
+
+    # Apply CLI overrides to instance config
+    if args.pipeline is not None or args.instance_meta is not None:
+        instance_config = config.setdefault("instance", {})
+        if args.pipeline is not None:
+            instance_config["sharding"] = "pipeline"
+            instance_config["min_nodes"] = args.pipeline
+            instance_config["max_nodes"] = args.pipeline
+            logger.info(f"CLI override: pipeline={args.pipeline} nodes")
+            # Limit concurrency for pipeline to avoid GPU timeouts
+            if args.pipeline >= 2:
+                lm_eval_config = config.setdefault("lm_eval", {})
+                lm_eval_config["num_concurrent"] = 4
+                logger.info("CLI override: num_concurrent=4 (pipeline>=2)")
+        if args.instance_meta is not None:
+            instance_config["instance_meta"] = args.instance_meta
+            logger.info(f"CLI override: instance_meta={args.instance_meta}")
+
+    # Check HuggingFace token if required
+    if not check_hf_token(config):
+        return 1
+
+    # Setup instance and resolve model
+    instance_id: str | None = None
+    preview: dict[str, Any] | None = None
+    client: ExoClient | None = None
+
+    if args.skip_instance_setup:
+        # Use model name as-is when skipping instance setup
+        full_model_id = args.model
+        logger.info(f"Using model: {full_model_id} (instance setup skipped)")
+    else:
+        client = ExoClient(args.host, args.port, timeout_s=args.timeout)
+
+        # Resolve model
+        try:
+            short_id, full_model_id = resolve_model_short_id(client, args.model)
+            logger.info(f"Resolved model: {short_id} -> {full_model_id}")
+        except Exception as e:
+            logger.error(f"Failed to resolve model: {e}")
+            return 1
+
+        instance_id, preview = setup_instance(
+            client, full_model_id, config, args.dry_run
+        )
+        if instance_id is None and not args.dry_run:
+            return 1
+
+    try:
+        # Run evaluation
+        if eval_type == "lm_eval":
+            return_code = run_lm_eval(
+                config,
+                args.host,
+                args.port,
+                full_model_id,
+                args.output,
+                args.limit,
+                args.dry_run,
+            )
+        elif eval_type == "swe_bench":
+            return_code = run_swe_bench(
+                config,
+                args.host,
+                args.port,
+                full_model_id,
+                args.output,
+                args.dry_run,
+            )
+        elif eval_type == "custom":
+            return_code = run_custom_eval(
+                config,
+                args.host,
+                args.port,
+                full_model_id,
+                args.output,
+                args.dry_run,
+            )
+        else:
+            logger.error(f"Unknown eval type: {eval_type}")
+            return 1
+
+        # Write metadata if output path specified and not dry-run
+        output_path = args.output or config.get(eval_type, {}).get("output_path")
+        if output_path and not args.dry_run:
+            write_results_metadata(
+                output_path,
+                config,
+                args.host,
+                args.port,
+                full_model_id,
+                eval_type,
+                return_code,
+                preview,
+            )
+
+        return return_code
+
+    finally:
+        # Teardown instance
+        if instance_id and client and not args.skip_instance_setup and not args.dry_run:
+            teardown_instance(client, instance_id)
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

bench/lm_eval_patched.py

@@ -0,0 +1,145 @@
+"""Patched lm_eval runner that fixes bugs in the upstream library.
+
+Fixes:
+- UnboundLocalError on `outputs` in TemplateAPI.amodel_call when API returns error
+- Prevents eval crash on transient API failures (returns None instead of raising)
+- Compatibility with transformers 5.x (missing AutoModelForVision2Seq)
+- sock_read timeout causing connection drops with large request queues
+
+Usage: python -m bench.lm_eval_patched [lm_eval args...]
+"""
+
+# ruff: noqa: I001, E402
+# pyright: reportMissingTypeStubs=false, reportUnknownVariableType=false
+# pyright: reportUnknownMemberType=false, reportAny=false, reportUnknownArgumentType=false
+# pyright: reportPrivateUsage=false, reportUnknownLambdaType=false
+
+# MUST patch transformers BEFORE any lm_eval imports
+# AutoModelForVision2Seq/AutoModelForImageTextToText were removed in transformers 5.0
+# Patch the lazy module's __getattr__ to return stubs for missing classes
+from transformers.utils import import_utils
+
+_original_getattr = import_utils._LazyModule.__getattr__
+
+
+def _patched_getattr(self: object, name: str) -> object:
+    if name in ("AutoModelForVision2Seq", "AutoModelForImageTextToText"):
+        return type(name, (), {})  # Return a stub class
+    return _original_getattr(self, name)  # type: ignore
+
+
+import_utils._LazyModule.__getattr__ = _patched_getattr
+
+import functools
+from typing import Any
+
+
+def _patch_amodel_call() -> None:
+    """Monkey-patch TemplateAPI.amodel_call to handle the unbound `outputs` variable bug."""
+    from lm_eval.models.api_models import TemplateAPI
+
+    original: Any = TemplateAPI.amodel_call
+
+    @functools.wraps(original)
+    async def patched_amodel_call(self: Any, *args: Any, **kwargs: Any) -> Any:
+        try:
+            return await original(self, *args, **kwargs)
+        except (UnboundLocalError, Exception):
+            # Return one empty-string result per request in the batch so the
+            # reorderer doesn't assert on missing coverage.
+            messages = kwargs.get("messages") or (args[2] if len(args) > 2 else [])
+            return [""] * max(len(messages), 1)
+
+    TemplateAPI.amodel_call = patched_amodel_call
+
+
+def _patch_client_timeout() -> None:
+    """Patch TemplateAPI.get_batched_requests to disable sock_read timeout.
+
+    By default, aiohttp's ClientTimeout can have a sock_read timeout that causes
+    connections to drop if no data is received for a while. With large request
+    queues, requests may wait a long time before processing starts, causing
+    spurious connection drops and retries that pile up requests.
+    """
+    from aiohttp import ClientSession, ClientTimeout, TCPConnector
+
+    from lm_eval.models.api_models import TemplateAPI
+
+    original_get_batched: Any = TemplateAPI.get_batched_requests
+
+    @functools.wraps(original_get_batched)
+    async def patched_get_batched_requests(self: Any, *args: Any, **kwargs: Any) -> Any:
+        # Override the timeout to explicitly disable sock_read timeout
+        # This prevents connection drops when requests are queued for a long time
+        original_timeout = getattr(self, "timeout", 604800)
+        conn = TCPConnector(limit=self._concurrent, ssl=self.verify_certificate)
+        timeout = ClientTimeout(
+            total=original_timeout, sock_read=None, sock_connect=None
+        )
+
+        async with ClientSession(connector=conn, timeout=timeout) as session:
+            # Call the internal async logic with our session
+            return await _run_batched_requests_with_session(
+                self, session, *args, **kwargs
+            )
+
+    async def _run_batched_requests_with_session(
+        self: Any,
+        session: ClientSession,
+        requests: Any,
+        cache_keys: Any = None,
+        ctxlens: Any = None,
+        **kwargs: Any,
+    ) -> Any:
+        import asyncio
+        import copy
+        import logging
+
+        from tqdm.asyncio import tqdm_asyncio
+        from tenacity import retry, stop_after_attempt, wait_exponential
+        from lm_eval.models.utils import chunks
+
+        eval_logger = logging.getLogger("lm_eval.models.api_models")
+        ctxlens = ctxlens if ctxlens else [None] * len(requests)
+        sem = asyncio.Semaphore(self._concurrent)
+
+        retry_: Any = retry(
+            stop=stop_after_attempt(self.max_retries),
+            wait=wait_exponential(multiplier=0.5, min=1, max=10),
+            reraise=True,
+            before_sleep=lambda retry_state: eval_logger.info(
+                f"Retry attempt {retry_state.attempt_number}"
+            ),
+        )(self.amodel_call)
+
+        tasks = [
+            asyncio.create_task(
+                retry_(
+                    session=session,
+                    sem=sem,
+                    messages=message,
+                    cache_keys=cache_key,
+                    ctxlens=ctxlen,
+                    gen_kwargs=copy.deepcopy(kwargs.get("gen_kwargs")),
+                    **{k: v for k, v in kwargs.items() if k != "gen_kwargs"},
+                )
+            )
+            for message, cache_key, ctxlen in zip(
+                chunks(requests, n=self._batch_size),
+                chunks(cache_keys, n=self._batch_size),
+                chunks(ctxlens, n=self._batch_size),
+                strict=True,
+            )
+        ]
+
+        return await tqdm_asyncio.gather(*tasks, desc="Requesting API")
+
+    TemplateAPI.get_batched_requests = patched_get_batched_requests
+
+
+if __name__ == "__main__":
+    _patch_amodel_call()
+    _patch_client_timeout()
+    from lm_eval.__main__ import cli_evaluate
+
+    cli_evaluate()

bench/stats_dashboard.html

@@ -0,0 +1,290 @@
+<!DOCTYPE html>
+<html lang="en">
+<head>
+<meta charset="UTF-8">
+<meta name="viewport" content="width=device-width, initial-scale=1.0">
+<title>exo Usage Stats</title>
+<style>
+  * { margin: 0; padding: 0; box-sizing: border-box; }
+  body {
+    font-family: -apple-system, BlinkMacSystemFont, 'SF Mono', 'Menlo', monospace;
+    background: #1a1a2e;
+    color: #e0e0e0;
+    padding: 24px;
+    min-height: 100vh;
+  }
+  .header {
+    display: flex;
+    justify-content: space-between;
+    align-items: center;
+    margin-bottom: 24px;
+    padding-bottom: 16px;
+    border-bottom: 1px solid #333;
+  }
+  .header h1 {
+    font-size: 20px;
+    font-weight: 600;
+    color: #fff;
+  }
+  .status {
+    display: flex;
+    align-items: center;
+    gap: 8px;
+    font-size: 13px;
+    color: #888;
+  }
+  .status-dot {
+    width: 8px;
+    height: 8px;
+    border-radius: 50%;
+    background: #666;
+  }
+  .status-dot.connected { background: #4caf50; }
+  .status-dot.error { background: #f44336; }
+  .config {
+    margin-bottom: 24px;
+    display: flex;
+    align-items: center;
+    gap: 8px;
+  }
+  .config label {
+    font-size: 12px;
+    color: #888;
+  }
+  .config input {
+    background: #252540;
+    border: 1px solid #444;
+    border-radius: 4px;
+    color: #e0e0e0;
+    padding: 4px 8px;
+    font-size: 13px;
+    font-family: inherit;
+    width: 280px;
+  }
+  .section {
+    background: #252540;
+    border-radius: 8px;
+    padding: 20px;
+    margin-bottom: 16px;
+  }
+  .section h2 {
+    font-size: 14px;
+    font-weight: 600;
+    color: #aaa;
+    text-transform: uppercase;
+    letter-spacing: 0.5px;
+    margin-bottom: 16px;
+  }
+  .stat-grid {
+    display: grid;
+    grid-template-columns: repeat(auto-fit, minmax(200px, 1fr));
+    gap: 16px;
+  }
+  .stat-card {
+    background: #1a1a2e;
+    border-radius: 6px;
+    padding: 16px;
+  }
+  .stat-label {
+    font-size: 11px;
+    color: #888;
+    text-transform: uppercase;
+    letter-spacing: 0.5px;
+    margin-bottom: 4px;
+  }
+  .stat-value {
+    font-size: 28px;
+    font-weight: 700;
+    color: #fff;
+  }
+  .stat-rate {
+    font-size: 12px;
+    color: #4caf50;
+    margin-top: 4px;
+  }
+  table {
+    width: 100%;
+    border-collapse: collapse;
+    font-size: 13px;
+  }
+  th {
+    text-align: left;
+    padding: 8px 12px;
+    color: #888;
+    font-weight: 500;
+    border-bottom: 1px solid #333;
+    font-size: 11px;
+    text-transform: uppercase;
+    letter-spacing: 0.5px;
+  }
+  td {
+    padding: 8px 12px;
+    border-bottom: 1px solid #2a2a45;
+  }
+  td.num {
+    text-align: right;
+    font-variant-numeric: tabular-nums;
+  }
+  .model-name {
+    color: #7c9eff;
+    max-width: 300px;
+    overflow: hidden;
+    text-overflow: ellipsis;
+    white-space: nowrap;
+  }
+  .empty-state {
+    color: #666;
+    font-style: italic;
+    padding: 16px 0;
+  }
+</style>
+</head>
+<body>
+  <div class="header">
+    <h1>exo Usage Stats</h1>
+    <div class="status">
+      <div class="status-dot" id="statusDot"></div>
+      <span id="statusText">connecting...</span>
+    </div>
+  </div>
+
+  <div class="config">
+    <label for="baseUrl">Base URL:</label>
+    <input type="text" id="baseUrl" value="http://mac8-1:52415">
+  </div>
+
+  <div class="section">
+    <h2>Totals</h2>
+    <div class="stat-grid">
+      <div class="stat-card">
+        <div class="stat-label">Requests</div>
+        <div class="stat-value" id="totalRequests">0</div>
+      </div>
+      <div class="stat-card">
+        <div class="stat-label">Prompt Tokens</div>
+        <div class="stat-value" id="totalPrompt">0</div>
+        <div class="stat-rate" id="promptRate"></div>
+      </div>
+      <div class="stat-card">
+        <div class="stat-label">Completion Tokens</div>
+        <div class="stat-value" id="totalCompletion">0</div>
+        <div class="stat-rate" id="completionRate"></div>
+      </div>
+      <div class="stat-card">
+        <div class="stat-label">Reasoning Tokens</div>
+        <div class="stat-value" id="totalReasoning">0</div>
+      </div>
+      <div class="stat-card">
+        <div class="stat-label">Total Tokens</div>
+        <div class="stat-value" id="totalTokens">0</div>
+        <div class="stat-rate" id="totalRate"></div>
+      </div>
+    </div>
+  </div>
+
+  <div class="section">
+    <h2>Per-Model Breakdown</h2>
+    <div id="modelTable">
+      <div class="empty-state">No data yet</div>
+    </div>
+  </div>
+
+<script>
+
+  function fmt(n) {
+    return n.toLocaleString();
+  }
+
+  // Track first non-zero timestamp for overall average rate
+  let firstSeenTime = null;
+  let firstSeenTokens = { prompt: 0, completion: 0, total: 0 };
+
+  function setRate(id, currentTokens, tokenType) {
+    const el = document.getElementById(id);
+    if (firstSeenTime === null || currentTokens <= firstSeenTokens[tokenType]) {
+      el.textContent = '';
+      return;
+    }
+    const elapsed = (performance.now() / 1000) - firstSeenTime;
+    if (elapsed <= 0) { el.textContent = ''; return; }
+    const delta = currentTokens - firstSeenTokens[tokenType];
+    const avg = delta / elapsed;
+    el.textContent = fmt(Math.round(avg)) + ' tok/s avg';
+  }
+
+  function renderModelTable(byModel) {
+    const container = document.getElementById('modelTable');
+    const models = Object.entries(byModel);
+    if (models.length === 0) {
+      container.innerHTML = '<div class="empty-state">No data yet</div>';
+      return;
+    }
+    let html = '<table><thead><tr>';
+    html += '<th>Model</th><th style="text-align:right">Requests</th>';
+    html += '<th style="text-align:right">Prompt</th>';
+    html += '<th style="text-align:right">Completion</th>';
+    html += '<th style="text-align:right">Reasoning</th>';
+    html += '<th style="text-align:right">Total</th>';
+    html += '</tr></thead><tbody>';
+    for (const [name, counters] of models) {
+      const total = (counters.prompt_tokens || 0) + (counters.completion_tokens || 0);
+      html += '<tr>';
+      html += `<td class="model-name" title="${name}">${name}</td>`;
+      html += `<td class="num">${fmt(counters.requests || 0)}</td>`;
+      html += `<td class="num">${fmt(counters.prompt_tokens || 0)}</td>`;
+      html += `<td class="num">${fmt(counters.completion_tokens || 0)}</td>`;
+      html += `<td class="num">${fmt(counters.reasoning_tokens || 0)}</td>`;
+      html += `<td class="num">${fmt(total)}</td>`;
+      html += '</tr>';
+    }
+    html += '</tbody></table>';
+    container.innerHTML = html;
+  }
+
+  async function poll() {
+    const baseUrl = document.getElementById('baseUrl').value.replace(/\/+$/, '');
+    const dot = document.getElementById('statusDot');
+    const text = document.getElementById('statusText');
+
+    try {
+      const resp = await fetch(baseUrl + '/v1/usage');
+      if (!resp.ok) throw new Error(`HTTP ${resp.status}`);
+      const data = await resp.json();
+
+      dot.className = 'status-dot connected';
+      text.textContent = 'connected';
+
+
+      document.getElementById('totalRequests').textContent = fmt(data.total_requests || 0);
+      document.getElementById('totalPrompt').textContent = fmt(data.total_prompt_tokens || 0);
+      document.getElementById('totalCompletion').textContent = fmt(data.total_completion_tokens || 0);
+      document.getElementById('totalReasoning').textContent = fmt(data.total_reasoning_tokens || 0);
+      document.getElementById('totalTokens').textContent = fmt(data.total_tokens || 0);
+
+      // Record first non-zero reading as baseline
+      if (firstSeenTime === null && (data.total_tokens || 0) > 0) {
+        firstSeenTime = performance.now() / 1000;
+        firstSeenTokens = {
+          prompt: data.total_prompt_tokens || 0,
+          completion: data.total_completion_tokens || 0,
+          total: data.total_tokens || 0,
+        };
+      }
+
+      setRate('promptRate', data.total_prompt_tokens || 0, 'prompt');
+      setRate('completionRate', data.total_completion_tokens || 0, 'completion');
+      setRate('totalRate', data.total_tokens || 0, 'total');
+
+      renderModelTable(data.by_model || {});
+
+    } catch (e) {
+      dot.className = 'status-dot error';
+      text.textContent = e.message || 'error';
+    }
+  }
+
+  poll();
+  setInterval(poll, 1000);
+</script>
+</body>
+</html>

dashboard/package-lock.json

@@ -865,6 +865,7 @@
 			"integrity": "sha512-oH8tXw7EZnie8FdOWYrF7Yn4IKrqTFHhXvl8YxXxbKwTMcD/5NNCryUSEXRk2ZR4ojnub0P8rNrsVGHXWqIDtA==",
 			"dev": true,
 			"license": "MIT",
+			"peer": true,
 			"dependencies": {
 				"@standard-schema/spec": "^1.0.0",
 				"@sveltejs/acorn-typescript": "^1.0.5",
@@ -904,6 +905,7 @@
 			"integrity": "sha512-Y1Cs7hhTc+a5E9Va/xwKlAJoariQyHY+5zBgCZg4PFWNYQ1nMN9sjK1zhw1gK69DuqVP++sht/1GZg1aRwmAXQ==",
 			"dev": true,
 			"license": "MIT",
+			"peer": true,
 			"dependencies": {
 				"@sveltejs/vite-plugin-svelte-inspector": "^4.0.1",
 				"debug": "^4.4.1",
@@ -1520,6 +1522,7 @@
 			"integrity": "sha512-LCCV0HdSZZZb34qifBsyWlUmok6W7ouER+oQIGBScS8EsZsQbrtFTUrDX4hOl+CS6p7cnNC4td+qrSVGSCTUfQ==",
 			"dev": true,
 			"license": "MIT",
+			"peer": true,
 			"dependencies": {
 				"undici-types": "~6.21.0"
 			}
@@ -1529,6 +1532,7 @@
 			"resolved": "https://registry.npmjs.org/acorn/-/acorn-8.15.0.tgz",
 			"integrity": "sha512-NZyJarBfL7nWwIq+FDL6Zp/yHEhePMNnnJ0y3qfieCrmNvYct8uvtiV41UvlSe6apAfk0fY1FbWx+NwfmpvtTg==",
 			"license": "MIT",
+			"peer": true,
 			"bin": {
 				"acorn": "bin/acorn"
 			},
@@ -1941,6 +1945,7 @@
 			"integrity": "sha512-fmTRWbNMmsmWq6xJV8D19U/gw/bwrHfNXxrIN+HfZgnzqTHp9jOmKMhsTUjXOJnZOdZY9Q28y4yebKzqDKlxlQ==",
 			"dev": true,
 			"license": "ISC",
+			"peer": true,
 			"engines": {
 				"node": ">=12"
 			}
@@ -2648,6 +2653,7 @@
 			"integrity": "sha512-5gTmgEY/sqK6gFXLIsQNH19lWb4ebPDLA4SdLP7dsWkIXHWlG66oPuVvXSGFPppYZz8ZDZq0dYYrbHfBCVUb1Q==",
 			"dev": true,
 			"license": "MIT",
+			"peer": true,
 			"engines": {
 				"node": ">=12"
 			},
@@ -2690,6 +2696,7 @@
 			"integrity": "sha512-UOnG6LftzbdaHZcKoPFtOcCKztrQ57WkHDeRD9t/PTQtmT0NHSeWWepj6pS0z/N7+08BHFDQVUrfmfMRcZwbMg==",
 			"dev": true,
 			"license": "MIT",
+			"peer": true,
 			"bin": {
 				"prettier": "bin/prettier.cjs"
 			},
@@ -2862,6 +2869,7 @@
 			"resolved": "https://registry.npmjs.org/svelte/-/svelte-5.45.3.tgz",
 			"integrity": "sha512-ngKXNhNvwPzF43QqEhDOue7TQTrG09em1sd4HBxVF0Wr2gopAmdEWan+rgbdgK4fhBtSOTJO8bYU4chUG7VXZQ==",
 			"license": "MIT",
+			"peer": true,
 			"dependencies": {
 				"@jridgewell/remapping": "^2.3.4",
 				"@jridgewell/sourcemap-codec": "^1.5.0",
@@ -3006,6 +3014,7 @@
 			"integrity": "sha512-jl1vZzPDinLr9eUt3J/t7V6FgNEw9QjvBPdysz9KfQDD41fQrC2Y4vKQdiaUpFT4bXlb1RHhLpp8wtm6M5TgSw==",
 			"dev": true,
 			"license": "Apache-2.0",
+			"peer": true,
 			"bin": {
 				"tsc": "bin/tsc",
 				"tsserver": "bin/tsserver"
@@ -3027,6 +3036,7 @@
 			"integrity": "sha512-+Oxm7q9hDoLMyJOYfUYBuHQo+dkAloi33apOPP56pzj+vsdJDzr+j1NISE5pyaAuKL4A3UD34qd0lx5+kfKp2g==",
 			"dev": true,
 			"license": "MIT",
+			"peer": true,
 			"dependencies": {
 				"esbuild": "^0.25.0",
 				"fdir": "^6.4.4",

dashboard/parts.nix

@@ -3,6 +3,61 @@
   perSystem =
     { pkgs, lib, ... }:
     let
+      # Filter source to ONLY include package.json and package-lock.json
+      # This ensures prettier-svelte only rebuilds when lockfiles change
+      dashboardLockfileSrc = lib.cleanSourceWith {
+        src = inputs.self;
+        filter =
+          path: type:
+          let
+            baseName = builtins.baseNameOf path;
+            isDashboardDir = baseName == "dashboard" && type == "directory";
+            isPackageFile =
+              (lib.hasInfix "/dashboard/" path || lib.hasSuffix "/dashboard" (builtins.dirOf path))
+              && (baseName == "package.json" || baseName == "package-lock.json");
+          in
+          isDashboardDir || isPackageFile;
+      };
+
+      # Stub source with lockfiles and minimal files for build to succeed
+      # This allows prettier-svelte to avoid rebuilding when dashboard source changes
+      dashboardStubSrc = pkgs.runCommand "dashboard-stub-src" { } ''
+        mkdir -p $out
+        cp ${dashboardLockfileSrc}/dashboard/package.json $out/
+        cp ${dashboardLockfileSrc}/dashboard/package-lock.json $out/
+        # Minimal files so vite build succeeds (produces empty output)
+        echo '<!DOCTYPE html><html><head></head><body></body></html>' > $out/index.html
+        mkdir -p $out/src
+        touch $out/src/app.html
+      '';
+
+      # Deps-only build using stub source (for prettier-svelte)
+      # Only rebuilds when package.json or package-lock.json change
+      dashboardDeps = inputs.dream2nix.lib.evalModules {
+        packageSets.nixpkgs = pkgs;
+        modules = [
+          ./dashboard.nix
+          {
+            paths.projectRoot = inputs.self;
+            paths.projectRootFile = "flake.nix";
+            paths.package = inputs.self + "/dashboard";
+          }
+          {
+            deps.dashboardSrc = lib.mkForce dashboardStubSrc;
+          }
+          # Override build phases to skip the actual build - just need node_modules
+          {
+            mkDerivation = {
+              buildPhase = lib.mkForce "true";
+              installPhase = lib.mkForce ''
+                runHook preInstall
+                runHook postInstall
+              '';
+            };
+          }
+        ];
+      };
+
       # Filter source to only include dashboard directory
       dashboardSrc = lib.cleanSourceWith {
         src = inputs.self;
@@ -42,11 +97,12 @@
       '';
 
       # Prettier with svelte plugin for treefmt
+      # Uses dashboardDeps instead of dashboardFull to avoid rebuilding on source changes
       packages.prettier-svelte = pkgs.writeShellScriptBin "prettier-svelte" ''
-        export NODE_PATH="${dashboardFull}/lib/node_modules/exo-dashboard/node_modules"
+        export NODE_PATH="${dashboardDeps}/lib/node_modules/exo-dashboard/node_modules"
         exec ${pkgs.nodejs}/bin/node \
-          ${dashboardFull}/lib/node_modules/exo-dashboard/node_modules/prettier/bin/prettier.cjs \
-          --plugin "${dashboardFull}/lib/node_modules/exo-dashboard/node_modules/prettier-plugin-svelte/plugin.js" \
+          ${dashboardDeps}/lib/node_modules/exo-dashboard/node_modules/prettier/bin/prettier.cjs \
+          --plugin "${dashboardDeps}/lib/node_modules/exo-dashboard/node_modules/prettier-plugin-svelte/plugin.js" \
           "$@"
       '';
     };

dashboard/src/lib/components/ChatForm.svelte

@@ -89,7 +89,10 @@
 
   const isImageModel = $derived(() => {
     if (!currentModel) return false;
-    return modelSupportsTextToImage(currentModel);
+    return (
+      modelSupportsTextToImage(currentModel) ||
+      modelSupportsImageEditing(currentModel)
+    );
   });
 
   const isEditOnlyWithoutImage = $derived(
@@ -646,6 +649,23 @@
             </svg>
             <span>EDIT</span>
           </span>
+        {:else if isEditOnlyWithoutImage}
+          <span class="inline-flex items-center gap-1.5">
+            <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="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"
+              />
+            </svg>
+            <span>EDIT</span>
+          </span>
         {:else if isImageModel()}
           <span class="inline-flex items-center gap-1.5">
             <svg

dashboard/src/lib/components/ImageParamsPanel.svelte

@@ -110,6 +110,36 @@
     setImageGenerationParams({ negativePrompt: value || null });
   }
 
+  function handleNumImagesChange(event: Event) {
+    const input = event.target as HTMLInputElement;
+    const value = input.value.trim();
+    if (value === "") {
+      setImageGenerationParams({ numImages: 1 });
+    } else {
+      const num = parseInt(value, 10);
+      if (!isNaN(num) && num >= 1) {
+        setImageGenerationParams({ numImages: num });
+      }
+    }
+  }
+
+  function handleStreamChange(enabled: boolean) {
+    setImageGenerationParams({ stream: enabled });
+  }
+
+  function handlePartialImagesChange(event: Event) {
+    const input = event.target as HTMLInputElement;
+    const value = input.value.trim();
+    if (value === "") {
+      setImageGenerationParams({ partialImages: 0 });
+    } else {
+      const num = parseInt(value, 10);
+      if (!isNaN(num) && num >= 0) {
+        setImageGenerationParams({ partialImages: num });
+      }
+    }
+  }
+
   function clearSteps() {
     setImageGenerationParams({ numInferenceSteps: null });
   }
@@ -134,90 +164,92 @@
 <div class="border-b border-exo-medium-gray/30 px-3 py-2">
   <!-- Basic params row -->
   <div class="flex items-center gap-3 flex-wrap">
-    <!-- Size -->
-    <div class="flex items-center gap-1.5">
-      <span class="text-xs text-exo-light-gray uppercase tracking-wider"
-        >SIZE:</span
-      >
-      <div class="relative">
-        <button
-          bind:this={sizeButtonRef}
-          type="button"
-          onclick={() => (isSizeDropdownOpen = !isSizeDropdownOpen)}
-          class="bg-exo-medium-gray/50 border border-exo-yellow/30 rounded pl-2 pr-6 py-1 text-xs font-mono text-exo-yellow cursor-pointer transition-all duration-200 hover:border-exo-yellow/50 focus:outline-none focus:border-exo-yellow/70 {isSizeDropdownOpen
-            ? 'border-exo-yellow/70'
-            : ''}"
+    <!-- Size (hidden in edit mode - output size comes from input image) -->
+    {#if !isEditMode}
+      <div class="flex items-center gap-1.5">
+        <span class="text-xs text-exo-light-gray uppercase tracking-wider"
+          >SIZE:</span
         >
-          {params.size}
-        </button>
-        <div
-          class="absolute right-1.5 top-1/2 -translate-y-1/2 pointer-events-none transition-transform duration-200 {isSizeDropdownOpen
-            ? 'rotate-180'
-            : ''}"
-        >
-          <svg
-            class="w-3 h-3 text-exo-yellow/60"
-            fill="none"
-            viewBox="0 0 24 24"
-            stroke="currentColor"
+        <div class="relative">
+          <button
+            bind:this={sizeButtonRef}
+            type="button"
+            onclick={() => (isSizeDropdownOpen = !isSizeDropdownOpen)}
+            class="bg-exo-medium-gray/50 border border-exo-yellow/30 rounded pl-2 pr-6 py-1 text-xs font-mono text-exo-yellow cursor-pointer transition-all duration-200 hover:border-exo-yellow/50 focus:outline-none focus:border-exo-yellow/70 {isSizeDropdownOpen
+              ? 'border-exo-yellow/70'
+              : ''}"
           >
-            <path
-              stroke-linecap="round"
-              stroke-linejoin="round"
-              stroke-width="2"
-              d="M19 9l-7 7-7-7"
-            />
-          </svg>
-        </div>
-      </div>
-
-      {#if isSizeDropdownOpen}
-        <!-- Backdrop to close dropdown -->
-        <button
-          type="button"
-          class="fixed inset-0 z-[9998] cursor-default"
-          onclick={() => (isSizeDropdownOpen = 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 min-w-max"
-          style="bottom: calc(100vh - {sizeDropdownPosition()
-            .top}px + 4px); left: {sizeDropdownPosition().left}px;"
-        >
-          <div class="py-1">
-            {#each sizeOptions as size}
-              <button
-                type="button"
-                onclick={() => selectSize(size)}
-                class="w-full px-3 py-1.5 text-left text-xs font-mono tracking-wide transition-colors duration-100 flex items-center gap-2 {params.size ===
-                size
-                  ? 'bg-transparent text-exo-yellow'
-                  : 'text-exo-light-gray hover:text-exo-yellow'}"
-              >
-                {#if params.size === size}
-                  <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}
-                <span>{size}</span>
-              </button>
-            {/each}
+            {params.size}
+          </button>
+          <div
+            class="absolute right-1.5 top-1/2 -translate-y-1/2 pointer-events-none transition-transform duration-200 {isSizeDropdownOpen
+              ? 'rotate-180'
+              : ''}"
+          >
+            <svg
+              class="w-3 h-3 text-exo-yellow/60"
+              fill="none"
+              viewBox="0 0 24 24"
+              stroke="currentColor"
+            >
+              <path
+                stroke-linecap="round"
+                stroke-linejoin="round"
+                stroke-width="2"
+                d="M19 9l-7 7-7-7"
+              />
+            </svg>
           </div>
         </div>
-      {/if}
-    </div>
+
+        {#if isSizeDropdownOpen}
+          <!-- Backdrop to close dropdown -->
+          <button
+            type="button"
+            class="fixed inset-0 z-[9998] cursor-default"
+            onclick={() => (isSizeDropdownOpen = 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 min-w-max"
+            style="bottom: calc(100vh - {sizeDropdownPosition()
+              .top}px + 4px); left: {sizeDropdownPosition().left}px;"
+          >
+            <div class="py-1">
+              {#each sizeOptions as size}
+                <button
+                  type="button"
+                  onclick={() => selectSize(size)}
+                  class="w-full px-3 py-1.5 text-left text-xs font-mono tracking-wide transition-colors duration-100 flex items-center gap-2 {params.size ===
+                  size
+                    ? 'bg-transparent text-exo-yellow'
+                    : 'text-exo-light-gray hover:text-exo-yellow'}"
+                >
+                  {#if params.size === size}
+                    <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}
+                  <span>{size}</span>
+                </button>
+              {/each}
+            </div>
+          </div>
+        {/if}
+      </div>
+    {/if}
 
     <!-- Quality -->
     <div class="flex items-center gap-1.5">
@@ -325,6 +357,59 @@
       </div>
     </div>
 
+    <!-- Number of Images (not in edit mode) -->
+    {#if !isEditMode}
+      <div class="flex items-center gap-1.5">
+        <span class="text-xs text-exo-light-gray uppercase tracking-wider"
+          >IMAGES:</span
+        >
+        <input
+          type="number"
+          min="1"
+          value={params.numImages}
+          oninput={handleNumImagesChange}
+          class="w-12 bg-exo-medium-gray/50 border border-exo-yellow/30 rounded px-2 py-1 text-xs font-mono text-exo-yellow text-center transition-all duration-200 hover:border-exo-yellow/50 focus:outline-none focus:border-exo-yellow/70"
+        />
+      </div>
+    {/if}
+
+    <!-- Stream toggle -->
+    <div class="flex items-center gap-1.5">
+      <span class="text-xs text-exo-light-gray uppercase tracking-wider"
+        >STREAM:</span
+      >
+      <button
+        type="button"
+        onclick={() => handleStreamChange(!params.stream)}
+        class="w-8 h-4 rounded-full transition-all duration-200 cursor-pointer relative {params.stream
+          ? 'bg-exo-yellow'
+          : 'bg-exo-medium-gray/50 border border-exo-yellow/30'}"
+        title={params.stream ? "Streaming enabled" : "Streaming disabled"}
+      >
+        <div
+          class="absolute top-0.5 w-3 h-3 rounded-full transition-all duration-200 {params.stream
+            ? 'right-0.5 bg-exo-black'
+            : 'left-0.5 bg-exo-light-gray'}"
+        ></div>
+      </button>
+    </div>
+
+    <!-- Partial Images (only when streaming) -->
+    {#if params.stream}
+      <div class="flex items-center gap-1.5">
+        <span class="text-xs text-exo-light-gray uppercase tracking-wider"
+          >PARTIALS:</span
+        >
+        <input
+          type="number"
+          min="0"
+          value={params.partialImages}
+          oninput={handlePartialImagesChange}
+          class="w-12 bg-exo-medium-gray/50 border border-exo-yellow/30 rounded px-2 py-1 text-xs font-mono text-exo-yellow text-center transition-all duration-200 hover:border-exo-yellow/50 focus:outline-none focus:border-exo-yellow/70"
+        />
+      </div>
+    {/if}
+
     <!-- Input Fidelity (edit mode only) -->
     {#if isEditMode}
       <div class="flex items-center gap-1.5">

dashboard/src/lib/components/TopologyGraph.svelte

@@ -5,22 +5,32 @@
     topologyData,
     isTopologyMinimized,
     debugMode,
+    nodeThunderboltBridge,
     type NodeInfo,
   } from "$lib/stores/app.svelte";
 
   interface Props {
     class?: string;
     highlightedNodes?: Set<string>;
+    filteredNodes?: Set<string>;
+    onNodeClick?: (nodeId: string) => void;
   }
 
-  let { class: className = "", highlightedNodes = new Set() }: Props = $props();
+  let {
+    class: className = "",
+    highlightedNodes = new Set(),
+    filteredNodes = new Set(),
+    onNodeClick,
+  }: Props = $props();
 
   let svgContainer: SVGSVGElement | undefined = $state();
   let resizeObserver: ResizeObserver | undefined;
+  let hoveredNodeId = $state<string | null>(null);
 
   const isMinimized = $derived(isTopologyMinimized());
   const data = $derived(topologyData());
   const debugEnabled = $derived(debugMode());
+  const tbBridgeData = $derived(nodeThunderboltBridge());
 
   function getNodeLabel(nodeId: string): string {
     const node = data?.nodes?.[nodeId];
@@ -522,10 +532,72 @@
         }
       }
 
+      let iconBaseWidth = nodeRadius * 1.2;
+      let iconBaseHeight = nodeRadius * 1.0;
+      const clipPathId = `clip-${nodeInfo.id.replace(/[^a-zA-Z0-9]/g, "-")}`;
+
+      const modelLower = modelId.toLowerCase();
+
+      // Check node states for styling
+      const isHighlighted = highlightedNodes.has(nodeInfo.id);
+      const isInFilter =
+        filteredNodes.size > 0 && filteredNodes.has(nodeInfo.id);
+      const isFilteredOut =
+        filteredNodes.size > 0 && !filteredNodes.has(nodeInfo.id);
+      const isHovered = hoveredNodeId === nodeInfo.id && !isInFilter;
+
+      // Holographic wireframe colors - bright yellow for filter, subtle yellow for hover, grey for filtered out
+      const wireColor = isInFilter
+        ? "rgba(255,215,0,1)" // Bright yellow for filter selection
+        : isHovered
+          ? "rgba(255,215,0,0.7)" // Subtle yellow for hover
+          : isHighlighted
+            ? "rgba(255,215,0,0.9)" // Yellow for instance highlight
+            : isFilteredOut
+              ? "rgba(140,140,140,0.6)" // Grey for filtered out
+              : "rgba(179,179,179,0.8)"; // Default
+      const wireColorBright = "rgba(255,255,255,0.9)";
+      const fillColor = isInFilter
+        ? "rgba(255,215,0,0.25)"
+        : isHovered
+          ? "rgba(255,215,0,0.12)"
+          : isHighlighted
+            ? "rgba(255,215,0,0.15)"
+            : "rgba(255,215,0,0.08)";
+      const strokeWidth = isInFilter
+        ? 3
+        : isHovered
+          ? 2
+          : isHighlighted
+            ? 2.5
+            : 1.5;
+      const screenFill = "rgba(0,20,40,0.9)";
+      const glowColor = "rgba(255,215,0,0.3)";
+
       const nodeG = nodesGroup
         .append("g")
         .attr("class", "graph-node")
-        .style("cursor", "pointer");
+        .style("cursor", onNodeClick ? "pointer" : "default")
+        .style("opacity", isFilteredOut ? 0.5 : 1);
+
+      // Add click and hover handlers - hover just updates state, styling is applied during render
+      nodeG
+        .on("click", (event: MouseEvent) => {
+          if (onNodeClick) {
+            event.stopPropagation();
+            onNodeClick(nodeInfo.id);
+          }
+        })
+        .on("mouseenter", () => {
+          if (onNodeClick) {
+            hoveredNodeId = nodeInfo.id;
+          }
+        })
+        .on("mouseleave", () => {
+          if (hoveredNodeId === nodeInfo.id) {
+            hoveredNodeId = null;
+          }
+        });
 
       // Add tooltip
       nodeG
@@ -534,27 +606,6 @@
           `${friendlyName}\nID: ${nodeInfo.id.slice(-8)}\nMemory: ${formatBytes(ramUsed)}/${formatBytes(ramTotal)}`,
         );
 
-      let iconBaseWidth = nodeRadius * 1.2;
-      let iconBaseHeight = nodeRadius * 1.0;
-      const clipPathId = `clip-${nodeInfo.id.replace(/[^a-zA-Z0-9]/g, "-")}`;
-
-      const modelLower = modelId.toLowerCase();
-
-      // Check if this node should be highlighted (from hovered instance)
-      const isHighlighted = highlightedNodes.has(nodeInfo.id);
-
-      // Holographic wireframe colors - yellow border when highlighted
-      const wireColor = isHighlighted
-        ? "rgba(255,215,0,0.9)"
-        : "rgba(179,179,179,0.8)";
-      const wireColorBright = "rgba(255,255,255,0.9)";
-      const fillColor = isHighlighted
-        ? "rgba(255,215,0,0.15)"
-        : "rgba(255,215,0,0.08)";
-      const strokeWidth = isHighlighted ? 2.5 : 1.5;
-      const screenFill = "rgba(0,20,40,0.9)";
-      const glowColor = "rgba(255,215,0,0.3)";
-
       if (modelLower === "mac studio") {
         // Mac Studio - classic cube with memory fill
         iconBaseWidth = nodeRadius * 1.25;
@@ -579,6 +630,7 @@
         // Main body (uniform color)
         nodeG
           .append("rect")
+          .attr("class", "node-outline")
           .attr("x", x)
           .attr("y", y)
           .attr("width", iconBaseWidth)
@@ -661,6 +713,7 @@
         // Main body (uniform color)
         nodeG
           .append("rect")
+          .attr("class", "node-outline")
           .attr("x", x)
           .attr("y", y)
           .attr("width", iconBaseWidth)
@@ -738,6 +791,7 @@
         // Screen outer frame
         nodeG
           .append("rect")
+          .attr("class", "node-outline")
           .attr("x", screenX)
           .attr("y", y)
           .attr("width", screenWidth)
@@ -846,6 +900,7 @@
         // Main shape
         nodeG
           .append("polygon")
+          .attr("class", "node-outline")
           .attr("points", hexPoints)
           .attr("fill", fillColor)
           .attr("stroke", wireColor)
@@ -1064,11 +1119,41 @@
           .attr("fill", "rgba(179,179,179,0.7)")
           .text(` (${ramUsagePercent.toFixed(0)}%)`);
       }
+
+      // Debug mode: Show TB bridge status
+      if (debugEnabled) {
+        const tbStatus = tbBridgeData[nodeInfo.id];
+        if (tbStatus) {
+          const tbY =
+            nodeInfo.y +
+            iconBaseHeight / 2 +
+            (showFullLabels ? 32 : showCompactLabels ? 26 : 22);
+          const tbFontSize = showFullLabels ? 9 : 7;
+          const tbColor = tbStatus.enabled
+            ? "rgba(234,179,8,0.9)"
+            : "rgba(100,100,100,0.7)";
+          const tbText = tbStatus.enabled ? "TB:ON" : "TB:OFF";
+          nodeG
+            .append("text")
+            .attr("x", nodeInfo.x)
+            .attr("y", tbY)
+            .attr("text-anchor", "middle")
+            .attr("fill", tbColor)
+            .attr("font-size", tbFontSize)
+            .attr("font-family", "SF Mono, Monaco, monospace")
+            .text(tbText);
+        }
+      }
     });
   }
 
   $effect(() => {
-    if (data) {
+    // Track all reactive dependencies that affect rendering
+    const _data = data;
+    const _hoveredNodeId = hoveredNodeId;
+    const _filteredNodes = filteredNodes;
+    const _highlightedNodes = highlightedNodes;
+    if (_data) {
       renderGraph();
     }
   });
@@ -1091,12 +1176,8 @@
 
 <style>
   :global(.graph-node) {
-    transition:
-      transform 0.2s ease,
-      opacity 0.2s ease;
-  }
-  :global(.graph-node:hover) {
-    filter: brightness(1.1);
+    /* Only transition opacity for filtered-out nodes, no transition on hover stroke changes */
+    transition: opacity 0.2s ease;
   }
   :global(.graph-link) {
     stroke: var(--exo-light-gray, #b3b3b3);

dashboard/src/routes/+page.svelte

@@ -19,6 +19,9 @@
     selectedPreviewModelId,
     isLoadingPreviews,
     selectPreviewModel,
+    togglePreviewNodeFilter,
+    clearPreviewNodeFilter,
+    previewNodeFilter,
     createConversation,
     setSelectedChatModel,
     selectedChatModel,
@@ -28,6 +31,8 @@
     toggleTopologyOnlyMode,
     chatSidebarVisible,
     toggleChatSidebarVisible,
+    thunderboltBridgeCycles,
+    nodeThunderboltBridge,
     type DownloadProgress,
     type PlacementPreview,
   } from "$lib/stores/app.svelte";
@@ -49,6 +54,41 @@
   const debugEnabled = $derived(debugMode());
   const topologyOnlyEnabled = $derived(topologyOnlyMode());
   const sidebarVisible = $derived(chatSidebarVisible());
+  const tbBridgeCycles = $derived(thunderboltBridgeCycles());
+  const tbBridgeData = $derived(nodeThunderboltBridge());
+  const nodeFilter = $derived(previewNodeFilter());
+
+  // Helper to get friendly node name from node ID
+  function getNodeName(nodeId: string): string {
+    const node = data?.nodes?.[nodeId];
+    return node?.friendly_name || nodeId.slice(0, 8) + "...";
+  }
+
+  // Helper to get the thunderbolt bridge service name from a cycle
+  function getTbBridgeServiceName(cycle: string[]): string {
+    // Try to find service name from any node in the cycle
+    for (const nodeId of cycle) {
+      const nodeData = tbBridgeData?.[nodeId];
+      if (nodeData?.serviceName) {
+        return nodeData.serviceName;
+      }
+    }
+    return "Thunderbolt Bridge"; // Fallback if no service name found
+  }
+
+  // Copy to clipboard state and function
+  let copiedCommand = $state(false);
+  async function copyToClipboard(text: string) {
+    try {
+      await navigator.clipboard.writeText(text);
+      copiedCommand = true;
+      setTimeout(() => {
+        copiedCommand = false;
+      }, 2000);
+    } catch (err) {
+      console.error("Failed to copy:", err);
+    }
+  }
 
   let mounted = $state(false);
 
@@ -90,6 +130,15 @@
       model.tasks.includes("ImageToImage")
     );
   }
+
+  // Helper to check if a model supports image editing
+  function modelSupportsImageEditing(modelId: string): boolean {
+    const model = models.find(
+      (m) => m.id === modelId || m.hugging_face_id === modelId,
+    );
+    if (!model?.tasks) return false;
+    return model.tasks.includes("ImageToImage");
+  }
   let selectedSharding = $state<"Pipeline" | "Tensor">("Pipeline");
   type InstanceMeta = "MlxRing" | "MlxIbv" | "MlxJaccl";
 
@@ -181,6 +230,9 @@
   // Preview card hover state for highlighting nodes in topology
   let hoveredPreviewNodes = $state<Set<string>>(new Set());
 
+  // Computed: Check if filter is active (from store)
+  const isFilterActive = $derived(() => nodeFilter.size > 0);
+
   // Helper to unwrap tagged instance for hover highlighting
   function unwrapInstanceNodes(instanceWrapped: unknown): Set<string> {
     if (!instanceWrapped || typeof instanceWrapped !== "object")
@@ -732,6 +784,8 @@
     instanceWrapped: unknown,
   ): {
     isDownloading: boolean;
+    isFailed: boolean;
+    errorMessage: string | null;
     progress: DownloadProgress | null;
     statusText: string;
     perNode: Array<{
@@ -743,6 +797,8 @@
     if (!downloadsData || Object.keys(downloadsData).length === 0) {
       return {
         isDownloading: false,
+        isFailed: false,
+        errorMessage: null,
         progress: null,
         statusText: "RUNNING",
         perNode: [],
@@ -754,6 +810,8 @@
     if (!instance || typeof instance !== "object") {
       return {
         isDownloading: false,
+        isFailed: false,
+        errorMessage: null,
         progress: null,
         statusText: "PREPARING",
         perNode: [],
@@ -809,6 +867,26 @@
           downloadKind
         ] as Record<string, unknown>;
 
+        // Handle DownloadFailed - return immediately with error info
+        if (downloadKind === "DownloadFailed") {
+          const downloadModelId = extractModelIdFromDownload(downloadPayload);
+          if (
+            instanceModelId &&
+            downloadModelId &&
+            downloadModelId === instanceModelId
+          ) {
+            return {
+              isDownloading: false,
+              isFailed: true,
+              errorMessage:
+                (downloadPayload.errorMessage as string) || "Download failed",
+              progress: null,
+              statusText: "FAILED",
+              perNode: [],
+            };
+          }
+        }
+
         if (downloadKind !== "DownloadOngoing") continue;
         if (!downloadPayload) continue;
 
@@ -844,6 +922,8 @@
       const statusInfo = deriveInstanceStatus(instanceWrapped);
       return {
         isDownloading: false,
+        isFailed: statusInfo.statusText === "FAILED",
+        errorMessage: null,
         progress: null,
         statusText: statusInfo.statusText,
         perNode: [],
@@ -856,6 +936,8 @@
 
     return {
       isDownloading: true,
+      isFailed: false,
+      errorMessage: null,
       progress: {
         totalBytes,
         downloadedBytes,
@@ -1451,7 +1533,7 @@
 
   // Get ALL filtered previews based on current settings (matching minimum nodes)
   // Note: previewsData already contains previews for the selected model (fetched via API)
-  // We filter by sharding/instance type and min nodes, returning ALL eligible previews
+  // Backend handles node_ids filtering, we filter by sharding/instance type and min nodes
   const filteredPreviews = $derived(() => {
     if (!selectedModelId || previewsData.length === 0) return [];
 
@@ -1584,7 +1666,86 @@
           <TopologyGraph
             class="w-full h-full"
             highlightedNodes={highlightedNodes()}
+            filteredNodes={nodeFilter}
+            onNodeClick={togglePreviewNodeFilter}
           />
+
+          <!-- Thunderbolt Bridge Cycle Warning -->
+          {#if tbBridgeCycles.length > 0}
+            {@const cycle = tbBridgeCycles[0]}
+            {@const serviceName = getTbBridgeServiceName(cycle)}
+            {@const disableCmd = `sudo networksetup -setnetworkserviceenabled "${serviceName}" off`}
+            <div class="absolute top-4 left-4 group" role="alert">
+              <div
+                class="flex items-center gap-2 px-3 py-2 rounded border border-yellow-500/50 bg-yellow-500/10 backdrop-blur-sm cursor-help"
+              >
+                <svg
+                  class="w-5 h-5 text-yellow-400 flex-shrink-0"
+                  fill="none"
+                  viewBox="0 0 24 24"
+                  stroke="currentColor"
+                  stroke-width="2"
+                >
+                  <path
+                    stroke-linecap="round"
+                    stroke-linejoin="round"
+                    d="M12 9v2m0 4h.01m-6.938 4h13.856c1.54 0 2.502-1.667 1.732-3L13.732 4c-.77-1.333-2.694-1.333-3.464 0L3.34 16c-.77 1.333.192 3 1.732 3z"
+                  />
+                </svg>
+                <span class="text-sm font-mono text-yellow-200">
+                  THUNDERBOLT BRIDGE CYCLE DETECTED
+                </span>
+              </div>
+
+              <!-- Tooltip on hover -->
+              <div
+                class="absolute top-full left-0 mt-2 w-80 p-3 rounded border border-yellow-500/30 bg-exo-dark-gray/95 backdrop-blur-sm opacity-0 invisible group-hover:opacity-100 group-hover:visible transition-all duration-200 z-50 shadow-lg"
+              >
+                <p class="text-xs text-white/80 mb-2">
+                  A network routing cycle was detected between nodes connected
+                  via Thunderbolt Bridge. This can cause connectivity issues.
+                </p>
+                <p class="text-xs text-white/60 mb-2">
+                  <span class="text-yellow-300">Affected nodes:</span>
+                  {cycle.map(getNodeName).join(" → ")}
+                </p>
+                <p class="text-xs text-white/60 mb-1">
+                  <span class="text-yellow-300">To fix:</span> Disable the Thunderbolt
+                  Bridge on one of the affected nodes:
+                </p>
+                <button
+                  type="button"
+                  onclick={() => copyToClipboard(disableCmd)}
+                  class="w-full flex items-center gap-2 text-[10px] font-mono bg-exo-black/60 px-2 py-1.5 rounded text-exo-yellow break-all text-left hover:bg-exo-black/80 transition-colors cursor-pointer group/copy"
+                  title="Click to copy"
+                >
+                  <span class="flex-1">{disableCmd}</span>
+                  <svg
+                    class="w-3.5 h-3.5 flex-shrink-0 text-white/40 group-hover/copy:text-exo-yellow transition-colors"
+                    fill="none"
+                    viewBox="0 0 24 24"
+                    stroke="currentColor"
+                    stroke-width="2"
+                  >
+                    {#if copiedCommand}
+                      <path
+                        stroke-linecap="round"
+                        stroke-linejoin="round"
+                        d="M5 13l4 4L19 7"
+                      />
+                    {:else}
+                      <path
+                        stroke-linecap="round"
+                        stroke-linejoin="round"
+                        d="M8 16H6a2 2 0 01-2-2V6a2 2 0 012-2h8a2 2 0 012 2v2m-6 12h8a2 2 0 002-2v-8a2 2 0 00-2-2h-8a2 2 0 00-2 2v8a2 2 0 002 2z"
+                      />
+                    {/if}
+                  </svg>
+                </button>
+              </div>
+            </div>
+          {/if}
+
           <!-- Exit topology-only mode button -->
           <button
             type="button"
@@ -1624,7 +1785,111 @@
             <TopologyGraph
               class="w-full h-full"
               highlightedNodes={highlightedNodes()}
+              filteredNodes={nodeFilter}
+              onNodeClick={togglePreviewNodeFilter}
             />
+
+            <!-- Thunderbolt Bridge Cycle Warning -->
+            {#if tbBridgeCycles.length > 0}
+              {@const cycle = tbBridgeCycles[0]}
+              {@const serviceName = getTbBridgeServiceName(cycle)}
+              {@const disableCmd = `sudo networksetup -setnetworkserviceenabled "${serviceName}" off`}
+              <div class="absolute top-4 left-4 group" role="alert">
+                <div
+                  class="flex items-center gap-2 px-3 py-2 rounded border border-yellow-500/50 bg-yellow-500/10 backdrop-blur-sm cursor-help"
+                >
+                  <svg
+                    class="w-5 h-5 text-yellow-400 flex-shrink-0"
+                    fill="none"
+                    viewBox="0 0 24 24"
+                    stroke="currentColor"
+                    stroke-width="2"
+                  >
+                    <path
+                      stroke-linecap="round"
+                      stroke-linejoin="round"
+                      d="M12 9v2m0 4h.01m-6.938 4h13.856c1.54 0 2.502-1.667 1.732-3L13.732 4c-.77-1.333-2.694-1.333-3.464 0L3.34 16c-.77 1.333.192 3 1.732 3z"
+                    />
+                  </svg>
+                  <span class="text-sm font-mono text-yellow-200">
+                    THUNDERBOLT BRIDGE CYCLE DETECTED
+                  </span>
+                </div>
+
+                <!-- Tooltip on hover -->
+                <div
+                  class="absolute top-full left-0 mt-2 w-80 p-3 rounded border border-yellow-500/30 bg-exo-dark-gray/95 backdrop-blur-sm opacity-0 invisible group-hover:opacity-100 group-hover:visible transition-all duration-200 z-50 shadow-lg"
+                >
+                  <p class="text-xs text-white/80 mb-2">
+                    A network routing cycle was detected between nodes connected
+                    via Thunderbolt Bridge. This can cause connectivity issues.
+                  </p>
+                  <p class="text-xs text-white/60 mb-2">
+                    <span class="text-yellow-300">Affected nodes:</span>
+                    {cycle.map(getNodeName).join(" → ")}
+                  </p>
+                  <p class="text-xs text-white/60 mb-1">
+                    <span class="text-yellow-300">To fix:</span> Disable the Thunderbolt
+                    Bridge on one of the affected nodes:
+                  </p>
+                  <button
+                    type="button"
+                    onclick={() => copyToClipboard(disableCmd)}
+                    class="w-full flex items-center gap-2 text-[10px] font-mono bg-exo-black/60 px-2 py-1.5 rounded text-exo-yellow break-all text-left hover:bg-exo-black/80 transition-colors cursor-pointer group/copy"
+                    title="Click to copy"
+                  >
+                    <span class="flex-1">{disableCmd}</span>
+                    <svg
+                      class="w-3.5 h-3.5 flex-shrink-0 text-white/40 group-hover/copy:text-exo-yellow transition-colors"
+                      fill="none"
+                      viewBox="0 0 24 24"
+                      stroke="currentColor"
+                      stroke-width="2"
+                    >
+                      {#if copiedCommand}
+                        <path
+                          stroke-linecap="round"
+                          stroke-linejoin="round"
+                          d="M5 13l4 4L19 7"
+                        />
+                      {:else}
+                        <path
+                          stroke-linecap="round"
+                          stroke-linejoin="round"
+                          d="M8 16H6a2 2 0 01-2-2V6a2 2 0 012-2h8a2 2 0 012 2v2m-6 12h8a2 2 0 002-2v-8a2 2 0 00-2-2h-8a2 2 0 00-2 2v8a2 2 0 002 2z"
+                        />
+                      {/if}
+                    </svg>
+                  </button>
+                </div>
+              </div>
+            {/if}
+
+            <!-- Node Filter Indicator (top-right corner) -->
+            {#if isFilterActive()}
+              <button
+                onclick={clearPreviewNodeFilter}
+                class="absolute top-2 right-2 flex items-center gap-1.5 px-2 py-1 bg-exo-dark-gray/80 border border-exo-yellow/40 rounded text-exo-yellow hover:border-exo-yellow/60 transition-colors cursor-pointer backdrop-blur-sm"
+                title="Clear filter"
+              >
+                <span class="text-[10px] font-mono tracking-wider">
+                  FILTER: {nodeFilter.size}
+                </span>
+                <svg
+                  class="w-3 h-3"
+                  fill="none"
+                  viewBox="0 0 24 24"
+                  stroke="currentColor"
+                  stroke-width="2"
+                >
+                  <path
+                    stroke-linecap="round"
+                    stroke-linejoin="round"
+                    d="M6 18L18 6M6 6l12 12"
+                  />
+                </svg>
+              </button>
+            {/if}
           </div>
 
           <!-- Chat Input - Below topology -->
@@ -2061,6 +2326,13 @@
                           >
                             {downloadInfo.statusText}
                           </div>
+                          {#if downloadInfo.isFailed && downloadInfo.errorMessage}
+                            <div
+                              class="text-xs text-red-400/80 font-mono mt-1 break-words"
+                            >
+                              {downloadInfo.errorMessage}
+                            </div>
+                          {/if}
                         {/if}
                       </div>
                     </div>
@@ -2106,6 +2378,9 @@
                     {@const isImageModel = modelSupportsImageGeneration(
                       foundModel.id,
                     )}
+                    {@const isImageEditModel = modelSupportsImageEditing(
+                      foundModel.id,
+                    )}
                     <span
                       class="flex items-center justify-between gap-2 w-full pr-4"
                     >
@@ -2132,6 +2407,22 @@
                             <polyline points="21 15 16 10 5 21" />
                           </svg>
                         {/if}
+                        {#if isImageEditModel}
+                          <svg
+                            class="w-4 h-4 flex-shrink-0 text-exo-yellow"
+                            fill="none"
+                            viewBox="0 0 24 24"
+                            stroke="currentColor"
+                            stroke-width="2"
+                          >
+                            <path
+                              d="M11 4H4a2 2 0 0 0-2 2v14a2 2 0 0 0 2 2h14a2 2 0 0 0 2-2v-7"
+                            />
+                            <path
+                              d="M18.5 2.5a2.121 2.121 0 0 1 3 3L12 15l-4 1 1-4 9.5-9.5z"
+                            />
+                          </svg>
+                        {/if}
                         <span class="truncate"
                           >{foundModel.name || foundModel.id}</span
                         >
@@ -2204,6 +2495,9 @@
                       {@const isImageModel = modelSupportsImageGeneration(
                         model.id,
                       )}
+                      {@const isImageEditModel = modelSupportsImageEditing(
+                        model.id,
+                      )}
                       <button
                         type="button"
                         onclick={() => {
@@ -2244,6 +2538,23 @@
                               <polyline points="21 15 16 10 5 21" />
                             </svg>
                           {/if}
+                          {#if isImageEditModel}
+                            <svg
+                              class="w-4 h-4 flex-shrink-0 text-exo-yellow"
+                              fill="none"
+                              viewBox="0 0 24 24"
+                              stroke="currentColor"
+                              stroke-width="2"
+                              aria-label="Image editing model"
+                            >
+                              <path
+                                d="M11 4H4a2 2 0 0 0-2 2v14a2 2 0 0 0 2 2h14a2 2 0 0 0 2-2v-7"
+                              />
+                              <path
+                                d="M18.5 2.5a2.121 2.121 0 0 1 3 3L12 15l-4 1 1-4 9.5-9.5z"
+                              />
+                            </svg>
+                          {/if}
                           <span class="truncate">{model.name || model.id}</span>
                         </span>
                         <span
@@ -2564,7 +2875,36 @@
               <div
                 class="relative aspect-square bg-exo-dark-gray rounded-lg overflow-hidden"
               >
-                <TopologyGraph highlightedNodes={highlightedNodes()} />
+                <TopologyGraph
+                  highlightedNodes={highlightedNodes()}
+                  filteredNodes={nodeFilter}
+                  onNodeClick={togglePreviewNodeFilter}
+                />
+
+                <!-- Thunderbolt Bridge Cycle Warning (compact) -->
+                {#if tbBridgeCycles.length > 0}
+                  <div
+                    class="absolute top-2 left-2 flex items-center gap-1.5 px-2 py-1 rounded border border-yellow-500/50 bg-yellow-500/10 backdrop-blur-sm"
+                    title="Thunderbolt Bridge cycle detected - click to view details"
+                  >
+                    <svg
+                      class="w-3.5 h-3.5 text-yellow-400"
+                      fill="none"
+                      viewBox="0 0 24 24"
+                      stroke="currentColor"
+                      stroke-width="2"
+                    >
+                      <path
+                        stroke-linecap="round"
+                        stroke-linejoin="round"
+                        d="M12 9v2m0 4h.01m-6.938 4h13.856c1.54 0 2.502-1.667 1.732-3L13.732 4c-.77-1.333-2.694-1.333-3.464 0L3.34 16c-.77 1.333.192 3 1.732 3z"
+                      />
+                    </svg>
+                    <span class="text-[10px] font-mono text-yellow-200"
+                      >TB CYCLE</span
+                    >
+                  </div>
+                {/if}
               </div>
             </button>
 
@@ -2993,6 +3333,13 @@
                             >
                               {downloadInfo.statusText}
                             </div>
+                            {#if downloadInfo.isFailed && downloadInfo.errorMessage}
+                              <div
+                                class="text-xs text-red-400/80 font-mono mt-1 break-words"
+                              >
+                                {downloadInfo.errorMessage}
+                              </div>
+                            {/if}
                           {/if}
                         </div>
                       </div>

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

@@ -6,6 +6,8 @@
     type DownloadProgress,
     refreshState,
     lastUpdate as lastUpdateStore,
+    startDownload,
+    deleteDownload,
   } from "$lib/stores/app.svelte";
   import HeaderNav from "$lib/components/HeaderNav.svelte";
 
@@ -28,6 +30,7 @@
     etaMs: number;
     status: "completed" | "downloading";
     files: FileProgress[];
+    shardMetadata?: Record<string, unknown>;
   };
 
   type NodeEntry = {
@@ -172,33 +175,6 @@
   }
 
   let downloadOverview = $state<NodeEntry[]>([]);
-  let models = $state<Array<{ id: string; storage_size_megabytes?: number }>>(
-    [],
-  );
-
-  async function fetchModels() {
-    try {
-      const response = await fetch("/models");
-      if (response.ok) {
-        const data = await response.json();
-        models = data.data || [];
-      }
-    } catch (error) {
-      console.error("Failed to fetch models:", error);
-    }
-  }
-
-  function getModelTotalBytes(
-    modelId: string,
-    downloadTotalBytes: number,
-  ): number {
-    if (downloadTotalBytes > 0) return downloadTotalBytes;
-    const model = models.find((m) => m.id === modelId);
-    if (model?.storage_size_megabytes) {
-      return model.storage_size_megabytes * 1024 * 1024;
-    }
-    return 0;
-  }
 
   $effect(() => {
     try {
@@ -296,6 +272,12 @@
             }
           }
 
+          // Extract shard_metadata for use with download actions
+          const shardMetadata = (downloadPayload.shard_metadata ??
+            downloadPayload.shardMetadata) as
+            | Record<string, unknown>
+            | undefined;
+
           const entry: ModelEntry = {
             modelId,
             prettyName,
@@ -312,6 +294,7 @@
                 ? "completed"
                 : "downloading",
             files,
+            shardMetadata,
           };
 
           const existing = modelMap.get(modelId);
@@ -373,7 +356,6 @@
   onMount(() => {
     // Ensure we fetch at least once when visiting downloads directly
     refreshState();
-    fetchModels();
   });
 </script>
 
@@ -482,7 +464,7 @@
                     {#if model.status !== "completed"}
                       <div class="text-[11px] text-exo-light-gray font-mono">
                         {formatBytes(model.downloadedBytes)} / {formatBytes(
-                          getModelTotalBytes(model.modelId, model.totalBytes),
+                          model.totalBytes,
                         )}
                       </div>
                     {/if}
@@ -497,6 +479,52 @@
                     >
                       {pct.toFixed(1)}%
                     </span>
+                    {#if model.status !== "completed" && model.shardMetadata}
+                      <button
+                        type="button"
+                        class="text-exo-light-gray hover:text-exo-yellow transition-colors"
+                        onclick={() =>
+                          startDownload(node.nodeId, model.shardMetadata!)}
+                        title="Start download"
+                      >
+                        <svg
+                          class="w-4 h-4"
+                          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>
+                    {/if}
+                    {#if model.status === "completed"}
+                      <button
+                        type="button"
+                        class="text-exo-light-gray hover:text-red-400 transition-colors"
+                        onclick={() =>
+                          deleteDownload(node.nodeId, model.modelId)}
+                        title="Delete download"
+                      >
+                        <svg
+                          class="w-4 h-4"
+                          viewBox="0 0 20 20"
+                          fill="none"
+                          stroke="currentColor"
+                          stroke-width="2"
+                        >
+                          <path
+                            d="M4 6h12M8 6V4h4v2m1 0v10a1 1 0 01-1 1H8a1 1 0 01-1-1V6h6"
+                            stroke-linecap="round"
+                            stroke-linejoin="round"
+                          ></path>
+                        </svg>
+                      </button>
+                    {/if}
                     <button
                       type="button"
                       class="text-exo-light-gray hover:text-exo-yellow transition-colors"

pyproject.toml

@@ -13,20 +13,21 @@ dependencies = [
     "filelock>=3.18.0",
     "rustworkx>=0.17.1",
     "huggingface-hub>=0.33.4",
+    "typer", # for huggingface-cli
     "psutil>=7.0.0",
     "loguru>=0.7.3",
     "exo_pyo3_bindings", # rust bindings
     "anyio==4.11.0",
     "mlx==0.30.3; sys_platform == 'darwin'",
     "mlx[cpu]==0.30.3; sys_platform == 'linux'",
-    "mlx-lm @ git+https://github.com/AlexCheema/mlx-lm.git@fix-transformers-5.0.0rc2",
+    "mlx-lm==0.30.5",
     "tiktoken>=0.12.0", # required for kimi k2 tokenizer
     "hypercorn>=0.18.0",
     "openai-harmony>=0.0.8",
     "httpx>=0.28.1",
     "tomlkit>=0.14.0",
     "pillow>=11.0,<12.0", # compatibility with mflux
-    "mflux>=0.14.2",
+    "mflux==0.15.4",
     "python-multipart>=0.0.21",
 ]
 
@@ -34,7 +35,7 @@ dependencies = [
 exo-master = "exo.master.main:main"
 exo-worker = "exo.worker.main:main"
 exo = "exo.main:main"
-exo-rsh = "exo.rsh.client:main"
+exo-eval = "bench.exo_eval:main"
 
 # dependencies only required for development
 [dependency-groups]
@@ -52,6 +53,9 @@ dev = [
 # cuda = [
 #     "mlx[cuda]==0.26.3",
 # ]
+eval = [
+    "lm_eval[api]",
+]
 
 ###
 # workspace configuration
@@ -67,6 +71,7 @@ exo_pyo3_bindings = { workspace = true }
 # Uncomment to use local mlx/mlx-lm development versions:
 # mlx = { path = "/Users/Shared/mlx", editable=true }
 # mlx-lm = { path = "/Users/Shared/mlx-lm", editable=true }
+mlx-lm = { git = "https://github.com/davidmcc73/mlx-lm.git", branch = "main" }
 
 [build-system]
 requires = ["uv_build>=0.8.9,<0.9.0"]

src/exo/cli/__init__.py

@@ -1,32 +0,0 @@
-"""Exo CLI - SLURM-compatible job management commands."""
-
-
-def run_subcommand(command: str, args: list[str]) -> int:
-    """Route to the appropriate subcommand handler.
-
-    Args:
-        command: The subcommand name (sbatch, squeue, scancel, salloc)
-        args: Command line arguments for the subcommand
-
-    Returns:
-        Exit code from the subcommand
-    """
-    if command == "sbatch":
-        from exo.cli.sbatch import main
-
-        return main(args)
-    elif command == "squeue":
-        from exo.cli.squeue import main
-
-        return main(args)
-    elif command == "scancel":
-        from exo.cli.scancel import main
-
-        return main(args)
-    elif command == "salloc":
-        from exo.cli.salloc import main
-
-        return main(args)
-    else:
-        print(f"Unknown subcommand: {command}")
-        return 1

src/exo/cli/common.py

@@ -1,118 +0,0 @@
-"""Common utilities for Exo CLI commands."""
-
-import json
-import os
-import urllib.request
-from typing import Any
-from urllib.error import HTTPError, URLError
-
-# Default API endpoint
-DEFAULT_API_HOST = "localhost"
-DEFAULT_API_PORT = 52415
-
-
-def get_api_base() -> str:
-    """Get the API base URL from environment or defaults."""
-    host = os.environ.get("EXO_API_HOST", DEFAULT_API_HOST)
-    port = os.environ.get("EXO_API_PORT", str(DEFAULT_API_PORT))
-    return f"http://{host}:{port}"
-
-
-def api_request(
-    method: str,
-    path: str,
-    data: dict[str, Any] | None = None,
-) -> dict[str, Any] | list[Any]:
-    """Make an API request to the Exo server.
-
-    Args:
-        method: HTTP method (GET, POST, DELETE, etc.)
-        path: API path (e.g., "/flash/instances")
-        data: Optional JSON data for POST/PUT requests
-
-    Returns:
-        Parsed JSON response
-
-    Raises:
-        SystemExit: On connection or HTTP errors
-    """
-    url = f"{get_api_base()}{path}"
-
-    request_data = None
-    if data is not None:
-        request_data = json.dumps(data).encode("utf-8")
-
-    req = urllib.request.Request(
-        url,
-        data=request_data,
-        method=method,
-    )
-    req.add_header("Content-Type", "application/json")
-
-    try:
-        with urllib.request.urlopen(req, timeout=30) as response:  # pyright: ignore[reportAny]
-            body: str = response.read().decode("utf-8")  # pyright: ignore[reportAny]
-            if body:
-                return json.loads(body)  # pyright: ignore[reportAny]
-            return {}
-    except HTTPError as e:
-        error_body = e.read().decode("utf-8") if e.fp else ""
-        print(f"API error: {e.code} {e.reason}")
-        if error_body:
-            try:
-                error_json: dict[str, str] = json.loads(error_body)  # pyright: ignore[reportAny]
-                if "detail" in error_json:
-                    print(f"  {error_json['detail']}")
-            except json.JSONDecodeError:
-                print(f"  {error_body}")
-        raise SystemExit(1) from None
-    except URLError as e:
-        print(f"Connection error: {e.reason}")
-        print(f"Is Exo running at {get_api_base()}?")
-        raise SystemExit(1) from None
-
-
-def truncate_id(instance_id: str, length: int = 8) -> str:
-    """Truncate a UUID for display.
-
-    Args:
-        instance_id: Full UUID string
-        length: Number of characters to keep
-
-    Returns:
-        Truncated ID without hyphens
-    """
-    return instance_id.replace("-", "")[:length]
-
-
-def format_table(headers: list[str], rows: list[list[str]]) -> str:
-    """Format data as a simple text table.
-
-    Args:
-        headers: Column headers
-        rows: List of rows, each row is a list of column values
-
-    Returns:
-        Formatted table string
-    """
-    if not rows:
-        return "  ".join(f"{h:<10}" for h in headers)
-
-    # Calculate column widths
-    widths = [len(h) for h in headers]
-    for row in rows:
-        for i, cell in enumerate(row):
-            if i < len(widths):
-                widths[i] = max(widths[i], len(cell))
-
-    # Build format string
-    fmt = "  ".join(f"{{:<{w}}}" for w in widths)
-
-    # Format output
-    lines = [fmt.format(*headers)]
-    for row in rows:
-        # Pad row if needed
-        padded = row + [""] * (len(headers) - len(row))
-        lines.append(fmt.format(*padded[: len(headers)]))
-
-    return "\n".join(lines)

src/exo/cli/salloc.py

@@ -1,100 +0,0 @@
-"""salloc - Allocate nodes for interactive use.
-
-Usage:
-    exo salloc [options] [-- command [args...]]
-
-Options:
-    -N, --nodes N       Number of nodes to allocate (default: 1)
-    --hosts HOSTS       Comma-separated list of hostnames
-
-If a command is provided after --, it will be executed with
-SLURM-like environment variables set:
-    SLURM_JOB_NODELIST  - Comma-separated list of allocated nodes
-    SLURM_NNODES        - Number of allocated nodes
-
-Examples:
-    exo salloc --nodes=2 --hosts=node1,node2 -- mpirun ./my_program
-    exo salloc --hosts=localhost -- bash
-"""
-
-import argparse
-import os
-import subprocess
-import sys
-
-
-def main(args: list[str]) -> int:
-    """Main entry point for salloc command."""
-    # Split args at -- if present
-    cmd_args: list[str] = []
-    salloc_args = args
-
-    if "--" in args:
-        idx = args.index("--")
-        salloc_args = args[:idx]
-        cmd_args = args[idx + 1 :]
-
-    parser = argparse.ArgumentParser(
-        prog="exo salloc",
-        description="Allocate nodes for interactive use",
-    )
-    parser.add_argument(
-        "-N",
-        "--nodes",
-        type=int,
-        default=1,
-        help="Number of nodes to allocate (default: 1)",
-    )
-    parser.add_argument(
-        "--hosts",
-        help="Comma-separated list of hostnames (required)",
-    )
-
-    parsed = parser.parse_args(salloc_args)
-
-    nodes: int = parsed.nodes  # pyright: ignore[reportAny]
-    hosts: str | None = parsed.hosts  # pyright: ignore[reportAny]
-
-    # Require explicit hosts since we can't discover them from topology
-    if not hosts:
-        print("Error: --hosts is required (e.g., --hosts=node1,node2)", file=sys.stderr)
-        print("       The Exo topology doesn't expose hostnames.", file=sys.stderr)
-        return 1
-
-    host_list = [h.strip() for h in hosts.split(",") if h.strip()]
-
-    if len(host_list) < nodes:
-        print(
-            f"Error: Requested {nodes} nodes but only {len(host_list)} hosts provided",
-            file=sys.stderr,
-        )
-        return 1
-
-    # Use first N hosts
-    allocated_hosts = host_list[:nodes]
-    nodelist = ",".join(allocated_hosts)
-
-    # Set environment variables
-    env = os.environ.copy()
-    env["SLURM_JOB_NODELIST"] = nodelist
-    env["SLURM_NNODES"] = str(nodes)
-
-    print(f"salloc: Granted job allocation on {nodes} node(s)")
-    print(f"salloc: Nodes: {nodelist}")
-
-    if cmd_args:
-        # Run the command
-        print(f"salloc: Running: {' '.join(cmd_args)}")
-        result = subprocess.run(cmd_args, env=env)
-        return result.returncode
-    else:
-        # Start interactive shell
-        shell = os.environ.get("SHELL", "/bin/bash")
-        print(f"salloc: Starting shell {shell}")
-        print("salloc: Use 'exit' to release allocation")
-        result = subprocess.run([shell], env=env)
-        return result.returncode
-
-
-if __name__ == "__main__":
-    sys.exit(main(sys.argv[1:]))

src/exo/cli/sbatch.py

@@ -1,233 +0,0 @@
-"""sbatch - Submit a batch job to Exo.
-
-Usage:
-    exo sbatch [options] <script|executable>
-    exo sbatch --job-name=NAME --nodes=N <executable>
-
-Options:
-    -J, --job-name NAME         Job name
-    -N, --nodes N               Number of nodes (default: 1)
-    --ntasks-per-node N         Tasks per node (default: 1)
-    -D, --chdir DIR             Working directory
-    --hosts HOSTS               Comma-separated list of hostnames
-
-Job scripts can contain #SBATCH directives:
-    #!/bin/bash
-    #SBATCH --job-name=Sod2D
-    #SBATCH --nodes=2
-    #SBATCH --chdir=/path/to/workdir
-
-    /path/to/flash4
-"""
-
-import argparse
-import os
-import re
-import sys
-
-from exo.cli.common import api_request, truncate_id
-
-
-def parse_job_script(script_path: str) -> tuple[dict[str, str], str | None]:
-    """Parse a job script for #SBATCH directives and executable.
-
-    Args:
-        script_path: Path to the job script
-
-    Returns:
-        Tuple of (directives dict, executable path or None)
-    """
-    directives: dict[str, str] = {}
-    executable: str | None = None
-
-    with open(script_path, "r") as f:
-        for line in f:
-            line = line.strip()
-
-            # Parse #SBATCH directives
-            if line.startswith("#SBATCH"):
-                # Handle both --option=value and --option value formats
-                match = re.match(r"#SBATCH\s+(-\w|--[\w-]+)(?:=|\s+)(.+)", line)
-                if match:
-                    opt, val = match.groups()
-                    directives[opt.lstrip("-")] = val.strip()
-                continue
-
-            # Skip comments and empty lines
-            if line.startswith("#") or not line:
-                continue
-
-            # First non-comment, non-directive line is the executable
-            if executable is None:
-                # Handle lines like "/path/to/flash4" or "srun /path/to/flash4"
-                parts = line.split()
-                if parts:
-                    # Skip srun/mpirun prefixes if present
-                    for part in parts:
-                        if not part.startswith("-") and "/" in part:
-                            executable = part
-                            break
-                    if executable is None and parts:
-                        executable = parts[-1]  # Last token
-
-    return directives, executable
-
-
-def main(args: list[str]) -> int:
-    """Main entry point for sbatch command."""
-    parser = argparse.ArgumentParser(
-        prog="exo sbatch",
-        description="Submit a batch job to Exo",
-    )
-    parser.add_argument(
-        "script",
-        help="Job script or executable path",
-    )
-    parser.add_argument(
-        "-J",
-        "--job-name",
-        dest="job_name",
-        help="Job name",
-    )
-    parser.add_argument(
-        "-N",
-        "--nodes",
-        type=int,
-        default=1,
-        help="Number of nodes (default: 1)",
-    )
-    parser.add_argument(
-        "--ntasks-per-node",
-        type=int,
-        default=1,
-        help="Tasks per node (default: 1)",
-    )
-    parser.add_argument(
-        "-D",
-        "--chdir",
-        help="Working directory",
-    )
-    parser.add_argument(
-        "--hosts",
-        help="Comma-separated list of hostnames",
-    )
-
-    parsed = parser.parse_args(args)
-
-    # Extract typed values from namespace
-    script_path: str = parsed.script  # pyright: ignore[reportAny]
-    arg_job_name: str | None = parsed.job_name  # pyright: ignore[reportAny]
-    arg_nodes: int = parsed.nodes  # pyright: ignore[reportAny]
-    arg_ntasks: int = parsed.ntasks_per_node  # pyright: ignore[reportAny]
-    arg_chdir: str | None = parsed.chdir  # pyright: ignore[reportAny]
-    arg_hosts: str | None = parsed.hosts  # pyright: ignore[reportAny]
-
-    # Determine if input is a script or direct executable
-    executable: str | None = None
-    directives: dict[str, str] = {}
-
-    if os.path.isfile(script_path):
-        # Check if it's a binary file (executable) or text script
-        is_binary = False
-        try:
-            with open(script_path, "rb") as f:
-                chunk = f.read(512)
-                # Binary files typically contain null bytes
-                is_binary = b"\x00" in chunk
-        except OSError:
-            pass
-
-        if is_binary:
-            # It's a binary executable
-            executable = script_path
-        else:
-            # Try to read as text
-            try:
-                with open(script_path, "r") as f:
-                    first_line = f.readline()
-                    f.seek(0)
-                    content = f.read(1024)
-
-                if first_line.startswith("#!") or "#SBATCH" in content:
-                    # It's a job script - parse it
-                    directives, executable = parse_job_script(script_path)
-                else:
-                    # It's an executable (text but no shebang/directives)
-                    executable = script_path
-            except UnicodeDecodeError:
-                # Can't read as text - treat as binary executable
-                executable = script_path
-    else:
-        # Not a file - treat as executable path
-        executable = script_path
-
-    if executable is None:
-        print("Error: No executable found in job script", file=sys.stderr)
-        return 1
-
-    # Build job parameters - CLI args override script directives
-    job_name = arg_job_name or directives.get("job-name") or directives.get("J")
-    if not job_name:
-        # Generate name from executable
-        job_name = os.path.basename(executable).replace(".", "_")
-
-    nodes = arg_nodes
-    if "nodes" in directives:
-        nodes = int(directives["nodes"])
-    if "N" in directives:
-        nodes = int(directives["N"])
-    if arg_nodes != 1:  # CLI override
-        nodes = arg_nodes
-
-    ntasks = arg_ntasks
-    if "ntasks-per-node" in directives:
-        ntasks = int(directives["ntasks-per-node"])
-    if arg_ntasks != 1:  # CLI override
-        ntasks = arg_ntasks
-
-    workdir = arg_chdir or directives.get("chdir") or directives.get("D")
-    if not workdir:
-        workdir = os.getcwd()
-
-    hosts = arg_hosts or directives.get("hosts") or ""
-
-    # Resolve executable to absolute path
-    if not os.path.isabs(executable):
-        executable = os.path.abspath(os.path.join(workdir, executable))
-
-    # Submit job via API using query parameters
-    from urllib.parse import urlencode
-
-    params = {
-        "simulation_name": job_name,
-        "flash_executable_path": executable,
-        "parameter_file_path": "",  # FLASH par file - use default
-        "working_directory": workdir,
-        "ranks_per_node": str(ntasks),
-        "min_nodes": str(nodes),
-        "hosts": hosts,
-    }
-
-    query_string = urlencode(params)
-    result = api_request("POST", f"/flash/launch?{query_string}")
-
-    # Print job submission confirmation
-    if isinstance(result, dict):
-        instance_id_val = result.get("instance_id")
-
-        if instance_id_val is not None:
-            job_id = truncate_id(str(instance_id_val))  # pyright: ignore[reportAny]
-            print(f"Submitted batch job {job_id}")
-        else:
-            # Instance created asynchronously - user should check squeue
-            print("Job submitted successfully")
-            print("Use 'exo squeue' to view job ID")
-    else:
-        print("Job submitted successfully")
-        print("Use 'exo squeue' to view job ID")
-
-    return 0
-
-
-if __name__ == "__main__":
-    sys.exit(main(sys.argv[1:]))

src/exo/cli/scancel.py

@@ -1,95 +0,0 @@
-"""scancel - Cancel jobs in the Exo queue.
-
-Usage:
-    exo scancel <jobid> [<jobid>...]
-
-Arguments:
-    jobid   Job ID (or prefix) to cancel. Can specify multiple.
-
-Examples:
-    exo scancel abc123          # Cancel job starting with abc123
-    exo scancel abc123 def456   # Cancel multiple jobs
-"""
-
-import argparse
-import sys
-from typing import Any, cast
-
-from exo.cli.common import api_request, truncate_id
-
-
-def main(args: list[str]) -> int:
-    """Main entry point for scancel command."""
-    parser = argparse.ArgumentParser(
-        prog="exo scancel",
-        description="Cancel jobs in the Exo queue",
-    )
-    parser.add_argument(
-        "jobids",
-        nargs="+",
-        help="Job ID(s) to cancel",
-    )
-
-    parsed = parser.parse_args(args)
-    jobids: list[str] = parsed.jobids  # pyright: ignore[reportAny]
-
-    # Fetch current jobs to resolve partial IDs
-    result = api_request("GET", "/flash/instances")
-    if isinstance(result, list):
-        instances = cast(list[dict[str, Any]], result)
-    else:
-        instances = cast(list[dict[str, Any]], result.get("instances", []))
-
-    # Build lookup of full IDs
-    id_map: dict[str, str] = {}
-    for inst in instances:
-        iid = inst.get("instance_id", "")  # pyright: ignore[reportAny]
-        full_id = str(iid) if iid else ""  # pyright: ignore[reportAny]
-        if full_id:
-            # Map both full ID and truncated versions
-            normalized = full_id.replace("-", "").lower()
-            id_map[normalized] = full_id
-            # Also map prefixes
-            for length in range(4, len(normalized) + 1):
-                prefix = normalized[:length]
-                if prefix not in id_map:
-                    id_map[prefix] = full_id
-
-    cancelled = 0
-    errors = 0
-
-    for jobid in jobids:
-        search = jobid.lower().replace("-", "")
-
-        # Find matching full ID
-        full_id = id_map.get(search)
-        if not full_id:
-            # Try prefix match
-            matches = [fid for key, fid in id_map.items() if key.startswith(search)]
-            if len(matches) == 1:
-                full_id = matches[0]
-            elif len(matches) > 1:
-                print(f"Ambiguous job ID: {jobid} matches multiple jobs")
-                errors += 1
-                continue
-            else:
-                print(f"Job not found: {jobid}")
-                errors += 1
-                continue
-
-        # Cancel the job
-        try:
-            api_request("DELETE", f"/flash/{full_id}")
-            print(f"Job {truncate_id(full_id)} cancelled")
-            cancelled += 1
-        except SystemExit:
-            print(f"Failed to cancel job {truncate_id(full_id)}")
-            errors += 1
-
-    if errors > 0 and cancelled == 0:
-        return 1
-    return 0
-
-
-if __name__ == "__main__":
-    sys.exit(main(sys.argv[1:]))

src/exo/cli/squeue.py

@@ -1,165 +0,0 @@
-"""squeue - View the Exo job queue.
-
-Usage:
-    exo squeue [options]
-
-Options:
-    -l, --long      Show detailed output
-    -j, --job ID    Show only this job
-
-Output columns:
-    JOBID   - Job identifier (truncated UUID)
-    NAME    - Job name
-    NODES   - Number of nodes
-    STATE   - Job state (PENDING, RUNNING, FAILED, etc.)
-"""
-
-import argparse
-import sys
-from typing import Any, cast
-
-from exo.cli.common import api_request, format_table, truncate_id
-
-# Map Exo runner statuses to SLURM-like states
-STATUS_MAP: dict[str, str] = {
-    "RunnerIdle": "PENDING",
-    "RunnerConnecting": "CONFIGURING",
-    "RunnerConnected": "CONFIGURING",
-    "RunnerLoading": "CONFIGURING",
-    "RunnerLoaded": "CONFIGURING",
-    "RunnerWarmingUp": "CONFIGURING",
-    "RunnerReady": "COMPLETING",
-    "RunnerRunning": "RUNNING",
-    "RunnerShuttingDown": "COMPLETING",
-    "RunnerShutdown": "COMPLETED",
-    "RunnerFailed": "FAILED",
-}
-
-
-def get_job_state(runner_statuses: dict[str, Any]) -> str:
-    """Determine overall job state from runner statuses."""
-    if not runner_statuses:
-        return "PENDING"
-
-    states: set[str] = set()
-    for status_val in runner_statuses.values():  # pyright: ignore[reportAny]
-        if isinstance(status_val, dict):
-            # Extract status type from discriminated union
-            type_val = status_val.get("type", "RunnerIdle")  # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType]
-            status_type = str(type_val) if type_val else "RunnerIdle"  # pyright: ignore[reportUnknownArgumentType]
-        elif isinstance(status_val, str):
-            status_type = status_val
-        else:
-            status_type = "RunnerIdle"
-        # Strip parentheses from status strings like "RunnerRunning()"
-        if status_type.endswith("()"):
-            status_type = status_type[:-2]
-        states.add(STATUS_MAP.get(status_type, "UNKNOWN"))
-
-    # Priority order for overall state
-    if "FAILED" in states:
-        return "FAILED"
-    if "RUNNING" in states:
-        return "RUNNING"
-    if "CONFIGURING" in states:
-        return "CONFIGURING"
-    if "COMPLETING" in states:
-        return "COMPLETING"
-    if "COMPLETED" in states:
-        return "COMPLETED"
-    if "PENDING" in states:
-        return "PENDING"
-    return "UNKNOWN"
-
-
-def main(args: list[str]) -> int:
-    """Main entry point for squeue command."""
-    parser = argparse.ArgumentParser(
-        prog="exo squeue",
-        description="View the Exo job queue",
-    )
-    parser.add_argument(
-        "-l",
-        "--long",
-        action="store_true",
-        help="Show detailed output",
-    )
-    parser.add_argument(
-        "-j",
-        "--job",
-        help="Show only this job ID",
-    )
-
-    parsed = parser.parse_args(args)
-
-    # Extract typed values
-    long_format: bool = parsed.long  # pyright: ignore[reportAny]
-    job_filter: str | None = parsed.job  # pyright: ignore[reportAny]
-
-    # Fetch jobs from API - returns list directly
-    result = api_request("GET", "/flash/instances")
-    # API returns list directly, not {"instances": [...]}
-    if isinstance(result, list):
-        instances = cast(list[dict[str, Any]], result)
-    else:
-        instances = cast(list[dict[str, Any]], result.get("instances", []))
-
-    if not instances:
-        # No jobs - just print header
-        if long_format:
-            print("JOBID           NAME            NODES  RANKS  STATE        WORKDIR")
-        else:
-            print("JOBID       NAME            NODES  STATE")
-        return 0
-
-    # Filter by job ID if specified
-    if job_filter:
-        search = job_filter.lower()
-        filtered: list[dict[str, Any]] = []
-        for i in instances:
-            iid = i.get("instance_id", "")  # pyright: ignore[reportAny]
-            if search in str(iid).lower().replace("-", ""):  # pyright: ignore[reportAny]
-                filtered.append(i)
-        instances = filtered
-
-    # Build table
-    rows: list[list[str]] = []
-
-    if long_format:
-        headers = ["JOBID", "NAME", "NODES", "RANKS", "STATE", "WORKDIR"]
-        for inst in instances:
-            iid_val = inst.get("instance_id", "")  # pyright: ignore[reportAny]
-            instance_id = str(iid_val) if iid_val else ""  # pyright: ignore[reportAny]
-            job_id = truncate_id(instance_id, 12)
-            name_val = inst.get("simulation_name", "")  # pyright: ignore[reportAny]
-            name = (str(name_val) if name_val else "")[:15]  # pyright: ignore[reportAny]
-            runner_statuses = cast(dict[str, Any], inst.get("runner_statuses", {}))
-            nodes = str(len(runner_statuses))
-            ranks_val = inst.get("total_ranks", 0)  # pyright: ignore[reportAny]
-            ranks = str(ranks_val) if ranks_val else "0"  # pyright: ignore[reportAny]
-            state = get_job_state(runner_statuses)
-            workdir_val = inst.get("working_directory", "")  # pyright: ignore[reportAny]
-            workdir = str(workdir_val) if workdir_val else ""  # pyright: ignore[reportAny]
-            # Truncate workdir for display
-            if len(workdir) > 30:
-                workdir = "..." + workdir[-27:]
-            rows.append([job_id, name, nodes, ranks, state, workdir])
-    else:
-        headers = ["JOBID", "NAME", "NODES", "STATE"]
-        for inst in instances:
-            iid_val = inst.get("instance_id", "")  # pyright: ignore[reportAny]
-            instance_id = str(iid_val) if iid_val else ""  # pyright: ignore[reportAny]
-            job_id = truncate_id(instance_id, 8)
-            name_val = inst.get("simulation_name", "")  # pyright: ignore[reportAny]
-            name = (str(name_val) if name_val else "")[:15]  # pyright: ignore[reportAny]
-            runner_statuses = cast(dict[str, Any], inst.get("runner_statuses", {}))
-            nodes = str(len(runner_statuses))
-            state = get_job_state(runner_statuses)
-            rows.append([job_id, name, nodes, state])
-
-    print(format_table(headers, rows))
-    return 0
-
-
-if __name__ == "__main__":
-    sys.exit(main(sys.argv[1:]))

src/exo/download/coordinator.py

@@ -0,0 +1,284 @@
+import asyncio
+from dataclasses import dataclass, field
+from typing import Iterator
+
+import anyio
+from anyio import current_time
+from anyio.abc import TaskGroup
+from loguru import logger
+
+from exo.download.download_utils import (
+    RepoDownloadProgress,
+    delete_model,
+    map_repo_download_progress_to_download_progress_data,
+)
+from exo.download.shard_downloader import ShardDownloader
+from exo.shared.models.model_cards import ModelId
+from exo.shared.types.commands import (
+    DeleteDownload,
+    ForwarderDownloadCommand,
+    StartDownload,
+)
+from exo.shared.types.common import NodeId, SessionId
+from exo.shared.types.events import (
+    Event,
+    ForwarderEvent,
+    NodeDownloadProgress,
+)
+from exo.shared.types.worker.downloads import (
+    DownloadCompleted,
+    DownloadFailed,
+    DownloadOngoing,
+    DownloadPending,
+    DownloadProgress,
+)
+from exo.shared.types.worker.shards import ShardMetadata
+from exo.utils.channels import Receiver, Sender, channel
+
+
+@dataclass
+class DownloadCoordinator:
+    node_id: NodeId
+    session_id: SessionId
+    shard_downloader: ShardDownloader
+    download_command_receiver: Receiver[ForwarderDownloadCommand]
+    local_event_sender: Sender[ForwarderEvent]
+    event_index_counter: Iterator[int]
+
+    # Local state
+    download_status: dict[ModelId, DownloadProgress] = field(default_factory=dict)
+    active_downloads: dict[ModelId, asyncio.Task[None]] = field(default_factory=dict)
+
+    # 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)
+
+    def __post_init__(self) -> None:
+        self.event_sender, self.event_receiver = channel[Event]()
+        self._tg = anyio.create_task_group()
+
+    async def run(self) -> None:
+        logger.info("Starting DownloadCoordinator")
+        async with self._tg as tg:
+            tg.start_soon(self._command_processor)
+            tg.start_soon(self._forward_events)
+            tg.start_soon(self._emit_existing_download_progress)
+
+    def shutdown(self) -> None:
+        self._tg.cancel_scope.cancel()
+
+    async def _command_processor(self) -> None:
+        with self.download_command_receiver as commands:
+            async for cmd in commands:
+                # Only process commands targeting this node
+                if cmd.command.target_node_id != self.node_id:
+                    continue
+
+                match cmd.command:
+                    case StartDownload(shard_metadata=shard):
+                        await self._start_download(shard)
+                    case DeleteDownload(model_id=model_id):
+                        await self._delete_download(model_id)
+
+    async def _start_download(self, shard: ShardMetadata) -> None:
+        model_id = shard.model_card.model_id
+
+        # Check if already downloading or complete
+        if model_id in self.download_status:
+            status = self.download_status[model_id]
+            if isinstance(status, (DownloadOngoing, DownloadCompleted)):
+                logger.debug(
+                    f"Download for {model_id} already in progress or complete, skipping"
+                )
+                return
+
+        # Emit pending status
+        progress = DownloadPending(shard_metadata=shard, node_id=self.node_id)
+        self.download_status[model_id] = progress
+        await self.event_sender.send(NodeDownloadProgress(download_progress=progress))
+
+        # Check initial status from downloader
+        initial_progress = (
+            await self.shard_downloader.get_shard_download_status_for_shard(shard)
+        )
+
+        if initial_progress.status == "complete":
+            completed = DownloadCompleted(
+                shard_metadata=shard,
+                node_id=self.node_id,
+                total_bytes=initial_progress.total_bytes,
+            )
+            self.download_status[model_id] = completed
+            await self.event_sender.send(
+                NodeDownloadProgress(download_progress=completed)
+            )
+            return
+
+        # Start actual download
+        self._start_download_task(shard, initial_progress)
+
+    def _start_download_task(
+        self, shard: ShardMetadata, initial_progress: RepoDownloadProgress
+    ) -> None:
+        model_id = shard.model_card.model_id
+
+        # Emit ongoing status
+        status = DownloadOngoing(
+            node_id=self.node_id,
+            shard_metadata=shard,
+            download_progress=map_repo_download_progress_to_download_progress_data(
+                initial_progress
+            ),
+        )
+        self.download_status[model_id] = status
+        self.event_sender.send_nowait(NodeDownloadProgress(download_progress=status))
+
+        last_progress_time = 0.0
+        throttle_interval_secs = 1.0
+
+        async def download_progress_callback(
+            callback_shard: ShardMetadata, progress: RepoDownloadProgress
+        ) -> None:
+            nonlocal last_progress_time
+
+            if progress.status == "complete":
+                completed = DownloadCompleted(
+                    shard_metadata=callback_shard,
+                    node_id=self.node_id,
+                    total_bytes=progress.total_bytes,
+                )
+                self.download_status[callback_shard.model_card.model_id] = completed
+                await self.event_sender.send(
+                    NodeDownloadProgress(download_progress=completed)
+                )
+                # Clean up active download tracking
+                if callback_shard.model_card.model_id in self.active_downloads:
+                    del self.active_downloads[callback_shard.model_card.model_id]
+            elif (
+                progress.status == "in_progress"
+                and current_time() - last_progress_time > throttle_interval_secs
+            ):
+                ongoing = DownloadOngoing(
+                    node_id=self.node_id,
+                    shard_metadata=callback_shard,
+                    download_progress=map_repo_download_progress_to_download_progress_data(
+                        progress
+                    ),
+                )
+                self.download_status[callback_shard.model_card.model_id] = ongoing
+                await self.event_sender.send(
+                    NodeDownloadProgress(download_progress=ongoing)
+                )
+                last_progress_time = current_time()
+
+        self.shard_downloader.on_progress(download_progress_callback)
+
+        async def download_wrapper() -> None:
+            try:
+                await self.shard_downloader.ensure_shard(shard)
+            except Exception as e:
+                logger.error(f"Download failed for {model_id}: {e}")
+                failed = DownloadFailed(
+                    shard_metadata=shard,
+                    node_id=self.node_id,
+                    error_message=str(e),
+                )
+                self.download_status[model_id] = failed
+                await self.event_sender.send(
+                    NodeDownloadProgress(download_progress=failed)
+                )
+            finally:
+                if model_id in self.active_downloads:
+                    del self.active_downloads[model_id]
+
+        task = asyncio.create_task(download_wrapper())
+        self.active_downloads[model_id] = task
+
+    async def _delete_download(self, model_id: ModelId) -> None:
+        # Cancel if active
+        if model_id in self.active_downloads:
+            logger.info(f"Cancelling active download for {model_id} before deletion")
+            self.active_downloads[model_id].cancel()
+            del self.active_downloads[model_id]
+
+        # Delete from disk
+        logger.info(f"Deleting model files for {model_id}")
+        deleted = await delete_model(model_id)
+
+        if deleted:
+            logger.info(f"Successfully deleted model {model_id}")
+        else:
+            logger.warning(f"Model {model_id} was not found on disk")
+
+        # Emit pending status to reset UI state, then remove from local tracking
+        if model_id in self.download_status:
+            current_status = self.download_status[model_id]
+            pending = DownloadPending(
+                shard_metadata=current_status.shard_metadata,
+                node_id=self.node_id,
+            )
+            await self.event_sender.send(
+                NodeDownloadProgress(download_progress=pending)
+            )
+            del self.download_status[model_id]
+
+    async def _forward_events(self) -> None:
+        with self.event_receiver as events:
+            async for event in events:
+                idx = next(self.event_index_counter)
+                fe = ForwarderEvent(
+                    origin_idx=idx,
+                    origin=self.node_id,
+                    session=self.session_id,
+                    event=event,
+                )
+                logger.debug(
+                    f"DownloadCoordinator published event {idx}: {str(event)[:100]}"
+                )
+                await self.local_event_sender.send(fe)
+
+    async def _emit_existing_download_progress(self) -> None:
+        try:
+            while True:
+                logger.info(
+                    "DownloadCoordinator: Fetching and emitting existing download progress..."
+                )
+                async for (
+                    _,
+                    progress,
+                ) in self.shard_downloader.get_shard_download_status():
+                    if progress.status == "complete":
+                        status: DownloadProgress = DownloadCompleted(
+                            node_id=self.node_id,
+                            shard_metadata=progress.shard,
+                            total_bytes=progress.total_bytes,
+                        )
+                    elif progress.status in ["in_progress", "not_started"]:
+                        if progress.downloaded_bytes_this_session.in_bytes == 0:
+                            status = DownloadPending(
+                                node_id=self.node_id, shard_metadata=progress.shard
+                            )
+                        else:
+                            status = DownloadOngoing(
+                                node_id=self.node_id,
+                                shard_metadata=progress.shard,
+                                download_progress=map_repo_download_progress_to_download_progress_data(
+                                    progress
+                                ),
+                            )
+                    else:
+                        continue
+
+                    self.download_status[progress.shard.model_card.model_id] = status
+                    await self.event_sender.send(
+                        NodeDownloadProgress(download_progress=status)
+                    )
+                logger.info(
+                    "DownloadCoordinator: Done emitting existing download progress."
+                )
+                await anyio.sleep(5 * 60)  # 5 minutes
+        except Exception as e:
+            logger.error(
+                f"DownloadCoordinator: Error emitting existing download progress: {e}"
+            )

src/exo/download/download_utils.py

@@ -24,7 +24,15 @@ from pydantic import (
     TypeAdapter,
 )
 
+from exo.download.huggingface_utils import (
+    filter_repo_objects,
+    get_allow_patterns,
+    get_auth_headers,
+    get_hf_endpoint,
+    get_hf_token,
+)
 from exo.shared.constants import EXO_MODELS_DIR
+from exo.shared.models.model_cards import ModelTask
 from exo.shared.types.common import ModelId
 from exo.shared.types.memory import Memory
 from exo.shared.types.worker.downloads import (
@@ -35,12 +43,27 @@ from exo.shared.types.worker.downloads import (
     RepoFileDownloadProgress,
 )
 from exo.shared.types.worker.shards import ShardMetadata
-from exo.worker.download.huggingface_utils import (
-    filter_repo_objects,
-    get_allow_patterns,
-    get_auth_headers,
-    get_hf_endpoint,
-)
+
+
+class HuggingFaceAuthenticationError(Exception):
+    """Raised when HuggingFace returns 401/403 for a model download."""
+
+
+async def _build_auth_error_message(status_code: int, model_id: ModelId) -> str:
+    token = await get_hf_token()
+    if status_code == 401 and token is None:
+        return (
+            f"Model '{model_id}' requires authentication. "
+            f"Set HF_TOKEN in the app's Advanced settings, set the HF_TOKEN environment variable, or run `hf auth login`. "
+            f"Get a token at https://huggingface.co/settings/tokens"
+        )
+    elif status_code == 403:
+        return (
+            f"Access denied to '{model_id}'. "
+            f"Please accept the model terms at https://huggingface.co/{model_id}"
+        )
+    else:
+        return f"Authentication failed for '{model_id}' (HTTP {status_code})"
 
 
 def trim_etag(etag: str) -> str:
@@ -98,11 +121,20 @@ async def ensure_models_dir() -> Path:
 
 
 async def delete_model(model_id: ModelId) -> bool:
-    model_dir = await ensure_models_dir() / model_id.normalize()
-    if not await aios.path.exists(model_dir):
-        return False
-    await asyncio.to_thread(shutil.rmtree, model_dir, ignore_errors=False)
-    return True
+    models_dir = await ensure_models_dir()
+    model_dir = models_dir / model_id.normalize()
+    cache_dir = models_dir / "caches" / model_id.normalize()
+
+    deleted = False
+    if await aios.path.exists(model_dir):
+        await asyncio.to_thread(shutil.rmtree, model_dir, ignore_errors=False)
+        deleted = True
+
+    # Also clear cache
+    if await aios.path.exists(cache_dir):
+        await asyncio.to_thread(shutil.rmtree, cache_dir, ignore_errors=False)
+
+    return deleted
 
 
 async def seed_models(seed_dir: str | Path):
@@ -128,16 +160,28 @@ async def fetch_file_list_with_cache(
     target_dir = (await ensure_models_dir()) / "caches" / model_id.normalize()
     await aios.makedirs(target_dir, exist_ok=True)
     cache_file = target_dir / f"{model_id.normalize()}--{revision}--file_list.json"
-    if await aios.path.exists(cache_file):
-        async with aiofiles.open(cache_file, "r") as f:
-            return TypeAdapter(list[FileListEntry]).validate_json(await f.read())
-    file_list = await fetch_file_list_with_retry(
-        model_id, revision, recursive=recursive
-    )
-    await aios.makedirs(cache_file.parent, exist_ok=True)
-    async with aiofiles.open(cache_file, "w") as f:
-        await f.write(TypeAdapter(list[FileListEntry]).dump_json(file_list).decode())
-    return file_list
+
+    # Always try fresh first
+    try:
+        file_list = await fetch_file_list_with_retry(
+            model_id, revision, recursive=recursive
+        )
+        # Update cache with fresh data
+        async with aiofiles.open(cache_file, "w") as f:
+            await f.write(
+                TypeAdapter(list[FileListEntry]).dump_json(file_list).decode()
+            )
+        return file_list
+    except Exception as e:
+        # Fetch failed - try cache fallback
+        if await aios.path.exists(cache_file):
+            logger.warning(
+                f"Failed to fetch file list for {model_id}, using cached data: {e}"
+            )
+            async with aiofiles.open(cache_file, "r") as f:
+                return TypeAdapter(list[FileListEntry]).validate_json(await f.read())
+        # No cache available, propagate the error
+        raise
 
 
 async def fetch_file_list_with_retry(
@@ -147,6 +191,8 @@ async def fetch_file_list_with_retry(
     for attempt in range(n_attempts):
         try:
             return await _fetch_file_list(model_id, revision, path, recursive)
+        except HuggingFaceAuthenticationError:
+            raise
         except Exception as e:
             if attempt == n_attempts - 1:
                 raise e
@@ -167,6 +213,9 @@ async def _fetch_file_list(
         create_http_session(timeout_profile="short") as session,
         session.get(url, headers=headers) as response,
     ):
+        if response.status in [401, 403]:
+            msg = await _build_auth_error_message(response.status, model_id)
+            raise HuggingFaceAuthenticationError(msg)
         if response.status == 200:
             data_json = await response.text()
             data = TypeAdapter(list[FileListEntry]).validate_json(data_json)
@@ -256,6 +305,9 @@ async def file_meta(
             # Otherwise, follow the redirect to get authoritative size/hash
             redirected_location = r.headers.get("location")
             return await file_meta(model_id, revision, path, redirected_location)
+        if r.status in [401, 403]:
+            msg = await _build_auth_error_message(r.status, model_id)
+            raise HuggingFaceAuthenticationError(msg)
         content_length = int(
             r.headers.get("x-linked-size") or r.headers.get("content-length") or 0
         )
@@ -279,6 +331,8 @@ async def download_file_with_retry(
             return await _download_file(
                 model_id, revision, path, target_dir, on_progress
             )
+        except HuggingFaceAuthenticationError:
+            raise
         except Exception as e:
             if isinstance(e, FileNotFoundError) or attempt == n_attempts - 1:
                 raise e
@@ -299,8 +353,28 @@ async def _download_file(
     target_dir: Path,
     on_progress: Callable[[int, int, bool], None] = lambda _, __, ___: None,
 ) -> Path:
-    if await aios.path.exists(target_dir / path):
-        return target_dir / path
+    target_path = target_dir / path
+
+    if await aios.path.exists(target_path):
+        local_size = (await aios.stat(target_path)).st_size
+
+        # Try to verify against remote, but allow offline operation
+        try:
+            remote_size, _ = await file_meta(model_id, revision, path)
+            if local_size != remote_size:
+                logger.info(
+                    f"File {path} size mismatch (local={local_size}, remote={remote_size}), re-downloading"
+                )
+                await aios.remove(target_path)
+            else:
+                return target_path
+        except Exception as e:
+            # Offline or network error - trust local file
+            logger.debug(
+                f"Could not verify {path} against remote (offline?): {e}, using local file"
+            )
+            return target_path
+
     await aios.makedirs((target_dir / path).parent, exist_ok=True)
     length, etag = await file_meta(model_id, revision, path)
     remote_hash = etag[:-5] if etag.endswith("-gzip") else etag
@@ -322,6 +396,9 @@ async def _download_file(
         ):
             if r.status == 404:
                 raise FileNotFoundError(f"File not found: {url}")
+            if r.status in [401, 403]:
+                msg = await _build_auth_error_message(r.status, model_id)
+                raise HuggingFaceAuthenticationError(msg)
             assert r.status in [200, 206], (
                 f"Failed to download {path} from {url}: {r.status}"
             )
@@ -446,6 +523,11 @@ async def resolve_allow_patterns(shard: ShardMetadata) -> list[str]:
         return ["*"]
 
 
+def is_image_model(shard: ShardMetadata) -> bool:
+    tasks = shard.model_card.tasks
+    return ModelTask.TextToImage in tasks or ModelTask.ImageToImage in tasks
+
+
 async def get_downloaded_size(path: Path) -> int:
     partial_path = path.with_suffix(path.suffix + ".partial")
     if await aios.path.exists(path):
@@ -463,7 +545,7 @@ async def download_shard(
     allow_patterns: list[str] | None = None,
 ) -> tuple[Path, RepoDownloadProgress]:
     if not skip_download:
-        logger.info(f"Downloading {shard.model_card.model_id=}")
+        logger.debug(f"Downloading {shard.model_card.model_id=}")
 
     revision = "main"
     target_dir = await ensure_models_dir() / str(shard.model_card.model_id).replace(
@@ -476,7 +558,7 @@ async def download_shard(
         allow_patterns = await resolve_allow_patterns(shard)
 
     if not skip_download:
-        logger.info(f"Downloading {shard.model_card.model_id=} with {allow_patterns=}")
+        logger.debug(f"Downloading {shard.model_card.model_id=} with {allow_patterns=}")
 
     all_start_time = time.time()
     file_list = await fetch_file_list_with_cache(
@@ -487,22 +569,40 @@ async def download_shard(
             file_list, allow_patterns=allow_patterns, key=lambda x: x.path
         )
     )
+
+    # For image models, skip root-level safetensors files since weights
+    # are stored in component subdirectories (e.g., transformer/, vae/)
+    if is_image_model(shard):
+        filtered_file_list = [
+            f
+            for f in filtered_file_list
+            if "/" in f.path or not f.path.endswith(".safetensors")
+        ]
     file_progress: dict[str, RepoFileDownloadProgress] = {}
 
     async def on_progress_wrapper(
         file: FileListEntry, curr_bytes: int, total_bytes: int, is_renamed: bool
     ) -> None:
-        start_time = (
-            file_progress[file.path].start_time
-            if file.path in file_progress
-            else time.time()
-        )
-        downloaded_this_session = (
-            file_progress[file.path].downloaded_this_session.in_bytes
-            + (curr_bytes - file_progress[file.path].downloaded.in_bytes)
-            if file.path in file_progress
-            else curr_bytes
+        previous_progress = file_progress.get(file.path)
+
+        # Detect re-download: curr_bytes < previous downloaded means file was deleted and restarted
+        is_redownload = (
+            previous_progress is not None
+            and curr_bytes < previous_progress.downloaded.in_bytes
         )
+
+        if is_redownload or previous_progress is None:
+            # Fresh download or re-download: reset tracking
+            start_time = time.time()
+            downloaded_this_session = curr_bytes
+        else:
+            # Continuing download: accumulate
+            start_time = previous_progress.start_time
+            downloaded_this_session = (
+                previous_progress.downloaded_this_session.in_bytes
+                + (curr_bytes - previous_progress.downloaded.in_bytes)
+            )
+
         speed = (
             downloaded_this_session / (time.time() - start_time)
             if time.time() - start_time > 0

src/exo/download/huggingface_utils.py

@@ -68,7 +68,11 @@ def get_hf_home() -> Path:
 
 
 async def get_hf_token() -> str | None:
-    """Retrieve the Hugging Face token from the user's HF_HOME directory."""
+    """Retrieve the Hugging Face token from HF_TOKEN env var or HF_HOME directory."""
+    # Check environment variable first
+    if token := os.environ.get("HF_TOKEN"):
+        return token
+    # Fall back to file-based token
     token_path = get_hf_home() / "token"
     if await aios.path.exists(token_path):
         async with aiofiles.open(token_path, "r") as f:

src/exo/download/impl_shard_downloader.py

@@ -3,13 +3,15 @@ from collections.abc import Awaitable
 from pathlib import Path
 from typing import AsyncIterator, Callable
 
+from loguru import logger
+
+from exo.download.download_utils import RepoDownloadProgress, download_shard
+from exo.download.shard_downloader import ShardDownloader
 from exo.shared.models.model_cards import MODEL_CARDS, ModelCard, ModelId
 from exo.shared.types.worker.shards import (
     PipelineShardMetadata,
     ShardMetadata,
 )
-from exo.worker.download.download_utils import RepoDownloadProgress, download_shard
-from exo.worker.download.shard_downloader import ShardDownloader
 
 
 def exo_shard_downloader(max_parallel_downloads: int = 8) -> ShardDownloader:
@@ -19,7 +21,7 @@ def exo_shard_downloader(max_parallel_downloads: int = 8) -> ShardDownloader:
 
 
 async def build_base_shard(model_id: ModelId) -> ShardMetadata:
-    model_card = await ModelCard.load(model_id)
+    model_card = await ModelCard.from_hf(model_id)
     return PipelineShardMetadata(
         model_card=model_card,
         device_rank=0,
@@ -166,7 +168,7 @@ class ResumableShardDownloader(ShardDownloader):
                 yield await task
             # TODO: except Exception
             except Exception as e:
-                print("Error downloading shard:", e)
+                logger.error("Error downloading shard:", e)
 
     async def get_shard_download_status_for_shard(
         self, shard: ShardMetadata

src/exo/download/shard_downloader.py

@@ -5,13 +5,13 @@ from datetime import timedelta
 from pathlib import Path
 from typing import AsyncIterator, Callable
 
+from exo.download.download_utils import RepoDownloadProgress
 from exo.shared.models.model_cards import ModelCard, ModelId, ModelTask
 from exo.shared.types.memory import Memory
 from exo.shared.types.worker.shards import (
     PipelineShardMetadata,
     ShardMetadata,
 )
-from exo.worker.download.download_utils import RepoDownloadProgress
 
 
 # TODO: the PipelineShardMetadata getting reinstantiated is a bit messy. Should this be a classmethod?

src/exo/download/tests/test_download_verification.py

@@ -0,0 +1,451 @@
+"""Tests for download verification and cache behavior."""
+
+import time
+from collections.abc import AsyncIterator
+from datetime import timedelta
+from pathlib import Path
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import aiofiles
+import aiofiles.os as aios
+import pytest
+from pydantic import TypeAdapter
+
+from exo.download.download_utils import (
+    delete_model,
+    fetch_file_list_with_cache,
+)
+from exo.shared.types.common import ModelId
+from exo.shared.types.memory import Memory
+from exo.shared.types.worker.downloads import FileListEntry, RepoFileDownloadProgress
+
+
+@pytest.fixture
+def model_id() -> ModelId:
+    return ModelId("test-org/test-model")
+
+
+@pytest.fixture
+async def temp_models_dir(tmp_path: Path) -> AsyncIterator[Path]:
+    """Set up a temporary models directory for testing."""
+    models_dir = tmp_path / "models"
+    await aios.makedirs(models_dir, exist_ok=True)
+    with patch("exo.download.download_utils.EXO_MODELS_DIR", models_dir):
+        yield models_dir
+
+
+class TestFileVerification:
+    """Tests for file size verification in _download_file."""
+
+    async def test_redownload_when_file_size_changes_upstream(
+        self, model_id: ModelId, tmp_path: Path
+    ) -> None:
+        """Test that files with mismatched sizes are re-downloaded."""
+        # Import inside test to allow patching
+        from exo.download.download_utils import (
+            _download_file,  # pyright: ignore[reportPrivateUsage]
+        )
+
+        target_dir = tmp_path / "downloads"
+        await aios.makedirs(target_dir, exist_ok=True)
+
+        # Create a local file with wrong size
+        local_file = target_dir / "test.safetensors"
+        async with aiofiles.open(local_file, "wb") as f:
+            await f.write(b"local content")  # 13 bytes
+
+        remote_size = 1000  # Different from local
+        remote_hash = "abc123"
+
+        with (
+            patch(
+                "exo.download.download_utils.file_meta",
+                new_callable=AsyncMock,
+                return_value=(remote_size, remote_hash),
+            ) as mock_file_meta,
+            patch(
+                "exo.download.download_utils.create_http_session"
+            ) as mock_session_factory,
+        ):
+            # Set up mock HTTP response for re-download
+            mock_response = MagicMock()
+            mock_response.status = 200
+            mock_response.content.read = AsyncMock(  # pyright: ignore[reportAny]
+                side_effect=[b"x" * remote_size, b""]
+            )
+
+            mock_session = MagicMock()
+            mock_session.get.return_value.__aenter__ = AsyncMock(  # pyright: ignore[reportAny]
+                return_value=mock_response
+            )
+            mock_session.get.return_value.__aexit__ = AsyncMock(  # pyright: ignore[reportAny]
+                return_value=None
+            )
+            mock_session_factory.return_value.__aenter__ = AsyncMock(  # pyright: ignore[reportAny]
+                return_value=mock_session
+            )
+            mock_session_factory.return_value.__aexit__ = AsyncMock(  # pyright: ignore[reportAny]
+                return_value=None
+            )
+
+            # Mock calc_hash to return the expected hash
+            with patch(
+                "exo.download.download_utils.calc_hash",
+                new_callable=AsyncMock,
+                return_value=remote_hash,
+            ):
+                await _download_file(model_id, "main", "test.safetensors", target_dir)
+
+            # file_meta should be called twice: once for verification, once for download
+            assert mock_file_meta.call_count == 2
+
+    async def test_skip_download_when_file_size_matches(
+        self, model_id: ModelId, tmp_path: Path
+    ) -> None:
+        """Test that files with matching sizes are not re-downloaded."""
+        from exo.download.download_utils import (
+            _download_file,  # pyright: ignore[reportPrivateUsage]
+        )
+
+        target_dir = tmp_path / "downloads"
+        await aios.makedirs(target_dir, exist_ok=True)
+
+        # Create a local file
+        local_file = target_dir / "test.safetensors"
+        local_content = b"local content"
+        async with aiofiles.open(local_file, "wb") as f:
+            await f.write(local_content)
+
+        remote_size = len(local_content)  # Same as local
+        remote_hash = "abc123"
+
+        with (
+            patch(
+                "exo.download.download_utils.file_meta",
+                new_callable=AsyncMock,
+                return_value=(remote_size, remote_hash),
+            ) as mock_file_meta,
+            patch(
+                "exo.download.download_utils.create_http_session"
+            ) as mock_session_factory,
+        ):
+            result = await _download_file(
+                model_id, "main", "test.safetensors", target_dir
+            )
+
+            # Should return immediately without downloading
+            assert result == local_file
+            mock_file_meta.assert_called_once()
+            mock_session_factory.assert_not_called()
+
+    async def test_offline_fallback_uses_local_file(
+        self, model_id: ModelId, tmp_path: Path
+    ) -> None:
+        """Test that local files are used when network is unavailable."""
+        from exo.download.download_utils import (
+            _download_file,  # pyright: ignore[reportPrivateUsage]
+        )
+
+        target_dir = tmp_path / "downloads"
+        await aios.makedirs(target_dir, exist_ok=True)
+
+        # Create a local file
+        local_file = target_dir / "test.safetensors"
+        async with aiofiles.open(local_file, "wb") as f:
+            await f.write(b"local content")
+
+        with (
+            patch(
+                "exo.download.download_utils.file_meta",
+                new_callable=AsyncMock,
+                side_effect=Exception("Network error"),
+            ),
+            patch(
+                "exo.download.download_utils.create_http_session"
+            ) as mock_session_factory,
+        ):
+            result = await _download_file(
+                model_id, "main", "test.safetensors", target_dir
+            )
+
+            # Should return local file without attempting download
+            assert result == local_file
+            mock_session_factory.assert_not_called()
+
+
+class TestFileListCache:
+    """Tests for file list caching behavior."""
+
+    async def test_fetch_fresh_and_update_cache(
+        self, model_id: ModelId, tmp_path: Path
+    ) -> None:
+        """Test that fresh data is fetched and cache is updated."""
+        models_dir = tmp_path / "models"
+
+        file_list = [
+            FileListEntry(type="file", path="model.safetensors", size=1000),
+            FileListEntry(type="file", path="config.json", size=100),
+        ]
+
+        with (
+            patch("exo.download.download_utils.EXO_MODELS_DIR", models_dir),
+            patch(
+                "exo.download.download_utils.fetch_file_list_with_retry",
+                new_callable=AsyncMock,
+                return_value=file_list,
+            ) as mock_fetch,
+        ):
+            result = await fetch_file_list_with_cache(model_id, "main")
+
+            assert result == file_list
+            mock_fetch.assert_called_once()
+
+            # Verify cache was written
+            cache_file = (
+                models_dir
+                / "caches"
+                / model_id.normalize()
+                / f"{model_id.normalize()}--main--file_list.json"
+            )
+            assert await aios.path.exists(cache_file)
+
+            async with aiofiles.open(cache_file, "r") as f:
+                cached_data = TypeAdapter(list[FileListEntry]).validate_json(
+                    await f.read()
+                )
+            assert cached_data == file_list
+
+    async def test_fallback_to_cache_when_fetch_fails(
+        self, model_id: ModelId, tmp_path: Path
+    ) -> None:
+        """Test that cached data is used when fetch fails."""
+        models_dir = tmp_path / "models"
+        cache_dir = models_dir / "caches" / model_id.normalize()
+        await aios.makedirs(cache_dir, exist_ok=True)
+
+        # Create cache file
+        cached_file_list = [
+            FileListEntry(type="file", path="model.safetensors", size=1000),
+        ]
+        cache_file = cache_dir / f"{model_id.normalize()}--main--file_list.json"
+        async with aiofiles.open(cache_file, "w") as f:
+            await f.write(
+                TypeAdapter(list[FileListEntry]).dump_json(cached_file_list).decode()
+            )
+
+        with (
+            patch("exo.download.download_utils.EXO_MODELS_DIR", models_dir),
+            patch(
+                "exo.download.download_utils.fetch_file_list_with_retry",
+                new_callable=AsyncMock,
+                side_effect=Exception("Network error"),
+            ),
+        ):
+            result = await fetch_file_list_with_cache(model_id, "main")
+
+            assert result == cached_file_list
+
+    async def test_error_propagates_when_no_cache(
+        self, model_id: ModelId, tmp_path: Path
+    ) -> None:
+        """Test that errors propagate when fetch fails and no cache exists."""
+        models_dir = tmp_path / "models"
+
+        with (
+            patch("exo.download.download_utils.EXO_MODELS_DIR", models_dir),
+            patch(
+                "exo.download.download_utils.fetch_file_list_with_retry",
+                new_callable=AsyncMock,
+                side_effect=Exception("Network error"),
+            ),
+            pytest.raises(Exception, match="Network error"),
+        ):
+            await fetch_file_list_with_cache(model_id, "main")
+
+
+class TestModelDeletion:
+    """Tests for model deletion including cache cleanup."""
+
+    async def test_delete_model_clears_cache(
+        self, model_id: ModelId, tmp_path: Path
+    ) -> None:
+        """Test that deleting a model also deletes its cache."""
+        models_dir = tmp_path / "models"
+        model_dir = models_dir / model_id.normalize()
+        cache_dir = models_dir / "caches" / model_id.normalize()
+
+        # Create model and cache directories
+        await aios.makedirs(model_dir, exist_ok=True)
+        await aios.makedirs(cache_dir, exist_ok=True)
+
+        # Add some files
+        async with aiofiles.open(model_dir / "model.safetensors", "w") as f:
+            await f.write("model data")
+        async with aiofiles.open(cache_dir / "file_list.json", "w") as f:
+            await f.write("[]")
+
+        with patch("exo.download.download_utils.EXO_MODELS_DIR", models_dir):
+            result = await delete_model(model_id)
+
+            assert result is True
+            assert not await aios.path.exists(model_dir)
+            assert not await aios.path.exists(cache_dir)
+
+    async def test_delete_model_only_cache_exists(
+        self, model_id: ModelId, tmp_path: Path
+    ) -> None:
+        """Test deleting when only cache exists (model already deleted)."""
+        models_dir = tmp_path / "models"
+        cache_dir = models_dir / "caches" / model_id.normalize()
+
+        # Only create cache directory
+        await aios.makedirs(cache_dir, exist_ok=True)
+        async with aiofiles.open(cache_dir / "file_list.json", "w") as f:
+            await f.write("[]")
+
+        with patch("exo.download.download_utils.EXO_MODELS_DIR", models_dir):
+            result = await delete_model(model_id)
+
+            # Returns False because model dir didn't exist
+            assert result is False
+            # But cache should still be cleaned up
+            assert not await aios.path.exists(cache_dir)
+
+    async def test_delete_nonexistent_model(
+        self, model_id: ModelId, tmp_path: Path
+    ) -> None:
+        """Test deleting a model that doesn't exist."""
+        models_dir = tmp_path / "models"
+        await aios.makedirs(models_dir, exist_ok=True)
+
+        with patch("exo.download.download_utils.EXO_MODELS_DIR", models_dir):
+            result = await delete_model(model_id)
+
+            assert result is False
+
+
+class TestProgressResetOnRedownload:
+    """Tests for progress tracking when files are re-downloaded."""
+
+    async def test_progress_resets_correctly_on_redownload(
+        self, model_id: ModelId
+    ) -> None:
+        """Test that progress tracking resets when a file is re-downloaded.
+
+        When a file is deleted and re-downloaded (due to size mismatch),
+        the progress tracking should reset rather than calculating negative
+        downloaded_this_session values.
+        """
+        # Simulate file_progress dict as it exists in download_shard
+        file_progress: dict[str, RepoFileDownloadProgress] = {}
+
+        # Initialize with old file progress (simulating existing large file)
+        old_file_size = 1_500_000_000  # 1.5 GB
+        file_progress["model.safetensors"] = RepoFileDownloadProgress(
+            repo_id=model_id,
+            repo_revision="main",
+            file_path="model.safetensors",
+            downloaded=Memory.from_bytes(old_file_size),
+            downloaded_this_session=Memory.from_bytes(0),
+            total=Memory.from_bytes(old_file_size),
+            speed=0,
+            eta=timedelta(0),
+            status="not_started",
+            start_time=time.time() - 10,  # Started 10 seconds ago
+        )
+
+        # Simulate the logic from on_progress_wrapper after re-download starts
+        # This is the exact logic from the fixed on_progress_wrapper
+        curr_bytes = 100_000  # 100 KB - new download just started
+        previous_progress = file_progress.get("model.safetensors")
+
+        # Detect re-download: curr_bytes < previous downloaded
+        is_redownload = (
+            previous_progress is not None
+            and curr_bytes < previous_progress.downloaded.in_bytes
+        )
+
+        if is_redownload or previous_progress is None:
+            # Fresh download or re-download: reset tracking
+            start_time = time.time()
+            downloaded_this_session = curr_bytes
+        else:
+            # Continuing download: accumulate
+            start_time = previous_progress.start_time
+            downloaded_this_session = (
+                previous_progress.downloaded_this_session.in_bytes
+                + (curr_bytes - previous_progress.downloaded.in_bytes)
+            )
+
+        # Key assertions
+        assert is_redownload is True, "Should detect re-download scenario"
+        assert downloaded_this_session == curr_bytes, (
+            "downloaded_this_session should equal curr_bytes on re-download"
+        )
+        assert downloaded_this_session > 0, (
+            "downloaded_this_session should be positive, not negative"
+        )
+
+        # Calculate speed (should be positive)
+        elapsed = time.time() - start_time
+        speed = downloaded_this_session / elapsed if elapsed > 0 else 0
+        assert speed >= 0, "Speed should be non-negative"
+
+    async def test_progress_accumulates_on_continuing_download(
+        self, model_id: ModelId
+    ) -> None:
+        """Test that progress accumulates correctly for continuing downloads.
+
+        When a download continues from where it left off (resume),
+        the progress should accumulate correctly.
+        """
+        file_progress: dict[str, RepoFileDownloadProgress] = {}
+
+        # Initialize with partial download progress
+        initial_downloaded = 500_000  # 500 KB already downloaded
+        start_time = time.time() - 5  # Started 5 seconds ago
+        file_progress["model.safetensors"] = RepoFileDownloadProgress(
+            repo_id=model_id,
+            repo_revision="main",
+            file_path="model.safetensors",
+            downloaded=Memory.from_bytes(initial_downloaded),
+            downloaded_this_session=Memory.from_bytes(initial_downloaded),
+            total=Memory.from_bytes(1_000_000),
+            speed=100_000,
+            eta=timedelta(seconds=5),
+            status="in_progress",
+            start_time=start_time,
+        )
+
+        # Progress callback with more bytes downloaded
+        curr_bytes = 600_000  # 600 KB - continuing download
+        previous_progress = file_progress.get("model.safetensors")
+
+        # This is NOT a re-download (curr_bytes > previous downloaded)
+        is_redownload = (
+            previous_progress is not None
+            and curr_bytes < previous_progress.downloaded.in_bytes
+        )
+
+        if is_redownload or previous_progress is None:
+            downloaded_this_session = curr_bytes
+            used_start_time = time.time()
+        else:
+            used_start_time = previous_progress.start_time
+            downloaded_this_session = (
+                previous_progress.downloaded_this_session.in_bytes
+                + (curr_bytes - previous_progress.downloaded.in_bytes)
+            )
+
+        # Key assertions
+        assert is_redownload is False, (
+            "Should NOT detect re-download for continuing download"
+        )
+        assert used_start_time == start_time, "Should preserve original start_time"
+        expected_session = initial_downloaded + (curr_bytes - initial_downloaded)
+        assert downloaded_this_session == expected_session, (
+            f"Should accumulate: {downloaded_this_session} == {expected_session}"
+        )
+        assert downloaded_this_session == 600_000, (
+            "downloaded_this_session should equal total downloaded so far"
+        )

src/exo/main.py

@@ -1,10 +1,11 @@
 import argparse
+import itertools
 import multiprocessing as mp
 import os
 import resource
 import signal
 from dataclasses import dataclass, field
-from typing import Self
+from typing import Iterator, Self
 
 import anyio
 from anyio.abc import TaskGroup
@@ -12,6 +13,8 @@ from loguru import logger
 from pydantic import PositiveInt
 
 import exo.routing.topics as topics
+from exo.download.coordinator import DownloadCoordinator
+from exo.download.impl_shard_downloader import exo_shard_downloader
 from exo.master.api import API  # TODO: should API be in master?
 from exo.master.main import Master
 from exo.routing.router import Router, get_node_id_keypair
@@ -21,7 +24,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.worker.download.impl_shard_downloader import exo_shard_downloader
 from exo.worker.main import Worker
 
 
@@ -29,6 +31,7 @@ from exo.worker.main import Worker
 @dataclass
 class Node:
     router: Router
+    download_coordinator: DownloadCoordinator | None
     worker: Worker | None
     election: Election  # Every node participates in election, as we do want a node to become master even if it isn't a master candidate if no master candidates are present.
     election_result_receiver: Receiver[ElectionResult]
@@ -36,6 +39,7 @@ class Node:
     api: API | None
 
     node_id: NodeId
+    event_index_counter: Iterator[int]
     _tg: TaskGroup = field(init=False, default_factory=anyio.create_task_group)
 
     @classmethod
@@ -49,8 +53,26 @@ class Node:
         await router.register_topic(topics.COMMANDS)
         await router.register_topic(topics.ELECTION_MESSAGES)
         await router.register_topic(topics.CONNECTION_MESSAGES)
+        await router.register_topic(topics.DOWNLOAD_COMMANDS)
 
         logger.info(f"Starting node {node_id}")
+
+        # Create shared event index counter for Worker and DownloadCoordinator
+        event_index_counter = itertools.count()
+
+        # Create DownloadCoordinator (unless --no-downloads)
+        if not args.no_downloads:
+            download_coordinator = DownloadCoordinator(
+                node_id,
+                session_id,
+                exo_shard_downloader(),
+                download_command_receiver=router.receiver(topics.DOWNLOAD_COMMANDS),
+                local_event_sender=router.sender(topics.LOCAL_EVENTS),
+                event_index_counter=event_index_counter,
+            )
+        else:
+            download_coordinator = None
+
         if args.spawn_api:
             api = API(
                 node_id,
@@ -58,6 +80,7 @@ class Node:
                 port=args.api_port,
                 global_event_receiver=router.receiver(topics.GLOBAL_EVENTS),
                 command_sender=router.sender(topics.COMMANDS),
+                download_command_sender=router.sender(topics.DOWNLOAD_COMMANDS),
                 election_receiver=router.receiver(topics.ELECTION_MESSAGES),
             )
         else:
@@ -67,11 +90,12 @@ class Node:
             worker = Worker(
                 node_id,
                 session_id,
-                exo_shard_downloader(),
                 connection_message_receiver=router.receiver(topics.CONNECTION_MESSAGES),
                 global_event_receiver=router.receiver(topics.GLOBAL_EVENTS),
                 local_event_sender=router.sender(topics.LOCAL_EVENTS),
                 command_sender=router.sender(topics.COMMANDS),
+                download_command_sender=router.sender(topics.DOWNLOAD_COMMANDS),
+                event_index_counter=event_index_counter,
             )
         else:
             worker = None
@@ -99,13 +123,25 @@ class Node:
             election_result_sender=er_send,
         )
 
-        return cls(router, worker, election, er_recv, master, api, node_id)
+        return cls(
+            router,
+            download_coordinator,
+            worker,
+            election,
+            er_recv,
+            master,
+            api,
+            node_id,
+            event_index_counter,
+        )
 
     async def run(self):
         async with self._tg as tg:
             signal.signal(signal.SIGINT, lambda _, __: self.shutdown())
             tg.start_soon(self.router.run)
             tg.start_soon(self.election.run)
+            if self.download_coordinator:
+                tg.start_soon(self.download_coordinator.run)
             if self.worker:
                 tg.start_soon(self.worker.run)
             if self.master:
@@ -170,13 +206,27 @@ class Node:
                     )
                 if result.is_new_master:
                     await anyio.sleep(0)
+                    # Fresh counter for new session (buffer expects indices from 0)
+                    self.event_index_counter = itertools.count()
+                    if self.download_coordinator:
+                        self.download_coordinator.shutdown()
+                        self.download_coordinator = DownloadCoordinator(
+                            self.node_id,
+                            result.session_id,
+                            exo_shard_downloader(),
+                            download_command_receiver=self.router.receiver(
+                                topics.DOWNLOAD_COMMANDS
+                            ),
+                            local_event_sender=self.router.sender(topics.LOCAL_EVENTS),
+                            event_index_counter=self.event_index_counter,
+                        )
+                        self._tg.start_soon(self.download_coordinator.run)
                     if self.worker:
                         self.worker.shutdown()
                         # TODO: add profiling etc to resource monitor
                         self.worker = Worker(
                             self.node_id,
                             result.session_id,
-                            exo_shard_downloader(),
                             connection_message_receiver=self.router.receiver(
                                 topics.CONNECTION_MESSAGES
                             ),
@@ -185,6 +235,10 @@ class Node:
                             ),
                             local_event_sender=self.router.sender(topics.LOCAL_EVENTS),
                             command_sender=self.router.sender(topics.COMMANDS),
+                            download_command_sender=self.router.sender(
+                                topics.DOWNLOAD_COMMANDS
+                            ),
+                            event_index_counter=self.event_index_counter,
                         )
                         self._tg.start_soon(self.worker.run)
                     if self.api:
@@ -195,14 +249,6 @@ class Node:
 
 
 def main():
-    # Check for SLURM-compatible subcommands first
-    import sys
-
-    if len(sys.argv) > 1 and sys.argv[1] in ("sbatch", "squeue", "scancel", "salloc"):
-        from exo.cli import run_subcommand
-
-        sys.exit(run_subcommand(sys.argv[1], sys.argv[2:]))
-
     args = Args.parse()
     soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE)
     resource.setrlimit(resource.RLIMIT_NOFILE, (max(soft, 65535), hard))
@@ -213,11 +259,6 @@ def main():
     logger.info("Starting EXO")
     logger.info(f"EXO_LIBP2P_NAMESPACE: {os.getenv('EXO_LIBP2P_NAMESPACE')}")
 
-    # Discover and register plugins
-    from exo.plugins.registry import discover_plugins
-
-    discover_plugins()
-
     # Set FAST_SYNCH override env var for runner subprocesses
     if args.fast_synch is True:
         os.environ["EXO_FAST_SYNCH"] = "on"
@@ -239,6 +280,7 @@ class Args(CamelCaseModel):
     api_port: PositiveInt = 52415
     tb_only: bool = False
     no_worker: bool = False
+    no_downloads: bool = False
     fast_synch: bool | None = None  # None = auto, True = force on, False = force off
 
     @classmethod
@@ -281,6 +323,11 @@ class Args(CamelCaseModel):
             "--no-worker",
             action="store_true",
         )
+        parser.add_argument(
+            "--no-downloads",
+            action="store_true",
+            help="Disable the download coordinator (node won't download models)",
+        )
         fast_synch_group = parser.add_mutually_exclusive_group()
         fast_synch_group.add_argument(
             "--fast-synch",

src/exo/master/api.py

@@ -1,16 +1,17 @@
-import asyncio
 import base64
+import contextlib
 import json
-import os
+import re
 import time
 from collections.abc import AsyncGenerator
 from http import HTTPStatus
-from typing import Any, Callable, Literal, Optional, cast
+from typing import Annotated, Any, Literal, cast
+from uuid import uuid4
 
 import anyio
 from anyio import BrokenResourceError, create_task_group
 from anyio.abc import TaskGroup
-from fastapi import FastAPI, File, Form, HTTPException, Request, UploadFile
+from fastapi import FastAPI, File, Form, HTTPException, Query, Request, UploadFile
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import FileResponse, JSONResponse, StreamingResponse
 from fastapi.staticfiles import StaticFiles
@@ -18,12 +19,14 @@ from hypercorn.asyncio import serve  # pyright: ignore[reportUnknownVariableType
 from hypercorn.config import Config
 from hypercorn.typing import ASGIFramework
 from loguru import logger
-from pydantic import BaseModel
 
 from exo.master.image_store import ImageStore
 from exo.master.placement import place_instance as get_instance_placements
 from exo.shared.apply import apply
-from exo.shared.constants import EXO_IMAGE_CACHE_DIR, EXO_MAX_CHUNK_SIZE
+from exo.shared.constants import (
+    EXO_IMAGE_CACHE_DIR,
+    EXO_MAX_CHUNK_SIZE,
+)
 from exo.shared.election import ElectionMessage
 from exo.shared.logging import InterceptLogger
 from exo.shared.models.model_cards import (
@@ -32,6 +35,7 @@ from exo.shared.models.model_cards import (
     ModelId,
 )
 from exo.shared.types.api import (
+    AdvancedImageParams,
     BenchChatCompletionResponse,
     BenchChatCompletionTaskParams,
     BenchImageGenerationResponse,
@@ -39,8 +43,10 @@ from exo.shared.types.api import (
     ChatCompletionChoice,
     ChatCompletionMessage,
     ChatCompletionResponse,
+    CompletionTokensDetails,
     CreateInstanceParams,
     CreateInstanceResponse,
+    DeleteDownloadResponse,
     DeleteInstanceResponse,
     ErrorInfo,
     ErrorResponse,
@@ -53,72 +59,97 @@ from exo.shared.types.api import (
     ImageGenerationTaskParams,
     ImageListItem,
     ImageListResponse,
+    Logprobs,
+    LogprobsContentItem,
     ModelList,
     ModelListModel,
     PlaceInstanceParams,
     PlacementPreview,
     PlacementPreviewResponse,
+    StartDownloadParams,
+    StartDownloadResponse,
     StreamingChoiceResponse,
+    ToolCall,
+    Usage,
+)
+from exo.shared.types.chunks import (
+    CompletionChunk,
+    ErrorChunk,
+    ImageChunk,
+    InputImageChunk,
+    TokenChunk,
+    ToolCallChunk,
 )
-from exo.shared.types.chunks import ImageChunk, InputImageChunk, TokenChunk
 from exo.shared.types.commands import (
-    BaseCommand,
     ChatCompletion,
+    Command,
     CreateInstance,
+    DeleteDownload,
     DeleteInstance,
+    DownloadCommand,
     ForwarderCommand,
+    ForwarderDownloadCommand,
     ImageEdits,
     ImageGeneration,
     PlaceInstance,
     SendInputChunk,
+    StartDownload,
     TaskFinished,
 )
 from exo.shared.types.common import CommandId, Id, NodeId, SessionId
 from exo.shared.types.events import (
-    BaseEvent,
     ChunkGenerated,
+    Event,
     ForwarderEvent,
     IndexedEvent,
 )
 from exo.shared.types.memory import Memory
 from exo.shared.types.state import State
 from exo.shared.types.tasks import ChatCompletionTaskParams
-from exo.shared.types.worker.instances import (
-    BaseInstance,
-    Instance,
-    InstanceId,
-    InstanceMeta,
-)
+from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
 from exo.shared.types.worker.shards import Sharding
 from exo.utils.banner import print_startup_banner
 from exo.utils.channels import Receiver, Sender, channel
 from exo.utils.dashboard_path import find_dashboard
 from exo.utils.event_buffer import OrderedBuffer
 
+_THINK_TAG_RE = re.compile(r"<think>.*?</think>", re.DOTALL)
+
+
+def _strip_think_tags(text: str) -> str:
+    """Strip <think>...</think> blocks from response text.
+
+    These tags are an artifact of GPT-OSS channel parsing, not part of the
+    model's intended output. The OpenAI API content field should not contain them.
+    """
+    return _THINK_TAG_RE.sub("", text).lstrip()
+
 
 def _format_to_content_type(image_format: Literal["png", "jpeg", "webp"] | None) -> str:
     return f"image/{image_format or 'png'}"
 
 
-class ExecuteRequest(BaseModel):
-    """Request to execute a command."""
-
-    command: list[str]
-    cwd: Optional[str] = None
-    env: Optional[dict[str, str]] = None
-
-
-class ExecuteResponse(BaseModel):
-    """Response from command execution."""
-
-    exit_code: int
-    stdout: str
-    stderr: str
+def _build_logprobs(chunk: TokenChunk) -> Logprobs:
+    """Convert flat logprob fields to OpenAI Logprobs format."""
+    return Logprobs(
+        content=[
+            LogprobsContentItem(
+                token=chunk.text,
+                logprob=chunk.logprob if chunk.logprob is not None else 0.0,
+                bytes=list(chunk.text.encode("utf-8")),
+                top_logprobs=chunk.top_logprobs or [],
+            )
+        ]
+    )
 
 
 def chunk_to_response(
-    chunk: TokenChunk, command_id: CommandId
+    chunk: TokenChunk | ToolCallChunk, command_id: CommandId
 ) -> ChatCompletionResponse:
+    logprobs: Logprobs | None = None
+    if isinstance(chunk, TokenChunk) and chunk.logprob is not None:
+        logprobs = _build_logprobs(chunk)
+
     return ChatCompletionResponse(
         id=command_id,
         created=int(time.time()),
@@ -126,7 +157,20 @@ def chunk_to_response(
         choices=[
             StreamingChoiceResponse(
                 index=0,
-                delta=ChatCompletionMessage(role="assistant", content=chunk.text),
+                delta=ChatCompletionMessage(role="assistant", content=chunk.text)
+                if isinstance(chunk, TokenChunk)
+                else ChatCompletionMessage(
+                    role="assistant",
+                    tool_calls=[
+                        ToolCall(
+                            id=str(uuid4()),
+                            index=i,
+                            function=tool,
+                        )
+                        for i, tool in enumerate(chunk.tool_calls)
+                    ],
+                ),
+                logprobs=logprobs,
                 finish_reason=chunk.finish_reason,
             )
         ],
@@ -155,15 +199,17 @@ class API:
         # Ideally this would be a MasterForwarderEvent but type system says no :(
         global_event_receiver: Receiver[ForwarderEvent],
         command_sender: Sender[ForwarderCommand],
+        download_command_sender: Sender[ForwarderDownloadCommand],
         # This lets us pause the API if an election is running
         election_receiver: Receiver[ElectionMessage],
     ) -> None:
         self.state = State()
-        self._event_log: list[BaseEvent] = []
+        self._event_log: list[Event] = []
         self.command_sender = command_sender
+        self.download_command_sender = download_command_sender
         self.global_event_receiver = global_event_receiver
         self.election_receiver = election_receiver
-        self.event_buffer: OrderedBuffer[BaseEvent] = OrderedBuffer[BaseEvent]()
+        self.event_buffer: OrderedBuffer[Event] = OrderedBuffer[Event]()
         self.node_id: NodeId = node_id
         self.session_id: SessionId = session_id
         self.last_completed_election: int = 0
@@ -186,16 +232,24 @@ class API:
             name="dashboard",
         )
 
-        self._chat_completion_queues: dict[CommandId, Sender[TokenChunk]] = {}
-        self._image_generation_queues: dict[CommandId, Sender[ImageChunk]] = {}
+        self._chat_completion_queues: dict[
+            CommandId,
+            Sender[TokenChunk | ErrorChunk | ToolCallChunk | CompletionChunk],
+        ] = {}
+        self._image_generation_queues: dict[
+            CommandId, Sender[ImageChunk | ErrorChunk]
+        ] = {}
         self._image_store = ImageStore(EXO_IMAGE_CACHE_DIR)
         self._tg: TaskGroup | None = None
 
+        # Accumulated usage stats per instance (keyed by model id)
+        self._usage_by_model: dict[str, dict[str, int]] = {}
+
     def reset(self, new_session_id: SessionId, result_clock: int):
         logger.info("Resetting API State")
         self.state = State()
         self.session_id = new_session_id
-        self.event_buffer = OrderedBuffer[BaseEvent]()
+        self.event_buffer = OrderedBuffer[Event]()
         self._chat_completion_queues = {}
         self._image_generation_queues = {}
         self.unpause(result_clock)
@@ -255,129 +309,50 @@ class API:
         self.app.get("/images/{image_id}")(self.get_image)
         self.app.get("/state")(lambda: self.state)
         self.app.get("/events")(lambda: self._event_log)
-        self.app.post("/execute")(self.execute)
+        self.app.post("/download/start")(self.start_download)
+        self.app.delete("/download/{node_id}/{model_id:path}")(self.delete_download)
+        self.app.get("/v1/usage")(self.get_usage)
 
-        # Register plugin routes
-        self._setup_plugin_routes()
+    def get_usage(self) -> dict[str, Any]:
+        """Return accumulated token usage per model instance."""
+        total_requests = 0
+        total_prompt = 0
+        total_completion = 0
+        total_reasoning = 0
+        for counters in self._usage_by_model.values():
+            total_requests += counters.get("requests", 0)
+            total_prompt += counters.get("prompt_tokens", 0)
+            total_completion += counters.get("completion_tokens", 0)
+            total_reasoning += counters.get("reasoning_tokens", 0)
+        return {
+            "total_requests": total_requests,
+            "total_prompt_tokens": total_prompt,
+            "total_completion_tokens": total_completion,
+            "total_reasoning_tokens": total_reasoning,
+            "total_tokens": total_prompt + total_completion,
+            "by_model": self._usage_by_model,
+        }
 
-    def _setup_plugin_routes(self) -> None:
-        """Register API routes from all plugins."""
-        from exo.plugins.registry import PluginRegistry
-
-        registry = PluginRegistry.get()
-
-        for plugin in registry.all_plugins():
-            for method, path, handler in plugin.get_api_routes():
-                # Create a wrapper that injects PluginContext
-                # We need to capture handler in closure properly
-                self._register_plugin_route(method, path, handler)
-
-    def _register_plugin_route(
+    def _accumulate_usage(
         self,
-        method: str,
-        path: str,
-        handler: Callable[..., Any],
+        model: str,
+        prompt_tokens: int,
+        completion_tokens: int,
+        reasoning_tokens: int,
     ) -> None:
-        """Register a single plugin route with proper closure."""
-        import functools
-        import inspect
-
-        from exo.plugins.context import PluginContext
-
-        # Get the original handler's signature (excluding ctx)
-        sig = inspect.signature(handler)
-        params = [p for p in sig.parameters.values() if p.name != "ctx"]
-        new_sig = sig.replace(parameters=params)
-
-        @functools.wraps(handler)
-        async def route_wrapper(**kwargs: Any) -> Any:  # pyright: ignore[reportAny]
-            ctx = PluginContext(
-                state=self.state,
-                send_command=self._send,
-                node_id=self.node_id,
-            )
-            return await handler(ctx, **kwargs)  # pyright: ignore[reportAny]
-
-        # Override the signature for FastAPI
-        route_wrapper.__signature__ = new_sig  # type: ignore[attr-defined]
-
-        # Register the route
-        if method == "get":
-            self.app.get(path)(route_wrapper)
-        elif method == "post":
-            self.app.post(path)(route_wrapper)
-        elif method == "delete":
-            self.app.delete(path)(route_wrapper)
-        elif method == "put":
-            self.app.put(path)(route_wrapper)
-
-        logger.info(f"Registered plugin route: {method.upper()} {path}")
-
-    async def execute(self, request: ExecuteRequest) -> ExecuteResponse:
-        """Execute a command locally. Used by exo-rsh for MPI remote execution."""
-        cmd_str = " ".join(request.command)
-        logger.info(f"Executing: {cmd_str}")
-
-        try:
-            # Build environment
-            env = os.environ.copy()
-            if request.env:
-                env.update(request.env)
-
-            # Check if command contains shell metacharacters
-            # If so, run through shell. mpirun sends complex commands like:
-            # "VAR=value;export VAR;/path/to/prted --args"
-            needs_shell = any(c in cmd_str for c in ";|&$`")
-
-            # Commands with --daemonize (e.g., prted) fork a child that inherits
-            # stdout/stderr pipe fds. Using PIPE would cause communicate() to hang
-            # because the daemon child never closes them. Use DEVNULL instead.
-            is_daemonize = "--daemonize" in cmd_str
-            out_mode = (
-                asyncio.subprocess.DEVNULL if is_daemonize else asyncio.subprocess.PIPE
-            )
-
-            if needs_shell:
-                process = await asyncio.create_subprocess_shell(
-                    cmd_str,
-                    stdout=out_mode,
-                    stderr=out_mode,
-                    cwd=request.cwd,
-                    env=env,
-                )
-            else:
-                process = await asyncio.create_subprocess_exec(
-                    *request.command,
-                    stdout=out_mode,
-                    stderr=out_mode,
-                    cwd=request.cwd,
-                    env=env,
-                )
-
-            if is_daemonize:
-                await process.wait()
-                exit_code = process.returncode or 0
-                logger.info(f"Daemonized command completed with exit code {exit_code}")
-                return ExecuteResponse(exit_code=exit_code, stdout="", stderr="")
-
-            stdout, stderr = await process.communicate()
-            exit_code = process.returncode or 0
-
-            logger.info(f"Command completed with exit code {exit_code}")
-
-            return ExecuteResponse(
-                exit_code=exit_code,
-                stdout=stdout.decode("utf-8", errors="replace"),
-                stderr=stderr.decode("utf-8", errors="replace"),
-            )
-
-        except FileNotFoundError:
-            logger.error(f"Command not found: {request.command[0]}")
-            return ExecuteResponse(
-                exit_code=127,
-                stdout="",
-                stderr=f"Command not found: {request.command[0]}",
-            )
+        """Accumulate usage stats for a model instance."""
+        if model not in self._usage_by_model:
+            self._usage_by_model[model] = {
+                "requests": 0,
+                "prompt_tokens": 0,
+                "completion_tokens": 0,
+                "reasoning_tokens": 0,
+            }
+        counters = self._usage_by_model[model]
+        counters["requests"] += 1
+        counters["prompt_tokens"] += prompt_tokens
+        counters["completion_tokens"] += completion_tokens
+        counters["reasoning_tokens"] += reasoning_tokens
 
     async def place_instance(self, payload: PlaceInstanceParams):
         command = PlaceInstance(
@@ -425,7 +400,7 @@ class API:
         sharding: Sharding = Sharding.Pipeline,
         instance_meta: InstanceMeta = InstanceMeta.MlxRing,
         min_nodes: int = 1,
-    ) -> BaseInstance:
+    ) -> Instance:
         model_card = await resolve_model_card(model_id)
 
         try:
@@ -457,10 +432,14 @@ class API:
         return placements[new_ids[0]]
 
     async def get_placement_previews(
-        self, model_id: ModelId
+        self,
+        model_id: ModelId,
+        node_ids: Annotated[list[NodeId] | None, Query()] = None,
     ) -> PlacementPreviewResponse:
         seen: set[tuple[ModelId, Sharding, InstanceMeta, int]] = set()
         previews: list[PlacementPreview] = []
+        required_nodes = set(node_ids) if node_ids else None
+
         if len(list(self.state.topology.list_nodes())) == 0:
             return PlacementPreviewResponse(previews=[])
 
@@ -496,6 +475,7 @@ class API:
                         node_network=self.state.node_network,
                         topology=self.state.topology,
                         current_instances=self.state.instances,
+                        required_nodes=required_nodes,
                     )
                 except ValueError as exc:
                     if (model_card.model_id, sharding, instance_meta, 0) not in seen:
@@ -534,14 +514,16 @@ class API:
 
                 instance = new_instances[0]
                 shard_assignments = instance.shard_assignments
-                node_ids = list(shard_assignments.node_to_runner.keys())
+                placement_node_ids = list(shard_assignments.node_to_runner.keys())
 
                 memory_delta_by_node: dict[str, int] = {}
-                if node_ids:
+                if placement_node_ids:
                     total_bytes = model_card.storage_size.in_bytes
-                    per_node = total_bytes // len(node_ids)
-                    remainder = total_bytes % len(node_ids)
-                    for index, node_id in enumerate(sorted(node_ids, key=str)):
+                    per_node = total_bytes // len(placement_node_ids)
+                    remainder = total_bytes % len(placement_node_ids)
+                    for index, node_id in enumerate(
+                        sorted(placement_node_ids, key=str)
+                    ):
                         extra = 1 if index < remainder else 0
                         memory_delta_by_node[str(node_id)] = per_node + extra
 
@@ -549,23 +531,30 @@ class API:
                     model_card.model_id,
                     sharding,
                     instance_meta,
-                    len(node_ids),
+                    len(placement_node_ids),
                 ) not in seen:
                     previews.append(
                         PlacementPreview(
                             model_id=model_card.model_id,
                             sharding=sharding,
                             instance_meta=instance_meta,
-                            instance=cast(Instance, instance),
+                            instance=instance,
                             memory_delta_by_node=memory_delta_by_node or None,
                             error=None,
                         )
                     )
-                seen.add((model_card.model_id, sharding, instance_meta, len(node_ids)))
+                seen.add(
+                    (
+                        model_card.model_id,
+                        sharding,
+                        instance_meta,
+                        len(placement_node_ids),
+                    )
+                )
 
         return PlacementPreviewResponse(previews=previews)
 
-    def get_instance(self, instance_id: InstanceId) -> BaseInstance:
+    def get_instance(self, instance_id: InstanceId) -> Instance:
         if instance_id not in self.state.instances:
             raise HTTPException(status_code=404, detail="Instance not found")
         return self.state.instances[instance_id]
@@ -585,31 +574,42 @@ class API:
         )
 
     async def _chat_chunk_stream(
-        self, command_id: CommandId
-    ) -> AsyncGenerator[TokenChunk, None]:
-        """Yield `TokenChunk`s for a given command until completion."""
+        self, command_id: CommandId, timeout: float = 60000.0
+    ) -> AsyncGenerator[TokenChunk | ErrorChunk | ToolCallChunk, None]:
+        """Yield `TokenChunk`s for a given command until completion.
+
+        Args:
+            timeout: Max seconds to wait for the next chunk before aborting.
+        """
 
         try:
-            self._chat_completion_queues[command_id], recv = channel[TokenChunk]()
+            self._chat_completion_queues[command_id], recv = channel[
+                TokenChunk | ErrorChunk | ToolCallChunk
+            ]()
 
             with recv as token_chunks:
-                async for chunk in token_chunks:
-                    yield chunk
-                    if chunk.finish_reason is not None:
-                        break
+                with anyio.fail_after(timeout):
+                    async for chunk in token_chunks:
+                        yield chunk
+                        if chunk.finish_reason is not None:
+                            break
 
         except anyio.get_cancelled_exc_class():
-            # TODO: TaskCancelled
-            """
-            self.command_sender.send_nowait(
-                ForwarderCommand(origin=self.node_id, command=command)
-            )
-            """
             raise
+        except TimeoutError:
+            logger.warning(
+                f"Chat completion timed out after {timeout}s (command_id={command_id})"
+            )
+            yield ErrorChunk(
+                model=ModelId("unknown"),
+                finish_reason="error",
+                error_message=f"Request timed out after {timeout}s",
+            )
         finally:
             command = TaskFinished(finished_command_id=command_id)
             await self._send(command)
-            del self._chat_completion_queues[command_id]
+            if command_id in self._chat_completion_queues:
+                del self._chat_completion_queues[command_id]
 
     async def _generate_chat_stream(
         self, command_id: CommandId
@@ -617,7 +617,8 @@ class API:
         """Generate chat completion stream as JSON strings."""
 
         async for chunk in self._chat_chunk_stream(command_id):
-            if chunk.finish_reason == "error":
+            assert not isinstance(chunk, ImageChunk)
+            if isinstance(chunk, ErrorChunk):
                 error_response = ErrorResponse(
                     error=ErrorInfo(
                         message=chunk.error_message or "Internal server error",
@@ -637,6 +638,15 @@ class API:
             yield f"data: {chunk_response.model_dump_json()}\n\n"
 
             if chunk.finish_reason is not None:
+                # Accumulate usage stats from the final chunk
+                if isinstance(chunk, TokenChunk) and chunk.stats is not None:
+                    s = chunk.stats
+                    self._accumulate_usage(
+                        model=chunk.model,
+                        prompt_tokens=s.prompt_tokens,
+                        completion_tokens=s.generation_tokens,
+                        reasoning_tokens=s.reasoning_tokens,
+                    )
                 yield "data: [DONE]\n\n"
 
     async def _collect_chat_completion(
@@ -645,11 +655,16 @@ class API:
         """Collect all token chunks for a chat completion and return a single response."""
 
         text_parts: list[str] = []
+        tool_calls: list[ToolCall] = []
+        logprobs_items: list[LogprobsContentItem] = []
         model: str | None = None
         finish_reason: FinishReason | None = None
+        stats: GenerationStats | None = None
 
         async for chunk in self._chat_chunk_stream(command_id):
-            if chunk.finish_reason == "error":
+            # Skip CompletionChunk - it's for the legacy completions API
+
+            if isinstance(chunk, ErrorChunk):
                 raise HTTPException(
                     status_code=500,
                     detail=chunk.error_message or "Internal server error",
@@ -658,14 +673,59 @@ class API:
             if model is None:
                 model = chunk.model
 
-            text_parts.append(chunk.text)
+            if isinstance(chunk, TokenChunk):
+                text_parts.append(chunk.text)
+                if chunk.stats is not None:
+                    stats = chunk.stats
+                if chunk.logprob is not None:
+                    lp = _build_logprobs(chunk)
+                    if lp.content:
+                        if len(lp.content) != 1:
+                            logger.warning(
+                                f"Expected 1 logprobs content item per chunk, got {len(lp.content)}"
+                            )
+                        logprobs_items.append(lp.content[0])
+
+            if isinstance(chunk, ToolCallChunk):
+                tool_calls.extend(
+                    ToolCall(
+                        id=str(uuid4()),
+                        index=i,
+                        function=tool,
+                    )
+                    for i, tool in enumerate(chunk.tool_calls)
+                )
 
             if chunk.finish_reason is not None:
                 finish_reason = chunk.finish_reason
 
-        combined_text = "".join(text_parts)
+        combined_text = _strip_think_tags("".join(text_parts))
         assert model is not None
 
+        logprobs: Logprobs | None = None
+        if logprobs_items:
+            logprobs = Logprobs(content=logprobs_items)
+
+        usage: Usage | None = None
+        if stats is not None:
+            completion_tokens = stats.generation_tokens
+            usage = Usage(
+                prompt_tokens=stats.prompt_tokens,
+                completion_tokens=completion_tokens,
+                total_tokens=stats.prompt_tokens + completion_tokens,
+                completion_tokens_details=CompletionTokensDetails(
+                    reasoning_tokens=stats.reasoning_tokens,
+                )
+                if stats.reasoning_tokens > 0
+                else None,
+            )
+            self._accumulate_usage(
+                model=model or "unknown",
+                prompt_tokens=stats.prompt_tokens,
+                completion_tokens=completion_tokens,
+                reasoning_tokens=stats.reasoning_tokens,
+            )
+
         return ChatCompletionResponse(
             id=command_id,
             created=int(time.time()),
@@ -676,23 +736,27 @@ class API:
                     message=ChatCompletionMessage(
                         role="assistant",
                         content=combined_text,
+                        tool_calls=tool_calls,
                     ),
+                    logprobs=logprobs,
                     finish_reason=finish_reason,
                 )
             ],
+            usage=usage,
         )
 
     async def _collect_chat_completion_with_stats(
         self, command_id: CommandId
     ) -> BenchChatCompletionResponse:
         text_parts: list[str] = []
+        tool_calls: list[ToolCall] = []
         model: str | None = None
         finish_reason: FinishReason | None = None
 
         stats: GenerationStats | None = None
 
         async for chunk in self._chat_chunk_stream(command_id):
-            if chunk.finish_reason == "error":
+            if isinstance(chunk, ErrorChunk):
                 raise HTTPException(
                     status_code=500,
                     detail=chunk.error_message or "Internal server error",
@@ -701,13 +765,25 @@ class API:
             if model is None:
                 model = chunk.model
 
-            text_parts.append(chunk.text)
-            stats = chunk.stats or stats
+            if isinstance(chunk, TokenChunk):
+                text_parts.append(chunk.text)
+                stats = chunk.stats or stats
+
+            if isinstance(chunk, ToolCallChunk):
+                tool_calls.extend(
+                    ToolCall(
+                        id=str(uuid4()),
+                        index=i,
+                        function=tool,
+                    )
+                    for i, tool in enumerate(chunk.tool_calls)
+                )
+                stats = chunk.stats or stats
 
             if chunk.finish_reason is not None:
                 finish_reason = chunk.finish_reason
 
-        combined_text = "".join(text_parts)
+        combined_text = _strip_think_tags("".join(text_parts))
         assert model is not None
 
         resp = BenchChatCompletionResponse(
@@ -718,7 +794,7 @@ class API:
                 ChatCompletionChoice(
                     index=0,
                     message=ChatCompletionMessage(
-                        role="assistant", content=combined_text
+                        role="assistant", content=combined_text, tool_calls=tool_calls
                     ),
                     finish_reason=finish_reason,
                 )
@@ -758,7 +834,14 @@ class API:
                 media_type="text/event-stream",
             )
 
-        return await self._collect_chat_completion(command.command_id)
+        try:
+            return await self._collect_chat_completion(command.command_id)
+        except BaseException:
+            # Ensure task cleanup if handler is cancelled before _chat_chunk_stream's finally runs
+            with contextlib.suppress(Exception):
+                await self._send(TaskFinished(finished_command_id=command.command_id))
+            self._chat_completion_queues.pop(command.command_id, None)
+            raise
 
     async def bench_chat_completions(
         self, payload: BenchChatCompletionTaskParams
@@ -876,7 +959,9 @@ class API:
         images_complete = 0
 
         try:
-            self._image_generation_queues[command_id], recv = channel[ImageChunk]()
+            self._image_generation_queues[command_id], recv = channel[
+                ImageChunk | ErrorChunk
+            ]()
 
             with recv as chunks:
                 async for chunk in chunks:
@@ -916,6 +1001,7 @@ class API:
                             # Yield partial image event (always use b64_json for partials)
                             event_data = {
                                 "type": "partial",
+                                "image_index": chunk.image_index,
                                 "partial_index": partial_idx,
                                 "total_partials": total_partials,
                                 "format": str(chunk.format),
@@ -985,7 +1071,9 @@ class API:
         stats: ImageGenerationStats | None = None
 
         try:
-            self._image_generation_queues[command_id], recv = channel[ImageChunk]()
+            self._image_generation_queues[command_id], recv = channel[
+                ImageChunk | ErrorChunk
+            ]()
 
             while images_complete < num_images:
                 with recv as chunks:
@@ -1103,6 +1191,9 @@ class API:
         stream: bool,
         partial_images: int,
         bench: bool,
+        quality: Literal["high", "medium", "low"],
+        output_format: Literal["png", "jpeg", "webp"],
+        advanced_params: AdvancedImageParams | None,
     ) -> ImageEdits:
         """Prepare and send an image edits command with chunked image upload."""
         resolved_model = await self._validate_image_model(model)
@@ -1131,6 +1222,9 @@ class API:
                 stream=stream,
                 partial_images=partial_images,
                 bench=bench,
+                quality=quality,
+                output_format=output_format,
+                advanced_params=advanced_params,
             ),
         )
 
@@ -1141,7 +1235,6 @@ class API:
             await self._send(
                 SendInputChunk(
                     chunk=InputImageChunk(
-                        idx=chunk_index,
                         model=resolved_model,
                         command_id=command.command_id,
                         data=chunk_data,
@@ -1166,12 +1259,22 @@ class API:
         input_fidelity: Literal["low", "high"] = Form("low"),
         stream: str = Form("false"),
         partial_images: str = Form("0"),
+        quality: Literal["high", "medium", "low"] = Form("medium"),
+        output_format: Literal["png", "jpeg", "webp"] = Form("png"),
+        advanced_params: str | None = Form(None),
     ) -> ImageGenerationResponse | StreamingResponse:
         """Handle image editing requests (img2img)."""
         # Parse string form values to proper types
         stream_bool = stream.lower() in ("true", "1", "yes")
         partial_images_int = int(partial_images) if partial_images.isdigit() else 0
 
+        parsed_advanced_params: AdvancedImageParams | None = None
+        if advanced_params:
+            with contextlib.suppress(Exception):
+                parsed_advanced_params = AdvancedImageParams.model_validate_json(
+                    advanced_params
+                )
+
         command = await self._send_image_edits_command(
             image=image,
             prompt=prompt,
@@ -1183,6 +1286,9 @@ class API:
             stream=stream_bool,
             partial_images=partial_images_int,
             bench=False,
+            quality=quality,
+            output_format=output_format,
+            advanced_params=parsed_advanced_params,
         )
 
         if stream_bool and partial_images_int > 0:
@@ -1213,8 +1319,18 @@ class API:
         size: str = Form("1024x1024"),
         response_format: Literal["url", "b64_json"] = Form("b64_json"),
         input_fidelity: Literal["low", "high"] = Form("low"),
+        quality: Literal["high", "medium", "low"] = Form("medium"),
+        output_format: Literal["png", "jpeg", "webp"] = Form("png"),
+        advanced_params: str | None = Form(None),
     ) -> BenchImageGenerationResponse:
         """Handle benchmark image editing requests with generation stats."""
+        parsed_advanced_params: AdvancedImageParams | None = None
+        if advanced_params:
+            with contextlib.suppress(Exception):
+                parsed_advanced_params = AdvancedImageParams.model_validate_json(
+                    advanced_params
+                )
+
         command = await self._send_image_edits_command(
             image=image,
             prompt=prompt,
@@ -1226,6 +1342,9 @@ class API:
             stream=False,
             partial_images=0,
             bench=True,
+            quality=quality,
+            output_format=output_format,
+            advanced_params=parsed_advanced_params,
         )
 
         return await self._collect_image_generation_with_stats(
@@ -1295,27 +1414,26 @@ class API:
                 for idx, event in self.event_buffer.drain_indexed():
                     self._event_log.append(event)
                     self.state = apply(self.state, IndexedEvent(event=event, idx=idx))
+
                     if isinstance(event, ChunkGenerated):
-                        if event.command_id in self._chat_completion_queues:
-                            assert isinstance(event.chunk, TokenChunk)
-                            queue = self._chat_completion_queues.get(event.command_id)
-                            if queue is not None:
-                                try:
-                                    await queue.send(event.chunk)
-                                except BrokenResourceError:
-                                    self._chat_completion_queues.pop(
-                                        event.command_id, None
-                                    )
-                        elif event.command_id in self._image_generation_queues:
+                        if queue := self._image_generation_queues.get(
+                            event.command_id, None
+                        ):
                             assert isinstance(event.chunk, ImageChunk)
-                            queue = self._image_generation_queues.get(event.command_id)
-                            if queue is not None:
-                                try:
-                                    await queue.send(event.chunk)
-                                except BrokenResourceError:
-                                    self._image_generation_queues.pop(
-                                        event.command_id, None
-                                    )
+                            try:
+                                await queue.send(event.chunk)
+                            except BrokenResourceError:
+                                self._image_generation_queues.pop(
+                                    event.command_id, None
+                                )
+                        if queue := self._chat_completion_queues.get(
+                            event.command_id, None
+                        ):
+                            assert not isinstance(event.chunk, ImageChunk)
+                            try:
+                                await queue.send(event.chunk)
+                            except BrokenResourceError:
+                                self._chat_completion_queues.pop(event.command_id, None)
 
     async def _pause_on_new_election(self):
         with self.election_receiver as ems:
@@ -1332,9 +1450,34 @@ class API:
             if removed > 0:
                 logger.debug(f"Cleaned up {removed} expired images")
 
-    async def _send(self, command: BaseCommand):
+    async def _send(self, command: Command):
         while self.paused:
             await self.paused_ev.wait()
         await self.command_sender.send(
             ForwarderCommand(origin=self.node_id, command=command)
         )
+
+    async def _send_download(self, command: DownloadCommand):
+        await self.download_command_sender.send(
+            ForwarderDownloadCommand(origin=self.node_id, command=command)
+        )
+
+    async def start_download(
+        self, payload: StartDownloadParams
+    ) -> StartDownloadResponse:
+        command = StartDownload(
+            target_node_id=payload.target_node_id,
+            shard_metadata=payload.shard_metadata,
+        )
+        await self._send_download(command)
+        return StartDownloadResponse(command_id=command.command_id)
+
+    async def delete_download(
+        self, node_id: NodeId, model_id: ModelId
+    ) -> DeleteDownloadResponse:
+        command = DeleteDownload(
+            target_node_id=node_id,
+            model_id=ModelId(model_id),
+        )
+        await self._send_download(command)
+        return DeleteDownloadResponse(command_id=command.command_id)

src/exo/master/main.py

@@ -10,10 +10,10 @@ from exo.master.placement import (
     get_transition_events,
     place_instance,
 )
-from exo.plugins.registry import PluginRegistry
 from exo.shared.apply import apply
 from exo.shared.types.commands import (
     ChatCompletion,
+    Completion,
     CreateInstance,
     DeleteInstance,
     ForwarderCommand,
@@ -27,7 +27,6 @@ from exo.shared.types.commands import (
 )
 from exo.shared.types.common import CommandId, NodeId, SessionId
 from exo.shared.types.events import (
-    BaseEvent,
     Event,
     ForwarderEvent,
     IndexedEvent,
@@ -42,6 +41,9 @@ from exo.shared.types.state import State
 from exo.shared.types.tasks import (
     ChatCompletion as ChatCompletionTask,
 )
+from exo.shared.types.tasks import (
+    Completion as CompletionTask,
+)
 from exo.shared.types.tasks import (
     ImageEdits as ImageEditsTask,
 )
@@ -85,9 +87,9 @@ class Master:
         self._loopback_event_sender: Sender[ForwarderEvent] = (
             local_event_receiver.clone_sender()
         )
-        self._multi_buffer = MultiSourceBuffer[NodeId, BaseEvent]()
+        self._multi_buffer = MultiSourceBuffer[NodeId, Event]()
         # TODO: not have this
-        self._event_log: list[BaseEvent] = []
+        self._event_log: list[Event] = []
 
     async def run(self):
         logger.info("Starting Master")
@@ -160,6 +162,48 @@ class Master:
                                 )
                             )
 
+                            self.command_task_mapping[command.command_id] = task_id
+                        case Completion():
+                            for instance in self.state.instances.values():
+                                if (
+                                    instance.shard_assignments.model_id
+                                    == command.request_params.model
+                                ):
+                                    task_count = sum(
+                                        1
+                                        for task in self.state.tasks.values()
+                                        if task.instance_id == instance.instance_id
+                                    )
+                                    instance_task_counts[instance.instance_id] = (
+                                        task_count
+                                    )
+
+                            if not instance_task_counts:
+                                raise ValueError(
+                                    f"No instance found for model {command.request_params.model}"
+                                )
+
+                            available_instance_ids = sorted(
+                                instance_task_counts.keys(),
+                                key=lambda instance_id: instance_task_counts[
+                                    instance_id
+                                ],
+                            )
+
+                            task_id = TaskId()
+                            generated_events.append(
+                                TaskCreated(
+                                    task_id=task_id,
+                                    task=CompletionTask(
+                                        task_id=task_id,
+                                        command_id=command.command_id,
+                                        instance_id=available_instance_ids[0],
+                                        task_status=TaskStatus.Pending,
+                                        task_params=command.request_params,
+                                    ),
+                                )
+                            )
+
                             self.command_task_mapping[command.command_id] = task_id
                         case ImageGeneration():
                             for instance in self.state.instances.values():
@@ -281,34 +325,21 @@ class Master:
                                 )
                             )
                         case TaskFinished():
-                            generated_events.append(
-                                TaskDeleted(
-                                    task_id=self.command_task_mapping[
-                                        command.finished_command_id
-                                    ]
-                                )
+                            task_id = self.command_task_mapping.pop(
+                                command.finished_command_id, None
                             )
-                            if command.finished_command_id in self.command_task_mapping:
-                                del self.command_task_mapping[
-                                    command.finished_command_id
-                                ]
+                            if task_id is not None:
+                                generated_events.append(TaskDeleted(task_id=task_id))
+                            else:
+                                logger.debug(
+                                    f"TaskFinished for unknown command_id={command.finished_command_id} (already cleaned up)"
+                                )
                         case RequestEventLog():
                             # We should just be able to send everything, since other buffers will ignore old messages
                             for i in range(command.since_idx, len(self._event_log)):
                                 await self._send_event(
                                     IndexedEvent(idx=i, event=self._event_log[i])
                                 )
-                        case _:
-                            # Check if a plugin handles this command
-                            registry = PluginRegistry.get()
-                            plugin = registry.get_plugin_for_command(command)
-                            if plugin is not None:
-                                events = plugin.process_command(
-                                    command,
-                                    self.state.topology,
-                                    self.state.instances,
-                                )
-                                generated_events.extend(events)
                     for event in generated_events:
                         await self.event_sender.send(event)
                 except ValueError as e:

src/exo/master/placement.py

@@ -24,7 +24,7 @@ from exo.shared.types.events import Event, InstanceCreated, InstanceDeleted
 from exo.shared.types.memory import Memory
 from exo.shared.types.profiling import MemoryUsage, NodeNetworkInfo
 from exo.shared.types.worker.instances import (
-    BaseInstance,
+    Instance,
     InstanceId,
     InstanceMeta,
     MlxJacclInstance,
@@ -35,14 +35,14 @@ from exo.shared.types.worker.shards import Sharding
 
 def random_ephemeral_port() -> int:
     port = random.randint(49153, 65535)
-    return port - 1 if port <= 52415 else 52414
+    return port - 1 if port <= 52415 else port
 
 
 def add_instance_to_placements(
     command: CreateInstance,
     topology: Topology,
-    current_instances: Mapping[InstanceId, BaseInstance],
-) -> Mapping[InstanceId, BaseInstance]:
+    current_instances: Mapping[InstanceId, Instance],
+) -> Mapping[InstanceId, Instance]:
     # TODO: validate against topology
 
     return {**current_instances, command.instance.instance_id: command.instance}
@@ -51,12 +51,21 @@ def add_instance_to_placements(
 def place_instance(
     command: PlaceInstance,
     topology: Topology,
-    current_instances: Mapping[InstanceId, BaseInstance],
+    current_instances: Mapping[InstanceId, Instance],
     node_memory: Mapping[NodeId, MemoryUsage],
     node_network: Mapping[NodeId, NodeNetworkInfo],
-) -> dict[InstanceId, BaseInstance]:
+    required_nodes: set[NodeId] | None = None,
+) -> dict[InstanceId, Instance]:
     cycles = topology.get_cycles()
     candidate_cycles = list(filter(lambda it: len(it) >= command.min_nodes, cycles))
+
+    # Filter to cycles containing all required nodes (subset matching)
+    if required_nodes:
+        candidate_cycles = [
+            cycle
+            for cycle in candidate_cycles
+            if required_nodes.issubset(cycle.node_ids)
+        ]
     cycles_with_sufficient_memory = filter_cycles_by_memory(
         candidate_cycles, node_memory, command.model_card.storage_size
     )
@@ -87,12 +96,12 @@ def place_instance(
 
     smallest_cycles = get_smallest_cycles(cycles_with_sufficient_memory)
 
-    smallest_tb_cycles = [
-        cycle for cycle in smallest_cycles if topology.is_thunderbolt_cycle(cycle)
+    smallest_rdma_cycles = [
+        cycle for cycle in smallest_cycles if topology.is_rdma_cycle(cycle)
     ]
 
-    if smallest_tb_cycles != []:
-        smallest_cycles = smallest_tb_cycles
+    if command.instance_meta == InstanceMeta.MlxJaccl and smallest_rdma_cycles != []:
+        smallest_cycles = smallest_rdma_cycles
 
     cycles_with_leaf_nodes: list[Cycle] = [
         cycle
@@ -159,8 +168,8 @@ def place_instance(
 
 def delete_instance(
     command: DeleteInstance,
-    current_instances: Mapping[InstanceId, BaseInstance],
-) -> dict[InstanceId, BaseInstance]:
+    current_instances: Mapping[InstanceId, Instance],
+) -> dict[InstanceId, Instance]:
     target_instances = dict(deepcopy(current_instances))
     if command.instance_id in target_instances:
         del target_instances[command.instance_id]
@@ -169,8 +178,8 @@ def delete_instance(
 
 
 def get_transition_events(
-    current_instances: Mapping[InstanceId, BaseInstance],
-    target_instances: Mapping[InstanceId, BaseInstance],
+    current_instances: Mapping[InstanceId, Instance],
+    target_instances: Mapping[InstanceId, Instance],
 ) -> Sequence[Event]:
     events: list[Event] = []
 

src/exo/master/placement_utils.py

@@ -197,49 +197,6 @@ def get_shard_assignments(
             )
 
 
-def get_hosts_from_subgraph(cycle_digraph: Topology) -> list[Host]:
-    cycles = cycle_digraph.get_cycles()
-    expected_length = len(list(cycle_digraph.list_nodes()))
-    cycles = [cycle for cycle in cycles if len(cycle) == expected_length]
-    if not cycles:
-        if expected_length > 1:
-            logger.warning(
-                f"No cycles of length {expected_length} found even though chosen subgraph contained {expected_length} nodes"
-            )
-        return []
-
-    cycle = cycles[0]
-
-    get_thunderbolt = False
-    if cycle_digraph.is_thunderbolt_cycle(cycle):
-        get_thunderbolt = True
-
-    logger.debug(f"Using thunderbolt cycle: {get_thunderbolt}")
-
-    hosts: list[Host] = []
-    for i in range(len(cycle)):
-        current_node = cycle.node_ids[i]
-        next_node = cycle.node_ids[(i + 1) % len(cycle)]
-
-        for connection in cycle_digraph.get_all_connections_between(
-            source=current_node, sink=next_node
-        ):
-            if not isinstance(connection, SocketConnection):
-                continue
-
-            if get_thunderbolt and not connection.is_thunderbolt():
-                continue
-
-            host = Host(
-                ip=connection.sink_multiaddr.ip_address,
-                port=connection.sink_multiaddr.port,
-            )
-            hosts.append(host)
-            break
-
-    return hosts
-
-
 def get_mlx_jaccl_devices_matrix(
     selected_cycle: list[NodeId],
     cycle_digraph: Topology,
@@ -265,9 +222,6 @@ def get_mlx_jaccl_devices_matrix(
                     matrix[i][j] = conn.source_rdma_iface
                     break
             else:
-                logger.warning(
-                    f"Failed to find interface name between {node_i} and {node_j}"
-                )
                 raise ValueError(
                     "Current jaccl backend requires all-to-all RDMA connections"
                 )
@@ -279,22 +233,11 @@ def _find_connection_ip(
     node_i: NodeId,
     node_j: NodeId,
     cycle_digraph: Topology,
-) -> Generator[tuple[str, bool]]:
+) -> Generator[str, None, None]:
     """Find all IP addresses that connect node i to node j."""
     for connection in cycle_digraph.get_all_connections_between(node_i, node_j):
         if isinstance(connection, SocketConnection):
-            yield connection.sink_multiaddr.ip_address, connection.is_thunderbolt()
-
-
-def _find_interface_name_for_ip(
-    ip_address: str, node_network: NodeNetworkInfo
-) -> str | None:
-    """Find the interface name for an IP address on a node (any interface)."""
-    for interface in node_network.interfaces:
-        if interface.ip_address == ip_address:
-            return interface.name
-
-    return None
+            yield connection.sink_multiaddr.ip_address
 
 
 def _find_ip_prioritised(
@@ -303,43 +246,25 @@ def _find_ip_prioritised(
     cycle_digraph: Topology,
     node_network: Mapping[NodeId, NodeNetworkInfo],
 ) -> str | None:
-    # TODO: Actually prioritize in the correct Ethernet > Wifi > Non-TB > TB order.
     """Find an IP address between nodes with prioritization.
 
-    Priority order:
-    1. en0 (Ethernet on Mac Studio, WiFi on MacBook)
-    2. en1 (WiFi on Mac Studio, Ethernet on MacBook)
-    3. Non-Thunderbolt connections
-    4. Any other IP address
+    Priority: ethernet > wifi > unknown > thunderbolt
     """
     ips = list(_find_connection_ip(node_id, other_node_id, cycle_digraph))
-    # We expect a unique iface -> ip mapping
-    iface_map = {
-        _find_interface_name_for_ip(
-            ip, node_network.get(other_node_id, NodeNetworkInfo())
-        ): ip
-        for ip, _ in ips
+    if not ips:
+        return None
+    other_network = node_network.get(other_node_id, NodeNetworkInfo())
+    ip_to_type = {
+        iface.ip_address: iface.interface_type for iface in other_network.interfaces
     }
-
-    en0_ip = iface_map.get("en0")
-    if en0_ip:
-        return en0_ip
-
-    en1_ip = iface_map.get("en1")
-    if en1_ip:
-        return en1_ip
-
-    non_thunderbolt_ip = next(
-        (ip for (ip, is_thunderbolt) in ips if not is_thunderbolt), None
-    )
-
-    if non_thunderbolt_ip:
-        return non_thunderbolt_ip
-
-    if ips:
-        return ips[0][0]
-
-    return None
+    priority = {
+        "ethernet": 0,
+        "wifi": 1,
+        "unknown": 2,
+        "maybe_ethernet": 3,
+        "thunderbolt": 4,
+    }
+    return min(ips, key=lambda ip: priority.get(ip_to_type.get(ip, "unknown"), 2))
 
 
 def get_mlx_ring_hosts_by_node(
@@ -381,9 +306,6 @@ def get_mlx_ring_hosts_by_node(
                 node_id, other_node_id, cycle_digraph, node_network
             )
             if connection_ip is None:
-                logger.warning(
-                    f"Failed to find prioritised connection IP between {node_id} and {other_node_id}"
-                )
                 raise ValueError(
                     "MLX ring backend requires connectivity between neighbouring nodes"
                 )
@@ -416,9 +338,6 @@ def get_mlx_jaccl_coordinators(
         if ip is not None:
             return ip
 
-        logger.warning(
-            f"Failed to find directly connected ip between {n} and {coordinator}"
-        )
         raise ValueError(
             "Current jaccl backend requires all participating devices to be able to communicate"
         )

src/exo/master/tests/test_api_error_handling.py

@@ -1,13 +1,9 @@
 # pyright: reportUnusedFunction=false, reportAny=false
-from typing import Any, get_args
+from typing import Any
 
 from fastapi import FastAPI, HTTPException
 from fastapi.testclient import TestClient
 
-from exo.shared.types.api import ErrorInfo, ErrorResponse, FinishReason
-from exo.shared.types.chunks import ImageChunk, TokenChunk
-from exo.worker.tests.constants import MODEL_A_ID
-
 
 def test_http_exception_handler_formats_openai_style() -> None:
     """Test that HTTPException is converted to OpenAI-style error format."""
@@ -48,95 +44,3 @@ def test_http_exception_handler_formats_openai_style() -> None:
     assert data["error"]["message"] == "Resource not found"
     assert data["error"]["type"] == "Not Found"
     assert data["error"]["code"] == 404
-
-
-def test_finish_reason_includes_error() -> None:
-    valid_reasons = get_args(FinishReason)
-    assert "error" in valid_reasons
-
-
-def test_token_chunk_with_error_fields() -> None:
-    chunk = TokenChunk(
-        idx=0,
-        model=MODEL_A_ID,
-        text="",
-        token_id=0,
-        finish_reason="error",
-        error_message="Something went wrong",
-    )
-
-    assert chunk.finish_reason == "error"
-    assert chunk.error_message == "Something went wrong"
-
-
-def test_token_chunk_without_error() -> None:
-    chunk = TokenChunk(
-        idx=1,
-        model=MODEL_A_ID,
-        text="Hello",
-        token_id=42,
-        finish_reason=None,
-    )
-
-    assert chunk.finish_reason is None
-    assert chunk.error_message is None
-
-
-def test_error_response_construction() -> None:
-    error_response = ErrorResponse(
-        error=ErrorInfo(
-            message="Generation failed",
-            type="InternalServerError",
-            code=500,
-        )
-    )
-
-    assert error_response.error.message == "Generation failed"
-    assert error_response.error.code == 500
-
-
-def test_normal_finish_reasons_still_work() -> None:
-    for reason in ["stop", "length", "tool_calls", "content_filter", "function_call"]:
-        chunk = TokenChunk(
-            idx=0,
-            model=MODEL_A_ID,
-            text="done",
-            token_id=100,
-            finish_reason=reason,  # type: ignore[arg-type]
-        )
-        assert chunk.finish_reason == reason
-
-
-def test_image_chunk_with_error_fields() -> None:
-    chunk = ImageChunk(
-        idx=0,
-        model=MODEL_A_ID,
-        data="",
-        chunk_index=0,
-        total_chunks=1,
-        image_index=0,
-        finish_reason="error",
-        error_message="Image generation failed",
-    )
-
-    assert chunk.finish_reason == "error"
-    assert chunk.error_message == "Image generation failed"
-    assert chunk.data == ""
-    assert chunk.chunk_index == 0
-    assert chunk.total_chunks == 1
-    assert chunk.image_index == 0
-
-
-def test_image_chunk_without_error() -> None:
-    chunk = ImageChunk(
-        idx=0,
-        model=MODEL_A_ID,
-        data="base64encodeddata",
-        chunk_index=0,
-        total_chunks=1,
-        image_index=0,
-    )
-
-    assert chunk.finish_reason is None
-    assert chunk.error_message is None
-    assert chunk.data == "base64encodeddata"

src/exo/master/tests/test_placement_utils.py

@@ -3,7 +3,6 @@ import pytest
 from exo.master.placement_utils import (
     allocate_layers_proportionally,
     filter_cycles_by_memory,
-    get_hosts_from_subgraph,
     get_mlx_jaccl_coordinators,
     get_shard_assignments,
     get_smallest_cycles,
@@ -14,7 +13,7 @@ from exo.master.tests.conftest import (
 )
 from exo.shared.models.model_cards import ModelCard, ModelId, ModelTask
 from exo.shared.topology import Topology
-from exo.shared.types.common import Host, NodeId
+from exo.shared.types.common import NodeId
 from exo.shared.types.memory import Memory
 from exo.shared.types.profiling import (
     NetworkInterfaceInfo,
@@ -273,45 +272,6 @@ def test_get_shard_assignments(
     )
 
 
-def test_get_hosts_from_subgraph():
-    # arrange
-    node_a_id = NodeId()
-    node_b_id = NodeId()
-    node_c_id = NodeId()
-    topology = Topology()
-
-    topology.add_node(node_a_id)
-    topology.add_node(node_b_id)
-    topology.add_node(node_c_id)
-
-    connection1 = Connection(
-        source=node_a_id, sink=node_b_id, edge=create_socket_connection(1)
-    )
-    connection2 = Connection(
-        source=node_b_id, sink=node_c_id, edge=create_socket_connection(2)
-    )
-    connection3 = Connection(
-        source=node_c_id, sink=node_a_id, edge=create_socket_connection(3)
-    )
-
-    topology.add_connection(connection1)
-    topology.add_connection(connection2)
-    topology.add_connection(connection3)
-
-    # act
-    hosts = get_hosts_from_subgraph(topology)
-
-    # assert
-    assert len(hosts) == 3
-    expected_hosts = [
-        Host(ip="169.254.0.1", port=1234),
-        Host(ip="169.254.0.2", port=1234),
-        Host(ip="169.254.0.3", port=1234),
-    ]
-    for expected_host in expected_hosts:
-        assert expected_host in hosts
-
-
 def test_get_mlx_jaccl_coordinators():
     # arrange
     node_a_id = NodeId()

src/exo/plugins/__init__.py

@@ -1,24 +0,0 @@
-"""Exo Plugin System.
-
-This module provides the plugin architecture for extending exo with custom
-workload types (simulations, ML frameworks, etc.) without modifying core code.
-"""
-
-from exo.plugins.base import EXOPlugin, PluginCommand, PluginInstance
-from exo.plugins.registry import PluginRegistry, discover_plugins
-from exo.plugins.type_registry import (
-    command_registry,
-    event_registry,
-    instance_registry,
-)
-
-__all__ = [
-    "EXOPlugin",
-    "PluginCommand",
-    "PluginInstance",
-    "PluginRegistry",
-    "discover_plugins",
-    "command_registry",
-    "event_registry",
-    "instance_registry",
-]

src/exo/plugins/base.py

@@ -1,171 +0,0 @@
-"""Base classes and protocols for Exo plugins."""
-
-from abc import ABC, abstractmethod
-from collections.abc import Callable, Mapping, Sequence
-from typing import TYPE_CHECKING, Any
-
-from pydantic import Field
-
-from exo.shared.types.common import CommandId
-from exo.shared.types.events import Event
-from exo.shared.types.tasks import Task
-from exo.shared.types.worker.instances import InstanceId
-from exo.shared.types.worker.runners import RunnerId
-from exo.utils.pydantic_ext import TaggedModel
-
-if TYPE_CHECKING:
-    from exo.shared.topology import Topology
-    from exo.shared.types.worker.instances import BaseInstance, BoundInstance
-    from exo.utils.channels import MpReceiver, MpSender
-    from exo.worker.runner.runner_supervisor import RunnerSupervisor
-
-
-class PluginCommand(TaggedModel):
-    """Base class for plugin-defined commands.
-
-    All plugin commands must inherit from this class. Commands are serialized
-    with their class name as a tag for routing.
-    """
-
-    command_id: CommandId = Field(default_factory=CommandId)
-
-
-class PluginInstance(TaggedModel):
-    """Base class for plugin-defined instances.
-
-    All plugin instances must inherit from this class. Plugins are expected
-    to define their own instance type with workload-specific fields.
-    """
-
-    instance_id: InstanceId
-
-
-class EXOPlugin(ABC):
-    """Protocol that all exo plugins must implement.
-
-    A plugin provides:
-    - Custom command types for API -> Master communication
-    - Custom instance types representing running workloads
-    - Placement logic for distributing work across nodes
-    - Planning logic for local task scheduling
-    - Runner implementation for executing work
-    """
-
-    @property
-    @abstractmethod
-    def name(self) -> str:
-        """Unique identifier for this plugin (e.g., 'flash', 'pytorch', 'mpi')."""
-        ...
-
-    @property
-    @abstractmethod
-    def version(self) -> str:
-        """Semantic version string (e.g., '1.0.0')."""
-        ...
-
-    # ========== Type Registration ==========
-
-    @abstractmethod
-    def get_command_types(self) -> Sequence[type]:
-        """Return command types this plugin handles.
-
-        These commands are routed to this plugin's process_command method.
-        Can return core BaseCommand types or PluginCommand types.
-        """
-        ...
-
-    @abstractmethod
-    def get_instance_type(self) -> type:
-        """Return the instance type this plugin creates.
-
-        This instance type is used for routing in planning and runner bootstrap.
-        Can return core Instance types or PluginInstance types.
-        """
-        ...
-
-    # ========== API Routes ==========
-
-    @abstractmethod
-    def get_api_routes(
-        self,
-    ) -> Sequence[tuple[str, str, Callable[..., Any]]]:
-        """Return FastAPI routes to register.
-
-        Each tuple: (method, path, handler)
-        Example: [('post', '/flash/launch', self.launch_handler)]
-
-        Handlers receive a PluginContext with access to:
-        - state: Current State object
-        - send_command: Async function to send commands
-        - node_id: Current node's ID
-        """
-        ...
-
-    # ========== Master Command Handling ==========
-
-    @abstractmethod
-    def handles_command(self, command: Any) -> bool:  # pyright: ignore[reportAny]
-        """Return True if this plugin handles the given command type."""
-        ...
-
-    @abstractmethod
-    def process_command(
-        self,
-        command: Any,  # pyright: ignore[reportAny]
-        topology: "Topology",
-        current_instances: Mapping[InstanceId, "BaseInstance"],
-    ) -> Sequence[Event]:
-        """Process a command and return events to emit.
-
-        Typically creates placement and returns InstanceCreated/InstanceDeleted events.
-
-        Args:
-            command: The command to process
-            topology: Current cluster topology
-            current_instances: Currently running instances
-
-        Returns:
-            Sequence of events to emit (e.g., InstanceCreated, InstanceDeleted)
-        """
-        ...
-
-    # ========== Worker Planning ==========
-
-    @abstractmethod
-    def handles_instance(self, instance: object) -> bool:
-        """Return True if this plugin manages the given instance type."""
-        ...
-
-    @abstractmethod
-    def plan_task(
-        self,
-        runners: Mapping[RunnerId, "RunnerSupervisor"],
-        instances: Mapping[InstanceId, "BaseInstance"],
-    ) -> Task | None:
-        """Plan the next task for plugin instances.
-
-        Called during each planning cycle.
-        Return None if no task is needed.
-        """
-        ...
-
-    @abstractmethod
-    def should_skip_download(self, instance: object) -> bool:
-        """Return True if this instance type doesn't need model downloads."""
-        ...
-
-    # ========== Runner Bootstrap ==========
-
-    @abstractmethod
-    def create_runner(
-        self,
-        bound_instance: "BoundInstance",
-        event_sender: "MpSender[Event]",
-        task_receiver: "MpReceiver[Task]",
-    ) -> None:
-        """Entry point for the runner process.
-
-        Called in a subprocess to execute the actual workload.
-        This function should block until the workload completes.
-        """
-        ...

src/exo/plugins/context.py

@@ -1,21 +0,0 @@
-"""Context objects passed to plugin handlers."""
-
-from collections.abc import Awaitable, Callable
-from dataclasses import dataclass
-
-from exo.shared.types.commands import BaseCommand
-from exo.shared.types.common import NodeId
-from exo.shared.types.state import State
-
-
-@dataclass
-class PluginContext:
-    """Context provided to plugin API handlers.
-
-    This gives plugins access to the current state and the ability to send
-    commands without direct access to internal API components.
-    """
-
-    state: State
-    send_command: Callable[[BaseCommand], Awaitable[None]]
-    node_id: NodeId

src/exo/plugins/implementations/__init__.py

@@ -1,5 +0,0 @@
-"""Plugin implementations directory.
-
-Each subdirectory should contain a plugin with a register() function
-that returns an EXOPlugin instance.
-"""

src/exo/plugins/implementations/flash/__init__.py

@@ -1,15 +0,0 @@
-"""FLASH Plugin - MPI-based simulation support for Exo."""
-
-from exo.plugins.implementations.flash.plugin import FLASHPlugin
-from exo.plugins.implementations.flash.types import (
-    FLASHInstance,
-    LaunchFLASH,
-    StopFLASH,
-)
-
-__all__ = ["FLASHPlugin", "FLASHInstance", "LaunchFLASH", "StopFLASH", "register"]
-
-
-def register() -> FLASHPlugin:
-    """Entry point for plugin discovery."""
-    return FLASHPlugin()

src/exo/plugins/implementations/flash/api_handlers.py

@@ -1,109 +0,0 @@
-"""FLASH plugin API handlers."""
-
-from typing import Any
-
-from fastapi import HTTPException
-
-from exo.plugins.context import PluginContext
-from exo.plugins.implementations.flash.types import (
-    FLASHInstance,
-    LaunchFLASH,
-    StopFLASH,
-)
-
-
-async def handle_launch_flash(
-    ctx: PluginContext,
-    simulation_name: str,
-    flash_executable_path: str,
-    working_directory: str,
-    parameter_file_path: str = "",
-    ranks_per_node: int = 1,
-    min_nodes: int = 1,
-    hosts: str = "",
-) -> dict[str, str]:
-    """Launch a FLASH MPI simulation across the cluster.
-
-    Args:
-        ctx: Plugin context with state and send_command
-        simulation_name: Name of the simulation
-        flash_executable_path: Path to the FLASH executable
-        working_directory: Working directory for the simulation
-        parameter_file_path: Path to parameter file (optional)
-        ranks_per_node: Number of MPI ranks per node
-        min_nodes: Minimum number of nodes required
-        hosts: Optional comma-separated hostnames (e.g., "s14,james21-1").
-               If not provided, IPs are discovered from topology edges.
-    """
-    command = LaunchFLASH(
-        simulation_name=simulation_name,
-        flash_executable_path=flash_executable_path,
-        parameter_file_path=parameter_file_path,
-        working_directory=working_directory,
-        ranks_per_node=ranks_per_node,
-        min_nodes=min_nodes,
-        hosts=hosts,
-    )
-    await ctx.send_command(command)
-
-    return {
-        "message": "FLASH launch command received",
-        "command_id": str(command.command_id),
-        "simulation_name": simulation_name,
-    }
-
-
-async def handle_stop_flash(
-    ctx: PluginContext,
-    instance_id: str,
-) -> dict[str, str]:
-    """Stop a running FLASH simulation."""
-    from exo.shared.types.worker.instances import InstanceId
-
-    inst_id = InstanceId(instance_id)
-
-    if inst_id not in ctx.state.instances:
-        raise HTTPException(status_code=404, detail="Instance not found")
-
-    instance = ctx.state.instances[inst_id]
-    if not isinstance(instance, FLASHInstance):
-        raise HTTPException(
-            status_code=400, detail="Instance is not a FLASH simulation"
-        )
-
-    command = StopFLASH(instance_id=inst_id)
-    await ctx.send_command(command)
-
-    return {
-        "message": "Stop command received",
-        "command_id": str(command.command_id),
-        "instance_id": str(instance_id),
-    }
-
-
-async def handle_list_flash_instances(ctx: PluginContext) -> list[dict[str, Any]]:
-    """List all FLASH simulation instances."""
-    flash_instances: list[dict[str, Any]] = []
-    for instance_id, instance in ctx.state.instances.items():
-        if isinstance(instance, FLASHInstance):
-            # Get runner statuses for this instance
-            runner_statuses: dict[str, str | None] = {}
-            for (
-                node_id,
-                runner_id,
-            ) in instance.shard_assignments.node_to_runner.items():
-                runner_status = ctx.state.runners.get(runner_id)
-                runner_statuses[str(node_id)] = (
-                    str(runner_status) if runner_status else None
-                )
-
-            flash_instances.append(
-                {
-                    "instance_id": str(instance_id),
-                    "simulation_name": instance.simulation_name,
-                    "total_ranks": instance.total_ranks,
-                    "working_directory": instance.working_directory,
-                    "runner_statuses": runner_statuses,
-                }
-            )
-    return flash_instances

src/exo/plugins/implementations/flash/placement.py

@@ -1,167 +0,0 @@
-"""FLASH plugin placement logic."""
-
-from collections.abc import Mapping
-from copy import deepcopy
-
-from loguru import logger
-
-from exo.plugins.implementations.flash.types import FLASHInstance, LaunchFLASH
-from exo.shared.models.model_cards import ModelCard
-from exo.shared.topology import Topology
-from exo.shared.types.common import Host, ModelId, NodeId
-from exo.shared.types.memory import Memory
-from exo.shared.types.topology import SocketConnection
-from exo.shared.types.worker.instances import BaseInstance, InstanceId
-from exo.shared.types.worker.runners import (
-    RunnerId,
-    ShardAssignments,
-)
-from exo.shared.types.worker.shards import PipelineShardMetadata
-
-
-def place_flash_instance(
-    command: LaunchFLASH,
-    topology: Topology,
-    current_instances: Mapping[InstanceId, BaseInstance],
-) -> dict[InstanceId, BaseInstance]:
-    """Place a FLASH simulation instance across available nodes.
-
-    Unlike MLX instances which use ring/JACCL topology for tensor parallelism,
-    FLASH instances use MPI for communication. We just need to provide the
-    node IPs so the runner can generate an MPI hostfile.
-    """
-    instance_id = InstanceId()
-    target_instances: dict[InstanceId, BaseInstance] = dict(deepcopy(current_instances))
-
-    all_nodes = list(topology.list_nodes())
-
-    if len(all_nodes) < command.min_nodes:
-        raise ValueError(
-            f"Not enough nodes: need {command.min_nodes}, have {len(all_nodes)}"
-        )
-
-    # Select nodes (take the first min_nodes)
-    selected_nodes = all_nodes[: command.min_nodes]
-
-    logger.info(
-        f"Placing FLASH instance '{command.simulation_name}' on {len(selected_nodes)} nodes"
-    )
-
-    # Build shard assignments (one runner per node for FLASH)
-    runner_to_shard: dict[RunnerId, PipelineShardMetadata] = {}
-    node_to_runner: dict[NodeId, RunnerId] = {}
-
-    # Create a dummy ModelCard for FLASH (required by ShardMetadata interface)
-    flash_model_card = ModelCard(
-        model_id=ModelId(command.simulation_name),
-        storage_size=Memory(in_bytes=0),
-        n_layers=1,
-        hidden_size=1,
-        supports_tensor=False,
-        tasks=[],
-    )
-
-    for i, node_id in enumerate(selected_nodes):
-        runner_id = RunnerId()
-        node_to_runner[node_id] = runner_id
-        runner_to_shard[runner_id] = PipelineShardMetadata(
-            device_rank=i,
-            world_size=len(selected_nodes),
-            model_card=flash_model_card,
-            start_layer=0,
-            end_layer=1,
-            n_layers=1,
-        )
-
-    shard_assignments = ShardAssignments(
-        model_id=ModelId(command.simulation_name),
-        runner_to_shard=runner_to_shard,
-        node_to_runner=node_to_runner,
-    )
-
-    # Build hosts_by_node - get hostnames/IPs for MPI hostfile generation
-    hosts_by_node: dict[NodeId, list[Host]] = {}
-
-    # If explicit hosts are provided, use them directly
-    if command.hosts:
-        explicit_hosts = [h.strip() for h in command.hosts.split(",") if h.strip()]
-        logger.info(f"FLASH placement: explicit hosts provided: {explicit_hosts}")
-        for i, node_id in enumerate(selected_nodes):
-            if i < len(explicit_hosts):
-                hosts_by_node[node_id] = [Host(ip=explicit_hosts[i], port=0)]
-                logger.info(
-                    f"FLASH placement: node {node_id} (rank {i}) -> IP {explicit_hosts[i]}"
-                )
-            else:
-                logger.warning(
-                    f"Not enough hosts provided for node {i}, using localhost"
-                )
-                hosts_by_node[node_id] = [Host(ip="127.0.0.1", port=0)]
-        logger.info(
-            f"FLASH placement: coordinator will be rank 0 at IP {explicit_hosts[0]}"
-        )
-    else:
-        # Get each node's own IP by looking at how OTHER nodes see it.
-        # out_edges(A) gives edges A->B with B's IP in sink_multiaddr,
-        # so to find A's IP we look at edges B->A from any other node.
-        for node_id in selected_nodes:
-            candidate_ips: set[str] = set()
-
-            for other_node in all_nodes:
-                if other_node == node_id:
-                    continue
-                for conn in topology.out_edges(other_node):
-                    if conn.sink == node_id and isinstance(conn.edge, SocketConnection):
-                        ip = conn.edge.sink_multiaddr.ip_address
-                        # Skip link-local and localhost addresses
-                        if not ip.startswith("169.254.") and not ip.startswith("127."):
-                            candidate_ips.add(ip)
-
-            # Prefer private network IPs (10.x, 192.168.x) over Tailscale CGNAT (100.64-127.x)
-            chosen_ip: str | None = None
-            for ip in candidate_ips:
-                if ip.startswith(("10.", "192.168.")):
-                    chosen_ip = ip
-                    break
-            if chosen_ip is None and candidate_ips:
-                chosen_ip = next(iter(candidate_ips))
-
-            if chosen_ip:
-                hosts_by_node[node_id] = [Host(ip=chosen_ip, port=0)]
-            else:
-                logger.warning(
-                    f"Could not determine IP for node {node_id}, using localhost"
-                )
-                hosts_by_node[node_id] = [Host(ip="127.0.0.1", port=0)]
-
-            logger.info(
-                f"FLASH placement: node {node_id} -> IP {hosts_by_node[node_id][0].ip}"
-                f" (candidates: {candidate_ips})"
-            )
-
-    total_ranks = len(selected_nodes) * command.ranks_per_node
-
-    # Determine coordinator IP - first node's first host IP
-    first_node_id: NodeId = next(iter(hosts_by_node.keys()))
-    coordinator_ip: str = (
-        hosts_by_node[first_node_id][0].ip
-        if hosts_by_node[first_node_id]
-        else "127.0.0.1"
-    )
-
-    target_instances[instance_id] = FLASHInstance(
-        instance_id=instance_id,
-        shard_assignments=shard_assignments,
-        hosts_by_node=hosts_by_node,
-        flash_executable_path=command.flash_executable_path,
-        parameter_file_path=command.parameter_file_path,
-        working_directory=command.working_directory,
-        ranks_per_node=command.ranks_per_node,
-        total_ranks=total_ranks,
-        simulation_name=command.simulation_name,
-        coordinator_ip=coordinator_ip,
-    )
-
-    logger.info(f"Created FLASH instance {instance_id} with {total_ranks} total ranks")
-
-    return target_instances

src/exo/plugins/implementations/flash/planning.py

@@ -1,37 +0,0 @@
-"""FLASH plugin planning logic."""
-
-from collections.abc import Mapping
-
-from exo.plugins.implementations.flash.types import FLASHInstance
-from exo.shared.types.tasks import LoadModel, Task
-from exo.shared.types.worker.instances import BaseInstance, InstanceId
-from exo.shared.types.worker.runners import RunnerId, RunnerIdle
-from exo.worker.runner.runner_supervisor import RunnerSupervisor
-
-
-def plan_flash(
-    runners: Mapping[RunnerId, RunnerSupervisor],
-    instances: Mapping[InstanceId, BaseInstance],
-) -> Task | None:
-    """Plan tasks specifically for FLASH instances.
-
-    FLASH instances have a simpler lifecycle:
-    - CreateRunner (handled by core _create_runner)
-    - LoadModel (starts the simulation immediately)
-    - Shutdown (handled by core _kill_runner)
-
-    This function handles the LoadModel step for FLASH instances,
-    skipping the MLX-specific download/init/warmup steps.
-    """
-    for runner in runners.values():
-        instance = runner.bound_instance.instance
-
-        # Only handle FLASH instances
-        if not isinstance(instance, FLASHInstance):
-            continue
-
-        # If runner is idle, emit LoadModel to start the simulation
-        if isinstance(runner.status, RunnerIdle):
-            return LoadModel(instance_id=instance.instance_id)
-
-    return None

src/exo/plugins/implementations/flash/plugin.py

@@ -1,98 +0,0 @@
-"""FLASH Plugin - Main plugin class."""
-
-from collections.abc import Callable, Mapping, Sequence
-from typing import Any
-
-from exo.plugins.base import EXOPlugin
-from exo.plugins.implementations.flash.api_handlers import (
-    handle_launch_flash,
-    handle_list_flash_instances,
-    handle_stop_flash,
-)
-from exo.plugins.implementations.flash.placement import place_flash_instance
-from exo.plugins.implementations.flash.planning import plan_flash
-from exo.plugins.implementations.flash.runner import main as flash_runner_main
-from exo.plugins.implementations.flash.types import (
-    FLASHInstance,
-    LaunchFLASH,
-    StopFLASH,
-)
-from exo.shared.topology import Topology
-from exo.shared.types.commands import DeleteInstance
-from exo.shared.types.events import Event
-from exo.shared.types.tasks import Task
-from exo.shared.types.worker.instances import BaseInstance, BoundInstance, InstanceId
-from exo.shared.types.worker.runners import RunnerId
-from exo.utils.channels import MpReceiver, MpSender
-from exo.worker.runner.runner_supervisor import RunnerSupervisor
-
-
-class FLASHPlugin(EXOPlugin):
-    """Plugin for FLASH MPI simulations."""
-
-    @property
-    def name(self) -> str:
-        return "flash"
-
-    @property
-    def version(self) -> str:
-        return "1.0.0"
-
-    def get_command_types(self) -> Sequence[type]:
-        return [LaunchFLASH, StopFLASH]
-
-    def get_instance_type(self) -> type:
-        return FLASHInstance
-
-    def get_api_routes(
-        self,
-    ) -> Sequence[tuple[str, str, Callable[..., Any]]]:
-        return [
-            ("post", "/flash/launch", handle_launch_flash),
-            ("delete", "/flash/{instance_id}", handle_stop_flash),
-            ("get", "/flash/instances", handle_list_flash_instances),
-        ]
-
-    def handles_command(self, command: Any) -> bool:  # pyright: ignore[reportAny]
-        return isinstance(command, (LaunchFLASH, StopFLASH))
-
-    def process_command(
-        self,
-        command: Any,  # pyright: ignore[reportAny]
-        topology: Topology,
-        current_instances: Mapping[InstanceId, BaseInstance],
-    ) -> Sequence[Event]:
-        from exo.master.placement import delete_instance, get_transition_events
-
-        if isinstance(command, LaunchFLASH):
-            placement = place_flash_instance(command, topology, current_instances)
-            return list(get_transition_events(current_instances, placement))
-        elif isinstance(command, StopFLASH):
-            placement = delete_instance(
-                DeleteInstance(instance_id=command.instance_id),
-                current_instances,
-            )
-            return list(get_transition_events(current_instances, placement))
-        return []
-
-    def handles_instance(self, instance: object) -> bool:
-        return isinstance(instance, FLASHInstance)
-
-    def plan_task(
-        self,
-        runners: Mapping[RunnerId, RunnerSupervisor],
-        instances: Mapping[InstanceId, BaseInstance],
-    ) -> Task | None:
-        return plan_flash(runners, instances)
-
-    def should_skip_download(self, instance: object) -> bool:
-        # FLASH instances don't need model downloads
-        return True
-
-    def create_runner(
-        self,
-        bound_instance: BoundInstance,
-        event_sender: MpSender[Event],
-        task_receiver: MpReceiver[Task],
-    ) -> None:
-        flash_runner_main(bound_instance, event_sender, task_receiver)

src/exo/plugins/implementations/flash/runner.py

@@ -1,305 +0,0 @@
-"""FLASH MPI Runner - spawns and monitors FLASH simulations.
-
-Exo-native distributed MPI:
-- Exo handles node discovery and coordination
-- Coordinator generates hostfile from Exo topology
-- mpirun uses exo-rsh (no SSH required) to spawn on remote nodes
-- exo-rsh connects to each node's Exo API (/execute endpoint) for remote execution
-- Workers just report ready and wait
-"""
-# ruff: noqa: I001 - Import order intentional (plugin types before shared types)
-
-import os
-import shutil
-import socket
-import subprocess
-import threading
-
-from loguru import logger
-
-from exo.shared.types.events import (
-    Event,
-    RunnerStatusUpdated,
-    TaskAcknowledged,
-    TaskStatusUpdated,
-)
-from exo.shared.types.tasks import (
-    LoadModel,
-    Shutdown,
-    Task,
-    TaskStatus,
-)
-from exo.plugins.implementations.flash.types import FLASHInstance
-from exo.shared.types.worker.instances import BoundInstance
-from exo.shared.types.worker.runners import (
-    RunnerFailed,
-    RunnerIdle,
-    RunnerLoading,
-    RunnerReady,
-    RunnerRunning,
-    RunnerShutdown,
-    RunnerShuttingDown,
-    RunnerStatus,
-)
-from exo.utils.channels import MpReceiver, MpSender
-
-# Find mpirun in PATH, fallback to common locations
-MPIRUN_PATH = shutil.which("mpirun") or "/opt/homebrew/bin/mpirun"
-
-# exo-rsh is installed as console script by exo package
-_exo_rsh_path = shutil.which("exo-rsh")
-if not _exo_rsh_path:
-    raise RuntimeError("exo-rsh not found in PATH - this should be installed with exo")
-EXO_RSH_PATH: str = _exo_rsh_path
-
-
-def get_my_rank(instance: FLASHInstance, my_node_id: str) -> int:
-    """Determine this node's rank based on position in hosts_by_node."""
-    for i, node_id in enumerate(instance.hosts_by_node.keys()):
-        if str(node_id) == str(my_node_id):
-            return i
-    return -1
-
-
-def get_coordinator_host(instance: FLASHInstance) -> str:
-    """Get the IP of the coordinator node."""
-    return instance.coordinator_ip
-
-
-def resolve_host(host: str) -> str:
-    """Resolve host string to a usable hostname for MPI hostfile.
-
-    Accepts either an IP address or hostname. For IPs, attempts to resolve
-    to a hostname via DNS/mDNS. Hostnames are returned as-is after validation.
-    """
-    # Check if input is already a hostname (not an IP)
-    try:
-        socket.inet_aton(host)
-        is_ip = True
-    except socket.error:
-        is_ip = False
-
-    if not is_ip:
-        # Already a hostname, verify it resolves and return as-is
-        try:
-            socket.gethostbyname(host)
-            return host
-        except socket.gaierror:
-            logger.warning(f"Hostname {host} does not resolve, using anyway")
-            return host
-
-    # It's an IP address, try to resolve to hostname
-    try:
-        hostname, _, _ = socket.gethostbyaddr(host)
-        hostname = hostname.split(".")[0]
-        logger.info(f"Resolved {host} to {hostname}")
-        return hostname
-    except socket.herror:
-        pass
-
-    # Fall back to IP
-    logger.warning(f"Could not resolve {host} to hostname, using IP directly")
-    return host
-
-
-def generate_hostfile(instance: FLASHInstance, working_dir: str) -> str:
-    """Generate MPI hostfile from instance topology."""
-    hostfile_path = os.path.join(working_dir, "flash_hosts.txt")
-    with open(hostfile_path, "w") as f:
-        for _node_id, hosts in instance.hosts_by_node.items():
-            if hosts:
-                host = resolve_host(hosts[0].ip)
-                f.write(f"{host} slots={instance.ranks_per_node}\n")
-    logger.info(f"Generated hostfile at {hostfile_path}")
-    with open(hostfile_path, "r") as f:
-        logger.info(f"Hostfile contents:\n{f.read()}")
-    return hostfile_path
-
-
-def main(
-    bound_instance: BoundInstance,
-    event_sender: MpSender[Event],
-    task_receiver: MpReceiver[Task],
-) -> None:
-    """Main FLASH runner loop.
-
-    Coordinator: generates hostfile and runs mpirun (uses exo-rsh instead of SSH)
-    Workers: just report ready and wait for mpirun to spawn processes on them
-    """
-    assert isinstance(bound_instance.instance, FLASHInstance)
-    instance = bound_instance.instance
-    runner_id = bound_instance.bound_runner_id
-    my_node_id = str(bound_instance.bound_node_id)
-
-    logger.info(f"FLASH runner starting for simulation: {instance.simulation_name}")
-
-    my_rank = get_my_rank(instance, my_node_id)
-    world_size = len(instance.hosts_by_node)
-    is_coordinator = my_rank == 0
-    coordinator_ip = get_coordinator_host(instance)
-
-    logger.info(
-        f"FLASH node: rank={my_rank}, world_size={world_size}, coordinator={is_coordinator}"
-    )
-    logger.info(f"FLASH coordinator IP: {coordinator_ip}")
-
-    process: subprocess.Popen[bytes] | None = None
-    current_status: RunnerStatus = RunnerIdle()
-    shutdown_requested = False
-
-    event_sender.send(
-        RunnerStatusUpdated(runner_id=runner_id, runner_status=current_status)
-    )
-
-    def monitor_output(proc: subprocess.Popen[bytes]) -> None:
-        """Monitor FLASH stdout for progress updates."""
-        if proc.stdout is None:
-            return
-        for line in iter(proc.stdout.readline, b""):
-            if shutdown_requested:
-                break
-            try:
-                decoded: str = line.decode("utf-8", errors="replace").strip()
-                if decoded:
-                    logger.info(f"[FLASH] {decoded}")
-            except Exception as e:
-                logger.warning(f"Error parsing FLASH output: {e}")
-
-    with task_receiver as tasks:
-        for task in tasks:
-            event_sender.send(
-                TaskStatusUpdated(task_id=task.task_id, task_status=TaskStatus.Running)
-            )
-            event_sender.send(TaskAcknowledged(task_id=task.task_id))
-
-            match task:
-                case LoadModel() if isinstance(current_status, RunnerIdle):
-                    current_status = RunnerLoading()
-                    logger.info("Starting FLASH simulation")
-                    event_sender.send(
-                        RunnerStatusUpdated(
-                            runner_id=runner_id, runner_status=current_status
-                        )
-                    )
-
-                    try:
-                        if is_coordinator:
-                            # Coordinator: generate hostfile and run mpirun
-                            hostfile = generate_hostfile(
-                                instance, instance.working_directory
-                            )
-
-                            iface = instance.network_interface
-                            cmd = [
-                                MPIRUN_PATH,
-                                "-np",
-                                str(instance.total_ranks),
-                                "--hostfile",
-                                hostfile,
-                                "--wdir",
-                                instance.working_directory,
-                                "--oversubscribe",
-                                "--allow-run-as-root",
-                                "--mca",
-                                "btl",
-                                "tcp,self",
-                                "--mca",
-                                "btl_tcp_if_include",
-                                iface,
-                                "--mca",
-                                "oob_tcp_if_include",
-                                iface,
-                                "--mca",
-                                "plm_rsh_no_tree_spawn",
-                                "1",
-                            ]
-
-                            # Use exo-rsh for remote execution (no SSH needed)
-                            cmd.extend(["--mca", "plm_rsh_agent", EXO_RSH_PATH])
-
-                            cmd.append(instance.flash_executable_path)
-
-                            logger.info(f"FLASH distributed launch: {' '.join(cmd)}")
-
-                            process = subprocess.Popen(
-                                cmd,
-                                cwd=instance.working_directory,
-                                stdout=subprocess.PIPE,
-                                stderr=subprocess.STDOUT,
-                            )
-
-                            monitor_thread = threading.Thread(
-                                target=monitor_output, args=(process,), daemon=True
-                            )
-                            monitor_thread.start()
-
-                            current_status = RunnerRunning()
-                            logger.info(
-                                f"FLASH running on {world_size} nodes with {instance.total_ranks} ranks"
-                            )
-
-                        else:
-                            # Worker: mpirun on coordinator will use exo-rsh to spawn processes here
-                            logger.info(
-                                f"Worker {my_rank}: Ready for mpirun to spawn processes via exo-rsh"
-                            )
-                            current_status = RunnerRunning()
-
-                    except Exception as e:
-                        logger.error(f"Failed to start FLASH: {e}")
-                        import traceback
-
-                        logger.error(traceback.format_exc())
-                        current_status = RunnerFailed(error_message=str(e))
-
-                case Shutdown():
-                    shutdown_requested = True
-                    current_status = RunnerShuttingDown()
-                    logger.info("FLASH runner shutting down")
-                    event_sender.send(
-                        RunnerStatusUpdated(
-                            runner_id=runner_id, runner_status=current_status
-                        )
-                    )
-
-                    if process and process.poll() is None:
-                        logger.info("Terminating FLASH simulation")
-                        process.terminate()
-                        try:
-                            process.wait(timeout=10)
-                        except subprocess.TimeoutExpired:
-                            logger.warning("FLASH didn't terminate, killing")
-                            process.kill()
-                            process.wait()
-
-                    current_status = RunnerShutdown()
-
-                case _:
-                    if process and process.poll() is not None:
-                        exit_code = process.returncode
-                        if exit_code == 0:
-                            logger.info("FLASH simulation completed successfully")
-                            current_status = RunnerReady()
-                        else:
-                            logger.error(
-                                f"FLASH simulation failed with code {exit_code}"
-                            )
-                            current_status = RunnerFailed(
-                                error_message=f"Exit code {exit_code}"
-                            )
-
-            event_sender.send(
-                TaskStatusUpdated(task_id=task.task_id, task_status=TaskStatus.Complete)
-            )
-            event_sender.send(
-                RunnerStatusUpdated(runner_id=runner_id, runner_status=current_status)
-            )
-
-            if isinstance(current_status, RunnerShutdown):
-                break
-
-    if process and process.poll() is None:
-        process.terminate()
-        process.wait(timeout=5)
-
-    logger.info("FLASH runner exiting")

src/exo/plugins/implementations/flash/types.py

@@ -1,62 +0,0 @@
-"""FLASH plugin types - commands and instances."""
-
-from exo.plugins.type_registry import command_registry, instance_registry
-from exo.shared.types.commands import BaseCommand
-from exo.shared.types.common import Host, NodeId
-from exo.shared.types.worker.instances import BaseInstance, InstanceId
-from exo.shared.types.worker.runners import RunnerId
-from exo.shared.types.worker.shards import ShardMetadata
-
-# ============================================================================
-# Commands
-# ============================================================================
-
-
-@command_registry.register
-class LaunchFLASH(BaseCommand):
-    """Command to launch a FLASH MPI simulation."""
-
-    simulation_name: str
-    flash_executable_path: str
-    parameter_file_path: str
-    working_directory: str
-    ranks_per_node: int = 1
-    min_nodes: int = 1
-    # Optional: explicit hostnames for MPI (e.g., "s14,james21-1")
-    # Used when topology edges don't contain IP addresses
-    hosts: str = ""
-
-
-@command_registry.register
-class StopFLASH(BaseCommand):
-    """Command to stop a running FLASH simulation."""
-
-    instance_id: InstanceId
-
-
-# ============================================================================
-# Instances
-# ============================================================================
-
-
-@instance_registry.register
-class FLASHInstance(BaseInstance):
-    """Instance for FLASH MPI simulation.
-
-    Unlike MLX instances which do tensor parallelism, FLASH instances
-    coordinate MPI processes across nodes. Each node runs one or more
-    MPI ranks of the FLASH simulation.
-    """
-
-    hosts_by_node: dict[NodeId, list[Host]]
-    flash_executable_path: str
-    parameter_file_path: str
-    working_directory: str
-    ranks_per_node: int = 1
-    total_ranks: int
-    simulation_name: str
-    coordinator_ip: str
-    network_interface: str = "en0"  # Network interface for MPI (e.g., en0, eth0)
-
-    def shard(self, runner_id: RunnerId) -> ShardMetadata | None:
-        return self.shard_assignments.runner_to_shard.get(runner_id, None)

src/exo/plugins/registry.py

@@ -1,110 +0,0 @@
-"""Plugin registry for discovering and managing plugins."""
-
-from collections.abc import Callable, Sequence
-from typing import Any
-
-from loguru import logger
-
-from exo.plugins.base import EXOPlugin
-
-
-class PluginRegistry:
-    """Central registry for all plugins."""
-
-    _instance: "PluginRegistry | None" = None
-
-    def __init__(self) -> None:
-        self._plugins: dict[str, EXOPlugin] = {}
-        self._command_handlers: dict[type, EXOPlugin] = {}
-        self._instance_handlers: dict[type, EXOPlugin] = {}
-
-    @classmethod
-    def get(cls) -> "PluginRegistry":
-        """Get the singleton registry instance."""
-        if cls._instance is None:
-            cls._instance = cls()
-        return cls._instance
-
-    @classmethod
-    def reset(cls) -> None:
-        """Reset the singleton instance (useful for testing)."""
-        cls._instance = None
-
-    def register(self, plugin: EXOPlugin) -> None:
-        """Register a plugin and its types."""
-        if plugin.name in self._plugins:
-            raise ValueError(f"Plugin '{plugin.name}' already registered")
-
-        logger.info(f"Registering plugin: {plugin.name} v{plugin.version}")
-
-        self._plugins[plugin.name] = plugin
-
-        # Register command handlers
-        for cmd_type in plugin.get_command_types():
-            self._command_handlers[cmd_type] = plugin
-            logger.debug(f"  Registered command: {cmd_type.__name__}")
-
-        # Register instance handler
-        instance_type = plugin.get_instance_type()
-        self._instance_handlers[instance_type] = plugin
-        logger.debug(f"  Registered instance: {instance_type.__name__}")
-
-    def get_plugin(self, name: str) -> EXOPlugin | None:
-        """Get a plugin by name."""
-        return self._plugins.get(name)
-
-    def get_plugin_for_command(self, command: object) -> EXOPlugin | None:
-        """Get the plugin that handles a command."""
-        for plugin in self._plugins.values():
-            if plugin.handles_command(command):
-                return plugin
-        return None
-
-    def get_plugin_for_instance(self, instance: object) -> EXOPlugin | None:
-        """Get the plugin that manages an instance."""
-        for plugin in self._plugins.values():
-            if plugin.handles_instance(instance):
-                return plugin
-        return None
-
-    def all_plugins(self) -> Sequence[EXOPlugin]:
-        """Get all registered plugins."""
-        return list(self._plugins.values())
-
-    def get_all_api_routes(
-        self,
-    ) -> Sequence[tuple[str, str, Callable[..., Any], EXOPlugin]]:
-        """Get all API routes from all plugins."""
-        routes: list[tuple[str, str, Callable[..., Any], EXOPlugin]] = []
-        for plugin in self._plugins.values():
-            for method, path, handler in plugin.get_api_routes():
-                routes.append((method, path, handler, plugin))
-        return routes
-
-
-def discover_plugins() -> None:
-    """Auto-discover and register plugins from the implementations directory.
-
-    Plugins should have a register() function that returns an EXOPlugin instance.
-    """
-    import importlib
-    import pkgutil
-
-    registry = PluginRegistry.get()
-
-    try:
-        import exo.plugins.implementations as impl_package
-
-        for _, module_name, _ in pkgutil.iter_modules(impl_package.__path__):
-            try:
-                module = importlib.import_module(
-                    f"exo.plugins.implementations.{module_name}"
-                )
-                if hasattr(module, "register"):
-                    plugin = module.register()  # pyright: ignore[reportAny]
-                    if plugin is not None:
-                        registry.register(plugin)  # pyright: ignore[reportAny]
-            except Exception as e:
-                logger.warning(f"Failed to load plugin {module_name}: {e}")
-    except ImportError:
-        logger.debug("No plugin implementations package found")

src/exo/plugins/type_registry.py

@@ -1,84 +0,0 @@
-"""Dynamic type registry for plugin types.
-
-This module provides a registry system that allows plugins to register their
-command and instance types dynamically, eliminating the need for static union
-types and avoiding circular imports.
-"""
-
-from typing import TypeVar
-
-from loguru import logger
-
-from exo.utils.pydantic_ext import CamelCaseModel
-
-# TypeVar for preserving exact types through the register decorator
-_TCls = TypeVar("_TCls", bound=type[CamelCaseModel])
-
-
-class TypeRegistry[T: CamelCaseModel]:
-    """Registry for dynamically registered Pydantic types.
-
-    Enables plugins to register their types at import time. Deserialization
-    uses the class name from the tagged JSON format to look up the correct type.
-    """
-
-    def __init__(self, name: str) -> None:
-        self._name = name
-        self._types: dict[str, type[T]] = {}
-
-    def register(self, cls: _TCls) -> _TCls:
-        """Decorator to register a type with this registry.
-
-        Preserves the exact type through the decorator for proper type checking.
-        """
-        self._types[cls.__name__] = cls  # type: ignore[assignment]
-        logger.debug(f"{self._name}: registered {cls.__name__}")
-        return cls
-
-    def get(self, name: str) -> type[T] | None:
-        """Look up a type by class name."""
-        return self._types.get(name)
-
-    def all_types(self) -> dict[str, type[T]]:
-        """Return all registered types."""
-        return dict(self._types)
-
-    def deserialize(self, data: dict[str, dict[str, object]] | CamelCaseModel) -> T:
-        """Deserialize dict to the appropriate registered type.
-
-        Supports two formats:
-        1. Tagged format: {"ClassName": {...fields...}} - used for network serialization
-        2. Flat format: {...fields...} - used for API requests, tries each type
-        """
-        # If already deserialized (e.g., from Pydantic), return as-is
-        if isinstance(data, CamelCaseModel):
-            return data  # type: ignore[return-value]
-
-        # Check for tagged format: single key that matches a registered type
-        if len(data) == 1:
-            class_name: str = next(iter(data.keys()))
-            cls = self._types.get(class_name)
-            if cls is not None:
-                return cls.model_validate(data[class_name], strict=False)
-
-        # Flat format: try each registered type, use first that validates
-        errors: list[str] = []
-        for type_name, cls in self._types.items():
-            try:
-                return cls.model_validate(data, strict=False)
-            except Exception as e:  # noqa: BLE001
-                errors.append(f"{type_name}: {e}")
-
-        # None matched - provide helpful error
-        available = ", ".join(self._types.keys())
-        raise ValueError(
-            f"{self._name}: could not deserialize data. "
-            f"Available types: {available}. Errors: {'; '.join(errors[:3])}"
-        )
-
-
-# Global registries for commands, instances, events, and tasks
-command_registry: TypeRegistry[CamelCaseModel] = TypeRegistry("CommandRegistry")
-instance_registry: TypeRegistry[CamelCaseModel] = TypeRegistry("InstanceRegistry")
-event_registry: TypeRegistry[CamelCaseModel] = TypeRegistry("EventRegistry")
-task_registry: TypeRegistry[CamelCaseModel] = TypeRegistry("TaskRegistry")

src/exo/routing/topics.py

@@ -3,7 +3,7 @@ from enum import Enum
 
 from exo.routing.connection_message import ConnectionMessage
 from exo.shared.election import ElectionMessage
-from exo.shared.types.commands import ForwarderCommand
+from exo.shared.types.commands import ForwarderCommand, ForwarderDownloadCommand
 from exo.shared.types.events import (
     ForwarderEvent,
 )
@@ -30,7 +30,7 @@ class TypedTopic[T: CamelCaseModel]:
 
     @staticmethod
     def serialize(t: T) -> bytes:
-        return t.model_dump_json(by_alias=True, serialize_as_any=True).encode("utf-8")
+        return t.model_dump_json().encode("utf-8")
 
     def deserialize(self, b: bytes) -> T:
         return self.model_type.model_validate_json(b.decode("utf-8"))
@@ -45,3 +45,6 @@ ELECTION_MESSAGES = TypedTopic(
 CONNECTION_MESSAGES = TypedTopic(
     "connection_messages", PublishPolicy.Never, ConnectionMessage
 )
+DOWNLOAD_COMMANDS = TypedTopic(
+    "download_commands", PublishPolicy.Always, ForwarderDownloadCommand
+)

src/exo/rsh/__init__.py

@@ -1,13 +0,0 @@
-"""Exo RSH - Remote Shell for MPI without SSH.
-
-This module provides a remote execution mechanism that allows mpirun to spawn
-processes on remote nodes without requiring SSH setup. It works by:
-
-1. Each Exo node runs an API server on port 52415 with an /execute endpoint
-2. The exo-rsh script acts as a drop-in replacement for ssh
-3. When mpirun calls "exo-rsh hostname command", it HTTP POSTs to the target's /execute
-4. The target executes the command and returns output
-
-Usage:
-    mpirun --mca plm_rsh_agent exo-rsh -np 4 --hostfile hosts.txt ./program
-"""

src/exo/rsh/client.py

@@ -1,101 +0,0 @@
-#!/usr/bin/env python3
-"""exo-rsh - Remote shell client for MPI.
-
-This script is called by mpirun as a replacement for ssh.
-Usage: exo-rsh [ssh-options...] hostname command [args...]
-
-It connects to the target node's Exo API (port 52415) and executes the command.
-"""
-
-import json
-import socket
-import sys
-from typing import Any, cast
-from urllib.error import URLError
-from urllib.request import Request, urlopen
-
-# Use the same port as Exo's API server
-EXO_API_PORT = 52415
-
-
-def resolve_hostname(hostname: str) -> str:
-    """Resolve hostname to IP address."""
-    try:
-        return socket.gethostbyname(hostname)
-    except socket.gaierror:
-        # If resolution fails, try using the hostname directly
-        return hostname
-
-
-def main():
-    # Parse arguments - mpirun calls us like: exo-rsh [options] hostname command [args...]
-    # SSH options we might see: -x (disable X11), -o options, etc.
-    args = sys.argv[1:]
-
-    # Skip SSH-style options
-    hostname = None
-    command_start = 0
-
-    i = 0
-    while i < len(args):
-        arg = args[i]
-        if arg.startswith("-"):
-            # Skip option and its value if needed
-            if arg in ("-o", "-i", "-l", "-p", "-F"):
-                i += 2  # Skip option and its argument
-                continue
-            i += 1
-            continue
-        else:
-            # First non-option is the hostname
-            hostname = arg
-            command_start = i + 1
-            break
-        i += 1
-
-    if hostname is None or command_start >= len(args):
-        print("Usage: exo-rsh [options] hostname command [args...]", file=sys.stderr)
-        sys.exit(1)
-
-    command = args[command_start:]
-
-    # Resolve hostname to IP
-    ip = resolve_hostname(hostname)
-
-    # Make request to Exo API
-    url = f"http://{ip}:{EXO_API_PORT}/execute"
-    data = json.dumps({"command": command}).encode("utf-8")
-
-    try:
-        req = Request(url, data=data, headers={"Content-Type": "application/json"})
-        with urlopen(req, timeout=300) as response:  # pyright: ignore[reportAny]
-            response_body: bytes = cast(bytes, response.read())  # pyright: ignore[reportAny]
-            result: dict[str, Any] = json.loads(response_body.decode("utf-8"))  # pyright: ignore[reportAny]
-
-        # Output stdout/stderr
-        stdout: str = cast(str, result.get("stdout", ""))
-        stderr: str = cast(str, result.get("stderr", ""))
-        exit_code: int = cast(int, result.get("exit_code", 0))
-
-        if stdout:
-            sys.stdout.write(stdout)
-            sys.stdout.flush()
-        if stderr:
-            sys.stderr.write(stderr)
-            sys.stderr.flush()
-
-        sys.exit(exit_code)
-
-    except URLError as e:
-        print(
-            f"exo-rsh: Failed to connect to {hostname}:{EXO_API_PORT}: {e}",
-            file=sys.stderr,
-        )
-        sys.exit(255)
-    except Exception as e:
-        print(f"exo-rsh: Error: {e}", file=sys.stderr)
-        sys.exit(1)
-
-
-if __name__ == "__main__":
-    main()

src/exo/shared/apply.py

@@ -6,8 +6,8 @@ from loguru import logger
 
 from exo.shared.types.common import NodeId
 from exo.shared.types.events import (
-    BaseEvent,
     ChunkGenerated,
+    Event,
     IndexedEvent,
     InputChunkReceived,
     InstanceCreated,
@@ -30,12 +30,13 @@ from exo.shared.types.profiling import (
     NodeIdentity,
     NodeNetworkInfo,
     NodeThunderboltInfo,
+    ThunderboltBridgeStatus,
 )
 from exo.shared.types.state import State
-from exo.shared.types.tasks import BaseTask, TaskId, TaskStatus
+from exo.shared.types.tasks import Task, TaskId, TaskStatus
 from exo.shared.types.topology import Connection, RDMAConnection
 from exo.shared.types.worker.downloads import DownloadProgress
-from exo.shared.types.worker.instances import BaseInstance, InstanceId
+from exo.shared.types.worker.instances import Instance, InstanceId
 from exo.shared.types.worker.runners import RunnerId, RunnerStatus
 from exo.utils.info_gatherer.info_gatherer import (
     MacmonMetrics,
@@ -46,10 +47,11 @@ from exo.utils.info_gatherer.info_gatherer import (
     NodeConfig,
     NodeNetworkInterfaces,
     StaticNodeInformation,
+    ThunderboltBridgeInfo,
 )
 
 
-def event_apply(event: BaseEvent, state: State) -> State:
+def event_apply(event: Event, state: State) -> State:
     """Apply an event to state."""
     match event:
         case (
@@ -82,10 +84,6 @@ def event_apply(event: BaseEvent, state: State) -> State:
             return apply_topology_edge_created(event, state)
         case TopologyEdgeDeleted():
             return apply_topology_edge_deleted(event, state)
-        case _:
-            # Unknown event types from plugins are ignored
-            logger.debug(f"Ignoring unknown event type: {type(event).__name__}")
-            return state
 
 
 def apply(state: State, event: IndexedEvent) -> State:
@@ -126,12 +124,12 @@ def apply_node_download_progress(event: NodeDownloadProgress, state: State) -> S
 
 
 def apply_task_created(event: TaskCreated, state: State) -> State:
-    new_tasks: Mapping[TaskId, BaseTask] = {**state.tasks, event.task_id: event.task}
+    new_tasks: Mapping[TaskId, Task] = {**state.tasks, event.task_id: event.task}
     return state.model_copy(update={"tasks": new_tasks})
 
 
 def apply_task_deleted(event: TaskDeleted, state: State) -> State:
-    new_tasks: Mapping[TaskId, BaseTask] = {
+    new_tasks: Mapping[TaskId, Task] = {
         tid: task for tid, task in state.tasks.items() if tid != event.task_id
     }
     return state.model_copy(update={"tasks": new_tasks})
@@ -150,7 +148,7 @@ def apply_task_status_updated(event: TaskStatusUpdated, state: State) -> State:
         update["error_message"] = None
 
     updated_task = state.tasks[event.task_id].model_copy(update=update)
-    new_tasks: Mapping[TaskId, BaseTask] = {**state.tasks, event.task_id: updated_task}
+    new_tasks: Mapping[TaskId, Task] = {**state.tasks, event.task_id: updated_task}
     return state.model_copy(update={"tasks": new_tasks})
 
 
@@ -162,13 +160,13 @@ def apply_task_failed(event: TaskFailed, state: State) -> State:
     updated_task = state.tasks[event.task_id].model_copy(
         update={"error_type": event.error_type, "error_message": event.error_message}
     )
-    new_tasks: Mapping[TaskId, BaseTask] = {**state.tasks, event.task_id: updated_task}
+    new_tasks: Mapping[TaskId, Task] = {**state.tasks, event.task_id: updated_task}
     return state.model_copy(update={"tasks": new_tasks})
 
 
 def apply_instance_created(event: InstanceCreated, state: State) -> State:
     instance = event.instance
-    new_instances: Mapping[InstanceId, BaseInstance] = {
+    new_instances: Mapping[InstanceId, Instance] = {
         **state.instances,
         instance.instance_id: instance,
     }
@@ -176,7 +174,7 @@ def apply_instance_created(event: InstanceCreated, state: State) -> State:
 
 
 def apply_instance_deleted(event: InstanceDeleted, state: State) -> State:
-    new_instances: Mapping[InstanceId, BaseInstance] = {
+    new_instances: Mapping[InstanceId, Instance] = {
         iid: inst for iid, inst in state.instances.items() if iid != event.instance_id
     }
     return state.model_copy(update={"instances": new_instances})
@@ -229,6 +227,21 @@ def apply_node_timed_out(event: NodeTimedOut, state: State) -> State:
         for key, value in state.node_thunderbolt.items()
         if key != event.node_id
     }
+    node_thunderbolt_bridge = {
+        key: value
+        for key, value in state.node_thunderbolt_bridge.items()
+        if key != event.node_id
+    }
+    # Only recompute cycles if the leaving node had TB bridge enabled
+    leaving_node_status = state.node_thunderbolt_bridge.get(event.node_id)
+    leaving_node_had_tb_enabled = (
+        leaving_node_status is not None and leaving_node_status.enabled
+    )
+    thunderbolt_bridge_cycles = (
+        topology.get_thunderbolt_bridge_cycles(node_thunderbolt_bridge, node_network)
+        if leaving_node_had_tb_enabled
+        else [list(cycle) for cycle in state.thunderbolt_bridge_cycles]
+    )
     return state.model_copy(
         update={
             "downloads": downloads,
@@ -239,6 +252,8 @@ def apply_node_timed_out(event: NodeTimedOut, state: State) -> State:
             "node_system": node_system,
             "node_network": node_network,
             "node_thunderbolt": node_thunderbolt,
+            "node_thunderbolt_bridge": node_thunderbolt_bridge,
+            "thunderbolt_bridge_cycles": thunderbolt_bridge_cycles,
         }
     )
 
@@ -316,6 +331,22 @@ def apply_node_gathered_info(event: NodeGatheredInfo, state: State) -> State:
                 if tb_conn.sink_uuid in conn_map
             ]
             topology.replace_all_out_rdma_connections(event.node_id, as_rdma_conns)
+        case ThunderboltBridgeInfo():
+            new_tb_bridge: dict[NodeId, ThunderboltBridgeStatus] = {
+                **state.node_thunderbolt_bridge,
+                event.node_id: info.status,
+            }
+            update["node_thunderbolt_bridge"] = new_tb_bridge
+            # Only recompute cycles if the enabled status changed
+            old_status = state.node_thunderbolt_bridge.get(event.node_id)
+            old_enabled = old_status.enabled if old_status else False
+            new_enabled = info.status.enabled
+            if old_enabled != new_enabled:
+                update["thunderbolt_bridge_cycles"] = (
+                    topology.get_thunderbolt_bridge_cycles(
+                        new_tb_bridge, state.node_network
+                    )
+                )
 
     return state.model_copy(update=update)
 

src/exo/shared/constants.py

@@ -49,3 +49,7 @@ LIBP2P_COMMANDS_TOPIC = "commands"
 EXO_MAX_CHUNK_SIZE = 512 * 1024
 
 EXO_IMAGE_CACHE_DIR = EXO_CACHE_HOME / "images"
+
+EXO_ENABLE_IMAGE_MODELS = (
+    os.getenv("EXO_ENABLE_IMAGE_MODELS", "false").lower() == "true"
+)

src/exo/shared/models/model_cards.py

@@ -9,6 +9,7 @@ from huggingface_hub import model_info
 from loguru import logger
 from pydantic import BaseModel, Field, PositiveInt, field_validator
 
+from exo.shared.constants import EXO_ENABLE_IMAGE_MODELS
 from exo.shared.types.common import ModelId
 from exo.shared.types.memory import Memory
 from exo.utils.pydantic_ext import CamelCaseModel
@@ -410,161 +411,294 @@ MODEL_CARDS: dict[str, ModelCard] = {
         supports_tensor=True,
         tasks=[ModelTask.TextGeneration],
     ),
-    # Image models commented out - feature not stable (see https://github.com/exo-explore/exo/issues/1242)
-    # "flux1-schnell": ModelCard(
-    #     model_id=ModelId("black-forest-labs/FLUX.1-schnell"),
-    #     storage_size=Memory.from_bytes(23782357120 + 9524621312),
-    #     n_layers=57,
-    #     hidden_size=1,
-    #     supports_tensor=False,
-    #     tasks=[ModelTask.TextToImage],
-    #     components=[
-    #         ComponentInfo(
-    #             component_name="text_encoder",
-    #             component_path="text_encoder/",
-    #             storage_size=Memory.from_kb(0),
-    #             n_layers=12,
-    #             can_shard=False,
-    #             safetensors_index_filename=None,  # Single file
-    #         ),
-    #         ComponentInfo(
-    #             component_name="text_encoder_2",
-    #             component_path="text_encoder_2/",
-    #             storage_size=Memory.from_bytes(9524621312),
-    #             n_layers=24,
-    #             can_shard=False,
-    #             safetensors_index_filename="model.safetensors.index.json",
-    #         ),
-    #         ComponentInfo(
-    #             component_name="transformer",
-    #             component_path="transformer/",
-    #             storage_size=Memory.from_bytes(23782357120),
-    #             n_layers=57,  # 19 transformer_blocks + 38 single_transformer_blocks
-    #             can_shard=True,
-    #             safetensors_index_filename="diffusion_pytorch_model.safetensors.index.json",
-    #         ),
-    #         ComponentInfo(
-    #             component_name="vae",
-    #             component_path="vae/",
-    #             storage_size=Memory.from_kb(0),
-    #             n_layers=None,
-    #             can_shard=False,
-    #             safetensors_index_filename=None,
-    #         ),
-    #     ],
-    # ),
-    # "flux1-dev": ModelCard(
-    #     model_id=ModelId("black-forest-labs/FLUX.1-dev"),
-    #     storage_size=Memory.from_bytes(23782357120 + 9524621312),
-    #     n_layers=57,
-    #     hidden_size=1,
-    #     supports_tensor=False,
-    #     tasks=[ModelTask.TextToImage, ModelTask.ImageToImage],
-    #     components=[
-    #         ComponentInfo(
-    #             component_name="text_encoder",
-    #             component_path="text_encoder/",
-    #             storage_size=Memory.from_kb(0),
-    #             n_layers=12,
-    #             can_shard=False,
-    #             safetensors_index_filename=None,  # Single file
-    #         ),
-    #         ComponentInfo(
-    #             component_name="text_encoder_2",
-    #             component_path="text_encoder_2/",
-    #             storage_size=Memory.from_bytes(9524621312),
-    #             n_layers=24,
-    #             can_shard=False,
-    #             safetensors_index_filename="model.safetensors.index.json",
-    #         ),
-    #         ComponentInfo(
-    #             component_name="transformer",
-    #             component_path="transformer/",
-    #             storage_size=Memory.from_bytes(23802816640),
-    #             n_layers=57,  # 19 transformer_blocks + 38 single_transformer_blocks
-    #             can_shard=True,
-    #             safetensors_index_filename="diffusion_pytorch_model.safetensors.index.json",
-    #         ),
-    #         ComponentInfo(
-    #             component_name="vae",
-    #             component_path="vae/",
-    #             storage_size=Memory.from_kb(0),
-    #             n_layers=None,
-    #             can_shard=False,
-    #             safetensors_index_filename=None,
-    #         ),
-    #     ],
-    # ),
-    # "qwen-image": ModelCard(
-    #     model_id=ModelId("Qwen/Qwen-Image"),
-    #     storage_size=Memory.from_bytes(16584333312 + 40860802176),
-    #     n_layers=60,  # Qwen has 60 transformer blocks (all joint-style)
-    #     hidden_size=1,
-    #     supports_tensor=False,
-    #     tasks=[ModelTask.TextToImage, ModelTask.ImageToImage],
-    #     components=[
-    #         ComponentInfo(
-    #             component_name="text_encoder",
-    #             component_path="text_encoder/",
-    #             storage_size=Memory.from_kb(16584333312),
-    #             n_layers=12,
-    #             can_shard=False,
-    #             safetensors_index_filename=None,  # Single file
-    #         ),
-    #         ComponentInfo(
-    #             component_name="transformer",
-    #             component_path="transformer/",
-    #             storage_size=Memory.from_bytes(40860802176),
-    #             n_layers=60,
-    #             can_shard=True,
-    #             safetensors_index_filename="diffusion_pytorch_model.safetensors.index.json",
-    #         ),
-    #         ComponentInfo(
-    #             component_name="vae",
-    #             component_path="vae/",
-    #             storage_size=Memory.from_kb(0),
-    #             n_layers=None,
-    #             can_shard=False,
-    #             safetensors_index_filename=None,
-    #         ),
-    #     ],
-    # ),
-    # "qwen-image-edit-2509": ModelCard(
-    #     model_id=ModelId("Qwen/Qwen-Image-Edit-2509"),
-    #     storage_size=Memory.from_bytes(16584333312 + 40860802176),
-    #     n_layers=60,  # Qwen has 60 transformer blocks (all joint-style)
-    #     hidden_size=1,
-    #     supports_tensor=False,
-    #     tasks=[ModelTask.ImageToImage],
-    #     components=[
-    #         ComponentInfo(
-    #             component_name="text_encoder",
-    #             component_path="text_encoder/",
-    #             storage_size=Memory.from_kb(16584333312),
-    #             n_layers=12,
-    #             can_shard=False,
-    #             safetensors_index_filename=None,  # Single file
-    #         ),
-    #         ComponentInfo(
-    #             component_name="transformer",
-    #             component_path="transformer/",
-    #             storage_size=Memory.from_bytes(40860802176),
-    #             n_layers=60,
-    #             can_shard=True,
-    #             safetensors_index_filename="diffusion_pytorch_model.safetensors.index.json",
-    #         ),
-    #         ComponentInfo(
-    #             component_name="vae",
-    #             component_path="vae/",
-    #             storage_size=Memory.from_kb(0),
-    #             n_layers=None,
-    #             can_shard=False,
-    #             safetensors_index_filename=None,
-    #         ),
-    #     ],
-    # ),
 }
 
+_IMAGE_BASE_MODEL_CARDS: dict[str, ModelCard] = {
+    "flux1-schnell": ModelCard(
+        model_id=ModelId("exolabs/FLUX.1-schnell"),
+        storage_size=Memory.from_bytes(23782357120 + 9524621312),
+        n_layers=57,
+        hidden_size=1,
+        supports_tensor=False,
+        tasks=[ModelTask.TextToImage],
+        components=[
+            ComponentInfo(
+                component_name="text_encoder",
+                component_path="text_encoder/",
+                storage_size=Memory.from_kb(0),
+                n_layers=12,
+                can_shard=False,
+                safetensors_index_filename=None,
+            ),
+            ComponentInfo(
+                component_name="text_encoder_2",
+                component_path="text_encoder_2/",
+                storage_size=Memory.from_bytes(9524621312),
+                n_layers=24,
+                can_shard=False,
+                safetensors_index_filename="model.safetensors.index.json",
+            ),
+            ComponentInfo(
+                component_name="transformer",
+                component_path="transformer/",
+                storage_size=Memory.from_bytes(23782357120),
+                n_layers=57,
+                can_shard=True,
+                safetensors_index_filename="diffusion_pytorch_model.safetensors.index.json",
+            ),
+            ComponentInfo(
+                component_name="vae",
+                component_path="vae/",
+                storage_size=Memory.from_kb(0),
+                n_layers=None,
+                can_shard=False,
+                safetensors_index_filename=None,
+            ),
+        ],
+    ),
+    "flux1-dev": ModelCard(
+        model_id=ModelId("exolabs/FLUX.1-dev"),
+        storage_size=Memory.from_bytes(23782357120 + 9524621312),
+        n_layers=57,
+        hidden_size=1,
+        supports_tensor=False,
+        tasks=[ModelTask.TextToImage],
+        components=[
+            ComponentInfo(
+                component_name="text_encoder",
+                component_path="text_encoder/",
+                storage_size=Memory.from_kb(0),
+                n_layers=12,
+                can_shard=False,
+                safetensors_index_filename=None,
+            ),
+            ComponentInfo(
+                component_name="text_encoder_2",
+                component_path="text_encoder_2/",
+                storage_size=Memory.from_bytes(9524621312),
+                n_layers=24,
+                can_shard=False,
+                safetensors_index_filename="model.safetensors.index.json",
+            ),
+            ComponentInfo(
+                component_name="transformer",
+                component_path="transformer/",
+                storage_size=Memory.from_bytes(23802816640),
+                n_layers=57,
+                can_shard=True,
+                safetensors_index_filename="diffusion_pytorch_model.safetensors.index.json",
+            ),
+            ComponentInfo(
+                component_name="vae",
+                component_path="vae/",
+                storage_size=Memory.from_kb(0),
+                n_layers=None,
+                can_shard=False,
+                safetensors_index_filename=None,
+            ),
+        ],
+    ),
+    "flux1-krea-dev": ModelCard(
+        model_id=ModelId("exolabs/FLUX.1-Krea-dev"),
+        storage_size=Memory.from_bytes(23802816640 + 9524621312),  # Same as dev
+        n_layers=57,
+        hidden_size=1,
+        supports_tensor=False,
+        tasks=[ModelTask.TextToImage],
+        components=[
+            ComponentInfo(
+                component_name="text_encoder",
+                component_path="text_encoder/",
+                storage_size=Memory.from_kb(0),
+                n_layers=12,
+                can_shard=False,
+                safetensors_index_filename=None,
+            ),
+            ComponentInfo(
+                component_name="text_encoder_2",
+                component_path="text_encoder_2/",
+                storage_size=Memory.from_bytes(9524621312),
+                n_layers=24,
+                can_shard=False,
+                safetensors_index_filename="model.safetensors.index.json",
+            ),
+            ComponentInfo(
+                component_name="transformer",
+                component_path="transformer/",
+                storage_size=Memory.from_bytes(23802816640),
+                n_layers=57,
+                can_shard=True,
+                safetensors_index_filename="diffusion_pytorch_model.safetensors.index.json",
+            ),
+            ComponentInfo(
+                component_name="vae",
+                component_path="vae/",
+                storage_size=Memory.from_kb(0),
+                n_layers=None,
+                can_shard=False,
+                safetensors_index_filename=None,
+            ),
+        ],
+    ),
+    "qwen-image": ModelCard(
+        model_id=ModelId("exolabs/Qwen-Image"),
+        storage_size=Memory.from_bytes(16584333312 + 40860802176),
+        n_layers=60,
+        hidden_size=1,
+        supports_tensor=False,
+        tasks=[ModelTask.TextToImage],
+        components=[
+            ComponentInfo(
+                component_name="text_encoder",
+                component_path="text_encoder/",
+                storage_size=Memory.from_bytes(16584333312),
+                n_layers=12,
+                can_shard=False,
+                safetensors_index_filename=None,
+            ),
+            ComponentInfo(
+                component_name="transformer",
+                component_path="transformer/",
+                storage_size=Memory.from_bytes(40860802176),
+                n_layers=60,
+                can_shard=True,
+                safetensors_index_filename="diffusion_pytorch_model.safetensors.index.json",
+            ),
+            ComponentInfo(
+                component_name="vae",
+                component_path="vae/",
+                storage_size=Memory.from_kb(0),
+                n_layers=None,
+                can_shard=False,
+                safetensors_index_filename=None,
+            ),
+        ],
+    ),
+    "qwen-image-edit-2509": ModelCard(
+        model_id=ModelId("exolabs/Qwen-Image-Edit-2509"),
+        storage_size=Memory.from_bytes(16584333312 + 40860802176),
+        n_layers=60,
+        hidden_size=1,
+        supports_tensor=False,
+        tasks=[ModelTask.ImageToImage],
+        components=[
+            ComponentInfo(
+                component_name="text_encoder",
+                component_path="text_encoder/",
+                storage_size=Memory.from_bytes(16584333312),
+                n_layers=12,
+                can_shard=False,
+                safetensors_index_filename=None,
+            ),
+            ComponentInfo(
+                component_name="transformer",
+                component_path="transformer/",
+                storage_size=Memory.from_bytes(40860802176),
+                n_layers=60,
+                can_shard=True,
+                safetensors_index_filename="diffusion_pytorch_model.safetensors.index.json",
+            ),
+            ComponentInfo(
+                component_name="vae",
+                component_path="vae/",
+                storage_size=Memory.from_kb(0),
+                n_layers=None,
+                can_shard=False,
+                safetensors_index_filename=None,
+            ),
+        ],
+    ),
+}
+
+
+def _generate_image_model_quant_variants(
+    base_name: str,
+    base_card: ModelCard,
+) -> dict[str, ModelCard]:
+    """Create quantized variants of an image model card.
+
+    Only the transformer component is quantized; text encoders stay at bf16.
+    Sizes are calculated exactly from the base card's component sizes.
+    """
+    if base_card.components is None:
+        raise ValueError(f"Image model {base_name} must have components defined")
+
+    # quantizations = [8, 6, 5, 4, 3]
+    quantizations = [8, 4]
+
+    num_transformer_bytes = next(
+        c.storage_size.in_bytes
+        for c in base_card.components
+        if c.component_name == "transformer"
+    )
+
+    transformer_bytes = Memory.from_bytes(num_transformer_bytes)
+
+    remaining_bytes = Memory.from_bytes(
+        sum(
+            c.storage_size.in_bytes
+            for c in base_card.components
+            if c.component_name != "transformer"
+        )
+    )
+
+    def with_transformer_size(new_size: Memory) -> list[ComponentInfo]:
+        assert base_card.components is not None
+        return [
+            ComponentInfo(
+                component_name=c.component_name,
+                component_path=c.component_path,
+                storage_size=new_size
+                if c.component_name == "transformer"
+                else c.storage_size,
+                n_layers=c.n_layers,
+                can_shard=c.can_shard,
+                safetensors_index_filename=c.safetensors_index_filename,
+            )
+            for c in base_card.components
+        ]
+
+    variants = {
+        base_name: ModelCard(
+            model_id=base_card.model_id,
+            storage_size=transformer_bytes + remaining_bytes,
+            n_layers=base_card.n_layers,
+            hidden_size=base_card.hidden_size,
+            supports_tensor=base_card.supports_tensor,
+            tasks=base_card.tasks,
+            components=with_transformer_size(transformer_bytes),
+        )
+    }
+
+    for quant in quantizations:
+        quant_transformer_bytes = Memory.from_bytes(
+            (num_transformer_bytes * quant) // 16
+        )
+        total_bytes = remaining_bytes + quant_transformer_bytes
+
+        model_id = ModelId(base_card.model_id + f"-{quant}bit")
+
+        variants[f"{base_name}-{quant}bit"] = ModelCard(
+            model_id=model_id,
+            storage_size=total_bytes,
+            n_layers=base_card.n_layers,
+            hidden_size=base_card.hidden_size,
+            supports_tensor=base_card.supports_tensor,
+            tasks=base_card.tasks,
+            components=with_transformer_size(quant_transformer_bytes),
+        )
+
+    return variants
+
+
+_image_model_cards: dict[str, ModelCard] = {}
+for _base_name, _base_card in _IMAGE_BASE_MODEL_CARDS.items():
+    _image_model_cards |= _generate_image_model_quant_variants(_base_name, _base_card)
+_IMAGE_MODEL_CARDS = _image_model_cards
+
+if EXO_ENABLE_IMAGE_MODELS:
+    MODEL_CARDS.update(_IMAGE_MODEL_CARDS)
+
 
 class ConfigData(BaseModel):
     model_config = {"extra": "ignore"}  # Allow unknown fields
@@ -615,7 +749,7 @@ class ConfigData(BaseModel):
 
 async def get_config_data(model_id: ModelId) -> ConfigData:
     """Downloads and parses config.json for a model."""
-    from exo.worker.download.download_utils import (
+    from exo.download.download_utils import (
         download_file_with_retry,
         ensure_models_dir,
     )
@@ -627,7 +761,7 @@ async def get_config_data(model_id: ModelId) -> ConfigData:
         "main",
         "config.json",
         target_dir,
-        lambda curr_bytes, total_bytes, is_renamed: logger.info(
+        lambda curr_bytes, total_bytes, is_renamed: logger.debug(
             f"Downloading config.json for {model_id}: {curr_bytes}/{total_bytes} ({is_renamed=})"
         ),
     )
@@ -637,11 +771,11 @@ async def get_config_data(model_id: ModelId) -> ConfigData:
 
 async def get_safetensors_size(model_id: ModelId) -> Memory:
     """Gets model size from safetensors index or falls back to HF API."""
-    from exo.shared.types.worker.downloads import ModelSafetensorsIndex
-    from exo.worker.download.download_utils import (
+    from exo.download.download_utils import (
         download_file_with_retry,
         ensure_models_dir,
     )
+    from exo.shared.types.worker.downloads import ModelSafetensorsIndex
 
     target_dir = (await ensure_models_dir()) / model_id.normalize()
     await aios.makedirs(target_dir, exist_ok=True)
@@ -650,7 +784,7 @@ async def get_safetensors_size(model_id: ModelId) -> Memory:
         "main",
         "model.safetensors.index.json",
         target_dir,
-        lambda curr_bytes, total_bytes, is_renamed: logger.info(
+        lambda curr_bytes, total_bytes, is_renamed: logger.debug(
             f"Downloading model.safetensors.index.json for {model_id}: {curr_bytes}/{total_bytes} ({is_renamed=})"
         ),
     )

src/exo/shared/topology.py

@@ -7,6 +7,11 @@ import rustworkx as rx
 from pydantic import BaseModel, ConfigDict
 
 from exo.shared.types.common import NodeId
+from exo.shared.types.profiling import (
+    InterfaceType,
+    NodeNetworkInfo,
+    ThunderboltBridgeStatus,
+)
 from exo.shared.types.topology import (
     Connection,
     Cycle,
@@ -188,24 +193,25 @@ class Topology:
             cycles.append(Cycle(node_ids=[node_id]))
         return cycles
 
-    def get_cycles_tb(self) -> list[Cycle]:
-        tb_edges = [
+    def get_rdma_cycles(self) -> list[Cycle]:
+        rdma_edges = [
             (u, v, conn)
             for u, v, conn in self._graph.weighted_edge_list()
-            if conn.is_thunderbolt()
+            if isinstance(conn, RDMAConnection)
         ]
 
-        tb_graph: rx.PyDiGraph[NodeId, SocketConnection] = rx.PyDiGraph()
-        tb_graph.add_nodes_from(self._graph.nodes())
+        rdma_graph: rx.PyDiGraph[NodeId, SocketConnection | RDMAConnection] = (
+            rx.PyDiGraph()
+        )
+        rdma_graph.add_nodes_from(self._graph.nodes())
 
-        for u, v, conn in tb_edges:
-            if isinstance(conn, SocketConnection):
-                tb_graph.add_edge(u, v, conn)
+        for u, v, conn in rdma_edges:
+            rdma_graph.add_edge(u, v, conn)
 
-        cycle_idxs = rx.simple_cycles(tb_graph)
+        cycle_idxs = rx.simple_cycles(rdma_graph)
         cycles: list[Cycle] = []
         for cycle_idx in cycle_idxs:
-            cycle = Cycle(node_ids=[tb_graph[idx] for idx in cycle_idx])
+            cycle = Cycle(node_ids=[rdma_graph[idx] for idx in cycle_idx])
             cycles.append(cycle)
 
         return cycles
@@ -219,18 +225,83 @@ class Topology:
                 topology.add_connection(connection)
         return topology
 
-    def is_thunderbolt_cycle(self, cycle: Cycle) -> bool:
+    def is_rdma_cycle(self, cycle: Cycle) -> bool:
         node_idxs = [node for node in cycle]
         rx_idxs = [self._vertex_indices[idx] for idx in node_idxs]
         for rid in rx_idxs:
             for neighbor_rid in self._graph.neighbors(rid):
                 if neighbor_rid not in rx_idxs:
                     continue
-                has_tb = False
+                has_rdma = False
                 for edge in self._graph.get_all_edge_data(rid, neighbor_rid):
-                    if edge.is_thunderbolt():
-                        has_tb = True
+                    if isinstance(edge, RDMAConnection):
+                        has_rdma = True
                         break
-                if not has_tb:
+                if not has_rdma:
                     return False
         return True
+
+    def get_thunderbolt_bridge_cycles(
+        self,
+        node_tb_bridge_status: Mapping[NodeId, ThunderboltBridgeStatus],
+        node_network: Mapping[NodeId, NodeNetworkInfo],
+    ) -> list[list[NodeId]]:
+        """
+        Find cycles in the Thunderbolt topology where all nodes have TB bridge enabled.
+        Only returns cycles with >=2 nodes (2+ machines in a loop), as
+        1 node doesn't cause the broadcast storm problem.
+        """
+        enabled_nodes = {
+            node_id
+            for node_id, status in node_tb_bridge_status.items()
+            if status.enabled
+        }
+
+        if len(enabled_nodes) < 2:
+            return []
+
+        thunderbolt_ips = _get_ips_with_interface_type(
+            enabled_nodes, node_network, "thunderbolt"
+        )
+
+        # Build subgraph with only TB bridge enabled nodes and thunderbolt connections
+        graph: rx.PyDiGraph[NodeId, SocketConnection | RDMAConnection] = rx.PyDiGraph()
+        node_to_idx: dict[NodeId, int] = {}
+
+        for node_id in enabled_nodes:
+            if node_id in self._vertex_indices:
+                node_to_idx[node_id] = graph.add_node(node_id)
+
+        for u, v, conn in self._graph.weighted_edge_list():
+            source_id, sink_id = self._graph[u], self._graph[v]
+            if source_id not in node_to_idx or sink_id not in node_to_idx:
+                continue
+            # Include connection if it's over a thunderbolt interface
+            if (
+                isinstance(conn, SocketConnection)
+                and conn.sink_multiaddr.ip_address in thunderbolt_ips
+            ):
+                graph.add_edge(node_to_idx[source_id], node_to_idx[sink_id], conn)
+            if isinstance(conn, RDMAConnection):
+                graph.add_edge(node_to_idx[source_id], node_to_idx[sink_id], conn)
+
+        return [
+            [graph[idx] for idx in cycle]
+            for cycle in rx.simple_cycles(graph)
+            if len(cycle) >= 2
+        ]
+
+
+def _get_ips_with_interface_type(
+    node_ids: set[NodeId],
+    node_network: Mapping[NodeId, NodeNetworkInfo],
+    interface_type: InterfaceType,
+) -> set[str]:
+    """Get all IP addresses on interfaces of the specified type for the given nodes."""
+    ips: set[str] = set()
+    for node_id in node_ids:
+        network_info = node_network.get(node_id, NodeNetworkInfo())
+        for iface in network_info.interfaces:
+            if iface.interface_type == interface_type:
+                ips.add(iface.ip_address)
+    return ips

src/exo/shared/types/api.py

@@ -1,22 +1,17 @@
 import time
 from collections.abc import Generator
-from typing import Annotated, Any, Literal, cast
+from typing import Annotated, Any, Literal
 
 from fastapi import UploadFile
 from pydantic import BaseModel, Field, field_validator
 from pydantic_core import PydanticUseDefault
 
-from exo.plugins.type_registry import instance_registry
 from exo.shared.models.model_cards import ModelCard, ModelId
-from exo.shared.types.common import CommandId
+from exo.shared.types.common import CommandId, NodeId
 from exo.shared.types.memory import Memory
-from exo.shared.types.worker.instances import (
-    BaseInstance,
-    Instance,
-    InstanceId,
-    InstanceMeta,
-)
-from exo.shared.types.worker.shards import Sharding
+from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
+from exo.shared.types.worker.shards import Sharding, ShardMetadata
+from exo.utils.pydantic_ext import CamelCaseModel
 
 FinishReason = Literal[
     "stop", "length", "tool_calls", "content_filter", "function_call", "error"
@@ -60,6 +55,18 @@ class ChatCompletionMessageText(BaseModel):
     text: str
 
 
+class ToolCallItem(BaseModel):
+    name: str
+    arguments: str
+
+
+class ToolCall(BaseModel):
+    id: str
+    index: int | None = None
+    type: Literal["function"] = "function"
+    function: ToolCallItem
+
+
 class ChatCompletionMessage(BaseModel):
     role: Literal["system", "user", "assistant", "developer", "tool", "function"]
     content: (
@@ -67,7 +74,7 @@ class ChatCompletionMessage(BaseModel):
     ) = None
     thinking: str | None = None  # Added for GPT-OSS harmony format support
     name: str | None = None
-    tool_calls: list[dict[str, Any]] | None = None
+    tool_calls: list[ToolCall] | None = None
     tool_call_id: str | None = None
     function_call: dict[str, Any] | None = None
 
@@ -91,6 +98,8 @@ class LogprobsContentItem(BaseModel):
 
 class Logprobs(BaseModel):
     content: list[LogprobsContentItem] | None = None
+    # This will always be null for open source models, but exists for OpenAI API
+    refusal: list[LogprobsContentItem] | None = None
 
 
 class PromptTokensDetails(BaseModel):
@@ -143,6 +152,7 @@ class GenerationStats(BaseModel):
     generation_tps: float
     prompt_tokens: int
     generation_tokens: int
+    reasoning_tokens: int = 0
     peak_memory_usage: Memory
 
 
@@ -163,6 +173,52 @@ class BenchChatCompletionResponse(ChatCompletionResponse):
     generation_stats: GenerationStats | None = None
 
 
+# Legacy Completions API types (for lm_eval compatibility)
+class CompletionLogprobs(BaseModel):
+    """Logprobs in the legacy completions format."""
+
+    tokens: list[str]
+    token_logprobs: list[float | None]
+    top_logprobs: list[dict[str, float]]
+    text_offset: list[int]
+
+
+class CompletionChoice(BaseModel):
+    text: str
+    index: int
+    logprobs: CompletionLogprobs | None = None
+    finish_reason: FinishReason | None = None
+
+
+class CompletionResponse(BaseModel):
+    id: str
+    object: Literal["text_completion"] = "text_completion"
+    created: int
+    model: str
+    choices: list[CompletionChoice]
+    usage: Usage | None = None
+
+
+class CompletionTaskParams(BaseModel):
+    """Parameters for the legacy /v1/completions endpoint."""
+
+    model: str
+    # Prompt can be: string, list of strings, list of token IDs, or list of token ID lists
+    prompt: str | list[str] | list[int] | list[list[int]]
+    max_tokens: int | None = 16
+    temperature: float | None = 1.0
+    top_p: float | None = 1.0
+    n: int | None = 1
+    stream: bool = False
+    logprobs: int | None = None
+    echo: bool = False
+    stop: str | list[str] | None = None
+    presence_penalty: float | None = None
+    frequency_penalty: float | None = None
+    seed: int | None = None
+    user: str | None = None
+
+
 class ChatCompletionTaskParams(BaseModel):
     model: str
     frequency_penalty: float | None = None
@@ -206,12 +262,6 @@ class PlaceInstanceParams(BaseModel):
 class CreateInstanceParams(BaseModel):
     instance: Instance
 
-    @field_validator("instance", mode="before")
-    @classmethod
-    def validate_instance(cls, v: Any) -> BaseInstance:  # noqa: ANN401  # pyright: ignore[reportAny]
-        """Validate instance using registry to handle both tagged and flat formats."""
-        return cast(BaseInstance, instance_registry.deserialize(v))  # pyright: ignore[reportAny]
-
 
 class PlacementPreview(BaseModel):
     model_id: ModelId
@@ -352,3 +402,16 @@ class ImageListItem(BaseModel, frozen=True):
 
 class ImageListResponse(BaseModel, frozen=True):
     data: list[ImageListItem]
+
+
+class StartDownloadParams(CamelCaseModel):
+    target_node_id: NodeId
+    shard_metadata: ShardMetadata
+
+
+class StartDownloadResponse(CamelCaseModel):
+    command_id: CommandId
+
+
+class DeleteDownloadResponse(CamelCaseModel):
+    command_id: CommandId

src/exo/shared/types/chunks.py

@@ -1,31 +1,48 @@
 from collections.abc import Generator
-from enum import Enum
 from typing import Any, Literal
 
 from exo.shared.models.model_cards import ModelId
-from exo.shared.types.api import GenerationStats, ImageGenerationStats
+from exo.shared.types.api import GenerationStats, ImageGenerationStats, TopLogprobItem
 from exo.utils.pydantic_ext import TaggedModel
 
 from .api import FinishReason
 from .common import CommandId
-
-
-class ChunkType(str, Enum):
-    Token = "Token"
-    Image = "Image"
+from .worker.runner_response import ToolCallItem
 
 
 class BaseChunk(TaggedModel):
-    idx: int
     model: ModelId
 
 
 class TokenChunk(BaseChunk):
     text: str
     token_id: int
-    finish_reason: FinishReason | None = None
+    logprob: float | None = None
+    top_logprobs: list[TopLogprobItem] | None = None
+    finish_reason: Literal["stop", "length", "content_filter"] | None = None
     stats: GenerationStats | None = None
-    error_message: str | None = None
+
+
+class ErrorChunk(BaseChunk):
+    error_message: str
+    finish_reason: Literal["error"] = "error"
+
+
+class ToolCallChunk(BaseChunk):
+    tool_calls: list[ToolCallItem]
+    finish_reason: Literal["tool_calls"] = "tool_calls"
+    stats: GenerationStats | None = None
+
+
+class CompletionChunk(BaseChunk):
+    """Chunk for legacy completions API with full logprobs for all tokens."""
+
+    text: str
+    tokens: list[str]
+    token_logprobs: list[float | None]
+    top_logprobs: list[dict[str, float]]
+    text_offset: list[int]
+    finish_reason: FinishReason | None = None
 
 
 class ImageChunk(BaseChunk):
@@ -63,4 +80,4 @@ class InputImageChunk(BaseChunk):
                 yield name, value
 
 
-GenerationChunk = TokenChunk | ImageChunk
+GenerationChunk = TokenChunk | CompletionChunk | ImageChunk | ToolCallChunk | ErrorChunk

src/exo/shared/types/commands.py

@@ -1,54 +1,45 @@
-"""Command types for exo.
+from pydantic import Field
 
-Commands are registered dynamically via the command_registry, allowing plugins
-to add their own command types without modifying this file.
-"""
-
-from typing import Any, cast
-
-from pydantic import Field, field_validator
-
-from exo.plugins.type_registry import command_registry
-from exo.shared.models.model_cards import ModelCard
+from exo.shared.models.model_cards import ModelCard, ModelId
 from exo.shared.types.api import (
     ChatCompletionTaskParams,
+    CompletionTaskParams,
     ImageEditsInternalParams,
     ImageGenerationTaskParams,
 )
 from exo.shared.types.chunks import InputImageChunk
 from exo.shared.types.common import CommandId, NodeId
 from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
-from exo.shared.types.worker.shards import Sharding
+from exo.shared.types.worker.shards import Sharding, ShardMetadata
 from exo.utils.pydantic_ext import CamelCaseModel, TaggedModel
 
 
 class BaseCommand(TaggedModel):
-    """Base class for all commands."""
-
     command_id: CommandId = Field(default_factory=CommandId)
 
 
-@command_registry.register
 class TestCommand(BaseCommand):
     __test__ = False
 
 
-@command_registry.register
 class ChatCompletion(BaseCommand):
     request_params: ChatCompletionTaskParams
 
 
-@command_registry.register
+class Completion(BaseCommand):
+    """Legacy completions API command for scoring/generation."""
+
+    request_params: CompletionTaskParams
+
+
 class ImageGeneration(BaseCommand):
     request_params: ImageGenerationTaskParams
 
 
-@command_registry.register
 class ImageEdits(BaseCommand):
     request_params: ImageEditsInternalParams
 
 
-@command_registry.register
 class PlaceInstance(BaseCommand):
     model_card: ModelCard
     sharding: Sharding
@@ -56,38 +47,46 @@ class PlaceInstance(BaseCommand):
     min_nodes: int
 
 
-@command_registry.register
 class CreateInstance(BaseCommand):
     instance: Instance
 
 
-@command_registry.register
 class DeleteInstance(BaseCommand):
     instance_id: InstanceId
 
 
-@command_registry.register
 class TaskFinished(BaseCommand):
     finished_command_id: CommandId
 
 
-@command_registry.register
 class SendInputChunk(BaseCommand):
     """Command to send an input image chunk (converted to event by master)."""
 
     chunk: InputImageChunk
 
 
-@command_registry.register
 class RequestEventLog(BaseCommand):
     since_idx: int
 
 
-# Union type for core commands - used by ForwarderCommand for network deserialization
+class StartDownload(BaseCommand):
+    target_node_id: NodeId
+    shard_metadata: ShardMetadata
+
+
+class DeleteDownload(BaseCommand):
+    target_node_id: NodeId
+    model_id: ModelId
+
+
+DownloadCommand = StartDownload | DeleteDownload
+
+
 Command = (
     TestCommand
     | RequestEventLog
     | ChatCompletion
+    | Completion
     | ImageGeneration
     | ImageEdits
     | PlaceInstance
@@ -99,14 +98,10 @@ Command = (
 
 
 class ForwarderCommand(CamelCaseModel):
-    """Wrapper for commands that includes origin node."""
-
     origin: NodeId
-    command: BaseCommand
+    command: Command
 
-    @field_validator("command", mode="before")
-    @classmethod
-    def validate_command(cls, v: Any) -> BaseCommand:  # noqa: ANN401  # pyright: ignore[reportAny]
-        """Validate command, using registry for plugin commands not in Command union."""
-        # First try the registry (handles both core and plugin commands)
-        return cast(BaseCommand, command_registry.deserialize(v))  # pyright: ignore[reportAny]
+
+class ForwarderDownloadCommand(CamelCaseModel):
+    origin: NodeId
+    command: DownloadCommand

src/exo/shared/types/events.py

@@ -1,15 +1,13 @@
 from datetime import datetime
-from typing import Any, cast
 
-from pydantic import Field, field_validator
+from pydantic import Field
 
-from exo.plugins.type_registry import event_registry, instance_registry, task_registry
 from exo.shared.topology import Connection
 from exo.shared.types.chunks import GenerationChunk, InputImageChunk
 from exo.shared.types.common import CommandId, Id, NodeId, SessionId
-from exo.shared.types.tasks import BaseTask, TaskId, TaskStatus
+from exo.shared.types.tasks import Task, TaskId, TaskStatus
 from exo.shared.types.worker.downloads import DownloadProgress
-from exo.shared.types.worker.instances import BaseInstance, InstanceId
+from exo.shared.types.worker.instances import Instance, InstanceId
 from exo.shared.types.worker.runners import RunnerId, RunnerStatus
 from exo.utils.info_gatherer.info_gatherer import GatheredInfo
 from exo.utils.pydantic_ext import CamelCaseModel, TaggedModel
@@ -27,53 +25,36 @@ class BaseEvent(TaggedModel):
     _master_time_stamp: None | datetime = None
 
 
-@event_registry.register
 class TestEvent(BaseEvent):
     __test__ = False
 
 
-@event_registry.register
 class TaskCreated(BaseEvent):
     task_id: TaskId
-    task: BaseTask
-
-    @field_validator("task", mode="before")
-    @classmethod
-    def validate_task(cls, v: Any) -> BaseTask:  # noqa: ANN401  # pyright: ignore[reportAny]
-        return cast(BaseTask, task_registry.deserialize(v))  # pyright: ignore[reportAny]
+    task: Task
 
 
-@event_registry.register
 class TaskAcknowledged(BaseEvent):
     task_id: TaskId
 
 
-@event_registry.register
 class TaskDeleted(BaseEvent):
     task_id: TaskId
 
 
-@event_registry.register
 class TaskStatusUpdated(BaseEvent):
     task_id: TaskId
     task_status: TaskStatus
 
 
-@event_registry.register
 class TaskFailed(BaseEvent):
     task_id: TaskId
     error_type: str
     error_message: str
 
 
-@event_registry.register
 class InstanceCreated(BaseEvent):
-    instance: BaseInstance
-
-    @field_validator("instance", mode="before")
-    @classmethod
-    def validate_instance(cls, v: Any) -> BaseInstance:  # noqa: ANN401  # pyright: ignore[reportAny]
-        return cast(BaseInstance, instance_registry.deserialize(v))  # pyright: ignore[reportAny]
+    instance: Instance
 
     def __eq__(self, other: object) -> bool:
         if isinstance(other, InstanceCreated):
@@ -82,63 +63,52 @@ class InstanceCreated(BaseEvent):
         return False
 
 
-@event_registry.register
 class InstanceDeleted(BaseEvent):
     instance_id: InstanceId
 
 
-@event_registry.register
 class RunnerStatusUpdated(BaseEvent):
     runner_id: RunnerId
     runner_status: RunnerStatus
 
 
-@event_registry.register
 class RunnerDeleted(BaseEvent):
     runner_id: RunnerId
 
 
-@event_registry.register
 class NodeTimedOut(BaseEvent):
     node_id: NodeId
 
 
 # TODO: bikeshed this name
-@event_registry.register
 class NodeGatheredInfo(BaseEvent):
     node_id: NodeId
     when: str  # this is a manually cast datetime overrode by the master when the event is indexed, rather than the local time on the device
     info: GatheredInfo
 
 
-@event_registry.register
 class NodeDownloadProgress(BaseEvent):
     download_progress: DownloadProgress
 
 
-@event_registry.register
 class ChunkGenerated(BaseEvent):
     command_id: CommandId
     chunk: GenerationChunk
 
 
-@event_registry.register
 class InputChunkReceived(BaseEvent):
     command_id: CommandId
     chunk: InputImageChunk
 
 
-@event_registry.register
 class TopologyEdgeCreated(BaseEvent):
     conn: Connection
 
 
-@event_registry.register
 class TopologyEdgeDeleted(BaseEvent):
     conn: Connection
 
 
-# Union type for Pydantic validation - tries each type in order
 Event = (
     TestEvent
     | TaskCreated
@@ -164,12 +134,7 @@ class IndexedEvent(CamelCaseModel):
     """An event indexed by the master, with a globally unique index"""
 
     idx: int = Field(ge=0)
-    event: BaseEvent
-
-    @field_validator("event", mode="before")
-    @classmethod
-    def validate_event(cls, v: Any) -> BaseEvent:  # noqa: ANN401  # pyright: ignore[reportAny]
-        return cast(BaseEvent, event_registry.deserialize(v))  # pyright: ignore[reportAny]
+    event: Event
 
 
 class ForwarderEvent(CamelCaseModel):
@@ -178,9 +143,4 @@ class ForwarderEvent(CamelCaseModel):
     origin_idx: int = Field(ge=0)
     origin: NodeId
     session: SessionId
-    event: BaseEvent
-
-    @field_validator("event", mode="before")
-    @classmethod
-    def validate_event(cls, v: Any) -> BaseEvent:  # noqa: ANN401  # pyright: ignore[reportAny]
-        return cast(BaseEvent, event_registry.deserialize(v))  # pyright: ignore[reportAny]
+    event: Event

src/exo/shared/types/mlx.py

@@ -0,0 +1,12 @@
+"""Shared types for MLX-related functionality."""
+
+from collections.abc import Sequence
+
+from mlx_lm.models.cache import (
+    KVCache,
+    QuantizedKVCache,
+    RotatingKVCache,
+)
+
+# This list contains one cache entry per transformer layer
+KVCacheType = Sequence[KVCache | RotatingKVCache | QuantizedKVCache]

src/exo/shared/types/profiling.py

@@ -1,5 +1,5 @@
 from collections.abc import Sequence
-from typing import Self
+from typing import Literal, Self
 
 import psutil
 
@@ -48,9 +48,13 @@ class SystemPerformanceProfile(CamelCaseModel):
     ecpu_usage: float = 0.0
 
 
+InterfaceType = Literal["wifi", "ethernet", "maybe_ethernet", "thunderbolt", "unknown"]
+
+
 class NetworkInterfaceInfo(CamelCaseModel):
     name: str
     ip_address: str
+    interface_type: InterfaceType = "unknown"
 
 
 class NodeIdentity(CamelCaseModel):
@@ -71,3 +75,11 @@ class NodeThunderboltInfo(CamelCaseModel):
     """Thunderbolt interface identifiers for a node."""
 
     interfaces: Sequence[ThunderboltIdentifier] = []
+
+
+class ThunderboltBridgeStatus(CamelCaseModel):
+    """Whether the Thunderbolt Bridge network service is enabled on this node."""
+
+    enabled: bool
+    exists: bool
+    service_name: str | None = None

src/exo/shared/types/state.py

@@ -13,10 +13,11 @@ from exo.shared.types.profiling import (
     NodeNetworkInfo,
     NodeThunderboltInfo,
     SystemPerformanceProfile,
+    ThunderboltBridgeStatus,
 )
-from exo.shared.types.tasks import BaseTask, TaskId
+from exo.shared.types.tasks import Task, TaskId
 from exo.shared.types.worker.downloads import DownloadProgress
-from exo.shared.types.worker.instances import BaseInstance, InstanceId
+from exo.shared.types.worker.instances import Instance, InstanceId
 from exo.shared.types.worker.runners import RunnerId, RunnerStatus
 from exo.utils.pydantic_ext import CamelCaseModel
 
@@ -37,10 +38,10 @@ class State(CamelCaseModel):
         strict=True,
         arbitrary_types_allowed=True,
     )
-    instances: Mapping[InstanceId, BaseInstance] = {}
+    instances: Mapping[InstanceId, Instance] = {}
     runners: Mapping[RunnerId, RunnerStatus] = {}
     downloads: Mapping[NodeId, Sequence[DownloadProgress]] = {}
-    tasks: Mapping[TaskId, BaseTask] = {}
+    tasks: Mapping[TaskId, Task] = {}
     last_seen: Mapping[NodeId, datetime] = {}
     topology: Topology = Field(default_factory=Topology)
     last_event_applied_idx: int = Field(default=-1, ge=-1)
@@ -51,16 +52,10 @@ class State(CamelCaseModel):
     node_system: Mapping[NodeId, SystemPerformanceProfile] = {}
     node_network: Mapping[NodeId, NodeNetworkInfo] = {}
     node_thunderbolt: Mapping[NodeId, NodeThunderboltInfo] = {}
+    node_thunderbolt_bridge: Mapping[NodeId, ThunderboltBridgeStatus] = {}
 
-    @field_serializer("instances", mode="plain")
-    def _encode_instances(
-        self, value: Mapping[InstanceId, BaseInstance]
-    ) -> dict[str, Any]:
-        """Serialize instances with full subclass fields."""
-        return {
-            str(k): v.model_dump(by_alias=True, serialize_as_any=True)
-            for k, v in value.items()
-        }
+    # Detected cycles where all nodes have Thunderbolt bridge enabled (>2 nodes)
+    thunderbolt_bridge_cycles: Sequence[Sequence[NodeId]] = []
 
     @field_serializer("topology", mode="plain")
     def _encode_topology(self, value: Topology) -> TopologySnapshot:

src/exo/shared/types/tasks.py

@@ -2,9 +2,9 @@ from enum import Enum
 
 from pydantic import Field
 
-from exo.plugins.type_registry import task_registry
 from exo.shared.types.api import (
     ChatCompletionTaskParams,
+    CompletionTaskParams,
     ImageEditsInternalParams,
     ImageGenerationTaskParams,
 )
@@ -33,32 +33,26 @@ class BaseTask(TaggedModel):
     instance_id: InstanceId
 
 
-@task_registry.register
 class CreateRunner(BaseTask):  # emitted by Worker
     bound_instance: BoundInstance
 
 
-@task_registry.register
 class DownloadModel(BaseTask):  # emitted by Worker
     shard_metadata: ShardMetadata
 
 
-@task_registry.register
 class LoadModel(BaseTask):  # emitted by Worker
     pass
 
 
-@task_registry.register
 class ConnectToGroup(BaseTask):  # emitted by Worker
     pass
 
 
-@task_registry.register
 class StartWarmup(BaseTask):  # emitted by Worker
     pass
 
 
-@task_registry.register
 class ChatCompletion(BaseTask):  # emitted by Master
     command_id: CommandId
     task_params: ChatCompletionTaskParams
@@ -67,7 +61,16 @@ class ChatCompletion(BaseTask):  # emitted by Master
     error_message: str | None = Field(default=None)
 
 
-@task_registry.register
+class Completion(BaseTask):
+    """Legacy completions task for scoring tokens with echo=True."""
+
+    command_id: CommandId
+    task_params: CompletionTaskParams
+
+    error_type: str | None = Field(default=None)
+    error_message: str | None = Field(default=None)
+
+
 class ImageGeneration(BaseTask):  # emitted by Master
     command_id: CommandId
     task_params: ImageGenerationTaskParams
@@ -76,7 +79,6 @@ class ImageGeneration(BaseTask):  # emitted by Master
     error_message: str | None = Field(default=None)
 
 
-@task_registry.register
 class ImageEdits(BaseTask):  # emitted by Master
     command_id: CommandId
     task_params: ImageEditsInternalParams
@@ -85,12 +87,10 @@ class ImageEdits(BaseTask):  # emitted by Master
     error_message: str | None = Field(default=None)
 
 
-@task_registry.register
 class Shutdown(BaseTask):  # emitted by Worker
     runner_id: RunnerId
 
 
-# Union type for Pydantic validation - tries each type in order
 Task = (
     CreateRunner
     | DownloadModel
@@ -98,6 +98,7 @@ Task = (
     | LoadModel
     | StartWarmup
     | ChatCompletion
+    | Completion
     | ImageGeneration
     | ImageEdits
     | Shutdown

src/exo/shared/types/topology.py

@@ -21,9 +21,6 @@ class RDMAConnection(FrozenModel):
     source_rdma_iface: str
     sink_rdma_iface: str
 
-    def is_thunderbolt(self) -> bool:
-        return True
-
 
 class SocketConnection(FrozenModel):
     sink_multiaddr: Multiaddr
@@ -31,9 +28,6 @@ class SocketConnection(FrozenModel):
     def __hash__(self):
         return hash(self.sink_multiaddr.ip_address)
 
-    def is_thunderbolt(self) -> bool:
-        return str(self.sink_multiaddr.ipv4_address).startswith("169.254")
-
 
 class Connection(FrozenModel):
     source: NodeId

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

@@ -1,15 +1,7 @@
-"""Instance types for exo.
-
-Instances are registered dynamically via the instance_registry, allowing plugins
-to add their own instance types without modifying this file.
-"""
-
 from enum import Enum
-from typing import Any, cast
 
-from pydantic import field_validator, model_validator
+from pydantic import model_validator
 
-from exo.plugins.type_registry import instance_registry
 from exo.shared.types.common import Host, Id, NodeId
 from exo.shared.types.worker.runners import RunnerId, ShardAssignments, ShardMetadata
 from exo.utils.pydantic_ext import CamelCaseModel, TaggedModel
@@ -25,8 +17,6 @@ class InstanceMeta(str, Enum):
 
 
 class BaseInstance(TaggedModel):
-    """Base class for all instance types."""
-
     instance_id: InstanceId
     shard_assignments: ShardAssignments
 
@@ -34,36 +24,25 @@ class BaseInstance(TaggedModel):
         return self.shard_assignments.runner_to_shard.get(runner_id, None)
 
 
-@instance_registry.register
 class MlxRingInstance(BaseInstance):
     hosts_by_node: dict[NodeId, list[Host]]
     ephemeral_port: int
 
 
-@instance_registry.register
 class MlxJacclInstance(BaseInstance):
     jaccl_devices: list[list[str | None]]
     jaccl_coordinators: dict[NodeId, str]
 
 
-# Union type for Pydantic validation - tries each type in order
-# This is used by API endpoints (dashboard) which send flat format
+# TODO: Single node instance
 Instance = MlxRingInstance | MlxJacclInstance
 
 
 class BoundInstance(CamelCaseModel):
-    """An instance bound to a specific runner on a specific node."""
-
-    instance: BaseInstance
+    instance: Instance
     bound_runner_id: RunnerId
     bound_node_id: NodeId
 
-    @field_validator("instance", mode="before")
-    @classmethod
-    def validate_instance(cls, v: Any) -> BaseInstance:  # noqa: ANN401  # pyright: ignore[reportAny]
-        """Validate instance using registry to handle both tagged and flat formats."""
-        return cast(BaseInstance, instance_registry.deserialize(v))  # pyright: ignore[reportAny]
-
     @property
     def bound_shard(self) -> ShardMetadata:
         shard = self.instance.shard(self.bound_runner_id)

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

@@ -1,7 +1,13 @@
 from collections.abc import Generator
 from typing import Any, Literal
 
-from exo.shared.types.api import FinishReason, GenerationStats, ImageGenerationStats
+from exo.shared.types.api import (
+    FinishReason,
+    GenerationStats,
+    ImageGenerationStats,
+    ToolCallItem,
+    TopLogprobItem,
+)
 from exo.utils.pydantic_ext import TaggedModel
 
 
@@ -9,14 +15,11 @@ class BaseRunnerResponse(TaggedModel):
     pass
 
 
-class TokenizedResponse(BaseRunnerResponse):
-    prompt_tokens: int
-
-
 class GenerationResponse(BaseRunnerResponse):
     text: str
     token: int
-    # logprobs: list[float] | None = None # too big. we can change to be top-k
+    logprob: float | None = None
+    top_logprobs: list[TopLogprobItem] | None = None
     finish_reason: FinishReason | None = None
     stats: GenerationStats | None = None
 
@@ -25,6 +28,7 @@ class ImageGenerationResponse(BaseRunnerResponse):
     image_data: bytes
     format: Literal["png", "jpeg", "webp"] = "png"
     stats: ImageGenerationStats | None = None
+    image_index: int = 0
 
     def __repr_args__(self) -> Generator[tuple[str, Any], None, None]:
         for name, value in super().__repr_args__():  # pyright: ignore[reportAny]
@@ -39,6 +43,7 @@ class PartialImageResponse(BaseRunnerResponse):
     format: Literal["png", "jpeg", "webp"] = "png"
     partial_index: int
     total_partials: int
+    image_index: int = 0
 
     def __repr_args__(self) -> Generator[tuple[str, Any], None, None]:
         for name, value in super().__repr_args__():  # pyright: ignore[reportAny]
@@ -48,5 +53,9 @@ class PartialImageResponse(BaseRunnerResponse):
                 yield name, value
 
 
+class ToolCallResponse(BaseRunnerResponse):
+    tool_calls: list[ToolCallItem]
+
+
 class FinishedResponse(BaseRunnerResponse):
     pass

src/exo/utils/channels.py

@@ -194,6 +194,22 @@ class MpReceiver[T]:
                 raise EndOfStream from None
             return item
 
+    def receive_with_timeout(self, timeout: float) -> T | None:
+        """Receive with timeout, returns None if no message within timeout."""
+        if self._state.closed.is_set():
+            raise ClosedResourceError
+
+        try:
+            item = self._state.buffer.get(block=True, timeout=timeout)
+            if isinstance(item, _MpEndOfStream):
+                self.close()
+                raise EndOfStream
+            return item
+        except Empty:
+            return None
+        except ValueError as e:
+            raise ClosedResourceError from e
+
     # nb: this function will not cancel particularly well
     async def receive_async(self) -> T:
         return await to_thread.run_sync(self.receive, limiter=CapacityLimiter(1))

src/exo/utils/info_gatherer/info_gatherer.py

@@ -19,6 +19,7 @@ from exo.shared.types.memory import Memory
 from exo.shared.types.profiling import (
     MemoryUsage,
     NetworkInterfaceInfo,
+    ThunderboltBridgeStatus,
 )
 from exo.shared.types.thunderbolt import (
     ThunderboltConnection,
@@ -34,6 +35,142 @@ from .system_info import get_friendly_name, get_model_and_chip, get_network_inte
 IS_DARWIN = sys.platform == "darwin"
 
 
+async def _get_thunderbolt_devices() -> set[str] | None:
+    """Get Thunderbolt interface device names (e.g., en2, en3) from hardware ports.
+
+    Returns None if the networksetup command fails.
+    """
+    result = await anyio.run_process(
+        ["networksetup", "-listallhardwareports"],
+        check=False,
+    )
+    if result.returncode != 0:
+        logger.warning(
+            f"networksetup -listallhardwareports failed with code "
+            f"{result.returncode}: {result.stderr.decode()}"
+        )
+        return None
+
+    output = result.stdout.decode()
+    thunderbolt_devices: set[str] = set()
+    current_port: str | None = None
+
+    for line in output.splitlines():
+        line = line.strip()
+        if line.startswith("Hardware Port:"):
+            current_port = line.split(":", 1)[1].strip()
+        elif line.startswith("Device:") and current_port:
+            device = line.split(":", 1)[1].strip()
+            if "thunderbolt" in current_port.lower():
+                thunderbolt_devices.add(device)
+            current_port = None
+
+    return thunderbolt_devices
+
+
+async def _get_bridge_services() -> dict[str, str] | None:
+    """Get mapping of bridge device -> service name from network service order.
+
+    Returns None if the networksetup command fails.
+    """
+    result = await anyio.run_process(
+        ["networksetup", "-listnetworkserviceorder"],
+        check=False,
+    )
+    if result.returncode != 0:
+        logger.warning(
+            f"networksetup -listnetworkserviceorder failed with code "
+            f"{result.returncode}: {result.stderr.decode()}"
+        )
+        return None
+
+    # Parse service order to find bridge devices and their service names
+    # Format: "(1) Service Name\n(Hardware Port: ..., Device: bridge0)\n"
+    service_order_output = result.stdout.decode()
+    bridge_services: dict[str, str] = {}  # device -> service name
+    current_service: str | None = None
+
+    for line in service_order_output.splitlines():
+        line = line.strip()
+        # Match "(N) Service Name" or "(*) Service Name" (disabled)
+        # but NOT "(Hardware Port: ...)" lines
+        if (
+            line
+            and line.startswith("(")
+            and ")" in line
+            and not line.startswith("(Hardware Port:")
+        ):
+            paren_end = line.index(")")
+            if paren_end + 2 <= len(line):
+                current_service = line[paren_end + 2 :]
+        # Match "(Hardware Port: ..., Device: bridgeX)"
+        elif current_service and "Device: bridge" in line:
+            # Extract device name from "..., Device: bridge0)"
+            device_start = line.find("Device: ") + len("Device: ")
+            device_end = line.find(")", device_start)
+            if device_end > device_start:
+                device = line[device_start:device_end]
+                bridge_services[device] = current_service
+
+    return bridge_services
+
+
+async def _get_bridge_members(bridge_device: str) -> set[str]:
+    """Get member interfaces of a bridge device via ifconfig."""
+    result = await anyio.run_process(
+        ["ifconfig", bridge_device],
+        check=False,
+    )
+    if result.returncode != 0:
+        logger.debug(f"ifconfig {bridge_device} failed with code {result.returncode}")
+        return set()
+
+    members: set[str] = set()
+    ifconfig_output = result.stdout.decode()
+    for line in ifconfig_output.splitlines():
+        line = line.strip()
+        if line.startswith("member:"):
+            parts = line.split()
+            if len(parts) > 1:
+                members.add(parts[1])
+
+    return members
+
+
+async def _find_thunderbolt_bridge(
+    bridge_services: dict[str, str], thunderbolt_devices: set[str]
+) -> str | None:
+    """Find the service name of a bridge containing Thunderbolt interfaces.
+
+    Returns the service name if found, None otherwise.
+    """
+    for bridge_device, service_name in bridge_services.items():
+        members = await _get_bridge_members(bridge_device)
+        if members & thunderbolt_devices:  # intersection is non-empty
+            return service_name
+    return None
+
+
+async def _is_service_enabled(service_name: str) -> bool | None:
+    """Check if a network service is enabled.
+
+    Returns True if enabled, False if disabled, None on error.
+    """
+    result = await anyio.run_process(
+        ["networksetup", "-getnetworkserviceenabled", service_name],
+        check=False,
+    )
+    if result.returncode != 0:
+        logger.warning(
+            f"networksetup -getnetworkserviceenabled '{service_name}' "
+            f"failed with code {result.returncode}: {result.stderr.decode()}"
+        )
+        return None
+
+    stdout = result.stdout.decode().strip().lower()
+    return stdout == "enabled"
+
+
 class StaticNodeInformation(TaggedModel):
     """Node information that should NEVER change, to be gathered once at startup"""
 
@@ -58,6 +195,66 @@ class MacThunderboltConnections(TaggedModel):
     conns: Sequence[ThunderboltConnection]
 
 
+class ThunderboltBridgeInfo(TaggedModel):
+    status: ThunderboltBridgeStatus
+
+    @classmethod
+    async def gather(cls) -> Self | None:
+        """Check if a Thunderbolt Bridge network service is enabled on this node.
+
+        Detection approach:
+        1. Find all Thunderbolt interface devices (en2, en3, etc.) from hardware ports
+        2. Find bridge devices from network service order (not hardware ports, as
+           bridges may not appear there)
+        3. Check each bridge's members via ifconfig
+        4. If a bridge contains Thunderbolt interfaces, it's a Thunderbolt Bridge
+        5. Check if that network service is enabled
+        """
+        if not IS_DARWIN:
+            return None
+
+        def _no_bridge_status() -> Self:
+            return cls(
+                status=ThunderboltBridgeStatus(
+                    enabled=False, exists=False, service_name=None
+                )
+            )
+
+        try:
+            tb_devices = await _get_thunderbolt_devices()
+            if tb_devices is None:
+                return _no_bridge_status()
+
+            bridge_services = await _get_bridge_services()
+            if not bridge_services:
+                return _no_bridge_status()
+
+            tb_service_name = await _find_thunderbolt_bridge(
+                bridge_services, tb_devices
+            )
+            if not tb_service_name:
+                return _no_bridge_status()
+
+            enabled = await _is_service_enabled(tb_service_name)
+            if enabled is None:
+                return cls(
+                    status=ThunderboltBridgeStatus(
+                        enabled=False, exists=True, service_name=tb_service_name
+                    )
+                )
+
+            return cls(
+                status=ThunderboltBridgeStatus(
+                    enabled=enabled,
+                    exists=True,
+                    service_name=tb_service_name,
+                )
+            )
+        except Exception as e:
+            logger.warning(f"Failed to gather Thunderbolt Bridge info: {e}")
+            return None
+
+
 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"""
 
@@ -111,6 +308,7 @@ GatheredInfo = (
     | NodeNetworkInterfaces
     | MacThunderboltIdentifiers
     | MacThunderboltConnections
+    | ThunderboltBridgeInfo
     | NodeConfig
     | MiscData
     | StaticNodeInformation
@@ -125,6 +323,7 @@ class InfoGatherer:
     system_profiler_interval: float | None = 5 if IS_DARWIN else None
     memory_poll_rate: float | None = None if IS_DARWIN else 1
     macmon_interval: float | None = 1 if IS_DARWIN else None
+    thunderbolt_bridge_poll_interval: float | None = 10 if IS_DARWIN else None
     _tg: TaskGroup = field(init=False, default_factory=create_task_group)
 
     async def run(self):
@@ -133,6 +332,7 @@ class InfoGatherer:
                 if (macmon_path := shutil.which("macmon")) is not None:
                     tg.start_soon(self._monitor_macmon, macmon_path)
                 tg.start_soon(self._monitor_system_profiler_thunderbolt_data)
+                tg.start_soon(self._monitor_thunderbolt_bridge_status)
             tg.start_soon(self._watch_system_info)
             tg.start_soon(self._monitor_memory_usage)
             tg.start_soon(self._monitor_misc)
@@ -149,13 +349,8 @@ class InfoGatherer:
     async def _monitor_misc(self):
         if self.misc_poll_interval is None:
             return
-        prev = await MiscData.gather()
-        await self.info_sender.send(prev)
         while True:
-            curr = await MiscData.gather()
-            if prev != curr:
-                prev = curr
-                await self.info_sender.send(curr)
+            await self.info_sender.send(await MiscData.gather())
             await anyio.sleep(self.misc_poll_interval)
 
     async def _monitor_system_profiler_thunderbolt_data(self):
@@ -165,15 +360,12 @@ class InfoGatherer:
         if iface_map is None:
             return
 
-        old_idents = []
         while True:
             data = await ThunderboltConnectivity.gather()
             assert data is not None
 
             idents = [it for i in data if (it := i.ident(iface_map)) is not None]
-            if idents != old_idents:
-                await self.info_sender.send(MacThunderboltIdentifiers(idents=idents))
-            old_idents = idents
+            await self.info_sender.send(MacThunderboltIdentifiers(idents=idents))
 
             conns = [it for i in data if (it := i.conn()) is not None]
             await self.info_sender.send(MacThunderboltConnections(conns=conns))
@@ -198,14 +390,20 @@ class InfoGatherer:
     async def _watch_system_info(self):
         if self.interface_watcher_interval is None:
             return
-        old_nics = []
         while True:
-            nics = get_network_interfaces()
-            if nics != old_nics:
-                old_nics = nics
-                await self.info_sender.send(NodeNetworkInterfaces(ifaces=nics))
+            nics = await get_network_interfaces()
+            await self.info_sender.send(NodeNetworkInterfaces(ifaces=nics))
             await anyio.sleep(self.interface_watcher_interval)
 
+    async def _monitor_thunderbolt_bridge_status(self):
+        if self.thunderbolt_bridge_poll_interval is None:
+            return
+        while True:
+            curr = await ThunderboltBridgeInfo.gather()
+            if curr is not None:
+                await self.info_sender.send(curr)
+            await anyio.sleep(self.thunderbolt_bridge_poll_interval)
+
     async def _monitor_macmon(self, macmon_path: str):
         if self.macmon_interval is None:
             return

src/exo/utils/info_gatherer/system_info.py

@@ -5,7 +5,7 @@ from subprocess import CalledProcessError
 import psutil
 from anyio import run_process
 
-from exo.shared.types.profiling import NetworkInterfaceInfo
+from exo.shared.types.profiling import InterfaceType, NetworkInterfaceInfo
 
 
 async def get_friendly_name() -> str:
@@ -16,8 +16,7 @@ async def get_friendly_name() -> str:
     """
     hostname = socket.gethostname()
 
-    # TODO: better non mac support
-    if sys.platform != "darwin":  # 'darwin' is the platform name for macOS
+    if sys.platform != "darwin":
         return hostname
 
     try:
@@ -28,7 +27,41 @@ async def get_friendly_name() -> str:
     return process.stdout.decode("utf-8", errors="replace").strip() or hostname
 
 
-def get_network_interfaces() -> list[NetworkInterfaceInfo]:
+async def _get_interface_types_from_networksetup() -> dict[str, InterfaceType]:
+    """Parse networksetup -listallhardwareports to get interface types."""
+    if sys.platform != "darwin":
+        return {}
+
+    try:
+        result = await run_process(["networksetup", "-listallhardwareports"])
+    except CalledProcessError:
+        return {}
+
+    types: dict[str, InterfaceType] = {}
+    current_type: InterfaceType = "unknown"
+
+    for line in result.stdout.decode().splitlines():
+        if line.startswith("Hardware Port:"):
+            port_name = line.split(":", 1)[1].strip()
+            if "Wi-Fi" in port_name:
+                current_type = "wifi"
+            elif "Ethernet" in port_name or "LAN" in port_name:
+                current_type = "ethernet"
+            elif port_name.startswith("Thunderbolt"):
+                current_type = "thunderbolt"
+            else:
+                current_type = "unknown"
+        elif line.startswith("Device:"):
+            device = line.split(":", 1)[1].strip()
+            # enX is ethernet adapters or thunderbolt - these must be deprioritised
+            if device.startswith("en") and device not in ["en0", "en1"]:
+                current_type = "maybe_ethernet"
+            types[device] = current_type
+
+    return types
+
+
+async def get_network_interfaces() -> list[NetworkInterfaceInfo]:
     """
     Retrieves detailed network interface information on macOS.
     Parses output from 'networksetup -listallhardwareports' and 'ifconfig'
@@ -36,13 +69,18 @@ def get_network_interfaces() -> list[NetworkInterfaceInfo]:
     Returns a list of NetworkInterfaceInfo objects.
     """
     interfaces_info: list[NetworkInterfaceInfo] = []
+    interface_types = await _get_interface_types_from_networksetup()
 
     for iface, services in psutil.net_if_addrs().items():
         for service in services:
             match service.family:
                 case socket.AF_INET | socket.AF_INET6:
                     interfaces_info.append(
-                        NetworkInterfaceInfo(name=iface, ip_address=service.address)
+                        NetworkInterfaceInfo(
+                            name=iface,
+                            ip_address=service.address,
+                            interface_type=interface_types.get(iface, "unknown"),
+                        )
                     )
                 case _:
                     pass

src/exo/utils/keyed_backoff.py

@@ -0,0 +1,32 @@
+import time
+from typing import Generic, TypeVar
+
+K = TypeVar("K")
+
+
+class KeyedBackoff(Generic[K]):
+    """Tracks exponential backoff state per key."""
+
+    def __init__(self, base: float = 0.5, cap: float = 10.0):
+        self._base = base
+        self._cap = cap
+        self._attempts: dict[K, int] = {}
+        self._last_time: dict[K, float] = {}
+
+    def should_proceed(self, key: K) -> bool:
+        """Returns True if enough time has elapsed since last attempt."""
+        now = time.monotonic()
+        last = self._last_time.get(key, 0.0)
+        attempts = self._attempts.get(key, 0)
+        delay = min(self._cap, self._base * (2.0**attempts))
+        return now - last >= delay
+
+    def record_attempt(self, key: K) -> None:
+        """Record that an attempt was made for this key."""
+        self._last_time[key] = time.monotonic()
+        self._attempts[key] = self._attempts.get(key, 0) + 1
+
+    def reset(self, key: K) -> None:
+        """Reset backoff state for a key (e.g., on success)."""
+        self._attempts.pop(key, None)
+        self._last_time.pop(key, None)

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

@@ -6,10 +6,10 @@ import mlx.core as mx
 from mflux.models.common.config.config import Config
 from PIL import Image
 
+from exo.download.download_utils import build_model_path
 from exo.shared.types.api import AdvancedImageParams
 from exo.shared.types.worker.instances import BoundInstance
 from exo.shared.types.worker.shards import PipelineShardMetadata
-from exo.worker.download.download_utils import build_model_path
 from exo.worker.engines.image.config import ImageModelConfig
 from exo.worker.engines.image.models import (
     create_adapter_for_model,
@@ -140,6 +140,7 @@ class DistributedImageModel:
             width=width,
             image_path=image_path,
             model_config=self._adapter.model.model_config,  # pyright: ignore[reportAny]
+            guidance=guidance_override if guidance_override is not None else 4.0,
         )
 
         num_sync_steps = self._config.get_num_sync_steps(steps)

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

@@ -75,19 +75,20 @@ def generate_image(
     intermediate images, then ImageGenerationResponse for the final image.
 
     Yields:
-        PartialImageResponse for intermediate images (if partial_images > 0)
-        ImageGenerationResponse for the final complete image
+        PartialImageResponse for intermediate images (if partial_images > 0, first image only)
+        ImageGenerationResponse for final complete images
     """
     width, height = parse_size(task.size)
     quality: Literal["low", "medium", "high"] = task.quality or "medium"
 
     advanced_params = task.advanced_params
     if advanced_params is not None and advanced_params.seed is not None:
-        seed = advanced_params.seed
+        base_seed = advanced_params.seed
     else:
-        seed = random.randint(0, 2**32 - 1)
+        base_seed = random.randint(0, 2**32 - 1)
 
     is_bench = getattr(task, "bench", False)
+    num_images = task.n or 1
 
     generation_start_time: float = 0.0
 
@@ -95,7 +96,11 @@ def generate_image(
         mx.reset_peak_memory()
         generation_start_time = time.perf_counter()
 
-    partial_images = task.partial_images or (3 if task.stream else 0)
+    partial_images = (
+        task.partial_images
+        if task.partial_images is not None
+        else (3 if task.stream else 0)
+    )
 
     image_path: Path | None = None
 
@@ -105,72 +110,81 @@ def generate_image(
             image_path = Path(tmpdir) / "input.png"
             image_path.write_bytes(base64.b64decode(task.image_data))
 
-        # Iterate over generator results
-        for result in model.generate(
-            prompt=task.prompt,
-            height=height,
-            width=width,
-            quality=quality,
-            seed=seed,
-            image_path=image_path,
-            partial_images=partial_images,
-            advanced_params=advanced_params,
-        ):
-            if isinstance(result, tuple):
-                # Partial image: (Image, partial_index, total_partials)
-                image, partial_idx, total_partials = result
-                buffer = io.BytesIO()
-                image_format = task.output_format.upper()
-                if image_format == "JPG":
-                    image_format = "JPEG"
-                if image_format == "JPEG" and image.mode == "RGBA":
-                    image = image.convert("RGB")
-                image.save(buffer, format=image_format)
+        for image_num in range(num_images):
+            # Increment seed for each image to ensure unique results
+            current_seed = base_seed + image_num
 
-                yield PartialImageResponse(
-                    image_data=buffer.getvalue(),
-                    format=task.output_format,
-                    partial_index=partial_idx,
-                    total_partials=total_partials,
-                )
-            else:
-                image = result
+            for result in model.generate(
+                prompt=task.prompt,
+                height=height,
+                width=width,
+                quality=quality,
+                seed=current_seed,
+                image_path=image_path,
+                partial_images=partial_images,
+                advanced_params=advanced_params,
+            ):
+                if isinstance(result, tuple):
+                    # Partial image: (Image, partial_index, total_partials)
+                    image, partial_idx, total_partials = result
+                    buffer = io.BytesIO()
+                    image_format = task.output_format.upper()
+                    if image_format == "JPG":
+                        image_format = "JPEG"
+                    if image_format == "JPEG" and image.mode == "RGBA":
+                        image = image.convert("RGB")
+                    image.save(buffer, format=image_format)
 
-                stats: ImageGenerationStats | None = None
-                if is_bench:
-                    generation_end_time = time.perf_counter()
-                    total_generation_time = generation_end_time - generation_start_time
-
-                    num_inference_steps = model.get_steps_for_quality(quality)
-
-                    seconds_per_step = (
-                        total_generation_time / num_inference_steps
-                        if num_inference_steps > 0
-                        else 0.0
+                    yield PartialImageResponse(
+                        image_data=buffer.getvalue(),
+                        format=task.output_format,
+                        partial_index=partial_idx,
+                        total_partials=total_partials,
+                        image_index=image_num,
                     )
+                else:
+                    image = result
 
-                    peak_memory_gb = mx.get_peak_memory() / (1024**3)
+                    # Only include stats on the final image
+                    stats: ImageGenerationStats | None = None
+                    if is_bench and image_num == num_images - 1:
+                        generation_end_time = time.perf_counter()
+                        total_generation_time = (
+                            generation_end_time - generation_start_time
+                        )
 
-                    stats = ImageGenerationStats(
-                        seconds_per_step=seconds_per_step,
-                        total_generation_time=total_generation_time,
-                        num_inference_steps=num_inference_steps,
-                        num_images=task.n or 1,
-                        image_width=width,
-                        image_height=height,
-                        peak_memory_usage=Memory.from_gb(peak_memory_gb),
+                        num_inference_steps = model.get_steps_for_quality(quality)
+                        total_steps = num_inference_steps * num_images
+
+                        seconds_per_step = (
+                            total_generation_time / total_steps
+                            if total_steps > 0
+                            else 0.0
+                        )
+
+                        peak_memory_gb = mx.get_peak_memory() / (1024**3)
+
+                        stats = ImageGenerationStats(
+                            seconds_per_step=seconds_per_step,
+                            total_generation_time=total_generation_time,
+                            num_inference_steps=num_inference_steps,
+                            num_images=num_images,
+                            image_width=width,
+                            image_height=height,
+                            peak_memory_usage=Memory.from_gb(peak_memory_gb),
+                        )
+
+                    buffer = io.BytesIO()
+                    image_format = task.output_format.upper()
+                    if image_format == "JPG":
+                        image_format = "JPEG"
+                    if image_format == "JPEG" and image.mode == "RGBA":
+                        image = image.convert("RGB")
+                    image.save(buffer, format=image_format)
+
+                    yield ImageGenerationResponse(
+                        image_data=buffer.getvalue(),
+                        format=task.output_format,
+                        stats=stats,
+                        image_index=image_num,
                     )
-
-                buffer = io.BytesIO()
-                image_format = task.output_format.upper()
-                if image_format == "JPG":
-                    image_format = "JPEG"
-                if image_format == "JPEG" and image.mode == "RGBA":
-                    image = image.convert("RGB")
-                image.save(buffer, format=image_format)
-
-                yield ImageGenerationResponse(
-                    image_data=buffer.getvalue(),
-                    format=task.output_format,
-                    stats=stats,
-                )

src/exo/worker/engines/image/models/__init__.py

@@ -33,6 +33,7 @@ _ADAPTER_REGISTRY: dict[str, AdapterFactory] = {
 # Config registry: maps model ID patterns to configs
 _CONFIG_REGISTRY: dict[str, ImageModelConfig] = {
     "flux.1-schnell": FLUX_SCHNELL_CONFIG,
+    "flux.1-krea-dev": FLUX_DEV_CONFIG,  # Must come before "flux.1-dev" for pattern matching
     "flux.1-dev": FLUX_DEV_CONFIG,
     "qwen-image-edit": QWEN_IMAGE_EDIT_CONFIG,  # Must come before "qwen-image" for pattern matching
     "qwen-image": QWEN_IMAGE_CONFIG,

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

@@ -13,12 +13,17 @@ from mlx.nn.layers.distributed import (
     shard_linear,
     sum_gradients,
 )
+from mlx_lm.models.base import (
+    scaled_dot_product_attention,  # pyright: ignore[reportUnknownVariableType]
+)
 from mlx_lm.models.deepseek_v3 import DeepseekV3MLP
 from mlx_lm.models.deepseek_v3 import Model as DeepseekV3Model
 from mlx_lm.models.deepseek_v32 import DeepseekV32MLP
 from mlx_lm.models.deepseek_v32 import Model as DeepseekV32Model
 from mlx_lm.models.glm4_moe import Model as Glm4MoeModel
 from mlx_lm.models.glm4_moe import MoE
+from mlx_lm.models.glm4_moe_lite import Glm4MoeLiteDecoderLayer, Glm4MoeLiteMLP
+from mlx_lm.models.glm4_moe_lite import Model as GLM4MoeLiteModel
 from mlx_lm.models.gpt_oss import GptOssMoeModel
 from mlx_lm.models.gpt_oss import Model as GptOssModel
 from mlx_lm.models.llama import Model as LlamaModel
@@ -27,7 +32,8 @@ 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 Qwen3MoeSparseMoeBlock
 from mlx_lm.models.qwen3_next import Model as Qwen3NextModel
-from mlx_lm.models.qwen3_next import Qwen3NextSparseMoeBlock
+from mlx_lm.models.qwen3_next import Qwen3NextDecoderLayer, Qwen3NextSparseMoeBlock
+from transformers.models.qwen3.modeling_qwen3 import Qwen3DecoderLayer
 
 from exo.shared.logging import logger
 from exo.shared.types.worker.shards import PipelineShardMetadata
@@ -100,6 +106,16 @@ class CustomMlxLayer(nn.Module):
                 return getattr(original_layer, name)
 
 
+class EvalCheckpointLayer(CustomMlxLayer):
+    """Wraps a layer to force evaluation of its output, breaking up the computation graph
+    to prevent Metal command buffer timeouts with large batches in pipeline parallel."""
+
+    def __call__(self, x: mx.array, *args: object, **kwargs: object) -> mx.array:
+        output = self.original_layer(x, *args, **kwargs)
+        mx.eval(output)
+        return output
+
+
 class PipelineFirstLayer(CustomMlxLayer):
     def __init__(
         self,
@@ -137,14 +153,20 @@ class PipelineLastLayer(CustomMlxLayer):
         ).arguments.get("cache", None)
 
         output: mx.array = self.original_layer(x, *args, **kwargs)
+        mx.eval(output)
 
         if self.r != self.s - 1:
             output = mx.distributed.send(
                 output, (self.r + 1) % self.s, group=self.group
             )
+            mx.async_eval(output)
             if cache is not None:
                 cache.keys = mx.depends(cache.keys, output)  # type: ignore[reportUnknownMemberType]
 
+        output = mx.distributed.all_gather(output, group=self.group)[
+            -output.shape[0] :
+        ]  # type :ignore
+
         return output
 
 
@@ -195,6 +217,9 @@ def pipeline_auto_parallel(
 
     layers = layers[start_layer:end_layer]
     layers[0] = PipelineFirstLayer(layers[0], device_rank, group=group)
+    # Wrap intermediate layers with eval checkpoints to prevent GPU timeout
+    for i in range(1, len(layers) - 1):
+        layers[i] = EvalCheckpointLayer(layers[i])
     layers[-1] = PipelineLastLayer(
         layers[-1],
         device_rank,
@@ -248,14 +273,14 @@ def patch_pipeline_model[T](model: T, group: mx.distributed.Group) -> T:
             "cache", None
         )
 
+        # Evaluate logits before all_gather to break the computation graph
+        # and prevent Metal command buffer timeouts with large batches
+        mx.eval(logits)
+
         # Add dependency to last cache entry to ensure distributed ops are evaluated
         if cache is not None:
             cache[-1].state = mx.depends(cache[-1].state, logits)  # type: ignore
 
-        logits = mx.distributed.all_gather(logits, group=group)[
-            -logits.shape[0] :
-        ]  # type :ignore
-
         return logits
 
     cls.__call__ = patched_call
@@ -334,15 +359,7 @@ def tensor_auto_parallel(
         group=group,
     )
 
-    if hasattr(model, "shard") and not isinstance(model, GptOssModel):
-        try:
-            model.shard(group)  # type: ignore
-            return patch_tensor_model(model)
-        except (AttributeError, TypeError, NameError):
-            pass
-
     if isinstance(model, (LlamaModel, Ministral3Model)):
-        logger.warning("shouldn't be hit - upstream sharding exists")
         tensor_parallel_sharding_strategy = LlamaShardingStrategy(
             group,
             all_to_sharded_linear,
@@ -351,7 +368,6 @@ def tensor_auto_parallel(
             sharded_to_all_linear_in_place,
         )
     elif isinstance(model, (DeepseekV3Model, DeepseekV32Model)):
-        logger.warning("shouldn't be hit - upstream sharding exists")
         tensor_parallel_sharding_strategy = DeepSeekShardingStrategy(
             group,
             all_to_sharded_linear,
@@ -367,6 +383,14 @@ def tensor_auto_parallel(
             all_to_sharded_linear_in_place,
             sharded_to_all_linear_in_place,
         )
+    elif isinstance(model, GLM4MoeLiteModel):
+        tensor_parallel_sharding_strategy = GLM4MoeLiteShardingStrategy(
+            group,
+            all_to_sharded_linear,
+            sharded_to_all_linear,
+            all_to_sharded_linear_in_place,
+            sharded_to_all_linear_in_place,
+        )
     elif isinstance(model, (Qwen3MoeModel, Glm4MoeModel, Qwen3NextModel)):
         tensor_parallel_sharding_strategy = QwenShardingStrategy(
             group,
@@ -441,7 +465,7 @@ class LlamaShardingStrategy(TensorParallelShardingStrategy):
             layer.mlp.gate_proj = self.all_to_sharded_linear(layer.mlp.gate_proj)
             layer.mlp.down_proj = self.sharded_to_all_linear(layer.mlp.down_proj)
             layer.mlp.up_proj = self.all_to_sharded_linear(layer.mlp.up_proj)
-
+            mx.eval(layer)
         return model
 
 
@@ -496,6 +520,9 @@ class DeepSeekShardingStrategy(TensorParallelShardingStrategy):
                 layer.self_attn.kv_b_proj
             )
             layer.self_attn.o_proj = self.sharded_to_all_linear(layer.self_attn.o_proj)
+            # Store pre-shard head count and group for context parallelism
+            layer.self_attn.context_parallel_total_heads = layer.self_attn.num_heads
+            layer.self_attn._cp_group = self.group
             layer.self_attn.num_heads //= self.N
 
             # Shard the MLP
@@ -516,6 +543,12 @@ class DeepSeekShardingStrategy(TensorParallelShardingStrategy):
                 layer.mlp = ShardedDeepseekV3MoE(layer.mlp)  # type: ignore
                 layer.mlp.sharding_group = self.group
 
+            mx.eval(layer)
+
+        # Store group for context parallelism
+        if hasattr(model, "model"):
+            model.model._cp_group = self.group
+
         return model
 
 
@@ -533,6 +566,158 @@ class ShardedDeepseekV3MoE(CustomMlxLayer):
         return y
 
 
+class GLM4MoeLiteShardingStrategy(TensorParallelShardingStrategy):
+    def shard_model(
+        self,
+        model: nn.Module,
+        timeout_seconds: float,
+        on_timeout: TimeoutCallback | None,
+    ) -> nn.Module:
+        model = cast(GLM4MoeLiteModel, model)
+        for layer in model.layers:  # type: ignore
+            layer = cast(Glm4MoeLiteDecoderLayer, layer)
+            eval_with_timeout(
+                layer.parameters(),
+                timeout_seconds / len(model.layers),  # type: ignore
+                on_timeout,
+            )
+            if layer.self_attn.q_lora_rank is None:  # type: ignore
+                layer.self_attn.q_proj = self.all_to_sharded_linear(
+                    layer.self_attn.q_proj
+                )
+            else:
+                layer.self_attn.q_b_proj = self.all_to_sharded_linear(
+                    layer.self_attn.q_b_proj
+                )
+
+            layer.self_attn.o_proj = self.sharded_to_all_linear(layer.self_attn.o_proj)
+            layer.self_attn.num_heads //= self.N
+
+            # Logic from upstream mlx
+            num_heads = layer.self_attn.num_heads
+            sh = self.group.rank() * num_heads
+            eh = sh + num_heads
+
+            def shard_heads(w: mx.array, sh: int = sh, eh: int = eh) -> mx.array:
+                return w[sh:eh]
+
+            layer.self_attn.embed_q.apply(shard_heads)
+            layer.self_attn.unembed_out.apply(shard_heads)
+
+            if isinstance(layer.mlp, Glm4MoeLiteMLP):
+                layer.mlp.gate_proj = self.all_to_sharded_linear(layer.mlp.gate_proj)
+                layer.mlp.down_proj = self.sharded_to_all_linear(layer.mlp.down_proj)
+                layer.mlp.up_proj = self.all_to_sharded_linear(layer.mlp.up_proj)
+
+            else:
+                if getattr(layer.mlp, "shared_experts", None) is not None:
+                    self.all_to_sharded_linear_in_place(
+                        layer.mlp.shared_experts.gate_proj
+                    )
+                    self.sharded_to_all_linear_in_place(
+                        layer.mlp.shared_experts.down_proj
+                    )
+                    self.all_to_sharded_linear_in_place(
+                        layer.mlp.shared_experts.up_proj
+                    )
+                self.all_to_sharded_linear_in_place(layer.mlp.switch_mlp.gate_proj)
+                self.sharded_to_all_linear_in_place(layer.mlp.switch_mlp.down_proj)
+                self.all_to_sharded_linear_in_place(layer.mlp.switch_mlp.up_proj)
+                layer.mlp = ShardedGLM4MoeLiteMoE(layer.mlp)  # type: ignore
+                layer.mlp.sharding_group = self.group  # type: ignore
+            mx.eval(layer)
+
+        return model
+
+
+class ShardedGLM4MoeLiteMoE(CustomMlxLayer):
+    def __init__(self, layer: _LayerCallable):
+        super().__init__(layer)
+        self.sharding_group: mx.distributed.Group | None = None
+
+    def __call__(self, x: mx.array) -> mx.array:
+        if self.sharding_group is not None:
+            x = sum_gradients(self.sharding_group)(x)
+        y = self.original_layer.__call__(x)
+        if self.sharding_group is not None:
+            y = mx.distributed.all_sum(y, group=self.sharding_group)
+        return y
+
+
+class WrappedMiniMaxAttention(CustomMlxLayer):
+    def __init__(self, layer: _LayerCallable, group: mx.distributed.Group):
+        super().__init__(layer)
+        self.group = group
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: mx.array | Any = None,
+        cache: Any | None = None,
+    ) -> mx.array:
+        B, L, _ = x.shape
+
+        queries, keys, values = self.q_proj(x), self.k_proj(x), self.v_proj(x)  # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType]
+
+        if getattr(self, "use_qk_norm", False):
+            q_dim = queries.shape[-1]  # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType]
+            k_dim = keys.shape[-1]  # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType]
+            N = self.group.size()
+
+            qk = mx.concatenate([queries, keys], axis=-1)  # (B, L, q_dim + k_dim)
+            qk = mx.distributed.all_gather(
+                qk, group=self.group
+            )  # (N*B, L, q_dim + k_dim)
+
+            # Reshape to separate rank contributions: (N, B, L, q_dim + k_dim)
+            # Then transpose to (B, L, N, q_dim + k_dim) and merge N into feature dim
+            qk = qk.reshape(N, B, L, q_dim + k_dim).transpose(1, 2, 0, 3)  # pyright: ignore[reportUnknownMemberType,reportUnknownArgumentType]
+            queries = qk[..., :q_dim].reshape(
+                B, L, -1
+            )  # (B, L, N * q_dim)  # pyright: ignore[reportUnknownMemberType]
+            keys = qk[..., q_dim:].reshape(
+                B, L, -1
+            )  # (B, L, N * k_dim)  # pyright: ignore[reportUnknownMemberType]
+
+            queries = self.q_norm(queries)  # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType]
+            keys = self.k_norm(keys)  # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType]
+
+            # Split back and take this rank's portion
+            queries = mx.split(queries, N, axis=-1)[self.group.rank()]
+            keys = mx.split(keys, N, axis=-1)[self.group.rank()]
+
+        queries = queries.reshape(B, L, self.num_attention_heads, -1).transpose(  # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType,reportUnknownArgumentType]
+            0, 2, 1, 3
+        )
+        keys = keys.reshape(B, L, self.num_key_value_heads, -1).transpose(  # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType,reportUnknownArgumentType]
+            0, 2, 1, 3
+        )
+        values = values.reshape(B, L, self.num_key_value_heads, -1).transpose(  # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType]
+            0, 2, 1, 3
+        )
+
+        if cache is not None:
+            queries = self.rope(queries, offset=cache.offset)  # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType,reportAny]
+            keys = self.rope(keys, offset=cache.offset)  # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType,reportAny]
+            keys, values = cache.update_and_fetch(keys, values)  # pyright: ignore[reportAny]
+        else:
+            queries = self.rope(queries)  # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType]
+            keys = self.rope(keys)  # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType]
+
+        output = scaled_dot_product_attention(
+            queries,
+            keys,
+            values,
+            cache=cache,
+            scale=self.scale,
+            mask=mask,  # pyright: ignore[reportUnknownMemberType,reportUnknownArgumentType]
+        )
+
+        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)  # pyright: ignore[reportUnknownMemberType]
+
+        return self.o_proj(output)  # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType]
+
+
 class MiniMaxShardingStrategy(TensorParallelShardingStrategy):
     def shard_model(
         self,
@@ -550,9 +735,12 @@ class MiniMaxShardingStrategy(TensorParallelShardingStrategy):
             layer.self_attn.k_proj = self.all_to_sharded_linear(layer.self_attn.k_proj)
             layer.self_attn.v_proj = self.all_to_sharded_linear(layer.self_attn.v_proj)
             layer.self_attn.o_proj = self.sharded_to_all_linear(layer.self_attn.o_proj)
+
             layer.self_attn.num_attention_heads //= self.N
             layer.self_attn.num_key_value_heads //= self.N
 
+            layer.self_attn = WrappedMiniMaxAttention(layer.self_attn, self.group)  # pyright: ignore[reportAttributeAccessIssue,reportArgumentType]
+
             # Shard the MoE. Shard in place since the MoE should be responsible
             # for aggregating the results.
             self.all_to_sharded_linear_in_place(
@@ -566,7 +754,7 @@ class MiniMaxShardingStrategy(TensorParallelShardingStrategy):
             )
             layer.block_sparse_moe = ShardedQwenMoE(layer.block_sparse_moe)  # pyright: ignore[reportAttributeAccessIssue, reportArgumentType]
             layer.block_sparse_moe.sharding_group = self.group  # pyright: ignore[reportAttributeAccessIssue]
-
+            mx.eval(layer)
         return model
 
 
@@ -577,18 +765,32 @@ class QwenShardingStrategy(TensorParallelShardingStrategy):
         timeout_seconds: float,
         on_timeout: TimeoutCallback | None,
     ) -> nn.Module:
-        model = cast(Qwen3MoeModel, model)
+        model = cast(Qwen3MoeModel | Qwen3NextModel, model)
         for layer in model.layers:
             eval_with_timeout(
                 layer.parameters(), timeout_seconds / len(model.layers), on_timeout
             )
             # Shard the self attention
-            layer.self_attn.q_proj = self.all_to_sharded_linear(layer.self_attn.q_proj)
-            layer.self_attn.k_proj = self.all_to_sharded_linear(layer.self_attn.k_proj)
-            layer.self_attn.v_proj = self.all_to_sharded_linear(layer.self_attn.v_proj)
-            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
+            if isinstance(layer, Qwen3DecoderLayer) or hasattr(layer, "self_attn"):
+                layer.self_attn.q_proj = self.all_to_sharded_linear(
+                    layer.self_attn.q_proj
+                )
+                layer.self_attn.k_proj = self.all_to_sharded_linear(
+                    layer.self_attn.k_proj
+                )
+                layer.self_attn.v_proj = self.all_to_sharded_linear(
+                    layer.self_attn.v_proj
+                )
+                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) and hasattr(
+                    layer, "linear_attn"
+                )
+                # These layers are fast so we don't shard. This may change in future.
 
             # Shard the MoE. Shard in place since the MoE should be responsible
             # for aggregating the results.
@@ -607,6 +809,7 @@ class QwenShardingStrategy(TensorParallelShardingStrategy):
                 layer.mlp.down_proj = self.sharded_to_all_linear(layer.mlp.down_proj)
                 layer.mlp.up_proj = self.all_to_sharded_linear(layer.mlp.up_proj)
 
+            mx.eval(layer)
         return model
 
 
@@ -661,7 +864,7 @@ class GptOssShardingStrategy(TensorParallelShardingStrategy):
 
             layer.mlp = ShardedGptOssMoE(layer.mlp)  # type: ignore
             layer.mlp.sharding_group = self.group  # pyright: ignore[reportAttributeAccessIssue]
-
+            mx.eval(layer)
         return model
 
 

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

@@ -1,39 +1,81 @@
-# type: ignore
-# TODO: Fix this file, including types!
+import os
 from copy import deepcopy
-from typing import Callable
+from typing import Any, cast
 
 import mlx.core as mx
-from mlx_lm import stream_generate
-from mlx_lm.models.cache import _BaseCache, trim_prompt_cache
+from mlx_lm.models.cache import (
+    KVCache,
+    QuantizedKVCache,
+    RotatingKVCache,
+    trim_prompt_cache,
+)
+from mlx_lm.models.gpt_oss import Model as GptOssModel
 from mlx_lm.tokenizer_utils import TokenizerWrapper
 
+from exo.shared.types.mlx import KVCacheType
 from exo.worker.engines.mlx import Model
-from exo.worker.engines.mlx.constants import KEEP_KV_SIZE, KV_BITS, KV_GROUP_SIZE
-from exo.worker.engines.mlx.utils_mlx import make_kv_cache
+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
+_DEFAULT_MEMORY_THRESHOLD = 0.85
+_MEMORY_THRESHOLD = float(
+    os.environ.get("EXO_MEMORY_THRESHOLD", _DEFAULT_MEMORY_THRESHOLD)
+)
 
 
 class KVPrefixCache:
-    def __init__(self):
-        # Only one prefix cache per runner.
+    def __init__(self, tokenizer: TokenizerWrapper):
         self.prompts: list[mx.array] = []  # mx array of tokens (ints)
-        self.caches: list[list[_BaseCache]] = []
+        self.caches: list[KVCacheType] = []
+        self._last_used: list[int] = []  # monotonic counter of last access per entry
+        self._access_counter: int = 0
+        self._tokenizer: TokenizerWrapper = tokenizer
 
-    def add_kv_cache(
-        self, tokenizer: TokenizerWrapper, prompt: str, cache: list[_BaseCache]
-    ):
-        tokenized_prompt = self.encode_prompt(tokenizer, prompt)
+    def clear(self):
+        """Clear all cached prompts and caches."""
+        self.prompts.clear()
+        self.caches.clear()
+        self._last_used.clear()
+
+    def add_kv_cache(self, prompt: str, cache: KVCacheType):
+        """Add a new cache entry. Evicts LRU entries if memory is high."""
+        self._evict_if_needed()
+        tokenized_prompt = encode_prompt(self._tokenizer, prompt)
         self.prompts.append(tokenized_prompt)
         self.caches.append(deepcopy(cache))
+        self._access_counter += 1
+        self._last_used.append(self._access_counter)
+        logger.info(f"KV cache added: {len(tokenized_prompt)} tokens")
+
+    def update_kv_cache(
+        self,
+        index: int,
+        prompt: str,
+        cache: KVCacheType,
+    ):
+        """Update an existing cache entry in-place."""
+        tokenized_prompt = encode_prompt(self._tokenizer, prompt)
+        self.prompts[index] = tokenized_prompt
+        self.caches[index] = deepcopy(cache)
+        self._access_counter += 1
+        self._last_used[index] = self._access_counter
+        logger.info(f"KV cache updated (index {index}): {len(tokenized_prompt)} tokens")
 
     def get_kv_cache(
         self,
         model: Model,
-        tokenizer: TokenizerWrapper,
-        sampler: Callable[[mx.array], mx.array],
         prompt: str,
-    ) -> list[_BaseCache]:
-        tokenized_prompt = self.encode_prompt(tokenizer, prompt)
+    ) -> tuple[KVCacheType, mx.array, int | None]:
+        """Get KV cache for prompt, returning remaining tokens to prefill.
+
+        Returns:
+            Tuple of (cache, remaining_tokens, matched_index) where:
+            - cache: KV cache to use for generation
+            - remaining_tokens: tokens that still need prefilling
+            - matched_index: index of the matched entry (None if no match)
+        """
+        tokenized_prompt = encode_prompt(self._tokenizer, prompt)
         max_length = len(tokenized_prompt)
 
         best_snapshot_index, best_snapshot_length = None, 0
@@ -42,63 +84,127 @@ class KVPrefixCache:
             length = _get_prefix_length(tokenized_prompt, cached_prompt)
 
             if length == max_length:
-                return self.caches[i]
+                # Exact match - cached prompt starts with our entire prompt
+                # Trim cache to prompt length - 1, return last token for stream_generate
+                prompt_cache = deepcopy(self.caches[i])
+                cached_length = _cache_length(self.caches[i])
+                tokens_to_trim = cached_length - (max_length - 1)
+                if tokens_to_trim > 0:
+                    trim_prompt_cache(cast(list[Any], prompt_cache), tokens_to_trim)
+                self._access_counter += 1
+                self._last_used[i] = self._access_counter
+                logger.info(f"KV cache exact match: {max_length} tokens (instant)")
+                return prompt_cache, tokenized_prompt[-1:], i
 
             if length > best_snapshot_length:
                 best_snapshot_index, best_snapshot_length = i, length
 
         if best_snapshot_index is not None:
-            prompt_cache = deepcopy(self.caches[best_snapshot_index])
-            trim_prompt_cache(prompt_cache, max_length - best_snapshot_length)
-            tokenized_prompt = tokenized_prompt[best_snapshot_index:]
-
-        else:
-            prompt_cache = make_kv_cache(
-                model,
-                # max_kv_size=MAX_KV_SIZE,
-                # keep=KEEP_KV_SIZE
+            new_tokens = max_length - best_snapshot_length
+            logger.info(
+                f"KV cache prefix match: {best_snapshot_length}/{max_length} tokens "
+                f"(reusing {best_snapshot_length}, need to prefill {new_tokens})"
             )
 
-        prefill(model, tokenizer, sampler, tokenized_prompt, prompt_cache)
+            prompt_cache = deepcopy(self.caches[best_snapshot_index])
 
-        return prompt_cache
+            # Trim removes tokens from the end, so we trim (cached_length - prefix_length) to keep the prefix
+            cached_length = _cache_length(self.caches[best_snapshot_index])
+            tokens_to_trim = cached_length - best_snapshot_length
+            if tokens_to_trim > 0:
+                trim_prompt_cache(cast(list[Any], prompt_cache), tokens_to_trim)
 
-    def encode_prompt(self, tokenizer: TokenizerWrapper, prompt: str) -> mx.array:
-        add_special_tokens = tokenizer.bos_token is None or not prompt.startswith(
-            tokenizer.bos_token
-        )
-        tokenized_prompt = tokenizer.encode(
-            prompt, add_special_tokens=add_special_tokens
-        )
-        return mx.array(tokenized_prompt)
+            self._access_counter += 1
+            self._last_used[best_snapshot_index] = self._access_counter
+            remaining_tokens = tokenized_prompt[best_snapshot_length:]
+            return prompt_cache, remaining_tokens, best_snapshot_index
+
+        else:
+            prompt_cache = make_kv_cache(model)
+            if len(self.prompts) == 0:
+                logger.info(f"KV cache empty, need to prefill {max_length} tokens")
+            else:
+                logger.info(
+                    f"KV cache no prefix match, need to prefill {max_length} tokens"
+                )
+
+            return prompt_cache, tokenized_prompt, None
+
+    def _evict_if_needed(self):
+        """Evict least recently used entries while memory pressure is high."""
+        if len(self.caches) == 0:
+            return
+
+        active: int = mx.metal.get_active_memory()
+        limit = int(mx.metal.device_info()["max_recommended_working_set_size"])
+        if active < limit * _MEMORY_THRESHOLD:
+            return
+
+        # Evict LRU entries until below threshold or only one entry left
+        while len(self.caches) > 0:
+            lru_index = self._last_used.index(min(self._last_used))
+            evicted_tokens = len(self.prompts[lru_index])
+            self.prompts.pop(lru_index)
+            self.caches.pop(lru_index)
+            self._last_used.pop(lru_index)
+            logger.info(
+                f"KV cache evicted LRU entry ({evicted_tokens} tokens) due to memory pressure"
+            )
+
+            active = mx.metal.get_active_memory()
+            if active < limit * _MEMORY_THRESHOLD:
+                break
+
+
+def encode_prompt(tokenizer: TokenizerWrapper, prompt: str) -> mx.array:
+    """Encode a prompt string to token array.
+
+    For chat-templated prompts (which have their own structure markers like
+    <|im_user|>, <|im_middle|>, etc.), we should NOT add BOS/EOS tokens as
+    that would corrupt the prompt structure.
+    """
+    # Chat templates define their own structure - don't add BOS/EOS
+    tokenized_prompt = tokenizer.encode(prompt, add_special_tokens=False)
+    return mx.array(tokenized_prompt)
+
+
+def _cache_length(cache: KVCacheType) -> int:
+    """Get the number of tokens in a KV cache."""
+    # Use .offset attribute which all cache types have (len() not implemented in older QuantizedKVCache)
+    return max(c.offset for c in cache)  # type: ignore
 
 
 def _get_prefix_length(prompt: mx.array, cached_prompt: mx.array) -> int:
-    n = min(int(prompt.shape[0]), int(cached_prompt.shape[0]), KEEP_KV_SIZE)
+    """Find the length of the common prefix between two token arrays."""
+    n = min(int(prompt.shape[0]), int(cached_prompt.shape[0]))
     if n == 0:
         return 0
 
-    equal = (prompt[:n] == cached_prompt[:n]).astype(mx.int32)
+    equal = mx.equal(prompt[:n], cached_prompt[:n]).astype(mx.int32)
     prefix_mask = mx.cumprod(equal)  # stays 1 until first mismatch, then 0 forever
     return int(mx.sum(prefix_mask).item())
 
 
-def prefill(
-    model: Model,
-    tokenizer: TokenizerWrapper,
-    sampler: Callable[[mx.array], mx.array],
-    prompt: mx.array,
-    cache: list[_BaseCache],
-) -> None:
-    for _ in stream_generate(
-        model=model,
-        tokenizer=tokenizer,
-        prompt=prompt,
-        max_tokens=0,
-        sampler=sampler,
-        prompt_cache=cache,
-        prefill_step_size=2048,
-        kv_group_size=KV_GROUP_SIZE,
-        kv_bits=KV_BITS,
-    ):
-        pass
+def make_kv_cache(
+    model: Model, max_kv_size: int | None = None, keep: int = 0
+) -> KVCacheType:
+    assert hasattr(model, "layers")
+
+    # TODO: Do this for all models
+    if hasattr(model, "make_cache") and isinstance(model, GptOssModel):
+        logger.info("Using MLX LM's make cache")
+        return model.make_cache()  # type: ignore
+
+    if max_kv_size 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)
+                for _ in model.layers
+            ]
+    else:
+        logger.info(f"Using rotating KV cache with {max_kv_size=} with {keep=}")
+        return [RotatingKVCache(max_size=max_kv_size, keep=keep) for _ in model.layers]

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

@@ -4,7 +4,7 @@
 KV_GROUP_SIZE: int | None = 32
 KV_BITS: int | None = None
 ATTENTION_KV_BITS: int | None = 4
-MAX_TOKENS: int = 8192
+MAX_TOKENS: int = 32168
 MAX_KV_SIZE: int | None = 3200
 KEEP_KV_SIZE: int | None = 1600
 QUANTIZE_MODEL_MODE: str | None = "affine"

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

@@ -1,48 +1,92 @@
+import time
 from typing import Any, Callable, Generator, cast, get_args
 
 import mlx.core as mx
 from mlx_lm.generate import stream_generate
-from mlx_lm.models.cache import KVCache
+from mlx_lm.models.cache import KVCache, trim_prompt_cache
 from mlx_lm.sample_utils import make_sampler
 from mlx_lm.tokenizer_utils import TokenizerWrapper
 
-# from exo.engines.mlx.cache import KVPrefixCache
 from exo.shared.types.api import (
     BenchChatCompletionTaskParams,
     ChatCompletionMessage,
     FinishReason,
     GenerationStats,
+    TopLogprobItem,
 )
 from exo.shared.types.memory import Memory
+from exo.shared.types.mlx import KVCacheType
 from exo.shared.types.tasks import ChatCompletionTaskParams
 from exo.shared.types.worker.runner_response import (
     GenerationResponse,
 )
 from exo.worker.engines.mlx import Model
+from exo.worker.engines.mlx.cache import KVPrefixCache, encode_prompt, make_kv_cache
 from exo.worker.engines.mlx.constants import KV_BITS, KV_GROUP_SIZE, MAX_TOKENS
 from exo.worker.engines.mlx.utils_mlx import (
     apply_chat_template,
-    make_kv_cache,
     mx_barrier,
 )
 from exo.worker.runner.bootstrap import logger
 
 generation_stream = mx.new_stream(mx.default_device())
 
+_MIN_PREFIX_HIT_TO_UPDATE = 1000
 
-def maybe_quantize_kv_cache(
-    prompt_cache: list[KVCache | Any],
-    quantized_kv_start: int,
-    kv_group_size: int,
-    kv_bits: int | None,
-) -> None:
-    if kv_bits is None:
-        return
-    for e, c in enumerate(prompt_cache):
-        if (
-            hasattr(c, "to_quantized") and c.offset >= quantized_kv_start  # type: ignore
-        ):
-            prompt_cache[e] = c.to_quantized(group_size=kv_group_size, bits=kv_bits)
+
+def prefill(
+    model: Model,
+    tokenizer: TokenizerWrapper,
+    sampler: Callable[[mx.array], mx.array],
+    prompt_tokens: mx.array,
+    cache: KVCacheType,
+) -> float:
+    """Prefill the KV cache with prompt tokens.
+
+    This runs the model over the prompt tokens to populate the cache,
+    then trims off the extra generated token.
+
+    Returns:
+        tokens_per_sec
+    """
+    num_tokens = len(prompt_tokens)
+    if num_tokens == 0:
+        return 0.0
+
+    logger.debug(f"Prefilling {num_tokens} tokens...")
+    start_time = time.perf_counter()
+
+    def progress_callback(processed: int, total: int) -> None:
+        elapsed = time.time() - start_time
+        tok_per_sec = processed / elapsed if elapsed > 0 else 0
+        logger.debug(
+            f"Prefill progress: {processed}/{total} tokens ({tok_per_sec:.1f} tok/s)"
+        )
+
+    # Use max_tokens=1 because max_tokens=0 does not work.
+    # We just throw away the generated token - we only care about filling the cache
+    for _ in stream_generate(
+        model=model,
+        tokenizer=tokenizer,
+        prompt=prompt_tokens,
+        max_tokens=1,
+        sampler=sampler,
+        prompt_cache=cache,
+        prefill_step_size=2048,
+        kv_group_size=KV_GROUP_SIZE,
+        kv_bits=KV_BITS,
+        prompt_progress_callback=progress_callback,
+    ):
+        break  # Stop after first iteration - cache is now filled
+    trim_prompt_cache(cast(list[Any], cache), 1)
+
+    elapsed = time.perf_counter() - start_time
+    tokens_per_sec = num_tokens / elapsed if elapsed > 0 else 0.0
+    logger.debug(
+        f"Prefill complete: {num_tokens} tokens in {elapsed:.2f}s "
+        f"({tokens_per_sec:.1f} tok/s)"
+    )
+    return tokens_per_sec
 
 
 def warmup_inference(
@@ -115,11 +159,212 @@ def eos_ids_from_tokenizer(tokenizer: TokenizerWrapper) -> list[int]:
     return eos
 
 
+def extract_top_logprobs(
+    logprobs_array: mx.array,
+    selected_token: int,
+    tokenizer: TokenizerWrapper,
+    top_k: int | None,
+) -> tuple[float, list[TopLogprobItem]]:
+    """Extract the selected token's logprob and top-k alternatives.
+
+    top k an be set to None to return all the logprobs
+    """
+    selected_logprob = float(logprobs_array[selected_token].item())
+
+    if top_k == 0:
+        return selected_logprob, []
+
+    vocab_size = logprobs_array.shape[0]
+
+    if top_k is None:
+        sorted_indices = mx.argsort(-logprobs_array)
+        mx.eval(sorted_indices)
+        indices_list: list[int] = cast(list[int], sorted_indices.tolist())
+    else:
+        k = min(top_k, vocab_size)
+        top_indices = mx.argpartition(-logprobs_array, kth=k - 1)[:k]
+        top_logprobs_values = logprobs_array[top_indices]
+        sorted_order = mx.argsort(-top_logprobs_values)
+        top_indices = top_indices[sorted_order]
+        mx.eval(top_indices)
+        indices_list = cast(list[int], top_indices.tolist())
+
+    top_logprob_items: list[TopLogprobItem] = []
+    for token_id in indices_list:
+        logprob_value = float(logprobs_array[token_id].item())
+        token_str = tokenizer.decode([token_id])
+
+        top_logprob_items.append(
+            TopLogprobItem(
+                token=token_str,
+                logprob=logprob_value,
+                bytes=list(token_str.encode("utf-8")),
+            )
+        )
+
+    return selected_logprob, top_logprob_items
+
+
+def score_tokens(
+    model: Model,
+    tokenizer: TokenizerWrapper,
+    tokens: list[int],
+    top_k: int | None = None,
+) -> list[tuple[float, list[TopLogprobItem]]]:
+    """Score a sequence of tokens, returning logprobs for each token.
+
+    This is used for the completions API with echo=True, where we need
+    logprobs for the prompt tokens (not just generated tokens).
+
+    Args:
+        model: The MLX model.
+        tokenizer: The tokenizer.
+        tokens: List of token IDs to score.
+        top_k: Number of top logprobs to return per position.
+               If None, returns all logprobs.
+
+    Returns:
+        List of (token_logprob, top_logprobs) tuples for each token position.
+        The first position has no logprob (no previous context), so returns (0.0, []).
+    """
+    if len(tokens) == 0:
+        return []
+
+    # First token has no previous context to condition on
+    results: list[tuple[float, list[TopLogprobItem]]] = [(0.0, [])]
+
+    if len(tokens) == 1:
+        return results
+
+    # Create an empty KV cache for the forward pass
+    cache = make_kv_cache(model=model)
+
+    # Convert to MLX array and run forward pass
+    input_tokens = mx.array(tokens[:-1])[None]  # All tokens except last, batched
+
+    # Run the model to get logits for all positions
+    # The model returns logits with shape [1, seq_len, vocab_size]
+    logits: mx.array = model(input_tokens, cache=cast(list[KVCache], cache))
+    logits = logits.squeeze(0)  # Shape: [seq_len, vocab_size]
+
+    # Convert to log probabilities
+    logprobs_all: mx.array = logits - mx.logsumexp(logits, axis=-1, keepdims=True)
+
+    mx.eval(logprobs_all)
+
+    # For each position, extract the logprob of the actual next token
+    for i in range(len(tokens) - 1):
+        next_token = tokens[i + 1]
+        logprobs_at_position: mx.array = logprobs_all[i]
+
+        logprob, top_logprobs_items = extract_top_logprobs(
+            logprobs_array=logprobs_at_position,
+            selected_token=next_token,
+            tokenizer=tokenizer,
+            top_k=top_k,
+        )
+        results.append((logprob, top_logprobs_items))
+
+    return results
+
+
+def score_tokens_batched(
+    model: Model,
+    tokenizer: TokenizerWrapper,
+    token_sequences: list[list[int]],
+    top_k: int | None = None,
+) -> list[list[tuple[float, list[TopLogprobItem]]]]:
+    """Score multiple token sequences in a single batched forward pass.
+
+    This is significantly faster than calling score_tokens() multiple times
+    because it batches the forward pass across all sequences.
+
+    Args:
+        model: The MLX model.
+        tokenizer: The tokenizer.
+        token_sequences: List of token ID sequences to score.
+        top_k: Number of top logprobs to return per position.
+
+    Returns:
+        List of results for each sequence. Each result is a list of
+        (token_logprob, top_logprobs) tuples for each token position.
+    """
+    if not token_sequences:
+        return []
+
+    # Handle empty sequences and single-token sequences
+    results: list[list[tuple[float, list[TopLogprobItem]]]] = []
+    non_empty_indices: list[int] = []
+    non_empty_sequences: list[list[int]] = []
+
+    for i, tokens in enumerate(token_sequences):
+        if len(tokens) == 0:
+            results.append([])
+        elif len(tokens) == 1:
+            results.append([(0.0, [])])
+        else:
+            results.append([])  # Placeholder, will be filled later
+            non_empty_indices.append(i)
+            non_empty_sequences.append(tokens)
+
+    if not non_empty_sequences:
+        return results
+
+    # Find max sequence length (excluding last token since we predict it)
+    max_len = max(len(seq) - 1 for seq in non_empty_sequences)
+
+    # Get pad token (use eos_token_id or 0)
+    pad_token_id = getattr(tokenizer, "pad_token_id", None)
+    if pad_token_id is None:
+        pad_token_id = getattr(tokenizer, "eos_token_id", 0)
+
+    # Pad sequences and create attention mask
+    batch_size = len(non_empty_sequences)
+    padded_inputs = mx.full((batch_size, max_len), pad_token_id, dtype=mx.int32)
+    seq_lengths: list[int] = []
+
+    for i, tokens in enumerate(non_empty_sequences):
+        input_len = len(tokens) - 1  # Exclude last token
+        padded_inputs[i, :input_len] = mx.array(tokens[:-1], dtype=mx.int32)
+        seq_lengths.append(input_len)
+
+    # Run batched forward pass (no KV cache for scoring)
+    # The model accepts [batch_size, seq_len] and returns [batch_size, seq_len, vocab_size]
+    logits = model(padded_inputs, cache=None)
+
+    # Convert to log probabilities - logits shape: [batch, seq_len, vocab]
+    logprobs_all = logits - mx.logsumexp(logits, axis=-1, keepdims=True)
+    mx.eval(logprobs_all)
+
+    # Extract results for each sequence
+    for batch_idx, (orig_idx, tokens, seq_len) in enumerate(
+        zip(non_empty_indices, non_empty_sequences, seq_lengths, strict=True)
+    ):
+        seq_results: list[tuple[float, list[TopLogprobItem]]] = [(0.0, [])]
+
+        for pos in range(seq_len):
+            next_token = tokens[pos + 1]
+            logprobs_at_position: mx.array = logprobs_all[batch_idx, pos]
+
+            logprob, top_logprobs_items = extract_top_logprobs(
+                logprobs_array=logprobs_at_position,
+                selected_token=next_token,
+                tokenizer=tokenizer,
+                top_k=top_k,
+            )
+            seq_results.append((logprob, top_logprobs_items))
+
+        results[orig_idx] = seq_results
+
+    return results
+
+
 def mlx_generate(
     model: Model,
     tokenizer: TokenizerWrapper,
     task: ChatCompletionTaskParams,
     prompt: str,
+    kv_prefix_cache: KVPrefixCache | None = None,
 ) -> Generator[GenerationResponse]:
     # Ensure that generation stats only contains peak memory for this generation
     mx.reset_peak_memory()
@@ -131,7 +376,22 @@ def mlx_generate(
     if task.seed is not None:
         mx.random.seed(task.seed)
 
-    caches = make_kv_cache(model=model)
+    # Do not use the prefix cache if we are trying to do benchmarks.
+    if is_bench:
+        kv_prefix_cache = None
+
+    # Use prefix cache if available, otherwise create fresh cache
+    prefix_hit_length = 0
+    matched_index: int | None = None
+    if kv_prefix_cache is None:
+        caches = make_kv_cache(model=model)
+        prompt_tokens = encode_prompt(tokenizer, prompt)
+    else:
+        caches, prompt_tokens, matched_index = kv_prefix_cache.get_kv_cache(
+            model, prompt
+        )
+        all_prompt_tokens = encode_prompt(tokenizer, prompt)
+        prefix_hit_length = len(all_prompt_tokens) - len(prompt_tokens)
 
     logits_processors: list[Callable[[mx.array, mx.array], mx.array]] = []
     if is_bench:
@@ -144,11 +404,23 @@ def mlx_generate(
         top_p=task.top_p if task.top_p is not None else 1.0,
     )
 
+    # Prefill cache with all tokens except the last one
+    prefill_tps = prefill(model, tokenizer, sampler, prompt_tokens[:-1], caches)
+
+    # stream_generate starts from the last token
+    last_token = prompt_tokens[-1:]
+
+    # Determine if we need logprobs
+    should_extract_logprobs = task.logprobs is True
+    top_k = task.top_logprobs if task.top_logprobs is not None else 0
+
     max_tokens = task.max_tokens or MAX_TOKENS
+    generated_text_parts: list[str] = []
+    generation_start_time = time.perf_counter()
     for out in stream_generate(
         model=model,
         tokenizer=tokenizer,
-        prompt=prompt,
+        prompt=last_token,
         max_tokens=max_tokens,
         sampler=sampler,
         logits_processors=logits_processors,
@@ -158,12 +430,13 @@ def mlx_generate(
         kv_group_size=KV_GROUP_SIZE,
         kv_bits=KV_BITS,
     ):
+        generated_text_parts.append(out.text)
         logger.info(out.text)
 
         stats: GenerationStats | None = None
         if out.finish_reason is not None:
             stats = GenerationStats(
-                prompt_tps=float(out.prompt_tps),
+                prompt_tps=float(prefill_tps or out.prompt_tps),
                 generation_tps=float(out.generation_tps),
                 prompt_tokens=int(out.prompt_tokens),
                 generation_tokens=int(out.generation_tokens),
@@ -177,14 +450,47 @@ def mlx_generate(
                     f"Model generated unexpected finish_reason: {out.finish_reason}"
                 )
 
+        # Extract logprobs if requested
+        logprob: float | None = None
+        top_logprobs: list[TopLogprobItem] | None = None
+        if should_extract_logprobs:
+            logprob, top_logprobs = extract_top_logprobs(
+                logprobs_array=out.logprobs,
+                selected_token=out.token,
+                tokenizer=tokenizer,
+                top_k=top_k,
+            )
+
         yield GenerationResponse(
             text=out.text,
             token=out.token,
+            logprob=logprob,
+            top_logprobs=top_logprobs,
             finish_reason=cast(FinishReason | None, out.finish_reason),
             stats=stats,
         )
 
         if out.finish_reason is not None:
+            # Log generation stats
+            generation_elapsed = time.perf_counter() - generation_start_time
+            generated_tokens = len(generated_text_parts)
+            generation_tps = (
+                generated_tokens / generation_elapsed if generation_elapsed > 0 else 0.0
+            )
+            logger.debug(
+                f"Generation complete: prefill {prompt_tokens} tokens @ "
+                f"{prefill_tps:.1f} tok/s, generated {generated_tokens} tokens @ "
+                f"{generation_tps:.1f} tok/s"
+            )
+            if kv_prefix_cache is not None:
+                full_prompt = prompt + "".join(generated_text_parts)
+                if (
+                    matched_index is not None
+                    and prefix_hit_length >= _MIN_PREFIX_HIT_TO_UPDATE
+                ):
+                    kv_prefix_cache.update_kv_cache(matched_index, full_prompt, caches)
+                else:
+                    kv_prefix_cache.add_kv_cache(full_prompt, caches)
             break
 
         # TODO: Do we want an mx_barrier?

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

@@ -18,15 +18,12 @@ try:
 except ImportError:
     pass  # transformers < 5.0 or bytes_to_unicode not available
 
-from mlx_lm.models.cache import KVCache, QuantizedKVCache, RotatingKVCache
+from mlx_lm.models.cache import KVCache
 from mlx_lm.models.deepseek_v3 import DeepseekV3Model
-from mlx_lm.models.gpt_oss import Model as GptOssModel
 from mlx_lm.tokenizer_utils import TokenizerWrapper
 
 from exo.shared.models.model_cards import ModelId
 from exo.worker.engines.mlx.constants import (
-    CACHE_GROUP_SIZE,
-    KV_CACHE_BITS,
     TRUST_REMOTE_CODE,
 )
 
@@ -41,6 +38,7 @@ import mlx.nn as nn
 from mlx_lm.utils import load_model
 from pydantic import RootModel
 
+from exo.download.download_utils import build_model_path
 from exo.shared.types.api import ChatCompletionMessageText
 from exo.shared.types.common import Host
 from exo.shared.types.memory import Memory
@@ -55,7 +53,6 @@ from exo.shared.types.worker.shards import (
     ShardMetadata,
     TensorShardMetadata,
 )
-from exo.worker.download.download_utils import build_model_path
 from exo.worker.engines.mlx import Model
 from exo.worker.engines.mlx.auto_parallel import (
     TimeoutCallback,
@@ -178,11 +175,6 @@ def mlx_distributed_init(
                 os.environ["MLX_JACCL_COORDINATOR"] = jaccl_coordinator
                 group = mx.distributed.init(backend="jaccl", strict=True)
 
-            case _:
-                raise ValueError(
-                    f"Unsupported instance type for MLX distributed: {type(bound_instance.instance)}"
-                )
-
         logger.info(f"Rank {rank} mlx distributed initialization complete")
 
         return group
@@ -370,12 +362,35 @@ def load_tokenizer_for_model_id(
     return tokenizer
 
 
+def _normalize_tool_calls(msg_dict: dict[str, Any]) -> None:
+    """
+    Normalize tool_calls in a message dict.
+
+    OpenAI format has tool_calls[].function.arguments as a JSON string,
+    but some chat templates (e.g., GLM) expect it as a dict.
+    """
+    tool_calls = msg_dict.get("tool_calls")
+    if not tool_calls or not isinstance(tool_calls, list):
+        return
+
+    for tc in tool_calls:  # pyright: ignore[reportUnknownVariableType]
+        if not isinstance(tc, dict):
+            continue
+        func = tc.get("function")  # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType]
+        if not isinstance(func, dict):
+            continue
+        args = func.get("arguments")  # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType]
+        if isinstance(args, str):
+            with contextlib.suppress(json.JSONDecodeError):
+                func["arguments"] = json.loads(args)
+
+
 def apply_chat_template(
     tokenizer: TokenizerWrapper,
     chat_task_data: ChatCompletionTaskParams,
 ) -> str:
-    # Now we can properly access the messages
     messages = chat_task_data.messages
+    tools = chat_task_data.tools
 
     formatted_messages: list[dict[str, Any]] = []
     for message in messages:
@@ -387,19 +402,27 @@ def apply_chat_template(
                 continue
 
             message.content = "\n".join(c.text for c in message.content).strip()
-        if message.content is None and message.thinking is None:
+        if (
+            message.content is None
+            and message.thinking is None
+            and message.tool_calls is None
+        ):
             continue
 
         # Null values are not valid when applying templates in tokenizer
-        formatted_messages.append(
-            {k: v for k, v in message.model_dump().items() if v is not None}  # type: ignore
-        )
+        dumped: dict[str, Any] = message.model_dump()
+        msg_dict: dict[str, Any] = {k: v for k, v in dumped.items() if v is not None}  # pyright: ignore[reportAny]
+
+        # Parse tool_calls arguments from JSON string to dict for templates that expect dicts
+        _normalize_tool_calls(msg_dict)
+
+        formatted_messages.append(msg_dict)
 
     prompt: str = tokenizer.apply_chat_template(
         formatted_messages,
         tokenize=False,
         add_generation_prompt=True,
-        tools=chat_task_data.tools,
+        tools=tools,
     )
 
     logger.info(prompt)
@@ -440,31 +463,6 @@ class NullKVCache(KVCache):
         raise NotImplementedError("We should not be setting a NullKVCache.")
 
 
-def make_kv_cache(
-    model: Model, max_kv_size: int | None = None, keep: int = 0
-) -> list[KVCache | RotatingKVCache | QuantizedKVCache]:
-    assert hasattr(model, "layers")
-
-    # TODO: Do this for all models
-    if hasattr(model, "make_cache") and isinstance(model, GptOssModel):
-        logger.info("Using MLX LM's make cache")
-        return model.make_cache()  # type: ignore
-
-    if max_kv_size 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)
-                for _ in model.layers
-            ]
-    else:
-        logger.info(f"Using rotating KV cache with {max_kv_size=} with {keep=}")
-        return [RotatingKVCache(max_size=max_kv_size, keep=keep) for _ in model.layers]
-
-
 def mlx_force_oom(size: int = 40000) -> None:
     """
     Force an Out-Of-Memory (OOM) error in MLX by performing large tensor operations.

src/exo/worker/main.py

@@ -1,8 +1,9 @@
 from datetime import datetime, timezone
 from random import random
+from typing import Iterator
 
 import anyio
-from anyio import CancelScope, create_task_group, current_time, fail_after
+from anyio import CancelScope, create_task_group, fail_after
 from anyio.abc import TaskGroup
 from loguru import logger
 
@@ -10,16 +11,19 @@ from exo.routing.connection_message import ConnectionMessage, ConnectionMessageT
 from exo.shared.apply import apply
 from exo.shared.models.model_cards import ModelId
 from exo.shared.types.api import ImageEditsInternalParams
-from exo.shared.types.commands import ForwarderCommand, RequestEventLog
+from exo.shared.types.commands import (
+    ForwarderCommand,
+    ForwarderDownloadCommand,
+    RequestEventLog,
+    StartDownload,
+)
 from exo.shared.types.common import CommandId, NodeId, SessionId
 from exo.shared.types.events import (
-    BaseEvent,
     Event,
     EventId,
     ForwarderEvent,
     IndexedEvent,
     InputChunkReceived,
-    NodeDownloadProgress,
     NodeGatheredInfo,
     TaskCreated,
     TaskStatusUpdated,
@@ -29,30 +33,21 @@ from exo.shared.types.events import (
 from exo.shared.types.multiaddr import Multiaddr
 from exo.shared.types.state import State
 from exo.shared.types.tasks import (
-    BaseTask,
+    ChatCompletion,
     CreateRunner,
     DownloadModel,
     ImageEdits,
     Shutdown,
+    Task,
     TaskStatus,
 )
 from exo.shared.types.topology import Connection, SocketConnection
-from exo.shared.types.worker.downloads import (
-    DownloadCompleted,
-    DownloadOngoing,
-    DownloadPending,
-    DownloadProgress,
-)
 from exo.shared.types.worker.runners import RunnerId
-from exo.shared.types.worker.shards import ShardMetadata
 from exo.utils.channels import Receiver, Sender, channel
 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.worker.download.download_utils import (
-    map_repo_download_progress_to_download_progress_data,
-)
-from exo.worker.download.shard_downloader import RepoDownloadProgress, ShardDownloader
+from exo.utils.keyed_backoff import KeyedBackoff
 from exo.worker.plan import plan
 from exo.worker.runner.runner_supervisor import RunnerSupervisor
 
@@ -62,7 +57,6 @@ class Worker:
         self,
         node_id: NodeId,
         session_id: SessionId,
-        shard_downloader: ShardDownloader,
         *,
         connection_message_receiver: Receiver[ConnectionMessage],
         global_event_receiver: Receiver[ForwarderEvent],
@@ -70,23 +64,22 @@ class Worker:
         # This is for requesting updates. It doesn't need to be a general command sender right now,
         # but I think it's the correct way to be thinking about commands
         command_sender: Sender[ForwarderCommand],
+        download_command_sender: Sender[ForwarderDownloadCommand],
+        event_index_counter: Iterator[int],
     ):
         self.node_id: NodeId = node_id
         self.session_id: SessionId = session_id
 
-        self.shard_downloader: ShardDownloader = shard_downloader
-        self._pending_downloads: dict[RunnerId, ShardMetadata] = {}
-
         self.global_event_receiver = global_event_receiver
         self.local_event_sender = local_event_sender
-        self.local_event_index = 0
+        self.event_index_counter = event_index_counter
         self.command_sender = command_sender
+        self.download_command_sender = download_command_sender
         self.connection_message_receiver = connection_message_receiver
-        self.event_buffer = OrderedBuffer[BaseEvent]()
+        self.event_buffer = OrderedBuffer[Event]()
         self.out_for_delivery: dict[EventId, ForwarderEvent] = {}
 
         self.state: State = State()
-        self.download_status: dict[ModelId, DownloadProgress] = {}
         self.runners: dict[RunnerId, RunnerSupervisor] = {}
         self._tg: TaskGroup = create_task_group()
 
@@ -101,6 +94,8 @@ class Worker:
         self.input_chunk_buffer: dict[CommandId, dict[int, str]] = {}
         self.input_chunk_counts: dict[CommandId, int] = {}
 
+        self._download_backoff: KeyedBackoff[ModelId] = KeyedBackoff(base=0.5, cap=10.0)
+
     async def run(self):
         logger.info("Starting Worker")
 
@@ -111,7 +106,6 @@ class Worker:
             tg.start_soon(info_gatherer.run)
             tg.start_soon(self._forward_info, info_recv)
             tg.start_soon(self.plan_step)
-            tg.start_soon(self._emit_existing_download_progress)
             tg.start_soon(self._connection_message_event_writer)
             tg.start_soon(self._resend_out_for_delivery)
             tg.start_soon(self._event_applier)
@@ -121,6 +115,7 @@ class Worker:
         # Actual shutdown code - waits for all tasks to complete before executing.
         self.local_event_sender.close()
         self.command_sender.close()
+        self.download_command_sender.close()
         for runner in self.runners.values():
             runner.shutdown()
 
@@ -179,11 +174,9 @@ class Worker:
     async def plan_step(self):
         while True:
             await anyio.sleep(0.1)
-            # 3. based on the updated state, we plan & execute an operation.
-            task: BaseTask | None = plan(
+            task: Task | None = plan(
                 self.node_id,
                 self.runners,
-                self.download_status,
                 self.state.downloads,
                 self.state.instances,
                 self.state.runners,
@@ -192,8 +185,10 @@ class Worker:
                 self.input_chunk_counts,
             )
             if task is None:
+                # Only sleep when there's nothing to do - allows rapid task dispatch
+                await anyio.sleep(0.01)
                 continue
-            logger.info(f"Worker plan: {task.__class__.__name__}")
+            logger.debug(f"Worker plan: {task.__class__.__name__}")
             assert task.task_status
             await self.event_sender.send(TaskCreated(task_id=task.task_id, task=task))
 
@@ -207,42 +202,26 @@ class Worker:
                         )
                     )
                 case DownloadModel(shard_metadata=shard):
-                    if shard.model_card.model_id not in self.download_status:
-                        progress = DownloadPending(
-                            shard_metadata=shard, node_id=self.node_id
-                        )
-                        self.download_status[shard.model_card.model_id] = progress
-                        await self.event_sender.send(
-                            NodeDownloadProgress(download_progress=progress)
-                        )
-                    initial_progress = (
-                        await self.shard_downloader.get_shard_download_status_for_shard(
-                            shard
+                    model_id = shard.model_card.model_id
+                    if not self._download_backoff.should_proceed(model_id):
+                        continue
+
+                    self._download_backoff.record_attempt(model_id)
+
+                    await self.download_command_sender.send(
+                        ForwarderDownloadCommand(
+                            origin=self.node_id,
+                            command=StartDownload(
+                                target_node_id=self.node_id,
+                                shard_metadata=shard,
+                            ),
                         )
                     )
-                    if initial_progress.status == "complete":
-                        progress = DownloadCompleted(
-                            shard_metadata=shard,
-                            node_id=self.node_id,
-                            total_bytes=initial_progress.total_bytes,
+                    await self.event_sender.send(
+                        TaskStatusUpdated(
+                            task_id=task.task_id, task_status=TaskStatus.Running
                         )
-                        self.download_status[shard.model_card.model_id] = progress
-                        await self.event_sender.send(
-                            NodeDownloadProgress(download_progress=progress)
-                        )
-                        await self.event_sender.send(
-                            TaskStatusUpdated(
-                                task_id=task.task_id,
-                                task_status=TaskStatus.Complete,
-                            )
-                        )
-                    else:
-                        await self.event_sender.send(
-                            TaskStatusUpdated(
-                                task_id=task.task_id, task_status=TaskStatus.Running
-                            )
-                        )
-                        self._handle_shard_download_process(task, initial_progress)
+                    )
                 case Shutdown(runner_id=runner_id):
                     try:
                         with fail_after(3):
@@ -293,13 +272,19 @@ class Worker:
                     await self.runners[self._task_to_runner_id(task)].start_task(
                         modified_task
                     )
+                case ChatCompletion():
+                    # Don't wait for acknowledgment for batchable inference tasks
+                    # This allows multiple tasks to reach the runner for batching
+                    await self.runners[self._task_to_runner_id(task)].start_task(
+                        task, wait_for_ack=False
+                    )
                 case task:
                     await self.runners[self._task_to_runner_id(task)].start_task(task)
 
     def shutdown(self):
         self._tg.cancel_scope.cancel()
 
-    def _task_to_runner_id(self, task: BaseTask):
+    def _task_to_runner_id(self, task: Task):
         instance = self.state.instances[task.instance_id]
         return instance.shard_assignments.node_to_runner[self.node_id]
 
@@ -387,78 +372,17 @@ class Worker:
         self._tg.start_soon(runner.run)
         return runner
 
-    def _handle_shard_download_process(
-        self,
-        task: DownloadModel,
-        initial_progress: RepoDownloadProgress,
-    ):
-        """Manages the shard download process with progress tracking."""
-        status = DownloadOngoing(
-            node_id=self.node_id,
-            shard_metadata=task.shard_metadata,
-            download_progress=map_repo_download_progress_to_download_progress_data(
-                initial_progress
-            ),
-        )
-        self.download_status[task.shard_metadata.model_card.model_id] = status
-        self.event_sender.send_nowait(NodeDownloadProgress(download_progress=status))
-
-        last_progress_time = 0.0
-        throttle_interval_secs = 1.0
-
-        async def download_progress_callback(
-            shard: ShardMetadata, progress: RepoDownloadProgress
-        ) -> None:
-            nonlocal self
-            nonlocal last_progress_time
-            if progress.status == "complete":
-                status = DownloadCompleted(
-                    shard_metadata=shard,
-                    node_id=self.node_id,
-                    total_bytes=progress.total_bytes,
-                )
-                self.download_status[shard.model_card.model_id] = status
-                await self.event_sender.send(
-                    NodeDownloadProgress(download_progress=status)
-                )
-                await self.event_sender.send(
-                    TaskStatusUpdated(
-                        task_id=task.task_id, task_status=TaskStatus.Complete
-                    )
-                )
-            elif (
-                progress.status == "in_progress"
-                and current_time() - last_progress_time > throttle_interval_secs
-            ):
-                status = DownloadOngoing(
-                    node_id=self.node_id,
-                    shard_metadata=shard,
-                    download_progress=map_repo_download_progress_to_download_progress_data(
-                        progress
-                    ),
-                )
-                self.download_status[shard.model_card.model_id] = status
-                await self.event_sender.send(
-                    NodeDownloadProgress(download_progress=status)
-                )
-                last_progress_time = current_time()
-
-        self.shard_downloader.on_progress(download_progress_callback)
-        self._tg.start_soon(self.shard_downloader.ensure_shard, task.shard_metadata)
-
     async def _forward_events(self) -> None:
         with self.event_receiver as events:
             async for event in events:
+                idx = next(self.event_index_counter)
                 fe = ForwarderEvent(
-                    origin_idx=self.local_event_index,
+                    origin_idx=idx,
                     origin=self.node_id,
                     session=self.session_id,
                     event=event,
                 )
-                logger.debug(
-                    f"Worker published event {self.local_event_index}: {str(event)[:100]}"
-                )
-                self.local_event_index += 1
+                logger.debug(f"Worker published event {idx}: {str(event)[:100]}")
                 await self.local_event_sender.send(fe)
                 self.out_for_delivery[event.event_id] = fe
 
@@ -506,42 +430,3 @@ class Worker:
                     await self.event_sender.send(TopologyEdgeDeleted(conn=conn))
 
             await anyio.sleep(10)
-
-    async def _emit_existing_download_progress(self) -> None:
-        try:
-            while True:
-                logger.debug("Fetching and emitting existing download progress...")
-                async for (
-                    _,
-                    progress,
-                ) in self.shard_downloader.get_shard_download_status():
-                    if progress.status == "complete":
-                        status = DownloadCompleted(
-                            node_id=self.node_id,
-                            shard_metadata=progress.shard,
-                            total_bytes=progress.total_bytes,
-                        )
-                    elif progress.status in ["in_progress", "not_started"]:
-                        if progress.downloaded_bytes_this_session.in_bytes == 0:
-                            status = DownloadPending(
-                                node_id=self.node_id, shard_metadata=progress.shard
-                            )
-                        else:
-                            status = DownloadOngoing(
-                                node_id=self.node_id,
-                                shard_metadata=progress.shard,
-                                download_progress=map_repo_download_progress_to_download_progress_data(
-                                    progress
-                                ),
-                            )
-                    else:
-                        continue
-
-                    self.download_status[progress.shard.model_card.model_id] = status
-                    await self.event_sender.send(
-                        NodeDownloadProgress(download_progress=status)
-                    )
-                logger.debug("Done emitting existing download progress.")
-                await anyio.sleep(5 * 60)  # 5 minutes
-        except Exception as e:
-            logger.error(f"Error emitting existing download progress: {e}")

src/exo/worker/plan.py

@@ -2,11 +2,10 @@
 
 from collections.abc import Mapping, Sequence
 
-from exo.shared.models.model_cards import ModelId
 from exo.shared.types.common import CommandId, NodeId
 from exo.shared.types.tasks import (
-    BaseTask,
     ChatCompletion,
+    Completion,
     ConnectToGroup,
     CreateRunner,
     DownloadModel,
@@ -15,19 +14,17 @@ from exo.shared.types.tasks import (
     LoadModel,
     Shutdown,
     StartWarmup,
+    Task,
     TaskId,
     TaskStatus,
 )
 from exo.shared.types.worker.downloads import (
     DownloadCompleted,
+    DownloadFailed,
     DownloadOngoing,
     DownloadProgress,
 )
-from exo.shared.types.worker.instances import (
-    BaseInstance,
-    BoundInstance,
-    InstanceId,
-)
+from exo.shared.types.worker.instances import BoundInstance, Instance, InstanceId
 from exo.shared.types.worker.runners import (
     RunnerConnected,
     RunnerConnecting,
@@ -48,31 +45,18 @@ def plan(
     node_id: NodeId,
     # Runners is expected to be FRESH and so should not come from state
     runners: Mapping[RunnerId, RunnerSupervisor],
-    # DL_status is expected to be FRESH and so should not come from state
-    download_status: Mapping[ModelId, DownloadProgress],
-    # gdls is not expected to be fresh
     global_download_status: Mapping[NodeId, Sequence[DownloadProgress]],
-    instances: Mapping[InstanceId, BaseInstance],
+    instances: Mapping[InstanceId, Instance],
     all_runners: Mapping[RunnerId, RunnerStatus],  # all global
-    tasks: Mapping[TaskId, BaseTask],
+    tasks: Mapping[TaskId, Task],
     input_chunk_buffer: Mapping[CommandId, dict[int, str]] | None = None,
     input_chunk_counts: Mapping[CommandId, int] | None = None,
-) -> BaseTask | None:
-    from exo.plugins.registry import PluginRegistry
-
-    registry = PluginRegistry.get()
-
-    # Check plugin tasks first
-    for plugin in registry.all_plugins():
-        task = plugin.plan_task(runners, instances)
-        if task is not None:
-            return task
-
+) -> Task | None:
     # Python short circuiting OR logic should evaluate these sequentially.
     return (
         _kill_runner(runners, all_runners, instances)
         or _create_runner(node_id, runners, instances)
-        or _model_needs_download(runners, download_status)
+        or _model_needs_download(node_id, runners, global_download_status)
         or _init_distributed_backend(runners, all_runners)
         or _load_model(runners, all_runners, global_download_status)
         or _ready_to_warmup(runners, all_runners)
@@ -83,7 +67,7 @@ def plan(
 def _kill_runner(
     runners: Mapping[RunnerId, RunnerSupervisor],
     all_runners: Mapping[RunnerId, RunnerStatus],
-    instances: Mapping[InstanceId, BaseInstance],
+    instances: Mapping[InstanceId, Instance],
 ) -> Shutdown | None:
     for runner in runners.values():
         runner_id = runner.bound_instance.bound_runner_id
@@ -106,7 +90,7 @@ def _kill_runner(
 def _create_runner(
     node_id: NodeId,
     runners: Mapping[RunnerId, RunnerSupervisor],
-    instances: Mapping[InstanceId, BaseInstance],
+    instances: Mapping[InstanceId, Instance],
 ) -> CreateRunner | None:
     for instance in instances.values():
         runner_id = instance.shard_assignments.node_to_runner.get(node_id, None)
@@ -128,26 +112,22 @@ def _create_runner(
 
 
 def _model_needs_download(
+    node_id: NodeId,
     runners: Mapping[RunnerId, RunnerSupervisor],
-    download_status: Mapping[ModelId, DownloadProgress],
+    global_download_status: Mapping[NodeId, Sequence[DownloadProgress]],
 ) -> DownloadModel | None:
-    from exo.plugins.registry import PluginRegistry
-
-    registry = PluginRegistry.get()
+    local_downloads = global_download_status.get(node_id, [])
+    download_status = {
+        dp.shard_metadata.model_card.model_id: dp for dp in local_downloads
+    }
 
     for runner in runners.values():
-        instance = runner.bound_instance.instance
-
-        # Check if any plugin wants to skip download for this instance
-        plugin = registry.get_plugin_for_instance(instance)
-        if plugin is not None and plugin.should_skip_download(instance):
-            continue
-
         model_id = runner.bound_instance.bound_shard.model_card.model_id
         if isinstance(runner.status, RunnerIdle) and (
             model_id not in download_status
             or not isinstance(
-                download_status[model_id], (DownloadOngoing, DownloadCompleted)
+                download_status[model_id],
+                (DownloadOngoing, DownloadCompleted, DownloadFailed),
             )
         ):
             # We don't invalidate download_status randomly in case a file gets deleted on disk
@@ -289,14 +269,16 @@ def _ready_to_warmup(
 
 def _pending_tasks(
     runners: Mapping[RunnerId, RunnerSupervisor],
-    tasks: Mapping[TaskId, BaseTask],
+    tasks: Mapping[TaskId, Task],
     all_runners: Mapping[RunnerId, RunnerStatus],
     input_chunk_buffer: Mapping[CommandId, dict[int, str]] | None = None,
-) -> BaseTask | None:
+) -> Task | None:
     for task in tasks.values():
-        # for now, just forward chat completions
+        # for now, just forward chat completions and completions
         # TODO(ciaran): do this better!
-        if not isinstance(task, (ChatCompletion, ImageGeneration, ImageEdits)):
+        if not isinstance(
+            task, (ChatCompletion, Completion, ImageGeneration, ImageEdits)
+        ):
             continue
         if task.task_status not in (TaskStatus.Pending, TaskStatus.Running):
             continue
@@ -319,9 +301,14 @@ def _pending_tasks(
             if task.task_id in runner.completed:
                 continue
 
+            # Skip tasks already sent to runner (waiting for completion)
+            if task.task_id in runner.sent:
+                continue
+
             # TODO: Check ordering aligns with MLX distributeds expectations.
 
-            if isinstance(runner.status, RunnerReady) and all(
+            # Allow sending tasks when runner is Ready OR Running (for batching)
+            if isinstance(runner.status, (RunnerReady, RunnerRunning)) and all(
                 isinstance(all_runners[global_runner_id], (RunnerReady, RunnerRunning))
                 for global_runner_id in runner.bound_instance.instance.shard_assignments.runner_to_shard
             ):

src/exo/worker/runner/batched_handler.py

@@ -0,0 +1,662 @@
+"""Batched inference handler for processing multiple ChatCompletion requests concurrently."""
+
+import time
+from collections.abc import Generator
+from dataclasses import dataclass, field
+from typing import Any, Callable, Literal
+
+import mlx.core as mx
+from mlx_lm.generate import BatchGenerator
+from mlx_lm.models.gpt_oss import Model as GptOssModel
+from mlx_lm.sample_utils import make_sampler
+from mlx_lm.tokenizer_utils import TokenizerWrapper
+from openai_harmony import (  # pyright: ignore[reportMissingTypeStubs]
+    HarmonyEncodingName,
+    Role,
+    StreamableParser,
+    load_harmony_encoding,
+)
+
+from exo.shared.models.model_cards import ModelId
+from exo.shared.types.api import (
+    GenerationStats,
+    TopLogprobItem,
+)
+from exo.shared.types.chunks import ErrorChunk, TokenChunk
+from exo.shared.types.common import CommandId
+from exo.shared.types.events import ChunkGenerated, Event
+from exo.shared.types.memory import Memory
+from exo.shared.types.tasks import ChatCompletion
+from exo.worker.engines.mlx import Model
+from exo.worker.engines.mlx.constants import MAX_TOKENS
+from exo.worker.engines.mlx.generator.generate import extract_top_logprobs
+from exo.worker.engines.mlx.utils_mlx import apply_chat_template
+from exo.worker.runner.bootstrap import logger
+from exo.worker.runner.pipelined_generator import PipelinedGenerator, PipelinedResponse
+
+# Type alias for the finish_reason values TokenChunk accepts
+TokenFinishReason = Literal["stop", "length", "content_filter"]
+
+
+@dataclass
+class PendingRequest:
+    """A request waiting to be added to the batch."""
+
+    task: ChatCompletion
+    prompt: str
+    max_tokens: int
+    sampler: Callable[[mx.array], mx.array]
+    should_extract_logprobs: bool
+    top_k: int
+
+
+@dataclass
+class ActiveRequest:
+    """A request currently being processed in the batch."""
+
+    command_id: CommandId
+    should_extract_logprobs: bool
+    top_k: int
+    harmony_parser: Any | None = None  # StreamableParser for GPT-OSS models
+    in_thinking: bool = False  # Currently in thinking/reasoning section
+    tokens_generated: int = 0
+    reasoning_tokens: int = 0
+    prompt_tokens: int = 0
+    start_time: float = field(default_factory=time.perf_counter)
+
+
+class BatchedInferenceHandler:
+    """
+    Handles batched inference for multiple ChatCompletion requests.
+
+    Uses MLX-LM's BatchGenerator to process multiple requests concurrently,
+    improving throughput for scenarios with multiple concurrent requests.
+    """
+
+    def __init__(
+        self,
+        model: Model,
+        tokenizer: TokenizerWrapper,
+        model_id: ModelId,
+        device_rank: int,
+        world_size: int = 1,
+        max_batch_size: int = 32,
+    ):
+        self.model = model
+        self.tokenizer = tokenizer
+        self.model_id = model_id
+        self.device_rank = device_rank
+        self.world_size = world_size
+        self.max_batch_size = max_batch_size
+
+        # Model-specific thinking/reasoning detection
+        self.is_gpt_oss = isinstance(model, GptOssModel)
+        self._harmony_encoding: Any | None = None
+        if self.is_gpt_oss:
+            self._harmony_encoding = load_harmony_encoding(
+                HarmonyEncodingName.HARMONY_GPT_OSS
+            )
+            logger.info("GPT-OSS model detected, enabling harmony stream parsing")
+
+        # Detect <think></think> tokens from tokenizer (works for any model)
+        self._think_start_token: int | None = None
+        self._think_end_token: int | None = None
+        think_start: int | None = tokenizer.think_start_id  # pyright: ignore[reportAny]
+        if not self.is_gpt_oss and think_start is not None:
+            self._think_start_token = think_start
+            self._think_end_token = tokenizer.think_end_id  # pyright: ignore[reportAny]
+            logger.info(
+                f"Detected <think></think> tokens ({self._think_start_token}/{self._think_end_token}), enabling reasoning tracking"
+            )
+
+        # Pending requests waiting to be batched
+        self.pending: list[PendingRequest] = []
+
+        # Active batch generator and request tracking
+        self.batch_generator: BatchGenerator | None = None
+        self.pipelined_generator: PipelinedGenerator | None = None
+        self.uid_to_request: dict[int, ActiveRequest] = {}
+
+        # Use pipelined generator for multi-device pipeline parallelism
+        self.use_pipelined = world_size > 1
+        if self.use_pipelined:
+            logger.info(
+                f"Using PipelinedGenerator with {world_size} streams for pipeline overlap"
+            )
+
+        # EOS tokens for the model
+        self.stop_tokens: set[int] = set()
+        eos_ids: list[int] | None = getattr(tokenizer, "eos_token_ids", None)
+        if eos_ids:
+            self.stop_tokens = set(eos_ids)
+
+    @property
+    def is_active(self) -> bool:
+        """Check if there's an active batch being processed."""
+        if self.use_pipelined:
+            return (
+                self.pipelined_generator is not None
+                and self.pipelined_generator.has_active
+            )
+        return self.batch_generator is not None and len(self.uid_to_request) > 0
+
+    @property
+    def has_pending(self) -> bool:
+        """Check if there are pending requests waiting to be batched."""
+        return len(self.pending) > 0
+
+    @property
+    def current_batch_size(self) -> int:
+        """Current number of active requests in the batch."""
+        return len(self.uid_to_request)
+
+    def add_request(self, task: ChatCompletion) -> None:
+        """Add a ChatCompletion request to the pending batch."""
+        task_params = task.task_params
+
+        # Build prompt
+        prompt = apply_chat_template(self.tokenizer, task_params)
+
+        # Determine max tokens
+        max_tokens = task_params.max_tokens or MAX_TOKENS
+
+        # Create sampler for this request
+        sampler = make_sampler(
+            temp=task_params.temperature
+            if task_params.temperature is not None
+            else 0.7,
+            top_p=task_params.top_p if task_params.top_p is not None else 1.0,
+        )
+
+        # Logprobs configuration
+        should_extract_logprobs = task_params.logprobs is True
+        top_k = task_params.top_logprobs if task_params.top_logprobs is not None else 0
+
+        pending_request = PendingRequest(
+            task=task,
+            prompt=prompt,
+            max_tokens=max_tokens,
+            sampler=sampler,
+            should_extract_logprobs=should_extract_logprobs,
+            top_k=top_k,
+        )
+
+        self.pending.append(pending_request)
+
+        logger.info(
+            f"Added request to batch queue (pending={len(self.pending)}, active={self.current_batch_size})"
+        )
+
+    def flush(self) -> None:
+        """Start processing pending requests by adding them to the batch/pipelined generator."""
+        if not self.has_pending:
+            return
+
+        # Determine how many requests to flush (up to available slots)
+        available_slots = self.max_batch_size - self.current_batch_size
+        requests_to_flush = self.pending[:available_slots]
+        self.pending = self.pending[available_slots:]
+
+        # Prepare batch data - tokenize prompts
+        tokenized_prompts: list[list[int]] = []
+        max_tokens_list: list[int] = []
+        samplers: list[Callable[[mx.array], mx.array]] = []
+        prompt_token_counts: list[int] = []
+
+        for req in requests_to_flush:
+            tokens = self.tokenizer.encode(req.prompt)
+            tokenized_prompts.append(tokens)
+            max_tokens_list.append(req.max_tokens)
+            samplers.append(req.sampler)
+            prompt_token_counts.append(len(tokens))
+
+        if self.use_pipelined:
+            self._flush_pipelined(
+                requests_to_flush,
+                tokenized_prompts,
+                max_tokens_list,
+                samplers,
+                prompt_token_counts,
+            )
+        else:
+            self._flush_batch(
+                requests_to_flush,
+                tokenized_prompts,
+                max_tokens_list,
+                samplers,
+                prompt_token_counts,
+            )
+
+    def _flush_pipelined(
+        self,
+        requests_to_flush: list[PendingRequest],
+        tokenized_prompts: list[list[int]],
+        max_tokens_list: list[int],
+        samplers: list[Callable[[mx.array], mx.array]],
+        prompt_token_counts: list[int],
+    ) -> None:
+        """Flush using PipelinedGenerator (multi-stream pipeline overlap)."""
+        if self.pipelined_generator is None:
+            logger.info(
+                f"Creating PipelinedGenerator for {len(requests_to_flush)} requests ({self.world_size} streams)"
+            )
+            mx.reset_peak_memory()
+            self.pipelined_generator = PipelinedGenerator(
+                model=self.model,
+                world_size=self.world_size,
+                stop_tokens=self.stop_tokens if self.stop_tokens else None,
+                max_tokens=MAX_TOKENS,
+            )
+        else:
+            logger.info(
+                f"Adding {len(requests_to_flush)} requests to PipelinedGenerator"
+            )
+
+        uids = self.pipelined_generator.insert(
+            prompts=tokenized_prompts,
+            max_tokens=max_tokens_list,
+            samplers=samplers,
+        )
+
+        for uid, req, prompt_tokens, tokens in zip(
+            uids, requests_to_flush, prompt_token_counts, tokenized_prompts, strict=True
+        ):
+            parser = None
+            if self.is_gpt_oss and self._harmony_encoding is not None:
+                parser = StreamableParser(self._harmony_encoding, role=Role.ASSISTANT)  # pyright: ignore[reportAny]
+            # Check if prompt contains <think> token - if so, model is already in thinking mode
+            starts_in_thinking = (
+                self._think_start_token is not None
+                and self._think_start_token in tokens
+            )
+            self.uid_to_request[uid] = ActiveRequest(
+                command_id=req.task.command_id,
+                should_extract_logprobs=req.should_extract_logprobs,
+                top_k=req.top_k,
+                prompt_tokens=prompt_tokens,
+                harmony_parser=parser,
+                in_thinking=starts_in_thinking,
+            )
+
+        logger.info(
+            f"Flushed {len(requests_to_flush)} requests into pipelined generator (active={self.pipelined_generator.active_count}, uids={list(self.uid_to_request.keys())})"
+        )
+
+    def _flush_batch(
+        self,
+        requests_to_flush: list[PendingRequest],
+        tokenized_prompts: list[list[int]],
+        max_tokens_list: list[int],
+        samplers: list[Callable[[mx.array], mx.array]],
+        prompt_token_counts: list[int],
+    ) -> None:
+        """Flush using BatchGenerator (single-stream, for non-pipeline instances)."""
+        if self.batch_generator is None:
+            logger.info(
+                f"Creating new BatchGenerator for {len(requests_to_flush)} requests"
+            )
+            mx.reset_peak_memory()
+            self.batch_generator = BatchGenerator(
+                model=self.model,
+                max_tokens=MAX_TOKENS,
+                stop_tokens=self.stop_tokens if self.stop_tokens else None,
+                prefill_batch_size=1,
+            )
+        else:
+            logger.info(
+                f"Adding {len(requests_to_flush)} requests to existing BatchGenerator"
+            )
+
+        # Insert into batch generator
+        uids: list[int] = self.batch_generator.insert(  # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType]
+            prompts=tokenized_prompts,
+            max_tokens=max_tokens_list,
+            samplers=samplers,  # pyright: ignore[reportCallIssue]
+        )
+
+        for uid, req, prompt_tokens, tokens in zip(
+            uids, requests_to_flush, prompt_token_counts, tokenized_prompts, strict=True
+        ):  # pyright: ignore[reportUnknownArgumentType]
+            parser = None
+            if self.is_gpt_oss and self._harmony_encoding is not None:
+                parser = StreamableParser(self._harmony_encoding, role=Role.ASSISTANT)  # pyright: ignore[reportAny]
+            # Check if prompt contains <think> token - if so, model is already in thinking mode
+            starts_in_thinking = (
+                self._think_start_token is not None
+                and self._think_start_token in tokens
+            )
+            self.uid_to_request[uid] = ActiveRequest(
+                command_id=req.task.command_id,
+                should_extract_logprobs=req.should_extract_logprobs,
+                top_k=req.top_k,
+                prompt_tokens=prompt_tokens,
+                harmony_parser=parser,
+                in_thinking=starts_in_thinking,
+            )
+
+        logger.info(
+            f"Flushed {len(requests_to_flush)} requests into batch (active={self.current_batch_size}, uids={list(self.uid_to_request.keys())})"
+        )
+
+    def step(self) -> Generator[Event, None, None]:
+        """
+        Process one generation step and yield ChunkGenerated events.
+
+        Returns a generator of events for completed tokens across all active requests.
+        """
+        if self.use_pipelined:
+            yield from self._step_pipelined()
+            return
+
+        if self.batch_generator is None or not self.uid_to_request:
+            return
+
+        # Get next tokens for all active requests
+        # BatchGenerator.next() returns list of Response objects
+        logger.debug(
+            f"BatchGenerator.next() called (active_uids={list(self.uid_to_request.keys())})"
+        )
+        responses: list[Any] = self.batch_generator.next()  # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType]
+        logger.debug(f"BatchGenerator.next() returned {len(responses)} responses")  # pyright: ignore[reportUnknownArgumentType]
+
+        completed_uids: list[int] = []
+
+        for response in responses:  # pyright: ignore[reportUnknownVariableType]
+            uid: int = response.uid  # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType]
+            if uid not in self.uid_to_request:
+                logger.warning(f"Received response for unknown uid: {uid}")
+                continue
+
+            active_request = self.uid_to_request[uid]
+            active_request.tokens_generated += 1
+
+            # Extract response fields with explicit typing
+            resp_token: int = response.token  # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType]
+            resp_finish_reason: str | None = response.finish_reason  # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType]
+            resp_logprobs: mx.array = response.logprobs  # pyright: ignore[reportUnknownMemberType,reportUnknownVariableType]
+
+            # Only emit events from device_rank 0
+            if self.device_rank != 0:
+                if resp_finish_reason is not None:
+                    completed_uids.append(uid)  # pyright: ignore[reportUnknownArgumentType]
+                continue
+
+            # Decode token to text
+            # Skip emitting EOS token text (e.g., <|eot_id|>)
+            if resp_token in self.stop_tokens:
+                token_text = ""
+            else:
+                token_text = self.tokenizer.decode([resp_token])
+
+            # Handle thinking/reasoning token tracking
+            if active_request.harmony_parser is not None:
+                # GPT-OSS: Use harmony parser for channel-based thinking detection
+                parser = active_request.harmony_parser  # pyright: ignore[reportAny]
+                parser.process(resp_token)  # pyright: ignore[reportAny]
+                delta: str | None = parser.last_content_delta  # pyright: ignore[reportAny]
+                channel: str = parser.current_channel  # pyright: ignore[reportAny]
+
+                # Track reasoning tokens (analysis channel = thinking)
+                if channel == "analysis":
+                    active_request.reasoning_tokens += 1
+
+                # Handle thinking tag transitions
+                prefix = ""
+                if channel == "analysis" and not active_request.in_thinking:
+                    active_request.in_thinking = True
+                    prefix = "<think>"
+                elif channel != "analysis" and active_request.in_thinking:
+                    active_request.in_thinking = False
+                    prefix = "</think>"
+
+                if resp_finish_reason is not None and active_request.in_thinking:
+                    # Close thinking tag on finish
+                    prefix = "</think>"
+                    active_request.in_thinking = False
+
+                effective_delta = delta or ""
+                token_text = (
+                    prefix + effective_delta if (prefix or effective_delta) else ""
+                )
+                # Skip empty tokens (channel markers with no content delta)
+                if not token_text and resp_finish_reason is None:
+                    continue
+            elif self._think_start_token is not None:
+                # MiniMax: Track <think>/</ think> tokens directly
+                if resp_token == self._think_start_token:
+                    active_request.in_thinking = True
+                elif resp_token == self._think_end_token:
+                    active_request.in_thinking = False
+                elif active_request.in_thinking:
+                    active_request.reasoning_tokens += 1
+
+            # Extract logprobs if requested
+            logprob: float | None = None
+            top_logprobs: list[TopLogprobItem] | None = None
+            if active_request.should_extract_logprobs:
+                logprob, top_logprobs = extract_top_logprobs(
+                    logprobs_array=resp_logprobs,  # pyright: ignore[reportUnknownArgumentType]
+                    selected_token=resp_token,  # pyright: ignore[reportUnknownArgumentType]
+                    tokenizer=self.tokenizer,
+                    top_k=active_request.top_k,
+                )
+
+            # Build stats for final token
+            stats: GenerationStats | None = None
+            finish_reason: TokenFinishReason | None = None
+            if resp_finish_reason is not None:
+                elapsed_time = time.perf_counter() - active_request.start_time
+                prompt_tps = active_request.prompt_tokens / max(elapsed_time, 0.001)
+                generation_tps = active_request.tokens_generated / max(
+                    elapsed_time, 0.001
+                )
+
+                # Get peak memory
+                peak_memory_bytes = 0
+                if mx.metal.is_available():
+                    peak_memory_bytes = mx.metal.get_peak_memory()
+
+                stats = GenerationStats(
+                    prompt_tps=prompt_tps,
+                    generation_tps=generation_tps,
+                    prompt_tokens=active_request.prompt_tokens,
+                    generation_tokens=active_request.tokens_generated,
+                    reasoning_tokens=active_request.reasoning_tokens,
+                    peak_memory_usage=Memory.from_bytes(peak_memory_bytes),
+                )
+
+                # Map finish reason to the narrower type TokenChunk expects
+                if resp_finish_reason == "stop":
+                    finish_reason = "stop"
+                elif resp_finish_reason == "length":
+                    finish_reason = "length"
+                elif resp_finish_reason == "content_filter":
+                    finish_reason = "content_filter"
+                else:
+                    # Unknown finish reasons default to "stop"
+                    logger.warning(
+                        f"Unknown finish_reason: {resp_finish_reason}, mapping to 'stop'"
+                    )
+                    finish_reason = "stop"
+
+                completed_uids.append(uid)  # pyright: ignore[reportUnknownArgumentType]
+
+            yield ChunkGenerated(
+                command_id=active_request.command_id,
+                chunk=TokenChunk(
+                    model=self.model_id,
+                    text=token_text,
+                    token_id=resp_token,  # pyright: ignore[reportUnknownArgumentType]
+                    logprob=logprob,
+                    top_logprobs=top_logprobs,
+                    finish_reason=finish_reason,
+                    stats=stats,
+                ),
+            )
+
+        # Clean up completed requests
+        for uid in completed_uids:
+            del self.uid_to_request[uid]
+
+    def _step_pipelined(self) -> Generator[Event, None, None]:
+        """Process one generation step using the multi-stream PipelinedGenerator."""
+        if self.pipelined_generator is None or not self.uid_to_request:
+            return
+
+        logger.debug(
+            f"PipelinedGenerator.next() called (active={self.pipelined_generator.active_count})"
+        )
+        responses: list[PipelinedResponse] = self.pipelined_generator.next()
+        logger.debug(f"PipelinedGenerator.next() returned {len(responses)} responses")
+
+        completed_uids: list[int] = []
+
+        for response in responses:
+            uid = response.uid
+            if uid not in self.uid_to_request:
+                logger.warning(f"Received response for unknown uid: {uid}")
+                continue
+
+            active_request = self.uid_to_request[uid]
+            active_request.tokens_generated += 1
+
+            resp_token: int = response.token
+            resp_finish_reason: str | None = response.finish_reason
+            resp_logprobs: mx.array = response.logprobs
+
+            # Only emit events from device_rank 0
+            if self.device_rank != 0:
+                if resp_finish_reason is not None:
+                    completed_uids.append(uid)
+                continue
+
+            # Decode token to text
+            # Skip emitting EOS token text (e.g., <|eot_id|>)
+            if resp_token in self.stop_tokens:
+                token_text = ""
+            else:
+                token_text = self.tokenizer.decode([resp_token])
+
+            # Handle thinking/reasoning token tracking
+            if active_request.harmony_parser is not None:
+                # GPT-OSS: Use harmony parser for channel-based thinking detection
+                parser = active_request.harmony_parser  # pyright: ignore[reportAny]
+                parser.process(resp_token)  # pyright: ignore[reportAny]
+                delta: str | None = parser.last_content_delta  # pyright: ignore[reportAny]
+                channel: str = parser.current_channel  # pyright: ignore[reportAny]
+
+                if channel == "analysis":
+                    active_request.reasoning_tokens += 1
+
+                prefix = ""
+                if channel == "analysis" and not active_request.in_thinking:
+                    active_request.in_thinking = True
+                    prefix = "<think>"
+                elif channel != "analysis" and active_request.in_thinking:
+                    active_request.in_thinking = False
+                    prefix = "</think>"
+
+                if resp_finish_reason is not None and active_request.in_thinking:
+                    prefix = "</think>"
+                    active_request.in_thinking = False
+
+                effective_delta = delta or ""
+                token_text = (
+                    prefix + effective_delta if (prefix or effective_delta) else ""
+                )
+                if not token_text and resp_finish_reason is None:
+                    continue
+            elif self._think_start_token is not None:
+                # MiniMax: Track <think>/</think> tokens directly
+                if resp_token == self._think_start_token:
+                    active_request.in_thinking = True
+                elif resp_token == self._think_end_token:
+                    active_request.in_thinking = False
+                elif active_request.in_thinking:
+                    active_request.reasoning_tokens += 1
+
+            # Extract logprobs if requested
+            logprob: float | None = None
+            top_logprobs: list[TopLogprobItem] | None = None
+            if active_request.should_extract_logprobs:
+                logprob, top_logprobs = extract_top_logprobs(
+                    logprobs_array=resp_logprobs,
+                    selected_token=resp_token,
+                    tokenizer=self.tokenizer,
+                    top_k=active_request.top_k,
+                )
+
+            # Build stats for final token
+            stats: GenerationStats | None = None
+            finish_reason: TokenFinishReason | None = None
+            if resp_finish_reason is not None:
+                elapsed_time = time.perf_counter() - active_request.start_time
+                prompt_tps = active_request.prompt_tokens / max(elapsed_time, 0.001)
+                generation_tps = active_request.tokens_generated / max(
+                    elapsed_time, 0.001
+                )
+
+                peak_memory_bytes = 0
+                if mx.metal.is_available():
+                    peak_memory_bytes = mx.metal.get_peak_memory()
+
+                stats = GenerationStats(
+                    prompt_tps=prompt_tps,
+                    generation_tps=generation_tps,
+                    prompt_tokens=active_request.prompt_tokens,
+                    generation_tokens=active_request.tokens_generated,
+                    reasoning_tokens=active_request.reasoning_tokens,
+                    peak_memory_usage=Memory.from_bytes(peak_memory_bytes),
+                )
+
+                if resp_finish_reason == "stop":
+                    finish_reason = "stop"
+                elif resp_finish_reason == "length":
+                    finish_reason = "length"
+                else:
+                    finish_reason = "stop"
+
+                completed_uids.append(uid)
+
+            yield ChunkGenerated(
+                command_id=active_request.command_id,
+                chunk=TokenChunk(
+                    model=self.model_id,
+                    text=token_text,
+                    token_id=resp_token,
+                    logprob=logprob,
+                    top_logprobs=top_logprobs,
+                    finish_reason=finish_reason,
+                    stats=stats,
+                ),
+            )
+
+        for uid in completed_uids:
+            del self.uid_to_request[uid]
+
+    def emit_error(self, command_id: CommandId, error_message: str) -> Event:
+        """Create an error event for a failed request."""
+        return ChunkGenerated(
+            command_id=command_id,
+            chunk=ErrorChunk(
+                model=self.model_id,
+                finish_reason="error",
+                error_message=error_message,
+            ),
+        )
+
+    def _close_generator(self) -> None:
+        """Close and clean up the batch/pipelined generator."""
+        if self.batch_generator is not None:
+            self.batch_generator.close()  # pyright: ignore[reportUnknownMemberType,reportAttributeAccessIssue]
+            self.batch_generator = None
+        if self.pipelined_generator is not None:
+            self.pipelined_generator.close()
+            self.pipelined_generator = None
+        self.uid_to_request.clear()
+        logger.info("Generator closed")
+
+    def close(self) -> None:
+        """Close the handler and clean up resources."""
+        self._close_generator()
+        self.pending.clear()

src/exo/worker/runner/batched_scoring_handler.py

@@ -0,0 +1,200 @@
+"""Batched scoring handler for processing multiple Completion requests concurrently."""
+
+import time
+from dataclasses import dataclass, field
+
+from mlx_lm.tokenizer_utils import TokenizerWrapper
+
+from exo.shared.models.model_cards import ModelId
+from exo.shared.types.api import TopLogprobItem
+from exo.shared.types.chunks import CompletionChunk, ErrorChunk
+from exo.shared.types.events import ChunkGenerated, Event
+from exo.shared.types.tasks import Completion
+from exo.worker.engines.mlx import Model
+from exo.worker.engines.mlx.generator.generate import score_tokens_batched
+from exo.worker.runner.bootstrap import logger
+
+
+@dataclass
+class PendingScoringRequest:
+    """A scoring request waiting to be batched."""
+
+    task: Completion
+    tokens: list[int]
+    prompt_text: str
+    top_k: int | None
+    echo: bool
+
+
+@dataclass
+class BatchedScoringHandler:
+    """
+    Handles batched scoring for multiple Completion requests.
+
+    Collects multiple scoring requests and processes them in a single
+    batched forward pass for improved throughput.
+    """
+
+    model: Model
+    tokenizer: TokenizerWrapper
+    model_id: ModelId
+    device_rank: int
+    max_batch_size: int = 32
+    batch_timeout_ms: int = 10
+
+    pending: list[PendingScoringRequest] = field(default_factory=list)
+    pending_start_time: float | None = None
+
+    @property
+    def has_pending(self) -> bool:
+        """Check if there are pending requests."""
+        return len(self.pending) > 0
+
+    def add_request(
+        self,
+        task: Completion,
+        tokens: list[int],
+        prompt_text: str,
+    ) -> None:
+        """Add a Completion request to the pending batch."""
+        task_params = task.task_params
+        top_k = task_params.logprobs
+
+        self.pending.append(
+            PendingScoringRequest(
+                task=task,
+                tokens=tokens,
+                prompt_text=prompt_text,
+                top_k=top_k,
+                echo=task_params.echo,
+            )
+        )
+
+        if self.pending_start_time is None:
+            self.pending_start_time = time.perf_counter()
+
+        logger.debug(f"Added scoring request to batch (pending={len(self.pending)})")
+
+    def should_flush(self) -> bool:
+        """Check if the batch should be flushed."""
+        if not self.has_pending:
+            return False
+
+        # Flush if batch is full
+        if len(self.pending) >= self.max_batch_size:
+            return True
+
+        # Flush if timeout reached
+        if self.pending_start_time is not None:
+            elapsed_ms = (time.perf_counter() - self.pending_start_time) * 1000
+            if elapsed_ms >= self.batch_timeout_ms:
+                return True
+
+        return False
+
+    def flush(self) -> list[Event]:
+        """Process all pending requests and return events."""
+        if not self.has_pending:
+            return []
+
+        requests = self.pending
+        self.pending = []
+        self.pending_start_time = None
+
+        logger.info(f"Processing batch of {len(requests)} scoring requests")
+
+        # Collect all token sequences
+        token_sequences = [req.tokens for req in requests]
+
+        # Get common top_k (use first request's top_k, they should all be the same)
+        top_k = requests[0].top_k if requests else None
+
+        try:
+            # Run batched scoring
+            all_results = score_tokens_batched(
+                model=self.model,
+                tokenizer=self.tokenizer,
+                token_sequences=token_sequences,
+                top_k=top_k,
+            )
+
+            # Generate events for each request
+            events: list[Event] = []
+            for req, logprob_results in zip(requests, all_results, strict=True):
+                if self.device_rank != 0:
+                    continue
+
+                event = self._build_completion_event(req, logprob_results)
+                events.append(event)
+
+            logger.info(f"Batch scoring complete ({len(events)} events)")
+            return events
+
+        except Exception as e:
+            # Return error events for all requests
+            logger.error(f"Batch scoring failed: {e}")
+            events = []
+            for req in requests:
+                if self.device_rank == 0:
+                    events.append(
+                        ChunkGenerated(
+                            command_id=req.task.command_id,
+                            chunk=ErrorChunk(
+                                model=self.model_id,
+                                finish_reason="error",
+                                error_message=str(e),
+                            ),
+                        )
+                    )
+            return events
+
+    def _build_completion_event(
+        self,
+        req: PendingScoringRequest,
+        logprob_results: list[tuple[float, list[TopLogprobItem]]],
+    ) -> Event:
+        """Build a ChunkGenerated event for a completed scoring request."""
+        tokens = req.tokens
+        tokenizer = self.tokenizer
+
+        # Build response in completions format
+        token_strings: list[str] = []
+        token_logprobs: list[float | None] = []
+        top_logprobs: list[dict[str, float]] = []
+        text_offset: list[int] = []
+
+        offset = 0
+        for i, token_id in enumerate(tokens):
+            token_str = tokenizer.decode([token_id])
+            token_strings.append(token_str)
+
+            if i < len(logprob_results):
+                logprob, top_items = logprob_results[i]
+                # First token has no logprob (None in OpenAI format)
+                token_logprobs.append(logprob if i > 0 else None)
+                top_lp_dict = {item.token: item.logprob for item in top_items}
+                top_logprobs.append(top_lp_dict)
+            else:
+                token_logprobs.append(None)
+                top_logprobs.append({})
+
+            text_offset.append(offset)
+            offset += len(token_str)
+
+        return ChunkGenerated(
+            command_id=req.task.command_id,
+            chunk=CompletionChunk(
+                model=self.model_id,
+                text=req.prompt_text if req.echo else "",
+                tokens=token_strings,
+                token_logprobs=token_logprobs,
+                top_logprobs=top_logprobs,
+                text_offset=text_offset,
+                finish_reason="stop",
+            ),
+        )
+
+    def close(self) -> None:
+        """Clean up resources."""
+        self.pending.clear()
+        self.pending_start_time = None

src/exo/worker/runner/bootstrap.py

@@ -4,10 +4,7 @@ import loguru
 
 from exo.shared.types.events import Event, RunnerStatusUpdated
 from exo.shared.types.tasks import Task
-from exo.shared.types.worker.instances import (
-    BoundInstance,
-    MlxJacclInstance,
-)
+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
 
@@ -20,7 +17,6 @@ def entrypoint(
     task_receiver: MpReceiver[Task],
     _logger: "loguru.Logger",
 ) -> None:
-    # Set FAST_SYNCH based on env var or JACCL device count
     fast_synch_override = os.environ.get("EXO_FAST_SYNCH")
     if fast_synch_override == "on" or (
         fast_synch_override != "off"
@@ -38,26 +34,11 @@ def entrypoint(
 
     logger.info(f"Fast synch flag: {os.environ['MLX_METAL_FAST_SYNCH']}")
 
-    # Route based on instance type (plugins or default MLX)
+    # Import main after setting global logger - this lets us just import logger from this module
     try:
-        from exo.plugins.registry import PluginRegistry, discover_plugins
+        from exo.worker.runner.runner import main
 
-        # Discover plugins in subprocess (they aren't inherited from main process)
-        discover_plugins()
-
-        registry = PluginRegistry.get()
-        instance = bound_instance.instance
-
-        # Check if a plugin handles this instance type
-        plugin = registry.get_plugin_for_instance(instance)
-        if plugin is not None:
-            # Delegate to plugin runner
-            plugin.create_runner(bound_instance, event_sender, task_receiver)
-        else:
-            # MLX runner (default)
-            from exo.worker.runner.runner import main
-
-            main(bound_instance, event_sender, task_receiver)
+        main(bound_instance, event_sender, task_receiver)
     except ClosedResourceError:
         logger.warning("Runner communication closed unexpectedly")
     except Exception as e:

src/exo/worker/runner/pipelined_generator.py

@@ -0,0 +1,334 @@
+"""Multi-stream pipelined batch generator for pipeline-parallel inference.
+
+When a model is split across N ranks (pipeline parallelism), each rank's GPU is idle
+for (N-1)/N of each step while waiting for other ranks to compute their layers.
+
+This module fills the pipeline bubble by splitting sequences into N micro-batch groups
+and processing each group on a different MLX stream. The GPU can overlap one stream's
+network communication (send/recv/all_gather) with another stream's compute.
+"""
+
+# pyright: reportUnknownMemberType=false, reportUnknownVariableType=false
+# pyright: reportUnknownArgumentType=false, reportAny=false
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Any
+
+import mlx.core as mx
+import mlx.nn as nn
+from mlx_lm.models.cache import make_prompt_cache
+
+
+@dataclass
+class MicroBatch:
+    """State for one micro-batch group of sequences."""
+
+    uids: list[int]
+    y: mx.array  # Last sampled tokens [batch]
+    logprobs: list[mx.array]  # Logprobs for each sequence
+    max_tokens: list[int]
+    num_tokens: list[int]
+    cache: list[Any]  # KV cache (list of layer caches)
+    samplers: list[Callable[[mx.array], mx.array]]
+    tokens: list[mx.array]  # All tokens generated so far per sequence
+
+    def __len__(self) -> int:
+        return len(self.uids)
+
+
+@dataclass
+class PipelinedResponse:
+    """Response from one generation step."""
+
+    uid: int
+    token: int
+    logprobs: mx.array
+    finish_reason: str | None
+    cache: list[Any] | None = None
+
+
+@dataclass
+class PendingPrompt:
+    """A prompt waiting to be prefilled."""
+
+    uid: int
+    tokens: list[int]
+    max_tokens: int
+    sampler: Callable[[mx.array], mx.array]
+
+
+class PipelinedGenerator:
+    """
+    Multi-stream batch generator that fills pipeline bubbles.
+
+    Splits active sequences into `world_size` micro-batch groups, each processed
+    on its own MLX stream. During mx.eval(), the GPU overlaps network operations
+    on one stream with compute on another.
+    """
+
+    def __init__(
+        self,
+        model: nn.Module,
+        world_size: int,
+        stop_tokens: set[int] | None = None,
+        max_tokens: int = 4096,
+    ):
+        self.model = model
+        self.world_size = world_size
+        self.stop_tokens = stop_tokens or set()
+        self.max_tokens_default = max_tokens
+
+        # Create one stream per pipeline stage
+        self.streams = [mx.new_stream(mx.default_device()) for _ in range(world_size)]
+
+        # Micro-batch groups (one per stream)
+        self.micro_batches: list[MicroBatch | None] = [None] * world_size
+
+        # Pending prompts to be inserted
+        self.pending_prompts: list[PendingPrompt] = []
+
+        # UID counter
+        self._next_uid = 0
+
+    @property
+    def active_count(self) -> int:
+        """Total number of active sequences across all micro-batches."""
+        return sum(len(mb) for mb in self.micro_batches if mb is not None)
+
+    @property
+    def has_active(self) -> bool:
+        return self.active_count > 0 or len(self.pending_prompts) > 0
+
+    def insert(
+        self,
+        prompts: list[list[int]],
+        max_tokens: list[int],
+        samplers: list[Callable[[mx.array], mx.array]],
+    ) -> list[int]:
+        """Queue prompts for processing. Returns assigned UIDs."""
+        uids: list[int] = []
+        for prompt, mt, sampler in zip(prompts, max_tokens, samplers, strict=True):
+            uid = self._next_uid
+            self._next_uid += 1
+            self.pending_prompts.append(
+                PendingPrompt(uid=uid, tokens=prompt, max_tokens=mt, sampler=sampler)
+            )
+            uids.append(uid)
+        return uids
+
+    def _prefill_group(self, group_idx: int, prompts: list[PendingPrompt]) -> None:
+        """Prefill a group of prompts and create a MicroBatch."""
+        if not prompts:
+            return
+
+        stream = self.streams[group_idx]
+
+        with mx.stream(stream):
+            # Create per-sequence caches
+            caches = [make_prompt_cache(self.model) for _ in prompts]
+
+            # Tokenize and prefill each sequence
+            all_y: list[mx.array] = []
+            all_logprobs: list[mx.array] = []
+            all_samplers: list[Callable[[mx.array], mx.array]] = []
+            all_tokens: list[mx.array] = []
+
+            for prompt_info, cache in zip(prompts, caches, strict=True):
+                tokens = mx.array(prompt_info.tokens)
+                # Run prefill (process all tokens except last)
+                if len(prompt_info.tokens) > 1:
+                    self.model(tokens[:-1][None, :], cache=cache)
+                    mx.eval([c.state for c in cache])
+
+                # Process last token to get first generation logits
+                last_token = tokens[-1:][None, :]
+                logits = self.model(last_token, cache=cache)
+                logits = logits[:, -1, :]
+                logprobs = logits - mx.logsumexp(logits, axis=-1, keepdims=True)
+                sampled = prompt_info.sampler(logprobs)
+
+                all_y.append(sampled.squeeze(0))
+                all_logprobs.append(logprobs.squeeze(0))
+                all_samplers.append(prompt_info.sampler)
+                all_tokens.append(tokens)
+
+            mx.eval(*all_y, *all_logprobs)
+
+        # Create micro-batch
+        batch = MicroBatch(
+            uids=[p.uid for p in prompts],
+            y=mx.stack(all_y),
+            logprobs=all_logprobs,
+            max_tokens=[p.max_tokens for p in prompts],
+            num_tokens=[0] * len(prompts),
+            cache=caches,
+            samplers=all_samplers,
+            tokens=all_tokens,
+        )
+
+        if self.micro_batches[group_idx] is None:
+            self.micro_batches[group_idx] = batch
+        else:
+            # Extend existing micro-batch (would need cache merging - for now replace)
+            self.micro_batches[group_idx] = batch
+
+    def _prefill_pending(self) -> None:
+        """Distribute pending prompts across micro-batch groups and prefill."""
+        if not self.pending_prompts:
+            return
+
+        # Distribute round-robin across groups
+        groups: list[list[PendingPrompt]] = [[] for _ in range(self.world_size)]
+        for i, prompt in enumerate(self.pending_prompts):
+            groups[i % self.world_size].append(prompt)
+        self.pending_prompts.clear()
+
+        for group_idx, group_prompts in enumerate(groups):
+            if group_prompts:
+                self._prefill_group(group_idx, group_prompts)
+
+    def _step_all(self) -> None:
+        """
+        Run one generation step across all micro-batch groups on different streams.
+
+        This is where pipeline overlap happens: each group's model forward pass
+        runs on its own stream, and mx.eval() allows the GPU to overlap network
+        ops (send/recv/all_gather) from one stream with compute from another.
+
+        Each sequence is processed individually with its own KV cache, but all
+        lazy graphs across streams are evaluated together for GPU overlap.
+        """
+        # Build computation graphs on each stream (lazy, no evaluation yet)
+        # Each micro-batch group processes its sequences on its own stream.
+        all_sampled: list[mx.array] = []
+        all_logprobs: list[mx.array] = []
+        # Track which (group_idx, seq_idx) each result corresponds to
+        result_map: list[tuple[int, int]] = []
+
+        for i, mb in enumerate(self.micro_batches):
+            if mb is None or len(mb) == 0:
+                continue
+
+            with mx.stream(self.streams[i]):
+                for e in range(len(mb)):
+                    # Process each sequence individually with its own cache
+                    input_token = mb.y[e : e + 1][None, :]  # [1, 1]
+
+                    # Forward pass (lazy graph construction)
+                    # For pipeline models, this includes send/recv/all_gather ops
+                    logits = self.model(input_token, cache=mb.cache[e])
+                    logits = logits[:, -1, :]  # [1, vocab]
+
+                    # Compute logprobs
+                    logprobs = logits - mx.logsumexp(logits, axis=-1, keepdims=True)
+
+                    # Sample
+                    sampled = mb.samplers[e](logprobs)
+
+                    all_sampled.append(sampled.squeeze(0))
+                    all_logprobs.append(logprobs.squeeze(0))
+                    result_map.append((i, e))
+
+        if not result_map:
+            return
+
+        # Evaluate ALL streams together - this is where overlap happens!
+        # The GPU can execute stream0's all_gather while computing stream1's layers.
+        mx.eval(*all_sampled, *all_logprobs)
+
+        # Update micro-batch states with results
+        # Group results by micro-batch for efficient update
+        group_results: dict[int, list[int]] = {}
+        for idx, (group_idx, _seq_idx) in enumerate(result_map):
+            group_results.setdefault(group_idx, []).append(idx)
+
+        for group_idx, result_indices in group_results.items():
+            mb = self.micro_batches[group_idx]
+            assert mb is not None
+            group_sampled = [all_sampled[idx] for idx in result_indices]
+            group_logprobs = [all_logprobs[idx] for idx in result_indices]
+            mb.y = mx.stack(group_sampled)
+            mb.logprobs = group_logprobs
+            for e, idx in enumerate(result_indices):
+                mb.tokens[e] = mx.concatenate([mb.tokens[e], all_sampled[idx][None]])
+
+    def next(self) -> list[PipelinedResponse]:
+        """
+        Run one generation step and return responses.
+
+        Returns a PipelinedResponse for each active sequence (across all groups).
+        Finished sequences are removed from their micro-batch.
+        """
+        # Prefill any pending prompts first
+        self._prefill_pending()
+
+        if not self.has_active:
+            return []
+
+        # Run the multi-stream forward pass
+        self._step_all()
+
+        # Collect responses and filter completed sequences
+        responses: list[PipelinedResponse] = []
+
+        for group_idx, mb in enumerate(self.micro_batches):
+            if mb is None or len(mb) == 0:
+                continue
+
+            keep_idx: list[int] = []
+            end_idx: list[int] = []
+
+            for e in range(len(mb)):
+                token = int(mb.y[e].item())
+                uid = mb.uids[e]
+                num_tok = mb.num_tokens[e] + 1
+                max_tok = mb.max_tokens[e]
+                mb.num_tokens[e] = num_tok
+
+                if token in self.stop_tokens:
+                    finish_reason = "stop"
+                    end_idx.append(e)
+                elif num_tok >= max_tok:
+                    finish_reason = "length"
+                    end_idx.append(e)
+                else:
+                    finish_reason = None
+                    keep_idx.append(e)
+
+                responses.append(
+                    PipelinedResponse(
+                        uid=uid,
+                        token=token,
+                        logprobs=mb.logprobs[e],
+                        finish_reason=finish_reason,
+                    )
+                )
+
+            # Remove finished sequences
+            if end_idx:
+                if keep_idx:
+                    # Filter the micro-batch to keep only active sequences
+                    mb.uids = [mb.uids[i] for i in keep_idx]
+                    mb.y = mb.y[mx.array(keep_idx)]
+                    mb.logprobs = [mb.logprobs[i] for i in keep_idx]
+                    mb.max_tokens = [mb.max_tokens[i] for i in keep_idx]
+                    mb.num_tokens = [mb.num_tokens[i] for i in keep_idx]
+                    mb.samplers = [mb.samplers[i] for i in keep_idx]
+                    mb.tokens = [mb.tokens[i] for i in keep_idx]
+                    # Cache filtering: trim batch dimension
+                    for c in mb.cache:
+                        if hasattr(c, "keys") and c.keys is not None:
+                            c.keys = c.keys[mx.array(keep_idx)]
+                            c.values = c.values[mx.array(keep_idx)]
+                else:
+                    self.micro_batches[group_idx] = None
+
+        return responses
+
+    def close(self) -> None:
+        """Clean up resources."""
+        self.micro_batches = [None] * self.world_size
+        self.pending_prompts.clear()

src/exo/worker/runner/runner_supervisor.py

@@ -20,7 +20,7 @@ from exo.shared.types.events import (
     TaskAcknowledged,
     TaskStatusUpdated,
 )
-from exo.shared.types.tasks import BaseTask, TaskId, TaskStatus
+from exo.shared.types.tasks import Task, TaskId, TaskStatus
 from exo.shared.types.worker.instances import BoundInstance
 from exo.shared.types.worker.runners import (
     RunnerConnecting,
@@ -47,11 +47,14 @@ class RunnerSupervisor:
     runner_process: Process
     initialize_timeout: float
     _ev_recv: MpReceiver[Event]
-    _task_sender: MpSender[BaseTask]
+    _task_sender: MpSender[Task]
     _event_sender: Sender[Event]
     _tg: TaskGroup | None = field(default=None, init=False)
     status: RunnerStatus = field(default_factory=RunnerIdle, init=False)
     pending: dict[TaskId, anyio.Event] = field(default_factory=dict, init=False)
+    sent: set[TaskId] = field(
+        default_factory=set, init=False
+    )  # Tasks sent to runner (not yet completed)
     completed: set[TaskId] = field(default_factory=set, init=False)
 
     @classmethod
@@ -64,7 +67,7 @@ class RunnerSupervisor:
     ) -> Self:
         ev_send, ev_recv = mp_channel[Event]()
         # A task is kind of a runner command
-        task_sender, task_recv = mp_channel[BaseTask]()
+        task_sender, task_recv = mp_channel[Task]()
 
         runner_process = Process(
             target=entrypoint,
@@ -126,21 +129,39 @@ class RunnerSupervisor:
         assert self._tg
         self._tg.cancel_scope.cancel()
 
-    async def start_task(self, task: BaseTask):
+    async def start_task(self, task: Task, wait_for_ack: bool = True):
+        """
+        Send a task to the runner.
+
+        Args:
+            task: The task to send.
+            wait_for_ack: If True, wait for TaskAcknowledged before returning.
+                          If False, return immediately after sending (for batching).
+        """
         if task.task_id in self.completed:
-            logger.info(
-                f"Skipping invalid task {task} as it has already been completed"
+            logger.debug(
+                f"Skipping task {task.task_id} as it has already been completed"
             )
+            return
+        if task.task_id in self.sent:
+            logger.debug(f"Task {task.task_id} already sent, skipping duplicate")
+            return
+        if task.task_id in self.pending:
+            logger.debug(f"Task {task.task_id} already pending, skipping duplicate")
+            return
         logger.info(f"Starting task {task}")
         event = anyio.Event()
         self.pending[task.task_id] = event
+        self.sent.add(task.task_id)
         try:
             self._task_sender.send(task)
         except ClosedResourceError:
             logger.warning(f"Task {task} dropped, runner closed communication.")
+            self.sent.discard(task.task_id)
             return
-        await event.wait()
-        logger.info(f"Finished task {task}")
+        if wait_for_ack:
+            await event.wait()
+            logger.info(f"Finished task {task}")
 
     async def _forward_events(self):
         with self._ev_recv as events:
@@ -149,7 +170,11 @@ class RunnerSupervisor:
                     if isinstance(event, RunnerStatusUpdated):
                         self.status = event.runner_status
                     if isinstance(event, TaskAcknowledged):
-                        self.pending.pop(event.task_id).set()
+                        # Use pop with default to handle tasks sent with wait_for_ack=False
+                        # that may have already been removed or never added
+                        pending_event = self.pending.pop(event.task_id, None)
+                        if pending_event:
+                            pending_event.set()
                         continue
                     if (
                         isinstance(event, TaskStatusUpdated)
@@ -167,6 +192,7 @@ class RunnerSupervisor:
                             ),
                         )
                         self.completed.add(event.task_id)
+                        self.sent.discard(event.task_id)
                     await self._event_sender.send(event)
             except (ClosedResourceError, BrokenResourceError) as e:
                 await self._check_runner(e)

src/exo/worker/tests/unittests/conftest.py

@@ -20,6 +20,7 @@ class FakeRunnerSupervisor:
     bound_instance: BoundInstance
     status: RunnerStatus
     completed: set[TaskId] = field(default_factory=set)
+    sent: set[TaskId] = field(default_factory=set)
 
 
 class OtherTask(BaseTask):

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

@@ -0,0 +1,545 @@
+# type: ignore
+import time
+from typing import cast
+from unittest.mock import patch
+
+import mlx.core as mx
+import pytest
+from mlx_lm.models.cache import KVCache
+from mlx_lm.sample_utils import make_sampler
+
+from exo.shared.types.api import ChatCompletionMessage
+from exo.shared.types.common import ModelId
+from exo.shared.types.tasks import ChatCompletionTaskParams
+from exo.worker.engines.mlx import Model
+from exo.worker.engines.mlx.cache import (
+    KVPrefixCache,
+    _cache_length,
+    _get_prefix_length,
+    encode_prompt,
+    make_kv_cache,
+)
+from exo.worker.engines.mlx.generator.generate import mlx_generate, prefill
+from exo.worker.engines.mlx.utils_mlx import apply_chat_template
+from exo.worker.tests.unittests.test_mlx.conftest import (
+    DEFAULT_GPT_OSS_CONFIG,
+    DEFAULT_GPT_OSS_MODEL_ID,
+)
+
+
+def _check_model_exists() -> bool:
+    return DEFAULT_GPT_OSS_CONFIG.model_path.exists()
+
+
+class TestGetPrefixLength:
+    def test_identical_arrays(self):
+        a = mx.array([1, 2, 3, 4, 5])
+        b = mx.array([1, 2, 3, 4, 5])
+        assert _get_prefix_length(a, b) == 5
+
+    def test_no_common_prefix(self):
+        a = mx.array([1, 2, 3])
+        b = mx.array([4, 5, 6])
+        assert _get_prefix_length(a, b) == 0
+
+    def test_partial_prefix(self):
+        a = mx.array([1, 2, 3, 4, 5])
+        b = mx.array([1, 2, 3, 7, 8])
+        assert _get_prefix_length(a, b) == 3
+
+    def test_prompt_longer_than_cached(self):
+        a = mx.array([1, 2, 3, 4, 5])
+        b = mx.array([1, 2, 3])
+        assert _get_prefix_length(a, b) == 3
+
+    def test_cached_longer_than_prompt(self):
+        a = mx.array([1, 2, 3])
+        b = mx.array([1, 2, 3, 4, 5])
+        assert _get_prefix_length(a, b) == 3
+
+    def test_single_token_match(self):
+        a = mx.array([1, 2, 3])
+        b = mx.array([1, 5, 6])
+        assert _get_prefix_length(a, b) == 1
+
+    def test_empty_prompt(self):
+        a = mx.array([]).astype(mx.int32)
+        b = mx.array([1, 2, 3])
+        assert _get_prefix_length(a, b) == 0
+
+    def test_empty_cached(self):
+        a = mx.array([1, 2, 3])
+        b = mx.array([]).astype(mx.int32)
+        assert _get_prefix_length(a, b) == 0
+
+    def test_both_empty(self):
+        a = mx.array([]).astype(mx.int32)
+        b = mx.array([]).astype(mx.int32)
+        assert _get_prefix_length(a, b) == 0
+
+
+class TestKVPrefix:
+    @pytest.fixture
+    def mock_tokenizer(self):
+        """Create a minimal mock tokenizer for tests that don't need real tokenization."""
+        from unittest.mock import MagicMock
+
+        tokenizer = MagicMock()
+        tokenizer.encode.return_value = [1, 2, 3]
+        return tokenizer
+
+    def test_starts_empty(self, mock_tokenizer):
+        cache = KVPrefixCache(mock_tokenizer)
+        assert len(cache.prompts) == 0
+        assert len(cache.caches) == 0
+
+    def test_clear_empties_cache(self, mock_tokenizer):
+        cache = KVPrefixCache(mock_tokenizer)
+        cache.prompts.append(mx.array([1, 2, 3]))
+        cache.caches.append([KVCache()])
+        cache.clear()
+        assert len(cache.prompts) == 0
+        assert len(cache.caches) == 0
+
+    def test_clear_on_empty_cache(self, mock_tokenizer):
+        cache = KVPrefixCache(mock_tokenizer)
+        cache.clear()
+        assert len(cache.prompts) == 0
+
+
+def _load_gpt_oss() -> tuple[Model, object]:
+    from mlx_lm.utils import load_model
+
+    from exo.worker.engines.mlx.utils_mlx import load_tokenizer_for_model_id
+
+    model_path = DEFAULT_GPT_OSS_CONFIG.model_path
+    model_id = ModelId(DEFAULT_GPT_OSS_MODEL_ID)
+
+    model, _ = load_model(model_path, lazy=False)
+    tokenizer = load_tokenizer_for_model_id(model_id, model_path)
+    return cast(Model, model), tokenizer
+
+
+@pytest.mark.slow
+@pytest.mark.skipif(
+    not _check_model_exists(),
+    reason=f"GPT-OSS model not found at {DEFAULT_GPT_OSS_CONFIG.model_path}",
+)
+class TestKVPrefixCacheWithModel:
+    @pytest.fixture(scope="class")
+    def model_and_tokenizer(self):
+        model, tokenizer = _load_gpt_oss()
+        return model, tokenizer
+
+    def test_prefill_populates_cache(self, model_and_tokenizer):
+        model, tokenizer = model_and_tokenizer
+
+        task = ChatCompletionTaskParams(
+            model=DEFAULT_GPT_OSS_MODEL_ID,
+            messages=[ChatCompletionMessage(role="user", content="Hello!!")],
+            max_tokens=1,
+        )
+        prompt = apply_chat_template(tokenizer, task)
+        tokens = encode_prompt(tokenizer, prompt)
+        cache = make_kv_cache(model)
+
+        prefill(model, tokenizer, make_sampler(0.0), tokens, cache)
+
+        # Cache should now hold the prompt tokens
+        assert _cache_length(cache) == len(tokens)
+
+    def test_add_and_get_exact_match(self, model_and_tokenizer):
+        model, tokenizer = model_and_tokenizer
+
+        task = ChatCompletionTaskParams(
+            model=DEFAULT_GPT_OSS_MODEL_ID,
+            messages=[ChatCompletionMessage(role="user", content="Test exact")],
+            max_tokens=1,
+        )
+        prompt = apply_chat_template(tokenizer, task)
+        tokens = encode_prompt(tokenizer, prompt)
+        cache = make_kv_cache(model)
+
+        prefill(model, tokenizer, make_sampler(0.0), tokens, cache)
+
+        kv_prefix_cache = KVPrefixCache(tokenizer)
+        kv_prefix_cache.add_kv_cache(prompt, cache)
+
+        assert len(kv_prefix_cache.prompts) == 1
+        stored_length = _cache_length(kv_prefix_cache.caches[0])
+        assert stored_length > 0
+
+        # Retrieve with same prompt: exact match
+        result_cache, remaining_tokens, matched_index = kv_prefix_cache.get_kv_cache(
+            model, prompt
+        )
+        assert matched_index == 0
+
+        # Exact match returns only last token
+        assert len(remaining_tokens) == 1
+        assert mx.array_equal(remaining_tokens, tokens[-1:])
+
+    def test_add_and_get_prefix_match(self, model_and_tokenizer):
+        """get_kv_cache with a longer prompt sharing prefix should return partial match."""
+        model, tokenizer = model_and_tokenizer
+
+        short_task = ChatCompletionTaskParams(
+            model=DEFAULT_GPT_OSS_MODEL_ID,
+            messages=[ChatCompletionMessage(role="user", content="Hi")],
+            max_tokens=1,
+        )
+        short_prompt = apply_chat_template(tokenizer, short_task)
+        short_tokens = encode_prompt(tokenizer, short_prompt)
+        cache = make_kv_cache(model)
+
+        prefill(model, tokenizer, make_sampler(0.0), short_tokens, cache)
+
+        kv_prefix_cache = KVPrefixCache(tokenizer)
+        kv_prefix_cache.add_kv_cache(short_prompt, cache)
+
+        # Query with longer prompt that shares the chat template prefix
+        long_task = ChatCompletionTaskParams(
+            model=DEFAULT_GPT_OSS_MODEL_ID,
+            messages=[
+                ChatCompletionMessage(role="user", content="Hi there, how are you?")
+            ],
+            max_tokens=1,
+        )
+        long_prompt = apply_chat_template(tokenizer, long_task)
+        long_tokens = encode_prompt(tokenizer, long_prompt)
+
+        # The prompts share a prefix (chat template preamble + "Hi")
+        expected_prefix = _get_prefix_length(long_tokens, short_tokens)
+        assert expected_prefix > 0, (
+            "Prompts should share a prefix from the chat template"
+        )
+
+        result_cache, remaining_tokens, matched_index = kv_prefix_cache.get_kv_cache(
+            model, long_prompt
+        )
+        assert matched_index == 0
+
+        # remaining_tokens should be the suffix after the shared prefix
+        assert len(remaining_tokens) == len(long_tokens) - expected_prefix
+        assert mx.array_equal(remaining_tokens, long_tokens[expected_prefix:])
+
+    def test_stored_cache_not_mutated_after_get_and_generation(
+        self, model_and_tokenizer
+    ):
+        """Getting a cache and then mutating it (as generation does) must not corrupt stored cache."""
+        model, tokenizer = model_and_tokenizer
+
+        task = ChatCompletionTaskParams(
+            model=DEFAULT_GPT_OSS_MODEL_ID,
+            messages=[ChatCompletionMessage(role="user", content="Mutation test")],
+            max_tokens=1,
+        )
+        prompt = apply_chat_template(tokenizer, task)
+        tokens = encode_prompt(tokenizer, prompt)
+        cache = make_kv_cache(model)
+
+        prefill(model, tokenizer, make_sampler(0.0), tokens, cache)
+
+        kv_prefix_cache = KVPrefixCache(tokenizer)
+        kv_prefix_cache.add_kv_cache(prompt, cache)
+
+        stored_length = _cache_length(kv_prefix_cache.caches[0])
+
+        # Get cache and mutate it (simulating what generation does)
+        result_cache, _, matched_index = kv_prefix_cache.get_kv_cache(model, prompt)
+        assert matched_index == 0
+
+        # Simulate generation: feed many additional tokens through the cache
+        head_dim = result_cache[0].keys.shape[-1]
+        num_heads = result_cache[0].keys.shape[1]
+        extra_keys = mx.random.normal((1, num_heads, 50, head_dim))
+        extra_values = mx.random.normal((1, num_heads, 50, head_dim))
+        for layer_cache in result_cache:
+            layer_cache.update_and_fetch(extra_keys, extra_values)
+        mx.eval([c.keys for c in result_cache])
+
+        # Stored cache must be unchanged
+        assert _cache_length(kv_prefix_cache.caches[0]) == stored_length
+
+    def test_stored_cache_survives_repeated_get_mutate_cycles(
+        self, model_and_tokenizer
+    ):
+        """Multiple get+mutate cycles (like repeated user requests) must not corrupt cache."""
+        model, tokenizer = model_and_tokenizer
+
+        task = ChatCompletionTaskParams(
+            model=DEFAULT_GPT_OSS_MODEL_ID,
+            messages=[ChatCompletionMessage(role="user", content="Repeat test")],
+            max_tokens=1,
+        )
+        prompt = apply_chat_template(tokenizer, task)
+        tokens = encode_prompt(tokenizer, prompt)
+        cache = make_kv_cache(model)
+
+        prefill(model, tokenizer, make_sampler(0.0), tokens, cache)
+
+        kv_prefix_cache = KVPrefixCache(tokenizer)
+        kv_prefix_cache.add_kv_cache(prompt, cache)
+
+        stored_length = _cache_length(kv_prefix_cache.caches[0])
+
+        for i in range(3):
+            result_cache, _, _ = kv_prefix_cache.get_kv_cache(model, prompt)
+
+            head_dim = result_cache[0].keys.shape[-1]
+            num_heads = result_cache[0].keys.shape[1]
+            extra = mx.random.normal((1, num_heads, 30, head_dim))
+            for layer_cache in result_cache:
+                layer_cache.update_and_fetch(extra, extra)
+            mx.eval([c.keys for c in result_cache])
+
+            assert _cache_length(kv_prefix_cache.caches[0]) == stored_length, (
+                f"Failed on loop {i}"
+            )
+
+    def test_mlx_generate_populates_cache(self, model_and_tokenizer):
+        """mlx_generate should save the cache after generation completes."""
+        model, tokenizer = model_and_tokenizer
+
+        kv_prefix_cache = KVPrefixCache(tokenizer)
+        task = ChatCompletionTaskParams(
+            model=DEFAULT_GPT_OSS_MODEL_ID,
+            messages=[ChatCompletionMessage(role="user", content="Hello")],
+            max_tokens=5,
+        )
+        prompt = apply_chat_template(tokenizer, task)
+        prompt_tokens = encode_prompt(tokenizer, prompt)
+
+        # Consume the entire generator so the cache-saving code after yield runs
+        generated_tokens = 0
+        for _response in mlx_generate(
+            model=model,
+            tokenizer=tokenizer,
+            task=task,
+            prompt=prompt,
+            kv_prefix_cache=kv_prefix_cache,
+        ):
+            generated_tokens += 1
+
+        assert len(kv_prefix_cache.prompts) == 1
+        assert len(kv_prefix_cache.caches) == 1
+        # Cache should contain prompt + generated tokens
+        expected_length = len(prompt_tokens) + generated_tokens
+        assert _cache_length(kv_prefix_cache.caches[0]) == expected_length
+
+    def test_mlx_generate_second_call_gets_prefix_hit(self, model_and_tokenizer):
+        """Second mlx_generate call with same prompt should get a prefix hit from stored cache."""
+        model, tokenizer = model_and_tokenizer
+
+        kv_prefix_cache = KVPrefixCache(tokenizer)
+        task = ChatCompletionTaskParams(
+            model=DEFAULT_GPT_OSS_MODEL_ID,
+            messages=[ChatCompletionMessage(role="user", content="Reuse test")],
+            max_tokens=5,
+        )
+        prompt = apply_chat_template(tokenizer, task)
+        prompt_tokens = encode_prompt(tokenizer, prompt)
+
+        # First generation populates cache
+        for _response in mlx_generate(
+            model=model,
+            tokenizer=tokenizer,
+            task=task,
+            prompt=prompt,
+            kv_prefix_cache=kv_prefix_cache,
+        ):
+            pass
+
+        assert len(kv_prefix_cache.prompts) == 1
+
+        # Second call should find a prefix match (the stored cache contains
+        # prompt + generated tokens, which shares the prompt prefix)
+        result_cache, remaining_tokens, matched_index = kv_prefix_cache.get_kv_cache(
+            model, prompt
+        )
+        # The stored cache is longer than the prompt (it includes generated tokens),
+        # so this is a prefix match where our prompt is fully contained
+        assert matched_index == 0
+        # Exact match: remaining_tokens is just the last token
+        assert len(remaining_tokens) == 1
+        assert mx.array_equal(remaining_tokens, prompt_tokens[-1:])
+
+    def test_mlx_generate_long_prompt_updates_cache_in_place(self, model_and_tokenizer):
+        """With a prompt > 1000 tokens, second generation should update the cache entry in-place."""
+        model, tokenizer = model_and_tokenizer
+
+        kv_prefix_cache = KVPrefixCache(tokenizer)
+
+        # Build a long user message (> 1000 tokens) to exceed _MIN_PREFIX_HIT_TO_UPDATE
+        base_text = "The quick brown fox jumps over the lazy dog. "
+        base_tokens = tokenizer.encode(base_text)
+        repeats = (1200 // len(base_tokens)) + 2
+        long_content = base_text * repeats
+
+        task1 = ChatCompletionTaskParams(
+            model=DEFAULT_GPT_OSS_MODEL_ID,
+            messages=[ChatCompletionMessage(role="user", content=long_content)],
+            max_tokens=5,
+        )
+        prompt1 = apply_chat_template(tokenizer, task1)
+        prompt1_tokens = encode_prompt(tokenizer, prompt1)
+        assert len(prompt1_tokens) > 1000, (
+            "Prompt must exceed _MIN_PREFIX_HIT_TO_UPDATE"
+        )
+
+        # First generation populates the cache (must prefill all tokens)
+        t0 = time.perf_counter()
+        for _response in mlx_generate(
+            model=model,
+            tokenizer=tokenizer,
+            task=task1,
+            prompt=prompt1,
+            kv_prefix_cache=kv_prefix_cache,
+        ):
+            pass
+        first_gen_time = time.perf_counter() - t0
+
+        assert len(kv_prefix_cache.prompts) == 1
+        first_cache_length = _cache_length(kv_prefix_cache.caches[0])
+
+        # Second generation: same long prompt + extra content (simulating multi-turn)
+        task2 = ChatCompletionTaskParams(
+            model=DEFAULT_GPT_OSS_MODEL_ID,
+            messages=[
+                ChatCompletionMessage(role="user", content=long_content),
+                ChatCompletionMessage(role="assistant", content="Sure, I can help."),
+                ChatCompletionMessage(role="user", content="Tell me more."),
+            ],
+            max_tokens=5,
+        )
+        prompt2 = apply_chat_template(tokenizer, task2)
+        prompt2_tokens = encode_prompt(tokenizer, prompt2)
+
+        # Verify the prompts share a long prefix
+        prefix_len = _get_prefix_length(prompt2_tokens, prompt1_tokens)
+        assert prefix_len > 1000, "Prompts must share > 1000 token prefix"
+
+        # Second generation should reuse the cached prefix (only prefill new tokens)
+        t0 = time.perf_counter()
+        for _response in mlx_generate(
+            model=model,
+            tokenizer=tokenizer,
+            task=task2,
+            prompt=prompt2,
+            kv_prefix_cache=kv_prefix_cache,
+        ):
+            pass
+        second_gen_time = time.perf_counter() - t0
+
+        # Second generation should be significantly faster due to prefix cache hit - hopefully not flaky
+        assert second_gen_time < first_gen_time * 0.5, (
+            f"Expected prefix cache speedup: "
+            f"first={first_gen_time:.2f}s, second={second_gen_time:.2f}s"
+        )
+
+        # With prefix_hit > 1000, should update in-place (not add a second entry)
+        assert len(kv_prefix_cache.prompts) == 1
+        # Updated cache should be longer (prompt2 + generated > prompt1 + generated)
+        updated_cache_length = _cache_length(kv_prefix_cache.caches[0])
+        assert updated_cache_length > first_cache_length
+
+    def test_mlx_generate_stored_cache_not_mutated(self, model_and_tokenizer):
+        """After mlx_generate saves a cache, a second generation must not corrupt the stored copy."""
+        model, tokenizer = model_and_tokenizer
+
+        kv_prefix_cache = KVPrefixCache(tokenizer)
+        task = ChatCompletionTaskParams(
+            model=DEFAULT_GPT_OSS_MODEL_ID,
+            messages=[ChatCompletionMessage(role="user", content="Immutable test")],
+            max_tokens=5,
+        )
+        prompt = apply_chat_template(tokenizer, task)
+
+        # First generation populates cache
+        for _response in mlx_generate(
+            model=model,
+            tokenizer=tokenizer,
+            task=task,
+            prompt=prompt,
+            kv_prefix_cache=kv_prefix_cache,
+        ):
+            pass
+
+        first_cache_length = _cache_length(kv_prefix_cache.caches[0])
+
+        # Second generation gets the cache and mutates it during generation
+        for _response in mlx_generate(
+            model=model,
+            tokenizer=tokenizer,
+            task=task,
+            prompt=prompt,
+            kv_prefix_cache=kv_prefix_cache,
+        ):
+            pass
+
+        # The first stored cache must not have been mutated by the second generation
+        assert _cache_length(kv_prefix_cache.caches[0]) == first_cache_length
+
+    def test_evicts_lru_entry_under_memory_pressure(self, model_and_tokenizer):
+        """Under memory pressure, adding a new cache entry evicts the least recently used one."""
+        model, tokenizer = model_and_tokenizer
+
+        kv_prefix_cache = KVPrefixCache(tokenizer)
+
+        # Add three cache entries with different prompts
+        prompts = ["First entry", "Second entry", "Third entry"]
+        for i, content in enumerate(prompts):
+            task = ChatCompletionTaskParams(
+                model=DEFAULT_GPT_OSS_MODEL_ID,
+                messages=[ChatCompletionMessage(role="user", content=content)],
+                max_tokens=1,
+            )
+            prompt = apply_chat_template(tokenizer, task)
+            tokens = encode_prompt(tokenizer, prompt)
+            cache = make_kv_cache(model)
+            prefill(model, tokenizer, make_sampler(0.0), tokens, cache)
+            kv_prefix_cache.add_kv_cache(prompt, cache)
+            # Stagger _last_used so LRU order is deterministic
+            kv_prefix_cache._last_used[i] = float(i)
+
+        assert len(kv_prefix_cache.prompts) == 3
+
+        # Access the third entry to make it most recently used
+        kv_prefix_cache._last_used[2] = 100.0
+        # Entry 0 (_last_used=0.0) is LRU, entry 1 (_last_used=1.0) is next
+
+        # Simulate memory pressure: active memory exceeds threshold
+        fake_limit = 1000
+        fake_active = int(fake_limit * 0.90)  # Above _MEMORY_THRESHOLD (0.85)
+
+        with (
+            patch(
+                "exo.worker.engines.mlx.cache.mx.metal.get_active_memory",
+                return_value=fake_active,
+            ),
+            patch(
+                "exo.worker.engines.mlx.cache.mx.metal.device_info",
+                return_value={"max_recommended_working_set_size": fake_limit},
+            ),
+        ):
+            # Trigger eviction by adding a new entry
+            task = ChatCompletionTaskParams(
+                model=DEFAULT_GPT_OSS_MODEL_ID,
+                messages=[ChatCompletionMessage(role="user", content="New entry")],
+                max_tokens=1,
+            )
+            prompt = apply_chat_template(tokenizer, task)
+            tokens = encode_prompt(tokenizer, prompt)
+            cache = make_kv_cache(model)
+            prefill(model, tokenizer, make_sampler(0.0), tokens, cache)
+            kv_prefix_cache.add_kv_cache(prompt, cache)
+
+        # LRU entries should have been evicted (entries 0, 1, 2 in order of _last_used)
+        # Since fake_active stays above threshold after each eviction (we don't change it),
+        # all old entries get evicted, leaving only the newly added one
+        assert len(kv_prefix_cache.prompts) == 1
+        # The surviving entry should be the newly added one
+        new_tokens = encode_prompt(tokenizer, prompt)
+        assert _get_prefix_length(kv_prefix_cache.prompts[0], new_tokens) == len(
+            new_tokens
+        )

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

@@ -11,12 +11,12 @@ from pathlib import Path
 
 import pytest
 
-from exo.shared.models.model_cards import MODEL_CARDS, ModelCard, ModelId
-from exo.worker.download.download_utils import (
+from exo.download.download_utils import (
     download_file_with_retry,
     ensure_models_dir,
     fetch_file_list_with_cache,
 )
+from exo.shared.models.model_cards import MODEL_CARDS, ModelCard, ModelId
 from exo.worker.engines.mlx.utils_mlx import (
     get_eos_token_ids_for_model,
     load_tokenizer_for_model_id,

src/exo/worker/tests/unittests/test_plan/test_download_and_loading.py

@@ -1,5 +1,5 @@
 import exo.worker.plan as plan_mod
-from exo.shared.types.common import ModelId, NodeId
+from exo.shared.types.common import NodeId
 from exo.shared.types.memory import Memory
 from exo.shared.types.tasks import LoadModel
 from exo.shared.types.worker.downloads import DownloadCompleted, DownloadProgress
@@ -45,13 +45,9 @@ def test_plan_requests_download_when_waiting_and_shard_not_downloaded():
     instances = {INSTANCE_1_ID: instance}
     all_runners = {RUNNER_1_ID: RunnerIdle()}
 
-    # No entry for this shard -> should trigger DownloadModel
-    download_status: dict[ModelId, DownloadProgress] = {}
-
     result = plan_mod.plan(
         node_id=NODE_A,
         runners=runners,  # type: ignore
-        download_status=download_status,
         global_download_status={NODE_A: []},
         instances=instances,
         all_runners=all_runners,
@@ -92,14 +88,6 @@ def test_plan_loads_model_when_all_shards_downloaded_and_waiting():
         RUNNER_2_ID: RunnerConnected(),
     }
 
-    # Local node has already marked its shard as downloaded (not actually used by _load_model)
-    local_download_status = {
-        MODEL_A_ID: DownloadCompleted(
-            shard_metadata=shard1, node_id=NODE_A, total_bytes=Memory()
-        )
-    }
-
-    # Global view has completed downloads for both nodes
     global_download_status = {
         NODE_A: [
             DownloadCompleted(
@@ -116,7 +104,6 @@ def test_plan_loads_model_when_all_shards_downloaded_and_waiting():
     result = plan_mod.plan(
         node_id=NODE_A,
         runners=runners,  # type: ignore
-        download_status=local_download_status,
         global_download_status=global_download_status,
         instances=instances,
         all_runners=all_runners,
@@ -148,30 +135,26 @@ def test_plan_does_not_request_download_when_shard_already_downloaded():
     instances = {INSTANCE_1_ID: instance}
     all_runners = {RUNNER_1_ID: RunnerIdle()}
 
-    # Local status claims the shard is downloaded already
-    local_download_status = {
-        MODEL_A_ID: DownloadCompleted(
-            shard_metadata=shard, node_id=NODE_A, total_bytes=Memory()
-        )
-    }
-
-    # Global view hasn't caught up yet (no completed shards recorded for NODE_A)
+    # Global state shows shard is downloaded for NODE_A
     global_download_status: dict[NodeId, list[DownloadProgress]] = {
-        NODE_A: [],
+        NODE_A: [
+            DownloadCompleted(
+                shard_metadata=shard, node_id=NODE_A, total_bytes=Memory()
+            )
+        ],
         NODE_B: [],
     }
 
     result = plan_mod.plan(
         node_id=NODE_A,
         runners=runners,  # type: ignore
-        download_status=local_download_status,
         global_download_status=global_download_status,
         instances=instances,
         all_runners=all_runners,
         tasks={},
     )
 
-    assert result is None
+    assert not isinstance(result, plan_mod.DownloadModel)
 
 
 def test_plan_does_not_load_model_until_all_shards_downloaded_globally():
@@ -202,12 +185,6 @@ def test_plan_does_not_load_model_until_all_shards_downloaded_globally():
         RUNNER_2_ID: RunnerConnected(),
     }
 
-    # Only NODE_A's shard is recorded as downloaded globally
-    local_download_status = {
-        MODEL_A_ID: DownloadCompleted(
-            shard_metadata=shard1, node_id=NODE_A, total_bytes=Memory()
-        )
-    }
     global_download_status = {
         NODE_A: [
             DownloadCompleted(
@@ -220,7 +197,6 @@ def test_plan_does_not_load_model_until_all_shards_downloaded_globally():
     result = plan_mod.plan(
         node_id=NODE_A,
         runners=runners,  # type: ignore
-        download_status=local_download_status,
         global_download_status=global_download_status,
         instances=instances,
         all_runners=all_runners,
@@ -245,7 +221,6 @@ def test_plan_does_not_load_model_until_all_shards_downloaded_globally():
     result = plan_mod.plan(
         node_id=NODE_A,
         runners=runners,  # type: ignore
-        download_status=local_download_status,
         global_download_status=global_download_status,
         instances=instances,
         all_runners=all_runners,