cache completed download paths

dd41026c5b6bfca7dbac83092dec970226cafed6 · 2024-08-07 18:16:11 +0100 · Alex Cheema

Files touched

M exo/download/hf/hf_shard_download.py

Diff

commit dd41026c5b6bfca7dbac83092dec970226cafed6
Author: Alex Cheema <alexcheema123@gmail.com>
Date:   Wed Aug 7 18:16:11 2024 +0100

    cache completed download paths
---
 exo/download/hf/hf_shard_download.py | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/exo/download/hf/hf_shard_download.py b/exo/download/hf/hf_shard_download.py
index b5607b40..ce615dfe 100644
--- a/exo/download/hf/hf_shard_download.py
+++ b/exo/download/hf/hf_shard_download.py
@@ -11,9 +11,13 @@ from exo.helpers import AsyncCallbackSystem, DEBUG
 class HFShardDownloader(ShardDownloader):
     def __init__(self):
         self.active_downloads: Dict[Shard, asyncio.Task] = {}
+        self.completed_downloads: Dict[Shard, Path] = {}
         self._on_progress = AsyncCallbackSystem[str, Tuple[Shard, RepoProgressEvent]]()
 
     async def ensure_shard(self, shard: Shard) -> Path:
+        if shard in self.completed_downloads:
+            return self.completed_downloads[shard]
+
         # If a download on this shard is already in progress, keep that one
         for active_shard in self.active_downloads:
             if active_shard == shard:
@@ -39,7 +43,9 @@ class HFShardDownloader(ShardDownloader):
         download_task = asyncio.create_task(self._download_shard(shard))
         self.active_downloads[shard] = download_task
         try:
-            return await download_task
+            path = await download_task
+            self.completed_downloads[shard] = path
+            return path
         finally:
             # Ensure the task is removed even if an exception occurs
             print(f"Removing download task for {shard}: {shard in self.active_downloads}")

← 70648873 disable prefix matching on prompts. causes subsequent reques · back to Exo · fix inference engine test 09a9abc0 →