compress: thread-local zstd contexts; bench: --concurrency flag

zstandard.ZstdCompressor and ZstdDecompressor instances each carry an
internal libzstd context that is NOT thread-safe. Calling .compress()
or .decompress() on a single shared instance from multiple threads
corrupts the context and raises:

    ZstdError: decompression error: Data corruption detected

Surfaced when bench/qa_sweep.py learned a --concurrency flag and ran
4 (question, mode) cells in parallel. ~19% of retrievals failed on
zstd corruption before the fix. The prior comment claiming the
singletons were 'stateless across calls — safe to share across
threads' was wrong.

Replace the module-level singletons with threading.local() caches.
Each thread reuses its own ZstdCompressor / ZstdDecompressor; no
contention across threads. Init cost is negligible vs decompression.

bench/qa_sweep.py:

  --concurrency N (default 1) parallelizes (question, mode) CELLS
  using a ThreadPoolExecutor. Samples within a cell stay sequential
  so burn-then-insert against a single cache_key never races itself.
  Lock-protected JSONL writes & progress prints. Exception in any
  worker surfaces via fut.result().

vLLM handles concurrent requests well; 4-8 is a reasonable starting
point. With --n 2 --concurrency 4, expect ~10 min wall-clock for the
full 426-run sweep against hermes.ai.unturf.com (vs ~2h sequential).

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(

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)