Skip to content

Run the beta integration suite on seven pytest-xdist workers - #983

Merged
anth-volk merged 2 commits into
mainfrom
beta-integration-test-workers
Oct 29, 2026
Merged

anth-volk merged 2 commits into
mainfrom
beta-integration-test-workers

Conversation

@anth-volk

Copy link
Copy Markdown
Contributor

Fixes #980

Same-Repository Draft PR Check

  • This PR is a draft.
  • This PR is from a branch in PolicyEngine/policyengine-sim-api, not a personal fork.

Changelog

  • This PR adds a changelog fragment, changelog.d/<name>.<type>.md, whose type is breaking, added, changed, fixed or removed. (changelog.d/980.changed.md)

Summary

The beta deploy's integration step ran the complete suite (35 tests) in one pytest process: 17.4 minutes median over 16 deploys (14.5 to 37.2). Six beta_only tests submit calculations to beta and take almost all of that time; the longest, test_calculate_default_model, takes about 4.9 minutes. This PR runs the beta suite on pytest-xdist workers so the six calculations run at the same time.

  • .github/scripts/simulation-run-integ-tests.sh, beta branch: uv run --frozen pytest tests/simulation/ -v -n 7 --dist loadgroup --durations=0 --durations-min=1.0. The worker count is BETA_PYTEST_WORKERS=7. pytest-xdist 3.8.0 is already in the test extra and uv.lock of projects/policyengine-apis-integ, which the script installs with uv sync --extra test --frozen.
  • projects/policyengine-apis-integ/tests/simulation/conftest.py: a pytest_collection_modifyitems hook puts every test in that directory without beta_only into one xdist_group, short_tests. Without --dist loadgroup (PR CI, the prod branch) the group has no effect.
  • The prod branch of the script is unchanged and stays in one process. No workflow runs it (rule sim.contracts.full-integration-beta-only: prod passes run_full_integration: false), and it runs no calculation tests.
  • The workflow files are unchanged.

Worker count and distribution mode

-n 7 --dist loadgroup, with the 29 short tests in one group and each of the six calculation tests a work unit of its own. That is seven work units, so xdist gives each of the seven workers exactly one unit and then tells each worker to stop after it. Every calculation starts at once on its own worker, and none waits behind another. The 29 short tests run one after another on the seventh worker in seconds.

