scripts/run_math_ink_06_p_formula_release.py

"""P Formula 3-seed 학습부터 single-student export까지 실패 폐쇄로 오케스트레이션한다."""

from __future__ import annotations

import argparse
from datetime import datetime, timezone
import json
from pathlib import Path
import subprocess
import sys
from typing import Any, Sequence

PROJECT_ROOT = Path(__file__).parents[1]
SOURCE_ROOT = PROJECT_ROOT / "src"
for path in (PROJECT_ROOT, SOURCE_ROOT):
    if str(path) not in sys.path:
        sys.path.insert(0, str(path))

from math_grid_drawer.research.external_corpus import read_jsonl
from math_grid_drawer.research.p_formula_gate06 import audit_p_formula_records06
from scripts.train_math_ink_06_p_formula_adapter import _file_sha25606


REQUIRED_SEEDS06 = (17, 31, 47)


def parse_seed_adapters06(values: Sequence[str]) -> dict[int, Path]:
    """필요 변수: `seed=checkpoint` 문자열. 작동 원리: 17·31·47의 중복 없는 adapter mapping을 만든다."""

    adapters: dict[int, Path] = {}
    for value in values:
        seed_text, separator, path_text = value.partition("=")
        if not separator or not path_text:
            raise ValueError("adapter는 `17=path/to/adapter.pt` 형식이어야 합니다.")
        try:
            seed = int(seed_text)
        except ValueError as error:
            raise ValueError(f"adapter seed가 정수가 아닙니다: {seed_text}") from error
        if seed in adapters:
            raise ValueError(f"adapter seed가 중복되었습니다: {seed}")
        adapters[seed] = Path(path_text)
    if set(adapters) != set(REQUIRED_SEEDS06):
        raise ValueError("seed 17·31·47 adapter가 정확히 하나씩 필요합니다.")
    return adapters


def build_p_formula_release_commands06(
    *,
    python: Path,
    data: Path,
    adapters: dict[int, Path],
    output: Path,
    recipe: dict[str, Any],
    convert_litert: bool,
) -> list[list[str]]:
    """필요 변수: 실행기·동일 P corpus·seed adapter·recipe. 작동 원리: 학습→요약→증류→export 순서의 argv를 고정한다."""

    training = recipe["training"]
    seed_gate = recipe["seed_gate"]
    source_minimum = int(recipe["input"]["minimum_independent_sources"])
    commands: list[list[str]] = []
    for seed in REQUIRED_SEEDS06:
        commands.append([
            str(python),
            str(PROJECT_ROOT / "scripts/train_math_ink_06_p_formula_adapter.py"),
            "--data", str(data),
            "--adapter", str(adapters[seed]),
            "--output", str(output / f"seed{seed}"),
            "--seed", str(seed),
            "--epochs", str(training["epochs"]),
            "--batch-size", str(training["batch_size"]),
            "--learning-rate", str(training["learning_rate"]),
            "--weight-decay", str(training["weight_decay"]),
            "--exact-loss-weight", str(training["exact_loss_weight"]),
            "--context-dropout", str(training["context_dropout"]),
            "--hidden-size", str(recipe["model"]["hidden_size"]),
            "--patience", str(training["patience"]),
            "--minimum-independent-sources", str(source_minimum),
            "--top1-minimum", str(seed_gate["exact_top1_minimum"]),
            "--top5-minimum", str(seed_gate["exact_top5_minimum"]),
            "--macro-f1-minimum", str(seed_gate["macro_f1_minimum"]),
            "--writer-floor-minimum", str(seed_gate["writer_floor_minimum"]),
            "--missing-drop-maximum-pp",
            str(seed_gate["missing_metadata_drop_maximum_pp"]),
            "--device", str(training["device"]),
        ])
    summary = output / "three_seed_summary.json"
    summary_command = [
        str(python),
        str(PROJECT_ROOT / "scripts/summarize_math_ink_06_p_formula_seeds.py"),
    ]
    for seed in REQUIRED_SEEDS06:
        summary_command.extend(["--report", str(output / f"seed{seed}/report.json")])
    summary_command.extend(["--output", str(summary)])
    commands.append(summary_command)
    distillation = recipe["distillation"]
    student_dir = output / "student"
    distill_command = [
        str(python),
        str(PROJECT_ROOT / "scripts/distill_math_ink_06_p_formula_student.py"),
        "--data", str(data),
    ]
    for seed in REQUIRED_SEEDS06:
        distill_command.extend([
            "--teacher-report",
            str(output / f"seed{seed}/report.json"),
        ])
    distill_command.extend([
        "--summary", str(summary),
        "--student-adapter", str(adapters[int(distillation["student_seed"])]),
        "--output", str(student_dir),
        "--seed", str(distillation["student_seed"]),
        "--epochs", str(training["epochs"]),
        "--batch-size", str(training["batch_size"]),
        "--learning-rate", str(training["learning_rate"]),
        "--weight-decay", str(training["weight_decay"]),
        "--hidden-size", str(distillation["student_formula_adapter_hidden_size"]),
        "--temperature", str(distillation["temperature"]),
        "--teacher-exact-weight", str(distillation["teacher_exact_weight"]),
        "--teacher-family-weight", str(distillation["teacher_family_weight"]),
        "--hard-exact-weight", str(distillation["hard_exact_weight"]),
        "--hard-family-weight", str(distillation["hard_family_weight"]),
        "--patience", str(training["patience"]),
        "--minimum-independent-sources", str(source_minimum),
        "--distillation-regression-maximum-pp",
        str(distillation["maximum_teacher_to_student_regression_pp"]),
        "--device", str(training["device"]),
    ])
    commands.append(distill_command)
    export_command = [
        str(python),
        str(PROJECT_ROOT / "scripts/export_math_ink_06_p_formula_student.py"),
        "--student-checkpoint", str(student_dir / "p_formula_student_adapter.pt"),
        "--data", str(data),
        "--output", str(output / "export"),
    ]
    if convert_litert:
        export_command.append("--convert-litert")
    commands.append(export_command)
    return commands


