Cut the CPU the Python tests spend in garbage collection

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).

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 "<id(system)>_<path hash>_<file name>" and registers it in
 # sys.modules, where it stays after the system is garbage collected. Every
@@ -114,24 +201,31 @@ 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.
 
     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)
+    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:

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()

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):

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."
+    )