urgk

python/vllm_engine/src/vllm_engine/builder.py

@@ -1,30 +1,37 @@
+import contextlib
+import os
 from dataclasses import dataclass
-from typing import Self, Callable
+from typing import Callable, Self
+
 from exo_core.constants import EXO_MODELS_DIR
-from exo_core.engine import EngineBuilder, Engine
-from exo_core.types.common import ModelId
+from exo_core.engine import Engine, EngineBuilder
+from exo_core.types.chunks import ErrorChunk, PrefillProgressChunk
+from exo_core.types.common import CommandId, ModelId
 from exo_core.types.instances import BoundInstance
-from exo_core.types.tasks import TextGeneration
-from exo_core.types.runner_response import GenerationResponse
-from vllm_engine.vllm_generator import VllmBatchEngine
-from vllm_engine.vllm_generator import load_vllm_engine
+from exo_core.types.runner_response import GenerationResponse, ToolCallResponse
+from exo_core.types.tasks import TaskId, TextGeneration
+from exo_core.utils.channels import MpReceiver, MpSender
+from loguru import logger
+from mlx_engine.batch_generator import BatchGenerator
+
+from vllm_engine.vllm_generator import VllmBatchEngine, load_vllm_engine
 
 
 @dataclass
-class VllmBuilder(EngineBuilder[BoundInstance, TextGeneration, GenerationResponse]):
+class VllmBuilder(EngineBuilder[BoundInstance, TextGeneration, GenerationResponse | ToolCallResponse]):
     model_id: ModelId
     model_path: str
     trust_remote_code: bool
     cancel_receiver: MpReceiver[TaskId]
-    event_sender: MpSender[Event]
+    event_sender: MpSender[tuple[CommandId, ErrorChunk | PrefillProgressChunk]]
     bound_instance: BoundInstance
 
     @classmethod
     def create(
         cls,
         bound_instance: BoundInstance,
-        event_sender: MpSender[Event],
-        cancel_receiver: MpReceiver[TaskId],
+    cancel_receiver: MpReceiver[TaskId],
+    event_sender: MpSender[tuple[CommandId, ErrorChunk | PrefillProgressChunk]],
     ) -> Self:
         mid = bound_instance.instance.shard_assignments.model_id
         return cls(

python/vllm_engine/src/vllm_engine/growable_cache.py

@@ -1,9 +1,8 @@
 import torch
+from loguru import logger
 from mlx_engine.cache import KVPrefixCache
 from vllm.v1.worker.gpu_model_runner import GPUModelRunner
 
-from loguru import logger
-
 INITIAL_FRACTION = 0.05
 GROWTH_HEADROOM_BYTES = 512 * 1024 * 1024
 MIN_GROWTH_BLOCKS = 16

python/vllm_engine/src/vllm_engine/vllm_generator.py

@@ -8,28 +8,27 @@ from collections.abc import Callable, Generator
 from dataclasses import dataclass, field
 
 import torch
-from mlx_engine.cache import KVPrefixCache
-from mlx_engine.utils_mlx import get_eos_token_ids_for_model
+from exo_core.tokenizers.tool_parsers import ToolParser, infer_tool_parser
 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 exo_core.engine import Engine
-from loguru import logger
-from vllm.engine.arg_utils import EngineArgs
-from vllm.v1.attention.backends.registry import AttentionBackendEnum
-from vllm.sampling_params import SamplingParams
-from vllm.v1.engine.llm_engine import LLMEngine
-from vllm.v1.kv_cache_interface import KVCacheConfig
-
 from exo_core.types.runner_response import (
     CompletionTokensDetails,
+    GenerationResponse,
     GenerationStats,
     PromptTokensDetails,
     Usage,
 )
-from exo.worker.runner.llm_inference.tool_parsers import ToolParser, infer_tool_parser
+from exo_core.types.tasks import TaskId
+from exo_core.types.text_generation import TextGenerationTaskParams
+from exo_core.utils.memory import Memory
+from loguru import logger
+from mlx_engine.cache import KVPrefixCache
+from mlx_engine.utils_mlx import get_eos_token_ids_for_model
+from vllm.engine.arg_utils import EngineArgs
+from vllm.sampling_params import SamplingParams
+from vllm.v1.attention.backends.registry import AttentionBackendEnum
+from vllm.v1.engine.llm_engine import LLMEngine
+from vllm.v1.kv_cache_interface import KVCacheConfig
+
 from vllm_engine.growable_cache import (
     get_model_runner,
     patch_vllm,