diff --git a/.github/workflows/code_changes.yaml b/.github/workflows/code_changes.yaml index 6a0008f2ee..15d8b0a87b 100644 --- a/.github/workflows/code_changes.yaml +++ b/.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 diff --git a/.github/workflows/pr_code_changes.yaml b/.github/workflows/pr_code_changes.yaml index 9b7b313430..f039cfc1dd 100644 --- a/.github/workflows/pr_code_changes.yaml +++ b/.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 diff --git a/.github/workflows/pr_docs_changes.yaml b/.github/workflows/pr_docs_changes.yaml index 5add446350..ce7fceb931 100644 --- a/.github/workflows/pr_docs_changes.yaml +++ b/.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 diff --git a/Makefile b/Makefile index 764c102331..5723221068 100644 --- a/Makefile +++ b/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 diff --git a/changelog.d/ci-linux-parallel-tests.changed.md b/changelog.d/ci-linux-parallel-tests.changed.md new file mode 100644 index 0000000000..e6c1f87c91 --- /dev/null +++ b/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. diff --git a/policyengine_uk/tests/conftest.py b/policyengine_uk/tests/conftest.py index 56ee3b207b..deee0cc756 100644 --- a/policyengine_uk/tests/conftest.py +++ b/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 "__" 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() diff --git a/policyengine_uk/tests/test_latest_data_smoke.py b/policyengine_uk/tests/test_latest_data_smoke.py index e80e96ec2e..aa9639cab9 100644 --- a/policyengine_uk/tests/test_latest_data_smoke.py +++ b/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: diff --git a/policyengine_uk/tests/test_variable_module_release.py b/policyengine_uk/tests/test_variable_module_release.py new file mode 100644 index 0000000000..61d2047ccd --- /dev/null +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index 6453b944b6..e670f561cf 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -61,6 +61,7 @@ dev = [ "coverage", "hypothesis", "pytest-cov", + "pytest-xdist", "ruff>=0.9.0", "towncrier>=24.8.0", "furo<2023", diff --git a/uv.lock b/uv.lock index 47b04e0136..3200859772 100644 --- a/uv.lock +++ b/uv.lock @@ -69,7 +69,7 @@ resolution-markers = [ "python_full_version >= '3.14'", ] dependencies = [ - { name = "cffi", marker = "python_full_version >= '3.14'" }, + { name = "cffi" }, ] sdist = { url = "https://files.pythonhosted.org/packages/b9/e9/184b8ccce6683b0aa2fbb7ba5683ea4b9c5763f1356347f1312c32e3c66e/argon2-cffi-bindings-21.2.0.tar.gz", hash = "sha256:bb89ceffa6c791807d1305ceb77dbfacc5aa499891d2c55661c6459651fc39e3", size = 1779911, upload-time = "2021-12-01T08:52:55.68Z" } wheels = [ @@ -94,7 +94,7 @@ resolution-markers = [ "python_full_version < '3.12'", ] dependencies = [ - { name = "cffi", marker = "python_full_version < '3.14'" }, + { name = "cffi" }, ] sdist = { url = "https://files.pythonhosted.org/packages/5c/2d/db8af0df73c1cf454f71b2bbe5e356b8c1f8041c979f505b3d3186e520a9/argon2_cffi_bindings-25.1.0.tar.gz", hash = "sha256:b957f3e6ea4d55d820e40ff76f450952807013d361a65d7f28acc0acbf29229d", size = 1783441, upload-time = "2025-07-30T10:02:05.147Z" } wheels = [ @@ -483,6 +483,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/05/d1/8952806fbf9583004ab479d8f58a9496c3d35f6b6009ddd458bdd9978eaf/dpath-2.2.0-py3-none-any.whl", hash = "sha256:b330a375ded0a0d2ed404440f6c6a715deae5313af40bbb01c8a41d891900576", size = 17618, upload-time = "2024-06-12T22:08:01.881Z" }, ] +[[package]] +name = "execnet" +version = "2.1.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/bf/89/780e11f9588d9e7128a3f87788354c7946a9cbb1401ad38a48c4db9a4f07/execnet-2.1.2.tar.gz", hash = "sha256:63d83bfdd9a23e35b9c6a3261412324f964c2ec8dcd8d3c6916ee9373e0befcd", size = 166622, upload-time = "2025-11-12T09:56:37.75Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ab/84/02fc1827e8cdded4aa65baef11296a9bbe595c474f0d6d758af082d849fd/execnet-2.1.2-py3-none-any.whl", hash = "sha256:67fba928dd5a544b783f6056f449e5e3931a5c378b128bc18501f7ea79e296ec", size = 40708, upload-time = "2025-11-12T09:56:36.333Z" }, +] + [[package]] name = "executing" version = "2.2.0" @@ -1471,7 +1480,7 @@ wheels = [ [[package]] name = "policyengine-uk" -version = "2.102.3" +version = "2.107.0" source = { editable = "." } dependencies = [ { name = "microdf-python" }, @@ -1487,6 +1496,7 @@ dev = [ { name = "hypothesis" }, { name = "jupyter-book" }, { name = "pytest-cov" }, + { name = "pytest-xdist" }, { name = "rich" }, { name = "ruff" }, { name = "setuptools" }, @@ -1507,6 +1517,7 @@ requires-dist = [ { name = "policyengine-core", specifier = ">=3.32.9" }, { name = "pydantic", specifier = ">=2.11.7" }, { name = "pytest-cov", marker = "extra == 'dev'" }, + { name = "pytest-xdist", marker = "extra == 'dev'" }, { name = "rich", marker = "extra == 'dev'" }, { name = "ruff", marker = "extra == 'dev'", specifier = ">=0.9.0" }, { name = "setuptools", marker = "extra == 'dev'" }, @@ -1710,6 +1721,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/bc/16/4ea354101abb1287856baa4af2732be351c7bee728065aed451b678153fd/pytest_cov-6.2.1-py3-none-any.whl", hash = "sha256:f5bc4c23f42f1cdd23c70b1dab1bbaef4fc505ba950d53e0081d0730dd7e86d5", size = 24644, upload-time = "2025-06-12T10:47:45.932Z" }, ] +[[package]] +name = "pytest-xdist" +version = "3.8.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "execnet" }, + { name = "pytest" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/78/b4/439b179d1ff526791eb921115fca8e44e596a13efeda518b9d845a619450/pytest_xdist-3.8.0.tar.gz", hash = "sha256:7e578125ec9bc6050861aa93f2d59f1d8d085595d6551c2c90b6f4fad8d3a9f1", size = 88069, upload-time = "2025-07-01T13:30:59.346Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ca/31/d4e37e9e550c2b92a9cbc2e4d0b7420a27224968580b5a447f420847c975/pytest_xdist-3.8.0-py3-none-any.whl", hash = "sha256:202ca578cfeb7370784a8c33d6d05bc6e13b4f25b5053c30a152269fd10f0b88", size = 46396, upload-time = "2025-07-01T13:30:56.632Z" }, +] + [[package]] name = "python-dateutil" version = "2.9.0.post0"