diff --git a/.github/actions/typecheck/action.yml b/.github/actions/typecheck/action.yml deleted file mode 100644 index cd52d6e3..00000000 --- a/.github/actions/typecheck/action.yml +++ /dev/null @@ -1,12 +0,0 @@ -name: Type Check - -description: "Run type checker" - -runs: - using: "composite" - steps: - - name: Run type checker - run: | - nix --extra-experimental-features nix-command --extra-experimental-features flakes develop -c just sync - nix --extra-experimental-features nix-command --extra-experimental-features flakes develop -c just check - shell: bash diff --git a/.github/workflows/pipeline.yml b/.github/workflows/pipeline.yml index 7c5768ac..c2589453 100644 --- a/.github/workflows/pipeline.yml +++ b/.github/workflows/pipeline.yml @@ -26,73 +26,14 @@ jobs: name: exo authToken: "${{ secrets.CACHIX_AUTH_TOKEN }}" - - name: Configure git user - run: | - git config --local user.email "github-actions@users.noreply.github.com" - git config --local user.name "github-actions bot" - shell: bash + - name: Load nix develop environment + run: nix run github:nicknovitski/nix-develop/v1 - - name: Pull LFS files - run: | - echo "Pulling Git LFS files..." - git lfs pull - shell: bash + - name: Sync dependencies + run: uv sync --all-packages - - name: Setup Nix Environment - run: | - echo "Checking for nix installation..." - - # Check if nix binary exists directly - if [ -f /nix/var/nix/profiles/default/bin/nix ]; then - echo "Found nix binary at /nix/var/nix/profiles/default/bin/nix" - export PATH="/nix/var/nix/profiles/default/bin:$PATH" - echo "PATH=$PATH" >> $GITHUB_ENV - nix --version - elif [ -f /nix/var/nix/profiles/default/etc/profile.d/nix-daemon.sh ]; then - echo "Found nix profile script, sourcing..." - source /nix/var/nix/profiles/default/etc/profile.d/nix-daemon.sh - nix --version - elif command -v nix >/dev/null 2>&1; then - echo "Nix already in PATH" - nix --version - else - echo "Nix not found. Debugging info:" - echo "Contents of /nix/var/nix/profiles/default/:" - ls -la /nix/var/nix/profiles/default/ 2>/dev/null || echo "Directory not found" - echo "Contents of /nix/var/nix/profiles/default/bin/:" - ls -la /nix/var/nix/profiles/default/bin/ 2>/dev/null || echo "Directory not found" - exit 1 - fi - shell: bash - - - name: Configure basedpyright include for local MLX - run: | - RUNNER_LABELS='${{ toJSON(runner.labels) }}' - if echo "$RUNNER_LABELS" | grep -q "local_mlx"; then - if [ -d "/Users/Shared/mlx" ]; then - echo "Updating [tool.basedpyright].include to use /Users/Shared/mlx" - awk ' - BEGIN { in=0 } - /^\[tool\.basedpyright\]/ { in=1; print; next } - in && /^\[/ { in=0 } # next section - in && /^[ \t]*include[ \t]*=/ { - print "include = [\"/Users/Shared/mlx\"]" - next - } - { print } - ' pyproject.toml > pyproject.toml.tmp && mv pyproject.toml.tmp pyproject.toml - - echo "New [tool.basedpyright] section:" - sed -n '/^\[tool\.basedpyright\]/,/^\[/p' pyproject.toml | sed '$d' || true - else - echo "local_mlx tag present but /Users/Shared/mlx not found; leaving pyproject unchanged." - fi - else - echo "Runner does not have 'local_mlx' tag; leaving pyproject unchanged." - fi - shell: bash - - - uses: ./.github/actions/typecheck + - name: Run type checker + run: uv run basedpyright --project pyproject.toml nix: name: Build and check (${{ matrix.system }}) @@ -123,6 +64,63 @@ jobs: name: exo authToken: "${{ secrets.CACHIX_AUTH_TOKEN }}" + - name: Build Metal packages (macOS only) + if: runner.os == 'macOS' + run: | + # Try to build metal-toolchain first (may succeed via cachix cache hit) + if nix build .#metal-toolchain 2>/dev/null; then + echo "metal-toolchain built successfully (likely cache hit)" + else + echo "metal-toolchain build failed, extracting from Xcode..." + + NAR_HASH="sha256-ayR5mXN4sZAddwKEG2OszGRF93k9ZFc7H0yi2xbylQw=" + NAR_NAME="metal-toolchain-17C48.nar" + + # Use RUNNER_TEMP to avoid /tmp symlink issues on macOS + WORK_DIR="${RUNNER_TEMP}/metal-work" + mkdir -p "$WORK_DIR" + + # Download the Metal toolchain component + xcodebuild -downloadComponent MetalToolchain + + # Find and mount the DMG + DMG_PATH=$(find /System/Library/AssetsV2/com_apple_MobileAsset_MetalToolchain -name '*.dmg' 2>/dev/null | head -1) + if [ -z "$DMG_PATH" ]; then + echo "Error: Could not find Metal toolchain DMG" + exit 1 + fi + + echo "Found DMG at: $DMG_PATH" + hdiutil attach "$DMG_PATH" -mountpoint "${WORK_DIR}/metal-dmg" + + # Copy the toolchain + cp -R "${WORK_DIR}/metal-dmg/Metal.xctoolchain" "${WORK_DIR}/metal-export" + hdiutil detach "${WORK_DIR}/metal-dmg" + + # Create NAR and add to store + nix nar pack "${WORK_DIR}/metal-export" > "${WORK_DIR}/${NAR_NAME}" + STORE_PATH=$(nix store add --mode flat "${WORK_DIR}/${NAR_NAME}") + echo "Added NAR to store: $STORE_PATH" + + # Verify the hash matches + ACTUAL_HASH=$(nix hash file "${WORK_DIR}/${NAR_NAME}") + if [ "$ACTUAL_HASH" != "$NAR_HASH" ]; then + echo "Warning: NAR hash mismatch!" + echo "Expected: $NAR_HASH" + echo "Actual: $ACTUAL_HASH" + echo "The metal-toolchain.nix may need updating" + fi + + # Clean up + rm -rf "$WORK_DIR" + + # Retry the build now that NAR is in store + nix build .#metal-toolchain + fi + + # Build mlx (depends on metal-toolchain) + nix build .#mlx + - name: Build all Nix outputs run: | nix flake show --json | jq -r ' @@ -134,3 +132,16 @@ jobs: - name: Run nix flake check run: nix flake check + + - name: Run pytest (macOS only) + if: runner.os == 'macOS' + run: | + # Build the test environment (requires relaxed sandbox for uv2nix on macOS) + TEST_ENV=$(nix build '.#exo-test-env' --option sandbox relaxed --print-out-paths) + + # Run pytest outside sandbox (needs GPU access for MLX) + export HOME="$RUNNER_TEMP" + export EXO_TESTS=1 + export EXO_DASHBOARD_DIR="$PWD/dashboard/" + export EXO_RESOURCES_DIR="$PWD/resources" + $TEST_ENV/bin/python -m pytest src -m "not slow" --import-mode=importlib diff --git a/.gitignore b/.gitignore index a4231085..139fc326 100644 --- a/.gitignore +++ b/.gitignore @@ -28,3 +28,7 @@ target/ dashboard/build/ dashboard/node_modules/ dashboard/.svelte-kit/ + +# host config snapshots +hosts_*.json +.swp diff --git a/.mlx_typings/mlx_lm/tokenizer_utils.pyi b/.mlx_typings/mlx_lm/tokenizer_utils.pyi index 251e3d28..83eb4e33 100644 --- a/.mlx_typings/mlx_lm/tokenizer_utils.pyi +++ b/.mlx_typings/mlx_lm/tokenizer_utils.pyi @@ -108,6 +108,7 @@ class TokenizerWrapper: _tokenizer: PreTrainedTokenizerFast eos_token_id: int | None eos_token: str | None + eos_token_ids: list[int] | set[int] | None bos_token_id: int | None bos_token: str | None vocab_size: int @@ -117,7 +118,7 @@ class TokenizerWrapper: self, tokenizer: Any, detokenizer_class: Any = ..., - eos_token_ids: list[int] | None = ..., + eos_token_ids: list[int] | set[int] | None = ..., chat_template: Any = ..., tool_parser: Any = ..., tool_call_start: str | None = ..., diff --git a/README.md b/README.md index b2ab43b2..58d41c37 100644 --- a/README.md +++ b/README.md @@ -5,7 +5,7 @@ exo logo -exo: Run your own AI cluster at home with everyday devices. Maintained by [exo labs](https://x.com/exolabs). +exo: Run frontier AI locally. Maintained by [exo labs](https://x.com/exolabs).

