gotta move the cache soon

python/exo_core/pyproject.toml

@@ -5,8 +5,23 @@ description = "Add your description here"
 readme = "README.md"
 authors = [{ name = "Evan", email = "[email protected]" }]
 requires-python = ">=3.13"
-dependencies = ["pydantic"]
+dependencies = ["pydantic>=2.13.0b2"]
 
 [build-system]
 requires = ["uv_build>=0.9.24,<0.10.0"]
 build-backend = "uv_build"
+
+[tool.basedpyright]
+typeCheckingMode = "strict"
+failOnWarnings = true
+
+reportAny = "error"
+reportUnknownVariableType = "error"
+reportUnknownParameterType = "error"
+reportMissingParameterType = "error"
+reportMissingTypeStubs = "error"
+reportInvalidCast = "error"
+reportUnnecessaryCast = "error"
+reportUnnecessaryTypeIgnoreComment = "error"
+
+pythonVersion = "3.13"

python/exo_core/src/exo_core/constants.py

@@ -0,0 +1,85 @@
+import os
+import sys
+from pathlib import Path
+
+from exo_core.utils.paths import find_dashboard, find_resources
+
+_EXO_HOME_ENV = os.environ.get("EXO_HOME", None)
+
+
+def _get_xdg_dir(env_var: str, fallback: str) -> Path:
+    """Get XDG directory, prioritising EXO_HOME environment variable if its set. On non-Linux platforms, default to ~/.exo."""
+
+    if _EXO_HOME_ENV is not None:
+        return Path.home() / _EXO_HOME_ENV
+
+    if sys.platform != "linux":
+        return Path.home() / ".exo"
+
+    xdg_value = os.environ.get(env_var, None)
+    if xdg_value is not None:
+        return Path(xdg_value) / "exo"
+    return Path.home() / fallback / "exo"
+
+
+EXO_CONFIG_HOME = _get_xdg_dir("XDG_CONFIG_HOME", ".config")
+EXO_DATA_HOME = _get_xdg_dir("XDG_DATA_HOME", ".local/share")
+EXO_CACHE_HOME = _get_xdg_dir("XDG_CACHE_HOME", ".cache")
+
+# Models directory (data)
+_EXO_MODELS_DIR_ENV = os.environ.get("EXO_MODELS_DIR", None)
+EXO_MODELS_DIR = (
+    EXO_DATA_HOME / "models"
+    if _EXO_MODELS_DIR_ENV is None
+    else Path.home() / _EXO_MODELS_DIR_ENV
+)
+
+# Read-only search path for pre-downloaded models (colon-separated directories)
+_EXO_MODELS_PATH_ENV = os.environ.get("EXO_MODELS_PATH", None)
+EXO_MODELS_PATH: tuple[Path, ...] | None = (
+    tuple(Path(p).expanduser() for p in _EXO_MODELS_PATH_ENV.split(":") if p)
+    if _EXO_MODELS_PATH_ENV is not None
+    else None
+)
+
+_RESOURCES_DIR_ENV = os.environ.get("EXO_RESOURCES_DIR", None)
+RESOURCES_DIR = (
+    find_resources() if _RESOURCES_DIR_ENV is None else Path.home() / _RESOURCES_DIR_ENV
+)
+_DASHBOARD_DIR_ENV = os.environ.get("EXO_DASHBOARD_DIR", None)
+DASHBOARD_DIR = (
+    find_dashboard() if _DASHBOARD_DIR_ENV is None else Path.home() / _DASHBOARD_DIR_ENV
+)
+
+# Log files (data/logs or cache)
+EXO_LOG_DIR = EXO_CACHE_HOME / "exo_log"
+EXO_LOG = EXO_LOG_DIR / "exo.log"
+EXO_TEST_LOG = EXO_CACHE_HOME / "exo_test.log"
+
+# Identity (config)
+EXO_NODE_ID_KEYPAIR = EXO_CONFIG_HOME / "node_id.keypair"
+EXO_CONFIG_FILE = EXO_CONFIG_HOME / "config.toml"
+
+# libp2p topics for event forwarding
+LIBP2P_LOCAL_EVENTS_TOPIC = "worker_events"
+LIBP2P_GLOBAL_EVENTS_TOPIC = "global_events"
+LIBP2P_ELECTION_MESSAGES_TOPIC = "election_message"
+LIBP2P_COMMANDS_TOPIC = "commands"
+
+EXO_MAX_CHUNK_SIZE = 512 * 1024
+
+EXO_CUSTOM_MODEL_CARDS_DIR = EXO_DATA_HOME / "custom_model_cards"
+
+EXO_EVENT_LOG_DIR = EXO_DATA_HOME / "event_log"
+EXO_IMAGE_CACHE_DIR = EXO_CACHE_HOME / "images"
+EXO_TRACING_CACHE_DIR = EXO_CACHE_HOME / "traces"
+
+EXO_ENABLE_IMAGE_MODELS = (
+    os.getenv("EXO_ENABLE_IMAGE_MODELS", "false").lower() == "true"
+)
+
+EXO_OFFLINE = os.getenv("EXO_OFFLINE", "false").lower() == "true"
+
+EXO_TRACING_ENABLED = os.getenv("EXO_TRACING_ENABLED", "false").lower() == "true"
+
+EXO_MAX_CONCURRENT_REQUESTS = int(os.getenv("EXO_MAX_CONCURRENT_REQUESTS", "8"))

python/exo_core/src/exo_core/engine.py

@@ -2,8 +2,7 @@ from abc import ABC, abstractmethod
 from collections.abc import Callable, Iterable
 from typing import Self
 
-
-class TaskId(str): ...
+from exo_core.types.tasks import TaskId
 
 
 class Cancelled: ...

python/exo_core/src/exo_core/model_cards.py

@@ -1,31 +1,23 @@
 from enum import Enum
-from typing import Annotated, Any
 
-import aiofiles
 import aiofiles.os as aios
 import tomlkit
 from anyio import Path, open_file
-from exo.shared.types.common import ModelId
-from exo.shared.types.memory import Memory
-from exo.utils.pydantic_ext import CamelCaseModel
-from huggingface_hub import model_info
-from loguru import logger
 from pydantic import (
-    AliasChoices,
-    BaseModel,
-    Field,
     PositiveInt,
     ValidationError,
     field_validator,
-    model_validator,
 )
 from tomlkit.exceptions import TOMLKitError
 
-from exo.shared.constants import (
+from exo_core.constants import (
     EXO_CUSTOM_MODEL_CARDS_DIR,
     EXO_ENABLE_IMAGE_MODELS,
     RESOURCES_DIR,
 )
+from exo_core.models import CamelCaseModel
+from exo_core.types.common import ModelId
+from exo_core.utils.memory import Memory
 
 # kinda ugly...
 # TODO: load search path from config.toml
@@ -127,6 +119,9 @@ class ModelCard(CamelCaseModel):
     @staticmethod
     async def fetch_from_hf(model_id: ModelId) -> "ModelCard":
         """Fetches storage size and number of layers for a Hugging Face model, returns Pydantic ModelMeta."""
+        # needed to prevent circular imports
+        from exo_core.utils.downloads import fetch_config_data, fetch_safetensors_size
+
         # TODO: failure if files do not exist
         config_data = await fetch_config_data(model_id)
         num_layers = config_data.layer_count
@@ -165,119 +160,3 @@ def is_custom_card(model_id: ModelId) -> bool:
         ModelId(model_id).normalize() + ".toml"
     )
     return os.path.isfile(str(card_path))
-
-
-class ConfigData(BaseModel):
-    model_config = {"extra": "ignore"}  # Allow unknown fields
-
-    architectures: list[str] | None = None
-    hidden_size: Annotated[int, Field(ge=0)] | None = None
-    num_key_value_heads: PositiveInt | None = None
-    layer_count: int = Field(
-        validation_alias=AliasChoices(
-            "num_hidden_layers",
-            "num_layers",
-            "n_layer",
-            "n_layers",
-            "num_decoder_layers",
-            "decoder_layers",
-        )
-    )
-
-    @property
-    def supports_tensor(self) -> bool:
-        return self.architectures in [
-            ["Glm4MoeLiteForCausalLM"],
-            ["GlmMoeDsaForCausalLM"],
-            ["DeepseekV32ForCausalLM"],
-            ["DeepseekV3ForCausalLM"],
-            ["Qwen3NextForCausalLM"],
-            ["Qwen3MoeForCausalLM"],
-            ["Qwen3_5MoeForConditionalGeneration"],
-            ["Qwen3_5ForConditionalGeneration"],
-            ["MiniMaxM2ForCausalLM"],
-            ["LlamaForCausalLM"],
-            ["GptOssForCausalLM"],
-            ["Step3p5ForCausalLM"],
-            ["NemotronHForCausalLM"],
-        ]
-
-    @model_validator(mode="before")
-    @classmethod
-    def defer_to_text_config(cls, data: dict[str, Any]):
-        text_config = data.get("text_config")
-        if text_config is None:
-            return data
-
-        for field in [
-            "architectures",
-            "hidden_size",
-            "num_key_value_heads",
-            "num_hidden_layers",
-            "num_layers",
-            "n_layer",
-            "n_layers",
-            "num_decoder_layers",
-            "decoder_layers",
-        ]:
-            if (val := text_config.get(field)) is not None:  # pyright: ignore[reportAny]
-                data[field] = val
-
-        return data
-
-
-async def fetch_config_data(model_id: ModelId) -> ConfigData:
-    """Downloads and parses config.json for a model."""
-    from exo.download.download_utils import (
-        download_file_with_retry,
-        ensure_models_dir,
-    )
-
-    target_dir = (await ensure_models_dir()) / model_id.normalize()
-    await aios.makedirs(target_dir, exist_ok=True)
-    config_path = await download_file_with_retry(
-        model_id,
-        "main",
-        "config.json",
-        target_dir,
-        lambda curr_bytes, total_bytes, is_renamed: logger.debug(
-            f"Downloading config.json for {model_id}: {curr_bytes}/{total_bytes} ({is_renamed=})"
-        ),
-    )
-    async with aiofiles.open(config_path, "r") as f:
-        return ConfigData.model_validate_json(await f.read())
-
-
-async def fetch_safetensors_size(model_id: ModelId) -> Memory:
-    """Gets model size from safetensors index or falls back to HF API."""
-    from exo.download.download_utils import (
-        download_file_with_retry,
-        ensure_models_dir,
-    )
-    from exo.shared.types.worker.downloads import ModelSafetensorsIndex
-
-    target_dir = (await ensure_models_dir()) / model_id.normalize()
-    await aios.makedirs(target_dir, exist_ok=True)
-    try:
-        index_path = await download_file_with_retry(
-            model_id,
-            "main",
-            "model.safetensors.index.json",
-            target_dir,
-            lambda curr_bytes, total_bytes, is_renamed: logger.debug(
-                f"Downloading model.safetensors.index.json for {model_id}: {curr_bytes}/{total_bytes} ({is_renamed=})"
-            ),
-        )
-        async with aiofiles.open(index_path, "r") as f:
-            index_data = ModelSafetensorsIndex.model_validate_json(await f.read())
-
-        metadata = index_data.metadata
-        if metadata is not None:
-            return Memory.from_bytes(metadata.total_size)
-    except FileNotFoundError:
-        pass
-
-    info = model_info(model_id)
-    if info.safetensors is None:
-        raise ValueError(f"No safetensors info found for {model_id}")
-    return Memory.from_bytes(info.safetensors.total)

python/exo_core/src/exo_core/models.py

@@ -1,4 +1,4 @@
-# pyright: reportAny=false, reportUnknownArgumentType=false, reportUnknownVariableType=false
+# pyright: reportAny=false
 
 from typing import Any, Self
 
@@ -42,7 +42,7 @@ class TaggedModel(CamelCaseModel):
     @model_validator(mode="wrap")
     @classmethod
     def _validate(cls, v: Any, handler: ValidatorFunctionWrapHandler) -> Self:
-        if isinstance(v, dict) and len(v) == 1 and cls.__name__ in v:
+        if isinstance(v, dict) and len(v) == 1 and cls.__name__ in v:  # pyright: ignore[reportUnknownArgumentType]
             return handler(v[cls.__name__])
 
         return handler(v)

python/exo_core/src/exo_core/types/chunks.py

