Split the CI Test job into shards and cut its runtime

.github/check_test_shards.py

@@ -0,0 +1,195 @@
+"""Check that the shards of the Test job ran every test exactly once.
+
+CI splits ``pytest policyengine_uk/tests`` over several runners with
+``--shard K/N`` (see policyengine_uk/tests/conftest.py). This script compares
+what the shards ran with what an unsharded ``pytest --collect-only -q`` lists:
+
+    python .github/check_test_shards.py \
+        --collected collected.txt --junit junit-1.xml junit-2.xml junit-3.xml
+
+It fails if a collected test is missing from every shard, ran in more than one,
+ran without being collected, or failed. So splitting the suite cannot drop a
+test without the Test job going red.
+
+Standard library only: the job that runs it does not install the package.
+"""
+
+import argparse
+import re
+import sys
+import xml.etree.ElementTree as ElementTree
+from collections import Counter
+from pathlib import Path
+
+_COLLECTED_SUMMARY = re.compile(r"^(\d+) tests? collected\b(.*)$")
+_NOTHING_COLLECTED = re.compile(r"^no tests collected\b")
+
+
+class ShardCheckError(Exception):
+    """An input that cannot be checked, as opposed to a check that failed."""
+
+
+def junit_key(nodeid: str) -> tuple[str, str]:
+    """The ``(classname, name)`` pytest gives ``nodeid`` in JUnit XML.
+
+    Follows _pytest.junitxml.mangle_test_address; a test compares the two.
+    pytest also rewrites characters that XML does not allow, such as control
+    characters. Its node ids do not contain them, because it escapes them in
+    parameter ids. If one ever did, the names would differ and the check
+    would fail, not pass.
+    """
+    path, bracket, params = nodeid.partition("[")
+    names = path.split("::")
+    names[0] = re.sub(r"\.py$", "", names[0].replace("/", "."))
+    names[-1] += bracket + params
+    return ".".join(names[:-1]), names[-1]
+
+
+def read_collected(text: str) -> list[str]:
+    """Node ids from the output of ``pytest --collect-only -q``.
+
+    The listing is one node id per line, a blank line, then any warnings and
+    ``N tests collected``. The ids found must number N, so a listing this
+    cannot parse is an error, never a shorter list.
+    """
+    ids = []
+    total = None
+    listing_ended = False
+    for line in text.splitlines():
+        summary = _COLLECTED_SUMMARY.match(line)
+        if summary:
+            if "error" in summary.group(2):
+                raise ShardCheckError(f"collection reported errors: {line!r}")
+            total = int(summary.group(1))
+        elif _NOTHING_COLLECTED.match(line):
+            total = 0
+        elif not line.strip():
+            listing_ended = listing_ended or bool(ids)
+        elif "::" in line and not listing_ended:
+            ids.append(line)
+    if total is None:
+        raise ShardCheckError("no 'N tests collected' line in the collection listing")
+    if len(ids) != total:
+        raise ShardCheckError(
+            f"the collection listing says {total} tests but lists {len(ids)}"
+        )
+    if len(set(ids)) != len(ids):
+        raise ShardCheckError("the collection listing repeats a node id")
+    return ids
+
+
+def read_junit(path: Path) -> list[tuple[tuple[str, str], str]]:
+    """``((classname, name), outcome)`` for each test case in a JUnit file.
+
+    The outcome is ``failed`` for a failure or an error, else ``skipped`` for a
+    skip or an expected failure, else ``passed``.
+    """
+    cases = []
+    for case in ElementTree.parse(path).getroot().iter("testcase"):
+        children = {child.tag for child in case}
+        if children & {"failure", "error"}:
+            outcome = "failed"
+        elif "skipped" in children:
+            outcome = "skipped"
+        else:
+            outcome = "passed"
+        cases.append(((case.get("classname", ""), case.get("name", "")), outcome))
+    return cases
+
+
+def check(collected: list[str], shards: list[list[tuple[tuple[str, str], str]]]):
+    """Compare the collected node ids with what each shard ran.
+
+    Returns ``(problems, summary)``: ``problems`` is empty when every collected
+    test ran in exactly one shard and none failed.
+    """
+    expected = {}
+    for nodeid in collected:
+        key = junit_key(nodeid)
+        if key in expected:
+            raise ShardCheckError(
+                f"{nodeid!r} and {expected[key]!r} have the same JUnit name"
+            )
+        expected[key] = nodeid
+
+    runs = Counter()
+    outcomes = Counter()
+    failed = []
+    for number, cases in enumerate(shards, start=1):
+        for key, outcome in cases:
+            runs[key] += 1
+            outcomes[outcome] += 1
+            if outcome == "failed":
+                failed.append((number, key))
+
+    def name(key):
+        return expected.get(key) or f"{key[0]}::{key[1]}"
+
+    problems = []
+    missing = sorted(expected[key] for key in expected if key not in runs)
+    if missing:
+        problems.append((f"{len(missing)} collected tests ran in no shard", missing))
+    repeated = sorted(name(key) for key, count in runs.items() if count > 1)
+    if repeated:
+        problems.append((f"{len(repeated)} tests ran more than once", repeated))
+    extra = sorted(name(key) for key in runs if key not in expected)
+    if extra:
+        problems.append((f"{len(extra)} tests ran but were not collected", extra))
+    if failed:
+        problems.append(
+            (
+                f"{len(failed)} tests failed",
+                [f"shard {number}: {name(key)}" for number, key in failed],
+            )
+        )
+
+    sizes = "/".join(str(len(cases)) for cases in shards)
+    summary = (
+        f"{len(collected)} tests collected; {len(shards)} shards ran "
+        f"{sizes} = {sum(len(cases) for cases in shards)} "
+        f"({outcomes['passed']} passed, {outcomes['skipped']} skipped or "
+        f"expected failures, {outcomes['failed']} failed)"
+    )
+    return problems, summary
+
+
+def main(argv=None) -> int:
+    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
+    parser.add_argument(
+        "--collected",
+        required=True,
+        type=Path,
+        help="output of an unsharded `pytest --collect-only -q`",
+    )
+    parser.add_argument(
+        "--junit",
+        required=True,
+        nargs="+",
+        type=Path,
+        help="the JUnit XML file of each shard",
+    )
+    args = parser.parse_args(argv)
+
+    try:
+        collected = read_collected(args.collected.read_text(encoding="utf-8"))
+        shards = [read_junit(path) for path in args.junit]
+        problems, summary = check(collected, shards)
+    except (ShardCheckError, OSError, ElementTree.ParseError) as error:
+        print(f"Cannot check the shards: {error}", file=sys.stderr)
+        return 2
+
+    print(summary)
+    for title, names in problems:
+        print(f"\n{title}:")
+        for line in names[:50]:
+            print(f"  {line}")
+        if len(names) > 50:
+            print(f"  ... and {len(names) - 50} more")
+    if problems:
+        return 1
+    print("Every collected test ran in exactly one shard.")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

