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).
This commit is contained in:
parent
9810fd6218
commit
0870af60af
2 changed files with 88 additions and 31 deletions
|
|
@ -17,13 +17,32 @@ discipline is enforced at the application layer through these helpers.
|
||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import threading
|
||||||
|
|
||||||
import zstandard
|
import zstandard
|
||||||
|
|
||||||
# Module-level singletons. The ZstdCompressor / ZstdDecompressor objects are
|
# python-zstandard's ZstdCompressor / ZstdDecompressor instances carry
|
||||||
# stateless across calls — safe to share across threads in this codebase
|
# internal state (the underlying libzstd context) and are NOT thread-safe;
|
||||||
# (we don't run multi-threaded ingests inside one process).
|
# calling .compress() / .decompress() on a shared instance from multiple
|
||||||
_COMPRESSOR = zstandard.ZstdCompressor(level=3)
|
# threads corrupts that context and raises ZstdError. Thread-local caches
|
||||||
_DECOMPRESSOR = zstandard.ZstdDecompressor()
|
# 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 frame magic (4 bytes). RFC 8478 §3.1.1.
|
||||||
_ZSTD_MAGIC = b"\x28\xb5\x2f\xfd"
|
_ZSTD_MAGIC = b"\x28\xb5\x2f\xfd"
|
||||||
|
|
@ -55,7 +74,7 @@ def pack_chunk(text: str) -> bytes | str:
|
||||||
raw = text.encode("utf-8")
|
raw = text.encode("utf-8")
|
||||||
if len(raw) < _MIN_COMPRESS_BYTES:
|
if len(raw) < _MIN_COMPRESS_BYTES:
|
||||||
return text
|
return text
|
||||||
compressed = _COMPRESSOR.compress(raw)
|
compressed = _compressor().compress(raw)
|
||||||
# Defensive: if the entropy is near-incompressible (already-compressed
|
# Defensive: if the entropy is near-incompressible (already-compressed
|
||||||
# data, very short repeats), keep the smaller representation.
|
# data, very short repeats), keep the smaller representation.
|
||||||
if len(compressed) >= len(raw):
|
if len(compressed) >= len(raw):
|
||||||
|
|
@ -81,7 +100,7 @@ def unpack_chunk(value: object) -> str | None:
|
||||||
if isinstance(value, (bytes, bytearray, memoryview)):
|
if isinstance(value, (bytes, bytearray, memoryview)):
|
||||||
b = bytes(value)
|
b = bytes(value)
|
||||||
if is_compressed(b):
|
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.
|
# Legacy or non-compressed BLOB cell — try UTF-8 decode.
|
||||||
return b.decode("utf-8")
|
return b.decode("utf-8")
|
||||||
raise TypeError(
|
raise TypeError(
|
||||||
|
|
|
||||||
|
|
@ -29,7 +29,9 @@ import json
|
||||||
import os
|
import os
|
||||||
import sys
|
import sys
|
||||||
import time
|
import time
|
||||||
|
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
from threading import Lock
|
||||||
|
|
||||||
# Defer aborist imports until argparse runs so `--help` works without
|
# Defer aborist imports until argparse runs so `--help` works without
|
||||||
# the package installed.
|
# the package installed.
|
||||||
|
|
@ -471,6 +473,11 @@ def main(argv: list[str] | None = None) -> int:
|
||||||
"record so Hermes nondeterminism becomes the variance source")
|
"record so Hermes nondeterminism becomes the variance source")
|
||||||
ap.add_argument("--limit", type=int, default=0,
|
ap.add_argument("--limit", type=int, default=0,
|
||||||
help="truncate question list to N; 0 = all")
|
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(
|
ap.add_argument("--endpoint", default=os.environ.get(
|
||||||
"ABORIST_LLM_ENDPOINT", "https://hermes.ai.unturf.com/v1"))
|
"ABORIST_LLM_ENDPOINT", "https://hermes.ai.unturf.com/v1"))
|
||||||
ap.add_argument("--model", default=os.environ.get(
|
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} "
|
print(f"[bench] shards_dir={args.shards_dir} top_k={args.top_k} "
|
||||||
f"burn=always (per-sample, forces fresh inference)")
|
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(f"[bench] writing → {jsonl_path}")
|
||||||
print()
|
print()
|
||||||
|
|
||||||
|
cells = [(q, mode) for q in questions for mode in modes]
|
||||||
|
total_cells = len(cells)
|
||||||
rows: list[dict] = []
|
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:
|
with jsonl_path.open("w", encoding="utf-8") as f:
|
||||||
for q in questions:
|
def _process(qm: tuple[str, str]) -> None:
|
||||||
for mode in modes:
|
q, mode = qm
|
||||||
for sample_idx in range(args.n_samples):
|
cell_rows = _run_cell(q, mode)
|
||||||
tag = f"{mode:<22} #{sample_idx + 1}/{args.n_samples}"
|
with write_lock:
|
||||||
print(f" [{tag}] {q[:60]}", flush=True)
|
for r in cell_rows:
|
||||||
row = _run_one(
|
f.write(json.dumps(r, ensure_ascii=False) + "\n")
|
||||||
question=q,
|
rows.extend(cell_rows)
|
||||||
answer_mode=mode,
|
f.flush()
|
||||||
shards_dir=args.shards_dir,
|
with print_lock:
|
||||||
qa_db=args.qa_db,
|
done[0] += 1
|
||||||
top_k=args.top_k,
|
last = cell_rows[-1]
|
||||||
burn=True,
|
tag = (
|
||||||
endpoint=args.endpoint,
|
f"err: {last['error']}" if last["error"]
|
||||||
model=args.model,
|
else f"{last['audit_mode']} {last['n_verified']}/{last['n_quotes']} "
|
||||||
)
|
f"{last['elapsed_s']:.1f}s"
|
||||||
row["sample_idx"] = sample_idx
|
)
|
||||||
rows.append(row)
|
print(
|
||||||
f.write(json.dumps(row, ensure_ascii=False) + "\n")
|
f" [{done[0]:>3}/{total_cells}] {mode:<22} "
|
||||||
f.flush()
|
f"{q[:50]:<50} → {tag}",
|
||||||
tag = (
|
flush=True,
|
||||||
f"err: {row['error']}" if row["error"]
|
)
|
||||||
else f"{row['audit_mode']} {row['n_verified']}/{row['n_quotes']} {row['elapsed_s']:.1f}s"
|
|
||||||
)
|
if args.concurrency <= 1:
|
||||||
print(f" → {tag}", flush=True)
|
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)
|
summary = _summarize(rows)
|
||||||
md = _render_markdown(rows, summary, stamp, modes, questions, args.n_samples)
|
md = _render_markdown(rows, summary, stamp, modes, questions, args.n_samples)
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue