From 6860f9d2439c59ef0d03444a11f614dcc2a0ffba Mon Sep 17 00:00:00 2001
From: Jarek Potiuk <jarek@potiuk.com>
Date: Sun, 23 Jun 2024 07:45:37 +0200
Subject: [PATCH] Optimize startup time for parallel tests

When parallell tests start on CI, they are running parallel docker
compose's and on a clean machine in CI this means that every parallel
test is pulling the images needed to run tests. This means that
backend images are pulled in parallell by all starting parallel runs.

This PR optimizes this step - before running the tests in parallel
we run `docker compose pull` once with the same compose files as
tests - this will pull the necessary images only once.

It should save a few seconds and save a lot of unnecessary traffic
for CI tests - where same image is pulled multiple times - especially
for self-hosted runners of ours where we run 8 docker compose instances
in parallel.
---
 .../commands/testing_commands.py              | 22 ++++++++++++++-----
 .../src/airflow_breeze/params/shell_params.py |  6 +++--
 2 files changed, 20 insertions(+), 8 deletions(-)

diff --git a/dev/breeze/src/airflow_breeze/commands/testing_commands.py b/dev/breeze/src/airflow_breeze/commands/testing_commands.py
index bb00cb1110ea7..76e60fa7b23aa 100644
--- a/dev/breeze/src/airflow_breeze/commands/testing_commands.py
+++ b/dev/breeze/src/airflow_breeze/commands/testing_commands.py
@@ -153,14 +153,11 @@ def _run_test(
     shell_params: ShellParams,
     extra_pytest_args: tuple,
     python_version: str,
-    db_reset: bool,
     output: Output | None,
     test_timeout: int,
     output_outside_the_group: bool = False,
     skip_docker_compose_down: bool = False,
 ) -> tuple[int, str]:
-    shell_params.run_tests = True
-    shell_params.db_reset = db_reset
     if "[" in shell_params.test_type and not shell_params.test_type.startswith("Providers"):
         get_console(output=output).print(
             "[error]Only 'Providers' test type can specify actual tests with \\[\\][/]"
@@ -321,7 +318,6 @@ def _run_tests_in_pool(
                         "shell_params": shell_params.clone_with_test(test_type=test_type),
                         "extra_pytest_args": extra_pytest_args,
                         "python_version": shell_params.python,
-                        "db_reset": db_reset,
                         "output": outputs[index],
                         "test_timeout": test_timeout,
                         "skip_docker_compose_down": skip_docker_compose_down,
@@ -341,6 +337,17 @@ def _run_tests_in_pool(
     )
 
 
+def pull_images_for_docker_compose(shell_params: ShellParams):
+    get_console().print("Pulling images once before parallel run\n")
+    env = shell_params.env_variables_for_docker_commands
+    pull_cmd = [
+        "docker",
+        "compose",
+        "pull",
+    ]
+    run_command(pull_cmd, output=None, check=False, env=env)
+
+
 def run_tests_in_parallel(
     shell_params: ShellParams,
     extra_pytest_args: tuple,
@@ -363,6 +370,7 @@ def run_tests_in_parallel(
     get_console().print(f"[info]Skip docker-compose down: {skip_docker_compose_down}")
     get_console().print("[info]Shell params:")
     get_console().print(shell_params.__dict__)
+    pull_images_for_docker_compose(shell_params)
     _run_tests_in_pool(
         tests_to_run=shell_params.parallel_test_types_list,
         parallelism=parallelism,
@@ -752,6 +760,8 @@ def _run_test_command(
         use_airflow_version=use_airflow_version,
         use_packages_from_dist=use_packages_from_dist,
         use_xdist=use_xdist,
+        run_tests=True,
+        db_reset=db_reset,
     )
     rebuild_or_pull_ci_image_if_needed(command_params=shell_params)
     fix_ownership_using_docker()
@@ -798,7 +808,6 @@ def _run_test_command(
             shell_params=shell_params,
             extra_pytest_args=extra_pytest_args,
             python_version=python,
-            db_reset=db_reset,
             output=None,
             test_timeout=test_timeout,
             output_outside_the_group=True,
@@ -866,6 +875,8 @@ def integration_tests(
         skip_provider_tests=skip_provider_tests,
         test_type="Integration",
         force_sa_warnings=force_sa_warnings,
+        run_tests=True,
+        db_reset=db_reset,
     )
     fix_ownership_using_docker()
     cleanup_python_generated_files()
@@ -874,7 +885,6 @@ def integration_tests(
         shell_params=shell_params,
         extra_pytest_args=extra_pytest_args,
         python_version=python,
-        db_reset=db_reset,
         output=None,
         test_timeout=test_timeout,
         output_outside_the_group=True,
diff --git a/dev/breeze/src/airflow_breeze/params/shell_params.py b/dev/breeze/src/airflow_breeze/params/shell_params.py
index dd1bdbbfa5a84..f195bf4c968fd 100644
--- a/dev/breeze/src/airflow_breeze/params/shell_params.py
+++ b/dev/breeze/src/airflow_breeze/params/shell_params.py
@@ -459,11 +459,13 @@ def add_docker_in_docker(self, compose_file_list: list[Path]):
     def rootless_docker(self) -> bool:
         return is_docker_rootless()
 
-    @cached_property
+    @property
     def env_variables_for_docker_commands(self) -> dict[str, str]:
         """
         Constructs environment variables needed by the docker-compose command, based on Shell parameters
-        passed to it.
+        passed to it. We cannot cache this property because it can be run few times after modifying shell
+        params - for example when we first run "pull" on images before tests anda then run tests - each
+        separately with different test types.
 
         This is the only place where you need to add environment variables if you want to pass them to
         docker or docker-compose.