fix shard download

7ec660bba619284be58c0483fa5d28c18f117929 · 2024-08-07 14:10:49 +0100 · Alex Cheema

Files touched

M exo/api/chatgpt_api.py
M exo/download/hf/hf_shard_download.py

Diff

commit 7ec660bba619284be58c0483fa5d28c18f117929
Author: Alex Cheema <alexcheema123@gmail.com>
Date:   Wed Aug 7 14:10:49 2024 +0100

    fix shard download
---
 exo/api/chatgpt_api.py               | 1 -
 exo/download/hf/hf_shard_download.py | 5 +++--
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/exo/api/chatgpt_api.py b/exo/api/chatgpt_api.py
index 7b7695dc..000b865f 100644
--- a/exo/api/chatgpt_api.py
+++ b/exo/api/chatgpt_api.py
@@ -93,7 +93,6 @@ async def resolve_tokenizer(model_id: str):
     return processor
   except Exception as e:
     if DEBUG >= 4: print(f"Failed to load processor for {model_id}. Error: {e}")
-
     if DEBUG >= 4: print(traceback.format_exc())
 
   try:
diff --git a/exo/download/hf/hf_shard_download.py b/exo/download/hf/hf_shard_download.py
index 820ca065..b5607b40 100644
--- a/exo/download/hf/hf_shard_download.py
+++ b/exo/download/hf/hf_shard_download.py
@@ -15,15 +15,16 @@ class HFShardDownloader(ShardDownloader):
 
     async def ensure_shard(self, shard: Shard) -> Path:
         # If a download on this shard is already in progress, keep that one
-        for active_shard, task in self.active_downloads.values():
+        for active_shard in self.active_downloads:
             if active_shard == shard:
                 if DEBUG >= 2: print(f"Download already in progress for {shard}. Keeping that one.")
-                return await task
+                return await self.active_downloads[shard]
 
         # Cancel any downloads for this model_id on a different shard
         existing_active_shards = [active_shard for active_shard in self.active_downloads.keys() if active_shard.model_id == shard.model_id]
         for active_shard in existing_active_shards:
             if DEBUG >= 2: print(f"Cancelling download for {active_shard} (replacing with {shard})")
+            task = self.active_downloads[active_shard]
             task.cancel()
             try:
                 await task

← 29f15459 init active_downloadsa · back to Exo · handle writing responses errors b181f8aa →