diff --git a/frontend/src/pages/Upload.tsx b/frontend/src/pages/Upload.tsx index d620913f..69b932d6 100644 --- a/frontend/src/pages/Upload.tsx +++ b/frontend/src/pages/Upload.tsx @@ -72,6 +72,8 @@ const Upload = () => { const [uploadSuccess, setUploadSuccess] = useState(false); const [showDeleteConfirm, setShowDeleteConfirm] = useState(false); const [isDeleting, setIsDeleting] = useState(false); + const [deleteFromHub, setDeleteFromHub] = useState(false); + const [deleteConfirmText, setDeleteConfirmText] = useState(""); // Load actual dataset information from backend React.useEffect(() => { @@ -231,13 +233,18 @@ const Upload = () => { try { const response = await fetchWithHeaders(`${baseUrl}/delete-dataset`, { method: "POST", - body: JSON.stringify({ dataset_repo_id: datasetInfo.dataset_repo_id }), + body: JSON.stringify({ + dataset_repo_id: datasetInfo.dataset_repo_id, + delete_from_hub: deleteFromHub, + }), }); const data = await response.json(); if (response.ok && data.success) { toast({ title: "Dataset Deleted", - description: `${datasetInfo.dataset_repo_id} has been removed from disk.`, + description: deleteFromHub + ? `${datasetInfo.dataset_repo_id} has been removed from disk and the Hub.` + : `${datasetInfo.dataset_repo_id} has been removed from disk.`, }); navigate("/"); } else { @@ -256,6 +263,8 @@ const Upload = () => { } finally { setIsDeleting(false); setShowDeleteConfirm(false); + setDeleteFromHub(false); + setDeleteConfirmText(""); } }; @@ -536,7 +545,16 @@ const Upload = () => { )} - + { + setShowDeleteConfirm(open); + if (!open) { + setDeleteFromHub(false); + setDeleteConfirmText(""); + } + }} + > Delete dataset from disk? @@ -544,13 +562,42 @@ const Upload = () => { This permanently removes {datasetInfo.dataset_repo_id} from your local cache. This action cannot be undone. + {isAlreadyOnHub && ( +
+
+ setDeleteFromHub(checked === true)} + /> + +
+ {deleteFromHub && ( +
+ + setDeleteConfirmText(e.target.value)} + className="bg-gray-800 border-gray-700 text-white" + /> +
+ )} +
+ )} Keep dataset {isDeleting ? "Deleting…" : "Delete"} diff --git a/lelab/record.py b/lelab/record.py index 07e13f27..ef237832 100644 --- a/lelab/record.py +++ b/lelab/record.py @@ -89,6 +89,7 @@ class UploadRequest(BaseModel): class DatasetInfoRequest(BaseModel): dataset_repo_id: str + delete_from_hub: bool = False def _platform_backend(): @@ -591,6 +592,18 @@ def handle_delete_dataset(request: DatasetInfoRequest) -> dict[str, Any]: if not target.exists(): return {"success": False, "message": f"Dataset not found on disk: {repo_id}"} + if request.delete_from_hub: + from huggingface_hub import delete_repo + from huggingface_hub.errors import RepositoryNotFoundError + + try: + delete_repo(repo_id, repo_type="dataset") + except RepositoryNotFoundError: + pass + except Exception as e: + logger.error(f"Failed to delete {repo_id} from the Hub: {e}") + return {"success": False, "message": f"Failed to delete dataset from the Hub: {e}"} + try: shutil.rmtree(target) except Exception as e: diff --git a/tests/test_record.py b/tests/test_record.py index 486eb959..50bf3d15 100644 --- a/tests/test_record.py +++ b/tests/test_record.py @@ -15,6 +15,8 @@ from __future__ import annotations +from unittest.mock import MagicMock + import pytest @@ -119,3 +121,61 @@ def test_build_camera_configs_skips_non_opencv_type() -> None: configs = _build_camera_configs(cameras, Cv2Backends.ANY) assert configs == {} + + +def test_handle_delete_dataset_calls_delete_repo_when_requested( + tmp_lerobot_home, monkeypatch: pytest.MonkeyPatch +) -> None: + from lelab.record import DatasetInfoRequest, handle_delete_dataset + + monkeypatch.setattr("lerobot.utils.constants.HF_LEROBOT_HOME", tmp_lerobot_home) + dataset_dir = tmp_lerobot_home / "user" / "dataset" + dataset_dir.mkdir(parents=True) + + spy = MagicMock() + monkeypatch.setattr("huggingface_hub.delete_repo", spy) + + result = handle_delete_dataset(DatasetInfoRequest(dataset_repo_id="user/dataset", delete_from_hub=True)) + + assert result["success"] is True + spy.assert_called_once_with("user/dataset", repo_type="dataset") + assert not dataset_dir.exists() + + +def test_handle_delete_dataset_skips_hub_call_by_default( + tmp_lerobot_home, monkeypatch: pytest.MonkeyPatch +) -> None: + from lelab.record import DatasetInfoRequest, handle_delete_dataset + + monkeypatch.setattr("lerobot.utils.constants.HF_LEROBOT_HOME", tmp_lerobot_home) + dataset_dir = tmp_lerobot_home / "user" / "dataset" + dataset_dir.mkdir(parents=True) + + spy = MagicMock() + monkeypatch.setattr("huggingface_hub.delete_repo", spy) + + result = handle_delete_dataset(DatasetInfoRequest(dataset_repo_id="user/dataset")) + + assert result["success"] is True + spy.assert_not_called() + assert not dataset_dir.exists() + + +def test_handle_delete_dataset_hub_failure_keeps_local_copy( + tmp_lerobot_home, monkeypatch: pytest.MonkeyPatch +) -> None: + from lelab.record import DatasetInfoRequest, handle_delete_dataset + + monkeypatch.setattr("lerobot.utils.constants.HF_LEROBOT_HOME", tmp_lerobot_home) + dataset_dir = tmp_lerobot_home / "user" / "dataset" + dataset_dir.mkdir(parents=True) + + def raise_error(*args, **kwargs): + raise RuntimeError("network error") + + monkeypatch.setattr("huggingface_hub.delete_repo", raise_error) + + result = handle_delete_dataset(DatasetInfoRequest(dataset_repo_id="user/dataset", delete_from_hub=True)) + + assert result["success"] is False + assert dataset_dir.exists()