Repository: basnijholt/adaptive-scheduler

Issue:
save_dataframe does not write atomically
This is problematic, as data may be corrupted if the writing process dies during the write.

Discussion:

Patch:
diff --git a/adaptive_scheduler/utils.py b/adaptive_scheduler/utils.py
index c03d7fd..df412ca 100644
--- a/adaptive_scheduler/utils.py
+++ b/adaptive_scheduler/utils.py
@@ -14,14 +14,25 @@ import random
 import shutil
 import tempfile
 import time
+import uuid
 import warnings
 from concurrent.futures import ThreadPoolExecutor
-from contextlib import suppress
+from contextlib import contextmanager, suppress
 from datetime import datetime, timedelta, timezone
 from inspect import signature
 from multiprocessing import Manager
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Callable, List, Literal, Union
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    List,
+    Literal,
+    Union,
+)
+from typing import (
+    get_args as get_type_args,
+)
 
 import adaptive
 import cloudpickle
@@ -806,6 +817,7 @@ def save_dataframe(
     format: _DATAFRAME_FORMATS = "pickle",  # noqa: A002
     save_kwargs: dict[str, Any] | None = None,
     expand_dicts: bool = True,
+    atomically: bool = True,
     **to_dataframe_kwargs: Any,
 ) -> Callable[[adaptive.BaseLearner], None]:
     """Save the learner's data to disk as pandas.DataFrame."""
