gotta move the cache soon

python/vllm_engine/main.py

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

python/vllm_engine/pyproject.toml

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

python/vllm_engine/src/vllm_engine/prompt_format.py

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

python/vllm_engine/src/vllm_engine/vllm_generator.py

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