shmovin

pyproject.toml

@@ -16,6 +16,7 @@ dependencies = [
   "psutil>=7.0.0",
   "loguru>=0.7.3",
   "exo_pyo3_bindings",                       # rust bindings
+  "exo_core",
   "anyio==4.11.0",
   "mlx; sys_platform == 'darwin'",
   "mlx==0.30.6; sys_platform == 'linux'",
@@ -47,14 +48,18 @@ dev = [
 
 [project.optional-dependencies]
 mlx = ["mlx_runner"]
-cuda = ["vllm_runner"]
+cuda = ["vllm_runner; sys_platform == 'linux' and platform_machine == 'aarch64'"]
 
 ###
 # workspace configuration
 ###
 
 [tool.uv.workspace]
-members = ["rust/exo_pyo3_bindings", "bench"]
+members = [
+    "rust/exo_pyo3_bindings",
+    "bench",
+    "python/exo_core",
+]
 
 [tool.uv.sources]
 exo_pyo3_bindings = { workspace = true }
@@ -122,6 +127,8 @@ root = "src"
 required-version = ">=0.8.6"
 prerelease = "allow"
 environments = ["sys_platform == 'darwin'", "sys_platform == 'linux'"]
+conflicts = [[{ package = "exo", extra = "cuda" }, { package = "exo-bench" }]]
+
 
 ###
 # ruff configuration
@@ -129,7 +136,6 @@ environments = ["sys_platform == 'darwin'", "sys_platform == 'linux'"]
 
 [tool.ruff]
 extend-exclude = [
-  "shared/protobufs/**",
   "*mlx_typings/**",
   "*cuda_typings/**",
   "rust/exo_pyo3_bindings/**",

python/exo_core/pyproject.toml

@@ -0,0 +1,14 @@
+[project]
+name = "exo-core"
+version = "0.1.0"
+description = "Add your description here"
+readme = "README.md"
+authors = [
+    { name = "Evan", email = "[email protected]" }
+]
+requires-python = ">=3.13"
+dependencies = []
+
+[build-system]
+requires = ["uv_build>=0.9.24,<0.10.0"]
+build-backend = "uv_build"

python/exo_core/src/exo_core/engine.py

@@ -0,0 +1,65 @@
+from abc import ABC, abstractmethod
+from collections.abc import Callable, Iterable
+from typing import Self
+
+class TaskId(str): ...
+class Cancelled: ...
+class Finished: ...
+
+CANCEL_ALL_TASKS = TaskId("CANCEL_TALL_TASKS")
+
+class Engine[TaskType, ResponseType](ABC):
+    _cancelled_tasks: set[TaskId]
+
+    def should_cancel(self, task_id: TaskId) -> bool:
+        return (
+            task_id in self._cancelled_tasks
+            or CANCEL_ALL_TASKS in self._cancelled_tasks
+        )
+
+    def cancel_task(self, task_id: TaskId):
+        self._cancelled_tasks.add(task_id)
+
+    @abstractmethod
+    def warmup(self) -> None: ...
+
+    @abstractmethod
+    def submit(
+        self,
+        task: TaskType,
+    ) -> None: ...
+
+    @abstractmethod
+    def step(
+        self,
+    ) -> Iterable[
+        tuple[TaskId, ResponseType | Cancelled | Finished]
+    ]: ...
+
+    @abstractmethod
+    def close(self) -> None: ...
+
+class EngineBuilder[SetupType, TaskType, ResponseType](ABC):
+    @classmethod
+    @abstractmethod
+    def create(
+        cls,
+        bound_instance: SetupType,
+    ) -> Self: ...
+
+    @abstractmethod
+    def connect(self) -> None: ...
+
+    @abstractmethod
+    def load(
+        self,
+        on_timeout: Callable[[], None],
+        on_layer_loaded: Callable[[int, int], None],
+    ) -> None: ...
+
+    @abstractmethod
+    def build(self) -> Engine[TaskType, ResponseType]: ...
+
+    @abstractmethod
+    def close(self) -> None: ...
+

python/mlx_engine/pyproject.toml

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

python/mlx_runner/main.py

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

python/vllm_engine/pyproject.toml

@@ -0,0 +1,38 @@
+[project]
+name = "vllm-runner"
+version = "0.1.0"
+description = "Add your description here"
+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'",
+  "mlx-cuda-13==0.30.6; sys_platform == 'linux'",
+  "fastsafetensors>=0.1.10; sys_platform == 'linux'",
+]
+
+[tool.uv]
+environments = [
+    "sys_platform == 'linux' and platform_machine == 'aarch64'"
+]
+
+
+[tool.uv.sources]
+vllm = { git = "https://github.com/hmellor/vllm.git", rev = "b99bedc737166ae5ca98cb9e3534b96e0c8c69aa" }
+torch = [{ index = "pytorch-cu130", marker = "platform_machine == 'aarch64'" },
+  { index = "pytorch-cpu", marker = "platform_machine == 'x86_64'" },
+]
+
+[[tool.uv.index]]
+name = "pytorch-cu130"
+url = "https://download.pytorch.org/whl/cu130"
+explicit = true
+
+[[tool.uv.index]]
+name = "pytorch-cpu"
+url = "https://download.pytorch.org/whl/cpu"
+explicit = true
+
+[build-system]
+requires = ["uv_build>=0.8.9,<0.9.0"]
+build-backend = "uv_build"

python/vllm_runner/pyproject.toml

@@ -1,25 +0,0 @@
-[project]
-name = "vllm-runner"
-version = "0.1.0"
-description = "Add your description here"
-readme = "README.md"
-requires-python = ">=3.13"
-dependencies = [
-  "torch>=2.10.0; sys_platform == 'linux'",
-  "vllm>=0.13.0; sys_platform == 'linux'",
-  "mlx-cuda-13==0.30.6; sys_platform == 'linux'",
-  "fastsafetensors>=0.1.10; sys_platform == 'linux'",
-]
-
-[tool.uv.sources]
-vllm = { git = "https://github.com/hmellor/vllm.git", branch = "transformers-v5" }
-torch = [{ index = "pytorch-cu130", marker = "sys_platform == 'linux'" }]
-
-[[tool.uv.index]]
-name = "pytorch-cu130"
-url = "https://download.pytorch.org/whl/cu130"
-explicit = true
-
-[build-system]
-requires = ["uv_build>=0.8.9,<0.9.0"]
-build-backend = "uv_build"