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..76101c289 --- /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. 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 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/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 b7073081d..79e7d94ca 100644 --- a/policyengine_core/taxbenefitsystems/tax_benefit_system.py +++ b/policyengine_core/taxbenefitsystems/tax_benefit_system.py @@ -50,6 +50,7 @@ ) from policyengine_core.periods import Instant, Period from policyengine_core.populations import GroupPopulation, Population +from policyengine_core.tracers import TracingParameterNodeAtInstant from policyengine_core.variables import Variable log = logging.getLogger(__name__) @@ -559,10 +560,23 @@ def get_parameters_at_instant(self, instant: Instant) -> ParameterNodeAtInstant: ) ) + parameters = self.parameters parameters_at_instant = self._parameters_at_instant_cache.get(instant) - if parameters_at_instant is None and self.parameters is not None: - parameters_at_instant = self.parameters.get_at_instant(str(instant)) + if parameters_at_instant is None and parameters is not None: + parameters_at_instant = parameters.get_at_instant(str(instant)) + if isinstance(parameters_at_instant, TracingParameterNodeAtInstant): + parameters_at_instant = parameters_at_instant.parameter_node_at_instant self._parameters_at_instant_cache[instant] = parameters_at_instant + if parameters_at_instant is not None and getattr(parameters, "trace", False): + # Wrapped on each read and never cached, as the tree's own reads + # are (``ParameterNode._get_at_instant``): a cached wrapper kept + # the tracer of whichever simulation first read this instant + # alive for the life of the system, and recorded later reads into + # it, while a node cached before tracing was turned on went + # untraced. + return TracingParameterNodeAtInstant( + parameters_at_instant, parameters.tracer, parameters.branch_name + ) return parameters_at_instant def get_package_metadata(self) -> dict: diff --git a/policyengine_core/tools/test_runner.py b/policyengine_core/tools/test_runner.py index 6b8d4aefc..c8df6c8cf 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): """ @@ -94,6 +103,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. """ @@ -183,6 +203,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", "") @@ -245,15 +266,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) @@ -288,6 +316,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() @@ -419,30 +474,71 @@ 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 + + # 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). 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 current_tax_benefit_system = baseline.clone() @@ -459,10 +555,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_reform_cache_property.py b/tests/core/tools/test_runner/test_reform_cache_property.py new file mode 100644 index 000000000..7abde339b --- /dev/null +++ b/tests/core/tools/test_runner/test_reform_cache_property.py @@ -0,0 +1,106 @@ +"""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. +""" + +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"]), + # 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): + 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..fd4795505 --- /dev/null +++ b/tests/core/tools/test_runner/test_runner_memory.py @@ -0,0 +1,372 @@ +"""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_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