diff --git a/.github/scripts/simulation-run-integ-tests.sh b/.github/scripts/simulation-run-integ-tests.sh index 26cc52cde..a09a4fc49 100755 --- a/.github/scripts/simulation-run-integ-tests.sh +++ b/.github/scripts/simulation-run-integ-tests.sh @@ -120,9 +120,22 @@ if [ -n "$UK_VERSION" ]; then export simulation_integ_test_uk_model_version="$UK_VERSION" fi +# Beta runs the suite on pytest-xdist workers (in the project's test extra). +# Each beta_only test submits a calculation to beta and waits minutes for it; +# every other test takes seconds. tests/simulation/conftest.py puts the other +# tests in one xdist_group, and --dist loadgroup runs a group on one worker, so +# the suite is one work unit per calculation test plus one for the rest. With +# one worker per unit (six calculation tests, plus one), every calculation +# starts at once and none waits behind another on the same worker. The workers +# inherit the simulation_integ_test_* variables exported above. --durations +# lists every test phase that took a second or more. +BETA_PYTEST_WORKERS=7 + if [ "$ENVIRONMENT" = "beta" ]; then - echo "Running all simulation integration tests (including beta_only)" - uv run --frozen pytest tests/simulation/ -v + echo "Running all simulation integration tests (including beta_only) on $BETA_PYTEST_WORKERS pytest-xdist workers" + uv run --frozen pytest tests/simulation/ -v \ + -n "$BETA_PYTEST_WORKERS" --dist loadgroup \ + --durations=0 --durations-min=1.0 else echo "Running simulation integration tests (excluding beta_only)" uv run --frozen pytest tests/simulation/ -v -m "not beta_only" diff --git a/changelog.d/980.changed.md b/changelog.d/980.changed.md new file mode 100644 index 000000000..6018d2e4e --- /dev/null +++ b/changelog.d/980.changed.md @@ -0,0 +1 @@ +The beta deployment now runs its complete integration suite on seven pytest-xdist workers, so the six calculation tests run at the same time instead of one after another. The step should take about as long as its slowest test, about 5 minutes, instead of about 17. diff --git a/projects/policyengine-apis-integ/tests/simulation/conftest.py b/projects/policyengine-apis-integ/tests/simulation/conftest.py index eb29b2991..f04ea7f53 100644 --- a/projects/policyengine-apis-integ/tests/simulation/conftest.py +++ b/projects/policyengine-apis-integ/tests/simulation/conftest.py @@ -2,6 +2,7 @@ import time from dataclasses import dataclass from http import HTTPStatus +from pathlib import Path import httpx import pytest @@ -21,6 +22,29 @@ } BUDGET_WINDOW_REGION = "state/ut" BUDGET_WINDOW_MAX_PARALLEL = 2 +# simulation-run-integ-tests.sh runs the beta suite with pytest-xdist and +# --dist loadgroup, which runs all tests of one xdist_group on one worker. +SHORT_TESTS_XDIST_GROUP = "short_tests" +_SIMULATION_TESTS_DIRECTORY = Path(__file__).parent + + +@pytest.hookimpl(tryfirst=True) +def pytest_collection_modifyitems(items: list[pytest.Item]) -> None: + """Put every test here without ``beta_only`` in one xdist group. + + Each ``beta_only`` test submits a calculation and waits minutes for it, and + stays a work unit of its own; the other tests take seconds. With one worker + per unit, every calculation starts at once on its own worker. This runs + first because pytest-xdist reads the groups in its own + ``pytest_collection_modifyitems``. Without ``--dist loadgroup`` the group + changes nothing. + """ + for item in items: + if ( + item.path.is_relative_to(_SIMULATION_TESTS_DIRECTORY) + and item.get_closest_marker("beta_only") is None + ): + item.add_marker(pytest.mark.xdist_group(SHORT_TESTS_XDIST_GROUP)) class Settings(BaseSettings): diff --git a/projects/policyengine-apis-integ/tests/simulation/test_parallel_beta_run.py b/projects/policyengine-apis-integ/tests/simulation/test_parallel_beta_run.py new file mode 100644 index 000000000..9aa6c7152 --- /dev/null +++ b/projects/policyengine-apis-integ/tests/simulation/test_parallel_beta_run.py @@ -0,0 +1,133 @@ +"""How the beta deploy spreads this suite over pytest-xdist workers. + +``.github/scripts/simulation-run-integ-tests.sh`` runs the beta suite with +``--dist loadgroup``: xdist runs all tests of one ``xdist_group`` on one +worker, and every test without a group is a work unit of its own. Each +``beta_only`` test submits a calculation to beta and waits minutes for it; +every other test here takes seconds. So the other tests share one group, each +calculation test is a unit of its own, and the run has one worker per unit: +every calculation starts at once, and none waits behind another on the same +worker. +""" + +import ast +import re +from pathlib import Path + +import pytest + +SIMULATION_TESTS_DIRECTORY = Path(__file__).parent +REPOSITORY_ROOT = SIMULATION_TESTS_DIRECTORY.parents[3] +RUN_SCRIPT = REPOSITORY_ROOT / ".github" / "scripts" / "simulation-run-integ-tests.sh" +SHORT_TESTS_GROUP = "short_tests" +# The region types policyengine.py's bundle manifest routes to the UK +# local-area dataset (its region_datasets entries for UK constituencies and +# local authorities). +UK_LOCAL_AREA_REGION_PREFIXES = ("constituency/", "local_authority/") + + +def _xdist_group_names(item: pytest.Item) -> list[str]: + return [ + str(mark.args[0] if mark.args else mark.kwargs.get("name", "default")) + for mark in item.iter_markers("xdist_group") + ] + + +def _test_functions() -> list[ast.FunctionDef]: + """Return the module-level test functions of every test module here.""" + test_functions: list[ast.FunctionDef] = [] + for module_path in sorted(SIMULATION_TESTS_DIRECTORY.glob("test_*.py")): + module = ast.parse(module_path.read_text(encoding="utf-8")) + test_functions.extend( + node + for node in module.body + if isinstance(node, ast.FunctionDef) and node.name.startswith("test_") + ) + return test_functions + + +def _is_beta_only(test_function: ast.FunctionDef) -> bool: + return any( + ast.unparse(decorator) == "pytest.mark.beta_only" + for decorator in test_function.decorator_list + ) + + +def _requests_uk_local_area_dataset(test_function: ast.FunctionDef) -> bool: + """Return whether a request dict in the test routes to the local-area dataset.""" + for node in ast.walk(test_function): + if not isinstance(node, ast.Dict): + continue + for key, value in zip(node.keys, node.values, strict=True): + if not ( + isinstance(key, ast.Constant) + and isinstance(value, ast.Constant) + and isinstance(value.value, str) + ): + continue + if key.value == "breakdown" and value.value == "local_areas": + return True + if key.value == "region" and value.value.startswith( + UK_LOCAL_AREA_REGION_PREFIXES + ): + return True + return False + + +def test_short_tests_share_one_group_and_each_calculation_is_its_own_unit( + request: pytest.FixtureRequest, +): + """Every collected test here without beta_only, and only those, is grouped.""" + collected_here = [ + item + for item in request.session.items + if item.path.is_relative_to(SIMULATION_TESTS_DIRECTORY) + ] + assert request.node in collected_here + + for item in collected_here: + expected_groups = ( + [] if item.get_closest_marker("beta_only") else [SHORT_TESTS_GROUP] + ) + assert _xdist_group_names(item) == expected_groups, item.nodeid + + +def test_the_beta_run_has_one_worker_per_work_unit(): + """The worker count is the number of calculation tests plus one.""" + worker_setting = re.search( + r"^BETA_PYTEST_WORKERS=(\d+)$", + RUN_SCRIPT.read_text(encoding="utf-8"), + flags=re.MULTILINE, + ) + assert worker_setting is not None, f"BETA_PYTEST_WORKERS is not set in {RUN_SCRIPT}" + calculation_tests = [ + test_function.name + for test_function in _test_functions() + if _is_beta_only(test_function) + ] + + assert int(worker_setting.group(1)) == len(calculation_tests) + 1, ( + "Give each beta_only calculation test its own worker, plus one worker " + f"for the {SHORT_TESTS_GROUP} group: {calculation_tests}" + ) + + +def test_one_test_downloads_the_uk_local_area_dataset(): + """Only one test makes a worker download the UK local-area dataset. + + The default UK dataset is in the executor's Modal image. A local-area + breakdown, or a constituency or local-authority region, makes its worker + container download ``microcosm_uk_2024_25_dense.h5`` from Hugging Face + during the request, and policyengine.py fails on the first HTTP 429. With + one such test, the parallel beta run downloads the file in one container, + as the serial run did. A second such test would download it in a second + container at the same time: give both the same ``xdist_group`` so they run + one after the other. + """ + downloading_tests = [ + test_function.name + for test_function in _test_functions() + if _requests_uk_local_area_dataset(test_function) + ] + + assert len(downloading_tests) == 1, downloading_tests diff --git a/projects/policyengine-simulation-executor/tests/test_modal_scripts.py b/projects/policyengine-simulation-executor/tests/test_modal_scripts.py index d427800d9..3c14f66da 100644 --- a/projects/policyengine-simulation-executor/tests/test_modal_scripts.py +++ b/projects/policyengine-simulation-executor/tests/test_modal_scripts.py @@ -6,6 +6,7 @@ import os import subprocess +from dataclasses import dataclass from pathlib import Path import pytest @@ -15,6 +16,17 @@ FAKE_HF_TOKEN = "hf_test" +INTEGRATION_TEST_AUTH_VARIABLES = ( + "SIMULATION_TEST_AUTH_REQUIRED", + "SIMULATION_TEST_AUTH_ISSUER", + "SIMULATION_TEST_AUTH_AUDIENCE", + "SIMULATION_TEST_AUTH_CLIENT_ID", + "SIMULATION_TEST_AUTH_CLIENT_SECRET", +) +# The prefix of the variables the integration tests' Settings read. +INTEGRATION_SETTINGS_PREFIX = "simulation_integ_test_" +MINTED_ACCESS_TOKEN = "minted-test-access-token" + def _run_bash( *arguments: str, @@ -32,6 +44,79 @@ def _run_bash( ) +@dataclass(frozen=True) +class IntegrationScriptRun: + """One run of simulation-run-integ-tests.sh against fake uv and curl.""" + + result: subprocess.CompletedProcess[str] + # The arguments of each uv call, in order. + uv_calls: list[str] + # The integration-test settings the pytest process received. + pytest_environment: dict[str, str] + + +def _run_integration_tests_script( + tmp_path: Path, + environment: str, + *, + extra_environment: dict[str, str] | None = None, +) -> IntegrationScriptRun: + """Run the script with fake uv and curl first on PATH, from the repository root. + + The fake uv records its arguments and, for ``uv run``, the + ``simulation_integ_test_*`` variables it received. pytest-xdist workers + are child processes of that pytest process and inherit its environment. + The fake curl answers the token request with ``MINTED_ACCESS_TOKEN``. + """ + uv_calls_log = tmp_path / "uv_calls.log" + pytest_environment_log = tmp_path / "pytest_environment.log" + fake_bin = tmp_path / "bin" + fake_bin.mkdir() + fake_uv = fake_bin / "uv" + fake_uv.write_text( + "#!/bin/bash\n" + 'printf "%s\\n" "$*" >> "$UV_CALLS_LOG"\n' + 'if [ "$1" = "run" ]; then\n' + f" env | grep '^{INTEGRATION_SETTINGS_PREFIX}' | sort" + ' > "$PYTEST_ENVIRONMENT_LOG"\n' + "fi\n" + ) + fake_curl = fake_bin / "curl" + fake_curl.write_text( + f"#!/bin/bash\nprintf '%s' '{{\"access_token\": \"{MINTED_ACCESS_TOKEN}\"}}'\n" + ) + for fake_tool in (fake_uv, fake_curl): + fake_tool.chmod(0o755) + + script_environment = { + name: value + for name, value in os.environ.items() + if name not in INTEGRATION_TEST_AUTH_VARIABLES + and not name.startswith(INTEGRATION_SETTINGS_PREFIX) + } + script_environment["PATH"] = f"{fake_bin}:{script_environment['PATH']}" + script_environment["UV_CALLS_LOG"] = str(uv_calls_log) + script_environment["PYTEST_ENVIRONMENT_LOG"] = str(pytest_environment_log) + script_environment.update(extra_environment or {}) + + result = _run_bash( + str(SCRIPTS_DIR / "simulation-run-integ-tests.sh"), + environment, + "https://candidate.example.test", + "1.690.7", + "2.88.20", + env=script_environment, + cwd=REPO_ROOT, + ) + uv_calls = uv_calls_log.read_text().splitlines() if uv_calls_log.exists() else [] + pytest_environment: dict[str, str] = {} + if pytest_environment_log.exists(): + for setting in pytest_environment_log.read_text().splitlines(): + name, _, value = setting.partition("=") + pytest_environment[name] = value + return IntegrationScriptRun(result, uv_calls, pytest_environment) + + class TestModalExtractVersions: """Tests for modal-extract-versions.sh.""" @@ -1019,6 +1104,63 @@ def test_exports_us_and_uk_model_versions_to_integration_tests(self, tmp_path): assert "us=1.690.7" in log assert "uk=2.88.20" in log + @pytest.mark.parametrize( + ("environment", "pytest_invocation"), + [ + # Seven pytest-xdist workers: one for each of the six beta_only + # calculation tests, and one for the xdist_group that holds every + # other test (tests/simulation/conftest.py). --dist loadgroup keeps + # a group on one worker; --durations lists every test phase that + # took a second or more. + ( + "beta", + "run --frozen pytest tests/simulation/ -v -n 7 --dist loadgroup " + "--durations=0 --durations-min=1.0", + ), + # No workflow runs the prod branch (rule + # sim.contracts.full-integration-beta-only), and it runs no + # calculation tests, so it stays in one process. + ("prod", "run --frozen pytest tests/simulation/ -v -m not beta_only"), + ], + ) + def test_runs_the_environment_s_pytest_invocation( + self, tmp_path, environment, pytest_invocation + ): + """Each environment runs exactly its pytest invocation after the sync.""" + run = _run_integration_tests_script(tmp_path, environment) + + assert run.result.returncode == 0, f"Script failed: {run.result.stderr}" + assert run.uv_calls == ["sync --extra test --frozen", pytest_invocation] + + def test_beta_pytest_receives_the_minted_token_and_deployment_settings( + self, tmp_path + ): + """Every setting the tests read reaches pytest as an exported variable. + + pytest-xdist starts each worker as a child process of this pytest + process, so the workers see the same variables. + """ + run = _run_integration_tests_script( + tmp_path, + "beta", + extra_environment={ + "SIMULATION_TEST_AUTH_REQUIRED": "1", + "SIMULATION_TEST_AUTH_ISSUER": "https://tenant.example.test/", + "SIMULATION_TEST_AUTH_AUDIENCE": "https://simulation.example.test", + "SIMULATION_TEST_AUTH_CLIENT_ID": "test-client-id", + "SIMULATION_TEST_AUTH_CLIENT_SECRET": "test-client-secret", + }, + ) + + assert run.result.returncode == 0, f"Script failed: {run.result.stderr}" + assert run.pytest_environment == { + "simulation_integ_test_access_token": MINTED_ACCESS_TOKEN, + "simulation_integ_test_base_url": "https://candidate.example.test", + "simulation_integ_test_gateway_auth_required": "1", + "simulation_integ_test_uk_model_version": "2.88.20", + "simulation_integ_test_us_model_version": "1.690.7", + } + class TestAllScriptsHaveShebang: """Verify all scripts have proper shebang and error handling."""