From 133d3325057b35a0b54034aab0ba70f2dad51ff5 Mon Sep 17 00:00:00 2001 From: Akshat Patel Date: Thu, 3 Sep 2026 20:16:06 -0400 Subject: [PATCH] fix(lerobot): read episode metadata from every v3 shard (fixes #293) Dataset v3 splits meta/episodes into chunk-XXX/file-YYY.parquet shards by size. The importer cached each shard under its basename, so a second chunk directory reusing file-000.parquet looked already downloaded and its episodes vanished while the first shard was read twice. It also discovered video columns and video windows from the first shard only, so every episode in a later shard converted without its window and fell back to the whole chunk video. Keep the tree path below meta/episodes in the cache, refusing entries that would escape it, and read episode rows, video columns, and video windows from one relation over every shard. Bump CONVERTER_VERSION to v5 so outputs prepared under the old behavior do not share an identity with corrected ones. Measured on lerobot/droid_1.0.1 (v3.0, 95,658 episodes, seven metadata shards) with its first two shards: 14,152 of 28,305 episodes had no video window before, 0 after. --- docs/how-to/import-lerobot-v3.md | 6 +- src/hflow/importers/lerobot.py | 110 ++++++++++++++++++++----------- tests/test_lerobot_converter.py | 99 ++++++++++++++++++++++++++++ 3 files changed, 176 insertions(+), 39 deletions(-) diff --git a/docs/how-to/import-lerobot-v3.md b/docs/how-to/import-lerobot-v3.md index 7e965140..01617c04 100644 --- a/docs/how-to/import-lerobot-v3.md +++ b/docs/how-to/import-lerobot-v3.md @@ -84,8 +84,10 @@ channels named `/observation.state` and `/action`. | Language tensors, depth maps, or arbitrary nested features | Not supported | The importer reads `meta/info.json`, the episode index, Parquet feature -columns, frame rate, episode boundaries, and video-path templates. It refuses -an unsupported or missing required feature before publishing an episode. +columns, frame rate, episode boundaries, and video-path templates. Episode +metadata split across several `meta/episodes` shards is read in full, so +every episode keeps its own video window. It refuses an unsupported or +missing required feature before publishing an episode. Every output records the source repository, resolved commit, source episode index, task, embodiment, importer version, and FFmpeg build. Video is encoded diff --git a/src/hflow/importers/lerobot.py b/src/hflow/importers/lerobot.py index a85bf590..8ab25710 100644 --- a/src/hflow/importers/lerobot.py +++ b/src/hflow/importers/lerobot.py @@ -25,7 +25,7 @@ import urllib.request from collections.abc import Sequence from dataclasses import dataclass -from pathlib import Path +from pathlib import Path, PurePosixPath from typing import NotRequired, TypedDict from urllib.parse import urlsplit @@ -42,8 +42,14 @@ DEFAULT_OUTPUT_DIR = Path("./data/lerobot_pusht") DEFAULT_CAMERA_KEY = "observation.image" -CONVERTER_VERSION = "lerobot-converter-v4" +# "v5": episode metadata is read from every meta/episodes shard, not only the +# first, and shards are cached under their own chunk directory (#293). +# Multi-shard corpora previously published episodes with the wrong video +# window or the wrong source episode, so their outputs must not share an +# identity with the corrected ones. +CONVERTER_VERSION = "lerobot-converter-v5" PRESENTATION_TIMESTAMP_EPSILON_S = 0.050 +EPISODE_METADATA_TREE_PREFIX = PurePosixPath("meta/episodes") # Timestamp handling NANOSECONDS_PER_SECOND = 1_000_000_000 @@ -371,6 +377,32 @@ def _download_file(url: str, destination_path: Path, chunk_size: int = 1 << 20) raise +def _episode_metadata_cache_path( + episodes_metadata_directory: Path, tree_entry_path: str, *, repo_id: str +) -> Path: + """Where one ``meta/episodes`` tree entry lands in the local cache. + + The path below ``meta/episodes`` is kept rather than flattened to its + basename. Dataset v3 shards episode metadata as + ``chunk-XXX/file-YYY.parquet`` and reuses file names across chunk + directories, so two shards flattened to one cache file would make the + second look already downloaded and its episodes vanish (#293). The tree + listing is remote input, so an entry that would land outside the + metadata directory is refused rather than joined. + """ + tree_path = PurePosixPath(tree_entry_path) + try: + relative_path = tree_path.relative_to(EPISODE_METADATA_TREE_PREFIX) + except ValueError: + relative_path = None + if relative_path is None or not relative_path.parts or ".." in relative_path.parts: + raise ValueError( + f"Hugging Face tree response for {repo_id} lists {tree_entry_path!r}, which is " + f"not a file below {EPISODE_METADATA_TREE_PREFIX}/" + ) + return episodes_metadata_directory / relative_path + + def _ensure_source_archive(dataset_source: DatasetSource, cache_dir: Path) -> _SourceArchive: """Download the corpus parquets and video chunks needed for the given episodes.""" import duckdb @@ -410,7 +442,9 @@ def _ensure_source_archive(dataset_source: DatasetSource, cache_dir: Path) -> _S entries = _hf_tree(dataset_source.repo_id, dataset_source.revision, "meta/episodes") for entry in entries: if entry.get("type") == "file" and entry["path"].endswith(".parquet"): - destination_path = episodes_metadata_directory / Path(entry["path"]).name + destination_path = _episode_metadata_cache_path( + episodes_metadata_directory, entry["path"], repo_id=dataset_source.repo_id + ) if not destination_path.exists(): _download_file(f"{dataset_base_url}/{entry['path']}", destination_path) episode_metadata_files.append(destination_path) @@ -418,46 +452,49 @@ def _ensure_source_archive(dataset_source: DatasetSource, cache_dir: Path) -> _S if not episode_metadata_files: raise RuntimeError("no meta/episodes parquet files found") - # Index of per-episode data windows across chunks + # Index of per-episode data windows across chunks. Every shard is one + # relation: v3 splits episode metadata by size, so an episode and its + # video window can sit in any file, not only the first (#293). + episode_metadata_relation = ( + "read_parquet([" + + ", ".join( + "'" + str(episode_metadata_file).replace("'", "''") + "'" + for episode_metadata_file in episode_metadata_files + ) + + "], union_by_name=true)" + ) connection = duckdb.connect() try: episode_rows: list[_EpisodeRow] = [] - for episode_metadata_file in episode_metadata_files: - parquet_episode_rows = connection.execute( - f""" - SELECT "episode_index", "tasks", "length", - "data/chunk_index", "data/file_index", - "dataset_from_index", "dataset_to_index" - FROM read_parquet('{str(episode_metadata_file).replace("'", "''")}') - ORDER BY "episode_index" - """ - ).fetchall() - for parquet_episode_row in parquet_episode_rows: - tasks = parquet_episode_row[1] - if isinstance(tasks, list): - task = str(tasks[0]) if tasks else "" - else: - task = str(tasks or "") - episode_rows.append( - { - "episode_index": int(parquet_episode_row[0]), - "task": task, - "length": int(parquet_episode_row[2]), - "data_chunk": str(parquet_episode_row[3]).split("/")[-1], - "data_file": str(parquet_episode_row[4]).split("/")[-1], - "data_from": int(parquet_episode_row[5]), - "data_to": int(parquet_episode_row[6]), - } - ) - episode_rows.sort(key=lambda episode: episode["episode_index"]) + parquet_episode_rows = connection.execute( + f""" + SELECT "episode_index", "tasks", "length", + "data/chunk_index", "data/file_index", + "dataset_from_index", "dataset_to_index" + FROM {episode_metadata_relation} + ORDER BY "episode_index" + """ + ).fetchall() + for parquet_episode_row in parquet_episode_rows: + tasks = parquet_episode_row[1] + task = (str(tasks[0]) if tasks else "") if isinstance(tasks, list) else str(tasks or "") + episode_rows.append( + { + "episode_index": int(parquet_episode_row[0]), + "task": task, + "length": int(parquet_episode_row[2]), + "data_chunk": str(parquet_episode_row[3]).split("/")[-1], + "data_file": str(parquet_episode_row[4]).split("/")[-1], + "data_from": int(parquet_episode_row[5]), + "data_to": int(parquet_episode_row[6]), + } + ) # Video window columns: videos//{chunk_index,file_index,from_timestamp,to_timestamp} flattened_column_names = [ column_description[0] for column_description in connection.execute( - "SELECT * FROM read_parquet('" - + str(episode_metadata_files[0]).replace("'", "''") - + "') LIMIT 1" + f"SELECT * FROM {episode_metadata_relation} LIMIT 1" ).description ] video_keys = sorted( @@ -479,9 +516,8 @@ def _ensure_source_archive(dataset_source: DatasetSource, cache_dir: Path) -> _S f'"videos/{camera_key}/to_timestamp" as "vto_{camera_key}",', ] video_window_select_sql = "episode_index, " + " ".join(video_window_selectors).rstrip(",") - first_episode_metadata_file = str(episode_metadata_files[0]).replace("'", "''") video_window_rows = connection.execute( - f"SELECT {video_window_select_sql} FROM read_parquet('{first_episode_metadata_file}')" + f"SELECT {video_window_select_sql} FROM {episode_metadata_relation}" ).fetchall() video_window_column_names = [ column_description[0] for column_description in connection.description diff --git a/tests/test_lerobot_converter.py b/tests/test_lerobot_converter.py index 1a1b60d6..aa8bb240 100755 --- a/tests/test_lerobot_converter.py +++ b/tests/test_lerobot_converter.py @@ -259,6 +259,105 @@ def fake_dl(url: str, dest: Path, **kw: object) -> None: ] == pytest.approx(2.0) +@pytest.mark.parametrize( + "second_shard_path", + [ + # Distinct basenames in one chunk directory: how lerobot/droid_1.0.1 + # ships its seven metadata shards. + "meta/episodes/chunk-000/file-001.parquet", + # The same basename in the next chunk directory, which a cache keyed + # by basename alone would collapse onto the first shard. + "meta/episodes/chunk-001/file-000.parquet", + ], +) +def test_index_discovery_reads_every_metadata_shard( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, second_shard_path: str +) -> None: + """Episodes and video windows come from every ``meta/episodes`` shard (#293). + + The corpus is split the way Dataset v3 shards it: a different pair of + episodes in each file and a distinct video window per episode. + """ + corpus = _build_fake_corpus(tmp_path) + single_shard = tmp_path / "meta" / "episodes" / "chunk-000" / "file-000.parquet" + shard_paths = ("meta/episodes/chunk-000/file-000.parquet", second_shard_path) + import duckdb + + conn = duckdb.connect() + conn.execute( + "CREATE TABLE all_episodes AS SELECT * FROM read_parquet('" + + str(single_shard).replace("'", "''") + + "')" + ) + for shard_path, episode_indexes in zip(shard_paths, ((0, 1), (2, 3)), strict=True): + shard_file = tmp_path / shard_path + shard_file.parent.mkdir(parents=True, exist_ok=True) + conn.execute( + f"COPY (SELECT * FROM all_episodes WHERE episode_index IN {episode_indexes}) " + f"TO '{str(shard_file).replace(chr(39), chr(39) * 2)}' (FORMAT parquet)" + ) + conn.close() + + monkeypatch.setattr( + prep, "_hf_repo_info", lambda repo, rev: {"sha": rev, "license": "apache-2.0"} + ) + monkeypatch.setattr(prep, "_fetch_info_json", lambda repo, rev, cache: corpus["info"]) + monkeypatch.setattr( + prep, + "_hf_tree", + lambda repo, rev, path: ( + [{"path": shard_path, "type": "file"} for shard_path in shard_paths] + if "episodes" in path + else [{"path": "meta/info.json", "type": "file"}] + ), + ) + downloaded_destinations: dict[str, Path] = {} + + def fake_download(url: str, dest: Path, **kw: object) -> None: + relative_path = url.split("/resolve/abc/", 1)[1] + assert relative_path not in downloaded_destinations, f"downloaded twice: {url}" + downloaded_destinations[relative_path] = dest + dest.parent.mkdir(parents=True, exist_ok=True) + shutil.copy(tmp_path / relative_path, dest) + + monkeypatch.setattr(prep, "_download_file", fake_download) + + ds = prep.DatasetSource(repo_id="fake/repo", revision="abc", license="apache-2.0") + cache_dir = tmp_path / "cache" + found = prep._ensure_source_archive(ds, cache_dir) + # A second discovery reuses each shard's own cache entry instead of + # re-downloading or colliding with the other shard. + prep._ensure_source_archive(ds, cache_dir) + + assert set(downloaded_destinations) == set(shard_paths) + assert len(set(downloaded_destinations.values())) == len(shard_paths) + assert [episode["episode_index"] for episode in found["episodes"]] == [0, 1, 2, 3] + assert set(found["video_keys"]) == {"observation.images.up", "observation.images.side"} + for episode in found["episodes"]: + episode_index = episode["episode_index"] + assert episode["length"] == 60 + episode_index * 5 + for camera_key in found["video_keys"]: + window = episode["video_windows"][camera_key] + assert window["chunk_index"] == "chunk-000" + assert window["to_timestamp"] == pytest.approx(2.0 + episode_index * 0.2) + + +@pytest.mark.parametrize( + "tree_entry_path", + [ + "meta/episodes/../../data/chunk-000/file-000.parquet", + "/tmp/probe-293-absolute.parquet", + "meta/episodes", + "meta/other/file-000.parquet", + ], +) +def test_episode_metadata_cache_path_refuses_entries_outside_the_metadata_tree( + tmp_path: Path, tree_entry_path: str +) -> None: + with pytest.raises(ValueError, match="not a file below meta/episodes/"): + prep._episode_metadata_cache_path(tmp_path, tree_entry_path, repo_id="fake/repo") + + def test_video_cache_distinguishes_file_indices_and_reuses_same_source( tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: