Harden production dataset streaming pipeline

This commit is contained in:
Pepijn
2026-07-27 16:07:48 +02:00
parent c0e9b0bbff
commit 79cfb52a71
14 changed files with 397 additions and 56 deletions
+3 -1
View File
@@ -178,7 +178,9 @@ To capture comparable data-pipeline results on a training host:
uv run python scripts/bench_streaming_dataset.py \ uv run python scripts/bench_streaming_dataset.py \
--repo-id=yaak-ai/L2D-v3 \ --repo-id=yaak-ai/L2D-v3 \
--batch-size=16 \ --batch-size=16 \
--num-workers=4 \ --num-workers=1 \
--fetch-workers=4 \
--decode-threads=2 \
--summary-json=streaming-benchmark.json --summary-json=streaming-benchmark.json
``` ```
+16 -6
View File
@@ -58,12 +58,17 @@ Subset sidecars cannot be published.
The production defaults are: The production defaults are:
| Option | Default | Meaning | | Option | Default | Meaning |
| ----------------------------- | ------: | --------------------------------------------------- | | ----------------------------------- | ------: | --------------------------------------------------- |
| `streaming_episode_pool_size` | 32 | Maximum complete episodes mixed by each rank | | `streaming_episode_pool_size` | 32 | Maximum complete episodes mixed by each rank |
| `streaming_prefetch_episodes` | 8 | Episodes fetched ahead of the active pool | | `streaming_prefetch_episodes` | 8 | Episodes fetched ahead of the active pool |
| `streaming_byte_budget_gb` | 8 | Maximum synthesized MP4 bytes per rank | | `streaming_byte_budget_gb` | 8 | Maximum synthesized MP4 bytes per rank |
| `streaming_data_root` | unset | Optional local, Hub, bucket, or fsspec payload root | | `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 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 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_episode_pool_size=16 \
--dataset.streaming_prefetch_episodes=4 \ --dataset.streaming_prefetch_episodes=4 \
--dataset.streaming_byte_budget_gb=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 \ --num_workers=4 \
--policy.type=act \ --policy.type=act \
--output_dir=outputs/train/act_streaming --output_dir=outputs/train/act_streaming
@@ -109,6 +117,8 @@ uv run python scripts/bench_streaming_dataset.py \
--batch-size=16 \ --batch-size=16 \
--num-workers=1 \ --num-workers=1 \
--fetch-workers=4 \ --fetch-workers=4 \
--decode-threads=2 \
--decoded-queue-size=8 \
--summary-json=streaming-benchmark.json --summary-json=streaming-benchmark.json
``` ```
+6 -1
View File
@@ -70,7 +70,7 @@ def parse_args() -> argparse.Namespace:
type=int, type=int,
default=8, default=8,
help="Concurrent camera-fetch jobs. Total connections ~= workers x range-subranges; " 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( parser.add_argument(
"--range-subranges", "--range-subranges",
@@ -443,6 +443,7 @@ def run_exact_coverage_stream(
decode_pool.shutdown(wait=True) decode_pool.shutdown(wait=True)
elapsed = time.perf_counter() - start elapsed = time.perf_counter() - start
replacements = max(0, pool.admitted_count - min(pool_size, len(order)))
result = { result = {
"target_samples_s": target_samples_s, "target_samples_s": target_samples_s,
"actual_samples_s": samples_done / elapsed if elapsed > 0 else float("inf"), "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), "shard_total_frames": float(total_frames),
"epoch_complete": 1.0 if epoch_complete else 0.0, "epoch_complete": 1.0 if epoch_complete else 0.0,
"prefetch_ahead": float(prefetch_ahead), "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), "batch_size": float(batch_size),
"decode_workers": float(decode_workers), "decode_workers": float(decode_workers),
"kept_up": 1.0 "kept_up": 1.0
+15
View File
@@ -50,6 +50,11 @@ def parse_args() -> argparse.Namespace:
parser.add_argument("--episode-pool-size", type=int, default=32) parser.add_argument("--episode-pool-size", type=int, default=32)
parser.add_argument("--prefetch-episodes", type=int, default=8) parser.add_argument("--prefetch-episodes", type=int, default=8)
parser.add_argument("--byte-budget-gb", type=float, default=8.0) 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("--warmup-batches", type=int, default=8)
parser.add_argument("--measure-batches", type=int, default=128) parser.add_argument("--measure-batches", type=int, default=128)
parser.add_argument("--summary-json", type=Path, default=None) parser.add_argument("--summary-json", type=Path, default=None)
@@ -105,6 +110,11 @@ def main() -> None:
episode_pool_size=args.episode_pool_size, episode_pool_size=args.episode_pool_size,
prefetch_episodes=args.prefetch_episodes, prefetch_episodes=args.prefetch_episodes,
byte_budget_gb=args.byte_budget_gb, 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), max_num_shards=max(1, args.fetch_workers),
return_uint8=True, return_uint8=True,
) )
@@ -184,6 +194,11 @@ def main() -> None:
"episode_pool_size": args.episode_pool_size, "episode_pool_size": args.episode_pool_size,
"prefetch_episodes": args.prefetch_episodes, "prefetch_episodes": args.prefetch_episodes,
"byte_budget_gb": args.byte_budget_gb, "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, "warmup_batches": args.warmup_batches,
"measure_batches": args.measure_batches, "measure_batches": args.measure_batches,
}, },
+18
View File
@@ -52,6 +52,14 @@ class DatasetConfig:
streaming_prefetch_episodes: int = 8 streaming_prefetch_episodes: int = 8
# Hard per-rank cap for synthesized episode-video bytes. # Hard per-rank cap for synthesized episode-video bytes.
streaming_byte_budget_gb: float = 8.0 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). # Fraction of episodes held out per task for offline evaluation (0.0 = disabled).
eval_split: float = 0.0 eval_split: float = 0.0
@@ -68,6 +76,16 @@ class DatasetConfig:
raise ValueError("streaming_prefetch_episodes must be non-negative") raise ValueError("streaming_prefetch_episodes must be non-negative")
if self.streaming_byte_budget_gb <= 0: if self.streaming_byte_budget_gb <= 0:
raise ValueError("streaming_byte_budget_gb must be positive") 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 self.episodes is not None:
if any(ep < 0 for ep in self.episodes): if any(ep < 0 for ep in self.episodes):
raise ValueError( raise ValueError(
+10
View File
@@ -119,6 +119,11 @@ def make_dataset(
episode_pool_size=cfg.dataset.streaming_episode_pool_size, episode_pool_size=cfg.dataset.streaming_episode_pool_size,
prefetch_episodes=cfg.dataset.streaming_prefetch_episodes, prefetch_episodes=cfg.dataset.streaming_prefetch_episodes,
byte_budget_gb=cfg.dataset.streaming_byte_budget_gb, 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, repeat=True,
) )
else: else:
@@ -210,6 +215,11 @@ def make_train_eval_datasets(
episode_pool_size=cfg.dataset.streaming_episode_pool_size, episode_pool_size=cfg.dataset.streaming_episode_pool_size,
prefetch_episodes=cfg.dataset.streaming_prefetch_episodes, prefetch_episodes=cfg.dataset.streaming_prefetch_episodes,
byte_budget_gb=cfg.dataset.streaming_byte_budget_gb, 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, repeat=True,
) )
else: else:
+118 -39
View File
@@ -15,6 +15,7 @@
# limitations under the License. # limitations under the License.
import io import io
import os import os
from collections import deque
from collections.abc import Callable, Iterator, Mapping from collections.abc import Callable, Iterator, Mapping
from concurrent.futures import Future, ThreadPoolExecutor from concurrent.futures import Future, ThreadPoolExecutor
from pathlib import Path from pathlib import Path
@@ -122,6 +123,11 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
repeat: bool = False, repeat: bool = False,
*, *,
token: str | bool | None = None, 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. """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. episode_pool_size (int | None, optional): Number of complete episodes in the sampling pool.
prefetch_episodes (int, optional): Episodes prefetched beyond the active 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. 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 repeat (bool, optional): Repeat rank-local exact-coverage epochs without yielding a
short final training batch. The training factory enables this; direct iteration is short final training batch. The training factory enables this; direct iteration is
finite by default. finite by default.
@@ -196,9 +209,24 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
raise ValueError("prefetch_episodes must be non-negative") raise ValueError("prefetch_episodes must be non-negative")
if byte_budget_gb <= 0: if byte_budget_gb <= 0:
raise ValueError("byte_budget_gb must be positive") 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.episode_pool_size = episode_pool_size or min(buffer_size, 32)
self.prefetch_episodes = prefetch_episodes self.prefetch_episodes = prefetch_episodes
self.byte_budget = int(byte_budget_gb * 1024**3) 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.repeat = repeat
self._next_epoch = 0 self._next_epoch = 0
self._active_epoch = 0 self._active_epoch = 0
@@ -356,8 +384,13 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
columns=self._projected_columns, columns=self._projected_columns,
token=self._streaming_io_token, 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]] = {} episode_futures: dict[int, Future[datasets.Dataset]] = {}
decoded_futures: deque[Future[dict]] = deque()
scheduled_episodes: set[int] = set() scheduled_episodes: set[int] = set()
retained_video_episodes: set[int] = set() retained_video_episodes: set[int] = set()
if video_cache is not None: if video_cache is not None:
@@ -368,7 +401,7 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
def submit(episode_index: int) -> Future[datasets.Dataset]: def submit(episode_index: int) -> Future[datasets.Dataset]:
future = episode_futures.get(episode_index) future = episode_futures.get(episode_index)
if future is None: 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 episode_futures[episode_index] = future
return future return future
@@ -382,44 +415,83 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
video_cache.submit_prefetch(episode_index) video_cache.submit_prefetch(episode_index)
scheduled_episodes.add(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: try:
while True: schedule_frontier()
try: planner_exhausted = False
episode_index, frame_index = next(planner) while decoded_futures or not planner_exhausted:
except StopIteration: while not planner_exhausted and len(decoded_futures) < self.decoded_queue_size:
self._active_epoch = epoch + 1 if self.shuffle else 0 try:
self._state_offset = 0 episode_index, frame_index = next(planner)
break except StopIteration:
planner_exhausted = True
break
episode_dataset = submit(episode_index).result() episode_future = submit(episode_index)
item = self._make_episode_item( if video_cache is not None:
episode_dataset, video_cache.retain_episode(episode_index)
episode_index, try:
frame_index, decoded_future = decode_executor.submit(
video_cache=video_cache, decode_item,
) episode_future,
episode_index,
for evicted_episode in planner.evicted: frame_index,
episode_futures.pop(evicted_episode, None) )
if video_cache is not None and evicted_episode in retained_video_episodes: except Exception:
video_cache.release_episode(evicted_episode) if video_cache is not None:
retained_video_episodes.remove(evicted_episode) video_cache.release_episode(episode_index)
if video_cache is not None: raise
for admitted_episode in planner.newly_admitted: if video_cache is not None:
if admitted_episode not in retained_video_episodes: decoded_future.add_done_callback(
video_cache.retain_episode(admitted_episode) lambda _future, retained_episode=episode_index, retained_cache=video_cache: (
retained_video_episodes.add(admitted_episode) retained_cache.release_episode(retained_episode)
planner.evicted.clear() )
planner.newly_admitted.clear() )
schedule_frontier() 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 self._state_offset += 1
yield item yield item
self._active_epoch = epoch + 1 if self.shuffle else 0
self._state_offset = 0
finally: finally:
for future in decoded_futures:
future.cancel()
decode_executor.shutdown(wait=True, cancel_futures=True)
for future in episode_futures.values(): for future in episode_futures.values():
future.cancel() future.cancel()
executor.shutdown(wait=True, cancel_futures=True) parquet_executor.shutdown(wait=True, cancel_futures=True)
if video_cache is not None: if video_cache is not None:
video_cache.close() video_cache.close()
@@ -483,14 +555,15 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
sidecar_path=self._sidecar_path, sidecar_path=self._sidecar_path,
token=self._streaming_io_token, token=self._streaming_io_token,
) )
decoder_limit = max(1, min(64, self.episode_pool_size * max(1, len(self.meta.video_keys))))
return EpisodeByteCache( return EpisodeByteCache(
manifest, manifest,
self._data_root, self._data_root,
byte_budget=self.byte_budget, byte_budget=self.byte_budget,
workers=workers, workers=workers,
range_backend=range_backend, 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, video_backend=self._video_backend,
tolerance_s=self.tolerance_s, tolerance_s=self.tolerance_s,
token=self._streaming_io_token, token=self._streaming_io_token,
@@ -503,6 +576,7 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
frame_index: int, frame_index: int,
*, *,
video_cache: EpisodeByteCache | None, video_cache: EpisodeByteCache | None,
apply_image_transforms: bool = True,
) -> dict: ) -> dict:
item = episode_dataset[frame_index] item = episode_dataset[frame_index]
episode = self.meta.episodes[episode_index] episode = self.meta.episodes[episode_index]
@@ -563,11 +637,8 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
frames = frames.to(torch.float32) / 255.0 frames = frames.to(torch.float32) / 255.0
item[video_key] = frames.squeeze(0) item[video_key] = frames.squeeze(0)
if self.image_transforms is not None: if apply_image_transforms:
for camera_key in self.meta.camera_keys: self._apply_image_transforms(item)
if camera_key in self.meta.depth_keys:
continue
item[camera_key] = self.image_transforms(item[camera_key])
for key, stored_unit in self._image_depth_units.items(): 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: 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 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]: def state_dict(self) -> dict[str, int]:
return { return {
"epoch": self._active_epoch, "epoch": self._active_epoch,
+21 -8
View File
@@ -75,6 +75,7 @@ class EpisodeByteCache:
self._pool = ThreadPoolExecutor(max_workers=workers) self._pool = ThreadPoolExecutor(max_workers=workers)
self._cache: OrderedDict[tuple[int, str], dict[str, Any]] = OrderedDict() self._cache: OrderedDict[tuple[int, str], dict[str, Any]] = OrderedDict()
self._decoders: OrderedDict[tuple[int, 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._futures: dict[tuple[int, str], Future[dict[str, Any]]] = {}
self._retained_episodes: dict[int, int] = {} self._retained_episodes: dict[int, int] = {}
self._decoder_fallback_count = 0 self._decoder_fallback_count = 0
@@ -96,6 +97,7 @@ class EpisodeByteCache:
decoders = list(self._decoders.values()) decoders = list(self._decoders.values())
self._cache.clear() self._cache.clear()
self._decoders.clear() self._decoders.clear()
self._decoder_locks.clear()
self._futures.clear() self._futures.clear()
self._retained_episodes.clear() self._retained_episodes.clear()
self._fallback_decoders.clear() self._fallback_decoders.clear()
@@ -186,6 +188,7 @@ class EpisodeByteCache:
self._decoders[key] = decoder self._decoders[key] = decoder
while len(self._decoders) > self.max_open_decoders: while len(self._decoders) > self.max_open_decoders:
evicted_key, evicted_decoder = self._decoders.popitem(last=False) evicted_key, evicted_decoder = self._decoders.popitem(last=False)
self._decoder_locks.pop(evicted_key, None)
self._fallback_decoders.discard(evicted_key) self._fallback_decoders.discard(evicted_key)
_close_decoder(evicted_decoder) _close_decoder(evicted_decoder)
return decoder return decoder
@@ -227,15 +230,23 @@ class EpisodeByteCache:
decoder, release = self._decoder_for_frames(episode_index, camera_key) decoder, release = self._decoder_for_frames(episode_index, camera_key)
with self._lock: with self._lock:
uses_pyav_timestamps = self.video_backend == "pyav" and key not in self._fallback_decoders uses_pyav_timestamps = self.video_backend == "pyav" and key not in self._fallback_decoders
decode_lock = self._decoder_locks.setdefault(key, threading.Lock())
try: try:
if uses_pyav_timestamps: with decode_lock:
return decoder.get_frames_played_at(local_ts, tolerance_s=self.tolerance_s).data if uses_pyav_timestamps:
metadata = decoder.metadata return decoder.get_frames_played_at(local_ts, tolerance_s=self.tolerance_s).data
fps = getattr(metadata, "average_fps", None) metadata = decoder.metadata
if fps is None: fps = getattr(metadata, "average_fps", None)
duration = max(getattr(metadata, "end_stream_seconds", 0.0), 1e-9) if fps is None:
fps = metadata.num_frames / duration duration = max(getattr(metadata, "end_stream_seconds", 0.0), 1e-9)
return decoder.get_frames_at(indices=[round(ts * fps) for ts in local_ts]).data 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: finally:
if release is not None: if release is not None:
release() release()
@@ -257,6 +268,7 @@ class EpisodeByteCache:
with self._lock: with self._lock:
if self._decoders.get(key) is decoder: if self._decoders.get(key) is decoder:
self._decoders.pop(key) self._decoders.pop(key)
self._decoder_locks.pop(key, None)
self._fallback_decoders.discard(key) self._fallback_decoders.discard(key)
continue continue
return decoder, decoder.release return decoder, decoder.release
@@ -330,6 +342,7 @@ class EpisodeByteCache:
entry = self._cache.pop(key) entry = self._cache.pop(key)
self._bytes -= len(entry["bytes"]) self._bytes -= len(entry["bytes"])
decoder = self._decoders.pop(key, None) decoder = self._decoders.pop(key, None)
self._decoder_locks.pop(key, None)
self._fallback_decoders.discard(key) self._fallback_decoders.discard(key)
if decoder is not None: if decoder is not None:
_close_decoder(decoder) _close_decoder(decoder)
+27 -1
View File
@@ -11,8 +11,11 @@
from __future__ import annotations from __future__ import annotations
import posixpath import posixpath
import threading
import time
from collections.abc import Sequence from collections.abc import Sequence
from pathlib import Path from pathlib import Path
from typing import Any
import fsspec import fsspec
import pyarrow as pa import pyarrow as pa
@@ -29,6 +32,8 @@ class EpisodeParquetReader:
*, *,
columns: Sequence[str], columns: Sequence[str],
token: str | bool | None = None, token: str | bool | None = None,
max_retries: int = 4,
retry_backoff_s: float = 0.05,
): ):
if not columns: if not columns:
raise ValueError("EpisodeParquetReader requires at least one projected column") 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") self.columns if "episode_index" in self.columns else (*self.columns, "episode_index")
) )
data_root_str = str(data_root) 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 {} 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) 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") raise ValueError(f"Episode {episode_index} must contain at least one row")
path = posixpath.join(self._root_path.rstrip("/"), str(relative_path).lstrip("/")) 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) parquet = pq.ParquetFile(source)
available = set(parquet.schema_arrow.names) available = set(parquet.schema_arrow.names)
missing = sorted(set(self._read_columns) - available) missing = sorted(set(self._read_columns) - available)
@@ -71,6 +83,20 @@ class EpisodeParquetReader:
table = table.drop_columns(["episode_index"]) table = table.drop_columns(["episode_index"])
return table 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 @staticmethod
def _matching_row_group(parquet: pq.ParquetFile, episode_index: int) -> int | None: def _matching_row_group(parquet: pq.ParquetFile, episode_index: int) -> int | None:
episode_column = next( episode_column = next(
+5
View File
@@ -44,6 +44,11 @@ def test_dataset_config_empty_episodes_ok():
("streaming_episode_pool_size", 0, "episode_pool_size"), ("streaming_episode_pool_size", 0, "episode_pool_size"),
("streaming_prefetch_episodes", -1, "prefetch_episodes"), ("streaming_prefetch_episodes", -1, "prefetch_episodes"),
("streaming_byte_budget_gb", 0, "byte_budget_gb"), ("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): def test_dataset_config_rejects_invalid_streaming_resource_limits(field, value, message):
@@ -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") 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()
@@ -11,6 +11,7 @@
import json import json
import struct import struct
import threading import threading
import time
from concurrent.futures import ThreadPoolExecutor from concurrent.futures import ThreadPoolExecutor
import numpy as np 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 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): def test_releasing_episode_allows_immediate_eviction(monkeypatch, tmp_path):
with _fake_cache(monkeypatch, tmp_path, byte_budget=5) as cache: with _fake_cache(monkeypatch, tmp_path, byte_budget=5) as cache:
cache.retain_episode(0) cache.retain_episode(0)
+10
View File
@@ -38,6 +38,11 @@ def test_factory_wires_production_streaming_settings(monkeypatch):
streaming_episode_pool_size=7, streaming_episode_pool_size=7,
streaming_prefetch_episodes=3, streaming_prefetch_episodes=3,
streaming_byte_budget_gb=2.5, 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( cfg = SimpleNamespace(
dataset=dataset_config, 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"]["episode_pool_size"] == 7
assert captured["kwargs"]["prefetch_episodes"] == 3 assert captured["kwargs"]["prefetch_episodes"] == 3
assert captured["kwargs"]["byte_budget_gb"] == 2.5 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"]["max_num_shards"] == 1
assert captured["kwargs"]["video_backend"] == "pyav" assert captured["kwargs"]["video_backend"] == "pyav"
assert captured["kwargs"]["return_uint8"] is True assert captured["kwargs"]["return_uint8"] is True
@@ -10,6 +10,8 @@
from __future__ import annotations from __future__ import annotations
import threading
import time
from itertools import islice from itertools import islice
from pathlib import Path 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"])]) _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"]) @pytest.mark.parametrize("video_backend", ["torchcodec", "pyav"])
def test_streaming_rgb_video_matches_map_style( def test_streaming_rgb_video_matches_map_style(
tmp_path: Path, tmp_path: Path,