Skip to content

Optimize startup time for parallel tests #40388

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
potiuk merged 1 commit into from
Jun 23, 2024
Merged

Optimize startup time for parallel tests #40388

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 16 additions & 6 deletions dev/breeze/src/airflow_breeze/commands/testing_commands.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Comment thread
potiuk marked this conversation as resolved.
Outdated
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 \\[\\][/]"
Expand Down Expand Up @@ -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,
Expand All @@ -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,
Expand All @@ -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,
Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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()
Expand All @@ -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,
Expand Down
6 changes: 4 additions & 2 deletions dev/breeze/src/airflow_breeze/params/shell_params.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down