Fail fast DDP startup and HF checkpoint readiness

Author: Dieg0Code

src/training/checkpointing.py

@@ -258,6 +258,23 @@ def init_hf_checkpointer() -> HuggingFaceCheckpointer | None:
     )
 
 
+def ensure_hf_ready(checkpointer: HuggingFaceCheckpointer | None) -> None:
+    """Fail fast when HF checkpointing was requested but cannot be initialized."""
+    if not cfg_bool("hf_enabled"):
+        return
+    if checkpointer is not None:
+        log(
+            "HF checkpointing enabled: "
+            f"repo={cfg_str('hf_repo_id').strip()} run_id={cfg_str('hf_run_id').strip()}",
+        )
+        return
+    token_env = cfg_str("hf_token_env")
+    raise RuntimeError(
+        "HF checkpointing requested (--hf) but initialization failed. "
+        f"Check token env '{token_env}', hf_repo_id, and hf_run_id before starting.",
+    )
+
+
 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:
@@ -269,6 +286,7 @@ def should_save_iteration_checkpoint(iteration: int, total_iterations: int, save
     "HuggingFaceCheckpointer",
     "cleanup_local_checkpoints",
     "cleanup_old_log_versions",
+    "ensure_hf_ready",
     "init_hf_checkpointer",
     "should_save_iteration_checkpoint",
 ]

src/training/config_runtime.py

@@ -69,6 +69,7 @@
     "trainer_strategy": "auto",
     "trainer_precision": "bf16-mixed",
     "trainer_benchmark": True,
+    "ddp_timeout_seconds": 180,
     "mcts_use_amp": True,
     "mcts_cache_size": 100_000,
     "mcts_leaf_batch_size": 32,
@@ -121,6 +122,7 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--keep-log-versions", type=int, default=None)
     parser.add_argument("--devices", type=int, default=None)
     parser.add_argument("--strategy", default=None)
