Accelerate CI image materialization with same-run snapshots

.github/actions/load_image/action.yml

@@ -0,0 +1,116 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+---
+name: 'Load image (CI or PROD)'
+description: >
+  Loads the current image for one Python version: from the CI image snapshot when it fits this
+  runner's daemon, otherwise from the image stash.
+inputs:
+  python:
+    description: 'Python version of the image to load'
+    required: true
+  platform:
+    description: 'Platform of the image - linux/amd64 or linux/arm64'
+    required: true
+  image-type:
+    description: 'Which image type to load (ci/prod)'
+    default: "ci"
+  image-stash-ref:
+    description: 'Ref whose image stash to restore (empty = the per-branch one)'
+    required: false
+    default: ""
+runs:
+  using: "composite"
+  steps:
+    # The metadata is a few bytes, so a snapshot that cannot fit is rejected before its download.
+    # A pattern, unlike a name, downloads nothing without an error when the builder made no snapshot.
+    # Self-hosted runners restore the stash from their own storage, which a run artifact cannot beat.
+    - name: "Download CI image snapshot metadata ${{ inputs.platform }}:${{ inputs.python }}"
+      id: snapshot-meta
+      continue-on-error: true
+      uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c  # v8.0.1
+      with:
+        pattern: >-
+          ci-image-snapshot-v1-${{ inputs.python }}-${{
+          inputs.platform == 'linux/amd64' && 'amd64' || 'arm64' }}-meta
+        merge-multiple: true
+        path: "/mnt/"
+      if: >
+        github.event_name == 'pull_request' && runner.environment == 'github-hosted' &&
+        inputs.image-type == 'ci' && inputs.image-stash-ref == ''
+    - name: "Check CI image snapshot ${{ inputs.platform }}:${{ inputs.python }}"
+      id: snapshot-check
+      env:
+        PLATFORM: ${{ inputs.platform }}
+        PYTHON: ${{ inputs.python }}
+      run: |
+        if ./scripts/ci/docker_data_root_snapshot.sh check \
+            "/mnt/ci-image-snapshot-${PLATFORM//\//_}-${PYTHON}.tar.zst"; then
+          echo "fits=true" >> "${GITHUB_OUTPUT}"
+        else
+          echo "Falling back to loading the image stash"
+        fi
+      shell: bash
+      if: steps.snapshot-meta.outcome == 'success'
+    - name: "Download CI image snapshot ${{ inputs.platform }}:${{ inputs.python }}"
+      id: snapshot-download
+      continue-on-error: true
+      uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c  # v8.0.1
+      with:
+        name: >-
+          ci-image-snapshot-v1-${{ inputs.python }}-${{
+          inputs.platform == 'linux/amd64' && 'amd64' || 'arm64' }}
+        path: "/mnt/"
+      if: steps.snapshot-check.outputs.fits == 'true'
+    # Falls back to the image stash below whenever the snapshot cannot be restored.
+    - name: "Unpack CI image snapshot ${{ inputs.platform }}:${{ inputs.python }}"
+      id: snapshot
+      env:
+        PLATFORM: ${{ inputs.platform }}
+        PYTHON: ${{ inputs.python }}
+      run: |
+        if breeze ci-image load --platform "${PLATFORM}" --python "${PYTHON}" \
+            --snapshot-file "/mnt/ci-image-snapshot-${PLATFORM//\//_}-${PYTHON}.tar.zst"; then
+          echo "restored=true" >> "${GITHUB_OUTPUT}"
+        else
+          echo "Falling back to loading the image stash"
+        fi
+      shell: bash
+      if: steps.snapshot-download.outcome == 'success'
+    - name: >
+        Restore ${{ inputs.image-type }} docker image ${{ inputs.platform }}:${{ inputs.python }}
+        ${{ inputs.image-stash-ref != '' && format('built for ref {0}', inputs.image-stash-ref) || '' }}
+      uses: ./.github/actions/stash-restore
+      with:
+        key: "${{ inputs.image-type }}-image-save-v3-${{ inputs.platform }}-${{ inputs.python }}\
+          ${{ inputs.image-stash-ref != '' && format('-{0}', inputs.image-stash-ref) || '' }}"
+        path: "/mnt/"
+        # Must stay true: runs on RunsOn (canaries and release-branch pushes) save these image
+        # stashes as S3 digests only.
+        only-current-branch: 'true'
+        fail-on-download: 'true'
+      if: steps.snapshot.outputs.restored != 'true'
+    - name: "Load ${{ inputs.image-type }} image ${{ inputs.platform }}:${{ inputs.python }}"
+      env:
+        PLATFORM: ${{ inputs.platform }}
+        PYTHON: ${{ inputs.python }}
+        IMAGE_TYPE: ${{ inputs.image-type }}
+      run: >
+        breeze ${IMAGE_TYPE}-image load --platform "${PLATFORM}" --python "${PYTHON}" --image-file-dir "/mnt"
+      shell: bash
+      if: steps.snapshot.outputs.restored != 'true'

.github/actions/prepare_breeze_and_image/action.yml

@@ -79,23 +79,10 @@ runs:
       run: |
         echo "Checking free space!"
         df -H
-    - name: >
-        Restore ${{ inputs.image-type }} docker image ${{ inputs.platform }}:${{ inputs.python }}
-        ${{ inputs.image-stash-ref != '' && format('built for ref {0}', inputs.image-stash-ref) || '' }}
-      uses: ./.github/actions/stash-restore
-      with:
-        key: "${{ inputs.image-type }}-image-save-v3-${{ inputs.platform }}-${{ inputs.python }}\
-          ${{ inputs.image-stash-ref != '' && format('-{0}', inputs.image-stash-ref) || '' }}"
-        path: "/mnt/"
-        # Must stay true: runs on RunsOn (canaries and release-branch pushes) save these image
-        # stashes as S3 digests only.
-        only-current-branch: 'true'
-        fail-on-download: 'true'
     - name: "Load ${{ inputs.image-type }} image ${{ inputs.platform }}:${{ inputs.python }}"
