examples/benchmark_sidecar_cpu.py

from __future__ import annotations

import argparse
import hashlib
import json
import platform
import statistics
import sys
import time
from collections import Counter
from dataclasses import dataclass
from pathlib import Path

from deberta_ime import CandidateSource, SidecarClient
from deberta_ime.deberta import MODEL_ID, MODEL_REVISION
from deberta_ime.mozc import MOZC_REVISION, MozcDictionaryIndex
from deberta_ime.profiles import MOZC_SIDECAR_SOURCE_ID


@dataclass(frozen=True, slots=True)
class BenchmarkCase:
    reading: str
    candidates: tuple[dict[str, object], ...]
    left_context: tuple[str, ...]
    right_context: tuple[str, ...]


def _parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description=(
            "Measure warmup-excluded CPU latency through the real local sidecar process."
        )
    )
    parser.add_argument("--index", type=Path, required=True)
    parser.add_argument("--model-cache-dir", type=Path, required=True)
    parser.add_argument("--output", type=Path)
    parser.add_argument("--iterations", type=int, default=50)
    parser.add_argument("--warm-requests", type=int, default=5)
    parser.add_argument("--request-timeout-seconds", type=float, default=2.0)
    parser.add_argument("--source-commit", default="unrecorded")
    return parser


def _percentile(values: list[float], fraction: float) -> float:
    ordered = sorted(values)
    if not ordered:
        return 0.0
    position = (len(ordered) - 1) * fraction
    lower = int(position)
    upper = min(lower + 1, len(ordered) - 1)
    remainder = position - lower
    return ordered[lower] * (1.0 - remainder) + ordered[upper] * remainder


def _sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as stream:
        while chunk := stream.read(1024 * 1024):
            digest.update(chunk)
    return digest.hexdigest()


def _candidate_records(
    index: MozcDictionaryIndex, reading: str
) -> tuple[dict[str, object], ...]:
    return tuple(
        {"surface": item.surface, "prior_score": item.prior_score}
        for item in index.lookup(reading, limit=8)
    )


def run(argv: list[str] | None = None) -> int:
    args = _parser().parse_args(argv)
    if args.iterations < 1:
        raise ValueError("--iterations must be positive")
    if args.warm_requests < 0:
        raise ValueError("--warm-requests must not be negative")
    index_path = args.index.resolve()
    cache_path = args.model_cache_dir.resolve()
    with MozcDictionaryIndex(index_path) as index:
        cases = {
            "incremental": BenchmarkCase(
                reading="きょう",
                candidates=_candidate_records(index, "きょう"),
                left_context=("予定", "は"),
                right_context=(),
            ),
            "bidirectional": BenchmarkCase(
                reading="はし",
                candidates=_candidate_records(index, "はし"),
                left_context=("川", "に", "架かる"),
                right_context=("を", "渡る"),
            ),
        }
    if any(len(case.candidates) < 2 for case in cases.values()):
        raise ValueError("benchmark readings need at least two indexed candidates")

    command = (
        sys.executable,
        "-m",
        "deberta_ime",
        "serve",
        "--offline",
        "--cache-dir",
        str(cache_path),
    )
    source = CandidateSource(MOZC_SIDECAR_SOURCE_ID, MOZC_REVISION)
    measured: dict[str, object] = {}
    with SidecarClient(command, request_timeout_seconds=args.request_timeout_seconds) as client:
        health_before = client.health(timeout_seconds=10.0)
        warmup_started = time.perf_counter()
        warmup = client.warmup(timeout_seconds=120.0)
        warmup_ms = (time.perf_counter() - warmup_started) * 1000.0
        if not health_before.ok or not warmup.ok or not warmup.model_loaded:
            raise RuntimeError(
                f"sidecar setup failed: health={health_before.reason}, warmup={warmup.reason}"
            )
        for mode, case in cases.items():
            for _ in range(args.warm_requests):
                client.rerank(
                    reading=case.reading,
                    candidates=case.candidates,
                    candidate_source=source,
                    mode=mode,
                    left_context=case.left_context,
                    right_context=case.right_context,
                )
            latencies: list[float] = []
            reasons: Counter[str] = Counter()
            changed = 0
            measured_started = time.perf_counter()
            for _ in range(args.iterations):
                request_started = time.perf_counter()
                result = client.rerank(
                    reading=case.reading,
                    candidates=case.candidates,
                    candidate_source=source,
                    mode=mode,
                    left_context=case.left_context,
                    right_context=case.right_context,
                )
                latencies.append((time.perf_counter() - request_started) * 1000.0)
                reasons[result.reason] += 1
                changed += int(result.changed)
            elapsed = time.perf_counter() - measured_started
            measured[mode] = {
                "iterations": args.iterations,
                "excluded_warm_requests": args.warm_requests,
                "candidate_count": len(case.candidates),
                "elapsed_seconds": elapsed,
                "requests_per_second": args.iterations / elapsed,
                "latency_ms": {
                    "mean": statistics.fmean(latencies),
                    "p50": _percentile(latencies, 0.50),
                    "p95": _percentile(latencies, 0.95),
                    "min": min(latencies),
                    "max": max(latencies),
                },
                "changed": changed,
                "reasons": dict(sorted(reasons.items())),
            }
        health_after = client.health(timeout_seconds=10.0)

    payload = {
        "schema": "deberta-ime-sidecar-cpu-v1",
        "evidence_state": "LOCAL_PASS",
        "source_commit": args.source_commit,
        "protocol": {
            "transport": "local-child-stdio",
            "encoding": "utf-8",
            "schema_version": 1,
        },
        "runtime": {
            "platform": platform.platform(),
            "python": platform.python_version(),
            "processor": platform.processor(),
            "device": "cpu",
            "offline": True,
        },
        "model": {"id": MODEL_ID, "revision": MODEL_REVISION},
        "candidate_source": {
            "id": MOZC_SIDECAR_SOURCE_ID,
            "revision": MOZC_REVISION,
        },
        "index": {
            "path": args.index.as_posix(),
            "size": index_path.stat().st_size,
            "sha256": _sha256(index_path),
        },
        "health_before": {
            "ok": health_before.ok,
            "model_loaded": health_before.model_loaded,
            "reason": health_before.reason,
        },
        "warmup": {
            "ok": warmup.ok,
            "model_loaded": warmup.model_loaded,
            "reason": warmup.reason,
            "elapsed_ms": warmup_ms,
        },
        "steady_state": measured,
        "health_after": {
            "ok": health_after.ok,
            "model_loaded": health_after.model_loaded,
            "reason": health_after.reason,
        },
        "claims": {
            "includes_model_warmup_in_latency": False,
            "same_process_survived_all_measured_requests": health_after.model_loaded,
            "tsf_or_device_tested": False,
            "public_pass": False,
            "human_go": False,
        },
    }
    serialized = json.dumps(payload, ensure_ascii=False, indent=2) + "\n"
    if args.output is None:
        sys.stdout.reconfigure(encoding="utf-8")
        sys.stdout.write(serialized)
    else:
        args.output.parent.mkdir(parents=True, exist_ok=True)
        args.output.write_text(serialized, encoding="utf-8", newline="\n")
    return 0


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