diff --git a/src/lerobot/datasets/dataset_writer.py b/src/lerobot/datasets/dataset_writer.py index 5693e9bb2..90cbf5f2d 100644 --- a/src/lerobot/datasets/dataset_writer.py +++ b/src/lerobot/datasets/dataset_writer.py @@ -386,6 +386,9 @@ class DatasetWriter: self._episodes_since_last_encoding = 0 if episode_data is None: + # Post-save cleanup deliberately does not go through clear_episode_buffer(): + # staging frames of video cameras must survive here — the (possibly batched) + # encoder still needs them and deletes them once each video is written. if len(self._meta.image_keys) > 0: self._delete_camera_frame_dirs(self._meta.image_keys) self.episode_buffer = self._create_episode_buffer() diff --git a/tests/datasets/test_dataset_writer.py b/tests/datasets/test_dataset_writer.py index bd51a2687..db0282b4e 100644 --- a/tests/datasets/test_dataset_writer.py +++ b/tests/datasets/test_dataset_writer.py @@ -236,6 +236,42 @@ def test_clear_removes_video_frame_staging_dir(tmp_path): assert not video_staging_dir.exists() +def test_batched_encoding_staging_survives_save(tmp_path): + """The post-save clear must NOT delete video staging frames. + + With ``batch_encoding_size > 1`` the frames of already-saved episodes stay + on disk until the batch encode runs; the encoder deletes them afterwards. + A blanket switch of the post-save cleanup to ``camera_keys`` (as done in the + discard path) would silently break batched encoding. + """ + video_key = "observation.images.cam" + features = { + video_key: { + "dtype": "video", + "shape": (64, 96, 3), + "names": ["height", "width", "channels"], + }, + "action": {"dtype": "float32", "shape": (2,), "names": None}, + } + dataset = LeRobotDataset.create( + repo_id=DUMMY_REPO_ID, + fps=DEFAULT_FPS, + features=features, + root=tmp_path / "ds", + use_videos=True, + batch_encoding_size=2, + ) + for _ in range(3): + dataset.add_frame(_make_frame(features)) + + staging_dir = dataset.writer._get_image_file_dir(0, video_key) + assert staging_dir.is_dir() + + dataset.save_episode() # first of a batch of 2: no encoding yet + + assert staging_dir.is_dir() and any(staging_dir.iterdir()) + + def test_finalize_is_idempotent(tmp_path): """Calling finalize() twice does not raise.""" dataset = LeRobotDataset.create(