Make self-play training resilient to HF Space restarts

The L40S full run was getting killed at ~step 84/120 of the generator
phase because uvicorn (PID 1) caught a SIGTERM from the Space and the
backgrounded training subshell died with the container, with no
recovery path. The answerer reward graph was empty because training
never reached the answerer phase.

scripts/space_start.sh: invert the lifecycle. uvicorn now runs in the
background (so HF healthchecks pass) and self-play training runs in
the foreground as the script's blocking child. A SIGTERM from the
platform now propagates to the trainer instead of cleanly tearing down
PID 1, and the EXIT trap stops uvicorn afterwards.

src/osint_env/training/self_play.py: _train_grpo_phase now resumes
work on every restart. It (1) skips the phase entirely if final_model
is already on disk, (2) auto-detects the latest local checkpoint-* and
calls trainer.train(resume_from_checkpoint=...) when full state
(optimizer.pt + trainer_state.json) is present, (3) falls back to
downloading the latest checkpoint for this phase from
OSINT_HF_CHECKPOINT_REPO_ID via huggingface_hub.snapshot_download and
warm-starts from those weights when only inference state is available.

config/self_play_training_hf_l40s_full.json: drop logging_steps from 5
to 1 (matches the smoke run so the W&B reward graph is dense for both
GeneratorRewardFunction and AnswererRewardFunction) and drop save_steps
from 120 to 30 so a mid-phase restart loses at most ~30 steps.

Made-with: Cursor

config/self_play_training_hf_l40s_full.json

@@ -40,8 +40,8 @@
     "num_iterations": 1,
     "loss_type": "dapo",
     "scale_rewards": "group",
-    "logging_steps": 5,
-    "save_steps": 120,
+    "logging_steps": 1,
+    "save_steps": 30,
     "save_total_limit": 1,
     "optim": "adamw_torch_fused",
     "bf16": true,
@@ -72,8 +72,8 @@
     "num_iterations": 1,
     "loss_type": "dapo",
     "scale_rewards": "group",
-    "logging_steps": 5,
-    "save_steps": 120,
+    "logging_steps": 1,
+    "save_steps": 30,
     "save_total_limit": 1,
     "optim": "adamw_torch_fused",
     "bf16": true,

scripts/space_start.sh

@@ -13,7 +13,37 @@ TRAIN_CONFIG_PATH="${TRAIN_SELF_PLAY_CONFIG_PATH:-config/self_play_training_hf_l
 TRAIN_OUTPUT_DIR="${TRAIN_SELF_PLAY_OUTPUT_DIR:-}"
 RUN_FLAG="${RUN_SELF_PLAY_TRAINING:-1}"
 DRY_RUN_FLAG="${RUN_SELF_PLAY_DRY_RUN:-0}"
