From d6ba7874d580fa44ec0fc77c0be7e6c60cb76f3c Mon Sep 17 00:00:00 2001 From: "ravindu.somawansa" Date: Mon, 19 Oct 2026 23:09:54 +0200 Subject: [PATCH 1/2] feat(datasets): move episodes out of a local dataset into a new one --- .../dataset/MoveEpisodesDialog.test.tsx | 74 +++++++++ .../components/dataset/MoveEpisodesDialog.tsx | 144 ++++++++++++++++++ frontend/src/hooks/useEpisodes.ts | 9 +- frontend/src/lib/datasetApi.ts | 19 +++ frontend/src/lib/replayApi.ts | 21 +++ frontend/src/pages/EditDataset.tsx | 41 ++++- lelab/datasets.py | 67 +++++++- lelab/server.py | 32 +++- tests/test_datasets.py | 139 +++++++++++++++++ 9 files changed, 535 insertions(+), 11 deletions(-) create mode 100644 frontend/src/components/dataset/MoveEpisodesDialog.test.tsx create mode 100644 frontend/src/components/dataset/MoveEpisodesDialog.tsx diff --git a/frontend/src/components/dataset/MoveEpisodesDialog.test.tsx b/frontend/src/components/dataset/MoveEpisodesDialog.test.tsx new file mode 100644 index 000000000..1a12bb9be --- /dev/null +++ b/frontend/src/components/dataset/MoveEpisodesDialog.test.tsx @@ -0,0 +1,74 @@ +// @vitest-environment jsdom +import React, { act } from "react"; +import { createRoot } from "react-dom/client"; +import { expect, test } from "vitest"; + +import { parseEpisodeList } from "@/lib/datasetApi"; + +import MoveEpisodesDialog from "./MoveEpisodesDialog"; + +test("parses episode lists and ranges", () => { + expect(parseEpisodeList("3, 1, 5-7, 3", 10)).toEqual([1, 3, 5, 6, 7]); + expect(parseEpisodeList(" 2 - 4 ,", 10)).toEqual([2, 3, 4]); + expect(parseEpisodeList("", 10)).toBeNull(); + expect(parseEpisodeList("10", 10)).toBeNull(); + expect(parseEpisodeList("4-2", 10)).toBeNull(); + expect(parseEpisodeList("1, x", 10)).toBeNull(); +}); + +test("opens on the browsed episode and refuses to move every episode", async () => { + Object.assign(globalThis, { IS_REACT_ACT_ENVIRONMENT: true }); + const host = document.createElement("div"); + document.body.append(host); + const root = createRoot(host); + const moved: [number[], string][] = []; + + try { + await act(async () => { + root.render( + {}} + onMove={async (indices, name) => { + moved.push([indices, name]); + }} + />, + ); + }); + + const episodes = document.getElementById("move-episodes") as HTMLInputElement; + const name = document.getElementById("moved-dataset-name") as HTMLInputElement; + const moveButton = () => + [...document.body.querySelectorAll("button")].find((b) => + b.textContent?.startsWith("Move "), + ) as HTMLButtonElement; + const type = async (input: HTMLInputElement, value: string) => { + const setter = Object.getOwnPropertyDescriptor(HTMLInputElement.prototype, "value")!.set!; + await act(async () => { + setter.call(input, value); + input.dispatchEvent(new Event("input", { bubbles: true })); + }); + }; + + expect(episodes.value).toBe("1"); + expect(moveButton().disabled).toBe(true); + + await type(name, "rejected takes"); + expect(name.value).toBe("rejected_takes"); + expect(moveButton().disabled).toBe(false); + + await type(episodes, "0-1"); + expect(document.body.textContent).toContain("At least one episode must stay in the dataset."); + expect(moveButton().disabled).toBe(true); + + await type(episodes, "1"); + await act(async () => moveButton().click()); + expect(moved).toEqual([[[1], "rejected_takes"]]); + } finally { + await act(async () => root.unmount()); + host.remove(); + } +}); diff --git a/frontend/src/components/dataset/MoveEpisodesDialog.tsx b/frontend/src/components/dataset/MoveEpisodesDialog.tsx new file mode 100644 index 000000000..273ae5c04 --- /dev/null +++ b/frontend/src/components/dataset/MoveEpisodesDialog.tsx @@ -0,0 +1,144 @@ +import { useEffect, useState } from "react"; +import { Loader2 } from "lucide-react"; + +import { Button } from "@/components/ui/button"; +import { + Dialog, + DialogContent, + DialogDescription, + DialogFooter, + DialogHeader, + DialogTitle, +} from "@/components/ui/dialog"; +import { Input } from "@/components/ui/input"; +import { Label } from "@/components/ui/label"; +import { parseEpisodeList } from "@/lib/datasetApi"; + +interface MoveEpisodesDialogProps { + repoId: string; + totalEpisodes: number; + currentEpisode: number | null; + open: boolean; + onOpenChange: (open: boolean) => void; + onMove: (episodeIndices: number[], outputName: string) => Promise; +} + +export default function MoveEpisodesDialog({ + repoId, + totalEpisodes, + currentEpisode, + open, + onOpenChange, + onMove, +}: MoveEpisodesDialogProps) { + const [episodesText, setEpisodesText] = useState(""); + const [outputName, setOutputName] = useState(""); + const [moving, setMoving] = useState(false); + const [error, setError] = useState(null); + + useEffect(() => { + if (open) { + setEpisodesText(currentEpisode === null ? "" : String(currentEpisode)); + setOutputName(""); + setMoving(false); + setError(null); + } + }, [open, currentEpisode]); + + const episodes = parseEpisodeList(episodesText, totalEpisodes); + const movesEverything = episodes !== null && episodes.length >= totalEpisodes; + const canMove = + episodes !== null && !movesEverything && /^[A-Za-z0-9._-]+$/.test(outputName) && !moving; + + const handleOpenChange = (nextOpen: boolean) => { + if (moving) return; + onOpenChange(nextOpen); + }; + + const handleMove = async () => { + if (!canMove || episodes === null) return; + setMoving(true); + try { + await onMove(episodes, outputName); + onOpenChange(false); + } catch (caughtError) { + setError(caughtError instanceof Error ? caughtError.message : "Could not move the episodes."); + } finally { + setMoving(false); + } + }; + + return ( + + + + Move episodes to a new dataset + + The episodes leave {repoId} and are saved locally as a new dataset. The + episodes left in {repoId} are renumbered. + + + +
+
+ + { + setError(null); + setEpisodesText(event.target.value); + }} + placeholder="3, 7, 10-12" + disabled={moving} + className="border-gray-700 bg-gray-800 text-white" + /> + {episodesText.trim() && episodes === null && ( +

+ Use episode numbers from 0 to {totalEpisodes - 1}, e.g. 3, 7, 10-12. +

+ )} + {movesEverything && ( +

At least one episode must stay in the dataset.

+ )} +
+ +
+ + { + setError(null); + setOutputName(event.target.value.replace(/[^A-Za-z0-9._-]/g, "_")); + }} + placeholder="rejected_takes" + disabled={moving} + className="border-gray-700 bg-gray-800 text-white" + /> +

