Add HF upload timeout and backlog fail-fast guards

Author: Dieg0Code

src/training/checkpointing.py

@@ -3,6 +3,7 @@
 import json
 import os
 from concurrent.futures import Future
+from concurrent.futures import TimeoutError as FutureTimeoutError
 from datetime import datetime, timezone
 from pathlib import Path
 from typing import TYPE_CHECKING, Any
@@ -296,6 +297,29 @@ def drain_completed_hf_uploads(
     return pending
 
 
+def wait_for_hf_uploads(
+    futures: list[Future[None]],
+    *,
+    timeout_s: float,
+    fail_on_error: bool,
+) -> None:
+    """Wait for pending HF uploads with a bounded timeout per future."""
+    bounded_timeout = max(1.0, float(timeout_s))
+    for future in futures:
+        try:
+            future.result(timeout=bounded_timeout)
+        except FutureTimeoutError as exc:
+            if fail_on_error:
+                raise RuntimeError(
+                    f"HF upload timed out after {bounded_timeout:.1f}s per future.",
+                ) from exc
+            log(f"HF upload timed out (continuing): {exc}")
+        except Exception as exc:
+            if fail_on_error:
+                raise RuntimeError("HF upload future failed during shutdown wait.") from exc
+            log(f"HF upload failed during shutdown (continuing): {exc}")
+
+
 def should_save_iteration_checkpoint(iteration: int, total_iterations: int, save_every: int) -> bool:
     """Always persist the final iteration even when it is not divisible by save_every."""
     if iteration >= total_iterations:
@@ -311,4 +335,5 @@ def should_save_iteration_checkpoint(iteration: int, total_iterations: int, save
     "ensure_hf_ready",
     "init_hf_checkpointer",
     "should_save_iteration_checkpoint",
+    "wait_for_hf_uploads",
 ]

src/training/config_runtime.py

@@ -58,6 +58,8 @@
     "hf_run_id": "policy_spatial_v1",
     "hf_token_env": "HF_TOKEN",
     "hf_local_dir": "hf_checkpoints",
+    "max_pending_hf_uploads": 2,
+    "hf_upload_future_timeout_s": 120.0,
     "show_progress_bar": False,
     "trainer_log_every_n_steps": 99_999,
     "monitor_log_every": 5,
@@ -170,6 +172,8 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--hf", action="store_true")
     parser.add_argument("--hf-repo-id", default=None)
     parser.add_argument("--hf-run-id", default=None)
+    parser.add_argument("--max-pending-hf-uploads", type=int, default=None)
+    parser.add_argument("--hf-upload-timeout-s", type=float, default=None)
     return parser.parse_args()
 
 
@@ -285,6 +289,10 @@ def apply_cli_overrides(args: argparse.Namespace) -> None:
         CONFIG["hf_repo_id"] = args.hf_repo_id
     if args.hf_run_id is not None:
         CONFIG["hf_run_id"] = args.hf_run_id.strip()
+    if args.max_pending_hf_uploads is not None:
+        CONFIG["max_pending_hf_uploads"] = max(1, args.max_pending_hf_uploads)
+    if args.hf_upload_timeout_s is not None:
+        CONFIG["hf_upload_future_timeout_s"] = max(1.0, args.hf_upload_timeout_s)
 
 
 def cfg_int(key: str) -> int:
@@ -344,6 +352,10 @@ def validate_config() -> None:
         raise ValueError("CONFIG['mcts_cache_size'] must be >= 0.")
     if cfg_int("ddp_timeout_seconds") <= 0:
         raise ValueError("CONFIG['ddp_timeout_seconds'] must be > 0.")
+    if cfg_int("max_pending_hf_uploads") <= 0:
+        raise ValueError("CONFIG['max_pending_hf_uploads'] must be > 0.")
+    if cfg_float("hf_upload_future_timeout_s") <= 0.0:
+        raise ValueError("CONFIG['hf_upload_future_timeout_s'] must be > 0.")
 
     opp_sum = (
         cfg_float("opponent_self_prob")

tests/test_training_checkpointing.py

@@ -8,6 +8,7 @@
     drain_completed_hf_uploads,
     ensure_hf_ready,
     should_save_iteration_checkpoint,
+    wait_for_hf_uploads,
 )
 from training.config_runtime import CONFIG
 
@@ -55,6 +56,15 @@ def test_drain_completed_hf_uploads_continues_when_fail_fast_disabled(self) -> N
         remaining = drain_completed_hf_uploads([failed], fail_on_error=False)
         self.assertEqual(remaining, [])
 
+    def test_wait_for_hf_uploads_raises_timeout_in_fail_fast_mode(self) -> None:
+        pending: Future[None] = Future()
+        with self.assertRaises(RuntimeError):
+            wait_for_hf_uploads([pending], timeout_s=1.0, fail_on_error=True)
+
+    def test_wait_for_hf_uploads_tolerates_timeout_when_configured(self) -> None:
+        pending: Future[None] = Future()
+        wait_for_hf_uploads([pending], timeout_s=1.0, fail_on_error=False)
+
     def test_should_save_iteration_checkpoint_on_schedule(self) -> None:
         self.assertTrue(
             should_save_iteration_checkpoint(

train.py

@@ -30,6 +30,7 @@
     ensure_hf_ready,
     init_hf_checkpointer,
     should_save_iteration_checkpoint,
+    wait_for_hf_uploads,
 )
 from training.config_runtime import (  # noqa: E402
     CONFIG,
@@ -352,6 +353,8 @@ def main() -> None:
         for iteration in range(start_iteration + 1, iterations + 1):
             if hf_upload_futures:
                 hf_upload_futures = drain_completed_hf_uploads(hf_upload_futures, fail_on_error=cfg_bool("fail_on_hf_upload_error"))
+                if len(hf_upload_futures) > cfg_int("max_pending_hf_uploads"):
+                    raise RuntimeError("HF upload backlog is growing; aborting early.")
             epoch_pulse.set_iteration(iteration)
             selfplay_start = time.perf_counter()
             selfplay_stats = execute_self_play(
@@ -458,6 +461,8 @@ def main() -> None:
                         )
                         hf_upload_futures.append(future)
                         monitor.log_warning(iteration=iteration, message=f"HF upload queued for iteration {iteration}.")
+                        if len(hf_upload_futures) > cfg_int("max_pending_hf_uploads"):
+                            raise RuntimeError("HF upload backlog exceeded configured threshold.")
                     else:
                         hf_checkpointer.upload_checkpoint_files(
                             iteration=iteration,
@@ -482,12 +487,14 @@ def main() -> None:
             )
     finally:
         if hf_upload_executor is not None:
-            for future in hf_upload_futures:
-                try:
-                    future.result()
-                except Exception:
-                    log("A queued HF upload failed.")
-            hf_upload_executor.shutdown(wait=True)
-
+            try:
+                wait_for_hf_uploads(
+                    hf_upload_futures,
+                    timeout_s=cfg_float("hf_upload_future_timeout_s"),
+                    fail_on_error=cfg_bool("fail_on_hf_upload_error"),
+                )
+            except Exception as exc:
+                log(f"HF upload wait failed: {exc}")
+            hf_upload_executor.shutdown(wait=False, cancel_futures=True)
 if __name__ == "__main__":
     main()