@@ -10,7 +10,7 @@ from exo.api.types import (
     Usage,
 )
 from exo_core.model_cards import ModelId
-from exo_core.pydantic import TaggedModel
+from exo_core.models import TaggedModel
 
 from .common import CommandId
 

python/exo_core/src/exo_core/types/common.py

@@ -4,7 +4,7 @@ from uuid import uuid4
 from pydantic import GetCoreSchemaHandler, field_validator
 from pydantic_core import core_schema
 
-from exo_core.pydantic import CamelCaseModel
+from exo_core.models import CamelCaseModel
 
 
 class Id(str):

python/exo_core/src/exo_core/types/downloads.py

@@ -0,0 +1,165 @@
+from datetime import timedelta
+from typing import Annotated, Any, Literal
+
+from pydantic import (
+    AliasChoices,
+    BaseModel,
+    ConfigDict,
+    Field,
+    PositiveInt,
+    model_validator,
+)
+
+from exo_core.models import CamelCaseModel, TaggedModel
+from exo_core.types.common import NodeId
+from exo_core.types.shards import ShardMetadata
+from exo_core.utils.memory import Memory
+
+
+class DownloadProgressData(CamelCaseModel):
+    total: Memory
+    downloaded: Memory
+    downloaded_this_session: Memory
+
+    completed_files: int
+    total_files: int
+
+    speed: float
+    eta_ms: int
+
+    files: dict[str, "DownloadProgressData"]
+
+
+class BaseDownloadProgress(TaggedModel):
+    node_id: NodeId
+    shard_metadata: ShardMetadata
+    model_directory: str = ""
+
+
+class DownloadPending(BaseDownloadProgress):
+    downloaded: Memory = Memory()
+    total: Memory = Memory()
+
+
+class DownloadCompleted(BaseDownloadProgress):
+    total: Memory
+    read_only: bool = False
+
+
+class DownloadFailed(BaseDownloadProgress):
+    error_message: str
+
+
+class DownloadOngoing(BaseDownloadProgress):
+    download_progress: DownloadProgressData
+
+
+DownloadProgress = (
+    DownloadPending | DownloadCompleted | DownloadFailed | DownloadOngoing
+)
+
+
+class ModelSafetensorsIndexMetadata(BaseModel):
+    total_size: PositiveInt
+
+
+class ModelSafetensorsIndex(BaseModel):
+    metadata: ModelSafetensorsIndexMetadata | None
+    weight_map: dict[str, str]
+
+
+class FileListEntry(BaseModel):
+    type: Literal["file", "directory"]
+    path: str
+    size: int | None = None
+
+
+class RepoFileDownloadProgress(BaseModel):
+    repo_id: str
+    repo_revision: str
+    file_path: str
+    downloaded: Memory
+    downloaded_this_session: Memory
+    total: Memory
+    speed: float
+    eta: timedelta
+    status: Literal["not_started", "in_progress", "complete"]
+    start_time: float
+
+    model_config = ConfigDict(frozen=True)
+
+
+class RepoDownloadProgress(BaseModel):
+    repo_id: str
+    repo_revision: str
+    shard: ShardMetadata
+    completed_files: int
+    total_files: int
+    downloaded: Memory
+    downloaded_this_session: Memory
+    total: Memory
+    overall_speed: float
+    overall_eta: timedelta
+    status: Literal["not_started", "in_progress", "complete"]
+    file_progress: dict[str, RepoFileDownloadProgress] = Field(default_factory=dict)
+
+    model_config = ConfigDict(frozen=True)
+
+
+class ConfigData(BaseModel):
+    model_config = {"extra": "ignore"}  # Allow unknown fields
+
+    architectures: list[str] | None = None
+    hidden_size: Annotated[int, Field(ge=0)] | None = None
+    num_key_value_heads: PositiveInt | None = None
+    layer_count: int = Field(
+        validation_alias=AliasChoices(
+            "num_hidden_layers",
+            "num_layers",
+            "n_layer",
+            "n_layers",
+            "num_decoder_layers",
+            "decoder_layers",
+        )
+    )
+
+    @property
+    def supports_tensor(self) -> bool:
+        return self.architectures in [
+            ["Glm4MoeLiteForCausalLM"],
+            ["GlmMoeDsaForCausalLM"],
+            ["DeepseekV32ForCausalLM"],
+            ["DeepseekV3ForCausalLM"],
+            ["Qwen3NextForCausalLM"],
+            ["Qwen3MoeForCausalLM"],
+            ["Qwen3_5MoeForConditionalGeneration"],
+            ["Qwen3_5ForConditionalGeneration"],
+            ["MiniMaxM2ForCausalLM"],
+            ["LlamaForCausalLM"],
+            ["GptOssForCausalLM"],
+            ["Step3p5ForCausalLM"],
+            ["NemotronHForCausalLM"],
+        ]
+
+    @model_validator(mode="before")
+    @classmethod
+    def defer_to_text_config(cls, data: dict[str, Any]):
+        text_config = data.get("text_config")
+        if text_config is None:
+            return data
+
+        for field in [
+            "architectures",
+            "hidden_size",
+            "num_key_value_heads",
+            "num_hidden_layers",
+            "num_layers",
+            "n_layer",
+            "n_layers",
+            "num_decoder_layers",
+            "decoder_layers",
+        ]:
+            if (val := text_config.get(field)) is not None:  # pyright: ignore[reportAny]
+                data[field] = val
+
+        return data

python/exo_core/src/exo_core/types/instances.py

@@ -1,11 +1,11 @@
 from enum import Enum
 
-from exo.shared.types.common import Host, Id, NodeId
-from exo.shared.types.worker.runners import RunnerId, ShardAssignments, ShardMetadata
 from pydantic import model_validator
 
 from exo_core.model_cards import ModelTask
-from exo_core.pydantic_ext import CamelCaseModel, TaggedModel
+from exo_core.models import CamelCaseModel, TaggedModel
+from exo_core.types.common import Host, Id, NodeId
+from exo_core.types.runners import RunnerId, ShardAssignments, ShardMetadata
 
 
 class InstanceId(Id):

python/exo_core/src/exo_core/types/runner_response.py

@@ -1,8 +1,6 @@
 from collections.abc import Generator
 from typing import Any, Literal
 