.github/workflows/code_changes.yaml

@@ -31,19 +31,20 @@ jobs:
           uses: astral-sh/setup-uv@v8.1.0
         - name: Check formatting
           run: uvx ruff format --check .
-  Test:
+  # The tests run as four jobs: the YAML policy tests, and the Python tests
+  # split over three runners. As one job they took 59 minutes on 10 October
+  # 2026 (run 38075384588: 2.6 minutes of YAML tests, then 55.8 minutes of
+  # pytest on four workers), and runs were cancelled at a 60-minute limit
+  # (issues #2099 and #2257).
+  Test-YAML:
+      name: Test (YAML)
       # Linux arm64 keeps the CPU architecture of the macOS runners this job
       # used to run on. GitHub-hosted macOS runners are scarce, so PR runs
       # queued for hours; the parameter tree and every test result are
       # identical on both.
       runs-on: ubuntu-24.04-arm
-      # The whole job took about 18 minutes when this limit was set; by
-      # October 2026 it takes about an hour (57 to 60 minutes on 9 October,
-      # and every run on 10 October was cancelled at a 60-minute limit, one
-      # at 99% with every test passing; issue #2257). 120 minutes lets a
-      # slow run finish and still stops a hung one long before the 6-hour
-      # GitHub limit. The suite's runtime needs its own attention.
-      timeout-minutes: 120
+      # About 4 minutes on 10 October (run 38085920460).
+      timeout-minutes: 30
       permissions:
         contents: "read"
         id-token: "write"
@@ -61,7 +62,70 @@ jobs:
             with:
               workload_identity_provider: "projects/322898545428/locations/global/workloadIdentityPools/policyengine-research-id-pool/providers/prod-github-provider"
               service_account: "policyengine-research@policyengine-research.iam.gserviceaccount.com"
