Skip to content

Let simulations be deep-copied and unpickled - #567

Open
MaxGhenis wants to merge 1 commit into
masterfrom
fix/simulation-pickle-recursion
Open

MaxGhenis wants to merge 1 commit into
masterfrom
fix/simulation-pickle-recursion

Conversation

@MaxGhenis

@MaxGhenis MaxGhenis commented Oct 2, 2026 •

Copy link
Copy Markdown
Contributor

Summary

On master (3.32.12, b78b0ba) no simulation can be copied or unpickled:

import copy, pickle
from policyengine_core.country_template import CountryTaxBenefitSystem
from policyengine_core.simulations import SimulationBuilder

s = SimulationBuilder().build_from_entities(
    CountryTaxBenefitSystem(),
    {"persons": {"a": {"salary": {"2025-01": 1000}}}, "households": {"h": {"parents": ["a"]}}},
)
s.calculate("income_tax", "2025-01")
copy.deepcopy(s)                 # RecursionError
pickle.loads(pickle.dumps(s))    # RecursionError (dumps succeeds; loads fails)
copy.copy(s.persons)             # RecursionError

Found by the adversarial review of #561. Same result on Python 3.11, 3.12, 3.13 and 3.14.

Cause

copy and pickle rebuild an object by creating an empty instance and probing it for __setstate__ before restoring its __dict__. Population.__getattr__ answers every missing attribute through get_projector_from_shortcut, which reads population.entity. On the empty instance entity is missing too, so that read calls __getattr__("entity"), which reads population.entity again, until the interpreter raises RecursionError. hasattr only swallows AttributeError, so the error escapes.

Fixing that exposed three more defects on the same path, all fixed here:

Where Defect on master
Population.__getattr__ Recurses on an unfilled instance (the report above).
VectorialParameterNodeAtInstant.__getattr__ Same recursion through self.vector when unpickled. It also forwards __deepcopy__ to the vector, and copy.deepcopy looks that method up on the instance, so a deep copy came back as a bare numpy.recarray. These nodes are cached on the parameter nodes a tax-benefit system keeps, so a copied simulation would carry the wrong type.
TracingParameterNodeAtInstant.__getattr__ Same recursion through self.parameter_node_at_instant.
EnumArray numpy's __reduce__ rebuilds the array without possible_values, so an unpickled array could not be decoded or compared with an enum item (AttributeError). Formulas that read array.possible_values failed on an unpickled simulation.

Reform.__getattr__ has the same shape but terminates: TaxBenefitSystem defines baseline = None on the class. Dataset.__getattr__ was already guarded.

Changes

  • Population, VectorialParameterNodeAtInstant, TracingParameterNodeAtInstant: __getattr__ raises AttributeError for the instance's own attribute(s) it reads, so a lookup on an unfilled instance ends. The vectorial node also stops forwarding __deepcopy__. No other lookup changes: projector shortcuts and vector attributes resolve as before.
  • EnumArray.__reduce__ carries the enum by module and qualified name. Unpickling resolves it when this process can find it, and otherwise returns the array with possible_values = None. Pickling the enum by reference instead would have made arrays that cross a process boundary today fail to unpickle at all (see limits). Pickles written by earlier releases still load.

Invariants

Stated and tested (tests/core/test_simulation_copy_pickle.py, tests/core/test_simulation_copy_pickle_property.py):

  1. A copy is a simulation in its own right (differential). For any country-template situation (1-4 people, any partition into households, any enum input, any subset of 7 variables calculated beforehand, with or without a branch holding its own input), copy.deepcopy(s) and pickle.loads(pickle.dumps(s)) calculate what a freshly built simulation calculates, for 8 variables and for the branch.
  2. Independence. Inputs set and values calculated on the copy never appear in the original, and the original still calculates what it did.
  3. Termination. Any attribute lookup on an unfilled Population, GroupPopulation, vectorial node or tracing node raises AttributeError.
  4. Type preservation. A copied vectorial node is a vectorial node; a copied tracing wrapper is a tracing wrapper around the same kind of node; a copied EnumArray is an EnumArray with the same values and dtype.
  5. Enum round trip. possible_values after unpickling is the same enum when the process can find it and None otherwise; unpickling never raises because of the enum.

