Update download progress on pause
This commit is contained in:
@@ -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(
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user