litellm/scripts/benchmark_ocr_callbacks.py

#!/usr/bin/env python3
"""Measure serial sync/async OCR latency through a loopback HTTP provider

Run each callback mode in a fresh process against an installed release wheel:
python -I scripts/benchmark_ocr_callbacks.py --callbacks none --label before \
    --expected-transport rust --iterations 200 --warmup 20 --output before-none.json
Repeat with --callbacks noop and with the candidate wheel in a separate venv
"""

from __future__ import annotations

import argparse
import asyncio
import base64
import hashlib
import importlib.metadata
import json
import statistics
import sys
import threading
import time
from collections.abc import Sequence
from dataclasses import asdict, dataclass
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
from typing import Final

SIZES: Final = (
    1024,
    4 * 1024,
    16 * 1024,
    64 * 1024,
    256 * 1024,
    1024 * 1024,
)
MODEL: Final = "mistral/mistral-ocr-latest"
EXPECTED_MARKDOWN: Final = "mock remote OCR response"
RESPONSE: Final = json.dumps(
    {
        "pages": [{"index": 0, "markdown": EXPECTED_MARKDOWN, "images": [], "dimensions": None}],
        "model": "mistral-ocr-latest",
        "usage_info": {"pages_processed": 1},
    },
    separators=(",", ":"),
).encode()


class Server(ThreadingHTTPServer):
    daemon_threads = True

    def __init__(self) -> None:
        super().__init__(("127.0.0.1", 0), Handler)


class Handler(BaseHTTPRequestHandler):
    protocol_version = "HTTP/1.1"

    def do_POST(self) -> None:
        length: Final = int(self.headers["Content-Length"])
        body: Final = self.rfile.read(length)
        request: Final = json.loads(body)
        if self.path != "/v1/ocr" or request.get("model") != "mistral-ocr-latest":
            self.send_error(400)
            return
        document: Final = request.get("document", {})
        if not isinstance(document, dict) or not str(document.get("document_url", "")).startswith(
            "data:application/pdf;base64,"
        ):
            self.send_error(400)
            return
        self.send_response(200)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(RESPONSE)))
        self.end_headers()
        self.wfile.write(RESPONSE)

    def log_message(self, format: str, *args: object) -> None:
        return


@dataclass(frozen=True, slots=True)
class Result:
    label: str
    mode: str
    size: int
    iterations: int
    median_ms: float
    mean_ms: float
    p95_ms: float
    requests_per_second: float


def document(size: int) -> dict[str, str]:
    payload: Final = b"%PDF-1.4\n" + b"x" * max(0, size - 9)
    encoded: Final = base64.b64encode(payload[:size]).decode("ascii")
    return {"type": "document_url", "document_url": f"data:application/pdf;base64,{encoded}"}


def percentile(values: Sequence[float], quantile: float) -> float:
    ordered: Final = sorted(values)
    index: Final = min(len(ordered) - 1, round((len(ordered) - 1) * quantile))
    return ordered[index]


def verify(response: object, expected_transport: str) -> None:
    from litellm.litellm_core_utils.add_retry_fallback_headers import get_hidden_params_dict

    pages: Final = getattr(response, "pages", ())
    if len(pages) != 1 or getattr(pages[0], "markdown", None) != EXPECTED_MARKDOWN:
        raise RuntimeError(f"unexpected OCR response: {response!r}")
    additional_headers: Final = get_hidden_params_dict(response).get("additional_headers")
    rust_served: Final = isinstance(additional_headers, dict) and additional_headers.get("x-litellm-rust") == "true"
    if rust_served != (expected_transport == "rust"):
        raise RuntimeError(f"unexpected transport for {expected_transport}: rust_marker={rust_served}")


def summarize(label: str, mode: str, size: int, samples: Sequence[float]) -> Result:
    median: Final = statistics.median(samples)
    return Result(
        label=label,
        mode=mode,
        size=size,
        iterations=len(samples),
        median_ms=median * 1000,
        mean_ms=statistics.fmean(samples) * 1000,
        p95_ms=percentile(samples, 0.95) * 1000,
        requests_per_second=1 / median,
    )


def sync_samples(
    litellm: object, url: str, request_document: dict[str, str], count: int, expected_transport: str
) -> tuple[float, ...]:
    samples: list[float] = []
    for _ in range(count):
        started: Final = time.perf_counter()
        response: Final = litellm.ocr(
            model=MODEL, document=request_document, api_base=url, api_key="mock-key", timeout=30
        )
        samples.append(time.perf_counter() - started)
        verify(response, expected_transport)
    return tuple(samples)


async def async_samples(
    litellm: object, url: str, request_document: dict[str, str], count: int, expected_transport: str
) -> tuple[float, ...]:
    samples: list[float] = []
    for _ in range(count):
        started: Final = time.perf_counter()
        response: Final = await litellm.aocr(
            model=MODEL, document=request_document, api_base=url, api_key="mock-key", timeout=30
        )
        samples.append(time.perf_counter() - started)
        verify(response, expected_transport)
    return tuple(samples)


