Prepare Build Cache in Breeze2 (#22344)

GitOrigin-RevId: dc75f5d8768c8a42df29c86beb519de282539e1f

dev/breeze/src/airflow_breeze/breeze.py

@@ -181,6 +181,12 @@ def main():
     help='The basic apt runtime dependencies to use when building the images.',
     envvar='RUNTIME_APT_DEPS',
 )
+option_ci_flag = click.option(
+    '--ci',
+    help='Enabling this option will off the pip progress bar',
+    is_flag=True,
+    envvar='CI',
+)
 
 
 @main.command()
@@ -282,6 +288,8 @@ def shell(
 @click.option('--build-cache', help='Cache option')
 @option_platform
 @option_debian_version
+@click.option('--prepare-buildx-cache', is_flag=True)
+@option_ci_flag
 @option_upgrade_to_newer_dependencies
 def build_ci_image(
     verbose: bool,
@@ -302,6 +310,8 @@ def build_ci_image(
     build_cache: Optional[str],
     platform: Optional[str],
     debian_version: Optional[str],
+    prepare_buildx_cache: bool,
+    ci: bool,
     upgrade_to_newer_dependencies: str = "false",
 ):
     """Builds docker CI image without entering the container."""
@@ -331,6 +341,8 @@ def build_ci_image(
         docker_cache=build_cache,
         platform=platform,
         debian_version=debian_version,
+        prepare_buildx_cache=prepare_buildx_cache,
+        ci=ci,
         upgrade_to_newer_dependencies=upgrade_to_newer_dependencies,
     )
 
@@ -380,6 +392,7 @@ def build_ci_image(
 )
 @option_image_tag
 @click.option('--github-token', envvar='GITHUB_TOKEN')
+@option_ci_flag
 def build_prod_image(
     verbose: bool,
     cleanup_docker_context_files: bool,
@@ -417,6 +430,7 @@ def build_prod_image(
     install_from_docker_context_files: bool,
     image_tag: Optional[str],
     github_token: Optional[str],
+    ci: bool,
     upgrade_to_newer_dependencies: str = "false",
 ):
     """Builds docker Production image without entering the container."""
@@ -463,6 +477,7 @@ def build_prod_image(
         install_docker_context_files=install_from_docker_context_files,
         image_tag=image_tag,
         github_token=github_token,
+        ci=ci,
     )
 
 

dev/breeze/src/airflow_breeze/ci/build_image.py

@@ -72,8 +72,13 @@ def construct_arguments_docker_command(ci_image: BuildParams) -> List[str]:
 
 def construct_docker_command(ci_image: BuildParams) -> List[str]:
     arguments = construct_arguments_docker_command(ci_image)
+    build_command = ci_image.check_buildx_plugin_build_command()
+    build_flags = ci_image.extra_docker_ci_flags
     final_command = []
-    final_command.extend(["docker", "buildx", "build", "--builder", "default", "--progress=tty", "--pull"])
+    final_command.extend(["docker"])
+    final_command.extend(build_command)
+    final_command.extend(build_flags)
+    final_command.extend(["--pull"])
     final_command.extend(arguments)
     final_command.extend(["-t", ci_image.airflow_ci_image_name, "--target", "main", "."])
     final_command.extend(["-f", 'Dockerfile.ci'])

dev/breeze/src/airflow_breeze/ci/build_params.py

@@ -15,12 +15,15 @@
 # specific language governing permissions and limitations
 # under the License.
 import os
+import sys
 from dataclasses import dataclass
 from datetime import datetime
 from typing import List, Optional
 
 from airflow_breeze.branch_defaults import AIRFLOW_BRANCH, DEFAULT_AIRFLOW_CONSTRAINTS_BRANCH
+from airflow_breeze.console import console
 from airflow_breeze.global_constants import get_airflow_version
+from airflow_breeze.utils.docker_command_utils import check_if_buildx_plugin_available
 from airflow_breeze.utils.run_utils import run_command
 
 
@@ -59,6 +62,8 @@ class BuildParams:
     additional_runtime_apt_env: str = ""
     platform: str = f"linux/{os.uname().machine}"
     debian_version: str = "bullseye"
+    prepare_buildx_cache: bool = False
+    ci: bool = False
 
     @property
     def airflow_image_name(self):
@@ -106,18 +111,64 @@ def commit_sha(self):
         output = run_command(['git', 'rev-parse', 'HEAD'], capture_output=True, text=True)
         return output.stdout.strip()
 
+    @property
+    def airflow_version(self):
+        return get_airflow_version()
+
+    def check_buildx_plugin_build_command(self):
+        build_command_param = []
+        is_buildx_available = check_if_buildx_plugin_available(True)
+        if is_buildx_available:
+            if self.prepare_buildx_cache:
+                build_command_param.extend(
+                    ["buildx", "build", "--builder", "airflow_cache", "--progress=tty"]
+                )
+                cmd = ['docker', 'buildx', 'inspect', 'airflow_cache']
+                output = run_command(cmd, verbose=True, text=True)
+                if output.returncode != 0:
+                    next_cmd = ['docker', 'buildx', 'create', '--name', 'airflow_cache']
+                    run_command(next_cmd, verbose=True, text=True)
+            else:
+                build_command_param.extend(["buildx", "build", "--builder", "default", "--progress=tty"])
+        else:
+            if self.prepare_buildx_cache:
+                console.print(
+                    '\n[red] Buildx cli plugin is not available and you need it to prepare buildx cache. \n'
+                )
+                console.print(
+                    '[red] Please install it following https://docs.docker.com/buildx/working-with-buildx/ \n'
+                )
+                sys.exit()
+            build_command_param.append("build")
+        return build_command_param
+
     @property
     def docker_cache_ci_directive(self) -> List:
         docker_cache_ci_directive = []
+
         if self.docker_cache == "pulled":
-            docker_cache_ci_directive.append("--cache-from")
-            docker_cache_ci_directive.append(self.airflow_ci_image_name)
+            docker_cache_ci_directive.append(f"--cache-from={self.airflow_ci_image_name}")
         elif self.docker_cache == "disabled":
             docker_cache_ci_directive.append("--no-cache")
         else:
-            pass
+            docker_cache_ci_directive = []
+
+        if self.prepare_buildx_cache:
+            docker_cache_ci_directive.extend(["--cache-to=type=inline,mode=max", "--push"])
         return docker_cache_ci_directive
 
     @property
-    def airflow_version(self):
-        return get_airflow_version()
+    def extra_docker_ci_flags(self) -> List[str]:
+        extra_ci_flags = []
+        if self.ci:
+            extra_ci_flags.extend(
+                [
+                    "--build-arg",
+                    "PIP_PROGRESS_BAR=off",
+                ]
+            )
+        if self.airflow_constraints_location is not None and len(self.airflow_constraints_location) > 0:
+            extra_ci_flags.extend(
+                ["--build-arg", f"AIRFLOW_CONSTRAINTS_LOCATION={self.airflow_constraints_location}"]
+            )
+        return extra_ci_flags

dev/breeze/src/airflow_breeze/prod/prod_params.py

@@ -47,17 +47,17 @@ class ProdParams:
     install_docker_context_files: bool
     disable_pypi_when_building: bool
     disable_pip_cache: bool
+    skip_installing_airflow_providers_from_sources: bool
+    cleanup_docker_context_files: bool
+    prepare_buildx_cache: bool
     upgrade_to_newer_dependencies: str = "false"
-    skip_installing_airflow_providers_from_sources: bool = False
-    cleanup_docker_context_files: bool = False
-    prepare_buildx_cache: bool = False
     airflow_version: str = get_airflow_version()
     python_version: str = "3.7"
     airflow_branch_for_pypi_preloading: str = AIRFLOW_BRANCH
     install_airflow_reference: str = ""
     install_airflow_version: str = ""
     default_constraints_branch = DEFAULT_AIRFLOW_CONSTRAINTS_BRANCH
-    ci: str = "false"
+    ci: bool = False
     build_id: int = 0
     airflow_constraints: str = "constraints-source-providers"
     github_repository: str = "apache/airflow"