Merge branch 'main' into sami/flash

.github/workflows/build-app.yml

@@ -1,16 +1,5 @@
 name: Build EXO macOS DMG
 
-# Release workflow:
-# 1. Create a draft GitHub Release with the tag name (e.g. v1.0.0) and write release notes in markdown
-# 2. Push the tag: git tag v1.0.0 && git push origin v1.0.0
-# 3. This workflow builds, signs, and notarizes the DMG
-# 4. Release notes are embedded in appcast.xml for Sparkle (rendered as markdown)
-# 5. DMG and appcast.xml are uploaded to S3
-# 6. The draft GitHub Release is published with the DMG attached
-#
-# For alpha releases (e.g. v1.0.0-alpha.1): draft release and notes are optional.
-# If no draft exists, a release is auto-created with generated notes.
-
 on:
   workflow_dispatch:
   push:
@@ -22,10 +11,8 @@ on:
 jobs:
   build-macos-app:
     runs-on: "macos-26"
-    permissions:
-      contents: write
     env:
-      SPARKLE_VERSION: 2.9.0-beta.1
+      SPARKLE_VERSION: 2.8.1
       SPARKLE_DOWNLOAD_PREFIX: ${{ secrets.SPARKLE_DOWNLOAD_PREFIX }}
       SPARKLE_FEED_URL: ${{ secrets.SPARKLE_FEED_URL }}
       SPARKLE_ED25519_PUBLIC: ${{ secrets.SPARKLE_ED25519_PUBLIC }}
@@ -100,52 +87,6 @@ jobs:
             exit 1
           fi
 
-      - name: Fetch and validate release notes
-        if: github.ref_type == 'tag'
-        env:
-          GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
-        run: |
-          # Find draft release by name using gh release list (more reliable with default token)
-          echo "Looking for draft release named '$GITHUB_REF_NAME'..."
-          DRAFT_EXISTS=$(gh release list --json name,isDraft --jq ".[] | select(.isDraft == true) | select(.name == \"$GITHUB_REF_NAME\") | .name" 2>/dev/null || echo "")
-
-          if [[ -z "$DRAFT_EXISTS" ]]; then
-            if [[ "$IS_ALPHA" == "true" ]]; then
-              echo "No draft release found for alpha tag $GITHUB_REF_NAME (optional for alphas)"
-              echo "HAS_RELEASE_NOTES=false" >> $GITHUB_ENV
-              exit 0
-            fi
-            echo "ERROR: No draft release found for tag $GITHUB_REF_NAME"
-            echo "Please create a draft release with release notes before pushing the tag."
-            exit 1
-          fi
-
-          # Fetch full release details via API to get body and ID
-          echo "Found draft release, fetching details..."
-          RELEASE_JSON=$(gh api repos/${{ github.repository }}/releases --jq ".[] | select(.draft == true) | select(.name == \"$GITHUB_REF_NAME\")" 2>/dev/null || echo "")
-
-          # Extract release notes
-          NOTES=$(echo "$RELEASE_JSON" | jq -r '.body // ""')
-          if [[ -z "$NOTES" || "$NOTES" == "null" ]]; then
-            if [[ "$IS_ALPHA" == "true" ]]; then
-              echo "Draft release has no notes (optional for alphas)"
-              echo "HAS_RELEASE_NOTES=false" >> $GITHUB_ENV
-              exit 0
-            fi
-            echo "ERROR: Draft release exists but has no release notes"
-            echo "Please add release notes to the draft release before pushing the tag."
-            exit 1
-          fi
-
-          # Save release ID for later publishing
-          RELEASE_ID=$(echo "$RELEASE_JSON" | jq -r '.id')
-          echo "DRAFT_RELEASE_ID=$RELEASE_ID" >> $GITHUB_ENV
-          echo "HAS_RELEASE_NOTES=true" >> $GITHUB_ENV
-
-          echo "Found draft release (ID: $RELEASE_ID), saving release notes..."
-          echo "$NOTES" > /tmp/release_notes.md
-          echo "RELEASE_NOTES_FILE=/tmp/release_notes.md" >> $GITHUB_ENV
-
       # ============================================================
       # Install dependencies
       # ============================================================
@@ -363,28 +304,6 @@ jobs:
             $CHANNEL_FLAG \
             .
 
-      - name: Inject release notes into appcast
-        if: github.ref_type == 'tag' && env.HAS_RELEASE_NOTES == 'true'
-        env:
-          RELEASE_VERSION: ${{ env.RELEASE_VERSION }}
-        run: |
-          # Inject markdown release notes with sparkle:format="markdown" (Sparkle 2.9+)
-          export NOTES=$(cat "$RELEASE_NOTES_FILE")
-
-          # Insert description after the enclosure tag for this version
-          awk '
-            /<enclosure[^>]*>/ && index($0, ENVIRON["RELEASE_VERSION"]) {
-              print
-              print "            <description sparkle:format=\"markdown\"><![CDATA["
-              print ENVIRON["NOTES"]
-              print "            ]]></description>"
-              next
-            }
-            { print }
-          ' output/appcast.xml > output/appcast.xml.tmp && mv output/appcast.xml.tmp output/appcast.xml
-
-          echo "Injected markdown release notes for version $RELEASE_VERSION"
-
       # ============================================================
       # Upload artifacts
       # ============================================================
@@ -417,26 +336,3 @@ jobs:
             aws s3 cp "$DMG_NAME" "s3://${SPARKLE_S3_BUCKET}/${PREFIX}EXO-latest.dmg"
             aws s3 cp appcast.xml "s3://${SPARKLE_S3_BUCKET}/${PREFIX}appcast.xml" --content-type application/xml --cache-control no-cache
           fi
-
-      - name: Publish GitHub Release
-        if: github.ref_type == 'tag'
-        env:
-          GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
-        run: |
-          DMG_PATH="output/EXO-${RELEASE_VERSION}.dmg"
-
-          if [[ "$HAS_RELEASE_NOTES" == "true" ]]; then
-            # Update the draft release with the tag and upload DMG
-            gh api --method PATCH "repos/${{ github.repository }}/releases/$DRAFT_RELEASE_ID" \
-              -f tag_name="$GITHUB_REF_NAME" \
-              -F draft=false
-            gh release upload "$GITHUB_REF_NAME" "$DMG_PATH" --clobber
-            echo "Published release $GITHUB_REF_NAME with DMG attached"
-          else
-            # Alpha without draft release - create one with auto-generated notes
-            gh release create "$GITHUB_REF_NAME" "$DMG_PATH" \
-              --title "$GITHUB_REF_NAME" \
-              --generate-notes \
-              --prerelease
-            echo "Created alpha release $GITHUB_REF_NAME with auto-generated notes"
-          fi

AGENTS.md

