Fix local model launch preparation handling

CONTRIBUTING.md

@@ -11,9 +11,18 @@ To run EXO from source:
   ```bash
   brew install uv
   ```
-- [macmon](https://github.com/vladkens/macmon) (for hardware monitoring on Apple Silicon)
+- [rust](https://github.com/rust-lang/rustup) (to build Rust bindings, nightly for now)
   ```bash
-  brew install macmon
+  curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh
+  rustup toolchain install nightly
+  ```
+- [macmon](https://github.com/vladkens/macmon) (for hardware monitoring on Apple Silicon)
+  Use the pinned fork revision used by this repo instead of Homebrew `macmon`.
+  ```bash
+  cargo install --git https://github.com/swiftraccoon/macmon \
+    --rev 9154d234f763fbeffdcb4135d0bbbaf80609699b \
+    macmon \
+    --force
   ```
 
 ```bash

README.md

@@ -95,11 +95,10 @@ Then restart the Nix daemon: `sudo launchctl kickstart -k system/org.nixos.nix-d
   /bin/bash -c "$(curl -fsSL https://raw.githubusercontent.com/Homebrew/install/HEAD/install.sh)"
   ```
 - [uv](https://github.com/astral-sh/uv) (for Python dependency management)
-- [macmon](https://github.com/vladkens/macmon) (for hardware monitoring on Apple Silicon)
 - [node](https://github.com/nodejs/node) (for building the dashboard)
 
   ```bash
-  brew install uv macmon node
+  brew install uv node
   ```
 - [rust](https://github.com/rust-lang/rustup) (to build Rust bindings, nightly for now)
 
@@ -107,6 +106,17 @@ Then restart the Nix daemon: `sudo launchctl kickstart -k system/org.nixos.nix-d
   curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh
   rustup toolchain install nightly
   ```
+- [macmon](https://github.com/vladkens/macmon) (for hardware monitoring on Apple Silicon)
+
+  Install the pinned fork revision used by this repo instead of Homebrew `macmon`.
+  Homebrew `macmon 0.6.1` still crashes on Apple M5.
+
+  ```bash
+  cargo install --git https://github.com/swiftraccoon/macmon \
+    --rev 9154d234f763fbeffdcb4135d0bbbaf80609699b \
+    macmon \
+    --force
+  ```
 
 Clone the repo, build the dashboard, and run exo:
 

flake.nix

@@ -72,7 +72,7 @@
       ];
 
       perSystem =
-        { config, self', inputs', pkgs, lib, system, ... }:
+        { config, self', pkgs, lib, system, ... }:
         let
           # Use pinned nixpkgs for swift-format (swift is broken on x86_64-linux in newer nixpkgs)
           pkgsSwift = import inputs.nixpkgs-swift { inherit system; };
@@ -84,6 +84,17 @@
             config.allowUnfreePredicate = pkg: (pkg.pname or "") == "metal-toolchain";
             overlays = [
               (import ./nix/apple-sdk-overlay.nix)
+              (final: prev: {
+                macmon = prev.macmon.overrideAttrs (_: {
+                  version = "git";
+                  src = final.fetchFromGitHub {
+                    owner = "swiftraccoon";
+                    repo = "macmon";
+                    rev = "9154d234f763fbeffdcb4135d0bbbaf80609699b";
+                    hash = "sha256-CwhilKNbs5XL9/tF5DMwyPBlE/hpmjGNTuxQ36sM50M=";
+                  };
+                });
+              })
             ];
           };
           treefmt = {

packaging/pyinstaller/exo.spec

@@ -71,7 +71,9 @@ MACMON_PATH = shutil.which("macmon")
 if MACMON_PATH is None:
     raise SystemExit(
         "macmon binary not found in PATH. "
-        "Install it via: brew install macmon"
+        "Install the pinned fork used by exo via: "
+        "cargo install --git https://github.com/swiftraccoon/macmon "
+        "--rev 9154d234f763fbeffdcb4135d0bbbaf80609699b macmon --force"
     )
 
 BINARIES: list[tuple[str, str]] = [
@@ -120,4 +122,3 @@ coll = COLLECT(
     upx_exclude=[],
     name="exo",
 )
-

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:

src/exo/utils/info_gatherer/info_gatherer.py

@@ -287,8 +287,8 @@ class ThunderboltBridgeInfo(TaggedModel):
                     service_name=tb_service_name,
                 )
             )
-        except Exception as e:
-            logger.warning(f"Failed to gather Thunderbolt Bridge info: {e}")
+        except Exception:
+            logger.opt(exception=True).warning("Failed to gather Thunderbolt Bridge info")
             return None
 
 
@@ -382,18 +382,69 @@ class InfoGatherer:
     rdma_ctl_poll_interval: float | None = 10 if IS_DARWIN else None
     disk_poll_interval: float | None = 30
     _tg: TaskGroup = field(init=False, default_factory=TaskGroup)
+    _psutil_memory_fallback_enabled: bool = field(init=False, default=False)
+
+    def _enable_psutil_memory_fallback(self, reason: str) -> None:
+        if not self._psutil_memory_fallback_enabled:
+            logger.warning(reason)
+            self._psutil_memory_fallback_enabled = True
+        self.memory_poll_rate = 1
+
+    def _get_macmon_path(self) -> str | None:
+        return os.getenv("EXO_MACMON_PATH") or shutil.which("macmon")
+
+    async def _can_read_macmon_metrics(self, macmon_path: str) -> bool:
+        try:
+            with fail_after(5):
+                proc = await anyio.run_process(
+                    [macmon_path, "pipe", "--samples", "1", "--interval", "100"],
+                    check=False,
+                )
+        except Exception:
+            logger.opt(exception=True).warning(
+                f"Failed to validate macmon at {macmon_path}"
+            )
+            return False
+
+        if proc.returncode != 0:
+            stderr = proc.stderr.decode("utf-8", errors="replace").strip()
+            logger.warning(
+                f"macmon preflight failed with return code {proc.returncode}: "
+                f"{stderr or 'no stderr'}"
+            )
+            return False
+
+        stdout = proc.stdout.decode("utf-8", errors="replace").strip()
+        if not stdout:
+            logger.warning("macmon preflight returned no metrics")
+            return False
+
+        try:
+            MacmonMetrics.from_raw_json(stdout.splitlines()[0])
+        except ValidationError:
+            logger.opt(exception=True).warning(
+                "macmon preflight returned unexpected metrics JSON"
+            )
+            return False
+
+        return True
 
     async def run(self):
         async with self._tg as tg:
             if IS_DARWIN:
-                if (macmon_path := shutil.which("macmon")) is not None:
-                    tg.start_soon(self._monitor_macmon, macmon_path)
+                if (macmon_path := self._get_macmon_path()) is not None:
+                    if await self._can_read_macmon_metrics(macmon_path):
+                        tg.start_soon(self._monitor_macmon, macmon_path)
+                    else:
+                        self._enable_psutil_memory_fallback(
+                            f"macmon at {macmon_path} is unusable, falling back "
+                            f"to psutil memory monitoring"
+                        )
                 else:
-                    # macmon not installed — fall back to psutil for memory
-                    logger.warning(
-                        "macmon not found, falling back to psutil for memory monitoring"
+                    self._enable_psutil_memory_fallback(
+                        "macmon not found, falling back to psutil for memory "
+                        "monitoring"
                     )
-                    self.memory_poll_rate = 1
                 tg.start_soon(self._monitor_system_profiler_thunderbolt_data)
                 tg.start_soon(self._monitor_thunderbolt_bridge_status)
                 tg.start_soon(self._monitor_rdma_ctl_status)
@@ -417,8 +468,8 @@ class InfoGatherer:
             try:
                 with fail_after(30):
                     await self.info_sender.send(await StaticNodeInformation.gather())
-            except Exception as e:
-                logger.warning(f"Error gathering static node info: {e}")
+            except Exception:
+                logger.opt(exception=True).warning("Error gathering static node info")
             await anyio.sleep(self.static_info_poll_interval)
 
     async def _monitor_misc(self):
@@ -428,8 +479,8 @@ class InfoGatherer:
             try:
                 with fail_after(10):
                     await self.info_sender.send(await MiscData.gather())
-            except Exception as e:
-                logger.warning(f"Error gathering misc data: {e}")
+            except Exception:
+                logger.opt(exception=True).warning("Error gathering misc data")
             await anyio.sleep(self.misc_poll_interval)
 
     async def _monitor_system_profiler_thunderbolt_data(self):
@@ -455,8 +506,8 @@ class InfoGatherer:
 
                     conns = [it for i in data if (it := i.conn()) is not None]
                     await self.info_sender.send(MacThunderboltConnections(conns=conns))
-            except Exception as e:
-                logger.warning(f"Error gathering Thunderbolt data: {e}")
+            except Exception:
+                logger.opt(exception=True).warning("Error gathering Thunderbolt data")
             await anyio.sleep(self.system_profiler_interval)
 
     async def _monitor_memory_usage(self):
@@ -466,16 +517,18 @@ class InfoGatherer:
             if override_memory_env
             else None
         )
-        if self.memory_poll_rate is None:
-            return
         while True:
+            poll_rate = self.memory_poll_rate
+            if poll_rate is None:
+                await anyio.sleep(1)
+                continue
             try:
                 await self.info_sender.send(
                     MemoryUsage.from_psutil(override_memory=override_memory)
                 )
-            except Exception as e:
-                logger.warning(f"Error gathering memory usage: {e}")
-            await anyio.sleep(self.memory_poll_rate)
+            except Exception:
+                logger.opt(exception=True).warning("Error gathering memory usage")
+            await anyio.sleep(poll_rate)
 
     async def _watch_system_info(self):
         if self.interface_watcher_interval is None:
@@ -485,8 +538,8 @@ class InfoGatherer:
                 with fail_after(10):
                     nics = await get_network_interfaces()
                     await self.info_sender.send(NodeNetworkInterfaces(ifaces=nics))
-            except Exception as e:
-                logger.warning(f"Error gathering network interfaces: {e}")
+            except Exception:
+                logger.opt(exception=True).warning("Error gathering network interfaces")
             await anyio.sleep(self.interface_watcher_interval)
 
     async def _monitor_thunderbolt_bridge_status(self):
@@ -498,8 +551,8 @@ class InfoGatherer:
                     curr = await ThunderboltBridgeInfo.gather()
                     if curr is not None:
                         await self.info_sender.send(curr)
-            except Exception as e:
-                logger.warning(f"Error gathering Thunderbolt Bridge status: {e}")
+            except Exception:
+                logger.opt(exception=True).warning("Error gathering Thunderbolt Bridge status")
             await anyio.sleep(self.thunderbolt_bridge_poll_interval)
 
     async def _monitor_rdma_ctl_status(self):
@@ -510,8 +563,8 @@ class InfoGatherer:
                 curr = await RdmaCtlStatus.gather()
                 if curr is not None:
                     await self.info_sender.send(curr)
-            except Exception as e:
-                logger.warning(f"Error gathering RDMA ctl status: {e}")
+            except Exception:
+                logger.opt(exception=True).warning("Error gathering RDMA ctl status")
             await anyio.sleep(self.rdma_ctl_poll_interval)
 
     async def _monitor_disk_usage(self):
@@ -521,8 +574,8 @@ class InfoGatherer:
             try:
                 with fail_after(5):
                     await self.info_sender.send(await NodeDiskUsage.gather())
-            except Exception as e:
-                logger.warning(f"Error gathering disk usage: {e}")
+            except Exception:
+                logger.opt(exception=True).warning("Error gathering disk usage")
             await anyio.sleep(self.disk_poll_interval)
 
     async def _monitor_macmon(self, macmon_path: str):
@@ -554,6 +607,10 @@ class InfoGatherer:
                 logger.warning(
                     f"MacMon produced no output for {read_timeout}s, restarting"
                 )
+                self._enable_psutil_memory_fallback(
+                    "MacMon produced no output, falling back to psutil memory "
+                    "monitoring"
+                )
             except CalledProcessError as e:
                 stderr_msg = "no stderr"
                 stderr_output = cast(bytes | str | None, e.stderr)
@@ -566,6 +623,12 @@ class InfoGatherer:
                 logger.warning(
                     f"MacMon failed with return code {e.returncode}: {stderr_msg}"
                 )
-            except Exception as e:
-                logger.warning(f"Error in macmon monitor: {e}")
+                self._enable_psutil_memory_fallback(
+                    "MacMon failed, falling back to psutil memory monitoring"
+                )
+            except Exception:
+                logger.opt(exception=True).warning("Error in macmon monitor")
+                self._enable_psutil_memory_fallback(
+                    "MacMon crashed, falling back to psutil memory monitoring"
+                )
             await anyio.sleep(self.macmon_interval)