From 7a3867aae8a8255385ed5bb3a831b44f225f15b5 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Thu, 1 Oct 2026 15:48:05 -0400 Subject: [PATCH 1/9] Give each disk-backed holder storage its own directory A holder created lazily in a branch made an OnDiskStorage that owned the family's shared / directory, so garbage-collecting the branch deleted the parent's files. The shared directory also let a Simulation.clone() or a same-named branch under another parent overwrite the parent's .npy files, and let a parent's later write change what its branch read. Each storage now writes only into a directory of its own (created on first write inside the simulation's data_storage_dir), a clone keeps alive every directory it reads from, and a file a clone reads is never overwritten. Directories are removed by finalizers on reference-counted StorageDirectory handles, only in the process that created them. Co-Authored-By: Claude Opus 5.5 --- .gitignore | 3 +- .../disk-storage-branch-ownership.fixed.md | 1 + .../data_storage/on_disk_storage.py | 222 ++++++-- policyengine_core/holders/holder.py | 19 +- policyengine_core/simulations/simulation.py | 23 +- pyproject.toml | 1 + .../core/test_branch_scoped_delete_arrays.py | 40 +- tests/core/test_disk_storage_ownership.py | 478 ++++++++++++++++++ uv.lock | 96 +++- 9 files changed, 817 insertions(+), 66 deletions(-) create mode 100644 changelog.d/disk-storage-branch-ownership.fixed.md create mode 100644 tests/core/test_disk_storage_ownership.py diff --git a/.gitignore b/.gitignore index b98f81448..6c46608d2 100644 --- a/.gitignore +++ b/.gitignore @@ -13,4 +13,5 @@ docs/_build/ .venv .vscode venv -.uv/ \ No newline at end of file +.uv/ +.hypothesis/ 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..b36b5e9fb --- /dev/null +++ b/changelog.d/disk-storage-branch-ownership.fixed.md @@ -0,0 +1 @@ +Disk-backed holder storage (`MemoryConfig`) now gives every storage its own directory, keeps a directory while any clone reads its files and never overwrites a file a clone reads, so garbage-collecting a branch no longer deletes its parent's stored values and clones or same-named branches no longer overwrite each other's values. diff --git a/policyengine_core/data_storage/on_disk_storage.py b/policyengine_core/data_storage/on_disk_storage.py index 3563c805d..3298ef04d 100644 --- a/policyengine_core/data_storage/on_disk_storage.py +++ b/policyengine_core/data_storage/on_disk_storage.py @@ -1,5 +1,9 @@ +import itertools import os import shutil +import tempfile +import weakref +from typing import Dict, Optional, Set, Tuple, Union import numpy from numpy.typing import ArrayLike @@ -9,41 +13,170 @@ from policyengine_core.periods import Period +def _remove_directory(path: str, pid: int) -> None: + # A forked child inherits the finalizer, but the directory belongs to the + # process that created it. + if os.getpid() == pid: + shutil.rmtree(path, ignore_errors=True) + + +class StorageDirectory: + """A directory removed, with its contents, once nothing references it. + + Whatever may read a file in the directory holds a reference to this + object, so the directory stays on disk exactly as long as one of them is + alive. A directory created inside another holds that one as ``parent``. + """ + + def __init__(self, path: str, parent: Optional["StorageDirectory"] = None): + self.path = path + self.parent = parent + self._pid = os.getpid() + self._finalizer = weakref.finalize(self, _remove_directory, path, self._pid) + + @classmethod + def create( + cls, + prefix: str, + parent: Union["StorageDirectory", str, None] = None, + ) -> "StorageDirectory": + """Create a new, uniquely named directory inside ``parent``. + + ``parent`` is another :class:`StorageDirectory`, a path, or ``None`` + for the system temporary directory. + """ + if isinstance(parent, StorageDirectory): + parent_path = parent.path + else: + parent_path = parent + parent = None + if parent_path is not None: + os.makedirs(parent_path, exist_ok=True) + return cls(tempfile.mkdtemp(prefix=prefix, dir=parent_path), parent=parent) + + @property + def preserve(self) -> bool: + """Whether the directory stays on disk after nothing references it.""" + return not self._finalizer.alive + + @preserve.setter + def preserve(self, preserve: bool) -> None: + if preserve: + self._finalizer.detach() + elif not self._finalizer.alive: + self._finalizer = weakref.finalize( + self, _remove_directory, self.path, self._pid + ) + + +def _parse_file_name(file_name: str) -> Tuple[str, int]: + """Split a stored file's name, without ``.npy``, into its key and version. + + A key's first file is ``{key}.npy`` and later ones ``{key}.{n}.npy``. A + key is ``{branch_name}_{period}`` and a period's string form contains no + ``.``, so a final ``.`` followed only by digits is a version. + """ + key, separator, version = file_name.rpartition(".") + if separator and version.isdigit(): + return key, int(version) + return file_name, 0 + + class OnDiskStorage: """ - Low-level class responsible for storing and retrieving calculated vectors on disk + Low-level class responsible for storing and retrieving calculated vectors on disk. + + Each storage writes only into its own directory, ``storage_dir``. A + storage created with ``storage_dir=None`` makes that directory on its + first write, inside the directory given by :meth:`temporary` (or the + system temporary directory). + + A clone (:meth:`clone`) reads the files its source had stored when it was + cloned, and writes into a new directory of its own, so neither storage + can change or remove a value the other reads. Once a file has been + shared with a clone, storing its key again writes a new file instead of + overwriting it. + + Unless ``preserve_storage_dir`` is set, a storage's directory is removed + once the storage and every clone that may read a file in it are gone. """ def __init__( self, - storage_dir: str, + storage_dir: Optional[str], is_eternal: bool = False, preserve_storage_dir: bool = False, ): - self._files = {} - self._enums = {} + self._files: Dict[str, str] = {} + self._enums: Dict[str, type] = {} self.is_eternal = is_eternal - self.preserve_storage_dir = preserve_storage_dir self.storage_dir = storage_dir + self._preserve_storage_dir = preserve_storage_dir + self._directory: Optional[StorageDirectory] = None + self._name = "storage" + if storage_dir is not None: + self._directory = StorageDirectory(storage_dir) + self._directory.preserve = preserve_storage_dir + self._name = os.path.basename(os.path.normpath(storage_dir)) + # Where ``storage_dir`` is created on the first write, if it is None. + self._parent_directory: Union[StorageDirectory, str, None] = None + # Other storages' directories that files in ``_files`` may be in. + self._read_directories: Tuple[StorageDirectory, ...] = () + # Files this storage wrote or restored, and those of them a clone + # may read, which must not be overwritten. + self._written_files: Set[str] = set() + self._shared_files: Set[str] = set() + self._versions = itertools.count(1) + + @classmethod + def temporary( + cls, + name: str, + directory: Union[StorageDirectory, str, None] = None, + is_eternal: bool = False, + ) -> "OnDiskStorage": + """Create a storage whose directory is made on its first write. + + The directory is created inside ``directory`` (a + :class:`StorageDirectory`, which the storage keeps alive, a path, or + ``None`` for the system temporary directory) with a unique name + starting with ``name``, and removed once nothing reads it. + """ + storage = cls(None, is_eternal=is_eternal) + storage._name = name + storage._parent_directory = directory + return storage + + @property + def preserve_storage_dir(self) -> bool: + """Whether this storage's directory stays on disk after it is gone.""" + return self._preserve_storage_dir + + @preserve_storage_dir.setter + def preserve_storage_dir(self, preserve: bool) -> None: + self._preserve_storage_dir = preserve + if self._directory is not None: + self._directory.preserve = preserve 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. + """Create a storage that starts with this storage's values. + + The clone gets copies of the file and enum mappings, so deleting or + rewiring entries through one storage does not change the other. It + reads the same ``.npy`` files, which stay on disk while it may read + them, and writes into a directory of its own, created on its first + write next to this storage's. From now on, this storage writes a new + file for a key rather than overwrite one the clone reads. """ - clone = OnDiskStorage( - self.storage_dir, - is_eternal=self.is_eternal, - preserve_storage_dir=True, + clone = OnDiskStorage.temporary( + self._name, self._parent_directory, is_eternal=self.is_eternal ) clone._files = self._files.copy() clone._enums = self._enums.copy() - clone._storage_dir_owner = getattr(self, "_storage_dir_owner", self) + clone._read_directories = self._read_directories + if self._directory is not None: + clone._read_directories += (self._directory,) + self._shared_files.update(self._files.values()) return clone def _decode_file(self, file: str) -> ArrayLike: @@ -63,6 +196,30 @@ def get(self, period: Period, branch_name: str = "default") -> ArrayLike: return None return self._decode_file(values) + def _get_storage_dir(self) -> str: + if self.storage_dir is None: + self._directory = StorageDirectory.create( + prefix=f"{self._name}_", parent=self._parent_directory + ) + self._directory.preserve = self._preserve_storage_dir + self.storage_dir = self._directory.path + return self.storage_dir + + def _get_path_to_write(self, key: str) -> str: + """Return the file to store a new value for ``key`` in. + + That is the key's current file or ``{key}.npy`` in this storage's + directory, unless a clone may read it; then a new numbered file. + """ + current = self._files.get(key) + if current in self._written_files and current not in self._shared_files: + return current + storage_dir = self._get_storage_dir() + path = os.path.join(storage_dir, f"{key}.npy") + while path in self._shared_files: + path = os.path.join(storage_dir, f"{key}.{next(self._versions)}.npy") + return path + def put( self, value: ArrayLike, period: Period, branch_name: str = "default" ) -> None: @@ -71,14 +228,19 @@ def put( period = periods.period(period) filename = f"{branch_name}_{period}" - path = os.path.join(self.storage_dir, filename) + ".npy" + path = self._get_path_to_write(filename) 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._written_files.add(path) self._files[filename] = path def delete(self, period: Period = None, branch_name: str = "default") -> None: + # Deleting forgets the mapping only. The file stays on disk until the + # directory is removed, since a clone may still read it. if period is None: # Only wipe files belonging to the requested branch (previously # this wiped every branch regardless of ``branch_name`` — same @@ -112,20 +274,16 @@ def get_known_branch_periods(self) -> list: ] def restore(self) -> None: + """Map each key to its latest file in ``storage_dir``.""" self._files = files = {} - # Restore self._files from content of storage_dir. + if self.storage_dir is None: + return + versions = {} for filename in os.listdir(self.storage_dir): if not filename.endswith(".npy"): continue - path = os.path.join(self.storage_dir, filename) - filename_core = filename.rsplit(".", 1)[0] - files[filename_core] = path - - def __del__(self) -> None: - if self.preserve_storage_dir: - return - shutil.rmtree(self.storage_dir) # Remove the holder temporary files - # If the simulation temporary directory is empty, remove it - parent_dir = os.path.abspath(os.path.join(self.storage_dir, os.pardir)) - if not os.listdir(parent_dir): - shutil.rmtree(parent_dir) + key, version = _parse_file_name(filename[: -len(".npy")]) + if version >= versions.get(key, -1): + versions[key] = version + files[key] = os.path.join(self.storage_dir, filename) + self._written_files.update(files.values()) diff --git a/policyengine_core/holders/holder.py b/policyengine_core/holders/holder.py index 52ae81295..347f68570 100644 --- a/policyengine_core/holders/holder.py +++ b/policyengine_core/holders/holder.py @@ -74,6 +74,23 @@ def clone(self, population: "Population") -> "Holder": def create_disk_storage( self, directory: str = None, preserve: bool = False ) -> OnDiskStorage: + """Create on-disk storage for this variable. + + With ``directory``, or ``preserve``, the storage uses + ``/`` (``directory`` defaults to the + simulation's ``data_storage_dir``). Otherwise it makes a directory of + its own inside ``data_storage_dir`` on its first write, so holders + for the same variable in different simulations of a family, or in a + branch that created its own holder, never share or remove each + other's files. + """ + is_eternal = self.variable.definition_period == periods.ETERNITY + if directory is None and not preserve: + return OnDiskStorage.temporary( + self.variable.name, + self.simulation._get_data_storage_directory(), + is_eternal=is_eternal, + ) if directory is None: directory = self.simulation.data_storage_dir storage_dir = os.path.join(directory, self.variable.name) @@ -81,7 +98,7 @@ def create_disk_storage( os.mkdir(storage_dir) return OnDiskStorage( storage_dir, - is_eternal=(self.variable.definition_period == periods.ETERNITY), + is_eternal=is_eternal, preserve_storage_dir=preserve, ) diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index 566407a59..fe6f54053 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -1,5 +1,4 @@ import hashlib -import tempfile from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Dict, List, Optional, Type, Union @@ -11,6 +10,7 @@ from policyengine_core import commons, periods from policyengine_core.data.dataset import Dataset +from policyengine_core.data_storage.on_disk_storage import StorageDirectory from policyengine_core.entities.entity import Entity from policyengine_core.enums import Enum, EnumArray from policyengine_core.errors import CycleError, SpiralError @@ -189,6 +189,7 @@ def __init__( self.max_spiral_loops: int = 10 self.memory_config: MemoryConfig = None self._data_storage_dir: str = None + self._data_storage_directory: Optional[StorageDirectory] = None self.branches: Dict[str, Simulation] = {} self.has_axes = False @@ -588,15 +589,21 @@ def data_storage_dir(self) -> str: Temporary folder used to store intermediate calculation data in case the memory is saturated """ if self._data_storage_dir is None: - self._data_storage_dir = tempfile.mkdtemp(prefix="openfisca_") - message = [ - ( - "Intermediate results will be stored on disk in {} in case of memory overflow." - ).format(self._data_storage_dir), - "You should remove this directory once you're done with your simulation.", - ] + # Removed once this simulation, its clones and every storage + # inside it are gone. + self._data_storage_directory = StorageDirectory.create(prefix="openfisca_") + self._data_storage_dir = self._data_storage_directory.path return self._data_storage_dir + def _get_data_storage_directory(self) -> Union[StorageDirectory, str]: + """Return ``data_storage_dir``, as the object that keeps it on disk + if the simulation created it.""" + path = self.data_storage_dir + directory = getattr(self, "_data_storage_directory", None) + if directory is not None and directory.path == path: + return directory + return path + # ----- Calculation methods ----- # def calculate( diff --git a/pyproject.toml b/pyproject.toml index 4f102bfb5..abc4bb467 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -43,6 +43,7 @@ dev = [ "ruff>=0.9.0", "coverage", "furo", + "hypothesis>=6,<7", "jupyter-book<1", "markupsafe==2.0.1", "mypy==0.991", diff --git a/tests/core/test_branch_scoped_delete_arrays.py b/tests/core/test_branch_scoped_delete_arrays.py index 5014bc62b..eed852cd1 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,30 @@ 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 into its own 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 + 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") + + assert clone.storage_dir != storage.storage_dir + (clone_file,) = clone._files.values() + assert os.path.dirname(clone_file) == clone.storage_dir + 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 +144,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_disk_storage_ownership.py b/tests/core/test_disk_storage_ownership.py new file mode 100644 index 000000000..363cdde74 --- /dev/null +++ b/tests/core/test_disk_storage_ownership.py @@ -0,0 +1,478 @@ +"""Disk-backed holder storage keeps the values each simulation reads. + +Every ``OnDiskStorage`` writes into a directory of its own and reads the files +it inherited from the storage it was cloned from. A directory stays on disk +while any storage may read a file in it, and a file a clone reads is never +overwritten. These tests pin that for branches, clones and holders that a +branch creates itself, and check, as a property over random sequences of +operations, that storing values on disk gives the same results as storing +them in memory. +""" + +from __future__ import annotations + +import contextlib +import gc +import os +import sys +import warnings +from dataclasses import dataclass +from typing import List, Optional + +import numpy as np +import pytest +from hypothesis import HealthCheck, given, settings +from hypothesis import strategies as st + +from policyengine_core.country_template import CountryTaxBenefitSystem +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 + +PERIOD = "2017-01" + + +def _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 + + +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 = _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(): + """``Simulation.clone`` keeps the branch name, but not the files.""" + simulation = _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]) + + +def test_branches_with_the_same_name_keep_their_own_values(): + simulation = _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 = _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 = _simulation() + 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_branch_directory_goes_with_the_branch(): + simulation = _simulation() + simulation.set_input("salary", PERIOD, np.array([1_000.0])) + branch = simulation.get_branch("branch") + branch.set_input("salary", PERIOD, np.array([2.0])) + root_dir = simulation.get_holder("salary")._disk_storage.storage_dir + branch_dir = branch.get_holder("salary")._disk_storage.storage_dir + assert branch_dir != root_dir + assert os.path.dirname(branch_dir) == simulation.data_storage_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 = _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) + + +def test_a_data_storage_dir_set_by_the_user_stays(tmp_path): + directory = tmp_path / "simulation" + simulation = _simulation() + simulation._data_storage_dir = str(directory) + simulation.set_input("salary", PERIOD, np.array([1.0])) + storage_dir = simulation.get_holder("salary")._disk_storage.storage_dir + assert os.path.dirname(storage_dir) == str(directory) + + del simulation + gc.collect() + + assert not os.path.exists(storage_dir) + assert directory.is_dir() + + +def test_a_key_stored_again_overwrites_its_file_until_a_clone_reads_it(tmp_path): + storage = OnDiskStorage.temporary("salary", str(tmp_path)) + storage.put(np.array([1.0]), PERIOD) + storage.put(np.array([2.0]), PERIOD) + assert len(os.listdir(storage.storage_dir)) == 1 + + clone = storage.clone() + storage.put(np.array([3.0]), PERIOD) + storage.put(np.array([4.0]), PERIOD) + + assert len(os.listdir(storage.storage_dir)) == 2 + np.testing.assert_array_equal(clone.get(PERIOD), [2.0]) + np.testing.assert_array_equal(storage.get(PERIOD), [4.0]) + + +def test_a_clone_keeps_the_directories_it_reads(tmp_path): + storage = OnDiskStorage.temporary("salary", str(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_the_latest_file_of_each_key(tmp_path): + 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), [2.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]) + + +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_plain_array_stored_over_an_enum_reads_back_plain(tmp_path): + storage = OnDiskStorage.temporary("status", str(tmp_path)) + 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) + + +@pytest.mark.skipif(not hasattr(os, "fork"), reason="needs os.fork") +def test_a_forked_child_leaves_the_parent_directories(tmp_path): + storage = OnDiskStorage.temporary("salary", str(tmp_path)) + storage.put(np.array([1.0]), PERIOD) + storage_dir = storage.storage_dir + + with warnings.catch_warnings(): + # Python warns that forking a multi-threaded process may deadlock. + warnings.simplefilter("ignore", DeprecationWarning) + pid = os.fork() + if pid == 0: + try: + del storage + gc.collect() + finally: + os._exit(0) + os.waitpid(pid, 0) + + assert os.path.isdir(storage_dir) + np.testing.assert_array_equal(storage.get(PERIOD), [1.0]) + + +# Differential property: on disk as in memory. + +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(_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 == "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("delete_arrays"), _index, st.sampled_from(INPUTS + CALCULATED) + ), + st.tuples(st.just("new_holders"), _index), + st.tuples(st.just("drop"), _index), + ), + max_size=30, +) + + +@settings( + max_examples=75, + deadline=None, + suppress_health_check=[HealthCheck.too_slow], +) +@given(operations=_operations) +def test_disk_storage_gives_the_results_memory_storage_gives(operations): + """For any sequence of branching, cloning, inputs, calculations 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 directory is 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) + + directory = on_disk.members[0].simulation._data_storage_dir + del in_memory, on_disk, member_in_memory, member_on_disk + gc.collect() + + assert errors == [] + assert directory is None or not os.path.exists(directory) diff --git a/uv.lock b/uv.lock index 2608e8504..222deb407 100644 --- a/uv.lock +++ b/uv.lock @@ -130,7 +130,7 @@ name = "cffi" version = "2.0.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "pycparser" }, + { name = "pycparser", marker = "implementation_name != 'PyPy'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/eb/56/b1ba7935a17738ae8453301356628e8147c79dbb825bcbc73dc7401f9846/cffi-2.0.0.tar.gz", hash = "sha256:44d1b5909021139fe36001ae048dbdde8214afa20200eda0f64c068cac5d5529", size = 523588, upload-time = "2025-09-08T23:24:04.541Z" } wheels = [ @@ -682,6 +682,94 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/d5/ae/2f6d96b4e6c5478d87d606a1934b5d436c4a2bce6bb7c6fdece891c128e3/huggingface_hub-1.4.1-py3-none-any.whl", hash = "sha256:9931d075fb7a79af5abc487106414ec5fba2c0ae86104c0c62fd6cae38873d18", size = 553326, upload-time = "2026-02-06T09:20:00.728Z" }, ] +[[package]] +name = "hypothesis" +version = "6.168.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "sortedcontainers" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/09/b7/13118bbc45d6d8b9d04e2de779e2a4ff23145ea39efa692b994fb874ca72/hypothesis-6.168.3.tar.gz", hash = "sha256:a43388f9067678fef6e13bdff325b6cfa6961a590498bb37f7ff31589c83bc75", size = 511022, upload-time = "2026-09-28T05:20:58.499Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2b/cd/746a2fc1e5e5ba34f07ea6ee1ad07525a3dde5ff4836db3a7980af49a891/hypothesis-6.168.3-cp310-abi3-macosx_10_12_x86_64.whl", hash = "sha256:d20972ca134e652e9928ecd200966a8a2857adc2812c1d74b12a872e5cd503eb", size = 791536, upload-time = "2026-09-28T05:18:57.162Z" }, + { url = "https://files.pythonhosted.org/packages/c3/0f/7a158e377b69556c8e12c25c8fd811d0103e54c24a12dab1f3b9819f202e/hypothesis-6.168.3-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:eafbec09d3e87d13d8242411d1f5868b5e879f1bd6d95e233e9ff80c527bea1c", size = 787313, upload-time = "2026-09-28T05:20:18.122Z" }, + { url = "https://files.pythonhosted.org/packages/45/c6/7df9104ee359e8fcd77791781c47e70794964ca69673665e0de2a6590f4c/hypothesis-6.168.3-cp310-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:73c5627497968cc62d140e9fde1ea21e12cac7b64dc419fea8286cd34ff1ad4e", size = 1124036, upload-time = "2026-09-28T05:20:04.372Z" }, + { url = "https://files.pythonhosted.org/packages/92/0f/8a61715404a73b9a82cb1f976803428d603cdea976a9c30266c72ad6a7ce/hypothesis-6.168.3-cp310-abi3-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:29dc56e6dc6eeb0aeb08ac463f847279bde2336c4786faf7ee0af4b93cd031a7", size = 1147875, upload-time = "2026-09-28T05:19:15.495Z" }, + { url = "https://files.pythonhosted.org/packages/e4/3f/2b16e95cc3b9a069afa23830a012aefe377ce4a457a227e5a5026eb827ff/hypothesis-6.168.3-cp310-abi3-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:753bb501f8560d2e321ed3b62596b4496c56f15668b0a0669231ccc3e6c4e80d", size = 1149489, upload-time = "2026-09-28T05:19:18.886Z" }, + { url = "https://files.pythonhosted.org/packages/f6/e7/d7cf6dc068bb02732b2a6e38a35a0dcb7f2137608198fc79d740f69d197c/hypothesis-6.168.3-cp310-abi3-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:71ab606c472449cb872ef2a7acaec679ee4f651b7f0a477a04ebc85d8933ef8c", size = 1191992, upload-time = "2026-09-28T05:19:21.83Z" }, + { url = "https://files.pythonhosted.org/packages/5f/26/f1e5b25dec15998e8375d722825ec3dae1075cbe7cb1f74221b2312a2e33/hypothesis-6.168.3-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:26928956c54748e4dfa587333741ab750246123b93484955646ff316cca27eef", size = 1169911, upload-time = "2026-09-28T05:19:27.996Z" }, + { url = "https://files.pythonhosted.org/packages/03/36/901234a49147e6bf8a5b9e54b775706a223c2b52befea1106ea2aaf15d48/hypothesis-6.168.3-cp310-abi3-manylinux_2_31_riscv64.whl", hash = "sha256:0bdfc53c041b61c854fa3761735736b991bc2bddafee45a361cfe4d43027c1ef", size = 1129406, upload-time = "2026-09-28T05:18:42.632Z" }, + { url = "https://files.pythonhosted.org/packages/55/c1/5fc9ff91ec9fba6bb527833386c66812524b5cc75fe7d4a4f9122ee1c420/hypothesis-6.168.3-cp310-abi3-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:650528e2b1e2a45e95c2df624d4b9364b8ac5a027ba84e1eea2bc5398dd01dcd", size = 1160399, upload-time = "2026-09-28T05:20:23.911Z" }, + { url = "https://files.pythonhosted.org/packages/9b/d2/49ad1ef5c55547c8abfd0fc571b3493806dda530fd5cb3d8263dbad86a72/hypothesis-6.168.3-cp310-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:209dc54cdb1b4d6d7020ad8d09e44c49756361b7183adacea4d6f5705085b595", size = 1299906, upload-time = "2026-09-28T05:18:58.417Z" }, + { url = "https://files.pythonhosted.org/packages/0d/5f/f98a26094c4f462c4215807bed0f9a6b5508b27a3c329bcba1bc8adc7ac7/hypothesis-6.168.3-cp310-abi3-musllinux_1_2_armv7l.whl", hash = "sha256:af8ca98cd11dc7f9427bb90483abcd9688d4df2e8e63893a30a2423b027ebb11", size = 1425538, upload-time = "2026-09-28T05:20:49.074Z" }, + { url = "https://files.pythonhosted.org/packages/d0/b8/c5b4406e3a6591e51a42f85469351d1d824e0a54b84167aec9b2353759fe/hypothesis-6.168.3-cp310-abi3-musllinux_1_2_i686.whl", hash = "sha256:f8122cfdd0bc0ba843063effa22ac310bdebdb3a1319bf1e63f14baa119c72ab", size = 1377084, upload-time = "2026-09-28T05:20:06.618Z" }, + { url = "https://files.pythonhosted.org/packages/e0/01/0d84a8ea469d024c15f602799c0b8410fbe03d99fdb0370449b33c3c916f/hypothesis-6.168.3-cp310-abi3-musllinux_1_2_ppc64le.whl", hash = "sha256:4bedbb379eab34f792af7ee9a05aae04e9c08bbb52e0f90f5a2110d4fe4b2fbd", size = 1281162, upload-time = "2026-09-28T05:19:42.768Z" }, + { url = "https://files.pythonhosted.org/packages/22/b5/ea6435038de1a795f005cb603531e8ccc053febc50180f3ab55583c9c60e/hypothesis-6.168.3-cp310-abi3-musllinux_1_2_riscv64.whl", hash = "sha256:5a953115b9f5c95133ab2d04efffeec96e5658c3207923dca7285f7db3e6bbef", size = 1300433, upload-time = "2026-09-28T05:18:55.634Z" }, + { url = "https://files.pythonhosted.org/packages/6c/78/fa6c77d9f64b6d592e69e582d83cbce7d7e56897faa4762885a2123b0166/hypothesis-6.168.3-cp310-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:6173558e676ad25ed1e20507fa4024a0d816dd90f715f77006e5a28f19109026", size = 1336273, upload-time = "2026-09-28T05:19:08.273Z" }, + { url = "https://files.pythonhosted.org/packages/f1/56/4acd0d778bab818c9ea336fcba768bf43bc70db7f84cc373579298df4641/hypothesis-6.168.3-cp310-abi3-win32.whl", hash = "sha256:dc66390fb12d80585aa9222bf538ce8b7aa22cf5d118250647355a1c9e8f62f4", size = 678211, upload-time = "2026-09-28T05:19:23.355Z" }, + { url = "https://files.pythonhosted.org/packages/b0/db/0a02b146ad1c30f9716362f2e1a379dfd6b0f7be03ac1e22c9da24c9e2a9/hypothesis-6.168.3-cp310-abi3-win_amd64.whl", hash = "sha256:92325b276360fe86c5bf71a568c0d53a6d140b0de36dcf029f9164a17803bb24", size = 684906, upload-time = "2026-09-28T05:18:35.099Z" }, + { url = "https://files.pythonhosted.org/packages/a3/c0/958deaf726848f96f52250740bf39f13f476b068e22f73e2b358eaa07532/hypothesis-6.168.3-cp310-abi3-win_arm64.whl", hash = "sha256:3cf6f1eeaf41cd8d60cf1f88fde905ca1dd77c906929a507d6ac7f66f2ccba2a", size = 683292, upload-time = "2026-09-28T05:19:57.404Z" }, + { url = "https://files.pythonhosted.org/packages/f4/ab/443208c2988a24eaafbfbd6791b897ad9a8a1dd7f870d5b455d7d15a28ad/hypothesis-6.168.3-cp311-cp311-macosx_10_12_x86_64.whl", hash = "sha256:9029775dc2e25e3e8e596c4306926f0079158353631c7c52c285bbc5c8073c2b", size = 792239, upload-time = "2026-09-28T05:19:59.144Z" }, + { url = "https://files.pythonhosted.org/packages/cf/7b/619c1a76c52de2d3cfd68a2f9876c09464d63097db5bf86a28b62dc37bb1/hypothesis-6.168.3-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:0e1d91530cefdb9e0b6467c203eb509c61a11d70480a37db59cb53a2f9913102", size = 788113, upload-time = "2026-09-28T05:19:17.145Z" }, + { url = "https://files.pythonhosted.org/packages/ea/dd/33841e8833105f850214a486d83740ba52041892ca94e687a05829c0ce83/hypothesis-6.168.3-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:0df00cbe8133aa220308d11fa28e3add996cf5102a39b279c1d711015fd9114b", size = 1124241, upload-time = "2026-09-28T05:18:37.589Z" }, + { url = "https://files.pythonhosted.org/packages/4b/43/c84a7900feced6a420e824ab61d316e4f59a6887d991bc143ecc294a09ed/hypothesis-6.168.3-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:fd7f75a2e23288ee82ee965a952473d9c5447c2cbc1afc94be09d2402201774e", size = 1170406, upload-time = "2026-09-28T05:19:46.365Z" }, + { url = "https://files.pythonhosted.org/packages/47/5d/6e2a32b9c3020af7a06b00fe446a2c5155ae76ac468c0b68d7188e524002/hypothesis-6.168.3-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:bc2b1c37631f2231dd0d077225aa5908fc2bd34a4c3faf475053b0bc2eb24c95", size = 1300339, upload-time = "2026-09-28T05:20:31.968Z" }, + { url = "https://files.pythonhosted.org/packages/f9/78/19f3c9eb9a5c215b9fbde57be18ef7c1c78afa509b460d4e8f86026d741a/hypothesis-6.168.3-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:73960a58f6efc8cbc57d8ad0647955f7c5b25f3d562e3eaa57afa1c69894f7aa", size = 1336484, upload-time = "2026-09-28T05:19:02.841Z" }, + { url = "https://files.pythonhosted.org/packages/88/bb/d46c4519e989e022667705d8a56e83a09ad4ea50eae07ca4bb0f4708dd23/hypothesis-6.168.3-cp311-cp311-win_amd64.whl", hash = "sha256:3f1122223759acc0fe5301c1ae505d762ee61af91db7b191a791ccdbcd965462", size = 684646, upload-time = "2026-09-28T05:20:00.842Z" }, + { url = "https://files.pythonhosted.org/packages/7a/56/44a17266bcf1add5fba7c3df5c3920396a5e924d917fb0534a6f9e2b9b9a/hypothesis-6.168.3-cp312-cp312-macosx_10_12_x86_64.whl", hash = "sha256:bff12e036b67abc5ded682f724f4d7afaff5223190270a0763695f218c079068", size = 793313, upload-time = "2026-09-28T05:20:22.045Z" }, + { url = "https://files.pythonhosted.org/packages/d8/18/1a25b11e9133cb54c02ee6cf99a61296ad461aabc7aae3399e1d0ff1f095/hypothesis-6.168.3-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:c73c6188056e6dc110260e39d141c69c3d02f23e3ba1d8610c7a8b6f2510ec96", size = 784855, upload-time = "2026-09-28T05:19:34.496Z" }, + { url = "https://files.pythonhosted.org/packages/8b/da/728d3210f0b0089337f37206a100d31ef546d1f99839e709f0a9ca443f8e/hypothesis-6.168.3-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:bb1063cb794765097b8c0add598bd34a904f44dc15db81275dd1f0ed1ac567ed", size = 1123043, upload-time = "2026-09-28T05:19:05.468Z" }, + { url = "https://files.pythonhosted.org/packages/53/86/bfb9ec344a8b3ef2fa6f59bb3d30c5b3a7d86d594107670539c34ee53d1b/hypothesis-6.168.3-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:9b2d47f5d9060b049039bef0b267e88480f6468f95ee394e514e84a81ccdc5f5", size = 1169121, upload-time = "2026-09-28T05:20:40.512Z" }, + { url = "https://files.pythonhosted.org/packages/3b/e0/5a76dbafbcaf8fff075f09377cddefa17289ef05dbf51878dde88385b966/hypothesis-6.168.3-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:8367f623a98cbf90f33fd2a43b100671b152b5975bc343a578290945f52f7018", size = 1298749, upload-time = "2026-09-28T05:19:32.883Z" }, + { url = "https://files.pythonhosted.org/packages/28/15/0e65af96b353f2c4bdc283c53ed0980e5ac5009afe97ef3b23c77b364664/hypothesis-6.168.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:fd81241f4cd76fb18d481e87b9688b30a0cb5c30841730513323ba1c7aba1837", size = 1335271, upload-time = "2026-09-28T05:20:54.005Z" }, + { url = "https://files.pythonhosted.org/packages/f5/c5/37eab2155664aa544df321f6d3940b89c6702d79a72fdc98d7b5a2286f8c/hypothesis-6.168.3-cp312-cp312-win_amd64.whl", hash = "sha256:377438de53afb94347d9845b7d90db5905c6d2a8d612108e9938e0f80503bb90", size = 682211, upload-time = "2026-09-28T05:18:51.206Z" }, + { url = "https://files.pythonhosted.org/packages/0c/a2/6787da846d929e52fc3344d803299c45782dbeac287528af08380e984bf8/hypothesis-6.168.3-cp313-cp313-macosx_10_12_x86_64.whl", hash = "sha256:1b230a850de63334c16654a34a2d547e0179d36b9071d4439b3e7237f6d077e7", size = 793254, upload-time = "2026-09-28T05:19:36.294Z" }, + { url = "https://files.pythonhosted.org/packages/89/ba/2893f5ca42501f4d562cba3229c8994c3a3e5fac66ef60c1e0e0b5220a5d/hypothesis-6.168.3-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:d63b0226cd3e0d8bdd97c3384b22a21934ed4d53246c1c93575dada616672499", size = 784699, upload-time = "2026-09-28T05:18:54.194Z" }, + { url = "https://files.pythonhosted.org/packages/b7/aa/7d7349daf75b71f6f35876f8de115e974c5c04d600e86df1b2779b0883ab/hypothesis-6.168.3-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:0369f5df055f96e117ab12e5f249668ff144731ab5280bc7a205fdf81f990b89", size = 1123016, upload-time = "2026-09-28T05:20:36.372Z" }, + { url = "https://files.pythonhosted.org/packages/00/f0/7774e1ea072708ea46cb25c4aeee5f9978b6c271764809069e8a6789858e/hypothesis-6.168.3-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:f076bcd0f77fdcdb7797826c879099573d03a02e228ebe649ea917081b962ac0", size = 1168989, upload-time = "2026-09-28T05:20:08.44Z" }, + { url = "https://files.pythonhosted.org/packages/a7/7d/113992abed9efbd7944e3a496a382da58152dce126fddf6419ff31a0a57d/hypothesis-6.168.3-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:b823ba1fcec8da730f29316d010b06d3f7e0c3828dcf91020e24c55e7d24652a", size = 1298626, upload-time = "2026-09-28T05:19:40.986Z" }, + { url = "https://files.pythonhosted.org/packages/d5/65/3659fa5e733027e5b37a27486e57f4053535e40853d23501c422bb275043/hypothesis-6.168.3-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:dd2849c269d674e4618590f3b48d443bd4c06b5aef3d8d869086c2e6d213d248", size = 1335184, upload-time = "2026-09-28T05:20:20.087Z" }, + { url = "https://files.pythonhosted.org/packages/b4/6c/35aab2221b5ea65125340f236e1a765a9e9f3b28d89a389ea0033fed29f7/hypothesis-6.168.3-cp313-cp313-win_amd64.whl", hash = "sha256:3ef7d26f5789e691401d5f87eafed9bd2763f0dbe47af6d6b66012a509404766", size = 682173, upload-time = "2026-09-28T05:19:12.762Z" }, + { url = "https://files.pythonhosted.org/packages/6e/79/27b0cb56ff5d2bc92458bf6fcebdb1f0dc01f57f043a178e9f9d74498169/hypothesis-6.168.3-cp314-cp314-macosx_10_12_x86_64.whl", hash = "sha256:e2d4c68729a13df9af4998d2652cfb5d541c5609b88c306880a5dfb284938ec2", size = 793287, upload-time = "2026-09-28T05:20:38.522Z" }, + { url = "https://files.pythonhosted.org/packages/e8/bd/5342f95c3bc36586777ef8cb8eba87ac7acf0184210327fb0e6d8654e295/hypothesis-6.168.3-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:b345f818083ec99966a43ca4f7b38feb62c6920ce28572bc2bde4948a8b7eaba", size = 784857, upload-time = "2026-09-28T05:20:42.609Z" }, + { url = "https://files.pythonhosted.org/packages/4a/5f/fc774241518a0588d362680a3084275058aab86d6bc02a2c61e1073142d0/hypothesis-6.168.3-cp314-cp314-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:0608c610fc002978fc5de0f471770d8817e9455cb8649a47983c9f8bccfc1897", size = 1123330, upload-time = "2026-09-28T05:19:29.735Z" }, + { url = "https://files.pythonhosted.org/packages/b0/e6/131f16775a3dca5f4fe27f0d6ad9a9851600098a54a872d163f6ebf0b664/hypothesis-6.168.3-cp314-cp314-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:9dcf6448b1ecc37f2b23f2d1b3ddfc9ff6b6910614c15a642dc82f419b96037b", size = 1169178, upload-time = "2026-09-28T05:19:53.436Z" }, + { url = "https://files.pythonhosted.org/packages/06/61/c9f5bff8b73321c9fd12f2d69666c6e4861b97b60f5f24fdd6ac293007fd/hypothesis-6.168.3-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:ae4f9f094041dcce5119ebd7bab71062743ab02b6e654d056b370beda78c19e2", size = 1299127, upload-time = "2026-09-28T05:20:14.486Z" }, + { url = "https://files.pythonhosted.org/packages/e4/6d/d4618f8ab12dd76c4d58ea052252759aa2bf0c4e71f4b502ba5d577e0b9f/hypothesis-6.168.3-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:d479985fe73af97badfb72dc6d20c6a353e486a36f6d065f600026e0cf954a86", size = 1335390, upload-time = "2026-09-28T05:20:10.326Z" }, + { url = "https://files.pythonhosted.org/packages/95/b6/727161e17cc297df1aeea945de05c532033b66751625bec192bc5b5ff707/hypothesis-6.168.3-cp314-cp314-pyemscripten_2026_0_wasm32.whl", hash = "sha256:785e2c45f8c08e274b4bf1ccae97f1a4e09407e9a68e80790cfab1d17a4a45fa", size = 624291, upload-time = "2026-09-28T05:18:36.376Z" }, + { url = "https://files.pythonhosted.org/packages/77/23/2f9b506392a0b8712440bcb781abc5713be9ffac1303b6d10f3669050882/hypothesis-6.168.3-cp314-cp314-win_amd64.whl", hash = "sha256:320920b1e3dae8611eee8a03d063cf2187446f2a17c38cfb8a7fc1466f71eee2", size = 682064, upload-time = "2026-09-28T05:20:16.336Z" }, + { url = "https://files.pythonhosted.org/packages/50/93/efabfd95eb2b69c0c9fa1e1c83c8290aa2715d46baeb5a7764faf09c133a/hypothesis-6.168.3-cp314-cp314t-macosx_10_12_x86_64.whl", hash = "sha256:35380baa981108a7f60c4eab71e46acd8d8f58440520346a1a6aba06dca7e074", size = 791875, upload-time = "2026-09-28T05:19:37.789Z" }, + { url = "https://files.pythonhosted.org/packages/1b/fc/2a0ada1623a9a33048bbf9f00a7c92513398c6b8865cba5efb854262ce30/hypothesis-6.168.3-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:f0aaaed00438fa6673d856aaee12a4a8afbde82a9f3c5ff487cacfb064239afe", size = 783412, upload-time = "2026-09-28T05:18:48.266Z" }, + { url = "https://files.pythonhosted.org/packages/41/31/73c615d37eeb12da0209555612c32ced628d860267c23cfd2691c6207aea/hypothesis-6.168.3-cp314-cp314t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f27e6df1576bf838e7f484d4ae5cba92114497ca30afb917056bf1ea33e95675", size = 1121662, upload-time = "2026-09-28T05:19:26.225Z" }, + { url = "https://files.pythonhosted.org/packages/92/54/1893344a3b8bbf83fa7c9bb7e96f6836580213854f46671cb749f7b4706c/hypothesis-6.168.3-cp314-cp314t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:3bc85014577982ec6d2e266edc5cd6e7a7d3674c648791ca983c67a52f89a0c4", size = 1167761, upload-time = "2026-09-28T05:18:49.85Z" }, + { url = "https://files.pythonhosted.org/packages/26/97/a61d82968febd25a0f08ffbbfabfaeead66f68fc5fa73f3c9ece1c36fc8f/hypothesis-6.168.3-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:ccfc29505aa1cdcc254cb9cd701fd0811fef5480addd97d4127a00df023410f7", size = 1297309, upload-time = "2026-09-28T05:19:51.711Z" }, + { url = "https://files.pythonhosted.org/packages/bc/01/fe8e4cf6d39d02f4efadaa351427fa13722086cfc54a2da3579741c5177a/hypothesis-6.168.3-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:32d0699566aaa93f9e97a44705d7164386f1e91de78d277f53bada35e78bcc96", size = 1334260, upload-time = "2026-09-28T05:19:00.121Z" }, + { url = "https://files.pythonhosted.org/packages/39/6a/09177ebc62f4778f94379ecade6a51299d3d8c4f4b75ba636bf3a0202364/hypothesis-6.168.3-cp314-cp314t-win_amd64.whl", hash = "sha256:28d88fa174ecbd4ecbd7bb290f06d0db084a3971c3f511ae2830b65b5e25500f", size = 681992, upload-time = "2026-09-28T05:19:48.043Z" }, + { url = "https://files.pythonhosted.org/packages/f9/f6/890bf33d608cd63348d3146a3ca83362fbf09e589c5e5230a00c403070f4/hypothesis-6.168.3-cp315-abi3.abi3t-macosx_10_12_x86_64.whl", hash = "sha256:61f5782d807b1e6aa5c9beef1054cb2037e7d1add48a64972cd6ee447778781f", size = 791231, upload-time = "2026-09-28T05:18:43.998Z" }, + { url = "https://files.pythonhosted.org/packages/03/dd/fc36b204f8aa7437c676576d9437f0422de46aa8a729e6cd5bbb0224e759/hypothesis-6.168.3-cp315-abi3.abi3t-macosx_11_0_arm64.whl", hash = "sha256:7aacf3cf40c7ce8f9e4924d5b57bc0b068beafdfed2347bcf160a28b4cfbba7b", size = 783171, upload-time = "2026-09-28T05:19:20.353Z" }, + { url = "https://files.pythonhosted.org/packages/5c/d2/3cd087d577db67c404991d7723e64d7cd75570cdc82a0f7329aed6f2086c/hypothesis-6.168.3-cp315-abi3.abi3t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:01768a03a30dc54df7fe457c0b34016c84d598fa00d5965ededab93ba4eb3408", size = 1121152, upload-time = "2026-09-28T05:18:39.095Z" }, + { url = "https://files.pythonhosted.org/packages/42/6f/a05a68cc66e7a22701c50bf94fb9ddeb36f4f5d0de3cc130066709eda621/hypothesis-6.168.3-cp315-abi3.abi3t-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:1a8a4ffc6c6e6e577f2bfbcfebf7cffbb310283ba2532c729570a0a752cebaaa", size = 1144069, upload-time = "2026-09-28T05:20:44.705Z" }, + { url = "https://files.pythonhosted.org/packages/4a/f4/fdd7a093fb2fb3a0ce24256c92ad5bf67936f4669bd06d5ceae4298150ec/hypothesis-6.168.3-cp315-abi3.abi3t-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:b987d73eba95183a7e59cca6d1925c588aa4307922d9852cc2b4d282a2ae4128", size = 1146692, upload-time = "2026-09-28T05:20:25.83Z" }, + { url = "https://files.pythonhosted.org/packages/f1/c8/6dbd4377e935505ae4fc8ee4c7b18c69994c4bbaca015447c470218bdbee/hypothesis-6.168.3-cp315-abi3.abi3t-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:d4569c39bd97d9573e946429ed676f3b55a7c8ed80920addd67d384a47d59b38", size = 1189480, upload-time = "2026-09-28T05:20:46.734Z" }, + { url = "https://files.pythonhosted.org/packages/ba/bf/ff288b496b690000d2686dcfa7f67855e4c5a6dddb464ed8e4880be83bb2/hypothesis-6.168.3-cp315-abi3.abi3t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6042b8707a4b25bbbfe20b110258fb7549a68951e5525608a8ce06091039e9fc", size = 1167109, upload-time = "2026-09-28T05:19:24.758Z" }, + { url = "https://files.pythonhosted.org/packages/d9/d1/3fee2bc29fc747fabbb27e174515592e390809ca465d3bbbabbcfa3235a8/hypothesis-6.168.3-cp315-abi3.abi3t-manylinux_2_31_riscv64.whl", hash = "sha256:f413629de94d38a7a2ad259697a6752526143cb49c2e2d7bdba19381c80693f6", size = 1126823, upload-time = "2026-09-28T05:18:52.636Z" }, + { url = "https://files.pythonhosted.org/packages/58/6a/585294392fa6a6d9719446a042a777ba98ecca265d9e48cbedc10c88df0f/hypothesis-6.168.3-cp315-abi3.abi3t-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:f071737e4e775bebba07e319eb645a880d1e0186b4d24bad23255e849d86d483", size = 1155773, upload-time = "2026-09-28T05:20:51.268Z" }, + { url = "https://files.pythonhosted.org/packages/cd/4d/fa283ff79996debf1ff08593e01f8b773eb8211b19b3d6d5d491317a65bf/hypothesis-6.168.3-cp315-abi3.abi3t-musllinux_1_2_aarch64.whl", hash = "sha256:f59a3912858d0609c26054aa1c474847c797937f5e0e1c77f0d3f08032e50f09", size = 1296671, upload-time = "2026-09-28T05:18:46.937Z" }, + { url = "https://files.pythonhosted.org/packages/42/05/98f9c2f628da5afabb6c3e8df5d2d299d2d50bf306b14a6e26d3954c4951/hypothesis-6.168.3-cp315-abi3.abi3t-musllinux_1_2_armv7l.whl", hash = "sha256:f09c05a23a8025dd5cad07c2ed299a46778e72d49ff0dd5bf353bcc0f7dc1580", size = 1422001, upload-time = "2026-09-28T05:19:04.188Z" }, + { url = "https://files.pythonhosted.org/packages/b8/c9/f2109ade29a7ec284d55bca328d784743e8b13e321576fd34ee7e65f99af/hypothesis-6.168.3-cp315-abi3.abi3t-musllinux_1_2_i686.whl", hash = "sha256:44ace770bda3a0301739fc5a413c764de790c739df1f1a7218dd049f8594d9f5", size = 1374182, upload-time = "2026-09-28T05:20:27.764Z" }, + { url = "https://files.pythonhosted.org/packages/ce/6d/e05d5f72441564a3bebc71fa155deafd0fd3b015d6014ec8a00edf6e42bb/hypothesis-6.168.3-cp315-abi3.abi3t-musllinux_1_2_ppc64le.whl", hash = "sha256:03b131043608f94a2578a2a896a1a72092acb7079815eef5b8513b08447e0e62", size = 1278402, upload-time = "2026-09-28T05:20:56.332Z" }, + { url = "https://files.pythonhosted.org/packages/96/d7/6988a7f1f69c5c530687dce03a32c157e2d878e3a2e32af9269ce016b78a/hypothesis-6.168.3-cp315-abi3.abi3t-musllinux_1_2_riscv64.whl", hash = "sha256:e9784aca26eddfe99b03a0292320db742cc8a74200ef864fdce949f526f973cd", size = 1297770, upload-time = "2026-09-28T05:19:11.387Z" }, + { url = "https://files.pythonhosted.org/packages/a2/28/5bb82b60b836bd2329e1fe01ad94efc2b14dfa806f8e50ed14b795d78770/hypothesis-6.168.3-cp315-abi3.abi3t-musllinux_1_2_x86_64.whl", hash = "sha256:2b52ac363096232bebc2add117e9178d91f4248f4cdb919fd1026b5f86a4bb16", size = 1333977, upload-time = "2026-09-28T05:20:02.483Z" }, + { url = "https://files.pythonhosted.org/packages/8e/7d/d419841b8f65481ea1a50c4ba36f2670d48b8c7a51e4e8586089564cf7f6/hypothesis-6.168.3-cp315-abi3.abi3t-win32.whl", hash = "sha256:d28e3a6b511a74ce37df5274b51f36c2b274fea7365e00b16f0c81e22acd5957", size = 675394, upload-time = "2026-09-28T05:19:55.237Z" }, + { url = "https://files.pythonhosted.org/packages/d8/d2/1de6a2ad100e44817f2e9a8e8ba3eeaf4769621f4f87f2d4966c4971fcc4/hypothesis-6.168.3-cp315-abi3.abi3t-win_amd64.whl", hash = "sha256:7b9638789548361a57d984f56619ac694a914c328181d911271618409261ff4a", size = 681699, upload-time = "2026-09-28T05:19:09.957Z" }, + { url = "https://files.pythonhosted.org/packages/c0/79/be3370fd02734d6b1d950183ae9224580d19fa8d356b95348d60f1e3eeb7/hypothesis-6.168.3-cp315-abi3.abi3t-win_arm64.whl", hash = "sha256:65d78e4357ec48ed2c67825f06740ee3599be4cfe770a6092bed07108d679ac5", size = 679849, upload-time = "2026-09-28T05:19:06.884Z" }, + { url = "https://files.pythonhosted.org/packages/3f/0e/27960db1d45a94497e1681fb9d7b9140b44da4e655ba20f0d82f95d10230/hypothesis-6.168.3-pp311-pypy311_pp73-macosx_10_12_x86_64.whl", hash = "sha256:101531b4ccf7fa12965a6b295d10b64c2f8a8fab61e1f9b18c55335d8fb02575", size = 793135, upload-time = "2026-09-28T05:19:14.056Z" }, + { url = "https://files.pythonhosted.org/packages/9a/cb/be08379286c1925fcd651b92ed835f470a83b859b5c484b3431da9d851f5/hypothesis-6.168.3-pp311-pypy311_pp73-macosx_11_0_arm64.whl", hash = "sha256:41e0120814de5c3a6b58d8cb27cb11b8873ab33fba130cd09855a27f5a12a1ca", size = 788995, upload-time = "2026-09-28T05:20:34.164Z" }, + { url = "https://files.pythonhosted.org/packages/38/7b/6d885d139b2ed76e2f4908cf10eae56acc73bf05cbef286357c4559c9a04/hypothesis-6.168.3-pp311-pypy311_pp73-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:9d120009d145697909f2f0c532f2d568558ddd6990707f8cba6576bfb9770696", size = 1124974, upload-time = "2026-09-28T05:19:49.903Z" }, + { url = "https://files.pythonhosted.org/packages/05/cc/b302ae809b48254830cae57801d0af0b66783df515b5dd56b9f21c8dfa86/hypothesis-6.168.3-pp311-pypy311_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:818b3d09bba60ef90463e47a4bc87275dba1e8aeb184d3102888440ca7c7837d", size = 1171967, upload-time = "2026-09-28T05:20:12.295Z" }, + { url = "https://files.pythonhosted.org/packages/2c/ab/51031080a6600f5b07adf5f9dfa117345c2adda07a5c24bda523e70bd1e9/hypothesis-6.168.3-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:9e3c36eba80e089ecfa28ec08049723223cb41d3cfe4d9988169cf98440e402b", size = 685641, upload-time = "2026-09-28T05:19:44.609Z" }, +] + [[package]] name = "idna" version = "3.11" @@ -1324,7 +1412,7 @@ name = "pexpect" version = "4.9.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "ptyprocess" }, + { name = "ptyprocess", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/42/92/cc564bf6381ff43ce1f4d06852fc19a2f11d180f23dc32d9588bee2f149d/pexpect-4.9.0.tar.gz", hash = "sha256:ee7d41123f3c9911050ea2c2dac107568dc43b2d3b0c7557a33212c398ead30f", size = 166450, upload-time = "2023-11-25T09:07:26.339Z" } wheels = [ @@ -1364,7 +1452,7 @@ wheels = [ [[package]] name = "policyengine-core" -version = "3.32.9" +version = "3.32.11" source = { editable = "." } dependencies = [ { name = "dpath" }, @@ -1390,6 +1478,7 @@ dev = [ { name = "build" }, { name = "coverage" }, { name = "furo" }, + { name = "hypothesis" }, { name = "jupyter-book" }, { name = "markupsafe" }, { name = "mypy" }, @@ -1414,6 +1503,7 @@ requires-dist = [ { name = "furo", marker = "extra == 'dev'" }, { name = "h5py", specifier = ">=3,<4" }, { name = "huggingface-hub", specifier = ">=0.25.1" }, + { name = "hypothesis", marker = "extra == 'dev'", specifier = ">=6,<7" }, { name = "ipython", specifier = ">=8,<9" }, { name = "jupyter-book", marker = "extra == 'dev'", specifier = "<1" }, { name = "markupsafe", marker = "extra == 'dev'", specifier = "==2.0.1" }, From f8a19a842fc0418bd4cd920dc4b5f5b58443c16f Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Thu, 1 Oct 2026 15:50:54 -0400 Subject: [PATCH 2/9] Create the root's storage directory before branching in the mkdir regression test Co-Authored-By: Claude Opus 5.5 --- tests/core/test_disk_storage_ownership.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/core/test_disk_storage_ownership.py b/tests/core/test_disk_storage_ownership.py index 363cdde74..ebe2767d5 100644 --- a/tests/core/test_disk_storage_ownership.py +++ b/tests/core/test_disk_storage_ownership.py @@ -124,6 +124,8 @@ def test_branch_keeps_the_values_it_started_with(on_disk): def test_simulation_stores_on_disk_after_its_branch_storage_is_removed(): simulation = _simulation() + # Branches made once the directory exists share it. + simulation.data_storage_dir branch = simulation.get_branch("branch") branch.set_input("salary", PERIOD, np.array([1.0])) del branch From ecea197f99c939d90b9460520fa67253140474e3 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Thu, 1 Oct 2026 15:58:55 -0400 Subject: [PATCH 3/9] Test that derivative and apply_reform agree on disk and in memory Simulation.derivative perturbs a clone's input; with the shared directory it overwrote the simulation's own input file and returned 0. Co-Authored-By: Claude Opus 5.5 --- tests/core/test_disk_storage_ownership.py | 25 ++++++++++++++++++++++- 1 file changed, 24 insertions(+), 1 deletion(-) diff --git a/tests/core/test_disk_storage_ownership.py b/tests/core/test_disk_storage_ownership.py index ebe2767d5..16d794898 100644 --- a/tests/core/test_disk_storage_ownership.py +++ b/tests/core/test_disk_storage_ownership.py @@ -94,6 +94,18 @@ def test_clone_input_leaves_the_source_input(): 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 = _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 = _simulation() # Branches made after the root has stored on disk share its directory. @@ -369,6 +381,14 @@ def apply(self, operation: tuple): 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) @@ -431,6 +451,8 @@ def _outcome(calculate): 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) ), @@ -448,7 +470,8 @@ def _outcome(calculate): ) @given(operations=_operations) def test_disk_storage_gives_the_results_memory_storage_gives(operations): - """For any sequence of branching, cloning, inputs, calculations and drops: + """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; From 0adf21748ccc0e7b454c23874225e5322526dd85 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Thu, 1 Oct 2026 16:15:49 -0400 Subject: [PATCH 4/9] Move the Hypothesis property into its own module that skips without Hypothesis The country-package smoke job installs no dev dependencies, so a module-level hypothesis import failed its collection. Shared helpers move to tests/fixtures/disk_storage.py. Co-Authored-By: Claude Opus 5.5 --- tests/core/test_disk_storage_differential.py | 205 +++++++++++++++ tests/core/test_disk_storage_ownership.py | 249 ++----------------- tests/fixtures/disk_storage.py | 36 +++ 3 files changed, 256 insertions(+), 234 deletions(-) create mode 100644 tests/core/test_disk_storage_differential.py create mode 100644 tests/fixtures/disk_storage.py diff --git a/tests/core/test_disk_storage_differential.py b/tests/core/test_disk_storage_differential.py new file mode 100644 index 000000000..93d54be64 --- /dev/null +++ b/tests/core/test_disk_storage_differential.py @@ -0,0 +1,205 @@ +"""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.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 directory is 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) + + directory = on_disk.members[0].simulation._data_storage_dir + del in_memory, on_disk, member_in_memory, member_on_disk + gc.collect() + + assert errors == [] + assert directory is None or not os.path.exists(directory) diff --git a/tests/core/test_disk_storage_ownership.py b/tests/core/test_disk_storage_ownership.py index 16d794898..d78205f8d 100644 --- a/tests/core/test_disk_storage_ownership.py +++ b/tests/core/test_disk_storage_ownership.py @@ -3,63 +3,29 @@ Every ``OnDiskStorage`` writes into a directory of its own and reads the files it inherited from the storage it was cloned from. A directory stays on disk while any storage may read a file in it, and a file a clone reads is never -overwritten. These tests pin that for branches, clones and holders that a -branch creates itself, and check, as a property over random sequences of -operations, that storing values on disk gives the same results as storing -them in memory. +overwritten. These tests pin that for branches, clones, ``derivative`` and +holders that a branch creates itself; ``test_disk_storage_differential.py`` +checks it as a property over random sequences of operations. """ from __future__ import annotations -import contextlib import gc import os -import sys import warnings -from dataclasses import dataclass -from typing import List, Optional import numpy as np import pytest -from hypothesis import HealthCheck, given, settings -from hypothesis import strategies as st from policyengine_core.country_template import CountryTaxBenefitSystem 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.simulations import Simulation +from tests.fixtures.disk_storage import build_simulation, unraisable_exceptions PERIOD = "2017-01" -def _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 - - def _drop_branch(simulation: Simulation, name: str) -> None: del simulation.branches[name] gc.collect() @@ -67,7 +33,7 @@ def _drop_branch(simulation: Simulation, name: str) -> None: def test_branch_holder_leaves_the_parent_files(): """A holder a branch creates keeps its hands off the parent's files.""" - simulation = _simulation() + 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. @@ -84,7 +50,7 @@ def test_branch_holder_leaves_the_parent_files(): def test_clone_input_leaves_the_source_input(): """``Simulation.clone`` keeps the branch name, but not the files.""" - simulation = _simulation() + simulation = build_simulation() simulation.set_input("salary", PERIOD, np.array([1_000.0])) clone = simulation.clone() @@ -97,7 +63,7 @@ def test_clone_input_leaves_the_source_input(): @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 = _simulation(on_disk) + simulation = build_simulation(on_disk) simulation.set_input("salary", PERIOD, np.array([3_000.0])) derivative = simulation.derivative("income_tax", "salary", PERIOD, delta=100) @@ -107,7 +73,7 @@ def test_derivative_leaves_the_simulation_inputs(on_disk): def test_branches_with_the_same_name_keep_their_own_values(): - simulation = _simulation() + 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") @@ -123,7 +89,7 @@ def test_branches_with_the_same_name_keep_their_own_values(): @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 = _simulation(on_disk) + simulation = build_simulation(on_disk) simulation.set_input("salary", PERIOD, np.array([1_000.0])) branch = simulation.get_branch("branch") @@ -135,7 +101,7 @@ def test_branch_keeps_the_values_it_started_with(on_disk): def test_simulation_stores_on_disk_after_its_branch_storage_is_removed(): - simulation = _simulation() + simulation = build_simulation() # Branches made once the directory exists share it. simulation.data_storage_dir branch = simulation.get_branch("branch") @@ -149,7 +115,7 @@ def test_simulation_stores_on_disk_after_its_branch_storage_is_removed(): def test_branch_directory_goes_with_the_branch(): - simulation = _simulation() + simulation = build_simulation() simulation.set_input("salary", PERIOD, np.array([1_000.0])) branch = simulation.get_branch("branch") branch.set_input("salary", PERIOD, np.array([2.0])) @@ -166,8 +132,8 @@ def test_branch_directory_goes_with_the_branch(): def test_releasing_a_family_removes_its_directory_without_errors(): - with _unraisable_exceptions() as errors: - simulation = _simulation() + 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) @@ -186,7 +152,7 @@ def test_releasing_a_family_removes_its_directory_without_errors(): def test_a_data_storage_dir_set_by_the_user_stays(tmp_path): directory = tmp_path / "simulation" - simulation = _simulation() + simulation = build_simulation() simulation._data_storage_dir = str(directory) simulation.set_input("salary", PERIOD, np.array([1.0])) storage_dir = simulation.get_holder("salary")._disk_storage.storage_dir @@ -316,188 +282,3 @@ def test_a_forked_child_leaves_the_parent_directories(tmp_path): assert os.path.isdir(storage_dir) np.testing.assert_array_equal(storage.get(PERIOD), [1.0]) - - -# Differential property: on disk as in memory. - -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(_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, -) - - -@settings( - max_examples=75, - deadline=None, - suppress_health_check=[HealthCheck.too_slow], -) -@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 directory is 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) - - directory = on_disk.members[0].simulation._data_storage_dir - del in_memory, on_disk, member_in_memory, member_on_disk - gc.collect() - - assert errors == [] - assert directory is None or not os.path.exists(directory) 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 From 76839e96b061a0385b12fb832ae934a21fe7e6a1 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 02:00:37 -0400 Subject: [PATCH 5/9] Treat restored files as shared so storing their keys writes a new file Another storage, possibly in another process, may have restored the same directory and read those files. Co-Authored-By: Claude Opus 5.5 --- .../data_storage/on_disk_storage.py | 7 ++++++- tests/core/test_disk_storage_ownership.py | 20 +++++++++++++++++++ 2 files changed, 26 insertions(+), 1 deletion(-) diff --git a/policyengine_core/data_storage/on_disk_storage.py b/policyengine_core/data_storage/on_disk_storage.py index 3298ef04d..5263db820 100644 --- a/policyengine_core/data_storage/on_disk_storage.py +++ b/policyengine_core/data_storage/on_disk_storage.py @@ -274,7 +274,11 @@ def get_known_branch_periods(self) -> list: ] def restore(self) -> None: - """Map each key to its latest file in ``storage_dir``.""" + """Map each key to its latest file in ``storage_dir``. + + Other storages, including ones in other processes, may read the + restored files, so storing one of their keys again writes a new file. + """ self._files = files = {} if self.storage_dir is None: return @@ -287,3 +291,4 @@ def restore(self) -> None: versions[key] = version files[key] = os.path.join(self.storage_dir, filename) self._written_files.update(files.values()) + self._shared_files.update(files.values()) diff --git a/tests/core/test_disk_storage_ownership.py b/tests/core/test_disk_storage_ownership.py index d78205f8d..50f699095 100644 --- a/tests/core/test_disk_storage_ownership.py +++ b/tests/core/test_disk_storage_ownership.py @@ -222,6 +222,26 @@ def test_restore_reads_the_latest_file_of_each_key(tmp_path): np.testing.assert_array_equal(clone.get(PERIOD), [1.0]) +def test_storing_after_restore_leaves_the_restored_files(tmp_path): + """Another storage that restored the same directory keeps its values.""" + 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() + second = OnDiskStorage(str(storage_dir), preserve_storage_dir=True) + second.restore() + + first.put(np.array([2.0]), PERIOD) + + np.testing.assert_array_equal(second.get(PERIOD), [1.0]) + 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_storage_removes_only_its_own_directory(tmp_path): storage_dir = tmp_path / "parent" / "salary" storage_dir.mkdir(parents=True) From c7a036945f12b4587c0f44317bd821cdc20bc319 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 04:45:46 -0400 Subject: [PATCH 6/9] Count file versions with an int so storages copy and pickle on Python 3.14 Python 3.14 removed pickle/copy support from itertools objects, so itertools.count made copy.deepcopy and pickle of an OnDiskStorage raise. Found by the independent review. Co-Authored-By: Claude Opus 5.5 --- .../data_storage/on_disk_storage.py | 8 ++++--- tests/core/test_disk_storage_ownership.py | 23 +++++++++++++++++++ 2 files changed, 28 insertions(+), 3 deletions(-) diff --git a/policyengine_core/data_storage/on_disk_storage.py b/policyengine_core/data_storage/on_disk_storage.py index 5263db820..3d5b67a89 100644 --- a/policyengine_core/data_storage/on_disk_storage.py +++ b/policyengine_core/data_storage/on_disk_storage.py @@ -1,4 +1,3 @@ -import itertools import os import shutil import tempfile @@ -126,7 +125,9 @@ def __init__( # may read, which must not be overwritten. self._written_files: Set[str] = set() self._shared_files: Set[str] = set() - self._versions = itertools.count(1) + # The last number given to a ``{key}.{n}.npy`` file. A plain integer, + # not ``itertools.count``, so storages still copy and pickle on Python 3.14. + self._last_version = 0 @classmethod def temporary( @@ -217,7 +218,8 @@ def _get_path_to_write(self, key: str) -> str: storage_dir = self._get_storage_dir() path = os.path.join(storage_dir, f"{key}.npy") while path in self._shared_files: - path = os.path.join(storage_dir, f"{key}.{next(self._versions)}.npy") + self._last_version += 1 + path = os.path.join(storage_dir, f"{key}.{self._last_version}.npy") return path def put( diff --git a/tests/core/test_disk_storage_ownership.py b/tests/core/test_disk_storage_ownership.py index 50f699095..b88f1553d 100644 --- a/tests/core/test_disk_storage_ownership.py +++ b/tests/core/test_disk_storage_ownership.py @@ -10,8 +10,10 @@ from __future__ import annotations +import copy import gc import os +import pickle import warnings import numpy as np @@ -271,6 +273,27 @@ def test_preserved_storage_directory_stays(tmp_path, preserve_when): assert (storage_dir / f"default_{PERIOD}.npy").is_file() +@pytest.mark.parametrize("copy_storage", [copy.copy, copy.deepcopy, "pickle"]) +def test_copies_of_a_storage_read_its_values_and_keep_its_directory( + tmp_path, copy_storage +): + storage = OnDiskStorage.temporary("salary", str(tmp_path)) + storage.put(np.array([1.0]), PERIOD) + clone = storage.clone() + storage.put(np.array([2.0]), PERIOD) + if copy_storage == "pickle": + copied = pickle.loads(pickle.dumps(storage)) + else: + copied = copy_storage(storage) + + np.testing.assert_array_equal(copied.get(PERIOD), [2.0]) + del copied + gc.collect() + + assert os.path.isdir(storage.storage_dir) + np.testing.assert_array_equal(clone.get(PERIOD), [1.0]) + + def test_a_plain_array_stored_over_an_enum_reads_back_plain(tmp_path): storage = OnDiskStorage.temporary("status", str(tmp_path)) status = CountryTaxBenefitSystem().variables["housing_occupancy_status"] From a9a0b4057a084fa0f85a518a3647b3d80bf9eb35 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 07:33:08 -0400 Subject: [PATCH 7/9] Track live readers per file; create new files exclusively; preserve parents Round-3 review (REQUEST_CHANGES) found that a preserved storage inside the simulation's directory was removed with it, and that writing after restore() could overwrite an older clone's file or make a later restore read a stale value. The earlier attack script also showed permanent "shared" marks made every re-store after any clone (such as derivative's) add a file. - StorageDirectory keeps, per file, a WeakSet of the storages that may read it besides its creator (clones, copies, restorers); one handle per path per process. - A storage overwrites only the newest file it created itself, since the last fork, that no live reader may read (and, outside its own temporary directory, that is still the key's highest-numbered file). Otherwise it creates a new file with O_EXCL, numbered above every file for the key. - Preserving a directory preserves its parents; create_disk_storage(preserve= True) inside data_storage_dir takes the simulation's handle as parent. - copy/deepcopy return the same directory handle and pickles find the live one, so copies keep the directory; copies never overwrite the original's files. Co-Authored-By: Claude Opus 5.5 --- .../data_storage/on_disk_storage.py | 226 ++++++++++++++---- policyengine_core/holders/holder.py | 11 +- tests/core/test_disk_storage_ownership.py | 191 ++++++++++++++- 3 files changed, 377 insertions(+), 51 deletions(-) diff --git a/policyengine_core/data_storage/on_disk_storage.py b/policyengine_core/data_storage/on_disk_storage.py index 3d5b67a89..7581f577f 100644 --- a/policyengine_core/data_storage/on_disk_storage.py +++ b/policyengine_core/data_storage/on_disk_storage.py @@ -2,7 +2,7 @@ import shutil import tempfile import weakref -from typing import Dict, Optional, Set, Tuple, Union +from typing import Dict, Optional, Tuple, Union import numpy from numpy.typing import ArrayLike @@ -11,6 +11,25 @@ from policyengine_core.enums import EnumArray from policyengine_core.periods import Period +# Incremented in the parent and the child at every fork. A file created before +# a fork may be read by the other process, so neither overwrites it in place. +_fork_generation = 0 + + +def _after_fork() -> None: + global _fork_generation + _fork_generation += 1 + + +if hasattr(os, "register_at_fork"): + os.register_at_fork(after_in_parent=_after_fork, after_in_child=_after_fork) + +# The live handle for each directory path, so that every storage in this +# process that uses a directory shares one handle, and so one set of readers. +_live_directories: "weakref.WeakValueDictionary[str, StorageDirectory]" = ( + weakref.WeakValueDictionary() +) + def _remove_directory(path: str, pid: int) -> None: # A forked child inherits the finalizer, but the directory belongs to the @@ -19,19 +38,61 @@ def _remove_directory(path: str, pid: int) -> None: shutil.rmtree(path, ignore_errors=True) +def _create_exclusively(path: str) -> bool: + """Create an empty file at ``path``, or return False if one exists.""" + try: + os.close(os.open(path, os.O_CREAT | os.O_EXCL | os.O_WRONLY)) + except FileExistsError: + return False + return True + + +def _find_directory(path: str, pid: int) -> "StorageDirectory": + """Unpickle a directory handle. + + In the process that made it, this is the live handle, which keeps the + directory. Elsewhere it is a handle that never removes the directory. + """ + if os.getpid() == pid: + directory = _live_directories.get(path) + if directory is not None: + return directory + directory = StorageDirectory(path) + directory.preserve = True + return directory + + class StorageDirectory: """A directory removed, with its contents, once nothing references it. Whatever may read a file in the directory holds a reference to this object, so the directory stays on disk exactly as long as one of them is - alive. A directory created inside another holds that one as ``parent``. + alive. A directory created inside another holds that one as ``parent``, + and preserving a directory preserves its parents. + + ``readers`` records, for each file, the storages that may read it other + than the one that created it, so that storage knows when it may + overwrite the file. Copying a handle, with ``copy`` or ``deepcopy``, + returns the handle itself. """ def __init__(self, path: str, parent: Optional["StorageDirectory"] = None): - self.path = path + self.path = os.path.normpath(path) self.parent = parent + self.readers: Dict[str, "weakref.WeakSet[OnDiskStorage]"] = {} self._pid = os.getpid() - self._finalizer = weakref.finalize(self, _remove_directory, path, self._pid) + self._finalizer = weakref.finalize( + self, _remove_directory, self.path, self._pid + ) + _live_directories[self.path] = self + + @classmethod + def for_path(cls, path: str) -> "StorageDirectory": + """Return this process's handle for ``path``, creating one if needed.""" + directory = _live_directories.get(os.path.normpath(path)) + if directory is not None and directory._pid == os.getpid(): + return directory + return cls(path) @classmethod def create( @@ -62,11 +123,35 @@ def preserve(self) -> bool: def preserve(self, preserve: bool) -> None: if preserve: self._finalizer.detach() + if self.parent is not None: + self.parent.preserve = True elif not self._finalizer.alive: self._finalizer = weakref.finalize( self, _remove_directory, self.path, self._pid ) + def add_reader(self, path: str, storage: "OnDiskStorage") -> None: + self.readers.setdefault(path, weakref.WeakSet()).add(storage) + + def is_read(self, path: str) -> bool: + """Whether a live storage other than its creator may read ``path``.""" + readers = self.readers.get(path) + if readers is None: + return False + if not readers: + del self.readers[path] + return False + return True + + def __copy__(self) -> "StorageDirectory": + return self + + def __deepcopy__(self, memo: dict) -> "StorageDirectory": + return self + + def __reduce__(self): + return _find_directory, (self.path, self._pid) + def _parse_file_name(file_name: str) -> Tuple[str, int]: """Split a stored file's name, without ``.npy``, into its key and version. @@ -91,13 +176,15 @@ class OnDiskStorage: system temporary directory). A clone (:meth:`clone`) reads the files its source had stored when it was - cloned, and writes into a new directory of its own, so neither storage - can change or remove a value the other reads. Once a file has been - shared with a clone, storing its key again writes a new file instead of - overwriting it. + cloned, and writes into a new directory of its own. A storage overwrites + a file in place only if it created the file itself, since the last fork, + the file is still the key's newest, and no live clone, copy or storage + that restored the directory may read it. Otherwise it writes a new file, + created exclusively and numbered above every existing file for the key, + so :meth:`restore` finds each key's latest value. Unless ``preserve_storage_dir`` is set, a storage's directory is removed - once the storage and every clone that may read a file in it are gone. + once the storage and every storage that may read a file in it are gone. """ def __init__( @@ -114,20 +201,21 @@ def __init__( self._directory: Optional[StorageDirectory] = None self._name = "storage" if storage_dir is not None: - self._directory = StorageDirectory(storage_dir) - self._directory.preserve = preserve_storage_dir - self._name = os.path.basename(os.path.normpath(storage_dir)) + self._directory = StorageDirectory.for_path(storage_dir) + if preserve_storage_dir: + self._directory.preserve = True + self.storage_dir = self._directory.path + self._name = os.path.basename(self._directory.path) # Where ``storage_dir`` is created on the first write, if it is None. self._parent_directory: Union[StorageDirectory, str, None] = None # Other storages' directories that files in ``_files`` may be in. self._read_directories: Tuple[StorageDirectory, ...] = () - # Files this storage wrote or restored, and those of them a clone - # may read, which must not be overwritten. - self._written_files: Set[str] = set() - self._shared_files: Set[str] = set() - # The last number given to a ``{key}.{n}.npy`` file. A plain integer, - # not ``itertools.count``, so storages still copy and pickle on Python 3.14. - self._last_version = 0 + # The newest file this storage created for each key, and the fork + # generation it was created in: the only file it may overwrite. + self._created: Dict[str, Tuple[str, int]] = {} + # Whether the directory is this storage's own temporary one, which + # nothing restores. + self._temporary = False @classmethod def temporary( @@ -146,6 +234,7 @@ def temporary( storage = cls(None, is_eternal=is_eternal) storage._name = name storage._parent_directory = directory + storage._temporary = True return storage @property @@ -165,9 +254,9 @@ def clone(self) -> "OnDiskStorage": The clone gets copies of the file and enum mappings, so deleting or rewiring entries through one storage does not change the other. It reads the same ``.npy`` files, which stay on disk while it may read - them, and writes into a directory of its own, created on its first - write next to this storage's. From now on, this storage writes a new - file for a key rather than overwrite one the clone reads. + them and are not overwritten while it is alive, and writes into a + directory of its own, created on its first write next to this + storage's. """ clone = OnDiskStorage.temporary( self._name, self._parent_directory, is_eternal=self.is_eternal @@ -177,9 +266,33 @@ def clone(self) -> "OnDiskStorage": clone._read_directories = self._read_directories if self._directory is not None: clone._read_directories += (self._directory,) - self._shared_files.update(self._files.values()) + clone._register_as_reader() return clone + def _register_as_reader(self) -> None: + """Record this storage as a reader of every file it maps.""" + directories = { + directory.path: directory for directory in self._read_directories + } + if self._directory is not None: + directories[self._directory.path] = self._directory + for path in self._files.values(): + directory = directories.get(os.path.dirname(path)) + if directory is not None: + directory.add_reader(path, self) + + def __getstate__(self) -> dict: + state = self.__dict__.copy() + state["_files"] = dict(self._files) + state["_enums"] = dict(self._enums) + # A copy never overwrites a file the original created. + state["_created"] = {} + return state + + def __setstate__(self, state: dict) -> None: + self.__dict__.update(state) + self._register_as_reader() + def _decode_file(self, file: str) -> ArrayLike: enum = self._enums.get(file) if enum is not None: @@ -206,20 +319,53 @@ def _get_storage_dir(self) -> str: self.storage_dir = self._directory.path return self.storage_dir - def _get_path_to_write(self, key: str) -> str: - """Return the file to store a new value for ``key`` in. + def _newest_version(self, key: str) -> int: + """The highest version of any file for ``key`` in ``storage_dir``.""" + versions = [ + version + for file_key, version in ( + _parse_file_name(file_name[: -len(".npy")]) + for file_name in os.listdir(self.storage_dir) + if file_name.endswith(".npy") + ) + if file_key == key + ] + return max(versions, default=-1) - That is the key's current file or ``{key}.npy`` in this storage's - directory, unless a clone may read it; then a new numbered file. + def _create_file(self, key: str) -> str: + """Create an empty file for a new value of ``key`` and return its path. + + The file is ``{key}.npy`` if there is none yet, and otherwise + ``{key}.{n}.npy`` numbered above every file for the key. """ - current = self._files.get(key) - if current in self._written_files and current not in self._shared_files: - return current storage_dir = self._get_storage_dir() path = os.path.join(storage_dir, f"{key}.npy") - while path in self._shared_files: - self._last_version += 1 - path = os.path.join(storage_dir, f"{key}.{self._last_version}.npy") + if _create_exclusively(path): + return path + version = max(self._newest_version(key), 0) + while True: + version += 1 + path = os.path.join(storage_dir, f"{key}.{version}.npy") + if _create_exclusively(path): + return path + + def _may_overwrite(self, key: str, path: str, generation: int) -> bool: + if generation != _fork_generation or self._directory.is_read(path): + return False + if self._temporary: + return True + # In a directory another storage may write and restore, the file + # must still be the key's newest. + version = _parse_file_name(os.path.basename(path)[: -len(".npy")])[1] + return version >= self._newest_version(key) + + def _get_path_to_write(self, key: str) -> str: + """Return the file to store a new value for ``key`` in.""" + created = self._created.get(key) + if created is not None and self._may_overwrite(key, *created): + return created[0] + path = self._create_file(key) + self._created[key] = (path, _fork_generation) return path def put( @@ -237,12 +383,11 @@ def put( else: self._enums.pop(path, None) numpy.save(path, value) - self._written_files.add(path) self._files[filename] = path def delete(self, period: Period = None, branch_name: str = "default") -> None: # Deleting forgets the mapping only. The file stays on disk until the - # directory is removed, since a clone may still read it. + # directory is removed, since another storage may still read it. if period is None: # Only wipe files belonging to the requested branch (previously # this wiped every branch regardless of ``branch_name`` — same @@ -276,12 +421,14 @@ def get_known_branch_periods(self) -> list: ] def restore(self) -> None: - """Map each key to its latest file in ``storage_dir``. + """Map each key to its highest-numbered, and so latest, file in + ``storage_dir``. - Other storages, including ones in other processes, may read the - restored files, so storing one of their keys again writes a new file. + The storage becomes a reader of those files, and stores new values in + new files, since other storages may read the restored ones. """ self._files = files = {} + self._created = {} if self.storage_dir is None: return versions = {} @@ -292,5 +439,4 @@ def restore(self) -> None: if version >= versions.get(key, -1): versions[key] = version files[key] = os.path.join(self.storage_dir, filename) - self._written_files.update(files.values()) - self._shared_files.update(files.values()) + self._register_as_reader() diff --git a/policyengine_core/holders/holder.py b/policyengine_core/holders/holder.py index a3a5b2e95..4d0ed8798 100644 --- a/policyengine_core/holders/holder.py +++ b/policyengine_core/holders/holder.py @@ -8,6 +8,7 @@ from policyengine_core import commons, periods, tools from policyengine_core.data_storage import InMemoryStorage, OnDiskStorage +from policyengine_core.data_storage.on_disk_storage import StorageDirectory from policyengine_core.enums import Enum from policyengine_core.errors import PeriodMismatchError from policyengine_core.periods import Period @@ -99,16 +100,24 @@ def create_disk_storage( self.simulation._get_data_storage_directory(), is_eternal=is_eternal, ) + data_storage_directory = None if directory is None: directory = self.simulation.data_storage_dir + data_storage_directory = self.simulation._get_data_storage_directory() storage_dir = os.path.join(directory, self.variable.name) if not os.path.isdir(storage_dir): os.mkdir(storage_dir) - return OnDiskStorage( + storage = OnDiskStorage( storage_dir, is_eternal=is_eternal, preserve_storage_dir=preserve, ) + if isinstance(data_storage_directory, StorageDirectory): + # Keep the simulation's directory while the storage is alive, and + # after it too if the storage is preserved. + storage._directory.parent = data_storage_directory + storage.preserve_storage_dir = preserve + return storage def delete_arrays( self, period: Period = None, branch_name: str = "default" diff --git a/tests/core/test_disk_storage_ownership.py b/tests/core/test_disk_storage_ownership.py index b88f1553d..7ca8d884f 100644 --- a/tests/core/test_disk_storage_ownership.py +++ b/tests/core/test_disk_storage_ownership.py @@ -14,15 +14,19 @@ import gc import os import pickle +import shutil import warnings import numpy as np import pytest -from policyengine_core.country_template import CountryTaxBenefitSystem +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.simulations import Simulation +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" @@ -181,6 +185,73 @@ def test_a_key_stored_again_overwrites_its_file_until_a_clone_reads_it(tmp_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 len(os.listdir(storage.storage_dir)) == 2 + np.testing.assert_array_equal(storage.get(PERIOD), [5.0]) + + +def test_reforms_after_a_derivative_reuse_the_files(): + """Once ``derivative``'s clone is gone, recalculating overwrites in place.""" + 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 + + 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] + ) + + assert len(os.listdir(storage.storage_dir)) == 1 + + +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 + + for value in (2.0, 3.0, 4.0): + simulation.delete_arrays("salary", PERIOD) + simulation.set_input("salary", PERIOD, np.array([value])) + + assert len(os.listdir(storage.storage_dir)) == 2 + 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 = OnDiskStorage.temporary("salary", str(tmp_path)) @@ -244,6 +315,47 @@ def test_storing_after_restore_leaves_the_restored_files(tmp_path): np.testing.assert_array_equal(latest.get(PERIOD), [2.0]) +def test_a_write_after_restore_leaves_older_readers_and_is_restored_next(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) + clone = writer.clone() + writer.put(np.array([2.0]), PERIOD) + restorer = OnDiskStorage(str(storage_dir), preserve_storage_dir=True) + restorer.restore() + + 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), [2.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), [3.0]) + + +def test_two_storages_writing_after_restore_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() + + 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]) + latest = OnDiskStorage(str(storage_dir), preserve_storage_dir=True) + latest.restore() + np.testing.assert_array_equal(latest.get(PERIOD), [3.0]) + + def test_storage_removes_only_its_own_directory(tmp_path): storage_dir = tmp_path / "parent" / "salary" storage_dir.mkdir(parents=True) @@ -273,25 +385,58 @@ def test_preserved_storage_directory_stays(tmp_path, preserve_when): 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 + + del simulation, storage + gc.collect() + + try: + assert os.listdir(storage_dir) == [f"default_{PERIOD}.npy"] + finally: + shutil.rmtree(data_storage_dir, ignore_errors=True) + + @pytest.mark.parametrize("copy_storage", [copy.copy, copy.deepcopy, "pickle"]) -def test_copies_of_a_storage_read_its_values_and_keep_its_directory( - tmp_path, copy_storage -): +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 = OnDiskStorage.temporary("salary", str(tmp_path)) storage.put(np.array([1.0]), PERIOD) - clone = storage.clone() - storage.put(np.array([2.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]) - del copied + np.testing.assert_array_equal(storage.get(PERIOD), [3.0]) + + storage_dir = storage.storage_dir + del storage gc.collect() - assert os.path.isdir(storage.storage_dir) - np.testing.assert_array_equal(clone.get(PERIOD), [1.0]) + 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): @@ -325,3 +470,29 @@ def test_a_forked_child_leaves_the_parent_directories(tmp_path): assert os.path.isdir(storage_dir) np.testing.assert_array_equal(storage.get(PERIOD), [1.0]) + + +@pytest.mark.skipif(not hasattr(os, "fork"), reason="needs os.fork") +def test_a_forked_child_and_its_parent_keep_their_own_values(tmp_path): + storage = OnDiskStorage.temporary("salary", str(tmp_path)) + storage.put(np.array([1.0]), PERIOD) + parent_wrote_read, parent_wrote_write = os.pipe() + child_read_read, child_read_write = os.pipe() + + with warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + pid = os.fork() + if pid == 0: + try: + storage.put(np.array([2.0]), PERIOD) + os.read(parent_wrote_read, 1) + os.write(child_read_write, str(storage.get(PERIOD)[0]).encode()) + finally: + os._exit(0) + storage.put(np.array([3.0]), PERIOD) + os.write(parent_wrote_write, b"x") + child_value = float(os.read(child_read_read, 32).decode()) + os.waitpid(pid, 0) + + assert child_value == 2.0 + np.testing.assert_array_equal(storage.get(PERIOD), [3.0]) From 651ca466c8ff08bd7ae875615b1580df2d423bb1 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Wed, 7 Oct 2026 17:10:29 -0400 Subject: [PATCH 8/9] Isolate replacement holders within the #585 disk storage design --- .../disk-storage-branch-ownership.fixed.md | 2 +- docs/python_api/data_storage.md | 14 ++ .../data_storage/on_disk_storage.py | 2 + .../data_storage/storage_directory.py | 11 +- policyengine_core/holders/holder.py | 22 +- policyengine_core/simulations/simulation.py | 7 +- .../core/test_branch_scoped_delete_arrays.py | 7 +- tests/core/test_data_storage_dir.py | 24 ++- tests/core/test_data_storage_dir_property.py | 3 +- tests/core/test_disk_storage_differential.py | 26 ++- tests/core/test_disk_storage_ownership.py | 203 ++++++++++-------- 11 files changed, 200 insertions(+), 121 deletions(-) diff --git a/changelog.d/disk-storage-branch-ownership.fixed.md b/changelog.d/disk-storage-branch-ownership.fixed.md index b36b5e9fb..2605202f8 100644 --- a/changelog.d/disk-storage-branch-ownership.fixed.md +++ b/changelog.d/disk-storage-branch-ownership.fixed.md @@ -1 +1 @@ -Disk-backed holder storage (`MemoryConfig`) now gives every storage its own directory, keeps a directory while any clone reads its files and never overwrites a file a clone reads, so garbage-collecting a branch no longer deletes its parent's stored values and clones or same-named branches no longer overwrite each other's values. +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..fdf400b5c 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,9 @@ 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. + ```{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 eed852cd1..5db33dd90 100644 --- a/tests/core/test_branch_scoped_delete_arrays.py +++ b/tests/core/test_branch_scoped_delete_arrays.py @@ -84,7 +84,7 @@ def test_delete_arrays_purges_only_visible_branch_names(): def test_on_disk_storage_clone_copies_metadata_and_writes_elsewhere(tmp_path): - """A cloned disk view owns its metadata and writes into its own directory.""" + """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) @@ -102,9 +102,10 @@ def test_on_disk_storage_clone_copies_metadata_and_writes_elsewhere(tmp_path): clone.put(np.asarray([13.0]), PERIOD, "default") - assert clone.storage_dir != storage.storage_dir (clone_file,) = clone._files.values() - assert os.path.dirname(clone_file) == clone.storage_dir + (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]) 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 index 93d54be64..a19a5fa11 100644 --- a/tests/core/test_disk_storage_differential.py +++ b/tests/core/test_disk_storage_differential.py @@ -170,6 +170,22 @@ def _outcome(calculate): 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, @@ -177,7 +193,7 @@ def test_disk_storage_gives_the_results_memory_storage_gives(operations): - 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 directory is removed, and no + - once the simulations are gone, their directories are removed, and no finalizer raised. """ with unraisable_exceptions() as errors: @@ -197,9 +213,13 @@ def test_disk_storage_gives_the_results_memory_storage_gives(operations): lambda: member_in_memory.simulation.calculate(variable, period) ), (variable, period) - directory = on_disk.members[0].simulation._data_storage_dir + 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 is None or not os.path.exists(directory) + 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 index 7ca8d884f..420dead34 100644 --- a/tests/core/test_disk_storage_ownership.py +++ b/tests/core/test_disk_storage_ownership.py @@ -1,11 +1,10 @@ """Disk-backed holder storage keeps the values each simulation reads. -Every ``OnDiskStorage`` writes into a directory of its own and reads the files -it inherited from the storage it was cloned from. A directory stays on disk -while any storage may read a file in it, and a file a clone reads is never -overwritten. These tests pin that for branches, clones, ``derivative`` and -holders that a branch creates itself; ``test_disk_storage_differential.py`` -checks it as a property over random sequences of operations. +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 @@ -15,7 +14,7 @@ import os import pickle import shutil -import warnings +import tempfile import numpy as np import pytest @@ -32,6 +31,10 @@ 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() @@ -55,7 +58,7 @@ def test_branch_holder_leaves_the_parent_files(): def test_clone_input_leaves_the_source_input(): - """``Simulation.clone`` keeps the branch name, but not the files.""" + """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() @@ -108,7 +111,7 @@ def test_branch_keeps_the_values_it_started_with(on_disk): def test_simulation_stores_on_disk_after_its_branch_storage_is_removed(): simulation = build_simulation() - # Branches made once the directory exists share it. + # 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])) @@ -120,15 +123,16 @@ def test_simulation_stores_on_disk_after_its_branch_storage_is_removed(): np.testing.assert_array_equal(simulation.calculate("salary", PERIOD), [2.0]) -def test_branch_directory_goes_with_the_branch(): +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") - branch.set_input("salary", PERIOD, np.array([2.0])) + # 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("salary")._disk_storage.storage_dir + branch_dir = branch.get_holder("income_tax")._disk_storage.storage_dir assert branch_dir != root_dir - assert os.path.dirname(branch_dir) == simulation.data_storage_dir del branch _drop_branch(simulation, "branch") @@ -156,32 +160,81 @@ def test_releasing_a_family_removes_its_directory_without_errors(): assert not os.path.exists(directory) -def test_a_data_storage_dir_set_by_the_user_stays(tmp_path): +@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) - simulation.set_input("salary", PERIOD, np.array([1.0])) - storage_dir = simulation.get_holder("salary")._disk_storage.storage_dir - assert os.path.dirname(storage_dir) == str(directory) + storer = simulation.clone() if clone_first else simulation + assert not directory.exists() - del simulation + # 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 = OnDiskStorage.temporary("salary", str(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 len(os.listdir(storage.storage_dir)) == 1 + 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 len(os.listdir(storage.storage_dir)) == 2 + 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]) @@ -189,12 +242,12 @@ def test_a_key_stored_again_overwrites_its_file_until_a_clone_reads_it(tmp_path) gc.collect() storage.put(np.array([5.0]), PERIOD) - assert len(os.listdir(storage.storage_dir)) == 2 + 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(): - """Once ``derivative``'s clone is gone, recalculating overwrites in place.""" + """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) @@ -202,13 +255,17 @@ def test_reforms_after_a_derivative_reuse_the_files(): 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] ) - - assert len(os.listdir(storage.storage_dir)) == 1 + 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(): @@ -216,12 +273,16 @@ def test_a_live_branch_adds_at_most_one_file_per_key(): 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(os.listdir(storage.storage_dir)) == 2 + 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]) @@ -254,7 +315,7 @@ class carried(Variable): def test_a_clone_keeps_the_directories_it_reads(tmp_path): - storage = OnDiskStorage.temporary("salary", str(tmp_path)) + storage = _storage(tmp_path) storage.put(np.array([1.0]), PERIOD) storage_dir = storage.storage_dir grandchild = storage.clone().clone() @@ -271,7 +332,8 @@ def test_a_clone_keeps_the_directories_it_reads(tmp_path): assert not os.path.exists(storage_dir) -def test_restore_reads_the_latest_file_of_each_key(tmp_path): +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) @@ -289,53 +351,50 @@ def test_restore_reads_the_latest_file_of_each_key(tmp_path): "default_2017-01", "default_2017-02", ] - np.testing.assert_array_equal(restored.get(PERIOD), [2.0]) + 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_leaves_the_restored_files(tmp_path): - """Another storage that restored the same directory keeps its values.""" +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() - second = OnDiskStorage(str(storage_dir), preserve_storage_dir=True) - second.restore() first.put(np.array([2.0]), PERIOD) - np.testing.assert_array_equal(second.get(PERIOD), [1.0]) 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_older_readers_and_is_restored_next(tmp_path): +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) - clone = writer.clone() - writer.put(np.array([2.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), [2.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), [3.0]) + np.testing.assert_array_equal(latest.get(PERIOD), [1.0]) -def test_two_storages_writing_after_restore_keep_their_own_values(tmp_path): +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( @@ -345,15 +404,21 @@ def test_two_storages_writing_after_restore_keep_their_own_values(tmp_path): 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), [3.0]) + np.testing.assert_array_equal(latest.get(PERIOD), [1.0]) def test_storage_removes_only_its_own_directory(tmp_path): @@ -406,12 +471,14 @@ def test_preserving_a_holder_storage_keeps_it_after_the_family_goes(): 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.listdir(storage_dir) == [f"default_{PERIOD}.npy"] + 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) @@ -419,7 +486,7 @@ def test_preserving_a_holder_storage_keeps_it_after_the_family_goes(): @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 = OnDiskStorage.temporary("salary", str(tmp_path)) + storage = _storage(tmp_path) storage.put(np.array([1.0]), PERIOD) if copy_storage == "pickle": copied = pickle.loads(pickle.dumps(storage)) @@ -440,7 +507,7 @@ def test_a_copy_of_a_storage_is_independent_of_it(tmp_path, copy_storage): def test_a_plain_array_stored_over_an_enum_reads_back_plain(tmp_path): - storage = OnDiskStorage.temporary("status", str(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) @@ -448,51 +515,3 @@ def test_a_plain_array_stored_over_an_enum_reads_back_plain(tmp_path): storage.put(np.array([3]), PERIOD) assert not isinstance(storage.get(PERIOD), EnumArray) - - -@pytest.mark.skipif(not hasattr(os, "fork"), reason="needs os.fork") -def test_a_forked_child_leaves_the_parent_directories(tmp_path): - storage = OnDiskStorage.temporary("salary", str(tmp_path)) - storage.put(np.array([1.0]), PERIOD) - storage_dir = storage.storage_dir - - with warnings.catch_warnings(): - # Python warns that forking a multi-threaded process may deadlock. - warnings.simplefilter("ignore", DeprecationWarning) - pid = os.fork() - if pid == 0: - try: - del storage - gc.collect() - finally: - os._exit(0) - os.waitpid(pid, 0) - - assert os.path.isdir(storage_dir) - np.testing.assert_array_equal(storage.get(PERIOD), [1.0]) - - -@pytest.mark.skipif(not hasattr(os, "fork"), reason="needs os.fork") -def test_a_forked_child_and_its_parent_keep_their_own_values(tmp_path): - storage = OnDiskStorage.temporary("salary", str(tmp_path)) - storage.put(np.array([1.0]), PERIOD) - parent_wrote_read, parent_wrote_write = os.pipe() - child_read_read, child_read_write = os.pipe() - - with warnings.catch_warnings(): - warnings.simplefilter("ignore", DeprecationWarning) - pid = os.fork() - if pid == 0: - try: - storage.put(np.array([2.0]), PERIOD) - os.read(parent_wrote_read, 1) - os.write(child_read_write, str(storage.get(PERIOD)[0]).encode()) - finally: - os._exit(0) - storage.put(np.array([3.0]), PERIOD) - os.write(parent_wrote_write, b"x") - child_value = float(os.read(child_read_read, 32).decode()) - os.waitpid(pid, 0) - - assert child_value == 2.0 - np.testing.assert_array_equal(storage.get(PERIOD), [3.0]) From c616364cd65d8a79ef10bbe275e94b91ef6167b4 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Thu, 8 Oct 2026 19:26:31 -0400 Subject: [PATCH 9/9] Document cleanup of abandoned holder storage directories --- docs/python_api/data_storage.md | 1 + 1 file changed, 1 insertion(+) diff --git a/docs/python_api/data_storage.md b/docs/python_api/data_storage.md index fdf400b5c..cfaf5ba4d 100644 --- a/docs/python_api/data_storage.md +++ b/docs/python_api/data_storage.md @@ -35,6 +35,7 @@ writing when its existing values must remain isolated. 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