Stop Rest after cancellation, run Contrib after Baseline at 11 jobs, and share YAML runners by memory

.github/workflows/baseline-tests.yaml

@@ -11,12 +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 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: 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

.github/workflows/contrib-tests.yaml

@@ -11,8 +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
+      # 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

.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

.github/workflows/push.yaml

@@ -175,12 +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 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: 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
@@ -255,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
@@ -265,8 +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
+      # 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
@@ -380,8 +395,12 @@ jobs:
     # Historical group times: 32-46 minutes plus setup.
     # This runner also retains the legacy summary check, keeping PRs at 27 runners.
     timeout-minutes: 90
+    # 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: |
-      always()
+      !cancelled()
+      && (needs.RestPython.result != 'skipped')
       && (github.repository == 'PolicyEngine/policyengine-us')
       && (github.event.head_commit.message == 'Update PolicyEngine US')
     env:

.github/workflows/python-tests.yaml

@@ -66,7 +66,9 @@ 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.

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
@@ -212,8 +216,9 @@ test-yaml-no-structural-other-irs:
 	$(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.
+	# 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.
@@ -234,8 +239,9 @@ 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.
+	# 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

changelog.d/ci-sequential-suites-household-memory.fixed.md

@@ -0,0 +1 @@
+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.