mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-24 18:26:11 +00:00
Pepijn
9c82c39c7b
* feat(annotate): run lerobot-annotate on HF Jobs via --job.target Annotation needed a hand-edited launcher script (examples/annotations/run_hf_job.py) to reach a GPU: users copied it, rewrote the embedded CMD string for their dataset, and ran it with `python`. Fold that into the CLI instead, mirroring `lerobot-train`: `lerobot-annotate --job.target=h200` submits the exact command you'd run locally. - AnnotationJobConfig extends JobConfig with the annotation runtime's defaults (vllm/vllm-openai image, 2h cap) plus --job.lerobot_ref, so an unmerged branch can be exercised remotely without editing a script. - lerobot.jobs.annotate builds the pod command by replaying the user's own CLI flags (minus --job.*/--root, with --repo_id re-emitted from the config) after a setup prelude that installs lerobot on top of the vLLM image. Job monitoring, log tailing and Ctrl-C-detaches reuse the training submitter's plumbing. - Remote runs require --repo_id; a local-only dataset is pushed privately first. The generated pod command is byte-for-byte the script's old CMD. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * fix(annotate): reject client-side config files on remote runs draccus exposes `--config_path` plus a `--<field>` config-file arg for every nested dataclass (`--vlm`, `--plan`, `--job`, ...). All name files on the client's disk, so forwarding them to the pod silently dropped whatever settings they carried. Reject them up front instead. Bare `--job` also slipped past the `--job.` prefix filter, so a `--job=cfg.yaml` holding `target: h200` would have reached the pod and had the job submit a job of its own, recursively. It is dropped from the forwarded args as well. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * refactor(jobs): share the submit-and-follow loop between both submitters `submit_annotate_to_hf` reused the leaf helpers (`_poll_until_done`, `_tail_logs`, `_pod_forwarded_args`) but duplicated the orchestration around them: ~40 of the 50 lines that spawn the poll/log threads, install the Ctrl-C-detaches handler and raise on a non-COMPLETED stage were identical in both files. Extract that into `follow_job(job_id, *, detach, success_marker=None) -> bool`, returning True when the job finished and False when we stopped watching without a verdict (detach or Ctrl-C). Training keeps its model-pushed marker by passing it in; annotation has no equivalent line (the CLI keeps working after the upload log to write the card and tag) so its completion stays stage-based. Kept in hf.py rather than a new module so every existing monkeypatch target in test_hf.py still resolves. Behaviour change: a training run whose job reaches COMPLETED without the marker matching now prints its completion line instead of returning silently. The marker was already documented as an optimisation with a stage-based fallback; the fallback just never reported success. Tests: adds annotate coverage for the non-detach path (completion and failure) — previously only ever exercised with detach=true — plus a detach short-circuit test. Both new annotate tests verified to fail under a mutation that stubs out follow_job. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
246 lines
10 KiB
Python
246 lines
10 KiB
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 shlex
|
|
import sys
|
|
from unittest.mock import MagicMock
|
|
|
|
import draccus
|
|
import pytest
|
|
|
|
pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
|
|
|
|
from lerobot.annotations.steerable_pipeline.config import (
|
|
DEFAULT_ANNOTATE_JOB_IMAGE,
|
|
AnnotationJobConfig,
|
|
AnnotationPipelineConfig,
|
|
)
|
|
from lerobot.jobs.annotate import build_pod_command, build_pod_setup, submit_annotate_to_hf
|
|
|
|
|
|
def _parse(*args):
|
|
return draccus.parse(AnnotationPipelineConfig, args=list(args))
|
|
|
|
|
|
def _set_argv(monkeypatch, *args):
|
|
monkeypatch.setattr(sys, "argv", ["lerobot-annotate", *args])
|
|
|
|
|
|
# --- config ----------------------------------------------------------------
|
|
|
|
|
|
def test_annotation_job_defaults_are_local_with_vllm_image():
|
|
cfg = AnnotationJobConfig()
|
|
assert cfg.target is None
|
|
assert cfg.is_remote is False
|
|
assert cfg.image == DEFAULT_ANNOTATE_JOB_IMAGE
|
|
assert cfg.timeout == "2h"
|
|
assert cfg.lerobot_ref == "main"
|
|
|
|
|
|
def test_annotation_config_parses_job_target():
|
|
cfg = _parse("--repo_id", "u/d", "--job.target", "h200")
|
|
assert cfg.job.target == "h200"
|
|
assert cfg.job.is_remote is True
|
|
|
|
|
|
def test_annotation_config_defaults_to_local():
|
|
assert _parse("--repo_id", "u/d").job.is_remote is False
|
|
|
|
|
|
# --- pod command -----------------------------------------------------------
|
|
|
|
|
|
def test_pod_setup_installs_requested_ref():
|
|
setup = build_pod_setup("my-branch")
|
|
assert "git+https://github.com/huggingface/lerobot.git@my-branch" in setup
|
|
# The vLLM image has neither ffmpeg (video decode) nor lerobot's pinned deps.
|
|
assert "ffmpeg" in setup
|
|
assert "'draccus==0.10.0'" in setup
|
|
|
|
|
|
def _annotate_argv(command):
|
|
"""Extract the `lerobot-annotate ...` argv from a `bash -c` pod command."""
|
|
assert command[:2] == ["bash", "-c"]
|
|
_setup, _, annotate = command[2].rpartition(" && ")
|
|
return shlex.split(annotate)
|
|
|
|
|
|
def test_pod_command_forwards_user_flags_and_pins_local_target():
|
|
command = build_pod_command(
|
|
"u/d",
|
|
"main",
|
|
["--repo_id=u/d", "--new_repo_id=u/d_annotated", "--push_to_hub=true", "--job.target=h200"],
|
|
)
|
|
argv = _annotate_argv(command)
|
|
assert argv[0] == "lerobot-annotate"
|
|
# --job.* is client-side orchestration; the pod must not re-dispatch itself.
|
|
assert not any(a.startswith("--job.") for a in argv[1:-1])
|
|
assert argv[-1] == "--job.target=local"
|
|
assert "--new_repo_id=u/d_annotated" in argv
|
|
assert "--push_to_hub=true" in argv
|
|
|
|
|
|
def test_pod_command_replaces_host_local_root_with_repo_id():
|
|
"""--root points at a directory only the client has; the pod resolves by repo_id."""
|
|
command = build_pod_command("u/d", "main", ["--root", "/home/me/datasets/d", "--seed=7"])
|
|
argv = _annotate_argv(command)
|
|
assert "--root" not in argv
|
|
assert "/home/me/datasets/d" not in argv
|
|
assert argv.count("--repo_id=u/d") == 1
|
|
assert "--seed=7" in argv
|
|
|
|
|
|
def test_pod_command_does_not_duplicate_repo_id():
|
|
command = build_pod_command("u/d", "main", ["--repo_id", "u/d"])
|
|
assert _annotate_argv(command).count("--repo_id=u/d") == 1
|
|
|
|
|
|
def test_pod_command_quotes_flags_containing_spaces_and_json():
|
|
"""serve_command and chat_template_kwargs must survive the trip through `bash -c`."""
|
|
serve = "--vlm.serve_command=vllm serve Qwen/Qwen3.6-27B --max-model-len 32768 --port {port}"
|
|
kwargs = '--vlm.chat_template_kwargs={"enable_thinking": false}'
|
|
command = build_pod_command("u/d", "main", [serve, kwargs])
|
|
argv = _annotate_argv(command)
|
|
assert serve in argv
|
|
assert kwargs in argv
|
|
|
|
|
|
# --- submission ------------------------------------------------------------
|
|
|
|
|
|
def test_submit_requires_login(monkeypatch):
|
|
monkeypatch.setattr("lerobot.jobs.annotate.get_token", lambda: None)
|
|
with pytest.raises(RuntimeError, match="hf auth login"):
|
|
submit_annotate_to_hf(_parse("--repo_id", "u/d", "--job.target", "h200"))
|
|
|
|
|
|
def test_submit_requires_repo_id(monkeypatch):
|
|
"""A remote run over --root alone can't work: the pod can't see the client's disk."""
|
|
monkeypatch.setattr("lerobot.jobs.annotate.get_token", lambda: "tok")
|
|
cfg = _parse("--root", "/tmp/d", "--job.target", "h200")
|
|
with pytest.raises(ValueError, match="--repo_id"):
|
|
submit_annotate_to_hf(cfg)
|
|
|
|
|
|
@pytest.mark.parametrize("arg", ["--config_path=annotate.yaml", "--vlm=vlm.yaml", "--job=job.yaml"])
|
|
def test_submit_rejects_local_config_files(monkeypatch, arg):
|
|
"""draccus takes a config file for the whole config and for each nested one; the
|
|
pod can read none of them, so a remote run must refuse rather than drop them."""
|
|
monkeypatch.setattr("lerobot.jobs.annotate.get_token", lambda: "tok")
|
|
_set_argv(monkeypatch, arg, "--job.target=h200")
|
|
cfg = _parse("--repo_id", "u/d", "--job.target", "h200")
|
|
with pytest.raises(ValueError, match="cannot read config files"):
|
|
submit_annotate_to_hf(cfg)
|
|
|
|
|
|
def test_pod_command_drops_bare_job_config_file_arg():
|
|
"""`--job` isn't caught by the `--job.` prefix, and could carry a remote target
|
|
that would make the pod submit a job of its own — recursively."""
|
|
argv = _annotate_argv(build_pod_command("u/d", "main", ["--job", "job.yaml", "--seed=7"]))
|
|
assert "--job" not in argv
|
|
assert "job.yaml" not in argv
|
|
assert argv[-1] == "--job.target=local"
|
|
|
|
|
|
def test_submit_dispatches_job(monkeypatch):
|
|
monkeypatch.setattr("lerobot.jobs.annotate.get_token", lambda: "tok")
|
|
monkeypatch.setattr("lerobot.jobs.annotate.HfApi", lambda token=None: MagicMock())
|
|
monkeypatch.setattr("lerobot.jobs.annotate.ensure_dataset_available", lambda *a, **kw: None)
|
|
|
|
run_job_calls = []
|
|
|
|
def fake_run_job(**kwargs):
|
|
run_job_calls.append(kwargs)
|
|
return MagicMock(id="job-123")
|
|
|
|
monkeypatch.setattr("lerobot.jobs.annotate.run_job", fake_run_job)
|
|
_set_argv(monkeypatch, "--repo_id=u/d", "--push_to_hub=true", "--job.target=h200", "--job.detach=true")
|
|
|
|
cfg = _parse("--repo_id", "u/d", "--push_to_hub", "true", "--job.target", "h200", "--job.detach", "true")
|
|
submit_annotate_to_hf(cfg)
|
|
|
|
assert len(run_job_calls) == 1
|
|
call = run_job_calls[0]
|
|
assert call["flavor"] == "h200"
|
|
assert call["image"] == DEFAULT_ANNOTATE_JOB_IMAGE
|
|
assert call["timeout"] == "2h"
|
|
# The Hub token is forwarded so the pod can pull a private dataset and push the result.
|
|
assert call["secrets"]["HF_TOKEN"] == "tok"
|
|
assert call["labels"].get("lerobot") == "true"
|
|
argv = _annotate_argv(call["command"])
|
|
assert argv[0] == "lerobot-annotate"
|
|
assert "--push_to_hub=true" in argv
|
|
|
|
|
|
@pytest.mark.timeout(15)
|
|
def test_submit_follows_job_to_completion(monkeypatch, capsys):
|
|
"""Non-detach path must stream logs and RETURN (not hang) once the job is terminal.
|
|
|
|
Exercises the `follow_job` helper shared with the training submitter from the
|
|
annotation side, which is why the job-state patches target `lerobot.jobs.hf`.
|
|
Asserting on the completion message and not merely on "didn't hang" is what makes
|
|
this fail if `follow_job` ever reports detached-without-a-verdict instead.
|
|
"""
|
|
monkeypatch.setattr("lerobot.jobs.annotate.get_token", lambda: "tok")
|
|
monkeypatch.setattr("lerobot.jobs.annotate.HfApi", lambda token=None: MagicMock())
|
|
monkeypatch.setattr("lerobot.jobs.annotate.ensure_dataset_available", lambda *a, **kw: None)
|
|
monkeypatch.setattr("lerobot.jobs.annotate.run_job", lambda **kw: MagicMock(id="job-1", url="http://x"))
|
|
monkeypatch.setattr(
|
|
"lerobot.jobs.hf.inspect_job",
|
|
lambda job_id: MagicMock(status=MagicMock(stage=MagicMock(value="COMPLETED"), message=None)),
|
|
)
|
|
monkeypatch.setattr("lerobot.jobs.hf.fetch_job_logs", lambda job_id, follow=True: iter(()))
|
|
_set_argv(monkeypatch, "--repo_id=u/d", "--job.target=h200")
|
|
|
|
submit_annotate_to_hf(_parse("--repo_id", "u/d", "--push_to_hub", "true", "--job.target", "h200"))
|
|
assert "Annotation complete" in capsys.readouterr().out
|
|
|
|
|
|
@pytest.mark.timeout(15)
|
|
def test_submit_raises_when_job_fails(monkeypatch):
|
|
"""A job that ends in a non-COMPLETED stage must surface as an error, not a silent return."""
|
|
monkeypatch.setattr("lerobot.jobs.annotate.get_token", lambda: "tok")
|
|
monkeypatch.setattr("lerobot.jobs.annotate.HfApi", lambda token=None: MagicMock())
|
|
monkeypatch.setattr("lerobot.jobs.annotate.ensure_dataset_available", lambda *a, **kw: None)
|
|
monkeypatch.setattr("lerobot.jobs.annotate.run_job", lambda **kw: MagicMock(id="job-1", url=None))
|
|
monkeypatch.setattr(
|
|
"lerobot.jobs.hf.inspect_job",
|
|
lambda job_id: MagicMock(status=MagicMock(stage=MagicMock(value="ERROR"), message="Job timeout")),
|
|
)
|
|
monkeypatch.setattr("lerobot.jobs.hf.fetch_job_logs", lambda job_id, follow=True: iter(()))
|
|
_set_argv(monkeypatch, "--repo_id=u/d", "--job.target=h200")
|
|
|
|
with pytest.raises(RuntimeError, match="stage=ERROR .Job timeout."):
|
|
submit_annotate_to_hf(_parse("--repo_id", "u/d", "--job.target", "h200"))
|
|
|
|
|
|
def test_submit_ensures_dataset_is_on_the_hub(monkeypatch):
|
|
"""A local-only dataset is pushed (privately) before the job can reach it by repo_id."""
|
|
monkeypatch.setattr("lerobot.jobs.annotate.get_token", lambda: "tok")
|
|
monkeypatch.setattr("lerobot.jobs.annotate.HfApi", lambda token=None: MagicMock())
|
|
monkeypatch.setattr("lerobot.jobs.annotate.run_job", lambda **kw: MagicMock(id="job-1"))
|
|
|
|
seen = []
|
|
monkeypatch.setattr(
|
|
"lerobot.jobs.annotate.ensure_dataset_available",
|
|
lambda repo_id, *, api, tags=None: seen.append((repo_id, tags)),
|
|
)
|
|
_set_argv(monkeypatch, "--repo_id=u/d", "--job.target=h200", "--job.detach=true")
|
|
|
|
submit_annotate_to_hf(
|
|
_parse("--repo_id", "u/d", "--job.target", "h200", "--job.detach", "true", "--job.tags", '["lelab"]')
|
|
)
|
|
assert seen == [("u/d", ["lerobot", "lelab"])]
|