Improve caching strategy across the board of CI workflow

potiuk · potiuk · commit c7041f9e1d5f · 2024-12-30T23:46:13.000+01:00
We are using various caches in our build and so far - due to the way how "standard" caching works, PRs from forks could not effectively use the cache from main Airflow repository - because caches are not shared with other repositories - so the PRs builds could only use cache effectively when they were rebased and continued running from the same fork. This PR improves caching strategy using "stash" action from the ASF. Unlike `cache` - the action uses artifacts to store cache, and that makes it possible for the stash action to use such cache uploaded from `main` canary builds in PRs coming from the fork. As part of this change all the places where setup-python was used and breeze installed afterwards were reviewed and updated to use only breeze installation action (it already installs python) and this action has been improved to use UV caching effectively. Overall this PR should decrease setup overhead for many jobs across the CI workflow. Follow-up after #45266
diff --git a/.github/actions/breeze/action.yml b/.github/actions/breeze/action.yml
@@ -33,8 +33,11 @@ runs:
       uses: actions/setup-python@v5
       with:
         python-version: ${{ inputs.python-version }}
-        cache: 'pip'
-        cache-dependency-path: ./dev/breeze/pyproject.toml
+    - name: "Restore uv cache for breeze"
+      uses: apache/infrastructure-actions/stash/restore@c94b890bbedc2fc61466d28e6bd9966bc6c6643c
+      with:
+        key: "breeze-uv-cache-${{ inputs.python-version }}-${{ hashFiles('dev/breeze/pyproject.toml') }}"
+        path: ~/.cache/uv
     - name: "Install Breeze"
       shell: bash
       run: ./scripts/ci/install_breeze.sh
diff --git a/.github/actions/install-pre-commit/action.yml b/.github/actions/install-pre-commit/action.yml
@@ -19,31 +19,31 @@
 name: 'Install pre-commit'
 description: 'Installs pre-commit and related packages'
 inputs:
+  # TODO(potiuk): automate update of these versions
   python-version:
     description: 'Python version to use'
-    default: 3.9
+    default: "3.9"
   uv-version:
     description: 'uv version to use'
-    default: 0.5.5
+    default: "0.5.13"
   pre-commit-version:
     description: 'pre-commit version to use'
-    default: 4.0.1
+    default: "4.0.1"
   pre-commit-uv-version:
     description: 'pre-commit-uv version to use'
-    default: 4.1.4
+    default: "4.1.4"
 runs:
   using: "composite"
   steps:
+    - name: "Restore pre-commit cache"
+      uses: apache/infrastructure-actions/stash/restore@c94b890bbedc2fc61466d28e6bd9966bc6c6643c
+      with:
+        key: "pre-commit-cache-${{ inputs.python-version }}-${{ hashFiles('.pre-commit-config.yaml') }}"
+        path: ~/.cache/pre-commit/
     - name: Install pre-commit, uv, and pre-commit-uv
       shell: bash
       run: |
         pip install uv==${{inputs.uv-version}} || true
         uv tool install pre-commit==${{inputs.pre-commit-version}} --with uv==${{inputs.uv-version}} \
         --with pre-commit-uv==${{inputs.pre-commit-uv-version}}
-    - name: Cache pre-commit envs
-      uses: actions/cache@v4
-      with:
-        path: ~/.cache/pre-commit
-        key: "pre-commit-${{inputs.python-version}}-${{ hashFiles('.pre-commit-config.yaml') }}"
-        restore-keys: |
-          pre-commit-${{inputs.python-version}}-
+        pre-commit install-hooks
diff --git a/.github/workflows/basic-tests.yml b/.github/workflows/basic-tests.yml
@@ -72,16 +72,11 @@ jobs:
           persist-credentials: false
       - name: "Cleanup docker"
         run: ./scripts/ci/cleanup_docker.sh
-      - uses: actions/setup-python@v5
-        with:
-          python-version: "${{ inputs.default-python-version }}"
-          cache: 'pip'
-          cache-dependency-path: ./dev/breeze/pyproject.toml
-      - run: pip install --editable ./dev/breeze/
+      - name: "Install Breeze"
+        uses: ./.github/actions/breeze
       - run: python -m pytest -n auto --color=yes
         working-directory: ./dev/breeze/
 
