Stop the test shards' memory growing with every Simulation under sysmon coverage

changelog.d/2330.fixed.md

@@ -0,0 +1,3 @@
+- Stop the Python tests' memory growing with every Simulation under sys.monitoring coverage: each test process compiles each variable file once and every tax-benefit system runs that code (#2330).
+- Release dead tax-benefit systems' variable modules in the Python tests once eight more systems own modules than were alive at the last release, checked as each new system starts adding its variables as well as after each test.
+- Hand freed memory back to the system after each test file (glibc malloc_trim), and stop keeping the unpinned latest dataset cached after its smoke test.

policyengine_uk/tests/conftest.py

@@ -1,10 +1,12 @@
 import argparse
+import ctypes
 import gc
 import hashlib
 import os
 import re
 import sys
 import weakref
+from importlib.machinery import SourceFileLoader
 
 import pytest
 from policyengine_core.taxbenefitsystems import TaxBenefitSystem
@@ -152,6 +154,10 @@ def pytest_configure(config):
     _owners_at_start.update(_variable_module_owners())
     _threshold_before = gc.get_threshold()
     gc.set_threshold(GC_FIRST_GENERATION_THRESHOLD, *_threshold_before[1:])
+    # The variable code and module hooks defined below. pytest_configure, unlike
+    # pytest_sessionstart, also reaches a conftest that pytest finds only while
+    # collecting (a bare `pytest` from the repository root).
+    _install_variable_hooks()
 
 
 @pytest.hookimpl(trylast=True)
@@ -182,6 +188,7 @@ def pytest_unconfigure(config):
     _owners_at_start.clear()
     if _threshold_before is not None:
         gc.set_threshold(*_threshold_before)
+    _remove_variable_hooks()
 
 
 # policyengine-core loads every variable file of each new tax-benefit system as
@@ -230,7 +237,127 @@ def release_dead_variable_modules(max_owners: int = 8) -> int:
     return released
 
 
+# A test or a new system triggers the release only once more than ``max_new``
+# systems own modules beyond those alive at the last release, so a process
+# that keeps more than eight systems alive (an lru_cache of systems in a test
+# module) doesn't scan every tracked object after every test and every build.
+_live_owners = 0
+
+
+def release_after_new_systems(max_new: int = 8) -> int:
+    """Release dead systems' variable modules once more than ``max_new`` systems
+    own modules beyond those alive at the last release. Returns the number
+    removed."""
+    global _live_owners
+    if len(_variable_module_owners()) <= _live_owners + max_new:
+        return 0
+    released = release_dead_variable_modules(max_owners=0)
+    _live_owners = len(_variable_module_owners())
+    return released
+
+
 @pytest.fixture(autouse=True)
 def _release_variable_modules_of_dead_systems():
     yield
