From 62e8110e97aee1330c8ae60d444ee4e71abc3342 Mon Sep 17 00:00:00 2001
From: Alex Cheema <41707476+AlexCheema@users.noreply.github.com>
Date: Wed, 11 Feb 2026 15:50:43 -0800
Subject: [PATCH] fix: prevent DownloadModel TaskCreated event flood (#1452)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

## Motivation

When a model download fails repeatedly (e.g. `ContentLengthError` on a
large model like `zai-org/GLM-5`), the download coordinator accumulates
duplicate progress callbacks — one per retry cycle. Each callback
independently throttles at 1 event/sec, so after N retries, every
download progress tick generates N events instead of 1. After an hour of
failures (~60 retry cycles), this produces ~60 `NodeDownloadProgress`
events/sec, overwhelming the master, delaying heartbeats, and causing
the node to time itself out.

### The callback accumulation cycle
1. `_start_download_task()` calls
`shard_downloader.on_progress(callback)` which **appends** to a list
2. Download fails → `DownloadFailed` status set, but old callback stays
in the list
3. 60s later: `_emit_existing_download_progress()` scans disk → resets
status to `DownloadPending`
4. Worker sends new `StartDownload` → coordinator accepts (guard didn't
check `DownloadFailed`)
5. `_start_download_task()` appends **another** callback
6. Each callback has its own throttle → N callbacks = N events per
progress tick

## Changes

### Commit 1: `src/exo/worker/main.py`
Move the `DownloadModel` backoff check **before** `TaskCreated` emission
in `plan_step()`. Previously `TaskCreated` was emitted unconditionally
every 0.1s even when backoff blocked the download command.

### Commit 2: `src/exo/download/coordinator.py`
1. **Register progress callback once** in `__post_init__` instead of
per-download in `_start_download_task()`. Uses a per-model throttle dict
instead of per-callback closure variables.
2. **Add `DownloadFailed` to the `_start_download()` guard** so
redundant `_start_download_task()` calls don't happen. Retries still
work because `_emit_existing_download_progress` resets `DownloadFailed`
→ `DownloadPending` by scanning disk every 60s.

## Why It Works

The root cause was callbacks accumulating in
`ResumableShardDownloader.on_progress_callbacks` (a list that only
appends, never clears). By registering one callback per coordinator
lifetime and guarding against re-entry on `DownloadFailed`, we ensure
exactly one progress event per model per progress tick regardless of how
many retry cycles have occurred.

## Test Plan

### Manual Testing
- Verified the download retry flow: failed download → 60s scan resets
status → new `StartDownload` accepted → download retries with single
callback

### Automated Testing
- `uv run basedpyright` — 0 errors
- `uv run ruff check` — passes
- `uv run pytest` — 188 passed

---------

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
---
 src/exo/download/coordinator.py | 93 ++++++++++++++++++---------------
 src/exo/worker/main.py          | 11 ++--
 2 files changed, 58 insertions(+), 46 deletions(-)

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(