-
   tests-ui:
     timeout-minutes: 10
     name: React UI tests
@@ -108,15 +103,24 @@ jobs:
           node-version: 21
           cache: 'pnpm'
           cache-dependency-path: 'airflow/ui/pnpm-lock.yaml'
-      - name: "Cache eslint"
-        uses: actions/cache@v4
+      - name: "Restore eslint cache (ui)"
+        uses: apache/infrastructure-actions/stash/restore@c94b890bbedc2fc61466d28e6bd9966bc6c6643c
         with:
-          path: 'airflow/ui/node_modules'
+          path: 'airflow/ui/node_modules/'
           key: ${{ runner.os }}-ui-node-modules-${{ hashFiles('airflow/ui/**/pnpm-lock.yaml') }}
+        id: restore-eslint-cache
       - run: cd airflow/ui && pnpm install --frozen-lockfile
       - run: cd airflow/ui && pnpm test
         env:
           FORCE_COLOR: 2
+      - name: "Save eslint cache (ui)"
+        uses: apache/infrastructure-actions/stash/save@c94b890bbedc2fc61466d28e6bd9966bc6c6643c
+        with:
+          path: 'airflow/ui/node_modules/'
+          key: ${{ runner.os }}-ui-node-modules-${{ hashFiles('airflow/ui/**/pnpm-lock.yaml') }}
+          if-no-files-found: 'error'
+          retention-days: '2'
+        if: steps.restore-eslint-cache.outputs.cache-hit != 'true'
 
   tests-www:
     timeout-minutes: 10
@@ -137,15 +141,53 @@ jobs:
         uses: actions/setup-node@v4
         with:
           node-version: 21
-      - name: "Cache eslint"
-        uses: actions/cache@v4
+      - name: "Restore eslint cache (www)"
+        uses: apache/infrastructure-actions/stash/restore@c94b890bbedc2fc61466d28e6bd9966bc6c6643c
         with:
-          path: 'airflow/www/node_modules'
+          path: 'airflow/www/node_modules/'
           key: ${{ runner.os }}-www-node-modules-${{ hashFiles('airflow/www/**/yarn.lock') }}
+        id: restore-eslint-cache
       - run: yarn --cwd airflow/www/ install --frozen-lockfile --non-interactive
       - run: yarn --cwd airflow/www/ run test
         env:
           FORCE_COLOR: 2
+      - name: "Save eslint cache (www)"
+        uses: apache/infrastructure-actions/stash/save@c94b890bbedc2fc61466d28e6bd9966bc6c6643c
+        with:
+          path: 'airflow/www/node_modules/'
+          key: ${{ runner.os }}-www-node-modules-${{ hashFiles('airflow/www/**/yarn.lock') }}
+          if-no-files-found: 'error'
+          retention-days: '2'
+        if: steps.restore-eslint-cache.outputs.cache-hit != 'true'
+
+  install-pre-commit:
+    timeout-minutes: 5
+    name: "Install pre-commit for cache"
+    runs-on: ${{ fromJSON(inputs.runs-on-as-json-default) }}
+    env:
+      PYTHON_MAJOR_MINOR_VERSION: "${{ inputs.default-python-version }}"
+    if: inputs.basic-checks-only == 'true'
+    steps:
+      - name: "Install Breeze"
+        uses: ./.github/actions/breeze
+        id: breeze
+      - name: "Install pre-commit"
+        uses: ./.github/actions/install-pre-commit
+        id: pre-commit
+        with:
+          python-version: ${{steps.breeze.outputs.host-python-version}}
+      # Saving pre-commit cache should happen only in one job in the entire workflow - because otherwise
+      # it will cause 409 conflict errors - see https://github.com/actions/upload-artifact/issues/478
+      # the way it works with airflow - even if the same action is in "ci-image-tests" the if condition
+      # above `if: inputs.basic-checks-only == 'true'` will prevent it from running in the other job
+      - name: "Save pre-commit cache"
+        uses: apache/infrastructure-actions/stash/save@c94b890bbedc2fc61466d28e6bd9966bc6c6643c
+        with:
+          # yamllint disable rule:line-length
+          key: "pre-commit-cache-${{ steps.breeze.outputs.host-python-version }}-${{ hashFiles('.pre-commit-config.yaml') }}"
+          path: ~/.cache/pre-commit/
+          if-no-files-found: 'error'
+          retention-days: '2'
 
   # Those checks are run if no image needs to be built for checks. This is for simple changes that
   # Do not touch any of the python code or any of the important files that might require building
