Run the tests that need the dataset in their own CI shard

.github/test_durations.py

@@ -28,8 +28,9 @@
 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)? *$"
+    # The job's once-a-minute memory report can land on the same line. It
+    # counts MiB; logs from before November 2026 label it MB.
+    r"(?: memory: \d+ Mi?B used, \d+ Mi?B available)? *$"
 )
 # pytest-xdist logs this once its workers are up; a worker's first test is
 # timed from here.

.github/workflows/test-suite.yaml

@@ -4,10 +4,10 @@
 # publish. All three run the same jobs in the same locked environment.
 #
 # Jobs: the YAML policy tests (which also list every Python test), the Python
-# tests split over seven runners, and Result, which fails unless every listed
+# tests split over eight runners, and Result, which fails unless every listed
 # Python test ran in exactly one shard and passed, apart from the outcomes
 # accepted in .github/test_outcome_allowlist.toml. A caller that names its job
-# Test gets the checks "Test / YAML", "Test / Python K/7" and "Test / Result";
+# Test gets the checks "Test / YAML", "Test / Python K/8" and "Test / Result";
 # "Test / Result" is the one to require.
 #
 # Runners. These jobs, and the check jobs that pull requests, merge queue
@@ -143,7 +143,8 @@ jobs:
           retention-days: 7
 
   Python:
-    # Shard K of N runs the tests whose node id hashes to K (--shard in
+    # The last shard runs the tests that need the dataset. Each other shard K
+    # of N runs the remaining 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.
@@ -182,6 +183,18 @@ jobs:
     # cancelled shard 1/6 with its tests at 99% (run 41191216150, first
     # attempt). With seven, the listed slow tests come to about 115 minutes
     # of worker time per shard instead of 134.
+    #
+    # Eight, not seven, for memory: the eighth runs the tests that need the
+    # dataset, with one worker, so only one simulation of the dataset runs at
+    # a time, and the other seven run none of them. In merge queue run
+    # 42212587066 the five shards that built the fixture in
+    # test_aggregate_impact_certified.py peaked at 14.2 to 15.3 GiB, and the
+    # two without it at 10.9 and 12.6 GiB. On a merge of 28 pull requests with
+    # those tests on all seven shards, runs 42218280596 and 42218280615 lost
+    # eight of their 14 shard jobs within minutes at 15.1 to 15.4 GiB, and
+    # shard 6/7 of the first lost its runner at 15.2 GiB 51 minutes in, with
+    # test_latest_data_smoke.py running. With the certified tests on one shard
+    # of seven, run 42231098419 lost that shard at 15.2 GiB.
     name: Python ${{ matrix.shard }}/${{ strategy.job-total }}
     if: ${{ !inputs.skip }}
     runs-on: ${{ vars.CI_RUNNER_ARM != '' && (vars.CI_RUNNER_EVENTS == '' || contains(format(',{0},', vars.CI_RUNNER_EVENTS), format(',{0},', github.event_name))) && vars.CI_RUNNER_ARM || 'ubuntu-24.04-arm' }}
@@ -203,7 +216,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        shard: [1, 2, 3, 4, 5, 6, 7]
+        shard: [1, 2, 3, 4, 5, 6, 7, 8]
     steps:
       - name: Checkout repo
         uses: actions/checkout@v6
@@ -230,6 +243,7 @@ jobs:
           POLICYENGINE_UK_REQUIRE_TEST_DATA: ${{ inputs.require-data && '1' || '' }}
           SHARD: ${{ matrix.shard }}/${{ strategy.job-total }}
           OUTCOMES: outcomes-${{ matrix.shard }}.jsonl
+          WORKERS: ${{ matrix.shard == strategy.job-total && 1 || 4 }}
           # Measure coverage with sys.monitoring instead of the C trace
           # function: the same lines are reported, at less cost per line.
           COVERAGE_CORE: sysmon
@@ -238,9 +252,10 @@ jobs:
             export HUGGING_FACE_TOKEN="$HF_TOKEN_VALUE"
           fi
           # Log memory each minute, so a runner lost to memory pressure
-          # ("The runner has received a shutdown signal") shows it.
+          # ("The runner has received a shutdown signal") shows it. free -m
+          # counts MiB.
           (while sleep 60; do
-            free -m | awk 'NR == 2 { print "memory: " $3 " MB used, " $7 " MB available" }'
+            free -m | awk 'NR == 2 { print "memory: " $3 " MiB used, " $7 " MiB available" }'
           done) &
           trap 'kill $! 2>/dev/null' EXIT
           coverage=()
@@ -250,14 +265,15 @@ jobs:
           # .github/pytest_outcome_log.py records each test's outcome for
           # Result. xfail_strict makes an xfail marker without strict= fail
           # when its test passes, so a fixed bug cannot hide behind it.
