From 9c82c39c7b541e9c5bd8340abb7c9d8803c98744 Mon Sep 17 00:00:00 2001 From: Pepijn <138571049+pkooij@users.noreply.github.com> Date: Thu, 23 Jul 2026 10:30:33 +0200 Subject: [PATCH] feat(annotate): run lerobot-annotate on HF Jobs via --job.target (#4095) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * 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 * fix(annotate): reject client-side config files on remote runs draccus exposes `--config_path` plus a `--` 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 * 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 --------- Co-authored-by: Claude Opus 4.8 --- docs/source/annotation_pipeline.mdx | 71 +++-- examples/annotations/run_hf_job.py | 80 ------ .../annotations/steerable_pipeline/config.py | 28 ++ .../steerable_pipeline/executor.py | 4 +- .../steerable_pipeline/vlm_client.py | 17 +- src/lerobot/jobs/__init__.py | 3 +- src/lerobot/jobs/annotate.py | 176 +++++++++++++ src/lerobot/jobs/hf.py | 123 +++++---- src/lerobot/scripts/lerobot_annotate.py | 17 +- tests/jobs/test_annotate.py | 245 ++++++++++++++++++ tests/jobs/test_hf.py | 14 + 11 files changed, 616 insertions(+), 162 deletions(-) delete mode 100644 examples/annotations/run_hf_job.py create mode 100644 src/lerobot/jobs/annotate.py create mode 100644 tests/jobs/test_annotate.py diff --git a/docs/source/annotation_pipeline.mdx b/docs/source/annotation_pipeline.mdx index 3c389663a..45396ff67 100644 --- a/docs/source/annotation_pipeline.mdx +++ b/docs/source/annotation_pipeline.mdx @@ -89,8 +89,8 @@ subtask. The resulting spans are then stitched into a gap-free, full-episode cover, so **every frame has exactly one active subtask**. See -[`run_hf_job.py`](https://github.com/huggingface/lerobot/blob/main/examples/annotations/run_hf_job.py) -for the production settings (single camera, timestamped contact sheets, +[Running on Hugging Face Jobs](#running-on-hugging-face-jobs) for the +production settings (single camera, timestamped contact sheets, auto-windowed subtask generation). ### Tools @@ -110,28 +110,67 @@ not-yet-implemented. ## Running on Hugging Face Jobs -Annotation runs on [Hugging Face Jobs](https://huggingface.co/docs/hub/en/jobs). -The repo ships a launcher script you copy and tweak for your dataset: +Annotating a real dataset needs a GPU big enough to serve the VLM, so +`lerobot-annotate` can dispatch itself to +[Hugging Face Jobs](https://huggingface.co/docs/hub/en/jobs) — same as +`lerobot-train`. Add `--job.target=` to the exact command you'd +run locally and it runs on that hardware instead: ```bash -HF_TOKEN=hf_... uv run python examples/annotations/run_hf_job.py +hf auth login # once + +uv run lerobot-annotate \ + --repo_id=user/my_dataset \ + --new_repo_id=user/my_dataset_annotated \ + --push_to_hub=true \ + --vlm.model_id=Qwen/Qwen3.6-27B \ + --vlm.num_gpus=1 \ + --vlm.serve_command="vllm serve Qwen/Qwen3.6-27B --tensor-parallel-size 1 \ + --max-model-len 32768 --gpu-memory-utilization 0.8 \ + --uvicorn-log-level warning --port {port}" \ + --vlm.serve_ready_timeout_s=1800 \ + --vlm.chat_template_kwargs='{"enable_thinking": false}' \ + --job.target=h200 ``` -[`run_hf_job.py`](https://github.com/huggingface/lerobot/blob/main/examples/annotations/run_hf_job.py) -starts a single-GPU `h200` job (bump it to `h200x4` for big datasets) -that: +That submits a single-GPU `h200` job that: -1. installs `lerobot` (from `main`) plus the annotation extras, -2. boots one vLLM server per GPU (using the `vllm/vllm-openai` image) and - drives it over the OpenAI-compatible API, -3. runs the `plan` / `interjections` / `vqa` modules across the dataset - with `lerobot-annotate`, +1. starts from the `vllm/vllm-openai` image and installs `lerobot` on top, +2. boots one vLLM server per GPU and drives it over the OpenAI-compatible API, +3. runs the `plan` / `interjections` / `vqa` modules across the dataset, 4. with `--push_to_hub=true`, uploads the result to `--new_repo_id` (or back to `--repo_id` in place if you leave that unset). -To use a different dataset, model, or hub repo, edit the `CMD` block in -the script. Every flag there maps directly to a `lerobot-annotate` flag -(run `lerobot-annotate --help` for the full list). +The command streams the job's logs; `Ctrl-C` detaches without cancelling +it. List the available flavors and their pricing with `hf jobs hardware`. + + + +Qwen3.6 ships with thinking enabled, which eats the token budget the +annotator needs for its JSON answer — `--vlm.chat_template_kwargs='{"enable_thinking": false}'` +turns it off. Without `--push_to_hub=true` the annotated dataset is +discarded when the pod exits. + + + +### Job options + +| Flag | Default | What it does | +| ------------------- | ------------------------- | ------------------------------------------------------------------------------- | +| `--job.target` | `local` | HF Jobs flavor to run on (e.g. `h200`, `h200x4`). Omitted/`local` runs here. | +| `--job.image` | `vllm/vllm-openai:latest` | Runtime image for the pod. | +| `--job.timeout` | `2h` | Wall-clock cap. Raise it for large datasets. | +| `--job.detach` | `false` | Submit and exit instead of streaming logs. | +| `--job.lerobot_ref` | `main` | Git ref of lerobot installed on the pod — point it at a branch to test changes. | +| `--job.tags` | `[]` | Extra tags on the job and on any dataset it pushes (`lerobot` is always added). | + +For a bigger dataset, scale to `h200x4` and raise +`--vlm.parallel_servers` / `--vlm.num_gpus` to match, and give the job +more headroom with e.g. `--job.timeout=8h`. + +Remote runs need `--repo_id` (the pod pulls the dataset from the Hub; +`--root` names a directory only your machine has). A dataset that exists +only in your local cache is pushed to a **private** repo first. ## Key options diff --git a/examples/annotations/run_hf_job.py b/examples/annotations/run_hf_job.py deleted file mode 100644 index 506a45a12..000000000 --- a/examples/annotations/run_hf_job.py +++ /dev/null @@ -1,80 +0,0 @@ -#!/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. -"""Launch ``lerobot-annotate`` on a Hugging Face job (vllm + Qwen3.6-27B VLM). - -Spawns one single-GPU ``h200`` job that: - - 1. installs ``lerobot`` from ``main`` plus the annotation extras, - 2. boots one vllm server with Qwen3.6-27B (dense VLM), - 3. runs the plan / interjections / vqa modules across the dataset - in free-form mode (each episode generates its own subtasks + - memory), - 4. uploads the annotated dataset to ``--new_repo_id`` (when set) - or back to ``--repo_id``. - -Usage: - - HF_TOKEN=hf_... uv run python examples/annotations/run_hf_job.py - -Adjust ``CMD`` (dataset, model, hub repo) and ``flavor`` below for your -run. For larger datasets, scale to ``h200x4`` and raise -``--vlm.parallel_servers`` / ``--vlm.num_gpus`` to match. -""" - -import os - -from huggingface_hub import get_token, run_job - -token = os.environ.get("HF_TOKEN") or get_token() -if not token: - raise RuntimeError("No HF token. Run `huggingface-cli login` or `export HF_TOKEN=hf_...`") - -CMD = ( - "apt-get update -qq && apt-get install -y -qq git ffmpeg && " - "pip install --no-deps " - "'lerobot @ git+https://github.com/huggingface/lerobot.git@main' && " - # Pins mirror pyproject.toml — unpinned installs pull av 18 / datasets 5 / - # draccus 0.11, which break lerobot at import time. - "pip install --upgrade-strategy only-if-needed " - "'datasets>=4.7.0,<5.0.0' 'pyarrow>=21.0.0,<30.0.0' 'av>=15.0.0,<16.0.0' 'draccus==0.10.0' " - "'pandas>=2.0.0,<3.0.0' jsonlines gymnasium torchcodec mergedeep pyyaml-include toml typing-inspect " - "openai && " - "export VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS=0 && " - "export VLLM_VIDEO_BACKEND=pyav && " - "lerobot-annotate " - "--repo_id=pepijn223/robocasa_pretrain_human300_v4 " - "--new_repo_id=pepijn223/robocasa_pretrain_human300_v4_annotated " - "--push_to_hub=true " - "--vlm.backend=openai " - "--vlm.model_id=Qwen/Qwen3.6-27B " - "--vlm.num_gpus=1 " - '--vlm.serve_command="vllm serve Qwen/Qwen3.6-27B ' - "--tensor-parallel-size 1 --max-model-len 32768 " - '--gpu-memory-utilization 0.8 --uvicorn-log-level warning --port {port}" ' - "--vlm.serve_ready_timeout_s=1800 " - # Qwen3.6 ships with thinking on; annotation wants plain JSON answers. - "--vlm.chat_template_kwargs='{\"enable_thinking\": false}'" -) - -job = run_job( - image="vllm/vllm-openai:latest", - command=["bash", "-c", CMD], - flavor="h200", - secrets={"HF_TOKEN": token}, - timeout="2h", -) -print(f"Job URL: {job.url}") -print(f"Job ID: {job.id}") diff --git a/src/lerobot/annotations/steerable_pipeline/config.py b/src/lerobot/annotations/steerable_pipeline/config.py index 73f41f4f7..eebd2cdeb 100644 --- a/src/lerobot/annotations/steerable_pipeline/config.py +++ b/src/lerobot/annotations/steerable_pipeline/config.py @@ -20,6 +20,29 @@ from dataclasses import dataclass, field from pathlib import Path from typing import Any +from lerobot.configs.default import JobConfig + +# The annotation pipeline boots its own vLLM server, so the pod starts from the +# official vLLM runtime rather than the prebuilt `lerobot-gpu` training image; +# `lerobot` is pip-installed on top (see `lerobot.jobs.annotate`). +DEFAULT_ANNOTATE_JOB_IMAGE = "vllm/vllm-openai:latest" + + +@dataclass +class AnnotationJobConfig(JobConfig): + """`JobConfig` with the annotation runtime's defaults. + + Adds `lerobot_ref` because the vLLM image ships no lerobot: the pod installs + it from git, and the ref decides which code actually annotates. Point it at a + branch/tag/SHA to try unmerged changes remotely. + """ + + image: str = DEFAULT_ANNOTATE_JOB_IMAGE + # Annotation is a bounded pass over a dataset; a tighter cap than training's + # "2d" keeps a wedged vLLM server from burning a day of GPU time. + timeout: str | None = "2h" + lerobot_ref: str = "main" + @dataclass class PlanConfig: @@ -207,6 +230,11 @@ class AnnotationPipelineConfig: vlm: VlmConfig = field(default_factory=VlmConfig) executor: ExecutorConfig = field(default_factory=ExecutorConfig) + # Where the annotation runs: omitted / "local" annotates on this machine, any + # other value is an HF Jobs flavor (e.g. "h200") and submits the run there. + # List flavors + pricing with `hf jobs hardware`. + job: AnnotationJobConfig = field(default_factory=AnnotationJobConfig) + skip_validation: bool = False only_episodes: tuple[int, ...] | None = None diff --git a/src/lerobot/annotations/steerable_pipeline/executor.py b/src/lerobot/annotations/steerable_pipeline/executor.py index 69d10bc89..0be496f07 100644 --- a/src/lerobot/annotations/steerable_pipeline/executor.py +++ b/src/lerobot/annotations/steerable_pipeline/executor.py @@ -30,8 +30,8 @@ Phase 3 is why the ``plan`` module must be re-entered after the timestamps. Distributed execution is provided by Hugging Face Jobs (see -``examples/annotations/run_hf_job.py``); the runner inside the job -invokes ``lerobot-annotate`` which uses this in-process executor. +``lerobot.jobs.annotate``, reached via ``--job.target=``); the pod +inside the job invokes ``lerobot-annotate`` which uses this in-process executor. Episode-level concurrency is controlled by ``ExecutorConfig.episode_parallelism``. """ diff --git a/src/lerobot/annotations/steerable_pipeline/vlm_client.py b/src/lerobot/annotations/steerable_pipeline/vlm_client.py index 4d77786b0..7b0c8609c 100644 --- a/src/lerobot/annotations/steerable_pipeline/vlm_client.py +++ b/src/lerobot/annotations/steerable_pipeline/vlm_client.py @@ -194,12 +194,13 @@ def make_vlm_client(config: VlmConfig) -> VlmClient: """Build the shared VLM client. Only the ``openai`` backend is supported for now. The shipped workflow - is Hugging Face Jobs (``examples/annotations/run_hf_job.py``): it boots - a vLLM server inside the ``vllm/vllm-openai`` image and the pipeline - talks to it over the OpenAI-compatible API (``--vlm.backend=openai``, - optionally auto-spawning the server via ``auto_serve`` / - ``serve_command``). The former in-process ``vllm`` / ``transformers`` - backends were removed to keep the support surface to the HF Jobs path. + is Hugging Face Jobs (``lerobot-annotate --job.target=``): it + boots a vLLM server inside the ``vllm/vllm-openai`` image and the + pipeline talks to it over the OpenAI-compatible API + (``--vlm.backend=openai``, optionally auto-spawning the server via + ``auto_serve`` / ``serve_command``). The former in-process ``vllm`` / + ``transformers`` backends were removed to keep the support surface to + the HF Jobs path. For ``stub``, construct :class:`StubVlmClient` directly with a responder callable; it is rejected here to make accidental misuse obvious. @@ -213,8 +214,8 @@ def make_vlm_client(config: VlmConfig) -> VlmClient: if config.backend in {"vllm", "transformers"}: raise ValueError( f"backend={config.backend!r} (in-process local model) is not supported for now — " - "only backend='openai' (the Hugging Face Jobs flow) is. Run the pipeline via " - "examples/annotations/run_hf_job.py, which serves the model with vLLM in the " + "only backend='openai' (the Hugging Face Jobs flow) is. Run the pipeline with " + "`lerobot-annotate --job.target=`, which serves the model with vLLM in the " "vllm/vllm-openai image and talks to it over the OpenAI-compatible API." ) raise ValueError(f"Unknown VLM backend: {config.backend!r}") diff --git a/src/lerobot/jobs/__init__.py b/src/lerobot/jobs/__init__.py index 674b98b85..4f9bd9174 100644 --- a/src/lerobot/jobs/__init__.py +++ b/src/lerobot/jobs/__init__.py @@ -18,6 +18,7 @@ from lerobot.utils.import_utils import require_package # guard the optional dependency here so importing this package fails loudly if it's missing. require_package("datasets", extra="dataset") +from .annotate import submit_annotate_to_hf from .hf import submit_to_hf -__all__ = ["submit_to_hf"] +__all__ = ["submit_annotate_to_hf", "submit_to_hf"] diff --git a/src/lerobot/jobs/annotate.py b/src/lerobot/jobs/annotate.py new file mode 100644 index 000000000..d7baad024 --- /dev/null +++ b/src/lerobot/jobs/annotate.py @@ -0,0 +1,176 @@ +# 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. +"""Run ``lerobot-annotate`` on HF Jobs (HuggingFace GPUs). + +Same shape as the training submitter in ``hf.py``, with one difference: the +annotation pipeline serves its own VLM, so the pod starts from the official +``vllm/vllm-openai`` image (which has no lerobot) instead of the prebuilt +``lerobot-gpu`` image, and installs lerobot on top before running. + +Because there is no config repo to stage, the pod replays the user's own CLI +flags — everything except the client-only ``--job.*`` and the host-local +``--root``, which is replaced by ``--repo_id`` so the pod pulls the dataset +from the Hub. +""" + +from __future__ import annotations + +import shlex +import sys +from dataclasses import is_dataclass +from typing import TYPE_CHECKING + +from huggingface_hub import HfApi, get_token, run_job + +from .dataset import ensure_dataset_available + +# Package-internal reuse of the training submitter's job plumbing: following a +# submitted job and forwarding argv are identical for annotation runs. +from .hf import _pod_forwarded_args, follow_job, resolve_job_tags + +if TYPE_CHECKING: + from lerobot.annotations.steerable_pipeline.config import AnnotationPipelineConfig + +LEROBOT_GIT_URL = "https://github.com/huggingface/lerobot.git" + +# Mirrors the pins in pyproject.toml. The vLLM image resolves dependencies on its +# own otherwise, and pulls av 18 / datasets 5 / draccus 0.11 — each of which breaks +# lerobot at import time. `--upgrade-strategy only-if-needed` keeps vLLM's own +# (torch, transformers, ...) pins intact. +_RUNTIME_REQUIREMENTS = ( + "'datasets>=4.7.0,<5.0.0' 'pyarrow>=21.0.0,<30.0.0' 'av>=15.0.0,<16.0.0' 'draccus==0.10.0' " + "'pandas>=2.0.0,<3.0.0' jsonlines gymnasium torchcodec mergedeep pyyaml-include toml typing-inspect " + "openai" +) + +# Flags the submitter resolves itself instead of forwarding verbatim: `--root` +# names a directory only this machine has, `--repo_id` is re-emitted from the +# config, and the config-file args name local files (rejected up front by +# `submit_annotate_to_hf`). `--job.*` is dropped separately, by prefix; bare +# `--job` is not, hence its entry here — it is the one arg that could smuggle a +# remote `target` onto the pod and have the job recursively submit itself. +_SUBMITTER_OWNED_ARGS = ("--root", "--repo_id", "--config_path", "--job") + + +def _local_config_file_args(cfg: AnnotationPipelineConfig) -> list[str]: + """The CLI args that name a config file on the client's disk. + + draccus exposes ``--config_path`` for the whole config plus a ``--`` + for every nested dataclass (``--vlm``, ``--plan``, ``--job``, ...). The pod has + none of those files, so a remote run has to reject them rather than silently + drop the settings they carry. + """ + return ["--config_path", *(f"--{name}" for name in vars(cfg) if is_dataclass(getattr(cfg, name)))] + + +def build_pod_setup(lerobot_ref: str) -> str: + """Shell prelude that turns the vLLM image into a ``lerobot-annotate`` runtime.""" + spec = f"lerobot @ git+{LEROBOT_GIT_URL}@{lerobot_ref}" + return ( + # git to install from the repo, ffmpeg to decode the dataset's videos. + "apt-get update -qq && apt-get install -y -qq git ffmpeg && " + f"pip install --no-deps {shlex.quote(spec)} && " + f"pip install --upgrade-strategy only-if-needed {_RUNTIME_REQUIREMENTS} && " + # vLLM's cudagraph memory estimate over-reserves and starves the KV cache; + # PyAV is the video backend the server can decode our frames with. + "export VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS=0 && " + "export VLLM_VIDEO_BACKEND=pyav" + ) + + +def build_pod_command(repo_id: str, lerobot_ref: str, argv: list[str]) -> list[str]: + """Build the ``bash -c`` command the pod runs: setup prelude, then annotation. + + ``argv`` is the user's CLI (``sys.argv[1:]``) minus the flags in + ``_SUBMITTER_OWNED_ARGS``; ``--repo_id`` is re-added from the config so the pod + always annotates the dataset we just made sure is reachable on the Hub. + ``--job.target=local`` stops the pod from re-dispatching to itself. + """ + forwarded = _pod_forwarded_args(argv, drop_names=_SUBMITTER_OWNED_ARGS, drop_prefixes=("--job.",)) + annotate = shlex.join(["lerobot-annotate", f"--repo_id={repo_id}", *forwarded, "--job.target=local"]) + return ["bash", "-c", f"{build_pod_setup(lerobot_ref)} && {annotate}"] + + +def submit_annotate_to_hf(cfg: AnnotationPipelineConfig) -> None: + """Submit an annotation run to HF Jobs infrastructure. + + Resolves credentials, makes sure the source dataset is reachable from the pod, + submits the job, then tails its logs until the job reaches a terminal stage — + or returns immediately with ``--job.detach``. Ctrl-C detaches without + cancelling the remote job. + """ + token = get_token() + if not token: + raise RuntimeError("Not logged in to Hugging Face. Run `hf auth login` first.") + + if cfg.repo_id is None: + raise ValueError( + "Remote annotation requires --repo_id: the pod downloads the dataset from the Hub, " + "and --root only names a directory on this machine." + ) + + argv = sys.argv[1:] + passed = {tok.split("=", 1)[0] for tok in argv} + used_config_files = sorted(passed.intersection(_local_config_file_args(cfg))) + if used_config_files: + raise ValueError( + f"{', '.join(used_config_files)} cannot be used with a remote --job.target: the pod " + "cannot read config files from this machine. Pass the settings as CLI flags instead." + ) + + if not cfg.push_to_hub: + # The pod's filesystem is discarded when the job ends, so without a push the + # run produces nothing. Warn rather than fail: a smoke test over + # --only_episodes that only inspects the logs is a legitimate use. + print( + "WARNING: --push_to_hub is off. The annotated dataset lives only on the pod and is " + "discarded when the job ends. Pass --push_to_hub=true to keep the result." + ) + + api = HfApi(token=token) + tags = resolve_job_tags(cfg.job.tags) + ensure_dataset_available(cfg.repo_id, api=api, tags=tags) + + command = build_pod_command(cfg.repo_id, cfg.job.lerobot_ref, argv) + + print(f"Submitting job to HF Jobs (flavor={cfg.job.target}, image={cfg.job.image}) ...") + job_info = run_job( + image=cfg.job.image, + command=command, + flavor=cfg.job.target, + secrets={"HF_TOKEN": token}, + timeout=cfg.job.timeout, + # HF Jobs labels are key/value; expose each tag as a queryable label. + labels=dict.fromkeys(tags, "true"), + ) + job_id = job_info.id + job_url = getattr(job_info, "url", None) + print(f"Job submitted: {job_id}") + if job_url: + print(f" Job page: {job_url}") + target_repo_id = cfg.new_repo_id or cfg.repo_id + if cfg.push_to_hub: + print(f" Dataset repo: https://huggingface.co/datasets/{target_repo_id}") + print(f" Monitor: hf jobs logs {job_id}") + print(f" Cancel: hf jobs cancel {job_id}") + + # No success marker: `lerobot-annotate` keeps working after the upload log line + # (dataset card, version tag), so completion has to be stage-based. + if not follow_job(job_id, detach=cfg.job.detach): + return + + if cfg.push_to_hub: + print(f"\nAnnotation complete — dataset pushed to https://huggingface.co/datasets/{target_repo_id}") + else: + print("\nAnnotation complete. Note: --push_to_hub was off, so the result stayed on the pod.") diff --git a/src/lerobot/jobs/hf.py b/src/lerobot/jobs/hf.py index 645666412..80d40ad62 100644 --- a/src/lerobot/jobs/hf.py +++ b/src/lerobot/jobs/hf.py @@ -223,6 +223,74 @@ def _poll_until_done( return None +def follow_job(job_id: str, *, detach: bool = False, success_marker: str | None = None) -> bool: + """Watch a submitted job to the end, streaming its logs to stdout. + + Returns True when the job finished successfully and False when we stopped watching + without a verdict — `detach`, or the user pressing Ctrl-C, which detaches rather than + cancelling the remote job. Raises RuntimeError when the job reaches a terminal stage + other than COMPLETED. + + `success_marker` finishes as soon as that string appears in the logs instead of waiting + out the platform's post-run finalization (~30s). Callers that have a log line meaning + "the artifact is on the Hub" should pass it; without one, completion is stage-based. + """ + if detach: + return False + + done = threading.Event() + detached = threading.Event() + marker_seen = threading.Event() + stage_holder: dict[str, str | None] = {} + + def _poll() -> None: + stage_holder["stage"] = _poll_until_done(job_id, done, status_holder=stage_holder) + + poll_thread = threading.Thread(target=_poll, daemon=True) + poll_thread.start() + log_thread = threading.Thread( + target=_tail_logs, args=(job_id, done, success_marker, marker_seen), daemon=True + ) + log_thread.start() + + def _detach(sig, frame): + detached.set() + done.set() + print("\nDetached. Job is still running.") + print(f" Monitor: hf jobs logs {job_id}") + print(f" Cancel: hf jobs cancel {job_id}") + + # signal.signal only works on the main thread; when called from a worker thread + # (e.g. an orchestration framework) skip the Ctrl-C-detaches-instead-of-cancels + # handler rather than crashing with ValueError. + install_sigint = threading.current_thread() is threading.main_thread() + original_sigint = signal.getsignal(signal.SIGINT) if install_sigint else None + if install_sigint: + signal.signal(signal.SIGINT, _detach) + try: + # Timeout-based join so SIGINT is delivered to the main thread promptly. + while poll_thread.is_alive(): + poll_thread.join(timeout=0.5) + log_thread.join(timeout=5) + finally: + if install_sigint: + signal.signal(signal.SIGINT, original_sigint) + + if detached.is_set(): + return False + if marker_seen.is_set(): + return True + + stage = stage_holder.get("stage") + if stage != "COMPLETED": + message = stage_holder.get("message") + detail = f" ({message})" if message else "" + raise RuntimeError( + f"Job {job_id} ended with stage={stage}{detail}. Check logs: hf jobs logs {job_id}" + ) + return True + + def _pod_forwarded_args( argv: list[str], drop_names: tuple[str, ...] = (), drop_prefixes: tuple[str, ...] = () ) -> list[str]: @@ -362,64 +430,11 @@ def submit_to_hf(cfg: TrainPipelineConfig) -> None: print(f" Monitor: hf jobs logs {job_id}") print(f" Cancel: hf jobs cancel {job_id}") - if cfg.job.detach: - return - - done = threading.Event() - detached = threading.Event() - pushed_ok = threading.Event() - stage_holder: dict[str, str | None] = {} - - def _poll() -> None: - stage_holder["stage"] = _poll_until_done(job_id, done, status_holder=stage_holder) - - poll_thread = threading.Thread(target=_poll, daemon=True) - poll_thread.start() # Finish as soon as the model is pushed, rather than waiting out the platform's # post-run finalization before the job stage flips to COMPLETED. This matches the # exact log line emitted by PreTrainedPolicy.push_model_to_hub — the two must stay # in sync. If it ever stops matching we just fall back to stage-based completion # (~30s slower), so the contract is an optimization, not a correctness requirement. success_marker = f"Model pushed to https://huggingface.co/{repo_id}" - log_thread = threading.Thread( - target=_tail_logs, args=(job_id, done, success_marker, pushed_ok), daemon=True - ) - log_thread.start() - - def _detach(sig, frame): - detached.set() - done.set() - print("\nDetached. Job is still running.") - print(f" Monitor: hf jobs logs {job_id}") - print(f" Cancel: hf jobs cancel {job_id}") - - # signal.signal only works on the main thread; when called from a worker thread - # (e.g. an orchestration framework) skip the Ctrl-C-detaches-instead-of-cancels - # handler rather than crashing with ValueError. - install_sigint = threading.current_thread() is threading.main_thread() - original_sigint = signal.getsignal(signal.SIGINT) if install_sigint else None - if install_sigint: - signal.signal(signal.SIGINT, _detach) - try: - # Timeout-based join so SIGINT is delivered to the main thread promptly. - while poll_thread.is_alive(): - poll_thread.join(timeout=0.5) - log_thread.join(timeout=5) - finally: - if install_sigint: - signal.signal(signal.SIGINT, original_sigint) - - if detached.is_set(): - return - - if pushed_ok.is_set(): + if follow_job(job_id, detach=cfg.job.detach, success_marker=success_marker): print(f"\nTraining complete — model pushed to https://huggingface.co/{repo_id}") - return - - stage = stage_holder.get("stage") - if stage != "COMPLETED": - message = stage_holder.get("message") - detail = f" ({message})" if message else "" - raise RuntimeError( - f"Job {job_id} ended with stage={stage}{detail}. Check logs: hf jobs logs {job_id}" - ) diff --git a/src/lerobot/scripts/lerobot_annotate.py b/src/lerobot/scripts/lerobot_annotate.py index 41fede908..5594542d7 100644 --- a/src/lerobot/scripts/lerobot_annotate.py +++ b/src/lerobot/scripts/lerobot_annotate.py @@ -24,7 +24,14 @@ Example: --root=/path/to/dataset \\ --vlm.model_id=Qwen/Qwen2.5-VL-7B-Instruct -For distributed runs, see ``examples/annotations/run_hf_job.py``. +Pass ``--job.target=`` to run the same command on a Hugging Face +Jobs GPU instead of this machine (see ``lerobot.jobs.annotate``): + + uv run lerobot-annotate \\ + --repo_id=user/dataset \\ + --new_repo_id=user/dataset_annotated \\ + --push_to_hub=true \\ + --job.target=h200 """ import logging @@ -69,6 +76,14 @@ def _resolve_root(cfg: AnnotationPipelineConfig) -> Path: def annotate(cfg: AnnotationPipelineConfig) -> None: """Run the steerable annotation pipeline against a dataset.""" logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") + + if cfg.job.is_remote: + # Imported lazily: the submitter pulls in LeRobotDataset (the `dataset` + # extra), which a local annotation run over --root doesn't need. + from lerobot.jobs.annotate import submit_annotate_to_hf + + return submit_annotate_to_hf(cfg) + root = _resolve_root(cfg) logger.info("annotate: root=%s", root) diff --git a/tests/jobs/test_annotate.py b/tests/jobs/test_annotate.py new file mode 100644 index 000000000..8c3db9292 --- /dev/null +++ b/tests/jobs/test_annotate.py @@ -0,0 +1,245 @@ +# 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"])] diff --git a/tests/jobs/test_hf.py b/tests/jobs/test_hf.py index 3b275cb95..4804a5929 100644 --- a/tests/jobs/test_hf.py +++ b/tests/jobs/test_hf.py @@ -29,12 +29,26 @@ from lerobot.jobs.hf import ( _poll_until_done, build_remote_config_file, build_repo_id, + follow_job, resolve_job_tags, resolve_wandb_api_key, submit_to_hf, ) +def test_follow_job_detach_returns_without_watching(monkeypatch): + """`detach` must short-circuit before any polling or log streaming starts.""" + + def _boom(*a, **kw): + raise AssertionError("detach must not touch the job") + + monkeypatch.setattr("lerobot.jobs.hf.inspect_job", _boom) + monkeypatch.setattr("lerobot.jobs.hf.fetch_job_logs", _boom) + # False = "stopped watching without a verdict", so callers stay quiet rather than + # claiming success for a job that is still running. + assert follow_job("job-1", detach=True) is False + + def test_resolve_job_tags_always_includes_lerobot_and_dedups(): assert resolve_job_tags(None) == ["lerobot"] assert resolve_job_tags([]) == ["lerobot"]