fix: only shutdown runners that error unrecoverably

previously we would shutdown a runner that had previously failed,
leading to loops of create -> shutdown -> create.

now when a runner critically errors, we set its status to Shutdown, and
delete it. a failed runner is expected

src/exo/shared/types/worker/runners.py

@@ -63,7 +63,7 @@ class RunnerShutdown(BaseRunnerStatus):
 
 
 class RunnerFailed(BaseRunnerStatus):
-    error_message: str | None = None
+    error_message: str
 
 
 RunnerStatus = (

src/exo/worker/plan.py

@@ -55,7 +55,7 @@ def plan(
     # Python short circuiting OR logic should evaluate these sequentially.
     return (
         _cancel_tasks(runners, tasks)
-        or _kill_runner(runners, all_runners, instances)
+        or _kill_runner(runners, instances)
         or _create_runner(node_id, runners, instances)
         or _model_needs_download(node_id, runners, global_download_status)
         or _init_distributed_backend(runners, all_runners)
@@ -67,25 +67,14 @@ def plan(
 
 def _kill_runner(
     runners: Mapping[RunnerId, RunnerSupervisor],
-    all_runners: Mapping[RunnerId, RunnerStatus],
     instances: Mapping[InstanceId, Instance],
 ) -> Shutdown | None:
     for runner in runners.values():
         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)
-
-        for (
-            global_runner_id
-        ) in runner.bound_instance.instance.shard_assignments.node_to_runner.values():
-            if runner_id == global_runner_id:
-                continue
-
-            if isinstance(all_runners.get(global_runner_id, None), RunnerFailed):
-                return Shutdown(
-                    instance_id=instance_id,
-                    runner_id=runner_id,
-                )
+        if isinstance(runner.status, RunnerFailed):
+            return Shutdown(instance_id=instance_id, runner_id=runner_id)
 
 
 def _create_runner(

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

@@ -39,7 +39,6 @@ from exo.shared.types.worker.runner_response import (
 from exo.shared.types.worker.runners import (
     RunnerConnected,
     RunnerConnecting,
-    RunnerFailed,
     RunnerIdle,
     RunnerLoaded,
     RunnerLoading,
@@ -256,9 +255,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

@@ -146,9 +146,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

@@ -24,6 +24,7 @@ from exo.shared.types.tasks import (
     CANCEL_ALL_TASKS,
     ImageEdits,
     ImageGeneration,
+    Shutdown,
     Task,
     TaskId,
     TaskStatus,
@@ -169,7 +170,8 @@ class RunnerSupervisor:
             await self._task_sender.send_async(task)
         except ClosedResourceError:
             self.in_progress.pop(task.task_id, None)
-            logger.warning(f"Task {task} dropped, runner closed communication.")
+            if not isinstance(task, Shutdown):
+                logger.warning(f"Task {task} dropped, runner closed communication.")
             return
         await event.wait()