Speed-up constraints generation (apache#52449)

Constraints generation was slow because we run them in a loop and
we tried to run them all on a single machine - trying to utilize
the fact that we only have to build airflow and provider packages
once. But those are pretty fast, comparing to constraint generation
and it's much better to parallelize the constraint jobs and run
them on separatae workers. This will speed up constraint generation
delays that will allow building PROD images and kubernetes checks
faster.

Author: potiuk

.github/actions/breeze/action.yml

@@ -41,6 +41,8 @@ runs:
     - name: "Install Breeze"
       shell: bash
       run: ./scripts/ci/install_breeze.sh
+      env:
+        PYTHON_VERSION: ${{ inputs.python-version }}
     - name: "Free space"
       shell: bash
       run: breeze ci free-space

.github/workflows/ci.yml

@@ -225,6 +225,7 @@ jobs:
     permissions:
       contents: read
       packages: write
+      id-token: write
     if: needs.build-info.outputs.canary-run == 'true'
     with:
       runs-on-as-json-default: ${{ needs.build-info.outputs.runs-on-as-json-default }}
@@ -252,6 +253,8 @@ jobs:
     with:
       runs-on-as-json-public: ${{ needs.build-info.outputs.runs-on-as-json-public }}
       python-versions-list-as-string: ${{ needs.build-info.outputs.python-versions-list-as-string }}
+      python-versions: ${{ needs.build-info.outputs.python-versions }}
+      generate-pypi-constraints: "true"
       # generate no providers constraints only in canary builds - they take quite some time to generate
       # they are not needed for regular builds, they are only needed to update constraints in canaries
       generate-no-providers-constraints: ${{ needs.build-info.outputs.canary-run }}

.github/workflows/finalize-tests.yml

@@ -118,7 +118,7 @@ jobs:
       - name: "Download constraints from the constraints generated by build CI image"
         uses: actions/download-artifact@v4
         with:
-          name: constraints
+          pattern: constraints-*
           path: ./files
       - name: "Diff in constraints for Python: ${{ inputs.python-versions-list-as-string }}"
         run: ./scripts/ci/constraints/ci_diff_constraints.sh

.github/workflows/generate-constraints.yml

@@ -28,6 +28,10 @@ on:  # yamllint disable-line rule:truthy
         description: "Stringified array of all Python versions to test - separated by spaces."
         required: true
         type: string
+      python-versions:
+        description: "JSON-formatted array of Python versions to generate constraints for"
+        required: true
+        type: string
       generate-no-providers-constraints:
         description: "Whether to generate constraints without providers (true/false)"
         required: true
@@ -36,6 +40,10 @@ on:  # yamllint disable-line rule:truthy
         description: "Space-separated list of providers that should be installed from context files"
         required: true
         type: string
+      generate-pypi-constraints:
+        description: "Whether to generate PyPI constraints (true/false)"
+        required: true
+        type: string
       debug-resources:
         description: "Whether to run in debug mode (true/false)"
         required: true
@@ -45,19 +53,22 @@ on:  # yamllint disable-line rule:truthy
         required: true
         type: string
 jobs:
-  generate-constraints:
+  generate-constraints-matrix:
     permissions:
       contents: read
     timeout-minutes: 70
-    name: Generate constraints ${{ inputs.python-versions-list-as-string }}
+    name: Generate constraints for ${{ inputs.python-versions-list-as-string }}
     runs-on: ${{ fromJSON(inputs.runs-on-as-json-public) }}
+    strategy:
+      matrix:
+        python-version: ${{ fromJson(inputs.python-versions) }}
     env:
       DEBUG_RESOURCES: ${{ inputs.debug-resources }}
       GITHUB_REPOSITORY: ${{ github.repository }}
       GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
       GITHUB_USERNAME: ${{ github.actor }}
       INCLUDE_SUCCESS_OUTPUTS: "true"
-      PYTHON_VERSIONS: ${{ inputs.python-versions-list-as-string }}
+      PYTHON_VERSION: ${{ matrix.python-version }}
       VERBOSE: "true"
       VERSION_SUFFIX_FOR_PYPI: "dev0"
     steps:
@@ -68,70 +79,56 @@ jobs:
         uses: actions/checkout@v4
         with:
           persist-credentials: false
-      - name: "Cleanup docker"
-        run: ./scripts/ci/cleanup_docker.sh
-        shell: bash
-      - name: "Install Breeze"
-        uses: ./.github/actions/breeze
-        with:
-          use-uv: ${{ inputs.use-uv }}
-        id: breeze
-      - name: "Prepare all CI images: ${{ inputs.python-versions-list-as-string}}"
-        uses: ./.github/actions/prepare_all_ci_images
+      - name: "Prepare breeze & CI image: ${{ matrix.python-version }}"
+        uses: ./.github/actions/prepare_breeze_and_image
         with:
           platform: "linux/amd64"
-          python-versions-list-as-string: ${{ inputs.python-versions-list-as-string }}
-          docker-volume-location: ""  # TODO(jscheffl): Understand why it fails here and fix it
-      - name: "Verify all CI images ${{ inputs.python-versions-list-as-string }}"
-        run: breeze ci-image verify --run-in-parallel
+          python: ${{ matrix.python-version }}
+          use-uv: ${{ inputs.use-uv }}
       - name: "Source constraints"
         shell: bash
         run: >
-          breeze release-management generate-constraints --run-in-parallel
+          breeze release-management generate-constraints
           --airflow-constraints-mode constraints-source-providers --answer yes
+          --python "${PYTHON_VERSION}"
       - name: "No providers constraints"
         shell: bash
         timeout-minutes: 25
         run: >
-          breeze release-management generate-constraints --run-in-parallel
-          --airflow-constraints-mode constraints-no-providers --answer yes --parallelism 3
-        # The no providers constraints are only needed when we want to update constraints (in canary builds)
-        # They slow down the start of PROD image builds so we want to only run them when needed.
+          breeze release-management generate-constraints
+          --airflow-constraints-mode constraints-no-providers --answer yes
+          --python "${PYTHON_VERSION}"
         if: inputs.generate-no-providers-constraints == 'true'
-      - name: "Prepare chicken-eggs provider packages"
-        # In case of provider packages which use latest dev0 version of providers, we should prepare them
-        # from the source code, not from the PyPI because they have apache-airflow>=X.Y.Z dependency
-        # And when we prepare them from sources they will have apache-airflow>=X.Y.Z.dev0
+      - name: "Prepare updated provider distributions"
         shell: bash
-        env:
-          CHICKEN_EGG_PROVIDERS: ${{ inputs.chicken-egg-providers }}
         run: >
-          breeze release-management prepare-provider-packages --include-not-ready-providers
-          --package-format wheel --version-suffix-for-pypi dev0
-          ${CHICKEN_EGG_PROVIDERS}
-        if: inputs.chicken-egg-providers != ''
-      - name: "PyPI constraints"
+          breeze release-management prepare-provider-packages
+          --include-not-ready-providers --package-format wheel
+        if: inputs.generate-pypi-constraints == 'true'
+      - name: "Prepare airflow distributions"
         shell: bash
-        timeout-minutes: 25
-        env:
-          CHICKEN_EGG_PROVIDERS: ${{ inputs.chicken-egg-providers }}
         run: >
-          breeze release-management generate-constraints --run-in-parallel
-          --airflow-constraints-mode constraints --answer yes
-          --chicken-egg-providers "${CHICKEN_EGG_PROVIDERS}" --parallelism 3
-      - name: "Dependency upgrade summary"
+          breeze release-management prepare-airflow-packages --distribution-format wheel
+        if: inputs.generate-pypi-constraints == 'true'
+      - name: "PyPI constraints"
         shell: bash
-        env:
-          PYTHON_VERSIONS: ${{ env.PYTHON_VERSIONS }}
+        timeout-minutes: 25
         run: |
-          for PYTHON_VERSION in $PYTHON_VERSIONS; do
-            echo "Summarizing Python $PYTHON_VERSION"
-            cat "files/constraints-${PYTHON_VERSION}"/*.md >> $GITHUB_STEP_SUMMARY || true
-          done
+          breeze release-management generate-constraints --airflow-constraints-mode constraints \
+            --answer yes --python "${PYTHON_VERSION}"
+        if: inputs.generate-pypi-constraints == 'true'
       - name: "Upload constraint artifacts"
-        uses: actions/upload-artifact@v4
+        uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02  # v4.6.2
         with:
-          name: constraints
-          path: ./files/constraints-*/constraints-*.txt
+          name: constraints-${{ matrix.python-version }}
+          path: ./files/constraints-${{ matrix.python-version }}/constraints-*.txt
           retention-days: 7
           if-no-files-found: error
+      - name: "Dependency upgrade summary"
+        shell: bash
+        env:
+          PYTHON_VERSION: ${{ matrix.python-version }}
+        run: |
+          echo "Summarizing Python $PYTHON_VERSION"
+          cat "files/constraints-${PYTHON_VERSION}"/*.md >> $GITHUB_STEP_SUMMARY || true
+          df -H

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

@@ -233,7 +233,7 @@ jobs:
       - name: "Download constraints"
         uses: actions/download-artifact@v4
         with:
-          name: constraints
+          pattern: constraints-*
           path: ./docker-context-files
       - name: "Login to ghcr.io"
         env:

airflow/www/static/js/types/api-generated.ts

@@ -2532,11 +2532,8 @@ export interface components {
       | "dummy"
       | "all_skipped"
       | "always";
-    /**
-     * @description Weight rule.
-     * @enum {string}
-     */
-    WeightRule: "downstream" | "upstream" | "absolute";
+    /** @description Weight rule. One of 'downstream', 'upstream', 'absolute', or the path of the custom priority weight strategy class. */
+    WeightRule: string;
     /**
      * @description Health status
      * @enum {string|null}

dev/breeze/src/airflow_breeze/commands/release_management_commands.py

@@ -136,7 +136,6 @@
     run_with_pool,
 )
 from airflow_breeze.utils.path_utils import (
-    AIRFLOW_PROVIDERS_DIR,
     AIRFLOW_SOURCES_ROOT,
     CONSTRAINTS_CACHE_DIR,
     DIST_DIR,
@@ -871,43 +870,23 @@ def prepare_provider_packages(
             get_console().print()
             with ci_group(f"Preparing provider package [special]{provider_id}"):
                 get_console().print()
-                new_provider_root_dir = AIRFLOW_PROVIDERS_DIR.joinpath(*provider_id.split("."))
-                if (new_provider_root_dir / "provider.yaml").exists():
-                    get_console().print(
-                        f"[info]Provider {provider_id} is a new-style provider. "
-                        f"Skipping package generation as it is not needed for new-style providers."
-                    )
-                    cleanup_build_remnants(new_provider_root_dir)
-                    build_provider_package(
-                        provider_id=provider_id,
-                        package_format=package_format,
-                        target_provider_root_sources_path=new_provider_root_dir,
-                    )
-                    move_built_packages_and_cleanup(
-                        new_provider_root_dir,
-                        DIST_DIR,
-                        skip_cleanup=skip_deleting_generated_files,
-                        delete_only_build_and_dist_folders=True,
-                    )
-                else:
-                    # TODO(potiuk) - remove this once all providers are new-style
-                    target_provider_root_sources_path = copy_provider_sources_to_target(provider_id)
-                    generate_build_files(
-                        provider_id=provider_id,
-                        version_suffix=package_version,
-                        target_provider_root_sources_path=target_provider_root_sources_path,
-                    )
-                    cleanup_build_remnants(target_provider_root_sources_path)
-                    build_provider_package(
-                        provider_id=provider_id,
-                        package_format=package_format,
-                        target_provider_root_sources_path=target_provider_root_sources_path,
-                    )
-                    move_built_packages_and_cleanup(
-                        target_provider_root_sources_path,
-                        DIST_DIR,
-                        skip_cleanup=skip_deleting_generated_files,
-                    )
+                target_provider_root_sources_path = copy_provider_sources_to_target(provider_id)
+                generate_build_files(
+                    provider_id=provider_id,
+                    version_suffix=package_version,
+                    target_provider_root_sources_path=target_provider_root_sources_path,
+                )
+                cleanup_build_remnants(target_provider_root_sources_path)
+                build_provider_package(
+                    provider_id=provider_id,
+                    package_format=package_format,
+                    target_provider_root_sources_path=target_provider_root_sources_path,
+                )
+                move_built_packages_and_cleanup(
+                    target_provider_root_sources_path,
+                    DIST_DIR,
+                    skip_cleanup=skip_deleting_generated_files,
+                )
         except PrepareReleasePackageTagExistException:
             skipped_as_already_released_packages.append(provider_id)
         except PrepareReleasePackageWrongSetupException:

dev/breeze/src/airflow_breeze/prepare_providers/provider_packages.py

@@ -35,7 +35,7 @@
     render_template,
     tag_exists_for_provider,
 )
-from airflow_breeze.utils.path_utils import AIRFLOW_PROVIDERS_DIR, AIRFLOW_SOURCES_ROOT
+from airflow_breeze.utils.path_utils import AIRFLOW_SOURCES_ROOT
 from airflow_breeze.utils.run_utils import run_command
 from airflow_breeze.utils.version_utils import is_local_package_version
 
@@ -84,7 +84,7 @@ def copy_provider_sources_to_target(provider_id: str) -> Path:
     source_provider_sources_path = Path(AIRFLOW_SOURCES_ROOT / "airflow" / "providers").joinpath(
         *provider_id.split(".")
     )
-    relative_provider_path = source_provider_sources_path.relative_to(AIRFLOW_PROVIDERS_DIR)
+    relative_provider_path = source_provider_sources_path.relative_to(AIRFLOW_SOURCES_ROOT)
     target_providers_sub_folder = target_provider_root_path / relative_provider_path
     get_console().print(
         f"[info]Copying provider sources: {source_provider_sources_path} -> {target_providers_sub_folder}"

dev/breeze/src/airflow_breeze/utils/selective_checks.py

@@ -1120,6 +1120,7 @@ def skip_pre_commits(self) -> str:
                     "check-provider-yaml-valid",
                     "lint-helm-chart",
                     "validate-operators-init",
+                    "kubeconform",
                 )
             )
 

dev/breeze/tests/test_selective_checks.py

@@ -788,7 +788,7 @@ def test_full_test_needed_when_scripts_changes(files: tuple[str, ...], expected_
                     "docs-build": "true",
                     "docs-list-as-string": "apache-airflow docker-stack",
                     "full-tests-needed": "true",
-                    "skip-pre-commits": "check-airflow-provider-compatibility,check-extra-packages-references,check-provider-yaml-valid,identity,lint-helm-chart,mypy-airflow,mypy-dev,mypy-docs,mypy-providers,validate-operators-init",
+                    "skip-pre-commits": "check-airflow-provider-compatibility,check-extra-packages-references,check-provider-yaml-valid,identity,kubeconform,lint-helm-chart,mypy-airflow,mypy-dev,mypy-docs,mypy-providers,validate-operators-init",
                     "upgrade-to-newer-dependencies": "false",
                     "core-test-types-list-as-string": "API Always CLI Core Operators Other "
                     "Serialization WWW",
@@ -972,7 +972,7 @@ def test_expected_output_pull_request_v2_7(
                 "needs-helm-tests": "false",
                 "run-tests": "true",
                 "docs-build": "true",
-                "skip-pre-commits": "check-airflow-provider-compatibility,check-extra-packages-references,check-provider-yaml-valid,identity,lint-helm-chart,mypy-airflow,mypy-dev,mypy-docs,mypy-providers,validate-operators-init",
+                "skip-pre-commits": "check-airflow-provider-compatibility,check-extra-packages-references,check-provider-yaml-valid,identity,kubeconform,lint-helm-chart,mypy-airflow,mypy-dev,mypy-docs,mypy-providers,validate-operators-init",
                 "docs-list-as-string": "apache-airflow docker-stack",
                 "upgrade-to-newer-dependencies": "true",
                 "core-test-types-list-as-string": "API Always CLI Core Operators Other Serialization WWW",