[object Object]

← back to Exo

shard_downloader: make on_progress callback async

fb0151630def96323bb95215d43de1f599a9dea7 · 2026-01-15 13:15:18 +0000 · Jake Hillion

The on_progress callback was synchronous but always invoked from async
contexts, forcing the use of send_nowait() which could raise WouldBlock
if the channel buffer was full, potentially dropping progress updates.

Changed the callback type from `Callable[[ShardMetadata,
RepoDownloadProgress], None]` to return a coroutine, updated all
implementations to be async, and replaced send_nowait() with await
send() in the worker's download progress handler.

This allows proper backpressure handling when sending download progress
events through the channel, eliminating the "Footgun!" that was
previously documented in the code.

Test plan:
- Built a DMG and ran it on one node. All existing models showed as
  downloaded.
- Downloaded a new model. The progress bar on the download page worked.
- Downloaded another new model. The progress bar on the home page
  worked.

Files touched

Diff

commit fb0151630def96323bb95215d43de1f599a9dea7
Author: Jake Hillion <jake@hillion.co.uk>
Date:   Thu Jan 15 13:15:18 2026 +0000

    shard_downloader: make on_progress callback async
    
    The on_progress callback was synchronous but always invoked from async
    contexts, forcing the use of send_nowait() which could raise WouldBlock
    if the channel buffer was full, potentially dropping progress updates.
    
    Changed the callback type from `Callable[[ShardMetadata,
    RepoDownloadProgress], None]` to return a coroutine, updated all
    implementations to be async, and replaced send_nowait() with await
    send() in the worker's download progress handler.
    
    This allows proper backpressure handling when sending download progress
    events through the channel, eliminating the "Footgun!" that was
    previously documented in the code.
    
    Test plan:
    - Built a DMG and ran it on one node. All existing models showed as
      downloaded.
    - Downloaded a new model. The progress bar on the download page worked.
    - Downloaded another new model. The progress bar on the home page
      worked.
---
 src/exo/worker/download/download_utils.py        | 22 +++++++++++++++-------
 src/exo/worker/download/impl_shard_downloader.py | 16 ++++++++++------
 src/exo/worker/download/shard_downloader.py      |  7 +++++--
 src/exo/worker/main.py                           | 10 ++++------
 4 files changed, 34 insertions(+), 21 deletions(-)

diff --git a/src/exo/worker/download/download_utils.py b/src/exo/worker/download/download_utils.py
index b672a936..05d2ff73 100644
--- a/src/exo/worker/download/download_utils.py
+++ b/src/exo/worker/download/download_utils.py
@@ -5,6 +5,7 @@ import shutil
 import ssl
 import time
 import traceback
+from collections.abc import Awaitable
 from datetime import timedelta
 from pathlib import Path
 from typing import Callable, Literal