-
+
+          - name: Install package
+            run: uv pip install -e .[dev] --system
+          - name: Install policyengine
+            run: uv pip install policyengine --system
+          - name: Install JB
+            run: uv pip install "jupyter-book>=2.0.0a0" --system
+          - name: UV sync
+            run: uv sync --extra dev
+          - name: Run YAML tests
+            run: make test-yaml
+          - name: List the Python tests
+            # Every Python test, collected without --shard. The Test job
+            # checks that the shards below ran each of these exactly once.
+            run: |
+              # Only export token if non-empty (defensive against missing secrets)
+              if [ -n "$HF_TOKEN_VALUE" ]; then
+                export HUGGING_FACE_TOKEN="$HF_TOKEN_VALUE"
+              fi
+              pytest policyengine_uk/tests --collect-only -q -p no:cacheprovider > collected.txt
+              tail -n 1 collected.txt
+            env:
+              HF_TOKEN_VALUE: ${{ secrets.HUGGING_FACE_TOKEN }}
+          - name: Upload the list of Python tests
+            uses: actions/upload-artifact@v7
+            with:
+              name: test-collected
+              path: collected.txt
+              if-no-files-found: error
+              overwrite: true
+              retention-days: 7
+  Test-Python:
+      # Shard K of N runs the tests whose node id hashes to K (--shard in
+      # policyengine_uk/tests/conftest.py), so each test is in exactly one
+      # shard. N is the number of jobs this matrix makes: to add a shard, add
+      # a number to the list.
+      name: Test (Python ${{ matrix.shard }}/${{ strategy.job-total }})
+      runs-on: ubuntu-24.04-arm
+      # The shards took 15, 20 and 15 minutes on 10 October (run
+      # 38085920460). 60 minutes leaves room for a slow runner and still
+      # stops a hung shard.
+      timeout-minutes: 60
+      strategy:
+        fail-fast: false
+        matrix:
+          shard: [1, 2, 3]
+      permissions:
+        contents: "read"
+        id-token: "write"
+      steps:
+          - name: Checkout repo
+            uses: actions/checkout@v6
+          - name: Install uv
+            uses: astral-sh/setup-uv@v8.1.0
+
+          - name: Set up Python
+            uses: actions/setup-python@v6
+            with:
+                python-version: '3.13'
+          - uses: "google-github-actions/auth@v3"
+            with:
+              workload_identity_provider: "projects/322898545428/locations/global/workloadIdentityPools/policyengine-research-id-pool/providers/prod-github-provider"
+              service_account: "policyengine-research@policyengine-research.iam.gserviceaccount.com"
+
           - name: Install package
             run: uv pip install -e .[dev] --system
           - name: Install policyengine
@@ -78,12 +142,54 @@ jobs:
               fi
               # Four pytest-xdist workers, one per core; together they peak at
               # about 9 GB of the runner's 16 GB.
-              make test PYTEST_ARGS="-n 4 --dist worksteal"
+              make test-python PYTEST_ARGS="-n 4 --dist worksteal --shard ${{ matrix.shard }}/${{ strategy.job-total }} --junitxml=junit-${{ matrix.shard }}.xml"
             env:
               HF_TOKEN_VALUE: ${{ secrets.HUGGING_FACE_TOKEN }}
+              # Measure coverage with sys.monitoring instead of the C trace
+              # function: the same lines are reported, at less cost per line.
+              COVERAGE_CORE: sysmon
+          - name: Upload the test results
+            if: ${{ !cancelled() }}
+            uses: actions/upload-artifact@v7
+            with:
+              name: test-junit-${{ matrix.shard }}
+              path: junit-${{ matrix.shard }}.xml
+              if-no-files-found: error
+              overwrite: true
+              retention-days: 7
           - name: Upload coverage to Codecov
             uses: codecov/codecov-action@v6
             with:
               files: ./coverage.xml
               fail_ci_if_error: false
               verbose: true
+  Test:
+      # One result for the whole suite. It fails if any test job failed, or
+      # if the shards did not run every collected Python test exactly once.
+      needs: [Test-YAML, Test-Python]
+      if: ${{ !cancelled() }}
+      # The shards' runners, not ubuntu-latest: on 10 October this 5-second
+      # job waited 32 minutes for an ubuntu-latest runner after the shards
+      # had finished (run 38085920460).
+      runs-on: ubuntu-24.04-arm
+      timeout-minutes: 10
+      steps:
+          - name: Check that every test job passed
+            run: |
+              echo "YAML: ${{ needs.Test-YAML.result }}; Python: ${{ needs.Test-Python.result }}"
+              test "${{ needs.Test-YAML.result }}" = success
+              test "${{ needs.Test-Python.result }}" = success
+          - name: Checkout repo
+            if: ${{ !cancelled() }}
+            uses: actions/checkout@v6
+          - name: Download the test lists and results
+            if: ${{ !cancelled() }}
+            uses: actions/download-artifact@v8
+            with:
+              pattern: test-*
+              merge-multiple: true
+          - name: Check that every Python test ran in exactly one shard
+            # Runs even when a test job failed, so its report names the
+            # failed tests.
+            if: ${{ !cancelled() }}
+            run: python3 .github/check_test_shards.py --collected collected.txt --junit junit-*.xml