Merge pull request #70 from aqlaboratory/jandom/2025-12/build/docker-layering-improvements

build(docker): layering and caching improvements

Author: jandom

.github/workflows/ci-test-reusable.yml

@@ -0,0 +1,103 @@
+name: Reusable Docker Test
+
+on:
+  # Can only be called by another workflow, not directly by the user
+  workflow_call:
+    inputs:
+      cuda_base_image_tag:
+        description: 'CUDA base image tag (e.g., 12.2.2-cudnn8-devel-ubuntu22.04)'
+        required: true
+        type: string
+
+env:
+  REGISTRY: ghcr.io
+  IMAGE_NAME: ${{ github.repository }}/openfold3-docker
+
+jobs:
+  start-aws-runner:
+    runs-on: ubuntu-latest
+    permissions:
+      id-token: write
+      contents: read
+    outputs:
+      mapping: ${{ steps.aws-start.outputs.mapping }}
+      instances: ${{ steps.aws-start.outputs.instances }}
+    steps:
+      - name: Configure AWS credentials
+        uses: aws-actions/configure-aws-credentials@v5
+        with:
+          role-to-assume: arn:aws:iam::203627415330:role/of-gha-runner
+          aws-region: us-east-1
+      - name: Create cloud runner
+        id: aws-start
+        uses: omsf/start-aws-gha-runner@v1.1.0
+        with:
+          aws_image_id: ami-0754c6e75b3b97dcd  # Deep Learning AMI Neuron (Ubuntu 22.04)
+          aws_instance_type: t3.2xlarge
+          aws_home_dir: /home/ubuntu
+          aws_root_device_size: 200
+        env:
+          GH_PAT: ${{ secrets.GH_PAT }}
+
+  test-openfold-docker:
+    runs-on: ${{ fromJSON(needs.start-aws-runner.outputs.instances) }}
+    needs:
+      - start-aws-runner
+    permissions:
+      contents: read
+      packages: write
+    steps:
+      - uses: actions/checkout@v6
+
+      - name: Log in to GHCR
+        uses: docker/login-action@v3
+        with:
+          registry: ${{ env.REGISTRY }}
+          username: ${{ github.actor }}
+          password: ${{ secrets.GITHUB_TOKEN }}
+
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v3
+
+      - name: Build and push test image
+        uses: docker/build-push-action@v6
+        with:
+          context: .
+          file: docker/Dockerfile
+          target: test
+          push: true
+          build-args: |
+            CUDA_BASE_IMAGE_TAG=${{ inputs.cuda_base_image_tag }}
+          tags: |
+            ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}:test-${{ inputs.cuda_base_image_tag }}-${{ github.sha }}
+          cache-from: type=registry,ref=${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}:cache-${{ inputs.cuda_base_image_tag }}
+          cache-to: type=registry,ref=${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}:cache-${{ inputs.cuda_base_image_tag }},mode=max
+
+      - name: Run unit tests
+        run: |
+          docker run \
+            -v ${{ github.workspace }}:/opt/openfold3 \
+            ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}:test-${{ inputs.cuda_base_image_tag }}-${{ github.sha }} \
+            pytest openfold3/tests -vvv
+
+  stop-aws-runner:
+    runs-on: ubuntu-latest
+    permissions:
+        id-token: write
+        contents: read
+    needs:
+      - start-aws-runner
+      - test-openfold-docker
+    if: ${{ always() }}
+    steps:
+      - name: Configure AWS credentials
+        uses: aws-actions/configure-aws-credentials@v5
+        with:
+          role-to-assume: arn:aws:iam::203627415330:role/of-gha-runner
+          aws-region: us-east-1
+      - name: Stop instances
+        uses: omsf/stop-aws-gha-runner@v1.0.0
+        with:
+          instance_mapping: ${{ needs.start-aws-runner.outputs.mapping }}
+        env:
+          GH_PAT: ${{ secrets.GH_PAT }}

.github/workflows/ci-test.yml

@@ -1,4 +1,4 @@
-name: Run Tests with Docker image 
+name: Run Tests with Docker image
 
 on:
   pull_request_target:
@@ -8,80 +8,25 @@ on:
   workflow_dispatch:
 
 jobs:
-  start-aws-runner:
-    runs-on: ubuntu-latest
+  test:
     if: |
       (github.event_name == 'push') ||
       (github.event_name == 'workflow_dispatch') ||