@@ -525,7 +526,7 @@ async def download_progress_for_local_path(
 
 async def download_shard(
     shard: ShardMetadata,
-    on_progress: Callable[[ShardMetadata, RepoDownloadProgress], None],
+    on_progress: Callable[[ShardMetadata, RepoDownloadProgress], Awaitable[None]],
     max_parallel_downloads: int = 8,
     skip_download: bool = False,
     allow_patterns: list[str] | None = None,
@@ -566,9 +567,9 @@ async def download_shard(
     )
     file_progress: dict[str, RepoFileDownloadProgress] = {}
 
-    def on_progress_wrapper(
+    async def on_progress_wrapper(
         file: FileListEntry, curr_bytes: int, total_bytes: int, is_renamed: bool
-    ):
+    ) -> None:
         start_time = (
             file_progress[file.path].start_time
             if file.path in file_progress
@@ -604,7 +605,7 @@ async def download_shard(
             else "in_progress",
             start_time=start_time,
         )
-        on_progress(
+        await on_progress(
             shard,
             calculate_repo_progress(
                 shard,
@@ -632,14 +633,21 @@ async def download_shard(
 
     semaphore = asyncio.Semaphore(max_parallel_downloads)
 
-    async def download_with_semaphore(file: FileListEntry):
+    def schedule_progress(
+        file: FileListEntry, curr_bytes: int, total_bytes: int, is_renamed: bool
+    ) -> None:
+        asyncio.create_task(
+            on_progress_wrapper(file, curr_bytes, total_bytes, is_renamed)
+        )
+
+    async def download_with_semaphore(file: FileListEntry) -> None:
         async with semaphore:
             await download_file_with_retry(
                 str(shard.model_meta.model_id),
                 revision,
                 file.path,
                 target_dir,
-                lambda curr_bytes, total_bytes, is_renamed: on_progress_wrapper(
+                lambda curr_bytes, total_bytes, is_renamed: schedule_progress(
                     file, curr_bytes, total_bytes, is_renamed
                 ),
             )
@@ -651,7 +659,7 @@ async def download_shard(
     final_repo_progress = calculate_repo_progress(
         shard, str(shard.model_meta.model_id), revision, file_progress, all_start_time
     )
-    on_progress(shard, final_repo_progress)
+    await on_progress(shard, final_repo_progress)
     if gguf := next((f for f in filtered_file_list if f.path.endswith(".gguf")), None):
         return target_dir / gguf.path, final_repo_progress
     else:
diff --git a/src/exo/worker/download/impl_shard_downloader.py b/src/exo/worker/download/impl_shard_downloader.py
index 46f55ff9..1ca2e4e7 100644
--- a/src/exo/worker/download/impl_shard_downloader.py
+++ b/src/exo/worker/download/impl_shard_downloader.py
@@ -1,4 +1,5 @@
 import asyncio
+from collections.abc import Awaitable
 from pathlib import Path
 from typing import AsyncIterator, Callable
 
@@ -48,7 +49,8 @@ class SingletonShardDownloader(ShardDownloader):
         self.active_downloads: dict[ShardMetadata, asyncio.Task[Path]] = {}
 
     def on_progress(
-        self, callback: Callable[[ShardMetadata, RepoDownloadProgress], None]
+        self,
+        callback: Callable[[ShardMetadata, RepoDownloadProgress], Awaitable[None]],
     ) -> None:
         self.shard_downloader.on_progress(callback)
 
@@ -83,7 +85,8 @@ class CachedShardDownloader(ShardDownloader):
         self.cache: dict[tuple[str, ShardMetadata], Path] = {}
 
     def on_progress(
-        self, callback: Callable[[ShardMetadata, RepoDownloadProgress], None]
+        self,
+        callback: Callable[[ShardMetadata, RepoDownloadProgress], Awaitable[None]],
     ) -> None:
         self.shard_downloader.on_progress(callback)
 
@@ -113,17 +116,18 @@ class ResumableShardDownloader(ShardDownloader):
     def __init__(self, max_parallel_downloads: int = 8):
         self.max_parallel_downloads = max_parallel_downloads
         self.on_progress_callbacks: list[
-            Callable[[ShardMetadata, RepoDownloadProgress], None]
+            Callable[[ShardMetadata, RepoDownloadProgress], Awaitable[None]]
         ] = []
 
-    def on_progress_wrapper(
+    async def on_progress_wrapper(
         self, shard: ShardMetadata, progress: RepoDownloadProgress
     ) -> None:
         for callback in self.on_progress_callbacks:
-            callback(shard, progress)
+            await callback(shard, progress)
 
     def on_progress(
-        self, callback: Callable[[ShardMetadata, RepoDownloadProgress], None]
+        self,
+        callback: Callable[[ShardMetadata, RepoDownloadProgress], Awaitable[None]],
     ) -> None:
         self.on_progress_callbacks.append(callback)
 
diff --git a/src/exo/worker/download/shard_downloader.py b/src/exo/worker/download/shard_downloader.py
index dcc94a31..60a4692e 100644
--- a/src/exo/worker/download/shard_downloader.py
+++ b/src/exo/worker/download/shard_downloader.py
@@ -1,4 +1,5 @@
 from abc import ABC, abstractmethod
+from collections.abc import Awaitable
 from copy import copy
 from datetime import timedelta
 from pathlib import Path
@@ -31,7 +32,8 @@ class ShardDownloader(ABC):
 
     @abstractmethod
     def on_progress(
-        self, callback: Callable[[ShardMetadata, RepoDownloadProgress], None]
+        self,
+        callback: Callable[[ShardMetadata, RepoDownloadProgress], Awaitable[None]],
     ) -> None:
         pass
 
@@ -59,7 +61,8 @@ class NoopShardDownloader(ShardDownloader):
         return Path("/tmp/noop_shard")
 
     def on_progress(
-        self, callback: Callable[[ShardMetadata, RepoDownloadProgress], None]
+        self,
+        callback: Callable[[ShardMetadata, RepoDownloadProgress], Awaitable[None]],
     ) -> None:
         pass
 
diff --git a/src/exo/worker/main.py b/src/exo/worker/main.py
index 28b89e01..66ed0282 100644
--- a/src/exo/worker/main.py
+++ b/src/exo/worker/main.py
@@ -359,8 +359,7 @@ class Worker:
         last_progress_time = 0.0
         throttle_interval_secs = 1.0
 
-        # TODO: i hate callbacks
-        def download_progress_callback(
+        async def download_progress_callback(
             shard: ShardMetadata, progress: RepoDownloadProgress
         ) -> None:
             nonlocal self
@@ -372,11 +371,10 @@ class Worker:
                     total_bytes=progress.total_bytes,
                 )
                 self.download_status[shard.model_meta.model_id] = status
-                # Footgun!
-                self.event_sender.send_nowait(
+                await self.event_sender.send(
                     NodeDownloadProgress(download_progress=status)
                 )
-                self.event_sender.send_nowait(
+                await self.event_sender.send(
                     TaskStatusUpdated(
                         task_id=task.task_id, task_status=TaskStatus.Complete
                     )
@@ -393,7 +391,7 @@ class Worker:
                     ),
                 )
                 self.download_status[shard.model_meta.model_id] = status
-                self.event_sender.send_nowait(
+                await self.event_sender.send(
                     NodeDownloadProgress(download_progress=status)
                 )
                 last_progress_time = current_time()

← 346b13e2 Enhance LaTeX rendering in dashboard markdown (#1197)  ·  back to Exo  ·  Update README.md with some changes from release 1.0.61 (#115 ce3ad391 →