Add test to CI

replicate · Oct 12, 2024 · 7268db5 · 7268db5
1 parent a0802da
commit 7268db5
Show file tree

Hide file tree

Showing 18 changed files with 266 additions and 106 deletions.
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -5,8 +5,29 @@ on:
   pull_request:
 
 jobs:
+  ruff:
+    name: Ruff
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: astral-sh/setup-uv@v3
+      - run: uv tool run ruff check
+      - run: uv tool run ruff format --check
+      - run: uv run --python 3.12 --with mypy mypy src
+
+  test:
+    name: Test
+    runs-on: ubuntu-latest-4-cores
+    steps:
+      - uses: actions/checkout@v4
+      - uses: astral-sh/setup-uv@v3
+      - run: ./script/build.sh
+      - run: ./script/test.sh
+
   build:
     name: Build + release
+    needs:
+      - test
     permissions:
       contents: read
       id-token: write

diff --git a/ruff.toml b/ruff.toml
@@ -0,0 +1,2 @@
+[format]
+quote-style = 'single'
diff --git a/build.sh → script/build.sh b/build.sh → script/build.sh
@@ -4,4 +4,5 @@
 
 set -euo pipefail
 
+cd "$(git rev-parse --show-toplevel)"
 docker build --tag monobase:latest --platform=linux/amd64 .
diff --git a/script/check.sh b/script/check.sh
@@ -0,0 +1,12 @@
+#!/bin/bash
+
+# Lint and format
+
+set -euo pipefail
+
+MONOBASE_PYTHON='3.12'
+
+cd "$(git rev-parse --show-toplevel)"
+uv tool run ruff check --fix
+uv tool run ruff format
+uv run --python "$MONOBASE_PYTHON" --with mypy mypy src
diff --git a/mini.sh → script/mini.sh b/mini.sh → script/mini.sh
@@ -4,4 +4,5 @@
 
 set -euo pipefail
 
+cd "$(git rev-parse --show-toplevel)"
 docker build --file mini.Dockerfile --tag monobase:mini --platform=linux/amd64 .
diff --git a/test.sh → script/test.sh b/test.sh → script/test.sh
@@ -4,10 +4,12 @@
 
 set -euo pipefail
 
-BUILDER_PYTHON="3.12"
+MONOBASE_PYTHON='3.12'
+
+cd "$(git rev-parse --show-toplevel)"
 
 # Build test requirements
-uv run --python "$BUILDER_PYTHON" src/update.py --environment test
+uv run --python "$MONOBASE_PYTHON" src/update.py --environment test
 
 # Build test PREFIX
 mkdir -p monobase cache

diff --git a/script/update.sh b/script/update.sh
@@ -0,0 +1,10 @@
+#!/bin/bash
+
+# Update production generations
+
+set -euo pipefail
+
+MONOBASE_PYTHON='3.12'
+
+cd "$(git rev-parse --show-toplevel)"
+uv run --python "$MONOBASE_PYTHON" src/update.py --environment prod "$@"
diff --git a/src/build.py b/src/build.py
@@ -8,23 +8,39 @@
 from cuda import install_cuda, install_cudnn
 from optimize import optimize_ld_cache, optimize_rdfind
 from prune import clean_uv_cache, prune_cuda, prune_old_gen, prune_uv_cache
-from util import Version, add_arguments, is_done, logger, mark_done
+from util import (
+    add_arguments,
+    desc_version,
+    desc_version_key,
+    is_done,
+    logger,
+    mark_done,
+)
 from uv import install_venv
 
 parser = argparse.ArgumentParser(description='Build monobase enviroment')
 add_arguments(parser)
