runner failling should imply the runner is dead

this simplifies some restart logic that caused the runners to loop.
we only create runners if we failed or the instance has no failed
runners (the runners were revived successfully)

src/exo/shared/constants.py

@@ -97,3 +97,5 @@ 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 = 5

src/exo/utils/keyed_backoff.py

@@ -1,10 +1,9 @@
 import time
-from typing import Generic, TypeVar
-
-K = TypeVar("K")
+from typing import final
 
 
-class KeyedBackoff(Generic[K]):
+@final
+class KeyedBackoff[K]:
     """Tracks exponential backoff state per key."""
 
     def __init__(self, base: float = 0.5, cap: float = 10.0):
@@ -26,6 +25,10 @@ class KeyedBackoff(Generic[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,9 +9,11 @@ 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,
@@ -23,6 +25,7 @@ from exo.shared.types.events import (
     Event,
     IndexedEvent,
     InputChunkReceived,
+    InstanceDeleted,
     NodeDownloadProgress,
     NodeGatheredInfo,
     TaskCreated,
@@ -37,6 +40,7 @@ from exo.shared.types.tasks import (
     CreateRunner,
     DownloadModel,
     ImageEdits,
+    LoadModel,
     Shutdown,
     Task,
     TaskStatus,
@@ -44,6 +48,7 @@ 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
@@ -84,6 +89,9 @@ 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):
@@ -127,6 +135,9 @@ 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
@@ -156,15 +167,24 @@ class Worker:
                 self.state.runners,
                 self.state.tasks,
                 self.input_chunk_buffer,
+                self._instance_backoff,
+                self._download_backoff,
             )
             if task is None:
                 continue
 
-            # 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):
+            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),
+                        )
+                    )
                     continue
 
             logger.info(f"Worker plan: {task.__class__.__name__}")
@@ -175,6 +195,7 @@ 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
@@ -328,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)
 

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, NodeId
+from exo.shared.types.common import CommandId, ModelId, NodeId
 from exo.shared.types.tasks import (
     CancelTask,
     ConnectToGroup,
@@ -39,6 +39,7 @@ from exo.shared.types.worker.runners import (
     RunnerStatus,
     RunnerWarmingUp,
 )
+from exo.utils.keyed_backoff import KeyedBackoff
 from exo.worker.runner.runner_supervisor import RunnerSupervisor
 
 
@@ -50,18 +51,22 @@ 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]] | None = None,
+    input_chunk_buffer: Mapping[CommandId, Mapping[int, InputImageChunk]],
+    instance_backoff: KeyedBackoff[InstanceId],
+    download_backoff: KeyedBackoff[ModelId],
 ) -> 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)
-        or _model_needs_download(node_id, runners, global_download_status)
+        or _create_runner(node_id, runners, all_runners, instances, instance_backoff)
+        or _model_needs_download(
+            node_id, runners, global_download_status, download_backoff
+        )
         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 {})
+        or _pending_tasks(runners, tasks, all_runners, input_chunk_buffer)
     )
 
 
@@ -74,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
@@ -91,7 +101,9 @@ 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():
         runner_id = instance.shard_assignments.node_to_runner.get(node_id, None)
@@ -101,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,
@@ -116,6 +138,7 @@ 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 = {
@@ -124,12 +147,16 @@ 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(
@@ -272,7 +299,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]] | None,
+    input_chunk_buffer: Mapping[CommandId, Mapping[int, InputImageChunk]],
 ) -> Task | None:
     for task in tasks.values():
         # for now, just forward chat completions
@@ -287,7 +314,6 @@ 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/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/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_plan/test_download_and_loading.py

@@ -8,6 +8,7 @@ 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,
@@ -52,6 +53,9 @@ 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)
@@ -104,6 +108,9 @@ 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)
@@ -146,6 +153,9 @@ 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)
@@ -193,6 +203,9 @@ 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
@@ -213,6 +226,9 @@ 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,6 +9,7 @@ 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,
@@ -52,6 +53,9 @@ 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)
@@ -91,6 +95,9 @@ 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)
@@ -122,6 +129,9 @@ 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
@@ -160,6 +170,9 @@ 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
@@ -189,6 +202,9 @@ 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,6 +9,7 @@ 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,
@@ -71,6 +72,9 @@ 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
@@ -120,6 +124,9 @@ 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
@@ -166,6 +173,9 @@ 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
@@ -230,6 +240,9 @@ 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
@@ -269,6 +282,9 @@ 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,6 +7,7 @@ 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,
@@ -61,6 +62,9 @@ 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)
@@ -102,6 +106,9 @@ 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)
@@ -142,6 +149,9 @@ 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
@@ -186,6 +196,9 @@ 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
@@ -202,6 +215,9 @@ 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)
@@ -245,6 +261,9 @@ 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)
@@ -287,6 +306,9 @@ 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
@@ -328,6 +350,9 @@ 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