Nvidia-runners S3 cache implementation (#2053)

Steboss · web-flow · commit 82f7d107e755 · 2026-05-20T16:15:28.000+01:00
diff --git a/.github/actions/build-container/action.yml b/.github/actions/build-container/action.yml
@@ -54,6 +54,18 @@ inputs:
     description: "URL of the Bazel remote cache to use for building the image"
     required: true
     default: ""
+  docker-build-allow:
+    description: "Additional BuildKit entitlements to allow during docker build"
+    required: false
+    default: ""
+  docker-build-network:
+    description: "Networking mode for RUN instructions during docker build"
+    required: false
+    default: ""
+  registry-build-context-name:
+    description: "Optional named build context to source from a runner registry for the final image handoff"
+    required: false
+    default: ""
 
 outputs:
   DOCKER_TAG_MEALKIT:
@@ -72,6 +84,7 @@ runs:
       run: |
         echo 'UPLD_IMAGE=ghcr.io/nvidia/jax-toolbox-internal' >> $GITHUB_ENV
         echo "BADGE_FILENAME_FULL=${{ inputs.BADGE_FILENAME }}-${{ inputs.ARCHITECTURE }}.json" >> $GITHUB_ENV
+        echo "BUILDX_BUILDER_NAME=jax-toolbox-${{ inputs.CONTAINER_NAME }}-${{ inputs.ARCHITECTURE }}-${{ inputs.docker-build-network }}" >> $GITHUB_ENV
 
     - name: Setup SSH
       id: setup-ssh
@@ -90,9 +103,14 @@ runs:
     - name: Set up Docker Buildx
       uses: docker/setup-buildx-action@v4
       with:
+        name: ${{ env.BUILDX_BUILDER_NAME }}
         driver-opts: |
-          image=moby/buildkit:v0.12.1
-        version: v0.30.1
+          image=moby/buildkit:v0.13.2
+          network=${{ inputs.docker-build-network }}
+        keep-state: true
+        # TODO: check whether we need the oci.worker specification
+
+
 
     - name: Download nsys-jax version.py
       uses: actions/download-artifact@v8
@@ -124,7 +142,9 @@ runs:
       id: mealkit-build
       uses: docker/build-push-action@v7
       with:
+        builder: ${{ env.BUILDX_BUILDER_NAME }}
         context: ${{ inputs.DOCKER_CONTEXT }}
+        load: ${{inputs.registry-build-context-name == ''}}
         push: true
         file: ${{ inputs.DOCKERFILE }}
         platforms: linux/${{ inputs.ARCHITECTURE }}
@@ -134,6 +154,8 @@ runs:
         ssh: default
         secret-files: |
           "SSH_KNOWN_HOSTS=${{ steps.setup-ssh.outputs.known-hosts-file }}"
+        allow: ${{ inputs.docker-build-allow }}
+        network: ${{ inputs.docker-build-network }}
         build-args: |
           BASE_IMAGE=${{ inputs.BASE_IMAGE }}
           BAZEL_CACHE=${{ inputs.bazel-remote-cache-url }}
@@ -157,7 +179,9 @@ runs:
       id: final-build
       uses: docker/build-push-action@v7
       with:
+        builder: ${{ env.BUILDX_BUILDER_NAME }}
         context: ${{ inputs.DOCKER_CONTEXT }}
+        build-contexts: ${{ inputs.registry-build-context-name != '' && format('{0}=docker-image://{1}@{2}', inputs.registry-build-context-name, env.UPLD_IMAGE, steps.mealkit-build.outputs.digest) || '' }}
         push: true
         file: ${{ inputs.DOCKERFILE }}
         platforms: linux/${{ inputs.ARCHITECTURE }}
@@ -167,6 +191,8 @@ runs:
         ssh: default
         secret-files: |
           "SSH_KNOWN_HOSTS=${{ steps.setup-ssh.outputs.known-hosts-file }}"
+        allow: ${{ inputs.docker-build-allow }}
+        network: ${{ inputs.docker-build-network }}
         build-args: |
           BASE_IMAGE=${{ inputs.BASE_IMAGE }}
           BAZEL_CACHE=${{ inputs.bazel-remote-cache-url }}
diff --git a/.github/workflows/_ci.yaml b/.github/workflows/_ci.yaml
@@ -43,6 +43,7 @@ concurrency:
 permissions:
   contents: read  # to fetch code
   actions:  write # to cancel previous workflows
+  id-token: write # to assume the AWS role for bazel-remote cache
   packages: write # to upload container
 
 jobs:
@@ -65,10 +66,90 @@ jobs:
 
   build-jax:
     needs: build-base
-    runs-on: [self-hosted, "${{ inputs.ARCHITECTURE }}", "large"]
+    runs-on: ${{ inputs.ARCHITECTURE == 'amd64' && 'linux-amd64-cpu32m' || 'linux-arm64-cpu32m' }}
     steps:
         - name: Checkout repository
           uses: actions/checkout@v6
+        - name: Configure AWS credentials via OIDC
+          uses: aws-actions/configure-aws-credentials@v4
+          with:
+            role-to-assume: ${{ secrets.S3_CACHE_ROLE }}
+            aws-region: ${{ secrets.AWS_REGION}}
+        - name: Start bazel-remote sidecar
+          env:
+            BAZEL_REMOTE_BUCKET: ${{ secrets.BAZEL_REMOTE_BUCKET }}
+            BAZEL_REMOTE_AWS_DIR: /tmp/bazel-remote-aws
+            BAZEL_REMOTE_CONTAINER: bazel-remote-cache-${{ inputs.ARCHITECTURE }}
+            BAZEL_REMOTE_DISK_DIR: /tmp/bazel-remote-cache
+            BAZEL_REMOTE_DIAGNOSTICS_DIR: /tmp/bazel-remote-diagnostics
+            BAZEL_REMOTE_PORT: 9090
+            BAZEL_REMOTE_PREFIX: bazel-remote/${{ inputs.ARCHITECTURE }}
+            BAZEL_REMOTE_REGION: ${{ secrets.AWS_REGION }}
+          run: |
+            set -euxo pipefail
+            : "${AWS_ACCESS_KEY_ID:?missing AWS_ACCESS_KEY_ID}"
+            : "${AWS_SECRET_ACCESS_KEY:?missing AWS_SECRET_ACCESS_KEY}"
+            : "${AWS_SESSION_TOKEN:?missing AWS_SESSION_TOKEN}"
+            BAZEL_REMOTE_AWS_CREDENTIALS="${BAZEL_REMOTE_AWS_DIR}/credentials"
+            mkdir -p "${BAZEL_REMOTE_DIAGNOSTICS_DIR}" "${BAZEL_REMOTE_DISK_DIR}" "${BAZEL_REMOTE_AWS_DIR}"
+            BAZEL_REMOTE_UID="$(id -u)"
+            BAZEL_REMOTE_GID="$(id -g)"
+            set +x
+            umask 077
+            {
+              printf '[default]\n'
+              printf 'aws_access_key_id=%s\n' "${AWS_ACCESS_KEY_ID}"
+              printf 'aws_secret_access_key=%s\n' "${AWS_SECRET_ACCESS_KEY}"
+              printf 'aws_session_token=%s\n' "${AWS_SESSION_TOKEN}"
+            } > "${BAZEL_REMOTE_AWS_CREDENTIALS}"
+            set -x
+            docker rm -f "${BAZEL_REMOTE_CONTAINER}" >/dev/null 2>&1 || true
+            docker run -d \
+              --name "${BAZEL_REMOTE_CONTAINER}" \
+              --user "${BAZEL_REMOTE_UID}:${BAZEL_REMOTE_GID}" \
+              --publish "127.0.0.1:${BAZEL_REMOTE_PORT}:8080" \
+              --volume "${BAZEL_REMOTE_DISK_DIR}:/data" \
+              --volume "${BAZEL_REMOTE_AWS_CREDENTIALS}:/aws-config/credentials:ro" \
+              buchgr/bazel-remote-cache \
+              --dir /data \
+              --max_size 50 \
+              --http_address=0.0.0.0:8080 \
+              --s3.auth_method=aws_credentials_file \
+              --s3.aws_profile=default \
+              --s3.aws_shared_credentials_file=/aws-config/credentials \
+              --s3.bucket="${BAZEL_REMOTE_BUCKET}" \
+              --s3.endpoint="s3.${BAZEL_REMOTE_REGION}.amazonaws.com" \
+              --s3.prefix="${BAZEL_REMOTE_PREFIX}" \
+              --s3.region="${BAZEL_REMOTE_REGION}" \
+              --s3.update_timestamps
+        - name: Verify bazel-remote sidecar
+          env:
+            BAZEL_REMOTE_CONTAINER: bazel-remote-cache-${{ inputs.ARCHITECTURE }}
+            BAZEL_REMOTE_DIAGNOSTICS_DIR: /tmp/bazel-remote-diagnostics
+            BAZEL_REMOTE_PORT: 9090
+          run: |
+            set -euxo pipefail
+            ready=0
+            for _ in $(seq 1 30); do
+              if curl -sSf "http://127.0.0.1:${BAZEL_REMOTE_PORT}/status" > "${BAZEL_REMOTE_DIAGNOSTICS_DIR}/status-before.json"; then
+                ready=1
+                break
+              fi
+              container_running="$(docker inspect -f '{{.State.Running}}' "${BAZEL_REMOTE_CONTAINER}" 2>/dev/null || echo false)"
+              if [[ "${container_running}" != "true" ]]; then
+                docker ps -a
+                docker logs "${BAZEL_REMOTE_CONTAINER}" || true
+                exit 1
+              fi
+              sleep 1
+            done
+            if [[ "${ready}" != "1" ]]; then
+              docker ps -a
+              docker logs "${BAZEL_REMOTE_CONTAINER}" || true
+              exit 1
+            fi
+            curl -sSIf "http://127.0.0.1:${BAZEL_REMOTE_PORT}/cas/e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855" \
+              > "${BAZEL_REMOTE_DIAGNOSTICS_DIR}/empty-cas-head.txt"
         - name: Build JAX container
           id: build-jax
           uses: ./.github/actions/build-container
@@ -84,19 +165,44 @@ jobs:
             ssh-private-key: ${{ secrets.SSH_PRIVATE_KEY }}
             ssh-known-hosts: ${{ vars.SSH_KNOWN_HOSTS }}
             github-token: ${{ secrets.GITHUB_TOKEN }}
-            bazel-remote-cache-url: ${{ vars.BAZEL_REMOTE_CACHE_URL }}
+            bazel-remote-cache-url: http://127.0.0.1:9090
+            docker-build-allow: network.host
+            docker-build-network: host
+            registry-build-context-name: mealkit
             EXTRA_BUILD_ARGS: |
               URLREF_JAX=${{ fromJson(inputs.SOURCE_URLREFS).JAX }}
               URLREF_XLA=${{ fromJson(inputs.SOURCE_URLREFS).XLA }}
               URLREF_FLAX=${{ fromJson(inputs.SOURCE_URLREFS).FLAX }}
               URLREF_TRANSFORMER_ENGINE=${{ fromJson(inputs.SOURCE_URLREFS).TRANSFORMER_ENGINE }}
+        - name: Collect bazel-remote diagnostics
+          if: always()
+          env:
+            BAZEL_REMOTE_CONTAINER: bazel-remote-cache-${{ inputs.ARCHITECTURE }}
+            BAZEL_REMOTE_DIAGNOSTICS_DIR: /tmp/bazel-remote-diagnostics
+            BAZEL_REMOTE_PORT: 9090
+          run: |
+            set -euxo pipefail
+            mkdir -p "${BAZEL_REMOTE_DIAGNOSTICS_DIR}"
+            curl -sSf "http://127.0.0.1:${BAZEL_REMOTE_PORT}/status" > "${BAZEL_REMOTE_DIAGNOSTICS_DIR}/status-after.json" || true
+            docker logs "${BAZEL_REMOTE_CONTAINER}" > "${BAZEL_REMOTE_DIAGNOSTICS_DIR}/container.log" 2>&1 || true
+        - name: Upload bazel-remote diagnostics
+          if: always()
+          uses: actions/upload-artifact@v6
+          with:
+            name: bazel-remote-diagnostics-${{ inputs.ARCHITECTURE }}
+            path: /tmp/bazel-remote-diagnostics
+        - name: Stop bazel-remote sidecar
+          if: always()
+          run: |
+            docker rm -f "bazel-remote-cache-${{ inputs.ARCHITECTURE }}" || true
+            rm -rf /tmp/bazel-remote-aws
     outputs:
       DOCKER_TAG_MEALKIT: ${{ steps.build-jax.outputs.DOCKER_TAG_MEALKIT }}
       DOCKER_TAG_FINAL:   ${{ steps.build-jax.outputs.DOCKER_TAG_FINAL }}
 
   build-equinox:
     needs: build-jax
-    runs-on: [self-hosted, "${{ inputs.ARCHITECTURE }}", "small"]
+    runs-on: ${{ inputs.ARCHITECTURE == 'amd64' && 'linux-amd64-cpu16m' || 'linux-arm64-cpu16m' }}
     outputs:
       DOCKER_TAG_MEALKIT: ${{ steps.build-equinox.outputs.DOCKER_TAG_MEALKIT }}
       DOCKER_TAG_FINAL: ${{ steps.build-equinox.outputs.DOCKER_TAG_FINAL }}
@@ -118,13 +224,12 @@ jobs:
           ssh-private-key: ${{ secrets.SSH_PRIVATE_KEY }}
           ssh-known-hosts: ${{ vars.SSH_KNOWN_HOSTS }}
           github-token: ${{ secrets.GITHUB_TOKEN }}
-          bazel-remote-cache-url: ${{ vars.BAZEL_REMOTE_CACHE_URL }}
           EXTRA_BUILD_ARGS: |
             URLREF_EQUINOX=${{ fromJson(inputs.SOURCE_URLREFS).EQUINOX }}
 
   build-maxtext:
     needs: build-jax
-    runs-on: [self-hosted, "${{ inputs.ARCHITECTURE }}", "small"]
+    runs-on: ${{ inputs.ARCHITECTURE == 'amd64' && 'linux-amd64-cpu16m' || 'linux-arm64-cpu16m' }}
     outputs:
       DOCKER_TAG_MEALKIT: ${{ steps.build-maxtext.outputs.DOCKER_TAG_MEALKIT }}
       DOCKER_TAG_FINAL:   ${{ steps.build-maxtext.outputs.DOCKER_TAG_FINAL }}
@@ -146,13 +251,12 @@ jobs:
           ssh-private-key: ${{ secrets.SSH_PRIVATE_KEY }}
           ssh-known-hosts: ${{ vars.SSH_KNOWN_HOSTS }}
           github-token: ${{ secrets.GITHUB_TOKEN }}
-          bazel-remote-cache-url: ${{ vars.BAZEL_REMOTE_CACHE_URL }}
           EXTRA_BUILD_ARGS: |
             URLREF_MAXTEXT=${{ fromJson(inputs.SOURCE_URLREFS).MAXTEXT }}
 
   build-torchax:
     needs: build-jax
-    runs-on: [self-hosted, "${{ inputs.ARCHITECTURE }}", "small"]
+    runs-on: ${{ inputs.ARCHITECTURE == 'amd64' && 'linux-amd64-cpu16m' || 'linux-arm64-cpu16m' }}
     outputs:
       DOCKER_TAG_MEALKIT: ${{ steps.build-torchax.outputs.DOCKER_TAG_MEALKIT }}
       DOCKER_TAG_FINAL:   ${{ steps.build-torchax.outputs.DOCKER_TAG_FINAL }}
@@ -174,13 +278,12 @@ jobs:
           ssh-private-key: ${{ secrets.SSH_PRIVATE_KEY }}
           ssh-known-hosts: ${{ vars.SSH_KNOWN_HOSTS }}
           github-token: ${{ secrets.GITHUB_TOKEN }}
-          bazel-remote-cache-url: ${{ vars.BAZEL_REMOTE_CACHE_URL }}
           EXTRA_BUILD_ARGS: |
             URLREF_TORCHAX=${{ fromJson(inputs.SOURCE_URLREFS).TORCHAX }}
 
   build-axlearn:
     needs: build-jax
-    runs-on: [self-hosted, "${{ inputs.ARCHITECTURE }}", "large"]
+    runs-on: ${{ inputs.ARCHITECTURE == 'amd64' && 'linux-amd64-cpu16m' || 'linux-arm64-cpu16m' }}
     outputs:
       DOCKER_TAG_MEALKIT: ${{ steps.build-axlearn.outputs.DOCKER_TAG_MEALKIT }}
       DOCKER_TAG_FINAL:   ${{ steps.build-axlearn.outputs.DOCKER_TAG_FINAL }}
@@ -202,7 +305,6 @@ jobs:
           ssh-private-key: ${{ secrets.SSH_PRIVATE_KEY }}
           ssh-known-hosts: ${{ vars.SSH_KNOWN_HOSTS }}
           github-token: ${{ secrets.GITHUB_TOKEN }}
-          bazel-remote-cache-url: ${{ vars.BAZEL_REMOTE_CACHE_URL }}
           EXTRA_BUILD_ARGS: |
             URLREF_AXLEARN=${{ fromJson(inputs.SOURCE_URLREFS).AXLEARN }}
 
diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml
@@ -66,6 +66,7 @@ concurrency:
 permissions:
   contents: write       # to fetch code and push branch
   actions:  write       # to cancel previous workflows
+  id-token: write       # to assume the AWS role for bazel-remote cache
   packages: write       # to upload container
   pull-requests: write  # to make pull request for manifest bump
 
