Preflight CI image snapshot metadata before download

.github/actions/prepare_breeze_and_image/action.yml

@@ -49,6 +49,10 @@ inputs:
       other than the branch tip is stashed under a key of its own by ci-image-build.yml.
     required: false
     default: ""
+  restore-snapshot:
+    description: 'Whether to attempt the optional snapshot from this workflow run'
+    required: false
+    default: 'true'
 outputs:
   host-python-version:
     description: Python version used in host
@@ -79,6 +83,54 @@ runs:
       run: |
         echo "Checking free space!"
         df -H
+    - name: "Preflight CI image snapshot metadata ${{ inputs.platform }}:${{ inputs.python }}"
+      continue-on-error: true
+      uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c  # v8.0.1
+      with:
+        name: >-
+          ci-image-snapshot-meta-v2-${{ inputs.python }}-${{
+          inputs.platform == 'linux/amd64' && 'amd64' || 'arm64' }}
+        path: "/mnt/"
+      id: snapshot-metadata
+      if: >
+        github.event_name == 'pull_request' && inputs.restore-snapshot == 'true' &&
+        inputs.image-type == 'ci' && inputs.image-stash-ref == ''
+    - name: "Check CI image snapshot compatibility ${{ inputs.platform }}:${{ inputs.python }}"
+      id: snapshot-preflight
+      continue-on-error: true
+      env:
+        PLATFORM: ${{ inputs.platform }}
+        PYTHON: ${{ inputs.python }}
+      run: |
+        ./scripts/ci/docker_data_root_snapshot.sh preflight \
+          "/mnt/ci-image-snapshot-${PLATFORM//\//_}-${PYTHON}.tar.zst.meta"
+      shell: bash
+      if: steps.snapshot-metadata.outcome == 'success'
+    - name: "Download CI image snapshot ${{ inputs.platform }}:${{ inputs.python }}"
+      continue-on-error: true
+      uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c  # v8.0.1
+      with:
+        name: >-
+          ci-image-snapshot-v2-${{ inputs.python }}-${{
+          inputs.platform == 'linux/amd64' && 'amd64' || 'arm64' }}
+        path: "/mnt/"
+      id: restore-snapshot
+      if: steps.snapshot-preflight.outcome == 'success'
+    # Falls back to the image stash below whenever the snapshot does not fit this runner's daemon.
+    - name: "Unpack CI image snapshot ${{ inputs.platform }}:${{ inputs.python }}"
+      id: snapshot
+      env:
+        PLATFORM: ${{ inputs.platform }}
+        PYTHON: ${{ inputs.python }}
+      run: |
+        if ./scripts/ci/docker_data_root_snapshot.sh restore \
+            "/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.restore-snapshot.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) || '' }}
@@ -90,6 +142,7 @@ runs:
         path: "/mnt/"
         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 }}
@@ -98,3 +151,4 @@ runs:
       run: >
         breeze ${IMAGE_TYPE}-image load --platform "${PLATFORM}" --python "${PYTHON}" --image-file-dir "/mnt"
       shell: bash
+      if: steps.snapshot.outputs.restored != 'true'

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

@@ -401,3 +401,72 @@ jobs:
           steps.stashed-image.outputs.reusable != 'true'
       - name: "Check disk space after build"
         run: df -H
