Cut PR CI from hours to about 20 minutes: Linux arm64, 4 xdist workers, cancel superseded runs

.github/workflows/code_changes.yaml

@@ -22,7 +22,16 @@ jobs:
         - name: Check formatting
           run: uvx ruff format --check .
   Test:
-      runs-on: macos-latest
+      # 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 takes about 18 minutes (YAML tests 2, pytest 15, the
+      # rest under 1). More than three times that means a run is hung or
+      # the suite needs attention, so fail it rather than wait for the
+      # 6-hour GitHub limit.
+      timeout-minutes: 60
       permissions:
         contents: "read"
         id-token: "write"
@@ -55,7 +64,9 @@ jobs:
               if [ -n "$HF_TOKEN_VALUE" ]; then
                 export HUGGING_FACE_TOKEN="$HF_TOKEN_VALUE"
               fi
-              make test
+              # 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"
             env:
               HF_TOKEN_VALUE: ${{ secrets.HUGGING_FACE_TOKEN }}
           - name: Upload coverage to Codecov

.github/workflows/pr_code_changes.yaml

@@ -14,6 +14,13 @@ on:
       - uv.lock
   workflow_dispatch:
 
+# A new push to a pull request supersedes the run on its previous head, so
+# cancel that run instead of letting it hold runners. Manual dispatches are
+# grouped by ref and never cancelled.
+concurrency:
+  group: pr-code-changes-${{ github.event.pull_request.number || github.ref }}
+  cancel-in-progress: ${{ github.event_name == 'pull_request' }}
+
 jobs:
   Python-Compat:
     name: Install + smoke-import (py${{ matrix.python-version }})
@@ -43,7 +50,16 @@ jobs:
         - name: Check formatting
           run: uvx ruff format --check .
   Test:
-      runs-on: macos-latest
+      # 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 takes about 18 minutes (YAML tests 2, pytest 15, the
+      # rest under 1). More than three times that means a run is hung or
+      # the suite needs attention, so fail it rather than wait for the
+      # 6-hour GitHub limit.
+      timeout-minutes: 60
       permissions:
         contents: "read"
         id-token: "write"
@@ -74,7 +90,9 @@ jobs:
               if [ -n "$HF_TOKEN_VALUE" ]; then
                 export HUGGING_FACE_TOKEN="$HF_TOKEN_VALUE"
               fi
-              make test
+              # 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"
             env:
               HF_TOKEN_VALUE: ${{ secrets.HUGGING_FACE_TOKEN }}
           - name: Upload coverage to Codecov

.github/workflows/pr_docs_changes.yaml

