Add exo eval

Add ChatCompletion tool calling support
https://platform.openai.com/docs/api-reference/chat/get
2026-01-27 15:33:26 -05:00 · 2026-01-16 12:55:43 +00:00 · 2026-01-16 11:36:08 +00:00
40 changed files with 5585 additions and 2190 deletions
--- a/.github/workflows/build-app.yml
+++ b/.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
--- a/.github/workflows/pipeline.yml
+++ b/.github/workflows/pipeline.yml
@@ -113,29 +113,6 @@ jobs:
        with:
          lfs: false

-      - name: Select Xcode
-        if: startsWith(matrix.runner, 'macos-')
-        run: |
-          XCODE_BASEDIR="$(printf '%s\n' /Applications/Xcode_*.app | sort -V | tail -n 1)"
-          [[ -z "$XCODE_BASEDIR" ]] && exit 1
-
-          ls -ld "/Applications/Xcode.app"
-          sudo /usr/bin/xcode-select -s "$XCODE_BASEDIR"
-          /usr/bin/xcode-select -p || true
-          /usr/bin/xcrun --toolchain default --find xcodebuild || true
-
-      - name: Install Metal toolchain component
-        if: startsWith(matrix.runner, 'macos-')
-        run: |
-          set -e
-          if ! xcrun --find metal >/dev/null 2>&1; then
-            sudo xcodebuild -downloadComponent MetalToolchain
-          fi
-          xcrun --find metal
-          xcrun --find metallib
-          echo "GH_OVERRIDE_METAL=$(xcrun --find metal)" >> $GITHUB_ENV
-          echo "GH_OVERRIDE_METALLIB=$(xcrun --find metallib)" >> $GITHUB_ENV
-
      - uses: cachix/install-nix-action@v31
        with:
          nix_path: nixpkgs=channel:nixos-unstable
@@ -147,9 +124,6 @@ jobs:
          authToken: "${{ secrets.CACHIX_AUTH_TOKEN }}"

      - name: Build all Nix outputs
-        env:
-          GH_OVERRIDE_METAL: ${{ env.GH_OVERRIDE_METAL }}
-          GH_OVERRIDE_METALLIB: ${{ env.GH_OVERRIDE_METALLIB }}
        run: |
          nix flake show --json | jq -r '
            [
--- a/AGENTS.md
+++ b/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
--- a/README.md
+++ b/README.md
@@ -27,22 +27,13 @@ exo connects all your devices into an AI cluster. Not only does exo enable runni
 - **Tensor Parallelism**: exo supports sharding models, for up to 1.8x speedup on 2 devices and 3.2x speedup on 4 devices.
 - **MLX Support**: exo uses [MLX](https://github.com/ml-explore/mlx) as an inference backend and [MLX distributed](https://ml-explore.github.io/mlx/build/html/usage/distributed.html) for distributed communication.

-## Dashboard
-
-exo includes a built-in dashboard for managing your cluster and chatting with models.
-
-<p align="center">
-  <img src="docs/imgs/dashboard-cluster-view.png" alt="exo dashboard - cluster view showing 4 x M3 Ultra Mac Studio with DeepSeek v3.1 and Kimi-K2-Thinking loaded" width="80%" />
-</p>
-<p align="center"><em>4 × 512GB M3 Ultra Mac Studio running DeepSeek v3.1 (8-bit) and Kimi-K2-Thinking (4-bit)</em></p>
-
 ## Benchmarks

 <details>
  <summary>Qwen3-235B (8-bit) on 4 × M3 Ultra Mac Studio with Tensor Parallel RDMA</summary>
  <img src="docs/benchmarks/jeffgeerling/mac-studio-cluster-ai-full-1-qwen3-235b.jpeg" alt="Benchmark - Qwen3-235B (8-bit) on 4 × M3 Ultra Mac Studio with Tensor Parallel RDMA" width="80%" />
  <p>
-    <strong>Source:</strong> <a href="https://www.jeffgeerling.com/blog/2025/15-tb-vram-on-mac-studio-rdma-over-thunderbolt-5">Jeff Geerling: 15 TB VRAM on Mac Studio – RDMA over Thunderbolt 5</a>
+    <strong>Source:</strong> <a href="https://www.jeffgeerling.com/blog/2025/15-tb-vram-on-mac-studio-rdma-over-thunderbolt-5">Jeff Geerling: 15 TB VRAM on Mac Studio – RDMA over Thunderbolt 5</a>
  </p>
 </details>

@@ -50,7 +41,7 @@ exo includes a built-in dashboard for managing your cluster and chatting with mo
  <summary>DeepSeek v3.1 671B (8-bit) on 4 × M3 Ultra Mac Studio with Tensor Parallel RDMA</summary>
  <img src="docs/benchmarks/jeffgeerling/mac-studio-cluster-ai-full-2-deepseek-3.1-671b.jpeg" alt="Benchmark - DeepSeek v3.1 671B (8-bit) on 4 × M3 Ultra Mac Studio with Tensor Parallel RDMA" width="80%" />
  <p>
-    <strong>Source:</strong> <a href="https://www.jeffgeerling.com/blog/2025/15-tb-vram-on-mac-studio-rdma-over-thunderbolt-5">Jeff Geerling: 15 TB VRAM on Mac Studio – RDMA over Thunderbolt 5</a>
+    <strong>Source:</strong> <a href="https://www.jeffgeerling.com/blog/2025/15-tb-vram-on-mac-studio-rdma-over-thunderbolt-5">Jeff Geerling: 15 TB VRAM on Mac Studio – RDMA over Thunderbolt 5</a>
  </p>
 </details>

@@ -58,7 +49,7 @@ exo includes a built-in dashboard for managing your cluster and chatting with mo
  <summary>Kimi K2 Thinking (native 4-bit) on 4 × M3 Ultra Mac Studio with Tensor Parallel RDMA</summary>
  <img src="docs/benchmarks/jeffgeerling/mac-studio-cluster-ai-full-3-kimi-k2-thinking.jpeg" alt="Benchmark - Kimi K2 Thinking (native 4-bit) on 4 × M3 Ultra Mac Studio with Tensor Parallel RDMA" width="80%" />
  <p>
-    <strong>Source:</strong> <a href="https://www.jeffgeerling.com/blog/2025/15-tb-vram-on-mac-studio-rdma-over-thunderbolt-5">Jeff Geerling: 15 TB VRAM on Mac Studio – RDMA over Thunderbolt 5</a>
+    <strong>Source:</strong> <a href="https://www.jeffgeerling.com/blog/2025/15-tb-vram-on-mac-studio-rdma-over-thunderbolt-5">Jeff Geerling: 15 TB VRAM on Mac Studio – RDMA over Thunderbolt 5</a>
  </p>
 </details>

@@ -163,24 +154,6 @@ This starts the exo dashboard and API at http://localhost:52415/

 **Important note for Linux users:** Currently, exo runs on CPU on Linux. GPU support for Linux platforms is under development. If you'd like to see support for your specific Linux hardware, please [search for existing feature requests](https://github.com/exo-explore/exo/issues) or create a new one.

-**Configuration Options:**
-
- `--no-worker`: Run exo without the worker component. Useful for coordinator-only nodes that handle networking and orchestration but don't execute inference tasks. This is helpful for machines without sufficient GPU resources but with good network connectivity.
-
-  ```bash
-  uv run exo --no-worker
-  ```
-
-**File Locations (Linux):**
-
-exo follows the [XDG Base Directory Specification](https://specifications.freedesktop.org/basedir-spec/basedir-spec-latest.html) on Linux:
-
- **Configuration files**: `~/.config/exo/` (or `$XDG_CONFIG_HOME/exo/`)
- **Data files**: `~/.local/share/exo/` (or `$XDG_DATA_HOME/exo/`)
- **Cache files**: `~/.cache/exo/` (or `$XDG_CACHE_HOME/exo/`)
-
-You can override these locations by setting the corresponding XDG environment variables.
-
 ### macOS App

 exo ships a macOS app that runs in the background on your Mac.
@@ -193,19 +166,6 @@ Download the latest build here: [EXO-latest.dmg](https://assets.exolabs.net/EXO-

 The app will ask for permission to modify system settings and install a new Network profile. Improvements to this are being worked on.

-**Custom Namespace for Cluster Isolation:**
-
-The macOS app includes a custom namespace feature that allows you to isolate your exo cluster from others on the same network. This is configured through the `EXO_LIBP2P_NAMESPACE` setting:
-
- **Use cases**:
-  - Running multiple separate exo clusters on the same network
-  - Isolating development/testing clusters from production clusters
-  - Preventing accidental cluster joining
-
- **Configuration**: Access this setting in the app's Advanced settings (or set the `EXO_LIBP2P_NAMESPACE` environment variable when running from source)
-
-The namespace is logged on startup for debugging purposes.
-
 #### Uninstalling the macOS App

 The recommended way to uninstall is through the app itself: click the menu bar icon → Advanced → Uninstall. This cleanly removes all system components.
@@ -352,52 +312,6 @@ For further details, see:

 ---

-## Benchmarking
-
-The `exo-bench` tool measures model prefill and token generation speed across different placement configurations. This helps you optimize model performance and validate improvements.
-
-**Prerequisites:**
- Nodes should be running with `uv run exo` before benchmarking
- The tool uses the `/bench/chat/completions` endpoint
-
-**Basic usage:**
-
-```bash
-uv run bench/exo_bench.py \
-  --model llama-3.2-1b \
-  --pp 128,256,512 \
-  --tg 128,256
-```
-
-**Key parameters:**
-
- `--model`: Model to benchmark (short ID or HuggingFace ID)
- `--pp`: Prompt size hints (comma-separated integers)
- `--tg`: Generation lengths (comma-separated integers)
- `--max-nodes`: Limit placements to N nodes (default: 4)
- `--instance-meta`: Filter by `ring`, `jaccl`, or `both` (default: both)
- `--sharding`: Filter by `pipeline`, `tensor`, or `both` (default: both)
- `--repeat`: Number of repetitions per configuration (default: 1)
- `--warmup`: Warmup runs per placement (default: 0)
- `--json-out`: Output file for results (default: bench/results.json)
-
-**Example with filters:**
-
-```bash
-uv run bench/exo_bench.py \
-  --model llama-3.2-1b \
-  --pp 128,512 \
-  --tg 128 \
-  --max-nodes 2 \
-  --sharding tensor \
-  --repeat 3 \
-  --json-out my-results.json
-```
-
-The tool outputs performance metrics including prompt tokens per second (prompt_tps), generation tokens per second (generation_tps), and peak memory usage for each configuration.
-
---
-
 ## Hardware Accelerator Support

 On macOS, exo uses the GPU. On Linux, exo currently runs on CPU. We are working on extending hardware accelerator support. If you'd like support for a new hardware platform, please [search for an existing feature request](https://github.com/exo-explore/exo/issues) and add a thumbs up so we know what hardware is important to the community.
@@ -406,4 +320,4 @@ On macOS, exo uses the GPU. On Linux, exo currently runs on CPU. We are working

 ## Contributing

-See [CONTRIBUTING.md](CONTRIBUTING.md) for guidelines on how to contribute to exo.
+See [CONTRIBUTING.md](CONTRIBUTING.md) for guidelines on how to contribute to exo.
--- a/app/EXO/EXO.xcodeproj/project.pbxproj
+++ b/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 */
--- a/app/EXO/EXO.xcodeproj/project.xcworkspace/xcshareddata/swiftpm/Package.resolved
+++ b/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"
      }
    }
  ],
--- a/app/EXO/EXO/ContentView.swift
+++ b/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
        }