+
+  # Run checkout code in a separate job with no branch-cache publications. This optional
+  # producer hands the snapshot only to consumers of the same workflow run.
+  snapshot-ci-images:
+    name: "Snapshot CI ${{ inputs.platform }} image ${{ matrix.python-version }}"
+    needs: build-ci-images
+    if: >
+      github.event_name == 'pull_request' &&
+      inputs.upload-image-artifact == 'true' && inputs.image-stash-ref == '' &&
+      (inputs.checkout-ref == '' || inputs.checkout-ref == github.sha)
+    continue-on-error: true
+    timeout-minutes: 20
+    runs-on: ${{ fromJSON(inputs.runners) }}
+    permissions:
+      contents: read
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version: ${{ fromJSON(inputs.python-versions) }}
+    env:
+      PYTHON_MAJOR_MINOR_VERSION: ${{ matrix.python-version }}
+      GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+    steps:
+      - name: "Checkout sources"
+        uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1  # v7.0.1
+        with:
+          ref: ${{ github.sha }}
+          persist-credentials: false
+      - name: "Prepare authoritative CI image"
+        uses: ./.github/actions/prepare_breeze_and_image
+        with:
+          platform: ${{ inputs.platform }}
+          python: ${{ matrix.python-version }}
+          use-uv: ${{ inputs.use-uv }}
+          make-mnt-writeable-and-cleanup: 'true'
+          restore-snapshot: 'false'
+      - name: "Snapshot CI image ${{ inputs.platform }}:${{ env.PYTHON_MAJOR_MINOR_VERSION }}"
+        id: snapshot-export
+        continue-on-error: true
+        env:
+          PLATFORM: ${{ inputs.platform }}
+        run: >
+          ./scripts/ci/docker_data_root_snapshot.sh create
+          "/mnt/ci-image-snapshot-${PLATFORM//\//_}-${PYTHON_MAJOR_MINOR_VERSION}.tar.zst"
+        shell: bash
+      - 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-v2-${{ env.PYTHON_MAJOR_MINOR_VERSION }}-${{
+            inputs.platform == 'linux/amd64' && 'amd64' || 'arm64' }}
+          path: "/mnt/ci-image-snapshot-*-${{ env.PYTHON_MAJOR_MINOR_VERSION }}.tar.zst*"
+          if-no-files-found: 'error'
+          compression-level: '0'
+          retention-days: '2'
+        if: steps.snapshot-export.outcome == 'success'
+      - 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-meta-v2-${{ env.PYTHON_MAJOR_MINOR_VERSION }}-${{
+            inputs.platform == 'linux/amd64' && 'amd64' || 'arm64' }}
+          path: "/mnt/ci-image-snapshot-*-${{ env.PYTHON_MAJOR_MINOR_VERSION }}.tar.zst.meta"
+          if-no-files-found: 'error'
+          compression-level: '9'
+          retention-days: '2'
+        if: steps.snapshot-export.outcome == 'success'

scripts/ci/docker_data_root_snapshot.sh

