Merge branch 'main' into leo/truncate-long-logs

src/exo/api/main.py

@@ -172,11 +172,21 @@ from exo.shared.types.events import (
     ChunkGenerated,
     Event,
     IndexedEvent,
+    InstanceDeleted,
     TracesMerged,
 )
 from exo.shared.types.memory import Memory
 from exo.shared.types.state import State
 from exo.shared.types.text_generation import Base64Image, TextGenerationTaskParams
+from exo.shared.types.tasks import (
+    ImageEdits as ImageEditsTask,
+)
+from exo.shared.types.tasks import (
+    ImageGeneration as ImageGenerationTask,
+)
+from exo.shared.types.tasks import (
+    TextGeneration as TextGenerationTask,
+)
 from exo.shared.types.worker.downloads import DownloadCompleted
 from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
 from exo.shared.types.worker.shards import Sharding
@@ -1822,9 +1832,25 @@ class API:
                             await queue.send(event.chunk)
                         except (BrokenResourceError, ClosedResourceError):
                             self._text_generation_queues.pop(event.command_id, None)
+                if isinstance(event, InstanceDeleted):
+                    self._close_streams_for_instance(event.instance_id)
                 if isinstance(event, TracesMerged):
                     self._save_merged_trace(event)
 
+    def _close_streams_for_instance(self, instance_id: InstanceId) -> None:
+        """Close any active generation streams for commands running on the given instance."""
+        for task in self.state.tasks.values():
+            if task.instance_id != instance_id:
+                continue
+            if not isinstance(
+                task, (TextGenerationTask, ImageGenerationTask, ImageEditsTask)
+            ):
+                continue
+            if sender := self._text_generation_queues.pop(task.command_id, None):
+                sender.close()
+            if sender := self._image_generation_queues.pop(task.command_id, None):
+                sender.close()
+
     def _save_merged_trace(self, event: TracesMerged) -> None:
         traces = [
             TraceEvent(

src/exo/api/tests/test_instance_deleted_stream_cleanup.py

@@ -0,0 +1,93 @@
+# pyright: reportUnusedFunction=false, reportAny=false
+"""Tests that InstanceDeleted events close active generation streams."""
+
+from unittest.mock import MagicMock
+
+from exo.api.main import API
+from exo.api.types import ImageGenerationTaskParams
+from exo.shared.types.common import CommandId, ModelId
+from exo.shared.types.state import State
+from exo.shared.types.tasks import ImageGeneration, TextGeneration
+from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
+from exo.shared.types.worker.instances import InstanceId
+
+
+def _make_api_with_state(state: State) -> API:
+    """Create a minimal API instance with pre-set state."""
+    api = object.__new__(API)
+    api.state = state
+    api._text_generation_queues = {}  # pyright: ignore[reportPrivateUsage]
+    api._image_generation_queues = {}  # pyright: ignore[reportPrivateUsage]
+    return api
+
+
+def _make_text_gen_task(
+    instance_id: InstanceId, command_id: CommandId
+) -> TextGeneration:
+    return TextGeneration(
+        instance_id=instance_id,
+        command_id=command_id,
+        task_params=TextGenerationTaskParams(
+            model=ModelId("test-model"),
+            input=[InputMessage(role="user", content="hello")],
+        ),
+    )
+
+
+def test_close_streams_for_deleted_instance() -> None:
+    """Deleting an instance closes the text generation sender for commands on that instance."""
+    instance_id = InstanceId("inst-1")
+    command_id = CommandId("cmd-1")
+    task = _make_text_gen_task(instance_id, command_id)
+
+    state = State(tasks={task.task_id: task})
+    api = _make_api_with_state(state)
+
+    sender = MagicMock()
+    api._text_generation_queues[command_id] = sender  # pyright: ignore[reportPrivateUsage]
+
+    api._close_streams_for_instance(instance_id)  # pyright: ignore[reportPrivateUsage]
+
+    sender.close.assert_called_once()
+    assert command_id not in api._text_generation_queues  # pyright: ignore[reportPrivateUsage]
+
+
+def test_close_streams_ignores_unrelated_instances() -> None:
+    """Deleting an instance does NOT close streams for commands on other instances."""
+    target_id = InstanceId("inst-delete")
+    other_id = InstanceId("inst-keep")
+    other_cmd = CommandId("cmd-keep")
+    other_task = _make_text_gen_task(other_id, other_cmd)
+
+    state = State(tasks={other_task.task_id: other_task})
+    api = _make_api_with_state(state)
+
+    sender = MagicMock()
+    api._text_generation_queues[other_cmd] = sender  # pyright: ignore[reportPrivateUsage]
+
+    api._close_streams_for_instance(target_id)  # pyright: ignore[reportPrivateUsage]
+
+    sender.close.assert_not_called()
+    assert other_cmd in api._text_generation_queues  # pyright: ignore[reportPrivateUsage]
+
+
+def test_close_streams_for_deleted_instance_image_generation() -> None:
+    """Deleting an instance closes the image generation sender for commands on that instance."""
+    instance_id = InstanceId("inst-img")
+    command_id = CommandId("cmd-img")
+    task = ImageGeneration(
+        instance_id=instance_id,
+        command_id=command_id,
+        task_params=ImageGenerationTaskParams(prompt="a cat", model="test-model"),
+    )
+
+    state = State(tasks={task.task_id: task})
+    api = _make_api_with_state(state)
+
+    sender = MagicMock()
+    api._image_generation_queues[command_id] = sender  # pyright: ignore[reportPrivateUsage]
+
+    api._close_streams_for_instance(instance_id)  # pyright: ignore[reportPrivateUsage]
+
+    sender.close.assert_called_once()
+    assert command_id not in api._image_generation_queues  # pyright: ignore[reportPrivateUsage]

src/exo/shared/constants.py

@@ -98,4 +98,4 @@ EXO_TRACING_ENABLED = os.getenv("EXO_TRACING_ENABLED", "false").lower() == "true
 
 EXO_MAX_CONCURRENT_REQUESTS = int(os.getenv("EXO_MAX_CONCURRENT_REQUESTS", "8"))
 
-EXO_MAX_INSTANCE_RETRIES = 3
+EXO_MAX_INSTANCE_RETRIES = 5

src/exo/worker/engines/mlx/generator/batch_generate.py

@@ -45,7 +45,6 @@ from exo.worker.engines.mlx.patches.opt_batch_gen import (
     take_ready_topk,
 )
 from exo.worker.engines.mlx.utils_mlx import (
-    detect_thinking_prompt_suffix,
     fix_unmatched_think_end_tokens,
     system_prompt_token_count,
 )
@@ -82,8 +81,6 @@ class _EngineTask:
     potential_stop_sequence_text: str = ""
     completion_tokens: int = 0
     generation_start_time: float = 0.0
-    in_thinking: bool = False
-    reasoning_tokens: int = 0
     prefill_tps: float = 0.0
     media_regions: list[MediaRegion] = field(default_factory=list)
     first_gen_token_time: float | None = None
@@ -273,7 +270,6 @@ class ExoBatchGenerator:
             generation_start_time=time.perf_counter(),
             prefill_tps=_prefill_tps,
             media_regions=media_regions,
-            in_thinking=detect_thinking_prompt_suffix(prompt, self.tokenizer),
         )
 
         return uid
@@ -323,15 +319,6 @@ class ExoBatchGenerator:
             state.generated_text_parts.append(text)
             state.potential_stop_sequence_text += text
 
-            think_start = self.tokenizer.think_start
-            think_end = self.tokenizer.think_end
-            if think_start is not None and text == think_start:
-                state.in_thinking = True
-            elif think_end is not None and text == think_end:
-                state.in_thinking = False
-            if state.in_thinking:
-                state.reasoning_tokens += 1
-
             finish_reason: FinishReason | None = cast(
                 FinishReason | None, response.finish_reason
             )
@@ -402,7 +389,7 @@ class ExoBatchGenerator:
                         cached_tokens=state.prefix_hit_length
                     ),
                     completion_tokens_details=CompletionTokensDetails(
-                        reasoning_tokens=state.reasoning_tokens
+                        reasoning_tokens=0
                     ),
                 )
 

