← back to Exo
fix race condition with downloads where it cancels the download before renaming
d46c7e6a76bf342ef7dd38d59e6f9e863899b097 · 2025-10-31 02:03:23 +0000 · Alex Cheema
Files touched
M src/exo/engines/mlx/auto_parallel.pyM src/exo/engines/mlx/utils_mlx.pyM src/exo/shared/constants.pyM src/exo/shared/models/model_meta.pyM src/exo/worker/download/download_utils.pyM src/exo/worker/main.pyM src/exo/worker/runner/generate.py
Diff
commit d46c7e6a76bf342ef7dd38d59e6f9e863899b097
Author: Alex Cheema <41707476+AlexCheema@users.noreply.github.com>
Date: Fri Oct 31 02:03:23 2025 +0000
fix race condition with downloads where it cancels the download before renaming
---
src/exo/engines/mlx/auto_parallel.py | 2 +-
src/exo/engines/mlx/utils_mlx.py | 2 +-
src/exo/shared/constants.py | 4 ++++
src/exo/shared/models/model_meta.py | 8 ++++----
src/exo/worker/download/download_utils.py | 28 ++++++++++++++++------------
src/exo/worker/main.py | 4 ++--
src/exo/worker/runner/generate.py | 1 -
7 files changed, 28 insertions(+), 21 deletions(-)
diff --git a/src/exo/engines/mlx/auto_parallel.py b/src/exo/engines/mlx/auto_parallel.py
index 293a4da5..625d37dc 100644
--- a/src/exo/engines/mlx/auto_parallel.py
+++ b/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]
diff --git a/src/exo/engines/mlx/utils_mlx.py b/src/exo/engines/mlx/utils_mlx.py
index b7b97ac3..bef55c66 100644
--- a/src/exo/engines/mlx/utils_mlx.py
+++ b/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
diff --git a/src/exo/shared/constants.py b/src/exo/shared/constants.py
index 2961c686..489b871a 100644
--- a/src/exo/shared/constants.py
+++ b/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"
diff --git a/src/exo/shared/models/model_meta.py b/src/exo/shared/models/model_meta.py
index 9ed1f151..b9bb470a 100644
--- a/src/exo/shared/models/model_meta.py
+++ b/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:
diff --git a/src/exo/worker/download/download_utils.py b/src/exo/worker/download/download_utils.py
index 217da9a4..2a7f6cf1 100644
--- a/src/exo/worker/download/download_utils.py
+++ b/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
),
)
diff --git a/src/exo/worker/main.py b/src/exo/worker/main.py
index 3cc66f5d..e4374dd5 100644
--- a/src/exo/worker/main.py
+++ b/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):
diff --git a/src/exo/worker/runner/generate.py b/src/exo/worker/runner/generate.py
index e8e15c96..d0bbe700 100644
--- a/src/exo/worker/runner/generate.py
+++ b/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())
← 91c635ca Update mlx and mlx-lm packages
·
back to Exo
·
Squash merge merging_clusters into tensor_parallel94 3b409647 →