Fix local model launch preparation handling

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,

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)
 

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."""
 

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: