scripts/measure_latency.py

"""Latency harness for the fal masked-inpaint beautify engine.

Answers the operational question we measured cleanly on 2026-06-18: when the fal
container is warm the customer waits ~14s (default) — the 10-20s target is already
met warm; the real risk is the cold start. This harness re-measures on demand and
breaks the wait into every stage so regressions are obvious.

It calls the SAME production core (`beautify_with_fal`) plus the same watermark step
the worker applies, so the "customer wait" number is realistic. It runs each preset
N times, flags run 1 as possibly-cold, and reports the warm median + the warm/cold
telemetry fields (preset, provider latency, local-model-preloaded, first-load cost).

Gated like every paid path: refuses with no FAL_KEY (no network call). Each run is a
real (small) fal request and costs ~$0.04. Output (timings only, no images) goes
under runtime/ (gitignored). Pilot Ready: NOT CONFIRMED.

Usage (via wrapper so the key is the last arg):
  scripts/measure_latency.ps1 -Source "runtime/private-inputs/man-01.jpg" -Runs 3 -Token <FAL_KEY>
"""
from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from app.services.fal_client import FalUnavailable, fal_inpaint_model, fal_real_enabled  # noqa: E402
from app.services.falinpaint_beautify import (  # noqa: E402
    TIMING_STAGES,
    beautify_with_fal,
    face_model_preloaded,
    preload_face_model,
)
from app.services.watermark import apply_ai_watermark  # noqa: E402

# Preset definitions mirror the CLI/worker: "default" = validated quality recipe,
# "fast" = the speed preset (-Fast): fewer steps + smaller size, slight quality drop.
PRESETS = {
    "default": {"steps": 28, "max_size": 1024},
    "fast": {"steps": 16, "max_size": 832},
}

# Engine stages + the pipeline watermark step (applied here as the worker does).
STAGE_KEYS = (*TIMING_STAGES, "watermark")


def median(xs: list[float]) -> float:
    """Median of a list (0.0 for empty). No numpy dependency."""
    s = sorted(x for x in xs if x is not None)
    n = len(s)
    if n == 0:
        return 0.0
    mid = n // 2
    return s[mid] if n % 2 else (s[mid - 1] + s[mid]) / 2.0


def summarize(runs: list[dict], key: str) -> dict:
    """min/median/mean/max of `key` over the OK runs given (excludes failures)."""
    vals = [r["timings"].get(key) for r in runs
            if r.get("ok") and isinstance(r.get("timings"), dict)
            and r["timings"].get(key) is not None]
    if not vals:
        return {"n": 0}
    return {
        "n": len(vals),
        "min": round(min(vals), 2),
        "median": round(median(vals), 2),
        "mean": round(sum(vals) / len(vals), 2),
        "max": round(max(vals), 2),
    }


def build_run_record(metrics: dict, watermark_s: float) -> dict:
    """Per-run latency record. Carries the spec judgement fields VERBATIM
    (provider_latency_ms / total_latency_ms / preset_name / local_model_preloaded)
    so a parser greps the JSON for the exact spec names. Pure (no I/O)."""
    timings = dict(metrics.get("timings_s", {}))
    timings["watermark"] = round(watermark_s, 3)
    total_ms = metrics.get("total_latency_ms", 0)
    return {
        "ok": True,
        "timings": timings,
        "preset_name": metrics.get("preset_name"),
        "local_model_preloaded": metrics.get("local_model_preloaded"),
        "provider_latency_ms": metrics.get("provider_latency_ms"),
        "total_latency_ms": total_ms,
        "gender": metrics.get("gender_detected"),
        "customer_wait_s": round(total_ms / 1000.0 + watermark_s, 3),
    }


def build_preset_summary(runs: list[dict], preset_name: str, preloaded: bool) -> dict:
    """Per-preset summary. Emits the spec judgement fields under their exact names
    (provider_latency_ms, total_latency_ms via warm median; face_detect_first_load_ms
    from run 1; estimated_customer_wait_sec). Pure."""
    warm = [r for r in runs if r.get("index", 0) >= 2 and r.get("ok")]
    first = next((r for r in runs if r.get("index") == 1 and r.get("ok")), None)
    first_face_ms = (round(first["timings"].get("face_detect", 0.0) * 1000)
                     if first else None)
    warm_provider = [r["provider_latency_ms"] for r in warm
                     if r.get("provider_latency_ms") is not None]
    warm_total_ms = [r["total_latency_ms"] for r in warm
                     if r.get("total_latency_ms") is not None]
    warm_wait = [r["customer_wait_s"] for r in warm if r.get("customer_wait_s") is not None]
    return {
        "preset_name": preset_name,
        "local_model_preloaded": preloaded,
        "cold_first_run_total_s": (first or {}).get("timings", {}).get("total") if first else None,
        "face_detect_first_load_ms": first_face_ms,
        "warm_total_s": summarize(warm, "total"),
        "provider_latency_ms": round(median(warm_provider)) if warm_provider else None,
        "total_latency_ms": round(median(warm_total_ms)) if warm_total_ms else None,
        "estimated_customer_wait_sec": round(median(warm_wait), 2) if warm_wait else None,
    }


