Write spike clusters atomically to avoid truncation on crash

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.

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`.
 
 ## [2.7.0] 2025-12-10
 

phylib/io/model.py

@@ -13,6 +13,7 @@
 from operator import itemgetter
 from pathlib import Path
 import shutil
+import tempfile
 
 import numpy as np
 # from numpy.lib.format import open_memmap
@@ -62,6 +63,36 @@ 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.
+
+    """
+    path = Path(path)
+    # NOTE: the temporary file must not be named `*.npy`, otherwise a leftover file would be
+    # picked up by the `spike_*.npy` glob in `_load_spike_attributes()`.
+    tmp = tempfile.NamedTemporaryFile(
+        delete=False, dir=str(path.parent), prefix='.' + path.name + '.', suffix='.tmp')
+    try:
+        with tmp:
+            # NOTE: np.save() appends `.npy` to a file *name* that lacks it, but not to an
+            # already-open file object, which is what is passed here.
+            np.save(tmp, arr)
+            tmp.flush()
+            os.fsync(tmp.fileno())
+        os.replace(tmp.name, path)
+    except BaseException:
+        try:
+            os.remove(tmp.name)
+        except OSError:  # pragma: no cover
+            pass
+        raise
+
+
 def from_sparse(data, cols, channel_ids):
     """Convert a sparse structure into a dense one.
 
@@ -1357,7 +1388,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

@@ -14,6 +14,7 @@
 
 # from phylib.utils import Bunch
 from phylib.utils.testing import captured_output
+from .. import model as model_module
 from ..model import from_sparse, load_model
 
 logger = logging.getLogger(__name__)
@@ -135,6 +136,45 @@ 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()
+    m.save_spike_clusters(original)
+    assert not (m.dir_path / 'spike_clusters.npy.npy').exists()
+    ae(np.load(path), original)
+
+    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 test_model_reload_curated_dataset_without_templates(template_path):
     model = load_model(template_path)
     model.save_spikes_subset_waveforms(10, 16)