diff --git a/docs/how-to/import-lerobot-v3.md b/docs/how-to/import-lerobot-v3.md index 7e965140..3a5b0988 100644 --- a/docs/how-to/import-lerobot-v3.md +++ b/docs/how-to/import-lerobot-v3.md @@ -29,7 +29,19 @@ data/lerobot_pusht/ └── prepared-manifest.json # source commit and import summary ``` -Re-running against the same output directory reuses downloaded source files. +Re-running against the same output directory reuses downloaded source files. It also +reuses a completed canonical episode when the resolved dataset commit, source episode, +selected camera set, converter version, and canonical pipeline version all match the +requested import. This makes an all-episodes import resumable at episode boundaries: if a +later episode fails, retrying the same import keeps already completed episode bytes and +converts only the remaining work. + +An existing landing file is reused only when it is a readable MCAP with a complete summary +and matching identity metadata. Truncated, unreadable, or identity-mismatched files are +converted again and atomically replaced only after the new episode succeeds. A branch or tag +that resolves to a different commit is a different import identity, even when the revision +name is unchanged. + Run `uv run hflow doctor ./data/lerobot_pusht/landing/*.mcap` to print the canonical-format report. @@ -49,7 +61,11 @@ Durable outputs land as `landing/*.mcap` and `prepared-manifest.json` under that prefix. Hugging Face downloads stay in the local mirror under `_lerobot_cache/` (`HFLOW_MIRROR_DIR`, or `$XDG_CACHE_HOME/hflow/mirrors`) and are never uploaded into the bucket. The success manifest is published only -after every selected episode object has been written. +after every selected episode object has been written or verified as reusable. Its schema-v3 +episode list records each completed source episode and landing output key alongside the +resolved dataset, camera selection, converter version, and canonical pipeline version. A +failed mid-batch attempt therefore leaves completed episode objects in place but does not +publish a success manifest for an incomplete selected set. For a gated or private repository, export a read token as `HF_TOKEN` (or `HUGGING_FACE_HUB_TOKEN`) before running the command. HFlow sends the token diff --git a/src/hflow/importers/lerobot.py b/src/hflow/importers/lerobot.py index a85bf590..c0516d7e 100644 --- a/src/hflow/importers/lerobot.py +++ b/src/hflow/importers/lerobot.py @@ -29,11 +29,13 @@ from typing import NotRequired, TypedDict from urllib.parse import urlsplit +from mcap.exceptions import McapError from mcap.writer import Writer as McapWriter +from hflow.episode import Episode from hflow.ffmpeg import ffmpeg_path, ffmpeg_version, ffprobe_path from hflow.storage import LocalStorageRoot, StorageRoot, parse_storage_root -from hflow.transform import TransformConfig, write_canonical_episode +from hflow.transform import TransformConfig, compute_pipeline_version, write_canonical_episode logger = logging.getLogger(__name__) @@ -44,6 +46,7 @@ CONVERTER_VERSION = "lerobot-converter-v4" PRESENTATION_TIMESTAMP_EPSILON_S = 0.050 +_LEROBOT_TRANSFORM_CONFIG = TransformConfig(gop_seconds=1.0) # Timestamp handling NANOSECONDS_PER_SECOND = 1_000_000_000 @@ -76,6 +79,58 @@ class DatasetSource: license: str +@dataclass(frozen=True) +class _CompletedEpisodeIdentity: + dataset_repo: str + source_revision: str + source_episode_index: int + camera_topics: frozenset[str] + converter_version: str + canonical_pipeline_version: str + + +def _landing_relative_key(episode_index: int) -> str: + return f"landing/lerobot_episode_{episode_index + 1:04d}.mcap" + + +def _completed_episode_uri( + storage: StorageRoot, + identity: _CompletedEpisodeIdentity, +) -> str | None: + landing_relative_key = _landing_relative_key(identity.source_episode_index) + try: + candidate_path = storage.fetch(landing_relative_key) + except FileNotFoundError: + return None + + try: + with Episode(candidate_path) as episode: + camera_topics = frozenset(episode.cameras) + metadata_records = episode.metadata_records + except (McapError, UnicodeDecodeError, ValueError, struct.error) as error: + logger.info( + "existing LeRobot episode %s is not reusable: %s", + storage.uri_for(landing_relative_key), + error, + ) + return None + + episode_metadata = metadata_records.get("episode/v1", {}) + source_provenance = metadata_records.get("source-provenance/v1", {}) + canonical_provenance = metadata_records.get("provenance/v1", {}) + matches_identity = ( + episode_metadata.get("source_dataset") == identity.dataset_repo + and episode_metadata.get("source_revision") == identity.source_revision + and episode_metadata.get("source_episode_index") == str(identity.source_episode_index) + and source_provenance.get("converter_version") == identity.converter_version + and camera_topics == identity.camera_topics + and canonical_provenance.get("pipeline_version") == identity.canonical_pipeline_version + ) + if not matches_identity: + return None + return storage.uri_for(landing_relative_key) + + class _DatasetRepositoryInformation(TypedDict): sha: str license: str @@ -593,7 +648,8 @@ def import_lerobot_dataset( prefix (``s3://``, ``gs://``, ``az://``). Hugging Face downloads and MCAP construction use local staging under the root's workspace. Durable outputs are published as ``landing/*.mcap`` plus ``prepared-manifest.json`` after - every selected episode succeeds. The source cache (``_lerobot_cache``) stays + every selected episode is either verified as reusable or converted successfully. + The source cache (``_lerobot_cache``) stays in the workspace -- the local directory itself, or the bucket mirror under ``HFLOW_MIRROR_DIR`` -- and is never uploaded into a bucket root. @@ -675,11 +731,26 @@ def import_lerobot_dataset( [episode_index] if episode_index is not None else list(range(len(episode_rows))) ) published_episode_uris: list[str] = [] + completed_episodes: list[dict[str, int | str]] = [] dataset_source = source_archive["dataset"] + canonical_pipeline_version = compute_pipeline_version(_LEROBOT_TRANSFORM_CONFIG) + camera_topics = frozenset(f"/{camera_key}" for camera_key in resolved_camera_keys) for selected_episode_index in selected_episode_indexes: - published_episode_uris.append( - _convert_single_episode( + landing_relative_key = _landing_relative_key(selected_episode_index) + completed_episode_uri = _completed_episode_uri( + storage, + _CompletedEpisodeIdentity( + dataset_repo=dataset_source.repo_id, + source_revision=dataset_source.revision, + source_episode_index=selected_episode_index, + camera_topics=camera_topics, + converter_version=CONVERTER_VERSION, + canonical_pipeline_version=canonical_pipeline_version, + ), + ) + if completed_episode_uri is None: + completed_episode_uri = _convert_single_episode( source_archive=source_archive, dataset_source=dataset_source, storage=storage, @@ -688,11 +759,17 @@ def import_lerobot_dataset( numeric_schemas=numeric_schemas, frames_per_second=int(source_archive["fps"]), ) + published_episode_uris.append(completed_episode_uri) + completed_episodes.append( + { + "source_episode_index": selected_episode_index, + "output_key": landing_relative_key, + } ) manifest_contents = json.dumps( { - "schema_version": 2, + "schema_version": 3, "dataset": { "repo_id": dataset_source.repo_id, "revision": dataset_source.revision, @@ -701,6 +778,8 @@ def import_lerobot_dataset( "camera_keys": list(resolved_camera_keys), "episodes_converted": len(published_episode_uris), "converter_version": CONVERTER_VERSION, + "canonical_pipeline_version": canonical_pipeline_version, + "completed_episodes": completed_episodes, }, indent=2, ) @@ -894,7 +973,7 @@ def _feature_rows(feature_name: str) -> list | None: # Write MCAP into local staging, then publish the complete file. output_file_name = f"lerobot_episode_{episode_index + 1:04d}.mcap" - landing_relative_key = f"landing/{output_file_name}" + landing_relative_key = _landing_relative_key(episode_index) from foxglove_schemas_protobuf.CompressedVideo_pb2 import CompressedVideo from mcap_protobuf.schema import build_file_descriptor_set @@ -1000,7 +1079,7 @@ def _feature_rows(feature_name: str) -> list | None: write_canonical_episode( source_episode_path, canonical_episode_path, - TransformConfig(gop_seconds=1.0), + _LEROBOT_TRANSFORM_CONFIG, source_uri=source_uri, ) published_uri = storage.publish(canonical_episode_path, landing_relative_key) diff --git a/tests/test_lerobot_converter.py b/tests/test_lerobot_converter.py index 1a1b60d6..6ee09a5c 100755 --- a/tests/test_lerobot_converter.py +++ b/tests/test_lerobot_converter.py @@ -16,7 +16,9 @@ import pytest import hflow.importers.lerobot as prep +from hflow._grouped_mcap_writer import GroupedMcapWriter from hflow.cli import main as cli_main +from hflow.episode import Episode from hflow.storage import LocalStorageRoot, StorageRoot _DERIVE = prep._derive_numeric_schema @@ -960,6 +962,72 @@ def _stub_single_episode_source_archive( } +def _write_completed_test_episode( + output_path: Path, + *, + dataset_source: prep.DatasetSource, + episode_index: int, + camera_keys: tuple[str, ...] = (prep.DEFAULT_CAMERA_KEY,), + source_revision: str | None = None, + conversion_marker: str = "test", +) -> None: + """Write the minimal readable MCAP needed to exercise resume identity.""" + + with GroupedMcapWriter(output_path) as writer: + video_schema_id = writer.register_schema( + name="foxglove.CompressedVideo", + encoding="protobuf", + data=b"", + ) + for camera_key in camera_keys: + writer.register_channel( + f"/{camera_key}", + message_encoding="protobuf", + schema_id=video_schema_id, + group="cameras", + ) + writer.add_metadata( + "episode/v1", + { + "source_dataset": dataset_source.repo_id, + "source_revision": source_revision or dataset_source.revision, + "source_episode_index": str(episode_index), + }, + ) + writer.add_metadata( + "source-provenance/v1", + {"converter_version": prep.CONVERTER_VERSION}, + ) + writer.add_metadata( + "provenance/v1", + { + "schema_version": "1", + "pipeline_version": prep.compute_pipeline_version(prep._LEROBOT_TRANSFORM_CONFIG), + }, + ) + writer.add_metadata("test-conversion/v1", {"marker": conversion_marker}) + + +def _publish_completed_test_episode( + *, + storage: StorageRoot, + staging_dir: Path, + dataset_source: prep.DatasetSource, + episode_index: int, + camera_keys: tuple[str, ...], + conversion_marker: str, +) -> str: + staged = staging_dir / f"staged-{episode_index}-{conversion_marker}.mcap" + _write_completed_test_episode( + staged, + dataset_source=dataset_source, + episode_index=episode_index, + camera_keys=camera_keys, + conversion_marker=conversion_marker, + ) + return storage.publish(staged, prep._landing_relative_key(episode_index)) + + def _install_publish_through_convert(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> list[str]: """Exercise StorageRoot.publish without running the video converter.""" @@ -1044,7 +1112,7 @@ def test_import_publishes_into_a_bucket_data_root_without_uploading_cache( assert (data_root.mirror / "_lerobot_cache" / "abc").is_dir() -def test_import_skips_bucket_manifest_when_an_episode_publish_fails( +def test_import_resumes_bucket_after_an_episode_publish_fails( tmp_path: Path, bucket_over_tmp: tuple[object, Path], monkeypatch: pytest.MonkeyPatch, @@ -1066,33 +1134,39 @@ def ensure_two_episodes(dataset_source: prep.DatasetSource, cache_dir: Path) -> ] return archive - convert_calls = 0 + first_failure_pending = True + conversion_attempt = 0 - def fail_on_second_episode( + def fail_once_on_second_episode( *, source_archive: object, - dataset_source: object, + dataset_source: prep.DatasetSource, storage: StorageRoot, episode_index: int, - camera_keys: object, + camera_keys: tuple[str, ...], numeric_schemas: object, frames_per_second: object, ) -> str: - nonlocal convert_calls - del source_archive, dataset_source, camera_keys, numeric_schemas, frames_per_second - convert_calls += 1 - if episode_index == 1: + nonlocal conversion_attempt, first_failure_pending + del source_archive, numeric_schemas, frames_per_second + conversion_attempt += 1 + if episode_index == 1 and first_failure_pending: + first_failure_pending = False raise RuntimeError("forced publish failure") - relative_key = f"landing/lerobot_episode_{episode_index + 1:04d}.mcap" - staged = tmp_path / f"staged-{episode_index}.mcap" - staged.write_bytes(b"first") - return storage.publish(staged, relative_key) + return _publish_completed_test_episode( + storage=storage, + staging_dir=tmp_path, + dataset_source=dataset_source, + episode_index=episode_index, + camera_keys=camera_keys, + conversion_marker=f"attempt-{conversion_attempt}", + ) monkeypatch.setattr( prep, "_hf_repo_info", lambda repo, revision: {"sha": "abc", "license": "apache-2.0"} ) monkeypatch.setattr(prep, "_ensure_source_archive", ensure_two_episodes) - monkeypatch.setattr(prep, "_convert_single_episode", fail_on_second_episode) + monkeypatch.setattr(prep, "_convert_single_episode", fail_once_on_second_episode) with pytest.raises(RuntimeError, match="forced publish failure"): prep.import_lerobot_dataset( @@ -1101,7 +1175,138 @@ def fail_on_second_episode( output_dir=data_root, ) - assert convert_calls == 2 - assert (remote_dir / "landing" / "lerobot_episode_0001.mcap").is_file() + first_episode_path = remote_dir / "landing" / "lerobot_episode_0001.mcap" + first_episode_bytes = first_episode_path.read_bytes() + assert first_episode_path.is_file() assert not (remote_dir / "prepared-manifest.json").exists() assert "prepared-manifest.json" not in data_root.list_names() + + episode_uris = prep.import_lerobot_dataset( + dataset_repo="fake/repo", + revision="main", + output_dir=data_root, + ) + + assert first_episode_path.read_bytes() == first_episode_bytes + assert episode_uris == [ + f"{data_root.url}/landing/lerobot_episode_0001.mcap", + f"{data_root.url}/landing/lerobot_episode_0002.mcap", + ] + assert (remote_dir / "landing" / "lerobot_episode_0002.mcap").is_file() + manifest = json.loads((remote_dir / "prepared-manifest.json").read_text()) + assert manifest == { + "schema_version": 3, + "dataset": { + "repo_id": "fake/repo", + "revision": "abc", + "license": "apache-2.0", + }, + "camera_keys": [prep.DEFAULT_CAMERA_KEY], + "episodes_converted": 2, + "converter_version": prep.CONVERTER_VERSION, + "canonical_pipeline_version": prep.compute_pipeline_version(prep._LEROBOT_TRANSFORM_CONFIG), + "completed_episodes": [ + { + "source_episode_index": 0, + "output_key": "landing/lerobot_episode_0001.mcap", + }, + { + "source_episode_index": 1, + "output_key": "landing/lerobot_episode_0002.mcap", + }, + ], + } + + +@pytest.mark.parametrize( + "candidate_kind", + ["invalid-magic", "truncated", "no-summary", "identity-mismatch"], +) +def test_import_replaces_unreadable_or_identity_mismatched_completed_candidate( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, candidate_kind: str +) -> None: + output_dir = tmp_path / "out" + dataset_source = prep.DatasetSource(repo_id="fake/repo", revision="abc", license="apache-2.0") + landing_path = output_dir / "landing" / "lerobot_episode_0001.mcap" + landing_path.parent.mkdir(parents=True) + + if candidate_kind == "invalid-magic": + landing_path.write_bytes(b"not an mcap") + elif candidate_kind == "truncated": + complete_path = tmp_path / "complete-before-truncation.mcap" + _write_completed_test_episode(complete_path, dataset_source=dataset_source, episode_index=0) + landing_path.write_bytes(complete_path.read_bytes()[:-4]) + elif candidate_kind == "no-summary": + from mcap.writer import IndexType + + with landing_path.open("wb") as stream: + writer = prep.McapWriter( + stream, + index_types=IndexType.NONE, + repeat_channels=False, + repeat_schemas=False, + use_statistics=False, + use_summary_offsets=False, + ) + writer.start(profile="", library="test") + writer.finish() + elif candidate_kind == "identity-mismatch": + _write_completed_test_episode( + landing_path, + dataset_source=dataset_source, + episode_index=0, + source_revision="different-revision", + conversion_marker="stale", + ) + else: + raise AssertionError(f"unknown candidate kind: {candidate_kind}") + + candidate_bytes = landing_path.read_bytes() + + monkeypatch.setattr( + prep, "_hf_repo_info", lambda repo, revision: {"sha": "abc", "license": "apache-2.0"} + ) + monkeypatch.setattr(prep, "_ensure_source_archive", _stub_single_episode_source_archive) + + def replace_candidate( + *, + source_archive: object, + dataset_source: prep.DatasetSource, + storage: StorageRoot, + episode_index: int, + camera_keys: tuple[str, ...], + numeric_schemas: object, + frames_per_second: object, + ) -> str: + del source_archive, numeric_schemas, frames_per_second + return _publish_completed_test_episode( + storage=storage, + staging_dir=tmp_path, + dataset_source=dataset_source, + episode_index=episode_index, + camera_keys=camera_keys, + conversion_marker="replacement", + ) + + monkeypatch.setattr(prep, "_convert_single_episode", replace_candidate) + + episode_uris = prep.import_lerobot_dataset( + dataset_repo="fake/repo", + revision="main", + output_dir=output_dir, + episode_index=0, + ) + + assert episode_uris == [str(landing_path.resolve())] + assert landing_path.read_bytes() != candidate_bytes + with Episode(landing_path) as episode: + assert episode.metadata_records["episode/v1"]["source_revision"] == "abc" + assert episode.metadata_records["test-conversion/v1"] == {"marker": "replacement"} + manifest = json.loads((output_dir / "prepared-manifest.json").read_text()) + assert manifest["schema_version"] == 3 + assert manifest["completed_episodes"] == [ + { + "source_episode_index": 0, + "output_key": "landing/lerobot_episode_0001.mcap", + } + ]