searchi/tools/embeddinggemma2/query_context_benchmark.py

"""CPU-only paired real-worker context startup/memory measurements.

No downloads or GPU execution. Windows reports private commitment and resident
working set separately; Linux reports PSS/private resident pages. Neither should
be summed as model allocations across workers. Requires an idle reference host.
"""
import argparse
import ctypes
import hashlib
import json
import math
import os
from pathlib import Path
import platform
import statistics
import time
from benchmark import ROOT, Session, vectors


def memory(pid):
    if os.name == "nt":
        from ctypes import wintypes
        class Counters(ctypes.Structure):
            _fields_ = [("cb", wintypes.DWORD), ("PageFaultCount", wintypes.DWORD)] + [(name, ctypes.c_size_t) for name in ["PeakWorkingSetSize", "WorkingSetSize", "QuotaPeakPagedPoolUsage", "QuotaPagedPoolUsage", "QuotaPeakNonPagedPoolUsage", "QuotaNonPagedPoolUsage", "PagefileUsage", "PeakPagefileUsage", "PrivateUsage"]]
        kernel = ctypes.WinDLL("kernel32", use_last_error=True)
        kernel.OpenProcess.restype = wintypes.HANDLE
        kernel.OpenProcess.argtypes = [wintypes.DWORD, wintypes.BOOL, wintypes.DWORD]
        kernel.CloseHandle.argtypes = [wintypes.HANDLE]
        psapi = ctypes.WinDLL("psapi", use_last_error=True)
        psapi.GetProcessMemoryInfo.argtypes = [wintypes.HANDLE, ctypes.POINTER(Counters), wintypes.DWORD]
        handle = kernel.OpenProcess(0x1000 | 0x10, False, pid)
        if not handle:
            raise ctypes.WinError(ctypes.get_last_error())
        try:
            info = Counters()
            info.cb = ctypes.sizeof(info)
            if not psapi.GetProcessMemoryInfo(handle, ctypes.byref(info), info.cb):
                raise ctypes.WinError(ctypes.get_last_error())
            return {"private_commit_bytes": info.PrivateUsage, "working_set_bytes": info.WorkingSetSize, "page_fault_count": info.PageFaultCount}
        finally:
            kernel.CloseHandle(handle)
    values = {}
    for line in Path(f"/proc/{pid}/smaps_rollup").read_text().splitlines():
        parts = line.split()
        if parts[0] in {"Pss:", "Private_Clean:", "Private_Dirty:", "Rss:"}:
            values[parts[0][:-1] + "_bytes"] = int(parts[1]) * 1024
    return values


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--worker", type=Path, required=True)
    parser.add_argument("--models", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--repeats", type=int, default=3)
    args = parser.parse_args()
    if args.output.exists():
        parser.error("output must be new; preserve previous evidence")
    if not 1 <= args.repeats <= 20:
        parser.error("repeats must be 1..20")
    manifest = json.loads((ROOT / "tools/embeddinggemma2/model-manifest.json").read_text())
    entry = manifest["files"][0]
    model = args.models / entry["path"]
    started = time.perf_counter()
    assert model.stat().st_size == entry["bytes"]
    with model.open("rb") as stream:
        assert hashlib.file_digest(stream, "sha256").hexdigest() == entry["sha256"]
    verification_ms = (time.perf_counter() - started) * 1000
    samples = []
    for repeat in range(args.repeats):
        # Alternate ordering to expose warm filesystem/cache order effects.
        for role in (["indexing", "query"] if repeat % 2 == 0 else ["query", "indexing"]):
            session = Session([str(args.worker.resolve()), "--model", str(model.resolve()), "--threads", "1", "--device", "cpu", "--role", role])
            try:
                hello = session.next()
                startup_ms = (time.perf_counter() - session.started) * 1000
                assert hello["device"] == "cpu" and not hello["vulkan_module_loaded"]
                assert hello["role"] == role and hello["context_tokens"] == (512 if role == "query" else 2048)
                sample = {"repeat": repeat, "role": role, "hello": hello, "startup_ms": startup_ms, "ready_memory": memory(session.process.pid)}
                response, first_ms = session.worker({"op": "query", "text": "cancel the supplier partnership", "debug_performance": True})
                vectors(response, 1)
                sample.update(first_query_ms=first_ms, worker=response["performance"], after_query_memory=memory(session.process.pid))
                warm = []
                # Distinct inputs exercise native query inference without a Rust
                # query cache. Keep first input separate from these warmed calls.
                for index in range(20):
                    response, duration = session.worker({"op": "query", "text": f"supplier cancellation notice example {index}", "debug_performance": True})
                    vectors(response, 1)
                    warm.append({"round_trip_ms": duration, "worker": response["performance"]})
                ordered = sorted(item["round_trip_ms"] for item in warm)
                sample["warm_queries"] = {"samples": warm, "p95_ms": ordered[math.ceil(len(ordered) * .95) - 1]}
                if role == "query":
                    session.send({"op": "documents", "texts": ["test"], "title": "none"})
                    rejected = session.next()
                    assert not rejected["ok"] and not rejected["fatal"]
                    assert "only supports" in rejected["error"]
                    vectors(session.worker({"op": "code_query", "text": "retry failed requests"})[0], 1)
                samples.append(sample)
            finally:
                session.close()
    report = {"platform": platform.platform(), "cpu_count": os.cpu_count(), "verification_ms": verification_ms, "cache_note": "Cached model bytes; startup excludes Rust verification and desktop preparation. No GPU executed. Repeated fresh processes, not cold disk.", "worker_sha256": hashlib.sha256(args.worker.read_bytes()).hexdigest(), "samples": samples}
    args.output.parent.mkdir(parents=True, exist_ok=True)
    with args.output.open("x", encoding="utf-8") as output:
        output.write(json.dumps(report, indent=2) + "\n")
    for role in ["indexing", "query"]:
        rows = [s for s in samples if s["role"] == role]
        print(role, "startup median ms", statistics.median(s["startup_ms"] for s in rows), "ready memory", rows[0]["ready_memory"], "after query", rows[0]["after_query_memory"])


if __name__ == "__main__":
    main()