russell@unturf.com
0870af60af
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).
108 lines
3.7 KiB
Python
108 lines
3.7 KiB
Python
"""Transparent zstd compression for chunk content.
|
|
|
|
Backward-compatible: reads detect the zstd magic byte sequence
|
|
(\\x28\\xb5\\x2f\\xfd) and decompress; plaintext rows from before this
|
|
feature was introduced pass through unchanged.
|
|
|
|
Why zstd over zlib: ~3-5x compression on natural language vs zlib's
|
|
~2-3x, and zstd decompresses ~5-10x faster than zlib. Level 3 is the
|
|
default — higher levels (5-9) gain only ~5-10% extra ratio at
|
|
significant compression-time CPU cost. Decompression speed is
|
|
level-independent (~250 MB/s on modern x86).
|
|
|
|
Storage in SQLite: declared TEXT columns happily hold BLOB cells thanks
|
|
to SQLite's dynamic typing. We don't change the schema; the column type
|
|
discipline is enforced at the application layer through these helpers.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import threading
|
|
|
|
import zstandard
|
|
|
|
# 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"
|
|
|
|
# Below this UTF-8 byte length, compression overhead dominates and the
|
|
# compressed form is larger than the source. Empty / tiny chunks pass
|
|
# through as plaintext.
|
|
_MIN_COMPRESS_BYTES = 64
|
|
|
|
|
|
def is_compressed(value: object) -> bool:
|
|
"""True if value looks like a zstd frame (magic bytes match)."""
|
|
return isinstance(value, (bytes, bytearray, memoryview)) and bytes(value[:4]) == _ZSTD_MAGIC
|
|
|
|
|
|
def pack_chunk(text: str) -> bytes | str:
|
|
"""Compress text for storage in `chunks.content`.
|
|
|
|
Returns:
|
|
bytes — zstd-compressed UTF-8 if compression saves space.
|
|
str — original text, unmodified, when too small to benefit.
|
|
|
|
SQLite stores either as the column's natural cell type (TEXT for str,
|
|
BLOB for bytes); both round-trip correctly through the Python sqlite3
|
|
binding when read back.
|
|
"""
|
|
if not text:
|
|
return text
|
|
raw = text.encode("utf-8")
|
|
if len(raw) < _MIN_COMPRESS_BYTES:
|
|
return text
|
|
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):
|
|
return text
|
|
return compressed
|
|
|
|
|
|
def unpack_chunk(value: object) -> str | None:
|
|
"""Decompress a chunk content value if compressed; otherwise return as-is.
|
|
|
|
Inputs:
|
|
None -> None (cold-tier or null content)
|
|
str -> str (legacy plaintext rows)
|
|
bytes-like -> str (decoded UTF-8; decompressed first if zstd-framed)
|
|
|
|
Raises TypeError on unexpected input shapes so a bad row surfaces loudly
|
|
instead of silently returning garbage.
|
|
"""
|
|
if value is None:
|
|
return None
|
|
if isinstance(value, str):
|
|
return value
|
|
if isinstance(value, (bytes, bytearray, memoryview)):
|
|
b = bytes(value)
|
|
if is_compressed(b):
|
|
return _decompressor().decompress(b).decode("utf-8")
|
|
# Legacy or non-compressed BLOB cell — try UTF-8 decode.
|
|
return b.decode("utf-8")
|
|
raise TypeError(
|
|
f"unexpected chunk content type: {type(value).__name__}"
|
|
)
|