training/pipeline.py

"""
PERMANENCE training pipeline orchestrator.

Runs the four stages in order, passing artifacts between them. Each stage
can also be invoked in isolation via ``python -m training.stages.stage_N_*``.

Usage:
    python -m training.pipeline                        # full pipeline
    python -m training.pipeline --from gate            # skip SFT
    python -m training.pipeline --only sft             # SFT alone
    python -m training.pipeline --config my.yaml       # custom config

Exit codes:
    0 — all requested stages passed
    2 — a stage failed (status.ok=false)
    3 — fatal error (exception)

Stage outputs live under ``training/artifacts/<stage>/`` so you can inspect
status.json after any stage and decide whether to proceed.
"""
from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path
from typing import Callable, Dict, List, Tuple


_ROOT = Path(__file__).resolve().parent.parent
if str(_ROOT) not in sys.path:
    sys.path.insert(0, str(_ROOT))

from training.config import TrainingConfig, load_simple_yaml  # noqa: E402


STAGES: List[str] = ["sft", "gate", "grpo", "eval"]

ARTIFACTS_ROOT = _ROOT / "training" / "artifacts"


def _run_stage(
    name: str,
    config: TrainingConfig,
) -> Tuple[bool, Dict[str, object]]:
    """Import and invoke a stage's ``run_*`` function. Returns (ok, status)."""
    if name == "sft":
        from training.stages.stage_1_sft import run_sft
        status = run_sft(config)
    elif name == "gate":
        from training.stages.stage_2_gate import run_gate
        status = run_gate(config)
    elif name == "grpo":
        from training.stages.stage_3_grpo import run_grpo
        status = run_grpo(config)
    elif name == "eval":
        from training.stages.stage_4_eval import run_eval
        status = run_eval(config)
    else:
        raise ValueError(f"unknown stage: {name}")
    return bool(status.get("ok", False)), status


def run_pipeline(
    config: TrainingConfig,
    stages_to_run: List[str],
    bail_on_failure: bool = True,
) -> Dict[str, object]:
    """Run the requested stages in order. Returns a summary dict."""
    summary: Dict[str, object] = {"config_model": config.model_name, "stages": {}}
    for s in stages_to_run:
        print(f"\n━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━")
        print(f"▶ STAGE: {s}")
        print("━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━")
        try:
            ok, status = _run_stage(s, config)
        except Exception as exc:
            print(f"✗ Stage {s} raised: {exc}")
            summary["stages"][s] = {"ok": False, "error": str(exc)[:500]}
            if bail_on_failure:
                summary["final_status"] = "fatal"
                return summary
            continue
        summary["stages"][s] = status
        print(f"{'✓' if ok else '✗'} Stage {s}: {json.dumps(status, indent=2, default=str)}")
        if not ok and bail_on_failure:
            summary["final_status"] = f"failed_at_{s}"
            return summary
    summary["final_status"] = "completed"
    return summary


def main() -> int:
    parser = argparse.ArgumentParser(description="PERMANENCE training pipeline")
    parser.add_argument("--config", default=str(_ROOT / "training" / "config.yaml"))
    parser.add_argument(
        "--from",
        dest="from_stage",
        choices=STAGES,
        help="Start from this stage (skip earlier stages; assumes their artifacts exist)",
    )
    parser.add_argument(
        "--only",
        dest="only_stage",
        choices=STAGES,
        help="Run only this stage and exit",
    )
    parser.add_argument(
        "--no-bail",
        action="store_true",
        help="Continue through stages even if one fails (for post-mortem)",
    )
    args = parser.parse_args()

    cfg_map = load_simple_yaml(args.config)
    cfg = TrainingConfig.from_mapping(cfg_map)

    if args.only_stage:
        stages_to_run = [args.only_stage]
    elif args.from_stage:
        start_idx = STAGES.index(args.from_stage)
        stages_to_run = STAGES[start_idx:]
    else:
        stages_to_run = list(STAGES)

    ARTIFACTS_ROOT.mkdir(parents=True, exist_ok=True)
    summary = run_pipeline(cfg, stages_to_run, bail_on_failure=not args.no_bail)
    (ARTIFACTS_ROOT / "pipeline_summary.json").write_text(json.dumps(summary, indent=2, default=str))
    final = summary.get("final_status", "unknown")
    print(f"\n═══ PIPELINE {str(final).upper()} ═══")
    print(f"Summary → {ARTIFACTS_ROOT}/pipeline_summary.json")
    return 0 if final == "completed" else 2


if __name__ == "__main__":
    sys.exit(main())