scripts/build_math_ink_06_p_mobile_colab_bundle.py

"""민감한 P corpus를 공개본과 분리한 제품용 LiteRT Colab bundle을 만든다."""

from __future__ import annotations

import argparse
from datetime import datetime, timezone
from hashlib import sha256
import json
from pathlib import Path, PurePosixPath
from zipfile import ZIP_DEFLATED, ZIP_STORED, ZipFile, ZipInfo

import torch

from scripts.export_math_ink_06_p_formula_student import (
    _file_sha25606,
    validate_p_formula_student_artifacts06,
    validate_p_formula_student_export06,
)


PROJECT_ROOT = Path(__file__).parents[1]
MANIFEST_NAME_06 = "P_MOBILE_COLAB_BUNDLE_MANIFEST.json"


def _bundle_files06(
    *,
    base_checkpoint: Path,
    adapter_checkpoint: Path,
    student_checkpoint: Path,
    p_data: Path,
) -> dict[str, Path]:
    """필요 변수: P student 전체 lineage. 작동 원리: product pair exporter의 닫힌 최소 파일 집합을 만든다."""

    data_name = "private/p_formula.jsonl.gz" if p_data.suffix == ".gz" else "private/p_formula.jsonl"
    return {
        "pyproject.toml": PROJECT_ROOT / "pyproject.toml",
        "src/math_grid_drawer/__init__.py": PROJECT_ROOT / "research/colab/litert_bundle_package_init.py",
        "src/math_grid_drawer/research/__init__.py": PROJECT_ROOT / "src/math_grid_drawer/research/__init__.py",
        "src/math_grid_drawer/research/ink06_canonical.py": PROJECT_ROOT / "src/math_grid_drawer/research/ink06_canonical.py",
        "src/math_grid_drawer/research/ink06_export.py": PROJECT_ROOT / "src/math_grid_drawer/research/ink06_export.py",
        "src/math_grid_drawer/research/math_ink_06.py": PROJECT_ROOT / "src/math_grid_drawer/research/math_ink_06.py",
        "src/math_grid_drawer/research/p_formula_dataset06.py": PROJECT_ROOT / "src/math_grid_drawer/research/p_formula_dataset06.py",
        "src/math_grid_drawer/research/p_formula_gate06.py": PROJECT_ROOT / "src/math_grid_drawer/research/p_formula_gate06.py",
        "src/math_grid_drawer/research/raster_skeleton06.py": PROJECT_ROOT / "src/math_grid_drawer/research/raster_skeleton06.py",
        "src/math_grid_drawer/research/skeleton_adapter06.py": PROJECT_ROOT / "src/math_grid_drawer/research/skeleton_adapter06.py",
        "src/math_grid_drawer/research/trajectory_sequence.py": PROJECT_ROOT / "src/math_grid_drawer/research/trajectory_sequence.py",
        "scripts/export_math_ink_06_litert.py": PROJECT_ROOT / "scripts/export_math_ink_06_litert.py",
        "scripts/export_math_ink_06_p_formula_student.py": PROJECT_ROOT / "scripts/export_math_ink_06_p_formula_student.py",
        "scripts/export_math_ink_06_p_mobile_pair.py": PROJECT_ROOT / "scripts/export_math_ink_06_p_mobile_pair.py",
        "artifacts/base_378.pt": base_checkpoint,
        "artifacts/online_adapter.pt": adapter_checkpoint,
        "artifacts/p_formula_student_adapter.pt": student_checkpoint,
        data_name: p_data,
    }


