surya-ocr/scripts/quant/run_all.py

"""Orchestrate the quantization sweep: baseline -> each method -> aggregate -> plot.

Each method is isolated: any exception becomes a status="failed" row so one bad
method never aborts the sweep.
"""
from __future__ import annotations

import argparse
import json
from pathlib import Path
from typing import Any, Dict, List

from scripts.quant.aggregate import build_row, rows_to_csv, rows_to_markdown
from scripts.quant.bbox_iou import bbox_iou_over_dirs
from scripts.quant.build_model import build_model
from scripts.quant.capture import capture_page
from scripts.quant.manifest import load_manifest
from scripts.quant.recipes import METHOD_SPECS, method_names, vllm_serve_args
from scripts.quant.serve import start_server, stop_server, wait_healthy

# Reuse the existing CER metric.
from scripts.cer_divergence import cer_over_dirs

OCR_URL = "http://127.0.0.1:5002/v1/api/ai/suya_ocr_vllm/"


def _dir_size_mb(path: Path) -> float:
    total = sum(f.stat().st_size for f in path.rglob("*") if f.is_file())
    return round(total / (1024 * 1024), 1)


def _measure_method(method: str, base_model: str, work_dir: Path,
                    eval_images: List[Path], reference_dir: Path) -> Dict[str, Any]:
    """Build -> serve -> capture -> metrics for one method. Raises on any failure.

    Assumes the API process (api.py) is already running and reads SURYA_INFERENCE_URL
    from the env to point at the per-method vLLM server (port set by the runbook).
    Returns a dict of measured summary fields (no method/status/t4 keys).
    """
    model_dir = work_dir / "models" / method
    out_capture = work_dir / "results" / method
    port = 8000

    built = build_model(method, base_model, model_dir, eval_images[: min(8, len(eval_images))])
    serve_args = vllm_serve_args(method, str(built), base_model, port)
    proc = start_server(serve_args, str(work_dir / f"vllm_{method}.log"))
    try:
        if not wait_healthy(port):
            raise RuntimeError(f"vLLM did not become healthy for {method}")
        latencies = []
        for image in eval_images:
            cap = capture_page(OCR_URL, image, out_capture, timeout=900)
            if cap["elapsed_seconds"] is not None:
                latencies.append(cap["elapsed_seconds"])
        cer = cer_over_dirs(reference_dir, out_capture)
        iou = bbox_iou_over_dirs(reference_dir, out_capture)
    finally:
        stop_server(port, proc)

    latencies.sort()
    n = len(latencies)
    size = _dir_size_mb(built) if METHOD_SPECS[method]["kind"] in ("compressor", "bnb") else None
    return {
        "mean_latency_s": round(sum(latencies) / n, 3) if n else None,
        "p50_latency_s": latencies[n // 2] if n else None,
        "p95_latency_s": latencies[min(n - 1, int(n * 0.95))] if n else None,
        "mean_cer": cer["mean_cer"],
        "max_cer": cer["max_cer"],
        "mean_bbox_iou": iou["mean_bbox_iou"],
        "mean_missed_lines": iou["mean_missed_lines"],
        "mean_extra_lines": iou["mean_extra_lines"],
        "model_size_mb": size,
    }


def run_method(method: str, base_model: str, work_dir: Path,
              eval_images: List[Path], reference_dir: Path) -> Dict[str, Any]:
    t4 = METHOD_SPECS[method]["t4_deployable"]
    try:
        metrics = _measure_method(method, base_model, work_dir, eval_images, reference_dir)
        return build_row(method=method, status="ok", t4_deployable=t4, **metrics)
    except Exception as exc:  # feasibility probe: record and continue
        return build_row(method=method, status="failed", t4_deployable=t4, error=str(exc))


def main() -> None:
    parser = argparse.ArgumentParser(description="Run the quantization benchmark sweep.")
    parser.add_argument("--base-model", default="datalab-to/surya-ocr-2")
    parser.add_argument("--manifest", type=Path, default=Path("eval_set/manifest.txt"))
    parser.add_argument("--image-root", type=Path, default=Path("."))
    parser.add_argument("--work-dir", type=Path, default=Path("results/quant"))
    parser.add_argument("--methods", nargs="*", default=method_names())
    args = parser.parse_args()

    eval_images = load_manifest(args.manifest, args.image_root)
    reference_dir = args.work_dir / "reference"
    args.work_dir.mkdir(parents=True, exist_ok=True)

    rows = []
    for method in args.methods:
        print(f"=== {method} ===", flush=True)
        rows.append(run_method(method, args.base_model, args.work_dir, eval_images, reference_dir))

    (args.work_dir / "summary.csv").write_text(rows_to_csv(rows), encoding="utf-8")
    (args.work_dir / "summary.json").write_text(json.dumps(rows, indent=2), encoding="utf-8")
    (args.work_dir / "summary.md").write_text(rows_to_markdown(rows), encoding="utf-8")
    print(f"wrote summary for {len(rows)} methods to {args.work_dir}")


if __name__ == "__main__":
    main()