scripts/run_vc02_onnx_accuracy_supplement.py

#!/usr/bin/env python3
"""Run, validate, checksum, and log the VC02 ONNX accuracy supplement."""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import platform
import shlex
import subprocess
import tempfile
from pathlib import Path
from typing import Any


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for block in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(block)
    return digest.hexdigest()


def atomic_json(path: Path, value: Any) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with tempfile.NamedTemporaryFile("w", encoding="utf-8", dir=path.parent, delete=False) as handle:
        json.dump(value, handle, indent=2, sort_keys=True)
        handle.write("\n")
        temporary = Path(handle.name)
    os.replace(temporary, path)


def atomic_text(path: Path, value: str) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with tempfile.NamedTemporaryFile("w", encoding="utf-8", dir=path.parent, delete=False) as handle:
        handle.write(value)
        temporary = Path(handle.name)
    os.replace(temporary, path)


def run_stage(name: str, command: list[str], root: Path, run_dir: Path, env: dict[str, str]) -> dict[str, Any]:
    command_path = run_dir / f"{name}.command.txt"
    stdout_path = run_dir / f"{name}.stdout.log"
    stderr_path = run_dir / f"{name}.stderr.log"
    exit_path = run_dir / f"{name}.exit_code.txt"
    atomic_text(command_path, shlex.join(command) + "\n")
    completed = subprocess.run(command, cwd=root, env=env, text=True, capture_output=True, check=False)
    atomic_text(stdout_path, completed.stdout)
    atomic_text(stderr_path, completed.stderr)
    atomic_text(exit_path, f"{completed.returncode}\n")
    return {
        "name": name,
        "status": "PASS" if completed.returncode == 0 else "FAIL",
        "exit_code": completed.returncode,
        "command": shlex.join(command),
        "command_log": str(command_path),
        "stdout_log": str(stdout_path),
        "stderr_log": str(stderr_path),
        "exit_code_log": str(exit_path),
        "command_log_sha256": sha256_file(command_path),
        "stdout_log_sha256": sha256_file(stdout_path),
        "stderr_log_sha256": sha256_file(stderr_path),
        "exit_code_log_sha256": sha256_file(exit_path),
    }


