From 8607c7354166888c55b0ba00da606362e2e1bd4a Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 08:35:41 -0400 Subject: [PATCH 1/9] Release per-case state in the YAML test runner and bound its reform cache pytest keeps every collected item until the session ends, and each YamlItem kept its simulation and the system it ran on. The runner also cached every reform system it built, a full clone of the baseline each, forever. A policyengine-us run of ~1,500 baseline files in one process reached 118 GB. Co-Authored-By: Claude Opus 5.5 --- changelog.d/fix-test-runner-memory.fixed.md | 1 + .../scripts/policyengine_command.py | 6 + policyengine_core/scripts/run_test.py | 1 + policyengine_core/tools/test_runner.py | 81 ++++++++-- .../test_runner/test_reform_cache_property.py | 98 ++++++++++++ .../tools/test_runner/test_runner_memory.py | 147 ++++++++++++++++++ tests/fixtures/yaml_runner_memory.py | 92 +++++++++++ 7 files changed, 414 insertions(+), 12 deletions(-) create mode 100644 changelog.d/fix-test-runner-memory.fixed.md create mode 100644 tests/core/tools/test_runner/test_reform_cache_property.py create mode 100644 tests/core/tools/test_runner/test_runner_memory.py create mode 100644 tests/fixtures/yaml_runner_memory.py diff --git a/changelog.d/fix-test-runner-memory.fixed.md b/changelog.d/fix-test-runner-memory.fixed.md new file mode 100644 index 000000000..d98ae6f8e --- /dev/null +++ b/changelog.d/fix-test-runner-memory.fixed.md @@ -0,0 +1 @@ +Stop the YAML test runner keeping every case's simulation and every reform system alive for the whole run: finished cases release their simulation, and the reform-system cache keeps the 2 most recently used (`--reform-cache-size`). diff --git a/policyengine_core/scripts/policyengine_command.py b/policyengine_core/scripts/policyengine_command.py index 7aa85c888..8112615a7 100644 --- a/policyengine_core/scripts/policyengine_command.py +++ b/policyengine_core/scripts/policyengine_command.py @@ -91,6 +91,12 @@ def build_test_parser(parser): default=None, help="variables to ignore. If specified, do not test the given variables.", ) + parser.add_argument( + "--reform-cache-size", + type=int, + default=None, + help="how many reform systems (built for tests with reforms, extensions or parameter inputs) to keep cached; each is a full copy of the tax-benefit system. Defaults to 2; 0 caches none.", + ) return parser diff --git a/policyengine_core/scripts/run_test.py b/policyengine_core/scripts/run_test.py index 5eb4c7c7e..123f49fa0 100644 --- a/policyengine_core/scripts/run_test.py +++ b/policyengine_core/scripts/run_test.py @@ -30,6 +30,7 @@ def main(parser): "only_variables": args.only_variables, "ignore_variables": args.ignore_variables, "visualize": args.visualize, + "reform_cache_size": args.reform_cache_size, } paths = [os.path.abspath(path) for path in args.path] diff --git a/policyengine_core/tools/test_runner.py b/policyengine_core/tools/test_runner.py index 6b8d4aefc..d34b8b7c7 100644 --- a/policyengine_core/tools/test_runner.py +++ b/policyengine_core/tools/test_runner.py @@ -67,9 +67,18 @@ def import_yaml(): # reuses after GC — a collected baseline could produce the same id as a # completely unrelated new baseline and hit a stale cache entry (bug H9). import weakref +from collections import OrderedDict _tax_benefit_system_cache: "weakref.WeakKeyDictionary" = weakref.WeakKeyDictionary() +# How many reform systems (cases with ``reforms``, ``extensions`` or dotted +# parameter inputs) each baseline keeps, least recently used evicted first. +# Every one is a full clone of the baseline's parameters and variables, which +# is over a gigabyte for policyengine-us, and an unbounded cache made a long run +# grow by that much for every distinct combination it met. The reform-free +# system that most cases share is always kept and does not count. +DEFAULT_REFORM_CACHE_SIZE = 2 + def run_tests(tax_benefit_system, paths, options=None): """ @@ -248,6 +257,7 @@ def apply(self): reform_key="=".join( [f"{key}:{value}" for key, value in parametric_reform_items] ), + cache_size=self.options.get("reform_cache_size"), ) verbose = self.options.get("verbose") performance_graph = self.options.get("performance_graph") @@ -288,6 +298,22 @@ def apply(self): if visualize: self.generate_variable_graph(tracer) + def teardown(self): + # pytest keeps every collected item until the session ends, so a + # finished case must not keep its simulation, or the system it ran on, + # alive: over a run of thousands of cases that retained every + # simulation and every reform system any case used. + system = self.tax_benefit_system + if ( + self.simulation is not None + and getattr(system, "simulation", None) is self.simulation + ): + # ``Simulation.__init__`` registers itself on its system, and that + # system is cached across cases. + system.simulation = None + self.simulation = None + self.tax_benefit_system = None + def print_computation_log(self, tracer): print("Computation log:") # noqa T001 tracer.print_computation_log() @@ -419,30 +445,61 @@ def pytest_collect_file(self, parent, file_path): ) +class _CachedSystems: + """One baseline's cached systems: the reform-free one, kept for the + baseline's lifetime, and an LRU of reform systems.""" + + def __init__(self): + self.reform_free = None + self.reforms = OrderedDict() + + def _get_tax_benefit_system( baseline, reforms, extensions, reform_key=None, + cache_size=None, ): + """Return the system a YAML case runs on. + + ``cache_size`` bounds how many reform systems stay cached for ``baseline`` + (default ``DEFAULT_REFORM_CACHE_SIZE``; 0 caches none). The reform-free + system is always cached. + """ + if cache_size is None: + cache_size = DEFAULT_REFORM_CACHE_SIZE + if cache_size < 0: + raise ValueError(f"cache_size must be 0 or more, got {cache_size}") if not isinstance(reforms, list): reforms = [reforms] if not isinstance(extensions, list): extensions = [extensions] - # Inner key is (reforms in order, reform_key, extensions as a frozenset). - # The outer cache is keyed on the baseline *object* so its entries are - # tied to the baseline's lifetime (see bug H9 above). + # The cache is keyed on the baseline *object* so its entries are tied to + # the baseline's lifetime (see bug H9 above). + cached_systems = _tax_benefit_system_cache.get(baseline) + if cached_systems is None: + cached_systems = _CachedSystems() + _tax_benefit_system_cache[baseline] = cached_systems + + if not reforms and not extensions and not reform_key: + if cached_systems.reform_free is None: + cached_systems.reform_free = baseline.clone() + return cached_systems.reform_free + + # Key: (reforms in order, reform_key, extensions as a frozenset). Inline + # reform classes are rebuilt for every case, so they key by their + # ``reform_key`` (the parameter values they set), not by identity. inner_key = ( ":".join([reform if isinstance(reform, str) else "" for reform in reforms]), reform_key, frozenset(extensions), ) - baseline_entry = _tax_benefit_system_cache.get(baseline) - if baseline_entry is not None: - cached = baseline_entry.get(inner_key) - if cached is not None: - return cached + cached = cached_systems.reforms.get(inner_key) + if cached is not None: + cached_systems.reforms.move_to_end(inner_key) + return cached current_tax_benefit_system = baseline.clone() @@ -459,10 +516,10 @@ def _get_tax_benefit_system( current_tax_benefit_system = current_tax_benefit_system.clone() current_tax_benefit_system.load_extension(extension) - if baseline_entry is None: - baseline_entry = {} - _tax_benefit_system_cache[baseline] = baseline_entry - baseline_entry[inner_key] = current_tax_benefit_system + if cache_size > 0: + cached_systems.reforms[inner_key] = current_tax_benefit_system + while len(cached_systems.reforms) > cache_size: + cached_systems.reforms.popitem(last=False) return current_tax_benefit_system diff --git a/tests/core/tools/test_runner/test_reform_cache_property.py b/tests/core/tools/test_runner/test_reform_cache_property.py new file mode 100644 index 000000000..ff4622c91 --- /dev/null +++ b/tests/core/tools/test_runner/test_reform_cache_property.py @@ -0,0 +1,98 @@ +"""The YAML runner's reform-system cache against a reference LRU. + +For any sequence of requests and any cache size, the cache holds at most +``cache_size`` reform systems, returns a cached object exactly when a plain +LRU of that size would, always hands back a system built for the reforms and +extensions requested, and keeps one reform-free system throughout. +``test_runner_memory.py`` pins the same behaviour with examples. +""" + +from __future__ import annotations + +from collections import OrderedDict + +import pytest + +# The country-package smoke job installs no dev dependencies. +hypothesis = pytest.importorskip("hypothesis") +st = hypothesis.strategies + +from policyengine_core.tools.test_runner import ( + _get_tax_benefit_system, + _tax_benefit_system_cache, +) + + +class StubSystem: + """Records what the runner applied to it, in order.""" + + def __init__(self, applied=()): + self.applied = tuple(applied) + + def clone(self): + return StubSystem(self.applied) + + def apply_reform(self, path): + return StubSystem(self.applied + (("reform", path),)) + + def load_extension(self, extension): + self.applied += (("extension", extension),) + + +requests = st.tuples( + st.lists(st.sampled_from(["a", "b", "c"]), max_size=2), + st.lists(st.sampled_from(["x", "y"]), max_size=2, unique=True), + st.sampled_from(["", "p:1", "p:2"]), +) + + +@hypothesis.given( + sequence=st.lists(requests, max_size=40), + cache_size=st.integers(min_value=0, max_value=3), +) +@hypothesis.settings(max_examples=300, deadline=None) +def test_cache_matches_reference_lru(sequence, cache_size): + baseline = StubSystem() + reference = OrderedDict() + last_returned = {} + reform_free = None + + for reforms, extensions, reform_key in sequence: + system = _get_tax_benefit_system( + baseline, + list(reforms), + list(extensions), + reform_key=reform_key, + cache_size=cache_size, + ) + assert system is not baseline + # Reforms apply in the order given; extension order is not part of + # the cache key, so only the set is fixed. + applied_reforms = [name for kind, name in system.applied if kind == "reform"] + applied_extensions = { + name for kind, name in system.applied if kind == "extension" + } + assert applied_reforms == list(reforms) + assert applied_extensions == set(extensions) + + if not reforms and not extensions and not reform_key: + if reform_free is None: + reform_free = system + assert system is reform_free + else: + key = (":".join(reforms), reform_key, frozenset(extensions)) + if key in reference: + reference.move_to_end(key) + assert system is last_returned[key] + else: + assert system is not last_returned.get(key) + if cache_size: + reference[key] = True + while len(reference) > cache_size: + reference.popitem(last=False) + last_returned[key] = system + + cached = _tax_benefit_system_cache[baseline] + assert len(cached.reforms) <= cache_size + assert list(cached.reforms) == list(reference) + assert cached.reform_free is reform_free diff --git a/tests/core/tools/test_runner/test_runner_memory.py b/tests/core/tools/test_runner/test_runner_memory.py new file mode 100644 index 000000000..17d805cf2 --- /dev/null +++ b/tests/core/tools/test_runner/test_runner_memory.py @@ -0,0 +1,147 @@ +"""The YAML runner must not keep per-case state alive across a run. + +pytest holds every collected item until the session ends. When each +``YamlItem`` kept its simulation, and the runner kept every reform system it +ever built, memory grew with every case: a policyengine-us run of ~1,500 +files in one process reached 118 GB. +""" + +import gc +import weakref + +import pytest + +from policyengine_core.country_template import CountryTaxBenefitSystem +from policyengine_core.reforms import Reform, set_parameter +from policyengine_core.tools import test_runner +from policyengine_core.tools.test_runner import ( + DEFAULT_REFORM_CACHE_SIZE, + _get_tax_benefit_system, + _tax_benefit_system_cache, +) +from tests.fixtures.yaml_runner_memory import run_with_probe, write_cases + + +def _rate_reform(rate): + """An inline reform class like the one the runner builds for a dotted + parameter input, and the ``reform_key`` it would get.""" + modifier = set_parameter( + "taxes.income_tax_rate", rate, return_modifier=True, period="year:2000:40" + ) + + class inline_reform_class(Reform): + def apply(self): + self.parameters = modifier(self.parameters) + + return [inline_reform_class], f"taxes.income_tax_rate:{rate}" + + +def _system_for_rate(baseline, rate, cache_size=None): + reforms, key = _rate_reform(rate) + return _get_tax_benefit_system( + baseline, reforms, [], reform_key=key, cache_size=cache_size + ) + + +def _rate(system): + return system.parameters.taxes.income_tax_rate("2017-01-01") + + +def test_reform_cache_keeps_at_most_cache_size_reform_systems(): + baseline = CountryTaxBenefitSystem() + refs = [weakref.ref(_system_for_rate(baseline, 0.2 + i / 100)) for i in range(8)] + gc.collect() + + cached = _tax_benefit_system_cache[baseline].reforms + assert len(cached) == DEFAULT_REFORM_CACHE_SIZE + # Evicted systems are freed, not just dropped from the cache. + alive = [ref() for ref in refs if ref() is not None] + assert len(alive) == DEFAULT_REFORM_CACHE_SIZE + assert set(map(id, alive)) == set(map(id, cached.values())) + + +def test_reform_cache_evicts_least_recently_used(): + baseline = CountryTaxBenefitSystem() + a = _system_for_rate(baseline, 0.21, cache_size=2) + b = _system_for_rate(baseline, 0.22, cache_size=2) + assert _system_for_rate(baseline, 0.21, cache_size=2) is a # hit; A is newest + _system_for_rate(baseline, 0.23, cache_size=2) # evicts B, not A + assert _system_for_rate(baseline, 0.21, cache_size=2) is a + rebuilt_b = _system_for_rate(baseline, 0.22, cache_size=2) + assert rebuilt_b is not b + assert _rate(rebuilt_b) == pytest.approx(0.22) + + +def test_reform_free_system_is_kept_through_evictions(): + baseline = CountryTaxBenefitSystem() + reform_free = _get_tax_benefit_system(baseline, [], [], reform_key="") + for i in range(6): + _system_for_rate(baseline, 0.3 + i / 100, cache_size=1) + assert _get_tax_benefit_system(baseline, [], [], reform_key="") is reform_free + assert _get_tax_benefit_system(baseline, [], []) is reform_free + assert reform_free is not baseline + assert _rate(reform_free) == pytest.approx(0.15) + assert _rate(baseline) == pytest.approx(0.15) + + +def test_cache_size_zero_builds_a_fresh_reform_system_each_time(): + baseline = CountryTaxBenefitSystem() + first = _system_for_rate(baseline, 0.25, cache_size=0) + second = _system_for_rate(baseline, 0.25, cache_size=0) + assert first is not second + assert _rate(first) == _rate(second) == pytest.approx(0.25) + assert len(_tax_benefit_system_cache[baseline].reforms) == 0 + + +def test_negative_cache_size_is_rejected(): + with pytest.raises(ValueError): + _system_for_rate(CountryTaxBenefitSystem(), 0.25, cache_size=-1) + + +def test_finished_cases_release_their_simulations(tmp_path): + path = write_cases(tmp_path / "cases.yaml", 60) + probe = run_with_probe(CountryTaxBenefitSystem(), path) + assert len(probe.outcomes) == 60 + # Every item is still referenced by the session here; none may keep the + # simulation it ran. + assert probe.live_simulations == 0 + + +def test_live_systems_do_not_grow_with_the_number_of_cases(tmp_path): + short = run_with_probe( + CountryTaxBenefitSystem(), write_cases(tmp_path / "short.yaml", 12) + ) + long = run_with_probe( + CountryTaxBenefitSystem(), write_cases(tmp_path / "long.yaml", 90) + ) + # 4 vs 30 distinct reform systems requested; the live count must not move. + assert long.live_systems == short.live_systems + + +@pytest.mark.parametrize("reform_every", [0, 3]) +def test_traced_memory_stays_flat_over_hundreds_of_cases(tmp_path, reform_every): + """A few hundred cases in one process: memory after the last case is + within a small bound of memory after the first hundred.""" + path = write_cases(tmp_path / "cases.yaml", 300, reform_every=reform_every) + probe = run_with_probe(CountryTaxBenefitSystem(), path) + assert len(probe.outcomes) == 300 + # One sample every 10 cases: index 9 is after case 100. + growth = max(probe.traced[10:]) - probe.traced[9] + # Measured over these 200 cases: 0.5 MB with the fix (pytest's own + # per-case reports); 4.0 MB (plain cases) and 18.4 MB (every third case a + # reform) without it. + assert growth < 2_000_000, f"traced memory grew {growth / 1e6:.1f} MB" + + +def test_reform_cache_size_option_reaches_the_cache(tmp_path, monkeypatch): + seen = [] + original = test_runner._get_tax_benefit_system + + def spy(*args, **kwargs): + seen.append(kwargs.get("cache_size")) + return original(*args, **kwargs) + + monkeypatch.setattr(test_runner, "_get_tax_benefit_system", spy) + path = write_cases(tmp_path / "cases.yaml", 3) + run_with_probe(CountryTaxBenefitSystem(), path, {"reform_cache_size": 0}) + assert seen == [0, 0, 0] diff --git a/tests/fixtures/yaml_runner_memory.py b/tests/fixtures/yaml_runner_memory.py new file mode 100644 index 000000000..db010e8c9 --- /dev/null +++ b/tests/fixtures/yaml_runner_memory.py @@ -0,0 +1,92 @@ +"""Helpers for the YAML runner memory tests: generated case files and a pytest +plugin that measures what a run keeps alive.""" + +import gc +import tracemalloc +from pathlib import Path + +import pytest +import yaml + +from policyengine_core.simulations import Simulation +from policyengine_core.taxbenefitsystems import TaxBenefitSystem +from policyengine_core.tools.test_runner import OpenFiscaPlugin + +RUNNER_ARGV = ["--capture", "no", "--maxfail", "0", "--tb", "short"] + + +def write_cases(path: Path, count: int, reform_every: int = 3) -> Path: + """Write ``count`` country-template income tax cases to ``path``. + + Every ``reform_every``-th case sets the income tax rate as a dotted + parameter input with a value no other case uses, so it needs a reform + system of its own; the others run on the reform-free system. ``0`` + writes no reform cases. + """ + cases = [] + for i in range(count): + salary = 1000 + i + case = { + "name": f"case {i}", + "period": "2017-01", + "input": {"salary": salary}, + "absolute_error_margin": 0.01, + } + if reform_every and i % reform_every == 0: + rate = round(0.10 + i / 10_000, 4) + case["input"]["taxes.income_tax_rate"] = rate + case["output"] = {"income_tax": salary * rate} + else: + case["output"] = {"income_tax": salary * 0.15} + cases.append(case) + path.write_text(yaml.safe_dump(cases, sort_keys=False)) + return path + + +def live(cls) -> int: + """Number of live instances of ``cls`` after a full collection.""" + gc.collect() + return sum(1 for o in gc.get_objects() if isinstance(o, cls)) + + +class MemoryProbe: + """pytest plugin recording, after every ``sample_every``-th case, the bytes + Python has allocated (``tracemalloc``, after a full collection) and, at the + end of the session, how many simulations and tax-benefit systems are + alive.""" + + def __init__(self, sample_every: int = 10): + self.sample_every = sample_every + self.traced = [] + self.outcomes = [] + self.live_simulations = None + self.live_systems = None + + def pytest_sessionstart(self, session): + tracemalloc.start() + + def pytest_runtest_logreport(self, report): + if report.when == "call": + self.outcomes.append(report.outcome) + + def pytest_runtest_logfinish(self, nodeid, location): + if len(self.outcomes) % self.sample_every == 0: + gc.collect() + self.traced.append(tracemalloc.get_traced_memory()[0]) + + def pytest_sessionfinish(self, session, exitstatus): + tracemalloc.stop() + self.live_simulations = live(Simulation) + self.live_systems = live(TaxBenefitSystem) + + +def run_with_probe(tax_benefit_system, path: Path, options=None) -> MemoryProbe: + """Run ``path`` the way ``run_tests`` does, with a ``MemoryProbe``.""" + probe = MemoryProbe() + exit_code = pytest.main( + [*RUNNER_ARGV, "-q", "-p", "no:cacheprovider", str(path)], + plugins=[OpenFiscaPlugin(tax_benefit_system, options or {}), probe], + ) + assert exit_code == 0, f"YAML run failed with exit code {exit_code}" + assert probe.outcomes and set(probe.outcomes) == {"passed"} + return probe From de672c142029052f85fac67b53902658b20a8db6 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 08:41:45 -0400 Subject: [PATCH 2/9] Document the runner's memory behaviour and test the command-line option Co-Authored-By: Claude Opus 5.5 --- docs/engineering/skills/testing.md | 23 +++++++++++++++++++ docs/python_api/tools.md | 3 +++ policyengine_core/tools/test_runner.py | 10 ++++++++ .../tools/test_runner/test_runner_memory.py | 23 +++++++++++++++++++ 4 files changed, 59 insertions(+) diff --git a/docs/engineering/skills/testing.md b/docs/engineering/skills/testing.md index 86e214485..91967a8ab 100644 --- a/docs/engineering/skills/testing.md +++ b/docs/engineering/skills/testing.md @@ -36,3 +36,26 @@ broader behavioral change, run the relevant focused tests and formatting check. - Prefer small synthetic fixtures for regression tests. - When fixing a bug, add a regression test that fails without the fix and passes with it unless the change is documentation-only. + +## YAML suites and memory + +`policyengine-core test -c ` runs every case in one +process. Its memory is bounded by design, and tests must keep it that way: + +- A finished case releases its simulation (`YamlItem.teardown`). pytest keeps + every collected item until the session ends, so anything stored on the item + lives for the whole run. +- A case that names `reforms` or `extensions`, or sets a parameter with a + dotted input key, runs on a system built for that combination: a full copy + of the baseline system. The runner caches the `--reform-cache-size` most + recently used of them (default 2) plus the reform-free system. Do not add a + cache to the runner that grows with the number of cases or combinations. +- `tests/core/tools/test_runner/test_runner_memory.py` holds the regression + tests: a few hundred cases in one process must not grow traced memory or + the number of live simulations and systems. + +Run one large suite at a time. A country package's whole YAML tree in one +process still holds the baseline system, the reform-free copy and the cached +reform systems at once; use the country package's own batch runner where it +has one (policyengine-us: `policyengine_us/tests/test_batched.py` and the +`make test-yaml-*` targets). diff --git a/docs/python_api/tools.md b/docs/python_api/tools.md index 4c8473ab3..47fe40021 100644 --- a/docs/python_api/tools.md +++ b/docs/python_api/tools.md @@ -8,6 +8,9 @@ The `policyengine_core.tools` module contains miscellaneous utility functions, i .. autofunction:: policyengine_core.tools.test_runner.run_tests ``` +The same runner backs `policyengine-core test -c `, +where `--reform-cache-size N` sets `reform_cache_size`. + ## dump_simulation ```{eval-rst} diff --git a/policyengine_core/tools/test_runner.py b/policyengine_core/tools/test_runner.py index d34b8b7c7..357c29cbe 100644 --- a/policyengine_core/tools/test_runner.py +++ b/policyengine_core/tools/test_runner.py @@ -103,6 +103,16 @@ def run_tests(tax_benefit_system, paths, options=None): +-------------------------------+-----------+ See :any:`openfisca_test` options doc + | name_filter | ``str`` | | +-------------------------------+-----------+-------------------------------------------+ + | reform_cache_size | ``int`` | Reform systems kept cached (default 2) | + +-------------------------------+-----------+-------------------------------------------+ + + **Memory**: a case that names ``reforms`` or ``extensions``, or sets a + parameter with a dotted input key, runs on a system built for that + combination, which is a full copy of ``tax_benefit_system``. The runner + keeps the ``reform_cache_size`` most recently used of them and the + reform-free system every other case shares, and each case releases its + simulation when it finishes, so memory does not grow with the number of + cases or of distinct combinations. """ diff --git a/tests/core/tools/test_runner/test_runner_memory.py b/tests/core/tools/test_runner/test_runner_memory.py index 17d805cf2..36c4ede02 100644 --- a/tests/core/tools/test_runner/test_runner_memory.py +++ b/tests/core/tools/test_runner/test_runner_memory.py @@ -145,3 +145,26 @@ def spy(*args, **kwargs): path = write_cases(tmp_path / "cases.yaml", 3) run_with_probe(CountryTaxBenefitSystem(), path, {"reform_cache_size": 0}) assert seen == [0, 0, 0] + + +@pytest.mark.parametrize( + "arguments, expected", [([], None), (["--reform-cache-size", "3"], 3)] +) +def test_command_line_passes_reform_cache_size(monkeypatch, arguments, expected): + from policyengine_core.scripts import policyengine_command, run_test + + seen = {} + + def fake_run_tests(tax_benefit_system, paths, options): + seen.update(options) + return 0 + + monkeypatch.setattr(run_test, "run_tests", fake_run_tests) + monkeypatch.setattr(run_test, "build_tax_benefit_system", lambda *args: object()) + monkeypatch.setattr( + "sys.argv", ["policyengine-core", "test", "cases.yaml", *arguments] + ) + with pytest.raises(SystemExit) as exit_info: + run_test.main(policyengine_command.get_parser()) + assert exit_info.value.code == 0 + assert seen["reform_cache_size"] == expected From cc18d92624643899cc3e7d6d81c2885f366e25be Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 09:54:14 -0400 Subject: [PATCH 3/9] Count only the simulations and systems a probed run creates CI's full suite failed test_finished_cases_release_their_simulations with one live Simulation left by an earlier test module in the same process. Co-Authored-By: Claude Opus 5.5 --- tests/fixtures/yaml_runner_memory.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/tests/fixtures/yaml_runner_memory.py b/tests/fixtures/yaml_runner_memory.py index db010e8c9..6e92c8616 100644 --- a/tests/fixtures/yaml_runner_memory.py +++ b/tests/fixtures/yaml_runner_memory.py @@ -3,6 +3,7 @@ import gc import tracemalloc +import weakref from pathlib import Path import pytest @@ -43,17 +44,17 @@ def write_cases(path: Path, count: int, reform_every: int = 3) -> Path: return path -def live(cls) -> int: - """Number of live instances of ``cls`` after a full collection.""" +def live(cls) -> weakref.WeakSet: + """The live instances of ``cls`` after a full collection.""" gc.collect() - return sum(1 for o in gc.get_objects() if isinstance(o, cls)) + return weakref.WeakSet(o for o in gc.get_objects() if isinstance(o, cls)) class MemoryProbe: """pytest plugin recording, after every ``sample_every``-th case, the bytes Python has allocated (``tracemalloc``, after a full collection) and, at the end of the session, how many simulations and tax-benefit systems are - alive.""" + alive that the run created (not ones other tests left behind).""" def __init__(self, sample_every: int = 10): self.sample_every = sample_every @@ -63,6 +64,8 @@ def __init__(self, sample_every: int = 10): self.live_systems = None def pytest_sessionstart(self, session): + self.simulations_before = live(Simulation) + self.systems_before = live(TaxBenefitSystem) tracemalloc.start() def pytest_runtest_logreport(self, report): @@ -76,8 +79,8 @@ def pytest_runtest_logfinish(self, nodeid, location): def pytest_sessionfinish(self, session, exitstatus): tracemalloc.stop() - self.live_simulations = live(Simulation) - self.live_systems = live(TaxBenefitSystem) + self.live_simulations = len(live(Simulation) - self.simulations_before) + self.live_systems = len(live(TaxBenefitSystem) - self.systems_before) def run_with_probe(tax_benefit_system, path: Path, options=None) -> MemoryProbe: From 760a001a5a25c6b885a57649a464b4ca91639e20 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 12:12:47 -0400 Subject: [PATCH 4/9] Apply each call's reform-cache bound first, key unkeyed class reforms by class Review findings on 8607c735: - the bound was enforced only after inserting a new system, so a later call (or run on the same baseline) asking for a smaller cache, or 0, kept what a larger bound had cached: every call now trims to its own bound first; - class reforms passed without a reform_key all keyed as "" and shared one cached system (also on master): they now key by the class; the runner's inline reform classes still key by their reform_key; - reforms key as a tuple rather than a ':'-joined string; - --reform-cache-size rejects a negative value when it is parsed, rather than failing every case. Co-Authored-By: Claude Opus 5.5 --- .../scripts/policyengine_command.py | 9 ++++- policyengine_core/tools/test_runner.py | 28 ++++++++++----- .../test_runner/test_reform_cache_property.py | 20 ++++++----- .../tools/test_runner/test_runner_memory.py | 36 +++++++++++++++++++ 4 files changed, 74 insertions(+), 19 deletions(-) diff --git a/policyengine_core/scripts/policyengine_command.py b/policyengine_core/scripts/policyengine_command.py index 8112615a7..2f90aa746 100644 --- a/policyengine_core/scripts/policyengine_command.py +++ b/policyengine_core/scripts/policyengine_command.py @@ -9,6 +9,13 @@ """ +def non_negative_int(value: str) -> int: + number = int(value) + if number < 0: + raise argparse.ArgumentTypeError(f"must be 0 or more, got {number}") + return number + + def get_parser(): parser = argparse.ArgumentParser() @@ -93,7 +100,7 @@ def build_test_parser(parser): ) parser.add_argument( "--reform-cache-size", - type=int, + type=non_negative_int, default=None, help="how many reform systems (built for tests with reforms, extensions or parameter inputs) to keep cached; each is a full copy of the tax-benefit system. Defaults to 2; 0 caches none.", ) diff --git a/policyengine_core/tools/test_runner.py b/policyengine_core/tools/test_runner.py index 357c29cbe..ea32b613d 100644 --- a/policyengine_core/tools/test_runner.py +++ b/policyengine_core/tools/test_runner.py @@ -493,22 +493,32 @@ def _get_tax_benefit_system( cached_systems = _CachedSystems() _tax_benefit_system_cache[baseline] = cached_systems + # This call's bound applies to what earlier calls cached under a larger + # one, least recently used first. + reform_systems = cached_systems.reforms + while len(reform_systems) > cache_size: + reform_systems.popitem(last=False) + if not reforms and not extensions and not reform_key: if cached_systems.reform_free is None: cached_systems.reform_free = baseline.clone() return cached_systems.reform_free - # Key: (reforms in order, reform_key, extensions as a frozenset). Inline - # reform classes are rebuilt for every case, so they key by their - # ``reform_key`` (the parameter values they set), not by identity. + # Key: (reforms in order, reform_key, extensions as a frozenset). The + # runner rebuilds its inline reform class for every case, so with a + # ``reform_key`` (the parameter values it sets) a class reform keys by + # that; without one, by the class itself. inner_key = ( - ":".join([reform if isinstance(reform, str) else "" for reform in reforms]), + tuple( + reform if isinstance(reform, str) or not reform_key else "" + for reform in reforms + ), reform_key, frozenset(extensions), ) - cached = cached_systems.reforms.get(inner_key) + cached = reform_systems.get(inner_key) if cached is not None: - cached_systems.reforms.move_to_end(inner_key) + reform_systems.move_to_end(inner_key) return cached current_tax_benefit_system = baseline.clone() @@ -527,9 +537,9 @@ def _get_tax_benefit_system( current_tax_benefit_system.load_extension(extension) if cache_size > 0: - cached_systems.reforms[inner_key] = current_tax_benefit_system - while len(cached_systems.reforms) > cache_size: - cached_systems.reforms.popitem(last=False) + reform_systems[inner_key] = current_tax_benefit_system + while len(reform_systems) > cache_size: + reform_systems.popitem(last=False) return current_tax_benefit_system diff --git a/tests/core/tools/test_runner/test_reform_cache_property.py b/tests/core/tools/test_runner/test_reform_cache_property.py index ff4622c91..8146c0577 100644 --- a/tests/core/tools/test_runner/test_reform_cache_property.py +++ b/tests/core/tools/test_runner/test_reform_cache_property.py @@ -1,7 +1,7 @@ """The YAML runner's reform-system cache against a reference LRU. -For any sequence of requests and any cache size, the cache holds at most -``cache_size`` reform systems, returns a cached object exactly when a plain +For any sequence of requests, each with its own cache size, the cache holds +at most that call's ``cache_size`` reform systems after the call, returns a cached object exactly when a plain LRU of that size would, always hands back a system built for the reforms and extensions requested, and keeps one reform-free system throughout. ``test_runner_memory.py`` pins the same behaviour with examples. @@ -43,21 +43,23 @@ def load_extension(self, extension): st.lists(st.sampled_from(["a", "b", "c"]), max_size=2), st.lists(st.sampled_from(["x", "y"]), max_size=2, unique=True), st.sampled_from(["", "p:1", "p:2"]), + # Each call's bound; runs on one baseline can ask for different ones. + st.integers(min_value=0, max_value=3), ) -@hypothesis.given( - sequence=st.lists(requests, max_size=40), - cache_size=st.integers(min_value=0, max_value=3), -) +@hypothesis.given(sequence=st.lists(requests, max_size=40)) @hypothesis.settings(max_examples=300, deadline=None) -def test_cache_matches_reference_lru(sequence, cache_size): +def test_cache_matches_reference_lru(sequence): baseline = StubSystem() reference = OrderedDict() last_returned = {} reform_free = None - for reforms, extensions, reform_key in sequence: + for reforms, extensions, reform_key, cache_size in sequence: + # The reference applies each call's bound first, oldest out first. + while len(reference) > cache_size: + reference.popitem(last=False) system = _get_tax_benefit_system( baseline, list(reforms), @@ -80,7 +82,7 @@ def test_cache_matches_reference_lru(sequence, cache_size): reform_free = system assert system is reform_free else: - key = (":".join(reforms), reform_key, frozenset(extensions)) + key = (tuple(reforms), reform_key, frozenset(extensions)) if key in reference: reference.move_to_end(key) assert system is last_returned[key] diff --git a/tests/core/tools/test_runner/test_runner_memory.py b/tests/core/tools/test_runner/test_runner_memory.py index 36c4ede02..1ee4eae27 100644 --- a/tests/core/tools/test_runner/test_runner_memory.py +++ b/tests/core/tools/test_runner/test_runner_memory.py @@ -93,6 +93,31 @@ def test_cache_size_zero_builds_a_fresh_reform_system_each_time(): assert len(_tax_benefit_system_cache[baseline].reforms) == 0 +def test_a_smaller_cache_size_applies_to_what_is_already_cached(): + baseline = CountryTaxBenefitSystem() + first = _system_for_rate(baseline, 0.21, cache_size=3) + _system_for_rate(baseline, 0.22, cache_size=3) + _system_for_rate(baseline, 0.23, cache_size=3) + # A cache hit under a smaller bound still trims, least recently used first. + assert _system_for_rate(baseline, 0.23, cache_size=1) is not None + assert len(_tax_benefit_system_cache[baseline].reforms) == 1 + # Any request applies its bound, a reform-free one too; 0 empties it. + _get_tax_benefit_system(baseline, [], [], cache_size=0) + assert len(_tax_benefit_system_cache[baseline].reforms) == 0 + assert _system_for_rate(baseline, 0.21, cache_size=0) is not first + + +def test_class_reforms_without_a_key_are_cached_apart(): + baseline = CountryTaxBenefitSystem() + (reform_a,), _ = _rate_reform(0.21) + (reform_b,), _ = _rate_reform(0.22) + system_a = _get_tax_benefit_system(baseline, [reform_a], []) + system_b = _get_tax_benefit_system(baseline, [reform_b], []) + assert _rate(system_a) == pytest.approx(0.21) + assert _rate(system_b) == pytest.approx(0.22) + assert _get_tax_benefit_system(baseline, [reform_a], []) is system_a + + def test_negative_cache_size_is_rejected(): with pytest.raises(ValueError): _system_for_rate(CountryTaxBenefitSystem(), 0.25, cache_size=-1) @@ -147,6 +172,17 @@ def spy(*args, **kwargs): assert seen == [0, 0, 0] +def test_command_line_rejects_a_negative_reform_cache_size(capsys): + from policyengine_core.scripts import policyengine_command + + with pytest.raises(SystemExit) as exit_info: + policyengine_command.get_parser().parse_args( + ["test", "cases.yaml", "--reform-cache-size", "-1"] + ) + assert exit_info.value.code == 2 + assert "must be 0 or more" in capsys.readouterr().err + + @pytest.mark.parametrize( "arguments, expected", [([], None), (["--reform-cache-size", "3"], 3)] ) From 3c7d1993ae18ea1591ac714205e7905cabd6d3fb Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sat, 3 Oct 2026 08:03:40 -0400 Subject: [PATCH 5/9] Release failed-build simulations and tracers per case; key inline values by type Round-2 review of 760a001a: - a case whose simulation failed to build left it registered on the cached system: teardown now restores the system's simulation to what it was before the case, whether or not the build finished; - traced cases kept one tracer (and its results) per date: ParameterNode cached a tracer-bound wrapper per instant. It now caches the plain node and wraps it on each traced read, which also traces nodes cached before tracing started; teardown resets the cached system's trace fields when they hold the case's tracer; - inline parameter values key by repr, so 0.2 and "0.2" no longer share a system. Co-Authored-By: Claude Opus 5.5 --- changelog.d/fix-test-runner-memory.fixed.md | 2 +- .../parameters/parameter_node.py | 18 +++-- policyengine_core/tools/test_runner.py | 31 +++++--- tests/core/test_parameter_tracing_cache.py | 37 +++++++++ .../tools/test_runner/test_runner_memory.py | 77 ++++++++++++++++++- tests/fixtures/yaml_runner_memory.py | 21 ++++- 6 files changed, 163 insertions(+), 23 deletions(-) create mode 100644 tests/core/test_parameter_tracing_cache.py diff --git a/changelog.d/fix-test-runner-memory.fixed.md b/changelog.d/fix-test-runner-memory.fixed.md index d98ae6f8e..5eb0c8e61 100644 --- a/changelog.d/fix-test-runner-memory.fixed.md +++ b/changelog.d/fix-test-runner-memory.fixed.md @@ -1 +1 @@ -Stop the YAML test runner keeping every case's simulation and every reform system alive for the whole run: finished cases release their simulation, and the reform-system cache keeps the 2 most recently used (`--reform-cache-size`). +Stop the YAML test runner keeping per-case state alive for the whole run: finished cases release their simulation (also when it failed to build) and any tracer, the reform-system cache keeps the 2 most recently used (`--reform-cache-size`), and a traced parameter tree no longer caches tracer-bound wrappers. diff --git a/policyengine_core/parameters/parameter_node.py b/policyengine_core/parameters/parameter_node.py index 0df0e387b..a14c3a3c8 100644 --- a/policyengine_core/parameters/parameter_node.py +++ b/policyengine_core/parameters/parameter_node.py @@ -214,17 +214,19 @@ def clone(self) -> "ParameterNode": return clone def _get_at_instant(self, instant: Instant) -> ParameterNodeAtInstant: - if instant in self._at_instant_cache: - return self._at_instant_cache[instant] - node_at_instant = ParameterNodeAtInstant(self.name, self, instant) + node_at_instant = self._at_instant_cache.get(instant) + if node_at_instant is None: + node_at_instant = ParameterNodeAtInstant(self.name, self, instant) + self._at_instant_cache[instant] = node_at_instant if self.trace: - at_instant = TracingParameterNodeAtInstant( + # Wrapped on each read and never cached: a cached wrapper kept the + # tracer of whichever simulation first read this instant alive for + # the life of the parameters, and recorded later reads into it, + # while a node cached before tracing was turned on went untraced. + return TracingParameterNodeAtInstant( node_at_instant, self.tracer, self.branch_name ) - else: - at_instant = node_at_instant - self._at_instant_cache[instant] = at_instant - return at_instant + return node_at_instant def attach_to_parent(self, parent: "ParameterNode"): self.parent = parent diff --git a/policyengine_core/tools/test_runner.py b/policyengine_core/tools/test_runner.py index ea32b613d..f163a7a15 100644 --- a/policyengine_core/tools/test_runner.py +++ b/policyengine_core/tools/test_runner.py @@ -202,6 +202,7 @@ def __init__(self, *, baseline_tax_benefit_system, test, options, **kwargs): self.test = test self.simulation = None self.tax_benefit_system = None + self._system_simulation = None def runtest(self): self.name = self.test.get("name", "") @@ -264,8 +265,10 @@ def apply(self): self.baseline_tax_benefit_system, reforms + inline_reform, self.test.get("extensions", []), + # ``repr`` keeps 0.2 and "0.2" apart: one is a valid value, the + # other is rejected when the system is built. reform_key="=".join( - [f"{key}:{value}" for key, value in parametric_reform_items] + [f"{key}:{value!r}" for key, value in parametric_reform_items] ), cache_size=self.options.get("reform_cache_size"), ) @@ -274,6 +277,10 @@ def apply(self): performance_tables = self.options.get("performance_tables") visualize = self.options.get("visualize") + # ``Simulation.__init__`` registers itself on the system before the + # builder returns, so note what was there to restore it at teardown + # even if building fails. + self._system_simulation = getattr(self.tax_benefit_system, "simulation", None) try: builder.set_default_period(period) self.simulation = builder.build_from_dict(self.tax_benefit_system, input) @@ -313,16 +320,22 @@ def teardown(self): # finished case must not keep its simulation, or the system it ran on, # alive: over a run of thousands of cases that retained every # simulation and every reform system any case used. - system = self.tax_benefit_system - if ( - self.simulation is not None - and getattr(system, "simulation", None) is self.simulation - ): - # ``Simulation.__init__`` registers itself on its system, and that - # system is cached across cases. - system.simulation = None + system, simulation = self.tax_benefit_system, self.simulation + if system is not None: + if getattr(system, "simulation", None) is not self._system_simulation: + # A simulation built for this case, including one whose build + # failed, registered itself on the cached system. + system.simulation = self._system_simulation + tracer = getattr(simulation, "tracer", None) + parameters = getattr(system, "parameters", None) + if tracer is not None and getattr(parameters, "tracer", None) is tracer: + # A traced case marks the cached system's parameters as traced. + parameters.trace = False + parameters.tracer = None + parameters.branch_name = None self.simulation = None self.tax_benefit_system = None + self._system_simulation = None def print_computation_log(self, tracer): print("Computation log:") # noqa T001 diff --git a/tests/core/test_parameter_tracing_cache.py b/tests/core/test_parameter_tracing_cache.py new file mode 100644 index 000000000..5e7b450b6 --- /dev/null +++ b/tests/core/test_parameter_tracing_cache.py @@ -0,0 +1,37 @@ +"""A traced parameter tree wraps its cached nodes on each read and caches only +the plain nodes, so no tracer outlives its simulation and no cached node goes +untraced.""" + +from policyengine_core.parameters import ParameterNode, ParameterNodeAtInstant +from policyengine_core.tracers import FullTracer, TracingParameterNodeAtInstant + + +def tree(): + return ParameterNode("", data={"x": {"values": {"2010-01-01": 1}}}) + + +def trace(node, tracer): + node.trace, node.tracer, node.branch_name = True, tracer, "default" + + +def test_a_node_cached_before_tracing_is_traced_once_tracing_starts(): + parameters = tree() + plain = parameters.get_at_instant("2017-01-01") + trace(parameters, FullTracer()) + traced = parameters.get_at_instant("2017-01-01") + assert isinstance(traced, TracingParameterNodeAtInstant) + assert traced.parameter_node_at_instant is plain + + +def test_the_cache_never_holds_a_tracer(): + parameters = tree() + for tracer in (FullTracer(), FullTracer()): + trace(parameters, tracer) + assert parameters.get_at_instant("2017-01-01").tracer is tracer + assert parameters.get_at_instant("2018-01-01").tracer is tracer + assert all( + type(node) is ParameterNodeAtInstant + for node in parameters._at_instant_cache.values() + ) + parameters.trace, parameters.tracer = False, None + assert type(parameters.get_at_instant("2017-01-01")) is ParameterNodeAtInstant diff --git a/tests/core/tools/test_runner/test_runner_memory.py b/tests/core/tools/test_runner/test_runner_memory.py index 1ee4eae27..8b6ceffc9 100644 --- a/tests/core/tools/test_runner/test_runner_memory.py +++ b/tests/core/tools/test_runner/test_runner_memory.py @@ -10,6 +10,7 @@ import weakref import pytest +import yaml from policyengine_core.country_template import CountryTaxBenefitSystem from policyengine_core.reforms import Reform, set_parameter @@ -19,7 +20,8 @@ _get_tax_benefit_system, _tax_benefit_system_cache, ) -from tests.fixtures.yaml_runner_memory import run_with_probe, write_cases +from policyengine_core.tools.test_runner import OpenFiscaPlugin +from tests.fixtures.yaml_runner_memory import RUNNER_ARGV, run_with_probe, write_cases def _rate_reform(rate): @@ -143,6 +145,79 @@ def test_live_systems_do_not_grow_with_the_number_of_cases(tmp_path): assert long.live_systems == short.live_systems +def test_traced_cases_release_their_tracers(tmp_path): + """Verbose (traced) cases at different dates keep no tracer: the cached + system's parameters used to cache one tracing wrapper, holding that case's + tracer and results, per date.""" + cases = [ + { + "name": f"case {i}", + "period": f"20{15 + i // 12}-{i % 12 + 1:02d}", + "input": {"salary": 1000}, + "output": {"income_tax": 150}, + } + for i in range(16) + ] + path = tmp_path / "traced.yaml" + path.write_text(yaml.safe_dump(cases)) + system = CountryTaxBenefitSystem() + probe = run_with_probe(system, path, {"verbose": True}) + assert probe.live_tracers == 0 + assert probe.live_simulations == 0 + reform_free = _tax_benefit_system_cache[system].reform_free + assert reform_free.parameters.trace is False + assert reform_free.parameters.tracer is None + + +def test_cases_whose_simulation_fails_to_build_release_it(tmp_path): + cases = [ + { + "name": f"bad input {i}", + "period": "2017-01", + "input": {"salary": "not a number"}, + "output": {"income_tax": 0}, + } + for i in range(8) + ] + path = tmp_path / "bad.yaml" + path.write_text(yaml.safe_dump(cases)) + system = CountryTaxBenefitSystem() + probe = run_with_probe(system, path, outcome="failed") + assert probe.live_simulations == 0 + assert ( + getattr(_tax_benefit_system_cache[system].reform_free, "simulation", None) + is None + ) + + +def test_inline_parameter_values_key_by_type(tmp_path): + """A number and the same digits as a string are different inputs: the + string must not reuse the system built for the number.""" + base = {"period": "2017-01", "input": {"salary": 1000}} + cases = [ + { + **base, + "name": "number", + "input": {**base["input"], "taxes.income_tax_rate": 0.2}, + "output": {"income_tax": 200}, + }, + { + **base, + "name": "string", + "input": {**base["input"], "taxes.income_tax_rate": "0.2"}, + "output": {"income_tax": 200}, + }, + ] + path = tmp_path / "typed.yaml" + path.write_text(yaml.safe_dump(cases)) + exit_code = pytest.main( + [*RUNNER_ARGV, "-q", "-p", "no:cacheprovider", str(path)], + plugins=[OpenFiscaPlugin(CountryTaxBenefitSystem(), {})], + ) + # With a value-only key the string case passed on the number's system. + assert exit_code == 1 + + @pytest.mark.parametrize("reform_every", [0, 3]) def test_traced_memory_stays_flat_over_hundreds_of_cases(tmp_path, reform_every): """A few hundred cases in one process: memory after the last case is diff --git a/tests/fixtures/yaml_runner_memory.py b/tests/fixtures/yaml_runner_memory.py index 6e92c8616..8f14dea16 100644 --- a/tests/fixtures/yaml_runner_memory.py +++ b/tests/fixtures/yaml_runner_memory.py @@ -2,6 +2,7 @@ plugin that measures what a run keeps alive.""" import gc +import sys import tracemalloc import weakref from pathlib import Path @@ -11,6 +12,7 @@ from policyengine_core.simulations import Simulation from policyengine_core.taxbenefitsystems import TaxBenefitSystem +from policyengine_core.tracers import FullTracer from policyengine_core.tools.test_runner import OpenFiscaPlugin RUNNER_ARGV = ["--capture", "no", "--maxfail", "0", "--tb", "short"] @@ -62,10 +64,12 @@ def __init__(self, sample_every: int = 10): self.outcomes = [] self.live_simulations = None self.live_systems = None + self.live_tracers = None def pytest_sessionstart(self, session): self.simulations_before = live(Simulation) self.systems_before = live(TaxBenefitSystem) + self.tracers_before = live(FullTracer) tracemalloc.start() def pytest_runtest_logreport(self, report): @@ -79,17 +83,26 @@ def pytest_runtest_logfinish(self, nodeid, location): def pytest_sessionfinish(self, session, exitstatus): tracemalloc.stop() + # pytest keeps the last failure's traceback in ``sys.last_*``; that one + # root is pytest's, not the runner's. + for name in ("last_type", "last_value", "last_traceback", "last_exc"): + if hasattr(sys, name): + delattr(sys, name) self.live_simulations = len(live(Simulation) - self.simulations_before) self.live_systems = len(live(TaxBenefitSystem) - self.systems_before) + self.live_tracers = len(live(FullTracer) - self.tracers_before) -def run_with_probe(tax_benefit_system, path: Path, options=None) -> MemoryProbe: - """Run ``path`` the way ``run_tests`` does, with a ``MemoryProbe``.""" +def run_with_probe( + tax_benefit_system, path: Path, options=None, outcome="passed" +) -> MemoryProbe: + """Run ``path`` the way ``run_tests`` does, with a ``MemoryProbe``; every + case must end with ``outcome``.""" probe = MemoryProbe() exit_code = pytest.main( [*RUNNER_ARGV, "-q", "-p", "no:cacheprovider", str(path)], plugins=[OpenFiscaPlugin(tax_benefit_system, options or {}), probe], ) - assert exit_code == 0, f"YAML run failed with exit code {exit_code}" - assert probe.outcomes and set(probe.outcomes) == {"passed"} + assert exit_code == (0 if outcome == "passed" else 1), exit_code + assert probe.outcomes and set(probe.outcomes) == {outcome} return probe From 206b619a6c0d572e20039f867508ff402a256b19 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Tue, 6 Oct 2026 09:27:04 -0400 Subject: [PATCH 6/9] Keep tracers out of the system's parameter-instant cache TaxBenefitSystem.get_parameters_at_instant cached whatever the parameter tree returned. While the tree is marked traced that is a tracing wrapper holding the reading simulation's tracer, so the system kept one tracer, and the results it recorded, per date; served a node cached before tracing started untraced; and recorded later reads, after tracing stopped or under a later tracer, in the first one. Since #574 core's simulations trace through a per-call view and never mark the shared tree, so the round-3 scenario (traced YAML cases at distinct dates whose formula calls get_parameters_at_instant) no longer leaks; its test fails on 3c7d1993 and passes on the merge. Code outside core still marks the tree (policyengine-us's isolate_parameter_tracing sets trace, tracer and branch_name on the root), and then the system cache kept 16 tracers after 16 traced cases at distinct dates. The system cache now caches the plain node and, while the tree is traced, wraps it for that read only, as ParameterNode._get_at_instant does. Tests: example tests across 16 distinct dates, a Hypothesis differential property of the system's path against the tree's, and two YAML runs at 16 distinct dates with weak references to tracers and recorded results. The docs now say each system still caches its parameters at every distinct instant read. Co-Authored-By: Claude Opus 5.5 --- changelog.d/fix-test-runner-memory.fixed.md | 2 +- docs/engineering/skills/testing.md | 9 +- .../taxbenefitsystems/tax_benefit_system.py | 18 ++- policyengine_core/tools/test_runner.py | 10 +- tests/core/test_parameter_tracing_cache.py | 78 +++++++++- .../test_parameter_tracing_cache_property.py | 147 ++++++++++++++++++ .../tools/test_runner/test_runner_memory.py | 55 ++++++- tests/fixtures/yaml_runner_memory.py | 78 ++++++++++ 8 files changed, 387 insertions(+), 10 deletions(-) create mode 100644 tests/core/test_parameter_tracing_cache_property.py diff --git a/changelog.d/fix-test-runner-memory.fixed.md b/changelog.d/fix-test-runner-memory.fixed.md index 5eb0c8e61..f17b1a2db 100644 --- a/changelog.d/fix-test-runner-memory.fixed.md +++ b/changelog.d/fix-test-runner-memory.fixed.md @@ -1 +1 @@ -Stop the YAML test runner keeping per-case state alive for the whole run: finished cases release their simulation (also when it failed to build) and any tracer, the reform-system cache keeps the 2 most recently used (`--reform-cache-size`), and a traced parameter tree no longer caches tracer-bound wrappers. +Stop the YAML test runner keeping per-case state alive for the whole run: finished cases release their simulation (also when it failed to build) and any tracer, the reform-system cache keeps the 2 most recently used (`--reform-cache-size`), and neither a traced parameter tree nor the tax-benefit system's cache of it at each instant (`get_parameters_at_instant`) keeps tracer-bound wrappers. diff --git a/docs/engineering/skills/testing.md b/docs/engineering/skills/testing.md index 91967a8ab..ecf1ef527 100644 --- a/docs/engineering/skills/testing.md +++ b/docs/engineering/skills/testing.md @@ -50,9 +50,16 @@ process. Its memory is bounded by design, and tests must keep it that way: of the baseline system. The runner caches the `--reform-cache-size` most recently used of them (default 2) plus the reform-free system. Do not add a cache to the runner that grows with the number of cases or combinations. +- Each system still caches its parameter tree at every distinct instant its + cases read (`ParameterNode.get_at_instant`, + `TaxBenefitSystem.get_parameters_at_instant`), so memory grows with the + number of distinct dates a suite reads, not with the number of cases. + Neither cache may hold a tracer: a read of a traced tree is wrapped for that + read only. - `tests/core/tools/test_runner/test_runner_memory.py` holds the regression tests: a few hundred cases in one process must not grow traced memory or - the number of live simulations and systems. + the number of live simulations and systems, and traced cases at distinct + dates must leave no tracer or recorded result alive. Run one large suite at a time. A country package's whole YAML tree in one process still holds the baseline system, the reform-free copy and the cached diff --git a/policyengine_core/taxbenefitsystems/tax_benefit_system.py b/policyengine_core/taxbenefitsystems/tax_benefit_system.py index b7073081d..79e7d94ca 100644 --- a/policyengine_core/taxbenefitsystems/tax_benefit_system.py +++ b/policyengine_core/taxbenefitsystems/tax_benefit_system.py @@ -50,6 +50,7 @@ ) from policyengine_core.periods import Instant, Period from policyengine_core.populations import GroupPopulation, Population +from policyengine_core.tracers import TracingParameterNodeAtInstant from policyengine_core.variables import Variable log = logging.getLogger(__name__) @@ -559,10 +560,23 @@ def get_parameters_at_instant(self, instant: Instant) -> ParameterNodeAtInstant: ) ) + parameters = self.parameters parameters_at_instant = self._parameters_at_instant_cache.get(instant) - if parameters_at_instant is None and self.parameters is not None: - parameters_at_instant = self.parameters.get_at_instant(str(instant)) + if parameters_at_instant is None and parameters is not None: + parameters_at_instant = parameters.get_at_instant(str(instant)) + if isinstance(parameters_at_instant, TracingParameterNodeAtInstant): + parameters_at_instant = parameters_at_instant.parameter_node_at_instant self._parameters_at_instant_cache[instant] = parameters_at_instant + if parameters_at_instant is not None and getattr(parameters, "trace", False): + # Wrapped on each read and never cached, as the tree's own reads + # are (``ParameterNode._get_at_instant``): a cached wrapper kept + # the tracer of whichever simulation first read this instant + # alive for the life of the system, and recorded later reads into + # it, while a node cached before tracing was turned on went + # untraced. + return TracingParameterNodeAtInstant( + parameters_at_instant, parameters.tracer, parameters.branch_name + ) return parameters_at_instant def get_package_metadata(self) -> dict: diff --git a/policyengine_core/tools/test_runner.py b/policyengine_core/tools/test_runner.py index f163a7a15..c8df6c8cf 100644 --- a/policyengine_core/tools/test_runner.py +++ b/policyengine_core/tools/test_runner.py @@ -112,7 +112,8 @@ def run_tests(tax_benefit_system, paths, options=None): keeps the ``reform_cache_size`` most recently used of them and the reform-free system every other case shares, and each case releases its simulation when it finishes, so memory does not grow with the number of - cases or of distinct combinations. + cases or of distinct combinations. Each system still caches its + parameters at every distinct instant the cases read. """ @@ -329,7 +330,12 @@ def teardown(self): tracer = getattr(simulation, "tracer", None) parameters = getattr(system, "parameters", None) if tracer is not None and getattr(parameters, "tracer", None) is tracer: - # A traced case marks the cached system's parameters as traced. + # Core's simulations trace through a per-call view and leave + # the tree unmarked, but code outside core can mark the cached + # system's tree with this case's tracer (policyengine-us's + # ``isolate_parameter_tracing`` does). Left marked, the tree + # would keep this case's tracer alive, and later untraced + # cases would record their parameter reads in it. parameters.trace = False parameters.tracer = None parameters.branch_name = None diff --git a/tests/core/test_parameter_tracing_cache.py b/tests/core/test_parameter_tracing_cache.py index 5e7b450b6..5b6d73b0e 100644 --- a/tests/core/test_parameter_tracing_cache.py +++ b/tests/core/test_parameter_tracing_cache.py @@ -1,17 +1,35 @@ """A traced parameter tree wraps its cached nodes on each read and caches only the plain nodes, so no tracer outlives its simulation and no cached node goes -untraced.""" +untraced. The tax-benefit system's own cache of the tree at each instant +(``TaxBenefitSystem.get_parameters_at_instant``) follows the same rule.""" +import gc +import weakref + +from policyengine_core.country_template.entities import entities from policyengine_core.parameters import ParameterNode, ParameterNodeAtInstant +from policyengine_core.taxbenefitsystems import TaxBenefitSystem from policyengine_core.tracers import FullTracer, TracingParameterNodeAtInstant +DATES = [f"{year}-01-01" for year in range(2011, 2027)] + def tree(): return ParameterNode("", data={"x": {"values": {"2010-01-01": 1}}}) -def trace(node, tracer): - node.trace, node.tracer, node.branch_name = True, tracer, "default" +def system(): + tax_benefit_system = TaxBenefitSystem(entities) + tax_benefit_system.parameters = tree() + return tax_benefit_system + + +def trace(node, tracer, branch_name="default"): + node.trace, node.tracer, node.branch_name = True, tracer, branch_name + + +def untrace(node): + node.trace, node.tracer, node.branch_name = False, None, None def test_a_node_cached_before_tracing_is_traced_once_tracing_starts(): @@ -35,3 +53,57 @@ def test_the_cache_never_holds_a_tracer(): ) parameters.trace, parameters.tracer = False, None assert type(parameters.get_at_instant("2017-01-01")) is ParameterNodeAtInstant + + +def test_the_system_cache_never_holds_a_tracer(): + """Reads at distinct dates, each under a tracer of its own, as traced + cases on one cached system make them: the system used to cache one + tracing wrapper per date, each keeping its tracer alive.""" + tax_benefit_system = system() + refs = [] + for date in DATES: + tracer = FullTracer() + refs.append(weakref.ref(tracer)) + trace(tax_benefit_system.parameters, tracer) + assert tax_benefit_system.get_parameters_at_instant(date).tracer is tracer + del tracer + untrace(tax_benefit_system.parameters) + gc.collect() + + cached = tax_benefit_system._parameters_at_instant_cache + assert len(cached) == len(DATES) + assert all(type(node) is ParameterNodeAtInstant for node in cached.values()) + assert [ref for ref in refs if ref() is not None] == [] + + +def test_a_system_node_cached_before_tracing_is_traced_once_tracing_starts(): + tax_benefit_system = system() + plain = tax_benefit_system.get_parameters_at_instant("2017-01-01") + tracer = FullTracer() + trace(tax_benefit_system.parameters, tracer, "branch") + traced = tax_benefit_system.get_parameters_at_instant("2017-01-01") + assert isinstance(traced, TracingParameterNodeAtInstant) + assert traced.parameter_node_at_instant is plain + assert (traced.tracer, traced.branch_name) == (tracer, "branch") + untrace(tax_benefit_system.parameters) + assert tax_benefit_system.get_parameters_at_instant("2017-01-01") is plain + + +def test_each_system_read_records_in_the_tracer_current_at_the_read(): + """A read through the system records in the tree's tracer at the time of + the read, and nowhere once tracing stops: it used to record every later + read in the tracer current when the instant was first read.""" + tax_benefit_system = system() + first, second = FullTracer(), FullTracer() + for tracer in (first, second): + tracer.record_calculation_start("y", "2017") + + trace(tax_benefit_system.parameters, first) + assert tax_benefit_system.get_parameters_at_instant("2017-01-01").x == 1 + trace(tax_benefit_system.parameters, second) + assert tax_benefit_system.get_parameters_at_instant("2017-01-01").x == 1 + untrace(tax_benefit_system.parameters) + assert tax_benefit_system.get_parameters_at_instant("2017-01-01").x == 1 + + assert len(first.trees[0].parameters) == 1 + assert len(second.trees[0].parameters) == 1 diff --git a/tests/core/test_parameter_tracing_cache_property.py b/tests/core/test_parameter_tracing_cache_property.py new file mode 100644 index 000000000..32918fb54 --- /dev/null +++ b/tests/core/test_parameter_tracing_cache_property.py @@ -0,0 +1,147 @@ +"""The parameter-instant caches under tracing, for any sequence of tracing +switches and reads at any dates. + +Two paths read the parameter tree at an instant: the tree's own +(``ParameterNode.get_at_instant``, cached in ``_at_instant_cache``) and the +tax-benefit system's (``TaxBenefitSystem.get_parameters_at_instant``, cached +in ``_parameters_at_instant_cache``). Invariants checked after every step: + +- both caches hold plain ``ParameterNodeAtInstant`` objects only, so neither + keeps a tracer; +- differential: both paths return the same plain node for an instant, traced + exactly when the tree is traced at the time of the read, with the tree's + tracer and branch name at that time; +- a traced read records in that tracer and in no other, and tracing never + changes a value read; +- once tracing stops and the caller drops its tracers, none is alive. + +``test_parameter_tracing_cache.py`` pins the same behaviour with examples. +""" + +from __future__ import annotations + +import gc +import weakref + +import pytest + +# The country-package smoke job installs no dev dependencies. +hypothesis = pytest.importorskip("hypothesis") +st = hypothesis.strategies + +from policyengine_core.country_template.entities import entities +from policyengine_core.parameters import ParameterNode, ParameterNodeAtInstant +from policyengine_core.taxbenefitsystems import TaxBenefitSystem +from policyengine_core.tracers import FullTracer, TracingParameterNodeAtInstant + +# The parameter's value from each start date on: the reference the reads are +# checked against. +X_VALUES = {"2010-01-01": 1, "2015-01-01": 2, "2020-01-01": 3} +DATES = ["2012-06-01", "2015-01-01", "2017-01-01", "2020-01-01", "2024-12-31"] +TRACERS = 3 + + +def expected_x(date: str) -> int: + return X_VALUES[max(start for start in X_VALUES if start <= date)] + + +def system() -> TaxBenefitSystem: + tax_benefit_system = TaxBenefitSystem(entities) + tax_benefit_system.parameters = ParameterNode( + "", data={"x": {"values": dict(X_VALUES)}} + ) + return tax_benefit_system + + +steps = st.lists( + st.one_of( + st.tuples( + st.just("trace"), + st.integers(min_value=0, max_value=TRACERS - 1), + st.sampled_from(["default", "branch"]), + ), + st.tuples(st.just("untrace")), + st.tuples(st.sampled_from(["system", "tree"]), st.sampled_from(DATES)), + ), + max_size=30, +) + + +def check_caches(tax_benefit_system: TaxBenefitSystem) -> None: + for cache in ( + tax_benefit_system._parameters_at_instant_cache, + tax_benefit_system.parameters._at_instant_cache, + ): + assert all(type(node) is ParameterNodeAtInstant for node in cache.values()) + + +def plain(node): + if isinstance(node, TracingParameterNodeAtInstant): + return node.parameter_node_at_instant + return node + + +@hypothesis.given(sequence=steps) +# ``too_slow`` times input generation by the wall clock, so it fails on a busy +# host, not on anything this property checks. +@hypothesis.settings( + max_examples=300, + deadline=None, + suppress_health_check=[hypothesis.HealthCheck.too_slow], +) +def test_caches_stay_tracer_free_and_both_paths_agree(sequence): + tax_benefit_system = system() + parameters = tax_benefit_system.parameters + tracers = [FullTracer() for _ in range(TRACERS)] + for tracer in tracers: + tracer.record_calculation_start("y", "2017") + refs = [weakref.ref(tracer) for tracer in tracers] + current = None # (tracer index, branch name) while the tree is traced + reads = [0] * TRACERS + + for step in sequence: + if step[0] == "trace": + _, index, branch_name = step + parameters.trace = True + parameters.tracer = tracers[index] + parameters.branch_name = branch_name + current = (index, branch_name) + elif step[0] == "untrace": + parameters.trace, parameters.tracer, parameters.branch_name = ( + False, + None, + None, + ) + current = None + else: + path, date = step + read = ( + tax_benefit_system.get_parameters_at_instant(date) + if path == "system" + else parameters.get_at_instant(date) + ) + other = ( + parameters.get_at_instant(date) + if path == "system" + else tax_benefit_system.get_parameters_at_instant(date) + ) + for node in (read, other): + if current is None: + assert type(node) is ParameterNodeAtInstant + else: + index, branch_name = current + assert isinstance(node, TracingParameterNodeAtInstant) + assert node.tracer is tracers[index] + assert node.branch_name == branch_name + assert plain(read) is plain(other) + assert read.x == expected_x(date) + if current is not None: + reads[current[0]] += 1 + del read, other, node + check_caches(tax_benefit_system) + assert [len(tracer.trees[0].parameters) for tracer in tracers] == reads + + parameters.trace, parameters.tracer, parameters.branch_name = False, None, None + del tracers, tracer + gc.collect() + assert [ref for ref in refs if ref() is not None] == [] diff --git a/tests/core/tools/test_runner/test_runner_memory.py b/tests/core/tools/test_runner/test_runner_memory.py index 8b6ceffc9..b232f2794 100644 --- a/tests/core/tools/test_runner/test_runner_memory.py +++ b/tests/core/tools/test_runner/test_runner_memory.py @@ -21,7 +21,15 @@ _tax_benefit_system_cache, ) from policyengine_core.tools.test_runner import OpenFiscaPlugin -from tests.fixtures.yaml_runner_memory import RUNNER_ARGV, run_with_probe, write_cases +from tests.fixtures.yaml_runner_memory import ( + RUNNER_ARGV, + mark_parameter_trees_traced, + record_results, + run_with_probe, + system_reading_its_parameters, + write_cases, + write_cases_at_distinct_dates, +) def _rate_reform(rate): @@ -169,6 +177,51 @@ def test_traced_cases_release_their_tracers(tmp_path): assert reform_free.parameters.tracer is None +def _run_cases_reading_the_systems_parameters(tmp_path, monkeypatch): + """Run 16 verbose cases at distinct dates whose formula reads + ``TaxBenefitSystem.get_parameters_at_instant``; return the probe, the + weak references to every result their tracers recorded, and the system + they ran on.""" + results = record_results(monkeypatch) + system = system_reading_its_parameters() + path = write_cases_at_distinct_dates(tmp_path / "traced.yaml", 16) + probe = run_with_probe(system, path, {"verbose": True}) + assert len(probe.outcomes) == 16 + assert results + return probe, results, _tax_benefit_system_cache[system].reform_free + + +def test_traced_cases_reading_the_systems_parameters_release_their_tracers( + tmp_path, monkeypatch +): + """The system caches the parameter tree at each instant it is asked for. + While core's simulations marked the tree traced, that cache kept a + tracing wrapper, with its case's tracer and results, for every date (16 + tracers and 32 result arrays after 16 dates).""" + probe, results, reform_free = _run_cases_reading_the_systems_parameters( + tmp_path, monkeypatch + ) + assert probe.live_tracers == 0 + assert [ref for ref in results if ref() is not None] == [] + assert probe.live_simulations == 0 + assert len(reform_free._parameters_at_instant_cache) == 16 + + +def test_cases_on_a_tree_marked_traced_release_their_tracers(tmp_path, monkeypatch): + """The same cases, with each simulation marking its system's parameter + tree traced, as code outside core can: the system's cache must still + keep only plain nodes, and teardown unmarks the tree.""" + mark_parameter_trees_traced(monkeypatch) + probe, results, reform_free = _run_cases_reading_the_systems_parameters( + tmp_path, monkeypatch + ) + assert probe.live_tracers == 0 + assert [ref for ref in results if ref() is not None] == [] + assert probe.live_simulations == 0 + assert reform_free.parameters.trace is False + assert reform_free.parameters.tracer is None + + def test_cases_whose_simulation_fails_to_build_release_it(tmp_path): cases = [ { diff --git a/tests/fixtures/yaml_runner_memory.py b/tests/fixtures/yaml_runner_memory.py index 8f14dea16..b541b126e 100644 --- a/tests/fixtures/yaml_runner_memory.py +++ b/tests/fixtures/yaml_runner_memory.py @@ -7,9 +7,14 @@ import weakref from pathlib import Path +import numpy as np import pytest import yaml +from policyengine_core.country_template import CountryTaxBenefitSystem +from policyengine_core.country_template.variables.taxes import ( + income_tax as template_income_tax, +) from policyengine_core.simulations import Simulation from policyengine_core.taxbenefitsystems import TaxBenefitSystem from policyengine_core.tracers import FullTracer @@ -46,6 +51,79 @@ def write_cases(path: Path, count: int, reform_every: int = 3) -> Path: return path +def write_cases_at_distinct_dates(path: Path, count: int, people: int = 256) -> Path: + """Write ``count`` income tax cases for ``people`` people each, every one + in a month no other case uses, so each reads the parameters at an instant + of its own.""" + cases = [ + { + "name": f"case {i}", + "period": f"{2017 + i // 12}-{1 + i % 12:02d}", + "input": {"salary": [1000] * people}, + "output": {"income_tax": [150] * people}, + } + for i in range(count) + ] + path.write_text(yaml.safe_dump(cases)) + return path + + +def system_reading_its_parameters() -> CountryTaxBenefitSystem: + """A country-template system whose income tax formula reads the rate + through the public ``TaxBenefitSystem.get_parameters_at_instant``, not + through the formula's ``parameters`` argument.""" + system = CountryTaxBenefitSystem() + + class income_tax(template_income_tax): + def formula(person, period, parameters): + system_parameters = ( + person.simulation.tax_benefit_system.get_parameters_at_instant( + period.start + ) + ) + return person("salary", period) * system_parameters.taxes.income_tax_rate + + system.update_variable(income_tax) + return system + + +def mark_parameter_trees_traced(monkeypatch) -> None: + """Have every simulation that switches tracing on mark its system's + parameter tree traced with its own tracer and branch name. + + Core's simulations trace per call and leave the tree untraced; code + outside core still marks it (policyengine-us's + ``isolate_parameter_tracing`` marks the root of a tree no one else reads + where it stands). + """ + original = Simulation.trace + + def set_trace(simulation, trace): + original.fset(simulation, trace) + if trace: + parameters = simulation.tax_benefit_system.parameters + parameters.trace = True + parameters.tracer = simulation.tracer + parameters.branch_name = simulation.branch_name + + monkeypatch.setattr(Simulation, "trace", property(original.fget, set_trace)) + + +def record_results(monkeypatch) -> list: + """Weak references to every array a ``FullTracer`` records as a + calculation result from now on.""" + refs = [] + original = FullTracer.record_calculation_result + + def record(tracer, value): + if isinstance(value, np.ndarray): + refs.append(weakref.ref(value)) + original(tracer, value) + + monkeypatch.setattr(FullTracer, "record_calculation_result", record) + return refs + + def live(cls) -> weakref.WeakSet: """The live instances of ``cls`` after a full collection.""" gc.collect() From 0e27751f5b6814e867a939697324cc8e35ec63ca Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Tue, 6 Oct 2026 09:27:21 -0400 Subject: [PATCH 7/9] Don't fail the reform-cache property on a slow host Hypothesis's too_slow health check times input generation by the wall clock. Round 3 saw test_cache_matches_reference_lru fail it under host load (53 s of generation) and then pass in isolation; it checks nothing the property is about. Co-Authored-By: Claude Opus 5.5 --- .../core/tools/test_runner/test_reform_cache_property.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/tests/core/tools/test_runner/test_reform_cache_property.py b/tests/core/tools/test_runner/test_reform_cache_property.py index 8146c0577..7abde339b 100644 --- a/tests/core/tools/test_runner/test_reform_cache_property.py +++ b/tests/core/tools/test_runner/test_reform_cache_property.py @@ -49,7 +49,13 @@ def load_extension(self, extension): @hypothesis.given(sequence=st.lists(requests, max_size=40)) -@hypothesis.settings(max_examples=300, deadline=None) +# ``too_slow`` times input generation by the wall clock, so it fails on a busy +# host, not on anything this property checks. +@hypothesis.settings( + max_examples=300, + deadline=None, + suppress_health_check=[hypothesis.HealthCheck.too_slow], +) def test_cache_matches_reference_lru(sequence): baseline = StubSystem() reference = OrderedDict() From 6b8ec162dcc40b5bbb4969552a07bf57ad62175b Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Wed, 7 Oct 2026 17:05:55 -0400 Subject: [PATCH 8/9] Test deepcopy with tracer-free parameter instant caches --- .../core/test_tracing_parameter_isolation.py | 31 +++++++++++++------ 1 file changed, 21 insertions(+), 10 deletions(-) diff --git a/tests/core/test_tracing_parameter_isolation.py b/tests/core/test_tracing_parameter_isolation.py index 8c4f5e95a..2660fe9ab 100644 --- a/tests/core/test_tracing_parameter_isolation.py +++ b/tests/core/test_tracing_parameter_isolation.py @@ -358,23 +358,34 @@ def test_tracing_parameter_node_can_be_copied(isolated_tax_benefit_system, dupli ) -def test_parameters_holding_traced_nodes_can_be_deep_copied( +def test_traced_parameters_can_be_deep_copied_without_caching_wrappers( isolated_tax_benefit_system, ): - """A tree whose ``trace`` flag is set caches traced nodes; copying it works.""" + """A traced tree caches plain nodes and returns a wrapper for each read.""" parameters = isolated_tax_benefit_system.parameters parameters.trace = True parameters.tracer = FullTracer() - parameters(INSTANT) - assert type(parameters._at_instant_cache[INSTANT]) is ( - TracingParameterNodeAtInstant - ) + traced = parameters(INSTANT) + cached = parameters._at_instant_cache[INSTANT] + assert type(cached) is ParameterNodeAtInstant + assert type(traced) is TracingParameterNodeAtInstant + assert traced.parameter_node_at_instant is cached + assert traced.tracer is parameters.tracer + assert parameters(INSTANT) is not traced duplicated = copy.deepcopy(parameters) - - assert duplicated(INSTANT).taxes.income_tax_rate == ( - parameters(INSTANT).taxes.income_tax_rate - ) + copied_read = duplicated(INSTANT) + + assert type(copied_read) is TracingParameterNodeAtInstant + assert copied_read.tracer is duplicated.tracer + assert copied_read.tracer is not parameters.tracer + assert copied_read.parameter_node_at_instant is not cached + assert copied_read.taxes.income_tax_rate == traced.taxes.income_tax_rate + for tree in (parameters, duplicated): + assert all( + type(node) is ParameterNodeAtInstant + for node in tree._at_instant_cache.values() + ) class _SlottedParameterNode(ParameterNode): From 91744d3e6c88b1a56f7e81f1001f8cdeba3c9af8 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Wed, 7 Oct 2026 17:07:34 -0400 Subject: [PATCH 9/9] Isolate runner memory checks from rerun plugin tracebacks --- changelog.d/fix-test-runner-memory.fixed.md | 2 +- .../tools/test_runner/test_runner_memory.py | 50 ++++++++++++++++--- tests/fixtures/yaml_runner_memory.py | 19 ++++++- 3 files changed, 63 insertions(+), 8 deletions(-) diff --git a/changelog.d/fix-test-runner-memory.fixed.md b/changelog.d/fix-test-runner-memory.fixed.md index f17b1a2db..76101c289 100644 --- a/changelog.d/fix-test-runner-memory.fixed.md +++ b/changelog.d/fix-test-runner-memory.fixed.md @@ -1 +1 @@ -Stop the YAML test runner keeping per-case state alive for the whole run: finished cases release their simulation (also when it failed to build) and any tracer, the reform-system cache keeps the 2 most recently used (`--reform-cache-size`), and neither a traced parameter tree nor the tax-benefit system's cache of it at each instant (`get_parameters_at_instant`) keeps tracer-bound wrappers. +Stop the YAML test runner keeping per-case state alive for the whole run: finished cases drop runner-owned simulation references (also when the simulation failed to build) and any tracer, the reform-system cache keeps the 2 most recently used (`--reform-cache-size`), and neither a traced parameter tree nor the tax-benefit system's cache of it at each instant (`get_parameters_at_instant`) keeps tracer-bound wrappers. diff --git a/tests/core/tools/test_runner/test_runner_memory.py b/tests/core/tools/test_runner/test_runner_memory.py index b232f2794..fd4795505 100644 --- a/tests/core/tools/test_runner/test_runner_memory.py +++ b/tests/core/tools/test_runner/test_runner_memory.py @@ -30,6 +30,7 @@ write_cases, write_cases_at_distinct_dates, ) +from tests.fixtures.tracing import build_simulation def _rate_reform(rate): @@ -222,12 +223,27 @@ def test_cases_on_a_tree_marked_traced_release_their_tracers(tmp_path, monkeypat assert reform_free.parameters.tracer is None -def test_cases_whose_simulation_fails_to_build_release_it(tmp_path): +@pytest.mark.parametrize( + "invalid_input", + [ + pytest.param({"salary": "not a number"}, id="variable-value"), + pytest.param( + {"persons": {"person": {"salary": {"2017-01": "not a number"}}}}, + id="entity-value", + ), + pytest.param({"unknown_variable": 1}, id="unknown-variable"), + ], +) +@pytest.mark.parametrize("existing_simulation", [False, True]) +def test_cases_whose_simulation_fails_to_build_release_it( + tmp_path, invalid_input, existing_simulation +): + """Both builder paths release failed cases and restore any prior backlink.""" cases = [ { "name": f"bad input {i}", "period": "2017-01", - "input": {"salary": "not a number"}, + "input": invalid_input, "output": {"income_tax": 0}, } for i in range(8) @@ -235,12 +251,34 @@ def test_cases_whose_simulation_fails_to_build_release_it(tmp_path): path = tmp_path / "bad.yaml" path.write_text(yaml.safe_dump(cases)) system = CountryTaxBenefitSystem() + reform_free = _get_tax_benefit_system(system, [], []) + previous = build_simulation(reform_free) if existing_simulation else None probe = run_with_probe(system, path, outcome="failed") + assert len(probe.outcomes) == len(probe.items) == 8 assert probe.live_simulations == 0 - assert ( - getattr(_tax_benefit_system_cache[system].reform_free, "simulation", None) - is None - ) + assert getattr(reform_free, "simulation", None) is previous + for item in probe.items: + assert item.simulation is None + assert item.tax_benefit_system is None + assert item._system_simulation is None + + +def test_successful_cases_restore_the_systems_previous_simulation(tmp_path): + """Successful and failed builds have the same ownership after teardown.""" + system = CountryTaxBenefitSystem() + reform_free = _get_tax_benefit_system(system, [], []) + previous = build_simulation(reform_free) + path = write_cases(tmp_path / "good.yaml", 3, reform_every=0) + + probe = run_with_probe(system, path) + + assert len(probe.outcomes) == len(probe.items) == 3 + assert probe.live_simulations == 0 + assert reform_free.simulation is previous + for item in probe.items: + assert item.simulation is None + assert item.tax_benefit_system is None + assert item._system_simulation is None def test_inline_parameter_values_key_by_type(tmp_path): diff --git a/tests/fixtures/yaml_runner_memory.py b/tests/fixtures/yaml_runner_memory.py index b541b126e..fb6264a0d 100644 --- a/tests/fixtures/yaml_runner_memory.py +++ b/tests/fixtures/yaml_runner_memory.py @@ -20,7 +20,20 @@ from policyengine_core.tracers import FullTracer from policyengine_core.tools.test_runner import OpenFiscaPlugin -RUNNER_ARGV = ["--capture", "no", "--maxfail", "0", "--tb", "short"] +# pytest-rerunfailures 16.7 stores failed ExceptionInfo objects on collected +# items even when no reruns are requested, retaining the builder's simulation +# through traceback locals. Measure the runner's ownership without this dev +# plugin; the outer test session can still rerun these tests normally. +RUNNER_ARGV = [ + "--capture", + "no", + "--maxfail", + "0", + "--tb", + "short", + "-p", + "no:rerunfailures", +] def write_cases(path: Path, count: int, reform_every: int = 3) -> Path: @@ -138,6 +151,7 @@ class MemoryProbe: def __init__(self, sample_every: int = 10): self.sample_every = sample_every + self.items = [] self.traced = [] self.outcomes = [] self.live_simulations = None @@ -161,6 +175,9 @@ def pytest_runtest_logfinish(self, nodeid, location): def pytest_sessionfinish(self, session, exitstatus): tracemalloc.stop() + # Keep the collected items alive even after pytest.main returns, so + # collecting the session itself cannot hide a runner-owned reference. + self.items = list(session.items) # pytest keeps the last failure's traceback in ``sys.last_*``; that one # root is pytest's, not the runner's. for name in ("last_type", "last_value", "last_traceback", "last_exc"):