Support installing coglet==latest

replicate · Jan 28, 2025 · 656a566 · 656a566
1 parent 8434a34
commit 656a566
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 11 deletions.
diff --git a/script/test-mini b/script/test-mini
@@ -14,7 +14,7 @@ uv run --python "$MONOBASE_PYTHON" python -m monobase.update --environment test
 # Build test PREFIX
 mkdir -p build/monobase build/cache build/root
 docker run --rm \
-    --env R8_COG_VERSION=0.11.3 \
+    --env R8_COG_VERSION=coglet \
     --env R8_CUDA_VERSION=12.4 \
     --env R8_CUDNN_VERSION=9 \
     --env R8_PYTHON_VERSION=3.12 \
@@ -32,7 +32,7 @@ docker run --rm \
     "$@"
 
 docker run --rm \
-    --env R8_COG_VERSION=0.11.3 \
+    --env R8_COG_VERSION=coglet \
     --env R8_CUDA_VERSION=12.4 \
     --env R8_CUDNN_VERSION=9 \
     --env R8_PYTHON_VERSION=3.12 \
@@ -112,10 +112,10 @@ if [ "$fail" -gt 0 ]; then
 fi
 
 read -r -d '' SCRIPT << EOF || :
-import sys, cog, torch
+import sys, coglet, torch
 assert sys.version.startswith('3.12.8'), f'sys.version is not 3.12.8: {sys.version}'
-assert cog.__file__.startswith('/srv/r8/monobase/cog'), f'cog is not pre-installed: {cog.__file__}'
-assert cog.__version__ == '0.11.3', f'cog.__version__ is not 0.11.3: {cog.__version__}'
+assert coglet.__file__.startswith('/srv/r8/monobase/cog'), f'coglet is not pre-installed: {coglet.__file__}'
+assert coglet.__version__ == '0.1.0a1', f'coglet.__version__ is not 0.1.0a1: {coglet.__version__}'
 assert torch.__version__ == '2.4.1+cu124', f'torch.__version__ is not 2.4.1+cu124: {torch.__version__}'
 print('PASS: venv versions')
 EOF
@@ -124,7 +124,6 @@ docker run --rm \
     --volume "$PWD/src/monobase:/opt/r8/monobase:ro" \
     --volume "$PWD/build/monobase:/srv/r8/monobase:ro" \
     --volume "$PWD/build/root:/root" \
-    --env R8_COG_VERSION=0.11.3 \
     --env R8_CUDA_VERSION=12.4 \
     --env R8_CUDNN_VERSION=9 \
     --env R8_PYTHON_VERSION=3.12 \
@@ -142,7 +141,6 @@ docker run --rm \
     --volume "$PWD/src/monobase:/opt/r8/monobase:ro" \
     --volume "$PWD/build/monobase:/srv/r8/monobase:ro" \
     --volume "$PWD/build/root:/root" \
-    --env R8_COG_VERSION=0.11.3 \
     --env R8_CUDA_VERSION=12.4 \
     --env R8_CUDNN_VERSION=9 \
     --env R8_PYTHON_VERSION=3.12 \

diff --git a/src/monobase/cog.py b/src/monobase/cog.py
@@ -49,14 +49,22 @@ def install_cog(
         pkg = 'coglet' if 'coglet' in cog_version else 'cog'
         cog_name = f'{pkg}{h}'
         spec = f'{pkg}@{cog_version}'
-    elif cog_version.startswith('coglet=='):
+    elif cog_version.startswith('coglet'):
         if python_version == '3.8':
             log.warning('cog-runtime does not support Python 3.8')
             return
         try:
-            tag = cog_version.lstrip('coglet==')
-            cog_name = f'coglet{tag}'
-            url = f'https://api.github.com/repos/replicate/cog-runtime/releases/tags/v{tag}'
+            if cog_version == 'coglet':
+                v = 'latest'
+                part = 'latest'
+            elif '==' in cog_version:
+                v = cog_version.split('==')[1].strip()
+                part = f'tags/v{v}'
+            else:
+                log.error(f'Unsupported cog version {cog_version}')
+                return
+            cog_name = f'coglet{v}'
+            url = f'https://api.github.com/repos/replicate/cog-runtime/releases/{part}'
             content = urllib.request.urlopen(url).read()
             blob = json.loads(content)
             whl = next(filter(lambda a: a['name'].endswith('.whl'), blob['assets']))