@@ -0,0 +1,175 @@
+#!/usr/bin/env bash
+# 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.
+
+# Moves the CI image between jobs as a copy of the Docker data directory that holds it, so a job
+# restores it with one extraction instead of `docker image load` unpacking and checksumming every
+# layer again.
+#
+#   docker_data_root_snapshot.sh create SNAPSHOT_FILE
+#       Keeps only the CI image in the daemon and writes its image store to SNAPSHOT_FILE, plus
+#       SNAPSHOT_FILE.meta naming the daemon that wrote it and the image.
+#   docker_data_root_snapshot.sh restore SNAPSHOT_FILE
+#       Restores the image. Incompatible snapshots leave the daemon untouched. Failed
+#       materialization cleans partial image state and restarts Docker for the normal stash path.
+#   docker_data_root_snapshot.sh preflight METADATA_FILE
+#       Checks compatibility without requiring the multi-gigabyte snapshot archive.
+set -euo pipefail
+
+DATA_ROOT="/var/lib/docker"
+DAEMON_STOPPED=false
+RESTORE_IN_PROGRESS=false
+
+function cleanup_daemon() {
+    local result=$?
+    trap - EXIT
+    if [[ "${RESTORE_IN_PROGRESS}" == true ]]; then
+        stop_daemon
+        remove_image_store
+        start_daemon
+    elif [[ "${DAEMON_STOPPED}" == true ]]; then
+        start_daemon
+    fi
+    exit "${result}"
+}
+trap cleanup_daemon EXIT
+
+function check_supported_daemon() {
+    local daemon=$1
+    local _ driver root
+    read -r _ driver _ root <<< "${daemon}"
+    if [[ "${driver}" != overlay2 || "${root}" != "${DATA_ROOT}" ]]; then
+        echo "Unsupported Docker image store: ${daemon}"
+        exit 3
+    fi
+}
+
+# A data directory is only readable by the daemon version and storage driver that wrote it.
+function daemon_fingerprint() {
+    docker info --format '{{.ServerVersion}} {{.Driver}} {{.Architecture}} {{.DockerRootDir}}'
+}
+
+function stop_daemon() {
+    DAEMON_STOPPED=true
+    sudo systemctl stop docker.socket docker
+}
+
+function start_daemon() {
+    sudo systemctl start docker
+    DAEMON_STOPPED=false
+}
+
+function remove_image_store() {
+    sudo rm -rf "${DATA_ROOT}/image" "${DATA_ROOT}/overlay2"
+}
+
+function create_snapshot() {
+    local snapshot_file="${1}"
+    local image_id
+    check_supported_daemon "$(daemon_fingerprint)"
+    image_id="$(docker images --quiet --filter "label=org.apache.airflow.image=airflow-ci" | sort -u)"
+    if [[ -z "${image_id}" || "${image_id}" == *$'\n'* ]]; then
+        echo "Expected exactly one CI image in the daemon, found: '${image_id}'" >&2
+        exit 1
+    fi
+    docker ps --all --quiet | xargs --no-run-if-empty docker rm --force >/dev/null
+    docker images --quiet | sort -u | { grep --invert-match --fixed-strings --line-regexp "${image_id}" || true; } \
+        | xargs --no-run-if-empty docker rmi --force >/dev/null
+    docker builder prune --all --force >/dev/null
+    printf '%s %s %s\n' "$(daemon_fingerprint)" "${image_id}" "$(git rev-parse HEAD)" > "${snapshot_file}.meta"
+    stop_daemon
+    sudo tar --directory "${DATA_ROOT}" --xattrs --acls --numeric-owner --create --file - image overlay2 \
+        | zstd -3 -T0 --quiet --force -o "${snapshot_file}"
+    start_daemon
+}
+
+function validate_metadata() {
+    local metadata_file="${1}"
+    local -a meta
+    local snapshot_fingerprint daemon existing_images existing_containers
+    if [[ ! -f "${metadata_file}" ]]; then
+        echo "No snapshot metadata at ${metadata_file}"
+        exit 2
+    fi
+    read -r -a meta < "${metadata_file}"
+    if [[ ${#meta[@]} != 6 ]]; then
+        echo "Invalid snapshot metadata"
+        exit 3
+    fi
+    if [[ "${meta[5]}" != "$(git rev-parse HEAD)" ]]; then
+        echo "Snapshot belongs to a different checkout revision"
+        exit 3
+    fi
+    snapshot_fingerprint="${meta[*]:0:4}"
+    daemon="$(daemon_fingerprint)"
+    check_supported_daemon "${daemon}"
+    if [[ "${snapshot_fingerprint}" != "${daemon}" ]]; then
+        echo "The snapshot was written by '${snapshot_fingerprint}', this daemon is '${daemon}'"
+        exit 3
+    fi
+    existing_images="$(docker images --all --quiet)"
+    existing_containers="$(docker ps --all --quiet)"
+    if [[ -n "${existing_images}" || -n "${existing_containers}" ]]; then
+        echo "The daemon already holds images, which restoring the snapshot would drop"
+        exit 3
+    fi
+}
+
+function preflight_snapshot() {
+    validate_metadata "${1}"
+}
+
+function restore_snapshot() {
+    local snapshot_file="${1}"
+    local -a meta
+    local image_id
+    if [[ ! -f "${snapshot_file}" || ! -f "${snapshot_file}.meta" ]]; then
+        echo "No snapshot at ${snapshot_file}"
+        exit 2
+    fi
+    validate_metadata "${snapshot_file}.meta"
+    read -r -a meta < "${snapshot_file}.meta"
+    image_id="${meta[4]}"
+    stop_daemon
+    RESTORE_IN_PROGRESS=true
+    remove_image_store
+    zstd -d -T0 --quiet --stdout "${snapshot_file}" \
+        | sudo tar --directory "${DATA_ROOT}" --xattrs --acls --numeric-owner --extract --file -
+    start_daemon
+    if ! docker run --rm --entrypoint /bin/bash "${image_id}" -c true; then
+        echo "The restored image ${image_id} does not run"
+        exit 4
+    fi
+    RESTORE_IN_PROGRESS=false
+    rm -f "${snapshot_file}" "${snapshot_file}.meta"
+}
+
+case "${1:-}" in
+    create)
+        create_snapshot "${2}"
+        ;;
+    restore)
+        restore_snapshot "${2}"
+        ;;
+    preflight)
+        preflight_snapshot "${2}"
+        ;;
+    *)
+        echo "Usage: ${0} create|preflight|restore FILE" >&2
+        exit 1
+        ;;
+esac