@@ -154,6 +196,7 @@ jobs:
     timeout-minutes: 30
     name: "Static checks: basic checks only"
     runs-on: ${{ fromJSON(inputs.runs-on-as-json-public) }}
+    needs: install-pre-commit
     if: inputs.basic-checks-only == 'true'
     steps:
       - name: "Cleanup repo"
@@ -165,18 +208,6 @@ jobs:
           persist-credentials: false
       - name: "Cleanup docker"
         run: ./scripts/ci/cleanup_docker.sh
-      - name: "Setup python"
-        uses: actions/setup-python@v5
-        with:
-          python-version: ${{ inputs.default-python-version }}
-          cache: 'pip'
-          cache-dependency-path: ./dev/breeze/pyproject.toml
-      - name: "Setup python"
-        uses: actions/setup-python@v5
-        with:
-          python-version: "${{ inputs.default-python-version }}"
-          cache: 'pip'
-          cache-dependency-path: ./dev/breeze/pyproject.toml
       - name: "Install Breeze"
         uses: ./.github/actions/breeze
         id: breeze
@@ -216,6 +247,7 @@ jobs:
     timeout-minutes: 45
     name: "Upgrade checks"
     runs-on: ${{ fromJSON(inputs.runs-on-as-json-public) }}
+    needs: install-pre-commit
     env:
       PYTHON_MAJOR_MINOR_VERSION: "${{ inputs.default-python-version }}"
     if: inputs.canary-run == 'true' && inputs.latest-versions-only != 'true'
@@ -229,12 +261,14 @@ jobs:
           persist-credentials: false
       - name: "Cleanup docker"
         run: ./scripts/ci/cleanup_docker.sh
-      # Install python from scratch. No cache used. We always want to have fresh version of everything
-      - uses: actions/setup-python@v5
+      - name: "Install Breeze"
+        uses: ./.github/actions/breeze
+        id: breeze
+      - name: "Install pre-commit"
+        uses: ./.github/actions/install-pre-commit
+        id: pre-commit
         with:
-          python-version: "${{ inputs.default-python-version }}"
-      - name: "Install latest pre-commit"
-        run: pip install pre-commit
+          python-version: ${{steps.breeze.outputs.host-python-version}}
       - name: "Autoupdate all pre-commits"
         run: pre-commit autoupdate
       - name: "Run automated upgrade for black"
diff --git a/.github/workflows/ci-image-build.yml b/.github/workflows/ci-image-build.yml
@@ -162,5 +162,5 @@ jobs:
           key: "ci-image-save-${{ inputs.platform }}-${{ env.PYTHON_MAJOR_MINOR_VERSION }}"
           path: "/tmp/ci-image-save-*-${{ env.PYTHON_MAJOR_MINOR_VERSION }}.tar"
           if-no-files-found: 'error'
-          retention-days: 2
+          retention-days: '2'
         if: inputs.upload-image-artifact == 'true'
diff --git a/.github/workflows/ci-image-checks.yml b/.github/workflows/ci-image-checks.yml
@@ -106,10 +106,40 @@ on:  # yamllint disable-line rule:truthy
         type: string
 
 jobs:
