diff --git a/examples/dataset/slurm_recompute_stats.py b/examples/dataset/slurm_recompute_stats.py index 4217722b1..a0eb7080e 100644 --- a/examples/dataset/slurm_recompute_stats.py +++ b/examples/dataset/slurm_recompute_stats.py @@ -17,36 +17,33 @@ """ SLURM-distributed recomputation of a LeRobotDataset's ``meta/stats.json``. -This is a modified copy of lerobot's examples/dataset/slurm_recompute_stats.py -(feat/recompute-stats-readonly-and-visual branch) with three additions relevant -to a shared HPC cluster: +Modified copy of lerobot's examples/dataset/slurm_recompute_stats.py +(feat/recompute-stats-readonly-and-visual branch) with cluster-friendly additions: - 1. --qos : pass a SLURM QoS through to every worker's sbatch. - 2. per-worker hf-mount : each worker mounts the read-only source dataset on - its OWN node's /scratch before loading it, injected - via datatrove's ``env_command`` hook. This keeps the - terabytes of reads node-local and lazy (nothing piles - up on /fsx) and keeps hub traffic on the CPU nodes. - 3. --chain-aggregate : submit ``aggregate`` with an afterok dependency on - ``compute`` so it only runs once all shards exist - (no manual squeue-wait, no gap/overlap race). + 1. --qos : pass a SLURM QoS through to every worker's sbatch. + 2. --venv-path : activate a venv on each worker before the python step. + 3. --env-command : raw shell snippet injected before the python step (e.g. to + export HF_LEROBOT_HOME). Overrides --venv-path if given. + 4. --chain-aggregate : submit ``aggregate`` with an afterok dependency on + ``compute`` so it only runs once all shards exist + (no manual squeue-wait, no gap/overlap race). + +Data access: no filesystem mount. Point HF_LEROBOT_HOME at a node-visible shared +cache (e.g. /fsx/$USER/.cache) so the dataset downloads once and all workers read +it. This is the download route; the source dataset is fetched from the Hub on the +CPU workers. IMPORTANT — how to run (do NOT sbatch this file): Run it as a normal python process on the LOGIN node. datatrove submits the - workers for you. Because the reference copy (--new-root) walks the source tree - on the login node, the source must also be mountable there — so mount once on - the login node too, before launching (see the mount snippet below). + workers for you. The reference copy (--new-root) is built on the login node and + references the shared HF cache, so /fsx must be visible there (it is). Requires: pip install 'lerobot[dataset]' datatrove Example (single command, compute then dependent aggregate): - # 0. Mount on the login node so the reference-copy walk can list the source. - /fsx/$USER/bin/hf-mount-nfs-x86_64-linux \ - repo datasets/behavior-1k/2026-challenge-demos /scratch/$USER/behavior-demos \ - --cache-dir /scratch/$USER/hfmount-cache --cache-size 100000000000 & + export HF_LEROBOT_HOME=/fsx/$USER/.cache - # 1. Launch. Each worker will mount the source on its own node via --mount-repo. python slurm_recompute_stats_patched.py compute \ --repo-id behavior-1k/2026-challenge-demos \ --new-root /fsx/$USER/behavior-1k_recomputed \ @@ -55,19 +52,17 @@ Example (single command, compute then dependent aggregate): --skip-image-video 0 \ --workers 250 \ --partition hopper-cpu \ - --qos \ + --qos normal \ --cpus-per-task 8 --mem-per-cpu 4G \ - --mount-repo datasets/behavior-1k/2026-challenge-demos \ - --hf-mount-bin /fsx/$USER/bin/hf-mount-nfs-x86_64-linux \ --venv-path /fsx/$USER/venvs/lerobot/bin/activate \ + --env-command 'export HF_LEROBOT_HOME=/fsx/'"$USER"'/.cache' \ --chain-aggregate -REHEARSE FIRST with --workers 2 and inspect one worker's log under --logs-dir to -confirm the mount came up and video decoding ran (not a silent hub download). +REHEARSE FIRST with --workers 2 --skip-image-video 1 and inspect one worker's log +under --logs-dir to confirm QoS was accepted and a numeric stats.json is written. """ import argparse -import os from pathlib import Path from datatrove.executor import LocalPipelineExecutor @@ -216,41 +211,6 @@ def _mem_gb(mem: str) -> int: return int(float(s)) -def _build_env_command(args) -> str | None: - """Construct the per-worker shell snippet datatrove runs before the python step. - - Mounts the read-only source dataset on THIS worker's node-local /scratch, waits - for it to come up, and fails LOUDLY (exit 1) if it doesn't — so a broken mount - surfaces as a failed job instead of a silent fall-back to downloading the dataset. - Also activates the venv. Returns None if --mount-repo was not requested (in which - case you must supply --root yourself and datatrove uses --venv-path only). - """ - if args.env_command: - return args.env_command - - lines = [] - if args.venv_path: - lines.append(f"source {args.venv_path}") - - if args.mount_repo: - if not args.hf_mount_bin: - raise SystemExit("--mount-repo requires --hf-mount-bin") - mnt = args.mount_point - cache = args.mount_cache_dir - lines += [ - f'MNT="{mnt}"', - f'CACHE="{cache}"', - 'mkdir -p "$MNT" "$CACHE"', - f'{args.hf_mount_bin} \\', - f' repo {args.mount_repo} "$MNT" \\', - f' --cache-dir "$CACHE" --cache-size {args.mount_cache_size} &', - 'for i in $(seq 1 60); do [ -f "$MNT/meta/info.json" ] && break; sleep 2; done', - '[ -f "$MNT/meta/info.json" ] || { echo "hf-mount failed to come up at $MNT" >&2; exit 1; }', - ] - - return "\n".join(lines) if lines else None - - def _make_executor( pipeline, logs_dir, @@ -264,6 +224,7 @@ def _make_executor( mem, qos=None, env_command=None, + venv_path=None, depends=None, ): kwargs = {"pipeline": pipeline, "logging_dir": str(Path(logs_dir) / job_name)} @@ -283,7 +244,9 @@ def _make_executor( if qos: kwargs["qos"] = qos # -> "#SBATCH --qos=" on every worker if env_command: - kwargs["env_command"] = env_command # per-worker mount + venv, runs before python + kwargs["env_command"] = env_command # raw snippet before python (overrides venv_path) + elif venv_path: + kwargs["venv_path"] = venv_path # datatrove sources this before the python step if depends is not None: kwargs["depends"] = depends # chains --dependency=afterok: return SlurmPipelineExecutor(**kwargs) @@ -305,9 +268,9 @@ def _maybe_reference_copy(repo_id, root, new_root): _reference_copy_dataset(src.root, new_root_path) -def _add_shared_args(p, user): +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("--root", type=str, default=None, help="Source dataset root (defaults to the Hub cache).") p.add_argument( "--new-root", type=str, @@ -320,7 +283,7 @@ def _add_shared_args(p, user): 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, e.g. 'hopper-cpu'.") - p.add_argument("--qos", type=str, default=None, help="SLURM QoS, e.g. 'high'. Passed to every worker.") + p.add_argument("--qos", type=str, default=None, help="SLURM QoS, e.g. 'normal'. Passed to every worker.") 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'.") p.add_argument( @@ -330,37 +293,17 @@ def _add_shared_args(p, user): help="Video decoding backend (e.g. 'pyav', 'torchcodec'). Defaults to the dataset's default; " "use 'pyav' if torchcodec fails to load locally.", ) - - # --- per-worker mount options (patch) --- + p.add_argument("--venv-path", type=str, default=None, help="venv activate script sourced on each worker.") p.add_argument( "--env-command", type=str, default=None, - help="Raw shell snippet injected into each worker's sbatch before the python step. " - "Overrides the auto-generated mount snippet if given.", + help="Raw shell snippet injected into each worker's sbatch before the python step " + "(e.g. to export HF_LEROBOT_HOME). Overrides --venv-path if given.", ) - p.add_argument( - "--mount-repo", - type=str, - default=None, - help="If set, each worker mounts this repo (e.g. 'datasets/user/name') on its own node " - "via hf-mount before loading the dataset. Auto-sets --root to --mount-point if --root unset.", - ) - p.add_argument("--hf-mount-bin", type=str, default=None, help="Path to the hf-mount NFS binary.") - p.add_argument("--venv-path", type=str, default=None, help="Path to a venv activate script to source.") - p.add_argument( - "--mount-point", - type=str, - default=f"/scratch/{user}/behavior-demos", - help="Node-local mount path (must be identical on every node).", - ) - p.add_argument("--mount-cache-dir", type=str, default=f"/scratch/{user}/hfmount-cache") - p.add_argument("--mount-cache-size", type=str, default="100000000000", help="hf-mount --cache-size bytes.") def main(): - user = os.environ.get("USER", "user") - parser = argparse.ArgumentParser( description="PATCHED SLURM-distributed LeRobotDataset stats recomputation", formatter_class=argparse.RawDescriptionHelpFormatter, @@ -368,7 +311,7 @@ def main(): 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, user) + _add_shared_args(cp) cp.add_argument("--workers", type=int, default=50, help="Number of parallel SLURM tasks.") cp.add_argument( "--skip-image-video", @@ -384,7 +327,7 @@ def main(): cp.add_argument("--push-to-hub", action="store_true", help="For the chained aggregate: push after done.") ap = sub.add_parser("aggregate", help="Merge shards into meta/stats.json.") - _add_shared_args(ap, user) + _add_shared_args(ap) ap.add_argument("--push-to-hub", action="store_true", help="Push the dataset after aggregation.") ap.add_argument( "--depends-job-id", @@ -396,19 +339,11 @@ def main(): args = parser.parse_args() slurm = args.slurm == 1 - # If a per-worker mount is requested and --root wasn't given, workers read from the mount. - if args.mount_repo and not args.root: - args.root = args.mount_point - - env_command = _build_env_command(args) - 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. NOTE: this walks the source - # tree, so the source must be mountable on the login node too. + # can all load --new-root without racing to build it. _maybe_reference_copy(args.repo_id, args.root, args.new_root) - compute_job_name = args.job_name or "recompute_stats_compute" compute_exec = _make_executor( pipeline=[ ComputeEpisodeStatsShards( @@ -421,7 +356,7 @@ def main(): ) ], logs_dir=args.logs_dir, - job_name=compute_job_name, + job_name=args.job_name or "recompute_stats_compute", slurm=slurm, workers=args.workers, tasks=args.workers, @@ -430,7 +365,8 @@ def main(): cpus=args.cpus_per_task, mem=args.mem_per_cpu, qos=args.qos, - env_command=env_command, + env_command=args.env_command, + venv_path=args.venv_path, ) if args.chain_aggregate and slurm: @@ -457,7 +393,8 @@ def main(): cpus=args.cpus_per_task, mem=args.mem_per_cpu, qos=args.qos, - env_command=env_command, # aggregate also needs the mount to load the dataset + env_command=args.env_command, + venv_path=args.venv_path, depends=compute_exec, ) aggregate_exec.run() @@ -485,7 +422,8 @@ def main(): cpus=args.cpus_per_task, mem=args.mem_per_cpu, qos=args.qos, - env_command=env_command, + env_command=args.env_command, + venv_path=args.venv_path, ) if args.depends_job_id is not None: aggregate_exec.depends_job_id = args.depends_job_id