download status in parallel, support async ensure shard with using shard_downloader instead

8c191050a2dd20dea129d441c1e7d98e22228609 · 2025-01-05 02:31:59 +0000 · Alex Cheema

Files touched

M exo/api/chatgpt_api.py

Diff

commit 8c191050a2dd20dea129d441c1e7d98e22228609
Author: Alex Cheema <alexcheema123@gmail.com>
Date:   Sun Jan 5 02:31:59 2025 +0000

    download status in parallel, support async ensure shard with using shard_downloader instead
---
 exo/api/chatgpt_api.py | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/exo/api/chatgpt_api.py b/exo/api/chatgpt_api.py
index 73762de7..4daa199d 100644
--- a/exo/api/chatgpt_api.py
+++ b/exo/api/chatgpt_api.py
@@ -245,7 +245,7 @@ class ChatGPTAPI:
         )
         await response.prepare(request)
 
-        for model_name, pretty in pretty_name.items():
+        async def process_model(model_name, pretty):
             if model_name in model_cards:
                 model_info = model_cards[model_name]
 
@@ -273,6 +273,12 @@ class ChatGPTAPI:
 
                         await response.write(f"data: {json.dumps(model_data)}\n\n".encode())
 
+        # Process all models in parallel
+        await asyncio.gather(*[
+            process_model(model_name, pretty)
+            for model_name, pretty in pretty_name.items()
+        ])
+
         await response.write(b"data: [DONE]\n\n")
         return response
 
@@ -562,7 +568,7 @@ class ChatGPTAPI:
       if model_name not in model_cards: return web.json_response({"error": f"Invalid model: {model_name}. Supported models: {list(model_cards.keys())}"}, status=400)
       shard = build_base_shard(model_name, self.inference_engine_classname)
       if not shard: return web.json_response({"error": f"Could not build shard for model {model_name}"}, status=400)
-      asyncio.create_task(self.node.inference_engine.ensure_shard(shard))
+      asyncio.create_task(self.node.inference_engine.shard_downloader.ensure_shard(shard))
 
       return web.json_response({
         "status": "success",

← fe50d4d3 Add --system-prompt to exo cli · back to Exo · fix args for ensure_shard 29244c63 →