From 6e24570c282e1c87f14c36e4ab3f833dde4fe286 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 08:35:41 -0400 Subject: [PATCH 01/23] 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 40f998f251bd89985139b31d883ce0f8549e17ff Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 08:41:45 -0400 Subject: [PATCH 02/23] 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 8bc5b43d06dac0c807fe4c5d66639e0d061e30bf Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 09:54:14 -0400 Subject: [PATCH 03/23] 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 857b6dbeadbd3db60167b336cc8ebec5a7a2da74 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 12:12:47 -0400 Subject: [PATCH 04/23] 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 73ab993d1dc1fa8e47e403d991b3d10e53498c9f Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sat, 3 Oct 2026 08:03:40 -0400 Subject: [PATCH 05/23] 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 ab8ca194ba4be12433a8be6f54ee15dc93a9c5e5 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Tue, 6 Oct 2026 09:27:04 -0400 Subject: [PATCH 06/23] 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 212de058d..74d01c23c 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 from policyengine_core.variables.variable import NUMERIC_VALUE_TYPES @@ -609,10 +610,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 0c051ca63d3eca9c00f1bf905636208dc16473aa Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Tue, 6 Oct 2026 09:27:21 -0400 Subject: [PATCH 07/23] 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 e1fc0d4b907249f1d7d0f0afc8729d3985c9321d Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Wed, 7 Oct 2026 17:05:55 -0400 Subject: [PATCH 08/23] 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 202697756406514dc52ff43d49af5e5c647d3556 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Wed, 7 Oct 2026 17:07:34 -0400 Subject: [PATCH 09/23] 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"): From 7e0cd814d08c1ca5bc94aa4663a6f2f69545e0c1 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 08:35:41 -0400 Subject: [PATCH 10/23] Resolve ParameterNodeAtInstant children on first read Building every descendant in the constructor made each instant a system is asked about cost a copy of the whole parameter tree, kept for the system's lifetime. For policyengine-us that is about 34,000 node objects, 16-33 MB and several seconds per instant, per tax-benefit system. Co-Authored-By: Claude Opus 5.5 --- .../lazy-parameter-node-at-instant.changed.md | 1 + .../parameters/parameter_node_at_instant.py | 81 +++++++- .../test_parameter_node_at_instant_lazy.py | 194 ++++++++++++++++++ ...test_parameter_node_at_instant_property.py | 91 ++++++++ tests/fixtures/parameter_nodes_at_instant.py | 119 +++++++++++ 5 files changed, 478 insertions(+), 8 deletions(-) create mode 100644 changelog.d/lazy-parameter-node-at-instant.changed.md create mode 100644 tests/core/parameters/test_parameter_node_at_instant_lazy.py create mode 100644 tests/core/parameters/test_parameter_node_at_instant_property.py create mode 100644 tests/fixtures/parameter_nodes_at_instant.py diff --git a/changelog.d/lazy-parameter-node-at-instant.changed.md b/changelog.d/lazy-parameter-node-at-instant.changed.md new file mode 100644 index 000000000..00e96589b --- /dev/null +++ b/changelog.d/lazy-parameter-node-at-instant.changed.md @@ -0,0 +1 @@ +Resolve `ParameterNodeAtInstant` children on first read instead of building the whole parameter tree for every instant, which cost policyengine-us 16 to 33 MB and several seconds per instant on every tax-benefit system. diff --git a/policyengine_core/parameters/parameter_node_at_instant.py b/policyengine_core/parameters/parameter_node_at_instant.py index b55889eb6..776550569 100644 --- a/policyengine_core/parameters/parameter_node_at_instant.py +++ b/policyengine_core/parameters/parameter_node_at_instant.py @@ -16,9 +16,23 @@ ) +# Attributes of the node itself. A lookup of one that is not set is an +# ordinary missing attribute, never a parameter. +_TECHNICAL_ATTRIBUTES = frozenset( + {"_name", "_instant_str", "_node", "_resolved", "_children", "_vectorial_node"} +) + + class ParameterNodeAtInstant: """ Parameter node of the legislation, at a given instant. + + Children resolve on first read and are kept from then on. Building every + descendant up front made each instant a model is asked about cost a copy of + the whole parameter tree (16 to 33 MB and several seconds per instant for + policyengine-us), which every tax-benefit system then kept for its + lifetime. A child that has not been read yet takes the value its parameter + has when it is read. """ def __init__(self, name: str, node: "ParameterNode", instant_str: str): @@ -31,23 +45,68 @@ def __init__(self, name: str, node: "ParameterNode", instant_str: str): # The "technical" attributes are hidden, so that the node children can be easily browsed with auto-completion without pollution self._name = name self._instant_str = instant_str - self._children = {} + # The node whose children are still to be resolved; ``None`` once + # ``_children`` holds all of them. + self._node = node + # Children resolved so far, by name. Becomes ``_children`` itself once + # every child is resolved. + self._resolved = {} - for child_name, child in node.children.items(): - child_at_instant = child._get_at_instant(instant_str) + def _resolve(self, child_name: str): + """Resolve one child, returning ``None`` if this node has no such + child at its instant.""" + node = self._node + child = None if node is None else node.children.get(child_name) + if child is None: + return None + child_at_instant = child._get_at_instant(self._instant_str) + if child_at_instant is not None: + self.add_child(child_name, child_at_instant) + return child_at_instant + + def _resolve_all(self) -> dict: + """Resolve every child, and return them in the parameter node's order.""" + resolved = self._resolved + children = {} + for child_name, child in self._node.children.items(): + if child_name in resolved: + children[child_name] = resolved[child_name] + continue + child_at_instant = child._get_at_instant(self._instant_str) if child_at_instant is not None: - self.add_child(child_name, child_at_instant) + children[child_name] = child_at_instant + if child_name not in _TECHNICAL_ATTRIBUTES: + setattr(self, child_name, child_at_instant) + for child_name, child_at_instant in resolved.items(): + # Children added with ``add_child`` that the node does not have. + children.setdefault(child_name, child_at_instant) + self._children = self._resolved = children + self._node = None + return children def add_child(self, child_name: str, child_at_instant: "ParameterNodeAtInstant"): - self._children[child_name] = child_at_instant - setattr(self, child_name, child_at_instant) + self._resolved[child_name] = child_at_instant + if child_name not in _TECHNICAL_ATTRIBUTES: + setattr(self, child_name, child_at_instant) def __getattr__(self, key: str): - if key == "_name": + # Reached only for a name that is not an instance attribute yet: a + # child that has not been read, ``_children`` before every child has + # been resolved, or a name this node does not have. + if key in _TECHNICAL_ATTRIBUTES: + if key == "_children" and self.__dict__.get("_node") is not None: + return self._resolve_all() raise AttributeError(key) + if self.__dict__.get("_node") is not None: + child_at_instant = self._resolve(key) + if child_at_instant is not None: + return child_at_instant param_name = helpers._compose_name(self._name, item_name=key) raise ParameterNotFoundError(param_name, self._instant_str) + def __dir__(self): + return sorted(set(super().__dir__()) | set(map(str, self._children))) + def __getitem__( self, key: str ) -> Union["ParameterNodeAtInstant", VectorialParameterNodeAtInstant]: @@ -68,7 +127,13 @@ def __getitem__( ) self._vectorial_node = vectorial return vectorial[key] - return self._children[key] + resolved = self._resolved + if key in resolved: + return resolved[key] + child_at_instant = self._resolve(key) + if child_at_instant is None: + raise KeyError(key) + return child_at_instant def __iter__(self) -> Iterable: return iter(self._children) diff --git a/tests/core/parameters/test_parameter_node_at_instant_lazy.py b/tests/core/parameters/test_parameter_node_at_instant_lazy.py new file mode 100644 index 000000000..81fbf2e1d --- /dev/null +++ b/tests/core/parameters/test_parameter_node_at_instant_lazy.py @@ -0,0 +1,194 @@ +"""``ParameterNodeAtInstant`` resolves its children on first read. + +It used to build every descendant in its constructor, so each instant a +system was asked about cost a copy of the whole parameter tree (about 34,000 +node objects and 16-33 MB per instant for policyengine-us) that the system +then kept. These tests pin that nothing is built before it is read, and that +everything read is what the up-front build held. +""" + +import copy +import gc + +import numpy as np +import pytest + +from policyengine_core.country_template import CountryTaxBenefitSystem +from policyengine_core.errors import ParameterNotFoundError +from policyengine_core.parameters import ParameterNodeAtInstant +from tests.fixtures.parameter_nodes_at_instant import ( + INSTANTS, + build_tree, + materialise, + ordered, + snapshot, +) + + +@pytest.mark.parametrize("instant", INSTANTS) +def test_fully_read_node_matches_the_up_front_build(instant): + tree = build_tree() + assert ordered(materialise(tree.get_at_instant(instant))) == ordered( + snapshot(tree, instant) + ) + + +@pytest.mark.parametrize( + "instant", ["2000-01-01", "2012-06-01", "2015-01-01", "2017-01-01", "2030-01-01"] +) +def test_country_template_matches_the_up_front_build(instant): + parameters = CountryTaxBenefitSystem().parameters + assert ordered(materialise(parameters.get_at_instant(instant))) == ordered( + snapshot(parameters, instant) + ) + + +def test_nothing_is_resolved_until_it_is_read(): + at_instant = build_tree().get_at_instant("2017-03-01") + assert at_instant._resolved == {} + + assert at_instant.group.x == 0.1 + # Only the path that was read: the root's ``group`` and the group's ``x``. + assert list(at_instant._resolved) == ["group"] + assert list(at_instant.group._resolved) == ["x"] + + +def test_asking_for_many_instants_builds_one_node_each(): + tree = build_tree() + gc.collect() + before = sum(isinstance(o, ParameterNodeAtInstant) for o in gc.get_objects()) + for year in range(2000, 2100): + tree.get_at_instant(f"{year}-01-01") + after = sum(isinstance(o, ParameterNodeAtInstant) for o in gc.get_objects()) + # The up-front build made 9 nodes per instant for this tree (900 here). + assert after - before == 100 + + +def test_a_read_child_is_resolved_once(): + at_instant = build_tree().get_at_instant("2017-03-01") + # A scale is rebuilt by every ``_get_at_instant`` call, so identity shows + # that the node kept what it resolved. + assert at_instant.scale is at_instant.scale + assert at_instant["scale"] is at_instant.scale + assert at_instant._children["scale"] is at_instant.scale + assert at_instant.group is at_instant["group"] + + +def test_children_are_in_parameter_order_whatever_was_read_first(): + tree = build_tree() + at_instant = tree.get_at_instant("2030-01-01") + at_instant.by_zone + at_instant["late"] + at_instant.group + assert list(at_instant) == list(snapshot(tree, "2030-01-01")) + assert list(at_instant._children) == list(snapshot(tree, "2030-01-01")) + + +def test_child_without_a_value_yet_is_absent(): + at_instant = build_tree().get_at_instant("2017-03-01") + with pytest.raises(ParameterNotFoundError): + at_instant.late + with pytest.raises(ParameterNotFoundError) as error: + at_instant.group.y + assert str(error.value) == ( + "The parameter 'group[y]' was not found in the 2017-03-01 " + "tax and benefit system." + ) + with pytest.raises(KeyError): + at_instant["late"] + assert not hasattr(at_instant, "late") + assert "late" not in at_instant + assert "late" not in at_instant._children + assert at_instant.group.inner._children == {} + + +def test_unknown_names_raise_as_before(): + at_instant = build_tree().get_at_instant("2017-03-01") + for read_first in (False, True): + if read_first: + at_instant._children + with pytest.raises(ParameterNotFoundError) as error: + at_instant.group.nope + assert str(error.value) == ( + "The parameter 'group[nope]' was not found in the 2017-03-01 " + "tax and benefit system." + ) + with pytest.raises(KeyError): + at_instant.group["nope"] + assert not hasattr(at_instant, "nope") + + +def test_vectorial_indexing_reads_the_whole_subtree(): + at_instant = build_tree().get_at_instant("2020-01-01") + zones = np.array(["z1", "z3", "z2"]) + np.testing.assert_array_equal(at_instant.by_zone[zones].amount, [200, 600, 400]) + np.testing.assert_allclose(at_instant.by_zone[zones].rate, [0.1, 0.3, 0.2]) + # Nothing outside the indexed subtree was resolved for it. + assert list(at_instant._resolved) == ["by_zone"] + + +def test_dir_lists_children_for_completion(): + at_instant = build_tree().get_at_instant("2017-03-01") + names = dir(at_instant) + assert {"flat_rate", "group", "scale", "by_zone"} <= set(names) + assert "late" not in names + + +def test_repr_shows_every_child(): + text = repr(build_tree().get_at_instant("2017-03-01")) + for name in ("flat_rate", "group", "scale", "by_zone", "z3"): + assert name in text + + +def test_add_child_before_and_after_a_full_read(): + at_instant = build_tree().get_at_instant("2017-03-01") + at_instant.add_child("extra", 1) + assert at_instant.extra == 1 and at_instant["extra"] == 1 + assert list(at_instant._children)[-1] == "extra" + at_instant.add_child("later", 2) + assert at_instant.later == 2 and at_instant._children["later"] == 2 + + +def test_a_node_asked_for_after_an_update_reads_the_update(): + tree = build_tree() + assert tree.get_at_instant("2017-03-01").flat_rate == 0.1 + tree.flat_rate.update(value=0.5, period="year:2017:1") + assert tree.get_at_instant("2017-03-01").flat_rate == 0.5 + + +def test_a_child_read_before_an_update_keeps_what_was_read(): + tree = build_tree() + held = tree.get_at_instant("2017-03-01") + assert held.flat_rate == 0.1 + tree.flat_rate.update(value=0.5, period="year:2017:1") + assert held.flat_rate == 0.1 + + +def test_a_child_not_read_before_an_update_reads_the_update(): + """Intended difference from the up-front build, which held 0.1 here: a + node held across a change to its parameters reads, for each child, the + value at the time that child is first read.""" + tree = build_tree() + held = tree.get_at_instant("2017-03-01") + tree.flat_rate.update(value=0.5, period="year:2017:1") + assert held.flat_rate == 0.5 + + +def test_deep_copied_tree_keeps_its_own_nodes_at_instant(): + tree = build_tree() + tree.get_at_instant("2017-03-01").group.x # cache a partly read node + clone = copy.deepcopy(tree) + cached = clone._at_instant_cache["2017-03-01"] + clone.group.y.update(value=0.9, period="year:2010:30") + clone.flat_rate.update(value=0.7, period="year:2010:30") + # The copy's cached node reads the copy, not the original tree. + assert cached.flat_rate == 0.7 + assert tree.get_at_instant("2017-03-01").flat_rate == 0.1 + assert not hasattr(tree.get_at_instant("2017-03-01").group, "y") + + +def test_copy_of_a_node_at_instant_reads_the_same_values(): + at_instant = build_tree().get_at_instant("2017-03-01") + at_instant.group + for duplicate in (copy.copy(at_instant), copy.deepcopy(at_instant)): + assert materialise(duplicate) == materialise(at_instant) diff --git a/tests/core/parameters/test_parameter_node_at_instant_property.py b/tests/core/parameters/test_parameter_node_at_instant_property.py new file mode 100644 index 000000000..5eeec9372 --- /dev/null +++ b/tests/core/parameters/test_parameter_node_at_instant_property.py @@ -0,0 +1,91 @@ +"""Reads of a lazily resolved ``ParameterNodeAtInstant`` in any order. + +For any instant and any sequence of reads (attribute, item, iteration, +``_children``, ``dir``, ``repr``, vectorial indexing, names that do not +exist), every read returns what the up-front build held, and the node read in +full afterwards equals that build. ``test_parameter_node_at_instant_lazy.py`` +pins the same behaviour with examples. +""" + +from __future__ import annotations + +import numpy as np +import pytest + +# The country-package smoke job installs no dev dependencies. +hypothesis = pytest.importorskip("hypothesis") +st = hypothesis.strategies + +from policyengine_core.errors import ParameterNotFoundError +from policyengine_core.parameters import ParameterNodeAtInstant +from tests.fixtures.parameter_nodes_at_instant import ( + INSTANTS, + build_tree, + describe, + materialise, + ordered, + snapshot, +) + +# Node paths in the fixture tree, and names to ask each for (present at some +# instants, absent at others, or never there). +PATHS = [(), ("group",), ("group", "inner"), ("by_zone",), ("by_zone", "z2")] +NAMES = ["flat_rate", "late", "group", "scale", "by_zone", "x", "y", "inner"] +NAMES += ["1", "2", "z1", "z2", "z3", "amount", "rate", "nope"] +KINDS = ["getattr", "getitem", "iter", "children", "dir", "repr", "vector"] + +reads = st.tuples( + st.sampled_from(PATHS), st.sampled_from(KINDS), st.sampled_from(NAMES) +) + + +def _comparable(value): + if isinstance(value, ParameterNodeAtInstant): + return materialise(value) + return describe(value) + + +@hypothesis.given( + instant=st.sampled_from(INSTANTS), sequence=st.lists(reads, max_size=30) +) +@hypothesis.settings(max_examples=400, deadline=None) +def test_reads_in_any_order_match_the_up_front_build(instant, sequence): + tree = build_tree() + expected_root = snapshot(tree, instant) + root = tree.get_at_instant(instant) + + for path, kind, name in sequence: + node, expected = root, expected_root + for step in path: + node, expected = node[step], expected[step] + + if kind == "getattr": + if name in expected: + assert _comparable(getattr(node, name)) == expected[name] + else: + with pytest.raises(ParameterNotFoundError): + getattr(node, name) + assert not hasattr(node, name) + elif kind == "getitem": + if name in expected: + assert _comparable(node[name]) == expected[name] + else: + with pytest.raises(KeyError): + node[name] + elif kind == "iter": + assert list(node) == list(expected) + assert (name in node) == (name in expected) + elif kind == "children": + assert ordered(materialise(node)) == ordered(expected) + elif kind == "dir": + assert set(expected) <= set(dir(node)) + elif kind == "repr": + assert all(child in repr(node) for child in expected) + elif kind == "vector" and path == ("by_zone",) and expected["z1"]: + zones = np.array(["z3", "z1"]) + np.testing.assert_array_equal( + node[zones].amount, + [expected["z3"]["amount"], expected["z1"]["amount"]], + ) + + assert ordered(materialise(root)) == ordered(expected_root) diff --git a/tests/fixtures/parameter_nodes_at_instant.py b/tests/fixtures/parameter_nodes_at_instant.py new file mode 100644 index 000000000..70198ed31 --- /dev/null +++ b/tests/fixtures/parameter_nodes_at_instant.py @@ -0,0 +1,119 @@ +"""A parameter tree and an up-front reference for the at-instant tests. + +``snapshot`` computes what a ``ParameterNodeAtInstant`` held when it built +every descendant in its constructor, straight from the parameter tree and +without using ``ParameterNodeAtInstant``. ``materialise`` reads the same +thing out of a ``ParameterNodeAtInstant``. +""" + +from policyengine_core.parameters import ( + Parameter, + ParameterNode, + ParameterNodeAtInstant, + ParameterScale, +) + +INSTANTS = [ + "2009-12-31", + "2010-01-01", + "2014-06-30", + "2015-01-01", + "2017-03-01", + "2018-01-01", + "2020-01-01", + "2029-12-31", + "2030-01-01", + "2035-07-01", +] + + +def _zone(amount, rate): + return { + "amount": {"values": {"2010-01-01": amount, "2020-01-01": amount * 2}}, + "rate": {"values": {"2010-01-01": rate}}, + } + + +TREE = { + "flat_rate": {"values": {"2015-01-01": 0.1, "2020-01-01": 0.2}}, + # Not defined before 2030, so absent from earlier instants. + "late": {"values": {"2030-01-01": 5}}, + "group": { + "x": {"values": {"2010-01-01": 0.1}}, + "y": {"values": {"2018-01-01": 0.2}}, + # Every child undefined before 2018: an empty node, not an absent one. + "inner": { + "1": {"values": {"2018-01-01": 10}}, + "2": {"values": {"2018-01-01": 20, "2020-01-01": 30}}, + }, + }, + "scale": { + "brackets": [ + { + "threshold": {"2010-01-01": {"value": 0}}, + "rate": {"2010-01-01": {"value": 0.1}}, + }, + { + "threshold": {"2010-01-01": {"value": 1000}}, + "rate": { + "2010-01-01": {"value": 0.2}, + "2020-01-01": {"value": 0.3}, + }, + }, + ] + }, + # Homogeneous children: vectorisable with ``node[array]``. + "by_zone": {"z1": _zone(100, 0.1), "z2": _zone(200, 0.2), "z3": _zone(300, 0.3)}, +} + + +def build_tree() -> ParameterNode: + return ParameterNode("", data=TREE) + + +def describe(value): + """A comparable form of one at-instant value.""" + if isinstance(value, (dict, int, float, bool, str)) or value is None: + return value + # Tax scales refuse ``==``. + return ( + type(value).__name__, + tuple(value.thresholds), + tuple(getattr(value, "rates", getattr(value, "amounts", ()))), + ) + + +def snapshot(node: ParameterNode, instant_str: str) -> dict: + """Every descendant's value at the instant, in the node's own order, + leaving out children that have no value yet.""" + result = {} + for name, child in node.children.items(): + if isinstance(child, ParameterScale): + value = describe(child._get_at_instant(instant_str)) + elif isinstance(child, Parameter): + value = child._get_at_instant(instant_str) + else: + value = snapshot(child, instant_str) + if value is not None: + result[name] = value + return result + + +def materialise(node_at_instant: ParameterNodeAtInstant) -> dict: + """The same nested mapping, read from a ``ParameterNodeAtInstant``.""" + return { + name: ( + materialise(child) + if isinstance(child, ParameterNodeAtInstant) + else describe(child) + ) + for name, child in node_at_instant._children.items() + } + + +def ordered(mapping: dict): + """A form that compares key order as well as content.""" + return [ + (key, ordered(value) if isinstance(value, dict) else value) + for key, value in mapping.items() + ] From 084f168aed9ce41c31a0dee3e1d7b2f31f20a29d Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 09:52:14 -0400 Subject: [PATCH 11/23] Make lazy at-instant resolution thread-safe, collision-proof and copy-coherent Review of 404a3a38 found five regressions against the up-front build: - concurrent first reads could raise KeyError/AttributeError or resolve a child twice: first resolution now runs under a module RLock; - children named like the new state or helpers (_node, _resolved, _resolve, ...) read as internal objects: state is name-mangled, helpers are module functions, and children named like the node's own attributes resolve at construction, as before; - a shallow copy and its original could resolve a scale twice: both now read one grow-only resolved map; - a child with no value at the instant was re-resolved on every read: its absence is cached; - deepcopy/pickle of a traced tree's raw node recursed through TracingParameterNodeAtInstant.__getattr__, which now raises AttributeError for its own unset fields. A node also drops its parameter node once every child has been read. Co-Authored-By: Claude Opus 5.5 --- .../parameters/parameter_node_at_instant.py | 198 ++++++++++++------ .../test_parameter_node_at_instant_lazy.py | 181 +++++++++++++++- 2 files changed, 310 insertions(+), 69 deletions(-) diff --git a/policyengine_core/parameters/parameter_node_at_instant.py b/policyengine_core/parameters/parameter_node_at_instant.py index 776550569..010d5f545 100644 --- a/policyengine_core/parameters/parameter_node_at_instant.py +++ b/policyengine_core/parameters/parameter_node_at_instant.py @@ -1,5 +1,6 @@ import os import sys +import threading from typing import TYPE_CHECKING, Iterable, Union import numpy @@ -16,11 +17,20 @@ ) -# Attributes of the node itself. A lookup of one that is not set is an -# ordinary missing attribute, never a parameter. -_TECHNICAL_ATTRIBUTES = frozenset( - {"_name", "_instant_str", "_node", "_resolved", "_children", "_vectorial_node"} -) +# A node's own state, under name-mangled keys so that no parameter child can +# shadow it. +_NODE = "_ParameterNodeAtInstant__node" +_INSTANT = "_ParameterNodeAtInstant__instant" +_RESOLVED = "_ParameterNodeAtInstant__resolved" +_ABSENT = "_ParameterNodeAtInstant__absent" +_CHILDREN = "_ParameterNodeAtInstant__children" +_MISSING = object() + +# Held while a child is resolved for the first time, so that concurrent +# readers never see a half-resolved node or resolve one child twice. Reads of +# a resolved child never take it. Reentrant because resolving a child can +# build that child's own node. +_RESOLUTION_LOCK = threading.RLock() class ParameterNodeAtInstant: @@ -32,7 +42,8 @@ class ParameterNodeAtInstant: the whole parameter tree (16 to 33 MB and several seconds per instant for policyengine-us), which every tax-benefit system then kept for its lifetime. A child that has not been read yet takes the value its parameter - has when it is read. + has when it is first read; a child read once, including one found to have + no value at this instant, keeps what that read found. """ def __init__(self, name: str, node: "ParameterNode", instant_str: str): @@ -45,67 +56,51 @@ def __init__(self, name: str, node: "ParameterNode", instant_str: str): # The "technical" attributes are hidden, so that the node children can be easily browsed with auto-completion without pollution self._name = name self._instant_str = instant_str - # The node whose children are still to be resolved; ``None`` once - # ``_children`` holds all of them. - self._node = node - # Children resolved so far, by name. Becomes ``_children`` itself once - # every child is resolved. - self._resolved = {} - - def _resolve(self, child_name: str): - """Resolve one child, returning ``None`` if this node has no such - child at its instant.""" - node = self._node - child = None if node is None else node.children.get(child_name) - if child is None: - return None - child_at_instant = child._get_at_instant(self._instant_str) - if child_at_instant is not None: - self.add_child(child_name, child_at_instant) - return child_at_instant - - def _resolve_all(self) -> dict: - """Resolve every child, and return them in the parameter node's order.""" - resolved = self._resolved - children = {} - for child_name, child in self._node.children.items(): - if child_name in resolved: - children[child_name] = resolved[child_name] - continue - child_at_instant = child._get_at_instant(self._instant_str) - if child_at_instant is not None: - children[child_name] = child_at_instant - if child_name not in _TECHNICAL_ATTRIBUTES: - setattr(self, child_name, child_at_instant) - for child_name, child_at_instant in resolved.items(): - # Children added with ``add_child`` that the node does not have. - children.setdefault(child_name, child_at_instant) - self._children = self._resolved = children - self._node = None - return children + state = self.__dict__ + # The parameter node children still resolve from; None once every + # child is resolved. + state[_NODE] = node + state[_INSTANT] = instant_str + # Children resolved so far. Only ever grows, and a shallow copy shares + # it, so every copy reads the same child objects. + state[_RESOLVED] = {} + # Children of the node that have no value at this instant. + state[_ABSENT] = set() + # A child whose name the node itself uses (``add_child``, ``_name``, + # ...) shadows that name, as when every child was set here. + for child_name in _SHADOWING_NAMES: + if child_name in node.children: + _resolve(self, child_name) def add_child(self, child_name: str, child_at_instant: "ParameterNodeAtInstant"): - self._resolved[child_name] = child_at_instant - if child_name not in _TECHNICAL_ATTRIBUTES: - setattr(self, child_name, child_at_instant) + with _RESOLUTION_LOCK: + _store(self, child_name, child_at_instant) def __getattr__(self, key: str): - # Reached only for a name that is not an instance attribute yet: a - # child that has not been read, ``_children`` before every child has - # been resolved, or a name this node does not have. - if key in _TECHNICAL_ATTRIBUTES: - if key == "_children" and self.__dict__.get("_node") is not None: - return self._resolve_all() + # Reached only for a name that is not an instance attribute: a child + # not read yet, ``_children`` before every child is resolved, or a name + # this node does not have. + if key.startswith("__") and key.endswith("__"): + # Protocol lookups (copy, pickle, numpy) are never parameters. + raise AttributeError(key) + if _RESOLVED not in self.__dict__ or key == "_name": + # An instance still being rebuilt by copy or pickle. raise AttributeError(key) - if self.__dict__.get("_node") is not None: - child_at_instant = self._resolve(key) - if child_at_instant is not None: - return child_at_instant + if key == "_children": + return _materialise(self) + child_at_instant = _resolve(self, key) + if child_at_instant is not _MISSING: + return child_at_instant param_name = helpers._compose_name(self._name, item_name=key) raise ParameterNotFoundError(param_name, self._instant_str) def __dir__(self): - return sorted(set(super().__dir__()) | set(map(str, self._children))) + self._children + return [ + name + for name in super().__dir__() + if not name.startswith("_ParameterNodeAtInstant__") + ] def __getitem__( self, key: str @@ -127,12 +122,11 @@ def __getitem__( ) self._vectorial_node = vectorial return vectorial[key] - resolved = self._resolved - if key in resolved: - return resolved[key] - child_at_instant = self._resolve(key) - if child_at_instant is None: - raise KeyError(key) + child_at_instant = self.__dict__[_RESOLVED].get(key, _MISSING) + if child_at_instant is _MISSING: + child_at_instant = _resolve(self, key) + if child_at_instant is _MISSING: + raise KeyError(key) return child_at_instant def __iter__(self) -> Iterable: @@ -146,3 +140,79 @@ def __repr__(self) -> str: ] ) return result + + +def _store(node_at_instant, child_name, child_at_instant): + """Record a child. The caller holds ``_RESOLUTION_LOCK``.""" + state = node_at_instant.__dict__ + state[_RESOLVED][child_name] = child_at_instant + state[_ABSENT].discard(child_name) + children = state.get(_CHILDREN) + if children is not None: + children[child_name] = child_at_instant + setattr(node_at_instant, child_name, child_at_instant) + + +def _resolve(node_at_instant, child_name): + """Return one child, resolving it if no read has yet; ``_MISSING`` if the + node has no such child or it has no value at this instant.""" + state = node_at_instant.__dict__ + with _RESOLUTION_LOCK: + resolved = state[_RESOLVED] + if child_name in resolved: + return resolved[child_name] + if child_name in state[_ABSENT]: + return _MISSING + node = state[_NODE] + child = None if node is None else node.children.get(child_name) + if child is None: + return _MISSING + child_at_instant = child._get_at_instant(state[_INSTANT]) + if child_at_instant is None: + state[_ABSENT].add(child_name) + else: + _store(node_at_instant, child_name, child_at_instant) + if len(resolved) + len(state[_ABSENT]) >= len(node.children): + # Every child has been read: drop the parameter node. + _materialise(node_at_instant) + return _MISSING if child_at_instant is None else child_at_instant + + +def _materialise(node_at_instant): + """Resolve every child; return them in the parameter node's order, with + any added by ``add_child`` after them.""" + state = node_at_instant.__dict__ + with _RESOLUTION_LOCK: + children = state.get(_CHILDREN) + if children is not None: + return children + node = state[_NODE] + children = {} + for child_name in node.children if node is not None else (): + child_at_instant = _resolve(node_at_instant, child_name) + if child_at_instant is not _MISSING: + children[child_name] = child_at_instant + if state.get(_CHILDREN) is not None: + # Resolving the last child completed the node already. + return state[_CHILDREN] + for child_name, child_at_instant in state[_RESOLVED].items(): + children.setdefault(child_name, child_at_instant) + state[_CHILDREN] = children + if "_children" not in state: + # Unless a child named ``_children`` took the name. + state["_children"] = children + # Every child is resolved: the parameter node is no longer needed. + state[_NODE] = None + return children + + +# Names a ParameterNodeAtInstant itself answers to, which a child of the same +# name shadows. +_SHADOWING_NAMES = frozenset( + {"_name", "_instant_str", "_children", "_vectorial_node"} + | { + name + for name in dir(ParameterNodeAtInstant) + if not name.startswith("__") and not name.startswith("_ParameterNodeAtInstant") + } +) diff --git a/tests/core/parameters/test_parameter_node_at_instant_lazy.py b/tests/core/parameters/test_parameter_node_at_instant_lazy.py index 81fbf2e1d..ff2604e02 100644 --- a/tests/core/parameters/test_parameter_node_at_instant_lazy.py +++ b/tests/core/parameters/test_parameter_node_at_instant_lazy.py @@ -9,13 +9,18 @@ import copy import gc +import pickle +import threading +import weakref import numpy as np import pytest from policyengine_core.country_template import CountryTaxBenefitSystem from policyengine_core.errors import ParameterNotFoundError -from policyengine_core.parameters import ParameterNodeAtInstant +from policyengine_core.parameters import ParameterNode, ParameterNodeAtInstant +from policyengine_core.parameters.parameter_node_at_instant import _RESOLVED +from policyengine_core.tracers import FullTracer from tests.fixtures.parameter_nodes_at_instant import ( INSTANTS, build_tree, @@ -43,14 +48,18 @@ def test_country_template_matches_the_up_front_build(instant): ) +def resolved(node_at_instant): + return list(node_at_instant.__dict__[_RESOLVED]) + + def test_nothing_is_resolved_until_it_is_read(): at_instant = build_tree().get_at_instant("2017-03-01") - assert at_instant._resolved == {} + assert resolved(at_instant) == [] assert at_instant.group.x == 0.1 # Only the path that was read: the root's ``group`` and the group's ``x``. - assert list(at_instant._resolved) == ["group"] - assert list(at_instant.group._resolved) == ["x"] + assert resolved(at_instant) == ["group"] + assert resolved(at_instant.group) == ["x"] def test_asking_for_many_instants_builds_one_node_each(): @@ -124,7 +133,7 @@ def test_vectorial_indexing_reads_the_whole_subtree(): np.testing.assert_array_equal(at_instant.by_zone[zones].amount, [200, 600, 400]) np.testing.assert_allclose(at_instant.by_zone[zones].rate, [0.1, 0.3, 0.2]) # Nothing outside the indexed subtree was resolved for it. - assert list(at_instant._resolved) == ["by_zone"] + assert resolved(at_instant) == ["by_zone"] def test_dir_lists_children_for_completion(): @@ -192,3 +201,165 @@ def test_copy_of_a_node_at_instant_reads_the_same_values(): at_instant.group for duplicate in (copy.copy(at_instant), copy.deepcopy(at_instant)): assert materialise(duplicate) == materialise(at_instant) + + +@pytest.mark.parametrize( + "child_name", + # Not ``add_child``: ParameterNode itself cannot hold a child of that name. + ["_name", "_instant_str", "_vectorial_node", "_node", "_resolved"] + + ["_resolve", "_resolve_all", "_store", "_materialise"], +) +def test_children_named_like_the_nodes_own_attributes(child_name): + """Such a child reads as its value whatever the read order, as when every + child was set up front, and leaves its siblings readable.""" + tree = ParameterNode( + "", + data={ + child_name: {"values": {"2010-01-01": 42}}, + "ok": {"values": {"2010-01-01": 7}}, + }, + ) + for read_name_first in (True, False): + tree._at_instant_cache.clear() + at_instant = tree.get_at_instant("2017-03-01") + if read_name_first: + assert getattr(at_instant, child_name) == 42 + assert at_instant.ok == 7 and at_instant["ok"] == 7 + else: + assert at_instant["ok"] == 7 and at_instant.ok == 7 + assert getattr(at_instant, child_name) == 42 + assert at_instant[child_name] == 42 + assert set(at_instant._children) == {child_name, "ok"} + + +class BlockingChild: + """A parameter child whose first resolution waits until released.""" + + def __init__(self): + self.calls = 0 + self.entered = threading.Event() + self.release = threading.Event() + + def _get_at_instant(self, instant): + self.calls += 1 + if self.calls == 1: + self.entered.set() + assert self.release.wait(10), "never released" + return object() + + +@pytest.mark.parametrize( + "second_read", + [ + lambda node: node.slow, + lambda node: node["slow"], + lambda node: node._children["slow"], + lambda node: [node[name] for name in node][0], + ], +) +def test_concurrent_first_reads_resolve_a_child_once(second_read): + tree = ParameterNode("", data={"other": {"values": {"2010-01-01": 1}}}) + slow = BlockingChild() + tree.children = {"slow": slow, **tree.children} + at_instant = tree.get_at_instant("2017-03-01") + results, errors = [], [] + + def read(reader): + try: + results.append(reader(at_instant)) + except Exception as error: # noqa: BLE001 - recorded for the assertion + errors.append(error) + + first = threading.Thread(target=read, args=(lambda node: node.slow,)) + first.start() + assert slow.entered.wait(10) + second = threading.Thread(target=read, args=(second_read,)) + second.start() + second.join(0.2) # let the second reader reach the child too + slow.release.set() + first.join(10) + second.join(10) + + assert not errors, errors + assert len(results) == 2 and results[0] is results[1] + assert slow.calls == 1 + + +def test_shallow_copies_read_the_same_children(): + held = build_tree().get_at_instant("2017-03-01") + duplicate = copy.copy(held) + original_scale = held.scale + # A scale is rebuilt by each resolution, so identity shows one resolution. + assert duplicate.scale is original_scale + assert duplicate["scale"] is original_scale + assert held["scale"] is original_scale + duplicate._children + assert held._children["scale"] is original_scale + assert held.by_zone is duplicate["by_zone"] + + +@pytest.mark.parametrize( + "duplicate", [copy.deepcopy, lambda n: pickle.loads(pickle.dumps(n))] +) +def test_traced_nodes_copy_and_pickle(duplicate): + tree = build_tree() + tree.trace = True + tree.tracer = FullTracer() + tree.branch_name = "default" + traced = tree.get_at_instant("2017-03-01") + traced.group.x # partly resolved, through the tracing wrapper + raw = traced.parameter_node_at_instant + for node in (raw, traced): + assert duplicate(node).flat_rate == 0.1 + assert duplicate(node).group.x == 0.1 + + +def test_a_child_without_a_value_is_resolved_once(): + tree = ParameterNode("", data={"late": {"values": {"2030-01-01": 1}}}) + held = tree.get_at_instant("2020-01-01") + calls = [] + original = tree.children["late"]._get_at_instant + + def counting(instant): + calls.append(instant) + return original(instant) + + tree.children["late"]._get_at_instant = counting + for _ in range(3): + with pytest.raises(ParameterNotFoundError): + held.late + with pytest.raises(KeyError): + held["late"] + assert list(held) == [] + assert len(calls) == 1 + + # A child already read keeps what that read found, value or none. + tree.children["late"].update(value=2, period="year:2020:1") + with pytest.raises(KeyError): + held["late"] + assert tree.get_at_instant("2020-01-01").late == 2 + + +def test_a_node_lets_go_of_its_parameter_node_once_every_child_is_read(): + tree = ParameterNode( + "", + data={ + "group": { + "x": {"values": {"2010-01-01": 1}}, + "late": {"values": {"2030-01-01": 2}}, + }, + "unused": {"values": {"2010-01-01": 3}}, + }, + ) + group = tree.get_at_instant("2017-03-01").group + tree_ref = weakref.ref(tree) + del tree + gc.collect() + # Partly read: the node still needs the tree to resolve the rest. + assert tree_ref() is not None + assert group.x == 1 + with pytest.raises(ParameterNotFoundError): + group.late + gc.collect() + assert tree_ref() is None + assert group.x == 1 and list(group) == ["x"] From 6dba3b4afdb5b3a0d7eff11b871d46d1e6b30a13 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 10:00:12 -0400 Subject: [PATCH 12/23] Compare list and array parameter values in the at-instant reference Co-Authored-By: Claude Opus 5.5 --- tests/fixtures/parameter_nodes_at_instant.py | 26 ++++++++++++++------ 1 file changed, 19 insertions(+), 7 deletions(-) diff --git a/tests/fixtures/parameter_nodes_at_instant.py b/tests/fixtures/parameter_nodes_at_instant.py index 70198ed31..7b991041b 100644 --- a/tests/fixtures/parameter_nodes_at_instant.py +++ b/tests/fixtures/parameter_nodes_at_instant.py @@ -6,6 +6,8 @@ thing out of a ``ParameterNodeAtInstant``. """ +import numpy as np + from policyengine_core.parameters import ( Parameter, ParameterNode, @@ -73,14 +75,24 @@ def build_tree() -> ParameterNode: def describe(value): """A comparable form of one at-instant value.""" - if isinstance(value, (dict, int, float, bool, str)) or value is None: + if isinstance(value, (int, float, bool, str)) or value is None: return value - # Tax scales refuse ``==``. - return ( - type(value).__name__, - tuple(value.thresholds), - tuple(getattr(value, "rates", getattr(value, "amounts", ()))), - ) + if isinstance(value, (list, tuple)): + return tuple(describe(item) for item in value) + if isinstance(value, dict): + return {key: describe(item) for key, item in value.items()} + if hasattr(value, "thresholds"): + # Tax scales refuse ``==``. + return ( + type(value).__name__, + tuple(value.thresholds), + tuple(getattr(value, "rates", getattr(value, "amounts", ()))), + ) + if isinstance(value, np.ndarray): + return ("ndarray", str(value.dtype), value.tolist()) + if isinstance(value, np.generic): + return value.item() + return (type(value).__name__, repr(value)) def snapshot(node: ParameterNode, instant_str: str) -> dict: From eccf4a2a65c1508d10bbccd8c244e95e8c287140 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sat, 3 Oct 2026 08:24:11 -0400 Subject: [PATCH 13/23] Keep lazy at-instant state in one shared object; fix round-2 review cases Round-2 review of a90b524a found reads that still differed from the up-front build. A node now keeps all of its own state in one _LazyChildren object under a single reserved name (a parameter of that name is rejected): - which children a node has, and their order, are fixed at construction (a child deleted from the tree later keeps its place); - completion counts only those children, so add_child extras can no longer send the first read into recursive completion; - resolution results and add_child entries are kept apart: items are shared by shallow copies, attributes belong to each copy, as before; - a copy installs what it reads as attributes (no lock on warm reads); - children named like attributes of the node's runtime type (dunders and subclass methods included) resolve at construction; a subclass overriding add_child gets the up-front build through it; - a reentrant read keeps the first recorded child; - nodes pickled by the up-front build load (__setstate__). Co-Authored-By: Claude Opus 5.5 --- .../parameters/parameter_node_at_instant.py | 259 ++++++++++-------- .../test_parameter_node_at_instant_lazy.py | 151 +++++++++- 2 files changed, 299 insertions(+), 111 deletions(-) diff --git a/policyengine_core/parameters/parameter_node_at_instant.py b/policyengine_core/parameters/parameter_node_at_instant.py index 010d5f545..43787a045 100644 --- a/policyengine_core/parameters/parameter_node_at_instant.py +++ b/policyengine_core/parameters/parameter_node_at_instant.py @@ -17,13 +17,11 @@ ) -# A node's own state, under name-mangled keys so that no parameter child can -# shadow it. -_NODE = "_ParameterNodeAtInstant__node" -_INSTANT = "_ParameterNodeAtInstant__instant" -_RESOLVED = "_ParameterNodeAtInstant__resolved" -_ABSENT = "_ParameterNodeAtInstant__absent" -_CHILDREN = "_ParameterNodeAtInstant__children" +# The one instance attribute a node keeps for itself; a parameter of this +# name is rejected. +_STATE = "_ParameterNodeAtInstant__state" +# Attributes a node sets on itself, which a child of the same name shadows. +_OWN_NAMES = frozenset({"_name", "_instant_str", "_children", "_vectorial_node"}) _MISSING = object() # Held while a child is resolved for the first time, so that concurrent @@ -33,6 +31,28 @@ _RESOLUTION_LOCK = threading.RLock() +class _LazyChildren: + """What a node at instant knows about its children. A shallow copy of the + node shares it, as copies of a node that built every child up front shared + their ``_children`` mapping.""" + + __slots__ = ("instant", "source", "resolved", "absent", "added", "children") + + def __init__(self, instant: str, source: dict): + self.instant = instant + # Each child's parameter node as at construction, in order; None once + # every child is resolved. + self.source = source + # Children resolved from ``source`` so far. + self.resolved = {} + # Children of ``source`` with no value at the instant. + self.absent = set() + # Children given with ``add_child``, in call order. + self.added = {} + # Every child in order, once all of ``source`` is resolved. + self.children = None + + class ParameterNodeAtInstant: """ Parameter node of the legislation, at a given instant. @@ -41,9 +61,11 @@ class ParameterNodeAtInstant: descendant up front made each instant a model is asked about cost a copy of the whole parameter tree (16 to 33 MB and several seconds per instant for policyengine-us), which every tax-benefit system then kept for its - lifetime. A child that has not been read yet takes the value its parameter - has when it is first read; a child read once, including one found to have - no value at this instant, keeps what that read found. + lifetime. Which children a node has, and in what order, is fixed when it is + built. A child's value is taken when the child is first read, so a node + held across an in-place change to its parameters reads the changed value + for a child it has not read yet; once read (including found to have no + value), a child keeps what that read found. """ def __init__(self, name: str, node: "ParameterNode", instant_str: str): @@ -56,51 +78,75 @@ def __init__(self, name: str, node: "ParameterNode", instant_str: str): # The "technical" attributes are hidden, so that the node children can be easily browsed with auto-completion without pollution self._name = name self._instant_str = instant_str - state = self.__dict__ - # The parameter node children still resolve from; None once every - # child is resolved. - state[_NODE] = node - state[_INSTANT] = instant_str - # Children resolved so far. Only ever grows, and a shallow copy shares - # it, so every copy reads the same child objects. - state[_RESOLVED] = {} - # Children of the node that have no value at this instant. - state[_ABSENT] = set() - # A child whose name the node itself uses (``add_child``, ``_name``, - # ...) shadows that name, as when every child was set here. - for child_name in _SHADOWING_NAMES: - if child_name in node.children: - _resolve(self, child_name) + source = dict(node.children) + if _STATE in source: + raise ValueError(f"{_STATE!r} cannot name a parameter.") + state = _LazyChildren(instant_str, source) + self.__dict__[_STATE] = state + cls = type(self) + if cls.add_child is not ParameterNodeAtInstant.add_child: + # A subclass's add_child sees every child, as before. + for child_name, child in source.items(): + child_at_instant = child._get_at_instant(instant_str) + if child_at_instant is not None: + self.add_child(child_name, child_at_instant) + with _RESOLUTION_LOCK: + state.source = None + state.children = dict(state.added) + return + for child_name in source: + # A child whose name the node itself answers to shadows it, as + # when every child was set here. + if child_name in _OWN_NAMES or hasattr(cls, child_name): + child_at_instant = _resolve(state, child_name) + if child_at_instant is not _MISSING: + setattr(self, child_name, child_at_instant) def add_child(self, child_name: str, child_at_instant: "ParameterNodeAtInstant"): + if child_name == _STATE: + raise ValueError(f"{_STATE!r} cannot name a parameter.") + state = self.__dict__[_STATE] with _RESOLUTION_LOCK: - _store(self, child_name, child_at_instant) + state.added[child_name] = child_at_instant + if state.children is not None: + state.children[child_name] = child_at_instant + setattr(self, child_name, child_at_instant) def __getattr__(self, key: str): # Reached only for a name that is not an instance attribute: a child - # not read yet, ``_children`` before every child is resolved, or a name - # this node does not have. - if key.startswith("__") and key.endswith("__"): - # Protocol lookups (copy, pickle, numpy) are never parameters. - raise AttributeError(key) - if _RESOLVED not in self.__dict__ or key == "_name": + # not read yet by this instance, ``_children`` before it is set, or a + # name this node does not have. + state = self.__dict__.get(_STATE) + if state is None or key == "_name": # An instance still being rebuilt by copy or pickle. raise AttributeError(key) if key == "_children": - return _materialise(self) - child_at_instant = _resolve(self, key) + return _materialise(self, state) + child_at_instant = _resolve(state, key) if child_at_instant is not _MISSING: + # Attributes hold what the parameters give; ``add_child`` on + # another copy changes only that copy's attribute, as before. + self.__dict__.setdefault(key, child_at_instant) return child_at_instant param_name = helpers._compose_name(self._name, item_name=key) raise ParameterNotFoundError(param_name, self._instant_str) def __dir__(self): - self._children - return [ - name - for name in super().__dir__() - if not name.startswith("_ParameterNodeAtInstant__") - ] + children = self._children + names = set(super().__dir__()) | set(map(str, children)) + names.discard(_STATE) + return list(names) + + def __setstate__(self, state: dict): + self.__dict__.update(state) + if _STATE not in state: + # Pickled when every child was built up front, all in ``_children``. + children = state.get("_children") + if children is None: + children = self.__dict__["_children"] = {} + lazy = _LazyChildren(state.get("_instant_str"), None) + lazy.resolved = lazy.children = children + self.__dict__[_STATE] = lazy def __getitem__( self, key: str @@ -122,11 +168,14 @@ def __getitem__( ) self._vectorial_node = vectorial return vectorial[key] - child_at_instant = self.__dict__[_RESOLVED].get(key, _MISSING) + state = self.__dict__[_STATE] + child_at_instant = state.added.get(key, _MISSING) if child_at_instant is _MISSING: - child_at_instant = _resolve(self, key) + child_at_instant = state.resolved.get(key, _MISSING) if child_at_instant is _MISSING: - raise KeyError(key) + child_at_instant = _resolve(state, key) + if child_at_instant is _MISSING: + raise KeyError(key) return child_at_instant def __iter__(self) -> Iterable: @@ -142,77 +191,69 @@ def __repr__(self) -> str: return result -def _store(node_at_instant, child_name, child_at_instant): - """Record a child. The caller holds ``_RESOLUTION_LOCK``.""" - state = node_at_instant.__dict__ - state[_RESOLVED][child_name] = child_at_instant - state[_ABSENT].discard(child_name) - children = state.get(_CHILDREN) - if children is not None: - children[child_name] = child_at_instant - setattr(node_at_instant, child_name, child_at_instant) - - -def _resolve(node_at_instant, child_name): - """Return one child, resolving it if no read has yet; ``_MISSING`` if the - node has no such child or it has no value at this instant.""" - state = node_at_instant.__dict__ +def _resolve(state: _LazyChildren, child_name): + """A child as the parameters give it, resolving it if no read has yet; + ``_MISSING`` if the node has no such child or it has no value at the + instant.""" with _RESOLUTION_LOCK: - resolved = state[_RESOLVED] - if child_name in resolved: - return resolved[child_name] - if child_name in state[_ABSENT]: - return _MISSING - node = state[_NODE] - child = None if node is None else node.children.get(child_name) + child_at_instant = state.resolved.get(child_name, _MISSING) + if child_at_instant is not _MISSING or child_name in state.absent: + return child_at_instant + source = state.source + child = None if source is None else source.get(child_name) if child is None: return _MISSING - child_at_instant = child._get_at_instant(state[_INSTANT]) - if child_at_instant is None: - state[_ABSENT].add(child_name) + child_at_instant = child._get_at_instant(state.instant) + # A read made while resolving (the same thread re-entering) may have + # recorded this child already; the first recorded result stands. + if child_name in state.absent: + return _MISSING + recorded = state.resolved.get(child_name, _MISSING) + if recorded is not _MISSING: + child_at_instant = recorded + elif child_at_instant is None: + state.absent.add(child_name) else: - _store(node_at_instant, child_name, child_at_instant) - if len(resolved) + len(state[_ABSENT]) >= len(node.children): - # Every child has been read: drop the parameter node. - _materialise(node_at_instant) + state.resolved[child_name] = child_at_instant + if state.source is not None and len(state.resolved) + len(state.absent) == len( + state.source + ): + _complete(state) return _MISSING if child_at_instant is None else child_at_instant -def _materialise(node_at_instant): - """Resolve every child; return them in the parameter node's order, with - any added by ``add_child`` after them.""" - state = node_at_instant.__dict__ - with _RESOLUTION_LOCK: - children = state.get(_CHILDREN) - if children is not None: - return children - node = state[_NODE] - children = {} - for child_name in node.children if node is not None else (): - child_at_instant = _resolve(node_at_instant, child_name) - if child_at_instant is not _MISSING: - children[child_name] = child_at_instant - if state.get(_CHILDREN) is not None: - # Resolving the last child completed the node already. - return state[_CHILDREN] - for child_name, child_at_instant in state[_RESOLVED].items(): - children.setdefault(child_name, child_at_instant) - state[_CHILDREN] = children - if "_children" not in state: - # Unless a child named ``_children`` took the name. - state["_children"] = children - # Every child is resolved: the parameter node is no longer needed. - state[_NODE] = None - return children - - -# Names a ParameterNodeAtInstant itself answers to, which a child of the same -# name shadows. -_SHADOWING_NAMES = frozenset( - {"_name", "_instant_str", "_children", "_vectorial_node"} - | { - name - for name in dir(ParameterNodeAtInstant) - if not name.startswith("__") and not name.startswith("_ParameterNodeAtInstant") +def _complete(state: _LazyChildren): + """Order every child once all of ``source`` is resolved; let go of the + parameter nodes. The caller holds ``_RESOLUTION_LOCK``.""" + resolved, added = state.resolved, state.added + children = { + name: added.get(name, resolved[name]) + for name in state.source + if name in resolved } -) + for name, child_at_instant in added.items(): + # Added children the parameters do not have (or have no value for) + # follow, in the order they were added. + children.setdefault(name, child_at_instant) + if not added: + state.resolved = children + state.children = children + state.source = None + state.absent = set() + + +def _materialise(node_at_instant: ParameterNodeAtInstant, state: _LazyChildren): + """Resolve every child and set ``_children``.""" + with _RESOLUTION_LOCK: + if state.children is None: + for child_name in list(state.source): + if state.children is not None: + break + _resolve(state, child_name) + if state.children is None: + _complete(state) + children = state.children + if "_children" not in node_at_instant.__dict__: + # Unless a child named ``_children`` took the name. + node_at_instant.__dict__["_children"] = children + return children diff --git a/tests/core/parameters/test_parameter_node_at_instant_lazy.py b/tests/core/parameters/test_parameter_node_at_instant_lazy.py index ff2604e02..6ef76d999 100644 --- a/tests/core/parameters/test_parameter_node_at_instant_lazy.py +++ b/tests/core/parameters/test_parameter_node_at_instant_lazy.py @@ -19,7 +19,7 @@ from policyengine_core.country_template import CountryTaxBenefitSystem from policyengine_core.errors import ParameterNotFoundError from policyengine_core.parameters import ParameterNode, ParameterNodeAtInstant -from policyengine_core.parameters.parameter_node_at_instant import _RESOLVED +from policyengine_core.parameters.parameter_node_at_instant import _STATE from policyengine_core.tracers import FullTracer from tests.fixtures.parameter_nodes_at_instant import ( INSTANTS, @@ -49,7 +49,7 @@ def test_country_template_matches_the_up_front_build(instant): def resolved(node_at_instant): - return list(node_at_instant.__dict__[_RESOLVED]) + return list(node_at_instant.__dict__[_STATE].resolved) def test_nothing_is_resolved_until_it_is_read(): @@ -363,3 +363,150 @@ def test_a_node_lets_go_of_its_parameter_node_once_every_child_is_read(): gc.collect() assert tree_ref() is None assert group.x == 1 and list(group) == ["x"] + + +def plain_tree(**values): + return ParameterNode( + "", data={name: {"values": {"2010-01-01": v}} for name, v in values.items()} + ) + + +def test_the_reserved_state_name_is_rejected(): + with pytest.raises(ValueError): + plain_tree(_ParameterNodeAtInstant__state=1).get_at_instant("2017-01-01") + at_instant = plain_tree(x=1).get_at_instant("2017-01-01") + with pytest.raises(ValueError): + at_instant.add_child("_ParameterNodeAtInstant__state", 2) + + +@pytest.mark.parametrize( + "child_name", + [ + "_ParameterNodeAtInstant__resolved", + "_ParameterNodeAtInstant__node", + "_ParameterNodeAtInstant__business", + "__custom__", + "__iter__", + "__dir__", + ], +) +def test_other_unusual_names_are_plain_children(child_name): + at_instant = plain_tree(**{child_name: 42, "ok": 7}).get_at_instant("2017-01-01") + assert getattr(at_instant, child_name) == 42 + assert at_instant[child_name] == 42 + assert at_instant.ok == 7 and at_instant["ok"] == 7 + assert child_name in dir(at_instant) + + +def test_many_added_children_do_not_break_the_first_read(): + names = {f"p{i}": i for i in range(2000)} + at_instant = plain_tree(**names).get_at_instant("2017-01-01") + for i in range(2000): + at_instant.add_child(f"extra{i}", i) + assert at_instant.p0 == 0 + assert len(at_instant._children) == 4000 + assert list(at_instant._children)[:2] == ["p0", "p1"] + + +def test_add_child_on_a_copy_changes_only_shared_items_as_before(): + tree = ParameterNode( + "", + data={ + "x": {"values": {"2010-01-01": 10}}, + "late": {"values": {"2030-01-01": 5}}, + }, + ) + held = tree.get_at_instant("2017-01-01") + with pytest.raises(ParameterNotFoundError): + held.late + duplicate = copy.copy(held) + duplicate.add_child("x", 99) + duplicate.add_child("late", 88) + duplicate.add_child("extra", 77) + # Items are one mapping shared by copies; attributes belong to each copy. + assert held["x"] == 99 and held.x == 10 + assert held["late"] == 88 + with pytest.raises(ParameterNotFoundError): + held.late + assert held["extra"] == 77 + with pytest.raises(ParameterNotFoundError): + held.extra + assert duplicate.x == 99 and duplicate.late == 88 and duplicate.extra == 77 + assert list(held._children) == ["x", "late", "extra"] + + +def test_a_copy_installs_what_it_reads_as_attributes(): + held = build_tree().get_at_instant("2017-03-01") + duplicate = copy.copy(held) + scale = held.scale + assert duplicate.scale is scale + assert duplicate.__dict__["scale"] is scale + + +class _MethodCollision(ParameterNodeAtInstant): + def policy(self): + return "method" + + +class _DoublingAddChild(ParameterNodeAtInstant): + def add_child(self, child_name, child_at_instant): + super().add_child(child_name, child_at_instant * 2) + + +def test_subclasses_keep_their_collisions_and_add_child_hooks(): + tree = plain_tree(policy=42, x=5) + collided = _MethodCollision("", tree, "2017-01-01") + assert collided.policy == 42 and collided["policy"] == 42 + doubled = _DoublingAddChild("", tree, "2017-01-01") + assert doubled.x == 10 and doubled["x"] == 10 + assert dict(doubled._children) == {"policy": 84, "x": 10} + + +class ReenteringChild: + """A child whose first resolution reads the same child again.""" + + def __init__(self): + self.node_at_instant = None + self.calls = 0 + + def _get_at_instant(self, instant): + self.calls += 1 + result = object() + if self.calls == 1: + self.inner = self.node_at_instant["again"] + return result + + +def test_a_reentrant_read_keeps_the_first_recorded_child(): + tree = plain_tree(other=1) + child = ReenteringChild() + tree.children = {"again": child, **tree.children} + at_instant = tree.get_at_instant("2017-01-01") + child.node_at_instant = at_instant + outer = at_instant["again"] + assert outer is child.inner + assert at_instant["again"] is outer and at_instant.again is outer + + +def test_children_and_their_order_are_fixed_when_the_node_is_built(): + tree = plain_tree(x=10, y=20, z=30) + held = tree.get_at_instant("2017-01-01") + assert held.x == 10 + del tree.children["x"] + tree.children["w"] = ParameterNode("", data={"v": {"values": {"2010-01-01": 1}}}) + assert list(held) == ["x", "y", "z"] + assert held.x == 10 + assert "w" not in held._children + + +def test_a_node_pickled_when_every_child_was_built_up_front_loads(): + """State as the up-front build pickled it: everything in ``_children``.""" + restored = ParameterNodeAtInstant.__new__(ParameterNodeAtInstant) + restored.__setstate__( + {"_name": "", "_instant_str": "2017-01-01", "_children": {"x": 10}, "x": 10} + ) + assert restored.x == 10 and restored["x"] == 10 and list(restored) == ["x"] + restored.add_child("y", 2) + assert restored["y"] == 2 and list(restored) == ["x", "y"] + with pytest.raises(KeyError): + restored["nope"] From 089761c5da80b61056ecaab85325834028b31dcc Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 13 Oct 2026 23:25:49 +0400 Subject: [PATCH 14/23] Let lazy parameter nodes at instant pickle under protocols 0 and 1 The lazy state object has __slots__ and no __getstate__, which pickle protocols 0 and 1 refuse. Master's #584 tests pickle a tree holding cached nodes at instant under those protocols, and failed on this branch once it was rebased onto master. A partly read node now round-trips under every protocol. Co-Authored-By: Claude Opus 5.5 --- .../parameters/parameter_node_at_instant.py | 9 +++++++++ .../test_parameter_node_at_instant_lazy.py | 12 ++++++++++++ 2 files changed, 21 insertions(+) diff --git a/policyengine_core/parameters/parameter_node_at_instant.py b/policyengine_core/parameters/parameter_node_at_instant.py index 43787a045..ceafdd58a 100644 --- a/policyengine_core/parameters/parameter_node_at_instant.py +++ b/policyengine_core/parameters/parameter_node_at_instant.py @@ -52,6 +52,15 @@ def __init__(self, instant: str, source: dict): # Every child in order, once all of ``source`` is resolved. self.children = None + # Pickle protocols 0 and 1 refuse a class with ``__slots__`` that does not + # define these, and a parameter tree pickles the nodes it has cached. + def __getstate__(self) -> dict: + return {name: getattr(self, name) for name in self.__slots__} + + def __setstate__(self, state: dict) -> None: + for name, value in state.items(): + setattr(self, name, value) + class ParameterNodeAtInstant: """ diff --git a/tests/core/parameters/test_parameter_node_at_instant_lazy.py b/tests/core/parameters/test_parameter_node_at_instant_lazy.py index 6ef76d999..6f76623b0 100644 --- a/tests/core/parameters/test_parameter_node_at_instant_lazy.py +++ b/tests/core/parameters/test_parameter_node_at_instant_lazy.py @@ -314,6 +314,18 @@ def test_traced_nodes_copy_and_pickle(duplicate): assert duplicate(node).group.x == 0.1 +@pytest.mark.parametrize("protocol", range(pickle.HIGHEST_PROTOCOL + 1)) +def test_a_partly_read_node_pickles_under_every_protocol(protocol): + """Protocols 0 and 1 refuse an object with ``__slots__`` and no + ``__getstate__``, and the lazy state is one.""" + at_instant = build_tree().get_at_instant("2017-03-01") + assert at_instant.group.x == 0.1 + restored = pickle.loads(pickle.dumps(at_instant, protocol=protocol)) + assert restored.group.x == 0.1 + assert restored.flat_rate == 0.1 # not read before pickling + assert ordered(materialise(restored)) == ordered(materialise(at_instant)) + + def test_a_child_without_a_value_is_resolved_once(): tree = ParameterNode("", data={"late": {"values": {"2030-01-01": 1}}}) held = tree.get_at_instant("2020-01-01") From d5010092c4645cc1463f0210743d3cce2cfc04b7 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 13 Oct 2026 22:22:27 +0400 Subject: [PATCH 15/23] Copy only set parameters for YAML inputs, add read-only arrays and forked batches A YAML case's dotted parameter inputs now run on a reform that copies only the nodes on the paths to the parameters it sets and shares the rest of the tree, with its cached at-instant values, with the system the case's reforms select. Before, each distinct set of values built a full clone of the baseline plus a second full copy of its parameter tree, about 4.5 s and 1.2 GB on policyengine-us, and evicted the structural reform systems the cache holds. The runner's `read_only_arrays` option (`--read-only-arrays`) stores every array read-only for the run, so a formula that writes into an array it read from another variable fails at that line. `forked_test_runner.run_batches_in_forks` runs test batches in children forked from one process that has built the tax-benefit system, so batches skip the import and build and share the parent's system in memory. Co-Authored-By: Claude Opus 5.5 --- changelog.d/test-runner-perf.added.md | 1 + docs/engineering/skills/testing.md | 25 +- .../data_storage/in_memory_storage.py | 12 + policyengine_core/parameters/path_copy.py | 92 +++++++ policyengine_core/reforms/reform.py | 11 +- .../scripts/policyengine_command.py | 8 +- policyengine_core/scripts/run_test.py | 1 + policyengine_core/tools/forked_test_runner.py | 228 ++++++++++++++++++ policyengine_core/tools/test_runner.py | 119 +++++---- tests/core/parameters/test_path_copy.py | 179 ++++++++++++++ .../test_runner/test_forked_test_runner.py | 147 +++++++++++ .../test_runner/test_parameter_inputs.py | 149 ++++++++++++ .../test_runner/test_read_only_arrays.py | 148 ++++++++++++ 13 files changed, 1064 insertions(+), 56 deletions(-) create mode 100644 changelog.d/test-runner-perf.added.md create mode 100644 policyengine_core/parameters/path_copy.py create mode 100644 policyengine_core/tools/forked_test_runner.py create mode 100644 tests/core/parameters/test_path_copy.py create mode 100644 tests/core/tools/test_runner/test_forked_test_runner.py create mode 100644 tests/core/tools/test_runner/test_parameter_inputs.py create mode 100644 tests/core/tools/test_runner/test_read_only_arrays.py diff --git a/changelog.d/test-runner-perf.added.md b/changelog.d/test-runner-perf.added.md new file mode 100644 index 000000000..73c92202b --- /dev/null +++ b/changelog.d/test-runner-perf.added.md @@ -0,0 +1 @@ +The YAML test runner sets a case's dotted parameter inputs on a reform that copies only those parameters' paths instead of the whole tree, adds a `--read-only-arrays` option that fails a formula writing into an array it read, and `policyengine_core.tools.forked_test_runner.run_batches_in_forks` runs test batches in children forked from one process that built the tax-benefit system. diff --git a/docs/engineering/skills/testing.md b/docs/engineering/skills/testing.md index ecf1ef527..915c70abe 100644 --- a/docs/engineering/skills/testing.md +++ b/docs/engineering/skills/testing.md @@ -45,11 +45,16 @@ 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. +- A case that names `reforms` or `extensions` 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. +- A case that sets a parameter with a dotted input key runs on a system of its + own (`_ParameterInputReform`) that copies only the nodes on the paths to the + parameters it sets (`policyengine_core.parameters.path_copy`) and shares the + rest of the tree, with its cached at-instant values, with the system its + `reforms` select. Nothing may change a shared node in place. - 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 @@ -66,3 +71,13 @@ 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). + +`policyengine_core.tools.forked_test_runner.run_batches_in_forks` runs batches +in children forked from a process that has already built the system, so a +batch skips the import and build, and the children share the parent's system +in memory until they write to it. `tests/core/tools/test_runner/test_forked_test_runner.py` +covers its exit codes, worker limit and timeout. + +`--read-only-arrays` (option `read_only_arrays`) stores every array read-only +for the run, so a formula that writes into an array it read from another +variable fails at that line (`test_read_only_arrays.py`). diff --git a/policyengine_core/data_storage/in_memory_storage.py b/policyengine_core/data_storage/in_memory_storage.py index 423d19b9a..3aa8ee015 100644 --- a/policyengine_core/data_storage/in_memory_storage.py +++ b/policyengine_core/data_storage/in_memory_storage.py @@ -85,6 +85,14 @@ class InMemoryStorage: # holder per branch for inputs most branches never change. _inputs: Union[Set[str], FrozenSet[str]] = _NO_INPUTS is_eternal: bool + # Whether every storage stores and hands out read-only views of numpy + # arrays, so that code writing into an array it read (``x += y`` on the + # result of ``tax_unit("x", period)``) raises at that line instead of + # changing ``x`` for every later reader. Off by default; the YAML test + # runner turns it on for a run with the ``read_only_arrays`` option. The + # caller's own array stays writeable: only the stored view is protected. + # Masked arrays are left as they are, since a view shares their mask. + read_only_arrays: bool = False def __init__(self, is_eternal: bool): self._arrays = {} @@ -182,6 +190,8 @@ def get(self, period: Period, branch_name: str = "default") -> ArrayLike: # First read of an array shared by ``clone``: from here on the # caller may write into it, so it needs to be this storage's own. values = values.copy() + if self.read_only_arrays: + values = _read_only_view(values) self._arrays[key] = values self._stop_sharing(key) return values @@ -304,6 +314,8 @@ def put( "string form is lossy (see policyengine-core#526)." ) key = f"{branch_name}:{period}" + if self.read_only_arrays and _can_share(value): + value = _read_only_view(value) self._arrays[key] = value self._stop_sharing(key) if derived: diff --git a/policyengine_core/parameters/path_copy.py b/policyengine_core/parameters/path_copy.py new file mode 100644 index 000000000..3e1987264 --- /dev/null +++ b/policyengine_core/parameters/path_copy.py @@ -0,0 +1,92 @@ +"""Copy only the parts of a parameter tree that a change will touch.""" + +from typing import Iterable, Set, Union + +from policyengine_core import commons + +from .parameter import Parameter +from .parameter_node import ParameterNode +from .parameter_scale import ParameterScale + +TreeNode = Union[ParameterNode, ParameterScale, Parameter] + + +def copy_paths(root: ParameterNode, paths: Iterable[str]) -> ParameterNode: + """Return a tree that shares every node of ``root`` except those on ``paths``. + + Each path names a parameter as ``set_parameter`` and + ``ParameterNode.get_child`` do: dotted child names, with ``name[i]`` for + bracket ``i`` of a scale. The returned root, every node between it and each + named parameter, and each named parameter are new; every other node, with + the at-instant values it has cached, is ``root``'s own. + + A full ``clone`` rebuilds all of a large tree (policyengine-us has over + 130,000 nodes) and drops every cached at-instant value, which costs + seconds and a copy of the tree. Copying the paths instead costs a few + nodes per path. + + Change the returned tree only at the named parameters, for example with + ``Parameter.update``: an update clears the at-instant caches of the new + nodes above it and marks them modified, and stops there. A node shared + with ``root`` still names its parent in ``root``'s tree, so changing it + would change ``root``'s tree too. + + Raises the ``ValueError`` of ``ParameterNode.get_child`` if a path names + no parameter, before copying anything. + """ + paths = list(paths) + for path in paths: + root.get_child(path) + + copied_root = _copy_node(root) + copies = {id(copied_root)} + for path in paths: + node = copied_root + for name in path.split("."): + if "[" in name: + scale_name, index = name.split("[") + index = int(index[:-1]) + scale = _copied_child(node, scale_name, copies) + bracket = scale.brackets[index] + if id(bracket) not in copies: + bracket = _copy_node(bracket) + copies.add(id(bracket)) + scale.brackets[index] = bracket + bracket.parent = scale + node = bracket + else: + node = _copied_child(node, name, copies) + return copied_root + + +def _copied_child(node: ParameterNode, name: str, copies: Set[int]) -> TreeNode: + """``node``'s child ``name``, replaced by a copy unless it already is one.""" + child = node.children[name] + if id(child) in copies: + return child + child = _copy_node(child) + copies.add(id(child)) + node.children[name] = child + setattr(node, name, child) + child.parent = node + return child + + +def _copy_node(node: TreeNode) -> TreeNode: + """A copy of ``node`` that shares its children and their caches. + + A parameter is cloned outright, as it has no children to share. A node or + scale gets its own children mapping or bracket list, holding the same + children, and a node an empty at-instant cache, since a value below it is + about to change. + """ + if isinstance(node, Parameter): + return node.clone() + copied = commons.empty_clone(node) + copied.__dict__ = node.__dict__.copy() + if isinstance(node, ParameterScale): + copied.brackets = list(node.brackets) + else: + copied.children = dict(node.children) + copied._at_instant_cache = {} + return copied diff --git a/policyengine_core/reforms/reform.py b/policyengine_core/reforms/reform.py index 9358ee5c1..8c5200ac5 100644 --- a/policyengine_core/reforms/reform.py +++ b/policyengine_core/reforms/reform.py @@ -157,7 +157,7 @@ def __init__(self, baseline: TaxBenefitSystem): # deep-copying, but plenty of reforms mutate ``self.parameters`` # directly (bug C4). if baseline.parameters is not None: - self.parameters = baseline.parameters.clone() + self.parameters = self._copy_parameters(baseline.parameters) else: self.parameters = None # Use a fresh at-instant cache so reform mutations aren't shadowed @@ -172,6 +172,15 @@ def __init__(self, baseline: TaxBenefitSystem): ) self.apply() + def _copy_parameters(self, baseline_parameters: ParameterNode) -> ParameterNode: + """The tree ``apply`` starts from and may change in place. + + A full clone of the baseline's tree, since a reform can change any + node. A reform that changes only parameters it can name in advance can + copy less (see ``policyengine_core.parameters.path_copy``). + """ + return baseline_parameters.clone() + def __getattr__(self, attribute): return getattr(self.baseline, attribute) diff --git a/policyengine_core/scripts/policyengine_command.py b/policyengine_core/scripts/policyengine_command.py index 2f90aa746..94835d323 100644 --- a/policyengine_core/scripts/policyengine_command.py +++ b/policyengine_core/scripts/policyengine_command.py @@ -102,7 +102,13 @@ def build_test_parser(parser): "--reform-cache-size", 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.", + help="how many reform systems (built for tests with reforms or extensions) to keep cached; each is a full copy of the tax-benefit system. Defaults to 2; 0 caches none.", + ) + parser.add_argument( + "--read-only-arrays", + action="store_true", + default=False, + help="store every calculated and input array read-only, so a formula that writes into an array it read from another variable fails at that line instead of changing that variable for later readers", ) return parser diff --git a/policyengine_core/scripts/run_test.py b/policyengine_core/scripts/run_test.py index 123f49fa0..2fcb40c74 100644 --- a/policyengine_core/scripts/run_test.py +++ b/policyengine_core/scripts/run_test.py @@ -31,6 +31,7 @@ def main(parser): "ignore_variables": args.ignore_variables, "visualize": args.visualize, "reform_cache_size": args.reform_cache_size, + "read_only_arrays": args.read_only_arrays, } paths = [os.path.abspath(path) for path in args.path] diff --git a/policyengine_core/tools/forked_test_runner.py b/policyengine_core/tools/forked_test_runner.py new file mode 100644 index 000000000..7e97de2d8 --- /dev/null +++ b/policyengine_core/tools/forked_test_runner.py @@ -0,0 +1,228 @@ +"""Run batches of YAML tests in processes forked from one that built the system. + +Starting a process for each batch of a large country suite repeats the same +work every time: policyengine-us takes about 40 seconds to import and build +its tax-benefit system before the first case runs, and over half a minute to +tear its interpreter down afterwards. Forking from a process that has already +built the system skips both, and the children share the parent's copy of the +system in memory until they write to it, so several batches can run at once +in the memory one used to take. +""" + +import gc +import os +import signal +import sys +import tempfile +import time +import traceback +from collections import deque +from dataclasses import dataclass +from typing import Callable, Dict, List, Optional, Sequence + +from policyengine_core.tools.test_runner import _get_tax_benefit_system, run_tests + +# pytest's exit code for an internal error, used when a child fails outside +# pytest (for example, opening its log). +PYTEST_INTERNAL_ERROR = 3 +# How long a child gets to exit after SIGTERM before it is killed. +TERMINATE_GRACE_SECONDS = 5 +# How often the parent checks on its children. +POLL_SECONDS = 0.1 + + +@dataclass +class BatchResult: + """How one batch ended. + + ``exit_code`` is pytest's exit code, negative for a child ended by a + signal (as ``subprocess`` reports it). ``peak_rss_mb`` is the child's + peak resident set, which counts the pages it still shares with the parent. + """ + + paths: List[str] + exit_code: int + elapsed: float + peak_rss_mb: Optional[float] + output: str + timed_out: bool = False + + +def run_batches_in_forks( + tax_benefit_system, + batches: Sequence[Sequence[str]], + options: Optional[dict] = None, + workers: int = 1, + timeout: Optional[float] = None, + on_finish: Optional[Callable[[int, BatchResult], None]] = None, + on_heartbeat: Optional[Callable[[Dict[int, float]], None]] = None, + heartbeat_seconds: float = 60, +) -> List[BatchResult]: + """Run each batch of test paths in its own child process, ``workers`` at once. + + Each child is forked from this process, runs ``run_tests`` on its batch + with ``tax_benefit_system`` and ``options``, and exits as soon as pytest + returns. Its output goes to a file and is returned whole in its result, + so concurrent batches never interleave. Batches start in the order given. + + ``timeout`` bounds each batch in seconds: a batch still running then is + sent SIGTERM, then SIGKILL, and reported with ``timed_out``. + ``on_finish(index, result)`` is called as each batch ends, and + ``on_heartbeat({index: seconds running})`` every ``heartbeat_seconds`` + while any runs. + + Call this from a process with no other threads: a thread holding a lock + when the process forks leaves that lock held forever in the child. That + includes a numerical library's thread pool, so set + ``OPENBLAS_NUM_THREADS=1`` (and ``OMP_NUM_THREADS``, ``MKL_NUM_THREADS``) + before the process imports numpy. + + Before forking, this builds the reform-free system the runner gives every + case without reforms (see ``_get_tax_benefit_system``), so the children + share it rather than each building their own, and moves every object the + parent holds to the garbage collector's permanent generation + (``gc.freeze``), so a child's collections do not write to, and copy, the + pages it shares with the parent. + """ + if not hasattr(os, "fork"): + raise NotImplementedError("Running batches in forks needs os.fork.") + if workers < 1: + raise ValueError(f"workers must be 1 or more, got {workers}") + if options is None: + options = {} + + _get_tax_benefit_system(tax_benefit_system, [], []) + + pending = deque(enumerate(batches)) + running: Dict[int, tuple] = {} + results: List[Optional[BatchResult]] = [None] * len(batches) + last_heartbeat = time.monotonic() + log_directory = tempfile.mkdtemp(prefix="policyengine-batches-") + gc.collect() + gc.freeze() + try: + while pending or running: + while pending and len(running) < workers: + index, paths = pending.popleft() + log_path = os.path.join(log_directory, f"batch-{index}.log") + sys.stdout.flush() + sys.stderr.flush() + pid = os.fork() + if pid == 0: + _run_child(tax_benefit_system, list(paths), options, log_path) + running[pid] = (index, list(paths), log_path, time.monotonic(), None) + + any_finished = False + for pid, (index, paths, log_path, started, terminated_at) in list( + running.items() + ): + waited_pid, status, usage = os.wait4(pid, os.WNOHANG) + now = time.monotonic() + if waited_pid == 0: + if terminated_at is None: + if timeout is not None and now - started >= timeout: + os.kill(pid, signal.SIGTERM) + running[pid] = (index, paths, log_path, started, now) + elif now - terminated_at >= TERMINATE_GRACE_SECONDS: + os.kill(pid, signal.SIGKILL) + continue + any_finished = True + del running[pid] + result = BatchResult( + paths=paths, + exit_code=os.waitstatus_to_exitcode(status), + elapsed=now - started, + peak_rss_mb=_peak_rss_mb(usage), + output=_read_and_remove(log_path), + timed_out=terminated_at is not None, + ) + results[index] = result + if on_finish is not None: + on_finish(index, result) + + now = time.monotonic() + if ( + on_heartbeat is not None + and running + and (now - last_heartbeat >= heartbeat_seconds) + ): + on_heartbeat( + { + index: now - started + for index, _, _, started, _ in running.values() + } + ) + last_heartbeat = now + if not any_finished: + time.sleep(POLL_SECONDS) + finally: + # Reached with children still running only if the parent was + # interrupted: do not leave them behind. + for pid in running: + try: + os.kill(pid, signal.SIGKILL) + os.waitpid(pid, 0) + except (ChildProcessError, ProcessLookupError): + pass + gc.unfreeze() + for name in os.listdir(log_directory): + os.remove(os.path.join(log_directory, name)) + os.rmdir(log_directory) + return results + + +def _run_child(tax_benefit_system, paths, options, log_path) -> None: + """Run one batch in a forked child, then exit without returning.""" + exit_code = PYTEST_INTERNAL_ERROR + try: + log = os.open(log_path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o644) + # The standard descriptors, and whatever ``sys.stdout`` and + # ``sys.stderr`` write to if something (pytest's capture) replaced them. + for descriptor in {1, 2, *_stream_descriptors()}: + os.dup2(log, descriptor) + os.close(log) + exit_code = int(run_tests(tax_benefit_system, paths, options)) + except BaseException: + traceback.print_exc() + finally: + try: + sys.stdout.flush() + sys.stderr.flush() + finally: + # Skip interpreter teardown: freeing a large country's systems + # one object at a time took over half a minute, and nothing the + # child holds needs releasing before the process goes. + os._exit(exit_code) + + +def _stream_descriptors() -> List[int]: + descriptors = [] + for stream in (sys.stdout, sys.stderr): + try: + descriptors.append(stream.fileno()) + except (AttributeError, OSError, ValueError): + pass + return descriptors + + +def _peak_rss_mb(usage) -> Optional[float]: + """A finished child's peak resident set in MB, from ``os.wait4``.""" + if usage is None or not usage.ru_maxrss: + return None + # Linux reports kilobytes and macOS bytes. + if sys.platform == "darwin": + return usage.ru_maxrss / 1e6 + return usage.ru_maxrss / 1e3 + + +def _read_and_remove(log_path: str) -> str: + try: + with open(log_path, encoding="utf-8", errors="replace") as log: + return log.read() + except FileNotFoundError: + return "" + finally: + try: + os.remove(log_path) + except FileNotFoundError: + pass diff --git a/policyengine_core/tools/test_runner.py b/policyengine_core/tools/test_runner.py index c8df6c8cf..12b2b36ed 100644 --- a/policyengine_core/tools/test_runner.py +++ b/policyengine_core/tools/test_runner.py @@ -16,12 +16,14 @@ eval_expression, ) from policyengine_core.simulations import SimulationBuilder +from policyengine_core.data_storage import InMemoryStorage from policyengine_core.errors import ( SituationParsingError, VariableNotFoundError, ) from policyengine_core.scripts import build_tax_benefit_system from policyengine_core.reforms import Reform, set_parameter +from policyengine_core.parameters.path_copy import copy_paths from policyengine_core.populations import ADD, DIVIDE log = logging.getLogger(__name__) @@ -71,12 +73,14 @@ def import_yaml(): _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. +# How many reform systems (cases with ``reforms`` or ``extensions``) 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. Dotted parameter inputs do not use +# the cache: each such case gets a ``_ParameterInputReform`` of its own, which +# shares nearly all of its tree with the cached system it is built on. DEFAULT_REFORM_CACHE_SIZE = 2 @@ -105,15 +109,20 @@ def run_tests(tax_benefit_system, paths, options=None): +-------------------------------+-----------+-------------------------------------------+ | reform_cache_size | ``int`` | Reform systems kept cached (default 2) | +-------------------------------+-----------+-------------------------------------------+ + | read_only_arrays | ``bool`` | Store cached arrays read-only, so a | + | | | formula writing into one it read fails | + +-------------------------------+-----------+-------------------------------------------+ - **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. Each system still caches its - parameters at every distinct instant the cases read. + **Memory**: a case that names ``reforms`` or ``extensions`` 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. + A case that sets a parameter with a dotted input key runs on a system of + its own that copies only the paths to the parameters it sets and shares + everything else with the system its reforms select. Each case releases its + simulation and that system when it finishes, so memory does not grow with + the number of cases or of distinct combinations. Each system still caches + its parameters at every distinct instant the cases read. """ @@ -131,10 +140,16 @@ def run_tests(tax_benefit_system, paths, options=None): if options is None: options = {} - return pytest.main( - [*argv, *paths], - plugins=[OpenFiscaPlugin(tax_benefit_system, options)], - ) + previously_read_only = InMemoryStorage.read_only_arrays + if options.get("read_only_arrays"): + InMemoryStorage.read_only_arrays = True + try: + return pytest.main( + [*argv, *paths], + plugins=[OpenFiscaPlugin(tax_benefit_system, options)], + ) + finally: + InMemoryStorage.read_only_arrays = previously_read_only class YamlFile(pytest.File): @@ -228,51 +243,29 @@ def runtest(self): unsafe_input = {} period = self.test.get("period") input = {} - inline_reforms = [] - parametric_reform_items = [] + parameter_inputs = [] for key, value in unsafe_input.items(): if "." in key: - inline_reforms += [ - set_parameter( - key, - value, - return_modifier=True, - period=f"year:2000:40", - ) - ] - parametric_reform_items.append((key, value)) + parameter_inputs.append((key, value)) else: input[key] = value - if len(inline_reforms) == 0: - inline_reform = [] - else: - - class inline_reform_class(Reform): - def apply(self): - for modifier in inline_reforms: - self.parameters = modifier(self.parameters) - - inline_reform = [inline_reform_class] - reforms = self.test.get("reforms", []) if isinstance(reforms, str): reforms = [reforms] - if not inline_reforms: - inline_reform = [] - self.tax_benefit_system = _get_tax_benefit_system( self.baseline_tax_benefit_system, - reforms + inline_reform, + reforms, 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!r}" for key, value in parametric_reform_items] - ), cache_size=self.options.get("reform_cache_size"), ) + if parameter_inputs: + # Built for this case alone, and cheap: it shares every parameter + # it does not set with the cached system it is built on. + self.tax_benefit_system = _ParameterInputReform( + self.tax_benefit_system, parameter_inputs + ) verbose = self.options.get("verbose") performance_graph = self.options.get("performance_graph") performance_tables = self.options.get("performance_tables") @@ -483,6 +476,34 @@ def __init__(self): self.reforms = OrderedDict() +class _ParameterInputReform(Reform): + """Sets a case's dotted parameter inputs on top of ``baseline``. + + Each input sets its parameter from 2000 for 40 years, as + ``set_parameter`` does. The reform copies only the nodes on the paths of + the parameters it sets (see ``policyengine_core.parameters.path_copy``), + so it shares the rest of ``baseline``'s tree, including the at-instant + values ``baseline`` has cached. A full copy of policyengine-us's tree took + seconds and over half a gigabyte for every distinct set of values, and + most cases that set parameters use values no other case does. + """ + + def __init__(self, baseline, parameter_inputs): + self.parameter_inputs = list(parameter_inputs) + super().__init__(baseline) + + def _copy_parameters(self, baseline_parameters): + return copy_paths( + baseline_parameters, [key for key, _ in self.parameter_inputs] + ) + + def apply(self): + for key, value in self.parameter_inputs: + set_parameter(key, value, period="year:2000:40", return_modifier=True)( + self.parameters + ) + + def _get_tax_benefit_system( baseline, reforms, diff --git a/tests/core/parameters/test_path_copy.py b/tests/core/parameters/test_path_copy.py new file mode 100644 index 000000000..67a30b007 --- /dev/null +++ b/tests/core/parameters/test_path_copy.py @@ -0,0 +1,179 @@ +"""``copy_paths`` copies only the nodes on the paths it is given.""" + +import pytest +from hypothesis import given, settings +from hypothesis import strategies as st + +from policyengine_core.country_template import CountryTaxBenefitSystem +from policyengine_core.errors import ParameterNotFoundError +from policyengine_core.parameters import Parameter +from policyengine_core.parameters.path_copy import copy_paths +from policyengine_core.reforms import set_parameter + +SYSTEM = CountryTaxBenefitSystem() +INSTANTS = ["2012-06-01", "2015-01-01", "2017-01-01", "2020-01-01", "2035-01-01"] + + +def _leaf_values(tree, instant): + """Every parameter's value at ``instant``, by name.""" + return { + descendant.name: descendant(instant) + for descendant in tree.get_descendants() + if isinstance(descendant, Parameter) + } + + +def _set(tree, path, value): + return set_parameter(path, value, period="year:2000:40", return_modifier=True)(tree) + + +def _leaf_paths(tree): + """A path to every parameter of ``tree`` in ``set_parameter``'s syntax.""" + paths = [] + + def visit(node, prefix): + if isinstance(node, Parameter): + paths.append(prefix) + elif hasattr(node, "brackets"): + for index, bracket in enumerate(node.brackets): + visit(bracket, f"{prefix}[{index}]") + else: + for name, child in node.children.items(): + visit(child, f"{prefix}.{name}" if prefix else name) + + visit(tree, "") + return [path for path in paths if not path.startswith("baseline.")] + + +LEAF_PATHS = _leaf_paths(SYSTEM.parameters) + + +def test_only_the_nodes_on_the_path_are_new(): + original = SYSTEM.parameters.clone() + copied = copy_paths(original, ["taxes.income_tax_rate"]) + + assert copied is not original + assert copied.taxes is not original.taxes + assert copied.taxes.income_tax_rate is not original.taxes.income_tax_rate + assert copied.benefits is original.benefits + assert copied.taxes.housing_tax is original.taxes.housing_tax + assert copied.children["taxes"] is copied.taxes + assert copied.taxes.parent is copied + assert copied.taxes.income_tax_rate.parent is copied.taxes + + +def test_a_bracket_path_copies_the_scale_and_the_bracket(): + original = SYSTEM.parameters.clone() + path = "taxes.social_security_contribution[1].rate" + copied = copy_paths(original, [path]) + + scale = copied.taxes.social_security_contribution + original_scale = original.taxes.social_security_contribution + assert scale is not original_scale + assert scale.brackets[1] is not original_scale.brackets[1] + assert scale.brackets[1].parent is scale + assert scale.brackets[0] is original_scale.brackets[0] + assert scale.brackets[1].threshold is original_scale.brackets[1].threshold + + _set(copied, path, 0.5) + assert copied.get_child(path)("2020-01-01") == 0.5 + assert original.get_child(path)("2020-01-01") == 0.06 + + +def test_updating_the_copy_leaves_the_original_and_its_caches_alone(): + original = SYSTEM.parameters.clone() + warm_node = original.get_at_instant("2017-01-01") + warm_benefits = original.benefits.get_at_instant("2017-01-01") + + copied = copy_paths(original, ["taxes.income_tax_rate"]) + _set(copied, "taxes.income_tax_rate", 0.3) + + assert copied.taxes.income_tax_rate("2017-01-01") == 0.3 + assert copied.get_at_instant("2017-01-01").taxes.income_tax_rate == 0.3 + assert original.taxes.income_tax_rate("2017-01-01") == 0.15 + assert original.get_at_instant("2017-01-01") is warm_node + assert warm_node.taxes.income_tax_rate == 0.15 + # A branch the change did not touch keeps the values it cached. + assert copied.benefits.get_at_instant("2017-01-01") is warm_benefits + + +def test_paths_with_a_common_prefix_share_their_copies(): + original = SYSTEM.parameters.clone() + copied = copy_paths( + original, ["taxes.housing_tax.rate", "taxes.housing_tax.minimal_amount"] + ) + housing_tax = copied.taxes.housing_tax + assert housing_tax is not original.taxes.housing_tax + assert housing_tax.rate.parent is housing_tax + assert housing_tax.minimal_amount.parent is housing_tax + + +def test_an_unknown_path_raises_before_anything_is_copied(): + original = SYSTEM.parameters.clone() + with pytest.raises(ValueError, match="Could not find the parameter"): + copy_paths(original, ["taxes.income_tax_rate", "taxes.no_such_tax"]) + + +@settings(max_examples=60, deadline=None) +@given( + changes=st.lists( + st.tuples( + st.sampled_from(LEAF_PATHS), + st.floats(min_value=0, max_value=1e6, allow_nan=False), + ), + min_size=1, + max_size=4, + unique_by=lambda change: change[0], + ), + warm=st.booleans(), +) +def test_copying_paths_matches_cloning_the_whole_tree(changes, warm): + """For any set of parameter changes, a tree that copies only their paths + reads exactly what a full clone with the same changes reads, and the + original reads what it did before, whether or not its at-instant caches + were warm when it was copied.""" + original = SYSTEM.parameters.clone() + before = {instant: _leaf_values(original, instant) for instant in INSTANTS} + if warm: + for instant in INSTANTS: + original.get_at_instant(instant) + + paths = [path for path, _ in changes] + copied = copy_paths(original, paths) + reference = original.clone() + for path, value in changes: + _set(copied, path, value) + _set(reference, path, value) + + for instant in INSTANTS: + assert _leaf_values(copied, instant) == _leaf_values(reference, instant) + assert _leaf_values(original, instant) == before[instant] + # Read through the at-instant trees too, which is how formulas read + # parameters and where a stale shared cache would show. + assert _values_at_instant(copied, instant) == _values_at_instant( + reference, instant + ) + assert _values_at_instant(original, instant) == { + path: before[instant][original.get_child(path).name] + for path in PLAIN_LEAF_PATHS + } + + +# Paths outside scales, which an at-instant tree reaches by attribute. +PLAIN_LEAF_PATHS = [path for path in LEAF_PATHS if "[" not in path] + + +def _values_at_instant(tree, instant): + """Each plain parameter's value read through ``tree``'s at-instant tree, + or ``None`` where that tree leaves it out (it has no value then).""" + node_at_instant = tree.get_at_instant(instant) + values = {} + for path in PLAIN_LEAF_PATHS: + value = node_at_instant + try: + for name in path.split("."): + value = getattr(value, name) + except ParameterNotFoundError: + value = None + values[path] = value + return values diff --git a/tests/core/tools/test_runner/test_forked_test_runner.py b/tests/core/tools/test_runner/test_forked_test_runner.py new file mode 100644 index 000000000..b9695d38f --- /dev/null +++ b/tests/core/tools/test_runner/test_forked_test_runner.py @@ -0,0 +1,147 @@ +"""Batches run in children forked from the process that built the system.""" + +import os +import time + +import pytest +import yaml + +from policyengine_core.country_template import CountryTaxBenefitSystem +from policyengine_core.tools import forked_test_runner +from policyengine_core.tools.forked_test_runner import ( + PYTEST_INTERNAL_ERROR, + run_batches_in_forks, +) +from policyengine_core.tools.test_runner import _tax_benefit_system_cache +from tests.fixtures.yaml_runner_memory import RUNNER_ARGV + +pytestmark = [ + pytest.mark.skipif(not hasattr(os, "fork"), reason="needs os.fork"), + # The test process has native threads it did not start (the BLAS + # library's pool, which Python counts), so every fork here warns. The + # children only run pytest on the country template. + pytest.mark.filterwarnings( + "ignore:This process .* is multi-threaded:DeprecationWarning" + ), +] + +OPTIONS = {} + + +@pytest.fixture(autouse=True) +def nested_pytest_arguments(monkeypatch): + """Run each child's session with the memory tests' arguments.""" + main = forked_test_runner.run_tests.__globals__["pytest"].main + + def nested_main(args, plugins): + return main( + [*args, "-q", "-p", "no:rerunfailures", "-p", "no:cacheprovider"], + plugins=plugins, + ) + + monkeypatch.setattr( + forked_test_runner.run_tests.__globals__["pytest"], "main", nested_main + ) + + +def _write_case(path, expected_income_tax): + path.write_text( + yaml.safe_dump( + [ + { + "name": path.stem, + "period": "2017-01", + "input": {"salary": 1000}, + "output": {"income_tax": expected_income_tax}, + } + ] + ) + ) + return str(path) + + +def test_each_batch_reports_its_own_exit_code_and_output(tmp_path): + passing = _write_case(tmp_path / "passing.yaml", 150) + failing = _write_case(tmp_path / "failing.yaml", 999) + finished = [] + + results = run_batches_in_forks( + CountryTaxBenefitSystem(), + [[passing], [failing], [passing, failing]], + OPTIONS, + workers=2, + on_finish=lambda index, result: finished.append(index), + ) + + assert [result.exit_code for result in results] == [0, 1, 1] + assert "1 passed" in results[0].output + assert "1 failed" in results[1].output + assert "1 failed, 1 passed" in results[2].output + assert sorted(finished) == [0, 1, 2] + assert all(result.elapsed > 0 for result in results) + assert not any(result.timed_out for result in results) + + +def test_the_reform_free_system_is_built_once_before_forking(tmp_path): + system = CountryTaxBenefitSystem() + run_batches_in_forks(system, [[_write_case(tmp_path / "case.yaml", 150)]]) + # Built in the parent, where every child shares it. + assert _tax_benefit_system_cache[system].reform_free is not None + + +def test_no_more_than_workers_batches_run_at_once(tmp_path, monkeypatch): + def record_while_running(tax_benefit_system, paths, options): + (marker,) = paths + with open(marker, "a") as log: + log.write(f"start {time.monotonic()}\n") + time.sleep(0.5) + with open(marker, "a") as log: + log.write(f"end {time.monotonic()}\n") + return 0 + + monkeypatch.setattr(forked_test_runner, "run_tests", record_while_running) + markers = [str(tmp_path / f"batch-{index}.log") for index in range(5)] + results = run_batches_in_forks( + CountryTaxBenefitSystem(), [[marker] for marker in markers], workers=2 + ) + + assert [result.exit_code for result in results] == [0] * 5 + events = [] + for marker in markers: + with open(marker) as log: + for line in log: + kind, moment = line.split() + events.append((float(moment), 1 if kind == "start" else -1)) + running = most = 0 + for _, change in sorted(events): + running += change + most = max(most, running) + assert most == 2 + + +def test_a_batch_over_its_timeout_is_stopped(monkeypatch): + monkeypatch.setattr( + forked_test_runner, "run_tests", lambda *args: time.sleep(60) or 0 + ) + started = time.monotonic() + (result,) = run_batches_in_forks( + CountryTaxBenefitSystem(), [["anything"]], timeout=0.5 + ) + assert result.timed_out + assert result.exit_code < 0 + assert time.monotonic() - started < 30 + + +def test_an_error_outside_pytest_is_reported_as_an_internal_error(monkeypatch): + def fail(*args): + raise RuntimeError("the batch could not start") + + monkeypatch.setattr(forked_test_runner, "run_tests", fail) + (result,) = run_batches_in_forks(CountryTaxBenefitSystem(), [["anything"]]) + assert result.exit_code == PYTEST_INTERNAL_ERROR + assert "the batch could not start" in result.output + + +def test_workers_must_be_positive(): + with pytest.raises(ValueError, match="workers must be 1 or more"): + run_batches_in_forks(CountryTaxBenefitSystem(), [], workers=0) diff --git a/tests/core/tools/test_runner/test_parameter_inputs.py b/tests/core/tools/test_runner/test_parameter_inputs.py new file mode 100644 index 000000000..833a034e4 --- /dev/null +++ b/tests/core/tools/test_runner/test_parameter_inputs.py @@ -0,0 +1,149 @@ +"""A case's dotted parameter inputs run on a system that copies only the +parameters they set, on top of the system the case's reforms select.""" + +import pytest +import yaml + +from policyengine_core.country_template import CountryTaxBenefitSystem +from policyengine_core.reforms import Reform, set_parameter +from policyengine_core.simulations import SimulationBuilder +from policyengine_core.tools.test_runner import ( + OpenFiscaPlugin, + _ParameterInputReform, + _get_tax_benefit_system, + _tax_benefit_system_cache, +) +from tests.fixtures.yaml_runner_memory import RUNNER_ARGV, write_cases + +SCALE_REFORM = ( + "policyengine_core.country_template.reforms." + "modify_social_security_taxation.modify_social_security_taxation" +) +PARAMETER_INPUTS = [ + ("taxes.income_tax_rate", 0.3), + ("taxes.housing_tax.minimal_amount", 350), + ("taxes.social_security_contribution[1].rate", 0.5), +] + + +def _full_clone_system(baseline, reforms, parameter_inputs): + """The system the runner built before: a full clone, the named reforms, + then a reform that sets the inputs on a full copy of the tree.""" + system = baseline.clone() + for reform in reforms: + system = system.apply_reform(reform) + + class inline_reform_class(Reform): + def apply(self): + for key, value in parameter_inputs: + self.parameters = set_parameter( + key, value, return_modifier=True, period="year:2000:40" + )(self.parameters) + + return inline_reform_class(system) + + +def _run(path, baseline): + return pytest.main( + [*RUNNER_ARGV, "-q", "-p", "no:cacheprovider", str(path)], + plugins=[OpenFiscaPlugin(baseline, {})], + ) + + +def test_the_system_shares_every_parameter_it_does_not_set(): + baseline = CountryTaxBenefitSystem() + reform_free = _get_tax_benefit_system(baseline, [], []) + system = _ParameterInputReform(reform_free, [("taxes.income_tax_rate", 0.3)]) + + assert system.baseline is reform_free + assert system.parameters.taxes.income_tax_rate("2017-01-01") == 0.3 + assert reform_free.parameters.taxes.income_tax_rate("2017-01-01") == 0.15 + assert system.parameters.benefits is reform_free.parameters.benefits + assert system.parameters.taxes.housing_tax is ( + reform_free.parameters.taxes.housing_tax + ) + + +@pytest.mark.parametrize("reforms", [[], [SCALE_REFORM]]) +@pytest.mark.parametrize("instant", ["2015-06-01", "2017-01-01", "2030-01-01"]) +def test_it_reads_what_a_full_clone_with_the_same_inputs_reads(reforms, instant): + """Including a bracket index, which counts the brackets the named reform + left: the inputs apply after the reforms, as they did on a full clone.""" + baseline = CountryTaxBenefitSystem() + structural = _get_tax_benefit_system(baseline, reforms, []) + system = _ParameterInputReform(structural, PARAMETER_INPUTS) + reference = _full_clone_system(baseline, reforms, PARAMETER_INPUTS) + + values = system.get_parameters_at_instant(instant) + expected = reference.get_parameters_at_instant(instant) + assert values.taxes.income_tax_rate == expected.taxes.income_tax_rate + assert ( + values.taxes.housing_tax.minimal_amount + == expected.taxes.housing_tax.minimal_amount + ) + scale = values.taxes.social_security_contribution + expected_scale = expected.taxes.social_security_contribution + assert list(scale.rates) == list(expected_scale.rates) + assert list(scale.thresholds) == list(expected_scale.thresholds) + + +def test_cases_give_the_same_results_as_on_a_full_clone(tmp_path): + baseline = CountryTaxBenefitSystem() + situation = {"salary": 50_000} + inputs = {**situation, **{key: value for key, value in PARAMETER_INPUTS}} + variables = ["income_tax", "social_security_contribution"] + cases = [] + for reforms in [[], [SCALE_REFORM]]: + reference = _full_clone_system(baseline, reforms, PARAMETER_INPUTS) + builder = SimulationBuilder() + builder.set_default_period("2017-01") + simulation = builder.build_from_dict(reference, dict(situation)) + expected = { + variable: float(simulation.calculate(variable, "2017-01")[0]) + for variable in variables + } + assert expected["income_tax"] == pytest.approx(15_000) + for variable in variables: + cases.append( + { + "name": f"{variable} with reforms {reforms}", + "period": "2017-01", + **({"reforms": reforms} if reforms else {}), + "input": inputs, + "output": {variable: expected[variable]}, + "absolute_error_margin": 0.01, + } + ) + # The reform changes the scale, so the two contributions must differ for + # this to show the inputs were applied on top of it. + assert cases[1]["output"] != cases[3]["output"] + path = tmp_path / "cases.yaml" + path.write_text(yaml.safe_dump(cases, sort_keys=False)) + assert _run(path, baseline) == 0 + + +def test_cases_with_parameter_inputs_add_nothing_to_the_reform_cache(tmp_path): + baseline = CountryTaxBenefitSystem() + path = write_cases(tmp_path / "cases.yaml", 30, reform_every=1) + assert _run(path, baseline) == 0 + cached = _tax_benefit_system_cache[baseline] + assert cached.reform_free is not None + assert len(cached.reforms) == 0 + + +def test_an_unknown_parameter_input_fails_its_case(tmp_path, capsys): + path = tmp_path / "unknown.yaml" + path.write_text( + yaml.safe_dump( + [ + { + "name": "unknown", + "period": "2017-01", + "input": {"salary": 1000, "taxes.no_such_tax": 1}, + "output": {"income_tax": 150}, + } + ] + ) + ) + assert _run(path, CountryTaxBenefitSystem()) == 1 + assert "Could not find the parameter taxes.no_such_tax" in capsys.readouterr().out diff --git a/tests/core/tools/test_runner/test_read_only_arrays.py b/tests/core/tools/test_runner/test_read_only_arrays.py new file mode 100644 index 000000000..42eda2c3a --- /dev/null +++ b/tests/core/tools/test_runner/test_read_only_arrays.py @@ -0,0 +1,148 @@ +"""With ``read_only_arrays``, a formula that writes into an array it read from +another variable fails at that line instead of changing that variable.""" + +import numpy as np +import pytest +import yaml + +from policyengine_core.country_template import CountryTaxBenefitSystem +from policyengine_core.country_template.entities import Person +from policyengine_core.data_storage import InMemoryStorage +from policyengine_core.periods import MONTH, period +from policyengine_core.tools.test_runner import OpenFiscaPlugin, run_tests +from policyengine_core.variables import Variable +from tests.fixtures.yaml_runner_memory import RUNNER_ARGV + +JANUARY = period("2017-01") + + +class salary_plus_one(Variable): + """Writes into the salary array it read: changes ``salary`` itself.""" + + value_type = float + entity = Person + label = "Salary plus one, written into the salary array" + definition_period = MONTH + + def formula(person, period, parameters): + salary = person("salary", period) + salary += 1 + return salary + + +@pytest.fixture +def read_only(monkeypatch): + monkeypatch.setattr(InMemoryStorage, "read_only_arrays", True) + + +def test_stored_arrays_stay_writeable_by_default(): + storage = InMemoryStorage(is_eternal=False) + storage.put(np.array([1.0, 2.0]), JANUARY) + stored = storage.get(JANUARY) + stored += 1 + assert list(storage.get(JANUARY)) == [2.0, 3.0] + + +def test_stored_arrays_are_read_only_and_the_callers_array_is_not(read_only): + storage = InMemoryStorage(is_eternal=False) + given = np.array([1.0, 2.0]) + storage.put(given, JANUARY) + stored = storage.get(JANUARY) + + with pytest.raises(ValueError, match="read-only"): + stored += 1 + assert given.flags.writeable + assert list(storage.get(JANUARY)) == [1.0, 2.0] + + +def test_an_array_a_clone_copies_on_first_read_is_read_only(read_only): + storage = InMemoryStorage(is_eternal=False) + storage.put(np.array([1.0, 2.0]), JANUARY) + clone = storage.clone(share_arrays=True) + + copied = clone.get(JANUARY) + with pytest.raises(ValueError, match="read-only"): + copied[0] = 5 + assert list(storage.get(JANUARY)) == [1.0, 2.0] + + +def _system_with_a_writing_formula(): + system = CountryTaxBenefitSystem() + system.add_variable(salary_plus_one) + return system + + +def _cases(tmp_path): + path = tmp_path / "writes.yaml" + path.write_text( + yaml.safe_dump( + [ + { + "name": "writes into salary", + "period": "2017-01", + "input": {"salary": 1000}, + "output": {"salary_plus_one": 1001}, + } + ] + ) + ) + return path + + +def test_a_formula_writing_into_a_cached_array_passes_without_the_option(tmp_path): + exit_code = pytest.main( + [*RUNNER_ARGV, "-q", "-p", "no:cacheprovider", str(_cases(tmp_path))], + plugins=[OpenFiscaPlugin(_system_with_a_writing_formula(), {})], + ) + assert exit_code == 0 + + +def test_the_option_fails_a_formula_writing_into_a_cached_array( + tmp_path, capsys, monkeypatch +): + monkeypatch.setattr( + "policyengine_core.tools.test_runner.pytest.main", + _main_without_plugins_from_outside(), + ) + exit_code = run_tests( + _system_with_a_writing_formula(), + str(_cases(tmp_path)), + {"read_only_arrays": True}, + ) + assert exit_code == 1 + assert "read-only" in capsys.readouterr().out + # The run restores the setting it found. + assert InMemoryStorage.read_only_arrays is False + + +def _main_without_plugins_from_outside(): + """``pytest.main`` with the arguments the memory tests use, so the nested + session does not load this session's rerun plugin.""" + main = pytest.main + + def nested_main(args, plugins): + return main( + [*args, "-p", "no:rerunfailures", "-p", "no:cacheprovider"], plugins=plugins + ) + + return nested_main + + +def test_command_line_passes_read_only_arrays(monkeypatch): + 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()) + for arguments, expected in [([], False), (["--read-only-arrays"], True)]: + monkeypatch.setattr( + "sys.argv", ["policyengine-core", "test", "cases.yaml", *arguments] + ) + with pytest.raises(SystemExit): + run_test.main(policyengine_command.get_parser()) + assert seen["read_only_arrays"] is expected From f51bd673ca4a00022a81a9f33e6d51e3ac68c28f Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 00:54:35 +0400 Subject: [PATCH 16/23] Build reform systems on the reform-free system, and evict before building A YAML case's reform system was built on a fresh full clone of the baseline, which the reform kept as its `baseline`, so every cached reform system held two copies of the parameter tree. A reform named by its path, or a Reform class, copies the tree, entities and variable registry before `apply` changes them and leaves the system it is applied to as it was, so these are now built on the shared reform-free system. A reform function is still applied to a clone. The cache also drops its least recently used systems before building a new one, not after, so the cache and the system being built never hold more than `reform_cache_size` reform systems. Co-Authored-By: Claude Opus 5.5 --- policyengine_core/tools/test_runner.py | 34 +++++++++-- .../test_runner/test_reform_cache_property.py | 12 ++-- .../test_runner/test_reform_system_size.py | 60 +++++++++++++++++++ 3 files changed, 96 insertions(+), 10 deletions(-) create mode 100644 tests/core/tools/test_runner/test_reform_system_size.py diff --git a/policyengine_core/tools/test_runner.py b/policyengine_core/tools/test_runner.py index 12b2b36ed..9470d6078 100644 --- a/policyengine_core/tools/test_runner.py +++ b/policyengine_core/tools/test_runner.py @@ -539,9 +539,9 @@ def _get_tax_benefit_system( while len(reform_systems) > cache_size: reform_systems.popitem(last=False) + if cached_systems.reform_free is None: + cached_systems.reform_free = baseline.clone() 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). The @@ -561,7 +561,23 @@ def _get_tax_benefit_system( reform_systems.move_to_end(inner_key) return cached - current_tax_benefit_system = baseline.clone() + # Make room before building, not after, so the cache and the system being + # built never add up to more than ``cache_size`` reform systems. + while reform_systems and len(reform_systems) >= cache_size: + reform_systems.popitem(last=False) + + if (reforms or extensions) and all( + _copies_its_baseline(reform) for reform in reforms + ): + # Every reform copies the parameters it changes and leaves the system + # it is applied to as it was, and each extension below is loaded into + # a clone, so build on the shared reform-free system: what comes back + # is a new system either way. A clone of the baseline here doubled + # every reform system's size, the clone kept alive as the reform's + # ``baseline``. + current_tax_benefit_system = cached_systems.reform_free + else: + current_tax_benefit_system = baseline.clone() for reform_path in reforms: if isinstance(reform_path, str): @@ -578,12 +594,20 @@ def _get_tax_benefit_system( if cache_size > 0: reform_systems[inner_key] = current_tax_benefit_system - while len(reform_systems) > cache_size: - reform_systems.popitem(last=False) return current_tax_benefit_system +def _copies_its_baseline(reform) -> bool: + """Whether applying ``reform`` leaves the system it is applied to as it + was: a reform named by its path (``apply_reform`` requires a ``Reform``) + or a ``Reform`` class, whose constructor copies the baseline's parameter + tree, entities and variable registry before ``apply`` changes them.""" + return isinstance(reform, str) or ( + isinstance(reform, type) and issubclass(reform, Reform) + ) + + def assert_near( value, target_value, 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 7abde339b..3ce2cdccd 100644 --- a/tests/core/tools/test_runner/test_reform_cache_property.py +++ b/tests/core/tools/test_runner/test_reform_cache_property.py @@ -3,7 +3,8 @@ 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. +extensions requested, and keeps one reform-free system throughout, which no +reform or extension changes although reforms are built on it. ``test_runner_memory.py`` pins the same behaviour with examples. """ @@ -83,9 +84,12 @@ def test_cache_matches_reference_lru(sequence): assert applied_reforms == list(reforms) assert applied_extensions == set(extensions) + cached = _tax_benefit_system_cache[baseline] + if reform_free is None: + reform_free = cached.reform_free + assert cached.reform_free is reform_free + assert reform_free.applied == () 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 = (tuple(reforms), reform_key, frozenset(extensions)) @@ -100,7 +104,5 @@ def test_cache_matches_reference_lru(sequence): 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_reform_system_size.py b/tests/core/tools/test_runner/test_reform_system_size.py new file mode 100644 index 000000000..f3843e9ba --- /dev/null +++ b/tests/core/tools/test_runner/test_reform_system_size.py @@ -0,0 +1,60 @@ +"""A reform system holds one copy of the parameter tree, and the cache never +holds more reform systems than its size, even while it builds one.""" + +import gc +import weakref + +from policyengine_core.country_template import CountryTaxBenefitSystem +from policyengine_core.reforms import Reform +from policyengine_core.tools.test_runner import _get_tax_benefit_system + +NEW_TAX = "policyengine_core.country_template.reforms.add_new_tax.add_new_tax" + + +def test_a_named_reform_is_built_on_the_reform_free_system_and_leaves_it_alone(): + baseline = CountryTaxBenefitSystem() + reform_free = _get_tax_benefit_system(baseline, [], []) + variables = dict(reform_free.variables) + parameters = reform_free.parameters + entities = list(reform_free.entities) + + system = _get_tax_benefit_system(baseline, [NEW_TAX], []) + + assert system.baseline is reform_free + assert system.parameters is not parameters + assert "new_tax" in system.variables + assert "new_tax" not in reform_free.variables + assert reform_free.variables == variables + assert reform_free.parameters is parameters + assert reform_free.entities == entities + assert all( + entity._tax_benefit_system is reform_free for entity in reform_free.entities + ) + + +def test_a_reform_function_is_applied_to_a_clone(): + """Only reforms that copy what they change are built on the shared system.""" + baseline = CountryTaxBenefitSystem() + reform_free = _get_tax_benefit_system(baseline, [], []) + received = [] + + def reform_function(system): + received.append(system) + return system + + _get_tax_benefit_system(baseline, [reform_function], [], reform_key="function") + assert received and received[0] is not reform_free + + +def test_the_cache_releases_a_system_before_building_the_one_replacing_it(): + baseline = CountryTaxBenefitSystem() + first = weakref.ref(_get_tax_benefit_system(baseline, [NEW_TAX], [], cache_size=1)) + alive_while_building = [] + + class probe(Reform): + def apply(self): + gc.collect() + alive_while_building.append(first() is not None) + + _get_tax_benefit_system(baseline, [probe], [], cache_size=1) + assert alive_while_building == [False] From 1091db04dcc45c1d8d75c180cec69540329ee8bf Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 00:54:45 +0400 Subject: [PATCH 17/23] Describe reform systems built on the reform-free system Co-Authored-By: Claude Opus 5.5 --- changelog.d/test-runner-perf.added.md | 2 +- docs/engineering/skills/testing.md | 10 ++++++---- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/changelog.d/test-runner-perf.added.md b/changelog.d/test-runner-perf.added.md index 73c92202b..900355107 100644 --- a/changelog.d/test-runner-perf.added.md +++ b/changelog.d/test-runner-perf.added.md @@ -1 +1 @@ -The YAML test runner sets a case's dotted parameter inputs on a reform that copies only those parameters' paths instead of the whole tree, adds a `--read-only-arrays` option that fails a formula writing into an array it read, and `policyengine_core.tools.forked_test_runner.run_batches_in_forks` runs test batches in children forked from one process that built the tax-benefit system. +The YAML test runner sets a case's dotted parameter inputs on a reform that copies only those parameters' paths instead of the whole tree, builds reform systems on the shared reform-free system so each holds one copy of the parameter tree, adds a `--read-only-arrays` option that fails a formula writing into an array it read, and `policyengine_core.tools.forked_test_runner.run_batches_in_forks` runs test batches in children forked from one process that built the tax-benefit system. diff --git a/docs/engineering/skills/testing.md b/docs/engineering/skills/testing.md index 915c70abe..b20f02e35 100644 --- a/docs/engineering/skills/testing.md +++ b/docs/engineering/skills/testing.md @@ -46,10 +46,12 @@ process. Its memory is bounded by design, and tests must keep it that way: 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` 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. + that combination on the reform-free system: a reform copies the parameter + tree before it changes it, so each holds one copy of the tree. The runner + caches the `--reform-cache-size` most recently used of them (default 2) plus + the reform-free system, and drops one before building another when full. + Do not add a cache to the runner that grows with the number of cases or + combinations. - A case that sets a parameter with a dotted input key runs on a system of its own (`_ParameterInputReform`) that copies only the nodes on the paths to the parameters it sets (`policyengine_core.parameters.path_copy`) and shares the From 78ca7d8154c8164f8f58f456c02829bc1bb9188c Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 18:11:46 +0400 Subject: [PATCH 18/23] Clear cached instants from the tree modify_parameters copies modify_parameters gives the modifier a deep copy of the baseline's parameter tree, and the copy kept every value the baseline had cached at an instant. A modifier that changes values or brackets in place (the country template's modify_social_security_taxation) then read the baseline's values at any instant the baseline had already read. The YAML runner now builds reforms on the reform-free system other cases have read, so the country template's reform cases failed after any earlier case read 2017. The same happens on master for a reform applied to a system that has already computed. #584's reform test checked that the copy kept the baseline's cached vectorial node. The copy now keeps no cache, so the test reads the values first and then checks the node the read built. It passes with and without this change. Co-Authored-By: Claude Opus 5.5 --- ...modify-parameters-cached-instants.fixed.md | 1 + policyengine_core/reforms/reform.py | 11 ++++++++ .../test_vectorial_parameter_node_copy.py | 6 ++--- tests/core/test_reforms.py | 27 +++++++++++++++++++ 4 files changed, 42 insertions(+), 3 deletions(-) create mode 100644 changelog.d/modify-parameters-cached-instants.fixed.md diff --git a/changelog.d/modify-parameters-cached-instants.fixed.md b/changelog.d/modify-parameters-cached-instants.fixed.md new file mode 100644 index 000000000..095dcf566 --- /dev/null +++ b/changelog.d/modify-parameters-cached-instants.fixed.md @@ -0,0 +1 @@ +`Reform.modify_parameters` reads a modifier's in-place changes to values or brackets even when the baseline had already cached its parameters at that instant. diff --git a/policyengine_core/reforms/reform.py b/policyengine_core/reforms/reform.py index 8c5200ac5..b441f5ab1 100644 --- a/policyengine_core/reforms/reform.py +++ b/policyengine_core/reforms/reform.py @@ -17,6 +17,13 @@ import requests +def _clear_instant_caches(parameters: ParameterNode) -> None: + """Drop every node's cached values at each instant.""" + for node in (parameters, *parameters.get_descendants()): + if isinstance(node, ParameterNode): + node._at_instant_cache = {} + + class classproperty(object): def __init__(self, f): self.f = f @@ -207,6 +214,10 @@ def modify_parameters( """ baseline_parameters = self.baseline.parameters baseline_parameters_copy = copy.deepcopy(baseline_parameters) + # The copy carries the values the baseline has cached at each + # instant, which a modifier that changes values or brackets in place + # does not clear. + _clear_instant_caches(baseline_parameters_copy) reform_parameters = modifier_function(baseline_parameters_copy) if not isinstance(reform_parameters, ParameterNode): return ValueError( diff --git a/tests/core/parameters/test_vectorial_parameter_node_copy.py b/tests/core/parameters/test_vectorial_parameter_node_copy.py index 98af4f9b2..c92da4e30 100644 --- a/tests/core/parameters/test_vectorial_parameter_node_copy.py +++ b/tests/core/parameters/test_vectorial_parameter_node_copy.py @@ -200,7 +200,7 @@ def test_copied_parameter_tree_keeps_its_cached_vectorial_node(copier): assert_near(rate[families][housings].z2, expected) -def test_reform_that_modifies_parameters_keeps_the_cached_vectorial_node( +def test_reform_that_modifies_parameters_reads_the_vectorial_node( isolated_tax_benefit_system, ): system = isolated_tax_benefit_system @@ -210,10 +210,10 @@ def test_reform_that_modifies_parameters_keeps_the_cached_vectorial_node( class KeepParameters(Reform): def apply(self): - # Deep-copies the baseline's tree, cache included. + # Deep-copies the baseline's tree, which has cached that node. self.modify_parameters(lambda parameters: parameters) general = KeepParameters(system).get_parameters_at_instant("2017-01-01").general - assert type(general._vectorial_node) is VectorialParameterNodeAtInstant assert_near(general[names], expected) + assert type(general._vectorial_node) is VectorialParameterNodeAtInstant diff --git a/tests/core/test_reforms.py b/tests/core/test_reforms.py index 711888b50..7c7eeb01b 100644 --- a/tests/core/test_reforms.py +++ b/tests/core/test_reforms.py @@ -3,7 +3,11 @@ import pytest from policyengine_core import periods +from policyengine_core.country_template import CountryTaxBenefitSystem from policyengine_core.country_template.entities import Household, Person +from policyengine_core.country_template.reforms.modify_social_security_taxation import ( + modify_social_security_taxation, +) from policyengine_core.model_api import * from policyengine_core.parameters import ParameterNode from policyengine_core.periods import Instant @@ -299,6 +303,29 @@ def apply(self): assert parameters_at_instant.new_node.new_param is True +def test_modify_parameters_reads_its_changes_after_the_baseline_cached_them(): + """``modify_parameters`` gives the modifier a deep copy of the baseline's + tree. A copy that kept the values the baseline had cached at each instant + read those, not a change the modifier made to values or brackets in + place.""" + + def contribution_scale(system): + return system.get_parameters_at_instant( + "2017-01-01" + ).taxes.social_security_contribution + + expected = contribution_scale( + modify_social_security_taxation(CountryTaxBenefitSystem()) + ) + read_baseline = CountryTaxBenefitSystem() + baseline_scale = contribution_scale(read_baseline) + assert baseline_scale.rates != expected.rates + + reformed = contribution_scale(modify_social_security_taxation(read_baseline)) + assert reformed.thresholds == expected.thresholds + assert reformed.rates == expected.rates + + def _income_tax_rate(reform, year): return reform.get_parameters_at_instant(Instant((year, 1, 1))).taxes.income_tax_rate From a7b9973dc3aba41b7d808bb82c2f679f7c7f2387 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 20:04:15 +0400 Subject: [PATCH 19/23] Give modify_parameters the reform's own tree, and clear cached instants after it Reform.modify_parameters deep-copied the baseline's tree. That copy kept the baseline out of a modifier's edits until #459 made Reform.__init__ clone the tree, so it had become a second full copy (5.7 s for policyengine-us's tree, against 3.9 s for the clone). Because each call started again from the baseline, it also dropped any change apply() had made before it, which apply_reform_set keeps. Reform now inherits TaxBenefitSystem.modify_parameters, which hands the modifier the system's own tree, so a reform instance and an in-place reform behave the same. That method now clears every node's cached instants, and the system's, after the modifier runs: a modifier that edits values, brackets or children in place, rather than through Parameter.update, read stale values on the in-place path too. Fixes #618. Co-Authored-By: Claude Opus 5.5 --- ...modify-parameters-cached-instants.fixed.md | 2 +- ...form-modify-parameters-own-tree.changed.md | 1 + policyengine_core/reforms/reform.py | 47 ++--------- .../taxbenefitsystems/tax_benefit_system.py | 22 ++++++ .../test_vectorial_parameter_node_copy.py | 5 +- tests/core/test_reforms.py | 77 ++++++++++++++++++- .../core/test_tracing_parameter_isolation.py | 5 +- 7 files changed, 109 insertions(+), 50 deletions(-) create mode 100644 changelog.d/reform-modify-parameters-own-tree.changed.md diff --git a/changelog.d/modify-parameters-cached-instants.fixed.md b/changelog.d/modify-parameters-cached-instants.fixed.md index 095dcf566..5be9d168e 100644 --- a/changelog.d/modify-parameters-cached-instants.fixed.md +++ b/changelog.d/modify-parameters-cached-instants.fixed.md @@ -1 +1 @@ -`Reform.modify_parameters` reads a modifier's in-place changes to values or brackets even when the baseline had already cached its parameters at that instant. +`modify_parameters` drops the values a parameter tree has cached at each instant once its modifier has run, so a modifier's in-place changes to values, brackets or children are read even where the system had already read those parameters. diff --git a/changelog.d/reform-modify-parameters-own-tree.changed.md b/changelog.d/reform-modify-parameters-own-tree.changed.md new file mode 100644 index 000000000..332266b20 --- /dev/null +++ b/changelog.d/reform-modify-parameters-own-tree.changed.md @@ -0,0 +1 @@ +`Reform.modify_parameters` no longer deep-copies the baseline's parameter tree. The modifier changes the reform's own copy, as it already did for a reform applied with `apply_reform_set`, so changes made earlier in `apply` are kept and building such a reform skips a full copy of the tree. diff --git a/policyengine_core/reforms/reform.py b/policyengine_core/reforms/reform.py index b441f5ab1..2465013cc 100644 --- a/policyengine_core/reforms/reform.py +++ b/policyengine_core/reforms/reform.py @@ -1,7 +1,6 @@ from __future__ import annotations -import copy -from typing import Callable, Union, TYPE_CHECKING +from typing import Union, TYPE_CHECKING from policyengine_core.parameters import ParameterNode, Parameter from policyengine_core.taxbenefitsystems import TaxBenefitSystem @@ -17,13 +16,6 @@ import requests -def _clear_instant_caches(parameters: ParameterNode) -> None: - """Drop every node's cached values at each instant.""" - for node in (parameters, *parameters.get_descendants()): - if isinstance(node, ParameterNode): - node._at_instant_cache = {} - - class classproperty(object): def __init__(self, f): self.f = f @@ -159,10 +151,9 @@ def __init__(self, baseline: TaxBenefitSystem): super().__init__(baseline.entities) self.baseline = baseline # Clone the parameter tree so in-place mutations inside ``apply()`` - # (``set_parameter(...)`` or ``self.parameters.path.update(...)``) - # don't leak into the baseline. ``modify_parameters()`` was already - # deep-copying, but plenty of reforms mutate ``self.parameters`` - # directly (bug C4). + # (``set_parameter(...)``, ``self.parameters.path.update(...)`` or a + # ``modify_parameters()`` modifier) don't leak into the baseline + # (bug C4). if baseline.parameters is not None: self.parameters = self._copy_parameters(baseline.parameters) else: @@ -180,7 +171,8 @@ def __init__(self, baseline: TaxBenefitSystem): self.apply() def _copy_parameters(self, baseline_parameters: ParameterNode) -> ParameterNode: - """The tree ``apply`` starts from and may change in place. + """The tree ``apply`` starts from and may change in place, directly or + through ``modify_parameters``. A full clone of the baseline's tree, since a reform can change any node. A reform that changes only parameters it can name in advance can @@ -202,33 +194,6 @@ def full_key(self) -> str: key = ".".join([baseline_full_key, key]) return key - def modify_parameters( - self, modifier_function: Callable[[ParameterNode], ParameterNode] - ) -> None: - """Make modifications on the parameters of the legislation. - - Call this function in `apply()` if the reform asks for legislation parameter modifications. - - Args: - modifier_function: A function that takes a :obj:`.ParameterNode` and should return an object of the same type. - """ - baseline_parameters = self.baseline.parameters - baseline_parameters_copy = copy.deepcopy(baseline_parameters) - # The copy carries the values the baseline has cached at each - # instant, which a modifier that changes values or brackets in place - # does not clear. - _clear_instant_caches(baseline_parameters_copy) - reform_parameters = modifier_function(baseline_parameters_copy) - if not isinstance(reform_parameters, ParameterNode): - return ValueError( - "modifier_function {} in module {} must return a ParameterNode".format( - modifier_function.__name__, - modifier_function.__module__, - ) - ) - self.parameters = reform_parameters - self._parameters_at_instant_cache = {} - @staticmethod def from_dict( parameter_values: dict, diff --git a/policyengine_core/taxbenefitsystems/tax_benefit_system.py b/policyengine_core/taxbenefitsystems/tax_benefit_system.py index 74d01c23c..198c67579 100644 --- a/policyengine_core/taxbenefitsystems/tax_benefit_system.py +++ b/policyengine_core/taxbenefitsystems/tax_benefit_system.py @@ -57,6 +57,16 @@ log = logging.getLogger(__name__) +def _clear_instant_caches(parameters) -> None: + """Drop the values every node of a parameter tree has cached at each + instant.""" + if not isinstance(parameters, ParameterNode): + return + for node in (parameters, *parameters.get_descendants()): + if isinstance(node, ParameterNode): + node._at_instant_cache = {} + + class TaxBenefitSystem: """ Represents the legislation. @@ -765,6 +775,11 @@ def modify_parameters( Call this function in `apply()` if the reform asks for legislation parameter modifications. + The modifier is given this system's own parameter tree, with any + change ``apply()`` has made so far, and may change it in place. A + reform built with ``Reform(baseline)`` has its own copy of the + baseline's tree, so the baseline is left unchanged. + Args: modifier_function: A function that takes a :obj:`.ParameterNode` and should return an object of the same type. """ @@ -778,6 +793,11 @@ def modify_parameters( ) else: reform_parameters = modifier_function(self.parameters) + # A modifier may change values, brackets or children in place + # rather than through ``update``, which is what drops the values a + # node has cached at each instant. + self._parameters_at_instant_cache = {} + _clear_instant_caches(self.parameters) if not isinstance(reform_parameters, ParameterNode): return ValueError( "modifier_function {} in module {} must return a ParameterNode".format( @@ -785,6 +805,8 @@ def modify_parameters( modifier_function.__module__, ) ) + if reform_parameters is not self.parameters: + _clear_instant_caches(reform_parameters) self.parameters = reform_parameters self._parameters_at_instant_cache = {} return self diff --git a/tests/core/parameters/test_vectorial_parameter_node_copy.py b/tests/core/parameters/test_vectorial_parameter_node_copy.py index c92da4e30..5372ae7d2 100644 --- a/tests/core/parameters/test_vectorial_parameter_node_copy.py +++ b/tests/core/parameters/test_vectorial_parameter_node_copy.py @@ -9,7 +9,7 @@ ``recarray``. ``ParameterNodeAtInstant`` caches its vectorial node, and -``Reform.modify_parameters`` deep-copies the baseline's parameter tree, so a +``Reform.modify_parameters`` deep-copied the baseline's parameter tree, so a reform found a ``recarray`` where its cached node should be. """ @@ -210,7 +210,8 @@ def test_reform_that_modifies_parameters_reads_the_vectorial_node( class KeepParameters(Reform): def apply(self): - # Deep-copies the baseline's tree, which has cached that node. + # Given the reform's own copy of the tree, not the baseline's, + # which has cached that node. self.modify_parameters(lambda parameters: parameters) general = KeepParameters(system).get_parameters_at_instant("2017-01-01").general diff --git a/tests/core/test_reforms.py b/tests/core/test_reforms.py index 7c7eeb01b..ecec9285c 100644 --- a/tests/core/test_reforms.py +++ b/tests/core/test_reforms.py @@ -304,10 +304,9 @@ def apply(self): def test_modify_parameters_reads_its_changes_after_the_baseline_cached_them(): - """``modify_parameters`` gives the modifier a deep copy of the baseline's - tree. A copy that kept the values the baseline had cached at each instant - read those, not a change the modifier made to values or brackets in - place.""" + """A modifier that changes values or brackets in place is read even where + the baseline had cached its parameters. The modifier used to get a deep + copy of the baseline's tree that kept those cached values.""" def contribution_scale(system): return system.get_parameters_at_instant( @@ -326,6 +325,76 @@ def contribution_scale(system): assert reformed.rates == expected.rates +def _set_income_tax_rate_in_place(parameters): + """Changes values without ``Parameter.update``, which drops cached values.""" + for value_at_instant in parameters.taxes.income_tax_rate.values_list: + value_at_instant.value = 0.5 + return parameters + + +def _double_basic_income(parameters): + basic_income = parameters.benefits.basic_income + basic_income.update(period="year:2017:1", value=2 * basic_income("2017-01-01")) + return parameters + + +class set_income_tax_rate_in_place(Reform): + def apply(self): + self.modify_parameters(_set_income_tax_rate_in_place) + + +class update_then_two_modifiers(Reform): + def apply(self): + self.parameters.taxes.income_tax_rate.update(period="year:2017:1", value=0.2) + self.modify_parameters(_double_basic_income) + self.modify_parameters(_double_basic_income) + + +def _income_tax_rate_and_basic_income(system): + parameters = system.get_parameters_at_instant("2017-01-01") + return parameters.taxes.income_tax_rate, parameters.benefits.basic_income + + +def test_an_in_place_reform_reads_its_changes_after_the_system_cached_them(): + """``apply_reform_set`` runs ``apply`` on the system itself, whose tree + may already hold its values at an instant.""" + system = CountryTaxBenefitSystem() + assert _income_tax_rate_and_basic_income(system) == (0.15, 600) + + system.apply_reform_set(set_income_tax_rate_in_place) + + assert _income_tax_rate_and_basic_income(system) == (0.5, 600) + + +def test_modify_parameters_leaves_the_baseline_unchanged(): + baseline = CountryTaxBenefitSystem() + + reform = set_income_tax_rate_in_place(baseline) + + assert _income_tax_rate_and_basic_income(reform) == (0.5, 600) + assert _income_tax_rate_and_basic_income(baseline) == (0.15, 600) + + +def test_modify_parameters_keeps_what_apply_changed_before_it(): + """The modifier used to get a fresh copy of the baseline's tree, which + dropped the update made before it and the first modifier's change.""" + reform = update_then_two_modifiers(CountryTaxBenefitSystem()) + + assert _income_tax_rate_and_basic_income(reform) == (0.2, 2400) + + +def test_a_reform_instance_and_an_in_place_reform_give_the_same_parameters(): + """``Reform(baseline)`` and ``apply_reform_set`` run the same ``apply``.""" + in_place = CountryTaxBenefitSystem() + in_place.apply_reform_set(update_then_two_modifiers) + + reform = update_then_two_modifiers(CountryTaxBenefitSystem()) + + assert _income_tax_rate_and_basic_income( + reform + ) == _income_tax_rate_and_basic_income(in_place) + + def _income_tax_rate(reform, year): return reform.get_parameters_at_instant(Instant((year, 1, 1))).taxes.income_tax_rate diff --git a/tests/core/test_tracing_parameter_isolation.py b/tests/core/test_tracing_parameter_isolation.py index 2660fe9ab..d727e0f5a 100644 --- a/tests/core/test_tracing_parameter_isolation.py +++ b/tests/core/test_tracing_parameter_isolation.py @@ -16,7 +16,7 @@ - ``ParameterNode.clone`` copies ``trace`` and ``tracer``, so every clone of the system traced as well; - copying a ``TracingParameterNodeAtInstant`` (``copy``, ``deepcopy`` or - ``pickle``, as ``Reform.modify_parameters`` does when it deep-copies the + ``pickle``, as ``Reform.modify_parameters`` did when it deep-copied the tree) recursed without end, because ``__getattr__`` read ``parameter_node_at_instant`` from an instance the copy protocol had not filled in yet. @@ -81,7 +81,8 @@ def test_traced_calculation_leaves_the_system_untraced(isolated_tax_benefit_syst def test_reform_after_a_traced_calculation_can_copy_the_parameters( isolated_tax_benefit_system, ): - """``Reform.modify_parameters`` deep-copies the baseline's parameter tree.""" + """A reform built after a traced calculation reads its own parameters + and leaves the system's unchanged.""" system = isolated_tax_benefit_system simulation = build_simulation(system, trace=True) simulation.calculate("income_tax", JANUARY) From 3089c3439bac0d1e3f5757190263fc22eeb53730 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 15 Oct 2026 01:40:25 +0400 Subject: [PATCH 20/23] Rebuild the system's parameters at an instant after the tree changes TaxBenefitSystem.get_parameters_at_instant cached each instant's node on top of the tree's own cache. Parameter.update clears the tree's caches up to the root but never the system's, so the system went on returning the node it had cached before the change: a consistent old snapshot when nodes were built up front, and with #570's lazy nodes a mix of old values for children read before the change and new ones for the rest. The system now uses a cached node only while the root still caches the same node, which costs a dictionary lookup per call. Co-Authored-By: Claude Opus 5.5 --- ...-at-instant-cache-follows-updates.fixed.md | 1 + .../taxbenefitsystems/tax_benefit_system.py | 14 ++++++ .../core/test_parameters_at_instant_cache.py | 47 +++++++++++++++++++ 3 files changed, 62 insertions(+) create mode 100644 changelog.d/parameters-at-instant-cache-follows-updates.fixed.md create mode 100644 tests/core/test_parameters_at_instant_cache.py diff --git a/changelog.d/parameters-at-instant-cache-follows-updates.fixed.md b/changelog.d/parameters-at-instant-cache-follows-updates.fixed.md new file mode 100644 index 000000000..78559e977 --- /dev/null +++ b/changelog.d/parameters-at-instant-cache-follows-updates.fixed.md @@ -0,0 +1 @@ +`TaxBenefitSystem.get_parameters_at_instant` returns the parameters as they are after a `Parameter.update`, instead of the node it had cached before the change. diff --git a/policyengine_core/taxbenefitsystems/tax_benefit_system.py b/policyengine_core/taxbenefitsystems/tax_benefit_system.py index 74d01c23c..d5395b896 100644 --- a/policyengine_core/taxbenefitsystems/tax_benefit_system.py +++ b/policyengine_core/taxbenefitsystems/tax_benefit_system.py @@ -612,6 +612,12 @@ 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 not None and not _tree_still_caches( + parameters, str(instant), parameters_at_instant + ): + # A change to the tree (``Parameter.update``, ``add_child``) + # clears the tree's own caches up to its root, not this one. + parameters_at_instant = None 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): @@ -798,3 +804,11 @@ def add_modelled_policy_metadata(self): with open(self.modelled_policies, "r") as f: modelled_policies = yaml.load(f, Loader=yaml.FullLoader) self.modelled_policies = modelled_policies + + +def _tree_still_caches(parameters, instant_str: str, parameters_at_instant) -> bool: + """Whether the root of ``parameters`` still caches this node at instant; + ``Parameter.update`` clears the caches from the parameter it changes up to + the root.""" + tree_cache = getattr(parameters, "_at_instant_cache", None) + return tree_cache is None or tree_cache.get(instant_str) is parameters_at_instant diff --git a/tests/core/test_parameters_at_instant_cache.py b/tests/core/test_parameters_at_instant_cache.py new file mode 100644 index 000000000..ecc63733c --- /dev/null +++ b/tests/core/test_parameters_at_instant_cache.py @@ -0,0 +1,47 @@ +"""``TaxBenefitSystem.get_parameters_at_instant`` follows changes to the tree. + +The system caches the parameters at each instant on top of the tree's own +cache. ``Parameter.update`` clears the tree's caches up to the root but not +the system's, which went on returning the node it cached before the change. +""" + +from policyengine_core.country_template import CountryTaxBenefitSystem + +INSTANT = "2015-01-01" + + +def test_repeated_requests_return_the_cached_node(): + system = CountryTaxBenefitSystem() + first = system.get_parameters_at_instant(INSTANT) + assert system.get_parameters_at_instant(INSTANT) is first + + +def test_a_request_after_an_update_reads_the_update(): + system = CountryTaxBenefitSystem() + before = system.get_parameters_at_instant(INSTANT) + assert before.taxes.income_tax_rate == 0.15 + + taxes = system.parameters.taxes + taxes.income_tax_rate.update(period="year:2015:1", value=0.99) + taxes.housing_tax.rate.update(period="year:2015:1", value=77) + + after = system.get_parameters_at_instant(INSTANT) + assert after is not before + # Read before the update and not: both as the tree now holds them. + assert (after.taxes.income_tax_rate, after.taxes.housing_tax.rate) == (0.99, 77) + assert system.get_parameters_at_instant(INSTANT) is after + + +def test_an_update_refreshes_every_cached_instant(): + """``update`` clears the root's cache at every instant, so the system + rebuilds each instant it had cached, not only the one the update names.""" + system = CountryTaxBenefitSystem() + other_instant = system.get_parameters_at_instant("2016-01-01") + system.get_parameters_at_instant(INSTANT) + + system.parameters.taxes.income_tax_rate.update(period="year:2015:1", value=0.99) + + assert system.get_parameters_at_instant("2016-01-01") is not other_instant + assert system.get_parameters_at_instant("2016-01-01").taxes.income_tax_rate == ( + system.parameters.taxes.income_tax_rate("2016-01-01") + ) From 62d63079e2d5c6a1f186d8a6f15d595bb1394860 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 15 Oct 2026 01:40:35 +0400 Subject: [PATCH 21/23] Keep nodes at instant consistent after changes, copies and forks Review found that lazy nodes made existing gaps in cache invalidation visible as results that depend on what was read before: - A freshly loaded scale never set its brackets' parent, so Parameter.update on a bracket's rate or threshold stopped clearing at the bracket. ParameterScale now sets it, as clone already did. - ParameterNode.add_child cleared no cache. It now clears its node's and every ancestor's. - New ParameterNode.clear_instant_caches clears a subtree and the nodes above it, for direct edits to values_list, brackets or children. TaxBenefitSystem calls it once its parameter operations, which edit values in place after reforms have run, are done. A Hypothesis property interleaves reads, updates (scale brackets included) and add_child, and checks that a node requested after the last change reads the tree as it is. Also: - Pickling or deep-copying a node resolves its remaining children first, so the copy is a snapshot, as before, rather than carrying the whole parameter tree through the unread children's parents. - A node is only built at a date. parameters(None), which Simulation._calculate reaches for eternal variables, used to fail on the first scale and cache nothing; a lazy node was cached at "None". - A forked child gets a fresh resolution lock, so a fork while another thread is resolving a child no longer deadlocks the child. Co-Authored-By: Claude Opus 5.5 --- .../parameter-cache-invalidation.fixed.md | 1 + .../parameters/parameter_node.py | 16 +++ .../parameters/parameter_node_at_instant.py | 44 ++++++- .../parameters/parameter_scale.py | 3 + .../taxbenefitsystems/tax_benefit_system.py | 4 + .../test_parameter_cache_invalidation.py | 118 ++++++++++++++++++ .../test_parameter_node_at_instant_lazy.py | 103 ++++++++++++++- ...test_parameter_node_at_instant_property.py | 64 +++++++++- 8 files changed, 346 insertions(+), 7 deletions(-) create mode 100644 changelog.d/parameter-cache-invalidation.fixed.md create mode 100644 tests/core/parameters/test_parameter_cache_invalidation.py diff --git a/changelog.d/parameter-cache-invalidation.fixed.md b/changelog.d/parameter-cache-invalidation.fixed.md new file mode 100644 index 000000000..937dcea13 --- /dev/null +++ b/changelog.d/parameter-cache-invalidation.fixed.md @@ -0,0 +1 @@ +`Parameter.update` on a bracket of a freshly loaded scale and `ParameterNode.add_child` clear the cached parameters of every node above them, `ParameterNode.clear_instant_caches` clears them after direct edits to values, brackets or children, and a tax-benefit system clears them once its parameter operations have run. diff --git a/policyengine_core/parameters/parameter_node.py b/policyengine_core/parameters/parameter_node.py index 0df0e387b..0bee65065 100644 --- a/policyengine_core/parameters/parameter_node.py +++ b/policyengine_core/parameters/parameter_node.py @@ -182,6 +182,8 @@ def add_child(self, name: str, child: Union["ParameterNode", Parameter]): self.children[name] = child setattr(self, name, child) child.parent = self + # Nodes at instant cached here and above were built without it. + self.clear_parent_cache() def __repr__(self) -> str: result = os.linesep.join( @@ -239,6 +241,20 @@ def mark_as_modified(self): if self.parent is not None: self.parent.mark_as_modified() + def clear_instant_caches(self) -> None: + """Drop the nodes at instant cached by this node, every node below it + and every node above it. + + ``Parameter.update`` and ``add_child`` clear what they change. Call + this after changing ``values_list``, a scale's ``brackets`` or + ``children`` directly, so that nodes at instant requested afterwards + read the change. + """ + self.clear_parent_cache() + for descendant in self.get_descendants(): + if isinstance(descendant, ParameterNode): + descendant._at_instant_cache.clear() + def get_child(self, path: str) -> "ParameterNode": node = self for name in path.split("."): diff --git a/policyengine_core/parameters/parameter_node_at_instant.py b/policyengine_core/parameters/parameter_node_at_instant.py index ceafdd58a..5699a92fe 100644 --- a/policyengine_core/parameters/parameter_node_at_instant.py +++ b/policyengine_core/parameters/parameter_node_at_instant.py @@ -5,7 +5,7 @@ import numpy -from policyengine_core import parameters, tools +from policyengine_core import parameters, periods, tools from policyengine_core.errors import ParameterNotFoundError from policyengine_core.parameters import helpers @@ -31,6 +31,18 @@ _RESOLUTION_LOCK = threading.RLock() +def _reset_resolution_lock() -> None: + """Give a forked child a lock no thread holds. A thread of the parent that + held it at the fork does not exist in the child, so nothing would release + it, and the child's first resolution would wait for ever.""" + global _RESOLUTION_LOCK + _RESOLUTION_LOCK = threading.RLock() + + +if hasattr(os, "register_at_fork"): + os.register_at_fork(after_in_child=_reset_resolution_lock) + + class _LazyChildren: """What a node at instant knows about its children. A shallow copy of the node shares it, as copies of a node that built every child up front shared @@ -55,6 +67,12 @@ def __init__(self, instant: str, source: dict): # Pickle protocols 0 and 1 refuse a class with ``__slots__`` that does not # define these, and a parameter tree pickles the nodes it has cached. def __getstate__(self) -> dict: + # A pickle or a deep copy is a snapshot, as when every child was built + # up front. Resolving what is left drops ``source``, whose parameter + # nodes reach the whole tree through their parents: the copy would + # otherwise carry the tree with it. + if self.source is not None: + _resolve_remaining(self) return {name: getattr(self, name) for name in self.__slots__} def __setstate__(self, state: dict) -> None: @@ -75,6 +93,14 @@ class ParameterNodeAtInstant: held across an in-place change to its parameters reads the changed value for a child it has not read yet; once read (including found to have no value), a child keeps what that read found. + + A node requested after a change reads the change: ``Parameter.update`` and + ``ParameterNode.add_child`` clear the cached nodes at instant of every + node above the one they change. After changing ``values_list``, a scale's + ``brackets`` or ``children`` directly, call + ``ParameterNode.clear_instant_caches``. A pickle or deep copy of a node + resolves every child first, so it is a snapshot that does not carry the + parameter tree. """ def __init__(self, name: str, node: "ParameterNode", instant_str: str): @@ -84,6 +110,12 @@ def __init__(self, name: str, node: "ParameterNode", instant_str: str): :param instant_str: A date in the format `YYYY-MM-DD`. """ + # Built up front, a node at an instant that is not a date (such as + # ``parameters(None)``, which reads "None") failed on its first scale + # and was never cached. Built lazily, it would be cached and fail, or + # read the latest values, only when a child is read or copied. + periods.instant(instant_str) + # The "technical" attributes are hidden, so that the node children can be easily browsed with auto-completion without pollution self._name = name self._instant_str = instant_str @@ -251,8 +283,8 @@ def _complete(state: _LazyChildren): state.absent = set() -def _materialise(node_at_instant: ParameterNodeAtInstant, state: _LazyChildren): - """Resolve every child and set ``_children``.""" +def _resolve_remaining(state: _LazyChildren): + """Resolve every child not read yet, and complete the state.""" with _RESOLUTION_LOCK: if state.children is None: for child_name in list(state.source): @@ -261,6 +293,12 @@ def _materialise(node_at_instant: ParameterNodeAtInstant, state: _LazyChildren): _resolve(state, child_name) if state.children is None: _complete(state) + + +def _materialise(node_at_instant: ParameterNodeAtInstant, state: _LazyChildren): + """Resolve every child and set ``_children``.""" + with _RESOLUTION_LOCK: + _resolve_remaining(state) children = state.children if "_children" not in node_at_instant.__dict__: # Unless a child named ``_children`` took the name. diff --git a/policyengine_core/parameters/parameter_scale.py b/policyengine_core/parameters/parameter_scale.py index 9b92a75b8..50d39d0d8 100644 --- a/policyengine_core/parameters/parameter_scale.py +++ b/policyengine_core/parameters/parameter_scale.py @@ -53,6 +53,9 @@ def __init__(self, name: str, data: dict, file_path: str): bracket = parameters.ParameterScaleBracket( name=bracket_name, data=bracket_data, file_path=file_path ) + # As ``clone`` does, so that ``Parameter.update`` on a bracket's + # rate or threshold clears the caches above the scale. + bracket.parent = self brackets.append(bracket) self.brackets: typing.List[parameters.ParameterScaleBracket] = brackets self.propagate_uprating() diff --git a/policyengine_core/taxbenefitsystems/tax_benefit_system.py b/policyengine_core/taxbenefitsystems/tax_benefit_system.py index 212de058d..29653ca19 100644 --- a/policyengine_core/taxbenefitsystems/tax_benefit_system.py +++ b/policyengine_core/taxbenefitsystems/tax_benefit_system.py @@ -154,6 +154,10 @@ def __init__(self, entities: Sequence[Entity] = None, reform=None) -> None: self.parameters = uprate_parameters(self.parameters) self.parameters = propagate_parameter_metadata(self.parameters) self.add_abolition_parameters() + # The operations above change values in place, after the reforms + # have run and possibly read the parameters. + self.parameters.clear_instant_caches() + self._parameters_at_instant_cache = {} self.add_modelled_policy_metadata() diff --git a/tests/core/parameters/test_parameter_cache_invalidation.py b/tests/core/parameters/test_parameter_cache_invalidation.py new file mode 100644 index 000000000..49ab9fe9e --- /dev/null +++ b/tests/core/parameters/test_parameter_cache_invalidation.py @@ -0,0 +1,118 @@ +"""A change to a parameter tree reaches every node at instant requested after it. + +``ParameterNode`` caches its node at each instant, and a node at instant +resolves its children on first read. A change therefore has to clear the +cache of every node above the one it changes, or a node requested afterwards +would read old values for children an earlier read had resolved and new ones +for the rest. These tests pin each way of changing a tree, with and without +an earlier read. +""" + +import pytest + +from policyengine_core.country_template import CountryTaxBenefitSystem +from policyengine_core.errors import ParameterNotFoundError +from policyengine_core.parameters import Parameter, ParameterNode +from policyengine_core.reforms import Reform +from tests.fixtures.parameter_nodes_at_instant import build_tree, materialise, snapshot + +INSTANT = "2017-03-01" + + +@pytest.mark.parametrize("read_before", [False, True]) +def test_a_bracket_update_reaches_a_node_requested_afterwards(read_before): + """A freshly loaded scale's brackets know their scale, so updating a + bracket's rate clears the caches above the scale too.""" + tree = build_tree() + if read_before: + assert tree.get_at_instant(INSTANT).scale.rates == [0.1, 0.2] + + rate = tree.children["scale"].brackets[1].children["rate"] + rate.update(period="year:2017:1", value=0.5) + + assert tree.get_at_instant(INSTANT).scale.rates == [0.1, 0.5] + + +@pytest.mark.parametrize("read_before", [False, True]) +def test_add_child_reaches_a_node_requested_afterwards(read_before): + tree = build_tree() + if read_before: + assert tree.get_at_instant(INSTANT).group.x == 0.1 + + tree.children["group"].add_child( + "added", Parameter("group.added", data={"values": {"2010-01-01": 9}}) + ) + + group = tree.get_at_instant(INSTANT).group + assert group.added == 9 + assert "added" in list(group) + + +@pytest.mark.parametrize("read_before", [False, True]) +def test_clear_instant_caches_after_a_direct_edit(read_before): + """``values_list`` edits bypass ``Parameter.update``; clearing the caches + afterwards makes them visible.""" + tree = build_tree() + if read_before: + assert tree.get_at_instant(INSTANT).group.x == 0.1 + + for value_at_instant in tree.children["group"].children["x"].values_list: + value_at_instant.value = 0.7 + tree.children["group"].clear_instant_caches() + + assert tree.get_at_instant(INSTANT).group.x == 0.7 + assert materialise(tree.get_at_instant(INSTANT)) == snapshot(tree, INSTANT) + + +def test_clear_instant_caches_clears_the_subtree_and_the_nodes_above(): + tree = build_tree() + materialise(tree.get_at_instant(INSTANT)) + group = tree.children["group"] + assert tree._at_instant_cache and group._at_instant_cache + + group.clear_instant_caches() + + assert not tree._at_instant_cache + assert not group._at_instant_cache + assert not group.children["inner"]._at_instant_cache + # A sibling's cache holds nothing that depends on the group. + assert tree.children["by_zone"]._at_instant_cache + + +def test_a_system_reads_values_changed_after_a_reform_read_them_during_construction(): + """Reforms applied while a system is built run before the parameter + operations, which change values in place; the system clears its caches + once they have run.""" + + class read_then_change_in_place(Reform): + def apply(self): + assert self.parameters(INSTANT).taxes.income_tax_rate == 0.15 + for value_at_instant in self.parameters.taxes.income_tax_rate.values_list: + value_at_instant.value = 0.4 + + system = CountryTaxBenefitSystem(reform=read_then_change_in_place) + + assert system.parameters(INSTANT).taxes.income_tax_rate == 0.4 + assert system.get_parameters_at_instant(INSTANT).taxes.income_tax_rate == 0.4 + + +def test_a_node_held_across_a_change_keeps_what_it_read(): + """The documented exception: a node requested before the change.""" + tree = build_tree() + held = tree.get_at_instant(INSTANT).group + assert held.x == 0.1 + + tree.children["group"].children["x"].update(period="year:2017:1", value=0.9) + + assert held.x == 0.1 + assert tree.get_at_instant(INSTANT).group.x == 0.9 + + +def test_add_child_still_rejects_a_name_in_use(): + tree = build_tree() + with pytest.raises(ValueError): + tree.children["group"].add_child( + "x", Parameter("group.x", data={"values": {"2010-01-01": 1}}) + ) + with pytest.raises(ParameterNotFoundError): + tree.get_at_instant(INSTANT).group.missing diff --git a/tests/core/parameters/test_parameter_node_at_instant_lazy.py b/tests/core/parameters/test_parameter_node_at_instant_lazy.py index 6f76623b0..5140b58c5 100644 --- a/tests/core/parameters/test_parameter_node_at_instant_lazy.py +++ b/tests/core/parameters/test_parameter_node_at_instant_lazy.py @@ -9,7 +9,10 @@ import copy import gc +import os import pickle +import select +import signal import threading import weakref @@ -188,14 +191,31 @@ def test_deep_copied_tree_keeps_its_own_nodes_at_instant(): tree.get_at_instant("2017-03-01").group.x # cache a partly read node clone = copy.deepcopy(tree) cached = clone._at_instant_cache["2017-03-01"] + # The copy's cached node is a snapshot: a change to the original tree + # does not reach it. + tree.flat_rate.update(value=0.5, period="year:2010:30") + assert cached.flat_rate == 0.1 + # A node requested from the copy after a change to the copy reads the + # copy, and the original is untouched. clone.group.y.update(value=0.9, period="year:2010:30") clone.flat_rate.update(value=0.7, period="year:2010:30") - # The copy's cached node reads the copy, not the original tree. - assert cached.flat_rate == 0.7 - assert tree.get_at_instant("2017-03-01").flat_rate == 0.1 + assert clone.get_at_instant("2017-03-01").flat_rate == 0.7 + assert clone.get_at_instant("2017-03-01").group.y == 0.9 + assert tree.get_at_instant("2017-03-01").flat_rate == 0.5 assert not hasattr(tree.get_at_instant("2017-03-01").group, "y") +def test_a_node_is_built_only_at_a_date(): + """Built up front, a node at "None" (``parameters(None)``) failed on its + first scale and was never cached; ``Simulation._calculate`` relies on + that when it checks an eternal variable's abolition.""" + tree = build_tree() + with pytest.raises(ValueError): + tree(None) + assert "None" not in tree._at_instant_cache + copy.deepcopy(tree) + + def test_copy_of_a_node_at_instant_reads_the_same_values(): at_instant = build_tree().get_at_instant("2017-03-01") at_instant.group @@ -326,6 +346,83 @@ def test_a_partly_read_node_pickles_under_every_protocol(protocol): assert ordered(materialise(restored)) == ordered(materialise(at_instant)) +@pytest.mark.parametrize( + "duplicate", [copy.deepcopy, lambda node: pickle.loads(pickle.dumps(node))] +) +def test_a_copied_node_is_a_snapshot_without_the_parameter_tree(duplicate): + """Unread children hold parameter nodes, whose parents reach the root: a + copy that kept them would carry, and keep alive, the whole tree.""" + tree = ParameterNode( + "", + data={ + "small": { + "x": {"values": {"2010-01-01": 1}}, + "y": {"values": {"2010-01-01": 2}}, + }, + "large": { + f"parameter_{index}": {"values": {"2010-01-01": index}} + for index in range(2000) + }, + }, + ) + small = tree.get_at_instant("2017-01-01").small + assert small.x == 1 + + copied = duplicate(small) + + assert copied.__dict__[_STATE].source is None + assert (copied.x, copied.y, list(copied)) == (1, 2, ["x", "y"]) + assert len(pickle.dumps(small)) * 20 < len(pickle.dumps(tree)) + tree_reference = weakref.ref(tree) + del tree, small + gc.collect() + assert tree_reference() is None + assert copied.y == 2 + + +@pytest.mark.skipif(not hasattr(os, "fork"), reason="needs os.fork") +@pytest.mark.filterwarnings("ignore::DeprecationWarning") +def test_a_forked_child_resolves_while_another_thread_was_resolving(): + """A thread holding the resolution lock at the fork does not exist in the + child, so the child needs a lock of its own.""" + resolving, release = threading.Event(), threading.Event() + + class SlowChild: + def _get_at_instant(self, instant): + resolving.set() + release.wait(10) + return 1 + + tree = ParameterNode("", data={"other": {"values": {"2010-01-01": 7}}}) + tree.children = {"slow": SlowChild(), **tree.children} + held = tree.get_at_instant("2017-01-01") + reader = threading.Thread(target=lambda: held.slow) + reader.start() + assert resolving.wait(10) + unrelated = ParameterNode("", data={"x": {"values": {"2010-01-01": 3}}}) + + read_end, write_end = os.pipe() + child = os.fork() + if child == 0: # pragma: no cover - runs in the forked child + try: + value = unrelated.get_at_instant("2017-01-01").x + os.write(write_end, str(value).encode()) + finally: + os._exit(0) + os.close(write_end) + try: + ready, _, _ = select.select([read_end], [], [], 10) + received = os.read(read_end, 16) if ready else b"" + finally: + release.set() + reader.join() + if not received: + os.kill(child, signal.SIGKILL) + os.waitpid(child, 0) + os.close(read_end) + assert received == b"3" + + def test_a_child_without_a_value_is_resolved_once(): tree = ParameterNode("", data={"late": {"values": {"2030-01-01": 1}}}) held = tree.get_at_instant("2020-01-01") diff --git a/tests/core/parameters/test_parameter_node_at_instant_property.py b/tests/core/parameters/test_parameter_node_at_instant_property.py index 5eeec9372..57d96ee53 100644 --- a/tests/core/parameters/test_parameter_node_at_instant_property.py +++ b/tests/core/parameters/test_parameter_node_at_instant_property.py @@ -17,7 +17,7 @@ st = hypothesis.strategies from policyengine_core.errors import ParameterNotFoundError -from policyengine_core.parameters import ParameterNodeAtInstant +from policyengine_core.parameters import Parameter, ParameterNodeAtInstant from tests.fixtures.parameter_nodes_at_instant import ( INSTANTS, build_tree, @@ -89,3 +89,65 @@ def test_reads_in_any_order_match_the_up_front_build(instant, sequence): ) assert ordered(materialise(root)) == ordered(expected_root) + + +# Parameters a change can update: plain values, a value deep in the tree, a +# zone of a vectorisable node, and a scale's brackets. +UPDATABLE = [ + lambda tree: tree.children["flat_rate"], + lambda tree: tree.children["group"].children["x"], + lambda tree: tree.children["group"].children["inner"].children["2"], + lambda tree: tree.children["by_zone"].children["z2"].children["amount"], + lambda tree: tree.children["scale"].brackets[0].children["threshold"], + lambda tree: tree.children["scale"].brackets[1].children["rate"], +] +# Nodes a change can add a child to. +EXTENDABLE = [(), ("group",), ("group", "inner")] + +steps = st.one_of( + st.tuples(st.just("read"), st.sampled_from(PATHS), st.sampled_from(NAMES)), + st.tuples( + st.just("update"), + st.integers(0, len(UPDATABLE) - 1), + st.sampled_from([0.5, 7, 2000]), + ), + st.tuples(st.just("add_child"), st.sampled_from(EXTENDABLE), st.just(None)), +) + + +@hypothesis.given( + instant=st.sampled_from(INSTANTS), sequence=st.lists(steps, max_size=25) +) +@hypothesis.settings(max_examples=300, deadline=None) +def test_a_node_requested_after_changes_reads_the_tree_as_it_is(instant, sequence): + """Reads, ``Parameter.update`` (scale brackets included) and + ``ParameterNode.add_child`` in any order: a node requested after the last + change reads exactly what the tree holds, whatever earlier reads resolved.""" + tree = build_tree() + added = 0 + for step, target, argument in sequence: + if step == "read": + node = tree.get_at_instant(instant) + expected = snapshot(tree, instant) + for name in target: + if name not in expected: + break + node, expected = node[name], expected[name] + else: + if argument in expected: + assert _comparable(getattr(node, argument)) == expected[argument] + elif step == "update": + UPDATABLE[target](tree).update(period="year:2000:100", value=argument) + else: + parent = tree + for name in target: + parent = parent.children[name] + added += 1 + parent.add_child( + f"added_{added}", + Parameter(f"added_{added}", data={"values": {"2000-01-01": added}}), + ) + + assert ordered(materialise(tree.get_at_instant(instant))) == ordered( + snapshot(tree, instant) + ) From b26aa0a7f76c9eb5487c468e30d390202fcf858a Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 15 Oct 2026 01:53:05 +0400 Subject: [PATCH 22/23] Clear parent caches without recursing, and stop at a cleared node add_child now clears the caches above the node it changes, which first happens while that node is still being loaded. A child named "parent" (policyengine-us has gov.hhs.medicaid.eligibility.categories.parent) replaces the node's parent attribute through add_child's setattr and points back at it until the node is attached, so the recursive walk looped until RecursionError and policyengine-us failed to load. clear_parent_cache now walks up iteratively and stops at a node it has already cleared; scales, which cache nothing, are passed through. Co-Authored-By: Claude Opus 5.5 --- .../parameters/parameter_node.py | 19 ++++++++++++--- .../test_parameter_cache_invalidation.py | 23 +++++++++++++++++++ 2 files changed, 39 insertions(+), 3 deletions(-) diff --git a/policyengine_core/parameters/parameter_node.py b/policyengine_core/parameters/parameter_node.py index 0bee65065..db758b4a5 100644 --- a/policyengine_core/parameters/parameter_node.py +++ b/policyengine_core/parameters/parameter_node.py @@ -232,9 +232,22 @@ def attach_to_parent(self, parent: "ParameterNode"): self.parent = parent def clear_parent_cache(self): - self._at_instant_cache.clear() - if self.parent is not None: - self.parent.clear_parent_cache() + """Clear the nodes at instant cached by this node and every node above + it (scales cache nothing and are passed through). + + Walks up without recursing and stops at a node it has already + cleared: while a node is loaded, a child named ``parent`` (as + policyengine-us's ``gov.hhs.medicaid.eligibility.categories.parent``) + replaces its ``parent`` attribute and points back at it until the node + is attached to its own parent. + """ + node, cleared = self, set() + while node is not None and id(node) not in cleared: + cleared.add(id(node)) + cache = getattr(node, "_at_instant_cache", None) + if cache is not None: + cache.clear() + node = getattr(node, "parent", None) def mark_as_modified(self): self.modified = True diff --git a/tests/core/parameters/test_parameter_cache_invalidation.py b/tests/core/parameters/test_parameter_cache_invalidation.py index 49ab9fe9e..ae25473f5 100644 --- a/tests/core/parameters/test_parameter_cache_invalidation.py +++ b/tests/core/parameters/test_parameter_cache_invalidation.py @@ -108,6 +108,29 @@ def test_a_node_held_across_a_change_keeps_what_it_read(): assert tree.get_at_instant(INSTANT).group.x == 0.9 +@pytest.mark.parametrize( + "parent_child", + [{"values": {"2010-01-01": 1}}, {"x": {"values": {"2010-01-01": 1}}}], + ids=["parameter", "node"], +) +def test_a_child_named_parent_does_not_trap_cache_clearing(parent_child): + """``add_child`` sets each child as an attribute, so a child named + ``parent`` replaces the node's ``parent`` and points back at the node + until the node is attached; clearing caches meanwhile must not loop.""" + tree = ParameterNode( + "", + data={ + "group": {"parent": parent_child, "later": {"values": {"2010-01-01": 2}}} + }, + ) + group = tree.children["group"] + assert group.parent is tree + + group.children["later"].update(period="year:2017:1", value=3) + + assert tree.get_at_instant(INSTANT).group.later == 3 + + def test_add_child_still_rejects_a_name_in_use(): tree = build_tree() with pytest.raises(ValueError): From b8aa6e6a4807eb65783c379741da7b72e01b122e Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 15 Oct 2026 01:53:10 +0400 Subject: [PATCH 23/23] Clear a modified tree's caches with ParameterNode.clear_instant_caches #570 now provides the walk, so modify_parameters uses it instead of its own. Co-Authored-By: Claude Opus 5.5 --- .../taxbenefitsystems/tax_benefit_system.py | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) diff --git a/policyengine_core/taxbenefitsystems/tax_benefit_system.py b/policyengine_core/taxbenefitsystems/tax_benefit_system.py index 64a9a0899..4ee9efd6e 100644 --- a/policyengine_core/taxbenefitsystems/tax_benefit_system.py +++ b/policyengine_core/taxbenefitsystems/tax_benefit_system.py @@ -58,13 +58,9 @@ def _clear_instant_caches(parameters) -> None: - """Drop the values every node of a parameter tree has cached at each - instant.""" - if not isinstance(parameters, ParameterNode): - return - for node in (parameters, *parameters.get_descendants()): - if isinstance(node, ParameterNode): - node._at_instant_cache = {} + """Drop the nodes at instant every node of a parameter tree has cached.""" + if isinstance(parameters, ParameterNode): + parameters.clear_instant_caches() class TaxBenefitSystem: