From 9643de72f4ef59ee07091fe547a4c98d072f3ac8 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 29 Oct 2026 23:39:27 +0400 Subject: [PATCH 1/2] Add failing tests for running the beta integration suite on xdist workers Pin the beta pytest invocation in simulation-run-integ-tests.sh: seven pytest-xdist workers with --dist loadgroup and --durations, while the prod branch stays in one process. Check that the settings the integration tests read, including the minted access token, reach the pytest process as exported variables, which pytest-xdist workers inherit. In the integration suite, check that every test without beta_only shares the short_tests xdist group and that each beta_only calculation test is a work unit of its own, that the beta run has one worker per work unit, and that only one test makes a worker download the UK local-area dataset. Refs #980 Co-Authored-By: Claude Opus 5.5 --- .../simulation/test_parallel_beta_run.py | 133 ++++++++++++++++ .../tests/test_modal_scripts.py | 142 ++++++++++++++++++ 2 files changed, 275 insertions(+) create mode 100644 projects/policyengine-apis-integ/tests/simulation/test_parallel_beta_run.py 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.""" From 0db4f26f13136e7a7a9b940cf7a2a842c3269b28 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 29 Oct 2026 23:40:38 +0400 Subject: [PATCH 2/2] Run the beta integration suite on seven pytest-xdist workers The beta deploy's integration step ran 35 tests in one process for about 17 minutes; six beta_only tests submit calculations and take almost all of it, the longest about 4.9 minutes. Run beta with -n 7 --dist loadgroup and --durations=0 --durations-min=1.0. The integration conftest puts every test without beta_only in one xdist_group, short_tests, so the suite is one work unit per calculation test plus one for the short tests. With one worker per unit, xdist gives each worker exactly one unit, so every calculation starts at once and none waits behind another on the same worker. Without groups, --dist load, worksteal and loadgroup each keep a second test queued on a busy worker, and in a scaled local trial they put two calculation tests on one worker in 3, 10 and 0 of 10 runs at eight workers. The prod branch, which no workflow runs, stays in one process. Fixes #980 Co-Authored-By: Claude Opus 5.5 --- .github/scripts/simulation-run-integ-tests.sh | 17 +++++++++++-- changelog.d/980.changed.md | 1 + .../tests/simulation/conftest.py | 24 +++++++++++++++++++ 3 files changed, 40 insertions(+), 2 deletions(-) create mode 100644 changelog.d/980.changed.md 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):