From 9e03c929a222176ee585571b14d4af4b4c2075a0 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 17:29:11 -0400 Subject: [PATCH 1/5] Give a storage a set of shared keys only while it shares an array 3.32.12 (#556) created an empty set in every InMemoryStorage. A simulation has a storage for every variable of the tax-benefit system, so that was about 224 bytes per variable per simulation: 1.3 MB for each policyengine-us simulation, branches included. The YAML test runner keeps every case's simulation, and policyengine-us's largest CI batches grew by 2.1 to 2.6 GB and ran out of memory on 16 GB runners. A storage that shares nothing now refers to one empty frozenset, and has a set of its own only while at least one key is shared. Fixes #577 Co-Authored-By: Claude Opus 5.5 --- changelog.d/lazy-shared-array-index.fixed.md | 1 + .../data_storage/in_memory_storage.py | 48 +++- tests/core/test_storage_shared_index.py | 204 +++++++++++++++ .../test_storage_shared_index_differential.py | 235 ++++++++++++++++++ 4 files changed, 477 insertions(+), 11 deletions(-) create mode 100644 changelog.d/lazy-shared-array-index.fixed.md create mode 100644 tests/core/test_storage_shared_index.py create mode 100644 tests/core/test_storage_shared_index_differential.py diff --git a/changelog.d/lazy-shared-array-index.fixed.md b/changelog.d/lazy-shared-array-index.fixed.md new file mode 100644 index 000000000..9fc21a326 --- /dev/null +++ b/changelog.d/lazy-shared-array-index.fixed.md @@ -0,0 +1 @@ +Stop giving every holder's in-memory storage an empty set of shared keys: a storage now has one only while it shares an array with the storage it was cloned from, which removes about 1.3 MB per simulation for a country with 6,000 variables (the 3.32.12 increase that pushed policyengine-us YAML test batches past their CI runners' memory). diff --git a/policyengine_core/data_storage/in_memory_storage.py b/policyengine_core/data_storage/in_memory_storage.py index a1bddf76b..2d60ce033 100644 --- a/policyengine_core/data_storage/in_memory_storage.py +++ b/policyengine_core/data_storage/in_memory_storage.py @@ -1,4 +1,4 @@ -from typing import Dict, Union +from typing import AbstractSet, Dict, Union import numpy from numpy.typing import ArrayLike @@ -32,21 +32,30 @@ def _read_only_view(array: numpy.ndarray) -> numpy.ndarray: return view +# What ``InMemoryStorage._shared`` is while a storage shares nothing. Every +# such storage refers to this one object instead of holding an empty set of +# its own: a simulation has a storage for each variable, nearly all of them +# share nothing, and an empty set is 216 bytes. +_NOTHING_SHARED: AbstractSet[str] = frozenset() + + class InMemoryStorage: """ Low-level class responsible for storing and retrieving calculated vectors in memory """ _arrays: Dict[Period, ArrayLike] + # Keys of ``_arrays`` whose array still belongs to the storage this one + # was cloned from with ``share_arrays``. ``get`` replaces each with a copy + # the first time it is read. It is ``_NOTHING_SHARED`` unless at least one + # key is shared, and a set of this storage's own otherwise. A key left + # here after code outside this class empties ``_arrays`` costs one extra + # copy at most. + _shared: AbstractSet[str] = _NOTHING_SHARED is_eternal: bool def __init__(self, is_eternal: bool): self._arrays = {} - # Keys of ``_arrays`` whose array still belongs to the storage this - # one was cloned from with ``share_arrays``. ``get`` replaces each - # with a copy the first time it is read. A key left here after code - # outside this class empties ``_arrays`` costs one extra copy at most. - self._shared = set() self.is_eternal = is_eternal def clone(self, share_arrays: bool = False) -> "InMemoryStorage": @@ -75,12 +84,15 @@ def clone(self, share_arrays: bool = False) -> "InMemoryStorage": """ clone = InMemoryStorage(self.is_eternal) if share_arrays: + shared = set() for key, array in self._arrays.items(): if _can_share(array): clone._arrays[key] = _read_only_view(array) - clone._shared.add(key) + shared.add(key) else: clone._arrays[key] = array.copy() + if shared: + clone._shared = shared else: clone._arrays = {key: array.copy() for key, array in self._arrays.items()} return clone @@ -98,9 +110,23 @@ def get(self, period: Period, branch_name: str = "default") -> ArrayLike: # caller may write into it, so it needs to be this storage's own. values = values.copy() self._arrays[key] = values - self._shared.discard(key) + self._stop_sharing(key) return values + def _stop_sharing(self, key: str) -> None: + """Record that ``key`` no longer refers to a shared array.""" + if key in self._shared: + self._shared.discard(key) + if not self._shared: + self._shared = _NOTHING_SHARED + + def _stop_sharing_dropped_keys(self) -> None: + """Forget the shared keys that ``_arrays`` no longer has.""" + if self._shared: + self._shared.intersection_update(self._arrays) + if not self._shared: + self._shared = _NOTHING_SHARED + def put( self, value: ArrayLike, period: Period, branch_name: str = "default" ) -> None: @@ -125,7 +151,7 @@ def put( ) key = f"{branch_name}:{period}" self._arrays[key] = value - self._shared.discard(key) + self._stop_sharing(key) def delete(self, period: Period = None, branch_name: str = "default") -> None: if period is None: @@ -137,7 +163,7 @@ def delete(self, period: Period = None, branch_name: str = "default") -> None: for period_item, value in self._arrays.items() if not period_item.startswith(branch_prefix) } - self._shared.intersection_update(self._arrays) + self._stop_sharing_dropped_keys() return if self.is_eternal: @@ -155,7 +181,7 @@ def delete(self, period: Period = None, branch_name: str = "default") -> None: and period.contains(periods.period(period_item.split(":", 1)[1])) ) } - self._shared.intersection_update(self._arrays) + self._stop_sharing_dropped_keys() def get_known_periods(self) -> list: # Split on the first colon only: an anchored period's string form diff --git a/tests/core/test_storage_shared_index.py b/tests/core/test_storage_shared_index.py new file mode 100644 index 000000000..a2c077c74 --- /dev/null +++ b/tests/core/test_storage_shared_index.py @@ -0,0 +1,204 @@ +"""A storage holds a set of shared keys only while it shares an array. + +``InMemoryStorage._shared`` records which arrays a storage still shares with +the one it was cloned from (see ``InMemoryStorage.clone``). When it was added, +every storage was given an empty set of its own. A simulation has a storage +for every variable of the tax-benefit system, whether or not the variable is +ever read, so that was one 216-byte set per variable per simulation: 1.3 MB +for each simulation of a country with 6,000 variables, and the YAML test +runner keeps thousands of simulations alive. + +A storage that shares nothing now refers to one empty object for all +storages, and has a set of its own only while at least one key is shared. +``test_storage_shared_index_differential.py`` checks random sequences of +operations against the storage that always had its own set. +""" + +from __future__ import annotations + +import numpy as np + +from policyengine_core import periods +from policyengine_core.data_storage import InMemoryStorage +from policyengine_core.data_storage.in_memory_storage import _NOTHING_SHARED +from tests.fixtures.branch_shared_arrays import build_simulation + +JANUARY = periods.period("2017-01") + + +def _storages(simulation): + return { + name: holder._memory_storage + for population in simulation.populations.values() + for name, holder in population._holders.items() + } + + +def _storage_with(*months): + storage = InMemoryStorage(is_eternal=False) + for index, month in enumerate(months): + storage.put(np.array([float(index), 1.0]), month) + return storage + + +def test_the_shared_nothing_index_is_empty_and_cannot_change(): + assert isinstance(_NOTHING_SHARED, frozenset) + assert not _NOTHING_SHARED + + +def test_new_storages_have_no_set_of_their_own(): + first = InMemoryStorage(is_eternal=False) + second = InMemoryStorage(is_eternal=True) + + assert first._shared is _NOTHING_SHARED + assert second._shared is _NOTHING_SHARED + + # Storing, reading and deleting without sharing never needs one. + first.put(np.array([1.0]), "2017-01") + first.get("2017-01") + first.delete("2017-01") + first.delete() + assert first._shared is _NOTHING_SHARED + + +def test_copied_storage_has_no_set_of_its_own(): + storage = _storage_with("2017-01", "2017-02") + + assert storage.clone()._shared is _NOTHING_SHARED + assert storage.clone(share_arrays=True).clone()._shared is _NOTHING_SHARED + + +def test_sharing_clone_of_an_empty_storage_has_no_set_of_its_own(): + storage = InMemoryStorage(is_eternal=False) + + assert storage.clone(share_arrays=True)._shared is _NOTHING_SHARED + + +def test_sharing_clone_with_only_copied_arrays_has_no_set_of_its_own(): + storage = InMemoryStorage(is_eternal=False) + storage.put(np.ma.masked_array([1.0, 2.0], mask=[False, True]), "2017-01") + + clone = storage.clone(share_arrays=True) + + # A masked array is copied straight away, so nothing is shared. + assert clone._shared is _NOTHING_SHARED + assert not np.shares_memory(clone._arrays["default:2017-01"], storage.get(JANUARY)) + + +def test_sharing_clones_each_have_their_own_set(): + storage = _storage_with("2017-01", "2017-02") + + first = storage.clone(share_arrays=True) + second = storage.clone(share_arrays=True) + + assert first._shared == second._shared == {"default:2017-01", "default:2017-02"} + assert first._shared is not second._shared + assert storage._shared is _NOTHING_SHARED + + first.get("2017-01") + assert first._shared == {"default:2017-02"} + assert second._shared == {"default:2017-01", "default:2017-02"} + + +def test_set_is_released_once_every_shared_array_has_been_read(): + clone = _storage_with("2017-01", "2017-02").clone(share_arrays=True) + + clone.get("2017-01") + assert clone._shared == {"default:2017-02"} + clone.get("2017-02") + + assert clone._shared is _NOTHING_SHARED + + +def test_set_is_released_once_every_shared_array_has_been_replaced(): + clone = _storage_with("2017-01").clone(share_arrays=True) + + clone.put(np.array([5.0, 6.0]), "2017-01") + + assert clone._shared is _NOTHING_SHARED + + +def test_set_is_released_once_every_shared_array_has_been_deleted(): + by_period = _storage_with("2017-01", "2017-02").clone(share_arrays=True) + by_period.delete("2017") + assert by_period._shared is _NOTHING_SHARED + + by_branch = _storage_with("2017-01", "2017-02").clone(share_arrays=True) + by_branch.delete() + assert by_branch._shared is _NOTHING_SHARED + + +def test_deleting_another_branch_keeps_what_is_still_shared(): + storage = _storage_with("2017-01") + storage.put(np.array([7.0, 8.0]), "2017-01", "other") + clone = storage.clone(share_arrays=True) + + clone.delete(branch_name="other") + + assert clone._shared == {"default:2017-01"} + + +def test_storage_keeps_working_after_releasing_its_set(): + storage = _storage_with("2017-01") + clone = storage.clone(share_arrays=True) + clone.get("2017-01") + assert clone._shared is _NOTHING_SHARED + + own = np.array([3.0, 4.0]) + clone.put(own, "2017-02") + assert clone.get("2017-02") is own + clone.delete("2017-02") + assert clone.get("2017-02") is None + + # A clone of it shares again, with a set of its own. + grandchild = clone.clone(share_arrays=True) + assert grandchild._shared == {"default:2017-01"} + assert clone._shared is _NOTHING_SHARED + + +def test_simulation_storages_have_no_sets_of_their_own(tax_benefit_system): + simulation = build_simulation(tax_benefit_system) + simulation.calculate("disposable_income", JANUARY) + + storages = _storages(simulation) + # One storage per variable, as in every simulation. + assert len(storages) == len(tax_benefit_system.variables) + assert all(storage._shared is _NOTHING_SHARED for storage in storages.values()) + + copy = simulation.clone() + assert all( + storage._shared is _NOTHING_SHARED for storage in _storages(copy).values() + ) + + +def test_branch_storages_have_a_set_only_for_the_arrays_they_share( + tax_benefit_system, +): + simulation = build_simulation(tax_benefit_system) + simulation.calculate("disposable_income", JANUARY) + with_arrays = { + name for name, storage in _storages(simulation).items() if storage._arrays + } + assert with_arrays and with_arrays < set(_storages(simulation)) + + branch = simulation.get_branch("branch") + + sets = [] + for name, storage in _storages(branch).items(): + if name in with_arrays: + assert storage._shared == set(storage._arrays), name + sets.append(storage._shared) + else: + assert storage._shared is _NOTHING_SHARED, name + # No two storages hold the same set. + assert len({id(shared) for shared in sets}) == len(sets) + + # Reading through the branch copies what it reads and releases the sets + # of the storages that end up sharing nothing. + branch.calculate("disposable_income", JANUARY) + for name, storage in _storages(branch).items(): + if not storage._shared: + assert storage._shared is _NOTHING_SHARED, name + assert all( + storage._shared is _NOTHING_SHARED for storage in _storages(simulation).values() + ) diff --git a/tests/core/test_storage_shared_index_differential.py b/tests/core/test_storage_shared_index_differential.py new file mode 100644 index 000000000..463ed207f --- /dev/null +++ b/tests/core/test_storage_shared_index_differential.py @@ -0,0 +1,235 @@ +"""A storage without a set of its own behaves like one that always had one. + +``_EagerIndexStorage`` below is ``InMemoryStorage`` as it was when every +storage owned a set of shared keys. Random sequences of operations run on +both, and after every step the two must hold the same arrays and the same +shared keys; at the end every read must agree. The storage under test must +also keep its own invariants: it refers to the one shared-nothing object +exactly when it shares nothing, and no two storages hold the same set. +``test_storage_shared_index.py`` pins the same behaviour with examples. +""" + +from __future__ import annotations + +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 import periods +from policyengine_core.data_storage import InMemoryStorage +from policyengine_core.data_storage.in_memory_storage import ( + _NOTHING_SHARED, + _can_share, + _read_only_view, +) + + +class _EagerIndexStorage: + """The storage with an empty set of shared keys in every instance.""" + + def __init__(self, is_eternal): + self._arrays = {} + self._shared = set() + self.is_eternal = is_eternal + + def clone(self, share_arrays=False): + clone = _EagerIndexStorage(self.is_eternal) + if share_arrays: + for key, array in self._arrays.items(): + if _can_share(array): + clone._arrays[key] = _read_only_view(array) + clone._shared.add(key) + else: + clone._arrays[key] = array.copy() + else: + clone._arrays = {key: array.copy() for key, array in self._arrays.items()} + return clone + + def get(self, period, branch_name="default"): + if self.is_eternal: + period = periods.period(periods.ETERNITY) + period = periods.period(period) + key = f"{branch_name}:{period}" + values = self._arrays.get(key) + if values is None: + return None + if key in self._shared: + values = values.copy() + self._arrays[key] = values + self._shared.discard(key) + return values + + def put(self, value, period, branch_name="default"): + if self.is_eternal: + period = periods.period(periods.ETERNITY) + period = periods.period(period) + key = f"{branch_name}:{period}" + self._arrays[key] = value + self._shared.discard(key) + + def delete(self, period=None, branch_name="default"): + if period is None: + branch_prefix = f"{branch_name}:" + self._arrays = { + period_item: value + for period_item, value in self._arrays.items() + if not period_item.startswith(branch_prefix) + } + self._shared.intersection_update(self._arrays) + return + if self.is_eternal: + period = periods.period(periods.ETERNITY) + period = periods.period(period) + self._arrays = { + period_item: value + for period_item, value in self._arrays.items() + if not ( + period_item.startswith(f"{branch_name}:") + and period.contains(periods.period(period_item.split(":", 1)[1])) + ) + } + self._shared.intersection_update(self._arrays) + + +PERIODS = ["2017-01", "2017-02", "2017"] +BRANCHES = ["default", "other"] +ARRAY_KINDS = ["plain", "read_only", "masked"] + + +def _make_array(kind, seed): + values = np.array([seed, seed + 1.0, seed + 2.0]) + if kind == "read_only": + values.flags.writeable = False + if kind == "masked": + return np.ma.masked_array(values, mask=[False, True, False]) + return values + + +_storage = st.integers(min_value=0, max_value=10**6) +_period = st.sampled_from(PERIODS) +_branch = st.sampled_from(BRANCHES) +_seed = st.integers(min_value=0, max_value=50).map(float) + +_operation = st.one_of( + st.tuples( + st.just("put"), + _storage, + _period, + _branch, + st.sampled_from(ARRAY_KINDS), + _seed, + ), + st.tuples(st.just("get"), _storage, _period, _branch), + st.tuples(st.just("write"), _storage, _period, _branch, _seed), + st.tuples(st.just("delete"), _storage, st.none() | _period, _branch), + st.tuples(st.just("clone"), _storage, st.booleans()), + st.tuples(st.just("clear"), _storage), +) + + +def _same_array(tested, reference): + if tested is None or reference is None: + return tested is None and reference is None + if type(tested) is not type(reference) or tested.dtype != reference.dtype: + return False + if isinstance(tested, np.ma.MaskedArray): + return np.array_equal(tested.mask, reference.mask) and np.array_equal( + tested.data, reference.data + ) + return np.array_equal(tested, reference) + + +def _apply(storages, operation): + """Run one operation on a list of storages and return what it read.""" + kind, index = operation[0], operation[1] % len(storages) + storage = storages[index] + if kind == "put": + _, _, period, branch, array_kind, seed = operation + storage.put(_make_array(array_kind, seed), period, branch) + return None + if kind == "get": + _, _, period, branch = operation + return storage.get(period, branch) + if kind == "write": + # Write in place into what a read returned, as formulas do. + _, _, period, branch, seed = operation + values = storage.get(period, branch) + if values is not None and values.flags.writeable: + values[0] = seed + return values + if kind == "delete": + _, _, period, branch = operation + storage.delete(period, branch) + return None + if kind == "clone": + storages.append(storage.clone(share_arrays=operation[2])) + return None + if kind == "clear": + # policyengine-uk clears a dropped branch's storage like this. + storage._arrays.clear() + return None + raise AssertionError(kind) + + +def _assert_same_state(tested, reference): + assert len(tested) == len(reference) + for index, (storage, expected) in enumerate(zip(tested, reference)): + assert storage._arrays.keys() == expected._arrays.keys(), index + for key, array in storage._arrays.items(): + assert _same_array(array, expected._arrays[key]), (index, key) + assert storage._shared == expected._shared, index + + +def _assert_index_invariants(storages, cleared): + own_sets = [] + for index, storage in enumerate(storages): + shared = storage._shared + # The one shared-nothing object, exactly when nothing is shared. + assert (shared is _NOTHING_SHARED) == (not shared), index + if shared is not _NOTHING_SHARED: + own_sets.append(shared) + if index not in cleared: + assert shared <= storage._arrays.keys(), index + for key in shared & storage._arrays.keys(): + # A shared array cannot be written through. + assert not storage._arrays[key].flags.writeable, (index, key) + # A set belongs to one storage. + assert len({id(shared) for shared in own_sets}) == len(own_sets) + assert not _NOTHING_SHARED + + +@hypothesis.settings( + max_examples=500, + deadline=None, + suppress_health_check=[hypothesis.HealthCheck.too_slow], +) +@hypothesis.given( + is_eternal=st.booleans(), + operations=st.lists(_operation, max_size=60), +) +def test_storage_without_its_own_set_matches_storage_with_one(is_eternal, operations): + tested = [InMemoryStorage(is_eternal)] + reference = [_EagerIndexStorage(is_eternal)] + cleared = set() + + for step, operation in enumerate(operations): + read = _apply(tested, operation) + expected = _apply(reference, operation) + assert _same_array(read, expected), (step, operation) + if operation[0] == "clear": + cleared.add(operation[1] % len(tested)) + _assert_same_state(tested, reference) + _assert_index_invariants(tested, cleared) + + # Every storage reads the same value for everything. + for index, (storage, expected) in enumerate(zip(tested, reference)): + for branch in BRANCHES: + for period in PERIODS: + assert _same_array( + storage.get(period, branch), expected.get(period, branch) + ), (index, branch, period) + _assert_same_state(tested, reference) + _assert_index_invariants(tested, cleared) From 028ad7494d4117a4be7c2183c37ed55e1ebf2cd8 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 17:36:56 -0400 Subject: [PATCH 2/5] Release the set by deleting the attribute; test copies, pickles and old pickles A storage that shares nothing now has no _shared attribute of its own and reads the class's shared-nothing object, so a deep copy or a pickle of it does too (review finding). Adds tests for deepcopy and pickle round trips, for a storage pickled by 3.32.12 with an empty set, and the property that a storage has the attribute exactly when it shares something. Co-Authored-By: Claude Opus 5.5 --- .../data_storage/in_memory_storage.py | 30 ++++++---- tests/core/test_storage_shared_index.py | 57 +++++++++++++++++++ .../test_storage_shared_index_differential.py | 4 +- 3 files changed, 79 insertions(+), 12 deletions(-) diff --git a/policyengine_core/data_storage/in_memory_storage.py b/policyengine_core/data_storage/in_memory_storage.py index 2d60ce033..2d64ed4f7 100644 --- a/policyengine_core/data_storage/in_memory_storage.py +++ b/policyengine_core/data_storage/in_memory_storage.py @@ -1,4 +1,4 @@ -from typing import AbstractSet, Dict, Union +from typing import Dict, FrozenSet, Set, Union import numpy from numpy.typing import ArrayLike @@ -36,7 +36,7 @@ def _read_only_view(array: numpy.ndarray) -> numpy.ndarray: # such storage refers to this one object instead of holding an empty set of # its own: a simulation has a storage for each variable, nearly all of them # share nothing, and an empty set is 216 bytes. -_NOTHING_SHARED: AbstractSet[str] = frozenset() +_NOTHING_SHARED: FrozenSet[str] = frozenset() class InMemoryStorage: @@ -47,11 +47,11 @@ class InMemoryStorage: _arrays: Dict[Period, ArrayLike] # Keys of ``_arrays`` whose array still belongs to the storage this one # was cloned from with ``share_arrays``. ``get`` replaces each with a copy - # the first time it is read. It is ``_NOTHING_SHARED`` unless at least one - # key is shared, and a set of this storage's own otherwise. A key left - # here after code outside this class empties ``_arrays`` costs one extra - # copy at most. - _shared: AbstractSet[str] = _NOTHING_SHARED + # the first time it is read. A storage has a ``_shared`` attribute of its + # own, a set, only while at least one key is shared; otherwise it reads + # this class attribute. A key left in the set after code outside this + # class empties ``_arrays`` costs one extra copy at most. + _shared: Union[Set[str], FrozenSet[str]] = _NOTHING_SHARED is_eternal: bool def __init__(self, is_eternal: bool): @@ -117,15 +117,23 @@ def _stop_sharing(self, key: str) -> None: """Record that ``key`` no longer refers to a shared array.""" if key in self._shared: self._shared.discard(key) - if not self._shared: - self._shared = _NOTHING_SHARED + self._release_empty_set() def _stop_sharing_dropped_keys(self) -> None: """Forget the shared keys that ``_arrays`` no longer has.""" if self._shared: self._shared.intersection_update(self._arrays) - if not self._shared: - self._shared = _NOTHING_SHARED + self._release_empty_set() + + def _release_empty_set(self) -> None: + """Go back to the class's shared-nothing object once nothing is shared. + + The attribute is deleted rather than set, so that a storage sharing + nothing never carries one, and a copy or pickle of it reads the class + attribute too. + """ + if not self._shared: + del self._shared def put( self, value: ArrayLike, period: Period, branch_name: str = "default" diff --git a/tests/core/test_storage_shared_index.py b/tests/core/test_storage_shared_index.py index a2c077c74..c47b2ad00 100644 --- a/tests/core/test_storage_shared_index.py +++ b/tests/core/test_storage_shared_index.py @@ -16,7 +16,11 @@ from __future__ import annotations +import copy +import pickle + import numpy as np +import pytest from policyengine_core import periods from policyengine_core.data_storage import InMemoryStorage @@ -156,6 +160,59 @@ def test_storage_keeps_working_after_releasing_its_set(): assert clone._shared is _NOTHING_SHARED +def test_storage_sharing_nothing_has_no_attribute_of_its_own(): + storage = _storage_with("2017-01") + assert "_shared" not in vars(storage) + + clone = storage.clone(share_arrays=True) + assert "_shared" in vars(clone) + clone.get("2017-01") + + # Released, not replaced by another empty object. + assert "_shared" not in vars(clone) + + +@pytest.mark.parametrize( + "duplicate", + [copy.deepcopy, lambda storage: pickle.loads(pickle.dumps(storage))], + ids=["deepcopy", "pickle"], +) +def test_copied_or_unpickled_storage_keeps_the_index_rules(duplicate): + storage = _storage_with("2017-01", "2017-02") + released = storage.clone(share_arrays=True) + released.get("2017-01") + released.get("2017-02") + sharing = storage.clone(share_arrays=True) + + for original in (storage, released): + twin = duplicate(original) + assert twin._shared is _NOTHING_SHARED + assert np.array_equal(twin.get("2017-01"), original.get("2017-01")) + + twin = duplicate(sharing) + assert twin._shared == sharing._shared == {"default:2017-01", "default:2017-02"} + assert twin._shared is not sharing._shared + assert np.array_equal(twin.get("2017-01"), [0.0, 1.0]) + assert twin._shared == {"default:2017-02"} + assert sharing._shared == {"default:2017-01", "default:2017-02"} + twin.get("2017-02") + assert twin._shared is _NOTHING_SHARED + + +def test_storage_from_a_pickle_with_an_empty_set_still_works(): + # 3.32.12 pickled every storage with an empty set of its own. + storage = _storage_with("2017-01") + storage.__dict__["_shared"] = set() + restored = pickle.loads(pickle.dumps(storage)) + assert restored._shared == set() + + restored.get("2017-01") + restored.put(np.array([2.0, 3.0]), "2017-02") + restored.delete("2017-02") + restored.delete() + assert restored.get("2017-01") is None + + def test_simulation_storages_have_no_sets_of_their_own(tax_benefit_system): simulation = build_simulation(tax_benefit_system) simulation.calculate("disposable_income", JANUARY) diff --git a/tests/core/test_storage_shared_index_differential.py b/tests/core/test_storage_shared_index_differential.py index 463ed207f..792ec7a34 100644 --- a/tests/core/test_storage_shared_index_differential.py +++ b/tests/core/test_storage_shared_index_differential.py @@ -187,8 +187,10 @@ def _assert_index_invariants(storages, cleared): own_sets = [] for index, storage in enumerate(storages): shared = storage._shared - # The one shared-nothing object, exactly when nothing is shared. + # The one shared-nothing object, exactly when nothing is shared, and + # an attribute of the storage's own exactly when something is. assert (shared is _NOTHING_SHARED) == (not shared), index + assert ("_shared" in vars(storage)) == bool(shared), index if shared is not _NOTHING_SHARED: own_sets.append(shared) if index not in cleared: From 15a11184d2e2dc18d2cf2efa5fe4c319269eb8ca Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 22:54:59 -0400 Subject: [PATCH 3/5] Keep concurrent first reads safe when the set is released Two threads making first reads of one storage at once, which clone's docstring allows at the cost of an extra copy, could raise AttributeError: one released the set (del self._shared) while the other was between its membership check and its discard, or between its truth test and intersection_update. Each method now reads the set once into a local and works on it, and the release is a single dict pop that never raises. Regression tests run the other thread's read inside the membership check and inside the truth test, and 2,000 threaded trials; all three fail on the previous head (found in review round 2). Co-Authored-By: Claude Opus 5.5 --- .../data_storage/in_memory_storage.py | 33 +++-- tests/core/test_storage_shared_index.py | 113 ++++++++++++++++++ 2 files changed, 134 insertions(+), 12 deletions(-) diff --git a/policyengine_core/data_storage/in_memory_storage.py b/policyengine_core/data_storage/in_memory_storage.py index 2d64ed4f7..5547001a0 100644 --- a/policyengine_core/data_storage/in_memory_storage.py +++ b/policyengine_core/data_storage/in_memory_storage.py @@ -113,27 +113,36 @@ def get(self, period: Period, branch_name: str = "default") -> ArrayLike: self._stop_sharing(key) return values + # Each method below reads ``self._shared`` once and works on that set. Two + # threads making first reads at once (which the docstring of ``clone`` + # allows, at the cost of an extra copy) can then each release the set + # without either one touching the shared-nothing object or an attribute + # the other has already removed. + def _stop_sharing(self, key: str) -> None: """Record that ``key`` no longer refers to a shared array.""" - if key in self._shared: - self._shared.discard(key) - self._release_empty_set() + shared = self._shared + if key in shared: + shared.discard(key) + self._release_if_empty(shared) def _stop_sharing_dropped_keys(self) -> None: """Forget the shared keys that ``_arrays`` no longer has.""" - if self._shared: - self._shared.intersection_update(self._arrays) - self._release_empty_set() + shared = self._shared + if shared: + shared.intersection_update(self._arrays) + self._release_if_empty(shared) - def _release_empty_set(self) -> None: + def _release_if_empty(self, shared: Set[str]) -> None: """Go back to the class's shared-nothing object once nothing is shared. - The attribute is deleted rather than set, so that a storage sharing - nothing never carries one, and a copy or pickle of it reads the class - attribute too. + The storage's own attribute is removed rather than reassigned, so that + a storage sharing nothing carries none, and a copy or pickle of it + reads the class attribute too. ``pop`` with a default is one step that + never raises, so it is safe when another thread has already removed it. """ - if not self._shared: - del self._shared + if not shared: + self.__dict__.pop("_shared", None) def put( self, value: ArrayLike, period: Period, branch_name: str = "default" diff --git a/tests/core/test_storage_shared_index.py b/tests/core/test_storage_shared_index.py index c47b2ad00..fa1047fa7 100644 --- a/tests/core/test_storage_shared_index.py +++ b/tests/core/test_storage_shared_index.py @@ -259,3 +259,116 @@ def test_branch_storages_have_a_set_only_for_the_arrays_they_share( assert all( storage._shared is _NOTHING_SHARED for storage in _storages(simulation).values() ) + + +# ----- Two threads making first reads at once ----- # +# +# ``InMemoryStorage.clone`` documents that two threads may make first reads of +# one storage at once, at the cost of an extra copy. Releasing the set must not +# turn that into an error. The first tests run one read inside another at the +# point where a thread switch would do it, so they do not depend on timing. + + +class _ReadInsideCheck(set): + """A set of shared keys that runs ``interleave`` inside a membership check.""" + + def __init__(self, keys, interleave, on_call): + super().__init__(keys) + self.interleave = interleave + self.on_call = on_call + self.calls = 0 + + def __contains__(self, key): + found = set.__contains__(self, key) + self.calls += 1 + if self.calls == self.on_call: + self.interleave() + return found + + +@pytest.mark.parametrize("on_call", [1, 2]) +def test_another_thread_releasing_the_set_mid_read_is_harmless(on_call): + storage = _storage_with("2017-01") + clone = storage.clone(share_arrays=True) + # The other thread reads the same last shared key to the end, during this + # read's membership check in ``get`` (call 1) or in ``_stop_sharing`` + # (call 2). + clone._shared = _ReadInsideCheck( + clone._shared, lambda: clone.get("2017-01"), on_call + ) + + values = clone.get("2017-01") + + assert np.array_equal(values, [0.0, 1.0]) + assert values.flags.writeable + assert clone._shared is _NOTHING_SHARED + assert np.array_equal(storage.get("2017-01"), [0.0, 1.0]) + + +class _ReadAfterTruthTest(set): + """A set of shared keys that runs ``interleave`` after its first truth test.""" + + def __init__(self, keys, interleave): + super().__init__(keys) + self.interleave = interleave + self.tested = False + + def __bool__(self): + found = set.__len__(self) > 0 + if not self.tested: + self.tested = True + self.interleave() + return found + + +def test_another_thread_releasing_the_set_mid_delete_is_harmless(): + storage = _storage_with("2017-01", "2017-02") + clone = storage.clone(share_arrays=True) + clone.get("2017-01") + assert clone._shared == {"default:2017-02"} + # This thread deletes January, which it owns. Between its finding the set + # non-empty and its dropping deleted keys from it, the other thread reads + # February, the last shared key, and releases the set. + clone._shared = _ReadAfterTruthTest(clone._shared, lambda: clone.get("2017-02")) + + clone.delete("2017-01") + + assert clone._shared is _NOTHING_SHARED + assert clone.get("2017-01") is None + assert np.array_equal(clone.get("2017-02"), [1.0, 1.0]) + assert np.array_equal(storage.get("2017-02"), [1.0, 1.0]) + + +def test_concurrent_first_reads_never_raise(): + import sys + import threading + + source = _storage_with("2017-01", "2017-02") + errors = [] + interval = sys.getswitchinterval() + sys.setswitchinterval(1e-6) + try: + for _ in range(2_000): + clone = source.clone(share_arrays=True) + barrier = threading.Barrier(2) + + def read(month, clone=clone, barrier=barrier): + barrier.wait() + try: + clone.get(month) + except Exception as error: # noqa: BLE001 + errors.append(error) + + threads = [ + threading.Thread(target=read, args=(month,)) + for month in ("2017-01", "2017-02") + ] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + assert clone._shared is _NOTHING_SHARED + finally: + sys.setswitchinterval(interval) + + assert errors == [] From c8e024de4087b5ff6515bef097032d1001f8e9ca Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Fri, 2 Oct 2026 23:07:19 -0400 Subject: [PATCH 4/5] Test that nothing but the storage itself keeps its set alive Review round 2's mutation run found that a source storage keeping a reference to each clone's set (a leak, with no change in behaviour) survived every test. This test kills it. Co-Authored-By: Claude Opus 5.5 --- tests/core/test_storage_shared_index.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/tests/core/test_storage_shared_index.py b/tests/core/test_storage_shared_index.py index fa1047fa7..a622516f6 100644 --- a/tests/core/test_storage_shared_index.py +++ b/tests/core/test_storage_shared_index.py @@ -104,6 +104,24 @@ def test_sharing_clones_each_have_their_own_set(): assert second._shared == {"default:2017-01", "default:2017-02"} +def test_only_the_storage_itself_refers_to_its_set(): + import gc + + storage = _storage_with("2017-01", "2017-02") + clones = [storage.clone(share_arrays=True) for _ in range(3)] + + for clone in clones: + referrers = [ + referrer + for referrer in gc.get_referrers(clone._shared) + # Python 3.13+ may hold attributes inline in the object itself. + if referrer is not clone and referrer is not clone.__dict__ + ] + # Nothing else, the source included, keeps the set alive after the + # clone releases it. + assert referrers == [], referrers + + def test_set_is_released_once_every_shared_array_has_been_read(): clone = _storage_with("2017-01", "2017-02").clone(share_arrays=True) From c9fdfa60ff8f0008bde228a8a2be6e4e742f265a Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Sat, 3 Oct 2026 07:25:05 -0400 Subject: [PATCH 5/5] Say what the concurrency comments actually guarantee The comments said clone's docstring allows overlapping first reads at the cost of an extra copy. It says a storage is not safe to read from several threads at once, and review round 3 reproduced overlapping-read races that 3.32.12 already has. What this code guarantees is narrower: overlapping first reads do not raise. Comment-only change. Co-Authored-By: Claude Opus 5.5 --- policyengine_core/data_storage/in_memory_storage.py | 10 +++++----- tests/core/test_storage_shared_index.py | 9 +++++---- 2 files changed, 10 insertions(+), 9 deletions(-) diff --git a/policyengine_core/data_storage/in_memory_storage.py b/policyengine_core/data_storage/in_memory_storage.py index 5547001a0..9b70e2e99 100644 --- a/policyengine_core/data_storage/in_memory_storage.py +++ b/policyengine_core/data_storage/in_memory_storage.py @@ -113,11 +113,11 @@ def get(self, period: Period, branch_name: str = "default") -> ArrayLike: self._stop_sharing(key) return values - # Each method below reads ``self._shared`` once and works on that set. Two - # threads making first reads at once (which the docstring of ``clone`` - # allows, at the cost of an extra copy) can then each release the set - # without either one touching the shared-nothing object or an attribute - # the other has already removed. + # Each method below reads ``self._shared`` once and works on that set, and + # releases it with one ``dict.pop`` that never raises. A storage is still + # not safe to read from several threads at once (see ``clone``), but + # overlapping first reads do not raise here: neither can mutate the + # shared-nothing object or remove an attribute the other already removed. def _stop_sharing(self, key: str) -> None: """Record that ``key`` no longer refers to a shared array.""" diff --git a/tests/core/test_storage_shared_index.py b/tests/core/test_storage_shared_index.py index a622516f6..c919eb959 100644 --- a/tests/core/test_storage_shared_index.py +++ b/tests/core/test_storage_shared_index.py @@ -281,10 +281,11 @@ def test_branch_storages_have_a_set_only_for_the_arrays_they_share( # ----- Two threads making first reads at once ----- # # -# ``InMemoryStorage.clone`` documents that two threads may make first reads of -# one storage at once, at the cost of an extra copy. Releasing the set must not -# turn that into an error. The first tests run one read inside another at the -# point where a thread switch would do it, so they do not depend on timing. +# ``InMemoryStorage.clone`` documents that a storage is not safe to read from +# several threads at once; on 3.32.12 overlapping first reads never raised. +# Releasing the set must not make them raise. The first tests run one read +# inside another at the point where a thread switch would do it, so they do +# not depend on timing. class _ReadInsideCheck(set):