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

@@ -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/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)
 

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/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):