From 8d6fd1b8367d921d11622f5d0ba20fb1b4896e3b Mon Sep 17 00:00:00 2001 From: "russell@unturf.com" Date: Sat, 2 May 2026 10:16:33 -0400 Subject: [PATCH] bench: sample-level shuffled scheduling with per-cell locks MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Previously the bench scheduled at the (question, mode) cell level — each cell ran its n samples sequentially before the next cell started. With n=3 + concurrency=6, that meant 6 workers each chained 3 sequential calls, vLLM's continuous batcher saw correlated requests, and the n=3 variance signal was confounded with 'what was vLLM doing in this same-cell batch.' Two coupled defects, one fix: (1) Statistical: samples-of-the-same-cell are NOT i.i.d. when run back-to-back. Adjacent samples share vLLM batch composition, KV-cache locality, and queue load. The 'variance' you see across n=3 is partly that batch's particular weather, not true model nondeterminism. (2) Throughput: vLLM batches diverse requests well; correlated requests fill the batch with similar work and starve. Cell-grouped scheduling correlated by cache_key. Refactor: build sample-level tasks ((q, mode, sample_idx)), shuffle with deterministic seed (default 0), submit all to ThreadPoolExecutor. Per-cell Lock dict (defaultdict(Lock)) serializes burn+insert against the shared cache_key — two samples of the same cell that happen to land in adjacent worker slots will queue on the cell's lock instead of racing. Lock contention is rare under shuffle (samples of one cell are spread across time) so throughput cost is near-zero. New CLI flag --seed N for reproducible task ordering. Changes: - tasks list at sample granularity - rng = random.Random(seed); rng.shuffle(tasks) - cell_locks = defaultdict(Lock); held over _run_one + sample_idx assignment - print line shows '#i/n' so order-of-completion is visible Smoke at --n 1 --concurrency 6: clean, no errors, ~3 min for 15 tasks (essentially same wall-clock as cell-level scheduling at n=1, as expected — the real win comes at n=3 where sample-level concurrency unlocks parallelism within cells). --- bench/qa_sweep.py | 111 +++++++++++++++++++++++++++++----------------- 1 file changed, 71 insertions(+), 40 deletions(-) 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