Skip to content

Trace parameter reads per call and stop test reruns leaking module fixtures - #574

Open
MaxGhenis wants to merge 1 commit into
masterfrom
fix/rerun-leaked-traced-system
Open

MaxGhenis wants to merge 1 commit into
masterfrom
fix/rerun-leaked-traced-system

Conversation

@MaxGhenis

Copy link
Copy Markdown
Contributor

Summary

On #560's Windows CI (runs 36967463690, 36968947271, 36969160069, 36972147369) a flaky test was rerun, and afterwards tests/core/test_parameters.py::test_get_at_instant got a TracingParameterNodeAtInstant and tests/core/test_reforms.py::test_modify_parameters overflowed the stack, both on the same system object. Three defects chained together. This PR fixes each one, so breaking any single link would already have prevented the failure.

1. A rerun leaked the module-scoped tax_benefit_system into later modules (test tooling)

make test runs pytest --reruns 2. The installed plugin is pytest-rerunfailures 14.0, the newest the old <15 pin allows, and the lock pairs it with pytest 9.1.1. Before a rerun it empties item.session._setupstate.stack (_remove_failed_setup_state_from_session) without running the finalizers on it. pytest registers a fixture's finalizer with its scope node only when it creates the value (FixtureDef.execute returns cached_result early otherwise). So the rerun module's tax_benefit_system is never torn down, and every later module gets the same cached object.

