mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-20 16:31:55 +00:00
feat(train): run training remotely on HF Jobs via --job.target (#3856)
* feat(train): add JobConfig group, save_checkpoint_to_hub flag, Hub checkpoint helper
Introduce a JobConfig draccus group on TrainPipelineConfig (--job.target/image/
timeout/detach/tags) whose is_remote property gates remote dispatch, plus a
save_checkpoint_to_hub flag and validation. Add push_checkpoint_to_hub(), which
uploads a saved checkpoint directory to the model repo under checkpoints/<step>/
and creates the repo idempotently (private propagates from policy.private).
* feat(train): run training remotely on HF Jobs via --job.target
When --job.target names a GPU flavor, train() dispatches to lerobot.jobs.submit_to_hf
instead of training locally: it authenticates, ensures the dataset is on the Hub
(pushing a local-only one privately), serializes a pod-compatible train_config.json
(strips client-only fields, points at the model repo), submits via HfApi.run_job
with HF_TOKEN/WANDB_API_KEY secrets, then streams logs and finishes when the model
is pushed. Wires push_checkpoint_to_hub into the training loop behind
save_checkpoint_to_hub, and tags jobs/datasets/model with 'lerobot' + --job.tags.
* docs(train): document remote training on HF Jobs
* test(train): skip remote-dispatch tests without the dataset extra
The module imports lerobot.scripts.lerobot_train, which eagerly pulls in
lerobot.datasets (dataset extra). The base fast-test CI tier runs without
that extra, so collection failed there. Guard with pytest.importorskip,
matching the existing tests/scripts dataset-extra tests.
* refactor(jobs): hoist huggingface_hub imports to module level in hf.py
huggingface_hub is a core dependency, so the per-function dynamic imports
had no lazy-loading rationale. Move them to a single module-level import
and update test monkeypatch targets to lerobot.jobs.hf.* accordingly.
* refactor(jobs): build remote config dict via cfg.to_dict()
TrainPipelineConfig.to_dict() already returns the canonical draccus
encoding, so the StringIO + draccus.dump + json.loads round-trip was
redundant. Use it directly and drop the now-unused io/draccus imports.
* refactor(train): use module-level HfApi import in push_checkpoint_to_hub
huggingface_hub is a core dependency; the in-function import was
unnecessary. Move HfApi to a module-level import and point the test
monkeypatches at lerobot.common.train_utils.HfApi.
* refactor(configs): export JobConfig from the configs package
Re-export JobConfig in lerobot/configs/__init__.py so external callers
import it as `from lerobot.configs import JobConfig`, matching the other
config classes. Adapt the train script and test imports.
* refactor(jobs): check dataset presence with api.repo_exists
Replace the dataset_info try/except RepositoryNotFoundError dance with a
direct api.repo_exists(repo_id, repo_type="dataset") call, dropping the
httpx/RepositoryNotFoundError test scaffolding.
* chore(jobs): annotate ensure_dataset_available api param as HfApi
Add the missing HfApi type hint via a TYPE_CHECKING import.
* refactor(jobs): use HF_LEROBOT_HOME constant for the local cache root
Resolve the local dataset cache via lerobot.utils.constants.HF_LEROBOT_HOME
instead of re-reading the env var by hand, dropping the os/Path imports.
Tests now patch the imported constant and assert on a stable message
substring (the previous "neither" match only passed by accident, matching
the test name embedded in the pytest tmp_path).
* chore(jobs): guard LeRobotDataset import with require_package
Surface a clear "install lerobot[dataset]" error if the datasets extra
is missing, instead of a raw ImportError, before pushing a local dataset.
* docs(configs): clarify the is_remote_target/is_remote split
Add a comment explaining why JobConfig keeps both the staticmethod (tests
a raw target string from argv before a config exists) and the property
(accessor for an existing config instance).
* docs(train): note how to pin a pushed model version for inference
Document --policy.pretrained_revision alongside --policy.path so a
specific Hub-pushed checkpoint (once --save_checkpoint_to_hub has
committed several) can be selected for inference.
* test(jobs): skip dataset import guard in base-deps test
The fast test env installs base deps only, so require_package('datasets')
raised ImportError before the mocked lerobot.datasets import was reached.
Monkeypatch the guard to a no-op so the unit test exercises the upload logic.
* fix(jobs): address claude review findings on remote training
Resolve the claude[bot] review on #3856:
- Reject reward-model training under --job.target with a clear error instead
of crashing on a None policy inside build_remote_config_file.
- Support --policy.path remote runs: validate() no longer requires repo_id for
remote runs (it is auto-generated in submit_to_hf), and repo_id/push_to_hub
are now set after validate() resolves the policy.
- Narrow the bare `except Exception` in _tail_logs/_poll_until_done to
(OSError, httpx.HTTPError) so programming errors surface instead of being
silently retried or counted as job failures.
- Install the SIGINT detach handler only on the main thread.
- Generate model repo timestamps in UTC.
* docs(jobs): document the model-pushed marker contract and orphaned repos
Follow-up to the claude[bot] review on #3856 (non-blocking observations):
- Cross-reference the "Model pushed to <url>" log line between its producer
(PreTrainedPolicy.push_model_to_hub) and the remote-run consumer in
submit_to_hf, noting the contract is an early-finish optimization that
falls back to status polling if it drifts.
- Note in the HF Jobs guide that a failed remote run leaves its model repo
on the Hub (it is not auto-deleted) and how to remove it.
* feat(train): tag each pushed checkpoint with its step
Address review feedback on #3856: pushing a checkpoint to the Hub now
also creates a tag named after the checkpoint step, so a checkpoint can
be recovered with --policy.pretrained_revision=<step> instead of having
to look up its commit sha.
* fix(jobs): hoist ensure_dataset_available to a module-level import
Addresses Caroline's review comment on PR #3856: the local import of
ensure_dataset_available inside submit_to_hf was vestigial. dataset.py
does not import hf.py, so there is no circular-import risk and no extra
load cost (its heavy deps stay lazy), so make it a top-level import.
* refactor(configs): untangle config_path/resume resolution in validate()
Split the re-parse HACK block in TrainPipelineConfig.validate() into focused
helpers (_resolve_pretrained_from_cli, _resolve_resume_checkpoint) that handle
the policy path, reward-model path, and resume config_path as separate,
readable units. Behavior-preserving.
* feat(train): resume training from a Hub checkpoint
Allow --config_path to be a Hub repo id when resuming, not only a local path.
The latest checkpoint under checkpoints/<step>/ is downloaded into a fresh local
run dir and resumed from there (optimizer, scheduler, RNG and data order
restored as for a local resume). TrainPipelineConfig.from_pretrained falls back
to the latest checkpoint's train_config.json when a repo has no root config
(an interrupted run that only pushed checkpoints). The download is skipped when
dispatching remotely so the executor (local machine or HF Jobs pod) performs it.
- add find_latest_hub_checkpoint (utils/hub) and resolve_resume_checkpoint
(common/train_utils), the symmetric download counterpart to
push_checkpoint_to_hub
- unit tests for both helpers and the from_pretrained fallback
* feat(jobs): resume a run on HF Jobs from a checkpoint
When --resume is set with a remote --job.target, submit_to_hf resumes from the
checkpoint repo instead of staging a fresh config. A Hub config_path is resumed
in place (its checkpoint config already targets that repo); a local config_path
has its checkpoint uploaded to a new private repo first and the run is forced to
push back to it. The pod command carries --job.target=local so the checkpoint's
saved job.target can't make the pod re-dispatch itself, and the user's CLI
overrides are forwarded so a remote resume matches the same local command.
ensure_dataset_available is hoisted before the resume/fresh branch since it
applies to both.
* docs(train): document resuming from a Hub checkpoint, locally and on jobs
Show that --config_path accepts a Hub repo id for --resume, and that adding
--job.target resumes on HF Jobs (uploading a local checkpoint/dataset first).
* fix(jobs): default remote job timeout to 2d instead of the platform default
HF Jobs applies its own short 30-minute timeout when none is sent, which
silently kills long training runs. Pass an explicit, generous 2d cap by
default; users can still override --job.timeout to fail fast or extend it.
* fix(jobs): drop --dataset.root on resume + restore keyboard-control docs
Address the latest Claude review on #3856:
- _build_resume_job no longer forwards --dataset.root to the pod (a
host-local path it can't read); the fresh-run path already nulls it in
build_remote_config_file, so this makes resume consistent. Add a unit
test for _pod_forwarded_args covering the drop in both flag forms.
- Restore the display-independent keyboard-control docs (n/r/q letter
equivalents + X11/Wayland/headless Tip) in il_robots.mdx that this
branch was stale on relative to main (#3875).
* fix(jobs): handle str-typed job stage from huggingface_hub
inspect_job's status.stage is an enum (with .value) in some
huggingface_hub versions and a plain str in others. The poller
assumed the enum shape, raising "'str' object has no attribute
'value'" on resume for users on the str-returning version.
Read it via getattr(..., "value", ...) so both shapes work, and
parametrize the poll test over enum and str stages so the str case
is actually exercised (the old mock only ever simulated the enum).
* refactor(jobs): use relative import for ensure_dataset_available
* refactor(train): hoist submit_to_hf import to module top
The `from lerobot.jobs import submit_to_hf` was a function-local import in
train(); it pulls no heavy/optional deps and has no circular-import risk, so
move it to the top-level import block.
* refactor(train): hoist _remote_target_in_argv imports to module top
Move `import sys` and `from lerobot.configs import JobConfig` out of the
function body and into the top-level import block.
* refactor(utils): use relative import for sibling constants in hub.py
`from lerobot.utils.constants import CHECKPOINTS_DIR` was the odd one out in
utils/ — sibling modules there are imported relatively (.constants, .errors,
.utils, ...). Match that convention.
* refactor(jobs): hoist LeRobotDataset import, guard dataset extra at package init
Move the `from lerobot.datasets import LeRobotDataset` import to the top of
dataset.py and relocate the `require_package("datasets", extra="dataset")`
guard to the jobs package __init__, per review feedback.
* test(jobs): skip test_hf if datasets extra is missing
lerobot.configs.train pulls in datasets at import time, so the module
fails to collect without lerobot[dataset]. Guard with importorskip,
matching the convention in tests/training/test_multi_gpu.py.
* test(jobs): skip test_dataset if datasets extra is missing
tests/jobs/test_dataset.py imports lerobot.jobs.dataset, which triggers
the require_package("datasets") guard in lerobot/jobs/__init__.py at
import time. Without lerobot[dataset] the module fails to collect in the
base CI tier. Guard with importorskip, same as test_hf.py.
This commit is contained in:
@@ -15,6 +15,7 @@
|
||||
# limitations under the License.
|
||||
from pathlib import Path
|
||||
|
||||
from huggingface_hub import HfApi, snapshot_download
|
||||
from torch.optim import Optimizer
|
||||
from torch.optim.lr_scheduler import LRScheduler
|
||||
|
||||
@@ -35,6 +36,7 @@ from lerobot.utils.constants import (
|
||||
TRAINING_STATE_DIR,
|
||||
TRAINING_STEP,
|
||||
)
|
||||
from lerobot.utils.hub import find_latest_hub_checkpoint
|
||||
from lerobot.utils.io_utils import load_json, write_json
|
||||
from lerobot.utils.random_utils import load_rng_state, save_rng_state
|
||||
|
||||
@@ -283,3 +285,61 @@ def load_fsdp_optimizer_state(model, optimizer, checkpoint_dir: Path) -> None:
|
||||
with FSDP.state_dict_type(model, StateDictType.FULL_STATE_DICT, state_cfg, optim_cfg):
|
||||
sharded_osd = FSDP.optim_state_dict_to_load(model=model, optim=optimizer, optim_state_dict=full_osd)
|
||||
optimizer.load_state_dict(sharded_osd)
|
||||
|
||||
|
||||
def push_checkpoint_to_hub(
|
||||
checkpoint_dir: Path,
|
||||
repo_id: str,
|
||||
*,
|
||||
private: bool | None = None,
|
||||
) -> None:
|
||||
"""Upload a saved checkpoint directory to the Hub under checkpoints/<name>/.
|
||||
|
||||
Called once per save step when save_checkpoint_to_hub is enabled, so a
|
||||
timed-out or crashed run still leaves recoverable checkpoints on the Hub.
|
||||
The model repo is created idempotently, and the commit is tagged with the
|
||||
checkpoint step so a checkpoint can be recovered with
|
||||
--policy.pretrained_revision=<step> instead of a commit sha.
|
||||
"""
|
||||
api = HfApi()
|
||||
api.create_repo(repo_id=repo_id, repo_type="model", private=private, exist_ok=True)
|
||||
commit = api.upload_folder(
|
||||
folder_path=str(checkpoint_dir),
|
||||
repo_id=repo_id,
|
||||
repo_type="model",
|
||||
path_in_repo=f"checkpoints/{checkpoint_dir.name}",
|
||||
commit_message=f"checkpoint {checkpoint_dir.name}",
|
||||
)
|
||||
api.create_tag(
|
||||
repo_id=repo_id,
|
||||
tag=checkpoint_dir.name,
|
||||
revision=commit.oid,
|
||||
repo_type="model",
|
||||
exist_ok=True,
|
||||
)
|
||||
|
||||
|
||||
def resolve_resume_checkpoint(repo_id: str, output_dir: Path) -> Path:
|
||||
"""Download the latest checkpoint of a Hub training repo into a local run dir.
|
||||
|
||||
The symmetric counterpart to `push_checkpoint_to_hub`: given a model repo holding
|
||||
`checkpoints/<step>/{pretrained_model,training_state}` subtrees, download the highest-numbered step
|
||||
into `output_dir/checkpoints/<step>/`, recreate the local `last` symlink, and return that local
|
||||
checkpoint dir. Used to resume training from the Hub on a machine (or HF Jobs pod) that does not
|
||||
have the original local run dir.
|
||||
"""
|
||||
latest = find_latest_hub_checkpoint(repo_id)
|
||||
if latest is None:
|
||||
raise FileNotFoundError(
|
||||
f"No checkpoint found in '{repo_id}' under '{CHECKPOINTS_DIR}/'. "
|
||||
"Was the run trained with --save_checkpoint_to_hub?"
|
||||
)
|
||||
snapshot_download(
|
||||
repo_id=repo_id,
|
||||
repo_type="model",
|
||||
allow_patterns=f"{latest}/*",
|
||||
local_dir=str(output_dir),
|
||||
)
|
||||
checkpoint_dir = output_dir / latest
|
||||
update_last_checkpoint(checkpoint_dir)
|
||||
return checkpoint_dir
|
||||
|
||||
@@ -22,7 +22,7 @@ Import them directly: ``from lerobot.configs.train import TrainPipelineConfig``
|
||||
"""
|
||||
|
||||
from .dataset import DatasetRecordConfig
|
||||
from .default import DatasetConfig, EvalConfig, PeftConfig, WandBConfig
|
||||
from .default import DatasetConfig, EvalConfig, JobConfig, PeftConfig, WandBConfig
|
||||
from .policies import PreTrainedConfig
|
||||
from .recipe import MessageTurn, TrainingRecipe, load_recipe
|
||||
from .types import (
|
||||
@@ -55,6 +55,7 @@ __all__ = [
|
||||
"DatasetRecordConfig",
|
||||
"DatasetConfig",
|
||||
"EvalConfig",
|
||||
"JobConfig",
|
||||
"MessageTurn",
|
||||
"PeftConfig",
|
||||
"PreTrainedConfig",
|
||||
|
||||
@@ -145,3 +145,35 @@ class PeftConfig:
|
||||
# If None, the PEFT library defaults to alpha=8, which may dampen high-rank adapters.
|
||||
# Common values are r (alpha == rank) or 2*r.
|
||||
lora_alpha: int | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class JobConfig:
|
||||
# Where training runs. None (omitted) or "local" runs on this machine.
|
||||
# Any other value is an HF Jobs flavor and submits the run to HF Jobs.
|
||||
# List available flavors + pricing with `hf jobs hardware` command.
|
||||
target: str | None = None
|
||||
# Runtime image for the remote job (ignored for local runs).
|
||||
image: str = "huggingface/lerobot-gpu:latest"
|
||||
# Max wall-clock for the remote job as an HF Jobs duration string (e.g. "2h").
|
||||
# Defaults to "2d": We pass an explicit, generous cap instead. Set a smaller
|
||||
# value to fail fast, or a larger one for long runs.
|
||||
timeout: str | None = "2d"
|
||||
# Submit and exit instead of streaming the job logs in the foreground.
|
||||
detach: bool = False
|
||||
# Extra tags attached to the HF job and to any dataset this run pushes to the
|
||||
# Hub. A "lerobot" tag is always added; e.g. --job.tags '["lelab"]' adds more.
|
||||
tags: list[str] = field(default_factory=list)
|
||||
|
||||
# Two entry points to the same predicate: the staticmethod tests a raw target string
|
||||
# straight from argv (before any JobConfig exists, to decide dispatch early), while the
|
||||
# property is the ergonomic accessor for code that already holds a config instance.
|
||||
@staticmethod
|
||||
def is_remote_target(target: str | None) -> bool:
|
||||
"""True when `target` names an HF Jobs flavor rather than a local run."""
|
||||
return target not in (None, "local")
|
||||
|
||||
@property
|
||||
def is_remote(self) -> bool:
|
||||
"""True when training should run on HF Jobs rather than this machine."""
|
||||
return self.is_remote_target(self.target)
|
||||
|
||||
+100
-43
@@ -26,11 +26,12 @@ from huggingface_hub.errors import HfHubHTTPError
|
||||
|
||||
from lerobot import envs
|
||||
from lerobot.optim import LRSchedulerConfig, OptimizerConfig
|
||||
from lerobot.utils.hub import HubMixin
|
||||
from lerobot.utils.constants import PRETRAINED_MODEL_DIR
|
||||
from lerobot.utils.hub import HubMixin, find_latest_hub_checkpoint
|
||||
from lerobot.utils.sample_weighting import SampleWeightingConfig
|
||||
|
||||
from . import parser
|
||||
from .default import DatasetConfig, EvalConfig, PeftConfig, WandBConfig
|
||||
from .default import DatasetConfig, EvalConfig, JobConfig, PeftConfig, WandBConfig
|
||||
from .policies import PreTrainedConfig
|
||||
from .rewards import RewardModelConfig
|
||||
|
||||
@@ -83,10 +84,11 @@ class TrainPipelineConfig(HubMixin):
|
||||
# with the same value for `dir` its contents will be overwritten unless you set `resume` to true.
|
||||
output_dir: Path | None = None
|
||||
job_name: str | None = None
|
||||
# Set `resume` to true to resume a previous run. In order for this to work, you will need to make sure
|
||||
# `dir` is the directory of an existing run with at least one checkpoint in it.
|
||||
# Note that when resuming a run, the default behavior is to use the configuration from the checkpoint,
|
||||
# regardless of what's provided with the training command at the time of resumption.
|
||||
# Set `resume` to true to resume a previous run. Pass `--config_path` pointing at either a local
|
||||
# checkpoint's train_config.json or a Hub repo id holding `checkpoints/<step>/` subtrees (the
|
||||
# latest checkpoint is downloaded and resumed from). Note that when resuming, the default behavior
|
||||
# is to use the configuration from the checkpoint, regardless of what's provided with the training
|
||||
# command at the time of resumption (CLI `--*` flags still override).
|
||||
resume: bool = False
|
||||
# `seed` is used for training (eg: model initialization, dataset shuffling)
|
||||
# AND for the evaluation environments.
|
||||
@@ -118,6 +120,13 @@ class TrainPipelineConfig(HubMixin):
|
||||
wandb: WandBConfig = field(default_factory=WandBConfig)
|
||||
peft: PeftConfig | None = None
|
||||
|
||||
# Where to run training (local default, or an HF Jobs flavor). See JobConfig.
|
||||
job: JobConfig = field(default_factory=JobConfig)
|
||||
# Push each saved checkpoint to the Hub (policy.repo_id) as it is written, not
|
||||
# just the final model (useful to monitor progress mid-run). Optional; the
|
||||
# final model is pushed regardless. Works the same locally and remotely.
|
||||
save_checkpoint_to_hub: bool = False
|
||||
|
||||
# Sample weighting configuration (e.g., for RA-BC training)
|
||||
sample_weighting: SampleWeightingConfig | None = None
|
||||
|
||||
@@ -137,10 +146,17 @@ class TrainPipelineConfig(HubMixin):
|
||||
return self.reward_model # type: ignore[return-value]
|
||||
return self.policy # type: ignore[return-value]
|
||||
|
||||
def validate(self) -> None:
|
||||
# HACK: We parse again the cli args here to get the pretrained paths if there was some.
|
||||
policy_path = parser.get_path_arg("policy")
|
||||
def _resolve_pretrained_from_cli(self) -> None:
|
||||
"""Resolve the pretrained source passed on the CLI into a loaded config.
|
||||
|
||||
The pretrained paths (`--policy.path`, `--reward_model.path`) and
|
||||
`--config_path` are only recoverable by re-reading the CLI args: draccus
|
||||
has already consumed them by the time `validate()` runs, so they are not
|
||||
reflected on `self`. Exactly one source applies, in priority order:
|
||||
reward-model path, policy path, then resume.
|
||||
"""
|
||||
reward_model_path = parser.get_path_arg("reward_model")
|
||||
policy_path = parser.get_path_arg("policy")
|
||||
|
||||
if reward_model_path:
|
||||
cli_overrides = parser.get_cli_overrides("reward_model")
|
||||
@@ -149,31 +165,54 @@ class TrainPipelineConfig(HubMixin):
|
||||
)
|
||||
self.reward_model.pretrained_path = str(Path(reward_model_path))
|
||||
elif policy_path:
|
||||
yaml_overrides = parser.get_yaml_overrides("policy")
|
||||
cli_overrides = parser.get_cli_overrides("policy") or []
|
||||
self.policy = PreTrainedConfig.from_pretrained(
|
||||
policy_path, cli_overrides=yaml_overrides + cli_overrides
|
||||
)
|
||||
overrides = parser.get_yaml_overrides("policy") + (parser.get_cli_overrides("policy") or [])
|
||||
self.policy = PreTrainedConfig.from_pretrained(policy_path, cli_overrides=overrides)
|
||||
self.policy.pretrained_path = Path(policy_path)
|
||||
elif self.resume:
|
||||
config_path = parser.parse_arg("config_path")
|
||||
if not config_path:
|
||||
raise ValueError(
|
||||
f"A config_path is expected when resuming a run. Please specify path to {TRAIN_CONFIG_NAME}"
|
||||
)
|
||||
self._resolve_resume_checkpoint()
|
||||
|
||||
if not Path(config_path).resolve().exists():
|
||||
raise NotADirectoryError(
|
||||
f"{config_path=} is expected to be a local path. "
|
||||
"Resuming from the hub is not supported for now."
|
||||
)
|
||||
def _resolve_resume_checkpoint(self) -> None:
|
||||
"""Point the trainable config at the checkpoint named by `--config_path`.
|
||||
|
||||
`config_path` is either a local path (to a checkpoint's train_config.json or its
|
||||
pretrained_model/ dir) or a Hub repo id. For a Hub repo, the latest checkpoint is downloaded
|
||||
into a fresh local run dir and resumed from there. The download is skipped when dispatching to
|
||||
an HF Job (`job.is_remote`): the pod performs it when it runs the resume locally, and
|
||||
`submit_to_hf` resolves the source repo for the remote command.
|
||||
"""
|
||||
config_path = parser.parse_arg("config_path")
|
||||
if not config_path:
|
||||
raise ValueError(
|
||||
f"A config_path is expected when resuming a run. Please specify path to {TRAIN_CONFIG_NAME}"
|
||||
)
|
||||
|
||||
if Path(config_path).resolve().exists():
|
||||
policy_dir = Path(config_path).parent
|
||||
if self.policy is not None:
|
||||
self.policy.pretrained_path = policy_dir
|
||||
if self.reward_model is not None:
|
||||
self.reward_model.pretrained_path = str(policy_dir)
|
||||
self.checkpoint_path = policy_dir.parent
|
||||
elif self.job.is_remote:
|
||||
return
|
||||
else:
|
||||
from lerobot.common.train_utils import resolve_resume_checkpoint
|
||||
|
||||
# `self.output_dir` was loaded from the checkpoint's config and points at the original
|
||||
# run's (now-absent) local dir. Resume into a fresh local dir instead, unless the user
|
||||
# passed --output_dir explicitly.
|
||||
cli_output_dir = parser.parse_arg("output_dir")
|
||||
if cli_output_dir:
|
||||
self.output_dir = Path(cli_output_dir)
|
||||
else:
|
||||
now = dt.datetime.now()
|
||||
self.output_dir = Path("outputs/train") / f"{now:%Y-%m-%d}/{now:%H-%M-%S}_resume"
|
||||
self.checkpoint_path = resolve_resume_checkpoint(config_path, self.output_dir)
|
||||
policy_dir = self.checkpoint_path / PRETRAINED_MODEL_DIR
|
||||
|
||||
if self.policy is not None:
|
||||
self.policy.pretrained_path = policy_dir
|
||||
if self.reward_model is not None:
|
||||
self.reward_model.pretrained_path = str(policy_dir)
|
||||
|
||||
def validate(self) -> None:
|
||||
self._resolve_pretrained_from_cli()
|
||||
|
||||
if self.policy is None and self.reward_model is None:
|
||||
raise ValueError(
|
||||
@@ -216,9 +255,19 @@ class TrainPipelineConfig(HubMixin):
|
||||
if self.eval_steps > 0 and self.dataset.eval_split == 0.0:
|
||||
raise ValueError("eval_steps > 0 requires dataset.eval_split > 0.0 to hold out eval data.")
|
||||
|
||||
if hasattr(active_cfg, "push_to_hub") and active_cfg.push_to_hub and not active_cfg.repo_id:
|
||||
# Remote runs auto-generate the repo_id in submit_to_hf (the policy may only be
|
||||
# resolved here, from --policy.path), so don't demand it up front for them.
|
||||
if (
|
||||
hasattr(active_cfg, "push_to_hub")
|
||||
and active_cfg.push_to_hub
|
||||
and not active_cfg.repo_id
|
||||
and not self.job.is_remote
|
||||
):
|
||||
raise ValueError("'repo_id' argument missing. Please specify it to push the model to the hub.")
|
||||
|
||||
if self.save_checkpoint_to_hub and not (self.policy is not None and self.policy.repo_id):
|
||||
raise ValueError("save_checkpoint_to_hub requires --policy.repo_id.")
|
||||
|
||||
@classmethod
|
||||
def __get_path_fields__(cls) -> list[str]:
|
||||
"""Keys for draccus pretrained-path loading."""
|
||||
@@ -255,22 +304,30 @@ class TrainPipelineConfig(HubMixin):
|
||||
elif Path(model_id).is_file():
|
||||
config_file = model_id
|
||||
else:
|
||||
dl_kwargs = {
|
||||
"repo_id": model_id,
|
||||
"revision": revision,
|
||||
"cache_dir": cache_dir,
|
||||
"force_download": force_download,
|
||||
"proxies": proxies,
|
||||
"resume_download": resume_download,
|
||||
"token": token,
|
||||
"local_files_only": local_files_only,
|
||||
}
|
||||
try:
|
||||
config_file = hf_hub_download(
|
||||
repo_id=model_id,
|
||||
filename=TRAIN_CONFIG_NAME,
|
||||
revision=revision,
|
||||
cache_dir=cache_dir,
|
||||
force_download=force_download,
|
||||
proxies=proxies,
|
||||
resume_download=resume_download,
|
||||
token=token,
|
||||
local_files_only=local_files_only,
|
||||
)
|
||||
config_file = hf_hub_download(filename=TRAIN_CONFIG_NAME, **dl_kwargs)
|
||||
except HfHubHTTPError as e:
|
||||
raise FileNotFoundError(
|
||||
f"{TRAIN_CONFIG_NAME} not found on the HuggingFace Hub in {model_id}"
|
||||
) from e
|
||||
# No root train_config.json: this is a repo of periodic checkpoints from an
|
||||
# interrupted run. Fall back to the latest checkpoint's config so the run can be
|
||||
# resumed straight from the repo with `--config_path=<repo>`.
|
||||
latest = find_latest_hub_checkpoint(model_id, token=token, revision=revision)
|
||||
if latest is None:
|
||||
raise FileNotFoundError(
|
||||
f"{TRAIN_CONFIG_NAME} not found on the HuggingFace Hub in {model_id}"
|
||||
) from e
|
||||
config_file = hf_hub_download(
|
||||
filename=f"{latest}/{PRETRAINED_MODEL_DIR}/{TRAIN_CONFIG_NAME}", **dl_kwargs
|
||||
)
|
||||
|
||||
cli_args = kwargs.pop("cli_args", [])
|
||||
# Legacy RA-BC migration only applies to framework-saved checkpoints (always JSON).
|
||||
|
||||
@@ -0,0 +1,23 @@
|
||||
# Copyright 2025 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.
|
||||
|
||||
from lerobot.utils.import_utils import require_package
|
||||
|
||||
# LeRobotDataset (imported at module top in dataset.py) pulls in heavy dataset deps;
|
||||
# guard the optional dependency here so importing this package fails loudly if it's missing.
|
||||
require_package("datasets", extra="dataset")
|
||||
|
||||
from .hf import submit_to_hf
|
||||
|
||||
__all__ = ["submit_to_hf"]
|
||||
@@ -0,0 +1,53 @@
|
||||
# Copyright 2025 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.
|
||||
"""Make a training dataset reachable from an HF Job pod.
|
||||
|
||||
The pod can't see the host's ~/.cache/huggingface/lerobot, so the dataset has to
|
||||
live on the Hub: the pod downloads it by repo_id at train time (the forwarded
|
||||
HF_TOKEN covers private datasets). A dataset already on the Hub is used as-is; a
|
||||
local-only dataset is pushed to a PRIVATE repo first (never public).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from lerobot.datasets import LeRobotDataset
|
||||
from lerobot.utils.constants import HF_LEROBOT_HOME
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from huggingface_hub import HfApi
|
||||
|
||||
|
||||
def ensure_dataset_available(repo_id: str, *, api: HfApi, tags: list[str] | None = None) -> None:
|
||||
"""Ensure repo_id resolves on the Hub, pushing a local-only dataset privately first.
|
||||
|
||||
`tags` are attached to the dataset only when we push it (an already-on-Hub
|
||||
dataset is left untouched). Raises RuntimeError if the dataset is neither on
|
||||
the Hub nor in the local cache.
|
||||
"""
|
||||
if api.repo_exists(repo_id, repo_type="dataset"):
|
||||
return
|
||||
|
||||
local_present = (HF_LEROBOT_HOME / repo_id / "meta" / "info.json").is_file()
|
||||
if not local_present:
|
||||
raise RuntimeError(
|
||||
f"Dataset '{repo_id}' is not in the local cache ({HF_LEROBOT_HOME}) and could not be "
|
||||
f"reached on the Hub — it may not exist, or be private and inaccessible with your "
|
||||
f"token. Record or download it first, or run `hf auth login`."
|
||||
)
|
||||
|
||||
print(f"[dataset] '{repo_id}' is local-only; pushing to a PRIVATE Hub repo...")
|
||||
LeRobotDataset(repo_id).push_to_hub(private=True, tags=tags)
|
||||
print(f"[dataset] '{repo_id}' uploaded (private). The job will download it by repo_id.")
|
||||
@@ -0,0 +1,425 @@
|
||||
# Copyright 2025 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 a lerobot training on HF Jobs (HuggingFace GPUs).
|
||||
|
||||
Ported and simplified from lelab's runners/hf_cloud.py: no UI log queue, no
|
||||
registry — just submit and stream to stdout.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
import datetime as dt
|
||||
import json
|
||||
import netrc
|
||||
import os
|
||||
import re
|
||||
import signal
|
||||
import sys
|
||||
import tempfile
|
||||
import threading
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import httpx
|
||||
from huggingface_hub import (
|
||||
HfApi,
|
||||
create_repo,
|
||||
fetch_job_logs,
|
||||
get_token,
|
||||
inspect_job,
|
||||
run_job,
|
||||
upload_file,
|
||||
)
|
||||
|
||||
from lerobot.common.train_utils import push_checkpoint_to_hub
|
||||
from lerobot.configs import parser
|
||||
|
||||
from .dataset import ensure_dataset_available
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from lerobot.configs.train import TrainPipelineConfig
|
||||
|
||||
_SLUG_RE = re.compile(r"[^a-zA-Z0-9._-]+")
|
||||
|
||||
_TERMINAL_STAGES = {"COMPLETED", "CANCELED", "ERROR", "DELETED"}
|
||||
|
||||
# huggingface_hub 1.x runs on httpx: transient HTTP/transport failures surface as
|
||||
# httpx.HTTPError and socket-level errors as OSError. Catching only these keeps real
|
||||
# bugs (TypeError, AttributeError, ...) from being silently retried or counted as
|
||||
# job failures.
|
||||
_TRANSIENT_NET_ERRORS = (OSError, httpx.HTTPError)
|
||||
|
||||
# Always attached to remote jobs and pushed datasets so LeRobot-originated work
|
||||
# is identifiable on the Hub; callers (e.g. LeLab) add their own via --job.tags.
|
||||
LEROBOT_TAG = "lerobot"
|
||||
|
||||
|
||||
def resolve_job_tags(extra: list[str] | None) -> list[str]:
|
||||
"""Return the tag list for a run: the lerobot tag plus any extras, deduped, order-stable."""
|
||||
tags = [LEROBOT_TAG, *(extra or [])]
|
||||
seen: set[str] = set()
|
||||
return [t for t in tags if not (t in seen or seen.add(t))]
|
||||
|
||||
|
||||
def resolve_wandb_api_key() -> str | None:
|
||||
"""Host's wandb key for forwarding to the job: $WANDB_API_KEY, else ~/.netrc."""
|
||||
key = os.environ.get("WANDB_API_KEY")
|
||||
if key:
|
||||
return key
|
||||
try:
|
||||
rc = netrc.netrc()
|
||||
except (FileNotFoundError, netrc.NetrcParseError, OSError):
|
||||
return None
|
||||
auth = rc.authenticators("api.wandb.ai")
|
||||
if auth is None:
|
||||
return None
|
||||
_login, _account, password = auth
|
||||
return password or None
|
||||
|
||||
|
||||
def build_repo_id(username: str, job_name: str, now: dt.datetime) -> str:
|
||||
"""Generate the model repo id for a remote run: <user>/<job_name>_<timestamp>."""
|
||||
slug = _SLUG_RE.sub("-", job_name).strip("-") or "train"
|
||||
stamp = now.strftime("%Y-%m-%d_%H-%M-%S")
|
||||
return f"{username}/{slug}_{stamp}"
|
||||
|
||||
|
||||
def build_remote_config_file(cfg, repo_id: str, dest: Path, tags: list[str] | None = None) -> Path:
|
||||
"""Write a train_config.json for the pod, with remote overrides applied.
|
||||
|
||||
The pod runs `lerobot-train --config_path=<dest>` and downloads the dataset
|
||||
by repo_id into its own cache. Client-only fields are stripped so the config
|
||||
is accepted by the trainer image: `job` (pure client orchestration) is always
|
||||
removed, and `save_checkpoint_to_hub` is removed unless explicitly enabled —
|
||||
older lerobot images reject unknown keys, so the default keeps the config
|
||||
compatible with the released `lerobot-gpu` image. `tags` are merged into
|
||||
policy.tags so the trained model the pod pushes carries them too.
|
||||
"""
|
||||
remote = copy.deepcopy(cfg)
|
||||
remote.policy.push_to_hub = True
|
||||
remote.policy.repo_id = repo_id
|
||||
# Don't pin the client's resolved device (e.g. "mps"); let the pod auto-detect its GPU.
|
||||
remote.policy.device = None
|
||||
# Drop any host-local dataset root; the pod resolves the dataset by repo_id.
|
||||
remote.dataset.root = None
|
||||
if tags:
|
||||
existing = list(remote.policy.tags or [])
|
||||
remote.policy.tags = existing + [t for t in tags if t not in existing]
|
||||
|
||||
# Encode to the canonical, pod-parseable dict, then drop the keys the released
|
||||
# trainer image doesn't know about.
|
||||
data = remote.to_dict()
|
||||
data.pop("job", None)
|
||||
if not remote.save_checkpoint_to_hub:
|
||||
data.pop("save_checkpoint_to_hub", None)
|
||||
|
||||
dest.parent.mkdir(parents=True, exist_ok=True)
|
||||
dest.write_text(json.dumps(data, indent=4))
|
||||
return dest
|
||||
|
||||
|
||||
def _stage_config_on_hub(cfg, repo_id: str, token: str, tags: list[str] | None = None) -> str:
|
||||
"""Upload train_config.json to the model repo and return the repo_id for --config_path."""
|
||||
create_repo(repo_id, repo_type="model", private=True, exist_ok=True, token=token)
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
config_path = build_remote_config_file(cfg, repo_id, Path(tmp) / "train_config.json", tags=tags)
|
||||
upload_file(
|
||||
path_or_fileobj=config_path,
|
||||
path_in_repo="train_config.json",
|
||||
repo_id=repo_id,
|
||||
repo_type="model",
|
||||
token=token,
|
||||
)
|
||||
return repo_id
|
||||
|
||||
|
||||
def _tail_logs(
|
||||
job_id: str,
|
||||
done: threading.Event,
|
||||
success_marker: str | None = None,
|
||||
success_event: threading.Event | None = None,
|
||||
) -> None:
|
||||
"""Stream job logs to stdout, reconnecting on dropped streams until done is set.
|
||||
|
||||
Each reconnect re-fetches the full buffered log, so we track how many lines
|
||||
were already printed and skip them — otherwise a fast-failing job's traceback
|
||||
gets reprinted on every reconnect.
|
||||
|
||||
When `success_marker` appears in a line, set `success_event` and `done` so the
|
||||
caller can finish as soon as the trained model lands on the Hub, rather than
|
||||
waiting out the platform's post-run finalization (which can add ~30s).
|
||||
"""
|
||||
printed = 0
|
||||
while not done.is_set():
|
||||
try:
|
||||
seen = 0
|
||||
for line in fetch_job_logs(job_id=job_id, follow=True):
|
||||
seen += 1
|
||||
if seen <= printed:
|
||||
continue # already shown on a previous connection
|
||||
printed = seen
|
||||
# fetch_job_logs yields SSE data without trailing newlines, so add one
|
||||
# per entry — otherwise all log lines concatenate onto a single line.
|
||||
print(line.rstrip("\n"), flush=True)
|
||||
if success_marker and success_event is not None and success_marker in line:
|
||||
success_event.set()
|
||||
done.set()
|
||||
return
|
||||
if done.is_set():
|
||||
return
|
||||
# Stream closed cleanly. Wait a moment so the status poller can mark
|
||||
# the job terminal before we reconnect (avoids re-tailing the buffer).
|
||||
if done.wait(3):
|
||||
return
|
||||
except _TRANSIENT_NET_ERRORS:
|
||||
if done.wait(2):
|
||||
return
|
||||
|
||||
|
||||
def _poll_until_done(
|
||||
job_id: str,
|
||||
done: threading.Event,
|
||||
poll_interval: float = 5.0,
|
||||
status_holder: dict | None = None,
|
||||
max_failures: int = 6,
|
||||
) -> str | None:
|
||||
"""Poll inspect_job until a terminal stage or until `done` is set.
|
||||
|
||||
Returns the terminal stage string, or None if `done` was set first (detach)
|
||||
or after `max_failures` consecutive inspect_job errors. When a terminal stage
|
||||
is reached and `status_holder` is given, records `status_holder["message"]`
|
||||
(the platform's status message, e.g. "Job timeout").
|
||||
"""
|
||||
failures = 0
|
||||
while not done.is_set():
|
||||
try:
|
||||
info = inspect_job(job_id=job_id)
|
||||
failures = 0
|
||||
# `stage` is an enum in some huggingface_hub versions and a plain str in others.
|
||||
stage = getattr(info.status.stage, "value", info.status.stage)
|
||||
if stage in _TERMINAL_STAGES:
|
||||
if status_holder is not None:
|
||||
status_holder["message"] = getattr(info.status, "message", None)
|
||||
done.set()
|
||||
return stage
|
||||
except _TRANSIENT_NET_ERRORS:
|
||||
failures += 1
|
||||
if failures >= max_failures:
|
||||
done.set()
|
||||
return None
|
||||
done.wait(poll_interval)
|
||||
return None
|
||||
|
||||
|
||||
def _pod_forwarded_args(
|
||||
argv: list[str], drop_names: tuple[str, ...] = (), drop_prefixes: tuple[str, ...] = ()
|
||||
) -> list[str]:
|
||||
"""User CLI overrides to replay on the pod, minus flags the submitter sets itself.
|
||||
|
||||
Handles both `--name=value` and `--name value` forms. Forwarding the user's overrides (e.g.
|
||||
`--steps`, `--save_checkpoint_to_hub`) makes a remote resume behave like the same local command.
|
||||
"""
|
||||
out: list[str] = []
|
||||
skip_next = False
|
||||
for i, tok in enumerate(argv):
|
||||
if skip_next:
|
||||
skip_next = False
|
||||
continue
|
||||
name = tok.split("=", 1)[0]
|
||||
if name in drop_names or any(name.startswith(p) for p in drop_prefixes):
|
||||
if "=" not in tok and i + 1 < len(argv) and not argv[i + 1].startswith("--"):
|
||||
skip_next = True # also drop the space-separated value
|
||||
continue
|
||||
out.append(tok)
|
||||
return out
|
||||
|
||||
|
||||
def _build_resume_job(cfg: TrainPipelineConfig, username: str) -> tuple[str, list[str]]:
|
||||
"""Resolve the model repo and pod command to resume a run on a job.
|
||||
|
||||
A Hub `config_path` is resumed from directly: its checkpoint config already targets that repo,
|
||||
so new checkpoints continue the lineage there. A local `config_path` has its checkpoint uploaded
|
||||
to a new PRIVATE repo first, and the resumed run is forced to push back to it. The pod command
|
||||
always carries `--job.target=local` so the checkpoint's saved `job.target` can't make the pod
|
||||
re-dispatch itself.
|
||||
"""
|
||||
config_path = parser.parse_arg("config_path")
|
||||
forwarded = _pod_forwarded_args(
|
||||
sys.argv[1:],
|
||||
drop_names=("--config_path", "--policy.repo_id", "--policy.push_to_hub", "--dataset.root"),
|
||||
drop_prefixes=("--job.",),
|
||||
)
|
||||
|
||||
if Path(config_path).exists():
|
||||
# Local checkpoint: stage it on the Hub so the pod can resume from it, and push back there.
|
||||
# Resolve so a `last` symlink uploads under its real step name (digit), which the pod's
|
||||
# latest-checkpoint lookup keys on.
|
||||
checkpoint_dir = Path(cfg.checkpoint_path).resolve()
|
||||
source_repo = build_repo_id(username, cfg.job_name or "train", dt.datetime.now(dt.UTC))
|
||||
push_checkpoint_to_hub(checkpoint_dir, source_repo, private=True)
|
||||
extra = [f"--policy.repo_id={source_repo}", "--policy.push_to_hub=true"]
|
||||
else:
|
||||
source_repo = config_path
|
||||
extra = []
|
||||
|
||||
command = [
|
||||
"lerobot-train",
|
||||
*forwarded,
|
||||
f"--config_path={source_repo}",
|
||||
"--job.target=local",
|
||||
*extra,
|
||||
]
|
||||
return source_repo, command
|
||||
|
||||
|
||||
def submit_to_hf(cfg: TrainPipelineConfig) -> None:
|
||||
"""Submit a training job to HF Jobs infrastructure.
|
||||
|
||||
Validates cfg, resolves credentials, ensures the dataset is on the Hub, then either stages a
|
||||
sanitized config (fresh run) or resumes from a checkpoint repo, submits the job, and tails logs
|
||||
until completion or detaches immediately. 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.")
|
||||
|
||||
api = HfApi(token=token)
|
||||
user_info = api.whoami(token=token)
|
||||
username = user_info["name"]
|
||||
|
||||
now = dt.datetime.now(dt.UTC)
|
||||
fresh_repo_id: str | None = None
|
||||
if not cfg.resume:
|
||||
# Resolve the model repo and mark it for push BEFORE validate(): validate() requires repo_id
|
||||
# to be set whenever push_to_hub is True. (A resume reuses the checkpoint's repo instead.)
|
||||
if cfg.policy is not None:
|
||||
base_name = cfg.job_name or cfg.policy.type
|
||||
fresh_repo_id = cfg.policy.repo_id or build_repo_id(username, base_name, now)
|
||||
cfg.policy.repo_id = fresh_repo_id
|
||||
cfg.policy.push_to_hub = True
|
||||
else:
|
||||
# Path-based policy is resolved inside validate(); fall back to a generic slug.
|
||||
fresh_repo_id = build_repo_id(username, cfg.job_name or "train", now)
|
||||
|
||||
cfg.validate()
|
||||
|
||||
if cfg.is_reward_model_training:
|
||||
raise ValueError(
|
||||
"Remote training via --job.target only supports policy training, not reward models. "
|
||||
"Run reward-model training locally."
|
||||
)
|
||||
|
||||
secrets: dict[str, str] = {"HF_TOKEN": token}
|
||||
if cfg.wandb.enable:
|
||||
wandb_key = resolve_wandb_api_key()
|
||||
if wandb_key is None:
|
||||
raise ValueError(
|
||||
"wandb is enabled but no WANDB_API_KEY found. "
|
||||
"Set it via `export WANDB_API_KEY=...` or add it to ~/.netrc."
|
||||
)
|
||||
secrets["WANDB_API_KEY"] = wandb_key
|
||||
|
||||
tags = resolve_job_tags(cfg.job.tags)
|
||||
# The dataset must be reachable from the pod for both fresh and resumed runs; a local-only
|
||||
# dataset is pushed PRIVATE here. Hoisted before the resume/fresh branch since it applies to both.
|
||||
ensure_dataset_available(cfg.dataset.repo_id, api=api, tags=tags)
|
||||
|
||||
if cfg.resume:
|
||||
repo_id, command = _build_resume_job(cfg, username)
|
||||
else:
|
||||
config_repo_id = _stage_config_on_hub(cfg, fresh_repo_id, token, tags=tags)
|
||||
repo_id = fresh_repo_id
|
||||
command = ["lerobot-train", f"--config_path={config_repo_id}"]
|
||||
|
||||
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=secrets,
|
||||
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}")
|
||||
print(f" Model repo: https://huggingface.co/{repo_id}")
|
||||
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():
|
||||
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}"
|
||||
)
|
||||
@@ -340,6 +340,9 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
|
||||
ignore_patterns=["*.tmp", "*.log"],
|
||||
)
|
||||
|
||||
# Contract: lerobot.jobs.hf.submit_to_hf watches for this exact
|
||||
# "Model pushed to <url>" line to end a remote run early. Keep the wording
|
||||
# and URL format in sync (it falls back to status polling if they drift).
|
||||
logging.info(f"Model pushed to {commit_info.repo_url.url}")
|
||||
|
||||
def generate_model_card(
|
||||
|
||||
@@ -20,6 +20,7 @@ Requires: pip install 'lerobot[training]' (includes dataset + accelerate + wand
|
||||
|
||||
import dataclasses
|
||||
import logging
|
||||
import sys
|
||||
import time
|
||||
from contextlib import nullcontext
|
||||
from pprint import pformat
|
||||
@@ -41,15 +42,17 @@ from lerobot.common.train_utils import (
|
||||
load_training_batch_size,
|
||||
load_training_num_processes,
|
||||
load_training_state,
|
||||
push_checkpoint_to_hub,
|
||||
save_checkpoint,
|
||||
update_last_checkpoint,
|
||||
)
|
||||
from lerobot.common.wandb_utils import WandBLogger
|
||||
from lerobot.configs import parser
|
||||
from lerobot.configs import JobConfig, parser
|
||||
from lerobot.configs.train import TrainPipelineConfig
|
||||
from lerobot.datasets import EpisodeAwareSampler, compute_sampler_state
|
||||
from lerobot.datasets.factory import make_train_eval_datasets
|
||||
from lerobot.envs import close_envs, make_env, make_env_pre_post_processors
|
||||
from lerobot.jobs import submit_to_hf
|
||||
from lerobot.optim.factory import make_optimizer_and_scheduler
|
||||
from lerobot.policies import PreTrainedPolicy, make_policy, make_pre_post_processors
|
||||
from lerobot.rewards import make_reward_pre_post_processors
|
||||
@@ -188,6 +191,9 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
|
||||
cfg: A `TrainPipelineConfig` object containing all training configurations.
|
||||
accelerator: Optional Accelerator instance. If None, one will be created automatically.
|
||||
"""
|
||||
if cfg.job.is_remote:
|
||||
return submit_to_hf(cfg)
|
||||
|
||||
from lerobot.utils.import_utils import require_package
|
||||
|
||||
require_package("accelerate", extra="training")
|
||||
@@ -655,6 +661,12 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
|
||||
optim_state_dict=optim_state_dict,
|
||||
)
|
||||
update_last_checkpoint(checkpoint_dir)
|
||||
if cfg.save_checkpoint_to_hub:
|
||||
push_checkpoint_to_hub(
|
||||
checkpoint_dir,
|
||||
cfg.policy.repo_id,
|
||||
private=cfg.policy.private,
|
||||
)
|
||||
if wandb_logger:
|
||||
wandb_logger.log_policy(checkpoint_dir)
|
||||
|
||||
@@ -735,8 +747,25 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
|
||||
accelerator.end_training()
|
||||
|
||||
|
||||
def _remote_target_in_argv() -> bool:
|
||||
"""True when the CLI requests a remote HF Jobs run (--job.target=<non-local>)."""
|
||||
target = None
|
||||
args = sys.argv[1:]
|
||||
for i, tok in enumerate(args):
|
||||
if tok == "--job.target" and i + 1 < len(args):
|
||||
target = args[i + 1]
|
||||
elif tok.startswith("--job.target="):
|
||||
target = tok.split("=", 1)[1]
|
||||
return JobConfig.is_remote_target(target)
|
||||
|
||||
|
||||
def main():
|
||||
register_third_party_plugins()
|
||||
if _remote_target_in_argv():
|
||||
# The policy device is resolved on the remote pod, not here, so silence the
|
||||
# client-side "Device '...' is not available" warning PreTrainedConfig emits
|
||||
# while parsing the config (it fires before train() can dispatch remotely).
|
||||
logging.getLogger("lerobot.configs.policies").setLevel(logging.ERROR)
|
||||
train()
|
||||
|
||||
|
||||
|
||||
@@ -20,9 +20,33 @@ from typing import Any, TypeVar
|
||||
from huggingface_hub import HfApi
|
||||
from huggingface_hub.utils import validate_hf_hub_args
|
||||
|
||||
from .constants import CHECKPOINTS_DIR
|
||||
|
||||
T = TypeVar("T", bound="HubMixin")
|
||||
|
||||
|
||||
def find_latest_hub_checkpoint(
|
||||
repo_id: str,
|
||||
*,
|
||||
token: str | bool | None = None,
|
||||
revision: str | None = None,
|
||||
) -> str | None:
|
||||
"""Repo-relative path of the most recent checkpoint in a training repo.
|
||||
|
||||
Training runs push checkpoints to ``checkpoints/<step>/`` (see
|
||||
``push_checkpoint_to_hub``). This lists those step dirs and returns
|
||||
``checkpoints/<highest-step>``, or ``None`` if the repo has no checkpoints.
|
||||
"""
|
||||
files = HfApi().list_repo_files(repo_id=repo_id, repo_type="model", revision=revision, token=token)
|
||||
prefix = f"{CHECKPOINTS_DIR}/"
|
||||
steps = {
|
||||
name for f in files if f.startswith(prefix) and (name := f[len(prefix) :].split("/", 1)[0]).isdigit()
|
||||
}
|
||||
if not steps:
|
||||
return None
|
||||
return f"{CHECKPOINTS_DIR}/{max(steps, key=int)}"
|
||||
|
||||
|
||||
class HubMixin:
|
||||
"""
|
||||
A Mixin containing the functionality to push an object to the hub.
|
||||
|
||||
Reference in New Issue
Block a user