refactor(dataset): make streaming pool rank-owned

src/lerobot/configs/default.py

@@ -46,11 +46,11 @@ class DatasetConfig:
     streaming: bool = False
     # Optional data-plane root used by training-time streaming. Metadata still resolves from repo_id/root.
     streaming_data_root: str | None = None
-    # Number of complete episodes mixed by the exact-coverage sampler in each DataLoader worker.
+    # Number of complete episodes mixed by the rank-level exact-coverage sampler.
     streaming_episode_pool_size: int = 32
     # Complete episodes fetched ahead of the current admission frontier.
     streaming_prefetch_episodes: int = 8
-    # Hard per-worker cap for synthesized episode-video bytes.
+    # Hard per-rank cap for synthesized episode-video bytes.
     streaming_byte_budget_gb: float = 8.0
     # Fraction of episodes held out per task for offline evaluation (0.0 = disabled).
     eval_split: float = 0.0

src/lerobot/datasets/streaming_dataset.py

@@ -15,7 +15,7 @@
 # limitations under the License.
 import io
 import os
-from collections.abc import Callable, Iterator
+from collections.abc import Callable, Iterator, Mapping
 from concurrent.futures import Future, ThreadPoolExecutor
 from pathlib import Path
 
@@ -41,12 +41,31 @@ from .utils import check_version_compatibility
 from .video_utils import decode_video_frames_pyav
 
 
+def _balanced_episode_shards(
+    episode_indices: list[int],
+    episode_frame_counts: Mapping[int, int],
+    *,
+    world_size: int,
+) -> list[list[int]]:
+    """Assign whole episodes deterministically with greedy frame-count balancing."""
+    if world_size <= 0:
+        raise ValueError("world_size must be positive")
+    shards: list[list[int]] = [[] for _ in range(world_size)]
+    shard_frames = [0] * world_size
+    for episode in sorted(episode_indices, key=lambda item: (-episode_frame_counts[item], item)):
+        rank = min(range(world_size), key=lambda item: (shard_frames[item], item))
+        shards[rank].append(episode)
+        shard_frames[rank] += episode_frame_counts[episode]
+    return shards
+
+
 class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
     """Episode-scoped streaming reader for LeRobot datasets.
 
-    Metadata is cached locally, while each worker reads only the Parquet rows and MP4 byte ranges
-    needed for the complete episodes it owns. Episode ownership is disjoint across distributed
-    ranks and dataloader workers, and every selected frame is yielded exactly once per iteration.
+    Metadata is cached locally, while each rank reads only the Parquet rows and MP4 byte ranges
+    needed for the complete episodes it owns. One DataLoader worker per rank owns the logical pool;
+    its bounded result queue provides decoded-batch prefetch without creating independent samplers.
+    Episode ownership is disjoint and every selected frame is yielded exactly once per iteration.
     MP4 sidecars are resolved automatically and built in a revision-keyed local cache when absent.
 
     Example:
@@ -123,8 +142,8 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
                 and fsspec URLs.
             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-worker upper bound for synthesized episode video bytes.
-            repeat (bool, optional): Repeat worker-local exact-coverage epochs without yielding a
+            byte_budget_gb (float, optional): Per-rank upper bound for synthesized episode video bytes.
+            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.
         """
@@ -266,39 +285,40 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
         self._active_epoch = epoch
 
         worker = torch.utils.data.get_worker_info()
-        worker_index = worker.id if worker is not None else 0
-        worker_count = worker.num_workers if worker is not None else 1
+        if worker is not None and worker.num_workers > 1:
+            raise RuntimeError(
+                "StreamingLeRobotDataset uses one rank-level sampling pool and supports at most "
+                "one DataLoader worker per rank"
+            )
         resume_offset = self._resume_offset
         self._resume_offset = 0
-        worker_resume_offset, logical_worker_index = self._worker_resume_state(
-            resume_offset,
-            batch_size=self._resume_batch_size,
-            worker_index=worker_index,
-            worker_count=worker_count,
-        )
-        consumer_episodes, _consumer_index, _consumer_count = self._consumer_episodes(
-            worker_id=logical_worker_index,
-            num_workers=worker_count,
-        )
+        consumer_episodes, _rank, _world_size = self._rank_episodes()
         consumer_frame_count = sum(self._episode_frame_count(episode) for episode in consumer_episodes)
         if consumer_frame_count:
-            worker_epoch_delta, worker_resume_offset = divmod(
-                worker_resume_offset,
-                consumer_frame_count,
-            )
+            worker_epoch_delta, resume_offset = divmod(resume_offset, consumer_frame_count)
         else:
             worker_epoch_delta = 0
         epoch += worker_epoch_delta
         self._active_epoch = epoch
         if self.shuffle:
             self._next_epoch = epoch + 1
+
+        max_workers = min(self.max_num_shards, max(1, self.episode_pool_size + self.prefetch_episodes))
+        video_cache = self._make_video_cache(consumer_episodes, max_workers)
+        episode_byte_sizes = (
+            {episode: video_cache.manifest.episode_byte_size(episode) for episode in consumer_episodes}
+            if video_cache is not None
+            else None
+        )
         planner = ExactCoveragePool(
             [(episode, self._episode_frame_count(episode)) for episode in consumer_episodes],
             pool_size=self.episode_pool_size,
             seed=self.seed,
             epoch=epoch,
+            episode_byte_sizes=episode_byte_sizes,
+            byte_budget=self.byte_budget if episode_byte_sizes is not None else None,
         )
-        for _ in range(worker_resume_offset):
+        for _ in range(resume_offset):
             try:
                 next(planner)
             except StopIteration:
@@ -307,11 +327,9 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
         planner.evicted.clear()
 
         parquet_reader = EpisodeParquetReader(self._data_root, columns=self._projected_columns)
-        max_workers = min(self.max_num_shards, max(1, self.episode_pool_size + self.prefetch_episodes))
         executor = ThreadPoolExecutor(max_workers=max_workers, thread_name_prefix="lerobot-parquet")
         episode_futures: dict[int, Future[datasets.Dataset]] = {}
-        scheduled_cursor = 0
-        video_cache = self._make_video_cache(consumer_episodes, max_workers)
+        scheduled_episodes: set[int] = set()
         retained_video_episodes: set[int] = set()
         if video_cache is not None:
             for episode_index in planner.resident:
@@ -326,17 +344,14 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
             return future
 
         def schedule_frontier() -> None:
-            nonlocal scheduled_cursor
-            target = min(
-                len(planner.admission_order),
-                planner.admitted_count + self.prefetch_episodes,
-            )
-            while scheduled_cursor < target:
-                episode_index = planner.admission_order[scheduled_cursor]
+            frontier = [*planner.resident, *planner.prefetch_candidates(self.prefetch_episodes)]
+            for episode_index in frontier:
+                if episode_index in scheduled_episodes:
+                    continue
                 submit(episode_index)
                 if video_cache is not None:
                     video_cache.submit_prefetch(episode_index)
-                scheduled_cursor += 1
+                scheduled_episodes.add(episode_index)
 
         schedule_frontier()
         try:
@@ -379,12 +394,7 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
             if video_cache is not None:
                 video_cache.close()
 
-    def _consumer_episodes(
-        self,
-        *,
-        worker_id: int | None = None,
-        num_workers: int | None = None,
-    ) -> tuple[list[int], int, int]:
+    def _rank_episodes(self) -> tuple[list[int], int, int]:
         if torch.distributed.is_available() and torch.distributed.is_initialized():
             rank = torch.distributed.get_rank()
             world_size = torch.distributed.get_world_size()