Why not plain -n 8 --dist load (the starting point), worksteal, or loadgroup without groups: xdist keeps a second test queued behind the running one on each worker, and that queued test cannot move to an idle worker (worksteal never takes a worker's running test or the one queued next). So two calculation tests can end up on one worker, one waiting minutes for the other. In a local trial that replaced each test with a sleep of its measured length (in seconds instead of minutes), in the suite's collection order, two calculation tests shared a worker in this many of 10 runs at eight workers: load 3, worksteal 10, loadgroup without groups 0 (and 1 of 4 in an earlier set). At six workers every mode did it in every run (4 of 4 each). With the short_tests group, each calculation ran on its own worker in all 10 runs (5 at seven workers, 5 at eight), and a run took 6.0 to 6.4 seconds at seven workers against 17.8 serially.

More than seven workers cannot shorten the step: there are only six calculation tests, and an eighth worker would have no work unit (xdist stops it at once). Fewer workers would queue a calculation behind another. The number of calculations that run at once on beta Modal is the number of calculation tests, six, for any count of seven or more.

test_the_beta_run_has_one_worker_per_work_unit fails if a seventh beta_only test is added without raising BETA_PYTEST_WORKERS, so that change, and its extra load on beta, is a decision someone makes.

Are the tests safe to run in parallel?

  • Fixtures: every fixture in tests/simulation/conftest.py is function-scoped; there are no session- or module-scoped fixtures. settings = Settings() is built at import from environment variables in each worker process and only read.
  • Shared state and files: the tests write no files. The helper tests patch endpoint_calls.job_status and time.sleep with monkeypatch, which is per test and per process; each xdist worker is a separate process.
  • Fixed IDs and names: the auth smoke tests only GET the fixed job IDs auth-smoke-probe-no-token and auth-smoke-probe-does-not-exist; nothing creates them. The calculation tests use job IDs the gateway assigns (Modal function call IDs). The stores they touch (Modal Dicts simulation-api-job-metadata, simulation-api-job-progress, and the budget-window job and seed Dicts) are keyed by those IDs, and Stage 12 shadow evaluation IDs are derived from them. test_calculate_us_state_region_model and test_calculate_specific_model send identical requests; nothing deduplicates or caches submissions by content, so they run as two independent jobs.
  • Ordering: no test reads another test's result or relies on running before or after it.
  • Tests that must not run at the same time: none found. The only candidate was the Hugging Face download below.

Load on beta Modal

Containers per calculation, from the Modal app definitions (projects/policyengine-simulation-executor/src/modal/app.py):

Test Containers
test_calculate_default_model (US national, segmented) 1 run_simulation + 20 run_simulation_segment
test_calculate_us_state_region_model, test_calculate_specific_model (state/ut) 1 run_simulation each
test_calculate_uk_model, test_calculate_uk_local_area_breakdown (UK is never segmented) 1 run_simulation each
test_budget_window_multi_year_batch_completes 1 run_budget_window_batch + up to 2 run_simulation children (max_parallel: 2)

At most 7 of run_simulation's max_containers=100, 20 of run_simulation_segment's 300, and 1 of run_budget_window_batch's 100 run at once, so no function limit in the repository would queue them. That is about 28 containers (about 217 vCPU and 870 GiB requested) at the peak, against 21 (168 vCPU) for the default-model test alone in the serial run. Workspace-level Modal limits are not in the repository.

  • Gateway (policyengine-simulation-gateway, web_app): no max_containers and no @modal.concurrent, so each container serves one request at a time and Modal adds containers. The tests make at most seven requests at once (one per worker), polling every 5 seconds.
  • Cloud Run entrypoint (beta): concurrency 80 per instance, at most 2 instances, 60-second request timeout. Entrypoint to gateway: 5-second connect and 25-second read timeouts (the defaults; the deploy does not set OLD_GATEWAY_*_TIMEOUT_SECONDS).
  • Stage 12 shadow comparisons, if the beta environment's STAGE12_ENABLED turns them on (its value was not read): the dispatch skips requests with breakdown and those without time_period, so at most the two state/ut requests qualify. Their v2 functions allow 10 containers each. A shadow run never changes the v1 response the tests check.
  • Time limits: each calculation test polls for up to 25 minutes (timeout_in_millis = 1_500_000; the HTTP client timeout is the same), and each Modal function's timeout is 3600 seconds. The integration job sets no timeout-minutes. A test that waits for containers waits within its own 25-minute limit, about five times the longest measured test. A submission is not retried: a 5xx on POST /simulate/economy/comparison fails its test, while status polls retry 502, 503 and 504.

Hugging Face downloads of microcosm_uk_2024_25_dense.h5

Parallel workers do not make more containers download the file:

  • The image build's policyengine bundle install downloads only each country's default dataset (microcosm_uk_2024_25 for the UK), so the default UK dataset is in the image.
  • The dense file is downloaded during a request, by ensure_datasets into /tmp/policyengine-alternate-data/<hash>/ of the container, only for a local-area breakdown or a constituency or local-authority region (the bundle's region_datasets).
  • In this suite only test_calculate_uk_local_area_breakdown sends such a request. It is one request, which runs in one run_simulation container (UK requests are never segmented). test_calculate_uk_model is UK-wide and uses the default dataset. Stage 12 shadow dispatch skips breakdown requests. US state datasets for 2025 to 2027 come from the image.

So each deploy still downloads the file in one container, as before. What changes is when: at the start of the step, while the other calculations' containers start, instead of about 12 minutes in. The 2026-10-27 429s came from several runs downloading the file at once, and policyengine.py's download calls raise_for_status() with no retry; this PR does not change either. Because the count stays at one, the PR does not group the UK tests. test_one_test_downloads_the_uk_local_area_dataset keeps that true: it fails if a second test requests a local-area breakdown or a constituency or local-authority region, and its docstring says to put such tests in one xdist_group.

Logs

Locally, with a failing test under seven workers, the output keeps: one [gwN] [ NN%] PASSED|FAILED <test> entry per test (in completion order, with the worker), the start of each test, a FAILURES section whose header names the worker, a slowest durations section listing each test phase of a second or more, and the short test summary. The short tests' node IDs carry an @short_tests suffix.

Testing

Red and green, re-run at both commits:

  • Red, 9643de72 (tests only): TestSimulationRunIntegTests 1 failed, 9 passed (the beta invocation); tests/simulation/test_parallel_beta_run.py 2 failed, 1 passed (the group check and the worker-count check fail; the Hugging Face check pins an existing condition and passes).
  • Green, 0db4f26f: TestSimulationRunIntegTests 10 passed; test_parallel_beta_run.py 3 passed.

New tests:

  • projects/policyengine-simulation-executor/tests/test_modal_scripts.py: each environment's exact pytest invocation (beta with workers, prod without), and that the minted access token and every simulation_integ_test_* setting reach the pytest process as exported variables (fake uv and curl first on PATH).
  • projects/policyengine-apis-integ/tests/simulation/test_parallel_beta_run.py: every collected test in the directory without beta_only, and only those, carries the short_tests group; BETA_PYTEST_WORKERS equals the number of beta_only tests plus one; only one test requests the UK local-area dataset.

Checks run locally:

  • The suite's quick tests (no deployed service: -m "not beta_only", without test_ping): serially 29 passed, 2 skipped (the auth smoke, which needs gateway auth); with -n 7 --dist loadgroup the same 29 passed, 2 skipped, all on one worker. With nothing listening at the base URL, the full suite under -n 7 --dist loadgroup ran each calculation test on its own worker (gw1 to gw6) and every other test on gw0.
  • An environment probe under -n 4: every worker saw all five exported simulation_integ_test_* variables.
  • projects/policyengine-simulation-executor: 1907 passed, 22 skipped. projects/policyengine-simulation-entry: 840 passed.
  • ruff 0.15.22 (the pr.yml pin): ruff format --check projects/*/src libs/*/src and on the changed test files, ruff check .: passed.
  • python scripts/run_quality_guards.py: every guard passed. python -m scripts.rules guard --base-ref origin/main: passed. python -m unittest discover -s scripts/rules/tests -t .: 263 tests OK. python -m scripts.rules render: no diff.
  • python -m scripts.rules pyright projects/policyengine-apis-integ projects/policyengine-simulation-executor: no errors.
  • .github/scripts/check-uv-locks.sh with uv 0.13.6: passed. shellcheck on the changed script: passed.
  • I did not run the beta integration tests against beta.

Environment variables

The changed step is "Run complete beta integration tests through all services" in the integration job of simulation-deploy.reusable.yml. The job declares environment: ${{ inputs.release_environment }}, which simulation-deploy.yml's beta job sets to beta, and that caller passes secrets: inherit. The workflow wiring is unchanged by this PR. Names were confirmed with read-only gh api calls; no values were read.

Variable Read by Provider How it reaches the step Confirmed
SIMULATION_TEST_AUTH_REQUIRED the script workflow env: literal "1" step env: workflow file
SIMULATION_TEST_AUTH_ISSUER the script (token URL) beta environment secret SIMULATION_ENTRYPOINT_AUTH_ISSUER step env: from secrets.*; the job's environment: beta listed in environments/beta/secrets; not a repository or organization secret
SIMULATION_TEST_AUTH_AUDIENCE the script beta environment secret SIMULATION_ENTRYPOINT_AUTH_AUDIENCE as above as above
SIMULATION_TEST_AUTH_CLIENT_ID the script beta environment secret SIMULATION_ENTRYPOINT_TEST_AUTH_CLIENT_ID as above as above
SIMULATION_TEST_AUTH_CLIENT_SECRET the script beta environment secret SIMULATION_ENTRYPOINT_TEST_AUTH_CLIENT_SECRET as above as above
arguments 1 to 4 the script inputs.release_environment (beta), needs.deploy_entrypoint.outputs.candidate_url, needs.prepare.outputs.us_version and uk_version run: arguments workflow file
simulation_integ_test_base_url, _gateway_auth_required, _access_token, _us_model_version, _uk_model_version the tests (Settings) exported by the script from the values above (the token is minted from the four secrets) the pytest process environment, which each xdist worker inherits as a child process the new script test, and a local probe under -n 4
simulation_integ_test_timeout_in_millis, _poll_interval_seconds the tests not set; the Settings defaults apply none conftest.py
PATH (uv, python3, curl) the script the runner, plus the job's setup-python and setup-uv steps job steps workflow file

pytest-xdist itself sets PYTEST_XDIST_WORKER and related variables in each worker; the tests do not read them. Nothing the tests read is set only in the controller process (for example by a pytest option or a fixture in the controller), so the workers see every value.

Unverified until the first deploy after merge

  • The step's actual duration. Estimate: about 5 to 6 minutes against 17.4 today, saving about 9 to 12 minutes per deploy; more if beta Modal starts containers more slowly with about 28 starting at once.
  • Beta Modal behaviour under that load: workspace-level limits (not in the repository), container start times, and whether the segmented default-model run respawns slow segments under contention.
  • The gateway during the opening burst of six submissions: a gateway container that takes more than the entrypoint's 25-second read timeout to start would fail a submission, which the tests do not retry.
  • Whether the step's Hugging Face 429 rate changes. The dense file is still downloaded once per deploy, but now alongside the other containers' start-up; whether those containers make their own Hugging Face requests on start is not established.
  • That the log output reads as it did locally, on GitHub's log viewer.
  • That the secret values are valid (only their names were checked; deploy Increase timeout for Simulation API Cloud Run service #109 on 2026-10-29 passed with the same wiring).

🤖 Generated with Claude Code

anth-volk and others added 2 commits October 29, 2026 23:39
…kers

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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
@anth-volk
anth-volk marked this pull request as ready for review October 29, 2026 20:00
@anth-volk
anth-volk merged commit 6612848 into main Oct 29, 2026
23 checks passed
@anth-volk
anth-volk deleted the beta-integration-test-workers branch October 29, 2026 20:04
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Run the beta integration suite on pytest-xdist workers during deploys

1 participant