From 2c6180432ddc26ae98a6302b6475d5d5d22b1e7c Mon Sep 17 00:00:00 2001 From: Ziming Date: Wed, 14 Oct 2026 19:47:18 -0400 Subject: [PATCH 1/4] Run Baseline and Contrib groups in sequence, stop Rest on cancel, and split household tests - Rest shard 2 used a job-level always(), which kept it running (or let it start) after a manual or concurrency cancellation. Use !cancelled() so it still runs after a shard-1 failure but stops when the run is cancelled. - Group the 11 Baseline groups onto 5 runners and the 9 Contrib groups onto 4, by name, each runner running its groups as sequential steps. Later groups use !cancelled() so an earlier failure does not hide their results. No job depends on another, so re-running a failed runner re-runs only that runner. PR CI drops from 27 to 16 runners, at most 14 at once; both suites still finish before the Rest shard chain. Release CI uses the same matrices. - Split the household Baseline tests: the root files and income/household run one file per subprocess and other subfolders one subprocess each. The former quarter batches peaked at 15.1 GB on the 16 GB runner; locally measured batches (within 1% of CI) now peak at 7.1 GB. - Update the runner-budget, Rest, make-test and batching guards. Co-Authored-By: Claude Opus 5.5 --- .github/workflows/baseline-tests.yaml | 72 +++++---- .github/workflows/contrib-tests.yaml | 59 +++++--- .github/workflows/push.yaml | 137 +++++++++++------- .github/workflows/python-tests.yaml | 6 +- Makefile | 20 +-- ...equential-suites-household-memory.fixed.md | 1 + .../test_batched_per_subdir_budget.py | 45 +++++- .../code_health/test_make_test_matches_ci.py | 17 ++- .../code_health/test_pr_runner_budget.py | 108 ++++++++++++-- .../code_health/test_rest_python_groups.py | 8 +- policyengine_us/tests/test_batched.py | 34 +++++ 11 files changed, 376 insertions(+), 131 deletions(-) create mode 100644 changelog.d/ci-sequential-suites-household-memory.fixed.md diff --git a/.github/workflows/baseline-tests.yaml b/.github/workflows/baseline-tests.yaml index b0ab85dcdd9..5384aecc38f 100644 --- a/.github/workflows/baseline-tests.yaml +++ b/.github/workflows/baseline-tests.yaml @@ -7,7 +7,10 @@ jobs: Baseline: name: Full Suite - Baseline (${{ matrix.group }}) runs-on: ubuntu-latest - timeout-minutes: 60 + # Two or three groups in sequence: historical medians about 54-66 + # minutes per runner (household + ssa-usda projected at about 60 after + # the household split; about 22 for rest-a + rest-b). + timeout-minutes: 120 env: PYTHONUNBUFFERED: 1 PYTEST_ADDOPTS: "--durations=25 -p no:unraisableexception -p no:threadexception" @@ -15,30 +18,37 @@ jobs: fail-fast: false matrix: include: - # One state per subprocess; retain four runners and one worker each. - # The former NY/OH/OK process peaked at 15.1 GB (run 36930199045). - - group: states-shard-1 - cmd: python policyengine_us/tests/test_batched.py policyengine_us/tests/policy/baseline/gov/states --mode per-subdir --workers 1 --shard 1/4 - - group: states-shard-2 - cmd: python policyengine_us/tests/test_batched.py policyengine_us/tests/policy/baseline/gov/states --mode per-subdir --workers 1 --shard 2/4 - - group: states-shard-3 - cmd: python policyengine_us/tests/test_batched.py policyengine_us/tests/policy/baseline/gov/states --mode per-subdir --workers 1 --shard 3/4 - - group: states-shard-4 - cmd: python policyengine_us/tests/test_batched.py policyengine_us/tests/policy/baseline/gov/states --mode per-subdir --workers 1 --shard 4/4 - - group: irs - cmd: make test-yaml-no-structural-other-irs - - group: household - cmd: make test-yaml-no-structural-other-household - - group: ssa-usda - cmd: make test-yaml-no-structural-other-ssa-usda - - group: rest-a - cmd: make test-yaml-no-structural-other-rest-a - - group: rest-b - cmd: make test-yaml-no-structural-other-rest-b - - group: contrib-hhs - cmd: make test-yaml-contrib-hhs - - group: reform - cmd: make test-yaml-reform + # Five runners, grouped by name, run their groups one after another + # with one worker each. The states shards still run one state per + # subprocess: the former NY/OH/OK process peaked at 15.1 GB (run + # 36930199045). + - group: states-shard-1 + states-shard-2 + group1: states-shard-1 + cmd1: python policyengine_us/tests/test_batched.py policyengine_us/tests/policy/baseline/gov/states --mode per-subdir --workers 1 --shard 1/4 + group2: states-shard-2 + cmd2: python policyengine_us/tests/test_batched.py policyengine_us/tests/policy/baseline/gov/states --mode per-subdir --workers 1 --shard 2/4 + - group: states-shard-3 + states-shard-4 + group1: states-shard-3 + cmd1: python policyengine_us/tests/test_batched.py policyengine_us/tests/policy/baseline/gov/states --mode per-subdir --workers 1 --shard 3/4 + group2: states-shard-4 + cmd2: python policyengine_us/tests/test_batched.py policyengine_us/tests/policy/baseline/gov/states --mode per-subdir --workers 1 --shard 4/4 + - group: household + ssa-usda + group1: household + cmd1: make test-yaml-no-structural-other-household + group2: ssa-usda + cmd2: make test-yaml-no-structural-other-ssa-usda + - group: contrib-hhs + irs + reform + group1: contrib-hhs + cmd1: make test-yaml-contrib-hhs + group2: irs + cmd2: make test-yaml-no-structural-other-irs + group3: reform + cmd3: make test-yaml-reform + - group: rest-a + rest-b + group1: rest-a + cmd1: make test-yaml-no-structural-other-rest-a + group2: rest-b + cmd2: make test-yaml-no-structural-other-rest-b steps: - name: Checkout repo uses: actions/checkout@v6 @@ -57,5 +67,13 @@ jobs: - name: Turn off default branching shell: bash run: bash ./update_itemization.sh - - name: Run Baseline YAML tests (${{ matrix.group }}) - run: uv run --no-sync ${{ matrix.cmd }} + - name: Run Baseline YAML tests (${{ matrix.group1 }}) + run: uv run --no-sync ${{ matrix.cmd1 }} + # Later groups still run after an earlier group fails, but not after + # the run is cancelled. + - name: Run Baseline YAML tests (${{ matrix.group2 }}) + if: ${{ !cancelled() && matrix.cmd2 }} + run: uv run --no-sync ${{ matrix.cmd2 }} + - name: Run Baseline YAML tests (${{ matrix.group3 }}) + if: ${{ !cancelled() && matrix.cmd3 }} + run: uv run --no-sync ${{ matrix.cmd3 }} diff --git a/.github/workflows/contrib-tests.yaml b/.github/workflows/contrib-tests.yaml index cd1c80fcd32..a34b47b53d3 100644 --- a/.github/workflows/contrib-tests.yaml +++ b/.github/workflows/contrib-tests.yaml @@ -7,7 +7,9 @@ jobs: Contrib: name: Full Suite - Contrib (${{ matrix.group }}) runs-on: ubuntu-latest - timeout-minutes: 60 + # Two or three groups in sequence: historical medians 47-74 minutes per + # runner. + timeout-minutes: 120 env: PYTHONUNBUFFERED: 1 PYTEST_ADDOPTS: "--durations=25 -p no:unraisableexception -p no:threadexception" @@ -15,24 +17,29 @@ jobs: fail-fast: false matrix: include: - - group: states-shard-1 - target: test-yaml-structural-heavy-shard-1 - - group: states-shard-2 - target: test-yaml-structural-heavy-shard-2 - - group: states-shard-3 - target: test-yaml-structural-heavy-shard-3 - - group: states-shard-4 - target: test-yaml-structural-heavy-shard-4 - - group: other-shard-1 - target: test-yaml-structural-other - - group: other-shard-2a - target: test-yaml-structural-other-shard-2a - - group: other-shard-2b - target: test-yaml-structural-other-shard-2b - - group: other-shard-3 - target: test-yaml-structural-other-shard-3 - - group: congress - target: test-yaml-structural-congress + # Four runners, grouped by name, run their groups one after another. + - group: states-shard-1 + states-shard-2 + group1: states-shard-1 + target1: test-yaml-structural-heavy-shard-1 + group2: states-shard-2 + target2: test-yaml-structural-heavy-shard-2 + - group: states-shard-3 + states-shard-4 + group1: states-shard-3 + target1: test-yaml-structural-heavy-shard-3 + group2: states-shard-4 + target2: test-yaml-structural-heavy-shard-4 + - group: other-shard-1 + other-shard-2a + group1: other-shard-1 + target1: test-yaml-structural-other + group2: other-shard-2a + target2: test-yaml-structural-other-shard-2a + - group: other-shard-2b + other-shard-3 + congress + group1: other-shard-2b + target1: test-yaml-structural-other-shard-2b + group2: other-shard-3 + target2: test-yaml-structural-other-shard-3 + group3: congress + target3: test-yaml-structural-congress steps: - name: Checkout repo uses: actions/checkout@v6 @@ -51,8 +58,16 @@ jobs: - name: Turn off default branching shell: bash run: bash ./update_itemization.sh - - name: Run Contrib YAML tests (${{ matrix.group }}) - run: uv run --no-sync make ${{ matrix.target }} + - name: Run Contrib YAML tests (${{ matrix.group1 }}) + run: uv run --no-sync make ${{ matrix.target1 }} - name: Run Contrib Python tests - if: ${{ matrix.group == 'other-shard-1' }} + if: ${{ matrix.group1 == 'other-shard-1' }} run: uv run --no-sync make test-policy-contrib-python + # Later groups still run after an earlier group fails, but not after + # the run is cancelled. + - name: Run Contrib YAML tests (${{ matrix.group2 }}) + if: ${{ !cancelled() && matrix.target2 }} + run: uv run --no-sync make ${{ matrix.target2 }} + - name: Run Contrib YAML tests (${{ matrix.group3 }}) + if: ${{ !cancelled() && matrix.target3 }} + run: uv run --no-sync make ${{ matrix.target3 }} diff --git a/.github/workflows/push.yaml b/.github/workflows/push.yaml index 4917f0024b3..2f5d61641c6 100644 --- a/.github/workflows/push.yaml +++ b/.github/workflows/push.yaml @@ -168,7 +168,10 @@ jobs: needs: ReleaseLock name: Full Suite - Baseline (${{ matrix.group }}) runs-on: ubuntu-latest - timeout-minutes: 60 + # Two or three groups in sequence: historical medians about 54-66 + # minutes per runner (household + ssa-usda projected at about 60 after + # the household split; about 22 for rest-a + rest-b). + timeout-minutes: 120 if: | (github.repository == 'PolicyEngine/policyengine-us') && (github.event.head_commit.message == 'Update PolicyEngine US') @@ -179,30 +182,37 @@ jobs: fail-fast: false matrix: include: - # One state per subprocess; retain four runners and one worker each. - # The former NY/OH/OK process peaked at 15.1 GB (run 36930199045). - - group: states-shard-1 - cmd: python policyengine_us/tests/test_batched.py policyengine_us/tests/policy/baseline/gov/states --mode per-subdir --workers 1 --shard 1/4 - - group: states-shard-2 - cmd: python policyengine_us/tests/test_batched.py policyengine_us/tests/policy/baseline/gov/states --mode per-subdir --workers 1 --shard 2/4 - - group: states-shard-3 - cmd: python policyengine_us/tests/test_batched.py policyengine_us/tests/policy/baseline/gov/states --mode per-subdir --workers 1 --shard 3/4 - - group: states-shard-4 - cmd: python policyengine_us/tests/test_batched.py policyengine_us/tests/policy/baseline/gov/states --mode per-subdir --workers 1 --shard 4/4 - - group: irs - cmd: make test-yaml-no-structural-other-irs - - group: household - cmd: make test-yaml-no-structural-other-household - - group: ssa-usda - cmd: make test-yaml-no-structural-other-ssa-usda - - group: rest-a - cmd: make test-yaml-no-structural-other-rest-a - - group: rest-b - cmd: make test-yaml-no-structural-other-rest-b - - group: contrib-hhs - cmd: make test-yaml-contrib-hhs - - group: reform - cmd: make test-yaml-reform + # Five runners, grouped by name, run their groups one after another + # with one worker each. The states shards still run one state per + # subprocess: the former NY/OH/OK process peaked at 15.1 GB (run + # 36930199045). + - group: states-shard-1 + states-shard-2 + group1: states-shard-1 + cmd1: python policyengine_us/tests/test_batched.py policyengine_us/tests/policy/baseline/gov/states --mode per-subdir --workers 1 --shard 1/4 + group2: states-shard-2 + cmd2: python policyengine_us/tests/test_batched.py policyengine_us/tests/policy/baseline/gov/states --mode per-subdir --workers 1 --shard 2/4 + - group: states-shard-3 + states-shard-4 + group1: states-shard-3 + cmd1: python policyengine_us/tests/test_batched.py policyengine_us/tests/policy/baseline/gov/states --mode per-subdir --workers 1 --shard 3/4 + group2: states-shard-4 + cmd2: python policyengine_us/tests/test_batched.py policyengine_us/tests/policy/baseline/gov/states --mode per-subdir --workers 1 --shard 4/4 + - group: household + ssa-usda + group1: household + cmd1: make test-yaml-no-structural-other-household + group2: ssa-usda + cmd2: make test-yaml-no-structural-other-ssa-usda + - group: contrib-hhs + irs + reform + group1: contrib-hhs + cmd1: make test-yaml-contrib-hhs + group2: irs + cmd2: make test-yaml-no-structural-other-irs + group3: reform + cmd3: make test-yaml-reform + - group: rest-a + rest-b + group1: rest-a + cmd1: make test-yaml-no-structural-other-rest-a + group2: rest-b + cmd2: make test-yaml-no-structural-other-rest-b steps: - name: Checkout repo uses: actions/checkout@v6 @@ -221,8 +231,16 @@ jobs: - name: Turn off default branching shell: bash run: bash ./update_itemization.sh - - name: Run Baseline YAML tests (${{ matrix.group }}) - run: uv run --no-sync ${{ matrix.cmd }} + - name: Run Baseline YAML tests (${{ matrix.group1 }}) + run: uv run --no-sync ${{ matrix.cmd1 }} + # Later groups still run after an earlier group fails, but not after + # the run is cancelled. + - name: Run Baseline YAML tests (${{ matrix.group2 }}) + if: ${{ !cancelled() && matrix.cmd2 }} + run: uv run --no-sync ${{ matrix.cmd2 }} + - name: Run Baseline YAML tests (${{ matrix.group3 }}) + if: ${{ !cancelled() && matrix.cmd3 }} + run: uv run --no-sync ${{ matrix.cmd3 }} HouseholdAPIPartners: needs: ReleaseLock name: Household API Partners @@ -258,7 +276,9 @@ jobs: needs: ReleaseLock name: Full Suite - Contrib (${{ matrix.group }}) runs-on: ubuntu-latest - timeout-minutes: 60 + # Two or three groups in sequence: historical medians 47-74 minutes per + # runner. + timeout-minutes: 120 if: | (github.repository == 'PolicyEngine/policyengine-us') && (github.event.head_commit.message == 'Update PolicyEngine US') @@ -269,24 +289,29 @@ jobs: fail-fast: false matrix: include: - - group: states-shard-1 - target: test-yaml-structural-heavy-shard-1 - - group: states-shard-2 - target: test-yaml-structural-heavy-shard-2 - - group: states-shard-3 - target: test-yaml-structural-heavy-shard-3 - - group: states-shard-4 - target: test-yaml-structural-heavy-shard-4 - - group: other-shard-1 - target: test-yaml-structural-other - - group: other-shard-2a - target: test-yaml-structural-other-shard-2a - - group: other-shard-2b - target: test-yaml-structural-other-shard-2b - - group: other-shard-3 - target: test-yaml-structural-other-shard-3 - - group: congress - target: test-yaml-structural-congress + # Four runners, grouped by name, run their groups one after another. + - group: states-shard-1 + states-shard-2 + group1: states-shard-1 + target1: test-yaml-structural-heavy-shard-1 + group2: states-shard-2 + target2: test-yaml-structural-heavy-shard-2 + - group: states-shard-3 + states-shard-4 + group1: states-shard-3 + target1: test-yaml-structural-heavy-shard-3 + group2: states-shard-4 + target2: test-yaml-structural-heavy-shard-4 + - group: other-shard-1 + other-shard-2a + group1: other-shard-1 + target1: test-yaml-structural-other + group2: other-shard-2a + target2: test-yaml-structural-other-shard-2a + - group: other-shard-2b + other-shard-3 + congress + group1: other-shard-2b + target1: test-yaml-structural-other-shard-2b + group2: other-shard-3 + target2: test-yaml-structural-other-shard-3 + group3: congress + target3: test-yaml-structural-congress steps: - name: Checkout repo uses: actions/checkout@v6 @@ -305,11 +330,19 @@ jobs: - name: Turn off default branching shell: bash run: bash ./update_itemization.sh - - name: Run Contrib YAML tests (${{ matrix.group }}) - run: uv run --no-sync make ${{ matrix.target }} + - name: Run Contrib YAML tests (${{ matrix.group1 }}) + run: uv run --no-sync make ${{ matrix.target1 }} - name: Run Contrib Python tests - if: ${{ matrix.group == 'other-shard-1' }} + if: ${{ matrix.group1 == 'other-shard-1' }} run: uv run --no-sync make test-policy-contrib-python + # Later groups still run after an earlier group fails, but not after + # the run is cancelled. + - name: Run Contrib YAML tests (${{ matrix.group2 }}) + if: ${{ !cancelled() && matrix.target2 }} + run: uv run --no-sync make ${{ matrix.target2 }} + - name: Run Contrib YAML tests (${{ matrix.group3 }}) + if: ${{ !cancelled() && matrix.target3 }} + run: uv run --no-sync make ${{ matrix.target3 }} RestPython: needs: ReleaseLock name: Full Suite - Rest (Python, shard-1 + variables) @@ -378,10 +411,12 @@ jobs: name: Full Suite - Rest (Python + variables) runs-on: ubuntu-latest # Historical group times: 32-46 minutes plus setup. - # This runner also retains the legacy summary check, keeping PRs at 27 runners. + # This runner also retains the legacy summary check, keeping PRs at 16 runners. timeout-minutes: 90 + # Runs after a shard-1 failure but stops when the run is cancelled. + # always() kept this runner alive after cancellation. if: | - always() + !cancelled() && (github.repository == 'PolicyEngine/policyengine-us') && (github.event.head_commit.message == 'Update PolicyEngine US') env: diff --git a/.github/workflows/python-tests.yaml b/.github/workflows/python-tests.yaml index 7d22abd7b6d..ce85b48be0d 100644 --- a/.github/workflows/python-tests.yaml +++ b/.github/workflows/python-tests.yaml @@ -66,11 +66,13 @@ jobs: rest-variables-resources.txt Rest: needs: [RestPython] - if: ${{ always() }} + # Runs after a shard-1 failure but stops when the run is cancelled. + # always() kept this runner alive after cancellation. + if: ${{ !cancelled() }} name: Full Suite - Rest (Python + variables) runs-on: ubuntu-latest # Historical group times: 32-46 minutes plus setup. - # This runner also retains the legacy summary check, keeping PRs at 27 runners. + # This runner also retains the legacy summary check, keeping PRs at 16 runners. timeout-minutes: 90 env: PYTHONUNBUFFERED: 1 diff --git a/Makefile b/Makefile index 842b3d70fb2..c1c65213df7 100644 --- a/Makefile +++ b/Makefile @@ -197,8 +197,9 @@ test-yaml-variables: $(BATCH) $(TESTS)/variables --batches 1 test-yaml-no-structural-states: # Release model/parameter caches between states. The former NY/OH/OK - # group peaked at 15.1 GB; keep one subprocess at a time on each of - # the existing four CI runners. Root-level state tests remain included. + # group peaked at 15.1 GB; keep one subprocess at a time. CI runs the + # four --shard i/4 slices two per runner, one after another. + # Root-level state tests remain included. $(BATCH) $(TESTS)/policy/baseline/gov/states --mode per-subdir --workers 1 test-yaml-no-structural-other: $(BATCH) $(TESTS)/policy/baseline --batches 2 --exclude states @@ -211,10 +212,11 @@ test-yaml-no-structural-other-irs: # 28698452678 — co-scheduling a second batch leaves no headroom. $(BATCH) $(TESTS)/policy/baseline/gov/irs --mode per-subdir --workers 1 test-yaml-no-structural-other-household: - # 4 batches (was 2), --workers 1: a 2-batch half of this folder (incl. - # the weights/MTR files) peaked 11.8 GB on CI run 28698452678. Quarter - # batches keep each subprocess <= ~6 GB. - $(BATCH) $(TESTS)/policy/baseline/household --batches 4 --workers 1 + # --workers 1: the batcher's household branch runs the root files and + # income/household one file per subprocess and each other subfolder as + # one, keeping each subprocess <= ~7 GB. The former quarter batches + # peaked at 15.1 GB on the 16 GB runner. + $(BATCH) $(TESTS)/policy/baseline/household --workers 1 test-yaml-no-structural-other-contrib: # ubi_center is microsim-heavy → per-file, one ~4 GB peak at a time. $(BATCH) $(TESTS)/policy/baseline/contrib/ubi_center --mode per-file --workers 1 @@ -237,9 +239,9 @@ test-yaml-no-structural-other-hhs: # hhs (~2.7 GB peak) rides along with the baseline-contrib runner, # which has spare headroom. $(BATCH) $(TESTS)/policy/baseline/gov/hhs --batches 1 -# baseline contrib + gov/hhs share one CI runner. policy/reform no longer -# rides along: its per-file peaks (~5.5-8 GB) plus this pair filled a job -# past the safety margin, so reform gets its own runner. +# baseline contrib + gov/hhs form one make target. policy/reform stays a +# separate target (per-file peaks ~5.5-8 GB): CI runs contrib-hhs, irs and +# reform as sequential steps on one runner, so only one peak is resident. test-yaml-contrib-hhs: test-yaml-no-structural-other-contrib test-yaml-no-structural-other-hhs test-yaml-no-structural-other-ssa: # revenue is heavy enough to need its own 2-batch split. --workers 1: diff --git a/changelog.d/ci-sequential-suites-household-memory.fixed.md b/changelog.d/ci-sequential-suites-household-memory.fixed.md new file mode 100644 index 00000000000..6722a554301 --- /dev/null +++ b/changelog.d/ci-sequential-suites-household-memory.fixed.md @@ -0,0 +1 @@ +Stop the second Rest shard when a CI run is cancelled, run the Baseline and Contrib groups in sequence on 5 and 4 runners (a PR now uses 16 runners, at most 14 at once, instead of 27), and split the household Baseline tests into subprocesses of at most about 7 GB to avoid out-of-memory runner shutdowns. diff --git a/policyengine_us/tests/code_health/test_batched_per_subdir_budget.py b/policyengine_us/tests/code_health/test_batched_per_subdir_budget.py index d02bfb5179d..62d2853d4d6 100644 --- a/policyengine_us/tests/code_health/test_batched_per_subdir_budget.py +++ b/policyengine_us/tests/code_health/test_batched_per_subdir_budget.py @@ -103,9 +103,52 @@ def test_explicit_batches_keep_both_yaml_extensions(tmp_path, mode): assert batched.count_yaml_files(tmp_path) == 2 +def test_household_runs_heavy_files_alone_and_other_folders_together(tmp_path): + household = tmp_path / "policy" / "baseline" / "household" + for relative in ( + "demographic/a.yaml", + "demographic/nested/b.yaml", + "income/household/c.yaml", + "income/household/d.yml", + "income/person/e.yaml", + "income/f.yaml", + "income/g.yaml", + "h.yaml", + "i.yml", + ): + _write_cases(household / relative, ["r1"]) + assert batched.split_into_batches(household, 4) == [ + [str(household / "demographic")], + [str(household / "income" / "household" / "c.yaml")], + [str(household / "income" / "household" / "d.yml")], + [str(household / "income" / "person")], + [str(household / "income" / "f.yaml"), str(household / "income" / "g.yaml")], + [str(household / "h.yaml")], + [str(household / "i.yml")], + ] + + +def test_household_batches_run_every_yaml_file_once(): + household = BATCHER.parent / "policy" / "baseline" / "household" + files = [ + file + for batch in batched.split_into_batches(household, 4) + for path in map(Path, batch) + for file in (batched.yaml_files(path) if path.is_dir() else [path]) + ] + assert sorted(files) == batched.yaml_files(household) + + @pytest.mark.parametrize( "relative", - ["policy/contrib", "contrib/states", "reform", "gov/states", "arbitrary"], + [ + "policy/contrib", + "contrib/states", + "reform", + "gov/states", + "baseline/household", + "arbitrary", + ], ) def test_auto_path_routes_keep_root_yml_files(tmp_path, relative): root = tmp_path / relative diff --git a/policyengine_us/tests/code_health/test_make_test_matches_ci.py b/policyengine_us/tests/code_health/test_make_test_matches_ci.py index 532945e5cf9..c08fb096b36 100644 --- a/policyengine_us/tests/code_health/test_make_test_matches_ci.py +++ b/policyengine_us/tests/code_health/test_make_test_matches_ci.py @@ -20,6 +20,8 @@ WORKFLOW = REPO_ROOT / ".github" / "workflows" / "pr.yaml" SHARD = re.compile(r"\s+--shard\s+\d+/\d+") MAKE_TARGET = re.compile(r"\bmake\s+([\w-]+)") +MATRIX_TARGET = re.compile(r"target\d*") +MATRIX_CMD = re.compile(r"cmd\d*") pytestmark = pytest.mark.skipif( shutil.which("make") is None or not WORKFLOW.exists(), @@ -60,13 +62,14 @@ def ci_suites(): for job in workflow_jobs(WORKFLOW): matrix = (job.get("strategy") or {}).get("matrix") or {} for entry in matrix.get("include") or []: - if "target" in entry: - targets.add(entry["target"]) - if "cmd" in entry: - cmd = entry["cmd"] - targets.update(MAKE_TARGET.findall(cmd)) - if "test_batched.py" in cmd: - commands.add(cmd.strip()) + # A runner can run several groups: target1, target2, cmd1, ... + for key, value in entry.items(): + if MATRIX_TARGET.fullmatch(key): + targets.add(value) + if MATRIX_CMD.fullmatch(key): + targets.update(MAKE_TARGET.findall(value)) + if "test_batched.py" in value: + commands.add(value.strip()) for step in job.get("steps") or []: for target in MAKE_TARGET.findall(step.get("run") or ""): if not target.startswith("$"): diff --git a/policyengine_us/tests/code_health/test_pr_runner_budget.py b/policyengine_us/tests/code_health/test_pr_runner_budget.py index 48fcb490cbb..a55b878f86d 100644 --- a/policyengine_us/tests/code_health/test_pr_runner_budget.py +++ b/policyengine_us/tests/code_health/test_pr_runner_budget.py @@ -14,6 +14,23 @@ } PYTHON_JOBS = GROUP_WORKFLOWS["PythonTests"]["jobs"] KEEP_RUNNING = "${{ !cancelled() }}" +PUSH_JOBS = yaml.safe_load((ROOT / ".github/workflows/push.yaml").read_text())["jobs"] +# Each runner runs its groups in this order, grouped by name. +GROUPED_SUITES = { + "Baseline": ( + ("states-shard-1", "states-shard-2"), + ("states-shard-3", "states-shard-4"), + ("household", "ssa-usda"), + ("contrib-hhs", "irs", "reform"), + ("rest-a", "rest-b"), + ), + "Contrib": ( + ("states-shard-1", "states-shard-2"), + ("states-shard-3", "states-shard-4"), + ("other-shard-1", "other-shard-2a"), + ("other-shard-2b", "other-shard-3", "congress"), + ), +} def action(name): @@ -30,6 +47,18 @@ def calls(job, name): ] +def runners(job): + matrix = job.get("strategy", {}).get("matrix", {}) + if "include" in matrix: + return len(matrix["include"]) + return math.prod(len(values) for values in matrix.values()) + + +def needs(job): + value = job.get("needs", []) + return [value] if isinstance(value, str) else value + + def test_pr_runner_budget_keeps_the_heavy_suites_separate(): counts = {} runner_jobs = {name: job for name, job in JOBS.items() if "uses" not in job} @@ -38,27 +67,87 @@ def test_pr_runner_budget_keeps_the_heavy_suites_separate(): assert name not in runner_jobs runner_jobs[name] = job for name, job in runner_jobs.items(): - matrix = job.get("strategy", {}).get("matrix", {}) - if "include" in matrix: - counts[name] = len(matrix["include"]) - else: - counts[name] = math.prod(len(values) for values in matrix.values()) + counts[name] = runners(job) # Rest hit the 90-minute timeout on #10024, #9913, #10139 and #10029 # in October 2026. One extra runner buys about 45 minutes of headroom: # two tested shards share variables and the legacy summary check. - assert sum(counts.values()) == 27 + # Baseline and Contrib run their groups in sequence on 5 and 4 runners + # (were 11 and 9), still finishing before the Rest shard chain. + assert sum(counts.values()) == 16 assert { name: counts[name] for name in ("Baseline", "Contrib", "RestPython", "Rest", "Microsimulation") } == { - "Baseline": 11, - "Contrib": 9, + "Baseline": 5, + "Contrib": 4, "RestPython": 1, "Rest": 1, "Microsimulation": 1, } +def test_a_pr_runs_at_most_14_runners_at_once(): + # Validation runs alone. Everything that needs only validation then starts + # together; Rest waits for RestPython. + first_wave = {} + for name, job in JOBS.items(): + if needs(job) != ["ReleaseLock"]: + continue + if "uses" not in job: + first_wave[name] = runners(job) + continue + for inner_name, inner in GROUP_WORKFLOWS[name]["jobs"].items(): + if not needs(inner): + first_wave[inner_name] = runners(inner) + assert "Rest" not in first_wave + assert sum(first_wave.values()) == 14 + + +def test_grouped_runners_run_every_group_once_and_stop_on_cancellation(): + for suite, runner_groups in GROUPED_SUITES.items(): + pr_job = GROUP_WORKFLOWS[suite]["jobs"][suite] + push_job = PUSH_JOBS[suite] + # Release CI must test exactly what PR CI tests. + assert push_job["strategy"] == pr_job["strategy"] + key = "cmd" if suite == "Baseline" else "target" + entries = pr_job["strategy"]["matrix"]["include"] + assert [ + tuple(entry[f"group{i}"] for i in (1, 2, 3) if f"group{i}" in entry) + for entry in entries + ] == list(runner_groups) + for entry, groups in zip(entries, runner_groups): + assert entry["group"] == " + ".join(groups) + assert set(entry) == {"group"} | { + f"{prefix}{i}" + for i in range(1, len(groups) + 1) + for prefix in ("group", key) + } + for job in (pr_job, push_job): + later = [ + step + for step in job["steps"] + if any(f"matrix.{key}{i}" in step.get("run", "") for i in (2, 3)) + ] + # A failed group must not hide the next group's results, and a + # cancelled run must stop the remaining groups. + assert [step["if"] for step in later] == [ + f"${{{{ !cancelled() && matrix.{key}{i} }}}}" for i in (2, 3) + ] + + +def test_no_test_job_keeps_running_after_cancellation(): + # A job-level always() also holds after a manual or cancel-in-progress + # cancellation, so the job keeps running or even starts. Jobs that wait for + # another suite use !cancelled() to run after its failure instead. + workflows = ["pr.yaml", "push.yaml"] + [ + Path(job["uses"]).name for job in JOBS.values() if "uses" in job + ] + for file in workflows: + jobs = yaml.safe_load((ROOT / ".github/workflows" / file).read_text())["jobs"] + for name, job in jobs.items(): + assert "always()" not in str(job.get("if", "")), (file, name) + + def test_suite_groups_have_no_extra_runners_or_triggers(): groups = { "Baseline": ("Baseline tests", "baseline-tests.yaml"), @@ -88,7 +177,8 @@ def test_suite_groups_have_no_extra_runners_or_triggers(): assert not job.get("continue-on-error") # The second tested runner also checks the first runner's result. There # is no third runner just for the summary, or a variables-only runner. - assert PYTHON_JOBS["Rest"]["if"] == "${{ always() }}" + # !cancelled(), not always(): always() kept it running after cancellation. + assert PYTHON_JOBS["Rest"]["if"] == KEEP_RUNNING assert PYTHON_JOBS["Rest"]["name"] == "Full Suite - Rest (Python + variables)" diff --git a/policyengine_us/tests/code_health/test_rest_python_groups.py b/policyengine_us/tests/code_health/test_rest_python_groups.py index f76b4e2818b..8cd9cf4c20f 100644 --- a/policyengine_us/tests/code_health/test_rest_python_groups.py +++ b/policyengine_us/tests/code_health/test_rest_python_groups.py @@ -385,19 +385,21 @@ def test_original_rest_check_aggregates_all_shards_without_hiding_failure(workfl summary = jobs["Rest"] assert summary["name"] == "Full Suite - Rest (Python + variables)" assert summary["needs"] == ["RestPython"] - assert "always()" in summary["if"] + # Shard 2 still runs after a shard-1 failure, but a cancelled run stops + # it: a job-level always() kept it running after cancellation. + assert "always()" not in summary["if"] assert "strategy" not in summary assert summary["timeout-minutes"] == 90 assert not summary.get("continue-on-error", False) assert not jobs["RestPython"].get("continue-on-error", False) if workflow.name == "push.yaml": assert " ".join(summary["if"].split()) == ( - "always() && (github.repository == 'PolicyEngine/policyengine-us') " + "!cancelled() && (github.repository == 'PolicyEngine/policyengine-us') " "&& (github.event.head_commit.message == 'Update PolicyEngine US')" ) assert {"RestPython", "Rest"} <= set(jobs["Publish"]["needs"]) else: - assert summary["if"] in ("always()", "${{ always() }}") + assert summary["if"] == "${{ !cancelled() }}" (step,) = [ candidate for candidate in summary["steps"] diff --git a/policyengine_us/tests/test_batched.py b/policyengine_us/tests/test_batched.py index 73f9e510640..314827036cc 100644 --- a/policyengine_us/tests/test_batched.py +++ b/policyengine_us/tests/test_batched.py @@ -385,6 +385,40 @@ def split_into_batches( batches.append([str(f) for f in root_files]) return batches + # Special handling for household tests. Household totals (net income, + # marginal tax rates, all-program benefits) load most of the model in + # every file, and memory accumulates across the files one subprocess + # runs. The former quarter batches peaked at 15.1 GB on the 16 GB runner + # and ended in "The runner has received a shutdown signal". Measured + # 2026-10-14 (local peak footprint, within 1% of CI's batch peaks): the + # 15 root files peaked at 13.3 GB together and income/household's 8 files + # at 11.4 GB, but no single file exceeded 6.2 GB. Those files run one per + # subprocess; every other subfolder stays one subprocess (at most 7.1 GB, + # income/spm_unit), with income split into its subfolders. + if str(base_path).endswith("baseline/household"): + income = base_path / "income" + one_per_file = {base_path, income / "household"} + batches = [] + + def visit(folder: Path) -> None: + for item in sorted(folder.iterdir()): + if not item.is_dir() or item.name in exclude: + continue + if item == income: + visit(item) + elif item in one_per_file: + batches.extend([str(f)] for f in yaml_files(item)) + else: + batches.extend(subdir_batches(item)) + loose = root_yaml_files(folder) + if folder in one_per_file: + batches.extend([str(f)] for f in loose) + elif loose: + batches.append([str(f) for f in loose]) + + visit(base_path) + return batches + # Special handling for baseline tests if "baseline" in str(base_path) and str(base_path).endswith("baseline"): states_path = base_path / "gov" / "states" From eaeecc1de696e92b30f57b4d90249de074642122 Mon Sep 17 00:00:00 2001 From: Ziming Date: Wed, 14 Oct 2026 21:44:55 -0400 Subject: [PATCH 2/4] Run gov/hhs Baseline tests one subprocess per subfolder As one process, baseline/gov/hhs peaked at 15.1 GB on the 16 GB runner in this PR's first CI run (38656284185), the same failure margin as the old household batches. Run each hhs subfolder in its own subprocess, with loose yamls as a trailing batch, like the irs suite. Co-Authored-By: Claude Opus 5.5 --- Makefile | 7 ++++--- changelog.d/ci-sequential-suites-household-memory.fixed.md | 2 +- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/Makefile b/Makefile index c1c65213df7..0944bfe7007 100644 --- a/Makefile +++ b/Makefile @@ -236,9 +236,10 @@ test-yaml-reform: # ~8 GB file runs solo with ~8 GB spare); new reform files auto-route. $(BATCH) $(TESTS)/policy/reform --mode per-file --workers 1 test-yaml-no-structural-other-hhs: - # hhs (~2.7 GB peak) rides along with the baseline-contrib runner, - # which has spare headroom. - $(BATCH) $(TESTS)/policy/baseline/gov/hhs --batches 1 + # One subprocess per hhs subfolder + trailing batch for loose yamls. + # As a single process, hhs peaked at 15.1 GB on the 16 GB runner (CI + # run 38656284185). + $(BATCH) $(TESTS)/policy/baseline/gov/hhs --mode per-subdir --workers 1 # baseline contrib + gov/hhs form one make target. policy/reform stays a # separate target (per-file peaks ~5.5-8 GB): CI runs contrib-hhs, irs and # reform as sequential steps on one runner, so only one peak is resident. diff --git a/changelog.d/ci-sequential-suites-household-memory.fixed.md b/changelog.d/ci-sequential-suites-household-memory.fixed.md index 6722a554301..b0e0ba5b301 100644 --- a/changelog.d/ci-sequential-suites-household-memory.fixed.md +++ b/changelog.d/ci-sequential-suites-household-memory.fixed.md @@ -1 +1 @@ -Stop the second Rest shard when a CI run is cancelled, run the Baseline and Contrib groups in sequence on 5 and 4 runners (a PR now uses 16 runners, at most 14 at once, instead of 27), and split the household Baseline tests into subprocesses of at most about 7 GB to avoid out-of-memory runner shutdowns. +Stop the second Rest shard when a CI run is cancelled, run the Baseline and Contrib groups in sequence on 5 and 4 runners (a PR now uses 16 runners, at most 14 at once, instead of 27), and split the household and gov/hhs Baseline tests into smaller subprocesses to avoid out-of-memory runner shutdowns. From ab2fc197981f485979fc045af0ad0159fe0ac065 Mon Sep 17 00:00:00 2001 From: Ziming Date: Wed, 14 Oct 2026 22:34:46 -0400 Subject: [PATCH 3/4] Replace fixed batch layouts with a memory budget and max-parallel lanes The grouped-runner layout passed CI but needed hand rebalancing as groups grow, and its per-file household branch and gov/hhs split were static fixes for memory that keeps growing. Replace them with mechanisms that adapt: - test_batched.py adds a pytest plugin to every YAML batch subprocess. At each file boundary it reads the process's peak memory (VmHWM; the peak physical footprint on macOS) and, past 8,000 MB, hands the files it has not started to a fresh subprocess. run_batch continues them, merges any inherited JUnit report, and passes only if every part passes. A file always runs whole. Household and gov/hhs return to their upstream layouts. - Baseline and Contrib keep one job per group and add max-parallel 5 and 4, so the Pull request workflow runs at most 14 jobs at once; GitHub starts the next group as one finishes, so lanes balance from real durations. - Release Rest shard 2 also skips when shard 1 was skipped. Measured against the full model locally, household ran as 8 subprocesses at the default budget (peaks 8.1-10.8 GB, previously 15.2 GB in one) and as 20 at a 3,000 MB budget, with exactly 644 cases each time. Co-Authored-By: Claude Opus 5.5 --- .github/workflows/baseline-tests.yaml | 75 +++---- .github/workflows/contrib-tests.yaml | 62 +++-- .github/workflows/push.yaml | 145 +++++------- .github/workflows/python-tests.yaml | 2 +- Makefile | 29 ++- ...equential-suites-household-memory.fixed.md | 2 +- .../test_batched_per_subdir_budget.py | 45 +--- .../test_batched_yaml_collection.py | 91 +++++++- .../code_health/test_make_test_matches_ci.py | 17 +- .../code_health/test_pr_runner_budget.py | 109 ++++----- .../code_health/test_rest_python_groups.py | 3 +- policyengine_us/tests/test_batched.py | 212 ++++++++++++++---- 12 files changed, 444 insertions(+), 348 deletions(-) diff --git a/.github/workflows/baseline-tests.yaml b/.github/workflows/baseline-tests.yaml index 5384aecc38f..785d170d36e 100644 --- a/.github/workflows/baseline-tests.yaml +++ b/.github/workflows/baseline-tests.yaml @@ -7,48 +7,41 @@ jobs: Baseline: name: Full Suite - Baseline (${{ matrix.group }}) runs-on: ubuntu-latest - # Two or three groups in sequence: historical medians about 54-66 - # minutes per runner (household + ssa-usda projected at about 60 after - # the household split; about 22 for rest-a + rest-b). - timeout-minutes: 120 + timeout-minutes: 60 env: PYTHONUNBUFFERED: 1 PYTEST_ADDOPTS: "--durations=25 -p no:unraisableexception -p no:threadexception" strategy: fail-fast: false + # At most 5 Baseline groups run at once; GitHub starts the next group on a + # fresh runner as each finishes, so the lanes balance from real durations. + max-parallel: 5 matrix: include: - # Five runners, grouped by name, run their groups one after another - # with one worker each. The states shards still run one state per - # subprocess: the former NY/OH/OK process peaked at 15.1 GB (run - # 36930199045). - - group: states-shard-1 + states-shard-2 - group1: states-shard-1 - cmd1: python policyengine_us/tests/test_batched.py policyengine_us/tests/policy/baseline/gov/states --mode per-subdir --workers 1 --shard 1/4 - group2: states-shard-2 - cmd2: python policyengine_us/tests/test_batched.py policyengine_us/tests/policy/baseline/gov/states --mode per-subdir --workers 1 --shard 2/4 - - group: states-shard-3 + states-shard-4 - group1: states-shard-3 - cmd1: python policyengine_us/tests/test_batched.py policyengine_us/tests/policy/baseline/gov/states --mode per-subdir --workers 1 --shard 3/4 - group2: states-shard-4 - cmd2: python policyengine_us/tests/test_batched.py policyengine_us/tests/policy/baseline/gov/states --mode per-subdir --workers 1 --shard 4/4 - - group: household + ssa-usda - group1: household - cmd1: make test-yaml-no-structural-other-household - group2: ssa-usda - cmd2: make test-yaml-no-structural-other-ssa-usda - - group: contrib-hhs + irs + reform - group1: contrib-hhs - cmd1: make test-yaml-contrib-hhs - group2: irs - cmd2: make test-yaml-no-structural-other-irs - group3: reform - cmd3: make test-yaml-reform - - group: rest-a + rest-b - group1: rest-a - cmd1: make test-yaml-no-structural-other-rest-a - group2: rest-b - cmd2: make test-yaml-no-structural-other-rest-b + # One state per subprocess; retain four runners and one worker each. + # The former NY/OH/OK process peaked at 15.1 GB (run 36930199045). + - group: states-shard-1 + cmd: python policyengine_us/tests/test_batched.py policyengine_us/tests/policy/baseline/gov/states --mode per-subdir --workers 1 --shard 1/4 + - group: states-shard-2 + cmd: python policyengine_us/tests/test_batched.py policyengine_us/tests/policy/baseline/gov/states --mode per-subdir --workers 1 --shard 2/4 + - group: states-shard-3 + cmd: python policyengine_us/tests/test_batched.py policyengine_us/tests/policy/baseline/gov/states --mode per-subdir --workers 1 --shard 3/4 + - group: states-shard-4 + cmd: python policyengine_us/tests/test_batched.py policyengine_us/tests/policy/baseline/gov/states --mode per-subdir --workers 1 --shard 4/4 + - group: irs + cmd: make test-yaml-no-structural-other-irs + - group: household + cmd: make test-yaml-no-structural-other-household + - group: ssa-usda + cmd: make test-yaml-no-structural-other-ssa-usda + - group: rest-a + cmd: make test-yaml-no-structural-other-rest-a + - group: rest-b + cmd: make test-yaml-no-structural-other-rest-b + - group: contrib-hhs + cmd: make test-yaml-contrib-hhs + - group: reform + cmd: make test-yaml-reform steps: - name: Checkout repo uses: actions/checkout@v6 @@ -67,13 +60,5 @@ jobs: - name: Turn off default branching shell: bash run: bash ./update_itemization.sh - - name: Run Baseline YAML tests (${{ matrix.group1 }}) - run: uv run --no-sync ${{ matrix.cmd1 }} - # Later groups still run after an earlier group fails, but not after - # the run is cancelled. - - name: Run Baseline YAML tests (${{ matrix.group2 }}) - if: ${{ !cancelled() && matrix.cmd2 }} - run: uv run --no-sync ${{ matrix.cmd2 }} - - name: Run Baseline YAML tests (${{ matrix.group3 }}) - if: ${{ !cancelled() && matrix.cmd3 }} - run: uv run --no-sync ${{ matrix.cmd3 }} + - name: Run Baseline YAML tests (${{ matrix.group }}) + run: uv run --no-sync ${{ matrix.cmd }} diff --git a/.github/workflows/contrib-tests.yaml b/.github/workflows/contrib-tests.yaml index a34b47b53d3..9f4c8c413fc 100644 --- a/.github/workflows/contrib-tests.yaml +++ b/.github/workflows/contrib-tests.yaml @@ -7,39 +7,35 @@ jobs: Contrib: name: Full Suite - Contrib (${{ matrix.group }}) runs-on: ubuntu-latest - # Two or three groups in sequence: historical medians 47-74 minutes per - # runner. - timeout-minutes: 120 + timeout-minutes: 60 env: PYTHONUNBUFFERED: 1 PYTEST_ADDOPTS: "--durations=25 -p no:unraisableexception -p no:threadexception" strategy: fail-fast: false + # At most 4 Contrib groups run at once; GitHub starts the next group on a + # fresh runner as each finishes, so the lanes balance from real durations. + max-parallel: 4 matrix: include: - # Four runners, grouped by name, run their groups one after another. - - group: states-shard-1 + states-shard-2 - group1: states-shard-1 - target1: test-yaml-structural-heavy-shard-1 - group2: states-shard-2 - target2: test-yaml-structural-heavy-shard-2 - - group: states-shard-3 + states-shard-4 - group1: states-shard-3 - target1: test-yaml-structural-heavy-shard-3 - group2: states-shard-4 - target2: test-yaml-structural-heavy-shard-4 - - group: other-shard-1 + other-shard-2a - group1: other-shard-1 - target1: test-yaml-structural-other - group2: other-shard-2a - target2: test-yaml-structural-other-shard-2a - - group: other-shard-2b + other-shard-3 + congress - group1: other-shard-2b - target1: test-yaml-structural-other-shard-2b - group2: other-shard-3 - target2: test-yaml-structural-other-shard-3 - group3: congress - target3: test-yaml-structural-congress + - group: states-shard-1 + target: test-yaml-structural-heavy-shard-1 + - group: states-shard-2 + target: test-yaml-structural-heavy-shard-2 + - group: states-shard-3 + target: test-yaml-structural-heavy-shard-3 + - group: states-shard-4 + target: test-yaml-structural-heavy-shard-4 + - group: other-shard-1 + target: test-yaml-structural-other + - group: other-shard-2a + target: test-yaml-structural-other-shard-2a + - group: other-shard-2b + target: test-yaml-structural-other-shard-2b + - group: other-shard-3 + target: test-yaml-structural-other-shard-3 + - group: congress + target: test-yaml-structural-congress steps: - name: Checkout repo uses: actions/checkout@v6 @@ -58,16 +54,8 @@ jobs: - name: Turn off default branching shell: bash run: bash ./update_itemization.sh - - name: Run Contrib YAML tests (${{ matrix.group1 }}) - run: uv run --no-sync make ${{ matrix.target1 }} + - name: Run Contrib YAML tests (${{ matrix.group }}) + run: uv run --no-sync make ${{ matrix.target }} - name: Run Contrib Python tests - if: ${{ matrix.group1 == 'other-shard-1' }} + if: ${{ matrix.group == 'other-shard-1' }} run: uv run --no-sync make test-policy-contrib-python - # Later groups still run after an earlier group fails, but not after - # the run is cancelled. - - name: Run Contrib YAML tests (${{ matrix.group2 }}) - if: ${{ !cancelled() && matrix.target2 }} - run: uv run --no-sync make ${{ matrix.target2 }} - - name: Run Contrib YAML tests (${{ matrix.group3 }}) - if: ${{ !cancelled() && matrix.target3 }} - run: uv run --no-sync make ${{ matrix.target3 }} diff --git a/.github/workflows/push.yaml b/.github/workflows/push.yaml index 2f5d61641c6..8d8dc9e3c25 100644 --- a/.github/workflows/push.yaml +++ b/.github/workflows/push.yaml @@ -168,10 +168,7 @@ jobs: needs: ReleaseLock name: Full Suite - Baseline (${{ matrix.group }}) runs-on: ubuntu-latest - # Two or three groups in sequence: historical medians about 54-66 - # minutes per runner (household + ssa-usda projected at about 60 after - # the household split; about 22 for rest-a + rest-b). - timeout-minutes: 120 + timeout-minutes: 60 if: | (github.repository == 'PolicyEngine/policyengine-us') && (github.event.head_commit.message == 'Update PolicyEngine US') @@ -180,39 +177,35 @@ jobs: PYTEST_ADDOPTS: "--durations=25 -p no:unraisableexception -p no:threadexception" strategy: fail-fast: false + # At most 5 Baseline groups run at once; GitHub starts the next group on a + # fresh runner as each finishes, so the lanes balance from real durations. + max-parallel: 5 matrix: include: - # Five runners, grouped by name, run their groups one after another - # with one worker each. The states shards still run one state per - # subprocess: the former NY/OH/OK process peaked at 15.1 GB (run - # 36930199045). - - group: states-shard-1 + states-shard-2 - group1: states-shard-1 - cmd1: python policyengine_us/tests/test_batched.py policyengine_us/tests/policy/baseline/gov/states --mode per-subdir --workers 1 --shard 1/4 - group2: states-shard-2 - cmd2: python policyengine_us/tests/test_batched.py policyengine_us/tests/policy/baseline/gov/states --mode per-subdir --workers 1 --shard 2/4 - - group: states-shard-3 + states-shard-4 - group1: states-shard-3 - cmd1: python policyengine_us/tests/test_batched.py policyengine_us/tests/policy/baseline/gov/states --mode per-subdir --workers 1 --shard 3/4 - group2: states-shard-4 - cmd2: python policyengine_us/tests/test_batched.py policyengine_us/tests/policy/baseline/gov/states --mode per-subdir --workers 1 --shard 4/4 - - group: household + ssa-usda - group1: household - cmd1: make test-yaml-no-structural-other-household - group2: ssa-usda - cmd2: make test-yaml-no-structural-other-ssa-usda - - group: contrib-hhs + irs + reform - group1: contrib-hhs - cmd1: make test-yaml-contrib-hhs - group2: irs - cmd2: make test-yaml-no-structural-other-irs - group3: reform - cmd3: make test-yaml-reform - - group: rest-a + rest-b - group1: rest-a - cmd1: make test-yaml-no-structural-other-rest-a - group2: rest-b - cmd2: make test-yaml-no-structural-other-rest-b + # One state per subprocess; retain four runners and one worker each. + # The former NY/OH/OK process peaked at 15.1 GB (run 36930199045). + - group: states-shard-1 + cmd: python policyengine_us/tests/test_batched.py policyengine_us/tests/policy/baseline/gov/states --mode per-subdir --workers 1 --shard 1/4 + - group: states-shard-2 + cmd: python policyengine_us/tests/test_batched.py policyengine_us/tests/policy/baseline/gov/states --mode per-subdir --workers 1 --shard 2/4 + - group: states-shard-3 + cmd: python policyengine_us/tests/test_batched.py policyengine_us/tests/policy/baseline/gov/states --mode per-subdir --workers 1 --shard 3/4 + - group: states-shard-4 + cmd: python policyengine_us/tests/test_batched.py policyengine_us/tests/policy/baseline/gov/states --mode per-subdir --workers 1 --shard 4/4 + - group: irs + cmd: make test-yaml-no-structural-other-irs + - group: household + cmd: make test-yaml-no-structural-other-household + - group: ssa-usda + cmd: make test-yaml-no-structural-other-ssa-usda + - group: rest-a + cmd: make test-yaml-no-structural-other-rest-a + - group: rest-b + cmd: make test-yaml-no-structural-other-rest-b + - group: contrib-hhs + cmd: make test-yaml-contrib-hhs + - group: reform + cmd: make test-yaml-reform steps: - name: Checkout repo uses: actions/checkout@v6 @@ -231,16 +224,8 @@ jobs: - name: Turn off default branching shell: bash run: bash ./update_itemization.sh - - name: Run Baseline YAML tests (${{ matrix.group1 }}) - run: uv run --no-sync ${{ matrix.cmd1 }} - # Later groups still run after an earlier group fails, but not after - # the run is cancelled. - - name: Run Baseline YAML tests (${{ matrix.group2 }}) - if: ${{ !cancelled() && matrix.cmd2 }} - run: uv run --no-sync ${{ matrix.cmd2 }} - - name: Run Baseline YAML tests (${{ matrix.group3 }}) - if: ${{ !cancelled() && matrix.cmd3 }} - run: uv run --no-sync ${{ matrix.cmd3 }} + - name: Run Baseline YAML tests (${{ matrix.group }}) + run: uv run --no-sync ${{ matrix.cmd }} HouseholdAPIPartners: needs: ReleaseLock name: Household API Partners @@ -276,9 +261,7 @@ jobs: needs: ReleaseLock name: Full Suite - Contrib (${{ matrix.group }}) runs-on: ubuntu-latest - # Two or three groups in sequence: historical medians 47-74 minutes per - # runner. - timeout-minutes: 120 + timeout-minutes: 60 if: | (github.repository == 'PolicyEngine/policyengine-us') && (github.event.head_commit.message == 'Update PolicyEngine US') @@ -287,31 +270,29 @@ jobs: PYTEST_ADDOPTS: "--durations=25 -p no:unraisableexception -p no:threadexception" strategy: fail-fast: false + # At most 4 Contrib groups run at once; GitHub starts the next group on a + # fresh runner as each finishes, so the lanes balance from real durations. + max-parallel: 4 matrix: include: - # Four runners, grouped by name, run their groups one after another. - - group: states-shard-1 + states-shard-2 - group1: states-shard-1 - target1: test-yaml-structural-heavy-shard-1 - group2: states-shard-2 - target2: test-yaml-structural-heavy-shard-2 - - group: states-shard-3 + states-shard-4 - group1: states-shard-3 - target1: test-yaml-structural-heavy-shard-3 - group2: states-shard-4 - target2: test-yaml-structural-heavy-shard-4 - - group: other-shard-1 + other-shard-2a - group1: other-shard-1 - target1: test-yaml-structural-other - group2: other-shard-2a - target2: test-yaml-structural-other-shard-2a - - group: other-shard-2b + other-shard-3 + congress - group1: other-shard-2b - target1: test-yaml-structural-other-shard-2b - group2: other-shard-3 - target2: test-yaml-structural-other-shard-3 - group3: congress - target3: test-yaml-structural-congress + - group: states-shard-1 + target: test-yaml-structural-heavy-shard-1 + - group: states-shard-2 + target: test-yaml-structural-heavy-shard-2 + - group: states-shard-3 + target: test-yaml-structural-heavy-shard-3 + - group: states-shard-4 + target: test-yaml-structural-heavy-shard-4 + - group: other-shard-1 + target: test-yaml-structural-other + - group: other-shard-2a + target: test-yaml-structural-other-shard-2a + - group: other-shard-2b + target: test-yaml-structural-other-shard-2b + - group: other-shard-3 + target: test-yaml-structural-other-shard-3 + - group: congress + target: test-yaml-structural-congress steps: - name: Checkout repo uses: actions/checkout@v6 @@ -330,19 +311,11 @@ jobs: - name: Turn off default branching shell: bash run: bash ./update_itemization.sh - - name: Run Contrib YAML tests (${{ matrix.group1 }}) - run: uv run --no-sync make ${{ matrix.target1 }} + - name: Run Contrib YAML tests (${{ matrix.group }}) + run: uv run --no-sync make ${{ matrix.target }} - name: Run Contrib Python tests - if: ${{ matrix.group1 == 'other-shard-1' }} + if: ${{ matrix.group == 'other-shard-1' }} run: uv run --no-sync make test-policy-contrib-python - # Later groups still run after an earlier group fails, but not after - # the run is cancelled. - - name: Run Contrib YAML tests (${{ matrix.group2 }}) - if: ${{ !cancelled() && matrix.target2 }} - run: uv run --no-sync make ${{ matrix.target2 }} - - name: Run Contrib YAML tests (${{ matrix.group3 }}) - if: ${{ !cancelled() && matrix.target3 }} - run: uv run --no-sync make ${{ matrix.target3 }} RestPython: needs: ReleaseLock name: Full Suite - Rest (Python, shard-1 + variables) @@ -411,12 +384,14 @@ jobs: name: Full Suite - Rest (Python + variables) runs-on: ubuntu-latest # Historical group times: 32-46 minutes plus setup. - # This runner also retains the legacy summary check, keeping PRs at 16 runners. + # This runner also retains the legacy summary check, keeping PRs at 27 runners. timeout-minutes: 90 - # Runs after a shard-1 failure but stops when the run is cancelled. - # always() kept this runner alive after cancellation. + # Runs after a shard-1 failure, but not when shard 1 was skipped (a failed + # release lock) or the run is cancelled. always() kept this runner alive + # after cancellation. if: | !cancelled() + && (needs.RestPython.result != 'skipped') && (github.repository == 'PolicyEngine/policyengine-us') && (github.event.head_commit.message == 'Update PolicyEngine US') env: diff --git a/.github/workflows/python-tests.yaml b/.github/workflows/python-tests.yaml index ce85b48be0d..82db686ed99 100644 --- a/.github/workflows/python-tests.yaml +++ b/.github/workflows/python-tests.yaml @@ -72,7 +72,7 @@ jobs: name: Full Suite - Rest (Python + variables) runs-on: ubuntu-latest # Historical group times: 32-46 minutes plus setup. - # This runner also retains the legacy summary check, keeping PRs at 16 runners. + # This runner also retains the legacy summary check, keeping PRs at 27 runners. timeout-minutes: 90 env: PYTHONUNBUFFERED: 1 diff --git a/Makefile b/Makefile index 0944bfe7007..19d419af1bd 100644 --- a/Makefile +++ b/Makefile @@ -197,9 +197,8 @@ test-yaml-variables: $(BATCH) $(TESTS)/variables --batches 1 test-yaml-no-structural-states: # Release model/parameter caches between states. The former NY/OH/OK - # group peaked at 15.1 GB; keep one subprocess at a time. CI runs the - # four --shard i/4 slices two per runner, one after another. - # Root-level state tests remain included. + # group peaked at 15.1 GB; keep one subprocess at a time on each of + # the existing four CI runners. Root-level state tests remain included. $(BATCH) $(TESTS)/policy/baseline/gov/states --mode per-subdir --workers 1 test-yaml-no-structural-other: $(BATCH) $(TESTS)/policy/baseline --batches 2 --exclude states @@ -212,11 +211,11 @@ test-yaml-no-structural-other-irs: # 28698452678 — co-scheduling a second batch leaves no headroom. $(BATCH) $(TESTS)/policy/baseline/gov/irs --mode per-subdir --workers 1 test-yaml-no-structural-other-household: - # --workers 1: the batcher's household branch runs the root files and - # income/household one file per subprocess and each other subfolder as - # one, keeping each subprocess <= ~7 GB. The former quarter batches - # peaked at 15.1 GB on the 16 GB runner. - $(BATCH) $(TESTS)/policy/baseline/household --workers 1 + # 4 batches (was 2), --workers 1: a 2-batch half of this folder (incl. + # the weights/MTR files) peaked 11.8 GB on CI run 28698452678, and by + # October 2026 the last quarter batch reached 15.1 GB. test_batched.py's + # memory budget now continues a batch in a fresh subprocess past ~8 GB. + $(BATCH) $(TESTS)/policy/baseline/household --batches 4 --workers 1 test-yaml-no-structural-other-contrib: # ubi_center is microsim-heavy → per-file, one ~4 GB peak at a time. $(BATCH) $(TESTS)/policy/baseline/contrib/ubi_center --mode per-file --workers 1 @@ -236,13 +235,13 @@ test-yaml-reform: # ~8 GB file runs solo with ~8 GB spare); new reform files auto-route. $(BATCH) $(TESTS)/policy/reform --mode per-file --workers 1 test-yaml-no-structural-other-hhs: - # One subprocess per hhs subfolder + trailing batch for loose yamls. - # As a single process, hhs peaked at 15.1 GB on the 16 GB runner (CI - # run 38656284185). - $(BATCH) $(TESTS)/policy/baseline/gov/hhs --mode per-subdir --workers 1 -# baseline contrib + gov/hhs form one make target. policy/reform stays a -# separate target (per-file peaks ~5.5-8 GB): CI runs contrib-hhs, irs and -# reform as sequential steps on one runner, so only one peak is resident. + # hhs rides along with the baseline-contrib runner. As one process it + # peaked at 15.1 GB (CI run 38656284185); test_batched.py's memory budget + # continues it in a fresh subprocess past ~8 GB. + $(BATCH) $(TESTS)/policy/baseline/gov/hhs --batches 1 +# baseline contrib + gov/hhs share one CI runner. policy/reform no longer +# rides along: its per-file peaks (~5.5-8 GB) plus this pair filled a job +# past the safety margin, so reform gets its own runner. test-yaml-contrib-hhs: test-yaml-no-structural-other-contrib test-yaml-no-structural-other-hhs test-yaml-no-structural-other-ssa: # revenue is heavy enough to need its own 2-batch split. --workers 1: diff --git a/changelog.d/ci-sequential-suites-household-memory.fixed.md b/changelog.d/ci-sequential-suites-household-memory.fixed.md index b0e0ba5b301..62e022227a3 100644 --- a/changelog.d/ci-sequential-suites-household-memory.fixed.md +++ b/changelog.d/ci-sequential-suites-household-memory.fixed.md @@ -1 +1 @@ -Stop the second Rest shard when a CI run is cancelled, run the Baseline and Contrib groups in sequence on 5 and 4 runners (a PR now uses 16 runners, at most 14 at once, instead of 27), and split the household and gov/hhs Baseline tests into smaller subprocesses to avoid out-of-memory runner shutdowns. +Stop the second Rest shard when a CI run is cancelled, run at most 5 Baseline and 4 Contrib groups at once so the Pull request workflow runs at most 14 jobs at a time instead of about 25, and continue YAML test batches in a fresh subprocess once one passes an 8 GB memory budget, so batches no longer need re-splitting by hand as tests grow. diff --git a/policyengine_us/tests/code_health/test_batched_per_subdir_budget.py b/policyengine_us/tests/code_health/test_batched_per_subdir_budget.py index 62d2853d4d6..d02bfb5179d 100644 --- a/policyengine_us/tests/code_health/test_batched_per_subdir_budget.py +++ b/policyengine_us/tests/code_health/test_batched_per_subdir_budget.py @@ -103,52 +103,9 @@ def test_explicit_batches_keep_both_yaml_extensions(tmp_path, mode): assert batched.count_yaml_files(tmp_path) == 2 -def test_household_runs_heavy_files_alone_and_other_folders_together(tmp_path): - household = tmp_path / "policy" / "baseline" / "household" - for relative in ( - "demographic/a.yaml", - "demographic/nested/b.yaml", - "income/household/c.yaml", - "income/household/d.yml", - "income/person/e.yaml", - "income/f.yaml", - "income/g.yaml", - "h.yaml", - "i.yml", - ): - _write_cases(household / relative, ["r1"]) - assert batched.split_into_batches(household, 4) == [ - [str(household / "demographic")], - [str(household / "income" / "household" / "c.yaml")], - [str(household / "income" / "household" / "d.yml")], - [str(household / "income" / "person")], - [str(household / "income" / "f.yaml"), str(household / "income" / "g.yaml")], - [str(household / "h.yaml")], - [str(household / "i.yml")], - ] - - -def test_household_batches_run_every_yaml_file_once(): - household = BATCHER.parent / "policy" / "baseline" / "household" - files = [ - file - for batch in batched.split_into_batches(household, 4) - for path in map(Path, batch) - for file in (batched.yaml_files(path) if path.is_dir() else [path]) - ] - assert sorted(files) == batched.yaml_files(household) - - @pytest.mark.parametrize( "relative", - [ - "policy/contrib", - "contrib/states", - "reform", - "gov/states", - "baseline/household", - "arbitrary", - ], + ["policy/contrib", "contrib/states", "reform", "gov/states", "arbitrary"], ) def test_auto_path_routes_keep_root_yml_files(tmp_path, relative): root = tmp_path / relative diff --git a/policyengine_us/tests/code_health/test_batched_yaml_collection.py b/policyengine_us/tests/code_health/test_batched_yaml_collection.py index d2e6d20234c..f48219dd780 100644 --- a/policyengine_us/tests/code_health/test_batched_yaml_collection.py +++ b/policyengine_us/tests/code_health/test_batched_yaml_collection.py @@ -1,11 +1,13 @@ -"""YAML batches must not rerun the separately routed Python suites. +"""YAML batches must not rerun the separately routed Python suites, and must +continue in a fresh subprocess once one passes the memory budget. Exercise the real Core CLI and pytest collector with a one-variable country -package, so this runner regression never constructs the full US model. +package, so these runner regressions never construct the full US model. """ import importlib.util import os +import re import shlex import subprocess import sys @@ -13,6 +15,7 @@ from pathlib import Path import pytest +import yaml BATCHER = Path(__file__).resolve().parents[1] / "test_batched.py" SPEC = importlib.util.spec_from_file_location("test_batched", BATCHER) @@ -107,3 +110,87 @@ def test_python_only_directory_is_not_a_passing_yaml_batch(small_country): def test_python_only_subdir_does_not_create_a_yaml_batch(small_country): assert batched.split_into_batches(small_country.parent, 2, mode="per-subdir") == [] + + +def _write_age_cases(path, count, failing=0): + """`count` cases for one person's age; the first `failing` of them fail.""" + path.write_text( + yaml.safe_dump( + [ + { + "name": f"{path.stem} {i}", + "period": 2026, + "input": {"age": 30}, + "output": {"age": 31 if i < failing else 30}, + } + for i in range(count) + ] + ) + ) + + +def _part_summaries(output): + return re.findall(r"=+ ((?:\d+ failed, )?\d+ passed)", output) + + +def test_memory_budget_continues_unstarted_files_in_fresh_subprocesses( + small_country, monkeypatch +): + for name in ("a", "b", "c"): + _write_age_cases(small_country / f"{name}.yaml", 2) + # Every process is past a zero budget, so each file gets its own part. + monkeypatch.setenv(batched.MEMORY_BUDGET_ENV, "0") + report = small_country.parent / "report.xml" + monkeypatch.setenv("PYTEST_ADDOPTS", f"--junitxml={report}") + + result = batched.run_batch([str(small_country)], "budget", stream=False) + + assert result["status"] == "passed", result["output"] + assert result["parts"] == 3, result["output"] + # Each file ran exactly once, both of its cases, with none repeated. + assert _part_summaries(result["output"]) == ["2 passed"] * 3, result["output"] + # Every part rewrites the inherited report; the batch merges them back. + assert len(ET.parse(report).findall(".//testcase")) == 6 + + +def test_memory_budget_leaves_a_batch_under_it_in_one_subprocess( + small_country, monkeypatch +): + for name in ("a", "b", "c"): + _write_age_cases(small_country / f"{name}.yaml", 2) + monkeypatch.delenv(batched.MEMORY_BUDGET_ENV, raising=False) + + result = batched.run_batch([str(small_country)], "budget", stream=False) + + assert result["status"] == "passed", result["output"] + assert result["parts"] == 1, result["output"] + assert _part_summaries(result["output"]) == ["6 passed"], result["output"] + + +def test_memory_budget_never_splits_a_file(small_country, monkeypatch): + _write_age_cases(small_country / "a.yaml", 3) + monkeypatch.setenv(batched.MEMORY_BUDGET_ENV, "0") + + result = batched.run_batch([str(small_country)], "budget", stream=False) + + assert result["parts"] == 1, result["output"] + assert _part_summaries(result["output"]) == ["3 passed"], result["output"] + + +def test_a_failure_before_the_budget_stop_fails_the_batch_but_runs_the_rest( + small_country, monkeypatch +): + _write_age_cases(small_country / "a.yaml", 2, failing=1) + for name in ("b", "c"): + _write_age_cases(small_country / f"{name}.yaml", 2) + monkeypatch.setenv(batched.MEMORY_BUDGET_ENV, "0") + + result = batched.run_batch([str(small_country)], "budget", stream=False) + + assert result["status"] == "failed", result["output"] + assert result["returncode"] != 0 + assert _part_summaries(result["output"]) == [ + "1 failed, 1 passed", + "2 passed", + "2 passed", + ], result["output"] diff --git a/policyengine_us/tests/code_health/test_make_test_matches_ci.py b/policyengine_us/tests/code_health/test_make_test_matches_ci.py index c08fb096b36..532945e5cf9 100644 --- a/policyengine_us/tests/code_health/test_make_test_matches_ci.py +++ b/policyengine_us/tests/code_health/test_make_test_matches_ci.py @@ -20,8 +20,6 @@ WORKFLOW = REPO_ROOT / ".github" / "workflows" / "pr.yaml" SHARD = re.compile(r"\s+--shard\s+\d+/\d+") MAKE_TARGET = re.compile(r"\bmake\s+([\w-]+)") -MATRIX_TARGET = re.compile(r"target\d*") -MATRIX_CMD = re.compile(r"cmd\d*") pytestmark = pytest.mark.skipif( shutil.which("make") is None or not WORKFLOW.exists(), @@ -62,14 +60,13 @@ def ci_suites(): for job in workflow_jobs(WORKFLOW): matrix = (job.get("strategy") or {}).get("matrix") or {} for entry in matrix.get("include") or []: - # A runner can run several groups: target1, target2, cmd1, ... - for key, value in entry.items(): - if MATRIX_TARGET.fullmatch(key): - targets.add(value) - if MATRIX_CMD.fullmatch(key): - targets.update(MAKE_TARGET.findall(value)) - if "test_batched.py" in value: - commands.add(value.strip()) + if "target" in entry: + targets.add(entry["target"]) + if "cmd" in entry: + cmd = entry["cmd"] + targets.update(MAKE_TARGET.findall(cmd)) + if "test_batched.py" in cmd: + commands.add(cmd.strip()) for step in job.get("steps") or []: for target in MAKE_TARGET.findall(step.get("run") or ""): if not target.startswith("$"): diff --git a/policyengine_us/tests/code_health/test_pr_runner_budget.py b/policyengine_us/tests/code_health/test_pr_runner_budget.py index a55b878f86d..209f69321d0 100644 --- a/policyengine_us/tests/code_health/test_pr_runner_budget.py +++ b/policyengine_us/tests/code_health/test_pr_runner_budget.py @@ -13,24 +13,8 @@ if "uses" in job } PYTHON_JOBS = GROUP_WORKFLOWS["PythonTests"]["jobs"] -KEEP_RUNNING = "${{ !cancelled() }}" PUSH_JOBS = yaml.safe_load((ROOT / ".github/workflows/push.yaml").read_text())["jobs"] -# Each runner runs its groups in this order, grouped by name. -GROUPED_SUITES = { - "Baseline": ( - ("states-shard-1", "states-shard-2"), - ("states-shard-3", "states-shard-4"), - ("household", "ssa-usda"), - ("contrib-hhs", "irs", "reform"), - ("rest-a", "rest-b"), - ), - "Contrib": ( - ("states-shard-1", "states-shard-2"), - ("states-shard-3", "states-shard-4"), - ("other-shard-1", "other-shard-2a"), - ("other-shard-2b", "other-shard-3", "congress"), - ), -} +KEEP_RUNNING = "${{ !cancelled() }}" def action(name): @@ -54,11 +38,6 @@ def runners(job): return math.prod(len(values) for values in matrix.values()) -def needs(job): - value = job.get("needs", []) - return [value] if isinstance(value, str) else value - - def test_pr_runner_budget_keeps_the_heavy_suites_separate(): counts = {} runner_jobs = {name: job for name, job in JOBS.items() if "uses" not in job} @@ -71,79 +50,71 @@ def test_pr_runner_budget_keeps_the_heavy_suites_separate(): # Rest hit the 90-minute timeout on #10024, #9913, #10139 and #10029 # in October 2026. One extra runner buys about 45 minutes of headroom: # two tested shards share variables and the legacy summary check. - # Baseline and Contrib run their groups in sequence on 5 and 4 runners - # (were 11 and 9), still finishing before the Rest shard chain. - assert sum(counts.values()) == 16 + assert sum(counts.values()) == 27 assert { name: counts[name] for name in ("Baseline", "Contrib", "RestPython", "Rest", "Microsimulation") } == { - "Baseline": 5, - "Contrib": 4, + "Baseline": 11, + "Contrib": 9, "RestPython": 1, "Rest": 1, "Microsimulation": 1, } -def test_a_pr_runs_at_most_14_runners_at_once(): - # Validation runs alone. Everything that needs only validation then starts - # together; Rest waits for RestPython. +def needs(job): + value = job.get("needs", []) + return [value] if isinstance(value, str) else value + + +def test_lanes_cap_a_pr_at_14_runners_at_once(): + # Within pr.yaml, validation runs alone; everything that needs only it starts + # together, and Rest waits for RestPython. Baseline and Contrib each keep + # one job per group but run at most max-parallel of them at once: GitHub + # starts the next group as one finishes, so the lanes rebalance from real + # durations on every run, with nothing to re-pair as groups grow. + assert ( + GROUP_WORKFLOWS["Baseline"]["jobs"]["Baseline"]["strategy"]["max-parallel"] == 5 + ) + assert ( + GROUP_WORKFLOWS["Contrib"]["jobs"]["Contrib"]["strategy"]["max-parallel"] == 4 + ) + + def concurrent(job): + return min(job.get("strategy", {}).get("max-parallel", math.inf), runners(job)) + first_wave = {} for name, job in JOBS.items(): if needs(job) != ["ReleaseLock"]: continue if "uses" not in job: - first_wave[name] = runners(job) + first_wave[name] = concurrent(job) continue for inner_name, inner in GROUP_WORKFLOWS[name]["jobs"].items(): if not needs(inner): - first_wave[inner_name] = runners(inner) + first_wave[inner_name] = concurrent(inner) assert "Rest" not in first_wave assert sum(first_wave.values()) == 14 -def test_grouped_runners_run_every_group_once_and_stop_on_cancellation(): - for suite, runner_groups in GROUPED_SUITES.items(): - pr_job = GROUP_WORKFLOWS[suite]["jobs"][suite] - push_job = PUSH_JOBS[suite] - # Release CI must test exactly what PR CI tests. - assert push_job["strategy"] == pr_job["strategy"] - key = "cmd" if suite == "Baseline" else "target" - entries = pr_job["strategy"]["matrix"]["include"] - assert [ - tuple(entry[f"group{i}"] for i in (1, 2, 3) if f"group{i}" in entry) - for entry in entries - ] == list(runner_groups) - for entry, groups in zip(entries, runner_groups): - assert entry["group"] == " + ".join(groups) - assert set(entry) == {"group"} | { - f"{prefix}{i}" - for i in range(1, len(groups) + 1) - for prefix in ("group", key) - } - for job in (pr_job, push_job): - later = [ - step - for step in job["steps"] - if any(f"matrix.{key}{i}" in step.get("run", "") for i in (2, 3)) - ] - # A failed group must not hide the next group's results, and a - # cancelled run must stop the remaining groups. - assert [step["if"] for step in later] == [ - f"${{{{ !cancelled() && matrix.{key}{i} }}}}" for i in (2, 3) - ] +def test_release_ci_runs_the_same_baseline_and_contrib_lanes(): + for name in ("Baseline", "Contrib"): + pr_job = GROUP_WORKFLOWS[name]["jobs"][name] + assert PUSH_JOBS[name]["strategy"] == pr_job["strategy"] + assert PUSH_JOBS[name]["steps"] == pr_job["steps"] def test_no_test_job_keeps_running_after_cancellation(): # A job-level always() also holds after a manual or cancel-in-progress - # cancellation, so the job keeps running or even starts. Jobs that wait for - # another suite use !cancelled() to run after its failure instead. - workflows = ["pr.yaml", "push.yaml"] + [ - Path(job["uses"]).name for job in JOBS.values() if "uses" in job - ] - for file in workflows: - jobs = yaml.safe_load((ROOT / ".github/workflows" / file).read_text())["jobs"] + # cancellation, so the job keeps running or even starts after it. A job + # that waits for another uses !cancelled() to run after its failure. + workflows = {"pr.yaml": JOBS, "push.yaml": PUSH_JOBS} | { + Path(job["uses"]).name: GROUP_WORKFLOWS[name]["jobs"] + for name, job in JOBS.items() + if "uses" in job + } + for file, jobs in workflows.items(): for name, job in jobs.items(): assert "always()" not in str(job.get("if", "")), (file, name) diff --git a/policyengine_us/tests/code_health/test_rest_python_groups.py b/policyengine_us/tests/code_health/test_rest_python_groups.py index 8cd9cf4c20f..73e52e9a801 100644 --- a/policyengine_us/tests/code_health/test_rest_python_groups.py +++ b/policyengine_us/tests/code_health/test_rest_python_groups.py @@ -394,7 +394,8 @@ def test_original_rest_check_aggregates_all_shards_without_hiding_failure(workfl assert not jobs["RestPython"].get("continue-on-error", False) if workflow.name == "push.yaml": assert " ".join(summary["if"].split()) == ( - "!cancelled() && (github.repository == 'PolicyEngine/policyengine-us') " + "!cancelled() && (needs.RestPython.result != 'skipped') " + "&& (github.repository == 'PolicyEngine/policyengine-us') " "&& (github.event.head_commit.message == 'Update PolicyEngine US')" ) assert {"RestPython", "Rest"} <= set(jobs["Publish"]["needs"]) diff --git a/policyengine_us/tests/test_batched.py b/policyengine_us/tests/test_batched.py index 314827036cc..240c61bf142 100644 --- a/policyengine_us/tests/test_batched.py +++ b/policyengine_us/tests/test_batched.py @@ -7,16 +7,19 @@ import subprocess import sys import os +import tempfile import gc import time import argparse import re import select +import shlex import threading from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor from concurrent.futures import wait as futures_wait from pathlib import Path from typing import FrozenSet, List, Dict, Optional, Tuple +from xml.etree import ElementTree import yaml @@ -385,40 +388,6 @@ def split_into_batches( batches.append([str(f) for f in root_files]) return batches - # Special handling for household tests. Household totals (net income, - # marginal tax rates, all-program benefits) load most of the model in - # every file, and memory accumulates across the files one subprocess - # runs. The former quarter batches peaked at 15.1 GB on the 16 GB runner - # and ended in "The runner has received a shutdown signal". Measured - # 2026-10-14 (local peak footprint, within 1% of CI's batch peaks): the - # 15 root files peaked at 13.3 GB together and income/household's 8 files - # at 11.4 GB, but no single file exceeded 6.2 GB. Those files run one per - # subprocess; every other subfolder stays one subprocess (at most 7.1 GB, - # income/spm_unit), with income split into its subfolders. - if str(base_path).endswith("baseline/household"): - income = base_path / "income" - one_per_file = {base_path, income / "household"} - batches = [] - - def visit(folder: Path) -> None: - for item in sorted(folder.iterdir()): - if not item.is_dir() or item.name in exclude: - continue - if item == income: - visit(item) - elif item in one_per_file: - batches.extend([str(f)] for f in yaml_files(item)) - else: - batches.extend(subdir_batches(item)) - loose = root_yaml_files(folder) - if folder in one_per_file: - batches.extend([str(f)] for f in loose) - elif loose: - batches.append([str(f) for f in loose]) - - visit(base_path) - return batches - # Special handling for baseline tests if "baseline" in str(base_path) and str(base_path).endswith("baseline"): states_path = base_path / "gov" / "states" @@ -523,7 +492,9 @@ def collect_gov_paths(folder_names): PYTEST_NO_TESTS_COLLECTED = 5 # How long a child gets to die after being asked to, before it is killed. TERMINATE_GRACE_SECONDS = 5 -# Overall budget for one batch, enforced inside the monitoring loop below. +# Budget for one subprocess, enforced inside the monitoring loop below. A +# batch the memory budget continues in fresh subprocesses gets a fresh budget +# for each part (see run_batch). BATCH_TIMEOUT_SECONDS = 1800 # How long the runner may take to hand back a status after printing its # summary. That gap is pytest's own session-finish hooks and unconfigure, not @@ -548,6 +519,86 @@ def collect_gov_paths(folder_names): # is still terminated exactly as promptly as before. BATCH_EXIT_MARKER = "__POLICYENGINE_BATCH_EXIT__" +# Memory accumulates across the YAML files one subprocess runs: household's 15 +# root files peaked at 13.3 GB together (local peak footprint, 2026-10-14, within +# 1% of CI's VmHWM on the batches compared), yet none passed 6.2 GB alone, and +# CI's household and gov/hhs batches each reached 15.1 GB on the 16 GB runner. +# Any fixed batch layout therefore outgrows the runner as tests are added. Instead, after each file the child checks its own peak memory and, +# once past this budget, hands the files it has not started to a fresh +# subprocess (run_batch continues them). The next file starts below the budget, +# so a subprocess peaks at about the budget plus one file's growth. A file that +# alone exceeds the budget still runs whole. Set the variable to override it, +# e.g. 0 for one file per subprocess. +MEMORY_BUDGET_MB = 8000 +MEMORY_BUDGET_ENV = "POLICYENGINE_BATCH_MEMORY_BUDGET_MB" +REMAINDER_ENV = "POLICYENGINE_BATCH_REMAINDER" + +# Child-side half of the memory budget. policyengine-core calls +# pytest.main(..., plugins=[...]) through the pytest module, so wrapping it +# adds this plugin without changing Core. +MEMORY_GUARD = """ +budget_mb = os.environ.get("POLICYENGINE_BATCH_MEMORY_BUDGET_MB") +remainder_path = os.environ.get("POLICYENGINE_BATCH_REMAINDER") +if budget_mb and remainder_path: + import pytest + + def peak_mb(): + try: + with open("/proc/self/status") as status: + for line in status: + if line.startswith("VmHWM:"): + return int(line.split()[1]) / 1024 + except OSError: + pass + if sys.platform == "darwin": + # ru_maxrss leaves out compressed memory, which reads 2-3x low on + # these suites. RUSAGE_INFO_V4 word 30 is the lifetime peak physical + # footprint that /usr/bin/time -l reports. + try: + import ctypes + + info = (ctypes.c_uint64 * 64)() + libc = ctypes.CDLL("/usr/lib/libSystem.B.dylib") + if libc.proc_pid_rusage(os.getpid(), 4, ctypes.byref(info)) == 0: + return info[30] / 2**20 + except (OSError, AttributeError): + pass + import resource + + peak = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss + # ru_maxrss is bytes on macOS and kilobytes on Linux. + return peak / 2**20 if sys.platform == "darwin" else peak / 1024 + + class MemoryBudget: + @pytest.hookimpl(wrapper=True) + def pytest_runtest_protocol(self, item, nextitem): + result = yield + # Only stop between files, so every file runs whole in one process. + if nextitem is None or nextitem.path == item.path: + return result + peak = peak_mb() + if peak < float(budget_mb): + return result + items = item.session.items + remaining = list( + dict.fromkeys(str(later.path) for later in items[items.index(nextitem) :]) + ) + with open(remainder_path, "w") as handle: + handle.write("\\n".join(remaining)) + pytest.exit( + "memory budget: peak %.0f MB after %s; %d files continue in a " + "fresh subprocess" % (peak, item.path.name, len(remaining)), + returncode=1 if item.session.testsfailed else 0, + ) + + run_pytest = pytest.main + + def main(args=None, plugins=None, **kwargs): + return run_pytest(args, [*(plugins or []), MemoryBudget()], **kwargs) + + pytest.main = main +""" + RUNNER_SHIM = f""" import os import runpy @@ -558,7 +609,7 @@ def collect_gov_paths(folder_names): os.environ["PYTEST_ADDOPTS"] = ( os.environ.get("PYTEST_ADDOPTS", "") + " --ignore-glob=*.py" ) - +{MEMORY_GUARD} code = 0 try: runpy.run_module( @@ -627,11 +678,95 @@ def batch_status(returncode: Optional[int]) -> str: return "passed" if returncode == 0 else "failed" +def junit_report_path() -> Optional[str]: + """The --junitxml file PYTEST_ADDOPTS asks every part to write, if any.""" + tokens = shlex.split(os.environ.get("PYTEST_ADDOPTS", "")) + path = None + for i, token in enumerate(tokens): + for flag in ("--junitxml", "--junit-xml"): + if token.startswith(flag + "="): + path = token.split("=", 1)[1] + elif token == flag and i + 1 < len(tokens): + path = tokens[i + 1] + return path + + +def report_mtime(path: Optional[str]) -> Optional[int]: + try: + return os.stat(path).st_mtime_ns if path else None + except FileNotFoundError: + return None + + def run_batch( test_paths: List[str], batch_name: str, stream: bool = True, command: Optional[List[str]] = None, +) -> Dict: + """Run a batch, continuing past the memory budget in fresh subprocesses. + + Each part runs through run_subprocess. A part that reaches + MEMORY_BUDGET_MB between files writes the files it has not started to a + remainder file and exits; they continue in the next part. The batch passes + only if every part passes. Each part rewrites an inherited --junitxml file, + so the parts' suites are merged back into it. + """ + budget = os.environ.get(MEMORY_BUDGET_ENV, str(MEMORY_BUDGET_MB)) + report = junit_report_path() + suites: List[ElementTree.Element] = [] + parts: List[Dict] = [] + paths = test_paths + while paths: + handle, remainder = tempfile.mkstemp(prefix="batch-remainder-") + os.close(handle) + os.unlink(remainder) + name = batch_name if not parts else f"{batch_name} (part {len(parts) + 1})" + env = dict(os.environ, **{MEMORY_BUDGET_ENV: budget, REMAINDER_ENV: remainder}) + written = report_mtime(report) + parts.append(run_subprocess(paths, name, stream, command, env)) + command = None + if report_mtime(report) not in (None, written): + suites.extend(ElementTree.parse(report).getroot().iter("testsuite")) + try: + with open(remainder) as handle_in: + paths = [line for line in handle_in.read().splitlines() if line] + os.unlink(remainder) + except FileNotFoundError: + paths = [] + if paths: + note = ( + f" Memory budget ({budget} MB) reached; continuing " + f"{len(paths)} files in a fresh subprocess" + ) + if stream: + print(note, flush=True) + else: + parts[-1]["output"] = parts[-1]["output"] + note + "\n" + if len(parts) > 1 and suites: + merged = ElementTree.Element("testsuites") + merged.extend(suites) + ElementTree.ElementTree(merged).write( + report, encoding="utf-8", xml_declaration=True + ) + failed = [part for part in parts if part["status"] != "passed"] + peaks = [part["peak_rss_mb"] for part in parts if part["peak_rss_mb"] is not None] + return { + "elapsed": sum(part["elapsed"] for part in parts), + "status": failed[0]["status"] if failed else "passed", + "returncode": failed[0]["returncode"] if failed else 0, + "peak_rss_mb": max(peaks) if peaks else None, + "parts": len(parts), + "output": "".join(part["output"] for part in parts), + } + + +def run_subprocess( + test_paths: List[str], + batch_name: str, + stream: bool = True, + command: Optional[List[str]] = None, + env: Optional[Dict[str, str]] = None, ) -> Dict: """Run a batch of tests in an isolated subprocess. @@ -673,6 +808,7 @@ def emit(text: str, end: str = "\n") -> None: stdout=subprocess.PIPE, stderr=subprocess.STDOUT, bufsize=0, + env=env, ) def sample_rss() -> None: @@ -816,9 +952,9 @@ def overdue() -> Optional[str]: # If the runner never reported, wait for the process with a timeout if reported_returncode is None and process.poll() is None: try: - # Wait up to 30 minutes total + # Wait out the rest of this subprocess's budget elapsed = time.time() - start_time - remaining_timeout = max(1800 - elapsed, 1) + remaining_timeout = max(BATCH_TIMEOUT_SECONDS - elapsed, 1) process.wait(timeout=remaining_timeout) except subprocess.TimeoutExpired: emit(f"\n ⏱️ Timeout - terminating process...") From 073595a093fc3b8e5e48b2f266a89ab06e2ca6ba Mon Sep 17 00:00:00 2001 From: Ziming Date: Thu, 15 Oct 2026 09:41:49 -0400 Subject: [PATCH 4/4] Run Contrib after Baseline and two memory-aware YAML workers per runner Baseline and Contrib each run 6 groups at once, Contrib once Baseline passes, so the Pull request workflow runs at most 11 jobs at a time. A Baseline failure skips Contrib, and re-running the failed group runs Contrib for the first time without repeating finished work. POLICYENGINE_BATCH_WORKERS=2 lets test_batched.py run a second batch subprocess per runner while MemAvailable allows: admission at 6 GB with a 20 s stagger, a 4 GB hand-off floor at file boundaries, and eviction of the newest subprocess below 2 GB with an exclusive rerun. Over three fork runs this used 32% fewer runner-minutes for the same 34,794 cases with no runner shutdowns; a fixed two-process split crashed a runner in both runs. Plain --workers N and local runs are unchanged. Co-Authored-By: Claude Opus 5.5 --- .github/workflows/baseline-tests.yaml | 13 +- .github/workflows/contrib-tests.yaml | 9 +- .github/workflows/pr.yaml | 6 +- .github/workflows/push.yaml | 25 +- Makefile | 8 +- ...equential-suites-household-memory.fixed.md | 2 +- .../test_batched_yaml_collection.py | 336 +++++++++++++++++- .../code_health/test_pr_runner_budget.py | 55 ++- policyengine_us/tests/test_batched.py | 283 ++++++++++++++- 9 files changed, 694 insertions(+), 43 deletions(-) diff --git a/.github/workflows/baseline-tests.yaml b/.github/workflows/baseline-tests.yaml index 785d170d36e..2bf3a9c8b5c 100644 --- a/.github/workflows/baseline-tests.yaml +++ b/.github/workflows/baseline-tests.yaml @@ -11,15 +11,20 @@ jobs: env: PYTHONUNBUFFERED: 1 PYTEST_ADDOPTS: "--durations=25 -p no:unraisableexception -p no:threadexception" + # Two batch subprocesses per runner while its memory allows (see + # test_batched.py): 32% fewer runner-minutes, same cases, no shutdowns. + POLICYENGINE_BATCH_WORKERS: 2 strategy: fail-fast: false - # At most 5 Baseline groups run at once; GitHub starts the next group on a + # At most 6 Baseline groups run at once; GitHub starts the next group on a # fresh runner as each finishes, so the lanes balance from real durations. - max-parallel: 5 + max-parallel: 6 matrix: include: - # One state per subprocess; retain four runners and one worker each. - # The former NY/OH/OK process peaked at 15.1 GB (run 36930199045). + # One state per batch subprocess, over four runners. The job's + # POLICYENGINE_BATCH_WORKERS overrides --workers 1, and runner memory + # decides whether two run at once. The former NY/OH/OK process peaked + # at 15.1 GB (run 36930199045). - group: states-shard-1 cmd: python policyengine_us/tests/test_batched.py policyengine_us/tests/policy/baseline/gov/states --mode per-subdir --workers 1 --shard 1/4 - group: states-shard-2 diff --git a/.github/workflows/contrib-tests.yaml b/.github/workflows/contrib-tests.yaml index 9f4c8c413fc..19ee053f607 100644 --- a/.github/workflows/contrib-tests.yaml +++ b/.github/workflows/contrib-tests.yaml @@ -11,11 +11,14 @@ jobs: env: PYTHONUNBUFFERED: 1 PYTEST_ADDOPTS: "--durations=25 -p no:unraisableexception -p no:threadexception" + # Two batch subprocesses per runner while its memory allows (see + # test_batched.py): 32% fewer runner-minutes, same cases, no shutdowns. + POLICYENGINE_BATCH_WORKERS: 2 strategy: fail-fast: false - # At most 4 Contrib groups run at once; GitHub starts the next group on a - # fresh runner as each finishes, so the lanes balance from real durations. - max-parallel: 4 + # Contrib starts after Baseline finishes, so its 6 lanes take the place of + # Baseline's 6. + max-parallel: 6 matrix: include: - group: states-shard-1 diff --git a/.github/workflows/pr.yaml b/.github/workflows/pr.yaml index a002c0abfcd..a793e99fcb2 100644 --- a/.github/workflows/pr.yaml +++ b/.github/workflows/pr.yaml @@ -273,9 +273,13 @@ jobs: run: bash ./update_itemization.sh - name: Run Household API Partners YAML tests run: uv run --no-sync make test-yaml-no-structural-other-partners + # Contrib starts once Baseline passes, on the same 6 lanes, so most PRs see + # their Baseline result first. A Baseline failure skips Contrib, and + # re-running the failed Baseline group then runs Contrib for the first time: + # no finished Contrib work is repeated. Contrib: name: Contrib tests - needs: ReleaseLock + needs: [ReleaseLock, Baseline] uses: ./.github/workflows/contrib-tests.yaml PythonTests: name: Python tests diff --git a/.github/workflows/push.yaml b/.github/workflows/push.yaml index 8d8dc9e3c25..595efb8d499 100644 --- a/.github/workflows/push.yaml +++ b/.github/workflows/push.yaml @@ -175,15 +175,20 @@ jobs: env: PYTHONUNBUFFERED: 1 PYTEST_ADDOPTS: "--durations=25 -p no:unraisableexception -p no:threadexception" + # Two batch subprocesses per runner while its memory allows (see + # test_batched.py): 32% fewer runner-minutes, same cases, no shutdowns. + POLICYENGINE_BATCH_WORKERS: 2 strategy: fail-fast: false - # At most 5 Baseline groups run at once; GitHub starts the next group on a + # At most 6 Baseline groups run at once; GitHub starts the next group on a # fresh runner as each finishes, so the lanes balance from real durations. - max-parallel: 5 + max-parallel: 6 matrix: include: - # One state per subprocess; retain four runners and one worker each. - # The former NY/OH/OK process peaked at 15.1 GB (run 36930199045). + # One state per batch subprocess, over four runners. The job's + # POLICYENGINE_BATCH_WORKERS overrides --workers 1, and runner memory + # decides whether two run at once. The former NY/OH/OK process peaked + # at 15.1 GB (run 36930199045). - group: states-shard-1 cmd: python policyengine_us/tests/test_batched.py policyengine_us/tests/policy/baseline/gov/states --mode per-subdir --workers 1 --shard 1/4 - group: states-shard-2 @@ -258,7 +263,8 @@ jobs: - name: Run Household API Partners YAML tests run: uv run --no-sync make test-yaml-no-structural-other-partners Contrib: - needs: ReleaseLock + # Waits for Baseline and skips if it fails, as in pr.yaml. + needs: [ReleaseLock, Baseline] name: Full Suite - Contrib (${{ matrix.group }}) runs-on: ubuntu-latest timeout-minutes: 60 @@ -268,11 +274,14 @@ jobs: env: PYTHONUNBUFFERED: 1 PYTEST_ADDOPTS: "--durations=25 -p no:unraisableexception -p no:threadexception" + # Two batch subprocesses per runner while its memory allows (see + # test_batched.py): 32% fewer runner-minutes, same cases, no shutdowns. + POLICYENGINE_BATCH_WORKERS: 2 strategy: fail-fast: false - # At most 4 Contrib groups run at once; GitHub starts the next group on a - # fresh runner as each finishes, so the lanes balance from real durations. - max-parallel: 4 + # Contrib starts after Baseline finishes, so its 6 lanes take the place of + # Baseline's 6. + max-parallel: 6 matrix: include: - group: states-shard-1 diff --git a/Makefile b/Makefile index 19d419af1bd..a3dc2fcbe44 100644 --- a/Makefile +++ b/Makefile @@ -13,6 +13,9 @@ # test growth, and --workers 1 everywhere except the partners target # (its largest batch measured 2.5 GB, so two-wide stays trivially safe). BATCH := python policyengine_us/tests/test_batched.py +# CI's Baseline and Contrib jobs set POLICYENGINE_BATCH_WORKERS=2, which +# overrides --workers below: test_batched.py then runs two batch subprocesses +# per runner while its memory allows. Local runs keep the counts given here. TESTS := policyengine_us/tests # Run SPM construction/isolation tests in their own process on the existing # Microsimulation CI runner, before the dataset tests. Rest excludes these @@ -197,8 +200,9 @@ test-yaml-variables: $(BATCH) $(TESTS)/variables --batches 1 test-yaml-no-structural-states: # Release model/parameter caches between states. The former NY/OH/OK - # group peaked at 15.1 GB; keep one subprocess at a time on each of - # the existing four CI runners. Root-level state tests remain included. + # group peaked at 15.1 GB; keep one subprocess at a time here. CI's four + # runners run up to two while memory allows (POLICYENGINE_BATCH_WORKERS). + # Root-level state tests remain included. $(BATCH) $(TESTS)/policy/baseline/gov/states --mode per-subdir --workers 1 test-yaml-no-structural-other: $(BATCH) $(TESTS)/policy/baseline --batches 2 --exclude states diff --git a/changelog.d/ci-sequential-suites-household-memory.fixed.md b/changelog.d/ci-sequential-suites-household-memory.fixed.md index 62e022227a3..4445a3eb2d9 100644 --- a/changelog.d/ci-sequential-suites-household-memory.fixed.md +++ b/changelog.d/ci-sequential-suites-household-memory.fixed.md @@ -1 +1 @@ -Stop the second Rest shard when a CI run is cancelled, run at most 5 Baseline and 4 Contrib groups at once so the Pull request workflow runs at most 14 jobs at a time instead of about 25, and continue YAML test batches in a fresh subprocess once one passes an 8 GB memory budget, so batches no longer need re-splitting by hand as tests grow. +Stop the second Rest shard when a CI run is cancelled; run at most 6 Baseline groups at once and start Contrib (also at most 6) once Baseline passes, so the Pull request workflow runs at most 11 jobs at a time instead of about 25; run two YAML batch subprocesses per Baseline and Contrib runner while its free memory allows, cutting their runner time by about a third; and continue YAML test batches in a fresh subprocess once one passes an 8 GB memory budget, so batches no longer need re-splitting by hand as tests grow. diff --git a/policyengine_us/tests/code_health/test_batched_yaml_collection.py b/policyengine_us/tests/code_health/test_batched_yaml_collection.py index f48219dd780..f1a0517727e 100644 --- a/policyengine_us/tests/code_health/test_batched_yaml_collection.py +++ b/policyengine_us/tests/code_health/test_batched_yaml_collection.py @@ -1,16 +1,22 @@ -"""YAML batches must not rerun the separately routed Python suites, and must -continue in a fresh subprocess once one passes the memory budget. +"""YAML batches must not rerun the separately routed Python suites, must +continue in a fresh subprocess once one passes the memory budget, and share a +runner only while its memory allows. Exercise the real Core CLI and pytest collector with a one-variable country package, so these runner regressions never construct the full US model. +Admission and eviction run against a scripted memory reading and stand-in +processes. """ import importlib.util +import io import os import re import shlex import subprocess import sys +import threading +import time import xml.etree.ElementTree as ET from pathlib import Path @@ -194,3 +200,329 @@ def test_a_failure_before_the_budget_stop_fails_the_batch_but_runs_the_rest( "2 passed", "2 passed", ], result["output"] + + +WAIT = 5 + + +@pytest.fixture +def memory(monkeypatch): + """The runner's available memory, set by each test.""" + reading = {"mb": 16000.0} + monkeypatch.setattr(batched, "available_mb", lambda: reading["mb"]) + monkeypatch.setattr(batched, "ADMIT_STAGGER_SECONDS", 0) + return reading + + +@pytest.fixture +def admission(): + gate = batched.Admission() + yield gate + gate.stopped = True + + +class Process: + def __init__(self): + self.killed = threading.Event() + + def kill(self): + self.killed.set() + + +class Worker: + """A thread holding one admission slot, rerunning alone if stopped. + + As in run_batch, a stopped part's own thread asks for the exclusive slot. + """ + + def __init__(self, admission, process=None): + self.admitted = threading.Event() + self.go = threading.Event() + self.released = threading.Event() + self.rerun_start = threading.Event() + self.rerun_admitted = threading.Event() + self.rerun_go = threading.Event() + self.evicted = None + + def run(): + admission.acquire(False) + if process is not None: + admission.started(process) + self.admitted.set() + self.go.wait(WAIT * 4) + self.evicted = admission.release() + self.released.set() + if self.evicted: + self.rerun_start.wait(WAIT * 4) + admission.acquire(True) + self.rerun_admitted.set() + self.rerun_go.wait(WAIT * 4) + admission.release() + + threading.Thread(target=run, daemon=True).start() + + +def test_parent_reads_the_runners_available_memory(monkeypatch): + def reading(text): + monkeypatch.setattr( + batched, "open", lambda *args, **kwargs: io.StringIO(text), raising=False + ) + return batched.available_mb() + + assert reading("MemTotal: 16384000 kB\nMemAvailable: 8192000 kB\n") == 8000.0 + assert reading("MemTotal: 16384000 kB\n") is None + assert reading("MemAvailable: junk kB\n") is None + + def missing(*args, **kwargs): + raise OSError + + monkeypatch.setattr(batched, "open", missing, raising=False) + assert batched.available_mb() is None + + +def test_one_subprocess_at_a_time_when_memory_cannot_be_read(admission, memory): + memory["mb"] = None + admission.acquire(False) + second = Worker(admission) + assert not second.admitted.wait(1) + admission.release() + assert second.admitted.wait(WAIT) + second.go.set() + + +def test_another_subprocess_starts_only_with_memory_to_spare(admission, memory): + admission.acquire(False) + memory["mb"] = batched.ADMIT_MB - 1 + second = Worker(admission) + assert not second.admitted.wait(1.5) + memory["mb"] = batched.ADMIT_MB + assert second.admitted.wait(WAIT) + assert admission.waits == 1 + second.go.set() + + +def test_subprocesses_start_one_stagger_apart(admission, memory, monkeypatch): + monkeypatch.setattr(batched, "ADMIT_STAGGER_SECONDS", 1.5) + admission.acquire(False) + second = Worker(admission) + assert not second.admitted.wait(0.75) + assert second.admitted.wait(WAIT) + second.go.set() + + +def test_low_memory_stops_the_newest_and_admits_nothing_until_its_rerun( + admission, memory +): + older_process, newer_process = Process(), Process() + older = Worker(admission, process=older_process) + assert older.admitted.wait(WAIT) + time.sleep(0.05) + newer = Worker(admission, process=newer_process) + assert newer.admitted.wait(WAIT) + + memory["mb"] = batched.EVICT_MB - 1 + assert newer_process.killed.wait(WAIT) + assert not older_process.killed.is_set() + memory["mb"] = 16000.0 + newer.go.set() + assert newer.released.wait(WAIT) + assert newer.evicted is True + + # Nothing new starts between the stop and the rerun, nor during the rerun. + other = Worker(admission) + assert not other.admitted.wait(1) + newer.rerun_start.set() + assert not newer.rerun_admitted.wait(1) # the older one is still running + older.go.set() + assert newer.rerun_admitted.wait(WAIT) + assert not other.admitted.wait(1) + newer.rerun_go.set() + assert other.admitted.wait(WAIT) + other.go.set() + assert older.evicted is False + assert admission.evictions == 1 + + +def test_a_lone_subprocess_is_never_stopped(admission, memory): + process = Process() + admission.acquire(False) + admission.started(process) + memory["mb"] = 1.0 + assert not process.killed.wait(1) + assert admission.release() is False + + +def test_a_subprocess_that_has_reported_is_not_stopped(admission, memory): + older, newer = Process(), Process() + admission.acquire(False) + admission.started(older) + + def report(): + admission.acquire(False) + admission.started(newer) + admission.finished() + + thread = threading.Thread(target=report, daemon=True) + thread.start() + thread.join(WAIT) + memory["mb"] = 1.0 + assert not newer.killed.wait(1) + assert not older.killed.is_set() + + +def test_split_into_one_batch_per_worker_keeps_whole_files(tmp_path): + folder = tmp_path / "suite" + (folder / "nested").mkdir(parents=True) + files = [folder / "a.yaml", folder / "b.yaml", folder / "nested" / "c.yml"] + for file in files: + file.write_text("[]\n") + batches = batched.fill_workers([[str(folder)]], 2) + assert sorted(len(batch) for batch in batches) == [1, 2] + assert sorted(path for batch in batches for path in batch) == sorted( + str(file) for file in files + ) + two = [str(files[0]), str(files[1])] + assert sorted(len(batch) for batch in batched.fill_workers([two], 2)) == [1, 1] + assert batched.fill_workers([[str(files[0])]], 2) == [[str(files[0])]] + + +def test_split_keeps_no_batches_and_a_folder_without_yaml(tmp_path): + assert batched.fill_workers([], 2) == [] + empty = tmp_path / "python-only" + empty.mkdir() + (empty / "test_x.py").write_text("") + # Kept as the folder, so the subprocess still exits 5 (no tests collected) + # and fails the batch, as it does with one worker. + assert batched.fill_workers([[str(empty)]], 2) == [[str(empty)]] + + +class Gate: + """Stands in for Admission in run_batch; release() reports evictions.""" + + def __init__(self, evictions=()): + self.calls = [] + self.processes = [] + self.evictions = list(evictions) + self.forgotten = 0 + + def acquire(self, exclusive): + self.calls.append(exclusive) + + def started(self, process): + self.processes.append(process) + + def finished(self): + pass + + def forget(self): + self.forgotten = self.forgotten + 1 + + def release(self): + return self.evictions.pop(0) if self.evictions else False + + +def _scripted(flag, first_run_reports): + """A part that is stopped (or reports) on its first run, passes after.""" + first = ( + f" print({batched.BATCH_EXIT_MARKER!r} + ' 0', flush=True)\n" + if first_run_reports + else " os.kill(os.getpid(), signal.SIGKILL)\n" + ) + return [ + sys.executable, + "-c", + "import os, pathlib, signal, time\n" + f"flag = pathlib.Path({str(flag)!r})\n" + "if not flag.exists():\n" + " flag.touch()\n" + " print('==== 1 passed in 0.10s ====', flush=True)\n" + + first + + " time.sleep(30)\n" + "print('==== 2 passed in 0.10s ====')\n" + f"print({batched.BATCH_EXIT_MARKER!r} + ' 0', flush=True)\n", + ] + + +def test_an_evicted_part_reruns_alone_and_counts_once(small_country, monkeypatch): + gate = Gate(evictions=[True]) + monkeypatch.setattr(batched, "ADMISSION", gate) + command = _scripted(small_country / "first-run", first_run_reports=False) + result = batched.run_batch(["a.yaml"], "evicted", stream=False, command=command) + assert gate.calls == [False, True] + assert len(gate.processes) == 2 + assert all(hasattr(process, "kill") for process in gate.processes) + assert result["status"] == "passed", result["output"] + assert result["parts"] == 1 + assert "rerunning its 1 files alone" in result["output"] + + +def test_a_part_stopped_after_it_reported_keeps_its_result(small_country, monkeypatch): + # The watchdog can pick a part just as it reports; its result stands and + # its rerun reservation is dropped, so admission is not held forever. + gate = Gate(evictions=[True]) + monkeypatch.setattr(batched, "ADMISSION", gate) + command = _scripted(small_country / "first-run", first_run_reports=True) + result = batched.run_batch(["a.yaml"], "reported", stream=False, command=command) + assert gate.calls == [False] + assert gate.forgotten == 1 + assert result["parts"] == 1 + assert "rerunning" not in result["output"] + + +def _run_batcher(small_country, *args, workers=None): + # The runner checks it is started from the repository root. + (small_country.parent / "policyengine_us").mkdir(exist_ok=True) + env = dict(os.environ) + if workers is not None: + env[batched.WORKERS_ENV] = workers + run = subprocess.run( + [sys.executable, str(BATCHER), str(small_country), *args], + capture_output=True, + text=True, + timeout=600, + env=env, + cwd=small_country.parent, + ) + return run.returncode, run.stdout + run.stderr + + +def test_workers_variable_runs_every_case_and_keeps_failures(small_country): + for name in ("a", "b", "c"): + _write_age_cases( + small_country / f"{name}.yaml", 2, failing=1 if name == "b" else 0 + ) + code, output = _run_batcher(small_country, "--batches", "1", workers="2") + assert code == 1, output + assert "Workers: 2" in output + assert "Split into 2 batches for 2 workers" in output + assert "[memory] admission waits:" in output + passed = sum(int(n) for n in re.findall(r"(\d+) passed", output)) + failed = sum(int(n) for n in re.findall(r"(\d+) failed", output)) + assert (passed, failed) == (5, 1), output + + +def test_plain_workers_flag_keeps_running_its_batches_at_once(small_country): + for name in ("a", "b"): + (small_country / name).mkdir() + _write_age_cases(small_country / name / "x.yaml", 1) + code, output = _run_batcher(small_country, "--mode", "per-subdir", "--workers", "2") + assert code == 0, output + assert "Running 2 batches with 2 workers" in output + assert "[memory]" not in output + assert "Split into" not in output + + +@pytest.mark.skipif( + not Path("/proc/meminfo").exists(), reason="reads the runner's MemAvailable" +) +def test_low_runner_memory_hands_off_at_file_boundaries(small_country, monkeypatch): + for name in ("a", "b", "c"): + _write_age_cases(small_country / f"{name}.yaml", 2) + monkeypatch.setattr(batched, "ADMISSION", Gate()) + monkeypatch.setattr(batched, "FLOOR_MB", 10**9) + files = sorted(str(f) for f in small_country.glob("*.yaml")) + result = batched.run_batch(files, "floor", stream=False) + assert result["status"] == "passed", result["output"] + assert result["parts"] == 3 + assert "runner available" in result["output"] + assert "or runner floor" in result["output"] diff --git a/policyengine_us/tests/code_health/test_pr_runner_budget.py b/policyengine_us/tests/code_health/test_pr_runner_budget.py index 209f69321d0..1378ebf57b0 100644 --- a/policyengine_us/tests/code_health/test_pr_runner_budget.py +++ b/policyengine_us/tests/code_health/test_pr_runner_budget.py @@ -68,22 +68,23 @@ def needs(job): return [value] if isinstance(value, str) else value -def test_lanes_cap_a_pr_at_14_runners_at_once(): +def concurrent(job): + return min(job.get("strategy", {}).get("max-parallel", math.inf), runners(job)) + + +def test_lanes_cap_a_pr_at_11_runners_at_once(): # Within pr.yaml, validation runs alone; everything that needs only it starts - # together, and Rest waits for RestPython. Baseline and Contrib each keep - # one job per group but run at most max-parallel of them at once: GitHub - # starts the next group as one finishes, so the lanes rebalance from real - # durations on every run, with nothing to re-pair as groups grow. + # together, Rest waits for RestPython, and Contrib waits for Baseline. + # Baseline and Contrib each keep one job per group but run at most + # max-parallel of them at once: GitHub starts the next group as one + # finishes, so the lanes rebalance from real durations on every run, with + # nothing to re-pair as groups grow. assert ( - GROUP_WORKFLOWS["Baseline"]["jobs"]["Baseline"]["strategy"]["max-parallel"] == 5 + GROUP_WORKFLOWS["Baseline"]["jobs"]["Baseline"]["strategy"]["max-parallel"] == 6 ) assert ( - GROUP_WORKFLOWS["Contrib"]["jobs"]["Contrib"]["strategy"]["max-parallel"] == 4 + GROUP_WORKFLOWS["Contrib"]["jobs"]["Contrib"]["strategy"]["max-parallel"] == 6 ) - - def concurrent(job): - return min(job.get("strategy", {}).get("max-parallel", math.inf), runners(job)) - first_wave = {} for name, job in JOBS.items(): if needs(job) != ["ReleaseLock"]: @@ -94,8 +95,31 @@ def concurrent(job): for inner_name, inner in GROUP_WORKFLOWS[name]["jobs"].items(): if not needs(inner): first_wave[inner_name] = concurrent(inner) - assert "Rest" not in first_wave - assert sum(first_wave.values()) == 14 + assert "Rest" not in first_wave and "Contrib" not in first_wave + assert sum(first_wave.values()) == 11 + # Contrib starts only after Baseline has finished, so at most Contrib's + # lanes plus every other first-wave job can be running then. + single = {name: n for name, n in first_wave.items() if name != "Baseline"} + contrib = concurrent(GROUP_WORKFLOWS["Contrib"]["jobs"]["Contrib"]) + assert contrib + sum(single.values()) <= 11 + + +def test_contrib_waits_for_baseline_and_skips_when_it_fails(): + # A plain needs (no if, or an if without a status function) skips Contrib + # when Baseline fails or is cancelled. Re-running the failed Baseline group + # then runs Contrib for the first time (fork run 38687143315); an if that + # let Contrib run anyway would rerun all of its finished groups. + assert needs(JOBS["Contrib"]) == ["ReleaseLock", "Baseline"] + assert "if" not in JOBS["Contrib"] + assert needs(PUSH_JOBS["Contrib"]) == ["ReleaseLock", "Baseline"] + for status in ("always()", "cancelled()", "failure()", "success()"): + assert status not in PUSH_JOBS["Contrib"]["if"] + + +def test_yaml_suites_run_two_memory_aware_workers_per_runner(): + for name in ("Baseline", "Contrib"): + for job in (GROUP_WORKFLOWS[name]["jobs"][name], PUSH_JOBS[name]): + assert job["env"]["POLICYENGINE_BATCH_WORKERS"] == 2 def test_release_ci_runs_the_same_baseline_and_contrib_lanes(): @@ -103,6 +127,7 @@ def test_release_ci_runs_the_same_baseline_and_contrib_lanes(): pr_job = GROUP_WORKFLOWS[name]["jobs"][name] assert PUSH_JOBS[name]["strategy"] == pr_job["strategy"] assert PUSH_JOBS[name]["steps"] == pr_job["steps"] + assert PUSH_JOBS[name]["env"] == pr_job["env"] def test_no_test_job_keeps_running_after_cancellation(): @@ -129,7 +154,9 @@ def test_suite_groups_have_no_extra_runners_or_triggers(): for name, (label, file) in groups.items(): assert JOBS[name] == { "name": label, - "needs": "ReleaseLock", + "needs": ["ReleaseLock", "Baseline"] + if name == "Contrib" + else "ReleaseLock", "uses": f"./.github/workflows/{file}", } # PyYAML parses the bare `on` key as True. No push/PR trigger here: diff --git a/policyengine_us/tests/test_batched.py b/policyengine_us/tests/test_batched.py index 240c61bf142..099c5867cfb 100644 --- a/policyengine_us/tests/test_batched.py +++ b/policyengine_us/tests/test_batched.py @@ -533,15 +533,69 @@ def collect_gov_paths(folder_names): MEMORY_BUDGET_ENV = "POLICYENGINE_BATCH_MEMORY_BUDGET_MB" REMAINDER_ENV = "POLICYENGINE_BATCH_REMAINDER" +# POLICYENGINE_BATCH_WORKERS lets several batch subprocesses share a runner +# while its memory allows; it overrides --workers. (A plain --workers N, as +# the partners target uses for its small batches, still runs N at once.) A +# 16 GB, 4-CPU CI runner running one subprocess is about 25% busy, but two +# fixed subprocesses ran it out of memory (fork runs 38687143268 and +# 38689387981: a runner shutdown in the crfb group in each, and 77-296 MB +# left in the contrib states groups). So with the variable above 1, memory +# decides how many run: +# - another subprocess starts only while the runner has ADMIT_MB available, +# and at most one starts per ADMIT_STAGGER_SECONDS, since a new one takes +# that long to load the model; one subprocess may always run; +# - a subprocess hands off its unstarted files at a file boundary (as for the +# memory budget) once the runner has less than FLOOR_MB available, and the +# remainder waits for admission; +# - if the runner falls below EVICT_MB mid-file, the newest of two or more +# subprocesses is stopped and its files rerun alone once the others finish; +# - batches are split, whole files only, until there is one per worker. +# Where available memory cannot be read (off Linux), one subprocess runs at a +# time. With two workers the 20 Baseline and Contrib groups averaged 345 +# runner-minutes per fork run against 508 with one (32% fewer over three runs; +# 33% over 38689387981 and 38710504296, which ran this eviction protocol), +# passed the same 34,794 cases, and never ran a runner out of memory. Three or +# four workers were no faster, because memory, not CPU, runs out first. +WORKERS_ENV = "POLICYENGINE_BATCH_WORKERS" +FLOOR_ENV = "POLICYENGINE_BATCH_MEMORY_FLOOR_MB" +ADMIT_MB = 6000 +FLOOR_MB = 4000 +EVICT_MB = 2000 +ADMIT_STAGGER_SECONDS = 20 + + +def available_mb() -> Optional[float]: + """The runner's MemAvailable in MB, or None off Linux.""" + try: + with open("/proc/meminfo") as meminfo: + for line in meminfo: + if line.startswith("MemAvailable:"): + return int(line.split()[1]) / 1024 + except (OSError, ValueError, IndexError): + return None + return None + + # Child-side half of the memory budget. policyengine-core calls # pytest.main(..., plugins=[...]) through the pytest module, so wrapping it # adds this plugin without changing Core. MEMORY_GUARD = """ budget_mb = os.environ.get("POLICYENGINE_BATCH_MEMORY_BUDGET_MB") remainder_path = os.environ.get("POLICYENGINE_BATCH_REMAINDER") +floor_mb = os.environ.get("POLICYENGINE_BATCH_MEMORY_FLOOR_MB") if budget_mb and remainder_path: import pytest + def runner_available_mb(): + try: + with open("/proc/meminfo") as meminfo: + for line in meminfo: + if line.startswith("MemAvailable:"): + return int(line.split()[1]) / 1024 + except OSError: + pass + return None + def peak_mb(): try: with open("/proc/self/status") as status: @@ -577,7 +631,10 @@ def pytest_runtest_protocol(self, item, nextitem): if nextitem is None or nextitem.path == item.path: return result peak = peak_mb() - if peak < float(budget_mb): + low = runner_available_mb() if floor_mb else None + if low is not None and low >= float(floor_mb): + low = None + if peak < float(budget_mb) and low is None: return result items = item.session.items remaining = list( @@ -585,9 +642,14 @@ def pytest_runtest_protocol(self, item, nextitem): ) with open(remainder_path, "w") as handle: handle.write("\\n".join(remaining)) + reason = ( + "peak %.0f MB" % peak + if low is None + else "runner available %.0f MB (own peak %.0f MB)" % (low, peak) + ) pytest.exit( - "memory budget: peak %.0f MB after %s; %d files continue in a " - "fresh subprocess" % (peak, item.path.name, len(remaining)), + "memory budget: %s after %s; %d files continue in a " + "fresh subprocess" % (reason, item.path.name, len(remaining)), returncode=1 if item.session.testsfailed else 0, ) @@ -698,6 +760,147 @@ def report_mtime(path: Optional[str]) -> Optional[int]: return None +class Admission: + """Starts batch subprocesses only while the runner has memory for them. + + One subprocess may always run, so a run never stalls. An evicted part + reruns exclusively: no new subprocess starts while it waits or runs. + """ + + def __init__(self, admit_mb: float = ADMIT_MB, evict_mb: float = EVICT_MB): + self.admit_mb = admit_mb + self.evict_mb = evict_mb + self.cond = threading.Condition() + self.running: Dict[int, list] = {} + self.evicted: set = set() + self.pending_exclusive: set = set() + self.exclusive_waiting = 0 + self.exclusive_holder: Optional[int] = None + self.last_admit = 0.0 + self.waits = 0 + self.evictions = 0 + self.stopped = False + threading.Thread(target=self.watch, daemon=True).start() + + def may_admit(self) -> bool: + if ( + self.exclusive_holder is not None + or self.exclusive_waiting + or self.pending_exclusive + ): + return False + if not self.running: + return True + if time.time() - self.last_admit < ADMIT_STAGGER_SECONDS: + return False + avail = available_mb() + return avail is not None and avail >= self.admit_mb + + def acquire(self, exclusive: bool) -> None: + ident = threading.get_ident() + with self.cond: + if exclusive: + self.exclusive_waiting = self.exclusive_waiting + 1 + self.pending_exclusive.discard(ident) + while self.running or self.exclusive_holder is not None: + self.cond.wait(1) + self.exclusive_waiting = self.exclusive_waiting - 1 + self.exclusive_holder = ident + elif not self.may_admit(): + self.waits = self.waits + 1 + while not self.may_admit(): + self.cond.wait(1) + self.running[ident] = [None, time.time()] + self.last_admit = time.time() + + def started(self, process) -> None: + with self.cond: + entry = self.running.get(threading.get_ident()) + if entry is not None: + entry[0] = process + + def finished(self) -> None: + """The part reported or exited: stop treating it as live.""" + self.started(None) + + def forget(self) -> None: + """Drop this thread's rerun reservation, if it will not rerun.""" + with self.cond: + self.pending_exclusive.discard(threading.get_ident()) + self.cond.notify_all() + + def release(self) -> bool: + ident = threading.get_ident() + with self.cond: + self.running.pop(ident, None) + evicted = ident in self.evicted + self.evicted.discard(ident) + if self.exclusive_holder == ident: + self.exclusive_holder = None + self.cond.notify_all() + return evicted + + def watch(self) -> None: + while not self.stopped: + time.sleep(0.2) + avail = available_mb() + if avail is None or avail >= self.evict_mb: + continue + with self.cond: + live = [ + (entry[1], ident, entry[0]) + for ident, entry in self.running.items() + if entry[0] is not None and ident not in self.evicted + ] + if len(live) < 2: + continue + _, ident, process = max(live, key=lambda item: item[0]) + self.evicted.add(ident) + self.pending_exclusive.add(ident) + self.evictions = self.evictions + 1 + print( + f"[memory] runner available {avail:.0f} MB < {self.evict_mb:.0f} MB " + f"with {len(live)} subprocesses; stopping the newest and " + "rerunning its files alone", + flush=True, + ) + try: + process.kill() + except OSError: + pass + time.sleep(2) + + +ADMISSION: Optional[Admission] = None + + +def fill_workers(batches: List[List[str]], workers: int) -> List[List[str]]: + """Split batches, whole files only, until there is one per worker.""" + if not batches: + return batches + expanded = [] + for batch in batches: + files = [ + str(f) + for p in batch + for f in (yaml_files(Path(p)) if Path(p).is_dir() else [Path(p)]) + ] + expanded.append(files or list(batch)) + splittable = [ + len(b) > 1 and all(p.endswith(YAML_SUFFIXES) for p in b) for b in expanded + ] + while len(expanded) < workers and any(splittable): + largest = max( + (i for i in range(len(expanded)) if splittable[i]), + key=lambda i: len(expanded[i]), + ) + files = expanded[largest] + half = len(files) // 2 + expanded[largest : largest + 1] = [files[:half], files[half:]] + splittable[largest : largest + 1] = [half > 1, len(files) - half > 1] + return expanded + + def run_batch( test_paths: List[str], batch_name: str, @@ -707,8 +910,10 @@ def run_batch( """Run a batch, continuing past the memory budget in fresh subprocesses. Each part runs through run_subprocess. A part that reaches - MEMORY_BUDGET_MB between files writes the files it has not started to a - remainder file and exits; they continue in the next part. The batch passes + MEMORY_BUDGET_MB between files (or, when subprocesses share the runner, + finds less than FLOOR_MB available) writes the files it has not started to + a remainder file and exits; they continue in the next part. A part stopped + mid-file for runner memory is discarded and rerun alone. The batch passes only if every part passes. Each part rewrites an inherited --junitxml file, so the parts' suites are merged back into it. """ @@ -717,14 +922,40 @@ def run_batch( suites: List[ElementTree.Element] = [] parts: List[Dict] = [] paths = test_paths + exclusive = False + discarded = "" while paths: handle, remainder = tempfile.mkstemp(prefix="batch-remainder-") os.close(handle) os.unlink(remainder) name = batch_name if not parts else f"{batch_name} (part {len(parts) + 1})" env = dict(os.environ, **{MEMORY_BUDGET_ENV: budget, REMAINDER_ENV: remainder}) + if ADMISSION is not None: + env[FLOOR_ENV] = str(FLOOR_MB) written = report_mtime(report) - parts.append(run_subprocess(paths, name, stream, command, env)) + if ADMISSION is not None: + ADMISSION.acquire(exclusive) + try: + part = run_subprocess(paths, name, stream, command, env) + finally: + evicted = ADMISSION.release() if ADMISSION is not None else False + if evicted and not part.get("reported"): + note = ( + f" [memory] {name} stopped mid-file for runner memory; " + f"rerunning its {batch_cost(paths)} files alone" + ) + if stream: + print(note, flush=True) + discarded = discarded + part["output"] + note + "\n" + if os.path.exists(remainder): + os.unlink(remainder) + exclusive = True + continue + if evicted: + # Stopped only after it reported: keep its result. + ADMISSION.forget() + exclusive = False + parts.append(part) command = None if report_mtime(report) not in (None, written): suites.extend(ElementTree.parse(report).getroot().iter("testsuite")) @@ -735,8 +966,11 @@ def run_batch( except FileNotFoundError: paths = [] if paths: + limit = f"Memory budget ({budget} MB)" + if ADMISSION is not None: + limit = f"{limit} or runner floor ({FLOOR_MB} MB available)" note = ( - f" Memory budget ({budget} MB) reached; continuing " + f" {limit} reached; continuing " f"{len(paths)} files in a fresh subprocess" ) if stream: @@ -757,7 +991,7 @@ def run_batch( "returncode": failed[0]["returncode"] if failed else 0, "peak_rss_mb": max(peaks) if peaks else None, "parts": len(parts), - "output": "".join(part["output"] for part in parts), + "output": discarded + "".join(part["output"] for part in parts), } @@ -810,6 +1044,8 @@ def emit(text: str, end: str = "\n") -> None: bufsize=0, env=env, ) + if ADMISSION is not None: + ADMISSION.started(process) def sample_rss() -> None: # VmHWM is monotonic, so keeping only the latest reading preserves @@ -867,6 +1103,8 @@ def overdue() -> Optional[str]: poll_result = process.poll() if poll_result is not None: # Process terminated + if ADMISSION is not None: + ADMISSION.finished() break give_up = overdue() if give_up is not None: @@ -938,6 +1176,8 @@ def overdue() -> Optional[str]: marker_match = BATCH_EXIT_PATTERN.search(output_text) if marker_match: reported_returncode = int(marker_match.group(1)) + if ADMISSION is not None: + ADMISSION.finished() emit( "\n Runner reported " f"{describe_exit_status(reported_returncode)}; " @@ -997,6 +1237,7 @@ def overdue() -> Optional[str]: } status = batch_status(returncode) + reported = reported_returncode is not None if test_completed and summary_reported_pass and status != "passed": # The case that used to be reported as a pass: pytest printed no # failure count (e.g. "1 passed, 1 error") yet finished nonzero. @@ -1016,6 +1257,7 @@ def overdue() -> Optional[str]: "returncode": returncode, "peak_rss_mb": peak_rss_mb, "output": "".join(buf), + "reported": reported, } except Exception as e: @@ -1069,9 +1311,18 @@ def run_one(idx: int, batch_paths: List[str]) -> Dict: name = f"Batch {idx}" with in_flight_lock: in_flight[name] = time.time() + # Output is buffered until the batch ends, so record where each batch + # starts: a job timeout or lost runner otherwise leaves no trace of it. + print( + f"[Batch {idx}/{total}] started: {batch_cost(batch_paths)} files " + f"from {batch_paths[0]}", + flush=True, + ) try: return run_batch(batch_paths, name, stream=False) finally: + if ADMISSION is not None: + ADMISSION.forget() with in_flight_lock: in_flight.pop(name, None) @@ -1179,6 +1430,9 @@ def main(): ) args = parser.parse_args() + scheduled = bool(os.environ.get(WORKERS_ENV)) + if scheduled: + args.workers = int(os.environ[WORKERS_ENV]) exclude_list = [x.strip() for x in args.exclude.split(",") if x.strip()] shard_idx = shard_count = None @@ -1235,6 +1489,13 @@ def main(): f"({len(batches)} of {all_batch_count} batches)" ) + global ADMISSION + if scheduled and args.workers > 1: + if len(batches) < args.workers: + batches = fill_workers(batches, args.workers) + print(f"Split into {len(batches)} batches for {args.workers} workers") + ADMISSION = Admission() + if len(batches) != args.batches: print(f"Actual batches: {len(batches)} (optimized for {total_tests} files)") else: @@ -1277,6 +1538,12 @@ def main(): f"{result['status']}" ) print(f"Workers: {args.workers}") + if ADMISSION is not None: + ADMISSION.stopped = True + print( + f"[memory] admission waits: {ADMISSION.waits}, " + f"evictions: {ADMISSION.evictions}" + ) print(f"Total time: {total_elapsed:.1f}s") if args.workers > 1: print(f"Wall time: {time.time() - run_start:.1f}s")