fix: prevent DownloadModel TaskCreated event flood (#1452)

## 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>
This commit is contained in:
Alex Cheema
2026-02-11 15:50:43 -08:00
committed by GitHub
parent 98773437f3
commit 62e8110e97
2 changed files with 58 additions and 46 deletions
+50 -43
View File
@@ -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,
+8 -3
View File
@@ -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(