move chunks to core

pyproject.toml

@@ -47,8 +47,8 @@ dev = [
 ]
 
 [project.optional-dependencies]
-mlx = ["mlx_runner"]
-cuda = ["vllm_runner; sys_platform == 'linux' and platform_machine == 'aarch64'"]
+mlx = ["mlx_engine"]
+cuda = ["vllm_engine; sys_platform == 'linux' and platform_machine == 'aarch64'"]
 
 ###
 # workspace configuration
@@ -63,8 +63,9 @@ members = [
 
 [tool.uv.sources]
 exo_pyo3_bindings = { workspace = true }
-mlx_runner = { path = "python/mlx_runner" }
-vllm_runner = { path = "python/vllm_runner" }
+exo_core = { workspace = true }
+mlx_engine = { path = "python/mlx_engine" }
+vllm_engine = { path = "python/vllm_engine" }
 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:

python/exo_core/pyproject.toml

@@ -7,7 +7,9 @@ authors = [
     { name = "Evan", email = "[email protected]" }
 ]
 requires-python = ">=3.13"
-dependencies = []
+dependencies = [
+    "pydantic",
+]
 
 [build-system]
 requires = ["uv_build>=0.9.24,<0.10.0"]

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

@@ -9,8 +9,8 @@ from exo.api.types import (
     TopLogprobItem,
     Usage,
 )
-from exo.shared.models.model_cards import ModelId
-from exo.utils.pydantic_ext import TaggedModel
+from exo_core.model_cards import ModelId
+from exo_core.pydantic import TaggedModel
 
 from .common import CommandId
 

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

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

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

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

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

@@ -2,10 +2,10 @@ from collections.abc import Mapping
 
 from pydantic import model_validator
 
-from exo.shared.models.model_cards import ModelId
-from exo.shared.types.common import Id, NodeId
-from exo.shared.types.worker.shards import ShardMetadata
-from exo.utils.pydantic_ext import CamelCaseModel, TaggedModel
+from exo_core.model_cards import ModelId
+from exo_core.types.common import Id, NodeId
+from exo_core.types.shards import ShardMetadata
+from exo_core.pydantic import CamelCaseModel, TaggedModel
 
 
 class RunnerId(Id):

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

@@ -3,8 +3,8 @@ from typing import TypeAlias, final
 
 from pydantic import Field
 
-from exo.shared.models.model_cards import ModelCard
-from exo.utils.pydantic_ext import TaggedModel
+from exo_core.model_cards import ModelCard
+from exo_core.pydantic import TaggedModel
 
 
 class Sharding(str, Enum):

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

@@ -6,12 +6,12 @@ from exo.api.types import (
     ImageEditsTaskParams,
     ImageGenerationTaskParams,
 )
-from exo.shared.types.common import CommandId, Id
-from exo.shared.types.text_generation import TextGenerationTaskParams
-from exo.shared.types.worker.instances import BoundInstance, InstanceId
-from exo.shared.types.worker.runners import RunnerId
-from exo.shared.types.worker.shards import ShardMetadata
-from exo.utils.pydantic_ext import TaggedModel
+from exo_core.types.common import CommandId, Id
+from exo_core.types.text_generation import TextGenerationTaskParams
+from exo_core.types.instances import BoundInstance, InstanceId
+from exo_core.types.runners import RunnerId
+from exo_core.types.shards import ShardMetadata
+from exo_core.pydantic import TaggedModel
 
 
 class TaskId(Id):

python/exo_core/src/exo_core/types/text_generation.py

@@ -8,7 +8,7 @@ from typing import Any, Literal
 
 from pydantic import BaseModel
 
-from exo.shared.types.common import ModelId
+from exo_core.types.common import ModelId
 
 MessageRole = Literal["user", "assistant", "system", "developer"]
 ReasoningEffort = Literal["none", "minimal", "low", "medium", "high", "xhigh"]

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

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

python/vllm_engine/pyproject.toml

@@ -1,5 +1,5 @@
 [project]
-name = "vllm-runner"
+name = "vllm-engine"
 version = "0.1.0"
 description = "Add your description here"
 readme = "README.md"