async def main() -> int:
    parser: Final = argparse.ArgumentParser(description="E2E OCR benchmark against a local remote-style HTTP server")
    parser.add_argument("--callbacks", choices=("none", "noop"), required=True)
    parser.add_argument("--label", required=True)
    parser.add_argument("--expected-transport", choices=("python", "rust"), required=True)
    parser.add_argument("--iterations", type=int, default=30)
    parser.add_argument("--warmup", type=int, default=5)
    parser.add_argument("--sizes", type=int, nargs="+", default=SIZES)
    parser.add_argument("--output", type=Path, required=True)
    args: Final = parser.parse_args()

    import litellm
    from litellm.integrations.custom_logger import CustomLogger

    class NoopCallback(CustomLogger):
        def __init__(self) -> None:
            super().__init__()
            self.pre_calls = 0
            self.sync_successes = 0
            self.async_successes = 0

        def log_pre_api_call(self, model, messages, kwargs):
            self.pre_calls += 1

        def log_success_event(self, kwargs, response_obj, start_time, end_time):
            self.sync_successes += 1

        async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
            self.async_successes += 1

    registry_names: Final = (
        "callbacks",
        "input_callback",
        "success_callback",
        "failure_callback",
        "_async_input_callback",
        "_async_success_callback",
        "_async_failure_callback",
    )
    if any(getattr(litellm, name) for name in registry_names):
        raise RuntimeError("benchmark requires initially empty callback registrations")
    callback: Final = NoopCallback()
    if args.callbacks == "noop":
        litellm.callbacks.append(callback)

    rust_toggle: Final = getattr(litellm, "rust", None)
    if callable(rust_toggle):
        rust_toggle(False)
    package: Final = Path(litellm.__file__).resolve()
    version: Final = importlib.metadata.version("litellm")
    native_path: str | None = None
    native_sha256: str | None = None
    try:
        from litellm.rust_bridge import _native

        native: Final = Path(_native.__file__).resolve()
        native_path = str(native)
        native_sha256 = hashlib.file_digest(native.open("rb"), "sha256").hexdigest()
    except ImportError:
        pass

    server: Final = Server()
    thread: Final = threading.Thread(target=server.serve_forever, daemon=True)
    thread.start()
    url: Final = f"http://127.0.0.1:{server.server_port}"
    results: list[Result] = []
    try:
        for size in args.sizes:
            request_document: Final = document(size)
            sync_samples(litellm, url, request_document, args.warmup, args.expected_transport)
            sync_result: Final = summarize(
                args.label,
                "sync",
                size,
                sync_samples(litellm, url, request_document, args.iterations, args.expected_transport),
            )
            results.append(sync_result)
            await async_samples(litellm, url, request_document, args.warmup, args.expected_transport)
            async_result: Final = summarize(
                args.label,
                "async",
                size,
                await async_samples(litellm, url, request_document, args.iterations, args.expected_transport),
            )
            results.append(async_result)
            sys.stdout.write(json.dumps(asdict(sync_result)) + "\n")
            sys.stdout.write(json.dumps(asdict(async_result)) + "\n")
            sys.stdout.flush()
    finally:
        server.shutdown()
        server.server_close()
        thread.join()

    from litellm.litellm_core_utils.litellm_logging import executor
    from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER

    await GLOBAL_LOGGING_WORKER.flush()
    await asyncio.to_thread(executor.shutdown, wait=True)
    per_mode: Final = len(args.sizes) * (args.iterations + args.warmup)
    if args.callbacks == "noop":
        if (callback.pre_calls, callback.sync_successes, callback.async_successes) != (
            2 * per_mode,
            per_mode,
            per_mode,
        ):
            raise RuntimeError(f"callback delivery mismatch: {vars(callback)}")
    elif any(getattr(litellm, name) for name in registry_names):
        raise RuntimeError("callback registrations appeared in the no-callback case")
    await GLOBAL_LOGGING_WORKER.stop()

    artifact: Final = {
        "label": args.label,
        "callbacks": args.callbacks,
        "python": sys.executable,
        "callback_counts": {
            "pre": callback.pre_calls,
            "sync_success": callback.sync_successes,
            "async_success": callback.async_successes,
        },
        "version": version,
        "package": str(package),
        "native": native_path,
        "native_sha256": native_sha256,
        "results": tuple(asdict(result) for result in results),
    }
    args.output.write_text(json.dumps(artifact, indent=2) + "\n")
    sys.stdout.write(json.dumps({key: artifact[key] for key in ("label", "version", "package")}) + "\n")
    sys.stdout.write(f"results={args.output}\n")
    sys.stdout.flush()
    return 0


if __name__ == "__main__":
    raise SystemExit(asyncio.run(main()))