@@ -40,31 +40,6 @@ uv run ruff check
 nix fmt
 ```
 
-## Pre-Commit Checks (REQUIRED)
-
-**IMPORTANT: Always run these checks before committing code. CI will fail if these don't pass.**
-
-```bash
-# 1. Type checking - MUST pass with 0 errors
-uv run basedpyright
-
-# 2. Linting - MUST pass
-uv run ruff check
-
-# 3. Formatting - MUST be applied
-nix fmt
-
-# 4. Tests - MUST pass
-uv run pytest
-```
-
-Run all checks in sequence:
-```bash
-uv run basedpyright && uv run ruff check && nix fmt && uv run pytest
-```
-
-If `nix fmt` changes any files, stage them before committing. The CI runs `nix flake check` which verifies formatting, linting, and runs Rust tests.
-
 ## Architecture
 
 ### Node Composition

Cargo.lock

@@ -4340,6 +4340,25 @@ dependencies = [
  "libc",
 ]
 
+[[package]]
+name = "system_custodian"
+version = "0.0.1"
+dependencies = [
+ "delegate",
+ "derive_more",
+ "either",
+ "extend",
+ "futures",
+ "futures-timer",
+ "impl-trait-for-tuples",
+ "keccak-const",
+ "log",
+ "thiserror 2.0.17",
+ "tokio",
+ "tracing-subscriber",
+ "util",
+]
+
 [[package]]
 name = "tagptr"
 version = "0.2.0"

Cargo.toml

@@ -3,6 +3,7 @@ resolver = "3"
 members = [
     "rust/networking",
     "rust/exo_pyo3_bindings",
+    "rust/system_custodian",
     "rust/util",
 ]
 
@@ -24,6 +25,7 @@ opt-level = 3
 [workspace.dependencies]
 ## Crate members as common dependencies
 networking = { path = "rust/networking" }
+system_custodian = { path = "rust/system_custodian" }
 util = { path = "rust/util" }
 
 # Proc-macro authoring tools

app/EXO/EXO.xcodeproj/project.pbxproj

@@ -585,7 +585,7 @@
 			repositoryURL = "https://github.com/sparkle-project/Sparkle.git";
 			requirement = {
 				kind = upToNextMajorVersion;
-				minimumVersion = 2.9.0-beta.1;
+				minimumVersion = 2.8.1;
 			};
 		};
 /* End XCRemoteSwiftPackageReference section */

app/EXO/EXO.xcodeproj/project.xcworkspace/xcshareddata/swiftpm/Package.resolved

@@ -6,8 +6,8 @@
       "kind" : "remoteSourceControl",
       "location" : "https://github.com/sparkle-project/Sparkle.git",
       "state" : {
-        "revision" : "e641adb41915a8409895e2e30666aa64e487b637",
-        "version" : "2.9.0-beta.1"
+        "revision" : "5581748cef2bae787496fe6d61139aebe0a451f6",
+        "version" : "2.8.1"
       }
     }
   ],

app/EXO/EXO/ContentView.swift

@@ -56,11 +56,6 @@ struct ContentView: View {
     }
 
     private var shouldShowLocalNetworkWarning: Bool {
-        // Show warning if local network is not working and EXO is running.
-        // The checker uses a longer timeout on first launch to allow time for
-        // the permission prompt, so this correctly handles both:
-        // 1. User denied permission on first launch
-        // 2. Permission broke after restart (macOS TCC bug)
         if case .notWorking = localNetworkChecker.status {
             return controller.status != .stopped
         }

app/EXO/EXO/Services/LocalNetworkChecker.swift

@@ -5,8 +5,8 @@ import os.log
 /// Checks if the app's local network permission is actually functional.
 ///
 /// macOS local network permission can appear enabled in System Preferences but not
-/// actually work after a restart. This service uses NWConnection to mDNS multicast
-/// to verify actual connectivity.
+/// actually work after a restart. This service detects this by creating a UDP
+/// connection to the mDNS multicast address (224.0.0.251:5353).
 @MainActor
 final class LocalNetworkChecker: ObservableObject {
     enum Status: Equatable {
@@ -35,43 +35,30 @@ final class LocalNetworkChecker: ObservableObject {
     }
 
     private static let logger = Logger(subsystem: "io.exo.EXO", category: "LocalNetworkChecker")
-    private static let hasCompletedInitialCheckKey = "LocalNetworkChecker.hasCompletedInitialCheck"
 
     @Published private(set) var status: Status = .unknown
+    @Published private(set) var lastConnectionState: String = "none"
 
     private var connection: NWConnection?
     private var checkTask: Task<Void, Never>?
 
-    /// Whether we've completed at least one check (stored in UserDefaults)
-    private var hasCompletedInitialCheck: Bool {
-        get { UserDefaults.standard.bool(forKey: Self.hasCompletedInitialCheckKey) }
-        set { UserDefaults.standard.set(newValue, forKey: Self.hasCompletedInitialCheckKey) }
-    }
-
     /// Checks if local network access is working.
     func check() {
         checkTask?.cancel()
         status = .checking
-
-        // Use longer timeout on first launch to allow time for permission prompt
-        let isFirstCheck = !hasCompletedInitialCheck
-        let timeout: UInt64 = isFirstCheck ? 30_000_000_000 : 3_000_000_000
+        lastConnectionState = "connecting"
 
         checkTask = Task { [weak self] in
             guard let self else { return }
-
-            Self.logger.info("Checking local network connectivity (first check: \(isFirstCheck))")
-            let result = await self.checkConnectivity(timeout: timeout)
+            let result = await self.performCheck()
             self.status = result
-            self.hasCompletedInitialCheck = true
-
             Self.logger.info("Local network check complete: \(result.displayText)")
         }
     }
 
-    /// Checks connectivity using NWConnection to mDNS multicast.
-    /// The connection attempt triggers the permission prompt if not yet shown.
-    private func checkConnectivity(timeout: UInt64) async -> Status {
+    private func performCheck() async -> Status {
+        Self.logger.info("Checking local network access via UDP multicast")
+
         connection?.cancel()
         connection = nil
 
@@ -97,7 +84,22 @@ final class LocalNetworkChecker: ObservableObject {
                 continuation.resume(returning: status)
             }
 
-            conn.stateUpdateHandler = { state in
+            conn.stateUpdateHandler = { [weak self] state in
+                let stateStr: String
+                switch state {
+                case .setup: stateStr = "setup"
+                case .preparing: stateStr = "preparing"
+                case .ready: stateStr = "ready"
+                case .waiting(let e): stateStr = "waiting(\(e))"
+                case .failed(let e): stateStr = "failed(\(e))"
+                case .cancelled: stateStr = "cancelled"
+                @unknown default: stateStr = "unknown"
+                }
+
+                Task { @MainActor in
+                    self?.lastConnectionState = stateStr
+                }
+
                 switch state {
                 case .ready:
                     resumeOnce(.working)
@@ -106,7 +108,6 @@ final class LocalNetworkChecker: ObservableObject {
                     if errorStr.contains("54") || errorStr.contains("ECONNRESET") {
                         resumeOnce(.notWorking(reason: "Connection blocked"))
                     }
-                // Otherwise keep waiting - might be showing permission prompt
                 case .failed(let error):
                     let errorStr = "\(error)"
                     if errorStr.contains("65") || errorStr.contains("EHOSTUNREACH")
@@ -126,7 +127,7 @@ final class LocalNetworkChecker: ObservableObject {
             conn.start(queue: .main)
 
             Task {
-                try? await Task.sleep(nanoseconds: timeout)
+                try? await Task.sleep(nanoseconds: 3_000_000_000)
                 let state = conn.state
                 switch state {
                 case .ready:

bench/exo_bench.py

@@ -3,7 +3,6 @@
 from __future__ import annotations
 
 import argparse
-import contextlib
 import http.client
 import json
 import os
@@ -27,7 +26,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 = 2400.0):
         self.host = host
         self.port = port
         self.timeout_s = timeout_s
@@ -105,46 +104,22 @@ def runner_ready(runner: dict[str, Any]) -> bool:
     return "RunnerReady" in runner
 
 
-def runner_failed(runner: dict[str, Any]) -> bool:
-    return "RunnerFailed" in runner
-
-
-def get_runner_failed_message(runner: dict[str, Any]) -> str | None:
-    if "RunnerFailed" in runner:
-        return runner["RunnerFailed"].get("errorMessage")
-    return None
-
-
 def wait_for_instance_ready(
     client: ExoClient, instance_id: str, timeout: float = 24000.0
 ) -> None:
     start_time = time.time()
-    instance_existed = False
     while time.time() - start_time < timeout:
         state = client.request_json("GET", "/state")
         instances = state.get("instances", {})
 
         if instance_id not in instances:
-            if instance_existed:
-                # Instance was deleted after being created - likely due to runner failure
-                raise RuntimeError(
-                    f"Instance {instance_id} was deleted (runner may have failed)"
-                )
             time.sleep(0.1)
             continue
 
-        instance_existed = True
         instance = instances[instance_id]
         runner_ids = runner_ids_from_instance(instance)
         runners = state.get("runners", {})
 
-        # Check for failed runners first
-        for rid in runner_ids:
-            runner = runners.get(rid, {})
-            if runner_failed(runner):
-                error_msg = get_runner_failed_message(runner) or "Unknown error"
-                raise RuntimeError(f"Runner {rid} failed: {error_msg}")
-
         if all(runner_ready(runners.get(rid, {})) for rid in runner_ids):
             return
 
@@ -266,9 +241,6 @@ class PromptSizer:
             ids = tokenizer.apply_chat_template(
                 messages, tokenize=True, add_generation_prompt=True
             )
-            # Fix for transformers 5.x
-            if hasattr(ids, "input_ids"):
-                ids = ids.input_ids
             return int(len(ids))
 
         return count_fn
@@ -324,12 +296,6 @@ def main() -> int:
         default=4,
         help="Only consider placements using <= this many nodes.",
     )
-    ap.add_argument(
-        "--min-nodes",
-        type=int,
-        default=1,
-        help="Only consider placements using >= this many nodes.",
-    )
     ap.add_argument(
         "--instance-meta", choices=["ring", "jaccl", "both"], default="both"
     )
@@ -351,7 +317,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=2400.0, help="HTTP timeout (seconds)."
     )
     ap.add_argument(
         "--json-out",
@@ -430,7 +396,7 @@ def main() -> int:
         ):
             continue
 
-        if args.min_nodes <= n <= args.max_nodes:
+        if 0 < n <= args.max_nodes:
             selected.append(p)
 
     if not selected:
@@ -472,13 +438,7 @@ def main() -> int:
         )
 
         client.request_json("POST", "/instance", body={"instance": instance})
-        try:
-            wait_for_instance_ready(client, instance_id)
-        except (RuntimeError, TimeoutError) as e:
-            logger.error(f"Failed to initialize placement: {e}")
-            with contextlib.suppress(ExoHttpError):
-                client.request_json("DELETE", f"/instance/{instance_id}")
-            continue
+        wait_for_instance_ready(client, instance_id)
 
         time.sleep(1)
 

dashboard/package-lock.json

@@ -863,7 +863,6 @@
 			"integrity": "sha512-oH8tXw7EZnie8FdOWYrF7Yn4IKrqTFHhXvl8YxXxbKwTMcD/5NNCryUSEXRk2ZR4ojnub0P8rNrsVGHXWqIDtA==",
 			"dev": true,
 			"license": "MIT",
-			"peer": true,
 			"dependencies": {
 				"@standard-schema/spec": "^1.0.0",
 				"@sveltejs/acorn-typescript": "^1.0.5",
@@ -903,7 +902,6 @@
 			"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,7 +1518,6 @@
 			"integrity": "sha512-LCCV0HdSZZZb34qifBsyWlUmok6W7ouER+oQIGBScS8EsZsQbrtFTUrDX4hOl+CS6p7cnNC4td+qrSVGSCTUfQ==",
 			"dev": true,
 			"license": "MIT",
-			"peer": true,
 			"dependencies": {
 				"undici-types": "~6.21.0"
 			}
@@ -1530,7 +1527,6 @@
 			"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"
 			},
@@ -1943,7 +1939,6 @@
 			"integrity": "sha512-fmTRWbNMmsmWq6xJV8D19U/gw/bwrHfNXxrIN+HfZgnzqTHp9jOmKMhsTUjXOJnZOdZY9Q28y4yebKzqDKlxlQ==",
 			"dev": true,
 			"license": "ISC",
-			"peer": true,
 			"engines": {
 				"node": ">=12"
 			}
@@ -2651,7 +2646,6 @@
 			"integrity": "sha512-5gTmgEY/sqK6gFXLIsQNH19lWb4ebPDLA4SdLP7dsWkIXHWlG66oPuVvXSGFPppYZz8ZDZq0dYYrbHfBCVUb1Q==",
 			"dev": true,
 			"license": "MIT",
-			"peer": true,
 			"engines": {
 				"node": ">=12"
 			},
@@ -2839,7 +2833,6 @@
 			"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",
@@ -2984,7 +2977,6 @@
 			"integrity": "sha512-jl1vZzPDinLr9eUt3J/t7V6FgNEw9QjvBPdysz9KfQDD41fQrC2Y4vKQdiaUpFT4bXlb1RHhLpp8wtm6M5TgSw==",
 			"dev": true,
 			"license": "Apache-2.0",
-			"peer": true,
 			"bin": {
 				"tsc": "bin/tsc",
 				"tsserver": "bin/tsserver"
@@ -3006,7 +2998,6 @@
 			"integrity": "sha512-+Oxm7q9hDoLMyJOYfUYBuHQo+dkAloi33apOPP56pzj+vsdJDzr+j1NISE5pyaAuKL4A3UD34qd0lx5+kfKp2g==",
 			"dev": true,
 			"license": "MIT",
-			"peer": true,
 			"dependencies": {
 				"esbuild": "^0.25.0",
 				"fdir": "^6.4.4",

dashboard/src/lib/components/ChatForm.svelte

@@ -60,39 +60,12 @@
 		return models;
 	});
 
-	// Track previous model IDs to detect newly added models (plain variable to avoid reactive loop)
-	let previousModelIds: Set<string> = new Set();
-
-	// Auto-select the first available model if none is selected, if current selection is stale, or if a new model is added
+	// Auto-select the first available model if none is selected
 	$effect(() => {
 		const models = availableModels();
-		const currentModelIds = new Set(models.map(m => m.id));
-
-		if (models.length > 0) {
-			// Find newly added models (in current but not in previous)
-			const newModels = models.filter(m => !previousModelIds.has(m.id));
-
-			// If no model selected, select the first available
-			if (!currentModel) {
-				setSelectedChatModel(models[0].id);
-			}
-			// If current model is stale (no longer has a running instance), reset to first available
-			else if (!models.some(m => m.id === currentModel)) {
-				setSelectedChatModel(models[0].id);
-			}
-			// If a new model was just added, select it
-			else if (newModels.length > 0 && previousModelIds.size > 0) {
-				setSelectedChatModel(newModels[0].id);
-			}
-		} else {
-			// No instances running - clear the selected model
-			if (currentModel) {
-				setSelectedChatModel('');
-			}
+		if (models.length > 0 && !currentModel) {
+			setSelectedChatModel(models[0].id);
 		}
-
-		// Update previous model IDs for next comparison
-		previousModelIds = currentModelIds;
 	});
 
 	function getInstanceModelId(instanceWrapped: unknown): string {

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

@@ -69,8 +69,6 @@ export interface Instance {
 		runnerToShard?: Record<string, unknown>;
 		nodeToRunner?: Record<string, string>;
 	};
-	draftModel?: string;
-	numDraftTokens?: number;
 }
 
 interface RawNodeProfile {

dashboard/src/routes/+page.svelte

@@ -47,7 +47,7 @@ const sidebarVisible = $derived(chatSidebarVisible());
 	let mounted = $state(false);
 
 	// Instance launch state
-	let models = $state<Array<{id: string, hugging_face_id?: string, name?: string, storage_size_megabytes?: number}>>([]);
+	let models = $state<Array<{id: string, name?: string, storage_size_megabytes?: number}>>([]);
 	let selectedSharding = $state<'Pipeline' | 'Tensor'>('Pipeline');
 	type InstanceMeta = 'MlxRing' | 'MlxIbv' | 'MlxJaccl';
 	
@@ -58,8 +58,6 @@ const sidebarVisible = $derived(chatSidebarVisible());
 		sharding: 'Pipeline' | 'Tensor';
 		instanceType: InstanceMeta;
 		minNodes: number;
-		draftModel: string | null;
-		numDraftTokens: number;
 	}
 	
 	function saveLaunchDefaults(): void {
@@ -68,8 +66,6 @@ const sidebarVisible = $derived(chatSidebarVisible());
 			sharding: selectedSharding,
 			instanceType: selectedInstanceType,
 			minNodes: selectedMinNodes,
-			draftModel: selectedDraftModel,
-			numDraftTokens: selectedNumDraftTokens,
 		};
 		try {
 			localStorage.setItem(LAUNCH_DEFAULTS_KEY, JSON.stringify(defaults));
@@ -92,36 +88,24 @@ const sidebarVisible = $derived(chatSidebarVisible());
 	function applyLaunchDefaults(availableModels: Array<{id: string}>, maxNodes: number): void {
 		const defaults = loadLaunchDefaults();
 		if (!defaults) return;
-
+		
 		// Apply sharding and instance type unconditionally
 		selectedSharding = defaults.sharding;
 		selectedInstanceType = defaults.instanceType;
-
+		
 		// Apply minNodes if valid (between 1 and maxNodes)
 		if (defaults.minNodes && defaults.minNodes >= 1 && defaults.minNodes <= maxNodes) {
 			selectedMinNodes = defaults.minNodes;
 		}
-
+		
 		// Only apply model if it exists in the available models
 		if (defaults.modelId && availableModels.some(m => m.id === defaults.modelId)) {
 			selectPreviewModel(defaults.modelId);
 		}
-
-		// Apply draft model if it exists in the available models (check against hugging_face_id)
-		if (defaults.draftModel && availableModels.some(m => (m as {hugging_face_id?: string}).hugging_face_id === defaults.draftModel)) {
-			selectedDraftModel = defaults.draftModel;
-		}
-
-		// Apply num draft tokens if valid
-		if (defaults.numDraftTokens && defaults.numDraftTokens >= 1 && defaults.numDraftTokens <= 10) {
-			selectedNumDraftTokens = defaults.numDraftTokens;
-		}
 	}
 	
 	let selectedInstanceType = $state<InstanceMeta>('MlxRing');
 	let selectedMinNodes = $state<number>(1);
-	let selectedDraftModel = $state<string | null>(null);
-	let selectedNumDraftTokens = $state<number>(4);
 	let minNodesInitialized = $state(false);
 	let launchingModelId = $state<string | null>(null);
 let instanceDownloadExpandedNodes = $state<Set<string>>(new Set());
@@ -129,8 +113,6 @@ let instanceDownloadExpandedNodes = $state<Set<string>>(new Set());
 	// Custom dropdown state
 	let isModelDropdownOpen = $state(false);
 	let modelDropdownSearch = $state('');
-	let isDraftModelDropdownOpen = $state(false);
-	let draftModelDropdownSearch = $state('');
 	
 	// Slider dragging state
 	let isDraggingSlider = $state(false);
@@ -380,39 +362,49 @@ function toggleInstanceDownloadDetails(nodeId: string): void {
 
 	async function launchInstance(modelId: string, specificPreview?: PlacementPreview | null) {
 		if (!modelId || launchingModelId) return;
-
+		
 		launchingModelId = modelId;
-
+		
 		try {
 			// Use the specific preview if provided, otherwise fall back to filtered preview
 			const preview = specificPreview ?? filteredPreview();
-
-			let response: Response;
-
-			// Use /place_instance endpoint - it handles placement and creation in one step
-			// This also supports draft_model for speculative decoding
-			const placePayload = {
-				model_id: modelId,
-				sharding: preview?.sharding ?? selectedSharding,
-				instance_meta: preview?.instance_meta ?? selectedInstanceType,
-				min_nodes: selectedMinNodes,
-				draft_model: selectedDraftModel,
-				num_draft_tokens: selectedDraftModel ? selectedNumDraftTokens : 4,
-			};
-
-			response = await fetch('/place_instance', {
+			
+			let instanceData: unknown;
+			
+			if (preview?.instance) {
+				// Use the instance from the preview
+				instanceData = preview.instance;
+			} else {
+				// Fallback: GET placement from API
+				const placementResponse = await fetch(
+					`/instance/placement?model_id=${encodeURIComponent(modelId)}&sharding=${selectedSharding}&instance_meta=${selectedInstanceType}&min_nodes=${selectedMinNodes}`
+				);
+				
+				if (!placementResponse.ok) {
+					const errorText = await placementResponse.text();
+					console.error('Failed to get placement:', errorText);
+					return;
+				}
+				
+				instanceData = await placementResponse.json();
+			}
+			
+			// POST the instance to create it
+			const response = await fetch('/instance', {
 				method: 'POST',
 				headers: { 'Content-Type': 'application/json' },
-				body: JSON.stringify(placePayload)
+				body: JSON.stringify({ instance: instanceData })
 			});
-
+			
 			if (!response.ok) {
 				const errorText = await response.text();
 				console.error('Failed to launch instance:', errorText);
 			} else {
-				// Always auto-select the newly launched model so the user chats to what they just launched
-				setSelectedChatModel(modelId);
-
+				// Auto-select the launched model only if no model is currently selected
+				if (!selectedChatModel()) {
+					setSelectedChatModel(modelId);
+				}
+				
 				// Scroll to the bottom of instances container to show the new instance
 				// Use multiple attempts to ensure DOM has updated with the new instance
 				const scrollToBottom = () => {
@@ -771,10 +763,6 @@ function toggleInstanceDownloadDetails(nodeId: string): void {
 	async function deleteInstance(instanceId: string) {
 		if (!confirm(`Delete instance ${instanceId.slice(0, 8)}...?`)) return;
 		
-		// Get the model ID of the instance being deleted before we delete it
-		const deletedInstanceModelId = getInstanceModelId(instanceData[instanceId]);
-		const wasSelected = selectedChatModel() === deletedInstanceModelId;
-		
 		try {
 			const response = await fetch(`/instance/${instanceId}`, {
 				method: 'DELETE',
@@ -783,24 +771,6 @@ function toggleInstanceDownloadDetails(nodeId: string): void {
 			
 			if (!response.ok) {
 				console.error('Failed to delete instance:', response.status);
-			} else if (wasSelected) {
-				// If we deleted the currently selected model, switch to another available model
-				// Find another instance that isn't the one we just deleted
-				const remainingInstances = Object.entries(instanceData).filter(([id]) => id !== instanceId);
-				if (remainingInstances.length > 0) {
-					// Select the last instance (most recently added, since objects preserve insertion order)
-					const [, lastInstance] = remainingInstances[remainingInstances.length - 1];
-					const newModelId = getInstanceModelId(lastInstance);
-					if (newModelId && newModelId !== 'Unknown' && newModelId !== 'Unknown Model') {
-						setSelectedChatModel(newModelId);
-					} else {
-						// Clear selection if no valid model found
-						setSelectedChatModel('');
-					}
-				} else {
-					// No more instances, clear the selection
-					setSelectedChatModel('');
-				}
 			}
 		} catch (error) {
 			console.error('Error deleting instance:', error);
@@ -826,34 +796,30 @@ function toggleInstanceDownloadDetails(nodeId: string): void {
 	}
 
 	// Get instance details: type (MLX Ring/IBV), sharding (Pipeline/Tensor), and node names
-	function getInstanceInfo(instanceWrapped: unknown): {
-		instanceType: string;
-		sharding: string;
+	function getInstanceInfo(instanceWrapped: unknown): { 
+		instanceType: string; 
+		sharding: string; 
 		nodeNames: string[];
 		nodeIds: string[];
 		nodeCount: number;
-		draftModel: string | null;
-		numDraftTokens: number | null;
 	} {
 		const [instanceTag, instance] = getTagged(instanceWrapped);
 		if (!instance || typeof instance !== 'object') {
-			return { instanceType: 'Unknown', sharding: 'Unknown', nodeNames: [], nodeIds: [], nodeCount: 0, draftModel: null, numDraftTokens: null };
+			return { instanceType: 'Unknown', sharding: 'Unknown', nodeNames: [], nodeIds: [], nodeCount: 0 };
 		}
-
+		
 		// Instance type from tag
 		let instanceType = 'Unknown';
 		if (instanceTag === 'MlxRingInstance') instanceType = 'MLX Ring';
 		else if (instanceTag === 'MlxIbvInstance' || instanceTag === 'MlxJacclInstance') instanceType = 'MLX RDMA';
-
-		const inst = instance as {
-			shardAssignments?: {
-				nodeToRunner?: Record<string, string>;
+		
+		const inst = instance as { 
+			shardAssignments?: { 
+				nodeToRunner?: Record<string, string>; 
 				runnerToShard?: Record<string, unknown>;
-			};
-			draftModel?: string;
-			numDraftTokens?: number;
+			} 
 		};
-
+		
 		// Sharding strategy from first shard
 		let sharding = 'Unknown';
 		const runnerToShard = inst.shardAssignments?.runnerToShard || {};
@@ -864,7 +830,7 @@ function toggleInstanceDownloadDetails(nodeId: string): void {
 			else if (shardTag === 'TensorShardMetadata') sharding = 'Tensor';
 			else if (shardTag === 'PrefillDecodeShardMetadata') sharding = 'Prefill/Decode';
 		}
-
+		
 		// Node names from topology
 		const nodeToRunner = inst.shardAssignments?.nodeToRunner || {};
 		const nodeIds = Object.keys(nodeToRunner);
@@ -872,12 +838,8 @@ function toggleInstanceDownloadDetails(nodeId: string): void {
 			const node = data?.nodes?.[nodeId];
 			return node?.friendly_name || nodeId.slice(0, 8);
 		});
-
-		// Draft model for speculative decoding
-		const draftModel = inst.draftModel ?? null;
-		const numDraftTokens = inst.numDraftTokens ?? null;
-
-		return { instanceType, sharding, nodeNames, nodeIds, nodeCount: nodeIds.length, draftModel, numDraftTokens };
+		
+		return { instanceType, sharding, nodeNames, nodeIds, nodeCount: nodeIds.length };
 	}
 
 	function formatLastUpdate(): string {
@@ -1363,9 +1325,6 @@ function toggleInstanceDownloadDetails(nodeId: string): void {
 											<div class="pl-2">
 												<div class="text-exo-yellow text-xs font-mono tracking-wide truncate">{getInstanceModelId(instance)}</div>
 												<div class="text-white/60 text-xs font-mono">Strategy: <span class="text-white/80">{instanceInfo.sharding} ({instanceInfo.instanceType})</span></div>
-												{#if instanceInfo.draftModel}
-													<div class="text-white/60 text-xs font-mono">Draft: <span class="text-cyan-400">{instanceInfo.draftModel.split('/').pop()}</span>{#if instanceInfo.numDraftTokens}<span class="text-white/40"> ({instanceInfo.numDraftTokens}t)</span>{/if}</div>
-												{/if}
 												{#if instanceModelId && instanceModelId !== 'Unknown' && instanceModelId !== 'Unknown Model'}
 													<a
 														class="inline-flex items-center gap-1 text-[11px] text-white/60 hover:text-exo-yellow transition-colors mt-1"
@@ -1699,80 +1658,8 @@ function toggleInstanceDownloadDetails(nodeId: string): void {
 									{/each}
 								</div>
 							</div>
-
-							<!-- Draft Model (Speculative Decoding) -->
-							<div>
-								<div class="text-xs text-white/70 font-mono mb-2">Draft Model (Speculative):</div>
-								<div class="relative">
-									<button
-										onclick={() => { isDraftModelDropdownOpen = !isDraftModelDropdownOpen; draftModelDropdownSearch = ''; }}
-										class="w-full px-3 py-2 text-left text-sm font-mono border rounded transition-all duration-200 cursor-pointer flex items-center justify-between gap-2 {selectedDraftModel ? 'bg-transparent text-exo-yellow border-exo-yellow' : 'bg-transparent text-white/50 border-exo-medium-gray/50 hover:border-exo-yellow/50'}"
-									>
-										<span class="truncate">{selectedDraftModel ? selectedDraftModel.split('/').pop() : 'None'}</span>
-										<svg class="w-4 h-4 flex-shrink-0 transition-transform {isDraftModelDropdownOpen ? 'rotate-180' : ''}" fill="none" stroke="currentColor" viewBox="0 0 24 24">
-											<path stroke-linecap="round" stroke-linejoin="round" stroke-width="2" d="M19 9l-7 7-7-7" />
-										</svg>
-									</button>
-									{#if isDraftModelDropdownOpen}
-										<!-- svelte-ignore a11y_no_static_element_interactions -->
-										<div
-											class="fixed inset-0 z-40"
-											onclick={() => isDraftModelDropdownOpen = false}
-											onkeydown={(e) => e.key === 'Escape' && (isDraftModelDropdownOpen = false)}
-										></div>
-										<div class="absolute top-full left-0 right-0 mt-1 bg-exo-dark-gray border border-exo-medium-gray/50 rounded shadow-lg z-50 max-h-48 overflow-hidden flex flex-col">
-											<div class="p-2 border-b border-exo-medium-gray/30">
-												<input
-													type="text"
-													bind:value={draftModelDropdownSearch}
-													placeholder="Search models..."
-													class="w-full px-2 py-1.5 text-sm font-mono bg-transparent border border-exo-medium-gray/50 rounded text-white/90 placeholder:text-white/30 focus:outline-none focus:border-exo-yellow/50"
-												/>
-											</div>
-											<div class="overflow-y-auto max-h-36">
-												<!-- None option -->
-												<button
-													onclick={() => { selectedDraftModel = null; isDraftModelDropdownOpen = false; saveLaunchDefaults(); }}
-													class="w-full px-3 py-2 text-left text-sm font-mono tracking-wide transition-colors duration-100 flex items-center gap-2 {selectedDraftModel === null ? 'bg-transparent text-exo-yellow cursor-pointer' : 'text-white/80 hover:text-exo-yellow cursor-pointer'}"
-												>
-													<span>None</span>
-												</button>
-												{#each models.filter(m => (m.name ?? m.id).toLowerCase().includes(draftModelDropdownSearch.toLowerCase()) && m.id !== selectedModelId) as model}
-													{@const sizeGB = (model.storage_size_megabytes ?? 0) / 1024}
-													{@const modelHfId = model.hugging_face_id ?? model.id}
-													<button
-														onclick={() => { selectedDraftModel = modelHfId; isDraftModelDropdownOpen = false; saveLaunchDefaults(); }}
-														class="w-full px-3 py-2 text-left text-sm font-mono tracking-wide transition-colors duration-100 flex items-center justify-between gap-2 {selectedDraftModel === modelHfId ? 'bg-transparent text-exo-yellow cursor-pointer' : 'text-white/80 hover:text-exo-yellow cursor-pointer'}"
-													>
-														<span class="truncate">{model.name || model.id}</span>
-														<span class="flex-shrink-0 text-xs text-white/50">
-															{sizeGB >= 1 ? sizeGB.toFixed(0) : sizeGB.toFixed(1)}GB
-														</span>
-													</button>
-												{:else}
-													<div class="px-3 py-2 text-xs text-white/50 font-mono">No models found</div>
-												{/each}
-											</div>
-										</div>
-									{/if}
-								</div>
-							</div>
-							<!-- Draft Tokens (only show when draft model selected) -->
-							{#if selectedDraftModel}
-								<div class="flex items-center gap-2 mt-2">
-									<span class="text-xs text-white/50 font-mono">Tokens:</span>
-									<div class="flex items-center gap-1">
-										{#each [2, 3, 4, 5, 6] as n}
-											<button
-												onclick={() => { selectedNumDraftTokens = n; saveLaunchDefaults(); }}
-												class="w-6 h-6 text-xs font-mono rounded transition-all {selectedNumDraftTokens === n ? 'bg-exo-yellow/20 text-exo-yellow border border-exo-yellow/50' : 'text-white/50 hover:text-white/80 border border-transparent'}"
-											>{n}</button>
-										{/each}
-									</div>
-								</div>
-							{/if}
 						</div>
-
+						
 						<!-- Selected Model Preview -->
 						<div class="space-y-3">
 							{#if models.length === 0}

justfile

@@ -1,5 +1,3 @@
-export NIX_CONFIG := "extra-experimental-features = nix-command flakes"
-
 fmt:
     nix fmt
 

pyproject.toml

@@ -23,13 +23,13 @@ dependencies = [
     "tiktoken>=0.12.0", # required for kimi k2 tokenizer
     "hypercorn>=0.18.0",
     "openai-harmony>=0.0.8",
-    "httpx>=0.28.1",
 ]
 
 [project.scripts]
 exo-master = "exo.master.main:main"
 exo-worker = "exo.worker.main:main"
 exo = "exo.main:main"
+exo-rsh = "exo.rsh.client:main"
 
 # dependencies only required for development
 [dependency-groups]

rust/parts.nix

@@ -81,6 +81,20 @@
 
       config = {
         packages = {
+          # The system_custodian binary
+          system_custodian = craneLib.buildPackage (
+            commonArgs
+            // {
+              inherit cargoArtifacts;
+              cargoExtraArgs = "-p system_custodian";
+
+              meta = {
+                description = "System custodian daemon for exo";
+                mainProgram = "system_custodian";
+              };
+            }
+          );
+
           # Python bindings wheel via maturin
           exo_pyo3_bindings = craneLib.buildPackage (
             commonArgs

rust/system_custodian/Cargo.toml

@@ -0,0 +1,47 @@
+[package]
+name = "system_custodian"
+version = { workspace = true }
+edition = { workspace = true }
+publish = false
+
+[lib]
+doctest = false
+name = "system_custodian"
+path = "src/lib.rs"
+
+[[bin]]
+path = "src/bin/main.rs"
+name = "system_custodian"
+doc = false
+
+[lints]
+workspace = true
+
+[dependencies]
+# datastructures
+either = { workspace = true }
+
+# macro dependencies
+extend = { workspace = true }
+delegate = { workspace = true }
+impl-trait-for-tuples = { workspace = true }
+derive_more = { workspace = true }
+
+# async
+tokio = { workspace = true, features = ["full"] }
+futures = { workspace = true }
+futures-timer = { workspace = true }
+
+# utility dependencies
+util = { workspace = true }
+thiserror = { workspace = true }
+#internment = { workspace = true }
+#recursion = { workspace = true }
+#generativity = { workspace = true }
+#itertools = { workspace = true }
+tracing-subscriber = { version = "0.3.19", features = ["default", "env-filter"] }
+keccak-const = { workspace = true }
+
+# tracing/logging
+log = { workspace = true }
+

rust/system_custodian/src/bin/main.rs

@@ -0,0 +1,4 @@
+//! TODO: documentation
+//!
+
+fn main() {}

rust/system_custodian/src/lib.rs

@@ -0,0 +1,69 @@
+//! This crate defines the logic of, and ways to interact with, Exo's **_System Custodian_** daemon.
+//!
+//! The **_System Custodian_** daemon is supposed to be a long-living process that precedes the
+//! launch of the Exo application, and responsible for ensuring the system (configuration, settings,
+//! etc.) is in an appropriate state to facilitate the running of Exo application.
+//! The **_System Custodian_** daemon shall expose a [D-Bus](https://www.freedesktop.org/wiki/Software/dbus/)
+//! service which Exo application use to _control & query_ it.
+//!
+//! # Lifecycle
+//! When the Exo application starts, it will _wake_ the **_System Custodian_** daemon for the
+//! duration of its lifetime, and after it has terminated the daemon will go back to sleep. When
+//! the daemon wakes up, it will configure the system into a state suitable for the Exo Application;
+//! When the daemon goes to sleep, it will revert those changes as much as it can in case they were
+//! destructive to the user's pre-existing configurations.
+//!
+//! # Responsibilities
+//! TODO: these are purely on MacOS, but change to be more broad
+//! The **_System Custodian_** daemon is responsible for using System Configuration framework to
+//!  1. duplicate the current network set
+//!  2. modify existing services to turn on IPv6 if not there
+//!  3. remove any bridge services & add any missing services that AREN'T bridge
+//! TODO: In the future:
+//!  1. run a dummy AWDL service to [allow for macOS peer-to-peer wireless networking](https://yggdrasil-network.github.io/2019/08/19/awdl.html)
+//!  2. toggle some GPU/memory configurations to speed up GPU (ask Alex what those configurations are)
+//!  3. if we ever decide to provide our **own network interfaces** that abstract over some userland
+//!     logic, this would be the place to spin that up.
+//!
+//! Then it will watch the SCDynamicStore for:
+//!  1. all __actual__ network interfaces -> collect information on them e.g. their BSD name, MAC
+//!     address, MTU, IPv6 addresses, etc. -> and set up watchers/notifiers to inform the DBus
+//!     interface of any changes
+//!  2. watch for any __undesirable__ changes to configuration and revert it
+//!
+//! It should somehow (probably through system sockets and/or BSD interface) trigger IPv6 NDP on
+//! each of the interfaces & also listen to/query for any changes on the OS routing cache??
+//! Basically emulate the `ping6 ff02::1%enX` and `ndp -an` commands BUT BETTER!!!
+//!  1. all that info should coalesce back to the overall state colleted -> should be queryable
+//!     over D-Bus
+//! TODO:
+//!  1. we might potentially add to this step a handshake of some kind...? To ensure that we can
+//!     ACTUALLY communicate with that machine over that link over e.g. TCP, UDP, etc. Will the
+//!     handshake require to know Node ID? Will the handshake require heartbeats? Who knows...
+//!  2. if we ever decide to write proprietary L2/L3 protocols for quicker communication,
+//!     e.g. [AF_NDRV](https://www.zerotier.com/blog/how-zerotier-eliminated-kernel-extensions-on-macos/)
+//!     for raw ethernet frame communication, or even a [custom thunderbolt PCIe driver](https://developer.apple.com/documentation/pcidriverkit/creating-custom-pcie-drivers-for-thunderbolt-devices),
+//!     then this would be the place to carry out discovery and propper handshakes with devices
+//!     on the other end of the link.
+//!
+
+// enable Rust-unstable features for convenience
+#![feature(trait_alias)]
+#![feature(stmt_expr_attributes)]
+#![feature(type_alias_impl_trait)]
+#![feature(specialization)]
+#![feature(unboxed_closures)]
+#![feature(const_trait_impl)]
+#![feature(fn_traits)]
+
+pub(crate) mod private {
+    // sealed traits support
+    pub trait Sealed {}
+    impl<T: ?Sized> Sealed for T {}
+}
+
+/// Namespace for all the type/trait aliases used by this crate.
+pub(crate) mod alias {}
+
+/// Namespace for crate-wide extension traits/methods
+pub(crate) mod ext {}

src/exo/main.py

@@ -205,14 +205,6 @@ def main():
     logger.info("Starting EXO")
     logger.info(f"EXO_LIBP2P_NAMESPACE: {os.getenv('EXO_LIBP2P_NAMESPACE')}")
 
-    # Set FAST_SYNCH override env var for runner subprocesses
-    if args.fast_synch is True:
-        os.environ["EXO_FAST_SYNCH"] = "on"
-        logger.info("FAST_SYNCH forced ON")
-    elif args.fast_synch is False:
-        os.environ["EXO_FAST_SYNCH"] = "off"
-        logger.info("FAST_SYNCH forced OFF")
-
     node = anyio.run(Node.create, args)
     anyio.run(node.run)
     logger.info("EXO Shutdown complete")
@@ -226,7 +218,6 @@ class Args(CamelCaseModel):
     api_port: PositiveInt = 52415
     tb_only: bool = False
     no_worker: bool = False
-    fast_synch: bool | None = None  # None = auto, True = force on, False = force off
 
     @classmethod
     def parse(cls) -> Self:
@@ -268,20 +259,6 @@ class Args(CamelCaseModel):
             "--no-worker",
             action="store_true",
         )
-        fast_synch_group = parser.add_mutually_exclusive_group()
-        fast_synch_group.add_argument(
-            "--fast-synch",
-            action="store_true",
-            dest="fast_synch",
-            default=None,
-            help="Force MLX FAST_SYNCH on (for JACCL backend)",
-        )
-        fast_synch_group.add_argument(
-            "--no-fast-synch",
-            action="store_false",
-            dest="fast_synch",
-            help="Force MLX FAST_SYNCH off",
-        )
 
         args = parser.parse_args()
         return cls(**vars(args))  # pyright: ignore[reportAny] - We are intentionally validating here, we can't do it statically

src/exo/master/api.py

@@ -1,19 +1,27 @@
+import asyncio
+import os
 import time
 from collections.abc import AsyncGenerator
-from http import HTTPStatus
-from typing import cast
+from typing import Any, Optional, cast
 
 import anyio
-from anyio import BrokenResourceError, create_task_group
+from anyio import create_task_group
 from anyio.abc import TaskGroup
-from fastapi import FastAPI, HTTPException, Request
+from fastapi import FastAPI, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
-from fastapi.responses import JSONResponse, StreamingResponse
+from fastapi.responses import StreamingResponse
 from fastapi.staticfiles import StaticFiles
 from hypercorn.asyncio import serve  # pyright: ignore[reportUnknownVariableType]
 from hypercorn.config import Config
 from hypercorn.typing import ASGIFramework
 from loguru import logger
+from openai_harmony import (  # pyright: ignore[reportMissingTypeStubs]
+    HarmonyEncodingName,
+    Role,
+    StreamableParser,
+    load_harmony_encoding,
+)
+from pydantic import BaseModel
 
 from exo.master.placement import place_instance as get_instance_placements
 from exo.shared.apply import apply
@@ -30,8 +38,6 @@ from exo.shared.types.api import (
     CreateInstanceParams,
     CreateInstanceResponse,
     DeleteInstanceResponse,
-    ErrorInfo,
-    ErrorResponse,
     FinishReason,
     GenerationStats,
     ModelList,
@@ -48,27 +54,47 @@ from exo.shared.types.commands import (
     CreateInstance,
     DeleteInstance,
     ForwarderCommand,
+    LaunchFLASH,
     PlaceInstance,
+    StopFLASH,
     TaskFinished,
 )
 from exo.shared.types.common import CommandId, NodeId, SessionId
-from exo.shared.types.events import (
-    ChunkGenerated,
-    Event,
-    ForwarderEvent,
-    IndexedEvent,
-)
+from exo.shared.types.events import ChunkGenerated, Event, ForwarderEvent, IndexedEvent
 from exo.shared.types.memory import Memory
 from exo.shared.types.models import ModelId, ModelMetadata
 from exo.shared.types.state import State
 from exo.shared.types.tasks import ChatCompletionTaskParams
-from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
+from exo.shared.types.worker.instances import (
+    FLASHInstance,
+    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
 
+encoding = load_harmony_encoding(HarmonyEncodingName.HARMONY_GPT_OSS)
+
+
+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 chunk_to_response(
     chunk: TokenChunk, command_id: CommandId
@@ -123,7 +149,6 @@ class API:
         self.paused_ev: anyio.Event = anyio.Event()
 
         self.app = FastAPI()
-        self._setup_exception_handlers()
         self._setup_cors()
         self._setup_routes()
 
@@ -154,20 +179,6 @@ class API:
         self.paused_ev.set()
         self.paused_ev = anyio.Event()
 
-    def _setup_exception_handlers(self) -> None:
-        @self.app.exception_handler(HTTPException)
-        async def http_exception_handler(  # pyright: ignore[reportUnusedFunction]
-            _: Request, exc: HTTPException
-        ) -> JSONResponse:
-            err = ErrorResponse(
-                error=ErrorInfo(
-                    message=exc.detail,
-                    type=HTTPStatus(exc.status_code).phrase,
-                    code=exc.status_code,
-                )
-            )
-            return JSONResponse(err.model_dump(), status_code=exc.status_code)
-
     def _setup_cors(self) -> None:
         self.app.add_middleware(
             CORSMiddleware,
@@ -193,6 +204,12 @@ class API:
         self.app.post("/bench/chat/completions")(self.bench_chat_completions)
         self.app.get("/state")(lambda: self.state)
         self.app.get("/events")(lambda: self._event_log)
+        # FLASH simulation endpoints
+        self.app.post("/flash/launch")(self.launch_flash)
+        self.app.delete("/flash/{instance_id}")(self.stop_flash)
+        self.app.get("/flash/instances")(self.list_flash_instances)
+        # Remote execution endpoint (used by exo-rsh for MPI)
+        self.app.post("/execute")(self.execute)
 
     async def place_instance(self, payload: PlaceInstanceParams):
         command = PlaceInstance(
@@ -200,8 +217,6 @@ class API:
             sharding=payload.sharding,
             instance_meta=payload.instance_meta,
             min_nodes=payload.min_nodes,
-            draft_model=payload.draft_model,
-            num_draft_tokens=payload.num_draft_tokens,
         )
         await self._send(command)
 
@@ -398,8 +413,35 @@ class API:
             instance_id=instance_id,
         )
 
+    async def _process_gpt_oss(self, token_chunks: Receiver[TokenChunk]):
+        stream = StreamableParser(encoding, role=Role.ASSISTANT)
+        thinking = False
+
+        async for chunk in token_chunks:
+            stream.process(chunk.token_id)
+
+            delta = stream.last_content_delta
+            ch = stream.current_channel
+
+            if ch == "analysis" and not thinking:
+                thinking = True
+                yield chunk.model_copy(update={"text": "<think>"})
+
+            if ch != "analysis" and thinking:
+                thinking = False
+                yield chunk.model_copy(update={"text": "</think>"})
+
+            if delta:
+                yield chunk.model_copy(update={"text": delta})
+
+            if chunk.finish_reason is not None:
+                if thinking:
+                    yield chunk.model_copy(update={"text": "</think>"})
+                yield chunk
+                break
+
     async def _chat_chunk_stream(
-        self, command_id: CommandId
+        self, command_id: CommandId, parse_gpt_oss: bool
     ) -> AsyncGenerator[TokenChunk, None]:
         """Yield `TokenChunk`s for a given command until completion."""
 
@@ -407,10 +449,16 @@ class API:
             self._chat_completion_queues[command_id], recv = channel[TokenChunk]()
 
             with recv as token_chunks:
-                async for chunk in token_chunks:
-                    yield chunk
-                    if chunk.finish_reason is not None:
-                        break
+                if parse_gpt_oss:
+                    async for chunk in self._process_gpt_oss(token_chunks):
+                        yield chunk
+                        if chunk.finish_reason is not None:
+                            break
+                else:
+                    async for chunk in token_chunks:
+                        yield chunk
+                        if chunk.finish_reason is not None:
+                            break
 
         except anyio.get_cancelled_exc_class():
             # TODO: TaskCancelled
@@ -426,23 +474,11 @@ class API:
             del self._chat_completion_queues[command_id]
 
     async def _generate_chat_stream(
-        self, command_id: CommandId
+        self, command_id: CommandId, parse_gpt_oss: bool
     ) -> AsyncGenerator[str, None]:
         """Generate chat completion stream as JSON strings."""
 
-        async for chunk in self._chat_chunk_stream(command_id):
-            if chunk.finish_reason == "error":
-                error_response = ErrorResponse(
-                    error=ErrorInfo(
-                        message=chunk.error_message or "Internal server error",
-                        type="InternalServerError",
-                        code=500,
-                    )
-                )
-                yield f"data: {error_response.model_dump_json()}\n\n"
-                yield "data: [DONE]\n\n"
-                return
-
+        async for chunk in self._chat_chunk_stream(command_id, parse_gpt_oss):
             chunk_response: ChatCompletionResponse = chunk_to_response(
                 chunk, command_id
             )
@@ -454,7 +490,7 @@ class API:
                 yield "data: [DONE]\n\n"
 
     async def _collect_chat_completion(
-        self, command_id: CommandId
+        self, command_id: CommandId, parse_gpt_oss: bool
     ) -> ChatCompletionResponse:
         """Collect all token chunks for a chat completion and return a single response."""
 
@@ -462,13 +498,7 @@ class API:
         model: str | None = None
         finish_reason: FinishReason | None = None
 
-        async for chunk in self._chat_chunk_stream(command_id):
-            if chunk.finish_reason == "error":
-                raise HTTPException(
-                    status_code=500,
-                    detail=chunk.error_message or "Internal server error",
-                )
-
+        async for chunk in self._chat_chunk_stream(command_id, parse_gpt_oss):
             if model is None:
                 model = chunk.model
 
@@ -497,7 +527,7 @@ class API:
         )
 
     async def _collect_chat_completion_with_stats(
-        self, command_id: CommandId
+        self, command_id: CommandId, parse_gpt_oss: bool
     ) -> BenchChatCompletionResponse:
         text_parts: list[str] = []
         model: str | None = None
@@ -505,13 +535,7 @@ class API:
 
         stats: GenerationStats | None = None
 
-        async for chunk in self._chat_chunk_stream(command_id):
-            if chunk.finish_reason == "error":
-                raise HTTPException(
-                    status_code=500,
-                    detail=chunk.error_message or "Internal server error",
-                )
-
+        async for chunk in self._chat_chunk_stream(command_id, parse_gpt_oss):
             if model is None:
                 model = chunk.model
 
@@ -552,6 +576,8 @@ class API:
         """Handle chat completions, supporting both streaming and non-streaming responses."""
         model_meta = await resolve_model_meta(payload.model)
         payload.model = model_meta.model_id
+        parse_gpt_oss = "gpt-oss" in model_meta.model_id.lower()
+        logger.info(f"{parse_gpt_oss=}")
 
         if not any(
             instance.shard_assignments.model_id == payload.model
@@ -568,16 +594,17 @@ class API:
         await self._send(command)
         if payload.stream:
             return StreamingResponse(
-                self._generate_chat_stream(command.command_id),
+                self._generate_chat_stream(command.command_id, parse_gpt_oss),
                 media_type="text/event-stream",
             )
 
-        return await self._collect_chat_completion(command.command_id)
+        return await self._collect_chat_completion(command.command_id, parse_gpt_oss)
 
     async def bench_chat_completions(
         self, payload: BenchChatCompletionTaskParams
     ) -> BenchChatCompletionResponse:
         model_meta = await resolve_model_meta(payload.model)
+        parse_gpt_oss = "gpt-oss" in model_meta.model_id.lower()
         payload.model = model_meta.model_id
 
         if not any(
@@ -594,7 +621,10 @@ class API:
         command = ChatCompletion(request_params=payload)
         await self._send(command)
 
-        response = await self._collect_chat_completion_with_stats(command.command_id)
+        response = await self._collect_chat_completion_with_stats(
+            command.command_id,
+            parse_gpt_oss,
+        )
         return response
 
     def _calculate_total_available_memory(self) -> Memory:
@@ -624,6 +654,145 @@ class API:
             ]
         )
 
+    async def launch_flash(
+        self,
+        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:
+            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 self._send(command)
+
+        return {
+            "message": "FLASH launch command received",
+            "command_id": str(command.command_id),
+            "simulation_name": simulation_name,
+        }
+
+    async def stop_flash(self, instance_id: InstanceId) -> dict[str, str]:
+        """Stop a running FLASH simulation."""
+        if instance_id not in self.state.instances:
+            raise HTTPException(status_code=404, detail="Instance not found")
+
+        instance = self.state.instances[instance_id]
+        if not isinstance(instance, FLASHInstance):
+            raise HTTPException(
+                status_code=400, detail="Instance is not a FLASH simulation"
+            )
+
+        command = StopFLASH(instance_id=instance_id)
+        await self._send(command)
+
+        return {
+            "message": "Stop command received",
+            "command_id": str(command.command_id),
+            "instance_id": str(instance_id),
+        }
+
+    async def list_flash_instances(self) -> list[dict[str, Any]]:
+        """List all FLASH simulation instances."""
+        flash_instances: list[dict[str, Any]] = []
+        for instance_id, instance in self.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 = self.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
+
+    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 ";|&$`")
+
+            if needs_shell:
+                process = await asyncio.create_subprocess_shell(
+                    cmd_str,
+                    stdout=asyncio.subprocess.PIPE,
+                    stderr=asyncio.subprocess.PIPE,
+                    cwd=request.cwd,
+                    env=env,
+                )
+            else:
+                process = await asyncio.create_subprocess_exec(
+                    *request.command,
+                    stdout=asyncio.subprocess.PIPE,
+                    stderr=asyncio.subprocess.PIPE,
+                    cwd=request.cwd,
+                    env=env,
+                )
+
+            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]}",
+            )
+        except Exception as e:
+            logger.error(f"Execution error: {e}")
+            return ExecuteResponse(
+                exit_code=1,
+                stdout="",
+                stderr=str(e),
+            )
+
     async def run(self):
         cfg = Config()
         cfg.bind = f"0.0.0.0:{self.port}"
