From ba4db0af9171672611a267bb963a31253c238c35 Mon Sep 17 00:00:00 2001 From: Arpit Gupta Date: Sun, 19 Jul 2026 21:48:02 -0400 Subject: [PATCH 1/2] feat(io): streaming metadata reads for cloud LeRobot-v3 & Zarr (forge inspect) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `forge inspect s3://…` / `gs://…` now reads metadata over the network with range requests instead of downloading the whole dataset. Inspecting a 464 MB cloud LeRobot-v3 dataset fetches ~3 KB (its info.json) in ~0.3s — a ~150,000x reduction in bytes moved — and returns the identical DatasetInfo. - forge/io/source.py: `DataSource` — a dataset root (local or fsspec-remote) with the small metadata reads streaming inspection needs (json, parquet footer/schema, exists, glob), fetched as range requests for remote sources. - lerobot-v3 reader: `_inspect_remote` reads only meta/info.json (which carries episode/frame counts, fps, robot, per-camera dims, and the action/observation schema, so the mp4 probe is skipped). Local inspect path is untouched; info.json→DatasetInfo mapping factored into `_populate_from_info` and shared. - zarr reader: opens the store natively over fsspec (`zarr.open("s3://…")` range-reads chunks); can_read checks markers remotely. - registry: `detect_format` keeps remote URIs as strings (non-streamable readers safely return False, so only streamable formats match without a download); `STREAMABLE_FORMATS = {lerobot-v3, zarr}`. - CLI: `forge inspect` streams remote streamable formats; falls back to download for `--deep`, `--generate-config`, or other formats/commands. Tests (tests/test_streaming.py): correctness over memory:// (right DatasetInfo, no temp download) for both formats, plus a moto-S3 byte-budget assertion proving inspect fetches < 1 MB even next to an 8 MB payload. Existing cloud-inspect test updated (inspect now streams, `--deep` still downloads). No other tests change. Co-Authored-By: Claude Opus 4.8 (1M context) --- CHANGELOG.md | 11 ++ README.md | 4 +- forge/cli.py | 55 +++++--- forge/formats/lerobot_v3/reader.py | 65 +++++++++- forge/formats/registry.py | 11 +- forge/formats/zarr/reader.py | 51 ++++++-- forge/io/README.md | 11 +- forge/io/__init__.py | 2 + forge/io/source.py | 70 +++++++++++ tests/test_io_paths.py | 16 ++- tests/test_streaming.py | 195 +++++++++++++++++++++++++++++ 11 files changed, 458 insertions(+), 33 deletions(-) create mode 100644 forge/io/source.py create mode 100644 tests/test_streaming.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 90585f6..bff049a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,17 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added +- **Streaming metadata reads for cloud datasets (LeRobot-v3, Zarr).** + `forge inspect s3://…` / `gs://…` now reads metadata over the network with + **range requests** instead of downloading the whole dataset — inspecting a + 464 MB cloud LeRobot-v3 dataset fetches ~3 KB (its `info.json`), a ~150,000× + reduction in bytes moved, in a fraction of a second. LeRobot-v3 reads + `info.json` (which carries episode/frame counts, fps, robot, per-camera dims, + and schema); Zarr opens the store natively over fsspec. `--deep` forces a full + download for per-episode stats; non-streamable formats and other commands are + unchanged (download-to-temp). New `forge.io.DataSource` helper; the readers + keep their local paths untouched. + - **Search- and selection-driven curation.** `forge curate` now selects episodes three ways — a SQL `--where` predicate, an explicit `--ids` list, or a `--from ` file — so semantic search and Forge Studio can diff --git a/README.md b/README.md index 4a7f494..6d7dc8c 100644 --- a/README.md +++ b/README.md @@ -286,6 +286,8 @@ forge inspect s3://my-bucket/datasets/run_0413 forge convert gs://lab-data/rosbags ./out --format lerobot-v3 ``` +**Streaming inspect.** For **LeRobot-v3** and **Zarr**, `forge inspect` reads metadata over the network with **range requests** instead of downloading — e.g. inspecting a 464 MB cloud LeRobot dataset fetches ~3 KB (its `info.json`), not the whole thing. Add `--deep` to force a full download when you need per-episode stats. + Auth uses each provider's standard credential chain (AWS env vars / profiles / IAM roles; GCP Application Default Credentials) — Forge never handles credentials itself. Catalogs can be read *and written* in the cloud; per-format conversion **outputs** are local-only for now. Full guide + connectivity troubleshooting: [forge/io/README.md](forge/io/README.md). ## Dataset registry @@ -334,7 +336,7 @@ Full reference: [docs/cli.md](docs/cli.md). Run `forge --help` for any - [x] Cloud storage — `s3://` / `gs://` on every command - [x] The catalog — ingest, SQL query, semantic search, dedup, curation, Studio - [ ] **Snapshots + export** — freeze a curated selection → LeRobot/RLDS for training -- [ ] **Streaming reads** — process cloud datasets without a full download +- [~] **Streaming reads** — `forge inspect` streams LeRobot-v3 / Zarr metadata (range reads, no download); ingest/scoring streaming next - [ ] **Dataset merging & splitting** — combine datasets; stratified train/val/test - [ ] **Depth / point-cloud support** · **GR00T writer** · **distributed conversion** diff --git a/forge/cli.py b/forge/cli.py index 5bd72e6..53ad10f 100644 --- a/forge/cli.py +++ b/forge/cli.py @@ -417,24 +417,47 @@ def inspect_cmd( from forge.inspect import InspectionOptions, Inspector - # Resolve HuggingFace URLs to local paths - try: - resolved_path = _resolve_dataset_path(path) - except Exception as e: - console.print(f"[red]Error downloading dataset:[/red] {e}") - raise typer.Exit(1) + from forge.io import is_remote_uri - options = InspectionOptions( - sample_episodes=samples, - deep_scan=deep, - ) - inspector = Inspector(options) + info = None - try: - info = inspector.inspect(resolved_path, format) - except ForgeError as e: - console.print(f"[red]Error:[/red] {e}") - raise typer.Exit(1) + # Streaming inspect for cloud lerobot-v3 / zarr: read metadata over the + # network (range reads) instead of downloading the whole dataset. Falls back + # to download for --deep, --generate-config, or non-streamable formats. + if is_remote_uri(path) and not deep and generate_config is None: + from forge.formats.registry import FormatRegistry + + try: + detected = format or FormatRegistry.detect_format(path) + except Exception: + detected = None + if detected in FormatRegistry.STREAMABLE_FORMATS: + try: + info = FormatRegistry.get_reader(detected).inspect(path) + console.print("[dim]streamed metadata — no download[/dim]") + except ForgeError as e: + console.print(f"[red]Error:[/red] {e}") + raise typer.Exit(1) + + if info is None: + # Resolve HuggingFace URLs / download cloud datasets to local paths + try: + resolved_path = _resolve_dataset_path(path) + except Exception as e: + console.print(f"[red]Error downloading dataset:[/red] {e}") + raise typer.Exit(1) + + options = InspectionOptions( + sample_episodes=samples, + deep_scan=deep, + ) + inspector = Inspector(options) + + try: + info = inspector.inspect(resolved_path, format) + except ForgeError as e: + console.print(f"[red]Error:[/red] {e}") + raise typer.Exit(1) if output == "json": # Convert to JSON-serializable dict diff --git a/forge/formats/lerobot_v3/reader.py b/forge/formats/lerobot_v3/reader.py index eafdef9..574c739 100644 --- a/forge/formats/lerobot_v3/reader.py +++ b/forge/formats/lerobot_v3/reader.py @@ -243,6 +243,11 @@ def can_read(cls, path: Path) -> bool: Returns: True if LeRobot v3 markers found. """ + from forge.io import is_remote_uri + + if is_remote_uri(path): + return cls._can_read_remote(str(path)) + if not path.exists() or not path.is_dir(): return False @@ -282,6 +287,25 @@ def can_read(cls, path: Path) -> bool: return False + @classmethod + def _can_read_remote(cls, uri: str) -> bool: + """Detect a v3 dataset on a remote ``s3://`` / ``gs://`` source (no download). + + Reads only ``meta/info.json`` and checks ``codebase_version``. + """ + from forge.io.paths import get_filesystem + + try: + fs, root = get_filesystem(uri) + info_path = f"{root.rstrip('/')}/meta/info.json" + if not fs.exists(info_path): + return False + with fs.open(info_path, "rb") as f: + version = str(json.loads(f.read()).get("codebase_version", "")).lstrip("v") + return int(version.split(".")[0]) >= 3 + except Exception: + return False + @classmethod def detect_version(cls, path: Path) -> str | None: """Detect specific LeRobot version. @@ -314,6 +338,11 @@ def inspect(self, path: Path) -> DatasetInfo: Returns: DatasetInfo with schema and metadata. """ + from forge.io import is_remote_uri + + if is_remote_uri(path): + return self._inspect_remote(str(path)) + path = Path(path) info = DatasetInfo(path=path, format="lerobot-v3") info.format_version = self.detect_version(path) @@ -335,6 +364,29 @@ def inspect(self, path: Path) -> DatasetInfo: return info + def _inspect_remote(self, uri: str) -> DatasetInfo: + """Stream a v3 dataset's metadata from the cloud — no download. + + Reads only ``meta/info.json`` (a few KB) via a range request. info.json + carries episode/frame counts, fps, robot, per-camera dims, and the + action/observation schema, so the mp4 probe (which needs a local file) + is skipped. ``has_language`` is inferred from the tasks table's presence. + """ + from forge.io import DataSource + + src = DataSource(uri) + if not src.exists("meta", "info.json"): + raise InspectionError(uri, "Missing meta/info.json (required for v3)") + + data = src.read_json("meta", "info.json") + info = DatasetInfo(path=uri, format="lerobot-v3") + info.format_version = str(data.get("codebase_version", "3.0")).lstrip("v") + self._populate_from_info(data, info) + info.has_language = src.exists("meta", "tasks.parquet") or src.exists( + "meta", "tasks.jsonl" + ) + return info + def _load_info_json(self, path: Path, info: DatasetInfo) -> None: """Load metadata from info.json.""" info_path = path / "meta" / "info.json" @@ -344,7 +396,17 @@ def _load_info_json(self, path: Path, info: DatasetInfo) -> None: try: with open(info_path) as f: data = json.load(f) + except json.JSONDecodeError as e: + raise InspectionError(path, f"Invalid JSON in info.json: {e}") + self._populate_from_info(data, info) + + def _populate_from_info(self, data: dict, info: DatasetInfo) -> None: + """Map a parsed info.json dict onto DatasetInfo (cameras, schema, fps…). + Shared by the local and streaming (remote) inspect paths — info.json + alone carries episode/frame counts, fps, robot, and per-camera dims. + """ + if data: # Required fields in v3 info.num_episodes = data.get("total_episodes", 0) info.total_frames = data.get("total_frames", 0) @@ -398,9 +460,6 @@ def _load_info_json(self, path: Path, info: DatasetInfo) -> None: # Check for timestamps info.has_timestamps = data.get("has_timestamps", True) - except json.JSONDecodeError as e: - raise InspectionError(path, f"Invalid JSON in info.json: {e}") - def _str_to_dtype(self, dtype_str: str) -> Dtype: """Convert string dtype to Forge Dtype.""" dtype_str = str(dtype_str).lower() diff --git a/forge/formats/registry.py b/forge/formats/registry.py index e4b49f8..40be500 100644 --- a/forge/formats/registry.py +++ b/forge/formats/registry.py @@ -114,6 +114,10 @@ def get_writer(cls, format_name: str) -> FormatWriter: ) return cls._writers[format_name]() + # Formats whose metadata can be inspected over the network without a full + # download (range reads). Used by `forge inspect` to stream cloud datasets. + STREAMABLE_FORMATS: frozenset[str] = frozenset({"lerobot-v3", "zarr"}) + # Priority order for format detection (more specific formats first) _detection_priority: list[str] = [ "groot", # Check GR00T before LeRobot (more specific) @@ -142,7 +146,12 @@ def detect_format(cls, path: Path | str) -> str: Raises: FormatDetectionError: If no reader can handle the path. """ - path = Path(path) + from forge.io import is_remote_uri + + # Keep remote URIs as strings (Path() would mangle "s3://"); readers' + # can_read handle both. Non-streaming readers safely return False for + # a remote URI, so only streamable formats match without a download. + path = str(path) if is_remote_uri(path) else Path(path) # Check in priority order first for format_name in cls._detection_priority: diff --git a/forge/formats/zarr/reader.py b/forge/formats/zarr/reader.py index 7bd9692..cdab930 100644 --- a/forge/formats/zarr/reader.py +++ b/forge/formats/zarr/reader.py @@ -91,12 +91,20 @@ def format_name(self) -> str: def can_read(cls, path: Path) -> bool: """Check for Zarr markers: .zarr directory, .zarr.zip file, or .zarray files. + Handles ``s3://`` / ``gs://`` sources (checked over fsspec) as well as + local paths. Zarr reads remote stores natively via range requests. + Args: - path: Path to potential Zarr dataset. + path: Path or URI to a potential Zarr dataset. Returns: True if Zarr markers found. """ + from forge.io import is_remote_uri + + if is_remote_uri(path): + return cls._can_read_remote(str(path)) + if not path.exists(): return False @@ -132,6 +140,33 @@ def can_read(cls, path: Path) -> bool: return False + @classmethod + def _can_read_remote(cls, uri: str) -> bool: + """Zarr-marker check for a remote ``s3://`` / ``gs://`` store (no download).""" + from forge.io.paths import get_filesystem + + if uri.endswith((".zarr", ".zarr.zip")): + return True + try: + fs, root = get_filesystem(uri) + root = root.rstrip("/") + for marker in (".zgroup", ".zarray", "zarr.json"): + if fs.exists(f"{root}/{marker}"): + return True + # nested structure (data/ + meta/) + if fs.exists(f"{root}/data") or fs.exists(f"{root}/meta"): + return True + except Exception: + return False + return False + + @staticmethod + def _store(path): + """Return the argument to ``zarr.open``: a URI (streamed) or local path.""" + from forge.io import is_remote_uri + + return str(path) if is_remote_uri(path) else str(Path(path)) + @classmethod def detect_version(cls, path: Path) -> str | None: """Detect Zarr format version. @@ -146,7 +181,7 @@ def detect_version(cls, path: Path) -> str | None: import zarr try: - root = zarr.open(str(path), mode="r") + root = zarr.open(cls._store(path), mode="r") # Check for version in attrs if hasattr(root, "attrs"): version = root.attrs.get("version", root.attrs.get("format_version")) @@ -172,12 +207,15 @@ def inspect(self, path: Path) -> DatasetInfo: _check_zarr() import zarr - path = Path(path) - info = DatasetInfo(path=path, format="zarr") + from forge.io import is_remote_uri + + # Keep the URI as-is for remote sources (Path() would mangle "s3://"). + info_path = path if is_remote_uri(path) else Path(path) + info = DatasetInfo(path=info_path, format="zarr") info.format_version = self.detect_version(path) try: - root = zarr.open(str(path), mode="r") + root = zarr.open(self._store(path), mode="r") except Exception as e: raise InspectionError(path, f"Failed to open Zarr store: {e}") @@ -338,8 +376,7 @@ def read_episodes(self, path: Path) -> Iterator[Episode]: import numpy as np import zarr - path = Path(path) - root = zarr.open(str(path), mode="r") + root = zarr.open(self._store(path), mode="r") # Get episode boundaries episode_ends = self._get_episode_ends(root) diff --git a/forge/io/README.md b/forge/io/README.md index e229e11..6ef3eba 100644 --- a/forge/io/README.md +++ b/forge/io/README.md @@ -23,10 +23,13 @@ source string ──▶ localize() ──▶ local Path ──▶ existing reade hf://org/repo (handled by forge.hub, not here) ``` -> **Design note:** remote sources are downloaded in full before processing — -> not range-read/streamed. This keeps readers that seek randomly (video -> keyframes, HDF5 chunks, rosbag indexes) working without modification. See the -> [CHANGELOG](../../CHANGELOG.md) for the streaming follow-up. +> **Design note:** most operations download the remote source in full before +> processing — this keeps readers that seek randomly (video keyframes, HDF5 +> chunks, rosbag indexes) working without modification. The exception is +> **`forge inspect` on LeRobot-v3 / Zarr**, which reads metadata over the +> network via **range requests** ([`DataSource`](source.py)) — inspecting a +> multi-GB cloud dataset fetches only kilobytes. Streaming the ingest/scoring +> paths (parquet column reads) is the next step. ## Install diff --git a/forge/io/__init__.py b/forge/io/__init__.py index 9b9a524..26640e2 100644 --- a/forge/io/__init__.py +++ b/forge/io/__init__.py @@ -17,8 +17,10 @@ is_remote_uri, localize, ) +from forge.io.source import DataSource __all__ = [ + "DataSource", "RemoteWriteNotSupportedError", "cleanup_localized", "get_filesystem", diff --git a/forge/io/source.py b/forge/io/source.py new file mode 100644 index 0000000..87610c9 --- /dev/null +++ b/forge/io/source.py @@ -0,0 +1,70 @@ +"""A dataset root that may be local or remote, for streaming metadata reads. + +Format readers use :class:`DataSource` to inspect a dataset without downloading +it: for ``s3://`` / ``gs://`` sources the small metadata files (``info.json``, +Parquet footers) are fetched with **range requests**, not a full copy. Local +paths go through the same interface (fsspec's local filesystem), so a reader can +share one code path — but readers keep their existing local logic and only reach +for this on the remote/streaming path. +""" + +from __future__ import annotations + +import json +from typing import TYPE_CHECKING + +from forge.io.paths import get_filesystem, is_remote_uri + +if TYPE_CHECKING: + import pyarrow as pa + + +class DataSource: + """A dataset root (local or fsspec-remote) with metadata-read helpers.""" + + def __init__(self, uri: str): + self.uri = str(uri) + self.is_remote = is_remote_uri(self.uri) + self.fs, self.root = get_filesystem(self.uri) + + def path(self, *parts: str) -> str: + """Join ``parts`` onto the root as an fs-native path.""" + return "/".join([self.root.rstrip("/"), *[p.strip("/") for p in parts]]) + + def exists(self, *parts: str) -> bool: + try: + return self.fs.exists(self.path(*parts)) + except FileNotFoundError: + return False + + def open(self, *parts: str): + """Open a file (binary) — a range-capable handle for remote sources.""" + return self.fs.open(self.path(*parts), "rb") + + def read_json(self, *parts: str) -> dict: + with self.open(*parts) as f: + return json.loads(f.read()) + + def read_text(self, *parts: str) -> str: + with self.open(*parts) as f: + return f.read().decode("utf-8") + + def glob(self, *parts: str) -> list[str]: + try: + return self.fs.glob(self.path(*parts)) + except FileNotFoundError: + return [] + + def parquet_schema(self, full_path: str) -> pa.Schema: + """Read a Parquet file's schema from its footer (range read; no full copy).""" + import pyarrow.parquet as pq + + with self.fs.open(full_path, "rb") as f: + return pq.read_schema(f) + + def read_parquet(self, full_path: str, **kwargs) -> pa.Table: + """Read a Parquet file (optionally ``columns=…``) via range reads.""" + import pyarrow.parquet as pq + + with self.fs.open(full_path, "rb") as f: + return pq.read_table(f, **kwargs) diff --git a/tests/test_io_paths.py b/tests/test_io_paths.py index 81d994e..86e8171 100644 --- a/tests/test_io_paths.py +++ b/tests/test_io_paths.py @@ -195,9 +195,23 @@ def test_inspect_resolves_cloud_uri(self, memory_fs): result = CliRunner().invoke(app, ["inspect", "memory:///cli-bkt/ds"]) assert result.exit_code == 0, result.output - assert "Fetching from cloud storage" in result.output + # lerobot-v3 is streamable, so inspect reads metadata over the network + # instead of downloading the dataset. + assert "streamed metadata" in result.output assert "lerobot" in result.output.lower() + def test_inspect_deep_downloads_cloud_uri(self, memory_fs): + # --deep needs the full data, so it falls back to download. + from typer.testing import CliRunner + + from forge.cli import app + + self._make_lerobot_v3_in_memory(memory_fs, "/cli-bkt-deep/ds") + result = CliRunner().invoke( + app, ["inspect", "memory:///cli-bkt-deep/ds", "--deep"] + ) + assert "Fetching from cloud storage" in result.output + def test_convert_rejects_cloud_output(self, memory_fs): from typer.testing import CliRunner diff --git a/tests/test_streaming.py b/tests/test_streaming.py new file mode 100644 index 0000000..31d8997 --- /dev/null +++ b/tests/test_streaming.py @@ -0,0 +1,195 @@ +"""Streaming metadata reads for cloud datasets (lerobot-v3, zarr). + +`forge inspect` on an ``s3://`` / ``gs://`` lerobot-v3 or zarr dataset reads +metadata via range requests instead of downloading the whole dataset. These +tests verify (a) it produces the right DatasetInfo over fsspec without a local +copy, and (b) — against a moto S3 mock — that it actually streams, by asserting +the bytes pulled off the wire stay under a tight budget. +""" + +from __future__ import annotations + +import json +from pathlib import Path + +import numpy as np +import pytest + +from forge.formats.registry import FormatRegistry +from forge.io import cleanup_localized +from forge.io.paths import _TEMP_ROOTS + + +def _lerobot_v3_info() -> dict: + return { + "codebase_version": "v3.0", + "robot_type": "franka", + "total_episodes": 42, + "total_frames": 5000, + "fps": 20, + "features": { + "observation.state": {"dtype": "float32", "shape": [7]}, + "action": {"dtype": "float32", "shape": [7]}, + "observation.images.wrist": {"dtype": "video", "shape": [240, 320, 3]}, + }, + } + + +@pytest.fixture +def memory_fs(): + fs = pytest.importorskip("fsspec").filesystem("memory") + try: + for p in fs.find("/"): + fs.rm(p) + except FileNotFoundError: + pass + yield fs + cleanup_localized() + + +def _write_mem(fs, path: str, data: bytes) -> None: + with fs.open(path, "wb") as f: + f.write(data) + + +class TestLeRobotV3Streaming: + def test_detect_and_inspect_over_memory(self, memory_fs): + _write_mem( + memory_fs, "/lr3/meta/info.json", json.dumps(_lerobot_v3_info()).encode() + ) + _write_mem(memory_fs, "/lr3/meta/tasks.jsonl", b'{"task": "pick"}\n') + + uri = "memory:///lr3" + assert FormatRegistry.detect_format(uri) == "lerobot-v3" + + before = len(_TEMP_ROOTS) + info = FormatRegistry.get_reader("lerobot-v3").inspect(uri) + + # correct metadata, straight from info.json + assert info.num_episodes == 42 + assert info.total_frames == 5000 + assert info.inferred_fps == 20 + assert info.inferred_robot_type == "franka" + assert info.action_schema.shape == (7,) + assert "wrist" in info.cameras + cam = info.cameras["wrist"] + assert (cam.height, cam.width) == (240, 320) + assert info.has_language is True + # ...and it did NOT download anything. + assert len(_TEMP_ROOTS) == before + + def test_not_v3_without_info(self, memory_fs): + # a bare prefix isn't a v3 dataset + _write_mem(memory_fs, "/empty/readme.txt", b"hi") + assert FormatRegistry.get_reader("lerobot-v3")._can_read_remote( + "memory:///empty" + ) is False + + +class TestZarrStreaming: + def _upload_zarr(self, fs, prefix: str, episodes=2, frames=10): + zarr = pytest.importorskip("zarr") + import tempfile + + tmp = Path(tempfile.mkdtemp()) / "z.zarr" + root = zarr.open(str(tmp), mode="w") + root.attrs["fps"] = 10 + root.attrs["robot_type"] = "ur5" + data = root.create_group("data") + mk = getattr(data, "create_array", None) or data.create_dataset + cam = mk("camera0_rgb", shape=(episodes * frames, 16, 16, 3), dtype="uint8") + cam[:] = 0 + act = mk("action", shape=(episodes * frames, 7), dtype="float32") + act[:] = 0 + meta = root.create_group("meta") + mm = getattr(meta, "create_array", None) or meta.create_dataset + ends = mm("episode_ends", shape=(episodes,), dtype="int64") + ends[:] = np.array([(i + 1) * frames for i in range(episodes)]) + for p in tmp.rglob("*"): + if p.is_file(): + with fs.open(f"{prefix}/{p.relative_to(tmp).as_posix()}", "wb") as f: + f.write(p.read_bytes()) + + def test_detect_and_inspect_over_memory(self, memory_fs): + pytest.importorskip("zarr") + self._upload_zarr(memory_fs, "/z.zarr", episodes=3, frames=8) + + uri = "memory:///z.zarr" + assert FormatRegistry.detect_format(uri) == "zarr" + + before = len(_TEMP_ROOTS) + info = FormatRegistry.get_reader("zarr").inspect(uri) + assert info.num_episodes == 3 + assert info.inferred_fps == 10 + assert info.inferred_robot_type == "ur5" + assert info.action_schema is not None + assert len(_TEMP_ROOTS) == before # no download + + +class TestStreamableRegistry: + def test_streamable_set(self): + assert FormatRegistry.STREAMABLE_FORMATS == frozenset({"lerobot-v3", "zarr"}) + + +# --------------------------------------------------------------------------- +# Byte-budget: prove it streams (range reads), not downloads — via moto S3. +# --------------------------------------------------------------------------- + + +@pytest.fixture +def moto_s3(): + pytest.importorskip("s3fs") + moto_server = pytest.importorskip("moto.server") + import s3fs + + server = moto_server.ThreadedMotoServer(ip_address="127.0.0.1", port=0) + server.start() + host, port = server.get_host_and_port() + endpoint = f"http://{host}:{port}" + + monkey = pytest.MonkeyPatch() + for k, v in { + "AWS_ACCESS_KEY_ID": "testing", "AWS_SECRET_ACCESS_KEY": "testing", + "AWS_DEFAULT_REGION": "us-east-1", "AWS_ENDPOINT_URL": endpoint, + "AWS_ENDPOINT_URL_S3": endpoint, + }.items(): + monkey.setenv(k, v) + s3fs.S3FileSystem.clear_instance_cache() + try: + yield endpoint + finally: + s3fs.S3FileSystem.clear_instance_cache() + monkey.undo() + server.stop() + + +class TestByteBudgetS3: + def test_lerobot_v3_inspect_streams_not_downloads(self, moto_s3): + import s3fs + from s3fs.core import S3File + + fs = s3fs.S3FileSystem() + fs.mkdir("stream-bucket") + # a small info.json + a deliberately "large" fake payload alongside it, + # to prove inspect never reads the bulk. + with fs.open("stream-bucket/ds/meta/info.json", "wb") as f: + f.write(json.dumps(_lerobot_v3_info()).encode()) + with fs.open("stream-bucket/ds/videos/big.bin", "wb") as f: + f.write(b"\0" * (8 * 1024 * 1024)) # 8 MB of payload + + fetched = {"n": 0} + orig = S3File._fetch_range + + def counting(self, start, end): + fetched["n"] += end - start + return orig(self, start, end) + + S3File._fetch_range = counting + try: + info = FormatRegistry.get_reader("lerobot-v3").inspect("s3://stream-bucket/ds") + finally: + S3File._fetch_range = orig + + assert info.num_episodes == 42 + # streamed: fetched only info.json-scale bytes, nowhere near the 8 MB payload + assert fetched["n"] < 1_000_000, f"expected a streaming read, got {fetched['n']} bytes" From c60bb09f0122b4642b097ffb6d51355fc8c2f8e3 Mon Sep 17 00:00:00 2001 From: Arpit Gupta Date: Sun, 19 Jul 2026 22:01:42 -0400 Subject: [PATCH 2/2] feat(io): streaming ingest + quality for cloud LeRobot-v3 & Zarr MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Extend streaming from inspect to the ingest/scoring path: `forge ingest s3://…` registers and quality-scores episodes by reading only the proprio parquet columns (state/action/timestamp) via range reads — never the videos. Ingesting a 464 MB cloud LeRobot-v3 dataset streams ~3 MB (vs 464 MB), in ~0.5s, producing identical episodes and quality scores. - lerobot-v3 reader: `_read_episodes_remote` yields proprio-only Episodes from parquet column reads (observation.state/action/timestamp), plus streamed tasks + episode metadata for language. state_dim/action_dim are left None to match the local reader so an episode's content_hash is identical whether streamed or downloaded — cross-path re-ingest stays idempotent (verified: 100/100 hashes match; streaming into a download-built catalog skips all 100). Zarr streams via its native fsspec store (already wired). - catalog/ingest.py: `_resolve_for_ingest` streams remote streamable formats instead of localizing; other sources unchanged. Embeddings/video quality still download frames. Tests: streaming ingest over memory:// (registers + scores, no temp download) and a moto byte-budget assertion (ingest reads < 2 MB with an 8 MB video payload alongside). No regressions. Co-Authored-By: Claude Opus 4.8 (1M context) --- CHANGELOG.md | 26 +++--- README.md | 4 +- forge/catalog/ingest.py | 24 +++++- forge/formats/lerobot_v3/reader.py | 129 +++++++++++++++++++++++++++++ tests/test_streaming.py | 83 +++++++++++++++++++ 5 files changed, 252 insertions(+), 14 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index bff049a..cedf16b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,16 +9,22 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added -- **Streaming metadata reads for cloud datasets (LeRobot-v3, Zarr).** - `forge inspect s3://…` / `gs://…` now reads metadata over the network with - **range requests** instead of downloading the whole dataset — inspecting a - 464 MB cloud LeRobot-v3 dataset fetches ~3 KB (its `info.json`), a ~150,000× - reduction in bytes moved, in a fraction of a second. LeRobot-v3 reads - `info.json` (which carries episode/frame counts, fps, robot, per-camera dims, - and schema); Zarr opens the store natively over fsspec. `--deep` forces a full - download for per-episode stats; non-streamable formats and other commands are - unchanged (download-to-temp). New `forge.io.DataSource` helper; the readers - keep their local paths untouched. +- **Streaming reads for cloud datasets (LeRobot-v3, Zarr).** `forge inspect` + **and `forge ingest`** on `s3://…` / `gs://…` now read over the network with + **range requests** instead of downloading the whole dataset. + - **Inspect** fetches metadata only — a 464 MB cloud LeRobot-v3 dataset yields + a full inspect from ~3 KB (its `info.json`), a ~150,000× reduction, in a + fraction of a second. + - **Ingest** (register + quality-score) reads only the **proprio parquet + columns** (state/action/timestamp) via column/range reads — **never the + videos**. Ingesting the same 464 MB dataset streams ~3 MB and produces + identical episodes and quality scores. Content hashes are identical whether + an episode is streamed or downloaded, so cross-path re-ingest is idempotent. + - LeRobot-v3 gets a proprio-only streaming `read_episodes`; Zarr streams via + its native fsspec store. `forge inspect --deep`, embeddings, and video + quality still fetch frames; other formats/commands are unchanged + (download-to-temp). New `forge.io.DataSource` helper; readers keep their + local paths untouched. - **Search- and selection-driven curation.** `forge curate` now selects episodes three ways — a SQL `--where` predicate, an explicit `--ids` list, or diff --git a/README.md b/README.md index 6d7dc8c..3d6f085 100644 --- a/README.md +++ b/README.md @@ -286,7 +286,7 @@ forge inspect s3://my-bucket/datasets/run_0413 forge convert gs://lab-data/rosbags ./out --format lerobot-v3 ``` -**Streaming inspect.** For **LeRobot-v3** and **Zarr**, `forge inspect` reads metadata over the network with **range requests** instead of downloading — e.g. inspecting a 464 MB cloud LeRobot dataset fetches ~3 KB (its `info.json`), not the whole thing. Add `--deep` to force a full download when you need per-episode stats. +**Streaming reads.** For **LeRobot-v3** and **Zarr**, `forge inspect` and `forge ingest` read over the network with **range requests** instead of downloading. Inspecting a 464 MB cloud LeRobot dataset fetches ~3 KB (its `info.json`); **ingesting it — register + quality-score every episode — reads only the proprio parquet columns (~3 MB), never the videos.** `forge inspect --deep` forces a full download when you need per-episode stats; embeddings/video quality still fetch frames. Auth uses each provider's standard credential chain (AWS env vars / profiles / IAM roles; GCP Application Default Credentials) — Forge never handles credentials itself. Catalogs can be read *and written* in the cloud; per-format conversion **outputs** are local-only for now. Full guide + connectivity troubleshooting: [forge/io/README.md](forge/io/README.md). @@ -336,7 +336,7 @@ Full reference: [docs/cli.md](docs/cli.md). Run `forge --help` for any - [x] Cloud storage — `s3://` / `gs://` on every command - [x] The catalog — ingest, SQL query, semantic search, dedup, curation, Studio - [ ] **Snapshots + export** — freeze a curated selection → LeRobot/RLDS for training -- [~] **Streaming reads** — `forge inspect` streams LeRobot-v3 / Zarr metadata (range reads, no download); ingest/scoring streaming next +- [~] **Streaming reads** — `forge inspect` + `forge ingest` stream LeRobot-v3 / Zarr (metadata + proprio via range reads, no video download); frame/embedding streaming next - [ ] **Dataset merging & splitting** — combine datasets; stratified train/val/test - [ ] **Depth / point-cloud support** · **GR00T writer** · **distributed conversion** diff --git a/forge/catalog/ingest.py b/forge/catalog/ingest.py index 2895683..4219305 100644 --- a/forge/catalog/ingest.py +++ b/forge/catalog/ingest.py @@ -60,6 +60,26 @@ def _resolve_source(uri: str) -> Path: return Path(uri) +def _resolve_for_ingest(uri: str): + """Resolve for ingest, streaming remote streamable formats instead of downloading. + + Returns a source usable by the readers — the URI itself for a remote + ``lerobot-v3``/``zarr`` dataset (metadata + proprio streamed via range + reads; no video download), or a local path otherwise. + """ + from forge.formats.registry import FormatRegistry + from forge.io import is_remote_uri + + if is_remote_uri(uri): + try: + fmt = FormatRegistry.detect_format(uri) + except Exception: + fmt = None + if fmt in FormatRegistry.STREAMABLE_FORMATS: + return uri + return _resolve_source(uri) + + def _episode_content_hash(episode: Episode, frames: list) -> str: """xxh3 hash of an episode's content: proprio stream + shape signature. @@ -245,7 +265,7 @@ def ingest( total = 0 for src in sources: try: - path = _resolve_source(src) + path = _resolve_for_ingest(src) fmt = FormatRegistry.detect_format(path) total += FormatRegistry.get_reader(fmt).inspect(path).num_episodes or 0 except Exception: @@ -284,7 +304,7 @@ def _flush() -> None: try: for src in sources: try: - path = _resolve_source(src) + path = _resolve_for_ingest(src) fmt = FormatRegistry.detect_format(path) reader = FormatRegistry.get_reader(fmt) dataset_info = None diff --git a/forge/formats/lerobot_v3/reader.py b/forge/formats/lerobot_v3/reader.py index 574c739..2c14bba 100644 --- a/forge/formats/lerobot_v3/reader.py +++ b/forge/formats/lerobot_v3/reader.py @@ -754,6 +754,12 @@ def read_episodes(self, path: Path) -> Iterator[Episode]: _check_pyarrow() import pyarrow.parquet as pq + from forge.io import is_remote_uri + + if is_remote_uri(path): + yield from self._read_episodes_remote(str(path)) + return + path = Path(path) # V3 uses data/train/ structure @@ -827,6 +833,129 @@ def read_episodes(self, path: Path) -> Iterator[Episode]: yield self._load_episode(path, pq_file, episode_id, language) + # ── Streaming (remote) episode reads ───────────────────────────────── + # Reads only the proprio columns (state/action/timestamp) from the data + # parquet via range/column reads — never touches the mp4 videos. This powers + # streaming ingest + quality (content hash + the proprio metrics) on cloud + # datasets without a full download. Embeddings/video still need the frames, + # so those paths localize; this yields proprio-only episodes. + + _PROPRIO_COLUMNS = ( + "episode_index", "index", "frame_index", "timestamp", + "observation.state", "action", + ) + + def _read_episodes_remote(self, uri: str) -> Iterator[Episode]: + import numpy as np + + from forge.io import DataSource + + src = DataSource(uri) + info = src.read_json("meta", "info.json") if src.exists("meta", "info.json") else {} + fps = info.get("fps") + tasks = self._stream_all_tasks(src) + episode_meta = self._stream_episode_metadata(src) + + files = ( + src.glob("data", "train", "**", "*.parquet") + or src.glob("data", "**", "*.parquet") + or src.glob("data", "*.parquet") + ) + yielded: set[int] = set() + for pq_file in sorted(files): + names = set(self.__parquet_names(src, pq_file)) + cols = [c for c in self._PROPRIO_COLUMNS if c in names] + if "episode_index" not in cols: + continue + df = src.read_parquet(pq_file, columns=cols).to_pandas() + for episode_idx in sorted(df["episode_index"].unique()): + episode_idx = int(episode_idx) + if episode_idx in yielded: + continue + yielded.add(episode_idx) + ep_df = df[df["episode_index"] == episode_idx] + language = self._episode_language( + episode_meta.get(episode_idx) if episode_meta else None, tasks + ) + yield self._proprio_episode(ep_df, episode_idx, fps, language, np) + + @staticmethod + def __parquet_names(src, pq_file) -> list[str]: + try: + return list(src.parquet_schema(pq_file).names) + except Exception: + return list(LeRobotV3Reader._PROPRIO_COLUMNS) + + def _proprio_episode(self, ep_df, episode_idx, fps, language, np) -> Episode: + """Build a proprio-only Episode (no images) from parquet columns.""" + + def _stack(col): + if col not in ep_df.columns: + return None + vals = ep_df[col].to_list() + return np.asarray(vals, dtype=np.float64) if len(vals) else None + + states = _stack("observation.state") + actions = _stack("action") + timestamps = ( + ep_df["timestamp"].to_numpy() if "timestamp" in ep_df.columns else None + ) + n = len(ep_df) + + def loader(): + for i in range(n): + yield Frame( + index=i, + state=states[i] if states is not None else None, + action=actions[i] if actions is not None else None, + timestamp=float(timestamps[i]) if timestamps is not None else None, + ) + + # Note: state_dim/action_dim are left None to match the local reader, + # so an episode's content_hash is identical whether streamed or + # downloaded (the hash's shape signature uses episode.state_dim). Ingest + # infers the dims from the frames for the episodes table regardless. + return Episode( + episode_id=f"episode_{episode_idx:06d}", + language_instruction=language, + fps=fps, + metadata={"num_frames": n}, + _frame_loader=loader, + ) + + def _stream_all_tasks(self, src) -> list[str]: + try: + if src.exists("meta", "tasks.parquet"): + df = src.read_parquet(src.path("meta", "tasks.parquet")).to_pandas() + col = "task" if "task" in df.columns else df.columns[-1] + return [str(t) for t in df[col].to_list()] + if src.exists("meta", "tasks.jsonl"): + return [ + json.loads(line).get("task", "") + for line in src.read_text("meta", "tasks.jsonl").splitlines() + if line.strip() + ] + except Exception: + pass + return [] + + def _stream_episode_metadata(self, src) -> dict[int, dict]: + meta: dict[int, dict] = {} + try: + for f in src.glob("meta", "episodes", "**", "*.parquet"): + df = src.read_parquet(f).to_pandas() + if "episode_index" not in df.columns: + continue + for rec in df.to_dict("records"): + meta[int(rec["episode_index"])] = rec + if not meta and src.exists("meta", "episodes.jsonl"): + for i, line in enumerate(src.read_text("meta", "episodes.jsonl").splitlines()): + if line.strip(): + meta[i] = json.loads(line) + except Exception: + pass + return meta + def _episode_language(self, meta: dict | None, tasks: list[str]) -> str | None: """Resolve one episode's language instruction from its metadata. diff --git a/tests/test_streaming.py b/tests/test_streaming.py index 31d8997..1603ddf 100644 --- a/tests/test_streaming.py +++ b/tests/test_streaming.py @@ -193,3 +193,86 @@ def counting(self, start, end): assert info.num_episodes == 42 # streamed: fetched only info.json-scale bytes, nowhere near the 8 MB payload assert fetched["n"] < 1_000_000, f"expected a streaming read, got {fetched['n']} bytes" + + +def _write_lerobot_v3(fs, root: str, *, episodes=2, frames=10) -> None: + """Write a minimal streamable v3 dataset (info.json + one data parquet).""" + import io + + import pyarrow as pa + import pyarrow.parquet as pq + + info = dict(_lerobot_v3_info()) + info["total_episodes"] = episodes + info["total_frames"] = episodes * frames + with fs.open(f"{root}/meta/info.json", "wb") as f: + f.write(json.dumps(info).encode()) + + n = episodes * frames + rng = np.random.RandomState(0) + tbl = pa.table( + { + "episode_index": [e for e in range(episodes) for _ in range(frames)], + "index": list(range(n)), + "timestamp": [float(i % frames) / 10 for i in range(n)], + "observation.state": [rng.randn(7).tolist() for _ in range(n)], + "action": [rng.randn(7).tolist() for _ in range(n)], + } + ) + buf = io.BytesIO() + pq.write_table(tbl, buf) + with fs.open(f"{root}/data/chunk-000/file-000.parquet", "wb") as f: + f.write(buf.getvalue()) + + +class TestStreamingIngest: + def test_ingest_streams_over_memory(self, memory_fs, tmp_path): + pytest.importorskip("duckdb") + from forge.catalog import Catalog + from forge.catalog.ingest import ingest as ingest_datasets + + _write_lerobot_v3(memory_fs, "/ds", episodes=3, frames=8) + cat = Catalog.init(str(tmp_path / "cat")) + + before = len(_TEMP_ROOTS) + stats = ingest_datasets(["memory:///ds"], cat) + assert stats.ingested == 3 and stats.failed == 0 + assert len(_TEMP_ROOTS) == before # streamed — nothing downloaded + # registered + quality-scored, all from parquet proprio columns + assert cat.sql("SELECT count(*) n FROM episodes").to_pylist()[0]["n"] == 3 + scored = cat.sql( + "SELECT count(*) n FROM v_latest_quality WHERE overall_score IS NOT NULL" + ).to_pylist()[0]["n"] + assert scored == 3 + + def test_ingest_byte_budget_ignores_video(self, moto_s3, tmp_path): + pytest.importorskip("duckdb") + import s3fs + from s3fs.core import S3File + + from forge.catalog import Catalog + from forge.catalog.ingest import ingest as ingest_datasets + + fs = s3fs.S3FileSystem() + fs.mkdir("ingest-bucket") + _write_lerobot_v3(fs, "ingest-bucket/ds", episodes=2, frames=10) + # a big "video" payload the proprio ingest must not read + with fs.open("ingest-bucket/ds/videos/cam/big.mp4", "wb") as f: + f.write(b"\0" * (8 * 1024 * 1024)) + + fetched = {"n": 0} + orig = S3File._fetch_range + + def counting(self, start, end): + fetched["n"] += end - start + return orig(self, start, end) + + S3File._fetch_range = counting + try: + cat = Catalog.init(str(tmp_path / "cat")) + stats = ingest_datasets(["s3://ingest-bucket/ds"], cat) + finally: + S3File._fetch_range = orig + + assert stats.ingested == 2 + assert fetched["n"] < 2_000_000, f"ingest should stream proprio, got {fetched['n']} bytes"