diff --git a/changelog.d/2316.changed.md b/changelog.d/2316.changed.md new file mode 100644 index 0000000000..ec4681ab4a --- /dev/null +++ b/changelog.d/2316.changed.md @@ -0,0 +1 @@ +- Cut the CPU the Python tests spend in garbage collection: each test process raises the first-generation GC threshold and freezes the objects that exist once the tests are collected (#2316). diff --git a/policyengine_uk/tests/conftest.py b/policyengine_uk/tests/conftest.py index ca783dd9ab..5f97662dbc 100644 --- a/policyengine_uk/tests/conftest.py +++ b/policyengine_uk/tests/conftest.py @@ -4,6 +4,7 @@ import os import re import sys +import weakref import pytest from policyengine_core.taxbenefitsystems import TaxBenefitSystem @@ -102,6 +103,92 @@ def pytest_collection_modifyitems(config, items): item.add_marker(skip_microsimulation) +# Garbage collection. Most of the suite's time is in the property tests (the 51 +# *_properties.py files took 74% of pytest worker time on 10 October 2026), and +# each of their examples builds a Simulation, which loads about 1,100 variable +# modules again and clones the parameter tree. Python's cycle collector starts +# a collection every 2,000 net allocations of tracked objects (the default in +# CPython 3.13, which CI uses), and its full collections walk every tracked +# object in the process, including the 430,000 or so present once +# policyengine_uk is imported. Each process that runs tests (every +# pytest-xdist worker) therefore: +# +# - starts a collection every 10,000 net allocations instead +# (GC_FIRST_GENERATION_THRESHOLD, set at start-up; the other thresholds keep +# the interpreter's defaults); +# - freezes every object that exists once the tests are collected +# (gc.freeze(), about 1.6 million objects per worker): the modules, the +# tax-benefit systems built at import and the collected tests. The collector +# never scans frozen objects, so frozen objects in reference cycles that +# later become garbage are not freed until they are unfrozen (objects outside +# cycles still go with their last reference). A tax-benefit system and its +# variable modules are such cycles, so no system may be built and dropped +# while the tests are collected (test_variable_module_release.py checks +# this); build one in a fixture or on first use, not at import. Objects +# created by tests are collected as before. +# +# Measured on CI on 11 October 2026 (runs and numbers in PR #2317): the +# workers' CPU time, summed over the three shards, fell from 204 to 172 +# minutes, and their time in the collector from 83 to 55; the slowest shard +# took 20 minutes instead of 24. The runner's memory was already full on two +# shards under every setting; on the third, the largest worker's peak rose from +# 5.3 to 5.7 GiB. Freezing alone saved nothing, the threshold alone saved 13 +# minutes of worker CPU, and with a threshold of 25,000 two of the three +# runners shut down. Every test had the same outcome under each setting: this +# changes when memory is reclaimed, not which tests run or what they compute. +GC_FIRST_GENERATION_THRESHOLD = 10_000 + +# gc.get_objects() leaves out frozen objects, so release_dead_variable_modules +# looks up the systems that were frozen here. +_FROZEN_SYSTEMS = weakref.WeakSet() +# Owners of variable modules whose system was built and had died by the +# freeze, and the owners already loaded when the session started (a process +# that runs pytest more than once keeps the previous session's). +_DEAD_AT_FREEZE = set() +_owners_at_start = set() +_heap_frozen = False +_threshold_before = None + + +def pytest_configure(config): + global _threshold_before + _owners_at_start.update(_variable_module_owners()) + _threshold_before = gc.get_threshold() + gc.set_threshold(GC_FIRST_GENERATION_THRESHOLD, *_threshold_before[1:]) + + +@pytest.hookimpl(trylast=True) +def pytest_collection_finish(session): + global _heap_frozen + if _heap_frozen: + return + gc.collect() + systems = [obj for obj in gc.get_objects() if isinstance(obj, TaxBenefitSystem)] + _FROZEN_SYSTEMS.update(systems) + _DEAD_AT_FREEZE.update( + set(_variable_module_owners()) + - {str(id(system)) for system in systems} + - _owners_at_start + ) + gc.freeze() + _heap_frozen = True + + +def pytest_unconfigure(config): + # Undo the freeze and threshold for a process that runs pytest in-process + # (an IDE, a REPL). gc.unfreeze() also unfreezes anything frozen before + # the session. + global _heap_frozen + if _heap_frozen: + gc.unfreeze() + _FROZEN_SYSTEMS.clear() + _heap_frozen = False + _DEAD_AT_FREEZE.clear() + _owners_at_start.clear() + if _threshold_before is not None: + gc.set_threshold(*_threshold_before) + + # 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 @@ -114,6 +201,16 @@ def pytest_collection_modifyitems(config, items): _VARIABLE_MODULE_NAME = re.compile(r"(\d+)_-?\d+_") +def _variable_module_owners() -> dict: + """Map each owner id in loaded variable module names to those names.""" + owners = {} + for name in list(sys.modules): + match = _VARIABLE_MODULE_NAME.match(name) + if match: + owners.setdefault(match.group(1), []).append(name) + return owners + + def release_dead_variable_modules(max_owners: int = 8) -> int: """Remove dead tax-benefit systems' variable modules from sys.modules. @@ -121,17 +218,14 @@ def release_dead_variable_modules(max_owners: int = 8) -> int: 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) + owners = _variable_module_owners() if len(owners) <= max_owners: return 0 gc.collect() live = { str(id(obj)) for obj in gc.get_objects() if isinstance(obj, TaxBenefitSystem) } + live.update(str(id(system)) for system in _FROZEN_SYSTEMS) released = 0 for owner, names in owners.items(): if owner not in live: diff --git a/policyengine_uk/tests/test_council_tax_uprating.py b/policyengine_uk/tests/test_council_tax_uprating.py index 318a124427..92ec4d0922 100644 --- a/policyengine_uk/tests/test_council_tax_uprating.py +++ b/policyengine_uk/tests/test_council_tax_uprating.py @@ -1,3 +1,4 @@ +import functools from types import SimpleNamespace import numpy as np @@ -20,10 +21,12 @@ COUNCIL_TAX_GROWTH = parameters.gov.economic_assumptions.yoy_growth.obr.council_tax +@functools.cache def _country_of_each_region(): # The model's own region-to-country mapping (the `country` variable), so # the expected growth below comes from a second path to the one the - # uprating code takes. + # uprating code takes. Built on first use, not at import: a system built + # and dropped while the tests are collected would be frozen by conftest.py. people = {f"person_{region}": {"age": 40} for region in REGIONS} households = { f"household_{region}": {"members": [f"person_{region}"], "region": region} @@ -34,14 +37,11 @@ def _country_of_each_region(): return dict(zip(REGIONS, countries)) -COUNTRY = _country_of_each_region() - - def _expected_growth(region, year): # Council tax exists in England, Scotland and Wales; OBR forecasts it for # each. Northern Ireland levies domestic rates instead, so its council tax # has no growth. An unknown region takes England's. - country = COUNTRY[region] + country = _country_of_each_region()[region] if country == "NORTHERN_IRELAND": return 0.0 series = "england" if country == "UNKNOWN" else country.lower() diff --git a/policyengine_uk/tests/test_uc_childcare_work_condition_properties.py b/policyengine_uk/tests/test_uc_childcare_work_condition_properties.py index 7b9547ff31..a1dfb4c39a 100644 --- a/policyengine_uk/tests/test_uc_childcare_work_condition_properties.py +++ b/policyengine_uk/tests/test_uc_childcare_work_condition_properties.py @@ -38,7 +38,7 @@ from hypothesis import HealthCheck, given, settings from hypothesis import strategies as st -from policyengine_uk import Simulation +from policyengine_uk import Simulation, parameters YEAR = 2025 PROPERTY_SETTINGS = settings( @@ -144,12 +144,9 @@ def calculate(units, with_dependants=True): } -def min_care_hours(): - sim = Simulation(situation={"people": {"p": {"age": {YEAR: 30}}}}) - return sim.tax_benefit_system.parameters(YEAR).gov.dwp.carers_allowance.min_hours - - -MIN_CARE_HOURS = min_care_hours() +# Read from the package's parameters, not a Simulation's: a system built and +# dropped while the tests are collected would be frozen by conftest.py. +MIN_CARE_HOURS = parameters(YEAR).gov.dwp.carers_allowance.min_hours def meets_claimant_limb(member): diff --git a/policyengine_uk/tests/test_variable_module_release.py b/policyengine_uk/tests/test_variable_module_release.py index 3d52036c73..241798a645 100644 --- a/policyengine_uk/tests/test_variable_module_release.py +++ b/policyengine_uk/tests/test_variable_module_release.py @@ -8,10 +8,18 @@ note to remove the workaround, rather than failing every PR after a core upgrade. If core keeps leaking but names the modules differently, it fails, because the release would then silently do nothing. + +The conftest also freezes the objects that exist once the tests are collected +(gc.freeze()). gc.get_objects() leaves frozen objects out, so the release must +still count a frozen system as live. Frozen objects in reference cycles that +later become garbage are not freed until they are unfrozen, so no system may be +built and dropped while the tests are collected: its variable modules would be +frozen, then released and leaked for the rest of the session. """ import gc import sys +import weakref from pathlib import Path import pytest @@ -72,3 +80,42 @@ def test_dead_systems_modules_are_released_and_live_ones_kept(request): 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 + + +def test_frozen_systems_keep_their_modules(request): + conftest = _conftest(request) + # Built when policyengine_uk was imported, during collection, so frozen. + from policyengine_uk.tax_benefit_system import system + + assert gc.get_threshold()[0] == conftest.GC_FIRST_GENERATION_THRESHOLD + assert gc.get_freeze_count() > 0 + assert not any(obj is system for obj in gc.get_objects()) + module = type(system.variables["income_tax"]).__module__ + built_later = Simulation(situation=SITUATION).tax_benefit_system + if type(built_later.variables["income_tax"]).__module__ not in sys.modules: + pytest.skip("policyengine-core no longer keeps variable modules") + + conftest.release_dead_variable_modules(max_owners=0) + + kept = module in sys.modules + assert kept, f"{module}, of a live frozen system, was released" + + +def test_systems_built_after_the_freeze_are_still_collected(): + simulation = Simulation(situation=SITUATION) + system = weakref.ref(simulation.tax_benefit_system) + assert any(obj is system() for obj in gc.get_objects()) + del simulation + gc.collect() + assert system() is None + + +def test_no_system_dies_before_the_freeze(request): + dead = _conftest(request)._DEAD_AT_FREEZE + assert not dead, ( + f"{len(dead)} tax-benefit system(s) were built and dropped while the " + "tests were collected, so conftest.py froze their variable modules " + "and they will never be freed. A test module probably builds a " + "Simulation or CountryTaxBenefitSystem at import: build it in a " + "fixture or on first use instead." + )