diff --git a/bench/qa_sweep.py b/bench/qa_sweep.py index 1824bc4..ac6ef08 100644 --- a/bench/qa_sweep.py +++ b/bench/qa_sweep.py @@ -27,8 +27,10 @@ import argparse import datetime as _dt import json import os +import random import sys import time +from collections import defaultdict from concurrent.futures import ThreadPoolExecutor, as_completed from pathlib import Path from threading import Lock @@ -474,10 +476,20 @@ def main(argv: list[str] | None = None) -> int: ap.add_argument("--limit", type=int, default=0, help="truncate question list to N; 0 = all") ap.add_argument("--concurrency", type=int, default=1, - help="parallel (question, mode) cells; samples within " - "a cell stay sequential so burn+insert on a single " - "cache_key never races. vLLM handles concurrent " - "requests well; 4-8 is a safe starting point") + help="parallel sample-level tasks. Each (question, " + "mode, sample_idx) is one task; tasks are shuffled " + "via --seed and dispatched concurrently. Per-cell " + "Lock serializes burn+insert on the shared " + "cache_key so two samples of the same cell never " + "race. vLLM handles concurrent requests well; " + "4-8 is a safe starting point.") + ap.add_argument("--seed", type=int, default=0, + help="deterministic shuffle seed for sample-level " + "task ordering. Same seed = same task order = " + "reproducible bench. Random ordering is the " + "design — uncorrelated samples give true i.i.d. " + "variance estimation and feed vLLM's continuous " + "batcher a diverse request stream.") ap.add_argument("--endpoint", default=os.environ.get( "ABORIST_LLM_ENDPOINT", "https://hermes.ai.unturf.com/v1")) ap.add_argument("--model", default=os.environ.get( @@ -513,64 +525,83 @@ def main(argv: list[str] | None = None) -> int: ) print(f"[bench] shards_dir={args.shards_dir} top_k={args.top_k} " f"burn=always (per-sample, forces fresh inference)") - print(f"[bench] concurrency={args.concurrency} (cells in parallel; " - f"samples within a cell stay sequential)") + print(f"[bench] concurrency={args.concurrency} seed={args.seed} " + f"(sample-level tasks; per-cell Lock serializes cache_key writes)") print(f"[bench] writing → {jsonl_path}") print() - cells = [(q, mode) for q in questions for mode in modes] - total_cells = len(cells) + # Build sample-level task list. Each task is one (question, mode, + # sample_idx); shuffle so samples for the same cell get spread + # across time. Two effects: + # - True i.i.d. n-sample variance: consecutive samples of the + # same cell don't share vLLM batch composition or KV-cache + # locality, so what looks like model nondeterminism actually + # IS model nondeterminism. + # - vLLM's continuous batcher gets a diverse request stream; + # under the prior cell-grouped scheduling, batches were + # correlated and the engine couldn't fill efficiently. + tasks = [ + (q, mode, sample_idx) + for q in questions + for mode in modes + for sample_idx in range(args.n_samples) + ] + rng = random.Random(args.seed) + rng.shuffle(tasks) + total_tasks = len(tasks) + + # Per-cell lock — burn+insert against a shared cache_key must + # serialize. Two samples of the SAME (q, mode) ending up in + # adjacent worker slots could race; the lock makes that path + # safe without bottlenecking unrelated cells. Lock contention is + # rare under shuffle (samples spread out), so the throughput cost + # is near-zero while correctness is preserved. + cell_locks: dict[tuple[str, str], Lock] = defaultdict(Lock) + rows: list[dict] = [] write_lock = Lock() print_lock = Lock() done = [0] - def _run_cell(q: str, mode: str) -> list[dict]: - cell_rows: list[dict] = [] - for sample_idx in range(args.n_samples): - row = _run_one( - question=q, - answer_mode=mode, - shards_dir=args.shards_dir, - qa_db=args.qa_db, - top_k=args.top_k, - burn=True, - endpoint=args.endpoint, - model=args.model, - ) - row["sample_idx"] = sample_idx - cell_rows.append(row) - return cell_rows - with jsonl_path.open("w", encoding="utf-8") as f: - def _process(qm: tuple[str, str]) -> None: - q, mode = qm - cell_rows = _run_cell(q, mode) + def _process(task: tuple[str, str, int]) -> None: + q, mode, sample_idx = task + with cell_locks[(q, mode)]: + row = _run_one( + question=q, + answer_mode=mode, + shards_dir=args.shards_dir, + qa_db=args.qa_db, + top_k=args.top_k, + burn=True, + endpoint=args.endpoint, + model=args.model, + ) + row["sample_idx"] = sample_idx with write_lock: - for r in cell_rows: - f.write(json.dumps(r, ensure_ascii=False) + "\n") - rows.extend(cell_rows) + f.write(json.dumps(row, ensure_ascii=False) + "\n") + rows.append(row) f.flush() with print_lock: done[0] += 1 - last = cell_rows[-1] tag = ( - f"err: {last['error']}" if last["error"] - else f"{last['audit_mode']} {last['n_verified']}/{last['n_quotes']} " - f"{last['elapsed_s']:.1f}s" + f"err: {row['error']}" if row["error"] + else f"{row['audit_mode']} {row['n_verified']}/{row['n_quotes']} " + f"{row['elapsed_s']:.1f}s" ) print( - f" [{done[0]:>3}/{total_cells}] {mode:<22} " - f"{q[:50]:<50} → {tag}", + f" [{done[0]:>3}/{total_tasks}] {mode:<22} " + f"#{sample_idx + 1}/{args.n_samples} " + f"{q[:42]:<42} → {tag}", flush=True, ) if args.concurrency <= 1: - for qm in cells: - _process(qm) + for task in tasks: + _process(task) else: with ThreadPoolExecutor(max_workers=args.concurrency) as pool: - futures = [pool.submit(_process, qm) for qm in cells] + futures = [pool.submit(_process, t) for t in tasks] for fut in as_completed(futures): fut.result() # surfaces any exception