diff --git a/.github/workflows/push.yaml b/.github/workflows/push.yaml index c27b07cd6cc..abcff5545f5 100644 --- a/.github/workflows/push.yaml +++ b/.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 diff --git a/.github/workflows/python-tests.yaml b/.github/workflows/python-tests.yaml index 4e682b6ed72..ad6a88a0b01 100644 --- a/.github/workflows/python-tests.yaml +++ b/.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: diff --git a/Makefile b/Makefile index a3dc2fcbe44..0b15d4afae5 100644 --- a/Makefile +++ b/Makefile @@ -45,26 +45,12 @@ REST_SPM_TESTS := $(TESTS)/core/test_spm_policy_family.py \ # SSI resource deeming retains a reformed system at import and builds another # in its grid fixture. Core reached 13.3 GB in CI run 38141803955 after this # and three smaller modules arrived; isolate SSI rather than filling heavy-a -# (8.1 GB including the pending capital-gains MTR change). These groups remain -# sequential subprocesses within each Rest Python shard. +# (8.1 GB including the pending capital-gains MTR change). +# CI does not run this target. tests/rest_python_shards.py reads these groups +# from its dry run and spreads their files over parallel runners by recent +# CI times, splitting a long group into chunks. A process still runs files +# of one group only, in this order, so none holds more than its group here. REST_PYTHON_GROUPS := cached-arrays heavy-a heavy-b core-heavy core policy-parameters policy remaining-heavy remaining -# CI runs two shards of these existing isolated processes. Across seven -# recent runs (including 38326922385 and 38289790999), their summed group -# times were 32-45 and 32-46 minutes; variables adds about a minute to shard 1. -# Both stay near 45 minutes, while the old serial job hit its 90-minute limit. -# No shard runs more than one Python group at a time. Locally, leaving -# REST_PYTHON_SHARD unset still runs every group in the original order. -REST_PYTHON_SHARD_1_GROUPS := cached-arrays core-heavy policy-parameters remaining-heavy remaining -REST_PYTHON_SHARD_2_GROUPS := heavy-a heavy-b core policy -REST_PYTHON_SHARD ?= -ifeq ($(REST_PYTHON_SHARD),) -REST_PYTHON_SELECTED_GROUPS := $(REST_PYTHON_GROUPS) -else -REST_PYTHON_SELECTED_GROUPS := $(REST_PYTHON_SHARD_$(REST_PYTHON_SHARD)_GROUPS) -ifeq ($(strip $(REST_PYTHON_SELECTED_GROUPS)),) -$(error REST_PYTHON_SHARD must be 1 or 2) -endif -endif REST_CACHED_ARRAY_TESTS := $(TESTS)/test_formulas_do_not_write_into_cached_arrays.py REST_HEAVY_A_TESTS := $(TESTS)/test_md_poverty_line_credit_invariants.py \ $(TESTS)/core/test_ald_determinism.py \ @@ -97,10 +83,11 @@ REST_LISTED_TESTS := $(REST_CACHED_ARRAY_TESTS) $(REST_HEAVY_A_TESTS) $(REST_HEA $(REST_CORE_HEAVY_TESTS) $(REST_POLICY_PARAMETER_TESTS) $(REST_REMAINING_HEAVY_TESTS) REST_IGNORES := $(addprefix --ignore=,$(TESTS)/policy/contrib $(TESTS)/microsimulation $(REST_SPM_TESTS)) # $(call rest_pytest,,,) is the pytest -# command for one group. CI sets REST_REPORT_DIR, and each group then writes -# its own JUnit XML and GNU time report there: rest-python-.xml and -# rest-python--resources.txt. Unset, it is plain pytest, since macOS -# /usr/bin/time has no -v. +# command for one group. With REST_REPORT_DIR set, each group writes its own +# JUnit XML and GNU time report there: rest-python-.xml and +# rest-python--resources.txt; rest_python_shards.py names the groups +# from these files. Unset, it is plain pytest, since macOS /usr/bin/time has +# no -v. rest_pytest = $(strip $(if $(REST_REPORT_DIR),/usr/bin/time -v -o $(REST_REPORT_DIR)/rest-python-$(1)-resources.txt) \ pytest $(2) --maxfail=0 $(REST_IGNORES) $(addprefix --ignore=,$(3)) \ $(if $(REST_REPORT_DIR),--junitxml=$(REST_REPORT_DIR)/rest-python-$(1).xml)) @@ -257,7 +244,8 @@ test-yaml-no-structural-other-usda: # The single USDA process peaked at 14.9 GB in CI before a later run # shut down in WIC. Release accumulated model and parameter caches # between eight sequential YAML batches on the same runner. USDA Python - # tests remain covered by test-other-python-rest in the Rest job. + # tests stay in test-other-python-rest's policy group, which CI spreads + # over the planned Rest Python shards. $(BATCH) $(TESTS)/policy/baseline/gov/usda --batches 8 --workers 1 test-yaml-no-structural-other-ssa-usda: test-yaml-no-structural-other-ssa test-yaml-no-structural-other-usda test-yaml-no-structural-other-rest-a: @@ -304,7 +292,7 @@ test-other-python-spm: pytest $(REST_SPM_TESTS) --maxfail=0 # Every group runs even after one fails; the target fails if any group did. test-other-python-rest: - @failed=""; for group in $(REST_PYTHON_SELECTED_GROUPS); do \ + @failed=""; for group in $(REST_PYTHON_GROUPS); do \ echo "=== Rest Python group: $$group ==="; \ $(MAKE) --no-print-directory test-other-python-rest-$$group || failed="$$failed $$group"; \ done; \ diff --git a/changelog.d/ci-plan-rest-python-shards.changed.md b/changelog.d/ci-plan-rest-python-shards.changed.md new file mode 100644 index 00000000000..e3ade13c801 --- /dev/null +++ b/changelog.d/ci-plan-rest-python-shards.changed.md @@ -0,0 +1 @@ +Run the Rest Python tests on parallel CI shards that a plan job balances from the latest release run's test times, splitting long groups, instead of two chained runners that took about 130 minutes. 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..24bb5fef2ac 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,9 @@ WORKFLOW = REPO_ROOT / ".github" / "workflows" / "pr.yaml" SHARD = re.compile(r"\s+--shard\s+\d+/\d+") MAKE_TARGET = re.compile(r"\bmake\s+([\w-]+)") +# CI's planned Rest Python shards run this target's groups; the planner reads +# them from its dry run (code_health/test_rest_python_groups.py). +PLANNED_TARGETS = {"rest_python_shards.py run": "test-other-python-rest"} pytestmark = pytest.mark.skipif( shutil.which("make") is None or not WORKFLOW.exists(), @@ -59,7 +62,9 @@ def ci_suites(): targets, commands = set(), set() for job in workflow_jobs(WORKFLOW): matrix = (job.get("strategy") or {}).get("matrix") or {} - for entry in matrix.get("include") or []: + # A planned matrix is an expression; its steps name the target. + entries = [] if isinstance(matrix, str) else matrix.get("include") or [] + for entry in entries: if "target" in entry: targets.add(entry["target"]) if "cmd" in entry: @@ -71,6 +76,9 @@ def ci_suites(): for target in MAKE_TARGET.findall(step.get("run") or ""): if not target.startswith("$"): targets.add(target) + for command, target in PLANNED_TARGETS.items(): + if command in (step.get("run") or ""): + targets.add(target) return {t for t in targets if t.startswith("test")}, commands 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 91e51a5d024..6ac523b5ece 100644 --- a/policyengine_us/tests/code_health/test_pr_runner_budget.py +++ b/policyengine_us/tests/code_health/test_pr_runner_budget.py @@ -1,11 +1,17 @@ """Consolidating short checks must preserve coverage and failure isolation.""" +import importlib.util import math from pathlib import Path import yaml ROOT = Path(__file__).resolve().parents[3] +SPEC = importlib.util.spec_from_file_location( + "rest_python_shards", ROOT / "policyengine_us/tests/rest_python_shards.py" +) +shards = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(shards) JOBS = yaml.safe_load((ROOT / ".github/workflows/pr.yaml").read_text())["jobs"] GROUP_WORKFLOWS = { name: yaml.safe_load((ROOT / job["uses"]).read_text()) @@ -18,6 +24,7 @@ KEEP_RUNNING_WITH_SYSTEM_CA = ( "${{ !cancelled() && steps.system-ca.outcome == 'success' }}" ) +PLANNED_MATRIX = "${{ fromJSON(needs.RestPlan.outputs.matrix) }}" def action(name): @@ -34,7 +41,14 @@ def calls(job, name): ] +def planned(job): + return job.get("strategy", {}).get("matrix") == PLANNED_MATRIX + + def runners(job): + """A job's runners; a planned matrix counts its largest possible plan.""" + if planned(job): + return shards.MAX_SHARDS matrix = job.get("strategy", {}).get("matrix", {}) if "include" in matrix: return len(matrix["include"]) @@ -48,22 +62,27 @@ def test_pr_runner_budget_keeps_the_heavy_suites_separate(): for name, job in workflow["jobs"].items(): assert name not in runner_jobs runner_jobs[name] = job + assert {name for name, job in runner_jobs.items() if planned(job)} == {"RestPython"} for name, job in runner_jobs.items(): - counts[name] = runners(job) + if not planned(job): + 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. + # in October 2026, and by run 40014592006 its two chained runners took + # 130 minutes. RestPlan now spreads the groups over parallel shards (five + # from the October reports, at most MAX_SHARDS), and Rest checks their + # results on a runner of its own. assert sum(counts.values()) == 27 assert { name: counts[name] - for name in ("Baseline", "Contrib", "RestPython", "Rest", "Microsimulation") + for name in ("Baseline", "Contrib", "RestPlan", "Rest", "Microsimulation") } == { "Baseline": 11, "Contrib": 9, - "RestPython": 1, + "RestPlan": 1, "Rest": 1, "Microsimulation": 1, } + assert (shards.MIN_SHARDS, shards.MAX_SHARDS) == (4, 12) def needs(job): @@ -75,11 +94,12 @@ def concurrent(job): return min(job.get("strategy", {}).get("max-parallel", math.inf), runners(job)) -def test_uncapped_matrices_keep_a_pr_at_16_runners_at_once(): +def test_uncapped_matrices_keep_a_pr_at_27_runners_at_once(): # Within pr.yaml, validation runs alone; everything that needs only it starts - # together, Rest waits for RestPython, and Contrib waits for Baseline. - # The Enterprise runner capacity allows every group in each matrix to - # start together. The existing dependencies still bound the peak. + # together, the planned Python shards follow RestPlan, Rest waits for them, + # and Contrib waits for Baseline. The Enterprise runner capacity allows + # every group in each matrix to start together. The dependencies and the + # planner's MAX_SHARDS still bound the peak. for name in ("Baseline", "Contrib"): for job in (GROUP_WORKFLOWS[name]["jobs"][name], PUSH_JOBS[name]): assert "max-parallel" not in job["strategy"] @@ -93,13 +113,21 @@ def test_uncapped_matrices_keep_a_pr_at_16_runners_at_once(): 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 and "Contrib" not in first_wave + assert not {"Rest", "RestPython", "Contrib"} & set(first_wave) assert sum(first_wave.values()) == 16 - # Contrib starts only after Baseline has finished, so at most Contrib's - # groups plus every other first-wave job can be running then. - single = {name: n for name, n in first_wave.items() if name != "Baseline"} + # The shards start once RestPlan has finished. Contrib starts only after + # Baseline has finished, so at most Contrib's groups, the shards and every + # other first-wave job can be running then. + python = concurrent(PYTHON_JOBS["RestPython"]) + assert python == shards.MAX_SHARDS + assert sum(first_wave.values()) - first_wave["RestPlan"] + python <= 27 + single = { + name: n + for name, n in first_wave.items() + if name not in ("Baseline", "RestPlan") + } contrib = concurrent(GROUP_WORKFLOWS["Contrib"]["jobs"]["Contrib"]) - assert contrib + sum(single.values()) <= 16 + assert contrib + python + sum(single.values()) <= 27 def test_contrib_waits_for_baseline_and_skips_when_it_fails(): @@ -160,22 +188,29 @@ def test_suite_groups_have_no_extra_runners_or_triggers(): # PyYAML parses the bare `on` key as True. No push/PR trigger here: # the draft guard must control the only route to these runners. assert GROUP_WORKFLOWS[name][True] == {"workflow_call": None} - assert set(PYTHON_JOBS) == {"RestPython", "Rest", "Microsimulation"} + assert set(PYTHON_JOBS) == {"RestPlan", "RestPython", "Rest", "Microsimulation"} assert {name: job["timeout-minutes"] for name, job in PYTHON_JOBS.items()} == { + "RestPlan": 10, "RestPython": 90, - "Rest": 90, + "Rest": 5, "Microsimulation": 60, } + expected_needs = { + "RestPlan": [], + "RestPython": ["RestPlan"], + "Rest": ["RestPlan", "RestPython"], + "Microsimulation": [], + } for name, job in PYTHON_JOBS.items(): assert job["runs-on"] == "ubuntu-latest" - assert job.get("needs", []) == (["RestPython"] if name == "Rest" else []) - assert not job.get("strategy") + assert job.get("needs", []) == expected_needs[name] + assert bool(job.get("strategy")) == (name == "RestPython") 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. + # Only the planned shards form a matrix. The summary runner checks the + # plan's and every shard's result; there is no variables-only runner. # !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)" + assert PYTHON_JOBS["Rest"]["name"] == "Check Rest Python shards" def test_sequential_variants_preserve_failures_and_allow_later_results(): 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 293b2a7f947..aa1ce8ab3e7 100644 --- a/policyengine_us/tests/code_health/test_rest_python_groups.py +++ b/policyengine_us/tests/code_health/test_rest_python_groups.py @@ -4,29 +4,32 @@ and the SPM files in one pytest process. That process peaked at 15.7 GB RSS on the 16 GB runner (CI run 37198184683), and later runs lost the runner partway through the step. `make test-other-python-rest` now runs the same files as one -pytest process per group in REST_PYTHON_GROUPS. CI assigns those groups to -two Python shards on separate runners, with the variables YAML in the first -shard. The second shard retains the original Rest check and waits for the -first shard so it can fail if either shard fails. +pytest process per group in REST_PYTHON_GROUPS. In CI, rest_python_shards.py +plans those groups' files over parallel shards from a recent run's times, +splitting a long group into chunks. A process runs files of one group only, +and Rest, which runs no tests, fails unless every shard succeeded. Paths and --ignore flags can drop a file, or run it twice, while every test -still passes. So this guard expands the workflows' shard commands through -`make -n`, asks pytest which files each group collects, and checks that each -file of the old single command runs exactly once. A new test file lands in -core or policy when it sits under those folders and in remaining otherwise; -the check fails if a change to the groups or shards leaves one out. Running -the target with a stub pytest checks that every group runs and that any -failure fails the target. The original Rest check aggregates all shard -results, and must fail on failed, cancelled or skipped jobs. +still passes. So this guard expands the groups through `make -n`, asks pytest +which files each group collects, and checks that each file of the old single +command runs exactly once. A new test file lands in core or policy when it +sits under those folders and in remaining otherwise; the check fails if a +change to the groups leaves one out. Plans made with no, measured, skewed and +tiny times must then run each group's files once, in the group's order, and +the workflows must run and upload each shard of the plan. Running the target +and a shard with a stub pytest checks that every process runs and that any +failure fails them. The Rest check aggregates the plan and shards, and must +fail on failed, cancelled or skipped jobs. pytest lists the files here through its own discovery (paths, --ignore, -python_files), but nothing is imported: running this file as a script swaps -each test module for a placeholder. It also skips conftest.py files, because -importing policyengine_us/conftest.py imports the whole model. +python_files), but nothing is imported: rest_python_shards.py's list-files +mode swaps each test module for a placeholder. It also skips conftest.py +files, because importing policyengine_us/conftest.py imports the whole model. test_conftest_files_cannot_change_which_files_are_collected keeps that safe. """ import ast +import importlib.util import json import os import re @@ -34,7 +37,7 @@ import shutil import subprocess import sys -import tempfile +import zlib from collections import defaultdict from fnmatch import fnmatch from pathlib import Path @@ -44,23 +47,26 @@ REPO = Path(__file__).resolve().parents[3] TESTS = REPO / "policyengine_us" / "tests" +PLANNER = TESTS / "rest_python_shards.py" +SPEC = importlib.util.spec_from_file_location("rest_python_shards", PLANNER) +shards = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(shards) WORKFLOWS = [ REPO / ".github" / "workflows" / name for name in ("python-tests.yaml", "push.yaml") ] KEEP_RUNNING_WITH_SYSTEM_CA = ( "${{ !cancelled() && steps.system-ca.outcome == 'success' }}" ) -# CI runs this file under a shard of `make test-other-python-rest`, and that -# make exports these. Passed on, they would restrict the guard's reference to -# one shard, or carry its report directory into the stub run and overwrite -# the CI groups' GNU time reports. -PARENT_MAKE_VARIABLES = ( - "MAKEFLAGS", - "MFLAGS", - "MAKELEVEL", - "REST_REPORT_DIR", - "REST_PYTHON_SHARD", +RELEASE_ONLY = ( + "(github.repository == 'PolicyEngine/policyengine-us') " + "&& (github.event.head_commit.message == 'Update PolicyEngine US')" ) +# CI runs this file in a Rest shard, and a local `make test` exports these. +# Passed on, they would add report options to the guard's reference commands, +# or carry a report directory into the stub run and overwrite real reports. +PARENT_MAKE_VARIABLES = ("MAKEFLAGS", "MFLAGS", "MAKELEVEL", "REST_REPORT_DIR") +REPORT_DIR = "reports" +PLANNER_COMMAND = "python policyengine_us/tests/rest_python_shards.py" def dry_run(target, *overrides, prefixes=("pytest ", "/usr/bin/time ")): @@ -89,18 +95,7 @@ def collected_files(args, cwd=REPO): A command that collects nothing (pytest's exit 5) lists no files, which partition_problems reports by group. """ - env = {k: v for k, v in os.environ.items() if k != "PYTEST_ADDOPTS"} - with tempfile.TemporaryDirectory() as tmp: - listing = Path(tmp) / "files.json" - result = subprocess.run( - [sys.executable, __file__, str(listing), *args], - cwd=cwd, - env=env, - capture_output=True, - text=True, - ) - assert result.returncode in (0, 5), (args, result.stdout + result.stderr) - return json.loads(listing.read_text()) + return shards.collected_files(args, cwd) def partition_problems(reference, groups): @@ -131,17 +126,16 @@ def workflow_jobs(workflow): return yaml.safe_load(workflow.read_text())["jobs"] -def rest_steps(workflow, job="RestPython"): - return workflow_jobs(workflow)[job]["steps"] +def single(steps, predicate): + (step,) = [step for step in steps if predicate(step)] + return step def assert_system_ca_precedes(steps, runtime_step): """Continuations require a successful, unconditional CA setup first.""" - (setup,) = [ - step - for step in steps - if step.get("uses") == "./.github/actions/setup-system-ca" - ] + setup = single( + steps, lambda step: step.get("uses") == "./.github/actions/setup-system-ca" + ) assert setup["id"] == "system-ca" assert "if" not in setup assert not setup.get("continue-on-error", False) @@ -150,63 +144,6 @@ def assert_system_ca_precedes(steps, runtime_step): assert steps.index(setup) < steps.index(runtime_step) -def ci_python_shards(workflow): - """The real Make overrides and artifact configuration of every shard.""" - jobs = workflow_jobs(workflow) - calls = [ - name - for name, config in jobs.items() - for step in config.get("steps", []) - if "make test-other-python-rest" in step.get("run", "") - ] - assert sorted(calls) == ["Rest", "RestPython"], (workflow.name, calls) - shards = {} - for job_name, selector in (("RestPython", "1"), ("Rest", "2")): - job = jobs[job_name] - assert "strategy" not in job, f"{workflow.name}: extra {job_name} runners" - (run,) = [ - step - for step in job["steps"] - if "make test-other-python-rest" in step.get("run", "") - ] - (upload,) = [ - step - for step in job["steps"] - if step.get("uses", "").startswith("actions/upload-artifact") - and "rest-python-*.xml" in step["with"]["path"].split() - ] - argv = shlex.split(run["run"]) - assert argv[:5] == ["uv", "run", "--no-sync", "make", "test-other-python-rest"] - selectors = [ - arg.removeprefix("REST_PYTHON_SHARD=") - for arg in argv[5:] - if arg.startswith("REST_PYTHON_SHARD=") - ] - assert selectors == [selector], (workflow.name, job_name, argv) - assert all(arg.startswith("REST_") and "=" in arg for arg in argv[5:]), argv - shards[f"shard-{selector}"] = { - "overrides": argv[5:], - "artifact": upload["with"]["name"], - "patterns": upload["with"]["path"].split(), - } - return shards - - -def ci_report_dir(workflow): - """The REST_REPORT_DIR each workflow shard passes to the groups.""" - dirs = [] - for shard in ci_python_shards(workflow).values(): - values = [ - arg.removeprefix("REST_REPORT_DIR=") - for arg in shard["overrides"] - if arg.startswith("REST_REPORT_DIR=") - ] - assert len(values) == 1, f"{workflow.name} runs without one REST_REPORT_DIR" - dirs.extend(values) - assert len(set(dirs)) == 1, dirs - return dirs[0] - - def reported_groups(report_dir, *overrides): """Each group's pytest arguments and report files, in run order. @@ -225,6 +162,7 @@ def reported_groups(report_dir, *overrides): assert match, xml group = match.group(1) assert group not in groups, f"{group} runs twice" + assert argv[3] == f"{report_dir}/rest-python-{group}-resources.txt", argv groups[group] = { "args": [arg for arg in argv[5:] if arg != junit[0]], "reports": [xml, argv[3]], @@ -233,15 +171,8 @@ def reported_groups(report_dir, *overrides): @pytest.fixture(scope="module") -def report_dir(): - dirs = {ci_report_dir(workflow) for workflow in WORKFLOWS} - assert len(dirs) == 1, f"The workflows disagree on REST_REPORT_DIR: {dirs}" - return dirs.pop() - - -@pytest.fixture(scope="module") -def groups(report_dir): - return reported_groups(report_dir) +def groups(): + return reported_groups(REPORT_DIR) @pytest.fixture(scope="module") @@ -277,72 +208,517 @@ def test_groups_collect_each_file_of_the_old_single_process_once( assert not problems, "\n".join(problems) -@pytest.mark.parametrize("workflow", WORKFLOWS, ids=lambda path: path.name) -def test_ci_shards_run_each_group_and_each_python_file_once( - workflow, groups, group_files, reference, report_dir +def test_planner_reads_the_groups_the_makefile_runs(groups): + assert shards.make_groups() == { + group: spec["args"] for group, spec in groups.items() + } + + +def stable_seconds(file, low, high): + """A per-file time in [low, high) that depends only on the file name.""" + return low + (zlib.crc32(file.encode()) % 1000) / 1000 * (high - low) + + +def timing_scenarios(group_files): + files = [f for group in group_files.values() for f in group] + copies = "policyengine_us/tests/core/test_behavioral_response_measurement_copies.py" + measured = {f: stable_seconds(f, 0.5, 60) for f in files} + skewed = dict(measured) + skewed[files[0]] = 40 * 60 + if copies in skewed: + skewed[copies] = 27 * 60 + return { + "no reports": {}, + "measured": measured, + "partly measured": { + f: s for f, s in measured.items() if zlib.crc32(f.encode()) % 3 + }, + "one file past the target": skewed, + "tiny": {f: 0.1 for f in files}, + "huge": {f: 600 for f in files}, + } + + +@pytest.mark.parametrize( + "scenario", + [ + "no reports", + "measured", + "partly measured", + "one file past the target", + "tiny", + "huge", + ], +) +def test_every_plan_runs_each_group_file_once_in_group_order( + group_files, reference, scenario ): + seconds = timing_scenarios(group_files)[scenario] + plan = shards.make_plan(group_files, seconds) + assert shards.MIN_SHARDS <= len(plan) <= shards.MAX_SHARDS + assert [shard["shard"] for shard in plan] == list(range(1, len(plan) + 1)) + assert sum(shard["variables"] for shard in plan) == 1 selected = {} - for shard, config in ci_python_shards(workflow).items(): - shard_groups = reported_groups(report_dir, *config["overrides"]) - assert shard_groups, f"{workflow.name}: {shard} runs no groups" - for group, spec in shard_groups.items(): - assert group in groups, f"{workflow.name}: unknown group {group}" - assert spec == groups[group], (workflow.name, shard, group, spec) - selected[f"{shard}/{group}"] = group_files[group] + names = [] + for shard in plan: + assert shard["processes"], f"shard {shard['shard']} runs nothing" + for process in shard["processes"]: + names.append(process["name"]) + group = group_files[process["group"]] + assert process["name"] == process["group"] or re.fullmatch( + rf"{re.escape(process['group'])}-\d+", process["name"] + ) + # One group per process, in the order its group's process ran. + positions = [group.index(file) for file in process["files"]] + assert positions == sorted(positions), process["name"] + selected[f"shard {shard['shard']}/{process['name']}"] = process["files"] + # Process names become report names, which must not collide. + assert len(names) == len(set(names)), names problems = partition_problems(reference, selected) - assert not problems, f"{workflow.name}:\n" + "\n".join(problems) + assert not problems, f"{scenario}:\n" + "\n".join(problems) + + +def test_a_long_group_splits_and_its_longest_file_runs_alone(): + groups = { + "core": [f"core/test_{i:03d}.py" for i in range(115)], + "remaining": [f"test_{i:03d}.py" for i in range(87)], + "heavy": ["test_heavy.py"], + } + seconds = {f: 12 for f in groups["core"]} | {f: 15 for f in groups["remaining"]} + seconds["core/test_050.py"] = 27 * 60 + seconds["test_heavy.py"] = 6 * 60 + plan = shards.make_plan(groups, seconds, overhead=60, variables=70) + processes = {p["name"]: p for shard in plan for p in shard["processes"]} + # 50 core minutes split at 25; the 27-minute file fills a chunk alone. + assert sorted(processes) == ["core-1", "core-2", "heavy", "remaining"] + assert ["core/test_050.py"] in [p["files"] for p in processes.values()] + # Groups shorter than the target keep one process. + assert processes["remaining"]["files"] == groups["remaining"] + + +def rest_shaped_groups(): + """Groups and times shaped like the October 2026 Rest reports.""" + sizes = { + "cached-arrays": (1, 390), + "heavy-a": (4, 85), + "heavy-b": (4, 150), + "core-heavy": (3, 140), + "core": (115, 12), + "policy-parameters": (3, 70), + "policy": (25, 33), + "remaining-heavy": (6, 80), + "remaining": (87, 15), + } + groups = { + group: [f"{group}/test_{i:03d}.py" for i in range(count)] + for group, (count, _) in sizes.items() + } + seconds = { + file: sizes[group][1] for group, files in groups.items() for file in files + } + seconds["core/test_050.py"] = 27 * 60 + return groups, seconds + + +def test_the_plan_uses_the_fewest_shards_that_keep_the_longest_near_the_bound(): + groups, seconds = rest_shaped_groups() + plan = shards.make_plan(groups, seconds, overhead=60, variables=70) + cost = shards.estimates(groups, seconds) + weights = { + p["name"]: 60 + sum(cost[f] for f in p["files"]) + for shard in plan + for p in shard["processes"] + } + bound = shards.MAKESPAN_TOLERANCE * max( + shards.TARGET_MINUTES * 60, max(weights.values()) + ) + loads = [sum(weights[p["name"]] for p in shard["processes"]) for shard in plan] + assert max(loads) <= bound + # One shard fewer cannot meet the bound, so the plan uses no extra runner. + _, fewer = shards.longest_first(list(weights), len(plan) - 1, weights.get) + assert max(fewer) > bound + assert len(plan) == 6 + # The variables YAML goes to the least-loaded shard. + (marked,) = [i for i, shard in enumerate(plan) if shard["variables"]] + assert loads[marked] == min(loads) + assert plan[marked]["minutes"] == round((loads[marked] + 70) / 60, 1) + + +def test_plans_without_times_count_files_evenly_and_stay_complete(): + groups = {"a": [f"a/test_{i:03d}.py" for i in range(300)], "b": ["b/test_0.py"]} + plan = shards.make_plan(groups) + processes = [p for shard in plan for p in shard["processes"]] + files = sorted(f for p in processes for f in p["files"]) + assert files == sorted(groups["a"] + groups["b"]) + # 300 files at DEFAULT_FILE_SECONDS make 150 minutes: six equal chunks, + # 26 minutes each with process start. b cannot join one of them within + # the tolerance, so it gets a seventh shard, which also runs variables. + chunks = [len(p["files"]) for p in processes if p["group"] == "a"] + assert chunks == [50] * 6 + assert len(plan) == 7 + assert max(shard["minutes"] for shard in plan) == 26.0 + (marked,) = [shard for shard in plan if shard["variables"]] + assert [p["group"] for p in marked["processes"]] == ["b"] + + +def test_the_planner_refuses_group_options_a_shard_cannot_pass_on(tmp_path): + command = ( + "\t/usr/bin/time -v -o $(REST_REPORT_DIR)/rest-python-a-resources.txt " + "pytest a --maxfail=0 --ignore=a/b{extra} " + "--junitxml=$(REST_REPORT_DIR)/rest-python-a.xml\n" + ) + makefile = tmp_path / "Makefile" + makefile.write_text("test-other-python-rest:\n" + command.format(extra="")) + assert shards.make_groups(tmp_path) == {"a": ["a", "--maxfail=0", "--ignore=a/b"]} + makefile.write_text( + "test-other-python-rest:\n" + command.format(extra=" -p no:randomly") + ) + with pytest.raises(RuntimeError, match=r"cannot pass on \['-p'\]"): + shards.make_groups(tmp_path) + + +def test_the_plan_command_writes_a_complete_matrix_and_summary(tmp_path, group_files): + files = [f for group in group_files.values() for f in group] + write_report( + tmp_path / "timings" / "rest-python-test-timings-shard-1", + "rest-python-core", + [(file.removesuffix(".py").replace("/", "."), 30) for file in files[:20]], + wall="12:00", + ) + output = tmp_path / "github-output" + summary = tmp_path / "summary.md" + result = subprocess.run( + [ + sys.executable, + str(PLANNER), + "plan", + "--timings", + str(tmp_path / "timings"), + "--source", + "release run 1", + "--github-output", + str(output), + "--summary", + str(summary), + ], + cwd=REPO, + env={k: v for k, v in os.environ.items() if k not in PARENT_MAKE_VARIABLES}, + capture_output=True, + text=True, + ) + assert result.returncode == 0, result.stdout + result.stderr + (line,) = output.read_text().splitlines() + assert line.startswith("matrix=") + entries = json.loads(line.removeprefix("matrix="))["include"] + planned = [ + file + for entry in entries + for process in json.loads(entry["processes"]) + for file in process["files"] + ] + assert sorted(planned) == sorted(files) + assert sum(entry["variables"] for entry in entries) == 1 + assert summary.read_text().startswith( + f"Rest Python plan: {len(entries)} shards, times from release run 1 " + f"(20 of {len(files)} files timed)." + ) + + +def test_unknown_files_count_as_their_groups_median(): + groups = {"a": ["a1.py", "a2.py", "a3.py", "new.py"], "b": ["b1.py", "b2.py"]} + seconds = {"a1.py": 10, "a2.py": 20, "a3.py": 300, "b1.py": 1, "b2.py": 3} + assert shards.estimates(groups, seconds)["new.py"] == 20 + assert shards.estimates({"c": ["c1.py"]} | groups, seconds)["c1.py"] == 10 + + +@pytest.mark.parametrize( + "change, message", + [ + (lambda plan: plan[0]["processes"][0]["files"].pop(), "missing"), + ( + lambda plan: plan[1]["processes"][0]["files"].append("a/test_0.py"), + "outside its group", + ), + ( + lambda plan: plan[0]["processes"][0]["files"].reverse(), + "outside its group", + ), + ( + lambda plan: plan[1]["processes"].append( + dict(plan[0]["processes"][0], name="a-2") + ), + "repeated", + ), + (lambda plan: plan[1]["processes"].clear(), "runs nothing"), + (lambda plan: plan[1].update(variables=True), "exactly one shard"), + ( + lambda plan: plan[1]["processes"][0].update(name="a"), + "runs twice", + ), + ], +) +def test_the_plan_check_rejects_dropped_repeated_or_mixed_files(change, message): + groups = {"a": ["a/test_0.py", "a/test_1.py"], "b": ["b/test_0.py"]} + plan = [ + { + "shard": 1, + "variables": True, + "processes": [{"name": "a", "group": "a", "files": list(groups["a"])}], + }, + { + "shard": 2, + "variables": False, + "processes": [{"name": "b", "group": "b", "files": list(groups["b"])}], + }, + ] + shards.check_plan(groups, plan) + change(plan) + with pytest.raises(AssertionError, match=message): + shards.check_plan(groups, plan) + + +def write_report(directory, name, cases, wall=None): + """A JUnit report of (classname, seconds) cases and its GNU time report.""" + directory.mkdir(parents=True, exist_ok=True) + body = "".join( + f'' + for i, (classname, time) in enumerate(cases) + ) + (directory / f"{name}.xml").write_text( + f'{body}' + ) + if wall is not None: + (directory / f"{name}-resources.txt").write_text( + "\tCommand being timed: pytest\n" + f"\tElapsed (wall clock) time (h:mm:ss or m:ss): {wall}\n" + "\tMaximum resident set size (kbytes): 1000\n" + ) + + +def test_timings_come_from_each_shards_junit_and_gnu_time_reports(tmp_path): + files = ["pkg/tests/test_a.py", "pkg/tests/core/test_b.py", "pkg/tests/test_c.py"] + write_report( + tmp_path / "rest-python-test-timings-shard-1", + "rest-python-core-1", + [ + ("pkg.tests.test_a", 10), + ("pkg.tests.test_a.TestClass", 5.5), + ("pkg.tests.core.test_b", 20), + ("pkg.tests.test_removed", 7), + ], + wall="1:02.5", + ) + write_report( + tmp_path / "rest-python-test-timings-shard-2", + "rest-python-remaining", + [("pkg.tests.core.test_b", 1)], + wall="1:00:01", + ) + write_report(tmp_path / "rest-variables-test-timings", "rest-variables", [], "1:10") + # An earlier attempt's leftover report covers files a later shard ran + # too: each file counts once, at its longer time, not twice. + write_report( + tmp_path / "rest-python-test-timings-shard-6", + "rest-python-core-3", + [("pkg.tests.test_a", 12), ("pkg.tests.core.test_b", 30)], + ) + seconds, overhead, variables = shards.read_timings(tmp_path, files) + assert seconds == {"pkg/tests/test_a.py": 15.5, "pkg/tests/core/test_b.py": 30} + # 62.5 - 42.5 and 3601 - 1 seconds outside tests: the median of the two. + assert overhead == pytest.approx((20 + 3600) / 2) + assert variables == 70 + assert shards.read_timings(tmp_path / "missing", files) == ({}, None, None) + + +def test_the_matrix_carries_each_shards_processes_as_json(): + groups = {"a": ["a/test_0.py", "a/test_1.py"], "b": ["b/test 0's.py"]} + plan = shards.make_plan(groups, min_shards=2) + entries = shards.matrix(plan)["include"] + assert [entry["shard"] for entry in entries] == [1, 2] + assert [entry["variables"] for entry in entries] == [ + shard["variables"] for shard in plan + ] + assert [json.loads(entry["processes"]) for entry in entries] == [ + shard["processes"] for shard in plan + ] + + +def ci_jobs(workflow): + jobs = workflow_jobs(workflow) + calls = [ + name + for name, config in jobs.items() + for step in config.get("steps", []) + if PLANNER_COMMAND in step.get("run", "") + or "make test-other-python-rest" in step.get("run", "") + ] + assert sorted(calls) == ["RestPlan", "RestPython"], (workflow.name, calls) + return jobs @pytest.mark.parametrize("workflow", WORKFLOWS, ids=lambda path: path.name) -def test_each_group_writes_reports_its_ci_shard_uploads(workflow, report_dir): - artifacts = [] - reports = [] - for config in ci_python_shards(workflow).values(): - artifacts.append(config["artifact"]) - shard_reports = [ - report - for spec in reported_groups(report_dir, *config["overrides"]).values() - for report in spec["reports"] - ] - reports.extend(shard_reports) - patterns = [os.path.normpath(p) for p in config["patterns"]] - missing = [ - report - for report in shard_reports - if not any(fnmatch(os.path.normpath(report), p) for p in patterns) - ] - assert not missing, f"{workflow.name} does not upload {missing}" - assert len(set(reports)) == len(reports), reports - variable_uploads = [ - step - for step in rest_steps(workflow, "RestPython") - if step.get("uses", "").startswith("actions/upload-artifact") - and "rest-variables.xml" in step["with"]["path"].split() +def test_ci_plans_from_recent_reports_then_runs_and_uploads_each_shard(workflow): + jobs = ci_jobs(workflow) + plan_job, shard_job = jobs["RestPlan"], jobs["RestPython"] + assert plan_job["permissions"] == {"actions": "read", "contents": "read"} + assert plan_job["outputs"] == {"matrix": "${{ steps.plan.outputs.matrix }}"} + download = single(plan_job["steps"], lambda step: step.get("id") == "timings") + plan = single(plan_job["steps"], lambda step: step.get("id") == "plan") + assert plan_job["steps"].index(download) < plan_job["steps"].index(plan) + assert download["env"] == {"GH_TOKEN": "${{ github.token }}"} + assert plan["env"] == {"TIMINGS_SOURCE": "${{ steps.timings.outputs.source }}"} + assert shlex.split(plan["run"]) == [ + "uv", + "run", + "--no-sync", + *PLANNER_COMMAND.split(), + "plan", + "--timings", + "rest-timings", + "--source", + "$TIMINGS_SOURCE", + "--github-output", + "$GITHUB_OUTPUT", + "--summary", + "$GITHUB_STEP_SUMMARY", ] - assert len(variable_uploads) == 1 - artifacts.append(variable_uploads[0]["with"]["name"]) - assert len(set(artifacts)) == len(artifacts), artifacts - for job in ("RestPython", "Rest"): - steps = rest_steps(workflow, job) - (run,) = [ - step - for step in steps - if "make test-other-python-rest" in step.get("run", "") - ] - assert run["if"] == KEEP_RUNNING_WITH_SYSTEM_CA - assert_system_ca_precedes(steps, run) - (upload,) = [ - step - for step in steps - if step.get("uses", "").startswith("actions/upload-artifact") + assert not any(step.get("continue-on-error") for step in plan_job["steps"]) + + assert shard_job["needs"] == ["RestPlan"] + assert shard_job["strategy"] == { + "fail-fast": False, + "matrix": "${{ fromJSON(needs.RestPlan.outputs.matrix) }}", + } + assert "${{ matrix.shard }}" in shard_job["name"] + steps = shard_job["steps"] + run = single(steps, lambda step: PLANNER_COMMAND in step.get("run", "")) + assert shlex.split(run["run"]) == [ + "uv", + "run", + "--no-sync", + *PLANNER_COMMAND.split(), + "run", + "--report-dir", + ".", + ] + assert run["env"]["REST_PYTHON_PROCESSES"] == "${{ matrix.processes }}" + assert run["if"] == KEEP_RUNNING_WITH_SYSTEM_CA + assert_system_ca_precedes(steps, run) + upload = single( + steps, + lambda step: ( + step.get("uses", "").startswith("actions/upload-artifact") and "rest-python-*.xml" in step["with"]["path"].split() - ] - assert upload["if"] == "always()" - assert not any(step.get("continue-on-error", False) for step in steps) + ), + ) + assert upload["if"] == "always()" + assert ( + upload["with"]["name"] == "rest-python-test-timings-shard-${{ matrix.shard }}" + ) + patterns = upload["with"]["path"].split() + # Every report a planned process writes is uploaded, under a name the + # next plan downloads and reads. + reports = [ + f"./rest-python-{name}{suffix}" + for name in ("core", "core-12", "remaining-heavy") + for suffix in (".xml", "-resources.txt") + ] + for report in reports: + assert any(fnmatch(os.path.normpath(report), p) for p in patterns), report + assert fnmatch("rest-python-core-12.xml", "rest-python-*.xml") + shard_artifact = upload["with"]["name"].replace("${{ matrix.shard }}", "12") + assert "--pattern 'rest-python-test-timings-*'" in download["run"] + assert fnmatch(shard_artifact, "rest-python-test-timings-*") + assert not any(step.get("continue-on-error") for step in steps) + if workflow.name == "push.yaml": + for job in (plan_job, shard_job): + assert " ".join(job["if"].split()) == RELEASE_ONLY + assert plan_job["needs"] == "ReleaseLock" + else: + assert "if" not in shard_job and "if" not in plan_job + assert "needs" not in plan_job + + +STUB_GH = """#!/bin/sh +echo "$*" >> "$STUB_CALLS" +case "$1 $2" in + "run list") + [ "$STUB_LIST" = fail ] && exit 1 + printf '%s\\n' $STUB_RUNS + ;; + "run download") + for good in $STUB_GOOD; do + if [ "$3" = "$good" ]; then + mkdir -p "rest-timings/$3/rest-python-test-timings-shard-1" + exit 0 + fi + done + mkdir -p "rest-timings/$3/partial" + exit 1 + ;; +esac +""" @pytest.mark.parametrize("workflow", WORKFLOWS, ids=lambda path: path.name) -def test_variable_yaml_selection_runs_once_in_the_first_python_shard(workflow): +@pytest.mark.parametrize( + "runs, good, listing, source, tried", + [ + ("101 102", "101 102", "ok", "release run 101", ["101"]), + ("101 102", "102", "ok", "release run 102", ["101", "102"]), + ("101 102 103 104", "104", "ok", None, ["101", "102", "103"]), + ("", "", "ok", None, []), + ("101", "101", "fail", None, []), + ], +) +def test_the_timing_download_never_fails_the_plan( + tmp_path, workflow, runs, good, listing, source, tried +): + download = single( + workflow_jobs(workflow)["RestPlan"]["steps"], + lambda step: step.get("id") == "timings", + ) + assert download["shell"] == "bash" + stub = tmp_path / "bin" / "gh" + stub.parent.mkdir() + stub.write_text(STUB_GH) + stub.chmod(0o755) + output = tmp_path / "output.txt" + calls = tmp_path / "calls.txt" + calls.touch() + env = dict( + os.environ, + PATH=f"{stub.parent}{os.pathsep}{os.environ['PATH']}", + GITHUB_OUTPUT=str(output), + GITHUB_REPOSITORY="PolicyEngine/policyengine-us", + STUB_CALLS=str(calls), + STUB_RUNS=runs, + STUB_GOOD=good, + STUB_LIST=listing, + ) + result = subprocess.run( + ["bash", "-e", "-c", download["run"]], + cwd=tmp_path, + env=env, + capture_output=True, + text=True, + ) + assert result.returncode == 0, result.stdout + result.stderr + written = output.read_text() if output.exists() else "" + assert written == (f"source={source}\n" if source else "") + downloads = [ + line.split()[2] for line in calls.read_text().splitlines() if "download" in line + ] + assert downloads == tried + # A failed download leaves nothing for the plan to read. + kept = sorted(p.name for p in (tmp_path / "rest-timings").glob("*")) + assert kept == ([source.split()[-1]] if source else []) + if not source: + assert "::warning::No Rest Python timing reports" in result.stdout + + +@pytest.mark.parametrize("workflow", WORKFLOWS, ids=lambda path: path.name) +def test_variable_yaml_selection_runs_once_on_the_planned_shard(workflow): jobs = workflow_jobs(workflow) calls = [ (job, step) @@ -352,9 +728,12 @@ def test_variable_yaml_selection_runs_once_in_the_first_python_shard(workflow): ] assert len(calls) == 1, (workflow.name, calls) job, step = calls[0] + # The plan marks exactly one shard (see the plan tests above). assert job == "RestPython" - assert "strategy" not in jobs[job], "variables YAML must run exactly once" - assert step["if"] == KEEP_RUNNING_WITH_SYSTEM_CA + assert step["if"] == ( + "${{ !cancelled() && steps.system-ca.outcome == 'success' " + "&& matrix.variables }}" + ) assert_system_ca_precedes(jobs[job]["steps"], step) assert shlex.split(step["run"]) == [ "/usr/bin/time", @@ -375,17 +754,19 @@ def test_variable_yaml_selection_runs_once_in_the_first_python_shard(workflow): "--batches", "1", ] - assert not any( - "make test-yaml-variables" in candidate.get("run", "") - for candidate in rest_steps(workflow, "Rest") + upload = single( + jobs[job]["steps"], + lambda candidate: ( + candidate.get("uses", "").startswith("actions/upload-artifact") + and "rest-variables.xml" in candidate["with"]["path"].split() + ), ) - (upload,) = [ - candidate - for candidate in jobs[job]["steps"] - if candidate.get("uses", "").startswith("actions/upload-artifact") - and "rest-variables.xml" in candidate["with"]["path"].split() - ] - assert upload["if"] == "always()" + assert upload["if"] == "${{ always() && matrix.variables }}" + assert upload["with"]["name"] == "rest-variables-test-timings" + download = single( + jobs["RestPlan"]["steps"], lambda candidate: candidate.get("id") == "timings" + ) + assert "--pattern rest-variables-test-timings" in download["run"] junit = [ arg.removeprefix("--junitxml=") for arg in shlex.split(step["env"]["PYTEST_ADDOPTS"]) @@ -400,58 +781,37 @@ def test_variable_yaml_selection_runs_once_in_the_first_python_shard(workflow): @pytest.mark.parametrize("workflow", WORKFLOWS, ids=lambda path: path.name) -def test_original_rest_check_aggregates_all_shards_without_hiding_failure(workflow): +def test_the_rest_check_aggregates_the_plan_and_shards_without_hiding_failure( + workflow, +): jobs = workflow_jobs(workflow) summary = jobs["Rest"] - assert summary["name"] == "Full Suite - Rest (Python + variables)" - assert summary["needs"] == ["RestPython"] - # 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 summary["name"] == "Check Rest Python shards" + assert summary["needs"] == ["RestPlan", "RestPython"] + # It runs after a failed plan or shard, 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["RestPlan"].get("continue-on-error", False) assert not jobs["RestPython"].get("continue-on-error", False) if workflow.name == "push.yaml": assert " ".join(summary["if"].split()) == ( - "!cancelled() && (needs.RestPython.result != 'skipped') " - "&& (github.repository == 'PolicyEngine/policyengine-us') " - "&& (github.event.head_commit.message == 'Update PolicyEngine US')" + "!cancelled() && (needs.RestPlan.result != 'skipped') && " + RELEASE_ONLY ) - assert {"RestPython", "Rest"} <= set(jobs["Publish"]["needs"]) + assert {"RestPlan", "RestPython", "Rest"} <= set(jobs["Publish"]["needs"]) else: assert summary["if"] == "${{ !cancelled() }}" - (step,) = [ - candidate - for candidate in summary["steps"] - if "PYTHON_RESULT" in candidate.get("env", {}) - ] - (python_step,) = [ - candidate - for candidate in summary["steps"] - if candidate.get("id") == "python_tests" - ] - assert "make test-other-python-rest" in python_step["run"] - assert "REST_PYTHON_SHARD=2" in shlex.split(python_step["run"]) - assert python_step["if"] == KEEP_RUNNING_WITH_SYSTEM_CA - assert_system_ca_precedes(summary["steps"], python_step) - # This runner executes shard 2 before checking shard 1. Its own failures - # must remain failures, and cancellation skipping its Python step must - # fail the check even if shard 1 succeeded. - assert summary["steps"][-1] == step - assert step["if"] in ("always()", "${{ always() }}") + (step,) = summary["steps"] assert step["shell"] == "bash" assert step["env"] == { + "PLAN_RESULT": "${{ needs.RestPlan.result }}", "PYTHON_RESULT": "${{ needs.RestPython.result }}", - "SHARD_RESULT": "${{ steps.python_tests.outcome }}", } - assert not any( - candidate.get("continue-on-error", False) - for candidate in summary["steps"] + jobs["RestPython"]["steps"] - ) - for python in ("success", "failure", "cancelled", "skipped"): - for shard in ("success", "failure", "cancelled", "skipped"): - env = dict(os.environ, PYTHON_RESULT=python, SHARD_RESULT=shard) + outcomes = ("success", "failure", "cancelled", "skipped") + for plan in outcomes: + for python in outcomes: + env = dict(os.environ, PLAN_RESULT=plan, PYTHON_RESULT=python) result = subprocess.run( ["bash", "-e", "-c", step["run"]], cwd=REPO, @@ -459,10 +819,10 @@ def test_original_rest_check_aggregates_all_shards_without_hiding_failure(workfl capture_output=True, text=True, ) - assert (result.returncode == 0) == (python == shard == "success"), ( + assert (result.returncode == 0) == (plan == python == "success"), ( workflow.name, + plan, python, - shard, result.stdout + result.stderr, ) @@ -520,27 +880,13 @@ def test_groups_run_as_plain_pytest_without_a_report_dir(groups): assert [argv[1:] for argv in plain] == [spec["args"] for spec in groups.values()] -def test_inherited_ci_shard_does_not_restrict_the_reference(monkeypatch, groups): - monkeypatch.setenv("REST_PYTHON_SHARD", "2") +def test_an_inherited_report_dir_does_not_change_the_reference(monkeypatch, groups): monkeypatch.setenv("REST_REPORT_DIR", "inherited-reports") plain = dry_run("test-other-python-rest") assert [argv[1:] for argv in plain] == [spec["args"] for spec in groups.values()] - - -@pytest.mark.parametrize("shard", ["0", "3", "unknown"]) -def test_invalid_shard_cannot_pass_without_running_any_groups(shard): - if shutil.which("make") is None: - pytest.skip("make is not installed") - env = {k: v for k, v in os.environ.items() if k not in PARENT_MAKE_VARIABLES} - result = subprocess.run( - ["make", "-n", "test-other-python-rest", f"REST_PYTHON_SHARD={shard}"], - cwd=REPO, - env=env, - capture_output=True, - text=True, - ) - assert result.returncode != 0 - assert "REST_PYTHON_SHARD must be 1 or 2" in result.stderr + assert shards.make_groups() == { + group: spec["args"] for group, spec in groups.items() + } def import_name(file): @@ -565,11 +911,12 @@ def test_test_files_keep_distinct_import_names(reference): assert not clashes, clashes -# Stands in for pytest on PATH: records its arguments, then exits with the -# next outcome in STUB_OUTCOMES, or kills itself with SIGKILL, as the kernel's -# out-of-memory killer does, for "killed". +# Stands in for pytest on PATH: records its arguments, joined with "|" so a +# split argument shows, then exits with the next outcome in STUB_OUTCOMES, or +# kills itself with SIGKILL, as the kernel's out-of-memory killer does, for +# "killed". STUB_PYTEST = """#!/bin/sh -echo "$*" >> "$STUB_CALLS" +(IFS='|'; echo "$*") >> "$STUB_CALLS" n=$(wc -l < "$STUB_CALLS" | tr -d ' ') outcome=$(echo "$STUB_OUTCOMES" | cut -d ' ' -f "$n") if [ "$outcome" = killed ]; then kill -9 $$; fi @@ -577,15 +924,7 @@ def test_test_files_keep_distinct_import_names(reference): """ -@pytest.mark.parametrize("scenario", ["one-fails-one-killed", "all-fail", "all-pass"]) -def test_every_group_runs_and_any_failure_fails_the_target(tmp_path, groups, scenario): - outcomes = ["0"] * len(groups) - if scenario == "one-fails-one-killed": - assert len(groups) >= 4, "need groups before and after the failures" - outcomes[1] = "1" - outcomes[3] = "killed" - elif scenario == "all-fail": - outcomes = ["killed"] + ["1"] * (len(groups) - 1) +def stub_environment(tmp_path, outcomes): stub = tmp_path / "bin" / "pytest" stub.parent.mkdir() stub.write_text(STUB_PYTEST) @@ -595,6 +934,24 @@ def test_every_group_runs_and_any_failure_fails_the_target(tmp_path, groups, sce env["PATH"] = f"{stub.parent}{os.pathsep}{env['PATH']}" env["STUB_CALLS"] = str(calls) env["STUB_OUTCOMES"] = " ".join(outcomes) + return env, calls + + +def scenario_outcomes(scenario, count): + outcomes = ["0"] * count + if scenario == "one-fails-one-killed": + assert count >= 4, "need processes before and after the failures" + outcomes[1] = "1" + outcomes[3] = "killed" + elif scenario == "all-fail": + outcomes = ["killed"] + ["1"] * (count - 1) + return outcomes + + +@pytest.mark.parametrize("scenario", ["one-fails-one-killed", "all-fail", "all-pass"]) +def test_every_group_runs_and_any_failure_fails_the_target(tmp_path, groups, scenario): + outcomes = scenario_outcomes(scenario, len(groups)) + env, calls = stub_environment(tmp_path, outcomes) result = subprocess.run( ["make", "--no-print-directory", "test-other-python-rest"], cwd=REPO, @@ -603,7 +960,7 @@ def test_every_group_runs_and_any_failure_fails_the_target(tmp_path, groups, sce text=True, ) # Every group ran, with the arguments the dry runs above check. - ran = [" ".join(spec["args"]) for spec in groups.values()] + ran = ["|".join(spec["args"]) for spec in groups.values()] assert calls.read_text().splitlines() == ran failed = [group for group, outcome in zip(groups, outcomes) if outcome != "0"] assert (result.returncode != 0) == bool(failed), result.stdout + result.stderr @@ -611,6 +968,88 @@ def test_every_group_runs_and_any_failure_fails_the_target(tmp_path, groups, sce assert f"Failed Rest Python groups: {' '.join(failed)}" in result.stdout +@pytest.mark.parametrize("scenario", ["one-fails-one-killed", "all-fail", "all-pass"]) +def test_every_planned_process_runs_and_any_failure_fails_the_shard(tmp_path, scenario): + processes = [ + {"name": f"g{i}", "group": f"g{i}", "files": [f"g{i}/test_a.py", "t b.py"]} + for i in range(5) + ] + outcomes = scenario_outcomes(scenario, len(processes)) + env, calls = stub_environment(tmp_path, outcomes) + env["REST_PYTHON_PROCESSES"] = json.dumps(processes) + result = subprocess.run( + [sys.executable, str(PLANNER), "run"], + cwd=REPO, + env=env, + capture_output=True, + text=True, + ) + assert calls.read_text().splitlines() == [ + "|".join([*p["files"], "--maxfail=0"]) for p in processes + ] + failed = [p["name"] for p, outcome in zip(processes, outcomes) if outcome != "0"] + assert (result.returncode != 0) == bool(failed), result.stdout + result.stderr + if failed: + assert f"Failed Rest Python groups: {' '.join(failed)}" in result.stdout + + +def test_a_shard_without_processes_fails(tmp_path): + env, calls = stub_environment(tmp_path, []) + for value in (None, "", "[]"): + if value is None: + env.pop("REST_PYTHON_PROCESSES", None) + else: + env["REST_PYTHON_PROCESSES"] = value + result = subprocess.run( + [sys.executable, str(PLANNER), "run"], + cwd=REPO, + env=env, + capture_output=True, + text=True, + ) + assert result.returncode == 1, (value, result.stdout + result.stderr) + assert not calls.exists() + + +def test_each_process_writes_its_own_reports_in_the_report_dir(monkeypatch): + commands = [] + monkeypatch.setattr( + shards.subprocess, + "run", + lambda command: ( + commands.append(command) or subprocess.CompletedProcess(command, 0) + ), + ) + processes = [ + {"name": "core-2", "group": "core", "files": ["core/test_a.py"]}, + {"name": "policy", "group": "policy", "files": ["p/test_b.py", "p/test_c.py"]}, + ] + assert shards.run_processes(processes, ".") == 0 + assert commands == [ + [ + "/usr/bin/time", + "-v", + "-o", + "./rest-python-core-2-resources.txt", + "pytest", + "core/test_a.py", + "--maxfail=0", + "--junitxml=./rest-python-core-2.xml", + ], + [ + "/usr/bin/time", + "-v", + "-o", + "./rest-python-policy-resources.txt", + "pytest", + "p/test_b.py", + "p/test_c.py", + "--maxfail=0", + "--junitxml=./rest-python-policy.xml", + ], + ] + + def test_conftest_files_cannot_change_which_files_are_collected(old_command): """The listing skips conftest.py; that is exact only while no conftest.py the Rest groups load defines collection hooks, collect_ignore or plugins.""" @@ -678,52 +1117,3 @@ def test_the_check_reports_a_file_in_two_groups(tmp_path): "b/test_two.py is collected by everything but was not in the old " "single process", ] - - -# --- Run as a script: list the files pytest collects, importing none --------- - - -class _Unimported(pytest.File): - """A test file pytest would import, standing in as one placeholder item.""" - - def collect(self): - yield _Placeholder.from_parent(self, name=self.path.name) - - -class _Placeholder(pytest.Item): - def runtest(self): - raise AssertionError("listed for collection only") - - -class _FileListing: - def __init__(self): - self.files = [] - - @pytest.hookimpl(tryfirst=True) - def pytest_pycollect_makemodule(self, module_path, parent): - return _Unimported.from_parent(parent, path=module_path) - - def pytest_collection_finish(self, session): - # Read the kept items, not the modules built: for a file argument - # pytest builds a module for each file beside it, then keeps one. - self.files = [ - Path(os.path.relpath(item.path)).as_posix() for item in session.items - ] - - -if __name__ == "__main__": - output, *pytest_args = sys.argv[1:] - listing = _FileListing() - status = pytest.main( - [ - *pytest_args, - "--collect-only", - "-q", - "--noconftest", - "-p", - "no:cacheprovider", - ], - plugins=[listing], - ) - Path(output).write_text(json.dumps(listing.files)) - sys.exit(status) diff --git a/policyengine_us/tests/rest_python_shards.py b/policyengine_us/tests/rest_python_shards.py new file mode 100644 index 00000000000..ac41129c5b8 --- /dev/null +++ b/policyengine_us/tests/rest_python_shards.py @@ -0,0 +1,505 @@ +#!/usr/bin/env python +"""Plan the Rest Python tests over parallel CI shards, and run one shard. + +`make test-other-python-rest` runs one pytest process per group in +REST_PYTHON_GROUPS (Makefile): modules that retain full systems sit in their +own groups, so no process outgrows the 16 GB runner. CI used to run those +groups on two runners, the second waiting for the first. By October 2026 one +file, test_behavioral_response_measurement_copies.py (27 minutes), had taken +the core group to 50 minutes and the two runners to 130 (CI run 40014592006), +the longest path of every PR. + +`plan` reads the groups from the Makefile and pytest's own collection, and +each file's time from the JUnit reports of a recent CI run. A group longer +than the target splits into chunks of about the target, longest files first, +each chunk keeping the group's file order. Chunks and whole groups then go to +the least-loaded shard, longest first, over as few shards as keep the longest +near the target or the longest single chunk. A process still runs files of +one group only, so none holds more than its group did. A file without a +recorded time counts as its group's median. Without any reports every file +counts the same, and the plan still runs each file once, only less evenly. + +`run` runs one shard's processes in turn, each with its own JUnit and GNU +time report when given a report directory, and fails if any of them fails. +code_health/test_rest_python_groups.py checks that every plan runs each file +of the groups exactly once. +""" + +import argparse +import json +import math +import os +import re +import shlex +import statistics +import subprocess +import sys +import tempfile +from collections import defaultdict +from pathlib import Path +from xml.etree import ElementTree + +REPO = Path(__file__).resolve().parents[2] + +# Each shard aims at about this many minutes. The Baseline and Contrib chain +# takes about 70, so the Python shards finish well inside it. +TARGET_MINUTES = 25 +# Bounds on the shard count whatever the reports say: missing or tiny times +# cannot pile every group onto one runner. +MIN_SHARDS = 4 +MAX_SHARDS = 12 +# A plan whose longest shard is within this of the bound is good enough; +# another runner would only shave a minute or two. +MAKESPAN_TOLERANCE = 1.05 +# Interpreter start, model import and collection, per process, when no +# report measures it (CI run 40014592006: 50-105 s per group). +DEFAULT_PROCESS_SECONDS = 60 +# Per file when no report covers any file of the group: the Rest groups +# averaged 27 s per file in CI run 40014592006. +DEFAULT_FILE_SECONDS = 30 +DEFAULT_VARIABLES_SECONDS = 90 + +# A make run passes these on to the make the planner starts, where they +# would add or remove report options from the groups' commands. +PARENT_MAKE_VARIABLES = ("MAKEFLAGS", "MFLAGS", "MAKELEVEL", "REST_REPORT_DIR") +REPORT_MARKER = "@rest-report-dir@" + + +def make_groups(cwd=REPO): + """Each Rest group's pytest arguments, in run order, from make's dry run. + + The group name comes from its JUnit file, rest-python-.xml. + """ + env = {k: v for k, v in os.environ.items() if k not in PARENT_MAKE_VARIABLES} + result = subprocess.run( + [ + "make", + "-n", + "--no-print-directory", + "test-other-python-rest", + f"REST_REPORT_DIR={REPORT_MARKER}", + ], + cwd=cwd, + env=env, + capture_output=True, + text=True, + ) + if result.returncode != 0: + raise RuntimeError(result.stdout + result.stderr) + groups = {} + junit = re.compile(rf"--junitxml={re.escape(REPORT_MARKER)}/rest-python-(.+)\.xml") + for line in result.stdout.splitlines(): + if not line.startswith("/usr/bin/time "): + continue + argv = shlex.split(line) + if argv[:3] != ["/usr/bin/time", "-v", "-o"] or argv[4] != "pytest": + raise RuntimeError(f"unexpected group command: {line}") + names = [m.group(1) for arg in argv[5:] if (m := junit.fullmatch(arg))] + if len(names) != 1 or names[0] in groups: + raise RuntimeError(f"cannot name the group of: {line}") + args = [arg for arg in argv[5:] if not junit.fullmatch(arg)] + # A shard passes pytest only the collected files and --maxfail=0, so + # any other option (-p, -m, --deselect...) would be lost in CI. + unsupported = [ + arg + for arg in args + if arg.startswith("-") + and arg != "--maxfail=0" + and not arg.startswith("--ignore=") + ] + if unsupported: + raise RuntimeError(f"shards cannot pass on {unsupported}: {line}") + groups[names[0]] = args + if not groups: + raise RuntimeError("make test-other-python-rest runs no groups") + return groups + + +def collected_files(args, cwd=REPO): + """The test files pytest collects for these arguments, relative to cwd. + + Nothing is imported: the `list-files` mode below swaps each test module + for a placeholder and skips conftest.py files. A command that collects + nothing (pytest's exit 5) lists no files. + """ + env = {k: v for k, v in os.environ.items() if k != "PYTEST_ADDOPTS"} + with tempfile.TemporaryDirectory() as tmp: + listing = Path(tmp) / "files.json" + result = subprocess.run( + [sys.executable, __file__, "list-files", str(listing), *args], + cwd=cwd, + env=env, + capture_output=True, + text=True, + ) + if result.returncode not in (0, 5): + raise RuntimeError(f"{args}:\n{result.stdout}{result.stderr}") + return json.loads(listing.read_text()) + + +def group_files(cwd=REPO): + """Each group's test files, in the order its pytest process runs them.""" + return { + group: collected_files(args, cwd) for group, args in make_groups(cwd).items() + } + + +# --- Times from CI reports ----------------------------------------------------- + + +def file_of(classname, files): + """The test file of a JUnit classname, which is its node ID's path with + dots for slashes, followed by any class names.""" + parts = classname.split(".") + for end in range(len(parts), 0, -1): + candidate = "/".join(parts[:end]) + ".py" + if candidate in files: + return candidate + return None + + +def wall_seconds(resources): + """The elapsed time in a GNU time -v report, or None without one.""" + if not resources.exists(): + return None + for line in resources.read_text().splitlines(): + if "Elapsed (wall clock) time" in line: + fields = [float(x) for x in line.rsplit(" ", 1)[1].split(":")] + return sum(x * 60**i for i, x in enumerate(reversed(fields))) + return None + + +def read_timings(directory, files): + """Seconds per test file, per process outside its tests, and for the + variables YAML, from the reports under a directory. + + Files no report covers are absent; the other two are None when no report + measures them. A re-run release run also keeps its earlier attempts' + reports under names the later plan did not reuse, so a file two reports + cover counts once, at the longer time. + """ + known = set(files) + seconds = {} + overheads = [] + variables = None + if directory is not None and Path(directory).is_dir(): + for xml in sorted(Path(directory).rglob("rest-python-*.xml")): + try: + root = ElementTree.parse(xml).getroot() + except ElementTree.ParseError: + continue + total = 0.0 + report = defaultdict(float) + for case in root.iter("testcase"): + time = float(case.get("time") or 0) + total = total + time + file = file_of(case.get("classname", ""), known) + if file is not None: + report[file] = report[file] + time + for file, time in report.items(): + seconds[file] = max(time, seconds.get(file, 0.0)) + wall = wall_seconds(xml.with_name(f"{xml.stem}-resources.txt")) + if wall is not None: + overheads.append(max(0.0, wall - total)) + for resources in Path(directory).rglob("rest-variables-resources.txt"): + variables = wall_seconds(resources) + overhead = statistics.median(overheads) if overheads else None + return seconds, overhead, variables + + +# --- Planning ------------------------------------------------------------------ + + +def estimates(groups, seconds): + """Seconds per file: its recorded time, else its group's median, else the + median over all groups, else DEFAULT_FILE_SECONDS.""" + recorded = [seconds[f] for files in groups.values() for f in files if f in seconds] + fallback = statistics.median(recorded) if recorded else DEFAULT_FILE_SECONDS + result = {} + for files in groups.values(): + known = [seconds[f] for f in files if f in seconds] + group_fallback = statistics.median(known) if known else fallback + for file in files: + result[file] = seconds.get(file, group_fallback) + return result + + +def longest_first(items, bins, weight): + """Deal items, longest first, each to the least-loaded of `bins` bins. + + Ties go to the lowest bin, so the result depends only on the inputs. + """ + loads = [0.0] * bins + contents = [[] for _ in range(bins)] + for item in sorted(items, key=lambda item: -weight(item)): + target = min(range(bins), key=lambda i: (loads[i], i)) + loads[target] = loads[target] + weight(item) + contents[target].append(item) + return contents, loads + + +def chunk_group(group, files, cost, target_seconds): + """Split a group longer than the target into chunks of about the target. + + Each chunk keeps the group's file order, so a chunk's process meets its + files in the order the whole group's process did. + """ + total = sum(cost[f] for f in files) + count = max(1, min(len(files), math.ceil(total / target_seconds))) + if count == 1: + return [(group, list(files))] + order = {f: i for i, f in enumerate(files)} + contents, _ = longest_first(files, count, cost.__getitem__) + chunks = [sorted(chunk, key=order.__getitem__) for chunk in contents if chunk] + chunks.sort(key=lambda chunk: order[chunk[0]]) + return [(f"{group}-{i}", chunk) for i, chunk in enumerate(chunks, start=1)] + + +def make_plan( + groups, + seconds=None, + overhead=None, + variables=None, + target_minutes=TARGET_MINUTES, + min_shards=MIN_SHARDS, + max_shards=MAX_SHARDS, +): + """Shards of processes that run every group's files once. + + Each process runs files of one group, in that group's order. The shard + count is the smallest in [min_shards, max_shards] whose longest shard is + within MAKESPAN_TOLERANCE of the target, or of the longest process when + that is longer. The variables YAML runs on the least-loaded shard. + """ + seconds = seconds or {} + overhead = DEFAULT_PROCESS_SECONDS if overhead is None else overhead + variables = DEFAULT_VARIABLES_SECONDS if variables is None else variables + target = target_minutes * 60 + cost = estimates(groups, seconds) + processes = [] + for group, files in groups.items(): + for name, chunk in chunk_group(group, files, cost, target): + estimate = overhead + sum(cost[f] for f in chunk) + processes.append( + {"name": name, "group": group, "files": chunk, "seconds": estimate} + ) + if not processes: + raise ValueError("no test files to plan") + weight = lambda process: process["seconds"] # noqa: E731 + bound = MAKESPAN_TOLERANCE * max(target, max(map(weight, processes))) + upper = max(1, min(max_shards, len(processes))) + count = min(min_shards, upper) + while True: + contents, loads = longest_first(processes, count, weight) + if max(loads) <= bound or count == upper: + break + count = count + 1 + lightest = min(range(count), key=lambda i: (loads[i], i)) + loads[lightest] = loads[lightest] + variables + order = {p["name"]: i for i, p in enumerate(processes)} + shards = [ + { + "shard": i + 1, + "minutes": round(loads[i] / 60, 1), + "variables": i == lightest, + "processes": [ + {"name": p["name"], "group": p["group"], "files": p["files"]} + for p in sorted(contents[i], key=lambda p: order[p["name"]]) + ], + } + for i in range(count) + ] + check_plan(groups, shards) + return shards + + +def check_plan(groups, shards): + """Raise unless the shards run each group's files once, in group order, + with no process mixing groups and no shard idle.""" + seen = defaultdict(list) + names = set() + for shard in shards: + if not shard["processes"]: + raise AssertionError(f"shard {shard['shard']} runs nothing") + for process in shard["processes"]: + if process["name"] in names: + raise AssertionError(f"process {process['name']} runs twice") + names.add(process["name"]) + group = groups.get(process["group"]) + if group is None or not process["files"]: + raise AssertionError(f"process {process['name']} has no group files") + order = [group.index(f) if f in group else -1 for f in process["files"]] + if -1 in order or order != sorted(order): + raise AssertionError( + f"process {process['name']} runs files outside its group's order" + ) + for file in process["files"]: + seen[file].append(process["name"]) + expected = [f for files in groups.values() for f in files] + missing = sorted(set(expected) - set(seen)) + repeated = sorted(f for f, owners in seen.items() if len(owners) > 1) + if missing or repeated or len(expected) != len(set(expected)): + raise AssertionError(f"missing: {missing}; repeated: {repeated}") + if sum(shard["variables"] for shard in shards) != 1: + raise AssertionError("the variables YAML must run on exactly one shard") + + +def matrix(shards): + """The GitHub Actions matrix: each shard's processes as a JSON string.""" + return { + "include": [ + { + "shard": shard["shard"], + "variables": shard["variables"], + "processes": json.dumps(shard["processes"], separators=(",", ":")), + } + for shard in shards + ] + } + + +def summary(shards, source): + lines = [ + f"Rest Python plan: {len(shards)} shards, times from {source}.", + "", + "| Shard | Estimated minutes | Processes (files) |", + "|---:|---:|---|", + ] + for shard in shards: + processes = ", ".join( + f"{p['name']} ({len(p['files'])})" for p in shard["processes"] + ) + if shard["variables"]: + processes = processes + ", variables YAML" + lines.append(f"| {shard['shard']} | {shard['minutes']} | {processes} |") + return "\n".join(lines) + "\n" + + +# --- Running a shard ----------------------------------------------------------- + + +def run_processes(processes, report_dir=None): + """Run each process in turn; 1 if any failed, was killed or is missing.""" + if not processes: + print("This shard has no Rest Python processes to run.") + return 1 + failed = [] + for process in processes: + name = process["name"] + print(f"=== Rest Python group: {name} ({len(process['files'])} files) ===") + sys.stdout.flush() + command = ["pytest", *process["files"], "--maxfail=0"] + if report_dir is not None: + command = [ + "/usr/bin/time", + "-v", + "-o", + f"{report_dir}/rest-python-{name}-resources.txt", + *command, + f"--junitxml={report_dir}/rest-python-{name}.xml", + ] + if subprocess.run(command).returncode != 0: + failed.append(name) + if failed: + print(f"Failed Rest Python groups: {' '.join(failed)}") + return 1 + return 0 + + +# --- Listing collected files without importing them ---------------------------- + + +def list_files(output, pytest_args): + import pytest + + class Unimported(pytest.File): + """A test file pytest would import, standing in as one item.""" + + def collect(self): + yield Placeholder.from_parent(self, name=self.path.name) + + class Placeholder(pytest.Item): + def runtest(self): + raise AssertionError("listed for collection only") + + class FileListing: + def __init__(self): + self.files = [] + + @pytest.hookimpl(tryfirst=True) + def pytest_pycollect_makemodule(self, module_path, parent): + return Unimported.from_parent(parent, path=module_path) + + def pytest_collection_finish(self, session): + # Read the kept items, not the modules built: for a file argument + # pytest builds a module for each file beside it, then keeps one. + self.files = [ + Path(os.path.relpath(item.path)).as_posix() for item in session.items + ] + + listing = FileListing() + status = pytest.main( + [ + *pytest_args, + "--collect-only", + "-q", + "--noconftest", + "-p", + "no:cacheprovider", + ], + plugins=[listing], + ) + Path(output).write_text(json.dumps(listing.files)) + return int(status) + + +def main(argv=None): + argv = sys.argv[1:] if argv is None else argv + if argv[:1] == ["list-files"]: + return list_files(argv[1], argv[2:]) + parser = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) + commands = parser.add_subparsers(dest="command", required=True) + plan = commands.add_parser("plan", help="plan the shards") + plan.add_argument("--timings", help="directory of CI timing reports") + plan.add_argument("--source", default="no reports", help="where they came from") + plan.add_argument("--github-output", help="file to append matrix= to") + plan.add_argument("--summary", help="file to append a Markdown summary to") + plan.add_argument("--target-minutes", type=float, default=TARGET_MINUTES) + run = commands.add_parser("run", help="run one shard's processes") + run.add_argument( + "--processes-env", + default="REST_PYTHON_PROCESSES", + help="environment variable holding the shard's processes as JSON", + ) + run.add_argument("--report-dir", help="write JUnit and GNU time reports here") + args = parser.parse_args(argv) + + if args.command == "run": + processes = json.loads(os.environ.get(args.processes_env) or "[]") + return run_processes(processes, args.report_dir) + + groups = group_files() + files = [f for group in groups.values() for f in group] + seconds, overhead, variables = read_timings(args.timings, files) + source = args.source + if not seconds: + print("::warning::No Rest Python timing reports; every file counts the same.") + source = "no reports (every file counts the same)" + else: + covered = sum(f in seconds for f in files) + source = f"{source} ({covered} of {len(files)} files timed)" + shards = make_plan(groups, seconds, overhead, variables, args.target_minutes) + text = summary(shards, source) + print(text) + print(json.dumps(shards, indent=1)) + if args.github_output: + with open(args.github_output, "a") as output: + output.write(f"matrix={json.dumps(matrix(shards))}\n") + if args.summary: + with open(args.summary, "a") as output: + output.write(text) + return 0 + + +if __name__ == "__main__": + sys.exit(main())