-parser.add_argument('--prefix', metavar='PATH', default='/srv/r8/monobase',
-                    help='prefix for monobase')
-parser.add_argument('--cache', metavar='PATH', default='/var/cache/monobase',
-                    help='cache for monobase')
-parser.add_argument('--prune-old-gen', default=False, action='store_true',
-                    help='prune old generations')
-parser.add_argument('--prune-cuda', default=False, action='store_true',
-                    help='prune unused CUDAs and CuDNNs')
-parser.add_argument('--prune-uv-cache', default=True, action='store_true',
-                    help='prune uv cache')
-parser.add_argument('--clean-uv-cache', default=False, action='store_true',
-                    help='clean uv cache')
+parser.add_argument(
+    '--prefix', metavar='PATH', default='/srv/r8/monobase', help='prefix for monobase'
+)
+parser.add_argument(
+    '--cache', metavar='PATH', default='/var/cache/monobase', help='cache for monobase'
+)
+parser.add_argument(
+    '--prune-old-gen', default=False, action='store_true', help='prune old generations'
+)
+parser.add_argument(
+    '--prune-cuda',
+    default=False,
+    action='store_true',
+    help='prune unused CUDAs and CuDNNs',
+)
+parser.add_argument(
+    '--prune-uv-cache', default=True, action='store_true', help='prune uv cache'
+)
+parser.add_argument(
+    '--clean-uv-cache', default=False, action='store_true', help='clean uv cache'
+)
 
 
 def build_generation(args: argparse.Namespace, mg: MonoGen) -> None:
@@ -35,26 +51,26 @@ def build_generation(args: argparse.Namespace, mg: MonoGen) -> None:
     logger.info(f'Building monobase generation {mg.id}...')
     os.makedirs(gdir, exist_ok=True)
 
-    for k, v in sorted(mg.cuda.items(), key=lambda kv: Version.parse(kv[0]), reverse=True):
+    for k, v in desc_version_key(mg.cuda):
         src = install_cuda(args, v)
         dst = f'{gdir}/cuda{k}'
         os.symlink(os.path.relpath(src, gdir), dst)
         logger.info(f'CUDA symlinked in {dst}')
 
     cuda_major_p = re.compile(r'\.\d+$')
     cuda_majors = set(cuda_major_p.sub('', k) for k in mg.cuda.keys())
-    for k, v in sorted(mg.cudnn.items(), key=lambda kv: Version.parse(kv[0]), reverse=True):
-        for m in sorted(cuda_majors, key=Version.parse, reverse=True):
+    for k, v in desc_version_key(mg.cudnn):
+        for m in desc_version(cuda_majors):
             src = install_cudnn(args, v, m)
             dst = f'{gdir}/cudnn{k}-cuda{m}'
             os.symlink(os.path.relpath(src, gdir), dst)
             logger.info(f'CuDNN symlinked in {dst}')
 
     suffix = '' if args.environment == 'prod' else f'-{args.environment}'
     rdir = os.path.join('/opt/r8/monobase', f'requirements{suffix}', 'g%05d' % mg.id)
-    for p, pf in sorted(mg.python.items(), key=lambda kv: Version.parse(kv[0]), reverse=True):
-        for t in sorted(mg.torch, key=Version.parse, reverse=True):
-            for c in sorted(mg.cuda.keys(), key=Version.parse, reverse=True):
+    for p, pf in desc_version_key(mg.python):
+        for t in desc_version(mg.torch):
+            for c in desc_version(mg.cuda.keys()):
                 install_venv(args, rdir, gdir, p, pf, t, c)
 
     optimize_ld_cache(args, gdir, mg)

diff --git a/src/cuda.py b/src/cuda.py
@@ -44,7 +44,8 @@ def build_cudas() -> dict[str, Cuda]:
 
 def build_cudnns() -> dict[str, CuDNN]:
     p = re.compile(
-            r'^cudnn-linux-x86_64-(?P<cudnn>[^_]+)_cuda(?P<cuda_major>[^-]+)-archive.tar.xz$')
+        r'^cudnn-linux-x86_64-(?P<cudnn>[^_]+)_cuda(?P<cuda_major>[^-]+)-archive.tar.xz$'
+    )
     cudnns = {}
     for u in cudnn_urls:
         url = urllib.parse.urlparse(u)
@@ -76,9 +77,15 @@ def install_cuda(args: argparse.Namespace, version: str) -> str:
 
     logger.info(f'Installing CUDA {version}...')
     cmd = [
-        '/bin/sh', file, f'--installpath={cdir}',
-        '--toolkit', '--override', '--silent',
-        '--no-opengl-libs', '--no-man-page', '--no-drm'
+        '/bin/sh',
+        file,
+        f'--installpath={cdir}',
+        '--toolkit',
+        '--override',
+        '--silent',
+        '--no-opengl-libs',
+        '--no-man-page',
+        '--no-drm',
     ]
     subprocess.run(cmd, check=True)
 
@@ -114,7 +121,13 @@ def install_cudnn(args: argparse.Namespace, version: str, cuda_major: str) -> st
     file = os.path.join(args.cache, cudnn.filename)
     if not os.path.exists(file):
         logger.info(f'Downloading CuDNN {key}...')
-        cmd = [f'{args.prefix}/bin/pget', '--pid-file', '/tmp/pget.pid', cudnn.url, file]
+        cmd = [
+            f'{args.prefix}/bin/pget',
+            '--pid-file',
+            '/tmp/pget.pid',
+            cudnn.url,
+            file,
+        ]
         subprocess.run(cmd, check=True)
 
     logger.info(f'Installing CuDNN {key}...')

diff --git a/src/entrypoint.sh b/src/entrypoint.sh
@@ -4,7 +4,7 @@
 
 set -euo pipefail
 
-BUILDER_PYTHON='3.12'
+MONOBASE_PYTHON='3.12'
 DONE_FILE='/opt/r8/monobase/.done'
 
 UV_URL='https://github.com/astral-sh/uv/releases/latest/download/uv-x86_64-unknown-linux-gnu.tar.gz'
@@ -38,7 +38,7 @@ builder() {
     cp /opt/r8/monobase/pget "$MONOBASE_PREFIX/bin/pget"
 
     log "Running builder..."
-    uv run --python "$BUILDER_PYTHON" /opt/r8/monobase/build.py "$@"
+    uv run --python "$MONOBASE_PYTHON" /opt/r8/monobase/build.py "$@"
 
     # Inside K8S
     # Write done file to signal pod ready

diff --git a/src/monogen.py b/src/monogen.py
@@ -18,7 +18,7 @@ class MonoGen:
         cudnn={'9': '9.1.0.70'},
         python={'3.12': '3.12.6'},
         torch=['2.4.1', '2.6.0.dev20240918'],
-        pip_pkgs=['cog==0.9.23', 'opencv-python==4.10.0.84']
+        pip_pkgs=['cog==0.9.23', 'opencv-python==4.10.0.84'],
     ),
 ]
 
@@ -47,11 +47,18 @@ class MonoGen:
             '3.12': '3.12.6',
         },
         torch=[
-            '2.0.0', '2.0.1',
-            '2.1.0', '2.1.1', '2.1.2',
-            '2.2.0', '2.2.1', '2.2.2',
-            '2.3.0', '2.3.1',
-            '2.4.0', '2.4.1',
+            '2.0.0',
+            '2.0.1',
+            '2.1.0',
+            '2.1.1',
+            '2.1.2',
+            '2.2.0',
+            '2.2.1',
+            '2.2.2',
+            '2.3.0',
+            '2.3.1',
+            '2.4.0',
+            '2.4.1',
         ],
         pip_pkgs=[
             'https://github.com/replicate/cog/archive/refs/heads/add-waiting-env.zip',
@@ -77,11 +84,18 @@ class MonoGen:
             '3.12': '3.12.6',
         },
         torch=[
-            '2.0.0', '2.0.1',
-            '2.1.0', '2.1.1', '2.1.2',
-            '2.2.0', '2.2.1', '2.2.2',
-            '2.3.0', '2.3.1',
-            '2.4.0', '2.4.1',
+            '2.0.0',
+            '2.0.1',
+            '2.1.0',
+            '2.1.1',
+            '2.1.2',
+            '2.2.0',
+            '2.2.1',
+            '2.2.2',
+            '2.3.0',
+            '2.3.1',
+            '2.4.0',
+            '2.4.1',
             # Nightly
             '2.6.0.dev20240918',
         ],
