[airflow] 05/22: Make sure builder is created as needed before warming up (#24875)

[po...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

potiuk pushed a commit to branch v2-3-test
in repository https://gitbox.apache.org/repos/asf/airflow.git

commit fa4c440bee4ce67bbccb4c68c6888d00ca1194e5
Author: Jarek Potiuk <ja...@polidea.com>
AuthorDate: Fri Jul 8 01:33:30 2022 +0200

    Make sure builder is created as needed before warming up (#24875)
    
    In case of parallel cache building, we want to warm-up the builder and make
    sure the builder is created. It might or might not be created by
    in case of ARM builds (and then it will point to ARM builder) but in
    case we run AMD builds, we should attempt to check and re-create it if
    needed.
    
    (cherry picked from commit 6c256c8b00f0ad780c59eaea84bf45f57f4e3d4a)
---
 .../src/airflow_breeze/commands/ci_image_commands.py   | 10 ++++++++--
 .../commands/production_image_commands.py              | 10 +++++++---
 .../src/airflow_breeze/utils/docker_command_utils.py   | 18 +++++++++++-------
 3 files changed, 26 insertions(+), 12 deletions(-)

diff --git a/dev/breeze/src/airflow_breeze/commands/ci_image_commands.py b/dev/breeze/src/airflow_breeze/commands/ci_image_commands.py
index dbc6acda2f..bc0eca063f 100644
--- a/dev/breeze/src/airflow_breeze/commands/ci_image_commands.py
+++ b/dev/breeze/src/airflow_breeze/commands/ci_image_commands.py
@@ -74,6 +74,7 @@ from airflow_breeze.utils.confirm import STANDARD_TIMEOUT, Answer, user_confirm
 from airflow_breeze.utils.console import get_console
 from airflow_breeze.utils.docker_command_utils import (
     build_cache,
+    make_sure_builder_configured,
     perform_environment_checks,
     prepare_docker_build_command,
     prepare_docker_build_from_input,
@@ -228,6 +229,11 @@ def run_build_in_parallel(
     pool.close()
 
 
+def start_building(params: BuildCiParams, dry_run: bool, verbose: bool):
+    check_if_image_building_is_needed(params, dry_run=dry_run, verbose=verbose)
+    make_sure_builder_configured(parallel=True, params=params, dry_run=dry_run, verbose=verbose)
+
+
 @main.command(name='build-image')
 @option_github_repository
 @option_verbose
@@ -298,7 +304,7 @@ def build_image(
             params.python = python
             params.answer = answer
             params_list.append(params)
-        check_if_image_building_is_needed(params_list[0], dry_run=dry_run, verbose=verbose)
+        start_building(params_list[0], dry_run, verbose)
         run_build_in_parallel(
             image_params_list=params_list,
             python_version_list=python_version_list,
@@ -308,7 +314,7 @@ def build_image(
         )
     else:
         params = BuildCiParams(**parameters_passed)
-        check_if_image_building_is_needed(params, dry_run=dry_run, verbose=verbose)
+        start_building(params, dry_run, verbose)
         run_build(ci_image_params=params)
 
 
diff --git a/dev/breeze/src/airflow_breeze/commands/production_image_commands.py b/dev/breeze/src/airflow_breeze/commands/production_image_commands.py
index a9442ba639..2ee164645d 100644
--- a/dev/breeze/src/airflow_breeze/commands/production_image_commands.py
+++ b/dev/breeze/src/airflow_breeze/commands/production_image_commands.py
@@ -73,6 +73,7 @@ from airflow_breeze.utils.console import get_console
 from airflow_breeze.utils.custom_param_types import BetterChoice
 from airflow_breeze.utils.docker_command_utils import (
     build_cache,
+    make_sure_builder_configured,
     perform_environment_checks,
     prepare_docker_build_command,
     prepare_docker_build_from_input,
@@ -200,7 +201,10 @@ PRODUCTION_IMAGE_TOOLS_PARAMETERS = {
 }
 
 
-def start_building(prod_image_params: BuildProdParams, dry_run: bool, verbose: bool):
+def start_building(parallel: bool, prod_image_params: BuildProdParams, dry_run: bool, verbose: bool):
+    make_sure_builder_configured(
+        parallel=parallel, params=prod_image_params, dry_run=dry_run, verbose=verbose
+    )
     if prod_image_params.cleanup_context:
         clean_docker_context_files(verbose=verbose, dry_run=dry_run)
     check_docker_context_files(prod_image_params.install_packages_from_context)
@@ -346,7 +350,7 @@ def build_prod_image(
             params.python = python
             params.answer = answer
             params_list.append(params)
-        start_building(prod_image_params=params_list[0], dry_run=dry_run, verbose=verbose)
+        start_building(parallel=True, prod_image_params=params_list[0], dry_run=dry_run, verbose=verbose)
         run_build_in_parallel(
             image_params_list=params_list,
             python_version_list=python_version_list,
@@ -356,7 +360,7 @@ def build_prod_image(
         )
     else:
         params = BuildProdParams(**parameters_passed)
-        start_building(prod_image_params=params, dry_run=dry_run, verbose=verbose)
+        start_building(parallel=False, prod_image_params=params, dry_run=dry_run, verbose=verbose)
         run_build(prod_image_params=params)
 
 
diff --git a/dev/breeze/src/airflow_breeze/utils/docker_command_utils.py b/dev/breeze/src/airflow_breeze/utils/docker_command_utils.py
index d646ba0fab..df540a5bb8 100644
--- a/dev/breeze/src/airflow_breeze/utils/docker_command_utils.py
+++ b/dev/breeze/src/airflow_breeze/utils/docker_command_utils.py
@@ -496,13 +496,6 @@ def build_cache(
     build_command_result: Union[CompletedProcess, CalledProcessError] = CompletedProcess(
         args=[], returncode=0
     )
-    cmd = ['docker', 'buildx', 'inspect', 'airflow_cache']
-    buildx_command_result = run_command(
-        cmd, verbose=verbose, dry_run=dry_run, text=True, check=False, enabled_output_group=not parallel
-    )
-    if buildx_command_result and buildx_command_result.returncode != 0:
-        next_cmd = ['docker', 'buildx', 'create', '--name', 'airflow_cache']
-        run_command(next_cmd, verbose=verbose, text=True, check=False, enabled_output_group=not parallel)
     for platform in image_params.platforms:
         platform_image_params = deepcopy(image_params)
         # override the platform in the copied params to only be single platform per run
@@ -523,6 +516,17 @@ def build_cache(
     return build_command_result
 
 
+def make_sure_builder_configured(parallel: bool, params: CommonBuildParams, dry_run: bool, verbose: bool):
+    if params.builder != 'default':
+        cmd = ['docker', 'buildx', 'inspect', params.builder]
+        buildx_command_result = run_command(
+            cmd, verbose=verbose, dry_run=dry_run, text=True, check=False, enabled_output_group=not parallel
+        )
+        if buildx_command_result and buildx_command_result.returncode != 0:
+            next_cmd = ['docker', 'buildx', 'create', '--name', params.builder]
+            run_command(next_cmd, verbose=verbose, text=True, check=False, enabled_output_group=not parallel)
+
+
 def set_value_to_default_if_not_set(env: Dict[str, str], name: str, default: str):
     """
     Set value of name parameter to default (indexed by name) if not set.