Skip to content

Let simulations be deep-copied and unpickled - #567

Merged
MaxGhenis merged 3 commits into
masterfrom
fix/simulation-pickle-recursion
Oct 8, 2026
Merged

MaxGhenis merged 3 commits into
masterfrom
fix/simulation-pickle-recursion

Conversation

@MaxGhenis

@MaxGhenis MaxGhenis commented Oct 2, 2026 •

Copy link
Copy Markdown
Contributor

Summary

On the original base (3.32.12, b78b0ba), the original author reported that simulations could not 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 the original base (b78b0ba)
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.__getattr__ rejects its missing backing attributes and copy/pickle protocol probes before projector lookup. This also handles subclasses with a property-backed entity during reconstruction; ordinary projector shortcuts still resolve.
  • Master is merged without rewriting history. The vectorial parameter and tracing wrapper implementations are exactly master's versions, as requested in the resolution comment. This retains Let vectorial parameter nodes be copied and pickled #584's field access and child-node construction fixes and Trace parameter reads per call and stop test reruns leaking module fixtures #574's protocol guard, including __slots__.
  • EnumArray.__reduce__ carries the concrete array class and the enum's module and qualified name. Reconstruction preserves the subclass without invoking its constructor. If enum lookup raises ImportError or AttributeError, possible_values is None; other import/initialization errors propagate.
  • Legacy NumPy payloads lacking possible_values still load and can now be re-pickled with that metadata set to None. Reconstruction also accepts the earlier two-argument reducer payload.

Invariants

The existing simulation differential/property coverage is preserved; these checks are defined in tests/core/test_simulation_copy_pickle.py and tests/core/test_simulation_copy_pickle_property.py:

  1. A deep-copied or same-process unpickled simulation calculates the same country-template outputs as a freshly built simulation, including warmed caches and branch inputs.
  2. Inputs and calculated values written on a copy do not reach the original.
  3. Copy/pickle probes on unfilled populations terminate with AttributeError, including property-backed population subclasses during reconstruction.
  4. Enum array copy and pickle preserve the concrete class, encoded indices, dtype, shape, and resolvable enum metadata. Copies have independent array storage. A focused Hypothesis test compares these results with a plain NumPy array reference.
  5. Enum lookup falls back to None on ImportError/AttributeError; legacy arrays without metadata remain serializable.

The new regression cases cover EnumArray subclasses and legacy re-pickling under protocols 0–5, shallow/deep copying, property-backed populations, slotted scalar/vectorial tracing nodes under protocols 0–5, and a legitimate parameter field named vector. This sweep could not execute them because the required test wrapper is blocked by the sandbox.

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.

Validation

Current repair head: f310a8e.

  • Focused Ruff formatting and lint checks passed for the seven affected Python files; git diff --check passed.
  • Required foreground test wrapper attempted before and after repair. Both attempts exited 75 before pytest: the sandbox denied ps, so heavy_run.sh refused its nice-level check. Zero tests executed; the wrapper was not bypassed.
  • Intended targeted command: ~/reviews/us-hub/scripts/heavy_run.sh core567-final uv run --no-sync pytest tests/core/test_simulation_copy_pickle.py tests/core/test_simulation_copy_pickle_property.py tests/core/parameters/test_vectorial_parameter_node_copy.py tests/core/parameters/test_vectorial_parameter_node_copy_property.py tests/core/test_tracing_parameter_isolation.py tests/core/test_tracing_parameter_isolation_properties.py tests/core/test_tracers.py tests/core/enums/test_enum.py tests/core/test_projectors.py tests/core/parameters_fancy_indexing/test_fancy_indexing.py -q -p no:cacheprovider -n 2.
  • Full suites, country-template CLI, documentation build, mypy, and downstream microsimulations were not run in this bounded sweep. CI must validate the new head; previous green checks belong to b6c68eb.

Downstream impact

Re-run pending at f310a8e. The hub's existing downstream A/B gate remains outstanding. This sweep ran no microsimulation and changed no partner baseline expectations. #570 must land after #567 and retain protocols 0–5 support.

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>
@MaxGhenis

Copy link
Copy Markdown
Contributor Author

#574 and #584 have merged and now cover two of this PR's hunks, so this branch conflicts with master in two files:

Resolution: take master's version of both files. Master's guard covers what this PR's two hunks did (the unfilled-instance recursion and the delegated __deepcopy__), plus the rest of the copy and pickle protocol names. #584 also fixes the child-node constructor call that this PR's hunk left with one argument.

Checked locally: I merged master (042adb5) into this branch at b6c68eb and took master's side of those two files, with no other conflicts. This PR's test_simulation_copy_pickle.py and test_simulation_copy_pickle_property.py pass together with #584's and #574's tests and test_tracers.py: 197 passed on Python 3.11.

The Population and EnumArray changes here are untouched by either PR and still needed. I have not pushed anything to this branch.

MaxGhenis and others added 2 commits October 7, 2026 16:56
Keep master parameter and tracing implementations while retaining Population and EnumArray repairs.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Guard Population protocol probes before projector lookup. Add regressions for protocols 0-5, property-backed populations, slotted tracing nodes, vector fields, and NumPy differential copy behavior.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@MaxGhenis
MaxGhenis merged commit 0946ba3 into master Oct 8, 2026
19 checks passed
@MaxGhenis
MaxGhenis deleted the fix/simulation-pickle-recursion branch October 8, 2026 15:22
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