scripts/benchmark_onnx.py

"""Benchmark LightGBM joblib vs ONNX Runtime inference.

Loads N real feature rows from `data/reference_dataset.parquet`, runs each
sample one at a time through both backends (simulating the per-request path),
and reports p50/p95/p99 latency plus numerical equivalence.

Run:
    uv run python scripts/benchmark_onnx.py --n 1000
"""

from __future__ import annotations

import argparse
import json
import time
from pathlib import Path
from statistics import median

import joblib
import numpy as np
import onnxruntime as ort
import pandas as pd

ROOT = Path(__file__).resolve().parents[1]
DEFAULT_MODEL_PATH = ROOT / "models" / "model.joblib"
DEFAULT_ONNX_PATH = ROOT / "models" / "model.onnx"
DEFAULT_FEATURE_NAMES_PATH = ROOT / "models" / "feature_names.json"
DEFAULT_REFERENCE_PATH = ROOT / "data" / "reference_dataset.parquet"


def percentile(samples: list[float], pct: float) -> float:
    """Linear-interpolated percentile (pct in 0..100)."""
    return float(np.percentile(samples, pct))


def fmt_ms(seconds: float) -> str:
    return f"{seconds * 1000:.3f} ms"


def bench_lightgbm(model, samples_df: pd.DataFrame) -> tuple[list[float], np.ndarray]:
    """Run predict_proba one row at a time. Returns (latencies_s, probas)."""
    latencies: list[float] = []
    probas = np.empty(len(samples_df), dtype=np.float64)
    for i in range(len(samples_df)):
        row = samples_df.iloc[[i]]
        t0 = time.perf_counter()
        p = model.predict_proba(row)[0, 1]
        latencies.append(time.perf_counter() - t0)
        probas[i] = p
    return latencies, probas


def bench_onnx(
    session: ort.InferenceSession,
    input_name: str,
    output_name: str,
    samples_array: np.ndarray,
) -> tuple[list[float], np.ndarray]:
    """Run ONNX inference one row at a time. Returns (latencies_s, probas)."""
    latencies: list[float] = []
    probas = np.empty(len(samples_array), dtype=np.float64)
    for i in range(len(samples_array)):
        # Reshape to (1, n_features); float32 is required by the ONNX graph.
        row = samples_array[i : i + 1]
        t0 = time.perf_counter()
        out = session.run([output_name], {input_name: row})[0]
        latencies.append(time.perf_counter() - t0)
        # zipmap=False → out shape is (1, 2): [prob_class_0, prob_class_1].
        probas[i] = out[0, 1]
    return latencies, probas


def summarise(name: str, latencies: list[float]) -> dict[str, float]:
    return {
        "backend": name,
        "n": len(latencies),
        "p50_ms": percentile(latencies, 50) * 1000,
        "p95_ms": percentile(latencies, 95) * 1000,
        "p99_ms": percentile(latencies, 99) * 1000,
        "mean_ms": float(np.mean(latencies)) * 1000,
        "total_s": sum(latencies),
    }


def print_table(rows: list[dict[str, float]]) -> None:
    print(
        f"{'Backend':<14}{'n':>6}{'p50':>12}{'p95':>12}{'p99':>12}{'mean':>12}"
    )
    for r in rows:
        print(
            f"{r['backend']:<14}{r['n']:>6}"
            f"{r['p50_ms']:>10.3f}ms"
            f"{r['p95_ms']:>10.3f}ms"
            f"{r['p99_ms']:>10.3f}ms"
            f"{r['mean_ms']:>10.3f}ms"
        )


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--n", type=int, default=1000, help="Number of inferences")
    parser.add_argument("--warmup", type=int, default=20)
    parser.add_argument("--model", type=Path, default=DEFAULT_MODEL_PATH)
    parser.add_argument("--onnx", type=Path, default=DEFAULT_ONNX_PATH)
    parser.add_argument(
        "--feature-names", type=Path, default=DEFAULT_FEATURE_NAMES_PATH
    )
    parser.add_argument("--reference", type=Path, default=DEFAULT_REFERENCE_PATH)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument(
        "--out",
        type=Path,
        default=None,
        help="Optional path to write a JSON report (e.g. profiling/benchmark_onnx.json).",
    )
    args = parser.parse_args()

    feature_names = json.loads(args.feature_names.read_text())

    # Sample N rows from the reference dataset, aligned to model feature order.
    reference = pd.read_parquet(args.reference)
    rng = np.random.default_rng(args.seed)
    idx = rng.choice(len(reference), size=args.n + args.warmup, replace=True)
    samples_df = reference.iloc[idx][feature_names].reset_index(drop=True)
    samples_array = samples_df.to_numpy(dtype=np.float32)

    # --- Load both backends ---
    print("Loading LightGBM model...")
    model = joblib.load(args.model)
    raw_model = model.get_raw_model() if hasattr(model, "get_raw_model") else model

    print("Loading ONNX session...")
    session = ort.InferenceSession(
        str(args.onnx), providers=["CPUExecutionProvider"]
    )
    input_name = session.get_inputs()[0].name
    output_name = session.get_outputs()[1].name  # probabilities output

    # --- Warmup (excluded from stats) ---
    print(f"Warmup x{args.warmup}...")
    bench_lightgbm(raw_model, samples_df.iloc[: args.warmup])
    bench_onnx(session, input_name, output_name, samples_array[: args.warmup])

    # --- Measured runs ---
    print(f"Benchmarking LightGBM (n={args.n})...")
    lgbm_lat, lgbm_proba = bench_lightgbm(raw_model, samples_df.iloc[args.warmup :])

    print(f"Benchmarking ONNX     (n={args.n})...")
    onnx_lat, onnx_proba = bench_onnx(
        session, input_name, output_name, samples_array[args.warmup :]
    )

    # --- Numerical equivalence ---
    diff = np.abs(lgbm_proba - onnx_proba)
    print()
    print("=== Numerical equivalence ===")
    print(f"max  |delta_proba| = {diff.max():.2e}")
    print(f"mean |delta_proba| = {diff.mean():.2e}")
    print(f"# rows with |delta| > 1e-5: {(diff > 1e-5).sum()} / {len(diff)}")

    # --- Latency summary ---
    print()
    print("=== Latency (single-row predict_proba) ===")
    rows = [
        summarise("LightGBM", lgbm_lat),
        summarise("ONNX-Runtime", onnx_lat),
    ]
    print_table(rows)

    speedup = median(lgbm_lat) / median(onnx_lat) if median(onnx_lat) > 0 else float("inf")
    gain_pct = (1 - median(onnx_lat) / median(lgbm_lat)) * 100
    print()
    print(f"Speed-up (p50): x{speedup:.2f}")
    print(f"Gain (p50): {gain_pct:.1f}%")

    if args.out is not None:
        report = {
            "config": {
                "n": args.n,
                "warmup": args.warmup,
                "seed": args.seed,
                "model_joblib": str(args.model),
                "model_onnx": str(args.onnx),
                "reference": str(args.reference),
            },
            "equivalence": {
                "max_abs_delta": float(diff.max()),
                "mean_abs_delta": float(diff.mean()),
                "rows_above_1e-5": int((diff > 1e-5).sum()),
                "n_compared": int(len(diff)),
            },
            "latency": rows,
            "speedup_p50": speedup,
            "gain_p50_pct": gain_pct,
        }
        args.out.parent.mkdir(parents=True, exist_ok=True)
        args.out.write_text(json.dumps(report, indent=2))
        print()
        print(f"Wrote report → {args.out}")


if __name__ == "__main__":
    main()