Build and publish cuda enabled docker image (#291)

.github/workflows/build-and-push-images.yml

@@ -0,0 +1,186 @@
+name: Build and push images
+
+on:
+  push:
+    branches:
+      - master
+
+env:
+  CUDA_ARCHS: '89,120'
+
+jobs:
+  image_meta:
+    name: Get image metadata
+    runs-on: ubuntu-latest
+    outputs:
+      sha_tag: ${{ steps.meta.outputs.sha_tag }}
+      registry: ${{ steps.meta.outputs.registry }}
+    steps:
+      - name: Get image metadata
+        id: meta
+        run: |
+          GIT_SHA="${{ github.sha }}"
+          echo "sha_tag=${GIT_SHA:0:7}" >> $GITHUB_OUTPUT
+          echo "registry=ghcr.io/${{ github.repository }}" >> $GITHUB_OUTPUT
+
+  build_and_push:
+    name: Build and push docker image (${{ matrix.zkvm }})
+    needs: image_meta
+    runs-on: ubuntu-latest
+    permissions:
+      contents: read
+      packages: write
+    strategy:
+      fail-fast: false
+      matrix:
+        zkvm:
+          - airbender
+          - jolt
+          - miden
+          - nexus
+          - openvm
+          - pico
+          - risc0
+          - sp1
+          - ziren
+          - zisk
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v4
+
+      - name: Free up disk space
+        run: bash .github/scripts/free-up-disk-space.sh
+
+      - name: Log in to GitHub Container Registry
+        uses: docker/login-action@v3
+        with:
+          registry: ghcr.io
+          username: ${{ github.actor }}
+          password: ${{ secrets.GITHUB_TOKEN }}
+
+      - name: Build ere-base and ere-base-${{ matrix.zkvm }} images
+        run: |
+          bash .github/scripts/build-image.sh \
+            --zkvm ${{ matrix.zkvm }} \
+            --registry ${{ needs.image_meta.outputs.registry }} \
+            --tag ${{ needs.image_meta.outputs.sha_tag }} \
+            --base
+
+      - name: Push ere-base and ere-base-${{ matrix.zkvm }} images
+        run: |
+          docker push ${{ needs.image_meta.outputs.registry }}/ere-base-${{ matrix.zkvm }}:${{ needs.image_meta.outputs.sha_tag }}
+
+      - name: Build ere-compiler-${{ matrix.zkvm }} and ere-server-${{ matrix.zkvm }} images
+        run: |
+          bash .github/scripts/build-image.sh \
+            --zkvm ${{ matrix.zkvm }} \
+            --registry ${{ needs.image_meta.outputs.registry }} \
+            --tag ${{ needs.image_meta.outputs.sha_tag }} \
+            --compiler \
+            --server
+
+      - name: Push ere-compiler-${{ matrix.zkvm }} and ere-server-${{ matrix.zkvm }} images
+        run: |
+          docker push ${{ needs.image_meta.outputs.registry }}/ere-compiler-${{ matrix.zkvm }}:${{ needs.image_meta.outputs.sha_tag }}
+          docker push ${{ needs.image_meta.outputs.registry }}/ere-server-${{ matrix.zkvm }}:${{ needs.image_meta.outputs.sha_tag }}
+
+  build_and_push_cuda:
+    name: Build and push CUDA docker image (${{ matrix.zkvm }})
+    needs: image_meta
+    runs-on: ubuntu-latest
+    permissions:
+      contents: read
+      packages: write
+    strategy:
+      fail-fast: false
+      matrix:
+        zkvm:
+          - airbender
+          - openvm
+          - risc0
+          - sp1
+          - zisk
+        include:
+          - zkvm: zisk
+            cuda_archs: '120'
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v4
+
+      - name: Free up disk space
+        run: bash .github/scripts/free-up-disk-space.sh
+
+      - name: Log in to GitHub Container Registry
+        uses: docker/login-action@v3
+        with:
+          registry: ghcr.io
+          username: ${{ github.actor }}
+          password: ${{ secrets.GITHUB_TOKEN }}
+
+      - name: Build ere-base and ere-base-${{ matrix.zkvm }} images with CUDA enabled
+        run: |
+          bash .github/scripts/build-image.sh \
+            --zkvm ${{ matrix.zkvm }} \
+            --registry ${{ needs.image_meta.outputs.registry }} \
+            --tag ${{ needs.image_meta.outputs.sha_tag }}-cuda \
+            --base \
+            --cuda-archs '${{ matrix.cuda_archs || env.CUDA_ARCHS }}'
+
+      - name: Push ere-base and ere-base-${{ matrix.zkvm }} images with CUDA enabled
+        run: |
+          docker push ${{ needs.image_meta.outputs.registry }}/ere-base-${{ matrix.zkvm }}:${{ needs.image_meta.outputs.sha_tag }}-cuda
+
+      - name: Build ere-server-${{ matrix.zkvm }} image with CUDA enabled
+        run: |
+          bash .github/scripts/build-image.sh \
+            --zkvm ${{ matrix.zkvm }} \
+            --registry ${{ needs.image_meta.outputs.registry }} \
+            --tag ${{ needs.image_meta.outputs.sha_tag }}-cuda \
+            --server \
+            --cuda-archs '${{ matrix.cuda_archs || env.CUDA_ARCHS }}'
+
+      - name: Push ere-server-${{ matrix.zkvm }} image with CUDA enabled
+        run: |
+          docker push ${{ needs.image_meta.outputs.registry }}/ere-server-${{ matrix.zkvm }}:${{ needs.image_meta.outputs.sha_tag }}-cuda
+
+  build_and_push_cluster:
+    name: Build and push cluster docker image (${{ matrix.zkvm }})
+    needs: image_meta
+    runs-on: ubuntu-latest
+    permissions:
+      contents: read
+      packages: write
+    strategy:
+      fail-fast: false
+      matrix:
+        zkvm:
+          - zisk
+        include:
+          - zkvm: zisk
+            cuda_archs: '120'
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v4
+
+      - name: Free up disk space
+        run: bash .github/scripts/free-up-disk-space.sh
+
+      - name: Log in to GitHub Container Registry
+        uses: docker/login-action@v3
+        with:
+          registry: ghcr.io
+          username: ${{ github.actor }}
+          password: ${{ secrets.GITHUB_TOKEN }}
+
+      - name: Build ere-cluster-${{ matrix.zkvm }} image with CUDA enabled
+        run: |
+          bash .github/scripts/build-image.sh \
+            --zkvm ${{ matrix.zkvm }} \
+            --registry ${{ needs.image_meta.outputs.registry }} \
+            --tag ${{ needs.image_meta.outputs.sha_tag }}-cuda \
+            --cluster \
+            --cuda-archs '${{ matrix.cuda_archs || env.CUDA_ARCHS }}'
+
+      - name: Push ere-cluster-${{ matrix.zkvm }} image with CUDA enabled
+        run: |
+          docker push ${{ needs.image_meta.outputs.registry }}/ere-cluster-${{ matrix.zkvm }}:${{ needs.image_meta.outputs.sha_tag }}-cuda

.github/workflows/check-fmt.yml

@@ -2,7 +2,7 @@ name: Check formatting
 
 concurrency:
   group: ${{ github.workflow }}-${{ github.ref }}
-  cancel-in-progress: true
+  cancel-in-progress: ${{ github.event_name == 'pull_request' }}
 
 on:
   push:

.github/workflows/push-semver-tag.yml

@@ -0,0 +1,90 @@
+name: Push SemVer tag
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: ${{ github.event_name == 'pull_request' }}
+
+on:
+  push:
+    tags:
+      - 'v*'
+
+jobs:
+  image_meta:
+    name: Get image metadata
+    runs-on: ubuntu-latest
+    outputs:
+      sha_tag: ${{ steps.meta.outputs.sha_tag }}
+      semver_tag: ${{ steps.meta.outputs.semver_tag }}
+      registry: ${{ steps.meta.outputs.registry }}
+    steps:
+      - name: Get image metadata
+        id: meta
+        run: |
+          GIT_SHA="${{ github.sha }}"
+          GIT_TAG="${{ github.ref_name }}"
+          echo "sha_tag=${GIT_SHA:0:7}" >> $GITHUB_OUTPUT
+          echo "semver_tag=${GIT_TAG#v}" >> $GITHUB_OUTPUT
+          echo "registry=ghcr.io/${{ github.repository }}" >> $GITHUB_OUTPUT
+
+  create_semver_tag:
+    name: Tag images with SemVer (${{ matrix.zkvm }})
+    needs: image_meta
+    runs-on: ubuntu-latest
+    permissions:
+      contents: read
+      packages: write
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - { zkvm: airbender, cuda: true , cluster: false }
+          - { zkvm: jolt,      cuda: false, cluster: false }
+          - { zkvm: miden,     cuda: false, cluster: false }
+          - { zkvm: nexus,     cuda: false, cluster: false }
+          - { zkvm: openvm,    cuda: true , cluster: false }
+          - { zkvm: pico,      cuda: false, cluster: false }
+          - { zkvm: risc0,     cuda: true , cluster: false }
+          - { zkvm: sp1,       cuda: true , cluster: false }
+          - { zkvm: ziren,     cuda: false, cluster: false }
+          - { zkvm: zisk,      cuda: true , cluster: true  }
+    steps:
+      - name: Log in to GitHub Container Registry
+        uses: docker/login-action@v3
+        with:
+          registry: ghcr.io
+          username: ${{ github.actor }}
+          password: ${{ secrets.GITHUB_TOKEN }}
+
+      - name: Add SemVer tag to images
+        run: |
+          SHA_TAG="${{ needs.image_meta.outputs.sha_tag }}"
+          SEMVER_TAG="${{ needs.image_meta.outputs.semver_tag }}"
+
+          for IMAGE_KIND in base compiler server; do
+            IMAGE="${{ needs.image_meta.outputs.registry }}/ere-$IMAGE_KIND-${{ matrix.zkvm }}"
+            echo "Tagging $IMAGE:$SHA_TAG as $IMAGE:$SEMVER_TAG"
+            docker buildx imagetools create "$IMAGE:$SHA_TAG" --tag "$IMAGE:$SEMVER_TAG"
+          done
+
+      - name: Add SemVer tag to images with CUDA enabled
+        if: matrix.cuda
+        run: |
+          SHA_TAG="${{ needs.image_meta.outputs.sha_tag }}-cuda"
+          SEMVER_TAG="${{ needs.image_meta.outputs.semver_tag }}-cuda"
+
+          for IMAGE_KIND in base server; do
+            IMAGE="${{ needs.image_meta.outputs.registry }}/ere-$IMAGE_KIND-${{ matrix.zkvm }}"
+            echo "Tagging $IMAGE:$SHA_TAG as $IMAGE:$SEMVER_TAG"
+            docker buildx imagetools create "$IMAGE:$SHA_TAG" --tag "$IMAGE:$SEMVER_TAG"
+          done
+
+      - name: Add SemVer tag to cluster image with CUDA enabled
+        if: matrix.cluster
+        run: |
+          SHA_TAG="${{ needs.image_meta.outputs.sha_tag }}-cuda"
+          SEMVER_TAG="${{ needs.image_meta.outputs.semver_tag }}-cuda"
+
+          IMAGE="${{ needs.image_meta.outputs.registry }}/ere-cluster-${{ matrix.zkvm }}"
+          echo "Tagging $IMAGE:$SHA_TAG as $IMAGE:$SEMVER_TAG"
+          docker buildx imagetools create "$IMAGE:$SHA_TAG" --tag "$IMAGE:$SEMVER_TAG"

.github/workflows/test-common.yml

@@ -2,7 +2,7 @@ name: Test and clippy common crates
 
 concurrency:
   group: ${{ github.workflow }}-${{ github.ref }}
-  cancel-in-progress: true
+  cancel-in-progress: ${{ github.event_name == 'pull_request' }}
 
 on:
   push:

.github/workflows/test-zkvm-airbender.yml

@@ -2,22 +2,19 @@ name: Test and clippy Airbender
 
 concurrency:
   group: ${{ github.workflow }}-${{ github.ref }}
-  cancel-in-progress: true
+  cancel-in-progress: ${{ github.event_name == 'pull_request' }}
 
 on:
-  push:
-    branches:
-      - master
-    tags:
-      - 'v*'
+  workflow_run:
+    workflows: ['Build and push images']
+    types: [completed]
   pull_request:
 
 jobs:
   test:
+    if: github.event_name == 'pull_request' || github.event.workflow_run.conclusion == 'success'
     uses: ./.github/workflows/test-zkvm.yml
-    permissions:
-      contents: read
-      packages: write
     with:
       zkvm: airbender
       skip_prove_test: true
+      cuda: true

.github/workflows/test-zkvm-jolt.yml

@@ -2,21 +2,17 @@ name: Test and clippy Jolt
 
 concurrency:
   group: ${{ github.workflow }}-${{ github.ref }}
-  cancel-in-progress: true
+  cancel-in-progress: ${{ github.event_name == 'pull_request' }}
 
 on:
-  push:
-    branches:
-      - master
-    tags:
-      - 'v*'
+  workflow_run:
+    workflows: ['Build and push images']
+    types: [completed]
   pull_request:
 
 jobs:
   test:
+    if: github.event_name == 'pull_request' || github.event.workflow_run.conclusion == 'success'
     uses: ./.github/workflows/test-zkvm.yml
-    permissions:
-      contents: read
-      packages: write
     with:
       zkvm: jolt

.github/workflows/test-zkvm-miden.yml

@@ -2,21 +2,17 @@ name: Test and clippy Miden
 
 concurrency:
   group: ${{ github.workflow }}-${{ github.ref }}
-  cancel-in-progress: true
+  cancel-in-progress: ${{ github.event_name == 'pull_request' }}
 
 on:
-  push:
-    branches:
-      - master
-    tags:
-      - 'v*'
+  workflow_run:
+    workflows: ['Build and push images']
+    types: [completed]
   pull_request:
 
 jobs:
   test:
+    if: github.event_name == 'pull_request' || github.event.workflow_run.conclusion == 'success'
     uses: ./.github/workflows/test-zkvm.yml
-    permissions:
-      contents: read
-      packages: write
     with:
       zkvm: miden

.github/workflows/test-zkvm-nexus.yml

@@ -2,21 +2,17 @@ name: Test and clippy Nexus
 
 concurrency:
   group: ${{ github.workflow }}-${{ github.ref }}
-  cancel-in-progress: true
+  cancel-in-progress: ${{ github.event_name == 'pull_request' }}
 
 on:
-  push:
-    branches:
-      - master
-    tags:
-      - 'v*'
+  workflow_run:
+    workflows: ['Build and push images']
+    types: [completed]
   pull_request:
 
 jobs:
   test:
+    if: github.event_name == 'pull_request' || github.event.workflow_run.conclusion == 'success'
     uses: ./.github/workflows/test-zkvm.yml
-    permissions:
-      contents: read
-      packages: write
     with:
       zkvm: nexus

.github/workflows/test-zkvm-openvm.yml

@@ -2,21 +2,18 @@ name: Test and clippy OpenVM
 
 concurrency:
   group: ${{ github.workflow }}-${{ github.ref }}
-  cancel-in-progress: true
+  cancel-in-progress: ${{ github.event_name == 'pull_request' }}
 
 on:
-  push:
-    branches:
-      - master
-    tags:
-      - 'v*'
+  workflow_run:
+    workflows: ['Build and push images']
+    types: [completed]
   pull_request:
 
 jobs:
   test:
+    if: github.event_name == 'pull_request' || github.event.workflow_run.conclusion == 'success'
     uses: ./.github/workflows/test-zkvm.yml
-    permissions:
-      contents: read
-      packages: write
     with:
       zkvm: openvm
+      cuda: true

.github/workflows/test-zkvm-pico.yml

@@ -2,21 +2,17 @@ name: Test and clippy Pico
 
 concurrency:
   group: ${{ github.workflow }}-${{ github.ref }}
-  cancel-in-progress: true
+  cancel-in-progress: ${{ github.event_name == 'pull_request' }}
 
 on:
-  push:
-    branches:
-      - master
-    tags:
-      - 'v*'
+  workflow_run:
+    workflows: ['Build and push images']
+    types: [completed]
   pull_request:
 
 jobs:
   test:
+    if: github.event_name == 'pull_request' || github.event.workflow_run.conclusion == 'success'
     uses: ./.github/workflows/test-zkvm.yml
-    permissions:
-      contents: read
-      packages: write
     with:
       zkvm: pico

.github/workflows/test-zkvm-risc0.yml

@@ -2,21 +2,18 @@ name: Test and clippy Risc0
 
 concurrency:
   group: ${{ github.workflow }}-${{ github.ref }}
-  cancel-in-progress: true
+  cancel-in-progress: ${{ github.event_name == 'pull_request' }}
 
 on:
-  push:
-    branches:
-      - master
-    tags:
-      - 'v*'
+  workflow_run:
+    workflows: ['Build and push images']
+    types: [completed]
   pull_request:
 
 jobs:
   test:
+    if: github.event_name == 'pull_request' || github.event.workflow_run.conclusion == 'success'
     uses: ./.github/workflows/test-zkvm.yml
-    permissions:
-      contents: read
-      packages: write
     with:
       zkvm: risc0
+      cuda: true

.github/workflows/test-zkvm-sp1.yml

@@ -2,21 +2,18 @@ name: Test and clippy SP1
 
 concurrency:
   group: ${{ github.workflow }}-${{ github.ref }}
-  cancel-in-progress: true
+  cancel-in-progress: ${{ github.event_name == 'pull_request' }}
 
 on:
-  push:
-    branches:
-      - master
-    tags:
-      - 'v*'
+  workflow_run:
+    workflows: ['Build and push images']
+    types: [completed]
   pull_request:
 
 jobs:
   test:
+    if: github.event_name == 'pull_request' || github.event.workflow_run.conclusion == 'success'
     uses: ./.github/workflows/test-zkvm.yml
-    permissions:
-      contents: read
-      packages: write
     with:
       zkvm: sp1
+      cuda: true

.github/workflows/test-zkvm-ziren.yml

@@ -2,21 +2,17 @@ name: Test and clippy Ziren
 
 concurrency:
   group: ${{ github.workflow }}-${{ github.ref }}
-  cancel-in-progress: true
+  cancel-in-progress: ${{ github.event_name == 'pull_request' }}
 
 on:
-  push:
-    branches:
-      - master
-    tags:
-      - 'v*'
+  workflow_run:
+    workflows: ['Build and push images']
+    types: [completed]
   pull_request:
 
 jobs:
   test:
+    if: github.event_name == 'pull_request' || github.event.workflow_run.conclusion == 'success'
     uses: ./.github/workflows/test-zkvm.yml
-    permissions:
-      contents: read
-      packages: write
     with:
       zkvm: ziren

.github/workflows/test-zkvm-zisk.yml

@@ -2,22 +2,21 @@ name: Test and clippy ZisK
 
 concurrency:
   group: ${{ github.workflow }}-${{ github.ref }}
-  cancel-in-progress: true
+  cancel-in-progress: ${{ github.event_name == 'pull_request' }}
 
 on:
-  push:
-    branches:
-      - master
-    tags:
-      - 'v*'
+  workflow_run:
+    workflows: ['Build and push images']
+    types: [completed]
   pull_request:
 
 jobs:
   test:
+    if: github.event_name == 'pull_request' || github.event.workflow_run.conclusion == 'success'
     uses: ./.github/workflows/test-zkvm.yml
-    permissions:
-      contents: read
-      packages: write
     with:
       zkvm: zisk
+      cuda: true
+      cuda_archs: '120'
+      cluster: true
       skip_prove_test: true

.github/workflows/test-zkvm.yml

@@ -7,6 +7,21 @@ on:
         description: 'zkVM to test'
         required: true
         type: string
+      cuda:
+        description: 'Whether to build CUDA-enabled images'
+        required: false
+        type: boolean
+        default: false
+      cuda_archs:
+        description: 'Comma-separated CUDA archs to gencode'
+        required: false
+        type: string
+        default: '89,120'
+      cluster:
+        description: 'Whether to build cluster image'
+        required: false
+        type: boolean
+        default: false
       # Remove when we use larger runners, currently only needed to skip for zisk
       skip_prove_test:
         description: 'Whether to skip prove test and ere-dockerized test or not'
@@ -18,73 +33,20 @@ env:
   CARGO_TERM_COLOR: always
 
 jobs:
-  create_semver_image_tag:
-    name: Tag image with SemVer
-    if: startsWith(github.ref, 'refs/tags/')
+  image_meta:
+    name: Get image metadata
     runs-on: ubuntu-latest
-    permissions:
-      contents: read
-      packages: write
-    steps:
-      - name: Log in to GitHub Container Registry
-        uses: docker/login-action@v3
-        with:
-          registry: ghcr.io
-          username: ${{ github.actor }}
-          password: ${{ secrets.GITHUB_TOKEN }}
-
-      - name: Add SemVer tag to images
-        run: |
-          GIT_SHA="${{ github.sha }}"
-          GIT_SHA_TAG="${GIT_SHA:0:7}"
-
-          GIT_TAG="${{ github.ref_name }}"
-          SEMVER_TAG="${GIT_TAG#v}"
-
-          IMAGE_REGISTRY="ghcr.io/${{ github.repository }}"
-
-          for IMAGE in \
-            "ere-base" \
-            "ere-base-${{ inputs.zkvm }}" \
-            "ere-compiler-${{ inputs.zkvm }}" \
-            "ere-server-${{ inputs.zkvm }}"
-          do
-            echo "Tagging $IMAGE_REGISTRY/$IMAGE:$GIT_SHA_TAG as $SEMVER_TAG"
-            docker buildx imagetools create \
-              "$IMAGE_REGISTRY/$IMAGE:$GIT_SHA_TAG" \
-              --tag "$IMAGE_REGISTRY/$IMAGE:$SEMVER_TAG"
-          done
-
-  build_image:
-    name: Build image
-    if: github.event_name == 'pull_request' || (github.event_name == 'push' && github.ref == 'refs/heads/master')
-    runs-on: ubuntu-latest
-    permissions:
-      contents: read
-      packages: write
     outputs:
+      dockerfile_changed: ${{ steps.changed_files.outputs.any_changed }}
       image_registry: ${{ steps.image_meta.outputs.image_registry }}
       image_tag: ${{ steps.image_meta.outputs.image_tag }}
       cached_image_tag: ${{ steps.image_meta.outputs.cached_image_tag }}
-      base_image: ${{ steps.image_meta.outputs.base_image }}
       base_zkvm_image: ${{ steps.image_meta.outputs.base_zkvm_image }}
-      compiler_zkvm_image: ${{ steps.image_meta.outputs.compiler_zkvm_image }}
-      server_zkvm_image: ${{ steps.image_meta.outputs.server_zkvm_image }}
     steps:
       - name: Checkout repository
         uses: actions/checkout@v4
-
-      - name: Free up disk space
-        if: github.event_name == 'push'
-        run: bash .github/scripts/free-up-disk-space.sh
-
-      - name: Log in to GitHub Container Registry
-        if: github.event_name == 'push'
-        uses: docker/login-action@v3
         with:
-          registry: ghcr.io
-          username: ${{ github.actor }}
-          password: ${{ secrets.GITHUB_TOKEN }}
+          ref: ${{ github.event.workflow_run.head_sha || github.sha }}
 
       - name: Check Dockerfile changes
         id: changed_files
@@ -97,7 +59,11 @@ jobs:
       - name: Get image metadata
         id: image_meta
         run: |
-          GIT_SHA="${{ github.sha }}"
+          if [ "${{ github.event_name }}" == "workflow_run" ]; then
+            GIT_SHA="${{ github.event.workflow_run.head_sha }}"
+          else
+            GIT_SHA="${{ github.sha }}"
+          fi
           IMAGE_TAG="${GIT_SHA:0:7}"
 
           CACHED_IMAGE_TAG=""
@@ -107,53 +73,17 @@ jobs:
           fi
 
           IMAGE_REGISTRY="ghcr.io/${{ github.repository }}"
-          BASE_IMAGE="$IMAGE_REGISTRY/ere-base:$IMAGE_TAG"
           BASE_ZKVM_IMAGE="$IMAGE_REGISTRY/ere-base-${{ inputs.zkvm }}:$IMAGE_TAG"
-          COMPILER_ZKVM_IMAGE="$IMAGE_REGISTRY/ere-compiler-${{ inputs.zkvm }}:$IMAGE_TAG"
-          SERVER_ZKVM_IMAGE="$IMAGE_REGISTRY/ere-server-${{ inputs.zkvm }}:$IMAGE_TAG"
 
           echo "image_registry=$IMAGE_REGISTRY" >> $GITHUB_OUTPUT
           echo "image_tag=$IMAGE_TAG" >> $GITHUB_OUTPUT
           echo "cached_image_tag=$CACHED_IMAGE_TAG" >> $GITHUB_OUTPUT
-          echo "base_image=$BASE_IMAGE" >> $GITHUB_OUTPUT
           echo "base_zkvm_image=$BASE_ZKVM_IMAGE" >> $GITHUB_OUTPUT
-          echo "compiler_zkvm_image=$COMPILER_ZKVM_IMAGE" >> $GITHUB_OUTPUT
-          echo "server_zkvm_image=$SERVER_ZKVM_IMAGE" >> $GITHUB_OUTPUT
 
-      - name: Build and push ere-base and ere-base-${{ inputs.zkvm }} images
-        if: github.event_name == 'push'
-        run: |
-          bash .github/scripts/build-image.sh \
-            --zkvm ${{ inputs.zkvm }} \
-            --registry ${{ steps.image_meta.outputs.image_registry }} \
-            --tag ${{ steps.image_meta.outputs.image_tag }} \
-            --base
-          docker push ${{ steps.image_meta.outputs.base_image }}
-          docker push ${{ steps.image_meta.outputs.base_zkvm_image }}
-
-      - name: Build and push ere-compiler-${{ inputs.zkvm }} image
-        if: github.event_name == 'push'
-        run: |
-          bash .github/scripts/build-image.sh \
-            --zkvm ${{ inputs.zkvm }} \
-            --registry ${{ steps.image_meta.outputs.image_registry }} \
-            --tag ${{ steps.image_meta.outputs.image_tag }} \
-            --compiler
-          docker push ${{ steps.image_meta.outputs.compiler_zkvm_image }}
-
-      - name: Build and push ere-server-${{ inputs.zkvm }} image
-        if: github.event_name == 'push'
-        run: |
-          bash .github/scripts/build-image.sh \
-            --zkvm ${{ inputs.zkvm }} \
-            --registry ${{ steps.image_meta.outputs.image_registry }} \
-            --tag ${{ steps.image_meta.outputs.image_tag }} \
-            --server
-          docker push ${{ steps.image_meta.outputs.server_zkvm_image }}
-
-  clippy_via_docker:
-    name: Clippy via Docker
-    needs: build_image
+  build_image_cuda_check:
+    name: Build image with CUDA enabled
+    if: inputs.cuda && github.event_name == 'pull_request'
+    needs: image_meta
     runs-on: ubuntu-latest
     steps:
       - name: Checkout repository
@@ -162,23 +92,69 @@ jobs:
       - name: Free up disk space
         run: bash .github/scripts/free-up-disk-space.sh
 
+      - name: Pull or build ere-base and ere-base-${{ inputs.zkvm }} images with CUDA enabled
+        run: |
+          CACHED_TAG="${{ needs.image_meta.outputs.cached_image_tag }}"
+          if [ -n "$CACHED_TAG" ]; then
+            CACHED_TAG="${CACHED_TAG}-cuda"
+          fi
+          bash .github/scripts/pull-or-build-base-zkvm-image.sh \
+            --zkvm ${{ inputs.zkvm }} \
+            --registry ${{ needs.image_meta.outputs.image_registry }} \
+            --tag ${{ needs.image_meta.outputs.image_tag }}-cuda \
+            --cached-tag "$CACHED_TAG" \
+            --cuda-archs '${{ inputs.cuda_archs }}'
+
+      - name: Build ere-server-${{ inputs.zkvm }} image with CUDA enabled
+        run: |
+          bash .github/scripts/build-image.sh \
+            --zkvm ${{ inputs.zkvm }} \
+            --registry ${{ needs.image_meta.outputs.image_registry }} \
+            --tag ${{ needs.image_meta.outputs.image_tag }}-cuda \
+            --server \
+            --cuda-archs '${{ inputs.cuda_archs }}'
+
+      - name: Build ere-cluster-${{ inputs.zkvm }} image with CUDA enabled
+        if: ${{ inputs.cluster && needs.image_meta.outputs.dockerfile_changed == 'true' }}
+        run: |
+          bash .github/scripts/build-image.sh \
+            --zkvm ${{ inputs.zkvm }} \
+            --registry ${{ needs.image_meta.outputs.image_registry }} \
+            --tag ${{ needs.image_meta.outputs.image_tag }}-cuda \
+            --cluster \
+            --cuda-archs '${{ inputs.cuda_archs }}'
+
+  clippy_via_docker:
+    name: Clippy via Docker
+    needs: image_meta
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v4
+        with:
+          ref: ${{ github.event.workflow_run.head_sha || github.sha }}
+
+      - name: Free up disk space
+        run: bash .github/scripts/free-up-disk-space.sh
+
       - name: Cache dependencies
         uses: Swatinem/rust-cache@v2
         with:
           key: rust-${{ inputs.zkvm }}-${{ hashFiles('Cargo.lock') }}
 
+      - name: Pull published base zkvm image
+        if: github.event_name == 'workflow_run'
+        run: |
+          docker pull ${{ needs.image_meta.outputs.base_zkvm_image }}
+
       - name: Pull base zkvm image or build locally
         if: github.event_name == 'pull_request'
         run: |
           bash .github/scripts/pull-or-build-base-zkvm-image.sh \
             --zkvm ${{ inputs.zkvm }} \
-            --registry ${{ needs.build_image.outputs.image_registry }} \
-            --tag ${{ needs.build_image.outputs.image_tag }} \
-            --cached-tag "${{ needs.build_image.outputs.cached_image_tag }}"
-
-      - name: Pull base zkvm image
-        if: github.event_name == 'push'
-        run: docker pull ${{ needs.build_image.outputs.base_zkvm_image }}
+            --registry ${{ needs.image_meta.outputs.image_registry }} \
+            --tag ${{ needs.image_meta.outputs.image_tag }} \
+            --cached-tag "${{ needs.image_meta.outputs.cached_image_tag }}"
 
       - name: Run cargo clippy for ere-${{ inputs.zkvm }} via Docker
         run: |
@@ -189,7 +165,7 @@ jobs:
             --volume $HOME/.cargo/registry:/usr/local/cargo/registry \
             --volume $HOME/.cargo/git:/usr/local/cargo/git \
             --workdir /ere \
-            ${{ needs.build_image.outputs.base_zkvm_image }} \
+            ${{ needs.image_meta.outputs.base_zkvm_image }} \
             /bin/bash"
 
           cat <<EOF | $DOCKER_CMD
@@ -207,11 +183,13 @@ jobs:
 
   test_via_docker:
     name: Test via Docker
-    needs: build_image
+    needs: image_meta
     runs-on: ubuntu-latest
     steps:
       - name: Checkout repository
         uses: actions/checkout@v4
+        with:
+          ref: ${{ github.event.workflow_run.head_sha || github.sha }}
 
       - name: Free up disk space
         run: bash .github/scripts/free-up-disk-space.sh
@@ -221,18 +199,19 @@ jobs:
         with:
           key: rust-${{ inputs.zkvm }}-${{ hashFiles('Cargo.lock') }}
 
+      - name: Pull published base zkvm image
+        if: github.event_name == 'workflow_run'
+        run: |
+          docker pull ${{ needs.image_meta.outputs.base_zkvm_image }}
+
       - name: Pull base zkvm image or build locally
         if: github.event_name == 'pull_request'
         run: |
           bash .github/scripts/pull-or-build-base-zkvm-image.sh \
             --zkvm ${{ inputs.zkvm }} \
-            --registry ${{ needs.build_image.outputs.image_registry }} \
-            --tag ${{ needs.build_image.outputs.image_tag }} \
-            --cached-tag "${{ needs.build_image.outputs.cached_image_tag }}"
-
-      - name: Pull base zkvm image
-        if: github.event_name == 'push'
-        run: docker pull ${{ needs.build_image.outputs.base_zkvm_image }}
+            --registry ${{ needs.image_meta.outputs.image_registry }} \
+            --tag ${{ needs.image_meta.outputs.image_tag }} \
+            --cached-tag "${{ needs.image_meta.outputs.cached_image_tag }}"
 
       - name: Run cargo test for ere-${{ inputs.zkvm }} via Docker
         run: |
@@ -243,7 +222,7 @@ jobs:
             --volume $HOME/.cargo/registry:/usr/local/cargo/registry \
             --volume $HOME/.cargo/git:/usr/local/cargo/git \
             --workdir /ere \
-            ${{ needs.build_image.outputs.base_zkvm_image }} \
+            ${{ needs.image_meta.outputs.base_zkvm_image }} \
             /bin/bash"
 
           cat <<EOF | $DOCKER_CMD
@@ -259,11 +238,13 @@ jobs:
 
   test_ere_dockerized:
     name: Test ere-dockerized with the selected zkVM
-    needs: build_image
+    needs: image_meta
     runs-on: ubuntu-latest
     steps:
       - name: Checkout repository
         uses: actions/checkout@v4
+        with:
+          ref: ${{ github.event.workflow_run.head_sha || github.sha }}
 
       - name: Free up disk space
         run: bash .github/scripts/free-up-disk-space.sh
@@ -278,32 +259,32 @@ jobs:
         with:
           key: rust-${{ inputs.zkvm }}-${{ hashFiles('Cargo.lock') }}
 
+      - name: Pull published images
+        if: github.event_name == 'workflow_run'
+        run: |
+          docker pull ${{ needs.image_meta.outputs.image_registry }}/ere-compiler-${{ inputs.zkvm }}:${{ needs.image_meta.outputs.image_tag }}
+          docker pull ${{ needs.image_meta.outputs.image_registry }}/ere-server-${{ inputs.zkvm }}:${{ needs.image_meta.outputs.image_tag }}
+
       - name: Pull images or build locally
         if: github.event_name == 'pull_request'
         run: |
           bash .github/scripts/pull-or-build-base-zkvm-image.sh \
             --zkvm ${{ inputs.zkvm }} \
-            --registry ${{ needs.build_image.outputs.image_registry }} \
-            --tag ${{ needs.build_image.outputs.image_tag }} \
-            --cached-tag "${{ needs.build_image.outputs.cached_image_tag }}"
+            --registry ${{ needs.image_meta.outputs.image_registry }} \
+            --tag ${{ needs.image_meta.outputs.image_tag }} \
+            --cached-tag "${{ needs.image_meta.outputs.cached_image_tag }}"
 
           # Build ere-compiler-${{ inputs.zkvm }} and ere-server-${{ inputs.zkvm }}
           bash .github/scripts/build-image.sh \
             --zkvm ${{ inputs.zkvm }} \
-            --registry ${{ needs.build_image.outputs.image_registry }} \
-            --tag ${{ needs.build_image.outputs.image_tag }} \
+            --registry ${{ needs.image_meta.outputs.image_registry }} \
+            --tag ${{ needs.image_meta.outputs.image_tag }} \
             --compiler \
             --server
 
-      - name: Pull ere-compiler-${{ inputs.zkvm }} and ere-server-${{ inputs.zkvm }} images
-        if: github.event_name == 'push'
-        run: |
-          docker pull ${{ needs.build_image.outputs.compiler_zkvm_image }}
-          docker pull ${{ needs.build_image.outputs.server_zkvm_image }}
-
       - name: Run cargo test for ere-${{ inputs.zkvm }} via ere-dockerized
         env:
-          ERE_IMAGE_REGISTRY: ${{ needs.build_image.outputs.image_registry }}
+          ERE_IMAGE_REGISTRY: ${{ needs.image_meta.outputs.image_registry }}
         run: |
           cargo test --release --package ere-dockerized \
             -- ${{ inputs.zkvm }} ${{ inputs.skip_prove_test && '--skip prove' || ''  }} --test-threads=1