-    release_dead_variable_modules()
+    release_after_new_systems()
+
+
+# Coverage on CI is measured with sys.monitoring (COVERAGE_CORE=sysmon). That
+# core keeps every code object it has seen alive until the process exits, so
+# that it can key its records by id(code) (SysMonitor.code_objects in
+# coverage/sysmon.py). policyengine-core compiles a system's variable files
+# again for every new system, into new code objects, and every Simulation
+# builds a new system, so each Simulation left about 2,500 code objects that
+# were never freed. By the end of a CI shard each pytest-xdist worker's
+# coverage held 0.8 to 1.5 million of them, and the shards peaked at 14.4 to
+# 15.3 GiB of the runner's 15.6, using up to 1.8 GiB of swap (run 38159894687).
+#
+# A file's code object does not depend on the system that runs it, so each
+# process compiles each variable file once and every later system runs that
+# code object. Systems still get their own modules, classes and functions;
+# only the immutable code is shared. In the same run this lowered the shards'
+# peaks to 10.2 to 12.2 GiB with no swap, cut their pytest time by 15 to 21%,
+# and left every test's outcome unchanged. Remove _get_variable_code once
+# policyengine-core shares the code itself (PolicyEngine/policyengine-core#608).
+_variable_code = {}
+_get_code = SourceFileLoader.get_code
+_get_code_was_own = "get_code" in SourceFileLoader.__dict__
+
+
+def _get_variable_code(loader, fullname):
+    """``SourceFileLoader.get_code``, compiling each variable file only once."""
+    if not _VARIABLE_MODULE_NAME.match(fullname):
+        return _get_code(loader, fullname)
+    path = loader.get_filename(fullname)
+    stat = os.stat(path)
+    key = (path, stat.st_mtime_ns, stat.st_size)
+    code = _variable_code.get(key)
+    if code is None:
+        code = _variable_code[key] = _get_code(loader, fullname)
+    return code
+
+
+# Dead systems' variable modules are also released as each new system starts
+# adding its variables, not only after each test: in that run one property
+# test had 165 systems' modules loaded at once, and Hypothesis can build many
+# more while it shrinks a failing example. With an earlier form of this
+# release, which scanned at every new system once more than eight systems
+# owned modules, the shards peaked at 10.0-10.4 GiB (run 38165445604), against
+# 10.2-12.2 with code sharing alone (run 38159894687). Keep this until
+# policyengine-core stops leaking the modules
+# (PolicyEngine/policyengine-core#520).
+_add_variables_from_file = TaxBenefitSystem.add_variables_from_file
+_adding_system = None
+
+
+def _add_variables_from_file_releasing(system, file_path):
+    """``TaxBenefitSystem.add_variables_from_file``, releasing dead systems'
+    modules when a new system starts adding its variables."""
+    global _adding_system
+    if id(system) != _adding_system:
+        # A system adds all its variable files in turn, so this is a new one.
+        _adding_system = id(system)
+        release_after_new_systems()
+    return _add_variables_from_file(system, file_path)
+
+
+def _install_variable_hooks():
+    SourceFileLoader.get_code = _get_variable_code
+    TaxBenefitSystem.add_variables_from_file = _add_variables_from_file_releasing
+
+
+def _remove_variable_hooks():
+    if _get_code_was_own:
+        SourceFileLoader.get_code = _get_code
+    elif "get_code" in SourceFileLoader.__dict__:
+        # Inherited from SourceLoader: leave no copy on the subclass.
+        del SourceFileLoader.get_code
+    TaxBenefitSystem.add_variables_from_file = _add_variables_from_file
+
+
+# Freed memory stays in the worker: on CI the worker that ran the
+# microsimulation tests kept about 4 GiB of RSS after their simulations were
+# gone (run 38159894687), though locally only the 0.74 GiB cached dataset was
+# still referenced once they were deleted (tracemalloc). So after each test
+# file's fixtures are torn down, ask glibc to return free heap memory to the
+# system (malloc_trim(0)); memory still awaiting the cycle collector goes back
+# at a later file. Elsewhere this does nothing. With this, and
+# test_latest_data_smoke.py no longer keeping its dataset cached, the shards
+# peaked at 8.0 to 8.2 GiB instead of 10.0 to 10.4, in about the same time
+# (pytest wall 776 / 952 / 637 s against 788 / 902 / 581), and two shards
+# would peak at 8.1 (run 38165445604). With everything in this file the three
+# shards peaked at 7.6 to 8.7 GiB (run 38173323996).
+try:
+    _malloc_trim = ctypes.CDLL("libc.so.6").malloc_trim
+    _malloc_trim.argtypes = [ctypes.c_size_t]
+    _malloc_trim.restype = ctypes.c_int
+except (OSError, AttributeError):
+    _malloc_trim = None
+
+
+@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.
+    if _malloc_trim is not None and (nextitem is None or nextitem.path != item.path):
+        _malloc_trim(0)

policyengine_uk/tests/test_latest_data_smoke.py

@@ -16,10 +16,14 @@
 
 from __future__ import annotations
 
+import os
+from collections.abc import Iterator
+
 import numpy as np
 import pytest
 
 from policyengine_uk import Microsimulation
