← back to Exo
Emit download progress on start, and change downloads to be keyed by model_id (#1044)
1ec550dff1d2901e82a8538b67741d867d44b6a3 · 2025-12-31 01:18:10 +0000 · Alex Cheema
## Motivation
We added a download page to the dashboard which shows the currently
download status of each model on each node. Users have reported this to
be extremely useful.
However, we don't currently fetch the download progress on start, so it
doesn't show any model's download status.
## Changes
Fetch and emit model download status on start of worker, and
periodically every 5 mins.
Also to support this, I changed download_status to be keyed by model_id
instead of shard, since we want download_status of each model, not each
shard.
## Why It Works
The dashboard already implements the correct functionality, we just
weren't populating the download status in the state. Now it gets
populated and shows correctly.
## Test Plan
### Manual Testing
On a cluster of 2 x 512GB M3 Ultra Mac Studio, I launched an instance
onto one node that hadn't been downloaded. I checked the download page
and it showed the in progress download. I downloaded it to completion,
restarted exo on both nodes, and then opened the download page and it
showed the model as 100% downloaded and other models as 0% that hadn't
been downloaded.
---------
Co-authored-by: Evan <evanev7@gmail.com>
Files touched
M src/exo/worker/download/download_utils.pyM src/exo/worker/main.pyM src/exo/worker/plan.pyM src/exo/worker/tests/unittests/test_plan/test_download_and_loading.py
Diff
commit 1ec550dff1d2901e82a8538b67741d867d44b6a3
Author: Alex Cheema <41707476+AlexCheema@users.noreply.github.com>
Date: Wed Dec 31 01:18:10 2025 +0000
Emit download progress on start, and change downloads to be keyed by model_id (#1044)
## Motivation
We added a download page to the dashboard which shows the currently
download status of each model on each node. Users have reported this to
be extremely useful.
However, we don't currently fetch the download progress on start, so it
doesn't show any model's download status.
## Changes
Fetch and emit model download status on start of worker, and
periodically every 5 mins.
Also to support this, I changed download_status to be keyed by model_id
instead of shard, since we want download_status of each model, not each
shard.
## Why It Works
The dashboard already implements the correct functionality, we just
weren't populating the download status in the state. Now it gets
populated and shows correctly.
## Test Plan
### Manual Testing
On a cluster of 2 x 512GB M3 Ultra Mac Studio, I launched an instance
onto one node that hadn't been downloaded. I checked the download page
and it showed the in progress download. I downloaded it to completion,
restarted exo on both nodes, and then opened the download page and it
showed the model as 100% downloaded and other models as 0% that hadn't
been downloaded.
---------
Co-authored-by: Evan <evanev7@gmail.com>
---
src/exo/worker/download/download_utils.py | 5 ++
src/exo/worker/main.py | 53 +++++++++++++++++++---
src/exo/worker/plan.py | 13 +++---
.../test_plan/test_download_and_loading.py | 10 ++--
4 files changed, 63 insertions(+), 18 deletions(-)
diff --git a/src/exo/worker/download/download_utils.py b/src/exo/worker/download/download_utils.py
index 8d033471..a16311f3 100644
--- a/src/exo/worker/download/download_utils.py
+++ b/src/exo/worker/download/download_utils.py
@@ -450,6 +450,11 @@ async def get_weight_map(repo_id: str, revision: str = "main") -> dict[str, str]
async def resolve_allow_patterns(shard: ShardMetadata) -> list[str]:
+ # TODO: 'Smart' downloads are disabled because:
+ # (i) We don't handle all kinds of files;
+ # (ii) We don't have sticky sessions.
+ # (iii) Tensor parallel requires all files.
+ return ["*"]
try:
weight_map = await get_weight_map(str(shard.model_meta.model_id))
return get_allow_patterns(weight_map, shard)
diff --git a/src/exo/worker/main.py b/src/exo/worker/main.py
index a9560853..e10df75e 100644
--- a/src/exo/worker/main.py
+++ b/src/exo/worker/main.py
@@ -23,6 +23,7 @@ from exo.shared.types.events import (
TopologyEdgeCreated,
TopologyEdgeDeleted,
)
+from exo.shared.types.models import ModelId
from exo.shared.types.multiaddr import Multiaddr
from exo.shared.types.profiling import MemoryPerformanceProfile, NodePerformanceProfile
from exo.shared.types.state import State
@@ -83,7 +84,7 @@ class Worker:
self.out_for_delivery: dict[EventId, ForwarderEvent] = {}
self.state: State = State()
- self.download_status: dict[ShardMetadata, DownloadProgress] = {}
+ self.download_status: dict[ModelId, DownloadProgress] = {}
self.runners: dict[RunnerId, RunnerSupervisor] = {}
self._tg: TaskGroup | None = None
@@ -128,6 +129,7 @@ class Worker:
tg.start_soon(start_polling_node_metrics, resource_monitor_callback)
tg.start_soon(start_polling_memory_metrics, memory_monitor_callback)
+ tg.start_soon(self._emit_existing_download_progress)
tg.start_soon(self._connection_message_event_writer)
tg.start_soon(self._resend_out_for_delivery)
tg.start_soon(self._event_applier)
@@ -200,11 +202,11 @@ class Worker:
)
)
case DownloadModel(shard_metadata=shard):
- if shard not in self.download_status:
+ if shard.model_meta.model_id not in self.download_status:
progress = DownloadPending(
shard_metadata=shard, node_id=self.node_id
)
- self.download_status[shard] = progress
+ self.download_status[shard.model_meta.model_id] = progress
await self.event_sender.send(
NodeDownloadProgress(download_progress=progress)
)
@@ -217,7 +219,7 @@ class Worker:
progress = DownloadCompleted(
shard_metadata=shard, node_id=self.node_id
)
- self.download_status[shard] = progress
+ self.download_status[shard.model_meta.model_id] = progress
await self.event_sender.send(
NodeDownloadProgress(download_progress=progress)
)
@@ -349,7 +351,7 @@ class Worker:
initial_progress
),
)
- self.download_status[task.shard_metadata] = status
+ self.download_status[task.shard_metadata.model_meta.model_id] = status
self.event_sender.send_nowait(NodeDownloadProgress(download_progress=status))
last_progress_time = 0.0
@@ -363,7 +365,7 @@ class Worker:
nonlocal last_progress_time
if progress.status == "complete":
status = DownloadCompleted(shard_metadata=shard, node_id=self.node_id)
- self.download_status[shard] = status
+ self.download_status[shard.model_meta.model_id] = status
# Footgun!
self.event_sender.send_nowait(
NodeDownloadProgress(download_progress=status)
@@ -384,7 +386,7 @@ class Worker:
progress
),
)
- self.download_status[shard] = status
+ self.download_status[shard.model_meta.model_id] = status
self.event_sender.send_nowait(
NodeDownloadProgress(download_progress=status)
)
@@ -444,3 +446,40 @@ class Worker:
await self.event_sender.send(TopologyEdgeDeleted(edge=conn))
await anyio.sleep(10)
+
+ async def _emit_existing_download_progress(self) -> None:
+ try:
+ while True:
+ logger.info("Fetching and emitting existing download progress...")
+ async for (
+ _,
+ progress,
+ ) in self.shard_downloader.get_shard_download_status():
+ if progress.status == "complete":
+ status = DownloadCompleted(
+ node_id=self.node_id, shard_metadata=progress.shard
+ )
+ elif progress.status in ["in_progress", "not_started"]:
+ if progress.downloaded_bytes_this_session.in_bytes == 0:
+ status = DownloadPending(
+ node_id=self.node_id, shard_metadata=progress.shard
+ )
+ else:
+ status = DownloadOngoing(
+ node_id=self.node_id,
+ shard_metadata=progress.shard,
+ download_progress=map_repo_download_progress_to_download_progress_data(
+ progress
+ ),
+ )
+ else:
+ continue
+
+ self.download_status[progress.shard.model_meta.model_id] = status
+ await self.event_sender.send(
+ NodeDownloadProgress(download_progress=status)
+ )
+ logger.info("Done emitting existing download progress.")
+ await anyio.sleep(5 * 60) # 5 minutes
+ except Exception as e:
+ logger.error(f"Error emitting existing download progress: {e}")
diff --git a/src/exo/worker/plan.py b/src/exo/worker/plan.py
index 6c621e72..d72f2e2d 100644
--- a/src/exo/worker/plan.py
+++ b/src/exo/worker/plan.py
@@ -3,6 +3,7 @@
from collections.abc import Mapping, Sequence
from exo.shared.types.common import NodeId
+from exo.shared.types.models import ModelId
from exo.shared.types.tasks import (
ChatCompletion,
ConnectToGroup,
@@ -34,7 +35,6 @@ from exo.shared.types.worker.runners import (
RunnerStatus,
RunnerWarmingUp,
)
-from exo.shared.types.worker.shards import ShardMetadata
from exo.worker.runner.runner_supervisor import RunnerSupervisor
@@ -43,7 +43,7 @@ def plan(
# Runners is expected to be FRESH and so should not come from state
runners: Mapping[RunnerId, RunnerSupervisor],
# DL_status is expected to be FRESH and so should not come from state
- download_status: Mapping[ShardMetadata, DownloadProgress],
+ download_status: Mapping[ModelId, DownloadProgress],
# gdls is not expected to be fresh
global_download_status: Mapping[NodeId, Sequence[DownloadProgress]],
instances: Mapping[InstanceId, Instance],
@@ -111,13 +111,14 @@ def _create_runner(
def _model_needs_download(
runners: Mapping[RunnerId, RunnerSupervisor],
- download_status: Mapping[ShardMetadata, DownloadProgress],
+ download_status: Mapping[ModelId, DownloadProgress],
) -> DownloadModel | None:
for runner in runners.values():
+ model_id = runner.bound_instance.bound_shard.model_meta.model_id
if isinstance(runner.status, RunnerIdle) and (
- not isinstance(
- download_status.get(runner.bound_instance.bound_shard, None),
- (DownloadOngoing, DownloadCompleted),
+ model_id not in download_status
+ or not isinstance(
+ download_status[model_id], (DownloadOngoing, DownloadCompleted)
)
):
# We don't invalidate download_status randomly in case a file gets deleted on disk
diff --git a/src/exo/worker/tests/unittests/test_plan/test_download_and_loading.py b/src/exo/worker/tests/unittests/test_plan/test_download_and_loading.py
index 49d3b05c..cedee493 100644
--- a/src/exo/worker/tests/unittests/test_plan/test_download_and_loading.py
+++ b/src/exo/worker/tests/unittests/test_plan/test_download_and_loading.py
@@ -1,5 +1,6 @@
import exo.worker.plan as plan_mod
from exo.shared.types.common import NodeId
+from exo.shared.types.models import ModelId
from exo.shared.types.tasks import LoadModel
from exo.shared.types.worker.downloads import DownloadCompleted, DownloadProgress
from exo.shared.types.worker.instances import BoundInstance
@@ -7,7 +8,6 @@ from exo.shared.types.worker.runners import (
RunnerConnected,
RunnerIdle,
)
-from exo.shared.types.worker.shards import ShardMetadata
from exo.worker.tests.constants import (
INSTANCE_1_ID,
MODEL_A_ID,
@@ -46,7 +46,7 @@ def test_plan_requests_download_when_waiting_and_shard_not_downloaded():
all_runners = {RUNNER_1_ID: RunnerIdle()}
# No entry for this shard -> should trigger DownloadModel
- download_status: dict[ShardMetadata, DownloadProgress] = {}
+ download_status: dict[ModelId, DownloadProgress] = {}
result = plan_mod.plan(
node_id=NODE_A,
@@ -94,7 +94,7 @@ def test_plan_loads_model_when_all_shards_downloaded_and_waiting():
# Local node has already marked its shard as downloaded (not actually used by _load_model)
local_download_status = {
- shard1: DownloadCompleted(shard_metadata=shard1, node_id=NODE_A) # type: ignore[reportUnhashable]
+ MODEL_A_ID: DownloadCompleted(shard_metadata=shard1, node_id=NODE_A)
}
# Global view has completed downloads for both nodes
@@ -140,7 +140,7 @@ def test_plan_does_not_request_download_when_shard_already_downloaded():
# Local status claims the shard is downloaded already
local_download_status = {
- shard: DownloadCompleted(shard_metadata=shard, node_id=NODE_A) # type: ignore[reportUnhashable]
+ MODEL_A_ID: DownloadCompleted(shard_metadata=shard, node_id=NODE_A)
}
# Global view hasn't caught up yet (no completed shards recorded for NODE_A)
@@ -192,7 +192,7 @@ def test_plan_does_not_load_model_until_all_shards_downloaded_globally():
# Only NODE_A's shard is recorded as downloaded globally
local_download_status = {
- shard1: DownloadCompleted(shard_metadata=shard1, node_id=NODE_A) # type: ignore[reportUnhashable]
+ MODEL_A_ID: DownloadCompleted(shard_metadata=shard1, node_id=NODE_A)
}
global_download_status = {
NODE_A: [DownloadCompleted(shard_metadata=shard1, node_id=NODE_A)],
← 283c0e39 Placement filters for tensor parallel supports_tensor, tens
·
back to Exo
·
Fix tests broken by 283c (#1063) c1be5184 →