admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-07-08 18:41:54 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor(streaming): trim video_utils to the minimal readahead cap

Browse Source 
Drop the transient-IO retry layer and the decoder-cache observability counters from
video_utils.py, keeping only the fsspec readahead cache that bounds per-handle RAM for
remote (hf://) decoders. Remove the now-orphaned instrumentation from StreamingLeRobotDataset
(video_decode_device/NVDEC, shared cache-counter tensor, video_decoder_cache_stats(),
timing_stats()). Retry is deferred to a separate, focused PR.

Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
pepijn
2026-06-12 09:50:43 +00:00
parent 674c990a39
commit 7bcd5a1502
2 changed files with 15 additions and 203 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/lerobot/datasets/streaming_dataset.py
+3 -50
View File
@@ -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/<owner>/<name>``). 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]]:
src/lerobot/datasets/video_utils.py
+12 -153
View File
@@ -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())