diff --git a/frontend/src/components/dataset/EpisodeList.tsx b/frontend/src/components/dataset/EpisodeList.tsx index 120a24ac..36186b60 100644 --- a/frontend/src/components/dataset/EpisodeList.tsx +++ b/frontend/src/components/dataset/EpisodeList.tsx @@ -1,13 +1,16 @@ import React, { useEffect, useRef } from "react"; +import { Pencil } from "lucide-react"; import { EpisodeSummary, formatDuration } from "@/lib/datasetApi"; interface EpisodeListProps { episodes: EpisodeSummary[]; selected: number | null; onSelect: (episodeIndex: number) => void; + /** Shows an edit button on the selected episode's task when set. */ + onEditTask?: (episodeIndex: number) => void; } -const EpisodeList: React.FC = ({ episodes, selected, onSelect }) => { +const EpisodeList: React.FC = ({ episodes, selected, onSelect, onEditTask }) => { const activeRef = useRef(null); // Follow the selection. Paging with Previous/Next, or opening a link straight @@ -39,29 +42,43 @@ const EpisodeList: React.FC = ({ episodes, selected, onSelect {episodes.map((ep) => { const active = ep.episode_index === selected; return ( - + + {ep.episode_index} + + + {ep.tasks.length > 0 ? ep.tasks.join(", ") : "—"} + + + {formatDuration(ep.duration_s)} + + + {active && onEditTask && ( + + )} + ); })} diff --git a/frontend/src/components/dataset/EpisodeTaskDialog.test.tsx b/frontend/src/components/dataset/EpisodeTaskDialog.test.tsx new file mode 100644 index 00000000..9b1e34f4 --- /dev/null +++ b/frontend/src/components/dataset/EpisodeTaskDialog.test.tsx @@ -0,0 +1,61 @@ +// @vitest-environment jsdom +import React, { act } from "react"; +import { createRoot } from "react-dom/client"; +import { expect, test } from "vitest"; + +import EpisodeTaskDialog from "./EpisodeTaskDialog"; + +test("saves the task for the browsed episode or for 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 saved: [string, number | null][] = []; + + try { + await act(async () => { + root.render( + {}} + onSave={(task, episodeIndex) => { + saved.push([task, episodeIndex]); + }} + />, + ); + }); + + const input = document.getElementById("episode-task") as HTMLInputElement; + const saveButton = () => + [...document.body.querySelectorAll("button")].find( + (b) => b.textContent === "Save task", + ) as HTMLButtonElement; + const type = async (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(input.value).toBe("pick the cube"); + + await type(" "); + expect(saveButton().disabled).toBe(true); + + await type(" pick the bowl "); + await act(async () => saveButton().click()); + expect(saved).toEqual([["pick the bowl", 3]]); + + await act(async () => (document.getElementById("episode-task-all") as HTMLButtonElement).click()); + await act(async () => saveButton().click()); + expect(saved.at(-1)).toEqual(["pick the bowl", null]); + } finally { + await act(async () => root.unmount()); + host.remove(); + } +}); diff --git a/frontend/src/components/dataset/EpisodeTaskDialog.tsx b/frontend/src/components/dataset/EpisodeTaskDialog.tsx new file mode 100644 index 00000000..8ca08ecb --- /dev/null +++ b/frontend/src/components/dataset/EpisodeTaskDialog.tsx @@ -0,0 +1,106 @@ +import { useEffect, useState } from "react"; + +import { Button } from "@/components/ui/button"; +import { Checkbox } from "@/components/ui/checkbox"; +import { + Dialog, + DialogContent, + DialogDescription, + DialogFooter, + DialogHeader, + DialogTitle, +} from "@/components/ui/dialog"; +import { Input } from "@/components/ui/input"; +import { Label } from "@/components/ui/label"; + +interface EpisodeTaskDialogProps { + repoId: string; + episodeIndex: number; + currentTask: string; + totalEpisodes: number; + open: boolean; + onOpenChange: (open: boolean) => void; + /** `episodeIndex` is null when the task applies to every episode. */ + onSave: (task: string, episodeIndex: number | null) => void; +} + +export default function EpisodeTaskDialog({ + repoId, + episodeIndex, + currentTask, + totalEpisodes, + open, + onOpenChange, + onSave, +}: EpisodeTaskDialogProps) { + const [task, setTask] = useState(""); + const [allEpisodes, setAllEpisodes] = useState(false); + + useEffect(() => { + if (open) { + setTask(currentTask); + setAllEpisodes(false); + } + }, [open, currentTask]); + + const canSave = task.trim() !== ""; + + const handleSave = () => { + if (!canSave) return; + onSave(task.trim(), allEpisodes ? null : episodeIndex); + onOpenChange(false); + }; + + return ( + + + + Edit task + + The task is the instruction a policy is trained on, saved with each episode of{" "} + {repoId}. + + + +
+
+ + setTask(event.target.value)} + onKeyDown={(event) => { + if (event.key === "Enter") handleSave(); + }} + className="border-gray-700 bg-gray-800 text-white" + /> +
+ {totalEpisodes > 1 && ( +
+ setAllEpisodes(checked === true)} + className="border-gray-500" + /> + +
+ )} +
+ + + + + + + + ); +} diff --git a/frontend/src/hooks/useEpisodes.ts b/frontend/src/hooks/useEpisodes.ts index 791c6ddd..97875a1b 100644 --- a/frontend/src/hooks/useEpisodes.ts +++ b/frontend/src/hooks/useEpisodes.ts @@ -49,9 +49,23 @@ export const useEpisodes = (repoId: string | null) => { const refresh = useCallback(() => setVersion((v) => v + 1), []); + /** Show a task edit right away, before the backend has written it. */ + const applyTask = useCallback( + (task: string, episodeIndex: number | null) => + setData((d) => + d && { + ...d, + episodes: d.episodes.map((ep) => + episodeIndex === null || ep.episode_index === episodeIndex ? { ...ep, tasks: [task] } : ep, + ), + }, + ), + [], + ); + // 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, refresh }; + return { data: data?.repo_id === repoId ? data : null, loading, error, refresh, applyTask }; }; /** One episode's detail. Either arg null → idle. */ diff --git a/frontend/src/lib/replayApi.ts b/frontend/src/lib/replayApi.ts index 4e2c8b3f..8d793ef4 100644 --- a/frontend/src/lib/replayApi.ts +++ b/frontend/src/lib/replayApi.ts @@ -84,3 +84,36 @@ export async function moveEpisodes( action: "Move episodes", }); } + +export type EpisodeTaskStatus = + | { pending: true } + | { pending: false; local_error?: string | null; hub_error?: string | null; hub_updated?: boolean }; + +/** Queue a task edit for one episode, or for every episode when `episodeIndex` is null. */ +export async function setEpisodeTask( + baseUrl: string, + fetcher: Fetcher, + repoId: string, + task: string, + episodeIndex: number | null, +): Promise { + await apiRequest(baseUrl, fetcher, "/datasets/episode-task", { + method: "POST", + body: { repo_id: repoId, task, episode_index: episodeIndex }, + action: "Edit task", + }); +} + +/** Whether the queued task edits of a dataset are still being written, and how they ended. */ +export async function getEpisodeTaskStatus( + baseUrl: string, + fetcher: Fetcher, + repoId: string, +): Promise { + return apiRequest( + baseUrl, + fetcher, + `/datasets/episode-task-status?repo_id=${encodeURIComponent(repoId)}`, + { action: "Task edit status" }, + ); +} diff --git a/frontend/src/pages/EditDataset.tsx b/frontend/src/pages/EditDataset.tsx index 1d0f5c10..7c097f32 100644 --- a/frontend/src/pages/EditDataset.tsx +++ b/frontend/src/pages/EditDataset.tsx @@ -16,6 +16,7 @@ import { import { Button } from "@/components/ui/button"; import DeleteDatasetDialog from "@/components/dataset/DeleteDatasetDialog"; import EpisodeList from "@/components/dataset/EpisodeList"; +import EpisodeTaskDialog from "@/components/dataset/EpisodeTaskDialog"; import EpisodeViewer from "@/components/dataset/EpisodeViewer"; import MergeDatasetsDialog from "@/components/dataset/MergeDatasetsDialog"; import MoveEpisodesDialog from "@/components/dataset/MoveEpisodesDialog"; @@ -29,7 +30,7 @@ import { useRobots } from "@/hooks/useRobots"; import { useToast } from "@/hooks/use-toast"; import { formatDuration } from "@/lib/datasetApi"; import { openHubViewer } from "@/lib/hubViewer"; -import { DatasetItem, MoveDestination, mergeLocalDatasets, moveEpisodes } from "@/lib/replayApi"; +import { DatasetItem, MoveDestination, mergeLocalDatasets, getEpisodeTaskStatus, moveEpisodes, setEpisodeTask } from "@/lib/replayApi"; // The robot_type LeRobot stores for the SO-101 follower. const FOLLOWER_ROBOT_TYPE = "so_follower"; @@ -58,6 +59,7 @@ const EditDataset = () => { const [mergeOpen, setMergeOpen] = useState(false); const [deleteOpen, setDeleteOpen] = useState(false); const [moveOpen, setMoveOpen] = useState(false); + const [taskOpen, setTaskOpen] = 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"), @@ -70,7 +72,7 @@ const EditDataset = () => { const selectedSource = datasets.find((d) => d.repo_id === repoId)?.source; const selectedIsLocal = selectedSource === "local" || selectedSource === "both"; - const { data, loading, error, refresh: refreshEpisodes } = useEpisodes(repoId); + const { data, loading, error, refresh: refreshEpisodes, applyTask } = useEpisodes(repoId); const { detail, loading: detailLoading, @@ -125,6 +127,57 @@ const EditDataset = () => { description: `${result.moved_episodes} moved to ${result.repo_id}, ${result.remaining_episodes} left in ${repoId}.`, }); }; + // The dataset whose queued task edits are still being written; polled until the backend reports it idle. + const [savingRepo, setSavingRepo] = useState(null); + const handleSaveTask = (task: string, targetEpisode: number | null) => { + if (!repoId) return; + const edited = repoId; + applyTask(task, targetEpisode); + setEpisodeTask(baseUrl, fetchWithHeaders, edited, task, targetEpisode).then( + () => setSavingRepo(edited), + (e) => { + refreshEpisodes(); + toast({ + title: "Task not saved", + description: e instanceof Error ? e.message : "Could not save the task.", + variant: "destructive", + }); + }, + ); + }; + useEffect(() => { + if (!savingRepo) return; + const timer = window.setInterval(async () => { + try { + const status = await getEpisodeTaskStatus(baseUrl, fetchWithHeaders, savingRepo); + if (status.pending) return; + setSavingRepo(null); + if (status.local_error) { + refreshEpisodes(); + toast({ title: "Task not saved", description: status.local_error, variant: "destructive" }); + } else if (status.hub_error) { + toast({ + title: "Task saved locally only", + description: `The Hub copy was not updated: ${status.hub_error}.`, + variant: "destructive", + }); + } else { + toast({ + title: "Tasks saved", + description: status.hub_updated ? "The Hub copy is updated too." : undefined, + }); + } + } catch (e) { + setSavingRepo(null); + toast({ + title: "Could not check the task edits", + description: e instanceof Error ? e.message : undefined, + variant: "destructive", + }); + } + }, 700); + return () => window.clearInterval(timer); + }, [savingRepo, baseUrl, fetchWithHeaders, refreshEpisodes, toast]); const setEpisode = (next: number) => { if (!repoId) return; setSearchParams(withApi({ dataset: repoId, episode: String(next) })); @@ -303,6 +356,7 @@ const EditDataset = () => { {detail && } + {savingRepo === repoId && Saving tasks…}
@@ -310,6 +364,7 @@ const EditDataset = () => { episodes={data.episodes} selected={episodeIndex} onSelect={setEpisode} + onEditTask={selectedIsLocal ? () => setTaskOpen(true) : undefined} />
@@ -355,6 +410,18 @@ const EditDataset = () => { /> )} + {repoId && selectedIsLocal && data?.success && episodeIndex !== null && ( + ep.episode_index === episodeIndex)?.tasks[0] ?? ""} + totalEpisodes={data.total_episodes} + open={taskOpen} + onOpenChange={setTaskOpen} + onSave={handleSaveTask} + /> + )} + {repoId && selectedIsLocal && ( dict[str, Any]: episode_indices = sorted(set(episode_indices)) if not episode_indices: raise ValueError("Select at least one episode to move") @@ -338,6 +351,204 @@ def handle_move_episodes( } +_HUB_MIRRORED_GLOBS = ["meta/**", "data/**", "videos/**"] + + +def _hub_failure_reason(exc: Exception, repo_id: str) -> str: + status = getattr(getattr(exc, "response", None), "status_code", None) + if status == 401: + return "you are not logged in to the Hub" + if status == 403: + return f"you have no write access to {repo_id}" + return str(exc) or type(exc).__name__ + + +def _push_edit_to_hub(repo_id: str, root: Path) -> tuple[bool, str | None]: + """Mirror the local copy of the dataset folders to its Hub repo when it has one. + + `meta/`, `data/` and `videos/` are uploaded and the files the local copy no + longer has are deleted, in one commit, so a Hub copy that had diverged (episodes + moved out, episodes recorded without a push) ends up consistent. Unchanged files + are dropped from the commit by huggingface_hub, so videos are hashed, not re-sent. + The card and `.gitattributes` are outside those globs and stay as they are. + + Returns (pushed, error). + """ + if repo_id.split("/")[0] == LOCAL_DATASET_NAMESPACE: + return False, None + api = shared_hf_api() + try: + if not api.repo_exists(repo_id, repo_type="dataset"): + return False, None + api.upload_folder( + repo_id=repo_id, + repo_type="dataset", + folder_path=root, + allow_patterns=_HUB_MIRRORED_GLOBS, + delete_patterns=_HUB_MIRRORED_GLOBS, + commit_message="Edit episode tasks", + ) + except Exception as exc: + logger.warning("Hub update of %s failed: %s", repo_id, exc) + return False, _hub_failure_reason(exc, repo_id) + return True, None + + +def _rewrite_tasks( + repo_id: str, source_dir: Path, default_task: str | None, episode_tasks: dict[int, str] +) -> None: + """Apply task edits with one `modify_tasks` call on a staged copy swapped in on success. + + `modify_tasks` rewrites the dataset in place, so a failure here would leave + the source half written. Only `meta/` and `data/` are rewritten, so everything + else (videos, images) is hard-linked into the copy where the filesystem allows it. + """ + + def copy_or_link(src: str, dst: str) -> None: + if Path(src).relative_to(source_dir).parts[0] in ("meta", "data"): + shutil.copy2(src, dst) + else: + try: + os.link(src, dst) + except OSError: # e.g. a filesystem without hard links + shutil.copy2(src, dst) + + staging = Path(tempfile.mkdtemp(prefix=f".{source_dir.name}.task-", dir=source_dir.parent)) + edited_dir = staging / "edited" + try: + shutil.copytree(source_dir, edited_dir, copy_function=copy_or_link) + dataset = LeRobotDataset(repo_id=repo_id, root=edited_dir) + modify_tasks(dataset, new_task=default_task, episode_tasks=episode_tasks or None) + source_dir.rename(staging / "old") + try: + edited_dir.rename(source_dir) + except OSError: + (staging / "old").rename(source_dir) + raise + finally: + shutil.rmtree(staging, ignore_errors=True) + + +class _TaskEdits: + """Task edits of one dataset, applied in order by a single worker thread.""" + + def __init__(self) -> None: + # An "apply to all" edit sets `default_task` and drops every per-episode edit queued before it. + self.default_task: str | None = None + self.episode_tasks: dict[int, str] = {} + self.running = False + self.dirty = False # written locally, not yet mirrored to the Hub + self.local_error: str | None = None + self.hub_error: str | None = None + self.hub_updated = False + + def queued(self) -> bool: + return self.default_task is not None or bool(self.episode_tasks) + + +_task_edits: dict[str, _TaskEdits] = {} +_task_edits_lock = threading.Lock() +_moving: set[str] = set() + + +def _task_edits_busy(repo_id: str) -> bool: + edits = _task_edits.get(repo_id) + return edits is not None and (edits.running or edits.queued()) + + +@contextmanager +def _exclusive_rewrite(*repo_ids: str): + """Keep task edits and moves from rewriting the same dataset at the same time.""" + with _task_edits_lock: + for repo_id in repo_ids: + if repo_id in _moving or _task_edits_busy(repo_id): + raise ValueError(f"{repo_id} is being edited, try again in a moment") + _moving.update(repo_ids) + try: + yield + finally: + with _task_edits_lock: + _moving.difference_update(repo_ids) + + +def _task_edits_worker(repo_id: str, source_dir: Path, edits: _TaskEdits) -> None: + while True: + with _task_edits_lock: + if not edits.queued(): + edits.running = False + return + default_task, episode_tasks = edits.default_task, edits.episode_tasks + edits.default_task, edits.episode_tasks = None, {} + try: + _rewrite_tasks(repo_id, source_dir, default_task, episode_tasks) + edits.dirty = True + except Exception as exc: + logger.exception("Task edit of %s failed", repo_id) + edits.local_error = str(exc) or type(exc).__name__ + # One Hub commit per burst: push only once nothing more is queued. + if edits.dirty and not edits.queued(): + edits.hub_updated, edits.hub_error = _push_edit_to_hub(repo_id, source_dir) + edits.dirty = False + + +def handle_set_episode_task(repo_id: str, task: str, episode_index: int | None = None) -> dict[str, Any]: + """Queue a task edit for one episode of a local dataset, or for every episode when `episode_index` is None. + + The edit is validated now and written in the background: edits queued while a + rewrite runs are merged into the next one (the latest wins per episode, and an + edit of every episode replaces what was queued before it). Once the queue is + idle the dataset is mirrored to the Hub if it lives there too. Poll + `handle_episode_task_status` for the outcome. + """ + task = task.strip() + if not task: + raise ValueError("The task cannot be empty") + + try: + episode_media.ensure_episode_index(repo_id) + except episode_media.UnreadableDatasetError as exc: + raise ValueError(str(exc)) from exc + source_dir = episode_media.resolve_dataset_dir(repo_id) + + total = episode_media.read_info(source_dir).get("total_episodes", 0) + if episode_index is not None and not 0 <= episode_index < total: + raise ValueError(f"Episode {episode_index} is not in {repo_id}") + + with _task_edits_lock: + if repo_id in _moving: + raise ValueError(f"{repo_id} is being edited, try again in a moment") + edits = _task_edits.setdefault(repo_id, _TaskEdits()) + if episode_index is None: + edits.default_task, edits.episode_tasks = task, {} + else: + edits.episode_tasks[episode_index] = task + if not edits.running: + edits.running = True + threading.Thread( + target=_task_edits_worker, args=(repo_id, source_dir, edits), daemon=True + ).start() + return {"success": True, "repo_id": repo_id, "queued": True} + + +def handle_episode_task_status(repo_id: str) -> dict[str, Any]: + """Progress of the queued task edits of a dataset; the outcome is reported once, when it is idle.""" + with _task_edits_lock: + edits = _task_edits.get(repo_id) + if edits is None: + return {"success": True, "pending": False} + if edits.running or edits.queued(): + return {"success": True, "pending": True} + outcome = { + "success": True, + "pending": False, + "local_error": edits.local_error, + "hub_error": edits.hub_error, + "hub_updated": edits.hub_updated, + } + del _task_edits[repo_id] + return outcome + + # ── episode browsing ──────────────────────────────────────────────────────── # # Everything below reads a *local* dataset only. Browsing a Hub-only dataset diff --git a/lelab/server.py b/lelab/server.py index 1c636d5b..dbb69ba6 100644 --- a/lelab/server.py +++ b/lelab/server.py @@ -154,6 +154,13 @@ class MergeDatasetsRequest(BaseModel): output_name: str +class EpisodeTaskRequest(BaseModel): + repo_id: str + task: Annotated[str, Field(min_length=1)] + # None sets the task of every episode. + episode_index: Annotated[int, Field(ge=0)] | None = None + + class MoveEpisodesRequest(BaseModel): repo_id: str episode_indices: list[Annotated[int, Field(ge=0)]] = Field(min_length=1) @@ -463,6 +470,23 @@ def merge_local_datasets(request: MergeDatasetsRequest): raise HTTPException(status_code=400, detail=str(exc)) from exc +@app.post("/datasets/episode-task") +def set_episode_task(request: EpisodeTaskRequest): + """Queue a task edit for one episode, or for every episode, of a local dataset.""" + if _record.recording_active: + raise HTTPException(status_code=409, detail="Stop the recording before editing a dataset") + try: + return dataset_browser.handle_set_episode_task(request.repo_id, request.task, request.episode_index) + except (DatasetNotFoundError, ValueError) as exc: + raise HTTPException(status_code=400, detail=str(exc)) from exc + + +@app.get("/datasets/episode-task-status") +def episode_task_status(repo_id: str): + """Whether queued task edits of a dataset are still being written, and how they ended.""" + return dataset_browser.handle_episode_task_status(repo_id) + + @app.post("/datasets/move-episodes") def move_dataset_episodes(request: MoveEpisodesRequest): """Move episodes out of a local dataset into a new or an existing local dataset.""" diff --git a/tests/test_datasets.py b/tests/test_datasets.py index b6e89f18..2cbf006e 100644 --- a/tests/test_datasets.py +++ b/tests/test_datasets.py @@ -16,7 +16,8 @@ from __future__ import annotations from pathlib import Path -from unittest.mock import patch +from types import SimpleNamespace +from unittest.mock import MagicMock, patch import pytest from fastapi.testclient import TestClient @@ -472,6 +473,264 @@ def test_move_episodes_route_rejects_bad_indices_before_editing( assert response.status_code == 422 +@pytest.fixture +def fake_hub(monkeypatch: pytest.MonkeyPatch) -> MagicMock: + """A Hub API where no repo exists; the tests that edit tasks must never reach the real Hub.""" + from lelab import datasets as datasets_mod + + api = MagicMock() + api.repo_exists.return_value = False + monkeypatch.setattr(datasets_mod, "shared_hf_api", lambda: api) + return api + + +def _edit_tasks(repo_id: str, *edits: tuple[str, int | None]) -> dict: + """Queue edits, wait for the worker to go idle and return the reported outcome.""" + import time + + from lelab.datasets import handle_episode_task_status, handle_set_episode_task + + for task, episode_index in edits: + handle_set_episode_task(repo_id, task, episode_index) + deadline = time.monotonic() + 30 + while time.monotonic() < deadline: + status = handle_episode_task_status(repo_id) + if not status["pending"]: + return status + time.sleep(0.02) + pytest.fail("task edits did not finish") + + +@pytest.fixture(autouse=True) +def _idle_task_edits(): + from lelab import datasets as datasets_mod + + yield + datasets_mod._task_edits.clear() + datasets_mod._moving.clear() + + +def test_set_episode_task_relabels_one_episode(tmp_lerobot_home: Path, fake_hub: MagicMock) -> None: + root = _record_dataset("alice/demo", episodes=3) + # Stands in for the videos, which the edit must link rather than copy. + (root / "videos").mkdir() + (root / "videos" / "clip.mp4").write_bytes(b"frames") + inode = (root / "videos" / "clip.mp4").stat().st_ino + + status = _edit_tasks("alice/demo", (" pick the bowl ", 1)) + + assert status["local_error"] is None + assert status["hub_updated"] is False + assert _episode_tasks("alice/demo") == ["task 0", "pick the bowl", "task 2"] + assert (root / "videos" / "clip.mp4").stat().st_ino == inode + assert not list((tmp_lerobot_home / "alice").glob(".demo.task-*")) + + +def test_set_episode_task_relabels_every_episode(tmp_lerobot_home: Path, fake_hub: MagicMock) -> None: + _record_dataset("alice/demo", episodes=3) + + _edit_tasks("alice/demo", ("pick the bowl", None)) + + assert _episode_tasks("alice/demo") == ["pick the bowl"] * 3 + + +def test_set_episode_task_returns_before_the_write( + tmp_lerobot_home: Path, fake_hub: MagicMock, monkeypatch: pytest.MonkeyPatch +) -> None: + import threading + + from lelab import datasets as datasets_mod + + _record_dataset("alice/demo", episodes=2) + gate = threading.Event() + real_rewrite = datasets_mod._rewrite_tasks + monkeypatch.setattr(datasets_mod, "_rewrite_tasks", lambda *args: (gate.wait(30), real_rewrite(*args))[1]) + + result = datasets_mod.handle_set_episode_task("alice/demo", "pick the bowl", 0) + + assert result["queued"] is True + assert datasets_mod.handle_episode_task_status("alice/demo") == {"success": True, "pending": True} + assert _episode_tasks("alice/demo") == ["task 0", "task 1"] + gate.set() + _edit_tasks("alice/demo") + assert _episode_tasks("alice/demo") == ["pick the bowl", "task 1"] + + +def test_edits_queued_during_a_rewrite_are_merged_into_one_and_pushed_once( + tmp_lerobot_home: Path, fake_hub: MagicMock, monkeypatch: pytest.MonkeyPatch +) -> None: + import threading + + from lelab import datasets as datasets_mod + + _record_dataset("alice/demo", episodes=4) + fake_hub.repo_exists.return_value = True + started, gate = threading.Event(), threading.Event() + calls = [] + real_rewrite = datasets_mod._rewrite_tasks + + def gated_rewrite(repo_id, source_dir, default_task, episode_tasks): + calls.append((default_task, dict(episode_tasks))) + started.set() + gate.wait(30) + real_rewrite(repo_id, source_dir, default_task, episode_tasks) + + monkeypatch.setattr(datasets_mod, "_rewrite_tasks", gated_rewrite) + + datasets_mod.handle_set_episode_task("alice/demo", "first", 0) + assert started.wait(30) + datasets_mod.handle_set_episode_task("alice/demo", "second", 1) + datasets_mod.handle_set_episode_task("alice/demo", "everything", None) + datasets_mod.handle_set_episode_task("alice/demo", "third", 2) + datasets_mod.handle_set_episode_task("alice/demo", "fourth", 2) + gate.set() + status = _edit_tasks("alice/demo") + + # "everything" replaces the queued per-episode edits; later edits override it; latest wins. + assert calls == [(None, {0: "first"}), ("everything", {2: "fourth"})] + assert _episode_tasks("alice/demo") == ["everything", "everything", "fourth", "everything"] + assert status["hub_updated"] is True + fake_hub.upload_folder.assert_called_once() + kwargs = fake_hub.upload_folder.call_args.kwargs + assert kwargs["allow_patterns"] == ["meta/**", "data/**", "videos/**"] + assert kwargs["delete_patterns"] == ["meta/**", "data/**", "videos/**"] + + +def test_set_episode_task_does_not_contact_the_hub_for_a_local_namespace( + tmp_lerobot_home: Path, fake_hub: MagicMock +) -> None: + _record_dataset("local/demo", episodes=2) + + status = _edit_tasks("local/demo", ("pick the bowl", None)) + + assert status["hub_updated"] is False + fake_hub.repo_exists.assert_not_called() + fake_hub.upload_folder.assert_not_called() + + +@pytest.mark.parametrize( + ("status_code", "reason"), + [(401, "not logged in"), (403, "no write access to alice/demo"), (500, "hub exploded")], +) +def test_set_episode_task_keeps_the_local_edit_when_the_hub_push_fails( + tmp_lerobot_home: Path, fake_hub: MagicMock, status_code: int, reason: str +) -> None: + _record_dataset("alice/demo", episodes=2) + fake_hub.repo_exists.return_value = True + error = RuntimeError("hub exploded") + error.response = SimpleNamespace(status_code=status_code) + fake_hub.upload_folder.side_effect = error + + status = _edit_tasks("alice/demo", ("pick the bowl", None)) + + assert status["hub_updated"] is False + assert reason in status["hub_error"] + assert status["local_error"] is None + assert _episode_tasks("alice/demo") == ["pick the bowl"] * 2 + + +@pytest.mark.parametrize( + ("task", "episode_index", "message"), + [(" ", 0, "cannot be empty"), ("pick", 3, "Episode 3 is not in alice/demo")], +) +def test_set_episode_task_rejects_invalid_request( + tmp_lerobot_home: Path, task: str, episode_index: int, message: str +) -> None: + from lelab.datasets import handle_set_episode_task + + _record_dataset("alice/demo", episodes=3) + with pytest.raises(ValueError, match=message): + handle_set_episode_task("alice/demo", task, episode_index) + assert _episode_tasks("alice/demo") == ["task 0", "task 1", "task 2"] + + +def test_set_episode_task_reports_a_failed_write_and_leaves_the_source_untouched( + tmp_lerobot_home: Path, fake_hub: MagicMock, monkeypatch: pytest.MonkeyPatch +) -> None: + from lelab import datasets as datasets_mod + + _record_dataset("alice/demo", episodes=2) + fake_hub.repo_exists.return_value = True + real_modify = datasets_mod.modify_tasks + + def failing_modify(dataset, **kwargs): + real_modify(dataset, **kwargs) + raise RuntimeError("disk full") + + monkeypatch.setattr(datasets_mod, "modify_tasks", failing_modify) + + status = _edit_tasks("alice/demo", ("pick the bowl", None)) + + assert status["local_error"] == "disk full" + assert status["hub_updated"] is False + fake_hub.upload_folder.assert_not_called() + assert _episode_tasks("alice/demo") == ["task 0", "task 1"] + assert not list((tmp_lerobot_home / "alice").glob(".demo.task-*")) + + +def test_moves_and_task_edits_of_the_same_dataset_exclude_each_other( + tmp_lerobot_home: Path, fake_hub: MagicMock, monkeypatch: pytest.MonkeyPatch +) -> None: + import threading + + from lelab import datasets as datasets_mod + + _record_dataset("alice/demo", episodes=3) + gate = threading.Event() + real_rewrite = datasets_mod._rewrite_tasks + monkeypatch.setattr(datasets_mod, "_rewrite_tasks", lambda *args: (gate.wait(30), real_rewrite(*args))[1]) + datasets_mod.handle_set_episode_task("alice/demo", "pick the bowl", 0) + + with pytest.raises(ValueError, match="being edited"): + datasets_mod.handle_move_episodes("alice/demo", [1], "rejected") + gate.set() + _edit_tasks("alice/demo") + + # Mid-move, an edit of the same dataset is refused. + with datasets_mod._exclusive_rewrite("alice/demo"), pytest.raises(ValueError, match="being edited"): + datasets_mod.handle_set_episode_task("alice/demo", "x", 0) + + +def test_episode_task_route_refuses_while_recording( + client: TestClient, monkeypatch: pytest.MonkeyPatch +) -> None: + from lelab import datasets as datasets_mod, record + + monkeypatch.setattr(record, "recording_active", True) + monkeypatch.setattr( + datasets_mod, "handle_set_episode_task", lambda *_: pytest.fail("must not edit while recording") + ) + + response = client.post("/datasets/episode-task", json={"repo_id": "alice/demo", "task": "pick"}) + assert response.status_code == 409 + + +@pytest.mark.parametrize("body", [{"task": ""}, {"task": "pick", "episode_index": -1}]) +def test_episode_task_route_rejects_bad_request_before_editing( + client: TestClient, monkeypatch: pytest.MonkeyPatch, body: dict +) -> None: + from lelab import datasets as datasets_mod + + monkeypatch.setattr(datasets_mod, "handle_set_episode_task", lambda *_: pytest.fail("must not edit")) + response = client.post("/datasets/episode-task", json={"repo_id": "alice/demo", **body}) + assert response.status_code == 422 + + +def test_episode_task_route_maps_validation_error( + client: TestClient, monkeypatch: pytest.MonkeyPatch +) -> None: + from lelab import datasets as datasets_mod + + monkeypatch.setattr( + datasets_mod, "handle_set_episode_task", lambda *_: (_ for _ in ()).throw(ValueError("bad input")) + ) + response = client.post( + "/datasets/episode-task", json={"repo_id": "alice/demo", "task": "pick", "episode_index": 0} + ) + assert response.status_code == 400 + assert response.json()["detail"] == "bad input" + + # ── episode browsing ──────────────────────────────────────────────────────── # # The handlers are covered end-to-end through the routes, so the query-param