+ Saved locally as local/{outputName || "name"}. +

+
+ {error &&

{error}

} +
+ + + + + + + + ); +} diff --git a/frontend/src/hooks/useEpisodes.ts b/frontend/src/hooks/useEpisodes.ts index 4158ab4d6..791c6dddf 100644 --- a/frontend/src/hooks/useEpisodes.ts +++ b/frontend/src/hooks/useEpisodes.ts @@ -1,4 +1,4 @@ -import { useEffect, useState } from "react"; +import { useCallback, useEffect, useState } from "react"; import { useApi } from "@/contexts/ApiContext"; import { EpisodeDetail, @@ -13,6 +13,7 @@ export const useEpisodes = (repoId: string | null) => { const [data, setData] = useState(null); const [loading, setLoading] = useState(false); const [error, setError] = useState(null); + const [version, setVersion] = useState(0); useEffect(() => { setData(null); @@ -44,11 +45,13 @@ export const useEpisodes = (repoId: string | null) => { if (!controller.signal.aborted) setLoading(false); }); return () => controller.abort(); - }, [baseUrl, fetchWithHeaders, repoId]); + }, [baseUrl, fetchWithHeaders, repoId, version]); + + const refresh = useCallback(() => setVersion((v) => v + 1), []); // Effects run after render: do not let the new dataset auto-select an // episode from the old list before this effect clears it. - return { data: data?.repo_id === repoId ? data : null, loading, error }; + return { data: data?.repo_id === repoId ? data : null, loading, error, refresh }; }; /** One episode's detail. Either arg null → idle. */ diff --git a/frontend/src/lib/datasetApi.ts b/frontend/src/lib/datasetApi.ts index 4ac27a8af..1acfb5fa5 100644 --- a/frontend/src/lib/datasetApi.ts +++ b/frontend/src/lib/datasetApi.ts @@ -189,3 +189,22 @@ export function formatDuration(seconds: number | null | undefined): string { const s = Math.round(seconds % 60); return `${m}m ${String(s).padStart(2, "0")}s`; } + +/** + * Parse "3, 7, 10-12" into sorted episode indices. Returns null when the text + * is not a list of indices and ranges below `total`. + */ +export function parseEpisodeList(text: string, total: number): number[] | null { + const indices = new Set(); + for (const part of text.split(",")) { + const item = part.trim(); + if (!item) continue; + const match = /^(\d+)(?:\s*-\s*(\d+))?$/.exec(item); + if (!match) return null; + const start = Number(match[1]); + const end = match[2] === undefined ? start : Number(match[2]); + if (end < start || end >= total) return null; + for (let i = start; i <= end; i++) indices.add(i); + } + return indices.size > 0 ? [...indices].sort((a, b) => a - b) : null; +} diff --git a/frontend/src/lib/replayApi.ts b/frontend/src/lib/replayApi.ts index 21c591ec2..4c2190939 100644 --- a/frontend/src/lib/replayApi.ts +++ b/frontend/src/lib/replayApi.ts @@ -16,6 +16,13 @@ export interface MergeDatasetsResult { total_frames: number; } +export interface MoveEpisodesResult { + success: true; + repo_id: string; + moved_episodes: number; + remaining_episodes: number; +} + export async function listDatasets( baseUrl: string, fetcher: Fetcher, @@ -55,3 +62,17 @@ export async function mergeLocalDatasets( action: "Merge datasets", }); } + +export async function moveEpisodes( + baseUrl: string, + fetcher: Fetcher, + repoId: string, + episodeIndices: number[], + outputName: string, +): Promise { + return apiRequest(baseUrl, fetcher, "/datasets/move-episodes", { + method: "POST", + body: { repo_id: repoId, episode_indices: episodeIndices, output_name: outputName }, + action: "Move episodes", + }); +} diff --git a/frontend/src/pages/EditDataset.tsx b/frontend/src/pages/EditDataset.tsx index 53ec628a1..ba09b2686 100644 --- a/frontend/src/pages/EditDataset.tsx +++ b/frontend/src/pages/EditDataset.tsx @@ -9,6 +9,7 @@ import { ExternalLink, Loader2, Merge, + Scissors, Trash2, Upload, } from "lucide-react"; @@ -17,6 +18,7 @@ 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 MoveEpisodesDialog from "@/components/dataset/MoveEpisodesDialog"; import DatasetPicker from "@/components/landing/DatasetPicker"; import RecordingModal from "@/components/landing/RecordingModal"; import { useApi } from "@/contexts/ApiContext"; @@ -27,7 +29,7 @@ import { useRobots } from "@/hooks/useRobots"; import { useToast } from "@/hooks/use-toast"; import { formatDuration } from "@/lib/datasetApi"; import { openHubViewer } from "@/lib/hubViewer"; -import { DatasetItem, mergeLocalDatasets } from "@/lib/replayApi"; +import { DatasetItem, mergeLocalDatasets, moveEpisodes } from "@/lib/replayApi"; // The robot_type LeRobot stores for the SO-101 follower. const FOLLOWER_ROBOT_TYPE = "so_follower"; @@ -55,6 +57,7 @@ const EditDataset = () => { const { toast } = useToast(); const [mergeOpen, setMergeOpen] = useState(false); const [deleteOpen, setDeleteOpen] = useState(false); + const [moveOpen, setMoveOpen] = 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"), @@ -67,7 +70,7 @@ const EditDataset = () => { const selectedSource = datasets.find((d) => d.repo_id === repoId)?.source; const selectedIsLocal = selectedSource === "local" || selectedSource === "both"; - const { data, loading, error } = useEpisodes(repoId); + const { data, loading, error, refresh: refreshEpisodes } = useEpisodes(repoId); const { detail, loading: detailLoading, @@ -110,6 +113,18 @@ const EditDataset = () => { setDataset(result.repo_id); toast({ title: "Datasets merged", description: `${result.num_episodes} episodes are ready to browse.` }); }; + const handleMove = async (episodeIndices: number[], outputName: string) => { + if (!repoId) return; + const result = await moveEpisodes(baseUrl, fetchWithHeaders, repoId, episodeIndices, outputName); + refresh(); + // The episodes left behind are renumbered, so reopen the dataset from its first one. + setSearchParams(withApi({ dataset: repoId })); + refreshEpisodes(); + toast({ + title: "Episodes moved", + description: `${result.moved_episodes} moved to ${result.repo_id}, ${result.remaining_episodes} left in ${repoId}.`, + }); + }; const setEpisode = (next: number) => { if (!repoId) return; setSearchParams(withApi({ dataset: repoId, episode: String(next) })); @@ -176,6 +191,17 @@ const EditDataset = () => { Merge )} + {repoId && selectedIsLocal && data?.success && data.total_episodes > 1 && ( + + )} {repoId && selectedSource === "both" && ( + + {destination !== null && ( +

