diff --git a/pyproject.toml b/pyproject.toml index 7c79bd6d..3ddcd7cb 100644 --- a/pyproject.toml +++ b/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: diff --git a/python/exo_core/pyproject.toml b/python/exo_core/pyproject.toml index 7c15df50..0418c1ba 100644 --- a/python/exo_core/pyproject.toml +++ b/python/exo_core/pyproject.toml @@ -7,7 +7,9 @@ authors = [ { name = "Evan", email = "evanev7@gmail.com" } ] requires-python = ">=3.13" -dependencies = [] +dependencies = [ + "pydantic", +] [build-system] requires = ["uv_build>=0.9.24,<0.10.0"] diff --git a/src/exo/shared/models/model_cards.py b/python/exo_core/src/exo_core/model_cards.py similarity index 100% rename from src/exo/shared/models/model_cards.py rename to python/exo_core/src/exo_core/model_cards.py diff --git a/src/exo/utils/pydantic_ext.py b/python/exo_core/src/exo_core/pydantic.py similarity index 100% rename from src/exo/utils/pydantic_ext.py rename to python/exo_core/src/exo_core/pydantic.py diff --git a/python/exo_core/src/exo_core/types/__init__.py b/python/exo_core/src/exo_core/types/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/exo/shared/types/chunks.py b/python/exo_core/src/exo_core/types/chunks.py similarity index 96% rename from src/exo/shared/types/chunks.py rename to python/exo_core/src/exo_core/types/chunks.py index 9faf0cd6..9146b1d3 100644 --- a/src/exo/shared/types/chunks.py +++ b/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 diff --git a/src/exo/shared/types/common.py b/python/exo_core/src/exo_core/types/common.py similarity index 96% rename from src/exo/shared/types/common.py rename to python/exo_core/src/exo_core/types/common.py index b5674b4b..295274ae 100644 --- a/src/exo/shared/types/common.py +++ b/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): diff --git a/src/exo/shared/types/worker/instances.py b/python/exo_core/src/exo_core/types/instances.py similarity index 94% rename from src/exo/shared/types/worker/instances.py rename to python/exo_core/src/exo_core/types/instances.py index ee21f4f2..a19b842e 100644 --- a/src/exo/shared/types/worker/instances.py +++ b/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): diff --git a/src/exo/shared/types/worker/runner_response.py b/python/exo_core/src/exo_core/types/runner_response.py similarity index 100% rename from src/exo/shared/types/worker/runner_response.py rename to python/exo_core/src/exo_core/types/runner_response.py diff --git a/src/exo/shared/types/worker/runners.py b/python/exo_core/src/exo_core/types/runners.py similarity index 88% rename from src/exo/shared/types/worker/runners.py rename to python/exo_core/src/exo_core/types/runners.py index 1ac68947..fcd82c22 100644 --- a/src/exo/shared/types/worker/runners.py +++ b/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): diff --git a/src/exo/shared/types/worker/shards.py b/python/exo_core/src/exo_core/types/shards.py similarity index 95% rename from src/exo/shared/types/worker/shards.py rename to python/exo_core/src/exo_core/types/shards.py index 59a6c54e..3f7628e2 100644 --- a/src/exo/shared/types/worker/shards.py +++ b/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): diff --git a/src/exo/shared/types/tasks.py b/python/exo_core/src/exo_core/types/tasks.py similarity index 84% rename from src/exo/shared/types/tasks.py rename to python/exo_core/src/exo_core/types/tasks.py index e764ec33..65e78447 100644 --- a/src/exo/shared/types/tasks.py +++ b/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): diff --git a/src/exo/shared/types/text_generation.py b/python/exo_core/src/exo_core/types/text_generation.py similarity index 98% rename from src/exo/shared/types/text_generation.py rename to python/exo_core/src/exo_core/types/text_generation.py index 3c2bcb48..b3424fb9 100644 --- a/src/exo/shared/types/text_generation.py +++ b/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"] diff --git a/python/exo_core/src/exo_core/utils/__init__.py b/python/exo_core/src/exo_core/utils/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/exo/shared/types/memory.py b/python/exo_core/src/exo_core/utils/memory.py similarity index 99% rename from src/exo/shared/types/memory.py rename to python/exo_core/src/exo_core/utils/memory.py index 2684d9b9..728c283a 100644 --- a/src/exo/shared/types/memory.py +++ b/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): diff --git a/python/vllm_engine/pyproject.toml b/python/vllm_engine/pyproject.toml index d745152a..b20558a3 100644 --- a/python/vllm_engine/pyproject.toml +++ b/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"