fix(record): remove the empty dataset dir a failed recording leaves behind

lelab/episode_media.py

@@ -68,9 +68,17 @@ def lerobot_cache_root() -> Path:
     Reads the environment at call time rather than importing lerobot's
     ``HF_LEROBOT_HOME`` constant, which is frozen at import. Same source of
     truth as ``datasets._lerobot_cache_root``, and it keeps this module free of
-    a torch-pulling import it doesn't otherwise need.
+    a torch-pulling import it doesn't otherwise need. Mirrors lerobot's
+    resolution: ``HF_LEROBOT_HOME``, else ``$HF_HOME/lerobot``, where ``HF_HOME``
+    defaults to ``$XDG_CACHE_HOME/huggingface`` (``~/.cache/huggingface``).
     """
-    return Path(os.environ.get("HF_LEROBOT_HOME", "~/.cache/huggingface/lerobot")).expanduser().resolve()
+    explicit = os.environ.get("HF_LEROBOT_HOME")
+    if explicit:
+        return Path(explicit).expanduser().resolve()
+    hf_home = os.environ.get("HF_HOME")
+    if not hf_home:
+        hf_home = os.path.join(os.environ.get("XDG_CACHE_HOME") or "~/.cache", "huggingface")
+    return (Path(hf_home).expanduser() / "lerobot").resolve()
 
 
 def resolve_dataset_dir(repo_id: str) -> Path:

lelab/record.py

@@ -284,8 +284,12 @@ def recording_worker():
             recording_start_time = time.time()
             current_episode = 1
             saved_episodes = 0
+            # Assume the dir isn't ours until the probe says otherwise, so a
+            # failing probe can never make cleanup delete existing data.
+            dir_preexisted = True
 
             try:
+                dir_preexisted = _dataset_dir_preexisted(request.dataset_repo_id)
                 logger.info(
                     "Recording session started: dataset=%s task=%r episodes=%d",
                     request.dataset_repo_id,
@@ -331,6 +335,9 @@ def recording_worker():
                     "dataset_repo_id": request.dataset_repo_id,
                     "saved_episodes": saved_episodes,
                 }
+                _cleanup_failed_recording(
+                    request.dataset_repo_id, request.resume, saved_episodes, dir_preexisted
+                )
             finally:
                 if current_phase != "error":
                     current_phase = "completed"
@@ -573,19 +580,73 @@ def handle_get_dataset_info(request: DatasetInfoRequest) -> dict[str, Any]:
         }
 
 
+def _resolve_dataset_dir(repo_id: str):
+    """Resolve `repo_id` to a path strictly inside the local dataset cache, or None.
+
+    Shared traversal guard for anything that deletes a dataset directory by
+    repo_id coming from a request rather than a local scan. Uses
+    `episode_media.lerobot_cache_root()`, which reads the cache root from the
+    environment at call time rather than importing lerobot's `HF_LEROBOT_HOME`
+    constant (frozen at import, so it wouldn't see a root set after startup).
+    """
+    from . import episode_media
+
+    root = episode_media.lerobot_cache_root()
+    target = (root / repo_id).resolve()
+    if target == root or root not in target.parents:
+        return None
+    return target
+
+
+def _dataset_dir_preexisted(dataset_repo_id: str) -> bool:
+    """Whether the dataset dir already exists; True (conservative) if it can't be probed."""
+    if not dataset_repo_id:
+        return False
+    try:
+        target = _resolve_dataset_dir(dataset_repo_id)
+        return target is not None and target.exists()
+    except Exception as e:
+        logger.warning(f"Could not probe dataset directory for {dataset_repo_id}: {e}")
+        return True
+
+
+def _cleanup_failed_recording(
+    dataset_repo_id: str, resume: bool, saved_episodes: int, dir_preexisted: bool
+) -> None:
+    """Remove a dataset dir a failed recording created but never saved into.
+
+    A dataset dir is written (meta/info.json etc.) as soon as recording
+    starts, before the first episode is captured. If the session errors out
+    before saving anything and wasn't resuming an existing dataset, that dir
+    is an empty leftover the failure produced, not a partial recording worth
+    keeping — remove it instead of leaving it in the local cache and dataset
+    picker. Best-effort: logs and swallows failures rather than masking the
+    original recording error.
+
+    `dir_preexisted` is whether the dir already existed before this attempt
+    started. The dataset name only has second resolution, so a collision makes
+    `LeRobotDataset.create()` fail with FileExistsError on a dir this attempt
+    doesn't own — never delete it.
+    """
+    if resume or dir_preexisted or saved_episodes != 0 or not dataset_repo_id:
+        return
+    target = _resolve_dataset_dir(dataset_repo_id)
+    if target is None or not target.exists():
+        return
+    try:
+        shutil.rmtree(target)
+        logger.info(f"Removed empty dataset directory {target} after a failed recording")
+    except Exception as e:
+        logger.warning(f"Could not remove empty dataset directory {target}: {e}")
+
+
 def handle_delete_dataset(request: DatasetInfoRequest) -> dict[str, Any]:
     """Remove a recorded dataset's directory from local disk."""
     global last_recording_info
-    from pathlib import Path
-
-    from lerobot.utils.constants import HF_LEROBOT_HOME
 
     repo_id = request.dataset_repo_id
-    root = Path(HF_LEROBOT_HOME).resolve()
-    target = (root / repo_id).resolve()
-
-    # Reject path traversal: target must stay strictly inside HF_LEROBOT_HOME.
-    if target == root or root not in target.parents:
+    target = _resolve_dataset_dir(repo_id)
+    if target is None:
         return {"success": False, "message": "Invalid dataset path"}
 
     if not target.exists():

tests/test_record.py

@@ -46,6 +46,171 @@ def test_handle_stop_recording_when_idle_returns_dict(tmp_lerobot_home) -> None:
     assert isinstance(result, dict)
 
 
+def test_resolve_dataset_dir_rejects_traversal(tmp_lerobot_home) -> None:
+    from lelab.record import _resolve_dataset_dir
+
+    assert _resolve_dataset_dir("../../etc") is None
+    assert _resolve_dataset_dir(".") is None
+
+
+def test_resolve_dataset_dir_accepts_nested_repo_id(tmp_lerobot_home) -> None:
+    from lelab.record import _resolve_dataset_dir
+
+    target = _resolve_dataset_dir("alice/pusht")
+    assert target == tmp_lerobot_home / "alice" / "pusht"
+
+
+def test_handle_delete_dataset_rejects_traversal(tmp_lerobot_home) -> None:
+    from lelab.record import DatasetInfoRequest, handle_delete_dataset
+
+    result = handle_delete_dataset(DatasetInfoRequest(dataset_repo_id="../../etc"))
+    assert result["success"] is False
+
+
+def test_handle_delete_dataset_reports_missing(tmp_lerobot_home) -> None:
+    from lelab.record import DatasetInfoRequest, handle_delete_dataset
+
+    result = handle_delete_dataset(DatasetInfoRequest(dataset_repo_id="nope/nope"))
+    assert result["success"] is False
+
+
+def test_handle_delete_dataset_removes_directory(tmp_lerobot_home) -> None:
+    from lelab.record import DatasetInfoRequest, handle_delete_dataset
+
+    dataset_dir = tmp_lerobot_home / "alice" / "pusht"
+    (dataset_dir / "meta").mkdir(parents=True)
+    (dataset_dir / "meta" / "info.json").write_text("{}")
+
+    result = handle_delete_dataset(DatasetInfoRequest(dataset_repo_id="alice/pusht"))
+    assert result["success"] is True
+    assert not dataset_dir.exists()
+
+
+def test_cleanup_failed_recording_removes_empty_new_dataset(tmp_lerobot_home) -> None:
+    from lelab.record import _cleanup_failed_recording
+
+    dataset_dir = tmp_lerobot_home / "alice" / "pusht"
+    (dataset_dir / "meta").mkdir(parents=True)
+    (dataset_dir / "meta" / "info.json").write_text("{}")
+
+    _cleanup_failed_recording("alice/pusht", resume=False, saved_episodes=0, dir_preexisted=False)
+    assert not dataset_dir.exists()
+
+
+def test_cleanup_failed_recording_keeps_dataset_with_saved_episodes(tmp_lerobot_home) -> None:
+    from lelab.record import _cleanup_failed_recording
+
+    dataset_dir = tmp_lerobot_home / "alice" / "pusht"
+    (dataset_dir / "meta").mkdir(parents=True)
+    (dataset_dir / "meta" / "info.json").write_text("{}")
+
+    _cleanup_failed_recording("alice/pusht", resume=False, saved_episodes=2, dir_preexisted=False)
+    assert dataset_dir.exists()
+
+
+def test_cleanup_failed_recording_keeps_resumed_dataset(tmp_lerobot_home) -> None:
+    from lelab.record import _cleanup_failed_recording
+
+    dataset_dir = tmp_lerobot_home / "alice" / "pusht"
+    (dataset_dir / "meta").mkdir(parents=True)
+    (dataset_dir / "meta" / "info.json").write_text("{}")
+
+    _cleanup_failed_recording("alice/pusht", resume=True, saved_episodes=0, dir_preexisted=False)
+    assert dataset_dir.exists()
+
+
+def test_cleanup_failed_recording_keeps_preexisting_dir(tmp_lerobot_home) -> None:
+    from lelab.record import _cleanup_failed_recording
+
+    # Second-resolution name collision: create() raised FileExistsError on a dir
+    # this attempt doesn't own, so its contents must survive.
+    dataset_dir = tmp_lerobot_home / "alice" / "pusht"
+    (dataset_dir / "meta").mkdir(parents=True)
+    (dataset_dir / "meta" / "info.json").write_text("{}")
+
+    _cleanup_failed_recording("alice/pusht", resume=False, saved_episodes=0, dir_preexisted=True)
+    assert (dataset_dir / "meta" / "info.json").exists()
+
+
+def test_resolve_dataset_dir_follows_hf_home_when_lerobot_home_unset(
+    monkeypatch: pytest.MonkeyPatch, tmp_path
+) -> None:
+    from lelab.record import _resolve_dataset_dir
+
+    monkeypatch.delenv("HF_LEROBOT_HOME", raising=False)
+    monkeypatch.setenv("HF_HOME", str(tmp_path))
+
+    assert _resolve_dataset_dir("alice/pusht") == (tmp_path / "lerobot" / "alice" / "pusht").resolve()
+
+
+def test_cleanup_failed_recording_uses_hf_home_cache(monkeypatch: pytest.MonkeyPatch, tmp_path) -> None:
+    from lelab.record import _cleanup_failed_recording
+
+    monkeypatch.delenv("HF_LEROBOT_HOME", raising=False)
+    monkeypatch.setenv("HF_HOME", str(tmp_path))
+    dataset_dir = tmp_path / "lerobot" / "alice" / "pusht"
+    (dataset_dir / "meta").mkdir(parents=True)
+
+    _cleanup_failed_recording("alice/pusht", resume=False, saved_episodes=0, dir_preexisted=False)
+    assert not dataset_dir.exists()
+
+
+def test_dataset_dir_preexisted_defaults_true_when_probe_fails(
+    monkeypatch: pytest.MonkeyPatch, tmp_lerobot_home
+) -> None:
+    from lelab import record
+
+    def boom(repo_id):
+        raise OSError("Too many levels of symbolic links")
+
+    monkeypatch.setattr(record, "_resolve_dataset_dir", boom)
+    assert record._dataset_dir_preexisted("alice/pusht") is True
+
+
+def test_worker_releases_active_state_when_path_probe_fails(
+    monkeypatch: pytest.MonkeyPatch, tmp_lerobot_home
+) -> None:
+    from unittest.mock import MagicMock
+
+    from lelab import record
+
+    def boom(repo_id):
+        raise OSError("Too many levels of symbolic links")
+
+    def fail_recording(cfg, events):
+        raise RuntimeError("recording failed")
+
+    monkeypatch.setattr(record, "_resolve_dataset_dir", boom)
+    monkeypatch.setattr(record, "create_record_config", lambda request: MagicMock())
+    monkeypatch.setattr(record, "record_with_web_events", fail_recording)
+
+    request = record.RecordingRequest(
+        leader_port="/dev/null",
+        follower_port="/dev/null",
+        leader_config="l",
+        follower_config="f",
+        dataset_repo_id="alice/pusht",
+        single_task="t",
+    )
+    result = record.handle_start_recording(request)
+    assert result["success"] is True
+    record.recording_thread.join(timeout=10)
+
+    assert not record.recording_thread.is_alive()
+    assert record.recording_active is False
+    # The next start isn't rejected as "already active".
+    assert record.handle_start_recording(request)["success"] is True
+    record.recording_thread.join(timeout=10)
+    assert record.recording_active is False
+
+
+def test_cleanup_failed_recording_tolerates_missing_dir(tmp_lerobot_home) -> None:
+    from lelab.record import _cleanup_failed_recording
+
+    # No dataset dir was ever created (e.g. failure before LeRobotDataset.create()) — no-op, no crash.
+    _cleanup_failed_recording("alice/never-created", resume=False, saved_episodes=0, dir_preexisted=False)
+
+
 def test_create_record_config_pins_dshow_on_windows(monkeypatch: pytest.MonkeyPatch) -> None:
     """On Windows, recording must use the DSHOW backend so a camera_index opens
     the same device /available-cameras enumerated (via pygrabber, DSHOW order).