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).
This commit is contained in:
Nicolas Rabault
2026-06-22 15:43:52 +02:00
parent 73782447f2
commit 71c827f892
7 changed files with 179 additions and 2 deletions
View File
+17
View File
@@ -0,0 +1,17 @@
# 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.
# Importing concrete policy configs registers their draccus `--policy.type`
# choices (e.g. "act") so tests can parse them.
from lerobot.policies.act.configuration_act import ACTConfig # noqa: F401
+64
View File
@@ -0,0 +1,64 @@
# 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.
import draccus
import pytest
from lerobot.configs.default import JobConfig
from lerobot.configs.train import TrainPipelineConfig
def test_jobconfig_defaults_are_local():
cfg = JobConfig()
assert cfg.target is None
assert cfg.is_remote is False
assert cfg.image == "huggingface/lerobot-gpu:latest"
assert cfg.timeout is None
assert cfg.detach is False
def test_jobconfig_local_string_is_not_remote():
assert JobConfig(target="local").is_remote is False
def test_jobconfig_flavor_is_remote():
assert JobConfig(target="a10g-small").is_remote is True
def test_train_config_parses_job_target():
parsed = draccus.parse(
TrainPipelineConfig,
args=["--dataset.repo_id", "u/d", "--policy.type", "act", "--job.target", "a10g-small"],
)
assert parsed.job.target == "a10g-small"
assert parsed.job.is_remote is True
assert parsed.save_checkpoint_to_hub is False
def test_save_checkpoint_to_hub_requires_repo_id():
cfg = draccus.parse(
TrainPipelineConfig,
args=[
"--dataset.repo_id",
"u/d",
"--policy.type",
"act",
"--policy.push_to_hub",
"false",
"--save_checkpoint_to_hub",
"true",
],
)
with pytest.raises(ValueError, match="requires --policy.repo_id"):
cfg.validate()