+  install-pre-commit:
+    timeout-minutes: 5
+    name: "Install pre-commit for cache"
+    runs-on: ${{ fromJSON(inputs.runs-on-as-json-default) }}
+    env:
+      PYTHON_MAJOR_MINOR_VERSION: "${{ inputs.default-python-version }}"
+    if: inputs.basic-checks-only == 'false'
+    steps:
+      - name: "Install Breeze"
+        uses: ./.github/actions/breeze
+        id: breeze
+      - name: "Install pre-commit"
+        uses: ./.github/actions/install-pre-commit
+        id: pre-commit
+        with:
+          python-version: ${{steps.breeze.outputs.host-python-version}}
+      # Saving pre-commit cache should happen only in one job in the entire workflow - because otherwise
+      # it will cause 409 conflict errors - see https://github.com/actions/upload-artifact/issues/478
+      # the way it works with airflow - even if the same action is in "basic-tests" the if condition
+      # above `if: inputs.basic-checks-only == 'false'` will prevent it from running in the other job
+      - name: "Save pre-commit cache"
+        uses: apache/infrastructure-actions/stash/save@c94b890bbedc2fc61466d28e6bd9966bc6c6643c
+        with:
+          # yamllint disable rule:line-length
+          key: "pre-commit-cache-${{ steps.breeze.outputs.host-python-version }}-${{ hashFiles('.pre-commit-config.yaml') }}"
+          path: ~/.cache/pre-commit/
+          if-no-files-found: 'error'
+          retention-days: '2'
+
   static-checks:
     timeout-minutes: 45
     name: "Static checks"
     runs-on: ${{ fromJSON(inputs.runs-on-as-json-default) }}
+    needs: install-pre-commit
     env:
       PYTHON_MAJOR_MINOR_VERSION: "${{ inputs.default-python-version }}"
       UPGRADE_TO_NEWER_DEPENDENCIES: "${{ inputs.upgrade-to-newer-dependencies }}"
@@ -123,12 +153,6 @@ jobs:
         uses: actions/checkout@v4
         with:
           persist-credentials: false
-      - name: "Setup python"
-        uses: actions/setup-python@v5
-        with:
-          python-version: ${{ inputs.default-python-version }}
-          cache: 'pip'
-          cache-dependency-path: ./dev/breeze/pyproject.toml
       - name: "Prepare breeze & CI image: ${{ inputs.default-python-version }}"
         uses: ./.github/actions/prepare_breeze_and_image
         with:
@@ -154,6 +178,7 @@ jobs:
     timeout-minutes: 45
     name: "MyPy checks"
     runs-on: ${{ fromJSON(inputs.runs-on-as-json-default) }}
+    needs: install-pre-commit
     if: inputs.needs-mypy == 'true'
     strategy:
       fail-fast: false
@@ -221,24 +246,30 @@ jobs:
         with:
           platform: "linux/amd64"
           python: ${{ inputs.default-python-version }}
-      - uses: actions/cache@v4
-        id: cache-doc-inventories
+      - name: "Restore docs inventory cache"
+        uses: apache/infrastructure-actions/stash/restore@c94b890bbedc2fc61466d28e6bd9966bc6c6643c
         with:
           path: ./docs/_inventory_cache/
-          key: docs-inventory-${{ hashFiles('pyproject.toml;') }}
-          restore-keys: |
-            docs-inventory-${{ hashFiles('pyproject.toml;') }}
-            docs-inventory-
+          key: docs-inventory-${{ hashFiles('pyproject.toml') }}
+        id: restore-eslint-cache
       - name: "Building docs with ${{ matrix.flag }} flag"
         run: >
           breeze build-docs ${{ inputs.docs-list-as-string }} ${{ matrix.flag }}
+      - name: "Save docs inventory cache"
+        uses: apache/infrastructure-actions/stash/save@c94b890bbedc2fc61466d28e6bd9966bc6c6643c
+        with:
+          path: ./docs/_inventory_cache/
+          key: docs-inventory-${{ hashFiles('pyproject.toml') }}
+          if-no-files-found: 'error'
+          retention-days: '2'
+        if: steps.restore-eslint-cache.outputs.cache-hit != 'true'
       - name: "Upload build docs"
         uses: actions/upload-artifact@v4
         with:
           name: airflow-docs
           path: './docs/_build'
-          retention-days: 7
-          if-no-files-found: error
+          retention-days: '7'
+          if-no-files-found: 'error'
         if: matrix.flag == '--docs-only'
 
   publish-docs:
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -160,6 +160,13 @@ jobs:
           persist-credentials: false
       - name: "Install Breeze"
         uses: ./.github/actions/breeze
