From 7bee7fb9e3fab3566c62d920bb66a42ee69cc7ab Mon Sep 17 00:00:00 2001 From: CarolinePascal Date: Fri, 10 Jul 2026 14:47:18 +0200 Subject: [PATCH] feat(datasets): distribute stats recomputation across SLURM workers 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. --- examples/dataset/slurm_recompute_stats.py | 281 ++++++++++++++++++++++ src/lerobot/datasets/__init__.py | 4 + src/lerobot/datasets/dataset_tools.py | 158 +++++++----- 3 files changed, 389 insertions(+), 54 deletions(-) create mode 100644 examples/dataset/slurm_recompute_stats.py diff --git a/examples/dataset/slurm_recompute_stats.py b/examples/dataset/slurm_recompute_stats.py new file mode 100644 index 000000000..b4b942338 --- /dev/null +++ b/examples/dataset/slurm_recompute_stats.py @@ -0,0 +1,281 @@ +#!/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() diff --git a/src/lerobot/datasets/__init__.py b/src/lerobot/datasets/__init__.py index 7715a115e..9c6342898 100644 --- a/src/lerobot/datasets/__init__.py +++ b/src/lerobot/datasets/__init__.py @@ -25,6 +25,8 @@ from .compute_stats import DEFAULT_QUANTILES, aggregate_stats, get_feature_stats from .dataset_metadata import CODEBASE_VERSION, LeRobotDatasetMetadata from .dataset_tools import ( add_features, + aggregate_episode_stats, + compute_dataset_episode_stats, convert_image_to_video_dataset, delete_episodes, merge_datasets, @@ -78,8 +80,10 @@ __all__ = [ "detect_available_encoders_pyav", "add_features", "aggregate_datasets", + "aggregate_episode_stats", "aggregate_pipeline_dataset_features", "aggregate_stats", + "compute_dataset_episode_stats", "convert_image_to_video_dataset", "create_initial_features", "compute_sampler_state", diff --git a/src/lerobot/datasets/dataset_tools.py b/src/lerobot/datasets/dataset_tools.py index ebc53d1c8..1a4286175 100644 --- a/src/lerobot/datasets/dataset_tools.py +++ b/src/lerobot/datasets/dataset_tools.py @@ -1639,6 +1639,76 @@ def _compute_visual_episode_stats( return ep_stats +def compute_dataset_episode_stats( + dataset: LeRobotDataset, + episode_indices: list[int] | None = None, + skip_image_video: bool = True, + drop_keys: list[str] | None = None, +) -> list[dict]: + """Compute per-episode statistics for a subset of episodes. + + This is the shardable unit of work behind :func:`recompute_stats`: distribute + ``episode_indices`` across workers (e.g. ``list(range(n))[rank::world_size]``), + then combine the concatenated results with :func:`aggregate_stats`. + + Args: + dataset: The LeRobotDataset to compute stats for. + episode_indices: Episodes to process. When ``None``, all episodes are processed. + skip_image_video: If True (default), only numeric features are computed. If False, + image/video stats are also computed by sampling and decoding frames. + drop_keys: Feature keys to exclude (e.g. ``action`` when it is computed separately + in relative-action space). + + Returns: + A list of per-episode stat dicts, one per processed episode. + """ + features = dataset.meta.features + meta_keys = {"index", "episode_index", "task_index", "frame_index", "timestamp"} + drop = set(drop_keys or []) + features_to_compute = { + k: v + for k, v in features.items() + if v["dtype"] != "string" + and k not in meta_keys + and k not in drop + and (not skip_image_video or v["dtype"] not in ["image", "video"]) + } + numeric_keys = [k for k, v in features_to_compute.items() if v["dtype"] not in ["image", "video"]] + visual_keys = [k for k, v in features_to_compute.items() if v["dtype"] in ["image", "video"]] + + if dataset.meta.episodes is None: + dataset.meta.episodes = load_episodes(dataset.meta.root) + + if episode_indices is None: + episode_indices = list(range(dataset.meta.total_episodes)) + + # Group requested episodes by their data parquet file so each file is read once. + file_to_episodes: dict[Path, list[int]] = {} + for ep_idx in episode_indices: + file_to_episodes.setdefault(dataset.meta.get_data_file_path(ep_idx), []).append(ep_idx) + + all_episode_stats = [] + for src_path, eps in tqdm(sorted(file_to_episodes.items()), desc="Computing stats from data files"): + df = pd.read_parquet(dataset.root / src_path) if numeric_keys else None + for ep_idx in sorted(eps): + episode_data = {} + if numeric_keys: + ep_df = df[df["episode_index"] == ep_idx] + for key in numeric_keys: + if key in ep_df.columns: + values = ep_df[key].values + episode_data[key] = ( + np.stack(values) if hasattr(values[0], "__len__") else np.array(values) + ) + + ep_stats = compute_episode_stats(episode_data, features_to_compute) + if visual_keys: + ep_stats.update(_compute_visual_episode_stats(dataset, int(ep_idx), visual_keys)) + all_episode_stats.append(ep_stats) + + return all_episode_stats + + def recompute_stats( dataset: LeRobotDataset, skip_image_video: bool = True, @@ -1670,24 +1740,12 @@ def recompute_stats( The same dataset with updated stats. """ features = dataset.meta.features - meta_keys = {"index", "episode_index", "task_index", "frame_index", "timestamp"} - numeric_features = { - k: v - for k, v in features.items() - if v["dtype"] not in ["image", "video", "string"] and k not in meta_keys - } - - if skip_image_video: - features_to_compute = numeric_features - else: - features_to_compute = { - k: v for k, v in features.items() if v["dtype"] != "string" and k not in meta_keys - } # When relative_action is enabled, compute action stats via chunk-based sampling # (matching what the model sees during training) and skip action in the # per-episode pass below. relative_action_stats = None + drop_keys = None if relative_action and ACTION in features and OBS_STATE in features: if relative_exclude_joints is None: relative_exclude_joints = ["gripper"] @@ -1698,58 +1756,50 @@ def recompute_stats( exclude_joints=relative_exclude_joints, num_workers=num_workers, ) - features_to_compute.pop(ACTION, None) + drop_keys = [ACTION] - logging.info(f"Recomputing stats for features: {list(features_to_compute.keys())}") + all_episode_stats = compute_dataset_episode_stats( + dataset, skip_image_video=skip_image_video, drop_keys=drop_keys + ) - data_dir = dataset.root / DATA_DIR - parquet_files = sorted(data_dir.glob("*/*.parquet")) - if not parquet_files: - raise ValueError(f"No parquet files found in {data_dir}") - - all_episode_stats = [] - numeric_keys = [k for k, v in features_to_compute.items() if v["dtype"] not in ["image", "video"]] - visual_keys = [k for k, v in features_to_compute.items() if v["dtype"] in ["image", "video"]] - - for parquet_path in tqdm(parquet_files, desc="Computing stats from data files"): - df = pd.read_parquet(parquet_path) - - for ep_idx in sorted(df["episode_index"].unique()): - ep_df = df[df["episode_index"] == ep_idx] - episode_data = {} - for key in numeric_keys: - if key in ep_df.columns: - values = ep_df[key].values - if hasattr(values[0], "__len__"): - episode_data[key] = np.stack(values) - else: - episode_data[key] = np.array(values) - - ep_stats = compute_episode_stats(episode_data, features_to_compute) - if visual_keys: - ep_stats.update(_compute_visual_episode_stats(dataset, int(ep_idx), visual_keys)) - all_episode_stats.append(ep_stats) - - if features_to_compute and not all_episode_stats: + new_stats = aggregate_episode_stats( + dataset, all_episode_stats, extra_stats={ACTION: relative_action_stats} if relative_action_stats else None + ) + if new_stats is None: logging.warning("No episode stats computed") - return dataset + else: + logging.info("Stats recomputed successfully") + return dataset + + +def aggregate_episode_stats( + dataset: LeRobotDataset, + all_episode_stats: list[dict], + extra_stats: dict | None = None, +) -> dict | None: + """Aggregate per-episode stats, merge with existing stats, and write ``stats.json``. + + Companion to :func:`compute_dataset_episode_stats` for the distributed workflow: pass + the concatenation of every worker's per-episode stats. ``extra_stats`` lets callers + inject feature stats computed outside the per-episode pass (e.g. relative-action stats). + + Returns the written stats dict, or ``None`` if there was nothing to aggregate. + """ + if not all_episode_stats and not extra_stats: + return None new_stats = aggregate_stats(all_episode_stats) if all_episode_stats else {} + if extra_stats: + new_stats.update(extra_stats) - if relative_action_stats is not None: - new_stats[ACTION] = relative_action_stats - - # Merge: keep existing stats for features we didn't recompute + # Merge: keep existing stats for features we didn't recompute. if dataset.meta.stats: for key, value in dataset.meta.stats.items(): - if key not in new_stats: - new_stats[key] = value + new_stats.setdefault(key, value) write_stats(new_stats, dataset.root) dataset.meta.stats = new_stats - - logging.info("Stats recomputed successfully") - return dataset + return new_stats def convert_image_to_video_dataset(