+    parser.add_argument("--ddp-timeout-s", type=int, default=None)
     parser.add_argument(
         "--precision",
         choices=["16-mixed", "bf16-mixed", "32-true"],
@@ -206,6 +208,8 @@ def apply_cli_overrides(args: argparse.Namespace) -> None:
         CONFIG["trainer_devices"] = max(1, args.devices)
     if args.strategy is not None:
         CONFIG["trainer_strategy"] = args.strategy
+    if args.ddp_timeout_s is not None:
+        CONFIG["ddp_timeout_seconds"] = max(30, args.ddp_timeout_s)
     if args.precision is not None:
         CONFIG["trainer_precision"] = args.precision
     if args.num_workers is not None:
@@ -330,6 +334,8 @@ def validate_config() -> None:
         raise ValueError("CONFIG['value_loss_coeff'] must be >= 0.")
     if cfg_int("mcts_cache_size") < 0:
         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.")
 
     opp_sum = (
         cfg_float("opponent_self_prob")

src/training/trainer_runtime.py

@@ -1,10 +1,13 @@
 from __future__ import annotations
 
+from datetime import timedelta
+
 import pytorch_lightning as pl
 import torch
 from pytorch_lightning import Callback
 from pytorch_lightning.callbacks import LearningRateMonitor, ModelCheckpoint
 from pytorch_lightning.loggers import TensorBoardLogger
+from pytorch_lightning.strategies import DDPStrategy
 
 from training.config_runtime import (
     TrainerPrecision,
@@ -43,6 +46,15 @@ def is_ddp_rendezvous_timeout(exc: BaseException) -> bool:
     )
 
 
+def resolve_trainer_strategy(strategy: str) -> str | DDPStrategy:
+    timeout = timedelta(seconds=max(30, cfg_int("ddp_timeout_seconds")))
+    if strategy == "ddp":
+        return DDPStrategy(timeout=timeout, start_method="popen")
+    if strategy == "ddp_spawn":
+        return DDPStrategy(timeout=timeout, start_method="spawn")
+    return strategy
+
+
 def build_trainer(
     *,
     epochs: int,
@@ -59,11 +71,12 @@ def build_trainer(
     callbacks: list[Callback] = [checkpoint_callback, lr_monitor]
     if extra_callbacks is not None:
         callbacks.extend(extra_callbacks)
+    resolved_strategy = resolve_trainer_strategy(strategy)
     return pl.Trainer(
         max_epochs=epochs,
         accelerator=accelerator,
         devices=devices,
-        strategy=strategy,
+        strategy=resolved_strategy,
         precision=precision,
         benchmark=benchmark,
         callbacks=callbacks,
@@ -107,4 +120,5 @@ def export_onnx(model: torch.nn.Module, path: str, device: str) -> None:
     "is_ddp_rendezvous_timeout",
     "resolve_trainer_hw",
     "resolve_trainer_precision",
+    "resolve_trainer_strategy",
 ]

tests/test_training_checkpointing.py

@@ -4,17 +4,36 @@
 
 from training.checkpointing import (
     HuggingFaceCheckpointer,
+    ensure_hf_ready,
     should_save_iteration_checkpoint,
 )
+from training.config_runtime import CONFIG
 
 
 class TestTrainingCheckpointing(unittest.TestCase):
+    def setUp(self) -> None:
+        self._backup = dict(CONFIG)
+
+    def tearDown(self) -> None:
+        CONFIG.clear()
+        CONFIG.update(self._backup)
+
     def test_repo_path_is_namespaced_by_run_id(self) -> None:
         checkpointer = object.__new__(HuggingFaceCheckpointer)
         checkpointer.run_id = "policy_spatial_v1"
         repo_path = checkpointer._repo_path("model_iter_040.pt")
         self.assertEqual(repo_path, "runs/policy_spatial_v1/model_iter_040.pt")
 
+    def test_ensure_hf_ready_raises_when_hf_enabled_without_checkpointer(self) -> None:
+        CONFIG["hf_enabled"] = True
+        CONFIG["hf_token_env"] = "HF_TOKEN"  # noqa: S105 - test fixture value, not a secret.
+        with self.assertRaises(RuntimeError):
+            ensure_hf_ready(None)
+
+    def test_ensure_hf_ready_noop_when_hf_disabled(self) -> None:
+        CONFIG["hf_enabled"] = False
+        ensure_hf_ready(None)
+
     def test_should_save_iteration_checkpoint_on_schedule(self) -> None:
         self.assertTrue(
             should_save_iteration_checkpoint(

tests/test_training_trainer_runtime.py

@@ -0,0 +1,42 @@
+from __future__ import annotations
+
+import unittest
+from datetime import timedelta
+
+from pytorch_lightning.strategies import DDPStrategy
+
+from training.config_runtime import CONFIG
+from training.trainer_runtime import resolve_trainer_strategy
+
+
+class TestTrainingTrainerRuntime(unittest.TestCase):
+    def setUp(self) -> None:
+        self._backup = dict(CONFIG)
+
+    def tearDown(self) -> None:
+        CONFIG.clear()
+        CONFIG.update(self._backup)
+
+    def test_resolve_trainer_strategy_ddp_uses_configured_timeout(self) -> None:
+        CONFIG["ddp_timeout_seconds"] = 75
+        resolved = resolve_trainer_strategy("ddp")
+        self.assertIsInstance(resolved, DDPStrategy)
+        strategy = resolved
+        self.assertEqual(strategy._timeout, timedelta(seconds=75))
+        self.assertEqual(strategy._start_method, "popen")
+
+    def test_resolve_trainer_strategy_ddp_spawn_uses_spawn_start_method(self) -> None:
+        CONFIG["ddp_timeout_seconds"] = 90
+        resolved = resolve_trainer_strategy("ddp_spawn")
+        self.assertIsInstance(resolved, DDPStrategy)
+        strategy = resolved
+        self.assertEqual(strategy._timeout, timedelta(seconds=90))
+        self.assertEqual(strategy._start_method, "spawn")
+
+    def test_resolve_trainer_strategy_passthrough_for_auto(self) -> None:
+        resolved = resolve_trainer_strategy("auto")
+        self.assertEqual(resolved, "auto")
+
+
+if __name__ == "__main__":
+    unittest.main()

train.py

@@ -26,6 +26,7 @@
 from training.checkpointing import (  # noqa: E402
     cleanup_local_checkpoints,
     cleanup_old_log_versions,
+    ensure_hf_ready,
     init_hf_checkpointer,
     should_save_iteration_checkpoint,
 )
@@ -198,11 +199,11 @@ def _run_warmup_if_needed(
     optimizer_transfer: OptimizerStateTransfer,
     monitor: TrainingMonitor,
     epoch_pulse: EpochPulseCallback,
-) -> None:
+) -> tuple[str, int, str, TrainerPrecision]:
     warmup_games = cfg_int("warmup_games")
     warmup_epochs = cfg_int("warmup_epochs")
     if start_iteration != 0 or warmup_games <= 0 or warmup_epochs <= 0:
-        return
+        return trainer_accelerator, trainer_devices, trainer_strategy, trainer_precision
 
     # Warmup seeds the policy with legal, sensible moves before self-play noise.
     warmup_rng = torch.Generator().manual_seed(cfg_int("seed"))
@@ -216,24 +217,28 @@ def _run_warmup_if_needed(
     monitor.log_warmup(examples=len(warmup_examples), games=warmup_games)
     train_loader = _build_train_loader(buffer, device=device)
     val_loader = _build_val_loader(buffer, device=device)
-    warmup_trainer = build_trainer(
+    (
+        _warmup_trainer,
+        trainer_accelerator,
+        trainer_devices,
+        trainer_strategy,
+        trainer_precision,
+    ) = _fit_with_ddp_fallback(
+        system=system,
+        train_loader=train_loader,
+        val_loader=val_loader,
         epochs=warmup_epochs,
-        accelerator=trainer_accelerator,
-        devices=trainer_devices,
-        strategy=trainer_strategy,
-        precision=trainer_precision,
-        benchmark=cfg_bool("trainer_benchmark"),
+        trainer_accelerator=trainer_accelerator,
+        trainer_devices=trainer_devices,
+        trainer_strategy=trainer_strategy,
+        trainer_precision=trainer_precision,
         checkpoint_callback=checkpoint_callback,
         lr_monitor=lr_monitor,
         logger=logger,
-        extra_callbacks=[optimizer_transfer, epoch_pulse],
-    )
-    system.train()
-    warmup_trainer.fit(
-        model=system,
-        train_dataloaders=train_loader,
-        val_dataloaders=val_loader,
+        optimizer_transfer=optimizer_transfer,
+        epoch_pulse=epoch_pulse,
     )
+    return trainer_accelerator, trainer_devices, trainer_strategy, trainer_precision
 
 
 def main() -> None:
@@ -286,6 +291,7 @@ def main() -> None:
     buffer = ReplayBuffer(capacity=cfg_int("buffer_size"))
 
     hf_checkpointer = init_hf_checkpointer()
+    ensure_hf_ready(hf_checkpointer)
     hf_upload_executor: ThreadPoolExecutor | None = None
     hf_upload_futures: list[Future[None]] = []
     if hf_checkpointer is not None:
@@ -324,7 +330,7 @@ def main() -> None:
         pulse_every=cfg_int("epoch_pulse_every"),
     )
 
-    _run_warmup_if_needed(
+    trainer_accelerator, trainer_devices, trainer_strategy, trainer_precision = _run_warmup_if_needed(
         start_iteration=start_iteration,
         system=system,
         buffer=buffer,
@@ -469,18 +475,12 @@ def main() -> None:
                             message=f"HF checkpoint uploaded for iteration {iteration}.",
                         )
                 except (OSError, ValueError):
-                    monitor.log_warning(
-                        iteration=iteration,
-                        message="HF upload failed for this iteration.",
-                    )
+                    monitor.log_warning(iteration=iteration, message="HF upload failed for this iteration.")
             if cfg_bool("export_onnx"):
                 try:
                     export_onnx(system.model, cfg_str("onnx_path"), device=device)
                 except (OSError, RuntimeError, ValueError):
-                    monitor.log_warning(
-                        iteration=iteration,
-                        message="ONNX export failed for this iteration.",
-                    )
+                    monitor.log_warning(iteration=iteration, message="ONNX export failed for this iteration.")
 
             cleanup_old_log_versions(
                 log_dir=log_dir,