@@ -393,49 +403,23 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
             world_size = int(os.environ.get("WORLD_SIZE", "1"))
         if world_size <= 0 or rank < 0 or rank >= world_size:
             raise ValueError(f"Invalid distributed rank/world size: rank={rank}, world_size={world_size}")
-
-        worker = torch.utils.data.get_worker_info()
-        if worker_id is None:
-            worker_id = worker.id if worker is not None else 0
-        if num_workers is None:
-            num_workers = worker.num_workers if worker is not None else 1
-        consumer_index = rank * num_workers + worker_id
-        consumer_count = world_size * num_workers
-        return self._selected_episodes[consumer_index::consumer_count], consumer_index, consumer_count
-
-    @staticmethod
-    def _worker_resume_state(
-        offset: int,
-        *,
-        batch_size: int,
-        worker_index: int,
-        worker_count: int,
-    ) -> tuple[int, int]:
-        completed_batches, partial_batch = divmod(offset, batch_size)
-        completed_cycles, next_logical_worker = divmod(completed_batches, worker_count)
-        logical_worker = (worker_index + next_logical_worker) % worker_count
-        worker_batches = completed_cycles + int(logical_worker < next_logical_worker)
-        worker_offset = worker_batches * batch_size
-        if partial_batch and logical_worker == next_logical_worker:
-            worker_offset += partial_batch
-        return worker_offset, logical_worker
+        counts = {episode: self._episode_frame_count(episode) for episode in self._selected_episodes}
+        shards = _balanced_episode_shards(self._selected_episodes, counts, world_size=world_size)
+        return shards[rank], rank, world_size
 
     def _episode_frame_count(self, episode_index: int) -> int:
         episode = self.meta.episodes[episode_index]
         return int(episode["dataset_to_index"] - episode["dataset_from_index"])
 
     def num_frames_for_rank(self, rank: int, world_size: int, num_workers: int) -> int:
-        """Return frames owned by one training rank under episode-scoped worker sharding."""
+        """Return frames owned by one training rank under balanced whole-episode sharding."""
         if world_size <= 0 or rank < 0 or rank >= world_size:
             raise ValueError(f"Invalid distributed rank/world size: rank={rank}, world_size={world_size}")