-from exo.utils.pydantic_ext import TaggedModel
-
 from exo.api.types import (
     FinishReason,
     GenerationStats,
@@ -11,6 +9,7 @@ from exo.api.types import (
     TopLogprobItem,
     Usage,
 )
+from exo_core.models import TaggedModel
 
 
 class BaseRunnerResponse(TaggedModel):

python/exo_core/src/exo_core/types/runners.py

@@ -3,7 +3,7 @@ from collections.abc import Mapping
 from pydantic import model_validator
 
 from exo_core.model_cards import ModelId
-from exo_core.pydantic import CamelCaseModel, TaggedModel
+from exo_core.models import CamelCaseModel, TaggedModel
 from exo_core.types.common import Id, NodeId
 from exo_core.types.shards import ShardMetadata
 

python/exo_core/src/exo_core/types/shards.py

@@ -4,7 +4,7 @@ from typing import TypeAlias, final
 from pydantic import Field
 
 from exo_core.model_cards import ModelCard
-from exo_core.pydantic import TaggedModel
+from exo_core.models import TaggedModel
 
 
 class Sharding(str, Enum):

python/exo_core/src/exo_core/types/tasks.py

@@ -6,7 +6,7 @@ from exo.api.types import (
     ImageEditsTaskParams,
     ImageGenerationTaskParams,
 )
-from exo_core.pydantic import TaggedModel
+from exo_core.models import TaggedModel
 from exo_core.types.common import CommandId, Id
 from exo_core.types.instances import BoundInstance, InstanceId
 from exo_core.types.runners import RunnerId

python/exo_core/src/exo_core/utils/downloads.py

@@ -0,0 +1,933 @@
+import asyncio
+import hashlib
+import os
+import shutil
+import ssl
+import time
+import traceback
+from collections.abc import Awaitable
+from datetime import timedelta
+from pathlib import Path
+from typing import Callable, Literal
+from urllib.parse import urljoin
+
+import aiofiles
+import aiofiles.os as aios
+import aiohttp
+import certifi
+from huggingface_hub import (
+    model_info,
+    snapshot_download,  # pyright: ignore[reportUnknownVariableType]
+)
+from loguru import logger
+from pydantic import (
+    TypeAdapter,
+)
+
+from exo_core.constants import EXO_MODELS_DIR, EXO_MODELS_PATH
+from exo_core.types.common import ModelId
+from exo_core.types.downloads import (
+    ConfigData,
+    DownloadProgressData,
+    FileListEntry,
+    ModelSafetensorsIndex,
+    RepoDownloadProgress,
+    RepoFileDownloadProgress,
+)
+from exo_core.types.shards import ShardMetadata
+from exo_core.utils.huggingface import (
+    filter_repo_objects,
+    get_allow_patterns,
+    get_auth_headers,
+    get_hf_endpoint,
+    get_hf_token,
+)
+from exo_core.utils.memory import Memory
+
+
+class HuggingFaceAuthenticationError(Exception):
+    """Raised when HuggingFace returns 401/403 for a model download."""
+
+
+class HuggingFaceRateLimitError(Exception):
+    """429 Huggingface code"""
+
+
+async def _build_auth_error_message(status_code: int, model_id: ModelId) -> str:
+    token = await get_hf_token()
+    if status_code == 401 and token is None:
+        return (
+            f"Model '{model_id}' requires authentication. "
+            f"Set HF_TOKEN in the app's Advanced settings, set the HF_TOKEN environment variable, or run `hf auth login`. "
+            f"Get a token at https://huggingface.co/settings/tokens"
+        )
+    elif status_code == 403:
+        return (
+            f"Access denied to '{model_id}'. "
+            f"Please accept the model terms at https://huggingface.co/{model_id}"
+        )
+    else:
+        return f"Authentication failed for '{model_id}' (HTTP {status_code})"
+
+
+def trim_etag(etag: str) -> str:
+    if (etag[0] == '"' and etag[-1] == '"') or (etag[0] == "'" and etag[-1] == "'"):
+        return etag[1:-1]
+    return etag
+
+
+def map_repo_file_download_progress_to_download_progress_data(
+    repo_file_download_progress: RepoFileDownloadProgress,
+) -> DownloadProgressData:
+    return DownloadProgressData(
+        downloaded=repo_file_download_progress.downloaded,
+        downloaded_this_session=repo_file_download_progress.downloaded_this_session,
+        total=repo_file_download_progress.total,
+        completed_files=1 if repo_file_download_progress.status == "complete" else 0,
+        total_files=1,
+        speed=repo_file_download_progress.speed,
+        eta_ms=int(repo_file_download_progress.eta.total_seconds() * 1000),
+        files={},
+    )
+
+
+def map_repo_download_progress_to_download_progress_data(
+    repo_download_progress: RepoDownloadProgress,
+) -> DownloadProgressData:
+    return DownloadProgressData(
+        total=repo_download_progress.total,
+        downloaded=repo_download_progress.downloaded,
+        downloaded_this_session=repo_download_progress.downloaded_this_session,
+        completed_files=repo_download_progress.completed_files,
+        total_files=repo_download_progress.total_files,
+        speed=repo_download_progress.overall_speed,
+        eta_ms=int(repo_download_progress.overall_eta.total_seconds() * 1000),
+        files={
+            file_path: map_repo_file_download_progress_to_download_progress_data(
+                file_progress
+            )
+            for file_path, file_progress in repo_download_progress.file_progress.items()
+        },
+    )
+
+
+def resolve_model_in_path(model_id: ModelId) -> Path | None:
+    """Search EXO_MODELS_PATH directories for a pre-existing model.
+
+    Checks each directory for the normalized name (org--model).  A candidate
+    is only returned if ``is_model_directory_complete`` confirms all weight
+    files are present.
+    """
+    if EXO_MODELS_PATH is None:
+        return None
+    normalized = model_id.normalize()
+    for search_dir in EXO_MODELS_PATH:
+        candidate = search_dir / normalized
+        if candidate.is_dir() and is_model_directory_complete(candidate):
+            return candidate
+    return None
+
+
+def build_model_path(model_id: ModelId) -> Path:
+    found = resolve_model_in_path(model_id)
+    if found is not None:
+        return found
+    return EXO_MODELS_DIR / model_id.normalize()
+
+
+async def resolve_model_path_for_repo(model_id: ModelId) -> Path:
+    return (await ensure_models_dir()) / model_id.normalize()
+
+
+async def ensure_models_dir() -> Path:
+    await aios.makedirs(EXO_MODELS_DIR, exist_ok=True)
+    return EXO_MODELS_DIR
+
+
+async def delete_model(model_id: ModelId) -> bool:
+    models_dir = await ensure_models_dir()
+    model_dir = models_dir / model_id.normalize()
+    cache_dir = models_dir / "caches" / model_id.normalize()
+
+    deleted = False
+    if await aios.path.exists(model_dir):
+        await asyncio.to_thread(shutil.rmtree, model_dir, ignore_errors=False)
+        deleted = True
+
+    # Also clear cache
+    if await aios.path.exists(cache_dir):
+        await asyncio.to_thread(shutil.rmtree, cache_dir, ignore_errors=False)
+
+    return deleted
+
+
+async def seed_models(seed_dir: str | Path):
+    """Move models from resources folder to EXO_MODELS_DIR."""
+    source_dir = Path(seed_dir)
+    dest_dir = await ensure_models_dir()
+    for path in source_dir.iterdir():
+        if path.is_dir() and path.name.startswith("models--"):
+            dest_path = dest_dir / path.name
+            if await aios.path.exists(dest_path):
+                logger.info("Skipping moving model to .cache directory")
+            else:
+                try:
+                    await aios.rename(str(path), str(dest_path))
+                except Exception:
+                    logger.error(f"Error seeding model {path} to {dest_path}")
+                    logger.error(traceback.format_exc())
+
+
+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``.
+    """
+    index_files = list(model_dir.glob("**/*.safetensors.index.json"))
+    if not index_files:
+        return None
+
+    entries_by_path: dict[str, FileListEntry] = {}
+
+    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 Exception:
+            continue
+
+    return list(entries_by_path.values())
+
+
+def is_model_directory_complete(model_dir: Path) -> bool:
+    """Check if a model directory contains all required weight files."""
+    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)
+
+
+async def _build_file_list_from_local_directory(
+    model_id: ModelId,
+    recursive: bool = False,
+) -> 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.
+    """
+    model_dir = (await ensure_models_dir()) / model_id.normalize()
+    if not await aios.path.exists(model_dir):
+        return None
+
+    file_list = await asyncio.to_thread(_scan_model_directory, model_dir, recursive)
+    if not file_list:
+        return None
+    return file_list
+
+
+_fetched_file_lists_this_session: set[str] = set()
+
+
+async def fetch_file_list_with_cache(
+    model_id: ModelId,
+    revision: str = "main",
+    recursive: bool = False,
+    skip_internet: bool = False,
+    on_connection_lost: Callable[[], None] = lambda: None,
+) -> list[FileListEntry]:
+    target_dir = (await ensure_models_dir()) / "caches" / model_id.normalize()
+    await aios.makedirs(target_dir, exist_ok=True)
+    cache_file = target_dir / f"{model_id.normalize()}--{revision}--file_list.json"
+    cache_key = f"{model_id.normalize()}--{revision}"
+
+    if cache_key in _fetched_file_lists_this_session and await aios.path.exists(
+        cache_file
+    ):
+        async with aiofiles.open(cache_file, "r") as f:
+            return TypeAdapter(list[FileListEntry]).validate_json(await f.read())
+
+    if skip_internet:
+        if await aios.path.exists(cache_file):
+            async with aiofiles.open(cache_file, "r") as f:
+                return TypeAdapter(list[FileListEntry]).validate_json(await f.read())
+        local_file_list = await _build_file_list_from_local_directory(
+            model_id, recursive
+        )
+        if local_file_list is not None:
+            logger.warning(
+                f"No internet and no cached file list for {model_id} - using local file list"
+            )
+            return local_file_list
+        raise FileNotFoundError(
+            f"No internet connection and no cached file list for {model_id}"
+        )
+
+    try:
+        file_list = await fetch_file_list_with_retry(
+            model_id,
+            revision,
+            recursive=recursive,
+            on_connection_lost=on_connection_lost,
+        )
+        async with aiofiles.open(cache_file, "w") as f:
+            await f.write(
+                TypeAdapter(list[FileListEntry]).dump_json(file_list).decode()
+            )
+        _fetched_file_lists_this_session.add(cache_key)
+        return file_list
+    except Exception as e:
+        logger.opt(exception=e).warning(
+            "Ran into exception when fetching file list from HF."
+        )
+
+        if await aios.path.exists(cache_file):
+            logger.warning(
+                f"No cached file list for {model_id} - using local file list"
+            )
+            async with aiofiles.open(cache_file, "r") as f:
+                return TypeAdapter(list[FileListEntry]).validate_json(await f.read())
+        local_file_list = await _build_file_list_from_local_directory(
+            model_id, recursive
+        )
+        if local_file_list is not None:
+            logger.warning(
+                f"Failed to fetch file list for {model_id} and no cache exists, "
+            )
+            return local_file_list
+        raise FileNotFoundError(f"Failed to fetch file list for {model_id}: {e}") from e
+
+
+async def fetch_file_list_with_retry(
+    model_id: ModelId,
+    revision: str = "main",
+    path: str = "",
+    recursive: bool = False,
+    on_connection_lost: Callable[[], None] = lambda: None,
+) -> list[FileListEntry]:
+    n_attempts = 3
+    for attempt in range(n_attempts):
+        try:
+            return await _fetch_file_list(model_id, revision, path, recursive)
+        except HuggingFaceAuthenticationError:
+            raise
+        except Exception as e:
+            on_connection_lost()
+            if attempt == n_attempts - 1:
+                raise e
+            await asyncio.sleep(2.0**attempt)
+    raise Exception(
+        f"Failed to fetch file list for {model_id=} {revision=} {path=} {recursive=}"
+    )
+
+
+async def _fetch_file_list(
+    model_id: ModelId, revision: str = "main", path: str = "", recursive: bool = False
+) -> list[FileListEntry]:
+    api_url = f"{get_hf_endpoint()}/api/models/{model_id}/tree/{revision}"
+    url = f"{api_url}/{path}" if path else api_url
+
+    headers = await get_download_headers()
+    async with (
+        create_http_session(timeout_profile="short") as session,
+        session.get(url, headers=headers) as response,
+    ):
+        if response.status in [401, 403]:
+            msg = await _build_auth_error_message(response.status, model_id)
+            raise HuggingFaceAuthenticationError(msg)
+        elif response.status == 429:
+            raise HuggingFaceRateLimitError(
+                f"Couldn't download {model_id} because of HuggingFace rate limit."
+            )
+        elif response.status == 200:
+            data_json = await response.text()
+            data = TypeAdapter(list[FileListEntry]).validate_json(data_json)
+            files: list[FileListEntry] = []
+            for item in data:
+                if item.type == "file":
+                    files.append(FileListEntry.model_validate(item))
+                elif item.type == "directory" and recursive:
+                    subfiles = await _fetch_file_list(
+                        model_id, revision, item.path, recursive
+                    )
+                    files.extend(subfiles)
+            return files
+        else:
+            raise Exception(f"Failed to fetch file list: {response.status}")
+
+
+async def get_download_headers() -> dict[str, str]:
+    return {**(await get_auth_headers()), "Accept-Encoding": "identity"}
+
+
+def create_http_session(
+    auto_decompress: bool = False,
+    timeout_profile: Literal["short", "long"] = "long",
+) -> aiohttp.ClientSession:
+    if timeout_profile == "short":
+        total_timeout = 30
+        connect_timeout = 10
+        sock_read_timeout = 30
+        sock_connect_timeout = 10
+    else:
+        total_timeout = 1800
+        connect_timeout = 60
+        sock_read_timeout = 60
+        sock_connect_timeout = 60
+
+    ssl_context = ssl.create_default_context(
+        cafile=os.getenv("SSL_CERT_FILE") or certifi.where()
+    )
+    connector = aiohttp.TCPConnector(ssl=ssl_context)
+
+    return aiohttp.ClientSession(
+        auto_decompress=auto_decompress,
+        connector=connector,
+        proxy=os.getenv("HTTPS_PROXY") or os.getenv("HTTP_PROXY") or None,
+        timeout=aiohttp.ClientTimeout(
+            total=total_timeout,
+            connect=connect_timeout,
+            sock_read=sock_read_timeout,
+            sock_connect=sock_connect_timeout,
+        ),
+    )
+
+
+async def calc_hash(path: Path, hash_type: Literal["sha1", "sha256"] = "sha1") -> str:
+    hasher = hashlib.sha1() if hash_type == "sha1" else hashlib.sha256()
+    if hash_type == "sha1":
+        header = f"blob {(await aios.stat(path)).st_size}\0".encode()
+        hasher.update(header)
+    async with aiofiles.open(path, "rb") as f:
+        while chunk := await f.read(8 * 1024 * 1024):
+            hasher.update(chunk)
+    return hasher.hexdigest()
+
+
+async def file_meta(
+    model_id: ModelId, revision: str, path: str, redirected_location: str | None = None
+) -> tuple[int, str]:
+    url = (
+        urljoin(f"{get_hf_endpoint()}/{model_id}/resolve/{revision}/", path)
+        if redirected_location is None
+        else f"{get_hf_endpoint()}{redirected_location}"
+    )
+    headers = await get_download_headers()
+    async with (
+        create_http_session(timeout_profile="short") as session,
+        session.head(url, headers=headers) as r,
+    ):
+        if r.status == 307:
+            # On redirect, only trust Hugging Face's x-linked-* headers.
+            x_linked_size = r.headers.get("x-linked-size")
+            x_linked_etag = r.headers.get("x-linked-etag")
+            if x_linked_size and x_linked_etag:
+                content_length = int(x_linked_size)
+                etag = trim_etag(x_linked_etag)
+                return content_length, etag
+            # Otherwise, follow the redirect to get authoritative size/hash
+            redirected_location = r.headers.get("location")
+            return await file_meta(model_id, revision, path, redirected_location)
+        if r.status in [401, 403]:
+            msg = await _build_auth_error_message(r.status, model_id)
+            raise HuggingFaceAuthenticationError(msg)
+        content_length = int(
+            r.headers.get("x-linked-size") or r.headers.get("content-length") or 0
+        )
+        etag = r.headers.get("x-linked-etag") or r.headers.get("etag")
+        assert content_length > 0, f"No content length for {url}"
+        assert etag is not None, f"No remote hash for {url}"
+        etag = trim_etag(etag)
+        return content_length, etag
+
+
+async def download_file_with_retry(
+    model_id: ModelId,
+    revision: str,
+    path: str,
+    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, skip_internet
+            )
+        except HuggingFaceAuthenticationError:
+            raise
+        except FileNotFoundError:
+            raise
+        except HuggingFaceRateLimitError as e:
+            if attempt == n_attempts - 1:
+                raise e
+            logger.error(
+                f"Download error on attempt {attempt}/{n_attempts} for {model_id=} {revision=} {path=} {target_dir=}"
+            )
+            logger.error(traceback.format_exc())
+            await asyncio.sleep(2.0**attempt)
+        except Exception as e:
+            if attempt == n_attempts - 1:
+                on_connection_lost()
+                raise e
+            logger.error(
+                f"Download error on attempt {attempt + 1}/{n_attempts} for {model_id=} {revision=} {path=} {target_dir=}"
+            )
+            logger.error(traceback.format_exc())
+            await asyncio.sleep(2.0**attempt)
+    raise Exception(
+        f"Failed to download file {model_id=} {revision=} {path=} {target_dir=}"
+    )
+
+
+async def _download_file(
+    model_id: ModelId,
+    revision: str,
+    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
+        try:
+            remote_size, _ = await file_meta(model_id, revision, path)
+            if local_size != remote_size:
+                logger.info(
+                    f"File {path} size mismatch (local={local_size}, remote={remote_size}), re-downloading"
+                )
+                await aios.remove(target_path)
+            else:
+                return target_path
+        except Exception as e:
+            # Offline or network error - trust local file
+            logger.debug(
+                f"Could not verify {path} against remote (offline?): {e}, using local 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
+    partial_path = target_dir / f"{path}.partial"
+    resume_byte_pos = (
+        (await aios.stat(partial_path)).st_size
+        if (await aios.path.exists(partial_path))
+        else None
+    )
+    if resume_byte_pos != length:
+        url = urljoin(f"{get_hf_endpoint()}/{model_id}/resolve/{revision}/", path)
+        headers = await get_download_headers()
+        if resume_byte_pos:
+            headers["Range"] = f"bytes={resume_byte_pos}-"
+        n_read = resume_byte_pos or 0
+        async with (
+            create_http_session(timeout_profile="long") as session,
+            session.get(url, headers=headers) as r,
+        ):
+            if r.status == 404:
+                raise FileNotFoundError(f"File not found: {url}")
+            if r.status in [401, 403]:
+                msg = await _build_auth_error_message(r.status, model_id)
+                raise HuggingFaceAuthenticationError(msg)
+            assert r.status in [200, 206], (
+                f"Failed to download {path} from {url}: {r.status}"
+            )
+            async with aiofiles.open(
+                partial_path, "ab" if resume_byte_pos else "wb"
+            ) as f:
+                while chunk := await r.content.read(8 * 1024 * 1024):
+                    n_read = n_read + (await f.write(chunk))
+                    on_progress(n_read, length, False)
+
+    final_hash = await calc_hash(
+        partial_path, hash_type="sha256" if len(remote_hash) == 64 else "sha1"
+    )
+    integrity = final_hash == remote_hash
+    if not integrity:
+        try:
+            await aios.remove(partial_path)
+        except Exception as e:
+            logger.error(f"Error removing partial file {partial_path}: {e}")
+        raise Exception(
+            f"Downloaded file {target_dir / path} has hash {final_hash} but remote hash is {remote_hash}"
+        )
+    await aios.rename(partial_path, target_dir / path)
+    on_progress(length, length, True)
+    return target_dir / path
+
+
+def calculate_repo_progress(
+    shard: ShardMetadata,
+    model_id: ModelId,
+    revision: str,
+    file_progress: dict[str, RepoFileDownloadProgress],
+    all_start_time: float,
+) -> RepoDownloadProgress:
+    all_total = sum((p.total for p in file_progress.values()), Memory.from_bytes(0))
+    all_downloaded = sum(
+        (p.downloaded for p in file_progress.values()), Memory.from_bytes(0)
+    )
+    all_downloaded_this_session = sum(
+        (p.downloaded_this_session for p in file_progress.values()),
+        Memory.from_bytes(0),
+    )
+    elapsed_time = time.time() - all_start_time
+    all_speed = (
+        all_downloaded_this_session.in_bytes / elapsed_time if elapsed_time > 0 else 0
+    )
+    all_eta = (
+        timedelta(seconds=(all_total - all_downloaded).in_bytes / all_speed)
+        if all_speed > 0
+        else timedelta(seconds=0)
+    )
+    status = (
+        "complete"
+        if all(p.status == "complete" for p in file_progress.values())
+        else "in_progress"
+        if any(p.status == "in_progress" for p in file_progress.values())
+        else "not_started"
+    )
+    return RepoDownloadProgress(
+        repo_id=model_id,
+        repo_revision=revision,
+        shard=shard,
+        completed_files=len(
+            [p for p in file_progress.values() if p.downloaded == p.total]
+        ),
+        total_files=len(file_progress),
+        downloaded=all_downloaded,
+        downloaded_this_session=all_downloaded_this_session,
+        total=all_total,
+        overall_speed=all_speed,
+        overall_eta=all_eta,
+        status=status,
+        file_progress=file_progress,
+    )
+
+
+async def get_weight_map(model_id: ModelId, revision: str = "main") -> dict[str, str]:
+    target_dir = (await ensure_models_dir()) / model_id.normalize()
+    await aios.makedirs(target_dir, exist_ok=True)
+
+    index_files_dir = snapshot_download(
+        repo_id=model_id,
+        local_dir=target_dir,
+        allow_patterns="*.safetensors.index.json",
+    )
+
+    index_files = list(Path(index_files_dir).glob("**/*.safetensors.index.json"))
+
+    weight_map: dict[str, str] = {}
+
+    for index_file in index_files:
+        relative_dir = index_file.parent.relative_to(index_files_dir)
+
+        async with aiofiles.open(index_file, "r") as f:
+            index_data = ModelSafetensorsIndex.model_validate_json(await f.read())
+
+            if relative_dir != Path("."):
+                prefixed_weight_map = {
+                    f"{relative_dir}/{key}": str(relative_dir / value)
+                    for key, value in index_data.weight_map.items()
+                }
+                weight_map = weight_map | prefixed_weight_map
+            else:
+                weight_map = weight_map | index_data.weight_map
+
+    return weight_map
+
+
+async def resolve_allow_patterns(shard: ShardMetadata) -> list[str]:
+    # TODO: 'Smart' downloads are disabled because:
+    #  (i) We don't handle all kinds of files;
+    # (ii) We don't have sticky sessions.
+    # (iii) Tensor parallel requires all files.
+    return ["*"]
+    try:
+        weight_map = await get_weight_map(str(shard.model_card.model_id))
+        return get_allow_patterns(weight_map, shard)
+    except Exception:
+        logger.error(f"Error getting weight map for {shard.model_card.model_id=}")
+        logger.error(traceback.format_exc())
+        return ["*"]
+
+
+def _is_image_model(shard: ShardMetadata) -> bool:
+    from exo_core.model_cards import ModelTask
+
+    tasks = shard.model_card.tasks
+    return ModelTask.TextToImage in tasks or ModelTask.ImageToImage in tasks
+
+
+async def get_downloaded_size(path: Path) -> int:
+    partial_path = path.with_suffix(path.suffix + ".partial")
+    if await aios.path.exists(path):
+        return (await aios.stat(path)).st_size
+    if await aios.path.exists(partial_path):
+        return (await aios.stat(partial_path)).st_size
+    return 0
+
+
+async def download_shard(
+    shard: ShardMetadata,
+    on_progress: Callable[[ShardMetadata, RepoDownloadProgress], Awaitable[None]],
+    max_parallel_downloads: int = 8,
+    skip_download: bool = False,
+    skip_internet: bool = False,
+    allow_patterns: list[str] | None = None,
+    on_connection_lost: Callable[[], None] = lambda: None,
+) -> tuple[Path, RepoDownloadProgress]:
+    if not skip_download:
+        logger.debug(f"Downloading {shard.model_card.model_id=}")
+
+    revision = "main"
+    target_dir = await ensure_models_dir() / str(shard.model_card.model_id).replace(
+        "/", "--"
+    )
+    if not skip_download:
+        await aios.makedirs(target_dir, exist_ok=True)
+
+    if not allow_patterns:
+        allow_patterns = await resolve_allow_patterns(shard)
+
+    if not skip_download:
+        logger.debug(f"Downloading {shard.model_card.model_id=} with {allow_patterns=}")
+
+    all_start_time = time.time()
+    file_list = await fetch_file_list_with_cache(
+        shard.model_card.model_id,
+        revision,
+        recursive=True,
+        skip_internet=skip_internet,
+        on_connection_lost=on_connection_lost,
+    )
+    filtered_file_list = list(
+        filter_repo_objects(
+            file_list,
+            allow_patterns=allow_patterns,
+            ignore_patterns=["original/*", "metal/*"],
+            key=lambda x: x.path,
+        )
+    )
+
+    # For image models, skip root-level safetensors files since weights
+    # are stored in component subdirectories (e.g., transformer/, vae/)
+    if _is_image_model(shard):
+        filtered_file_list = [
+            f
+            for f in filtered_file_list
+            if "/" in f.path or not f.path.endswith(".safetensors")
+        ]
+    file_progress: dict[str, RepoFileDownloadProgress] = {}
+
+    async def on_progress_wrapper(
+        file: FileListEntry, curr_bytes: int, total_bytes: int, is_renamed: bool
+    ) -> None:
+        previous_progress = file_progress.get(file.path)
+
+        # Detect re-download: curr_bytes < previous downloaded means file was deleted and restarted
+        is_redownload = (
+            previous_progress is not None
+            and curr_bytes < previous_progress.downloaded.in_bytes
+        )
+
+        if is_redownload or previous_progress is None:
+            # Fresh download or re-download: reset tracking
+            start_time = time.time()
+            downloaded_this_session = curr_bytes
+        else:
+            # Continuing download: accumulate
+            start_time = previous_progress.start_time
+            downloaded_this_session = (
+                previous_progress.downloaded_this_session.in_bytes
+                + (curr_bytes - previous_progress.downloaded.in_bytes)
+            )
+
+        speed = (
+            downloaded_this_session / (time.time() - start_time)
+            if time.time() - start_time > 0
+            else 0
+        )
+        eta = (
+            timedelta(seconds=(total_bytes - curr_bytes) / speed)
+            if speed > 0
+            else timedelta(seconds=0)
+        )
+        file_progress[file.path] = RepoFileDownloadProgress(
+            repo_id=shard.model_card.model_id,
+            repo_revision=revision,
+            file_path=file.path,
+            downloaded=Memory.from_bytes(curr_bytes),
+            downloaded_this_session=Memory.from_bytes(downloaded_this_session),
+            total=Memory.from_bytes(total_bytes),
+            speed=speed,
+            eta=eta,
+            status="complete"
+            if curr_bytes == total_bytes and is_renamed
+            else "in_progress",
+            start_time=start_time,
+        )
+        await on_progress(
+            shard,
+            calculate_repo_progress(
+                shard,
+                shard.model_card.model_id,
+                revision,
+                file_progress,
+                all_start_time,
+            ),
+        )
+
+    for file in filtered_file_list:
+        downloaded_bytes = await get_downloaded_size(target_dir / file.path)
+        final_file_exists = await aios.path.exists(target_dir / file.path)
+        file_progress[file.path] = RepoFileDownloadProgress(
+            repo_id=shard.model_card.model_id,
+            repo_revision=revision,
+            file_path=file.path,
+            downloaded=Memory.from_bytes(downloaded_bytes),
+            downloaded_this_session=Memory.from_bytes(0),
+            total=Memory.from_bytes(file.size or 0),
+            speed=0,
+            eta=timedelta(0),
+            status="complete"
+            if final_file_exists and downloaded_bytes == file.size
+            else "not_started",
+            start_time=time.time(),
+        )
+
+    semaphore = asyncio.Semaphore(max_parallel_downloads)
+
+    def schedule_progress(
+        file: FileListEntry, curr_bytes: int, total_bytes: int, is_renamed: bool
+    ) -> None:
+        asyncio.create_task(
+            on_progress_wrapper(file, curr_bytes, total_bytes, is_renamed)
+        )
+
+    async def download_with_semaphore(file: FileListEntry) -> None:
+        async with semaphore:
+            await download_file_with_retry(
+                shard.model_card.model_id,
+                revision,
+                file.path,
+                target_dir,
+                lambda curr_bytes, total_bytes, is_renamed: schedule_progress(
+                    file, curr_bytes, total_bytes, is_renamed
+                ),
+                on_connection_lost=on_connection_lost,
+                skip_internet=skip_internet,
+            )
+
+    if not skip_download:
+        await asyncio.gather(
+            *[download_with_semaphore(file) for file in filtered_file_list]
+        )
+    final_repo_progress = calculate_repo_progress(
+        shard, shard.model_card.model_id, revision, file_progress, all_start_time
+    )
+    await on_progress(shard, final_repo_progress)
+    if gguf := next((f for f in filtered_file_list if f.path.endswith(".gguf")), None):
+        return target_dir / gguf.path, final_repo_progress
+    else:
+        return target_dir, final_repo_progress
+
+
+async def fetch_config_data(model_id: ModelId) -> ConfigData:
+    """Downloads and parses config.json for a model."""
+
+    target_dir = (await ensure_models_dir()) / model_id.normalize()
+    await aios.makedirs(target_dir, exist_ok=True)
+    config_path = await download_file_with_retry(
+        model_id,
+        "main",
+        "config.json",
+        target_dir,
+        lambda curr_bytes, total_bytes, is_renamed: logger.debug(
+            f"Downloading config.json for {model_id}: {curr_bytes}/{total_bytes} ({is_renamed=})"
+        ),
+    )
+    async with aiofiles.open(config_path, "r") as f:
+        return ConfigData.model_validate_json(await f.read())
+
+
+async def fetch_safetensors_size(model_id: ModelId) -> Memory:
+    """Gets model size from safetensors index or falls back to HF API."""
+    target_dir = (await ensure_models_dir()) / model_id.normalize()
+    await aios.makedirs(target_dir, exist_ok=True)
+    try:
+        index_path = await download_file_with_retry(
+            model_id,
+            "main",
+            "model.safetensors.index.json",
+            target_dir,
+            lambda curr_bytes, total_bytes, is_renamed: logger.debug(
+                f"Downloading model.safetensors.index.json for {model_id}: {curr_bytes}/{total_bytes} ({is_renamed=})"
+            ),
+        )
+        async with aiofiles.open(index_path, "r") as f:
+            index_data = ModelSafetensorsIndex.model_validate_json(await f.read())
+
+        metadata = index_data.metadata
+        if metadata is not None:
+            return Memory.from_bytes(metadata.total_size)
+    except FileNotFoundError:
+        pass
+
+    info = model_info(model_id)
+    if info.safetensors is None:
+        raise ValueError(f"No safetensors info found for {model_id}")
+    return Memory.from_bytes(info.safetensors.total)

python/exo_core/src/exo_core/utils/huggingface.py

@@ -0,0 +1,171 @@
+import os
+from fnmatch import fnmatch
+from pathlib import Path
+from typing import Callable, Generator, Iterable
+
+import aiofiles
+import aiofiles.os as aios
+from loguru import logger
+
+from exo_core.types.shards import ShardMetadata
+
+
+def filter_repo_objects[T](
+    items: Iterable[T],
+    *,
+    allow_patterns: list[str] | str | None = None,
+    ignore_patterns: list[str] | str | None = None,
+    key: Callable[[T], str] | None = None,
+) -> Generator[T, None, None]:
+    if isinstance(allow_patterns, str):
+        allow_patterns = [allow_patterns]
+    if isinstance(ignore_patterns, str):
+        ignore_patterns = [ignore_patterns]
+    if allow_patterns is not None:
+        allow_patterns = [_add_wildcard_to_directories(p) for p in allow_patterns]
+    if ignore_patterns is not None:
+        ignore_patterns = [_add_wildcard_to_directories(p) for p in ignore_patterns]
+
+    if key is None:
+
+        def _identity(item: T) -> str:
+            if isinstance(item, str):
+                return item
+            if isinstance(item, Path):
+                return str(item)
+            raise ValueError(
+                f"Please provide `key` argument in `filter_repo_objects`: `{item}` is not a string."
+            )
+
+        key = _identity
+
+    for item in items:
+        path = key(item)
+        if allow_patterns is not None and not any(
+            fnmatch(path, r) for r in allow_patterns
+        ):
+            continue
+        if ignore_patterns is not None and any(
+            fnmatch(path, r) for r in ignore_patterns
+        ):
+            continue
+        yield item
+
+
+def _add_wildcard_to_directories(pattern: str) -> str:
+    if pattern[-1] == "/":
+        return pattern + "*"
+    return pattern
+
+
+def get_hf_endpoint() -> str:
+    return os.environ.get("HF_ENDPOINT", "https://huggingface.co")
+
+
+def get_hf_home() -> Path:
+    """Get the Hugging Face home directory."""
+    return Path(os.environ.get("HF_HOME", Path.home() / ".cache" / "huggingface"))
+
+
+async def get_hf_token() -> str | None:
+    """Retrieve the Hugging Face token from HF_TOKEN env var or HF_HOME directory."""
+    # Check environment variable first
+    if token := os.environ.get("HF_TOKEN"):
+        return token
+    # Fall back to file-based token
+    token_path = get_hf_home() / "token"
+    if await aios.path.exists(token_path):
+        async with aiofiles.open(token_path, "r") as f:
+            return (await f.read()).strip()
+    return None
+
+
+async def get_auth_headers() -> dict[str, str]:
+    """Get authentication headers if a token is available."""
+    token = await get_hf_token()
+    if token:
+        return {"Authorization": f"Bearer {token}"}
+    return {}
+
+
+def extract_layer_num(tensor_name: str) -> int | None:
+    # This is a simple example and might need to be adjusted based on the actual naming convention
+    parts = tensor_name.split(".")
+    for part in parts:
+        if part.isdigit():
+            return int(part)
+    return None
+
+
+def get_allow_patterns(weight_map: dict[str, str], shard: ShardMetadata) -> list[str]:
+    default_patterns = set(
+        [
+            "*.json",
+            "*.py",
+            "tokenizer.model",
+            "tiktoken.model",
+            "*/spiece.model",
+            "*.tiktoken",
+            "*.txt",
+            "*.jinja",
+        ]
+    )
+    shard_specific_patterns: set[str] = set()
+
+    if shard.model_card.components is not None:
+        shardable_component = next(
+            (c for c in shard.model_card.components if c.can_shard), None
+        )
+
+        if weight_map and shardable_component:
+            for tensor_name, filename in weight_map.items():
+                # Strip component prefix from tensor name (added by weight map namespacing)
+                # E.g., "transformer/blocks.0.weight" -> "blocks.0.weight"
+                if "/" in tensor_name:
+                    _, tensor_name_no_prefix = tensor_name.split("/", 1)
+                else:
+                    tensor_name_no_prefix = tensor_name
+
+                # Determine which component this file belongs to from filename
+                component_path = Path(filename).parts[0] if "/" in filename else None
+
+                if component_path == shardable_component.component_path.rstrip("/"):
+                    layer_num = extract_layer_num(tensor_name_no_prefix)
+                    if (
+                        layer_num is not None
+                        and shard.start_layer <= layer_num < shard.end_layer
+                    ):
+                        shard_specific_patterns.add(filename)
+
+                    if shard.is_first_layer or shard.is_last_layer:
+                        shard_specific_patterns.add(filename)
+                else:
+                    shard_specific_patterns.add(filename)
+
+        else:
+            shard_specific_patterns = set(["*.safetensors"])
+
+        # TODO(ciaran): temporary - Include all files from non-shardable components that have no index file
+        for component in shard.model_card.components:
+            if not component.can_shard and component.safetensors_index_filename is None:
+                component_pattern = f"{component.component_path.rstrip('/')}/*"
+                shard_specific_patterns.add(component_pattern)
+    else:
+        if weight_map:
+            for tensor_name, filename in weight_map.items():
+                layer_num = extract_layer_num(tensor_name)
+                if (
+                    layer_num is not None
+                    and shard.start_layer <= layer_num < shard.end_layer
+                ):
+                    shard_specific_patterns.add(filename)
+            layer_independent_files = set(
+                [v for k, v in weight_map.items() if extract_layer_num(k) is None]
+            )
+            shard_specific_patterns.update(layer_independent_files)
+            logger.debug(f"get_allow_patterns {shard=} {layer_independent_files=}")
+        else:
+            shard_specific_patterns = set(["*.safetensors"])
+
+    logger.info(f"get_allow_patterns {shard=} {shard_specific_patterns=}")
+    return list(default_patterns | shard_specific_patterns)

python/exo_core/src/exo_core/utils/memory.py

@@ -1,7 +1,7 @@
 from math import ceil
 from typing import Self, overload
 
-from exo_core.pydantic import FrozenModel
+from exo_core.models import FrozenModel
 
 
 class Memory(FrozenModel):

python/exo_core/src/exo_core/utils/paths.py

@@ -0,0 +1,59 @@
+import sys
+from pathlib import Path
+from typing import cast
+
+
+def find_resources() -> Path:
+    resources = _find_resources_in_repo() or _find_resources_in_bundle()
+    if resources is None:
+        raise FileNotFoundError(
+            "Unable to locate resources. Did you clone the repo properly?"
+        )
+    return resources
+
+
+def _find_resources_in_repo() -> Path | None:
+    current_module = Path(__file__).resolve()
+    for parent in current_module.parents:
+        build = parent / "resources"
+        if build.is_dir():
+            return build
+    return None
+
+
+def _find_resources_in_bundle() -> Path | None:
+    frozen_root = cast(str | None, getattr(sys, "_MEIPASS", None))
+    if frozen_root is None:
+        return None
+    candidate = Path(frozen_root) / "resources"
+    if candidate.is_dir():
+        return candidate
+    return None
+
+
+def find_dashboard() -> Path:
+    dashboard = _find_dashboard_in_repo() or _find_dashboard_in_bundle()
+    if not dashboard:
+        raise FileNotFoundError(
+            "Unable to locate dashboard assets - you probably forgot to run `cd dashboard && npm install && npm run build && cd ..`"
+        )
+    return dashboard
+
+
+def _find_dashboard_in_repo() -> Path | None:
+    current_module = Path(__file__).resolve()
+    for parent in current_module.parents:
+        build = parent / "dashboard" / "build"
+        if build.is_dir() and (build / "index.html").exists():
+            return build
+    return None
+
+
+def _find_dashboard_in_bundle() -> Path | None:
+    frozen_root = cast(str | None, getattr(sys, "_MEIPASS", None))
+    if frozen_root is None:
+        return None
+    candidate = Path(frozen_root) / "dashboard"
+    if candidate.is_dir():
+        return candidate
+    return None

python/mlx_engine/pyproject.toml

@@ -4,8 +4,11 @@ version = "0.1.0"
 description = "Add your description here"
 readme = "README.md"
 requires-python = ">=3.13"
-dependencies = []
+dependencies = ["exo_core"]
 
 [build-system]
 requires = ["uv_build>=0.8.9,<0.9.0"]
 build-backend = "uv_build"
+
+[tool.uv.sources]
+exo_core = { workspace = true }

python/mlx_engine/uv.lock

@@ -0,0 +1,123 @@
+version = 1
+revision = 3
+requires-python = ">=3.13"
+
+[[package]]
+name = "annotated-types"
+version = "0.7.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/ee/67/531ea369ba64dcff5ec9c3402f9f51bf748cec26dde048a2f973a4eea7f5/annotated_types-0.7.0.tar.gz", hash = "sha256:aff07c09a53a08bc8cfccb9c85b05f1aa9a2a6f23728d790723543408344ce89", size = 16081, upload-time = "2024-05-20T21:33:25.928Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/78/b6/6307fbef88d9b5ee7421e68d78a9f162e0da4900bc5f5793f6d3d0e34fb8/annotated_types-0.7.0-py3-none-any.whl", hash = "sha256:1f02e8b43a8fbbc3f3e0d4f0f4bfc8131bcb4eebe8849b8e5c773f3a1c582a53", size = 13643, upload-time = "2024-05-20T21:33:24.1Z" },
+]
+
+[[package]]
+name = "exo-core"
+version = "0.1.0"
+source = { directory = "../exo_core" }
+dependencies = [
+    { name = "pydantic" },
+]
+
+[package.metadata]
+requires-dist = [{ name = "pydantic" }]
+
+[[package]]
+name = "mlx-engine"
+version = "0.1.0"
+source = { editable = "." }
+dependencies = [
+    { name = "exo-core" },
+]
+
+[package.metadata]
+requires-dist = [{ name = "exo-core", directory = "../exo_core" }]
+
+[[package]]
+name = "pydantic"
+version = "2.12.5"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "annotated-types" },
+    { name = "pydantic-core" },
+    { name = "typing-extensions" },
+    { name = "typing-inspection" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/69/44/36f1a6e523abc58ae5f928898e4aca2e0ea509b5aa6f6f392a5d882be928/pydantic-2.12.5.tar.gz", hash = "sha256:4d351024c75c0f085a9febbb665ce8c0c6ec5d30e903bdb6394b7ede26aebb49", size = 821591, upload-time = "2025-11-26T15:11:46.471Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/5a/87/b70ad306ebb6f9b585f114d0ac2137d792b48be34d732d60e597c2f8465a/pydantic-2.12.5-py3-none-any.whl", hash = "sha256:e561593fccf61e8a20fc46dfc2dfe075b8be7d0188df33f221ad1f0139180f9d", size = 463580, upload-time = "2025-11-26T15:11:44.605Z" },
+]
+
+[[package]]
+name = "pydantic-core"
+version = "2.41.5"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "typing-extensions" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/71/70/23b021c950c2addd24ec408e9ab05d59b035b39d97cdc1130e1bce647bb6/pydantic_core-2.41.5.tar.gz", hash = "sha256:08daa51ea16ad373ffd5e7606252cc32f07bc72b28284b6bc9c6df804816476e", size = 460952, upload-time = "2025-11-04T13:43:49.098Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/87/06/8806241ff1f70d9939f9af039c6c35f2360cf16e93c2ca76f184e76b1564/pydantic_core-2.41.5-cp313-cp313-macosx_10_12_x86_64.whl", hash = "sha256:941103c9be18ac8daf7b7adca8228f8ed6bb7a1849020f643b3a14d15b1924d9", size = 2120403, upload-time = "2025-11-04T13:40:25.248Z" },
+    { url = "https://files.pythonhosted.org/packages/94/02/abfa0e0bda67faa65fef1c84971c7e45928e108fe24333c81f3bfe35d5f5/pydantic_core-2.41.5-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:112e305c3314f40c93998e567879e887a3160bb8689ef3d2c04b6cc62c33ac34", size = 1896206, upload-time = "2025-11-04T13:40:27.099Z" },
+    { url = "https://files.pythonhosted.org/packages/15/df/a4c740c0943e93e6500f9eb23f4ca7ec9bf71b19e608ae5b579678c8d02f/pydantic_core-2.41.5-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:0cbaad15cb0c90aa221d43c00e77bb33c93e8d36e0bf74760cd00e732d10a6a0", size = 1919307, upload-time = "2025-11-04T13:40:29.806Z" },
+    { url = "https://files.pythonhosted.org/packages/9a/e3/6324802931ae1d123528988e0e86587c2072ac2e5394b4bc2bc34b61ff6e/pydantic_core-2.41.5-cp313-cp313-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:03ca43e12fab6023fc79d28ca6b39b05f794ad08ec2feccc59a339b02f2b3d33", size = 2063258, upload-time = "2025-11-04T13:40:33.544Z" },
+    { url = "https://files.pythonhosted.org/packages/c9/d4/2230d7151d4957dd79c3044ea26346c148c98fbf0ee6ebd41056f2d62ab5/pydantic_core-2.41.5-cp313-cp313-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:dc799088c08fa04e43144b164feb0c13f9a0bc40503f8df3e9fde58a3c0c101e", size = 2214917, upload-time = "2025-11-04T13:40:35.479Z" },
+    { url = "https://files.pythonhosted.org/packages/e6/9f/eaac5df17a3672fef0081b6c1bb0b82b33ee89aa5cec0d7b05f52fd4a1fa/pydantic_core-2.41.5-cp313-cp313-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:97aeba56665b4c3235a0e52b2c2f5ae9cd071b8a8310ad27bddb3f7fb30e9aa2", size = 2332186, upload-time = "2025-11-04T13:40:37.436Z" },
+    { url = "https://files.pythonhosted.org/packages/cf/4e/35a80cae583a37cf15604b44240e45c05e04e86f9cfd766623149297e971/pydantic_core-2.41.5-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:406bf18d345822d6c21366031003612b9c77b3e29ffdb0f612367352aab7d586", size = 2073164, upload-time = "2025-11-04T13:40:40.289Z" },
+    { url = "https://files.pythonhosted.org/packages/bf/e3/f6e262673c6140dd3305d144d032f7bd5f7497d3871c1428521f19f9efa2/pydantic_core-2.41.5-cp313-cp313-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:b93590ae81f7010dbe380cdeab6f515902ebcbefe0b9327cc4804d74e93ae69d", size = 2179146, upload-time = "2025-11-04T13:40:42.809Z" },
+    { url = "https://files.pythonhosted.org/packages/75/c7/20bd7fc05f0c6ea2056a4565c6f36f8968c0924f19b7d97bbfea55780e73/pydantic_core-2.41.5-cp313-cp313-musllinux_1_1_aarch64.whl", hash = "sha256:01a3d0ab748ee531f4ea6c3e48ad9dac84ddba4b0d82291f87248f2f9de8d740", size = 2137788, upload-time = "2025-11-04T13:40:44.752Z" },
+    { url = "https://files.pythonhosted.org/packages/3a/8d/34318ef985c45196e004bc46c6eab2eda437e744c124ef0dbe1ff2c9d06b/pydantic_core-2.41.5-cp313-cp313-musllinux_1_1_armv7l.whl", hash = "sha256:6561e94ba9dacc9c61bce40e2d6bdc3bfaa0259d3ff36ace3b1e6901936d2e3e", size = 2340133, upload-time = "2025-11-04T13:40:46.66Z" },
+    { url = "https://files.pythonhosted.org/packages/9c/59/013626bf8c78a5a5d9350d12e7697d3d4de951a75565496abd40ccd46bee/pydantic_core-2.41.5-cp313-cp313-musllinux_1_1_x86_64.whl", hash = "sha256:915c3d10f81bec3a74fbd4faebe8391013ba61e5a1a8d48c4455b923bdda7858", size = 2324852, upload-time = "2025-11-04T13:40:48.575Z" },
+    { url = "https://files.pythonhosted.org/packages/1a/d9/c248c103856f807ef70c18a4f986693a46a8ffe1602e5d361485da502d20/pydantic_core-2.41.5-cp313-cp313-win32.whl", hash = "sha256:650ae77860b45cfa6e2cdafc42618ceafab3a2d9a3811fcfbd3bbf8ac3c40d36", size = 1994679, upload-time = "2025-11-04T13:40:50.619Z" },
+    { url = "https://files.pythonhosted.org/packages/9e/8b/341991b158ddab181cff136acd2552c9f35bd30380422a639c0671e99a91/pydantic_core-2.41.5-cp313-cp313-win_amd64.whl", hash = "sha256:79ec52ec461e99e13791ec6508c722742ad745571f234ea6255bed38c6480f11", size = 2019766, upload-time = "2025-11-04T13:40:52.631Z" },
+    { url = "https://files.pythonhosted.org/packages/73/7d/f2f9db34af103bea3e09735bb40b021788a5e834c81eedb541991badf8f5/pydantic_core-2.41.5-cp313-cp313-win_arm64.whl", hash = "sha256:3f84d5c1b4ab906093bdc1ff10484838aca54ef08de4afa9de0f5f14d69639cd", size = 1981005, upload-time = "2025-11-04T13:40:54.734Z" },
+    { url = "https://files.pythonhosted.org/packages/ea/28/46b7c5c9635ae96ea0fbb779e271a38129df2550f763937659ee6c5dbc65/pydantic_core-2.41.5-cp314-cp314-macosx_10_12_x86_64.whl", hash = "sha256:3f37a19d7ebcdd20b96485056ba9e8b304e27d9904d233d7b1015db320e51f0a", size = 2119622, upload-time = "2025-11-04T13:40:56.68Z" },
+    { url = "https://files.pythonhosted.org/packages/74/1a/145646e5687e8d9a1e8d09acb278c8535ebe9e972e1f162ed338a622f193/pydantic_core-2.41.5-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:1d1d9764366c73f996edd17abb6d9d7649a7eb690006ab6adbda117717099b14", size = 1891725, upload-time = "2025-11-04T13:40:58.807Z" },
+    { url = "https://files.pythonhosted.org/packages/23/04/e89c29e267b8060b40dca97bfc64a19b2a3cf99018167ea1677d96368273/pydantic_core-2.41.5-cp314-cp314-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:25e1c2af0fce638d5f1988b686f3b3ea8cd7de5f244ca147c777769e798a9cd1", size = 1915040, upload-time = "2025-11-04T13:41:00.853Z" },
+    { url = "https://files.pythonhosted.org/packages/84/a3/15a82ac7bd97992a82257f777b3583d3e84bdb06ba6858f745daa2ec8a85/pydantic_core-2.41.5-cp314-cp314-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:506d766a8727beef16b7adaeb8ee6217c64fc813646b424d0804d67c16eddb66", size = 2063691, upload-time = "2025-11-04T13:41:03.504Z" },
+    { url = "https://files.pythonhosted.org/packages/74/9b/0046701313c6ef08c0c1cf0e028c67c770a4e1275ca73131563c5f2a310a/pydantic_core-2.41.5-cp314-cp314-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:4819fa52133c9aa3c387b3328f25c1facc356491e6135b459f1de698ff64d869", size = 2213897, upload-time = "2025-11-04T13:41:05.804Z" },
+    { url = "https://files.pythonhosted.org/packages/8a/cd/6bac76ecd1b27e75a95ca3a9a559c643b3afcd2dd62086d4b7a32a18b169/pydantic_core-2.41.5-cp314-cp314-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:2b761d210c9ea91feda40d25b4efe82a1707da2ef62901466a42492c028553a2", size = 2333302, upload-time = "2025-11-04T13:41:07.809Z" },
+    { url = "https://files.pythonhosted.org/packages/4c/d2/ef2074dc020dd6e109611a8be4449b98cd25e1b9b8a303c2f0fca2f2bcf7/pydantic_core-2.41.5-cp314-cp314-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:22f0fb8c1c583a3b6f24df2470833b40207e907b90c928cc8d3594b76f874375", size = 2064877, upload-time = "2025-11-04T13:41:09.827Z" },
+    { url = "https://files.pythonhosted.org/packages/18/66/e9db17a9a763d72f03de903883c057b2592c09509ccfe468187f2a2eef29/pydantic_core-2.41.5-cp314-cp314-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:2782c870e99878c634505236d81e5443092fba820f0373997ff75f90f68cd553", size = 2180680, upload-time = "2025-11-04T13:41:12.379Z" },
+    { url = "https://files.pythonhosted.org/packages/d3/9e/3ce66cebb929f3ced22be85d4c2399b8e85b622db77dad36b73c5387f8f8/pydantic_core-2.41.5-cp314-cp314-musllinux_1_1_aarch64.whl", hash = "sha256:0177272f88ab8312479336e1d777f6b124537d47f2123f89cb37e0accea97f90", size = 2138960, upload-time = "2025-11-04T13:41:14.627Z" },
+    { url = "https://files.pythonhosted.org/packages/a6/62/205a998f4327d2079326b01abee48e502ea739d174f0a89295c481a2272e/pydantic_core-2.41.5-cp314-cp314-musllinux_1_1_armv7l.whl", hash = "sha256:63510af5e38f8955b8ee5687740d6ebf7c2a0886d15a6d65c32814613681bc07", size = 2339102, upload-time = "2025-11-04T13:41:16.868Z" },
+    { url = "https://files.pythonhosted.org/packages/3c/0d/f05e79471e889d74d3d88f5bd20d0ed189ad94c2423d81ff8d0000aab4ff/pydantic_core-2.41.5-cp314-cp314-musllinux_1_1_x86_64.whl", hash = "sha256:e56ba91f47764cc14f1daacd723e3e82d1a89d783f0f5afe9c364b8bb491ccdb", size = 2326039, upload-time = "2025-11-04T13:41:18.934Z" },
+    { url = "https://files.pythonhosted.org/packages/ec/e1/e08a6208bb100da7e0c4b288eed624a703f4d129bde2da475721a80cab32/pydantic_core-2.41.5-cp314-cp314-win32.whl", hash = "sha256:aec5cf2fd867b4ff45b9959f8b20ea3993fc93e63c7363fe6851424c8a7e7c23", size = 1995126, upload-time = "2025-11-04T13:41:21.418Z" },
+    { url = "https://files.pythonhosted.org/packages/48/5d/56ba7b24e9557f99c9237e29f5c09913c81eeb2f3217e40e922353668092/pydantic_core-2.41.5-cp314-cp314-win_amd64.whl", hash = "sha256:8e7c86f27c585ef37c35e56a96363ab8de4e549a95512445b85c96d3e2f7c1bf", size = 2015489, upload-time = "2025-11-04T13:41:24.076Z" },
+    { url = "https://files.pythonhosted.org/packages/4e/bb/f7a190991ec9e3e0ba22e4993d8755bbc4a32925c0b5b42775c03e8148f9/pydantic_core-2.41.5-cp314-cp314-win_arm64.whl", hash = "sha256:e672ba74fbc2dc8eea59fb6d4aed6845e6905fc2a8afe93175d94a83ba2a01a0", size = 1977288, upload-time = "2025-11-04T13:41:26.33Z" },
+    { url = "https://files.pythonhosted.org/packages/92/ed/77542d0c51538e32e15afe7899d79efce4b81eee631d99850edc2f5e9349/pydantic_core-2.41.5-cp314-cp314t-macosx_10_12_x86_64.whl", hash = "sha256:8566def80554c3faa0e65ac30ab0932b9e3a5cd7f8323764303d468e5c37595a", size = 2120255, upload-time = "2025-11-04T13:41:28.569Z" },
+    { url = "https://files.pythonhosted.org/packages/bb/3d/6913dde84d5be21e284439676168b28d8bbba5600d838b9dca99de0fad71/pydantic_core-2.41.5-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:b80aa5095cd3109962a298ce14110ae16b8c1aece8b72f9dafe81cf597ad80b3", size = 1863760, upload-time = "2025-11-04T13:41:31.055Z" },
+    { url = "https://files.pythonhosted.org/packages/5a/f0/e5e6b99d4191da102f2b0eb9687aaa7f5bea5d9964071a84effc3e40f997/pydantic_core-2.41.5-cp314-cp314t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:3006c3dd9ba34b0c094c544c6006cc79e87d8612999f1a5d43b769b89181f23c", size = 1878092, upload-time = "2025-11-04T13:41:33.21Z" },
+    { url = "https://files.pythonhosted.org/packages/71/48/36fb760642d568925953bcc8116455513d6e34c4beaa37544118c36aba6d/pydantic_core-2.41.5-cp314-cp314t-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:72f6c8b11857a856bcfa48c86f5368439f74453563f951e473514579d44aa612", size = 2053385, upload-time = "2025-11-04T13:41:35.508Z" },
+    { url = "https://files.pythonhosted.org/packages/20/25/92dc684dd8eb75a234bc1c764b4210cf2646479d54b47bf46061657292a8/pydantic_core-2.41.5-cp314-cp314t-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:5cb1b2f9742240e4bb26b652a5aeb840aa4b417c7748b6f8387927bc6e45e40d", size = 2218832, upload-time = "2025-11-04T13:41:37.732Z" },
+    { url = "https://files.pythonhosted.org/packages/e2/09/f53e0b05023d3e30357d82eb35835d0f6340ca344720a4599cd663dca599/pydantic_core-2.41.5-cp314-cp314t-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:bd3d54f38609ff308209bd43acea66061494157703364ae40c951f83ba99a1a9", size = 2327585, upload-time = "2025-11-04T13:41:40Z" },
+    { url = "https://files.pythonhosted.org/packages/aa/4e/2ae1aa85d6af35a39b236b1b1641de73f5a6ac4d5a7509f77b814885760c/pydantic_core-2.41.5-cp314-cp314t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:2ff4321e56e879ee8d2a879501c8e469414d948f4aba74a2d4593184eb326660", size = 2041078, upload-time = "2025-11-04T13:41:42.323Z" },
+    { url = "https://files.pythonhosted.org/packages/cd/13/2e215f17f0ef326fc72afe94776edb77525142c693767fc347ed6288728d/pydantic_core-2.41.5-cp314-cp314t-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:d0d2568a8c11bf8225044aa94409e21da0cb09dcdafe9ecd10250b2baad531a9", size = 2173914, upload-time = "2025-11-04T13:41:45.221Z" },
+    { url = "https://files.pythonhosted.org/packages/02/7a/f999a6dcbcd0e5660bc348a3991c8915ce6599f4f2c6ac22f01d7a10816c/pydantic_core-2.41.5-cp314-cp314t-musllinux_1_1_aarch64.whl", hash = "sha256:a39455728aabd58ceabb03c90e12f71fd30fa69615760a075b9fec596456ccc3", size = 2129560, upload-time = "2025-11-04T13:41:47.474Z" },
+    { url = "https://files.pythonhosted.org/packages/3a/b1/6c990ac65e3b4c079a4fb9f5b05f5b013afa0f4ed6780a3dd236d2cbdc64/pydantic_core-2.41.5-cp314-cp314t-musllinux_1_1_armv7l.whl", hash = "sha256:239edca560d05757817c13dc17c50766136d21f7cd0fac50295499ae24f90fdf", size = 2329244, upload-time = "2025-11-04T13:41:49.992Z" },
+    { url = "https://files.pythonhosted.org/packages/d9/02/3c562f3a51afd4d88fff8dffb1771b30cfdfd79befd9883ee094f5b6c0d8/pydantic_core-2.41.5-cp314-cp314t-musllinux_1_1_x86_64.whl", hash = "sha256:2a5e06546e19f24c6a96a129142a75cee553cc018ffee48a460059b1185f4470", size = 2331955, upload-time = "2025-11-04T13:41:54.079Z" },
+    { url = "https://files.pythonhosted.org/packages/5c/96/5fb7d8c3c17bc8c62fdb031c47d77a1af698f1d7a406b0f79aaa1338f9ad/pydantic_core-2.41.5-cp314-cp314t-win32.whl", hash = "sha256:b4ececa40ac28afa90871c2cc2b9ffd2ff0bf749380fbdf57d165fd23da353aa", size = 1988906, upload-time = "2025-11-04T13:41:56.606Z" },
+    { url = "https://files.pythonhosted.org/packages/22/ed/182129d83032702912c2e2d8bbe33c036f342cc735737064668585dac28f/pydantic_core-2.41.5-cp314-cp314t-win_amd64.whl", hash = "sha256:80aa89cad80b32a912a65332f64a4450ed00966111b6615ca6816153d3585a8c", size = 1981607, upload-time = "2025-11-04T13:41:58.889Z" },
+    { url = "https://files.pythonhosted.org/packages/9f/ed/068e41660b832bb0b1aa5b58011dea2a3fe0ba7861ff38c4d4904c1c1a99/pydantic_core-2.41.5-cp314-cp314t-win_arm64.whl", hash = "sha256:35b44f37a3199f771c3eaa53051bc8a70cd7b54f333531c59e29fd4db5d15008", size = 1974769, upload-time = "2025-11-04T13:42:01.186Z" },
+]
+
+[[package]]
+name = "typing-extensions"
+version = "4.15.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/72/94/1a15dd82efb362ac84269196e94cf00f187f7ed21c242792a923cdb1c61f/typing_extensions-4.15.0.tar.gz", hash = "sha256:0cea48d173cc12fa28ecabc3b837ea3cf6f38c6d1136f85cbaaf598984861466", size = 109391, upload-time = "2025-08-25T13:49:26.313Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/18/67/36e9267722cc04a6b9f15c7f3441c2363321a3ea07da7ae0c0707beb2a9c/typing_extensions-4.15.0-py3-none-any.whl", hash = "sha256:f0fa19c6845758ab08074a0cfa8b7aecb71c999ca73d62883bc25cc018c4e548", size = 44614, upload-time = "2025-08-25T13:49:24.86Z" },
+]
+
+[[package]]
+name = "typing-inspection"
+version = "0.4.2"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "typing-extensions" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/55/e3/70399cb7dd41c10ac53367ae42139cf4b1ca5f36bb3dc6c9d33acdb43655/typing_inspection-0.4.2.tar.gz", hash = "sha256:ba561c48a67c5958007083d386c3295464928b01faa735ab8547c5692e87f464", size = 75949, upload-time = "2025-10-01T02:14:41.687Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/dc/9b/47798a6c91d8bdb567fe2698fe81e0c6b7cb7ef4d13da4114b41d239f65d/typing_inspection-0.4.2-py3-none-any.whl", hash = "sha256:4ed1cacbdc298c220f1bd249ed5287caa16f34d44ef4e9c3d0cbad5b521545e7", size = 14611, upload-time = "2025-10-01T02:14:40.154Z" },
+]

