From a92236561f50c6279782fa95626d2a607ee95e26 Mon Sep 17 00:00:00 2001
From: Max Ghenis <max@maxghenis.com>
Date: Fri, 30 Oct 2026 17:26:03 -0400
Subject: [PATCH 1/4] Pause automatic collection while the YAML runner builds a
 reform system

A cache miss in _get_tax_benefit_system clones the whole baseline. On
policyengine-us the full collections that its allocations trigger walk
every object both systems hold and free almost nothing. Pause automatic
collection across the clone, reforms and extensions, and restore the
caller's setting in a finally. A cache hit does not touch collection.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 ...se-gc-during-yaml-system-builds.changed.md |   1 +
 policyengine_core/tools/test_runner.py        |  47 ++-
 tests/core/tools/test_runner/test_gc_pause.py | 314 ++++++++++++++++++
 3 files changed, 350 insertions(+), 12 deletions(-)
 create mode 100644 changelog.d/pause-gc-during-yaml-system-builds.changed.md
 create mode 100644 tests/core/tools/test_runner/test_gc_pause.py

diff --git a/changelog.d/pause-gc-during-yaml-system-builds.changed.md b/changelog.d/pause-gc-during-yaml-system-builds.changed.md
new file mode 100644
index 000000000..c646bc023
--- /dev/null
+++ b/changelog.d/pause-gc-during-yaml-system-builds.changed.md
@@ -0,0 +1 @@
+Pause automatic garbage collection while the YAML test runner builds the system for a reform combination, and restore the caller's setting afterwards: full collections that freed almost nothing had taken about half of each build on policyengine-us.
diff --git a/policyengine_core/tools/test_runner.py b/policyengine_core/tools/test_runner.py
index 6b8d4aefc..4b6d92fb4 100644
--- a/policyengine_core/tools/test_runner.py
+++ b/policyengine_core/tools/test_runner.py
@@ -1,10 +1,12 @@
 # -*- coding: utf-8 -*-
 
+import gc
 import logging
 import sys
 import os
 import traceback
 import textwrap
+from contextlib import contextmanager
 from typing import Dict, List
 
 import pytest
@@ -419,6 +421,25 @@ def pytest_collect_file(self, parent, file_path):
             )
 
 
