diff --git a/src/exo/download/download_utils.py b/src/exo/download/download_utils.py index 55a8f6c662..3fb4ebce90 100644 --- a/src/exo/download/download_utils.py +++ b/src/exo/download/download_utils.py @@ -281,25 +281,128 @@ async def seed_models(seed_dir: str | Path): logger.error(traceback.format_exc()) +def has_partial_files(model_dir: Path) -> bool: + """Check if the directory contains any in-progress '.partial' download files.""" + if not model_dir.is_dir(): + return False + for _, _, filenames in os.walk(model_dir): + for filename in filenames: + if filename.endswith(".partial"): + return True + return False + + +def is_local_model_folder_ready(model_dir: Path) -> bool: + """Checks if a model directory is non-empty and contains no '.partial' download files. + + Implements official TODO.md item #8: allows loading pre-copied or offline model directories + without requiring internet connection to verify a remote index or cache. + """ + if not model_dir.is_dir(): + return False + + has_files = False + has_recognized_model_artifact = False + + for dirpath, _, filenames in os.walk(model_dir): + for filename in filenames: + if filename.endswith(".partial"): + return False # Incomplete download fragment present + has_files = True + fn_lower = filename.lower() + if ( + fn_lower.endswith( + ( + ".safetensors", + ".bin", + ".pt", + ".gguf", + ".onnx", + ".model", + ) + ) + or fn_lower + in ( + "config.json", + "model_card.toml", + "tokenizer.json", + "params.json", + ) + ): + has_recognized_model_artifact = True + + return has_files and has_recognized_model_artifact + + def _scan_model_directory( model_dir: Path, recursive: bool = False ) -> list[FileListEntry] | None: """Scan a local model directory and build a file list. - Requires at least one ``*.safetensors.index.json``. Every weight file - referenced by the index that is missing on disk gets ``size=None``. + Supports both: + 1. Standard repos with a ``*.safetensors.index.json``. + 2. Offline/copied model folders without index json (TODO #8) as long as the + folder is non-empty and has no ``.partial`` files. """ index_files = list(model_dir.glob("**/*.safetensors.index.json")) - if not index_files: - return None - entries_by_path: dict[str, FileListEntry] = {} + if index_files: + # Case A: Sharded model with index file + if recursive: + for dirpath, _, filenames in os.walk(model_dir): + for filename in filenames: + if filename.endswith(".partial"): + continue + full_path = Path(dirpath) / filename + rel_path = str(full_path.relative_to(model_dir)) + entries_by_path[rel_path] = FileListEntry( + type="file", + path=rel_path, + size=full_path.stat().st_size, + ) + else: + for item in model_dir.iterdir(): + if item.is_file() and not item.name.endswith(".partial"): + entries_by_path[item.name] = FileListEntry( + type="file", + path=item.name, + size=item.stat().st_size, + ) + + # Add expected weight files from index that haven't been downloaded yet + for index_file in index_files: + try: + index_data = ModelSafetensorsIndex.model_validate_json( + index_file.read_text() + ) + relative_dir = index_file.parent.relative_to(model_dir) + for filename in set(index_data.weight_map.values()): + rel_path = ( + str(relative_dir / filename) + if relative_dir != Path(".") + else filename + ) + if rel_path not in entries_by_path: + entries_by_path[rel_path] = FileListEntry( + type="file", + path=rel_path, + size=None, + ) + except (OSError, ValueError): + continue + + return list(entries_by_path.values()) + + # Case B (TODO.md #8): Model directory without index json (e.g. single safetensors, GGUF, or copied model) + if not is_local_model_folder_ready(model_dir): + return None + if recursive: for dirpath, _, filenames in os.walk(model_dir): for filename in filenames: if filename.endswith(".partial"): - continue + return None # Incomplete download fragment full_path = Path(dirpath) / filename rel_path = str(full_path.relative_to(model_dir)) entries_by_path[rel_path] = FileListEntry( @@ -309,36 +412,16 @@ def _scan_model_directory( ) else: for item in model_dir.iterdir(): - if item.is_file() and not item.name.endswith(".partial"): + if item.is_file(): + if item.name.endswith(".partial"): + return None entries_by_path[item.name] = FileListEntry( type="file", path=item.name, size=item.stat().st_size, ) - # Add expected weight files from index that haven't been downloaded yet - for index_file in index_files: - try: - index_data = ModelSafetensorsIndex.model_validate_json( - index_file.read_text() - ) - relative_dir = index_file.parent.relative_to(model_dir) - for filename in set(index_data.weight_map.values()): - rel_path = ( - str(relative_dir / filename) - if relative_dir != Path(".") - else filename - ) - if rel_path not in entries_by_path: - entries_by_path[rel_path] = FileListEntry( - type="file", - path=rel_path, - size=None, - ) - except Exception: - continue - - return list(entries_by_path.values()) + return list(entries_by_path.values()) if entries_by_path else None def is_model_directory_complete(model_dir: Path, card: ModelCard | None = None) -> bool: @@ -369,9 +452,7 @@ async def _build_file_list_from_local_directory( ) -> list[FileListEntry] | None: """Build a file list from locally existing model files. - We can only figure out the files we need from safetensors index, so - a local directory must contain a *.safetensors.index.json and - safetensors listed there. + Supports both index-based models and standalone non-empty model folders with no .partial files (TODO #8). """ normalized = model_id.normalize() for search_dir in (*EXO_MODELS_READ_ONLY_DIRS, *EXO_MODELS_DIRS): @@ -416,6 +497,13 @@ async def fetch_file_list_with_cache( logger.warning( f"No internet and no cached file list for {model_id} - using local file list" ) + try: + async with aiofiles.open(cache_file, "w") as f: + await f.write( + TypeAdapter(list[FileListEntry]).dump_json(local_file_list).decode() + ) + except OSError as write_err: + logger.debug(f"Could not persist offline file list cache {cache_file}: {write_err}") return local_file_list raise FileNotFoundError( f"No internet connection and no cached file list for {model_id}" @@ -451,6 +539,13 @@ async def fetch_file_list_with_cache( logger.warning( f"Failed to fetch file list for {model_id} and no cache exists, using local file list" ) + try: + async with aiofiles.open(cache_file, "w") as f: + await f.write( + TypeAdapter(list[FileListEntry]).dump_json(local_file_list).decode() + ) + except OSError as write_err: + logger.debug(f"Could not persist offline file list cache {cache_file}: {write_err}") return local_file_list raise FileNotFoundError(f"Failed to fetch file list for {model_id}: {e}") from e diff --git a/src/exo/download/tests/test_offline_model_loading.py b/src/exo/download/tests/test_offline_model_loading.py new file mode 100644 index 0000000000..ee4de3bf81 --- /dev/null +++ b/src/exo/download/tests/test_offline_model_loading.py @@ -0,0 +1,141 @@ +import json +from pathlib import Path + +import pytest + +from exo.download.download_utils import ( + _scan_model_directory, + fetch_file_list_with_cache, + has_partial_files, + is_local_model_folder_ready, + is_model_directory_complete, +) +from exo.shared.models.model_cards import ModelId + + +def test_has_partial_files(tmp_path: Path): + """Verify detection of in-progress .partial download fragments.""" + model_dir = tmp_path / "test_model" + model_dir.mkdir() + assert not has_partial_files(model_dir) + + # Add standard files + (model_dir / "config.json").write_text("{}") + (model_dir / "model.safetensors").write_bytes(b"dummy weight data") + assert not has_partial_files(model_dir) + + # Add a .partial file in a nested subdirectory + nested = model_dir / "sub" + nested.mkdir() + (nested / "shard_02.safetensors.partial").write_bytes(b"partial data") + assert has_partial_files(model_dir) + + +def test_is_local_model_folder_ready(tmp_path: Path): + """Verify non-empty folder validation and requirement of no .partial files (TODO #8).""" + model_dir = tmp_path / "llama-offline" + model_dir.mkdir() + + # Empty folder is not ready + assert not is_local_model_folder_ready(model_dir) + + # Folder with only random unrelated text is not recognized + (model_dir / "notes.txt").write_text("hello") + assert not is_local_model_folder_ready(model_dir) + + # Folder with config.json and weights + (model_dir / "config.json").write_text('{"architectures": ["LlamaForCausalLM"]}') + (model_dir / "model.safetensors").write_bytes(b"XYZ" * 100) + assert is_local_model_folder_ready(model_dir) + + # If a .partial file is present, it must be rejected as incomplete + (model_dir / "download.partial").write_bytes(b"incomplete") + assert not is_local_model_folder_ready(model_dir) + + +def test_scan_model_directory_offline_copy_without_index(tmp_path: Path): + """Verify scanning offline pre-copied model directories without an index json.""" + model_dir = tmp_path / "qwen-offline" + model_dir.mkdir() + + config_path = model_dir / "config.json" + config_path.write_text('{"vocab_size": 32000}') + + weight_path = model_dir / "model.safetensors" + weight_path.write_bytes(b"\x00\x01\x02\x03" * 256) + + # Must build valid FileListEntry list even without *.safetensors.index.json + file_list = _scan_model_directory(model_dir, recursive=True) + assert file_list is not None + assert len(file_list) == 2 + + paths = {f.path for f in file_list} + assert "config.json" in paths + assert "model.safetensors" in paths + + for f in file_list: + assert f.size is not None + assert f.size > 0 + + # Directory must be considered complete + assert is_model_directory_complete(model_dir) + + +def test_scan_model_directory_with_partial_file_fails(tmp_path: Path): + """Verify directory with in-progress .partial file is rejected as incomplete.""" + model_dir = tmp_path / "incomplete-model" + model_dir.mkdir() + + (model_dir / "config.json").write_text("{}") + (model_dir / "weights.safetensors.partial").write_bytes(b"incomplete bytes") + + file_list = _scan_model_directory(model_dir, recursive=True) + assert file_list is None + assert not is_model_directory_complete(model_dir) + + +@pytest.mark.asyncio +async def test_fetch_file_list_with_cache_offline_fallback(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + """Verify fetch_file_list_with_cache succeeds offline with no pre-existing cache file.""" + # Point models dirs and cache dir to temporary path + models_root = tmp_path / "models" + models_root.mkdir() + cache_root = tmp_path / "cache" + cache_root.mkdir() + + model_id = ModelId("meta-llama/Llama-Offline-Test") + normalized_name = model_id.normalize() + local_model_dir = models_root / normalized_name + local_model_dir.mkdir(parents=True) + + (local_model_dir / "config.json").write_text('{"hidden_size": 2048}') + (local_model_dir / "model.safetensors").write_bytes(b"test weights" * 50) + + # Monkeypatch search dirs and cache dir + monkeypatch.setattr("exo.download.download_utils.EXO_MODELS_DIRS", [models_root]) + monkeypatch.setattr("exo.download.download_utils.EXO_MODELS_READ_ONLY_DIRS", []) + + async def _mock_ensure_cache_dir(_model_id: ModelId) -> Path: + return cache_root + + monkeypatch.setattr("exo.download.download_utils.ensure_cache_dir", _mock_ensure_cache_dir) + + # Call with skip_internet=True and no pre-existing cache file + entries = await fetch_file_list_with_cache( + model_id=model_id, + revision="main", + recursive=True, + skip_internet=True, + ) + + assert entries is not None + assert len(entries) == 2 + paths = {e.path for e in entries} + assert "config.json" in paths + assert "model.safetensors" in paths + + # Verify that a local cache file was automatically constructed + cache_files = list(cache_root.glob("*.json")) + assert len(cache_files) == 1 + cached_content = json.loads(cache_files[0].read_text()) + assert len(cached_content) == 2