add --download-quick-check flag to bypass the hf api calls / remote file checks

e6902b2fcf1c02a206aeff0c8639ee6cae65999e · 2024-08-09 15:40:12 +0100 · Alex Cheema
Files touched

M exo/download/hf/hf_shard_download.py
M main.py
Diff

commit e6902b2fcf1c02a206aeff0c8639ee6cae65999e
Author: Alex Cheema <alexcheema123@gmail.com>
Date:   Fri Aug 9 15:40:12 2024 +0100

    add --download-quick-check flag to bypass the hf api calls / remote file checks
---
 exo/download/hf/hf_shard_download.py | 11 +++++++++--
 main.py                              |  3 ++-
 2 files changed, 11 insertions(+), 3 deletions(-)

diff --git a/exo/download/hf/hf_shard_download.py b/exo/download/hf/hf_shard_download.py
index 176ddd32..5f178282 100644
--- a/exo/download/hf/hf_shard_download.py
+++ b/exo/download/hf/hf_shard_download.py
@@ -5,11 +5,12 @@ from typing import Dict, List, Tuple
 from exo.inference.shard import Shard
 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
+from exo.download.hf.hf_helpers import download_repo_files, RepoProgressEvent, get_weight_map, get_allow_patterns, get_repo_root
 from exo.helpers import AsyncCallbackSystem, DEBUG
 
 class HFShardDownloader(ShardDownloader):
-    def __init__(self):
+    def __init__(self, quick_check: bool = False):
+        self.quick_check = quick_check
         self.active_downloads: Dict[Shard, asyncio.Task] = {}
         self.completed_downloads: Dict[Shard, Path] = {}
         self._on_progress = AsyncCallbackSystem[str, Tuple[Shard, RepoProgressEvent]]()
@@ -17,6 +18,12 @@ class HFShardDownloader(ShardDownloader):
     async def ensure_shard(self, shard: Shard) -> Path:
         if shard in self.completed_downloads:
             return self.completed_downloads[shard]
+        if self.quick_check:
+            repo_root = get_repo_root(shard.model_id)
+            snapshots_dir = repo_root / "snapshots"
+            if snapshots_dir.exists():
+                most_recent_dir = max(snapshots_dir.iterdir(), key=lambda x: x.stat().st_mtime)
+                return most_recent_dir
 
         # If a download on this shard is already in progress, keep that one
         for active_shard in self.active_downloads:
diff --git a/main.py b/main.py
index 79a7f338..e66aa689 100644
--- a/main.py
+++ b/main.py
@@ -20,6 +20,7 @@ parser.add_argument("--node-id", type=str, default=None, help="Node ID")
 parser.add_argument("--node-host", type=str, default="0.0.0.0", help="Node host")
 parser.add_argument("--node-port", type=int, default=None, help="Node port")
 parser.add_argument("--listen-port", type=int, default=5678, help="Listening port for discovery")
+parser.add_argument("--download-quick-check", action="store_true", help="Quick check local path for shard download")
 parser.add_argument("--prometheus-client-port", type=int, default=None, help="Prometheus client port")
 parser.add_argument("--broadcast-port", type=int, default=5678, help="Broadcast port for discovery")
 parser.add_argument("--discovery-timeout", type=int, default=30, help="Discovery timeout in seconds")
@@ -36,7 +37,7 @@ print_yellow_exo()
 system_info = get_system_info()
 print(f"Detected system: {system_info}")
 
-shard_downloader: ShardDownloader = HFShardDownloader()
+shard_downloader: ShardDownloader = HFShardDownloader(quick_check=args.download_quick_check)
 inference_engine_name = args.inference_engine or ("mlx" if system_info == "Apple Silicon Mac" else "tinygrad")
 inference_engine = get_inference_engine(inference_engine_name, shard_downloader)
 print(f"Using inference engine: {inference_engine.__class__.__name__} with shard downloader: {shard_downloader.__class__.__name__}")
← 84afdbcb add --download-quick-check flag to bypass the hf api calls / · back to Exo · tinygrad select model size 31641d10 →