Run the beta integration suite on seven pytest-xdist workers

.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"

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.

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):

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

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."""