@@ -816,30 +828,62 @@ def save_dataframe(
         if expand_dicts:
             df = expand_dict_columns(df)
         fname_df = fname_to_dataframe(fname, format=format)
-        if format == "parquet":
-            df.to_parquet(fname_df, **save_kwargs)
-        elif format == "csv":
-            df.to_csv(fname_df, **save_kwargs)
-        elif format == "hdf":
+
+        if format not in get_type_args(_DATAFRAME_FORMATS):
+            msg = f"Unknown format {format}"
+            raise ValueError(msg)
+        do_save = getattr(df, f"to_{format}")
+
+        if format == "hdf":
             assert save_kwargs is not None  # for mypy
             if "key" not in save_kwargs:
                 save_kwargs["key"] = "data"
-            df.to_hdf(fname_df, **save_kwargs)
-        elif format == "pickle":
-            df.to_pickle(fname_df, **save_kwargs)
-        elif format == "feather":
-            df.to_feather(fname_df, **save_kwargs)
-        elif format == "excel":
-            df.to_excel(fname_df, **save_kwargs)
-        elif format == "json":
-            df.to_json(fname_df, **save_kwargs)
+
+        if atomically:
+            with atomic_write(fname_df, return_path=True) as fname_temp:
+                do_save(fname_temp, **save_kwargs)
         else:
-            msg = f"Unknown format {format}."
-            raise ValueError(msg)
+            do_save(fname_df, **save_kwargs)
 
     return save
 
 
+@contextmanager
+def atomic_write(
+    dest: os.PathLike,
+    mode: str = "w",
+    *args: Any,
+    return_path: bool = False,
+    **kwargs: Any,
+) -> Any:
+    """Write atomically to 'dest', using a temporary file in the same directory.
+
+    This function has the same signature as 'open', except that the default
+    mode is 'w', not 'r', and there is an additional keyword-only parameter,
+    'return_path'. If 'return_path=True' then a Path pointing to the (as yet
+    nonexistant) temporary file is yielded, rather than a file handle.
+    This is useful when calling libraries that expect a path, rather than an
+    open file handle.
+    """
+    temp_dest = Path(dest).with_suffix(f".temp.{os.getpid()}.{uuid.uuid4()}")
+    try:
+        # First create an empty file; this ensures we have the same semantics
+        # as 'open(..., mode="w")'.
+        temp_dest.open("w").close()
+        # Now give control back to the caller.
+        if return_path:
+            yield temp_dest
+        else:
+            with temp_dest.open(mode, *args, **kwargs) as fp:
+                yield fp
+        # Atomically change 'dest' to point to the 'temp_dest' inode.
+        os.replace(temp_dest, dest)  # noqa: PTH105
+    except Exception:
+        with suppress(FileNotFoundError):
+            os.remove(temp_dest)  # noqa: PTH107
+        raise
+
+
 _DATAFRAME_FORMATS = Literal[
     "parquet",
     "csv",
@@ -867,7 +911,7 @@ def expand_dict_columns(df: pd.DataFrame) -> pd.DataFrame:
     return df
 
 
-def load_dataframes(  # noqa: PLR0912
+def load_dataframes(
     fnames: list[str] | list[list[str]] | list[Path] | list[list[Path]],
     *,
     concat: bool = True,
@@ -876,37 +920,30 @@ def load_dataframes(  # noqa: PLR0912
 ) -> pd.DataFrame | list[pd.DataFrame]:
     """Load a list of dataframes from disk."""
     read_kwargs = read_kwargs or {}
+    if format == "hdf" and "key" not in read_kwargs:
+        read_kwargs["key"] = "data"
+
+    if format not in get_type_args(_DATAFRAME_FORMATS):
+        msg = f"Unknown format {format}."
+        raise ValueError(msg)
+
+    do_read = getattr(pd, f"read_{format}")
+
     dfs = []
     for fn in fnames:
         fn_df = fname_to_dataframe(fn, format=format)
         if not os.path.exists(fn_df):  # noqa: PTH110
             continue
         try:
-            if format == "parquet":
-                df = pd.read_parquet(fn_df, **read_kwargs)
-            elif format == "csv":
-                df = pd.read_csv(fn_df, **read_kwargs)
-            elif format == "hdf":
-                if "key" not in read_kwargs:
-                    read_kwargs["key"] = "data"
-                df = pd.read_hdf(fn_df, **read_kwargs)
-            elif format == "pickle":
-                df = pd.read_pickle(fn_df, **read_kwargs)  # noqa: S301
-            elif format == "feather":
-                df = pd.read_feather(fn_df, **read_kwargs)
-            elif format == "excel":
-                df = pd.read_excel(fn_df, **read_kwargs)
-            elif format == "json":
-                df = pd.read_json(fn_df, **read_kwargs)
-            else:
-                msg = f"Unknown format {format}."
-                raise ValueError(msg)  # noqa: TRY301
+            df = do_read(fn_df, **read_kwargs)
         except Exception:  # noqa: BLE001
             msg = f"`{fn}`'s DataFrame ({fn_df}) could not be read."
             console.print(msg)
             continue
+
         df["fname"] = len(df) * [fn]
         dfs.append(df)
+
     if concat:
         if dfs:
             return pd.concat(dfs, axis=0)


Test patch:
diff --git a/tests/test_utils.py b/tests/test_utils.py
index 074f233..200f137 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -4,6 +4,7 @@ from __future__ import annotations
 import platform
 import sys
 import time
+import typing
 from concurrent.futures import ProcessPoolExecutor
 from datetime import datetime, timedelta
 from pathlib import Path
@@ -335,25 +336,36 @@ def test_cloudpickle_learners(tmp_path: Path) -> None:
         assert fname_learner.exists()
 
 
-def test_save_dataframe(tmp_path: Path) -> None:
+@pytest.mark.parametrize("atomically", [False, True])
+@pytest.mark.parametrize("fmt", ["pickle", "csv", "json"])
+def test_save_dataframe(
+    tmp_path: Path,
+    atomically: bool,  # noqa: FBT001
+    fmt: str,
+) -> None:
     """Test `utils.save_dataframe`."""
+    fmt = typing.cast(utils._DATAFRAME_FORMATS, fmt)
     learner = adaptive.Learner1D(lambda x: x, bounds=(-10, 10))
-    fname = str(tmp_path / ("test.pickle"))
-    save_df = utils.save_dataframe(fname)
+    fname = str(tmp_path / f"test.{fmt}")
+    save_df = utils.save_dataframe(fname, atomically=atomically, format=fmt)
 
     learner.ask(10)
     learner.tell(10, 42)
 
     save_df(learner)
 
-    assert utils.fname_to_dataframe(fname).exists()
+    assert utils.fname_to_dataframe(fname, format=fmt).exists()
 
 
-def test_load_dataframes(tmp_path: Path) -> None:
+@pytest.mark.parametrize("atomically", [False, True])
+def test_load_dataframes(
+    tmp_path: Path,
+    atomically: bool,  # noqa: FBT001
+) -> None:
     """Test `utils.load_dataframes`."""
     learner = adaptive.Learner1D(lambda x: x, bounds=(-10, 10))
     fname = str(tmp_path / "test.pickle")
-    save_df = utils.save_dataframe(fname)
+    save_df = utils.save_dataframe(fname, atomically=atomically)
 
     learner.ask(10)
     learner.tell(10, 42)
@@ -433,13 +445,97 @@ def test_fname_to_dataframe_with_folder() -> None:
     assert df_fname == Path("test_folder/dataframe.test.parquet")
 
 
-def test_load_dataframes_with_folder(tmp_path: Path) -> None:
+def test_atomic_write(tmp_path: Path) -> None:
+    """Ensure atomic_write works when operated in a basic manner."""
+    path = tmp_path / "testfile"
+    content = "this is some content"
+
+    # Works correctly in basic mode, with no file existing
+    with utils.atomic_write(path) as fp:
+        fp.write(content)
+    with path.open() as fp:
+        assert content == fp.read()
+
+    # Works correctly in basic mode, with the file existing
+    assert path.exists()
+    assert path.is_file()
+    content = "this is some additional content"
+    with utils.atomic_write(path) as fp:
+        fp.write(content)
+    with path.open() as fp:
+        assert content == fp.read()
+
+    # Works correctly when 'return_path' is used.
+    content = "even more content"
+    with utils.atomic_write(path, return_path=True) as tmp_path, tmp_path.open(
+        "w",
+    ) as fp:
+        fp.write(content)
+    with path.open() as fp:
+        assert content == fp.read()
+
+
+def test_atomic_write_no_write(tmp_path: Path) -> None:
+    """Ensure atomic_write creates an empty file, if we do nothing.
+
+    This gives 'atomic_write' the same semantics as 'open', when
+    the file does not exist.
+    """
+    path = tmp_path / "testfile"
+
+    assert not path.exists()  # Sanity check
+
+    with utils.atomic_write(path) as _:
+        pass
+    assert path.stat().st_size == 0
+
+    with utils.atomic_write(path, return_path=True) as _:
+        pass
+    assert path.stat().st_size == 0
+
+    #
+    with utils.atomic_write(path) as fp:
+        fp.write("content")
+    assert path.stat().st_size > 0
+
+    with utils.atomic_write(path) as _:
+        pass
+    assert path.stat().st_size == 0
+
+    with utils.atomic_write(path) as fp:
+        fp.write("content")
+    assert path.stat().st_size > 0
+
+    with utils.atomic_write(path, return_path=True) as _:
+        pass
+    assert path.stat().st_size == 0
+
+
+def test_atomic_write_nested(tmp_path: Path) -> None:
+    """Ensure nested calls to atomic_write on the same file work as expected."""
+    path = tmp_path / "testfile"
+    with utils.atomic_write(path, mode="w") as fp, utils.atomic_write(
+        path,
+        mode="w",
+    ) as fp2:
+        fp.write("one")
+        fp2.write("two")
+    # Outer call wins
+    with path.open() as fp:
+        assert fp.read() == "one"
+
+
+@pytest.mark.parametrize("atomically", [False, True])
+def test_load_dataframes_with_folder(
+    tmp_path: Path,
+    atomically: bool,  # noqa: FBT001
+) -> None:
     """Test `utils.load_dataframes` with `folder`."""
     folder = tmp_path / "test_folder"
     folder.mkdir()
     learner = adaptive.Learner1D(lambda x: x, bounds=(-10, 10))
     fname = str(folder / "test.pickle")
-    save_df = utils.save_dataframe(fname)
+    save_df = utils.save_dataframe(fname, atomically=atomically)
 
     learner.ask(10)
     learner.tell(10, 42)