mirror of
https://github.com/huggingface/lerobot.git
synced 2026-05-15 00:29:52 +00:00
Pepijn
143 lines
4.4 KiB
Python
143 lines
4.4 KiB
Python
#!/usr/bin/env python
|
|
|
|
# Copyright 2026 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.
|
|
|
|
import json
|
|
|
|
from benchmarks.run_benchmark_matrix import (
|
|
PlannedJob,
|
|
compute_gradient_accumulation_steps,
|
|
plan_jobs,
|
|
render_sbatch_script,
|
|
write_manifest,
|
|
)
|
|
|
|
|
|
def _one_job(job_list: list[PlannedJob]) -> PlannedJob:
|
|
assert len(job_list) == 1
|
|
return job_list[0]
|
|
|
|
|
|
def test_compute_gradient_accumulation_steps_for_fixed_effective_batch():
|
|
assert compute_gradient_accumulation_steps(
|
|
effective_batch_size=256,
|
|
num_gpus=8,
|
|
microbatch_per_gpu=32,
|
|
) == 1
|
|
assert compute_gradient_accumulation_steps(
|
|
effective_batch_size=256,
|
|
num_gpus=4,
|
|
microbatch_per_gpu=32,
|
|
) == 2
|
|
assert compute_gradient_accumulation_steps(
|
|
effective_batch_size=256,
|
|
num_gpus=1,
|
|
microbatch_per_gpu=32,
|
|
) == 8
|
|
|
|
|
|
def test_plan_jobs_filters_libero_plus_only(tmp_path):
|
|
jobs = plan_jobs(
|
|
output_dir=tmp_path,
|
|
hub_org="lerobot",
|
|
results_repo="lerobot/benchmark-history",
|
|
policies=["pi0", "act"],
|
|
benchmarks=["libero_plus"],
|
|
)
|
|
|
|
assert [job.benchmark for job in jobs] == ["libero_plus", "libero_plus"]
|
|
assert [job.policy for job in jobs] == ["pi0", "act"]
|
|
|
|
|
|
def test_plan_jobs_includes_libero_plus_and_robomme(tmp_path):
|
|
jobs = plan_jobs(
|
|
output_dir=tmp_path,
|
|
hub_org="lerobot",
|
|
results_repo="lerobot/benchmark-history",
|
|
policies=["pi0"],
|
|
benchmarks=["libero_plus", "robomme"],
|
|
)
|
|
|
|
assert [job.benchmark for job in jobs] == ["libero_plus", "robomme"]
|
|
assert jobs[0].effective_batch_size == 256
|
|
assert jobs[1].effective_batch_size == 256
|
|
|
|
|
|
def test_plan_jobs_sets_expected_gpu_and_accumulation(tmp_path):
|
|
jobs = plan_jobs(
|
|
output_dir=tmp_path,
|
|
hub_org="lerobot",
|
|
results_repo="lerobot/benchmark-history",
|
|
policies=["pi0", "xvla", "act"],
|
|
benchmarks=["robomme"],
|
|
)
|
|
by_policy = {job.policy: job for job in jobs}
|
|
|
|
assert by_policy["pi0"].num_gpus == 8
|
|
assert by_policy["pi0"].gradient_accumulation_steps == 1
|
|
assert by_policy["xvla"].num_gpus == 4
|
|
assert by_policy["xvla"].gradient_accumulation_steps == 2
|
|
assert by_policy["act"].num_gpus == 1
|
|
assert by_policy["act"].gradient_accumulation_steps == 8
|
|
|
|
|
|
def test_render_sbatch_script_contains_train_eval_and_publish(tmp_path):
|
|
job = _one_job(
|
|
plan_jobs(
|
|
output_dir=tmp_path,
|
|
hub_org="lerobot",
|
|
results_repo="lerobot/benchmark-history",
|
|
policies=["pi0_fast"],
|
|
benchmarks=["robomme"],
|
|
)
|
|
)
|
|
|
|
script = render_sbatch_script(
|
|
job=job,
|
|
output_dir=tmp_path,
|
|
results_repo_id="lerobot/benchmark-history",
|
|
git_commit="deadbeef",
|
|
)
|
|
|
|
assert "docker/Dockerfile" not in script
|
|
assert "lerobot-benchmark-robomme:latest" in script
|
|
assert '--dataset.repo_id="lerobot/robomme"' in script
|
|
assert '--env.type="robomme"' in script
|
|
assert "--gradient_accumulation_steps=1" in script
|
|
assert "lerobot-train-tokenizer" in script
|
|
assert "benchmarks/publish_benchmark_result.py" in script
|
|
|
|
|
|
def test_write_manifest_records_job_metadata(tmp_path):
|
|
jobs = plan_jobs(
|
|
output_dir=tmp_path,
|
|
hub_org="lerobot",
|
|
results_repo="lerobot/benchmark-history",
|
|
policies=["pi0"],
|
|
benchmarks=["libero_plus", "robomme"],
|
|
)
|
|
manifest_path = write_manifest(
|
|
output_dir=tmp_path,
|
|
jobs=jobs,
|
|
git_commit="deadbeef",
|
|
hub_org="lerobot",
|
|
results_repo="lerobot/benchmark-history",
|
|
)
|
|
|
|
manifest = json.loads(manifest_path.read_text())
|
|
assert manifest["git_commit"] == "deadbeef"
|
|
assert manifest["results_repo"] == "lerobot/benchmark-history"
|
|
assert [job["benchmark"] for job in manifest["jobs"]] == ["libero_plus", "robomme"]
|