Stop the YAML test runner keeping every case's simulation and reform system

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.

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 <paths> -c <country_package>` 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).

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 <paths> -c <country_package>`,
+where `--reform-cache-size N` sets `reform_cache_size`.
+
 ## dump_simulation
 
 ```{eval-rst}

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

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
 

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]

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:

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