Discord @@ -107,6 +107,10 @@ uv run exo This starts the exo dashboard and API at http://localhost:52415/ + +*Please view the section on RDMA to enable this feature on MacOS >=26.2!* + + ### Run from Source (Linux) **Prerequisites:** @@ -230,7 +234,7 @@ This removes: RDMA is a new capability added to macOS 26.2. It works on any Mac with Thunderbolt 5 (M4 Pro Mac Mini, M4 Max Mac Studio, M4 Max MacBook Pro, M3 Ultra Mac Studio). -Note that on Mac Studio, you cannot use the Thunderbolt 5 port next to the Ethernet port. +Please refer to the caveats for immediate troubleshooting. To enable RDMA on macOS, follow these steps: @@ -247,6 +251,14 @@ To enable RDMA on macOS, follow these steps: After that, RDMA will be enabled in macOS and exo will take care of the rest. +**Important Caveats** + +1. Devices that wish to be part of an RDMA cluster must be connected to all other devices in the cluster. +2. The cables must support TB5. +3. On a Mac Studio, you cannot use the Thunderbolt 5 port next to the Ethernet port. +4. If running from source, please use the script found at `tmp/set_rdma_network_config.sh`, which will disable Thunderbolt Bridge and set dhcp on each RDMA port. +5. RDMA ports may be unable to discover each other on different versions of MacOS. Please ensure that OS versions match exactly (even beta version numbers) on all devices. + --- ### Using the API diff --git a/app/EXO/EXO.xcodeproj/project.pbxproj b/app/EXO/EXO.xcodeproj/project.pbxproj index bd593086..4427289a 100644 --- a/app/EXO/EXO.xcodeproj/project.pbxproj +++ b/app/EXO/EXO.xcodeproj/project.pbxproj @@ -342,6 +342,8 @@ SDKROOT = macosx; SWIFT_ACTIVE_COMPILATION_CONDITIONS = "DEBUG $(inherited)"; SWIFT_OPTIMIZATION_LEVEL = "-Onone"; + SWIFT_TREAT_WARNINGS_AS_ERRORS = YES; + GCC_TREAT_WARNINGS_AS_ERRORS = YES; }; name = Debug; }; @@ -397,6 +399,8 @@ MTL_FAST_MATH = YES; SDKROOT = macosx; SWIFT_COMPILATION_MODE = wholemodule; + SWIFT_TREAT_WARNINGS_AS_ERRORS = YES; + GCC_TREAT_WARNINGS_AS_ERRORS = YES; }; name = Release; }; diff --git a/app/EXO/EXO/ContentView.swift b/app/EXO/EXO/ContentView.swift index a604fa46..d00bbf5c 100644 --- a/app/EXO/EXO/ContentView.swift +++ b/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 = [] @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 diff --git a/app/EXO/EXO/EXOApp.swift b/app/EXO/EXO/EXOApp.swift index e52edafe..3aff58c3 100644 --- a/app/EXO/EXO/EXOApp.swift +++ b/app/EXO/EXO/EXOApp.swift @@ -14,13 +14,13 @@ import SwiftUI import UserNotifications import os.log -@main struct EXOApp: App { @StateObject private var controller: ExoProcessController @StateObject private var stateService: ClusterStateService @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 +41,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 +61,7 @@ struct EXOApp: App { .environmentObject(networkStatusService) .environmentObject(localNetworkChecker) .environmentObject(updater) + .environmentObject(thunderboltBridgeService) } label: { menuBarIcon } @@ -130,6 +134,7 @@ struct EXOApp: App { "Failed to register EXO for launch at login: \(error.localizedDescription)") } } + } /// Helper for managing EXO's launch-at-login registration @@ -219,7 +224,7 @@ private final class ExoUpdaterDelegate: NSObject, SPUUpdaterDelegate { } } - private func showNotification(title: String, body: String) { + nonisolated private func showNotification(title: String, body: String) { let center = UNUserNotificationCenter.current() let content = UNMutableNotificationContent() content.title = title diff --git a/app/EXO/EXO/ExoProcessController.swift b/app/EXO/EXO/ExoProcessController.swift index 69d8e02a..2dec3868 100644 --- a/app/EXO/EXO/ExoProcessController.swift +++ b/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 { diff --git a/app/EXO/EXO/Models/ClusterState.swift b/app/EXO/EXO/Models/ClusterState.swift index b82bf7ea..ee4c3816 100644 --- a/app/EXO/EXO/Models/ClusterState.swift +++ b/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 } @@ -201,7 +293,7 @@ struct ClusterTask { let modelName: String? let promptPreview: String? let errorMessage: String? - let parameters: ChatCompletionTaskParameters? + let parameters: TextGenerationTaskParameters? var sortPriority: Int { switch status { @@ -238,12 +330,12 @@ struct ClusterTaskPayload: Decodable { let taskStatus: TaskStatus? let instanceId: String? let commandId: String? - let taskParams: ChatCompletionTaskParameters? + let taskParams: TextGenerationTaskParameters? let errorType: String? let errorMessage: String? } -struct ChatCompletionTaskParameters: Decodable, Equatable { +struct TextGenerationTaskParameters: Decodable, Equatable { let model: String? let messages: [ChatCompletionMessage]? let maxTokens: Int? @@ -282,7 +374,7 @@ extension ClusterTask { guard let id = payload.taskId else { return nil } let status = payload.taskStatus ?? .unknown switch kindKey { - case "ChatCompletion": + case "TextGeneration": self.init( id: id, status: status, diff --git a/app/EXO/EXO/Services/BugReportService.swift b/app/EXO/EXO/Services/BugReportService.swift index 70180855..f6d842e5 100644 --- a/app/EXO/EXO/Services/BugReportService.swift +++ b/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"]) diff --git a/app/EXO/EXO/Services/LocalNetworkChecker.swift b/app/EXO/EXO/Services/LocalNetworkChecker.swift index 9129030b..c6ed4a97 100644 --- a/app/EXO/EXO/Services/LocalNetworkChecker.swift +++ b/app/EXO/EXO/Services/LocalNetworkChecker.swift @@ -41,6 +41,7 @@ final class LocalNetworkChecker: ObservableObject { private var connection: NWConnection? private var checkTask: Task? + private var periodicTask: Task? /// 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() diff --git a/app/EXO/EXO/Services/NetworkSetupHelper.swift b/app/EXO/EXO/Services/NetworkSetupHelper.swift index 28931ead..5428ee8e 100644 --- a/app/EXO/EXO/Services/NetworkSetupHelper.swift +++ b/app/EXO/EXO/Services/NetworkSetupHelper.swift @@ -7,14 +7,20 @@ 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 set -euo pipefail + # Wait for macOS to finish network setup after boot + sleep 20 + PREFS="/Library/Preferences/SystemConfiguration/preferences.plist" # Remove bridge0 interface @@ -28,19 +34,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 configure local networking. This will disable Thunderbolt Bridge (preventing packet storms) and install a Network Location.\n\nYou will be prompted for your 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 +119,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 +157,7 @@ enum NetworkSetupHelper { return true } - private static func installLaunchDaemon() async throws { + private static func installLaunchDaemon() throws { let installerScript = makeInstallerScript() try runShellAsAdmin(installerScript) } @@ -151,8 +168,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 +212,137 @@ enum NetworkSetupHelper { 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 >/dev/null 2>&1 || true + + # Delete the exo network location if it exists + networksetup -listlocations 2>/dev/null | grep -q '^exo$' && { + networksetup -deletelocation exo >/dev/null 2>&1 || 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) } + ') || true + [ -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}') || true + [ -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 } + ') || true + if [ -n "$service_name" ]; then + networksetup -setnetworkserviceenabled "$service_name" on 2>/dev/null || true + return 0 + fi + fi + done + done + return 0 + } + find_and_enable_thunderbolt_bridge || true + + echo "EXO network components removed successfully" + """ + } + + /// Direct install without GUI (requires root). + /// Returns true on success, false on failure. + static func installDirectly() -> Bool { + let script = makeInstallerScript() + return runShellDirectly(script) + } + + /// Direct uninstall without GUI (requires root). + /// Returns true on success, false on failure. + static func uninstallDirectly() -> Bool { + let script = makeUninstallScript() + return runShellDirectly(script) + } + + /// Run a shell script directly via Process (no AppleScript, requires root). + /// Returns true on success, false on failure. + private static func runShellDirectly(_ script: String) -> Bool { + let process = Process() + process.executableURL = URL(fileURLWithPath: "/bin/bash") + process.arguments = ["-c", script] + + let outputPipe = Pipe() + let errorPipe = Pipe() + process.standardOutput = outputPipe + process.standardError = errorPipe + + do { + try process.run() + process.waitUntilExit() + + let outputData = outputPipe.fileHandleForReading.readDataToEndOfFile() + let errorData = errorPipe.fileHandleForReading.readDataToEndOfFile() + + if let output = String(data: outputData, encoding: .utf8), !output.isEmpty { + print(output) + } + if let errorOutput = String(data: errorData, encoding: .utf8), !errorOutput.isEmpty { + fputs(errorOutput, stderr) + } + + if process.terminationStatus == 0 { + logger.info("Shell script completed successfully") + return true + } else { + logger.error("Shell script failed with exit code \(process.terminationStatus)") + return false + } + } catch { + logger.error( + "Failed to run shell script: \(error.localizedDescription, privacy: .public)") + fputs("Error: \(error.localizedDescription)\n", stderr) + return false + } + } + private static func runShellAsAdmin(_ script: String) throws { let escapedScript = script diff --git a/app/EXO/EXO/Services/NetworkStatusService.swift b/app/EXO/EXO/Services/NetworkStatusService.swift index dea77a1f..5d80ea39 100644 --- a/app/EXO/EXO/Services/NetworkStatusService.swift +++ b/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? { diff --git a/app/EXO/EXO/Services/ThunderboltBridgeDetector.swift b/app/EXO/EXO/Services/ThunderboltBridgeDetector.swift new file mode 100644 index 00000000..ae03b9a7 --- /dev/null +++ b/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? { + let result = runCommand(["networksetup", "-listallhardwareports"]) + guard result.exitCode == 0 else { + logger.warning("networksetup -listallhardwareports failed: \(result.error)") + return nil + } + + var thunderboltDevices: Set = [] + 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[.. Set { + let result = runCommand(["ifconfig", bridgeDevice]) + guard result.exitCode == 0 else { + logger.debug("ifconfig \(bridgeDevice) failed") + return [] + } + + var members: Set = [] + 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) + ) + } +} diff --git a/app/EXO/EXO/Services/ThunderboltBridgeService.swift b/app/EXO/EXO/Services/ThunderboltBridgeService.swift new file mode 100644 index 00000000..112fb991 --- /dev/null +++ b/app/EXO/EXO/Services/ThunderboltBridgeService.swift @@ -0,0 +1,261 @@ +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() + 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" + status = rightName.withCString { nameCString in + var item = AuthorizationItem( + name: nameCString, + valueLength: 0, + value: nil, + flags: 0 + ) + return withUnsafeMutablePointer(to: &item) { itemPointer in + var rights = AuthorizationRights(count: 1, items: itemPointer) + return 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" + } + } +} diff --git a/app/EXO/EXO/ViewModels/InstanceViewModel.swift b/app/EXO/EXO/ViewModels/InstanceViewModel.swift index 17b694ec..6bcac3f6 100644 --- a/app/EXO/EXO/ViewModels/InstanceViewModel.swift +++ b/app/EXO/EXO/ViewModels/InstanceViewModel.swift @@ -216,7 +216,7 @@ struct InstanceTaskViewModel: Identifiable, Equatable { let promptPreview: String? let errorMessage: String? let subtitle: String? - let parameters: ChatCompletionTaskParameters? + let parameters: TextGenerationTaskParameters? var title: String { switch kind { diff --git a/app/EXO/EXO/ViewModels/NodeViewModel.swift b/app/EXO/EXO/ViewModels/NodeViewModel.swift index e52670f0..6713491c 100644 --- a/app/EXO/EXO/ViewModels/NodeViewModel.swift +++ b/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 } diff --git a/app/EXO/EXO/main.swift b/app/EXO/EXO/main.swift new file mode 100644 index 00000000..9383981f --- /dev/null +++ b/app/EXO/EXO/main.swift @@ -0,0 +1,85 @@ +// +// main.swift +// EXO +// +// Created by Jake Hillion on 2026-02-03. +// + +import Foundation + +/// Command line options for the EXO app +enum CLICommand { + case install + case uninstall + case help + case none +} + +/// Parse command line arguments to determine the CLI command +func parseArguments() -> CLICommand { + let args = CommandLine.arguments + if args.contains("--help") || args.contains("-h") { + return .help + } + if args.contains("--install") { + return .install + } + if args.contains("--uninstall") { + return .uninstall + } + return .none +} + +/// Print usage information +func printUsage() { + let programName = (CommandLine.arguments.first as NSString?)?.lastPathComponent ?? "EXO" + print( + """ + Usage: \(programName) [OPTIONS] + + Options: + --install Install EXO network configuration (requires root) + --uninstall Uninstall EXO network configuration (requires root) + --help, -h Show this help message + + When run without options, starts the normal GUI application. + + Examples: + sudo \(programName) --install Install network components as root + sudo \(programName) --uninstall Remove network components as root + """) +} + +/// Check if running as root +func isRunningAsRoot() -> Bool { + return getuid() == 0 +} + +// Main entry point +let command = parseArguments() + +switch command { +case .help: + printUsage() + exit(0) + +case .install: + if !isRunningAsRoot() { + fputs("Error: --install requires root privileges. Run with sudo.\n", stderr) + exit(1) + } + let success = NetworkSetupHelper.installDirectly() + exit(success ? 0 : 1) + +case .uninstall: + if !isRunningAsRoot() { + fputs("Error: --uninstall requires root privileges. Run with sudo.\n", stderr) + exit(1) + } + let success = NetworkSetupHelper.uninstallDirectly() + exit(success ? 0 : 1) + +case .none: + // Start normal GUI application + EXOApp.main() +} diff --git a/app/EXO/uninstall-exo.sh b/app/EXO/uninstall-exo.sh index 7cc60925..c51f33a4 100755 --- a/app/EXO/uninstall-exo.sh +++ b/app/EXO/uninstall-exo.sh @@ -29,21 +29,21 @@ YELLOW='\033[1;33m' NC='\033[0m' # No Color echo_info() { - echo -e "${GREEN}[INFO]${NC} $1" + echo -e "${GREEN}[INFO]${NC} $1" } echo_warn() { - echo -e "${YELLOW}[WARN]${NC} $1" + echo -e "${YELLOW}[WARN]${NC} $1" } echo_error() { - echo -e "${RED}[ERROR]${NC} $1" + echo -e "${RED}[ERROR]${NC} $1" } # Check if running as root if [[ $EUID -ne 0 ]]; then - echo_error "This script must be run as root (use sudo)" - exit 1 + echo_error "This script must be run as root (use sudo)" + exit 1 fi echo "" @@ -55,64 +55,64 @@ echo "" # Unload the LaunchDaemon if running echo_info "Stopping network setup daemon..." if launchctl list | grep -q "$LABEL"; then - launchctl bootout system/"$LABEL" 2>/dev/null || true - echo_info "Daemon stopped" + launchctl bootout system/"$LABEL" 2>/dev/null || true + echo_info "Daemon stopped" else - echo_warn "Daemon was not running" + echo_warn "Daemon was not running" fi # Remove LaunchDaemon plist -if [[ -f "$PLIST_DEST" ]]; then - rm -f "$PLIST_DEST" - echo_info "Removed LaunchDaemon plist" +if [[ -f $PLIST_DEST ]]; then + rm -f "$PLIST_DEST" + echo_info "Removed LaunchDaemon plist" else - echo_warn "LaunchDaemon plist not found (already removed?)" + echo_warn "LaunchDaemon plist not found (already removed?)" fi # Remove the script and parent directory -if [[ -f "$SCRIPT_DEST" ]]; then - rm -f "$SCRIPT_DEST" - echo_info "Removed network setup script" +if [[ -f $SCRIPT_DEST ]]; then + rm -f "$SCRIPT_DEST" + echo_info "Removed network setup script" else - echo_warn "Network setup script not found (already removed?)" + echo_warn "Network setup script not found (already removed?)" fi # Remove EXO directory if empty if [[ -d "/Library/Application Support/EXO" ]]; then - rmdir "/Library/Application Support/EXO" 2>/dev/null && \ - echo_info "Removed EXO support directory" || \ - echo_warn "EXO support directory not empty, leaving in place" + rmdir "/Library/Application Support/EXO" 2>/dev/null && + echo_info "Removed EXO support directory" || + echo_warn "EXO support directory not empty, leaving in place" fi # Remove log files -if [[ -f "$LOG_OUT" ]] || [[ -f "$LOG_ERR" ]]; then - rm -f "$LOG_OUT" "$LOG_ERR" - echo_info "Removed log files" +if [[ -f $LOG_OUT ]] || [[ -f $LOG_ERR ]]; then + rm -f "$LOG_OUT" "$LOG_ERR" + echo_info "Removed log files" else - echo_warn "Log files not found (already removed?)" + echo_warn "Log files not found (already removed?)" fi # Switch back to Automatic network location echo_info "Restoring network configuration..." if networksetup -listlocations | grep -q "^Automatic$"; then - networksetup -switchtolocation Automatic 2>/dev/null || true - echo_info "Switched to Automatic network location" + networksetup -switchtolocation Automatic 2>/dev/null || true + echo_info "Switched to Automatic network location" else - echo_warn "Automatic network location not found" + echo_warn "Automatic network location not found" fi # Delete the exo network location if it exists if networksetup -listlocations | grep -q "^exo$"; then - networksetup -deletelocation exo 2>/dev/null || true - echo_info "Deleted 'exo' network location" + networksetup -deletelocation exo 2>/dev/null || true + echo_info "Deleted 'exo' network location" else - echo_warn "'exo' network location not found (already removed?)" + echo_warn "'exo' network location not found (already removed?)" fi # Re-enable Thunderbolt Bridge if it exists if networksetup -listnetworkservices 2>/dev/null | grep -q "Thunderbolt Bridge"; then - networksetup -setnetworkserviceenabled "Thunderbolt Bridge" on 2>/dev/null || true - echo_info "Re-enabled Thunderbolt Bridge" + networksetup -setnetworkserviceenabled "Thunderbolt Bridge" on 2>/dev/null || true + echo_info "Re-enabled Thunderbolt Bridge" fi # Note about launch at login registration @@ -124,14 +124,14 @@ echo_warn " System Settings → General → Login Items → Remove EXO" # Check if EXO.app exists in common locations APP_FOUND=false for app_path in "/Applications/EXO.app" "$HOME/Applications/EXO.app"; do - if [[ -d "$app_path" ]]; then - if [[ "$APP_FOUND" == false ]]; then - echo "" - APP_FOUND=true - fi - echo_warn "EXO.app found at: $app_path" - echo_warn "You may want to move it to Trash manually." + if [[ -d $app_path ]]; then + if [[ $APP_FOUND == false ]]; then + echo "" + APP_FOUND=true fi + echo_warn "EXO.app found at: $app_path" + echo_warn "You may want to move it to Trash manually." + fi done echo "" @@ -151,4 +151,3 @@ echo "" echo "Manual step required:" echo " Remove EXO from Login Items in System Settings → General → Login Items" echo "" - diff --git a/bench/exo_bench.py b/bench/exo_bench.py index c8bbba5c..83390b0b 100644 --- a/bench/exo_bench.py +++ b/bench/exo_bench.py @@ -5,10 +5,13 @@ from __future__ import annotations import argparse import contextlib import http.client +import itertools import json import os +import sys import time from collections.abc import Callable +from pathlib import Path from statistics import mean from typing import Any from urllib.parse import urlencode @@ -16,6 +19,84 @@ from urllib.parse import urlencode from loguru import logger from transformers import AutoTokenizer +# Monkey-patch for transformers 5.x compatibility +# Kimi's tokenization_kimi.py imports bytes_to_unicode from the old location +# which was moved in transformers 5.0.0rc2 +try: + import transformers.models.gpt2.tokenization_gpt2 as gpt2_tokenization + from transformers.convert_slow_tokenizer import bytes_to_unicode + + if not hasattr(gpt2_tokenization, "bytes_to_unicode"): + gpt2_tokenization.bytes_to_unicode = bytes_to_unicode # type: ignore[attr-defined] +except ImportError: + pass # transformers < 5.0 or bytes_to_unicode not available + + +def load_tokenizer_for_bench(model_id: str) -> Any: + """ + Load tokenizer for benchmarking, with special handling for Kimi models. + + Kimi uses a custom TikTokenTokenizer that transformers 5.x can't load via AutoTokenizer. + This function replicates the logic from utils_mlx.py for bench compatibility. + """ + model_id_lower = model_id.lower() + + if "kimi-k2" in model_id_lower: + import importlib.util + import types + + from huggingface_hub import snapshot_download + + # Download/get the model path + model_path = Path( + snapshot_download( + model_id, + allow_patterns=["*.json", "*.py", "*.tiktoken"], + ) + ) + + sys.path.insert(0, str(model_path)) + + # Load tool_declaration_ts first (tokenization_kimi imports it with relative import) + tool_decl_path = model_path / "tool_declaration_ts.py" + if tool_decl_path.exists(): + spec = importlib.util.spec_from_file_location( + "tool_declaration_ts", tool_decl_path + ) + if spec and spec.loader: + tool_decl_module = importlib.util.module_from_spec(spec) + sys.modules["tool_declaration_ts"] = tool_decl_module + spec.loader.exec_module(tool_decl_module) + + # Load tokenization_kimi with patched source (convert relative to absolute import) + tok_path = model_path / "tokenization_kimi.py" + source = tok_path.read_text() + source = source.replace("from .tool_declaration_ts", "from tool_declaration_ts") + spec = importlib.util.spec_from_file_location("tokenization_kimi", tok_path) + if spec: + tok_module = types.ModuleType("tokenization_kimi") + tok_module.__file__ = str(tok_path) + sys.modules["tokenization_kimi"] = tok_module + exec(compile(source, tok_path, "exec"), tok_module.__dict__) # noqa: S102 + TikTokenTokenizer = tok_module.TikTokenTokenizer # noqa: N806 + else: + from tokenization_kimi import TikTokenTokenizer # type: ignore[import-not-found] # noqa: I001 + + hf_tokenizer: Any = TikTokenTokenizer.from_pretrained(model_path) + + # Patch encode to use internal tiktoken model directly + # transformers 5.x has a bug in the encode->pad path for slow tokenizers + def _patched_encode(text: str, **kwargs: object) -> list[int]: + # Pass allowed_special="all" to handle special tokens like <|im_user|> + return list(hf_tokenizer.model.encode(text, allowed_special="all")) + + hf_tokenizer.encode = _patched_encode + + return hf_tokenizer + + # Default: use AutoTokenizer + return AutoTokenizer.from_pretrained(model_id, trust_remote_code=True) + class ExoHttpError(RuntimeError): def __init__(self, status: int, reason: str, body_preview: str): @@ -24,7 +105,7 @@ class ExoHttpError(RuntimeError): class ExoClient: - def __init__(self, host: str, port: int, timeout_s: float = 600.0): + def __init__(self, host: str, port: int, timeout_s: float = 7200.0): self.host = host self.port = port self.timeout_s = timeout_s @@ -180,14 +261,7 @@ def parse_int_list(values: list[str]) -> list[int]: part = part.strip() if part: items.append(int(part)) - - seen: set[int] = set() - out: list[int] = [] - for x in items: - if x not in seen: - out.append(x) - seen.add(x) - return out + return items def resolve_model_short_id(client: ExoClient, model_arg: str) -> tuple[str, str]: @@ -240,7 +314,11 @@ def run_one_completion( stats = out.get("generation_stats") - preview = (out.get("choices") or [{}])[0]["message"]["content"][:200] + # Extract preview, handling None content (common for thinking models) + choices = out.get("choices") or [{}] + message = choices[0].get("message", {}) if choices else {} + content = message.get("content") or "" + preview = content[:200] if content else "" return { "elapsed_s": elapsed, @@ -277,12 +355,29 @@ class PromptSizer: f"Target ({target}) is smaller than template overhead ({self.base_tokens})." ) - content = "" - tok = self.count_fn(content) + # Estimate tokens per atom using a sample + sample_count = 100 + sample_content = self.atom * sample_count + sample_tokens = self.count_fn(sample_content) - self.base_tokens + tokens_per_atom = sample_tokens / sample_count - while tok < target: - content += self.atom - tok = self.count_fn(content) + # Estimate starting point + needed_tokens = target - self.base_tokens + estimated_atoms = int(needed_tokens / tokens_per_atom) + + # Binary search to find exact atom count + low, high = 0, estimated_atoms * 2 + 100 + while low < high: + mid = (low + high) // 2 + tok = self.count_fn(self.atom * mid) + if tok < target: + low = mid + 1 + else: + high = mid + + content = self.atom * low + tok = self.count_fn(content) + logger.info(f"{tok=}") if tok != target: raise RuntimeError( @@ -348,7 +443,7 @@ def main() -> int: help="Warmup runs per placement (uses first pp/tg).", ) ap.add_argument( - "--timeout", type=float, default=600.0, help="HTTP timeout (seconds)." + "--timeout", type=float, default=7200.0, help="HTTP timeout (seconds)." ) ap.add_argument( "--json-out", @@ -358,6 +453,11 @@ def main() -> int: ap.add_argument( "--dry-run", action="store_true", help="List selected placements and exit." ) + ap.add_argument( + "--all-combinations", + action="store_true", + help="Force all pp×tg combinations (cartesian product) even when lists have equal length.", + ) args = ap.parse_args() pp_list = parse_int_list(args.pp) @@ -369,6 +469,15 @@ def main() -> int: logger.error("--repeat must be >= 1") return 2 + # Log pairing mode + use_combinations = args.all_combinations or len(pp_list) != len(tg_list) + if use_combinations: + logger.info( + f"pp/tg mode: combinations (product) - {len(pp_list) * len(tg_list)} pairs" + ) + else: + logger.info(f"pp/tg mode: tandem (zip) - {len(pp_list)} pairs") + client = ExoClient(args.host, args.port, timeout_s=args.timeout) short_id, full_model_id = resolve_model_short_id(client, args.model) @@ -377,10 +486,7 @@ def main() -> int: ) previews = previews_resp.get("previews") or [] - tokenizer = AutoTokenizer.from_pretrained( - full_model_id, - trust_remote_code=True, - ) + tokenizer = load_tokenizer_for_bench(full_model_id) if tokenizer is None: raise RuntimeError("[exo-bench] tokenizer load failed") @@ -486,60 +592,55 @@ def main() -> int: ) logger.debug(f" warmup {i + 1}/{args.warmup} done") - for pp in pp_list: - # if ( - # pp * n_nodes > 2048 - # and "ring" in instance_meta.lower() - # and "tensor" in sharding.lower() - # ): - # model_card = MODEL_CARDS[short_id] - # if model_card.metadata.storage_size > Memory.from_gb(10): - # logger.info( - # f"Skipping tensor ring as this is too slow for model of size {model_card.metadata.storage_size} on {n_nodes=}" - # ) - # continue - for tg in tg_list: - runs: list[dict[str, Any]] = [] - for r in range(args.repeat): - time.sleep(3) - try: - row, actual_pp_tokens = run_one_completion( - client, full_model_id, pp, tg, prompt_sizer - ) - except Exception as e: - logger.error(e) - continue - row.update( - { - "model_short_id": short_id, - "model_id": full_model_id, - "placement_sharding": sharding, - "placement_instance_meta": instance_meta, - "placement_nodes": n_nodes, - "instance_id": instance_id, - "pp_tokens": actual_pp_tokens, - "tg": tg, - "repeat_index": r, - } - ) - runs.append(row) - all_rows.append(row) + # If pp and tg lists have same length, run in tandem (zip) + # Otherwise (or if --all-combinations), run all combinations (cartesian product) + if use_combinations: + pp_tg_pairs = list(itertools.product(pp_list, tg_list)) + else: + pp_tg_pairs = list(zip(pp_list, tg_list, strict=True)) - if runs: - prompt_tps = mean(x["stats"]["prompt_tps"] for x in runs) - gen_tps = mean(x["stats"]["generation_tps"] for x in runs) - ptok = mean(x["stats"]["prompt_tokens"] for x in runs) - gtok = mean(x["stats"]["generation_tokens"] for x in runs) - peak = mean( - x["stats"]["peak_memory_usage"]["inBytes"] for x in runs + for pp, tg in pp_tg_pairs: + runs: list[dict[str, Any]] = [] + for r in range(args.repeat): + time.sleep(3) + try: + row, actual_pp_tokens = run_one_completion( + client, full_model_id, pp, tg, prompt_sizer ) + except Exception as e: + logger.error(e) + continue + row.update( + { + "model_short_id": short_id, + "model_id": full_model_id, + "placement_sharding": sharding, + "placement_instance_meta": instance_meta, + "placement_nodes": n_nodes, + "instance_id": instance_id, + "pp_tokens": actual_pp_tokens, + "tg": tg, + "repeat_index": r, + } + ) + runs.append(row) + all_rows.append(row) - logger.info( - f"prompt_tps={prompt_tps:.2f} gen_tps={gen_tps:.2f} " - f"prompt_tokens={ptok} gen_tokens={gtok} " - f"peak_memory={format_peak_memory(peak)}\n" - ) - time.sleep(2) + if runs: + prompt_tps = mean(x["stats"]["prompt_tps"] for x in runs) + gen_tps = mean(x["stats"]["generation_tps"] for x in runs) + ptok = mean(x["stats"]["prompt_tokens"] for x in runs) + gtok = mean(x["stats"]["generation_tokens"] for x in runs) + peak = mean( + x["stats"]["peak_memory_usage"]["inBytes"] for x in runs + ) + + logger.info( + f"prompt_tps={prompt_tps:.2f} gen_tps={gen_tps:.2f} " + f"prompt_tokens={ptok} gen_tokens={gtok} " + f"peak_memory={format_peak_memory(peak)}\n" + ) + time.sleep(2) finally: try: client.request_json("DELETE", f"/instance/{instance_id}") diff --git a/dashboard/parts.nix b/dashboard/parts.nix index 487078d5..80df5c7e 100644 --- a/dashboard/parts.nix +++ b/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 '' > $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" \ "$@" ''; }; diff --git a/dashboard/src/lib/components/ChatForm.svelte b/dashboard/src/lib/components/ChatForm.svelte index 42a9bb4c..6801287d 100644 --- a/dashboard/src/lib/components/ChatForm.svelte +++ b/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 @@ EDIT + {:else if isEditOnlyWithoutImage} + + + + + EDIT + {:else if isImageModel()} (null); let expandedThinkingMessageIds = $state>(new Set()); + // Uncertainty heatmap toggle + let heatmapMessageIds = $state>(new Set()); + + function toggleHeatmap(messageId: string) { + const next = new Set(heatmapMessageIds); + if (next.has(messageId)) { + next.delete(messageId); + } else { + next.add(messageId); + } + heatmapMessageIds = next; + } + + function isHeatmapVisible(messageId: string): boolean { + return heatmapMessageIds.has(messageId); + } + function formatTimestamp(timestamp: number): string { return new Date(timestamp).toLocaleTimeString("en-US", { hour12: false, @@ -548,13 +567,23 @@ > {:else if message.content || (loading && !message.attachments?.some((a) => a.type === "generated-image"))} - - {#if loading && !message.content} - + {#if isHeatmapVisible(message.id) && message.tokens && message.tokens.length > 0} + + regenerateFromToken(message.id, tokenIndex)} + /> + {:else} + + {#if loading && !message.content} + + {/if} {/if} {/if} @@ -629,6 +658,35 @@ {/if} + + {#if message.role === "assistant" && message.tokens && message.tokens.length > 0} + + {/if} + {#if message.role === "assistant" && isLastAssistantMessage(message.id) && !loading} + + + {#if hasFavorites} + + {/if} + + + + +

