← back to Exo
downloads: refactor to run at node level
9357503c6fcfb23bc2ca3cf57c945a5796a3319d · 2026-01-19 18:59:32 +0000 · Jake Hillion
The Worker previously owned the ShardDownloader directly via dependency
injection, which prevented --no-worker nodes from downloading and made
it impossible for multiple Workers to share a single downloader instance.
Moved download functionality to a new DownloadCoordinator component at
the Node level that communicates via the DOWNLOAD_COMMANDS pub/sub topic.
Workers now send StartDownload commands instead of calling the downloader
directly, and receive progress updates through the event-sourced state.
This decouples downloads from the Worker lifecycle and enables future
features like UI-triggered downloads to specific nodes and multi-worker
download sharing.
Test plan:
- Mostly tested in the next PR that adds explicit downloads/deletions to
the dashboard.
- Started a model that isn't downloaded - it works.
Files touched
A src/exo/download/coordinator.pyR099 src/exo/worker/download/download_utils.py src/exo/download/download_utils.pyR100 src/exo/worker/download/huggingface_utils.py src/exo/download/huggingface_utils.pyR097 src/exo/worker/download/impl_shard_downloader.py src/exo/download/impl_shard_downloader.pyR097 src/exo/worker/download/shard_downloader.py src/exo/download/shard_downloader.pyM src/exo/main.pyM src/exo/routing/topics.pyM src/exo/shared/models/model_cards.pyM src/exo/shared/types/commands.pyA src/exo/utils/keyed_backoff.pyM src/exo/worker/engines/image/distributed_model.pyM src/exo/worker/engines/mlx/utils_mlx.pyM src/exo/worker/main.pyM src/exo/worker/plan.pyM src/exo/worker/tests/unittests/test_mlx/test_tokenizers.pyM src/exo/worker/tests/unittests/test_plan/test_download_and_loading.pyM src/exo/worker/tests/unittests/test_plan/test_runner_lifecycle.pyM src/exo/worker/tests/unittests/test_plan/test_task_forwarding.pyM src/exo/worker/tests/unittests/test_plan/test_warmup.pyM tests/headless_runner.py
Diff
commit 9357503c6fcfb23bc2ca3cf57c945a5796a3319d
Author: Jake Hillion <jake@hillion.co.uk>
Date: Mon Jan 19 18:59:32 2026 +0000
downloads: refactor to run at node level
The Worker previously owned the ShardDownloader directly via dependency
injection, which prevented --no-worker nodes from downloading and made
it impossible for multiple Workers to share a single downloader instance.
Moved download functionality to a new DownloadCoordinator component at
the Node level that communicates via the DOWNLOAD_COMMANDS pub/sub topic.
Workers now send StartDownload commands instead of calling the downloader
directly, and receive progress updates through the event-sourced state.
This decouples downloads from the Worker lifecycle and enables future
features like UI-triggered downloads to specific nodes and multi-worker
download sharing.
Test plan:
- Mostly tested in the next PR that adds explicit downloads/deletions to
the dashboard.
- Started a model that isn't downloaded - it works.
---
src/exo/download/coordinator.py | 284 +++++++++++++++++++++
src/exo/{worker => }/download/download_utils.py | 14 +-
src/exo/{worker => }/download/huggingface_utils.py | 0
.../{worker => }/download/impl_shard_downloader.py | 4 +-
src/exo/{worker => }/download/shard_downloader.py | 2 +-
src/exo/main.py | 69 ++++-
src/exo/routing/topics.py | 5 +-
src/exo/shared/models/model_cards.py | 6 +-
src/exo/shared/types/commands.py | 22 +-
src/exo/utils/keyed_backoff.py | 32 +++
src/exo/worker/engines/image/distributed_model.py | 2 +-
src/exo/worker/engines/mlx/utils_mlx.py | 2 +-
src/exo/worker/main.py | 222 +++-------------
src/exo/worker/plan.py | 14 +-
.../tests/unittests/test_mlx/test_tokenizers.py | 4 +-
.../test_plan/test_download_and_loading.py | 39 +--
.../unittests/test_plan/test_runner_lifecycle.py | 11 +-
.../unittests/test_plan/test_task_forwarding.py | 5 -
.../tests/unittests/test_plan/test_warmup.py | 8 -
tests/headless_runner.py | 8 +-
20 files changed, 479 insertions(+), 274 deletions(-)
diff --git a/src/exo/download/coordinator.py b/src/exo/download/coordinator.py
new file mode 100644
index 00000000..c2f7b9e9
--- /dev/null
+++ b/src/exo/download/coordinator.py
@@ -0,0 +1,284 @@
+import asyncio
+from dataclasses import dataclass, field
+from typing import Iterator
+
+import anyio
+from anyio import current_time
+from anyio.abc import TaskGroup
+from loguru import logger
+
+from exo.download.download_utils import (
+ RepoDownloadProgress,
+ delete_model,
+ map_repo_download_progress_to_download_progress_data,
+)
+from exo.download.shard_downloader import ShardDownloader
+from exo.shared.models.model_cards import ModelId
+from exo.shared.types.commands import (
+ DeleteDownload,
+ ForwarderDownloadCommand,
+ StartDownload,
+)
+from exo.shared.types.common import NodeId, SessionId
+from exo.shared.types.events import (
+ Event,
+ ForwarderEvent,
+ NodeDownloadProgress,
+)
+from exo.shared.types.worker.downloads import (
+ DownloadCompleted,
+ DownloadFailed,
+ DownloadOngoing,
+ DownloadPending,
+ DownloadProgress,
+)
+from exo.shared.types.worker.shards import ShardMetadata
+from exo.utils.channels import Receiver, Sender, channel
+
+
+@dataclass
+class DownloadCoordinator:
+ node_id: NodeId
+ session_id: SessionId
+ shard_downloader: ShardDownloader
+ download_command_receiver: Receiver[ForwarderDownloadCommand]
+ local_event_sender: Sender[ForwarderEvent]
+ event_index_counter: Iterator[int]
+
+ # Local state
+ download_status: dict[ModelId, DownloadProgress] = field(default_factory=dict)
+ active_downloads: dict[ModelId, asyncio.Task[None]] = field(default_factory=dict)
+
+ # Internal event channel for forwarding (initialized in __post_init__)
+ event_sender: Sender[Event] = field(init=False)
+ event_receiver: Receiver[Event] = field(init=False)
+ _tg: TaskGroup = field(init=False)
+
+ def __post_init__(self) -> None:
+ self.event_sender, self.event_receiver = channel[Event]()
+ self._tg = anyio.create_task_group()
+
+ async def run(self) -> None:
+ logger.info("Starting DownloadCoordinator")
+ async with self._tg as tg:
+ tg.start_soon(self._command_processor)
+ tg.start_soon(self._forward_events)
+ tg.start_soon(self._emit_existing_download_progress)
+
+ def shutdown(self) -> None:
+ self._tg.cancel_scope.cancel()
+
+ async def _command_processor(self) -> None:
+ with self.download_command_receiver as commands:
+ async for cmd in commands:
+ # Only process commands targeting this node
+ if cmd.command.target_node_id != self.node_id:
+ continue
+
+ match cmd.command:
+ case StartDownload(shard_metadata=shard):
+ await self._start_download(shard)
+ case DeleteDownload(model_id=model_id):
+ await self._delete_download(model_id)
+
+ async def _start_download(self, shard: ShardMetadata) -> None:
+ model_id = shard.model_card.model_id
+
+ # Check if already downloading or complete
+ if model_id in self.download_status:
+ status = self.download_status[model_id]
+ if isinstance(status, (DownloadOngoing, DownloadCompleted)):
+ logger.debug(
+ f"Download for {model_id} already in progress or complete, skipping"
+ )
+ return
+
+ # Emit pending status
+ progress = DownloadPending(shard_metadata=shard, node_id=self.node_id)
+ self.download_status[model_id] = progress
+ await self.event_sender.send(NodeDownloadProgress(download_progress=progress))
+
+ # Check initial status from downloader
+ initial_progress = (
+ await self.shard_downloader.get_shard_download_status_for_shard(shard)
+ )
+
+ if initial_progress.status == "complete":
+ completed = DownloadCompleted(
+ shard_metadata=shard,
+ node_id=self.node_id,
+ total_bytes=initial_progress.total_bytes,
+ )
+ self.download_status[model_id] = completed
+ await self.event_sender.send(
+ NodeDownloadProgress(download_progress=completed)
+ )
+ return
+
+ # Start actual download
+ self._start_download_task(shard, initial_progress)
+
+ def _start_download_task(
+ self, shard: ShardMetadata, initial_progress: RepoDownloadProgress
+ ) -> None:
+ model_id = shard.model_card.model_id
+
+ # Emit ongoing status
+ status = DownloadOngoing(
+ node_id=self.node_id,
+ shard_metadata=shard,
+ download_progress=map_repo_download_progress_to_download_progress_data(
+ initial_progress
+ ),
+ )
+ self.download_status[model_id] = status
+ self.event_sender.send_nowait(NodeDownloadProgress(download_progress=status))
+
+ last_progress_time = 0.0
+ throttle_interval_secs = 1.0
+
+ async def download_progress_callback(
+ callback_shard: ShardMetadata, progress: RepoDownloadProgress
+ ) -> None:
+ nonlocal last_progress_time
+
+ if progress.status == "complete":
+ completed = DownloadCompleted(
+ shard_metadata=callback_shard,
+ node_id=self.node_id,
+ total_bytes=progress.total_bytes,
+ )
+ self.download_status[callback_shard.model_card.model_id] = completed
+ await self.event_sender.send(
+ NodeDownloadProgress(download_progress=completed)
+ )
+ # Clean up active download tracking
+ if callback_shard.model_card.model_id in self.active_downloads:
+ del self.active_downloads[callback_shard.model_card.model_id]
+ elif (
+ progress.status == "in_progress"
+ and current_time() - last_progress_time > throttle_interval_secs
+ ):
+ ongoing = DownloadOngoing(
+ node_id=self.node_id,
+ shard_metadata=callback_shard,
+ download_progress=map_repo_download_progress_to_download_progress_data(
+ progress
+ ),
+ )
+ self.download_status[callback_shard.model_card.model_id] = ongoing
+ await self.event_sender.send(
+ NodeDownloadProgress(download_progress=ongoing)
+ )
+ last_progress_time = current_time()
+
+ self.shard_downloader.on_progress(download_progress_callback)
+
+ async def download_wrapper() -> None:
+ try:
+ await self.shard_downloader.ensure_shard(shard)
+ except Exception as e:
+ logger.error(f"Download failed for {model_id}: {e}")
+ failed = DownloadFailed(
+ shard_metadata=shard,
+ node_id=self.node_id,
+ error_message=str(e),
+ )
+ self.download_status[model_id] = failed
+ await self.event_sender.send(
+ NodeDownloadProgress(download_progress=failed)
+ )
+ finally:
+ if model_id in self.active_downloads:
+ del self.active_downloads[model_id]
+
+ task = asyncio.create_task(download_wrapper())
+ self.active_downloads[model_id] = task
+
+ async def _delete_download(self, model_id: ModelId) -> None:
+ # Cancel if active
+ if model_id in self.active_downloads:
+ logger.info(f"Cancelling active download for {model_id} before deletion")
+ self.active_downloads[model_id].cancel()
+ del self.active_downloads[model_id]
+
+ # Delete from disk
+ logger.info(f"Deleting model files for {model_id}")
+ deleted = await delete_model(model_id)
+
+ if deleted:
+ logger.info(f"Successfully deleted model {model_id}")
+ else:
+ logger.warning(f"Model {model_id} was not found on disk")
+
+ # Emit pending status to reset UI state, then remove from local tracking
+ if model_id in self.download_status:
+ current_status = self.download_status[model_id]
+ pending = DownloadPending(
+ shard_metadata=current_status.shard_metadata,
+ node_id=self.node_id,
+ )
+ await self.event_sender.send(
+ NodeDownloadProgress(download_progress=pending)
+ )
+ del self.download_status[model_id]
+
+ async def _forward_events(self) -> None:
+ with self.event_receiver as events:
+ async for event in events:
+ idx = next(self.event_index_counter)
+ fe = ForwarderEvent(
+ origin_idx=idx,
+ origin=self.node_id,
+ session=self.session_id,
+ event=event,
+ )
+ logger.debug(
+ f"DownloadCoordinator published event {idx}: {str(event)[:100]}"
+ )
+ await self.local_event_sender.send(fe)
+
+ async def _emit_existing_download_progress(self) -> None:
+ try:
+ while True:
+ logger.info(
+ "DownloadCoordinator: Fetching and emitting existing download progress..."
+ )
+ async for (
+ _,
+ progress,
+ ) in self.shard_downloader.get_shard_download_status():
+ if progress.status == "complete":
+ status: DownloadProgress = DownloadCompleted(
+ node_id=self.node_id,
+ shard_metadata=progress.shard,
+ total_bytes=progress.total_bytes,
+ )
+ elif progress.status in ["in_progress", "not_started"]:
+ if progress.downloaded_bytes_this_session.in_bytes == 0:
+ status = DownloadPending(
+ node_id=self.node_id, shard_metadata=progress.shard
+ )
+ else:
+ status = DownloadOngoing(
+ node_id=self.node_id,
+ shard_metadata=progress.shard,
+ download_progress=map_repo_download_progress_to_download_progress_data(
+ progress
+ ),
+ )
+ else:
+ continue
+
+ self.download_status[progress.shard.model_card.model_id] = status
+ await self.event_sender.send(
+ NodeDownloadProgress(download_progress=status)
+ )
+ logger.info(
+ "DownloadCoordinator: Done emitting existing download progress."
+ )
+ await anyio.sleep(5 * 60) # 5 minutes
+ except Exception as e:
+ logger.error(
+ f"DownloadCoordinator: Error emitting existing download progress: {e}"
+ )
diff --git a/src/exo/worker/download/download_utils.py b/src/exo/download/download_utils.py
similarity index 99%
rename from src/exo/worker/download/download_utils.py
rename to src/exo/download/download_utils.py
index 8816d701..f08aa0ed 100644
--- a/src/exo/worker/download/download_utils.py
+++ b/src/exo/download/download_utils.py
@@ -24,6 +24,13 @@ from pydantic import (
TypeAdapter,
)
+from exo.download.huggingface_utils import (
+ filter_repo_objects,
+ get_allow_patterns,
+ get_auth_headers,
+ get_hf_endpoint,
+ get_hf_token,
+)
from exo.shared.constants import EXO_MODELS_DIR
from exo.shared.types.common import ModelId
from exo.shared.types.memory import Memory
@@ -35,13 +42,6 @@ from exo.shared.types.worker.downloads import (
RepoFileDownloadProgress,
)
from exo.shared.types.worker.shards import ShardMetadata
-from exo.worker.download.huggingface_utils import (
- filter_repo_objects,
- get_allow_patterns,
- get_auth_headers,
- get_hf_endpoint,
- get_hf_token,
-)
class HuggingFaceAuthenticationError(Exception):
diff --git a/src/exo/worker/download/huggingface_utils.py b/src/exo/download/huggingface_utils.py
similarity index 100%
rename from src/exo/worker/download/huggingface_utils.py
rename to src/exo/download/huggingface_utils.py
diff --git a/src/exo/worker/download/impl_shard_downloader.py b/src/exo/download/impl_shard_downloader.py
similarity index 97%
rename from src/exo/worker/download/impl_shard_downloader.py
rename to src/exo/download/impl_shard_downloader.py
index b253ad9d..24f13afe 100644
--- a/src/exo/worker/download/impl_shard_downloader.py
+++ b/src/exo/download/impl_shard_downloader.py
@@ -5,13 +5,13 @@ from typing import AsyncIterator, Callable
from loguru import logger
+from exo.download.download_utils import RepoDownloadProgress, download_shard
+from exo.download.shard_downloader import ShardDownloader
from exo.shared.models.model_cards import MODEL_CARDS, ModelCard, ModelId
from exo.shared.types.worker.shards import (
PipelineShardMetadata,
ShardMetadata,
)
-from exo.worker.download.download_utils import RepoDownloadProgress, download_shard
-from exo.worker.download.shard_downloader import ShardDownloader
def exo_shard_downloader(max_parallel_downloads: int = 8) -> ShardDownloader:
diff --git a/src/exo/worker/download/shard_downloader.py b/src/exo/download/shard_downloader.py
similarity index 97%
rename from src/exo/worker/download/shard_downloader.py
rename to src/exo/download/shard_downloader.py
index deb0aaee..30c11d25 100644
--- a/src/exo/worker/download/shard_downloader.py
+++ b/src/exo/download/shard_downloader.py
@@ -5,13 +5,13 @@ from datetime import timedelta
from pathlib import Path
from typing import AsyncIterator, Callable
+from exo.download.download_utils import RepoDownloadProgress
from exo.shared.models.model_cards import ModelCard, ModelId, ModelTask
from exo.shared.types.memory import Memory
from exo.shared.types.worker.shards import (
PipelineShardMetadata,
ShardMetadata,
)
-from exo.worker.download.download_utils import RepoDownloadProgress
# TODO: the PipelineShardMetadata getting reinstantiated is a bit messy. Should this be a classmethod?
diff --git a/src/exo/main.py b/src/exo/main.py
index 6425595d..8223383e 100644
--- a/src/exo/main.py
+++ b/src/exo/main.py
@@ -1,10 +1,11 @@
import argparse
+import itertools
import multiprocessing as mp
import os
import resource
import signal
from dataclasses import dataclass, field
-from typing import Self
+from typing import Iterator, Self
import anyio
from anyio.abc import TaskGroup
@@ -12,6 +13,8 @@ from loguru import logger
from pydantic import PositiveInt
import exo.routing.topics as topics
+from exo.download.coordinator import DownloadCoordinator
+from exo.download.impl_shard_downloader import exo_shard_downloader
from exo.master.api import API # TODO: should API be in master?
from exo.master.main import Master
from exo.routing.router import Router, get_node_id_keypair
@@ -21,7 +24,6 @@ from exo.shared.logging import logger_cleanup, logger_setup
from exo.shared.types.common import NodeId, SessionId
from exo.utils.channels import Receiver, channel
from exo.utils.pydantic_ext import CamelCaseModel
-from exo.worker.download.impl_shard_downloader import exo_shard_downloader
from exo.worker.main import Worker
@@ -29,6 +31,7 @@ from exo.worker.main import Worker
@dataclass
class Node:
router: Router
+ download_coordinator: DownloadCoordinator | None
worker: Worker | None
election: Election # Every node participates in election, as we do want a node to become master even if it isn't a master candidate if no master candidates are present.
election_result_receiver: Receiver[ElectionResult]
@@ -36,6 +39,7 @@ class Node:
api: API | None
node_id: NodeId
+ event_index_counter: Iterator[int]
_tg: TaskGroup = field(init=False, default_factory=anyio.create_task_group)
@classmethod
@@ -49,8 +53,26 @@ class Node:
await router.register_topic(topics.COMMANDS)
await router.register_topic(topics.ELECTION_MESSAGES)
await router.register_topic(topics.CONNECTION_MESSAGES)
+ await router.register_topic(topics.DOWNLOAD_COMMANDS)
logger.info(f"Starting node {node_id}")
+
+ # Create shared event index counter for Worker and DownloadCoordinator
+ event_index_counter = itertools.count()
+
+ # Create DownloadCoordinator (unless --no-downloads)
+ if not args.no_downloads:
+ download_coordinator = DownloadCoordinator(
+ node_id,
+ session_id,
+ exo_shard_downloader(),
+ download_command_receiver=router.receiver(topics.DOWNLOAD_COMMANDS),
+ local_event_sender=router.sender(topics.LOCAL_EVENTS),
+ event_index_counter=event_index_counter,
+ )
+ else:
+ download_coordinator = None
+
if args.spawn_api:
api = API(
node_id,
@@ -67,11 +89,12 @@ class Node:
worker = Worker(
node_id,
session_id,
- exo_shard_downloader(),
connection_message_receiver=router.receiver(topics.CONNECTION_MESSAGES),
global_event_receiver=router.receiver(topics.GLOBAL_EVENTS),
local_event_sender=router.sender(topics.LOCAL_EVENTS),
command_sender=router.sender(topics.COMMANDS),
+ download_command_sender=router.sender(topics.DOWNLOAD_COMMANDS),
+ event_index_counter=event_index_counter,
)
else:
worker = None
@@ -99,13 +122,25 @@ class Node:
election_result_sender=er_send,
)
- return cls(router, worker, election, er_recv, master, api, node_id)
+ return cls(
+ router,
+ download_coordinator,
+ worker,
+ election,
+ er_recv,
+ master,
+ api,
+ node_id,
+ event_index_counter,
+ )
async def run(self):
async with self._tg as tg:
signal.signal(signal.SIGINT, lambda _, __: self.shutdown())
tg.start_soon(self.router.run)
tg.start_soon(self.election.run)
+ if self.download_coordinator:
+ tg.start_soon(self.download_coordinator.run)
if self.worker:
tg.start_soon(self.worker.run)
if self.master:
@@ -170,13 +205,27 @@ class Node:
)
if result.is_new_master:
await anyio.sleep(0)
+ # Fresh counter for new session (buffer expects indices from 0)
+ self.event_index_counter = itertools.count()
+ if self.download_coordinator:
+ self.download_coordinator.shutdown()
+ self.download_coordinator = DownloadCoordinator(
+ self.node_id,
+ result.session_id,
+ exo_shard_downloader(),
+ download_command_receiver=self.router.receiver(
+ topics.DOWNLOAD_COMMANDS
+ ),
+ local_event_sender=self.router.sender(topics.LOCAL_EVENTS),
+ event_index_counter=self.event_index_counter,
+ )
+ self._tg.start_soon(self.download_coordinator.run)
if self.worker:
self.worker.shutdown()
# TODO: add profiling etc to resource monitor
self.worker = Worker(
self.node_id,
result.session_id,
- exo_shard_downloader(),
connection_message_receiver=self.router.receiver(
topics.CONNECTION_MESSAGES
),
@@ -185,6 +234,10 @@ class Node:
),
local_event_sender=self.router.sender(topics.LOCAL_EVENTS),
command_sender=self.router.sender(topics.COMMANDS),
+ download_command_sender=self.router.sender(
+ topics.DOWNLOAD_COMMANDS
+ ),
+ event_index_counter=self.event_index_counter,
)
self._tg.start_soon(self.worker.run)
if self.api:
@@ -226,6 +279,7 @@ class Args(CamelCaseModel):
api_port: PositiveInt = 52415
tb_only: bool = False
no_worker: bool = False
+ no_downloads: bool = False
fast_synch: bool | None = None # None = auto, True = force on, False = force off
@classmethod
@@ -268,6 +322,11 @@ class Args(CamelCaseModel):
"--no-worker",
action="store_true",
)
+ parser.add_argument(
+ "--no-downloads",
+ action="store_true",
+ help="Disable the download coordinator (node won't download models)",
+ )
fast_synch_group = parser.add_mutually_exclusive_group()
fast_synch_group.add_argument(
"--fast-synch",
diff --git a/src/exo/routing/topics.py b/src/exo/routing/topics.py
index 50f1c9af..5a122a95 100644
--- a/src/exo/routing/topics.py
+++ b/src/exo/routing/topics.py
@@ -3,7 +3,7 @@ from enum import Enum
from exo.routing.connection_message import ConnectionMessage
from exo.shared.election import ElectionMessage
-from exo.shared.types.commands import ForwarderCommand
+from exo.shared.types.commands import ForwarderCommand, ForwarderDownloadCommand
from exo.shared.types.events import (
ForwarderEvent,
)
@@ -45,3 +45,6 @@ ELECTION_MESSAGES = TypedTopic(
CONNECTION_MESSAGES = TypedTopic(
"connection_messages", PublishPolicy.Never, ConnectionMessage
)
+DOWNLOAD_COMMANDS = TypedTopic(
+ "download_commands", PublishPolicy.Always, ForwarderDownloadCommand
+)
diff --git a/src/exo/shared/models/model_cards.py b/src/exo/shared/models/model_cards.py
index eaad0154..35e077ba 100644
--- a/src/exo/shared/models/model_cards.py
+++ b/src/exo/shared/models/model_cards.py
@@ -621,7 +621,7 @@ class ConfigData(BaseModel):
async def get_config_data(model_id: ModelId) -> ConfigData:
"""Downloads and parses config.json for a model."""
- from exo.worker.download.download_utils import (
+ from exo.download.download_utils import (
download_file_with_retry,
ensure_models_dir,
)
@@ -643,11 +643,11 @@ async def get_config_data(model_id: ModelId) -> ConfigData:
async def get_safetensors_size(model_id: ModelId) -> Memory:
"""Gets model size from safetensors index or falls back to HF API."""
- from exo.shared.types.worker.downloads import ModelSafetensorsIndex
- from exo.worker.download.download_utils import (
+ from exo.download.download_utils import (
download_file_with_retry,
ensure_models_dir,
)
+ from exo.shared.types.worker.downloads import ModelSafetensorsIndex
target_dir = (await ensure_models_dir()) / model_id.normalize()
await aios.makedirs(target_dir, exist_ok=True)
diff --git a/src/exo/shared/types/commands.py b/src/exo/shared/types/commands.py
index a0ddb968..ed086f35 100644
--- a/src/exo/shared/types/commands.py
+++ b/src/exo/shared/types/commands.py
@@ -1,6 +1,6 @@
from pydantic import Field
-from exo.shared.models.model_cards import ModelCard
+from exo.shared.models.model_cards import ModelCard, ModelId
from exo.shared.types.api import (
ChatCompletionTaskParams,
ImageEditsInternalParams,
@@ -9,7 +9,7 @@ from exo.shared.types.api import (
from exo.shared.types.chunks import InputImageChunk
from exo.shared.types.common import CommandId, NodeId
from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
-from exo.shared.types.worker.shards import Sharding
+from exo.shared.types.worker.shards import Sharding, ShardMetadata
from exo.utils.pydantic_ext import CamelCaseModel, TaggedModel
@@ -62,6 +62,19 @@ class RequestEventLog(BaseCommand):
since_idx: int
+class StartDownload(BaseCommand):
+ target_node_id: NodeId
+ shard_metadata: ShardMetadata
+
+
+class DeleteDownload(BaseCommand):
+ target_node_id: NodeId
+ model_id: ModelId
+
+
+DownloadCommand = StartDownload | DeleteDownload
+
+
Command = (
TestCommand
| RequestEventLog
@@ -79,3 +92,8 @@ Command = (
class ForwarderCommand(CamelCaseModel):
origin: NodeId
command: Command
+
+
+class ForwarderDownloadCommand(CamelCaseModel):
+ origin: NodeId
+ command: DownloadCommand
diff --git a/src/exo/utils/keyed_backoff.py b/src/exo/utils/keyed_backoff.py
new file mode 100644
index 00000000..b481d43e
--- /dev/null
+++ b/src/exo/utils/keyed_backoff.py
@@ -0,0 +1,32 @@
+import time
+from typing import Generic, TypeVar
+
+K = TypeVar("K")
+
+
+class KeyedBackoff(Generic[K]):
+ """Tracks exponential backoff state per key."""
+
+ def __init__(self, base: float = 0.5, cap: float = 10.0):
+ self._base = base
+ self._cap = cap
+ self._attempts: dict[K, int] = {}
+ self._last_time: dict[K, float] = {}
+
+ def should_proceed(self, key: K) -> bool:
+ """Returns True if enough time has elapsed since last attempt."""
+ now = time.monotonic()
+ last = self._last_time.get(key, 0.0)
+ attempts = self._attempts.get(key, 0)
+ delay = min(self._cap, self._base * (2.0**attempts))
+ return now - last >= delay
+
+ def record_attempt(self, key: K) -> None:
+ """Record that an attempt was made for this key."""
+ self._last_time[key] = time.monotonic()
+ self._attempts[key] = self._attempts.get(key, 0) + 1
+
+ def reset(self, key: K) -> None:
+ """Reset backoff state for a key (e.g., on success)."""
+ self._attempts.pop(key, None)
+ self._last_time.pop(key, None)
diff --git a/src/exo/worker/engines/image/distributed_model.py b/src/exo/worker/engines/image/distributed_model.py
index 9c127f36..45fd0e29 100644
--- a/src/exo/worker/engines/image/distributed_model.py
+++ b/src/exo/worker/engines/image/distributed_model.py
@@ -6,10 +6,10 @@ import mlx.core as mx
from mflux.models.common.config.config import Config
from PIL import Image
+from exo.download.download_utils import build_model_path
from exo.shared.types.api import AdvancedImageParams
from exo.shared.types.worker.instances import BoundInstance
from exo.shared.types.worker.shards import PipelineShardMetadata
-from exo.worker.download.download_utils import build_model_path
from exo.worker.engines.image.config import ImageModelConfig
from exo.worker.engines.image.models import (
create_adapter_for_model,
diff --git a/src/exo/worker/engines/mlx/utils_mlx.py b/src/exo/worker/engines/mlx/utils_mlx.py
index 0703d6b5..cb0e72b3 100644
--- a/src/exo/worker/engines/mlx/utils_mlx.py
+++ b/src/exo/worker/engines/mlx/utils_mlx.py
@@ -41,6 +41,7 @@ import mlx.nn as nn
from mlx_lm.utils import load_model
from pydantic import RootModel
+from exo.download.download_utils import build_model_path
from exo.shared.types.api import ChatCompletionMessageText
from exo.shared.types.common import Host
from exo.shared.types.memory import Memory
@@ -55,7 +56,6 @@ from exo.shared.types.worker.shards import (
ShardMetadata,
TensorShardMetadata,
)
-from exo.worker.download.download_utils import build_model_path
from exo.worker.engines.mlx import Model
from exo.worker.engines.mlx.auto_parallel import (
TimeoutCallback,
diff --git a/src/exo/worker/main.py b/src/exo/worker/main.py
index e861884a..3473be41 100644
--- a/src/exo/worker/main.py
+++ b/src/exo/worker/main.py
@@ -1,8 +1,9 @@
from datetime import datetime, timezone
from random import random
+from typing import Iterator
import anyio
-from anyio import CancelScope, create_task_group, current_time, fail_after
+from anyio import CancelScope, create_task_group, fail_after
from anyio.abc import TaskGroup
from loguru import logger
@@ -10,7 +11,12 @@ from exo.routing.connection_message import ConnectionMessage, ConnectionMessageT
from exo.shared.apply import apply
from exo.shared.models.model_cards import ModelId
from exo.shared.types.api import ImageEditsInternalParams
-from exo.shared.types.commands import ForwarderCommand, RequestEventLog
+from exo.shared.types.commands import (
+ ForwarderCommand,
+ ForwarderDownloadCommand,
+ RequestEventLog,
+ StartDownload,
+)
from exo.shared.types.common import CommandId, NodeId, SessionId
from exo.shared.types.events import (
Event,
@@ -18,7 +24,6 @@ from exo.shared.types.events import (
ForwarderEvent,
IndexedEvent,
InputChunkReceived,
- NodeDownloadProgress,
NodeGatheredInfo,
TaskCreated,
TaskStatusUpdated,
@@ -36,23 +41,12 @@ from exo.shared.types.tasks import (
TaskStatus,
)
from exo.shared.types.topology import Connection, SocketConnection
-from exo.shared.types.worker.downloads import (
- DownloadCompleted,
- DownloadFailed,
- DownloadOngoing,
- DownloadPending,
- DownloadProgress,
-)
from exo.shared.types.worker.runners import RunnerId
-from exo.shared.types.worker.shards import ShardMetadata
from exo.utils.channels import Receiver, Sender, channel
from exo.utils.event_buffer import OrderedBuffer
from exo.utils.info_gatherer.info_gatherer import GatheredInfo, InfoGatherer
from exo.utils.info_gatherer.net_profile import check_reachable
-from exo.worker.download.download_utils import (
- map_repo_download_progress_to_download_progress_data,
-)
-from exo.worker.download.shard_downloader import RepoDownloadProgress, ShardDownloader
+from exo.utils.keyed_backoff import KeyedBackoff
from exo.worker.plan import plan
from exo.worker.runner.runner_supervisor import RunnerSupervisor
@@ -62,7 +56,6 @@ class Worker:
self,
node_id: NodeId,
session_id: SessionId,
- shard_downloader: ShardDownloader,
*,
connection_message_receiver: Receiver[ConnectionMessage],
global_event_receiver: Receiver[ForwarderEvent],
@@ -70,23 +63,22 @@ class Worker:
# This is for requesting updates. It doesn't need to be a general command sender right now,
# but I think it's the correct way to be thinking about commands
command_sender: Sender[ForwarderCommand],
+ download_command_sender: Sender[ForwarderDownloadCommand],
+ event_index_counter: Iterator[int],
):
self.node_id: NodeId = node_id
self.session_id: SessionId = session_id
- self.shard_downloader: ShardDownloader = shard_downloader
- self._pending_downloads: dict[RunnerId, ShardMetadata] = {}
-
self.global_event_receiver = global_event_receiver
self.local_event_sender = local_event_sender
- self.local_event_index = 0
+ self.event_index_counter = event_index_counter
self.command_sender = command_sender
+ self.download_command_sender = download_command_sender
self.connection_message_receiver = connection_message_receiver
self.event_buffer = OrderedBuffer[Event]()
self.out_for_delivery: dict[EventId, ForwarderEvent] = {}
self.state: State = State()
- self.download_status: dict[ModelId, DownloadProgress] = {}
self.runners: dict[RunnerId, RunnerSupervisor] = {}
self._tg: TaskGroup = create_task_group()
@@ -101,6 +93,8 @@ class Worker:
self.input_chunk_buffer: dict[CommandId, dict[int, str]] = {}
self.input_chunk_counts: dict[CommandId, int] = {}
+ self._download_backoff: KeyedBackoff[ModelId] = KeyedBackoff(base=0.5, cap=10.0)
+
async def run(self):
logger.info("Starting Worker")
@@ -111,7 +105,6 @@ class Worker:
tg.start_soon(info_gatherer.run)
tg.start_soon(self._forward_info, info_recv)
tg.start_soon(self.plan_step)
- tg.start_soon(self._emit_existing_download_progress)
tg.start_soon(self._connection_message_event_writer)
tg.start_soon(self._resend_out_for_delivery)
tg.start_soon(self._event_applier)
@@ -121,6 +114,7 @@ class Worker:
# Actual shutdown code - waits for all tasks to complete before executing.
self.local_event_sender.close()
self.command_sender.close()
+ self.download_command_sender.close()
for runner in self.runners.values():
runner.shutdown()
@@ -179,11 +173,9 @@ class Worker:
async def plan_step(self):
while True:
await anyio.sleep(0.1)
- # 3. based on the updated state, we plan & execute an operation.
task: Task | None = plan(
self.node_id,
self.runners,
- self.download_status,
self.state.downloads,
self.state.instances,
self.state.runners,
@@ -207,42 +199,26 @@ class Worker:
)
)
case DownloadModel(shard_metadata=shard):
- if shard.model_card.model_id not in self.download_status:
- progress = DownloadPending(
- shard_metadata=shard, node_id=self.node_id
- )
- self.download_status[shard.model_card.model_id] = progress
- await self.event_sender.send(
- NodeDownloadProgress(download_progress=progress)
- )
- initial_progress = (
- await self.shard_downloader.get_shard_download_status_for_shard(
- shard
+ model_id = shard.model_card.model_id
+ if not self._download_backoff.should_proceed(model_id):
+ continue
+
+ self._download_backoff.record_attempt(model_id)
+
+ await self.download_command_sender.send(
+ ForwarderDownloadCommand(
+ origin=self.node_id,
+ command=StartDownload(
+ target_node_id=self.node_id,
+ shard_metadata=shard,
+ ),
)
)
- if initial_progress.status == "complete":
- progress = DownloadCompleted(
- shard_metadata=shard,
- node_id=self.node_id,
- total_bytes=initial_progress.total_bytes,
- )
- self.download_status[shard.model_card.model_id] = progress
- await self.event_sender.send(
- NodeDownloadProgress(download_progress=progress)
- )
- await self.event_sender.send(
- TaskStatusUpdated(
- task_id=task.task_id,
- task_status=TaskStatus.Complete,
- )
- )
- else:
- await self.event_sender.send(
- TaskStatusUpdated(
- task_id=task.task_id, task_status=TaskStatus.Running
- )
+ await self.event_sender.send(
+ TaskStatusUpdated(
+ task_id=task.task_id, task_status=TaskStatus.Running
)
- self._handle_shard_download_process(task, initial_progress)
+ )
case Shutdown(runner_id=runner_id):
try:
with fail_after(3):
@@ -387,104 +363,17 @@ class Worker:
self._tg.start_soon(runner.run)
return runner
- def _handle_shard_download_process(
- self,
- task: DownloadModel,
- initial_progress: RepoDownloadProgress,
- ):
- """Manages the shard download process with progress tracking."""
- status = DownloadOngoing(
- node_id=self.node_id,
- shard_metadata=task.shard_metadata,
- download_progress=map_repo_download_progress_to_download_progress_data(
- initial_progress
- ),
- )
- self.download_status[task.shard_metadata.model_card.model_id] = status
- self.event_sender.send_nowait(NodeDownloadProgress(download_progress=status))
-
- last_progress_time = 0.0
- throttle_interval_secs = 1.0
-
- async def download_progress_callback(
- shard: ShardMetadata, progress: RepoDownloadProgress
- ) -> None:
- nonlocal self
- nonlocal last_progress_time
- if progress.status == "complete":
- status = DownloadCompleted(
- shard_metadata=shard,
- node_id=self.node_id,
- total_bytes=progress.total_bytes,
- )
- self.download_status[shard.model_card.model_id] = status
- await self.event_sender.send(
- NodeDownloadProgress(download_progress=status)
- )
- await self.event_sender.send(
- TaskStatusUpdated(
- task_id=task.task_id, task_status=TaskStatus.Complete
- )
- )
- elif (
- progress.status == "in_progress"
- and current_time() - last_progress_time > throttle_interval_secs
- ):
- status = DownloadOngoing(
- node_id=self.node_id,
- shard_metadata=shard,
- download_progress=map_repo_download_progress_to_download_progress_data(
- progress
- ),
- )
- self.download_status[shard.model_card.model_id] = status
- await self.event_sender.send(
- NodeDownloadProgress(download_progress=status)
- )
- last_progress_time = current_time()
-
- self.shard_downloader.on_progress(download_progress_callback)
-
- async def download_with_error_handling() -> None:
- try:
- await self.shard_downloader.ensure_shard(task.shard_metadata)
- except Exception as e:
- error_message = str(e)
- logger.error(
- f"Download failed for {task.shard_metadata.model_card.model_id}: {error_message}"
- )
- failed_status = DownloadFailed(
- node_id=self.node_id,
- shard_metadata=task.shard_metadata,
- error_message=error_message,
- )
- self.download_status[task.shard_metadata.model_card.model_id] = (
- failed_status
- )
- await self.event_sender.send(
- NodeDownloadProgress(download_progress=failed_status)
- )
- await self.event_sender.send(
- TaskStatusUpdated(
- task_id=task.task_id, task_status=TaskStatus.Failed
- )
- )
-
- self._tg.start_soon(download_with_error_handling)
-
async def _forward_events(self) -> None:
with self.event_receiver as events:
async for event in events:
+ idx = next(self.event_index_counter)
fe = ForwarderEvent(
- origin_idx=self.local_event_index,
+ origin_idx=idx,
origin=self.node_id,
session=self.session_id,
event=event,
)
- logger.debug(
- f"Worker published event {self.local_event_index}: {str(event)[:100]}"
- )
- self.local_event_index += 1
+ logger.debug(f"Worker published event {idx}: {str(event)[:100]}")
await self.local_event_sender.send(fe)
self.out_for_delivery[event.event_id] = fe
@@ -532,42 +421,3 @@ class Worker:
await self.event_sender.send(TopologyEdgeDeleted(conn=conn))
await anyio.sleep(10)
-
- async def _emit_existing_download_progress(self) -> None:
- try:
- while True:
- logger.debug("Fetching and emitting existing download progress...")
- async for (
- _,
- progress,
- ) in self.shard_downloader.get_shard_download_status():
- if progress.status == "complete":
- status = DownloadCompleted(
- node_id=self.node_id,
- shard_metadata=progress.shard,
- total_bytes=progress.total_bytes,
- )
- elif progress.status in ["in_progress", "not_started"]:
- if progress.downloaded_bytes_this_session.in_bytes == 0:
- status = DownloadPending(
- node_id=self.node_id, shard_metadata=progress.shard
- )
- else:
- status = DownloadOngoing(
- node_id=self.node_id,
- shard_metadata=progress.shard,
- download_progress=map_repo_download_progress_to_download_progress_data(
- progress
- ),
- )
- else:
- continue
-
- self.download_status[progress.shard.model_card.model_id] = status
- await self.event_sender.send(
- NodeDownloadProgress(download_progress=status)
- )
- logger.debug("Done emitting existing download progress.")
- await anyio.sleep(5 * 60) # 5 minutes
- except Exception as e:
- logger.error(f"Error emitting existing download progress: {e}")
diff --git a/src/exo/worker/plan.py b/src/exo/worker/plan.py
index 7877b96c..cb6fbbf3 100644
--- a/src/exo/worker/plan.py
+++ b/src/exo/worker/plan.py
@@ -2,7 +2,6 @@
from collections.abc import Mapping, Sequence
-from exo.shared.models.model_cards import ModelId
from exo.shared.types.common import CommandId, NodeId
from exo.shared.types.tasks import (
ChatCompletion,
@@ -45,9 +44,6 @@ def plan(
node_id: NodeId,
# Runners is expected to be FRESH and so should not come from state
runners: Mapping[RunnerId, RunnerSupervisor],
- # DL_status is expected to be FRESH and so should not come from state
- download_status: Mapping[ModelId, DownloadProgress],
- # gdls is not expected to be fresh
global_download_status: Mapping[NodeId, Sequence[DownloadProgress]],
instances: Mapping[InstanceId, Instance],
all_runners: Mapping[RunnerId, RunnerStatus], # all global
@@ -59,7 +55,7 @@ def plan(
return (
_kill_runner(runners, all_runners, instances)
or _create_runner(node_id, runners, instances)
- or _model_needs_download(runners, download_status)
+ or _model_needs_download(node_id, runners, global_download_status)
or _init_distributed_backend(runners, all_runners)
or _load_model(runners, all_runners, global_download_status)
or _ready_to_warmup(runners, all_runners)
@@ -115,9 +111,15 @@ def _create_runner(
def _model_needs_download(
+ node_id: NodeId,
runners: Mapping[RunnerId, RunnerSupervisor],
- download_status: Mapping[ModelId, DownloadProgress],
+ global_download_status: Mapping[NodeId, Sequence[DownloadProgress]],
) -> DownloadModel | None:
+ local_downloads = global_download_status.get(node_id, [])
+ download_status = {
+ dp.shard_metadata.model_card.model_id: dp for dp in local_downloads
+ }
+
for runner in runners.values():
model_id = runner.bound_instance.bound_shard.model_card.model_id
if isinstance(runner.status, RunnerIdle) and (
diff --git a/src/exo/worker/tests/unittests/test_mlx/test_tokenizers.py b/src/exo/worker/tests/unittests/test_mlx/test_tokenizers.py
index d0b6707b..1f51b703 100644
--- a/src/exo/worker/tests/unittests/test_mlx/test_tokenizers.py
+++ b/src/exo/worker/tests/unittests/test_mlx/test_tokenizers.py
@@ -11,12 +11,12 @@ from pathlib import Path
import pytest
-from exo.shared.models.model_cards import MODEL_CARDS, ModelCard, ModelId
-from exo.worker.download.download_utils import (
+from exo.download.download_utils import (
download_file_with_retry,
ensure_models_dir,
fetch_file_list_with_cache,
)
+from exo.shared.models.model_cards import MODEL_CARDS, ModelCard, ModelId
from exo.worker.engines.mlx.utils_mlx import (
get_eos_token_ids_for_model,
load_tokenizer_for_model_id,
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 20cf72df..28ecd8e4 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
@@ -1,5 +1,5 @@
import exo.worker.plan as plan_mod
-from exo.shared.types.common import ModelId, NodeId
+from exo.shared.types.common import NodeId
from exo.shared.types.memory import Memory
from exo.shared.types.tasks import LoadModel
from exo.shared.types.worker.downloads import DownloadCompleted, DownloadProgress
@@ -45,13 +45,9 @@ def test_plan_requests_download_when_waiting_and_shard_not_downloaded():
instances = {INSTANCE_1_ID: instance}
all_runners = {RUNNER_1_ID: RunnerIdle()}
- # No entry for this shard -> should trigger DownloadModel
- download_status: dict[ModelId, DownloadProgress] = {}
-
result = plan_mod.plan(
node_id=NODE_A,
runners=runners, # type: ignore
- download_status=download_status,
global_download_status={NODE_A: []},
instances=instances,
all_runners=all_runners,
@@ -92,14 +88,6 @@ def test_plan_loads_model_when_all_shards_downloaded_and_waiting():
RUNNER_2_ID: RunnerConnected(),
}
- # Local node has already marked its shard as downloaded (not actually used by _load_model)
- local_download_status = {
- MODEL_A_ID: DownloadCompleted(
- shard_metadata=shard1, node_id=NODE_A, total_bytes=Memory()
- )
- }
-
- # Global view has completed downloads for both nodes
global_download_status = {
NODE_A: [
DownloadCompleted(
@@ -116,7 +104,6 @@ def test_plan_loads_model_when_all_shards_downloaded_and_waiting():
result = plan_mod.plan(
node_id=NODE_A,
runners=runners, # type: ignore
- download_status=local_download_status,
global_download_status=global_download_status,
instances=instances,
all_runners=all_runners,
@@ -148,23 +135,19 @@ def test_plan_does_not_request_download_when_shard_already_downloaded():
instances = {INSTANCE_1_ID: instance}
all_runners = {RUNNER_1_ID: RunnerIdle()}
- # Local status claims the shard is downloaded already
- local_download_status = {
- MODEL_A_ID: DownloadCompleted(
- shard_metadata=shard, node_id=NODE_A, total_bytes=Memory()
- )
- }
-
- # Global view hasn't caught up yet (no completed shards recorded for NODE_A)
+ # Global state shows shard is downloaded for NODE_A
global_download_status: dict[NodeId, list[DownloadProgress]] = {
- NODE_A: [],
+ NODE_A: [
+ DownloadCompleted(
+ shard_metadata=shard, node_id=NODE_A, total_bytes=Memory()
+ )
+ ],
NODE_B: [],
}
result = plan_mod.plan(
node_id=NODE_A,
runners=runners, # type: ignore
- download_status=local_download_status,
global_download_status=global_download_status,
instances=instances,
all_runners=all_runners,
@@ -202,12 +185,6 @@ def test_plan_does_not_load_model_until_all_shards_downloaded_globally():
RUNNER_2_ID: RunnerConnected(),
}
- # Only NODE_A's shard is recorded as downloaded globally
- local_download_status = {
- MODEL_A_ID: DownloadCompleted(
- shard_metadata=shard1, node_id=NODE_A, total_bytes=Memory()
- )
- }
global_download_status = {
NODE_A: [
DownloadCompleted(
@@ -220,7 +197,6 @@ def test_plan_does_not_load_model_until_all_shards_downloaded_globally():
result = plan_mod.plan(
node_id=NODE_A,
runners=runners, # type: ignore
- download_status=local_download_status,
global_download_status=global_download_status,
instances=instances,
all_runners=all_runners,
@@ -245,7 +221,6 @@ def test_plan_does_not_load_model_until_all_shards_downloaded_globally():
result = plan_mod.plan(
node_id=NODE_A,
runners=runners, # type: ignore
- download_status=local_download_status,
global_download_status=global_download_status,
instances=instances,
all_runners=all_runners,
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 944cb6db..a1b87ca1 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
@@ -47,8 +47,7 @@ def test_plan_kills_runner_when_instance_missing():
result = plan_mod.plan(
node_id=NODE_A,
- runners=runners, # type: ignore
- download_status={},
+ runners=runners, # type: ignore[arg-type]
global_download_status={NODE_A: []},
instances=instances,
all_runners=all_runners,
@@ -87,8 +86,7 @@ def test_plan_kills_runner_when_sibling_failed():
result = plan_mod.plan(
node_id=NODE_A,
- runners=runners, # type: ignore
- download_status={},
+ runners=runners, # type: ignore[arg-type]
global_download_status={NODE_A: []},
instances=instances,
all_runners=all_runners,
@@ -120,7 +118,6 @@ def test_plan_creates_runner_when_missing_for_node():
result = plan_mod.plan(
node_id=NODE_A,
runners=runners,
- download_status={},
global_download_status={NODE_A: []},
instances=instances,
all_runners=all_runners,
@@ -158,8 +155,7 @@ def test_plan_does_not_create_runner_when_supervisor_already_present():
result = plan_mod.plan(
node_id=NODE_A,
- runners=runners, # type: ignore
- download_status={},
+ runners=runners, # type: ignore[arg-type]
global_download_status={NODE_A: []},
instances=instances,
all_runners=all_runners,
@@ -189,7 +185,6 @@ def test_plan_does_not_create_runner_for_unassigned_node():
result = plan_mod.plan(
node_id=NODE_A,
runners=runners, # type: ignore
- download_status={},
global_download_status={NODE_A: []},
instances=instances,
all_runners=all_runners,
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 5b3aa5ec..6c8fea8d 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
@@ -65,7 +65,6 @@ def test_plan_forwards_pending_chat_completion_when_runner_ready():
result = plan_mod.plan(
node_id=NODE_A,
runners=runners, # type: ignore
- download_status={},
global_download_status={NODE_A: []},
instances=instances,
all_runners=all_runners,
@@ -113,7 +112,6 @@ def test_plan_does_not_forward_chat_completion_if_any_runner_not_ready():
result = plan_mod.plan(
node_id=NODE_A,
runners=runners, # type: ignore
- download_status={},
global_download_status={NODE_A: [], NODE_B: []},
instances=instances,
all_runners=all_runners,
@@ -158,7 +156,6 @@ def test_plan_does_not_forward_tasks_for_other_instances():
result = plan_mod.plan(
node_id=NODE_A,
runners=runners, # type: ignore
- download_status={},
global_download_status={NODE_A: []},
instances=instances,
all_runners=all_runners,
@@ -221,7 +218,6 @@ def test_plan_ignores_non_pending_or_non_chat_tasks():
result = plan_mod.plan(
node_id=NODE_A,
runners=runners, # type: ignore
- download_status={},
global_download_status={NODE_A: [], NODE_B: []},
instances=instances,
all_runners=all_runners,
@@ -261,7 +257,6 @@ def test_plan_returns_none_when_nothing_to_do():
result = plan_mod.plan(
node_id=NODE_A,
runners=runners, # type: ignore
- download_status={},
global_download_status={NODE_A: [], NODE_B: []},
instances=instances,
all_runners=all_runners,
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 b42a5afd..52848446 100644
--- a/src/exo/worker/tests/unittests/test_plan/test_warmup.py
+++ b/src/exo/worker/tests/unittests/test_plan/test_warmup.py
@@ -57,7 +57,6 @@ def test_plan_starts_warmup_for_accepting_rank_when_all_loaded_or_warming():
result = plan_mod.plan(
node_id=NODE_B,
runners=runners, # type: ignore
- download_status={},
global_download_status={NODE_A: []},
instances=instances,
all_runners=all_runners,
@@ -99,7 +98,6 @@ def test_plan_starts_warmup_for_rank_zero_after_others_warming():
result = plan_mod.plan(
node_id=NODE_A,
runners=runners, # type: ignore
- download_status={},
global_download_status={NODE_A: []},
instances=instances,
all_runners=all_runners,
@@ -140,7 +138,6 @@ def test_plan_does_not_start_warmup_for_non_zero_rank_until_all_loaded_or_warmin
result = plan_mod.plan(
node_id=NODE_B,
runners=runners, # type: ignore
- download_status={},
global_download_status={NODE_A: [], NODE_B: []},
instances=instances,
all_runners=all_runners,
@@ -185,7 +182,6 @@ def test_plan_does_not_start_warmup_for_rank_zero_until_others_warming():
result = plan_mod.plan(
node_id=NODE_A,
runners=runners, # type: ignore
- download_status={},
global_download_status={NODE_A: []},
instances=instances,
all_runners=all_runners,
@@ -202,7 +198,6 @@ def test_plan_does_not_start_warmup_for_rank_zero_until_others_warming():
result = plan_mod.plan(
node_id=NODE_A,
runners=runners, # type: ignore
- download_status={},
global_download_status={NODE_A: []},
instances=instances,
all_runners=all_runners,
@@ -246,7 +241,6 @@ def test_plan_starts_warmup_for_connecting_rank_after_others_warming():
result = plan_mod.plan(
node_id=NODE_B,
runners=runners, # type: ignore
- download_status={},
global_download_status={NODE_B: []},
instances=instances,
all_runners=all_runners,
@@ -289,7 +283,6 @@ def test_plan_does_not_start_warmup_for_accepting_rank_until_all_loaded_or_warmi
result = plan_mod.plan(
node_id=NODE_A,
runners=runners, # type: ignore
- download_status={},
global_download_status={NODE_A: [], NODE_B: []},
instances=instances,
all_runners=all_runners,
@@ -331,7 +324,6 @@ def test_plan_does_not_start_warmup_for_connecting_rank_until_others_warming():
result = plan_mod.plan(
node_id=NODE_A,
runners=runners, # type: ignore
- download_status={},
global_download_status={NODE_A: [], NODE_B: []},
instances=instances,
all_runners=all_runners,
diff --git a/tests/headless_runner.py b/tests/headless_runner.py
index 354186a5..30fd19d4 100644
--- a/tests/headless_runner.py
+++ b/tests/headless_runner.py
@@ -11,6 +11,10 @@ from hypercorn.asyncio import serve # pyright: ignore[reportUnknownVariableType
from loguru import logger
from pydantic import BaseModel
+from exo.download.impl_shard_downloader import (
+ build_full_shard,
+ exo_shard_downloader,
+)
from exo.shared.logging import InterceptLogger, logger_setup
from exo.shared.models.model_cards import MODEL_CARDS, ModelId
from exo.shared.types.api import ChatCompletionMessage, ChatCompletionTaskParams
@@ -36,10 +40,6 @@ from exo.shared.types.worker.runners import RunnerId, ShardAssignments
from exo.shared.types.worker.shards import PipelineShardMetadata, TensorShardMetadata
from exo.utils.channels import MpReceiver, MpSender, channel, mp_channel
from exo.utils.info_gatherer.info_gatherer import GatheredInfo, InfoGatherer
-from exo.worker.download.impl_shard_downloader import (
- build_full_shard,
- exo_shard_downloader,
-)
from exo.worker.runner.bootstrap import entrypoint
← ba199408 Fix regenerate for image models (#1263)
·
back to Exo
·
downloads: add download and delete buttons to downloads UI 6dbbe779 →