From b12cd1b1862fe368f5ee41c6b0df4a0a5af40c73 Mon Sep 17 00:00:00 2001
From: ciaranbor <81697641+ciaranbor@users.noreply.github.com>
Date: Wed, 8 Apr 2026 16:14:28 +0100
Subject: [PATCH 1/3] Cancel SSE keep-alive when instance is deleted (#1828)
## Motivation
When a model instance is deleted (e.g. node disconnect, manual
teardown), any in-flight SSE streaming connections for that instance
hang indefinitely. The API never closes the response stream, so clients
block forever waiting for more chunks.
## Changes
- Listen for `InstanceDeleted` events in the API event loop
- Add `_close_streams_for_instance()` to find and close any active
text/image generation queues tied to tasks on the deleted instance
- Add unit tests covering text gen, image gen, and
unrelated-instance-not-closed scenarios
## Why It Works
When an instance is deleted, we iterate `state.tasks` to find commands
running on that instance, then close and remove their send-side queue
handles. This causes the SSE generator to terminate, unblocking the
client.
## Test Plan
### Manual Testing
- This was causing issues for me on another branch (integration tests).
Including this fix solved the issue
### Automated Testing
- `test_instance_deleted_stream_cleanup.py`: 3 tests covering text gen
cleanup, image gen cleanup, and ensuring unrelated streams are not
affected
---
src/exo/api/main.py | 26 ++++++
.../test_instance_deleted_stream_cleanup.py | 93 +++++++++++++++++++
2 files changed, 119 insertions(+)
create mode 100644 src/exo/api/tests/test_instance_deleted_stream_cleanup.py
diff --git a/src/exo/api/main.py b/src/exo/api/main.py
index 3c183f4f7..73488ff07 100644
--- a/src/exo/api/main.py
+++ b/src/exo/api/main.py
@@ -172,10 +172,20 @@ 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.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.text_generation import TextGenerationTaskParams
from exo.shared.types.worker.downloads import DownloadCompleted
from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
@@ -1818,9 +1828,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(
diff --git a/src/exo/api/tests/test_instance_deleted_stream_cleanup.py b/src/exo/api/tests/test_instance_deleted_stream_cleanup.py
new file mode 100644
index 000000000..843554b9f
--- /dev/null
+++ b/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]
From e2e17eafb73f64cfe8088c2007d272a72fd0835e Mon Sep 17 00:00:00 2001
From: ciaranbor <81697641+ciaranbor@users.noreply.github.com>
Date: Thu, 9 Apr 2026 12:29:35 +0100
Subject: [PATCH 2/3] Fix reasoning_tokens counting for multi-token thinking
tag models (#1848)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
## Motivation
`reasoning_tokens` is always 0 in usage stats, even when thinking
content streams correctly via `reasoning_content` SSE deltas. The MLX
generators had their own thinking detection comparing individual
detokenized tokens against think tags — this never fires for models
where tags span multiple tokens (e.g. gpt-oss-120b) or are already in
the prompt.
## Changes
- Removed broken per-token thinking detection from `batch_generate.py`
and `generate.py`
- Added `_count_reasoning_tokens` wrapper in `model_output_parsers.py`
that counts `is_thinking=True` responses and patches the total into
Usage on the final response
- Wired it as the outermost stage of `apply_all_parsers`, so it works
regardless of which parser sets `is_thinking`
- Added 3 tests covering `parse_thinking_models` and `parse_gpt_oss`
paths
## Why It Works
The parser pipeline already correctly sets `is_thinking` on each
response. Counting at the output of `apply_all_parsers` means one
counting point that works for all model types, replacing the duplicate
broken logic in two generators.
## Test Plan
### Manual Testing
- 4-node cluster, `mlx-community/gpt-oss-120b-MXFP4-Q8`
- Main branch: `reasoning_tokens: 0` — fix branch: `reasoning_tokens:
25`
### Automated Testing
- 3 new tests: explicit think tags, `starts_in_thinking=True`, and
gpt-oss Harmony analysis channel
---
.../engines/mlx/generator/batch_generate.py | 15 +---
.../worker/engines/mlx/generator/generate.py | 17 +----
.../llm_inference/model_output_parsers.py | 60 ++++++++++-----
.../test_runner/test_finish_reason_sse.py | 73 +++++++++++++++++++
.../test_runner/test_parse_gpt_oss.py | 55 +++++++++++++-
5 files changed, 168 insertions(+), 52 deletions(-)
diff --git a/src/exo/worker/engines/mlx/generator/batch_generate.py b/src/exo/worker/engines/mlx/generator/batch_generate.py
index 42e02eea5..1d47ad40c 100644
--- a/src/exo/worker/engines/mlx/generator/batch_generate.py
+++ b/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
),
)
diff --git a/src/exo/worker/engines/mlx/generator/generate.py b/src/exo/worker/engines/mlx/generator/generate.py
index eef11101c..43a38bca7 100644
--- a/src/exo/worker/engines/mlx/generator/generate.py
+++ b/src/exo/worker/engines/mlx/generator/generate.py
@@ -53,7 +53,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,
@@ -597,11 +596,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)
@@ -623,13 +617,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(
@@ -673,9 +660,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
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 1242909a3..9f380a770 100644
--- a/src/exo/worker/runner/llm_inference/model_output_parsers.py
+++ b/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
diff --git a/src/exo/worker/tests/unittests/test_runner/test_finish_reason_sse.py b/src/exo/worker/tests/unittests/test_runner/test_finish_reason_sse.py
index 907ccddb6..0aee03ba6 100644
--- a/src/exo/worker/tests/unittests/test_runner/test_finish_reason_sse.py
+++ b/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("", 0),
+ _make_response("let me", 1),
+ _make_response(" think", 2),
+ _make_response("", 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_end="",
+ 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("", 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_end="",
+ 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) ──────────────────────────────────
diff --git a/src/exo/worker/tests/unittests/test_runner/test_parse_gpt_oss.py b/src/exo/worker/tests/unittests/test_runner/test_parse_gpt_oss.py
index ba3313361..28bbb7c84 100644
--- a/src/exo/worker/tests/unittests/test_runner/test_parse_gpt_oss.py
+++ b/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
+ )
From f2e6b1ef763a8cadeb68b59d3cce5baba4425b1c Mon Sep 17 00:00:00 2001
From: Evan Quiney
Date: Thu, 9 Apr 2026 12:34:35 +0100
Subject: [PATCH 3/3] prevent some crash loops (#1827)
extension to #1763 that prevents crash looping in some common scenarios.
---
src/exo/shared/constants.py | 2 +-
src/exo/worker/main.py | 7 ++++++
src/exo/worker/plan.py | 25 ++++++++++++++-----
src/exo/worker/runner/image_models/runner.py | 5 +---
src/exo/worker/runner/llm_inference/runner.py | 4 +--
src/exo/worker/runner/runner_supervisor.py | 4 +--
6 files changed, 30 insertions(+), 17 deletions(-)
diff --git a/src/exo/shared/constants.py b/src/exo/shared/constants.py
index 12bc16344..01f315bc3 100644
--- a/src/exo/shared/constants.py
+++ b/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
diff --git a/src/exo/worker/main.py b/src/exo/worker/main.py
index 08cc1e0ec..e3da3bdb9 100644
--- a/src/exo/worker/main.py
+++ b/src/exo/worker/main.py
@@ -40,6 +40,7 @@ from exo.shared.types.tasks import (
CreateRunner,
DownloadModel,
ImageEdits,
+ LoadModel,
Shutdown,
Task,
TaskStatus,
@@ -348,6 +349,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)
diff --git a/src/exo/worker/plan.py b/src/exo/worker/plan.py
index 713481185..07aeeab8d 100644
--- a/src/exo/worker/plan.py
+++ b/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,
diff --git a/src/exo/worker/runner/image_models/runner.py b/src/exo/worker/runner/image_models/runner.py
index 4fdecec2a..2eb90baec 100644
--- a/src/exo/worker/runner/image_models/runner.py
+++ b/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)
diff --git a/src/exo/worker/runner/llm_inference/runner.py b/src/exo/worker/runner/llm_inference/runner.py
index 5c19ec733..53fc2f197 100644
--- a/src/exo/worker/runner/llm_inference/runner.py
+++ b/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())
diff --git a/src/exo/worker/runner/runner_supervisor.py b/src/exo/worker/runner/runner_supervisor.py
index c4b4dc5af..b55c0bd45 100644
--- a/src/exo/worker/runner/runner_supervisor.py
+++ b/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):