From 2ebe6216b431643d7a7dc16f3914e8dba2db9cd6 Mon Sep 17 00:00:00 2001 From: vskiwi <141816715+vskiwi@users.noreply.github.com> Date: Wed, 18 Feb 2026 19:18:09 +0300 Subject: [PATCH] feat: add explicit --offline mode for air-gapped clusters (#1525) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Motivation Closes #1510 There is currently no reliable way to run exo on an air-gapped or offline cluster where models are pre-staged on local disks. The two existing mechanisms — `--no-downloads` and `HF_HUB_OFFLINE=1` — each cover only a subset of the problem: 1. **`--no-downloads` blocks model loading**: When passed, `DownloadCoordinator` is not created. No `NodeDownloadProgress` events are ever emitted, so `_model_needs_download()` in `plan.py` perpetually returns `DownloadModel`, short-circuiting `_load_model()` and preventing the model from ever being loaded. 2. **`HF_HUB_OFFLINE=1` doesn't cover exo's aiohttp code**: exo's download pipeline primarily uses raw `aiohttp` for HTTP operations (file list fetching, file downloads, HEAD verification), not the `huggingface_hub` library. These calls will attempt connections and time out on air-gapped networks. 3. **`skip_internet` is not propagated to `download_file_with_retry()`**: Even when `internet_connection = False`, the `_download_file()` function still makes HTTP HEAD calls via `file_meta()` to verify local files and unconditionally attempts downloads for missing files. ## Changes ### `src/exo/main.py` - Add `--offline` flag to `Args` with env var detection (`EXO_OFFLINE=1`, `HF_HUB_OFFLINE=1`) - Pass `offline` to `DownloadCoordinator` at creation and re-creation (election loop) ### `src/exo/download/coordinator.py` - Add `offline: bool = False` field - In offline mode: set `internet_connection = False` immediately in `__post_init__`, skip `_test_internet_connection()` ping (avoids 3s timeout), skip `_check_internet_connection` periodic loop - In `_start_download()`: if model is not fully available locally, emit `DownloadFailed` with clear message instead of starting a download task ### `src/exo/download/download_utils.py` - Add `skip_internet: bool` parameter to `download_file_with_retry()` and `_download_file()` - When `skip_internet=True` in `_download_file()`: return local file immediately without HTTP HEAD verification; raise `FileNotFoundError` for missing files - Propagate `skip_internet` from `download_shard()` to `download_file_with_retry()` ### `src/exo/download/tests/test_offline_mode.py` (new) - 8 tests covering `_download_file`, `download_file_with_retry`, and `fetch_file_list_with_cache` in offline mode ## Why It Works Unlike `--no-downloads` which disables `DownloadCoordinator` entirely, `--offline` keeps the coordinator running in a restricted mode. The existing `_emit_existing_download_progress()` disk scanner still runs every 60 seconds, emitting `DownloadCompleted` events for pre-staged models. These events flow through the event-sourcing pipeline and populate `state.downloads`, which unblocks `_model_needs_download()` in `plan.py` — no changes to the planning logic required. ``` --offline flag → DownloadCoordinator (offline mode) → Skip 1.1.1.1 ping, internet_connection = False → _emit_existing_download_progress scans disk → Emits DownloadCompleted for pre-staged models → _model_needs_download sees DownloadCompleted → _load_model proceeds normally ``` ## Test Plan ### Automated Testing - `ruff check` — passes - 8 new tests in `test_offline_mode.py` — all pass - 11 existing download tests in `test_download_verification.py` — all pass (no regressions) ### Manual Testing 1. Pre-stage a model on disk (e.g., `~/.exo/models/mlx-community--Qwen3-0.6B-4bit/`) 2. Start exo with `--offline` (or `EXO_OFFLINE=1`) 3. Place an instance via API or dashboard 4. Verify: model loads into memory and inference works without any network calls ### Environment - macOS (Apple Silicon), multi-node cluster with Thunderbolt interconnect - Models pre-staged via rsync / NFS mount --- src/exo/download/coordinator.py | 27 ++- src/exo/download/download_utils.py | 13 +- src/exo/download/tests/test_offline_mode.py | 230 ++++++++++++++++++++ src/exo/main.py | 13 ++ 4 files changed, 279 insertions(+), 4 deletions(-) create mode 100644 src/exo/download/tests/test_offline_mode.py diff --git a/src/exo/download/coordinator.py b/src/exo/download/coordinator.py index db13ccef..899e4f14 100644 --- a/src/exo/download/coordinator.py +++ b/src/exo/download/coordinator.py @@ -47,6 +47,7 @@ class DownloadCoordinator: download_command_receiver: Receiver[ForwarderDownloadCommand] local_event_sender: Sender[ForwarderEvent] event_index_counter: Iterator[int] + offline: bool = False # Local state download_status: dict[ModelId, DownloadProgress] = field(default_factory=dict) @@ -62,6 +63,8 @@ class DownloadCoordinator: def __post_init__(self) -> None: self.event_sender, self.event_receiver = channel[Event]() + if self.offline: + self.shard_downloader.set_internet_connection(False) self.shard_downloader.on_progress(self._download_progress_callback) def _model_dir(self, model_id: ModelId) -> str: @@ -107,13 +110,17 @@ class DownloadCoordinator: self._last_progress_time[model_id] = current_time() async def run(self) -> None: - logger.info("Starting DownloadCoordinator") - self._test_internet_connection() + logger.info( + f"Starting DownloadCoordinator{' (offline mode)' if self.offline else ''}" + ) + if not self.offline: + self._test_internet_connection() async with self._tg as tg: tg.start_soon(self._command_processor) tg.start_soon(self._forward_events) tg.start_soon(self._emit_existing_download_progress) - tg.start_soon(self._check_internet_connection) + if not self.offline: + tg.start_soon(self._check_internet_connection) def _test_internet_connection(self) -> None: try: @@ -202,6 +209,20 @@ class DownloadCoordinator: ) return + if self.offline: + logger.warning( + f"Offline mode: model {model_id} is not fully available locally, cannot download" + ) + failed = DownloadFailed( + shard_metadata=shard, + node_id=self.node_id, + error_message=f"Model files not found locally in offline mode: {model_id}", + model_directory=self._model_dir(model_id), + ) + self.download_status[model_id] = failed + await self.event_sender.send(NodeDownloadProgress(download_progress=failed)) + return + # Start actual download self._start_download_task(shard, initial_progress) diff --git a/src/exo/download/download_utils.py b/src/exo/download/download_utils.py index 7974d504..5691e5dd 100644 --- a/src/exo/download/download_utils.py +++ b/src/exo/download/download_utils.py @@ -448,12 +448,13 @@ async def download_file_with_retry( target_dir: Path, on_progress: Callable[[int, int, bool], None] = lambda _, __, ___: None, on_connection_lost: Callable[[], None] = lambda: None, + skip_internet: bool = False, ) -> Path: n_attempts = 3 for attempt in range(n_attempts): try: return await _download_file( - model_id, revision, path, target_dir, on_progress + model_id, revision, path, target_dir, on_progress, skip_internet ) except HuggingFaceAuthenticationError: raise @@ -487,10 +488,14 @@ async def _download_file( path: str, target_dir: Path, on_progress: Callable[[int, int, bool], None] = lambda _, __, ___: None, + skip_internet: bool = False, ) -> Path: target_path = target_dir / path if await aios.path.exists(target_path): + if skip_internet: + return target_path + local_size = (await aios.stat(target_path)).st_size # Try to verify against remote, but allow offline operation @@ -510,6 +515,11 @@ async def _download_file( ) return target_path + if skip_internet: + raise FileNotFoundError( + f"File {path} not found locally and cannot download in offline mode" + ) + await aios.makedirs((target_dir / path).parent, exist_ok=True) length, etag = await file_meta(model_id, revision, path) remote_hash = etag[:-5] if etag.endswith("-gzip") else etag @@ -814,6 +824,7 @@ async def download_shard( file, curr_bytes, total_bytes, is_renamed ), on_connection_lost=on_connection_lost, + skip_internet=skip_internet, ) if not skip_download: diff --git a/src/exo/download/tests/test_offline_mode.py b/src/exo/download/tests/test_offline_mode.py new file mode 100644 index 00000000..15210c3f --- /dev/null +++ b/src/exo/download/tests/test_offline_mode.py @@ -0,0 +1,230 @@ +"""Tests for offline/air-gapped mode.""" + +from collections.abc import AsyncIterator +from pathlib import Path +from unittest.mock import AsyncMock, patch + +import aiofiles +import aiofiles.os as aios +import pytest + +from exo.download.download_utils import ( + _download_file, # pyright: ignore[reportPrivateUsage] + download_file_with_retry, + fetch_file_list_with_cache, +) +from exo.shared.types.common import ModelId +from exo.shared.types.worker.downloads import FileListEntry + + +@pytest.fixture +def model_id() -> ModelId: + return ModelId("test-org/test-model") + + +@pytest.fixture +async def temp_models_dir(tmp_path: Path) -> AsyncIterator[Path]: + models_dir = tmp_path / "models" + await aios.makedirs(models_dir, exist_ok=True) + with patch("exo.download.download_utils.EXO_MODELS_DIR", models_dir): + yield models_dir + + +class TestDownloadFileOffline: + """Tests for _download_file with skip_internet=True.""" + + async def test_returns_local_file_without_http_verification( + self, model_id: ModelId, tmp_path: Path + ) -> None: + """When skip_internet=True and file exists locally, return it immediately + without making any HTTP calls (no file_meta verification).""" + target_dir = tmp_path / "downloads" + await aios.makedirs(target_dir, exist_ok=True) + + local_file = target_dir / "model.safetensors" + async with aiofiles.open(local_file, "wb") as f: + await f.write(b"model weights data") + + with patch( + "exo.download.download_utils.file_meta", + new_callable=AsyncMock, + ) as mock_file_meta: + result = await _download_file( + model_id, + "main", + "model.safetensors", + target_dir, + skip_internet=True, + ) + + assert result == local_file + mock_file_meta.assert_not_called() + + async def test_raises_file_not_found_for_missing_file( + self, model_id: ModelId, tmp_path: Path + ) -> None: + """When skip_internet=True and file does NOT exist locally, + raise FileNotFoundError instead of attempting download.""" + target_dir = tmp_path / "downloads" + await aios.makedirs(target_dir, exist_ok=True) + + with pytest.raises(FileNotFoundError, match="offline mode"): + await _download_file( + model_id, + "main", + "missing_model.safetensors", + target_dir, + skip_internet=True, + ) + + async def test_returns_local_file_in_subdirectory( + self, model_id: ModelId, tmp_path: Path + ) -> None: + """When skip_internet=True and file exists in a subdirectory, + return it without HTTP calls.""" + target_dir = tmp_path / "downloads" + subdir = target_dir / "transformer" + await aios.makedirs(subdir, exist_ok=True) + + local_file = subdir / "diffusion_pytorch_model.safetensors" + async with aiofiles.open(local_file, "wb") as f: + await f.write(b"weights") + + with patch( + "exo.download.download_utils.file_meta", + new_callable=AsyncMock, + ) as mock_file_meta: + result = await _download_file( + model_id, + "main", + "transformer/diffusion_pytorch_model.safetensors", + target_dir, + skip_internet=True, + ) + + assert result == local_file + mock_file_meta.assert_not_called() + + +class TestDownloadFileWithRetryOffline: + """Tests for download_file_with_retry with skip_internet=True.""" + + async def test_propagates_skip_internet_to_download_file( + self, model_id: ModelId, tmp_path: Path + ) -> None: + """Verify skip_internet is passed through to _download_file.""" + target_dir = tmp_path / "downloads" + await aios.makedirs(target_dir, exist_ok=True) + + local_file = target_dir / "config.json" + async with aiofiles.open(local_file, "wb") as f: + await f.write(b'{"model_type": "qwen2"}') + + with patch( + "exo.download.download_utils.file_meta", + new_callable=AsyncMock, + ) as mock_file_meta: + result = await download_file_with_retry( + model_id, + "main", + "config.json", + target_dir, + skip_internet=True, + ) + + assert result == local_file + mock_file_meta.assert_not_called() + + async def test_file_not_found_does_not_retry( + self, model_id: ModelId, tmp_path: Path + ) -> None: + """FileNotFoundError from offline mode should not trigger retries.""" + target_dir = tmp_path / "downloads" + await aios.makedirs(target_dir, exist_ok=True) + + with pytest.raises(FileNotFoundError): + await download_file_with_retry( + model_id, + "main", + "nonexistent.safetensors", + target_dir, + skip_internet=True, + ) + + +class TestFetchFileListOffline: + """Tests for fetch_file_list_with_cache with skip_internet=True.""" + + async def test_uses_cached_file_list( + self, model_id: ModelId, temp_models_dir: Path + ) -> None: + """When skip_internet=True and cache file exists, use it without network.""" + from pydantic import TypeAdapter + + cache_dir = temp_models_dir / "caches" / model_id.normalize() + await aios.makedirs(cache_dir, exist_ok=True) + + cached_list = [ + FileListEntry(type="file", path="model.safetensors", size=1000), + FileListEntry(type="file", path="config.json", size=200), + ] + cache_file = cache_dir / f"{model_id.normalize()}--main--file_list.json" + async with aiofiles.open(cache_file, "w") as f: + await f.write( + TypeAdapter(list[FileListEntry]).dump_json(cached_list).decode() + ) + + with patch( + "exo.download.download_utils.fetch_file_list_with_retry", + new_callable=AsyncMock, + ) as mock_fetch: + result = await fetch_file_list_with_cache( + model_id, "main", skip_internet=True + ) + + assert result == cached_list + mock_fetch.assert_not_called() + + async def test_falls_back_to_local_directory_scan( + self, model_id: ModelId, temp_models_dir: Path + ) -> None: + """When skip_internet=True and no cache but local files exist, + build file list from local directory.""" + import json + + model_dir = temp_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"}') + + index_data = { + "metadata": {}, + "weight_map": {"model.layers.0.weight": "model.safetensors"}, + } + async with aiofiles.open(model_dir / "model.safetensors.index.json", "w") as f: + await f.write(json.dumps(index_data)) + + async with aiofiles.open(model_dir / "model.safetensors", "wb") as f: + await f.write(b"x" * 500) + + with patch( + "exo.download.download_utils.fetch_file_list_with_retry", + new_callable=AsyncMock, + ) as mock_fetch: + result = await fetch_file_list_with_cache( + model_id, "main", skip_internet=True + ) + + mock_fetch.assert_not_called() + paths = {entry.path for entry in result} + assert "config.json" in paths + assert "model.safetensors" in paths + + async def test_raises_when_no_cache_and_no_local_files( + self, model_id: ModelId, temp_models_dir: Path + ) -> None: + """When skip_internet=True and neither cache nor local files exist, + raise FileNotFoundError.""" + with pytest.raises(FileNotFoundError, match="No internet"): + await fetch_file_list_with_cache(model_id, "main", skip_internet=True) diff --git a/src/exo/main.py b/src/exo/main.py index 1d358975..ec203181 100644 --- a/src/exo/main.py +++ b/src/exo/main.py @@ -39,6 +39,7 @@ class Node: node_id: NodeId event_index_counter: Iterator[int] + offline: bool _tg: TaskGroup = field(init=False, default_factory=anyio.create_task_group) @classmethod @@ -68,6 +69,7 @@ class Node: download_command_receiver=router.receiver(topics.DOWNLOAD_COMMANDS), local_event_sender=router.sender(topics.LOCAL_EVENTS), event_index_counter=event_index_counter, + offline=args.offline, ) else: download_coordinator = None @@ -132,6 +134,7 @@ class Node: api, node_id, event_index_counter, + args.offline, ) async def run(self): @@ -222,6 +225,7 @@ class Node: ), local_event_sender=self.router.sender(topics.LOCAL_EVENTS), event_index_counter=self.event_index_counter, + offline=self.offline, ) self._tg.start_soon(self.download_coordinator.run) if self.worker: @@ -260,6 +264,9 @@ def main(): logger.info("Starting EXO") logger.info(f"EXO_LIBP2P_NAMESPACE: {os.getenv('EXO_LIBP2P_NAMESPACE')}") + if args.offline: + logger.info("Running in OFFLINE mode — no internet checks, local models only") + # Set FAST_SYNCH override env var for runner subprocesses if args.fast_synch is True: os.environ["EXO_FAST_SYNCH"] = "on" @@ -282,6 +289,7 @@ class Args(CamelCaseModel): tb_only: bool = False no_worker: bool = False no_downloads: bool = False + offline: bool = False fast_synch: bool | None = None # None = auto, True = force on, False = force off @classmethod @@ -329,6 +337,11 @@ class Args(CamelCaseModel): action="store_true", help="Disable the download coordinator (node won't download models)", ) + parser.add_argument( + "--offline", + action="store_true", + help="Run in offline/air-gapped mode: skip internet checks, use only pre-staged local models", + ) fast_synch_group = parser.add_mutually_exclusive_group() fast_synch_group.add_argument( "--fast-synch",