Write spike clusters atomically to avoid truncation on crash (#60)

np.save() opens the destination for writing and truncates it immediately,
so a process that dies, is killed, or fills the disk part way through
save_spike_clusters() destroys the existing spike_clusters.npy and leaves
an incomplete one in its place. That is unrecoverable loss of a manual
curation session, as reported in cortex-lab/phy#1249.

Write the array to a temporary file in the same directory as the
destination, fsync it, and move it into place with os.replace(), which is
atomic when both paths are on the same filesystem. If the write fails the
temporary file is removed and the existing file is untouched.

Co-authored-by: Cyrille Rossant <rossant@users.noreply.github.com>

Author: adityasingh2400

CHANGELOG.md

@@ -14,6 +14,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Development checkouts now report a development version with a Git commit suffix.
 - Template datasets without waveform templates can be reloaded after cluster assignments change,
   and stored spike-waveform subsets provide their waveform sample count.
+- Cluster assignments are written atomically, so a crash during a save no longer truncates `spike_clusters.npy`.
 - TSV, JSON, text and `params.py` files are written atomically, so a crash during a save no longer
   truncates the file it was replacing, for instance `cluster_group.tsv`.
 - #57 a blank `dat_path` in `params.py` is now read as "no raw data file" instead of resolving to

phylib/io/model.py

@@ -24,7 +24,7 @@
     get_ephys_reader, RandomEphysReader, extract_waveforms,
     get_spike_waveforms, export_waveforms)
 from phylib.utils import Bunch
-from phylib.utils._misc import _write_tsv_simple, read_tsv, read_python
+from phylib.utils._misc import _atomic_open, _write_tsv_simple, read_tsv, read_python
 from phylib.utils.geometry import linear_positions
 
 logger = logging.getLogger(__name__)
@@ -62,6 +62,22 @@ def write_array(name, arr):
     np.save(name, arr)
 
 
+def _save_npy_atomic(path, arr):
+    """Save an array to a npy file without ever truncating the existing file.
+
+    The array is written to a temporary file in the same directory as the destination, and only
+    then moved into place with `os.replace()`, which is atomic as long as both paths are on the
+    same filesystem. If anything goes wrong while writing, the temporary file is removed and the
+    existing file is left untouched.
+
+    """
+    # NOTE: np.save() appends `.npy` to a file *name* that lacks it, but not to an already-open
+    # file object. The shared writer also preserves existing permissions and removes its hidden
+    # sibling temporary file if np.save() raises.
+    with _atomic_open(path, mode='wb') as f:
+        np.save(f, arr)
+
+
 def from_sparse(data, cols, channel_ids):
     """Convert a sparse structure into a dense one.
 
@@ -1354,7 +1370,7 @@ def save_spike_clusters(self, spike_clusters):
         """Save the spike clusters."""
         path = self._find_path('spike_clusters.npy', 'spikes.clusters.npy', multiple_ok=False)
         logger.debug("Save spike clusters to `%s`.", path)
-        np.save(path, spike_clusters)
+        _save_npy_atomic(path, spike_clusters)
 
     def save_spikes_subset_waveforms(self, max_n_spikes_per_template=None, max_n_channels=None,
                                      sample2unit=1.):

phylib/io/tests/test_model.py

@@ -7,13 +7,16 @@
 #------------------------------------------------------------------------------
 
 import logging
+import os
+import stat
 
 import numpy as np
 from numpy.testing import assert_equal as ae
 from pytest import raises
 
 # from phylib.utils import Bunch
 from phylib.utils.testing import captured_output
+from .. import model as model_module
 from ..model import from_sparse, get_template_params, load_model
 
 logger = logging.getLogger(__name__)
@@ -135,6 +138,50 @@ def test_model_save(template_model_full):
     m.save_spike_clusters(m.spike_clusters)
 
 
+def test_model_save_spike_clusters_atomic(template_model, monkeypatch):
+    # Regression test for cortex-lab/phy#1249: a save that dies half-way through must not
+    # leave spike_clusters.npy truncated or empty.
+    m = template_model
+    path = m.dir_path / 'spike_clusters.npy'
+
+    # Happy path: a successful save round-trips and does not create `spike_clusters.npy.npy`.
+    original = m.spike_clusters.copy()
+    if os.name != 'nt':
+        path.chmod(0o640)
+    m.save_spike_clusters(original)
+    assert not (m.dir_path / 'spike_clusters.npy.npy').exists()
+    ae(np.load(path), original)
+    if os.name != 'nt':
+        # Shared lab datasets must not become private to the user who saved them.
+        assert stat.S_IMODE(path.stat().st_mode) == 0o640
+
+    original_bytes = path.read_bytes()
+    assert original_bytes
+    files_before = sorted(p.name for p in m.dir_path.iterdir())
+
+    def _partial_save(file, arr, *args, **kwargs):
+        # Mimic np.save() dying half-way through: whatever it was given is opened for writing
+        # (which truncates it) and a few bytes are written before the failure.
+        if hasattr(file, 'write'):
+            file.write(b'\x93NUMPY')
+        else:
+            with open(str(file), 'wb') as f:
+                f.write(b'\x93NUMPY')
+        raise IOError("No space left on device")
+
+    monkeypatch.setattr(model_module.np, 'save', _partial_save)
+    with raises(IOError):
+        m.save_spike_clusters(original + 1)
+
+    # The original file must still be there, byte for byte.
+    assert path.exists()
+    assert path.read_bytes() == original_bytes
+    ae(np.load(path), original)
+
+    # No temporary file must be left behind.
+    assert sorted(p.name for p in m.dir_path.iterdir()) == files_before
+
+
 def _set_dat_path(params_path, value):
     """Replace the `dat_path` assignment of a `params.py` file by a literal value."""
     lines = params_path.read_text().splitlines()