def _run_once(src_bytes: bytes, preset_opts: dict) -> dict:
    """One beautify call + watermark (the real customer path). Never raises."""
    import time

    try:
        final, metrics, _mask = beautify_with_fal(src_bytes, **preset_opts)
        t = time.perf_counter()
        apply_ai_watermark(final)
        wm = round(time.perf_counter() - t, 3)
        return build_run_record(metrics, wm)
    except FalUnavailable as exc:
        return {"ok": False, "error": str(exc)}
    except Exception as exc:  # noqa: BLE001
        return {"ok": False, "error": f"{type(exc).__name__}: {exc}"}


def _fmt_stage_row(label: str, t: dict) -> str:
    cells = "  ".join(f"{k}={t.get(k, '-')!s:>7}" for k in STAGE_KEYS)
    return f"  {label:<14} {cells}"


def main(argv: list[str] | None = None) -> int:
    ap = argparse.ArgumentParser(description="fal beautify latency harness")
    ap.add_argument("--source", required=True, help="a real customer-style photo (gitignored input)")
    ap.add_argument("--runs", type=int, default=3, help="runs per preset (run 1 = possibly cold)")
    ap.add_argument("--preset", choices=("default", "fast", "both"), default="both")
    ap.add_argument("--preload", action="store_true",
                    help="pre-load the face model first (measures the preloaded worker case)")
    ap.add_argument("--tag", default="lat-01")
    ap.add_argument("--evidence-root", default="runtime/gemini-smoke-evidence")
    args = ap.parse_args(argv)

    src = Path(args.source)
    if not src.exists():
        print(f"REFUSED: source not found: {src}")
        return 2
    if not fal_real_enabled():
        print("REFUSED: FAL_KEY not set (no network call made).")
        return 2

    if args.preload:
        print(f"preload: {preload_face_model()}")
    presets = ["default", "fast"] if args.preset == "both" else [args.preset]
    src_bytes = src.read_bytes()
    out_dir = Path(args.evidence_root) / "gemini-smoke" / "falinpaint" / "_latency"
    out_dir.mkdir(parents=True, exist_ok=True)

    print(f"latency: model={fal_inpaint_model()} runs={args.runs} presets={presets} "
          f"preloaded={face_model_preloaded()}")
    print("note: run 1 of the session may be COLD (fal spins the container up); "
          "runs 2+ are the warm steady state.\n")

    results: dict = {"model": fal_inpaint_model(), "runs_per_preset": args.runs,
                     "local_model_preloaded": face_model_preloaded(),
                     "presets": {}, "pilot_ready": "NOT CONFIRMED"}

    for preset in presets:
        opts = PRESETS[preset]
        runs: list[dict] = []
        print(f"[{preset}] steps={opts['steps']} max_size={opts['max_size']}")
        for i in range(1, args.runs + 1):
            tag = "cold?" if i == 1 else "warm"
            r = _run_once(src_bytes, opts)
            r["index"] = i
            r["phase"] = tag
            runs.append(r)
            if r["ok"]:
                print(_fmt_stage_row(f"run {i} ({tag})", r["timings"]))
            else:
                print(f"  run {i} ({tag})  FAILED: {r['error']}")

        summary = build_preset_summary(runs, preset, face_model_preloaded())
        results["presets"][preset] = {"opts": opts, "runs": runs, "summary": summary}

        wt = summary["warm_total_s"]
        if wt.get("n"):
            print(f"  -> warm total median {wt['median']}s (min {wt['min']}/max {wt['max']}, "
                  f"n={wt['n']}); provider(fal) ~{summary['provider_latency_ms']}ms; "
                  f"est customer wait ~{summary['estimated_customer_wait_sec']}s; "
                  f"cold first-run {summary['cold_first_run_total_s']}s; "
                  f"first face-load {summary['face_detect_first_load_ms']}ms")
        print()

    out = out_dir / f"{args.tag}.json"
    out.write_text(json.dumps(results, indent=2, ensure_ascii=False), encoding="utf-8")
    print(f"latency: wrote {out}")
    print("Pilot Ready: NOT CONFIRMED.")
    return 0


if __name__ == "__main__":
    try:
        raise SystemExit(main())
    except KeyboardInterrupt:
        print("\nlatency: stopped.")