Plan the Rest Python tests over parallel CI shards from recent timings

.github/workflows/push.yaml

@@ -330,23 +330,19 @@ jobs:
       - name: Run Contrib Python tests
         if: ${{ matrix.group == 'other-shard-1' }}
         run: uv run --no-sync make test-policy-contrib-python
-  RestPython:
+  RestPlan:
     needs: ReleaseLock
-    name: Full Suite - Rest (Python, shard-1 + variables)
+    name: Plan Rest Python shards
     runs-on: ubuntu-latest
-    # Historical group times, including variables: 33-46 minutes plus setup.
-    timeout-minutes: 90
+    timeout-minutes: 10
     if: |
       (github.repository == 'PolicyEngine/policyengine-us')
       && (github.event.head_commit.message == 'Update PolicyEngine US')
-    env:
-      PYTHONUNBUFFERED: 1
-      OMP_NUM_THREADS: 4
-      OPENBLAS_NUM_THREADS: 4
-      MKL_NUM_THREADS: 4
-      VECLIB_MAXIMUM_THREADS: 4
-      NUMEXPR_NUM_THREADS: 4
-      OMP_WAIT_POLICY: PASSIVE
+    permissions:
+      actions: read
+      contents: read
+    outputs:
+      matrix: ${{ steps.plan.outputs.matrix }}
     steps:
       - name: Checkout repo
         uses: actions/checkout@v6
@@ -361,56 +357,51 @@ jobs:
         uses: astral-sh/setup-uv@v8.1.0
         with:
           version: "0.13.4"
-      - name: Resolve latest compatible PolicyEngine Core
-        run: uv lock --upgrade-package policyengine-core
       - name: Install dependencies
         run: uv sync --locked --extra dev
-      - name: Turn off default branching
+      # A successful release run on main has every shard's JUnit and GNU time
+      # reports; a failed one can lack a lost runner's, which would skew the
+      # estimates. Without any, the plan counts each file the same.
+      - name: Download recent Rest Python timings
+        id: timings
         shell: bash
-        run: bash ./update_itemization.sh
-      # Each group remains a separate, sequential process on its runner.
-      # SPM simulation/API contracts stay in the Microsimulation job.
-      - name: Run remaining Python tests
-        if: ${{ !cancelled() && steps.system-ca.outcome == 'success' }}
         env:
-          PYTEST_ADDOPTS: "--durations=0 -v -p no:unraisableexception -p no:threadexception"
-        run: uv run --no-sync make test-other-python-rest REST_PYTHON_SHARD=1 REST_REPORT_DIR=.
-      - name: Upload Rest Python timing and memory reports
-        if: always()
-        uses: actions/upload-artifact@v4
-        with:
-          name: rest-python-test-timings-shard-1
-          path: |
-            rest-python-*.xml
-            rest-python-*-resources.txt
-      - name: Run tests/variables YAML tests
-        if: ${{ !cancelled() && steps.system-ca.outcome == 'success' }}
+          GH_TOKEN: ${{ github.token }}
+        run: |
+          set -uo pipefail
+          runs=$(gh run list --repo "$GITHUB_REPOSITORY" --workflow push.yaml \
+            --branch main --status success --limit 40 \
+            --json databaseId,displayTitle \
+            --jq '.[] | select(.displayTitle == "Update PolicyEngine US") | .databaseId') || runs=""
+          for run in $(echo "$runs" | head -n 3); do
+            if gh run download "$run" --repo "$GITHUB_REPOSITORY" \
+              --dir "rest-timings/$run" \
+              --pattern 'rest-python-test-timings-*' \
+              --pattern rest-variables-test-timings; then
+              echo "source=release run $run" >> "$GITHUB_OUTPUT"
+              exit 0
+            fi
+            rm -rf "rest-timings/$run"
+          done
+          echo "::warning::No Rest Python timing reports found; the plan counts every file the same."
+      - name: Plan shards
+        id: plan
         env:
-          PYTEST_ADDOPTS: "--durations=0 --junitxml=rest-variables.xml -p no:unraisableexception -p no:threadexception"
-        run: /usr/bin/time -v -o rest-variables-resources.txt uv run --no-sync make test-yaml-variables
-      - name: Upload Rest variables timing and memory reports
-        if: always()
-        uses: actions/upload-artifact@v4
-        with:
-          name: rest-variables-test-timings
-          path: |
-            rest-variables.xml
-            rest-variables-resources.txt
-  Rest:
-    needs: [RestPython]
-    name: Full Suite - Rest (Python + variables)
+          TIMINGS_SOURCE: ${{ steps.timings.outputs.source }}
+        run: uv run --no-sync python policyengine_us/tests/rest_python_shards.py plan --timings rest-timings --source "$TIMINGS_SOURCE" --github-output "$GITHUB_OUTPUT" --summary "$GITHUB_STEP_SUMMARY"
+  RestPython:
+    needs: [RestPlan]
+    name: Full Suite - Rest (Python shard ${{ matrix.shard }})
     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.
