Push adapter at every checkpoint (--push-checkpoints) so cancel/timeout cannot lose a run

train_securecoder.py

@@ -802,6 +802,32 @@ def train(args, model, tokenizer, records: list[dict]):
     )
 
     started = time.time()
+
+    if args.push_checkpoints and args.output_repo:
+        # Push the *adapter* at every checkpoint save, so a cancel or a timeout
+        # still leaves a resumable adapter on the hub. Repository creation races
+        # with `save_and_push` below, hence exist_ok=True.
+        from huggingface_hub import HfApi
+        from transformers import TrainerCallback
+
+        repo_id, api = args.output_repo, HfApi()
+        if os.environ.get("HF_TOKEN"):
+            api.create_repo(repo_id, repo_type="model", private=False, exist_ok=True)
+
+        class _PushAdapterOnSave(TrainerCallback):
+            def on_save(self, _args, state, control, **kwargs):
+                if state.global_step <= 0:
+                    return
+                try:
+                    _upload_adapter_files(api, f"checkpoints/checkpoint-{state.global_step}",
+                                          repo_id)
+                    log.info("checkpoint adapter @step %d pushed to %s (%.1f min in)",
+                             state.global_step, repo_id, (time.time() - started) / 60)
+                except Exception as exc:  # noqa: BLE001 - never break training
+                    log.warning("checkpoint push failed (continuing): %s", exc)
+
+        trainer.add_callback(_PushAdapterOnSave())
+
     stats = trainer.train()
     elapsed = time.time() - started
     log.info("training finished in %.1f min (final loss %.4f)",
@@ -819,6 +845,28 @@ def train(args, model, tokenizer, records: list[dict]):
     return trainer, stats, elapsed
 
 
+def _upload_adapter_files(api, folder: str, repo_id: str) -> None:
+    """Upload just the LoRA adapter (and any tokenizer files) from ``folder``
+    to the repo root, so the repo always holds the latest *resumable* adapter.
+
+    A canceled or timed-out job then leaves behind something ``--resume-adapter``
+    can continue from, instead of losing the whole run.
+    """
+    try:
+        names = os.listdir(folder)
+    except OSError:
+        return
+    for name in sorted(names):
+        wanted = (name.startswith(("adapter_", "tokenizer", "chat_template"))
+                  or name in ("special_tokens_map.json", "added_tokens.json"))
+        if not wanted:
+            continue
+        path = os.path.join(folder, name)
+        if os.path.isfile(path):
+            api.upload_file(path_or_fileobj=path, path_in_repo=name,
+                            repo_id=repo_id, repo_type="model")
+
+
 def _push_adapter(api, model, args):
     """Unsloth's push_to_hub signature varies between releases (one version
     rejects ``tokenizer=``), so fall back to uploading the saved folder - the
@@ -867,6 +915,9 @@ def parse_args(argv=None):
 
     p.add_argument("--base-model", default="Qwen/Qwen3-Coder-30B-A3B-Instruct")
     p.add_argument("--output-repo", default=None, help="Hub repo for the LoRA adapter")
+    p.add_argument("--push-checkpoints", action="store_true",
+                   help="push the adapter to --output-repo at every save step, so a "
+                        "canceled/timed-out job still leaves a resumable adapter")
     p.add_argument("--resume-adapter", default=None,
                    help="existing LoRA repo to keep training (does not init a fresh adapter)")
     p.add_argument("--merge-repo", default=None, help="optional Hub repo for a 16-bit merge")