Prevent failed instance retries (#1763)

## Motivation

Currently, when a runner fails, the master retries the instance. Most of
the time, this causes a loop over failure. Retries need backoff and a
cap.

## Changes

- src/exo/worker/main.py: Before creating a runner, check an exponential
backoff timer per instance. After EXO_MAX_INSTANCE_RETRIES failures,
send DeleteInstance to permanently remove the instance. Record attempts
on Shutdown; reset on InstanceDeleted.
- src/exo/utils/keyed_backoff.py: Add attempts() method to query retry
count
- src/exo/shared/constants.py: Add EXO_MAX_INSTANCE_RETRIES = 3.

## Why It Works

The worker gates CreateRunner tasks behind a KeyedBackoff, adding
exponential delay (2s base, 30s cap) between retries. After 3 failures
the worker sends DeleteInstance, stopping retries entirely. The backoff
resets when the instance is deleted, so a fresh placement starts clean.

---------

Co-authored-by: Evan <evanev7@gmail.com>

src/exo/api/main.py

@@ -169,20 +169,10 @@ 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
@@ -1824,25 +1814,9 @@ 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

@@ -1,93 +0,0 @@
-# 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,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 = 3

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,
@@ -44,6 +47,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 +88,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 +134,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 +166,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 +194,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

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, 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)
     )
 
 
@@ -92,8 +97,12 @@ 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
@@ -116,6 +125,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 +134,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 +286,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 +301,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/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