-          # Four pytest-xdist workers, one per core. Their memory grows with
-          # each test they run (#2330): with three shards they peaked at
-          # 15.6 GB of the runner's 16 GB (run 38156967136).
+          # Four pytest-xdist workers, one per core, and one on the last shard.
+          # Their memory grows with each test they run (#2330): with three
+          # shards they peaked at 15.3 GiB of the 15.6 GiB the runner reports
+          # as used and available (run 38156967136).
           PYTHONPATH="$GITHUB_WORKSPACE/.github" pytest "$TESTS" \
             --rootdir "$ROOTDIR" -c pyproject.toml -p no:cacheprovider \
             -p pytest_outcome_log --outcome-log="$OUTCOMES" \
             -o xfail_strict=true --maxfail=0 -v \
-            -n 4 --dist worksteal --shard "$SHARD" "${coverage[@]}"
+            -n "$WORKERS" --dist worksteal --shard "$SHARD" "${coverage[@]}"
       - name: Upload the test outcomes
         if: ${{ !cancelled() }}
         uses: actions/upload-artifact@v7

Makefile

@@ -24,7 +24,7 @@ test-yaml:
 
 # Every Python test, serially. To spread them over four cores add
 # PYTEST_ARGS="-n 4 --dist worksteal" (pytest-xdist); to run what one CI
-# shard runs, add "--shard K/7" (policyengine_uk/tests/conftest.py). CI runs
+# shard runs, add "--shard K/8" (policyengine_uk/tests/conftest.py). CI runs
 # these tests through .github/workflows/test-suite.yaml, in the locked
 # environment (uv sync --frozen --extra dev), and also checks each test's
 # outcome against .github/test_outcome_allowlist.toml.

changelog.d/data-tests-own-shard.changed.md

@@ -0,0 +1 @@
+Run the tests that need the dataset in their own CI shard, with one worker, so the other seven shards no longer build simulations of the whole dataset.

docs/engineering/skills/testing.md

@@ -86,15 +86,17 @@ the same suite,
 tree; a release tests the wheel it is about to publish, installed on its own.
 
 The suite runs the YAML policy tests in one job and splits the Python tests
-over seven runners. Each runner runs one shard:
+over eight runners. Each runner runs one shard:
 
 ```bash
-pytest policyengine_uk/tests -n 4 --dist worksteal --shard K/7 -o xfail_strict=true
+pytest policyengine_uk/tests -n 4 --dist worksteal --shard K/8 -o xfail_strict=true
 ```
 
-`--shard K/N` is defined in `policyengine_uk/tests/conftest.py`. It keeps the
-tests whose node id hashes to shard K of N, so every test is in exactly one
-shard, including a test added later. Each shard records every test's outcome
+`--shard K/N` is defined in `policyengine_uk/tests/conftest.py`. Shard N of N
+keeps the tests with the `microsimulation` marker, which need the dataset, and
+CI runs it with one worker. Every other shard keeps the remaining tests whose
+node id hashes to it, so every test is in exactly one shard, including a test
+added later. Each shard records every test's outcome
 with `.github/pytest_outcome_log.py`. The final job, `Test / Result`, runs
 `.github/check_test_shards.py`. It fails unless every collected test ran in
 exactly one shard and passed. The exceptions are outcomes accepted in

policyengine_uk/tests/code_health/test_ci_test_shards.py

@@ -40,6 +40,9 @@
     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.
+13. With two or more shards, the last runs exactly the tests that need the
+    dataset, and ``shard_of`` over one shard fewer places the others. The
+    workflow numbers its shards from 1 to N and gives the last one one worker.
 """
 
 import ast
@@ -245,6 +248,56 @@ def test_a_missing_table_leaves_every_test_to_the_hash(conftest, tmp_path):
     assert conftest.read_shard_durations(tmp_path / "absent.toml") == {}
 
 
+# Invariant 13: the last shard runs the tests that need the dataset.
+
+
+class _CollectedTest:
+    """What item_shard reads from a pytest item."""
+
+    def __init__(self, nodeid, needs_data):
+        self.nodeid = nodeid
+        self.needs_data = needs_data
+
+    def get_closest_marker(self, name):
+        return object() if self.needs_data and name == "microsimulation" else None
+
+
+@PROPERTY_SETTINGS
+@given(
+    collection=st.dictionaries(nodeids, st.booleans(), max_size=40),
+    shards=st.integers(2, 8),
+)
+def test_the_last_shard_runs_exactly_the_tests_that_need_the_dataset(
+    conftest, collection, shards
+):
+    for nodeid, needs_data in collection.items():
+        shard = conftest.item_shard(_CollectedTest(nodeid, needs_data), shards)
+        if needs_data:
+            assert shard == shards
+        else:
+            assert shard == conftest.shard_of(nodeid, shards - 1) < shards
+
+
+@PROPERTY_SETTINGS
+@given(nodeid=nodeids, needs_data=st.booleans())
+def test_one_shard_runs_every_test(conftest, nodeid, needs_data):
+    assert conftest.item_shard(_CollectedTest(nodeid, needs_data), 1) == 1
+
+
+def test_the_workflow_gives_the_last_shard_one_worker():
+    import yaml
+
+    workflow = GITHUB_DIR / "workflows" / "test-suite.yaml"
+    job = yaml.safe_load(workflow.read_text(encoding="utf-8"))["jobs"]["Python"]
+    shards = job["strategy"]["matrix"]["shard"]
+    assert shards == list(range(1, len(shards) + 1))
+    [step] = [step for step in job["steps"] if step.get("name") == "Run the tests"]
+    assert step["env"]["WORKERS"] == (
+        "${{ matrix.shard == strategy.job-total && 1 || 4 }}"
+    )
+    assert '-n "$WORKERS"' in step["run"]
+
+
 # Invariant 4: the collection listing.
 
 
@@ -545,6 +598,15 @@ def test_never_runs():
     raise AssertionError
 """
 
