mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-24 10:16:09 +00:00
fix(streaming): worker-exact resume arithmetic and multi-worker resume test
The fast-forward skip assumed every DataLoader worker delivers batches; workers that own no shards yield nothing and are stopped, so the batch round-robin runs over min(num_workers, num_shards) active workers. Use that effective count (shard-less workers skip nothing). Adds a resume test under num_workers=2 asserting exact continuation. Note: the test fixtures write a single parquet file regardless of data_files_size_in_mb, so worker-splitting tests exercise the degenerate single-shard layout; multi-shard behavior is covered by the rank-level split_dataset_by_node tests. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
@@ -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`,
|
`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
|
`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`
|
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)
|
## Bucket sources & prewarming (manual)
|
||||||
|
|
||||||
|
|||||||
@@ -463,7 +463,9 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
|
|||||||
epoch = self._epoch
|
epoch = self._epoch
|
||||||
self._epoch += 1
|
self._epoch += 1
|
||||||
rng = self._consumer_rng(epoch, worker_id)
|
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).
|
# 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]
|
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"]),
|
"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:
|
if self._resume_state is None:
|
||||||
return
|
return
|
||||||
batches = self._resume_state["batches_consumed"]
|
batches = self._resume_state["batches_consumed"]
|
||||||
batch_size = self._resume_state["batch_size"]
|
batch_size = self._resume_state["batch_size"]
|
||||||
self._resume_state = None
|
self._resume_state = None
|
||||||
# DataLoader assigns batch j to worker j % num_workers.
|
if worker_id >= active_workers:
|
||||||
my_batches = batches // num_workers + (1 if batches % num_workers > worker_id else 0)
|
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
|
self._ff_remaining = my_batches * batch_size
|
||||||
if self._ff_remaining:
|
if self._ff_remaining:
|
||||||
logger.info(
|
logger.info(
|
||||||
|
|||||||
@@ -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 sorted(shuffled) == sorted(ordered), "shuffling changed the set of frames"
|
||||||
assert shuffled != ordered, "shuffle did not decorrelate output order"
|
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"
|
||||||
|
)
|
||||||
|
|||||||
Reference in New Issue
Block a user