diff --git a/src/lerobot/datasets/streaming_dataset.py b/src/lerobot/datasets/streaming_dataset.py index a3c337dfb..e841df89b 100644 --- a/src/lerobot/datasets/streaming_dataset.py +++ b/src/lerobot/datasets/streaming_dataset.py @@ -14,7 +14,6 @@ # See the License for the specific language governing permissions and # limitations under the License. import logging -import time from collections.abc import Callable, Iterator from pathlib import Path @@ -109,7 +108,6 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset): world_size: int | None = None, video_decoder_cache_size: int | None = None, data_files_root: str | None = None, - video_decode_device: str = "cpu", ): """Initialize a StreamingLeRobotDataset. @@ -149,8 +147,6 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset): data_files_root (str | None, optional): fsspec root holding the bulk ``data/`` and ``videos/`` trees (e.g. ``hf://buckets//``). When set, parquet and video bytes are read from there while metadata still loads from ``repo_id`` on the Hub. - video_decode_device (str, optional): Device for torchcodec decode. ``"cuda"`` offloads to - NVDEC (needs a CUDA torchcodec build and ``spawn`` DataLoader workers). """ super().__init__() self.repo_id = repo_id @@ -184,14 +180,9 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset): self.rank, self.world_size = self._resolve_distributed(rank, world_size) self.video_decoder_cache_size = video_decoder_cache_size self.data_files_root = data_files_root.rstrip("/") if data_files_root else None - self.video_decode_device = video_decode_device # We cache the video decoders to avoid re-initializing them at each frame (avoiding a ~10x slowdown) self.video_decoder_cache = None - # Shared [hits, misses, evictions, decode_ns, fetch_ns] tensor so DataLoader workers aggregate - # decoder-cache stats and component timings into one place the main process can read after - # iteration (see video_decoder_cache_stats() / timing_stats()). - self._cache_counters = torch.zeros(5, dtype=torch.int64).share_memory_() self._epoch = 0 self._in_flight_epoch = 0 @@ -357,19 +348,11 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset): def _make_video_decoder_cache(self) -> VideoDecoderCache: """Size the decoder cache to the pool's working set (pool episodes x cameras), capped at 128.""" if self.video_decoder_cache_size is not None: - return VideoDecoderCache( - max_size=self.video_decoder_cache_size, - counters=self._cache_counters, - device=self.video_decode_device, - ) + return VideoDecoderCache(max_size=self.video_decoder_cache_size) num_cameras = len(self.meta.video_keys) if num_cameras == 0: - return VideoDecoderCache(counters=self._cache_counters, device=self.video_decode_device) - return VideoDecoderCache( - max_size=min((self.episode_pool_size + 1) * num_cameras, 128), - counters=self._cache_counters, - device=self.video_decode_device, - ) + return VideoDecoderCache() + return VideoDecoderCache(max_size=min((self.episode_pool_size + 1) * num_cameras, 128)) def __iter__(self) -> Iterator[dict[str, torch.Tensor]]: # `datasets` reshuffles (and re-permutes shard order) per epoch from (seed, epoch); @@ -383,13 +366,10 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset): iterator = iter(self._pipeline) while True: - fetch_start = time.perf_counter_ns() try: row = next(iterator) except StopIteration: return - finally: - self._cache_counters[4] += time.perf_counter_ns() - fetch_start yield self._finalize_sample(row) def _finalize_sample(self, row: dict) -> dict: @@ -416,9 +396,7 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset): query_timestamps = self._get_query_timestamps( current_ts, self.delta_indices, episode_boundaries_ts ) - decode_start = time.perf_counter_ns() video_frames = self._query_videos(query_timestamps, ep_idx) - self._cache_counters[3] += time.perf_counter_ns() - decode_start if self.image_transforms is not None: for cam in self.meta.camera_keys: @@ -451,31 +429,6 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset): self._epoch = int(state_dict.get("epoch", 0)) self._pipeline.load_state_dict(state_dict["pipeline"]) - def video_decoder_cache_stats(self) -> dict[str, int | float]: - """Decoder-cache reuse aggregated across DataLoader workers via the shared counter tensor. - - Unlike ``self.video_decoder_cache.stats()`` (which only reflects the main process), this sums - hits/misses/evictions over every worker. Counts are lock-free across processes, so treat them as - approximate; the ``hit_rate`` ratio is preserved. - """ - hits, misses, evictions = (int(x) for x in self._cache_counters[:3].tolist()) - total = hits + misses - return { - "hits": hits, - "misses": misses, - "evictions": evictions, - "hit_rate": round(hits / total, 4) if total else 0.0, - } - - def timing_stats(self) -> dict[str, float]: - """Cumulative seconds spent in video decode and in the upstream tabular pipeline (parquet - fetch + grouping + shuffles + explode), summed across DataLoader workers via the shared - counter tensor. These overlap in wall-clock (workers run in parallel), so compare them to - ``num_workers x wallclock`` for time fractions. - """ - decode_ns, fetch_ns = (int(x) for x in self._cache_counters[3:5].tolist()) - return {"decode_s_total": round(decode_ns / 1e9, 2), "fetch_s_total": round(fetch_ns / 1e9, 2)} - def _make_timestamps_from_indices( self, start_ts: float, indices: dict[str, list[int]] | None = None ) -> dict[str, list[float]]: diff --git a/src/lerobot/datasets/video_utils.py b/src/lerobot/datasets/video_utils.py index ee801e052..30fda72d1 100644 --- a/src/lerobot/datasets/video_utils.py +++ b/src/lerobot/datasets/video_utils.py @@ -22,7 +22,6 @@ import queue import shutil import tempfile import threading -import time import warnings from collections import OrderedDict from dataclasses import asdict, dataclass, field @@ -48,92 +47,6 @@ from lerobot.utils.import_utils import get_safe_default_video_backend logger = logging.getLogger(__name__) -DEFAULT_REMOTE_IO_MAX_RETRIES = 5 -"""Retry budget for transient hf:// / fsspec / httpx transport errors during streaming video decode. - -Streaming a dataset from an HF bucket/CDN issues many small range requests and occasionally hits a -transient transport failure (timeout, dropped connection, 408/5xx). The right response is to rebuild -the connection and retry rather than crash the DataLoader worker. Override via -``LEROBOT_REMOTE_IO_MAX_RETRIES``; set to ``0`` to disable retries (fail fast). -""" - -# Transient transport failures from the hf:// -> fsspec -> httpx stack. We match on text because the -# concrete exception types live in optional deps (httpx, huggingface_hub) and vary across versions. -# "client has been closed" is the important one: once a shared httpx client is closed by a single -# failed read, every subsequent read in that worker fails until the fsspec instance cache is cleared. -_RETRYABLE_TRANSPORT_FRAGMENTS = ( - "client has been closed", - "server disconnected", - "remoteprotocolerror", - "unexpected_eof", - "eof occurred in violation of protocol", - "connection reset", - "connection aborted", - "connection broken", - "incompleteread", - "read operation timed out", - "timed out", - "request time-out", - "408", - "502", - "503", - "504", -) - - -def _remote_io_max_retries() -> int: - raw = os.environ.get("LEROBOT_REMOTE_IO_MAX_RETRIES") - if raw is None: - return DEFAULT_REMOTE_IO_MAX_RETRIES - try: - return max(0, int(raw)) - except ValueError as e: - raise ValueError(f"LEROBOT_REMOTE_IO_MAX_RETRIES must be an integer; got {raw!r}") from e - - -def _is_retryable_transport_error(exc: BaseException) -> bool: - """True if ``exc`` looks like a transient remote-IO failure worth retrying (vs a real bug).""" - text = f"{type(exc).__name__}: {exc}".lower() - return any(fragment in text for fragment in _RETRYABLE_TRANSPORT_FRAGMENTS) - - -def _recover_remote_io(decoder_cache: "VideoDecoderCache", video_path: str) -> None: - """Drop the dead decoder for ``video_path`` and force a fresh fsspec client before a retry. - - fsspec caches one filesystem instance per (protocol, args), and that instance owns the httpx - client a failed read may have closed. Clearing the instance cache makes the next ``fsspec.open`` - build a new client, which is what breaks the "client has been closed" cascade. - """ - decoder_cache.invalidate(video_path) - with contextlib.suppress(Exception): - fsspec.AbstractFileSystem.clear_instance_cache() - - -def _retry_remote_io(operation, on_retry, max_retries: int, base_delay: float = 0.5, max_delay: float = 10.0): - """Run ``operation()``, retrying transient transport errors after ``on_retry()`` + capped backoff. - - Non-transport errors (decode / index / timestamp issues) propagate immediately so real bugs are - never masked by retries. - """ - attempt = 0 - while True: - try: - return operation() - except Exception as e: - if attempt >= max_retries or not _is_retryable_transport_error(e): - raise - attempt += 1 - logger.warning( - "Transient remote-IO error (%s: %s); rebuilding connection and retrying (%d/%d).", - type(e).__name__, - e, - attempt, - max_retries, - ) - on_retry() - time.sleep(min(base_delay * 2 ** (attempt - 1), max_delay)) - - def decode_video_frames( video_path: Path | str, timestamps: list[float], @@ -329,12 +242,7 @@ class VideoDecoderCache: _SENTINEL: ClassVar[object] = object() - def __init__( - self, - max_size: int | None | object = _SENTINEL, - counters: "torch.Tensor | None" = None, - device: str = "cpu", - ): + def __init__(self, max_size: int | None | object = _SENTINEL): if max_size is VideoDecoderCache._SENTINEL: max_size = _default_max_cache_size() if max_size is not None and max_size <= 0: @@ -342,18 +250,6 @@ class VideoDecoderCache: self.max_size: int | None = max_size # type: ignore[assignment] self._cache: OrderedDict[str, tuple[Any, Any]] = OrderedDict() self._lock = Lock() - # Decode device for the underlying torchcodec VideoDecoder. "cuda" offloads H.264/H.265 decode to - # the GPU's dedicated NVDEC engine (independent of the SMs used for training); requires a - # CUDA-enabled torchcodec/FFmpeg build. See https://developer.nvidia.com/video-codec-sdk. - self.device = device - # Observability counters (cheap, updated under the lock) for benchmarking decoder reuse. - self.hits = 0 - self.misses = 0 - self.evictions = 0 - # Optional shared [hits, misses, evictions] tensor so DataLoader workers aggregate into one place - # (the per-worker `self.*` ints are invisible to the main process). Lock-free across processes, so - # treat the aggregate as approximate; the hit-rate ratio is preserved. - self._counters = counters def __contains__(self, video_path: object) -> bool: with self._lock: @@ -375,21 +271,15 @@ class VideoDecoderCache: entry = self._cache.get(video_path) if entry is not None: self._cache.move_to_end(video_path) - self.hits += 1 - if self._counters is not None: - self._counters[0] += 1 return entry[0] - self.misses += 1 - if self._counters is not None: - self._counters[1] += 1 # Bound per-handle buffering: with many decoders kept open at once (one per camera per active # shard, across all workers), the default fsspec read cache balloons RAM on remote backends # like hf:// buckets. A small readahead cache caps each handle's footprint without hurting the # mostly-sequential reads torchcodec issues. file_handle = fsspec.open(video_path, cache_type="readahead", block_size=2**20).__enter__() try: - decoder = VideoDecoder(file_handle, seek_mode="approximate", device=self.device) + decoder = VideoDecoder(file_handle, seek_mode="approximate") except Exception: file_handle.close() raise @@ -401,9 +291,6 @@ class VideoDecoderCache: if self.max_size is not None: while len(self._cache) > self.max_size: _evicted_path, (_evicted_decoder, evicted_handle) = self._cache.popitem(last=False) - self.evictions += 1 - if self._counters is not None: - self._counters[2] += 1 with contextlib.suppress(Exception): evicted_handle.close() @@ -417,35 +304,11 @@ class VideoDecoderCache: file_handle.close() self._cache.clear() - def invalidate(self, video_path: str) -> None: - """Drop and close the cached decoder for a path whose connection went bad. - - After a transport error the cached ``fsspec`` handle (and the httpx client behind it) is dead; - removing the entry forces the next :meth:`get_decoder` to re-open a fresh handle. - """ - with self._lock: - entry = self._cache.pop(str(video_path), None) - if entry is not None: - with contextlib.suppress(Exception): - entry[1].close() - def size(self) -> int: """Return the number of cached decoders.""" with self._lock: return len(self._cache) - def stats(self) -> dict[str, int | float]: - """Return reuse counters (hits/misses/evictions, hit rate, current size) for benchmarking.""" - with self._lock: - total = self.hits + self.misses - return { - "hits": self.hits, - "misses": self.misses, - "evictions": self.evictions, - "hit_rate": self.hits / total if total else 0.0, - "size": len(self._cache), - } - class FrameTimestampError(ValueError): """Helper error to indicate the retrieved timestamps exceed the queried ones""" @@ -484,24 +347,20 @@ def decode_video_frames_torchcodec( if decoder_cache is None: decoder_cache = _default_decoder_cache - def _decode_frames(): - # Both opening the decoder and reading frames go over the network for hf:// paths, so wrap the - # whole unit: a transient transport error retries by dropping the dead handle and rebuilding - # the connection (see _retry_remote_io / _recover_remote_io) instead of killing the worker. - decoder = decoder_cache.get_decoder(str(video_path)) - average_fps = decoder.metadata.average_fps - frame_indices = [round(ts * average_fps) for ts in timestamps] - return decoder.get_frames_at(indices=frame_indices) - - frames_batch = _retry_remote_io( - _decode_frames, - on_retry=lambda: _recover_remote_io(decoder_cache, str(video_path)), - max_retries=_remote_io_max_retries(), - ) + # Use cached decoder instead of creating new one each time + decoder = decoder_cache.get_decoder(str(video_path)) loaded_ts = [] loaded_frames = [] + # get metadata for frame information + metadata = decoder.metadata + average_fps = metadata.average_fps + # convert timestamps to frame indices + frame_indices = [round(ts * average_fps) for ts in timestamps] + # retrieve frames based on indices + frames_batch = decoder.get_frames_at(indices=frame_indices) + for frame, pts in zip(frames_batch.data, frames_batch.pts_seconds, strict=True): loaded_frames.append(frame) loaded_ts.append(pts.item())