-        workers = max(1, num_workers)
-        consumer_count = world_size * workers
-        consumer_indices = set(range(rank * workers, (rank + 1) * workers))
-        return sum(
-            self._episode_frame_count(episode)
-            for position, episode in enumerate(self._selected_episodes)
-            if position % consumer_count in consumer_indices
-        )
+        if num_workers > 1:
+            raise ValueError("Rank-level streaming supports at most one DataLoader worker per rank")
+        counts = {episode: self._episode_frame_count(episode) for episode in self._selected_episodes}
+        shards = _balanced_episode_shards(self._selected_episodes, counts, world_size=world_size)
+        return sum(counts[episode] for episode in shards[rank])
 
     def _load_episode_dataset(
         self,

src/lerobot/scripts/lerobot_train.py

@@ -460,22 +460,16 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
     else:
         shuffle = True
         sampler = None
-        train_num_workers = cfg.num_workers
-        if train_num_workers > 0:
-            max_nonempty_workers = dataset.num_episodes // accelerator.num_processes
-            if max_nonempty_workers == 0:
-                raise ValueError(
-                    "At least one distributed rank owns no streaming episode. "
-                    "Reduce the distributed world size."
-                )
-            train_num_workers = min(train_num_workers, max_nonempty_workers)
-            if train_num_workers != cfg.num_workers and is_main_process:
-                logging.warning(
-                    "Reducing streaming DataLoader workers from %d to %d so every rank/worker "
-                    "owns at least one complete episode.",
-                    cfg.num_workers,
-                    train_num_workers,
-                )
+        # One dedicated DataLoader process owns the rank-level planner/cache. Its bounded result
+        # queue provides decoded-batch prefetch; max_num_shards controls the internal Parquet/range
+        # fetch executors independently.
+        train_num_workers = min(cfg.num_workers, 1)
+        if cfg.num_workers > 1 and is_main_process:
+            logging.info(
+                "Using one streaming DataLoader worker per rank; %d configured workers remain "
+                "available as the dataset's internal fetch concurrency.",
+                cfg.num_workers,
+            )
         rank_frames = dataset.num_frames_for_rank(
             accelerator.process_index,
             accelerator.num_processes,

src/lerobot/streaming/episode_video.py

@@ -17,7 +17,7 @@ import posixpath
 import threading
 import time
 from collections import OrderedDict
-from collections.abc import Sequence
+from collections.abc import Mapping, Sequence
 from concurrent.futures import Future, ThreadPoolExecutor, as_completed
 from dataclasses import dataclass
 from datetime import UTC, datetime
@@ -33,7 +33,13 @@ import numpy as np
 from huggingface_hub import HfApi, HfFileSystem, constants
 from huggingface_hub.utils import get_session, hf_raise_for_status
 
-from lerobot.streaming.mp4 import Mp4Index, Mp4SampleSlice, fetch_mp4_index, synthesize_mp4
+from lerobot.streaming.mp4 import (
+    Mp4Index,
+    Mp4SampleSlice,
+    fetch_mp4_index,
+    synthesize_mp4,
+    synthesized_mp4_size,
+)
 
 if TYPE_CHECKING:
     from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata
@@ -1041,6 +1047,16 @@ class EpisodeVideoManifest:
             source_start_pts=span.source_start_pts,
         )
 
+    def episode_byte_size(self, episode_index: int) -> int:
+        """Exact synthesized video bytes retained while an episode is active."""
+        return sum(
+            synthesized_mp4_size(
+                self.mp4_index(episode_index, camera_key),
+                self.sample_slice(episode_index, camera_key),
+            )
+            for camera_key in self.video_keys
+        )
+
 
 class ExactCoveragePool:
     """Deterministic, exactly-once frame coverage over a byte-cache episode pool.
@@ -1052,7 +1068,8 @@ class ExactCoveragePool:
     reproducible.
 
     Mechanics (this is the "evict only when all frames sampled" model):
-      - Episodes are admitted in a seeded global permutation; the first ``pool_size`` fill the pool.
+      - Episodes are admitted in a seeded global permutation until either ``pool_size`` or the
+        optional indexed-byte budget is reached.
       - Each resident episode carries a seeded shuffle of its own frame indices.
       - Each draw picks a resident episode with probability proportional to its *remaining* frames
         (i.e. a uniform draw over all remaining frames in the pool, the map-style ideal) and pops
@@ -1065,9 +1082,9 @@ class ExactCoveragePool:
     is fully unit-testable. It yields ``(episode_index, frame_index)``; map to a decode timestamp
     with ``frame_index / max(frame_count - 1, 1)``.
 
-    Determinism: the order is a pure function of ``(seed, epoch)`` and the episode->frame_count
-    map. Resume is a deterministic fast-forward: re-instantiate with the same seed/epoch and skip
-    ``n`` samples (tabular only, no decode).
+    Determinism: the order is a pure function of ``(seed, epoch)``, the episode frame counts, and
+    optional byte sizes/budget. Resume is a deterministic fast-forward: re-instantiate with the
+    same inputs and skip ``n`` samples (tabular only, no decode).
     """
 
     def __init__(
@@ -1077,34 +1094,71 @@ class ExactCoveragePool:
         *,
         seed: int,
         epoch: int = 0,
+        episode_byte_sizes: Mapping[int, int] | None = None,
+        byte_budget: int | None = None,
     ):
         self._counts = {int(ep): int(n) for ep, n in episode_frame_counts if int(n) > 0}
         self._rng = np.random.default_rng([seed, epoch])
         order = np.array(sorted(self._counts), dtype=np.int64)
         self._rng.shuffle(order)
-        # Full admission order is exposed so a caller can prefetch episodes ahead of when they
-        # enter the sampling pool (a freshly admitted episode is immediately eligible to be drawn,
-        # so its bytes must already be resident).
+        self.pool_size = max(1, pool_size)
+        self._byte_budget = byte_budget
+        if byte_budget is not None and byte_budget <= 0:
+            raise ValueError("byte_budget must be positive")
+        if byte_budget is not None and episode_byte_sizes is None:
+            raise ValueError("episode_byte_sizes are required when byte_budget is set")
+        self._byte_sizes = {
+            episode: int(episode_byte_sizes[episode]) if episode_byte_sizes is not None else 0
+            for episode in self._counts
+        }
+        if any(size < 0 for size in self._byte_sizes.values()):
+            raise ValueError("episode byte sizes must be non-negative")
+        if byte_budget is not None:
+            oversized = next(
+                ((episode, size) for episode, size in self._byte_sizes.items() if size > byte_budget),
+                None,
+            )
+            if oversized is not None:
+                episode, size = oversized
+                raise ValueError(
+                    f"Episode {episode} requires {size} bytes, exceeding the byte budget {byte_budget}"
+                )
+
+        # Preserve the full seeded order for benchmark/tooling compatibility. Byte-aware admission
+        # may temporarily skip an entry, but every episode remains in this deterministic frontier.
         self.admission_order: list[int] = order.tolist()
-        self._admit_cursor = 0
-        self._remaining: dict[int, list[int]] = {}
+        self._pending: list[int] = list(self.admission_order)
+        self._admitted_count = 0
+        self._remaining: dict[int, tuple[np.ndarray, int]] = {}
         self._remaining_total = 0
+        self._resident_bytes = 0
         self.newly_admitted: list[int] = []
         self.evicted: list[int] = []
-        for _ in range(max(1, pool_size)):
-            self._admit_one()
+        self._admit_available()
 
-    def _admit_one(self) -> None:
-        if self._admit_cursor >= len(self.admission_order):
-            return
-        ep = self.admission_order[self._admit_cursor]
-        self._admit_cursor += 1
-        n = self._counts[ep]
-        frames = np.arange(n, dtype=np.int64)
-        self._rng.shuffle(frames)
-        self._remaining[ep] = frames.tolist()
-        self._remaining_total += n
-        self.newly_admitted.append(ep)
+    def _admit_available(self) -> None:
+        while len(self._remaining) < self.pool_size and self._pending:
+            available_bytes = None if self._byte_budget is None else self._byte_budget - self._resident_bytes
+            pending_index = next(
+                (
+                    index
+                    for index, episode in enumerate(self._pending)
+                    if available_bytes is None or self._byte_sizes[episode] <= available_bytes
+                ),
+                None,
+            )
+            if pending_index is None:
+                return
+
+            episode = self._pending.pop(pending_index)
+            frame_count = self._counts[episode]
+            frames = np.arange(frame_count, dtype=np.int64)
+            self._rng.shuffle(frames)
+            self._remaining[episode] = (frames, frame_count)
+            self._remaining_total += frame_count
+            self._resident_bytes += self._byte_sizes[episode]
+            self._admitted_count += 1
+            self.newly_admitted.append(episode)
 
     @property
     def remaining_total(self) -> int:
@@ -1113,12 +1167,22 @@ class ExactCoveragePool:
     @property
     def admitted_count(self) -> int:
         """Number of episodes pulled from the admission order so far (pool fills + rotations)."""
-        return self._admit_cursor
+        return self._admitted_count
 
     @property
     def resident(self) -> list[int]:
         return list(self._remaining)
 
+    @property
+    def resident_bytes(self) -> int:
+        return self._resident_bytes
+
+    def prefetch_candidates(self, count: int) -> list[int]:
+        """Return the next deterministic pending frontier without admitting it."""
+        if count <= 0:
+            return []
+        return self._pending[:count]
+
     def __iter__(self) -> ExactCoveragePool:
         return self
 
@@ -1129,18 +1193,22 @@ class ExactCoveragePool:
         # remaining count. O(pool_size) per draw (~1024) -> negligible next to decode.
         target = int(self._rng.integers(self._remaining_total))
         chosen = None
-        for ep, frames in self._remaining.items():
-            if target < len(frames):
+        for ep, (_frames, remaining) in self._remaining.items():
+            if target < remaining:
                 chosen = ep
                 break
-            target -= len(frames)
-        frames = self._remaining[chosen]
-        frame_index = frames.pop()
+            target -= remaining
+        frames, remaining = self._remaining[chosen]
+        remaining -= 1
+        frame_index = int(frames[remaining])
         self._remaining_total -= 1
-        if not frames:
+        if remaining == 0:
             del self._remaining[chosen]
             self.evicted.append(chosen)
-            self._admit_one()
+            self._resident_bytes -= self._byte_sizes[chosen]
+            self._admit_available()
+        else:
+            self._remaining[chosen] = (frames, remaining)
         return chosen, frame_index
 
 

src/lerobot/streaming/mp4.py

@@ -329,10 +329,49 @@ def synthesize_mp4(index: Mp4Index, sample_slice: Mp4SampleSlice, mdat_payload:
     sync = index.sync_samples[(index.sync_samples >= lo) & (index.sync_samples < hi)] - lo + 1
     moov = _make_moov(index, durations, sizes, rel_offsets, sync, mdat_data_offset=0)
     header_size = len(index.ftyp) + len(moov)
-    moov = _make_moov(index, durations, sizes, rel_offsets, sync, mdat_data_offset=header_size + 8)
+    mdat_header_size = 8 if len(mdat_payload) + 8 <= 0xFFFFFFFF else 16
+    moov = _make_moov(
+        index,
+        durations,
+        sizes,
+        rel_offsets,
+        sync,
+        mdat_data_offset=header_size + mdat_header_size,
+    )
     return index.ftyp + moov + _box(b"mdat", mdat_payload)
 
 
+def synthesized_mp4_size(index: Mp4Index, sample_slice: Mp4SampleSlice) -> int:
+    """Return the exact synthesized mini-MP4 size without fetching its media payload."""
+    lo = sample_slice.sample_lo
+    hi = sample_slice.sample_hi + 1
+    if lo < 0 or hi > len(index.sample_sizes) or lo >= hi:
+        raise ValueError(f"Invalid sample range [{lo}, {hi}) for {index.file_path}")
+
+    offsets = index.sample_offsets[lo:hi]
+    sizes = index.sample_sizes[lo:hi]
+    rel_offsets = offsets - sample_slice.byte_offset
+    if int(rel_offsets.min()) != 0:
+        raise ValueError("Sample slice must start at the minimum referenced sample offset")
+    if int((rel_offsets + sizes).max()) > sample_slice.byte_length:
+        raise ValueError("Sample slice does not cover all referenced samples")
+
+    durations = index.sample_durations[lo:hi]
+    sync = index.sync_samples[(index.sync_samples >= lo) & (index.sync_samples < hi)] - lo + 1
+    moov = _make_moov(index, durations, sizes, rel_offsets, sync, mdat_data_offset=0)
+    header_size = len(index.ftyp) + len(moov)
+    mdat_header_size = 8 if sample_slice.byte_length + 8 <= 0xFFFFFFFF else 16
+    moov = _make_moov(
+        index,
+        durations,
+        sizes,
+        rel_offsets,
+        sync,
+        mdat_data_offset=header_size + mdat_header_size,
+    )
+    return len(index.ftyp) + len(moov) + mdat_header_size + sample_slice.byte_length
+
+
 def iter_boxes(
     data: bytes,
     start: int,