From 781428a17671415d71b0af51b73acb3fdcfa0fcb Mon Sep 17 00:00:00 2001 From: Alex Cheema Date: Tue, 24 Mar 2026 16:54:01 -0700 Subject: [PATCH] Fix local model launch preparation handling --- src/exo/download/coordinator.py | 18 ++++++ src/exo/download/download_utils.py | 11 ++++ .../tests/test_download_verification.py | 64 +++++++++++++++++++ src/exo/download/tests/test_re_download.py | 51 +++++++++++++++ 4 files changed, 144 insertions(+) diff --git a/src/exo/download/coordinator.py b/src/exo/download/coordinator.py index 5c55970f..c5e409fc 100644 --- a/src/exo/download/coordinator.py +++ b/src/exo/download/coordinator.py @@ -7,6 +7,7 @@ from loguru import logger from exo.download.download_utils import ( RepoDownloadProgress, delete_model, + is_model_directory_complete, map_repo_download_progress_to_download_progress_data, resolve_model_in_path, ) @@ -168,6 +169,23 @@ class DownloadCoordinator: ) return + local_model_dir = EXO_MODELS_DIR / model_id.normalize() + if local_model_dir.is_dir() and is_model_directory_complete(local_model_dir): + logger.info( + f"DownloadCoordinator: Model {model_id} already complete at {local_model_dir}" + ) + completed = DownloadCompleted( + shard_metadata=shard, + node_id=self.node_id, + total=shard.model_card.storage_size, + model_directory=str(local_model_dir), + ) + self.download_status[model_id] = completed + await self.event_sender.send( + NodeDownloadProgress(download_progress=completed) + ) + return + # Emit pending status progress = DownloadPending( shard_metadata=shard, diff --git a/src/exo/download/download_utils.py b/src/exo/download/download_utils.py index 3f6f1dc9..1ea128d7 100644 --- a/src/exo/download/download_utils.py +++ b/src/exo/download/download_utils.py @@ -239,6 +239,17 @@ def _scan_model_directory( def is_model_directory_complete(model_dir: Path) -> bool: """Check if a model directory contains all required weight files.""" + index_files = list(model_dir.glob("**/*.safetensors.index.json")) + if not index_files: + return False + + for index_file in index_files: + try: + ModelSafetensorsIndex.model_validate_json(index_file.read_text()) + except Exception: + logger.warning(f"Failed to parse model index {index_file}") + return False + file_list = _scan_model_directory(model_dir, recursive=True) return file_list is not None and all(f.size is not None for f in file_list) diff --git a/src/exo/download/tests/test_download_verification.py b/src/exo/download/tests/test_download_verification.py index 2d8d076d..0db92071 100644 --- a/src/exo/download/tests/test_download_verification.py +++ b/src/exo/download/tests/test_download_verification.py @@ -14,6 +14,7 @@ from pydantic import TypeAdapter from exo.download.download_utils import ( delete_model, fetch_file_list_with_cache, + is_model_directory_complete, ) from exo.shared.types.common import ModelId from exo.shared.types.memory import Memory @@ -324,6 +325,69 @@ class TestModelDeletion: assert result is False +class TestModelDirectoryComplete: + """Tests for local completeness checks used to skip download probing.""" + + async def test_returns_false_when_only_partial_weight_exists( + self, model_id: ModelId, tmp_path: Path + ) -> None: + import json + + model_dir = tmp_path / model_id.normalize() + await aios.makedirs(model_dir, exist_ok=True) + + async with aiofiles.open(model_dir / "model.safetensors.index.json", "w") as f: + await f.write( + json.dumps( + { + "metadata": {"total_size": 256}, + "weight_map": {"model.layers.0.weight": "model.safetensors"}, + } + ) + ) + async with aiofiles.open(model_dir / "model.safetensors.partial", "wb") as f: + await f.write(b"x" * 128) + + assert is_model_directory_complete(model_dir) is False + + async def test_returns_true_when_final_weight_exists_even_with_stale_partial( + self, model_id: ModelId, tmp_path: Path + ) -> None: + import json + + model_dir = tmp_path / model_id.normalize() + await aios.makedirs(model_dir, exist_ok=True) + + async with aiofiles.open(model_dir / "model.safetensors.index.json", "w") as f: + await f.write( + json.dumps( + { + "metadata": {"total_size": 256}, + "weight_map": {"model.layers.0.weight": "model.safetensors"}, + } + ) + ) + async with aiofiles.open(model_dir / "model.safetensors", "wb") as f: + await f.write(b"x" * 256) + async with aiofiles.open(model_dir / "model.safetensors.partial", "wb") as f: + await f.write(b"x" * 128) + + assert is_model_directory_complete(model_dir) is True + + async def test_returns_false_when_index_is_invalid( + self, model_id: ModelId, tmp_path: Path + ) -> None: + model_dir = tmp_path / model_id.normalize() + await aios.makedirs(model_dir, exist_ok=True) + + async with aiofiles.open(model_dir / "model.safetensors.index.json", "w") as f: + await f.write("{not valid json") + async with aiofiles.open(model_dir / "model.safetensors", "wb") as f: + await f.write(b"x" * 256) + + assert is_model_directory_complete(model_dir) is False + + class TestProgressResetOnRedownload: """Tests for progress tracking when files are re-downloaded.""" diff --git a/src/exo/download/tests/test_re_download.py b/src/exo/download/tests/test_re_download.py index 3f159814..320a36ee 100644 --- a/src/exo/download/tests/test_re_download.py +++ b/src/exo/download/tests/test_re_download.py @@ -2,12 +2,16 @@ import asyncio import contextlib +import json from collections.abc import AsyncIterator, Awaitable from datetime import timedelta from pathlib import Path from typing import Callable from unittest.mock import AsyncMock, patch +import aiofiles +import aiofiles.os as aios + from exo.download.coordinator import DownloadCoordinator from exo.download.download_utils import RepoDownloadProgress from exo.download.impl_shard_downloader import SingletonShardDownloader @@ -192,6 +196,53 @@ async def test_re_download_after_delete_completes() -> None: await coordinator_task +async def test_start_download_uses_complete_local_model_without_status_probe( + tmp_path: Path, +) -> None: + _, cmd_recv = channel[ForwarderDownloadCommand]() + event_send, event_recv = channel[Event]() + + fake_downloader = FakeShardDownloader() + fake_downloader.get_shard_download_status_for_shard = AsyncMock( + side_effect=AssertionError("status probe should be skipped for complete models") + ) + + coordinator = DownloadCoordinator( + node_id=NODE_ID, + shard_downloader=SingletonShardDownloader(fake_downloader), + download_command_receiver=cmd_recv, + event_sender=event_send, + ) + shard = _make_shard() + + models_dir = tmp_path / "models" + model_dir = models_dir / MODEL_ID.normalize() + await aios.makedirs(model_dir, exist_ok=True) + + async with aiofiles.open(model_dir / "config.json", "w") as f: + await f.write('{"model_type":"qwen2"}') + async with aiofiles.open(model_dir / "model.safetensors.index.json", "w") as f: + await f.write( + json.dumps( + { + "metadata": {"total_size": 128}, + "weight_map": {"model.layers.0.weight": "model.safetensors"}, + } + ) + ) + async with aiofiles.open(model_dir / "model.safetensors", "wb") as f: + await f.write(b"x" * 128) + + with patch("exo.download.coordinator.EXO_MODELS_DIR", models_dir): + await coordinator._start_download(shard) # pyright: ignore[reportPrivateUsage] + + event = await event_recv.receive() + assert isinstance(event, NodeDownloadProgress) + assert isinstance(event.download_progress, DownloadCompleted) + assert event.download_progress.model_directory == str(model_dir) + assert event.download_progress.total == shard.model_card.storage_size + + async def _wait_for_download_completed( event_recv: Receiver[Event], model_id: ModelId, timeout: float = 2.0 ) -> DownloadCompleted | None: