Cancel SSE keep-alive when instance is deleted

src/exo/api/main.py

@@ -169,10 +169,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
@@ -1814,9 +1824,25 @@ class API:
                             await queue.send(event.chunk)
                         except BrokenResourceError:
                             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

@@ -97,5 +97,3 @@ EXO_OFFLINE = os.getenv("EXO_OFFLINE", "false").lower() == "true"
 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

src/exo/utils/keyed_backoff.py

@@ -1,9 +1,10 @@
 import time
-from typing import final
+from typing import Generic, TypeVar
+
+K = TypeVar("K")
 
 
-@final
-class KeyedBackoff[K]:
+class KeyedBackoff(Generic[K]):
     """Tracks exponential backoff state per key."""
 
     def __init__(self, base: float = 0.5, cap: float = 10.0):
@@ -25,10 +26,6 @@ class KeyedBackoff[K]:
         self._last_time[key] = time.monotonic()
         self._attempts[key] = self._attempts.get(key, 0) + 1
 
-    def attempts(self, key: K) -> int:
-        """Return the number of recorded attempts for a key."""
-        return self._attempts.get(key, 0)
-
     def reset(self, key: K) -> None:
         """Reset backoff state for a key (e.g., on success)."""
         self._attempts.pop(key, None)

src/exo/worker/main.py

@@ -9,11 +9,9 @@ from loguru import logger
 from exo.api.types import ImageEditsTaskParams
 from exo.download.download_utils import is_read_only_model_dir, resolve_existing_model
 from exo.shared.apply import apply
-from exo.shared.constants import EXO_MAX_INSTANCE_RETRIES
 from exo.shared.models.model_cards import ModelId, add_to_card_cache, delete_custom_card
 from exo.shared.types.chunks import InputImageChunk
 from exo.shared.types.commands import (
-    DeleteInstance,
     ForwarderCommand,
     ForwarderDownloadCommand,
     StartDownload,
@@ -25,7 +23,6 @@ from exo.shared.types.events import (
     Event,
     IndexedEvent,
     InputChunkReceived,
-    InstanceDeleted,
     NodeDownloadProgress,
     NodeGatheredInfo,
     TaskCreated,
@@ -47,7 +44,6 @@ from exo.shared.types.tasks import (
 )
 from exo.shared.types.topology import Connection, SocketConnection
 from exo.shared.types.worker.downloads import DownloadCompleted
-from exo.shared.types.worker.instances import InstanceId
 from exo.shared.types.worker.runners import RunnerId
 from exo.utils.channels import Receiver, Sender, channel
 from exo.utils.info_gatherer.info_gatherer import GatheredInfo, InfoGatherer
@@ -88,9 +84,6 @@ class Worker:
         self.image_cache: dict[str, str] = {}
 
         self._download_backoff: KeyedBackoff[ModelId] = KeyedBackoff(base=0.5, cap=10.0)
-        self._instance_backoff: KeyedBackoff[InstanceId] = KeyedBackoff(
-            base=0.5, cap=10.0
-        )
         self._stopped: anyio.Event = anyio.Event()
 
     async def run(self):
@@ -134,9 +127,6 @@ class Worker:
                 self.state = apply(self.state, event=event)
                 event = event.event
 
-                if isinstance(event, InstanceDeleted):
-                    self._instance_backoff.reset(event.instance_id)
-
                 # Buffer input image chunks for image editing
                 if isinstance(event, InputChunkReceived):
                     cmd_id = event.command_id
@@ -166,24 +156,15 @@ class Worker:
                 self.state.runners,
                 self.state.tasks,
                 self.input_chunk_buffer,
-                self._instance_backoff,
-                self._download_backoff,
             )
             if task is None:
                 continue
 
-            if isinstance(task, CreateRunner):
-                iid = task.instance_id
-                if self._instance_backoff.attempts(iid) >= EXO_MAX_INSTANCE_RETRIES:
-                    logger.warning(
-                        f"Instance {iid} exceeded {EXO_MAX_INSTANCE_RETRIES} retries, requesting deletion"
-                    )
-                    await self.command_sender.send(
-                        ForwarderCommand(
-                            origin=self._system_id,
-                            command=DeleteInstance(instance_id=iid),
-                        )
-                    )
+            # Gate DownloadModel on backoff BEFORE emitting TaskCreated
+            # to prevent flooding the event log with useless events
+            if isinstance(task, DownloadModel):
+                model_id = task.shard_metadata.model_card.model_id
+                if not self._download_backoff.should_proceed(model_id):
                     continue
 
             logger.info(f"Worker plan: {task.__class__.__name__}")
@@ -194,7 +175,6 @@ class Worker:
             match task:
                 case CreateRunner():
                     self._create_supervisor(task)
-                    self._instance_backoff.record_attempt(task.instance_id)
                     await self.event_sender.send(
                         TaskStatusUpdated(
                             task_id=task.task_id, task_status=TaskStatus.Complete

src/exo/worker/plan.py

@@ -3,7 +3,7 @@
 from collections.abc import Mapping, Sequence
 
 from exo.shared.types.chunks import InputImageChunk
-from exo.shared.types.common import CommandId, ModelId, NodeId
+from exo.shared.types.common import CommandId, NodeId
 from exo.shared.types.tasks import (
     CancelTask,
     ConnectToGroup,
@@ -39,7 +39,6 @@ from exo.shared.types.worker.runners import (
     RunnerStatus,
     RunnerWarmingUp,
 )
-from exo.utils.keyed_backoff import KeyedBackoff
 from exo.worker.runner.runner_supervisor import RunnerSupervisor
 
 
@@ -51,22 +50,18 @@ def plan(
     instances: Mapping[InstanceId, Instance],
     all_runners: Mapping[RunnerId, RunnerStatus],  # all global
     tasks: Mapping[TaskId, Task],
-    input_chunk_buffer: Mapping[CommandId, Mapping[int, InputImageChunk]],
-    instance_backoff: KeyedBackoff[InstanceId],
-    download_backoff: KeyedBackoff[ModelId],
+    input_chunk_buffer: Mapping[CommandId, Mapping[int, InputImageChunk]] | None = None,
 ) -> Task | None:
     # Python short circuiting OR logic should evaluate these sequentially.
     return (
         _cancel_tasks(runners, tasks)
         or _kill_runner(runners, all_runners, instances)
-        or _create_runner(node_id, runners, instances, instance_backoff)
-        or _model_needs_download(
-            node_id, runners, global_download_status, download_backoff
-        )
+        or _create_runner(node_id, runners, instances)
+        or _model_needs_download(node_id, runners, global_download_status)
         or _init_distributed_backend(runners, all_runners)
         or _load_model(runners, all_runners, global_download_status)
         or _ready_to_warmup(runners, all_runners)
-        or _pending_tasks(runners, tasks, all_runners, input_chunk_buffer)
+        or _pending_tasks(runners, tasks, all_runners, input_chunk_buffer or {})
     )
 
 
@@ -97,12 +92,8 @@ def _create_runner(
     node_id: NodeId,
     runners: Mapping[RunnerId, RunnerSupervisor],
     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
@@ -125,7 +116,6 @@ def _model_needs_download(
     node_id: NodeId,
     runners: Mapping[RunnerId, RunnerSupervisor],
     global_download_status: Mapping[NodeId, Sequence[DownloadProgress]],
-    download_backoff: KeyedBackoff[ModelId],
 ) -> DownloadModel | None:
     local_downloads = global_download_status.get(node_id, [])
     download_status = {
@@ -134,16 +124,12 @@ def _model_needs_download(
 
     for runner in runners.values():
         model_id = runner.bound_instance.bound_shard.model_card.model_id
-        if (
-            isinstance(runner.status, RunnerIdle)
-            and (
-                model_id not in download_status
-                or not isinstance(
-                    download_status[model_id],
-                    (DownloadOngoing, DownloadCompleted, DownloadFailed),
-                )
+        if isinstance(runner.status, RunnerIdle) and (
+            model_id not in download_status
+            or not isinstance(
+                download_status[model_id],
+                (DownloadOngoing, DownloadCompleted, DownloadFailed),
             )
-            and download_backoff.should_proceed(model_id)
         ):
             # We don't invalidate download_status randomly in case a file gets deleted on disk
             return DownloadModel(
@@ -286,7 +272,7 @@ def _pending_tasks(
     runners: Mapping[RunnerId, RunnerSupervisor],
     tasks: Mapping[TaskId, Task],
     all_runners: Mapping[RunnerId, RunnerStatus],
-    input_chunk_buffer: Mapping[CommandId, Mapping[int, InputImageChunk]],
+    input_chunk_buffer: Mapping[CommandId, Mapping[int, InputImageChunk]] | None,
 ) -> Task | None:
     for task in tasks.values():
         # for now, just forward chat completions
@@ -301,6 +287,7 @@ def _pending_tasks(
         if isinstance(task, (ImageEdits, TextGeneration)):
             expected_image_chunks = task.task_params.total_input_chunks
         if expected_image_chunks > 0:
+            assert input_chunk_buffer is not None
             cmd_id = task.command_id
             received = len(input_chunk_buffer.get(cmd_id, {}))
             if received < expected_image_chunks:

src/exo/worker/tests/unittests/test_plan/test_download_and_loading.py

@@ -8,7 +8,6 @@ from exo.shared.types.worker.runners import (
     RunnerConnected,
     RunnerIdle,
 )
-from exo.utils.keyed_backoff import KeyedBackoff
 from exo.worker.tests.constants import (
     INSTANCE_1_ID,
     MODEL_A_ID,
@@ -53,9 +52,6 @@ def test_plan_requests_download_when_waiting_and_shard_not_downloaded():
         instances=instances,
         all_runners=all_runners,
         tasks={},
-        input_chunk_buffer={},
-        instance_backoff=KeyedBackoff(),
-        download_backoff=KeyedBackoff(),
     )
 
     assert isinstance(result, plan_mod.DownloadModel)
@@ -108,9 +104,6 @@ def test_plan_loads_model_when_all_shards_downloaded_and_waiting():
         instances=instances,
         all_runners=all_runners,
         tasks={},
-        input_chunk_buffer={},
-        instance_backoff=KeyedBackoff(),
-        download_backoff=KeyedBackoff(),
     )
 
     assert isinstance(result, LoadModel)
@@ -153,9 +146,6 @@ def test_plan_does_not_request_download_when_shard_already_downloaded():
         instances=instances,
         all_runners=all_runners,
         tasks={},
-        input_chunk_buffer={},
-        instance_backoff=KeyedBackoff(),
-        download_backoff=KeyedBackoff(),
     )
 
     assert not isinstance(result, plan_mod.DownloadModel)
@@ -203,9 +193,6 @@ def test_plan_does_not_load_model_until_all_shards_downloaded_globally():
         instances=instances,
         all_runners=all_runners,
         tasks={},
-        input_chunk_buffer={},
-        instance_backoff=KeyedBackoff(),
-        download_backoff=KeyedBackoff(),
     )
 
     assert result is None
@@ -226,9 +213,6 @@ def test_plan_does_not_load_model_until_all_shards_downloaded_globally():
         instances=instances,
         all_runners=all_runners,
         tasks={},
-        input_chunk_buffer={},
-        instance_backoff=KeyedBackoff(),
-        download_backoff=KeyedBackoff(),
     )
 
     assert result is not None

src/exo/worker/tests/unittests/test_plan/test_runner_lifecycle.py

@@ -9,7 +9,6 @@ from exo.shared.types.worker.runners import (
     RunnerReady,
     RunnerStatus,
 )
-from exo.utils.keyed_backoff import KeyedBackoff
 from exo.worker.tests.constants import (
     INSTANCE_1_ID,
     MODEL_A_ID,
@@ -53,9 +52,6 @@ def test_plan_kills_runner_when_instance_missing():
         instances=instances,
         all_runners=all_runners,
         tasks={},
-        input_chunk_buffer={},
-        instance_backoff=KeyedBackoff(),
-        download_backoff=KeyedBackoff(),
     )
 
     assert isinstance(result, Shutdown)
@@ -95,9 +91,6 @@ def test_plan_kills_runner_when_sibling_failed():
         instances=instances,
         all_runners=all_runners,
         tasks={},
-        input_chunk_buffer={},
-        instance_backoff=KeyedBackoff(),
-        download_backoff=KeyedBackoff(),
     )
 
     assert isinstance(result, Shutdown)
@@ -129,9 +122,6 @@ def test_plan_creates_runner_when_missing_for_node():
         instances=instances,
         all_runners=all_runners,
         tasks={},
-        input_chunk_buffer={},
-        instance_backoff=KeyedBackoff(),
-        download_backoff=KeyedBackoff(),
     )
 
     # We patched plan_mod.CreateRunner → CreateRunner
@@ -170,9 +160,6 @@ def test_plan_does_not_create_runner_when_supervisor_already_present():
         instances=instances,
         all_runners=all_runners,
         tasks={},
-        input_chunk_buffer={},
-        instance_backoff=KeyedBackoff(),
-        download_backoff=KeyedBackoff(),
     )
 
     assert result is None
@@ -202,9 +189,6 @@ def test_plan_does_not_create_runner_for_unassigned_node():
         instances=instances,
         all_runners=all_runners,
         tasks={},
-        input_chunk_buffer={},
-        instance_backoff=KeyedBackoff(),
-        download_backoff=KeyedBackoff(),
     )
 
     assert result is None

src/exo/worker/tests/unittests/test_plan/test_task_forwarding.py

@@ -9,7 +9,6 @@ from exo.shared.types.worker.runners import (
     RunnerReady,
     RunnerRunning,
 )
-from exo.utils.keyed_backoff import KeyedBackoff
 from exo.worker.tests.constants import (
     COMMAND_1_ID,
     INSTANCE_1_ID,
@@ -72,9 +71,6 @@ def test_plan_forwards_pending_chat_completion_when_runner_ready():
         instances=instances,
         all_runners=all_runners,
         tasks={TASK_1_ID: task},
-        input_chunk_buffer={},
-        instance_backoff=KeyedBackoff(),
-        download_backoff=KeyedBackoff(),
     )
 
     assert result is task
@@ -124,9 +120,6 @@ def test_plan_does_not_forward_chat_completion_if_any_runner_not_ready():
         instances=instances,
         all_runners=all_runners,
         tasks={TASK_1_ID: task},
-        input_chunk_buffer={},
-        instance_backoff=KeyedBackoff(),
-        download_backoff=KeyedBackoff(),
     )
 
     assert result is None
@@ -173,9 +166,6 @@ def test_plan_does_not_forward_tasks_for_other_instances():
         instances=instances,
         all_runners=all_runners,
         tasks={foreign_task.task_id: foreign_task},
-        input_chunk_buffer={},
-        instance_backoff=KeyedBackoff(),
-        download_backoff=KeyedBackoff(),
     )
 
     assert result is None
@@ -240,9 +230,6 @@ def test_plan_ignores_non_pending_or_non_chat_tasks():
         instances=instances,
         all_runners=all_runners,
         tasks={TASK_1_ID: completed_task, other_task_id: other_task},
-        input_chunk_buffer={},
-        instance_backoff=KeyedBackoff(),
-        download_backoff=KeyedBackoff(),
     )
 
     assert result is None
