feat(streaming): defer video decode, episode-pool shuffle, and remote-IO retries

- streaming_dataset: defer torchcodec decode until a sample leaves the shuffle
  buffer (buffer now holds ~KB tabular rows, not MB of pixels) and add an opt-in
  episode-pool shuffle (episode_pool_size) with exact in-episode delta lookups;
  expose decode/fetch timing_stats.
- video_utils: retry transient hf:///fsspec/httpx transport errors during
  streaming decode (LEROBOT_REMOTE_IO_MAX_RETRIES).
- dataset_tools: write multiple ~32MB row groups with a page index to bound
  per-shard streaming memory.
- benchmarks/slurm: streaming benchmark + matrix submitter updates.

Co-authored-by: Cursor <cursoragent@cursor.com>

benchmarks/streaming/benchmark_streaming.py

@@ -36,7 +36,9 @@ is whatever ``--repo_id``/``--root`` point at. See the README for bucket prewarm
 import argparse
 import csv
 import json
+import os
 import statistics
+import threading
 import time
 from pathlib import Path
 
@@ -47,6 +49,60 @@ from lerobot.datasets import LeRobotDatasetMetadata, StreamingLeRobotDataset
 from lerobot.utils.constants import ACTION
 
 
+def _tree_rss_bytes() -> int:
+    """Sum RSS of this process and all its descendants via /proc (Linux only; 0 elsewhere).
+
+    DataLoader workers are separate processes, so the parent's own RSS misses most of the pipeline's
+    memory. Walking the process tree captures the real footprint (parquet buffers + decoders + shuffle).
+    """
+    try:
+        children: dict[int, list[int]] = {}
+        for entry in os.listdir("/proc"):
+            if not entry.isdigit():
+                continue
+            try:
+                with open(f"/proc/{entry}/stat") as f:
+                    ppid = int(f.read().split(") ", 1)[1].split()[1])
+                children.setdefault(ppid, []).append(int(entry))
+            except (OSError, ValueError, IndexError):
+                pass
+        total, stack = 0, [os.getpid()]
+        while stack:
+            cur = stack.pop()
+            try:
+                with open(f"/proc/{cur}/statm") as f:
+                    total += int(f.read().split()[1]) * os.sysconf("SC_PAGE_SIZE")
+            except (OSError, ValueError, IndexError):
+                pass
+            stack.extend(children.get(cur, []))
+        return total
+    except OSError:
+        return 0
+
+
+class PeakRSSSampler:
+    """Background thread tracking peak process-tree RSS for the duration of the `with` block."""
+
+    def __init__(self, interval_s: float = 0.5):
+        self.interval_s = interval_s
+        self.peak_bytes = 0
+        self._stop = threading.Event()
+        self._thread = threading.Thread(target=self._run, daemon=True)
+
+    def _run(self) -> None:
+        while not self._stop.is_set():
+            self.peak_bytes = max(self.peak_bytes, _tree_rss_bytes())
+            self._stop.wait(self.interval_s)
+
+    def __enter__(self) -> "PeakRSSSampler":
+        self._thread.start()
+        return self
+
+    def __exit__(self, *exc) -> None:
+        self._stop.set()
+        self._thread.join(timeout=2)
+
+
 def parse_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser(description=__doc__)
     parser.add_argument("--repo_id", type=str, required=True)
@@ -62,8 +118,30 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--source", type=str, default="hub", help="Label only: hub | bucket | warmed_bucket.")
     parser.add_argument("--batch_size", type=int, default=64)
     parser.add_argument("--num_workers", type=int, default=8)
+    parser.add_argument(
+        "--prefetch_factor",
+        type=int,
+        default=2,
+        help="DataLoader batches prefetched per worker. Higher hides IO/decode latency but raises RAM "
+        "(prefetch_factor x num_workers x batch_size decoded frames held in flight). Ignored if num_workers=0.",
+    )
     parser.add_argument("--buffer_size", type=int, default=2000)