-      (github.event_name == 'pull_request_target' && 
+      (github.event_name == 'pull_request_target' &&
        github.event.action == 'labeled' &&
        github.event.label.name == 'safe-to-test')
     permissions:
       id-token: write
       contents: read
-    outputs:
-      mapping: ${{ steps.aws-start.outputs.mapping }}
-      instances: ${{ steps.aws-start.outputs.instances }}
-    steps:
-      - name: Configure AWS credentials
-        uses: aws-actions/configure-aws-credentials@v5
-        with:
-          role-to-assume: arn:aws:iam::203627415330:role/of-gha-runner 
-          aws-region: us-east-1
-      - name: Create cloud runner
-        id: aws-start
-        uses: omsf/start-aws-gha-runner@v1.1.0
-        with:
-          aws_image_id: ami-0754c6e75b3b97dcd  # Deep Learning AMI Neuron (Ubuntu 22.04)
-          aws_instance_type: t3.2xlarge 
-          aws_home_dir: /home/ubuntu
-        env:
-          GH_PAT: ${{ secrets.GH_PAT }}
-    
-  test-openfold-docker:
-    runs-on: ${{ fromJSON(needs.start-aws-runner.outputs.instances) }}
-    needs:
-      - start-aws-runner
-    steps:
-      - uses: actions/checkout@v6
-        with:
-          ref: ${{ github.event_name == 'pull_request_target' && 
-                   format('refs/pull/{0}/merge', github.event.pull_request.number) || 
-                   github.ref }}
-
-      - name: Log in to Docker Hub
-        uses: docker/login-action@v3.6.0
-        with:
-          username: ${{ secrets.DOCKERHUB_USERNAME }}
-          password: ${{ secrets.DOCKERHUB_TOKEN }}
-      
-      - name: Pull pre-built Docker image
-        run: docker pull openfoldconsortium/openfold3:stable
-      
-      - name: Build test layers on Docker image 
-        run: docker build -t openfold3-test-runner -f openfold3/tests/Dockerfile . 
-      
-      - name: Run unit tests
-        run: docker run -v ${{ github.workspace }}:/opt/openfold3 openfold3-test-runner:latest pytest openfold3/tests
-
-  stop-aws-runner:
-    runs-on: ubuntu-latest
-    permissions:
-        id-token: write
-        contents: read
-    needs:
-      - start-aws-runner
-      - test-openfold-docker 
-    if: ${{ always() }}
-    steps:
-      - name: Configure AWS credentials
-        uses: aws-actions/configure-aws-credentials@v5
-        with:
-          role-to-assume: arn:aws:iam::203627415330:role/of-gha-runner 
-          aws-region: us-east-1
-      - name: Stop instances
-        uses: omsf/stop-aws-gha-runner@v1.0.0
-        with:
-          instance_mapping: ${{ needs.start-aws-runner.outputs.mapping }}
-        env:
-          GH_PAT: ${{ secrets.GH_PAT }}
+      packages: write
+    strategy:
+      matrix:
+        include:
+          - cuda_base_image_tag: "12.1.1-cudnn8-devel-ubuntu22.04"
+    concurrency:
+      group: ${{ github.workflow }}-${{ github.head_ref || github.ref }}-${{ matrix.cuda_base_image_tag }}
+      cancel-in-progress: true
+    uses: ./.github/workflows/ci-test-reusable.yml
+    with:
+      cuda_base_image_tag: ${{ matrix.cuda_base_image_tag }}
+    secrets: inherit

Build_instructions_blackwell.md docker/Build_instructions_blackwell.mdBuild_instructions_blackwell.md renamed to docker/Build_instructions_blackwell.md

@@ -0,0 +1,35 @@
+## Production images
+
+TODO
+
+For Blackwell image build, see [Build_instructions_blackwell.md](Build_instructions_blackwell.md)
+
+## Development images
+
+These images are the biggest but come with all the build tooling, needed to compile things at runtime (Deepspeed)
+
+```
+docker build \
+    -f docker/Dockerfile \
+    --target devel \
+    -t openfold-docker:devel .
+```
+
+## Test images
+
+Build the test image
+```
+docker build \
+    -f docker/development/Dockerfile \
+    --target test \
+    -t openfold-docker:test .
+```
+
+Run the unit tests
+```
+docker run \
+    --rm \
+    -v $(pwd -P):/opt/openfold3 \
+    -t openfold-docker:test \
+    pytest openfold3/tests -vvv
+```

docker/DOCKER.md

@@ -1,5 +1,6 @@
 # Full performance multi-stage build with complete CUDA toolchain
-FROM nvidia/cuda:12.1.1-cudnn8-devel-ubuntu22.04 AS builder
+ARG CUDA_BASE_IMAGE_TAG=12.2.2-cudnn8-devel-ubuntu22.04
+FROM nvidia/cuda:${CUDA_BASE_IMAGE_TAG} AS builder
 
 # Install complete build dependencies including CUDA compiler tools
 RUN apt-get update && apt-get install -y \
@@ -13,29 +14,35 @@ RUN apt-get update && apt-get install -y \
     && rm -rf /var/lib/apt/lists/*
 
 # Install miniforge
+# FIXME this needs to be pinned, with more recent versions (25.11.0-1) the package resolution is stuck
 RUN wget -P /tmp \
     "https://github.com/conda-forge/miniforge/releases/download/25.3.1-0/Miniforge3-Linux-x86_64.sh" \
     && bash /tmp/Miniforge3-Linux-x86_64.sh -b -p /opt/conda \
     && rm /tmp/Miniforge3-Linux-x86_64.sh
 
 ENV PATH=/opt/conda/bin:$PATH
+ENV CONDA_PREFIX=/opt/conda
 
 # Copy and install dependencies with aggressive cleanup
 COPY environments/production.yml /opt/openfold3/environment.yml
 RUN mamba env update -n base --file /opt/openfold3/environment.yml \
     && mamba clean --all --yes \
     && conda clean --all --yes
 
-# Copy the entire source tree
-COPY . /opt/openfold3/
+# Copy the minimal set of files needed to install the package
+COPY setup.py /opt/openfold3/
+COPY pyproject.toml /opt/openfold3/
+COPY openfold3/__init__.py /opt/openfold3/openfold3/
+COPY scripts/ /opt/openfold3/scripts/
 
 # Install third party dependencies
 WORKDIR /opt/
 RUN /opt/openfold3/scripts/install_third_party_dependencies.sh
 
 # Install the package
 WORKDIR /opt/openfold3
-RUN python3 setup.py install
+# even `pip install --no-build-isolation` not actually working here, needs investigation
+RUN python setup.py install
 
 # Set CUDA architecture for compilation (adjust based on your GPU)
 ENV TORCH_CUDA_ARCH_LIST="8.0;8.6;9.0"
@@ -44,10 +51,11 @@ ENV TORCH_CUDA_ARCH_LIST="8.0;8.6;9.0"
 # RUN python3 -c "import deepspeed; deepspeed.ops.op_builder.EvoformerAttnBuilder().load()" || \
 #     python3 -c "import deepspeed; print('DeepSpeed ops loaded successfully')"
 
-# Runtime stage - use devel image for full CUDA support
-FROM nvidia/cuda:12.1.1-cudnn8-devel-ubuntu22.04 AS runtime
+# Devel stage - use devel image for full CUDA support
+ARG CUDA_BASE_IMAGE_TAG=12.2.2-cudnn8-devel-ubuntu22.04
+FROM nvidia/cuda:${CUDA_BASE_IMAGE_TAG} AS devel
 
-# Install runtime dependencies
+# Install devel dependencies
 RUN apt-get update && apt-get install -y \
     libopenmpi3 \
     libaio1 \
@@ -85,7 +93,16 @@ ENV KMP_AFFINITY=none
 ENV LIBRARY_PATH=/opt/conda/lib:$LIBRARY_PATH
 ENV LD_LIBRARY_PATH=/opt/conda/lib:$LD_LIBRARY_PATH
 
-# Copy the entire source tree to the runtime image
-COPY --from=builder /opt/openfold3 /opt/openfold3
+# Copy the entire source tree directly (at the very end for optimal caching)
+COPY . /opt/openfold3
 
 WORKDIR /opt/openfold3
+
+# Test stage - build on devel layer with test dependencies
+FROM devel AS test
+
+COPY environments/requirements-test.txt /opt/openfold3/requirements-test.txt
+
+WORKDIR /opt/openfold3
+RUN pip install -r requirements-test.txt
+RUN pip install --no-deps --editable .

Dockerfile docker/DockerfileDockerfile renamed to docker/Dockerfile

@@ -1,14 +1,19 @@
 #!/bin/bash
 
-python setup.py install
+# -e (errexit): Exit immediately if any command returns a non-zero exit status
+# -u (nounset): Exit if you try to use an uninitialized variable
+# -o pipefail: Exit if any command in a pipeline fails (not just the last one)
+set -euo pipefail
 
+# These are necessary for subsequent (runtime) compilation of Deepspeed
 echo "Download CUTLASS, required for Deepspeed Evoformer attention kernel"
 git clone https://github.com/NVIDIA/cutlass --branch v3.6.0 --depth 1
 conda env config vars set CUTLASS_PATH=$PWD/cutlass
 
 # This setting is used to fix a worker assignment issue during data loading
 conda env config vars set KMP_AFFINITY=none
 
+# These will only be available outside of this script if it's sourced in the current shell
 export LIBRARY_PATH=$CONDA_PREFIX/lib:$LIBRARY_PATH
 export LD_LIBRARY_PATH=$CONDA_PREFIX/lib:$LD_LIBRARY_PATH
 export PATH=$PATH:/sbin  # TODO: Check if this is necessary, or is NERSC-specific