@@ -282,9 +269,6 @@ def test_plan_returns_none_when_nothing_to_do():
         instances=instances,
         all_runners=all_runners,
         tasks={},
-        input_chunk_buffer={},
-        instance_backoff=KeyedBackoff(),
-        download_backoff=KeyedBackoff(),
     )
 
     assert result is None

src/exo/worker/tests/unittests/test_plan/test_warmup.py

@@ -7,7 +7,6 @@ from exo.shared.types.worker.runners import (
     RunnerLoading,
     RunnerWarmingUp,
 )
-from exo.utils.keyed_backoff import KeyedBackoff
 from exo.worker.tests.constants import (
     INSTANCE_1_ID,
     MODEL_A_ID,
@@ -62,9 +61,6 @@ def test_plan_starts_warmup_for_accepting_rank_when_all_loaded_or_warming():
         instances=instances,
         all_runners=all_runners,
         tasks={},
-        input_chunk_buffer={},
-        instance_backoff=KeyedBackoff(),
-        download_backoff=KeyedBackoff(),
     )
 
     assert isinstance(result, StartWarmup)
@@ -106,9 +102,6 @@ def test_plan_starts_warmup_for_rank_zero_after_others_warming():
         instances=instances,
         all_runners=all_runners,
         tasks={},
-        input_chunk_buffer={},
-        instance_backoff=KeyedBackoff(),
-        download_backoff=KeyedBackoff(),
     )
 
     assert isinstance(result, StartWarmup)
