From cffb3722a497b452da3e63336652a5f0fbfecaa0 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sun, 11 Oct 2026 12:27:08 -0400 Subject: [PATCH 1/4] Freeze the import-time heap and collect less often in the tests Each process that runs tests raises the first-generation GC threshold from 2,000 to 10,000 at start-up and calls gc.freeze() once the tests are collected, so collections stop walking the ~430,000 objects present after importing policyengine_uk. release_dead_variable_modules now also counts the frozen tax-benefit systems as live, since gc.get_objects() leaves frozen objects out; without that it would drop the import-time system's variable modules. Fixes #2316 Co-Authored-By: Claude Opus 5.5 --- changelog.d/2316.changed.md | 1 + policyengine_uk/tests/conftest.py | 49 +++++++++++++++++++ .../tests/test_variable_module_release.py | 34 +++++++++++++ 3 files changed, 84 insertions(+) create mode 100644 changelog.d/2316.changed.md 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..883c02fdad 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,53 @@ 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 (its default +# threshold), 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_THRESHOLD, +# set at start-up; the older generations keep Python's defaults); +# - freezes every object that exists once the tests are collected +# (gc.freeze()): the modules, the import-time tax-benefit systems and the +# collected tests, all of which live for the whole session. The collector +# never scans or frees frozen objects. Objects created by tests are collected +# as before. +# +# On test_ni_class_1_age_properties.py (52 examples, no coverage) this cut CPU +# time from 88-96 seconds to 61, and raised peak memory from 0.79 GB to 0.89 GB +# (measured locally on 10 October 2026; freezing alone gave 77 seconds, the +# threshold alone 67). It changes when memory is reclaimed, not which tests run +# or what they compute. +GC_THRESHOLD = (10_000, 10, 10) + +# gc.get_objects() leaves out frozen objects, so release_dead_variable_modules +# looks up the systems that were frozen here. +_FROZEN_SYSTEMS = weakref.WeakSet() +_heap_frozen = False + + +def pytest_configure(config): + gc.set_threshold(*GC_THRESHOLD) + + +@pytest.hookimpl(trylast=True) +def pytest_collection_finish(session): + global _heap_frozen + if _heap_frozen: + return + gc.collect() + _FROZEN_SYSTEMS.update( + obj for obj in gc.get_objects() if isinstance(obj, TaxBenefitSystem) + ) + gc.freeze() + _heap_frozen = True + + # 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 @@ -132,6 +180,7 @@ def release_dead_variable_modules(max_owners: int = 8) -> int: 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_variable_module_release.py b/policyengine_uk/tests/test_variable_module_release.py index 3d52036c73..9cf75de1fd 100644 --- a/policyengine_uk/tests/test_variable_module_release.py +++ b/policyengine_uk/tests/test_variable_module_release.py @@ -8,10 +8,16 @@ 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, and systems the tests build afterwards +must still be collected. """ import gc import sys +import weakref from pathlib import Path import pytest @@ -72,3 +78,31 @@ 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() == conftest.GC_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 From 057132c1ea6f9325416d00762493c905d37b48c6 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sun, 11 Oct 2026 14:35:33 -0400 Subject: [PATCH 2/4] Cite the CI measurement of the GC settings in conftest.py Co-Authored-By: Claude Opus 5.5 --- policyengine_uk/tests/conftest.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/policyengine_uk/tests/conftest.py b/policyengine_uk/tests/conftest.py index 883c02fdad..b9385a1ae8 100644 --- a/policyengine_uk/tests/conftest.py +++ b/policyengine_uk/tests/conftest.py @@ -120,11 +120,16 @@ def pytest_collection_modifyitems(config, items): # never scans or frees frozen objects. Objects created by tests are collected # as before. # -# On test_ni_class_1_age_properties.py (52 examples, no coverage) this cut CPU -# time from 88-96 seconds to 61, and raised peak memory from 0.79 GB to 0.89 GB -# (measured locally on 10 October 2026; freezing alone gave 77 seconds, the -# threshold alone 67). It changes when memory is reclaimed, not which tests run -# or what they compute. +# Measured on CI on 11 October 2026 (PR #2317: runs 38155711978 and +# 38158313028, each running the three shards under several settings at once +# from one commit): the workers' CPU time, summed over the 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. Peak runner memory did not +# change (15.3 GiB of 15.6). In the first run, freezing alone saved nothing, +# the threshold alone saved 13 minutes, and with a threshold of 25,000 two of +# the three shards stalled. Every test had the same outcome under each +# setting: this changes when memory is reclaimed, not which tests run or what +# they compute. GC_THRESHOLD = (10_000, 10, 10) # gc.get_objects() leaves out frozen objects, so release_dead_variable_modules From 50ebd887793dcd0fab4b99b6c622640d85ee7701 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sun, 11 Oct 2026 15:35:05 -0400 Subject: [PATCH 3/4] Keep systems alive past the freeze, guard it, and restore GC state at exit Review findings on #2317: - Two test modules built a Simulation at import (council tax uprating's region-to-country map, the childcare test's minimum care hours). Their systems died during collection, so their variable modules were frozen, then released and never freed (about 19,000 objects per worker). Read the hours from the package's parameters and build the map on first use. - conftest.py records any variable-module owner whose system died before the freeze, and test_no_system_dies_before_the_freeze fails on one. - Set only the first-generation threshold and keep the interpreter's other thresholds (CPython 3.14.4 defaults to (2000, 10, 0)). - Undo the freeze and threshold at unconfigure, for in-process runs. - Correct the comment: about 1.6 million objects are frozen per worker, and on the shard whose runner memory isn't full the largest worker's peak rose from 5.3 to 5.7 GiB. Co-Authored-By: Claude Opus 5.5 --- policyengine_uk/tests/conftest.py | 93 +++++++++++++------ .../tests/test_council_tax_uprating.py | 10 +- ..._uc_childcare_work_condition_properties.py | 11 +-- .../tests/test_variable_module_release.py | 18 +++- 4 files changed, 88 insertions(+), 44 deletions(-) diff --git a/policyengine_uk/tests/conftest.py b/policyengine_uk/tests/conftest.py index b9385a1ae8..1ec459c814 100644 --- a/policyengine_uk/tests/conftest.py +++ b/policyengine_uk/tests/conftest.py @@ -107,39 +107,51 @@ def pytest_collection_modifyitems(config, items): # *_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 (its default -# threshold), 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: +# 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_THRESHOLD, -# set at start-up; the older generations keep Python's defaults); +# - 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()): the modules, the import-time tax-benefit systems and the -# collected tests, all of which live for the whole session. The collector -# never scans or frees frozen objects. Objects created by tests are collected -# as before. +# (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 a frozen object that later becomes garbage +# is never freed. No tax-benefit system may therefore die before the freeze +# (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 (PR #2317: runs 38155711978 and -# 38158313028, each running the three shards under several settings at once -# from one commit): the workers' CPU time, summed over the 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. Peak runner memory did not -# change (15.3 GiB of 15.6). In the first run, freezing alone saved nothing, -# the threshold alone saved 13 minutes, and with a threshold of 25,000 two of -# the three shards stalled. Every test had the same outcome under each -# setting: this changes when memory is reclaimed, not which tests run or what -# they compute. -GC_THRESHOLD = (10_000, 10, 10) +# 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): - gc.set_threshold(*GC_THRESHOLD) + 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) @@ -148,13 +160,30 @@ def pytest_collection_finish(session): if _heap_frozen: return gc.collect() - _FROZEN_SYSTEMS.update( - obj for obj in gc.get_objects() if isinstance(obj, TaxBenefitSystem) + 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): + # Leave a process that runs pytest in-process (an IDE, a REPL) as it was. + 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 @@ -167,6 +196,16 @@ def pytest_collection_finish(session): _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. @@ -174,11 +213,7 @@ 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() 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 9cf75de1fd..a0300c4538 100644 --- a/policyengine_uk/tests/test_variable_module_release.py +++ b/policyengine_uk/tests/test_variable_module_release.py @@ -11,8 +11,9 @@ 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, and systems the tests build afterwards -must still be collected. +still count a frozen system as live. A frozen object that later becomes +garbage is never freed, so no system may die before the freeze: its variable +modules would be frozen, then released and leaked. """ import gc @@ -85,7 +86,7 @@ def test_frozen_systems_keep_their_modules(request): # Built when policyengine_uk was imported, during collection, so frozen. from policyengine_uk.tax_benefit_system import system - assert gc.get_threshold() == conftest.GC_THRESHOLD + 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__ @@ -106,3 +107,14 @@ def test_systems_built_after_the_freeze_are_still_collected(): 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." + ) From 37cadb47192adb0110a796cddf5ef2cb2ce6d795 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Mon, 12 Oct 2026 06:53:46 -0400 Subject: [PATCH 4/4] Say which frozen objects the freeze keeps, and what unconfigure undoes From the GPT-6.1 Sol review of a8ff60fa0: frozen objects outside reference cycles are still freed by reference counting, and gc.unfreeze() also unfreezes anything frozen before the session. Comments only. Co-Authored-By: Claude Opus 5.5 --- policyengine_uk/tests/conftest.py | 15 ++++++++++----- .../tests/test_variable_module_release.py | 7 ++++--- 2 files changed, 14 insertions(+), 8 deletions(-) diff --git a/policyengine_uk/tests/conftest.py b/policyengine_uk/tests/conftest.py index 1ec459c814..5f97662dbc 100644 --- a/policyengine_uk/tests/conftest.py +++ b/policyengine_uk/tests/conftest.py @@ -119,10 +119,13 @@ def pytest_collection_modifyitems(config, items): # - 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 a frozen object that later becomes garbage -# is never freed. No tax-benefit system may therefore die before the freeze -# (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. +# 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 @@ -172,7 +175,9 @@ def pytest_collection_finish(session): def pytest_unconfigure(config): - # Leave a process that runs pytest in-process (an IDE, a REPL) as it was. + # 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() diff --git a/policyengine_uk/tests/test_variable_module_release.py b/policyengine_uk/tests/test_variable_module_release.py index a0300c4538..241798a645 100644 --- a/policyengine_uk/tests/test_variable_module_release.py +++ b/policyengine_uk/tests/test_variable_module_release.py @@ -11,9 +11,10 @@ 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. A frozen object that later becomes -garbage is never freed, so no system may die before the freeze: its variable -modules would be frozen, then released and leaked. +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