+SAMPLE_DATA_TEST = """
+import pytest
+
+
+@pytest.mark.microsimulation
+def test_needs_the_dataset():
+    pass
+"""
+
 FAILING_SUITE = """
 import pytest
 
@@ -608,6 +670,11 @@ def test_passes():
 nodeid = "suite/test_skipped_module.py"
 outcome = "collect-skipped"
 why = "sample"
+
+[[allow]]
+nodeid = "suite/test_latest_data_smoke.py::test_needs_the_dataset"
+outcome = "skipped"
+why = "sample"
 """
 
 
@@ -728,19 +795,32 @@ def test_a_skipped_subtest_is_not_hidden_by_its_tests_pass(tmp_path):
 
 @pytest.mark.parametrize("xdist", [(), ("-n", "2", "--dist", "worksteal")])
 def test_real_pytest_shards_pass_the_check_only_when_allowed(tmp_path, xdist):
-    _suite(tmp_path, test_sample=SAMPLE_SUITE, test_skipped_module=SKIPPED_MODULE)
+    # conftest.py verifies no dataset for a test in this file (UNPINNED_BY_DESIGN),
+    # so the run downloads nothing.
+    _suite(
+        tmp_path,
+        test_sample=SAMPLE_SUITE,
+        test_skipped_module=SKIPPED_MODULE,
+        test_latest_data_smoke=SAMPLE_DATA_TEST,
+    )
     listing = _pytest(tmp_path, "suite", "--collect-only", "-q")
     assert listing.returncode == 0, listing.stdout + listing.stderr
     collected = tmp_path / "collected.txt"
     collected.write_text(listing.stdout, encoding="utf-8")
-    assert len(checker.read_collected(listing.stdout)) == 43
+    assert len(checker.read_collected(listing.stdout)) == 44
 
     logs = _run_shards(tmp_path, 3, *xdist)
     sizes = [
         sum(not r.is_collector for r in checker.read_outcomes(log.read_text()))
         for log in logs
     ]
-    assert sum(sizes) == 43 and all(sizes)
+    assert sum(sizes) == 44 and all(sizes)
+    # The last shard runs the test that needs the dataset and nothing else.
+    assert [
+        r.nodeid
+        for r in checker.read_outcomes(logs[-1].read_text())
+        if not r.is_collector
+    ] == ["suite/test_latest_data_smoke.py::test_needs_the_dataset"]
     allowlist = tmp_path / "allow.toml"
     allowlist.write_text(SAMPLE_ALLOWLIST, encoding="utf-8")
 

policyengine_uk/tests/conftest.py

@@ -175,7 +175,7 @@ def pytest_report_header(config):
     ]
     shard = config.getoption("shard", default=None)
     if shard is not None:
-        lines.append(f"shard: {shard[0]} of {shard[1]}, chosen by node id")
+        lines.append(f"shard: {shard[0]} of {shard[1]}, by marker and node id")
     return lines
 
 
@@ -193,12 +193,17 @@ 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 (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).
+# test is a function of its marker and 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
+# The last shard runs the tests that need the dataset, which build simulations
+# of all of it, and the other shards share the rest. CI gives the last shard
+# one worker, so only one of those simulations runs at a time
+# (.github/workflows/test-suite.yaml has the measurements).
+#
+# Most other 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
@@ -257,6 +262,17 @@ def shard_of(nodeid: str, shards: int) -> int:
     return int.from_bytes(digest[:8], "big") % shards + 1
 
 
+def item_shard(item, shards: int) -> int:
+    """The shard, from 1 to ``shards``, that runs a collected test.
+
+    The last of two or more shards runs the tests that need the dataset, and
+    ``shard_of`` places the others on the remaining shards.
+    """
+    if shards == 1:
+        return 1
+    return shards if needs_data(item) else shard_of(item.nodeid, shards - 1)
+
+
 def parse_shard(value: str) -> tuple[int, int]:
     """Parse ``K/N`` into ``(K, N)``, requiring 1 <= K <= N."""
     match = re.fullmatch(r"([0-9]+)/([0-9]+)", value)
@@ -286,7 +302,7 @@ def pytest_collection_modifyitems(config, items):
         this_shard, shards = shard
         selected, deselected = [], []
         for item in items:
-            if shard_of(item.nodeid, shards) == this_shard:
+            if item_shard(item, shards) == this_shard:
                 selected.append(item)
             else:
                 deselected.append(item)