mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
## 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>
109 lines
3.5 KiB
Python
109 lines
3.5 KiB
Python
"""Phase 3 (#1171): off-path BackgroundCompressor.
|
|
|
|
A single per-process drain compresses enqueued work with NO request-coupled
|
|
deadline and stores the result -- so the request path never blocks on ML.
|
|
Tests use asyncio.run so they do not depend on a pytest-asyncio config.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
|
|
from headroom.proxy.background_compression import BackgroundCompressor
|
|
|
|
|
|
async def _passthrough_executor(fn):
|
|
return fn()
|
|
|
|
|
|
def test_enqueue_compresses_and_stores():
|
|
async def main():
|
|
stored: list[str] = []
|
|
bc = BackgroundCompressor(_passthrough_executor)
|
|
await bc.start()
|
|
ok = bc.enqueue("k1", lambda: "COMPRESSED", lambda r: stored.append(r))
|
|
await asyncio.wait_for(bc._queue.join(), timeout=2)
|
|
await bc.stop()
|
|
return ok, stored, bc.stats()
|
|
|
|
ok, stored, stats = asyncio.run(main())
|
|
assert ok is True
|
|
assert stored == ["COMPRESSED"]
|
|
assert stats["processed"] == 1
|
|
assert stats["pending"] == 0
|
|
|
|
|
|
def test_dedup_skips_inflight_key():
|
|
async def main():
|
|
calls: list[int] = []
|
|
|
|
async def slow_executor(fn):
|
|
await asyncio.sleep(0.05)
|
|
return fn()
|
|
|
|
bc = BackgroundCompressor(slow_executor)
|
|
await bc.start()
|
|
first = bc.enqueue("same", lambda: calls.append(1), lambda r: None)
|
|
second = bc.enqueue("same", lambda: calls.append(1), lambda r: None)
|
|
await asyncio.wait_for(bc._queue.join(), timeout=2)
|
|
await bc.stop()
|
|
return first, second, len(calls)
|
|
|
|
first, second, n = asyncio.run(main())
|
|
assert first is True
|
|
assert second is False # duplicate key skipped
|
|
assert n == 1
|
|
|
|
|
|
def test_fail_open_on_compress_error():
|
|
async def main():
|
|
stored: list[str] = []
|
|
bc = BackgroundCompressor(_passthrough_executor)
|
|
await bc.start()
|
|
|
|
def boom():
|
|
raise RuntimeError("kompress exploded")
|
|
|
|
bc.enqueue("bad", boom, lambda r: stored.append(r))
|
|
bc.enqueue("good", lambda: "ok", lambda r: stored.append(r))
|
|
await asyncio.wait_for(bc._queue.join(), timeout=2)
|
|
await bc.stop()
|
|
return stored, bc.stats()
|
|
|
|
stored, stats = asyncio.run(main())
|
|
assert stored == ["ok"] # failing job did not store; later job still processed
|
|
assert stats["errors"] == 1
|
|
assert stats["processed"] == 1
|
|
assert stats["pending"] == 0 # key released even on error
|
|
|
|
|
|
def test_queue_full_drops_without_raising():
|
|
async def main():
|
|
# Never start the drain, so the queue fills and stays full.
|
|
bc = BackgroundCompressor(_passthrough_executor, max_queue=2)
|
|
results = [bc.enqueue(f"k{i}", lambda: None, lambda r: None) for i in range(5)]
|
|
return results, bc.stats()
|
|
|
|
results, stats = asyncio.run(main())
|
|
assert results[:2] == [True, True]
|
|
assert results[2:] == [False, False, False] # overflow dropped, no exception
|
|
assert stats["dropped"] == 3
|
|
|
|
|
|
def test_stop_drains_remaining_jobs():
|
|
async def main():
|
|
stored: list[str] = []
|
|
|
|
async def slow_executor(fn):
|
|
await asyncio.sleep(0.02)
|
|
return fn()
|
|
|
|
bc = BackgroundCompressor(slow_executor)
|
|
await bc.start()
|
|
for i in range(4):
|
|
bc.enqueue(f"k{i}", (lambda i=i: f"c{i}"), lambda r: stored.append(r))
|
|
await bc.stop() # drain=True by default -> waits for queue.join()
|
|
return sorted(stored)
|
|
|
|
stored = asyncio.run(main())
|
|
assert stored == ["c0", "c1", "c2", "c3"]
|