engineering/arborist
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

bench: sample-level shuffled scheduling with per-cell locks

Browse source 
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).
This commit is contained in:
russell@unturf.com 2026-05-02 10:16:33 -04:00
parent 3ac13975a4
commit 8d6fd1b836
No known key found for this signature in database
1 changed files with 71 additions and 40 deletions
Download patch file Download diff file Expand all files Collapse all files

111
bench/qa_sweep.py
View file

@ -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