diff --git a/src/lerobot/datasets/streaming_dataset.py b/src/lerobot/datasets/streaming_dataset.py index c03e19c61..2f69867e1 100644 --- a/src/lerobot/datasets/streaming_dataset.py +++ b/src/lerobot/datasets/streaming_dataset.py @@ -400,6 +400,21 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset): state = _mix64(state ^ _mix64(salt)) return np.random.default_rng(state) + def _shard_order(self, epoch: int, num_shards: int) -> list[int]: + """Seeded permutation of this rank's shard indices, re-drawn every epoch. + + In a sub-epoch run over a corpus consolidated source-by-source, index-order shard + assignment means training on whatever the first N% of files contains; permuting the + shard order turns that into a uniform sample of files. Seeded by (seed, epoch, rank) + only — every DataLoader worker of the rank must agree on this list, because workers + stride it and disagreement would create overlapping shard assignments. + """ + order = list(range(num_shards)) + if self.shuffle: + state = _mix64(self.seed) ^ _mix64(0x5EED5EED) ^ _mix64(self.rank) ^ _mix64(epoch) + np.random.default_rng(_mix64(state)).shuffle(order) + return order + 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: @@ -481,7 +496,9 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset): ds = split_dataset_by_node(ds, rank=self.rank, world_size=self.world_size) num_shards = ds.num_shards if self.max_num_shards is None else min(ds.num_shards, self.max_num_shards) - shard_indices = list(range(num_shards)) + epoch = self._epoch + self._epoch += 1 + shard_indices = self._shard_order(epoch, num_shards) # DataLoader workers within this rank further split the shards so they don't yield duplicates. worker_info = torch.utils.data.get_worker_info() @@ -498,8 +515,6 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset): prefetcher = self._make_prefetcher() self._prefetcher = prefetcher - epoch = self._epoch - self._epoch += 1 rng = self._consumer_rng(epoch, worker_id) # Workers beyond the shard count yield nothing and are stopped by the DataLoader, so the # batch round-robin effectively runs over min(num_workers, num_shards) active workers. diff --git a/tests/datasets/test_streaming_native.py b/tests/datasets/test_streaming_native.py index f616bfe8d..2856cc4ff 100644 --- a/tests/datasets/test_streaming_native.py +++ b/tests/datasets/test_streaming_native.py @@ -378,3 +378,27 @@ def test_episode_grouping_native_and_fallback_agree(tmp_path, lerobot_dataset_fa return native = episode_signature(True) assert native == fallback + + +def test_shard_order_permutation_properties(tmp_path, lerobot_dataset_factory): + """Shard order: a valid permutation, deterministic per (seed, epoch, rank), worker-independent + (workers stride the same list, so it must not depend on worker id), reshuffled across epochs, + and identity when shuffle is off.""" + repo_id = f"{DUMMY_REPO_ID}-shardorder" + _make_local_dataset(lerobot_dataset_factory, tmp_path / "ds", repo_id, total_episodes=4, total_frames=80) + + ds = StreamingLeRobotDataset(repo_id=repo_id, root=tmp_path / "ds", shuffle=True, seed=5) + num_shards = 32 + order_epoch0 = ds._shard_order(0, num_shards) + assert sorted(order_epoch0) == list(range(num_shards)) + assert ds._shard_order(0, num_shards) == order_epoch0 # deterministic + assert ds._shard_order(1, num_shards) != order_epoch0 # reshuffles per epoch + assert order_epoch0 != list(range(num_shards)) # actually permuted (P=1/32! of false alarm) + + other_rank = StreamingLeRobotDataset( + repo_id=repo_id, root=tmp_path / "ds", shuffle=True, seed=5, rank=1, world_size=2 + ) + assert other_rank._shard_order(0, num_shards) != order_epoch0 # ranks decorrelated + + unshuffled = StreamingLeRobotDataset(repo_id=repo_id, root=tmp_path / "ds", shuffle=False, seed=5) + assert unshuffled._shard_order(0, num_shards) == list(range(num_shards))