+from policyengine_uk.simulation import _url_dataset_cache
 
 
 LATEST_DATASET_URL = (
@@ -29,7 +33,7 @@
 
 
 @pytest.fixture(scope="module")
-def sim() -> Microsimulation:
+def sim() -> Iterator[Microsimulation]:
     """Simulation built against the unpinned latest dataset.
 
     Overrides any pinned-version dataset set in conftest.py so the test
@@ -39,7 +43,13 @@ def sim() -> Microsimulation:
     """
     with pytest.MonkeyPatch.context() as monkeypatch:
         monkeypatch.setenv("POLICYENGINE_UK_DEFAULT_DATASET", LATEST_DATASET_URL)
-        return Microsimulation()
+        simulation = Microsimulation()
+    yield simulation
+    # Only this file uses the unpinned dataset, so don't keep it cached beside
+    # the pinned one for the rest of the test process (a cached copy of the
+    # pinned 2023-24 dataset held 0.74 GiB, measured locally).
+    if os.environ.get("POLICYENGINE_UK_DEFAULT_DATASET") != LATEST_DATASET_URL:
+        _url_dataset_cache.pop(LATEST_DATASET_URL, None)
 
 
 def _weighted(sim: Microsimulation, variable: str, period: int = YEAR) -> float:

policyengine_uk/tests/test_malloc_trim.py

@@ -0,0 +1,43 @@
+"""The conftest hands freed memory back to the system after each test file.
+
+It calls glibc's malloc_trim(0) once a file's last test has been torn down,
+and only then. With glibc it must find the function, or the release would
+silently do nothing.
+"""
+
+import platform
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+
+
+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")
+
+
+def test_malloc_trim_is_found_with_glibc(request):
+    if platform.libc_ver()[0] != "glibc":
+        pytest.skip("malloc_trim is part of glibc")
+    assert _conftest(request)._malloc_trim is not None
+
+
+@pytest.mark.parametrize(
+    "next_file, trims",
+    [("test_a.py", False), ("test_b.py", True), (None, True)],
+    ids=["same file", "another file", "last test"],
+)
+def test_trims_only_after_a_files_last_test(request, monkeypatch, next_file, trims):
+    conftest = _conftest(request)
+    calls = []
+    monkeypatch.setattr(conftest, "_malloc_trim", calls.append)
+    item = SimpleNamespace(path=Path("test_a.py"))
+    nextitem = None if next_file is None else SimpleNamespace(path=Path(next_file))
+
+    conftest.pytest_runtest_teardown(item, nextitem)
+
+    assert calls == ([0] if trims else [])

policyengine_uk/tests/test_variable_code_sharing.py

@@ -0,0 +1,138 @@
+"""The conftest compiles each variable file once and shares the code.
+
+policyengine-core compiles a system's variable files again for every new
+system. Coverage's sys.monitoring core, which CI uses, keeps every code object
+it has seen alive, so each Simulation used to leave its code objects behind.
+The conftest makes every system run the code object compiled the first time.
+
+Systems must still get their own modules, classes and functions, run the code
+of the file's current source, and compute the same results. Once
+policyengine-core shares the code itself (PolicyEngine/policyengine-core#608),
+the tests of the workaround skip with a note to remove it.
+"""
+
+import sys
+from importlib.machinery import SourceFileLoader
+from pathlib import Path
+
+import pytest
+
+from policyengine_uk import Simulation
+
+SITUATION = {
+    "people": {"adult": {"age": {2026: 40}, "employment_income": {2026: 30_000}}},
+    "benunits": {"benunit": {"members": ["adult"]}},
+    "households": {"household": {"members": ["adult"]}},
+}
+
+
+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")
+
+
+@pytest.fixture(scope="module")
+def core_shares_code(request):
+    """Whether policyengine-core itself gives two systems the same code."""
+    from policyengine_uk.tax_benefit_system import CountryTaxBenefitSystem
+
+    with pytest.MonkeyPatch.context() as monkeypatch:
+        monkeypatch.setattr(SourceFileLoader, "get_code", _conftest(request)._get_code)
+        first, second = CountryTaxBenefitSystem(), CountryTaxBenefitSystem()
+    first = first.variables["income_tax"].get_formula().__code__
+    second = second.variables["income_tax"].get_formula().__code__
+    return first is second
+
+
+def _skip_if_core_shares(core_shares_code):
+    if core_shares_code:
+        pytest.skip(
+            "policyengine-core now shares variable code between systems; "
+            "remove _get_variable_code from conftest.py and the tests of it"
+        )
+
+
+def _income_tax_formula(simulation):
+    variable = simulation.tax_benefit_system.variables["income_tax"]
+    return variable.get_formula()
+
+
+def _coverage_code_objects():
+    """Code objects coverage's sys.monitoring core holds, or None if inactive."""
+    try:
+        import coverage
+    except ImportError:
+        return None
+    cov = coverage.Coverage.current()
+    collector = getattr(cov, "_collector", None)
+    if collector is None:
+        return None
+    tracers = getattr(collector, "tracers", ())
+    lists = [getattr(tracer, "code_objects", None) for tracer in tracers]
+    lists = [objects for objects in lists if objects is not None]
+    return sum(len(objects) for objects in lists) if lists else None
+
+
+def test_systems_share_code_but_not_modules_classes_or_functions():
+    first = Simulation(situation=SITUATION)
+    second = Simulation(situation=SITUATION)
+    first_class = type(first.tax_benefit_system.variables["income_tax"])
+    second_class = type(second.tax_benefit_system.variables["income_tax"])
+    first_formula = _income_tax_formula(first)
+    second_formula = _income_tax_formula(second)
+
+    assert first_class is not second_class
+    assert (
+        sys.modules[first_class.__module__]
+        is not (sys.modules[second_class.__module__])
+    )
+    assert first_formula is not second_formula
+    assert first_formula.__globals__ is not second_formula.__globals__
+    assert first_formula.__code__ is second_formula.__code__
+    assert first.calculate("income_tax", 2026) == second.calculate("income_tax", 2026)
+
+
+def test_shared_code_is_compiled_from_the_current_source(request, core_shares_code):
+    _skip_if_core_shares(core_shares_code)
+    Simulation(situation=SITUATION)
+    shared = _conftest(request)._variable_code
+    variables_dir = Path(sys.modules["policyengine_uk"].__file__).parent / "variables"
+    # Core reads a directory's __init__.py for metadata and runs every other file.
+    files = {
+        str(path) for path in variables_dir.rglob("*.py") if path.name != "__init__.py"
+    }
+
+    assert shared, "no variable code was shared"
+    for (path, mtime_ns, size), code in shared.items():
+        stat = Path(path).stat()
+        assert (stat.st_mtime_ns, stat.st_size) == (mtime_ns, size), path
+        source = Path(path).read_text(encoding="utf-8")
+        assert code == compile(source, path, "exec", dont_inherit=True), path
+    assert {path for path, _, _ in shared} >= files
+
+
+def test_policyengine_core_still_compiles_each_system_again(core_shares_code):
+    _skip_if_core_shares(core_shares_code)
+
+
+def test_a_simulation_leaves_coverage_no_new_code_objects(
+    request, monkeypatch, core_shares_code
+):
+    _skip_if_core_shares(core_shares_code)
+    if _coverage_code_objects() is None:
+        pytest.skip("coverage is not measuring with sys.monitoring")
+    Simulation(situation=SITUATION).calculate("income_tax", 2026)
+
+    before = _coverage_code_objects()
+    Simulation(situation=SITUATION).calculate("income_tax", 2026)
+    shared = _coverage_code_objects() - before
+    monkeypatch.setattr(SourceFileLoader, "get_code", _conftest(request)._get_code)
+    Simulation(situation=SITUATION).calculate("income_tax", 2026)
+    unshared = _coverage_code_objects() - before - shared
+
+    # Without sharing, a Simulation leaves about 2,500 code objects.
+    assert unshared > 1_000, unshared
+    assert shared < 100, shared