diff --git a/changelog.d/fix-simulation-copy-pickle-recursion.fixed.md b/changelog.d/fix-simulation-copy-pickle-recursion.fixed.md new file mode 100644 index 00000000..1332e922 --- /dev/null +++ b/changelog.d/fix-simulation-copy-pickle-recursion.fixed.md @@ -0,0 +1 @@ +Stop `copy.deepcopy(simulation)`, `copy.copy(population)` and unpickling a simulation in the process that pickled it from raising `RecursionError`, keep vectorial parameter nodes as nodes when deep-copied, and keep an `EnumArray`'s `possible_values` through pickling. diff --git a/policyengine_core/enums/enum_array.py b/policyengine_core/enums/enum_array.py index c518add8..36be3f03 100644 --- a/policyengine_core/enums/enum_array.py +++ b/policyengine_core/enums/enum_array.py @@ -1,7 +1,9 @@ from __future__ import annotations +import importlib +import operator import typing -from typing import Any, NoReturn, Optional, Type +from typing import Any, NoReturn, Optional, Tuple, Type import numpy @@ -9,6 +11,31 @@ from policyengine_core.enums import Enum +def _restore_enum_array( + array: numpy.ndarray, enum_name: Optional[Tuple[str, str]] +) -> EnumArray: + """Rebuild a pickled EnumArray, with its enum if this process can find it. + + ``enum_name`` is the enum's module and qualified name, or ``None`` for an + array that had no enum. + + Tax-benefit systems load variable files under module names that exist only + in the process that loaded them, so an enum defined in one cannot be found + from another process. The array then comes back with ``possible_values`` + unset (``None``), where pickling the enum by reference would fail to + unpickle the array at all. + """ + possible_values = None + if enum_name is not None: + module_name, qualified_name = enum_name + try: + module = importlib.import_module(module_name) + possible_values = operator.attrgetter(qualified_name)(module) + except (ImportError, AttributeError): + pass + return EnumArray(array, possible_values) + + class EnumArray(numpy.ndarray): """ Numpy array subclass representing an array of enum items. @@ -35,6 +62,15 @@ def __array_finalize__(self, obj: Optional[numpy.int_]) -> None: self.possible_values = getattr(obj, "possible_values", None) + def __reduce__(self) -> tuple: + # ndarray's own ``__reduce__`` rebuilds the array without + # ``possible_values``, so an unpickled EnumArray could be neither + # decoded nor compared with an enum item. The enum travels by name + # rather than by reference; see ``_restore_enum_array``. + enum = self.possible_values + name = None if enum is None else (enum.__module__, enum.__qualname__) + return _restore_enum_array, (self.view(numpy.ndarray), name) + def __eq__(self, other: Any) -> bool: # When comparing to an item of self.possible_values, use the item index # to speed up the comparison. diff --git a/policyengine_core/parameters/vectorial_parameter_node_at_instant.py b/policyengine_core/parameters/vectorial_parameter_node_at_instant.py index fcc562ef..81f1a9d3 100644 --- a/policyengine_core/parameters/vectorial_parameter_node_at_instant.py +++ b/policyengine_core/parameters/vectorial_parameter_node_at_instant.py @@ -209,6 +209,14 @@ def __init__(self, name: str, vector: ArrayLike, instant_str: str): self._instant_str = instant_str def __getattr__(self, attribute: str) -> Any: + # ``vector`` is missing while copy or pickle rebuilds a node, and + # looking it up would recurse here. ``copy.deepcopy`` looks + # ``__deepcopy__`` up on the instance, and the vector's would copy the + # vector alone and hand back a bare ``numpy.recarray``. + if attribute in ("vector", "__deepcopy__"): + raise AttributeError( + f"{type(self).__name__!s} has no attribute {attribute!r}" + ) result = getattr(self.vector, attribute) if isinstance(result, numpy.recarray): return VectorialParameterNodeAtInstant(result) diff --git a/policyengine_core/populations/population.py b/policyengine_core/populations/population.py index d2e8df1d..d391b810 100644 --- a/policyengine_core/populations/population.py +++ b/policyengine_core/populations/population.py @@ -53,6 +53,14 @@ def filled_array(self, value: Any, dtype: Any = None) -> numpy.ndarray: return numpy.full(self.count, value, dtype) def __getattr__(self, attribute: str) -> Any: + # The shortcut lookup reads ``self.entity`` and ``self.simulation``. + # They are missing while copy or pickle rebuilds a population (both + # probe the new, empty instance for ``__setstate__`` before restoring + # its ``__dict__``), and looking them up would recurse here. + if attribute in ("entity", "simulation"): + raise AttributeError( + f"{type(self).__name__!s} has no attribute {attribute!r}" + ) projector = projectors.get_projector_from_shortcut(self, attribute) if not projector: raise AttributeError( diff --git a/policyengine_core/tracers/tracing_parameter_node_at_instant.py b/policyengine_core/tracers/tracing_parameter_node_at_instant.py index d4bf1746..5d3fde49 100644 --- a/policyengine_core/tracers/tracing_parameter_node_at_instant.py +++ b/policyengine_core/tracers/tracing_parameter_node_at_instant.py @@ -36,6 +36,10 @@ def __getattr__( self, key: str, ) -> Union[TracingParameterNodeAtInstant, Child]: + # ``parameter_node_at_instant`` is missing while copy or pickle + # rebuilds a wrapper, and looking it up would recurse here. + if key == "parameter_node_at_instant": + raise AttributeError(f"{type(self).__name__!s} has no attribute {key!r}") child = getattr(self.parameter_node_at_instant, key) return self.get_traced_child(child, key) diff --git a/tests/core/test_simulation_copy_pickle.py b/tests/core/test_simulation_copy_pickle.py new file mode 100644 index 00000000..89fc212a --- /dev/null +++ b/tests/core/test_simulation_copy_pickle.py @@ -0,0 +1,341 @@ +"""Simulations survive ``copy.deepcopy`` and a pickle round trip. + +``copy`` and ``pickle`` rebuild an object by creating an empty instance and +then probing it for ``__setstate__`` before its ``__dict__`` is restored. A +``__getattr__`` that reads one of the instance's own attributes recurses on +that probe until the interpreter raises ``RecursionError``. +``Population.__getattr__`` did this through its projector lookup (it reads +``self.entity``), so ``copy.deepcopy(simulation)``, +``pickle.loads(pickle.dumps(simulation))`` and even ``copy.copy(population)`` +failed on every simulation. The parameter wrappers that forward attribute +lookups (vectorial and tracing nodes) had the same recursion, and the +vectorial node also forwarded ``__deepcopy__`` to its numpy vector, so a deep +copy came back as a bare ``numpy.recarray``. Behind the recursion, a pickled +``EnumArray`` lost its ``possible_values``. + +Pickles are for the process that wrote them: a tax-benefit system loads +variable files under module names no other process has. +""" + +from __future__ import annotations + +import copy +import io +import pickle + +import numpy as np +import pytest + +from policyengine_core.country_template import CountryTaxBenefitSystem +from policyengine_core.enums import Enum, EnumArray +from policyengine_core.parameters import VectorialParameterNodeAtInstant +from policyengine_core.populations import GroupPopulation, Population +from policyengine_core.tools import assert_near +from policyengine_core.tracers import FullTracer, TracingParameterNodeAtInstant +from tests.fixtures.simulation_copy import ( + COPIERS, + FEB, + JAN, + DoubleIncomeTaxRate, + build_simulation, + rate_node, +) + +# Population: lookups on an instance that copy or pickle has not filled yet. + + +@pytest.mark.parametrize("population_class", [Population, GroupPopulation]) +def test_lookups_on_an_unfilled_population_raise_attribute_error(population_class): + unfilled = population_class.__new__(population_class) + + assert getattr(unfilled, "__setstate__", None) is None + assert not hasattr(unfilled, "household") + with pytest.raises(AttributeError): + unfilled.entity + with pytest.raises(AttributeError): + unfilled.simulation + + +def test_copy_protocol_lookups_on_a_population_find_nothing(): + simulation = build_simulation() + + assert getattr(simulation.persons, "__deepcopy__", None) is None + assert getattr(simulation.household, "__setstate__", None) is None + + +def test_population_projector_shortcuts_still_resolve(): + simulation = build_simulation() + + assert_near( + simulation.persons.household("rent", JAN), simulation.household("rent", JAN) + ) + with pytest.raises(AttributeError, match="not a known attribute"): + simulation.persons.not_an_entity + + +def test_shallow_copy_of_a_population_shares_its_holders(): + simulation = build_simulation() + + population = copy.copy(simulation.persons) + + assert population.entity is simulation.persons.entity + assert population._holders is simulation.persons._holders + + +# Simulations. + + +@pytest.mark.parametrize("copier", COPIERS.values(), ids=COPIERS.keys()) +def test_copied_simulation_calculates_what_the_original_does(copier): + simulation = build_simulation(salaries=(1000, 2500)) + expected = simulation.calculate("income_tax", JAN) + + copied = copier(simulation) + + assert_near(copied.calculate("income_tax", JAN), expected) + assert_near( + copied.calculate("disposable_income", JAN), + simulation.calculate("disposable_income", JAN), + ) + assert copied.persons is not simulation.persons + assert copied.persons.simulation is copied + assert copied.household.simulation is copied + assert copied.household.members is copied.persons + assert copied.populations["person"] is copied.persons + + +@pytest.mark.parametrize("copier", COPIERS.values(), ids=COPIERS.keys()) +def test_inputs_set_on_a_copy_do_not_reach_the_original(copier): + simulation = build_simulation() + simulation.calculate("income_tax", JAN) + + copied = copier(simulation) + copied.set_input("salary", FEB, [2000]) + + assert_near(copied.calculate("income_tax", FEB), [300]) + assert simulation.persons.get_holder("salary").get_array(FEB) is None + assert simulation.persons.get_holder("income_tax").get_array(FEB) is None + + +@pytest.mark.parametrize("copier", COPIERS.values(), ids=COPIERS.keys()) +def test_copied_simulation_keeps_its_branches(copier): + simulation = build_simulation() + simulation.get_branch("child").set_input("salary", FEB, [2000]) + + copied = copier(simulation) + + child = copied.branches["child"] + assert child.parent_branch is copied + assert child is not simulation.branches["child"] + assert_near(child.calculate("income_tax", FEB), [300]) + + +@pytest.mark.parametrize("copier", COPIERS.values(), ids=COPIERS.keys()) +def test_copied_reform_simulation_keeps_the_reform(copier): + reformed = DoubleIncomeTaxRate(CountryTaxBenefitSystem()) + simulation = build_simulation(tax_benefit_system=reformed) + expected = simulation.calculate("income_tax", JAN) + + copied = copier(simulation) + + assert isinstance(copied.tax_benefit_system, DoubleIncomeTaxRate) + assert_near(copied.calculate("income_tax", JAN), expected) + assert_near(expected, [300]) + + +@pytest.mark.parametrize("copier", COPIERS.values(), ids=COPIERS.keys()) +def test_copied_traced_simulation_still_traces(copier): + simulation = build_simulation() + simulation.trace = True + simulation.calculate("income_tax", JAN) + + copied = copier(simulation) + copied.calculate("disposable_income", JAN) + + assert copied.tracer is not simulation.tracer + assert "disposable_income<2025-01, (default)>" in copied.tracer.get_flat_trace() + + +@pytest.mark.parametrize("copier", COPIERS.values(), ids=COPIERS.keys()) +def test_copied_simulation_keeps_cached_vectorial_parameter_nodes(copier): + tax_benefit_system = CountryTaxBenefitSystem() + tax_benefit_system.parameters.add_child("rate", rate_node()) + simulation = build_simulation(tax_benefit_system=tax_benefit_system) + statuses = np.array(["couple", "single"]) + parameters = simulation.tax_benefit_system.get_parameters_at_instant("2015-01-01") + assert_near(parameters.rate[statuses].owner, [500, 100]) # caches the node + + copied = copier(simulation) + + rate = copied.tax_benefit_system.get_parameters_at_instant("2015-01-01").rate + assert isinstance(rate._vectorial_node, VectorialParameterNodeAtInstant) + assert_near(rate[statuses].tenant, [700, 300]) + + +@pytest.mark.parametrize("copier", COPIERS.values(), ids=COPIERS.keys()) +def test_copied_simulation_keeps_enum_inputs(copier): + simulation = build_simulation( + salaries=(1000, 1000), + households=[[0], [1]], + occupancy=["owner", "free_lodger"], + ) + + copied = copier(simulation) + + status = copied.household("housing_occupancy_status", JAN) + assert list(status.decode_to_str()) == ["owner", "free_lodger"] + # The formula reads the enum from the stored array's ``possible_values``. + housing_tax = copied.calculate("housing_tax", "2025") + assert_near(housing_tax, simulation.calculate("housing_tax", "2025")) + assert housing_tax[0] > 0 + assert housing_tax[1] == 0 + + +# Enum arrays. + + +@pytest.mark.parametrize("protocol", range(pickle.HIGHEST_PROTOCOL + 1)) +def test_pickled_enum_array_keeps_its_possible_values(protocol): + status = build_simulation(occupancy=["free_lodger"]).household( + "housing_occupancy_status", JAN + ) + + restored = pickle.loads(pickle.dumps(status, protocol=protocol)) + + assert type(restored) is EnumArray + assert restored.possible_values is status.possible_values + assert list(restored.decode_to_str()) == ["free_lodger"] + assert (restored == status.possible_values.free_lodger).all() + assert restored.dtype == status.dtype + + +def test_pickled_enum_array_without_possible_values_stays_without(): + restored = pickle.loads(pickle.dumps(EnumArray(np.array([0, 1])))) + + assert type(restored) is EnumArray + assert restored.possible_values is None + assert restored.tolist() == [0, 1] + + +@pytest.mark.parametrize( + "module_name, qualified_name", + [ + ("module_only_another_process_loaded", "HousingOccupancyStatus"), + ("policyengine_core.enums", "NoSuchEnum"), + ], +) +def test_enum_array_whose_enum_cannot_be_found_unpickles_without_it( + module_name, qualified_name +): + """An enum from a variable file has a module name only its process knows.""" + + class Unreachable(Enum): + a = "a" + b = "b" + + Unreachable.__module__ = module_name + Unreachable.__qualname__ = qualified_name + array = EnumArray(np.array([1, 0, 1]), Unreachable) + + restored = pickle.loads(pickle.dumps(array)) + + assert type(restored) is EnumArray + assert restored.possible_values is None + assert restored.tolist() == [1, 0, 1] + + +def test_enum_array_pickled_by_an_earlier_release_still_loads(): + """Earlier releases pickled an EnumArray as numpy pickles any ndarray.""" + status = build_simulation(occupancy=["owner"]).household( + "housing_occupancy_status", JAN + ) + + class EarlierPickler(pickle.Pickler): + def reducer_override(self, obj): + if isinstance(obj, EnumArray): + return np.ndarray.__reduce__(obj) + return NotImplemented + + buffer = io.BytesIO() + EarlierPickler(buffer).dump(status) + + restored = pickle.loads(buffer.getvalue()) + + assert type(restored) is EnumArray + assert restored.tolist() == status.tolist() + + +def test_copied_enum_array_keeps_its_possible_values(): + status = build_simulation(occupancy=["owner"]).household( + "housing_occupancy_status", JAN + ) + + for copied in (copy.copy(status), copy.deepcopy(status)): + assert type(copied) is EnumArray + assert copied.possible_values is status.possible_values + assert list(copied.decode_to_str()) == ["owner"] + assert not np.shares_memory(copied, status) + + +# Parameter wrappers that forward attribute lookups. + + +@pytest.mark.parametrize("copier", COPIERS.values(), ids=COPIERS.keys()) +def test_copied_vectorial_parameter_node_is_a_vectorial_node(copier): + node = rate_node()("2015-01-01")[np.array(["single", "couple", "single"])] + + copied = copier(node) + + assert type(copied) is VectorialParameterNodeAtInstant + assert copied._name == node._name + assert copied._instant_str == node._instant_str + assert_near(copied.owner, [100, 500, 100]) + assert_near(copied[np.array(["owner", "tenant", "tenant"])], [100, 700, 300]) + assert not np.shares_memory(copied.vector, node.vector) + + +def test_vectorial_parameter_node_still_forwards_to_its_vector(): + node = rate_node()("2015-01-01")[np.array(["single", "couple"])] + + assert set(node.dtype.names) == {"owner", "tenant"} + assert node.shape == (2,) + + +@pytest.mark.parametrize("copier", COPIERS.values(), ids=COPIERS.keys()) +def test_copied_tracing_parameter_node_still_traces(copier): + tracer = FullTracer() + tracer.record_calculation_start("v", "2015-01") + node = TracingParameterNodeAtInstant( + rate_node()("2015-01-01"), tracer, branch_name="default" + ) + + copied = copier(node) + + assert type(copied) is TracingParameterNodeAtInstant + assert copied.single.owner == 100 + assert copied.tracer.trees[0].parameters[-1].name == "rate.single.owner" + assert node.tracer.trees[0].parameters == [] + + +@pytest.mark.parametrize("copier", COPIERS.values(), ids=COPIERS.keys()) +def test_copied_tracing_wrapper_of_a_vectorial_node_keeps_both_types(copier): + vectorial = rate_node()("2015-01-01")[np.array(["couple", "single"])] + node = TracingParameterNodeAtInstant(vectorial, FullTracer(), "default") + + copied = copier(node) + + assert type(copied) is TracingParameterNodeAtInstant + assert type(copied.parameter_node_at_instant) is VectorialParameterNodeAtInstant + assert_near(copied.parameter_node_at_instant.owner, [500, 100]) + + +@pytest.mark.parametrize( + "wrapper_class", [VectorialParameterNodeAtInstant, TracingParameterNodeAtInstant] +) +def test_lookups_on_an_unfilled_parameter_wrapper_raise_attribute_error( + wrapper_class, +): + unfilled = wrapper_class.__new__(wrapper_class) + + assert getattr(unfilled, "__setstate__", None) is None + assert not hasattr(unfilled, "anything") diff --git a/tests/core/test_simulation_copy_pickle_property.py b/tests/core/test_simulation_copy_pickle_property.py new file mode 100644 index 00000000..84dea200 --- /dev/null +++ b/tests/core/test_simulation_copy_pickle_property.py @@ -0,0 +1,99 @@ +"""A copied simulation is a simulation in its own right. + +For random situations, warm caches and branches, a ``copy.deepcopy`` or +pickle round trip of a simulation calculates what a freshly built simulation +does, and writing to the copy leaves the original unchanged. +``test_simulation_copy_pickle.py`` pins the same behaviour with examples. +""" + +from __future__ import annotations + +import pytest + +# The country-package smoke job installs no dev dependencies. +hypothesis = pytest.importorskip("hypothesis") +st = hypothesis.strategies + +from policyengine_core.country_template import CountryTaxBenefitSystem +from policyengine_core.tools import assert_near +from tests.fixtures.simulation_copy import COPIERS, FEB, JAN, build_simulation + +YEAR = "2025" +VARIABLES = ( + "income_tax", + "social_security_contribution", + "basic_income", + "disposable_income", + "household_income", + "total_taxes", + "total_benefits", +) + + +@st.composite +def situations(draw): + salaries = draw(st.lists(st.integers(0, 20_000), min_size=1, max_size=4)) + count = len(salaries) + order = draw(st.permutations(range(count))) + cuts = sorted(draw(st.sets(st.integers(1, count - 1)))) if count > 1 else [] + bounds = [0, *cuts, count] + households = [list(order[a:b]) for a, b in zip(bounds, bounds[1:])] + occupancy = draw( + st.lists( + st.sampled_from([None, "owner", "tenant", "free_lodger", "homeless"]), + min_size=len(households), + max_size=len(households), + ) + ) + warmed = draw(st.lists(st.sampled_from(VARIABLES), unique=True)) + branch_salary = draw(st.none() | st.integers(0, 20_000)) + return salaries, households, occupancy, warmed, branch_salary + + +@hypothesis.settings( + max_examples=60, + deadline=None, + suppress_health_check=[hypothesis.HealthCheck.too_slow], +) +@hypothesis.given(situation=situations(), copier_name=st.sampled_from(sorted(COPIERS))) +def test_copy_calculates_like_a_fresh_simulation_and_stays_independent( + situation, copier_name +): + """For any situation and warm cache, a copy matches a freshly built + simulation on every variable, and writing to the copy leaves the original + unchanged.""" + salaries, households, occupancy, warmed, branch_salary = situation + tax_benefit_system = CountryTaxBenefitSystem() + original = build_simulation(tax_benefit_system, salaries, households, occupancy) + for variable in warmed: + original.calculate(variable, JAN) + if branch_salary is not None: + original.get_branch("child").set_input( + "salary", FEB, [branch_salary] * len(salaries) + ) + + copied = COPIERS[copier_name](original) + fresh = build_simulation(tax_benefit_system, salaries, households, occupancy) + + for variable in VARIABLES: + assert_near(copied.calculate(variable, JAN), fresh.calculate(variable, JAN)) + assert_near( + copied.calculate("housing_tax", YEAR), fresh.calculate("housing_tax", YEAR) + ) + assert list(copied.household("housing_occupancy_status", JAN).decode_to_str()) == [ + status or "tenant" for status in occupancy + ] + if branch_salary is not None: + fresh_child = fresh.get_branch("child") + fresh_child.set_input("salary", FEB, [branch_salary] * len(salaries)) + assert_near( + copied.branches["child"].calculate("income_tax", FEB), + fresh_child.calculate("income_tax", FEB), + ) + + copied.set_input("salary", FEB, [1] * len(salaries)) + copied.calculate("disposable_income", FEB) + assert original.persons.get_holder("salary").get_array(FEB) is None + assert original.persons.get_holder("disposable_income").get_array(FEB) is None + for variable in warmed: + assert_near(original.calculate(variable, JAN), fresh.calculate(variable, JAN)) diff --git a/tests/fixtures/simulation_copy.py b/tests/fixtures/simulation_copy.py new file mode 100644 index 00000000..596e4672 --- /dev/null +++ b/tests/fixtures/simulation_copy.py @@ -0,0 +1,81 @@ +"""Helpers for the simulation copy and pickle tests.""" + +from __future__ import annotations + +import copy +import pickle + +from policyengine_core.country_template import CountryTaxBenefitSystem +from policyengine_core.parameters import ParameterNode +from policyengine_core.reforms import Reform +from policyengine_core.simulations import SimulationBuilder + +JAN = "2025-01" +FEB = "2025-02" + +COPIERS = { + "deepcopy": copy.deepcopy, + "pickle": lambda value: pickle.loads(pickle.dumps(value)), +} + + +def build_simulation( + tax_benefit_system=None, salaries=(1000,), households=None, occupancy=None +): + """One household per entry of ``households`` (lists of person indices). + + ``occupancy`` gives each household's ``housing_occupancy_status``. + """ + tax_benefit_system = tax_benefit_system or CountryTaxBenefitSystem() + names = [f"p{i}" for i in range(len(salaries))] + households = households or [list(range(len(salaries)))] + occupancy = occupancy or [None] * len(households) + return SimulationBuilder().build_from_entities( + tax_benefit_system, + { + "persons": { + name: {"salary": {JAN: salary}} for name, salary in zip(names, salaries) + }, + "households": { + f"h{i}": { + "parents": [names[j] for j in members[:2]], + "children": [names[j] for j in members[2:]], + "accommodation_size": {JAN: 100}, + **( + {"housing_occupancy_status": {JAN: occupancy[i]}} + if occupancy[i] + else {} + ), + } + for i, members in enumerate(households) + }, + }, + ) + + +def rate_node() -> ParameterNode: + return ParameterNode( + "rate", + data={ + "single": { + "owner": {"values": {"2015-01-01": 100}}, + "tenant": {"values": {"2015-01-01": 300}}, + }, + "couple": { + "owner": {"values": {"2015-01-01": 500}}, + "tenant": {"values": {"2015-01-01": 700}}, + }, + }, + ) + + +class DoubleIncomeTaxRate(Reform): + """Module level, so pickle can find the reform class by name.""" + + def apply(self): + def double(parameters): + rate = parameters.taxes.income_tax_rate + rate.update(period="year:2015:20", value=rate("2015-01-01") * 2) + return parameters + + self.modify_parameters(double)