code/self_healing_finetuning/training/start_policy.py

from __future__ import annotations

import json
from pathlib import Path
from typing import Any

from n21.settings import SHFT_WORKSPACE_ROOT


VALID_START_POLICIES = {"bootstrap", "continue-best"}


def _load_json(path: Path) -> dict[str, Any] | None:
    if not path.exists():
        return None
    try:
        return json.loads(path.read_text(encoding="utf-8-sig"))
    except (OSError, json.JSONDecodeError):
        return None


def resolve_training_start(
    *,
    release_id: str | None,
    model_candidate: str,
    start_policy: str = "bootstrap",
    adapter_bootstrap: bool = True,
) -> dict[str, Any]:
    policy = (start_policy or "bootstrap").strip().lower()
    if policy not in VALID_START_POLICIES:
        raise ValueError(f"invalid finetune start policy: {start_policy}; expected one of {sorted(VALID_START_POLICIES)}")

    if policy == "bootstrap":
        start_adapter = model_candidate if adapter_bootstrap else None
        return {
            "policy": "bootstrap",
            "start_adapter": start_adapter,
            "bootstrap_adapter": model_candidate,
            "continued_from_run_id": None,
            "continued_from_adapter": None,
            "source": "approved_adapter_bootstrap" if adapter_bootstrap else "fresh_base_model_lora",
            "adapter_bootstrap": adapter_bootstrap,
            "rationale": (
                "Start this role adapter from the approved bootstrap adapter."
                if adapter_bootstrap
                else "Start a fresh LoRA adapter on the selected base model."
            ),
        }

    if not release_id:
        raise ValueError("continue-best start policy requires --release-id")
    best_path = SHFT_WORKSPACE_ROOT / "best_runs" / f"{release_id}.json"
    best_report = _load_json(best_path)
    best_run = best_report.get("best_run") if isinstance(best_report, dict) else None
    if not isinstance(best_run, dict) or not best_run.get("run_id"):
        return {
            "policy": "continue-best",
            "start_adapter": model_candidate if adapter_bootstrap else None,
            "bootstrap_adapter": model_candidate,
            "continued_from_run_id": None,
            "continued_from_adapter": None,
            "source": "no_best_recorded_fallback_bootstrap" if adapter_bootstrap else "no_best_recorded_fresh_base_model_lora",
            "adapter_bootstrap": adapter_bootstrap,
            "best_run_report": str(best_path),
            "rationale": (
                "continue-best was requested, but no best measured checkpoint is recorded yet for this release; "
                + (
                    "start from the approved bootstrap adapter for the first measured run."
                    if adapter_bootstrap
                    else "start a fresh LoRA adapter on the selected base model."
                )
            ),
        }
    best_run_id = str(best_run["run_id"])
    adapter_path = f"/artifacts/runs/{best_run_id}/adapter"
    return {
        "policy": "continue-best",
        "start_adapter": adapter_path,
        "bootstrap_adapter": model_candidate,
        "continued_from_run_id": best_run_id,
        "continued_from_adapter": adapter_path,
        "source": "best_measured_checkpoint",
        "adapter_bootstrap": adapter_bootstrap,
        "best_run_report": str(best_path),
        "best_run_metrics": {
            "candidate_aggregate": best_run.get("candidate_aggregate"),
            "candidate_critical_pass_rate": best_run.get("candidate_critical_pass_rate"),
            "pairwise_win_rate": best_run.get("pairwise_win_rate"),
            "pairwise_loss_rate": best_run.get("pairwise_loss_rate"),
        },
        "rationale": "Continue from the best measured accepted adapter for this release; failed/non-best runs are ignored.",
    }