python/nanobind_cmake.patch

@@ -1,25 +0,0 @@
-diff --git a/cpp/nanobind/CMakeLists.txt b/cpp/nanobind/CMakeLists.txt
-index ebfd3da..6ef06e3 100644
---- a/cpp/nanobind/CMakeLists.txt
-+++ b/cpp/nanobind/CMakeLists.txt
-@@ -18,8 +18,18 @@ target_sources(python_methods PRIVATE python_methods.cc)
- target_link_libraries(python_methods PUBLIC xgrammar)
- 
- # Any code that uses nanobind directly lives here
--nanobind_add_module(xgrammar_bindings LTO nanobind.cc)
--target_link_libraries(xgrammar_bindings PRIVATE python_methods)
-+nanobind_build_library(nanobind)
-+add_library(xgrammar_bindings MODULE nanobind.cc)
-+target_link_libraries(xgrammar_bindings PRIVATE
-+  python_methods
-+  nanobind
-+)
-+nanobind_opt_size(xgrammar_bindings)
-+nanobind_lto(xgrammar_bindings)
-+nanobind_set_visibility(xgrammar_bindings)
-+nanobind_extension(xgrammar_bindings)
-+nanobind_compile_options(xgrammar_bindings)
-+nanobind_link_options(xgrammar_bindings)
- 
- if(DEFINED SKBUILD_PROJECT_NAME)
-   # Building wheel through scikit-build-core