+    parser.add_argument(
+        "--max_num_shards",
+        type=int,
+        default=16,
+        help="Cap on concurrently-open stream shards. Each open shard holds ~one parquet row group in "
+        "RAM; reading from an hf:// bucket buffers ~5x more per shard than hf:// datasets, so lower this "
+        "(e.g. to num_workers) for bucket sources to avoid OOM. All data is still covered via re-sharding.",
+    )
     parser.add_argument("--video_decoder_cache_size", type=int, default=None)
+    parser.add_argument(
+        "--episode_pool_size",
+        type=int,
+        default=None,
+        help="A3 shuffle: keep this many full episodes live and sample frames uniformly across them "
+        "(mixing radius = this many episodes). Unset = default per-shard reservoir shuffle.",
+    )
     parser.add_argument(
         "--video_decode_device",
         type=str,
@@ -87,8 +165,10 @@ def build_dataset(args: argparse.Namespace, meta: LeRobotDatasetMetadata) -> Str
         data_files_root=args.data_files_root,
         delta_timestamps=delta_timestamps,
         buffer_size=args.buffer_size,
+        max_num_shards=args.max_num_shards,
         video_decoder_cache_size=args.video_decoder_cache_size,
         video_decode_device=args.video_decode_device,
+        episode_pool_size=args.episode_pool_size,
         tolerance_s=1e-3,
     )
 
@@ -116,37 +196,43 @@ def main() -> None:
         # tensors errors). Pin only when decode is on CPU and we copy to a CUDA device.
         pin_memory=device.type == "cuda" and not gpu_decode,
         drop_last=True,
-        prefetch_factor=2 if args.num_workers > 0 else None,
+        prefetch_factor=args.prefetch_factor if args.num_workers > 0 else None,
         # CUDA cannot initialize in forked workers; NVDEC decode in workers needs the spawn start method.
         multiprocessing_context="spawn" if gpu_decode and args.num_workers > 0 else None,
     )
 
     sample_latencies_ms: list[float] = []
+    episodes_per_batch: list[int] = []  # shuffle-randomness proxy: distinct episodes within a batch
     frames = 0
     first_batch_latency_s = None
     steady_start = None  # wall-clock start of the post-warmup measurement window
 
     t_start = time.perf_counter()
     t_prev = t_start
-    for i, batch in enumerate(loader):
-        # Dummy consume: move tensors to the device, mimicking what a real trainer would do.
-        for value in batch.values():
-            if torch.is_tensor(value):
-                value.to(device, non_blocking=device.type == "cuda")
-        now = time.perf_counter()
-        if first_batch_latency_s is None:
-            first_batch_latency_s = now - t_start
+    with PeakRSSSampler() as rss:
+        for i, batch in enumerate(loader):
+            # Dummy consume: move tensors to the device, mimicking what a real trainer would do.
+            for value in batch.values():
+                if torch.is_tensor(value):
+                    value.to(device, non_blocking=device.type == "cuda")
+            now = time.perf_counter()
+            if first_batch_latency_s is None:
+                first_batch_latency_s = now - t_start
 
-        if i == args.warmup_batches:
-            # Start the steady window here; the slow first batch and the prefetch queue it filled are
-            # excluded so throughput reflects sustained production, not draining a pre-filled queue.
-            steady_start = now
-        elif i > args.warmup_batches:
-            sample_latencies_ms.append((now - t_prev) / args.batch_size * 1000.0)
-            frames += args.batch_size
-        t_prev = now
-        if i + 1 >= args.num_batches:
-            break
+            if i == args.warmup_batches:
+                # Start the steady window here; the slow first batch and the prefetch queue it filled are
+                # excluded so throughput reflects sustained production, not draining a pre-filled queue.
+                steady_start = now
+            elif i > args.warmup_batches:
+                sample_latencies_ms.append((now - t_prev) / args.batch_size * 1000.0)
+                frames += args.batch_size
+                ep = batch.get("episode_index")
+                if torch.is_tensor(ep):
+                    episodes_per_batch.append(int(torch.unique(ep).numel()))
+            t_prev = now
+            if i + 1 >= args.num_batches:
+                break
+    peak_rss_gb = round(rss.peak_bytes / 1e9, 2) if rss.peak_bytes else None
 
     now = time.perf_counter()
     elapsed = now - t_start