+      - name: "Save uv cache for breeze"
+        uses: apache/infrastructure-actions/stash/save@c94b890bbedc2fc61466d28e6bd9966bc6c6643c
+        with:
+          key: "breeze-uv-cache-${{ inputs.python-version }}-${{ hashFiles('dev/breeze/pyproject.toml') }}"
+          path: ~/.cache/uv/
+          if-no-files-found: 'error'
+          retention-days: '2'
       - name: "Get information about the Workflow"
         id: source-run-info
         run: breeze ci get-workflow-info 2>> ${GITHUB_OUTPUT}
@@ -171,7 +178,6 @@ jobs:
           PR_LABELS: "${{ steps.source-run-info.outputs.pr-labels }}"
           COMMIT_REF: "${{ github.sha }}"
           VERBOSE: "false"
-
         run: breeze ci selective-check 2>> ${GITHUB_OUTPUT}
       - name: env
         run: printenv
diff --git a/.github/workflows/k8s-tests.yml b/.github/workflows/k8s-tests.yml
@@ -85,13 +85,14 @@ jobs:
           platform: ${{ inputs.platform }}
           image-type: "prod"
           python: ${{ env.PYTHON_MAJOR_MINOR_VERSION }}
-      - name: "Cache bin folder with tools for kubernetes testing"
-        uses: actions/cache@v4
+      - name: "Restore cache folder with tools for kubernetes testing"
+        uses: apache/infrastructure-actions/stash/restore@c94b890bbedc2fc61466d28e6bd9966bc6c6643c
         with:
-          path: ".build/.k8s-env"
           key: "\
             k8s-env-${{ steps.breeze.outputs.host-python-version }}-\
             ${{ hashFiles('scripts/ci/kubernetes/k8s_requirements.txt','hatch_build.py') }}"
+          path: ".build/.k8s-env"
+        id: restore-k8s-env-cache
       - name: "\
           Run complete K8S tests ${{ matrix.executor }}-${{ env.PYTHON_MAJOR_MINOR_VERSION }}-\
           ${{env.KUBERNETES_VERSION}}-${{ matrix.use-standard-naming }}"
@@ -100,6 +101,16 @@ jobs:
           EXECUTOR: ${{ matrix.executor }}
           USE_STANDARD_NAMING: ${{ matrix.use-standard-naming }}
           VERBOSE: "false"
+      - name: "Save cache folder with tools for kubernetes testing"
+        uses: apache/infrastructure-actions/stash/save@c94b890bbedc2fc61466d28e6bd9966bc6c6643c
+        with:
+          key: "\
+            k8s-env-${{ steps.breeze.outputs.host-python-version }}-\
+            ${{ hashFiles('scripts/ci/kubernetes/k8s_requirements.txt','hatch_build.py') }}"
+          path: ".build/.k8s-env"
+          if-no-files-found: 'error'
+          retention-days: '2'
+        if: steps.restore-k8s-env-cache.outputs.cache-hit != 'true'
       - name: "\
           Upload KinD logs on failure ${{ matrix.executor }}-${{ matrix.kubernetes-combo }}-\
           ${{ matrix.use-standard-naming }}"
@@ -110,7 +121,7 @@ jobs:
             kind-logs-${{ matrix.kubernetes-combo }}-${{ matrix.executor }}-\
             ${{ matrix.use-standard-naming }}"
           path: /tmp/kind_logs_*
-          retention-days: 7
+          retention-days: '7'
       - name: "\
           Upload test resource logs on failure ${{ matrix.executor }}-${{ matrix.kubernetes-combo }}-\
           ${{ matrix.use-standard-naming }}"
@@ -121,7 +132,7 @@ jobs:
             k8s-test-resources-${{ matrix.kubernetes-combo }}-${{ matrix.executor }}-\
             ${{ matrix.use-standard-naming }}"
           path: /tmp/k8s_test_resources_*
-          retention-days: 7
+          retention-days: '7'
       - name: "Delete clusters just in case they are left"
         run: breeze k8s delete-cluster --all
         if: always()
diff --git a/.github/workflows/prod-image-build.yml b/.github/workflows/prod-image-build.yml