+@contextmanager
+def _automatic_gc_paused():
+    """Pause automatic garbage collection, then restore the caller's setting.
+
+    Building a reform system copies the whole baseline. Its allocations
+    trigger full collections that walk every object the baseline and its
+    copies hold and free almost nothing: on policyengine-us they took about
+    half of each build. Collection resumes as before once the build returns
+    or raises, and a caller that had disabled it keeps it disabled.
+    """
+    was_enabled = gc.isenabled()
+    gc.disable()
+    try:
+        yield
+    finally:
+        if was_enabled:
+            gc.enable()
+
+
 def _get_tax_benefit_system(
     baseline,
     reforms,
@@ -444,20 +465,22 @@ def _get_tax_benefit_system(
         if cached is not None:
             return cached
 
-    current_tax_benefit_system = baseline.clone()
+    # Only a cache miss builds a system, so only a miss pauses collection.
+    with _automatic_gc_paused():
+        current_tax_benefit_system = baseline.clone()
 
-    for reform_path in reforms:
-        if isinstance(reform_path, str):
-            current_tax_benefit_system = current_tax_benefit_system.apply_reform(
-                reform_path
-            )
-        else:
-            current_tax_benefit_system = reform_path(current_tax_benefit_system)
-        current_tax_benefit_system._parameters_at_instant_cache = {}
+        for reform_path in reforms:
+            if isinstance(reform_path, str):
+                current_tax_benefit_system = current_tax_benefit_system.apply_reform(
+                    reform_path
+                )
+            else:
+                current_tax_benefit_system = reform_path(current_tax_benefit_system)
+            current_tax_benefit_system._parameters_at_instant_cache = {}
 
-    for extension in extensions:
-        current_tax_benefit_system = current_tax_benefit_system.clone()
-        current_tax_benefit_system.load_extension(extension)
+        for extension in extensions:
+            current_tax_benefit_system = current_tax_benefit_system.clone()
+            current_tax_benefit_system.load_extension(extension)
 
     if baseline_entry is None:
         baseline_entry = {}
diff --git a/tests/core/tools/test_runner/test_gc_pause.py b/tests/core/tools/test_runner/test_gc_pause.py
new file mode 100644
index 000000000..274675fd7
--- /dev/null
+++ b/tests/core/tools/test_runner/test_gc_pause.py
@@ -0,0 +1,314 @@
+"""The YAML runner pauses automatic collection while it builds a reform system.
+
+A cache miss in ``_get_tax_benefit_system`` builds with automatic collection
+paused and leaves ``gc.isenabled()`` as it found it, whether the build
+returns or raises and whether the caller had collection on or off. A cache
+hit does not touch collection. The pause changes no result: the cache still
+returns the identical system for a key it has seen and builds a new one for
+any other key, and YAML cases produce the same outcomes and values with or
+without it.
+"""
+
+import contextlib
+import gc
+import os
+
+import pytest
+
+from policyengine_core.country_template import CountryTaxBenefitSystem
+from policyengine_core.taxbenefitsystems import TaxBenefitSystem
+from policyengine_core.tools import test_runner
+from policyengine_core.tools.test_runner import (
+    _automatic_gc_paused,
+    _get_tax_benefit_system,
+)
+from tests.fixtures import yaml_tests
+
+# The country-package smoke job installs no dev dependencies.
+hypothesis = pytest.importorskip("hypothesis")
+st = hypothesis.strategies
+
+
+class BuildFailed(Exception):
+    pass
+
+
+class RecordingSystem:
+    """A stand-in system that records whether collection was enabled at each
+    build step, and can be told to raise at one of them."""
+
+    def __init__(self, log, fail_at=None, reforms=(), extensions=()):
+        self.log = log
+        self.fail_at = fail_at
+        self.reforms = reforms
+        self.extensions = extensions
+
+    def _step(self, name):
+        self.log.append((name, gc.isenabled()))
+        if self.fail_at == name:
+            raise BuildFailed(name)
+
+    def clone(self):
+        self._step("clone")
+        return RecordingSystem(self.log, self.fail_at, self.reforms, self.extensions)
+
+    def apply_reform(self, reform_path):
+        self._step("reform")
+        return RecordingSystem(
+            self.log, self.fail_at, self.reforms + (reform_path,), self.extensions
+        )
+
+    def load_extension(self, extension):
+        self._step("extension")
+        self.extensions = self.extensions + (extension,)
+
+
+def callable_reform(system):
+    system._step("reform")
+    return RecordingSystem(
+        system.log, system.fail_at, system.reforms + ("callable",), system.extensions
+    )
+
+
+@pytest.fixture
+def restore_gc():
+    was_enabled = gc.isenabled()
+    yield
+    if was_enabled:
+        gc.enable()
+    else:
+        gc.disable()
+
+
+def set_gc(enabled):
+    if enabled:
+        gc.enable()
+    else:
+        gc.disable()
+
+
+def test_a_cache_miss_builds_with_collection_paused(restore_gc):
+    gc.enable()
+    log = []
+    baseline = RecordingSystem(log)
+
+    system = _get_tax_benefit_system(
+        baseline, ["some.reform", callable_reform], ["some.extension"]
+    )
+
+    assert [name for name, _ in log] == [
+        "clone",
+        "reform",
+        "reform",
+        "clone",
+        "extension",
+    ]
+    assert [enabled for _, enabled in log] == [False] * 5
+    assert gc.isenabled()
+    assert system.reforms == ("some.reform", "callable")
+    assert system.extensions == ("some.extension",)
+
+
+SCENARIOS = ["miss", "hit", "clone_raises", "reform_raises", "extension_raises"]
+
+
+@pytest.mark.parametrize("enabled_before", [True, False])
+@pytest.mark.parametrize("scenario", SCENARIOS)
+def test_collection_state_after_a_call_equals_the_state_before(
+    restore_gc, enabled_before, scenario
+):
+    fail_at = {
+        "clone_raises": "clone",
+        "reform_raises": "reform",
+        "extension_raises": "extension",
+    }.get(scenario)
+    baseline = RecordingSystem([], fail_at=fail_at)
+    if scenario == "hit":
+        _get_tax_benefit_system(baseline, ["some.reform"], ["some.extension"])
+
+    set_gc(enabled_before)
+    if fail_at is None:
+        _get_tax_benefit_system(baseline, ["some.reform"], ["some.extension"])
+    else:
+        with pytest.raises(BuildFailed, match=fail_at):
+            _get_tax_benefit_system(baseline, ["some.reform"], ["some.extension"])
+    assert gc.isenabled() is enabled_before
+
+    # A failed build caches nothing, so the next call builds again.
+    log_length = len(baseline.log)
+    baseline.fail_at = None
+    _get_tax_benefit_system(baseline, ["some.reform"], ["some.extension"])
+    assert (len(baseline.log) > log_length) is (fail_at is not None)
+    assert gc.isenabled() is enabled_before
+
+
+def test_a_cache_hit_leaves_collection_alone(restore_gc, monkeypatch):
+    gc.enable()
+    baseline = RecordingSystem([])
+    system = _get_tax_benefit_system(baseline, ["some.reform"], [])
+    calls = []
+    for name in ("disable", "enable"):
+        real = getattr(gc, name)
+        monkeypatch.setattr(
+            gc, name, lambda real=real, name=name: calls.append(name) or real()
+        )
+
+    assert _get_tax_benefit_system(baseline, ["some.reform"], []) is system
+    assert calls == []
+
+
+def test_the_pause_restores_the_callers_state_when_the_body_raises(restore_gc):
+    for enabled_before in (True, False):
+        set_gc(enabled_before)
+        with pytest.raises(BuildFailed):
+            with _automatic_gc_paused():
+                assert not gc.isenabled()
+                raise BuildFailed("body")
+        assert gc.isenabled() is enabled_before
+
+
+REFORM_NAMES = ["a.reform", "b.reform", "c.reform"]
+EXTENSION_NAMES = ["x.extension", "y.extension"]
+
+request_strategy = st.tuples(
+    st.lists(st.sampled_from(REFORM_NAMES), max_size=3),
+    st.sampled_from([None, "", "gov.rate:1", "gov.rate:2"]),
+    st.lists(st.sampled_from(EXTENSION_NAMES), max_size=2),
+    st.booleans(),
+)
+
+
+@hypothesis.settings(
+    max_examples=200,
+    deadline=None,
+    suppress_health_check=[hypothesis.HealthCheck.function_scoped_fixture],
+)
+@hypothesis.given(requests=st.lists(request_strategy, max_size=25))
+def test_the_cache_matches_a_reference_dict_and_keeps_the_callers_state(
+    restore_gc, requests
+):
+    log = []
+    baseline = RecordingSystem(log)
+    reference = {}
+    for reforms, reform_key, extensions, enabled_before in requests:
+        key = (tuple(reforms), reform_key, frozenset(extensions))
+        clones_before = sum(name == "clone" for name, _ in log)
+        set_gc(enabled_before)
+
+        system = _get_tax_benefit_system(
+            baseline, list(reforms), list(extensions), reform_key=reform_key
+        )
+
+        assert gc.isenabled() is enabled_before
+        assert all(not enabled for _, enabled in log)
+        built = sum(name == "clone" for name, _ in log) > clones_before
+        if key in reference:
+            assert system is reference[key]
+            assert not built
+        else:
+            assert built
+            assert all(system is not other for other in reference.values())
+            assert system.reforms == tuple(reforms)
+            assert sorted(system.extensions) == sorted(extensions)
+            reference[key] = system
+
+
+yaml_tests_dir = os.path.dirname(yaml_tests.__file__)
+
+INLINE_PARAMETER_CASES = """
+- name: "A dotted parameter input builds a reform system"
+  period: 2017-01
+  input:
+    taxes.income_tax_rate: 0.2
+    salary: 2000
+  output:
+    income_tax: 400
+
+- name: "A second parameter value builds another"
+  period: 2017-01
+  input:
+    taxes.income_tax_rate: 0.25
+    salary: 2000
+  output:
+    income_tax: 500
+
+- name: "A failing case fails either way"
+  period: 2017-01
+  input:
+    taxes.income_tax_rate: 0.2
+    salary: 2000
+  output:
+    income_tax: 1
+"""
+
+
+def run_recording(paths, pause):
+    """Run the YAML runner on a fresh baseline, with the real pause or none.
+
+    Returns the exit code, each case's outcome, each computed value, and
+    whether collection was enabled at each system clone.
+    """
+    outcomes = []
+    values = []
+    gc_at_clones = []
+
+    class Outcomes:
+        def pytest_runtest_logreport(self, report):
+            if report.when == "call":
+                outcomes.append((report.nodeid, report.outcome))
+
+    assert_near = test_runner.assert_near
+    clone = TaxBenefitSystem.clone
+
+    def recording_assert_near(value, target_value, *args, **kwargs):
+        values.append((kwargs.get("message"), repr(value)))
+        return assert_near(value, target_value, *args, **kwargs)
+
+    def recording_clone(self):
+        gc_at_clones.append(gc.isenabled())
+        return clone(self)
+
+    baseline = CountryTaxBenefitSystem()
+    with pytest.MonkeyPatch.context() as patch:
+        patch.setattr(test_runner, "assert_near", recording_assert_near)
+        patch.setattr(TaxBenefitSystem, "clone", recording_clone)
+        if not pause:
+            patch.setattr(test_runner, "_automatic_gc_paused", contextlib.nullcontext)
+        # run_tests' own arguments, plus a plugin that records outcomes.
+        exit_code = pytest.main(
+            ["--capture", "no", "--maxfail", "0", "--tb", "no", "-q", *paths],
+            plugins=[test_runner.OpenFiscaPlugin(baseline, {}), Outcomes()],
+        )
+    return exit_code, outcomes, values, gc_at_clones
+
+
+def test_yaml_cases_have_the_same_outcomes_with_and_without_the_pause(
+    restore_gc, tmp_path
+):
+    gc.enable()
+    inline = tmp_path / "inline_parameters.yaml"
+    inline.write_text(INLINE_PARAMETER_CASES, encoding="utf-8")
+    paths = [
+        os.path.join(yaml_tests_dir, "test_with_reform.yaml"),
+        os.path.join(yaml_tests_dir, "test_with_reform_2.yaml"),
+        os.path.join(yaml_tests_dir, "test_with_extension.yaml"),
+        str(inline),
+    ]
+
+    paused = run_recording(paths, pause=True)
+    unpaused = run_recording(paths, pause=False)
+
+    exit_code, outcomes, values, gc_at_clones = paused
+    assert paused[:3] == unpaused[:3]
+    assert exit_code == pytest.ExitCode.TESTS_FAILED
+    assert [outcome for _, outcome in outcomes].count("failed") == 2
+    assert len(outcomes) == 8
+    assert len(values) >= 8
+    # Six builds (two reform lists, the empty combination, one extension and
+    # two parameter values) cloned the system with collection paused, and
+    # with the pause removed every clone ran with collection on.
+    assert len(gc_at_clones) >= 6
+    assert not any(gc_at_clones)
+    assert len(unpaused[3]) == len(gc_at_clones)
+    assert all(unpaused[3])
+    assert gc.isenabled()

From 569c56a715611626b35e0a994c051ae6e2a13d3f Mon Sep 17 00:00:00 2001
From: Max Ghenis <max@maxghenis.com>
Date: Fri, 30 Oct 2026 18:49:47 -0400
Subject: [PATCH 2/4] Collect once after a paused reform-system build

With collection only paused, the build's objects were walked three times
after it resumed: by a young, a middle and a full collection, all in the
next case. On policyengine-us that cost more than the collections the
pause avoided. Run one full collection when a build returns, for a caller
that had collection on, and test that a hit, a failed build and a caller
with collection off get none.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 ...se-gc-during-yaml-system-builds.changed.md |  2 +-
 policyengine_core/tools/test_runner.py        | 23 ++++---
 tests/core/tools/test_runner/test_gc_pause.py | 64 +++++++++++++------
 3 files changed, 60 insertions(+), 29 deletions(-)

diff --git a/changelog.d/pause-gc-during-yaml-system-builds.changed.md b/changelog.d/pause-gc-during-yaml-system-builds.changed.md
index c646bc023..526aefb4b 100644
--- a/changelog.d/pause-gc-during-yaml-system-builds.changed.md
+++ b/changelog.d/pause-gc-during-yaml-system-builds.changed.md
@@ -1 +1 @@
-Pause automatic garbage collection while the YAML test runner builds the system for a reform combination, and restore the caller's setting afterwards: full collections that freed almost nothing had taken about half of each build on policyengine-us.
+Pause automatic garbage collection while the YAML test runner builds the system for a reform combination, run one full collection when the build returns, and restore the caller's setting, so a cache-miss build no longer spends most of its time in repeated full collections that free almost nothing.
diff --git a/policyengine_core/tools/test_runner.py b/policyengine_core/tools/test_runner.py
index 4b6d92fb4..acdda6210 100644
--- a/policyengine_core/tools/test_runner.py
+++ b/policyengine_core/tools/test_runner.py
@@ -423,18 +423,25 @@ def pytest_collect_file(self, parent, file_path):
 
 @contextmanager
 def _automatic_gc_paused():
-    """Pause automatic garbage collection, then restore the caller's setting.
-
-    Building a reform system copies the whole baseline. Its allocations
-    trigger full collections that walk every object the baseline and its
-    copies hold and free almost nothing: on policyengine-us they took about
-    half of each build. Collection resumes as before once the build returns
-    or raises, and a caller that had disabled it keeps it disabled.
+    """Pause automatic garbage collection for a build, then collect once.
+
+    Building a reform system copies the whole baseline. On policyengine-us
+    the copy's allocations set off six or seven full collections per build,
+    each walking every object in the process and freeing almost nothing.
+    With collection only paused, the copy's objects were walked three times
+    once it resumed: by a young, a middle and a full collection. One full
+    collection at the end of a build that returns walks them once.
+
+    The caller's setting is restored when the build returns or raises. A
+    caller that had disabled collection gets no collection and keeps it
+    disabled.
     """
     was_enabled = gc.isenabled()
     gc.disable()
     try:
         yield
+        if was_enabled:
+            gc.collect()
     finally:
         if was_enabled:
             gc.enable()
@@ -465,7 +472,7 @@ def _get_tax_benefit_system(
         if cached is not None:
             return cached
 
-    # Only a cache miss builds a system, so only a miss pauses collection.
+    # Only a cache miss builds a system, so only a miss touches collection.
     with _automatic_gc_paused():
         current_tax_benefit_system = baseline.clone()
 
diff --git a/tests/core/tools/test_runner/test_gc_pause.py b/tests/core/tools/test_runner/test_gc_pause.py
index 274675fd7..8b8343f36 100644
--- a/tests/core/tools/test_runner/test_gc_pause.py
+++ b/tests/core/tools/test_runner/test_gc_pause.py
@@ -1,12 +1,13 @@
 """The YAML runner pauses automatic collection while it builds a reform system.
 
 A cache miss in ``_get_tax_benefit_system`` builds with automatic collection
-paused and leaves ``gc.isenabled()`` as it found it, whether the build
-returns or raises and whether the caller had collection on or off. A cache
-hit does not touch collection. The pause changes no result: the cache still
-returns the identical system for a key it has seen and builds a new one for
-any other key, and YAML cases produce the same outcomes and values with or
-without it.
+paused, then, if the caller had collection on and the build returned, runs
+one full collection. It leaves ``gc.isenabled()`` as it found it, whether the
+build returns or raises and whether the caller had collection on or off. A
+cache hit does not touch collection. The pause changes no result: the cache
+still returns the identical system for a key it has seen and builds a new
+one for any other key, and YAML cases produce the same outcomes and values
+with or without it.
 """
 
 import contextlib
@@ -80,6 +81,23 @@ def restore_gc():
         gc.disable()
 
 
+@pytest.fixture
+def gc_calls(restore_gc, monkeypatch):
+    """Every call to gc.disable, gc.enable and gc.collect, in order. Disabling
+    and enabling take effect; a recorded collection does not run, to keep
+    hundreds of builds fast."""
+    calls = []
+    for name in ("disable", "enable", "collect"):
+        real = getattr(gc, name) if name != "collect" else lambda *args: 0
+
+        def spy(*args, real=real, name=name):
+            calls.append(name)
+            return real(*args)
+
+        monkeypatch.setattr(gc, name, spy)
+    return calls
+
+
 def set_gc(enabled):
     if enabled:
         gc.enable()
@@ -87,10 +105,11 @@ def set_gc(enabled):
         gc.disable()
 
 
-def test_a_cache_miss_builds_with_collection_paused(restore_gc):
-    gc.enable()
+def test_a_cache_miss_builds_paused_then_collects_once(gc_calls):
+    set_gc(True)
     log = []
     baseline = RecordingSystem(log)
+    del gc_calls[:]
 
     system = _get_tax_benefit_system(
         baseline, ["some.reform", callable_reform], ["some.extension"]
@@ -104,6 +123,7 @@ def test_a_cache_miss_builds_with_collection_paused(restore_gc):
         "extension",
     ]
     assert [enabled for _, enabled in log] == [False] * 5
+    assert gc_calls == ["disable", "collect", "enable"]
     assert gc.isenabled()
     assert system.reforms == ("some.reform", "callable")
     assert system.extensions == ("some.extension",)
@@ -115,7 +135,7 @@ def test_a_cache_miss_builds_with_collection_paused(restore_gc):
 @pytest.mark.parametrize("enabled_before", [True, False])
 @pytest.mark.parametrize("scenario", SCENARIOS)
 def test_collection_state_after_a_call_equals_the_state_before(
-    restore_gc, enabled_before, scenario
+    gc_calls, enabled_before, scenario
 ):
     fail_at = {
         "clone_raises": "clone",
@@ -127,12 +147,15 @@ def test_collection_state_after_a_call_equals_the_state_before(
         _get_tax_benefit_system(baseline, ["some.reform"], ["some.extension"])
 
     set_gc(enabled_before)
+    del gc_calls[:]
     if fail_at is None:
         _get_tax_benefit_system(baseline, ["some.reform"], ["some.extension"])
     else:
         with pytest.raises(BuildFailed, match=fail_at):
             _get_tax_benefit_system(baseline, ["some.reform"], ["some.extension"])
     assert gc.isenabled() is enabled_before
+    # Only a build that returns, for a caller with collection on, collects.
+    assert gc_calls.count("collect") == int(scenario == "miss" and enabled_before)
 
     # A failed build caches nothing, so the next call builds again.
     log_length = len(baseline.log)
@@ -142,29 +165,27 @@ def test_collection_state_after_a_call_equals_the_state_before(
     assert gc.isenabled() is enabled_before
 
 
-def test_a_cache_hit_leaves_collection_alone(restore_gc, monkeypatch):
-    gc.enable()
+def test_a_cache_hit_leaves_collection_alone(gc_calls):
+    set_gc(True)
     baseline = RecordingSystem([])
     system = _get_tax_benefit_system(baseline, ["some.reform"], [])
-    calls = []
-    for name in ("disable", "enable"):
-        real = getattr(gc, name)
-        monkeypatch.setattr(
-            gc, name, lambda real=real, name=name: calls.append(name) or real()
-        )
+    del gc_calls[:]
 
     assert _get_tax_benefit_system(baseline, ["some.reform"], []) is system
-    assert calls == []
+    assert gc_calls == []
+    assert gc.isenabled()
 
 
-def test_the_pause_restores_the_callers_state_when_the_body_raises(restore_gc):
+def test_the_pause_restores_the_callers_state_when_the_body_raises(gc_calls):
     for enabled_before in (True, False):
         set_gc(enabled_before)
+        del gc_calls[:]
         with pytest.raises(BuildFailed):
             with _automatic_gc_paused():
                 assert not gc.isenabled()
                 raise BuildFailed("body")
         assert gc.isenabled() is enabled_before
+        assert "collect" not in gc_calls
 
 
 REFORM_NAMES = ["a.reform", "b.reform", "c.reform"]
@@ -185,7 +206,7 @@ def test_the_pause_restores_the_callers_state_when_the_body_raises(restore_gc):
 )
 @hypothesis.given(requests=st.lists(request_strategy, max_size=25))
 def test_the_cache_matches_a_reference_dict_and_keeps_the_callers_state(
-    restore_gc, requests
+    gc_calls, requests
 ):
     log = []
     baseline = RecordingSystem(log)
@@ -194,6 +215,7 @@ def test_the_cache_matches_a_reference_dict_and_keeps_the_callers_state(
         key = (tuple(reforms), reform_key, frozenset(extensions))
         clones_before = sum(name == "clone" for name, _ in log)
         set_gc(enabled_before)
+        del gc_calls[:]
 
         system = _get_tax_benefit_system(
             baseline, list(reforms), list(extensions), reform_key=reform_key
@@ -202,9 +224,11 @@ def test_the_cache_matches_a_reference_dict_and_keeps_the_callers_state(
         assert gc.isenabled() is enabled_before
         assert all(not enabled for _, enabled in log)
         built = sum(name == "clone" for name, _ in log) > clones_before
+        assert gc_calls.count("collect") == int(built and enabled_before)
         if key in reference:
             assert system is reference[key]
             assert not built
+            assert gc_calls == []
         else:
             assert built
             assert all(system is not other for other in reference.values())

From 4baa215ce8e976e17119f5ca344d60ab495b8165 Mon Sep 17 00:00:00 2001
From: Max Ghenis <mghenis@gmail.com>
Date: Sat, 31 Oct 2026 00:27:00 -0400
Subject: [PATCH 3/4] Force the post-build collection only while the cache
 holds few systems

A forced full collection after every cache miss walks every cached system
again, so a run of misses cost quadratic time: 100 distinct parameter
reforms on the country template took 6.4 s of CPU against master's 0.09 s.
Force it only while the cache holds fewer than _FULL_COLLECTION_CACHE_LIMIT
(3) systems, and leave later copies to the automatic collector, whose
quarter-of-the-heap rule keeps full collections rare. A miss past the limit
still pauses collection and restores the caller's setting.

The tests' spy now records each collection's generation and asserts a full
one, so gc.collect(1) no longer passes. New tests pin where the last forced
collection and the first skipped one fall for several limits, that a miss
past the limit still pauses and restores, and that the count covers every
baseline's entries; the Hypothesis cache test checks at most the limit's
number of forced collections per sequence.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 ...se-gc-during-yaml-system-builds.changed.md |   2 +-
 policyengine_core/tools/test_runner.py        |  37 ++++-
 tests/core/tools/test_runner/test_gc_pause.py | 145 +++++++++++++++---
 3 files changed, 155 insertions(+), 29 deletions(-)

diff --git a/changelog.d/pause-gc-during-yaml-system-builds.changed.md b/changelog.d/pause-gc-during-yaml-system-builds.changed.md
index 526aefb4b..723266125 100644
--- a/changelog.d/pause-gc-during-yaml-system-builds.changed.md
+++ b/changelog.d/pause-gc-during-yaml-system-builds.changed.md
@@ -1 +1 @@
-Pause automatic garbage collection while the YAML test runner builds the system for a reform combination, run one full collection when the build returns, and restore the caller's setting, so a cache-miss build no longer spends most of its time in repeated full collections that free almost nothing.
+Pause automatic garbage collection while the YAML test runner builds the system for a reform combination, run one full collection when the build returns while the cache holds fewer than three systems, and restore the caller's setting, so a cache-miss build no longer spends most of its time in repeated full collections that free almost nothing.
diff --git a/policyengine_core/tools/test_runner.py b/policyengine_core/tools/test_runner.py
index acdda6210..e34f107d6 100644
--- a/policyengine_core/tools/test_runner.py
+++ b/policyengine_core/tools/test_runner.py
@@ -421,26 +421,45 @@ def pytest_collect_file(self, parent, file_path):
             )
 
 
+# A cache miss forces a full collection only while the cache holds fewer than
+# this many systems: the forced collection walks every object in the process,
+# and in #643's padded-template model it stopped paying for itself once two
+# or three systems were cached, while the automatic collector's full
+# collections grew rarer as the heap grew.
+_FULL_COLLECTION_CACHE_LIMIT = 3
+
+
 @contextmanager
-def _automatic_gc_paused():
-    """Pause automatic garbage collection for a build, then collect once.
+def _automatic_gc_paused(collect=True):
+    """Pause automatic garbage collection for a build, then maybe collect once.
 
     Building a reform system copies the whole baseline. On policyengine-us
     the copy's allocations set off six or seven full collections per build,
     each walking every object in the process and freeing almost nothing.
-    With collection only paused, the copy's objects were walked three times
-    once it resumed: by a young, a middle and a full collection. One full
-    collection at the end of a build that returns walks them once.
+    Pausing collection skips them. When collection resumes, the copy's
+    objects pass through a young and a middle collection, and through a full
+    one only when CPython's growth rule allows it: an automatic full
+    collection waits until the objects moved into the oldest generation since
+    the last full collection number at least a quarter of those that survived
+    it.
+
+    With ``collect``, a build that returns ends with one ``gc.collect()``. That
+    full collection bypasses the growth rule and moves the copy straight into
+    the oldest generation, so its objects skip the young and middle
+    collections. It pays only while the copy is a large share of the heap, so
+    ``_get_tax_benefit_system`` asks for it only while the cache holds fewer
+    than ``_FULL_COLLECTION_CACHE_LIMIT`` systems and otherwise leaves the copy
+    to the automatic collector.
 
     The caller's setting is restored when the build returns or raises. A
     caller that had disabled collection gets no collection and keeps it
-    disabled.
+    disabled, and a build that raises gets none.
     """
     was_enabled = gc.isenabled()
     gc.disable()
     try:
         yield
-        if was_enabled:
+        if was_enabled and collect:
             gc.collect()
     finally:
         if was_enabled:
@@ -473,7 +492,9 @@ def _get_tax_benefit_system(
             return cached
 
     # Only a cache miss builds a system, so only a miss touches collection.
-    with _automatic_gc_paused():
+    # The count covers every baseline's entries.
+    cached_systems = sum(len(entry) for entry in _tax_benefit_system_cache.values())
+    with _automatic_gc_paused(collect=cached_systems < _FULL_COLLECTION_CACHE_LIMIT):
         current_tax_benefit_system = baseline.clone()
 
         for reform_path in reforms:
diff --git a/tests/core/tools/test_runner/test_gc_pause.py b/tests/core/tools/test_runner/test_gc_pause.py
index 8b8343f36..1d70a02ac 100644
--- a/tests/core/tools/test_runner/test_gc_pause.py
+++ b/tests/core/tools/test_runner/test_gc_pause.py
@@ -1,18 +1,21 @@
 """The YAML runner pauses automatic collection while it builds a reform system.
 
 A cache miss in ``_get_tax_benefit_system`` builds with automatic collection
-paused, then, if the caller had collection on and the build returned, runs
-one full collection. It leaves ``gc.isenabled()`` as it found it, whether the
-build returns or raises and whether the caller had collection on or off. A
-cache hit does not touch collection. The pause changes no result: the cache
-still returns the identical system for a key it has seen and builds a new
-one for any other key, and YAML cases produce the same outcomes and values
-with or without it.
+paused. If the caller had collection on, the build returned, and the cache
+held fewer than ``_FULL_COLLECTION_CACHE_LIMIT`` systems, it then runs one
+full collection; otherwise it leaves the new system to the automatic
+collector, so a run of misses forces at most that many collections. It leaves
+``gc.isenabled()`` as it found it, whether the build returns or raises and
+whether the caller had collection on or off. A cache hit does not touch
+collection. The pause changes no result: the cache still returns the
+identical system for a key it has seen and builds a new one for any other
+key, and YAML cases produce the same outcomes and values with or without it.
 """
 
 import contextlib
 import gc
 import os
+import weakref
 
 import pytest
 
@@ -81,23 +84,47 @@ def restore_gc():
         gc.disable()
 
 
+@pytest.fixture(autouse=True)
+def empty_cache(monkeypatch):
+    """Give each test its own empty system cache, so systems other tests
+    cached do not count toward the full-collection bound."""
+    monkeypatch.setattr(
+        test_runner, "_tax_benefit_system_cache", weakref.WeakKeyDictionary()
+    )
+
+
+# gc.collect()'s default generation, the oldest: a full collection.
+FULL = "collect(2)"
+LIMIT = test_runner._FULL_COLLECTION_CACHE_LIMIT
+
+
 @pytest.fixture
 def gc_calls(restore_gc, monkeypatch):
-    """Every call to gc.disable, gc.enable and gc.collect, in order. Disabling
-    and enabling take effect; a recorded collection does not run, to keep
-    hundreds of builds fast."""
+    """Every call to gc.disable, gc.enable and gc.collect, in order, with the
+    generation each collection asked for: ``FULL`` for ``gc.collect()``.
+    Disabling and enabling take effect; a recorded collection does not run,
+    to keep hundreds of builds fast."""
     calls = []
-    for name in ("disable", "enable", "collect"):
-        real = getattr(gc, name) if name != "collect" else lambda *args: 0
+    for name in ("disable", "enable"):
 
-        def spy(*args, real=real, name=name):
+        def spy(real=getattr(gc, name), name=name):
             calls.append(name)
-            return real(*args)
+            return real()
 
         monkeypatch.setattr(gc, name, spy)
+
+    def collect(generation=2):
+        calls.append(f"collect({generation})")
+        return 0
+
+    monkeypatch.setattr(gc, "collect", collect)
     return calls
 
 
+def collections(calls):
+    return [call for call in calls if call.startswith("collect")]
+
+
 def set_gc(enabled):
     if enabled:
         gc.enable()
@@ -123,7 +150,7 @@ def test_a_cache_miss_builds_paused_then_collects_once(gc_calls):
         "extension",
     ]
     assert [enabled for _, enabled in log] == [False] * 5
-    assert gc_calls == ["disable", "collect", "enable"]
+    assert gc_calls == ["disable", FULL, "enable"]
     assert gc.isenabled()
     assert system.reforms == ("some.reform", "callable")
     assert system.extensions == ("some.extension",)
@@ -155,7 +182,8 @@ def test_collection_state_after_a_call_equals_the_state_before(
             _get_tax_benefit_system(baseline, ["some.reform"], ["some.extension"])
     assert gc.isenabled() is enabled_before
     # Only a build that returns, for a caller with collection on, collects.
-    assert gc_calls.count("collect") == int(scenario == "miss" and enabled_before)
+    expected = [FULL] if scenario == "miss" and enabled_before else []
+    assert collections(gc_calls) == expected
 
     # A failed build caches nothing, so the next call builds again.
     log_length = len(baseline.log)
@@ -185,7 +213,71 @@ def test_the_pause_restores_the_callers_state_when_the_body_raises(gc_calls):
                 assert not gc.isenabled()
                 raise BuildFailed("body")
         assert gc.isenabled() is enabled_before
-        assert "collect" not in gc_calls
+        assert collections(gc_calls) == []
+
+
+def miss(baseline, number):
+    """A cache miss: a parameter key the cache has not seen."""
+    return _get_tax_benefit_system(baseline, [], [], reform_key=f"gov.rate:{number}")
+
+
+@pytest.mark.parametrize("limit", [0, 1, 2, LIMIT, LIMIT + 2])
+def test_misses_collect_only_while_the_cache_holds_fewer_systems_than_the_limit(
+    gc_calls, monkeypatch, limit
+):
+    monkeypatch.setattr(test_runner, "_FULL_COLLECTION_CACHE_LIMIT", limit)
+    set_gc(True)
+    baseline = RecordingSystem([])
+    collected = []
+    for number in range(limit + 3):
+        del gc_calls[:]
+        miss(baseline, number)
+        collected.append(collections(gc_calls))
+        assert gc.isenabled()
+    # The last forced collection is miss number limit - 1, with limit - 1
+    # systems cached; from miss number limit on, none.
+    assert collected == [[FULL]] * limit + [[]] * 3
+
+
+def test_a_miss_past_the_limit_still_pauses_and_restores_collection(gc_calls):
+    log = []
+    baseline = RecordingSystem(log)
+    for enabled_before in (True, False):
+        set_gc(True)
+        for number in range(LIMIT):
+            miss(baseline, number)
+        del log[:]
+        set_gc(enabled_before)
+        del gc_calls[:]
+
+        system = _get_tax_benefit_system(baseline, ["some.reform"], ["x.extension"])
+
+        assert [name for name, _ in log] == ["clone", "reform", "clone", "extension"]
+        assert not any(enabled for _, enabled in log)
+        assert gc_calls == (["disable", "enable"] if enabled_before else ["disable"])
+        assert gc.isenabled() is enabled_before
+        assert system.reforms == ("some.reform",)
+        baseline = RecordingSystem(log)
+
+
+def test_the_limit_counts_systems_cached_for_every_baseline(gc_calls, monkeypatch):
+    monkeypatch.setattr(test_runner, "_FULL_COLLECTION_CACHE_LIMIT", 3)
+    set_gc(True)
+    first = RecordingSystem([])
+    miss(first, 0)
+    miss(first, 1)
+    second = RecordingSystem([])
+    del gc_calls[:]
+    miss(second, 0)  # 2 cached
+    assert collections(gc_calls) == [FULL]
+    del gc_calls[:]
+    miss(second, 1)  # 3 cached
+    assert collections(gc_calls) == []
+    # A baseline that goes away takes its systems out of the count.
+    del first
+    del gc_calls[:]
+    miss(second, 2)  # 2 cached
+    assert collections(gc_calls) == [FULL]
 
 
 REFORM_NAMES = ["a.reform", "b.reform", "c.reform"]
@@ -206,11 +298,16 @@ def test_the_pause_restores_the_callers_state_when_the_body_raises(gc_calls):
 )
 @hypothesis.given(requests=st.lists(request_strategy, max_size=25))
 def test_the_cache_matches_a_reference_dict_and_keeps_the_callers_state(
-    gc_calls, requests
+    gc_calls, monkeypatch, requests
 ):
+    # Each example starts from an empty cache.
+    monkeypatch.setattr(
+        test_runner, "_tax_benefit_system_cache", weakref.WeakKeyDictionary()
+    )
     log = []
     baseline = RecordingSystem(log)
     reference = {}
+    forced = 0
     for reforms, reform_key, extensions, enabled_before in requests:
         key = (tuple(reforms), reform_key, frozenset(extensions))
         clones_before = sum(name == "clone" for name, _ in log)
@@ -224,7 +321,11 @@ def test_the_cache_matches_a_reference_dict_and_keeps_the_callers_state(
         assert gc.isenabled() is enabled_before
         assert all(not enabled for _, enabled in log)
         built = sum(name == "clone" for name, _ in log) > clones_before
-        assert gc_calls.count("collect") == int(built and enabled_before)
+        # The reference dict holds what the cache held before this call.
+        collects = built and enabled_before and len(reference) < LIMIT
+        assert collections(gc_calls) == ([FULL] if collects else [])
+        forced += collects
+        assert forced <= LIMIT
         if key in reference:
             assert system is reference[key]
             assert not built
@@ -297,7 +398,11 @@ def recording_clone(self):
         patch.setattr(test_runner, "assert_near", recording_assert_near)
         patch.setattr(TaxBenefitSystem, "clone", recording_clone)
         if not pause:
-            patch.setattr(test_runner, "_automatic_gc_paused", contextlib.nullcontext)
+            patch.setattr(
+                test_runner,
+                "_automatic_gc_paused",
+                lambda collect=True: contextlib.nullcontext(),
+            )
         # run_tests' own arguments, plus a plugin that records outcomes.
         exit_code = pytest.main(
             ["--capture", "no", "--maxfail", "0", "--tb", "no", "-q", *paths],

From 8a2d28440bb0583cc57d8a328723b0e026c5cd7a Mon Sep 17 00:00:00 2001
From: Max Ghenis <mghenis@gmail.com>
Date: Sat, 31 Oct 2026 01:49:59 -0400
Subject: [PATCH 4/4] Give the full-collection limit's measured reason

The padded-template model, sized to policyengine-us's proportions (a copy
adds about 17.3 million tracked objects to a heap of 3.9 million), gave the
least collection CPU over 16 and 32 misses with a limit of 3; limits from 1
to 8 came within 20%, and no limit cost about three times as much.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 policyengine_core/tools/test_runner.py | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/policyengine_core/tools/test_runner.py b/policyengine_core/tools/test_runner.py
index e34f107d6..80adfa50b 100644
--- a/policyengine_core/tools/test_runner.py
+++ b/policyengine_core/tools/test_runner.py
@@ -422,10 +422,12 @@ def pytest_collect_file(self, parent, file_path):
 
 
 # A cache miss forces a full collection only while the cache holds fewer than
-# this many systems: the forced collection walks every object in the process,
-# and in #643's padded-template model it stopped paying for itself once two
-# or three systems were cached, while the automatic collector's full
-# collections grew rarer as the heap grew.
+# this many systems. That collection walks every object in the process, so
+# each cached system makes it dearer, while the automatic collector's full
+# collections grow rarer as the heap grows. In PR #643's model of
+# policyengine-us-sized systems, 3 gave the least collection CPU over 16 and
+# 32 misses; limits from 1 to 8 came within 20%, and no limit cost about three
+# times as much.
 _FULL_COLLECTION_CACHE_LIMIT = 3