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:
@@ -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,
|
||||
|
||||
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user