diff --git a/benchmarks/streaming/README.md b/benchmarks/streaming/README.md index 18677fcda..cce36f329 100644 --- a/benchmarks/streaming/README.md +++ b/benchmarks/streaming/README.md @@ -74,7 +74,7 @@ frames/s/node at `--num_workers 3` (3 cameras, fps 20). `frames_per_s_node`, `samples_per_s`, `first_batch_latency_s`, `p50/p95/p99_sample_latency_ms`, `wallclock_s`, and `video_decoder_cache` (`hits`, `misses`, `evictions`, `hit_rate`, `size`). A low cache `hit_rate` with high `p99` is the decoder-thrash signature — raise `--video_decoder_cache_size` -or `--buffer_size`, or reduce `num_workers`. +or `--episode_pool_size`, or reduce `num_workers`. ## Bucket sources & prewarming (manual) diff --git a/src/lerobot/datasets/streaming_dataset.py b/src/lerobot/datasets/streaming_dataset.py index 787720803..882884be1 100644 --- a/src/lerobot/datasets/streaming_dataset.py +++ b/src/lerobot/datasets/streaming_dataset.py @@ -463,7 +463,9 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset): epoch = self._epoch self._epoch += 1 rng = self._consumer_rng(epoch, worker_id) - self._consume_resume_state(worker_id, num_workers) + # 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. + self._consume_resume_state(worker_id, min(num_workers, num_shards)) # Round-robin episode admission across this consumer's shard streams (deterministic). streams = [self._iter_shard_episodes(safe_shard(ds, idx, num_shards)) for idx in shard_indices] @@ -541,14 +543,16 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset): "batch_size": int(state_dict["batch_size"]), } - def _consume_resume_state(self, worker_id: int, num_workers: int) -> None: + def _consume_resume_state(self, worker_id: int, active_workers: int) -> None: if self._resume_state is None: return batches = self._resume_state["batches_consumed"] batch_size = self._resume_state["batch_size"] self._resume_state = None - # DataLoader assigns batch j to worker j % num_workers. - my_batches = batches // num_workers + (1 if batches % num_workers > worker_id else 0) + if worker_id >= active_workers: + return # this worker owns no shards and never delivered a batch + # The DataLoader assigns batch j to active worker j % active_workers. + my_batches = batches // active_workers + (1 if batches % active_workers > worker_id else 0) self._ff_remaining = my_batches * batch_size if self._ff_remaining: logger.info( diff --git a/tests/datasets/test_streaming_native.py b/tests/datasets/test_streaming_native.py index 1577cae38..47cb9a7cd 100644 --- a/tests/datasets/test_streaming_native.py +++ b/tests/datasets/test_streaming_native.py @@ -316,3 +316,40 @@ def test_shuffle_decorrelates_output_order(tmp_path, lerobot_dataset_factory): ) assert sorted(shuffled) == sorted(ordered), "shuffling changed the set of frames" assert shuffled != ordered, "shuffle did not decorrelate output order" + + +def test_fast_forward_resume_with_dataloader_workers(tmp_path, lerobot_dataset_factory): + """Resume must be exact under num_workers > 0: each worker re-derives its own skip.""" + from torch.utils.data import DataLoader + + repo_id = f"{DUMMY_REPO_ID}-resume-workers" + _make_local_dataset(lerobot_dataset_factory, tmp_path / "ds", repo_id, total_episodes=8, total_frames=120) + + num_workers = 2 + + def fresh_ds(): + return StreamingLeRobotDataset( + repo_id=repo_id, + root=tmp_path / "ds", + shuffle=True, + seed=11, + episode_pool_size=3, + max_num_shards=4, + ) + + def epoch_samples(ds): + # batch_size=None yields raw samples; the DataLoader round-robins them across workers, + # which is batch_size=1 in the resume arithmetic. + loader = DataLoader(ds, batch_size=None, num_workers=num_workers) + return [int(sample["index"]) for sample in loader] + + full = epoch_samples(fresh_ds()) + + samples_consumed = 17 + resumed_ds = fresh_ds() + resumed_ds.load_state_dict({"batches_consumed": samples_consumed, "batch_size": 1}) + resumed = epoch_samples(resumed_ds) + + assert resumed == full[samples_consumed:], ( + "fast-forward resume with DataLoader workers did not continue at the exact sample" + )