Simplify caching mechanisms for CI and PROD images

For a long time we had used a sophisticated mechanism to speed up
our CI jobs by building the images in "pull_request_target" workflow
and pushing them to GitHub registry. That however had several drawbacks:

* CI image was complex when it comes to layer setup (we had to pre-
  cache installed dependencies by installing them from branch tip

* The pull_request_target is a very dangerous workflow, we had a number
  of security problems with it (and it's difficult to debug)

* Caching of `pip` and `uv` was not used because it increased size of
  the image significantly

This PR significantly improves the caching mechanisms for the images
building of several advacements that were not possible before:

* The upload-artifacts@v4 action and improved stash action developed
  by @assignUser and published in "apache/infrastructure-actions"
  allows us to store all images (8GB per run) in artifacts rather
  than in registry - so we can do the image build once and share
  it with all the jobs.

* The uv speed is "enough" to allow occasional installation of Airlfow
  locally. This allows to utilize cache-mount and locally build uv
  cache, rather than rely on "remote" cache when we are building
  local images for breeze. The first time you build local breeze
  image it will take 2-5 more minutes (depending on your network
  speed, but because we can utilise cache mounts, every subsequent
  build should be very fast - even if all dependencies change. Using
  uv also allows to "always" reinstall airflow when you build the
  image even if single source file changed, because with cache
  it takes sub-seconds to reinstall airflow and all dependencies.

* the cache mounts are not included in the image size, and since we
  can export and import images in CI in artifacts and we do not
  need to rebuild them, the images shared as compressed artifacts are
  relatively small (2GB) - cache of `uv` is around 4GB on top of that
  so sharing image built in the "build image" job with other jobs
  in the same workflow is fast.

* we are still using registry cache for the "non-python" parts of
  the image - both CI and breeze image build speed benefit from using
  the image cache for system dependencies, database clients etc.

Fixes: #42999
Fixes: #43268

Author: potiuk

.github/actions/checkout_target_commit/action.yml

@@ -0,0 +1,71 @@
+# 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: 'Prepare all CI images'
+description: 'Recreates current python CI images from artifacts for all python versions'
+inputs:
+  python-versions-list-as-string:
+    description: 'Stringified array of all Python versions to test - separated by spaces.'
+    required: true
+  platform:
+    description: 'Platform for the build - linux/amd64 or linux/arm64'
+    required: true
+outputs:
+  host-python-version:
+    description: Python version used in host
+    value: ${{ steps.breeze.outputs.host-python-version }}
+runs:
+  using: "composite"
+  steps:
+    - name: "Cleanup docker"
+      run: ./scripts/ci/cleanup_docker.sh
+      shell: bash
+    # TODO: Currently we cannot loop through the list of python versions and have dynamic list of
+    #       tasks. Instead we hardcode all possible python versions and they - but
+    #       this should be implemented in stash action as list of keys to download.
+    #       That includes 3.8 - 3.12 as we are backporting it to v2-10-test branch
+    - name: "Restore CI docker image ${{ inputs.platform }}:3.8"
+      uses: ./.github/actions/prepare_single_ci_image
+      with:
+        platform: ${{ inputs.platform }}
+        python: "3.8"
+        python-versions-list-as-string: ${{ inputs.python-versions-list-as-string }}
+    - name: "Restore CI docker image ${{ inputs.platform }}:3.9"
+      uses: ./.github/actions/prepare_single_ci_image
+      with:
+        platform: ${{ inputs.platform }}
+        python: "3.9"
+        python-versions-list-as-string: ${{ inputs.python-versions-list-as-string }}
+    - name: "Restore CI docker image ${{ inputs.platform }}:3.10"
+      uses: ./.github/actions/prepare_single_ci_image
+      with:
+        platform: ${{ inputs.platform }}
+        python: "3.10"
+        python-versions-list-as-string: ${{ inputs.python-versions-list-as-string }}
+    - name: "Restore CI docker image ${{ inputs.platform }}:3.11"
+      uses: ./.github/actions/prepare_single_ci_image
+      with:
+        platform: ${{ inputs.platform }}
+        python: "3.11"
+        python-versions-list-as-string: ${{ inputs.python-versions-list-as-string }}
+    - name: "Restore CI docker image ${{ inputs.platform }}:3.12"
+      uses: ./.github/actions/prepare_single_ci_image
+      with:
+        platform: ${{ inputs.platform }}
+        python: "3.12"
+        python-versions-list-as-string: ${{ inputs.python-versions-list-as-string }}

.github/actions/prepare_all_ci_images/action.yml

@@ -16,30 +16,48 @@
 # under the License.
 #
 ---
-name: 'Prepare breeze && current python image'
-description: 'Installs breeze and pulls current python image'
+name: 'Prepare breeze && current image (CI or PROD)'
+description: 'Installs breeze and recreates current python image from artifact'
 inputs:
-  pull-image-type:
-    description: 'Which image to pull'
-    default: CI
+  python:
+    description: 'Python version for image to prepare'
+    required: true
+  image-type:
+    description: 'Which image type to prepare (CI/PROD)'
+    default: "CI"
+  platform:
+    description: 'Platform for the build - linux/amd64 or linux/arm64'
+    required: true
 outputs:
   host-python-version:
     description: Python version used in host
     value: ${{ steps.breeze.outputs.host-python-version }}
 runs:
   using: "composite"
   steps:
+    - name: "Cleanup docker"
+      run: ./scripts/ci/cleanup_docker.sh
+      shell: bash
     - name: "Install Breeze"
       uses: ./.github/actions/breeze
       id: breeze
-    - name: Login to ghcr.io
-      shell: bash
-      run: echo "${{ env.GITHUB_TOKEN }}" | docker login ghcr.io -u ${{ github.actor }} --password-stdin
-    - name: Pull CI image ${{ env.PYTHON_MAJOR_MINOR_VERSION }}:${{ env.IMAGE_TAG }}
+    - name: "Restore CI docker image ${{ inputs.platform }}:${{ inputs.python }}"
+      uses: apache/infrastructure-actions/stash/restore@c94b890bbedc2fc61466d28e6bd9966bc6c6643c
+      with:
+        key: "ci-image-save-${{ inputs.platform }}-${{ inputs.python }}"
+        path: "/tmp/"
+      if: inputs.image-type == 'CI'
+    - name: "Load CI image ${{ inputs.platform }}:${{ inputs.python }}"
+      run: breeze ci-image load --platform ${{ inputs.platform }} --python ${{ inputs.python }}
       shell: bash
-      run: breeze ci-image pull --tag-as-latest
-      if: inputs.pull-image-type == 'CI'
-    - name: Pull PROD image ${{ env.PYTHON_MAJOR_MINOR_VERSION }}:${{ env.IMAGE_TAG }}
+      if: inputs.image-type == 'CI'
+    - name: "Restore PROD docker image ${{ inputs.platform }}:${{ inputs.python }}"
+      uses: apache/infrastructure-actions/stash/restore@c94b890bbedc2fc61466d28e6bd9966bc6c6643c
+      with:
+        key: "prod-image-save-${{ inputs.platform }}-${{ inputs.python }}"
+        path: "/tmp/"
+      if: inputs.image-type == 'PROD'
+    - name: "Load PROD image ${{ inputs.platform }}:${{ inputs.python }}"
+      run: breeze prod-image load --platform ${{ inputs.platform }} --python ${{ inputs.python }}
       shell: bash
-      run: breeze prod-image pull --tag-as-latest
-      if: inputs.pull-image-type == 'PROD'
+      if: inputs.image-type == 'PROD'

.github/actions/prepare_breeze_and_image/action.yml

@@ -0,0 +1,47 @@
+# 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: 'Prepare single images'
+description: 'Recreates current python image from artifacts'
+inputs:
+  python:
+    description: 'Python version for image to prepare'
+    required: true
+  python-versions-list-as-string:
+    description: 'Stringified array of all Python versions to prepare - separated by spaces.'
+    required: true
+  platform:
+    description: 'Platform for the build - linux/amd64 or linux/arm64'
+    required: true
+outputs:
+  host-python-version:
+    description: Python version used in host
+    value: ${{ steps.breeze.outputs.host-python-version }}
+runs:
+  using: "composite"
+  steps:
+    - name: "Restore CI docker images ${{ inputs.platform }}:${{ inputs.python }}"
+      uses: apache/infrastructure-actions/stash/restore@c94b890bbedc2fc61466d28e6bd9966bc6c6643c
+      with:
+        key: "ci-image-save-${{ inputs.platform }}-${{ inputs.python }}"
+        path: "/tmp/"
+      if: contains(inputs.python-versions-list-as-string, inputs.python)
+    - name: "Load CI image ${{ inputs.platform }}:${{ inputs.python }}"
+      run: breeze ci-image load --platform "${{ inputs.platform }}" --python "${{ inputs.python }}"
+      shell: bash
+      if: contains(inputs.python-versions-list-as-string, inputs.python)

.github/actions/prepare_single_ci_image/action.yml

@@ -32,10 +32,6 @@ on:  # yamllint disable-line rule:truthy
         description: "The array of labels (in json form) determining self-hosted runners."
         required: true
         type: string
-      image-tag:
-        description: "Tag to set for the image"
-        required: true
-        type: string
       python-versions:
         description: "The list of python versions (stringified JSON array) to run the tests on."
         required: true
@@ -103,8 +99,6 @@ jobs:
       contents: read
       # This write is only given here for `push` events from "apache/airflow" repo. It is not given for PRs
       # from forks. This is to prevent malicious PRs from creating images in the "apache/airflow" repo.
-      # For regular build for PRS this "build-prod-images" workflow will be skipped anyway by the
-      # "in-workflow-build" condition
       packages: write
     secrets: inherit
     with:
@@ -159,7 +153,7 @@ jobs:
 #  # There is no point in running this one in "canary" run, because the above step is doing the
 #  # same build anyway.
 #  build-ci-arm-images:
-#    name: Build CI ARM images (in-workflow)
+#    name: Build CI ARM images
 #    uses: ./.github/workflows/ci-image-build.yml
 #    permissions:
 #      contents: read
@@ -169,7 +163,6 @@ jobs:
 #      push-image: "false"
 #      runs-on-as-json-public: ${{ inputs.runs-on-as-json-public }}
 #      runs-on-as-json-self-hosted: ${{ inputs.runs-on-as-json-self-hosted }}
-#      image-tag: ${{ inputs.image-tag }}
 #      python-versions: ${{ inputs.python-versions }}
 #      platform: "linux/arm64"
 #      branch: ${{ inputs.branch }}

.github/workflows/additional-ci-image-checks.yml

@@ -32,10 +32,6 @@ on:  # yamllint disable-line rule:truthy
         description: "Branch used to construct constraints URL from."
         required: true
         type: string
-      image-tag:
-        description: "Tag to set for the image"
-        required: true
-        type: string
       upgrade-to-newer-dependencies:
         description: "Whether to upgrade to newer dependencies (true/false)"
         required: true
@@ -70,7 +66,6 @@ jobs:
       default-python-version: ${{ inputs.default-python-version }}
       branch: ${{ inputs.default-branch }}
       use-uv: "false"
-      image-tag: ${{ inputs.image-tag }}
       build-provider-packages: ${{ inputs.default-branch == 'main' }}
       upgrade-to-newer-dependencies: ${{ inputs.upgrade-to-newer-dependencies }}
       chicken-egg-providers: ${{ inputs.chicken-egg-providers }}
@@ -88,7 +83,6 @@ jobs:
       default-python-version: ${{ inputs.default-python-version }}
       branch: ${{ inputs.default-branch }}
       use-uv: "false"
-      image-tag: ${{ inputs.image-tag }}
       build-provider-packages: ${{ inputs.default-branch == 'main' }}
       upgrade-to-newer-dependencies: ${{ inputs.upgrade-to-newer-dependencies }}
       chicken-egg-providers: ${{ inputs.chicken-egg-providers }}
@@ -117,36 +111,25 @@ jobs:
           persist-credentials: false
       - name: "Cleanup docker"
         run: ./scripts/ci/cleanup_docker.sh
-      - name: "Install Breeze"
-        uses: ./.github/actions/breeze
-      - name: Login to ghcr.io
-        shell: bash
-        run: echo "${{ env.GITHUB_TOKEN }}" | docker login ghcr.io -u ${{ github.actor }} --password-stdin
-      - name: Pull PROD image ${{ inputs.default-python-version}}:${{ inputs.image-tag }}
-        run: breeze prod-image pull --tag-as-latest
-        env:
-          PYTHON_MAJOR_MINOR_VERSION: "${{ inputs.default-python-version }}"
-          IMAGE_TAG: "${{ inputs.image-tag }}"
-      - name: "Setup python"
-        uses: actions/setup-python@v5
+      - name: "Prepare breeze & PROD image: ${{ inputs.default-python-version }}"
+        uses: ./.github/actions/prepare_breeze_and_image
         with:
-          python-version: ${{ inputs.default-python-version }}
-          cache: 'pip'
-          cache-dependency-path: ./dev/requirements.txt
+          platform: "linux/amd64"
+          image-type: "PROD"
+          python: ${{ inputs.default-python-version }}
       - name: "Test examples of PROD image building"
         run: "
           cd ./docker_tests && \
           python -m pip install -r requirements.txt && \
           TEST_IMAGE=\"ghcr.io/${{ github.repository }}/${{ inputs.default-branch }}\
-          /prod/python${{ inputs.default-python-version }}:${{ inputs.image-tag }}\" \
+          /prod/python${{ inputs.default-python-version }}\" \
           python -m pytest test_examples_of_prod_image_building.py -n auto --color=yes"
 
   test-docker-compose-quick-start:
     timeout-minutes: 60
-    name: "Docker-compose quick start with PROD image verifying"
+    name: "Docker Compose quick start with PROD image verifying"
     runs-on: ${{ fromJSON(inputs.runs-on-as-json-public) }}
     env:
-      IMAGE_TAG: "${{ inputs.image-tag }}"
       PYTHON_MAJOR_MINOR_VERSION: "${{ inputs.default-python-version }}"
       GITHUB_REPOSITORY: ${{ github.repository }}
       GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
@@ -161,14 +144,12 @@ jobs:
         with:
           fetch-depth: 2
           persist-credentials: false
-      - name: "Cleanup docker"
-        run: ./scripts/ci/cleanup_docker.sh
-      - name: "Install Breeze"
-        uses: ./.github/actions/breeze
-      - name: Login to ghcr.io
-        shell: bash
-        run: echo "${{ env.GITHUB_TOKEN }}" | docker login ghcr.io -u ${{ github.actor }} --password-stdin
-      - name: "Pull image ${{ inputs.default-python-version}}:${{ inputs.image-tag }}"
-        run: breeze prod-image pull --tag-as-latest
+      - name: "Prepare breeze & PROD image: ${{ env.PYTHON_MAJOR_MINOR_VERSION }}"
+        uses: ./.github/actions/prepare_breeze_and_image
+        with:
+          platform: "linux/amd64"
+          image-type: "PROD"
+          python: ${{ env.PYTHON_MAJOR_MINOR_VERSION }}
+        id: breeze
       - name: "Test docker-compose quick start"
         run: breeze testing docker-compose-tests