diff --git a/aborist/compress.py b/aborist/compress.py index 0f337a7..3709a3f 100644 --- a/aborist/compress.py +++ b/aborist/compress.py @@ -17,13 +17,32 @@ discipline is enforced at the application layer through these helpers. from __future__ import annotations +import threading + import zstandard -# Module-level singletons. The ZstdCompressor / ZstdDecompressor objects are -# stateless across calls — safe to share across threads in this codebase -# (we don't run multi-threaded ingests inside one process). -_COMPRESSOR = zstandard.ZstdCompressor(level=3) -_DECOMPRESSOR = zstandard.ZstdDecompressor() +# python-zstandard's ZstdCompressor / ZstdDecompressor instances carry +# internal state (the underlying libzstd context) and are NOT thread-safe; +# calling .compress() / .decompress() on a shared instance from multiple +# threads corrupts that context and raises ZstdError. Thread-local caches +# give us per-thread reuse without contention. +_TLS = threading.local() + + +def _compressor() -> zstandard.ZstdCompressor: + comp = getattr(_TLS, "comp", None) + if comp is None: + comp = zstandard.ZstdCompressor(level=3) + _TLS.comp = comp + return comp + + +def _decompressor() -> zstandard.ZstdDecompressor: + dec = getattr(_TLS, "dec", None) + if dec is None: + dec = zstandard.ZstdDecompressor() + _TLS.dec = dec + return dec # zstd frame magic (4 bytes). RFC 8478 §3.1.1. _ZSTD_MAGIC = b"\x28\xb5\x2f\xfd" @@ -55,7 +74,7 @@ def pack_chunk(text: str) -> bytes | str: raw = text.encode("utf-8") if len(raw) < _MIN_COMPRESS_BYTES: return text - compressed = _COMPRESSOR.compress(raw) + compressed = _compressor().compress(raw) # Defensive: if the entropy is near-incompressible (already-compressed # data, very short repeats), keep the smaller representation. if len(compressed) >= len(raw): @@ -81,7 +100,7 @@ def unpack_chunk(value: object) -> str | None: if isinstance(value, (bytes, bytearray, memoryview)): b = bytes(value) if is_compressed(b): - return _DECOMPRESSOR.decompress(b).decode("utf-8") + return _decompressor().decompress(b).decode("utf-8") # Legacy or non-compressed BLOB cell — try UTF-8 decode. return b.decode("utf-8") raise TypeError( diff --git a/bench/qa_sweep.py b/bench/qa_sweep.py index 0c7cd57..1824bc4 100644 --- a/bench/qa_sweep.py +++ b/bench/qa_sweep.py @@ -29,7 +29,9 @@ import json import os import sys import time +from concurrent.futures import ThreadPoolExecutor, as_completed from pathlib import Path +from threading import Lock # Defer aborist imports until argparse runs so `--help` works without # the package installed. @@ -471,6 +473,11 @@ def main(argv: list[str] | None = None) -> int: "record so Hermes nondeterminism becomes the variance source") 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") 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( @@ -506,35 +513,66 @@ 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] writing → {jsonl_path}") print() + cells = [(q, mode) for q in questions for mode in modes] + total_cells = len(cells) 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: - for q in questions: - for mode in modes: - for sample_idx in range(args.n_samples): - tag = f"{mode:<22} #{sample_idx + 1}/{args.n_samples}" - print(f" [{tag}] {q[:60]}", flush=True) - 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 - rows.append(row) - f.write(json.dumps(row, ensure_ascii=False) + "\n") - f.flush() - tag = ( - f"err: {row['error']}" if row["error"] - else f"{row['audit_mode']} {row['n_verified']}/{row['n_quotes']} {row['elapsed_s']:.1f}s" - ) - print(f" → {tag}", flush=True) + def _process(qm: tuple[str, str]) -> None: + q, mode = qm + cell_rows = _run_cell(q, mode) + with write_lock: + for r in cell_rows: + f.write(json.dumps(r, ensure_ascii=False) + "\n") + rows.extend(cell_rows) + 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" + ) + print( + f" [{done[0]:>3}/{total_cells}] {mode:<22} " + f"{q[:50]:<50} → {tag}", + flush=True, + ) + + if args.concurrency <= 1: + for qm in cells: + _process(qm) + else: + with ThreadPoolExecutor(max_workers=args.concurrency) as pool: + futures = [pool.submit(_process, qm) for qm in cells] + for fut in as_completed(futures): + fut.result() # surfaces any exception summary = _summarize(rows) md = _render_markdown(rows, summary, stamp, modes, questions, args.n_samples)