-      env:
-        PLATFORM: ${{ inputs.platform }}
-        PYTHON: ${{ inputs.python }}
-        IMAGE_TYPE: ${{ inputs.image-type }}
-      run: >
-        breeze ${IMAGE_TYPE}-image load --platform "${PLATFORM}" --python "${PYTHON}" --image-file-dir "/mnt"
-      shell: bash
+      uses: ./.github/actions/load_image
+      with:
+        python: ${{ inputs.python }}
+        platform: ${{ inputs.platform }}
+        image-type: ${{ inputs.image-type }}
+        image-stash-ref: ${{ inputs.image-stash-ref }}

.github/actions/prepare_single_ci_image/action.yml

@@ -35,20 +35,9 @@ inputs:
 runs:
   using: "composite"
   steps:
-    - name: "Restore CI docker images ${{ inputs.platform }}:${{ inputs.python }}"
-      uses: ./.github/actions/stash-restore
-      with:
-        key: ci-image-save-v3-${{ inputs.platform }}-${{ inputs.python }}
-        path: "/mnt/"
-        # Must stay true: runs on RunsOn (canaries and release-branch pushes) save these image
-        # stashes as S3 digests only.
-        only-current-branch: 'true'
-        fail-on-download: 'true'
-      if: contains(inputs.python-versions-list-as-string, inputs.python)
     - name: "Load CI image ${{ inputs.platform }}:${{ inputs.python }}"
-      env:
-        PLATFORM: ${{ inputs.platform }}
-        PYTHON: ${{ inputs.python }}
-      run: breeze ci-image load --platform "${PLATFORM}" --python "${PYTHON}" --image-file-dir "/mnt/"
-      shell: bash
+      uses: ./.github/actions/load_image
+      with:
+        python: ${{ inputs.python }}
+        platform: ${{ inputs.platform }}
       if: contains(inputs.python-versions-list-as-string, inputs.python)

.github/workflows/ci-image-build.yml

@@ -412,3 +412,53 @@ jobs:
           steps.stashed-image.outputs.reusable != 'true'
       - name: "Check disk space after build"
         run: df -H
+
+      # Snapshot only after cache publication: pruning removes the mount-cache image.
+      - name: "Snapshot CI image ${{ inputs.platform }}:${{ env.PYTHON_MAJOR_MINOR_VERSION }}"
+        id: snapshot-export
+        continue-on-error: true
+        timeout-minutes: 5
+        env:
+          PLATFORM: ${{ inputs.platform }}
+          # Docker is bind-mounted from /mnt on AMD; write the archive on the root disk.
+          SNAPSHOT_FILE: >-
+            ${{ runner.temp }}/ci-image-snapshot-${{
+            inputs.platform == 'linux/amd64' && 'linux_amd64' || 'linux_arm64'
+            }}-${{ env.PYTHON_MAJOR_MINOR_VERSION }}.tar.zst
+        run: >
+          breeze ci-image save --platform "${PLATFORM}" --python "${PYTHON_MAJOR_MINOR_VERSION}"
+          --snapshot-file "${SNAPSHOT_FILE}"
+        shell: bash
+        if: >
+          github.event_name == 'pull_request' && runner.environment == 'github-hosted' &&
+          inputs.upload-image-artifact == 'true' && inputs.image-stash-ref == '' &&
+          (inputs.checkout-ref == '' || inputs.checkout-ref == github.sha)
+      - name: "Upload CI image snapshot ${{ inputs.platform }}:${{ env.PYTHON_MAJOR_MINOR_VERSION }}"
+        continue-on-error: true
+        uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a  # v7.0.1
+        with:
+          name: >-
+            ci-image-snapshot-v1-${{ env.PYTHON_MAJOR_MINOR_VERSION }}-${{
+            inputs.platform == 'linux/amd64' && 'amd64' || 'arm64' }}
+          path: "${{ runner.temp }}/ci-image-snapshot-*-${{ env.PYTHON_MAJOR_MINOR_VERSION }}.tar.zst"
+          if-no-files-found: 'error'
+          compression-level: '0'
+          # Read only within this run; once it expires, consumers fall back to the stash.
+          retention-days: '1'
+          # A rerun replaces the snapshot of its earlier attempt instead of failing to upload.
+          overwrite: true
+        if: steps.snapshot-export.outcome == 'success'
+      # Uploaded even when the export fails, so its "unavailable" marker retires the snapshot of an
+      # earlier attempt; consumers read it before downloading the snapshot.
+      - name: "Upload CI image snapshot metadata ${{ inputs.platform }}:${{ env.PYTHON_MAJOR_MINOR_VERSION }}"
+        continue-on-error: true
+        uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a  # v7.0.1
+        with:
+          name: >-
+            ci-image-snapshot-v1-${{ env.PYTHON_MAJOR_MINOR_VERSION }}-${{
+            inputs.platform == 'linux/amd64' && 'amd64' || 'arm64' }}-meta
+          path: "${{ runner.temp }}/ci-image-snapshot-*-${{ env.PYTHON_MAJOR_MINOR_VERSION }}.tar.zst.meta"
+          if-no-files-found: 'error'
+          retention-days: '1'
+          overwrite: true
+        if: steps.snapshot-export.outcome != 'skipped'