diff --git a/frontend/src/components/dataset/DeleteDatasetDialog.test.tsx b/frontend/src/components/dataset/DeleteDatasetDialog.test.tsx new file mode 100644 index 00000000..fe87aca5 --- /dev/null +++ b/frontend/src/components/dataset/DeleteDatasetDialog.test.tsx @@ -0,0 +1,66 @@ +// @vitest-environment jsdom +import React, { act } from "react"; +import { createRoot } from "react-dom/client"; +import { expect, test, vi } from "vitest"; + +import DeleteDatasetDialog from "./DeleteDatasetDialog"; + +const fetchWithHeaders = vi.fn(async (_url: string, _init: { body: string }) => ({ ok: true, json: async () => ({ success: true }) })); + +vi.mock("@/contexts/ApiContext", () => ({ + useApi: () => ({ baseUrl: "http://api", fetchWithHeaders }), +})); +vi.mock("@/hooks/use-toast", () => ({ useToast: () => ({ toast: () => {} }) })); + +const render = async (onHub: boolean, onDeleted = () => {}) => { + Object.assign(globalThis, { IS_REACT_ACT_ENVIRONMENT: true }); + const host = document.createElement("div"); + document.body.append(host); + const root = createRoot(host); + await act(async () => { + root.render( + {}} + onDeleted={onDeleted} + />, + ); + }); + return root; +}; + +const deleteButton = () => + [...document.querySelectorAll("button")].find((b) => b.textContent === "Delete")!; + +test("a local-only dataset offers no Hub checkbox and deletes locally", async () => { + fetchWithHeaders.mockClear(); + const onDeleted = vi.fn(); + const root = await render(false, onDeleted); + expect(document.body.textContent).not.toContain("Also delete this dataset from the HuggingFace Hub"); + + await act(async () => deleteButton().click()); + + expect(JSON.parse(fetchWithHeaders.mock.calls[0][1].body)).toEqual({ + dataset_repo_id: "me/demo", + delete_from_hub: false, + }); + expect(onDeleted).toHaveBeenCalled(); + act(() => root.unmount()); +}); + +test("deleting from the Hub stays disabled until the repo id is typed", async () => { + const root = await render(true); + const checkbox = document.querySelector("#delete-from-hub")!; + await act(async () => checkbox.click()); + expect(deleteButton().disabled).toBe(true); + + const input = document.querySelector("input")!; + await act(async () => { + Object.getOwnPropertyDescriptor(HTMLInputElement.prototype, "value")!.set!.call(input, "me/demo"); + input.dispatchEvent(new Event("input", { bubbles: true })); + }); + expect(deleteButton().disabled).toBe(false); + act(() => root.unmount()); +}); diff --git a/frontend/src/components/dataset/DeleteDatasetDialog.tsx b/frontend/src/components/dataset/DeleteDatasetDialog.tsx new file mode 100644 index 00000000..bd77f6b0 --- /dev/null +++ b/frontend/src/components/dataset/DeleteDatasetDialog.tsx @@ -0,0 +1,139 @@ +import React, { useState } from "react"; +import { + AlertDialog, + AlertDialogAction, + AlertDialogCancel, + AlertDialogContent, + AlertDialogDescription, + AlertDialogFooter, + AlertDialogHeader, + AlertDialogTitle, +} from "@/components/ui/alert-dialog"; +import { Checkbox } from "@/components/ui/checkbox"; +import { Input } from "@/components/ui/input"; +import { Label } from "@/components/ui/label"; +import { useApi } from "@/contexts/ApiContext"; +import { useToast } from "@/hooks/use-toast"; + +interface DeleteDatasetDialogProps { + repoId: string; + onHub: boolean; + open: boolean; + onOpenChange: (open: boolean) => void; + onDeleted: () => void; +} + +/** Confirms and performs the removal of a local dataset, optionally from the Hub too. */ +const DeleteDatasetDialog: React.FC = ({ + repoId, + onHub, + open, + onOpenChange, + onDeleted, +}) => { + const { baseUrl, fetchWithHeaders } = useApi(); + const { toast } = useToast(); + const [isDeleting, setIsDeleting] = useState(false); + const [deleteFromHub, setDeleteFromHub] = useState(false); + const [confirmText, setConfirmText] = useState(""); + + const reset = () => { + setDeleteFromHub(false); + setConfirmText(""); + }; + + const handleDelete = async () => { + setIsDeleting(true); + try { + const response = await fetchWithHeaders(`${baseUrl}/delete-dataset`, { + method: "POST", + body: JSON.stringify({ dataset_repo_id: repoId, delete_from_hub: deleteFromHub }), + }); + const data = await response.json(); + if (response.ok && data.success) { + toast({ + title: "Dataset Deleted", + description: deleteFromHub + ? `${repoId} has been removed from disk and the Hub.` + : `${repoId} has been removed from disk.`, + }); + onDeleted(); + } else { + toast({ + title: "Delete Failed", + description: data.message || "Could not delete the dataset.", + variant: "destructive", + }); + } + } catch { + toast({ + title: "Connection Error", + description: "Could not connect to the backend server.", + variant: "destructive", + }); + } finally { + setIsDeleting(false); + onOpenChange(false); + reset(); + } + }; + + return ( + { + onOpenChange(next); + if (!next) reset(); + }} + > + + + Delete dataset from disk? + + This permanently removes {repoId} from your local cache. This action cannot be undone. + + + {onHub && ( +
+
+ setDeleteFromHub(checked === true)} + /> + +
+ {deleteFromHub && ( +
+ + setConfirmText(e.target.value)} + className="bg-gray-800 border-gray-700 text-white" + /> +
+ )} +
+ )} + + + Keep dataset + + + {isDeleting ? "Deleting…" : "Delete"} + + + + + ); +}; + +export default DeleteDatasetDialog; diff --git a/frontend/src/lib/hubViewer.ts b/frontend/src/lib/hubViewer.ts new file mode 100644 index 00000000..84ad3bf6 --- /dev/null +++ b/frontend/src/lib/hubViewer.ts @@ -0,0 +1,8 @@ +/** Open a Hub dataset in the visualize_dataset Space (private ones go through the login redirect). */ +export const openHubViewer = (repoId: string, isPrivate: boolean) => { + const spacePath = `/spaces/lerobot/visualize_dataset?path=${encodeURIComponent(`/${repoId}`)}`; + const target = isPrivate + ? `https://huggingface.co/login?next=${encodeURIComponent(spacePath)}` + : `https://huggingface.co${spacePath}`; + window.open(target, "_blank", "noopener,noreferrer"); +}; diff --git a/frontend/src/pages/EditDataset.test.tsx b/frontend/src/pages/EditDataset.test.tsx new file mode 100644 index 00000000..aceb3d0a --- /dev/null +++ b/frontend/src/pages/EditDataset.test.tsx @@ -0,0 +1,63 @@ +// @vitest-environment jsdom +import React, { act } from "react"; +import { createRoot } from "react-dom/client"; +import { MemoryRouter } from "react-router-dom"; +import { expect, test, vi } from "vitest"; + +import EditDataset from "./EditDataset"; + +const datasets = [ + { repo_id: "me/local", source: "local", last_modified: "", private: false }, + { repo_id: "me/pushed", source: "both", last_modified: "", private: false }, +]; + +vi.mock("@/hooks/useDatasets", () => ({ + useDatasets: () => ({ datasets, loading: false, refresh: () => {} }), +})); +vi.mock("@/hooks/useEpisodes", () => ({ + useEpisodes: () => ({ data: null, loading: false, error: null }), + useEpisodeDetail: () => ({ detail: null, loading: false, error: null }), +})); +vi.mock("@/hooks/useRobots", () => ({ useRobots: () => ({ selectedRecord: null }) })); +vi.mock("@/hooks/useRecording", () => ({ + useRecording: () => ({ openForNew: () => {}, openToAppend: () => {}, modalProps: {} }), +})); +vi.mock("@/hooks/use-toast", () => ({ useToast: () => ({ toast: () => {} }) })); +vi.mock("@/contexts/ApiContext", () => ({ + useApi: () => ({ baseUrl: "http://api", fetchWithHeaders: vi.fn() }), +})); +vi.mock("@/components/landing/RecordingModal", () => ({ default: () => null })); + +const renderAt = async (repoId: string) => { + Object.assign(globalThis, { IS_REACT_ACT_ENVIRONMENT: true }); + const host = document.createElement("div"); + document.body.append(host); + const root = createRoot(host); + await act(async () => { + root.render( + + + , + ); + }); + return { host, root }; +}; + +const buttonTexts = (host: HTMLElement) => + [...host.querySelectorAll("button, a")].map((e) => e.textContent?.trim()); + +test("a dataset not on the Hub offers Upload and Delete", async () => { + const { host, root } = await renderAt("me/local"); + expect(buttonTexts(host)).toContain("Upload"); + expect(host.querySelector('[aria-label="Delete dataset from disk"]')).not.toBeNull(); + expect(host.querySelector('a[href^="https://huggingface.co"]')).toBeNull(); + act(() => root.unmount()); +}); + +test("a dataset already on the Hub links to its Hub page instead of Upload", async () => { + const { host, root } = await renderAt("me/pushed"); + expect(buttonTexts(host)).not.toContain("Upload"); + expect(host.querySelector("a")?.getAttribute("href")).toBe("https://huggingface.co/datasets/me/pushed"); + expect(host.querySelector('[aria-label="Delete dataset from disk"]')).not.toBeNull(); + act(() => root.unmount()); +}); diff --git a/frontend/src/pages/EditDataset.tsx b/frontend/src/pages/EditDataset.tsx index 0b55ff80..53ec628a 100644 --- a/frontend/src/pages/EditDataset.tsx +++ b/frontend/src/pages/EditDataset.tsx @@ -1,17 +1,23 @@ import React, { useEffect, useMemo, useState } from "react"; import { useNavigate, useSearchParams } from "react-router-dom"; -import { AlertCircle, ArrowLeft, Database, Disc, Loader2, Merge, Upload } from "lucide-react"; import { - Select, - SelectContent, - SelectItem, - SelectTrigger, - SelectValue, -} from "@/components/ui/select"; + AlertCircle, + ArrowLeft, + ChevronsUpDown, + Database, + Disc, + ExternalLink, + Loader2, + Merge, + Trash2, + Upload, +} from "lucide-react"; import { Button } from "@/components/ui/button"; +import DeleteDatasetDialog from "@/components/dataset/DeleteDatasetDialog"; import EpisodeList from "@/components/dataset/EpisodeList"; import EpisodeViewer from "@/components/dataset/EpisodeViewer"; import MergeDatasetsDialog from "@/components/dataset/MergeDatasetsDialog"; +import DatasetPicker from "@/components/landing/DatasetPicker"; import RecordingModal from "@/components/landing/RecordingModal"; import { useApi } from "@/contexts/ApiContext"; import { useDatasets } from "@/hooks/useDatasets"; @@ -20,7 +26,8 @@ import { useRecording } from "@/hooks/useRecording"; import { useRobots } from "@/hooks/useRobots"; import { useToast } from "@/hooks/use-toast"; import { formatDuration } from "@/lib/datasetApi"; -import { mergeLocalDatasets } from "@/lib/replayApi"; +import { openHubViewer } from "@/lib/hubViewer"; +import { DatasetItem, mergeLocalDatasets } from "@/lib/replayApi"; // The robot_type LeRobot stores for the SO-101 follower. const FOLLOWER_ROBOT_TYPE = "so_follower"; @@ -47,6 +54,7 @@ const EditDataset = () => { const { baseUrl, fetchWithHeaders } = useApi(); const { toast } = useToast(); const [mergeOpen, setMergeOpen] = useState(false); + const [deleteOpen, setDeleteOpen] = useState(false); // A Hub-only dataset has no videos on disk to decode. const localDatasets = useMemo( () => datasets.filter((d) => d.source === "local" || d.source === "both"), @@ -91,6 +99,11 @@ const EditDataset = () => { return api ? { ...params, api } : params; }; const setDataset = (next: string) => setSearchParams(withApi({ dataset: next })); + // A Hub-only dataset has no videos on disk to browse, so it opens in the Hub viewer. + const handlePickExisting = (item: DatasetItem) => { + if (item.source === "hub") openHubViewer(item.repo_id, item.private); + else setDataset(item.repo_id); + }; const handleMerge = async (sourceRepoIds: string[], outputName: string) => { const result = await mergeLocalDatasets(baseUrl, fetchWithHeaders, sourceRepoIds, outputName); refresh(); @@ -133,20 +146,24 @@ const EditDataset = () => {
- + openHubViewer(id, true)} + onCreateNew={recording.openForNew} + > + +
{repoId && selectedIsLocal && ( )} - {/* Picking a dataset now lands here rather than on /upload, so the - upload + delete flow hangs off the page you browse from. */} - {repoId && ( + {repoId && selectedSource === "both" && ( + + )} + {repoId && selectedSource !== "both" && ( )} + {repoId && selectedIsLocal && ( + + )} {appendTarget && (
); diff --git a/frontend/src/pages/Landing.tsx b/frontend/src/pages/Landing.tsx index b27ca705..748f3e5e 100644 --- a/frontend/src/pages/Landing.tsx +++ b/frontend/src/pages/Landing.tsx @@ -14,6 +14,7 @@ import { useRobots } from "@/hooks/useRobots"; import { useDatasets } from "@/hooks/useDatasets"; import { useRecording } from "@/hooks/useRecording"; import { DatasetItem } from "@/lib/replayApi"; +import { openHubViewer } from "@/lib/hubViewer"; import { isHostedSpace } from "@/lib/isHostedSpace"; const ON_SPACE = isHostedSpace(); @@ -39,14 +40,6 @@ const Landing = () => { const handleTrainingClick = () => navigate("/training"); - const openHubViewer = (repoId: string, isPrivate: boolean) => { - const spacePath = `/spaces/lerobot/visualize_dataset?path=${encodeURIComponent(`/${repoId}`)}`; - const target = isPrivate - ? `https://huggingface.co/login?next=${encodeURIComponent(spacePath)}` - : `https://huggingface.co${spacePath}`; - window.open(target, "_blank", "noopener,noreferrer"); - }; - const handlePickExisting = (item: DatasetItem) => { // Anything with files on disk opens in the local episode browser; upload and // delete are reachable from there. A Hub-only dataset has no videos here to diff --git a/frontend/src/pages/Upload.tsx b/frontend/src/pages/Upload.tsx index d620913f..d74e5734 100644 --- a/frontend/src/pages/Upload.tsx +++ b/frontend/src/pages/Upload.tsx @@ -3,8 +3,8 @@ import { useNavigate, useLocation } from "react-router-dom"; import { Button } from "@/components/ui/button"; import { Input } from "@/components/ui/input"; import { Label } from "@/components/ui/label"; -import { useToast } from "@/hooks/use-toast"; import { Checkbox } from "@/components/ui/checkbox"; +import { useToast } from "@/hooks/use-toast"; import { ArrowLeft, Upload as UploadIcon, @@ -20,16 +20,7 @@ import { } from "lucide-react"; import { useApi } from "@/contexts/ApiContext"; import { DatasetSource } from "@/lib/replayApi"; -import { - AlertDialog, - AlertDialogAction, - AlertDialogCancel, - AlertDialogContent, - AlertDialogDescription, - AlertDialogFooter, - AlertDialogHeader, - AlertDialogTitle, -} from "@/components/ui/alert-dialog"; +import DeleteDatasetDialog from "@/components/dataset/DeleteDatasetDialog"; interface DatasetInfo { dataset_repo_id: string; @@ -71,7 +62,6 @@ const Upload = () => { const [isUploading, setIsUploading] = useState(false); const [uploadSuccess, setUploadSuccess] = useState(false); const [showDeleteConfirm, setShowDeleteConfirm] = useState(false); - const [isDeleting, setIsDeleting] = useState(false); // Load actual dataset information from backend React.useEffect(() => { @@ -225,40 +215,6 @@ const Upload = () => { navigate("/"); }; - const handleDeleteDataset = async () => { - if (!datasetInfo) return; - setIsDeleting(true); - try { - const response = await fetchWithHeaders(`${baseUrl}/delete-dataset`, { - method: "POST", - body: JSON.stringify({ dataset_repo_id: datasetInfo.dataset_repo_id }), - }); - const data = await response.json(); - if (response.ok && data.success) { - toast({ - title: "Dataset Deleted", - description: `${datasetInfo.dataset_repo_id} has been removed from disk.`, - }); - navigate("/"); - } else { - toast({ - title: "Delete Failed", - description: data.message || "Could not delete the dataset.", - variant: "destructive", - }); - } - } catch (error) { - toast({ - title: "Connection Error", - description: "Could not connect to the backend server.", - variant: "destructive", - }); - } finally { - setIsDeleting(false); - setShowDeleteConfirm(false); - } - }; - // Show loading state while fetching dataset info if (isLoadingDatasetInfo || !datasetInfo) { return ( @@ -291,7 +247,6 @@ const Upload = () => { onClick={() => setShowDeleteConfirm(true)} variant="outline" size="icon" - disabled={isDeleting} aria-label="Delete dataset from disk" className="border-red-500/40 text-red-400 hover:border-red-400 hover:text-red-300 hover:bg-red-500/10" > @@ -536,28 +491,13 @@ const Upload = () => { )} - - - - Delete dataset from disk? - - This permanently removes {datasetInfo.dataset_repo_id} from your local cache. This action cannot be undone. - - - - - Keep dataset - - - {isDeleting ? "Deleting…" : "Delete"} - - - - + navigate("/")} + /> ); }; diff --git a/lelab/record.py b/lelab/record.py index 0ad2ca92..f6257973 100644 --- a/lelab/record.py +++ b/lelab/record.py @@ -91,6 +91,7 @@ class UploadRequest(BaseModel): class DatasetInfoRequest(BaseModel): dataset_repo_id: str + delete_from_hub: bool = False def _normalize_dataset_repo_id(repo_id: str) -> str: @@ -693,6 +694,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"Hub deletion of {repo_id} failed: {e}") + return {"success": False, "message": f"Could not remove the dataset on the Hub: {e}"} + try: shutil.rmtree(target) except Exception as e: diff --git a/tests/test_record.py b/tests/test_record.py index 1c4aef18..dc97c8b8 100644 --- a/tests/test_record.py +++ b/tests/test_record.py @@ -16,6 +16,7 @@ from __future__ import annotations import shutil +from unittest.mock import MagicMock import pytest @@ -404,3 +405,61 @@ def test_resume_refuses_an_unreadable_dataset_before_the_robot( assert result["success"] is False assert "record it again" in result["message"] assert record.recording_active is False + + +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()