Release each YAML case's simulation and its file's reformed systems

changelog.d/2741.changed.md

@@ -0,0 +1 @@
+The YAML test run releases each case's simulation and its file's reformed systems once they have run, so its peak memory falls from 16.1 GiB to 1.65 GiB and CI's YAML job finishes well within its time limit.

policyengine_uk/tests/conftest.py

@@ -14,6 +14,10 @@
 
 import pytest
 from policyengine_core.taxbenefitsystems import TaxBenefitSystem
+from policyengine_core.tools.test_runner import YamlItem
+from policyengine_core.tools.test_runner import (
+    _tax_benefit_system_cache as _yaml_systems,
+)
 
 from policyengine_uk.utils.aggregate_impact.dataset import load_pin, materialise
 
@@ -562,9 +566,42 @@ def _remove_variable_hooks():
     _malloc_trim = None
 
 
+# policyengine-core's YAML runner (``policyengine-core test``, CI's YAML job)
+# keeps every case's simulation and reformed system until the process ends.
+# pytest holds each collected item for the whole session, a YamlItem keeps
+# ``simulation`` and ``tax_benefit_system`` after its case has run, and
+# ``_tax_benefit_system_cache`` keeps a full clone of the system for each
+# combination of dotted parameter inputs (``gov.x.y: value``) without ever
+# evicting one. On 2d437ebba the 5,235 YAML cases peaked at 16.1 GiB of RSS
+# on a Mac with memory to spare, and the process took 107 s to exit after
+# printing its summary. On CI's 16 GB arm64 runners the job took 30 to 45
+# minutes and hit its time limit after every case had passed (#2698; run
+# 40830740956). Those logs record no memory figures, so a shortage of memory
+# is the likely cause rather than a measured one. release_yaml_case drops
+# each case's references to its simulation and system after it runs, and a
+# file's dotted-input systems after its last case. With it the same run
+# peaked at 1.65 GiB and exited 5 s after its summary, and CI's YAML step
+# took 7 min 25 s (run 40851484802). Systems without dotted inputs stay
+# cached, since files share them. Remove this once policyengine-core releases
+# the same (PolicyEngine/policyengine-core#569).
+def release_yaml_case(item, nextitem) -> None:
+    """Drop a YAML case's simulation and system, and, after a file's last
+    case, the cached systems built for dotted parameter inputs."""
+    if not isinstance(item, YamlItem):
+        return
+    item.simulation = None
+    item.tax_benefit_system = None
+    if nextitem is None or nextitem.path != item.path:
+        # Keys are (reforms, dotted inputs, extensions).
+        for systems in _yaml_systems.values():
+            for key in [key for key in systems if key[1]]:
+                del systems[key]
+
+
 @pytest.hookimpl(trylast=True)
 def pytest_runtest_teardown(item, nextitem):
     # trylast: after pytest's own teardown, which finalizes the file's fixtures
     # when the next test is in another file.
+    release_yaml_case(item, nextitem)
     if _malloc_trim is not None and (nextitem is None or nextitem.path != item.path):
         _malloc_trim(0)

policyengine_uk/tests/test_yaml_case_release.py

@@ -0,0 +1,109 @@
+"""The conftest releases each YAML case's simulation and reformed systems.
+
+policyengine-core's YAML runner keeps every case's simulation, and every
+system it built for dotted parameter inputs, until the process ends. The
+conftest's ``release_yaml_case`` drops a case's simulation and system after
+it runs, and a file's dotted-input systems after its last case, so the YAML
+run's peak memory falls from 16.1 GiB to 1.65 GiB. If core starts releasing
+them itself (PolicyEngine/policyengine-core#569), the first test skips with
+a note to remove the workaround. If core keeps the cache but changes its
+keys, the last test fails, because the release would then silently do
+nothing.
+"""
+
+from contextlib import contextmanager
+from pathlib import Path
+
+import pytest
+from policyengine_core.tools import test_runner
+from policyengine_core.tools.test_runner import YamlItem
+
+PLAIN = ("", "", frozenset())
+NAMED = ("policyengine_uk.reforms.example", "", frozenset())
+DOTTED = ("", "gov.x:1", frozenset())
+NAMED_AND_DOTTED = ("policyengine_uk.reforms.example", "gov.x:1", frozenset())
+
+
+def _conftest(request):
+    path = Path(__file__).with_name("conftest.py")
+    for plugin in request.config.pluginmanager.get_plugins():
+        if getattr(plugin, "__file__", None) and Path(plugin.__file__) == path:
+            return plugin
+    raise AssertionError(f"{path} is not loaded as a plugin")
+
+
+class _Baseline:
+    """A stand-in baseline system, which the cache holds weakly."""
+
+    def clone(self):
+        return _Baseline()
+
+
+@contextmanager
+def _cached(keys):
+    """Cache a system under each key for a stand-in baseline, then remove it."""
+    baseline = _Baseline()
+    test_runner._tax_benefit_system_cache[baseline] = {key: object() for key in keys}
+    try:
+        yield test_runner._tax_benefit_system_cache[baseline]
+    finally:
+        test_runner._tax_benefit_system_cache.pop(baseline, None)
+
+
+def _case(path: str) -> YamlItem:
+    # A case as the runner leaves it once it has run.
+    case = object.__new__(YamlItem)
+    case.path = Path(path)
+    case.simulation = object()
+    case.tax_benefit_system = object()
+    return case
+
+
+def test_core_still_keeps_yaml_cases():
+    if "teardown" in vars(YamlItem):
+        pytest.skip(
+            "policyengine-core's YamlItem now has its own teardown "
+            "(PolicyEngine/policyengine-core#569). If it releases each case's "
+            "simulation and bounds _tax_benefit_system_cache, remove "
+            "release_yaml_case from conftest.py, and this file."
+        )
+
+
+def test_a_case_releases_its_simulation_and_system(request):
+    release_yaml_case = _conftest(request).release_yaml_case
+    case = _case("a.yaml")
+    release_yaml_case(case, _case("a.yaml"))
+    assert case.simulation is None
+    assert case.tax_benefit_system is None
+
+
+def test_a_files_last_case_drops_only_its_dotted_input_systems(request):
+    release_yaml_case = _conftest(request).release_yaml_case
+    with _cached([PLAIN, NAMED, DOTTED, NAMED_AND_DOTTED]) as systems:
+        # Later cases in the same file may reuse the file's systems.
+        release_yaml_case(_case("a.yaml"), _case("a.yaml"))
+        assert set(systems) == {PLAIN, NAMED, DOTTED, NAMED_AND_DOTTED}
+        release_yaml_case(_case("a.yaml"), _case("b.yaml"))
+        assert set(systems) == {PLAIN, NAMED}
+    with _cached([PLAIN, DOTTED]) as systems:
+        release_yaml_case(_case("a.yaml"), None)
+        assert set(systems) == {PLAIN}
+
+
+def test_other_tests_release_nothing(request):
+    release_yaml_case = _conftest(request).release_yaml_case
+    with _cached([PLAIN, DOTTED]) as systems:
+        release_yaml_case(request.node, None)
+        assert set(systems) == {PLAIN, DOTTED}
+
+
+def test_the_cache_keys_hold_the_dotted_inputs_second():
+    # release_yaml_case reads a key's second part as its dotted inputs.
+    with _cached([]) as systems:
+        (baseline,) = [
+            key
+            for key, value in test_runner._tax_benefit_system_cache.items()
+            if value is systems
+        ]
+        test_runner._get_tax_benefit_system(baseline, [], [], reform_key="gov.x:1")
+        assert set(systems) == {DOTTED}