+    # Plans aim at 25 minutes a shard (26-30 from the October 2026 reports);
+    # one made without reports can run longer.
     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: |
-      !cancelled()
-      && (needs.RestPython.result != 'skipped')
-      && (github.repository == 'PolicyEngine/policyengine-us')
+      (github.repository == 'PolicyEngine/policyengine-us')
       && (github.event.head_commit.message == 'Update PolicyEngine US')
+    strategy:
+      fail-fast: false
+      matrix: ${{ fromJSON(needs.RestPlan.outputs.matrix) }}
     env:
       PYTHONUNBUFFERED: 1
       OMP_NUM_THREADS: 4
@@ -440,31 +431,57 @@ jobs:
       - name: Turn off default branching
         shell: bash
         run: bash ./update_itemization.sh
-      # Each group remains a separate, sequential process on its runner.
+      # Each planned process runs files of one group, one process at a time.
       # SPM simulation/API contracts stay in the Microsimulation job.
       - name: Run remaining Python tests
-        id: python_tests
         if: ${{ !cancelled() && steps.system-ca.outcome == 'success' }}
         env:
           PYTEST_ADDOPTS: "--durations=0 -v -p no:unraisableexception -p no:threadexception"
-        run: uv run --no-sync make test-other-python-rest REST_PYTHON_SHARD=2 REST_REPORT_DIR=.
+          REST_PYTHON_PROCESSES: ${{ matrix.processes }}
+        run: uv run --no-sync python policyengine_us/tests/rest_python_shards.py run --report-dir .
       - name: Upload Rest Python timing and memory reports
         if: always()
         uses: actions/upload-artifact@v4
         with:
-          name: rest-python-test-timings-shard-2
+          name: rest-python-test-timings-shard-${{ matrix.shard }}
           path: |
             rest-python-*.xml
             rest-python-*-resources.txt
+      - name: Run tests/variables YAML tests
+        if: ${{ !cancelled() && steps.system-ca.outcome == 'success' && matrix.variables }}
+        env:
+          PYTEST_ADDOPTS: "--durations=0 --junitxml=rest-variables.xml -p no:unraisableexception -p no:threadexception"
+        run: /usr/bin/time -v -o rest-variables-resources.txt uv run --no-sync make test-yaml-variables
+      - name: Upload Rest variables timing and memory reports
+        if: ${{ always() && matrix.variables }}
+        uses: actions/upload-artifact@v4
+        with:
+          name: rest-variables-test-timings
+          path: |
+            rest-variables.xml
+            rest-variables-resources.txt
+  Rest:
+    needs: [RestPlan, RestPython]
+    # Runs after a failed plan or shard, but not when the plan was skipped (a
+    # failed release lock) or the run is cancelled. always() kept this runner
+    # alive after cancellation.
+    if: |
+      !cancelled()
+      && (needs.RestPlan.result != 'skipped')
+      && (github.repository == 'PolicyEngine/policyengine-us')
+      && (github.event.head_commit.message == 'Update PolicyEngine US')
+    name: Check Rest Python shards
+    runs-on: ubuntu-latest
+    timeout-minutes: 5
+    steps:
       - name: Check all Rest jobs succeeded
-        if: always()
         shell: bash
         env:
+          PLAN_RESULT: ${{ needs.RestPlan.result }}
           PYTHON_RESULT: ${{ needs.RestPython.result }}
-          SHARD_RESULT: ${{ steps.python_tests.outcome }}
         run: |
-          if [[ "$PYTHON_RESULT" != "success" || "$SHARD_RESULT" != "success" ]]; then
-            echo "::error::Rest shards did not all succeed: shard-1=$PYTHON_RESULT, shard-2=$SHARD_RESULT"
+          if [[ "$PLAN_RESULT" != "success" || "$PYTHON_RESULT" != "success" ]]; then
+            echo "::error::Rest jobs did not all succeed: plan=$PLAN_RESULT, shards=$PYTHON_RESULT"
             exit 1
           fi
   Microsimulation:
@@ -527,7 +544,7 @@ jobs:
     if: |
       (github.repository == 'PolicyEngine/policyengine-us')
       && (github.event.head_commit.message == 'Update PolicyEngine US')
