Deal the slowest tests to CI shards by duration, not hash

.github/test_durations.py

@@ -0,0 +1,90 @@
+"""Measure the slowest Python tests from CI logs, for shard_durations.toml.
+
+The shards of the Test job run with ``pytest -v -n 4``, so each finished test
+logs a line such as::
+
+    2026-10-24T10:57:20.4034260Z [gw0] [  0%] PASSED path.py::test_name
+
+GitHub stamps every line with the time it was written. A pytest-xdist worker
+runs one test at a time, so the time between two consecutive lines from the
+same worker is the second test's duration, set-up and tear-down included.
+This reads the logs of one or more shard jobs, saved with::
+
+    gh run view RUN --log --job JOB > RUN-K.log
+
+averages each test's duration over the logs it appears in, and prints the
+tests that took at least --min-seconds as the TOML table conftest.py reads.
+
+Standard library only, like the other scripts here.
+"""
+
+import argparse
+import re
+import statistics
+import sys
+from collections import defaultdict
+from datetime import datetime
+
+LINE = re.compile(
+    r"(\d{4}-\d\d-\d\dT\d\d:\d\d:\d\d(?:\.\d+)?)Z "
+    r"\[(gw\d+)\] \[\s*\d+%\] (?:PASSED|FAILED|SKIPPED|XFAIL|XPASS|ERROR) (.+?)"
+    # The job's once-a-minute memory report can land on the same line.
+    r"(?: memory: \d+ MB used, \d+ MB available)? *$"
+)
+# pytest-xdist logs this once its workers are up; a worker's first test is
+# timed from here.
+START = re.compile(
+    r"(\d{4}-\d\d-\d\dT\d\d:\d\d:\d\d(?:\.\d+)?)Z created: \d+/\d+ workers"
+)
+
+
+def _time(stamp: str) -> datetime:
+    # Seven fractional digits: datetime reads at most six.
+    whole, _, fraction = stamp.partition(".")
+    return datetime.fromisoformat(f"{whole}.{(fraction + '000000')[:6]}")
+
+
+def durations(lines) -> dict[str, float]:
+    """Each test's duration in seconds in one shard's log."""
+    start = None
+    last = {}
+    found = {}
+    for line in lines:
+        if start is None and (match := START.search(line)):
+            start = _time(match[1])
+        if match := LINE.search(line):
+            stamp, worker, nodeid = _time(match[1]), match[2], match[3]
+            previous = last.get(worker, start or stamp)
+            found[nodeid] = (stamp - previous).total_seconds()
+            last[worker] = stamp
+    return found
+
+
+def main(argv=None) -> int:
+    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
+    parser.add_argument("logs", nargs="+", help="saved logs of shard jobs")
+    parser.add_argument("--min-seconds", type=float, default=20)
+    args = parser.parse_args(argv)
+    seen = defaultdict(list)
+    for path in args.logs:
+        with open(path, encoding="utf-8", errors="replace") as log:
+            for nodeid, seconds in durations(log).items():
+                seen[nodeid].append(seconds)
+    mean = {nodeid: statistics.mean(values) for nodeid, values in seen.items()}
+    slow = sorted(
+        ((seconds, nodeid) for nodeid, seconds in mean.items()),
+        key=lambda item: (-item[0], item[1]),
+    )
+    print("[seconds]")
+    for seconds, nodeid in slow:
+        if seconds < args.min_seconds:
+            break
+        key = nodeid.replace("\\", "\\\\").replace('"', '\\"')
+        print(f'"{key}" = {round(seconds)}')
+    total = sum(mean.values()) / 60
+    print(f"# {len(mean)} tests, {total:.0f} worker-minutes", file=sys.stderr)
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

.github/workflows/test-suite.yaml

@@ -135,6 +135,13 @@ jobs:
     # shard. N is the number of jobs this matrix makes: to add a shard, add
     # a number to the list.
     #
