admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-07-11 12:01:52 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor(examples): replace hf-mount with HF_LEROBOT_HOME + venv/env-command in slurm stats script

Browse Source 
Drops the per-worker hf-mount machinery in favor of a shared HF_LEROBOT_HOME
cache plus --venv-path and --env-command hooks, which is simpler and avoids
node-local mount setup. Removes the now-unused os import.
This commit is contained in:
CarolinePascal
2026-07-10 18:36:13 +02:00
parent e1e9934a78
commit fbe9c11b60
1 changed files with 41 additions and 103 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/dataset/slurm_recompute_stats.py
+41 -103
View File
@@ -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 <your-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=<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:<compute jobid>
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