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
+34 -1
View File
@@ -15,7 +15,7 @@
# limitations under the License.
from pathlib import Path
from unittest.mock import Mock, patch
from unittest.mock import MagicMock, Mock, patch
from lerobot.common.train_utils import (
get_step_checkpoint_dir,
@@ -24,6 +24,7 @@ from lerobot.common.train_utils import (
load_training_num_processes,
load_training_state,
load_training_step,
push_checkpoint_to_hub,
save_checkpoint,
save_training_state,
save_training_step,
@@ -151,3 +152,35 @@ def test_load_training_state_skip_optimizer(tmp_path, optimizer, scheduler):
assert loaded_step == 10
assert loaded_optimizer is optimizer
assert loaded_scheduler is scheduler
def test_push_checkpoint_to_hub_creates_repo_and_uploads(tmp_path, monkeypatch):
import huggingface_hub
ckpt = tmp_path / "010000"
(ckpt / "pretrained_model").mkdir(parents=True)
api = MagicMock()
monkeypatch.setattr(huggingface_hub, "HfApi", lambda *a, **k: api)
push_checkpoint_to_hub(ckpt, "user/run", private=True)
api.create_repo.assert_called_once()
assert api.create_repo.call_args.kwargs["private"] is True
assert api.create_repo.call_args.kwargs["repo_type"] == "model"
api.upload_folder.assert_called_once()
kwargs = api.upload_folder.call_args.kwargs
assert kwargs["repo_id"] == "user/run"
assert kwargs["repo_type"] == "model"
assert kwargs["path_in_repo"] == "checkpoints/010000"
assert kwargs["folder_path"] == str(ckpt)
assert kwargs["commit_message"] == "checkpoint 010000"
def test_push_checkpoint_to_hub_defaults_to_hub_default_visibility(tmp_path, monkeypatch):
import huggingface_hub
ckpt = tmp_path / "010000"
(ckpt / "pretrained_model").mkdir(parents=True)
api = MagicMock()
monkeypatch.setattr(huggingface_hub, "HfApi", lambda *a, **k: api)
push_checkpoint_to_hub(ckpt, "user/run")
api.create_repo.assert_called_once()
assert api.create_repo.call_args.kwargs["private"] is None