From c5075d31c8597379c054bef2e1c181c09463834b Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 16:09:18 -0400 Subject: [PATCH 1/3] Trace parameter reads per call and stop reruns leaking module fixtures Simulation._run_formula no longer sets trace, tracer and branch_name on the shared parameter tree when a simulation traces. Its formulas get a per-call TracingParameterNode instead, so the tax-benefit system, its clones and every other simulation on it stay untraced, and each traced simulation and branch records its own parameter reads under its own branch name. TracingParameterNodeAtInstant (and the new wrapper) answer special names and lookups on unfilled instances with AttributeError, so copy, deepcopy and pickle no longer recurse (Reform.modify_parameters deep-copies the tree). The dev extra and the smoke job require pytest-rerunfailures>=16.2,<17: 14.0 empties pytest's setup stack before a rerun without running its finalizers, so a module-scoped fixture of the rerun module was handed to every later module (pytest-dev/pytest-rerunfailures#278 restores them from 15.0). Co-Authored-By: Claude Opus 5.5 --- .github/workflows/pr.yaml | 2 +- .../rerunfailures-fixture-teardown.changed.md | 1 + ...ed-simulation-parameter-isolation.fixed.md | 1 + docs/python_api/tracers.md | 9 + policyengine_core/simulations/simulation.py | 18 +- policyengine_core/tracers/__init__.py | 5 +- .../tracing_parameter_node_at_instant.py | 62 +++- pyproject.toml | 2 +- tests/core/test_rerun_fixture_isolation.py | 134 +++++++ .../core/test_tracing_parameter_isolation.py | 331 ++++++++++++++++++ ..._tracing_parameter_isolation_properties.py | 107 ++++++ tests/fixtures/tracing.py | 39 +++ uv.lock | 10 +- 13 files changed, 704 insertions(+), 17 deletions(-) create mode 100644 changelog.d/rerunfailures-fixture-teardown.changed.md create mode 100644 changelog.d/traced-simulation-parameter-isolation.fixed.md create mode 100644 tests/core/test_rerun_fixture_isolation.py create mode 100644 tests/core/test_tracing_parameter_isolation.py create mode 100644 tests/core/test_tracing_parameter_isolation_properties.py create mode 100644 tests/fixtures/tracing.py diff --git a/.github/workflows/pr.yaml b/.github/workflows/pr.yaml index 9492221b3..a52d09218 100644 --- a/.github/workflows/pr.yaml +++ b/.github/workflows/pr.yaml @@ -77,7 +77,7 @@ jobs: - name: Install Core and the legacy country smoke bundle # Resolve the checkout and all country dependencies together so a later # install cannot replace Core or pair a legacy US model with SPM 1.x. - run: uv pip install --system . -r .github/requirements/legacy-country-smoke.txt "pytest-rerunfailures>=10,<15" + run: uv pip install --system . -r .github/requirements/legacy-country-smoke.txt "pytest-rerunfailures>=16.2,<17" - name: Verify installed dependency compatibility run: uv pip check --system - name: Verify pytest plugins diff --git a/changelog.d/rerunfailures-fixture-teardown.changed.md b/changelog.d/rerunfailures-fixture-teardown.changed.md new file mode 100644 index 000000000..eb005f2fe --- /dev/null +++ b/changelog.d/rerunfailures-fixture-teardown.changed.md @@ -0,0 +1 @@ +The dev dependency on pytest-rerunfailures now requires 16.2 or later, whose reruns no longer leave a module-scoped fixture cached for later test modules. diff --git a/changelog.d/traced-simulation-parameter-isolation.fixed.md b/changelog.d/traced-simulation-parameter-isolation.fixed.md new file mode 100644 index 000000000..7a876b0e6 --- /dev/null +++ b/changelog.d/traced-simulation-parameter-isolation.fixed.md @@ -0,0 +1 @@ +Tracing a simulation no longer switches tracing on in the parameter tree of its tax-benefit system, so other simulations, branches and clones of the system stay untraced, every traced simulation and branch records the parameters its formulas read under its own tracer and branch name, and traced parameter nodes can be copied, deep-copied and pickled. diff --git a/docs/python_api/tracers.md b/docs/python_api/tracers.md index 5ab9b2e05..6c72e330a 100644 --- a/docs/python_api/tracers.md +++ b/docs/python_api/tracers.md @@ -56,6 +56,15 @@ The `policyengine_core.tracers` module contains classes used to represent tracer :show-inheritance: ``` +## TracingParameterNode + +```{eval-rst} +.. autoclass:: policyengine_core.tracers.tracing_parameter_node_at_instant.TracingParameterNode + :members: + :inherited-members: + :show-inheritance: +``` + ## TracingParameterNodeAtInstant ```{eval-rst} diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index 2652d9723..0da824745 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -22,7 +22,7 @@ from policyengine_core.tracers import ( FullTracer, SimpleTracer, - TracingParameterNodeAtInstant, + TracingParameterNode, ) import random from policyengine_core.tools.hugging_face import * @@ -1184,14 +1184,16 @@ def _run_formula( ) return values - if self.trace and not isinstance( - self.tax_benefit_system.parameters, TracingParameterNodeAtInstant - ): - # Soft-recast - self.tax_benefit_system.parameters.branch_name = self.branch_name - self.tax_benefit_system.parameters.trace = True - self.tax_benefit_system.parameters.tracer = self.tracer parameters_at = self.tax_benefit_system.parameters + if self.trace: + # Trace through a view of the parameter tree that belongs to this + # call. Switching tracing on in the tree itself would leave the + # shared tax-benefit system traced for every simulation, branch + # and clone that uses it afterwards, with this simulation's tracer + # and branch name cached in the tree's nodes at each instant. + parameters_at = TracingParameterNode( + parameters_at, self.tracer, self.branch_name + ) # A rules-engine formula must be a pure, deterministic function of its # inputs. Randomness is forbidden statically at variable registration diff --git a/policyengine_core/tracers/__init__.py b/policyengine_core/tracers/__init__.py index e6b07cbd5..c3df2b66a 100644 --- a/policyengine_core/tracers/__init__.py +++ b/policyengine_core/tracers/__init__.py @@ -5,4 +5,7 @@ from .variable_graph import VariableGraph from .simple_tracer import SimpleTracer from .trace_node import TraceNode -from .tracing_parameter_node_at_instant import TracingParameterNodeAtInstant +from .tracing_parameter_node_at_instant import ( + TracingParameterNode, + TracingParameterNodeAtInstant, +) diff --git a/policyengine_core/tracers/tracing_parameter_node_at_instant.py b/policyengine_core/tracers/tracing_parameter_node_at_instant.py index d4bf17466..117b5dfdd 100644 --- a/policyengine_core/tracers/tracing_parameter_node_at_instant.py +++ b/policyengine_core/tracers/tracing_parameter_node_at_instant.py @@ -21,6 +21,66 @@ Child = Union[ParameterNode, ArrayLike] +def _wrapped(wrapper: object, attribute: str, key: str) -> object: + """Return the object ``wrapper`` delegates the lookup of ``key`` to. + + ``__getattr__`` only runs when normal lookup fails. Two such lookups must + not reach the wrapped object: + + - special names, which ``copy``, ``deepcopy`` and ``pickle`` probe on an + instance (``__deepcopy__``, ``__setstate__``, ...). The wrapped + object's answer would act on the wrapped object, not on the wrapper; + - any name on an instance those protocols have created with ``__new__`` + and not filled in yet. It has no ``attribute``, so reading it here + would call ``__getattr__`` again, without end. + """ + if key.startswith("__") and key.endswith("__"): + raise AttributeError(key) + try: + return wrapper.__dict__[attribute] + except KeyError: + raise AttributeError(key) from None + + +class TracingParameterNode: + """The parameter tree as the formulas of one traced simulation see it. + + Calling it at an instant returns a :class:`TracingParameterNodeAtInstant` + that records every parameter a formula reads in ``tracer``, under + ``branch_name``. Any other attribute is read from the wrapped node. + + The wrapped node is never modified, so tracing one simulation does not + trace the tax-benefit system it shares with other simulations, branches + and clones. + """ + + def __init__( + self, + parameter_node: parameters.ParameterNode, + tracer: tracers.FullTracer, + branch_name: str, + ) -> None: + self.parameter_node = parameter_node + self.tracer = tracer + self.branch_name = branch_name + + def __call__(self, instant) -> TracingParameterNodeAtInstant: + return self.get_at_instant(instant) + + def get_at_instant(self, instant) -> TracingParameterNodeAtInstant: + node_at_instant = self.parameter_node.get_at_instant(instant) + if isinstance(node_at_instant, TracingParameterNodeAtInstant): + # The node traces by itself (its ``trace`` flag is set): record + # in this simulation's tracer, not in the one the node holds. + node_at_instant = node_at_instant.parameter_node_at_instant + return TracingParameterNodeAtInstant( + node_at_instant, self.tracer, self.branch_name + ) + + def __getattr__(self, key: str): + return getattr(_wrapped(self, "parameter_node", key), key) + + class TracingParameterNodeAtInstant: def __init__( self, @@ -36,7 +96,7 @@ def __getattr__( self, key: str, ) -> Union[TracingParameterNodeAtInstant, Child]: - child = getattr(self.parameter_node_at_instant, key) + child = getattr(_wrapped(self, "parameter_node_at_instant", key), key) return self.get_traced_child(child, key) def __getitem__( diff --git a/pyproject.toml b/pyproject.toml index 3cea68655..b5cc106fd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -48,7 +48,7 @@ dev = [ "markupsafe==2.0.1", "mypy==0.991", "pytest-cov", - "pytest-rerunfailures>=10,<15", + "pytest-rerunfailures>=16.2,<17", "sphinx==5.0.0", "sphinx-argparse==0.4.0", "sphinx-math-dollar==1.2.1", diff --git a/tests/core/test_rerun_fixture_isolation.py b/tests/core/test_rerun_fixture_isolation.py new file mode 100644 index 000000000..5844396e2 --- /dev/null +++ b/tests/core/test_rerun_fixture_isolation.py @@ -0,0 +1,134 @@ +"""A rerun must not carry a module-scoped fixture into later modules. + +``make test`` reruns failed tests (``--reruns 2``). pytest-rerunfailures 14.0 +empties pytest's setup stack before a rerun without running the finalizers on +it, and a cached fixture registers its finalizer only when it is first +created. So the module-scoped ``tax_benefit_system`` of the rerun test's +module was never torn down, and pytest handed the same cached object to every +later module, with whatever earlier modules had done to it. On Windows CI a +later module then traced it, and ``tests/core/test_parameters.py`` and +``tests/core/test_reforms.py`` failed on a system that was not theirs. + +pytest-rerunfailures 15.0 puts the finalizers back after a rerun +(pytest-dev/pytest-rerunfailures#278), and the dev dependency now requires a +release that does. This test runs a small suite with the installed plugin in a +subprocess: a test in the first module fails and is rerun, and the module +after it must get a fresh fixture, created after the first module's was torn +down. +""" + +import os +import subprocess +import sys +import textwrap +import xml.etree.ElementTree as ElementTree + +import pytest + +pytest.importorskip("pytest_rerunfailures") + +CONFTEST = """ +import pytest + +CREATED = [] + + +class System: + def __init__(self, module): + self.module = module + self.torn_down = False + + +@pytest.fixture(scope="module") +def system(request): + system = System(request.module.__name__) + CREATED.append(system) + yield system + system.torn_down = True +""" + +PASSES_ON_RERUN = """ +ATTEMPTS = [] + + +def test_rerun(system): + ATTEMPTS.append(system) + assert len(ATTEMPTS) > 1, "the first attempt fails" +""" + +FAILS_EVERY_ATTEMPT = """ +def test_rerun(system): + assert False, "every attempt fails" +""" + +LATER_MODULE = """ +from conftest import CREATED + + +def test_gets_a_fresh_fixture(system): + assert system.module == __name__ + earlier = [created for created in CREATED if created is not system] + assert earlier + assert all(created.torn_down for created in earlier) + assert all(created.module != __name__ for created in earlier) +""" + + +def _outcomes(junit_xml): + outcomes = {} + for case in ElementTree.parse(junit_xml).iter("testcase"): + name = f"{case.get('classname')}.{case.get('name')}" + failed = case.find("failure") is not None or case.find("error") is not None + outcomes[name] = "failed" if failed else "passed" + return outcomes + + +@pytest.mark.parametrize( + "rerun_module, expected_rerun_outcome", + [(PASSES_ON_RERUN, "passed"), (FAILS_EVERY_ATTEMPT, "failed")], + ids=["passes-on-rerun", "fails-every-attempt"], +) +def test_rerun_does_not_leak_module_fixture_into_later_modules( + tmp_path, rerun_module, expected_rerun_outcome +): + (tmp_path / "pytest.ini").write_text("[pytest]\n") + (tmp_path / "conftest.py").write_text(textwrap.dedent(CONFTEST)) + (tmp_path / "test_a_rerun.py").write_text(textwrap.dedent(rerun_module)) + (tmp_path / "test_b_later.py").write_text(textwrap.dedent(LATER_MODULE)) + junit_xml = tmp_path / "junit.xml" + + environment = { + key: value + for key, value in os.environ.items() + if not key.startswith(("PYTEST_", "COV_", "COVERAGE_")) + } + environment["PYTEST_DISABLE_PLUGIN_AUTOLOAD"] = "1" + completed = subprocess.run( + [ + sys.executable, + "-m", + "pytest", + "-p", + "pytest_rerunfailures", + "-p", + "no:cacheprovider", + "--reruns", + "1", + "--reruns-delay", + "0", + f"--junitxml={junit_xml}", + "-q", + ], + cwd=tmp_path, + env=environment, + capture_output=True, + text=True, + timeout=120, + ) + output = completed.stdout + completed.stderr + + assert _outcomes(junit_xml) == { + "test_a_rerun.test_rerun": expected_rerun_outcome, + "test_b_later.test_gets_a_fresh_fixture": "passed", + }, output + assert "1 rerun" in output, output diff --git a/tests/core/test_tracing_parameter_isolation.py b/tests/core/test_tracing_parameter_isolation.py new file mode 100644 index 000000000..076689607 --- /dev/null +++ b/tests/core/test_tracing_parameter_isolation.py @@ -0,0 +1,331 @@ +"""Tracing a simulation must not change the tax-benefit system it runs on. + +``Simulation._run_formula`` used to switch tracing on in the system's +parameter tree (setting ``trace``, ``tracer`` and ``branch_name`` on the root +``ParameterNode``) the first time a traced simulation ran a formula, and never +switched it off. From then on: + +- the root node caches one node per instant. ``Simulation._calculate`` reads + ``parameters(period)`` (the abolition check) before any formula runs, so at + an instant first read before tracing came on the cached node was untraced, + and the traced simulation recorded no parameter reads there. At an instant + first read afterwards it was a ``TracingParameterNodeAtInstant`` holding the + tracer and branch name of that moment, and every later simulation on the + system, traced or not, read through it: their own tracers missed the reads + and a branch's reads were filed under another branch's name; +- ``ParameterNode.clone`` copies ``trace`` and ``tracer``, so every clone of + the system traced as well; +- copying a ``TracingParameterNodeAtInstant`` (``copy``, ``deepcopy`` or + ``pickle``, as ``Reform.modify_parameters`` does when it deep-copies the + tree) recursed without end, because ``__getattr__`` read + ``parameter_node_at_instant`` from an instance the copy protocol had not + filled in yet. + +Formulas of a traced simulation now read parameters through a +``TracingParameterNode`` that belongs to the call, and both wrappers answer +lookups they cannot delegate with ``AttributeError``. + +Invariants checked here: tracing a simulation leaves the system's parameter +tree exactly as untraced (flags, cached nodes, clones, deep copies); each +traced simulation and branch records its own parameter reads, under its own +branch name; and tracing never changes a computed value. +""" + +import copy +import pickle + +import numpy as np +import pytest + +from policyengine_core.country_template import CountryTaxBenefitSystem +from policyengine_core.country_template.entities import Person +from policyengine_core.parameters import ParameterNodeAtInstant +from policyengine_core.periods import MONTH +from policyengine_core.reforms import Reform +from policyengine_core.tracers import ( + FullTracer, + TracingParameterNode, + TracingParameterNodeAtInstant, +) +from policyengine_core.variables import Variable +from tests.fixtures.tracing import ( + assert_untraced, + build_simulation, + parameter_reads, +) + +JANUARY = "2017-01" +INSTANT = "2017-01-01" + + +# ----- The system is left as it was ----- # + + +def test_traced_calculation_leaves_the_system_untraced(isolated_tax_benefit_system): + system = isolated_tax_benefit_system + simulation = build_simulation(system, trace=True) + simulation.calculate("income_tax", JANUARY) + # A second month: its instant is first read while the simulation traces. + simulation.calculate("income_tax", "2017-02") + + reads = parameter_reads(simulation.tracer.trees) + assert reads.count(("taxes.income_tax_rate", "default")) == 2 + assert_untraced(system.parameters) + assert type(system.get_parameters_at_instant(INSTANT)) is ParameterNodeAtInstant + assert_untraced(system.clone().parameters) + + +def test_reform_after_a_traced_calculation_can_copy_the_parameters( + isolated_tax_benefit_system, +): + """``Reform.modify_parameters`` deep-copies the baseline's parameter tree.""" + system = isolated_tax_benefit_system + simulation = build_simulation(system, trace=True) + simulation.calculate("income_tax", JANUARY) + # A second month: its instant is first read while the simulation traces. + simulation.calculate("income_tax", "2017-02") + + class double_income_tax_rate(Reform): + def apply(self): + def modify(parameters): + parameters.taxes.income_tax_rate.update( + period=f"year:{INSTANT}:1", + value=2 * parameters.taxes.income_tax_rate(INSTANT), + ) + return parameters + + self.modify_parameters(modify) + + reform = double_income_tax_rate(system) + + rate = system.parameters(INSTANT).taxes.income_tax_rate + assert reform.get_parameters_at_instant(INSTANT).taxes.income_tax_rate == ( + pytest.approx(2 * rate) + ) + assert system.parameters(INSTANT).taxes.income_tax_rate == rate + + +# ----- Each simulation and branch records its own reads ----- # + + +def test_each_traced_simulation_records_its_own_parameter_reads( + isolated_tax_benefit_system, +): + system = isolated_tax_benefit_system + first = build_simulation(system, trace=True) + first.calculate("income_tax", JANUARY) + first_reads = parameter_reads(first.tracer.trees) + + second = build_simulation(system, trace=True) + second.calculate("income_tax", JANUARY) + + assert ("taxes.income_tax_rate", "default") in parameter_reads(second.tracer.trees) + assert parameter_reads(first.tracer.trees) == first_reads + + +def test_branch_records_parameter_reads_under_its_own_name( + isolated_tax_benefit_system, +): + simulation = build_simulation(isolated_tax_benefit_system, trace=True) + # The parent reads parameters at the same instant first. + simulation.calculate("income_tax", JANUARY) + + branch = simulation.get_branch("policy") + branch.calculate("basic_income", JANUARY) + + reads = parameter_reads(simulation.tracer.trees) + assert ("benefits.basic_income", "policy") in reads + assert ("benefits.basic_income", "default") not in reads + + +def test_formulas_of_an_untraced_simulation_get_untraced_parameters( + isolated_tax_benefit_system, +): + system = isolated_tax_benefit_system + seen = [] + + class parameters_seen_by_formula(Variable): + value_type = float + entity = Person + definition_period = MONTH + label = "Parameters seen by the formula" + + def formula(person, period, parameters): + seen.append(parameters(period)) + return person.filled_array(0) + + system.add_variable(parameters_seen_by_formula) + + build_simulation(system, trace=True).calculate( + "parameters_seen_by_formula", JANUARY + ) + build_simulation(system).calculate("parameters_seen_by_formula", JANUARY) + + traced, untraced = seen + assert type(traced) is TracingParameterNodeAtInstant + assert type(untraced) is ParameterNodeAtInstant + + +def test_traced_formula_can_read_parameters_by_attribute(isolated_tax_benefit_system): + system = isolated_tax_benefit_system + + class income_tax_rate_by_attribute(Variable): + value_type = float + entity = Person + definition_period = MONTH + label = "Income tax rate read by attribute" + + def formula(person, period, parameters): + return person.filled_array(parameters.taxes.income_tax_rate(period)) + + system.add_variable(income_tax_rate_by_attribute) + + traced = build_simulation(system, trace=True).calculate( + "income_tax_rate_by_attribute", JANUARY + ) + untraced = build_simulation(system).calculate( + "income_tax_rate_by_attribute", JANUARY + ) + + np.testing.assert_array_equal(traced, untraced) + assert traced[0] == system.parameters(INSTANT).taxes.income_tax_rate + + +# ----- Tracing never changes a value ----- # + + +COMPUTED_VARIABLES = sorted( + name + for name, variable in CountryTaxBenefitSystem().variables.items() + if variable.formulas +) + + +@pytest.mark.parametrize("variable_name", COMPUTED_VARIABLES) +def test_traced_and_untraced_simulations_agree(tax_benefit_system, variable_name): + """Differential check over every formula of the country template.""" + variable = tax_benefit_system.get_variable(variable_name) + period = "2017" if variable.definition_period == "year" else JANUARY + + untraced = build_simulation(tax_benefit_system).calculate(variable_name, period) + traced_simulation = build_simulation(tax_benefit_system, trace=True) + traced = traced_simulation.calculate(variable_name, period) + + np.testing.assert_array_equal(traced, untraced) + assert_untraced(tax_benefit_system.parameters) + + +# ----- The tracing wrappers ----- # + + +def test_tracing_parameter_node_reads_through_to_the_node(isolated_tax_benefit_system): + parameters = isolated_tax_benefit_system.parameters + tracer = FullTracer() + view = TracingParameterNode(parameters, tracer, "policy") + + assert view.taxes is parameters.taxes + assert view.children is parameters.children + + at_instant = view(INSTANT) + assert type(at_instant) is TracingParameterNodeAtInstant + assert at_instant.parameter_node_at_instant is parameters(INSTANT) + assert at_instant.tracer is tracer + assert at_instant.branch_name == "policy" + assert type(view.get_at_instant(INSTANT)) is TracingParameterNodeAtInstant + assert_untraced(parameters) + + +def test_tracing_parameter_node_records_in_its_own_tracer_when_the_node_traces( + isolated_tax_benefit_system, +): + parameters = isolated_tax_benefit_system.parameters + parameters.trace = True + parameters.tracer = FullTracer() + parameters.branch_name = "elsewhere" + tracer = FullTracer() + + at_instant = TracingParameterNode(parameters, tracer, "policy")(INSTANT) + + assert at_instant.tracer is tracer + assert at_instant.branch_name == "policy" + assert type(at_instant.parameter_node_at_instant) is ParameterNodeAtInstant + + +def _round_trip_by_pickle(value): + return pickle.loads(pickle.dumps(value)) + + +DUPLICATORS = pytest.mark.parametrize( + "duplicate", + [copy.copy, copy.deepcopy, _round_trip_by_pickle], + ids=["copy", "deepcopy", "pickle"], +) + + +@DUPLICATORS +def test_tracing_parameter_node_at_instant_can_be_copied( + isolated_tax_benefit_system, duplicate +): + tracer = FullTracer() + tracer.record_calculation_start("income_tax", JANUARY) + traced = TracingParameterNodeAtInstant( + isolated_tax_benefit_system.parameters(INSTANT), tracer, "default" + ) + + duplicated = duplicate(traced) + + assert type(duplicated) is TracingParameterNodeAtInstant + assert duplicated.branch_name == "default" + assert duplicated.taxes.income_tax_rate == traced.taxes.income_tax_rate + assert parameter_reads(duplicated.tracer.trees)[-1] == ( + "taxes.income_tax_rate", + "default", + ) + + +@DUPLICATORS +def test_tracing_parameter_node_can_be_copied(isolated_tax_benefit_system, duplicate): + view = TracingParameterNode( + isolated_tax_benefit_system.parameters, FullTracer(), "default" + ) + + duplicated = duplicate(view) + + assert type(duplicated) is TracingParameterNode + assert duplicated.branch_name == "default" + assert duplicated(INSTANT).taxes.income_tax_rate == ( + view(INSTANT).taxes.income_tax_rate + ) + + +def test_parameters_holding_traced_nodes_can_be_deep_copied( + isolated_tax_benefit_system, +): + """A tree whose ``trace`` flag is set caches traced nodes; copying it works.""" + parameters = isolated_tax_benefit_system.parameters + parameters.trace = True + parameters.tracer = FullTracer() + parameters(INSTANT) + assert type(parameters._at_instant_cache[INSTANT]) is ( + TracingParameterNodeAtInstant + ) + + duplicated = copy.deepcopy(parameters) + + assert duplicated(INSTANT).taxes.income_tax_rate == ( + parameters(INSTANT).taxes.income_tax_rate + ) + + +@pytest.mark.parametrize( + "wrapper_class", [TracingParameterNode, TracingParameterNodeAtInstant] +) +def test_lookups_on_an_unfilled_wrapper_raise_attribute_error(wrapper_class): + """``copy`` and ``pickle`` probe instances they have only created with ``__new__``.""" + unfilled = wrapper_class.__new__(wrapper_class) + + with pytest.raises(AttributeError): + unfilled.taxes + assert getattr(unfilled, "__setstate__", None) is None + assert not hasattr(unfilled, "parameter_node_at_instant") + assert not hasattr(unfilled, "parameter_node") diff --git a/tests/core/test_tracing_parameter_isolation_properties.py b/tests/core/test_tracing_parameter_isolation_properties.py new file mode 100644 index 000000000..b1bf4b2e9 --- /dev/null +++ b/tests/core/test_tracing_parameter_isolation_properties.py @@ -0,0 +1,107 @@ +"""What a traced simulation records must not depend on what ran before it. + +Runs random sequences of calculations on one tax-benefit system: traced and +untraced simulations, branches of traced simulations, and traced simulations +on clones of the system. After every step: + +- the parameter tree of the system, and of every clone made of it, is as + untraced as a fresh system's; +- a traced simulation records exactly the parameter reads the same + calculation records when traced on a fresh system, and a branch records only + reads such a calculation makes, all under the branch's own name; +- every value equals the same calculation on a system nothing has traced. + +``test_tracing_parameter_isolation.py`` has the example tests. +""" + +import functools + +import numpy as np +import pytest + +from policyengine_core.country_template import CountryTaxBenefitSystem +from tests.fixtures.tracing import ( + assert_untraced, + build_simulation, + parameter_reads, +) + +hypothesis = pytest.importorskip("hypothesis") +st = hypothesis.strategies + +VARIABLES = sorted( + name + for name, variable in CountryTaxBenefitSystem().variables.items() + if variable.formulas +) +# Both dated formulas of ``basic_income`` (from 2015-12 and from 2016-12). +# Earlier months fail traced or not: parenting_allowance has no parameters yet. +MONTHS = ["2015-12", "2016-06", "2016-12", "2017-01", "2017-02"] +KINDS = ["traced", "untraced", "branch", "clone"] + +_REFERENCE = CountryTaxBenefitSystem() + + +def _period(variable_name, month): + variable = _REFERENCE.get_variable(variable_name) + return month[:4] if variable.definition_period == "year" else month + + +@functools.lru_cache(maxsize=None) +def _expected(variable_name, period): + """Value and parameter reads of the calculation on fresh systems.""" + value = build_simulation(CountryTaxBenefitSystem()).calculate(variable_name, period) + traced = build_simulation(CountryTaxBenefitSystem(), trace=True) + traced.calculate(variable_name, period) + return value, parameter_reads(traced.tracer.trees) + + +@hypothesis.settings( + max_examples=40, + deadline=None, + suppress_health_check=[hypothesis.HealthCheck.too_slow], +) +@hypothesis.given( + steps=st.lists( + st.tuples( + st.sampled_from(KINDS), + st.sampled_from(VARIABLES), + st.sampled_from(MONTHS), + ), + min_size=1, + max_size=6, + ) +) +def test_traced_record_does_not_depend_on_history(steps): + system = CountryTaxBenefitSystem() + systems = [system] + + for kind, variable_name, month in steps: + period = _period(variable_name, month) + expected_value, expected_reads = _expected(variable_name, period) + + if kind == "untraced": + value = build_simulation(system).calculate(variable_name, period) + elif kind == "traced": + simulation = build_simulation(system, trace=True) + value = simulation.calculate(variable_name, period) + assert parameter_reads(simulation.tracer.trees) == expected_reads + elif kind == "branch": + simulation = build_simulation(system, trace=True) + # The parent reads parameters at the same instant first. + simulation.calculate("income_tax", month) + trees_before = len(simulation.tracer.trees) + value = simulation.get_branch("policy").calculate(variable_name, period) + reads = parameter_reads(simulation.tracer.trees[trees_before:]) + assert {branch_name for _, branch_name in reads} <= {"policy"} + assert {name for name, _ in reads} <= {name for name, _ in expected_reads} + else: + clone = system.clone() + systems.append(clone) + simulation = build_simulation(clone, trace=True) + value = simulation.calculate(variable_name, period) + assert parameter_reads(simulation.tracer.trees) == expected_reads + + np.testing.assert_array_equal(value, expected_value) + for each_system in systems: + assert_untraced(each_system.parameters) diff --git a/tests/fixtures/tracing.py b/tests/fixtures/tracing.py new file mode 100644 index 000000000..36f0282c8 --- /dev/null +++ b/tests/fixtures/tracing.py @@ -0,0 +1,39 @@ +"""Helpers for tests of traced simulations.""" + +import copy + +from policyengine_core.parameters import ParameterNodeAtInstant +from policyengine_core.simulations import SimulationBuilder +from tests.fixtures.branch_shared_arrays import SITUATION + + +def build_simulation(tax_benefit_system, trace=False): + simulation = SimulationBuilder().build_from_entities( + tax_benefit_system, copy.deepcopy(SITUATION) + ) + simulation.trace = trace + return simulation + + +def parameter_reads(trees): + """Every parameter read recorded in ``trees``, as (name, branch name).""" + reads = [] + + def walk(node): + reads.extend((read.name, read.branch_name) for read in node.parameters) + for child in node.children: + walk(child) + + for tree in trees: + walk(tree) + return reads + + +def assert_untraced(parameters, instant="2017-01-01"): + """The parameter tree is as no simulation had ever traced it.""" + assert parameters.trace is False + assert parameters.tracer is None + assert parameters.branch_name is None + assert type(parameters(instant)) is ParameterNodeAtInstant + for node_at_instant in parameters._at_instant_cache.values(): + assert type(node_at_instant) is ParameterNodeAtInstant diff --git a/uv.lock b/uv.lock index 222deb407..8a837867c 100644 --- a/uv.lock +++ b/uv.lock @@ -1452,7 +1452,7 @@ wheels = [ [[package]] name = "policyengine-core" -version = "3.32.11" +version = "3.32.12" source = { editable = "." } dependencies = [ { name = "dpath" }, @@ -1516,7 +1516,7 @@ requires-dist = [ { name = "psutil", specifier = ">=6,<7" }, { name = "pytest", specifier = ">=8,<10" }, { name = "pytest-cov", marker = "extra == 'dev'" }, - { name = "pytest-rerunfailures", marker = "extra == 'dev'", specifier = ">=10,<15" }, + { name = "pytest-rerunfailures", marker = "extra == 'dev'", specifier = ">=16.2,<17" }, { name = "pyvis", specifier = ">=0.3.2" }, { name = "requests", specifier = ">=2,<3" }, { name = "ruff", marker = "extra == 'dev'", specifier = ">=0.9.0" }, @@ -1683,15 +1683,15 @@ wheels = [ [[package]] name = "pytest-rerunfailures" -version = "14.0" +version = "16.7" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "packaging" }, { name = "pytest" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/cc/a4/6de45fe850759e94aa9a55cda807c76245af1941047294df26c851dfb4a9/pytest-rerunfailures-14.0.tar.gz", hash = "sha256:4a400bcbcd3c7a4ad151ab8afac123d90eca3abe27f98725dc4d9702887d2e92", size = 21350, upload-time = "2024-03-13T08:21:39.444Z" } +sdist = { url = "https://files.pythonhosted.org/packages/d1/b0/6b5337f9d59b26b0069ea3d5e863c31dc04b69e51bbfb1cf2ea6328fba87/pytest_rerunfailures-16.7.tar.gz", hash = "sha256:6956ddfb65ca1d07e7e3d99e2c5359d82300f4cb062e6049563bd2c106f72d5c", size = 59666, upload-time = "2026-09-17T07:08:48.871Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/dc/e7/e75bd157331aecc190f5f8950d7ea3d2cf56c3c57fb44da70e60b221133f/pytest_rerunfailures-14.0-py3-none-any.whl", hash = "sha256:4197bdd2eaeffdbf50b5ea6e7236f47ff0e44d1def8dae08e409f536d84e7b32", size = 12709, upload-time = "2024-03-13T08:21:37.199Z" }, + { url = "https://files.pythonhosted.org/packages/c1/d3/07ea35102cf2020ddaaac368d2a6f0bcc63d6523fb8918805bac3eb8b98f/pytest_rerunfailures-16.7-py3-none-any.whl", hash = "sha256:edf1886209c2b7dafe35b5bf1708d6ec40ccf6c6b357f0f02807efcec0204c99", size = 23952, upload-time = "2026-09-17T07:08:47.635Z" }, ] [[package]] From d77a007ce7da7ae59c580c1081f2274face11bcb Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sat, 3 Oct 2026 07:36:06 -0400 Subject: [PATCH 2/3] Keep traced parameter nodes iterable and array-convertible Review of #574 found two regressions in traced mode: a formula iterating a parameter node, or testing membership with `in`, got KeyError: 0 (the wrapper had __getitem__ but no __iter__), and the blanket guard on special names stopped NumPy reading the array protocol through a wrapped vectorial node. The guard now covers only the copy and pickle protocol and unfilled instances; the wrappers define __iter__ and __repr__ from the node they wrap. The property test's branch step now has the parent calculate an input, so the branch must record exactly a fresh calculation's reads under its name. Co-Authored-By: Claude Opus 5.5 --- ...ed-simulation-parameter-isolation.fixed.md | 2 +- .../tracing_parameter_node_at_instant.py | 48 +++++++++++++---- .../core/test_tracing_parameter_isolation.py | 54 +++++++++++++++++++ ..._tracing_parameter_isolation_properties.py | 14 ++--- 4 files changed, 101 insertions(+), 17 deletions(-) diff --git a/changelog.d/traced-simulation-parameter-isolation.fixed.md b/changelog.d/traced-simulation-parameter-isolation.fixed.md index 7a876b0e6..a24a597ee 100644 --- a/changelog.d/traced-simulation-parameter-isolation.fixed.md +++ b/changelog.d/traced-simulation-parameter-isolation.fixed.md @@ -1 +1 @@ -Tracing a simulation no longer switches tracing on in the parameter tree of its tax-benefit system, so other simulations, branches and clones of the system stay untraced, every traced simulation and branch records the parameters its formulas read under its own tracer and branch name, and traced parameter nodes can be copied, deep-copied and pickled. +Tracing a simulation no longer switches tracing on in the parameter tree of its tax-benefit system, so other simulations, branches and clones of the system stay untraced, every traced simulation and branch records the parameters its formulas read under its own tracer and branch name, and traced parameter nodes can be copied, deep-copied, pickled, iterated and converted to arrays like untraced ones. diff --git a/policyengine_core/tracers/tracing_parameter_node_at_instant.py b/policyengine_core/tracers/tracing_parameter_node_at_instant.py index 117b5dfdd..76eaab01c 100644 --- a/policyengine_core/tracers/tracing_parameter_node_at_instant.py +++ b/policyengine_core/tracers/tracing_parameter_node_at_instant.py @@ -1,7 +1,7 @@ from __future__ import annotations import typing -from typing import Union +from typing import Iterator, Union import numpy @@ -21,20 +21,39 @@ Child = Union[ParameterNode, ArrayLike] +# Copying and pickling look these names up on an instance. Answered by the +# wrapped object, they would copy or restore the wrapped object, not the +# wrapper. +_COPY_PROTOCOL = frozenset( + { + "__copy__", + "__deepcopy__", + "__getstate__", + "__setstate__", + "__reduce__", + "__reduce_ex__", + "__getnewargs__", + "__getnewargs_ex__", + } +) + + def _wrapped(wrapper: object, attribute: str, key: str) -> object: """Return the object ``wrapper`` delegates the lookup of ``key`` to. ``__getattr__`` only runs when normal lookup fails. Two such lookups must not reach the wrapped object: - - special names, which ``copy``, ``deepcopy`` and ``pickle`` probe on an - instance (``__deepcopy__``, ``__setstate__``, ...). The wrapped - object's answer would act on the wrapped object, not on the wrapper; - - any name on an instance those protocols have created with ``__new__`` - and not filled in yet. It has no ``attribute``, so reading it here - would call ``__getattr__`` again, without end. + - the copy and pickle protocol (``__deepcopy__``, ``__setstate__``, ...); + - any name on an instance ``copy`` or ``pickle`` has created with + ``__new__`` and not filled in yet. It has no ``attribute``, so reading + it here would call ``__getattr__`` again, without end. + + Every other name is delegated, special names included: NumPy reads its + array protocol (``__array_interface__``, ...) from a vectorial node this + way. """ - if key.startswith("__") and key.endswith("__"): + if key in _COPY_PROTOCOL: raise AttributeError(key) try: return wrapper.__dict__[attribute] @@ -51,7 +70,8 @@ class TracingParameterNode: The wrapped node is never modified, so tracing one simulation does not trace the tax-benefit system it shares with other simulations, branches - and clones. + and clones. Like :class:`TracingParameterNodeAtInstant`, it is not an + instance of the class it wraps. """ def __init__( @@ -80,6 +100,9 @@ def get_at_instant(self, instant) -> TracingParameterNodeAtInstant: def __getattr__(self, key: str): return getattr(_wrapped(self, "parameter_node", key), key) + def __repr__(self) -> str: + return repr(self.parameter_node) + class TracingParameterNodeAtInstant: def __init__( @@ -106,6 +129,13 @@ def __getitem__( child = self.parameter_node_at_instant[key] return self.get_traced_child(child, key) + def __iter__(self) -> Iterator: + # Without it, ``iter`` and ``in`` fall back to ``__getitem__(0)``. + return iter(self.parameter_node_at_instant) + + def __repr__(self) -> str: + return repr(self.parameter_node_at_instant) + def get_traced_child( self, child: Child, diff --git a/tests/core/test_tracing_parameter_isolation.py b/tests/core/test_tracing_parameter_isolation.py index 076689607..0ae068286 100644 --- a/tests/core/test_tracing_parameter_isolation.py +++ b/tests/core/test_tracing_parameter_isolation.py @@ -48,6 +48,9 @@ TracingParameterNodeAtInstant, ) from policyengine_core.variables import Variable +from tests.core.parameters_fancy_indexing.test_fancy_indexing import ( + parameters as fancy_indexing_parameters, +) from tests.fixtures.tracing import ( assert_untraced, build_simulation, @@ -191,6 +194,30 @@ def formula(person, period, parameters): assert traced[0] == system.parameters(INSTANT).taxes.income_tax_rate +def test_traced_formula_can_iterate_a_parameter_node(isolated_tax_benefit_system): + """Iteration and ``in`` fell back to ``__getitem__(0)`` on a traced node.""" + system = isolated_tax_benefit_system + + class taxes_node_shape(Variable): + value_type = float + entity = Person + definition_period = MONTH + label = "Children of the taxes node, plus 10 if it has the rate" + + def formula(person, period, parameters): + taxes = parameters(period).taxes + has_rate = "income_tax_rate" in taxes + return person.filled_array(len(list(taxes)) + 10 * has_rate) + + system.add_variable(taxes_node_shape) + + traced = build_simulation(system, trace=True).calculate("taxes_node_shape", JANUARY) + untraced = build_simulation(system).calculate("taxes_node_shape", JANUARY) + + np.testing.assert_array_equal(traced, untraced) + assert untraced[0] == len(list(system.parameters(INSTANT).taxes)) + 10 + + # ----- Tracing never changes a value ----- # @@ -235,6 +262,33 @@ def test_tracing_parameter_node_reads_through_to_the_node(isolated_tax_benefit_s assert_untraced(parameters) +def test_tracing_wrappers_show_the_node_they_wrap(isolated_tax_benefit_system): + parameters = isolated_tax_benefit_system.parameters + view = TracingParameterNode(parameters, FullTracer(), "default") + + assert repr(view) == repr(parameters) + assert repr(view(INSTANT)) == repr(parameters(INSTANT)) + assert list(view(INSTANT).taxes) == list(parameters(INSTANT).taxes) + + +def test_traced_vectorial_node_converts_to_the_same_array(): + """NumPy reads the array protocol through the wrapper, as untraced.""" + rates = fancy_indexing_parameters("2015-01-01").rate + status = np.asarray(["owner", "tenant"]) + untraced = rates.single[status] + traced = TracingParameterNodeAtInstant(rates, FullTracer(), "default").single[ + status + ] + assert type(traced) is TracingParameterNodeAtInstant + + expected = np.asarray(untraced) + converted = np.asarray(traced) + + assert converted.dtype == expected.dtype + assert converted.shape == expected.shape + np.testing.assert_array_equal(converted, expected) + + def test_tracing_parameter_node_records_in_its_own_tracer_when_the_node_traces( isolated_tax_benefit_system, ): diff --git a/tests/core/test_tracing_parameter_isolation_properties.py b/tests/core/test_tracing_parameter_isolation_properties.py index b1bf4b2e9..9290d7766 100644 --- a/tests/core/test_tracing_parameter_isolation_properties.py +++ b/tests/core/test_tracing_parameter_isolation_properties.py @@ -6,9 +6,9 @@ - the parameter tree of the system, and of every clone made of it, is as untraced as a fresh system's; -- a traced simulation records exactly the parameter reads the same - calculation records when traced on a fresh system, and a branch records only - reads such a calculation makes, all under the branch's own name; +- a traced simulation, or a branch of one, records exactly the parameter + reads the same calculation records when traced on a fresh system, a + branch's all under its own name; - every value equals the same calculation on a system nothing has traced. ``test_tracing_parameter_isolation.py`` has the example tests. @@ -88,13 +88,13 @@ def test_traced_record_does_not_depend_on_history(steps): assert parameter_reads(simulation.tracer.trees) == expected_reads elif kind == "branch": simulation = build_simulation(system, trace=True) - # The parent reads parameters at the same instant first. - simulation.calculate("income_tax", month) + # The parent reads parameters at the same instant first. An input + # has no formula, so the branch reuses nothing that reads them. + simulation.calculate("salary", month) trees_before = len(simulation.tracer.trees) value = simulation.get_branch("policy").calculate(variable_name, period) reads = parameter_reads(simulation.tracer.trees[trees_before:]) - assert {branch_name for _, branch_name in reads} <= {"policy"} - assert {name for name, _ in reads} <= {name for name, _ in expected_reads} + assert reads == [(name, "policy") for name, _ in expected_reads] else: clone = system.clone() systems.append(clone) From 43c2b5bf9313d354cf3594578f44bab97d7669d7 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sat, 3 Oct 2026 09:15:58 -0400 Subject: [PATCH 3/3] Keep __slots__ off the tracing wrappers' delegation The delta review of #574 found that pickle protocols 0 and 1 read __slots__ from the instance: delegated to a slotted wrapped node, it made the wrapper look like a slotted class without __getstate__, and pickling raised TypeError. __slots__ describes the wrapper's own layout, so the guard answers it with AttributeError like the copy and pickle protocol. The copy tests now pickle with every protocol, and a slotted node is covered for both wrappers. Co-Authored-By: Claude Opus 5.5 --- .../tracing_parameter_node_at_instant.py | 7 +++- .../core/test_tracing_parameter_isolation.py | 37 ++++++++++++++++--- 2 files changed, 37 insertions(+), 7 deletions(-) diff --git a/policyengine_core/tracers/tracing_parameter_node_at_instant.py b/policyengine_core/tracers/tracing_parameter_node_at_instant.py index 76eaab01c..7b2e46db1 100644 --- a/policyengine_core/tracers/tracing_parameter_node_at_instant.py +++ b/policyengine_core/tracers/tracing_parameter_node_at_instant.py @@ -23,9 +23,11 @@ # Copying and pickling look these names up on an instance. Answered by the # wrapped object, they would copy or restore the wrapped object, not the -# wrapper. +# wrapper. ``__slots__`` describes the wrapper's own layout: pickle protocols +# 0 and 1 refuse an instance that reports slots but no ``__getstate__``. _COPY_PROTOCOL = frozenset( { + "__slots__", "__copy__", "__deepcopy__", "__getstate__", @@ -44,7 +46,8 @@ def _wrapped(wrapper: object, attribute: str, key: str) -> object: ``__getattr__`` only runs when normal lookup fails. Two such lookups must not reach the wrapped object: - - the copy and pickle protocol (``__deepcopy__``, ``__setstate__``, ...); + - the copy and pickle protocol (``__deepcopy__``, ``__setstate__``, + ``__slots__``, ...); - any name on an instance ``copy`` or ``pickle`` has created with ``__new__`` and not filled in yet. It has no ``attribute``, so reading it here would call ``__getattr__`` again, without end. diff --git a/tests/core/test_tracing_parameter_isolation.py b/tests/core/test_tracing_parameter_isolation.py index 0ae068286..8c4f5e95a 100644 --- a/tests/core/test_tracing_parameter_isolation.py +++ b/tests/core/test_tracing_parameter_isolation.py @@ -39,7 +39,7 @@ from policyengine_core.country_template import CountryTaxBenefitSystem from policyengine_core.country_template.entities import Person -from policyengine_core.parameters import ParameterNodeAtInstant +from policyengine_core.parameters import ParameterNode, ParameterNodeAtInstant from policyengine_core.periods import MONTH from policyengine_core.reforms import Reform from policyengine_core.tracers import ( @@ -305,14 +305,20 @@ def test_tracing_parameter_node_records_in_its_own_tracer_when_the_node_traces( assert type(at_instant.parameter_node_at_instant) is ParameterNodeAtInstant -def _round_trip_by_pickle(value): - return pickle.loads(pickle.dumps(value)) +def _round_trip_by_pickle(protocol): + def round_trip(value): + return pickle.loads(pickle.dumps(value, protocol=protocol)) + return round_trip + + +PICKLE_PROTOCOLS = range(pickle.HIGHEST_PROTOCOL + 1) DUPLICATORS = pytest.mark.parametrize( "duplicate", - [copy.copy, copy.deepcopy, _round_trip_by_pickle], - ids=["copy", "deepcopy", "pickle"], + [copy.copy, copy.deepcopy] + + [_round_trip_by_pickle(protocol) for protocol in PICKLE_PROTOCOLS], + ids=["copy", "deepcopy"] + [f"pickle-{protocol}" for protocol in PICKLE_PROTOCOLS], ) @@ -371,6 +377,27 @@ def test_parameters_holding_traced_nodes_can_be_deep_copied( ) +class _SlottedParameterNode(ParameterNode): + __slots__ = ("extra",) + + def __getstate__(self): + return self.__dict__ + + +@pytest.mark.parametrize("protocol", PICKLE_PROTOCOLS) +def test_wrappers_of_a_slotted_node_can_be_pickled(protocol): + """Protocols 0 and 1 read ``__slots__`` from the instance, not the class.""" + node = _SlottedParameterNode( + data={"rate": {"values": {"2015-01-01": {"value": 1}}}} + ) + view = TracingParameterNode(node, FullTracer(), "default") + + for wrapper in (view, view("2015-01-01")): + restored = pickle.loads(pickle.dumps(wrapper, protocol=protocol)) + assert type(restored) is type(wrapper) + assert pickle.loads(pickle.dumps(view, protocol=protocol))("2015-01-01").rate == 1 + + @pytest.mark.parametrize( "wrapper_class", [TracingParameterNode, TracingParameterNodeAtInstant] )