src/exo/worker/engines/mlx/generator/generate.py

@@ -57,7 +57,6 @@ from exo.worker.engines.mlx.constants import (
 )
 from exo.worker.engines.mlx.utils_mlx import (
     apply_chat_template,
-    detect_thinking_prompt_suffix,
     fix_unmatched_think_end_tokens,
     mx_barrier,
     system_prompt_token_count,
@@ -603,11 +602,6 @@ def mlx_generate(
     generated_text_parts: list[str] = []
     generation_start_time = time.perf_counter()
     usage: Usage | None = None
-    in_thinking = detect_thinking_prompt_suffix(prompt, tokenizer)
-    reasoning_tokens = 0
-    think_start = tokenizer.think_start
-    think_end = tokenizer.think_end
-
     logger.info("Starting decode")
     mx_barrier(group)
 
@@ -629,13 +623,6 @@ def mlx_generate(
         generated_text_parts.append(out.text)
         accumulated_text += out.text
 
-        if think_start is not None and out.text == think_start:
-            in_thinking = True
-        elif think_end is not None and out.text == think_end:
-            in_thinking = False
-        if in_thinking:
-            reasoning_tokens += 1
-
         # Check for stop sequences
         text = out.text
         finish_reason: FinishReason | None = cast(
@@ -679,9 +666,7 @@ def mlx_generate(
                 prompt_tokens_details=PromptTokensDetails(
                     cached_tokens=prefix_hit_length
                 ),
-                completion_tokens_details=CompletionTokensDetails(
-                    reasoning_tokens=reasoning_tokens
-                ),
+                completion_tokens_details=CompletionTokensDetails(reasoning_tokens=0),
             )
 
         # Extract logprobs from the full vocabulary logprobs array

src/exo/worker/main.py

@@ -40,6 +40,7 @@ from exo.shared.types.tasks import (
     CreateRunner,
     DownloadModel,
     ImageEdits,
+    LoadModel,
     Shutdown,
     Task,
     TaskStatus,
@@ -351,6 +352,12 @@ class Worker:
                     if cmd_id in self.input_chunk_counts:
                         del self.input_chunk_counts[cmd_id]
                     await self._start_runner_task(modified_task)
+                case LoadModel(instance_id=instance_id):
+                    if (instance := self.state.instances.get(instance_id)) is not None:
+                        model_id = instance.shard_assignments.model_id
+                        self._download_backoff.reset(model_id)
+
+                    await self._start_runner_task(task)
                 case task:
                     await self._start_runner_task(task)
 

src/exo/worker/plan.py

@@ -59,7 +59,7 @@ def plan(
     return (
         _cancel_tasks(runners, tasks)
         or _kill_runner(runners, all_runners, instances)
-        or _create_runner(node_id, runners, instances, instance_backoff)
+        or _create_runner(node_id, runners, all_runners, instances, instance_backoff)
         or _model_needs_download(
             node_id, runners, global_download_status, download_backoff
         )
@@ -79,6 +79,11 @@ def _kill_runner(
         runner_id = runner.bound_instance.bound_runner_id
         if (instance_id := runner.bound_instance.instance.instance_id) not in instances:
             return Shutdown(instance_id=instance_id, runner_id=runner_id)
+        if isinstance(runner.status, RunnerFailed):
+            return Shutdown(
+                instance_id=runner.bound_instance.instance.instance_id,
+                runner_id=runner_id,
+            )
 
         for (
             global_runner_id
@@ -96,13 +101,11 @@ def _kill_runner(
 def _create_runner(
     node_id: NodeId,
     runners: Mapping[RunnerId, RunnerSupervisor],
+    all_runners: Mapping[RunnerId, RunnerStatus],
     instances: Mapping[InstanceId, Instance],
     instance_backoff: KeyedBackoff[InstanceId],
 ) -> CreateRunner | None:
     for instance in instances.values():
-        if not instance_backoff.should_proceed(instance.instance_id):
-            continue
-
         runner_id = instance.shard_assignments.node_to_runner.get(node_id, None)
         if runner_id is None:
             continue
@@ -110,8 +113,18 @@ def _create_runner(
         if runner_id in runners:
             continue
 
-        shard = instance.shard(runner_id)
-        assert shard is not None
+        # don't create runners if any other nodes have runners that have failed - wait for them to fix themselves first.
+        instance_has_failed_runner = any(
+            isinstance(all_runners.get(remote_runner_id), RunnerFailed)
+            for remote_runner_id in instance.shard_assignments.node_to_runner.values()
+            if remote_runner_id != runner_id
+        )
+        we_have_failed_before = isinstance(all_runners.get(runner_id), RunnerFailed)
+        if instance_has_failed_runner and not we_have_failed_before:
+            continue
+
+        if not instance_backoff.should_proceed(instance.instance_id):
+            continue
 
         return CreateRunner(
             instance_id=instance.instance_id,

src/exo/worker/runner/image_models/runner.py

@@ -43,7 +43,6 @@ from exo.shared.types.worker.runner_response import (
 from exo.shared.types.worker.runners import (
     RunnerConnected,
     RunnerConnecting,
-    RunnerFailed,
     RunnerIdle,
     RunnerLoaded,
     RunnerLoading,
@@ -331,9 +330,7 @@ class Runner:
 
     def handle_task(self, task: Task):
         match task:
-            case ConnectToGroup() if isinstance(
-                self.current_status, (RunnerIdle, RunnerFailed)
-            ):
+            case ConnectToGroup() if isinstance(self.current_status, RunnerIdle):
                 logger.info("runner connecting")
                 self.update_status(RunnerConnecting())
                 self.acknowledge_task(task)

src/exo/worker/runner/llm_inference/model_output_parsers.py

@@ -30,6 +30,32 @@ def get_gpt_oss_encoding():
     return encoding
 
 
+def count_reasoning_tokens(
+    responses: Generator[GenerationResponse | ToolCallResponse | None],
+) -> Generator[GenerationResponse | ToolCallResponse | None]:
+    """Count tokens with is_thinking=True and patch the total into Usage on the final response."""
+    reasoning_tokens = 0
+    for response in responses:
+        if response is None:
+            yield None
+            continue
+        if isinstance(response, GenerationResponse) and response.is_thinking:
+            reasoning_tokens += 1
+        if response.usage is not None and reasoning_tokens > 0:
+            response = response.model_copy(
+                update={
+                    "usage": response.usage.model_copy(
+                        update={
+                            "completion_tokens_details": response.usage.completion_tokens_details.model_copy(
+                                update={"reasoning_tokens": reasoning_tokens}
+                            )
+                        }
+                    )
+                }
+            )
+        yield response
+
+
 def apply_all_parsers(
     receiver: Generator[GenerationResponse | None],
     prompt: str,
@@ -41,14 +67,6 @@ def apply_all_parsers(
 ) -> Generator[GenerationResponse | ToolCallResponse | None]:
     mlx_generator = receiver
 
-    if tokenizer.has_thinking:
-        mlx_generator = parse_thinking_models(
-            mlx_generator,
-            tokenizer.think_start,
-            tokenizer.think_end,
-            starts_in_thinking=detect_thinking_prompt_suffix(prompt, tokenizer),
-        )
-
     if issubclass(model_type, GptOssModel):
         mlx_generator = parse_gpt_oss(mlx_generator)
     elif (
@@ -56,10 +74,19 @@ def apply_all_parsers(
         and "deepseek" in model_id.normalize().lower()
     ):
         mlx_generator = parse_deepseek_v32(mlx_generator)
-    elif tool_parser:
-        mlx_generator = parse_tool_calls(mlx_generator, tool_parser, tools)
+    else:
+        if tokenizer.has_thinking:
+            mlx_generator = parse_thinking_models(
+                mlx_generator,
+                tokenizer.think_start,
+                tokenizer.think_end,
+                starts_in_thinking=detect_thinking_prompt_suffix(prompt, tokenizer),
+            )
 
-    return mlx_generator
+        if tool_parser:
+            mlx_generator = parse_tool_calls(mlx_generator, tool_parser, tools)
+
+    return count_reasoning_tokens(mlx_generator)
 
 
 def parse_gpt_oss(
@@ -67,7 +94,6 @@ def parse_gpt_oss(
 ) -> Generator[GenerationResponse | ToolCallResponse | None]:
     encoding = get_gpt_oss_encoding()
     stream = StreamableParser(encoding, role=Role.ASSISTANT)
-    thinking = False
     current_tool_name: str | None = None
     tool_arg_parts: list[str] = []
 
@@ -121,14 +147,10 @@ def parse_gpt_oss(
                 tool_arg_parts = []
             continue
 
-        if ch == "analysis" and not thinking:
-            thinking = True
-
-        if ch != "analysis" and thinking:
-            thinking = False
-
         if delta:
-            yield response.model_copy(update={"text": delta, "is_thinking": thinking})
+            yield response.model_copy(
+                update={"text": delta, "is_thinking": ch == "analysis"}
+            )
 
         if response.finish_reason is not None:
             yield response

src/exo/worker/runner/llm_inference/runner.py

@@ -149,9 +149,7 @@ class Runner:
         self.send_task_status(task.task_id, TaskStatus.Running)
 
         match task:
-            case ConnectToGroup() if isinstance(
-                self.current_status, (RunnerIdle, RunnerFailed)
-            ):
+            case ConnectToGroup() if isinstance(self.current_status, RunnerIdle):
                 assert isinstance(self.generator, Builder)
                 logger.info("runner connecting")
                 self.update_status(RunnerConnecting())

src/exo/worker/runner/runner_supervisor.py

@@ -276,9 +276,7 @@ class RunnerSupervisor:
                 await self._event_sender.send(
                     RunnerStatusUpdated(
                         runner_id=self.bound_instance.bound_runner_id,
-                        runner_status=RunnerFailed(
-                            error_message=f"Terminated ({cause})"
-                        ),
+                        runner_status=self.status,
                     )
                 )
         except (ClosedResourceError, BrokenResourceError):

src/exo/worker/tests/unittests/test_runner/test_finish_reason_sse.py

@@ -1,6 +1,7 @@
 from collections.abc import Generator
 from typing import Any
 
+from exo.api.types import CompletionTokensDetails, PromptTokensDetails, Usage
 from exo.shared.types.worker.runner_response import (
     FinishReason,
     GenerationResponse,
@@ -14,6 +15,7 @@ from exo.worker.engines.mlx.dsml_encoding import (
     TOOL_CALLS_START,
 )
 from exo.worker.runner.llm_inference.model_output_parsers import (
+    count_reasoning_tokens,
     parse_deepseek_v32,
     parse_thinking_models,
     parse_tool_calls,
@@ -243,6 +245,77 @@ class TestThinkingModelsFinishReason:
         )
         assert _got_finish(results)
 
+    def test_reasoning_tokens_counted(self):
+        """reasoning_tokens in Usage reflects the number of thinking tokens."""
+        usage = Usage(
+            prompt_tokens=10,
+            completion_tokens=4,
+            total_tokens=14,
+            prompt_tokens_details=PromptTokensDetails(cached_tokens=0),
+            completion_tokens_details=CompletionTokensDetails(reasoning_tokens=0),
+        )
+        tokens = [
+            _make_response("<think>", 0),
+            _make_response("let me", 1),
+            _make_response(" think", 2),
+            _make_response("</think>", 3),
+            GenerationResponse(text="42", token=4, finish_reason="stop", usage=usage),
+        ]
+        results = _step_until_finish(
+            count_reasoning_tokens(
+                parse_thinking_models(
+                    _queue_source(tokens),
+                    think_start="<think>",
+                    think_end="</think>",
+                    starts_in_thinking=False,
+                )
+            )
+        )
+        final = [
+            r
+            for r in results
+            if isinstance(r, GenerationResponse) and r.finish_reason is not None
+        ]
+        assert len(final) == 1
+        assert final[0].usage is not None
+        assert final[0].usage.completion_tokens_details.reasoning_tokens == 2
+
+    def test_reasoning_tokens_starts_in_thinking(self):
+        """reasoning_tokens counts correctly when starts_in_thinking=True."""
+        usage = Usage(
+            prompt_tokens=10,
+            completion_tokens=3,
+            total_tokens=13,
+            prompt_tokens_details=PromptTokensDetails(cached_tokens=0),
+            completion_tokens_details=CompletionTokensDetails(reasoning_tokens=0),
+        )
+        tokens = [
+            _make_response("hmm", 0),
+            _make_response("ok", 1),
+            _make_response("</think>", 2),
+            GenerationResponse(
+                text="answer", token=3, finish_reason="stop", usage=usage
+            ),
+        ]
+        results = _step_until_finish(
+            count_reasoning_tokens(
+                parse_thinking_models(
+                    _queue_source(tokens),
+                    think_start="<think>",
+                    think_end="</think>",
+                    starts_in_thinking=True,
+                )
+            )
+        )
+        final = [
+            r
+            for r in results
+            if isinstance(r, GenerationResponse) and r.finish_reason is not None
+        ]
+        assert len(final) == 1
+        assert final[0].usage is not None
+        assert final[0].usage.completion_tokens_details.reasoning_tokens == 2
+
 
 # ── parse_tool_calls (generic) ──────────────────────────────────
 

src/exo/worker/tests/unittests/test_runner/test_parse_gpt_oss.py

@@ -1,11 +1,19 @@
 from collections.abc import Generator
 
-from exo.api.types import FinishReason
+from exo.api.types import (
+    CompletionTokensDetails,
+    FinishReason,
+    PromptTokensDetails,
+    Usage,
+)
 from exo.shared.types.worker.runner_response import (
     GenerationResponse,
     ToolCallResponse,
 )
-from exo.worker.runner.llm_inference.model_output_parsers import parse_gpt_oss
+from exo.worker.runner.llm_inference.model_output_parsers import (
+    count_reasoning_tokens,
+    parse_gpt_oss,
+)
 
 # Token IDs from mlx-community/gpt-oss-20b-MXFP4-Q8 tokenizer.
 # These are stable since they come from the model's vocabulary.
@@ -85,6 +93,7 @@ THINKING_THEN_TOOL_TOKENS: list[tuple[int, str]] = [
 def _make_gen_responses(
     tokens: list[tuple[int, str]],
     last_finish_reason: FinishReason = "stop",
+    last_usage: Usage | None = None,
 ) -> list[GenerationResponse]:
     """Build GenerationResponse list from (token_id, text) pairs."""
     responses: list[GenerationResponse] = []
@@ -95,7 +104,7 @@ def _make_gen_responses(
                 text=text,
                 token=tid,
                 finish_reason=last_finish_reason if is_last else None,
-                usage=None,
+                usage=last_usage if is_last else None,
             )
         )
     return responses
@@ -251,3 +260,43 @@ class TestParseGptOssMaxTokensTruncation:
         # due to Harmony encoding, so we just check something was emitted)
         all_text = "".join(r.text for r in gen_responses)
         assert len(all_text) > 0
+
+
+class TestGptOssReasoningTokensCounted:
+    """count_reasoning_tokens must patch Usage when parse_gpt_oss emits thinking tokens."""
+
+    def test_thinking_then_text_counts_reasoning_tokens(self):
+        usage = Usage(
+            prompt_tokens=10,
+            completion_tokens=len(PLAIN_TEXT_TOKENS),
+            total_tokens=10 + len(PLAIN_TEXT_TOKENS),
+            prompt_tokens_details=PromptTokensDetails(cached_tokens=0),
+            completion_tokens_details=CompletionTokensDetails(reasoning_tokens=0),
+        )
+        responses = _make_gen_responses(PLAIN_TEXT_TOKENS, last_usage=usage)
+
+        def _gen() -> Generator[GenerationResponse, None, None]:
+            yield from responses
+
+        results = list(
+            x for x in count_reasoning_tokens(parse_gpt_oss(_gen())) if x is not None
+        )
+
+        # Verify thinking tokens were detected
+        thinking = [
+            r for r in results if isinstance(r, GenerationResponse) and r.is_thinking
+        ]
+        assert len(thinking) > 0
+
+        # Verify reasoning_tokens is patched on responses that carry Usage
+        with_usage = [
+            r
+            for r in results
+            if isinstance(r, GenerationResponse) and r.usage is not None
+        ]
+        assert len(with_usage) > 0
+        assert all(
+            r.usage is not None
+            and r.usage.completion_tokens_details.reasoning_tokens == len(thinking)
+            for r in with_usage
+        )