From fe2237541373dd0eee91923fd046367adc53ead3 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Thu, 29 Oct 2026 05:40:09 -0400 Subject: [PATCH 1/2] Release each YAML case's simulation and its file's reformed systems policyengine-core's YAML runner keeps every case's simulation, and a full system clone for each combination of dotted parameter inputs, until the process ends. On main the 5,235 YAML cases peaked at 16.5 GB and took 107 s to exit; on CI's 16 GB runners the job swaps and runs into its time limit after its tests pass (#2741, #2698). The conftest's teardown hook now drops each case's simulation and system, and a file's dotted-input systems after its last case. The same local run peaks at 1.7 GB and takes 188 s in all instead of 277. Remove it once policyengine-core releases the same (PolicyEngine/policyengine-core#569). Fixes #2741 Co-Authored-By: Claude Opus 5.5 --- changelog.d/2741.changed.md | 1 + policyengine_uk/tests/conftest.py | 35 ++++++ .../tests/test_yaml_case_release.py | 108 ++++++++++++++++++ 3 files changed, 144 insertions(+) create mode 100644 changelog.d/2741.changed.md create mode 100644 policyengine_uk/tests/test_yaml_case_release.py diff --git a/changelog.d/2741.changed.md b/changelog.d/2741.changed.md new file mode 100644 index 0000000000..31b6500746 --- /dev/null +++ b/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 CI's YAML job no longer runs out of memory. diff --git a/policyengine_uk/tests/conftest.py b/policyengine_uk/tests/conftest.py index 377bfa8088..5af1a25410 100644 --- a/policyengine_uk/tests/conftest.py +++ b/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,40 @@ 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.5 GB on a Mac +# with memory to spare, and the process took 107 s to exit after printing +# its summary. The 16 GB arm64 runners swap, so the job took 30 to 45 minutes +# and hit its time limit after its tests had passed (#2698; run 40830740956). +# Releasing each case's simulation and system after it runs, and a file's +# dotted-input systems after its last case, the same run peaked at 1.7 GB and +# took 188 s in all instead of 277: 179 s of tests against 168, and 5 s to +# exit. 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) diff --git a/policyengine_uk/tests/test_yaml_case_release.py b/policyengine_uk/tests/test_yaml_case_release.py new file mode 100644 index 0000000000..7c5ed9dad5 --- /dev/null +++ b/policyengine_uk/tests/test_yaml_case_release.py @@ -0,0 +1,108 @@ +"""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 CI's YAML +job stays within its runner's memory. 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} From 1325e52ff296613aef12112d56c32edc4f58ad6e Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Thu, 29 Oct 2026 06:28:08 -0400 Subject: [PATCH 2/2] State the YAML run's memory in GiB and the CI cause as likely Review of fe2237541: the peak files hold MiB, so the run peaked at 16.1 GiB on main and 1.65 GiB here, not 16.5 GB and 1.7 GB. The CI logs record no memory figures, so the comment now gives a memory shortage as the likely cause of the timeouts rather than a measured one, and adds this PR's CI YAML step (7 min 25 s). Co-Authored-By: Claude Opus 5.5 --- changelog.d/2741.changed.md | 2 +- policyengine_uk/tests/conftest.py | 22 ++++++++++--------- .../tests/test_yaml_case_release.py | 11 +++++----- 3 files changed, 19 insertions(+), 16 deletions(-) diff --git a/changelog.d/2741.changed.md b/changelog.d/2741.changed.md index 31b6500746..b4aa5ce3a3 100644 --- a/changelog.d/2741.changed.md +++ b/changelog.d/2741.changed.md @@ -1 +1 @@ -The YAML test run releases each case's simulation and its file's reformed systems once they have run, so CI's YAML job no longer runs out of memory. +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. diff --git a/policyengine_uk/tests/conftest.py b/policyengine_uk/tests/conftest.py index 5af1a25410..d945afd3e5 100644 --- a/policyengine_uk/tests/conftest.py +++ b/policyengine_uk/tests/conftest.py @@ -572,16 +572,18 @@ def _remove_variable_hooks(): # ``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.5 GB on a Mac -# with memory to spare, and the process took 107 s to exit after printing -# its summary. The 16 GB arm64 runners swap, so the job took 30 to 45 minutes -# and hit its time limit after its tests had passed (#2698; run 40830740956). -# Releasing each case's simulation and system after it runs, and a file's -# dotted-input systems after its last case, the same run peaked at 1.7 GB and -# took 188 s in all instead of 277: 179 s of tests against 168, and 5 s to -# exit. Systems without dotted inputs stay cached, since files share them. -# Remove this once policyengine-core releases the same -# (PolicyEngine/policyengine-core#569). +# 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.""" diff --git a/policyengine_uk/tests/test_yaml_case_release.py b/policyengine_uk/tests/test_yaml_case_release.py index 7c5ed9dad5..4f1b7b5d9a 100644 --- a/policyengine_uk/tests/test_yaml_case_release.py +++ b/policyengine_uk/tests/test_yaml_case_release.py @@ -3,11 +3,12 @@ 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 CI's YAML -job stays within its runner's memory. 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. +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