working worker bee

Author: earonesty

ai_worker/main.py

@@ -2,11 +2,9 @@
 import asyncio
 import json
 import multiprocessing
-import os
 from typing import Optional
 import logging as log
 
-import httpx
 import psutil
 import sseclient
 import websockets
@@ -17,9 +15,10 @@
 from fastapi.testclient import TestClient
 from starlette.responses import Response
 
+from gguf_loader.main import get_size
+
 APP_NAME= "gputopia"
-DEFAULT_COORDINATOR = "https://gputopia.ai/api/v1"
-DEFAULT_BASE_URL = "https://gputopia.ai/models"
+DEFAULT_COORDINATOR = "wss://gputopia.ai/api/v1"
 
 
 class Req(BaseModel):
@@ -31,8 +30,6 @@ class Config(BaseSettings):
     model_config = SettingsConfigDict(env_prefix=APP_NAME +'_worker', case_sensitive=False)
     auth_key: str = ""
     coordinator_url: str = DEFAULT_COORDINATOR
-    model_base_url: str = DEFAULT_BASE_URL
-    model_dir: str = os.path.expanduser('~/.ai-models')
 
 
 class WorkerMain:
@@ -53,13 +50,13 @@ async def run(self):
 
     async def guess_layers(self, model_path):
         # todo: read model file and compare to gpu resources
-        return 30
+        return 20
 
     async def load_model(self, name):
         if name == self.llama_model:
             return
         model_path = await self.get_model(name)
-        settings = LlamaSettings(model=model_path, n_gpu_layers=self.guess_layers(model_path), seed=-1, embedding=True, cache=True, port=8181)
+        settings = LlamaSettings(model=model_path, n_gpu_layers=await self.guess_layers(model_path), seed=-1, embedding=True, cache=True, port=8181)
         self.llama = create_llama_app(settings)
         self.llama_cli = TestClient(self.llama)
 
@@ -112,56 +109,26 @@ async def run_ws(self, ws):
                 ws.send(res.body.decode("urf-8"))
 
     async def get_model(self, name):
-        ret = self.get_local_model(name)
-        if ret:
-            return ret
         return await self.download_model(name)
 
-    def get_local_model(self, name):
-        dest = self.model_file_for(name)
-        if os.path.getsize(dest) > 0:
-            return dest
-        return None
-
-    def model_file_for(self, name):
-        return self.conf.model_dir + "/" + name.replace("/", ".")
-
     async def download_model(self, name):
-        url = self.conf.model_base_url + "/" + name.replace("/", ".")
-
-        async with httpx.AsyncClient() as client:
-            r = await client.head(url)
-            size = r.headers.get('Content-Length')
-            if not size:
-                params = self.get_model_params(name)
-                bits = self.get_model_bits(name)
-                # 70b * 4 bit = 35gb (roughly)
-                size = params * bits / 8
-
-            assert size, "unable to estimate model size, not downloading"
-
-            await self.free_up_space(size)
-
-            dest = self.model_file_for(name)
-
-            done = 0
-            with open(dest + ".tmp", "wb") as f:
-                async with client.stream("GET", url) as r:
-                    async for chunk in r.aiter_bytes():
-                        f.write(chunk)
-                        done += len(chunk)
-                        self.report_pct(name, done/size)
-            os.replace(dest + ".tmp", dest)
-            self.report_done(name)
-
-        return dest
+        # uses hf cache, so no need to handle here
+        from gguf_loader.main import download_gguf
+        size = get_size(name)
+        await self.free_up_space(size)
+        loop = asyncio.get_running_loop()
+        path = await loop.run_in_executor(None, lambda: download_gguf(name))
+        return path
 
     def report_done(self, name):
         print("\r", name, 100)
 
     def report_pct(self, name, pct):
         print("\r", name, pct, end='')
 
+    async def free_up_space(self, size):
+        pass
+
 
 def main():
     parser = argparse.ArgumentParser()
@@ -181,7 +148,6 @@ def main():
 
     conf = Config(**{k: v for k, v in vars(args).items() if v is not None})
 
-
     wm = WorkerMain(conf)
 
-    asyncio.run(wm.main())
+    asyncio.run(wm.run())

gguf_loader/main.py

@@ -22,7 +22,12 @@ def convert_to_gguf(file):
     return dest
 
 
-def download_gguf(name):
+def get_size(name):
+    typ, hf, fil = pick_file(name)
+    return fil["size"]
+
+
+def pick_file(name):
     parts = name.split(":", 1)
     if len(parts) == 1:
         hf, filt = parts[0], ""
@@ -50,26 +55,32 @@ def download_gguf(name):
             # this is all heuristics, but, imo it can be more than good enough
             raise ValueError("Need ggml or gguf")
 
-        base = os.path.basename(ggml[0]["name"])
+        return "ggml", hf, ggml[0]
 
-        log.debug("downloading...")
+    if len(gguf) > 1:
+        raise ValueError("Multiple files match, please specify a better filter")
 
-        # use hf so we get a nice cache
-        path = hf_hub_download(repo_id=hf, filename=base, resume_download=True)
+    return "gguf", hf, gguf[0]
 
-        return convert_to_gguf(path)
 
-    if len(gguf) > 1:
-        raise ValueError("Multiple files match, please specify a better filter")
+def download_gguf(name):
+    typ, repo_id, fil = pick_file(name)
+    if typ == "ggml":
+        base = os.path.basename(fil["name"])
+        log.debug("downloading...")
+        # use hf so we get a nice cache
+        path = hf_hub_download(repo_id=repo_id, filename=base, resume_download=True)
+        return convert_to_gguf(path)
 
-    base = os.path.basename(gguf[0]["name"])
+    base = os.path.basename(fil["name"])
     log.debug("downloading...")
-    return hf_hub_download(repo_id=hf, filename=base)
+    return hf_hub_download(repo_id=repo_id, filename=base)
 
 
 # Load environment variables from .env file
 load_dotenv()
 
+
 # Get AWS credentials from environment variables
 
 

pytest.ini

@@ -0,0 +1,2 @@
+[pytest]
+asyncio_mode = auto

tests/__pycache__/test_conn.cpython-311-pytest-7.4.0.pyc

@@ -1,7 +1,7 @@
 import json
 
 from ai_worker.main import WorkerMain, Config
-from gguf_loader.main import download_gguf, main as loader_main
+from gguf_loader.main import download_gguf, main as loader_main, get_size
 
 try:
     from pynvml.smi import nvidia_smi
@@ -21,8 +21,21 @@ def test_conn_str():
     assert js["vram"]
 
 
+async def test_wm():
+    wm = WorkerMain(Config())
+    await wm.load_model("TheBloke/WizardLM-7B-uncensored-GGML:q4_K_M")
+    res = wm.llama_cli.post("/v1/chat/completions", json=dict(
+        model=wm.llama_model,
+        messages=[
+            {"role": "system", "content": "You are a helpful assistant"},
+            {"role": "user", "content": "hello"},
+        ]
+    ))
+    assert res
+
 def test_download_model():
-    download_gguf("TheBloke/WizardLM-7B-uncensored-GGML:q4_K_M")
+    assert get_size("TheBloke/WizardLM-7B-uncensored-GGML:q4_K_M") > 0
+    assert download_gguf("TheBloke/WizardLM-7B-uncensored-GGML:q4_K_M")
 
 
 def test_download_main(capsys):