The Hypothesis property (1, 2) ran 2,000 examples locally. It is in its own module behind pytest.importorskip("hypothesis"), because the country-package smoke job installs no dev extras.

Each guard was checked by mutation: removing any one of them fails between 1 and 17 of the new tests.

Limits (not changed here)

  • Pickles are for the process that wrote them. add_variables_from_file registers each variable file under f"{id(self)}_{hash(path)}_{file_name}", a name no other process has, so a simulation (or tax-benefit system, or variable) unpickled elsewhere raises ModuleNotFoundError. Changing that naming is not small: Simulations pickled in one process cannot be unpickled in another #568.
  • policyengine-us systems still do not copy. Its spm_forecast_provider holds spm_calculator's SPMForecast, whose MappingProxyType fields cannot be copied or pickled: SPMForecast cannot be pickled or deep-copied (mappingproxy fields) spm-calculator#49. With the tax-benefit systems shared between original and copy, a policyengine-us (2.21.0) household simulation deep-copies and round-trips through pickle in the same process with this branch, and recalculates the same values.
  • A copy of a simulation whose holders store arrays on disk points at the same directory as the original.

Downstream use

git grep on the default branches of policyengine.py, policyengine-api, policyengine-api-v2, policyengine-household-api, policyengine-us, policyengine-uk, policyengine-canada, policyengine-us-data and policyengine-uk-data: nothing pickles or copies a simulation. Their deepcopy calls are on JSON-like dicts, and the us-data worker pools pass file paths and arrays and build one Microsimulation per worker. policyengine-us clones systems through its own clone_spm_system. #560 pickles InMemoryStorage and StoreHistory on their own; it touches none of these files.

Tests

  • uv run pytest tests on Python 3.14: 1183 passed, 4 skipped, 1 xfailed
  • New and neighbouring tests (enums, tracers, fancy indexing, projectors, entities) on Python 3.11: 126 passed
  • policyengine-core test policyengine_core/country_template/tests -c policyengine_core.country_template: 39 passed
  • ruff format --check ., ruff check .: clean
  • Smoke-style collection without Hypothesis installed (RUN_SMOKE_TESTS=0 python -m pytest -m smoke --collect-only): exit 0, property module skipped

Not run: make documentation, mypy.

axiom: n/a: engine infrastructure, no policy rule.

🤖 Generated with Claude Code

copy and pickle rebuild an object by creating an empty instance and
probing it for __setstate__ before restoring its __dict__.
Population.__getattr__ answered that probe through the projector lookup,
which reads self.entity; on the empty instance that read re-entered
__getattr__ until RecursionError. So copy.deepcopy(simulation),
pickle.loads(pickle.dumps(simulation)) and copy.copy(population) failed
on every simulation, on Python 3.11 to 3.14.

The same path had three more defects:

- VectorialParameterNodeAtInstant and TracingParameterNodeAtInstant
  recursed the same way through the attribute they forward to.
- The vectorial node forwarded __deepcopy__ to its numpy vector, so a
  deep copy came back as a bare recarray. These nodes are cached on the
  parameter nodes a tax-benefit system keeps.
- numpy's __reduce__ rebuilt an EnumArray without possible_values, so an
  unpickled enum array could be neither decoded nor compared with an
  enum item. EnumArray now pickles its enum by name and restores it when
  the process can find it, otherwise None.

Tests: examples for each defect, plus a Hypothesis property (in its own
importorskip module, for the smoke job) that a deep copy or pickle round
trip calculates what a freshly built simulation does and that writes to
the copy never reach the original.

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