@@ -109,11 +123,18 @@ class MonoGen:
             '3.12': '3.12.6',
         },
         torch=[
-            '2.0.0', '2.0.1',
-            '2.1.0', '2.1.1', '2.1.2',
-            '2.2.0', '2.2.1', '2.2.2',
-            '2.3.0', '2.3.1',
-            '2.4.0', '2.4.1',
+            '2.0.0',
+            '2.0.1',
+            '2.1.0',
+            '2.1.1',
+            '2.1.2',
+            '2.2.0',
+            '2.2.1',
+            '2.2.2',
+            '2.3.0',
+            '2.3.1',
+            '2.4.0',
+            '2.4.1',
             # Nightly
             '2.6.0.dev20240918',
         ],
@@ -141,11 +162,18 @@ class MonoGen:
             '3.12': '3.12.6',
         },
         torch=[
-            '2.0.0', '2.0.1',
-            '2.1.0', '2.1.1', '2.1.2',
-            '2.2.0', '2.2.1', '2.2.2',
-            '2.3.0', '2.3.1',
-            '2.4.0', '2.4.1',
+            '2.0.0',
+            '2.0.1',
+            '2.1.0',
+            '2.1.1',
+            '2.1.2',
+            '2.2.0',
+            '2.2.1',
+            '2.2.2',
+            '2.3.0',
+            '2.3.1',
+            '2.4.0',
+            '2.4.1',
             # Nightly
             '2.6.0.dev20240918',
         ],
@@ -173,7 +201,9 @@ def validate():
             for k, v in g.python.items():
                 assert v.startswith(f'{k}.'), f'[{env}] Python {v} is not {k}'
             for p in g.pip_pkgs:
-                assert '==' in p or p.startswith('https://'), f'PIP package {p} is not pinned'
+                assert '==' in p or p.startswith(
+                    'https://'
+                ), f'PIP package {p} is not pinned'
 
 
 validate()
diff --git a/src/optimize.py b/src/optimize.py
@@ -28,10 +28,15 @@ def optimize_ld_cache(args: argparse.Namespace, gdir: str, mg: MonoGen) -> None:
 def optimize_rdfind(args: argparse.Namespace, gdir: str, mg: MonoGen) -> None:
     logger.info(f'Running rdfind for generation {mg.id}...')
     cmd = [
-        'rdfind', '-minsize', str(1024*1024),
-        '-deterministic', 'true',
-        '-makehardlinks', 'true',
-        '-outputname', '/dev/null',
+        'rdfind',
+        '-minsize',
+        str(1024 * 1024),
+        '-deterministic',
+        'true',
+        '-makehardlinks',
+        'true',
+        '-outputname',
+        '/dev/null',
         f'{args.prefix}/uv/cache',
         f'{args.prefix}/cuda',
         gdir,

diff --git a/src/prune.py b/src/prune.py
@@ -16,11 +16,20 @@ def prune_old_gen(args: argparse.Namespace) -> None:
 
 def prune_cuda(args: argparse.Namespace) -> None:
     cmd = [
-        'find', f'{args.prefix}/monobase',
-        '-type', 'l',
-        '-maxdepth', '2',
-        '(', '-name', 'cuda*', '-or', '-name', 'cudnn*', ')',
-        '-print0'
+        'find',
+        f'{args.prefix}/monobase',
+        '-type',
+        'l',
+        '-maxdepth',
+        '2',
+        '(',
+        '-name',
+        'cuda*',
+        '-or',
+        '-name',
+        'cudnn*',
+        ')',
+        '-print0',
     ]
     proc = subprocess.run(cmd, check=True, capture_output=True, text=True)
     links = set(filter(None, proc.stdout.split('\0')))
Original file line number	Diff line number	Diff line change
Expand Up		@@ -4,4 +4,5 @@

		set -euo pipefail

		cd "$(git rev-parse --show-toplevel)"
		docker build --tag monobase:latest --platform=linux/amd64 .