-BACKGROUND_FLAG="${RUN_SELF_PLAY_BACKGROUND:-1}"
+SERVE_API_FLAG="${RUN_SPACE_API_SERVER:-1}"
+PORT_VALUE="${PORT:-7860}"
+UVICORN_LOG_PATH="${UVICORN_LOG_PATH:-/tmp/uvicorn.log}"
+
+UVICORN_PID=""
+
+_start_api_server_background() {
+  if ! _is_true "$SERVE_API_FLAG"; then
+    echo "[space_start] RUN_SPACE_API_SERVER disabled. Skipping API server."
+    return
+  fi
+  echo "[space_start] Starting API server in background on port ${PORT_VALUE} (logs: ${UVICORN_LOG_PATH})."
+  # API server runs in background ONLY for HF healthchecks. Training is the
+  # primary process. If HF infrastructure SIGTERMs the container we still
+  # want training to receive the signal and flush a final checkpoint, not
+  # to silently die because PID 1 (uvicorn previously) exited first.
+  uvicorn server:app --host 0.0.0.0 --port "${PORT_VALUE}" \
+    >"${UVICORN_LOG_PATH}" 2>&1 &
+  UVICORN_PID=$!
+  echo "[space_start] uvicorn pid=${UVICORN_PID}"
+}
+
+_stop_api_server() {
+  if [ -n "${UVICORN_PID}" ] && kill -0 "${UVICORN_PID}" 2>/dev/null; then
+    echo "[space_start] Stopping uvicorn pid=${UVICORN_PID}."
+    kill "${UVICORN_PID}" 2>/dev/null || true
+    wait "${UVICORN_PID}" 2>/dev/null || true
+  fi
+}
+
+trap '_stop_api_server' EXIT INT TERM
 
 _train_self_play() {
   if [ -n "${TRAIN_OUTPUT_DIR}" ]; then
@@ -27,7 +57,7 @@ _train_self_play() {
     # shellcheck disable=SC2086
     osint-env train-self-play --config "${ENV_CONFIG_PATH}" --train-config "${TRAIN_CONFIG_PATH}" ${OUTPUT_ARG} --dry-run
   else
-    echo "[space_start] Running self-play training."
+    echo "[space_start] Running self-play training (foreground)."
     # shellcheck disable=SC2086
     osint-env train-self-play --config "${ENV_CONFIG_PATH}" --train-config "${TRAIN_CONFIG_PATH}" ${OUTPUT_ARG}
   fi
@@ -47,15 +77,20 @@ if _is_true "$RUN_FLAG"; then
   if [ -n "${OSINT_HF_CHECKPOINT_REPO_ID:-}" ]; then
     echo "[space_start] HF checkpoint repo: ${OSINT_HF_CHECKPOINT_REPO_ID}"
   fi
-  if _is_true "$BACKGROUND_FLAG"; then
-    echo "[space_start] Launching self-play in background so the Space API can stay online."
-    _train_self_play &
-  else
-    _train_self_play
+  _start_api_server_background
+  # Run training in the FOREGROUND so the script (and therefore PID 1)
+  # blocks until training is finished. A graceful SIGTERM from HF will
+  # propagate to the training process via the shell's signal handling
+  # and the trap above will cleanly stop uvicorn afterwards.
+  _train_self_play
+  echo "[space_start] Training finished. Keeping API server alive for log inspection."
+  if [ -n "${UVICORN_PID}" ] && kill -0 "${UVICORN_PID}" 2>/dev/null; then
+    wait "${UVICORN_PID}"
   fi
 else
   echo "[space_start] RUN_SELF_PLAY_TRAINING disabled. Skipping self-play run."
+  _start_api_server_background
+  if [ -n "${UVICORN_PID}" ] && kill -0 "${UVICORN_PID}" 2>/dev/null; then
+    wait "${UVICORN_PID}"
+  fi
 fi
-
-echo "[space_start] Starting API server."
-exec uvicorn server:app --host 0.0.0.0 --port "${PORT:-7860}"

src/osint_env/training/self_play.py

@@ -179,6 +179,120 @@ def _require_training_stack() -> tuple[Any, Any, Any]:
     return Dataset, GRPOConfig, GRPOTrainer
 
 
+def _latest_local_checkpoint(output_dir: Path) -> Path | None:
+    if not output_dir.exists():
+        return None
+    candidates: list[tuple[int, Path]] = []
+    for path in output_dir.glob("checkpoint-*"):
+        if not path.is_dir():
+            continue
+        suffix = path.name.split("-", 1)[-1]
+        try:
+            step = int(suffix)
+        except ValueError:
+            continue
+        candidates.append((step, path))
+    if not candidates:
+        return None
+    candidates.sort(key=lambda item: item[0])
+    return candidates[-1][1]
+
+
+def _final_model_already_present(output_dir: Path) -> bool:
+    final_dir = output_dir / "final_model"
+    if not final_dir.is_dir():
+        return False
+    safetensors = list(final_dir.glob("*.safetensors"))
+    legacy_bin = list(final_dir.glob("pytorch_model*.bin"))
+    return bool(safetensors or legacy_bin)
+
+
+def _maybe_download_phase_checkpoints_from_hf(output_dir: Path, run_dir: Path) -> Path | None:
+    """If no local checkpoint exists, try to recover the latest checkpoint
+    for this phase from the HF Hub repo we already upload to. Returns the
+    local path of the restored ``checkpoint-*`` directory, or None.
+
+    Designed to make Space restarts non-destructive: training state is
+    pushed to ``OSINT_HF_CHECKPOINT_REPO_ID`` after every phase, so on a
+    fresh container we can pull it back and resume.
+    """
+    if _latest_local_checkpoint(output_dir) is not None:
+        return _latest_local_checkpoint(output_dir)
+    repo_id = _default_hf_checkpoint_repo_id(run_dir)
+    token = _resolve_hf_upload_token()
+    if not repo_id or not token:
+        return None
+
+    try:
+        from huggingface_hub import HfApi, snapshot_download
+    except ImportError:
+        return None
+
+    repo_type = str(os.getenv("OSINT_HF_CHECKPOINT_REPO_TYPE", "model")).strip() or "model"
+    api = HfApi(token=token)
+    try:
+        files = list(api.list_repo_files(repo_id=repo_id, repo_type=repo_type))
+    except Exception as exc:  # noqa: BLE001
+        print(
+            f"[self_play][resume] could not list files in {repo_id}: "
+            f"{type(exc).__name__}: {exc}",
+            flush=True,
+        )
+        return None
+
+    phase_prefix = _hf_relative_repo_path(output_dir, run_dir)
+    phase_prefix_clean = phase_prefix.strip("/") + "/"
+    candidate_steps: dict[int, list[str]] = {}
+    for remote_path in files:
+        if not remote_path.startswith(phase_prefix_clean):
+            continue
+        relative = remote_path[len(phase_prefix_clean) :]
+        parts = relative.split("/", 1)
+        if len(parts) < 2 or not parts[0].startswith("checkpoint-"):
+            continue
+        try:
+            step = int(parts[0].split("-", 1)[1])
+        except ValueError:
+            continue
+        candidate_steps.setdefault(step, []).append(remote_path)
+
+    if not candidate_steps:
+        return None
+
+    best_step = max(candidate_steps.keys())
+    target_local_dir = output_dir / f"checkpoint-{best_step}"
+    target_local_dir.mkdir(parents=True, exist_ok=True)
+    print(
+        f"[self_play][resume] downloading phase checkpoint from HF Hub: "
+        f"repo={repo_id} prefix={phase_prefix_clean}checkpoint-{best_step} "
+        f"-> {target_local_dir}",
+        flush=True,
+    )
+    try:
+        snapshot_download(
+            repo_id=repo_id,
+            repo_type=repo_type,
+            local_dir=str(output_dir),
+            allow_patterns=[f"{phase_prefix_clean}checkpoint-{best_step}/*"],
+            token=token,
+        )
+        downloaded_root = output_dir / phase_prefix_clean.rstrip("/") / f"checkpoint-{best_step}"
+        if downloaded_root.exists() and downloaded_root != target_local_dir:
+            for item in downloaded_root.iterdir():
+                dest = target_local_dir / item.name
+                if dest.exists():
+                    continue
+                item.replace(dest)
+        return target_local_dir if target_local_dir.exists() else None
+    except Exception as exc:  # noqa: BLE001
+        print(
+            f"[self_play][resume] failed to download checkpoint from HF Hub: "
+            f"{type(exc).__name__}: {exc}",
+            flush=True,
+        )
+        return None
+
+
 
 def _task_to_edge_json(task: TaskInstance) -> str:
     payload = [
@@ -731,6 +845,63 @@ def _train_grpo_phase(
     Dataset, GRPOConfig, GRPOTrainer = _require_training_stack()
 
     output_dir.mkdir(parents=True, exist_ok=True)
+
+    phase_label = str(run_name).strip() or str(output_dir.name)
+    reward_class_name = type(reward_function).__name__
+
+    # Output layout: <run_dir>/round_NNN/<phase_subdir>. Match the run_dir
+    # used by the corresponding HF upload helpers so resume paths line up
+    # with where checkpoints were written.
+    run_dir_for_resume = output_dir.parents[1] if len(output_dir.parents) >= 2 else output_dir.parent
+
+    if _final_model_already_present(output_dir):
+        final_dir = output_dir / "final_model"
+        print(
+            f"[self_play][resume] phase={phase_label} already has final_model at {final_dir}. "
+            f"Skipping retrain on Space restart.",
+            flush=True,
+        )
+        checkpoint_dirs = [str(path) for path in sorted(output_dir.glob("checkpoint-*")) if path.is_dir()]
+        return {
+            "model_path": str(final_dir),
+            "final_model_path": str(final_dir),
+            "phase_output_dir": str(output_dir),
+            "checkpoint_dirs": checkpoint_dirs,
+            "global_step": int(getattr(phase, "max_steps", 0) or 0),
+            "training_loss": 0.0,
+            "train_rows": len(rows),
+            "tuning_mode": str(tuning_mode).strip().lower() or "full",
+            "is_full_finetune": str(tuning_mode).strip().lower() != "lora",
+            "resumed_skipped": True,
+        }
+
+    resume_checkpoint = _latest_local_checkpoint(output_dir)
+    resume_is_full_state = bool(
+        resume_checkpoint is not None
+        and (resume_checkpoint / "optimizer.pt").exists()
+        and (resume_checkpoint / "trainer_state.json").exists()
+    )
+    if resume_checkpoint is None:
+        downloaded = _maybe_download_phase_checkpoints_from_hf(output_dir, run_dir_for_resume)
+        if downloaded is not None:
+            resume_checkpoint = downloaded
+            resume_is_full_state = bool(
+                (downloaded / "optimizer.pt").exists()
+                and (downloaded / "trainer_state.json").exists()
+            )
+
+    # If the resume checkpoint is weights-only (e.g. recovered from HF Hub
+    # which intentionally drops optimizer state to keep uploads small),
+    # warm-start the model from those weights and start a fresh trainer
+    # state. Better than restarting from the base model.
+    warm_start_from: Path | None = None
+    if resume_checkpoint is not None and not resume_is_full_state:
+        warm_start_from = resume_checkpoint
+        resume_checkpoint = None
+
+    if warm_start_from is not None:
+        model_name_or_path = str(warm_start_from)
+
     dataset = Dataset.from_list(rows)
     args = _safe_build_grpo_config(
         phase=phase,
@@ -753,8 +924,6 @@ def _train_grpo_phase(
             raise RuntimeError("Installed TRL version does not expose peft_config in GRPOTrainer.")
         trainer_kwargs["peft_config"] = _build_lora_config(lora)
 
-    phase_label = str(run_name).strip() or str(output_dir.name)
-    reward_class_name = type(reward_function).__name__
     print(
         f"[self_play] Starting phase: {phase_label} rows={len(rows)} "
         f"max_steps={phase.max_steps}",
@@ -769,6 +938,17 @@ def _train_grpo_phase(
         f"per_device_train_batch_size={phase.per_device_train_batch_size}",
         flush=True,
     )
+    if resume_checkpoint is not None:
+        print(
+            f"[self_play][resume] phase={phase_label} resuming (full state) from checkpoint={resume_checkpoint}",
+            flush=True,
+        )
+    elif warm_start_from is not None:
+        print(
+            f"[self_play][resume] phase={phase_label} warm-starting from weights only at {warm_start_from} "
+            f"(no optimizer state available)",
+            flush=True,
+        )
     strict_asserts = str(os.getenv("OSINT_TRAIN_STRICT_ASSERTS", "")).strip().lower() in {"1", "true", "yes", "on"}
     trainer = GRPOTrainer(**trainer_kwargs)
     tracked_params = [
@@ -780,7 +960,10 @@ def _train_grpo_phase(
         name: float(param.detach().float().abs().mean().item())
         for name, param in tracked_params
     }
-    train_output = trainer.train()
+    if resume_checkpoint is not None:
+        train_output = trainer.train(resume_from_checkpoint=str(resume_checkpoint))
+    else:
+        train_output = trainer.train()
 
     final_dir = output_dir / "final_model"
     trainer.save_model(str(final_dir))