Add exo eval

https://platform.openai.com/docs/api-reference/chat/get

.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

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)
 

bench/exo_eval.py

@@ -0,0 +1,378 @@
+#!/usr/bin/env python3
+# pyright: reportAny=false, reportUnknownMemberType=false, reportUnknownVariableType=false, reportUnknownArgumentType=false, reportMissingTypeStubs=false
+"""
+exo-eval: Run SWE-bench evaluation against exo using OpenHands SDK (local, no Docker).
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import subprocess
+import tempfile
+import time
+from dataclasses import dataclass
+from enum import Enum
+from pathlib import Path
+from typing import Any
+
+from datasets import load_dataset
+from loguru import logger
+from openhands.sdk import LLM, Agent, Conversation, Tool
+from openhands.tools.file_editor import FileEditorTool
+from openhands.tools.terminal import TerminalTool
+
+
+class EvalStatus(str, Enum):
+    Resolved = "Resolved"
+    Failed = "Failed"
+    Error = "Error"
+    Timeout = "Timeout"
+
+
+@dataclass
+class EvalResult:
+    instance_id: str
+    repo: str
+    status: EvalStatus
+    elapsed_seconds: float
+    tests_passed: list[str]
+    tests_failed: list[str]
+    error_message: str | None = None
+
+
+def load_swe_bench(
+    split: str = "lite",
+    limit: int | None = None,
+    instance_ids: list[str] | None = None,
+) -> list[dict[str, Any]]:
+    """Load SWE-bench dataset from HuggingFace."""
+    # SWE-bench Lite is a curated 300-instance subset
+    dataset_name = (
+        "princeton-nlp/SWE-bench_Lite" if split == "lite" else "princeton-nlp/SWE-bench"
+    )
+    actual_split = "test" if split == "lite" else split
+
+    ds = load_dataset(dataset_name, split=actual_split)
+    instances = [dict(row) for row in ds]
+
+    if instance_ids:
+        instances = [i for i in instances if i["instance_id"] in instance_ids]
+
+    if limit:
+        instances = instances[:limit]
+
+    return instances
+
+
+def clone_repo_at_commit(repo: str, commit: str, dest: Path) -> None:
+    """Clone a repo at a specific commit."""
+    repo_url = f"https://github.com/{repo}.git"
+
+    subprocess.run(
+        ["git", "clone", "--depth", "1", repo_url, str(dest)],
+        check=True,
+        capture_output=True,
+    )
+
+    subprocess.run(
+        ["git", "fetch", "--depth", "1", "origin", commit],
+        cwd=dest,
+        check=True,
+        capture_output=True,
+    )
+
+    subprocess.run(
+        ["git", "checkout", commit],
+        cwd=dest,
+        check=True,
+        capture_output=True,
+    )
+
+
+def build_agent_prompt(instance: dict[str, Any]) -> str:
+    """Build the prompt for the agent."""
+    return f"""You are a software engineer fixing a bug in the {instance['repo']} repository.
+
+## Problem Statement
+{instance['problem_statement']}
+
+## Instructions
+1. Explore the codebase to understand the issue
+2. Identify the files that need to be modified
+3. Make the necessary changes to fix the issue
+4. The fix should be minimal and targeted
+
+You have access to:
+- terminal: Run shell commands (git, grep, python, etc.)
+- file_editor: View and edit files
+
+Start by exploring the repository structure to understand where the relevant code is.
+"""
+
+
+def parse_fail_to_pass(fail_to_pass_str: str) -> list[str]:
+    """Parse the FAIL_TO_PASS field into a list of test names."""
+    try:
+        return json.loads(fail_to_pass_str)
+    except json.JSONDecodeError:
+        return [t.strip() for t in fail_to_pass_str.split(",") if t.strip()]
+
+
+def run_tests(workspace: Path, tests: list[str]) -> tuple[list[str], list[str]]:
+    """Run tests and return (passed, failed) lists."""
+    passed = []
+    failed = []
+
+    for test in tests:
+        try:
+            result = subprocess.run(
+                ["python", "-m", "pytest", "-xvs", test],
+                cwd=workspace,
+                capture_output=True,
+                timeout=300,
+            )
+            if result.returncode == 0:
+                passed.append(test)
+            else:
+                failed.append(test)
+        except subprocess.TimeoutExpired:
+            failed.append(test)
+
+    return passed, failed
+
+
+def run_single_eval(
+    instance: dict[str, Any],
+    host: str,
+    port: int,
+    model: str,
+    max_turns: int = 30,
+    timeout: float = 600.0,
+) -> EvalResult:
+    """Evaluate a single SWE-bench instance."""
+    instance_id = instance["instance_id"]
+    repo = instance["repo"]
+    base_commit = instance["base_commit"]
+    fail_to_pass = parse_fail_to_pass(instance["FAIL_TO_PASS"])
+
+    start_time = time.perf_counter()
+
+    try:
+        with tempfile.TemporaryDirectory() as tmpdir:
+            workspace = Path(tmpdir) / "repo"
+
+            # Clone repo at base commit
+            logger.info(f"Cloning {repo} at {base_commit[:8]}...")
+            clone_repo_at_commit(repo, base_commit, workspace)
+
+            # Setup OpenHands agent
+            llm = LLM(
+                model=f"openai/{model}",
+                base_url=f"http://{host}:{port}/v1",
+                api_key="not-needed",
+            )
+
+            agent = Agent(
+                llm=llm,
+                tools=[
+                    Tool(name=TerminalTool.name),
+                    Tool(name=FileEditorTool.name),
+                ],
+            )
+
+            # Run agent
+            conversation = Conversation(
+                agent=agent,
+                workspace=str(workspace),
+            )
+
+            logger.info(f"Running agent on {instance_id}...")
+            conversation.send_message(build_agent_prompt(instance))
+
+            for _turn in range(max_turns):
+                if time.perf_counter() - start_time > timeout:
+                    return EvalResult(
+                        instance_id=instance_id,
+                        repo=repo,
+                        status=EvalStatus.Timeout,
+                        elapsed_seconds=time.perf_counter() - start_time,
+                        tests_passed=[],
+                        tests_failed=fail_to_pass,
+                    )
+
+                result = conversation.run(max_turns=1)
+                if result.done:
+                    break
+
+            # Run tests to verify
+            logger.info(f"Running tests for {instance_id}...")
+            passed, failed = run_tests(workspace, fail_to_pass)
+
+            elapsed = time.perf_counter() - start_time
+            status = EvalStatus.Resolved if not failed else EvalStatus.Failed
+
+            return EvalResult(
+                instance_id=instance_id,
+                repo=repo,
+                status=status,
+                elapsed_seconds=elapsed,
+                tests_passed=passed,
+                tests_failed=failed,
+            )
+
+    except Exception as e:
+        return EvalResult(
+            instance_id=instance_id,
+            repo=repo,
+            status=EvalStatus.Error,
+            elapsed_seconds=time.perf_counter() - start_time,
+            tests_passed=[],
+            tests_failed=[],
+            error_message=str(e),
+        )
+
+
+def verify_exo_running(host: str, port: int, model: str) -> str:
+    """Verify exo is running and return full model ID."""
+    import http.client
+
+    conn = http.client.HTTPConnection(host, port, timeout=10)
+    conn.request("GET", "/models")
+    resp = conn.getresponse()
+
+    if resp.status != 200:
+        raise RuntimeError(f"exo not responding at {host}:{port}")
+
+    data = json.loads(resp.read())
+    for m in data.get("data", []):
+        if m.get("id") == model or m.get("hugging_face_id") == model:
+            return m.get("hugging_face_id") or m.get("id")
+
+    raise ValueError(f"Model '{model}' not found in exo")
+
+
+def main() -> int:
+    ap = argparse.ArgumentParser(
+        prog="exo-eval",
+        description="Run SWE-bench evaluation against exo (local, no Docker).",
+    )
+
+    ap.add_argument("--host", default=os.environ.get("EXO_HOST", "localhost"))
+    ap.add_argument(
+        "--port", type=int, default=int(os.environ.get("EXO_PORT", "52415"))
+    )
+    ap.add_argument("--model", required=True, help="exo model ID")
+    ap.add_argument(
+        "--split", default="lite", choices=["lite", "dev", "test", "train"]
+    )
+    ap.add_argument("--limit", type=int, default=10, help="Max instances")
+    ap.add_argument("--instance-ids", nargs="+", help="Specific instance IDs")
+    ap.add_argument("--max-turns", type=int, default=30)
+    ap.add_argument("--timeout", type=float, default=600.0)
+    ap.add_argument("--json-out", default="bench/eval_results.json")
+    ap.add_argument("-v", "--verbose", action="store_true")
+    ap.add_argument("--dry-run", action="store_true")
+
+    args = ap.parse_args()
+
+    # Load dataset first (doesn't require exo to be running)
+    logger.info(f"Loading SWE-bench {args.split} dataset...")
+    instances = load_swe_bench(
+        split=args.split,
+        limit=args.limit,
+        instance_ids=args.instance_ids,
+    )
+    logger.info(f"Loaded {len(instances)} instances")
+
+    if args.dry_run:
+        print(f"\nSWE-bench {args.split} instances ({len(instances)}):")
+        for inst in instances:
+            print(f"  {inst['instance_id']} ({inst['repo']})")
+        return 0
+
+    # Verify exo is running
+    model_id = verify_exo_running(args.host, args.port, args.model)
+    logger.info(f"Using model: {model_id}")
+
+    # Run evaluation
+    results: list[EvalResult] = []
+    for i, instance in enumerate(instances):
+        logger.info(f"[{i+1}/{len(instances)}] {instance['instance_id']}")
+
+        result = run_single_eval(
+            instance=instance,
+            host=args.host,
+            port=args.port,
+            model=model_id,
+            max_turns=args.max_turns,
+            timeout=args.timeout,
+        )
+        results.append(result)
+
+        logger.info(f"  Status: {result.status.value}")
+        if result.tests_passed:
+            logger.info(f"  Passed: {len(result.tests_passed)} tests")
+        if result.tests_failed:
+            logger.info(f"  Failed: {len(result.tests_failed)} tests")
+        if result.error_message:
+            logger.error(f"  Error: {result.error_message}")
+
+    # Compute summary
+    total = len(results)
+    resolved = sum(1 for r in results if r.status == EvalStatus.Resolved)
+    failed = sum(1 for r in results if r.status == EvalStatus.Failed)
+    errors = sum(1 for r in results if r.status == EvalStatus.Error)
+    timeouts = sum(1 for r in results if r.status == EvalStatus.Timeout)
+
+    summary = {
+        "model": model_id,
+        "split": args.split,
+        "total": total,
+        "resolved": resolved,
+        "resolved_rate": resolved / total if total else 0,
+        "failed": failed,
+        "errors": errors,
+        "timeouts": timeouts,
+    }
+
+    output = {
+        "summary": summary,
+        "results": [
+            {
+                "instance_id": r.instance_id,
+                "repo": r.repo,
+                "status": r.status.value,
+                "elapsed_seconds": r.elapsed_seconds,
+                "tests_passed": r.tests_passed,
+                "tests_failed": r.tests_failed,
+                "error_message": r.error_message,
+            }
+            for r in results
+        ],
+    }
+
+    Path(args.json_out).write_text(json.dumps(output, indent=2))
+    logger.info(f"Results written to {args.json_out}")
+
+    # Print summary
+    print("\n" + "=" * 60)
+    print("SWE-bench Evaluation Results")
+    print("=" * 60)
+    print(f"Model:    {model_id}")
+    print(f"Split:    {args.split}")
+    print(f"Total:    {total}")
+    if total:
+        print(f"Resolved: {resolved} ({resolved/total*100:.1f}%)")
+    else:
+        print("Resolved: 0")
+    print(f"Failed:   {failed}")
+    print(f"Errors:   {errors}")
+    print(f"Timeouts: {timeouts}")
+    print("=" * 60)
+
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

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';
 	
@@ -59,7 +59,7 @@ const sidebarVisible = $derived(chatSidebarVisible());
 		instanceType: InstanceMeta;
 		minNodes: number;
 	}
-
+	
 	function saveLaunchDefaults(): void {
 		const defaults: LaunchDefaults = {
 			modelId: selectedPreviewModelId(),
@@ -88,16 +88,16 @@ 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);
@@ -109,19 +109,11 @@ const sidebarVisible = $derived(chatSidebarVisible());
 	let minNodesInitialized = $state(false);
 	let launchingModelId = $state<string | null>(null);
 let instanceDownloadExpandedNodes = $state<Set<string>>(new Set());
-
-	// Draft model edit modal state
-	let editingDraftInstanceId = $state<string | null>(null);
-	let editDraftModel = $state<string | null>(null);
-	let editNumDraftTokens = $state<number>(4);
-	let isDraftEditDropdownOpen = $state(false);
-	let draftEditDropdownSearch = $state('');
-	let isSavingDraftModel = $state(false);
-
+	
 	// Custom dropdown state
 	let isModelDropdownOpen = $state(false);
 	let modelDropdownSearch = $state('');
-
+	
 	// Slider dragging state
 	let isDraggingSlider = $state(false);
 	let sliderTrackElement: HTMLDivElement | null = $state(null);
@@ -370,36 +362,47 @@ 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
-			const placePayload = {
-				model_id: modelId,
-				sharding: preview?.sharding ?? selectedSharding,
-				instance_meta: preview?.instance_meta ?? selectedInstanceType,
-				min_nodes: selectedMinNodes,
-			};
-
-			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);
-
+				
 				// 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 = () => {
@@ -794,52 +797,6 @@ function toggleInstanceDownloadDetails(nodeId: string): void {
 		}
 	}
 
-	// Open draft model edit modal for an instance
-	function openDraftModelEdit(instanceId: string, currentDraftModel: string | null, currentNumTokens: number | null) {
-		editingDraftInstanceId = instanceId;
-		editDraftModel = currentDraftModel;
-		editNumDraftTokens = currentNumTokens ?? 4;
-		isDraftEditDropdownOpen = false;
-		draftEditDropdownSearch = '';
-	}
-
-	// Close draft model edit modal
-	function closeDraftModelEdit() {
-		editingDraftInstanceId = null;
-		editDraftModel = null;
-		editNumDraftTokens = 4;
-		isDraftEditDropdownOpen = false;
-		draftEditDropdownSearch = '';
-	}
-
-	// Save draft model settings for an instance
-	async function saveDraftModel() {
-		if (!editingDraftInstanceId || isSavingDraftModel) return;
-
-		isSavingDraftModel = true;
-		try {
-			const response = await fetch(`/instance/${editingDraftInstanceId}/draft_model`, {
-				method: 'PUT',
-				headers: { 'Content-Type': 'application/json' },
-				body: JSON.stringify({
-					draft_model: editDraftModel,
-					num_draft_tokens: editNumDraftTokens,
-				})
-			});
-
-			if (!response.ok) {
-				const errorText = await response.text();
-				console.error('Failed to set draft model:', errorText);
-			} else {
-				closeDraftModelEdit();
-			}
-		} catch (error) {
-			console.error('Error setting draft model:', error);
-		} finally {
-			isSavingDraftModel = false;
-		}
-	}
-
 	// Helper to unwrap tagged unions like { MlxRingInstance: {...} }
 	function getTagged(obj: unknown): [string | null, unknown] {
 		if (!obj || typeof obj !== 'object') return [null, null];
@@ -859,34 +816,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 || {};
@@ -897,7 +850,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);
@@ -905,12 +858,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 {
@@ -1386,31 +1335,16 @@ function toggleInstanceDownloadDetails(nodeId: string): void {
 												<div class="w-1.5 h-1.5 {isDownloading ? 'bg-blue-400 animate-pulse' : isFailed ? 'bg-red-400' : isLoading ? 'bg-yellow-400 animate-pulse' : isReady ? 'bg-green-400' : 'bg-teal-400'} rounded-full shadow-[0_0_6px_currentColor]"></div>
 												<span class="text-exo-light-gray font-mono text-sm tracking-wider">{id.slice(0, 8).toUpperCase()}</span>
 											</div>
-											<div class="flex items-center gap-2">
-												<!-- Draft Model Button -->
-												<button
-													onclick={() => openDraftModelEdit(id, instanceInfo.draftModel, instanceInfo.numDraftTokens)}
-													class="p-1.5 font-mono border transition-all duration-200 cursor-pointer {instanceInfo.draftModel ? 'border-cyan-500/50 text-cyan-400 hover:bg-cyan-500/20 hover:border-cyan-500' : 'border-exo-medium-gray/50 text-white/40 hover:text-cyan-400 hover:border-cyan-500/50'}"
-													title={instanceInfo.draftModel ? `Draft: ${instanceInfo.draftModel.split('/').pop()} (${instanceInfo.numDraftTokens}t)` : 'Configure speculative decoding'}
-												>
-													<svg class="w-4 h-4" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round">
-														<path d="M13 2L3 14h9l-1 8 10-12h-9l1-8z"/>
-													</svg>
-												</button>
-												<button
-													onclick={() => deleteInstance(id)}
-													class="text-xs px-2 py-1 font-mono tracking-wider uppercase border border-red-500/30 text-red-400 hover:bg-red-500/20 hover:text-red-400 hover:border-red-500/50 transition-all duration-200 cursor-pointer"
-												>
-													DELETE
-												</button>
-											</div>
+											<button 
+												onclick={() => deleteInstance(id)}
+												class="text-xs px-2 py-1 font-mono tracking-wider uppercase border border-red-500/30 text-red-400 hover:bg-red-500/20 hover:text-red-400 hover:border-red-500/50 transition-all duration-200 cursor-pointer"
+											>
+												DELETE
+											</button>
 											</div>
 											<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"
@@ -1745,7 +1679,7 @@ function toggleInstanceDownloadDetails(nodeId: string): void {
 								</div>
 							</div>
 						</div>
-
+						
 						<!-- Selected Model Preview -->
 						<div class="space-y-3">
 							{#if models.length === 0}
@@ -1904,31 +1838,16 @@ function toggleInstanceDownloadDetails(nodeId: string): void {
 													<div class="w-1.5 h-1.5 {isDownloading ? 'bg-blue-400 animate-pulse' : isFailed ? 'bg-red-400' : isLoading ? 'bg-yellow-400 animate-pulse' : isReady ? 'bg-green-400' : 'bg-teal-400'} rounded-full shadow-[0_0_6px_currentColor]"></div>
 													<span class="text-exo-light-gray font-mono text-sm tracking-wider">{id.slice(0, 8).toUpperCase()}</span>
 												</div>
-												<div class="flex items-center gap-2">
-													<!-- Draft Model Button -->
-													<button
-														onclick={() => openDraftModelEdit(id, instanceInfo.draftModel, instanceInfo.numDraftTokens)}
-														class="p-1.5 font-mono border transition-all duration-200 cursor-pointer {instanceInfo.draftModel ? 'border-cyan-500/50 text-cyan-400 hover:bg-cyan-500/20 hover:border-cyan-500' : 'border-exo-medium-gray/50 text-white/40 hover:text-cyan-400 hover:border-cyan-500/50'}"
-														title={instanceInfo.draftModel ? `Draft: ${instanceInfo.draftModel.split('/').pop()} (${instanceInfo.numDraftTokens}t)` : 'Configure speculative decoding'}
-													>
-														<svg class="w-4 h-4" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round">
-															<path d="M13 2L3 14h9l-1 8 10-12h-9l1-8z"/>
-														</svg>
-													</button>
-													<button
-														onclick={() => deleteInstance(id)}
-														class="text-xs px-2 py-1 font-mono tracking-wider uppercase border border-red-500/30 text-red-400 hover:bg-red-500/20 hover:text-red-400 hover:border-red-500/50 transition-all duration-200 cursor-pointer"
-													>
-														DELETE
-													</button>
-												</div>
+												<button 
+													onclick={() => deleteInstance(id)}
+													class="text-xs px-2 py-1 font-mono tracking-wider uppercase border border-red-500/30 text-red-400 hover:bg-red-500/20 hover:text-red-400 hover:border-red-500/50 transition-all duration-200 cursor-pointer"
+												>
+													DELETE
+												</button>
 												</div>
 												<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"
@@ -2059,120 +1978,4 @@ function toggleInstanceDownloadDetails(nodeId: string): void {
 		{/if}
 	</main>
 
-	<!-- Draft Model Edit Modal -->
-	{#if editingDraftInstanceId}
-		<!-- svelte-ignore a11y_no_static_element_interactions -->
-		<div
-			class="fixed inset-0 z-50 flex items-center justify-center bg-black/70 backdrop-blur-sm"
-			onclick={closeDraftModelEdit}
-			onkeydown={(e) => e.key === 'Escape' && closeDraftModelEdit()}
-		>
-			<!-- svelte-ignore a11y_click_events_have_key_events -->
-			<div
-				class="bg-exo-dark-gray border border-exo-medium-gray/50 rounded-lg shadow-2xl p-6 w-full max-w-md mx-4"
-				onclick={(e) => e.stopPropagation()}
-			>
-				<div class="flex items-center justify-between mb-4">
-					<h3 class="text-lg font-mono text-exo-yellow tracking-wide">Speculative Decoding</h3>
-					<button
-						onclick={closeDraftModelEdit}
-						class="text-white/60 hover:text-white transition-colors cursor-pointer"
-						aria-label="Close"
-					>
-						<svg class="w-5 h-5" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2">
-							<path stroke-linecap="round" stroke-linejoin="round" d="M6 18L18 6M6 6l12 12" />
-						</svg>
-					</button>
-				</div>
-
-				<p class="text-white/60 text-sm font-mono mb-4">
-					Configure a draft model for faster generation. The draft model proposes tokens that the main model verifies.
-				</p>
-
-				<!-- Draft Model Dropdown -->
-				<div class="mb-4">
-					<div class="text-xs text-white/70 font-mono mb-2">Draft Model:</div>
-					<div class="relative">
-						<button
-							onclick={() => { isDraftEditDropdownOpen = !isDraftEditDropdownOpen; draftEditDropdownSearch = ''; }}
-							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 {editDraftModel ? 'bg-transparent text-cyan-400 border-cyan-500/50' : 'bg-transparent text-white/50 border-exo-medium-gray/50 hover:border-cyan-500/50'}"
-						>
-							<span class="truncate">{editDraftModel ? editDraftModel.split('/').pop() : 'None'}</span>
-							<svg class="w-4 h-4 flex-shrink-0 transition-transform {isDraftEditDropdownOpen ? '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 isDraftEditDropdownOpen}
-							<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={draftEditDropdownSearch}
-										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-cyan-500/50"
-									/>
-								</div>
-								<div class="overflow-y-auto max-h-36">
-									<!-- None option -->
-									<button
-										onclick={() => { editDraftModel = null; isDraftEditDropdownOpen = false; }}
-										class="w-full px-3 py-2 text-left text-sm font-mono tracking-wide transition-colors duration-100 flex items-center gap-2 {editDraftModel === null ? 'bg-transparent text-cyan-400 cursor-pointer' : 'text-white/80 hover:text-cyan-400 cursor-pointer'}"
-									>
-										<span>None (Disable)</span>
-									</button>
-									{#each models.filter(m => (m.name ?? m.id).toLowerCase().includes(draftEditDropdownSearch.toLowerCase())) as model}
-										{@const sizeGB = (model.storage_size_megabytes ?? 0) / 1024}
-										{@const modelHfId = model.hugging_face_id ?? model.id}
-										<button
-											onclick={() => { editDraftModel = modelHfId; isDraftEditDropdownOpen = false; }}
-											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 {editDraftModel === modelHfId ? 'bg-transparent text-cyan-400 cursor-pointer' : 'text-white/80 hover:text-cyan-400 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 -->
-				{#if editDraftModel}
-					<div class="mb-6">
-						<div class="text-xs text-white/70 font-mono mb-2">Draft Tokens per Iteration:</div>
-						<div class="flex items-center gap-2">
-							{#each [2, 3, 4, 5, 6] as n}
-								<button
-									onclick={() => editNumDraftTokens = n}
-									class="w-8 h-8 text-sm font-mono rounded transition-all {editNumDraftTokens === n ? 'bg-cyan-500/20 text-cyan-400 border border-cyan-500/50' : 'text-white/50 hover:text-white/80 border border-exo-medium-gray/50 hover:border-white/30'} cursor-pointer"
-								>{n}</button>
-							{/each}
-						</div>
-					</div>
-				{/if}
-
-				<!-- Action Buttons -->
-				<div class="flex items-center justify-end gap-3">
-					<button
-						onclick={closeDraftModelEdit}
-						class="px-4 py-2 text-sm font-mono text-white/70 hover:text-white transition-colors cursor-pointer"
-					>
-						Cancel
-					</button>
-					<button
-						onclick={saveDraftModel}
-						disabled={isSavingDraftModel}
-						class="px-4 py-2 text-sm font-mono border border-cyan-500/50 text-cyan-400 hover:bg-cyan-500/20 hover:border-cyan-500 transition-all disabled:opacity-50 disabled:cursor-not-allowed cursor-pointer"
-					>
-						{isSavingDraftModel ? 'Saving...' : 'Save'}
-					</button>
-				</div>
-			</div>
-		</div>
-	{/if}
 </div>

pyproject.toml

@@ -23,7 +23,9 @@ dependencies = [
     "tiktoken>=0.12.0", # required for kimi k2 tokenizer
     "hypercorn>=0.18.0",
     "openai-harmony>=0.0.8",
-    "httpx>=0.28.1",
+    "openhands-sdk>=0.1.0", # for exo-eval SWE-bench evaluation
+    "openhands-tools>=0.1.0", # tools for openhands agents
+    "datasets>=3.0.0", # for loading SWE-bench from HuggingFace
 ]
 
 [project.scripts]
@@ -126,6 +128,3 @@ env = [
   "EXO_TESTS=1"
 ]
 addopts = "-m 'not slow'"
-filterwarnings = [
-    "ignore:builtin type Swig:DeprecationWarning",
-]

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,14 +1,13 @@
 import time
 from collections.abc import AsyncGenerator
-from http import HTTPStatus
-from typing import cast
+from typing import Any, 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
@@ -30,8 +29,6 @@ from exo.shared.types.api import (
     CreateInstanceParams,
     CreateInstanceResponse,
     DeleteInstanceResponse,
-    ErrorInfo,
-    ErrorResponse,
     FinishReason,
     GenerationStats,
     ModelList,
@@ -39,8 +36,6 @@ from exo.shared.types.api import (
     PlaceInstanceParams,
     PlacementPreview,
     PlacementPreviewResponse,
-    SetDraftModelParams,
-    SetDraftModelResponse,
     StreamingChoiceResponse,
 )
 from exo.shared.types.chunks import TokenChunk
@@ -51,16 +46,10 @@ from exo.shared.types.commands import (
     DeleteInstance,
     ForwarderCommand,
     PlaceInstance,
-    SetInstanceDraftModel,
     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
@@ -83,7 +72,13 @@ def chunk_to_response(
         choices=[
             StreamingChoiceResponse(
                 index=0,
-                delta=ChatCompletionMessage(role="assistant", content=chunk.text),
+                delta=ChatCompletionMessage(
+                    role="assistant",
+                    content=chunk.text if chunk.text else None,
+                    tool_calls=[tc.model_dump() for tc in chunk.tool_calls]
+                    if chunk.tool_calls
+                    else None,
+                ),
                 finish_reason=chunk.finish_reason,
             )
         ],
@@ -126,7 +121,6 @@ class API:
         self.paused_ev: anyio.Event = anyio.Event()
 
         self.app = FastAPI()
-        self._setup_exception_handlers()
         self._setup_cors()
         self._setup_routes()
 
@@ -157,20 +151,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,
@@ -188,7 +168,6 @@ class API:
         self.app.get("/instance/previews")(self.get_placement_previews)
         self.app.get("/instance/{instance_id}")(self.get_instance)
         self.app.delete("/instance/{instance_id}")(self.delete_instance)
-        self.app.put("/instance/{instance_id}/draft_model")(self.set_draft_model)
         self.app.get("/models")(self.get_models)
         self.app.get("/v1/models")(self.get_models)
         self.app.post("/v1/chat/completions", response_model=None)(
@@ -204,8 +183,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)
 
@@ -402,24 +379,6 @@ class API:
             instance_id=instance_id,
         )
 
-    async def set_draft_model(
-        self, instance_id: InstanceId, payload: SetDraftModelParams
-    ) -> SetDraftModelResponse:
-        if instance_id not in self.state.instances:
-            raise HTTPException(status_code=404, detail="Instance not found")
-
-        command = SetInstanceDraftModel(
-            instance_id=instance_id,
-            draft_model=payload.draft_model,
-            num_draft_tokens=payload.num_draft_tokens,
-        )
-        await self._send(command)
-        return SetDraftModelResponse(
-            message="Command received.",
-            command_id=command.command_id,
-            instance_id=instance_id,
-        )
-
     async def _chat_chunk_stream(
         self, command_id: CommandId
     ) -> AsyncGenerator[TokenChunk, None]:
@@ -453,18 +412,6 @@ class API:
         """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
