diff --git a/CLAUDE.md b/CLAUDE.md index bd756681b..422532923 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -35,13 +35,12 @@ Run the Python tests with `pytest` (config in [pyproject.toml](pyproject.toml); [server.py](lelab/server.py) is a thin FastAPI router. Each feature lives in its own module that owns its global state and exposes `handle_*` functions plus a Pydantic request model: -- [record.py](lelab/record.py) — dataset recording (wraps `lerobot.record.record`); patches `lerobot.common.utils.control_utils` keyboard listener so frontend buttons replace arrow-key controls. +- [record.py](lelab/record.py) — dataset recording (wraps `lerobot.record.record`); patches `lerobot.common.utils.control_utils` keyboard listener so frontend buttons replace arrow-key controls. The worker is a daemon thread, so the server's shutdown hook calls `stop_recording_and_wait()` to let an active session finalize on Ctrl+C. - [teleoperate.py](lelab/teleoperate.py) — leader→follower teleoperation (wraps `lerobot.teleoperate`). - [calibrate.py](lelab/calibrate.py) — step-by-step web calibration with a `CalibrationManager` singleton and `_step_complete` threading.Event. - [train.py](lelab/train.py) — wraps the LeRobot training CLI as a subprocess (psutil for lifecycle, queue for log streaming). - [datasets.py](lelab/datasets.py) — dataset listing (local cache + Hub, merged under a `source` field) plus the episode-browsing handlers behind `/dataset-episodes`, `/dataset-episode`, `/dataset-frame`, `/dataset-thumbnails`, `/dataset-motion` and `/dataset-video`. -- [episode_media.py](lelab/episode_media.py) — reads the LeRobot v3.0 on-disk layout (`meta/episodes/*.parquet` + `videos/`) and decodes frames with PyAV. Deliberately does **not** go through `LeRobotDataset`, so browsing a dataset never builds a robot config or imports torch. Several episodes share one mp4, so locating an episode resolves both the file and the `from_/to_timestamp` window inside it — which is why `/dataset-video` is addressed by chunk/file rather than by episode. -- [dataset_repair.py](lelab/dataset_repair.py) — rebuilds `meta/episodes/` for a recording interrupted before `LeRobotDataset.finalize()` ran. Without an episode index LeRobot reads a local dataset as "not downloaded yet" and fetches it from the Hub, which 404s for a dataset that was never pushed. Call `repair_local_dataset(repo_id)` before opening a recorded dataset with `LeRobotDataset`. +- [episode_media.py](lelab/episode_media.py) — reads the LeRobot v3.0 on-disk layout (`meta/episodes/*.parquet` + `videos/`) and decodes frames with PyAV. Deliberately does **not** go through `LeRobotDataset`, so browsing a dataset never builds a robot config or imports torch. Several episodes share one mp4, so locating an episode resolves both the file and the `from_/to_timestamp` window inside it — which is why `/dataset-video` is addressed by chunk/file rather than by episode. Call `ensure_episode_index(repo_id)` before opening a local dataset with `LeRobotDataset`: LeRobot writes the index only in `finalize()`, and without it reads the dataset as "not downloaded yet" and fetches it from the Hub, which 404s for a dataset that was never pushed. - [utils/config.py](lelab/utils/config.py) — shared paths and persistence: calibration JSON, saved ports, saved config selections. **Import shared constants from here, do not hardcode paths in feature modules.** ### State model diff --git a/lelab/dataset_repair.py b/lelab/dataset_repair.py deleted file mode 100644 index 9e3ce6788..000000000 --- a/lelab/dataset_repair.py +++ /dev/null @@ -1,318 +0,0 @@ -# Copyright 2025 The HuggingFace Inc. team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Rebuild the episode index of a recording that was never finalized. - -``LeRobotDataset.finalize()`` is what writes ``meta/episodes/`` and the parquet -footers; frames, videos and ``meta/info.json`` are written as the session runs. -A recording interrupted before finalize therefore leaves a directory that looks -complete but carries no episode index, and ``LeRobotDataset`` reads that as -"not available locally" and tries to download it from the Hub — a 404 for a -dataset that was never pushed. - -Recovery is bounded by the parquet footer: a data file the writer was still -appending to when it died cannot be read at all, so its frames and every frame -after it are gone. -""" - -import json -import logging -import re -import shutil -from pathlib import Path -from typing import Any - -logger = logging.getLogger(__name__) - -_CHUNK_FILE_RE = re.compile(r"chunk-(\d+)/file-(\d+)\.") - - -class DatasetRepairError(Exception): - """The dataset is damaged beyond what can be rebuilt from disk.""" - - -def _chunk_file(path: Path) -> tuple[int, int]: - match = _CHUNK_FILE_RE.search(path.as_posix()) - if match is None: - raise DatasetRepairError(f"Unexpected dataset file layout: {path}") - return int(match.group(1)), int(match.group(2)) - - -def _video_files(root: Path, video_key: str) -> list[tuple[int, int, float]]: - """Return (chunk_index, file_index, duration_s) per encoded video file.""" - import av - - files = [] - for path in sorted((root / "videos" / video_key).rglob("*.mp4"), key=_chunk_file): - try: - with av.open(str(path)) as container: - duration = float(container.duration) / av.time_base if container.duration else 0.0 - except Exception as e: - # An unclosed mp4 has no moov atom; treat it as carrying no episodes. - logger.warning("Could not read %s (%s) — episodes in it are unrecoverable", path, e) - duration = 0.0 - chunk_index, file_index = _chunk_file(path) - files.append((chunk_index, file_index, duration)) - return files - - -def _episodes_from_data(root: Path) -> tuple[list[dict[str, Any]], list[Path]]: - """Reconstruct one row per episode from the readable data files. - - Also returns the trailing files that have no parquet footer — the reader - globs every file under ``data/``, so they have to be moved out of the way. - """ - import pandas as pd - import pyarrow.parquet as pq - - # Written with the first saved episode, so its absence means the session - # died before one completed. - if not (root / "meta" / "tasks.parquet").exists(): - return [], [] - - tasks = pd.read_parquet(root / "meta" / "tasks.parquet") - task_by_index = {int(index): task for task, index in tasks["task_index"].items()} - - data_files = sorted((root / "data").rglob("*.parquet"), key=_chunk_file) - rows: list[dict[str, Any]] = [] - next_index = 0 - for position, path in enumerate(data_files): - try: - frames = pq.read_table(path, columns=["episode_index", "task_index"]).to_pandas() - except Exception as e: - # No footer: the writer died mid-file. Nothing from here on is - # readable, since only the newest file is ever left open. - logger.warning("Data file %s is truncated (%s) — recovering the episodes before it", path, e) - return rows, data_files[position:] - - chunk_index, file_index = _chunk_file(path) - for episode_index, group in frames.groupby("episode_index", sort=True): - length = len(group) - rows.append( - { - "episode_index": int(episode_index), - "tasks": [task_by_index[int(i)] for i in sorted(group["task_index"].unique())], - "length": length, - "data/chunk_index": chunk_index, - "data/file_index": file_index, - "dataset_from_index": next_index, - "dataset_to_index": next_index + length, - "meta/episodes/chunk_index": 0, - "meta/episodes/file_index": 0, - } - ) - next_index += length - return rows, [] - - -def _assign_videos(rows: list[dict[str, Any]], root: Path, video_key: str, fps: int) -> int: - """Fill in the video columns for `video_key`, in place. - - Episodes are concatenated into a video file until it hits the size limit, - so each episode's window is its predecessors' durations within that file. - Returns how many leading rows are actually covered by encoded video. - """ - files = _video_files(root, video_key) - cursor = 0 - offset = 0.0 - tolerance = 1.0 / fps - - for covered, row in enumerate(rows): - duration = row["length"] / fps - while cursor < len(files) and offset + duration > files[cursor][2] + tolerance: - cursor += 1 - offset = 0.0 - if cursor >= len(files): - return covered - - chunk_index, file_index, _ = files[cursor] - row[f"videos/{video_key}/chunk_index"] = chunk_index - row[f"videos/{video_key}/file_index"] = file_index - row[f"videos/{video_key}/from_timestamp"] = offset - row[f"videos/{video_key}/to_timestamp"] = offset + duration - offset += duration - - return len(rows) - - -def _trim_data(root: Path, keep_through: int) -> int: - """Drop data rows belonging to episodes past `keep_through`. - - The reader loads every row under data/ and push_to_hub uploads the whole - directory, so rows the rebuilt index no longer references would otherwise - travel with the dataset. Returns the number of frames dropped. - """ - import pyarrow.compute as pc - import pyarrow.parquet as pq - - dropped = 0 - for path in sorted((root / "data").rglob("*.parquet"), key=_chunk_file): - table = pq.read_table(path) - kept = table.filter(pc.less_equal(table["episode_index"], keep_through)) - if kept.num_rows == table.num_rows: - continue - dropped += table.num_rows - kept.num_rows - if kept.num_rows == 0: - path.unlink() - else: - pq.write_table(kept, path, compression="snappy", use_dictionary=True) - return dropped - - -def _episode_video_frames(root: Path, video_key: str, row: dict[str, Any]) -> "np.ndarray": # noqa: F821 - """Decode a sample of the frames one episode references, as (N, C, H, W) in [0, 1].""" - import av - import numpy as np - - from lerobot.datasets.compute_stats import sample_indices - - chunk_index = row[f"videos/{video_key}/chunk_index"] - file_index = row[f"videos/{video_key}/file_index"] - start = row[f"videos/{video_key}/from_timestamp"] - end = row[f"videos/{video_key}/to_timestamp"] - - path = root / "videos" / video_key / f"chunk-{chunk_index:03d}" / f"file-{file_index:03d}.mp4" - with av.open(str(path)) as container: - frames = [ - frame - for frame in container.decode(container.streams.video[0]) - if start <= float(frame.time) < end - ] - - if not frames: - raise DatasetRepairError(f"No decodable frames for {video_key} in episode {row['episode_index']}") - - sampled = [frames[i] for i in sample_indices(len(frames))] - return np.stack([frame.to_ndarray(format="rgb24").transpose(2, 0, 1) for frame in sampled]) / 255.0 - - -def _recompute_stats(root: Path, info: dict[str, Any], rows: list[dict[str, Any]]) -> None: - """Rewrite meta/stats.json over the retained frames only. - - stats.json is aggregated as each episode is saved, so it still describes - episodes that were lost with the interruption, and training normalizes with - it. Rebuilt the way recording builds it, per episode then aggregated, so the - numbers match what the same episodes would have produced on their own. - """ - import numpy as np - import pyarrow.parquet as pq - - from lerobot.datasets.compute_stats import aggregate_stats, get_feature_stats - from lerobot.datasets.io_utils import write_stats - - columns: dict[str, np.ndarray] = {} # noqa: F821 - for path in sorted((root / "data").rglob("*.parquet"), key=_chunk_file): - table = pq.read_table(path) - for key in table.schema.names: - values = np.stack(table[key].to_numpy(zero_copy_only=False).tolist()) - columns[key] = np.concatenate([columns[key], values]) if key in columns else values - - per_episode = [] - for row in rows: - episode_stats = {} - for key, feature in info["features"].items(): - if feature["dtype"] in {"string", "language"}: - continue - - if feature["dtype"] in {"image", "video"}: - values = _episode_video_frames(root, key, row) - episode_stats[key] = { - name: value if name == "count" else np.squeeze(value, axis=0) - for name, value in get_feature_stats(values, axis=(0, 2, 3), keepdims=True).items() - } - elif key in columns: - values = columns[key][row["dataset_from_index"] : row["dataset_to_index"]] - episode_stats[key] = get_feature_stats(values, axis=0, keepdims=values.ndim == 1) - per_episode.append(episode_stats) - - write_stats(aggregate_stats(per_episode), root) - - -def _write_episodes(root: Path, rows: list[dict[str, Any]]) -> None: - import pyarrow as pa - import pyarrow.parquet as pq - - episodes_dir = root / "meta" / "episodes" - shutil.rmtree(episodes_dir, ignore_errors=True) - path = episodes_dir / "chunk-000" / "file-000.parquet" - path.parent.mkdir(parents=True, exist_ok=True) - pq.write_table(pa.Table.from_pylist(rows), path, compression="snappy", use_dictionary=True) - - -def repair_local_dataset(repo_id: str) -> str | None: - """Rebuild `repo_id`'s episode index when the recording never finalized. - - Returns a description of what was recovered, or None when the dataset is - absent or already readable. Raises DatasetRepairError when nothing is left - to recover. - """ - from lerobot.datasets.io_utils import load_episodes - from lerobot.utils.constants import HF_LEROBOT_HOME - - cache_root = Path(HF_LEROBOT_HOME).resolve() - root = (cache_root / repo_id).resolve() - - # repo_id reaches here straight from the request body, and this function - # rewrites info.json and renames files. Keep it inside the cache. - if root == cache_root or cache_root not in root.parents: - raise DatasetRepairError(f"Invalid dataset id: {repo_id}") - - if not (root / "meta" / "info.json").exists(): - return None - - try: - load_episodes(root) - return None - except Exception as e: - logger.info("%s has no usable episode index (%s) — rebuilding it from disk", repo_id, e) - - info = json.loads((root / "meta" / "info.json").read_text()) - fps = info["fps"] - recorded_episodes = info.get("total_episodes", 0) - - rows, truncated = _episodes_from_data(root) - for video_key, feature in info["features"].items(): - if feature["dtype"] == "video": - del rows[_assign_videos(rows, root, video_key, fps) :] - - if not rows: - raise DatasetRepairError( - f"{repo_id} was interrupted before any episode was fully written to disk, " - "and the partial files it left behind cannot be read. Re-record it." - ) - - # The reader loads every parquet file under data/, so an unreadable one - # would still break it. Keep the bytes, just out of the glob's reach. - for path in truncated: - path.rename(path.with_suffix(".parquet.unreadable")) - - lost = recorded_episodes - len(rows) - if lost > 0: - # Episodes the videos don't cover still have their frames in the - # readable shards; drop them so data, index and stats agree. - _trim_data(root, rows[-1]["episode_index"]) - _recompute_stats(root, info, rows) - - _write_episodes(root, rows) - - info["total_episodes"] = len(rows) - info["total_frames"] = sum(row["length"] for row in rows) - info["splits"] = {"train": f"0:{len(rows)}"} - (root / "meta" / "info.json").write_text(json.dumps(info, indent=4)) - - message = f"Recovered {len(rows)} episode(s) from an interrupted recording of {repo_id}" - if lost > 0: - message += f"; {lost} episode(s) were lost because the recording was cut short" - logger.info(message) - return message diff --git a/lelab/datasets.py b/lelab/datasets.py index 82620c5fa..e4b58d5f1 100644 --- a/lelab/datasets.py +++ b/lelab/datasets.py @@ -27,7 +27,6 @@ from lerobot.datasets.dataset_tools import merge_datasets from . import episode_media -from .dataset_repair import DatasetRepairError, repair_local_dataset from .utils.hf_auth import cached_whoami, shared_hf_api logger = logging.getLogger(__name__) @@ -203,13 +202,10 @@ def handle_merge_local_datasets(source_repo_ids: list[str], output_name: str) -> if output_dir.exists(): raise ValueError(f"A local dataset already exists at {output_repo_id}") - # An interrupted local recording can have data and info.json but no episode - # index. LeRobot then treats it as absent and attempts a Hub lookup, so make - # each source readable before constructing a dataset from its local path. try: for repo_id in source_repo_ids: - repair_local_dataset(repo_id) - except DatasetRepairError as exc: + episode_media.ensure_episode_index(repo_id) + except episode_media.UnreadableDatasetError as exc: raise ValueError(str(exc)) from exc source_dirs = [episode_media.resolve_dataset_dir(repo_id) for repo_id in source_repo_ids] diff --git a/lelab/episode_media.py b/lelab/episode_media.py index e94abd158..d33d96daf 100644 --- a/lelab/episode_media.py +++ b/lelab/episode_media.py @@ -59,6 +59,10 @@ class EpisodeNotFoundError(LookupError): """The requested episode/camera/frame combination doesn't exist on disk.""" +class UnreadableDatasetError(Exception): + """A local dataset LeRobot cannot open: it has no readable episode index.""" + + # ── dataset resolution ────────────────────────────────────────────────────── @@ -81,8 +85,8 @@ def lerobot_cache_root() -> Path: return (Path(hf_home).expanduser() / "lerobot").resolve() -def resolve_dataset_dir(repo_id: str) -> Path: - """Resolve ``repo_id`` to a local dataset directory. +def _cache_path(repo_id: str) -> Path: + """Resolve ``repo_id`` inside the cache root, whether or not it exists. Rejects path traversal the same way ``handle_delete_dataset`` does: the resolved target must stay strictly inside the cache root. ``repo_id`` comes @@ -96,6 +100,12 @@ def resolve_dataset_dir(repo_id: str) -> Path: raise DatasetNotFoundError(f"Invalid dataset path: {repo_id}") from e if target == root or root not in target.parents: raise DatasetNotFoundError(f"Invalid dataset path: {repo_id}") + return target + + +def resolve_dataset_dir(repo_id: str) -> Path: + """Resolve ``repo_id`` to an existing local dataset directory.""" + target = _cache_path(repo_id) if not (target / "meta" / "info.json").is_file(): raise DatasetNotFoundError(f"Dataset not found locally: {repo_id}") return target @@ -216,6 +226,26 @@ def read_episode_index(dataset_dir: Path) -> list[EpisodeIndexRow]: return rows +def ensure_episode_index(repo_id: str) -> None: + """Raise when a local dataset has no episode index LeRobot can read. + + LeRobot writes that index in ``finalize()``, so a recording killed before + it (power cut, kill -9) has none. LeRobot then takes the dataset for one not + downloaded yet and asks the Hub, which 404s for a dataset never pushed. + """ + # Checked before the Hub fallback, so an id escaping the cache never + # reaches LeRobot, whether or not something exists at its target. + dataset_dir = _cache_path(repo_id) + # Not in the local cache: LeRobot fetches it from the Hub. + if not (dataset_dir / "meta" / "info.json").is_file(): + return + if not read_episode_index(dataset_dir): + raise UnreadableDatasetError( + f"{repo_id} has no readable episodes, most likely because its recording was " + "interrupted. Delete it and record it again." + ) + + # ── video location ────────────────────────────────────────────────────────── diff --git a/lelab/record.py b/lelab/record.py index 766e2f86b..0ad2ca92e 100644 --- a/lelab/record.py +++ b/lelab/record.py @@ -31,7 +31,7 @@ from lerobot.scripts.lerobot_record import RecordConfig from lerobot.teleoperators.so_leader import SO101LeaderConfig -from .dataset_repair import DatasetRepairError, repair_local_dataset +from .episode_media import UnreadableDatasetError, ensure_episode_index from .utils.config import setup_calibration_files, with_lelab_tag from .utils.devices import safe_disconnect_device @@ -277,7 +277,7 @@ def handle_start_recording(request: RecordingRequest) -> dict[str, Any]: if not request.resume and request.dataset_repo_id: request.dataset_repo_id = f"{request.dataset_repo_id}_{datetime.now().strftime('%Y%m%d_%H%M%S')}" if request.resume: - repair_local_dataset(request.dataset_repo_id) + ensure_episode_index(request.dataset_repo_id) logger.info(f"Starting recording for dataset: {request.dataset_repo_id}") logger.info(f"Task: {request.single_task}") @@ -405,6 +405,27 @@ def handle_stop_recording() -> dict[str, Any]: } +def stop_recording_and_wait() -> None: + """Stop the active session and wait for it to finalize its dataset. + + The worker is a daemon thread, so a server exiting mid-session kills it + before finalize(): the data file and episode index stay unreadable and the + whole session is lost. + """ + thread = recording_thread + if thread is None or not thread.is_alive(): + return + handle_stop_recording() + # Covers finalize() encoding any videos still queued. + timeout_s = 60 + logger.info("Saving the recording before exiting (up to %ds)...", timeout_s) + thread.join(timeout_s) + if thread.is_alive(): + logger.error( + "Recording did not finish saving in time; %s may be unreadable", recording_config.dataset_repo_id + ) + + def handle_exit_early() -> dict[str, Any]: """Handle exit early request - replaces right arrow key""" if not recording_active or recording_events is None: @@ -575,7 +596,7 @@ def handle_get_dataset_info(request: DatasetInfoRequest) -> dict[str, Any]: try: from lerobot.datasets import LeRobotDataset - repair_local_dataset(request.dataset_repo_id) + ensure_episode_index(request.dataset_repo_id) dataset = LeRobotDataset(request.dataset_repo_id) return { @@ -588,8 +609,8 @@ def handle_get_dataset_info(request: DatasetInfoRequest) -> dict[str, Any]: "total_frames": dataset.num_frames, "robot_type": getattr(dataset.meta, "robot_type", "Unknown robot"), } - except DatasetRepairError as e: - logger.warning(f"Could not repair local dataset {request.dataset_repo_id}: {e}") + except UnreadableDatasetError as e: + logger.warning(str(e)) return {"success": False, "message": str(e)} except Exception as e: @@ -691,7 +712,7 @@ def handle_upload_dataset(request: UploadRequest) -> dict[str, Any]: # Import LeRobotDataset to load and upload the dataset from lerobot.datasets import LeRobotDataset - repair_local_dataset(request.dataset_repo_id) + ensure_episode_index(request.dataset_repo_id) logger.info(f"Loading dataset {request.dataset_repo_id} for upload") @@ -714,7 +735,7 @@ def handle_upload_dataset(request: UploadRequest) -> dict[str, Any]: "num_episodes": dataset.num_episodes, } - except DatasetRepairError as e: + except UnreadableDatasetError as e: logger.error(f"Cannot upload {request.dataset_repo_id}: {e}") return {"success": False, "message": str(e)} @@ -885,7 +906,9 @@ def record_with_web_events(cfg: RecordConfig, web_events: dict) -> LeRobotDatase # devices come up, while the user is still moving into position. settled_in = False - while saved_episodes < cfg.dataset.num_episodes: + # Checked up front too: each phase clears exit_early as it starts, so a + # Stop sent during device setup would otherwise sit out a full reset. + while saved_episodes < cfg.dataset.num_episodes and not web_events["stop_recording"]: if not settled_in and cfg.dataset.reset_time_s > 0: settled_in = True current_phase = "resetting" @@ -953,6 +976,11 @@ def record_with_web_events(cfg: RecordConfig, web_events: dict) -> LeRobotDatase web_events["exit_early"] = False dataset.clear_episode_buffer() + # Stop mid-episode lands here too: drop the take, but there is + # no next episode to reset for. + if web_events["stop_recording"]: + break + # Go through reset phase before re-recording (don't increment episode counters) # RESET PHASE - without dataset (matches original record.py exactly) current_phase = "resetting" diff --git a/lelab/scripts/lelab.py b/lelab/scripts/lelab.py index 045bccbcd..48dca7802 100644 --- a/lelab/scripts/lelab.py +++ b/lelab/scripts/lelab.py @@ -362,12 +362,20 @@ def _run_prod(*, no_open: bool = False, rebuild: bool = False) -> None: def _shutdown(_signum, _frame) -> None: logger.info("Shutting down LeLab...") try: - for child in psutil.Process().children(recursive=True): - with contextlib.suppress(psutil.NoSuchProcess): - child.terminate() - except Exception: - pass - os._exit(0) + # os._exit skips uvicorn's shutdown hook, which is what saves + # an active recording. Imported here: Ctrl+C can land while the + # server is still loading. + from lelab.record import stop_recording_and_wait + + stop_recording_and_wait() + finally: + try: + for child in psutil.Process().children(recursive=True): + with contextlib.suppress(psutil.NoSuchProcess): + child.terminate() + except Exception: + pass + os._exit(0) signal.signal(signal.SIGINT, _shutdown) for _name in ("SIGTERM", "SIGBREAK"): @@ -417,6 +425,10 @@ def _run_dev(*, no_open: bool = False) -> None: "--port", str(BACKEND_PORT), "--reload", + # Same as prod: an open /camera-feed stream would otherwise + # hold off the shutdown hook that saves an active recording. + "--timeout-graceful-shutdown", + "2", ], PROJECT_ROOT, env=os.environ.copy(), diff --git a/lelab/server.py b/lelab/server.py index 416e35351..d5fb7372c 100644 --- a/lelab/server.py +++ b/lelab/server.py @@ -1386,7 +1386,7 @@ async def shutdown_event(): """Clean up resources when FastAPI shuts down""" logger.info("🔄 FastAPI shutting down, cleaning up...") - # Stop any active recording - handled by recording module cleanup + _record.stop_recording_and_wait() if manager: manager.stop_broadcast_thread() diff --git a/tests/test_dataset_repair.py b/tests/test_dataset_repair.py deleted file mode 100644 index eef509811..000000000 --- a/tests/test_dataset_repair.py +++ /dev/null @@ -1,250 +0,0 @@ -# Copyright 2025 The HuggingFace Inc. team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Tests for lelab.dataset_repair — rebuilding the episode index of a -recording that was interrupted before `LeRobotDataset.finalize()` ran.""" - -from __future__ import annotations - -import json -import shutil -from pathlib import Path - -import numpy as np -import pytest - -REPO_ID = "repair_test/session" -EPISODES = 2 -FRAMES = 5 -FPS = 10 - - -def _record(repo_id: str, *, video: bool, episodes: int = EPISODES) -> Path: - """Write a small finished dataset and return its root.""" - from lerobot.datasets import LeRobotDataset - - features = { - "action": {"dtype": "float32", "shape": (2,), "names": ["a", "b"]}, - "observation.state": {"dtype": "float32", "shape": (2,), "names": ["a", "b"]}, - } - if video: - features["observation.images.cam"] = { - "dtype": "video", - "shape": (32, 32, 3), - "names": ["h", "w", "c"], - } - - dataset = LeRobotDataset.create(repo_id, fps=FPS, features=features, use_videos=video) - for episode in range(episodes): - for frame in range(FRAMES): - values = np.full(2, episode * FRAMES + frame, dtype=np.float32) - item = {"action": values, "observation.state": values, "task": "repair me"} - if video: - item["observation.images.cam"] = np.full((32, 32, 3), 100, dtype=np.uint8) - dataset.add_frame(item) - dataset.save_episode() - dataset.finalize() - return dataset.root - - -def test_chunk_file_parses_indices_from_path() -> None: - from lelab.dataset_repair import _chunk_file - - assert _chunk_file(Path("data/chunk-003/file-012.parquet")) == (3, 12) - - -def test_chunk_file_rejects_unexpected_layout() -> None: - from lelab.dataset_repair import DatasetRepairError, _chunk_file - - with pytest.raises(DatasetRepairError): - _chunk_file(Path("data/whatever.parquet")) - - -def test_absent_dataset_is_left_to_the_normal_hub_path(tmp_lerobot_home: Path) -> None: - from lelab.dataset_repair import repair_local_dataset - - assert repair_local_dataset("someone/never-recorded") is None - - -def test_readable_dataset_is_not_touched(tmp_lerobot_home: Path) -> None: - from lelab.dataset_repair import repair_local_dataset - - _record(REPO_ID, video=False) - - assert repair_local_dataset(REPO_ID) is None - - -@pytest.mark.parametrize("repo_id", ["../escape", "repair_test/../../escape", "/etc"]) -def test_repo_id_cannot_escape_the_cache_directory(tmp_lerobot_home: Path, repo_id: str) -> None: - from lelab.dataset_repair import DatasetRepairError, repair_local_dataset - - with pytest.raises(DatasetRepairError, match="Invalid dataset id"): - repair_local_dataset(repo_id) - - -@pytest.mark.parametrize("video", [False, True]) -def test_missing_episode_index_is_rebuilt_from_the_data_files(tmp_lerobot_home: Path, video: bool) -> None: - from lelab.dataset_repair import repair_local_dataset - from lerobot.datasets import LeRobotDataset - - root = _record(REPO_ID, video=video) - # What an un-finalized recording leaves behind: frames and videos on disk, - # no episode index. Loading it in this state is what reaches for the Hub - # and 404s, so the test never does it — it goes straight to the repair. - shutil.rmtree(root / "meta" / "episodes") - - message = repair_local_dataset(REPO_ID) - assert message is not None - assert f"Recovered {EPISODES} episode(s)" in message - - repaired = LeRobotDataset(REPO_ID) - assert repaired.num_episodes == EPISODES - assert repaired.num_frames == EPISODES * FRAMES - - last = repaired[repaired.num_frames - 1] - assert int(last["episode_index"]) == EPISODES - 1 - assert last["task"] == "repair me" - np.testing.assert_allclose(last["action"].numpy(), np.full(2, EPISODES * FRAMES - 1)) - if video: - assert tuple(last["observation.images.cam"].shape) == (3, 32, 32) - - -def test_repair_is_idempotent(tmp_lerobot_home: Path) -> None: - from lelab.dataset_repair import repair_local_dataset - - root = _record(REPO_ID, video=False) - shutil.rmtree(root / "meta" / "episodes") - - assert repair_local_dataset(REPO_ID) is not None - assert repair_local_dataset(REPO_ID) is None - - -def test_truncated_data_file_is_moved_out_of_the_readers_way(tmp_lerobot_home: Path) -> None: - from lelab.dataset_repair import repair_local_dataset - from lerobot.datasets import LeRobotDataset - - root = _record(REPO_ID, video=False) - shutil.rmtree(root / "meta" / "episodes") - - # A file the writer was still appending to when the process died: bytes - # present, parquet footer missing. - orphan = root / "data" / "chunk-000" / "file-001.parquet" - orphan.write_bytes(b"PAR1\x00\x01\x02 no footer here") - - assert repair_local_dataset(REPO_ID) is not None - assert not orphan.exists() - assert orphan.with_suffix(".parquet.unreadable").exists() - assert LeRobotDataset(REPO_ID).num_episodes == EPISODES - - -def test_episodes_without_video_lose_their_frames_and_stats_too(tmp_lerobot_home: Path) -> None: - """An episode the videos don't cover is dropped from the index, so its rows - and its contribution to stats.json have to go with it.""" - from lelab.dataset_repair import repair_local_dataset - from lerobot.datasets import LeRobotDataset - - broken = _record(REPO_ID, video=True, episodes=2) - # Same recording, one episode shorter: stands in both for the video that - # was never flushed and for the statistics the repair should end up with. - reference = _record("repair_test/reference", video=True, episodes=1) - - shutil.rmtree(broken / "meta" / "episodes") - shutil.copy( - next((reference / "videos").rglob("*.mp4")), - next((broken / "videos").rglob("*.mp4")), - ) - - message = repair_local_dataset(REPO_ID) - assert message is not None - assert "Recovered 1 episode(s)" in message - - repaired = LeRobotDataset(REPO_ID) - assert repaired.num_episodes == 1 - assert repaired.num_frames == FRAMES - # The dropped episode's rows are gone from the shards, not just unreferenced. - assert len(repaired.hf_dataset) == FRAMES - - repaired_stats = json.loads((broken / "meta" / "stats.json").read_text()) - reference_stats = json.loads((reference / "meta" / "stats.json").read_text()) - assert repaired_stats["action"]["max"] == reference_stats["action"]["max"] - assert repaired_stats["action"]["mean"] == reference_stats["action"]["mean"] - - -def test_unrecoverable_dataset_reports_instead_of_hitting_the_hub(tmp_lerobot_home: Path) -> None: - from lelab.dataset_repair import DatasetRepairError, repair_local_dataset - - root = _record(REPO_ID, video=False) - shutil.rmtree(root / "meta" / "episodes") - for path in (root / "data").rglob("*.parquet"): - path.write_bytes(b"PAR1 truncated") - - with pytest.raises(DatasetRepairError, match="Re-record it"): - repair_local_dataset(REPO_ID) - - -def test_repaired_image_stats_use_the_episode_video_chunk(tmp_lerobot_home: Path) -> None: - """Repeated file names across chunks must not change retained image statistics.""" - from lelab.dataset_repair import repair_local_dataset - from lerobot.datasets import LeRobotDataset - - video_key = "observation.images.cam" - features = { - "action": {"dtype": "float32", "shape": (2,), "names": ["a", "b"]}, - "observation.state": {"dtype": "float32", "shape": (2,), "names": ["a", "b"]}, - video_key: {"dtype": "video", "shape": (32, 32, 3), "names": ["h", "w", "c"]}, - } - - def record_chunks(repo_id: str, episodes: int) -> Path: - dataset = LeRobotDataset.create(repo_id, fps=FPS, features=features, use_videos=True) - # Exercise the real writer's rollover without a large recording. - dataset.meta.info["chunks_size"] = 1 - dataset.meta.info["video_files_size_in_mb"] = 0.001 - for episode in range(episodes): - for frame in range(FRAMES): - values = np.full(2, episode * FRAMES + frame, dtype=np.float32) - dataset.add_frame( - { - "action": values, - "observation.state": values, - "task": "repair me", - video_key: np.full((32, 32, 3), 30 + 90 * episode, dtype=np.uint8), - } - ) - dataset.save_episode() - dataset.finalize() - return dataset.root - - broken = record_chunks(REPO_ID, 3) - reference = record_chunks("repair_test/chunk_reference", 2) - video_paths = sorted((broken / "videos" / video_key).rglob("*.mp4")) - assert [path.parent.name for path in video_paths] == ["chunk-000", "chunk-001", "chunk-002"] - assert all(path.name == "file-000.mp4" for path in video_paths) - shutil.rmtree(broken / "meta" / "episodes") - video_paths[-1].unlink() - - message = repair_local_dataset(REPO_ID) - assert message is not None - assert "Recovered 2 episode(s)" in message - repaired = LeRobotDataset(REPO_ID, video_backend="pyav") - assert repaired.num_episodes == 2 - assert repaired.num_frames == 2 * FRAMES - assert float(repaired[0][video_key].mean()) < float(repaired[FRAMES][video_key].mean()) - - repaired_stats = json.loads((broken / "meta" / "stats.json").read_text())[video_key] - reference_stats = json.loads((reference / "meta" / "stats.json").read_text())[video_key] - assert repaired_stats.keys() == reference_stats.keys() - # Repair decodes lossy AV1 video; recording statistics use the source pixels. - for name in reference_stats: - if name != "count": - np.testing.assert_allclose(repaired_stats[name], reference_stats[name], atol=0.01) - np.testing.assert_array_equal(repaired_stats["count"], reference_stats["count"]) diff --git a/tests/test_datasets.py b/tests/test_datasets.py index 24b15b678..da9dd6ce5 100644 --- a/tests/test_datasets.py +++ b/tests/test_datasets.py @@ -140,10 +140,8 @@ def test_merge_local_datasets_creates_new_dataset_without_touching_sources( ) -> None: from lelab import datasets as datasets_mod - from .test_dataset_repair import _record - - _record("alice/first", video=False, episodes=1) - _record("bob/second", video=False, episodes=1) + _write_dataset(tmp_lerobot_home, "alice/first", episodes=[(0, 4)], cameras=()) + _write_dataset(tmp_lerobot_home, "bob/second", episodes=[(0, 4)], cameras=()) calls = [] class FakeDataset: @@ -174,59 +172,16 @@ def fake_merge(datasets, output_repo_id: str, output_dir: Path): assert (tmp_lerobot_home / "bob" / "second" / "meta" / "info.json").is_file() -def test_merge_local_datasets_repairs_an_interrupted_source_before_opening( - tmp_lerobot_home: Path, monkeypatch: pytest.MonkeyPatch -) -> None: - """A recorded-but-unfinalized source must stay local when it is merged.""" - from lelab import datasets as datasets_mod - - from .test_dataset_repair import _record - - first = _record("alice/first", video=False, episodes=1) - _record("bob/second", video=False, episodes=1) - # This is the on-disk state an interrupted recording leaves behind. With - # LeRobot v0.6.0, opening it before repair would fall through to the Hub. - import shutil - - shutil.rmtree(first / "meta" / "episodes") - - class FakeDataset: - def __init__(self, repo_id: str, root: Path) -> None: - assert (root / "meta" / "episodes").is_dir() - self.repo_id = repo_id - - class FakeMergedDataset: - num_episodes = 4 - num_frames = 4 - - def fake_merge(datasets, output_repo_id: str, output_dir: Path): - (output_dir / "meta").mkdir(parents=True) - (output_dir / "meta" / "info.json").write_text("{}") - return FakeMergedDataset() - - monkeypatch.setattr(datasets_mod, "LeRobotDataset", FakeDataset) - monkeypatch.setattr(datasets_mod, "merge_datasets", fake_merge) - - result = datasets_mod.handle_merge_local_datasets(["alice/first", "bob/second"], "combined") - - assert result["success"] is True - assert (first / "meta" / "episodes").is_dir() - - -def test_merge_local_datasets_reports_unrepairable_source( - tmp_lerobot_home: Path, monkeypatch: pytest.MonkeyPatch -) -> None: - from lelab import datasets as datasets_mod - from lelab.dataset_repair import DatasetRepairError +def test_merge_local_datasets_reports_unreadable_source(tmp_lerobot_home: Path) -> None: + """A recording interrupted before finalize has no episode index; LeRobot + would take it for a Hub dataset and 404.""" + from lelab.datasets import handle_merge_local_datasets - monkeypatch.setattr( - datasets_mod, - "repair_local_dataset", - lambda _repo_id: (_ for _ in ()).throw(DatasetRepairError("Re-record it")), - ) + _write_dataset(tmp_lerobot_home, "alice/first", episodes=[(0, 4)], cameras=()) + _make_dataset(tmp_lerobot_home, "bob/second") - with pytest.raises(ValueError, match="Re-record it"): - datasets_mod.handle_merge_local_datasets(["alice/first", "bob/second"], "combined") + with pytest.raises(ValueError, match="bob/second has no readable episodes"): + handle_merge_local_datasets(["alice/first", "bob/second"], "combined") @pytest.mark.parametrize( @@ -259,10 +214,8 @@ def test_merge_local_datasets_cleans_staging_output_on_failure( ) -> None: from lelab import datasets as datasets_mod - from .test_dataset_repair import _record - - _record("alice/first", video=False, episodes=1) - _record("bob/second", video=False, episodes=1) + _write_dataset(tmp_lerobot_home, "alice/first", episodes=[(0, 4)], cameras=()) + _write_dataset(tmp_lerobot_home, "bob/second", episodes=[(0, 4)], cameras=()) monkeypatch.setattr(datasets_mod, "LeRobotDataset", lambda repo_id, root: object()) def failing_merge(datasets, output_repo_id: str, output_dir: Path): diff --git a/tests/test_episode_media.py b/tests/test_episode_media.py index c0799053b..12f4f97a7 100644 --- a/tests/test_episode_media.py +++ b/tests/test_episode_media.py @@ -201,6 +201,55 @@ def test_read_episode_index_empty_when_meta_missing(tmp_lerobot_home: Path) -> N assert read_episode_index(d) == [] +def test_ensure_episode_index_accepts_finalized_dataset(dataset: str) -> None: + from lelab.episode_media import ensure_episode_index + + ensure_episode_index(dataset) + + +def test_ensure_episode_index_leaves_absent_dataset_to_lerobot(tmp_lerobot_home: Path) -> None: + from lelab.episode_media import ensure_episode_index + + ensure_episode_index("acme/on-the-hub") + + +@pytest.mark.parametrize( + "repo_id", ["../outside/ds", "../outside/absent", "{outside}/outside/ds", ".", "acme/link"] +) +def test_ensure_episode_index_refuses_ids_escaping_the_cache(tmp_lerobot_home: Path, repo_id: str) -> None: + """Even when nothing exists there, an escaping id must not reach LeRobot.""" + from lelab.episode_media import DatasetNotFoundError, ensure_episode_index + + outside = tmp_lerobot_home.parent + _write_dataset(outside, "outside/ds", episodes=[(0, 4)], cameras=()) + (tmp_lerobot_home / "acme").mkdir() + (tmp_lerobot_home / "acme" / "link").symlink_to(outside / "elsewhere") + + with pytest.raises(DatasetNotFoundError, match="Invalid dataset path"): + ensure_episode_index(repo_id.format(outside=outside)) + + +@pytest.mark.parametrize("damage", ["missing", "no_footer"]) +def test_ensure_episode_index_rejects_unfinalized_recording( + tmp_lerobot_home: Path, dataset: str, damage: str +) -> None: + """A recording killed before finalize has no index, or one without its + parquet footer. LeRobot would take either for a Hub dataset and 404.""" + import shutil + + from lelab.episode_media import UnreadableDatasetError, ensure_episode_index + + episodes_dir = tmp_lerobot_home / dataset / "meta" / "episodes" + if damage == "missing": + shutil.rmtree(episodes_dir) + else: + index = episodes_dir / "chunk-000" / "file-000.parquet" + index.write_bytes(index.read_bytes()[:-8]) + + with pytest.raises(UnreadableDatasetError, match="record it again"): + ensure_episode_index(dataset) + + # ── video location ────────────────────────────────────────────────────────── diff --git a/tests/test_record.py b/tests/test_record.py index c1293a1dc..1c4aef18b 100644 --- a/tests/test_record.py +++ b/tests/test_record.py @@ -60,6 +60,42 @@ def test_handle_stop_recording_when_idle_returns_dict(tmp_lerobot_home) -> None: assert isinstance(result, dict) +def test_stop_recording_and_wait_when_idle_returns() -> None: + from lelab.record import stop_recording_and_wait + + stop_recording_and_wait() + + +def test_stop_recording_and_wait_lets_the_session_finalize(monkeypatch: pytest.MonkeyPatch) -> None: + """Server shutdown must end the session through its teardown, not kill it.""" + import threading + import time + from types import SimpleNamespace + + from lelab import record + + events = {"stop_recording": False, "exit_early": False} + finalized = threading.Event() + + def session() -> None: + while not events["stop_recording"]: + time.sleep(0.01) + finalized.set() + + thread = threading.Thread(target=session, daemon=True) + monkeypatch.setattr(record, "recording_active", True) + monkeypatch.setattr(record, "recording_events", events) + monkeypatch.setattr(record, "recording_thread", thread) + monkeypatch.setattr(record, "recording_config", SimpleNamespace(dataset_repo_id="local/x")) + monkeypatch.setattr(record, "current_phase", "recording") + monkeypatch.setattr(record, "phase_start_time", None) + thread.start() + + record.stop_recording_and_wait() + + assert finalized.is_set() + + def test_resolve_dataset_dir_rejects_traversal(tmp_lerobot_home) -> None: from lelab.record import _resolve_dataset_dir @@ -300,6 +336,29 @@ def test_build_camera_configs_skips_non_opencv_type() -> None: assert configs == {} +EPISODES = 2 + + +def _record(repo_id: str): + """Write a small finished LeRobot dataset and return its root.""" + import numpy as np + + from lerobot.datasets import LeRobotDataset + + features = { + "action": {"dtype": "float32", "shape": (2,), "names": ["a", "b"]}, + "observation.state": {"dtype": "float32", "shape": (2,), "names": ["a", "b"]}, + } + dataset = LeRobotDataset.create(repo_id, fps=10, features=features, use_videos=False) + for _ in range(EPISODES): + for frame in range(5): + values = np.full(2, frame, dtype=np.float32) + dataset.add_frame({"action": values, "observation.state": values, "task": "pick"}) + dataset.save_episode() + dataset.finalize() + return dataset.root + + def _resume_request(repo_id: str): from lelab.record import RecordingRequest @@ -319,9 +378,8 @@ def test_resume_reopens_the_dataset_in_the_local_cache( ) -> None: import lelab.record as record from lerobot.datasets import LeRobotDataset - from tests.test_dataset_repair import EPISODES, _record - _record("local/ds", video=False) + _record("local/ds") monkeypatch.setattr(record, "setup_calibration_files", lambda leader, follower: ("leader", "follower")) config = record.create_record_config(_resume_request("local/ds")) @@ -329,13 +387,12 @@ def test_resume_reopens_the_dataset_in_the_local_cache( assert dataset.meta.total_episodes == EPISODES -def test_resume_refuses_an_unrecoverable_dataset_before_the_robot( +def test_resume_refuses_an_unreadable_dataset_before_the_robot( tmp_lerobot_home, monkeypatch: pytest.MonkeyPatch ) -> None: import lelab.record as record - from tests.test_dataset_repair import _record - root = _record("local/ds", video=False) + root = _record("local/ds") shutil.rmtree(root / "meta" / "episodes") for path in (root / "data").rglob("*.parquet"): path.write_bytes(b"PAR1 truncated") @@ -345,5 +402,5 @@ def test_resume_refuses_an_unrecoverable_dataset_before_the_robot( result = record.handle_start_recording(_resume_request("local/ds")) assert result["success"] is False - assert "Re-record it" in result["message"] + assert "record it again" in result["message"] assert record.recording_active is False diff --git a/tests/test_server.py b/tests/test_server.py index b9361237f..60b097bde 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -54,6 +54,18 @@ def test_app_exposes_required_endpoints() -> None: assert not missing, f"missing routes: {missing}" +def test_shutdown_lets_an_active_recording_finalize(monkeypatch: pytest.MonkeyPatch) -> None: + """Without this the daemon recording thread dies unfinalized on Ctrl+C.""" + from lelab import record + from lelab.server import app + + stop = MagicMock() + monkeypatch.setattr(record, "stop_recording_and_wait", stop) + with TestClient(app): + stop.assert_not_called() + stop.assert_called_once_with() + + def test_health_endpoint_returns_200(client: TestClient) -> None: response = client.get("/health") assert response.status_code == 200