diff --git a/changelog.d/disk-storage-branch-ownership.fixed.md b/changelog.d/disk-storage-branch-ownership.fixed.md new file mode 100644 index 000000000..2605202f8 --- /dev/null +++ b/changelog.d/disk-storage-branch-ownership.fixed.md @@ -0,0 +1 @@ +Give independently created temporary holder storages distinct directories, create caller-supplied storage directories on first use, and clear stale enum metadata when replacing enum values with plain arrays. diff --git a/docs/python_api/data_storage.md b/docs/python_api/data_storage.md index 4307934f6..cfaf5ba4d 100644 --- a/docs/python_api/data_storage.md +++ b/docs/python_api/data_storage.md @@ -13,6 +13,17 @@ The `policyengine_core.data_storage` module contains the classes that are used t ## OnDiskStorage +Each independently created temporary holder storage gets its own subdirectory, +so replacing a holder leaves files read by existing clones intact. Cloned +storages retain their source's directory and use new files for later writes. +Explicit `Holder.create_disk_storage(directory=...)` calls and preserved +storages keep the `directory/variable` layout used by simulation dumps. + +`OnDiskStorage.restore()` reads direct key files in its directory. It does not +read files in the `replaced` subdirectory or select the latest write across +independently constructed storages. Clone or copy a restored storage before +writing when its existing values must remain isolated. + ```{eval-rst} .. autoclass:: policyengine_core.data_storage.on_disk_storage.OnDiskStorage :members: @@ -22,6 +33,10 @@ The `policyengine_core.data_storage` module contains the classes that are used t ## TemporaryStorageDirectory +A caller-supplied `simulation._data_storage_dir` is created on first use if it +does not exist. It keeps the caller-owned lifetime described below. +Exits that skip finalizers (such as `os._exit` in multiprocessing workers) can leave uniquely named holder subdirectories in a caller-owned directory to accumulate across runs; the caller must remove them. + ```{eval-rst} .. automodule:: policyengine_core.data_storage.storage_directory :no-members: diff --git a/policyengine_core/data_storage/on_disk_storage.py b/policyengine_core/data_storage/on_disk_storage.py index 3c5dbbb3f..b90bef369 100644 --- a/policyengine_core/data_storage/on_disk_storage.py +++ b/policyengine_core/data_storage/on_disk_storage.py @@ -281,6 +281,8 @@ def put( if isinstance(value, EnumArray): self._enums[path] = value.possible_values value = value.view(numpy.ndarray) + else: + self._enums.pop(path, None) numpy.save(path, value) self._files[filename] = path self._own_paths[filename] = path diff --git a/policyengine_core/data_storage/storage_directory.py b/policyengine_core/data_storage/storage_directory.py index db043adf8..5db1caf37 100644 --- a/policyengine_core/data_storage/storage_directory.py +++ b/policyengine_core/data_storage/storage_directory.py @@ -1,7 +1,7 @@ """The temporary folder a simulation stores values on disk in. -A simulation with a ``MemoryConfig`` stores values on disk, each variable's -in a subfolder (an ``OnDiskStorage``) of the simulation's folder, +A simulation with a ``MemoryConfig`` stores values on disk, each independently +created holder's in its own subfolder (an ``OnDiskStorage``) of the simulation's folder, ``Simulation.data_storage_dir``. Unless the caller chose that folder, the simulation makes it, as a ``TemporaryStorageDirectory``. That object removes the folder, with everything in it but the subfolders of preserved storages @@ -17,7 +17,8 @@ storage is an ``OnDiskStorage.clone`` that reads the source's files and keeps alive the folder and the storage that removes their subfolder. A pickled or copied disk storage does the same. The disk storages a clone - makes itself go in a folder of its own. + makes itself go in a folder of its own. Replacing a holder creates a new + subfolder, so its writes and cleanup leave the old holder's files alone. Folders nest: a simulation given the folder, or a folder in it, as its own (``_data_storage_dir``) keeps the folder alive from then on, and so @@ -29,7 +30,7 @@ is collected first, with one exception, as before: a disk storage removes everything in its subfolder when collected, including files there that another storage reads without keeping it alive. Those are the files of a - second storage made to remove the same subfolder, and anything made + second storage explicitly made to remove the same subfolder, and anything made inside the subfolder. Nor does a storage write over a file another reads: storages cloned or copied from one another write a new file instead (see ``OnDiskStorage._path_to_write``), for the files any of them read back @@ -71,7 +72,7 @@ only the subfolders disk storages made in it, as before, unless it is in a folder a simulation made: that one removes it with everything else in it, once nothing keeps that one alive (above). A clone or forked process makes - its own folder inside it. + its own folder inside it. If it does not exist, it is created on first use. * **A process that ends without running the exit handlers that remove its folders** (killed, or ended with ``os._exit``, as a process ``multiprocessing`` forks is) leaves the folders it made that were still diff --git a/policyengine_core/holders/holder.py b/policyengine_core/holders/holder.py index 324e3d7e9..8e1f868e0 100644 --- a/policyengine_core/holders/holder.py +++ b/policyengine_core/holders/holder.py @@ -1,4 +1,5 @@ import os +import tempfile import warnings from typing import TYPE_CHECKING, Any, List, Tuple @@ -82,11 +83,22 @@ def clone(self, population: "Population", share_arrays: bool = False) -> "Holder def create_disk_storage( self, directory: str = None, preserve: bool = False ) -> OnDiskStorage: - if directory is None: - directory = self.simulation.data_storage_dir - storage_dir = os.path.join(directory, self.variable.name) - if not os.path.isdir(storage_dir): - os.mkdir(storage_dir) + """Create storage for this holder, or a named directory for a dump. + + Independently created temporary storages have distinct directories, + even when a clone still reads an earlier holder for this variable. + Explicit or preserved storages keep the ``directory/variable`` layout. + """ + if directory is None and not preserve: + storage_dir = tempfile.mkdtemp( + prefix=f"{self.variable.name}_", + dir=self.simulation.data_storage_dir, + ) + else: + if directory is None: + directory = self.simulation.data_storage_dir + storage_dir = os.path.join(directory, self.variable.name) + os.makedirs(storage_dir, exist_ok=True) # In the temporary folder the simulation made, the storage (and every # clone or copy of it) keeps the folder until it is collected, so the # folder outlives the simulation while a clone still reads it (see diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index 123a5f1ae..50bd72001 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -769,7 +769,8 @@ def data_storage_dir(self) -> str: Folder in which this simulation stores values on disk when memory is short (see ``MemoryConfig``). - Set ``_data_storage_dir`` to choose the folder: nothing removes it, + Set ``_data_storage_dir`` to choose the folder: it is created on first + use if it does not exist. Nothing removes it, unless it is in a temporary folder a simulation made, which this simulation then keeps alive, and which removes it with everything else in it once nothing keeps it. Otherwise this is a new temporary folder, @@ -801,10 +802,14 @@ def data_storage_dir(self) -> str: if os.path.isdir(inherited): self._storage_dir_parent = inherited if self._data_storage_dir is None: + if self._storage_dir_parent is not None: + os.makedirs(self._storage_dir_parent, exist_ok=True) self._storage_directory = TemporaryStorageDirectory( self._storage_dir_parent ) self._data_storage_dir = self._storage_directory.path + elif not self._made_data_storage_dir(): + os.makedirs(self._data_storage_dir, exist_ok=True) self._storage_dir_pid = pid return self._data_storage_dir diff --git a/tests/core/test_branch_scoped_delete_arrays.py b/tests/core/test_branch_scoped_delete_arrays.py index 5014bc62b..5db33dd90 100644 --- a/tests/core/test_branch_scoped_delete_arrays.py +++ b/tests/core/test_branch_scoped_delete_arrays.py @@ -3,6 +3,7 @@ from __future__ import annotations import gc +import os import numpy as np @@ -82,32 +83,31 @@ def test_delete_arrays_purges_only_visible_branch_names(): ) -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.""" +def test_on_disk_storage_clone_copies_metadata_and_writes_elsewhere(tmp_path): + """A cloned disk view owns its metadata and writes a separate file.""" 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 + clone = storage.clone() + + assert clone is not storage + assert clone.is_eternal == storage.is_eternal + assert clone._files == storage._files + assert clone._files is not storage._files + assert clone._enums == storage._enums + assert clone._enums is not storage._enums + + clone.put(np.asarray([13.0]), PERIOD, "default") + + (clone_file,) = clone._files.values() + (source_file,) = storage._files.values() + assert clone_file != source_file + assert os.path.isfile(clone_file) + np.testing.assert_array_equal(storage.get(PERIOD, "default"), [12.0]) + np.testing.assert_array_equal(clone.get(PERIOD, "default"), [13.0]) def test_on_disk_storage_clone_keeps_cleanup_owner_alive(tmp_path): @@ -145,7 +145,6 @@ def test_child_disk_delete_keeps_parent_view_intact(tmp_path): 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 diff --git a/tests/core/test_data_storage_dir.py b/tests/core/test_data_storage_dir.py index cd2836063..26a5acc44 100644 --- a/tests/core/test_data_storage_dir.py +++ b/tests/core/test_data_storage_dir.py @@ -110,12 +110,15 @@ def _run(script, tmp_path, **environment): def test_a_simulation_storing_on_disk_leaves_no_folder_once_collected(): simulation = disk_simulation() - _fill(simulation) + stored = _fill(simulation) folder = simulation.data_storage_dir - assert _files(folder) == [ - "disk_amount/default_2015.npy", - "disk_level/default_2015.npy", - ] + assert _files(folder) == sorted( + Path(simulation.persons._holders[variable]._disk_storage._files["default_2015"]) + .relative_to(folder) + .as_posix() + for variable in VARIABLES + ) + _assert_reads(simulation, stored) del simulation gc.collect() @@ -663,16 +666,19 @@ def test_a_disk_storage_preserving_its_folder_keeps_it_at_exit(tmp_path): SYSTEM.simulation = simulation simulation.set_input("disk_level", "2015", values("disk_level", 0)) simulation.set_input("disk_amount", "2015", values("disk_amount", 1)) - simulation.persons._holders["disk_amount"]._disk_storage.preserve_storage_dir = True + storage = simulation.persons._holders["disk_amount"]._disk_storage + storage.preserve_storage_dir = True print(simulation.data_storage_dir) + print(storage.storage_dir) """, tmp_path, ) - folder = printed.split()[-1] + folder, storage_dir = printed.splitlines()[-2:] try: - assert _files(folder) == ["disk_amount/default_2015.npy"] + relative = Path(storage_dir).relative_to(folder).as_posix() + assert _files(folder) == [f"{relative}/default_2015.npy"] np.testing.assert_array_equal( - _restored(os.path.join(folder, "disk_amount")).get("2015"), + _restored(storage_dir).get("2015"), values("disk_amount", 1), ) finally: diff --git a/tests/core/test_data_storage_dir_property.py b/tests/core/test_data_storage_dir_property.py index 01256c764..5b0ae15fe 100644 --- a/tests/core/test_data_storage_dir_property.py +++ b/tests/core/test_data_storage_dir_property.py @@ -17,8 +17,7 @@ simulation's folder, preserving it or not (``preserve_storage_dir``); * ``nest``: a new simulation is given a simulation's folder: a new subfolder of it, or the folder itself, in which case only its clone, - made before it stores anything, is kept (two simulations storing a - variable in one folder would share its subfolder). Its clones then make + made before it stores anything, is kept. Its clones then make their folders in a folder another simulation may have made; * ``preserve``: a disk storage a simulation or ``make`` made is set to preserve its folder; diff --git a/tests/core/test_disk_storage_differential.py b/tests/core/test_disk_storage_differential.py new file mode 100644 index 000000000..a19a5fa11 --- /dev/null +++ b/tests/core/test_disk_storage_differential.py @@ -0,0 +1,225 @@ +"""Storing holder values on disk gives the results storing them in memory gives. + +A Hypothesis property over random sequences of branching, cloning, inputs, +calculations, derivatives, reforms, holder replacement and drops, run on a +disk-backed and a memory-backed simulation family side by side. +""" + +from __future__ import annotations + +import gc +import os +from dataclasses import dataclass +from typing import List, Optional + +import numpy as np +import pytest + +# The country-package smoke job installs no dev dependencies. +hypothesis = pytest.importorskip("hypothesis") +st = hypothesis.strategies + +from policyengine_core.simulations import Simulation +from tests.fixtures.disk_storage import build_simulation, unraisable_exceptions + +MONTHS = ["2017-01", "2017-02"] +NAMES = ["a", "b", "c"] +INPUTS = ["salary", "rent", "housing_occupancy_status", "birth"] +CALCULATED = [ + "salary", + "income_tax", + "disposable_income", + "basic_income", + "housing_allowance", + "housing_occupancy_status", + "total_taxes", + "household_income", +] +HOUSING_STATUSES = ["owner", "tenant", "free_lodger", "homeless"] +COUNT = 2 + + +def _input_value(variable: str, seed: int) -> np.ndarray: + values = [seed, seed // 7 + 1] + if variable == "housing_occupancy_status": + return np.array([HOUSING_STATUSES[value % 4] for value in values]) + if variable == "birth": + return np.array( + [f"{1950 + value % 60}-01-01" for value in values], dtype="datetime64[D]" + ) + return np.array(values, dtype=float) + + +@dataclass +class _Member: + simulation: Simulation + parent: Optional["_Member"] + branch_name: Optional[str] + + +class _Family: + """A simulation and every branch or clone made from it, still in use.""" + + def __init__(self, on_disk: bool): + self.members = [_Member(build_simulation(on_disk, count=COUNT), None, None)] + + def pick(self, index: int) -> _Member: + return self.members[index % len(self.members)] + + def apply(self, operation: tuple): + kind, index, *arguments = operation + member = self.pick(index) + simulation = member.simulation + if kind == "branch": + (name,) = arguments + branch = simulation.get_branch(name) + if not any(other.simulation is branch for other in self.members): + self.members.append(_Member(branch, member, name)) + elif kind == "clone": + self.members.append(_Member(simulation.clone(), member, None)) + elif kind == "set": + variable, period, seed = arguments + simulation.set_input(variable, period, _input_value(variable, seed)) + elif kind == "calculate": + variable, period = arguments + return _outcome(lambda: simulation.calculate(variable, period)) + elif kind == "derivative": + (period,) = arguments + return _outcome( + lambda: simulation.derivative("income_tax", "salary", period, delta=100) + ) + elif kind == "reform": + (rate,) = arguments + simulation.apply_reform({"taxes.income_tax_rate": rate}) + elif kind == "delete_arrays": + (variable,) = arguments + simulation.delete_arrays(variable) + elif kind == "new_holders": + for population in simulation.populations.values(): + population._holders = {} + elif kind == "drop" and member.parent is not None: + self.drop(member) + + def drop(self, member: _Member) -> None: + dropped = [member] + for other in self.members: + ancestor = other.parent + while ancestor is not None and ancestor not in dropped: + ancestor = ancestor.parent + if ancestor is not None: + dropped.append(other) + self.members = [other for other in self.members if other not in dropped] + if member.branch_name is not None: + branches = member.parent.simulation.branches + if branches.get(member.branch_name) is member.simulation: + del branches[member.branch_name] + del dropped, member + gc.collect() + + def missing_files(self) -> List[str]: + return [ + path + for member in self.members + for population in member.simulation.populations.values() + for holder in population._holders.values() + if holder._disk_storage is not None + for path in holder._disk_storage._files.values() + if not os.path.isfile(path) + ] + + +def _outcome(calculate): + try: + return np.asarray(calculate()).tolist() + except Exception as error: # Both storages must fail the same way. + return type(error).__name__ + + +_index = st.integers(0, 15) +_operations = st.lists( + st.one_of( + st.tuples(st.just("branch"), _index, st.sampled_from(NAMES)), + st.tuples(st.just("clone"), _index), + st.tuples( + st.just("set"), + _index, + st.sampled_from(INPUTS), + st.sampled_from(MONTHS), + st.integers(0, 5_000), + ), + st.tuples( + st.just("calculate"), + _index, + st.sampled_from(CALCULATED), + st.sampled_from(MONTHS), + ), + st.tuples(st.just("derivative"), _index, st.sampled_from(MONTHS)), + st.tuples(st.just("reform"), _index, st.sampled_from([0.1, 0.2, 0.3])), + st.tuples( + st.just("delete_arrays"), _index, st.sampled_from(INPUTS + CALCULATED) + ), + st.tuples(st.just("new_holders"), _index), + st.tuples(st.just("drop"), _index), + ), + max_size=30, +) + + +@hypothesis.settings( + max_examples=75, + deadline=None, + suppress_health_check=[hypothesis.HealthCheck.too_slow], +) +@hypothesis.example( + operations=[ + ("set", 0, "salary", "2017-01", 1000), + ("clone", 0), + ("new_holders", 0), + ] +) +@hypothesis.example( + operations=[ + ("set", 0, "salary", "2017-01", 1000), + ("clone", 0), + ("new_holders", 0), + ("set", 0, "salary", "2017-01", 2000), + ("calculate", 1, "salary", "2017-01"), + ] +) +@hypothesis.given(operations=_operations) +def test_disk_storage_gives_the_results_memory_storage_gives(operations): + """For any sequence of branching, cloning, inputs, calculations, + derivatives, reforms and drops: + + - every calculation returns what it returns with values held in memory; + - every file a live storage maps is on disk; + - once the simulations are gone, their directories are removed, and no + finalizer raised. + """ + with unraisable_exceptions() as errors: + in_memory = _Family(on_disk=False) + on_disk = _Family(on_disk=True) + for operation in operations: + expected = in_memory.apply(operation) + assert on_disk.apply(operation) == expected, operation + assert on_disk.missing_files() == [], operation + + for member_in_memory, member_on_disk in zip(in_memory.members, on_disk.members): + for variable in CALCULATED: + for period in MONTHS: + assert _outcome( + lambda: member_on_disk.simulation.calculate(variable, period) + ) == _outcome( + lambda: member_in_memory.simulation.calculate(variable, period) + ), (variable, period) + + directories = { + member.simulation._data_storage_dir + for member in on_disk.members + if member.simulation._data_storage_dir is not None + } + del in_memory, on_disk, member_in_memory, member_on_disk + gc.collect() + + assert errors == [] + assert [directory for directory in directories if os.path.exists(directory)] == [] diff --git a/tests/core/test_disk_storage_ownership.py b/tests/core/test_disk_storage_ownership.py new file mode 100644 index 000000000..420dead34 --- /dev/null +++ b/tests/core/test_disk_storage_ownership.py @@ -0,0 +1,517 @@ +"""Disk-backed holder storage keeps the values each simulation reads. + +Each newly created transient holder gets its own directory. Cloned disk +storages share their source's directory and files, keep them alive, and write +new files whenever a clone still reads the old value. These tests pin that for +branches, clones, ``derivative`` and holders created after replacing a holder; +``test_disk_storage_differential.py`` checks it over random operation sequences. +""" + +from __future__ import annotations + +import copy +import gc +import os +import pickle +import shutil +import tempfile + +import numpy as np +import pytest + +from policyengine_core import periods +from policyengine_core.country_template import CountryTaxBenefitSystem, entities +from policyengine_core.data_storage import OnDiskStorage +from policyengine_core.enums import EnumArray +from policyengine_core.experimental import MemoryConfig +from policyengine_core.simulations import Simulation, SimulationBuilder +from policyengine_core.variables import Variable +from tests.fixtures.disk_storage import build_simulation, unraisable_exceptions + +PERIOD = "2017-01" + + +def _storage(tmp_path, name: str = "salary") -> OnDiskStorage: + return OnDiskStorage(tempfile.mkdtemp(prefix=f"{name}_", dir=tmp_path)) + + +def _drop_branch(simulation: Simulation, name: str) -> None: + del simulation.branches[name] + gc.collect() + + +def test_branch_holder_leaves_the_parent_files(): + """A holder a branch creates keeps its hands off the parent's files.""" + simulation = build_simulation() + simulation.set_input("salary", PERIOD, np.array([1_000.0])) + branch = simulation.get_branch("branch") + # The branch creates its own income_tax holder, then the root creates one. + branch.calculate("income_tax", PERIOD) + expected = simulation.calculate("income_tax", PERIOD).copy() + + del branch + _drop_branch(simulation, "branch") + + np.testing.assert_array_equal( + simulation.get_holder("income_tax").get_array(PERIOD), expected + ) + + +def test_clone_input_leaves_the_source_input(): + """A clone writing the same key leaves its source's value readable.""" + simulation = build_simulation() + simulation.set_input("salary", PERIOD, np.array([1_000.0])) + clone = simulation.clone() + + clone.set_input("salary", PERIOD, np.array([5.0])) + + np.testing.assert_array_equal(simulation.calculate("salary", PERIOD), [1_000.0]) + np.testing.assert_array_equal(clone.calculate("salary", PERIOD), [5.0]) + + +@pytest.mark.parametrize("on_disk", [False, True]) +def test_derivative_leaves_the_simulation_inputs(on_disk): + """``derivative`` perturbs a clone's input, not the simulation's own.""" + simulation = build_simulation(on_disk) + simulation.set_input("salary", PERIOD, np.array([3_000.0])) + + derivative = simulation.derivative("income_tax", "salary", PERIOD, delta=100) + + np.testing.assert_allclose(derivative, [0.15]) + np.testing.assert_array_equal(simulation.calculate("salary", PERIOD), [3_000.0]) + + +def test_branches_with_the_same_name_keep_their_own_values(): + simulation = build_simulation() + # Branches made after the root has stored on disk share its directory. + simulation.set_input("salary", PERIOD, np.array([1_000.0])) + first = simulation.get_branch("first").get_branch("nested") + second = simulation.get_branch("second").get_branch("nested") + + first.set_input("salary", PERIOD, np.array([1.0])) + second.set_input("salary", PERIOD, np.array([2.0])) + + np.testing.assert_array_equal(first.calculate("salary", PERIOD), [1.0]) + np.testing.assert_array_equal(second.calculate("salary", PERIOD), [2.0]) + + +@pytest.mark.parametrize("on_disk", [False, True]) +def test_branch_keeps_the_values_it_started_with(on_disk): + """A parent replacing a value after branching leaves the branch's value.""" + simulation = build_simulation(on_disk) + simulation.set_input("salary", PERIOD, np.array([1_000.0])) + branch = simulation.get_branch("branch") + + simulation.delete_arrays("salary", PERIOD) + simulation.set_input("salary", PERIOD, np.array([7.0])) + + np.testing.assert_array_equal(branch.calculate("salary", PERIOD), [1_000.0]) + np.testing.assert_array_equal(simulation.calculate("salary", PERIOD), [7.0]) + + +def test_simulation_stores_on_disk_after_its_branch_storage_is_removed(): + simulation = build_simulation() + # A branch creating a new holder gets a directory of its own. + simulation.data_storage_dir + branch = simulation.get_branch("branch") + branch.set_input("salary", PERIOD, np.array([1.0])) + del branch + _drop_branch(simulation, "branch") + + simulation.set_input("salary", PERIOD, np.array([2.0])) + + np.testing.assert_array_equal(simulation.calculate("salary", PERIOD), [2.0]) + + +def test_directory_of_a_holder_created_by_a_branch_goes_with_the_branch(): + simulation = build_simulation() + simulation.set_input("salary", PERIOD, np.array([1_000.0])) + branch = simulation.get_branch("branch") + # A copied salary storage shares the source directory, while this new + # holder uses the branch's own lazily created simulation directory. + branch.calculate("income_tax", PERIOD) + root_dir = simulation.get_holder("salary")._disk_storage.storage_dir + branch_dir = branch.get_holder("income_tax")._disk_storage.storage_dir + assert branch_dir != root_dir + + del branch + _drop_branch(simulation, "branch") + + assert not os.path.exists(branch_dir) + np.testing.assert_array_equal(simulation.calculate("salary", PERIOD), [1_000.0]) + + +def test_releasing_a_family_removes_its_directory_without_errors(): + with unraisable_exceptions() as errors: + simulation = build_simulation() + simulation.set_input("salary", PERIOD, np.array([1_000.0])) + branch = simulation.get_branch("branch") + branch.calculate("income_tax", PERIOD) + simulation.calculate("income_tax", PERIOD) + directory = simulation.data_storage_dir + assert os.path.isdir(directory) + + del branch + _drop_branch(simulation, "branch") + del simulation + gc.collect() + + assert errors == [] + assert not os.path.exists(directory) + + +@pytest.mark.parametrize("clone_first", [False, True], ids=["root", "clone"]) +def test_a_data_storage_dir_set_by_the_user_stays(tmp_path, clone_first): + directory = tmp_path / "simulation" + simulation = build_simulation() + simulation._data_storage_dir = str(directory) + storer = simulation.clone() if clone_first else simulation + assert not directory.exists() + + # A chosen folder, or a clone's chosen parent, may not exist at first use. + data_storage_dir = storer.data_storage_dir + assert directory.is_dir() + assert os.path.isdir(data_storage_dir) + if clone_first: + assert os.path.dirname(data_storage_dir) == str(directory) + else: + assert data_storage_dir == str(directory) + storer.set_input("salary", PERIOD, np.array([1.0])) + storage_dir = storer.get_holder("salary")._disk_storage.storage_dir + assert os.path.dirname(storage_dir) == data_storage_dir + np.testing.assert_array_equal(storer.calculate("salary", PERIOD), [1.0]) + + del storer, simulation + gc.collect() + + assert not os.path.exists(storage_dir) + assert directory.is_dir() + + +@pytest.mark.parametrize("preserve", [False, True]) +def test_explicit_holder_storage_keeps_the_variable_directory_layout( + tmp_path, preserve +): + simulation = build_simulation(on_disk=False) + directory = tmp_path / "dump" + storage = simulation.get_holder("salary").create_disk_storage( + directory=str(directory), preserve=preserve + ) + storage.put(np.array([1.0]), PERIOD) + + assert storage.storage_dir == str(directory / "salary") + np.testing.assert_array_equal( + np.load(directory / "salary" / f"default_{PERIOD}.npy"), [1.0] + ) + + +def test_preserved_holder_storage_keeps_the_variable_directory_layout(): + simulation = build_simulation(on_disk=False) + storage = simulation.get_holder("salary").create_disk_storage(preserve=True) + storage.put(np.array([1.0]), PERIOD) + directory = simulation.data_storage_dir + + try: + assert storage.storage_dir == os.path.join(directory, "salary") + np.testing.assert_array_equal( + np.load(os.path.join(directory, "salary", f"default_{PERIOD}.npy")), + [1.0], + ) + finally: + shutil.rmtree(directory, ignore_errors=True) + + +def test_a_key_stored_again_overwrites_its_file_until_a_clone_reads_it(tmp_path): + storage = _storage(tmp_path) + storage.put(np.array([1.0]), PERIOD) + first_path = storage._files[f"default_{PERIOD}"] + storage.put(np.array([2.0]), PERIOD) + assert storage._files[f"default_{PERIOD}"] == first_path + + clone = storage.clone() + storage.put(np.array([3.0]), PERIOD) + replacement_path = storage._files[f"default_{PERIOD}"] + storage.put(np.array([4.0]), PERIOD) + + assert replacement_path != first_path + assert storage._files[f"default_{PERIOD}"] == replacement_path + np.testing.assert_array_equal(clone.get(PERIOD), [2.0]) + np.testing.assert_array_equal(storage.get(PERIOD), [4.0]) + + del clone + gc.collect() + storage.put(np.array([5.0]), PERIOD) + + assert storage._files[f"default_{PERIOD}"] == replacement_path + np.testing.assert_array_equal(storage.get(PERIOD), [5.0]) + + +def test_reforms_after_a_derivative_reuse_the_files(): + """Repeated reforms reuse the storage's first unshared replacement.""" + simulation = build_simulation() + simulation.set_input("salary", PERIOD, np.array([3_000.0])) + simulation.calculate("income_tax", PERIOD) + simulation.derivative("income_tax", "salary", PERIOD, delta=100) + gc.collect() + storage = simulation.get_holder("income_tax")._disk_storage + + replacement_path = None + for rate in (0.1, 0.2, 0.3, 0.4): + simulation.apply_reform({"taxes.income_tax_rate": rate}) + np.testing.assert_allclose( + simulation.calculate("income_tax", PERIOD), [3_000.0 * rate] + ) + path = storage._files[f"default_{PERIOD}"] + if replacement_path is None: + replacement_path = path + else: + assert path == replacement_path + + +def test_a_live_branch_adds_at_most_one_file_per_key(): + simulation = build_simulation() + simulation.set_input("salary", PERIOD, np.array([1.0])) + branch = simulation.get_branch("branch") + storage = simulation.get_holder("salary")._disk_storage + source_path = storage._files[f"default_{PERIOD}"] + + replacement_paths = set() + for value in (2.0, 3.0, 4.0): + simulation.delete_arrays("salary", PERIOD) + simulation.set_input("salary", PERIOD, np.array([value])) + replacement_paths.add(storage._files[f"default_{PERIOD}"]) + + assert len(replacement_paths) == 1 + assert source_path not in replacement_paths + np.testing.assert_array_equal(branch.calculate("salary", PERIOD), [1.0]) + np.testing.assert_array_equal(simulation.calculate("salary", PERIOD), [4.0]) + + +def test_derivative_of_a_carried_over_input_leaves_it(): + """A year input carried over to later years is not perturbed by + ``derivative`` (from the carry-over-order review's repro).""" + + class carried(Variable): + value_type = float + entity = entities.Person + definition_period = periods.YEAR + label = "Carried-over input" + + tax_benefit_system = CountryTaxBenefitSystem() + tax_benefit_system.auto_carry_over_input_variables = True + simulation = SimulationBuilder().build_default_simulation( + tax_benefit_system, count=1 + ) + simulation.memory_config = MemoryConfig(max_memory_occupation=0) + # Its holder is created after memory_config, so it stores on disk. + tax_benefit_system.add_variables(carried) + simulation.set_input("carried", "2012", [100]) + + derivative = simulation.derivative("carried", "carried", "2012", delta=1) + + np.testing.assert_array_equal(derivative, [1.0]) + np.testing.assert_array_equal(simulation.get_array("carried", "2012"), [100.0]) + np.testing.assert_array_equal(simulation.calculate("carried", "2013"), [100.0]) + + +def test_a_clone_keeps_the_directories_it_reads(tmp_path): + storage = _storage(tmp_path) + storage.put(np.array([1.0]), PERIOD) + storage_dir = storage.storage_dir + grandchild = storage.clone().clone() + + del storage + gc.collect() + + assert os.path.isdir(storage_dir) + np.testing.assert_array_equal(grandchild.get(PERIOD), [1.0]) + + del grandchild + gc.collect() + + assert not os.path.exists(storage_dir) + + +def test_restore_reads_direct_key_files_after_a_clone_shares_them(tmp_path): + """Restore excludes replacements, as documented by ``OnDiskStorage``.""" + storage_dir = tmp_path / "salary" + storage_dir.mkdir() + storage = OnDiskStorage(str(storage_dir), preserve_storage_dir=True) + storage.put(np.array([1.0]), PERIOD) + storage.put(np.array([5.0]), PERIOD, "a.1") + clone = storage.clone() + storage.put(np.array([2.0]), PERIOD) + storage.put(np.array([3.0]), "2017-02") + + restored = OnDiskStorage(str(storage_dir), preserve_storage_dir=True) + restored.restore() + + assert sorted(restored._files) == [ + "a.1_2017-01", + "default_2017-01", + "default_2017-02", + ] + np.testing.assert_array_equal(restored.get(PERIOD), [1.0]) + np.testing.assert_array_equal(restored.get(PERIOD, "a.1"), [5.0]) + np.testing.assert_array_equal(restored.get("2017-02"), [3.0]) + np.testing.assert_array_equal(clone.get(PERIOD), [1.0]) + np.testing.assert_array_equal(storage.get(PERIOD), [2.0]) + + +def test_storing_after_restore_updates_the_direct_key_file(tmp_path): + """Without cloning, restore continues writing the direct key file.""" + storage_dir = tmp_path / "salary" + storage_dir.mkdir() + writer = OnDiskStorage(str(storage_dir), preserve_storage_dir=True) + writer.put(np.array([1.0]), PERIOD) + first = OnDiskStorage(str(storage_dir), preserve_storage_dir=True) + first.restore() + + first.put(np.array([2.0]), PERIOD) + + np.testing.assert_array_equal(first.get(PERIOD), [2.0]) + latest = OnDiskStorage(str(storage_dir), preserve_storage_dir=True) + latest.restore() + np.testing.assert_array_equal(latest.get(PERIOD), [2.0]) + + +def test_a_write_after_restore_leaves_the_restored_storage_clone(tmp_path): + storage_dir = tmp_path / "salary" + storage_dir.mkdir() + writer = OnDiskStorage(str(storage_dir), preserve_storage_dir=True) + writer.put(np.array([1.0]), PERIOD) + restorer = OnDiskStorage(str(storage_dir), preserve_storage_dir=True) + restorer.restore() + clone = restorer.clone() + + restorer.put(np.array([3.0]), PERIOD) + + np.testing.assert_array_equal(clone.get(PERIOD), [1.0]) + np.testing.assert_array_equal(writer.get(PERIOD), [1.0]) + np.testing.assert_array_equal(restorer.get(PERIOD), [3.0]) + latest = OnDiskStorage(str(storage_dir), preserve_storage_dir=True) + latest.restore() + np.testing.assert_array_equal(latest.get(PERIOD), [1.0]) + + +def test_two_restored_storage_families_keep_their_own_values(tmp_path): + storage_dir = tmp_path / "salary" + storage_dir.mkdir() + OnDiskStorage(str(storage_dir), preserve_storage_dir=True).put( + np.array([1.0]), PERIOD + ) + first = OnDiskStorage(str(storage_dir), preserve_storage_dir=True) + first.restore() + second = OnDiskStorage(str(storage_dir), preserve_storage_dir=True) + second.restore() + # Each family protects the file once a clone may read it. Independently + # restored storages alone have no cross-family snapshot guarantee. + first_clone = first.clone() + second_clone = second.clone() + + first.put(np.array([2.0]), PERIOD) + second.put(np.array([3.0]), PERIOD) + + np.testing.assert_array_equal(first.get(PERIOD), [2.0]) + np.testing.assert_array_equal(second.get(PERIOD), [3.0]) + np.testing.assert_array_equal(first_clone.get(PERIOD), [1.0]) + np.testing.assert_array_equal(second_clone.get(PERIOD), [1.0]) + latest = OnDiskStorage(str(storage_dir), preserve_storage_dir=True) + latest.restore() + np.testing.assert_array_equal(latest.get(PERIOD), [1.0]) + + +def test_storage_removes_only_its_own_directory(tmp_path): + storage_dir = tmp_path / "parent" / "salary" + storage_dir.mkdir(parents=True) + storage = OnDiskStorage(str(storage_dir)) + storage.put(np.array([1.0]), PERIOD) + + del storage + gc.collect() + + assert not storage_dir.exists() + assert (tmp_path / "parent").is_dir() + + +@pytest.mark.parametrize("preserve_when", ["created", "later"]) +def test_preserved_storage_directory_stays(tmp_path, preserve_when): + storage_dir = tmp_path / "salary" + storage_dir.mkdir() + storage = OnDiskStorage( + str(storage_dir), preserve_storage_dir=preserve_when == "created" + ) + storage.preserve_storage_dir = True + storage.put(np.array([1.0]), PERIOD) + + del storage + gc.collect() + + assert (storage_dir / f"default_{PERIOD}.npy").is_file() + + +def test_a_preserved_storage_in_the_simulation_directory_stays(): + simulation = build_simulation(on_disk=False) + storage = simulation.get_holder("salary").create_disk_storage(preserve=True) + storage.put(np.array([1.0]), PERIOD) + data_storage_dir = simulation.data_storage_dir + + del simulation + gc.collect() + + try: + np.testing.assert_array_equal(storage.get(PERIOD), [1.0]) + finally: + shutil.rmtree(data_storage_dir, ignore_errors=True) + + +def test_preserving_a_holder_storage_keeps_it_after_the_family_goes(): + simulation = build_simulation() + simulation.set_input("salary", PERIOD, np.array([1.0])) + storage = simulation.get_holder("salary")._disk_storage + storage.preserve_storage_dir = True + data_storage_dir, storage_dir = simulation.data_storage_dir, storage.storage_dir + path = storage._files[f"default_{PERIOD}"] + + del simulation, storage + gc.collect() + + try: + assert os.path.isdir(storage_dir) + np.testing.assert_array_equal(np.load(path), [1.0]) + finally: + shutil.rmtree(data_storage_dir, ignore_errors=True) + + +@pytest.mark.parametrize("copy_storage", [copy.copy, copy.deepcopy, "pickle"]) +def test_a_copy_of_a_storage_is_independent_of_it(tmp_path, copy_storage): + """A copy reads the original's files, keeps them, and writes its own.""" + storage = _storage(tmp_path) + storage.put(np.array([1.0]), PERIOD) + if copy_storage == "pickle": + copied = pickle.loads(pickle.dumps(storage)) + else: + copied = copy_storage(storage) + + copied.put(np.array([2.0]), PERIOD) + storage.put(np.array([3.0]), PERIOD) + np.testing.assert_array_equal(copied.get(PERIOD), [2.0]) + np.testing.assert_array_equal(storage.get(PERIOD), [3.0]) + + storage_dir = storage.storage_dir + del storage + gc.collect() + + assert os.path.isdir(storage_dir) + np.testing.assert_array_equal(copied.get(PERIOD), [2.0]) + + +def test_a_plain_array_stored_over_an_enum_reads_back_plain(tmp_path): + storage = _storage(tmp_path, "status") + status = CountryTaxBenefitSystem().variables["housing_occupancy_status"] + storage.put(status.possible_values.encode(np.array(["owner"])), PERIOD) + assert isinstance(storage.get(PERIOD), EnumArray) + + storage.put(np.array([3]), PERIOD) + + assert not isinstance(storage.get(PERIOD), EnumArray) diff --git a/tests/fixtures/disk_storage.py b/tests/fixtures/disk_storage.py new file mode 100644 index 000000000..1034baed9 --- /dev/null +++ b/tests/fixtures/disk_storage.py @@ -0,0 +1,36 @@ +"""Helpers for tests of disk-backed holder storage.""" + +import contextlib +import sys +from typing import List + +from policyengine_core.country_template import CountryTaxBenefitSystem +from policyengine_core.experimental import MemoryConfig +from policyengine_core.simulations import Simulation, SimulationBuilder + + +def build_simulation(on_disk: bool = True, count: int = 1) -> Simulation: + # A tax-benefit system keeps the last simulation built on it, so give + # each simulation its own for the tests that release one. + simulation = SimulationBuilder().build_default_simulation( + CountryTaxBenefitSystem(), count=count + ) + if on_disk: + simulation.memory_config = MemoryConfig(max_memory_occupation=0) + # Holders made before ``memory_config`` is set have no disk storage, so + # start without any: each holder is created where it is first used. + for population in simulation.populations.values(): + population._holders = {} + return simulation + + +@contextlib.contextmanager +def unraisable_exceptions(): + """Collect exceptions Python can only print, such as those in finalizers.""" + errors: List[str] = [] + previous_hook = sys.unraisablehook + sys.unraisablehook = lambda unraisable: errors.append(repr(unraisable.exc_value)) + try: + yield errors + finally: + sys.unraisablehook = previous_hook