diff --git a/lelab/record.py b/lelab/record.py index 07e13f27..abef1cdf 100644 --- a/lelab/record.py +++ b/lelab/record.py @@ -331,6 +331,7 @@ 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) finally: if current_phase != "error": current_phase = "completed" @@ -573,19 +574,54 @@ 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 _cleanup_failed_recording(dataset_repo_id: str, resume: bool, saved_episodes: int) -> 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. + """ + if resume 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(): diff --git a/tests/test_record.py b/tests/test_record.py index 486eb959..7e4e9d31 100644 --- a/tests/test_record.py +++ b/tests/test_record.py @@ -46,6 +46,86 @@ 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) + 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) + 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) + assert dataset_dir.exists() + + +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) + + 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).