fix race condition with downloads where it cancels the download before renaming

src/exo/engines/mlx/auto_parallel.py

@@ -1,4 +1,4 @@
-from typing import cast, override, Protocol, TYPE_CHECKING
+from typing import TYPE_CHECKING, Protocol, cast, override
 
 import mlx.core as mx
 import mlx.nn as nn  # pyright: ignore[reportMissingTypeStubs]

src/exo/engines/mlx/utils_mlx.py

@@ -3,10 +3,10 @@ import concurrent.futures
 import contextlib
 import os
 import resource
-from loguru import logger
 from asyncio import AbstractEventLoop
 from typing import Any, Callable, Optional, cast
 
+from loguru import logger
 from mlx_lm.models.cache import KVCache
 from mlx_lm.sample_utils import make_sampler
 from mlx_lm.tokenizer_utils import TokenizerWrapper as _TokenizerWrapper

src/exo/shared/constants.py

@@ -4,6 +4,10 @@ from pathlib import Path
 
 EXO_HOME_RELATIVE_PATH = os.environ.get("EXO_HOME", ".exo")
 EXO_HOME = Path.home() / EXO_HOME_RELATIVE_PATH
+
+EXO_MODELS_DIR_ENV = os.environ.get("EXO_MODELS_DIR")
+EXO_MODELS_DIR = Path(EXO_MODELS_DIR_ENV) if EXO_MODELS_DIR_ENV else EXO_HOME / "models"
+
 EXO_GLOBAL_EVENT_DB = EXO_HOME / "global_events.db"
 EXO_WORKER_EVENT_DB = EXO_HOME / "worker_events.db"
 EXO_MASTER_STATE = EXO_HOME / "master_state.json"

src/exo/shared/models/model_meta.py

@@ -58,8 +58,8 @@ async def get_config_data(model_id: str) -> ConfigData:
         "main",
         "config.json",
         target_dir,
-        lambda curr_bytes, total_bytes: logger.info(
-            f"Downloading config.json for {model_id}: {curr_bytes}/{total_bytes}"
+        lambda curr_bytes, total_bytes, is_renamed: logger.info(
+            f"Downloading config.json for {model_id}: {curr_bytes}/{total_bytes} ({is_renamed=})"
         ),
     )
     async with aiofiles.open(config_path, "r") as f:
@@ -75,8 +75,8 @@ async def get_safetensors_size(model_id: str) -> Memory:
         "main",
         "model.safetensors.index.json",
         target_dir,
-        lambda curr_bytes, total_bytes: logger.info(
-            f"Downloading model.safetensors.index.json for {model_id}: {curr_bytes}/{total_bytes}"
+        lambda curr_bytes, total_bytes, is_renamed: logger.info(
+            f"Downloading model.safetensors.index.json for {model_id}: {curr_bytes}/{total_bytes} ({is_renamed=})"
         ),
     )
     async with aiofiles.open(index_path, "r") as f:

src/exo/worker/download/download_utils.py

@@ -22,7 +22,7 @@ from pydantic import (
     TypeAdapter,
 )
 
-from exo.shared.constants import EXO_HOME
+from exo.shared.constants import EXO_HOME, EXO_MODELS_DIR
 from exo.shared.types.memory import Memory
 from exo.shared.types.worker.downloads import DownloadProgressData
 from exo.shared.types.worker.shards import ShardMetadata
@@ -123,7 +123,7 @@ def map_repo_download_progress_to_download_progress_data(
 
 
 def build_model_path(model_id: str) -> DirectoryPath:
-    return EXO_HOME / "models" / model_id.replace("/", "--")
+    return EXO_MODELS_DIR / model_id.replace("/", "--")
 
 
 async def resolve_model_path_for_repo(repo_id: str) -> Path:
@@ -150,9 +150,8 @@ async def has_exo_home_write_access() -> bool:
 
 
 async def ensure_models_dir() -> Path:
-    models_dir = EXO_HOME / "models"
-    await aios.makedirs(models_dir, exist_ok=True)
-    return models_dir
+    await aios.makedirs(EXO_MODELS_DIR, exist_ok=True)
+    return EXO_MODELS_DIR
 
 
 async def delete_model(repo_id: str) -> bool:
@@ -324,7 +323,7 @@ async def download_file_with_retry(
     revision: str,
     path: str,
     target_dir: Path,
-    on_progress: Callable[[int, int], None] = lambda _, __: None,
+    on_progress: Callable[[int, int, bool], None] = lambda _, __, ___: None,
 ) -> Path:
     n_attempts = 30
     for attempt in range(n_attempts):
@@ -350,7 +349,7 @@ async def _download_file(
     revision: str,
     path: str,
     target_dir: Path,
-    on_progress: Callable[[int, int], None] = lambda _, __: None,
+    on_progress: Callable[[int, int, bool], None] = lambda _, __, ___: None,
 ) -> Path:
     if await aios.path.exists(target_dir / path):
         return target_dir / path
@@ -383,7 +382,7 @@ async def _download_file(
             ) as f:
                 while chunk := await r.content.read(8 * 1024 * 1024):
                     n_read = n_read + (await f.write(chunk))
-                    on_progress(n_read, length)
+                    on_progress(n_read, length, False)
 
     final_hash = await calc_hash(
         partial_path, hash_type="sha256" if len(remote_hash) == 64 else "sha1"
@@ -398,6 +397,7 @@ async def _download_file(
             f"Downloaded file {target_dir / path} has hash {final_hash} but remote hash is {remote_hash}"
         )
     await aios.rename(partial_path, target_dir / path)
+    on_progress(length, length, True)
     return target_dir / path
 
 
@@ -570,7 +570,9 @@ async def download_shard(
     )
     file_progress: Dict[str, RepoFileDownloadProgress] = {}
 
-    def on_progress_wrapper(file: FileListEntry, curr_bytes: int, total_bytes: int):
+    def on_progress_wrapper(
+        file: FileListEntry, curr_bytes: int, total_bytes: int, is_renamed: bool
+    ):
         start_time = (
             file_progress[file.path].start_time
             if file.path in file_progress
@@ -601,7 +603,9 @@ async def download_shard(
             total=Memory.from_bytes(total_bytes),
             speed=speed,
             eta=eta,
-            status="complete" if curr_bytes == total_bytes else "in_progress",
+            status="complete"
+            if curr_bytes == total_bytes and is_renamed
+            else "in_progress",
             start_time=start_time,
         )
         on_progress(
@@ -639,8 +643,8 @@ async def download_shard(
                 revision,
                 file.path,
                 target_dir,
-                lambda curr_bytes, total_bytes: on_progress_wrapper(
-                    file, curr_bytes, total_bytes
+                lambda curr_bytes, total_bytes, is_renamed: on_progress_wrapper(
+                    file, curr_bytes, total_bytes, is_renamed
                 ),
             )
 

src/exo/worker/main.py

@@ -1,8 +1,8 @@
 import asyncio
+import time
 from asyncio import Queue
 from functools import partial
 from random import random
-import time
 from typing import AsyncGenerator, Optional
 
 import anyio
@@ -199,7 +199,7 @@ class Worker:
                     await self.event_publisher(event)
             except Exception as e:
                 logger.opt(exception=e).warning(
-                    f"Error occurred when executing task", flush=True
+                    "Error occurred when executing task", flush=True
                 )
 
                 if isinstance(op, ExecuteTaskOp):

src/exo/worker/runner/generate.py

@@ -30,7 +30,6 @@ from exo.shared.types.worker.communication import (
     runner_print,
 )
 
-
 generation_stream = mx.new_stream(mx.default_device())