@@ -656,14 +825,14 @@ 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 (
+                        isinstance(event, ChunkGenerated)
+                        and 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)
+                        await self._chat_completion_queues[event.command_id].send(
+                            event.chunk
+                        )
 
     async def _pause_on_new_election(self):
         with self.election_receiver as ems:

src/exo/master/main.py

@@ -8,6 +8,7 @@ from exo.master.placement import (
     add_instance_to_placements,
     delete_instance,
     get_transition_events,
+    place_flash_instance,
     place_instance,
 )
 from exo.shared.apply import apply
@@ -16,8 +17,10 @@ from exo.shared.types.commands import (
     CreateInstance,
     DeleteInstance,
     ForwarderCommand,
+    LaunchFLASH,
     PlaceInstance,
     RequestEventLog,
+    StopFLASH,
     TaskFinished,
     TestCommand,
 )
@@ -173,6 +176,26 @@ class Master:
                                 self.state.instances, placement
                             )
                             generated_events.extend(transition_events)
+                        case LaunchFLASH():
+                            placement = place_flash_instance(
+                                command,
+                                self.state.topology,
+                                self.state.instances,
+                            )
+                            transition_events = get_transition_events(
+                                self.state.instances, placement
+                            )
+                            generated_events.extend(transition_events)
+                        case StopFLASH():
+                            # Reuse delete_instance logic to stop FLASH simulation
+                            placement = delete_instance(
+                                DeleteInstance(instance_id=command.instance_id),
+                                self.state.instances,
+                            )
+                            transition_events = get_transition_events(
+                                self.state.instances, placement
+                            )
+                            generated_events.extend(transition_events)
                         case TaskFinished():
                             generated_events.append(
                                 TaskDeleted(

src/exo/master/placement.py

@@ -17,20 +17,24 @@ from exo.shared.topology import Topology
 from exo.shared.types.commands import (
     CreateInstance,
     DeleteInstance,
+    LaunchFLASH,
     PlaceInstance,
 )
+from exo.shared.types.common import Host, NodeId
 from exo.shared.types.events import Event, InstanceCreated, InstanceDeleted
 from exo.shared.types.memory import Memory
-from exo.shared.types.models import ModelId
+from exo.shared.types.models import ModelId, ModelMetadata
 from exo.shared.types.topology import NodeInfo
 from exo.shared.types.worker.instances import (
+    FLASHInstance,
     Instance,
     InstanceId,
     InstanceMeta,
     MlxJacclInstance,
     MlxRingInstance,
 )
-from exo.shared.types.worker.shards import Sharding
+from exo.shared.types.worker.runners import RunnerId, ShardAssignments
+from exo.shared.types.worker.shards import PipelineShardMetadata, Sharding
 
 
 def random_ephemeral_port() -> int:
@@ -151,8 +155,6 @@ def place_instance(
                 shard_assignments=shard_assignments,
                 ibv_devices=mlx_ibv_devices,
                 jaccl_coordinators=mlx_jaccl_coordinators,
-                draft_model=command.draft_model,
-                num_draft_tokens=command.num_draft_tokens,
             )
         case InstanceMeta.MlxRing:
             ephemeral_port = random_ephemeral_port()
@@ -166,9 +168,10 @@ def place_instance(
                 shard_assignments=shard_assignments,
                 hosts_by_node=hosts_by_node,
                 ephemeral_port=ephemeral_port,
-                draft_model=command.draft_model,
-                num_draft_tokens=command.num_draft_tokens,
             )
+        case InstanceMeta.FLASH:
+            # FLASH instances are handled by place_flash_instance()
+            raise ValueError("FLASH instances should use place_flash_instance()")
 
     return target_instances
 
@@ -184,6 +187,148 @@ def delete_instance(
     raise ValueError(f"Instance {command.instance_id} not found")
 
 
+def place_flash_instance(
+    command: LaunchFLASH,
+    topology: Topology,
+    current_instances: Mapping[InstanceId, Instance],
+) -> dict[InstanceId, Instance]:
+    """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(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 ModelMetadata for FLASH (required by ShardMetadata interface)
+    flash_model_meta = ModelMetadata(
+        model_id=ModelId(command.simulation_name),
+        pretty_name=f"FLASH: {command.simulation_name}",
+        storage_size=Memory(in_bytes=0),
+        n_layers=1,
+        hidden_size=1,
+        supports_tensor=False,
+    )
+
+    for i, node_info in enumerate(selected_nodes):
+        runner_id = RunnerId()
+        node_to_runner[node_info.node_id] = runner_id
+        runner_to_shard[runner_id] = PipelineShardMetadata(
+            device_rank=i,
+            world_size=len(selected_nodes),
+            model_meta=flash_model_meta,
+            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_info in enumerate(selected_nodes):
+            if i < len(explicit_hosts):
+                hosts_by_node[node_info.node_id] = [Host(ip=explicit_hosts[i], port=0)]
+                logger.info(
+                    f"FLASH placement: node {node_info.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_info.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:
+        # Try to get IPs from topology edges
+        for node_info in selected_nodes:
+            node_hosts: list[Host] = []
+
+            # Get IP from outgoing edges (connections to other nodes via mDNS discovery)
+            for _, edge_data in topology.out_edges(node_info.node_id):
+                if hasattr(edge_data, "send_back_multiaddr"):
+                    # Extract IP from multiaddr like /ip4/192.168.1.100/tcp/52415
+                    multiaddr = str(edge_data.send_back_multiaddr)
+                    if "/ip4/" in multiaddr:
+                        parts = multiaddr.split("/")
+                        try:
+                            ip_idx = parts.index("ip4") + 1
+                            ip = parts[ip_idx]
+                            # Skip link-local and localhost addresses
+                            if not ip.startswith("169.254.") and not ip.startswith(
+                                "127."
+                            ):
+                                node_hosts.append(Host(ip=ip, port=0))
+                                break
+                        except (ValueError, IndexError):
+                            pass
+
+            # Last resort: use localhost (will only work for single-node)
+            if not node_hosts:
+                logger.warning(
+                    f"Could not determine IP for node {node_info.node_id}, using localhost"
+                )
+                node_hosts.append(Host(ip="127.0.0.1", port=0))
+
+            hosts_by_node[node_info.node_id] = node_hosts
+
+    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
+
+
 def get_transition_events(
     current_instances: Mapping[InstanceId, Instance],
     target_instances: Mapping[InstanceId, Instance],

src/exo/master/tests/test_api_error_handling.py

@@ -1,107 +0,0 @@
-# pyright: reportUnusedFunction=false, reportAny=false
-from typing import Any, get_args
-
-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 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."""
-    from exo.master.api import API
-
-    app = FastAPI()
-
-    # Setup exception handler
-    api = object.__new__(API)
-    api.app = app
-    api._setup_exception_handlers()  # pyright: ignore[reportPrivateUsage]
-
-    # Add test routes that raise HTTPException
-    @app.get("/test-error")
-    async def _test_error() -> None:
-        raise HTTPException(status_code=500, detail="Test error message")
-
-    @app.get("/test-not-found")
-    async def _test_not_found() -> None:
-        raise HTTPException(status_code=404, detail="Resource not found")
-
-    client = TestClient(app)
-
-    # Test 500 error
-    response = client.get("/test-error")
-    assert response.status_code == 500
-    data: dict[str, Any] = response.json()
-    assert "error" in data
-    assert data["error"]["message"] == "Test error message"
-    assert data["error"]["type"] == "Internal Server Error"
-    assert data["error"]["code"] == 500
-
-    # Test 404 error
-    response = client.get("/test-not-found")
-    assert response.status_code == 404
-    data = response.json()
-    assert "error" in data
-    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

src/exo/rsh/__init__.py

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

@@ -0,0 +1,101 @@
+#!/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/logging.py

@@ -29,11 +29,6 @@ class _InterceptHandler(logging.Handler):
 
 def logger_setup(log_file: Path | None, verbosity: int = 0):
     """Set up logging for this process - formatting, file handles, verbosity and output"""
-
-    logging.getLogger("exo_pyo3_bindings").setLevel(logging.WARNING)
-    logging.getLogger("httpx").setLevel(logging.WARNING)
-    logging.getLogger("httpcore").setLevel(logging.WARNING)
-
     logger.remove()
 
     # replace all stdlib loggers with _InterceptHandlers that log to loguru

src/exo/shared/models/model_cards.py

@@ -14,6 +14,32 @@ class ModelCard(CamelCaseModel):
 
 MODEL_CARDS: dict[str, ModelCard] = {
     # deepseek v3
+    # "deepseek-v3-0324:4bit": ModelCard(
+    #     short_id="deepseek-v3-0324:4bit",
+    #     model_id="mlx-community/DeepSeek-V3-0324-4bit",
+    #     name="DeepSeek V3 0324 (4-bit)",
+    #     description="""DeepSeek V3 is a large language model trained on the DeepSeek V3 dataset.""",
+    #     tags=[],
+    #     metadata=ModelMetadata(
+    #         model_id=ModelId("mlx-community/DeepSeek-V3-0324-4bit"),
+    #         pretty_name="DeepSeek V3 0324 (4-bit)",
+    #         storage_size=Memory.from_kb(409706307),
+    #         n_layers=61,
+    #     ),
+    # ),
+    # "deepseek-v3-0324": ModelCard(
+    #     short_id="deepseek-v3-0324",
+    #     model_id="mlx-community/DeepSeek-v3-0324-8bit",
+    #     name="DeepSeek V3 0324 (8-bit)",
+    #     description="""DeepSeek V3 is a large language model trained on the DeepSeek V3 dataset.""",
+    #     tags=[],
+    #     metadata=ModelMetadata(
+    #         model_id=ModelId("mlx-community/DeepSeek-v3-0324-8bit"),
+    #         pretty_name="DeepSeek V3 0324 (8-bit)",
+    #         storage_size=Memory.from_kb(754706307),
+    #         n_layers=61,
+    #     ),
+    # ),
     "deepseek-v3.1-4bit": ModelCard(
         short_id="deepseek-v3.1-4bit",
         model_id=ModelId("mlx-community/DeepSeek-V3.1-4bit"),
@@ -44,6 +70,65 @@ MODEL_CARDS: dict[str, ModelCard] = {
             supports_tensor=True,
         ),
     ),
+    # "deepseek-v3.2": ModelCard(
+    #     short_id="deepseek-v3.2",
+    #     model_id=ModelId("mlx-community/DeepSeek-V3.2-8bit"),
+    #     name="DeepSeek V3.2 (8-bit)",
+    #     description="""DeepSeek V3.2 is a large language model trained on the DeepSeek V3.2 dataset.""",
+    #     tags=[],
+    #     metadata=ModelMetadata(
+    #         model_id=ModelId("mlx-community/DeepSeek-V3.2-8bit"),
+    #         pretty_name="DeepSeek V3.2 (8-bit)",
+    #         storage_size=Memory.from_kb(754706307),
+    #         n_layers=61,
+    #         hidden_size=7168,
+    #         supports_tensor=True,
+    #     ),
+    # ),
+    # "deepseek-v3.2-4bit": ModelCard(
+    #     short_id="deepseek-v3.2-4bit",
+    #     model_id=ModelId("mlx-community/DeepSeek-V3.2-4bit"),
+    #     name="DeepSeek V3.2 (4-bit)",
+    #     description="""DeepSeek V3.2 is a large language model trained on the DeepSeek V3.2 dataset.""",
+    #     tags=[],
+    #     metadata=ModelMetadata(
+    #         model_id=ModelId("mlx-community/DeepSeek-V3.2-4bit"),
+    #         pretty_name="DeepSeek V3.2 (4-bit)",
+    #         storage_size=Memory.from_kb(754706307 // 2),  # TODO !!!!!
+    #         n_layers=61,
+    #         hidden_size=7168,
+    #         supports_tensor=True,
+    #     ),
+    # ),
+    # deepseek r1
+    # "deepseek-r1-0528-4bit": ModelCard(
+    #     short_id="deepseek-r1-0528-4bit",
+    #     model_id="mlx-community/DeepSeek-R1-0528-4bit",
+    #     name="DeepSeek-R1-0528 (4-bit)",
+    #     description="""DeepSeek R1 is a large language model trained on the DeepSeek R1 dataset.""",
+    #     tags=[],
+    #     metadata=ModelMetadata(
+    #         model_id=ModelId("mlx-community/DeepSeek-R1-0528-4bit"),
+    #         pretty_name="DeepSeek R1 671B (4-bit)",
+    #         storage_size=Memory.from_kb(409706307),
+    #         n_layers=61,
+    #         hidden_size=7168,
+    #     ),
+    # ),
+    # "deepseek-r1-0528": ModelCard(
+    #     short_id="deepseek-r1-0528",
+    #     model_id="mlx-community/DeepSeek-R1-0528-8bit",
+    #     name="DeepSeek-R1-0528 (8-bit)",
+    #     description="""DeepSeek R1 is a large language model trained on the DeepSeek R1 dataset.""",
+    #     tags=[],
+    #     metadata=ModelMetadata(
+    #         model_id=ModelId("mlx-community/DeepSeek-R1-0528-8bit"),
+    #         pretty_name="DeepSeek R1 671B (8-bit)",
+    #         storage_size=Memory.from_bytes(754998771712),
+    #         n_layers=61,
+    # .       hidden_size=7168,
+    #     ),
+    # ),
     # kimi k2
     "kimi-k2-instruct-4bit": ModelCard(
         short_id="kimi-k2-instruct-4bit",
@@ -425,24 +510,23 @@ MODEL_CARDS: dict[str, ModelCard] = {
             supports_tensor=True,
         ),
     ),
-    "gpt-oss-20b-MXFP4-Q8": ModelCard(
-        short_id="gpt-oss-20b-MXFP4-Q8",
-        model_id=ModelId("mlx-community/gpt-oss-20b-MXFP4-Q8"),
-        name="GPT-OSS 20B (MXFP4-Q8, MLX)",
-        description="""OpenAI's GPT-OSS 20B is a medium-sized MoE model for lower-latency and local or specialized use cases; this variant is a 4-bit MLX conversion for Apple Silicon.""",
+    "gpt-oss-20b-4bit": ModelCard(
+        short_id="gpt-oss-20b-4bit",
+        model_id=ModelId("mlx-community/gpt-oss-20b-MXFP4-Q4"),
+        name="GPT-OSS 20B (MXFP4-Q4, MLX)",
+        description="""OpenAI's GPT-OSS 20B is a medium-sized MoE model for lower-latency and local or specialized use cases; this MLX variant uses MXFP4 4-bit quantization.""",
         tags=[],
         metadata=ModelMetadata(
-            model_id=ModelId("mlx-community/gpt-oss-20b-MXFP4-Q8"),
-            pretty_name="GPT-OSS 20B (MXFP4-Q8, MLX)",
+            model_id=ModelId("mlx-community/gpt-oss-20b-MXFP4-Q4"),
+            pretty_name="GPT-OSS 20B (MXFP4-Q4, MLX)",
             storage_size=Memory.from_kb(11_744_051),
             n_layers=24,
             hidden_size=2880,
             supports_tensor=True,
         ),
     ),
-    # glm 4.5
+    # Needs to be quantized g32 or g16.
     "glm-4.5-air-8bit": ModelCard(
-        # Needs to be quantized g32 or g16 to work with tensor parallel
         short_id="glm-4.5-air-8bit",
         model_id=ModelId("mlx-community/GLM-4.5-Air-8bit"),
         name="GLM 4.5 Air 8bit",
@@ -472,7 +556,6 @@ MODEL_CARDS: dict[str, ModelCard] = {
             supports_tensor=True,
         ),
     ),
-    # glm 4.7
     "glm-4.7-4bit": ModelCard(
         short_id="glm-4.7-4bit",
         model_id=ModelId("mlx-community/GLM-4.7-4bit"),
@@ -518,7 +601,6 @@ MODEL_CARDS: dict[str, ModelCard] = {
             supports_tensor=True,
         ),
     ),
-    # minimax-m2
     "minimax-m2.1-8bit": ModelCard(
         short_id="minimax-m2.1-8bit",
         model_id=ModelId("mlx-community/MiniMax-M2.1-8bit"),
@@ -549,4 +631,19 @@ MODEL_CARDS: dict[str, ModelCard] = {
             supports_tensor=True,
         ),
     ),
+    # "devstral-2-123b-instruct-2512-8bit": ModelCard(
+    #     short_id="devstral-2-123b-instruct-2512-8bit",
+    #     model_id=ModelId("mlx-community/Devstral-2-123B-Instruct-2512-8bit"),
+    #     name="Devstral 2 123B Instruct 2512 (8-bit, MLX)",
+    #     description="""Mistral AI's Devstral 2 123B Instruct (2512) is an agentic coding model.""",
+    #     tags=[],
+    #     metadata=ModelMetadata(
+    #         model_id=ModelId("mlx-community/Devstral-2-123B-Instruct-2512-8bit"),
+    #         pretty_name="Devstral 2 123B Instruct 2512 (8-bit, MLX)",
+    #         storage_size=Memory.from_kb(133_000_000),
+    #         n_layers=88,
+    #         hidden_size=12288,
+    #         supports_tensor=True,
+    #     ),
+    # ),
 }

src/exo/shared/types/api.py

@@ -11,21 +11,10 @@ from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
 from exo.shared.types.worker.shards import Sharding
 
 FinishReason = Literal[
-    "stop", "length", "tool_calls", "content_filter", "function_call", "error"
+    "stop", "length", "tool_calls", "content_filter", "function_call"
 ]
 
 
-class ErrorInfo(BaseModel):
-    message: str
-    type: str
-    param: str | None = None
-    code: int
-
-
-class ErrorResponse(BaseModel):
-    error: ErrorInfo
-
-
 class ModelListModel(BaseModel):
     id: str
     object: str = "model"
@@ -161,8 +150,6 @@ class ChatCompletionTaskParams(BaseModel):
     tool_choice: str | dict[str, Any] | None = None
     parallel_tool_calls: bool | None = None
     user: str | None = None
-    # Speculative decoding: tokens to draft per iteration (if instance has draft model)
-    num_draft_tokens: int = 3
 
 
 class BenchChatCompletionTaskParams(ChatCompletionTaskParams):
@@ -174,8 +161,6 @@ class PlaceInstanceParams(BaseModel):
     sharding: Sharding = Sharding.Pipeline
     instance_meta: InstanceMeta = InstanceMeta.MlxRing
     min_nodes: int = 1
-    draft_model: ModelId | None = None  # For speculative decoding
-    num_draft_tokens: int = 4  # Tokens to draft per iteration
 
     @field_validator("sharding", "instance_meta", mode="plain")
     @classmethod

src/exo/shared/types/chunks.py

@@ -22,7 +22,6 @@ class TokenChunk(BaseChunk):
     token_id: int
     finish_reason: FinishReason | None = None
     stats: GenerationStats | None = None
-    error_message: str | None = None
 
 
 class ImageChunk(BaseChunk):

src/exo/shared/types/commands.py

@@ -2,7 +2,7 @@ from pydantic import Field
 
 from exo.shared.types.api import ChatCompletionTaskParams
 from exo.shared.types.common import CommandId, NodeId
-from exo.shared.types.models import ModelId, ModelMetadata
+from exo.shared.types.models import ModelMetadata
 from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
 from exo.shared.types.worker.shards import Sharding
 from exo.utils.pydantic_ext import CamelCaseModel, TaggedModel
@@ -25,8 +25,6 @@ class PlaceInstance(BaseCommand):
     sharding: Sharding
     instance_meta: InstanceMeta
     min_nodes: int
-    draft_model: ModelId | None = None  # For speculative decoding
-    num_draft_tokens: int = 4  # Tokens to draft per iteration
 
 
 class CreateInstance(BaseCommand):
@@ -37,6 +35,26 @@ class DeleteInstance(BaseCommand):
     instance_id: InstanceId
 
 
+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 = ""
+
+
+class StopFLASH(BaseCommand):
+    """Command to stop a running FLASH simulation."""
+
+    instance_id: InstanceId
+
+
 class TaskFinished(BaseCommand):
     finished_command_id: CommandId
 
@@ -52,6 +70,8 @@ Command = (
     | PlaceInstance
     | CreateInstance
     | DeleteInstance
+    | LaunchFLASH
+    | StopFLASH
     | TaskFinished
 )
 

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

@@ -3,7 +3,6 @@ from enum import Enum
 from pydantic import model_validator
 
 from exo.shared.types.common import Host, Id, NodeId
-from exo.shared.types.models import ModelId
 from exo.shared.types.worker.runners import RunnerId, ShardAssignments, ShardMetadata
 from exo.utils.pydantic_ext import CamelCaseModel, TaggedModel
 
@@ -15,13 +14,12 @@ class InstanceId(Id):
 class InstanceMeta(str, Enum):
     MlxRing = "MlxRing"
     MlxJaccl = "MlxJaccl"
+    FLASH = "FLASH"
 
 
 class BaseInstance(TaggedModel):
     instance_id: InstanceId
     shard_assignments: ShardAssignments
-    draft_model: ModelId | None = None  # For speculative decoding (rank 0 only)
-    num_draft_tokens: int = 4  # Tokens to draft per iteration (when draft_model is set)
 
     def shard(self, runner_id: RunnerId) -> ShardMetadata | None:
         return self.shard_assignments.runner_to_shard.get(runner_id, None)
@@ -37,8 +35,27 @@ class MlxJacclInstance(BaseInstance):
     jaccl_coordinators: dict[NodeId, str]
 
 
+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)
+
+
 # TODO: Single node instance
-Instance = MlxRingInstance | MlxJacclInstance
+Instance = MlxRingInstance | MlxJacclInstance | FLASHInstance
 
 
 class BoundInstance(CamelCaseModel):

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

@@ -119,8 +119,6 @@ def mlx_generate(
     model: Model,
     tokenizer: TokenizerWrapper,
     task: ChatCompletionTaskParams,
-    draft_model: Model | None = None,
-    num_draft_tokens: int = 4,
 ) -> Generator[GenerationResponse]:
     # Ensure that generation stats only contains peak memory for this generation
     mx.reset_peak_memory()
@@ -137,6 +135,8 @@ def mlx_generate(
         chat_task_data=task,
     )
 
+    caches = make_kv_cache(model=model)
+
     logits_processors: list[Callable[[mx.array, mx.array], mx.array]] = []
     if is_bench:
         # Only sample length eos tokens
@@ -149,31 +149,19 @@ def mlx_generate(
     )
 
     max_tokens = task.max_tokens or MAX_TOKENS
-
-    # Build kwargs for stream_generate, conditionally adding draft model params
-    generate_kwargs: dict[str, object] = {
-        "model": model,
-        "tokenizer": tokenizer,
-        "prompt": prompt,
-        "max_tokens": max_tokens,
-        "sampler": sampler,
-        "logits_processors": logits_processors,
-        "prefill_step_size": 2048,
-        "kv_group_size": KV_GROUP_SIZE,
-        "kv_bits": KV_BITS,
-    }
-
-    # Add speculative decoding parameters if draft model is provided
-    # Note: When using draft_model, we let mlx_lm create its own trimmable cache
-    # as speculative decoding requires cache trimming capabilities
-    if draft_model is not None:
-        generate_kwargs["draft_model"] = draft_model
-        generate_kwargs["num_draft_tokens"] = num_draft_tokens
-    else:
-        # Only use custom cache for non-speculative generation
-        generate_kwargs["prompt_cache"] = make_kv_cache(model=model)
-
-    for out in stream_generate(**generate_kwargs):  # type: ignore[arg-type]
+    for out in stream_generate(
+        model=model,
+        tokenizer=tokenizer,
+        prompt=prompt,
+        max_tokens=max_tokens,
+        sampler=sampler,
+        logits_processors=logits_processors,
+        prompt_cache=caches,
+        # TODO: Dynamically change prefill step size to be the maximum possible without timing out.
+        prefill_step_size=2048,
+        kv_group_size=KV_GROUP_SIZE,
+        kv_bits=KV_BITS,
+    ):
         logger.info(out.text)
 
         stats: GenerationStats | None = None

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

@@ -2,9 +2,7 @@ import json
 import os
 import resource
 import sys
-import threading
 import time
-from collections.abc import Callable
 from pathlib import Path
 from typing import Any, cast
 
@@ -22,7 +20,6 @@ except ImportError:
 
 from mlx_lm.models.cache import KVCache, QuantizedKVCache, RotatingKVCache
 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.worker.engines.mlx.constants import (
@@ -84,45 +81,6 @@ def get_weights_size(model_shard_meta: ShardMetadata) -> Memory:
     )
 
 
-class ModelLoadingTimeoutError(Exception):
-    pass
-
-
-TimeoutCallback = Callable[[], None]
-
-
-def eval_with_timeout(
-    mlx_item: Any,  # pyright: ignore[reportAny]
-    timeout_seconds: float = 60.0,
-    on_timeout: TimeoutCallback | None = None,
-) -> None:
-    """Evaluate MLX item with a hard timeout.
-
-    If on_timeout callback is provided, it will be called before terminating
-    the process. This allows the runner to send a failure event before exit.
-    """
-    completed = threading.Event()
-
-    def watchdog() -> None:
-        if not completed.wait(timeout=timeout_seconds):
-            logger.error(
-                f"mlx_item evaluation timed out after {timeout_seconds:.0f}s. "
-                "This may indicate an issue with FAST_SYNCH and tensor parallel sharding. "
-                "Terminating process."
-            )
-            if on_timeout is not None:
-                on_timeout()
-            os._exit(1)
-
-    watchdog_thread = threading.Thread(target=watchdog, daemon=True)
-    watchdog_thread.start()
-
-    try:
-        mx.eval(mlx_item)  # pyright: ignore[reportAny]
-    finally:
-        completed.set()
-
-
 def mx_barrier(group: Group | None = None):
     mx.eval(
         mx.distributed.all_sum(
@@ -206,6 +164,11 @@ 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
@@ -229,9 +192,7 @@ def initialize_mlx(
 
 
 def load_mlx_items(
-    bound_instance: BoundInstance,
-    group: Group | None,
-    on_timeout: TimeoutCallback | None = None,
+    bound_instance: BoundInstance, group: Group | None
 ) -> tuple[Model, TokenizerWrapper]:
     if group is None:
         logger.info(f"Single device used for {bound_instance.instance}")
@@ -245,9 +206,7 @@ def load_mlx_items(
     else:
         logger.info("Starting distributed init")
         start_time = time.perf_counter()
-        model, tokenizer = shard_and_load(
-            bound_instance.bound_shard, group=group, on_timeout=on_timeout
-        )
+        model, tokenizer = shard_and_load(bound_instance.bound_shard, group=group)
         end_time = time.perf_counter()
         logger.info(
             f"Time taken to shard and load model: {(end_time - start_time):.2f}s"
@@ -258,31 +217,9 @@ def load_mlx_items(
     return cast(Model, model), tokenizer
 
 
-def load_draft_model(model_id: str) -> nn.Module:
-    """Load a draft model for speculative decoding (rank 0 only).
-
-    Draft models are small models (typically 0.5B-2B parameters) used to
-    generate candidate tokens quickly, which are then verified by the main
-    model in a single forward pass.
-
-    Assumes the model has already been downloaded by the worker.
-
-    Args:
-        model_id: HuggingFace model ID for the draft model
-
-    Returns:
-        The loaded draft model
-    """
-    model_path = build_model_path(model_id)
-    draft_model, _ = load_model(model_path, strict=True)
-    logger.info(f"Loaded draft model from {model_path}")
-    return draft_model
-
-
 def shard_and_load(
     shard_metadata: ShardMetadata,
     group: Group,
-    on_timeout: TimeoutCallback | None = None,
 ) -> tuple[nn.Module, TokenizerWrapper]:
     model_path = build_model_path(shard_metadata.model_meta.model_id)
 
@@ -319,15 +256,7 @@ def shard_and_load(
             logger.info(f"loading model from {model_path} with pipeline parallelism")
             model = pipeline_auto_parallel(model, group, shard_metadata)
 
-    # Estimate timeout based on model size
-    base_timeout = float(os.environ.get("EXO_MODEL_LOAD_TIMEOUT", "60"))
-    model_size_gb = get_weights_size(shard_metadata).in_bytes / (1024**3)
-    timeout_seconds = base_timeout + model_size_gb / 5
-    logger.info(
-        f"Evaluating model parameters with timeout of {timeout_seconds:.0f}s "
-        f"(model size: {model_size_gb:.1f}GB)"
-    )
-    eval_with_timeout(model.parameters(), timeout_seconds, on_timeout)
+    mx.eval(model.parameters())
 
     # TODO: Do we need this?
     mx.eval(model)
@@ -441,8 +370,6 @@ def apply_chat_template(
         tools=chat_task_data.tools,
     )
 
-    logger.info(prompt)
-
     return prompt
 
 
@@ -474,11 +401,6 @@ def make_kv_cache(
 ) -> 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")

src/exo/worker/plan.py

@@ -3,8 +3,7 @@
 from collections.abc import Mapping, Sequence
 
 from exo.shared.types.common import NodeId
-from exo.shared.types.memory import Memory
-from exo.shared.types.models import ModelId, ModelMetadata
+from exo.shared.types.models import ModelId
 from exo.shared.types.tasks import (
     ChatCompletion,
     ConnectToGroup,
@@ -22,7 +21,12 @@ from exo.shared.types.worker.downloads import (
     DownloadOngoing,
     DownloadProgress,
 )
-from exo.shared.types.worker.instances import BoundInstance, Instance, InstanceId
+from exo.shared.types.worker.instances import (
+    BoundInstance,
+    FLASHInstance,
+    Instance,
+    InstanceId,
+)
 from exo.shared.types.worker.runners import (
     RunnerConnected,
     RunnerConnecting,
@@ -36,7 +40,6 @@ from exo.shared.types.worker.runners import (
     RunnerStatus,
     RunnerWarmingUp,
 )
-from exo.shared.types.worker.shards import PipelineShardMetadata
 from exo.worker.runner.runner_supervisor import RunnerSupervisor
 
 
@@ -52,6 +55,11 @@ def plan(
     all_runners: Mapping[RunnerId, RunnerStatus],  # all global
     tasks: Mapping[TaskId, Task],
 ) -> Task | None:
+    # Check for FLASH instance tasks first
+    flash_task = _plan_flash(runners, instances)
+    if flash_task is not None:
+        return flash_task
+
     # Python short circuiting OR logic should evaluate these sequentially.
     return (
         _kill_runner(runners, all_runners, instances)
@@ -59,12 +67,39 @@ def plan(
         or _model_needs_download(runners, download_status)
         or _init_distributed_backend(runners, all_runners)
         or _load_model(runners, all_runners, global_download_status)
-        or _draft_model_needs_download(runners, download_status)
         or _ready_to_warmup(runners, all_runners)
         or _pending_tasks(runners, tasks, all_runners)
     )
 
 
+def _plan_flash(
+    runners: Mapping[RunnerId, RunnerSupervisor],
+    instances: Mapping[InstanceId, Instance],
+) -> Task | None:
+    """Plan tasks specifically for FLASH instances.
+
+    FLASH instances have a simpler lifecycle:
+    - CreateRunner (handled by _create_runner)
+    - LoadModel (starts the simulation immediately)
+    - Shutdown (handled by _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
+
+
 def _kill_runner(
     runners: Mapping[RunnerId, RunnerSupervisor],
     all_runners: Mapping[RunnerId, RunnerStatus],
@@ -117,6 +152,10 @@ def _model_needs_download(
     download_status: Mapping[ModelId, DownloadProgress],
 ) -> DownloadModel | None:
     for runner in runners.values():
+        # FLASH instances don't need model downloads
+        if isinstance(runner.bound_instance.instance, FLASHInstance):
+            continue
+
         model_id = runner.bound_instance.bound_shard.model_meta.model_id
         if isinstance(runner.status, RunnerIdle) and (
             model_id not in download_status
@@ -131,57 +170,6 @@ def _model_needs_download(
             )
 
 
-def _draft_model_needs_download(
-    runners: Mapping[RunnerId, RunnerSupervisor],
-    download_status: Mapping[ModelId, DownloadProgress],
-) -> DownloadModel | None:
-    """Check if draft model needs download (for speculative decoding).
-
-    Only rank 0 needs the draft model, and only after the main model is loaded.
-    """
-    for runner in runners.values():
-        instance = runner.bound_instance.instance
-        shard = runner.bound_instance.bound_shard
-
-        # Only check when runner is loaded and ready for warmup
-        if not isinstance(runner.status, RunnerLoaded):
-            continue
-
-        # Only rank 0 loads the draft model
-        if shard.device_rank != 0:
-            continue
-
-        # Check if instance has a draft model configured
-        draft_model_id = instance.draft_model
-        if draft_model_id is None:
-            continue
-
-        # Check if draft model needs download
-        if draft_model_id not in download_status or not isinstance(
-            download_status[draft_model_id], (DownloadOngoing, DownloadCompleted)
-        ):
-            # Create minimal shard metadata for draft model download
-            draft_shard = PipelineShardMetadata(
-                model_meta=ModelMetadata(
-                    model_id=draft_model_id,
-                    pretty_name=str(draft_model_id),
-                    storage_size=Memory.from_bytes(0),  # Unknown, will be determined during download
-                    n_layers=1,  # Placeholder
-                    hidden_size=1,  # Placeholder
-                    supports_tensor=False,
-                ),
-                device_rank=0,
-                world_size=1,
-                start_layer=0,
-                end_layer=1,
-                n_layers=1,
-            )
-            return DownloadModel(
-                instance_id=instance.instance_id,
-                shard_metadata=draft_shard,
-            )
-
-
 def _init_distributed_backend(
     runners: Mapping[RunnerId, RunnerSupervisor],
     all_runners: Mapping[RunnerId, RunnerStatus],

src/exo/worker/runner/bootstrap.py

@@ -4,7 +4,11 @@ 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,
+    FLASHInstance,
+    MlxJacclInstance,
+)
 from exo.shared.types.worker.runners import RunnerFailed
 from exo.utils.channels import ClosedResourceError, MpReceiver, MpSender
 
@@ -17,28 +21,27 @@ def entrypoint(
     task_receiver: MpReceiver[Task],
     _logger: "loguru.Logger",
 ) -> None:
-    fast_synch_override = os.environ.get("EXO_FAST_SYNCH")
-    if fast_synch_override == "on" or (
-        fast_synch_override != "off"
-        and (
-            isinstance(bound_instance.instance, MlxJacclInstance)
-            and len(bound_instance.instance.ibv_devices) >= 2
-        )
-    ):
-        os.environ["MLX_METAL_FAST_SYNCH"] = "1"
-    else:
-        os.environ["MLX_METAL_FAST_SYNCH"] = "0"
-
     global logger
     logger = _logger
 
-    logger.info(f"Fast synch flag: {os.environ['MLX_METAL_FAST_SYNCH']}")
-
-    # Import main after setting global logger - this lets us just import logger from this module
+    # Route based on instance type
     try:
-        from exo.worker.runner.runner import main
+        if isinstance(bound_instance.instance, FLASHInstance):
+            # FLASH MPI simulation runner
+            from exo.worker.runner.flash_runner import main
 
-        main(bound_instance, event_sender, task_receiver)
+            main(bound_instance, event_sender, task_receiver)
+        else:
+            # MLX runner (default)
+            if (
+                isinstance(bound_instance.instance, MlxJacclInstance)
+                and len(bound_instance.instance.ibv_devices) >= 2
+            ):
+                os.environ["MLX_METAL_FAST_SYNCH"] = "1"
+
+            from exo.worker.runner.runner import main
+
+            main(bound_instance, event_sender, task_receiver)
     except ClosedResourceError:
         logger.warning("Runner communication closed unexpectedly")
     except Exception as e:

src/exo/worker/runner/flash_runner.py

@@ -0,0 +1,301 @@
+"""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
+"""
+
+import os
+import shutil
+import socket
+import subprocess
+import threading
+
+from exo.shared.types.events import (
+    Event,
+    RunnerStatusUpdated,
+    TaskAcknowledged,
+    TaskStatusUpdated,
+)
+from exo.shared.types.tasks import (
+    LoadModel,
+    Shutdown,
+    Task,
+    TaskStatus,
+)
+from exo.shared.types.worker.instances import BoundInstance, FLASHInstance
+from exo.shared.types.worker.runners import (
+    RunnerFailed,
+    RunnerIdle,
+    RunnerLoading,
+    RunnerReady,
+    RunnerRunning,
+    RunnerShutdown,
+    RunnerShuttingDown,
+    RunnerStatus,
+)
+from exo.utils.channels import MpReceiver, MpSender
+from exo.worker.runner.bootstrap import logger
+
+# 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],
+):
+    """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",
+                                "--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/worker/runner/runner.py

@@ -1,21 +1,9 @@
 import time
-from collections.abc import Generator
-from contextlib import contextmanager
-from functools import cache
-from typing import cast
 
 import mlx.core as mx
-from mlx_lm.models.gpt_oss import Model as GptOssModel
-from openai_harmony import (  # pyright: ignore[reportMissingTypeStubs]
-    HarmonyEncodingName,
-    Role,
-    StreamableParser,
-    load_harmony_encoding,
-)
 
 from exo.shared.types.api import ChatCompletionMessageText
 from exo.shared.types.chunks import TokenChunk
-from exo.shared.types.common import CommandId
 from exo.shared.types.events import (
     ChunkGenerated,
     Event,
@@ -23,7 +11,6 @@ from exo.shared.types.events import (
     TaskAcknowledged,
     TaskStatusUpdated,
 )
-from exo.shared.types.models import ModelId
 from exo.shared.types.tasks import (
     ChatCompletion,
     ConnectToGroup,
@@ -52,44 +39,15 @@ from exo.shared.types.worker.runners import (
     RunnerWarmingUp,
 )
 from exo.utils.channels import MpReceiver, MpSender
-from exo.worker.engines.mlx import Model
 from exo.worker.engines.mlx.generator.generate import mlx_generate, warmup_inference
 from exo.worker.engines.mlx.utils_mlx import (
     initialize_mlx,
-    load_draft_model,
     load_mlx_items,
     mlx_force_oom,
 )
 from exo.worker.runner.bootstrap import logger
 
 
-@contextmanager
-def send_error_chunk_on_exception(
-    event_sender: MpSender[Event],
-    command_id: CommandId,
-    model_id: ModelId,
-    device_rank: int,
-):
-    try:
-        yield
-    except Exception as e:
-        logger.error(e)
-        if device_rank == 0:
-            event_sender.send(
-                ChunkGenerated(
-                    command_id=command_id,
-                    chunk=TokenChunk(
-                        idx=0,
-                        model=model_id,
-                        text="",
-                        token_id=0,
-                        finish_reason="error",
-                        error_message=str(e),
-                    ),
-                )
-            )
-
-
 def main(
     bound_instance: BoundInstance,
     event_sender: MpSender[Event],
@@ -111,7 +69,6 @@ def main(
     model = None
     tokenizer = None
     group = None
-    draft_model: Model | None = None  # Loaded during warmup if instance has draft_model
 
     current_status: RunnerStatus = RunnerIdle()
     logger.info("runner created")
@@ -152,20 +109,7 @@ def main(
                         )
                     )
 
-                    def on_model_load_timeout() -> None:
-                        event_sender.send(
-                            RunnerStatusUpdated(
-                                runner_id=runner_id,
-                                runner_status=RunnerFailed(
-                                    error_message="Model loading timed out"
-                                ),
-                            )
-                        )
-                        time.sleep(0.5)
-
-                    model, tokenizer = load_mlx_items(
-                        bound_instance, group, on_timeout=on_model_load_timeout
-                    )
+                    model, tokenizer = load_mlx_items(bound_instance, group)
 
                     current_status = RunnerLoaded()
                     logger.info("runner loaded")
@@ -180,20 +124,11 @@ def main(
                         )
                     )
 
-                    # Load draft model for speculative decoding (rank 0 only)
-                    if (
-                        instance.draft_model is not None
-                        and shard_metadata.device_rank == 0
-                    ):
-                        logger.info(f"Loading draft model: {instance.draft_model}")
-                        draft_model = cast(
-                            Model, load_draft_model(str(instance.draft_model))
-                        )
-
                     logger.info(f"warming up inference for instance: {instance}")
                     toks = warmup_inference(
-                        model=cast(Model, model),
+                        model=model,
                         tokenizer=tokenizer,
+                        # kv_prefix_cache=kv_prefix_cache,  # supply for warmup-time prefix caching
                     )
                     logger.info(f"warmed up by generating {toks} tokens")
                     logger.info(
@@ -204,6 +139,8 @@ def main(
                 case ChatCompletion(task_params=task_params, command_id=command_id) if (
                     isinstance(current_status, RunnerReady)
                 ):
+                    assert model
+                    assert tokenizer
                     logger.info(f"received chat request: {str(task)[:500]}")
                     current_status = RunnerRunning()
                     logger.info("runner running")
@@ -212,49 +149,33 @@ def main(
                             runner_id=runner_id, runner_status=current_status
                         )
                     )
-                    with send_error_chunk_on_exception(
-                        event_sender,
-                        command_id,
-                        shard_metadata.model_meta.model_id,
-                        shard_metadata.device_rank,
+                    assert task_params.messages[0].content is not None
+                    _check_for_debug_prompts(task_params.messages[0].content)
+
+                    # Generate responses using the actual MLX generation
+                    for response in mlx_generate(
+                        model=model,
+                        tokenizer=tokenizer,
+                        task=task_params,
                     ):
-                        assert model
-                        assert tokenizer
-                        assert task_params.messages[0].content is not None
-                        _check_for_debug_prompts(task_params.messages[0].content)
-
-                        # Generate responses (draft_model loaded at warmup if configured)
-                        mlx_generator = mlx_generate(
-                            model=cast(Model, model),
-                            tokenizer=tokenizer,
-                            task=task_params,
-                            draft_model=draft_model,
-                            num_draft_tokens=instance.num_draft_tokens,
-                        )
-
-                        # GPT-OSS specific parsing to match other model formats.
-                        if isinstance(model, GptOssModel):
-                            mlx_generator = parse_gpt_oss(mlx_generator)
-
-                        # TODO: Add tool call parser here
-
-                        for response in mlx_generator:
-                            match response:
-                                case GenerationResponse():
-                                    if shard_metadata.device_rank == 0:
-                                        event_sender.send(
-                                            ChunkGenerated(
-                                                command_id=command_id,
-                                                chunk=TokenChunk(
-                                                    idx=response.token,
-                                                    model=shard_metadata.model_meta.model_id,
-                                                    text=response.text,
-                                                    token_id=response.token,
-                                                    finish_reason=response.finish_reason,
-                                                    stats=response.stats,
-                                                ),
-                                            )
+                        match response:
+                            case GenerationResponse():
+                                if shard_metadata.device_rank == 0:
+                                    event_sender.send(
+                                        ChunkGenerated(
+                                            command_id=command_id,
+                                            chunk=TokenChunk(
+                                                idx=response.token,
+                                                model=shard_metadata.model_meta.model_id,
+                                                text=response.text,
+                                                token_id=response.token,
+                                                finish_reason=response.finish_reason,
+                                                stats=response.stats,
+                                            ),
                                         )
+                                    )
+                                # case TokenizedResponse():
+                                # TODO: something here ig
 
                     current_status = RunnerReady()
                     logger.info("runner ready")
@@ -278,7 +199,7 @@ def main(
                 RunnerStatusUpdated(runner_id=runner_id, runner_status=current_status)
             )
             if isinstance(current_status, RunnerShutdown):
-                del model, tokenizer, group, draft_model
+                del model, tokenizer, group
                 mx.clear_cache()
                 import gc
 
@@ -286,43 +207,6 @@ def main(
                 break
 
 
-@cache
-def get_gpt_oss_encoding():
-    encoding = load_harmony_encoding(HarmonyEncodingName.HARMONY_GPT_OSS)
-    return encoding
-
-
-def parse_gpt_oss(
-    responses: Generator[GenerationResponse],
-) -> Generator[GenerationResponse]:
-    encoding = get_gpt_oss_encoding()
-    stream = StreamableParser(encoding, role=Role.ASSISTANT)
-    thinking = False
-
-    for response in responses:
-        stream.process(response.token)
-
-        delta = stream.last_content_delta
-        ch = stream.current_channel
-
-        if ch == "analysis" and not thinking:
-            thinking = True
-            yield response.model_copy(update={"text": "<think>"})
-
-        if ch != "analysis" and thinking:
-            thinking = False
-            yield response.model_copy(update={"text": "</think>"})
-
-        if delta:
-            yield response.model_copy(update={"text": delta})
-
-        if response.finish_reason is not None:
-            if thinking:
-                yield response.model_copy(update={"text": "</think>"})
-            yield response
-            break
-
-
 EXO_RUNNER_MUST_FAIL = "EXO RUNNER MUST FAIL"
 EXO_RUNNER_MUST_OOM = "EXO RUNNER MUST OOM"
 EXO_RUNNER_MUST_TIMEOUT = "EXO RUNNER MUST TIMEOUT"

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

@@ -1,50 +0,0 @@
-# pyright: reportAny=false
-from unittest.mock import MagicMock
-
-from exo.shared.types.chunks import TokenChunk
-from exo.shared.types.common import CommandId
-from exo.shared.types.events import ChunkGenerated
-from exo.worker.runner.runner import send_error_chunk_on_exception
-from exo.worker.tests.constants import MODEL_A_ID
-
-
-def test_send_error_chunk_on_exception_no_error() -> None:
-    event_sender = MagicMock()
-    command_id = CommandId()
-
-    with send_error_chunk_on_exception(
-        event_sender, command_id, MODEL_A_ID, device_rank=0
-    ):
-        _ = 1 + 1
-
-    event_sender.send.assert_not_called()
-
-
-def test_send_error_chunk_on_exception_catches_error() -> None:
-    event_sender = MagicMock()
-    command_id = CommandId()
-
-    with send_error_chunk_on_exception(
-        event_sender, command_id, MODEL_A_ID, device_rank=0
-    ):
-        raise ValueError("test error")
-
-    event_sender.send.assert_called_once()
-    call_args = event_sender.send.call_args[0][0]
-    assert isinstance(call_args, ChunkGenerated)
-    assert call_args.command_id == command_id
-    assert isinstance(call_args.chunk, TokenChunk)
-    assert call_args.chunk.finish_reason == "error"
-    assert call_args.chunk.error_message == "test error"
-
-
-def test_send_error_chunk_on_exception_skips_non_rank_zero() -> None:
-    event_sender = MagicMock()
-    command_id = CommandId()
-
-    with send_error_chunk_on_exception(
-        event_sender, command_id, MODEL_A_ID, device_rank=1
-    ):
-        raise ValueError("test error")
-
-    event_sender.send.assert_not_called()

src/exo/worker/utils/net_profile.py

@@ -1,64 +1,49 @@
-import anyio
-import httpx
-from anyio import create_task_group
+import http.client
+
+from anyio import create_task_group, to_thread
 from loguru import logger
 
 from exo.shared.topology import Topology
 from exo.shared.types.common import NodeId
 
-REACHABILITY_ATTEMPTS = 3
-
 
 async def check_reachability(
     target_ip: str,
     expected_node_id: NodeId,
+    self_node_id: NodeId,
     out: dict[NodeId, set[str]],
-    client: httpx.AsyncClient,
 ) -> None:
     """Check if a node is reachable at the given IP and verify its identity."""
-    if ":" in target_ip:
-        # TODO: use real IpAddress types
-        target_ip = f"[{target_ip}]"
-    url = f"http://{target_ip}:52415/node_id"
 
-    remote_node_id = None
-    last_error = None
-
-    for _ in range(REACHABILITY_ATTEMPTS):
+    def _fetch_remote_node_id() -> NodeId | None:
+        connection = http.client.HTTPConnection(target_ip, 52415, timeout=1)
         try:
-            r = await client.get(url)
-            if r.status_code != 200:
-                await anyio.sleep(1)
-                continue
+            connection.request("GET", "/node_id")
+            response = connection.getresponse()
+            if response.status != 200:
+                return None
 
-            body = r.text.strip().strip('"')
-            if not body:
-                await anyio.sleep(1)
-                continue
+            body = response.read().decode("utf-8").strip()
 
-            remote_node_id = NodeId(body)
-            break
+            # Strip quotes if present (JSON string response)
+            if body.startswith('"') and body.endswith('"') and len(body) >= 2:
+                body = body[1:-1]
 
-        # expected failure cases
-        except (
-            httpx.TimeoutException,
-            httpx.NetworkError,
-        ):
-            await anyio.sleep(1)
-
-        # other failures should be logged on last attempt
-        except httpx.HTTPError as e:
-            last_error = e
-            await anyio.sleep(1)
-
-    if last_error is not None:
-        logger.warning(
-            f"connect error {type(last_error).__name__} from {target_ip} after {REACHABILITY_ATTEMPTS} attempts; treating as down"
-        )
+            return NodeId(body) or None
+        except OSError:
+            return None
+        except http.client.HTTPException:
+            return None
+        finally:
+            connection.close()
 
+    remote_node_id = await to_thread.run_sync(_fetch_remote_node_id)
     if remote_node_id is None:
         return
 
+    if remote_node_id == self_node_id:
+        return
+
     if remote_node_id != expected_node_id:
         logger.warning(
             f"Discovered node with unexpected node_id; "
@@ -76,33 +61,18 @@ async def check_reachable(
     topology: Topology, self_node_id: NodeId
 ) -> dict[NodeId, set[str]]:
     """Check which nodes are reachable and return their IPs."""
-
     reachable: dict[NodeId, set[str]] = {}
-
-    # these are intentionally httpx's defaults so we can tune them later
-    timeout = httpx.Timeout(timeout=5.0)
-    limits = httpx.Limits(
-        max_connections=100,
-        max_keepalive_connections=20,
-        keepalive_expiry=5,
-    )
-
-    async with (
-        httpx.AsyncClient(timeout=timeout, limits=limits) as client,
-        create_task_group() as tg,
-    ):
+    async with create_task_group() as tg:
         for node in topology.list_nodes():
             if not node.node_profile:
                 continue
-            if node.node_id == self_node_id:
-                continue
             for iface in node.node_profile.network_interfaces:
                 tg.start_soon(
                     check_reachability,
                     iface.ip_address,
                     node.node_id,
+                    self_node_id,
                     reachable,
-                    client,
                 )
 
     return reachable