+    # The slowest tests, listed with their CI durations in
+    # policyengine_uk/tests/shard_durations.toml, are dealt out by duration
+    # instead of hashed. With the hash alone, over 95 runs on 24 October
+    # shard 2/5 averaged 57 minutes (up to 70) and shard 1/5 49: a 22-minute
+    # property test and several of 4 to 7 minutes had hashed to shard 2.
+    # .github/test_durations.py regenerates the table from shard logs.
+    #
     # Four shards, not three, for memory: each pytest-xdist worker's memory
     # grows with the tests it runs, and with three shards of four workers
     # the 16 GB runners peaked at 15.6 GB, 283 MB short of running out

changelog.d/ci-balance-test-shards.changed.md

@@ -0,0 +1 @@
+Balance the CI Python test shards: the slowest tests, listed with their measured CI durations, are dealt out by duration instead of by hash, so no shard carries most of them.

policyengine_uk/tests/code_health/test_ci_test_shards.py

@@ -36,6 +36,10 @@
 11. ``changed_code.py`` calls a change code exactly when a path matches a
     pattern as a workflow ``paths`` filter would (``dir/**`` is every file
     under ``dir/``), and calls an empty change list code.
+12. The slow tests in ``shard_durations.toml`` are dealt out longest first to
+    the least loaded shard, whatever order the table lists them in, so no
+    shard's listed load exceeds another's by more than the longest test; the
+    table parses and each entry names a test function that exists.
 """
 
 import ast
@@ -171,6 +175,76 @@ def test_parse_shard_rejects_anything_else(conftest, value):
         conftest.parse_shard(value)
 
 
+# Invariant 12: the slowest tests are dealt out by duration.
+
+_durations = st.dictionaries(nodeids, st.floats(0, 3_000, allow_nan=False), max_size=40)
+
+
+@PROPERTY_SETTINGS
+@given(durations=_durations, shards=st.integers(1, 8), data=st.data())
+def test_balanced_shards_bound_the_gap_by_the_longest_test(
+    conftest, durations, shards, data
+):
+    placed = conftest.balanced_shards(durations, shards)
+    assert placed.keys() == durations.keys()
+    assert all(1 <= shard <= shards for shard in placed.values())
+    load = [0.0] * shards
+    for nodeid, shard in placed.items():
+        load[shard - 1] += durations[nodeid]
+    assert max(load) - min(load) <= max(durations.values(), default=0) + 1e-6
+    # The same placement whatever order the table lists the tests in.
+    reordered = dict(data.draw(st.permutations(list(durations.items()))))
+    assert conftest.balanced_shards(reordered, shards) == placed
+
+
+@PROPERTY_SETTINGS
+@given(ids=st.lists(nodeids, max_size=40, unique=True), shards=st.integers(1, 8))
+def test_shards_partition_a_collection_with_listed_tests(conftest, ids, shards):
+    listed = list(conftest.read_shard_durations())[:20]
+    collection = sorted(set(ids) | set(listed))
+    members = {
+        shard: [
+            nodeid
+            for nodeid in collection
+            if conftest.shard_of(nodeid, shards) == shard
+        ]
+        for shard in range(1, shards + 1)
+    }
+    assert sorted(sum(members.values(), [])) == collection
+
+
+def test_listed_tests_follow_the_table(conftest):
+    durations = conftest.read_shard_durations()
+    for shards in (1, 3, 5, 6):
+        placed = conftest.balanced_shards(durations, shards)
+        assert all(conftest.shard_of(n, shards) == k for n, k in placed.items())
+
+
+def test_the_checked_in_shard_durations_name_tests_that_exist(conftest):
+    durations = conftest.read_shard_durations()
+    assert durations
+    assert all(seconds > 0 for seconds in durations.values())
+    for nodeid in durations:
+        path, _, rest = nodeid.partition("::")
+        file = REPO_ROOT / path
+        assert file.exists(), nodeid
+        names = rest.split("::")
+        scope = ast.parse(file.read_text(encoding="utf-8")).body
+        for name in names[:-1]:
+            [cls] = [n for n in scope if isinstance(n, ast.ClassDef) and n.name == name]
+            scope = cls.body
+        functions = {
+            n.name
+            for n in scope
+            if isinstance(n, (ast.FunctionDef, ast.AsyncFunctionDef))
+        }
+        assert names[-1].split("[")[0] in functions, nodeid
+
+
+def test_a_missing_table_leaves_every_test_to_the_hash(conftest, tmp_path):
+    assert conftest.read_shard_durations(tmp_path / "absent.toml") == {}
+
+
 # Invariant 4: the collection listing.
 
 

policyengine_uk/tests/conftest.py

@@ -6,6 +6,7 @@
 import os
 import re
 import sys
+import tomllib
 import weakref
 from dataclasses import dataclass
 from importlib.machinery import SourceFileLoader
@@ -189,15 +190,66 @@ def cloned_uk_tax_benefit_system():
 
 
 # CI splits the suite over several runners with ``--shard K/N``. The shard of a
-# test is a function of its node id alone, so the shards partition whatever is
-# collected: every test runs in exactly one of them, including a test added
-# later. The Test job checks that on each run (.github/check_test_shards.py).
+# test is a function of its node id alone (and of the checked-in table below),
+# so the shards partition whatever is collected: every test runs in exactly
+# one of them, including a test added later. The Test job checks that on each
+# run (.github/check_test_shards.py).
+#
+# Most tests are placed by a hash of their node id. A hash spreads thousands
+# of short tests evenly, but not a few long ones: on 24 October 2026 one
+# property test took 22 minutes, and by chance it and several others of 4 to
+# 7 minutes hashed to shard 2 of 5, which then took 57 minutes on average
+# against 49 for shard 1 (95 runs). The tests in shard_durations.toml, the
+# slowest on CI, are instead dealt out by duration, longest first, each to the
+# shard with the least so far. A test missing from the table, or the table
+# missing, falls back to the hash.
+SHARD_DURATIONS = (
+    Path(importlib.util.find_spec("policyengine_uk").origin).parent
+    / "tests"
+    / "shard_durations.toml"
+)
+
+
+def read_shard_durations(path: Path = SHARD_DURATIONS) -> dict[str, float]:
+    """Seconds each listed test took on CI, by node id; empty without a table."""
+    try:
+        text = path.read_text(encoding="utf-8")
+    except FileNotFoundError:
+        return {}
+    return {
+        nodeid: float(seconds)
+        for nodeid, seconds in tomllib.loads(text).get("seconds", {}).items()
+    }
+
+
+def balanced_shards(durations: Mapping[str, float], shards: int) -> dict[str, int]:
+    """Deal tests to shards, longest first, each to the least loaded shard.
+
+    Ties go to the lower node id and the lower shard, so every process gives
+    the same answer.
+    """
+    load = [0.0] * shards
+    placed = {}
+    for nodeid, seconds in sorted(durations.items(), key=lambda kv: (-kv[1], kv[0])):
+        shard = min(range(shards), key=lambda k: (load[k], k))
+        load[shard] += seconds
+        placed[nodeid] = shard + 1
+    return placed
+
+
+_BALANCED: dict[int, dict[str, int]] = {}
+
+
 def shard_of(nodeid: str, shards: int) -> int:
     """The shard, from 1 to ``shards``, that runs the test ``nodeid``.
 
     Uses SHA-256 rather than ``hash``, which Python salts per process: every
     runner and every pytest-xdist worker must give the same answer.
     """
+    if shards not in _BALANCED:
+        _BALANCED[shards] = balanced_shards(read_shard_durations(), shards)
+    if nodeid in _BALANCED[shards]:
+        return _BALANCED[shards][nodeid]
     digest = hashlib.sha256(nodeid.encode("utf-8")).digest()
     return int.from_bytes(digest[:8], "big") % shards + 1