mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-11 20:11:48 +00:00
CarolinePascal
7bee7fb9e3
Expose the shardable unit of work behind recompute_stats: compute_dataset_episode_stats computes per-episode stats for an episode subset, and aggregate_episode_stats merges the concatenated shards (count-weighted) and writes stats.json. recompute_stats now composes these, so single-process behavior is unchanged. Add examples/dataset/slurm_recompute_stats.py, a datatrove compute/aggregate driver that shards episodes across workers and is read-only safe (reference-copies the source when --new-root is given). Most useful for the expensive image/video stats path.
282 lines
11 KiB
Python
282 lines
11 KiB
Python
#!/usr/bin/env python
|
||
|
||
# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
|
||
#
|
||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||
# you may not use this file except in compliance with the License.
|
||
# You may obtain a copy of the License at
|
||
#
|
||
# http://www.apache.org/licenses/LICENSE-2.0
|
||
#
|
||
# Unless required by applicable law or agreed to in writing, software
|
||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||
# See the License for the specific language governing permissions and
|
||
# limitations under the License.
|
||
|
||
"""
|
||
SLURM-distributed recomputation of a LeRobotDataset's ``meta/stats.json``.
|
||
|
||
Per-episode statistics are embarrassingly parallel, so we shard episodes across
|
||
workers, each computing stats for its subset, then a single worker aggregates all
|
||
shards (weighted by frame counts) and writes ``meta/stats.json``. This is mostly
|
||
useful when recomputing image/video stats (``--skip-image-video 0``), which decodes
|
||
frames and is far more expensive than the numeric-only path.
|
||
|
||
Requires: pip install 'lerobot[dataset]' datatrove
|
||
|
||
Two subcommands, each a separate SLURM submission:
|
||
|
||
compute – N workers, each writes per-episode stats for its episode shard
|
||
aggregate – 1 worker, merges shards into meta/stats.json (optionally push to hub)
|
||
|
||
The dataset is read-only during ``compute``. When ``--new-root`` is given, a
|
||
lightweight reference copy is made (large files symlinked, only meta/ copied) so a
|
||
read-only / mounted source dataset is never modified; stats land in ``--new-root``.
|
||
|
||
Usage:
|
||
# Recompute image/video stats for a mounted, read-only dataset with 50 workers.
|
||
python slurm_recompute_stats.py compute \\
|
||
--repo-id someone-else/their-dataset \\
|
||
--root /path/to/mounted/repo \\
|
||
--new-root /local/writable/their-dataset_recomputed \\
|
||
--skip-image-video 0 --workers 50 --partition cpu
|
||
|
||
python slurm_recompute_stats.py aggregate \\
|
||
--repo-id someone-else/their-dataset \\
|
||
--new-root /local/writable/their-dataset_recomputed \\
|
||
--partition cpu
|
||
"""
|
||
|
||
import argparse
|
||
from pathlib import Path
|
||
|
||
from datatrove.executor import LocalPipelineExecutor
|
||
from datatrove.executor.slurm import SlurmPipelineExecutor
|
||
from datatrove.pipeline.base import PipelineStep
|
||
|
||
SHARD_PATTERN = "episode_stats_{rank:05d}.pkl"
|
||
SHARD_GLOB = "episode_stats_*.pkl"
|
||
|
||
|
||
def _load_dataset(repo_id: str, root: str | None, new_root: str | None):
|
||
"""Load the (possibly reference-copied) dataset used for stats.
|
||
|
||
When ``new_root`` differs from the source, create a read-only-safe reference copy
|
||
once (only the aggregator's rank 0 or the first compute worker needs to; here every
|
||
rank just loads ``new_root`` if it already exists, else falls back to the source).
|
||
"""
|
||
from lerobot.datasets import LeRobotDataset
|
||
|
||
if new_root and Path(new_root).exists():
|
||
return LeRobotDataset(repo_id, root=new_root)
|
||
return LeRobotDataset(repo_id, root=root)
|
||
|
||
|
||
class ComputeEpisodeStatsShards(PipelineStep):
|
||
"""Each worker computes per-episode stats for its ``episodes[rank::world_size]`` shard."""
|
||
|
||
def __init__(self, repo_id, root, new_root, skip_image_video, shard_dir):
|
||
super().__init__()
|
||
self.repo_id = repo_id
|
||
self.root = root
|
||
self.new_root = new_root
|
||
self.skip_image_video = skip_image_video
|
||
self.shard_dir = shard_dir
|
||
|
||
def run(self, data=None, rank: int = 0, world_size: int = 1):
|
||
import logging
|
||
import pickle
|
||
|
||
from lerobot.datasets import compute_dataset_episode_stats
|
||
from lerobot.utils.utils import init_logging
|
||
|
||
init_logging()
|
||
dataset = _load_dataset(self.repo_id, self.root, self.new_root)
|
||
|
||
my_episodes = list(range(dataset.meta.total_episodes))[rank::world_size]
|
||
if not my_episodes:
|
||
logging.info(f"Rank {rank}: no episodes assigned")
|
||
return
|
||
logging.info(f"Rank {rank}: {len(my_episodes)} / {dataset.meta.total_episodes} episodes")
|
||
|
||
episode_stats = compute_dataset_episode_stats(
|
||
dataset,
|
||
episode_indices=my_episodes,
|
||
skip_image_video=self.skip_image_video,
|
||
)
|
||
|
||
shard_dir = Path(self.shard_dir)
|
||
shard_dir.mkdir(parents=True, exist_ok=True)
|
||
out = shard_dir / SHARD_PATTERN.format(rank=rank)
|
||
with open(out, "wb") as f:
|
||
pickle.dump(episode_stats, f)
|
||
logging.info(f"Rank {rank}: saved {len(episode_stats)} episode stats to {out}")
|
||
|
||
|
||
class AggregateEpisodeStats(PipelineStep):
|
||
"""Merge all per-episode stat shards into meta/stats.json."""
|
||
|
||
def __init__(self, repo_id, root, new_root, shard_dir, push_to_hub=False):
|
||
super().__init__()
|
||
self.repo_id = repo_id
|
||
self.root = root
|
||
self.new_root = new_root
|
||
self.shard_dir = shard_dir
|
||
self.push_to_hub = push_to_hub
|
||
|
||
def run(self, data=None, rank: int = 0, world_size: int = 1):
|
||
import logging
|
||
import pickle
|
||
|
||
from lerobot.datasets import aggregate_episode_stats
|
||
from lerobot.utils.utils import init_logging
|
||
|
||
init_logging()
|
||
if rank != 0:
|
||
return
|
||
|
||
shard_dir = Path(self.shard_dir)
|
||
shards = sorted(shard_dir.glob(SHARD_GLOB))
|
||
if not shards:
|
||
raise FileNotFoundError(f"No episode stat shards found in {shard_dir}")
|
||
|
||
all_episode_stats = []
|
||
for shard in shards:
|
||
with open(shard, "rb") as f:
|
||
all_episode_stats.extend(pickle.load(f))
|
||
logging.info(f"Aggregating {len(all_episode_stats)} episode stats from {len(shards)} shards")
|
||
|
||
dataset = _load_dataset(self.repo_id, self.root, self.new_root)
|
||
new_stats = aggregate_episode_stats(dataset, all_episode_stats)
|
||
if new_stats is None:
|
||
raise RuntimeError("Aggregation produced no stats")
|
||
logging.info(f"Wrote stats for features: {list(new_stats.keys())} to {dataset.root}")
|
||
|
||
if self.push_to_hub:
|
||
logging.info(f"Pushing {self.repo_id} to hub")
|
||
dataset.push_to_hub()
|
||
|
||
|
||
def _make_executor(pipeline, logs_dir, job_name, slurm, workers, tasks, time, partition, cpus, mem):
|
||
kwargs = {"pipeline": pipeline, "logging_dir": str(Path(logs_dir) / job_name)}
|
||
if slurm:
|
||
kwargs.update(
|
||
{
|
||
"job_name": job_name,
|
||
"tasks": tasks,
|
||
"workers": workers,
|
||
"time": time,
|
||
"partition": partition,
|
||
"cpus_per_task": cpus,
|
||
"sbatch_args": {"mem-per-cpu": mem},
|
||
}
|
||
)
|
||
return SlurmPipelineExecutor(**kwargs)
|
||
kwargs.update({"tasks": tasks, "workers": 1})
|
||
return LocalPipelineExecutor(**kwargs)
|
||
|
||
|
||
def _maybe_reference_copy(repo_id, root, new_root):
|
||
"""Create the read-only-safe reference copy once, before submitting workers."""
|
||
if not new_root:
|
||
return
|
||
from lerobot.datasets import LeRobotDataset
|
||
from lerobot.scripts.lerobot_edit_dataset import _reference_copy_dataset
|
||
|
||
new_root_path = Path(new_root)
|
||
if new_root_path.exists():
|
||
return
|
||
src = LeRobotDataset(repo_id, root=root)
|
||
_reference_copy_dataset(src.root, new_root_path)
|
||
|
||
|
||
def _add_shared_args(p):
|
||
p.add_argument("--repo-id", type=str, required=True, help="Dataset identifier, e.g. 'user/dataset'.")
|
||
p.add_argument("--root", type=str, default=None, help="Source dataset root (e.g. a mount).")
|
||
p.add_argument(
|
||
"--new-root",
|
||
type=str,
|
||
default=None,
|
||
help="Writable output root; a read-only-safe reference copy of --root. If omitted, stats "
|
||
"are written in place at --root.",
|
||
)
|
||
p.add_argument("--shard-dir", type=Path, default=Path("stats_shards"), help="Per-rank shard dir.")
|
||
p.add_argument("--logs-dir", type=Path, default=Path("logs"), help="datatrove logs dir.")
|
||
p.add_argument("--job-name", type=str, default=None, help="SLURM job name.")
|
||
p.add_argument("--slurm", type=int, default=1, help="1 = submit via SLURM; 0 = run locally.")
|
||
p.add_argument("--partition", type=str, default=None, help="SLURM partition.")
|
||
p.add_argument("--cpus-per-task", type=int, default=4, help="CPUs per SLURM task.")
|
||
p.add_argument("--mem-per-cpu", type=str, default="4G", help="Memory per CPU, e.g. '4G'.")
|
||
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser(
|
||
description="SLURM-distributed LeRobotDataset stats recomputation",
|
||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||
)
|
||
sub = parser.add_subparsers(dest="command", required=True)
|
||
|
||
cp = sub.add_parser("compute", help="Distribute per-episode stats across SLURM workers.")
|
||
_add_shared_args(cp)
|
||
cp.add_argument("--workers", type=int, default=50, help="Number of parallel SLURM tasks.")
|
||
cp.add_argument(
|
||
"--skip-image-video",
|
||
type=int,
|
||
default=1,
|
||
help="1 = numeric features only (fast); 0 = also recompute image/video stats (decodes frames).",
|
||
)
|
||
|
||
ap = sub.add_parser("aggregate", help="Merge shards into meta/stats.json.")
|
||
_add_shared_args(ap)
|
||
ap.add_argument("--push-to-hub", action="store_true", help="Push the dataset after aggregation.")
|
||
|
||
args = parser.parse_args()
|
||
slurm = args.slurm == 1
|
||
|
||
if args.command == "compute":
|
||
# The reference copy (if any) is created once on the submitting node so workers
|
||
# can all load --new-root without racing to build it.
|
||
_maybe_reference_copy(args.repo_id, args.root, args.new_root)
|
||
job_name = args.job_name or "recompute_stats_compute"
|
||
executor = _make_executor(
|
||
pipeline=[
|
||
ComputeEpisodeStatsShards(
|
||
args.repo_id, args.root, args.new_root, args.skip_image_video == 0, str(args.shard_dir)
|
||
)
|
||
],
|
||
logs_dir=args.logs_dir,
|
||
job_name=job_name,
|
||
slurm=slurm,
|
||
workers=args.workers,
|
||
tasks=args.workers,
|
||
time="24:00:00",
|
||
partition=args.partition,
|
||
cpus=args.cpus_per_task,
|
||
mem=args.mem_per_cpu,
|
||
)
|
||
else:
|
||
job_name = args.job_name or "recompute_stats_aggregate"
|
||
executor = _make_executor(
|
||
pipeline=[
|
||
AggregateEpisodeStats(
|
||
args.repo_id, args.root, args.new_root, str(args.shard_dir), args.push_to_hub
|
||
)
|
||
],
|
||
logs_dir=args.logs_dir,
|
||
job_name=job_name,
|
||
slurm=slurm,
|
||
workers=1,
|
||
tasks=1,
|
||
time="02:00:00",
|
||
partition=args.partition,
|
||
cpus=args.cpus_per_task,
|
||
mem=args.mem_per_cpu,
|
||
)
|
||
|
||
executor.run()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|