diff options
Diffstat (limited to 'scripts/python/p2_benchmark.py')
| -rw-r--r-- | scripts/python/p2_benchmark.py | 180 |
1 files changed, 180 insertions, 0 deletions
diff --git a/scripts/python/p2_benchmark.py b/scripts/python/p2_benchmark.py new file mode 100644 index 0000000..792673c --- /dev/null +++ b/scripts/python/p2_benchmark.py @@ -0,0 +1,180 @@ +#!/usr/bin/env python3 +from __future__ import annotations + +import argparse +import ctypes +import json +import shutil +import subprocess +import tempfile +import time +from pathlib import Path +from typing import Any + + +class ProcessMemoryCounters(ctypes.Structure): + _fields_ = [ + ("cb", ctypes.c_ulong), + ("PageFaultCount", ctypes.c_ulong), + ("PeakWorkingSetSize", ctypes.c_size_t), + ("WorkingSetSize", ctypes.c_size_t), + ("QuotaPeakPagedPoolUsage", ctypes.c_size_t), + ("QuotaPagedPoolUsage", ctypes.c_size_t), + ("QuotaPeakNonPagedPoolUsage", ctypes.c_size_t), + ("QuotaNonPagedPoolUsage", ctypes.c_size_t), + ("PagefileUsage", ctypes.c_size_t), + ("PeakPagefileUsage", ctypes.c_size_t), + ] + + +class Server: + def __init__(self, exe: Path, data_dir: Path) -> None: + self.process = subprocess.Popen( + [str(exe), "--data-dir", str(data_dir)], + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + encoding="utf-8", + bufsize=1, + ) + self.request_id = 1 + self.max_response_bytes = 0 + self.request( + "initialize", + { + "protocolVersion": "2024-11-05", + "capabilities": {}, + "clientInfo": {"name": "p2-benchmark", "version": "1"}, + }, + ) + self.notify("notifications/initialized", {}) + + def request(self, method: str, params: dict[str, Any]) -> dict[str, Any]: + assert self.process.stdin and self.process.stdout + request_id = self.request_id + self.request_id += 1 + self.process.stdin.write( + json.dumps( + {"jsonrpc": "2.0", "id": request_id, "method": method, "params": params}, + separators=(",", ":"), + ) + + "\n" + ) + self.process.stdin.flush() + line = self.process.stdout.readline() + self.max_response_bytes = max(self.max_response_bytes, len(line.encode("utf-8"))) + response = json.loads(line) + if "error" in response: + raise RuntimeError(response) + return response["result"] + + def call(self, name: str, arguments: dict[str, Any]) -> dict[str, Any]: + result = self.request("tools/call", {"name": name, "arguments": arguments}) + payload = json.loads(result["content"][0]["text"]) + if result.get("isError"): + raise RuntimeError(payload) + return payload + + def notify(self, method: str, params: dict[str, Any]) -> None: + assert self.process.stdin + self.process.stdin.write( + json.dumps({"jsonrpc": "2.0", "method": method, "params": params}) + "\n" + ) + self.process.stdin.flush() + + def peak_working_set(self) -> int: + counters = ProcessMemoryCounters() + counters.cb = ctypes.sizeof(counters) + handle = ctypes.windll.kernel32.OpenProcess(0x1000 | 0x0400, False, self.process.pid) + try: + if handle and ctypes.windll.psapi.GetProcessMemoryInfo( + handle, ctypes.byref(counters), counters.cb + ): + return int(counters.PeakWorkingSetSize) + return 0 + finally: + if handle: + ctypes.windll.kernel32.CloseHandle(handle) + + def close(self) -> None: + if self.process.stdin: + self.process.stdin.close() + code = self.process.wait(timeout=30) + if code != 0: + raise RuntimeError(self.process.stderr.read() if self.process.stderr else code) + + +def read_all(server: Server, session_id: str) -> tuple[int, int]: + cursor = None + count = pages = 0 + while True: + arguments: dict[str, Any] = {"session_id": session_id, "limit": 1000} + if cursor: + arguments["cursor"] = cursor + payload = server.call("read_messages", arguments) + count += len(payload["history"]) + pages += 1 + cursor = payload.get("nextCursor") + if not cursor: + return count, pages + + +def main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument("--server", type=Path, required=True) + parser.add_argument("--messages", type=int, choices=(10000, 50000), required=True) + parser.add_argument("--sessions", type=int, choices=(1, 100), required=True) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + data_dir = Path(tempfile.mkdtemp(prefix="tiny-convmem-p2-")) + try: + server = Server(args.server.resolve(), data_dir) + write_started = time.perf_counter() + for index in range(args.messages): + server.call( + "write_message", + { + "session_id": f"session-{index % args.sessions}", + "role": "user", + "content": f"history item {index}", + }, + ) + write_ms = (time.perf_counter() - write_started) * 1000 + server.close() + + restart_started = time.perf_counter() + server = Server(args.server.resolve(), data_dir) + restart_ms = (time.perf_counter() - restart_started) * 1000 + read_started = time.perf_counter() + total = pages = 0 + for session in range(args.sessions): + session_count, session_pages = read_all(server, f"session-{session}") + total += session_count + pages += session_pages + read_ms = (time.perf_counter() - read_started) * 1000 + rebuild_started = time.perf_counter() + server.call("rebuild_index", {}) + rebuild_ms = (time.perf_counter() - rebuild_started) * 1000 + result = { + "messages": args.messages, + "sessions": args.sessions, + "messages_read": total, + "pages": pages, + "write_ms": write_ms, + "restart_ms": restart_ms, + "full_read_ms": read_ms, + "rebuild_ms": rebuild_ms, + "peak_working_set_bytes": server.peak_working_set(), + "max_response_bytes": server.max_response_bytes, + } + server.close() + args.output.write_text(json.dumps(result, indent=2), encoding="utf-8") + print(json.dumps(result, indent=2)) + return 0 if total == args.messages and server.max_response_bytes <= 8 * 1024 * 1024 else 1 + finally: + shutil.rmtree(data_dir, ignore_errors=True) + + +if __name__ == "__main__": + raise SystemExit(main()) |