feat(annotate): run lerobot-annotate on HF Jobs via --job.target (#4095)

* 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>

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
 

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=<flavor>``); the pod
+inside the job invokes ``lerobot-annotate`` which uses this in-process executor.
 Episode-level concurrency is controlled by
 ``ExecutorConfig.episode_parallelism``.
 """

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=<flavor>``): 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=<flavor>`, 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}")

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"]

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 ``--<field>``
+    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.")

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}"
-        )

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=<flavor>`` 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)