Use measured query-only context across desktop CLI and MCP

tools/embeddinggemma2/query_context_benchmark.py

@@ -0,0 +1,98 @@
+"""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 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 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))
+                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)
+    args.output.write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8")
+    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()