benchmark.py

#!/usr/bin/env python3
"""
TurboQuant 4-bit disk-rescore benchmark for wiki DPR / e5.

Sweeps hnsw_ef over {50,100,150,200,256} with oversampling=1 for both rescore patterns:
  Pattern A  single-stage rescore
  Pattern B  two-stage prefetch + bounded rescore

Each operating point runs recall@100, closed-loop load (concurrency 4, 120 s), and
open-loop load (32 QPS target).

Run:
  export QDRANT_URL=... QDRANT_API_KEY=... ASYNC_SCORER=true
  python3 benchmark.py
  # or: scripts/run_profile.sh qdrant_2vcpu8gb
"""

import os, csv, json, time
import urllib.request
from pathlib import Path
from qdrant_client import QdrantClient, models
from queryset import load_queries
from runengine import SearchConfig, measure_recall, closed_loop, open_loop, LIMIT

COLLECTION = os.getenv("COLLECTION", "wiki_dpr_e5")
QDRANT_URL = os.environ["QDRANT_URL"]
QDRANT_API_KEY = os.getenv("QDRANT_API_KEY")
PREFER_GRPC = os.getenv("PREFER_GRPC", "0") == "1"
QUERIES_PATH = os.getenv("QUERIES", "queries.parquet")
RUN_PROFILE = os.getenv("RUN_PROFILE", "qdrant_benchmark")
N_QUERIES = int(os.getenv("N_QUERIES", "10000"))
WARMUP = int(os.getenv("WARMUP", "100"))
CONCURRENCY = int(os.getenv("CONCURRENCY", "4"))
DURATION = float(os.getenv("DURATION", "120"))
OPEN_QPS = float(os.getenv("OPEN_QPS", "32"))
ASYNC_SCORER = os.getenv("ASYNC_SCORER", "unknown")
HNSW_EFS = [int(x) for x in os.getenv("HNSW_EFS", "50,100,150,200,256").split(",")]
OVERSAMPLINGS = [float(x) for x in os.getenv("OVERSAMPLINGS", "1").split(",")]
OUT_DIR = Path(os.getenv("OUT_DIR", "results"))
_REST_FALLBACK_WARNED = False

DIM = 768
N_POINTS = 21_015_300


def get_client():
    return QdrantClient(
        url=QDRANT_URL, api_key=QDRANT_API_KEY, prefer_grpc=PREFER_GRPC, timeout=600
    )


def wait_green(client, label):
    while True:
        info = get_collection_info(client)
        status = info_value(info, "status")
        points = info_value(info, "points_count")
        indexed = info_value(info, "indexed_vectors_count")
        print(f"  [{label}] {status}  pts={points}  indexed={indexed}")
        if str(status).lower().endswith("green"):
            return info
        time.sleep(15)


def set_turboquant(client):
    config = {"turbo": {"always_ram": True, "bits": "bits4"}}
    print("\n[quant] switching to TurboQuant 4-bit ...")
    if current_quantization_matches(client, config):
        print("[quant] TurboQuant 4-bit already active; waiting for readiness")
        return wait_green(client, "S2-TurboQuant-4bit")
    client.update_collection(
        collection_name=COLLECTION, quantization_config=models.Disabled.DISABLED
    )
    time.sleep(5)
    patch_collection({"quantization_config": config})
    return wait_green(client, "S2-TurboQuant-4bit")


def patch_collection(payload):
    req = urllib.request.Request(
        QDRANT_URL.rstrip("/") + f"/collections/{COLLECTION}",
        data=json.dumps(payload).encode("utf-8"),
        headers={"Content-Type": "application/json", "api-key": QDRANT_API_KEY or ""},
        method="PATCH",
    )
    with urllib.request.urlopen(req, timeout=600) as resp:
        print(f"[patch] {resp.read().decode('utf-8', 'replace')}")


def get_collection_info(client):
    global _REST_FALLBACK_WARNED
    try:
        return client.get_collection(COLLECTION)
    except Exception as e:
        if not _REST_FALLBACK_WARNED:
            print(f"[collection] client parse failed; falling back to REST ({type(e).__name__})")
            _REST_FALLBACK_WARNED = True
        req = urllib.request.Request(
            QDRANT_URL.rstrip("/") + f"/collections/{COLLECTION}",
            headers={"api-key": QDRANT_API_KEY or ""},
        )
        with urllib.request.urlopen(req, timeout=60) as resp:
            body = json.loads(resp.read().decode("utf-8"))
        return body.get("result", body)


def current_quantization_matches(client, desired):
    info = get_collection_info(client)
    if not isinstance(info, dict):
        return False
    return info.get("config", {}).get("quantization_config") == desired


def info_value(info, key, default=None):
    if isinstance(info, dict):
        return info.get(key, default)
    return getattr(info, key, default)


def vector_on_disk(info):
    if isinstance(info, dict):
        vectors = info.get("config", {}).get("params", {}).get("vectors", {})
        return vectors.get("on_disk")
    return getattr(info.config.params.vectors, "on_disk", None)


def distance_value(info):
    if isinstance(info, dict):
        vectors = info.get("config", {}).get("params", {}).get("vectors", {})
        return vectors.get("distance")
    return str(getattr(info.config.params.vectors, "distance", None))


def warmup(client, queries, cfg, n):
    for q in queries[:n]:
        try:
            client.query_points(
                collection_name=COLLECTION,
                query=q["vector"],
                limit=LIMIT,
                with_payload=False,
                with_vectors=False,
            )
        except Exception:
            pass