-
             chunk_response: ChatCompletionResponse = chunk_to_response(
                 chunk, command_id
             )
@@ -483,19 +430,28 @@ class API:
         text_parts: list[str] = []
         model: str | None = None
         finish_reason: FinishReason | None = None
+        all_tool_calls: list[dict[str, Any]] = []
 
         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",
-                )
-
             if model is None:
                 model = chunk.model
 
             text_parts.append(chunk.text)
 
+            # Collect tool calls
+            if chunk.tool_calls:
+                for tc in chunk.tool_calls:
+                    all_tool_calls.append(
+                        {
+                            "id": tc.id,
+                            "type": tc.type,
+                            "function": {
+                                "name": tc.function.name,
+                                "arguments": tc.function.arguments,
+                            },
+                        }
+                    )
+
             if chunk.finish_reason is not None:
                 finish_reason = chunk.finish_reason
 
@@ -511,7 +467,8 @@ class API:
                     index=0,
                     message=ChatCompletionMessage(
                         role="assistant",
-                        content=combined_text,
+                        content=combined_text if combined_text else None,
+                        tool_calls=all_tool_calls if all_tool_calls else None,
                     ),
                     finish_reason=finish_reason,
                 )
@@ -524,22 +481,31 @@ class API:
         text_parts: list[str] = []
         model: str | None = None
         finish_reason: FinishReason | None = None
