tgrep-ai-skill/scripts/benchmark.py

#!/usr/bin/env python3
"""Measure full CLI indexed/fresh latency on deterministic synthetic corpora."""

import argparse
import hashlib
import json
import math
import os
from pathlib import Path
import platform
import random
import shutil
import statistics
import subprocess
import sys
import tempfile
import time
from typing import Any

SOURCE = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(SOURCE / "src"))
from local_search import config, engine  # noqa: E402

QUERIES = [
    ("rare_literal", "unique_benchmark_needle", False),
    ("distributed_literal", "batch_benchmark_marker", False),
    ("absent_literal", "absent_benchmark_xyz987", False),
    ("selective_regex", "unique_benchmark_[a-z]+", True),
    ("short_absent_literal", "ZQ", False),
]


def run(command: list[str], env: dict[str, str] | None = None) -> str:
    return subprocess.run(
        command, capture_output=True, text=True, env=env, timeout=120, check=True
    ).stdout.strip()


def corpus(path: Path, count: int) -> dict[str, Any]:
    """Create roughly 4 KiB per file, with known sparse and distributed matches."""
    path.mkdir()
    digest = hashlib.sha256()
    total = 0
    for number in range(count):
        folder = path / f"package_{number // 100:05d}"
        folder.mkdir(exist_ok=True)
        lines = [f"# synthetic module {number}\n"]
        if number == count // 2:
            lines.append("# unique_benchmark_needle\n")
        if number % 100 == 0:
            lines.append("# batch_benchmark_marker\n")
        for item in range(64):
            lines.append(
                f"def function_{number}_{item}(value): "
                f"return value + {number * 64 + item}\n"
            )
        data = "".join(lines).encode()
        (folder / f"module_{number:06d}.py").write_bytes(data)
        digest.update(data)
        total += len(data)
    return {"files": count, "bytes": total, "content_sha256": digest.hexdigest()}


def measure(
    env: dict[str, str], pattern: str, regex: bool, fresh: bool,
) -> tuple[float, list[tuple[str, int, str]]]:
    command = [
        sys.executable, "-m", "local_search.cli", "search", pattern,
        "--root", "benchmark", "--limit", "1000",
    ]
    if regex:
        command.append("--regex")
    if fresh:
        command.append("--fresh")
    started = time.perf_counter()
    result = subprocess.run(
        command, env=env, capture_output=True, text=True, timeout=120
    )
    elapsed = (time.perf_counter() - started) * 1000
    if result.returncode not in (0, 1):
        raise RuntimeError(f"Search failed: {result.stderr} {result.stdout}")
    response = json.loads(result.stdout)
    if response.get("truncated"):
        raise RuntimeError("Truncated output cannot establish exact match parity")
    reports = response["reports"]
    expected = "rg" if fresh else "tgrep"
    for report in reports:
        if report["backend"] != expected or report["warnings"]:
            raise RuntimeError(f"Unexpected backend or warning: {report}")
    matches = sorted(
        (item["path"], item["line"], item["text"])
        for item in response["matches"]
    )
    return elapsed, matches


def summarize(samples: list[float]) -> dict[str, Any]:
    ordered = sorted(samples)
    return {
        "samples_ms": samples,
        "median_ms": statistics.median(samples),
        "p95_ms": ordered[math.ceil(len(ordered) * 0.95) - 1],
        "min_ms": min(samples),
        "max_ms": max(samples),
    }