@@ -11,6 +11,13 @@ on:
       - .github/**
   workflow_dispatch:
 
+# A new push to a pull request supersedes the run on its previous head, so
+# cancel that run instead of letting it hold runners. Manual dispatches are
+# grouped by ref and never cancelled.
+concurrency:
+  group: pr-docs-changes-${{ github.event.pull_request.number || github.ref }}
+  cancel-in-progress: ${{ github.event_name == 'pull_request' }}
+
 jobs:
   Test:
       runs-on: ubuntu-latest

Makefile

@@ -1,4 +1,4 @@
-.PHONY: docs
+.PHONY: docs test test-yaml test-python
 
 all: install
 	pip install build
@@ -14,9 +14,15 @@ format:
 	ruff format .
 	ruff check .
 
-test:
+test: test-yaml test-python
+
+test-yaml:
 	policyengine-core test policyengine_uk/tests/policy -c policyengine_uk
-	pytest policyengine_uk/tests/ --cov=policyengine_uk --cov-report=xml --maxfail=0 -v
+
+# CI passes PYTEST_ARGS="-n 4 --dist worksteal" (pytest-xdist) to spread the
+# tests over the runner's four cores. Without it this runs serially, as before.
+test-python:
+	pytest policyengine_uk/tests/ --cov=policyengine_uk --cov-report=xml --maxfail=0 -v $(PYTEST_ARGS)
 
 test-all:
 	policyengine-core test policyengine_uk/tests/policy -c policyengine_uk

changelog.d/ci-linux-parallel-tests.changed.md

@@ -0,0 +1 @@
+Run the CI Test job on Linux arm64 with four pytest-xdist workers (pytest-xdist is now a dev dependency), cancel superseded pull request runs, and time the job out after 60 minutes.

policyengine_uk/tests/conftest.py

@@ -1,6 +1,10 @@
+import gc
 import os
+import re
+import sys
 
 import pytest
+from policyengine_core.taxbenefitsystems import TaxBenefitSystem
 
 DEFAULT_TEST_DATASET_URL = (
     "hf://policyengine/policyengine-uk-data-private/enhanced_frs_2023_24.h5@1.40.3"
@@ -26,3 +30,48 @@ def pytest_collection_modifyitems(config, items):
     for item in items:
         if "microsimulation" in item.keywords:
             item.add_marker(skip_microsimulation)
+
+
+# policyengine-core loads every variable file of each new tax-benefit system as
+# a module named "<id(system)>_<path hash>_<file name>" and registers it in
+# sys.modules, where it stays after the system is garbage collected. Every
+# Simulation builds its own system, so each one leaves about 1,000 modules
+# (roughly 10 MB) behind, and the property tests build hundreds: one test
+# process grew to 15.7 GB on a 16 GB CI runner. After each test, drop the
+# modules of systems that no longer exist; live systems keep theirs. Remove
+# this once policyengine-core stops leaking them
+# (PolicyEngine/policyengine-core#520).
+_VARIABLE_MODULE_NAME = re.compile(r"(\d+)_-?\d+_")
+
+
+def release_dead_variable_modules(max_owners: int = 8) -> int:
+    """Remove dead tax-benefit systems' variable modules from sys.modules.
+
+    Does nothing while at most ``max_owners`` systems own loaded variable
+    modules, since finding the live systems scans every tracked object.
+    Returns the number of modules removed.
+    """
+    owners = {}
+    for name in list(sys.modules):
+        match = _VARIABLE_MODULE_NAME.match(name)
+        if match:
+            owners.setdefault(match.group(1), []).append(name)
+    if len(owners) <= max_owners:
+        return 0
+    gc.collect()
+    live = {
+        str(id(obj)) for obj in gc.get_objects() if isinstance(obj, TaxBenefitSystem)
+    }
+    released = 0
+    for owner, names in owners.items():
+        if owner not in live:
+            for name in names:
+                sys.modules.pop(name, None)
+            released += len(names)
+    return released
+
+
+@pytest.fixture(autouse=True)
+def _release_variable_modules_of_dead_systems():
+    yield
+    release_dead_variable_modules()

policyengine_uk/tests/test_latest_data_smoke.py

@@ -16,8 +16,6 @@
 
 from __future__ import annotations
 
-import os
-
 import numpy as np
 import pytest
 
@@ -35,10 +33,13 @@ def sim() -> Microsimulation:
     """Simulation built against the unpinned latest dataset.
 
     Overrides any pinned-version dataset set in conftest.py so the test
-    exercises whatever is on HuggingFace ``main`` right now.
+    exercises whatever is on HuggingFace ``main`` right now. The override
+    lasts only while the simulation is built, so tests that run later in the
+    same process still get the pinned dataset whatever order they run in.
     """
-    os.environ["POLICYENGINE_UK_DEFAULT_DATASET"] = LATEST_DATASET_URL
-    return Microsimulation()
+    with pytest.MonkeyPatch.context() as monkeypatch:
+        monkeypatch.setenv("POLICYENGINE_UK_DEFAULT_DATASET", LATEST_DATASET_URL)
+        return Microsimulation()
 
 
 def _weighted(sim: Microsimulation, variable: str, period: int = YEAR) -> float:

policyengine_uk/tests/test_variable_module_release.py

@@ -0,0 +1,62 @@
+"""The conftest releases the variable modules of dead tax-benefit systems.
+
+policyengine-core registers each system's variable modules in sys.modules
+under the system's id and never removes them. The conftest's
+``release_dead_variable_modules`` drops those of systems that no longer exist,
+which keeps a long test process's memory flat. If core stops leaking the
+modules (PolicyEngine/policyengine-core#520), the precondition below fails
+and the workaround can go.
+"""
+
+import gc
+import sys
+from pathlib import Path
+
+from policyengine_uk import Simulation
+
+SITUATION = {
+    "people": {"adult": {"age": {2026: 40}, "employment_income": {2026: 30_000}}},
+    "benunits": {"benunit": {"members": ["adult"]}},
+    "households": {"household": {"members": ["adult"]}},
+}
+
+
+def _conftest(request):
+    path = Path(__file__).with_name("conftest.py")
+    for plugin in request.config.pluginmanager.get_plugins():
+        if getattr(plugin, "__file__", None) and Path(plugin.__file__) == path:
+            return plugin
+    raise AssertionError(f"{path} is not loaded as a plugin")
+
+
+def _modules_of(owner: str) -> list:
+    return [name for name in sys.modules if name.startswith(owner + "_")]
+
+
+def test_dead_systems_modules_are_released_and_live_ones_kept(request):
+    release_dead_variable_modules = _conftest(request).release_dead_variable_modules
+    kept = Simulation(situation=SITUATION)
+    kept_owner = str(id(kept.tax_benefit_system))
+    dropped = Simulation(situation=SITUATION)
+    dropped_owner = str(id(dropped.tax_benefit_system))
+    expected = dropped.calculate("income_tax", 2026)
+
+    assert _modules_of(dropped_owner), (
+        "policyengine-core no longer registers variable modules under the "
+        "system's id; the conftest workaround may no longer be needed"
+    )
+    del dropped
+    gc.collect()
+
+    released = release_dead_variable_modules(max_owners=0)
+
+    assert released >= 1
+    assert not _modules_of(dropped_owner)
+    assert _modules_of(kept_owner)
+    assert kept.calculate("income_tax", 2026) == expected
+    assert Simulation(situation=SITUATION).calculate("income_tax", 2026) == expected
+
+
+def test_release_waits_until_enough_systems_own_modules(request):
+    release_dead_variable_modules = _conftest(request).release_dead_variable_modules
+    assert release_dead_variable_modules(max_owners=10**9) == 0

pyproject.toml

@@ -61,6 +61,7 @@ dev = [
     "coverage",
     "hypothesis",
     "pytest-cov",
+    "pytest-xdist",
     "ruff>=0.9.0",
     "towncrier>=24.8.0",
     "furo<2023",