--- a/app/EXO/EXO/Services/LocalNetworkChecker.swift
+++ b/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:
--- a/bench/exo_bench.py
+++ b/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)

--- a/bench/exo_eval.py
+++ b/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())
--- a/dashboard/src/lib/components/MarkdownContent.svelte
+++ b/dashboard/src/lib/components/MarkdownContent.svelte
@@ -53,285 +53,62 @@
 	marked.use({ renderer });

 	/**
-	 * Unescape HTML entities that marked may have escaped
-	 */
-	function unescapeHtmlEntities(text: string): string {
-		return text
-			.replace(/&lt;/g, '<')
-			.replace(/&gt;/g, '>')
-			.replace(/&amp;/g, '&')
-			.replace(/&quot;/g, '"')
-			.replace(/&#39;/g, "'");
-	}
-
-	// Storage for math expressions extracted before markdown processing
-	const mathExpressions: Map<string, { content: string; displayMode: boolean }> = new Map();
-	let mathCounter = 0;
-
-	// Storage for HTML snippets that need protection from markdown
-	const htmlSnippets: Map<string, string> = new Map();
-	let htmlCounter = 0;
-
-	// Use alphanumeric placeholders that won't be interpreted as HTML tags
-	const MATH_PLACEHOLDER_PREFIX = 'MATHPLACEHOLDER';
-	const CODE_PLACEHOLDER_PREFIX = 'CODEPLACEHOLDER';
-	const HTML_PLACEHOLDER_PREFIX = 'HTMLPLACEHOLDER';
-
-	/**
-	 * Preprocess LaTeX: extract math, handle LaTeX document commands, and protect content
+	 * Preprocess LaTeX: convert \(...\) to $...$ and \[...\] to $$...$$
+	 * Also protect code blocks from LaTeX processing
 	 */
 	function preprocessLaTeX(text: string): string {
-		// Reset storage
-		mathExpressions.clear();
-		mathCounter = 0;
-		htmlSnippets.clear();
-		htmlCounter = 0;
-
-		// Protect code blocks first
+		// Protect code blocks
 		const codeBlocks: string[] = [];
 		let processed = text.replace(/```[\s\S]*?```|`[^`]+`/g, (match) => {
 			codeBlocks.push(match);
-			return `${CODE_PLACEHOLDER_PREFIX}${codeBlocks.length - 1}END`;
+			return `<<CODE_${codeBlocks.length - 1}>>`;
 		});

-		// Remove LaTeX document commands
-		processed = processed.replace(/\\documentclass(\[[^\]]*\])?\{[^}]*\}/g, '');
-		processed = processed.replace(/\\usepackage(\[[^\]]*\])?\{[^}]*\}/g, '');
-		processed = processed.replace(/\\begin\{document\}/g, '');
-		processed = processed.replace(/\\end\{document\}/g, '');
-		processed = processed.replace(/\\maketitle/g, '');
-		processed = processed.replace(/\\title\{[^}]*\}/g, '');
-		processed = processed.replace(/\\author\{[^}]*\}/g, '');
-		processed = processed.replace(/\\date\{[^}]*\}/g, '');
-
-		// Remove \require{...} commands (MathJax-specific, not supported by KaTeX)
-		processed = processed.replace(/\$\\require\{[^}]*\}\$/g, '');
-		processed = processed.replace(/\\require\{[^}]*\}/g, '');
-
-		// Remove unsupported LaTeX commands/environments (tikzpicture, figure, center, etc.)
-		processed = processed.replace(/\\begin\{tikzpicture\}[\s\S]*?\\end\{tikzpicture\}/g, () => {
-			const placeholder = `${HTML_PLACEHOLDER_PREFIX}${htmlCounter}END`;
-			htmlSnippets.set(placeholder, '<div class="latex-diagram-placeholder"><span class="latex-diagram-icon">📐</span><span class="latex-diagram-text">Diagram</span></div>');
-			htmlCounter++;
-			return placeholder;
-		});
-		processed = processed.replace(/\\begin\{figure\}[\s\S]*?\\end\{figure\}/g, () => {
-			const placeholder = `${HTML_PLACEHOLDER_PREFIX}${htmlCounter}END`;
-			htmlSnippets.set(placeholder, '<div class="latex-diagram-placeholder"><span class="latex-diagram-icon">🖼️</span><span class="latex-diagram-text">Figure</span></div>');
-			htmlCounter++;
-			return placeholder;
-		});
-		// Strip center environment (layout only, no content change)
-		processed = processed.replace(/\\begin\{center\}/g, '');
-		processed = processed.replace(/\\end\{center\}/g, '');
-		// Strip other layout environments
-		processed = processed.replace(/\\begin\{flushleft\}/g, '');
-		processed = processed.replace(/\\end\{flushleft\}/g, '');
-		processed = processed.replace(/\\begin\{flushright\}/g, '');
-		processed = processed.replace(/\\end\{flushright\}/g, '');
-		processed = processed.replace(/\\label\{[^}]*\}/g, '');
-		processed = processed.replace(/\\caption\{[^}]*\}/g, '');
-
-		// Protect escaped dollar signs (e.g., \$50 should become $50, not LaTeX)
-		processed = processed.replace(/\\\$/g, 'ESCAPEDDOLLARPLACEHOLDER');
-
-		// Convert LaTeX math environments to display math (both bare and wrapped in $...$)
-		const mathEnvs = ['align', 'align\\*', 'equation', 'equation\\*', 'gather', 'gather\\*', 'multline', 'multline\\*', 'eqnarray', 'eqnarray\\*', 'array', 'matrix', 'pmatrix', 'bmatrix', 'vmatrix', 'cases'];
-		for (const env of mathEnvs) {
-			// Handle $\begin{env}...\end{env}$ (with dollar signs, possibly multiline)
-			const wrappedRegex = new RegExp(`\\$\\\\begin\\{${env}\\}(\\{[^}]*\\})?([\\s\\S]*?)\\\\end\\{${env}\\}\\$`, 'g');
-			processed = processed.replace(wrappedRegex, (_, args, content) => {
-				const cleanEnv = env.replace('\\*', '*');
-				const mathContent = `\\begin{${cleanEnv}}${args || ''}${content}\\end{${cleanEnv}}`;
-				const placeholder = `${MATH_PLACEHOLDER_PREFIX}DISPLAY${mathCounter}END`;
-				mathExpressions.set(placeholder, { content: mathContent, displayMode: true });
-				mathCounter++;
-				return placeholder;
-			});
-
-			// Handle bare \begin{env}...\end{env} (without dollar signs)
-			const bareRegex = new RegExp(`\\\\begin\\{${env}\\}(\\{[^}]*\\})?([\\s\\S]*?)\\\\end\\{${env}\\}`, 'g');
-			processed = processed.replace(bareRegex, (_, args, content) => {
-				const cleanEnv = env.replace('\\*', '*');
-				const mathContent = `\\begin{${cleanEnv}}${args || ''}${content}\\end{${cleanEnv}}`;
-				const placeholder = `${MATH_PLACEHOLDER_PREFIX}DISPLAY${mathCounter}END`;
-				mathExpressions.set(placeholder, { content: mathContent, displayMode: true });
-				mathCounter++;
-				return placeholder;
-			});
-		}
-
-		// Convert LaTeX proof environments to styled blocks (use placeholders for HTML)
-		processed = processed.replace(
-			/\\begin\{proof\}([\s\S]*?)\\end\{proof\}/g,
-			(_, content) => {
-				const html = `<div class="latex-proof"><div class="latex-proof-header">Proof</div><div class="latex-proof-content">${content}</div></div>`;
-				const placeholder = `${HTML_PLACEHOLDER_PREFIX}${htmlCounter}END`;
-				htmlSnippets.set(placeholder, html);
-				htmlCounter++;
-				return placeholder;
-			}
-		);
-
-		// Convert LaTeX theorem-like environments
-		const theoremEnvs = ['theorem', 'lemma', 'corollary', 'proposition', 'definition', 'remark', 'example'];
-		for (const env of theoremEnvs) {
-			const envRegex = new RegExp(`\\\\begin\\{${env}\\}([\\s\\S]*?)\\\\end\\{${env}\\}`, 'gi');
-			const envName = env.charAt(0).toUpperCase() + env.slice(1);
-			processed = processed.replace(envRegex, (_, content) => {
-				const html = `<div class="latex-theorem"><div class="latex-theorem-header">${envName}</div><div class="latex-theorem-content">${content}</div></div>`;
-				const placeholder = `${HTML_PLACEHOLDER_PREFIX}${htmlCounter}END`;
-				htmlSnippets.set(placeholder, html);
-				htmlCounter++;
-				return placeholder;
-			});
-		}
-
-		// Convert LaTeX text formatting commands (use placeholders to protect from markdown)
-		processed = processed.replace(/\\emph\{([^}]*)\}/g, (_, content) => {
-			const placeholder = `${HTML_PLACEHOLDER_PREFIX}${htmlCounter}END`;
-			htmlSnippets.set(placeholder, `<em>${content}</em>`);
-			htmlCounter++;
-			return placeholder;
-		});
-		processed = processed.replace(/\\textit\{([^}]*)\}/g, (_, content) => {
-			const placeholder = `${HTML_PLACEHOLDER_PREFIX}${htmlCounter}END`;
-			htmlSnippets.set(placeholder, `<em>${content}</em>`);
-			htmlCounter++;
-			return placeholder;
-		});
-		processed = processed.replace(/\\textbf\{([^}]*)\}/g, (_, content) => {
-			const placeholder = `${HTML_PLACEHOLDER_PREFIX}${htmlCounter}END`;
-			htmlSnippets.set(placeholder, `<strong>${content}</strong>`);
-			htmlCounter++;
-			return placeholder;
-		});
-		processed = processed.replace(/\\texttt\{([^}]*)\}/g, (_, content) => {
-			const placeholder = `${HTML_PLACEHOLDER_PREFIX}${htmlCounter}END`;
-			htmlSnippets.set(placeholder, `<code class="inline-code">${content}</code>`);
-			htmlCounter++;
-			return placeholder;
-		});
-		processed = processed.replace(/\\underline\{([^}]*)\}/g, (_, content) => {
-			const placeholder = `${HTML_PLACEHOLDER_PREFIX}${htmlCounter}END`;
-			htmlSnippets.set(placeholder, `<u>${content}</u>`);
-			htmlCounter++;
-			return placeholder;
-		});
-
-		// Handle LaTeX line breaks and spacing
-		processed = processed.replace(/\\\\(?:\s*\n)?/g, '\n'); // \\ -> newline
-		processed = processed.replace(/\\newline/g, '\n');
-		processed = processed.replace(/\\par\b/g, '\n\n');
-		processed = processed.replace(/\\quad/g, ' ');
-		processed = processed.replace(/\\qquad/g, '  ');
-		processed = processed.replace(/~~/g, ' '); // non-breaking space
-
-		// Remove other common LaTeX commands that don't render
-		processed = processed.replace(/\\centering/g, '');
-		processed = processed.replace(/\\noindent/g, '');
-		processed = processed.replace(/\\hfill/g, '');
-		processed = processed.replace(/\\vspace\{[^}]*\}/g, '');
-		processed = processed.replace(/\\hspace\{[^}]*\}/g, ' ');
-
-		// Convert \(...\) to placeholder (display: false)
-		processed = processed.replace(/\\\(([\s\S]+?)\\\)/g, (_, content) => {
-			const placeholder = `${MATH_PLACEHOLDER_PREFIX}INLINE${mathCounter}END`;
-			mathExpressions.set(placeholder, { content, displayMode: false });
-			mathCounter++;
-			return placeholder;
-		});
-
-		// Convert \[...\] to placeholder (display: true)
-		processed = processed.replace(/\\\[([\s\S]*?)\\\]/g, (_, content) => {
-			const placeholder = `${MATH_PLACEHOLDER_PREFIX}DISPLAY${mathCounter}END`;
-			mathExpressions.set(placeholder, { content, displayMode: true });
-			mathCounter++;
-			return placeholder;
-		});
-
-		// Extract display math ($$...$$) BEFORE markdown processing
-		processed = processed.replace(/\$\$([\s\S]*?)\$\$/g, (_, content) => {
-			const placeholder = `${MATH_PLACEHOLDER_PREFIX}DISPLAY${mathCounter}END`;
-			mathExpressions.set(placeholder, { content: content.trim(), displayMode: true });
-			mathCounter++;
-			return placeholder;
-		});
-
-		// Extract inline math ($...$) BEFORE markdown processing
-		// Allow single-line only, skip currency patterns like $5 or $50
-		processed = processed.replace(/\$([^\$\n]+?)\$/g, (match, content) => {
-			if (/^\d/.test(content.trim())) {
-				return match; // Keep as-is for currency
-			}
-			const placeholder = `${MATH_PLACEHOLDER_PREFIX}INLINE${mathCounter}END`;
-			mathExpressions.set(placeholder, { content: content.trim(), displayMode: false });
-			mathCounter++;
-			return placeholder;
-		});
-
-		// Restore escaped dollar signs
-		processed = processed.replace(/ESCAPEDDOLLARPLACEHOLDER/g, '$');
+		// Convert \(...\) to $...$
+		processed = processed.replace(/\\\((.+?)\\\)/g, '$$$1$');
+		
+		// Convert \[...\] to $$...$$
+		processed = processed.replace(/\\\[([\s\S]*?)\\\]/g, '$$$$$1$$$$');

 		// Restore code blocks
-		processed = processed.replace(new RegExp(`${CODE_PLACEHOLDER_PREFIX}(\\d+)END`, 'g'), (_, index) => codeBlocks[parseInt(index)]);
-
-		// Clean up any remaining stray backslashes from unrecognized commands
-		processed = processed.replace(/\\(?=[a-zA-Z])/g, ''); // Remove \ before letters (unrecognized commands)
+		processed = processed.replace(/<<CODE_(\d+)>>/g, (_, index) => codeBlocks[parseInt(index)]);

 		return processed;
 	}

 	/**
-	 * Render math expressions with KaTeX and restore HTML placeholders
+	 * Render math expressions with KaTeX after HTML is generated
 	 */
 	function renderMath(html: string): string {
-		// Replace all math placeholders with rendered KaTeX
-		for (const [placeholder, { content, displayMode }] of mathExpressions) {
-			const escapedPlaceholder = placeholder.replace(/[.*+?^${}()|[\]\\]/g, '\\$&');
-			const regex = new RegExp(escapedPlaceholder, 'g');
+		// Render display math ($$...$$)
+		html = html.replace(/\$\$([\s\S]*?)\$\$/g, (_, math) => {
+			try {
+				return katex.renderToString(math.trim(), {
+					displayMode: true,
+					throwOnError: false,
+					output: 'html'
+				});
+			} catch {
+				return `<span class="math-error">$$${math}$$</span>`;
+			}
+		});

-			html = html.replace(regex, () => {
-				try {
-					const rendered = katex.renderToString(content, {
-						displayMode,
-						throwOnError: false,
-						output: 'html'
-					});
-
-					if (displayMode) {
-						return `
-							<div class="math-display-wrapper">
-								<div class="math-display-header">
-									<span class="math-label">LaTeX</span>
-									<button type="button" class="copy-math-btn" data-math-source="${encodeURIComponent(content)}" title="Copy LaTeX source">
-										<svg width="14" height="14" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round">
-											<rect width="14" height="14" x="8" y="8" rx="2" ry="2"/>
-											<path d="M4 16c-1.1 0-2-.9-2-2V4c0-1.1.9-2 2-2h10c1.1 0 2 .9 2 2"/>
-										</svg>
-									</button>
-								</div>
-								<div class="math-display-content">
-									${rendered}
-								</div>
-							</div>
-						`;
-					} else {
-						return `<span class="math-inline">${rendered}</span>`;
-					}
-				} catch {
-					const display = displayMode ? `$$${content}$$` : `$${content}$`;
-					return `<span class="math-error"><span class="math-error-icon">⚠</span> ${display}</span>`;
-				}
-			});
-		}
-
-		// Restore HTML placeholders (for \textbf, \emph, etc.)
-		for (const [placeholder, htmlContent] of htmlSnippets) {
-			const escapedPlaceholder = placeholder.replace(/[.*+?^${}()|[\]\\]/g, '\\$&');
-			const regex = new RegExp(escapedPlaceholder, 'g');
-			html = html.replace(regex, htmlContent);
-		}
+		// Render inline math ($...$) but avoid matching currency like $5
+		html = html.replace(/\$([^\$\n]+?)\$/g, (match, math) => {
+			// Skip if it looks like currency ($ followed by number)
+			if (/^\d/.test(math.trim())) {
+				return match;
+			}
+			try {
+				return katex.renderToString(math.trim(), {
+					displayMode: false,
+					throwOnError: false,
+					output: 'html'
+				});
+			} catch {
+				return `<span class="math-error">$${math}$</span>`;
+			}
+		});

 		return html;
 	}
@@ -377,50 +154,16 @@
 		}
 	}

-	async function handleMathCopyClick(event: Event) {
-		const target = event.currentTarget as HTMLButtonElement;
-		const encodedSource = target.getAttribute('data-math-source');
-		if (!encodedSource) return;
-
-		const source = decodeURIComponent(encodedSource);
-
-		try {
-			await navigator.clipboard.writeText(source);
-			// Show copied feedback
-			const originalHtml = target.innerHTML;
-			target.innerHTML = `
-				<svg width="14" height="14" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round">
-					<path d="M20 6L9 17l-5-5"/>
-				</svg>
-			`;
-			target.classList.add('copied');
-			setTimeout(() => {
-				target.innerHTML = originalHtml;
-				target.classList.remove('copied');
-			}, 2000);
-		} catch (error) {
-			console.error('Failed to copy math:', error);
-		}
-	}
-
 	function setupCopyButtons() {
 		if (!containerRef || !browser) return;

-		const codeButtons = containerRef.querySelectorAll<HTMLButtonElement>('.copy-code-btn');
-		for (const button of codeButtons) {
+		const buttons = containerRef.querySelectorAll<HTMLButtonElement>('.copy-code-btn');
+		for (const button of buttons) {
 			if (button.dataset.listenerBound !== 'true') {
 				button.dataset.listenerBound = 'true';
 				button.addEventListener('click', handleCopyClick);
 			}
 		}
-
-		const mathButtons = containerRef.querySelectorAll<HTMLButtonElement>('.copy-math-btn');
-		for (const button of mathButtons) {
-			if (button.dataset.listenerBound !== 'true') {
-				button.dataset.listenerBound = 'true';
-				button.addEventListener('click', handleMathCopyClick);
-			}
-		}
 	}

 	$effect(() => {
@@ -681,290 +424,28 @@
 		color: #60a5fa;
 	}

-	/* KaTeX math styling - Base */
+	/* KaTeX math styling */
 	.markdown-content :global(.katex) {
 		font-size: 1.1em;
-		color: oklch(0.9 0 0);
 	}

-	/* Display math container wrapper */
-	.markdown-content :global(.math-display-wrapper) {
+	.markdown-content :global(.katex-display) {
 		margin: 1rem 0;
-		border-radius: 0.5rem;
-		overflow: hidden;
-		border: 1px solid rgba(255, 215, 0, 0.15);
-		background: rgba(0, 0, 0, 0.3);
-		transition: border-color 0.2s ease, box-shadow 0.2s ease;
-	}
-
-	.markdown-content :global(.math-display-wrapper:hover) {
-		border-color: rgba(255, 215, 0, 0.25);
-		box-shadow: 0 0 12px rgba(255, 215, 0, 0.08);
-	}
-
-	/* Display math header - hidden by default, slides in on hover */
-	.markdown-content :global(.math-display-header) {
-		display: flex;
-		justify-content: space-between;
-		align-items: center;
-		padding: 0.375rem 0.75rem;
-		background: rgba(255, 215, 0, 0.03);
-		border-bottom: 1px solid rgba(255, 215, 0, 0.08);
-		opacity: 0;
-		max-height: 0;
-		padding-top: 0;
-		padding-bottom: 0;
-		overflow: hidden;
-		transition:
-			opacity 0.2s ease,
-			max-height 0.2s ease,
-			padding 0.2s ease;
-	}
-
-	.markdown-content :global(.math-display-wrapper:hover .math-display-header) {
-		opacity: 1;
-		max-height: 2.5rem;
-		padding: 0.375rem 0.75rem;
-	}
-
-	.markdown-content :global(.math-label) {
-		color: rgba(255, 215, 0, 0.7);
-		font-size: 0.65rem;
-		font-weight: 500;
-		text-transform: uppercase;
-		letter-spacing: 0.1em;
-		font-family: ui-monospace, SFMono-Regular, 'SF Mono', Monaco, Consolas, monospace;
-	}
-
-	.markdown-content :global(.copy-math-btn) {
-		display: flex;
-		align-items: center;
-		justify-content: center;
-		padding: 0.25rem;
-		background: transparent;
-		border: none;
-		color: var(--exo-light-gray, #9ca3af);
-		cursor: pointer;
-		transition: color 0.2s;
-		border-radius: 0.25rem;
-		opacity: 0;
-		transition:
-			color 0.2s,
-			opacity 0.15s ease;
-	}
-
-	.markdown-content :global(.math-display-wrapper:hover .copy-math-btn) {
-		opacity: 1;
-	}
-
-	.markdown-content :global(.copy-math-btn:hover) {
-		color: var(--exo-yellow, #ffd700);
-	}
-
-	.markdown-content :global(.copy-math-btn.copied) {
-		color: #22c55e;
-	}
-
-	/* Display math content area */
-	.markdown-content :global(.math-display-content) {
-		padding: 1rem 1.25rem;
 		overflow-x: auto;
 		overflow-y: hidden;
+		padding: 0.5rem 0;
 	}

-	/* Custom scrollbar for math overflow */
-	.markdown-content :global(.math-display-content::-webkit-scrollbar) {
-		height: 6px;
-	}
-
-	.markdown-content :global(.math-display-content::-webkit-scrollbar-track) {
-		background: rgba(255, 255, 255, 0.05);
-		border-radius: 3px;
-	}
-
-	.markdown-content :global(.math-display-content::-webkit-scrollbar-thumb) {
-		background: rgba(255, 215, 0, 0.2);
-		border-radius: 3px;
-	}
-
-	.markdown-content :global(.math-display-content::-webkit-scrollbar-thumb:hover) {
-		background: rgba(255, 215, 0, 0.35);
-	}
-
-	.markdown-content :global(.math-display-content .katex-display) {
-		margin: 0;
-		padding: 0;
-	}
-
-	.markdown-content :global(.math-display-content .katex-display > .katex) {
+	.markdown-content :global(.katex-display > .katex) {
 		text-align: center;
 	}

-	/* Inline math wrapper */
-	.markdown-content :global(.math-inline) {
-		display: inline;
-		padding: 0 0.125rem;
-		border-radius: 0.25rem;
-		transition: background-color 0.15s ease;
-	}
-
-	.markdown-content :global(.math-inline:hover) {
-		background: rgba(255, 215, 0, 0.05);
-	}
-
-	/* Dark theme KaTeX overrides */
-	.markdown-content :global(.katex .mord),
-	.markdown-content :global(.katex .minner),
-	.markdown-content :global(.katex .mop),
-	.markdown-content :global(.katex .mbin),
-	.markdown-content :global(.katex .mrel),
-	.markdown-content :global(.katex .mpunct) {
-		color: oklch(0.9 0 0);
-	}
-
-	/* Fraction lines and rules */
-	.markdown-content :global(.katex .frac-line),
-	.markdown-content :global(.katex .overline-line),
-	.markdown-content :global(.katex .underline-line),
-	.markdown-content :global(.katex .hline),
-	.markdown-content :global(.katex .rule) {
-		border-color: oklch(0.85 0 0) !important;
-		background: oklch(0.85 0 0);
-	}
-
-	/* Square roots and SVG elements */
-	.markdown-content :global(.katex .sqrt-line) {
-		border-color: oklch(0.85 0 0) !important;
-	}
-
-	.markdown-content :global(.katex svg) {
-		fill: oklch(0.85 0 0);
-		stroke: oklch(0.85 0 0);
-	}
-
-	.markdown-content :global(.katex svg path) {
-		stroke: oklch(0.85 0 0);
-	}
-
-	/* Delimiters (parentheses, brackets, braces) */
-	.markdown-content :global(.katex .delimsizing),
-	.markdown-content :global(.katex .delim-size1),
-	.markdown-content :global(.katex .delim-size2),
-	.markdown-content :global(.katex .delim-size3),
-	.markdown-content :global(.katex .delim-size4),
-	.markdown-content :global(.katex .mopen),
-	.markdown-content :global(.katex .mclose) {
-		color: oklch(0.75 0 0);
-	}
-
-	/* Math error styling */
 	.markdown-content :global(.math-error) {
-		display: inline-flex;
-		align-items: center;
-		gap: 0.375rem;
 		color: #f87171;
 		font-family: ui-monospace, SFMono-Regular, 'SF Mono', Monaco, Consolas, monospace;
 		font-size: 0.875em;
 		background: rgba(248, 113, 113, 0.1);
-		padding: 0.25rem 0.5rem;
+		padding: 0.125rem 0.25rem;
 		border-radius: 0.25rem;
-		border: 1px solid rgba(248, 113, 113, 0.2);
-	}
-
-	.markdown-content :global(.math-error-icon) {
-		font-size: 0.875em;
-		opacity: 0.9;
-	}
-
-	/* LaTeX proof environment */
-	.markdown-content :global(.latex-proof) {
-		margin: 1rem 0;
-		padding: 1rem 1.25rem;
-		background: rgba(255, 255, 255, 0.02);
-		border-left: 3px solid rgba(255, 215, 0, 0.4);
-		border-radius: 0 0.375rem 0.375rem 0;
-	}
-
-	.markdown-content :global(.latex-proof-header) {
-		font-weight: 600;
-		font-style: italic;
-		color: oklch(0.85 0 0);
-		margin-bottom: 0.5rem;
-	}
-
-	.markdown-content :global(.latex-proof-header::after) {
-		content: '.';
-	}
-
-	.markdown-content :global(.latex-proof-content) {
-		color: oklch(0.9 0 0);
-	}
-
-	.markdown-content :global(.latex-proof-content p:last-child) {
-		margin-bottom: 0;
-	}
-
-	/* QED symbol at end of proof */
-	.markdown-content :global(.latex-proof-content::after) {
-		content: '∎';
-		display: block;
-		text-align: right;
-		color: oklch(0.7 0 0);
-		margin-top: 0.5rem;
-	}
-
-	/* LaTeX theorem-like environments */
-	.markdown-content :global(.latex-theorem) {
-		margin: 1rem 0;
-		padding: 1rem 1.25rem;
-		background: rgba(255, 215, 0, 0.03);
-		border: 1px solid rgba(255, 215, 0, 0.15);
-		border-radius: 0.375rem;
-	}
-
-	.markdown-content :global(.latex-theorem-header) {
-		font-weight: 700;
-		color: var(--exo-yellow, #ffd700);
-		margin-bottom: 0.5rem;
-	}
-
-	.markdown-content :global(.latex-theorem-header::after) {
-		content: '.';
-	}
-
-	.markdown-content :global(.latex-theorem-content) {
-		color: oklch(0.9 0 0);
-		font-style: italic;
-	}
-
-	.markdown-content :global(.latex-theorem-content p:last-child) {
-		margin-bottom: 0;
-	}
-
-	/* LaTeX diagram/figure placeholder */
-	.markdown-content :global(.latex-diagram-placeholder) {
-		display: flex;
-		align-items: center;
-		justify-content: center;
-		gap: 0.5rem;
-		margin: 1rem 0;
-		padding: 1.5rem 2rem;
-		background: rgba(255, 255, 255, 0.02);
-		border: 1px dashed rgba(255, 215, 0, 0.25);
-		border-radius: 0.5rem;
-		color: rgba(255, 215, 0, 0.6);
-		font-size: 0.875rem;
-	}
-
-	.markdown-content :global(.latex-diagram-icon) {
-		font-size: 1.25rem;
-		opacity: 0.8;
-	}
-
-	.markdown-content :global(.latex-diagram-text) {
-		font-family: ui-monospace, SFMono-Regular, 'SF Mono', Monaco, Consolas, monospace;
-		font-size: 0.75rem;
-		text-transform: uppercase;
-		letter-spacing: 0.05em;
 	}
 </style>
--- a/docs/imgs/dashboard-cluster-view.png
+++ b/docs/imgs/dashboard-cluster-view.png
--- a/flake.nix
+++ b/flake.nix
@@ -87,12 +87,6 @@
            touch $out
          '';

-          packages =
-            if pkgs.stdenv.isDarwin then {
-              metal = pkgs.callPackage ./nix/metalWrapper.nix { metalVersion = "310"; };
-              mlx = pkgs.callPackage ./nix/mlx.nix {};
-            } else { };
-
          devShells.default = with pkgs; pkgs.mkShell {
            inputsFrom = [ self'.checks.cargo-build ];

@@ -130,7 +124,6 @@

            OPENSSL_NO_VENDOR = "1";

-
            shellHook = ''
              export LD_LIBRARY_PATH="$LD_LIBRARY_PATH:${python313}/lib"
              ${lib.optionalString stdenv.isLinux ''
--- a/nix/darwin-build-fixes.patch
+++ b/nix/darwin-build-fixes.patch
@@ -1,79 +0,0 @@
-diff --git a/CMakeLists.txt b/CMakeLists.txt
-index 0ed30932..d8528132 100644
--- a/CMakeLists.txt
-+++ b/CMakeLists.txt
-@@ -177,11 +177,7 @@ if(MLX_BUILD_METAL)
-     add_compile_definitions(MLX_METAL_DEBUG)
-   endif()
- 
-  # Throw an error if xcrun not found
-  execute_process(
-    COMMAND zsh "-c" "/usr/bin/xcrun -sdk macosx --show-sdk-version"
-    OUTPUT_VARIABLE MACOS_SDK_VERSION
-    OUTPUT_STRIP_TRAILING_WHITESPACE COMMAND_ERROR_IS_FATAL ANY)
-+  set(MACOS_SDK_VERSION @sdkVersion@)
- 
-   if(${MACOS_SDK_VERSION} LESS 14.0)
-     message(
-@@ -199,11 +195,8 @@ if(MLX_BUILD_METAL)
-     endif()
-     set(XCRUN_FLAGS "-mmacosx-version-min=${CMAKE_OSX_DEPLOYMENT_TARGET}")
-   endif()
-  execute_process(
-    COMMAND
-      zsh "-c"
-      "echo \"__METAL_VERSION__\" | xcrun -sdk macosx metal ${XCRUN_FLAGS} -E -x metal -P - | tail -1 | tr -d '\n'"
-    OUTPUT_VARIABLE MLX_METAL_VERSION COMMAND_ERROR_IS_FATAL ANY)
-+  set(
-+    MLX_METAL_VERSION @metalVersion@)
-   FetchContent_Declare(metal_cpp URL ${METAL_CPP_URL})
-   FetchContent_MakeAvailable(metal_cpp)
-   target_include_directories(
-diff --git a/cmake/extension.cmake b/cmake/extension.cmake
-index 13db804a..5b385132 100644
--- a/cmake/extension.cmake
-+++ b/cmake/extension.cmake
-@@ -36,7 +36,7 @@ macro(mlx_build_metallib)
-   add_custom_command(
-     OUTPUT ${MTLLIB_BUILD_TARGET}
-     COMMAND
-      xcrun -sdk macosx metal
-+      metal
-       "$<LIST:TRANSFORM,${MTLLIB_INCLUDE_DIRS},PREPEND,-I>"
-       ${MTLLIB_COMPILE_OPTIONS} ${MTLLIB_SOURCES} -o ${MTLLIB_BUILD_TARGET}
-     DEPENDS ${MTLLIB_DEPS} ${MTLLIB_SOURCES}
-diff --git a/mlx/backend/metal/kernels/CMakeLists.txt b/mlx/backend/metal/kernels/CMakeLists.txt
-index 262b0495..5c7446ad 100644
--- a/mlx/backend/metal/kernels/CMakeLists.txt
-+++ b/mlx/backend/metal/kernels/CMakeLists.txt
-@@ -29,7 +29,7 @@ function(build_kernel_base TARGET SRCFILE DEPS)
-                     "-mmacosx-version-min=${CMAKE_OSX_DEPLOYMENT_TARGET}")
-   endif()
-   add_custom_command(
-    COMMAND xcrun -sdk macosx metal ${METAL_FLAGS} -c ${SRCFILE}
-+    COMMAND metal ${METAL_FLAGS} -c ${SRCFILE}
-             -I${PROJECT_SOURCE_DIR} -o ${TARGET}.air
-     DEPENDS ${SRCFILE} ${DEPS} ${BASE_HEADERS}
-     OUTPUT ${TARGET}.air
-@@ -170,7 +170,7 @@ endif()
- 
- add_custom_command(
-   OUTPUT ${MLX_METAL_PATH}/mlx.metallib
-  COMMAND xcrun -sdk macosx metallib ${KERNEL_AIR} -o
-+  COMMAND metallib ${KERNEL_AIR} -o
-           ${MLX_METAL_PATH}/mlx.metallib
-   DEPENDS ${KERNEL_AIR}
-   COMMENT "Building mlx.metallib"
-diff --git a/mlx/backend/metal/make_compiled_preamble.sh b/mlx/backend/metal/make_compiled_preamble.sh
-index bb55ed3a..94ea7dd7 100644
--- a/mlx/backend/metal/make_compiled_preamble.sh
-+++ b/mlx/backend/metal/make_compiled_preamble.sh
-@@ -31,7 +31,7 @@ OUTPUT_FILE=${OUTPUT_DIR}/${SRC_NAME}.cpp
- mkdir -p "$OUTPUT_DIR"
- 
- # Use the metal compiler to get a list of headers (with depth)
-CCC="xcrun -sdk macosx metal -x metal"
-+CCC="metal -x metal"
- HDRS=$( $CCC -I"$SRC_DIR" -I"$JIT_INCLUDES" -DMLX_METAL_JIT -E -P -CC -C -H "$INPUT_FILE" $CFLAGS -w 2>&1 1>/dev/null )
- 
- # Remove any included system frameworks (for MetalPerformancePrimitive headers)
--- a/nix/metalWrapper.nix
+++ b/nix/metalWrapper.nix
@@ -1,25 +0,0 @@
-{ stdenvNoCC
-, metalVersion
-}:
-assert stdenvNoCC.isDarwin;
-stdenvNoCC.mkDerivation {
-  pname = "metal-wrapper-impure";
-  version = metalVersion;
-
-  __noChroot = true;
-  buildCommand = ''
-    mkdir -p $out/bin && cd $out/bin
-
-    METALLIB_PATH=''${GH_OVERRIDE_METALLIB:-$(/usr/bin/xcrun --sdk macosx -f metallib)}
-    METAL_PATH=''${GH_OVERRIDE_METAL:-"$(dirname "$METALLIB_PATH")/metal"}
-    echo "$METAL_PATH"
-    echo "$METALLIB_PATH"
-
-    ln -sf "$METAL_PATH" metal
-    ln -sf "$METALLIB_PATH" metallib
-
-    [[ -e $out/bin/metal ]] && [[ -e $out/bin/metallib ]] || { echo ":(" && exit 1; }
-    METAL_VERSION=$(echo __METAL_VERSION__ | "$METAL_PATH" -E -x metal -P - | tail -1 | tr -d '\n')
-    [[ "$METAL_VERSION" == "${metalVersion}" ]] || { echo "Metal version $METAL_VERSION is not ${metalVersion}" && exit 1; }
-  '';
-}
--- a/nix/mlx.nix
+++ b/nix/mlx.nix
@@ -1,158 +0,0 @@
-{ stdenv
-, lib
-, fetchFromGitHub
-, fetchPypi
-, pyprojectHook
-, pypaInstallHook
-, replaceVars
-, fetchzip
-, setuptools
-, cmake
-, nanobind
-, pybind11
-, nlohmann_json
-, apple-sdk_26
-, metal
-, numpy
-, pytestCheckHook
-, python
-, runCommand
-, fmt
-}:
-assert stdenv.isDarwin;
-let
-  # static dependencies included directly during compilation
-  gguf-tools = fetchFromGitHub {
-    owner = "antirez";
-    repo = "gguf-tools";
-    rev = "8fa6eb65236618e28fd7710a0fba565f7faa1848";
-    hash = "sha256-15FvyPOFqTOr5vdWQoPnZz+mYH919++EtghjozDlnSA=";
-  };
-
-  metal_cpp = fetchzip {
-    url = "https://developer.apple.com/metal/cpp/files/metal-cpp_26.zip";
-    hash = "sha256-7n2eI2lw/S+Us6l7YPAATKwcIbRRpaQ8VmES7S8ZjY8=";
-  };
-
-  mlx = stdenv.mkDerivation rec {
-    pname = "mlx";
-    version = "0.30.1";
-    pyproject = true;
-
-    src = fetchFromGitHub {
-      owner = "ml-explore";
-      repo = "mlx";
-      tag = "v${version}";
-      hash = "sha256-Vt0RH+70VBwUjXSfPTsNdRS3g0ookJHhzf2kvgEtgH8=";
-    };
-
-    patches = [
-      (replaceVars ./darwin-build-fixes.patch {
-        sdkVersion = apple-sdk_26.version;
-        metalVersion = metal.version;
-      })
-    ];
-
-    postPatch = ''
-      substituteInPlace pyproject.toml \
-        --replace-fail "nanobind==2.10.2" "nanobind"
-
-      substituteInPlace mlx/backend/cpu/jit_compiler.cpp \
-        --replace-fail "g++" "$CXX"
-    '';
-
-    dontUseCmakeConfigure = true;
-
-    enableParallelBuilding = true;
-
-    # Allows multiple cores to be used in Python builds.
-    postUnpack = ''
-      export MAKEFLAGS+="''${enableParallelBuilding:+-j$NIX_BUILD_CORES}"
-    '';
-
-    # updates the wrong fetcher rev attribute
-    passthru.skipBulkUpdate = true;
-
-    env = {
-      DEV_RELEASE = 1;
-      # NOTE The `metal` command-line utility used to build the Metal kernels is not open-source.
-      # this is what the xcode wrapper is for - it patches in the system metal cli
-      CMAKE_ARGS = toString [
-        (lib.cmakeBool "USE_SYSTEM_FMT" true)
-        (lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_GGUFLIB" "${gguf-tools}")
-        (lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_JSON" "${nlohmann_json.src}")
-        (lib.cmakeBool "FETCHCONTENT_FULLY_DISCONNECTED" true)
-        (lib.cmakeBool "MLX_BUILD_METAL" true)
-        (lib.cmakeOptionType "filepath" "METAL_LIB"
-          "${metal}/Metal.framework")
-        (lib.cmakeOptionType "filepath" "FETCHCONTENT_SOURCE_DIR_METAL_CPP" "${metal_cpp}")
-        (lib.cmakeOptionType "string" "CMAKE_OSX_DEPLOYMENT_TARGET" "${apple-sdk_26.version}")
-        (lib.cmakeOptionType "filepath" "CMAKE_OSX_SYSROOT" "${apple-sdk_26.passthru.sdkroot}")
-      ];
-      SDKROOT = apple-sdk_26.passthru.sdkroot;
-      MACOSX_DEPLOYMENT_TARGET = apple-sdk_26.version;
-    };
-
-    build-system = [
-      setuptools
-    ];
-
-    nativeBuildInputs = [
-      cmake
-      metal
-      pyprojectHook
-      pypaInstallHook
-    ];
-
-    buildInputs = [
-      fmt
-      gguf-tools
-      nanobind
-      pybind11
-      apple-sdk_26
-    ];
-
-    pythonImportsCheck = [ "mlx" ];
-
-    # Run the mlx Python test suite.
-    nativeCheckInputs = [
-      numpy
-      pytestCheckHook
-    ];
-
-    enabledTestPaths = [
-      "python/tests/"
-    ];
-
-    # Additional testing by executing the example Python scripts supplied with mlx
-    # using the version of the library we've built.
-    passthru.tests = {
-      mlxTest =
-        runCommand "run-mlx-examples"
-          {
-            buildInputs = [ mlx ];
-            nativeBuildInputs = [ python ];
-          }
-          ''
-            cp ${src}/examples/python/logistic_regression.py .
-            ${python.interpreter} logistic_regression.py
-            rm logistic_regression.py
-
-            cp ${src}/examples/python/linear_regression.py .
-            ${python.interpreter} linear_regression.py
-            rm linear_regression.py
-
-            touch $out
-          '';
-    };
-
-    meta = {
-      homepage = "https://github.com/ml-explore/mlx";
-      description = "Array framework for Apple silicon";
-      changelog = "https://github.com/ml-explore/mlx/releases/tag/${src.tag}";
-      license = lib.licenses.mit;
-      platforms = [ "x86_64-linux" "aarch64-linux" "aarch64-darwin" ];
-    };
-  };
-in
-mlx
--- a/pyproject.toml
+++ b/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",
-]
--- a/src/exo/main.py
+++ b/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
--- a/src/exo/master/api.py
+++ b/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,
@@ -52,12 +49,7 @@ from exo.shared.types.commands import (
    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
@@ -80,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,
            )
        ],
@@ -123,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()

@@ -154,21 +151,6 @@ class API:
        self.paused_ev.set()
        self.paused_ev = anyio.Event()

-    def _setup_exception_handlers(self) -> None:
-        self.app.exception_handler(HTTPException)(self.http_exception_handler)
-
-    async def http_exception_handler(
-        self, _: 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,
@@ -430,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
            )
@@ -460,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

@@ -488,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,
                )
@@ -501,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

@@ -531,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,
                )
@@ -655,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:
--- a/src/exo/master/placement_utils.py
+++ b/src/exo/master/placement_utils.py
@@ -49,83 +49,33 @@ def get_smallest_cycles(cycles: list[list[NodeInfo]]) -> list[list[NodeInfo]]:
    return [cycle for cycle in cycles if len(cycle) == min_nodes]


-def allocate_layers_proportionally(
-    total_layers: int,
-    memory_fractions: list[float],
-) -> list[int]:
-    n = len(memory_fractions)
-    if n == 0:
-        raise ValueError("Cannot allocate layers to an empty node list")
-    if total_layers < n:
-        raise ValueError(
-            f"Cannot distribute {total_layers} layers across {n} nodes "
-            "(need at least 1 layer per node)"
-        )
-
-    # Largest remainder: floor each, then distribute remainder by fractional part
-    raw = [f * total_layers for f in memory_fractions]
-    result = [int(r) for r in raw]
-    by_remainder = sorted(range(n), key=lambda i: raw[i] - result[i], reverse=True)
-    for i in range(total_layers - sum(result)):
-        result[by_remainder[i]] += 1
-
-    # Ensure minimum 1 per node by taking from the largest
-    for i in range(n):
-        if result[i] == 0:
-            max_idx = max(range(n), key=lambda j: result[j])
-            assert result[max_idx] > 1
-            result[max_idx] -= 1
-            result[i] = 1
-
-    return result
-
-
 def get_shard_assignments_for_pipeline_parallel(
    model_meta: ModelMetadata,
    selected_cycle: list[NodeWithProfile],
 ):
-    if not selected_cycle:
-        raise ValueError("Cannot create shard assignments for empty node cycle")
-
    cycle_memory = sum(
        (node.node_profile.memory.ram_available for node in selected_cycle),
        start=Memory(),
    )
-
-    if cycle_memory.in_bytes == 0:
-        raise ValueError("Cannot create shard assignments: total available memory is 0")
-
    total_layers = model_meta.n_layers
    world_size = len(selected_cycle)
    runner_to_shard: dict[RunnerId, ShardMetadata] = {}
    node_to_runner: dict[NodeId, RunnerId] = {}

-    layer_allocations = allocate_layers_proportionally(
-        total_layers=total_layers,
-        memory_fractions=[
-            node.node_profile.memory.ram_available.in_bytes / cycle_memory.in_bytes
-            for node in selected_cycle
-        ],
-    )
-
-    # Validate each node has sufficient memory for its assigned layers
-    memory_per_layer = model_meta.storage_size.in_bytes / total_layers
-    for i, (node, node_layers) in enumerate(
-        zip(selected_cycle, layer_allocations, strict=True)
-    ):
-        required_memory = node_layers * memory_per_layer
-        available_memory = node.node_profile.memory.ram_available.in_bytes
-        if required_memory > available_memory:
-            raise ValueError(
-                f"Node {i} ({node.node_id}) has insufficient memory: "
-                f"requires {required_memory / (1024**3):.2f} GB for {node_layers} layers, "
-                f"but only has {available_memory / (1024**3):.2f} GB available"
-            )
-
    layers_assigned = 0
-    for i, (node, node_layers) in enumerate(
-        zip(selected_cycle, layer_allocations, strict=True)
-    ):
+    for i, node in enumerate(selected_cycle):
+        if i == len(selected_cycle) - 1:
+            node_layers = total_layers - layers_assigned
+        else:
+            node_layers = round(
+                total_layers
+                * (
+                    node.node_profile.memory.ram_available.in_bytes
+                    / cycle_memory.in_bytes
+                )
+            )
+            node_layers = max(1, node_layers)
+
        runner_id = RunnerId()

        shard = PipelineShardMetadata(
--- a/src/exo/master/tests/test_api_error_handling.py
+++ b/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
--- a/src/exo/master/tests/test_placement.py
+++ b/src/exo/master/tests/test_placement.py
@@ -70,7 +70,7 @@ def place_instance_command(model_meta: ModelMetadata) -> PlaceInstance:
    [
        ((500, 500, 1000), 12, (3, 3, 6)),
        ((500, 500, 500), 12, (4, 4, 4)),
-        ((312, 468, 1092), 12, (2, 3, 7)),
+        ((312, 518, 1024), 12, (2, 3, 7)),
    ],
 )
 def test_get_instance_placements_create_instance(
--- a/src/exo/master/tests/test_placement_utils.py
+++ b/src/exo/master/tests/test_placement_utils.py
@@ -3,7 +3,6 @@ from typing import Callable
 import pytest

 from exo.master.placement_utils import (
-    allocate_layers_proportionally,
    filter_cycles_by_memory,
    get_hosts_from_subgraph,
    get_mlx_jaccl_coordinators,
@@ -166,9 +165,6 @@ def test_get_smallest_cycles(
        ((500, 500, 1000), 12, (3, 3, 6)),
        ((500, 500, 500), 12, (4, 4, 4)),
        ((312, 518, 1024), 12, (2, 3, 7)),
-        # Edge case: one node has ~90% of memory - should not over-allocate.
-        # Each node must have enough memory for at least 1 layer (50 KB = 1000/20).
-        ((900, 50, 50), 20, (18, 1, 1)),
    ],
 )
 def test_get_shard_assignments(
@@ -401,96 +397,3 @@ def test_get_mlx_jaccl_coordinators(
    assert coordinators[node_c_id] == (
        f"{conn_c_a.send_back_multiaddr.ip_address}:5000"
    ), "node_c should use the IP from conn_c_a"
-
-
-class TestAllocateLayersProportionally:
-    def test_empty_node_list_raises(self):
-        with pytest.raises(ValueError, match="empty node list"):
-            allocate_layers_proportionally(total_layers=10, memory_fractions=[])
-
-    def test_zero_layers_raises(self):
-        with pytest.raises(ValueError, match="need at least 1 layer per node"):
-            allocate_layers_proportionally(total_layers=0, memory_fractions=[0.5, 0.5])
-
-    def test_negative_layers_raises(self):
-        with pytest.raises(ValueError, match="need at least 1 layer per node"):
-            allocate_layers_proportionally(total_layers=-1, memory_fractions=[0.5, 0.5])
-
-    def test_fewer_layers_than_nodes_raises(self):
-        with pytest.raises(ValueError, match="need at least 1 layer per node"):
-            allocate_layers_proportionally(
-                total_layers=2, memory_fractions=[0.33, 0.33, 0.34]
-            )
-
-    def test_equal_distribution(self):
-        result = allocate_layers_proportionally(
-            total_layers=12, memory_fractions=[0.25, 0.25, 0.25, 0.25]
-        )
-        assert result == [3, 3, 3, 3]
-        assert sum(result) == 12
-
-    def test_proportional_distribution(self):
-        result = allocate_layers_proportionally(
-            total_layers=12, memory_fractions=[0.25, 0.25, 0.50]
-        )
-        assert result == [3, 3, 6]
-        assert sum(result) == 12
-
-    def test_extreme_imbalance_ensures_minimum(self):
-        result = allocate_layers_proportionally(
-            total_layers=20, memory_fractions=[0.975, 0.0125, 0.0125]
-        )
-        assert all(layers >= 1 for layers in result)
-        assert sum(result) == 20
-        # Small nodes get minimum 1 layer
-        assert result == [18, 1, 1]
-
-    def test_single_node_gets_all_layers(self):
-        result = allocate_layers_proportionally(total_layers=10, memory_fractions=[1.0])
-        assert result == [10]
-
-    def test_minimum_viable_allocation(self):
-        result = allocate_layers_proportionally(
-            total_layers=3, memory_fractions=[0.33, 0.33, 0.34]
-        )
-        assert result == [1, 1, 1]
-        assert sum(result) == 3
-
-
-def test_get_shard_assignments_insufficient_memory_raises(
-    topology: Topology,
-    create_node: Callable[[int, NodeId | None], NodeInfo],
-    create_connection: Callable[[NodeId, NodeId], Connection],
-):
-    """Test that ValueError is raised when a node has insufficient memory for its layers."""
-    node_a_id = NodeId()
-    node_b_id = NodeId()
-    node_c_id = NodeId()
-
-    # Node C has only 10 KB but would need 50 KB for 1 layer (1000 KB / 20 layers)
-    node_a = create_node(900 * 1024, node_a_id)
-    node_b = create_node(50 * 1024, node_b_id)
-    node_c = create_node(10 * 1024, node_c_id)  # Insufficient memory
-
-    topology.add_node(node_a)
-    topology.add_node(node_b)
-    topology.add_node(node_c)
-
-    topology.add_connection(create_connection(node_a_id, node_b_id))
-    topology.add_connection(create_connection(node_b_id, node_c_id))
-    topology.add_connection(create_connection(node_c_id, node_a_id))
-    topology.add_connection(create_connection(node_b_id, node_a_id))
-
-    model_meta = ModelMetadata(
-        model_id=ModelId("test-model"),
-        pretty_name="Test Model",
-        n_layers=20,
-        storage_size=Memory.from_kb(1000),
-        hidden_size=1000,
-        supports_tensor=True,
-    )
-    cycles = topology.get_cycles()
-    selected_cycle = cycles[0]
-
-    with pytest.raises(ValueError, match="insufficient memory"):
-        get_shard_assignments(model_meta, selected_cycle, Sharding.Pipeline)
--- a/src/exo/shared/logging.py
+++ b/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
--- a/src/exo/shared/types/api.py
+++ b/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"
--- a/src/exo/shared/types/chunks.py
+++ b/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):
--- a/src/exo/shared/types/worker/runner_response.py
+++ b/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):
--- a/src/exo/worker/download/download_utils.py
+++ b/src/exo/worker/download/download_utils.py
@@ -5,7 +5,6 @@ import shutil
 import ssl
 import time
 import traceback
-from collections.abc import Awaitable
 from datetime import timedelta
 from pathlib import Path
 from typing import Callable, Literal
@@ -246,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,
@@ -526,7 +522,7 @@ async def download_progress_for_local_path(

 async def download_shard(
    shard: ShardMetadata,
-    on_progress: Callable[[ShardMetadata, RepoDownloadProgress], Awaitable[None]],
+    on_progress: Callable[[ShardMetadata, RepoDownloadProgress], None],
    max_parallel_downloads: int = 8,
    skip_download: bool = False,
    allow_patterns: list[str] | None = None,
@@ -567,9 +563,9 @@ async def download_shard(
    )
    file_progress: dict[str, RepoFileDownloadProgress] = {}

-    async def on_progress_wrapper(
+    def on_progress_wrapper(
        file: FileListEntry, curr_bytes: int, total_bytes: int, is_renamed: bool
-    ) -> None:
+    ):
        start_time = (
            file_progress[file.path].start_time
            if file.path in file_progress
@@ -605,7 +601,7 @@ async def download_shard(
            else "in_progress",
            start_time=start_time,
        )
-        await on_progress(
+        on_progress(
            shard,
            calculate_repo_progress(
                shard,
@@ -633,21 +629,14 @@ async def download_shard(

    semaphore = asyncio.Semaphore(max_parallel_downloads)

-    def schedule_progress(
-        file: FileListEntry, curr_bytes: int, total_bytes: int, is_renamed: bool
-    ) -> None:
-        asyncio.create_task(
-            on_progress_wrapper(file, curr_bytes, total_bytes, is_renamed)
-        )
-
-    async def download_with_semaphore(file: FileListEntry) -> None:
+    async def download_with_semaphore(file: FileListEntry):
        async with semaphore:
            await download_file_with_retry(
                str(shard.model_meta.model_id),
                revision,
                file.path,
                target_dir,
-                lambda curr_bytes, total_bytes, is_renamed: schedule_progress(
+                lambda curr_bytes, total_bytes, is_renamed: on_progress_wrapper(
                    file, curr_bytes, total_bytes, is_renamed
                ),
            )
@@ -659,7 +648,7 @@ async def download_shard(
    final_repo_progress = calculate_repo_progress(
        shard, str(shard.model_meta.model_id), revision, file_progress, all_start_time
    )
-    await on_progress(shard, final_repo_progress)
+    on_progress(shard, final_repo_progress)
    if gguf := next((f for f in filtered_file_list if f.path.endswith(".gguf")), None):
        return target_dir / gguf.path, final_repo_progress
    else:
--- a/src/exo/worker/download/impl_shard_downloader.py
+++ b/src/exo/worker/download/impl_shard_downloader.py
@@ -1,5 +1,4 @@
 import asyncio
-from collections.abc import Awaitable
 from pathlib import Path
 from typing import AsyncIterator, Callable

@@ -49,8 +48,7 @@ class SingletonShardDownloader(ShardDownloader):
        self.active_downloads: dict[ShardMetadata, asyncio.Task[Path]] = {}

    def on_progress(
-        self,
-        callback: Callable[[ShardMetadata, RepoDownloadProgress], Awaitable[None]],
+        self, callback: Callable[[ShardMetadata, RepoDownloadProgress], None]
    ) -> None:
        self.shard_downloader.on_progress(callback)

@@ -85,8 +83,7 @@ class CachedShardDownloader(ShardDownloader):
        self.cache: dict[tuple[str, ShardMetadata], Path] = {}

    def on_progress(
-        self,
-        callback: Callable[[ShardMetadata, RepoDownloadProgress], Awaitable[None]],
+        self, callback: Callable[[ShardMetadata, RepoDownloadProgress], None]
    ) -> None:
        self.shard_downloader.on_progress(callback)

@@ -116,18 +113,17 @@ class ResumableShardDownloader(ShardDownloader):
    def __init__(self, max_parallel_downloads: int = 8):
        self.max_parallel_downloads = max_parallel_downloads
        self.on_progress_callbacks: list[
-            Callable[[ShardMetadata, RepoDownloadProgress], Awaitable[None]]
+            Callable[[ShardMetadata, RepoDownloadProgress], None]
        ] = []

-    async def on_progress_wrapper(
+    def on_progress_wrapper(
        self, shard: ShardMetadata, progress: RepoDownloadProgress
    ) -> None:
        for callback in self.on_progress_callbacks:
-            await callback(shard, progress)
+            callback(shard, progress)

    def on_progress(
-        self,
-        callback: Callable[[ShardMetadata, RepoDownloadProgress], Awaitable[None]],
+        self, callback: Callable[[ShardMetadata, RepoDownloadProgress], None]
    ) -> None:
        self.on_progress_callbacks.append(callback)

--- a/src/exo/worker/download/shard_downloader.py
+++ b/src/exo/worker/download/shard_downloader.py
@@ -1,5 +1,4 @@
 from abc import ABC, abstractmethod
-from collections.abc import Awaitable
 from copy import copy
 from datetime import timedelta
 from pathlib import Path
@@ -32,8 +31,7 @@ class ShardDownloader(ABC):

    @abstractmethod
    def on_progress(
-        self,
-        callback: Callable[[ShardMetadata, RepoDownloadProgress], Awaitable[None]],
+        self, callback: Callable[[ShardMetadata, RepoDownloadProgress], None]
    ) -> None:
        pass

@@ -61,8 +59,7 @@ class NoopShardDownloader(ShardDownloader):
        return Path("/tmp/noop_shard")

    def on_progress(
-        self,
-        callback: Callable[[ShardMetadata, RepoDownloadProgress], Awaitable[None]],
+        self, callback: Callable[[ShardMetadata, RepoDownloadProgress], None]
    ) -> None:
        pass

--- a/src/exo/worker/engines/mlx/auto_parallel.py
+++ b/src/exo/worker/engines/mlx/auto_parallel.py
@@ -46,11 +46,9 @@ class CustomMlxLayer(nn.Module):

    def __init__(self, original_layer: _LayerCallable):
        super().__init__()
+        # Set twice to avoid __setattr__ recursion
        object.__setattr__(self, "_original_layer", original_layer)
-
-    @property
-    def original_layer(self) -> _LayerCallable:
-        return cast(_LayerCallable, object.__getattribute__(self, "_original_layer"))
+        self.original_layer: _LayerCallable = original_layer

    # Calls __getattr__ for any attributes not found on nn.Module (e.g. use_sliding)
    if not TYPE_CHECKING:
@@ -60,7 +58,7 @@ class CustomMlxLayer(nn.Module):
                return super().__getattr__(name)
            except AttributeError:
                original_layer = object.__getattribute__(self, "_original_layer")
-                return getattr(original_layer, name)
+                return object.__getattribute__(original_layer, name)


 class PipelineFirstLayer(CustomMlxLayer):
@@ -170,21 +168,11 @@ def pipeline_auto_parallel(
        inner_model_instance.layer_types = inner_model_instance.layer_types[  # type: ignore
            start_layer:end_layer
        ]
-        # We can assume the model has at least one layer thanks to placement.
-        # If a layer type doesn't exist, we can set it to 0.
-        inner_model_instance.swa_idx = (
-            0
-            if "sliding_attention" not in inner_model_instance.layer_types  # type: ignore
-            else inner_model_instance.layer_types.index(  # type: ignore
-                "sliding_attention"
-            )
+        inner_model_instance.swa_idx = inner_model_instance.layer_types.index(  # type: ignore
+            "sliding_attention"
        )
-        inner_model_instance.ga_idx = (
-            0
-            if "full_attention" not in inner_model_instance.layer_types  # type: ignore
-            else inner_model_instance.layer_types.index(  # type: ignore
-                "full_attention"
-            )
+        inner_model_instance.ga_idx = inner_model_instance.layer_types.index(  # type: ignore
+            "full_attention"
        )

    _set_layers(model, layers)
--- a/src/exo/worker/engines/mlx/utils_mlx.py
+++ b/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"
@@ -261,7 +216,6 @@ def load_mlx_items(
 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)

@@ -298,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)
--- a/src/exo/worker/main.py
+++ b/src/exo/worker/main.py
@@ -359,7 +359,8 @@ class Worker:
        last_progress_time = 0.0
        throttle_interval_secs = 1.0

-        async def download_progress_callback(
+        # TODO: i hate callbacks
+        def download_progress_callback(
            shard: ShardMetadata, progress: RepoDownloadProgress
        ) -> None:
            nonlocal self
@@ -371,10 +372,11 @@ class Worker:
                    total_bytes=progress.total_bytes,
                )
                self.download_status[shard.model_meta.model_id] = status
-                await self.event_sender.send(
+                # Footgun!
+                self.event_sender.send_nowait(
                    NodeDownloadProgress(download_progress=status)
                )
-                await self.event_sender.send(
+                self.event_sender.send_nowait(
                    TaskStatusUpdated(
                        task_id=task.task_id, task_status=TaskStatus.Complete
                    )
@@ -391,7 +393,7 @@ class Worker:
                    ),
                )
                self.download_status[shard.model_meta.model_id] = status
-                await self.event_sender.send(
+                self.event_sender.send_nowait(
                    NodeDownloadProgress(download_progress=status)
                )
                last_progress_time = current_time()
--- a/src/exo/worker/runner/bootstrap.py
+++ b/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
--- a/src/exo/worker/runner/runner.py
+++ b/src/exo/worker/runner/runner.py
@@ -1,9 +1,13 @@
+import json
 import time
 from collections.abc import Generator
 from functools import cache
+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,
@@ -12,7 +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.chunks import TokenChunk, ToolCall, ToolCallFunction
 from exo.shared.types.events import (
    ChunkGenerated,
    Event,
@@ -67,7 +71,6 @@ def main(
        bound_instance.bound_runner_id,
        bound_instance.bound_shard,
    )
-    device_rank = shard_metadata.device_rank
    logger.info("hello from the runner")
    if getattr(shard_metadata, "immediate_exception", False):
        raise Exception("Fake exception - runner failed to spin up.")
@@ -119,20 +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
-                    )
+                    model, tokenizer = load_mlx_items(bound_instance, group)

                    current_status = RunnerLoaded()
                    logger.info("runner loaded")
@@ -162,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")
@@ -170,61 +162,45 @@ def main(
                            runner_id=runner_id, runner_status=current_status
                        )
                    )
-                    assert model
-                    assert tokenizer
                    assert task_params.messages[0].content is not None
+                    _check_for_debug_prompts(task_params.messages[0].content)

-                    try:
-                        _check_for_debug_prompts(task_params.messages[0].content)
+                    # Generate responses using the actual MLX generation
+                    mlx_generator = mlx_generate(
+                        model=model,
+                        tokenizer=tokenizer,
+                        task=task_params,
+                    )

-                        # 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)
+                    # Parse tool calls to place them in the tool calls section
+                    mlx_generator = parse_tool_calls(
+                        mlx_generator, tokenizer, task_params.tools
+                    )

-                        # TODO: Add tool call parser here
-
-                        for response in mlx_generator:
-                            match response:
-                                case GenerationResponse():
-                                    if 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,
+                                            ),
                                        )
-
-                    # can we make this more explicit?
-                    except Exception as e:
-                        if device_rank == 0:
-                            event_sender.send(
-                                ChunkGenerated(
-                                    command_id=command_id,
-                                    chunk=TokenChunk(
-                                        idx=0,
-                                        model=shard_metadata.model_meta.model_id,
-                                        text="",
-                                        token_id=0,
-                                        finish_reason="error",
-                                        error_message=str(e),
-                                    ),
-                                )
-                            )
-                        raise
+                                    )
+                                # case TokenizedResponse():
+                                # TODO: something here ig

                    current_status = RunnerReady()
                    logger.info("runner ready")
@@ -293,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"
--- a/src/exo/worker/tests/unittests/test_mlx/conftest.py
+++ b/src/exo/worker/tests/unittests/test_mlx/conftest.py
@@ -1,202 +0,0 @@
-# type: ignore
-from dataclasses import dataclass
-from pathlib import Path
-from typing import Any
-
-import mlx.core as mx
-import mlx.nn as nn
-
-from exo.shared.constants import EXO_MODELS_DIR
-
-
-class MockLayer(nn.Module):
-    def __init__(self) -> None:
-        super().__init__()
-        self.custom_attr = "test_value"
-        self.use_sliding = True
-
-    def __call__(self, x: mx.array, *args: object, **kwargs: object) -> mx.array:
-        return x * 2
-
-
-@dataclass(frozen=True)
-class PipelineTestConfig:
-    model_path: Path
-    total_layers: int
-    base_port: int
-    max_tokens: int
-
-
-def create_hostfile(world_size: int, base_port: int) -> tuple[str, list[str]]:
-    import json
-    import tempfile
-
-    hosts = [f"127.0.0.1:{base_port + i}" for i in range(world_size)]
-
-    with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f:
-        json.dump(hosts, f)
-        hostfile_path = f.name
-
-    return hostfile_path, hosts
-
-
-# Use GPT OSS 20b to test as it is a model with a lot of strange behaviour
-
-DEFAULT_GPT_OSS_CONFIG = PipelineTestConfig(
-    model_path=EXO_MODELS_DIR / "mlx-community--gpt-oss-20b-MXFP4-Q8",
-    total_layers=24,
-    base_port=29600,
-    max_tokens=200,
-)
-
-
-def run_gpt_oss_pipeline_device(
-    rank: int,
-    world_size: int,
-    hostfile_path: str,
-    model_path: Path,
-    layer_splits: list[tuple[int, int]],
-    prompt_tokens: int,
-    prefill_step_size: int,
-    result_queue: Any,  # pyright: ignore[reportAny]
-    max_tokens: int = 200,
-) -> None:
-    import os
-    import traceback
-
-    os.environ["MLX_HOSTFILE"] = hostfile_path
-    os.environ["MLX_RANK"] = str(rank)
-
-    import mlx.core as mlx_core
-    from mlx_lm import load, stream_generate
-
-    from exo.shared.types.memory import Memory
-    from exo.shared.types.models import ModelId, ModelMetadata
-    from exo.shared.types.worker.shards import PipelineShardMetadata
-    from exo.worker.engines.mlx.auto_parallel import pipeline_auto_parallel
-
-    try:
-        group = mlx_core.distributed.init(backend="ring", strict=True)
-
-        model, tokenizer = load(str(model_path))
-
-        # Generate a prompt of exact token length
-        base_text = "The quick brown fox jumps over the lazy dog. "
-        base_tokens = tokenizer.encode(base_text)
-        base_len = len(base_tokens)
-
-        # Build prompt with approximate target length
-        repeats = (prompt_tokens // base_len) + 2
-        long_text = base_text * repeats
-        tokens = tokenizer.encode(long_text)
-        # Truncate to exact target length
-        tokens = tokens[:prompt_tokens]
-        prompt_text = tokenizer.decode(tokens)
-
-        formatted_prompt = tokenizer.apply_chat_template(
-            [{"role": "user", "content": prompt_text}],
-            tokenize=False,
-            add_generation_prompt=True,
-        )
-
-        start_layer, end_layer = layer_splits[rank]
-
-        shard_meta = PipelineShardMetadata(
-            model_meta=ModelMetadata(
-                model_id=ModelId("mlx-community/gpt-oss-20b-MXFP4-Q8"),
-                pretty_name="GPT-OSS 20B",
-                storage_size=Memory.from_gb(12),
-                n_layers=24,
-                hidden_size=2880,
-                supports_tensor=False,
-            ),
-            device_rank=rank,
-            world_size=world_size,
-            start_layer=start_layer,
-            end_layer=end_layer,
-            n_layers=24,
-        )
-
-        model = pipeline_auto_parallel(model, group, shard_meta)
-
-        # Barrier before generation
-        barrier = mlx_core.distributed.all_sum(mlx_core.array([1.0]), group=group)
-        mlx_core.eval(barrier)
-
-        generated_text = ""
-        for response in stream_generate(
-            model=model,
-            tokenizer=tokenizer,
-            prompt=formatted_prompt,
-            max_tokens=max_tokens,
-            prefill_step_size=prefill_step_size,
-        ):
-            generated_text += response.text
-
-        result_queue.put((rank, True, generated_text))  # pyright: ignore[reportAny]
-
-    except Exception as e:
-        result_queue.put((rank, False, f"{e}\n{traceback.format_exc()}"))  # pyright: ignore[reportAny]
-
-
-def run_gpt_oss_tensor_parallel_device(
-    rank: int,
-    world_size: int,
-    hostfile_path: str,
-    model_path: Path,
-    prompt_tokens: int,
-    prefill_step_size: int,
-    result_queue: Any,  # pyright: ignore[reportAny]
-    max_tokens: int = 10,
-) -> None:
-    import os
-    import traceback
-
-    os.environ["MLX_HOSTFILE"] = hostfile_path
-    os.environ["MLX_RANK"] = str(rank)
-
-    import mlx.core as mlx_core
-    from mlx_lm import load, stream_generate
-
-    from exo.worker.engines.mlx.auto_parallel import tensor_auto_parallel
-
-    try:
-        group = mlx_core.distributed.init(backend="ring", strict=True)
-
-        model, tokenizer = load(str(model_path))
-
-        base_text = "The quick brown fox jumps over the lazy dog. "
-        base_tokens = tokenizer.encode(base_text)
-        base_len = len(base_tokens)
-
-        repeats = (prompt_tokens // base_len) + 2
-        long_text = base_text * repeats
-        tokens = tokenizer.encode(long_text)
-        tokens = tokens[:prompt_tokens]
-        prompt_text = tokenizer.decode(tokens)
-
-        formatted_prompt = tokenizer.apply_chat_template(
-            [{"role": "user", "content": prompt_text}],
-            tokenize=False,
-            add_generation_prompt=True,
-        )
-
-        model = tensor_auto_parallel(model, group)
-
-        barrier = mlx_core.distributed.all_sum(mlx_core.array([1.0]), group=group)
-        mlx_core.eval(barrier)
-
-        generated_text = ""
-        for response in stream_generate(
-            model=model,
-            tokenizer=tokenizer,
-            prompt=formatted_prompt,
-            max_tokens=max_tokens,
-            prefill_step_size=prefill_step_size,
-        ):
-            generated_text += response.text
-
-        result_queue.put((rank, True, generated_text))  # pyright: ignore[reportAny]
-
-    except Exception as e:
-        result_queue.put((rank, False, f"{e}\n{traceback.format_exc()}"))  # pyright: ignore[reportAny]
--- a/src/exo/worker/tests/unittests/test_mlx/test_auto_parallel.py
+++ b/src/exo/worker/tests/unittests/test_mlx/test_auto_parallel.py
@@ -1,137 +0,0 @@
-import multiprocessing as mp
-from typing import Any
-
-import mlx.core as mx
-import pytest
-
-from exo.worker.engines.mlx.auto_parallel import (
-    CustomMlxLayer,
-    PipelineFirstLayer,
-    PipelineLastLayer,
-)
-from exo.worker.tests.unittests.test_mlx.conftest import MockLayer
-
-
-def run_pipeline_device(
-    rank: int,
-    world_size: int,
-    hostfile_path: str,
-    result_queue: Any,  # pyright: ignore[reportAny]
-) -> None:
-    import os
-
-    os.environ["MLX_HOSTFILE"] = hostfile_path
-    os.environ["MLX_RANK"] = str(rank)
-
-    import mlx.core as mlx_core
-    import mlx.nn as mlx_nn
-
-    class MockLayerInner(mlx_nn.Module):
-        def __init__(self) -> None:
-            super().__init__()
-            self.custom_attr = "test_value"
-
-        def __call__(
-            self, x: mlx_core.array, *args: object, **kwargs: object
-        ) -> mlx_core.array:
-            return x * 2
-
-    try:
-        group = mlx_core.distributed.init(backend="ring", strict=True)
-
-        mock = MockLayerInner()
-        first = PipelineFirstLayer(mock, r=rank, group=group)
-        composed = PipelineLastLayer(first, r=rank, s=world_size, group=group)
-
-        x = mlx_core.ones((1, 4))
-        result = composed(x)
-        mlx_core.eval(result)
-
-        success = result.shape == x.shape
-        result_queue.put((rank, success, result))  # pyright: ignore[reportAny]
-    except Exception as e:
-        result_queue.put((rank, False, str(e)))  # pyright: ignore[reportAny]
-
-
-def test_single_wrapper_delegates_attributes() -> None:
-    mock = MockLayer()
-    wrapped = CustomMlxLayer(mock)
-
-    assert wrapped.custom_attr == "test_value"  # type: ignore[attr-defined]
-    assert wrapped.use_sliding is True  # type: ignore[attr-defined]
-
-
-def test_composed_wrappers_delegate_attributes() -> None:
-    mock = MockLayer()
-    group = mx.distributed.init()
-
-    first = PipelineFirstLayer(mock, r=0, group=group)
-    composed = PipelineLastLayer(first, r=0, s=1, group=group)
-
-    assert composed.custom_attr == "test_value"  # type: ignore[attr-defined]
-    assert composed.use_sliding is True  # type: ignore[attr-defined]
-
-
-def test_missing_attribute_raises() -> None:
-    mock = MockLayer()
-    wrapped = CustomMlxLayer(mock)
-
-    with pytest.raises(AttributeError):
-        _ = wrapped.nonexistent_attr  # type: ignore[attr-defined]
-
-
-def test_composed_call_works() -> None:
-    import json
-    import os
-    import tempfile
-
-    ctx = mp.get_context("spawn")
-
-    world_size = 2
-    base_port = 29500
-
-    hosts = [f"127.0.0.1:{base_port + i}" for i in range(world_size)]
-
-    with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f:
-        json.dump(hosts, f)
-        hostfile_path = f.name
-
-    try:
-        result_queue: Any = ctx.Queue()
-
-        processes: list[Any] = []
-        for rank in range(world_size):
-            p = ctx.Process(
-                target=run_pipeline_device,
-                args=(rank, world_size, hostfile_path, result_queue),
-            )
-            p.start()
-            processes.append(p)
-
-        for p in processes:  # pyright: ignore[reportAny]
-            p.join(timeout=10)  # pyright: ignore[reportAny]
-
-        results: dict[int, Any] = {}
-        errors: dict[int, str] = {}
-        while not result_queue.empty():  # pyright: ignore[reportAny]
-            rank, success, value = result_queue.get()  # pyright: ignore[reportAny]
-            if success:
-                results[rank] = value
-            else:
-                errors[rank] = value
-
-        assert len(results) == world_size, (
-            f"Expected {world_size} results, got {len(results)}. Errors: {errors}"
-        )
-
-        for rank in range(world_size):
-            assert rank in results, (
-                f"Device {rank} failed: {errors.get(rank, 'unknown')}"
-            )
-            result_array = results[rank]
-            # Both devices see the final result (4.0) after all_gather
-            assert (result_array == 4.0).all(), (
-                f"Device {rank}: expected 4.0, got {result_array}"
-            )
-    finally:
-        os.unlink(hostfile_path)
--- a/src/exo/worker/tests/unittests/test_runner/test_event_ordering.py
+++ b/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):
--- a/src/exo/worker/utils/net_profile.py
+++ b/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
--- a/uv.lock
+++ b/uv.lock
Author	SHA1	Message	Date
Ryuichi Leo Takashige	826da9512d	Add exo eval	2026-01-16 12:55:43 +00:00
Ryuichi Leo Takashige	4d9114b9b5	Add ChatCompletion tool calling support https://platform.openai.com/docs/api-reference/chat/get	2026-01-16 11:36:08 +00:00