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
 

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

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/FamilySidebar.svelte

@@ -42,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
     >
@@ -90,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
       >
@@ -115,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
       >
@@ -139,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
     >
@@ -165,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/stores/app.svelte.ts

@@ -2650,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);
@@ -2704,6 +2707,7 @@ class AppStore {
           "Content-Type": "application/json",
         },
         body: JSON.stringify(requestBody),
+        signal: abortController.signal,
       });
 
       if (!response.ok) {
@@ -2843,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();
     }
@@ -2914,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);
@@ -2975,6 +2995,7 @@ class AppStore {
       const apiResponse = await fetch("/v1/images/edits", {
         method: "POST",
         body: formData,
+        signal: abortController.signal,
       });
 
       if (!apiResponse.ok) {
@@ -3075,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();
     }

src/exo/api/main.py

@@ -129,6 +129,7 @@ from exo.shared.logging import InterceptLogger
 from exo.shared.models.model_cards import (
     ModelCard,
     ModelId,
+    add_to_card_cache,
     get_card,
     get_model_cards,
 )
@@ -414,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
@@ -476,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:
@@ -1585,6 +1588,10 @@ class API:
             )
         )
 
+        # 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,

src/exo/download/coordinator.py

@@ -1,17 +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_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,
@@ -24,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,
@@ -49,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)
@@ -56,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
@@ -66,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)
@@ -88,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(
@@ -100,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:
@@ -130,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(
@@ -149,18 +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,
+            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(
@@ -172,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))
@@ -183,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)
@@ -203,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))
@@ -224,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))
@@ -239,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(
@@ -256,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
@@ -285,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)
@@ -309,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,
                             )
@@ -335,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
@@ -346,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:
@@ -357,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,
@@ -367,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
@@ -253,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()
@@ -273,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}"
 
@@ -329,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
@@ -658,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,
@@ -730,23 +772,19 @@ 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()
     try:
         file_list = await fetch_file_list_with_cache(
-            shard.model_card.model_id,
+            model_id,
             revision,
             recursive=True,
             skip_internet=skip_internet,
@@ -754,7 +792,7 @@ async def download_shard(
         )
     except FileNotFoundError:
         not_started_progress = RepoDownloadProgress(
-            repo_id=str(shard.model_card.model_id),
+            repo_id=str(model_id),
             repo_revision=revision,
             shard=shard,
             completed_files=0,
@@ -767,10 +805,13 @@ async def download_shard(
             status="not_started",
             file_progress={},
         )
-        return target_dir, not_started_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,
         )
     )
 
@@ -782,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(
@@ -818,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),
@@ -846,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),
@@ -872,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,
@@ -888,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
@@ -25,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."""
 
@@ -188,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,
@@ -234,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,
@@ -252,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,
@@ -284,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
@@ -303,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
@@ -318,7 +317,10 @@ 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

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

@@ -186,7 +186,7 @@ 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

src/exo/main.py

@@ -228,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),
@@ -240,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,

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/models/model_cards.py

@@ -243,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",
@@ -265,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/utils/info_gatherer/info_gatherer.py

@@ -13,7 +13,7 @@ from anyio.streams.buffered import BufferedByteReceiveStream
 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,
@@ -328,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
             )
         )
 
@@ -371,20 +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 _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:
@@ -425,25 +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:
-                        logger.warning(
-                            f"macmon at {macmon_path} is unusable, falling back to psutil memory monitoring"
-                        )
-                else:
-                    logger.warning(
-                        "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:
@@ -452,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 as e:
                 logger.opt(exception=e).warning("Error gathering static node info")
-            await anyio.sleep(self.static_info_poll_interval)
+            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 as e:
                 logger.opt(exception=e).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
+            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):
@@ -499,17 +473,18 @@ class InfoGatherer:
                     await self.info_sender.send(MacThunderboltConnections(conns=conns))
             except Exception as e:
                 logger.opt(exception=e).warning("Error gathering Thunderbolt data")
-            await anyio.sleep(self.system_profiler_interval)
+            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
             if override_memory_env
             else None
         )
-        if self.memory_poll_rate is None:
-            return
         while True:
             try:
                 await self.info_sender.send(
@@ -517,11 +492,9 @@ class InfoGatherer:
                 )
             except Exception as e:
                 logger.opt(exception=e).warning("Error gathering memory usage")
-            await anyio.sleep(self.memory_poll_rate)
+            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):
@@ -529,11 +502,11 @@ class InfoGatherer:
                     await self.info_sender.send(NodeNetworkInterfaces(ifaces=nics))
             except Exception as e:
                 logger.opt(exception=e).warning("Error gathering network interfaces")
-            await anyio.sleep(self.interface_watcher_interval)
+            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):
@@ -544,11 +517,9 @@ class InfoGatherer:
                 logger.opt(exception=e).warning(
                     "Error gathering Thunderbolt Bridge status"
                 )
-            await anyio.sleep(self.thunderbolt_bridge_poll_interval)
+            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()
@@ -556,26 +527,36 @@ class InfoGatherer:
                     await self.info_sender.send(curr)
             except Exception as e:
                 logger.opt(exception=e).warning("Error gathering RDMA ctl status")
-            await anyio.sleep(self.rdma_ctl_poll_interval)
+            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 as e:
                 logger.opt(exception=e).warning("Error gathering disk usage")
-            await anyio.sleep(self.disk_poll_interval)
+            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(
@@ -583,7 +564,7 @@ class InfoGatherer:
                         macmon_path,
                         "pipe",
                         "--interval",
-                        str(self.macmon_interval * 1000),
+                        str(macmon_interval * 1000),
                     ]
                 ) as p:
                     if not p.stdout:
@@ -602,8 +583,7 @@ class InfoGatherer:
                 logger.warning(
                     f"MacMon produced no output for {read_timeout}s, restarting"
                 )
-                self.memory_poll_rate = 1
-                self._tg.start_soon(self._monitor_memory_usage)
+                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)
@@ -616,10 +596,8 @@ class InfoGatherer:
                 logger.warning(
                     f"MacMon failed with return code {e.returncode}: {stderr_msg}"
                 )
-                self.memory_poll_rate = 1
-                self._tg.start_soon(self._monitor_memory_usage)
+                self._tg.start_soon(self._monitor_memory_usage, 1)
             except Exception as e:
                 logger.opt(exception=e).warning("Error in macmon monitor")
-                self.memory_poll_rate = 1
-                self._tg.start_soon(self._monitor_memory_usage)
-            await anyio.sleep(self.macmon_interval)
+                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/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/main.py

@@ -2,11 +2,11 @@ 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, add_to_card_cache, delete_custom_card
 from exo.shared.types.commands import (
@@ -80,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")
@@ -102,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:
@@ -179,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(
@@ -191,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),
                                 )
                             )
                         )
@@ -280,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/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/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_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_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)
 

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("--", "/")