def benchmark(
    count: int, rounds: int, warmups: int, binaries: dict[str, str],
) -> dict[str, Any]:
    base = Path(tempfile.mkdtemp(prefix="local-search-benchmark-"))
    safe_to_remove = True
    try:
        metadata = corpus(base / "corpus", count)
        env = {
            **os.environ,
            "XDG_CONFIG_HOME": str(base / "config"),
            "XDG_DATA_HOME": str(base / "data"),
            "PYTHONPATH": str(SOURCE / "src"),
            **binaries,
        }
        previous = os.environ.copy()
        os.environ.update(env)
        root = None
        try:
            root = config.add_root("benchmark", base / "corpus")
            started = time.perf_counter()
            server = engine.start(root, rebuild=True)
            deadline = time.monotonic() + 120
            while (
                server.get("indexing", True)
                or server.get("reconcile_running")
                or server.get("reconcile_pending")
            ):
                if time.monotonic() > deadline:
                    raise RuntimeError("Index did not become ready")
                time.sleep(0.1)
                server = engine.rpc_status(config.index_dir(root)) or {}
            build_seconds = time.perf_counter() - started
            if server["num_files"] != count:
                raise RuntimeError("Unexpected indexed file count")
            results = []
            rng = random.Random(20260909)
            for name, pattern, regex in QUERIES:
                timings = {False: [], True: []}
                reference = None
                for iteration in range(warmups + rounds):
                    order = [False, True]
                    rng.shuffle(order)
                    for fresh in order:
                        elapsed, matches = measure(env, pattern, regex, fresh)
                        if reference is None:
                            reference = matches
                        if reference != matches:
                            raise RuntimeError(f"Match mismatch for {name}")
                        if iteration >= warmups:
                            timings[fresh].append(elapsed)
                expected_count = (
                    1 if name in ("rare_literal", "selective_regex")
                    else (count + 99) // 100 if name == "distributed_literal"
                    else 0
                )
                if len(reference) != expected_count:
                    raise RuntimeError(f"Unexpected match count for {name}")
                indexed = summarize(timings[False])
                fresh = summarize(timings[True])
                saving = (fresh["median_ms"] - indexed["median_ms"]) / 1000
                results.append({
                    "query": name, "pattern": pattern, "regex": regex,
                    "matches": len(reference), "exact_match_parity": True,
                    "indexed": indexed, "fresh": fresh,
                    "speedup": fresh["median_ms"] / indexed["median_ms"],
                    "build_break_even_queries": (
                        math.ceil(build_seconds / saving) if saving > 0 else None
                    ),
                })
                print(
                    f"{count} files / {name}: "
                    f"{indexed['median_ms']:.1f} vs {fresh['median_ms']:.1f} ms",
                    file=sys.stderr, flush=True,
                )
            directory = config.index_dir(root)
            index_bytes = sum(p.stat().st_size for p in directory.rglob("*")
                              if p.is_file())
            rss_kib = int(run(["ps", "-p", str(server["pid"]), "-o", "rss="]))
            return {
                **metadata, "build_and_start_seconds": build_seconds,
                "index_directory_bytes": index_bytes,
                "server_rss_kib_after_queries": rss_kib, "queries": results,
            }
        finally:
            try:
                if root is not None:
                    safe_to_remove = False
                    try:
                        engine.stop(root)
                    except BaseException:
                        print(f"Stop failed; preserving diagnostics: {base}",
                              file=sys.stderr)
                        raise
                    safe_to_remove = True
            finally:
                os.environ.clear()
                os.environ.update(previous)
    finally:
        if safe_to_remove:
            shutil.rmtree(base)


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--tgrep", required=True, type=Path)
    parser.add_argument("--rg", required=True, type=Path)
    parser.add_argument("--sizes", nargs="+", type=int, default=[1000, 20000])
    parser.add_argument("--rounds", type=int, default=9)
    parser.add_argument("--warmups", type=int, default=2)
    parser.add_argument("--output", required=True, type=Path)
    args = parser.parse_args()
    if (args.rounds < 3 or args.warmups < 1
            or any(n < 1 or n > 100000 for n in args.sizes)):
        parser.error("Use rounds >= 3, warmups >= 1 and sizes in 1..100000")
    binaries = {
        "LOCAL_SEARCH_TGREP": str(args.tgrep.resolve(strict=True)),
        "LOCAL_SEARCH_RG": str(args.rg.resolve(strict=True)),
    }
    report = {
        "schema": 1, "utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        "platform": platform.platform(), "machine": platform.machine(),
        "python": platform.python_version(), "logical_cpus": os.cpu_count(),
        "cpu_model": run(["sysctl", "-n", "machdep.cpu.brand_string"])
        if sys.platform == "darwin" else platform.processor(),
        "memory_bytes": int(run(["sysctl", "-n", "hw.memsize"]))
        if sys.platform == "darwin" else None,
        "source_commit": run(["git", "-C", str(SOURCE), "rev-parse", "HEAD"]),
        "benchmark_sha256": hashlib.sha256(Path(__file__).read_bytes()).hexdigest(),
        "tgrep_version": run([binaries["LOCAL_SEARCH_TGREP"], "--version"]),
        "rg_version": run([binaries["LOCAL_SEARCH_RG"], "--version"]),
        "rounds": args.rounds, "warmups": args.warmups,
        "method": "Full CLI wall time; paired randomized order; warm OS caches; "
                  "synthetic Python; identical filters; exact uncapped match parity; "
                  "no LLM; RSS is a point sample, not a peak; index build is one run",
        "corpora": [],
    }
    for count in args.sizes:
        report["corpora"].append(
            benchmark(count, args.rounds, args.warmups, binaries)
        )
        args.output.parent.mkdir(parents=True, exist_ok=True)
        args.output.write_text(json.dumps(report, indent=2) + "\n")


if __name__ == "__main__":
    main()