diff --git a/changelog.d/2357.fixed.md b/changelog.d/2357.fixed.md new file mode 100644 index 0000000000..cf545ad608 --- /dev/null +++ b/changelog.d/2357.fixed.md @@ -0,0 +1 @@ +Release the tax-benefit systems and simulations that five test modules cache once each module's tests finish, build the two-child limit and cliff property tests' reforms over one shared baseline system, and check that every test module cache is released (#2357). diff --git a/policyengine_uk/tests/code_health/test_test_module_caches_released.py b/policyengine_uk/tests/code_health/test_test_module_caches_released.py new file mode 100644 index 0000000000..d3059879f5 --- /dev/null +++ b/policyengine_uk/tests/code_health/test_test_module_caches_released.py @@ -0,0 +1,384 @@ +"""Test modules must release their module-level caches when they finish. + +Property tests cache tax-benefit systems and simulations with +``functools.lru_cache`` so that each Hypothesis example does not build them +again. A module-level cache outlives the module's tests: it keeps everything +it holds for the rest of the test process, and a system also keeps the last +simulation built from it (``Simulation.__init__`` sets +``tax_benefit_system.simulation``). One such cache kept 13 systems, about +0.5 GiB, in a CI worker for the rest of its run (run 38165445604, #2357). + +So a test module must call ``cache_clear()``, after the ``yield`` of a +module-scoped autouse fixture, on every functools cache that it: + +- defines at module level (a decorated function, or + ``name = lru_cache(...)(function)``); +- imports from another test file; or +- fills through a function it imports from another test file, where that + function reaches the cache through that file's own module-level functions. + +An imported cache counts separately because, under pytest's default (prepend) +import mode, pytest collects the files in this directory by bare name (it has +no ``__init__.py``), so ``from policyengine_uk.tests.X import f`` loads a second +copy of ``X`` with its own caches, which ``X``'s fixtures never clear. + +This is a static scan. It does not see cached methods, caches reached through +``import module`` rather than ``from module import name``, a cache re-exported +through a third test file or reached through a helper in one, or whether a +``cache_clear()`` written after the ``yield`` runs (one in an ``except`` branch +counts). +""" + +import ast +from collections.abc import Iterator +from pathlib import Path + +TESTS = Path(__file__).resolve().parents[1] +CACHES = {"lru_cache", "cache"} +FUNCTIONS = (ast.FunctionDef, ast.AsyncFunctionDef) + + +def _is_functools_cache(node: ast.AST) -> bool: + """``lru_cache``, ``cache``, ``functools.lru_cache``, called or not.""" + if isinstance(node, ast.Call): + node = node.func + if isinstance(node, ast.Attribute): + return ( + node.attr in CACHES + and isinstance(node.value, ast.Name) + and node.value.id == "functools" + ) + return isinstance(node, ast.Name) and node.id in CACHES + + +def module_functions(tree: ast.Module) -> dict[str, ast.AST]: + return {node.name: node for node in tree.body if isinstance(node, FUNCTIONS)} + + +def cache_bodies(tree: ast.Module) -> dict[str, ast.AST | None]: + """Module-level names bound to a functools cache, each with the module + function it wraps (None when it wraps something else).""" + functions = module_functions(tree) + caches: dict[str, ast.AST | None] = { + name: function + for name, function in functions.items() + if any(_is_functools_cache(d) for d in function.decorator_list) + } + for node in tree.body: + if not (isinstance(node, ast.Assign) and isinstance(node.value, ast.Call)): + continue + call = node.value + # name = lru_cache(maxsize=1)(function) or name = cache(function), but + # not name = lru_cache(maxsize=1), which is a decorator, not a cache. + if isinstance(call.func, ast.Call): + applied = _is_functools_cache(call.func) + else: + applied = ( + _is_functools_cache(call) + and bool(call.args) + and not isinstance(call.args[0], ast.Constant) + ) + if applied: + wrapped = call.args[0] if call.args else None + body = functions.get(wrapped.id) if isinstance(wrapped, ast.Name) else None + caches.update({t.id: body for t in node.targets if isinstance(t, ast.Name)}) + return caches + + +def cached_functions(tree: ast.Module) -> set[str]: + return set(cache_bodies(tree)) + + +def reached_caches(tree: ast.Module, names: set[str]) -> set[str]: + """Caches of this module that calling the named functions can fill: the + names themselves, and any cache their bodies reach through the module's + own functions.""" + caches = cache_bodies(tree) + bodies = module_functions(tree) | { + name: body for name, body in caches.items() if body is not None + } + seen: set[str] = set() + pending = [name for name in names if name in bodies or name in caches] + while pending: + name = pending.pop() + if name in seen: + continue + seen.add(name) + if name in bodies: + pending += [ + node.id + for node in ast.walk(bodies[name]) + if isinstance(node, ast.Name) + and (node.id in bodies or node.id in caches) + ] + return seen & set(caches) + + +def _own_nodes(function: ast.AST) -> Iterator[ast.AST]: + """The function's nodes, without the bodies of functions nested in it.""" + pending = list(ast.iter_child_nodes(function)) + while pending: + node = pending.pop() + yield node + if not isinstance(node, (*FUNCTIONS, ast.Lambda, ast.ClassDef)): + pending.extend(ast.iter_child_nodes(node)) + + +def _is_module_autouse_fixture(function: ast.AST) -> bool: + for decorator in getattr(function, "decorator_list", []): + if not isinstance(decorator, ast.Call): + continue + func = decorator.func + name = func.attr if isinstance(func, ast.Attribute) else getattr(func, "id", "") + keywords = { + keyword.arg: keyword.value.value + for keyword in decorator.keywords + if isinstance(keyword.value, ast.Constant) + } + if ( + name == "fixture" + and keywords.get("scope") == "module" + and keywords.get("autouse") is True + ): + return True + return False + + +def released_caches(tree: ast.Module) -> set[str]: + """Names whose cache_clear() comes after the yield of a module-scoped + autouse fixture, including a yield inside ``with`` or ``try``.""" + released = set() + for function in module_functions(tree).values(): + if not _is_module_autouse_fixture(function): + continue + nodes = list(_own_nodes(function)) + yields = [node.lineno for node in nodes if isinstance(node, ast.Yield)] + if not yields: + continue + for node in nodes: + if ( + isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and node.func.attr == "cache_clear" + and isinstance(node.func.value, ast.Name) + and node.lineno > min(yields) + ): + released.add(node.func.value.id) + return released + + +def imported_test_files( + tree: ast.Module, path: Path +) -> dict[Path, dict[str, set[str]]]: + """Other test files this module imports names from: for each file, the + imported names and everything this module calls each of them.""" + imports: dict[Path, dict[str, set[str]]] = {} + for node in tree.body: + if not isinstance(node, ast.ImportFrom): + continue + module = (node.module or "").removeprefix("policyengine_uk.tests.") + if not module: + continue + relative = Path(*module.split(".")).with_suffix(".py") + bases = [path.parents[node.level - 1]] if node.level else [TESTS, path.parent] + for base in bases: + if (base / relative).is_file(): + names = imports.setdefault(base / relative, {}) + for alias in node.names: + names.setdefault(alias.name, set()).add(alias.asname or alias.name) + break + return imports + + +def parse(path: Path) -> ast.Module: + return ast.parse(path.read_text(), filename=str(path)) + + +def unreleased_caches(path: Path) -> list[str]: + tree = parse(path) + # Each cache, as the names this module can clear it by. + caches = [{name} for name in cached_functions(tree)] + for imported, names in imported_test_files(tree, path).items(): + if imported == path: + continue + # A cache this module fills without importing it is reported under + # its own name: import it to clear it. + caches += [ + names.get(cache, {cache}) + for cache in reached_caches(parse(imported), set(names)) + ] + released = released_caches(tree) + return sorted({min(names) for names in caches if not names & released}) + + +def test_every_test_module_cache_is_released(): + findings = { + str(path.relative_to(TESTS)): missing + for path in sorted(TESTS.rglob("*.py")) + if (missing := unreleased_caches(path)) + } + assert not findings, ( + "Clear these caches after the yield of a module-scoped autouse " + f"fixture in the module that uses them: {findings}" + ) + + +SCANNED = """ +import functools +from functools import cache, lru_cache + +import pytest + + +@lru_cache(maxsize=None) +def systems(value): + return value + + +@cache +def simulation(): + return 1 + + +@functools.lru_cache +async def template(): + return 2 + + +def build(): + return 3 + + +assigned = lru_cache(maxsize=1)(build) +wrapped = cache(build) +not_a_functools_cache = registry.cache(build) + + +def helper(): + return indirect() + + +def indirect(): + return simulation() + + +def pure(): + return build() + + +def fills_simulation(): + return simulation() + + +through_a_wrapper = cache(fills_simulation) +a_decorator = lru_cache(maxsize=1) + + +@pytest.fixture(scope="module", autouse=True) +def release(): + systems.cache_clear() # before the yield: does not count + yield + + def never_called(): + wrapped.cache_clear() # in a nested function: does not count + + simulation.cache_clear() + + +@pytest.fixture(autouse=True, scope="module") +def release_after_a_yield_inside_with(): + with open(__file__): + try: + yield + finally: + template.cache_clear() +""" + + +def test_the_scan_finds_caches_and_releases(): + tree = ast.parse(SCANNED) + assert cached_functions(tree) == { + "systems", + "simulation", + "template", + "assigned", + "wrapped", + "through_a_wrapper", + } + assert released_caches(tree) == {"simulation", "template"} + + function_scoped = SCANNED.replace('scope="module"', 'scope="function"') + assert released_caches(ast.parse(function_scoped)) == set() + not_autouse = SCANNED.replace("autouse=True", "autouse=False") + assert released_caches(ast.parse(not_autouse)) == set() + + +def test_the_scan_follows_functions_to_the_caches_they_fill(): + tree = ast.parse(SCANNED) + assert reached_caches(tree, {"helper"}) == {"simulation"} + assert reached_caches(tree, {"systems", "pure"}) == {"systems"} + # An assigned cache also fills what the function it wraps reaches. + assert reached_caches(tree, {"through_a_wrapper"}) == { + "through_a_wrapper", + "simulation", + } + assert reached_caches(tree, {"pure", "not_defined_here"}) == set() + + +def test_the_scan_follows_imports_from_other_test_files(tmp_path): + (tmp_path / "test_source.py").write_text(SCANNED) + user = tmp_path / "test_user.py" + for imports, expected in [ + # An imported cache, under the name this module gives it. + ("from test_source import systems as cached_systems", ["cached_systems"]), + # A cache filled through an imported helper, which is not imported. + ("from test_source import helper", ["simulation"]), + ("from .test_source import helper, simulation", ["simulation"]), + ("from test_source import pure", []), + ]: + user.write_text(imports + "\n") + assert unreleased_caches(user) == expected, imports + + # A relative import from a subdirectory, which only the import's level + # resolves. + nested = tmp_path / "nested" / "test_user.py" + nested.parent.mkdir() + nested.write_text("from ..test_source import helper\n") + assert unreleased_caches(nested) == ["simulation"] + + # A cache is released under any of the names this module imports it by. + user.write_text( + "import pytest\n" + "from test_source import systems as a\n" + "from test_source import systems as b\n\n\n" + '@pytest.fixture(scope="module", autouse=True)\n' + "def release():\n" + " yield\n" + " a.cache_clear()\n" + ) + assert unreleased_caches(user) == [] + + user.write_text( + "import pytest\n" + "from test_source import helper, simulation as shared\n\n\n" + '@pytest.fixture(scope="module", autouse=True)\n' + "def release():\n" + " yield\n" + " shared.cache_clear()\n" + ) + assert unreleased_caches(user) == [] + + +def test_the_scan_resolves_this_directorys_imports(): + tree = ast.parse( + "from policyengine_uk.tests.test_state_pension_age import (\n" + " YEARS,\n" + " differential_simulation as cached,\n" + ")\n" + ) + assert imported_test_files(tree, TESTS / "test_x.py") == { + TESTS / "test_state_pension_age.py": { + "YEARS": {"YEARS"}, + "differential_simulation": {"cached"}, + } + } diff --git a/policyengine_uk/tests/test_cliff_evaluated_properties.py b/policyengine_uk/tests/test_cliff_evaluated_properties.py index 0e2144d3e8..14c46b230c 100644 --- a/policyengine_uk/tests/test_cliff_evaluated_properties.py +++ b/policyengine_uk/tests/test_cliff_evaluated_properties.py @@ -16,6 +16,7 @@ from functools import lru_cache import numpy as np +import pytest from hypothesis import HealthCheck, given, settings from hypothesis import strategies as st from policyengine_core.reforms import Reform, set_parameter @@ -32,7 +33,14 @@ ) -@lru_cache(maxsize=None) +@lru_cache(maxsize=1) +def baseline_system(): + return CountryTaxBenefitSystem() + + +# One reform per value scenarios() draws for mtr_adults. A reform clones only +# the baseline's parameters. +@lru_cache(maxsize=4) def tax_benefit_system(mtr_adults): modifier = set_parameter( MTR_ADULTS, mtr_adults, return_modifier=True, period="year:2000:100" @@ -42,7 +50,16 @@ class override(Reform): def apply(self): self.parameters = modifier(self.parameters) - return override(CountryTaxBenefitSystem()) + return override(baseline_system()) + + +@pytest.fixture(scope="module", autouse=True) +def release_systems(): + # Each system also keeps its last simulation alive. Drop them once this + # module's tests are done, instead of for the rest of the test process. + yield + tax_benefit_system.cache_clear() + baseline_system.cache_clear() earnings = st.sampled_from([0, 5_000, 20_000, 60_000, 120_000]) diff --git a/policyengine_uk/tests/test_state_pension_age.py b/policyengine_uk/tests/test_state_pension_age.py index a8cb7d948c..459c5af90e 100644 --- a/policyengine_uk/tests/test_state_pension_age.py +++ b/policyengine_uk/tests/test_state_pension_age.py @@ -256,6 +256,16 @@ def differential_simulation() -> tuple: return Simulation(situation=situation(people)), keys +@pytest.fixture(scope="module", autouse=True) +def release_cached_simulations(): + # The differential simulation keeps every year it has calculated. Drop + # the cached simulations once this module's tests are done, instead of + # for the rest of the test process. + yield + differential_simulation.cache_clear() + situation_template.cache_clear() + + @pytest.mark.parametrize("year", YEARS) def test_status_matches_statute_for_every_birth_date(year): """State Pension age status, type, the Savings Credit age test and the day diff --git a/policyengine_uk/tests/test_state_pension_credit_qualifying_age.py b/policyengine_uk/tests/test_state_pension_credit_qualifying_age.py index a1eea4e249..3dc872d96a 100644 --- a/policyengine_uk/tests/test_state_pension_credit_qualifying_age.py +++ b/policyengine_uk/tests/test_state_pension_credit_qualifying_age.py @@ -81,6 +81,19 @@ def property_simulation(people: dict): return sim +@pytest.fixture(scope="module", autouse=True) +def release_cached_simulations(): + # Under pytest's default import mode, the import above loads + # policyengine_uk.tests.test_state_pension_age, a separate module from the + # test_state_pension_age that pytest collects by bare name (this directory + # has no __init__.py). Its differential_simulation cache is filled here, so + # it is cleared here. Drop the cached simulations once this module's tests + # are done, instead of for the rest of the test process. + yield + differential_simulation.cache_clear() + one_person_units_template.cache_clear() + + @pytest.mark.parametrize("year", YEARS) def test_qualifying_age_matches_statute_for_every_birth_date(year): """For births every third day from 1945 to 1984 and on both sides of every diff --git a/policyengine_uk/tests/test_two_child_limit_age_exemption_properties.py b/policyengine_uk/tests/test_two_child_limit_age_exemption_properties.py index 315dc29705..cc636ab030 100644 --- a/policyengine_uk/tests/test_two_child_limit_age_exemption_properties.py +++ b/policyengine_uk/tests/test_two_child_limit_age_exemption_properties.py @@ -17,6 +17,7 @@ from functools import lru_cache import numpy as np +import pytest from hypothesis import HealthCheck, given, settings from hypothesis import strategies as st from policyengine_core.reforms import Reform, set_parameter @@ -38,12 +39,17 @@ TOLERANCE = 0.01 -@lru_cache(maxsize=None) -def tax_benefit_system(parameter, value): - """The UK system, with one parameter overridden for all years if given.""" - system = CountryTaxBenefitSystem() - if parameter is None: - return system +@lru_cache(maxsize=1) +def baseline_system(): + return CountryTaxBenefitSystem() + + +# Each example uses the repealed limit, then one threshold, so the last two +# reforms are the repealed limit's and the latest threshold's: a reform is +# built only when the threshold changes, instead of one system being held per +# threshold drawn. A reform clones only the baseline's parameters. +@lru_cache(maxsize=2) +def reformed_system(parameter, value): modifier = set_parameter( parameter, value, return_modifier=True, period="year:2000:100" ) @@ -52,7 +58,23 @@ class override(Reform): def apply(self): self.parameters = modifier(self.parameters) - return override(system) + return override(baseline_system()) + + +def tax_benefit_system(parameter, value): + """The UK system, with one parameter overridden for all years if given.""" + if parameter is None: + return baseline_system() + return reformed_system(parameter, value) + + +@pytest.fixture(scope="module", autouse=True) +def release_systems(): + # Each system also keeps its last simulation alive. Drop them once this + # module's tests are done, instead of for the rest of the test process. + yield + reformed_system.cache_clear() + baseline_system.cache_clear() @st.composite diff --git a/policyengine_uk/tests/test_uc_work_related_groups_properties.py b/policyengine_uk/tests/test_uc_work_related_groups_properties.py index 2b3c643ae7..573e820fe5 100644 --- a/policyengine_uk/tests/test_uc_work_related_groups_properties.py +++ b/policyengine_uk/tests/test_uc_work_related_groups_properties.py @@ -499,6 +499,14 @@ def qualifying_age_couples(): return Simulation(situation=situation), keys +@pytest.fixture(scope="module", autouse=True) +def release_cached_simulation(): + # Drop the cached simulation, and the system it built, once this module's + # tests are done, instead of for the rest of the test process. + yield + qualifying_age_couples.cache_clear() + + @pytest.mark.parametrize("year", QUALIFYING_AGE_YEARS) def test_regulation_89_1_a_follows_the_statute_for_every_birth_date(year): """The elder member is in section 19 exactly when the statute puts a woman