+        all_tool_calls: list[dict[str, Any]] = []
 
         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",
-                )
-
             if model is None:
                 model = chunk.model
 
             text_parts.append(chunk.text)
             stats = chunk.stats or stats
 
+            # Collect tool calls
+            if chunk.tool_calls:
+                for tc in chunk.tool_calls:
+                    all_tool_calls.append(
+                        {
+                            "id": tc.id,
+                            "type": tc.type,
+                            "function": {
+                                "name": tc.function.name,
+                                "arguments": tc.function.arguments,
+                            },
+                        }
+                    )
+
             if chunk.finish_reason is not None:
                 finish_reason = chunk.finish_reason
 
@@ -554,7 +520,9 @@ class API:
                 ChatCompletionChoice(
                     index=0,
                     message=ChatCompletionMessage(
-                        role="assistant", content=combined_text
+                        role="assistant",
+                        content=combined_text if combined_text else None,
+                        tool_calls=all_tool_calls if all_tool_calls else None,
                     ),
                     finish_reason=finish_reason,
                 )
@@ -678,14 +646,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

@@ -18,7 +18,6 @@ from exo.shared.types.commands import (
     ForwarderCommand,
     PlaceInstance,
     RequestEventLog,
-    SetInstanceDraftModel,
     TaskFinished,
     TestCommand,
 )
