Give each disk-backed holder storage its own directory

changelog.d/disk-storage-branch-ownership.fixed.md

@@ -0,0 +1 @@
+Give independently created temporary holder storages distinct directories, create caller-supplied storage directories on first use, and clear stale enum metadata when replacing enum values with plain arrays.

docs/python_api/data_storage.md

@@ -13,6 +13,17 @@ The `policyengine_core.data_storage` module contains the classes that are used t
 
 ## OnDiskStorage
 
+Each independently created temporary holder storage gets its own subdirectory,
+so replacing a holder leaves files read by existing clones intact. Cloned
+storages retain their source's directory and use new files for later writes.
+Explicit `Holder.create_disk_storage(directory=...)` calls and preserved
+storages keep the `directory/variable` layout used by simulation dumps.
+
+`OnDiskStorage.restore()` reads direct key files in its directory. It does not
+read files in the `replaced` subdirectory or select the latest write across
+independently constructed storages. Clone or copy a restored storage before
+writing when its existing values must remain isolated.
+
 ```{eval-rst}
 .. autoclass:: policyengine_core.data_storage.on_disk_storage.OnDiskStorage
     :members:
@@ -22,6 +33,10 @@ The `policyengine_core.data_storage` module contains the classes that are used t
 
 ## TemporaryStorageDirectory
 
+A caller-supplied `simulation._data_storage_dir` is created on first use if it
+does not exist. It keeps the caller-owned lifetime described below.
+Exits that skip finalizers (such as `os._exit` in multiprocessing workers) can leave uniquely named holder subdirectories in a caller-owned directory to accumulate across runs; the caller must remove them.
+
 ```{eval-rst}
 .. automodule:: policyengine_core.data_storage.storage_directory
     :no-members:

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

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

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

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
 

tests/core/test_branch_scoped_delete_arrays.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import gc
+import os
 
 import numpy as np
 
@@ -82,32 +83,31 @@ def test_delete_arrays_purges_only_visible_branch_names():
     )
 
 
-def test_on_disk_storage_clone_copies_metadata_without_owning_directory(tmp_path):
-    """A cloned disk view must own metadata, but not the shared directory."""
+def test_on_disk_storage_clone_copies_metadata_and_writes_elsewhere(tmp_path):
+    """A cloned disk view owns its metadata and writes a separate file."""
     storage_dir = tmp_path / "storage"
     storage_dir.mkdir()
     storage = OnDiskStorage(str(storage_dir), is_eternal=True)
     storage.put(np.asarray([12.0]), PERIOD, "default")
     storage._enums["sentinel"] = "enum"
 
-    try:
-        assert storage.preserve_storage_dir is False
-        clone = storage.clone()
-
-        assert clone is not storage
-        assert clone.storage_dir == storage.storage_dir
-        assert clone.is_eternal == storage.is_eternal
-        assert clone.preserve_storage_dir is True
-        assert clone._files == storage._files
-        assert clone._files is not storage._files
-        assert clone._enums == storage._enums
-        assert clone._enums is not storage._enums
-        assert clone._storage_dir_owner is storage
-        assert clone.clone()._storage_dir_owner is storage
-    finally:
-        # Let pytest's tmp_path cleanup remove the directory even if an
-        # assertion fails before the clone can prove its ownership setting.
-        storage.preserve_storage_dir = True
+    clone = storage.clone()
+
+    assert clone is not storage
+    assert clone.is_eternal == storage.is_eternal
+    assert clone._files == storage._files
+    assert clone._files is not storage._files
+    assert clone._enums == storage._enums
+    assert clone._enums is not storage._enums
+
+    clone.put(np.asarray([13.0]), PERIOD, "default")
+
+    (clone_file,) = clone._files.values()
+    (source_file,) = storage._files.values()
+    assert clone_file != source_file
+    assert os.path.isfile(clone_file)
+    np.testing.assert_array_equal(storage.get(PERIOD, "default"), [12.0])
+    np.testing.assert_array_equal(clone.get(PERIOD, "default"), [13.0])
 
 
 def test_on_disk_storage_clone_keeps_cleanup_owner_alive(tmp_path):
@@ -145,7 +145,6 @@ def test_child_disk_delete_keeps_parent_view_intact(tmp_path):
         child.delete_arrays("salary", PERIOD)
 
         assert child_storage is not parent_storage
-        assert child_storage.preserve_storage_dir is True
         np.testing.assert_array_equal(inherited_value, [3_000.0])
         assert disk_key in parent_storage._files
         assert disk_key not in child_storage._files

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:

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;