+ + + {#each families as family} + + {/each} + diff --git a/dashboard/src/lib/components/HuggingFaceResultItem.svelte b/dashboard/src/lib/components/HuggingFaceResultItem.svelte new file mode 100644 index 00000000..566d8e17 --- /dev/null +++ b/dashboard/src/lib/components/HuggingFaceResultItem.svelte @@ -0,0 +1,127 @@ + + +
+
+
+ {modelName} + {#if isAdded} + Added + {/if} +
+
+ {model.author} + + + + + {formatNumber(model.downloads)} + + + + + + {formatNumber(model.likes)} + +
+
+ +
+ {#if isAdded} + + {:else} + + {/if} +
+
diff --git a/dashboard/src/lib/components/ImageParamsPanel.svelte b/dashboard/src/lib/components/ImageParamsPanel.svelte index a7d23798..d3a9fbef 100644 --- a/dashboard/src/lib/components/ImageParamsPanel.svelte +++ b/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 @@
- -
- SIZE: -
- -
- + -
-
- - {#if isSizeDropdownOpen} - - - - -
-
- {#each sizeOptions as size} - - {/each} + {params.size} + +
+ + +
- {/if} -
+ + {#if isSizeDropdownOpen} + + + + +
+
+ {#each sizeOptions as size} + + {/each} +
+
+ {/if} +
+ {/if}
@@ -325,6 +357,59 @@
+ + {#if !isEditMode} +
+ IMAGES: + +
+ {/if} + + +
+ STREAM: + +
+ + + {#if params.stream} +
+ PARTIALS: + +
+ {/if} + {#if isEditMode}
diff --git a/dashboard/src/lib/components/ModelFilterPopover.svelte b/dashboard/src/lib/components/ModelFilterPopover.svelte new file mode 100644 index 00000000..5406618a --- /dev/null +++ b/dashboard/src/lib/components/ModelFilterPopover.svelte @@ -0,0 +1,182 @@ + + + + + +
e.stopPropagation()} + role="dialog" + aria-label="Filter options" +> +
+ +
+

Capabilities

+
+ {#each availableCapabilities as cap} + {@const isSelected = filters.capabilities.includes(cap.id)} + + {/each} +
+
+ + +
+

Model Size

+
+ {#each sizeRanges as range} + {@const isSelected = + filters.sizeRange && + filters.sizeRange.min === range.min && + filters.sizeRange.max === range.max} + + {/each} +
+
+ + + +
+
diff --git a/dashboard/src/lib/components/ModelPickerGroup.svelte b/dashboard/src/lib/components/ModelPickerGroup.svelte new file mode 100644 index 00000000..b3ad425b --- /dev/null +++ b/dashboard/src/lib/components/ModelPickerGroup.svelte @@ -0,0 +1,324 @@ + + +
+ +
{ + if (group.hasMultipleVariants) { + onToggleExpand(); + } else { + const modelId = group.variants[0]?.id; + if (modelId && canModelFit(modelId)) { + onSelectModel(modelId); + } + } + }} + role="button" + tabindex="0" + onkeydown={(e) => { + if (e.key === "Enter" || e.key === " ") { + e.preventDefault(); + if (group.hasMultipleVariants) { + onToggleExpand(); + } else { + const modelId = group.variants[0]?.id; + if (modelId && canModelFit(modelId)) { + onSelectModel(modelId); + } + } + } + }} + > + + {#if group.hasMultipleVariants} + + + + {:else} +
+ {/if} + + +
+
+ + {group.name} + + + {#each group.capabilities.filter((c) => c !== "text") as cap} + {#if cap === "thinking"} + + + + {:else if cap === "code"} + + + + {:else if cap === "vision"} + + + + + {:else if cap === "image_gen"} + + + + + + {/if} + {/each} +
+
+ + + {#if !group.hasMultipleVariants && group.smallestVariant?.storage_size_megabytes} + + {formatSize(group.smallestVariant.storage_size_megabytes)} + + {/if} + + + {#if group.hasMultipleVariants} + + {group.variants.length} variants + + {/if} + + + {#if isMainSelected} + + + + {/if} + + + + + + +
+ + + {#if isExpanded && group.hasMultipleVariants} +
+ {#each group.variants as variant} + {@const modelCanFit = canModelFit(variant.id)} + {@const isSelected = selectedModelId === variant.id} + + {/each} +
+ {/if} +
diff --git a/dashboard/src/lib/components/ModelPickerModal.svelte b/dashboard/src/lib/components/ModelPickerModal.svelte new file mode 100644 index 00000000..40827a0e --- /dev/null +++ b/dashboard/src/lib/components/ModelPickerModal.svelte @@ -0,0 +1,748 @@ + + + + +{#if isOpen} + + + + + + + + {#if infoGroup} +
(infoGroup = null)} + role="presentation" + >
+ + {/if} +{/if} diff --git a/dashboard/src/lib/components/TokenHeatmap.svelte b/dashboard/src/lib/components/TokenHeatmap.svelte new file mode 100644 index 00000000..d6de667c --- /dev/null +++ b/dashboard/src/lib/components/TokenHeatmap.svelte @@ -0,0 +1,236 @@ + + +
+ {#each tokens as tokenData, i (i)} + handleMouseEnter(e, tokenData, i)} + onmouseleave={handleMouseLeave}>{tokenData.token} + {/each} +
+ + +{#if hoveredToken} +
+
+ +
+ Token: + "{hoveredToken.token.token}" + {formatProbability(hoveredToken.token.probability)} +
+ +
+ logprob: {formatLogprob(hoveredToken.token.logprob)} +
+ + + {#if hoveredToken.token.topLogprobs.length > 0} +
+
Alternatives:
+ {#each hoveredToken.token.topLogprobs.slice(0, 5) as alt, idx (idx)} + {@const altProb = Math.exp(alt.logprob)} +
+ "{alt.token}" + {formatProbability(altProb)} +
+ {/each} +
+ {/if} + + + {#if onRegenerateFrom} + + {/if} +
+ +
+
+
+
+{/if} + + diff --git a/dashboard/src/lib/components/TopologyGraph.svelte b/dashboard/src/lib/components/TopologyGraph.svelte index 41748977..c0acc39d 100644 --- a/dashboard/src/lib/components/TopologyGraph.svelte +++ b/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; + filteredNodes?: Set; + 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(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 @@