gotta move the cache soon

2026-03-21 22:22:25 +00:00
parent 532a8f0b07
commit 8cd6191c70
124 changed files with 1550 additions and 1401 deletions
@@ -2,11 +2,11 @@
  "nodes": {
    "crane": {
      "locked": {
-        "lastModified": 1767744144,
-        "narHash": "sha256-9/9ntI0D+HbN4G0TrK3KmHbTvwgswz7p8IEJsWyef8Q=",
+        "lastModified": 1773857772,
+        "narHash": "sha256-5xsK26KRHf0WytBtsBnQYC/lTWDhQuT57HJ7SzuqZcM=",
        "owner": "ipetkov",
        "repo": "crane",
-        "rev": "2fb033290bf6b23f226d4c8b32f7f7a16b043d7e",
+        "rev": "b556d7bbae5ff86e378451511873dfd07e4504cd",
        "type": "github"
      },
      "original": {
@@ -47,11 +47,11 @@
        "rust-analyzer-src": "rust-analyzer-src"
      },
      "locked": {
-        "lastModified": 1768287139,
-        "narHash": "sha256-nsXFt0OzUi6K7dUzzJD5/v9e0Ic+fvclfIW936/43ZM=",
+        "lastModified": 1773990552,
+        "narHash": "sha256-5dJwbgX05hVWZj8Kf9Tmb3mYWUZzF0kn6GMYBJQtPDI=",
        "owner": "nix-community",
        "repo": "fenix",
-        "rev": "a4a3aa956931f90f35453cb519e4545e9ad7f773",
+        "rev": "a90c1a2e004e318a4463c9e87ae1f00788afba1b",
        "type": "github"
      },
      "original": {
@@ -83,11 +83,11 @@
        ]
      },
      "locked": {
-        "lastModified": 1768135262,
-        "narHash": "sha256-PVvu7OqHBGWN16zSi6tEmPwwHQ4rLPU9Plvs8/1TUBY=",
+        "lastModified": 1772408722,
+        "narHash": "sha256-rHuJtdcOjK7rAHpHphUb1iCvgkU3GpfvicLMwwnfMT0=",
        "owner": "hercules-ci",
        "repo": "flake-parts",
-        "rev": "80daad04eddbbf5a4d883996a73f3f542fa437ac",
+        "rev": "f20dc5d9b8027381c474144ecabc9034d6a839a3",
        "type": "github"
      },
      "original": {
@@ -164,11 +164,11 @@
        ]
      },
      "locked": {
-        "lastModified": 1763662255,
-        "narHash": "sha256-4bocaOyLa3AfiS8KrWjZQYu+IAta05u3gYZzZ6zXbT0=",
+        "lastModified": 1773870109,
+        "narHash": "sha256-ZoTdqZP03DcdoyxvpFHCAek4bkPUTUPUF3oCCgc3dP4=",
        "owner": "pyproject-nix",
        "repo": "build-system-pkgs",
-        "rev": "042904167604c681a090c07eb6967b4dd4dae88c",
+        "rev": "b6e74f433b02fa4b8a7965ee24680f4867e2926f",
        "type": "github"
      },
      "original": {
@@ -184,11 +184,11 @@
        ]
      },
      "locked": {
-        "lastModified": 1764134915,
-        "narHash": "sha256-xaKvtPx6YAnA3HQVp5LwyYG1MaN4LLehpQI8xEdBvBY=",
+        "lastModified": 1773909723,
+        "narHash": "sha256-HmcZQ/hMPHR22Ri/6Sl7Z0B5J8nZa9bRnZJtDFInM7I=",
        "owner": "pyproject-nix",
        "repo": "pyproject.nix",
-        "rev": "2c8df1383b32e5443c921f61224b198a2282a657",
+        "rev": "d37dcf34ac7194eac4b0d10520d01298c434267d",
        "type": "github"
      },
      "original": {
@@ -214,11 +214,11 @@
    "rust-analyzer-src": {
      "flake": false,
      "locked": {
-        "lastModified": 1768224240,
-        "narHash": "sha256-Pp1dDrXKPBUJReZnnDElFyHYn67XTd48zRhToheLjtk=",
+        "lastModified": 1773946486,
+        "narHash": "sha256-xZHZj3Bk0cf17U9kgP9M/CFnAEoaZeOoGnQggfJSeGQ=",
        "owner": "rust-lang",
        "repo": "rust-analyzer",
-        "rev": "725349602e525df37f377701e001fe8aab807878",
+        "rev": "304785c5bf672fe01adc4201bd01bf61c33d567f",
        "type": "github"
      },
      "original": {
@@ -257,11 +257,11 @@
        ]
      },
      "locked": {
-        "lastModified": 1768158989,
-        "narHash": "sha256-67vyT1+xClLldnumAzCTBvU0jLZ1YBcf4vANRWP3+Ak=",
+        "lastModified": 1773297127,
+        "narHash": "sha256-6E/yhXP7Oy/NbXtf1ktzmU8SdVqJQ09HC/48ebEGBpk=",
        "owner": "numtide",
        "repo": "treefmt-nix",
-        "rev": "e96d59dff5c0d7fddb9d113ba108f03c3ef99eca",
+        "rev": "71b125cd05fbfd78cab3e070b73544abe24c5016",
        "type": "github"
      },
      "original": {
@@ -280,11 +280,11 @@
        ]
      },
      "locked": {
-        "lastModified": 1767701098,
-        "narHash": "sha256-CJhKZnWb3gumR9oTRjFvCg/6lYTGbZRU7xtvcyWIRwU=",
+        "lastModified": 1773958975,
+        "narHash": "sha256-Lr2k67KFPxPLqMtCWvwfg30S8huAEpEY9UcXHfi1Q+4=",
        "owner": "pyproject-nix",
        "repo": "uv2nix",
-        "rev": "9d357f0d2ce6f5f35ec7959d7e704452352eb4da",
+        "rev": "ffd52b90f29babbc4f309c29f2a2cdd6547be443",
        "type": "github"
      },
      "original": {
@@ -72,18 +72,23 @@
      ];

      perSystem =
-        { config, self', inputs', pkgs, lib, system, ... }:
+        { config, self', pkgs, lib, system, ... }:
        let
          # Use pinned nixpkgs for swift-format (swift is broken on x86_64-linux in newer nixpkgs)
          pkgsSwift = import inputs.nixpkgs-swift { inherit system; };
        in
        {
+          _module.args.cudaPkgs = import inputs.nixpkgs
+            {
+              inherit system;
+              config.allowUnfreePredicate = pkg: builtins.elem (lib.getName pkg) [ "cuda-merged" "cuda_cuobjdump" "cuda_gdb" "cuda_nvcc" "cuda_nvdisasm" "cuda_nvprune" "cuda_cccl" "cuda_cudart" "cuda_cupti" "cuda_cuxxfilt" "cuda_nvml_dev" "cuda_nvrtc" "cuda_nvtx" "cuda_profiler_api" "cuda_sanitizer_api" "libcublas" "libcufft" "libcurand" "libcusolver" "libnvjitlink" "libcusparse" "libnpp" "cudnn" "libcusparse_lt" "libcufile" "libnvshmem" "libnvvm" "cuda_crt" ];
+              cudaSupport = true;
+              cudaCapabilities = [ "12.1" ];
+            };
          # Allow unfree for metal-toolchain (needed for Darwin Metal packages)
          _module.args.pkgs = import inputs.nixpkgs {
            inherit system;
-            config.allowUnfreePredicate = pkg: builtins.elem (lib.getName pkg) [ "metal-toolchain" "cuda-merged" "cuda_cuobjdump" "cuda_gdb" "cuda_nvcc" "cuda_nvdisasm" "cuda_nvprune" "cuda_cccl" "cuda_cudart" "cuda_cupti" "cuda_cuxxfilt" "cuda_nvml_dev" "cuda_nvrtc" "cuda_nvtx" "cuda_profiler_api" "cuda_sanitizer_api" "libcublas" "libcufft" "libcurand" "libcusolver" "libnvjitlink" "libcusparse" "libnpp" "cudnn" "libcusparse_lt" "libcufile" "libnvshmem" "libnvvm" "cuda_crt" ];
-            cudaSupport = true;
-            cudaCapabilities = [ "12.1" ];
+            config.allowUnfreePredicate = pkg: builtins.elem (lib.getName pkg) [ "metal-toolchain" ];
            overlays = lib.optionals (system == "aarch64-darwin") [
              (import ./nix/apple-sdk-overlay.nix)
            ];
@@ -133,7 +138,7 @@
            );
          devShells =
            {
-              default = with pkgs; pkgs.mkShell {
+              default = with pkgs; mkShell {
                inputsFrom = [ self'.checks.cargo-build ];

                packages =
@@ -142,7 +147,8 @@
                    config.treefmt.build.wrapper

                    # PYTHON
-                    python313
+                    self'.packages.editable-venv
+                    self'.packages.python
                    uv
                    ruff
                    basedpyright
@@ -168,10 +174,17 @@
                    macmon
                  ];

-                OPENSSL_NO_VENDOR = "1";
+                env = {
+                  UV_NO_SYNC = "1";
+                  UV_PYTHON = self'.packages.python.interpreter;
+                  UV_PYTHON_DOWNLOADS = "never";
+                  OPENSSL_NO_VENDOR = "1";
+                };

                shellHook = ''
-                  export LD_LIBRARY_PATH="$LD_LIBRARY_PATH:${python313}/lib"
+                  unset PYTHONPATH
+                  export REPO_ROOT=$(git rev-parse --show-toplevel)
+                  export LD_LIBRARY_PATH="$LD_LIBRARY_PATH:${self'.packages.python}/lib"
                  ${lib.optionalString stdenv.isLinux ''
                    export LD_LIBRARY_PATH="${openssl.out}/lib:$LD_LIBRARY_PATH"
                  ''}
@@ -30,6 +30,7 @@ dependencies = [
  "python-multipart>=0.0.21",
  "msgspec>=0.19.0",
  "zstandard>=0.23.0",
+  "exo_core",
 ]

 [project.scripts]
@@ -45,28 +46,30 @@ dev = [
  "pytest-env",
  "ruff>=0.11.13",
 ]
-build = [
-  "nanobind"
-]
+build = ["nanobind"]

 [project.optional-dependencies]
 mlx = ["mlx_engine"]
-cuda = [
-  "vllm_engine; sys_platform == 'linux' and platform_machine == 'aarch64'",
-]
+cuda = ["vllm_engine; sys_platform == 'linux'"]

 ###
 # workspace configuration
 ###

 [tool.uv.workspace]
-members = ["rust/exo_pyo3_bindings", "bench", "python/exo_core"]
+members = [
+  "rust/exo_pyo3_bindings",
+  "bench",
+  "python/exo_core",
+  "python/vllm_engine",
+  "python/mlx_engine",
+]

 [tool.uv.sources]
 exo_pyo3_bindings = { workspace = true }
 exo_core = { workspace = true }
-mlx_engine = { path = "python/mlx_engine" }
-vllm_engine = { path = "python/vllm_engine" }
+mlx_engine = { workspace = true }
+vllm_engine = { workspace = true }
 mlx = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git", branch = "address-rdma-gpu-locks", marker = "sys_platform == 'darwin'" }
 mlx-lm = { git = "https://github.com/rltakashige/mlx-lm", branch = "leo/eval-left-padding-in-batched-rotation" }
 # Uncomment to use local mlx/mlx-lm development versions:
@@ -132,7 +135,7 @@ environments = ["sys_platform == 'darwin'", "sys_platform == 'linux'"]
 conflicts = [[{ package = "exo", extra = "cuda" }, { package = "exo-bench" }]]

 [tool.uv.extra-build-dependencies]
-xgrammar = [ "nanobind" ]
+xgrammar = ["nanobind"]


 ###
@@ -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"
@@ -2,7 +2,7 @@ import os
 import sys
 from pathlib import Path

-from exo.utils.dashboard_path import find_dashboard, find_resources
+from exo_core.utils.paths import find_dashboard, find_resources

 _EXO_HOME_ENV = os.environ.get("EXO_HOME", None)

@@ -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: ...
@@ -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)
@@ -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)
@@ -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

@@ -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):
@@ -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
@@ -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):
@@ -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):
@@ -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

