Fix custom model add requiring two attempts + enlarge sidebar buttons (#1805)

## Motivation

Adding a custom model from the Hub tab shows "Added" toast but the model
doesn't appear in the All tab. You have to add it a second time for it
to work. Also, the "All" button in the model picker sidebar is too small
to read comfortably.

## Changes

**Race condition fix (`src/exo/api/main.py`):**
- Call `add_to_card_cache(card)` directly in `add_custom_model()` after
sending the `ForwarderCommand`, before the API response returns

**Sidebar sizing
(`dashboard/src/lib/components/FamilySidebar.svelte`):**
- Increased sidebar min-width from 72/64px to 80/72px
- Increased "All" icon from `w-5 h-5` to `w-6 h-6`
- Increased all sidebar labels from 9px to 11px

## Why It Works

`POST /models/add` sends a `ForwarderCommand(AddCustomModelCard)` and
returns immediately. The frontend then calls `GET /models` which reads
from `_card_cache`. But the cache was only updated by the worker event
handler after the event round-trips through the master — a race the
frontend almost always loses. By updating the cache directly in the API
handler, `GET /models` immediately reflects the new model. The worker's
later `add_to_card_cache` call is idempotent (dict key assignment).

## Test Plan

### Manual Testing
<!-- Hardware: any Mac -->
- Open model picker → Hub tab → add a custom model → verify it appears
in All tab on the first attempt
- Verify sidebar "All" button and other labels are visually larger and
readable

### Automated Testing
- `uv run basedpyright` passes with 0 errors
- `uv run ruff check` passes

---------

Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

.github/workflows/build-app.yml

@@ -159,7 +159,7 @@ jobs:
           fi
 
       - name: Install Homebrew packages
-        run: brew install just awscli macmon
+        run: brew install just awscli
 
       - name: Install UV
         uses: astral-sh/setup-uv@v6
@@ -243,6 +243,14 @@ jobs:
       # Build the bundle
       # ============================================================
 
+      - name: Add pinned macmon to PATH
+        run: |
+          MACMON_DIR=$(nix develop --command sh -c 'dirname $(which macmon)')
+          echo "Using macmon from: $MACMON_DIR"
+          echo "$MACMON_DIR" >> $GITHUB_PATH
+          # Remove any Homebrew macmon so PyInstaller can't accidentally pick it up
+          brew uninstall macmon 2>/dev/null || true
+
       - name: Build PyInstaller bundle
         run: uv run pyinstaller packaging/pyinstaller/exo.spec
 

.mlx_typings/mlx/core/__init__.pyi

@@ -2396,7 +2396,7 @@ def degrees(a: array, /, *, stream: Stream | Device | None = ...) -> array:
         array: The angles in degrees.
     """
 
-def depends(inputs: array | Sequence[array], dependencies: array | Sequence[array]):
+def depends[T](inputs: T, dependencies: array | Sequence[array]) -> T:
     """
     Insert dependencies between arrays in the graph. The outputs are
     identical to ``inputs`` but with dependencies on ``dependencies``.

.mlx_typings/mlx/nn/__init__.pyi

@@ -1,9 +1,5 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-from layers import *
-from utils import *
+from .layers import *
+from .utils import *
 
 from . import init as init
 from . import losses as losses

.mlx_typings/mlx/nn/layers/__init__.pyi

@@ -1,20 +1,16 @@
-"""
-This type stub file was generated by pyright.
-"""
-
-from activations import *
-from base import *
-from containers import *
-from convolution import *
-from convolution_transpose import *
-from distributed import *
-from dropout import *
-from embedding import *
-from linear import *
-from normalization import *
-from pooling import *
-from positional_encoding import *
-from quantized import *
-from recurrent import *
-from transformer import *
-from upsample import *
+from .activations import *
+from .base import *
+from .containers import *
+from .convolution import *
+from .convolution_transpose import *
+from .distributed import *
+from .dropout import *
+from .embedding import *
+from .linear import *
+from .normalization import *
+from .pooling import *
+from .positional_encoding import *
+from .quantized import *
+from .recurrent import *
+from .transformer import *
+from .upsample import *

.mlx_typings/mlx/nn/layers/base.pyi

@@ -53,7 +53,7 @@ class Module(dict):
         mx.eval(model.parameters())
     """
 
-    __call__: Callable
+    def __call__(self, *args: Any, **kwargs: Any) -> mx.array: ...
     def __init__(self) -> None:
         """Should be called by the subclasses of ``Module``."""
 

.mlx_typings/mlx_lm/generate.pyi

@@ -30,7 +30,7 @@ def str2bool(string):  # -> bool:
 def setup_arg_parser():  # -> ArgumentParser:
     """Set up and return the argument parser."""
 
-generation_stream = ...
+generation_stream: mx.Stream
 
 @contextlib.contextmanager
 def wired_limit(
@@ -266,12 +266,12 @@ def _merge_caches(caches: Any) -> List[Any]: ...
 class Batch:
     uids: List[int]
     y: mx.array
-    logprobs: mx.array
+    logprobs: List[mx.array] | mx.array
     max_tokens: List[int]
     num_tokens: List[int]
     cache: List[Any]
-    samplers: List[Any]
-    logits_processors: List[Any]
+    samplers: List[Callable[[mx.array], mx.array] | None]
+    logits_processors: List[List[Callable[[mx.array, mx.array], mx.array]]]
     tokens: List[mx.array]
     def __len__(self) -> int: ...
     def filter(self, keep_idx: List[int]) -> None: ...
@@ -279,13 +279,18 @@ class Batch:
     def extract_cache(self, idx: int) -> List[Any]: ...
 
 class BatchGenerator:
-    model: Any
+    model: nn.Module
+    sampler: Callable[[mx.array], mx.array]
+    stop_tokens: set[int]
     max_kv_size: Optional[int]
     prefill_step_size: int
+    completion_batch_size: int
+    prefill_batch_size: int
     unprocessed_prompts: List[Any]
     active_batch: Optional[Batch]
     prompt_progress_callback: Callable[[List[Tuple[int, int, int]]], None]
     _stats: BatchStats
+    _next_count: int
 
     @dataclass
     class Response:

.mlx_typings/mlx_lm/models/cache.pyi

@@ -88,8 +88,8 @@ def create_attention_mask(
 ) -> array | Literal["causal"] | None: ...
 
 class _BaseCache(Cache):
-    keys: mx.array
-    values: mx.array
+    keys: mx.array | None
+    values: mx.array | None
     offset: int
     @property
     def state(self) -> tuple[mx.array | None, mx.array | None]: ...
@@ -268,29 +268,14 @@ class CacheList(_BaseCache):
         """
 
 class BatchKVCache(_BaseCache):
-    step = ...
-    def __init__(self, left_padding: List[int]) -> None:
-        """
-        The BatchKV cache expects inputs to be left-padded.
-
-        E.g. the following prompts:
-
-            [1, 3, 5]
-            [7]
-            [2, 6, 8, 9]
-
-        Should be padded like so:
-
-            [0, 1, 3, 5]
-            [0, 0, 0, 7]
-            [2, 6, 8, 9]
-
-        And ``left_padding`` specifies the amount of padding for each.
-        In this case, ``left_padding = [1, 3, 0]``.
-        """
-
-    def update_and_fetch(self, keys, values):  # -> tuple[array | Any, array | Any]:
-        ...
+    step: int
+    keys: array | None
+    values: array | None
+    offset: array
+    left_padding: array
+    _idx: int
+    def __init__(self, left_padding: List[int]) -> None: ...
+    def update_and_fetch(self, keys: array, values: array) -> tuple[array, array]: ...
     @property
     def state(
         self,
@@ -316,12 +301,21 @@ class BatchKVCache(_BaseCache):
         """
 
 class BatchRotatingKVCache(_BaseCache):
-    step = ...
-    def __init__(self, max_size, left_padding: List[int]) -> None: ...
-    def update_and_fetch(
-        self, keys, values
-    ):  # -> tuple[array | Any, array | Any] | tuple[array | Any, array | Any | None]:
-        ...
+    step: int
+    keys: array | None
+    values: array | None
+    offset: array
+    left_padding: array
+    max_size: int
+    _idx: int
+    _offset: int
+    rotated: bool
+    _lengths: array | None
+    def __init__(self, max_size: int, left_padding: List[int]) -> None: ...
+    def _trim(self, trim_size: int, v: array, append: array | None = ...) -> array: ...
+    def _update_in_place(self, keys: array, values: array) -> tuple[array, array]: ...
+    def _update_concat(self, keys: array, values: array) -> tuple[array, array]: ...
+    def update_and_fetch(self, keys: array, values: array) -> tuple[array, array]: ...
     @property
     def state(
         self,

.mlx_typings/mlx_lm/models/gated_delta.pyi

@@ -0,0 +1,35 @@
+from typing import Optional
+
+import mlx.core as mx
+
+def compute_g(A_log: mx.array, a: mx.array, dt_bias: mx.array) -> mx.array: ...
+def gated_delta_update(
+    q: mx.array,
+    k: mx.array,
+    v: mx.array,
+    a: mx.array,
+    b: mx.array,
+    A_log: mx.array,
+    dt_bias: mx.array,
+    state: Optional[mx.array] = ...,
+    mask: Optional[mx.array] = ...,
+    use_kernel: bool = ...,
+) -> tuple[mx.array, mx.array]: ...
+def gated_delta_ops(
+    q: mx.array,
+    k: mx.array,
+    v: mx.array,
+    g: mx.array,
+    beta: mx.array,
+    state: Optional[mx.array] = ...,
+    mask: Optional[mx.array] = ...,
+) -> tuple[mx.array, mx.array]: ...
+def gated_delta_kernel(
+    q: mx.array,
+    k: mx.array,
+    v: mx.array,
+    g: mx.array,
+    beta: mx.array,
+    state: mx.array,
+    mask: Optional[mx.array] = ...,
+) -> tuple[mx.array, mx.array]: ...

.mlx_typings/mlx_lm/models/rope_utils.pyi

@@ -0,0 +1,51 @@
+from typing import Any, Optional
+
+import mlx.nn as nn
+
+class YarnRoPE(nn.Module):
+    def __init__(
+        self,
+        dims: int,
+        traditional: bool = ...,
+        max_position_embeddings: int = ...,
+        base: float = ...,
+        scaling_factor: float = ...,
+        original_max_position_embeddings: int = ...,
+        beta_fast: float = ...,
+        beta_slow: float = ...,
+        mscale: float = ...,
+        mscale_all_dim: float = ...,
+    ) -> None: ...
+
+class Llama3RoPE(nn.Module):
+    def __init__(
+        self,
+        dims: int,
+        traditional: bool = ...,
+        max_position_embeddings: int = ...,
+        base: float = ...,
+        scaling_factor: float = ...,
+        original_max_position_embeddings: int = ...,
+        low_freq_factor: float = ...,
+        high_freq_factor: float = ...,
+    ) -> None: ...
+
+class SuScaledRoPE(nn.Module):
+    def __init__(
+        self,
+        dims: int,
+        traditional: bool = ...,
+        max_position_embeddings: int = ...,
+        base: float = ...,
+        short_factor: Any = ...,
+        long_factor: Any = ...,
+        original_max_position_embeddings: int = ...,
+    ) -> None: ...
+
+def initialize_rope(
+    dims: int,
+    base: float = ...,
+    traditional: bool = ...,
+    scaling_config: Optional[dict[str, Any]] = ...,
+    max_position_embeddings: Optional[int] = ...,
+) -> nn.Module: ...

README.md

@@ -295,8 +295,9 @@ exo supports several environment variables for configuration:
 
 | Variable | Description | Default |
 |----------|-------------|---------|
-| `EXO_MODELS_PATH` | Colon-separated paths to search for pre-downloaded models (e.g., on NFS mounts or shared storage) | None |
-| `EXO_MODELS_DIR` | Directory where exo downloads and stores models | `~/.local/share/exo/models` (Linux) or `~/.exo/models` (macOS) |
+| `EXO_DEFAULT_MODELS_DIR` | Default directory for model downloads and caches. Always first in the writable dirs list. | `~/.local/share/exo/models` (Linux) or `~/.exo/models` (macOS) |
+| `EXO_MODELS_DIRS` | Colon-separated additional writable directories for model downloads. Checked in order after the default; first with enough free space is used. | None |
+| `EXO_MODELS_READ_ONLY_DIRS` | Colon-separated read-only directories to search for pre-downloaded models (e.g., NFS mounts, shared storage). Models here cannot be deleted. | None |
 | `EXO_OFFLINE` | Run without internet connection (uses only local models) | `false` |
 | `EXO_ENABLE_IMAGE_MODELS` | Enable image model support | `false` |
 | `EXO_LIBP2P_NAMESPACE` | Custom namespace for cluster isolation | None |
@@ -306,8 +307,11 @@ exo supports several environment variables for configuration:
 **Example usage:**
 
 ```bash
-# Use pre-downloaded models from NFS mount
-EXO_MODELS_PATH=/mnt/nfs/models:/opt/ai-models uv run exo
+# Use pre-downloaded models from NFS mount (read-only)
+EXO_MODELS_READ_ONLY_DIRS=/mnt/nfs/models:/opt/ai-models uv run exo
+
+# Download models to an external SSD (falls back to default dir if full)
+EXO_MODELS_DIRS=/Volumes/ExternalSSD/exo-models uv run exo
 
 # Run in offline mode
 EXO_OFFLINE=true uv run exo

app/EXO/EXO/ExoProcessController.swift

@@ -4,6 +4,7 @@ import Foundation
 
 private let customNamespaceKey = "EXOCustomNamespace"
 private let hfTokenKey = "EXOHFToken"
+private let hfEndpointKey = "EXOHFEndpoint"
 private let enableImageModelsKey = "EXOEnableImageModels"
 private let offlineModeKey = "EXOOfflineMode"
 private let onboardingCompletedKey = "EXOOnboardingCompleted"
@@ -53,6 +54,14 @@ final class ExoProcessController: ObservableObject {
             UserDefaults.standard.set(hfToken, forKey: hfTokenKey)
         }
     }
+    @Published var hfEndpoint: String = {
+        return UserDefaults.standard.string(forKey: hfEndpointKey) ?? ""
+    }()
+    {
+        didSet {
+            UserDefaults.standard.set(hfEndpoint, forKey: hfEndpointKey)
+        }
+    }
     @Published var enableImageModels: Bool = {
         return UserDefaults.standard.bool(forKey: enableImageModelsKey)
     }()
@@ -273,6 +282,9 @@ final class ExoProcessController: ObservableObject {
         if !hfToken.isEmpty {
             environment["HF_TOKEN"] = hfToken
         }
+        if !hfEndpoint.isEmpty {
+            environment["HF_ENDPOINT"] = hfEndpoint
+        }
         if enableImageModels {
             environment["EXO_ENABLE_IMAGE_MODELS"] = "true"
         }

app/EXO/EXO/Views/SettingsView.swift

@@ -12,6 +12,7 @@ struct SettingsView: View {
 
     @State private var pendingNamespace: String = ""
     @State private var pendingHFToken: String = ""
+    @State private var pendingHFEndpoint: String = ""
     @State private var pendingEnableImageModels = false
     @State private var pendingOfflineMode = false
     @State private var needsRestart = false
@@ -42,6 +43,7 @@ struct SettingsView: View {
         .onAppear {
             pendingNamespace = controller.customNamespace
             pendingHFToken = controller.hfToken
+            pendingHFEndpoint = controller.hfEndpoint
             pendingEnableImageModels = controller.enableImageModels
             pendingOfflineMode = controller.offlineMode
             needsRestart = false
@@ -74,6 +76,17 @@ struct SettingsView: View {
                     .foregroundColor(.secondary)
             }
 
+            Section {
+                LabeledContent("HuggingFace Endpoint") {
+                    TextField("default", text: $pendingHFEndpoint)
+                        .textFieldStyle(.roundedBorder)
+                        .frame(width: 200)
+                }
+                Text("Defaults to huggingface.co. Use a mirror (e.g. hf-mirror.com) for China.")
+                    .font(.caption)
+                    .foregroundColor(.secondary)
+            }
+
             Section {
                 Toggle("Offline Mode", isOn: $pendingOfflineMode)
                 Text("Skip internet checks and use only locally available models.")
@@ -454,6 +467,7 @@ struct SettingsView: View {
 
     private var hasGeneralChanges: Bool {
         pendingNamespace != controller.customNamespace || pendingHFToken != controller.hfToken
+            || pendingHFEndpoint != controller.hfEndpoint
             || pendingOfflineMode != controller.offlineMode
     }
 
@@ -464,6 +478,7 @@ struct SettingsView: View {
     private func applyGeneralSettings() {
         controller.customNamespace = pendingNamespace
         controller.hfToken = pendingHFToken
+        controller.hfEndpoint = pendingHFEndpoint
         controller.offlineMode = pendingOfflineMode
         restartIfRunning()
     }

bench/exo_bench.py

@@ -501,23 +501,21 @@ def main() -> int:
                                     for x, _ in batch_results
                                     if x["stats"]["generation_tps"] > 0
                                 ]
-                                agg_gen_tps = (
+                                per_req_tps = (
                                     mean(valid_gen_tps) if valid_gen_tps else 0.0
                                 )
-                                gen_tps = agg_gen_tps / concurrency
+                                agg_gen_tps = per_req_tps * concurrency
                                 logger.info(
                                     f"[concurrent {concurrency}x]  "
                                     f"agg_gen_tps={agg_gen_tps:.2f}  "
-                                    f"gen_tps={gen_tps:.2f}  "
+                                    f"per_req_tps={per_req_tps:.2f}  "
                                     f"errors={batch_errors}"
                                 )
 
                     if runs:
                         prompt_tps = mean(x["stats"]["prompt_tps"] for x in runs)
-                        gen_tps = mean(
-                            x["stats"]["generation_tps"] / x["concurrency"]
-                            for x in runs
-                        )
+                        per_req_tps = mean(x["stats"]["generation_tps"] for x in runs)
+                        gen_tps = per_req_tps * concurrency
                         ptok = mean(x["stats"]["prompt_tokens"] for x in runs)
                         gtok = mean(x["stats"]["generation_tokens"] for x in runs)
                         peak = mean(

bench/harness.py

@@ -377,7 +377,7 @@ def run_planning_phase(
                 f"have {avail // (1024**3)}GB. Use --danger-delete-downloads to free space."
             )
 
-        # Delete from smallest to largest (skip read-only models from EXO_MODELS_PATH)
+        # Delete from smallest to largest (skip read-only models)
         completed = [
             (
                 unwrap_instance(p["DownloadCompleted"]["shardMetadata"])["modelCard"][

dashboard/src/lib/components/FamilyLogos.svelte

@@ -88,6 +88,12 @@
       d="M22.012 0h1.032v.927H24v.968h-.956V3.78h-1.032V1.896h-1.878v-.97h1.878V0zM2.6 12.371V1.87h.969v10.502h-.97zm10.423.66h10.95v.918h-6.208v9.579h-4.742V13.03zM5.629 3.333v12.356H0v4.51h10.386V8L20.859 8l-.003-4.668-15.227.001z"
     />
   </svg>
+{:else if family === "nemotron"}
+  <svg class="w-6 h-6 {className}" viewBox="0 0 24 24" fill="currentColor">
+    <path
+      d="M8.948 8.798v-1.43a6.7 6.7 0 0 1 .424-.018c3.922-.124 6.493 3.374 6.493 3.374s-2.774 3.851-5.75 3.851c-.398 0-.787-.062-1.158-.185v-4.346c1.528.185 1.837.857 2.747 2.385l2.04-1.714s-1.492-1.952-4-1.952a6.016 6.016 0 0 0-.796.035m0-4.735v2.138l.424-.027c5.45-.185 9.01 4.47 9.01 4.47s-4.08 4.964-8.33 4.964c-.37 0-.733-.035-1.095-.097v1.325c.3.035.61.062.91.062 3.957 0 6.82-2.023 9.593-4.408.459.371 2.34 1.263 2.73 1.652-2.633 2.208-8.772 3.984-12.253 3.984-.335 0-.653-.018-.971-.053v1.864H24V4.063zm0 10.326v1.131c-3.657-.654-4.673-4.46-4.673-4.46s1.758-1.944 4.673-2.262v1.237H8.94c-1.528-.186-2.73 1.245-2.73 1.245s.68 2.412 2.739 3.11M2.456 10.9s2.164-3.197 6.5-3.533V6.201C4.153 6.59 0 10.653 0 10.653s2.35 6.802 8.948 7.42v-1.237c-4.84-.6-6.492-5.936-6.492-5.936z"
+    />
+  </svg>
 {:else}
   <svg class="w-6 h-6 {className}" viewBox="0 0 24 24" fill="currentColor">
     <path

dashboard/src/lib/components/FamilySidebar.svelte

@@ -31,6 +31,7 @@
     kimi: "Kimi",
     flux: "FLUX",
     "qwen-image": "Qwen Img",
+    nemotron: "NVIDIA",
   };
 
   function getFamilyName(family: string): string {
@@ -41,31 +42,20 @@
 </script>
 
 <div
-  class="flex flex-col gap-1 py-2 px-1 border-r border-exo-yellow/10 bg-exo-medium-gray/30 min-w-[72px] sm:min-w-[64px] overflow-y-auto scrollbar-hide"
+  class="flex flex-col gap-1 py-2 px-1 border-r border-exo-yellow/10 bg-exo-medium-gray/30 min-w-[80px] sm:min-w-[72px] overflow-y-auto scrollbar-hide"
 >
   <!-- All models (no filter) -->
   <button
     type="button"
     onclick={() => onSelect(null)}
-    class="group flex flex-col items-center justify-center p-2 sm:p-2 rounded transition-all duration-200 cursor-pointer min-h-[44px] sm:min-h-0 {selectedFamily ===
+    class="group flex items-center justify-center px-3 py-2.5 rounded transition-all duration-200 cursor-pointer min-h-[44px] sm:min-h-0 {selectedFamily ===
     null
       ? 'bg-exo-yellow/20 border-l-2 border-exo-yellow'
       : 'hover:bg-white/5 border-l-2 border-transparent'}"
     title="All models"
   >
-    <svg
-      class="w-5 h-5 {selectedFamily === null
-        ? 'text-exo-yellow'
-        : 'text-white/50 group-hover:text-white/70'}"
-      viewBox="0 0 24 24"
-      fill="currentColor"
-    >
-      <path
-        d="M4 8h4V4H4v4zm6 12h4v-4h-4v4zm-6 0h4v-4H4v4zm0-6h4v-4H4v4zm6 0h4v-4h-4v4zm6-10v4h4V4h-4zm-6 4h4V4h-4v4zm6 6h4v-4h-4v4zm0 6h4v-4h-4v4z"
-      />
-    </svg>
     <span
-      class="text-[9px] font-mono mt-0.5 {selectedFamily === null
+      class="text-[12px] font-mono font-medium {selectedFamily === null
         ? 'text-exo-yellow'
         : 'text-white/40 group-hover:text-white/60'}">All</span
     >
@@ -89,7 +79,7 @@
           : "text-white/50 group-hover:text-amber-400/70"}
       />
       <span
-        class="text-[9px] font-mono mt-0.5 {selectedFamily === 'favorites'
+        class="text-[11px] font-mono mt-0.5 {selectedFamily === 'favorites'
           ? 'text-amber-400'
           : 'text-white/40 group-hover:text-white/60'}">Faves</span
       >
@@ -114,7 +104,7 @@
           : "text-white/50 group-hover:text-white/70"}
       />
       <span
-        class="text-[9px] font-mono mt-0.5 {selectedFamily === 'recents'
+        class="text-[11px] font-mono mt-0.5 {selectedFamily === 'recents'
           ? 'text-exo-yellow'
           : 'text-white/40 group-hover:text-white/60'}">Recent</span
       >
@@ -138,7 +128,7 @@
         : "text-white/50 group-hover:text-orange-400/70"}
     />
     <span
-      class="text-[9px] font-mono mt-0.5 {selectedFamily === 'huggingface'
+      class="text-[11px] font-mono mt-0.5 {selectedFamily === 'huggingface'
         ? 'text-orange-400'
         : 'text-white/40 group-hover:text-white/60'}">Hub</span
     >
@@ -164,7 +154,7 @@
           : "text-white/50 group-hover:text-white/70"}
       />
       <span
-        class="text-[9px] font-mono mt-0.5 truncate max-w-full {selectedFamily ===
+        class="text-[11px] font-mono mt-0.5 truncate max-w-full {selectedFamily ===
         family
           ? 'text-exo-yellow'
           : 'text-white/40 group-hover:text-white/60'}"

dashboard/src/lib/components/ModelFilterPopover.svelte

@@ -73,7 +73,7 @@
 
 <!-- svelte-ignore a11y_no_static_element_interactions -->
 <div
-  class="filter-popover absolute right-0 top-full mt-2 w-64 bg-exo-dark-gray border border-exo-yellow/10 rounded-lg shadow-xl z-10"
+  class="filter-popover absolute right-0 top-full mt-2 w-64 bg-exo-dark-gray border border-exo-yellow/10 rounded-lg shadow-xl z-20"
   transition:fly={{ y: -10, duration: 200, easing: cubicOut }}
   onclick={(e) => e.stopPropagation()}
   role="dialog"

dashboard/src/lib/components/ModelPickerModal.svelte

@@ -459,6 +459,7 @@
       "llama",
       "flux",
       "qwen-image",
+      "nemotron",
     ];
     return Array.from(families).sort((a, b) => {
       const aIdx = familyOrder.indexOf(a);

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

@@ -1793,6 +1793,14 @@ class AppStore {
             this.persistConversation(targetConversationId);
           }
         },
+        {
+          generation_stats: (data) => {
+            const stats = data as { generation_tps: number };
+            if (stats.generation_tps > 0) {
+              this.tps = stats.generation_tps;
+            }
+          },
+        },
       );
 
       // Final update
@@ -1990,6 +1998,14 @@ class AppStore {
             this.persistConversation(targetConversationId);
           }
         },
+        {
+          generation_stats: (data) => {
+            const stats = data as { generation_tps: number };
+            if (stats.generation_tps > 0) {
+              this.tps = stats.generation_tps;
+            }
+          },
+        },
       );
 
       // Final cleanup of the message (if conversation still exists)
@@ -2397,7 +2413,7 @@ class AppStore {
 
       let streamedContent = "";
       let streamedThinking = "";
-
+      let serverTpsReceived = false;
       interface ChatCompletionChunk {
         choices?: Array<{
           delta?: { content?: string; reasoning_content?: string };
@@ -2462,7 +2478,6 @@ class AppStore {
             tokenCount += 1;
             this.totalTokens = tokenCount;
 
-            // Update real-time TPS during streaming
             if (firstTokenTime !== null && tokenCount > 1) {
               const elapsed = performance.now() - firstTokenTime;
               this.tps = (tokenCount / elapsed) * 1000;
@@ -2513,16 +2528,24 @@ class AppStore {
               startedAt: this.prefillProgress?.startedAt ?? performance.now(),
             };
           },
+          generation_stats: (data) => {
+            const stats = data as { generation_tps: number };
+
+            if (stats.generation_tps > 0) {
+              this.tps = stats.generation_tps;
+              serverTpsReceived = true;
+            }
+          },
         },
       );
 
       // Clear prefill progress after stream ends
       this.prefillProgress = null;
 
-      // Calculate final TPS
-      if (firstTokenTime !== null && tokenCount > 1) {
+      // Use server-side TPS if available, otherwise fall back to client-side
+      if (!serverTpsReceived && firstTokenTime !== null && tokenCount > 1) {
         const totalGenerationTime = performance.now() - firstTokenTime;
-        this.tps = (tokenCount / totalGenerationTime) * 1000; // tokens per second
+        this.tps = (tokenCount / totalGenerationTime) * 1000;
       }
 
       // Final cleanup of the message (if conversation still exists)
@@ -2627,6 +2650,9 @@ class AppStore {
     this.syncActiveMessagesIfNeeded(targetConversationId);
     this.saveConversationsToStorage();
 
+    const abortController = new AbortController();
+    this.currentAbortController = abortController;
+
     try {
       // Determine the model to use
       const model = this.getModelForRequest(modelId);
@@ -2681,6 +2707,7 @@ class AppStore {
           "Content-Type": "application/json",
         },
         body: JSON.stringify(requestBody),
+        signal: abortController.signal,
       });
 
       if (!response.ok) {
@@ -2820,14 +2847,27 @@ class AppStore {
         );
       }
     } catch (error) {
-      console.error("Error generating image:", error);
-      this.handleStreamingError(
-        error,
-        targetConversationId,
-        assistantMessage.id,
-        "Failed to generate image",
-      );
+      if (abortController.signal.aborted) {
+        this.updateConversationMessage(
+          targetConversationId,
+          assistantMessage.id,
+          (msg) => {
+            msg.content = "Cancelled";
+            msg.attachments = [];
+          },
+        );
+        this.syncActiveMessagesIfNeeded(targetConversationId);
+      } else {
+        console.error("Error generating image:", error);
+        this.handleStreamingError(
+          error,
+          targetConversationId,
+          assistantMessage.id,
+          "Failed to generate image",
+        );
+      }
     } finally {
+      this.currentAbortController = null;
       this.isLoading = false;
       this.saveConversationsToStorage();
     }
@@ -2891,6 +2931,9 @@ class AppStore {
     // Clear editing state
     this.editingImage = null;
 
+    const abortController = new AbortController();
+    this.currentAbortController = abortController;
+
     try {
       // Determine the model to use
       const model = this.getModelForRequest(modelId);
@@ -2952,6 +2995,7 @@ class AppStore {
       const apiResponse = await fetch("/v1/images/edits", {
         method: "POST",
         body: formData,
+        signal: abortController.signal,
       });
 
       if (!apiResponse.ok) {
@@ -3052,14 +3096,27 @@ class AppStore {
         );
       }
     } catch (error) {
-      console.error("Error editing image:", error);
-      this.handleStreamingError(
-        error,
-        targetConversationId,
-        assistantMessage.id,
-        "Failed to edit image",
-      );
+      if (abortController.signal.aborted) {
+        this.updateConversationMessage(
+          targetConversationId,
+          assistantMessage.id,
+          (msg) => {
+            msg.content = "Cancelled";
+            msg.attachments = [];
+          },
+        );
+        this.syncActiveMessagesIfNeeded(targetConversationId);
+      } else {
+        console.error("Error editing image:", error);
+        this.handleStreamingError(
+          error,
+          targetConversationId,
+          assistantMessage.id,
+          "Failed to edit image",
+        );
+      }
     } finally {
+      this.currentAbortController = null;
       this.isLoading = false;
       this.saveConversationsToStorage();
     }

dashboard/src/routes/+page.svelte

@@ -42,6 +42,7 @@
     setSelectedChatModel,
     selectedChatModel,
     sendMessage,
+    thinkingEnabled,
     generateImage,
     editImage,
     editingImage,
@@ -852,7 +853,7 @@
   ) {
     const model = selectedChatModel();
     if (!model) {
-      sendMessage(content, files, null);
+      sendMessage(content, files, thinkingEnabled());
       return;
     }
 
@@ -880,7 +881,7 @@
     }
 
     // Default: text chat
-    sendMessage(content, files, null);
+    sendMessage(content, files, thinkingEnabled());
   }
 
   let selectedSharding = $state<"Pipeline" | "Tensor">("Pipeline");
@@ -4574,7 +4575,7 @@
                 type="button"
                 onclick={() => {
                   completeOnboarding();
-                  sendMessage(chip);
+                  sendMessage(chip, undefined, thinkingEnabled());
                 }}
                 class="px-4 py-2 rounded-full border border-white/10 bg-white/5 text-sm text-white/60 hover:bg-white/10 hover:text-white/80 hover:border-white/20 transition-all duration-200 cursor-pointer"
               >
@@ -6069,7 +6070,7 @@
                           onclick={() => {
                             chatLaunchState = "idle";
                             selectedChatCategory = null;
-                            sendMessage(prompt);
+                            sendMessage(prompt, undefined, thinkingEnabled());
                           }}
                           class="text-left px-3 py-2.5 text-xs text-exo-light-gray hover:text-white font-mono rounded-lg border border-exo-medium-gray/30 hover:border-exo-yellow/30 bg-exo-dark-gray/30 hover:bg-exo-dark-gray/60 transition-all duration-200 cursor-pointer"
                         >

justfile

@@ -1,5 +1,8 @@
 export NIX_CONFIG := "extra-experimental-features = nix-command flakes"
 
+default: lint fmt
+all: lint fmt check
+
 fmt:
     treefmt || nix fmt
 
@@ -31,6 +34,10 @@ build-dashboard:
 package:
     uv run pyinstaller packaging/pyinstaller/exo.spec
 
+build-app: package
+    xcodebuild build -project app/EXO/EXO.xcodeproj -scheme EXO -configuration Debug -derivedDataPath app/EXO/build
+    @echo "\nBuild complete. Run with:\n  open {{justfile_directory()}}/app/EXO/build/Build/Products/Debug/EXO.app"
+
 clean:
     rm -rf **/__pycache__
     rm -rf target/

pyproject.toml

@@ -25,7 +25,7 @@ dependencies = [
   "openai-harmony>=0.0.8",
   "httpx>=0.28.1",
   "tomlkit>=0.14.0",
-  "mflux==0.16.9",
+  "mflux==0.17.2",
   "python-multipart>=0.0.21",
   "msgspec>=0.19.0",
   "zstandard>=0.23.0",
@@ -61,7 +61,7 @@ members = ["rust/exo_pyo3_bindings", "bench"]
 [tool.uv.sources]
 exo_pyo3_bindings = { workspace = true }
 mlx = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git", branch = "address-rdma-gpu-locks", marker = "sys_platform == 'darwin'" }
-mlx-lm = { git = "https://github.com/rltakashige/mlx-lm", branch = "leo/eval-left-padding-in-batched-rotation" }
+mlx-lm = { git = "https://github.com/rltakashige/mlx-lm", branch = "leo/fix-deepseek-v32-indexer" }
 # Uncomment to use local mlx/mlx-lm development versions:
 # mlx = { path = "/Users/Shared/mlx", editable=true }
 # mlx-lm = { path = "/Users/Shared/mlx-lm", editable=true }

resources/inference_model_cards/mlx-community--DeepSeek-V3.2-4bit.toml

@@ -0,0 +1,13 @@
+model_id = "mlx-community/DeepSeek-V3.2-4bit"
+n_layers = 61
+hidden_size = 7168
+num_key_value_heads = 128
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "deepseek"
+quantization = "4bit"
+base_model = "DeepSeek V3.2"
+capabilities = ["text", "thinking", "thinking_toggle"]
+
+[storage_size]
+in_bytes = 378086226621

resources/inference_model_cards/mlx-community--DeepSeek-V3.2-8bit.toml

@@ -0,0 +1,13 @@
+model_id = "mlx-community/DeepSeek-V3.2-8bit"
+n_layers = 61
+hidden_size = 7168
+num_key_value_heads = 128
+supports_tensor = true
+tasks = ["TextGeneration"]
+family = "deepseek"
+quantization = "8bit"
+base_model = "DeepSeek V3.2"
+capabilities = ["text", "thinking", "thinking_toggle"]
+
+[storage_size]
+in_bytes = 755957120916

rust/exo_pyo3_bindings/exo_pyo3_bindings.pyi

@@ -42,7 +42,7 @@ class MessageTooLargeError(builtins.Exception):
 
 @typing.final
 class NetworkingHandle:
-    def __new__(cls, identity: Keypair) -> NetworkingHandle: ...
+    def __new__(cls, identity: Keypair, bootstrap_peers: typing.Sequence[builtins.str], listen_port: builtins.int) -> NetworkingHandle: ...
     async def gossipsub_subscribe(self, topic: builtins.str) -> builtins.bool:
         r"""
         Subscribe to a `GossipSub` topic.

rust/exo_pyo3_bindings/src/networking.rs

@@ -180,7 +180,12 @@ impl PyNetworkingHandle {
     // ---- Lifecycle management methods ----
 
     #[new]
-    fn py_new(identity: Bound<'_, PyKeypair>) -> PyResult<Self> {
+    #[pyo3(signature = (identity, bootstrap_peers, listen_port))]
+    fn py_new(
+        identity: Bound<'_, PyKeypair>,
+        bootstrap_peers: Vec<String>,
+        listen_port: u16,
+    ) -> PyResult<Self> {
         // create communication channels
         let (to_swarm, from_client) = mpsc::channel(MPSC_CHANNEL_SIZE);
 
@@ -189,7 +194,9 @@ impl PyNetworkingHandle {
 
         // create networking swarm (within tokio context!! or it crashes)
         let _guard = pyo3_async_runtimes::tokio::get_runtime().enter();
-        let swarm = create_swarm(identity, from_client).pyerr()?.into_stream();
+        let swarm = create_swarm(identity, from_client, bootstrap_peers, listen_port)
+            .pyerr()?
+            .into_stream();
 
         Ok(Self {
             swarm: Arc::new(Mutex::new(swarm)),

rust/networking/examples/chatroom.rs

@@ -16,9 +16,14 @@ async fn main() {
     let (to_swarm, from_client) = mpsc::channel(20);
 
     // Configure swarm
-    let mut swarm = swarm::create_swarm(identity::Keypair::generate_ed25519(), from_client)
-        .expect("Swarm creation failed")
-        .into_stream();
+    let mut swarm = swarm::create_swarm(
+        identity::Keypair::generate_ed25519(),
+        from_client,
+        vec![],
+        0,
+    )
+    .expect("Swarm creation failed")
+    .into_stream();
 
     // Create a Gossipsub topic & subscribe
     let (tx, rx) = oneshot::channel();

rust/networking/src/discovery.rs

@@ -104,6 +104,7 @@ pub struct Behaviour {
     // state-tracking for managed behaviors & mDNS-discovered peers
     managed: managed::Behaviour,
     mdns_discovered: HashMap<PeerId, BTreeSet<Multiaddr>>,
+    bootstrap_peers: Vec<Multiaddr>,
 
     retry_delay: Delay, // retry interval
 
@@ -112,10 +113,11 @@ pub struct Behaviour {
 }
 
 impl Behaviour {
-    pub fn new(keypair: &identity::Keypair) -> io::Result<Self> {
+    pub fn new(keypair: &identity::Keypair, bootstrap_peers: Vec<Multiaddr>) -> io::Result<Self> {
         Ok(Self {
             managed: managed::Behaviour::new(keypair)?,
             mdns_discovered: HashMap::new(),
+            bootstrap_peers,
             retry_delay: Delay::new(RETRY_CONNECT_INTERVAL),
             pending_events: WakerDeque::new(),
         })
@@ -368,6 +370,12 @@ impl NetworkBehaviour for Behaviour {
                     self.dial(p, ma)
                 }
             }
+            // dial bootstrap peers (for environments where mDNS is unavailable)
+            for addr in &self.bootstrap_peers {
+                self.pending_events.push_back(ToSwarm::Dial {
+                    opts: DialOpts::unknown_peer_id().address(addr.clone()).build(),
+                })
+            }
             self.retry_delay.reset(RETRY_CONNECT_INTERVAL) // reset timeout
         }
 

rust/networking/src/swarm.rs

@@ -142,19 +142,29 @@ fn filter_swarm_event(event: SwarmEvent<BehaviourEvent>) -> Option<FromSwarm> {
     }
 }
 
-/// Create and configure a swarm which listens to all ports on OS
+/// Create and configure a swarm.
+///
+/// - `listen_port`: TCP port to listen on. `0` lets the OS assign one.
+/// - `bootstrap_peers`: multiaddrs to dial for environments without mDNS.
 pub fn create_swarm(
     keypair: identity::Keypair,
     from_client: mpsc::Receiver<ToSwarm>,
+    bootstrap_peers: Vec<String>,
+    listen_port: u16,
 ) -> alias::AnyResult<Swarm> {
+    let parsed_bootstrap_peers: Vec<libp2p::Multiaddr> = bootstrap_peers
+        .iter()
+        .filter(|s| !s.is_empty())
+        .filter_map(|s| s.parse().ok())
+        .collect();
+
     let mut swarm = SwarmBuilder::with_existing_identity(keypair)
         .with_tokio()
         .with_other_transport(tcp_transport)?
-        .with_behaviour(Behaviour::new)?
+        .with_behaviour(|keypair| Behaviour::new(keypair, parsed_bootstrap_peers))?
         .build();
 
-    // Listen on all interfaces and whatever port the OS assigns
-    swarm.listen_on("/ip4/0.0.0.0/tcp/0".parse()?)?;
+    swarm.listen_on(format!("/ip4/0.0.0.0/tcp/{listen_port}").parse()?)?;
     Ok(Swarm { swarm, from_client })
 }
 
@@ -246,9 +256,12 @@ mod behaviour {
     }
 
     impl Behaviour {
-        pub fn new(keypair: &identity::Keypair) -> alias::AnyResult<Self> {
+        pub fn new(
+            keypair: &identity::Keypair,
+            bootstrap_peers: Vec<libp2p::Multiaddr>,
+        ) -> alias::AnyResult<Self> {
             Ok(Self {
-                discovery: discovery::Behaviour::new(keypair)?,
+                discovery: discovery::Behaviour::new(keypair, bootstrap_peers)?,
                 gossipsub: gossipsub_behaviour(keypair),
             })
         }

rust/networking/tests/bootstrap_peers.rs

@@ -0,0 +1,107 @@
+use futures_lite::StreamExt;
+use networking::swarm::{FromSwarm, create_swarm};
+use std::time::Duration;
+use tokio::sync::mpsc;
+use tokio::time::timeout;
+
+/// Helper: find a free TCP port.
+fn free_port() -> u16 {
+    let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
+    listener.local_addr().unwrap().port()
+}
+
+/// Two nodes connect via bootstrap peers — no mDNS needed.
+///
+/// Node A listens on a fixed port. Node B bootstraps to A's address.
+/// We verify that B emits `FromSwarm::Discovered` for A's peer ID.
+#[tokio::test]
+async fn two_nodes_connect_via_bootstrap_peers() {
+    let port_a = free_port();
+
+    // Node A: listens on a known port, no bootstrap peers
+    let keypair_a = libp2p::identity::Keypair::generate_ed25519();
+    let peer_id_a = keypair_a.public().to_peer_id();
+    let (_tx_a, rx_a) = mpsc::channel(16);
+    let swarm_a = create_swarm(keypair_a, rx_a, vec![], port_a).expect("create swarm A");
+    let mut stream_a = swarm_a.into_stream();
+
+    // Node B: bootstraps to A's address
+    let keypair_b = libp2p::identity::Keypair::generate_ed25519();
+    let (_tx_b, rx_b) = mpsc::channel(16);
+    let swarm_b = create_swarm(
+        keypair_b,
+        rx_b,
+        vec![format!("/ip4/127.0.0.1/tcp/{port_a}")],
+        0,
+    )
+    .expect("create swarm B");
+    let mut stream_b = swarm_b.into_stream();
+
+    // Wait for B to discover A (connection established)
+    let connected = timeout(Duration::from_secs(10), async {
+        loop {
+            tokio::select! {
+                Some(event) = stream_a.next() => {
+                    // A will also see B connect, but we check from B's perspective
+                    let _ = event;
+                }
+                Some(event) = stream_b.next() => {
+                    if let FromSwarm::Discovered { peer_id } = event {
+                        if peer_id == peer_id_a {
+                            return true;
+                        }
+                    }
+                }
+            }
+        }
+    })
+    .await;
+
+    assert!(
+        connected.is_ok() && connected.unwrap(),
+        "Node B should discover Node A via bootstrap peer"
+    );
+}
+
+/// Empty bootstrap peers should work (backward compatible).
+#[tokio::test]
+async fn create_swarm_with_empty_bootstrap_peers() {
+    let keypair = libp2p::identity::Keypair::generate_ed25519();
+    let (_tx, rx) = mpsc::channel(16);
+    let swarm = create_swarm(keypair, rx, vec![], 0);
+    assert!(
+        swarm.is_ok(),
+        "create_swarm with no bootstrap peers should succeed"
+    );
+}
+
+/// Invalid multiaddr strings are silently filtered out.
+#[tokio::test]
+async fn create_swarm_ignores_invalid_bootstrap_addrs() {
+    let keypair = libp2p::identity::Keypair::generate_ed25519();
+    let (_tx, rx) = mpsc::channel(16);
+    let swarm = create_swarm(
+        keypair,
+        rx,
+        vec![
+            "not-a-valid-multiaddr".to_string(),
+            "".to_string(),
+            "/ip4/10.0.0.1/tcp/30000".to_string(), // valid
+        ],
+        0,
+    );
+    assert!(
+        swarm.is_ok(),
+        "create_swarm should succeed even with invalid bootstrap addrs"
+    );
+}
+
+/// Fixed listen port works correctly.
+#[tokio::test]
+async fn create_swarm_with_fixed_port() {
+    let port = free_port();
+    let keypair = libp2p::identity::Keypair::generate_ed25519();
+    let (_tx, rx) = mpsc::channel(16);
+    let swarm = create_swarm(keypair, rx, vec![], port);
+    assert!(swarm.is_ok(), "create_swarm with fixed port should succeed");
+}

src/exo/api/adapters/chat_completions.py

@@ -202,6 +202,8 @@ async def generate_chat_stream(
                     usage=last_usage,
                 )
                 yield f"data: {tool_response.model_dump_json()}\n\n"
+                if chunk.stats is not None:
+                    yield f": generation_stats {chunk.stats.model_dump_json()}\n\n"
                 yield "data: [DONE]\n\n"
                 return
 
@@ -216,7 +218,10 @@ async def generate_chat_stream(
                 yield f"data: {chunk_response.model_dump_json()}\n\n"
 
                 if chunk.finish_reason is not None:
+                    if chunk.stats is not None:
+                        yield f": generation_stats {chunk.stats.model_dump_json()}\n\n"
                     yield "data: [DONE]\n\n"
+                    return
 
 
 async def collect_chat_response(

src/exo/api/main.py

@@ -129,9 +129,9 @@ from exo.shared.logging import InterceptLogger
 from exo.shared.models.model_cards import (
     ModelCard,
     ModelId,
-    delete_custom_card,
+    add_to_card_cache,
+    get_card,
     get_model_cards,
-    is_custom_card,
 )
 from exo.shared.tracing import TraceEvent, compute_stats, export_trace, load_trace_file
 from exo.shared.types.chunks import (
@@ -143,8 +143,10 @@ from exo.shared.types.chunks import (
     ToolCallChunk,
 )
 from exo.shared.types.commands import (
+    AddCustomModelCard,
     Command,
     CreateInstance,
+    DeleteCustomModelCard,
     DeleteDownload,
     DeleteInstance,
     DownloadCommand,
@@ -413,6 +415,7 @@ class API:
                 node_network=self.state.node_network,
                 topology=self.state.topology,
                 current_instances=self.state.instances,
+                download_status=self.state.downloads,
             )
         except ValueError as exc:
             raise HTTPException(status_code=400, detail=str(exc)) from exc
@@ -475,6 +478,7 @@ class API:
                     topology=self.state.topology,
                     current_instances=self.state.instances,
                     required_nodes=required_nodes,
+                    download_status=self.state.downloads,
                 )
             except ValueError as exc:
                 if (model_card.model_id, sharding, instance_meta, 0) not in seen:
@@ -1558,7 +1562,7 @@ class API:
                     storage_size_megabytes=card.storage_size.in_mb,
                     supports_tensor=card.supports_tensor,
                     tasks=[task.value for task in card.tasks],
-                    is_custom=is_custom_card(card.model_id),
+                    is_custom=card.is_custom,
                     family=card.family,
                     quantization=card.quantization,
                     base_model=card.base_model,
@@ -1569,7 +1573,7 @@ class API:
         )
 
     async def add_custom_model(self, payload: AddCustomModelParams) -> ModelListModel:
-        """Fetch a model from HuggingFace and save as a custom model card."""
+        """Fetch a model from HuggingFace and save as a custom model card, then sync across the cluster."""
         try:
             card = await ModelCard.fetch_from_hf(payload.model_id)
         except Exception as exc:
@@ -1577,6 +1581,17 @@ class API:
                 status_code=400, detail=f"Failed to fetch model: {exc}"
             ) from exc
 
+        await self.command_sender.send(
+            ForwarderCommand(
+                origin=self._system_id,
+                command=AddCustomModelCard(model_card=card),
+            )
+        )
+
+        # Immediately update the local cache so the subsequent GET /models
+        # returns the new model without waiting for the event round-trip.
+        add_to_card_cache(card)
+
         return ModelListModel(
             id=card.model_id,
             hugging_face_id=card.model_id,
@@ -1590,10 +1605,18 @@ class API:
         )
 
     async def delete_custom_model(self, model_id: ModelId) -> JSONResponse:
-        """Delete a user-added custom model card."""
-        deleted = await delete_custom_card(model_id)
-        if not deleted:
+        """Delete a user-added custom model card and sync deletion across the cluster."""
+        card = get_card(model_id)
+        if card is None or not card.is_custom:
             raise HTTPException(status_code=404, detail="Custom model card not found")
+
+        await self.command_sender.send(
+            ForwarderCommand(
+                origin=self._system_id,
+                command=DeleteCustomModelCard(model_id=model_id),
+            )
+        )
+
         return JSONResponse(
             {"message": "Model card deleted", "model_id": str(model_id)}
         )

src/exo/download/coordinator.py

@@ -1,18 +1,21 @@
+from __future__ import annotations
+
 from dataclasses import dataclass, field
+from pathlib import Path
 
 import anyio
-from anyio import current_time
+from anyio import current_time, to_thread
 from loguru import logger
 
 from exo.download.download_utils import (
     RepoDownloadProgress,
     delete_model,
-    is_model_directory_complete,
+    is_read_only_model_dir,
     map_repo_download_progress_to_download_progress_data,
-    resolve_model_in_path,
+    resolve_existing_model,
 )
 from exo.download.shard_downloader import ShardDownloader
-from exo.shared.constants import EXO_MODELS_DIR, EXO_MODELS_PATH
+from exo.shared.constants import EXO_DEFAULT_MODELS_DIR, EXO_MODELS_READ_ONLY_DIRS
 from exo.shared.models.model_cards import ModelId, get_model_cards
 from exo.shared.types.commands import (
     CancelDownload,
@@ -25,6 +28,7 @@ from exo.shared.types.events import (
     Event,
     NodeDownloadProgress,
 )
+from exo.shared.types.memory import Memory
 from exo.shared.types.worker.downloads import (
     DownloadCompleted,
     DownloadFailed,
@@ -50,6 +54,7 @@ class DownloadCoordinator:
     active_downloads: dict[ModelId, anyio.CancelScope] = field(default_factory=dict)
 
     _tg: TaskGroup = field(init=False, default_factory=TaskGroup)
+    _stopped: anyio.Event = field(init=False, default_factory=anyio.Event)
 
     # Per-model throttle for download progress events
     _last_progress_time: dict[ModelId, float] = field(default_factory=dict)
@@ -57,8 +62,23 @@ class DownloadCoordinator:
     def __post_init__(self) -> None:
         self.shard_downloader.on_progress(self._download_progress_callback)
 
-    def _model_dir(self, model_id: ModelId) -> str:
-        return str(EXO_MODELS_DIR / model_id.normalize())
+    @staticmethod
+    def _default_model_dir(model_id: ModelId) -> str:
+        return str(EXO_DEFAULT_MODELS_DIR / model_id.normalize())
+
+    def _completed_from_path(
+        self,
+        shard: ShardMetadata,
+        found: Path,
+        total: Memory,
+    ) -> DownloadCompleted:
+        return DownloadCompleted(
+            shard_metadata=shard,
+            node_id=self.node_id,
+            total=total,
+            model_directory=str(found),
+            read_only=is_read_only_model_dir(found),
+        )
 
     async def _download_progress_callback(
         self, callback_shard: ShardMetadata, progress: RepoDownloadProgress
@@ -67,12 +87,18 @@ class DownloadCoordinator:
         throttle_interval_secs = 1.0
 
         if progress.status == "complete":
-            completed = DownloadCompleted(
-                shard_metadata=callback_shard,
-                node_id=self.node_id,
-                total=progress.total,
-                model_directory=self._model_dir(model_id),
-            )
+            found = await to_thread.run_sync(resolve_existing_model, model_id)
+            if found is not None:
+                completed = self._completed_from_path(
+                    callback_shard, found, progress.total
+                )
+            else:
+                completed = DownloadCompleted(
+                    shard_metadata=callback_shard,
+                    node_id=self.node_id,
+                    total=progress.total,
+                    model_directory=self._default_model_dir(model_id),
+                )
             self.download_status[model_id] = completed
             await self.event_sender.send(
                 NodeDownloadProgress(download_progress=completed)
@@ -89,7 +115,7 @@ class DownloadCoordinator:
                 download_progress=map_repo_download_progress_to_download_progress_data(
                     progress
                 ),
-                model_directory=self._model_dir(model_id),
+                model_directory=self._default_model_dir(model_id),
             )
             self.download_status[model_id] = ongoing
             await self.event_sender.send(
@@ -101,12 +127,16 @@ class DownloadCoordinator:
         logger.info(
             f"Starting DownloadCoordinator{' (offline mode)' if self.offline else ''}"
         )
-        async with self._tg as tg:
-            tg.start_soon(self._command_processor)
-            tg.start_soon(self._emit_existing_download_progress)
+        try:
+            async with self._tg as tg:
+                tg.start_soon(self._command_processor)
+                tg.start_soon(self._emit_existing_download_progress)
+        finally:
+            self._stopped.set()
 
-    def shutdown(self) -> None:
+    async def shutdown(self) -> None:
         self._tg.cancel_tasks()
+        await self._stopped.wait()
 
     async def _command_processor(self) -> None:
         with self.download_command_receiver as commands:
@@ -131,7 +161,7 @@ class DownloadCoordinator:
             pending = DownloadPending(
                 shard_metadata=current_status.shard_metadata,
                 node_id=self.node_id,
-                model_directory=self._model_dir(model_id),
+                model_directory=self._default_model_dir(model_id),
             )
             self.download_status[model_id] = pending
             await self.event_sender.send(
@@ -150,35 +180,12 @@ class DownloadCoordinator:
                 )
                 return
 
-        # Check EXO_MODELS_PATH for pre-downloaded models
-        found_path = resolve_model_in_path(model_id)
+        # Check all model directories for pre-existing complete models
+        found_path = await to_thread.run_sync(resolve_existing_model, model_id)
         if found_path is not None:
-            logger.info(
-                f"DownloadCoordinator: Model {model_id} found in EXO_MODELS_PATH at {found_path}"
-            )
-            completed = DownloadCompleted(
-                shard_metadata=shard,
-                node_id=self.node_id,
-                total=shard.model_card.storage_size,
-                model_directory=str(found_path),
-                read_only=True,
-            )
-            self.download_status[model_id] = completed
-            await self.event_sender.send(
-                NodeDownloadProgress(download_progress=completed)
-            )
-            return
-
-        local_model_dir = EXO_MODELS_DIR / model_id.normalize()
-        if local_model_dir.is_dir() and is_model_directory_complete(local_model_dir):
-            logger.info(
-                f"DownloadCoordinator: Model {model_id} already complete at {local_model_dir}"
-            )
-            completed = DownloadCompleted(
-                shard_metadata=shard,
-                node_id=self.node_id,
-                total=shard.model_card.storage_size,
-                model_directory=str(local_model_dir),
+            logger.info(f"DownloadCoordinator: Model {model_id} found at {found_path}")
+            completed = self._completed_from_path(
+                shard, found_path, shard.model_card.storage_size
             )
             self.download_status[model_id] = completed
             await self.event_sender.send(
@@ -190,7 +197,7 @@ class DownloadCoordinator:
         progress = DownloadPending(
             shard_metadata=shard,
             node_id=self.node_id,
-            model_directory=self._model_dir(model_id),
+            model_directory=self._default_model_dir(model_id),
         )
         self.download_status[model_id] = progress
         await self.event_sender.send(NodeDownloadProgress(download_progress=progress))
@@ -201,12 +208,18 @@ class DownloadCoordinator:
         )
 
         if initial_progress.status == "complete":
-            completed = DownloadCompleted(
-                shard_metadata=shard,
-                node_id=self.node_id,
-                total=initial_progress.total,
-                model_directory=self._model_dir(model_id),
-            )
+            found = await to_thread.run_sync(resolve_existing_model, model_id)
+            if found is not None:
+                completed = self._completed_from_path(
+                    shard, found, initial_progress.total
+                )
+            else:
+                completed = DownloadCompleted(
+                    shard_metadata=shard,
+                    node_id=self.node_id,
+                    total=initial_progress.total,
+                    model_directory=self._default_model_dir(model_id),
+                )
             self.download_status[model_id] = completed
             await self.event_sender.send(
                 NodeDownloadProgress(download_progress=completed)
@@ -221,7 +234,7 @@ class DownloadCoordinator:
                 shard_metadata=shard,
                 node_id=self.node_id,
                 error_message=f"Model files not found locally in offline mode: {model_id}",
-                model_directory=self._model_dir(model_id),
+                model_directory=self._default_model_dir(model_id),
             )
             self.download_status[model_id] = failed
             await self.event_sender.send(NodeDownloadProgress(download_progress=failed))
@@ -242,7 +255,7 @@ class DownloadCoordinator:
             download_progress=map_repo_download_progress_to_download_progress_data(
                 initial_progress
             ),
-            model_directory=self._model_dir(model_id),
+            model_directory=self._default_model_dir(model_id),
         )
         self.download_status[model_id] = status
         self.event_sender.send_nowait(NodeDownloadProgress(download_progress=status))
@@ -257,7 +270,7 @@ class DownloadCoordinator:
                     shard_metadata=shard,
                     node_id=self.node_id,
                     error_message=str(e),
-                    model_directory=self._model_dir(model_id),
+                    model_directory=self._default_model_dir(model_id),
                 )
                 self.download_status[model_id] = failed
                 await self.event_sender.send(
@@ -274,13 +287,11 @@ class DownloadCoordinator:
         self.active_downloads[model_id] = scope
 
     async def _delete_download(self, model_id: ModelId) -> None:
-        # Protect read-only models (from EXO_MODELS_PATH) from deletion
+        # Protect read-only models from deletion
         if model_id in self.download_status:
             current = self.download_status[model_id]
             if isinstance(current, DownloadCompleted) and current.read_only:
-                logger.warning(
-                    f"Refusing to delete read-only model {model_id} (from EXO_MODELS_PATH)"
-                )
+                logger.warning(f"Refusing to delete read-only model {model_id}")
                 return
 
         # Cancel if active
@@ -303,7 +314,7 @@ class DownloadCoordinator:
             pending = DownloadPending(
                 shard_metadata=current_status.shard_metadata,
                 node_id=self.node_id,
-                model_directory=self._model_dir(model_id),
+                model_directory=self._default_model_dir(model_id),
             )
             await self.event_sender.send(
                 NodeDownloadProgress(download_progress=pending)
@@ -327,22 +338,26 @@ class DownloadCoordinator:
                         continue
 
                     if progress.status == "complete":
-                        status: DownloadProgress = DownloadCompleted(
-                            node_id=self.node_id,
-                            shard_metadata=progress.shard,
-                            total=progress.total,
-                            model_directory=self._model_dir(
-                                progress.shard.model_card.model_id
-                            ),
+                        found = await to_thread.run_sync(
+                            resolve_existing_model, model_id
                         )
+                        if found is not None:
+                            status: DownloadProgress = self._completed_from_path(
+                                progress.shard, found, progress.total
+                            )
+                        else:
+                            status = DownloadCompleted(
+                                node_id=self.node_id,
+                                shard_metadata=progress.shard,
+                                total=progress.total,
+                                model_directory=self._default_model_dir(model_id),
+                            )
                     elif progress.status in ["in_progress", "not_started"]:
                         if progress.downloaded_this_session.in_bytes == 0:
                             status = DownloadPending(
                                 node_id=self.node_id,
                                 shard_metadata=progress.shard,
-                                model_directory=self._model_dir(
-                                    progress.shard.model_card.model_id
-                                ),
+                                model_directory=self._default_model_dir(model_id),
                                 downloaded=progress.downloaded,
                                 total=progress.total,
                             )
@@ -353,9 +368,7 @@ class DownloadCoordinator:
                                 download_progress=map_repo_download_progress_to_download_progress_data(
                                     progress
                                 ),
-                                model_directory=self._model_dir(
-                                    progress.shard.model_card.model_id
-                                ),
+                                model_directory=self._default_model_dir(model_id),
                             )
                     else:
                         continue
@@ -364,8 +377,8 @@ class DownloadCoordinator:
                     await self.event_sender.send(
                         NodeDownloadProgress(download_progress=status)
                     )
-                # Scan EXO_MODELS_PATH for pre-downloaded models
-                if EXO_MODELS_PATH is not None:
+                # Scan read-only directories for pre-downloaded models
+                if EXO_MODELS_READ_ONLY_DIRS:
                     for card in await get_model_cards():
                         mid = card.model_id
                         if mid in self.active_downloads:
@@ -375,8 +388,8 @@ class DownloadCoordinator:
                             (DownloadCompleted, DownloadOngoing, DownloadFailed),
                         ):
                             continue
-                        found = resolve_model_in_path(mid)
-                        if found is not None:
+                        found = await to_thread.run_sync(resolve_existing_model, mid)
+                        if found is not None and is_read_only_model_dir(found):
                             path_shard = PipelineShardMetadata(
                                 model_card=card,
                                 device_rank=0,
@@ -385,12 +398,10 @@ class DownloadCoordinator:
                                 end_layer=card.n_layers,
                                 n_layers=card.n_layers,
                             )
-                            path_completed: DownloadProgress = DownloadCompleted(
-                                node_id=self.node_id,
-                                shard_metadata=path_shard,
-                                total=card.storage_size,
-                                model_directory=str(found),
-                                read_only=True,
+                            path_completed: DownloadProgress = (
+                                self._completed_from_path(
+                                    path_shard, found, card.storage_size
+                                )
                             )
                             self.download_status[mid] = path_completed
                             await self.event_sender.send(

src/exo/download/download_utils.py

@@ -30,7 +30,11 @@ from exo.download.huggingface_utils import (
     get_hf_endpoint,
     get_hf_token,
 )
-from exo.shared.constants import EXO_MODELS_DIR, EXO_MODELS_PATH
+from exo.shared.constants import (
+    EXO_DEFAULT_MODELS_DIR,
+    EXO_MODELS_DIRS,
+    EXO_MODELS_READ_ONLY_DIRS,
+)
 from exo.shared.models.model_cards import ModelTask
 from exo.shared.types.common import ModelId
 from exo.shared.types.memory import Memory
@@ -110,50 +114,87 @@ def map_repo_download_progress_to_download_progress_data(
     )
 
 
-def resolve_model_in_path(model_id: ModelId) -> Path | None:
-    """Search EXO_MODELS_PATH directories for a pre-existing model.
+class InsufficientDiskSpaceError(Exception):
+    """Raised when no writable model directory has enough free space."""
 
-    Checks each directory for the normalized name (org--model).  A candidate
-    is only returned if ``is_model_directory_complete`` confirms all weight
-    files are present.
+
+def resolve_existing_model(model_id: ModelId) -> Path | None:
+    """Search all model directories for a complete, pre-existing model.
+
+    Checks read-only directories first, then writable directories.
+    A candidate is only returned if ``is_model_directory_complete`` confirms
+    all weight files are present.
     """
-    if EXO_MODELS_PATH is None:
-        return None
     normalized = model_id.normalize()
-    for search_dir in EXO_MODELS_PATH:
+    for search_dir in (*EXO_MODELS_READ_ONLY_DIRS, *EXO_MODELS_DIRS):
         candidate = search_dir / normalized
         if candidate.is_dir() and is_model_directory_complete(candidate):
             return candidate
     return None
 
 
+def is_read_only_model_dir(model_dir: Path) -> bool:
+    """Check if a model directory lives under a read-only models root."""
+    return any(model_dir.is_relative_to(d) for d in EXO_MODELS_READ_ONLY_DIRS)
+
+
 def build_model_path(model_id: ModelId) -> Path:
-    found = resolve_model_in_path(model_id)
+    found = resolve_existing_model(model_id)
     if found is not None:
         return found
-    return EXO_MODELS_DIR / model_id.normalize()
+    return EXO_DEFAULT_MODELS_DIR / model_id.normalize()
 
 
-async def resolve_model_path_for_repo(model_id: ModelId) -> Path:
-    return (await ensure_models_dir()) / model_id.normalize()
+def select_download_dir(required_bytes: int) -> Path:
+    """Pick the first writable model directory with enough free space.
+
+    Raises ``InsufficientDiskSpaceError`` if none have enough space.
+    """
+    for candidate_dir in EXO_MODELS_DIRS:
+        if not candidate_dir.exists():
+            continue
+        try:
+            usage = shutil.disk_usage(candidate_dir)
+            if usage.free >= required_bytes:
+                return candidate_dir
+        except OSError:
+            continue
+    raise InsufficientDiskSpaceError(
+        f"No writable model directory has {required_bytes / (1024**3):.1f} GiB free. "
+        f"Checked: {[str(d) for d in EXO_MODELS_DIRS]}"
+    )
 
 
-async def ensure_models_dir() -> Path:
-    await aios.makedirs(EXO_MODELS_DIR, exist_ok=True)
-    return EXO_MODELS_DIR
+async def resolve_model_dir(model_id: ModelId) -> Path:
+    """Return the directory for a model's files, creating it if needed.
+
+    Checks all model directories for an existing complete model first,
+    then falls back to the default models directory.
+    """
+    target = await asyncio.to_thread(build_model_path, model_id)
+    await aios.makedirs(target, exist_ok=True)
+    return target
+
+
+async def ensure_cache_dir(model_id: ModelId) -> Path:
+    """Return the cache directory for a model's metadata, creating it if needed."""
+    target = EXO_DEFAULT_MODELS_DIR / "caches" / model_id.normalize()
+    await aios.makedirs(target, exist_ok=True)
+    return target
 
 
 async def delete_model(model_id: ModelId) -> bool:
-    models_dir = await ensure_models_dir()
-    model_dir = models_dir / model_id.normalize()
-    cache_dir = models_dir / "caches" / model_id.normalize()
-
+    """Delete a model from writable directories. Skips read-only dirs."""
+    normalized = model_id.normalize()
     deleted = False
-    if await aios.path.exists(model_dir):
-        await asyncio.to_thread(shutil.rmtree, model_dir, ignore_errors=False)
-        deleted = True
+    for models_dir in EXO_MODELS_DIRS:
+        model_dir = models_dir / normalized
+        if await aios.path.exists(model_dir):
+            await asyncio.to_thread(shutil.rmtree, model_dir, ignore_errors=False)
+            deleted = True
 
-    # Also clear cache
+    # Clear cache from default dir
+    cache_dir = EXO_DEFAULT_MODELS_DIR / "caches" / normalized
     if await aios.path.exists(cache_dir):
         await asyncio.to_thread(shutil.rmtree, cache_dir, ignore_errors=False)
 
@@ -161,9 +202,10 @@ async def delete_model(model_id: ModelId) -> bool:
 
 
 async def seed_models(seed_dir: str | Path):
-    """Move models from resources folder to EXO_MODELS_DIR."""
+    """Move models from resources folder to the default models directory."""
     source_dir = Path(seed_dir)
-    dest_dir = await ensure_models_dir()
+    await aios.makedirs(EXO_DEFAULT_MODELS_DIR, exist_ok=True)
+    dest_dir = EXO_DEFAULT_MODELS_DIR
     for path in source_dir.iterdir():
         if path.is_dir() and path.name.startswith("models--"):
             dest_path = dest_dir / path.name
@@ -239,17 +281,6 @@ def _scan_model_directory(
 
 def is_model_directory_complete(model_dir: Path) -> bool:
     """Check if a model directory contains all required weight files."""
-    index_files = list(model_dir.glob("**/*.safetensors.index.json"))
-    if not index_files:
-        return False
-
-    for index_file in index_files:
-        try:
-            ModelSafetensorsIndex.model_validate_json(index_file.read_text())
-        except Exception:
-            logger.warning(f"Failed to parse model index {index_file}")
-            return False
-
     file_list = _scan_model_directory(model_dir, recursive=True)
     return file_list is not None and all(f.size is not None for f in file_list)
 
@@ -264,14 +295,16 @@ async def _build_file_list_from_local_directory(
     a local directory must contain a *.safetensors.index.json and
     safetensors listed there.
     """
-    model_dir = (await ensure_models_dir()) / model_id.normalize()
-    if not await aios.path.exists(model_dir):
-        return None
-
-    file_list = await asyncio.to_thread(_scan_model_directory, model_dir, recursive)
-    if not file_list:
-        return None
-    return file_list
+    normalized = model_id.normalize()
+    for search_dir in (*EXO_MODELS_READ_ONLY_DIRS, *EXO_MODELS_DIRS):
+        model_dir = search_dir / normalized
+        if await aios.path.exists(model_dir):
+            file_list = await asyncio.to_thread(
+                _scan_model_directory, model_dir, recursive
+            )
+            if file_list:
+                return file_list
+    return None
 
 
 _fetched_file_lists_this_session: set[str] = set()
@@ -284,8 +317,7 @@ async def fetch_file_list_with_cache(
     skip_internet: bool = False,
     on_connection_lost: Callable[[], None] = lambda: None,
 ) -> list[FileListEntry]:
-    target_dir = (await ensure_models_dir()) / "caches" / model_id.normalize()
-    await aios.makedirs(target_dir, exist_ok=True)
+    target_dir = await ensure_cache_dir(model_id)
     cache_file = target_dir / f"{model_id.normalize()}--{revision}--file_list.json"
     cache_key = f"{model_id.normalize()}--{revision}"
 
@@ -340,7 +372,7 @@ async def fetch_file_list_with_cache(
         )
         if local_file_list is not None:
             logger.warning(
-                f"Failed to fetch file list for {model_id} and no cache exists, "
+                f"Failed to fetch file list for {model_id} and no cache exists, using local file list"
             )
             return local_file_list
         raise FileNotFoundError(f"Failed to fetch file list for {model_id}: {e}") from e
@@ -669,8 +701,7 @@ def calculate_repo_progress(
 
 
 async def get_weight_map(model_id: ModelId, revision: str = "main") -> dict[str, str]:
-    target_dir = (await ensure_models_dir()) / model_id.normalize()
-    await aios.makedirs(target_dir, exist_ok=True)
+    target_dir = await resolve_model_dir(model_id)
 
     index_files_dir = snapshot_download(
         repo_id=model_id,
@@ -741,30 +772,46 @@ async def download_shard(
     if not skip_download:
         logger.debug(f"Downloading {shard.model_card.model_id=}")
 
+    model_id = shard.model_card.model_id
     revision = "main"
-    target_dir = await ensure_models_dir() / str(shard.model_card.model_id).replace(
-        "/", "--"
-    )
-    if not skip_download:
-        await aios.makedirs(target_dir, exist_ok=True)
 
     if not allow_patterns:
         allow_patterns = await resolve_allow_patterns(shard)
 
     if not skip_download:
-        logger.debug(f"Downloading {shard.model_card.model_id=} with {allow_patterns=}")
+        logger.debug(f"Downloading {model_id=} with {allow_patterns=}")
 
     all_start_time = time.time()
-    file_list = await fetch_file_list_with_cache(
-        shard.model_card.model_id,
-        revision,
-        recursive=True,
-        skip_internet=skip_internet,
-        on_connection_lost=on_connection_lost,
-    )
+    try:
+        file_list = await fetch_file_list_with_cache(
+            model_id,
+            revision,
+            recursive=True,
+            skip_internet=skip_internet,
+            on_connection_lost=on_connection_lost,
+        )
+    except FileNotFoundError:
+        not_started_progress = RepoDownloadProgress(
+            repo_id=str(model_id),
+            repo_revision=revision,
+            shard=shard,
+            completed_files=0,
+            total_files=0,
+            downloaded=Memory.from_bytes(0),
+            downloaded_this_session=Memory.from_bytes(0),
+            total=Memory.from_bytes(0),
+            overall_speed=0.0,
+            overall_eta=timedelta(0),
+            status="not_started",
+            file_progress={},
+        )
+        return EXO_DEFAULT_MODELS_DIR / model_id.normalize(), not_started_progress
     filtered_file_list = list(
         filter_repo_objects(
-            file_list, allow_patterns=allow_patterns, key=lambda x: x.path
+            file_list,
+            allow_patterns=allow_patterns,
+            ignore_patterns=["original/*", "metal/*"],
+            key=lambda x: x.path,
         )
     )
 
@@ -776,6 +823,15 @@ async def download_shard(
             for f in filtered_file_list
             if "/" in f.path or not f.path.endswith(".safetensors")
         ]
+
+    # Pick a writable directory with enough free space
+    total_size = sum(f.size or 0 for f in filtered_file_list)
+    models_dir = (
+        select_download_dir(total_size) if not skip_download else EXO_DEFAULT_MODELS_DIR
+    )
+    target_dir = models_dir / model_id.normalize()
+    if not skip_download:
+        await aios.makedirs(target_dir, exist_ok=True)
     file_progress: dict[str, RepoFileDownloadProgress] = {}
 
     async def on_progress_wrapper(
@@ -812,7 +868,7 @@ async def download_shard(
             else timedelta(seconds=0)
         )
         file_progress[file.path] = RepoFileDownloadProgress(
-            repo_id=shard.model_card.model_id,
+            repo_id=model_id,
             repo_revision=revision,
             file_path=file.path,
             downloaded=Memory.from_bytes(curr_bytes),
@@ -840,7 +896,7 @@ async def download_shard(
         downloaded_bytes = await get_downloaded_size(target_dir / file.path)
         final_file_exists = await aios.path.exists(target_dir / file.path)
         file_progress[file.path] = RepoFileDownloadProgress(
-            repo_id=shard.model_card.model_id,
+            repo_id=model_id,
             repo_revision=revision,
             file_path=file.path,
             downloaded=Memory.from_bytes(downloaded_bytes),
@@ -866,7 +922,7 @@ async def download_shard(
     async def download_with_semaphore(file: FileListEntry) -> None:
         async with semaphore:
             await download_file_with_retry(
-                shard.model_card.model_id,
+                model_id,
                 revision,
                 file.path,
                 target_dir,
@@ -882,7 +938,7 @@ async def download_shard(
             *[download_with_semaphore(file) for file in filtered_file_list]
         )
     final_repo_progress = calculate_repo_progress(
-        shard, shard.model_card.model_id, revision, file_progress, all_start_time
+        shard, model_id, revision, file_progress, all_start_time
     )
     await on_progress(shard, final_repo_progress)
     if gguf := next((f for f in filtered_file_list if f.path.endswith(".gguf")), None):

src/exo/download/tests/test_download_verification.py

@@ -1,7 +1,6 @@
 """Tests for download verification and cache behavior."""
 
 import time
-from collections.abc import AsyncIterator
 from datetime import timedelta
 from pathlib import Path
 from unittest.mock import AsyncMock, MagicMock, patch
@@ -14,7 +13,6 @@ from pydantic import TypeAdapter
 from exo.download.download_utils import (
     delete_model,
     fetch_file_list_with_cache,
-    is_model_directory_complete,
 )
 from exo.shared.types.common import ModelId
 from exo.shared.types.memory import Memory
@@ -26,15 +24,6 @@ def model_id() -> ModelId:
     return ModelId("test-org/test-model")
 
 
-@pytest.fixture
-async def temp_models_dir(tmp_path: Path) -> AsyncIterator[Path]:
-    """Set up a temporary models directory for testing."""
-    models_dir = tmp_path / "models"
-    await aios.makedirs(models_dir, exist_ok=True)
-    with patch("exo.download.download_utils.EXO_MODELS_DIR", models_dir):
-        yield models_dir
-
-
 class TestFileVerification:
     """Tests for file size verification in _download_file."""
 
@@ -189,7 +178,8 @@ class TestFileListCache:
         ]
 
         with (
-            patch("exo.download.download_utils.EXO_MODELS_DIR", models_dir),
+            patch("exo.download.download_utils.EXO_MODELS_DIRS", (models_dir,)),
+            patch("exo.download.download_utils.EXO_DEFAULT_MODELS_DIR", models_dir),
             patch(
                 "exo.download.download_utils.fetch_file_list_with_retry",
                 new_callable=AsyncMock,
@@ -235,7 +225,8 @@ class TestFileListCache:
             )
 
         with (
-            patch("exo.download.download_utils.EXO_MODELS_DIR", models_dir),
+            patch("exo.download.download_utils.EXO_MODELS_DIRS", (models_dir,)),
+            patch("exo.download.download_utils.EXO_DEFAULT_MODELS_DIR", models_dir),
             patch(
                 "exo.download.download_utils.fetch_file_list_with_retry",
                 new_callable=AsyncMock,
@@ -253,7 +244,8 @@ class TestFileListCache:
         models_dir = tmp_path / "models"
 
         with (
-            patch("exo.download.download_utils.EXO_MODELS_DIR", models_dir),
+            patch("exo.download.download_utils.EXO_MODELS_DIRS", (models_dir,)),
+            patch("exo.download.download_utils.EXO_DEFAULT_MODELS_DIR", models_dir),
             patch(
                 "exo.download.download_utils.fetch_file_list_with_retry",
                 new_callable=AsyncMock,
@@ -285,7 +277,10 @@ class TestModelDeletion:
         async with aiofiles.open(cache_dir / "file_list.json", "w") as f:
             await f.write("[]")
 
-        with patch("exo.download.download_utils.EXO_MODELS_DIR", models_dir):
+        with (
+            patch("exo.download.download_utils.EXO_MODELS_DIRS", (models_dir,)),
+            patch("exo.download.download_utils.EXO_DEFAULT_MODELS_DIR", models_dir),
+        ):
             result = await delete_model(model_id)
 
             assert result is True
@@ -304,7 +299,10 @@ class TestModelDeletion:
         async with aiofiles.open(cache_dir / "file_list.json", "w") as f:
             await f.write("[]")
 
-        with patch("exo.download.download_utils.EXO_MODELS_DIR", models_dir):
+        with (
+            patch("exo.download.download_utils.EXO_MODELS_DIRS", (models_dir,)),
+            patch("exo.download.download_utils.EXO_DEFAULT_MODELS_DIR", models_dir),
+        ):
             result = await delete_model(model_id)
 
             # Returns False because model dir didn't exist
@@ -319,75 +317,15 @@ class TestModelDeletion:
         models_dir = tmp_path / "models"
         await aios.makedirs(models_dir, exist_ok=True)
 
-        with patch("exo.download.download_utils.EXO_MODELS_DIR", models_dir):
+        with (
+            patch("exo.download.download_utils.EXO_MODELS_DIRS", (models_dir,)),
+            patch("exo.download.download_utils.EXO_DEFAULT_MODELS_DIR", models_dir),
+        ):
             result = await delete_model(model_id)
 
             assert result is False
 
 
-class TestModelDirectoryComplete:
-    """Tests for local completeness checks used to skip download probing."""
-
-    async def test_returns_false_when_only_partial_weight_exists(
-        self, model_id: ModelId, tmp_path: Path
-    ) -> None:
-        import json
-
-        model_dir = tmp_path / model_id.normalize()
-        await aios.makedirs(model_dir, exist_ok=True)
-
-        async with aiofiles.open(model_dir / "model.safetensors.index.json", "w") as f:
-            await f.write(
-                json.dumps(
-                    {
-                        "metadata": {"total_size": 256},
-                        "weight_map": {"model.layers.0.weight": "model.safetensors"},
-                    }
-                )
-            )
-        async with aiofiles.open(model_dir / "model.safetensors.partial", "wb") as f:
-            await f.write(b"x" * 128)
-
-        assert is_model_directory_complete(model_dir) is False
-
-    async def test_returns_true_when_final_weight_exists_even_with_stale_partial(
-        self, model_id: ModelId, tmp_path: Path
-    ) -> None:
-        import json
-
-        model_dir = tmp_path / model_id.normalize()
-        await aios.makedirs(model_dir, exist_ok=True)
-
-        async with aiofiles.open(model_dir / "model.safetensors.index.json", "w") as f:
-            await f.write(
-                json.dumps(
-                    {
-                        "metadata": {"total_size": 256},
-                        "weight_map": {"model.layers.0.weight": "model.safetensors"},
-                    }
-                )
-            )
-        async with aiofiles.open(model_dir / "model.safetensors", "wb") as f:
-            await f.write(b"x" * 256)
-        async with aiofiles.open(model_dir / "model.safetensors.partial", "wb") as f:
-            await f.write(b"x" * 128)
-
-        assert is_model_directory_complete(model_dir) is True
-
-    async def test_returns_false_when_index_is_invalid(
-        self, model_id: ModelId, tmp_path: Path
-    ) -> None:
-        model_dir = tmp_path / model_id.normalize()
-        await aios.makedirs(model_dir, exist_ok=True)
-
-        async with aiofiles.open(model_dir / "model.safetensors.index.json", "w") as f:
-            await f.write("{not valid json")
-        async with aiofiles.open(model_dir / "model.safetensors", "wb") as f:
-            await f.write(b"x" * 256)
-
-        assert is_model_directory_complete(model_dir) is False
-
-
 class TestProgressResetOnRedownload:
     """Tests for progress tracking when files are re-downloaded."""
 

src/exo/download/tests/test_model_dirs.py

@@ -0,0 +1,297 @@
+"""Tests for multi-directory model resolution, download target selection, and deletion."""
+
+import json
+import shutil
+from collections.abc import AsyncIterator
+from pathlib import Path
+from unittest.mock import patch
+
+import aiofiles
+import aiofiles.os as aios
+import pytest
+
+from exo.download.download_utils import (
+    InsufficientDiskSpaceError,
+    delete_model,
+    is_read_only_model_dir,
+    resolve_existing_model,
+    select_download_dir,
+)
+from exo.shared.types.common import ModelId
+
+MODEL_ID = ModelId("test-org/test-model")
+NORMALIZED = MODEL_ID.normalize()
+
+
+def _create_complete_model(model_dir: Path) -> None:
+    """Create a minimal complete model directory on disk."""
+    model_dir.mkdir(parents=True, exist_ok=True)
+    weight_map = {"layer.weight": "model.safetensors"}
+    index = {"metadata": {"total_size": 1024}, "weight_map": weight_map}
+    (model_dir / "model.safetensors.index.json").write_text(json.dumps(index))
+    (model_dir / "model.safetensors").write_bytes(b"weights")
+    (model_dir / "config.json").write_text('{"model_type": "test"}')
+
+
+def _create_incomplete_model(model_dir: Path) -> None:
+    """Create a model directory missing weight files."""
+    model_dir.mkdir(parents=True, exist_ok=True)
+    weight_map = {"layer.weight": "model.safetensors"}
+    index = {"metadata": {"total_size": 1024}, "weight_map": weight_map}
+    (model_dir / "model.safetensors.index.json").write_text(json.dumps(index))
+    # model.safetensors is missing
+
+
+# ---------------------------------------------------------------------------
+# resolve_existing_model
+# ---------------------------------------------------------------------------
+
+
+class TestResolveExistingModel:
+    def test_returns_none_when_no_dirs_have_model(self, tmp_path: Path) -> None:
+        writable = tmp_path / "writable"
+        writable.mkdir()
+        with (
+            patch("exo.download.download_utils.EXO_MODELS_READ_ONLY_DIRS", ()),
+            patch("exo.download.download_utils.EXO_MODELS_DIRS", (writable,)),
+        ):
+            assert resolve_existing_model(MODEL_ID) is None
+
+    def test_finds_model_in_writable_dir(self, tmp_path: Path) -> None:
+        writable = tmp_path / "writable"
+        _create_complete_model(writable / NORMALIZED)
+        with (
+            patch("exo.download.download_utils.EXO_MODELS_READ_ONLY_DIRS", ()),
+            patch("exo.download.download_utils.EXO_MODELS_DIRS", (writable,)),
+        ):
+            assert resolve_existing_model(MODEL_ID) == writable / NORMALIZED
+
+    def test_finds_model_in_read_only_dir(self, tmp_path: Path) -> None:
+        read_only = tmp_path / "readonly"
+        _create_complete_model(read_only / NORMALIZED)
+        writable = tmp_path / "writable"
+        writable.mkdir()
+        with (
+            patch(
+                "exo.download.download_utils.EXO_MODELS_READ_ONLY_DIRS", (read_only,)
+            ),
+            patch("exo.download.download_utils.EXO_MODELS_DIRS", (writable,)),
+        ):
+            assert resolve_existing_model(MODEL_ID) == read_only / NORMALIZED
+
+    def test_read_only_takes_priority_over_writable(self, tmp_path: Path) -> None:
+        read_only = tmp_path / "readonly"
+        _create_complete_model(read_only / NORMALIZED)
+        writable = tmp_path / "writable"
+        _create_complete_model(writable / NORMALIZED)
+        with (
+            patch(
+                "exo.download.download_utils.EXO_MODELS_READ_ONLY_DIRS", (read_only,)
+            ),
+            patch("exo.download.download_utils.EXO_MODELS_DIRS", (writable,)),
+        ):
+            result = resolve_existing_model(MODEL_ID)
+            assert result == read_only / NORMALIZED
+
+    def test_skips_incomplete_model(self, tmp_path: Path) -> None:
+        incomplete = tmp_path / "incomplete"
+        _create_incomplete_model(incomplete / NORMALIZED)
+        complete = tmp_path / "complete"
+        _create_complete_model(complete / NORMALIZED)
+        with (
+            patch(
+                "exo.download.download_utils.EXO_MODELS_READ_ONLY_DIRS", (incomplete,)
+            ),
+            patch("exo.download.download_utils.EXO_MODELS_DIRS", (complete,)),
+        ):
+            result = resolve_existing_model(MODEL_ID)
+            assert result == complete / NORMALIZED
+
+    def test_searches_multiple_read_only_dirs_in_order(self, tmp_path: Path) -> None:
+        ro1 = tmp_path / "ro1"
+        ro1.mkdir()
+        ro2 = tmp_path / "ro2"
+        _create_complete_model(ro2 / NORMALIZED)
+        writable = tmp_path / "writable"
+        writable.mkdir()
+        with (
+            patch("exo.download.download_utils.EXO_MODELS_READ_ONLY_DIRS", (ro1, ro2)),
+            patch("exo.download.download_utils.EXO_MODELS_DIRS", (writable,)),
+        ):
+            assert resolve_existing_model(MODEL_ID) == ro2 / NORMALIZED
+
+
+# ---------------------------------------------------------------------------
+# is_read_only_model_dir
+# ---------------------------------------------------------------------------
+
+
+class TestIsReadOnlyModelDir:
+    def test_path_under_read_only_dir(self, tmp_path: Path) -> None:
+        ro = tmp_path / "readonly"
+        with patch("exo.download.download_utils.EXO_MODELS_READ_ONLY_DIRS", (ro,)):
+            assert is_read_only_model_dir(ro / NORMALIZED) is True
+
+    def test_path_under_writable_dir(self, tmp_path: Path) -> None:
+        writable = tmp_path / "writable"
+        with patch("exo.download.download_utils.EXO_MODELS_READ_ONLY_DIRS", ()):
+            assert is_read_only_model_dir(writable / NORMALIZED) is False
+
+    def test_path_not_under_any_read_only_dir(self, tmp_path: Path) -> None:
+        ro = tmp_path / "readonly"
+        other = tmp_path / "other"
+        with patch("exo.download.download_utils.EXO_MODELS_READ_ONLY_DIRS", (ro,)):
+            assert is_read_only_model_dir(other / NORMALIZED) is False
+
+
+# ---------------------------------------------------------------------------
+# select_download_dir
+# ---------------------------------------------------------------------------
+
+
+class TestSelectDownloadDir:
+    def test_picks_first_dir_with_enough_space(self, tmp_path: Path) -> None:
+        dir1 = tmp_path / "dir1"
+        dir2 = tmp_path / "dir2"
+        dir1.mkdir()
+        dir2.mkdir()
+        # Both exist on same filesystem so both have space; first wins
+        with patch("exo.download.download_utils.EXO_MODELS_DIRS", (dir1, dir2)):
+            assert select_download_dir(1) == dir1
+
+    def test_skips_dir_without_enough_space(self, tmp_path: Path) -> None:
+        dir1 = tmp_path / "dir1"
+        dir2 = tmp_path / "dir2"
+        dir1.mkdir()
+        dir2.mkdir()
+
+        real_disk_usage = shutil.disk_usage
+
+        def mock_disk_usage(path: str | Path) -> object:
+            if Path(path).is_relative_to(dir1):
+                real = real_disk_usage(path)
+                return shutil._ntuple_diskusage(real.total, real.total, 0)  # pyright: ignore[reportPrivateUsage]
+            return real_disk_usage(path)
+
+        with (
+            patch("exo.download.download_utils.EXO_MODELS_DIRS", (dir1, dir2)),
+            patch("shutil.disk_usage", side_effect=mock_disk_usage),
+        ):
+            assert select_download_dir(1024) == dir2
+
+    def test_raises_when_no_dir_has_space(self, tmp_path: Path) -> None:
+        dir1 = tmp_path / "dir1"
+        dir1.mkdir()
+
+        real_disk_usage = shutil.disk_usage
+
+        def mock_disk_usage(path: str | Path) -> object:
+            real = real_disk_usage(path)
+            return shutil._ntuple_diskusage(real.total, real.total, 0)  # pyright: ignore[reportPrivateUsage]
+
+        with (
+            patch("exo.download.download_utils.EXO_MODELS_DIRS", (dir1,)),
+            patch("shutil.disk_usage", side_effect=mock_disk_usage),
+            pytest.raises(InsufficientDiskSpaceError),
+        ):
+            select_download_dir(1024)
+
+    def test_skips_nonexistent_dir(self, tmp_path: Path) -> None:
+        nonexistent = tmp_path / "does-not-exist"
+        with (
+            patch("exo.download.download_utils.EXO_MODELS_DIRS", (nonexistent,)),
+            pytest.raises(InsufficientDiskSpaceError),
+        ):
+            select_download_dir(1)
+
+    def test_skips_dir_raising_oserror(self, tmp_path: Path) -> None:
+        dir1 = tmp_path / "unmounted"
+        dir2 = tmp_path / "ok"
+        dir1.mkdir()
+        dir2.mkdir()
+
+        real_disk_usage = shutil.disk_usage
+
+        def mock_disk_usage(path: str | Path) -> object:
+            if Path(path).is_relative_to(dir1):
+                raise OSError("device not mounted")
+            return real_disk_usage(path)
+
+        with (
+            patch("exo.download.download_utils.EXO_MODELS_DIRS", (dir1, dir2)),
+            patch("shutil.disk_usage", side_effect=mock_disk_usage),
+        ):
+            assert select_download_dir(1) == dir2
+
+
+# ---------------------------------------------------------------------------
+# delete_model
+# ---------------------------------------------------------------------------
+
+
+class TestDeleteModel:
+    @pytest.fixture
+    async def dirs(self, tmp_path: Path) -> AsyncIterator[tuple[Path, Path, Path]]:
+        writable1 = tmp_path / "w1"
+        writable2 = tmp_path / "w2"
+        default = tmp_path / "default"
+        await aios.makedirs(writable1, exist_ok=True)
+        await aios.makedirs(writable2, exist_ok=True)
+        await aios.makedirs(default, exist_ok=True)
+        with (
+            patch(
+                "exo.download.download_utils.EXO_MODELS_DIRS",
+                (writable1, writable2, default),
+            ),
+            patch("exo.download.download_utils.EXO_DEFAULT_MODELS_DIR", default),
+        ):
+            yield writable1, writable2, default
+
+    async def test_deletes_from_writable_dir(
+        self, dirs: tuple[Path, Path, Path]
+    ) -> None:
+        w1, _, _ = dirs
+        model_dir = w1 / NORMALIZED
+        await aios.makedirs(model_dir, exist_ok=True)
+        async with aiofiles.open(model_dir / "weights.safetensors", "w") as f:
+            await f.write("data")
+
+        result = await delete_model(MODEL_ID)
+        assert result is True
+        assert not await aios.path.exists(model_dir)
+
+    async def test_deletes_from_multiple_writable_dirs(
+        self, dirs: tuple[Path, Path, Path]
+    ) -> None:
+        w1, w2, _ = dirs
+        model_dir1 = w1 / NORMALIZED
+        model_dir2 = w2 / NORMALIZED
+        await aios.makedirs(model_dir1, exist_ok=True)
+        await aios.makedirs(model_dir2, exist_ok=True)
+        async with aiofiles.open(model_dir1 / "w.safetensors", "w") as f:
+            await f.write("data")
+        async with aiofiles.open(model_dir2 / "w.safetensors", "w") as f:
+            await f.write("data")
+
+        result = await delete_model(MODEL_ID)
+        assert result is True
+        assert not await aios.path.exists(model_dir1)
+        assert not await aios.path.exists(model_dir2)
+
+    async def test_cleans_cache_from_default_dir(
+        self, dirs: tuple[Path, Path, Path]
+    ) -> None:
+        _, _, default = dirs
+        cache_dir = default / "caches" / NORMALIZED
+        await aios.makedirs(cache_dir, exist_ok=True)
+        async with aiofiles.open(cache_dir / "file_list.json", "w") as f:
+            await f.write("[]")
+
+        await delete_model(MODEL_ID)
+        assert not await aios.path.exists(cache_dir)
+
+    async def test_returns_false_when_model_not_found(
+        self, dirs: tuple[Path, Path, Path]
+    ) -> None:
+        result = await delete_model(MODEL_ID)
+        assert result is False

src/exo/download/tests/test_offline_mode.py

@@ -26,7 +26,10 @@ def model_id() -> ModelId:
 async def temp_models_dir(tmp_path: Path) -> AsyncIterator[Path]:
     models_dir = tmp_path / "models"
     await aios.makedirs(models_dir, exist_ok=True)
-    with patch("exo.download.download_utils.EXO_MODELS_DIR", models_dir):
+    with (
+        patch("exo.download.download_utils.EXO_MODELS_DIRS", (models_dir,)),
+        patch("exo.download.download_utils.EXO_DEFAULT_MODELS_DIR", models_dir),
+    ):
         yield models_dir
 
 

src/exo/download/tests/test_re_download.py

@@ -2,16 +2,12 @@
 
 import asyncio
 import contextlib
-import json
 from collections.abc import AsyncIterator, Awaitable
 from datetime import timedelta
 from pathlib import Path
 from typing import Callable
 from unittest.mock import AsyncMock, patch
 
-import aiofiles
-import aiofiles.os as aios
-
 from exo.download.coordinator import DownloadCoordinator
 from exo.download.download_utils import RepoDownloadProgress
 from exo.download.impl_shard_downloader import SingletonShardDownloader
@@ -190,59 +186,12 @@ async def test_re_download_after_delete_completes() -> None:
                 "Re-download after deletion should complete"
             )
         finally:
-            coordinator.shutdown()
+            await coordinator.shutdown()
             coordinator_task.cancel()
             with contextlib.suppress(asyncio.CancelledError):
                 await coordinator_task
 
 
-async def test_start_download_uses_complete_local_model_without_status_probe(
-    tmp_path: Path,
-) -> None:
-    _, cmd_recv = channel[ForwarderDownloadCommand]()
-    event_send, event_recv = channel[Event]()
-
-    fake_downloader = FakeShardDownloader()
-    fake_downloader.get_shard_download_status_for_shard = AsyncMock(
-        side_effect=AssertionError("status probe should be skipped for complete models")
-    )
-
-    coordinator = DownloadCoordinator(
-        node_id=NODE_ID,
-        shard_downloader=SingletonShardDownloader(fake_downloader),
-        download_command_receiver=cmd_recv,
-        event_sender=event_send,
-    )
-    shard = _make_shard()
-
-    models_dir = tmp_path / "models"
-    model_dir = models_dir / MODEL_ID.normalize()
-    await aios.makedirs(model_dir, exist_ok=True)
-
-    async with aiofiles.open(model_dir / "config.json", "w") as f:
-        await f.write('{"model_type":"qwen2"}')
-    async with aiofiles.open(model_dir / "model.safetensors.index.json", "w") as f:
-        await f.write(
-                json.dumps(
-                    {
-                        "metadata": {"total_size": 128},
-                        "weight_map": {"model.layers.0.weight": "model.safetensors"},
-                    }
-                )
-            )
-    async with aiofiles.open(model_dir / "model.safetensors", "wb") as f:
-        await f.write(b"x" * 128)
-
-    with patch("exo.download.coordinator.EXO_MODELS_DIR", models_dir):
-        await coordinator._start_download(shard)  # pyright: ignore[reportPrivateUsage]
-
-    event = await event_recv.receive()
-    assert isinstance(event, NodeDownloadProgress)
-    assert isinstance(event.download_progress, DownloadCompleted)
-    assert event.download_progress.model_directory == str(model_dir)
-    assert event.download_progress.total == shard.model_card.storage_size
-
-
 async def _wait_for_download_completed(
     event_recv: Receiver[Event], model_id: ModelId, timeout: float = 2.0
 ) -> DownloadCompleted | None:

src/exo/main.py

@@ -47,7 +47,11 @@ class Node:
         keypair = get_node_id_keypair()
         node_id = NodeId(keypair.to_node_id())
         session_id = SessionId(master_node_id=node_id, election_clock=0)
-        router = Router.create(keypair)
+        router = Router.create(
+            keypair,
+            bootstrap_peers=args.bootstrap_peers,
+            listen_port=args.libp2p_port,
+        )
         await router.register_topic(topics.GLOBAL_EVENTS)
         await router.register_topic(topics.LOCAL_EVENTS)
         await router.register_topic(topics.COMMANDS)
@@ -224,7 +228,7 @@ class Node:
                     )
                 if result.is_new_master:
                     if self.download_coordinator:
-                        self.download_coordinator.shutdown()
+                        await self.download_coordinator.shutdown()
                         self.download_coordinator = DownloadCoordinator(
                             self.node_id,
                             exo_shard_downloader(offline=self.offline),
@@ -236,7 +240,7 @@ class Node:
                         )
                         self._tg.start_soon(self.download_coordinator.run)
                     if self.worker:
-                        self.worker.shutdown()
+                        await self.worker.shutdown()
                         # TODO: add profiling etc to resource monitor
                         self.worker = Worker(
                             self.node_id,
@@ -264,12 +268,17 @@ def main():
     mp.set_start_method("spawn", force=True)
     # TODO: Refactor the current verbosity system
     logger_setup(EXO_LOG, args.verbosity)
-    logger.info("Starting EXO")
+    logger.info(f"{'=' * 40}")
+    logger.info(f"Starting EXO | pid={os.getpid()}")
+    logger.info(f"{'=' * 40}")
     logger.info(f"EXO_LIBP2P_NAMESPACE: {os.getenv('EXO_LIBP2P_NAMESPACE')}")
 
     if args.offline:
         logger.info("Running in OFFLINE mode — no internet checks, local models only")
 
+    if args.bootstrap_peers:
+        logger.info(f"Bootstrap peers: {args.bootstrap_peers}")
+
     if args.no_batch:
         os.environ["EXO_NO_BATCH"] = "1"
         logger.info("Continuous batching disabled (--no-batch)")
@@ -306,6 +315,8 @@ class Args(CamelCaseModel):
     offline: bool = os.getenv("EXO_OFFLINE", "false").lower() == "true"
     no_batch: bool = False
     fast_synch: bool | None = None  # None = auto, True = force on, False = force off
+    bootstrap_peers: list[str] = []
+    libp2p_port: int
 
     @classmethod
     def parse(cls) -> Self:
@@ -363,6 +374,22 @@ class Args(CamelCaseModel):
             action="store_true",
             help="Disable continuous batching, use sequential generation",
         )
+        parser.add_argument(
+            "--bootstrap-peers",
+            type=lambda s: [p for p in s.split(",") if p],
+            default=os.getenv("EXO_BOOTSTRAP_PEERS", "").split(",")
+            if os.getenv("EXO_BOOTSTRAP_PEERS")
+            else [],
+            dest="bootstrap_peers",
+            help="Comma-separated libp2p multiaddrs to dial on startup (env: EXO_BOOTSTRAP_PEERS)",
+        )
+        parser.add_argument(
+            "--libp2p-port",
+            type=int,
+            default=0,
+            dest="libp2p_port",
+            help="Fixed TCP port for libp2p to listen on (0 = OS-assigned).",
+        )
         fast_synch_group = parser.add_mutually_exclusive_group()
         fast_synch_group.add_argument(
             "--fast-synch",

src/exo/master/main.py

@@ -13,7 +13,9 @@ from exo.master.placement import (
 from exo.shared.apply import apply
 from exo.shared.constants import EXO_EVENT_LOG_DIR, EXO_TRACING_ENABLED
 from exo.shared.types.commands import (
+    AddCustomModelCard,
     CreateInstance,
+    DeleteCustomModelCard,
     DeleteInstance,
     ForwarderCommand,
     ForwarderDownloadCommand,
@@ -29,6 +31,8 @@ from exo.shared.types.commands import (
 )
 from exo.shared.types.common import CommandId, NodeId, SessionId, SystemId
 from exo.shared.types.events import (
+    CustomModelCardAdded,
+    CustomModelCardDeleted,
     Event,
     GlobalForwarderEvent,
     IndexedEvent,
@@ -294,6 +298,7 @@ class Master:
                                 self.state.instances,
                                 self.state.node_memory,
                                 self.state.node_network,
+                                download_status=self.state.downloads,
                             )
                             transition_events = get_transition_events(
                                 self.state.instances, placement, self.state.tasks
@@ -344,6 +349,14 @@ class Master:
                                     f"Finished command {command.finished_command_id} finished"
                                 )
 
+                        case AddCustomModelCard():
+                            generated_events.append(
+                                CustomModelCardAdded(model_card=command.model_card)
+                            )
+                        case DeleteCustomModelCard():
+                            generated_events.append(
+                                CustomModelCardDeleted(model_id=command.model_id)
+                            )
                         case RequestEventLog():
                             # We should just be able to send everything, since other buffers will ignore old messages
                             # rate limit to 1000 at a time

src/exo/master/placement.py

@@ -32,7 +32,10 @@ from exo.shared.types.memory import Memory
 from exo.shared.types.profiling import MemoryUsage, NodeNetworkInfo
 from exo.shared.types.tasks import Task, TaskId, TaskStatus
 from exo.shared.types.worker.downloads import (
+    DownloadCompleted,
+    DownloadFailed,
     DownloadOngoing,
+    DownloadPending,
     DownloadProgress,
 )
 from exo.shared.types.worker.instances import (
@@ -60,6 +63,45 @@ def add_instance_to_placements(
     return {**current_instances, command.instance.instance_id: command.instance}
 
 
+def _get_node_download_fraction(
+    node_id: NodeId,
+    model_id: ModelId,
+    download_status: Mapping[NodeId, Sequence[DownloadProgress]],
+) -> float:
+    """Return the download fraction (0.0–1.0) for a model on a given node."""
+    for progress in download_status.get(node_id, []):
+        if progress.shard_metadata.model_card.model_id != model_id:
+            continue
+        match progress:
+            case DownloadCompleted():
+                return 1.0
+            case DownloadOngoing():
+                total = progress.download_progress.total.in_bytes
+                return (
+                    progress.download_progress.downloaded.in_bytes / total
+                    if total > 0
+                    else 0.0
+                )
+            case DownloadPending():
+                total = progress.total.in_bytes
+                return progress.downloaded.in_bytes / total if total > 0 else 0.0
+            case DownloadFailed():
+                return 0.0
+    return 0.0
+
+
+def _cycle_download_score(
+    cycle: Cycle,
+    model_id: ModelId,
+    download_status: Mapping[NodeId, Sequence[DownloadProgress]],
+) -> float:
+    """Sum of download fractions across all nodes in a cycle."""
+    return sum(
+        _get_node_download_fraction(node_id, model_id, download_status)
+        for node_id in cycle
+    )
+
+
 def place_instance(
     command: PlaceInstance,
     topology: Topology,
@@ -67,6 +109,7 @@ def place_instance(
     node_memory: Mapping[NodeId, MemoryUsage],
     node_network: Mapping[NodeId, NodeNetworkInfo],
     required_nodes: set[NodeId] | None = None,
+    download_status: Mapping[NodeId, Sequence[DownloadProgress]] | None = None,
 ) -> dict[InstanceId, Instance]:
     cycles = topology.get_cycles()
     candidate_cycles = list(filter(lambda it: len(it) >= command.min_nodes, cycles))
@@ -130,11 +173,21 @@ def place_instance(
         if any(topology.node_is_leaf(node_id) for node_id in cycle)
     ]
 
+    resolved_download_status = download_status or {}
+    candidate_cycles = (
+        cycles_with_leaf_nodes if cycles_with_leaf_nodes != [] else smallest_cycles
+    )
+
     selected_cycle = max(
-        cycles_with_leaf_nodes if cycles_with_leaf_nodes != [] else smallest_cycles,
-        key=lambda cycle: sum(
-            (node_memory[node_id].ram_available for node_id in cycle),
-            start=Memory(),
+        candidate_cycles,
+        key=lambda cycle: (
+            _cycle_download_score(
+                cycle, command.model_card.model_id, resolved_download_status
+            ),
+            sum(
+                (node_memory[node_id].ram_available for node_id in cycle),
+                start=Memory(),
+            ),
         ),
     )
 

src/exo/master/tests/test_placement.py

@@ -25,6 +25,12 @@ from exo.shared.types.profiling import NetworkInterfaceInfo, NodeNetworkInfo
 from exo.shared.types.tasks import TaskId, TaskStatus, TextGeneration
 from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
 from exo.shared.types.topology import Connection, SocketConnection
+from exo.shared.types.worker.downloads import (
+    DownloadCompleted,
+    DownloadFailed,
+    DownloadOngoing,
+    DownloadProgressData,
+)
 from exo.shared.types.worker.instances import (
     Instance,
     InstanceId,
@@ -33,7 +39,7 @@ from exo.shared.types.worker.instances import (
     MlxRingInstance,
 )
 from exo.shared.types.worker.runners import ShardAssignments
-from exo.shared.types.worker.shards import Sharding
+from exo.shared.types.worker.shards import PipelineShardMetadata, Sharding
 
 
 @pytest.fixture
@@ -576,3 +582,183 @@ def test_get_transition_events_delete_instance_cancels_only_matching_tasks(
     assert cancel_events[0].task_status == TaskStatus.Cancelled
     assert len(delete_events) == 1
     assert delete_events[0].instance_id == instance_id_a
+
+
+def _make_shard_metadata(model_card: ModelCard) -> PipelineShardMetadata:
+    return PipelineShardMetadata(
+        model_card=model_card,
+        device_rank=0,
+        world_size=1,
+        start_layer=0,
+        end_layer=model_card.n_layers,
+        n_layers=model_card.n_layers,
+    )
+
+
+def test_placement_prefers_cycle_with_downloaded_model(
+    model_card: ModelCard,
+) -> None:
+    """When two cycles are otherwise equal, prefer the one with the model already downloaded."""
+    topology = Topology()
+
+    model_card.storage_size = Memory.from_bytes(500)
+
+    node_a = NodeId()
+    node_b = NodeId()
+
+    node_memory = {
+        node_a: create_node_memory(1000),
+        node_b: create_node_memory(1000),
+    }
+    node_network = {
+        node_a: create_node_network(),
+        node_b: create_node_network(),
+    }
+
+    topology.add_node(node_a)
+    topology.add_node(node_b)
+    # No connections between them — two single-node cycles
+
+    shard_meta = _make_shard_metadata(model_card)
+
+    # node_b has the model fully downloaded, node_a does not
+    download_status = {
+        node_b: [
+            DownloadCompleted(
+                node_id=node_b,
+                shard_metadata=shard_meta,
+                total=model_card.storage_size,
+            ),
+        ],
+    }
+
+    cic = place_instance_command(model_card)
+    placements = place_instance(
+        cic, topology, {}, node_memory, node_network, download_status=download_status
+    )
+
+    assert len(placements) == 1
+    instance = list(placements.values())[0]
+    assigned_nodes = set(instance.shard_assignments.node_to_runner.keys())
+    assert assigned_nodes == {node_b}
+
+
+def test_placement_prefers_cycle_with_higher_download_progress(
+    model_card: ModelCard,
+) -> None:
+    """When two cycles are otherwise equal, prefer the one with more download progress."""
+    topology = Topology()
+
+    model_card.storage_size = Memory.from_bytes(1000)
+
+    node_a = NodeId()
+    node_b = NodeId()
+
+    node_memory = {
+        node_a: create_node_memory(1000),
+        node_b: create_node_memory(1000),
+    }
+    node_network = {
+        node_a: create_node_network(),
+        node_b: create_node_network(),
+    }
+
+    topology.add_node(node_a)
+    topology.add_node(node_b)
+
+    shard_meta = _make_shard_metadata(model_card)
+
+    # node_a: 30% downloaded, node_b: 80% downloaded
+    download_status = {
+        node_a: [
+            DownloadOngoing(
+                node_id=node_a,
+                shard_metadata=shard_meta,
+                download_progress=DownloadProgressData(
+                    total=Memory.from_bytes(1000),
+                    downloaded=Memory.from_bytes(300),
+                    downloaded_this_session=Memory.from_bytes(300),
+                    completed_files=0,
+                    total_files=1,
+                    speed=0.0,
+                    eta_ms=0,
+                    files={},
+                ),
+            ),
+        ],
+        node_b: [
+            DownloadOngoing(
+                node_id=node_b,
+                shard_metadata=shard_meta,
+                download_progress=DownloadProgressData(
+                    total=Memory.from_bytes(1000),
+                    downloaded=Memory.from_bytes(800),
+                    downloaded_this_session=Memory.from_bytes(800),
+                    completed_files=0,
+                    total_files=1,
+                    speed=0.0,
+                    eta_ms=0,
+                    files={},
+                ),
+            ),
+        ],
+    }
+
+    cic = place_instance_command(model_card)
+    placements = place_instance(
+        cic, topology, {}, node_memory, node_network, download_status=download_status
+    )
+
+    assert len(placements) == 1
+    instance = list(placements.values())[0]
+    assigned_nodes = set(instance.shard_assignments.node_to_runner.keys())
+    assert assigned_nodes == {node_b}
+
+
+def test_placement_does_not_prefer_cycle_with_failed_download(
+    model_card: ModelCard,
+) -> None:
+    """A failed download should count as 0% — not preferred over a node with no download history."""
+    topology = Topology()
+
+    model_card.storage_size = Memory.from_bytes(500)
+
+    node_a = NodeId()
+    node_b = NodeId()
+
+    # node_a has slightly more RAM so it would win on the RAM tiebreaker
+    node_memory = {
+        node_a: create_node_memory(1001),
+        node_b: create_node_memory(1000),
+    }
+    node_network = {
+        node_a: create_node_network(),
+        node_b: create_node_network(),
+    }
+
+    topology.add_node(node_a)
+    topology.add_node(node_b)
+
+    shard_meta = _make_shard_metadata(model_card)
+
+    # node_b has a failed download — should not be preferred
+    download_status = {
+        node_b: [
+            DownloadFailed(
+                node_id=node_b,
+                shard_metadata=shard_meta,
+                error_message="connection reset",
+            ),
+        ],
+    }
+
+    cic = place_instance_command(model_card)
+    placements = place_instance(
+        cic, topology, {}, node_memory, node_network, download_status=download_status
+    )
+
+    assert len(placements) == 1
+    instance = list(placements.values())[0]
+    assigned_nodes = set(instance.shard_assignments.node_to_runner.keys())
+    # node_a should win on RAM tiebreaker since failed download scores 0.0
+    assert assigned_nodes == {node_a}

src/exo/routing/router.py

@@ -1,3 +1,4 @@
+from collections.abc import Sequence
 from copy import copy
 from itertools import count
 from math import inf
@@ -102,8 +103,15 @@ class TopicRouter[T: CamelCaseModel]:
 
 class Router:
     @classmethod
-    def create(cls, identity: Keypair) -> "Router":
-        return cls(handle=NetworkingHandle(identity))
+    def create(
+        cls,
+        identity: Keypair,
+        bootstrap_peers: Sequence[str] = (),
+        listen_port: int = 0,
+    ) -> "Router":
+        return cls(
+            handle=NetworkingHandle(identity, list(bootstrap_peers), listen_port)
+        )
 
     def __init__(self, handle: NetworkingHandle):
         self.topic_routers: dict[str, TopicRouter[CamelCaseModel]] = {}

src/exo/shared/apply.py

@@ -7,6 +7,8 @@ from loguru import logger
 from exo.shared.types.common import NodeId
 from exo.shared.types.events import (
     ChunkGenerated,
+    CustomModelCardAdded,
+    CustomModelCardDeleted,
     Event,
     IndexedEvent,
     InputChunkReceived,
@@ -65,6 +67,8 @@ def event_apply(event: Event, state: State) -> State:
             | InputChunkReceived()
             | TracesCollected()
             | TracesMerged()
+            | CustomModelCardAdded()
+            | CustomModelCardDeleted()
         ):  # Pass-through events that don't modify state
             return state
         case InstanceCreated():

src/exo/shared/constants.py

@@ -26,21 +26,35 @@ EXO_CONFIG_HOME = _get_xdg_dir("XDG_CONFIG_HOME", ".config")
 EXO_DATA_HOME = _get_xdg_dir("XDG_DATA_HOME", ".local/share")
 EXO_CACHE_HOME = _get_xdg_dir("XDG_CACHE_HOME", ".cache")
 
-# Models directory (data)
-_EXO_MODELS_DIR_ENV = os.environ.get("EXO_MODELS_DIR", None)
-EXO_MODELS_DIR = (
-    EXO_DATA_HOME / "models"
-    if _EXO_MODELS_DIR_ENV is None
-    else Path.home() / _EXO_MODELS_DIR_ENV
+# Default models directory (always included as first entry in writable dirs)
+_EXO_DEFAULT_MODELS_DIR_ENV = os.environ.get("EXO_DEFAULT_MODELS_DIR", None)
+EXO_DEFAULT_MODELS_DIR = (
+    Path(_EXO_DEFAULT_MODELS_DIR_ENV).expanduser()
+    if _EXO_DEFAULT_MODELS_DIR_ENV is not None
+    else EXO_DATA_HOME / "models"
 )
 
-# Read-only search path for pre-downloaded models (colon-separated directories)
-_EXO_MODELS_PATH_ENV = os.environ.get("EXO_MODELS_PATH", None)
-EXO_MODELS_PATH: tuple[Path, ...] | None = (
-    tuple(Path(p).expanduser() for p in _EXO_MODELS_PATH_ENV.split(":") if p)
-    if _EXO_MODELS_PATH_ENV is not None
-    else None
+
+def _parse_colon_dirs(env_var: str) -> tuple[Path, ...]:
+    raw = os.environ.get(env_var, None)
+    if raw is None:
+        return ()
+    return tuple(Path(p).expanduser() for p in raw.split(":") if p)
+
+
+# Read-only model directories (colon-separated). Never written to or deleted from.
+_EXO_MODELS_READ_ONLY_DIRS_ENV = _parse_colon_dirs("EXO_MODELS_READ_ONLY_DIRS")
+# Writable model directories (colon-separated). Default dir is always prepended.
+_EXO_MODELS_DIRS_ENV = _parse_colon_dirs("EXO_MODELS_DIRS")
+
+# If a directory appears in both lists, treat it as read-only.
+_read_only_set = frozenset(_EXO_MODELS_READ_ONLY_DIRS_ENV)
+EXO_MODELS_DIRS: tuple[Path, ...] = tuple(
+    d
+    for d in (EXO_DEFAULT_MODELS_DIR, *_EXO_MODELS_DIRS_ENV)
+    if d not in _read_only_set
 )
+EXO_MODELS_READ_ONLY_DIRS: tuple[Path, ...] = _EXO_MODELS_READ_ONLY_DIRS_ENV
 
 _RESOURCES_DIR_ENV = os.environ.get("EXO_RESOURCES_DIR", None)
 RESOURCES_DIR = (

src/exo/shared/logging.py

@@ -66,7 +66,7 @@ def logger_setup(log_file: Path | None, verbosity: int = 0):
     else:
         logger.add(
             sys.__stderr__,  # type: ignore
-            format="[ {time:HH:mm:ss.SSS} | <level>{level: <8}</level> | {name}:{function}:{line} ] <level>{message}</level>",
+            format="[ {time:YYYY-MM-DD HH:mm:ss.SSS} | <level>{level: <8}</level> | {name}:{function}:{line} ] <level>{message}</level>",
             level="DEBUG",
             colorize=True,
             enqueue=True,
@@ -76,7 +76,7 @@ def logger_setup(log_file: Path | None, verbosity: int = 0):
         logger.add(
             log_file,
             format="[ {time:YYYY-MM-DD HH:mm:ss.SSS} | {level: <8} | {name}:{function}:{line} ] {message}",
-            level="INFO",
+            level="DEBUG" if verbosity > 0 else "INFO",
             colorize=False,
             enqueue=True,
             rotation=lambda _, __: next(rotate_once),

src/exo/shared/models/model_cards.py

@@ -30,30 +30,42 @@ from exo.utils.pydantic_ext import CamelCaseModel
 # kinda ugly...
 # TODO: load search path from config.toml
 _custom_cards_dir = Path(str(EXO_CUSTOM_MODEL_CARDS_DIR))
-CARD_SEARCH_PATH = [
+_BUILTIN_CARD_DIRS = [
     Path(RESOURCES_DIR) / "inference_model_cards",
     Path(RESOURCES_DIR) / "image_model_cards",
-    _custom_cards_dir,
 ]
 
 _card_cache: dict[ModelId, "ModelCard"] = {}
 
 
-async def _refresh_card_cache():
-    for path in CARD_SEARCH_PATH:
-        async for toml_file in path.rglob("*.toml"):
-            try:
-                card = await ModelCard.load_from_path(toml_file)
-                if card.model_id not in _card_cache:
-                    _card_cache[card.model_id] = card
-            except (ValidationError, TOMLKitError):
-                pass
+async def _load_cards_from_dir(directory: Path, *, is_custom: bool) -> None:
+    """Load all TOML model cards from a directory into the cache."""
+    async for toml_file in directory.rglob("*.toml"):
+        try:
+            card = await ModelCard.load_from_path(toml_file)
+            if is_custom:
+                card = card.model_copy(update={"is_custom": True})
+            if card.model_id not in _card_cache:
+                _card_cache[card.model_id] = card
+        except (ValidationError, TOMLKitError):
+            pass
+
+
+async def _refresh_card_cache() -> None:
+    for path in _BUILTIN_CARD_DIRS:
+        await _load_cards_from_dir(path, is_custom=False)
+    await _load_cards_from_dir(_custom_cards_dir, is_custom=True)
 
 
 def _is_image_card(card: "ModelCard") -> bool:
     return any(t in (ModelTask.TextToImage, ModelTask.ImageToImage) for t in card.tasks)
 
 
+def get_card(model_id: ModelId) -> "ModelCard | None":
+    """Look up a single model card from the cache by ID."""
+    return _card_cache.get(model_id)
+
+
 async def get_model_cards() -> list["ModelCard"]:
     if len(_card_cache) == 0:
         await _refresh_card_cache()
@@ -92,6 +104,7 @@ class ModelCard(CamelCaseModel):
     capabilities: list[str] = []
     uses_cfg: bool = False
     trust_remote_code: bool = True
+    is_custom: bool = False
 
     @field_validator("tasks", mode="before")
     @classmethod
@@ -100,7 +113,7 @@ class ModelCard(CamelCaseModel):
 
     async def save(self, path: Path) -> None:
         async with await open_file(path, "w") as f:
-            py = self.model_dump(exclude_none=True)
+            py = self.model_dump(exclude_none=True, exclude={"is_custom"})
             data = tomlkit.dumps(py)  # pyright: ignore[reportUnknownMemberType]
             await f.write(data)
 
@@ -122,17 +135,24 @@ class ModelCard(CamelCaseModel):
         if (mc := _card_cache.get(model_id)) is not None:
             return mc
 
-        return await ModelCard.fetch_from_hf(model_id)
+        mc = await ModelCard.fetch_from_hf(model_id)
+        await mc.save_to_custom_dir()
+        _card_cache[model_id] = mc
+        return mc
 
     @staticmethod
     async def fetch_from_hf(model_id: ModelId) -> "ModelCard":
-        """Fetches storage size and number of layers for a Hugging Face model, returns Pydantic ModelMeta."""
+        """Fetches storage size and number of layers for a Hugging Face model, returns Pydantic ModelMeta.
+
+        This is a pure fetch — it does NOT save to disk or update the cache.
+        Persistence is handled by the event-sourcing layer (worker event handler).
+        """
         # TODO: failure if files do not exist
         config_data = await fetch_config_data(model_id)
         num_layers = config_data.layer_count
         mem_size_bytes = await fetch_safetensors_size(model_id)
 
-        mc = ModelCard(
+        return ModelCard(
             model_id=ModelId(model_id),
             storage_size=mem_size_bytes,
             n_layers=num_layers,
@@ -141,10 +161,13 @@ class ModelCard(CamelCaseModel):
             num_key_value_heads=config_data.num_key_value_heads,
             tasks=[ModelTask.TextGeneration],
             trust_remote_code=False,
+            is_custom=True,
         )
-        await mc.save_to_custom_dir()
-        _card_cache[model_id] = mc
-        return mc
+
+
+def add_to_card_cache(card: "ModelCard") -> None:
+    """Add or update a model card in the in-memory cache."""
+    _card_cache[card.model_id] = card
 
 
 async def delete_custom_card(model_id: ModelId) -> bool:
@@ -157,16 +180,6 @@ async def delete_custom_card(model_id: ModelId) -> bool:
     return False
 
 
-def is_custom_card(model_id: ModelId) -> bool:
-    """Check if a model card exists in the custom cards directory."""
-    import os
-
-    card_path = Path(str(EXO_CUSTOM_MODEL_CARDS_DIR)) / (
-        ModelId(model_id).normalize() + ".toml"
-    )
-    return os.path.isfile(str(card_path))
-
-
 class ConfigData(BaseModel):
     model_config = {"extra": "ignore"}  # Allow unknown fields
 
@@ -230,11 +243,10 @@ async def fetch_config_data(model_id: ModelId) -> ConfigData:
     """Downloads and parses config.json for a model."""
     from exo.download.download_utils import (
         download_file_with_retry,
-        ensure_models_dir,
+        resolve_model_dir,
     )
 
-    target_dir = (await ensure_models_dir()) / model_id.normalize()
-    await aios.makedirs(target_dir, exist_ok=True)
+    target_dir = await resolve_model_dir(model_id)
     config_path = await download_file_with_retry(
         model_id,
         "main",
@@ -252,12 +264,11 @@ async def fetch_safetensors_size(model_id: ModelId) -> Memory:
     """Gets model size from safetensors index or falls back to HF API."""
     from exo.download.download_utils import (
         download_file_with_retry,
-        ensure_models_dir,
+        resolve_model_dir,
     )
     from exo.shared.types.worker.downloads import ModelSafetensorsIndex
 
-    target_dir = (await ensure_models_dir()) / model_id.normalize()
-    await aios.makedirs(target_dir, exist_ok=True)
+    target_dir = await resolve_model_dir(model_id)
     index_path = await download_file_with_retry(
         model_id,
         "main",

src/exo/shared/tests/test_xdg_paths.py

@@ -105,9 +105,9 @@ def test_node_id_in_config_dir():
 
 
 def test_models_in_data_dir():
-    """Test that models directory is in the data directory."""
-    # Clear EXO_MODELS_DIR to test default behavior
-    env = {k: v for k, v in os.environ.items() if k != "EXO_MODELS_DIR"}
+    """Test that default models directory is in the data directory."""
+    # Clear EXO_MODELS_DIRS to test default behavior
+    env = {k: v for k, v in os.environ.items() if k != "EXO_MODELS_DIRS"}
     with mock.patch.dict(os.environ, env, clear=True):
         import importlib
 
@@ -115,4 +115,106 @@ def test_models_in_data_dir():
 
         importlib.reload(constants)
 
-        assert constants.EXO_MODELS_DIR.parent == constants.EXO_DATA_HOME
+        assert constants.EXO_DEFAULT_MODELS_DIR.parent == constants.EXO_DATA_HOME
+
+
+def test_default_dir_always_prepended_to_models_dirs():
+    """Test that the default models dir is always the first entry in EXO_MODELS_DIRS."""
+    env = {
+        k: v
+        for k, v in os.environ.items()
+        if k not in ("EXO_MODELS_DIRS", "EXO_MODELS_READ_ONLY_DIRS", "EXO_HOME")
+    }
+    env["EXO_MODELS_DIRS"] = "/tmp/custom-models"
+    with mock.patch.dict(os.environ, env, clear=True):
+        import importlib
+
+        import exo.shared.constants as constants
+
+        importlib.reload(constants)
+
+        assert constants.EXO_MODELS_DIRS[0] == constants.EXO_DEFAULT_MODELS_DIR
+        assert Path("/tmp/custom-models") in constants.EXO_MODELS_DIRS
+
+
+def test_default_models_dir_override():
+    """Test that EXO_DEFAULT_MODELS_DIR can be overridden via env var."""
+    env = {
+        k: v
+        for k, v in os.environ.items()
+        if k
+        not in (
+            "EXO_MODELS_DIRS",
+            "EXO_MODELS_READ_ONLY_DIRS",
+            "EXO_HOME",
+            "EXO_DEFAULT_MODELS_DIR",
+        )
+    }
+    env["EXO_DEFAULT_MODELS_DIR"] = "/Volumes/FastSSD/exo-models"
+    with mock.patch.dict(os.environ, env, clear=True):
+        import importlib
+
+        import exo.shared.constants as constants
+
+        importlib.reload(constants)
+
+        assert Path("/Volumes/FastSSD/exo-models") == constants.EXO_DEFAULT_MODELS_DIR
+        assert constants.EXO_MODELS_DIRS[0] == constants.EXO_DEFAULT_MODELS_DIR
+
+
+def test_default_dir_only_entry_when_env_unset():
+    """Test that EXO_MODELS_DIRS contains only the default when env var is not set."""
+    env = {
+        k: v
+        for k, v in os.environ.items()
+        if k not in ("EXO_MODELS_DIRS", "EXO_MODELS_READ_ONLY_DIRS", "EXO_HOME")
+    }
+    with mock.patch.dict(os.environ, env, clear=True):
+        import importlib
+
+        import exo.shared.constants as constants
+
+        importlib.reload(constants)
+
+        assert constants.EXO_MODELS_DIRS == (constants.EXO_DEFAULT_MODELS_DIR,)
+
+
+def test_overlap_between_dirs_and_read_only_dirs():
+    """Test that a directory in both lists is excluded from writable dirs."""
+    env = {
+        k: v
+        for k, v in os.environ.items()
+        if k not in ("EXO_MODELS_DIRS", "EXO_MODELS_READ_ONLY_DIRS", "EXO_HOME")
+    }
+    env["EXO_MODELS_DIRS"] = "/tmp/shared:/tmp/writable-only"
+    env["EXO_MODELS_READ_ONLY_DIRS"] = "/tmp/shared:/tmp/ro-only"
+    with mock.patch.dict(os.environ, env, clear=True):
+        import importlib
+
+        import exo.shared.constants as constants
+
+        importlib.reload(constants)
+
+        # /tmp/shared should be excluded from writable dirs
+        assert Path("/tmp/shared") not in constants.EXO_MODELS_DIRS
+        assert Path("/tmp/writable-only") in constants.EXO_MODELS_DIRS
+        # /tmp/shared should still be in read-only dirs
+        assert Path("/tmp/shared") in constants.EXO_MODELS_READ_ONLY_DIRS
+        assert Path("/tmp/ro-only") in constants.EXO_MODELS_READ_ONLY_DIRS
+
+
+def test_empty_read_only_dirs_when_unset():
+    """Test that EXO_MODELS_READ_ONLY_DIRS is empty when env var is not set."""
+    env = {
+        k: v
+        for k, v in os.environ.items()
+        if k not in ("EXO_MODELS_DIRS", "EXO_MODELS_READ_ONLY_DIRS", "EXO_HOME")
+    }
+    with mock.patch.dict(os.environ, env, clear=True):
+        import importlib
+
+        import exo.shared.constants as constants
+
+        importlib.reload(constants)
+
+        assert constants.EXO_MODELS_READ_ONLY_DIRS == ()

src/exo/shared/types/commands.py

@@ -81,6 +81,14 @@ class CancelDownload(BaseCommand):
     model_id: ModelId
 
 
+class AddCustomModelCard(BaseCommand):
+    model_card: ModelCard
+
+
+class DeleteCustomModelCard(BaseCommand):
+    model_id: ModelId
+
+
 DownloadCommand = StartDownload | DeleteDownload | CancelDownload
 
 
@@ -96,6 +104,8 @@ Command = (
     | TaskCancelled
     | TaskFinished
     | SendInputChunk
+    | AddCustomModelCard
+    | DeleteCustomModelCard
 )
 
 

src/exo/shared/types/events.py

@@ -3,9 +3,10 @@ from typing import final
 
 from pydantic import Field
 
+from exo.shared.models.model_cards import ModelCard
 from exo.shared.topology import Connection
 from exo.shared.types.chunks import GenerationChunk, InputImageChunk
-from exo.shared.types.common import CommandId, Id, NodeId, SessionId, SystemId
+from exo.shared.types.common import CommandId, Id, ModelId, NodeId, SessionId, SystemId
 from exo.shared.types.tasks import Task, TaskId, TaskStatus
 from exo.shared.types.worker.downloads import DownloadProgress
 from exo.shared.types.worker.instances import Instance, InstanceId
@@ -106,6 +107,14 @@ class TopologyEdgeDeleted(BaseEvent):
     conn: Connection
 
 
+class CustomModelCardAdded(BaseEvent):
+    model_card: ModelCard
+
+
+class CustomModelCardDeleted(BaseEvent):
+    model_id: ModelId
+
+
 @final
 class TraceEventData(FrozenModel):
     name: str
@@ -147,6 +156,8 @@ Event = (
     | TopologyEdgeDeleted
     | TracesCollected
     | TracesMerged
+    | CustomModelCardAdded
+    | CustomModelCardDeleted
 )
 
 

src/exo/utils/info_gatherer/info_gatherer.py

@@ -10,11 +10,10 @@ from typing import Self, cast
 import anyio
 from anyio import fail_after, open_process, to_thread
 from anyio.streams.buffered import BufferedByteReceiveStream
-from anyio.streams.text import TextReceiveStream
 from loguru import logger
 from pydantic import ValidationError
 
-from exo.shared.constants import EXO_CONFIG_FILE, EXO_MODELS_DIR
+from exo.shared.constants import EXO_CONFIG_FILE, EXO_DEFAULT_MODELS_DIR
 from exo.shared.types.memory import Memory
 from exo.shared.types.profiling import (
     DiskUsage,
@@ -287,8 +286,8 @@ class ThunderboltBridgeInfo(TaggedModel):
                     service_name=tb_service_name,
                 )
             )
-        except Exception:
-            logger.opt(exception=True).warning("Failed to gather Thunderbolt Bridge info")
+        except Exception as e:
+            logger.opt(exception=e).warning("Failed to gather Thunderbolt Bridge info")
             return None
 
 
@@ -329,7 +328,7 @@ class NodeDiskUsage(TaggedModel):
     async def gather(cls) -> Self:
         return cls(
             disk_usage=await to_thread.run_sync(
-                lambda: DiskUsage.from_path(EXO_MODELS_DIR)
+                DiskUsage.from_path, EXO_DEFAULT_MODELS_DIR
             )
         )
 
@@ -372,26 +371,8 @@ GatheredInfo = (
 @dataclass
 class InfoGatherer:
     info_sender: Sender[GatheredInfo]
-    interface_watcher_interval: float | None = 10
-    misc_poll_interval: float | None = 60
-    system_profiler_interval: float | None = 5 if IS_DARWIN else None
-    memory_poll_rate: float | None = None if IS_DARWIN else 1
-    macmon_interval: float | None = 1 if IS_DARWIN else None
-    thunderbolt_bridge_poll_interval: float | None = 10 if IS_DARWIN else None
-    static_info_poll_interval: float | None = 60
-    rdma_ctl_poll_interval: float | None = 10 if IS_DARWIN else None
-    disk_poll_interval: float | None = 30
     _tg: TaskGroup = field(init=False, default_factory=TaskGroup)
-    _psutil_memory_fallback_enabled: bool = field(init=False, default=False)
-
-    def _enable_psutil_memory_fallback(self, reason: str) -> None:
-        if not self._psutil_memory_fallback_enabled:
-            logger.warning(reason)
-            self._psutil_memory_fallback_enabled = True
-        self.memory_poll_rate = 1
-
-    def _get_macmon_path(self) -> str | None:
-        return os.getenv("EXO_MACMON_PATH") or shutil.which("macmon")
+    _psutil_enabled: bool = field(init=False, default=False)
 
     async def _can_read_macmon_metrics(self, macmon_path: str) -> bool:
         try:
@@ -400,8 +381,8 @@ class InfoGatherer:
                     [macmon_path, "pipe", "--samples", "1", "--interval", "100"],
                     check=False,
                 )
-        except Exception:
-            logger.opt(exception=True).warning(
+        except Exception as e:
+            logger.opt(exception=e).warning(
                 f"Failed to validate macmon at {macmon_path}"
             )
             return False
@@ -421,8 +402,8 @@ class InfoGatherer:
 
         try:
             MacmonMetrics.from_raw_json(stdout.splitlines()[0])
-        except ValidationError:
-            logger.opt(exception=True).warning(
+        except ValidationError as e:
+            logger.opt(exception=e).warning(
                 "macmon preflight returned unexpected metrics JSON"
             )
             return False
@@ -432,27 +413,16 @@ class InfoGatherer:
     async def run(self):
         async with self._tg as tg:
             if IS_DARWIN:
-                if (macmon_path := self._get_macmon_path()) is not None:
-                    if await self._can_read_macmon_metrics(macmon_path):
-                        tg.start_soon(self._monitor_macmon, macmon_path)
-                    else:
-                        self._enable_psutil_memory_fallback(
-                            f"macmon at {macmon_path} is unusable, falling back "
-                            f"to psutil memory monitoring"
-                        )
-                else:
-                    self._enable_psutil_memory_fallback(
-                        "macmon not found, falling back to psutil for memory "
-                        "monitoring"
-                    )
-                tg.start_soon(self._monitor_system_profiler_thunderbolt_data)
-                tg.start_soon(self._monitor_thunderbolt_bridge_status)
-                tg.start_soon(self._monitor_rdma_ctl_status)
-            tg.start_soon(self._watch_system_info)
-            tg.start_soon(self._monitor_memory_usage)
-            tg.start_soon(self._monitor_misc)
-            tg.start_soon(self._monitor_static_info)
-            tg.start_soon(self._monitor_disk_usage)
+                tg.start_soon(self._monitor_macmon, 1)
+                tg.start_soon(self._monitor_system_profiler_thunderbolt_data, 5)
+                tg.start_soon(self._monitor_thunderbolt_bridge_status, 10)
+                tg.start_soon(self._monitor_rdma_ctl_status, 10)
+            if not IS_DARWIN:
+                tg.start_soon(self._monitor_memory_usage, 1)
+            tg.start_soon(self._watch_system_info, 10)
+            tg.start_soon(self._monitor_misc, 60)
+            tg.start_soon(self._monitor_static_info, 60)
+            tg.start_soon(self._monitor_disk_usage, 30)
 
             nc = await NodeConfig.gather()
             if nc is not None:
@@ -461,32 +431,27 @@ class InfoGatherer:
     def shutdown(self):
         self._tg.cancel_tasks()
 
-    async def _monitor_static_info(self):
-        if self.static_info_poll_interval is None:
-            return
+    async def _monitor_static_info(self, static_info_poll_interval: float):
         while True:
             try:
                 with fail_after(30):
                     await self.info_sender.send(await StaticNodeInformation.gather())
-            except Exception:
-                logger.opt(exception=True).warning("Error gathering static node info")
-            await anyio.sleep(self.static_info_poll_interval)
+            except Exception as e:
+                logger.opt(exception=e).warning("Error gathering static node info")
+            await anyio.sleep(static_info_poll_interval)
 
-    async def _monitor_misc(self):
-        if self.misc_poll_interval is None:
-            return
+    async def _monitor_misc(self, misc_poll_interval: float):
         while True:
             try:
                 with fail_after(10):
                     await self.info_sender.send(await MiscData.gather())
-            except Exception:
-                logger.opt(exception=True).warning("Error gathering misc data")
-            await anyio.sleep(self.misc_poll_interval)
-
-    async def _monitor_system_profiler_thunderbolt_data(self):
-        if self.system_profiler_interval is None:
-            return
+            except Exception as e:
+                logger.opt(exception=e).warning("Error gathering misc data")
+            await anyio.sleep(misc_poll_interval)
 
+    async def _monitor_system_profiler_thunderbolt_data(
+        self, system_profiler_interval: float
+    ):
         while True:
             try:
                 with fail_after(30):
@@ -506,11 +471,14 @@ class InfoGatherer:
 
                     conns = [it for i in data if (it := i.conn()) is not None]
                     await self.info_sender.send(MacThunderboltConnections(conns=conns))
-            except Exception:
-                logger.opt(exception=True).warning("Error gathering Thunderbolt data")
-            await anyio.sleep(self.system_profiler_interval)
+            except Exception as e:
+                logger.opt(exception=e).warning("Error gathering Thunderbolt data")
+            await anyio.sleep(system_profiler_interval)
 
-    async def _monitor_memory_usage(self):
+    async def _monitor_memory_usage(self, memory_poll_rate: float):
+        if self._psutil_enabled:
+            return
+        self._psutil_enabled = True
         override_memory_env = os.getenv("OVERRIDE_MEMORY_MB")
         override_memory: int | None = (
             Memory.from_mb(int(override_memory_env)).in_bytes
@@ -518,73 +486,77 @@ class InfoGatherer:
             else None
         )
         while True:
-            poll_rate = self.memory_poll_rate
-            if poll_rate is None:
-                await anyio.sleep(1)
-                continue
             try:
                 await self.info_sender.send(
                     MemoryUsage.from_psutil(override_memory=override_memory)
                 )
-            except Exception:
-                logger.opt(exception=True).warning("Error gathering memory usage")
-            await anyio.sleep(poll_rate)
+            except Exception as e:
+                logger.opt(exception=e).warning("Error gathering memory usage")
+            await anyio.sleep(memory_poll_rate)
 
-    async def _watch_system_info(self):
-        if self.interface_watcher_interval is None:
-            return
+    async def _watch_system_info(self, interface_watcher_interval: float):
         while True:
             try:
                 with fail_after(10):
                     nics = await get_network_interfaces()
                     await self.info_sender.send(NodeNetworkInterfaces(ifaces=nics))
-            except Exception:
-                logger.opt(exception=True).warning("Error gathering network interfaces")
-            await anyio.sleep(self.interface_watcher_interval)
+            except Exception as e:
+                logger.opt(exception=e).warning("Error gathering network interfaces")
+            await anyio.sleep(interface_watcher_interval)
 
-    async def _monitor_thunderbolt_bridge_status(self):
-        if self.thunderbolt_bridge_poll_interval is None:
-            return
+    async def _monitor_thunderbolt_bridge_status(
+        self, thunderbolt_bridge_poll_interval: float
+    ):
         while True:
             try:
                 with fail_after(30):
                     curr = await ThunderboltBridgeInfo.gather()
                     if curr is not None:
                         await self.info_sender.send(curr)
-            except Exception:
-                logger.opt(exception=True).warning("Error gathering Thunderbolt Bridge status")
-            await anyio.sleep(self.thunderbolt_bridge_poll_interval)
+            except Exception as e:
+                logger.opt(exception=e).warning(
+                    "Error gathering Thunderbolt Bridge status"
+                )
+            await anyio.sleep(thunderbolt_bridge_poll_interval)
 
-    async def _monitor_rdma_ctl_status(self):
-        if self.rdma_ctl_poll_interval is None:
-            return
+    async def _monitor_rdma_ctl_status(self, rdma_ctl_poll_interval: float):
         while True:
             try:
                 curr = await RdmaCtlStatus.gather()
                 if curr is not None:
                     await self.info_sender.send(curr)
-            except Exception:
-                logger.opt(exception=True).warning("Error gathering RDMA ctl status")
-            await anyio.sleep(self.rdma_ctl_poll_interval)
+            except Exception as e:
+                logger.opt(exception=e).warning("Error gathering RDMA ctl status")
+            await anyio.sleep(rdma_ctl_poll_interval)
 
-    async def _monitor_disk_usage(self):
-        if self.disk_poll_interval is None:
-            return
+    async def _monitor_disk_usage(self, disk_poll_interval: float):
         while True:
             try:
                 with fail_after(5):
                     await self.info_sender.send(await NodeDiskUsage.gather())
-            except Exception:
-                logger.opt(exception=True).warning("Error gathering disk usage")
-            await anyio.sleep(self.disk_poll_interval)
+            except Exception as e:
+                logger.opt(exception=e).warning("Error gathering disk usage")
+            await anyio.sleep(disk_poll_interval)
 
-    async def _monitor_macmon(self, macmon_path: str):
-        if self.macmon_interval is None:
+    async def _monitor_macmon(self, macmon_interval: float):
+        if (
+            macmon_path := os.getenv("EXO_MACMON_PATH") or shutil.which("macmon")
+        ) is None:
+            logger.warning(
+                "macmon not found, falling back to psutil for memory monitoring"
+            )
+            self._tg.start_soon(self._monitor_memory_usage, 1)
+            return
+        if not await self._can_read_macmon_metrics(macmon_path):
+            logger.warning(
+                f"macmon at {macmon_path} is unusable, falling back to psutil memory monitoring"
+            )
+            self._tg.start_soon(self._monitor_memory_usage, 1)
             return
         # macmon pipe --interval [interval in ms]
         # Timeout: if macmon produces no output for this many seconds, restart it.
         # macmon writes every macmon_interval seconds, so 10x that is generous.
-        read_timeout = max(self.macmon_interval * 10, 30)
+        read_timeout = max(macmon_interval * 10, 30)
         while True:
             try:
                 async with await open_process(
@@ -592,25 +564,26 @@ class InfoGatherer:
                         macmon_path,
                         "pipe",
                         "--interval",
-                        str(self.macmon_interval * 1000),
+                        str(macmon_interval * 1000),
                     ]
                 ) as p:
                     if not p.stdout:
                         logger.critical("MacMon closed stdout")
                         return
-                    stream = TextReceiveStream(BufferedByteReceiveStream(p.stdout))
+                    stream = BufferedByteReceiveStream(p.stdout)
                     while True:
                         with fail_after(read_timeout):
-                            text = await stream.receive()
-                        await self.info_sender.send(MacmonMetrics.from_raw_json(text))
+                            data = await stream.receive_until(
+                                delimiter=b"\n", max_bytes=8 * 1024
+                            )
+                            text = data.decode("utf-8", errors="replace").strip()
+                            metrics = MacmonMetrics.from_raw_json(text)
+                        await self.info_sender.send(metrics)
             except TimeoutError:
                 logger.warning(
                     f"MacMon produced no output for {read_timeout}s, restarting"
                 )
-                self._enable_psutil_memory_fallback(
-                    "MacMon produced no output, falling back to psutil memory "
-                    "monitoring"
-                )
+                self._tg.start_soon(self._monitor_memory_usage, 1)
             except CalledProcessError as e:
                 stderr_msg = "no stderr"
                 stderr_output = cast(bytes | str | None, e.stderr)
@@ -623,12 +596,8 @@ class InfoGatherer:
                 logger.warning(
                     f"MacMon failed with return code {e.returncode}: {stderr_msg}"
                 )
-                self._enable_psutil_memory_fallback(
-                    "MacMon failed, falling back to psutil memory monitoring"
-                )
-            except Exception:
-                logger.opt(exception=True).warning("Error in macmon monitor")
-                self._enable_psutil_memory_fallback(
-                    "MacMon crashed, falling back to psutil memory monitoring"
-                )
-            await anyio.sleep(self.macmon_interval)
+                self._tg.start_soon(self._monitor_memory_usage, 1)
+            except Exception as e:
+                logger.opt(exception=e).warning("Error in macmon monitor")
+                self._tg.start_soon(self._monitor_memory_usage, 1)
+            await anyio.sleep(macmon_interval)

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

@@ -1,4 +1,4 @@
-from collections.abc import Generator
+from collections.abc import Callable, Generator
 from pathlib import Path
 from typing import Any, Literal, Optional
 
@@ -116,6 +116,7 @@ class DistributedImageModel:
         image_path: Path | None = None,
         partial_images: int = 0,
         advanced_params: AdvancedImageParams | None = None,
+        cancel_checker: Callable[[], bool] | None = None,
     ) -> Generator[Image.Image | tuple[Image.Image, int, int], None, None]:
         if (
             advanced_params is not None
@@ -163,6 +164,7 @@ class DistributedImageModel:
             guidance_override=guidance_override,
             negative_prompt=negative_prompt,
             num_sync_steps=num_sync_steps,
+            cancel_checker=cancel_checker,
         ):
             if isinstance(result, tuple):
                 # Partial image: (GeneratedImage, partial_index, total_partials)

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

@@ -3,6 +3,7 @@ import io
 import random
 import tempfile
 import time
+from collections.abc import Callable
 from pathlib import Path
 from typing import Generator, Literal
 
@@ -69,6 +70,7 @@ def warmup_image_generator(model: DistributedImageModel) -> Image.Image | None:
 def generate_image(
     model: DistributedImageModel,
     task: ImageGenerationTaskParams | ImageEditsTaskParams,
+    cancel_checker: Callable[[], bool] | None = None,
 ) -> Generator[ImageGenerationResponse | PartialImageResponse, None, None]:
     """Generate image(s), optionally yielding partial results.
 
@@ -127,6 +129,7 @@ def generate_image(
                 image_path=image_path,
                 partial_images=partial_images,
                 advanced_params=advanced_params,
+                cancel_checker=cancel_checker,
             ):
                 if isinstance(result, tuple):
                     # Partial image: (Image, partial_index, total_partials)

src/exo/worker/engines/image/models/qwen/wrappers.py

@@ -56,10 +56,10 @@ class QwenJointBlockWrapper(JointBlockWrapper[QwenTransformerBlock]):
         attn = self.block.attn
 
         img_mod_params = self.block.img_mod_linear(
-            self.block.img_mod_silu(text_embeddings)  # pyright: ignore[reportUnknownArgumentType]
+            self.block.img_mod_silu(text_embeddings)
         )
         txt_mod_params = self.block.txt_mod_linear(
-            self.block.txt_mod_silu(text_embeddings)  # pyright: ignore[reportUnknownArgumentType]
+            self.block.txt_mod_silu(text_embeddings)
         )
 
         img_mod1, img_mod2 = mx.split(img_mod_params, 2, axis=-1)

src/exo/worker/engines/image/pipeline/runner.py

@@ -1,4 +1,4 @@
-from collections.abc import Iterator
+from collections.abc import Callable, Iterator
 from dataclasses import dataclass
 from math import ceil
 from typing import Any, Optional, final
@@ -100,6 +100,8 @@ class DiffusionRunner:
         self.total_layers = config.total_blocks
 
         self._guidance_override: float | None = None
+        self._cancel_checker: Callable[[], bool] | None = None
+        self._cancelling: bool = False
 
         self._compute_assigned_blocks()
 
@@ -240,6 +242,43 @@ class DiffusionRunner:
     def is_distributed(self) -> bool:
         return self.group is not None
 
+    def _is_sentinel(self, tensor: mx.array) -> bool:
+        return bool(mx.all(mx.isnan(tensor)).item())
+
+    def _check_cancellation(self) -> None:
+        if self._cancelling:
+            return
+        if (
+            self.is_first_stage
+            and self._cancel_checker is not None
+            and self._cancel_checker()
+        ):
+            self._cancelling = True
+
+    def _send(self, data: mx.array, dst: int) -> mx.array:
+        assert self.group is not None
+        if self._cancelling:
+            data = mx.full(data.shape, float("nan"), dtype=data.dtype)
+        return mx.distributed.send(data, dst, group=self.group)
+
+    def _recv_and_check(self, result: mx.array) -> mx.array:
+        mx.eval(result)
+        if self._is_sentinel(result):
+            self._cancelling = True
+        return result
+
+    def _recv(self, shape: tuple[int, ...], dtype: mx.Dtype, src: int) -> mx.array:
+        assert self.group is not None
+        return self._recv_and_check(
+            mx.distributed.recv(shape, dtype, src, group=self.group)
+        )
+
+    def _recv_like(self, template: mx.array, src: int) -> mx.array:
+        assert self.group is not None
+        return self._recv_and_check(
+            mx.distributed.recv_like(template, src, group=self.group)
+        )
+
     def _get_effective_guidance_scale(self) -> float | None:
         if self._guidance_override is not None:
             return self._guidance_override
@@ -313,19 +352,13 @@ class DiffusionRunner:
         assert self.cfg_peer_rank is not None
 
         if is_positive:
-            noise = mx.distributed.send(noise, self.cfg_peer_rank, group=self.group)
+            noise = self._send(noise, self.cfg_peer_rank)
             mx.async_eval(noise)
-            noise_neg = mx.distributed.recv_like(
-                noise, self.cfg_peer_rank, group=self.group
-            )
-            mx.eval(noise_neg)
+            noise_neg = self._recv_like(noise, src=self.cfg_peer_rank)
             noise_pos = noise
         else:
-            noise_pos = mx.distributed.recv_like(
-                noise, self.cfg_peer_rank, group=self.group
-            )
-            mx.eval(noise_pos)
-            noise = mx.distributed.send(noise, self.cfg_peer_rank, group=self.group)
+            noise_pos = self._recv_like(noise, src=self.cfg_peer_rank)
+            noise = self._send(noise, self.cfg_peer_rank)
             mx.async_eval(noise)
             noise_neg = noise
 
@@ -432,6 +465,7 @@ class DiffusionRunner:
         guidance_override: float | None = None,
         negative_prompt: str | None = None,
         num_sync_steps: int = 1,
+        cancel_checker: Callable[[], bool] | None = None,
     ):
         """Primary entry point for image generation.
 
@@ -454,6 +488,8 @@ class DiffusionRunner:
             Final GeneratedImage
         """
         self._guidance_override = guidance_override
+        self._cancel_checker = cancel_checker
+        self._cancelling = False
         latents = self.adapter.create_latents(seed, runtime_config)
         prompt_data = self.adapter.encode_prompt(prompt, negative_prompt)
 
@@ -495,7 +531,7 @@ class DiffusionRunner:
             except StopIteration as e:
                 latents = e.value  # pyright: ignore[reportAny]
 
-        if self.is_last_stage:
+        if self.is_last_stage and not self._cancelling:
             yield self.adapter.decode_latents(latents, runtime_config, seed, prompt)  # pyright: ignore[reportAny]
 
     def _run_diffusion_loop(
@@ -524,7 +560,12 @@ class DiffusionRunner:
             latents=latents,
         )
 
+        t = -1  # default if time_steps is empty; drain condition uses t
         for t in time_steps:
+            self._check_cancellation()
+            if self._cancelling and self.group is None:
+                break
+
             try:
                 latents = self._diffusion_step(
                     t=t,
@@ -542,7 +583,7 @@ class DiffusionRunner:
 
                 mx.eval(latents)
 
-                if t in capture_steps and self.is_last_stage:
+                if t in capture_steps and self.is_last_stage and not self._cancelling:
                     yield (latents, t)
 
             except KeyboardInterrupt:  # noqa: PERF203
@@ -551,6 +592,24 @@ class DiffusionRunner:
                     f"Stopping image generation at step {t + 1}/{len(time_steps)}"
                 ) from None
 
+            if self._cancelling:
+                break
+
+        # Drain pending ring recvs after cancellation during async steps.
+        # The last stage sent patches during the final completed step, but
+        # the first stage will never enter the next step to recv them.
+        if (
+            self._cancelling
+            and self.is_first_stage
+            and not self.is_last_stage
+            and self.group is not None
+            and t >= runtime_config.init_time_step + num_sync_steps
+            and t != runtime_config.num_inference_steps - 1
+        ):
+            patch_latents_drain, _ = self._create_patches(latents, runtime_config)
+            for patch in patch_latents_drain:
+                self._recv_like(patch, src=self.last_pipeline_rank)
+
         ctx.after_loop(latents=latents)  # pyright: ignore[reportAny]
 
         return latents
@@ -777,19 +836,16 @@ class DiffusionRunner:
                     rank=self.rank,
                     category="comms",
                 ):
-                    hidden_states = mx.distributed.recv(
+                    hidden_states = self._recv(
                         (batch_size, num_img_tokens, hidden_dim),
                         dtype,
                         self.prev_pipeline_rank,
-                        group=self.group,
                     )
-                    encoder_hidden_states = mx.distributed.recv(
+                    encoder_hidden_states = self._recv(
                         (batch_size, text_seq_len, hidden_dim),
                         dtype,
                         self.prev_pipeline_rank,
-                        group=self.group,
                     )
-                    mx.eval(hidden_states, encoder_hidden_states)
 
             assert self.joint_block_wrappers is not None
             assert encoder_hidden_states is not None
@@ -825,9 +881,7 @@ class DiffusionRunner:
                     rank=self.rank,
                     category="comms",
                 ):
-                    concatenated = mx.distributed.send(
-                        concatenated, self.next_pipeline_rank, group=self.group
-                    )
+                    concatenated = self._send(concatenated, self.next_pipeline_rank)
                     mx.async_eval(concatenated)
 
         elif self.has_joint_blocks and not self.is_last_stage:
@@ -838,11 +892,9 @@ class DiffusionRunner:
                 rank=self.rank,
                 category="comms",
             ):
-                hidden_states = mx.distributed.send(
-                    hidden_states, self.next_pipeline_rank, group=self.group
-                )
-                encoder_hidden_states = mx.distributed.send(
-                    encoder_hidden_states, self.next_pipeline_rank, group=self.group
+                hidden_states = self._send(hidden_states, self.next_pipeline_rank)
+                encoder_hidden_states = self._send(
+                    encoder_hidden_states, self.next_pipeline_rank
                 )
                 mx.async_eval(hidden_states, encoder_hidden_states)
 
@@ -854,13 +906,11 @@ class DiffusionRunner:
                     rank=self.rank,
                     category="comms",
                 ):
-                    hidden_states = mx.distributed.recv(
+                    hidden_states = self._recv(
                         (batch_size, text_seq_len + num_img_tokens, hidden_dim),
                         dtype,
                         self.prev_pipeline_rank,
-                        group=self.group,
                     )
-                    mx.eval(hidden_states)
 
             assert self.single_block_wrappers is not None
             with trace(
@@ -886,9 +936,7 @@ class DiffusionRunner:
                     rank=self.rank,
                     category="comms",
                 ):
-                    hidden_states = mx.distributed.send(
-                        hidden_states, self.next_pipeline_rank, group=self.group
-                    )
+                    hidden_states = self._send(hidden_states, self.next_pipeline_rank)
                     mx.async_eval(hidden_states)
 
         hidden_states = hidden_states[:, text_seq_len:, ...]
@@ -961,16 +1009,11 @@ class DiffusionRunner:
             )
 
             if not self.is_first_stage:
-                hidden_states = mx.distributed.send(
-                    hidden_states, self.first_pipeline_rank, group=self.group
-                )
+                hidden_states = self._send(hidden_states, self.first_pipeline_rank)
                 mx.async_eval(hidden_states)
 
         elif self.is_first_stage:
-            hidden_states = mx.distributed.recv_like(
-                prev_latents, src=self.last_pipeline_rank, group=self.group
-            )
-            mx.eval(hidden_states)
+            hidden_states = self._recv_like(prev_latents, src=self.last_pipeline_rank)
 
         else:
             hidden_states = prev_latents
@@ -1006,10 +1049,7 @@ class DiffusionRunner:
                     rank=self.rank,
                     category="comms",
                 ):
-                    patch = mx.distributed.recv_like(
-                        patch, src=self.last_pipeline_rank, group=self.group
-                    )
-                    mx.eval(patch)
+                    patch = self._recv_like(patch, src=self.last_pipeline_rank)
 
             results: list[tuple[bool, mx.array]] = []
 
@@ -1066,10 +1106,9 @@ class DiffusionRunner:
                         rank=self.rank,
                         category="comms",
                     ):
-                        patch_latents[patch_idx] = mx.distributed.send(
+                        patch_latents[patch_idx] = self._send(
                             patch_latents[patch_idx],
                             self.first_pipeline_rank,
-                            group=self.group,
                         )
                         mx.async_eval(patch_latents[patch_idx])
 
@@ -1116,13 +1155,11 @@ class DiffusionRunner:
                     rank=self.rank,
                     category="comms",
                 ):
-                    patch = mx.distributed.recv(
+                    patch = self._recv(
                         (batch_size, patch_len, hidden_dim),
                         patch.dtype,
                         self.prev_pipeline_rank,
-                        group=self.group,
                     )
-                    mx.eval(patch)
 
                 if patch_idx == 0:
                     with trace(
@@ -1130,13 +1167,11 @@ class DiffusionRunner:
                         rank=self.rank,
                         category="comms",
                     ):
-                        encoder_hidden_states = mx.distributed.recv(
+                        encoder_hidden_states = self._recv(
                             (batch_size, text_seq_len, hidden_dim),
                             patch.dtype,
                             self.prev_pipeline_rank,
-                            group=self.group,
                         )
-                        mx.eval(encoder_hidden_states)
 
             if self.is_first_stage:
                 patch, encoder_hidden_states = self.adapter.compute_embeddings(
@@ -1175,9 +1210,7 @@ class DiffusionRunner:
                     rank=self.rank,
                     category="comms",
                 ):
-                    patch_concat = mx.distributed.send(
-                        patch_concat, self.next_pipeline_rank, group=self.group
-                    )
+                    patch_concat = self._send(patch_concat, self.next_pipeline_rank)
                     mx.async_eval(patch_concat)
 
         elif self.has_joint_blocks and not self.is_last_stage:
@@ -1187,9 +1220,7 @@ class DiffusionRunner:
                 rank=self.rank,
                 category="comms",
             ):
-                patch = mx.distributed.send(
-                    patch, self.next_pipeline_rank, group=self.group
-                )
+                patch = self._send(patch, self.next_pipeline_rank)
                 mx.async_eval(patch)
 
             if patch_idx == 0:
@@ -1199,8 +1230,8 @@ class DiffusionRunner:
                     rank=self.rank,
                     category="comms",
                 ):
-                    encoder_hidden_states = mx.distributed.send(
-                        encoder_hidden_states, self.next_pipeline_rank, group=self.group
+                    encoder_hidden_states = self._send(
+                        encoder_hidden_states, self.next_pipeline_rank
                     )
                     mx.async_eval(encoder_hidden_states)
 
@@ -1213,13 +1244,11 @@ class DiffusionRunner:
                     rank=self.rank,
                     category="comms",
                 ):
-                    patch = mx.distributed.recv(
+                    patch = self._recv(
                         (batch_size, text_seq_len + patch_len, hidden_dim),
                         patch.dtype,
                         self.prev_pipeline_rank,
-                        group=self.group,
                     )
-                    mx.eval(patch)
 
             assert self.single_block_wrappers is not None
             with trace(
@@ -1245,9 +1274,7 @@ class DiffusionRunner:
                     rank=self.rank,
                     category="comms",
                 ):
-                    patch = mx.distributed.send(
-                        patch, self.next_pipeline_rank, group=self.group
-                    )
+                    patch = self._send(patch, self.next_pipeline_rank)
                     mx.async_eval(patch)
 
         noise: mx.array | None = None

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

@@ -57,8 +57,8 @@ from mlx_lm.models.step3p5 import Model as Step35Model
 from mlx_lm.models.step3p5 import Step3p5MLP as Step35MLP
 from mlx_lm.models.step3p5 import Step3p5Model as Step35InnerModel
 
-from exo.shared.logging import logger
 from exo.shared.types.worker.shards import PipelineShardMetadata
+from exo.worker.runner.bootstrap import logger
 
 if TYPE_CHECKING:
     from mlx_lm.models.cache import Cache
@@ -480,7 +480,7 @@ def patch_tensor_model[T](model: T) -> T:
             last = cache[-1]  # pyright: ignore[reportAny]
             dep_cache = last[0] if hasattr(last, "caches") else last  # pyright: ignore[reportAny]
             if hasattr(dep_cache, "keys"):  # type: ignore
-                dep_cache.keys = mx.depends(dep_cache.keys, logits)  # pyright: ignore[reportAny,reportUnknownMemberType]
+                dep_cache.keys = mx.depends(dep_cache.keys, logits)  # pyright: ignore[reportAny]
 
         return logits
 

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

@@ -15,7 +15,28 @@ USER_TOKEN = "<\uff5cUser\uff5c>"
 ASSISTANT_TOKEN = "<\uff5cAssistant\uff5c>"
 TOOL_CALLS_START = f"<{DSML_TOKEN}function_calls>"
 TOOL_CALLS_END = f"</{DSML_TOKEN}function_calls>"
-encode_messages = deepseek_v32.encode_messages
+_ORPHAN_THINK_END = ASSISTANT_TOKEN + THINKING_END
+_FIXED_THINK_BLOCK = ASSISTANT_TOKEN + THINKING_START + "\n" + THINKING_END
+
+
+def encode_messages(
+    messages: list[dict[str, Any]],
+    thinking_mode: str = "thinking",
+    context: list[dict[str, Any]] | None = None,
+    drop_thinking: bool = True,
+    add_default_bos_token: bool = True,
+    tools: Any = None,  # pyright: ignore[reportAny]
+) -> str:
+    prompt: str = deepseek_v32.encode_messages(
+        messages,
+        thinking_mode=thinking_mode,
+        context=context,
+        drop_thinking=drop_thinking,
+        add_default_bos_token=add_default_bos_token,
+        tools=tools,
+    )
+    return prompt.replace(_ORPHAN_THINK_END, _FIXED_THINK_BLOCK)
+
 
 _INVOKE_PATTERN = re.compile(
     rf"<{re.escape(DSML_TOKEN)}invoke\s+name=\"([^\"]+)\">"

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

@@ -6,6 +6,9 @@ import mlx.core as mx
 from mlx_lm.generate import (
     BatchGenerator as MlxBatchGenerator,
 )
+from mlx_lm.generate import (
+    generation_stream,
+)
 from mlx_lm.models.cache import RotatingKVCache
 from mlx_lm.sample_utils import make_logits_processors, make_sampler
 from mlx_lm.tokenizer_utils import StreamingDetokenizer, TokenizerWrapper
@@ -63,6 +66,7 @@ class _EngineTask:
     potential_stop_sequence_text: str = ""
     completion_tokens: int = 0
     generation_start_time: float = 0.0
+    generation_time_at_start: float = 0.0
     in_thinking: bool = False
     reasoning_tokens: int = 0
     prefill_tps: float = 0.0
@@ -75,22 +79,23 @@ class ExoBatchGenerator:
     group: mx.distributed.Group | None
     kv_prefix_cache: KVPrefixCache | None
 
-    _exo_gen: MlxBatchGenerator = field(init=False)
+    _mlx_gen: MlxBatchGenerator = field(init=False)
     _active_tasks: dict[int, _EngineTask] = field(default_factory=dict, init=False)
 
     def __post_init__(self) -> None:
-        self._exo_gen = MlxBatchGenerator(
+        self._mlx_gen = MlxBatchGenerator(
             model=self.model,
             stop_tokens=set(eos_ids_from_tokenizer(self.tokenizer)),
             prefill_step_size=4096,
         )
+        self._mlx_gen._needs_topk = False  # pyright: ignore[reportAttributeAccessIssue]
 
     @property
     def has_work(self) -> bool:
         return (
             bool(self._active_tasks)
-            or bool(self._exo_gen.unprocessed_prompts)
-            or self._exo_gen.active_batch is not None
+            or bool(self._mlx_gen.unprocessed_prompts)
+            or self._mlx_gen.active_batch is not None
         )
 
     def submit(
@@ -188,7 +193,7 @@ class ExoBatchGenerator:
 
         max_tokens = task_params.max_output_tokens or MAX_TOKENS
 
-        uids = self._exo_gen.insert(
+        uids = self._mlx_gen.insert(
             prompts=[last_tokens.tolist()],
             max_tokens=[max_tokens],
             caches=[list(cache)],
@@ -211,6 +216,7 @@ class ExoBatchGenerator:
             on_generation_token=on_generation_token,
             generation_start_time=time.perf_counter(),
             prefill_tps=_prefill_tps,
+            generation_time_at_start=self._mlx_gen._stats.generation_time,
         )
 
         return uid
@@ -219,7 +225,12 @@ class ExoBatchGenerator:
         if not self.has_work:
             return []
 
-        responses = self._exo_gen.next()
+        self._mlx_gen._needs_topk = any(  # pyright: ignore[reportAttributeAccessIssue]
+            t.task_params.logprobs for t in self._active_tasks.values()
+        )
+        _step_tic = time.perf_counter()
+        responses = self._mlx_gen.next()
+        _next_elapsed = time.perf_counter() - _step_tic
 
         results: list[tuple[int, GenerationResponse]] = []
 
@@ -277,28 +288,31 @@ class ExoBatchGenerator:
             logprob: float | None = None
             top_logprobs: list[TopLogprobItem] | None = None
             if task_params.logprobs:
-                logprob, top_logprobs = extract_top_logprobs(
-                    logprobs=response.logprobs,
-                    tokenizer=self.tokenizer,
-                    top_logprobs=task_params.top_logprobs or DEFAULT_TOP_LOGPROBS,
-                    selected_token=response.token,
-                )
+                with mx.stream(generation_stream):
+                    logprob, top_logprobs = extract_top_logprobs(
+                        logprobs=response.logprobs,
+                        tokenizer=self.tokenizer,
+                        top_logprobs=task_params.top_logprobs or DEFAULT_TOP_LOGPROBS,
+                        selected_token=response.token,
+                        precomputed_indices=getattr(response, "_topk_indices", None),
+                        precomputed_values=getattr(response, "_topk_values", None),
+                        precomputed_selected=getattr(
+                            response, "_selected_logprob", None
+                        ),
+                    )
 
             stats: GenerationStats | None = None
             usage: Usage | None = None
             if is_done:
-                try:
-                    mlx_stats = self._exo_gen.stats()
-                    generation_tps = mlx_stats.generation_tps
-                except ZeroDivisionError:
-                    generation_elapsed = (
-                        time.perf_counter() - state.generation_start_time
-                    )
-                    generation_tps = (
-                        state.completion_tokens / generation_elapsed
-                        if generation_elapsed > 0
-                        else 0.0
-                    )
+                gen_time_delta = (
+                    self._mlx_gen._stats.generation_time
+                    - state.generation_time_at_start
+                )
+                generation_tps = (
+                    state.completion_tokens / gen_time_delta
+                    if gen_time_delta > 0
+                    else 0.0
+                )
 
                 stats = GenerationStats(
                     prompt_tps=state.prefill_tps,
@@ -345,15 +359,22 @@ class ExoBatchGenerator:
                     -max_stop_len:
                 ]
 
+        _step_elapsed = time.perf_counter() - _step_tic
+        _overhead = _step_elapsed - _next_elapsed
+        if self._mlx_gen._next_count % 64 == 0 and responses:
+            logger.debug(
+                f"step overhead: {_overhead * 1000:.2f}ms (next={_next_elapsed * 1000:.2f}ms total={_step_elapsed * 1000:.2f}ms)"
+            )
+
         return results
 
     def cancel(self, uids: list[int]) -> None:
-        self._exo_gen.remove(uids)
+        self._mlx_gen.remove(uids)
         for uid in uids:
             self._active_tasks.pop(uid, None)
 
     def close(self) -> None:
-        self._exo_gen.close()
+        self._mlx_gen.close()
 
     def _save_prefix_cache(
         self,
@@ -372,9 +393,8 @@ class ExoBatchGenerator:
                 if len(all_prompt_tokens) > 0
                 else 0.0
             )
-            if (
-                matched_index is not None
-                and hit_ratio >= _MIN_PREFIX_HIT_RATIO_TO_UPDATE
+            if matched_index is not None and (
+                prefix_hit_length > 1000 or hit_ratio >= _MIN_PREFIX_HIT_RATIO_TO_UPDATE
             ):
                 self.kv_prefix_cache.update_kv_cache(
                     matched_index,

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

@@ -179,7 +179,8 @@ def pipeline_parallel_prefill(
         flush_prefill_sends()
 
     assert _prompt_cache is not None
-    mx.eval([c.state for c in _prompt_cache])  # type: ignore
+    with mx.stream(generation_stream):
+        mx.eval([c.state for c in _prompt_cache])  # type: ignore
 
     # Final callback matching generate_step
     prompt_progress_callback(total, total)
@@ -312,52 +313,46 @@ def warmup_inference(
     model_id: ModelId,
 ) -> int:
     logger.info(f"warming up inference for instance: {model_id}")
-    t = time.monotonic()
 
     content = "Prompt to warm up the inference engine. Repeat this."
 
+    warmup_task_params = TextGenerationTaskParams(
+        model=model_id,
+        input=[InputMessage(role="user", content=content)],
+        max_output_tokens=50,
+        temperature=0.0,
+    )
+
     warmup_prompt = apply_chat_template(
         tokenizer=tokenizer,
-        task_params=TextGenerationTaskParams(
-            model=ModelId(""),
-            input=[InputMessage(role="user", content=content)],
-        ),
+        task_params=warmup_task_params,
     )
 
     tokens_generated = 0
 
-    cache = make_kv_cache(
-        model=model,
-    )
-
-    # Use a default sampler for warmup
-    sampler = make_sampler(temp=0.0)
-
     mx_barrier(group)
 
     logger.info("Generating warmup tokens")
-    for _r in stream_generate(
+
+    t = time.monotonic()
+
+    for _r in mlx_generate(
         model=model,
         tokenizer=tokenizer,
+        task=warmup_task_params,
         prompt=warmup_prompt,
-        max_tokens=50,
-        sampler=sampler,
-        prompt_cache=cache,
-        prefill_step_size=2048,
-        kv_group_size=KV_GROUP_SIZE,
-        kv_bits=KV_BITS,
+        kv_prefix_cache=None,
+        group=group,
     ):
-        logger.info("Generated warmup token: " + str(_r.text))
         tokens_generated += 1
 
-    logger.info("Generated ALL warmup tokens")
+    check_for_cancel_every = min(
+        math.ceil(tokens_generated / min(time.monotonic() - t, 0.001)), 100
+    )
 
     mx_barrier(group)
 
     logger.info(f"warmed up by generating {tokens_generated} tokens")
-    check_for_cancel_every = min(
-        math.ceil(tokens_generated / min(time.monotonic() - t, 0.001)), 100
-    )
     if group is not None:
         check_for_cancel_every = int(
             mx.max(
@@ -398,52 +393,44 @@ def extract_top_logprobs(
     tokenizer: TokenizerWrapper,
     top_logprobs: int,
     selected_token: int,
+    precomputed_indices: list[int] | None = None,
+    precomputed_values: list[float] | None = None,
+    precomputed_selected: float | None = None,
 ) -> tuple[float, list[TopLogprobItem]]:
-    """Extract the selected token's logprob and top alternative tokens.
-
-    Args:
-        logprobs: Full vocabulary logprobs array from MLX
-        tokenizer: Tokenizer for decoding token IDs to strings
-        top_logprobs: Number of top alternatives to return
-        selected_token: The token ID that was actually sampled
-
-    Returns:
-        Tuple of (selected_token_logprob, list of TopLogprobItem for top alternatives)
-    """
-    # Get the logprob of the selected token
-    selected_logprob = float(logprobs[selected_token].item())
-
-    # Get top indices (most probable tokens)
-    # mx.argpartition gives indices that would partition the array
-    # We negate logprobs since argpartition finds smallest, and we want largest
-    top_logprobs = min(top_logprobs, logprobs.shape[0])  # Don't exceed vocab size
-    top_indices = mx.argpartition(-logprobs, top_logprobs)[:top_logprobs]
-
-    # Get the actual logprob values for these indices
-    top_values = logprobs[top_indices]
-
-    # Sort by logprob (descending) for consistent ordering
-    sort_order = mx.argsort(-top_values)
-    top_indices = top_indices[sort_order]
-    top_values = top_values[sort_order]
+    if (
+        precomputed_indices is not None
+        and precomputed_values is not None
+        and precomputed_selected is not None
+    ):
+        top_indices_list: list[int] = precomputed_indices[:top_logprobs]
+        top_values_list: list[float] = precomputed_values[:top_logprobs]
+        selected_logprob = precomputed_selected
+    else:
+        selected_logprob_arr = logprobs[selected_token]
+        top_logprobs = min(top_logprobs, logprobs.shape[0] - 1)
+        top_indices = mx.argpartition(-logprobs, top_logprobs)[:top_logprobs]
+        top_values = logprobs[top_indices]
+        sort_order = mx.argsort(-top_values)
+        top_indices = top_indices[sort_order]
+        top_values = top_values[sort_order]
+        mx.eval(selected_logprob_arr, top_indices, top_values)
+        selected_logprob = float(selected_logprob_arr.item())
+        top_indices_list = top_indices.tolist()  # type: ignore
+        top_values_list = top_values.tolist()  # type: ignore
 
     # Convert to list of TopLogprobItem
     top_logprob_items: list[TopLogprobItem] = []
-    for i in range(top_logprobs):
-        token_id = int(top_indices[i].item())
-        token_logprob = float(top_values[i].item())
+    for token_id, token_logprob in zip(top_indices_list, top_values_list, strict=True):
         if math.isnan(token_logprob):
             continue
 
         # Decode token ID to string
         token_str = tokenizer.decode([token_id])
-        # Get byte representation
-        token_bytes = list(token_str.encode("utf-8"))
         top_logprob_items.append(
             TopLogprobItem(
                 token=token_str,
                 logprob=token_logprob,
-                bytes=token_bytes,
+                bytes=list(token_str.encode("utf-8")),
             )
         )
 
@@ -624,12 +611,13 @@ def mlx_generate(
         logprob: float | None = None
         top_logprobs: list[TopLogprobItem] | None = None
         if task.logprobs:
-            logprob, top_logprobs = extract_top_logprobs(
-                logprobs=out.logprobs,
-                tokenizer=tokenizer,
-                top_logprobs=task.top_logprobs or DEFAULT_TOP_LOGPROBS,
-                selected_token=out.token,
-            )
+            with mx.stream(generation_stream):
+                logprob, top_logprobs = extract_top_logprobs(
+                    logprobs=out.logprobs,
+                    tokenizer=tokenizer,
+                    top_logprobs=task.top_logprobs or DEFAULT_TOP_LOGPROBS,
+                    selected_token=out.token,
+                )
 
         if is_done:
             # Log generation stats
@@ -657,9 +645,9 @@ def mlx_generate(
                     if len(all_prompt_tokens) > 0
                     else 0.0
                 )
-                if (
-                    matched_index is not None
-                    and hit_ratio >= _MIN_PREFIX_HIT_RATIO_TO_UPDATE
+                if matched_index is not None and (
+                    prefix_hit_length > 1000
+                    or hit_ratio >= _MIN_PREFIX_HIT_RATIO_TO_UPDATE
                 ):
                     kv_prefix_cache.update_kv_cache(
                         matched_index,

src/exo/worker/engines/mlx/patches/__init__.py

@@ -0,0 +1,14 @@
+from exo.worker.engines.mlx.patches.opt_batch_gen import apply_batch_gen_patch
+from exo.worker.engines.mlx.patches.standard_yarn_rope import patch_yarn_rope
+
+_applied = False
+
+
+def apply_mlx_patches() -> None:
+    global _applied
+    if _applied:
+        return
+    _applied = True
+    patch_yarn_rope()
+    # patch_gdn_softplus()
+    apply_batch_gen_patch()

src/exo/worker/engines/mlx/patches/opt_batch_gen.py

@@ -0,0 +1,173 @@
+import time
+from typing import Any, cast
+
+import mlx.core as mx
+from mlx_lm.generate import BatchGenerator, generation_stream
+
+_PRECOMPUTE_TOP_K = 20
+
+_original_public_next = BatchGenerator.next
+
+_pending_topk_idx: mx.array | None = None
+_pending_topk_val: mx.array | None = None
+_pending_selected_lps: mx.array | None = None
+
+
+def _fast_next(self: BatchGenerator) -> list[BatchGenerator.Response]:
+    tic = time.perf_counter()
+    batch = self.active_batch
+    assert batch is not None
+    batch_size = len(batch)
+
+    prev_tokens = batch.y
+    prev_logprobs = batch.logprobs
+
+    has_processors = any(p for ps in batch.logits_processors for p in ps)
+    if has_processors:
+        for i, toks in enumerate(batch.tokens):
+            batch.tokens[i] = mx.concatenate([toks, prev_tokens[i : i + 1]])
+
+    logits = self.model(prev_tokens[:, None], cache=batch.cache)
+    logits = logits[:, -1, :]
+
+    if has_processors:
+        processed_logits: list[mx.array] = []
+        for e in range(batch_size):
+            sample_logits: mx.array = logits[e : e + 1]
+            for processor in batch.logits_processors[e]:
+                sample_logits = processor(batch.tokens[e], sample_logits)
+            processed_logits.append(sample_logits)
+        logits = mx.concatenate(processed_logits, axis=0)
+
+    logprobs = logits - mx.logsumexp(logits, axis=-1, keepdims=True)
+
+    if (
+        batch_size == 1
+        or any(batch.samplers)
+        and all(s is batch.samplers[0] for s in batch.samplers)
+    ):
+        sampler = batch.samplers[0] or self.sampler
+        batch.y = sampler(logprobs)
+    elif any(batch.samplers):
+        all_samples: list[mx.array] = []
+        for e in range(batch_size):
+            s = batch.samplers[e] or self.sampler
+            all_samples.append(s(logprobs[e : e + 1]))
+        batch.y = mx.concatenate(all_samples, axis=0)
+    else:
+        batch.y = self.sampler(logprobs)
+    batch.logprobs = list(logprobs)
+
+    global _pending_topk_idx, _pending_topk_val, _pending_selected_lps
+
+    emit_topk_indices: list[list[int]] = (
+        cast(list[list[int]], _pending_topk_idx.tolist())
+        if _pending_topk_idx is not None
+        else []
+    )
+    emit_topk_values: list[list[float]] = (
+        cast(list[list[float]], _pending_topk_val.tolist())
+        if _pending_topk_val is not None
+        else []
+    )
+    emit_selected_lps: list[float] = (
+        cast(list[float], _pending_selected_lps.tolist())
+        if _pending_selected_lps is not None
+        else []
+    )
+
+    needs_topk: bool = getattr(self, "_needs_topk", False)
+    if needs_topk:
+        k = min(_PRECOMPUTE_TOP_K, logprobs.shape[1])
+        _pending_topk_idx = mx.argpartition(-logprobs, k, axis=1)[:, :k]
+        _pending_topk_val = mx.take_along_axis(logprobs, _pending_topk_idx, axis=1)
+        sort_order = mx.argsort(-_pending_topk_val, axis=1)
+        _pending_topk_idx = mx.take_along_axis(_pending_topk_idx, sort_order, axis=1)
+        _pending_topk_val = mx.take_along_axis(_pending_topk_val, sort_order, axis=1)
+        _pending_selected_lps = logprobs[mx.arange(batch_size), batch.y]
+        mx.async_eval(
+            batch.y,
+            *batch.logprobs,
+            *batch.tokens,
+            _pending_topk_idx,
+            _pending_topk_val,
+            _pending_selected_lps,
+        )
+    else:
+        _pending_topk_idx = None
+        _pending_topk_val = None
+        _pending_selected_lps = None
+        mx.async_eval(batch.y, *batch.logprobs, *batch.tokens)
+
+    prev_token_list: list[int] = cast(list[int], prev_tokens.tolist())
+
+    toc = time.perf_counter()
+    self._stats.generation_time += toc - tic
+
+    keep_idx: list[int] = []
+    end_idx: list[int] = []
+    responses: list[Any] = []
+    stop_tokens = self.stop_tokens
+
+    for e in range(batch_size):
+        t = prev_token_list[e]
+        uid = batch.uids[e]
+        num_tok = batch.num_tokens[e] + 1
+        batch.num_tokens[e] = num_tok
+
+        if t in stop_tokens:
+            finish_reason = "stop"
+            end_idx.append(e)
+        elif num_tok >= batch.max_tokens[e]:
+            finish_reason = "length"
+            end_idx.append(e)
+        else:
+            finish_reason = None
+            keep_idx.append(e)
+
+        cache = None
+        if finish_reason is not None:
+            cache = batch.extract_cache(e)
+        response = self.Response(uid, t, prev_logprobs[e], finish_reason, cache)
+        if emit_topk_indices and e < len(emit_topk_indices):
+            response._topk_indices = emit_topk_indices[e]  # pyright: ignore[reportAttributeAccessIssue]
+            response._topk_values = emit_topk_values[e]  # pyright: ignore[reportAttributeAccessIssue]
+            response._selected_logprob = emit_selected_lps[e]  # pyright: ignore[reportAttributeAccessIssue]
+        responses.append(response)
+
+    if end_idx:
+        if keep_idx:
+            batch.filter(keep_idx)
+            if (
+                _pending_topk_idx is not None
+                and _pending_topk_val is not None
+                and _pending_selected_lps is not None
+            ):
+                ki = mx.array(keep_idx)
+                _pending_topk_idx = _pending_topk_idx[ki]
+                _pending_topk_val = _pending_topk_val[ki]
+                _pending_selected_lps = _pending_selected_lps[ki]
+        else:
+            self.active_batch = None
+            _pending_topk_idx = None
+            _pending_topk_val = None
+            _pending_selected_lps = None
+
+    self._next_count += 1
+    if self._next_count % 512 == 0:
+        mx.clear_cache()
+    self._stats.generation_tokens += len(responses)
+    return responses
+
+
+def _patched_public_next(self: BatchGenerator) -> list[BatchGenerator.Response]:
+    batch = self.active_batch
+    # Only do decode with fast_next
+    if batch is not None and not self.unprocessed_prompts:
+        with mx.stream(generation_stream):
+            return _fast_next(self)
+    return _original_public_next(self)
+
+
+def apply_batch_gen_patch() -> None:
+    BatchGenerator.next = _patched_public_next

src/exo/worker/engines/mlx/patches/standard_yarn_rope.py

@@ -0,0 +1,118 @@
+import math
+
+import mlx.core as mx
+from mlx_lm.models import rope_utils
+
+_original_YarnRoPE_init = rope_utils.YarnRoPE.__init__  # noqa: N816
+_original_initialize_rope = rope_utils.initialize_rope
+
+
+def _patched_yarn_init(
+    self: rope_utils.YarnRoPE,
+    dims: int,
+    traditional: bool = False,
+    max_position_embeddings: int = 2048,
+    base: float = 10000,
+    scaling_factor: float = 1.0,
+    original_max_position_embeddings: int = 4096,
+    beta_fast: float = 32,
+    beta_slow: float = 1,
+    mscale: float = 1,
+    mscale_all_dim: float = 0,
+    truncate: bool = True,
+) -> None:
+    """Patch mlx_lm's YarnRoPE to match vLLM's inverse-frequency blending formula for compatability."""
+
+    super(rope_utils.YarnRoPE, self).__init__()
+
+    def yarn_find_correction_dim(num_rotations: float) -> float:
+        return (
+            dims
+            * math.log(original_max_position_embeddings / (num_rotations * 2 * math.pi))
+        ) / (2 * math.log(base))
+
+    def yarn_find_correction_range() -> tuple[float, float]:
+        low: float = yarn_find_correction_dim(beta_fast)
+        high: float = yarn_find_correction_dim(beta_slow)
+        if truncate:
+            low = math.floor(low)
+            high = math.ceil(high)
+        return max(low, 0), min(high, dims - 1)
+
+    def yarn_get_mscale(scale: float = 1, ms: float = 1) -> float:
+        if scale <= 1:
+            return 1.0
+        return 0.1 * ms * math.log(scale) + 1.0
+
+    def yarn_linear_ramp_mask(min_val: float, max_val: float, dim: int) -> mx.array:
+        if min_val == max_val:
+            max_val += 0.001
+        linear_func = (mx.arange(dim, dtype=mx.float32) - min_val) / (max_val - min_val)
+        return mx.clip(linear_func, 0, 1)
+
+    self.mscale = yarn_get_mscale(scaling_factor, mscale) / yarn_get_mscale(
+        scaling_factor, mscale_all_dim
+    )
+    pos_freqs = base ** (mx.arange(0, dims, 2, dtype=mx.float32) / dims)
+    inv_freq_extrapolation = 1.0 / pos_freqs
+    inv_freq_interpolation = 1.0 / (scaling_factor * pos_freqs)
+    low, high = yarn_find_correction_range()
+    inv_freq_mask = 1.0 - yarn_linear_ramp_mask(low, high, dims // 2)
+    inv_freq = (
+        inv_freq_interpolation * (1 - inv_freq_mask)
+        + inv_freq_extrapolation * inv_freq_mask
+    )
+    self._freqs = 1.0 / inv_freq
+    self.dims = dims
+    self.traditional = traditional
+
+
+def _patched_initialize_rope(
+    dims: int,
+    base: float,
+    traditional: bool,
+    scaling_config: dict[str, str | int | float | bool] | None = None,
+    max_position_embeddings: int | None = None,
+) -> object:
+    rope_type = "default"
+    if scaling_config is not None:
+        rope_type = str(
+            scaling_config.get("type") or scaling_config.get("rope_type", "default")
+        )
+
+    # All the yarn rope types supported in mlx lm
+    if rope_type in ("yarn", "deepseek_yarn"):
+        assert scaling_config is not None
+        cfg = scaling_config
+
+        def _float(key: str, default: float) -> float:
+            v = cfg.get(key)
+            return float(v) if v is not None else default
+
+        def _int(key: str, default: int) -> int:
+            v = cfg.get(key)
+            return int(v) if v is not None else default
+
+        return rope_utils.YarnRoPE(
+            dims=dims,
+            max_position_embeddings=max_position_embeddings or 2048,
+            traditional=traditional,
+            scaling_factor=_float("factor", 1.0),
+            base=base,
+            original_max_position_embeddings=_int(
+                "original_max_position_embeddings", 4096
+            ),
+            beta_fast=_float("beta_fast", 32),
+            beta_slow=_float("beta_slow", 1),
+            mscale=_float("mscale", 1),
+            mscale_all_dim=_float("mscale_all_dim", 0),
+        )
+
+    return _original_initialize_rope(
+        dims, base, traditional, scaling_config, max_position_embeddings
+    )
+
+
+def patch_yarn_rope() -> None:
+    rope_utils.YarnRoPE.__init__ = _patched_yarn_init
+    rope_utils.initialize_rope = _patched_initialize_rope

src/exo/worker/engines/mlx/tests/test_extract_top_logprobs.py

@@ -0,0 +1,290 @@
+# type: ignore
+import math
+from unittest.mock import MagicMock
+
+import mlx.core as mx
+import mlx.nn as nn
+import pytest
+from mlx_lm.generate import BatchGenerator
+
+from exo.worker.engines.mlx.generator.generate import extract_top_logprobs
+from exo.worker.engines.mlx.patches.opt_batch_gen import (
+    _PRECOMPUTE_TOP_K,
+    apply_batch_gen_patch,
+)
+
+
+def _mock_tokenizer() -> MagicMock:
+    tok = MagicMock()
+    tok.decode = lambda ids: f"tok_{ids[0]}"
+    return tok
+
+
+def _make_logprobs(values: list[float]) -> mx.array:
+    arr = mx.array(values, dtype=mx.float32)
+    mx.eval(arr)
+    return arr
+
+
+class TestExtractTopLogprobsFallback:
+    def test_returns_correct_selected_logprob(self) -> None:
+        lp = _make_logprobs([-1.0, -2.0, -0.5, -3.0, -4.0])
+        selected, _ = extract_top_logprobs(
+            lp, _mock_tokenizer(), top_logprobs=3, selected_token=2
+        )
+        assert selected == pytest.approx(-0.5)
+
+    def test_returns_top_k_sorted_descending(self) -> None:
+        lp = _make_logprobs([-1.0, -2.0, -0.5, -3.0, -4.0])
+        _, items = extract_top_logprobs(
+            lp, _mock_tokenizer(), top_logprobs=3, selected_token=0
+        )
+        logprob_values = [item.logprob for item in items]
+        assert logprob_values == sorted(logprob_values, reverse=True)
+        assert len(items) == 3
+
+    def test_top_tokens_are_most_probable(self) -> None:
+        lp = _make_logprobs([-5.0, -1.0, -3.0, -0.1, -2.0])
+        _, items = extract_top_logprobs(
+            lp, _mock_tokenizer(), top_logprobs=2, selected_token=0
+        )
+        token_ids = [int(item.token.split("_")[1]) for item in items]
+        assert 3 in token_ids
+        assert 1 in token_ids
+
+    def test_top_logprobs_clamped_to_vocab_size(self) -> None:
+        lp = _make_logprobs([-1.0, -2.0, -3.0, -4.0, -5.0])
+        _, items = extract_top_logprobs(
+            lp, _mock_tokenizer(), top_logprobs=10, selected_token=0
+        )
+        assert len(items) == 4
+
+    def test_nan_logprobs_filtered(self) -> None:
+        lp = _make_logprobs([-1.0, float("nan"), -0.5])
+        _, items = extract_top_logprobs(
+            lp, _mock_tokenizer(), top_logprobs=3, selected_token=0
+        )
+        for item in items:
+            assert not math.isnan(item.logprob)
+
+    def test_token_bytes_correct(self) -> None:
+        tok = MagicMock()
+        tok.decode = lambda ids: "hello"
+        lp = _make_logprobs([-1.0, -2.0])
+        _, items = extract_top_logprobs(lp, tok, top_logprobs=2, selected_token=0)
+        assert items[0].bytes == list("hello".encode("utf-8"))
+
+
+class TestExtractTopLogprobsPrecomputed:
+    def test_uses_precomputed_data(self) -> None:
+        lp = _make_logprobs([-99.0])
+        selected, items = extract_top_logprobs(
+            lp,
+            _mock_tokenizer(),
+            top_logprobs=2,
+            selected_token=0,
+            precomputed_indices=[3, 1, 0],
+            precomputed_values=[-0.1, -1.0, -5.0],
+            precomputed_selected=-0.1,
+        )
+        assert selected == pytest.approx(-0.1)
+        assert len(items) == 2
+        assert items[0].token == "tok_3"
+        assert items[0].logprob == pytest.approx(-0.1)
+        assert items[1].token == "tok_1"
+        assert items[1].logprob == pytest.approx(-1.0)
+
+    def test_slices_precomputed_to_requested_k(self) -> None:
+        lp = _make_logprobs([-99.0])
+        _, items = extract_top_logprobs(
+            lp,
+            _mock_tokenizer(),
+            top_logprobs=1,
+            selected_token=0,
+            precomputed_indices=[3, 1, 0, 2, 4],
+            precomputed_values=[-0.1, -1.0, -2.0, -3.0, -4.0],
+            precomputed_selected=-0.1,
+        )
+        assert len(items) == 1
+        assert items[0].token == "tok_3"
+
+    def test_falls_back_when_precomputed_partial(self) -> None:
+        lp = _make_logprobs([-1.0, -2.0, -0.5])
+        selected, items = extract_top_logprobs(
+            lp,
+            _mock_tokenizer(),
+            top_logprobs=2,
+            selected_token=2,
+            precomputed_indices=[0, 2],
+            precomputed_values=None,
+            precomputed_selected=None,
+        )
+        assert selected == pytest.approx(-0.5)
+        assert len(items) == 2
+
+    def test_precomputed_matches_fallback(self) -> None:
+        lp = _make_logprobs([-1.0, -0.3, -2.5, -0.1, -4.0, -0.8, -3.0, -1.5])
+        tok = _mock_tokenizer()
+
+        selected_fb, items_fb = extract_top_logprobs(
+            lp, tok, top_logprobs=5, selected_token=1
+        )
+
+        pre_indices = [item.token.split("_")[1] for item in items_fb]
+        pre_indices_int = [int(x) for x in pre_indices]
+        pre_values = [item.logprob for item in items_fb]
+
+        selected_pc, items_pc = extract_top_logprobs(
+            lp,
+            tok,
+            top_logprobs=5,
+            selected_token=1,
+            precomputed_indices=pre_indices_int,
+            precomputed_values=pre_values,
+            precomputed_selected=selected_fb,
+        )
+
+        assert selected_pc == pytest.approx(selected_fb)
+        assert len(items_pc) == len(items_fb)
+        for a, b in zip(items_pc, items_fb, strict=True):
+            assert a.token == b.token
+            assert a.logprob == pytest.approx(b.logprob)
+
+
+def _tiny_model() -> nn.Module:
+    from mlx_lm.models.llama import Model, ModelArgs
+
+    mx.random.seed(42)
+    args = ModelArgs(
+        model_type="llama",
+        hidden_size=64,
+        num_hidden_layers=2,
+        intermediate_size=128,
+        num_attention_heads=2,
+        num_key_value_heads=1,
+        rms_norm_eps=1e-6,
+        vocab_size=256,
+        rope_theta=10000.0,
+        tie_word_embeddings=True,
+    )
+    model = Model(args)
+    mx.eval(model.parameters())
+    return model
+
+
+@pytest.mark.slow
+class TestBatchedTopKPrecompute:
+    @pytest.fixture(autouse=True)
+    def _reset_globals(self) -> None:
+        import exo.worker.engines.mlx.patches.opt_batch_gen as _mod
+
+        _mod._pending_topk_idx = None
+        _mod._pending_topk_val = None
+        _mod._pending_selected_lps = None
+
+    def _run_generator(
+        self, model: nn.Module, prompts: list[list[int]], steps: int, needs_topk: bool
+    ) -> list[list[BatchGenerator.Response]]:
+        apply_batch_gen_patch()
+        gen = BatchGenerator(model=model, stop_tokens=set(), prefill_step_size=512)
+        gen._needs_topk = needs_topk
+        gen.insert(prompts)
+        all_responses: list[list[BatchGenerator.Response]] = []
+        for _ in range(steps + len(prompts)):
+            responses = gen.next()
+            if responses:
+                all_responses.append(responses)
+            if gen.active_batch is None and not gen.unprocessed_prompts:
+                break
+        gen.close()
+        return all_responses
+
+    def test_precomputed_topk_attached_to_responses(self) -> None:
+        model = _tiny_model()
+        steps = self._run_generator(model, [[1, 2, 3]], 5, needs_topk=True)
+        found_precomputed = False
+        for step_responses in steps:
+            for resp in step_responses:
+                if hasattr(resp, "_topk_indices"):
+                    found_precomputed = True
+                    assert hasattr(resp, "_topk_values"), (
+                        "Response missing _topk_values"
+                    )
+                    assert hasattr(resp, "_selected_logprob"), (
+                        "Response missing _selected_logprob"
+                    )
+                    assert len(resp._topk_indices) == _PRECOMPUTE_TOP_K
+                    assert len(resp._topk_values) == _PRECOMPUTE_TOP_K
+        assert found_precomputed, "No responses had precomputed topk"
+
+    def test_no_topk_when_not_needed(self) -> None:
+        model = _tiny_model()
+        steps = self._run_generator(model, [[1, 2, 3]], 5, needs_topk=False)
+        for step_responses in steps:
+            for resp in step_responses:
+                assert not hasattr(resp, "_topk_indices")
+
+    def test_precomputed_matches_fallback_in_batch(self) -> None:
+        model = _tiny_model()
+        tok = _mock_tokenizer()
+        steps = self._run_generator(model, [[1, 2, 3]], 10, needs_topk=True)
+        for step_responses in steps[1:]:
+            for resp in step_responses:
+                if not hasattr(resp, "_topk_indices"):
+                    continue
+                selected_fb, items_fb = extract_top_logprobs(
+                    resp.logprobs, tok, top_logprobs=5, selected_token=resp.token
+                )
+                selected_pc, items_pc = extract_top_logprobs(
+                    resp.logprobs,
+                    tok,
+                    top_logprobs=5,
+                    selected_token=resp.token,
+                    precomputed_indices=resp._topk_indices,
+                    precomputed_values=resp._topk_values,
+                    precomputed_selected=resp._selected_logprob,
+                )
+                assert selected_pc == pytest.approx(selected_fb, abs=1e-5)
+                for a, b in zip(items_pc, items_fb, strict=True):
+                    assert a.token == b.token
+                    assert a.logprob == pytest.approx(b.logprob, abs=1e-5)
+
+    def test_topk_correct_after_batch_shrink(self) -> None:
+        model = _tiny_model()
+        tok = _mock_tokenizer()
+        apply_batch_gen_patch()
+        gen = BatchGenerator(
+            model=model, stop_tokens={0}, prefill_step_size=512, max_tokens=3
+        )
+        gen._needs_topk = True
+        gen.insert([[1, 2, 3], [4, 5, 6]], max_tokens=[3, 20])
+
+        seen_shrink = False
+        for _ in range(30):
+            responses = gen.next()
+            for resp in responses:
+                if resp.finish_reason is not None:
+                    seen_shrink = True
+                    continue
+                if not hasattr(resp, "_topk_indices"):
+                    continue
+                selected_fb, items_fb = extract_top_logprobs(
+                    resp.logprobs, tok, top_logprobs=5, selected_token=resp.token
+                )
+                selected_pc, _ = extract_top_logprobs(
+                    resp.logprobs,
+                    tok,
+                    top_logprobs=5,
+                    selected_token=resp.token,
+                    precomputed_indices=resp._topk_indices,
+                    precomputed_values=resp._topk_values,
+                    precomputed_selected=resp._selected_logprob,
+                )
+                assert selected_pc == pytest.approx(selected_fb, abs=1e-5), (
+                    f"Mismatch after batch shrink: precomputed={selected_pc}, fallback={selected_fb}"
+                )
+            if gen.active_batch is None and not gen.unprocessed_prompts:
+                break
+
+        gen.close()
+        assert seen_shrink, "Expected at least one request to finish (batch shrink)"

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

@@ -486,16 +486,7 @@ def _patch_lossy_chat_template(template: str) -> str | None:
 
 
 def _needs_dsml_encoding(task_params: TextGenerationTaskParams) -> bool:
-    if "deepseek-v3.2" not in task_params.model.lower():
-        return False
-    # Use DSML encoding when tools are provided or tool results are in the conversation
-    if task_params.tools:
-        return True
-    if task_params.chat_template_messages:
-        return any(
-            msg.get("role") == "tool" for msg in task_params.chat_template_messages
-        )
-    return False
+    return "deepseek-v3.2" in task_params.model.lower()
 
 
 def apply_chat_template(
@@ -514,8 +505,6 @@ def apply_chat_template(
     if task_params.chat_template_messages is not None:
         # Use pre-formatted messages that preserve tool_calls, thinking, etc.
         formatted_messages = list(task_params.chat_template_messages)
-        for msg in formatted_messages:
-            _normalize_tool_calls(msg)
     else:
         # Add system message (instructions) if present
         if task_params.instructions:
@@ -541,7 +530,10 @@ def apply_chat_template(
 
         prompt = encode_messages(
             messages=formatted_messages,
-            thinking_mode="thinking" if task_params.enable_thinking else "chat",
+            # Only use chat mode if enable thinking is explicitly Fakse.
+            thinking_mode="chat"
+            if task_params.enable_thinking is False
+            else "thinking",
             tools=task_params.tools,
         )
         if partial_assistant_content:
@@ -549,6 +541,9 @@ def apply_chat_template(
         logger.info(prompt)
         return prompt
 
+    for msg in formatted_messages:
+        _normalize_tool_calls(msg)
+
     extra_kwargs: dict[str, Any] = {}
     if task_params.enable_thinking is not None:
         # Qwen3 and GLM use "enable_thinking"; DeepSeek uses "thinking".
@@ -638,6 +633,7 @@ class NullKVCache(KVCache):
     @property
     def state(self) -> tuple[mx.array, mx.array]:
         # matches what mx.save_safetensors / mx.eval expect
+        assert self.keys is not None and self.values is not None
         return self.keys, self.values
 
     @state.setter
@@ -739,12 +735,9 @@ def _parse_kimi_tool_calls(text: str):
         if func_args_match is None:
             raise ValueError("No tool call arguments found.")
         func_args = func_args_match.group(1)
-        try:
-            arg_dct = json.loads(func_args)  # pyright: ignore[reportAny]
-        except Exception:
-            arg_dct = None
+        arg_dct = json.loads(func_args)  # pyright: ignore[reportAny]
 
-        return dict(id=tool_call_id, name=func_name, arguments=arg_dct)
+        return dict(id=tool_call_id, name=func_name, arguments=arg_dct)  # pyright: ignore[reportAny]
 
     tool_matches = _tool_call_split_regex.findall(text)
     if tool_matches:

src/exo/worker/main.py

@@ -2,13 +2,13 @@ from collections import defaultdict
 from datetime import datetime, timezone
 
 import anyio
-from anyio import fail_after
+from anyio import fail_after, to_thread
 from loguru import logger
 
 from exo.api.types import ImageEditsTaskParams
-from exo.download.download_utils import resolve_model_in_path
+from exo.download.download_utils import is_read_only_model_dir, resolve_existing_model
 from exo.shared.apply import apply
-from exo.shared.models.model_cards import ModelId
+from exo.shared.models.model_cards import ModelId, add_to_card_cache, delete_custom_card
 from exo.shared.types.commands import (
     ForwarderCommand,
     ForwarderDownloadCommand,
@@ -16,6 +16,8 @@ from exo.shared.types.commands import (
 )
 from exo.shared.types.common import CommandId, NodeId, SystemId
 from exo.shared.types.events import (
+    CustomModelCardAdded,
+    CustomModelCardDeleted,
     Event,
     IndexedEvent,
     InputChunkReceived,
@@ -78,6 +80,7 @@ class Worker:
         self.input_chunk_counts: dict[CommandId, int] = {}
 
         self._download_backoff: KeyedBackoff[ModelId] = KeyedBackoff(base=0.5, cap=10.0)
+        self._stopped: anyio.Event = anyio.Event()
 
     async def run(self):
         logger.info("Starting Worker")
@@ -100,6 +103,7 @@ class Worker:
             self.download_command_sender.close()
             for runner in self.runners.values():
                 runner.shutdown()
+            self._stopped.set()
 
     async def _forward_info(self, recv: Receiver[GatheredInfo]):
         with recv as info_stream:
@@ -130,6 +134,13 @@ class Worker:
                         event.chunk.data
                     )
 
+                if isinstance(event, CustomModelCardAdded):
+                    await event.model_card.save_to_custom_dir()
+                    add_to_card_cache(event.model_card)
+
+                if isinstance(event, CustomModelCardDeleted):
+                    await delete_custom_card(event.model_id)
+
     async def plan_step(self):
         while True:
             await anyio.sleep(0.1)
@@ -170,11 +181,11 @@ class Worker:
                     model_id = shard.model_card.model_id
                     self._download_backoff.record_attempt(model_id)
 
-                    found_path = resolve_model_in_path(model_id)
+                    found_path = await to_thread.run_sync(
+                        resolve_existing_model, model_id
+                    )
                     if found_path is not None:
-                        logger.info(
-                            f"Model {model_id} found in EXO_MODELS_PATH at {found_path}"
-                        )
+                        logger.info(f"Model {model_id} found at {found_path}")
                         await self.event_sender.send(
                             NodeDownloadProgress(
                                 download_progress=DownloadCompleted(
@@ -182,7 +193,7 @@ class Worker:
                                     shard_metadata=shard,
                                     model_directory=str(found_path),
                                     total=shard.model_card.storage_size,
-                                    read_only=True,
+                                    read_only=is_read_only_model_dir(found_path),
                                 )
                             )
                         )
@@ -271,8 +282,9 @@ class Worker:
                 case task:
                     await self._start_runner_task(task)
 
-    def shutdown(self):
+    async def shutdown(self):
         self._tg.cancel_tasks()
+        await self._stopped.wait()
 
     async def _start_runner_task(self, task: Task):
         if (instance := self.state.instances.get(task.instance_id)) is not None:

src/exo/worker/runner/bootstrap.py

@@ -8,6 +8,7 @@ from exo.shared.types.tasks import Task, TaskId
 from exo.shared.types.worker.instances import BoundInstance
 from exo.shared.types.worker.runners import RunnerFailed
 from exo.utils.channels import ClosedResourceError, MpReceiver, MpSender
+from exo.worker.engines.mlx.patches import apply_mlx_patches
 
 logger: "loguru.Logger" = loguru.logger
 
@@ -45,6 +46,8 @@ def entrypoint(
         else:
             from exo.worker.runner.llm_inference.runner import Runner
 
+            apply_mlx_patches()
+
             runner = Runner(
                 bound_instance, event_sender, task_receiver, cancel_receiver
             )

src/exo/worker/runner/image_models/runner.py

@@ -4,7 +4,11 @@ from typing import TYPE_CHECKING, Literal
 
 import mlx.core as mx
 
-from exo.api.types import ImageGenerationStats
+from exo.api.types import (
+    ImageEditsTaskParams,
+    ImageGenerationStats,
+    ImageGenerationTaskParams,
+)
 from exo.shared.constants import EXO_MAX_CHUNK_SIZE, EXO_TRACING_ENABLED
 from exo.shared.models.model_cards import ModelTask
 from exo.shared.tracing import clear_trace_buffer, get_trace_buffer
@@ -235,6 +239,77 @@ class Runner:
     def acknowledge_task(self, task: Task):
         self.event_sender.send(TaskAcknowledged(task_id=task.task_id))
 
+    def _check_cancelled(self, task_id: TaskId) -> bool:
+        for cancel_id in self.cancel_receiver.collect():
+            self.cancelled_tasks.add(cancel_id)
+        return (
+            task_id in self.cancelled_tasks or CANCEL_ALL_TASKS in self.cancelled_tasks
+        )
+
+    def _run_image_task(
+        self,
+        task: Task,
+        task_params: ImageGenerationTaskParams | ImageEditsTaskParams,
+        command_id: CommandId,
+    ) -> None:
+        assert self.image_model
+        logger.info(f"received image task: {str(task)[:500]}")
+        logger.info("runner running")
+        self.update_status(RunnerRunning())
+        self.acknowledge_task(task)
+
+        def cancel_checker() -> bool:
+            return self._check_cancelled(task.task_id)
+
+        try:
+            image_index = 0
+            for response in generate_image(
+                model=self.image_model,
+                task=task_params,
+                cancel_checker=cancel_checker,
+            ):
+                if _is_primary_output_node(self.shard_metadata):
+                    match response:
+                        case PartialImageResponse():
+                            logger.info(
+                                f"sending partial ImageChunk {response.partial_index}/{response.total_partials}"
+                            )
+                            _process_image_response(
+                                response,
+                                command_id,
+                                self.shard_metadata,
+                                self.event_sender,
+                                image_index,
+                            )
+                        case ImageGenerationResponse():
+                            logger.info("sending final ImageChunk")
+                            _process_image_response(
+                                response,
+                                command_id,
+                                self.shard_metadata,
+                                self.event_sender,
+                                image_index,
+                            )
+                            image_index += 1
+        except Exception as e:
+            if _is_primary_output_node(self.shard_metadata):
+                self.event_sender.send(
+                    ChunkGenerated(
+                        command_id=command_id,
+                        chunk=ErrorChunk(
+                            model=self.shard_metadata.model_card.model_id,
+                            finish_reason="error",
+                            error_message=str(e),
+                        ),
+                    )
+                )
+            raise
+        finally:
+            _send_traces_if_enabled(self.event_sender, task.task_id, self.device_rank)
+
+        self.current_status = RunnerReady()
+        logger.info("runner ready")
+
     def main(self):
         with self.task_receiver as tasks:
             for task in tasks:
@@ -306,124 +381,11 @@ class Runner:
                 self.current_status = RunnerReady()
                 logger.info("runner ready")
 
-            case ImageGeneration(task_params=task_params, command_id=command_id) if (
-                isinstance(self.current_status, RunnerReady)
-            ):
-                assert self.image_model
-                logger.info(f"received image generation request: {str(task)[:500]}")
-                logger.info("runner running")
-                self.update_status(RunnerRunning())
-                self.acknowledge_task(task)
-
-                try:
-                    image_index = 0
-                    for response in generate_image(
-                        model=self.image_model, task=task_params
-                    ):
-                        is_primary_output = _is_primary_output_node(self.shard_metadata)
-
-                        if is_primary_output:
-                            match response:
-                                case PartialImageResponse():
-                                    logger.info(
-                                        f"sending partial ImageChunk {response.partial_index}/{response.total_partials}"
-                                    )
-                                    _process_image_response(
-                                        response,
-                                        command_id,
-                                        self.shard_metadata,
-                                        self.event_sender,
-                                        image_index,
-                                    )
-                                case ImageGenerationResponse():
-                                    logger.info("sending final ImageChunk")
-                                    _process_image_response(
-                                        response,
-                                        command_id,
-                                        self.shard_metadata,
-                                        self.event_sender,
-                                        image_index,
-                                    )
-                                    image_index += 1
-                # can we make this more explicit?
-                except Exception as e:
-                    if _is_primary_output_node(self.shard_metadata):
-                        self.event_sender.send(
-                            ChunkGenerated(
-                                command_id=command_id,
-                                chunk=ErrorChunk(
-                                    model=self.shard_metadata.model_card.model_id,
-                                    finish_reason="error",
-                                    error_message=str(e),
-                                ),
-                            )
-                        )
-                    raise
-                finally:
-                    _send_traces_if_enabled(
-                        self.event_sender, task.task_id, self.device_rank
-                    )
-
-                self.current_status = RunnerReady()
-                logger.info("runner ready")
-
-            case ImageEdits(task_params=task_params, command_id=command_id) if (
-                isinstance(self.current_status, RunnerReady)
-            ):
-                assert self.image_model
-                logger.info(f"received image edits request: {str(task)[:500]}")
-                logger.info("runner running")
-                self.update_status(RunnerRunning())
-                self.acknowledge_task(task)
-
-                try:
-                    image_index = 0
-                    for response in generate_image(
-                        model=self.image_model, task=task_params
-                    ):
-                        if _is_primary_output_node(self.shard_metadata):
-                            match response:
-                                case PartialImageResponse():
-                                    logger.info(
-                                        f"sending partial ImageChunk {response.partial_index}/{response.total_partials}"
-                                    )
-                                    _process_image_response(
-                                        response,
-                                        command_id,
-                                        self.shard_metadata,
-                                        self.event_sender,
-                                        image_index,
-                                    )
-                                case ImageGenerationResponse():
-                                    logger.info("sending final ImageChunk")
-                                    _process_image_response(
-                                        response,
-                                        command_id,
-                                        self.shard_metadata,
-                                        self.event_sender,
-                                        image_index,
-                                    )
-                                    image_index += 1
-                except Exception as e:
-                    if _is_primary_output_node(self.shard_metadata):
-                        self.event_sender.send(
-                            ChunkGenerated(
-                                command_id=command_id,
-                                chunk=ErrorChunk(
-                                    model=self.shard_metadata.model_card.model_id,
-                                    finish_reason="error",
-                                    error_message=str(e),
-                                ),
-                            )
-                        )
-                    raise
-                finally:
-                    _send_traces_if_enabled(
-                        self.event_sender, task.task_id, self.device_rank
-                    )
-
-                self.current_status = RunnerReady()
-                logger.info("runner ready")
+            case (
+                ImageGeneration(task_params=task_params, command_id=command_id)
+                | ImageEdits(task_params=task_params, command_id=command_id)
+            ) if isinstance(self.current_status, RunnerReady):
+                self._run_image_task(task, task_params, command_id)
 
             case Shutdown():
                 logger.info("runner shutting down")

src/exo/worker/runner/llm_inference/batch_generator.py

@@ -195,21 +195,29 @@ class SequentialGenerator(InferenceGenerator):
         assert self._active is not None
 
         task, mlx_gen, queue, output_generator = self._active
-        response = None
+        output: list[
+            tuple[TaskId, GenerationResponse | ToolCallResponse | Cancelled | Finished]
+        ] = []
         try:
-            queue.push(next(mlx_gen))
-            response = next(output_generator)
+            response = next(mlx_gen)
+            queue.push(response)
+            # drain potentially many responses every time
+            while (parsed := next(output_generator, None)) is not None:
+                output.append((task.task_id, parsed))
+
         except (StopIteration, PrefillCancelled):
-            response = Finished()
+            output.append((task.task_id, Finished()))
             self._active = None
             if self._queue:
                 self._start_next()
+
         except Exception as e:
             self._send_error(task, e)
             self._active = None
             raise
+
         return itertools.chain(
-            [] if response is None else [(task.task_id, response)],
+            output,
             map(lambda task: (task, Cancelled()), self._cancelled_tasks),
         )
 
@@ -427,11 +435,11 @@ class BatchGenerator(InferenceGenerator):
 
             task, queue, output_generator = self._active_tasks[uid]
             queue.push(response)
-            parsed = next(output_generator)
-
-            if parsed is not None:
+            # If a generator fails to parse for some reason and returns early, we should not crash
+            while (parsed := next(output_generator, None)) is not None:
                 output.append((task.task_id, parsed))
 
+            # check if original response was terminal and append a Finished()
             if response.finish_reason is not None:
                 output.append((task.task_id, Finished()))
                 del self._active_tasks[uid]

src/exo/worker/runner/llm_inference/model_output_parsers.py

@@ -159,11 +159,42 @@ def parse_deepseek_v32(
     # Text accumulated during a tool call block
     tool_call_text = ""
 
+    def _try_parse_tool_call(
+        text: str, response: GenerationResponse
+    ) -> ToolCallResponse | GenerationResponse:
+        parsed = parse_dsml_output(text)
+        if parsed is not None:
+            return ToolCallResponse(
+                tool_calls=parsed, usage=response.usage, stats=response.stats
+            )
+        logger.warning(f"DSML tool call parsing failed for: {text}")
+        return response.model_copy(update={"text": text})
+
     for response in responses:
         if response is None:
             yield None
             continue
 
+        if response.finish_reason is not None:
+            yield from pending_buffer
+            pending_buffer.clear()
+            if in_tool_call:
+                tool_call_text += response.text
+                yield (
+                    _try_parse_tool_call(tool_call_text, response)
+                    if TOOL_CALLS_END in tool_call_text
+                    else response.model_copy(update={"text": tool_call_text})
+                )
+            elif TOOL_CALLS_START in response.text and TOOL_CALLS_END in response.text:
+                dsml_start = response.text.index(TOOL_CALLS_START)
+                before = response.text[:dsml_start]
+                if before:
+                    yield response.model_copy(update={"text": before})
+                yield _try_parse_tool_call(response.text[dsml_start:], response)
+            else:
+                yield response
+            break
+
         # ── Handle thinking tags ──
         if not thinking and THINKING_START in response.text:
             thinking = True
@@ -191,28 +222,7 @@ def parse_deepseek_v32(
         if in_tool_call:
             tool_call_text += response.text
             if TOOL_CALLS_END in tool_call_text:
-                # Parse the accumulated DSML block
-                parsed = parse_dsml_output(tool_call_text)
-                if parsed is not None:
-                    logger.info(f"parsed DSML tool calls: {parsed}")
-                    yield ToolCallResponse(
-                        tool_calls=parsed,
-                        usage=response.usage,
-                        stats=response.stats,
-                    )
-                else:
-                    logger.warning(
-                        f"DSML tool call parsing failed for: {tool_call_text}"
-                    )
-                    yield response.model_copy(update={"text": tool_call_text})
-                in_tool_call = False
-                tool_call_text = ""
-                continue
-
-            # EOS reached before end marker — yield buffered text as-is
-            if response.finish_reason is not None:
-                logger.info("DSML tool call parsing interrupted by EOS")
-                yield response.model_copy(update={"text": tool_call_text})
+                yield _try_parse_tool_call(tool_call_text, response)
                 in_tool_call = False
                 tool_call_text = ""
             continue
@@ -228,33 +238,22 @@ def parse_deepseek_v32(
             if pre_text:
                 # Flush pending buffer tokens that contributed text before the marker
                 for buf_resp in pending_buffer:
-                    if pre_text:
-                        chunk = buf_resp.text
-                        if len(chunk) <= len(pre_text):
-                            yield buf_resp
-                            pre_text = pre_text[len(chunk) :]
-                        else:
-                            yield buf_resp.model_copy(update={"text": pre_text})
-                            pre_text = ""
+                    if not pre_text:
+                        break
+                    chunk = buf_resp.text
+                    if len(chunk) <= len(pre_text):
+                        yield buf_resp
+                        pre_text = pre_text[len(chunk) :]
+                    else:
+                        yield buf_resp.model_copy(update={"text": pre_text})
+                        pre_text = ""
             pending_buffer = []
             tool_call_text = accumulated[start_idx:]
             accumulated = ""
 
             # Check if the end marker is already present (entire tool call in one token)
             if TOOL_CALLS_END in tool_call_text:
-                parsed = parse_dsml_output(tool_call_text)
-                if parsed is not None:
-                    logger.info(f"parsed DSML tool calls: {parsed}")
-                    yield ToolCallResponse(
-                        tool_calls=parsed,
-                        usage=response.usage,
-                        stats=response.stats,
-                    )
-                else:
-                    logger.warning(
-                        f"DSML tool call parsing failed for: {tool_call_text}"
-                    )
-                    yield response.model_copy(update={"text": tool_call_text})
+                yield _try_parse_tool_call(tool_call_text, response)
                 tool_call_text = ""
             else:
                 in_tool_call = True
@@ -267,15 +266,13 @@ def parse_deepseek_v32(
             continue
 
         # No partial match — flush all pending tokens and the current one
-        for buf_resp in pending_buffer:
-            yield buf_resp
-        pending_buffer = []
+        yield from pending_buffer
+        pending_buffer.clear()
         accumulated = ""
         yield response
 
     # Flush any remaining pending buffer at generator end
-    for buf_resp in pending_buffer:
-        yield buf_resp
+    yield from pending_buffer
 
 
 def _could_be_dsml_prefix(text: str) -> bool:
@@ -358,8 +355,10 @@ def parse_tool_calls(
 
             if parsed is None:
                 logger.warning(f"tool call parsing failed for text {combined}")
-                yield response.model_copy(update={"text": combined})
-                continue
+                yield response.model_copy(
+                    update={"text": combined, "token": 0, "finish_reason": "error"}
+                )
+                break
 
             yield ToolCallResponse(
                 tool_calls=parsed, usage=response.usage, stats=response.stats
@@ -374,6 +373,7 @@ def parse_tool_calls(
                 update={
                     "text": "".join(tool_call_text_parts),
                     "token": 0,
+                    "finish_reason": "error",
                 }
             )
             yield response

src/exo/worker/runner/llm_inference/runner.py

@@ -319,7 +319,9 @@ class Runner:
         return ExitCode.AllTasksComplete
 
     def send_response(
-        self, response: GenerationResponse | ToolCallResponse, command_id: CommandId
+        self,
+        response: GenerationResponse | ToolCallResponse,
+        command_id: CommandId,
     ):
         match response:
             case GenerationResponse():

src/exo/worker/runner/runner_supervisor.py

@@ -110,39 +110,45 @@ class RunnerSupervisor:
 
     async def run(self):
         self.runner_process.start()
-        async with self._tg as tg:
-            tg.start_soon(self._watch_runner)
-            tg.start_soon(self._forward_events)
+        try:
+            async with self._tg as tg:
+                tg.start_soon(self._watch_runner)
+                tg.start_soon(self._forward_events)
+        finally:
+            logger.info("Runner supervisor shutting down")
+            if not self._cancel_watch_runner.cancel_called:
+                self._cancel_watch_runner.cancel()
+            with contextlib.suppress(ClosedResourceError):
+                self._ev_recv.close()
+            with contextlib.suppress(ClosedResourceError):
+                self._task_sender.close()
+            with contextlib.suppress(ClosedResourceError):
+                self._event_sender.close()
+            with contextlib.suppress(ClosedResourceError):
+                self._cancel_sender.send(CANCEL_ALL_TASKS)
+            with contextlib.suppress(ClosedResourceError):
+                self._cancel_sender.close()
+
+            await to_thread.run_sync(self.runner_process.join, 5)
+
+            if self.runner_process.is_alive():
+                logger.warning(
+                    "Runner process didn't shutdown succesfully, terminating"
+                )
+                self.runner_process.terminate()
+                self.runner_process.join(timeout=5)
+                # This is overkill but it's not technically bad, just unnecessary.
+                if self.runner_process.is_alive():
+                    logger.critical("Runner process didn't respond to SIGTERM, killing")
+                    self.runner_process.kill()
+                    self.runner_process.join(timeout=5)
+            else:
+                logger.info("Runner process succesfully terminated")
+
+            self.runner_process.close()
 
     def shutdown(self):
-        logger.info("Runner supervisor shutting down")
         self._tg.cancel_tasks()
-        if not self._cancel_watch_runner.cancel_called:
-            self._cancel_watch_runner.cancel()
-        with contextlib.suppress(ClosedResourceError):
-            self._ev_recv.close()
-        with contextlib.suppress(ClosedResourceError):
-            self._task_sender.close()
-        with contextlib.suppress(ClosedResourceError):
-            self._event_sender.close()
-        with contextlib.suppress(ClosedResourceError):
-            self._cancel_sender.send(CANCEL_ALL_TASKS)
-        with contextlib.suppress(ClosedResourceError):
-            self._cancel_sender.close()
-        self.runner_process.join(5)
-        if not self.runner_process.is_alive():
-            logger.info("Runner process succesfully terminated")
-            return
-
-        # This is overkill but it's not technically bad, just unnecessary.
-        logger.warning("Runner process didn't shutdown succesfully, terminating")
-        self.runner_process.terminate()
-        self.runner_process.join(1)
-        if not self.runner_process.is_alive():
-            return
-
-        logger.critical("Runner process didn't respond to SIGTERM, killing")
-        self.runner_process.kill()
 
     async def start_task(self, task: Task):
         if task.task_id in self.pending:
@@ -218,12 +224,6 @@ class RunnerSupervisor:
             for tid in self.pending:
                 self.pending[tid].set()
 
-    def __del__(self) -> None:
-        if self.runner_process.is_alive():
-            logger.critical("RunnerSupervisor was not stopped cleanly.")
-            with contextlib.suppress(ValueError):
-                self.runner_process.kill()
-
     async def _watch_runner(self) -> None:
         with self._cancel_watch_runner:
             while True:

src/exo/worker/tests/unittests/test_image/test_cancellation.py

@@ -0,0 +1,433 @@
+# pyright: reportPrivateUsage=false
+"""Tests for image generation cancellation logic.
+
+Tests the NaN sentinel protocol, cancellation checking, and the
+image runner's cancel_checker integration.
+"""
+
+from collections.abc import Callable
+from unittest.mock import MagicMock
+
+import mlx.core as mx
+
+from exo.shared.types.tasks import CANCEL_ALL_TASKS, TaskId
+from exo.worker.engines.image.pipeline.runner import DiffusionRunner
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _make_runner() -> DiffusionRunner:
+    """Create a DiffusionRunner with minimal config for unit testing.
+
+    Uses a mock adapter and no distributed group (single-node).
+    """
+    mock_config = MagicMock()
+    mock_config.joint_block_count = 10
+    mock_config.single_block_count = 10
+    mock_config.total_blocks = 20
+    mock_config.guidance_scale = None
+
+    mock_adapter = MagicMock()
+
+    mock_shard = MagicMock()
+    mock_shard.device_rank = 0
+    mock_shard.world_size = 1
+    mock_shard.start_layer = 0
+    mock_shard.end_layer = 20
+
+    runner = DiffusionRunner(
+        config=mock_config,
+        adapter=mock_adapter,
+        group=None,
+        shard_metadata=mock_shard,
+    )
+    return runner
+
+
+class FakeCancelReceiver:
+    """Fake MpReceiver that returns pre-loaded items from collect()."""
+
+    def __init__(self, items: list[TaskId] | None = None):
+        self._items = list(items) if items else []
+
+    def collect(self) -> list[TaskId]:
+        result = self._items
+        self._items = []
+        return result
+
+
+class FakeImageRunner:
+    """Fake image runner for testing _check_cancelled logic."""
+
+    def __init__(self, cancel_items: list[TaskId] | None = None) -> None:
+        self.cancel_receiver = FakeCancelReceiver(cancel_items)
+        self.cancelled_tasks = set[TaskId]()
+
+    def _check_cancelled(self, task_id: TaskId) -> bool:
+        for cancel_id in self.cancel_receiver.collect():
+            self.cancelled_tasks.add(cancel_id)
+        return (
+            task_id in self.cancelled_tasks or CANCEL_ALL_TASKS in self.cancelled_tasks
+        )
+
+
+# ---------------------------------------------------------------------------
+# _is_sentinel
+# ---------------------------------------------------------------------------
+
+
+class TestIsSentinel:
+    def test_all_nan_is_sentinel(self) -> None:
+        runner = _make_runner()
+        tensor = mx.full((2, 3), float("nan"))
+        mx.eval(tensor)
+        assert runner._is_sentinel(tensor) is True
+
+    def test_all_zeros_is_not_sentinel(self) -> None:
+        runner = _make_runner()
+        tensor = mx.zeros((2, 3))
+        mx.eval(tensor)
+        assert runner._is_sentinel(tensor) is False
+
+    def test_mixed_nan_and_real_is_not_sentinel(self) -> None:
+        """A tensor with some NaN and some real values must NOT be a sentinel.
+        Using mx.any(isnan) would incorrectly flag this as a sentinel.
+        """
+        runner = _make_runner()
+        tensor = mx.array([float("nan"), 1.0, 2.0])
+        mx.eval(tensor)
+        assert runner._is_sentinel(tensor) is False
+
+    def test_single_element_nan(self) -> None:
+        runner = _make_runner()
+        tensor = mx.array([float("nan")])
+        mx.eval(tensor)
+        assert runner._is_sentinel(tensor) is True
+
+    def test_large_tensor_all_nan(self) -> None:
+        runner = _make_runner()
+        tensor = mx.full((64, 128, 32), float("nan"))
+        mx.eval(tensor)
+        assert runner._is_sentinel(tensor) is True
+
+    def test_real_data_not_sentinel(self) -> None:
+        runner = _make_runner()
+        tensor = mx.random.normal((4, 8))
+        mx.eval(tensor)
+        assert runner._is_sentinel(tensor) is False
+
+
+# ---------------------------------------------------------------------------
+# _check_cancellation
+# ---------------------------------------------------------------------------
+
+
+class TestCheckCancellation:
+    def test_first_stage_polls_checker(self) -> None:
+        runner = _make_runner()
+        assert runner.is_first_stage  # single-node is always first stage
+
+        checker: Callable[[], bool] = MagicMock(return_value=True)
+        runner._cancel_checker = checker
+
+        runner._check_cancellation()
+
+        checker.assert_called_once()
+        assert runner._cancelling is True
+
+    def test_checker_returning_false_does_not_cancel(self) -> None:
+        runner = _make_runner()
+        checker: Callable[[], bool] = MagicMock(return_value=False)
+        runner._cancel_checker = checker
+
+        runner._check_cancellation()
+
+        assert runner._cancelling is False
+
+    def test_no_checker_does_not_cancel(self) -> None:
+        runner = _make_runner()
+        runner._cancel_checker = None
+
+        runner._check_cancellation()
+
+        assert runner._cancelling is False
+
+    def test_already_cancelling_skips_checker(self) -> None:
+        runner = _make_runner()
+        runner._cancelling = True
+        checker: Callable[[], bool] = MagicMock(return_value=False)
+        runner._cancel_checker = checker
+
+        runner._check_cancellation()
+
+        checker.assert_not_called()
+        assert runner._cancelling is True  # stays True
+
+    def test_cancelling_flag_is_false_on_init(self) -> None:
+        """_cancelling defaults to False on a fresh runner."""
+        runner = _make_runner()
+        assert runner._cancelling is False
+
+
+# ---------------------------------------------------------------------------
+# _send wrapper
+# ---------------------------------------------------------------------------
+
+
+class TestSendWrapper:
+    def test_send_replaces_data_with_nan_when_cancelling(self) -> None:
+        """When _cancelling is True, _send should replace data with NaN."""
+        runner = _make_runner()
+        runner._cancelling = True
+        # _send asserts group is not None, so we need a mock group
+        runner.group = MagicMock()
+
+        data = mx.ones((2, 3))
+        mx.eval(data)
+
+        # Mock mx.distributed.send to capture what's sent
+        original_send = mx.distributed.send
+        sent_data: list[mx.array] = []
+
+        def mock_send(d: mx.array, dst: int, group: mx.distributed.Group) -> mx.array:
+            mx.eval(d)
+            sent_data.append(d)
+            return d
+
+        mx.distributed.send = mock_send
+        try:
+            runner._send(data, dst=1)
+            assert len(sent_data) == 1
+            mx.eval(sent_data[0])
+            assert mx.all(mx.isnan(sent_data[0])).item()
+            assert sent_data[0].shape == (2, 3)
+        finally:
+            mx.distributed.send = original_send
+
+    def test_send_passes_real_data_when_not_cancelling(self) -> None:
+        runner = _make_runner()
+        runner._cancelling = False
+        runner.group = MagicMock()
+
+        data = mx.ones((2, 3))
+        mx.eval(data)
+
+        sent_data: list[mx.array] = []
+
+        def mock_send(d: mx.array, dst: int, group: mx.distributed.Group) -> mx.array:
+            mx.eval(d)
+            sent_data.append(d)
+            return d
+
+        original_send = mx.distributed.send
+        mx.distributed.send = mock_send
+        try:
+            runner._send(data, dst=1)
+            assert len(sent_data) == 1
+            mx.eval(sent_data[0])
+            assert not mx.any(mx.isnan(sent_data[0])).item()
+        finally:
+            mx.distributed.send = original_send
+
+
+# ---------------------------------------------------------------------------
+# Image runner _check_cancelled
+# ---------------------------------------------------------------------------
+
+
+class TestImageRunnerCheckCancelled:
+    """Tests for the image runner's _check_cancelled method."""
+
+    def test_no_cancellation(self) -> None:
+        runner = FakeImageRunner()
+        assert runner._check_cancelled(TaskId("task-1")) is False
+
+    def test_specific_task_cancelled(self) -> None:
+        task_id = TaskId("task-1")
+        runner = FakeImageRunner([task_id])
+        assert runner._check_cancelled(task_id) is True
+
+    def test_different_task_not_cancelled(self) -> None:
+        runner = FakeImageRunner([TaskId("task-2")])
+        assert runner._check_cancelled(TaskId("task-1")) is False
+
+    def test_cancel_all_tasks(self) -> None:
+        runner = FakeImageRunner([CANCEL_ALL_TASKS])
+        assert runner._check_cancelled(TaskId("any-task")) is True
+
+    def test_collect_accumulates(self) -> None:
+        """Multiple collect() calls accumulate cancelled task IDs."""
+        runner = FakeImageRunner([TaskId("task-1")])
+        runner._check_cancelled(TaskId("task-1"))
+
+        # First collect drained the receiver, but task-1 is in cancelled_tasks
+        assert runner._check_cancelled(TaskId("task-1")) is True
+
+    def test_collect_empty_after_drain(self) -> None:
+        """After draining, collect returns empty and previous cancellations persist."""
+        runner = FakeImageRunner([TaskId("task-1")])
+
+        # First call drains
+        runner._check_cancelled(TaskId("other"))
+        # task-1 is now in cancelled_tasks but "other" was never cancelled
+        assert runner._check_cancelled(TaskId("other")) is False
+        assert runner._check_cancelled(TaskId("task-1")) is True
+
+
+# ---------------------------------------------------------------------------
+# Drain condition logic
+# ---------------------------------------------------------------------------
+
+
+class TestDrainCondition:
+    """Verify the drain condition evaluates correctly for various scenarios."""
+
+    def _should_drain(
+        self,
+        *,
+        cancelling: bool,
+        is_first_stage: bool,
+        is_last_stage: bool,
+        is_distributed: bool,
+        t: int,
+        init_time_step: int,
+        num_sync_steps: int,
+        num_inference_steps: int,
+    ) -> bool:
+        """Replicate the drain condition from _run_diffusion_loop."""
+        return (
+            cancelling
+            and is_first_stage
+            and not is_last_stage
+            and is_distributed
+            and t >= init_time_step + num_sync_steps
+            and t != num_inference_steps - 1
+        )
+
+    def test_no_drain_during_sync_step(self) -> None:
+        """Sync steps have no cross-timestep ring state."""
+        assert not self._should_drain(
+            cancelling=True,
+            is_first_stage=True,
+            is_last_stage=False,
+            is_distributed=True,
+            t=0,  # sync step
+            init_time_step=0,
+            num_sync_steps=2,
+            num_inference_steps=10,
+        )
+
+    def test_drain_during_async_step(self) -> None:
+        assert self._should_drain(
+            cancelling=True,
+            is_first_stage=True,
+            is_last_stage=False,
+            is_distributed=True,
+            t=3,  # async step
+            init_time_step=0,
+            num_sync_steps=2,
+            num_inference_steps=10,
+        )
+
+    def test_no_drain_on_last_step(self) -> None:
+        """Last step doesn't send, so nothing to drain."""
+        assert not self._should_drain(
+            cancelling=True,
+            is_first_stage=True,
+            is_last_stage=False,
+            is_distributed=True,
+            t=9,  # last step
+            init_time_step=0,
+            num_sync_steps=2,
+            num_inference_steps=10,
+        )
+
+    def test_no_drain_when_not_cancelling(self) -> None:
+        assert not self._should_drain(
+            cancelling=False,
+            is_first_stage=True,
+            is_last_stage=False,
+            is_distributed=True,
+            t=5,
+            init_time_step=0,
+            num_sync_steps=2,
+            num_inference_steps=10,
+        )
+
+    def test_no_drain_on_last_stage(self) -> None:
+        """Last stage is also first stage (single pipeline) — no ring."""
+        assert not self._should_drain(
+            cancelling=True,
+            is_first_stage=True,
+            is_last_stage=True,
+            is_distributed=True,
+            t=5,
+            init_time_step=0,
+            num_sync_steps=2,
+            num_inference_steps=10,
+        )
+
+    def test_no_drain_single_node(self) -> None:
+        assert not self._should_drain(
+            cancelling=True,
+            is_first_stage=True,
+            is_last_stage=False,
+            is_distributed=False,
+            t=5,
+            init_time_step=0,
+            num_sync_steps=2,
+            num_inference_steps=10,
+        )
+
+    def test_no_drain_not_first_stage(self) -> None:
+        """Only first stage needs to drain (it's the one receiving)."""
+        assert not self._should_drain(
+            cancelling=True,
+            is_first_stage=False,
+            is_last_stage=False,
+            is_distributed=True,
+            t=5,
+            init_time_step=0,
+            num_sync_steps=2,
+            num_inference_steps=10,
+        )
+
+    def test_drain_first_async_step(self) -> None:
+        """First async step: last stage sends, so drain is needed."""
+        assert self._should_drain(
+            cancelling=True,
+            is_first_stage=True,
+            is_last_stage=False,
+            is_distributed=True,
+            t=2,  # first async step (init=0, sync=2)
+            init_time_step=0,
+            num_sync_steps=2,
+            num_inference_steps=10,
+        )
+
+    def test_drain_with_nonzero_init_time_step(self) -> None:
+        """img2img can have init_time_step > 0."""
+        assert self._should_drain(
+            cancelling=True,
+            is_first_stage=True,
+            is_last_stage=False,
+            is_distributed=True,
+            t=5,
+            init_time_step=3,
+            num_sync_steps=1,
+            num_inference_steps=10,
+        )
+
+    def test_no_drain_sync_with_nonzero_init(self) -> None:
+        assert not self._should_drain(
+            cancelling=True,
+            is_first_stage=True,
+            is_last_stage=False,
+            is_distributed=True,
+            t=3,
+            init_time_step=3,
+            num_sync_steps=1,
+            num_inference_steps=10,
+        )

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

@@ -10,7 +10,7 @@ from typing import Any, cast
 import mlx.core as mx
 import mlx.nn as nn
 
-from exo.shared.constants import EXO_MODELS_DIR
+from exo.shared.constants import EXO_DEFAULT_MODELS_DIR
 from exo.shared.models.model_cards import ModelCard, ModelTask
 from exo.shared.types.common import ModelId
 from exo.shared.types.memory import Memory
@@ -52,7 +52,7 @@ def create_hostfile(world_size: int, base_port: int) -> tuple[str, list[str]]:
 # 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",
+    model_path=EXO_DEFAULT_MODELS_DIR / "mlx-community--gpt-oss-20b-MXFP4-Q8",
     total_layers=24,
     base_port=29600,
     max_tokens=200,

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

@@ -43,8 +43,8 @@ def run_pipeline_device(
 
         def __call__(self, x: mx.array, *args: object, **kwargs: object) -> mx.array:
             for layer in self.layers:
-                x = layer(x, *args, **kwargs)  # pyright: ignore[reportUnknownVariableType]
-            return x  # pyright: ignore[reportUnknownVariableType]
+                x = layer(x, *args, **kwargs)
+            return x
 
     try:
         group = mx.distributed.init(backend="ring", strict=True)

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

@@ -1,389 +0,0 @@
-import copy
-import gc
-import json
-import shutil
-import tempfile
-from pathlib import Path
-from typing import Any, cast
-
-import mlx.core as mx
-import pytest
-from mlx_lm.tokenizer_utils import TokenizerWrapper
-
-from exo.shared.types.common import ModelId
-from exo.shared.types.mlx import KVCacheType, Model
-from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
-from exo.worker.engines.mlx.cache import CacheSnapshot, KVPrefixCache, cache_length
-from exo.worker.engines.mlx.generator.batch_generate import ExoBatchGenerator
-from exo.worker.engines.mlx.generator.generate import mlx_generate
-from exo.worker.engines.mlx.utils_mlx import (
-    apply_chat_template,
-    load_tokenizer_for_model_id,
-)
-
-from .test_prefix_cache_architectures import (
-    ARCHITECTURES,
-    ArchSpec,
-    _arch_available,  # pyright: ignore[reportPrivateUsage]
-    _build_model,  # pyright: ignore[reportPrivateUsage]
-    _copy_tokenizer,  # pyright: ignore[reportPrivateUsage]
-    _find_snapshot,  # pyright: ignore[reportPrivateUsage]
-    _reduce_config,  # pyright: ignore[reportPrivateUsage]
-)
-
-
-def _make_task(
-    content: str = "Hello, what is 2+2?",
-    max_tokens: int = 10,
-    seed: int = 42,
-) -> TextGenerationTaskParams:
-    return TextGenerationTaskParams(
-        model=ModelId("test"),
-        input=[InputMessage(role="user", content=content)],
-        max_output_tokens=max_tokens,
-        temperature=0.7,
-        seed=seed,
-    )
-
-
-# ── Helpers ──────────────────────────────────────────────────────────────── #
-
-
-def _collect_mlx_generate(
-    model: Model,
-    tokenizer: TokenizerWrapper,
-    task: TextGenerationTaskParams,
-    kv_prefix_cache: KVPrefixCache | None,
-) -> list[int]:
-    """Run mlx_generate and collect output token IDs."""
-    prompt = apply_chat_template(tokenizer=tokenizer, task_params=task)
-    tokens: list[int] = []
-    for resp in mlx_generate(
-        model=model,
-        tokenizer=tokenizer,
-        task=task,
-        prompt=prompt,
-        kv_prefix_cache=kv_prefix_cache,
-        group=None,
-    ):
-        tokens.append(resp.token)
-        if resp.finish_reason is not None:
-            break
-    return tokens
-
-
-def _collect_batch_generate(
-    model: Model,
-    tokenizer: TokenizerWrapper,
-    task_params: TextGenerationTaskParams,
-    kv_prefix_cache: KVPrefixCache | None,
-) -> list[int]:
-    """Run ExoBatchGenerator and collect raw output token IDs"""
-    exo_gen = ExoBatchGenerator(
-        model=model,
-        tokenizer=tokenizer,
-        group=None,
-        kv_prefix_cache=kv_prefix_cache,
-    )
-
-    prompt = apply_chat_template(tokenizer=tokenizer, task_params=task_params)
-    exo_gen.submit(task_params=task_params, prompt=prompt)
-
-    tokens: list[int] = []
-    while exo_gen.has_work:
-        results = exo_gen.step()
-        for _uid, response in results:
-            tokens.append(response.token)
-
-    exo_gen.close()
-    return tokens
-
-
-def _assert_state_equal(sa: object, sb: object, label: str) -> None:
-    """Compare two state items, handling both plain arrays and tuples of arrays (CacheList)."""
-    if isinstance(sa, tuple):
-        assert isinstance(sb, tuple), f"{label}: type mismatch"
-        for k, (arr_a, arr_b) in enumerate(
-            zip(
-                cast(tuple[mx.array, ...], sa),
-                cast(tuple[mx.array, ...], sb),
-                strict=True,
-            )
-        ):
-            a_f = mx.array(arr_a).astype(mx.float32)
-            b_f = mx.array(arr_b).astype(mx.float32)
-            if a_f.size == 0:
-                assert b_f.size == 0, f"{label}[{k}]: size mismatch"
-                continue
-            diff = float(mx.max(mx.abs(a_f - b_f)).item())
-            assert diff == 0.0, f"{label}[{k}]: max diff {diff}"
-    else:
-        sa_f = mx.array(cast(mx.array, sa)).astype(mx.float32)
-        sb_f = mx.array(cast(mx.array, sb)).astype(mx.float32)
-        if sa_f.size == 0:
-            assert sb_f.size == 0, f"{label}: size mismatch"
-            return
-        diff = float(mx.max(mx.abs(sa_f - sb_f)).item())
-        assert diff == 0.0, f"{label}: max diff {diff}"
-
-
-def _compare_cache_arrays(
-    cache_a: KVCacheType,
-    cache_b: KVCacheType,
-    label: str = "",
-) -> None:
-    """Assert two KV caches have identical array values."""
-    assert len(cache_a) == len(cache_b), (
-        f"{label}Cache layer count: {len(cache_a)} vs {len(cache_b)}"
-    )
-    for i, (a, b) in enumerate(zip(cache_a, cache_b, strict=True)):
-        assert type(a) is type(b), (
-            f"{label}Layer {i}: type {type(a).__name__} vs {type(b).__name__}"
-        )
-        states_a = a.state
-        states_b = b.state
-        assert len(states_a) == len(states_b), (
-            f"{label}Layer {i}: state count {len(states_a)} vs {len(states_b)}"
-        )
-        for j, (sa, sb) in enumerate(zip(states_a, states_b, strict=True)):
-            if sa is None and sb is None:
-                continue
-            assert sa is not None and sb is not None, (
-                f"{label}Layer {i}, state {j}: one is None"
-            )
-            _assert_state_equal(sa, sb, f"{label}Layer {i}, state {j}")
-
-
-def _safe_state(cache: object) -> list[object]:
-    """Safely access .state on a cache object. Returns [] if uninitialized."""
-    # RotatingKVCache.state crashes when keys is None (uninitialized)
-    if getattr(cache, "keys", _SENTINEL) is None:
-        return []
-    try:
-        return list(cache.state)  # type: ignore[union-attr]
-    except (AttributeError, TypeError):
-        return []
-
-
-_SENTINEL = object()
-
-
-def _compare_snapshots(
-    snaps_a: list[CacheSnapshot] | None,
-    snaps_b: list[CacheSnapshot] | None,
-    label: str = "",
-) -> None:
-    """Assert two snapshot lists are identical."""
-    if snaps_a is None:
-        assert snaps_b is None, f"{label}One side has snapshots, other doesn't"
-        return
-    assert snaps_b is not None, f"{label}One side has snapshots, other doesn't"
-    assert len(snaps_a) == len(snaps_b), (
-        f"{label}Snapshot count: {len(snaps_a)} vs {len(snaps_b)}"
-    )
-    for k, (sa, sb) in enumerate(zip(snaps_a, snaps_b, strict=True)):
-        assert sa.token_count == sb.token_count, (
-            f"{label}Snapshot {k} token_count: {sa.token_count} vs {sb.token_count}"
-        )
-        for layer_i, (s1, s2) in enumerate(zip(sa.states, sb.states, strict=True)):
-            if s1 is None and s2 is None:
-                continue
-            assert s1 is not None and s2 is not None, (
-                f"{label}Snapshot {k}, layer {layer_i}: one state is None"
-            )
-            state_a = _safe_state(s1)
-            state_b = _safe_state(s2)
-            if not state_a and not state_b:
-                continue
-            assert len(state_a) == len(state_b), (
-                f"{label}Snapshot {k}, layer {layer_i}: state length mismatch"
-            )
-            for st_j, (arr_a, arr_b) in enumerate(zip(state_a, state_b, strict=True)):
-                if arr_a is None and arr_b is None:
-                    continue
-                assert arr_a is not None and arr_b is not None
-                _assert_state_equal(
-                    arr_a,
-                    arr_b,
-                    f"{label}Snapshot {k}, layer {layer_i}, state {st_j}",
-                )
-
-
-# ── Test class ────────────────────────────────────────────────────────────── #
-
-
-@pytest.mark.slow
-class TestBatchVsGenerate:
-    """Verify BatchGenerator matches mlx_generate for output tokens and prefix cache."""
-
-    @pytest.fixture(autouse=True)
-    def _cleanup(self):
-        yield
-        mx.clear_cache()
-        gc.collect()
-
-    @pytest.mark.parametrize(
-        "spec",
-        ARCHITECTURES,
-        ids=[a.name for a in ARCHITECTURES],
-    )
-    def test_same_output_and_cache(self, spec: ArchSpec) -> None:
-        if not _arch_available(spec):
-            pytest.skip(f"Model {spec.hub_name} not cached locally")
-
-        snapshot = _find_snapshot(spec.hub_name)
-        assert snapshot is not None
-
-        tmpdir = Path(tempfile.mkdtemp(prefix=f"exo_batchtest_{spec.name}_"))
-        try:
-            # Build reduced config
-            with open(snapshot / "config.json") as f:
-                cfg = cast(dict[str, Any], json.load(f))
-            reduced = _reduce_config(copy.deepcopy(cfg))
-            (tmpdir / "config.json").write_text(json.dumps(reduced))
-
-            # Copy tokenizer
-            tok_src = snapshot
-            if spec.tokenizer_hub is not None:
-                alt = _find_snapshot(spec.tokenizer_hub)
-                if alt is not None:
-                    tok_src = alt
-            _copy_tokenizer(tok_src, tmpdir)
-
-            # Load tokenizer, build model with random weights
-            model_id = ModelId(f"mlx-community/{spec.hub_name}")
-            tokenizer = load_tokenizer_for_model_id(model_id, tmpdir)
-            mx.random.seed(0)
-            model = _build_model(spec.module, reduced)
-
-            task = _make_task()
-
-            # ── Run mlx_generate path ──
-            # Seed is set inside mlx_generate/ExoBatchGenerator.submit from task.seed
-            kv_mlx = KVPrefixCache(None)
-            mlx_tokens = _collect_mlx_generate(model, tokenizer, task, kv_mlx)
-
-            # ── Run batch generator path ──
-            kv_batch = KVPrefixCache(None)
-            batch_tokens = _collect_batch_generate(model, tokenizer, task, kv_batch)
-
-            # ── Compare output tokens ──
-            assert len(mlx_tokens) > 0, "mlx_generate produced no tokens"
-            assert len(batch_tokens) > 0, "BatchGenerator produced no tokens"
-            assert mlx_tokens == batch_tokens, (
-                f"[{spec.name}] Token mismatch:\n"
-                f"  mlx_generate:    {mlx_tokens}\n"
-                f"  BatchGenerator:  {batch_tokens}"
-            )
-
-            # ── Compare prefix cache KV arrays ──
-            assert len(kv_mlx.caches) == 1, "mlx_generate didn't save to prefix cache"
-            assert len(kv_batch.caches) == 1, (
-                "BatchGenerator didn't save to prefix cache"
-            )
-
-            _compare_cache_arrays(
-                kv_mlx.caches[0],
-                kv_batch.caches[0],
-                label=f"[{spec.name}] ",
-            )
-
-            # ── Compare cache lengths ──
-            mlx_len = cache_length(kv_mlx.caches[0])
-            batch_len = cache_length(kv_batch.caches[0])
-            assert mlx_len == batch_len, (
-                f"[{spec.name}] Cache length: mlx={mlx_len} vs batch={batch_len}"
-            )
-
-            # ── Compare snapshots ──
-            _compare_snapshots(
-                kv_mlx._snapshots[0],  # pyright: ignore[reportPrivateUsage]
-                kv_batch._snapshots[0],  # pyright: ignore[reportPrivateUsage]
-                label=f"[{spec.name}] ",
-            )
-
-        finally:
-            shutil.rmtree(tmpdir, ignore_errors=True)
-
-    @pytest.mark.parametrize(
-        "spec",
-        ARCHITECTURES,
-        ids=[a.name for a in ARCHITECTURES],
-    )
-    def test_concurrent_batch_completes(self, spec: ArchSpec) -> None:
-        """Two requests processed concurrently must both complete without
-        crashing and produce non-empty output.
-
-        Note: batch decode logits are NOT bit-exact with sequential because
-        Metal's matmul kernel picks different reduction tiling for B=1 vs B=2
-        when L=1 (decode step). This introduces sub-ULP float16 diffs in
-        gate_proj/down_proj/lm_head which swiglu amplifies by |up_values|.
-        With random weights these accumulate into argmax flips; with trained
-        weights the diffs are absorbed and output matches exactly (verified
-        with real Llama-3.2-1B-Instruct-4bit weights).
-        """
-        if not _arch_available(spec):
-            pytest.skip(f"Model {spec.hub_name} not cached locally")
-
-        snapshot = _find_snapshot(spec.hub_name)
-        assert snapshot is not None
-
-        tmpdir = Path(tempfile.mkdtemp(prefix=f"exo_concurrent_{spec.name}_"))
-        try:
-            with open(snapshot / "config.json") as f:
-                cfg = cast(dict[str, Any], json.load(f))
-            reduced = _reduce_config(copy.deepcopy(cfg))
-            (tmpdir / "config.json").write_text(json.dumps(reduced))
-
-            tok_src = snapshot
-            if spec.tokenizer_hub is not None:
-                alt = _find_snapshot(spec.tokenizer_hub)
-                if alt is not None:
-                    tok_src = alt
-            _copy_tokenizer(tok_src, tmpdir)
-
-            model_id = ModelId(f"mlx-community/{spec.hub_name}")
-            tokenizer = load_tokenizer_for_model_id(model_id, tmpdir)
-            mx.random.seed(0)
-            model = _build_model(spec.module, reduced)
-
-            # Two different prompts → different prompt lengths.
-            task_a = _make_task(content="Hello, what is 2+2?", seed=42)
-            task_a = task_a.model_copy(update={"temperature": 0.0})
-            task_b = _make_task(
-                content="Write a short poem about the ocean and the sky.",
-                seed=99,
-            )
-            task_b = task_b.model_copy(update={"temperature": 0.0})
-
-            # ── Concurrent: submit both to one ExoBatchGenerator ──
-            exo_gen = ExoBatchGenerator(
-                model=model,
-                tokenizer=tokenizer,
-                group=None,
-                kv_prefix_cache=None,
-            )
-
-            prompt_a = apply_chat_template(tokenizer=tokenizer, task_params=task_a)
-            prompt_b = apply_chat_template(tokenizer=tokenizer, task_params=task_b)
-            uid_a = exo_gen.submit(task_params=task_a, prompt=prompt_a)
-            uid_b = exo_gen.submit(task_params=task_b, prompt=prompt_b)
-
-            batch_tokens: dict[int, list[int]] = {uid_a: [], uid_b: []}
-            finished: set[int] = set()
-            while exo_gen.has_work:
-                results = exo_gen.step()
-                for uid, response in results:
-                    batch_tokens[uid].append(response.token)
-                    if response.finish_reason is not None:
-                        finished.add(uid)
-
-            exo_gen.close()
-
-            # ── Verify both completed ──
-            assert len(batch_tokens[uid_a]) > 0, "No tokens for task A"
-            assert len(batch_tokens[uid_b]) > 0, "No tokens for task B"
-            assert uid_a in finished, "Task A never finished"
-            assert uid_b in finished, "Task B never finished"
-        finally:
-            shutil.rmtree(tmpdir, ignore_errors=True)

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

@@ -15,14 +15,14 @@ from typing import Any, cast
 
 import pytest
 
-from exo.shared.constants import EXO_MODELS_DIR
+from exo.shared.constants import EXO_DEFAULT_MODELS_DIR
 from exo.shared.models.model_cards import ModelCard, ModelTask
 from exo.shared.types.common import ModelId
 from exo.shared.types.memory import Memory
 from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
 
 MODEL_ID = "mlx-community/gpt-oss-20b-MXFP4-Q8"
-MODEL_PATH = EXO_MODELS_DIR / "mlx-community--gpt-oss-20b-MXFP4-Q8"
+MODEL_PATH = EXO_DEFAULT_MODELS_DIR / "mlx-community--gpt-oss-20b-MXFP4-Q8"
 TOTAL_LAYERS = 24
 MAX_TOKENS = 10
 SEED = 42

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

@@ -190,7 +190,7 @@ ARCHITECTURES: list[ArchSpec] = [
 
 def _arch_available(spec: ArchSpec) -> bool:
     snap = _find_snapshot(spec.hub_name)
-    if snap is None:
+    if snap is None or not (snap / "config.json").exists():
         return False
     if spec.tokenizer_hub is not None:
         return _find_snapshot(spec.tokenizer_hub) is not None

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

@@ -13,8 +13,8 @@ import pytest
 
 from exo.download.download_utils import (
     download_file_with_retry,
-    ensure_models_dir,
     fetch_file_list_with_cache,
+    resolve_model_dir,
 )
 from exo.shared.models.model_cards import ModelCard, ModelId, get_model_cards
 from exo.worker.engines.mlx.utils_mlx import (
@@ -53,8 +53,7 @@ def is_tokenizer_file(filename: str) -> bool:
 
 async def download_tokenizer_files(model_id: ModelId) -> Path:
     """Download only the tokenizer-related files for a model."""
-    target_dir = await ensure_models_dir() / model_id.normalize()
-    target_dir.mkdir(parents=True, exist_ok=True)
+    target_dir = await resolve_model_dir(model_id)
 
     file_list = await fetch_file_list_with_cache(model_id, "main", recursive=True)
 

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

@@ -2,6 +2,7 @@ import json
 from collections.abc import Generator
 from typing import Any
 
+from exo.shared.types.common import ModelId
 from exo.shared.types.worker.runner_response import (
     GenerationResponse,
     ToolCallResponse,
@@ -965,3 +966,70 @@ class TestE2EFullRoundTrip:
         assert "sunny" in final_text.lower()
         assert "5°C" in final_text
         assert "12°C" in final_text
+
+
+class TestMultiTurnThinkingPrompt:
+    def test_no_orphan_think_end_in_multiturn(self):
+        messages: list[dict[str, Any]] = [
+            {"role": "user", "content": "Hi!"},
+            {"role": "assistant", "content": "Hello! How can I help you today?"},
+            {"role": "user", "content": "Tell me about Paris."},
+        ]
+        prompt = encode_messages(messages, thinking_mode="thinking")
+        assistant_token = "<\uff5cAssistant\uff5c>"
+        parts = prompt.split(assistant_token)
+        for part in parts[1:]:
+            assert not part.startswith(THINKING_END), (
+                f"Orphan </think> without <think> after <Assistant>: ...{assistant_token}{part[:50]}"
+            )
+
+
+class TestApplyChatTemplateWithToolCalls:
+    def test_dsml_encoding_with_tool_calls_in_history(self):
+        from exo.shared.types.text_generation import (
+            InputMessage,
+            TextGenerationTaskParams,
+        )
+        from exo.worker.engines.mlx.utils_mlx import apply_chat_template
+
+        chat_template_messages: list[dict[str, Any]] = [
+            {"role": "system", "content": "You are a helpful assistant."},
+            {"role": "user", "content": "What's the weather?"},
+            {
+                "role": "assistant",
+                "content": "",
+                "tool_calls": [
+                    {
+                        "id": "call_1",
+                        "type": "function",
+                        "function": {
+                            "name": "get_weather",
+                            "arguments": '{"city": "Tokyo"}',
+                        },
+                    }
+                ],
+            },
+            {"role": "tool", "content": "Sunny, 25°C"},
+            {"role": "user", "content": "Thanks!"},
+        ]
+
+        from unittest.mock import MagicMock
+
+        tokenizer = MagicMock()
+        tokenizer.has_thinking = True
+        tokenizer.think_start = "<think>"
+        tokenizer.think_end = "</think>"
+
+        params = TextGenerationTaskParams(
+            model=ModelId("mlx-community/DeepSeek-V3.2-8bit"),
+            input=[InputMessage(role="user", content="Thanks!")],
+            instructions="You are a helpful assistant.",
+            enable_thinking=True,
+            chat_template_messages=chat_template_messages,
+            tools=_WEATHER_TOOLS,
+        )
+
+        prompt = apply_chat_template(tokenizer, params)
+        assert "get_weather" in prompt
+        assert "Tokyo" in prompt
+        assert "Sunny" in prompt

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

@@ -0,0 +1,332 @@
+from collections.abc import Generator
+from typing import Any
+
+from exo.shared.types.worker.runner_response import (
+    FinishReason,
+    GenerationResponse,
+    ToolCallResponse,
+)
+from exo.worker.engines.mlx.dsml_encoding import (
+    DSML_TOKEN,
+    THINKING_END,
+    THINKING_START,
+    TOOL_CALLS_END,
+    TOOL_CALLS_START,
+)
+from exo.worker.runner.llm_inference.model_output_parsers import (
+    parse_deepseek_v32,
+    parse_thinking_models,
+    parse_tool_calls,
+)
+from exo.worker.runner.llm_inference.tool_parsers import make_mlx_parser
+
+
+def _make_response(
+    text: str, token: int, finish_reason: FinishReason | None = None
+) -> GenerationResponse:
+    return GenerationResponse(
+        text=text, token=token, finish_reason=finish_reason, usage=None
+    )
+
+
+def _queue_source(
+    tokens: list[GenerationResponse],
+) -> Generator[GenerationResponse | None]:
+    for token in tokens:
+        yield token
+        yield None
+    while True:
+        yield None
+
+
+def _step_until_finish(
+    parser_gen: Generator[GenerationResponse | ToolCallResponse | None],
+    max_steps: int = 200,
+) -> list[GenerationResponse | ToolCallResponse]:
+    results: list[GenerationResponse | ToolCallResponse] = []
+    for _ in range(max_steps):
+        try:
+            result = next(parser_gen)
+        except StopIteration:
+            break
+        if result is None:
+            continue
+        results.append(result)
+        if isinstance(result, GenerationResponse) and result.finish_reason is not None:
+            return results
+        if isinstance(result, ToolCallResponse):
+            return results
+    return results
+
+
+def _got_finish(results: list[GenerationResponse | ToolCallResponse]) -> bool:
+    for r in results:
+        if isinstance(r, ToolCallResponse):
+            return True
+        if r.finish_reason is not None:
+            return True
+    return False
+
+
+# ── parse_deepseek_v32 ──────────────────────────────────────────
+
+
+class TestDeepSeekV32FinishReason:
+    def test_finish_reason_with_buffered_dsml_prefix(self):
+        tokens = [
+            _make_response("Hello! The answer is x", 0),
+            _make_response("<", 1),
+            _make_response("", 2, finish_reason="stop"),
+        ]
+        results = _step_until_finish(parse_deepseek_v32(_queue_source(tokens)))
+        assert _got_finish(results)
+        full_text = "".join(
+            r.text for r in results if isinstance(r, GenerationResponse)
+        )
+        assert "Hello" in full_text
+        assert "<" in full_text
+
+    def test_finish_reason_completes_tool_call_block(self):
+        tokens = [
+            _make_response(TOOL_CALLS_START, 0),
+            _make_response("\n", 1),
+            _make_response(f'<{DSML_TOKEN}invoke name="get_weather">\n', 2),
+            _make_response(
+                f'<{DSML_TOKEN}parameter name="city" string="true">Tokyo</{DSML_TOKEN}parameter>\n',
+                3,
+            ),
+            _make_response(f"</{DSML_TOKEN}invoke>\n", 4),
+            _make_response(TOOL_CALLS_END, 5, finish_reason="stop"),
+        ]
+        results = _step_until_finish(parse_deepseek_v32(_queue_source(tokens)))
+        tool_results = [r for r in results if isinstance(r, ToolCallResponse)]
+        assert len(tool_results) == 1
+        assert tool_results[0].tool_calls[0].name == "get_weather"
+
+    def test_finish_reason_mid_tool_call_before_close(self):
+        tokens = [
+            _make_response(TOOL_CALLS_START, 0),
+            _make_response("\n", 1),
+            _make_response(
+                f'<{DSML_TOKEN}invoke name="get_weather">\n', 2, finish_reason="stop"
+            ),
+        ]
+        results = _step_until_finish(parse_deepseek_v32(_queue_source(tokens)))
+        assert _got_finish(results)
+
+    def test_finish_reason_single_token_complete_dsml_block(self):
+        dsml_block = (
+            f"{TOOL_CALLS_START}\n"
+            f'<{DSML_TOKEN}invoke name="get_weather">\n'
+            f'<{DSML_TOKEN}parameter name="city" string="true">Tokyo</{DSML_TOKEN}parameter>\n'
+            f"</{DSML_TOKEN}invoke>\n"
+            f"{TOOL_CALLS_END}"
+        )
+        tokens = [_make_response(dsml_block, 0, finish_reason="stop")]
+        results = _step_until_finish(parse_deepseek_v32(_queue_source(tokens)))
+        tool_results = [r for r in results if isinstance(r, ToolCallResponse)]
+        assert len(tool_results) == 1
+        assert tool_results[0].tool_calls[0].name == "get_weather"
+
+    def test_finish_reason_during_thinking(self):
+        tokens = [
+            _make_response(THINKING_START, 0),
+            _make_response("I need to think about this", 1),
+            _make_response(" carefully before responding", 2, finish_reason="stop"),
+        ]
+        results = _step_until_finish(parse_deepseek_v32(_queue_source(tokens)))
+        assert _got_finish(results)
+
+    def test_finish_reason_after_thinking_then_tool_call(self):
+        tokens = [
+            _make_response(THINKING_START, 0),
+            _make_response("Let me check the weather.", 1),
+            _make_response(THINKING_END, 2),
+            _make_response("\n\n", 3),
+            _make_response(TOOL_CALLS_START, 4),
+            _make_response("\n", 5),
+            _make_response(f'<{DSML_TOKEN}invoke name="get_weather">\n', 6),
+            _make_response(
+                f'<{DSML_TOKEN}parameter name="city" string="true">NYC</{DSML_TOKEN}parameter>\n',
+                7,
+            ),
+            _make_response(f"</{DSML_TOKEN}invoke>\n", 8),
+            _make_response(TOOL_CALLS_END, 9, finish_reason="stop"),
+        ]
+        results = _step_until_finish(parse_deepseek_v32(_queue_source(tokens)))
+        tool_results = [r for r in results if isinstance(r, ToolCallResponse)]
+        assert len(tool_results) == 1
+        assert tool_results[0].tool_calls[0].name == "get_weather"
+
+    def test_finish_reason_normal_text_no_buffering(self):
+        tokens = [
+            _make_response("Hello", 0),
+            _make_response(" world", 1),
+            _make_response("!", 2, finish_reason="stop"),
+        ]
+        results = _step_until_finish(parse_deepseek_v32(_queue_source(tokens)))
+        assert _got_finish(results)
+        full_text = "".join(
+            r.text for r in results if isinstance(r, GenerationResponse)
+        )
+        assert full_text == "Hello world!"
+
+    def test_finish_reason_multiple_buffered_prefix_tokens(self):
+        tokens = [
+            _make_response("text ", 0),
+            _make_response("<", 1),
+            _make_response("not a tag", 2),
+            _make_response(" more<", 3),
+            _make_response("", 4, finish_reason="stop"),
+        ]
+        results = _step_until_finish(parse_deepseek_v32(_queue_source(tokens)))
+        assert _got_finish(results)
+
+
+# ── parse_thinking_models ────────────────────────────────────────
+
+
+class TestThinkingModelsFinishReason:
+    def test_finish_reason_during_thinking(self):
+        tokens = [
+            _make_response("<think>", 0),
+            _make_response("reasoning here", 1),
+            _make_response("more reasoning", 2, finish_reason="stop"),
+        ]
+        results = _step_until_finish(
+            parse_thinking_models(
+                _queue_source(tokens),
+                think_start="<think>",
+                think_end="</think>",
+                starts_in_thinking=False,
+            )
+        )
+        assert _got_finish(results)
+        last_gen = [
+            r
+            for r in results
+            if isinstance(r, GenerationResponse) and r.finish_reason is not None
+        ]
+        assert len(last_gen) == 1
+        assert last_gen[0].is_thinking is False
+
+    def test_finish_reason_after_thinking(self):
+        tokens = [
+            _make_response("<think>", 0),
+            _make_response("hmm", 1),
+            _make_response("</think>", 2),
+            _make_response("The answer is 42.", 3, finish_reason="stop"),
+        ]
+        results = _step_until_finish(
+            parse_thinking_models(
+                _queue_source(tokens),
+                think_start="<think>",
+                think_end="</think>",
+                starts_in_thinking=False,
+            )
+        )
+        assert _got_finish(results)
+
+    def test_finish_reason_starts_in_thinking(self):
+        tokens = [
+            _make_response("still thinking", 0),
+            _make_response("</think>", 1),
+            _make_response("done", 2, finish_reason="stop"),
+        ]
+        results = _step_until_finish(
+            parse_thinking_models(
+                _queue_source(tokens),
+                think_start="<think>",
+                think_end="</think>",
+                starts_in_thinking=True,
+            )
+        )
+        assert _got_finish(results)
+
+
+# ── parse_tool_calls (generic) ──────────────────────────────────
+
+
+def _dummy_parser_fn(text: str) -> dict[str, Any]:
+    return {"name": "test_fn", "arguments": {"arg": text}}
+
+
+_dummy_parser = make_mlx_parser("<tool_call>", "</tool_call>", _dummy_parser_fn)
+
+
+class TestGenericToolCallsFinishReason:
+    def test_finish_reason_after_complete_tool_call(self):
+        tokens = [
+            _make_response("<tool_call>", 0),
+            _make_response("body", 1),
+            _make_response("</tool_call>", 2),
+            _make_response("extra text", 3, finish_reason="stop"),
+        ]
+        results = _step_until_finish(
+            parse_tool_calls(
+                _queue_source(tokens),
+                _dummy_parser,
+                tools=None,
+            )
+        )
+        tool_results = [r for r in results if isinstance(r, ToolCallResponse)]
+        assert len(tool_results) == 1
+
+    def test_finish_reason_mid_tool_call_unclosed(self):
+        tokens = [
+            _make_response("<tool_call>", 0),
+            _make_response("partial content", 1, finish_reason="stop"),
+        ]
+        results = _step_until_finish(
+            parse_tool_calls(
+                _queue_source(tokens),
+                _dummy_parser,
+                tools=None,
+            )
+        )
+        assert _got_finish(results)
+
+    def test_finish_reason_no_tool_calls(self):
+        tokens = [
+            _make_response("Just", 0),
+            _make_response(" a", 1),
+            _make_response(" normal", 2),
+            _make_response(" response.", 3, finish_reason="stop"),
+        ]
+        results = _step_until_finish(
+            parse_tool_calls(
+                _queue_source(tokens),
+                _dummy_parser,
+                tools=None,
+            )
+        )
+        assert _got_finish(results)
+
+
+# ── Double parser chain (parse_thinking_models → parse_deepseek_v32) ──
+
+
+class TestBatchGeneratorSingleNext:
+    def test_finish_reason_with_buffered_tokens_drain_loop(self):
+        from exo.worker.runner.llm_inference.batch_generator import GeneratorQueue
+
+        queue: GeneratorQueue[GenerationResponse] = GeneratorQueue()
+        parser = parse_deepseek_v32(queue.gen())
+
+        tokens = [
+            _make_response("Hello ", 0),
+            _make_response(" `<", 1),
+            _make_response("", 2, finish_reason="stop"),
+        ]
+
+        collected: list[GenerationResponse | ToolCallResponse] = []
+        for token in tokens:
+            queue.push(token)
+            while (parsed := next(parser, None)) is not None:
+                collected.append(parsed)
+            if token.finish_reason is not None:
+                break
+
+        assert _got_finish(collected), (
+            f"No finish_reason in collected: {[(type(r).__name__, getattr(r, 'finish_reason', None) if isinstance(r, GenerationResponse) else 'tool') for r in collected]}"
+        )

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

@@ -87,6 +87,7 @@ class TestParseToolCalls:
         assert len(results) == 1
         assert isinstance(results[0], GenerationResponse)
         assert results[0].text == "<tool_call>bad content</tool_call>"
+        assert results[0].finish_reason == "error"
 
     def test_tool_schema_coerces_string_arguments_to_expected_types(self):
         """Tool argument values should be coerced using provided JSON schema."""

tests/auto_bench.sh

@@ -50,6 +50,6 @@ bench_runner="${hosts[0]}"
 mkdir -p "./bench/$commit"
 nix run .#exo-get-all-models-on-cluster -- "$bench_runner" | while IFS= read -r model; do
   echo "running bench for $model" 1>&2
-  ssh -Tn -o BatchMode=yes -o ServerAliveInterval=30 "$bench_runner@$bench_runner" "/nix/var/nix/profiles/default/bin/nix run github:exo-explore/exo/$commit#exo-bench -- --model $model --pp 128 4096 --tg 128 --stdout --skip-tensor-ring" >>"./bench/$commit/${model//\//--}.json"
+  ssh -Tn -o BatchMode=yes -o ServerAliveInterval=30 "$bench_runner@$bench_runner" "/nix/var/nix/profiles/default/bin/nix run github:exo-explore/exo/$commit#exo-bench -- --model $model --pp 128 4096 --tg 128 --concurrency 1 3 8 --stdout --skip-tensor-ring" >>"./bench/$commit/${model//\//--}.json"
   echo
 done

tests/headless_runner.py

@@ -9,7 +9,7 @@ from hypercorn.asyncio import serve  # pyright: ignore[reportUnknownVariableType
 from loguru import logger
 from pydantic import BaseModel
 
-from exo.shared.constants import EXO_MODELS_DIR
+from exo.shared.constants import EXO_DEFAULT_MODELS_DIR
 from exo.shared.models.model_cards import ModelCard, ModelId
 from exo.shared.types.chunks import TokenChunk
 from exo.shared.types.commands import CommandId
@@ -90,7 +90,7 @@ async def tb_detection():
 
 def list_models():
     sent = set[str]()
-    for path in EXO_MODELS_DIR.rglob("model-*.safetensors"):
+    for path in EXO_DEFAULT_MODELS_DIR.rglob("model-*.safetensors"):
         if "--" not in path.parent.name:
             continue
         name = path.parent.name.replace("--", "/")

tests/run_exo_on.sh

@@ -11,10 +11,17 @@ set -euo pipefail
   exit 1
 }
 
+upstream=$(git rev-parse --abbrev-ref --symbolic-full-name "@{u}" 2>/dev/null) || {
+  echo "No upstream"
+  exit 1
+}
 commit=$(git rev-parse HEAD)
-git fetch -q origin
-git branch -r --contains "$commit" | grep -qE '^\s*origin/' || {
-  echo "Not pushed to origin"
+remote=${upstream%%/*}
+remote_installable=$(git remote get-url "$remote" | sed -E "s#^(git@github.com:|https://github\.com/)([^/]+)/([^/]+)(\.git)?\$#github:\2/\3/$commit#")
+
+git fetch -q "$remote"
+git branch -r --contains "$commit" | grep -qE "^[[:space:]]*$remote/" || {
+  echo "Not pushed to $remote"
   exit 1
 }
 
@@ -35,7 +42,7 @@ i=0
 for host; do
   colour=${colours[i++ % 4]}
   ssh -T -o BatchMode=yes -o ServerAliveInterval=30 "$host@$host" \
-    "EXO_LIBP2P_NAMESPACE=$commit /nix/var/nix/profiles/default/bin/nix run github:exo-explore/exo/$commit" |&
+    "EXO_LIBP2P_NAMESPACE=$commit /nix/var/nix/profiles/default/bin/nix run $remote_installable" 2>&1 |
     awk -v p="${colour}[${host}]${reset}" '{ print p $0; fflush() }' &
 done
 

uv.lock

@@ -213,14 +213,20 @@ sdist = { url = "https://files.pythonhosted.org/packages/eb/56/b1ba7935a17738ae8
 wheels = [
     { url = "https://files.pythonhosted.org/packages/b0/1e/d22cc63332bd59b06481ceaac49d6c507598642e2230f201649058a7e704/cffi-2.0.0-cp313-cp313-manylinux1_i686.manylinux2014_i686.manylinux_2_17_i686.manylinux_2_5_i686.whl", hash = "sha256:07b271772c100085dd28b74fa0cd81c8fb1a3ba18b21e03d7c27f3436a10606b", size = 212446, upload-time = "2025-09-08T23:23:03.472Z" },
     { url = "https://files.pythonhosted.org/packages/a9/f5/a2c23eb03b61a0b8747f211eb716446c826ad66818ddc7810cc2cc19b3f2/cffi-2.0.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:d48a880098c96020b02d5a1f7d9251308510ce8858940e6fa99ece33f610838b", size = 220101, upload-time = "2025-09-08T23:23:04.792Z" },
+    { url = "https://files.pythonhosted.org/packages/f2/7f/e6647792fc5850d634695bc0e6ab4111ae88e89981d35ac269956605feba/cffi-2.0.0-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:f93fd8e5c8c0a4aa1f424d6173f14a892044054871c771f8566e4008eaa359d2", size = 207948, upload-time = "2025-09-08T23:23:06.127Z" },
+    { url = "https://files.pythonhosted.org/packages/cb/1e/a5a1bd6f1fb30f22573f76533de12a00bf274abcdc55c8edab639078abb6/cffi-2.0.0-cp313-cp313-manylinux2014_s390x.manylinux_2_17_s390x.whl", hash = "sha256:dd4f05f54a52fb558f1ba9f528228066954fee3ebe629fc1660d874d040ae5a3", size = 206422, upload-time = "2025-09-08T23:23:07.753Z" },
     { url = "https://files.pythonhosted.org/packages/98/df/0a1755e750013a2081e863e7cd37e0cdd02664372c754e5560099eb7aa44/cffi-2.0.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:c8d3b5532fc71b7a77c09192b4a5a200ea992702734a2e9279a37f2478236f26", size = 219499, upload-time = "2025-09-08T23:23:09.648Z" },
     { url = "https://files.pythonhosted.org/packages/50/e1/a969e687fcf9ea58e6e2a928ad5e2dd88cc12f6f0ab477e9971f2309b57c/cffi-2.0.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:d9b29c1f0ae438d5ee9acb31cadee00a58c46cc9c0b2f9038c6b0b3470877a8c", size = 222928, upload-time = "2025-09-08T23:23:10.928Z" },
     { url = "https://files.pythonhosted.org/packages/36/54/0362578dd2c9e557a28ac77698ed67323ed5b9775ca9d3fe73fe191bb5d8/cffi-2.0.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:6d50360be4546678fc1b79ffe7a66265e28667840010348dd69a314145807a1b", size = 221302, upload-time = "2025-09-08T23:23:12.42Z" },
     { url = "https://files.pythonhosted.org/packages/d6/43/0e822876f87ea8a4ef95442c3d766a06a51fc5298823f884ef87aaad168c/cffi-2.0.0-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:24b6f81f1983e6df8db3adc38562c83f7d4a0c36162885ec7f7b77c7dcbec97b", size = 220049, upload-time = "2025-09-08T23:23:20.853Z" },
+    { url = "https://files.pythonhosted.org/packages/b4/89/76799151d9c2d2d1ead63c2429da9ea9d7aac304603de0c6e8764e6e8e70/cffi-2.0.0-cp314-cp314-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:12873ca6cb9b0f0d3a0da705d6086fe911591737a59f28b7936bdfed27c0d47c", size = 207793, upload-time = "2025-09-08T23:23:22.08Z" },
+    { url = "https://files.pythonhosted.org/packages/bb/dd/3465b14bb9e24ee24cb88c9e3730f6de63111fffe513492bf8c808a3547e/cffi-2.0.0-cp314-cp314-manylinux2014_s390x.manylinux_2_17_s390x.whl", hash = "sha256:d9b97165e8aed9272a6bb17c01e3cc5871a594a446ebedc996e2397a1c1ea8ef", size = 206300, upload-time = "2025-09-08T23:23:23.314Z" },
     { url = "https://files.pythonhosted.org/packages/47/d9/d83e293854571c877a92da46fdec39158f8d7e68da75bf73581225d28e90/cffi-2.0.0-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:afb8db5439b81cf9c9d0c80404b60c3cc9c3add93e114dcae767f1477cb53775", size = 219244, upload-time = "2025-09-08T23:23:24.541Z" },
     { url = "https://files.pythonhosted.org/packages/2b/0f/1f177e3683aead2bb00f7679a16451d302c436b5cbf2505f0ea8146ef59e/cffi-2.0.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:737fe7d37e1a1bffe70bd5754ea763a62a066dc5913ca57e957824b72a85e205", size = 222828, upload-time = "2025-09-08T23:23:26.143Z" },
     { url = "https://files.pythonhosted.org/packages/c6/0f/cafacebd4b040e3119dcb32fed8bdef8dfe94da653155f9d0b9dc660166e/cffi-2.0.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:38100abb9d1b1435bc4cc340bb4489635dc2f0da7456590877030c9b3d40b0c1", size = 220926, upload-time = "2025-09-08T23:23:27.873Z" },
     { url = "https://files.pythonhosted.org/packages/be/b4/c56878d0d1755cf9caa54ba71e5d049479c52f9e4afc230f06822162ab2f/cffi-2.0.0-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:7cc09976e8b56f8cebd752f7113ad07752461f48a58cbba644139015ac24954c", size = 221593, upload-time = "2025-09-08T23:23:31.91Z" },
+    { url = "https://files.pythonhosted.org/packages/e0/0d/eb704606dfe8033e7128df5e90fee946bbcb64a04fcdaa97321309004000/cffi-2.0.0-cp314-cp314t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:92b68146a71df78564e4ef48af17551a5ddd142e5190cdf2c5624d0c3ff5b2e8", size = 209354, upload-time = "2025-09-08T23:23:33.214Z" },
+    { url = "https://files.pythonhosted.org/packages/d8/19/3c435d727b368ca475fb8742ab97c9cb13a0de600ce86f62eab7fa3eea60/cffi-2.0.0-cp314-cp314t-manylinux2014_s390x.manylinux_2_17_s390x.whl", hash = "sha256:b1e74d11748e7e98e2f426ab176d4ed720a64412b6a15054378afdb71e0f37dc", size = 208480, upload-time = "2025-09-08T23:23:34.495Z" },
     { url = "https://files.pythonhosted.org/packages/d0/44/681604464ed9541673e486521497406fadcc15b5217c3e326b061696899a/cffi-2.0.0-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:28a3a209b96630bca57cce802da70c266eb08c6e97e5afd61a75611ee6c64592", size = 221584, upload-time = "2025-09-08T23:23:36.096Z" },
     { url = "https://files.pythonhosted.org/packages/25/8e/342a504ff018a2825d395d44d63a767dd8ebc927ebda557fecdaca3ac33a/cffi-2.0.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:7553fb2090d71822f02c629afe6042c299edf91ba1bf94951165613553984512", size = 224443, upload-time = "2025-09-08T23:23:37.328Z" },
     { url = "https://files.pythonhosted.org/packages/e1/5e/b666bacbbc60fbf415ba9988324a132c9a7a0448a9a8f125074671c0f2c3/cffi-2.0.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:6c6c373cfc5c83a975506110d17457138c8c63016b563cc9ed6e056a82f13ce4", size = 223437, upload-time = "2025-09-08T23:23:38.945Z" },
@@ -344,8 +350,10 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/5c/49/498c86566a1d80e978b42f0d702795f69887005548c041636df6ae1ca64c/cryptography-46.0.3-cp311-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:01ca9ff2885f3acc98c29f1860552e37f6d7c7d013d7334ff2a9de43a449315d", size = 4450807, upload-time = "2025-10-15T23:16:56.414Z" },
     { url = "https://files.pythonhosted.org/packages/4b/0a/863a3604112174c8624a2ac3c038662d9e59970c7f926acdcfaed8d61142/cryptography-46.0.3-cp311-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:6eae65d4c3d33da080cff9c4ab1f711b15c1d9760809dad6ea763f3812d254cb", size = 4299615, upload-time = "2025-10-15T23:16:58.442Z" },
     { url = "https://files.pythonhosted.org/packages/64/02/b73a533f6b64a69f3cd3872acb6ebc12aef924d8d103133bb3ea750dc703/cryptography-46.0.3-cp311-abi3-manylinux_2_28_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:e5bf0ed4490068a2e72ac03d786693adeb909981cc596425d09032d372bcc849", size = 4016800, upload-time = "2025-10-15T23:17:00.378Z" },
+    { url = "https://files.pythonhosted.org/packages/25/d5/16e41afbfa450cde85a3b7ec599bebefaef16b5c6ba4ec49a3532336ed72/cryptography-46.0.3-cp311-abi3-manylinux_2_28_ppc64le.whl", hash = "sha256:5ecfccd2329e37e9b7112a888e76d9feca2347f12f37918facbb893d7bb88ee8", size = 4984707, upload-time = "2025-10-15T23:17:01.98Z" },
     { url = "https://files.pythonhosted.org/packages/c9/56/e7e69b427c3878352c2fb9b450bd0e19ed552753491d39d7d0a2f5226d41/cryptography-46.0.3-cp311-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:a2c0cd47381a3229c403062f764160d57d4d175e022c1df84e168c6251a22eec", size = 4482541, upload-time = "2025-10-15T23:17:04.078Z" },
     { url = "https://files.pythonhosted.org/packages/78/f6/50736d40d97e8483172f1bb6e698895b92a223dba513b0ca6f06b2365339/cryptography-46.0.3-cp311-abi3-manylinux_2_34_aarch64.whl", hash = "sha256:549e234ff32571b1f4076ac269fcce7a808d3bf98b76c8dd560e42dbc66d7d91", size = 4299464, upload-time = "2025-10-15T23:17:05.483Z" },
+    { url = "https://files.pythonhosted.org/packages/00/de/d8e26b1a855f19d9994a19c702fa2e93b0456beccbcfe437eda00e0701f2/cryptography-46.0.3-cp311-abi3-manylinux_2_34_ppc64le.whl", hash = "sha256:c0a7bb1a68a5d3471880e264621346c48665b3bf1c3759d682fc0864c540bd9e", size = 4950838, upload-time = "2025-10-15T23:17:07.425Z" },
     { url = "https://files.pythonhosted.org/packages/8f/29/798fc4ec461a1c9e9f735f2fc58741b0daae30688f41b2497dcbc9ed1355/cryptography-46.0.3-cp311-abi3-manylinux_2_34_x86_64.whl", hash = "sha256:10b01676fc208c3e6feeb25a8b83d81767e8059e1fe86e1dc62d10a3018fa926", size = 4481596, upload-time = "2025-10-15T23:17:09.343Z" },
     { url = "https://files.pythonhosted.org/packages/15/8d/03cd48b20a573adfff7652b76271078e3045b9f49387920e7f1f631d125e/cryptography-46.0.3-cp311-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:0abf1ffd6e57c67e92af68330d05760b7b7efb243aab8377e583284dbab72c71", size = 4426782, upload-time = "2025-10-15T23:17:11.22Z" },
     { url = "https://files.pythonhosted.org/packages/fa/b1/ebacbfe53317d55cf33165bda24c86523497a6881f339f9aae5c2e13e57b/cryptography-46.0.3-cp311-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:a04bee9ab6a4da801eb9b51f1b708a1b5b5c9eb48c03f74198464c66f0d344ac", size = 4698381, upload-time = "2025-10-15T23:17:12.829Z" },
@@ -353,8 +361,10 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/c5/fd/bc1daf8230eaa075184cbbf5f8cd00ba9db4fd32d63fb83da4671b72ed8a/cryptography-46.0.3-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:39b6755623145ad5eff1dab323f4eae2a32a77a7abef2c5089a04a3d04366715", size = 4435078, upload-time = "2025-10-15T23:17:23.042Z" },
     { url = "https://files.pythonhosted.org/packages/82/98/d3bd5407ce4c60017f8ff9e63ffee4200ab3e23fe05b765cab805a7db008/cryptography-46.0.3-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:db391fa7c66df6762ee3f00c95a89e6d428f4d60e7abc8328f4fe155b5ac6e54", size = 4293460, upload-time = "2025-10-15T23:17:24.885Z" },
     { url = "https://files.pythonhosted.org/packages/26/e9/e23e7900983c2b8af7a08098db406cf989d7f09caea7897e347598d4cd5b/cryptography-46.0.3-cp314-cp314t-manylinux_2_28_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:78a97cf6a8839a48c49271cdcbd5cf37ca2c1d6b7fdd86cc864f302b5e9bf459", size = 3995237, upload-time = "2025-10-15T23:17:26.449Z" },
+    { url = "https://files.pythonhosted.org/packages/91/15/af68c509d4a138cfe299d0d7ddb14afba15233223ebd933b4bbdbc7155d3/cryptography-46.0.3-cp314-cp314t-manylinux_2_28_ppc64le.whl", hash = "sha256:dfb781ff7eaa91a6f7fd41776ec37c5853c795d3b358d4896fdbb5df168af422", size = 4967344, upload-time = "2025-10-15T23:17:28.06Z" },
     { url = "https://files.pythonhosted.org/packages/ca/e3/8643d077c53868b681af077edf6b3cb58288b5423610f21c62aadcbe99f4/cryptography-46.0.3-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:6f61efb26e76c45c4a227835ddeae96d83624fb0d29eb5df5b96e14ed1a0afb7", size = 4466564, upload-time = "2025-10-15T23:17:29.665Z" },
     { url = "https://files.pythonhosted.org/packages/0e/43/c1e8726fa59c236ff477ff2b5dc071e54b21e5a1e51aa2cee1676f1c986f/cryptography-46.0.3-cp314-cp314t-manylinux_2_34_aarch64.whl", hash = "sha256:23b1a8f26e43f47ceb6d6a43115f33a5a37d57df4ea0ca295b780ae8546e8044", size = 4292415, upload-time = "2025-10-15T23:17:31.686Z" },
+    { url = "https://files.pythonhosted.org/packages/42/f9/2f8fefdb1aee8a8e3256a0568cffc4e6d517b256a2fe97a029b3f1b9fe7e/cryptography-46.0.3-cp314-cp314t-manylinux_2_34_ppc64le.whl", hash = "sha256:b419ae593c86b87014b9be7396b385491ad7f320bde96826d0dd174459e54665", size = 4931457, upload-time = "2025-10-15T23:17:33.478Z" },
     { url = "https://files.pythonhosted.org/packages/79/30/9b54127a9a778ccd6d27c3da7563e9f2d341826075ceab89ae3b41bf5be2/cryptography-46.0.3-cp314-cp314t-manylinux_2_34_x86_64.whl", hash = "sha256:50fc3343ac490c6b08c0cf0d704e881d0d660be923fd3076db3e932007e726e3", size = 4466074, upload-time = "2025-10-15T23:17:35.158Z" },
     { url = "https://files.pythonhosted.org/packages/ac/68/b4f4a10928e26c941b1b6a179143af9f4d27d88fe84a6a3c53592d2e76bf/cryptography-46.0.3-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:22d7e97932f511d6b0b04f2bfd818d73dcd5928db509460aaf48384778eb6d20", size = 4420569, upload-time = "2025-10-15T23:17:37.188Z" },
     { url = "https://files.pythonhosted.org/packages/a3/49/3746dab4c0d1979888f125226357d3262a6dd40e114ac29e3d2abdf1ec55/cryptography-46.0.3-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:d55f3dffadd674514ad19451161118fd010988540cee43d8bc20675e775925de", size = 4681941, upload-time = "2025-10-15T23:17:39.236Z" },
@@ -362,8 +372,10 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/26/42/fa8389d4478368743e24e61eea78846a0006caffaf72ea24a15159215a14/cryptography-46.0.3-cp38-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:15ab9b093e8f09daab0f2159bb7e47532596075139dd74365da52ecc9cb46c5d", size = 4440029, upload-time = "2025-10-15T23:17:49.837Z" },
     { url = "https://files.pythonhosted.org/packages/5f/eb/f483db0ec5ac040824f269e93dd2bd8a21ecd1027e77ad7bdf6914f2fd80/cryptography-46.0.3-cp38-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:46acf53b40ea38f9c6c229599a4a13f0d46a6c3fa9ef19fc1a124d62e338dfa0", size = 4297222, upload-time = "2025-10-15T23:17:51.357Z" },
     { url = "https://files.pythonhosted.org/packages/fd/cf/da9502c4e1912cb1da3807ea3618a6829bee8207456fbbeebc361ec38ba3/cryptography-46.0.3-cp38-abi3-manylinux_2_28_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:10ca84c4668d066a9878890047f03546f3ae0a6b8b39b697457b7757aaf18dbc", size = 4012280, upload-time = "2025-10-15T23:17:52.964Z" },
+    { url = "https://files.pythonhosted.org/packages/6b/8f/9adb86b93330e0df8b3dcf03eae67c33ba89958fc2e03862ef1ac2b42465/cryptography-46.0.3-cp38-abi3-manylinux_2_28_ppc64le.whl", hash = "sha256:36e627112085bb3b81b19fed209c05ce2a52ee8b15d161b7c643a7d5a88491f3", size = 4978958, upload-time = "2025-10-15T23:17:54.965Z" },
     { url = "https://files.pythonhosted.org/packages/d1/a0/5fa77988289c34bdb9f913f5606ecc9ada1adb5ae870bd0d1054a7021cc4/cryptography-46.0.3-cp38-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:1000713389b75c449a6e979ffc7dcc8ac90b437048766cef052d4d30b8220971", size = 4473714, upload-time = "2025-10-15T23:17:56.754Z" },
     { url = "https://files.pythonhosted.org/packages/14/e5/fc82d72a58d41c393697aa18c9abe5ae1214ff6f2a5c18ac470f92777895/cryptography-46.0.3-cp38-abi3-manylinux_2_34_aarch64.whl", hash = "sha256:b02cf04496f6576afffef5ddd04a0cb7d49cf6be16a9059d793a30b035f6b6ac", size = 4296970, upload-time = "2025-10-15T23:17:58.588Z" },
+    { url = "https://files.pythonhosted.org/packages/78/06/5663ed35438d0b09056973994f1aec467492b33bd31da36e468b01ec1097/cryptography-46.0.3-cp38-abi3-manylinux_2_34_ppc64le.whl", hash = "sha256:71e842ec9bc7abf543b47cf86b9a743baa95f4677d22baa4c7d5c69e49e9bc04", size = 4940236, upload-time = "2025-10-15T23:18:00.897Z" },
     { url = "https://files.pythonhosted.org/packages/fc/59/873633f3f2dcd8a053b8dd1d38f783043b5fce589c0f6988bf55ef57e43e/cryptography-46.0.3-cp38-abi3-manylinux_2_34_x86_64.whl", hash = "sha256:402b58fc32614f00980b66d6e56a5b4118e6cb362ae8f3fda141ba4689bd4506", size = 4472642, upload-time = "2025-10-15T23:18:02.749Z" },
     { url = "https://files.pythonhosted.org/packages/3d/39/8e71f3930e40f6877737d6f69248cf74d4e34b886a3967d32f919cc50d3b/cryptography-46.0.3-cp38-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:ef639cb3372f69ec44915fafcd6698b6cc78fbe0c2ea41be867f6ed612811963", size = 4423126, upload-time = "2025-10-15T23:18:04.85Z" },
     { url = "https://files.pythonhosted.org/packages/cd/c7/f65027c2810e14c3e7268353b1681932b87e5a48e65505d8cc17c99e36ae/cryptography-46.0.3-cp38-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:3b51b8ca4f1c6453d8829e1eb7299499ca7f313900dd4d89a24b8b87c0a780d4", size = 4686573, upload-time = "2025-10-15T23:18:06.908Z" },
@@ -473,7 +485,7 @@ dependencies = [
     { name = "loguru", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "mflux", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "mlx", version = "0.30.6", source = { registry = "https://pypi.org/simple" }, extra = ["cpu"], marker = "sys_platform == 'linux'" },
-    { name = "mlx", version = "0.30.7.dev20260303+257d5692", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#257d5692fc7af6bba3b8afaeb63c549b7d1e43d5" }, marker = "sys_platform == 'darwin'" },
+    { name = "mlx", version = "0.31.2.dev20260324+e5e64331", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#e5e64331830d9b04ae9082b843073f9c1fa7705e" }, marker = "sys_platform == 'darwin'" },
     { name = "mlx-lm", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "msgspec", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "openai-harmony", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
@@ -509,10 +521,10 @@ requires-dist = [
     { name = "huggingface-hub", specifier = ">=0.33.4" },
     { name = "hypercorn", specifier = ">=0.18.0" },
     { name = "loguru", specifier = ">=0.7.3" },
-    { name = "mflux", specifier = "==0.16.9" },
+    { name = "mflux", specifier = "==0.17.2" },
     { name = "mlx", marker = "sys_platform == 'darwin'", git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks" },
     { name = "mlx", extras = ["cpu"], marker = "sys_platform == 'linux'", specifier = "==0.30.6" },
-    { name = "mlx-lm", git = "https://github.com/rltakashige/mlx-lm?branch=leo%2Feval-left-padding-in-batched-rotation" },
+    { name = "mlx-lm", git = "https://github.com/rltakashige/mlx-lm?branch=leo%2Ffix-deepseek-v32-indexer" },
     { name = "msgspec", specifier = ">=0.19.0" },
     { name = "openai-harmony", specifier = ">=0.0.8" },
     { name = "psutil", specifier = ">=7.0.0" },
@@ -1342,7 +1354,7 @@ wheels = [
 
 [[package]]
 name = "mflux"
-version = "0.16.9"
+version = "0.17.2"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "filelock", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
@@ -1351,12 +1363,13 @@ dependencies = [
     { name = "huggingface-hub", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "matplotlib", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "mlx", version = "0.30.6", source = { registry = "https://pypi.org/simple" }, extra = ["cuda13"], marker = "sys_platform == 'linux'" },
-    { name = "mlx", version = "0.30.7.dev20260303+257d5692", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#257d5692fc7af6bba3b8afaeb63c549b7d1e43d5" }, marker = "sys_platform == 'darwin'" },
+    { name = "mlx", version = "0.31.2.dev20260324+e5e64331", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#e5e64331830d9b04ae9082b843073f9c1fa7705e" }, marker = "sys_platform == 'darwin'" },
     { name = "numpy", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "opencv-python", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "piexif", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "pillow", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "platformdirs", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
+    { name = "protobuf", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "regex", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "requests", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "safetensors", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
@@ -1369,9 +1382,9 @@ dependencies = [
     { name = "twine", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "urllib3", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/e3/dc/73dfccae395df22623fbb6f748657700fd9372be22904a0dafb825dc42d2/mflux-0.16.9.tar.gz", hash = "sha256:ff35e9386b5f026a6b97c786b3426e5341105de3356e55ec1b0c5951ff0ac8c8", size = 764296, upload-time = "2026-03-07T11:54:45.031Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/1b/75/65f791f54c9531b524813d18ff87a0d34866ef220a8fe2ad637437f3cb54/mflux-0.17.2.tar.gz", hash = "sha256:52dee2d27cf438a84648e5c1861b92ceb63a9ac06823d14452a78646a1d30ee7", size = 779264, upload-time = "2026-03-23T13:08:18.377Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/47/69/091df2d0f4a2677d5e2a8b326d568bd884d7ca37f15d3ff78e574aca2271/mflux-0.16.9-py3-none-any.whl", hash = "sha256:988e925653a024e81b46401fd5a2e423dd283428ba64949af4d9257e7df8ddd9", size = 1018804, upload-time = "2026-03-07T11:54:43.47Z" },
+    { url = "https://files.pythonhosted.org/packages/68/02/f94eca4e77b7d12685060461eb793cbc8c00e96cc7fe0ce376374201aed2/mflux-0.17.2-py3-none-any.whl", hash = "sha256:be1642b04847413c0a8ed1dae82ce1ca023e155b057d82a8301eca9c3fe08339", size = 1037451, upload-time = "2026-03-23T13:08:16.747Z" },
 ]
 
 [[package]]
@@ -1399,8 +1412,8 @@ cuda13 = [
 
 [[package]]
 name = "mlx"
-version = "0.30.7.dev20260303+257d5692"
-source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#257d5692fc7af6bba3b8afaeb63c549b7d1e43d5" }
+version = "0.31.2.dev20260324+e5e64331"
+source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#e5e64331830d9b04ae9082b843073f9c1fa7705e" }
 resolution-markers = [
     "python_full_version >= '3.14' and sys_platform == 'darwin'",
     "python_full_version < '3.14' and sys_platform == 'darwin'",
@@ -1432,11 +1445,11 @@ wheels = [
 
 [[package]]
 name = "mlx-lm"
-version = "0.31.0"
-source = { git = "https://github.com/rltakashige/mlx-lm?branch=leo%2Feval-left-padding-in-batched-rotation#5e0c484cfb5c68d281a71409927ee1bb75adaae2" }
+version = "0.31.2"
+source = { git = "https://github.com/rltakashige/mlx-lm?branch=leo%2Ffix-deepseek-v32-indexer#d388ff77858fec3b5d2e3b1d9502a7e2878b8109" }
 dependencies = [
     { name = "jinja2", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
-    { name = "mlx", version = "0.30.7.dev20260303+257d5692", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#257d5692fc7af6bba3b8afaeb63c549b7d1e43d5" }, marker = "sys_platform == 'darwin'" },
+    { name = "mlx", version = "0.31.2.dev20260324+e5e64331", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=address-rdma-gpu-locks#e5e64331830d9b04ae9082b843073f9c1fa7705e" }, marker = "sys_platform == 'darwin'" },
     { name = "numpy", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "protobuf", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "pyyaml", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
@@ -2091,15 +2104,14 @@ wheels = [
 
 [[package]]
 name = "protobuf"
-version = "6.33.3"
+version = "5.29.6"
 source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/cc/5c/f912bdebdd4af4160da6a2c2b1b3aaa1b8c578d0243ba8f694f93c7095f0/protobuf-6.33.3.tar.gz", hash = "sha256:c8794debeb402963fddff41a595e1f649bcd76616ba56c835645cab4539e810e", size = 444318, upload-time = "2026-01-09T23:05:02.79Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/7e/57/394a763c103e0edf87f0938dafcd918d53b4c011dfc5c8ae80f3b0452dbb/protobuf-5.29.6.tar.gz", hash = "sha256:da9ee6a5424b6b30fd5e45c5ea663aef540ca95f9ad99d1e887e819cdf9b8723", size = 425623, upload-time = "2026-02-04T22:54:40.584Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/ec/5d/0ef28dded98973a26443a6a7bc49bff6206be8c57dc1d1e28e6c1147b879/protobuf-6.33.3-cp39-abi3-macosx_10_9_universal2.whl", hash = "sha256:648b7b0144222eb06cf529a3d7b01333c5f30b4196773b682d388f04db373759", size = 427594, upload-time = "2026-01-09T23:04:53.358Z" },
-    { url = "https://files.pythonhosted.org/packages/c5/46/551c69b6ff1957bd703654342bfb776bb97db400bc80afc56fbb64e7c11d/protobuf-6.33.3-cp39-abi3-manylinux2014_aarch64.whl", hash = "sha256:08a6ca12f60ba99097dd3625ef4275280f99c9037990e47ce9368826b159b890", size = 324469, upload-time = "2026-01-09T23:04:54.332Z" },
-    { url = "https://files.pythonhosted.org/packages/ca/6d/ade1cca06c64a421ee9745e082671465ead28164c809efaf2c15bc93f9a0/protobuf-6.33.3-cp39-abi3-manylinux2014_s390x.whl", hash = "sha256:642fce7187526c98683c79a3ad68e5d646a5ef5eb004582fe123fc9a33a9456b", size = 339242, upload-time = "2026-01-09T23:04:55.347Z" },
-    { url = "https://files.pythonhosted.org/packages/38/8c/6522b8e543ece46f645911c3cebe361d8460134c0fee02ddcf70ebf32999/protobuf-6.33.3-cp39-abi3-manylinux2014_x86_64.whl", hash = "sha256:6fa9b5f4baa12257542273e5e6f3c3d3867b30bc2770c14ad9ac8315264bf986", size = 323298, upload-time = "2026-01-09T23:04:56.866Z" },
-    { url = "https://files.pythonhosted.org/packages/a6/b9/067b8a843569d5605ba6f7c039b9319720a974f82216cd623e13186d3078/protobuf-6.33.3-py3-none-any.whl", hash = "sha256:c2bf221076b0d463551efa2e1319f08d4cffcc5f0d864614ccd3d0e77a637794", size = 170518, upload-time = "2026-01-09T23:05:01.227Z" },
+    { url = "https://files.pythonhosted.org/packages/06/db/49b05966fd208ae3f44dcd33837b6243b4915c57561d730a43f881f24dea/protobuf-5.29.6-cp38-abi3-macosx_10_9_universal2.whl", hash = "sha256:b5a169e664b4057183a34bdc424540e86eea47560f3c123a0d64de4e137f9269", size = 418619, upload-time = "2026-02-04T22:54:30.266Z" },
+    { url = "https://files.pythonhosted.org/packages/b7/d7/48cbf6b0c3c39761e47a99cb483405f0fde2be22cf00d71ef316ce52b458/protobuf-5.29.6-cp38-abi3-manylinux2014_aarch64.whl", hash = "sha256:a8866b2cff111f0f863c1b3b9e7572dc7eaea23a7fae27f6fc613304046483e6", size = 320284, upload-time = "2026-02-04T22:54:31.782Z" },
+    { url = "https://files.pythonhosted.org/packages/e3/dd/cadd6ec43069247d91f6345fa7a0d2858bef6af366dbd7ba8f05d2c77d3b/protobuf-5.29.6-cp38-abi3-manylinux2014_x86_64.whl", hash = "sha256:e3387f44798ac1106af0233c04fb8abf543772ff241169946f698b3a9a3d3ab9", size = 320478, upload-time = "2026-02-04T22:54:32.909Z" },
+    { url = "https://files.pythonhosted.org/packages/5a/cb/e3065b447186cb70aa65acc70c86baf482d82bf75625bf5a2c4f6919c6a3/protobuf-5.29.6-py3-none-any.whl", hash = "sha256:6b9edb641441b2da9fa8f428760fc136a49cf97a52076010cf22a2ff73438a86", size = 173126, upload-time = "2026-02-04T22:54:39.462Z" },
 ]
 
 [[package]]