+ {"newName" in destination ? ( + <> + Saved locally as local/{destination.newName}. + + ) : ( + <> + Merged into {destination.existingRepoId} after its own episodes. It must + have the same fps, features and robot type. + + )} +

+ )} {error &&

{error}

} diff --git a/frontend/src/lib/replayApi.ts b/frontend/src/lib/replayApi.ts index 4c2190939..4e2c8b3ff 100644 --- a/frontend/src/lib/replayApi.ts +++ b/frontend/src/lib/replayApi.ts @@ -63,16 +63,24 @@ export async function mergeLocalDatasets( }); } +export type MoveDestination = { newName: string } | { existingRepoId: string }; + export async function moveEpisodes( baseUrl: string, fetcher: Fetcher, repoId: string, episodeIndices: number[], - outputName: string, + destination: MoveDestination, ): Promise { return apiRequest(baseUrl, fetcher, "/datasets/move-episodes", { method: "POST", - body: { repo_id: repoId, episode_indices: episodeIndices, output_name: outputName }, + body: { + repo_id: repoId, + episode_indices: episodeIndices, + ...("newName" in destination + ? { output_name: destination.newName } + : { target_repo_id: destination.existingRepoId }), + }, action: "Move episodes", }); } diff --git a/frontend/src/pages/EditDataset.tsx b/frontend/src/pages/EditDataset.tsx index ba09b2686..1d0f5c107 100644 --- a/frontend/src/pages/EditDataset.tsx +++ b/frontend/src/pages/EditDataset.tsx @@ -9,7 +9,7 @@ import { ExternalLink, Loader2, Merge, - Scissors, + SquareArrowOutUpRight, Trash2, Upload, } from "lucide-react"; @@ -29,7 +29,7 @@ import { useRobots } from "@/hooks/useRobots"; import { useToast } from "@/hooks/use-toast"; import { formatDuration } from "@/lib/datasetApi"; import { openHubViewer } from "@/lib/hubViewer"; -import { DatasetItem, mergeLocalDatasets, moveEpisodes } from "@/lib/replayApi"; +import { DatasetItem, MoveDestination, mergeLocalDatasets, moveEpisodes } from "@/lib/replayApi"; // The robot_type LeRobot stores for the SO-101 follower. const FOLLOWER_ROBOT_TYPE = "so_follower"; @@ -113,9 +113,9 @@ const EditDataset = () => { setDataset(result.repo_id); toast({ title: "Datasets merged", description: `${result.num_episodes} episodes are ready to browse.` }); }; - const handleMove = async (episodeIndices: number[], outputName: string) => { + const handleMove = async (episodeIndices: number[], destination: MoveDestination) => { if (!repoId) return; - const result = await moveEpisodes(baseUrl, fetchWithHeaders, repoId, episodeIndices, outputName); + const result = await moveEpisodes(baseUrl, fetchWithHeaders, repoId, episodeIndices, destination); refresh(); // The episodes left behind are renumbered, so reopen the dataset from its first one. setSearchParams(withApi({ dataset: repoId })); @@ -198,7 +198,7 @@ const EditDataset = () => { onClick={() => setMoveOpen(true)} className="h-9 gap-1.5 border-gray-800 bg-gray-950 text-xs text-gray-300 hover:bg-gray-900" > - + Move episodes )} @@ -345,6 +345,7 @@ const EditDataset = () => { {repoId && selectedIsLocal && data?.success && ( } -def handle_move_episodes(repo_id: str, episode_indices: list[int], output_name: str) -> dict[str, Any]: - """Move episodes out of a local dataset into a new local dataset. +def _swap_dirs(replacements: list[tuple[Path, Path]], backup_dir: Path) -> None: + """Replace each live dir with its staged copy; a failure restores the ones already swapped.""" + done: list[tuple[Path, Path | None]] = [] + try: + for i, (live, staged) in enumerate(replacements): + backup = backup_dir / f"old{i}" if live.exists() else None + if backup is not None: + live.rename(backup) + done.append((live, backup)) + staged.rename(live) + except OSError: + for live, backup in reversed(done): + shutil.rmtree(live, ignore_errors=True) + if backup is not None: + backup.rename(live) + raise + - Both datasets are written to a staging dir first, so a failure leaves the - source as it was. The source is then swapped for the copy without the moved - episodes, which renumbers the episodes left in it. +def handle_move_episodes( + repo_id: str, + episode_indices: list[int], + output_name: str | None = None, + target_repo_id: str | None = None, +) -> dict[str, Any]: + """Move episodes out of a local dataset into a new or an existing local dataset. + + Every rewritten dataset is built in a staging dir first, so a failure leaves + the source and the target as they were. The source is then swapped for the + copy without the moved episodes, which renumbers the episodes left in it. An + existing target is merged with the moved episodes, which keeps its own + episodes first and refuses a target whose fps, features or robot type differ. """ episode_indices = sorted(set(episode_indices)) if not episode_indices: raise ValueError("Select at least one episode to move") + if (output_name is None) == (target_repo_id is None): + raise ValueError("Choose either a new dataset name or an existing target dataset") - output_name = output_name.strip() - if not _DATASET_NAME_RE.fullmatch(output_name): - raise ValueError("Output name may contain only letters, numbers, '.', '_' and '-'") - - output_repo_id = f"{LOCAL_DATASET_NAMESPACE}/{output_name}" cache_root = _lerobot_cache_root() - output_dir = (cache_root / output_repo_id).resolve() - if cache_root not in output_dir.parents: - raise ValueError("Invalid output dataset path") - if output_dir.exists(): - raise ValueError(f"A local dataset already exists at {output_repo_id}") + target_dir: Path | None = None + if target_repo_id is not None: + if target_repo_id == repo_id: + raise ValueError("The target must be a different dataset than the source") + try: + episode_media.ensure_episode_index(target_repo_id) + except episode_media.UnreadableDatasetError as exc: + raise ValueError(str(exc)) from exc + target_dir = episode_media.resolve_dataset_dir(target_repo_id) + destination_repo_id = target_repo_id + destination_dir = target_dir + staging_prefix = f".{target_dir.name}.move-" + else: + output_name = output_name.strip() + if not _DATASET_NAME_RE.fullmatch(output_name): + raise ValueError("Output name may contain only letters, numbers, '.', '_' and '-'") + destination_repo_id = f"{LOCAL_DATASET_NAMESPACE}/{output_name}" + destination_dir = (cache_root / destination_repo_id).resolve() + if cache_root not in destination_dir.parents: + raise ValueError("Invalid output dataset path") + if destination_dir.exists(): + raise ValueError(f"A local dataset already exists at {destination_repo_id}") + staging_prefix = f".{output_name}.move-" try: episode_media.ensure_episode_index(repo_id) @@ -272,29 +312,28 @@ def handle_move_episodes(repo_id: str, episode_indices: list[int], output_name: if not kept: raise ValueError("Cannot move every episode; at least one must stay in the dataset") - output_dir.parent.mkdir(parents=True, exist_ok=True) - staging = Path(tempfile.mkdtemp(prefix=f".{output_name}.move-", dir=output_dir.parent)) + destination_dir.parent.mkdir(parents=True, exist_ok=True) + staging = Path(tempfile.mkdtemp(prefix=staging_prefix, dir=destination_dir.parent)) try: - moved = delete_episodes(dataset, kept, output_dir=staging / "moved", repo_id=output_repo_id) + moved = delete_episodes(dataset, kept, output_dir=staging / "moved", repo_id=destination_repo_id) delete_episodes(dataset, episode_indices, output_dir=staging / "kept", repo_id=repo_id) - (staging / "moved").rename(output_dir) - source_dir.rename(staging / "old") - try: - (staging / "kept").rename(source_dir) - except OSError: - (staging / "old").rename(source_dir) - shutil.rmtree(output_dir, ignore_errors=True) - raise - except Exception: - shutil.rmtree(staging, ignore_errors=True) - raise - else: + if target_dir is not None: + target = LeRobotDataset(repo_id=target_repo_id, root=target_dir) + merge_datasets([target, moved], output_repo_id=target_repo_id, output_dir=staging / "merged") + destination_staged = staging / "merged" + else: + destination_staged = staging / "moved" + _swap_dirs( + [(source_dir, staging / "kept"), (destination_dir, destination_staged)], + backup_dir=staging, + ) + finally: shutil.rmtree(staging, ignore_errors=True) return { "success": True, - "repo_id": output_repo_id, - "moved_episodes": moved.num_episodes, + "repo_id": destination_repo_id, + "moved_episodes": len(episode_indices), "remaining_episodes": len(kept), } diff --git a/lelab/server.py b/lelab/server.py index 77c6950aa..1c636d5b2 100644 --- a/lelab/server.py +++ b/lelab/server.py @@ -157,7 +157,8 @@ class MergeDatasetsRequest(BaseModel): class MoveEpisodesRequest(BaseModel): repo_id: str episode_indices: list[Annotated[int, Field(ge=0)]] = Field(min_length=1) - output_name: str + output_name: str | None = None + target_repo_id: str | None = None # Cache for HF Jobs hardware flavors (5-minute TTL) @@ -464,12 +465,12 @@ def merge_local_datasets(request: MergeDatasetsRequest): @app.post("/datasets/move-episodes") def move_dataset_episodes(request: MoveEpisodesRequest): - """Move episodes out of a local dataset into a new local dataset.""" + """Move episodes out of a local dataset into a new or an existing local dataset.""" if _record.recording_active: raise HTTPException(status_code=409, detail="Stop the recording before editing a dataset") try: return dataset_browser.handle_move_episodes( - request.repo_id, request.episode_indices, request.output_name + request.repo_id, request.episode_indices, request.output_name, request.target_repo_id ) except (DatasetNotFoundError, ValueError) as exc: raise HTTPException(status_code=400, detail=str(exc)) from exc diff --git a/tests/test_datasets.py b/tests/test_datasets.py index 64ce75071..b6e89f188 100644 --- a/tests/test_datasets.py +++ b/tests/test_datasets.py @@ -262,14 +262,14 @@ def test_merge_route_returns_result_and_maps_validation_error( assert response.json()["detail"] == "bad input" -def _record_dataset(repo_id: str, episodes: int) -> Path: +def _record_dataset(repo_id: str, episodes: int, fps: int = 10) -> Path: """Write a finished LeRobot dataset whose episode i has every value set to i.""" import numpy as np from lerobot.datasets import LeRobotDataset features = {"action": {"dtype": "float32", "shape": (1,), "names": ["a"]}} - dataset = LeRobotDataset.create(repo_id, fps=10, features=features, use_videos=False) + dataset = LeRobotDataset.create(repo_id, fps=fps, features=features, use_videos=False) for ep in range(episodes): for _ in range(3): dataset.add_frame({"action": np.full(1, ep, dtype=np.float32), "task": f"task {ep}"}) @@ -324,6 +324,60 @@ def test_move_episodes_rejects_invalid_request( assert _episode_tasks("alice/demo") == ["task 0", "task 1"] +def test_move_episodes_merges_into_an_existing_dataset(tmp_lerobot_home: Path) -> None: + from lelab.datasets import handle_move_episodes + + _record_dataset("alice/demo", episodes=4) + _record_dataset("bob/target", episodes=2) + + result = handle_move_episodes("alice/demo", [1, 3], target_repo_id="bob/target") + + assert result == { + "success": True, + "repo_id": "bob/target", + "moved_episodes": 2, + "remaining_episodes": 2, + } + assert _episode_tasks("bob/target") == ["task 0", "task 1", "task 1", "task 3"] + assert _episode_tasks("alice/demo") == ["task 0", "task 2"] + assert not list((tmp_lerobot_home / "bob").glob(".target.move-*")) + + +def test_move_episodes_refuses_an_incompatible_target(tmp_lerobot_home: Path) -> None: + from lelab.datasets import handle_move_episodes + + _record_dataset("alice/demo", episodes=3) + _record_dataset("bob/target", episodes=2, fps=20) + + with pytest.raises(ValueError, match="Same fps is expected"): + handle_move_episodes("alice/demo", [1], target_repo_id="bob/target") + + assert _episode_tasks("alice/demo") == ["task 0", "task 1", "task 2"] + assert _episode_tasks("bob/target") == ["task 0", "task 1"] + assert not list((tmp_lerobot_home / "bob").glob(".target.move-*")) + + +@pytest.mark.parametrize( + ("kwargs", "message"), + [ + ({}, "either a new dataset name or an existing target"), + ( + {"output_name": "x", "target_repo_id": "bob/target"}, + "either a new dataset name or an existing target", + ), + ({"target_repo_id": "alice/demo"}, "different dataset than the source"), + ], +) +def test_move_episodes_rejects_an_ambiguous_destination( + tmp_lerobot_home: Path, kwargs: dict[str, str], message: str +) -> None: + from lelab.datasets import handle_move_episodes + + _record_dataset("alice/demo", episodes=2) + with pytest.raises(ValueError, match=message): + handle_move_episodes("alice/demo", [0], **kwargs) + + def test_move_episodes_rejects_existing_output(tmp_lerobot_home: Path) -> None: from lelab.datasets import handle_move_episodes @@ -387,6 +441,23 @@ def test_move_episodes_route_maps_validation_error( assert response.json()["detail"] == "bad input" +def test_move_episodes_route_forwards_an_existing_target( + client: TestClient, monkeypatch: pytest.MonkeyPatch +) -> None: + from lelab import datasets as datasets_mod + + calls = [] + monkeypatch.setattr( + datasets_mod, "handle_move_episodes", lambda *args: calls.append(args) or {"success": True} + ) + response = client.post( + "/datasets/move-episodes", + json={"repo_id": "alice/demo", "episode_indices": [0], "target_repo_id": "bob/target"}, + ) + assert response.status_code == 200 + assert calls == [("alice/demo", [0], None, "bob/target")] + + @pytest.mark.parametrize("episode_indices", [[], [-1]]) def test_move_episodes_route_rejects_bad_indices_before_editing( client: TestClient, monkeypatch: pytest.MonkeyPatch, episode_indices: list[int]