Save demonstration files atomically

CHANGELOG.md

@@ -17,7 +17,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Fixed
 
-- None.
+- Demonstration safetensors are now saved atomically so an interrupted or
+  failed write does not corrupt an existing demo file.
 
 ## 4.2.0
 

demonstrations/demo.py

@@ -1,6 +1,7 @@
 """Class to hold demo."""
 from __future__ import annotations
 import logging
+import os
 import numpy as np
 from collections import defaultdict
 from copy import deepcopy
@@ -9,6 +10,7 @@
 from safetensors import safe_open
 from safetensors.numpy import save_file
 from typing import Optional, Any, Union, Iterable
+import uuid
 
 from gymnasium.core import ActType
 
@@ -282,7 +284,28 @@ def save(self, path: Union[str, Path]) -> Path:
             for key, val in timesteps[GYM_INFO_KEY].items()
         }
 
-        save_file(demo_dict, path, self.safetensor_metadata)
+        temporary_path = path.with_name(f".{path.name}.{uuid.uuid4().hex}.tmp")
+        try:
+            save_file(demo_dict, temporary_path, self.safetensor_metadata)
+            with temporary_path.open("rb") as temporary_file:
+                os.fsync(temporary_file.fileno())
+            os.replace(temporary_path, path)
+            try:
+                directory_fd = os.open(path.parent, os.O_RDONLY)
+            except OSError:
+                directory_fd = None
+            if directory_fd is not None:
+                try:
+                    os.fsync(directory_fd)
+                except OSError:
+                    pass
+                finally:
+                    os.close(directory_fd)
+        finally:
+            try:
+                temporary_path.unlink()
+            except FileNotFoundError:
+                pass
         logging.info(f"Saved {path}")
         return path
 

tests/test_demos.py

@@ -2,6 +2,7 @@
 import copy
 import pytest
 import tempfile
+from pathlib import Path
 
 from numpy.testing import assert_allclose
 
@@ -12,6 +13,7 @@
 from bigym.envs.manipulation import StackBlocks
 from bigym.utils.observation_config import CameraConfig, ObservationConfig
 from demonstrations.const import ACTION_KEY
+import demonstrations.demo as demo_module
 from demonstrations.demo import Demo
 from demonstrations.demo_recorder import DemoRecorder
 from demonstrations.demo_converter import DemoConverter
@@ -204,3 +206,50 @@ def test_long_running_demo():
         demo = Demo.from_safetensors(filepath)
         assert len(demo.timesteps) == 1000
         TestDemos.assert_replay_observations(env, demo)
+
+
+class _TestMetadata:
+    def ready_for_safetensors(self):
+        return {"source": "atomic-save-test"}
+
+
+def test_failed_demo_save_preserves_existing_file(monkeypatch, tmp_path):
+    target = tmp_path / "demo.safetensors"
+    target.write_bytes(b"existing-demo")
+    temporary_paths = []
+
+    def fail_after_partial_write(_tensors, path, _metadata):
+        temporary_path = Path(path)
+        temporary_paths.append(temporary_path)
+        temporary_path.write_bytes(b"partial-demo")
+        raise RuntimeError("simulated save failure")
+
+    monkeypatch.setattr(demo_module, "save_file", fail_after_partial_write)
+    demo = Demo(_TestMetadata())
+
+    with pytest.raises(RuntimeError, match="simulated save failure"):
+        demo.save(target)
+
+    assert target.read_bytes() == b"existing-demo"
+    assert len(temporary_paths) == 1
+    assert temporary_paths[0].parent == target.parent
+    assert not temporary_paths[0].exists()
+
+
+def test_successful_demo_save_atomically_replaces_existing_file(monkeypatch, tmp_path):
+    target = tmp_path / "demo.safetensors"
+    target.write_bytes(b"existing-demo")
+    temporary_paths = []
+
+    def save_complete_file(_tensors, path, _metadata):
+        temporary_path = Path(path)
+        temporary_paths.append(temporary_path)
+        temporary_path.write_bytes(b"complete-demo")
+
+    monkeypatch.setattr(demo_module, "save_file", save_complete_file)
+    demo = Demo(_TestMetadata())
+
+    assert demo.save(target) == target
+    assert target.read_bytes() == b"complete-demo"
+    assert len(temporary_paths) == 1
+    assert not temporary_paths[0].exists()