pytest-rerunfailures restores those finalizers from 15.0 (pytest-dev/pytest-rerunfailures#278, "fix compatibility with pytest 8.2 by restoring deleted finalizers"). 16.2 is the first release that declares pytest 9 support. Fix: the dev extra and the smoke job now require pytest-rerunfailures>=16.2,<17, and uv.lock moves to 16.7. The lock also syncs its stale policyengine-core version line (3.32.11 → 3.32.12, matching pyproject.toml).

2. Tracing switched tracing on in the shared parameter tree (library)

When a traced simulation ran a formula, Simulation._run_formula set trace, tracer and branch_name on tax_benefit_system.parameters (the "soft-recast") and never reset them. The consequences:

  • The system stayed traced. The root node caches one node per instant, and after the soft-recast that is a TracingParameterNodeAtInstant holding the tracer and branch name of that moment. Every later simulation on the system, traced or not, read through it. A second traced simulation's reads went to the first simulation's tracer, and a branch's reads were filed under another branch's name. ParameterNode.clone copies trace and tracer, so clones traced too.
  • Traces missed formula reads. Simulation._calculate reads parameters(period) (the abolition check) before any formula runs, which caches an untraced node at that instant. So on master a traced income_tax for 2017-01 records no parameter reads at all. The soft-recast only reached instants first read after it.

policyengine-us's SharedParameterPolicy (spm.py) lends one parameter tree to several systems and documents that writing to it "must not happen at all"; the soft-recast was such a write.

Fix: a traced simulation's formulas now get a TracingParameterNode, a per-call view of the tree. Calling it at an instant returns a TracingParameterNodeAtInstant with this simulation's tracer and branch name. Other attributes are read from the wrapped node, which is never modified. Untraced simulations are unchanged: they still get tax_benefit_system.parameters itself.

Behaviour change: traces now record the parameters a formula reads through its parameters argument, at every instant. Reads made elsewhere aren't recorded: the abolition check, and code that reaches simulation.tax_benefit_system.parameters directly. On master those were recorded only at instants first read after the soft-recast. policyengine.py's derive reads only node.children, so it is unaffected.

3. TracingParameterNodeAtInstant could not be copied (library)

copy.copy, copy.deepcopy and pickle all recursed. Each creates the instance with __new__ and then probes it (__setstate__, __deepcopy__, ...), and __getattr__ read self.parameter_node_at_instant, which isn't set yet, so it called itself without end. Reform.modify_parameters deep-copies the baseline's parameter tree, so once a traced node sat in the root's at-instant cache, any reform that called modify_parameters on that system crashed. Fix: both wrappers raise AttributeError for special (__dunder__) names and for any lookup on an instance that hasn't been filled in. Copy, deepcopy and pickle now round-trip.

Options considered

Option Verdict
(1) Function-scoped tax_benefit_system Not adopted. Measured: 228 tests in 25 modules use it; construction takes 7.2 ms median (p90 10.7 ms), so about 1.5 s more on a 40–70 s local suite. It would also need five module-scoped autouse fixtures (in test_cycles, test_formulas, test_reforms, test_opt_out_cache and test_calculate_output) made function-scoped, and it fails Hypothesis' function_scoped_fixture health check in test_branch_shared_arrays_differential.py. The plugin bump fixes the leak for every shared fixture at every scope; function scope would cover this one fixture only.
Bump pytest-rerunfailures Adopted. It fixes the cause, and 14.0 doesn't support the pytest 9 that CI installs.
(2) Stop tracing from changing the shared system Adopted, as a per-call TracingParameterNode rather than restoring flags afterwards. Restoring would still leave traced nodes in the at-instant caches, and it races when threads share a system.
(3) Make TracingParameterNodeAtInstant copy-safe Adopted.

Reproduction

  • End to end, before: I put a temporary module that fails once ahead of test_branch_shared_arrays.py and ran it with test_parameters.py and test_reforms.py under --reruns 2 with 14.0. All four modules got the same system object. test_get_at_instant failed (TracingParameterNodeAtInstant) and test_modify_parameters raised RecursionError, which is the CI failure.
  • After, with 16.7: each module gets its own system and all 92 tests pass.
  • After, with 14.0 kept but this PR's library code: the four modules still share one system object, but nothing traces it, so all 92 tests pass. The library fixes break the chain on their own.
  • tests/core/test_rerun_fixture_isolation.py runs a two-module suite in a subprocess with the installed plugin, once with a test that passes on its rerun and once with one that fails every attempt. With 14.0 both cases fail (assert 'test_a_rerun' == 'test_b_later'); with 15.0, 16.2 and 16.7 both pass.
  • tests/core/test_tracing_parameter_isolation.py: 28 of its 30 tests fail against master's library code. The other two are non-regression checks.
  • tests/core/test_tracing_parameter_isolation_properties.py: against master, Hypothesis shrinks the failure to a single traced calculation that leaves the system traced.

Invariants

These hold for every input and are stated and executed in the new tests:

  • Tracing a simulation leaves its system's parameter tree as an untraced one: no trace/tracer/branch_name, only plain cached at-instant nodes, and clones and deep copies plain too.
  • A traced simulation records exactly the parameter reads it would record on a fresh system, so the record never depends on history. A branch records only reads such a calculation makes, all under its own branch name. (Property test over random sequences of traced, untraced, branch and clone steps.)
  • Tracing never changes a value. A differential test compares traced and untraced runs for every formula variable of the country template, and the property test compares each step with an untraced reference system.
  • A rerun never hands a module-scoped fixture to a later module, and the earlier module's instance is torn down first.

Tests

  • New: tests/core/test_tracing_parameter_isolation.py (30), tests/core/test_tracing_parameter_isolation_properties.py (Hypothesis, 40 examples), tests/core/test_rerun_fixture_isolation.py (2), and helpers in tests/fixtures/tracing.py.
  • Full suite with the Makefile command (coverage run ... --reruns 2 --reruns-delay 5): running locally; this PR's CI is the authoritative run.
  • ruff format --check . and ruff check pass.

Follow-up

Iterating a traced parameter node (for name in node, "x" in node) raises KeyError: 0, because the wrapper has __getitem__ but no __iter__. It's out of scope here and filed as a separate task.

axiom: n/a: engine tracing and test tooling, no policy change

🤖 Generated with Claude Code

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 <noreply@anthropic.com>

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant