Fix GPT OSS by not retokenizing prompts

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(

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,
             )

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

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)