@@ -149,9 +142,6 @@ def test_plan_does_not_start_warmup_for_non_zero_rank_until_all_loaded_or_warmin
         instances=instances,
         all_runners=all_runners,
         tasks={},
-        input_chunk_buffer={},
-        instance_backoff=KeyedBackoff(),
-        download_backoff=KeyedBackoff(),
     )
 
     assert result is None
@@ -196,9 +186,6 @@ def test_plan_does_not_start_warmup_for_rank_zero_until_others_warming():
         instances=instances,
         all_runners=all_runners,
         tasks={},
-        input_chunk_buffer={},
-        instance_backoff=KeyedBackoff(),
-        download_backoff=KeyedBackoff(),
     )
 
     assert result is None
@@ -215,9 +202,6 @@ def test_plan_does_not_start_warmup_for_rank_zero_until_others_warming():
         instances=instances,
         all_runners=all_runners,
         tasks={},
-        input_chunk_buffer={},
-        instance_backoff=KeyedBackoff(),
-        download_backoff=KeyedBackoff(),
     )
 
     assert isinstance(result, StartWarmup)
@@ -261,9 +245,6 @@ def test_plan_starts_warmup_for_connecting_rank_after_others_warming():
         instances=instances,
         all_runners=all_runners,
         tasks={},
-        input_chunk_buffer={},
-        instance_backoff=KeyedBackoff(),
-        download_backoff=KeyedBackoff(),
     )
 
     assert isinstance(result, StartWarmup)
@@ -306,9 +287,6 @@ def test_plan_does_not_start_warmup_for_accepting_rank_until_all_loaded_or_warmi
         instances=instances,
         all_runners=all_runners,
         tasks={},
-        input_chunk_buffer={},
-        instance_backoff=KeyedBackoff(),
-        download_backoff=KeyedBackoff(),
     )
 
     assert result is None
@@ -350,9 +328,6 @@ def test_plan_does_not_start_warmup_for_connecting_rank_until_others_warming():
         instances=instances,
         all_runners=all_runners,
         tasks={},
-        input_chunk_buffer={},
-        instance_backoff=KeyedBackoff(),
-        download_backoff=KeyedBackoff(),
     )
 
     assert result is None