← back to Exo
modify get_shard_download_status to use helper function
dec79ac747ba79a657b8fe2244e697ab2095eaf5 · 2024-11-18 13:54:38 -0800 · cadenmackenzie
Files touched
M exo/download/hf/hf_shard_download.py
Diff
commit dec79ac747ba79a657b8fe2244e697ab2095eaf5
Author: cadenmackenzie <cadenmackenzie@gmail.com>
Date: Mon Nov 18 13:54:38 2024 -0800
modify get_shard_download_status to use helper function
---
exo/download/hf/hf_shard_download.py | 58 ++++++++++++++++++++++--------------
1 file changed, 36 insertions(+), 22 deletions(-)
diff --git a/exo/download/hf/hf_shard_download.py b/exo/download/hf/hf_shard_download.py
index 2b4ac6db..5b98e489 100644
--- a/exo/download/hf/hf_shard_download.py
+++ b/exo/download/hf/hf_shard_download.py
@@ -7,7 +7,9 @@ from exo.download.shard_download import ShardDownloader
from exo.download.download_progress import RepoProgressEvent
from exo.download.hf.hf_helpers import (
download_repo_files, RepoProgressEvent, get_weight_map,
- get_allow_patterns, get_repo_root, fetch_file_list, get_local_snapshot_dir
+ get_allow_patterns, get_repo_root, fetch_file_list,
+ get_local_snapshot_dir, get_file_download_percentage,
+ filter_repo_objects
)
from exo.helpers import AsyncCallbackSystem, DEBUG
from exo.models import model_cards, get_repo
@@ -105,32 +107,44 @@ class HFShardDownloader(ShardDownloader):
if DEBUG >= 2: print(f"No weight map found for {self.current_repo_id}")
return None
- # Get the patterns for this shard
+ # Get all files needed for this shard
patterns = get_allow_patterns(weight_map, self.current_shard)
- # First check which files exist locally
+ # Check download status for all relevant files
status = {}
- local_files = []
- local_sizes = {}
+ total_bytes = 0
+ downloaded_bytes = 0
- for pattern in patterns:
- if pattern.endswith('safetensors') or pattern.endswith('mlx'):
- file_path = snapshot_dir / pattern
- if await aios.path.exists(file_path):
- local_size = await aios.path.getsize(file_path)
- local_files.append(pattern)
- local_sizes[pattern] = local_size
-
- # Only fetch remote info if we found local files
- if local_files:
- async with aiohttp.ClientSession() as session:
- file_list = await fetch_file_list(session, self.current_repo_id, self.revision)
+ async with aiohttp.ClientSession() as session:
+ file_list = await fetch_file_list(session, self.current_repo_id, self.revision)
+ relevant_files = list(filter_repo_objects(file_list, allow_patterns=patterns, key=lambda x: x["path"]))
+
+ for file in relevant_files:
+ file_size = file["size"]
+ total_bytes += file_size
- for pattern in local_files:
- for file in file_list:
- if file["path"].endswith(pattern):
- status[pattern] = (local_sizes[pattern] / file["size"]) * 100
- break
+ percentage = await get_file_download_percentage(
+ session,
+ self.current_repo_id,
+ self.revision,
+ file["path"],
+ snapshot_dir
+ )
+ status[file["path"]] = percentage
+ downloaded_bytes += (file_size * (percentage / 100))
+
+ # Add overall progress weighted by file size
+ if total_bytes > 0:
+ status["overall"] = (downloaded_bytes / total_bytes) * 100
+ else:
+ status["overall"] = 0
+
+ if DEBUG >= 2:
+ print(f"Download calculation for {self.current_repo_id}:")
+ print(f"Total bytes: {total_bytes}")
+ print(f"Downloaded bytes: {downloaded_bytes}")
+ for file in relevant_files:
+ print(f"File {file['path']}: size={file['size']}, percentage={status[file['path']]}")
return status
← c61f40c6 adding helper funciton to check file download. also modifyin
·
back to Exo
·
modifying helper fucntion checking size to follow redirect f 4c6fda7c →