dryrun.py

#!/usr/bin/env python3
"""
Pre-flight check: validate Pattern A / Pattern B query shapes on a small query subset
before running the full ef sweep.

    export QDRANT_URL=... QDRANT_API_KEY=... ASYNC_SCORER=true
    python3 dryrun.py
"""

import os
from qdrant_client import QdrantClient
from queryset import load_queries
from runengine import SearchConfig, measure_recall, closed_loop, build_query_kwargs, LIMIT

COLLECTION = os.getenv("COLLECTION", "wiki_dpr_e5")
QDRANT_URL = os.environ["QDRANT_URL"]
QDRANT_API_KEY = os.getenv("QDRANT_API_KEY")
N = int(os.getenv("DRY_N", "20"))
ASYNC_SCORER = os.getenv("ASYNC_SCORER", "unknown")


def main():
    client = QdrantClient(url=QDRANT_URL, api_key=QDRANT_API_KEY, timeout=120)
    info = client.get_collection(COLLECTION)
    vp = info.config.params.vectors
    print(
        f"[collection] {COLLECTION} status={info.status} points={info.points_count} "
        f"on_disk={getattr(vp, 'on_disk', None)} quant={info.config.quantization_config}"
    )
    print(f"[async_scorer asserted] {ASYNC_SCORER}")

    queries = load_queries(os.getenv("QUERIES", "queries.parquet"), limit=N)
    print(
        f"[queries] {len(queries)} loaded, dim={len(queries[0]['vector'])}, "
        f"recall@{LIMIT} target"
    )

    cfgs = [
        ("single-stage Pattern A  ", SearchConfig("A", 128, rescore=True, oversampling=1.0)),
        ("two-stage Pattern B     ", SearchConfig("B", 128, rescore=True, oversampling=1.0)),
    ]
    for label, cfg in cfgs:
        res = client.query_points(
            **build_query_kwargs(COLLECTION, queries[0]["vector"], cfg)
        ).points
        rec = measure_recall(client, COLLECTION, queries, cfg)
        cl = closed_loop(client, COLLECTION, queries, cfg, concurrency=4, duration=5)
        print(
            f"[{label}] returned={len(res)} recall@{LIMIT}={rec} "
            f"qps(c4)={cl['qps']} p95={cl['p95_ms']}ms"
        )

    print("[dryrun] OK — proceed with scripts/run_profile.sh")


if __name__ == "__main__":
    main()