diff --git a/CHANGELOG.md b/CHANGELOG.md index 2230e87..6b143e2 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/demonstrations/demo.py b/demonstrations/demo.py index 42bcf3d..f365a65 100644 --- a/demonstrations/demo.py +++ b/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 diff --git a/tests/test_demos.py b/tests/test_demos.py index 670c587..a8421ae 100644 --- a/tests/test_demos.py +++ b/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()