@@ -154,6 +240,16 @@ def main() -> None:
     # gaps collapse to ~0 (the consumer drains a pre-filled queue) and overstate throughput by ~100x.
     steady_elapsed_s = (now - steady_start) if steady_start is not None else elapsed
     cache_stats = dataset.video_decoder_cache_stats()
+    timing = dataset.timing_stats()  # cumulative decode/fetch seconds summed across workers
+    # Image (camera frame) resolution as decoded, e.g. [C, H, W]. Read from the dataset feature contract.
+    image_shape = (
+        list(meta.features[meta.video_keys[0]]["shape"]) if meta.video_keys else None
+    )
+    # Decode/fetch overlap in wall-clock (workers run in parallel), so normalize against the total worker
+    # budget (num_workers x wallclock) to express each stage as a fraction of available worker time.
+    worker_budget_s = max(args.num_workers, 1) * elapsed
+    decode_pct = round(100 * timing["decode_s_total"] / worker_budget_s, 1) if worker_budget_s else None
+    fetch_pct = round(100 * timing["fetch_s_total"] / worker_budget_s, 1) if worker_budget_s else None
 
     # A 0-frame run is a failure, not a 0-throughput result: the pipeline produced no batches (decode
     # error swallowed in workers, all batches dropped by drop_last, etc.). Exit non-zero so the job is
@@ -172,11 +268,22 @@ def main() -> None:
         "mode": args.mode,
         "batch_size": args.batch_size,
         "num_workers": args.num_workers,
+        "prefetch_factor": args.prefetch_factor if args.num_workers > 0 else None,
         "buffer_size": args.buffer_size,
+        "episode_pool_size": args.episode_pool_size,
+        "episodes_per_batch_mean": round(statistics.mean(episodes_per_batch), 1)
+        if episodes_per_batch
+        else None,
+        # Fraction of a batch that is distinct episodes; ~1.0 ≈ map-style uniform, low ≈ correlated.
+        "shuffle_randomness_frac": round(statistics.mean(episodes_per_batch) / args.batch_size, 3)
+        if episodes_per_batch
+        else None,
         "num_cameras": len(meta.video_keys),
+        "image_shape": image_shape,
         "fps": meta.fps,
         "device": str(device),
         "video_decode_device": args.video_decode_device,
+        "peak_rss_gb": peak_rss_gb,
         "frames_measured": frames,
         "first_batch_latency_s": round(first_batch_latency_s or float("nan"), 4),
         "frames_per_s_node": round(frames / steady_elapsed_s, 2) if steady_elapsed_s else 0.0,
@@ -186,13 +293,23 @@ def main() -> None:
         else None,
         "p95_sample_latency_ms": round(percentile(sample_latencies_ms, 95), 3),
         "p99_sample_latency_ms": round(percentile(sample_latencies_ms, 99), 3),
+        "total_time_s": round(elapsed, 2),
+        "steady_time_s": round(steady_elapsed_s, 2),
         "wallclock_s": round(elapsed, 2),
+        "decode_s_total": timing["decode_s_total"],
+        "fetch_s_total": timing["fetch_s_total"],
+        "decode_pct_worker_time": decode_pct,
+        "fetch_pct_worker_time": fetch_pct,
         "video_decoder_cache": cache_stats,
     }
 
     out_dir = Path(args.out_dir)
     out_dir.mkdir(parents=True, exist_ok=True)
-    tag = f"{args.source}_{args.mode}_bs{args.batch_size}_w{args.num_workers}_{args.video_decode_device}"
+    pool_tag = f"_ep{args.episode_pool_size}" if args.episode_pool_size else ""
+    tag = (
+        f"{args.source}_{args.mode}_bs{args.batch_size}_w{args.num_workers}"
+        f"_pf{args.prefetch_factor}{pool_tag}_{args.video_decode_device}"
+    )
     (out_dir / f"{tag}.json").write_text(json.dumps(results, indent=2))
     flat = {k: (json.dumps(v) if isinstance(v, dict) else v) for k, v in results.items()}
     with open(out_dir / f"{tag}.csv", "w", newline="") as f: