← back to Exo
Rework model storage directory management (for external storage) (#1765)
15f1b61f4cc70af1e46e9ce7769d95b49166d110 · 2026-03-26 17:46:46 +0000 · ciaranbor
## Motivation
Replace confusing EXO_MODELS_DIR/EXO_MODELS_PATH with clearer
multi-directory support, enabling automatic download spillover across
volumes.
## Changes
- EXO_MODELS_DIRS: colon-separated writable dirs (default always
prepended, first with enough space wins)
- EXO_MODELS_READ_ONLY_DIRS: colon-separated read-only dirs (protected
from deletion)
- select_download_dir(): picks writable dir by free space
- resolve_existing_model(): unified lookup across all dirs
- is_read_only_model_dir(): path-based read-only detection instead of
hardcoded flag
- Updated coordinator, worker, model cards, tests
## Why It Works
Default dir always included so zero-config behavior is unchanged. Disk
space checked at download time for automatic spillover. Read-only status
derived from path, not hardcoded.
## Test Plan
### Manual Testing
- No env vars set → identical behavior
- EXO_MODELS_DIRS=/Volumes/SSD/models → downloads to external storage
- EXO_MODELS_READ_ONLY_DIRS=/mnt/nfs → models found, deletion blocked
### Automated Testing
- 4 new tests in test_xdg_paths.py (prepend, default-only, overlap,
empty read-only)
- Existing tests updated to patch new constants
Files touched
M README.mdM bench/harness.pyM src/exo/download/coordinator.pyM src/exo/download/download_utils.pyM src/exo/download/tests/test_download_verification.pyA src/exo/download/tests/test_model_dirs.pyM src/exo/download/tests/test_offline_mode.pyM src/exo/shared/constants.pyM src/exo/shared/models/model_cards.pyM src/exo/shared/tests/test_xdg_paths.pyM src/exo/utils/info_gatherer/info_gatherer.pyM src/exo/worker/main.pyM src/exo/worker/tests/unittests/test_mlx/conftest.pyM src/exo/worker/tests/unittests/test_mlx/test_pipeline_prefill_callbacks.pyM src/exo/worker/tests/unittests/test_mlx/test_tokenizers.pyM tests/headless_runner.py
Diff
commit 15f1b61f4cc70af1e46e9ce7769d95b49166d110
Author: ciaranbor <81697641+ciaranbor@users.noreply.github.com>
Date: Thu Mar 26 17:46:46 2026 +0000
Rework model storage directory management (for external storage) (#1765)
## Motivation
Replace confusing EXO_MODELS_DIR/EXO_MODELS_PATH with clearer
multi-directory support, enabling automatic download spillover across
volumes.
## Changes
- EXO_MODELS_DIRS: colon-separated writable dirs (default always
prepended, first with enough space wins)
- EXO_MODELS_READ_ONLY_DIRS: colon-separated read-only dirs (protected
from deletion)
- select_download_dir(): picks writable dir by free space
- resolve_existing_model(): unified lookup across all dirs
- is_read_only_model_dir(): path-based read-only detection instead of
hardcoded flag
- Updated coordinator, worker, model cards, tests
## Why It Works
Default dir always included so zero-config behavior is unchanged. Disk
space checked at download time for automatic spillover. Read-only status
derived from path, not hardcoded.
## Test Plan
### Manual Testing
- No env vars set → identical behavior
- EXO_MODELS_DIRS=/Volumes/SSD/models → downloads to external storage
- EXO_MODELS_READ_ONLY_DIRS=/mnt/nfs → models found, deletion blocked
### Automated Testing
- 4 new tests in test_xdg_paths.py (prepend, default-only, overlap,
empty read-only)
- Existing tests updated to patch new constants
---
README.md | 12 +-
bench/harness.py | 2 +-
src/exo/download/coordinator.py | 148 +++++-----
src/exo/download/download_utils.py | 153 +++++++----
.../download/tests/test_download_verification.py | 34 +--
src/exo/download/tests/test_model_dirs.py | 297 +++++++++++++++++++++
src/exo/download/tests/test_offline_mode.py | 5 +-
src/exo/shared/constants.py | 38 ++-
src/exo/shared/models/model_cards.py | 10 +-
src/exo/shared/tests/test_xdg_paths.py | 110 +++++++-
src/exo/utils/info_gatherer/info_gatherer.py | 4 +-
src/exo/worker/main.py | 14 +-
.../worker/tests/unittests/test_mlx/conftest.py | 4 +-
.../test_mlx/test_pipeline_prefill_callbacks.py | 4 +-
.../tests/unittests/test_mlx/test_tokenizers.py | 5 +-
tests/headless_runner.py | 4 +-
16 files changed, 667 insertions(+), 177 deletions(-)
diff --git a/README.md b/README.md
index 5104f14f..f12a1c53 100644
--- a/README.md
+++ b/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
diff --git a/bench/harness.py b/bench/harness.py
index 366de660..3d771de7 100644
--- a/bench/harness.py
+++ b/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"][
diff --git a/src/exo/download/coordinator.py b/src/exo/download/coordinator.py
index 06ff54ee..a7a77064 100644
--- a/src/exo/download/coordinator.py
+++ b/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,
@@ -57,8 +62,23 @@ class DownloadCoordinator:
def __post_init__(self) -> None:
self.shard_downloader.on_progress(self._download_progress_callback)
- def _model_dir(self, model_id: ModelId) -> str:
- return str(EXO_MODELS_DIR / model_id.normalize())
+ @staticmethod
+ def _default_model_dir(model_id: ModelId) -> str:
+ return str(EXO_DEFAULT_MODELS_DIR / model_id.normalize())
+
+ def _completed_from_path(
+ self,
+ shard: ShardMetadata,
+ found: Path,
+ total: Memory,
+ ) -> DownloadCompleted:
+ return DownloadCompleted(
+ shard_metadata=shard,
+ node_id=self.node_id,
+ total=total,
+ model_directory=str(found),
+ read_only=is_read_only_model_dir(found),
+ )
async def _download_progress_callback(
self, callback_shard: ShardMetadata, progress: RepoDownloadProgress
@@ -67,12 +87,18 @@ class DownloadCoordinator:
throttle_interval_secs = 1.0
if progress.status == "complete":
- completed = DownloadCompleted(
- shard_metadata=callback_shard,
- node_id=self.node_id,
- total=progress.total,
- model_directory=self._model_dir(model_id),
- )
+ found = await to_thread.run_sync(resolve_existing_model, model_id)
+ if found is not None:
+ completed = self._completed_from_path(
+ callback_shard, found, progress.total
+ )
+ else:
+ completed = DownloadCompleted(
+ shard_metadata=callback_shard,
+ node_id=self.node_id,
+ total=progress.total,
+ model_directory=self._default_model_dir(model_id),
+ )
self.download_status[model_id] = completed
await self.event_sender.send(
NodeDownloadProgress(download_progress=completed)
@@ -89,7 +115,7 @@ class DownloadCoordinator:
download_progress=map_repo_download_progress_to_download_progress_data(
progress
),
- model_directory=self._model_dir(model_id),
+ model_directory=self._default_model_dir(model_id),
)
self.download_status[model_id] = ongoing
await self.event_sender.send(
@@ -135,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(
@@ -154,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(
@@ -177,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))
@@ -188,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)
@@ -208,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))
@@ -229,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))
@@ -244,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(
@@ -261,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
@@ -290,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)
@@ -314,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,
)
@@ -340,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
@@ -351,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:
@@ -362,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,
@@ -372,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(
diff --git a/src/exo/download/download_utils.py b/src/exo/download/download_utils.py
index 86c092ab..818efca3 100644
--- a/src/exo/download/download_utils.py
+++ b/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."""
+
+
+def resolve_existing_model(model_id: ModelId) -> Path | None:
+ """Search all model directories for a complete, pre-existing model.
- 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.
+ 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.
-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()
+ 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
- deleted = False
- if await aios.path.exists(model_dir):
- await asyncio.to_thread(shutil.rmtree, model_dir, ignore_errors=False)
- deleted = True
- # Also clear cache
+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:
+ """Delete a model from writable directories. Skips read-only dirs."""
+ normalized = model_id.normalize()
+ deleted = False
+ 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
+
+ # 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,7 +805,7 @@ 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,
@@ -785,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(
@@ -821,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),
@@ -849,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),
@@ -875,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,
@@ -891,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):
diff --git a/src/exo/download/tests/test_download_verification.py b/src/exo/download/tests/test_download_verification.py
index 2d8d076d..2e04e3af 100644
--- a/src/exo/download/tests/test_download_verification.py
+++ b/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
diff --git a/src/exo/download/tests/test_model_dirs.py b/src/exo/download/tests/test_model_dirs.py
new file mode 100644
index 00000000..4064add7
--- /dev/null
+++ b/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
diff --git a/src/exo/download/tests/test_offline_mode.py b/src/exo/download/tests/test_offline_mode.py
index 15210c3f..9a94e520 100644
--- a/src/exo/download/tests/test_offline_mode.py
+++ b/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
diff --git a/src/exo/shared/constants.py b/src/exo/shared/constants.py
index 69534298..e6870a73 100644
--- a/src/exo/shared/constants.py
+++ b/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 = (
diff --git a/src/exo/shared/models/model_cards.py b/src/exo/shared/models/model_cards.py
index fef1efaa..7cc5a33c 100644
--- a/src/exo/shared/models/model_cards.py
+++ b/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",
diff --git a/src/exo/shared/tests/test_xdg_paths.py b/src/exo/shared/tests/test_xdg_paths.py
index 73f81d10..f3b82ebf 100644
--- a/src/exo/shared/tests/test_xdg_paths.py
+++ b/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 == ()
diff --git a/src/exo/utils/info_gatherer/info_gatherer.py b/src/exo/utils/info_gatherer/info_gatherer.py
index a79f9f70..9e75e143 100644
--- a/src/exo/utils/info_gatherer/info_gatherer.py
+++ b/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
)
)
diff --git a/src/exo/worker/main.py b/src/exo/worker/main.py
index 833f9134..30128594 100644
--- a/src/exo/worker/main.py
+++ b/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 (
@@ -181,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(
@@ -193,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),
)
)
)
diff --git a/src/exo/worker/tests/unittests/test_mlx/conftest.py b/src/exo/worker/tests/unittests/test_mlx/conftest.py
index a6bfdfc1..c09bfe1e 100644
--- a/src/exo/worker/tests/unittests/test_mlx/conftest.py
+++ b/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,
diff --git a/src/exo/worker/tests/unittests/test_mlx/test_pipeline_prefill_callbacks.py b/src/exo/worker/tests/unittests/test_mlx/test_pipeline_prefill_callbacks.py
index 89f1b66a..be723380 100644
--- a/src/exo/worker/tests/unittests/test_mlx/test_pipeline_prefill_callbacks.py
+++ b/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
diff --git a/src/exo/worker/tests/unittests/test_mlx/test_tokenizers.py b/src/exo/worker/tests/unittests/test_mlx/test_tokenizers.py
index ce06ab86..d6e3774b 100644
--- a/src/exo/worker/tests/unittests/test_mlx/test_tokenizers.py
+++ b/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)
diff --git a/tests/headless_runner.py b/tests/headless_runner.py
index 56fb2632..176f6fcf 100644
--- a/tests/headless_runner.py
+++ b/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("--", "/")
← 90343001 [Fix] Node hang on reelection (#1801)
·
back to Exo
·
Fix custom model add requiring two attempts + enlarge sideba 5327bdde →