def _write_json06(path: Path, value: dict[str, Any]) -> None:
    """필요 변수: 목표 경로·JSON 값. 작동 원리: UTF-8 임시 파일을 원자적으로 교체한다."""

    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_suffix(path.suffix + ".part")
    temporary.write_text(
        json.dumps(value, ensure_ascii=False, indent=2) + "\n",
        encoding="utf-8",
    )
    temporary.replace(path)


def _require_gate06(path: Path, keys: Sequence[str]) -> dict[str, Any]:
    """필요 변수: 단계 report·중첩 boolean 경로. 작동 원리: 산출물 부재나 false gate에서 즉시 중단한다."""

    if not path.is_file():
        raise RuntimeError(f"단계 report가 생성되지 않았습니다: {path}")
    report = json.loads(path.read_text(encoding="utf-8"))
    value: Any = report
    for key in keys:
        value = value[key]
    if value is not True:
        raise RuntimeError(f"정식 release gate가 실패했습니다: {path} / {'.'.join(keys)}")
    return report


def _parse_args() -> argparse.Namespace:
    """필요 변수: 실제 P corpus·세 adapter·출력. 작동 원리: 정식 또는 dry-run release CLI를 만든다."""

    parser = argparse.ArgumentParser(description="Run Math Ink 0.6 P Formula release loop")
    parser.add_argument("--data", type=Path, required=True)
    parser.add_argument("--adapter", action="append", required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument(
        "--recipe",
        type=Path,
        default=PROJECT_ROOT / "research/configs/MATH-INK-06-P-FORMULA-v1.json",
    )
    parser.add_argument("--python", type=Path, default=Path(sys.executable))
    parser.add_argument("--convert-litert", action="store_true")
    parser.add_argument("--dry-run", action="store_true")
    return parser.parse_args()


def main() -> None:
    """필요 변수: 승인 P corpus와 seed별 제품 adapter. 작동 원리: 각 gate를 확인한 뒤에만 다음 외부 프로세스를 실행한다."""

    args = _parse_args()
    adapters = parse_seed_adapters06(args.adapter)
    missing = [str(path) for path in (args.data, args.recipe, args.python, *adapters.values()) if not path.is_file()]
    if missing:
        raise FileNotFoundError(f"release 입력 파일이 없습니다: {missing}")
    recipe = json.loads(args.recipe.read_text(encoding="utf-8"))
    if recipe.get("recipe_id") != "MATH-INK-06-P-FORMULA-v1":
        raise ValueError("지원하지 않는 P Formula release recipe입니다.")
    records = list(read_jsonl(args.data))
    audit = audit_p_formula_records06(
        records,
        minimum_independent_sources=int(recipe["input"]["minimum_independent_sources"]),
    )
    if not audit["eligible_for_product_evaluation"]:
        raise ValueError("P Formula corpus가 정식 release preflight를 통과하지 못했습니다.")
    commands = build_p_formula_release_commands06(
        python=args.python,
        data=args.data,
        adapters=adapters,
        output=args.output,
        recipe=recipe,
        convert_litert=args.convert_litert,
    )
    plan = {
        "schema": "aiflow-math-ink-06-p-formula-release-plan-v1",
        "generated_at": datetime.now(timezone.utc).isoformat(),
        "data": str(args.data),
        "data_sha256": _file_sha25606(args.data),
        "recipe": str(args.recipe),
        "adapters": {str(seed): str(path) for seed, path in adapters.items()},
        "commands": commands,
        "dry_run": bool(args.dry_run),
        "product_validation": False,
    }
    _write_json06(args.output / "release_plan.json", plan)
    if args.dry_run:
        print(json.dumps(plan, ensure_ascii=False, indent=2))
        return
    for index, command in enumerate(commands):
        subprocess.run(command, cwd=PROJECT_ROOT, check=True)
        if index < 3:
            seed = REQUIRED_SEEDS06[index]
            _require_gate06(args.output / f"seed{seed}/report.json", ("seed_gate", "passed"))
        elif index == 3:
            _require_gate06(
                args.output / "three_seed_summary.json",
                ("decision", "student_distillation_allowed"),
            )
        elif index == 4:
            _require_gate06(
                args.output / "student/report.json",
                ("distillation_gate_passed",),
            )
        elif index == 5:
            _require_gate06(
                args.output / "export/export_manifest.json",
                ("torch_export_gate_passed",),
            )
    final = {
        **plan,
        "dry_run": False,
        "completed": True,
        "torch_export_gate_passed": True,
        "litert_requested": bool(args.convert_litert),
        "android_validation": False,
        "product_validation": False,
        "next_gate": "LiteRT parity and Android low/mid/high tier benchmark",
    }
    _write_json06(args.output / "release_report.json", final)


if __name__ == "__main__":
    main()