Add fail-fast guards for long-running training failures

Author: Dieg0Code

src/training/checkpointing.py

@@ -2,6 +2,7 @@
 
 import json
 import os
+from concurrent.futures import Future
 from datetime import datetime, timezone
 from pathlib import Path
 from typing import TYPE_CHECKING, Any
@@ -275,6 +276,26 @@ def ensure_hf_ready(checkpointer: HuggingFaceCheckpointer | None) -> None:
     )
 
 
+def drain_completed_hf_uploads(
+    futures: list[Future[None]],
+    *,
+    fail_on_error: bool,
+) -> list[Future[None]]:
+    """Return pending futures while surfacing completed upload failures early."""
+    pending: list[Future[None]] = []
+    for future in futures:
+        if not future.done():
+            pending.append(future)
+            continue
+        try:
+            future.result()
+        except Exception as exc:
+            if fail_on_error:
+                raise RuntimeError("HF upload future failed. Aborting early.") from exc
+            log(f"HF upload failed (continuing): {exc}")
+    return pending
+
+
 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:
@@ -286,6 +307,7 @@ def should_save_iteration_checkpoint(iteration: int, total_iterations: int, save
     "HuggingFaceCheckpointer",
     "cleanup_local_checkpoints",
     "cleanup_old_log_versions",
+    "drain_completed_hf_uploads",
     "ensure_hf_ready",
     "init_hf_checkpointer",
     "should_save_iteration_checkpoint",

src/training/config_runtime.py

@@ -73,6 +73,8 @@
     "mcts_use_amp": True,
     "mcts_cache_size": 100_000,
     "mcts_leaf_batch_size": 32,
+    "fail_on_selfplay_parallel_error": True,
+    "fail_on_hf_upload_error": True,
     "opponent_self_prob": 0.45,
     "opponent_heuristic_prob": 0.5,
     "opponent_random_prob": 0.05,
@@ -154,6 +156,8 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--eval-games", type=int, default=None)
     parser.add_argument("--eval-sims", type=int, default=None)
     parser.add_argument("--selfplay-workers", type=int, default=None)
+    parser.add_argument("--allow-selfplay-fallback", action="store_true")
+    parser.add_argument("--allow-hf-upload-errors", action="store_true")
     parser.add_argument("--warmup-games", type=int, default=None)
     parser.add_argument("--warmup-epochs", type=int, default=None)
     parser.add_argument(
@@ -256,6 +260,10 @@ def apply_cli_overrides(args: argparse.Namespace) -> None:
         CONFIG["eval_sims"] = max(8, args.eval_sims)
     if args.selfplay_workers is not None:
         CONFIG["selfplay_workers"] = max(1, args.selfplay_workers)
+    if args.allow_selfplay_fallback:
+        CONFIG["fail_on_selfplay_parallel_error"] = False
+    if args.allow_hf_upload_errors:
+        CONFIG["fail_on_hf_upload_error"] = False
     if args.warmup_games is not None:
         CONFIG["warmup_games"] = max(0, args.warmup_games)
     if args.warmup_epochs is not None:

src/training/selfplay_runtime.py

@@ -247,6 +247,15 @@ def history_to_examples(
     return examples
 
 
+def handle_parallel_selfplay_failure(exc: Exception) -> None:
+    if cfg_bool("fail_on_selfplay_parallel_error"):
+        raise RuntimeError(
+            "Process self-play failed with parallel workers. "
+            "Aborting instead of silently falling back to sequential mode.",
+        ) from exc
+    log(f"  Process self-play failed, falling back to sequential mode: {exc}")
+
+
 def execute_self_play(
     system: AtaxxZero,
     buffer: ReplayBuffer,
@@ -367,9 +376,7 @@ def execute_self_play(
             used_parallel = True
             log(f"  Self-play process workers active: {max_workers}", verbose_only=True)
         except Exception as exc:
-            log(
-                f"  Process self-play failed, falling back to sequential mode: {exc}",
-            )
+            handle_parallel_selfplay_failure(exc)
             episode_results.clear()
 
     if not used_parallel:
@@ -422,4 +429,5 @@ def execute_self_play(
 
 __all__ = [
     "execute_self_play",
+    "handle_parallel_selfplay_failure",
 ]

tests/test_training_checkpointing.py

@@ -1,9 +1,11 @@
 from __future__ import annotations
 
 import unittest
+from concurrent.futures import Future
 
 from training.checkpointing import (
     HuggingFaceCheckpointer,
+    drain_completed_hf_uploads,
     ensure_hf_ready,
     should_save_iteration_checkpoint,
 )
@@ -34,6 +36,25 @@ def test_ensure_hf_ready_noop_when_hf_disabled(self) -> None:
         CONFIG["hf_enabled"] = False
         ensure_hf_ready(None)
 
+    def test_drain_completed_hf_uploads_keeps_pending_only(self) -> None:
+        done: Future[None] = Future()
+        done.set_result(None)
+        pending: Future[None] = Future()
+        remaining = drain_completed_hf_uploads([done, pending], fail_on_error=True)
+        self.assertEqual(remaining, [pending])
+
+    def test_drain_completed_hf_uploads_raises_when_fail_fast_enabled(self) -> None:
+        failed: Future[None] = Future()
+        failed.set_exception(RuntimeError("upload failed"))
+        with self.assertRaises(RuntimeError):
+            drain_completed_hf_uploads([failed], fail_on_error=True)
+
+    def test_drain_completed_hf_uploads_continues_when_fail_fast_disabled(self) -> None:
+        failed: Future[None] = Future()
+        failed.set_exception(RuntimeError("upload failed"))
+        remaining = drain_completed_hf_uploads([failed], fail_on_error=False)
+        self.assertEqual(remaining, [])
+
     def test_should_save_iteration_checkpoint_on_schedule(self) -> None:
         self.assertTrue(
             should_save_iteration_checkpoint(

tests/test_training_selfplay_runtime.py

@@ -0,0 +1,28 @@
+from __future__ import annotations
+
+import unittest
+
+from training.config_runtime import CONFIG
+from training.selfplay_runtime import handle_parallel_selfplay_failure
+
+
+class TestTrainingSelfplayRuntime(unittest.TestCase):
+    def setUp(self) -> None:
+        self._backup = dict(CONFIG)
+
+    def tearDown(self) -> None:
+        CONFIG.clear()
+        CONFIG.update(self._backup)
+
+    def test_handle_parallel_selfplay_failure_raises_in_fail_fast_mode(self) -> None:
+        CONFIG["fail_on_selfplay_parallel_error"] = True
+        with self.assertRaises(RuntimeError):
+            handle_parallel_selfplay_failure(RuntimeError("pool broke"))
+
+    def test_handle_parallel_selfplay_failure_allows_fallback_when_configured(self) -> None:
+        CONFIG["fail_on_selfplay_parallel_error"] = False
+        handle_parallel_selfplay_failure(RuntimeError("pool broke"))
+
+
+if __name__ == "__main__":
+    unittest.main()

train.py

@@ -26,6 +26,7 @@
 from training.checkpointing import (  # noqa: E402
     cleanup_local_checkpoints,
     cleanup_old_log_versions,
+    drain_completed_hf_uploads,
     ensure_hf_ready,
     init_hf_checkpointer,
     should_save_iteration_checkpoint,
@@ -349,6 +350,8 @@ def main() -> None:
 
     try:
         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"))
             epoch_pulse.set_iteration(iteration)
             selfplay_start = time.perf_counter()
             selfplay_stats = execute_self_play(
@@ -358,6 +361,8 @@ def main() -> None:
                 device=device,
             )
             selfplay_s = time.perf_counter() - selfplay_start
+            if len(buffer) == 0:
+                raise RuntimeError("Replay buffer is empty after self-play; aborting early.")
 
             train_loader = _build_train_loader(buffer, device=device)
             val_loader = _build_val_loader(buffer, device=device)
@@ -409,10 +414,7 @@ def main() -> None:
                         best_path = checkpoint_dir / "best_eval.ckpt"
                         trainer.save_checkpoint(str(best_path))
                 except Exception as exc:
-                    monitor.log_warning(
-                        iteration=iteration,
-                        message=f"eval failed, continuing training: {exc}",
-                    )
+                    monitor.log_warning(iteration=iteration, message=f"eval failed, continuing training: {exc}")
 
             if not should_save_iteration_checkpoint(
                 iteration=iteration,
@@ -430,10 +432,7 @@ def main() -> None:
                     keep_last_n=cfg_int("keep_last_n_local_checkpoints"),
                 )
             except OSError:
-                monitor.log_warning(
-                    iteration=iteration,
-                    message="local checkpoint save failed.",
-                )
+                monitor.log_warning(iteration=iteration, message="local checkpoint save failed.")
 
             if hf_checkpointer is not None:
                 try:
@@ -458,10 +457,7 @@ def main() -> None:
                             keep_last_n=cfg_int("keep_last_n_hf_checkpoints"),
                         )
                         hf_upload_futures.append(future)
-                        monitor.log_warning(
-                            iteration=iteration,
-                            message=f"HF upload queued for iteration {iteration}.",
-                        )
+                        monitor.log_warning(iteration=iteration, message=f"HF upload queued for iteration {iteration}.")
                     else:
                         hf_checkpointer.upload_checkpoint_files(
                             iteration=iteration,
@@ -470,10 +466,7 @@ def main() -> None:
                             metadata_path=metadata_path,
                             keep_last_n=cfg_int("keep_last_n_hf_checkpoints"),
                         )
-                        monitor.log_warning(
-                            iteration=iteration,
-                            message=f"HF checkpoint uploaded for iteration {iteration}.",
-                        )
+                        monitor.log_warning(iteration=iteration, message=f"HF checkpoint uploaded for iteration {iteration}.")
                 except (OSError, ValueError):
                     monitor.log_warning(iteration=iteration, message="HF upload failed for this iteration.")
             if cfg_bool("export_onnx"):