diff --git a/changelog.d/branch-scoped-delete-arrays.fixed.md b/changelog.d/branch-scoped-delete-arrays.fixed.md new file mode 100644 index 00000000..93a27e76 --- /dev/null +++ b/changelog.d/branch-scoped-delete-arrays.fixed.md @@ -0,0 +1 @@ +Fixed nested simulation cache deletion so perturbation branches recompute inherited values without mutating parent storage. diff --git a/policyengine_core/data_storage/on_disk_storage.py b/policyengine_core/data_storage/on_disk_storage.py index ca4ee102..3563c805 100644 --- a/policyengine_core/data_storage/on_disk_storage.py +++ b/policyengine_core/data_storage/on_disk_storage.py @@ -26,6 +26,26 @@ def __init__( self.preserve_storage_dir = preserve_storage_dir self.storage_dir = storage_dir + def clone(self) -> "OnDiskStorage": + """Create a private metadata view over this storage directory. + + The file and enum mappings are copied so deleting or rewiring entries + through the clone does not mutate the source storage. The underlying + ``.npy`` files remain shared: writing the same ``{branch}_{period}`` + key from two views targets the same path and can overwrite the file. + Clones retain the original cleanup owner so the shared directory stays + alive, but never own cleanup themselves. + """ + clone = OnDiskStorage( + self.storage_dir, + is_eternal=self.is_eternal, + preserve_storage_dir=True, + ) + clone._files = self._files.copy() + clone._enums = self._enums.copy() + clone._storage_dir_owner = getattr(self, "_storage_dir_owner", self) + return clone + def _decode_file(self, file: str) -> ArrayLike: enum = self._enums.get(file) if enum is not None: diff --git a/policyengine_core/holders/holder.py b/policyengine_core/holders/holder.py index f6be8d04..52ae8129 100644 --- a/policyengine_core/holders/holder.py +++ b/policyengine_core/holders/holder.py @@ -57,10 +57,14 @@ def clone(self, population: "Population") -> "Holder": "formula", "simulation", "_memory_storage", + "_disk_storage", ): new_dict[key] = value new._memory_storage = self._memory_storage.clone() + new._disk_storage = ( + self._disk_storage.clone() if self._disk_storage is not None else None + ) new_dict["population"] = population new_dict["simulation"] = population.simulation diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index b7f52591..2f9fd75c 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -1261,10 +1261,14 @@ def get_memory_usage(self, variables: List[str] = None) -> dict: def delete_arrays(self, variable: str, period: Period = None) -> None: """ - Delete a variable's value for a given period + Delete a variable's values visible to this simulation branch. - :param variable: the variable to be set - :param period: the period for which the value should be deleted + The calling branch, each ancestor branch, and the default branch are + purged from this simulation's private holder storage. Other branch + names and the parent simulation's holder storage remain unchanged. + + :param variable: the variable whose cached values should be deleted + :param period: the period to delete, or all periods when omitted Example: @@ -1286,7 +1290,9 @@ def delete_arrays(self, variable: str, period: Period = None) -> None: >>> simulation.get_array('age', '2018-05') is None True """ - self.get_holder(variable).delete_arrays(period) + holder = self.get_holder(variable) + for branch_name in self._get_visible_branch_names(): + holder.delete_arrays(period, branch_name) _fast_cache = getattr(self, "_fast_cache", None) if period is None: if _fast_cache is not None: @@ -1631,7 +1637,7 @@ def _is_exportable_input_variable(self, variable_name: str) -> bool: return variable is not None and variable.is_input_variable() def _get_visible_branch_names(self) -> List[str]: - branch_names = [self.branch_name] + branch_names = [getattr(self, "branch_name", "default")] parent = getattr(self, "parent_branch", None) while parent is not None: branch_names.append(parent.branch_name) diff --git a/tests/core/test_branch_scoped_delete_arrays.py b/tests/core/test_branch_scoped_delete_arrays.py new file mode 100644 index 00000000..5014bc62 --- /dev/null +++ b/tests/core/test_branch_scoped_delete_arrays.py @@ -0,0 +1,176 @@ +"""Regression tests for deleting caches from nested simulation branches.""" + +from __future__ import annotations + +import gc + +import numpy as np + +from policyengine_core.country_template import CountryTaxBenefitSystem +from policyengine_core.data_storage import OnDiskStorage +from policyengine_core.simulations import SimulationBuilder + + +PERIOD = "2017-01" + + +def _salary_simulation(): + return SimulationBuilder().build_from_entities( + CountryTaxBenefitSystem(), + { + "persons": {"bill": {"salary": {PERIOD: 3_000}}}, + "households": {"household": {"parents": ["bill"]}}, + }, + ) + + +def _purge_formula_outputs(simulation): + for variable_name in simulation.tax_benefit_system.variables: + if variable_name not in simulation.input_variables: + simulation.delete_arrays(variable_name) + + +def test_nested_branch_recomputes_after_formula_output_purge(): + """A perturbed nested branch must not reuse its parent's formula output.""" + simulation = _salary_simulation() + measurement = simulation.get_branch("measurement") + parent_value = measurement.calculate("income_tax", PERIOD).copy() + perturbed = measurement.get_branch("perturbed") + + _purge_formula_outputs(perturbed) + perturbed.set_input("salary", PERIOD, np.asarray([6_000.0])) + recomputed_value = perturbed.calculate("income_tax", PERIOD) + + np.testing.assert_allclose(parent_value, [450.0], rtol=1e-6) + np.testing.assert_allclose(recomputed_value, [900.0], rtol=1e-6) + np.testing.assert_allclose( + measurement.calculate("income_tax", PERIOD), + parent_value, + ) + + +def test_delete_arrays_purges_only_visible_branch_names(): + """Root and nested purges must preserve unrelated dispatched values.""" + simulation = _salary_simulation() + root_holder = simulation.get_holder("salary") + dispatched_value = np.asarray([7_777.0]) + root_holder._memory_storage.put(dispatched_value, PERIOD, "dispatched") + + simulation.delete_arrays("salary", PERIOD) + + assert root_holder._memory_storage.get(PERIOD, "default") is None + np.testing.assert_array_equal( + root_holder._memory_storage.get(PERIOD, "dispatched"), + dispatched_value, + ) + + measurement = simulation.get_branch("measurement") + perturbed = measurement.get_branch("perturbed") + child_holder = perturbed.get_holder("salary") + child_holder._memory_storage.put(np.asarray([1.0]), PERIOD, "default") + child_holder._memory_storage.put(np.asarray([2.0]), PERIOD, "measurement") + child_holder._memory_storage.put(np.asarray([3.0]), PERIOD, "perturbed") + + perturbed.delete_arrays("salary", PERIOD) + + assert child_holder._memory_storage.get(PERIOD, "default") is None + assert child_holder._memory_storage.get(PERIOD, "measurement") is None + assert child_holder._memory_storage.get(PERIOD, "perturbed") is None + np.testing.assert_array_equal( + child_holder._memory_storage.get(PERIOD, "dispatched"), + dispatched_value, + ) + + +def test_on_disk_storage_clone_copies_metadata_without_owning_directory(tmp_path): + """A cloned disk view must own metadata, but not the shared directory.""" + storage_dir = tmp_path / "storage" + storage_dir.mkdir() + storage = OnDiskStorage(str(storage_dir), is_eternal=True) + storage.put(np.asarray([12.0]), PERIOD, "default") + storage._enums["sentinel"] = "enum" + + try: + assert storage.preserve_storage_dir is False + clone = storage.clone() + + assert clone is not storage + assert clone.storage_dir == storage.storage_dir + assert clone.is_eternal == storage.is_eternal + assert clone.preserve_storage_dir is True + assert clone._files == storage._files + assert clone._files is not storage._files + assert clone._enums == storage._enums + assert clone._enums is not storage._enums + assert clone._storage_dir_owner is storage + assert clone.clone()._storage_dir_owner is storage + finally: + # Let pytest's tmp_path cleanup remove the directory even if an + # assertion fails before the clone can prove its ownership setting. + storage.preserve_storage_dir = True + + +def test_on_disk_storage_clone_keeps_cleanup_owner_alive(tmp_path): + """A detached clone must remain readable after its source is released.""" + storage_dir = tmp_path / "storage" + storage_dir.mkdir() + storage = OnDiskStorage(str(storage_dir)) + storage.put(np.asarray([12.0]), PERIOD, "default") + clone = storage.clone() + + del storage + gc.collect() + + assert storage_dir.is_dir() + np.testing.assert_array_equal(clone.get(PERIOD, "default"), [12.0]) + + +def test_child_disk_delete_keeps_parent_view_intact(tmp_path): + """Deleting through a child must rewire only the child's disk mappings.""" + simulation = _salary_simulation() + parent_holder = simulation.get_holder("salary") + parent_holder._memory_storage.delete(PERIOD, "default") + parent_holder._disk_storage = parent_holder.create_disk_storage(str(tmp_path)) + parent_storage = parent_holder._disk_storage + parent_storage.put(np.asarray([3_000.0]), PERIOD, "default") + disk_key = f"default_{PERIOD}" + + try: + child = simulation.get_branch("measurement") + child_holder = child.get_holder("salary") + child_storage = child_holder._disk_storage + assert child_storage._files[disk_key] == parent_storage._files[disk_key] + inherited_value = child_holder.get_array(PERIOD, child.branch_name) + + child.delete_arrays("salary", PERIOD) + + assert child_storage is not parent_storage + assert child_storage.preserve_storage_dir is True + np.testing.assert_array_equal(inherited_value, [3_000.0]) + assert disk_key in parent_storage._files + assert disk_key not in child_storage._files + assert (tmp_path / "salary" / f"{disk_key}.npy").is_file() + np.testing.assert_array_equal( + parent_holder._disk_storage.get(PERIOD, "default"), + [3_000.0], + ) + finally: + # The path belongs to pytest for this test; avoid a later destructor + # racing tmp_path cleanup if the assertion above fails. + parent_storage.preserve_storage_dir = True + + +def test_derivative_recomputes_inside_named_branch(): + """A derivative clone must purge formula outputs visible to its branch.""" + simulation = _salary_simulation() + measurement = simulation.get_branch("measurement") + measurement.calculate("income_tax", PERIOD) + + derivative = measurement.derivative( + "income_tax", + "salary", + PERIOD, + delta=1, + ) + + np.testing.assert_allclose(derivative, [0.15], rtol=1e-4, atol=1e-6) diff --git a/tests/core/test_fast_cache_guards.py b/tests/core/test_fast_cache_guards.py index 62762fca..77dfcaba 100644 --- a/tests/core/test_fast_cache_guards.py +++ b/tests/core/test_fast_cache_guards.py @@ -60,7 +60,7 @@ def test_set_input_without_fast_cache_attribute(): def test_delete_arrays_without_fast_cache_attribute(): sim = _bare_simulation() sim.get_holder = lambda name: types.SimpleNamespace( - delete_arrays=lambda period: None + delete_arrays=lambda period, branch_name: None ) # No _fast_cache attribute — must not crash sim.delete_arrays("variable", period=None)