diff --git a/src/exo/shared/constants.py b/src/exo/shared/constants.py index e6870a73..12bc1634 100644 --- a/src/exo/shared/constants.py +++ b/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 diff --git a/src/exo/utils/keyed_backoff.py b/src/exo/utils/keyed_backoff.py index b481d43e..4d7c9a66 100644 --- a/src/exo/utils/keyed_backoff.py +++ b/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) diff --git a/src/exo/worker/main.py b/src/exo/worker/main.py index ded0481b..08cc1e0e 100644 --- a/src/exo/worker/main.py +++ b/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 diff --git a/src/exo/worker/plan.py b/src/exo/worker/plan.py index 0df10ca4..71348118 100644 --- a/src/exo/worker/plan.py +++ b/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: diff --git a/src/exo/worker/tests/unittests/test_plan/test_download_and_loading.py b/src/exo/worker/tests/unittests/test_plan/test_download_and_loading.py index abcb4939..461d8f12 100644 --- a/src/exo/worker/tests/unittests/test_plan/test_download_and_loading.py +++ b/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 diff --git a/src/exo/worker/tests/unittests/test_plan/test_runner_lifecycle.py b/src/exo/worker/tests/unittests/test_plan/test_runner_lifecycle.py index a1b87ca1..1ac9dee0 100644 --- a/src/exo/worker/tests/unittests/test_plan/test_runner_lifecycle.py +++ b/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 diff --git a/src/exo/worker/tests/unittests/test_plan/test_task_forwarding.py b/src/exo/worker/tests/unittests/test_plan/test_task_forwarding.py index 64be74e6..0aa901ce 100644 --- a/src/exo/worker/tests/unittests/test_plan/test_task_forwarding.py +++ b/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 diff --git a/src/exo/worker/tests/unittests/test_plan/test_warmup.py b/src/exo/worker/tests/unittests/test_plan/test_warmup.py index 52848446..612868e7 100644 --- a/src/exo/worker/tests/unittests/test_plan/test_warmup.py +++ b/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