def build_p_mobile_colab_bundle06(
    *,
    base_checkpoint: Path,
    adapter_checkpoint: Path,
    student_checkpoint: Path,
    p_data: Path,
    output: Path,
) -> dict:
    """필요 변수: P model/data·출력. 작동 원리: hash 검증 후 공개 금지 private Colab archive를 결정적으로 만든다."""

    for path in (base_checkpoint, adapter_checkpoint, student_checkpoint, p_data):
        if not path.is_file():
            raise FileNotFoundError(f"제품 Colab bundle 필수 파일이 없습니다: {path}")
    student = torch.load(student_checkpoint, map_location="cpu", weights_only=False)
    data_sha256 = _file_sha25606(p_data)
    validate_p_formula_student_export06(student, data_sha256=data_sha256)
    validate_p_formula_student_artifacts06(
        student,
        base_checkpoint=base_checkpoint,
        online_adapter=adapter_checkpoint,
    )
    files = _bundle_files06(
        base_checkpoint=base_checkpoint,
        adapter_checkpoint=adapter_checkpoint,
        student_checkpoint=student_checkpoint,
        p_data=p_data,
    )
    missing = [name for name, path in files.items() if not path.is_file()]
    if missing:
        raise FileNotFoundError(f"제품 Colab source가 없습니다: {missing}")
    payloads = {name: path.read_bytes() for name, path in files.items()}
    data_name = next(name for name in payloads if name.startswith("private/"))
    manifest = {
        "schema": "aiflow-math-ink-06-p-mobile-colab-bundle-v1",
        "generated_at": datetime.now(timezone.utc).isoformat(),
        "litert_torch_version": "0.9.1",
        "data_sha256": data_sha256,
        "teacher_seeds": [17, 31, 47],
        "privacy": {
            "contains_p_data": True,
            "public_upload_allowed": False,
            "server_upload_permitted": False,
        },
        "raster_output_count": 5,
        "command": [
            "python",
            "scripts/export_math_ink_06_p_mobile_pair.py",
            "--student-checkpoint",
            "artifacts/p_formula_student_adapter.pt",
            "--base-checkpoint",
            "artifacts/base_378.pt",
            "--adapter-checkpoint",
            "artifacts/online_adapter.pt",
            "--data",
            data_name,
            "--output",
            "outputs",
            "--convert-litert",
        ],
        "files": [
            {
                "path": name,
                "bytes": len(payload),
                "sha256": sha256(payload).hexdigest(),
            }
            for name, payload in payloads.items()
        ],
        "product_validation": False,
    }
    output.parent.mkdir(parents=True, exist_ok=True)
    temporary = output.with_suffix(output.suffix + ".part")
    with ZipFile(temporary, "w", allowZip64=True) as bundle:
        for name, payload in payloads.items():
            compression = (
                ZIP_STORED
                if PurePosixPath(name).suffix in {".pt", ".gz"}
                else ZIP_DEFLATED
            )
            info = ZipInfo(name, date_time=(1980, 1, 1, 0, 0, 0))
            info.compress_type = compression
            info.external_attr = 0o600 << 16 if name.startswith("private/") else 0o644 << 16
            bundle.writestr(info, payload)
        info = ZipInfo(MANIFEST_NAME_06, date_time=(1980, 1, 1, 0, 0, 0))
        info.compress_type = ZIP_DEFLATED
        info.external_attr = 0o600 << 16
        bundle.writestr(
            info,
            json.dumps(manifest, ensure_ascii=False, indent=2).encode("utf-8") + b"\n",
        )
    temporary.replace(output)
    manifest["bundle"] = str(output)
    manifest["bundle_bytes"] = output.stat().st_size
    manifest["bundle_sha256"] = sha256(output.read_bytes()).hexdigest()
    return manifest


def verify_p_mobile_colab_bundle06(bundle_path: Path) -> dict:
    """필요 변수: private bundle. 작동 원리: 안전 경로·hash·privacy fail-closed 계약을 추출 없이 검사한다."""

    failures = []
    with ZipFile(bundle_path) as bundle:
        manifest = json.loads(bundle.read(MANIFEST_NAME_06).decode("utf-8"))
        privacy = manifest.get("privacy") or {}
        if privacy != {
            "contains_p_data": True,
            "public_upload_allowed": False,
            "server_upload_permitted": False,
        }:
            failures.append({"path": MANIFEST_NAME_06, "reason": "privacy_contract"})
        if int(manifest.get("raster_output_count", 0)) != 5:
            failures.append({"path": MANIFEST_NAME_06, "reason": "raster_output_count"})
        names = set(bundle.namelist())
        for row in manifest.get("files", []):
            name = str(row.get("path") or "")
            path = PurePosixPath(name)
            if path.is_absolute() or ".." in path.parts:
                failures.append({"path": name, "reason": "unsafe_path"})
            elif name not in names:
                failures.append({"path": name, "reason": "missing"})
            else:
                payload = bundle.read(name)
                if len(payload) != int(row["bytes"]):
                    failures.append({"path": name, "reason": "bytes"})
                elif sha256(payload).hexdigest() != str(row["sha256"]):
                    failures.append({"path": name, "reason": "sha256"})
    return {
        "schema": manifest.get("schema"),
        "files": len(manifest.get("files", [])),
        "failures": failures,
        "passed": not failures,
        "public_upload_allowed": False,
        "product_validation": False,
    }


def main() -> None:
    """필요 변수: 실제 P lineage·출력. 작동 원리: private bundle 생성 직후 archive를 재검증한다."""

    parser = argparse.ArgumentParser(description="Build private P mobile Colab bundle")
    parser.add_argument("--base-checkpoint", type=Path, required=True)
    parser.add_argument("--adapter-checkpoint", type=Path, required=True)
    parser.add_argument("--student-checkpoint", type=Path, required=True)
    parser.add_argument("--p-data", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()
    report = build_p_mobile_colab_bundle06(
        base_checkpoint=args.base_checkpoint,
        adapter_checkpoint=args.adapter_checkpoint,
        student_checkpoint=args.student_checkpoint,
        p_data=args.p_data,
        output=args.output,
    )
    report["verification"] = verify_p_mobile_colab_bundle06(args.output)
    if not report["verification"]["passed"]:
        raise ValueError(f"제품 Colab bundle 검증 실패: {report['verification']['failures']}")
    report_path = args.output.with_suffix(args.output.suffix + ".manifest.json")
    report_path.write_text(
        json.dumps(report, ensure_ascii=False, indent=2) + "\n",
        encoding="utf-8",
    )
    print(json.dumps(report, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()