def provenance(client, info):
    try:
        ver = client.info().version
    except Exception:
        ver = None
    return {
        "collection": COLLECTION,
        "server_version": ver,
        "points_count": info_value(info, "points_count"),
        "indexed_vectors_count": info_value(info, "indexed_vectors_count"),
        "segments_count": info_value(info, "segments_count"),
        "vectors_on_disk": vector_on_disk(info),
        "distance": distance_value(info),
        "dim": DIM,
        "n_points": N_POINTS,
        "run_profile": RUN_PROFILE,
        "n_queries": N_QUERIES,
        "warmup": WARMUP,
        "concurrency": CONCURRENCY,
        "duration_s": DURATION,
        "open_qps_target": OPEN_QPS,
        "hnsw_efs": HNSW_EFS,
        "oversamplings": OVERSAMPLINGS,
        "async_scorer_asserted": ASYNC_SCORER,
        "patterns": ["A", "B"],
        "quantization": "turbo_4bit",
    }


def row_base(pattern, ef, oversampling):
    return {
        "scenario": "S2",
        "quant": "turbo_4bit",
        "pattern": pattern,
        "ef": ef,
        "oversampling": oversampling,
        "async_scorer": ASYNC_SCORER,
    }


def run_point(client, queries, cfg, label, base):
    print(f"  {label}  ef={cfg.ef} ov={cfg.oversampling} ...", end=" ", flush=True)
    t0 = time.time()
    rec = measure_recall(client, COLLECTION, queries, cfg)
    cl = closed_loop(client, COLLECTION, queries, cfg, CONCURRENCY, DURATION)
    ol = open_loop(client, COLLECTION, queries, cfg, OPEN_QPS, DURATION)
    print(
        f"recall={rec}  qps={cl['qps']}  avg={cl['avg_ms']}ms  p99={cl['p99_ms']}ms  "
        f"({time.time()-t0:.0f}s)"
    )
    return {
        **base,
        "recall": rec,
        "closed_qps": cl["qps"],
        "closed_avg_ms": cl["avg_ms"],
        "closed_p95_ms": cl["p95_ms"],
        "closed_p99_ms": cl["p99_ms"],
        "closed_count": cl["count"],
        "open_target_qps": ol["target_qps"],
        "open_achieved_qps": ol["achieved_qps"],
        "open_avg_ms": ol["avg_ms"],
        "open_p95_ms": ol["p95_ms"],
        "open_p99_ms": ol["p99_ms"],
    }


def run_matrix(client, queries, rows):
    set_turboquant(client)
    print(f"[S2] sweeping ef={HNSW_EFS} x ov={OVERSAMPLINGS} x patterns A,B")
    warmup(client, queries, SearchConfig("A", HNSW_EFS[-1], True, OVERSAMPLINGS[-1]), WARMUP)
    for ef in HNSW_EFS:
        for ov in OVERSAMPLINGS:
            for pat in ("A", "B"):
                cfg = SearchConfig(pat, ef, rescore=True, oversampling=ov)
                base = row_base(pat, ef, ov)
                rows.append(run_point(client, queries, cfg, f"pattern-{pat}", base))


def print_summary(rows):
    print("\n" + "=" * 72)
    print("SUMMARY  (nearest to recall@100 ≈ 0.96)")
    print("=" * 72)
    target = 0.96
    for pat in ("A", "B"):
        grp = [r for r in rows if r["pattern"] == pat]
        if not grp:
            continue
        closest = min(grp, key=lambda r: abs(r["recall"] - target))
        print(
            f"  pattern {pat} ef={closest['ef']} recall={closest['recall']:.4f}  "
            f"qps={closest['closed_qps']} avg={closest['closed_avg_ms']}ms  "
            f"p99={closest['closed_p99_ms']}ms"
        )


def main():
    OUT_DIR.mkdir(exist_ok=True)
    client = get_client()
    info0 = get_collection_info(client)
    prov = provenance(client, info0)
    print("[provenance]", json.dumps(prov, indent=2))

    points_count = info_value(info0, "points_count", 0) or 0
    if points_count < N_POINTS:
        print(f"[WARN] only {points_count}/{N_POINTS} points — load incomplete")
    if not prov["vectors_on_disk"]:
        print("[WARN] vectors NOT on_disk — rescore will not hit disk")

    queries = load_queries(QUERIES_PATH, limit=N_QUERIES)
    print(f"[queries] {len(queries)} loaded")

    rows = []
    run_matrix(client, queries, rows)
    _save(rows, prov, OUT_DIR)
    print_summary(rows)
    print(f"\n[done] {len(rows)} rows → {OUT_DIR}/results.json / .csv")


def _save(rows, prov, out):
    data = {"provenance": prov, "results": rows}
    (out / "results.json").write_text(json.dumps(data, indent=2))
    if not rows:
        return
    cols = [
        "scenario",
        "quant",
        "pattern",
        "ef",
        "oversampling",
        "async_scorer",
        "recall",
        "closed_qps",
        "closed_avg_ms",
        "closed_p95_ms",
        "closed_p99_ms",
        "closed_count",
        "open_target_qps",
        "open_achieved_qps",
        "open_avg_ms",
        "open_p95_ms",
        "open_p99_ms",
    ]
    with open(out / "results.csv", "w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=cols, extrasaction="ignore")
        w.writeheader()
        w.writerows(rows)


if __name__ == "__main__":
    main()