headroom/tests/parity/record_text_crusher.py
Zhenjia ZHOU 6c68ff4e9f
perf(compression): take large cold-start contexts off the synchronous kompress path (#1171) (#1298)
## Description

On a cold-start large context, kompress (ModernBERT ONNX) runs
**synchronously on the request thread** — ~200–300s for ~1M tokens. It
blows the 30s compression budget, leaks a non-preemptible worker, and
cascades (executor saturation → queue timeouts on healthy requests); on
timeout the request is forwarded **uncompressed** after eating 30s. This
adds four layered, **default-off, fail-open** mitigations so the request
path is never blocked on ML compression.

Closes #1171

## Type of Change

- [ ] Bug fix (non-breaking change that fixes an issue)
- [x] New feature (non-breaking change that adds functionality)
- [ ] Breaking change (fix or feature that would cause existing
functionality to change)
- [ ] Documentation update
- [x] Performance improvement
- [ ] Code refactoring (no functional changes)

## Changes Made

- **Phase 0 — size gate** (`HEADROOM_KOMPRESS_MAX_TOKENS`, default
50000): route oversized text away from ModernBERT (→ LogCompressor /
TextCrusher / passthrough) at the single `_try_ml_compressor` boundary.
- **Phase 1 — cooperative deadline**
(`HEADROOM_COMPRESSION_DEADLINE_MS`, default 20000): any kompress run
self-terminates at the next chunk boundary past the budget, keeping the
unprocessed tail verbatim.
- **Phase 2 — TextCrusher** (`HEADROOM_TEXT_CRUSHER`): a new **native
Rust** extractive prose compressor in
`crates/headroom-core/src/transforms/text_crusher/`, exposed via PyO3 as
`headroom._core.TextCrusher` with a thin Python wrapper. It **reuses the
shared `crate::relevance::BM25Scorer`** rather than reimplementing BM25,
and ships record/replay parity fixtures (mirroring the SmartCrusher
Rust-core + Python-shim pattern).
- **Phase 3 — off-path compression**
(`HEADROOM_BACKGROUND_COMPRESSION`): forward uncompressed immediately
and compress in a per-process background drain; a byte-identical cache
hit on a later turn means the request never blocks on ML.
- Benchmark (`benchmarks/text_crusher_quality_eval.py`), CHANGELOG
entry, and docstrings documenting the fail-open limits.

## Testing

- [x] Unit tests pass (`pytest`)
- [x] Linting passes (`ruff check`)
- [x] Type checking passes (`mypy`, new modules)
- [x] New tests added for new functionality
- [ ] Manual testing performed (Phase 0/1 gate-fire + deadline observed
on real traffic in earlier iterations; Phase 3 off-path is unit- +
byte-identity-tested, not yet live-validated)

### Test Output

```text
$ pytest tests/test_transforms/ tests/test_cache/ \
    tests/test_proxy/test_background_compression.py tests/test_proxy/test_phase3_byte_identity.py -q
501 passed, 37 skipped in 40.33s

$ cargo test -p headroom-core --lib text_crusher
test result: ok. 3 passed; 0 failed; 0 ignored; 0 measured; 834 filtered out

$ ruff check <changed files>
All checks passed!

$ mypy headroom/proxy/background_compression.py headroom/transforms/text_crusher.py
Success: no issues found in 2 source files
```

New coverage: size-gate incl. the strategy-dispatch funnel (KOMPRESS +
TEXT); partial-run deadline (chunk-0 compressed + chunk-1 verbatim
tail); BackgroundCompressor (dedup / queue-full / fail-open); Phase 3
byte-identity round-trip; TextCrusher unit + parity.

## Real Behavior Proof

- Environment: macOS, local dev — `uv` venv, Rust `_core` built via `uv
pip install -e .`.
- Exact command / steps: the `pytest` / `cargo test` / `ruff` / `mypy`
commands shown under Test Output; quality eval `python
benchmarks/text_crusher_quality_eval.py /tmp/squad_dev.json`.
- Observed result: 501 Python + 3 Rust tests pass; ruff + mypy clean on
changed/new modules. Quality eval: TextCrusher keeps ~94% of buried
SQuAD answers at 30% size vs ~36% truncate/random; self-contained speed
run ~333k words in ~76ms (one O(n) pass) — sub-second where ModernBERT
takes minutes (fast-vs-slow contrast, not a same-input run).
- Not tested: Phase 3 off-path on live traffic; multi-worker
(per-process by design — see Additional Notes).

## Review Readiness

- [x] I have performed a self-review
- [x] This PR is ready for human review

## Checklist

- [x] My code follows the project's style guidelines
- [x] I have performed a self-review of my code
- [x] I have commented my code, particularly in hard-to-understand areas
- [x] I have made corresponding changes to the documentation
- [x] My changes generate no new warnings
- [x] I have added tests that prove my fix is effective or that my
feature works
- [x] New and existing unit tests pass locally with my changes
- [x] I have updated the CHANGELOG.md if applicable

## Additional Notes

- **All four features are off by default and fail-open** — with the env
flags unset the paths are no-ops for realistic inputs; on any error the
request is forwarded (compressed if possible, else verbatim), never
dropped. A full background queue / duplicate key surfaces as
`deferred:dropped`.
- **Known limits (documented in `background_compression.py`):** Phase 3
is per-process, in-memory, and token-mode-only — these are
**lost-savings, never lost-correctness**, and consistent with the
project's existing per-process compression cache + sticky-session
multi-worker model. The startup multi-worker warning now names off-path
background compression.
- Phase 2 reuses the existing BM25 scorer; reuse did not improve
answer-retention over a Python prototype (query-awareness dominates) —
its value is the Rust speed + repo-conventional Rust-core/Python-shim
shape.

---------

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-06-23 10:48:06 -05:00

94 lines
3.5 KiB
Python

"""Record TextCrusher parity fixtures (Phase 2, #1171).
Locks the Rust core's ``compress`` output for a fixed set of deterministic
scenarios so a future change to the Rust algorithm is caught as a regression.
The Python wrapper delegates to ``headroom._core.TextCrusher``, so these
fixtures are recorded from (and verified against) the native implementation.
Re-record after an intentional algorithm change:
python tests/parity/record_text_crusher.py
"""
from __future__ import annotations
import hashlib
import json
import os
from headroom.transforms.text_crusher import TextCrusher
FIXTURE_DIR = os.path.join(os.path.dirname(__file__), "fixtures", "text_crusher")
def _prose(n: int) -> str:
return " ".join(
f"Sentence number {i} explains how distributed systems reconcile state across topic {i}."
for i in range(n)
)
def _redundant() -> str:
dup = "The quick brown fox jumps over the very lazy dog every single morning."
uniques = [f"A distinct fact about subsystem {i} is recorded plainly here." for i in range(8)]
return "\n".join([dup] * 10 + uniques)
def _salient() -> str:
return "\n".join(
[
"ERROR connection refused at host 10.0.0.42 after 3 retries.",
"The authentication module validated tokens against auth.registry before forwarding.",
"A traceback was logged with code 500 and request_id req-9182.",
"Just some generic filler text without any specific identifiers here.",
"More plain filler describing the overall behavior in vague terms.",
"Warning: cache hit ratio dropped to 71 percent during the spike.",
"Another unremarkable sentence with no salient tokens at all today.",
"The pipeline.apply call returned 42 kept rows out of 1000 total.",
]
)
# (label, content, context, target_ratio)
SCENARIOS: list[tuple[str, str, str, float | None]] = [
("plain_prose", _prose(30), "how do distributed systems reconcile state", 0.3),
("plain_prose_no_query", _prose(30), "", 0.5),
("redundant", _redundant(), "", 0.9),
("salient_heavy", _salient(), "authentication tokens errors", 0.4),
("short_passthrough", "one thing. two thing. three thing.", "", None),
(
"unicode",
" ".join(f"句子 {i} 描述了系统在主题 {i} 上的行为细节。" for i in range(12)),
"系统",
0.4,
),
]
def record() -> None:
os.makedirs(FIXTURE_DIR, exist_ok=True)
tc = TextCrusher()
for label, content, context, ratio in SCENARIOS:
r = tc.compress(content, context, ratio)
digest = hashlib.sha256(content.encode("utf-8")).hexdigest()
fixture = {
"transform": "text_crusher",
"label": label,
"input": {"content": content, "context": context, "target_ratio": ratio},
"output": {
"compressed": r.compressed,
"original_tokens": r.original_tokens,
"compressed_tokens": r.compressed_tokens,
"compression_ratio": r.compression_ratio,
"kept_segments": r.kept_segments,
"total_segments": r.total_segments,
},
"input_sha256": digest,
}
path = os.path.join(FIXTURE_DIR, f"{label}_{digest[:12]}.json")
with open(path, "w", encoding="utf-8") as fh:
json.dump(fixture, fh, indent=2, ensure_ascii=False)
print(f"wrote {path}")
if __name__ == "__main__":
record()