@@ -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):
@@ -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
@@ -15,11 +15,8 @@ import aiofiles
 import aiofiles.os as aios
 import aiohttp
 import certifi
-from exo.shared.models.model_cards import ModelTask
-from exo.shared.types.common import ModelId
-from exo.shared.types.memory import Memory
-from exo.shared.types.worker.shards import ShardMetadata
 from huggingface_hub import (
+    model_info,
    snapshot_download,  # pyright: ignore[reportUnknownVariableType]
 )
 from loguru import logger
@@ -27,21 +24,25 @@ from pydantic import (
    TypeAdapter,
 )

-from exo.download.huggingface_utils import (
-    filter_repo_objects,
-    get_allow_patterns,
-    get_auth_headers,
-    get_hf_endpoint,
-    get_hf_token,
-)
-from exo.shared.constants import EXO_MODELS_DIR, EXO_MODELS_PATH
-from exo.shared.types.worker.downloads import (
+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):
@@ -704,7 +705,9 @@ async def resolve_allow_patterns(shard: ShardMetadata) -> list[str]:
        return ["*"]


-def is_image_model(shard: ShardMetadata) -> bool:
+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

@@ -762,7 +765,7 @@ async def download_shard(

    # For image models, skip root-level safetensors files since weights
    # are stored in component subdirectories (e.g., transformer/, vae/)
-    if is_image_model(shard):
+    if _is_image_model(shard):
        filtered_file_list = [
            f
            for f in filtered_file_list
@@ -881,3 +884,50 @@ async def download_shard(
        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)
@@ -5,9 +5,10 @@ from typing import Callable, Generator, Iterable

 import aiofiles
 import aiofiles.os as aios
-from exo.shared.types.worker.shards import ShardMetadata
 from loguru import logger

+from exo_core.types.shards import ShardMetadata
+

 def filter_repo_objects[T](
    items: Iterable[T],
@@ -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):
@@ -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 }
@@ -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" },
+]
@@ -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 = {
@@ -1,6 +0,0 @@
-def main():
-    print("Hello from vllm-runner!")
-
-
-if __name__ == "__main__":
-    main()
@@ -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"
@@ -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
@@ -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]:
@@ -4,6 +4,19 @@ import time
 from collections.abc import AsyncGenerator
 from typing import Any