python/parts.nix

@@ -1,11 +1,11 @@
 { inputs, ... }:
 {
   perSystem =
-    { config, self', pkgs, lib, system, ... }:
+    { self', pkgs, cudaPkgs, lib, ... }:
     let
       # Load workspace from uv.lock
       workspace = inputs.uv2nix.lib.workspace.loadWorkspace {
-        workspaceRoot = inputs.self;
+        workspaceRoot = ../.;
       };
 
       # Create overlay from workspace
@@ -35,16 +35,9 @@
 
       inherit (pkgs.stdenv.hostPlatform) isDarwin isLinux;
 
-      python = pkgs.python313;
-
       # Overlay to provide build systems and custom packages
-      buildSystemsOverlay = final: prev:
+      cudaOverlay = final: prev:
         let
-          addSetupTools = (old: {
-            nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [
-              final.setuptools
-            ];
-          });
           torchLibs = [
             final.nvidia-cuda-runtime
             final.nvidia-cuda-nvrtc
@@ -147,7 +140,7 @@
             '';
           };
 
-          mergedCudaLibraries = with pkgs.cudaPackages_13; [
+          mergedCudaLibraries = with cudaPkgs.cudaPackages_13; [
             cuda_cudart # cuda_runtime.h, -lcudart
             cuda_cccl
             libcurand # curand_kernel.h
@@ -159,25 +152,16 @@
             libcublas
           ];
 
-	  cuda_cccl_compat = pkgs.runCommand "cuda-cccl-compat" {} ''
-	    mkdir -p $out/include
-	    ln -s ${pkgs.cudaPackages_13.cuda_cccl}/include $out/include/cccl
-	  '';
+          cuda_cccl_compat = pkgs.runCommand "cuda-cccl-compat" { } ''
+            mkdir -p $out/include
+            ln -s ${cudaPkgs.cudaPackages_13.cuda_cccl}/include $out/include/cccl
+          '';
           cudaToolkitRoot = pkgs.symlinkJoin {
-              name = "cuda-merged-exo";
-              paths = builtins.concatMap (p: [ (lib.getBin p) (lib.getLib p) (lib.getDev p) ]) (mergedCudaLibraries ++ [ pkgs.cudaPackages_13.cuda_nvcc  cuda_cccl_compat ]);
-            };
-
+            name = "cuda-merged-exo";
+            paths = builtins.concatMap (p: [ (lib.getBin p) (lib.getLib p) (lib.getDev p) ]) (mergedCudaLibraries ++ [ cudaPkgs.cudaPackages_13.cuda_nvcc cuda_cccl_compat ]);
+          };
         in
-
         {
-          # mlx-lm is a git dependency that needs setuptools
-          mlx-lm = prev.mlx-lm.overrideAttrs addSetupTools;
-          # rouge-score and sacrebleu don't declare setuptools as a build dependency
-          rouge-score = prev.rouge-score.overrideAttrs addSetupTools;
-          sacrebleu = prev.sacrebleu.overrideAttrs addSetupTools;
-          sqlitedict = prev.sqlitedict.overrideAttrs addSetupTools;
-          word2number = prev.word2number.overrideAttrs addSetupTools;
           fastsafetensors = prev.fastsafetensors.overrideAttrs (old: { nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [ final.setuptools final.pybind11 ]; });
           torch = prev.torch.overrideAttrs (old: {
             propagatedBuildInputs = (old.propagatedBuildInputs or [ ]) ++ torchLibs ++ [ final.typing-extensions final.numpy ];
@@ -203,16 +187,17 @@
               '';
               autoPatchelfIgnoreMissingDeps = (old.autoPatchelfIgnoreMissingDeps or [ ]) ++ [ "libcuda.so.1" ];
             });
-          xgrammar = prev.xgrammar.overrideAttrs (old: { nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [ final.setuptools final.scikit-build-core final.packaging final.pathspec pkgs.cmake final.nanobind ]; 
+          xgrammar = prev.xgrammar.overrideAttrs (old: {
+            nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [ final.setuptools final.scikit-build-core final.packaging final.pathspec pkgs.cmake final.nanobind ];
 
-  prePatch = ''
-cat cpp/nanobind/CMakeLists.txt
-'';
-  patches = (old.patches or [ ]) ++ [ ./nanobind_cmake.patch ];
-});
+            prePatch = ''
+              cat cpp/nanobind/CMakeLists.txt
+            '';
+            patches = (old.patches or [ ]) ++ [ ../nix/nanobind_cmake.patch ];
+          });
           vllm = prev.vllm.overrideAttrs (old: {
-            patches = (old.patches or [ ]) ++ [ ./vllm_uv2nix_cmake.patch ];
-            nativeBuildInputs = with pkgs.cudaPackages_13; (old.nativeBuildInputs or [ ]) ++ [
+            patches = (old.patches or [ ]) ++ [ ../nix/vllm_uv2nix_cmake.patch ];
+            nativeBuildInputs = with cudaPkgs.cudaPackages_13; (old.nativeBuildInputs or [ ]) ++ [
               final.setuptools
               final.setuptools-scm
               final.scikit-build-core
@@ -224,7 +209,7 @@ cat cpp/nanobind/CMakeLists.txt
               pkgs.ninja
               pkgs.autoAddDriverRunpath
             ];
-            buildInputs = with pkgs.cudaPackages_13; [
+            buildInputs = with cudaPkgs.cudaPackages_13; [
               libcufile
               cudnn
               nccl
@@ -250,7 +235,7 @@ cat cpp/nanobind/CMakeLists.txt
               "-DVLLM_FLASH_ATTN_SRC_DIR=${lib.getDev vllm-flash-attn}"
               "-DQUTLASS_SRC_DIR=${lib.getDev qutlass}"
               "-DTORCH_CUDA_ARCH_LIST=12.0;12.1"
-              "-DCUTLASS_NVCC_ARCHS_ENABLED=${pkgs.cudaPackages_13.flags.cmakeCudaArchitecturesString}"
+              "-DCUTLASS_NVCC_ARCHS_ENABLED=${cudaPkgs.cudaPackages_13.flags.cmakeCudaArchitecturesString}"
               "-DCUDA_HOME=${cudaToolkitRoot}"
               "-DCAFFE2_USE_CUDNN=ON"
               "-DCAFFE2_USE_CUFILE=ON"
@@ -258,6 +243,33 @@ cat cpp/nanobind/CMakeLists.txt
             ];
 
 
+          });
+        };
+
+      buildSystemsOverlay = final: prev:
+        let
+          addSetupTools = old: {
+            nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [
+              final.setuptools
+            ];
+          };
+
+        in
+        {
+          # mlx-lm is a git dependency that needs setuptools
+          mlx-lm = prev.mlx-lm.overrideAttrs addSetupTools;
+          # rouge-score and sacrebleu don't declare setuptools as a build dependency
+          rouge-score = prev.rouge-score.overrideAttrs addSetupTools;
+          sacrebleu = prev.sacrebleu.overrideAttrs addSetupTools;
+          sqlitedict = prev.sqlitedict.overrideAttrs addSetupTools;
+          word2number = prev.word2number.overrideAttrs addSetupTools;
+          vllm = prev.vllm.overrideAttrs (old: {
+            nativeBuildInputs = (old.nativeBuildInputs or [ ]) ++ [
+              final.setuptools
+              final.setuptools-scm
+            ];
+            propagatedBuildInputs = (old.propagatedBuildInputs or [ ]) ++ [ final.torch ];
+            VLLM_TARGET_DEVICE = "empty";
           });
         } // lib.optionalAttrs isDarwin {
           # Use our pure Nix-built MLX with Metal support (macOS only)
@@ -295,6 +307,7 @@ cat cpp/nanobind/CMakeLists.txt
         );
 
 
+      python = pkgs.python313;
 
       pythonSet = (pkgs.callPackage inputs.pyproject-nix.build.packages {
         inherit python;
@@ -307,6 +320,31 @@ cat cpp/nanobind/CMakeLists.txt
           linuxOverlay
         ]
       );
+      editablePythonSet = pythonSet.overrideScope (
+        workspace.mkEditablePyprojectOverlay { root = "$REPO_ROOT"; members = [ "exo" "exo_core" "vllm_engine" "mlx_engine" "exo_bench" ]; }
+      );
+      evenv = editablePythonSet.mkVirtualEnv "exo-dev-env"
+        {
+          exo = lib.optionals isDarwin [ "mlx" ];
+          exo-pyo3-bindings = [ ];
+          exo-bench = [ ];
+          mlx-engine = [ ];
+          vllm-engine = [ ];
+
+        }
+      ;
+      cudaPythonSet = (cudaPkgs.callPackage inputs.pyproject-nix.build.packages {
+        python = cudaPkgs.python313;
+      }).overrideScope (
+        lib.composeManyExtensions [
+          inputs.pyproject-build-systems.overlays.default
+          overlay
+          exoOverlay
+          buildSystemsOverlay
+          linuxOverlay
+          cudaOverlay
+        ]
+      );
       # mlx-cpu and mlx-cuda-13 both ship mlx/ site-packages files; keep first.
       # mlx-cpu/mlx-cuda-13 and nvidia-cudnn-cu12/cu13 ship overlapping files.
       venvCollisionPaths = lib.optionals isLinux [
@@ -320,8 +358,8 @@ cat cpp/nanobind/CMakeLists.txt
       }).overrideAttrs {
         venvIgnoreCollisions = venvCollisionPaths;
       };
-      exoCudaVenv = (pythonSet.mkVirtualEnv "exo-env" {
-        exo = [ "cuda" ];
+      exoCudaVenv = (cudaPythonSet.mkVirtualEnv "exo-env" {
+        exo = lib.optionals cudaPkgs.config.cudaSupport [ "cuda" ];
         exo-pyo3-bindings = [ ];
       }).overrideAttrs {
         venvIgnoreCollisions = venvCollisionPaths;
@@ -379,7 +417,7 @@ cat cpp/nanobind/CMakeLists.txt
             ${lib.optionalString isDarwin "--prefix PATH : ${pkgs.macmon}/bin"}
         '';
 
-      exoCudaPackage = pkgs.runCommand "exo"
+      exoCudaPackage = cudaPkgs.runCommand "exo"
         {
           nativeBuildInputs = [ pkgs.makeWrapper ];
         }
@@ -406,6 +444,8 @@ cat cpp/nanobind/CMakeLists.txt
         exo-eval = mkBenchScript "exo-eval" (inputs.self + /bench/exo_eval.py);
         exo-eval-tool-calls = mkBenchScript "exo-eval-tool-calls" (inputs.self + /bench/eval_tool_calls.py);
         exo-get-all-models-on-cluster = mkSimplePythonScript "exo-get-all-models-on-cluster" (inputs.self + /tests/get_all_models_on_cluster.py);
+        editable-venv = evenv;
+        inherit python;
       };
 
       checks = {

python/vllm_engine/main.py

@@ -1,6 +0,0 @@
-def main():
-    print("Hello from vllm-runner!")
-
-
-if __name__ == "__main__":
-    main()

python/vllm_engine/pyproject.toml

@@ -6,16 +6,18 @@ readme = "README.md"
 requires-python = ">=3.13"
 dependencies = [
   "torch>=2.10.0; sys_platform == 'linux'",
-  "vllm>=0.13.0; sys_platform == 'linux' and platform_machine == 'aarch64'",
+  "vllm>=0.13.0; sys_platform == 'linux'",
   "mlx-cuda-13==0.30.6; sys_platform == 'linux'",
   "fastsafetensors>=0.1.10; sys_platform == 'linux'",
+  "exo_core",
 ]
 
 [tool.uv]
-environments = ["sys_platform == 'linux' and platform_machine == 'aarch64'"]
+environments = ["sys_platform == 'linux'"]
 
 
 [tool.uv.sources]
+exo_core = { workspace = true }
 vllm = { git = "https://github.com/hmellor/vllm.git", rev = "b99bedc737166ae5ca98cb9e3534b96e0c8c69aa" }
 torch = [
   { index = "pytorch-cu130", marker = "platform_machine == 'aarch64'" },
@@ -35,3 +37,18 @@ explicit = true
 [build-system]
 requires = ["uv_build>=0.8.9,<0.9.0"]
 build-backend = "uv_build"
+
+[tool.basedpyright]
+typeCheckingMode = "strict"
+failOnWarnings = true
+
+reportAny = "error"
+reportUnknownVariableType = "error"
+reportUnknownParameterType = "error"
+reportMissingParameterType = "error"
+reportMissingTypeStubs = "error"
+reportInvalidCast = "error"
+reportUnnecessaryCast = "error"
+reportUnnecessaryTypeIgnoreComment = "error"
+
+pythonVersion = "3.13"

python/vllm_engine/src/vllm_engine/prompt_format.py

@@ -1,5 +1,5 @@
-from exo.shared.types.common import ModelId
-from exo.shared.types.text_generation import TextGenerationTaskParams
+from exo_core.types.common import ModelId
+from exo_core.types.text_generation import TextGenerationTaskParams
 from mlx_lm.tokenizer_utils import TokenizerWrapper
 from vllm.sampling_params import SamplingParams
 from vllm.v1.engine.llm_engine import LLMEngine

python/vllm_engine/src/vllm_engine/vllm_generator.py

@@ -8,36 +8,36 @@ from collections.abc import Callable, Generator
 from dataclasses import dataclass, field
 
 import torch
-from exo.shared.types.api import (
-    CompletionTokensDetails,
-    GenerationStats,
-    PromptTokensDetails,
-    Usage,
-)
-from exo.shared.types.common import ModelId
-from exo.shared.types.memory import Memory
-from exo.shared.types.tasks import TaskId
-from exo.shared.types.text_generation import TextGenerationTaskParams
-from exo.shared.types.worker.runner_response import GenerationResponse
-from exo.worker.engines.vllm.growable_cache import (
-    get_model_runner,
-    patch_vllm,
-    set_prefix_cache,
-)
-from exo.worker.engines.vllm.kv_cache import TorchKVCache
-from exo.worker.engines.vllm.prompt_format import (
-    format_vllm_prompt,
-    make_vllm_sampling_params,
-)
+from exo_core.types.common import ModelId
+from exo_core.types.runner_response import GenerationResponse
+from exo_core.types.tasks import TaskId
+from exo_core.types.text_generation import TextGenerationTaskParams
+from exo_core.utils.memory import Memory
 from vllm.engine.arg_utils import EngineArgs
 from vllm.sampling_params import SamplingParams
 from vllm.v1.engine.llm_engine import LLMEngine
 from vllm.v1.kv_cache_interface import KVCacheConfig
 
+from exo.api.types import (
+    CompletionTokensDetails,
+    GenerationStats,
+    PromptTokensDetails,
+    Usage,
+)
 from exo.worker.engines.mlx.cache import KVPrefixCache
 from exo.worker.engines.mlx.utils_mlx import get_eos_token_ids_for_model
 from exo.worker.runner.bootstrap import logger
 from exo.worker.runner.llm_inference.tool_parsers import ToolParser, infer_tool_parser
+from vllm_engine.growable_cache import (
+    get_model_runner,
+    patch_vllm,
+    set_prefix_cache,
+)
+from vllm_engine.kv_cache import TorchKVCache
+from vllm_engine.prompt_format import (
+    format_vllm_prompt,
+    make_vllm_sampling_params,
+)
 
 
 def _build_layer_groups(kv_cache_config: KVCacheConfig) -> list[int]:

python/vllm_uv2nix_cmake.patch

@@ -1,12 +0,0 @@
-diff --git a/setup.py b/setup.py
-index 68861fe4b..fd4738089 100644
---- a/setup.py
-+++ b/setup.py
-@@ -150,6 +150,7 @@ class cmake_build_ext(build_ext):
-         cmake_args = [
-             "-DCMAKE_BUILD_TYPE={}".format(cfg),
-             "-DVLLM_TARGET_DEVICE={}".format(VLLM_TARGET_DEVICE),
-+            *json.loads(os.environ.get("UV2NIX_CMAKE_FLAGS_JSON", "[]"))
-         ]
- 
-         verbose = envs.VERBOSE