From 79cfb52a71519533ded84b36619294e741f9ec04 Mon Sep 17 00:00:00 2001 From: Pepijn Date: Mon, 27 Jul 2026 16:07:48 +0200 Subject: [PATCH] Harden production dataset streaming pipeline --- docs/source/lerobot-dataset-v3.mdx | 4 +- docs/source/training_dataset_streaming.mdx | 22 ++- scripts/bench_episode_byte_cache.py | 7 +- scripts/bench_streaming_dataset.py | 15 ++ src/lerobot/configs/default.py | 18 ++ src/lerobot/datasets/factory.py | 10 ++ src/lerobot/datasets/streaming_dataset.py | 157 +++++++++++++----- src/lerobot/streaming/episode_cache.py | 29 +++- src/lerobot/streaming/episode_parquet.py | 28 +++- tests/configs/test_default.py | 5 + tests/datasets/test_episode_parquet_reader.py | 30 ++++ .../datasets/test_episode_video_streaming.py | 61 +++++++ tests/datasets/test_streaming_factory.py | 10 ++ tests/datasets/test_streaming_production.py | 57 +++++++ 14 files changed, 397 insertions(+), 56 deletions(-) diff --git a/docs/source/lerobot-dataset-v3.mdx b/docs/source/lerobot-dataset-v3.mdx index 0f672c255..b63b2a08d 100644 --- a/docs/source/lerobot-dataset-v3.mdx +++ b/docs/source/lerobot-dataset-v3.mdx @@ -178,7 +178,9 @@ To capture comparable data-pipeline results on a training host: uv run python scripts/bench_streaming_dataset.py \ --repo-id=yaak-ai/L2D-v3 \ --batch-size=16 \ - --num-workers=4 \ + --num-workers=1 \ + --fetch-workers=4 \ + --decode-threads=2 \ --summary-json=streaming-benchmark.json ``` diff --git a/docs/source/training_dataset_streaming.mdx b/docs/source/training_dataset_streaming.mdx index 477dcb354..7288637eb 100644 --- a/docs/source/training_dataset_streaming.mdx +++ b/docs/source/training_dataset_streaming.mdx @@ -58,12 +58,17 @@ Subset sidecars cannot be published. The production defaults are: -| Option | Default | Meaning | -| ----------------------------- | ------: | --------------------------------------------------- | -| `streaming_episode_pool_size` | 32 | Maximum complete episodes mixed by each rank | -| `streaming_prefetch_episodes` | 8 | Episodes fetched ahead of the active pool | -| `streaming_byte_budget_gb` | 8 | Maximum synthesized MP4 bytes per rank | -| `streaming_data_root` | unset | Optional local, Hub, bucket, or fsspec payload root | +| Option | Default | Meaning | +| ----------------------------------- | ------: | --------------------------------------------------- | +| `streaming_episode_pool_size` | 32 | Maximum complete episodes mixed by each rank | +| `streaming_prefetch_episodes` | 8 | Episodes fetched ahead of the active pool | +| `streaming_byte_budget_gb` | 8 | Maximum synthesized MP4 bytes per rank | +| `streaming_decode_threads` | 2 | Parallel sample assembly and video decode workers | +| `streaming_decoded_queue_size` | 8 | Decoded samples buffered ahead, in planner order | +| `streaming_max_open_decoders` | 64 | Independent open-decoder LRU cap per rank | +| `streaming_native_http_connections` | unset | Native HTTP connection cap per rank | +| `streaming_native_http_subranges` | 1 | Concurrent subranges per native HTTP range read | +| `streaming_data_root` | unset | Optional local, Hub, bucket, or fsspec payload root | The active episode set is capped by both episode count and the exact synthesized mini-MP4 sizes computed from the sidecar. An episode larger than the complete rank budget fails before training @@ -77,6 +82,9 @@ lerobot-train \ --dataset.streaming_episode_pool_size=16 \ --dataset.streaming_prefetch_episodes=4 \ --dataset.streaming_byte_budget_gb=4 \ + --dataset.streaming_decode_threads=2 \ + --dataset.streaming_decoded_queue_size=8 \ + --dataset.streaming_max_open_decoders=64 \ --num_workers=4 \ --policy.type=act \ --output_dir=outputs/train/act_streaming @@ -109,6 +117,8 @@ uv run python scripts/bench_streaming_dataset.py \ --batch-size=16 \ --num-workers=1 \ --fetch-workers=4 \ + --decode-threads=2 \ + --decoded-queue-size=8 \ --summary-json=streaming-benchmark.json ``` diff --git a/scripts/bench_episode_byte_cache.py b/scripts/bench_episode_byte_cache.py index 9000fa3f2..73ceb5468 100644 --- a/scripts/bench_episode_byte_cache.py +++ b/scripts/bench_episode_byte_cache.py @@ -70,7 +70,7 @@ def parse_args() -> argparse.Namespace: type=int, default=8, help="Concurrent camera-fetch jobs. Total connections ~= workers x range-subranges; " - "the HF bucket path saturates around 64 connections per host, so keep the product near 64.", + "measure the host-specific saturation point before increasing this value.", ) parser.add_argument( "--range-subranges", @@ -443,6 +443,7 @@ def run_exact_coverage_stream( decode_pool.shutdown(wait=True) elapsed = time.perf_counter() - start + replacements = max(0, pool.admitted_count - min(pool_size, len(order))) result = { "target_samples_s": target_samples_s, "actual_samples_s": samples_done / elapsed if elapsed > 0 else float("inf"), @@ -453,6 +454,10 @@ def run_exact_coverage_stream( "shard_total_frames": float(total_frames), "epoch_complete": 1.0 if epoch_complete else 0.0, "prefetch_ahead": float(prefetch_ahead), + "prefetch_episodes": float(prefetch_ahead), + "replacements": float(replacements), + "replacement_episodes_s": replacements / elapsed if elapsed > 0 else float("inf"), + "samples_per_episode": samples_done / replacements if replacements else float(samples_done), "batch_size": float(batch_size), "decode_workers": float(decode_workers), "kept_up": 1.0 diff --git a/scripts/bench_streaming_dataset.py b/scripts/bench_streaming_dataset.py index 3f163c629..c19d99727 100644 --- a/scripts/bench_streaming_dataset.py +++ b/scripts/bench_streaming_dataset.py @@ -50,6 +50,11 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--episode-pool-size", type=int, default=32) parser.add_argument("--prefetch-episodes", type=int, default=8) parser.add_argument("--byte-budget-gb", type=float, default=8.0) + parser.add_argument("--decode-threads", type=int, default=2) + parser.add_argument("--decoded-queue-size", type=int, default=8) + parser.add_argument("--max-open-decoders", type=int, default=64) + parser.add_argument("--native-http-connections", type=int, default=None) + parser.add_argument("--native-http-subranges", type=int, default=1) parser.add_argument("--warmup-batches", type=int, default=8) parser.add_argument("--measure-batches", type=int, default=128) parser.add_argument("--summary-json", type=Path, default=None) @@ -105,6 +110,11 @@ def main() -> None: episode_pool_size=args.episode_pool_size, prefetch_episodes=args.prefetch_episodes, byte_budget_gb=args.byte_budget_gb, + decode_threads=args.decode_threads, + decoded_queue_size=args.decoded_queue_size, + max_open_decoders=args.max_open_decoders, + native_http_connections=args.native_http_connections, + native_http_subranges=args.native_http_subranges, max_num_shards=max(1, args.fetch_workers), return_uint8=True, ) @@ -184,6 +194,11 @@ def main() -> None: "episode_pool_size": args.episode_pool_size, "prefetch_episodes": args.prefetch_episodes, "byte_budget_gb": args.byte_budget_gb, + "decode_threads": args.decode_threads, + "decoded_queue_size": args.decoded_queue_size, + "max_open_decoders": args.max_open_decoders, + "native_http_connections": args.native_http_connections, + "native_http_subranges": args.native_http_subranges, "warmup_batches": args.warmup_batches, "measure_batches": args.measure_batches, }, diff --git a/src/lerobot/configs/default.py b/src/lerobot/configs/default.py index cf3de3c85..9ad8f047e 100644 --- a/src/lerobot/configs/default.py +++ b/src/lerobot/configs/default.py @@ -52,6 +52,14 @@ class DatasetConfig: streaming_prefetch_episodes: int = 8 # Hard per-rank cap for synthesized episode-video bytes. streaming_byte_budget_gb: float = 8.0 + # Parallel sample assembly/decode workers and their bounded in-order result queue. + streaming_decode_threads: int = 2 + streaming_decoded_queue_size: int = 8 + # Independent decoder-state cap. + streaming_max_open_decoders: int = 64 + # Per-rank native HTTP limits. None preserves the fetcher's worker-derived default. + streaming_native_http_connections: int | None = None + streaming_native_http_subranges: int = 1 # Fraction of episodes held out per task for offline evaluation (0.0 = disabled). eval_split: float = 0.0 @@ -68,6 +76,16 @@ class DatasetConfig: raise ValueError("streaming_prefetch_episodes must be non-negative") if self.streaming_byte_budget_gb <= 0: raise ValueError("streaming_byte_budget_gb must be positive") + if self.streaming_decode_threads <= 0: + raise ValueError("streaming_decode_threads must be positive") + if self.streaming_decoded_queue_size <= 0: + raise ValueError("streaming_decoded_queue_size must be positive") + if self.streaming_max_open_decoders <= 0: + raise ValueError("streaming_max_open_decoders must be positive") + if self.streaming_native_http_connections is not None and self.streaming_native_http_connections <= 0: + raise ValueError("streaming_native_http_connections must be positive") + if self.streaming_native_http_subranges <= 0: + raise ValueError("streaming_native_http_subranges must be positive") if self.episodes is not None: if any(ep < 0 for ep in self.episodes): raise ValueError( diff --git a/src/lerobot/datasets/factory.py b/src/lerobot/datasets/factory.py index fb8b79382..f027f4cbb 100644 --- a/src/lerobot/datasets/factory.py +++ b/src/lerobot/datasets/factory.py @@ -119,6 +119,11 @@ def make_dataset( episode_pool_size=cfg.dataset.streaming_episode_pool_size, prefetch_episodes=cfg.dataset.streaming_prefetch_episodes, byte_budget_gb=cfg.dataset.streaming_byte_budget_gb, + decode_threads=cfg.dataset.streaming_decode_threads, + decoded_queue_size=cfg.dataset.streaming_decoded_queue_size, + max_open_decoders=cfg.dataset.streaming_max_open_decoders, + native_http_connections=cfg.dataset.streaming_native_http_connections, + native_http_subranges=cfg.dataset.streaming_native_http_subranges, repeat=True, ) else: @@ -210,6 +215,11 @@ def make_train_eval_datasets( episode_pool_size=cfg.dataset.streaming_episode_pool_size, prefetch_episodes=cfg.dataset.streaming_prefetch_episodes, byte_budget_gb=cfg.dataset.streaming_byte_budget_gb, + decode_threads=cfg.dataset.streaming_decode_threads, + decoded_queue_size=cfg.dataset.streaming_decoded_queue_size, + max_open_decoders=cfg.dataset.streaming_max_open_decoders, + native_http_connections=cfg.dataset.streaming_native_http_connections, + native_http_subranges=cfg.dataset.streaming_native_http_subranges, repeat=True, ) else: diff --git a/src/lerobot/datasets/streaming_dataset.py b/src/lerobot/datasets/streaming_dataset.py index be5d1cb40..f9f1bee2f 100644 --- a/src/lerobot/datasets/streaming_dataset.py +++ b/src/lerobot/datasets/streaming_dataset.py @@ -15,6 +15,7 @@ # limitations under the License. import io import os +from collections import deque from collections.abc import Callable, Iterator, Mapping from concurrent.futures import Future, ThreadPoolExecutor from pathlib import Path @@ -122,6 +123,11 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset): repeat: bool = False, *, token: str | bool | None = None, + decode_threads: int = 2, + decoded_queue_size: int = 8, + max_open_decoders: int = 64, + native_http_connections: int | None = None, + native_http_subranges: int = 1, ): """Initialize a StreamingLeRobotDataset. @@ -152,6 +158,13 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset): episode_pool_size (int | None, optional): Number of complete episodes in the sampling pool. prefetch_episodes (int, optional): Episodes prefetched beyond the active pool. byte_budget_gb (float, optional): Per-rank upper bound for synthesized episode video bytes. + decode_threads (int, optional): Parallel sample assembly and video decode workers. + decoded_queue_size (int, optional): Maximum number of decoded samples produced ahead + of the consumer. Results are yielded in exact planner order. + max_open_decoders (int, optional): Maximum number of open video decoders per rank. + native_http_connections (int | None, optional): Native HTTP connection limit per rank. + ``None`` preserves the fetcher's worker-derived default. + native_http_subranges (int, optional): Concurrent HTTP subranges per video range read. repeat (bool, optional): Repeat rank-local exact-coverage epochs without yielding a short final training batch. The training factory enables this; direct iteration is finite by default. @@ -196,9 +209,24 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset): raise ValueError("prefetch_episodes must be non-negative") if byte_budget_gb <= 0: raise ValueError("byte_budget_gb must be positive") + if decode_threads <= 0: + raise ValueError("decode_threads must be positive") + if decoded_queue_size <= 0: + raise ValueError("decoded_queue_size must be positive") + if max_open_decoders <= 0: + raise ValueError("max_open_decoders must be positive") + if native_http_connections is not None and native_http_connections <= 0: + raise ValueError("native_http_connections must be positive") + if native_http_subranges <= 0: + raise ValueError("native_http_subranges must be positive") self.episode_pool_size = episode_pool_size or min(buffer_size, 32) self.prefetch_episodes = prefetch_episodes self.byte_budget = int(byte_budget_gb * 1024**3) + self.decode_threads = decode_threads + self.decoded_queue_size = decoded_queue_size + self.max_open_decoders = max_open_decoders + self.native_http_connections = native_http_connections + self.native_http_subranges = native_http_subranges self.repeat = repeat self._next_epoch = 0 self._active_epoch = 0 @@ -356,8 +384,13 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset): columns=self._projected_columns, token=self._streaming_io_token, ) - executor = ThreadPoolExecutor(max_workers=max_workers, thread_name_prefix="lerobot-parquet") + parquet_executor = ThreadPoolExecutor(max_workers=max_workers, thread_name_prefix="lerobot-parquet") + decode_executor = ThreadPoolExecutor( + max_workers=self.decode_threads, + thread_name_prefix="lerobot-decode", + ) episode_futures: dict[int, Future[datasets.Dataset]] = {} + decoded_futures: deque[Future[dict]] = deque() scheduled_episodes: set[int] = set() retained_video_episodes: set[int] = set() if video_cache is not None: @@ -368,7 +401,7 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset): def submit(episode_index: int) -> Future[datasets.Dataset]: future = episode_futures.get(episode_index) if future is None: - future = executor.submit(self._load_episode_dataset, parquet_reader, episode_index) + future = parquet_executor.submit(self._load_episode_dataset, parquet_reader, episode_index) episode_futures[episode_index] = future return future @@ -382,44 +415,83 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset): video_cache.submit_prefetch(episode_index) scheduled_episodes.add(episode_index) - schedule_frontier() + def decode_item( + episode_future: Future[datasets.Dataset], + episode_index: int, + frame_index: int, + ) -> dict: + return self._make_episode_item( + episode_future.result(), + episode_index, + frame_index, + video_cache=video_cache, + apply_image_transforms=False, + ) + + def update_frontier() -> None: + for evicted_episode in planner.evicted: + episode_futures.pop(evicted_episode, None) + if video_cache is not None and evicted_episode in retained_video_episodes: + video_cache.release_episode(evicted_episode) + retained_video_episodes.remove(evicted_episode) + if video_cache is not None: + for admitted_episode in planner.newly_admitted: + if admitted_episode not in retained_video_episodes: + video_cache.retain_episode(admitted_episode) + retained_video_episodes.add(admitted_episode) + planner.evicted.clear() + planner.newly_admitted.clear() + schedule_frontier() + try: - while True: - try: - episode_index, frame_index = next(planner) - except StopIteration: - self._active_epoch = epoch + 1 if self.shuffle else 0 - self._state_offset = 0 - break + schedule_frontier() + planner_exhausted = False + while decoded_futures or not planner_exhausted: + while not planner_exhausted and len(decoded_futures) < self.decoded_queue_size: + try: + episode_index, frame_index = next(planner) + except StopIteration: + planner_exhausted = True + break - episode_dataset = submit(episode_index).result() - item = self._make_episode_item( - episode_dataset, - episode_index, - frame_index, - video_cache=video_cache, - ) - - for evicted_episode in planner.evicted: - episode_futures.pop(evicted_episode, None) - if video_cache is not None and evicted_episode in retained_video_episodes: - video_cache.release_episode(evicted_episode) - retained_video_episodes.remove(evicted_episode) - if video_cache is not None: - for admitted_episode in planner.newly_admitted: - if admitted_episode not in retained_video_episodes: - video_cache.retain_episode(admitted_episode) - retained_video_episodes.add(admitted_episode) - planner.evicted.clear() - planner.newly_admitted.clear() - schedule_frontier() + episode_future = submit(episode_index) + if video_cache is not None: + video_cache.retain_episode(episode_index) + try: + decoded_future = decode_executor.submit( + decode_item, + episode_future, + episode_index, + frame_index, + ) + except Exception: + if video_cache is not None: + video_cache.release_episode(episode_index) + raise + if video_cache is not None: + decoded_future.add_done_callback( + lambda _future, retained_episode=episode_index, retained_cache=video_cache: ( + retained_cache.release_episode(retained_episode) + ) + ) + decoded_futures.append(decoded_future) + update_frontier() + if not decoded_futures: + continue + item = decoded_futures.popleft().result() + self._apply_image_transforms(item) self._state_offset += 1 yield item + self._active_epoch = epoch + 1 if self.shuffle else 0 + self._state_offset = 0 finally: + for future in decoded_futures: + future.cancel() + decode_executor.shutdown(wait=True, cancel_futures=True) for future in episode_futures.values(): future.cancel() - executor.shutdown(wait=True, cancel_futures=True) + parquet_executor.shutdown(wait=True, cancel_futures=True) if video_cache is not None: video_cache.close() @@ -483,14 +555,15 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset): sidecar_path=self._sidecar_path, token=self._streaming_io_token, ) - decoder_limit = max(1, min(64, self.episode_pool_size * max(1, len(self.meta.video_keys)))) return EpisodeByteCache( manifest, self._data_root, byte_budget=self.byte_budget, workers=workers, range_backend=range_backend, - max_open_decoders=decoder_limit, + native_http_connections=self.native_http_connections, + native_http_subranges=self.native_http_subranges, + max_open_decoders=self.max_open_decoders, video_backend=self._video_backend, tolerance_s=self.tolerance_s, token=self._streaming_io_token, @@ -503,6 +576,7 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset): frame_index: int, *, video_cache: EpisodeByteCache | None, + apply_image_transforms: bool = True, ) -> dict: item = episode_dataset[frame_index] episode = self.meta.episodes[episode_index] @@ -563,11 +637,8 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset): frames = frames.to(torch.float32) / 255.0 item[video_key] = frames.squeeze(0) - if self.image_transforms is not None: - for camera_key in self.meta.camera_keys: - if camera_key in self.meta.depth_keys: - continue - item[camera_key] = self.image_transforms(item[camera_key]) + if apply_image_transforms: + self._apply_image_transforms(item) for key, stored_unit in self._image_depth_units.items(): if key in item and stored_unit is not None and stored_unit != self._depth_output_unit: @@ -585,6 +656,14 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset): ) return item + def _apply_image_transforms(self, item: dict) -> None: + if self.image_transforms is None: + return + for camera_key in self.meta.camera_keys: + if camera_key in self.meta.depth_keys: + continue + item[camera_key] = self.image_transforms(item[camera_key]) + def state_dict(self) -> dict[str, int]: return { "epoch": self._active_epoch, diff --git a/src/lerobot/streaming/episode_cache.py b/src/lerobot/streaming/episode_cache.py index 78dad4490..516bb1677 100644 --- a/src/lerobot/streaming/episode_cache.py +++ b/src/lerobot/streaming/episode_cache.py @@ -75,6 +75,7 @@ class EpisodeByteCache: self._pool = ThreadPoolExecutor(max_workers=workers) self._cache: OrderedDict[tuple[int, str], dict[str, Any]] = OrderedDict() self._decoders: OrderedDict[tuple[int, str], Any] = OrderedDict() + self._decoder_locks: dict[tuple[int, str], threading.Lock] = {} self._futures: dict[tuple[int, str], Future[dict[str, Any]]] = {} self._retained_episodes: dict[int, int] = {} self._decoder_fallback_count = 0 @@ -96,6 +97,7 @@ class EpisodeByteCache: decoders = list(self._decoders.values()) self._cache.clear() self._decoders.clear() + self._decoder_locks.clear() self._futures.clear() self._retained_episodes.clear() self._fallback_decoders.clear() @@ -186,6 +188,7 @@ class EpisodeByteCache: self._decoders[key] = decoder while len(self._decoders) > self.max_open_decoders: evicted_key, evicted_decoder = self._decoders.popitem(last=False) + self._decoder_locks.pop(evicted_key, None) self._fallback_decoders.discard(evicted_key) _close_decoder(evicted_decoder) return decoder @@ -227,15 +230,23 @@ class EpisodeByteCache: decoder, release = self._decoder_for_frames(episode_index, camera_key) with self._lock: uses_pyav_timestamps = self.video_backend == "pyav" and key not in self._fallback_decoders + decode_lock = self._decoder_locks.setdefault(key, threading.Lock()) try: - if uses_pyav_timestamps: - return decoder.get_frames_played_at(local_ts, tolerance_s=self.tolerance_s).data - metadata = decoder.metadata - fps = getattr(metadata, "average_fps", None) - if fps is None: - duration = max(getattr(metadata, "end_stream_seconds", 0.0), 1e-9) - fps = metadata.num_frames / duration - return decoder.get_frames_at(indices=[round(ts * fps) for ts in local_ts]).data + with decode_lock: + if uses_pyav_timestamps: + return decoder.get_frames_played_at(local_ts, tolerance_s=self.tolerance_s).data + metadata = decoder.metadata + fps = getattr(metadata, "average_fps", None) + if fps is None: + duration = max(getattr(metadata, "end_stream_seconds", 0.0), 1e-9) + fps = metadata.num_frames / duration + num_frames = getattr(metadata, "num_frames", None) + if num_frames is None: + duration = max(getattr(metadata, "end_stream_seconds", 0.0), 1e-9) + num_frames = round(duration * fps) + last_index = max(0, int(num_frames) - 1) + indices = [min(max(round(ts * fps), 0), last_index) for ts in local_ts] + return decoder.get_frames_at(indices=indices).data finally: if release is not None: release() @@ -257,6 +268,7 @@ class EpisodeByteCache: with self._lock: if self._decoders.get(key) is decoder: self._decoders.pop(key) + self._decoder_locks.pop(key, None) self._fallback_decoders.discard(key) continue return decoder, decoder.release @@ -330,6 +342,7 @@ class EpisodeByteCache: entry = self._cache.pop(key) self._bytes -= len(entry["bytes"]) decoder = self._decoders.pop(key, None) + self._decoder_locks.pop(key, None) self._fallback_decoders.discard(key) if decoder is not None: _close_decoder(decoder) diff --git a/src/lerobot/streaming/episode_parquet.py b/src/lerobot/streaming/episode_parquet.py index c01ac9494..36996f84a 100644 --- a/src/lerobot/streaming/episode_parquet.py +++ b/src/lerobot/streaming/episode_parquet.py @@ -11,8 +11,11 @@ from __future__ import annotations import posixpath +import threading +import time from collections.abc import Sequence from pathlib import Path +from typing import Any import fsspec import pyarrow as pa @@ -29,6 +32,8 @@ class EpisodeParquetReader: *, columns: Sequence[str], token: str | bool | None = None, + max_retries: int = 4, + retry_backoff_s: float = 0.05, ): if not columns: raise ValueError("EpisodeParquetReader requires at least one projected column") @@ -37,6 +42,13 @@ class EpisodeParquetReader: self.columns if "episode_index" in self.columns else (*self.columns, "episode_index") ) data_root_str = str(data_root) + if max_retries < 0: + raise ValueError("max_retries must be non-negative") + if retry_backoff_s < 0: + raise ValueError("retry_backoff_s must be non-negative") + self._max_retries = max_retries + self._retry_backoff_s = retry_backoff_s + self._open_lock = threading.Lock() if data_root_str.startswith("hf://") else None storage_options = {"token": token} if token is not None and data_root_str.startswith("hf://") else {} self._filesystem, self._root_path = fsspec.core.url_to_fs(data_root_str, **storage_options) @@ -51,7 +63,7 @@ class EpisodeParquetReader: raise ValueError(f"Episode {episode_index} must contain at least one row") path = posixpath.join(self._root_path.rstrip("/"), str(relative_path).lstrip("/")) - with self._filesystem.open(path, "rb") as source: + with self._open_with_retry(path) as source: parquet = pq.ParquetFile(source) available = set(parquet.schema_arrow.names) missing = sorted(set(self._read_columns) - available) @@ -71,6 +83,20 @@ class EpisodeParquetReader: table = table.drop_columns(["episode_index"]) return table + def _open_with_retry(self, path: str) -> Any: + for attempt in range(self._max_retries + 1): + try: + if self._open_lock is None: + return self._filesystem.open(path, "rb") + with self._open_lock: + return self._filesystem.open(path, "rb") + except FileNotFoundError: + if attempt == self._max_retries: + raise + self._filesystem.invalidate_cache(posixpath.dirname(path)) + time.sleep(self._retry_backoff_s * 2**attempt) + raise RuntimeError("unreachable") + @staticmethod def _matching_row_group(parquet: pq.ParquetFile, episode_index: int) -> int | None: episode_column = next( diff --git a/tests/configs/test_default.py b/tests/configs/test_default.py index 57945b75e..b5c4846c8 100644 --- a/tests/configs/test_default.py +++ b/tests/configs/test_default.py @@ -44,6 +44,11 @@ def test_dataset_config_empty_episodes_ok(): ("streaming_episode_pool_size", 0, "episode_pool_size"), ("streaming_prefetch_episodes", -1, "prefetch_episodes"), ("streaming_byte_budget_gb", 0, "byte_budget_gb"), + ("streaming_decode_threads", 0, "decode_threads"), + ("streaming_decoded_queue_size", 0, "decoded_queue_size"), + ("streaming_max_open_decoders", 0, "max_open_decoders"), + ("streaming_native_http_connections", 0, "native_http_connections"), + ("streaming_native_http_subranges", 0, "native_http_subranges"), ], ) def test_dataset_config_rejects_invalid_streaming_resource_limits(field, value, message): diff --git a/tests/datasets/test_episode_parquet_reader.py b/tests/datasets/test_episode_parquet_reader.py index 7649bc737..e873ee861 100644 --- a/tests/datasets/test_episode_parquet_reader.py +++ b/tests/datasets/test_episode_parquet_reader.py @@ -122,3 +122,33 @@ def test_reader_forwards_explicit_token_to_hf_filesystem(monkeypatch) -> None: ) url_to_fs.assert_called_once_with("hf://datasets/private@revision", token="hf_test_token") + + +def test_reader_retries_transient_remote_file_not_found(tmp_path: Path, monkeypatch) -> None: + path = tmp_path / "episode.parquet" + pq.write_table(_table([0, 0]), path) + filesystem = Mock() + attempts = 0 + + def open_remote(*_args, **_kwargs): + nonlocal attempts + attempts += 1 + if attempts == 1: + raise FileNotFoundError("partial HfFileSystem dircache") + return path.open("rb") + + filesystem.open.side_effect = open_remote + filesystem.invalidate_cache = Mock() + monkeypatch.setattr(fsspec.core, "url_to_fs", Mock(return_value=(filesystem, "root"))) + reader = EpisodeParquetReader( + "hf://datasets/private@revision", + columns=("episode_index", "frame_index"), + max_retries=1, + retry_backoff_s=0, + ) + + table = reader.read_episode("data/file.parquet", episode_index=0, expected_rows=2) + + assert len(table) == 2 + assert attempts == 2 + filesystem.invalidate_cache.assert_called_once() diff --git a/tests/datasets/test_episode_video_streaming.py b/tests/datasets/test_episode_video_streaming.py index 38e5ef056..ba872734d 100644 --- a/tests/datasets/test_episode_video_streaming.py +++ b/tests/datasets/test_episode_video_streaming.py @@ -11,6 +11,7 @@ import json import struct import threading +import time from concurrent.futures import ThreadPoolExecutor import numpy as np @@ -217,6 +218,66 @@ def test_decoder_falls_back_to_pyav_when_torchcodec_rejects_mini_mp4(monkeypatch assert cache.decoder_fallback_count == 1 +def test_torchcodec_frame_indices_are_clamped_to_decoder_bounds(monkeypatch, tmp_path): + requested_indices = [] + + class FakeDecoder: + metadata = type("Metadata", (), {"average_fps": 30.0, "num_frames": 10})() + + def get_frames_at(self, *, indices): + requested_indices.extend(indices) + return type("Frames", (), {"data": indices})() + + with _fake_cache(monkeypatch, tmp_path) as cache: + monkeypatch.setattr( + cache.manifest, + "lookup", + lambda *_args: type("Span", (), {"source_start_pts": 0.0})(), + ) + monkeypatch.setattr(cache, "_decoder_for_frames", lambda *_args: (FakeDecoder(), None)) + + cache.get_frames(0, "camera", [-0.1, 1.0]) + + assert requested_indices == [0, 9] + + +def test_frame_reads_serialize_access_to_each_decoder(monkeypatch, tmp_path): + state_lock = threading.Lock() + active = 0 + max_active = 0 + + class FakeDecoder: + metadata = type("Metadata", (), {"average_fps": 30.0, "num_frames": 10})() + + def get_frames_at(self, *, indices): + nonlocal active, max_active + with state_lock: + active += 1 + max_active = max(max_active, active) + try: + time.sleep(0.01) + return type("Frames", (), {"data": indices})() + finally: + with state_lock: + active -= 1 + + with _fake_cache(monkeypatch, tmp_path) as cache: + monkeypatch.setattr( + cache.manifest, + "lookup", + lambda *_args: type("Span", (), {"source_start_pts": 0.0})(), + ) + decoder = FakeDecoder() + monkeypatch.setattr(cache, "_decoder_for_frames", lambda *_args: (decoder, None)) + + with ThreadPoolExecutor(max_workers=2) as executor: + futures = [executor.submit(cache.get_frames, 0, "camera", [0.1]) for _ in range(2)] + for future in futures: + future.result() + + assert max_active == 1 + + def test_releasing_episode_allows_immediate_eviction(monkeypatch, tmp_path): with _fake_cache(monkeypatch, tmp_path, byte_budget=5) as cache: cache.retain_episode(0) diff --git a/tests/datasets/test_streaming_factory.py b/tests/datasets/test_streaming_factory.py index bb4d4c677..47a85a287 100644 --- a/tests/datasets/test_streaming_factory.py +++ b/tests/datasets/test_streaming_factory.py @@ -38,6 +38,11 @@ def test_factory_wires_production_streaming_settings(monkeypatch): streaming_episode_pool_size=7, streaming_prefetch_episodes=3, streaming_byte_budget_gb=2.5, + streaming_decode_threads=2, + streaming_decoded_queue_size=5, + streaming_max_open_decoders=17, + streaming_native_http_connections=9, + streaming_native_http_subranges=3, ) cfg = SimpleNamespace( dataset=dataset_config, @@ -54,6 +59,11 @@ def test_factory_wires_production_streaming_settings(monkeypatch): assert captured["kwargs"]["episode_pool_size"] == 7 assert captured["kwargs"]["prefetch_episodes"] == 3 assert captured["kwargs"]["byte_budget_gb"] == 2.5 + assert captured["kwargs"]["decode_threads"] == 2 + assert captured["kwargs"]["decoded_queue_size"] == 5 + assert captured["kwargs"]["max_open_decoders"] == 17 + assert captured["kwargs"]["native_http_connections"] == 9 + assert captured["kwargs"]["native_http_subranges"] == 3 assert captured["kwargs"]["max_num_shards"] == 1 assert captured["kwargs"]["video_backend"] == "pyav" assert captured["kwargs"]["return_uint8"] is True diff --git a/tests/datasets/test_streaming_production.py b/tests/datasets/test_streaming_production.py index 2ee4a4338..e876dd9f1 100644 --- a/tests/datasets/test_streaming_production.py +++ b/tests/datasets/test_streaming_production.py @@ -10,6 +10,8 @@ from __future__ import annotations +import threading +import time from itertools import islice from pathlib import Path @@ -61,6 +63,61 @@ def test_streaming_matches_map_style_with_exact_coverage(tmp_path: Path, lerobot _assert_item_equal(sample, map_dataset[int(sample["index"])]) +def test_parallel_decode_queue_preserves_planner_order( + tmp_path: Path, + lerobot_dataset_factory, + monkeypatch, +) -> None: + root = tmp_path / "dataset" + lerobot_dataset_factory( + root=root, + repo_id=DUMMY_REPO_ID, + total_episodes=4, + total_frames=40, + use_videos=False, + ) + expected = _indices( + StreamingLeRobotDataset( + DUMMY_REPO_ID, + root=root, + seed=17, + buffer_size=3, + decode_threads=1, + decoded_queue_size=1, + ) + ) + parallel = StreamingLeRobotDataset( + DUMMY_REPO_ID, + root=root, + seed=17, + buffer_size=3, + decode_threads=3, + decoded_queue_size=5, + ) + original_make_item = parallel._make_episode_item + state_lock = threading.Lock() + active = 0 + max_active = 0 + + def delayed_make_item(*args, **kwargs): + nonlocal active, max_active + with state_lock: + active += 1 + max_active = max(max_active, active) + try: + frame_index = int(args[2]) + time.sleep(0.005 if frame_index % 3 == 0 else 0.001) + return original_make_item(*args, **kwargs) + finally: + with state_lock: + active -= 1 + + monkeypatch.setattr(parallel, "_make_episode_item", delayed_make_item) + + assert _indices(parallel) == expected + assert 1 < max_active <= parallel.decode_threads + + @pytest.mark.parametrize("video_backend", ["torchcodec", "pyav"]) def test_streaming_rgb_video_matches_map_style( tmp_path: Path,