-    needs: [ReleaseLock, ReleaseLockTests, Baseline, HouseholdAPIPartners, Contrib, RestPython, Rest, Microsimulation]
+    needs: [ReleaseLock, ReleaseLockTests, Baseline, HouseholdAPIPartners, Contrib, RestPlan, RestPython, Rest, Microsimulation]
     # PyPI trusted publishing: PyPI accepts uploads only from this workflow in
     # the pypi environment, which is limited to main. No stored PyPI token.
     environment: pypi

.github/workflows/python-tests.yaml

@@ -1,23 +1,21 @@
 name: Python tests
 on:
   workflow_call:
-# Two sequential Rest runners keep each shard below its timeout.
-# Shard 1 includes variables; shard 2 also supplies the legacy summary check.
+# RestPlan spreads the Rest Python groups over parallel shards by the latest
+# release run's times (policyengine_us/tests/rest_python_shards.py); each
+# process still runs files of one Makefile group. Rest runs no tests: it
+# fails unless the plan and every shard succeeded.
 # Only the guarded PR workflow calls this file; it has no automatic triggers.
 jobs:
-  RestPython:
-    name: Full Suite - Rest (Python, shard-1 + variables)
+  RestPlan:
+    name: Plan Rest Python shards
     runs-on: ubuntu-latest
-    # Historical group times, including variables: 33-46 minutes plus setup.
-    timeout-minutes: 90
-    env:
-      PYTHONUNBUFFERED: 1
-      OMP_NUM_THREADS: 4
-      OPENBLAS_NUM_THREADS: 4
-      MKL_NUM_THREADS: 4
-      VECLIB_MAXIMUM_THREADS: 4
-      NUMEXPR_NUM_THREADS: 4
-      OMP_WAIT_POLICY: PASSIVE
+    timeout-minutes: 10
+    permissions:
+      actions: read
+      contents: read
+    outputs:
+      matrix: ${{ steps.plan.outputs.matrix }}
     steps:
       - name: Checkout repo
         uses: actions/checkout@v6
@@ -32,51 +30,48 @@ jobs:
         uses: astral-sh/setup-uv@v8.1.0
         with:
           version: "0.13.4"
-      - name: Resolve latest compatible PolicyEngine Core
-        run: uv lock --upgrade-package policyengine-core
       - name: Install dependencies
         run: uv sync --locked --extra dev
-      - name: Turn off default branching
+      # A successful release run on main has every shard's JUnit and GNU time
+      # reports; a failed one can lack a lost runner's, which would skew the
+      # estimates. Without any, the plan counts each file the same.
+      - name: Download recent Rest Python timings
+        id: timings
         shell: bash
-        run: bash ./update_itemization.sh
-      # Each group remains a separate, sequential process on its runner.
-      # SPM simulation/API contracts stay in the Microsimulation job.
-      - name: Run remaining Python tests
-        if: ${{ !cancelled() && steps.system-ca.outcome == 'success' }}
         env:
-          PYTEST_ADDOPTS: "--durations=0 -v -p no:unraisableexception -p no:threadexception"
-        run: uv run --no-sync make test-other-python-rest REST_PYTHON_SHARD=1 REST_REPORT_DIR=.
-      - name: Upload Rest Python timing and memory reports
-        if: always()
-        uses: actions/upload-artifact@v4
-        with:
-          name: rest-python-test-timings-shard-1
-          path: |
-            rest-python-*.xml
-            rest-python-*-resources.txt
-      - name: Run tests/variables YAML tests
-        if: ${{ !cancelled() && steps.system-ca.outcome == 'success' }}
+          GH_TOKEN: ${{ github.token }}
+        run: |
+          set -uo pipefail
+          runs=$(gh run list --repo "$GITHUB_REPOSITORY" --workflow push.yaml \
+            --branch main --status success --limit 40 \
+            --json databaseId,displayTitle \
+            --jq '.[] | select(.displayTitle == "Update PolicyEngine US") | .databaseId') || runs=""
+          for run in $(echo "$runs" | head -n 3); do
+            if gh run download "$run" --repo "$GITHUB_REPOSITORY" \
+              --dir "rest-timings/$run" \
+              --pattern 'rest-python-test-timings-*' \
+              --pattern rest-variables-test-timings; then
+              echo "source=release run $run" >> "$GITHUB_OUTPUT"
+              exit 0
+            fi
+            rm -rf "rest-timings/$run"
+          done
+          echo "::warning::No Rest Python timing reports found; the plan counts every file the same."
+      - name: Plan shards
+        id: plan
         env:
-          PYTEST_ADDOPTS: "--durations=0 --junitxml=rest-variables.xml -p no:unraisableexception -p no:threadexception"
-        run: /usr/bin/time -v -o rest-variables-resources.txt uv run --no-sync make test-yaml-variables
-      - name: Upload Rest variables timing and memory reports
-        if: always()
-        uses: actions/upload-artifact@v4
-        with:
-          name: rest-variables-test-timings
-          path: |
-            rest-variables.xml
-            rest-variables-resources.txt
-  Rest:
-    needs: [RestPython]
-    # 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)
+          TIMINGS_SOURCE: ${{ steps.timings.outputs.source }}
+        run: uv run --no-sync python policyengine_us/tests/rest_python_shards.py plan --timings rest-timings --source "$TIMINGS_SOURCE" --github-output "$GITHUB_OUTPUT" --summary "$GITHUB_STEP_SUMMARY"
+  RestPython:
+    needs: [RestPlan]
+    name: Full Suite - Rest (Python shard ${{ matrix.shard }})
     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.
+    # Plans aim at 25 minutes a shard (26-30 from the October 2026 reports);
+    # one made without reports can run longer.
     timeout-minutes: 90
+    strategy:
+      fail-fast: false
+      matrix: ${{ fromJSON(needs.RestPlan.outputs.matrix) }}
     env:
       PYTHONUNBUFFERED: 1
       OMP_NUM_THREADS: 4
@@ -106,31 +101,52 @@ jobs:
       - name: Turn off default branching
         shell: bash
         run: bash ./update_itemization.sh
-      # Each group remains a separate, sequential process on its runner.
+      # Each planned process runs files of one group, one process at a time.
       # SPM simulation/API contracts stay in the Microsimulation job.
       - name: Run remaining Python tests
-        id: python_tests
         if: ${{ !cancelled() && steps.system-ca.outcome == 'success' }}
         env:
           PYTEST_ADDOPTS: "--durations=0 -v -p no:unraisableexception -p no:threadexception"
-        run: uv run --no-sync make test-other-python-rest REST_PYTHON_SHARD=2 REST_REPORT_DIR=.
+          REST_PYTHON_PROCESSES: ${{ matrix.processes }}
+        run: uv run --no-sync python policyengine_us/tests/rest_python_shards.py run --report-dir .
       - name: Upload Rest Python timing and memory reports
         if: always()
         uses: actions/upload-artifact@v4
         with:
-          name: rest-python-test-timings-shard-2
+          name: rest-python-test-timings-shard-${{ matrix.shard }}
           path: |
             rest-python-*.xml
             rest-python-*-resources.txt
+      - name: Run tests/variables YAML tests
+        if: ${{ !cancelled() && steps.system-ca.outcome == 'success' && matrix.variables }}
+        env:
+          PYTEST_ADDOPTS: "--durations=0 --junitxml=rest-variables.xml -p no:unraisableexception -p no:threadexception"
+        run: /usr/bin/time -v -o rest-variables-resources.txt uv run --no-sync make test-yaml-variables
+      - name: Upload Rest variables timing and memory reports
+        if: ${{ always() && matrix.variables }}
+        uses: actions/upload-artifact@v4
+        with:
+          name: rest-variables-test-timings
+          path: |
+            rest-variables.xml
+            rest-variables-resources.txt
+  Rest:
+    needs: [RestPlan, RestPython]
+    # Runs after a failed plan or shard but stops when the run is cancelled.
+    # always() kept this runner alive after cancellation.
+    if: ${{ !cancelled() }}
+    name: Check Rest Python shards
+    runs-on: ubuntu-latest
+    timeout-minutes: 5
+    steps:
       - name: Check all Rest jobs succeeded
-        if: always()
         shell: bash
         env:
+          PLAN_RESULT: ${{ needs.RestPlan.result }}
           PYTHON_RESULT: ${{ needs.RestPython.result }}
-          SHARD_RESULT: ${{ steps.python_tests.outcome }}
         run: |
-          if [[ "$PYTHON_RESULT" != "success" || "$SHARD_RESULT" != "success" ]]; then
-            echo "::error::Rest shards did not all succeed: shard-1=$PYTHON_RESULT, shard-2=$SHARD_RESULT"
+          if [[ "$PLAN_RESULT" != "success" || "$PYTHON_RESULT" != "success" ]]; then
+            echo "::error::Rest jobs did not all succeed: plan=$PLAN_RESULT, shards=$PYTHON_RESULT"
             exit 1
           fi
   Microsimulation: