Atomic workspace saving and non-uniform coordinates in Web API transfers

datalab/gui/h5io.py

@@ -27,7 +27,7 @@
 from datalab.env import execenv
 from datalab.gui.panel.base import H5ImportBatch
 from datalab.h5 import H5Importer
-from datalab.h5.native import NativeH5Reader, NativeH5Writer
+from datalab.h5.native import NativeH5Reader, atomic_native_writer
 from datalab.utils.qthelpers import create_progress_bar, qt_try_loadsave_file
 
 if TYPE_CHECKING:
@@ -80,14 +80,15 @@ def __progbartitle(fname: str) -> str:
         return _("Loading data from %s...") % osp.basename(fname)
 
     def save_file(self, filename: str) -> None:
-        """Save all signals and images from DataLab model into a HDF5 file"""
-        writer = NativeH5Writer(filename)
-        try:
+        """Save all signals and images from DataLab model into a HDF5 file
+
+        The file is replaced only once completely written: if saving fails, the
+        previous file (if any) is left untouched.
+        """
+        with atomic_native_writer(filename) as writer:
             for panel in self.mainwindow.panels:
                 panel.serialize_to_hdf5(writer)
             self.mainwindow.provenance.save(writer.h5)
-        finally:
-            writer.close()
 
     @staticmethod
     def __read_provenance(

datalab/h5/history.py

@@ -13,7 +13,7 @@
 
 from datalab.config import Conf, _
 from datalab.gui.panel.base import H5ImportBatch
-from datalab.h5.native import NativeH5Reader, NativeH5Writer
+from datalab.h5.native import NativeH5Reader, NativeH5Writer, atomic_native_writer
 from datalab.history import HistorySession
 from datalab.objectmodel import get_uuid
 from datalab.utils.qthelpers import qt_try_loadsave_file, save_restore_stds
@@ -53,7 +53,7 @@ def save_to_dlhist_file(panel: HistoryPanel, filename: str | None = None) -> boo
         filename += ".dlhist"
     with qt_try_loadsave_file(panel.parentWidget(), filename, "save"):
         Conf.base_dir.set(filename)
-        with NativeH5Writer(filename) as writer:
+        with atomic_native_writer(filename) as writer:
             # Make the .dlhist file panel-contained: store the signal and
             # image panel objects (all of them) alongside the history, so
             # that reopening restores both the data objects and the history

datalab/h5/native.py

@@ -8,7 +8,13 @@
 
 from __future__ import annotations
 
+import contextlib
 import importlib
+import os
+import os.path as osp
+import stat
+import tempfile
+from collections.abc import Generator
 from typing import Any, Callable
 
 from guidata.io import HDF5Reader, HDF5Writer
@@ -73,6 +79,57 @@ def write(self, val: Any, group_name: str | None = None) -> None:
                 raise
 
 
+def _set_final_permissions(tmpname: str, filename: str) -> None:
+    """Give the temporary file the permissions the final file must have.
+
+    Args:
+        tmpname: Temporary file name (created with owner-only permissions).
+        filename: Final file name: its permissions are kept if it exists,
+         otherwise those of a newly created file are applied.
+    """
+    if osp.exists(filename):
+        mode = stat.S_IMODE(os.stat(filename).st_mode)
+    else:
+        umask = os.umask(0)
+        os.umask(umask)
+        mode = 0o666 & ~umask
+    os.chmod(tmpname, mode)
+
+
+@contextlib.contextmanager
+def atomic_native_writer(filename: str) -> Generator[NativeH5Writer, None, None]:
+    """Write a native DataLab HDF5 file without losing the previous one on failure.
+
+    The file is written to a temporary file of the same directory, which replaces
+    ``filename`` only once it has been completely written and closed. If writing
+    fails (serialization error, full disk, interruption...), the temporary file is
+    removed and the previous file, if any, is left untouched.
+
+    Args:
+        filename: HDF5 file name
+
+    Yields:
+        Writer of the temporary file
+    """
+    filename = osp.abspath(filename)
+    fd, tmpname = tempfile.mkstemp(
+        prefix=f".{osp.basename(filename)}.", suffix=".tmp", dir=osp.dirname(filename)
+    )
+    os.close(fd)
+    try:
+        writer = NativeH5Writer(tmpname)
+        try:
+            yield writer
+        finally:
+            writer.close()
+        _set_final_permissions(tmpname, filename)
+        os.replace(tmpname, filename)
+    except BaseException:
+        with contextlib.suppress(OSError):
+            os.remove(tmpname)
+        raise
+
+
 class NativeH5Reader(HDF5Reader):
     """DataLab signal/image objects HDF5 guidata dataset Writer class
 

datalab/tests/features/hdf5/atomic_save_unit_test.py

@@ -0,0 +1,160 @@
+# Copyright (c) DataLab Platform Developers, BSD 3-Clause license, see LICENSE file.
+
+"""
+Atomic HDF5 save unit tests
+---------------------------
+
+A failed save must never destroy the previous version of a workspace (or
+``.dlhist``) file: the file is written to a temporary file of the same directory,
+which replaces the previous file only once completely written.
+"""
+
+# guitest: show
+
+from __future__ import annotations
+
+import os
+import os.path as osp
+from types import SimpleNamespace
+
+import h5py
+import numpy as np
+import pytest
+from sigima.tests.data import create_paracetamol_signal
+
+from datalab.gui.h5io import H5InputOutput
+from datalab.h5.native import DATALAB_VERSION_NAME, atomic_native_writer
+from datalab.tests import datalab_test_app_context, helpers
+
+
+def _write_previous_file(filename: str) -> bytes:
+    """Write a previous version of a file and return its contents."""
+    with h5py.File(filename, "w") as h5file:
+        h5file["original"] = np.arange(5.0)
+    with open(filename, "rb") as fdesc:
+        return fdesc.read()
+
+
+def _read_bytes(filename: str) -> bytes:
+    """Return the contents of a file."""
+    with open(filename, "rb") as fdesc:
+        return fdesc.read()
+
+
+def _assert_no_temporary_file(dirname: str, filename: str) -> None:
+    """Check that no temporary file was left next to the saved file."""
+    assert sorted(os.listdir(dirname)) == [osp.basename(filename)]
+
+
+def test_atomic_writer_replaces_file_once_written(tmp_path) -> None:
+    """A successful write replaces the previous file."""
+    filename = str(tmp_path / "workspace.h5")
+    _write_previous_file(filename)
+    with atomic_native_writer(filename) as writer:
+        writer.h5["partial"] = np.arange(3.0)
+        assert not osp.samefile(writer.filename, filename)
+    with h5py.File(filename, "r") as h5file:
+        assert set(h5file) == {DATALAB_VERSION_NAME, "partial"}
+    _assert_no_temporary_file(str(tmp_path), filename)
+
+
+def test_atomic_writer_creates_new_file(tmp_path) -> None:
+    """A new file is created with the permissions of a newly created file."""
+    filename = str(tmp_path / "new.h5")
+    reference = str(tmp_path / "reference")
+    with open(reference, "wb"):
+        pass
+    with atomic_native_writer(filename) as writer:
+        writer.h5["data"] = np.arange(3.0)
+    assert os.stat(filename).st_mode == os.stat(reference).st_mode
+    os.remove(reference)
+    _assert_no_temporary_file(str(tmp_path), filename)
+
+
+def test_atomic_writer_keeps_previous_file_on_failure(tmp_path) -> None:
+    """A failed write leaves the previous file untouched and no temporary file."""
+    filename = str(tmp_path / "workspace.h5")
+    previous = _write_previous_file(filename)
+    with pytest.raises(OSError, match="disk full"):
+        with atomic_native_writer(filename) as writer:
+            writer.h5["partial"] = np.arange(3.0)
+            raise OSError("disk full")
+    assert _read_bytes(filename) == previous
+    _assert_no_temporary_file(str(tmp_path), filename)
+
+
+class _Panel:
+    """Panel double writing one dataset, then optionally failing."""
+
+    def __init__(self, name: str, fail: bool = False) -> None:
+        self.name = name
+        self.fail = fail
+
+    def serialize_to_hdf5(self, writer) -> None:
+        """Serialize the panel, like ``BaseDataPanel.serialize_to_hdf5``."""
+        writer.h5[self.name] = np.arange(3.0)
+        if self.fail:
+            raise OSError("write error")
+
+
+class _Provenance:
+    """Provenance double, optionally failing."""
+
+    def __init__(self, fail: bool = False) -> None:
+        self.fail = fail
+
+    def save(self, h5file: h5py.File) -> None:
+        """Save the provenance block, like ``ProvenanceService.save``."""
+        h5file["provenance"] = np.arange(2.0)
+        if self.fail:
+            raise ValueError("provenance error")
+
+
+@pytest.mark.parametrize(
+    "panels, provenance",
+    [
+        ([_Panel("signals"), _Panel("images", fail=True)], _Provenance()),
+        ([_Panel("signals"), _Panel("images")], _Provenance(fail=True)),
+    ],
+    ids=["panel", "provenance"],
+)
+def test_save_file_keeps_previous_workspace_on_failure(
+    tmp_path, panels: list[_Panel], provenance: _Provenance
+) -> None:
+    """Regression: saving truncated the previous workspace file before writing,
+    so that an error in a panel or in the provenance destroyed it."""
+    filename = str(tmp_path / "workspace.h5")
+    previous = _write_previous_file(filename)
+    h5io = H5InputOutput(SimpleNamespace(panels=panels, provenance=provenance))
+    with pytest.raises((OSError, ValueError)):
+        h5io.save_file(filename)
+    assert _read_bytes(filename) == previous
+    _assert_no_temporary_file(str(tmp_path), filename)
+
+
+def test_failed_save_keeps_workspace_modified(monkeypatch) -> None:
+    """After a failed save, the previous file is intact and the workspace is still
+    marked as modified."""
+    with helpers.WorkdirRestoringTempDir() as tmpdir:
+        filename = osp.join(tmpdir, "workspace.h5")
+        previous = _write_previous_file(filename)
+        with datalab_test_app_context(console=False) as win:
+            win.signalpanel.add_object(create_paracetamol_signal())
+            assert win.is_modified()
+
+            def fail(_writer) -> None:
+                raise OSError("write error")
+
+            monkeypatch.setattr(win.imagepanel, "serialize_to_hdf5", fail)
+            with pytest.raises(OSError, match="write error"):
+                win.save_h5_workspace(filename)
+            assert win.is_modified()
+            assert _read_bytes(filename) == previous
+            monkeypatch.undo()
+            win.save_h5_workspace(filename)
+            assert not win.is_modified()
+        assert not [name for name in os.listdir(tmpdir) if name.endswith(".tmp")]
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])

datalab/tests/scenarios/common.py

@@ -328,9 +328,9 @@ def run_image_computations(
             width=width_param,
             height=height_param,
             xmin=ima1.x0,
-            xmax=ima1.x0 + ima1.width,
+            xmax=ima1.x0 + ima1.width - ima1.dx,
             ymin=ima1.y0,
-            ymax=ima1.y0 + ima1.height,
+            ymax=ima1.y0 + ima1.height - ima1.dy,
         )
         panel.processor.run_feature("resampling", param)
 

datalab/tests/webapi_extra_test.py

@@ -363,6 +363,76 @@ def test_image_obj_metadata_with_numpy_origin(self) -> None:
         assert result.x0 == pytest.approx(1.5)
         assert result.dx == pytest.approx(0.1)
 
+    def test_image_non_uniform_coords_round_trip(self) -> None:
+        """Non-uniform image coordinates survive a round trip.
+
+        Regression: they were lost, the image being rebuilt with uniform
+        coordinates (its pixels survived, but not their physical positions).
+        """
+        obj = _make_image(shape=(3, 4))
+        obj.xunit, obj.yunit = "mm", "s"
+        obj.set_coords(np.array([1.0, 2.0, 4.0, 8.0]), np.array([10.0, 11.0, 15.0]))
+        result = deserialize_object_from_npz(serialize_object_to_npz(obj))
+        assert not result.is_uniform_coords
+        np.testing.assert_array_equal(result.xcoords, obj.xcoords)
+        np.testing.assert_array_equal(result.ycoords, obj.ycoords)
+        np.testing.assert_array_equal(result.data, obj.data)
+        assert (result.xunit, result.yunit) == ("mm", "s")
+
+    def test_image_uniform_coords_round_trip(self) -> None:
+        """Uniform images are stored without coordinate arrays."""
+        obj = _make_image(shape=(3, 4))
+        obj.set_uniform_coords(0.5, 2.0, -1.0, 3.0)
+        data = serialize_object_to_npz(obj)
+        with zipfile.ZipFile(io.BytesIO(data)) as zf:
+            assert sorted(zf.namelist()) == ["data.npy", "metadata.json"]
+            assert json.loads(zf.read("metadata.json"))["is_uniform_coords"] is True
+        result = deserialize_object_from_npz(data)
+        assert result.is_uniform_coords
+        assert (result.dx, result.dy, result.x0, result.y0) == (0.5, 2.0, -1.0, 3.0)
+
+    def test_image_archive_without_coords_flag_is_uniform(self) -> None:
+        """Archives written by earlier versions (no ``is_uniform_coords``) are read
+        with uniform coordinates."""
+        buf = io.BytesIO()
+        with zipfile.ZipFile(buf, "w") as zf:
+            arr = io.BytesIO()
+            np.save(arr, np.zeros((3, 4)))
+            zf.writestr("data.npy", arr.getvalue())
+            metadata = {"type": "image", "x0": 1.0, "y0": 2.0, "dx": 0.5, "dy": 0.25}
+            zf.writestr("metadata.json", json.dumps(metadata))
+        result = deserialize_object_from_npz(buf.getvalue())
+        assert result.is_uniform_coords
+        assert (result.dx, result.dy, result.x0, result.y0) == (0.5, 0.25, 1.0, 2.0)
+
+    @pytest.mark.parametrize(
+        "xcoords, ycoords, match",
+        [
+            (None, np.arange(3.0), "missing xcoords.npy"),
+            (np.arange(3.0), np.arange(3.0), "do not match image shape"),
+            (np.arange(4.0), np.arange(4.0), "do not match image shape"),
+        ],
+    )
+    def test_image_invalid_non_uniform_coords_raise(
+        self, xcoords, ycoords, match: str
+    ) -> None:
+        """Missing or mismatched coordinate arrays are rejected."""
+        buf = io.BytesIO()
+        with zipfile.ZipFile(buf, "w") as zf:
+            for name, array in (
+                ("data.npy", np.zeros((3, 4))),
+                ("xcoords.npy", xcoords),
+                ("ycoords.npy", ycoords),
+            ):
+                if array is not None:
+                    arr = io.BytesIO()
+                    np.save(arr, array)
+                    zf.writestr(name, arr.getvalue())
+            metadata = {"type": "image", "is_uniform_coords": False}
+            zf.writestr("metadata.json", json.dumps(metadata))
+        with pytest.raises(ValueError, match=match):
+            deserialize_object_from_npz(buf.getvalue())
+
 
 # =============================================================================
 # Routes integration tests
@@ -638,6 +708,29 @@ def test_get_data_missing_object(self, api_client) -> None:
         response = client.get("/api/v1/objects/nope/data", headers=_auth(token))
         assert response.status_code == 404
 
+    def test_non_uniform_image_round_trip_through_endpoints(self, api_client) -> None:
+        """An image with non-uniform coordinates keeps them when uploaded with
+        ``PUT data`` and downloaded with ``GET data``."""
+        client, token, adapter = api_client
+        obj = _make_image("NU", shape=(3, 4))
+        xcoords, ycoords = np.array([1.0, 2.0, 4.0, 8.0]), np.array([10.0, 11.0, 15.0])
+        obj.set_coords(xcoords, ycoords)
+        response = client.put(
+            "/api/v1/objects/NU/data",
+            content=serialize_object_to_npz(obj),
+            headers={**_auth(token), "Content-Type": "application/x-npz"},
+        )
+        assert response.status_code == 201
+        stored = adapter.get_object("NU")
+        assert not stored.is_uniform_coords
+        np.testing.assert_array_equal(stored.xcoords, xcoords)
+        response = client.get("/api/v1/objects/NU/data", headers=_auth(token))
+        assert response.status_code == 200
+        result = deserialize_object_from_npz(response.content)
+        assert not result.is_uniform_coords
+        np.testing.assert_array_equal(result.xcoords, xcoords)
+        np.testing.assert_array_equal(result.ycoords, ycoords)
+
     def test_put_creates_object(self, api_client) -> None:
         """``PUT data`` creates a new object from an NPZ payload."""
         client, token, adapter = api_client