Update download progress on pause

This commit is contained in:
ciaranbor
2026-04-01 19:52:54 +01:00
parent 89750121f4
commit bed9d9220d
2 changed files with 17 additions and 1 deletions
+7
View File
@@ -158,10 +158,17 @@ class DownloadCoordinator:
logger.info(f"Cancelling download for {model_id}")
self.active_downloads[model_id].cancel()
current_status = self.download_status[model_id]
downloaded = Memory()
total = Memory()
if isinstance(current_status, DownloadOngoing):
downloaded = current_status.download_progress.downloaded
total = current_status.download_progress.total
pending = DownloadPending(
shard_metadata=current_status.shard_metadata,
node_id=self.node_id,
model_directory=self._default_model_dir(model_id),
downloaded=downloaded,
total=total,
)
self.download_status[model_id] = pending
await self.event_sender.send(
+10 -1
View File
@@ -189,6 +189,14 @@ async def test_cancel_active_download_transitions_to_pending() -> None:
# Wait for the download to actually start (blocking in ensure_shard)
await asyncio.wait_for(slow_downloader.download_started.wait(), timeout=2.0)
# Drain any events emitted before the cancel (initial DownloadPending, DownloadOngoing)
while True:
try:
async with asyncio.timeout(0.1):
await event_recv.receive()
except TimeoutError:
break
# Cancel the download
await cmd_send.send(
ForwarderDownloadCommand(
@@ -197,10 +205,11 @@ async def test_cancel_active_download_transitions_to_pending() -> None:
)
)
# Should receive a DownloadPending event
# Should receive a DownloadPending event with preserved progress
pending = await _wait_for_pending(event_recv, MODEL_ID)
assert pending is not None, "Cancel should emit DownloadPending"
assert pending.shard_metadata.model_card.model_id == MODEL_ID
assert pending.total == Memory.from_mb(100), "Should preserve total bytes"
# Give coordinator time to clean up
await asyncio.sleep(0.05)