Minimal sleap-nn CUDA Base Image

.github/workflows/sleap_nn_v002_cuda_production.yml

@@ -0,0 +1,67 @@
+name: Build and Push sleap-nn-v002-cuda-production (Production Workflow)
+
+# Run on push to branches other than main for sleapnn_v002_cuda_v128
+on:
+  push:
+    branches:
+      - main
+    paths:
+      - sleapnn_v002_cuda_v128/** # Only run on changes to sleapnn_v002_cuda_v128
+      - .github/workflows/sleap_nn_v002_cuda_production.yml # Only run on changes to this workflow
+
+jobs:
+  build:
+    runs-on: ubuntu-latest # Only build on Ubuntu for now since Docker is not available on macOS runners
+    strategy:
+      matrix:
+        platform: [linux/amd64] # Only build amd64 for now
+      max-parallel: 2  # Build both architectures in parallel (if more than one)
+    outputs:
+      git_sha: ${{ steps.get_sha.outputs.sha }}
+      sanitized_platform: ${{ steps.sanitize_platform.outputs.sanitized_platform }}
+    steps:
+      # Step 1: Checkout the repository
+      - name: Checkout code
+        uses: actions/checkout@v4
+
+      # Step 2: Get Git SHA for tagging the image
+      - name: Get Git SHA
+        id: get_sha
+        run: echo "sha=$(git rev-parse HEAD)" >> $GITHUB_OUTPUT
+
+      - name: Debug Git SHA
+        run: echo "Git SHA ${{ steps.get_sha.outputs.sha }}"
+
+      # Step 3: Sanitize platform name for tagging
+      - name: Sanitize platform name
+        id: sanitize_platform
+        run: |
+          sanitized_platform="${{ matrix.platform }}" # Copy platform value
+          sanitized_platform="${sanitized_platform/\//-}" # Replace / with -
+          echo "sanitized_platform=$sanitized_platform" >> $GITHUB_OUTPUT
+
+      # Step 4: Set up Docker Buildx for multi-architecture builds
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v3
+        with:
+          driver: docker-container # Use a container driver for Buildx (default)
+
+      # Step 5: Authenticate to GitHub Container Registry
+      - name: Authenticate to GitHub Container Registry
+        run: echo "${{ secrets.GITHUB_TOKEN }}" | docker login ghcr.io -u ${{ github.actor }} --password-stdin
+
+      # Step 6: Build and push the Docker image to GitHub Container Registry
+      - name: Build and push Docker image
+        uses: docker/build-push-action@v6
+        with:
+          context: ./sleapnn_v002_cuda_v128 # Build context wrt the root of the repository
+          file: ./sleapnn_v002_cuda_v128/Dockerfile # Path to Dockerfile wrt the root of the repository
+          platforms: ${{ matrix.platform }}
+          push: true # Push the image to GitHub Container Registry
+          # Tags all include "-production" to differentiate from test images
+          tags: |
+            ghcr.io/${{ github.repository_owner }}/sleap-nn-cuda:latest
+            ghcr.io/${{ github.repository_owner }}/sleap-nn-cuda:${{ steps.sanitize_platform.outputs.sanitized_platform }}
+            ghcr.io/${{ github.repository_owner }}/sleap-nn-cuda:${{ steps.sanitize_platform.outputs.sanitized_platform }}-nvidia-cuda-12.8.0-cudnn-runtime-ubuntu24.04
+            ghcr.io/${{ github.repository_owner }}/sleap-nn-cuda:${{ steps.sanitize_platform.outputs.sanitized_platform }}-sleap-nn-0.0.2
+            ghcr.io/${{ github.repository_owner }}/sleap-nn-cuda:${{ steps.sanitize_platform.outputs.sanitized_platform }}-${{ steps.get_sha.outputs.sha }}

.github/workflows/sleap_nn_v002_cuda_test.yml

@@ -0,0 +1,76 @@
+name: Build and Push sleap-nn-v002-cuda-test (Test Workflow)
+
+# Run on push to branches other than main for sleapnn_v002_cuda_v128
+on:
+  push:
+    branches-ignore:
+      - main
+    paths:
+      - sleapnn_v002_cuda_v128/** # Only run on changes to sleapnn_v002_cuda_v128
+      - .github/workflows/sleap_nn_v002_cuda_test.yml # Only run on changes to this workflow
+
+jobs:
+  build:
+    runs-on: ubuntu-latest # Only build on Ubuntu for now since Docker is not available on macOS runners
+    strategy:
+      matrix:
+        platform: [linux/amd64] # Only build amd64 for now
+      max-parallel: 2  # Build both architectures in parallel (if more than one)
+    outputs:
+      git_sha: ${{ steps.get_sha.outputs.sha }}
+      sanitized_platform: ${{ steps.sanitize_platform.outputs.sanitized_platform }}
+    steps:
+      # Step 1: Checkout the repository
+      - name: Checkout code
+        uses: actions/checkout@v4
+
+      # Step 2: Get Git SHA for tagging the image
+      - name: Get Git SHA
+        id: get_sha
+        run: echo "sha=$(git rev-parse HEAD)" >> $GITHUB_OUTPUT
+
+      - name: Debug Git SHA
+        run: echo "Git SHA ${{ steps.get_sha.outputs.sha }}"
+
+      # Step 3: Sanitize platform name for tagging
+      - name: Sanitize platform name
+        id: sanitize_platform
+        run: |
+          sanitized_platform="${{ matrix.platform }}" # Copy platform value
+          sanitized_platform="${sanitized_platform/\//-}" # Replace / with -
+          echo "sanitized_platform=$sanitized_platform" >> $GITHUB_OUTPUT
+
+      # Step 4: Set up Docker Buildx for multi-architecture builds
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v3
+        with:
+          driver: docker-container # Use a container driver for Buildx (default)
+
+      # Step 4.5: Free up disk space
+      - name: Free up disk space
+        run: |
+          sudo rm -rf /usr/share/dotnet
+          sudo rm -rf /opt/ghc
+          sudo rm -rf /usr/local/share/boost
+          sudo rm -rf "$AGENT_TOOLSDIRECTORY"
+          sudo docker system prune -af
+          df -h
+
+      # Step 5: Authenticate to GitHub Container Registry
+      - name: Authenticate to GitHub Container Registry
+        run: echo "${{ secrets.GITHUB_TOKEN }}" | docker login ghcr.io -u ${{ github.actor }} --password-stdin
+
+      # Step 6: Build and push the Docker image to GitHub Container Registry
+      - name: Build and push Docker image
+        uses: docker/build-push-action@v6
+        with:
+          context: ./sleapnn_v002_cuda_v128 # Build context wrt the root of the repository
+          file: ./sleapnn_v002_cuda_v128/Dockerfile # Path to Dockerfile wrt the root of the repository
+          platforms: ${{ matrix.platform }}
+          push: true # Push the image to GitHub Container Registry
+          # Tags all include "-test" to differentiate from production images
+          tags: |
+            ghcr.io/${{ github.repository_owner }}/sleap-nn-cuda:${{ steps.sanitize_platform.outputs.sanitized_platform }}-test
+            ghcr.io/${{ github.repository_owner }}/sleap-nn-cuda:${{ steps.sanitize_platform.outputs.sanitized_platform }}-nvidia-cuda-12.8.0-cudnn-runtime-ubuntu24.04-test
+            ghcr.io/${{ github.repository_owner }}/sleap-nn-cuda:${{ steps.sanitize_platform.outputs.sanitized_platform }}-sleap-nn-0.0.2-test
+            ghcr.io/${{ github.repository_owner }}/sleap-nn-cuda:${{ steps.sanitize_platform.outputs.sanitized_platform }}-${{ steps.get_sha.outputs.sha }}-test

sleapnn_v002_cuda_v128/.devcontainer/devcontainer.json

@@ -0,0 +1,21 @@
+// For format details, see https://aka.ms/devcontainer.json. For config options, see the
+// README at: https://github.com/devcontainers/templates/tree/main/src/docker-existing-dockerfile
+{
+	"name": "sleap-nn v002 CUDA v12.8.0",
+	"build": {
+		"context": "..",
+		"dockerfile": "../Dockerfile"
+	},
+    "runArgs": [
+        "--gpus=all"
+    ],
+	"customizations": {
+		"vscode": {
+		"settings": {
+			"terminal.integrated.defaultProfile.linux": "bash"
+		}
+		}
+	},
+	"postCreateCommand": "echo 'Devcontainer ready for use!'",
+	"remoteUser": "root"
+	}

sleapnn_v002_cuda_v128/.dockerignore

@@ -0,0 +1,19 @@
+# Ignore Python cache
+__pycache__/
+*.pyc
+*.pyo
+
+# Ignore virtual environments
+env/
+venv/
+
+# Ignore version control and logs
+.git/
+*.log
+
+# Ignore Docker build artifacts
+docker/*.tmp
+
+README.md
+.gitignore
+terraform/*

sleapnn_v002_cuda_v128/Dockerfile

@@ -0,0 +1,69 @@
+# Base image with GPU support
+FROM nvidia/cuda:12.8.1-cudnn-devel-ubuntu24.04
+
+# Set working directory
+WORKDIR /app
+
+# Set non-interactive mode
+ENV DEBIAN_FRONTEND=noninteractive
+
+# Set user
+ENV USER=root
+
+# Set NVIDIA driver capabilities
+ENV NVIDIA_DRIVER_CAPABILITIES=all
+
+# Set QT debug environment variable to help debug issues with Qt plugins
+ENV QT_DEBUG_PLUGINS=1
+
+# Set virtual environment
+ENV VIRTUAL_ENV=/app/.venv
+
+# Install uv + dependencies
+# opencv requires opengl https://github.com/conda-forge/opencv-feedstock/issues/401
+# Default python3 is 3.8 in ubuntu 20.04 https://wiki.ubuntu.com/FocalFossa/ReleaseNotes#Python3_by_default
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    build-essential \
+    libgl1 \
+    libglx0 \
+    libglapi-mesa \
+    libegl-mesa0 \
+    libegl1 \
+    libopengl0 \
+    libglib2.0-0 \
+    libfontconfig1 \
+    libgssapi-krb5-2 \
+    libdbus-1-3 \
+    libx11-xcb1 \
+    libxkbcommon-x11-0 \
+    python3 \
+    python3-venv \
+    python3-pip \
+    git \
+    curl && \
+    apt-get clean && \
+    rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/*
+
+# Install uv
+RUN curl -LsSf https://astral.sh/uv/install.sh | sh \
+    && ln -s /root/.local/bin/uv /usr/local/bin/uv
+
+# Initialize uv project
+RUN uv init --python 3.13 --no-readme
+
+# Create virtual env w/ Python 3.13
+RUN uv venv
+
+# Set virtual environment path
+ENV PATH="$VIRTUAL_ENV/bin:$PATH" \
+    PYTHONDONTWRITEBYTECODE=1 \
+    PIP_NO_CACHE_DIR=1
+
+# Add MarkupSafe from git
+RUN uv add git+https://github.com/pallets/markupsafe@3.0.2
+
+# Clean up apt cache and unnecessary files to free space
+RUN apt-get clean && rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/*
+
+# Install sleap-nn w/ CUDA 12.8 support
+RUN uv add sleap-nn[torch] --index https://download.pytorch.org/whl/cu128 --index https://pypi.org/simple --no-cache

sleapnn_v002_cuda_v128/README.md

@@ -0,0 +1,109 @@
+# sleap-nn-cuda
+
+
+## Description
+This repo contains a DockerFile for a lightweight container (~6 GB) with the PyPI installation of SLEAP and all of its dependencies. The container repository is located at [https://hub.docker.com/repository/docker/eberrigan/sleap-cuda/general](https://hub.docker.com/repository/docker/eberrigan/sleap-cuda/general).
+
+The base image used is [nvidia/cuda:12.8.1-cudnn-devel-ubuntu24.04](https://hub.docker.com/layers/nvidia/cuda/12.8.1-cudnn-devel-ubuntu24.04/images/sha256-3986465b3dd3b4d602c07061f2cff417e0bfb24810129408d4eb12e111015a6c).
+- The Dockerfile is located at `./sleapnn_v002_cuda_v128/Dockerfile`.
+- The repo has CI set up in `.github/workflows` for building and pushing the image when making changes.
+  - The workflow uses the linux/amd64 platform to build. 
+- `.devcontainer/devcontainer.json` is convenient for developing inside a container made with the DockerFile using Visual Studio Code.
+
+
+## Installation
+
+**Make sure to have Docker Daemon running first**
+
+
+You can pull the image if you don't have it built locally, or need to update the latest, with
+
+```
+docker pull ghcr.io/talmolab/sleap-rtc-worker:latest
+```
+
+## Usage 
+
+Then, to run the image with gpus interactively:
+
+```
+docker run --gpus all -it ghcr.io/talmolab/sleap-rtc-worker:latest bash
+```
+
+and test with 
+
+```
+python -c "import sleap; sleap.versions()" && nvidia-smi
+```
+
+In general, use the syntax
+
+```
+docker run -v /path/on/host:/path/in/container [other options] image_name [command]
+```
+
+Note that host paths are absolute. 
+
+
+Use this syntax to give host permissions to mounted volumes
+```
+docker run -u $(id -u):$(id -g) -v /your/host/directory:/container/directory [options] your-image-name [command]
+```
+
+```
+docker run -u $(id -u):$(id -g) -v ./tests/data:/tests/data --gpus all -it ghcr.io/talmolab/sleap-rtc-worker:latest bash
+```
+
+Test:
+
+```
+ python3 -c "import sleap; print('SLEAP version:', sleap.__version__)"
+ nvidia-smi # Check that the GPUs are discoverable
+ sleap-nn train "tests/data/initial_config.json" "tests/data/dance.mp4.labels.slp" --video-paths "tests/data/dance.mp4"
+```
+
+**Notes:**
+
+- The `ghcr.io/talmolab/sleap-rtc-worker` is the Docker registry where the images are pulled from. This is only used when pulling images from the cloud, and not necesary when building/running locally.
+- `-it` ensures that you get an interactive terminal. The `i` stands for interactive, and `t` allocates a pseudo-TTY, which is what allows you to interact with the bash shell inside the container.
+- The `-v` or `--volume` option mounts the specified directory with the same level of access as the directory has on the host.
+- `bash` is the command that gets executed inside the container, which in this case is to start the bash shell.
+- Order of operations is 1. Pull (if needed): Get a pre-built image from a registry. 2. Run: Start a container from an image.
+
+## Contributing
+
+- Use the `devcontainer.json` to open the repo in a dev container using VS Code.
+  - There is some test data in the `tests` directory that will be automatically mounted for use since the working directory is the workspace.
+  - Rebuild the container when you make changes using `Dev Container: Rebuild Container`.
+
+- Please make a new branch, starting with your name, with any changes, and request a reviewer before merging with the main branch since this image will be used by others.
+- Please document using the same conventions (docstrings for each function and class, typing-hints, informative comments).
+- Tests are written in the pytest framework. Data used in the tests are defined as fixtures in `tests/fixtures/data.py` (https://docs.pytest.org/en/6.2.x/reference.html#fixtures-api).
+
+
+## Build
+To build and push via automated CI, just push changes to a branch. 
+- Pushes to `main` result in an image with the tag `latest`. 
+- Pushes to other branches have tags with `-test` appended. 
+- See `.github/workflows` for testing and production workflows.
+
+To test `test` images locally use after pushing the `test` images via CI:
+
+```
+docker pull eberrigan/sleap-cuda:linux-amd64-test
+```
+
+then 
+
+```
+docker run -v ./tests/data:/tests/data --gpus all -it eberrigan/sleap-cuda:linux-amd64-test bash
+```
+
+To build locally for testing you can use the command (from the root of the repo):
+
+```
+docker build --platform linux/amd64 ./sleap_cuda
+```
+
+## Support
+contact Elizabeth at eberrigan@salk.edu