← back to Exo
Fix download speed/ETA display for re-downloads (#1294)
bd4f0bf04813fe01ef0bffbd81cc26625241e6e6 · 2026-01-26 13:56:58 -0800 · Alex Cheema
## Motivation
After the download verification fix, when files are re-downloaded due to
upstream changes (size mismatch), the download progress displays
correctly (completion %, bytes, file counts), but speed shows 0 B/s and
ETA shows "--" for both overall and per-file progress.
## Changes
- Modified `on_progress_wrapper` in `src/exo/download/download_utils.py`
to detect re-download scenarios
- Added re-download detection: when `curr_bytes < previous_downloaded`,
the file was deleted and download restarted
- On re-download: reset `start_time` to current time and set
`downloaded_this_session = curr_bytes`
- Added two tests to `test_download_verification.py` covering
re-download and continuing download scenarios
## Why It Works
The bug occurred because:
1. `file_progress` is initialized with the OLD local file size (e.g.,
1.5GB)
2. When `_download_file` detects size mismatch, it deletes the file and
starts fresh
3. Progress callback receives small `curr_bytes` (e.g., 8KB) but
compares against old size
4. `downloaded_this_session = 0 + (8KB - 1.5GB) = -1.5GB` (negative!)
5. Negative session bytes → 0 or negative speed → ETA shows "--"
The fix detects when `curr_bytes < previous_downloaded` (indicating
re-download started) and resets tracking to treat it as a fresh
download.
## Test Plan
### Manual Testing
<!-- Hardware: (e.g., MacBook Pro M1 Max 32GB, Mac Mini M2 16GB,
connected via Thunderbolt 4) -->
<!-- What you did: -->
- Download a model, modify a file to change its size, restart exo,
verify speed/ETA display correctly during re-download
### Automated Testing
- Added `TestProgressResetOnRedownload` class with two tests:
- `test_progress_resets_correctly_on_redownload`: Verifies progress
resets correctly when re-download starts
- `test_progress_accumulates_on_continuing_download`: Verifies
continuing downloads still accumulate correctly
- All 11 download tests pass
- Type checking (basedpyright): 0 errors
- Linting (ruff): All checks passed
🤖 Generated with [Claude Code](https://claude.com/claude-code)
Co-authored-by: Claude Opus 4.5 <noreply@anthropic.com>
Files touched
M src/exo/download/download_utils.pyM src/exo/download/tests/test_download_verification.py
Diff
commit bd4f0bf04813fe01ef0bffbd81cc26625241e6e6
Author: Alex Cheema <41707476+AlexCheema@users.noreply.github.com>
Date: Mon Jan 26 13:56:58 2026 -0800
Fix download speed/ETA display for re-downloads (#1294)
## Motivation
After the download verification fix, when files are re-downloaded due to
upstream changes (size mismatch), the download progress displays
correctly (completion %, bytes, file counts), but speed shows 0 B/s and
ETA shows "--" for both overall and per-file progress.
## Changes
- Modified `on_progress_wrapper` in `src/exo/download/download_utils.py`
to detect re-download scenarios
- Added re-download detection: when `curr_bytes < previous_downloaded`,
the file was deleted and download restarted
- On re-download: reset `start_time` to current time and set
`downloaded_this_session = curr_bytes`
- Added two tests to `test_download_verification.py` covering
re-download and continuing download scenarios
## Why It Works
The bug occurred because:
1. `file_progress` is initialized with the OLD local file size (e.g.,
1.5GB)
2. When `_download_file` detects size mismatch, it deletes the file and
starts fresh
3. Progress callback receives small `curr_bytes` (e.g., 8KB) but
compares against old size
4. `downloaded_this_session = 0 + (8KB - 1.5GB) = -1.5GB` (negative!)
5. Negative session bytes → 0 or negative speed → ETA shows "--"
The fix detects when `curr_bytes < previous_downloaded` (indicating
re-download started) and resets tracking to treat it as a fresh
download.
## Test Plan
### Manual Testing
<!-- Hardware: (e.g., MacBook Pro M1 Max 32GB, Mac Mini M2 16GB,
connected via Thunderbolt 4) -->
<!-- What you did: -->
- Download a model, modify a file to change its size, restart exo,
verify speed/ETA display correctly during re-download
### Automated Testing
- Added `TestProgressResetOnRedownload` class with two tests:
- `test_progress_resets_correctly_on_redownload`: Verifies progress
resets correctly when re-download starts
- `test_progress_accumulates_on_continuing_download`: Verifies
continuing downloads still accumulate correctly
- All 11 download tests pass
- Type checking (basedpyright): 0 errors
- Linting (ruff): All checks passed
🤖 Generated with [Claude Code](https://claude.com/claude-code)
Co-authored-by: Claude Opus 4.5 <noreply@anthropic.com>
---
src/exo/download/download_utils.py | 29 +++--
.../download/tests/test_download_verification.py | 132 ++++++++++++++++++++-
2 files changed, 150 insertions(+), 11 deletions(-)
diff --git a/src/exo/download/download_utils.py b/src/exo/download/download_utils.py
index f10b95e8..6dec4718 100644
--- a/src/exo/download/download_utils.py
+++ b/src/exo/download/download_utils.py
@@ -583,17 +583,26 @@ async def download_shard(
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
- else time.time()
- )
- downloaded_this_session = (
- file_progress[file.path].downloaded_this_session.in_bytes
- + (curr_bytes - file_progress[file.path].downloaded.in_bytes)
- if file.path in file_progress
- else curr_bytes
+ previous_progress = file_progress.get(file.path)
+
+ # Detect re-download: curr_bytes < previous downloaded means file was deleted and restarted
+ is_redownload = (
+ previous_progress is not None
+ and curr_bytes < previous_progress.downloaded.in_bytes
)
+
+ if is_redownload or previous_progress is None:
+ # Fresh download or re-download: reset tracking
+ start_time = time.time()
+ downloaded_this_session = curr_bytes
+ else:
+ # Continuing download: accumulate
+ start_time = previous_progress.start_time
+ downloaded_this_session = (
+ previous_progress.downloaded_this_session.in_bytes
+ + (curr_bytes - previous_progress.downloaded.in_bytes)
+ )
+
speed = (
downloaded_this_session / (time.time() - start_time)
if time.time() - start_time > 0
diff --git a/src/exo/download/tests/test_download_verification.py b/src/exo/download/tests/test_download_verification.py
index 80bb6dc2..2d8d076d 100644
--- a/src/exo/download/tests/test_download_verification.py
+++ b/src/exo/download/tests/test_download_verification.py
@@ -1,6 +1,8 @@
"""Tests for download verification and cache behavior."""
+import time
from collections.abc import AsyncIterator
+from datetime import timedelta
from pathlib import Path
from unittest.mock import AsyncMock, MagicMock, patch
@@ -14,7 +16,8 @@ from exo.download.download_utils import (
fetch_file_list_with_cache,
)
from exo.shared.types.common import ModelId
-from exo.shared.types.worker.downloads import FileListEntry
+from exo.shared.types.memory import Memory
+from exo.shared.types.worker.downloads import FileListEntry, RepoFileDownloadProgress
@pytest.fixture
@@ -319,3 +322,130 @@ class TestModelDeletion:
result = await delete_model(model_id)
assert result is False
+
+
+class TestProgressResetOnRedownload:
+ """Tests for progress tracking when files are re-downloaded."""
+
+ async def test_progress_resets_correctly_on_redownload(
+ self, model_id: ModelId
+ ) -> None:
+ """Test that progress tracking resets when a file is re-downloaded.
+
+ When a file is deleted and re-downloaded (due to size mismatch),
+ the progress tracking should reset rather than calculating negative
+ downloaded_this_session values.
+ """
+ # Simulate file_progress dict as it exists in download_shard
+ file_progress: dict[str, RepoFileDownloadProgress] = {}
+
+ # Initialize with old file progress (simulating existing large file)
+ old_file_size = 1_500_000_000 # 1.5 GB
+ file_progress["model.safetensors"] = RepoFileDownloadProgress(
+ repo_id=model_id,
+ repo_revision="main",
+ file_path="model.safetensors",
+ downloaded=Memory.from_bytes(old_file_size),
+ downloaded_this_session=Memory.from_bytes(0),
+ total=Memory.from_bytes(old_file_size),
+ speed=0,
+ eta=timedelta(0),
+ status="not_started",
+ start_time=time.time() - 10, # Started 10 seconds ago
+ )
+
+ # Simulate the logic from on_progress_wrapper after re-download starts
+ # This is the exact logic from the fixed on_progress_wrapper
+ curr_bytes = 100_000 # 100 KB - new download just started
+ previous_progress = file_progress.get("model.safetensors")
+
+ # Detect re-download: curr_bytes < previous downloaded
+ is_redownload = (
+ previous_progress is not None
+ and curr_bytes < previous_progress.downloaded.in_bytes
+ )
+
+ if is_redownload or previous_progress is None:
+ # Fresh download or re-download: reset tracking
+ start_time = time.time()
+ downloaded_this_session = curr_bytes
+ else:
+ # Continuing download: accumulate
+ start_time = previous_progress.start_time
+ downloaded_this_session = (
+ previous_progress.downloaded_this_session.in_bytes
+ + (curr_bytes - previous_progress.downloaded.in_bytes)
+ )
+
+ # Key assertions
+ assert is_redownload is True, "Should detect re-download scenario"
+ assert downloaded_this_session == curr_bytes, (
+ "downloaded_this_session should equal curr_bytes on re-download"
+ )
+ assert downloaded_this_session > 0, (
+ "downloaded_this_session should be positive, not negative"
+ )
+
+ # Calculate speed (should be positive)
+ elapsed = time.time() - start_time
+ speed = downloaded_this_session / elapsed if elapsed > 0 else 0
+ assert speed >= 0, "Speed should be non-negative"
+
+ async def test_progress_accumulates_on_continuing_download(
+ self, model_id: ModelId
+ ) -> None:
+ """Test that progress accumulates correctly for continuing downloads.
+
+ When a download continues from where it left off (resume),
+ the progress should accumulate correctly.
+ """
+ file_progress: dict[str, RepoFileDownloadProgress] = {}
+
+ # Initialize with partial download progress
+ initial_downloaded = 500_000 # 500 KB already downloaded
+ start_time = time.time() - 5 # Started 5 seconds ago
+ file_progress["model.safetensors"] = RepoFileDownloadProgress(
+ repo_id=model_id,
+ repo_revision="main",
+ file_path="model.safetensors",
+ downloaded=Memory.from_bytes(initial_downloaded),
+ downloaded_this_session=Memory.from_bytes(initial_downloaded),
+ total=Memory.from_bytes(1_000_000),
+ speed=100_000,
+ eta=timedelta(seconds=5),
+ status="in_progress",
+ start_time=start_time,
+ )
+
+ # Progress callback with more bytes downloaded
+ curr_bytes = 600_000 # 600 KB - continuing download
+ previous_progress = file_progress.get("model.safetensors")
+
+ # This is NOT a re-download (curr_bytes > previous downloaded)
+ is_redownload = (
+ previous_progress is not None
+ and curr_bytes < previous_progress.downloaded.in_bytes
+ )
+
+ if is_redownload or previous_progress is None:
+ downloaded_this_session = curr_bytes
+ used_start_time = time.time()
+ else:
+ used_start_time = previous_progress.start_time
+ downloaded_this_session = (
+ previous_progress.downloaded_this_session.in_bytes
+ + (curr_bytes - previous_progress.downloaded.in_bytes)
+ )
+
+ # Key assertions
+ assert is_redownload is False, (
+ "Should NOT detect re-download for continuing download"
+ )
+ assert used_start_time == start_time, "Should preserve original start_time"
+ expected_session = initial_downloaded + (curr_bytes - initial_downloaded)
+ assert downloaded_this_session == expected_session, (
+ f"Should accumulate: {downloaded_this_session} == {expected_session}"
+ )
+ assert downloaded_this_session == 600_000, (
+ "downloaded_this_session should equal total downloaded so far"
+ )
← cd8c01b7 Fix kv prefix cache (#1262)
·
back to Exo
·
Add mlx lm style tensor sharding for Minimax (#1299) c55cbf67 →