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..612ddb868 --- /dev/null +++ b/changelog.d/fix-test-runner-memory.fixed.md @@ -0,0 +1 @@ +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. A miss on the reform-free system now builds with automatic collection paused too, and the full-collection limit counts each baseline's reform-free system plus its cached reform systems. diff --git a/docs/engineering/skills/testing.md b/docs/engineering/skills/testing.md index 86e214485..ecf1ef527 100644 --- a/docs/engineering/skills/testing.md +++ b/docs/engineering/skills/testing.md @@ -36,3 +36,33 @@ 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. +- 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, 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 +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/parameters/parameter_node.py b/policyengine_core/parameters/parameter_node.py index 8338716f2..aefeca6dd 100644 --- a/policyengine_core/parameters/parameter_node.py +++ b/policyengine_core/parameters/parameter_node.py @@ -215,17 +215,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/scripts/policyengine_command.py b/policyengine_core/scripts/policyengine_command.py index 7aa85c888..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() @@ -91,6 +98,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=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.", + ) 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/taxbenefitsystems/tax_benefit_system.py b/policyengine_core/taxbenefitsystems/tax_benefit_system.py index ccadd2a39..946a4e359 100644 --- a/policyengine_core/taxbenefitsystems/tax_benefit_system.py +++ b/policyengine_core/taxbenefitsystems/tax_benefit_system.py @@ -54,6 +54,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 @@ -659,10 +660,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 80adfa50b..bad744580 100644 --- a/policyengine_core/tools/test_runner.py +++ b/policyengine_core/tools/test_runner.py @@ -69,9 +69,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): """ @@ -96,6 +105,17 @@ 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. Each system still caches its + parameters at every distinct instant the cases read. """ @@ -185,6 +205,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", "") @@ -247,15 +268,22 @@ 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"), ) verbose = self.options.get("verbose") performance_graph = self.options.get("performance_graph") 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) @@ -290,6 +318,33 @@ 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, 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: + # 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 + self.simulation = None + self.tax_benefit_system = None + self._system_simulation = None + def print_computation_log(self, tracer): print("Computation log:") # noqa T001 tracer.print_computation_log() @@ -422,12 +477,23 @@ def pytest_collect_file(self, parent, file_path): # A cache miss forces a full collection only while the cache holds fewer than -# this many systems. That collection walks every object in the process, so -# each cached system makes it dearer, while the automatic collector's full -# collections grow rarer as the heap grows. In PR #643's model of -# policyengine-us-sized systems, 3 gave the least collection CPU over 16 and -# 32 misses; limits from 1 to 8 came within 20%, and no limit cost about three -# times as much. +# this many systems, counting each baseline's reform-free system and its +# reform systems across every baseline. That collection walks every object in +# the process, so each cached system makes it dearer, while the automatic +# collector's full collections grow rarer as the heap grows. In PR #643's +# model of policyengine-us-sized systems, 3 gave the least collection CPU over +# 16 and 32 misses; limits from 1 to 8 came within 20%, and no limit cost +# about three times as much. Treat 3 as a practical bound, not a proven best +# value: the model's ranking of limits moves with the heap's shape. +# +# While the cache only grows, a run therefore forces at most 3 collections. +# Eviction can keep it below 3 for good, and then every miss forces one. With +# the default reform cache, a baseline whose run builds the reform-free system +# and two reform systems holds 3 from then on, so later misses leave the copy +# to the automatic collector. It stays below 3 when the run has no reform-free +# case, when a call's ``reform_cache_size`` is 0 or 1, or after a test hook +# deletes reform systems. A forced collection walks at most two cached +# systems besides the new one and the baselines themselves. _FULL_COLLECTION_CACHE_LIMIT = 3 @@ -468,35 +534,116 @@ def _automatic_gc_paused(collect=True): gc.enable() +class _CachedSystems: + """One baseline's cached systems: the reform-free one, kept for the + baseline's lifetime, and an LRU of reform systems. + + Before the bound, each baseline's entry was a dict from a (reforms, + reform_key, extensions) key to a system, and country packages' test hooks + walk it to drop systems early: policyengine-uk's ``conftest.py`` iterates + each entry's keys and deletes those with a ``reform_key``. Iterating, + indexing, ``in`` and ``del`` therefore act on the reform systems, keyed + as before, so such hooks keep working. The reform-free system is not + among them. + """ + + def __init__(self): + self.reform_free = None + self.reforms = OrderedDict() + + def __iter__(self): + return iter(self.reforms) + + def __len__(self): + return len(self.reforms) + + def __contains__(self, key): + return key in self.reforms + + def __getitem__(self, key): + return self.reforms[key] + + def __delitem__(self, key): + del self.reforms[key] + + +def _cached_system_count(): + """How many systems the cache holds across every baseline: each + baseline's reform-free system, once built, plus its reform systems.""" + return sum( + (entry.reform_free is not None) + len(entry.reforms) + for entry in _tax_benefit_system_cache.values() + ) + + 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 + + # 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: + # A miss: build paused, as for a reform system below. + with _automatic_gc_paused( + collect=_cached_system_count() < _FULL_COLLECTION_CACHE_LIMIT + ): + reform_free = baseline.clone() + cached_systems.reform_free = reform_free + return cached_systems.reform_free + + # 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), ) - 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 = reform_systems.get(inner_key) + if cached is not None: + reform_systems.move_to_end(inner_key) + return cached # Only a cache miss builds a system, so only a miss touches collection. - # The count covers every baseline's entries. - cached_systems = sum(len(entry) for entry in _tax_benefit_system_cache.values()) - with _automatic_gc_paused(collect=cached_systems < _FULL_COLLECTION_CACHE_LIMIT): + # The count covers every baseline's systems, after this call's bound has + # trimmed this baseline's reform systems and before the new one goes in. + with _automatic_gc_paused( + collect=_cached_system_count() < _FULL_COLLECTION_CACHE_LIMIT + ): current_tax_benefit_system = baseline.clone() for reform_path in reforms: @@ -512,10 +659,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: + 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/test_parameter_tracing_cache.py b/tests/core/test_parameter_tracing_cache.py new file mode 100644 index 000000000..5b6d73b0e --- /dev/null +++ b/tests/core/test_parameter_tracing_cache.py @@ -0,0 +1,109 @@ +"""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. 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 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(): + 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 + + +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/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): diff --git a/tests/core/tools/test_runner/test_gc_pause.py b/tests/core/tools/test_runner/test_gc_pause.py index 1d70a02ac..f8f0ea308 100644 --- a/tests/core/tools/test_runner/test_gc_pause.py +++ b/tests/core/tools/test_runner/test_gc_pause.py @@ -1,13 +1,20 @@ """The YAML runner pauses automatic collection while it builds a reform system. -A cache miss in ``_get_tax_benefit_system`` builds with automatic collection -paused. If the caller had collection on, the build returned, and the cache -held fewer than ``_FULL_COLLECTION_CACHE_LIMIT`` systems, it then runs one -full collection; otherwise it leaves the new system to the automatic -collector, so a run of misses forces at most that many collections. It leaves -``gc.isenabled()`` as it found it, whether the build returns or raises and -whether the caller had collection on or off. A cache hit does not touch -collection. The pause changes no result: the cache still returns the +A cache miss in ``_get_tax_benefit_system``, for the reform-free system or a +reform system (including one the cache evicted, or any build with the reform +cache disabled), builds with automatic collection paused. If the caller had +collection on, the build returned, and the cache held fewer than +``_FULL_COLLECTION_CACHE_LIMIT`` systems (each baseline's reform-free system +plus its reform systems, across every baseline, after the call's bound trims +and before the new system goes in), it then runs one full collection; +otherwise it leaves the new system to the automatic collector. While the cache +only grows, a run of misses therefore forces at most that many collections. +Eviction can keep the cache below the limit for good (a ``reform_cache_size`` +of 0 or 1 on one baseline holds at most 2 systems), and then every miss forces +one. It leaves ``gc.isenabled()`` as it found it, whether the build returns or +raises (any ``BaseException``, including ``KeyboardInterrupt`` and a +``gc.collect()`` that raises) and whether the caller had collection on or +off. A cache hit does not touch collection. The pause changes no result: the cache still returns the identical system for a key it has seen and builds a new one for any other key, and YAML cases produce the same outcomes and values with or without it. """ @@ -16,6 +23,7 @@ import gc import os import weakref +from collections import OrderedDict import pytest @@ -39,27 +47,34 @@ class BuildFailed(Exception): class RecordingSystem: """A stand-in system that records whether collection was enabled at each - build step, and can be told to raise at one of them.""" + build step, and can be told to raise ``error`` at one of them.""" - def __init__(self, log, fail_at=None, reforms=(), extensions=()): + def __init__(self, log, fail_at=None, reforms=(), extensions=(), error=BuildFailed): self.log = log self.fail_at = fail_at self.reforms = reforms self.extensions = extensions + self.error = error def _step(self, name): self.log.append((name, gc.isenabled())) if self.fail_at == name: - raise BuildFailed(name) + raise self.error(name) def clone(self): self._step("clone") - return RecordingSystem(self.log, self.fail_at, self.reforms, self.extensions) + return RecordingSystem( + self.log, self.fail_at, self.reforms, self.extensions, self.error + ) def apply_reform(self, reform_path): self._step("reform") return RecordingSystem( - self.log, self.fail_at, self.reforms + (reform_path,), self.extensions + self.log, + self.fail_at, + self.reforms + (reform_path,), + self.extensions, + self.error, ) def load_extension(self, extension): @@ -70,7 +85,11 @@ def load_extension(self, extension): def callable_reform(system): system._step("reform") return RecordingSystem( - system.log, system.fail_at, system.reforms + ("callable",), system.extensions + system.log, + system.fail_at, + system.reforms + ("callable",), + system.extensions, + system.error, ) @@ -216,9 +235,15 @@ def test_the_pause_restores_the_callers_state_when_the_body_raises(gc_calls): assert collections(gc_calls) == [] -def miss(baseline, number): +def miss(baseline, number, cache_size=None): """A cache miss: a parameter key the cache has not seen.""" - return _get_tax_benefit_system(baseline, [], [], reform_key=f"gov.rate:{number}") + return _get_tax_benefit_system( + baseline, [], [], reform_key=f"gov.rate:{number}", cache_size=cache_size + ) + + +# A reform cache large enough that the tests below never evict. +NO_EVICTION = 100 @pytest.mark.parametrize("limit", [0, 1, 2, LIMIT, LIMIT + 2]) @@ -231,7 +256,7 @@ def test_misses_collect_only_while_the_cache_holds_fewer_systems_than_the_limit( collected = [] for number in range(limit + 3): del gc_calls[:] - miss(baseline, number) + miss(baseline, number, cache_size=NO_EVICTION) collected.append(collections(gc_calls)) assert gc.isenabled() # The last forced collection is miss number limit - 1, with limit - 1 @@ -245,12 +270,14 @@ def test_a_miss_past_the_limit_still_pauses_and_restores_collection(gc_calls): for enabled_before in (True, False): set_gc(True) for number in range(LIMIT): - miss(baseline, number) + miss(baseline, number, cache_size=NO_EVICTION) del log[:] set_gc(enabled_before) del gc_calls[:] - system = _get_tax_benefit_system(baseline, ["some.reform"], ["x.extension"]) + system = _get_tax_benefit_system( + baseline, ["some.reform"], ["x.extension"], cache_size=NO_EVICTION + ) assert [name for name, _ in log] == ["clone", "reform", "clone", "extension"] assert not any(enabled for _, enabled in log) @@ -280,6 +307,186 @@ def test_the_limit_counts_systems_cached_for_every_baseline(gc_calls, monkeypatc assert collections(gc_calls) == [FULL] +def reform_free(baseline, cache_size=None): + return _get_tax_benefit_system(baseline, [], [], cache_size=cache_size) + + +def test_a_reform_free_miss_builds_paused_and_a_hit_leaves_collection_alone( + gc_calls, +): + set_gc(True) + log = [] + baseline = RecordingSystem(log) + del gc_calls[:] + + system = reform_free(baseline) + + assert log == [("clone", False)] + assert gc_calls == ["disable", FULL, "enable"] + assert gc.isenabled() + del gc_calls[:] + assert reform_free(baseline) is system + assert reform_free(baseline, cache_size=0) is system + assert len(log) == 1 + assert gc_calls == [] + + +def test_a_count_of_two_collects_and_a_count_of_three_does_not(gc_calls): + """With the default reform cache, the reform-free system and one reform + system make 2 (a miss collects); with a second reform system, 3 (a miss + does not).""" + assert LIMIT == 3 + set_gc(True) + baseline = RecordingSystem([]) + reform_free(baseline) + miss(baseline, 0) + assert test_runner._cached_system_count() == 2 + del gc_calls[:] + miss(baseline, 1) + assert gc_calls == ["disable", FULL, "enable"] + assert test_runner._cached_system_count() == 3 + del gc_calls[:] + miss(baseline, 2) # Evicts miss 0 after the build. + assert gc_calls == ["disable", "enable"] + assert test_runner._cached_system_count() == 3 + + +def test_a_miss_on_an_evicted_system_rebuilds_paused(gc_calls): + set_gc(True) + log = [] + baseline = RecordingSystem(log) + first = miss(baseline, 0, cache_size=1) + miss(baseline, 1, cache_size=1) # Evicts miss 0. + del log[:] + del gc_calls[:] + + rebuilt = miss(baseline, 0, cache_size=1) + + assert rebuilt is not first + assert log == [("clone", False)] + # One reform system cached and no reform-free one: 1 < 3, so it collects. + assert gc_calls == ["disable", FULL, "enable"] + assert gc.isenabled() + + # With the reform-free system and two reform systems cached, the + # rebuild of an evicted one pauses but leaves collection to the + # automatic collector. + reform_free(baseline) + miss(baseline, 1, cache_size=2) + miss(baseline, 2, cache_size=2) # Evicts miss 0. + assert test_runner._cached_system_count() == 3 + del log[:] + del gc_calls[:] + miss(baseline, 0, cache_size=2) + assert log == [("clone", False)] + assert gc_calls == ["disable", "enable"] + assert gc.isenabled() + + +@pytest.mark.parametrize("enabled_before", [True, False]) +def test_with_the_reform_cache_disabled_every_miss_builds_paused( + gc_calls, enabled_before +): + set_gc(True) + log = [] + baseline = RecordingSystem(log) + reform_free(baseline, cache_size=0) + systems = [] + for _ in range(LIMIT + 2): + del log[:] + set_gc(enabled_before) + del gc_calls[:] + systems.append(miss(baseline, 0, cache_size=0)) + assert log == [("clone", False)] + assert gc.isenabled() is enabled_before + # The cache holds only the reform-free system, below the limit, so + # every miss collects for a caller with collection on. + expected = [FULL] if enabled_before else [] + assert collections(gc_calls) == expected + assert len({id(system) for system in systems}) == len(systems) + assert test_runner._cached_system_count() == 1 + + +def test_the_count_trims_to_this_calls_bound_first(gc_calls): + set_gc(True) + baseline = RecordingSystem([]) + for number in range(LIMIT): + miss(baseline, number, cache_size=NO_EVICTION) + assert test_runner._cached_system_count() == LIMIT + del gc_calls[:] + # A bound of 1 leaves one reform system before this miss counts, so + # the count is 1 and the miss collects. + miss(baseline, LIMIT, cache_size=1) + assert collections(gc_calls) == [FULL] + assert test_runner._cached_system_count() == 1 + + +class Interrupted(KeyboardInterrupt): + pass + + +@pytest.mark.parametrize("path", ["reform_free", "reform"]) +@pytest.mark.parametrize("enabled_before", [True, False]) +@pytest.mark.parametrize("error", [BuildFailed, Interrupted]) +def test_a_build_that_raises_restores_the_callers_state_on_both_paths( + gc_calls, path, enabled_before, error +): + baseline = RecordingSystem([], fail_at="clone", error=error) + set_gc(enabled_before) + del gc_calls[:] + with pytest.raises(error): + if path == "reform_free": + reform_free(baseline) + else: + miss(baseline, 0) + assert gc.isenabled() is enabled_before + assert collections(gc_calls) == [] + assert test_runner._cached_system_count() == 0 + + +@pytest.mark.parametrize("enabled_before", [True, False]) +def test_the_pause_restores_the_callers_state_after_a_keyboard_interrupt( + gc_calls, enabled_before +): + set_gc(enabled_before) + del gc_calls[:] + with pytest.raises(KeyboardInterrupt): + with _automatic_gc_paused(): + assert not gc.isenabled() + raise KeyboardInterrupt + assert gc.isenabled() is enabled_before + assert collections(gc_calls) == [] + + +class CollectFailed(Exception): + pass + + +@pytest.mark.parametrize("error", [CollectFailed, Interrupted]) +@pytest.mark.parametrize("path", ["pause", "reform_free", "reform"]) +def test_a_collection_that_raises_still_restores_collection( + gc_calls, monkeypatch, path, error +): + def failing_collect(generation=2): + gc_calls.append(f"collect({generation})") + raise error("collect") + + monkeypatch.setattr(gc, "collect", failing_collect) + set_gc(True) + baseline = RecordingSystem([]) + del gc_calls[:] + with pytest.raises(error): + if path == "pause": + with _automatic_gc_paused(): + pass + elif path == "reform_free": + reform_free(baseline) + else: + miss(baseline, 0) + assert gc_calls == ["disable", FULL, "enable"] + assert gc.isenabled() + + REFORM_NAMES = ["a.reform", "b.reform", "c.reform"] EXTENSION_NAMES = ["x.extension", "y.extension"] @@ -288,6 +495,8 @@ def test_the_limit_counts_systems_cached_for_every_baseline(gc_calls, monkeypatc st.sampled_from([None, "", "gov.rate:1", "gov.rate:2"]), st.lists(st.sampled_from(EXTENSION_NAMES), max_size=2), st.booleans(), + # Each call's reform-cache bound, from disabled to never evicting. + st.sampled_from([0, 1, 2, 3, None, NO_EVICTION]), ) @@ -297,7 +506,7 @@ def test_the_limit_counts_systems_cached_for_every_baseline(gc_calls, monkeypatc suppress_health_check=[hypothesis.HealthCheck.function_scoped_fixture], ) @hypothesis.given(requests=st.lists(request_strategy, max_size=25)) -def test_the_cache_matches_a_reference_dict_and_keeps_the_callers_state( +def test_the_cache_matches_a_reference_lru_and_keeps_the_callers_state( gc_calls, monkeypatch, requests ): # Each example starts from an empty cache. @@ -306,36 +515,71 @@ def test_the_cache_matches_a_reference_dict_and_keeps_the_callers_state( ) log = [] baseline = RecordingSystem(log) - reference = {} + # What the cache should hold: the reform-free system and an LRU of + # reform systems. + reform_free_system = None + reference = OrderedDict() forced = 0 - for reforms, reform_key, extensions, enabled_before in requests: + ever_dropped = False + for reforms, reform_key, extensions, enabled_before, cache_size in requests: + bound = ( + test_runner.DEFAULT_REFORM_CACHE_SIZE if cache_size is None else cache_size + ) + is_reform_free = not reforms and not extensions and not reform_key key = (tuple(reforms), reform_key, frozenset(extensions)) + # The call trims to its bound first; the count is taken after that. + while len(reference) > bound: + reference.popitem(last=False) + ever_dropped = True + count_before = (reform_free_system is not None) + len(reference) clones_before = sum(name == "clone" for name, _ in log) set_gc(enabled_before) del gc_calls[:] system = _get_tax_benefit_system( - baseline, list(reforms), list(extensions), reform_key=reform_key + baseline, + list(reforms), + list(extensions), + reform_key=reform_key, + cache_size=cache_size, ) assert gc.isenabled() is enabled_before assert all(not enabled for _, enabled in log) built = sum(name == "clone" for name, _ in log) > clones_before - # The reference dict holds what the cache held before this call. - collects = built and enabled_before and len(reference) < LIMIT + collects = built and enabled_before and count_before < LIMIT assert collections(gc_calls) == ([FULL] if collects else []) forced += collects - assert forced <= LIMIT - if key in reference: + # The bound on forced collections holds while the cache only grows. + if not ever_dropped: + assert forced <= LIMIT + if is_reform_free: + if reform_free_system is None: + assert built + reform_free_system = system + assert system is reform_free_system + assert not built or gc_calls[0] == "disable" + elif key in reference: + reference.move_to_end(key) assert system is reference[key] assert not built assert gc_calls == [] else: assert built + assert system is not reform_free_system assert all(system is not other for other in reference.values()) assert system.reforms == tuple(reforms) assert sorted(system.extensions) == sorted(extensions) - reference[key] = system + if bound > 0: + reference[key] = system + if len(reference) > bound: + reference.popitem(last=False) + ever_dropped = True + else: + ever_dropped = True + assert test_runner._cached_system_count() == ( + (reform_free_system is not None) + len(reference) + ) yaml_tests_dir = os.path.dirname(yaml_tests.__file__) 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..63951e469 --- /dev/null +++ b/tests/core/tools/test_runner/test_reform_cache_property.py @@ -0,0 +1,118 @@ +"""The YAML runner's reform-system cache against a reference LRU. + +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. + +A cache miss can end with a full ``gc.collect()`` (see ``test_gc_pause.py``), +which walks every object in the test process and never changes what the cache +returns. Each example replaces it with a no-op: real collections made this +file the slowest in the suite, about 20 minutes under CI's coverage tracing. +""" + +from __future__ import annotations + +import gc +from collections import OrderedDict +from unittest import mock + +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"]), + # 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)) +# ``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): + with mock.patch.object(gc, "collect", lambda generation=2: 0): + check_against_reference_lru(sequence) + + +def check_against_reference_lru(sequence): + baseline = StubSystem() + reference = OrderedDict() + last_returned = {} + reform_free = None + + 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), + 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 = (tuple(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..297e7580d --- /dev/null +++ b/tests/core/tools/test_runner/test_runner_memory.py @@ -0,0 +1,395 @@ +"""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 +import yaml + +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 policyengine_core.tools.test_runner import OpenFiscaPlugin +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, +) +from tests.fixtures.tracing import build_simulation + + +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_a_baseline_entry_still_reads_as_a_mapping_of_reform_systems(): + # policyengine-uk's conftest.py drops a file's dotted-input systems after + # its last case by walking each baseline's entry like the old dict: + # for key in [key for key in systems if key[1]]: del systems[key] + baseline = CountryTaxBenefitSystem() + reform_free = _get_tax_benefit_system(baseline, [], []) + first = _system_for_rate(baseline, 0.21) + _system_for_rate(baseline, 0.22) + systems = _tax_benefit_system_cache[baseline] + keys = list(systems) + assert len(systems) == len(keys) == 2 + assert all(key[1] for key in keys) + assert systems[keys[0]] is first + assert keys[0] in systems + + for key in [key for key in systems if key[1]]: + del systems[key] + + assert len(systems) == 0 and list(systems) == [] + assert _get_tax_benefit_system(baseline, [], []) is reform_free + assert _system_for_rate(baseline, 0.21) is not first + + +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_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) + + +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 + + +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 _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 + + +@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": invalid_input, + "output": {"income_tax": 0}, + } + for i in range(8) + ] + 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(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): + """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 + 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] + + +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)] +) +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 diff --git a/tests/fixtures/yaml_runner_memory.py b/tests/fixtures/yaml_runner_memory.py new file mode 100644 index 000000000..fb6264a0d --- /dev/null +++ b/tests/fixtures/yaml_runner_memory.py @@ -0,0 +1,203 @@ +"""Helpers for the YAML runner memory tests: generated case files and a pytest +plugin that measures what a run keeps alive.""" + +import gc +import sys +import tracemalloc +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 +from policyengine_core.tools.test_runner import OpenFiscaPlugin + +# 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: + """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 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() + 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 that the run created (not ones other tests left behind).""" + + def __init__(self, sample_every: int = 10): + self.sample_every = sample_every + self.items = [] + self.traced = [] + 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): + 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() + # 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"): + 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, 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 if outcome == "passed" else 1), exit_code + assert probe.outcomes and set(probe.outcomes) == {outcome} + return probe