def build_manifest(root: Path, output_dir: Path, report: Path, run_dir: Path) -> None:
    excluded = {output_dir / "artifact_manifest.json", output_dir / "artifacts.sha256"}
    files = sorted(path for path in output_dir.rglob("*") if path.is_file() and path not in excluded)
    files.append(report)
    files.extend([
        root / "configs/evaluation/vision_classification/VC02_onnx_runtime_quality_supplement.json",
        root / "scripts/stages/evaluate_vc02_onnx_accuracy.py",
        root / "scripts/validate_vc02_onnx_accuracy.py",
        root / "scripts/run_vc02_onnx_accuracy_supplement.py",
        root / "tests/test_vc02_onnx_accuracy_evaluator.py",
        root / "tests/test_vc02_onnx_accuracy.py",
    ])
    files.extend(sorted(path for path in run_dir.rglob("*") if path.is_file()))
    records = []
    for path in files:
        records.append({
            "path": str(path.relative_to(root)),
            "bytes": path.stat().st_size,
            "sha256": sha256_file(path),
        })
    summary = json.loads((output_dir / "quality_summary.json").read_text())
    validation = json.loads((output_dir / "validation.json").read_text())
    manifest = {
        "schema_version": "1.0",
        "model_id": "VC02",
        "stage": "ONNX_RUNTIME_Q1_SUPPLEMENT",
        "status": summary["status"],
        "independent_validation_status": validation["status"],
        "files": records,
        "policy": summary["policy"],
        "execution_logs": str(run_dir.relative_to(root)),
    }
    manifest_path = output_dir / "artifact_manifest.json"
    atomic_json(manifest_path, manifest)
    checksum_rows = records + [{
        "path": str(manifest_path.relative_to(root)),
        "sha256": sha256_file(manifest_path),
    }]
    atomic_text(output_dir / "artifacts.sha256", "".join(f"{row['sha256']}  {row['path']}\n" for row in checksum_rows))


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--repo-root", required=True, type=Path)
    parser.add_argument("--output-dir", required=True, type=Path)
    parser.add_argument("--run-dir", required=True, type=Path)
    parser.add_argument("--report", required=True, type=Path)
    args = parser.parse_args()
    root = args.repo_root.resolve()
    output_dir = args.output_dir.resolve()
    run_dir = args.run_dir.resolve()
    report = args.report.resolve()
    if output_dir.exists() and any(output_dir.iterdir()):
        raise FileExistsError(f"refusing non-empty output directory: {output_dir}")
    if run_dir.exists() and any(run_dir.iterdir()):
        raise FileExistsError(f"refusing non-empty run directory: {run_dir}")
    if report.exists():
        raise FileExistsError(f"refusing existing report: {report}")
    output_dir.mkdir(parents=True, exist_ok=True)
    run_dir.mkdir(parents=True, exist_ok=True)
    python = root / ".venv/bin/python"
    env = dict(os.environ)
    for name in ("OMP_NUM_THREADS", "OPENBLAS_NUM_THREADS", "MKL_NUM_THREADS", "NUMBA_NUM_THREADS"):
        env[name] = "1"
    stages = []
    stages.append(run_stage("evaluate", [
        str(python), "scripts/stages/evaluate_vc02_onnx_accuracy.py",
        "--repo-root", str(root),
        "--config", str(root / "configs/evaluation/vision_classification/VC02_onnx_runtime_quality_supplement.json"),
        "--output-dir", str(output_dir),
        "--report", str(report),
    ], root, run_dir, env))
    validator = root / "scripts/validate_vc02_onnx_accuracy.py"
    if stages[-1]["status"] == "PASS" and validator.is_file():
        stages.append(run_stage("validate", [
            str(python), str(validator),
            "--repo-root", str(root),
            "--result-dir", str(output_dir),
            "--output", str(output_dir / "validation.json"),
        ], root, run_dir, env))
    else:
        stages.append({"name": "validate", "status": "BLOCKED", "exit_code": None})
    tests = ["tests/test_vc02_onnx_accuracy_evaluator.py"]
    validator_test = root / "tests/test_vc02_onnx_accuracy.py"
    if validator_test.is_file():
        tests.append("tests/test_vc02_onnx_accuracy.py")
    stages.append(run_stage("targeted_tests", [str(python), "-m", "pytest", "-q", *tests], root, run_dir, env))
    compile_targets = [
        "scripts/stages/evaluate_vc02_onnx_accuracy.py",
        "scripts/run_vc02_onnx_accuracy_supplement.py",
    ]
    if validator.is_file():
        compile_targets.append("scripts/validate_vc02_onnx_accuracy.py")
    stages.append(run_stage("compileall", [str(python), "-m", "py_compile", *compile_targets], root, run_dir, env))
    ready = all(stage["status"] == "PASS" for stage in stages) and (output_dir / "validation.json").is_file()
    stages.append({"name": "checksum_manifest", "status": "PENDING" if ready else "BLOCKED", "exit_code": None})
    summary_status = None
    if (output_dir / "quality_summary.json").is_file():
        summary_status = json.loads((output_dir / "quality_summary.json").read_text()).get("status")
    validation_status = None
    if (output_dir / "validation.json").is_file():
        validation_status = json.loads((output_dir / "validation.json").read_text()).get("status")
    preliminary_overall = "PASS" if (
        summary_status in {"PASS", "PARTIAL"}
        and validation_status == "PASS"
        and all(stage["status"] == "PASS" for stage in stages[:-1])
    ) else "FAIL"
    execution_manifest = {
        "schema_version": "1.0",
        "model_id": "VC02",
        "stage": "ONNX_RUNTIME_Q1_SUPPLEMENT",
        "status": preliminary_overall,
        "result_status": summary_status,
        "validation_status": validation_status,
        "stages": stages,
        "environment": {"python": platform.python_version(), "platform": platform.platform(), "threads_per_runtime": 1},
        "policy": {
            "latency_benchmark": False,
            "converter_run": False,
            "mlir_lowering_or_codegen_run": False,
            "dataset_download": False,
            "training_calibration_or_quantization": False,
            "model_or_weight_modification": False,
        },
    }
    atomic_json(run_dir / "execution_manifest.json", execution_manifest)
    atomic_text(run_dir / "execution_manifest.sha256", f"{sha256_file(run_dir / 'execution_manifest.json')}  {run_dir.relative_to(root)}/execution_manifest.json\n")
    if preliminary_overall == "PASS":
        build_manifest(root, output_dir, report, run_dir)
        checksum_stage = run_stage("checksum_manifest", [
            "sha256sum", "-c", str(output_dir / "artifacts.sha256")
        ], root, run_dir, env)
    else:
        checksum_stage = {"name": "checksum_manifest", "status": "BLOCKED", "exit_code": None}
    final_overall = "PASS" if preliminary_overall == "PASS" and checksum_stage["status"] == "PASS" else "FAIL"
    final_record = {
        "status": final_overall,
        "result_status": summary_status,
        "validation_status": validation_status,
        "checksum_stage": checksum_stage,
        "execution_manifest_sha256": sha256_file(run_dir / "execution_manifest.json"),
        "artifact_manifest_sha256": sha256_file(output_dir / "artifact_manifest.json") if (output_dir / "artifact_manifest.json").is_file() else None,
        "artifacts_sha256_sha256": sha256_file(output_dir / "artifacts.sha256") if (output_dir / "artifacts.sha256").is_file() else None,
    }
    atomic_json(run_dir / "final_result.json", final_record)
    atomic_text(run_dir / "final_result.sha256", f"{sha256_file(run_dir / 'final_result.json')}  {run_dir.relative_to(root)}/final_result.json\n")
    log_files = sorted(
        path for path in run_dir.rglob("*")
        if path.is_file() and path.name not in {"logs.sha256"}
    )
    atomic_text(
        run_dir / "logs.sha256",
        "".join(f"{sha256_file(path)}  {path.relative_to(root)}\n" for path in log_files),
    )
    print(json.dumps({"status": final_overall, "result_status": summary_status, "validation_status": validation_status}, sort_keys=True))
    return 0 if final_overall == "PASS" else 1


if __name__ == "__main__":
    raise SystemExit(main())