From d8c45fafc367d400803c6680b0fe9076c1616765 Mon Sep 17 00:00:00 2001 From: "ravindu.somawansa" Date: Tue, 20 Oct 2026 09:34:06 +0200 Subject: [PATCH 1/5] feat(datasets): edit an episode's task in the dataset browser --- .../src/components/dataset/EpisodeList.tsx | 57 +++++--- .../dataset/EpisodeTaskDialog.test.tsx | 61 ++++++++ .../components/dataset/EpisodeTaskDialog.tsx | 130 ++++++++++++++++++ frontend/src/lib/replayApi.ts | 22 +++ frontend/src/pages/EditDataset.tsx | 29 +++- lelab/datasets.py | 60 +++++++- lelab/server.py | 18 +++ tests/test_datasets.py | 109 +++++++++++++++ 8 files changed, 464 insertions(+), 22 deletions(-) create mode 100644 frontend/src/components/dataset/EpisodeTaskDialog.test.tsx create mode 100644 frontend/src/components/dataset/EpisodeTaskDialog.tsx diff --git a/frontend/src/components/dataset/EpisodeList.tsx b/frontend/src/components/dataset/EpisodeList.tsx index 120a24aca..36186b60a 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 000000000..71c0ec3a6 --- /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={async (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 000000000..bc7d3bf02 --- /dev/null +++ b/frontend/src/components/dataset/EpisodeTaskDialog.tsx @@ -0,0 +1,130 @@ +import { useEffect, useState } from "react"; +import { Loader2 } from "lucide-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) => Promise; +} + +export default function EpisodeTaskDialog({ + repoId, + episodeIndex, + currentTask, + totalEpisodes, + open, + onOpenChange, + onSave, +}: EpisodeTaskDialogProps) { + const [task, setTask] = useState(""); + const [allEpisodes, setAllEpisodes] = useState(false); + const [saving, setSaving] = useState(false); + const [error, setError] = useState(null); + + useEffect(() => { + if (open) { + setTask(currentTask); + setAllEpisodes(false); + setSaving(false); + setError(null); + } + }, [open, currentTask]); + + const canSave = task.trim() !== "" && !saving; + + const handleOpenChange = (nextOpen: boolean) => { + if (saving) return; + onOpenChange(nextOpen); + }; + + const handleSave = async () => { + if (!canSave) return; + setSaving(true); + try { + await onSave(task.trim(), allEpisodes ? null : episodeIndex); + onOpenChange(false); + } catch (caughtError) { + setError(caughtError instanceof Error ? caughtError.message : "Could not save the task."); + } finally { + setSaving(false); + } + }; + + return ( + + + + Edit task + + The task is the instruction a policy is trained on, saved with each episode of{" "} + {repoId}. + + + +
+
+ + { + setError(null); + setTask(event.target.value); + }} + onKeyDown={(event) => { + if (event.key === "Enter") void handleSave(); + }} + disabled={saving} + className="border-gray-700 bg-gray-800 text-white" + /> +
+ {totalEpisodes > 1 && ( +
+ setAllEpisodes(checked === true)} + disabled={saving} + className="border-gray-500" + /> + +
+ )} + {error &&

{error}

} +
+ + + + + + + + ); +} diff --git a/frontend/src/lib/replayApi.ts b/frontend/src/lib/replayApi.ts index 4e2c8b3ff..c43e2b1e8 100644 --- a/frontend/src/lib/replayApi.ts +++ b/frontend/src/lib/replayApi.ts @@ -84,3 +84,25 @@ export async function moveEpisodes( action: "Move episodes", }); } + +export interface EpisodeTaskResult { + success: true; + repo_id: string; + task: string; + updated_episodes: number; +} + +/** Set the task of one episode, or of every episode when `episodeIndex` is null. */ +export async function setEpisodeTask( + baseUrl: string, + fetcher: Fetcher, + repoId: string, + task: string, + episodeIndex: number | null, +): Promise { + return apiRequest(baseUrl, fetcher, "/datasets/episode-task", { + method: "POST", + body: { repo_id: repoId, task, episode_index: episodeIndex }, + action: "Edit task", + }); +} diff --git a/frontend/src/pages/EditDataset.tsx b/frontend/src/pages/EditDataset.tsx index 1d0f5c107..e7a873163 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, 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"), @@ -125,6 +127,18 @@ const EditDataset = () => { description: `${result.moved_episodes} moved to ${result.repo_id}, ${result.remaining_episodes} left in ${repoId}.`, }); }; + const handleSaveTask = async (task: string, targetEpisode: number | null) => { + if (!repoId) return; + const result = await setEpisodeTask(baseUrl, fetchWithHeaders, repoId, task, targetEpisode); + refreshEpisodes(); + toast({ + title: "Task saved", + description: + targetEpisode === null + ? `All ${result.updated_episodes} episodes now have the task "${result.task}".` + : `Episode ${targetEpisode} now has the task "${result.task}".`, + }); + }; const setEpisode = (next: number) => { if (!repoId) return; setSearchParams(withApi({ dataset: repoId, episode: String(next) })); @@ -310,6 +324,7 @@ const EditDataset = () => { episodes={data.episodes} selected={episodeIndex} onSelect={setEpisode} + onEditTask={selectedIsLocal ? () => setTaskOpen(true) : undefined} />
@@ -355,6 +370,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]: + """Set the task of one episode of a local dataset, or of every episode when `episode_index` is None. + + LeRobot's `modify_tasks` rewrites the dataset in place, so it runs on a + staged copy that is swapped in afterwards: a failure leaves the source as it + was. Only `meta/` and `data/` are rewritten, so everything else (videos, + images) is hard-linked into the copy where the filesystem allows it. + """ + 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}") + + 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) + if episode_index is None: + modify_tasks(dataset, new_task=task) + else: + modify_tasks(dataset, episode_tasks={episode_index: task}) + 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) + + return { + "success": True, + "repo_id": repo_id, + "task": task, + "updated_episodes": total if episode_index is None else 1, + } + + # ── 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 1c636d5b2..16a519956 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,17 @@ 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): + """Set the task of one episode, or of 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.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 b6e89f188..e5dba6c1a 100644 --- a/tests/test_datasets.py +++ b/tests/test_datasets.py @@ -472,6 +472,115 @@ def test_move_episodes_route_rejects_bad_indices_before_editing( assert response.status_code == 422 +def test_set_episode_task_relabels_one_episode(tmp_lerobot_home: Path) -> None: + from lelab.datasets import handle_set_episode_task + + 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 + + result = handle_set_episode_task("alice/demo", " pick the bowl ", episode_index=1) + + assert result == { + "success": True, + "repo_id": "alice/demo", + "task": "pick the bowl", + "updated_episodes": 1, + } + 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) -> None: + from lelab.datasets import handle_set_episode_task + + _record_dataset("alice/demo", episodes=3) + + result = handle_set_episode_task("alice/demo", "pick the bowl") + + assert result["updated_episodes"] == 3 + assert _episode_tasks("alice/demo") == ["pick the bowl"] * 3 + + +@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_leaves_source_untouched_on_failure( + tmp_lerobot_home: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + from lelab import datasets as datasets_mod + + _record_dataset("alice/demo", episodes=2) + 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) + + with pytest.raises(RuntimeError, match="disk full"): + datasets_mod.handle_set_episode_task("alice/demo", "pick the bowl") + + assert _episode_tasks("alice/demo") == ["task 0", "task 1"] + assert not list((tmp_lerobot_home / "alice").glob(".demo.task-*")) + + +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 From 9256a5c83a2a3876ab144bbc82ec4b60eff280aa Mon Sep 17 00:00:00 2001 From: Nicolas Rabault Date: Sat, 31 Oct 2026 23:42:36 +0100 Subject: [PATCH 2/5] feat(datasets): update the Hub copy when a task edit hits a dataset that is also on the Hub After the edited copy is swapped in locally, upload meta/ and data/ to the dataset's Hub repo when it exists. A failed push keeps the local edit and the toast says the Hub copy was not updated, with the reason. --- frontend/src/lib/replayApi.ts | 2 + frontend/src/pages/EditDataset.tsx | 16 +++++-- lelab/datasets.py | 40 +++++++++++++++- tests/test_datasets.py | 75 ++++++++++++++++++++++++++++-- 4 files changed, 124 insertions(+), 9 deletions(-) diff --git a/frontend/src/lib/replayApi.ts b/frontend/src/lib/replayApi.ts index c43e2b1e8..e7d6bda6d 100644 --- a/frontend/src/lib/replayApi.ts +++ b/frontend/src/lib/replayApi.ts @@ -90,6 +90,8 @@ export interface EpisodeTaskResult { repo_id: string; task: string; updated_episodes: number; + hub_updated: boolean; + hub_error: string | null; } /** Set the task of one episode, or of every episode when `episodeIndex` is null. */ diff --git a/frontend/src/pages/EditDataset.tsx b/frontend/src/pages/EditDataset.tsx index e7a873163..a1ab1b2c0 100644 --- a/frontend/src/pages/EditDataset.tsx +++ b/frontend/src/pages/EditDataset.tsx @@ -131,12 +131,18 @@ const EditDataset = () => { if (!repoId) return; const result = await setEpisodeTask(baseUrl, fetchWithHeaders, repoId, task, targetEpisode); refreshEpisodes(); + const scope = + targetEpisode === null + ? `All ${result.updated_episodes} episodes now have the task "${result.task}".` + : `Episode ${targetEpisode} now has the task "${result.task}".`; toast({ - title: "Task saved", - description: - targetEpisode === null - ? `All ${result.updated_episodes} episodes now have the task "${result.task}".` - : `Episode ${targetEpisode} now has the task "${result.task}".`, + title: result.hub_error ? "Task saved locally only" : "Task saved", + description: result.hub_error + ? `${scope} The Hub copy was not updated: ${result.hub_error}.` + : result.hub_updated + ? `${scope} The Hub copy is updated too.` + : scope, + variant: result.hub_error ? "destructive" : undefined, }); }; const setEpisode = (next: number) => { diff --git a/lelab/datasets.py b/lelab/datasets.py index 2a00a3590..b84b3f576 100644 --- a/lelab/datasets.py +++ b/lelab/datasets.py @@ -339,13 +339,48 @@ def handle_move_episodes( } +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 an edit of `meta/` and `data/` to the dataset's Hub repo when it has one. + + Returns (pushed, error). Videos are untouched by the edit and are not sent. + """ + 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=["meta/**", "data/**"], + 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 handle_set_episode_task(repo_id: str, task: str, episode_index: int | None = None) -> dict[str, Any]: """Set the task of one episode of a local dataset, or of every episode when `episode_index` is None. LeRobot's `modify_tasks` rewrites the dataset in place, so it runs on a staged copy that is swapped in afterwards: a failure leaves the source as it was. Only `meta/` and `data/` are rewritten, so everything else (videos, - images) is hard-linked into the copy where the filesystem allows it. + images) is hard-linked into the copy where the filesystem allows it. Once the + local copy is swapped in, a dataset that also lives on the Hub is updated + there too; a failed push keeps the local edit and is reported in `hub_error`. """ task = task.strip() if not task: @@ -388,11 +423,14 @@ def copy_or_link(src: str, dst: str) -> None: finally: shutil.rmtree(staging, ignore_errors=True) + hub_updated, hub_error = _push_edit_to_hub(repo_id, source_dir) return { "success": True, "repo_id": repo_id, "task": task, "updated_episodes": total if episode_index is None else 1, + "hub_updated": hub_updated, + "hub_error": hub_error, } diff --git a/tests/test_datasets.py b/tests/test_datasets.py index e5dba6c1a..09e2114bc 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,7 +473,18 @@ def test_move_episodes_route_rejects_bad_indices_before_editing( assert response.status_code == 422 -def test_set_episode_task_relabels_one_episode(tmp_lerobot_home: Path) -> None: +@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 test_set_episode_task_relabels_one_episode(tmp_lerobot_home: Path, fake_hub: MagicMock) -> None: from lelab.datasets import handle_set_episode_task root = _record_dataset("alice/demo", episodes=3) @@ -488,13 +500,15 @@ def test_set_episode_task_relabels_one_episode(tmp_lerobot_home: Path) -> None: "repo_id": "alice/demo", "task": "pick the bowl", "updated_episodes": 1, + "hub_updated": False, + "hub_error": None, } 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) -> None: +def test_set_episode_task_relabels_every_episode(tmp_lerobot_home: Path, fake_hub: MagicMock) -> None: from lelab.datasets import handle_set_episode_task _record_dataset("alice/demo", episodes=3) @@ -505,6 +519,61 @@ def test_set_episode_task_relabels_every_episode(tmp_lerobot_home: Path) -> None assert _episode_tasks("alice/demo") == ["pick the bowl"] * 3 +def test_set_episode_task_pushes_only_meta_and_data_when_the_repo_is_on_the_hub( + tmp_lerobot_home: Path, fake_hub: MagicMock +) -> None: + from lelab.datasets import handle_set_episode_task + + root = _record_dataset("alice/demo", episodes=2) + fake_hub.repo_exists.return_value = True + + result = handle_set_episode_task("alice/demo", "pick the bowl") + + assert result["hub_updated"] is True + assert result["hub_error"] is None + fake_hub.upload_folder.assert_called_once() + kwargs = fake_hub.upload_folder.call_args.kwargs + assert kwargs["repo_id"] == "alice/demo" + assert kwargs["folder_path"] == root + assert kwargs["allow_patterns"] == ["meta/**", "data/**"] + + +def test_set_episode_task_does_not_contact_the_hub_for_a_local_namespace( + tmp_lerobot_home: Path, fake_hub: MagicMock +) -> None: + from lelab.datasets import handle_set_episode_task + + _record_dataset("local/demo", episodes=2) + + result = handle_set_episode_task("local/demo", "pick the bowl") + + assert result["hub_updated"] is False + fake_hub.repo_exists.assert_not_called() + fake_hub.upload_folder.assert_not_called() + + +@pytest.mark.parametrize( + ("status", "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: int, reason: str +) -> None: + from lelab.datasets import handle_set_episode_task + + _record_dataset("alice/demo", episodes=2) + fake_hub.repo_exists.return_value = True + error = RuntimeError("hub exploded") + error.response = SimpleNamespace(status_code=status) + fake_hub.upload_folder.side_effect = error + + result = handle_set_episode_task("alice/demo", "pick the bowl") + + assert result["hub_updated"] is False + assert reason in result["hub_error"] + 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")], From 3a076e6f3dcc4a2eaebfe70bd3a3a6b2beda20e6 Mon Sep 17 00:00:00 2001 From: Nicolas Rabault Date: Sat, 31 Oct 2026 23:46:23 +0100 Subject: [PATCH 3/5] fix(datasets): mirror the local dataset folders to the Hub after a task edit The Hub copy is mirrored from the local one: meta/, data/ and videos/ are uploaded and files the local copy no longer has are deleted, in one commit. Uploading only meta/ and data/ corrupted a Hub copy that had diverged, for example after episodes were moved out locally or recorded without a push. Unchanged files are dropped from the commit by huggingface_hub. --- lelab/datasets.py | 16 +++++++++++++--- tests/test_datasets.py | 5 +++-- 2 files changed, 16 insertions(+), 5 deletions(-) diff --git a/lelab/datasets.py b/lelab/datasets.py index b84b3f576..b4261c38e 100644 --- a/lelab/datasets.py +++ b/lelab/datasets.py @@ -339,6 +339,9 @@ 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: @@ -349,9 +352,15 @@ def _hub_failure_reason(exc: Exception, repo_id: str) -> str: def _push_edit_to_hub(repo_id: str, root: Path) -> tuple[bool, str | None]: - """Mirror an edit of `meta/` and `data/` to the dataset's Hub repo when it has one. + """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). Videos are untouched by the edit and are not sent. + Returns (pushed, error). """ if repo_id.split("/")[0] == LOCAL_DATASET_NAMESPACE: return False, None @@ -363,7 +372,8 @@ def _push_edit_to_hub(repo_id: str, root: Path) -> tuple[bool, str | None]: repo_id=repo_id, repo_type="dataset", folder_path=root, - allow_patterns=["meta/**", "data/**"], + allow_patterns=_HUB_MIRRORED_GLOBS, + delete_patterns=_HUB_MIRRORED_GLOBS, commit_message="Edit episode tasks", ) except Exception as exc: diff --git a/tests/test_datasets.py b/tests/test_datasets.py index 09e2114bc..da6146f77 100644 --- a/tests/test_datasets.py +++ b/tests/test_datasets.py @@ -519,7 +519,7 @@ def test_set_episode_task_relabels_every_episode(tmp_lerobot_home: Path, fake_hu assert _episode_tasks("alice/demo") == ["pick the bowl"] * 3 -def test_set_episode_task_pushes_only_meta_and_data_when_the_repo_is_on_the_hub( +def test_set_episode_task_mirrors_the_dataset_folders_when_the_repo_is_on_the_hub( tmp_lerobot_home: Path, fake_hub: MagicMock ) -> None: from lelab.datasets import handle_set_episode_task @@ -535,7 +535,8 @@ def test_set_episode_task_pushes_only_meta_and_data_when_the_repo_is_on_the_hub( kwargs = fake_hub.upload_folder.call_args.kwargs assert kwargs["repo_id"] == "alice/demo" assert kwargs["folder_path"] == root - assert kwargs["allow_patterns"] == ["meta/**", "data/**"] + 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( From 00cee46fce3e40f1e0fc898fbdf79280b9e5c2ef Mon Sep 17 00:00:00 2001 From: Nicolas Rabault Date: Mon, 2 Nov 2026 17:47:37 +0100 Subject: [PATCH 4/5] feat(datasets): write task edits in the background The edit endpoint queues the change and returns at once. One worker per dataset merges the edits queued while a rewrite runs into a single modify_tasks call and mirrors to the Hub once the queue is idle, so a burst of edits gives one Hub commit. The page shows the new task immediately, polls a status endpoint while edits are pending and reports a failed write or Hub push in a toast. Moving episodes is refused while task edits of the same dataset are pending. --- .../dataset/EpisodeTaskDialog.test.tsx | 2 +- .../components/dataset/EpisodeTaskDialog.tsx | 44 +---- frontend/src/hooks/useEpisodes.ts | 16 +- frontend/src/lib/replayApi.ts | 31 ++-- frontend/src/pages/EditDataset.tsx | 70 +++++-- lelab/datasets.py | 171 ++++++++++++++---- lelab/server.py | 8 +- tests/test_datasets.py | 167 ++++++++++++----- 8 files changed, 367 insertions(+), 142 deletions(-) diff --git a/frontend/src/components/dataset/EpisodeTaskDialog.test.tsx b/frontend/src/components/dataset/EpisodeTaskDialog.test.tsx index 71c0ec3a6..9b1e34f47 100644 --- a/frontend/src/components/dataset/EpisodeTaskDialog.test.tsx +++ b/frontend/src/components/dataset/EpisodeTaskDialog.test.tsx @@ -22,7 +22,7 @@ test("saves the task for the browsed episode or for every episode", async () => totalEpisodes={5} open onOpenChange={() => {}} - onSave={async (task, episodeIndex) => { + onSave={(task, episodeIndex) => { saved.push([task, episodeIndex]); }} />, diff --git a/frontend/src/components/dataset/EpisodeTaskDialog.tsx b/frontend/src/components/dataset/EpisodeTaskDialog.tsx index bc7d3bf02..8ca08ecb6 100644 --- a/frontend/src/components/dataset/EpisodeTaskDialog.tsx +++ b/frontend/src/components/dataset/EpisodeTaskDialog.tsx @@ -1,5 +1,4 @@ import { useEffect, useState } from "react"; -import { Loader2 } from "lucide-react"; import { Button } from "@/components/ui/button"; import { Checkbox } from "@/components/ui/checkbox"; @@ -22,7 +21,7 @@ interface EpisodeTaskDialogProps { open: boolean; onOpenChange: (open: boolean) => void; /** `episodeIndex` is null when the task applies to every episode. */ - onSave: (task: string, episodeIndex: number | null) => Promise; + onSave: (task: string, episodeIndex: number | null) => void; } export default function EpisodeTaskDialog({ @@ -36,40 +35,24 @@ export default function EpisodeTaskDialog({ }: EpisodeTaskDialogProps) { const [task, setTask] = useState(""); const [allEpisodes, setAllEpisodes] = useState(false); - const [saving, setSaving] = useState(false); - const [error, setError] = useState(null); useEffect(() => { if (open) { setTask(currentTask); setAllEpisodes(false); - setSaving(false); - setError(null); } }, [open, currentTask]); - const canSave = task.trim() !== "" && !saving; + const canSave = task.trim() !== ""; - const handleOpenChange = (nextOpen: boolean) => { - if (saving) return; - onOpenChange(nextOpen); - }; - - const handleSave = async () => { + const handleSave = () => { if (!canSave) return; - setSaving(true); - try { - await onSave(task.trim(), allEpisodes ? null : episodeIndex); - onOpenChange(false); - } catch (caughtError) { - setError(caughtError instanceof Error ? caughtError.message : "Could not save the task."); - } finally { - setSaving(false); - } + onSave(task.trim(), allEpisodes ? null : episodeIndex); + onOpenChange(false); }; return ( - + Edit task @@ -87,14 +70,10 @@ export default function EpisodeTaskDialog({ { - setError(null); - setTask(event.target.value); - }} + onChange={(event) => setTask(event.target.value)} onKeyDown={(event) => { - if (event.key === "Enter") void handleSave(); + if (event.key === "Enter") handleSave(); }} - disabled={saving} className="border-gray-700 bg-gray-800 text-white" />
@@ -104,7 +83,6 @@ export default function EpisodeTaskDialog({ id="episode-task-all" checked={allEpisodes} onCheckedChange={(checked) => setAllEpisodes(checked === true)} - disabled={saving} className="border-gray-500" /> )} - {error &&

{error}

} - diff --git a/frontend/src/hooks/useEpisodes.ts b/frontend/src/hooks/useEpisodes.ts index 791c6dddf..97875a1bc 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 e7d6bda6d..8d793ef4e 100644 --- a/frontend/src/lib/replayApi.ts +++ b/frontend/src/lib/replayApi.ts @@ -85,26 +85,35 @@ export async function moveEpisodes( }); } -export interface EpisodeTaskResult { - success: true; - repo_id: string; - task: string; - updated_episodes: number; - hub_updated: boolean; - hub_error: string | null; -} +export type EpisodeTaskStatus = + | { pending: true } + | { pending: false; local_error?: string | null; hub_error?: string | null; hub_updated?: boolean }; -/** Set the task of one episode, or of every episode when `episodeIndex` is null. */ +/** 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 { - return apiRequest(baseUrl, fetcher, "/datasets/episode-task", { +): 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 a1ab1b2c0..7c097f321 100644 --- a/frontend/src/pages/EditDataset.tsx +++ b/frontend/src/pages/EditDataset.tsx @@ -30,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, setEpisodeTask } 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"; @@ -72,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, @@ -127,24 +127,57 @@ const EditDataset = () => { description: `${result.moved_episodes} moved to ${result.repo_id}, ${result.remaining_episodes} left in ${repoId}.`, }); }; - const handleSaveTask = async (task: string, targetEpisode: number | null) => { + // 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 result = await setEpisodeTask(baseUrl, fetchWithHeaders, repoId, task, targetEpisode); - refreshEpisodes(); - const scope = - targetEpisode === null - ? `All ${result.updated_episodes} episodes now have the task "${result.task}".` - : `Episode ${targetEpisode} now has the task "${result.task}".`; - toast({ - title: result.hub_error ? "Task saved locally only" : "Task saved", - description: result.hub_error - ? `${scope} The Hub copy was not updated: ${result.hub_error}.` - : result.hub_updated - ? `${scope} The Hub copy is updated too.` - : scope, - variant: result.hub_error ? "destructive" : undefined, - }); + 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) })); @@ -323,6 +356,7 @@ const EditDataset = () => { {detail && } + {savingRepo === repoId && Saving tasks…}
diff --git a/lelab/datasets.py b/lelab/datasets.py index b4261c38e..61fb1d955 100644 --- a/lelab/datasets.py +++ b/lelab/datasets.py @@ -18,6 +18,8 @@ import re import shutil import tempfile +import threading +from contextlib import contextmanager from datetime import UTC, datetime from pathlib import Path from typing import Any @@ -267,6 +269,16 @@ def handle_move_episodes( 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. """ + with _exclusive_rewrite(*(r for r in (repo_id, target_repo_id) if r)): + return _move_episodes(repo_id, episode_indices, output_name, target_repo_id) + + +def _move_episodes( + repo_id: str, + episode_indices: list[int], + output_name: str | None = None, + target_repo_id: str | None = None, +) -> dict[str, Any]: episode_indices = sorted(set(episode_indices)) if not episode_indices: raise ValueError("Select at least one episode to move") @@ -382,29 +394,13 @@ def _push_edit_to_hub(repo_id: str, root: Path) -> tuple[bool, str | None]: return True, None -def handle_set_episode_task(repo_id: str, task: str, episode_index: int | None = None) -> dict[str, Any]: - """Set the task of one episode of a local dataset, or of every episode when `episode_index` is None. - - LeRobot's `modify_tasks` rewrites the dataset in place, so it runs on a - staged copy that is swapped in afterwards: a failure leaves the source as it - was. Only `meta/` and `data/` are rewritten, so everything else (videos, - images) is hard-linked into the copy where the filesystem allows it. Once the - local copy is swapped in, a dataset that also lives on the Hub is updated - there too; a failed push keeps the local edit and is reported in `hub_error`. - """ - task = task.strip() - if not task: - raise ValueError("The task cannot be empty") +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. - 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}") + `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"): @@ -420,10 +416,7 @@ def copy_or_link(src: str, dst: str) -> None: try: shutil.copytree(source_dir, edited_dir, copy_function=copy_or_link) dataset = LeRobotDataset(repo_id=repo_id, root=edited_dir) - if episode_index is None: - modify_tasks(dataset, new_task=task) - else: - modify_tasks(dataset, episode_tasks={episode_index: task}) + modify_tasks(dataset, new_task=default_task, episode_tasks=episode_tasks or None) source_dir.rename(staging / "old") try: edited_dir.rename(source_dir) @@ -433,15 +426,123 @@ def copy_or_link(src: str, dst: str) -> None: finally: shutil.rmtree(staging, ignore_errors=True) - hub_updated, hub_error = _push_edit_to_hub(repo_id, source_dir) - return { - "success": True, - "repo_id": repo_id, - "task": task, - "updated_episodes": total if episode_index is None else 1, - "hub_updated": hub_updated, - "hub_error": hub_error, - } + +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 ──────────────────────────────────────────────────────── diff --git a/lelab/server.py b/lelab/server.py index 16a519956..dbb69ba68 100644 --- a/lelab/server.py +++ b/lelab/server.py @@ -472,7 +472,7 @@ def merge_local_datasets(request: MergeDatasetsRequest): @app.post("/datasets/episode-task") def set_episode_task(request: EpisodeTaskRequest): - """Set the task of one episode, or of every episode, of a local dataset.""" + """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: @@ -481,6 +481,12 @@ def set_episode_task(request: EpisodeTaskRequest): 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 da6146f77..19d032a2a 100644 --- a/tests/test_datasets.py +++ b/tests/test_datasets.py @@ -484,57 +484,116 @@ def fake_hub(monkeypatch: pytest.MonkeyPatch) -> MagicMock: return api -def test_set_episode_task_relabels_one_episode(tmp_lerobot_home: Path, fake_hub: MagicMock) -> None: - from lelab.datasets import handle_set_episode_task +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 - result = handle_set_episode_task("alice/demo", " pick the bowl ", episode_index=1) + status = _edit_tasks("alice/demo", (" pick the bowl ", 1)) - assert result == { - "success": True, - "repo_id": "alice/demo", - "task": "pick the bowl", - "updated_episodes": 1, - "hub_updated": False, - "hub_error": None, - } + 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: - from lelab.datasets import handle_set_episode_task - _record_dataset("alice/demo", episodes=3) - result = handle_set_episode_task("alice/demo", "pick the bowl") + _edit_tasks("alice/demo", ("pick the bowl", None)) - assert result["updated_episodes"] == 3 assert _episode_tasks("alice/demo") == ["pick the bowl"] * 3 -def test_set_episode_task_mirrors_the_dataset_folders_when_the_repo_is_on_the_hub( - tmp_lerobot_home: Path, fake_hub: MagicMock +def test_set_episode_task_returns_before_the_write( + tmp_lerobot_home: Path, fake_hub: MagicMock, monkeypatch: pytest.MonkeyPatch ) -> None: - from lelab.datasets import handle_set_episode_task + import threading - root = _record_dataset("alice/demo", episodes=2) - fake_hub.repo_exists.return_value = True + 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"] - result = handle_set_episode_task("alice/demo", "pick the bowl") - assert result["hub_updated"] is True - assert result["hub_error"] is None +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["repo_id"] == "alice/demo" - assert kwargs["folder_path"] == root assert kwargs["allow_patterns"] == ["meta/**", "data/**", "videos/**"] assert kwargs["delete_patterns"] == ["meta/**", "data/**", "videos/**"] @@ -542,36 +601,33 @@ def test_set_episode_task_mirrors_the_dataset_folders_when_the_repo_is_on_the_hu def test_set_episode_task_does_not_contact_the_hub_for_a_local_namespace( tmp_lerobot_home: Path, fake_hub: MagicMock ) -> None: - from lelab.datasets import handle_set_episode_task - _record_dataset("local/demo", episodes=2) - result = handle_set_episode_task("local/demo", "pick the bowl") + status = _edit_tasks("local/demo", ("pick the bowl", None)) - assert result["hub_updated"] is False + assert status["hub_updated"] is False fake_hub.repo_exists.assert_not_called() fake_hub.upload_folder.assert_not_called() @pytest.mark.parametrize( - ("status", "reason"), + ("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: int, reason: str + tmp_lerobot_home: Path, fake_hub: MagicMock, status_code: int, reason: str ) -> None: - from lelab.datasets import handle_set_episode_task - _record_dataset("alice/demo", episodes=2) fake_hub.repo_exists.return_value = True error = RuntimeError("hub exploded") - error.response = SimpleNamespace(status_code=status) + error.response = SimpleNamespace(status_code=status_code) fake_hub.upload_folder.side_effect = error - result = handle_set_episode_task("alice/demo", "pick the bowl") + status = _edit_tasks("alice/demo", ("pick the bowl", None)) - assert result["hub_updated"] is False - assert reason in result["hub_error"] + 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 @@ -590,12 +646,13 @@ def test_set_episode_task_rejects_invalid_request( assert _episode_tasks("alice/demo") == ["task 0", "task 1", "task 2"] -def test_set_episode_task_leaves_source_untouched_on_failure( - tmp_lerobot_home: Path, monkeypatch: pytest.MonkeyPatch +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): @@ -604,13 +661,41 @@ def failing_modify(dataset, **kwargs): monkeypatch.setattr(datasets_mod, "modify_tasks", failing_modify) - with pytest.raises(RuntimeError, match="disk full"): - datasets_mod.handle_set_episode_task("alice/demo", "pick the bowl") + 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"): + with 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 42acb4fb16f5fdbefb1a846ce7788f0e15b41b29 Mon Sep 17 00:00:00 2001 From: Nicolas Rabault Date: Mon, 2 Nov 2026 23:04:38 +0100 Subject: [PATCH 5/5] style(datasets): apply ruff to the background task edits --- lelab/datasets.py | 8 ++++++-- tests/test_datasets.py | 13 ++++--------- 2 files changed, 10 insertions(+), 11 deletions(-) diff --git a/lelab/datasets.py b/lelab/datasets.py index 61fb1d955..c682f7669 100644 --- a/lelab/datasets.py +++ b/lelab/datasets.py @@ -394,7 +394,9 @@ def _push_edit_to_hub(repo_id: str, root: Path) -> tuple[bool, str | None]: return True, None -def _rewrite_tasks(repo_id: str, source_dir: Path, default_task: str | None, episode_tasks: dict[int, str]) -> 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 @@ -522,7 +524,9 @@ def handle_set_episode_task(repo_id: str, task: str, episode_index: int | None = 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() + threading.Thread( + target=_task_edits_worker, args=(repo_id, source_dir, edits), daemon=True + ).start() return {"success": True, "repo_id": repo_id, "queued": True} diff --git a/tests/test_datasets.py b/tests/test_datasets.py index 19d032a2a..2cbf006e0 100644 --- a/tests/test_datasets.py +++ b/tests/test_datasets.py @@ -544,9 +544,7 @@ def test_set_episode_task_returns_before_the_write( _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] - ) + 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) @@ -680,9 +678,7 @@ def test_moves_and_task_edits_of_the_same_dataset_exclude_each_other( _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] - ) + 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"): @@ -691,9 +687,8 @@ def test_moves_and_task_edits_of_the_same_dataset_exclude_each_other( _edit_tasks("alice/demo") # Mid-move, an edit of the same dataset is refused. - with datasets_mod._exclusive_rewrite("alice/demo"): - with pytest.raises(ValueError, match="being edited"): - datasets_mod.handle_set_episode_task("alice/demo", "x", 0) + 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(