@@ -28,7 +27,6 @@ from exo.shared.types.events import (
     ForwarderEvent,
     IndexedEvent,
     InstanceDeleted,
-    InstanceDraftModelUpdated,
     NodeTimedOut,
     TaskCreated,
     TaskDeleted,
@@ -175,14 +173,6 @@ class Master:
                                 self.state.instances, placement
                             )
                             generated_events.extend(transition_events)
-                        case SetInstanceDraftModel():
-                            generated_events.append(
-                                InstanceDraftModelUpdated(
-                                    instance_id=command.instance_id,
-                                    draft_model=command.draft_model,
-                                    num_draft_tokens=command.num_draft_tokens,
-                                )
-                            )
                         case TaskFinished():
                             generated_events.append(
                                 TaskDeleted(

src/exo/master/placement.py

@@ -3,6 +3,8 @@ from collections.abc import Mapping
 from copy import deepcopy
 from typing import Sequence
 
+from loguru import logger
+
 from exo.master.placement_utils import (
     filter_cycles_by_memory,
     get_mlx_ibv_devices_matrix,
@@ -53,6 +55,7 @@ def place_instance(
 ) -> dict[InstanceId, Instance]:
     all_nodes = list(topology.list_nodes())
 
+    logger.info("finding cycles:")
     cycles = topology.get_cycles()
     singleton_cycles = [[node] for node in all_nodes]
     candidate_cycles = list(
@@ -125,6 +128,10 @@ def place_instance(
     target_instances = dict(deepcopy(current_instances))
 
     if len(selected_cycle) == 1:
+        logger.warning(
+            "You have likely selected ibv for a single node instance; falling back to MlxRing"
+        )
+
         command.instance_meta = InstanceMeta.MlxRing
 
     # TODO: Single node instances
@@ -144,8 +151,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()
@@ -159,8 +164,6 @@ 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,
             )
 
     return target_instances

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/shared/apply.py

@@ -11,7 +11,6 @@ from exo.shared.types.events import (
     IndexedEvent,
     InstanceCreated,
     InstanceDeleted,
-    InstanceDraftModelUpdated,
     NodeCreated,
     NodeDownloadProgress,
     NodeMemoryMeasured,
@@ -48,8 +47,6 @@ def event_apply(event: Event, state: State) -> State:
             return apply_instance_created(event, state)
         case InstanceDeleted():
             return apply_instance_deleted(event, state)
-        case InstanceDraftModelUpdated():
-            return apply_instance_draft_model_updated(event, state)
         case NodeCreated():
             return apply_topology_node_created(event, state)
         case NodeTimedOut():
@@ -172,25 +169,6 @@ def apply_instance_deleted(event: InstanceDeleted, state: State) -> State:
     return state.model_copy(update={"instances": new_instances})
 
 
-def apply_instance_draft_model_updated(
-    event: InstanceDraftModelUpdated, state: State
-) -> State:
-    if event.instance_id not in state.instances:
-        return state
-    instance = state.instances[event.instance_id]
-    updated_instance = instance.model_copy(
-        update={
-            "draft_model": event.draft_model,
-            "num_draft_tokens": event.num_draft_tokens,
-        }
-    )
-    new_instances: Mapping[InstanceId, Instance] = {
-        **state.instances,
-        event.instance_id: updated_instance,
-    }
-    return state.model_copy(update={"instances": new_instances})
-
-
 def apply_runner_status_updated(event: RunnerStatusUpdated, state: State) -> State:
     new_runners: Mapping[RunnerId, RunnerStatus] = {
         **state.runners,

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/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
@@ -217,14 +202,3 @@ class DeleteInstanceResponse(BaseModel):
     message: str
     command_id: CommandId
     instance_id: InstanceId
-
-
-class SetDraftModelParams(BaseModel):
-    draft_model: ModelId | None = None  # None to disable speculative decoding
-    num_draft_tokens: int = 4
-
-
-class SetDraftModelResponse(BaseModel):
-    message: str
-    command_id: CommandId
-    instance_id: InstanceId

src/exo/shared/types/chunks.py

@@ -1,4 +1,7 @@
 from enum import Enum
+from typing import Literal
+
+from pydantic import BaseModel
 
 from exo.shared.types.api import GenerationStats
 from exo.utils.pydantic_ext import TaggedModel
@@ -12,6 +15,17 @@ class ChunkType(str, Enum):
     Image = "Image"
 
 
+class ToolCallFunction(BaseModel, frozen=True):
+    name: str
+    arguments: str
+
+
+class ToolCall(BaseModel, frozen=True):
+    id: str
+    type: Literal["function"] = "function"
+    function: ToolCallFunction
+
+
 class BaseChunk(TaggedModel):
     idx: int
     model: ModelId
@@ -22,7 +36,7 @@ class TokenChunk(BaseChunk):
     token_id: int
     finish_reason: FinishReason | None = None
     stats: GenerationStats | None = None
-    error_message: str | None = None
+    tool_calls: list[ToolCall] | 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,14 +35,6 @@ class DeleteInstance(BaseCommand):
     instance_id: InstanceId
 
 
-class SetInstanceDraftModel(BaseCommand):
-    """Set or update the draft model for an existing instance."""
-
-    instance_id: InstanceId
-    draft_model: ModelId | None  # None to disable speculative decoding
-    num_draft_tokens: int = 4
-
-
 class TaskFinished(BaseCommand):
     finished_command_id: CommandId
 
@@ -60,7 +50,6 @@ Command = (
     | PlaceInstance
     | CreateInstance
     | DeleteInstance
-    | SetInstanceDraftModel
     | TaskFinished
 )
 

src/exo/shared/types/events.py

@@ -5,7 +5,6 @@ from pydantic import Field
 from exo.shared.topology import Connection, NodePerformanceProfile
 from exo.shared.types.chunks import GenerationChunk
 from exo.shared.types.common import CommandId, Id, NodeId, SessionId
-from exo.shared.types.models import ModelId
 from exo.shared.types.profiling import MemoryPerformanceProfile
 from exo.shared.types.tasks import Task, TaskId, TaskStatus
 from exo.shared.types.worker.downloads import DownloadProgress
@@ -68,14 +67,6 @@ class InstanceDeleted(BaseEvent):
     instance_id: InstanceId
 
 
-class InstanceDraftModelUpdated(BaseEvent):
-    """Draft model updated on an existing instance."""
-
-    instance_id: InstanceId
-    draft_model: ModelId | None
-    num_draft_tokens: int
-
-
 class RunnerStatusUpdated(BaseEvent):
     runner_id: RunnerId
     runner_status: RunnerStatus
@@ -132,7 +123,6 @@ Event = (
     | TaskAcknowledged
     | InstanceCreated
     | InstanceDeleted
-    | InstanceDraftModelUpdated
     | RunnerStatusUpdated
     | RunnerDeleted
     | NodeCreated

src/exo/shared/types/tasks.py

@@ -36,12 +36,6 @@ class DownloadModel(BaseTask):  # emitted by Worker
     shard_metadata: ShardMetadata
 
 
-class DownloadDraftModel(BaseTask):  # emitted by Worker
-    """Download a draft model for speculative decoding (rank 0 only)."""
-
-    model_id: str  # HuggingFace model ID
-
-
 class LoadModel(BaseTask):  # emitted by Worker
     pass
 
@@ -66,21 +60,12 @@ class Shutdown(BaseTask):  # emitted by Worker
     runner_id: RunnerId
 
 
-class SetDraftModel(BaseTask):  # emitted by Worker
-    """Load or clear a draft model on an already-running instance."""
-
-    model_id: str | None  # HuggingFace model ID, or None to clear
-    num_draft_tokens: int = 4
-
-
 Task = (
     CreateRunner
     | DownloadModel
-    | DownloadDraftModel
     | ConnectToGroup
     | LoadModel
     | StartWarmup
     | ChatCompletion
     | Shutdown
-    | SetDraftModel
 )

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
 
@@ -20,8 +19,6 @@ class InstanceMeta(str, Enum):
 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)

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

@@ -1,4 +1,5 @@
 from exo.shared.types.api import FinishReason, GenerationStats
+from exo.shared.types.chunks import ToolCall
 from exo.utils.pydantic_ext import TaggedModel
 
 
@@ -16,6 +17,7 @@ class GenerationResponse(BaseRunnerResponse):
     # logprobs: list[float] | None = None # too big. we can change to be top-k
     finish_reason: FinishReason | None = None
     stats: GenerationStats | None = None
+    tool_calls: list[ToolCall] | None = None
 
 
 class FinishedResponse(BaseRunnerResponse):

src/exo/worker/download/download_utils.py

@@ -245,15 +245,12 @@ def create_http_session(
         sock_read_timeout = 1800
         sock_connect_timeout = 60
 
-    ssl_context = ssl.create_default_context(
-        cafile=os.getenv("SSL_CERT_FILE") or certifi.where()
-    )
+    ssl_context = ssl.create_default_context(cafile=certifi.where())
     connector = aiohttp.TCPConnector(ssl=ssl_context)
 
     return aiohttp.ClientSession(
         auto_decompress=auto_decompress,
         connector=connector,
-        proxy=os.getenv("HTTPS_PROXY") or os.getenv("HTTP_PROXY") or None,
         timeout=aiohttp.ClientTimeout(
             total=total_timeout,
             connect=connect_timeout,

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

@@ -48,8 +48,6 @@ def maybe_quantize_kv_cache(
 def warmup_inference(
     model: Model,
     tokenizer: TokenizerWrapper,
-    draft_model: Model | None = None,
-    num_draft_tokens: int = 4,
 ) -> int:
     content = "Prompt to warm up the inference engine. Repeat this."
 
@@ -68,30 +66,25 @@ def warmup_inference(
 
     tokens_generated = 0
 
+    cache = make_kv_cache(
+        model=model,
+    )
+
     # Use a default sampler for warmup
     sampler = make_sampler(temp=0.7)
 
-    generate_kwargs: dict[str, object] = {
-        "model": model,
-        "tokenizer": tokenizer,
-        "prompt": warmup_prompt,
-        "max_tokens": 50,
-        "sampler": sampler,
-        "prefill_step_size": 2048,
-        "kv_group_size": KV_GROUP_SIZE,
-        "kv_bits": KV_BITS,
-    }
-
-    # Warm up with draft model if provided (speculative decoding path)
-    if draft_model is not None:
-        logger.info("Warming up with speculative decoding (draft model)")
-        generate_kwargs["draft_model"] = draft_model
-        generate_kwargs["num_draft_tokens"] = num_draft_tokens
-    else:
-        generate_kwargs["prompt_cache"] = make_kv_cache(model=model)
-
     logger.info("Generating warmup tokens")
-    for _r in stream_generate(**generate_kwargs):  # type: ignore[arg-type]
+    for _r in stream_generate(
+        model=model,
+        tokenizer=tokenizer,
+        prompt=warmup_prompt,
+        max_tokens=50,
+        sampler=sampler,
+        prompt_cache=cache,
+        prefill_step_size=2048,
+        kv_group_size=KV_GROUP_SIZE,
+        kv_bits=KV_BITS,
+    ):
         logger.info("Generated warmup token: " + str(_r.text))
         tokens_generated += 1
 
@@ -126,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()
@@ -144,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
@@ -156,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
 
@@ -84,45 +82,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(
@@ -229,9 +188,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 +202,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 +213,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 +252,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)

src/exo/worker/main.py

@@ -29,9 +29,7 @@ from exo.shared.types.profiling import MemoryPerformanceProfile, NodePerformance
 from exo.shared.types.state import State
 from exo.shared.types.tasks import (
     CreateRunner,
-    DownloadDraftModel,
     DownloadModel,
-    SetDraftModel,
     Shutdown,
     Task,
     TaskStatus,
@@ -50,7 +48,6 @@ from exo.utils.event_buffer import OrderedBuffer
 from exo.worker.download.download_utils import (
     map_repo_download_progress_to_download_progress_data,
 )
-from exo.worker.download.impl_shard_downloader import build_full_shard
 from exo.worker.download.shard_downloader import RepoDownloadProgress, ShardDownloader
 from exo.worker.plan import plan
 from exo.worker.runner.runner_supervisor import RunnerSupervisor
@@ -205,10 +202,42 @@ class Worker:
                         )
                     )
                 case DownloadModel(shard_metadata=shard):
-                    await self._handle_download(shard, task)
-                case DownloadDraftModel(model_id=model_id):
-                    shard = await build_full_shard(model_id)
-                    await self._handle_download(shard, task)
+                    if shard.model_meta.model_id not in self.download_status:
+                        progress = DownloadPending(
+                            shard_metadata=shard, node_id=self.node_id
+                        )
+                        self.download_status[shard.model_meta.model_id] = progress
+                        await self.event_sender.send(
+                            NodeDownloadProgress(download_progress=progress)
+                        )
+                    initial_progress = (
+                        await self.shard_downloader.get_shard_download_status_for_shard(
+                            shard
+                        )
+                    )
+                    if initial_progress.status == "complete":
+                        progress = DownloadCompleted(
+                            shard_metadata=shard,
+                            node_id=self.node_id,
+                            total_bytes=initial_progress.total_bytes,
+                        )
+                        self.download_status[shard.model_meta.model_id] = progress
+                        await self.event_sender.send(
+                            NodeDownloadProgress(download_progress=progress)
+                        )
+                        await self.event_sender.send(
+                            TaskStatusUpdated(
+                                task_id=task.task_id,
+                                task_status=TaskStatus.Complete,
+                            )
+                        )
+                    else:
+                        await self.event_sender.send(
+                            TaskStatusUpdated(
+                                task_id=task.task_id, task_status=TaskStatus.Running
+                            )
+                        )
+                        self._handle_shard_download_process(task, initial_progress)
                 case Shutdown(runner_id=runner_id):
                     try:
                         with fail_after(3):
@@ -219,25 +248,6 @@ class Worker:
                                 task_id=task.task_id, task_status=TaskStatus.TimedOut
                             )
                         )
-                case SetDraftModel(
-                    model_id=draft_model_id, num_draft_tokens=num_tokens
-                ):
-                    runner = self.runners[self._task_to_runner_id(task)]
-                    await runner.start_task(task)
-                    # Update bound_instance to reflect new/cleared draft model
-                    updated_instance = runner.bound_instance.instance.model_copy(
-                        update={
-                            "draft_model": (
-                                ModelId(draft_model_id)
-                                if draft_model_id is not None
-                                else None
-                            ),
-                            "num_draft_tokens": num_tokens,
-                        }
-                    )
-                    runner.bound_instance = runner.bound_instance.model_copy(
-                        update={"instance": updated_instance}
-                    )
                 case task:
                     await self.runners[self._task_to_runner_id(task)].start_task(task)
 
@@ -330,46 +340,6 @@ class Worker:
         self._tg.start_soon(runner.run)
         return runner
 
-    async def _handle_download(self, shard: ShardMetadata, task: Task) -> None:
-        """Handle model download - shared logic for main and draft models."""
-        model_id = shard.model_meta.model_id
-
-        if model_id not in self.download_status:
-            progress = DownloadPending(shard_metadata=shard, node_id=self.node_id)
-            self.download_status[model_id] = progress
-            await self.event_sender.send(
-                NodeDownloadProgress(download_progress=progress)
-            )
-
-        initial_progress = (
-            await self.shard_downloader.get_shard_download_status_for_shard(shard)
-        )
-
-        if initial_progress.status == "complete":
-            progress = DownloadCompleted(
-                shard_metadata=shard,
-                node_id=self.node_id,
-                total_bytes=initial_progress.total_bytes,
-            )
-            self.download_status[model_id] = progress
-            await self.event_sender.send(
-                NodeDownloadProgress(download_progress=progress)
-            )
-            await self.event_sender.send(
-                TaskStatusUpdated(task_id=task.task_id, task_status=TaskStatus.Complete)
-            )
-        else:
-            await self.event_sender.send(
-                TaskStatusUpdated(task_id=task.task_id, task_status=TaskStatus.Running)
-            )
-            download_task = DownloadModel(
-                instance_id=task.instance_id,
-                shard_metadata=shard,
-                task_id=task.task_id,
-                task_status=task.task_status,
-            )
-            self._handle_shard_download_process(download_task, initial_progress)
-
     def _handle_shard_download_process(
         self,
         task: DownloadModel,

src/exo/worker/plan.py

@@ -8,10 +8,8 @@ from exo.shared.types.tasks import (
     ChatCompletion,
     ConnectToGroup,
     CreateRunner,
-    DownloadDraftModel,
     DownloadModel,
     LoadModel,
-    SetDraftModel,
     Shutdown,
     StartWarmup,
     Task,
@@ -40,16 +38,6 @@ from exo.shared.types.worker.runners import (
 from exo.worker.runner.runner_supervisor import RunnerSupervisor
 
 
-def _is_download_in_progress_or_complete(
-    model_id: ModelId,
-    download_status: Mapping[ModelId, DownloadProgress],
-) -> bool:
-    """Check if model download is in progress or complete."""
-    return model_id in download_status and isinstance(
-        download_status[model_id], (DownloadOngoing, DownloadCompleted)
-    )
-
-
 def plan(
     node_id: NodeId,
     # Runners is expected to be FRESH and so should not come from state
@@ -67,11 +55,9 @@ def plan(
         _kill_runner(runners, all_runners, instances)
         or _create_runner(node_id, runners, instances)
         or _model_needs_download(runners, download_status)
-        or _draft_model_needs_download(runners, download_status, instances)
         or _init_distributed_backend(runners, all_runners)
-        or _load_model(runners, all_runners, global_download_status, download_status)
+        or _load_model(runners, all_runners, global_download_status)
         or _ready_to_warmup(runners, all_runners)
-        or _set_draft_model(runners, instances, download_status)
         or _pending_tasks(runners, tasks, all_runners)
     )
 
@@ -129,9 +115,12 @@ def _model_needs_download(
 ) -> DownloadModel | None:
     for runner in runners.values():
         model_id = runner.bound_instance.bound_shard.model_meta.model_id
-        if isinstance(
-            runner.status, RunnerIdle
-        ) and not _is_download_in_progress_or_complete(model_id, download_status):
+        if isinstance(runner.status, RunnerIdle) and (
+            model_id not in download_status
+            or not isinstance(
+                download_status[model_id], (DownloadOngoing, DownloadCompleted)
+            )
+        ):
             # We don't invalidate download_status randomly in case a file gets deleted on disk
             return DownloadModel(
                 instance_id=runner.bound_instance.instance.instance_id,
@@ -139,43 +128,6 @@ def _model_needs_download(
             )
 
 
-def _draft_model_needs_download(
-    runners: Mapping[RunnerId, RunnerSupervisor],
-    download_status: Mapping[ModelId, DownloadProgress],
-    instances: Mapping[InstanceId, Instance],
-) -> DownloadDraftModel | None:
-    """Check if draft model needs download for rank 0 runner.
-
-    Triggers download when:
-    - RunnerIdle with draft model (initial setup)
-    - RunnerReady with new draft model (updated via API)
-    """
-    rank_0_runner = next(
-        (r for r in runners.values() if r.bound_instance.bound_shard.device_rank == 0),
-        None,
-    )
-    if rank_0_runner is None:
-        return None
-    if not isinstance(rank_0_runner.status, (RunnerIdle, RunnerReady)):
-        return None
-
-    # Use current instance state (may have been updated via API)
-    instance_id = rank_0_runner.bound_instance.instance.instance_id
-    current_instance = instances.get(instance_id)
-    if current_instance is None:
-        return None
-
-    draft_model_id = current_instance.draft_model
-    if draft_model_id is None:
-        return None
-    if _is_download_in_progress_or_complete(draft_model_id, download_status):
-        return None
-    return DownloadDraftModel(
-        instance_id=instance_id,
-        model_id=str(draft_model_id),
-    )
-
-
 def _init_distributed_backend(
     runners: Mapping[RunnerId, RunnerSupervisor],
     all_runners: Mapping[RunnerId, RunnerStatus],
@@ -230,12 +182,10 @@ def _load_model(
     runners: Mapping[RunnerId, RunnerSupervisor],
     all_runners: Mapping[RunnerId, RunnerStatus],
     global_download_status: Mapping[NodeId, Sequence[DownloadProgress]],
-    download_status: Mapping[ModelId, DownloadProgress],
 ) -> LoadModel | None:
     for runner in runners.values():
         instance = runner.bound_instance.instance
         shard_assignments = instance.shard_assignments
-        shard = runner.bound_instance.bound_shard
 
         all_local_downloads_complete = all(
             nid in global_download_status
@@ -249,14 +199,6 @@ def _load_model(
         if not all_local_downloads_complete:
             continue
 
-        # Rank 0 with draft model must wait for draft download before loading
-        if shard.device_rank == 0:
-            draft_model_id = instance.draft_model
-            if draft_model_id is not None and not isinstance(
-                download_status.get(draft_model_id), DownloadCompleted
-            ):
-                continue
-
         is_single_node_instance = len(instance.shard_assignments.runner_to_shard) == 1
         if is_single_node_instance and isinstance(runner.status, RunnerIdle):
             return LoadModel(instance_id=instance.instance_id)
@@ -316,53 +258,6 @@ def _ready_to_warmup(
     return None
 
 
-def _set_draft_model(
-    runners: Mapping[RunnerId, RunnerSupervisor],
-    instances: Mapping[InstanceId, Instance],
-    download_status: Mapping[ModelId, DownloadProgress],
-) -> SetDraftModel | None:
-    """Check if rank 0 runner needs to load or clear a draft model."""
-    rank_0_runner = next(
-        (r for r in runners.values() if r.bound_instance.bound_shard.device_rank == 0),
-        None,
-    )
-    if rank_0_runner is None:
-        return None
-    if not isinstance(rank_0_runner.status, RunnerReady):
-        return None
-
-    instance_id = rank_0_runner.bound_instance.instance.instance_id
-    current_instance = instances.get(instance_id)
-    if current_instance is None:
-        return None
-
-    # Compare runner's bound draft model vs current instance draft model
-    runner_draft_model = rank_0_runner.bound_instance.instance.draft_model
-    current_draft_model = current_instance.draft_model
-
-    if runner_draft_model == current_draft_model:
-        return None
-
-    # Draft model changed - need to update
-    if current_draft_model is None:
-        # Clear draft model
-        return SetDraftModel(
-            instance_id=instance_id,
-            model_id=None,
-            num_draft_tokens=4,
-        )
-
-    # Wait for draft model to be downloaded
-    if not isinstance(download_status.get(current_draft_model), DownloadCompleted):
-        return None
-
-    return SetDraftModel(
-        instance_id=instance_id,
-        model_id=str(current_draft_model),
-        num_draft_tokens=current_instance.num_draft_tokens,
-    )
-
-
 def _pending_tasks(
     runners: Mapping[RunnerId, RunnerSupervisor],
     tasks: Mapping[TaskId, Task],

src/exo/worker/runner/bootstrap.py

@@ -17,23 +17,15 @@ 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
-        )
+    if (
+        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
     try:
         from exo.worker.runner.runner import main

src/exo/worker/runner/runner.py

@@ -1,11 +1,13 @@
+import json
 import time
 from collections.abc import Generator
-from contextlib import contextmanager
 from functools import cache
-from typing import cast
+from typing import Any
+from uuid import uuid4
 
 import mlx.core as mx
 from mlx_lm.models.gpt_oss import Model as GptOssModel
+from mlx_lm.tokenizer_utils import TokenizerWrapper
 from openai_harmony import (  # pyright: ignore[reportMissingTypeStubs]
     HarmonyEncodingName,
     Role,
@@ -14,8 +16,7 @@ from openai_harmony import (  # pyright: ignore[reportMissingTypeStubs]
 )
 
 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.chunks import TokenChunk, ToolCall, ToolCallFunction
 from exo.shared.types.events import (
     ChunkGenerated,
     Event,
@@ -23,12 +24,10 @@ from exo.shared.types.events import (
     TaskAcknowledged,
     TaskStatusUpdated,
 )
-from exo.shared.types.models import ModelId
 from exo.shared.types.tasks import (
     ChatCompletion,
     ConnectToGroup,
     LoadModel,
-    SetDraftModel,
     Shutdown,
     StartWarmup,
     Task,
@@ -53,44 +52,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],
@@ -112,7 +82,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")
@@ -153,30 +122,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
-                    )
-
-                    # 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))
-                        )
+                    model, tokenizer = load_mlx_items(bound_instance, group)
 
                     current_status = RunnerLoaded()
                     logger.info("runner loaded")
@@ -193,10 +139,9 @@ def main(
 
                     logger.info(f"warming up inference for instance: {instance}")
                     toks = warmup_inference(
-                        model=cast(Model, model),
+                        model=model,
                         tokenizer=tokenizer,
-                        draft_model=draft_model,
-                        num_draft_tokens=instance.num_draft_tokens,
+                        # kv_prefix_cache=kv_prefix_cache,  # supply for warmup-time prefix caching
                     )
                     logger.info(f"warmed up by generating {toks} tokens")
                     logger.info(
@@ -207,6 +152,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")
@@ -215,96 +162,48 @@ 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 model
-                        assert tokenizer
-                        assert task_params.messages[0].content is not None
-                        _check_for_debug_prompts(task_params.messages[0].content)
+                    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,
-                        )
+                    # Generate responses using the actual MLX generation
+                    mlx_generator = mlx_generate(
+                        model=model,
+                        tokenizer=tokenizer,
+                        task=task_params,
+                    )
 
-                        # GPT-OSS specific parsing to match other model formats.
-                        if isinstance(model, GptOssModel):
-                            mlx_generator = parse_gpt_oss(mlx_generator)
+                    # 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
+                    # Parse tool calls to place them in the tool calls section
+                    mlx_generator = parse_tool_calls(
+                        mlx_generator, tokenizer, task_params.tools
+                    )
 
-                        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,
-                                                ),
-                                            )
+                    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,
+                                                tool_calls=response.tool_calls,
+                                            ),
                                         )
+                                    )
+                                # case TokenizedResponse():
+                                # TODO: something here ig
 
                     current_status = RunnerReady()
                     logger.info("runner ready")
-                case SetDraftModel(
-                    model_id=draft_model_id, num_draft_tokens=num_tokens
-                ) if isinstance(current_status, RunnerReady):
-                    current_status = RunnerWarmingUp()
-                    logger.info("runner warming up (setting draft model)")
-                    event_sender.send(
-                        RunnerStatusUpdated(
-                            runner_id=runner_id, runner_status=current_status
-                        )
-                    )
-                    assert model is not None
-                    assert tokenizer is not None
-
-                    if draft_model_id is None:
-                        # Clear draft model
-                        logger.info("Clearing draft model")
-                        draft_model = None
-                        instance = instance.model_copy(
-                            update={
-                                "draft_model": None,
-                                "num_draft_tokens": 4,
-                            }
-                        )
-                    else:
-                        # Load new draft model
-                        logger.info(f"Loading draft model: {draft_model_id}")
-                        draft_model = cast(Model, load_draft_model(draft_model_id))
-                        instance = instance.model_copy(
-                            update={
-                                "draft_model": ModelId(draft_model_id),
-                                "num_draft_tokens": num_tokens,
-                            }
-                        )
-                        # Warm up with speculative decoding
-                        logger.info("Warming up with new draft model")
-                        warmup_inference(
-                            model=cast(Model, model),
-                            tokenizer=tokenizer,
-                            draft_model=draft_model,
-                            num_draft_tokens=num_tokens,
-                        )
-                        logger.info("Draft model loaded and warmed up")
-
-                    current_status = RunnerReady()
                 case Shutdown():
                     current_status = RunnerShuttingDown()
                     logger.info("runner shutting down")
@@ -325,7 +224,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
 
@@ -370,6 +269,98 @@ def parse_gpt_oss(
             break
 
 
+def _generate_tool_call_id() -> str:
+    return f"call_{uuid4().hex[:24]}"
+
+
+def _parse_tool_call_content(
+    content: str,
+    tokenizer: TokenizerWrapper,
+    tools: list[dict[str, Any]] | None,
+) -> ToolCall | None:
+    content = content.strip()
+    if not content:
+        return None
+
+    tool_parser: Any = getattr(tokenizer, "tool_parser", None)
+    if tool_parser is None:
+        logger.warning("No tool_parser available for tokenizer")
+        return None
+
+    try:
+        parsed: dict[str, Any] = tool_parser(content, tools)  # pyright: ignore[reportAny]
+        if parsed and "name" in parsed:
+            arguments: Any = parsed.get("arguments", {})  # pyright: ignore[reportAny]
+            arguments_str: str = (
+                json.dumps(arguments)
+                if not isinstance(arguments, str)
+                else arguments
+            )
+            return ToolCall(
+                id=_generate_tool_call_id(),
+                type="function",
+                function=ToolCallFunction(
+                    name=str(parsed["name"]),  # pyright: ignore[reportAny]
+                    arguments=arguments_str,
+                ),
+            )
+    except Exception as e:
+        logger.warning(f"tool_parser failed: {e}")
+
+    return None
+
+
+def parse_tool_calls(
+    responses: Generator[GenerationResponse],
+    tokenizer: TokenizerWrapper,
+    tools: list[dict[str, Any]] | None,
+) -> Generator[GenerationResponse]:
+    has_tool_calling = getattr(tokenizer, "has_tool_calling", False)
+    if not has_tool_calling or tools is None:
+        yield from responses
+        return
+
+    tool_call_start: str | None = getattr(tokenizer, "tool_call_start", None)
+    tool_call_end: str | None = getattr(tokenizer, "tool_call_end", None)
+
+    if tool_call_start is None or tool_call_end is None:
+        yield from responses
+        return
+
+    in_tool_call = False
+    tool_call_buffer: list[str] = []
+    pending_tool_calls: list[ToolCall] = []
+
+    for response in responses:
+        if response.text == tool_call_start:
+            in_tool_call = True
+            tool_call_buffer = []
+            continue
+
+        if response.text == tool_call_end:
+            in_tool_call = False
+            parsed = _parse_tool_call_content(
+                "".join(tool_call_buffer), tokenizer, tools
+            )
+            if parsed is not None:
+                pending_tool_calls.append(parsed)
+            continue
+
+        if in_tool_call:
+            tool_call_buffer.append(response.text)
+            continue
+
+        if response.finish_reason is None or not pending_tool_calls:
+            yield response
+        else:
+            yield response.model_copy(
+                update={
+                    "finish_reason": "tool_calls",
+                    "tool_calls": pending_tool_calls if pending_tool_calls else None,
+                }
+            )
+
+
 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/tests/unittests/test_runner/test_event_ordering.py

@@ -121,21 +121,6 @@ def patch_out_mlx(monkeypatch: pytest.MonkeyPatch):
     monkeypatch.setattr(mlx_runner, "mlx_generate", fake_generate)
 
 
-# Use a fake event_sender to remove test flakiness.
-class EventCollector:
-    def __init__(self) -> None:
-        self.events: list[Event] = []
-
-    def send(self, event: Event) -> None:
-        self.events.append(event)
-
-    def close(self) -> None:
-        pass
-
-    def join(self) -> None:
-        pass
-
-
 def _run(tasks: Iterable[Task]):
     bound_instance = get_bound_mlx_ring_instance(
         instance_id=INSTANCE_1_ID,
@@ -145,20 +130,22 @@ def _run(tasks: Iterable[Task]):
     )
 
     task_sender, task_receiver = mp_channel[Task]()
-    event_sender = EventCollector()
+    event_sender, event_receiver = mp_channel[Event]()
 
-    with task_sender:
+    with task_sender, event_receiver:
         for t in tasks:
             task_sender.send(t)
 
         # worst monkeypatch known to man
         # this is some c++ nonsense
+        event_sender.close = nothin
+        event_sender.join = nothin
         task_receiver.close = nothin
         task_receiver.join = nothin
 
-        mlx_runner.main(bound_instance, event_sender, task_receiver)  # type: ignore[arg-type]
+        mlx_runner.main(bound_instance, event_sender, task_receiver)
 
-        return event_sender.events
+        return event_receiver.collect()
 
 
 def test_events_processed_in_correct_order(patch_out_mlx: pytest.MonkeyPatch):

src/exo/worker/utils/net_profile.py

@@ -1,64 +1,62 @@
-import anyio
-import httpx
-from anyio import create_task_group
+import http.client
+import time
+
+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
+BAD_STATUSLINE_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):
+    # TODO: use an async http client
+    def _fetch_remote_node_id(*, attempt: int = 1) -> NodeId | None:
+        connection = http.client.HTTPConnection(target_ip, 52415, timeout=3)
         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.BadStatusLine:
+            if attempt >= BAD_STATUSLINE_ATTEMPTS:
+                logger.warning(
+                    f"BadStatusLine from {target_ip}, after {attempt} attempts, assuming connection to {expected_node_id} has dropped"
+                )
+                return None
+            time.sleep(1)
+            return _fetch_remote_node_id(attempt=attempt + 1)
+        except http.client.HTTPException as e:
+            logger.warning(f"HTTPException from {target_ip}: {type(e).__name__}: {e}")
+            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 +74,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