diff --git a/src/exo/worker/engines/vllm/prompt_format.py b/src/exo/worker/engines/vllm/prompt_format.py index eb55350a..404e7967 100644 --- a/src/exo/worker/engines/vllm/prompt_format.py +++ b/src/exo/worker/engines/vllm/prompt_format.py @@ -14,7 +14,7 @@ from exo.worker.runner.bootstrap import logger def format_vllm_prompt( engine: LLMEngine, params: TextGenerationTaskParams -) -> tuple[str, int]: +) -> tuple[list[int], str, int]: tokenizer = engine.get_tokenizer() if params.chat_template_messages is not None: @@ -51,7 +51,7 @@ def format_vllm_prompt( if patched_template is not None: logger.info("Patched lossy chat template (removed inner_type length guard)") - result = tokenizer.apply_chat_template( + prompt_text: str = tokenizer.apply_chat_template( formatted_messages, tokenize=False, add_generation_prompt=True, @@ -59,17 +59,26 @@ def format_vllm_prompt( **({"chat_template": patched_template} if patched_template is not None else {}), **extra_kwargs, ) - assert isinstance(result, str) + assert isinstance(prompt_text, str) - if params.tools and schemas_lost_in_prompt(result, params.tools): + if params.tools and schemas_lost_in_prompt(prompt_text, params.tools): logger.warning("Chat template lost nested tool schemas even after patching") if partial_assistant_content: - result += partial_assistant_content + prompt_text += partial_assistant_content - prompt_token_count = len(tokenizer.encode(result)) + token_ids: list[int] = tokenizer.apply_chat_template( + formatted_messages, + tokenize=True, + add_generation_prompt=True, + tools=params.tools, + **({"chat_template": patched_template} if patched_template is not None else {}), + **extra_kwargs, + ) + if partial_assistant_content: + token_ids += tokenizer.encode(partial_assistant_content, add_special_tokens=False) - return result, prompt_token_count + return token_ids, prompt_text, len(token_ids) def make_vllm_sampling_params( diff --git a/src/exo/worker/engines/vllm/vllm_generator.py b/src/exo/worker/engines/vllm/vllm_generator.py index e260ec1c..6f02d17f 100644 --- a/src/exo/worker/engines/vllm/vllm_generator.py +++ b/src/exo/worker/engines/vllm/vllm_generator.py @@ -101,18 +101,22 @@ class VllmGenerator(InferenceGenerator): if self.should_cancel(task.task_id): self._cancelled_tasks.discard(task.task_id) return [(task.task_id, Cancelled())] - prompt, prompt_token_count = format_vllm_prompt( + token_ids, prompt_text, prompt_token_count = format_vllm_prompt( self.engine, task.task_params ) request_id = str(task.task_id) sampling_params = make_vllm_sampling_params(self.engine, task.task_params) - self.engine.add_request(request_id, prompt, sampling_params) + self.engine.add_request( + request_id, + {"prompt_token_ids": token_ids}, + sampling_params, + ) queue: GeneratorQueue[GenerationResponse] = GeneratorQueue() parsed_gen = apply_vllm_parsers( queue.gen(), self.model_id, - prompt, + prompt_text, self.tool_parser, task.task_params.tools, ) diff --git a/src/exo/worker/runner/llm_inference/model_output_parsers.py b/src/exo/worker/runner/llm_inference/model_output_parsers.py index 8f709777..68cd2000 100644 --- a/src/exo/worker/runner/llm_inference/model_output_parsers.py +++ b/src/exo/worker/runner/llm_inference/model_output_parsers.py @@ -105,7 +105,15 @@ def parse_gpt_oss( try: stream.process(response.token) except HarmonyError: - logger.error("Encountered critical Harmony Error, returning early") + logger.warning( + f"HarmonyError on token={response.token}, falling back to raw text passthrough" + ) + yield response + for remaining in responses: + if remaining is None: + yield None + else: + yield remaining return delta = stream.last_content_delta diff --git a/src/exo/worker/runner/llm_inference/runner.py b/src/exo/worker/runner/llm_inference/runner.py index e90fd7d2..681a4aca 100644 --- a/src/exo/worker/runner/llm_inference/runner.py +++ b/src/exo/worker/runner/llm_inference/runner.py @@ -164,10 +164,11 @@ class Runner: if isinstance(self.current_status, RunnerShutdown): break finally: - if isinstance(self.generator, Builder): - self.generator.shutdown_cleanup() - else: - self.generator.close() + if not isinstance(self.current_status, RunnerShutdown): + if isinstance(self.generator, Builder): + self.generator.shutdown_cleanup() + else: + self.generator.close() def handle_first_task(self, task: Task): self.send_task_status(task.task_id, TaskStatus.Running)