test(datasets): guard batched-encoding video staging from post-save cleanup (#4110)

The discard-path fix (#3683) deletes staging for all camera_keys in
clear_episode_buffer(). The post-save cleanup in save_episode() must keep
iterating image_keys only: with batch_encoding_size > 1 the video staging
frames of already-saved episodes stay on disk until the batch encoder
consumes and deletes them. Add a regression test pinning that behavior,
plus a comment explaining why the two paths differ.

Co-authored-by: Steven Palma <imstevenpmwork@ieee.org>

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()

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(