ver

Author: earonesty

.gitignore

@@ -1,3 +1,5 @@
+ai_worker/version.py
+
 # lib stuff
 CLBlast/
 build-cuda/

ai_worker/main.py

@@ -26,6 +26,7 @@
 from gguf_loader.main import get_size
 
 from .gguf_reader import GGUFReader
+from .version import VERSION
 
 APP_NAME = "gputopia"
 DEFAULT_COORDINATOR = "wss://queenbee.gputopia.ai/worker"
@@ -47,6 +48,7 @@ class GpuInfo(BaseModel):
 
 
 class ConnectMessage(BaseModel):
+    worker_version: str
     ln_url: str
     auth_key: str
     cpu_count: int
@@ -189,6 +191,7 @@ def _get_connect_info(self) -> ConnectMessage:
         disk_space = get_free_space_mb(".")
 
         connect_msg = ConnectMessage(
+            worker_version=VERSION,
             ln_url=self.conf.ln_url,
             auth_key=self.conf.auth_key,
             disk_space=int(disk_space),
@@ -301,11 +304,12 @@ async def free_up_space(self, size):
         pass
 
 
-def main():
+def main(argv=None):
     logging.basicConfig(level=logging.INFO,
                         format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
                         stream=sys.stdout)
     parser = argparse.ArgumentParser()
+    arg_names = []
     for name, field in Config.model_fields.items():
         description = field.description
         if field.default is not None and description is not None:
@@ -318,16 +322,23 @@ def main():
         )
         if field.annotation is bool:
             args.pop("type")
+        arg_names.append(name)
         parser.add_argument(f"--{name}", **args)
+        
+    parser.add_argument(f"--version", action="store_true")
 
-    args = parser.parse_args()
+    args = parser.parse_args(args=argv)
     if args.debug:
         log.setLevel(logging.DEBUG)
 
+    if args.version:
+        print(VERSION)
+        exit(0)
+
     if os.path.exists("gputopia-worker.ini"):
         logging.config.fileConfig("gputopia-worker.ini")
 
-    conf = Config(**{k: v for k, v in vars(args).items() if v is not None})
+    conf = Config(**{k: getattr(args, k) for k in arg_names if getattr(args, k) is not None})
 
     wm = WorkerMain(conf)
 

ai_worker/version.py

@@ -0,0 +1 @@
+VERSION = '0.1.3'

build-bin.sh

@@ -16,6 +16,8 @@ python -mvenv "build-$gpu"
 # python is absurd putting these in different places
 . build-$gpu/bin/activate 2> /dev/null || . build-$gpu/scripts/activate
 
+python build-version.py
+
 pip uninstall -y llama-cpp-python
 
 # windows/linux cache rm (poetry cache control is kinda blunt/broken)

build-version.py

@@ -0,0 +1,8 @@
+import toml
+
+with open("pyproject.toml", "r") as file:
+    pyproject_content = toml.load(file)
+    version = pyproject_content["tool"]["poetry"]["version"]
+
+with open("ai_worker/version.py", "w") as file:
+    file.write(f"VERSION = '{version}'\n")

poetry.lock

@@ -32,6 +32,7 @@ pytest-asyncio = "^0.21.1"
 pyinstaller = "^5.13.2"
 boto3 = "^1.28.43"
 ruff = "^0.0.288"
+toml = "^0.10.2"
 
 
 [tool.poetry.scripts]

pyproject.toml

@@ -1,12 +1,13 @@
 import asyncio
 import json
+import re
 import time
 from threading import Thread
 from typing import Any, Optional
 
 import pytest
 import websockets
-from ai_worker.main import WorkerMain, Config
+from ai_worker.main import WorkerMain, Config, main as worker_main
 from gguf_loader.main import download_gguf, main as loader_main, get_size
 
 from pynvml.smi import nvidia_smi
@@ -113,3 +114,21 @@ def test_download_main(capsys):
     loader_main(["TheBloke/WizardLM-7B-uncensored-GGML:q4_K_M"])
     oe = capsys.readouterr().out
     assert "q4_K_M" in oe
+
+
+def test_version(capsys):
+    try:
+        worker_main(["--version"])
+    except SystemExit:
+        pass
+    oe = capsys.readouterr().out
+    re.match(r"\d+\.\d+\.\d+", oe)
+
+
+def test_main(capsys):
+    try:
+        worker_main(["--version"])
+    except SystemExit:
+        pass
+    oe = capsys.readouterr().out
+    re.match(r"\d+\.\d+\.\d+", oe)