diff --git a/src/exo/download/coordinator.py b/src/exo/download/coordinator.py index b209c697..a05bd6f8 100644 --- a/src/exo/download/coordinator.py +++ b/src/exo/download/coordinator.py @@ -56,8 +56,49 @@ class DownloadCoordinator: event_receiver: Receiver[Event] = field(init=False) _tg: TaskGroup = field(init=False, default_factory=anyio.create_task_group) + # Per-model throttle for download progress events + _last_progress_time: dict[ModelId, float] = field(default_factory=dict) + def __post_init__(self) -> None: self.event_sender, self.event_receiver = channel[Event]() + self.shard_downloader.on_progress(self._download_progress_callback) + + async def _download_progress_callback( + self, callback_shard: ShardMetadata, progress: RepoDownloadProgress + ) -> None: + model_id = callback_shard.model_card.model_id + throttle_interval_secs = 1.0 + + if progress.status == "complete": + completed = DownloadCompleted( + shard_metadata=callback_shard, + node_id=self.node_id, + total_bytes=progress.total_bytes, + ) + self.download_status[model_id] = completed + await self.event_sender.send( + NodeDownloadProgress(download_progress=completed) + ) + if model_id in self.active_downloads: + del self.active_downloads[model_id] + self._last_progress_time.pop(model_id, None) + elif ( + progress.status == "in_progress" + and current_time() - self._last_progress_time.get(model_id, 0.0) + > 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[model_id] = ongoing + await self.event_sender.send( + NodeDownloadProgress(download_progress=ongoing) + ) + self._last_progress_time[model_id] = current_time() async def run(self) -> None: logger.info("Starting DownloadCoordinator") @@ -119,12 +160,12 @@ class DownloadCoordinator: async def _start_download(self, shard: ShardMetadata) -> None: model_id = shard.model_card.model_id - # Check if already downloading or complete + # Check if already downloading, complete, or recently failed if model_id in self.download_status: status = self.download_status[model_id] - if isinstance(status, (DownloadOngoing, DownloadCompleted)): + if isinstance(status, (DownloadOngoing, DownloadCompleted, DownloadFailed)): logger.debug( - f"Download for {model_id} already in progress or complete, skipping" + f"Download for {model_id} already in progress, complete, or failed, skipping" ) return @@ -169,46 +210,6 @@ class DownloadCoordinator: 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) @@ -283,6 +284,12 @@ class DownloadCoordinator: _, progress, ) in self.shard_downloader.get_shard_download_status(): + model_id = progress.shard.model_card.model_id + + # Active downloads emit progress via the callback — don't overwrite + if model_id in self.active_downloads: + continue + if progress.status == "complete": status: DownloadProgress = DownloadCompleted( node_id=self.node_id, diff --git a/src/exo/worker/main.py b/src/exo/worker/main.py index 4f059bfb..4abd3c43 100644 --- a/src/exo/worker/main.py +++ b/src/exo/worker/main.py @@ -184,6 +184,14 @@ class Worker: ) 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): + continue + logger.info(f"Worker plan: {task.__class__.__name__}") assert task.task_status await self.event_sender.send(TaskCreated(task_id=task.task_id, task=task)) @@ -199,9 +207,6 @@ class Worker: ) case DownloadModel(shard_metadata=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(