+from exo_core.types.chunks import (
+    ErrorChunk,
+    PrefillProgressChunk,
+    TokenChunk,
+    ToolCallChunk,
+)
+from exo_core.types.common import CommandId
+from exo_core.types.text_generation import (
+    InputMessage,
+    TextGenerationTaskParams,
+    resolve_reasoning_params,
+)
+
 from exo.api.types import (
    ChatCompletionChoice,
    ChatCompletionMessage,
@@ -19,18 +32,6 @@ from exo.api.types import (
    ToolCall,
    Usage,
 )
-from exo.shared.types.chunks import (
-    ErrorChunk,
-    PrefillProgressChunk,
-    TokenChunk,
-    ToolCallChunk,
-)
-from exo.shared.types.common import CommandId
-from exo.shared.types.text_generation import (
-    InputMessage,
-    TextGenerationTaskParams,
-    resolve_reasoning_params,
-)


 def chat_request_to_text_generation(
@@ -5,6 +5,15 @@ import re
 from collections.abc import AsyncGenerator
 from typing import Any

+from exo_core.types.chunks import (
+    ErrorChunk,
+    PrefillProgressChunk,
+    TokenChunk,
+    ToolCallChunk,
+)
+from exo_core.types.common import CommandId
+from exo_core.types.text_generation import InputMessage, TextGenerationTaskParams
+
 from exo.api.types import FinishReason, Usage
 from exo.api.types.claude_api import (
    ClaudeContentBlock,
@@ -29,14 +38,6 @@ from exo.api.types.claude_api import (
    ClaudeToolUseBlock,
    ClaudeUsage,
 )
-from exo.shared.types.chunks import (
-    ErrorChunk,
-    PrefillProgressChunk,
-    TokenChunk,
-    ToolCallChunk,
-)
-from exo.shared.types.common import CommandId
-from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams


 def finish_reason_to_claude_stop_reason(
@@ -2,6 +2,15 @@ import json
 from collections.abc import AsyncGenerator
 from typing import Any

+from exo_core.types.chunks import (
+    ErrorChunk,
+    PrefillProgressChunk,
+    TokenChunk,
+    ToolCallChunk,
+)
+from exo_core.types.common import CommandId
+from exo_core.types.text_generation import InputMessage, TextGenerationTaskParams
+
 from exo.api.types.ollama_api import (
    OllamaChatRequest,
    OllamaChatResponse,
@@ -12,14 +21,6 @@ from exo.api.types.ollama_api import (
    OllamaToolCall,
    OllamaToolFunction,
 )
-from exo.shared.types.chunks import (
-    ErrorChunk,
-    PrefillProgressChunk,
-    TokenChunk,
-    ToolCallChunk,
-)
-from exo.shared.types.common import CommandId
-from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams


 def _map_done_reason(
@@ -4,6 +4,19 @@ from collections.abc import AsyncGenerator
 from itertools import count
 from typing import Any

+from exo_core.types.chunks import (
+    ErrorChunk,
+    PrefillProgressChunk,
+    TokenChunk,
+    ToolCallChunk,
+)
+from exo_core.types.common import CommandId
+from exo_core.types.text_generation import (
+    InputMessage,
+    TextGenerationTaskParams,
+    resolve_reasoning_params,
+)
+
 from exo.api.types import Usage
 from exo.api.types.openai_responses import (
    FunctionCallInputItem,
@@ -35,18 +48,6 @@ from exo.api.types.openai_responses import (
    ResponseTextDoneEvent,
    ResponseUsage,
 )
-from exo.shared.types.chunks import (
-    ErrorChunk,
-    PrefillProgressChunk,
-    TokenChunk,
-    ToolCallChunk,
-)
-from exo.shared.types.common import CommandId
-from exo.shared.types.text_generation import (
-    InputMessage,
-    TextGenerationTaskParams,
-    resolve_reasoning_params,
-)


 def _format_sse(event: ResponsesStreamEvent) -> str:
@@ -12,6 +12,34 @@ from uuid import uuid4

 import anyio
 from anyio import BrokenResourceError
+from exo_core.constants import (
+    DASHBOARD_DIR,
+    EXO_CACHE_HOME,
+    EXO_EVENT_LOG_DIR,
+    EXO_IMAGE_CACHE_DIR,
+    EXO_MAX_CHUNK_SIZE,
+    EXO_TRACING_CACHE_DIR,
+)
+from exo_core.model_cards import (
+    ModelCard,
+    ModelId,
+    delete_custom_card,
+    get_model_cards,
+    is_custom_card,
+)
+from exo_core.types.chunks import (
+    ErrorChunk,
+    ImageChunk,
+    InputImageChunk,
+    PrefillProgressChunk,
+    TokenChunk,
+    ToolCallChunk,
+)
+from exo_core.types.common import CommandId, Id, NodeId, SystemId
+from exo_core.types.downloads import DownloadCompleted
+from exo_core.types.instances import Instance, InstanceId, InstanceMeta
+from exo_core.types.shards import Sharding
+from exo_core.utils.memory import Memory
 from fastapi import FastAPI, File, Form, HTTPException, Query, Request, UploadFile
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import FileResponse, JSONResponse, StreamingResponse
@@ -116,32 +144,9 @@ from exo.api.types.openai_responses import (
 from exo.master.image_store import ImageStore
 from exo.master.placement import place_instance as get_instance_placements
 from exo.shared.apply import apply
-from exo.shared.constants import (
-    DASHBOARD_DIR,
-    EXO_CACHE_HOME,
-    EXO_EVENT_LOG_DIR,
-    EXO_IMAGE_CACHE_DIR,
-    EXO_MAX_CHUNK_SIZE,
-    EXO_TRACING_CACHE_DIR,
-)
 from exo.shared.election import ElectionMessage
 from exo.shared.logging import InterceptLogger
-from exo.shared.models.model_cards import (
-    ModelCard,
-    ModelId,
-    delete_custom_card,
-    get_model_cards,
-    is_custom_card,
-)
 from exo.shared.tracing import TraceEvent, compute_stats, export_trace, load_trace_file
-from exo.shared.types.chunks import (
-    ErrorChunk,
-    ImageChunk,
-    InputImageChunk,
-    PrefillProgressChunk,
-    TokenChunk,
-    ToolCallChunk,
-)
 from exo.shared.types.commands import (
    Command,
    CreateInstance,
@@ -159,18 +164,13 @@ from exo.shared.types.commands import (
    TaskFinished,
    TextGeneration,
 )
-from exo.shared.types.common import CommandId, Id, NodeId, SystemId
 from exo.shared.types.events import (
    ChunkGenerated,
    Event,
    IndexedEvent,
    TracesMerged,
 )
-from exo.shared.types.memory import Memory
 from exo.shared.types.state import State
-from exo.shared.types.worker.downloads import DownloadCompleted
-from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
-from exo.shared.types.worker.shards import Sharding
 from exo.utils.banner import print_startup_banner
 from exo.utils.channels import Receiver, Sender, channel
 from exo.utils.disk_event_log import DiskEventLog
@@ -2,7 +2,7 @@
 from typing import Any
 from unittest.mock import AsyncMock, MagicMock

-from exo.shared.types.common import CommandId
+from exo_core.types.common import CommandId
 from fastapi import FastAPI
 from fastapi.testclient import TestClient

@@ -2,7 +2,7 @@

 import pydantic
 import pytest
-from exo.shared.types.common import ModelId
+from exo_core.types.common import ModelId

 from exo.api.adapters.claude import (
    claude_request_to_text_generation,
@@ -4,8 +4,8 @@ import json
 from collections.abc import AsyncGenerator
 from typing import Any, cast

-from exo.shared.types.chunks import ErrorChunk, TokenChunk, ToolCallChunk
-from exo.shared.types.common import CommandId, ModelId
+from exo_core.types.chunks import ErrorChunk, TokenChunk, ToolCallChunk
+from exo_core.types.common import CommandId, ModelId

 from exo.api.adapters.claude import (
    ClaudeMessagesResponse,
@@ -6,7 +6,7 @@ The responses adapter converts it to TextGenerationTaskParams for the pipeline.

 import pydantic
 import pytest
-from exo.shared.types.common import ModelId
+from exo_core.types.common import ModelId

 from exo.api.types.openai_responses import (
    ResponseInputMessage,
@@ -3,16 +3,15 @@ from collections.abc import Generator
 from typing import Annotated, Any, Literal, get_args
 from uuid import uuid4

+from exo_core.model_cards import ModelCard
+from exo_core.models import CamelCaseModel
+from exo_core.types.common import CommandId, ModelId, NodeId
+from exo_core.types.instances import Instance, InstanceId, InstanceMeta
+from exo_core.types.shards import Sharding, ShardMetadata
+from exo_core.types.text_generation import ReasoningEffort
+from exo_core.utils.memory import Memory
 from pydantic import BaseModel, Field, field_validator

-from exo.shared.models.model_cards import ModelCard, ModelId
-from exo.shared.types.common import CommandId, NodeId
-from exo.shared.types.memory import Memory
-from exo.shared.types.text_generation import ReasoningEffort
-from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
-from exo.shared.types.worker.shards import Sharding, ShardMetadata
-from exo.utils.pydantic_ext import CamelCaseModel
-
 FinishReason = Literal[
    "stop", "length", "tool_calls", "content_filter", "function_call", "error"
 ]
@@ -2,10 +2,9 @@

 from typing import Any, Literal

+from exo_core.types.common import ModelId
 from pydantic import BaseModel, Field

-from exo.shared.types.common import ModelId
-
 # Tool definition types
 ClaudeToolInputSchema = dict[str, Any]

@@ -1,10 +1,9 @@
 import time
 from typing import Any, Literal

+from exo_core.model_cards import ModelId
 from pydantic import BaseModel, Field

-from exo.shared.models.model_cards import ModelId
-
 # https://github.com/ollama/ollama/blob/main/docs/api.md

 OllamaRole = Literal["system", "user", "assistant", "tool"]
@@ -3,17 +3,16 @@
 These types model the OpenAI Responses API request/response format.
 ResponsesRequest is the API-level wire type; for the canonical internal
 task params type used by the inference pipeline, see
-``exo.shared.types.text_generation.TextGenerationTaskParams``.
+``exo_core.types.text_generation.TextGenerationTaskParams``.
 """

 import time
 from typing import Any, Literal

+from exo_core.types.common import ModelId
+from exo_core.types.text_generation import ReasoningEffort
 from pydantic import BaseModel, Field

-from exo.shared.types.common import ModelId
-from exo.shared.types.text_generation import ReasoningEffort
-
 # Type aliases
 ResponseStatus = Literal["completed", "failed", "in_progress", "incomplete"]
 ResponseRole = Literal["user", "assistant", "system", "developer"]
@@ -2,19 +2,26 @@ from dataclasses import dataclass, field

 import anyio
 from anyio import current_time
-from exo.shared.models.model_cards import ModelId, get_model_cards
-from exo.shared.types.common import NodeId
-from exo.shared.types.worker.shards import PipelineShardMetadata, ShardMetadata
-from loguru import logger
-
-from exo.download.download_utils import (
+from exo_core.constants import EXO_MODELS_DIR, EXO_MODELS_PATH
+from exo_core.model_cards import ModelId, get_model_cards
+from exo_core.types.common import NodeId
+from exo_core.types.downloads import (
+    DownloadCompleted,
+    DownloadFailed,
+    DownloadOngoing,
+    DownloadPending,
+    DownloadProgress,
+)
+from exo_core.types.shards import PipelineShardMetadata, ShardMetadata
+from exo_core.utils.downloads import (
    RepoDownloadProgress,
    delete_model,
    map_repo_download_progress_to_download_progress_data,
    resolve_model_in_path,
 )
+from loguru import logger
+
 from exo.download.shard_downloader import ShardDownloader
-from exo.shared.constants import EXO_MODELS_DIR, EXO_MODELS_PATH
 from exo.shared.types.commands import (
    CancelDownload,
    DeleteDownload,
@@ -25,13 +32,6 @@ from exo.shared.types.events import (
    Event,
    NodeDownloadProgress,
 )
-from exo.shared.types.worker.downloads import (
-    DownloadCompleted,
-    DownloadFailed,
-    DownloadOngoing,
-    DownloadPending,
-    DownloadProgress,
-)
 from exo.utils.channels import Receiver, Sender
 from exo.utils.task_group import TaskGroup

@@ -4,14 +4,14 @@ from collections.abc import Awaitable
 from pathlib import Path
 from typing import AsyncIterator, Callable

-from exo.shared.models.model_cards import ModelCard, ModelId, get_model_cards
-from exo.shared.types.worker.shards import (
+from exo_core.model_cards import ModelCard, ModelId, get_model_cards
+from exo_core.types.shards import (
    PipelineShardMetadata,
    ShardMetadata,
 )
+from exo_core.utils.downloads import RepoDownloadProgress, download_shard
 from loguru import logger

-from exo.download.download_utils import RepoDownloadProgress, download_shard
 from exo.download.shard_downloader import ShardDownloader


@@ -5,14 +5,13 @@ from datetime import timedelta
 from pathlib import Path
 from typing import AsyncIterator, Callable

-from exo.shared.models.model_cards import ModelCard, ModelId, ModelTask
-from exo.shared.types.memory import Memory
-from exo.shared.types.worker.shards import (
+from exo_core.model_cards import ModelCard, ModelId, ModelTask
+from exo_core.types.shards import (
    PipelineShardMetadata,
    ShardMetadata,
 )
-
-from exo.download.download_utils import RepoDownloadProgress
+from exo_core.utils.downloads import RepoDownloadProgress
+from exo_core.utils.memory import Memory


 # TODO: the PipelineShardMetadata getting reinstantiated is a bit messy. Should this be a classmethod?
@@ -9,15 +9,14 @@ from unittest.mock import AsyncMock, MagicMock, patch
 import aiofiles
 import aiofiles.os as aios
 import pytest
-from exo.shared.types.common import ModelId
-from exo.shared.types.memory import Memory
-from pydantic import TypeAdapter
-
-from exo.download.download_utils import (
+from exo_core.types.common import ModelId
+from exo_core.types.downloads import FileListEntry, RepoFileDownloadProgress
+from exo_core.utils.downloads import (
    delete_model,
    fetch_file_list_with_cache,
 )
-from exo.shared.types.worker.downloads import FileListEntry, RepoFileDownloadProgress
+from exo_core.utils.memory import Memory
+from pydantic import TypeAdapter


@pytest.fixture
@@ -30,7 +29,7 @@ async def temp_models_dir(tmp_path: Path) -> AsyncIterator[Path]:
    """Set up a temporary models directory for testing."""
    models_dir = tmp_path / "models"
    await aios.makedirs(models_dir, exist_ok=True)
-    with patch("exo.download.download_utils.EXO_MODELS_DIR", models_dir):
+    with patch("exo_core.utils.downloads.EXO_MODELS_DIR", models_dir):
        yield models_dir


@@ -42,7 +41,7 @@ class TestFileVerification:
    ) -> None:
        """Test that files with mismatched sizes are re-downloaded."""
        # Import inside test to allow patching
-        from exo.download.download_utils import (
+        from exo_core.utils.downloads import (
            _download_file,  # pyright: ignore[reportPrivateUsage]
        )

@@ -59,12 +58,12 @@ class TestFileVerification:

        with (
            patch(
-                "exo.download.download_utils.file_meta",
+                "exo_core.utils.downloads.file_meta",
                new_callable=AsyncMock,
                return_value=(remote_size, remote_hash),
            ) as mock_file_meta,
            patch(
-                "exo.download.download_utils.create_http_session"
+                "exo_core.utils.downloads.create_http_session"
            ) as mock_session_factory,
        ):
            # Set up mock HTTP response for re-download
@@ -90,7 +89,7 @@ class TestFileVerification:

            # Mock calc_hash to return the expected hash
            with patch(
-                "exo.download.download_utils.calc_hash",
+                "exo_core.utils.downloads.calc_hash",
                new_callable=AsyncMock,
                return_value=remote_hash,
            ):
@@ -103,7 +102,7 @@ class TestFileVerification:
        self, model_id: ModelId, tmp_path: Path
    ) -> None:
        """Test that files with matching sizes are not re-downloaded."""
-        from exo.download.download_utils import (
+        from exo_core.utils.downloads import (
            _download_file,  # pyright: ignore[reportPrivateUsage]
        )

@@ -121,12 +120,12 @@ class TestFileVerification:

        with (
            patch(
-                "exo.download.download_utils.file_meta",
+                "exo_core.utils.downloads.file_meta",
                new_callable=AsyncMock,
                return_value=(remote_size, remote_hash),
            ) as mock_file_meta,
            patch(
-                "exo.download.download_utils.create_http_session"
+                "exo_core.utils.downloads.create_http_session"
            ) as mock_session_factory,
        ):
            result = await _download_file(
@@ -142,7 +141,7 @@ class TestFileVerification:
        self, model_id: ModelId, tmp_path: Path
    ) -> None:
        """Test that local files are used when network is unavailable."""
-        from exo.download.download_utils import (
+        from exo_core.utils.downloads import (
            _download_file,  # pyright: ignore[reportPrivateUsage]
        )

@@ -156,12 +155,12 @@ class TestFileVerification:

        with (
            patch(
-                "exo.download.download_utils.file_meta",
+                "exo_core.utils.downloads.file_meta",
                new_callable=AsyncMock,
                side_effect=Exception("Network error"),
            ),
            patch(
-                "exo.download.download_utils.create_http_session"
+                "exo_core.utils.downloads.create_http_session"
            ) as mock_session_factory,
        ):
            result = await _download_file(
@@ -188,9 +187,9 @@ class TestFileListCache:
        ]

        with (
-            patch("exo.download.download_utils.EXO_MODELS_DIR", models_dir),
+            patch("exo_core.utils.downloads.EXO_MODELS_DIR", models_dir),
            patch(
-                "exo.download.download_utils.fetch_file_list_with_retry",
+                "exo_core.utils.downloads.fetch_file_list_with_retry",
                new_callable=AsyncMock,
                return_value=file_list,
            ) as mock_fetch,
@@ -234,9 +233,9 @@ class TestFileListCache:
            )

        with (
-            patch("exo.download.download_utils.EXO_MODELS_DIR", models_dir),
+            patch("exo_core.utils.downloads.EXO_MODELS_DIR", models_dir),
            patch(
-                "exo.download.download_utils.fetch_file_list_with_retry",
+                "exo_core.utils.downloads.fetch_file_list_with_retry",
                new_callable=AsyncMock,
                side_effect=Exception("Network error"),
            ),
@@ -252,9 +251,9 @@ class TestFileListCache:
        models_dir = tmp_path / "models"

        with (
-            patch("exo.download.download_utils.EXO_MODELS_DIR", models_dir),
+            patch("exo_core.utils.downloads.EXO_MODELS_DIR", models_dir),
            patch(
-                "exo.download.download_utils.fetch_file_list_with_retry",
+                "exo_core.utils.downloads.fetch_file_list_with_retry",
                new_callable=AsyncMock,
                side_effect=Exception("Network error"),
            ),
@@ -284,7 +283,7 @@ class TestModelDeletion:
        async with aiofiles.open(cache_dir / "file_list.json", "w") as f:
            await f.write("[]")

-        with patch("exo.download.download_utils.EXO_MODELS_DIR", models_dir):
+        with patch("exo_core.utils.downloads.EXO_MODELS_DIR", models_dir):
            result = await delete_model(model_id)

            assert result is True
@@ -303,7 +302,7 @@ class TestModelDeletion:
        async with aiofiles.open(cache_dir / "file_list.json", "w") as f:
            await f.write("[]")

-        with patch("exo.download.download_utils.EXO_MODELS_DIR", models_dir):
+        with patch("exo_core.utils.downloads.EXO_MODELS_DIR", models_dir):
            result = await delete_model(model_id)

            # Returns False because model dir didn't exist
@@ -318,7 +317,7 @@ class TestModelDeletion:
        models_dir = tmp_path / "models"
        await aios.makedirs(models_dir, exist_ok=True)

-        with patch("exo.download.download_utils.EXO_MODELS_DIR", models_dir):
+        with patch("exo_core.utils.downloads.EXO_MODELS_DIR", models_dir):
            result = await delete_model(model_id)

            assert result is False
@@ -7,14 +7,13 @@ from unittest.mock import AsyncMock, patch
 import aiofiles
 import aiofiles.os as aios
 import pytest
-from exo.shared.types.common import ModelId
-
-from exo.download.download_utils import (
+from exo_core.types.common import ModelId
+from exo_core.types.downloads import FileListEntry
+from exo_core.utils.downloads import (
    _download_file,  # pyright: ignore[reportPrivateUsage]
    download_file_with_retry,
    fetch_file_list_with_cache,
 )
-from exo.shared.types.worker.downloads import FileListEntry


@pytest.fixture
@@ -26,7 +25,7 @@ def model_id() -> ModelId:
 async def temp_models_dir(tmp_path: Path) -> AsyncIterator[Path]:
    models_dir = tmp_path / "models"
    await aios.makedirs(models_dir, exist_ok=True)
-    with patch("exo.download.download_utils.EXO_MODELS_DIR", models_dir):
+    with patch("exo_core.utils.downloads.EXO_MODELS_DIR", models_dir):
        yield models_dir


@@ -46,7 +45,7 @@ class TestDownloadFileOffline:
            await f.write(b"model weights data")

        with patch(
-            "exo.download.download_utils.file_meta",
+            "exo_core.utils.downloads.file_meta",
            new_callable=AsyncMock,
        ) as mock_file_meta:
            result = await _download_file(
@@ -91,7 +90,7 @@ class TestDownloadFileOffline:
            await f.write(b"weights")

        with patch(
-            "exo.download.download_utils.file_meta",
+            "exo_core.utils.downloads.file_meta",
            new_callable=AsyncMock,
        ) as mock_file_meta:
            result = await _download_file(
@@ -121,7 +120,7 @@ class TestDownloadFileWithRetryOffline:
            await f.write(b'{"model_type": "qwen2"}')

        with patch(
-            "exo.download.download_utils.file_meta",
+            "exo_core.utils.downloads.file_meta",
            new_callable=AsyncMock,
        ) as mock_file_meta:
            result = await download_file_with_retry(
@@ -175,7 +174,7 @@ class TestFetchFileListOffline:
            )

        with patch(
-            "exo.download.download_utils.fetch_file_list_with_retry",
+            "exo_core.utils.downloads.fetch_file_list_with_retry",
            new_callable=AsyncMock,
        ) as mock_fetch:
            result = await fetch_file_list_with_cache(
@@ -209,7 +208,7 @@ class TestFetchFileListOffline:
            await f.write(b"x" * 500)

        with patch(
-            "exo.download.download_utils.fetch_file_list_with_retry",
+            "exo_core.utils.downloads.fetch_file_list_with_retry",
            new_callable=AsyncMock,
        ) as mock_fetch:
            result = await fetch_file_list_with_cache(
@@ -8,13 +8,14 @@ from pathlib import Path
 from typing import Callable
 from unittest.mock import AsyncMock, patch

-from exo.shared.models.model_cards import ModelCard, ModelId, ModelTask
-from exo.shared.types.common import NodeId, SystemId
-from exo.shared.types.memory import Memory
-from exo.shared.types.worker.shards import PipelineShardMetadata, ShardMetadata
+from exo_core.model_cards import ModelCard, ModelId, ModelTask
+from exo_core.types.common import NodeId, SystemId
+from exo_core.types.downloads import DownloadCompleted
+from exo_core.types.shards import PipelineShardMetadata, ShardMetadata
+from exo_core.utils.downloads import RepoDownloadProgress
+from exo_core.utils.memory import Memory

 from exo.download.coordinator import DownloadCoordinator
-from exo.download.download_utils import RepoDownloadProgress
 from exo.download.impl_shard_downloader import SingletonShardDownloader
 from exo.download.shard_downloader import ShardDownloader
 from exo.shared.types.commands import (
@@ -23,7 +24,6 @@ from exo.shared.types.commands import (
    StartDownload,
 )
 from exo.shared.types.events import Event, NodeDownloadProgress
-from exo.shared.types.worker.downloads import DownloadCompleted
 from exo.utils.channels import Receiver, Sender, channel

 NODE_ID = NodeId("aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa")
@@ -7,6 +7,9 @@ from dataclasses import dataclass, field
 from typing import Self

 import anyio
+from exo_core.constants import EXO_LOG
+from exo_core.models import CamelCaseModel
+from exo_core.types.common import NodeId, SessionId
 from loguru import logger
 from pydantic import PositiveInt

@@ -17,12 +20,9 @@ from exo.download.impl_shard_downloader import exo_shard_downloader
 from exo.master.main import Master
 from exo.routing.event_router import EventRouter
 from exo.routing.router import Router, get_node_id_keypair
-from exo.shared.constants import EXO_LOG
 from exo.shared.election import Election, ElectionResult
 from exo.shared.logging import logger_cleanup, logger_setup
-from exo.shared.types.common import NodeId, SessionId
 from exo.utils.channels import Receiver, channel
-from exo.utils.pydantic_ext import CamelCaseModel
 from exo.utils.task_group import TaskGroup
 from exo.worker.main import Worker

@@ -1,10 +1,9 @@
 import time
 from pathlib import Path

+from exo_core.types.common import Id
 from pydantic import BaseModel

-from exo.shared.types.common import Id
-

 class StoredImage(BaseModel, frozen=True):
    image_id: Id
@@ -1,6 +1,22 @@
 from datetime import datetime, timedelta, timezone

 import anyio
+from exo_core.constants import EXO_EVENT_LOG_DIR, EXO_TRACING_ENABLED
+from exo_core.types.common import CommandId, NodeId, SessionId, SystemId
+from exo_core.types.instances import InstanceId
+from exo_core.types.tasks import (
+    ImageEdits as ImageEditsTask,
+)
+from exo_core.types.tasks import (
+    ImageGeneration as ImageGenerationTask,
+)
+from exo_core.types.tasks import (
+    TaskId,
+    TaskStatus,
+)
+from exo_core.types.tasks import (
+    TextGeneration as TextGenerationTask,
+)
 from loguru import logger

 from exo.master.placement import (
@@ -11,7 +27,6 @@ from exo.master.placement import (
    place_instance,
 )
 from exo.shared.apply import apply
-from exo.shared.constants import EXO_EVENT_LOG_DIR, EXO_TRACING_ENABLED
 from exo.shared.types.commands import (
    CreateInstance,
    DeleteInstance,
@@ -27,7 +42,6 @@ from exo.shared.types.commands import (
    TestCommand,
    TextGeneration,
 )
-from exo.shared.types.common import CommandId, NodeId, SessionId, SystemId
 from exo.shared.types.events import (
    Event,
    GlobalForwarderEvent,
@@ -45,20 +59,6 @@ from exo.shared.types.events import (
    TracesMerged,
 )
 from exo.shared.types.state import State
-from exo.shared.types.tasks import (
-    ImageEdits as ImageEditsTask,
-)
-from exo.shared.types.tasks import (
-    ImageGeneration as ImageGenerationTask,
-)
-from exo.shared.types.tasks import (
-    TaskId,
-    TaskStatus,
-)
-from exo.shared.types.tasks import (
-    TextGeneration as TextGenerationTask,
-)
-from exo.shared.types.worker.instances import InstanceId
 from exo.utils.channels import Receiver, Sender
 from exo.utils.disk_event_log import DiskEventLog
 from exo.utils.event_buffer import MultiSourceBuffer
@@ -3,6 +3,24 @@ from collections.abc import Mapping
 from copy import deepcopy
 from typing import Sequence

+from exo_core.model_cards import ModelId
+from exo_core.types.common import NodeId
+from exo_core.types.downloads import (
+    DownloadOngoing,
+    DownloadProgress,
+)
+from exo_core.types.instances import (
+    Instance,
+    InstanceId,
+    InstanceMeta,
+    MlxJacclInstance,
+    MlxRingInstance,
+    VllmInstance,
+)
+from exo_core.types.shards import Sharding
+from exo_core.types.tasks import Task, TaskId, TaskStatus
+from exo_core.utils.memory import Memory
+
 from exo.master.placement_utils import (
    Cycle,
    filter_cycles_by_memory,
@@ -12,7 +30,6 @@ from exo.master.placement_utils import (
    get_shard_assignments,
    get_smallest_cycles,
 )
-from exo.shared.models.model_cards import ModelId
 from exo.shared.topology import Topology
 from exo.shared.types.commands import (
    CancelDownload,
@@ -21,29 +38,13 @@ from exo.shared.types.commands import (
    DownloadCommand,
    PlaceInstance,
 )
-from exo.shared.types.common import NodeId
 from exo.shared.types.events import (
    Event,
    InstanceCreated,
    InstanceDeleted,
    TaskStatusUpdated,
 )
-from exo.shared.types.memory import Memory
 from exo.shared.types.profiling import MemoryUsage, NodeNetworkInfo
-from exo.shared.types.tasks import Task, TaskId, TaskStatus
-from exo.shared.types.worker.downloads import (
-    DownloadOngoing,
-    DownloadProgress,
-)
-from exo.shared.types.worker.instances import (
-    Instance,
-    InstanceId,
-    InstanceMeta,
-    MlxJacclInstance,
-    MlxRingInstance,
-    VllmInstance,
-)
-from exo.shared.types.worker.shards import Sharding


 def random_ephemeral_port() -> int:
@@ -1,21 +1,21 @@
 from collections.abc import Generator, Mapping

-from loguru import logger
-
-from exo.shared.models.model_cards import ModelCard
-from exo.shared.topology import Topology
-from exo.shared.types.common import Host, NodeId
-from exo.shared.types.memory import Memory
-from exo.shared.types.profiling import MemoryUsage, NodeNetworkInfo
-from exo.shared.types.topology import Cycle, RDMAConnection, SocketConnection
-from exo.shared.types.worker.runners import RunnerId, ShardAssignments
-from exo.shared.types.worker.shards import (
+from exo_core.model_cards import ModelCard
+from exo_core.types.common import Host, NodeId
+from exo_core.types.runners import RunnerId, ShardAssignments
+from exo_core.types.shards import (
    CfgShardMetadata,
    PipelineShardMetadata,
    Sharding,
    ShardMetadata,
    TensorShardMetadata,
 )
+from exo_core.utils.memory import Memory
+from loguru import logger
+
+from exo.shared.topology import Topology
+from exo.shared.types.profiling import MemoryUsage, NodeNetworkInfo
+from exo.shared.types.topology import Cycle, RDMAConnection, SocketConnection


 def filter_cycles_by_memory(
@@ -3,18 +3,18 @@ from typing import Sequence

 import anyio
 import pytest
-from exo.shared.models.model_cards import ModelCard, ModelTask
-from exo.shared.types.common import ModelId, NodeId, SessionId, SystemId
-from exo.shared.types.memory import Memory
-from exo.shared.types.tasks import TaskStatus
-from exo.shared.types.tasks import TextGeneration as TextGenerationTask
-from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
-from exo.shared.types.worker.instances import (
+from exo_core.model_cards import ModelCard, ModelTask
+from exo_core.types.common import ModelId, NodeId, SessionId, SystemId
+from exo_core.types.instances import (
    InstanceMeta,
    MlxRingInstance,
    ShardAssignments,
 )
-from exo.shared.types.worker.shards import PipelineShardMetadata, Sharding
+from exo_core.types.shards import PipelineShardMetadata, Sharding
+from exo_core.types.tasks import TaskStatus
+from exo_core.types.tasks import TextGeneration as TextGenerationTask
+from exo_core.types.text_generation import InputMessage, TextGenerationTaskParams
+from exo_core.utils.memory import Memory
 from loguru import logger

 from exo.master.main import Master
@@ -1,18 +1,18 @@
 import pytest
-from exo.shared.models.model_cards import ModelCard, ModelId, ModelTask
-from exo.shared.types.common import CommandId, NodeId
-from exo.shared.types.memory import Memory
-from exo.shared.types.tasks import TaskId, TaskStatus, TextGeneration
-from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
-from exo.shared.types.worker.instances import (
+from exo_core.model_cards import ModelCard, ModelId, ModelTask
+from exo_core.types.common import CommandId, NodeId
+from exo_core.types.instances import (
    Instance,
    InstanceId,
    InstanceMeta,
    MlxJacclInstance,
    MlxRingInstance,
 )
-from exo.shared.types.worker.runners import ShardAssignments
-from exo.shared.types.worker.shards import Sharding
+from exo_core.types.runners import ShardAssignments
+from exo_core.types.shards import Sharding
+from exo_core.types.tasks import TaskId, TaskStatus, TextGeneration
+from exo_core.types.text_generation import InputMessage, TextGenerationTaskParams
+from exo_core.utils.memory import Memory

 from exo.master.placement import (
    get_transition_events,
@@ -1,12 +1,12 @@
 import pytest
-from exo.shared.models.model_cards import ModelCard, ModelId, ModelTask
-from exo.shared.types.common import NodeId
-from exo.shared.types.memory import Memory
-from exo.shared.types.worker.shards import (
+from exo_core.model_cards import ModelCard, ModelId, ModelTask
+from exo_core.types.common import NodeId
+from exo_core.types.shards import (
    CfgShardMetadata,
    PipelineShardMetadata,
    Sharding,
 )
+from exo_core.utils.memory import Memory

 from exo.master.placement_utils import (
    allocate_layers_proportionally,
@@ -1,5 +1,5 @@
 import pytest
-from exo.shared.types.common import NodeId
+from exo_core.types.common import NodeId

 from exo.shared.topology import Topology
 from exo.shared.types.multiaddr import Multiaddr
@@ -1,8 +1,7 @@
+from exo_core.models import CamelCaseModel
+from exo_core.types.common import NodeId
 from exo_pyo3_bindings import PyFromSwarm

-from exo.shared.types.common import NodeId
-from exo.utils.pydantic_ext import CamelCaseModel
-
 """Serialisable types for Connection Updates/Messages"""


@@ -4,10 +4,10 @@ from random import random
 import anyio
 from anyio import BrokenResourceError, ClosedResourceError
 from anyio.abc import CancelScope
+from exo_core.types.common import SessionId, SystemId
 from loguru import logger

 from exo.shared.types.commands import ForwarderCommand, RequestEventLog
-from exo.shared.types.common import SessionId, SystemId
 from exo.shared.types.events import (
    Event,
    EventId,
@@ -11,6 +11,8 @@ from anyio import (
    move_on_after,
    sleep_forever,
 )
+from exo_core.constants import EXO_NODE_ID_KEYPAIR
+from exo_core.models import CamelCaseModel
 from exo_pyo3_bindings import (
    AllQueuesFullError,
    Keypair,
@@ -22,9 +24,7 @@ from exo_pyo3_bindings import (
 from filelock import FileLock
 from loguru import logger

-from exo.shared.constants import EXO_NODE_ID_KEYPAIR
 from exo.utils.channels import Receiver, Sender, channel
-from exo.utils.pydantic_ext import CamelCaseModel
 from exo.utils.task_group import TaskGroup

 from .connection_message import ConnectionMessage
@@ -1,6 +1,8 @@
 from dataclasses import dataclass
 from enum import Enum

+from exo_core.models import CamelCaseModel
+
 from exo.routing.connection_message import ConnectionMessage
 from exo.shared.election import ElectionMessage
 from exo.shared.types.commands import ForwarderCommand, ForwarderDownloadCommand
@@ -8,7 +10,6 @@ from exo.shared.types.events import (
    GlobalForwarderEvent,
    LocalForwarderEvent,
 )
-from exo.utils.pydantic_ext import CamelCaseModel


 class PublishPolicy(str, Enum):
@@ -2,9 +2,13 @@ import copy
 from collections.abc import Mapping, Sequence
 from datetime import datetime

+from exo_core.types.common import NodeId
+from exo_core.types.downloads import DownloadProgress
+from exo_core.types.instances import Instance, InstanceId
+from exo_core.types.runners import RunnerId, RunnerShutdown, RunnerStatus
+from exo_core.types.tasks import Task, TaskId, TaskStatus
 from loguru import logger

-from exo.shared.types.common import NodeId
 from exo.shared.types.events import (
    ChunkGenerated,
    Event,
@@ -35,11 +39,7 @@ from exo.shared.types.profiling import (
    ThunderboltBridgeStatus,
 )
 from exo.shared.types.state import State
-from exo.shared.types.tasks import Task, TaskId, TaskStatus
 from exo.shared.types.topology import Connection, RDMAConnection
-from exo.shared.types.worker.downloads import DownloadProgress
-from exo.shared.types.worker.instances import Instance, InstanceId
-from exo.shared.types.worker.runners import RunnerId, RunnerShutdown, RunnerStatus
 from exo.utils.info_gatherer.info_gatherer import (
    MacmonMetrics,
    MacThunderboltConnections,
@@ -6,13 +6,13 @@ from anyio import (
    Event,
    get_cancelled_exc_class,
 )
+from exo_core.models import CamelCaseModel
+from exo_core.types.common import NodeId, SessionId
 from loguru import logger

 from exo.routing.connection_message import ConnectionMessage
 from exo.shared.types.commands import ForwarderCommand
-from exo.shared.types.common import NodeId, SessionId
 from exo.utils.channels import Receiver, Sender
-from exo.utils.pydantic_ext import CamelCaseModel
 from exo.utils.task_group import TaskGroup

 DEFAULT_ELECTION_TIMEOUT = 3.0
@@ -5,12 +5,11 @@ from typing import Generator

 import pytest
 from _pytest.logging import LogCaptureFixture
+from exo_core.model_cards import ModelCard, ModelId, ModelTask
+from exo_core.types.shards import PipelineShardMetadata, ShardMetadata
+from exo_core.utils.memory import Memory
 from loguru import logger

-from exo.shared.models.model_cards import ModelCard, ModelId, ModelTask
-from exo.shared.types.memory import Memory
-from exo.shared.types.worker.shards import PipelineShardMetadata, ShardMetadata
-

@pytest.fixture(scope="session")
 def event_loop() -> Generator[asyncio.AbstractEventLoop, None, None]:
@@ -1,11 +1,11 @@
-from exo.shared.types.common import NodeId
-from exo.shared.types.memory import Memory
+from exo_core.types.common import NodeId
+from exo_core.types.downloads import DownloadCompleted
+from exo_core.utils.memory import Memory

 from exo.shared.apply import apply_node_download_progress
 from exo.shared.tests.conftest import get_pipeline_shard_metadata
 from exo.shared.types.events import NodeDownloadProgress
 from exo.shared.types.state import State
-from exo.shared.types.worker.downloads import DownloadCompleted
 from exo.worker.tests.constants import MODEL_A_ID, MODEL_B_ID


@@ -1,4 +1,4 @@
-from exo.shared.types.worker.runners import RunnerId, RunnerIdle, RunnerShutdown
+from exo_core.types.runners import RunnerId, RunnerIdle, RunnerShutdown

 from exo.shared.apply import apply_runner_status_updated
 from exo.shared.types.events import RunnerStatusUpdated
@@ -1,10 +1,10 @@
 import pytest
 from anyio import create_task_group, fail_after, move_on_after
+from exo_core.types.common import NodeId, SessionId, SystemId

 from exo.routing.connection_message import ConnectionMessage
 from exo.shared.election import Election, ElectionMessage, ElectionResult
 from exo.shared.types.commands import ForwarderCommand, TestCommand
-from exo.shared.types.common import NodeId, SessionId, SystemId
 from exo.utils.channels import channel

 # ======= #
@@ -7,11 +7,11 @@ from multiprocessing.queues import Queue as QueueT
 from multiprocessing.synchronize import Event as EventT
 from multiprocessing.synchronize import Semaphore as SemaphoreT

+from exo_core.constants import EXO_NODE_ID_KEYPAIR
 from loguru import logger
 from pytest import LogCaptureFixture, mark

 from exo.routing.router import get_node_id_keypair
-from exo.shared.constants import EXO_NODE_ID_KEYPAIR

 NUM_CONCURRENT_PROCS = 10

@@ -1,6 +1,5 @@
 import pytest
-
-from exo.shared.types.text_generation import resolve_reasoning_params
+from exo_core.types.text_generation import resolve_reasoning_params


 def test_both_none_returns_none_none() -> None:
@@ -1,4 +1,5 @@
-from exo.shared.types.common import NodeId
+from exo_core.types.common import NodeId
+
 from exo.shared.types.multiaddr import Multiaddr
 from exo.shared.types.state import State
 from exo.shared.types.topology import Connection, SocketConnection
@@ -23,7 +23,7 @@ def test_xdg_paths_on_linux():
        # Re-import to pick up mocked values
        import importlib

-        import exo.shared.constants as constants
+        import exo_core.constants as constants

        importlib.reload(constants)

@@ -46,7 +46,7 @@ def test_xdg_default_paths_on_linux():
    ):
        import importlib

-        import exo.shared.constants as constants
+        import exo_core.constants as constants

        importlib.reload(constants)

@@ -68,7 +68,7 @@ def test_legacy_exo_home_takes_precedence():
    ):
        import importlib

-        import exo.shared.constants as constants
+        import exo_core.constants as constants

        importlib.reload(constants)

@@ -87,7 +87,7 @@ def test_macos_uses_traditional_paths():
    ):
        import importlib

-        import exo.shared.constants as constants
+        import exo_core.constants as constants

        importlib.reload(constants)

@@ -99,7 +99,7 @@ def test_macos_uses_traditional_paths():

 def test_node_id_in_config_dir():
    """Test that node ID keypair is in the config directory."""
-    import exo.shared.constants as constants
+    import exo_core.constants as constants

    assert constants.EXO_NODE_ID_KEYPAIR.parent == constants.EXO_CONFIG_HOME

@@ -111,7 +111,7 @@ def test_models_in_data_dir():
    with mock.patch.dict(os.environ, env, clear=True):
        import importlib

-        import exo.shared.constants as constants
+        import exo_core.constants as constants

        importlib.reload(constants)

@@ -4,9 +4,9 @@ from dataclasses import dataclass, field
 from typing import Iterable

 import rustworkx as rx
+from exo_core.types.common import NodeId
 from pydantic import BaseModel, ConfigDict

-from exo.shared.types.common import NodeId
 from exo.shared.types.profiling import (
    InterfaceType,
    NodeNetworkInfo,
@@ -8,7 +8,8 @@ from dataclasses import dataclass, field
 from pathlib import Path
 from typing import cast, final

-from exo.shared.constants import EXO_TRACING_ENABLED
+from exo_core.constants import EXO_TRACING_ENABLED
+
 from exo.worker.runner.bootstrap import logger

 # Context variable to track the current trace category for hierarchical nesting
@@ -1,16 +1,16 @@
+from exo_core.model_cards import ModelCard
+from exo_core.models import CamelCaseModel, TaggedModel
+from exo_core.types.chunks import InputImageChunk
+from exo_core.types.common import CommandId, ModelId, NodeId, SystemId
+from exo_core.types.instances import Instance, InstanceId, InstanceMeta
+from exo_core.types.shards import Sharding, ShardMetadata
+from exo_core.types.text_generation import TextGenerationTaskParams
 from pydantic import Field

 from exo.api.types import (
    ImageEditsTaskParams,
    ImageGenerationTaskParams,
 )
-from exo.shared.models.model_cards import ModelCard, ModelId
-from exo.shared.types.chunks import InputImageChunk
-from exo.shared.types.common import CommandId, NodeId, SystemId
-from exo.shared.types.text_generation import TextGenerationTaskParams
-from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
-from exo.shared.types.worker.shards import Sharding, ShardMetadata
-from exo.utils.pydantic_ext import CamelCaseModel, TaggedModel


 class BaseCommand(TaggedModel):
@@ -1,17 +1,17 @@
 from datetime import datetime
 from typing import final

+from exo_core.models import CamelCaseModel, FrozenModel, TaggedModel
+from exo_core.types.chunks import GenerationChunk, InputImageChunk
+from exo_core.types.common import CommandId, Id, NodeId, SessionId, SystemId
+from exo_core.types.downloads import DownloadProgress
+from exo_core.types.instances import Instance, InstanceId
+from exo_core.types.runners import RunnerId, RunnerStatus
+from exo_core.types.tasks import Task, TaskId, TaskStatus
 from pydantic import Field

 from exo.shared.topology import Connection
-from exo.shared.types.chunks import GenerationChunk, InputImageChunk
-from exo.shared.types.common import CommandId, Id, NodeId, SessionId, SystemId
-from exo.shared.types.tasks import Task, TaskId, TaskStatus
-from exo.shared.types.worker.downloads import DownloadProgress
-from exo.shared.types.worker.instances import Instance, InstanceId
-from exo.shared.types.worker.runners import RunnerId, RunnerStatus
 from exo.utils.info_gatherer.info_gatherer import GatheredInfo
-from exo.utils.pydantic_ext import CamelCaseModel, FrozenModel, TaggedModel


 class EventId(Id):
@@ -4,10 +4,10 @@ from pathlib import Path
 from typing import Literal, Self

 import psutil
+from exo_core.models import CamelCaseModel
+from exo_core.utils.memory import Memory

-from exo.shared.types.memory import Memory
 from exo.shared.types.thunderbolt import ThunderboltIdentifier
-from exo.utils.pydantic_ext import CamelCaseModel


 class MemoryUsage(CamelCaseModel):
@@ -2,11 +2,16 @@ from collections.abc import Mapping, Sequence
 from datetime import datetime
 from typing import Any, cast

+from exo_core.models import CamelCaseModel
+from exo_core.types.common import NodeId
+from exo_core.types.downloads import DownloadProgress
+from exo_core.types.instances import Instance, InstanceId
+from exo_core.types.runners import RunnerId, RunnerStatus
+from exo_core.types.tasks import Task, TaskId
 from pydantic import ConfigDict, Field, field_serializer, field_validator
 from pydantic.alias_generators import to_camel

 from exo.shared.topology import Topology, TopologySnapshot
-from exo.shared.types.common import NodeId
 from exo.shared.types.profiling import (
    DiskUsage,
    MemoryUsage,
@@ -17,11 +22,6 @@ from exo.shared.types.profiling import (
    SystemPerformanceProfile,
    ThunderboltBridgeStatus,
 )
-from exo.shared.types.tasks import Task, TaskId
-from exo.shared.types.worker.downloads import DownloadProgress
-from exo.shared.types.worker.instances import Instance, InstanceId
-from exo.shared.types.worker.runners import RunnerId, RunnerStatus
-from exo.utils.pydantic_ext import CamelCaseModel


 class State(CamelCaseModel):
@@ -1,8 +1,7 @@
 import anyio
+from exo_core.models import CamelCaseModel
 from pydantic import BaseModel, Field

-from exo.utils.pydantic_ext import CamelCaseModel
-

 class ThunderboltConnection(CamelCaseModel):
    source_uuid: str
@@ -1,9 +1,10 @@
 from collections.abc import Iterator
 from dataclasses import dataclass

-from exo.shared.types.common import NodeId
+from exo_core.models import FrozenModel
+from exo_core.types.common import NodeId
+
 from exo.shared.types.multiaddr import Multiaddr
-from exo.utils.pydantic_ext import FrozenModel


@dataclass(frozen=True)
@@ -1,98 +0,0 @@
-from datetime import timedelta
-from typing import Literal
-
-from exo.shared.types.common import NodeId
-from exo.shared.types.memory import Memory
-from exo.shared.types.worker.shards import ShardMetadata
-from exo.utils.pydantic_ext import CamelCaseModel, TaggedModel
-from pydantic import BaseModel, ConfigDict, Field, PositiveInt
-
-
-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)
@@ -3,7 +3,7 @@ import os
 import sys
 import webbrowser

-from exo.shared.constants import EXO_CONFIG_HOME
+from exo_core.constants import EXO_CONFIG_HOME

 logger = logging.getLogger(__name__)

@@ -11,11 +11,12 @@ import anyio
 from anyio import fail_after, open_process, to_thread
 from anyio.streams.buffered import BufferedByteReceiveStream
 from anyio.streams.text import TextReceiveStream
+from exo_core.constants import EXO_CONFIG_FILE, EXO_MODELS_DIR
+from exo_core.models import TaggedModel
+from exo_core.utils.memory import Memory
 from loguru import logger
 from pydantic import ValidationError

-from exo.shared.constants import EXO_CONFIG_FILE, EXO_MODELS_DIR
-from exo.shared.types.memory import Memory
 from exo.shared.types.profiling import (
    DiskUsage,
    MemoryUsage,
@@ -28,7 +29,6 @@ from exo.shared.types.thunderbolt import (
    ThunderboltIdentifier,
 )
 from exo.utils.channels import Sender
-from exo.utils.pydantic_ext import TaggedModel
 from exo.utils.task_group import TaskGroup

 from .macmon import MacmonMetrics
@@ -1,9 +1,9 @@
 from typing import Self

+from exo_core.models import TaggedModel
 from pydantic import BaseModel

 from exo.shared.types.profiling import MemoryUsage, SystemPerformanceProfile
-from exo.utils.pydantic_ext import TaggedModel


 class _TempMetrics(BaseModel, extra="ignore"):
@@ -4,10 +4,10 @@ from collections.abc import AsyncGenerator, Mapping
 import anyio
 import httpx
 from anyio import create_task_group
+from exo_core.types.common import NodeId
 from loguru import logger

 from exo.shared.topology import Topology
-from exo.shared.types.common import NodeId
 from exo.shared.types.profiling import NodeNetworkInfo
 from exo.utils.channels import Sender, channel

@@ -1,6 +1,7 @@
 # pyright: reportMissingImports=false
+from exo_core.models import TaggedModel
+
 from exo.shared.types.profiling import SystemPerformanceProfile
-from exo.utils.pydantic_ext import TaggedModel

 try:
    from pynvml import (
@@ -4,9 +4,9 @@ from collections.abc import Callable, Mapping
 from typing import final

 import anyio
+from exo_core.types.common import NodeId

 from exo.api.types import NodePowerStats, PowerUsage
-from exo.shared.types.common import NodeId
 from exo.shared.types.profiling import SystemPerformanceProfile


@@ -2,7 +2,7 @@ from collections.abc import Mapping

 import anyio
 import pytest
-from exo.shared.types.common import NodeId
+from exo_core.types.common import NodeId

 from exo.api.types import PowerUsage
 from exo.shared.types.profiling import SystemPerformanceProfile
@@ -1,6 +1,6 @@
 import anyio
 import pytest
-from exo.utils.pydantic_ext import TaggedModel
+from exo_core.models import TaggedModel
 from pydantic import BaseModel, TypeAdapter, ValidationError


@@ -3,13 +3,13 @@ from pathlib import Path
 from typing import Any, Literal, Optional

 import mlx.core as mx
+from exo_core.types.instances import BoundInstance
+from exo_core.types.shards import CfgShardMetadata, PipelineShardMetadata
+from exo_core.utils.downloads import build_model_path
 from mflux.models.common.config.config import Config
 from PIL import Image

 from exo.api.types import AdvancedImageParams
-from exo.download.download_utils import build_model_path
-from exo.shared.types.worker.instances import BoundInstance
-from exo.shared.types.worker.shards import CfgShardMetadata, PipelineShardMetadata
 from exo.worker.engines.image.config import ImageModelConfig
 from exo.worker.engines.image.models import (
    create_adapter_for_model,
@@ -7,6 +7,11 @@ from pathlib import Path
 from typing import Generator, Literal

 import mlx.core as mx
+from exo_core.types.runner_response import (
+    ImageGenerationResponse,
+    PartialImageResponse,
+)
+from exo_core.utils.memory import Memory
 from PIL import Image

 from exo.api.types import (
@@ -16,11 +21,6 @@ from exo.api.types import (
    ImageGenerationTaskParams,
    ImageSize,
 )
-from exo.shared.types.memory import Memory
-from exo.shared.types.worker.runner_response import (
-    ImageGenerationResponse,
-    PartialImageResponse,
-)
 from exo.worker.engines.image.distributed_model import DistributedImageModel


@@ -4,16 +4,16 @@ from math import ceil
 from typing import Any, Optional, final

 import mlx.core as mx
+from exo_core.constants import EXO_TRACING_ENABLED
+from exo_core.types.shards import CfgShardMetadata, PipelineShardMetadata
 from mflux.models.common.config.config import Config
 from mflux.utils.exceptions import StopImageGenerationException
 from tqdm import tqdm

-from exo.shared.constants import EXO_TRACING_ENABLED
 from exo.shared.tracing import (
    clear_trace_buffer,
    trace,
 )
-from exo.shared.types.worker.shards import CfgShardMetadata, PipelineShardMetadata
 from exo.worker.engines.image.config import ImageModelConfig
 from exo.worker.engines.image.models.base import (
    ModelAdapter,
@@ -8,6 +8,7 @@ from typing import TYPE_CHECKING, Any, Literal, Protocol, cast

 import mlx.core as mx
 import mlx.nn as nn
+from exo_core.types.shards import PipelineShardMetadata
 from mlx.nn.layers.distributed import (
    shard_inplace,
    shard_linear,
@@ -58,7 +59,6 @@ from mlx_lm.models.step3p5 import Step3p5MLP as Step35MLP
 from mlx_lm.models.step3p5 import Step3p5Model as Step35InnerModel

 from exo.shared.logging import logger
-from exo.shared.types.worker.shards import PipelineShardMetadata

 if TYPE_CHECKING:
    from mlx_lm.models.cache import Cache
@@ -4,6 +4,7 @@ from typing import TYPE_CHECKING, cast

 import mlx.core as mx
 import psutil
+from exo_core.utils.memory import Memory
 from mlx_lm.models.cache import (
    ArraysCache,
    CacheList,
@@ -13,7 +14,6 @@ from mlx_lm.models.cache import (
 )
 from mlx_lm.tokenizer_utils import TokenizerWrapper

-from exo.shared.types.memory import Memory
 from exo.shared.types.mlx import KVCacheType, MLXCacheType, Model
 from exo.worker.engines.mlx.constants import CACHE_GROUP_SIZE, KV_CACHE_BITS
 from exo.worker.runner.bootstrap import logger
@@ -3,6 +3,11 @@ from dataclasses import dataclass, field
 from typing import Callable, cast

 import mlx.core as mx
+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 mlx_lm.generate import (
    BatchGenerator as MlxBatchGenerator,
 )
@@ -18,12 +23,7 @@ from exo.api.types import (
    TopLogprobItem,
    Usage,
 )
-from exo.shared.types.common import ModelId
-from exo.shared.types.memory import Memory
 from exo.shared.types.mlx import MLXCacheType, Model
-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.mlx.cache import (
    CacheSnapshot,
    KVPrefixCache,
@@ -5,6 +5,12 @@ from copy import deepcopy
 from typing import Callable, Generator, cast, get_args

 import mlx.core as mx
+from exo_core.types.common import ModelId
+from exo_core.types.runner_response import (
+    GenerationResponse,
+)
+from exo_core.types.text_generation import InputMessage, TextGenerationTaskParams
+from exo_core.utils.memory import Memory
 from mlx_lm.generate import (
    maybe_quantize_kv_cache,
    stream_generate,
@@ -21,13 +27,7 @@ from exo.api.types import (
    TopLogprobItem,
    Usage,
 )
-from exo.shared.types.common import ModelId
-from exo.shared.types.memory import Memory
 from exo.shared.types.mlx import MLXCacheType, Model
-from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
-from exo.shared.types.worker.runner_response import (
-    GenerationResponse,
-)
 from exo.worker.engines.mlx.auto_parallel import (
    PipelineFirstLayer,
    PipelineLastLayer,
@@ -19,11 +19,11 @@ try:
 except ImportError:
    pass  # transformers < 5.0 or bytes_to_unicode not available

+from exo_core.model_cards import ModelId
 from mlx_lm.models.cache import KVCache
 from mlx_lm.models.deepseek_v3 import DeepseekV3Model
 from mlx_lm.tokenizer_utils import TokenizerWrapper

-from exo.shared.models.model_cards import ModelId
 from exo.worker.engines.mlx.constants import TRUST_REMOTE_CODE

 try:
@@ -34,26 +34,26 @@ import contextlib

 import mlx.core as mx
 import mlx.nn as nn
-from mlx_lm.utils import load_model
-from pydantic import RootModel
-
-from exo.download.download_utils import build_model_path
-from exo.shared.types.common import Host
-from exo.shared.types.memory import Memory
-from exo.shared.types.mlx import Model
-from exo.shared.types.tasks import TaskId, TextGeneration
-from exo.shared.types.text_generation import TextGenerationTaskParams
-from exo.shared.types.worker.instances import (
+from exo_core.types.common import Host
+from exo_core.types.instances import (
    BoundInstance,
    MlxJacclInstance,
    MlxRingInstance,
 )
-from exo.shared.types.worker.shards import (
+from exo_core.types.shards import (
    CfgShardMetadata,
    PipelineShardMetadata,
    ShardMetadata,
    TensorShardMetadata,
 )
+from exo_core.types.tasks import TaskId, TextGeneration
+from exo_core.types.text_generation import TextGenerationTaskParams
+from exo_core.utils.downloads import build_model_path
+from exo_core.utils.memory import Memory
+from mlx_lm.utils import load_model
+from pydantic import RootModel
+
+from exo.shared.types.mlx import Model
 from exo.worker.engines.mlx.auto_parallel import (
    LayerLoadedCallback,
    TimeoutCallback,
@@ -3,18 +3,29 @@ from datetime import datetime, timezone

 import anyio
 from anyio import fail_after
+from exo_core.model_cards import ModelId
+from exo_core.types.common import CommandId, NodeId, SystemId
+from exo_core.types.downloads import DownloadCompleted
+from exo_core.types.runners import RunnerId
+from exo_core.types.tasks import (
+    CancelTask,
+    CreateRunner,
+    DownloadModel,
+    ImageEdits,
+    Shutdown,
+    Task,
+    TaskStatus,
+)
+from exo_core.utils.downloads import resolve_model_in_path
 from loguru import logger

 from exo.api.types import ImageEditsTaskParams
-from exo.download.download_utils import resolve_model_in_path
 from exo.shared.apply import apply
-from exo.shared.models.model_cards import ModelId
 from exo.shared.types.commands import (
    ForwarderCommand,
    ForwarderDownloadCommand,
    StartDownload,
 )
-from exo.shared.types.common import CommandId, NodeId, SystemId
 from exo.shared.types.events import (
    Event,
    IndexedEvent,
@@ -28,18 +39,7 @@ from exo.shared.types.events import (
 )
 from exo.shared.types.multiaddr import Multiaddr
 from exo.shared.types.state import State
-from exo.shared.types.tasks import (
-    CancelTask,
-    CreateRunner,
-    DownloadModel,
-    ImageEdits,
-    Shutdown,
-    Task,
-    TaskStatus,
-)
 from exo.shared.types.topology import Connection, SocketConnection
-from exo.shared.types.worker.downloads import DownloadCompleted
-from exo.shared.types.worker.runners import RunnerId
 from exo.utils.channels import Receiver, Sender, channel
 from exo.utils.info_gatherer.info_gatherer import GatheredInfo, InfoGatherer
 from exo.utils.info_gatherer.net_profile import check_reachable
@@ -2,8 +2,28 @@

 from collections.abc import Mapping, Sequence

-from exo.shared.types.common import CommandId, NodeId
-from exo.shared.types.tasks import (
+from exo_core.types.common import CommandId, NodeId
+from exo_core.types.downloads import (
+    DownloadCompleted,
+    DownloadFailed,
+    DownloadOngoing,
+    DownloadProgress,
+)
+from exo_core.types.instances import BoundInstance, Instance, InstanceId
+from exo_core.types.runners import (
+    RunnerConnected,
+    RunnerConnecting,
+    RunnerFailed,
+    RunnerId,
+    RunnerIdle,
+    RunnerLoaded,
+    RunnerLoading,
+    RunnerReady,
+    RunnerRunning,
+    RunnerStatus,
+    RunnerWarmingUp,
+)
+from exo_core.types.tasks import (
    CancelTask,
    ConnectToGroup,
    CreateRunner,
@@ -18,26 +38,7 @@ from exo.shared.types.tasks import (
    TaskStatus,
    TextGeneration,
 )
-from exo.shared.types.worker.downloads import (
-    DownloadCompleted,
-    DownloadFailed,
-    DownloadOngoing,
-    DownloadProgress,
-)
-from exo.shared.types.worker.instances import BoundInstance, Instance, InstanceId
-from exo.shared.types.worker.runners import (
-    RunnerConnected,
-    RunnerConnecting,
-    RunnerFailed,
-    RunnerId,
-    RunnerIdle,
-    RunnerLoaded,
-    RunnerLoading,
-    RunnerReady,
-    RunnerRunning,
-    RunnerStatus,
-    RunnerWarmingUp,
-)
+
 from exo.worker.runner.runner_supervisor import RunnerSupervisor


@@ -5,11 +5,11 @@ import sys
 from pathlib import Path

 import loguru
+from exo_core.types.instances import BoundInstance, VllmInstance
+from exo_core.types.runners import RunnerFailed
+from exo_core.types.tasks import Task, TaskId

 from exo.shared.types.events import Event, RunnerStatusUpdated
-from exo.shared.types.tasks import Task, TaskId
-from exo.shared.types.worker.instances import BoundInstance, VllmInstance
-from exo.shared.types.worker.runners import RunnerFailed
 from exo.utils.channels import ClosedResourceError, MpReceiver, MpSender

 logger: "loguru.Logger" = loguru.logger
@@ -72,7 +72,8 @@ def entrypoint(
            os.environ["VLLM_KV_CACHE_LAYOUT"] = "NHD"
            os.environ["VLLM_BATCH_INVARIANT"] = "1"
            _ensure_cuda_libs()
-            from exo.shared.constants import EXO_MODELS_DIR
+            from exo_core.constants import EXO_MODELS_DIR
+
            from exo.worker.runner.llm_inference.runner import Runner, VllmBuilder

            model_id = bound_instance.bound_shard.model_card.model_id
--- a/Show More
+++ b/Show More