Async checkpoint upload to avoid blocking training

lerobot/src/lerobot/scripts/lerobot_train.py

@@ -472,25 +472,31 @@ def train(cfg: TrainPipelineConfig, accelerator: Accelerator | None = None):
                     postprocessor=postprocessor,
                 )
                 update_last_checkpoint(checkpoint_dir)
-                # Upload checkpoint to HF then delete previous local copy
-                uploaded = False
+                # Upload checkpoint to HF async, delete previous after upload
                 if cfg.policy.push_to_hub:
-                    try:
-                        from huggingface_hub import HfApi
-                        api = HfApi()
-                        checkpoint_repo = cfg.policy.repo_id + "-checkpoints"
-                        api.create_repo(checkpoint_repo, exist_ok=True)
-                        api.upload_folder(
-                            folder_path=checkpoint_dir,
-                            repo_id=checkpoint_repo,
-                            path_in_repo=f"step_{step:06d}",
-                        )
-                        logging.info(f"Uploaded checkpoint step {step} to {checkpoint_repo}")
-                        uploaded = True
-                    except Exception as e:
-                        logging.warning(f"Failed to upload checkpoint step {step}: {e}")
-                if uploaded and prev_checkpoint_dir is not None and Path(prev_checkpoint_dir).exists():
-                    shutil.rmtree(prev_checkpoint_dir)
+                    import threading
+                    def _upload_and_cleanup(ckpt_dir, prev_dir, step_num):
+                        try:
+                            from huggingface_hub import HfApi
+                            api = HfApi()
+                            checkpoint_repo = cfg.policy.repo_id + "-checkpoints"
+                            api.create_repo(checkpoint_repo, exist_ok=True)
+                            api.upload_folder(
+                                folder_path=ckpt_dir,
+                                repo_id=checkpoint_repo,
+                                path_in_repo=f"step_{step_num:06d}",
+                            )
+                            logging.info(f"Uploaded checkpoint step {step_num} to {checkpoint_repo}")
+                            if prev_dir is not None and Path(prev_dir).exists():
+                                shutil.rmtree(prev_dir)
+                        except Exception as e:
+                            logging.warning(f"Failed to upload checkpoint step {step_num}: {e}")
+                    upload_thread = threading.Thread(
+                        target=_upload_and_cleanup,
+                        args=(checkpoint_dir, prev_checkpoint_dir, step),
+                        daemon=True,
+                    )
+                    upload_thread.start()
                 prev_checkpoint_dir = checkpoint_dir
                 if wandb_logger:
                     wandb_logger.log_policy(checkpoint_dir)