mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-24 10:16:09 +00:00
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>
This commit is contained in:
@@ -89,8 +89,8 @@ subtask.
|
|||||||
|
|
||||||
The resulting spans are then stitched into a gap-free, full-episode
|
The resulting spans are then stitched into a gap-free, full-episode
|
||||||
cover, so **every frame has exactly one active subtask**. See
|
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)
|
[Running on Hugging Face Jobs](#running-on-hugging-face-jobs) for the
|
||||||
for the production settings (single camera, timestamped contact sheets,
|
production settings (single camera, timestamped contact sheets,
|
||||||
auto-windowed subtask generation).
|
auto-windowed subtask generation).
|
||||||
|
|
||||||
### Tools
|
### Tools
|
||||||
@@ -110,28 +110,67 @@ not-yet-implemented.
|
|||||||
|
|
||||||
## Running on Hugging Face Jobs
|
## Running on Hugging Face Jobs
|
||||||
|
|
||||||
Annotation runs on [Hugging Face Jobs](https://huggingface.co/docs/hub/en/jobs).
|
Annotating a real dataset needs a GPU big enough to serve the VLM, so
|
||||||
The repo ships a launcher script you copy and tweak for your dataset:
|
`lerobot-annotate` can dispatch itself to
|
||||||
|
[Hugging Face Jobs](https://huggingface.co/docs/hub/en/jobs) — same as
|
||||||
|
`lerobot-train`. Add `--job.target=<flavor>` to the exact command you'd
|
||||||
|
run locally and it runs on that hardware instead:
|
||||||
|
|
||||||
```bash
|
```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)
|
That submits a single-GPU `h200` job that:
|
||||||
starts a single-GPU `h200` job (bump it to `h200x4` for big datasets)
|
|
||||||
that:
|
|
||||||
|
|
||||||
1. installs `lerobot` (from `main`) plus the annotation extras,
|
1. starts from the `vllm/vllm-openai` image and installs `lerobot` on top,
|
||||||
2. boots one vLLM server per GPU (using the `vllm/vllm-openai` image) and
|
2. boots one vLLM server per GPU and drives it over the OpenAI-compatible API,
|
||||||
drives it over the OpenAI-compatible API,
|
3. runs the `plan` / `interjections` / `vqa` modules across the dataset,
|
||||||
3. runs the `plan` / `interjections` / `vqa` modules across the dataset
|
|
||||||
with `lerobot-annotate`,
|
|
||||||
4. with `--push_to_hub=true`, uploads the result to `--new_repo_id` (or
|
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).
|
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 command streams the job's logs; `Ctrl-C` detaches without cancelling
|
||||||
the script. Every flag there maps directly to a `lerobot-annotate` flag
|
it. List the available flavors and their pricing with `hf jobs hardware`.
|
||||||
(run `lerobot-annotate --help` for the full list).
|
|
||||||
|
<Tip warning={true}>
|
||||||
|
|
||||||
|
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.
|
||||||
|
|
||||||
|
</Tip>
|
||||||
|
|
||||||
|
### 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
|
## Key options
|
||||||
|
|
||||||
|
|||||||
@@ -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}")
|
|
||||||
@@ -20,6 +20,29 @@ from dataclasses import dataclass, field
|
|||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from typing import Any
|
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
|
@dataclass
|
||||||
class PlanConfig:
|
class PlanConfig:
|
||||||
@@ -207,6 +230,11 @@ class AnnotationPipelineConfig:
|
|||||||
vlm: VlmConfig = field(default_factory=VlmConfig)
|
vlm: VlmConfig = field(default_factory=VlmConfig)
|
||||||
executor: ExecutorConfig = field(default_factory=ExecutorConfig)
|
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
|
skip_validation: bool = False
|
||||||
only_episodes: tuple[int, ...] | None = None
|
only_episodes: tuple[int, ...] | None = None
|
||||||
|
|
||||||
|
|||||||
@@ -30,8 +30,8 @@ Phase 3 is why the ``plan`` module must be re-entered after the
|
|||||||
timestamps.
|
timestamps.
|
||||||
|
|
||||||
Distributed execution is provided by Hugging Face Jobs (see
|
Distributed execution is provided by Hugging Face Jobs (see
|
||||||
``examples/annotations/run_hf_job.py``); the runner inside the job
|
``lerobot.jobs.annotate``, reached via ``--job.target=<flavor>``); the pod
|
||||||
invokes ``lerobot-annotate`` which uses this in-process executor.
|
inside the job invokes ``lerobot-annotate`` which uses this in-process executor.
|
||||||
Episode-level concurrency is controlled by
|
Episode-level concurrency is controlled by
|
||||||
``ExecutorConfig.episode_parallelism``.
|
``ExecutorConfig.episode_parallelism``.
|
||||||
"""
|
"""
|
||||||
|
|||||||
@@ -194,12 +194,13 @@ def make_vlm_client(config: VlmConfig) -> VlmClient:
|
|||||||
"""Build the shared VLM client.
|
"""Build the shared VLM client.
|
||||||
|
|
||||||
Only the ``openai`` backend is supported for now. The shipped workflow
|
Only the ``openai`` backend is supported for now. The shipped workflow
|
||||||
is Hugging Face Jobs (``examples/annotations/run_hf_job.py``): it boots
|
is Hugging Face Jobs (``lerobot-annotate --job.target=<flavor>``): it
|
||||||
a vLLM server inside the ``vllm/vllm-openai`` image and the pipeline
|
boots a vLLM server inside the ``vllm/vllm-openai`` image and the
|
||||||
talks to it over the OpenAI-compatible API (``--vlm.backend=openai``,
|
pipeline talks to it over the OpenAI-compatible API
|
||||||
optionally auto-spawning the server via ``auto_serve`` /
|
(``--vlm.backend=openai``, optionally auto-spawning the server via
|
||||||
``serve_command``). The former in-process ``vllm`` / ``transformers``
|
``auto_serve`` / ``serve_command``). The former in-process ``vllm`` /
|
||||||
backends were removed to keep the support surface to the HF Jobs path.
|
``transformers`` backends were removed to keep the support surface to
|
||||||
|
the HF Jobs path.
|
||||||
|
|
||||||
For ``stub``, construct :class:`StubVlmClient` directly with a responder
|
For ``stub``, construct :class:`StubVlmClient` directly with a responder
|
||||||
callable; it is rejected here to make accidental misuse obvious.
|
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"}:
|
if config.backend in {"vllm", "transformers"}:
|
||||||
raise ValueError(
|
raise ValueError(
|
||||||
f"backend={config.backend!r} (in-process local model) is not supported for now — "
|
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 "
|
"only backend='openai' (the Hugging Face Jobs flow) is. Run the pipeline with "
|
||||||
"examples/annotations/run_hf_job.py, which serves the model with vLLM in the "
|
"`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."
|
"vllm/vllm-openai image and talks to it over the OpenAI-compatible API."
|
||||||
)
|
)
|
||||||
raise ValueError(f"Unknown VLM backend: {config.backend!r}")
|
raise ValueError(f"Unknown VLM backend: {config.backend!r}")
|
||||||
|
|||||||
@@ -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.
|
# guard the optional dependency here so importing this package fails loudly if it's missing.
|
||||||
require_package("datasets", extra="dataset")
|
require_package("datasets", extra="dataset")
|
||||||
|
|
||||||
|
from .annotate import submit_annotate_to_hf
|
||||||
from .hf import submit_to_hf
|
from .hf import submit_to_hf
|
||||||
|
|
||||||
__all__ = ["submit_to_hf"]
|
__all__ = ["submit_annotate_to_hf", "submit_to_hf"]
|
||||||
|
|||||||
@@ -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.")
|
||||||
+69
-54
@@ -223,6 +223,74 @@ def _poll_until_done(
|
|||||||
return None
|
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(
|
def _pod_forwarded_args(
|
||||||
argv: list[str], drop_names: tuple[str, ...] = (), drop_prefixes: tuple[str, ...] = ()
|
argv: list[str], drop_names: tuple[str, ...] = (), drop_prefixes: tuple[str, ...] = ()
|
||||||
) -> list[str]:
|
) -> list[str]:
|
||||||
@@ -362,64 +430,11 @@ def submit_to_hf(cfg: TrainPipelineConfig) -> None:
|
|||||||
print(f" Monitor: hf jobs logs {job_id}")
|
print(f" Monitor: hf jobs logs {job_id}")
|
||||||
print(f" Cancel: hf jobs cancel {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
|
# 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
|
# 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
|
# 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
|
# 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.
|
# (~30s slower), so the contract is an optimization, not a correctness requirement.
|
||||||
success_marker = f"Model pushed to https://huggingface.co/{repo_id}"
|
success_marker = f"Model pushed to https://huggingface.co/{repo_id}"
|
||||||
log_thread = threading.Thread(
|
if follow_job(job_id, detach=cfg.job.detach, success_marker=success_marker):
|
||||||
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():
|
|
||||||
print(f"\nTraining complete — model pushed to https://huggingface.co/{repo_id}")
|
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}"
|
|
||||||
)
|
|
||||||
|
|||||||
@@ -24,7 +24,14 @@ Example:
|
|||||||
--root=/path/to/dataset \\
|
--root=/path/to/dataset \\
|
||||||
--vlm.model_id=Qwen/Qwen2.5-VL-7B-Instruct
|
--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
|
import logging
|
||||||
@@ -69,6 +76,14 @@ def _resolve_root(cfg: AnnotationPipelineConfig) -> Path:
|
|||||||
def annotate(cfg: AnnotationPipelineConfig) -> None:
|
def annotate(cfg: AnnotationPipelineConfig) -> None:
|
||||||
"""Run the steerable annotation pipeline against a dataset."""
|
"""Run the steerable annotation pipeline against a dataset."""
|
||||||
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
|
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)
|
root = _resolve_root(cfg)
|
||||||
logger.info("annotate: root=%s", root)
|
logger.info("annotate: root=%s", root)
|
||||||
|
|
||||||
|
|||||||
@@ -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"])]
|
||||||
@@ -29,12 +29,26 @@ from lerobot.jobs.hf import (
|
|||||||
_poll_until_done,
|
_poll_until_done,
|
||||||
build_remote_config_file,
|
build_remote_config_file,
|
||||||
build_repo_id,
|
build_repo_id,
|
||||||
|
follow_job,
|
||||||
resolve_job_tags,
|
resolve_job_tags,
|
||||||
resolve_wandb_api_key,
|
resolve_wandb_api_key,
|
||||||
submit_to_hf,
|
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():
|
def test_resolve_job_tags_always_includes_lerobot_and_dedups():
|
||||||
assert resolve_job_tags(None) == ["lerobot"]
|
assert resolve_job_tags(None) == ["lerobot"]
|
||||||
assert resolve_job_tags([]) == ["lerobot"]
|
assert resolve_job_tags([]) == ["lerobot"]
|
||||||
|
|||||||
Reference in New Issue
Block a user