fix(kompress): raise the default execution-slot wait (#2456)

## Description

Concurrent Kompress requests currently fail open after a 25 ms
execution-slot wait even though ordinary ONNX inference can hold the
single slot for hundreds of milliseconds. This raises the existing
default wait to 3000 ms while retaining concurrency one, the
`HEADROOM_KOMPRESS_EXECUTION_TIMEOUT_MS` override, the tighter acquire
and request budgets, and passthrough after a genuine timeout.

The reproduction and validated 3000 ms setting come from
https://github.com/headroomlabs-ai/headroom/issues/2451

Closes #2451

## Type of Change

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

## Changes Made

- Raise the default Kompress execution-slot wait from 25 ms to 3000 ms.
- Start the Kompress request deadline at call entry and carry it through
single-item acquire, single-to-batch delegation, and sequential-fallback
lineage.
- Cap the raised execution-slot wait by that live request deadline on
both single-item and batch acquire paths.
- Keep the per-backend default concurrency at one and preserve all
tighter time budgets.
- Add queued single-item, batch, request-deadline, carried-deadline
lineage, and router-watchdog lifecycle regressions at the same owner
layer that currently fails.
- Preserve the explicit short-timeout fail-open path.

## Testing

- [x] Unit tests pass (`uv run pytest tests/test_kompress_failsafe.py
tests/test_kompress_request_nonblocking.py
tests/test_content_router_single_item_deadline.py
tests/test_transforms/test_kompress_deadline.py -v`)
- [x] Linting passes (`uv run ruff check
headroom/transforms/kompress_compressor.py
tests/test_kompress_failsafe.py
tests/test_content_router_single_item_deadline.py
tests/test_transforms/test_kompress_deadline.py`)
- [x] Formatting passes (`uv run ruff format
headroom/transforms/kompress_compressor.py
tests/test_kompress_failsafe.py
tests/test_content_router_single_item_deadline.py
tests/test_transforms/test_kompress_deadline.py --check`)
- [x] New regression tests prove the saturation fix
- [ ] Manual testing performed

### Test Output

```text
uv run pytest tests/test_kompress_failsafe.py tests/test_kompress_request_nonblocking.py tests/test_content_router_single_item_deadline.py tests/test_transforms/test_kompress_deadline.py -v
37 passed in 4.22s

uv run ruff check headroom/transforms/kompress_compressor.py tests/test_kompress_failsafe.py tests/test_content_router_single_item_deadline.py tests/test_transforms/test_kompress_deadline.py
All checks passed!

uv run ruff format headroom/transforms/kompress_compressor.py tests/test_kompress_failsafe.py tests/test_content_router_single_item_deadline.py tests/test_transforms/test_kompress_deadline.py --check
4 files already formatted
```

## Real Behavior Proof

- Environment: worktree Python environment from `uv sync --extra dev`,
focused pytest with real Python threads and `threading.BoundedSemaphore`
- Exact command / steps: hold the sole execution slot with the
environment override unset, start queued single-item and batch
compression workers, wait until each worker proves it reached a blocked
acquire on the shared execution semaphore, release the slot, rerun the
explicit 1 ms timeout preservation case, then set
`HEADROOM_COMPRESSION_DEADLINE_MS=10` and repeat the held-slot
single-item and batch acquires plus a router single-cache-miss run whose
Kompress load sleeps past the request deadline.
- Observed result: The queued single-item and batch workers each proved
a real blocked acquire before release, then acquired after release and
compressed, while `HEADROOM_KOMPRESS_EXECUTION_TIMEOUT_MS=1` still
passed through promptly, the 10 ms request deadline capped the raised
default wait so both held-slot paths failed open before 200 ms without
reaching model inference, the single-to-batch and sequential-fallback
lineage regressions proved later branches inherit the original request
start instead of resetting it, and the router lifecycle proof showed the
carried deadline now allows slow Kompress load to start but still
expires before model inference after the outer request has already
failed open.
- Not tested: live ONNX proxy savings under sustained concurrent load

## 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] My changes generate no new warnings
- [x] I have added tests that prove my fix is effective
- [x] New and existing unit tests pass locally with my changes
- [x] I have updated the CHANGELOG.md if applicable

## Additional Notes

`CHANGELOG.md` stays unchanged because the release pipeline generates
changelog entries from conventional commits. The fail-open path from
#1430 stays intact; this change stops it from firing spuriously under
ordinary queueing.
This commit is contained in:
Rod Boev 2026-07-22 09:17:33 -04:00 committed by GitHub
parent a09ba6c087
commit 5bd2266f16
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 306 additions and 10 deletions

View file

@ -96,11 +96,12 @@ KOMPRESS_ONNX_INTER_THREADS_ENV = "HEADROOM_KOMPRESS_ONNX_INTER_THREADS"
KOMPRESS_COREML_CACHE_DIR_ENV = "HEADROOM_KOMPRESS_COREML_CACHE_DIR"
KOMPRESS_MAX_CONCURRENT_ENV = "HEADROOM_KOMPRESS_MAX_CONCURRENT"
KOMPRESS_EXECUTION_SEMAPHORE_WAIT_MS_ENV = "HEADROOM_KOMPRESS_EXECUTION_TIMEOUT_MS"
KOMPRESS_EXECUTION_SEMAPHORE_WAIT_MS_DEFAULT = 25
KOMPRESS_EXECUTION_SEMAPHORE_WAIT_MS_DEFAULT = 3000
KOMPRESS_BATCH_SIZE_ENV = "HEADROOM_KOMPRESS_BATCH_SIZE"
KOMPRESS_ACQUIRE_TIMEOUT_ENV = "HEADROOM_KOMPRESS_ACQUIRE_TIMEOUT_SECONDS"
KOMPRESS_TIME_BUDGET_ENV = "HEADROOM_KOMPRESS_TIME_BUDGET_SECONDS"
KOMPRESS_CANARY_THRESHOLD_ENV = "HEADROOM_KOMPRESS_CANARY_SECONDS"
KOMPRESS_REQUEST_DEADLINE_ENV = "HEADROOM_COMPRESSION_DEADLINE_MS"
# Both defaults sit well under the proxy's 30s compression-stage timeout so a
# slow model gives up (passthrough) before the request is abandoned. A thread
@ -175,6 +176,13 @@ def _execution_wait_budget_seconds() -> float:
return parsed / 1000.0
def _request_deadline_seconds() -> float:
try:
return max(0.0, float(os.environ.get(KOMPRESS_REQUEST_DEADLINE_ENV, "20000")) / 1000.0)
except ValueError:
return 20.0
def _acquire_execution_slot(
backend: str,
device_type: str,
@ -1199,6 +1207,7 @@ class KompressCompressor(Transform):
*,
allow_download: bool = True,
ccr_original: str | None = None,
_deadline_started_at: float | None = None,
) -> KompressResult:
"""Compress content using Kompress model.
@ -1223,6 +1232,7 @@ class KompressCompressor(Transform):
Returns:
KompressResult with compressed text.
"""
t_deadline = time.perf_counter() if _deadline_started_at is None else _deadline_started_at
words = content.split()
n_words = len(words)
@ -1238,13 +1248,7 @@ class KompressCompressor(Transform):
# Cached per instance: operator config, read once -- not per compress() call.
deadline_s = getattr(self, "_deadline_s", None)
if deadline_s is None:
try:
deadline_s = max(
0.0,
float(os.environ.get("HEADROOM_COMPRESSION_DEADLINE_MS", "20000")) / 1000.0,
)
except ValueError:
deadline_s = 20.0
deadline_s = _request_deadline_seconds()
self._deadline_s = deadline_s
try:
@ -1263,6 +1267,7 @@ class KompressCompressor(Transform):
target_ratio=[target_ratio],
batch_size=_batch_size(),
ccr_originals=[ccr_original],
_deadline_started_at=t_deadline,
)
if batch_result:
return batch_result[0]
@ -1271,7 +1276,6 @@ class KompressCompressor(Transform):
kept_ids: set[int] = set()
inference_ms = 0.0
chunk_count = 0
t_deadline = time.perf_counter()
acquire_timeout = _acquire_timeout_seconds()
budget = _time_budget_seconds()
@ -1327,12 +1331,29 @@ class KompressCompressor(Transform):
input_ids = input_ids.to(device)
attention_mask = attention_mask.to(device)
request_remaining: float | None = None
if deadline_s:
request_remaining = deadline_s - (time.perf_counter() - t_deadline)
if request_remaining <= 0:
kept_ids.update(range(chunk_start, n_words))
logger.warning(
"Kompress hit %.1fs deadline before acquire after %d/%d words "
"(%d chunks done); kept remainder verbatim to free the request "
"thread (#1171)",
deadline_s,
chunk_start,
n_words,
chunk_count,
)
break
acquire_bounds = [
bound
for bound in (
_execution_wait_budget_seconds(),
acquire_timeout,
remaining,
request_remaining,
)
if bound is not None
]
@ -1471,6 +1492,7 @@ class KompressCompressor(Transform):
batch_size: int = 32,
*,
ccr_originals: list[str | None] | None = None,
_deadline_started_at: float | None = None,
) -> list[KompressResult]:
"""Compress multiple texts. Uses batched inference on GPU, sequential on CPU.
@ -1531,6 +1553,7 @@ class KompressCompressor(Transform):
n = len(contents)
if n == 0:
return []
t_deadline = time.perf_counter() if _deadline_started_at is None else _deadline_started_at
# Normalize target_ratio to a per-text list
if isinstance(target_ratio, list):
@ -1572,6 +1595,7 @@ class KompressCompressor(Transform):
question=question,
target_ratio=r,
ccr_original=ccr_source,
_deadline_started_at=t_deadline,
)
for content, r, ccr_source in zip(contents, ratios, ccr_sources, strict=True)
]
@ -1608,6 +1632,10 @@ class KompressCompressor(Transform):
device_type = _model_device_type(model, backend)
kept_ids_per_text: dict[int, set[int]] = {i: set() for i in range(n) if results[i] is None}
inference_ms = 0.0
deadline_s = getattr(self, "_deadline_s", None)
if deadline_s is None:
deadline_s = _request_deadline_seconds()
self._deadline_s = deadline_s
acquire_timeout = _acquire_timeout_seconds()
budget = _time_budget_seconds()
@ -1637,6 +1665,9 @@ class KompressCompressor(Transform):
if remaining <= 0:
_bail_remaining("time budget exhausted", batch_start)
break
if deadline_s and (deadline_s - (time.perf_counter() - t_deadline)) <= 0:
_bail_remaining("request deadline exhausted", batch_start)
break
batch = chunk_queue[batch_start : batch_start + batch_size]
batch_word_lists = [c[2] for c in batch]
@ -1660,6 +1691,13 @@ class KompressCompressor(Transform):
input_ids = input_ids.to(device)
attention_mask = attention_mask.to(device)
request_remaining: float | None = None
if deadline_s:
request_remaining = deadline_s - (time.perf_counter() - t_deadline)
if request_remaining <= 0:
_bail_remaining("request deadline exhausted", batch_start)
break
# Single forward pass for all chunks in this batch.
acquire_bounds = [
bound
@ -1667,6 +1705,7 @@ class KompressCompressor(Transform):
_execution_wait_budget_seconds(),
acquire_timeout,
remaining,
request_remaining,
)
if bound is not None
]

View file

@ -2,6 +2,7 @@ from __future__ import annotations
import time
import headroom.transforms.kompress_compressor as kc
from headroom.transforms.content_detector import ContentType
from headroom.transforms.content_router import (
CompressionStrategy,
@ -10,6 +11,7 @@ from headroom.transforms.content_router import (
RouterCompressionResult,
RoutingDecision,
)
from headroom.transforms.kompress_compressor import KompressCompressor, KompressConfig
class _Tokenizer:
@ -48,7 +50,7 @@ def _messages() -> list[dict[str, str]]:
{"role": "assistant", "content": "frozen prefix content remains unchanged"},
{
"role": "assistant",
"content": "pending cache miss content takes the inline compression branch",
"content": "pending cache miss content takes the inline compression branch today",
},
]
@ -112,3 +114,71 @@ def test_single_cache_miss_preserves_disabled_deadline(monkeypatch):
)
assert result.messages[1]["content"] == "compressed output"
def test_single_cache_miss_deadline_starts_before_kompress_load(monkeypatch, caplog):
router = _router()
class _Encoding(dict):
def __init__(self, rows: list[list[str]]):
super().__init__(
input_ids=[[0] * len(row) for row in rows],
attention_mask=[[1] * len(row) for row in rows],
)
self._rows = rows
def word_ids(self, batch_index: int = 0):
return list(range(len(self._rows[batch_index])))
class _Tokenizer:
def count_text(self, content: str) -> int:
return len(content.split())
def __call__(self, words, **_kwargs):
rows = words if words and isinstance(words[0], list) else [words]
return _Encoding(rows)
class _Model:
def __init__(self):
self.calls = 0
def get_keep_mask(self, input_ids, attention_mask):
self.calls += 1
return [[i % 2 == 0 for i in range(len(row))] for row in input_ids]
model = _Model()
compressor = KompressCompressor(config=KompressConfig(enable_ccr=False))
monkeypatch.setattr(compressor, "_should_batch_single_content", lambda *a, **k: False)
load_state = {"calls": 0}
def _slow_load(*_args, **_kwargs):
load_state["calls"] += 1
time.sleep(0.05)
return model, _Tokenizer(), "onnx"
monkeypatch.setattr(kc, "_load_kompress", _slow_load)
monkeypatch.setattr(
router,
"compress",
lambda content, *, context="", bias=1.0: _compression_result(
content,
compressor.compress(content).compressed,
),
)
monkeypatch.setenv("HEADROOM_COMPRESSION_DEADLINE_MS", "10")
started = time.perf_counter()
result = router.apply(
_messages(),
_Tokenizer(),
frozen_message_count=1,
min_tokens_to_compress=1,
)
elapsed = time.perf_counter() - started
time.sleep(0.1)
assert elapsed < 0.12
assert result.messages[1]["content"] == _messages()[1]["content"]
assert "failing open via PASSTHROUGH" in caplog.text
assert load_state["calls"] == 1
assert model.calls == 0

View file

@ -11,6 +11,7 @@ No ML dependencies — the model/tokenizer are fakes injected via
``_load_kompress``.
"""
import threading
import time
import pytest
@ -19,6 +20,8 @@ import headroom.transforms.kompress_compressor as kc
from headroom.transforms.kompress_compressor import (
KOMPRESS_ACQUIRE_TIMEOUT_ENV,
KOMPRESS_CANARY_THRESHOLD_ENV,
KOMPRESS_EXECUTION_SEMAPHORE_WAIT_MS_ENV,
KOMPRESS_REQUEST_DEADLINE_ENV,
KOMPRESS_TIME_BUDGET_ENV,
KompressCompressor,
KompressConfig,
@ -81,6 +84,8 @@ def _reset_module_state(monkeypatch):
KOMPRESS_ACQUIRE_TIMEOUT_ENV,
KOMPRESS_TIME_BUDGET_ENV,
KOMPRESS_CANARY_THRESHOLD_ENV,
KOMPRESS_EXECUTION_SEMAPHORE_WAIT_MS_ENV,
KOMPRESS_REQUEST_DEADLINE_ENV,
):
monkeypatch.delenv(env, raising=False)
yield
@ -98,6 +103,31 @@ def _make_compressor(monkeypatch, model: FakeModel, **config_kwargs) -> Kompress
return compressor
def _make_block_tracking_semaphore(monkeypatch):
blocked = threading.Event()
class TrackingSemaphore:
def __init__(self):
self._inner = threading.BoundedSemaphore(1)
def acquire(self, blocking=True, timeout=None):
if not blocking:
return self._inner.acquire(blocking=False)
if not self._inner.acquire(blocking=False):
blocked.set()
if timeout is None:
return self._inner.acquire()
return self._inner.acquire(timeout=timeout)
return True
def release(self):
self._inner.release()
semaphore = TrackingSemaphore()
monkeypatch.setattr(kc, "_execution_semaphore", lambda *_args, **_kwargs: semaphore)
return semaphore, blocked
CONTENT_40_WORDS = " ".join(f"word{i}" for i in range(40))
@ -179,6 +209,163 @@ def test_stuck_semaphore_batch_passes_through(monkeypatch):
assert [r.compressed for r in results] == contents # all passthrough, no data loss
def test_default_wait_allows_queued_single(monkeypatch):
compressor = _make_compressor(monkeypatch, FakeModel())
stuck, blocked = _make_block_tracking_semaphore(monkeypatch)
assert stuck.acquire(timeout=0)
finished = threading.Event()
result_holder = {}
def _run():
result_holder["result"] = compressor.compress(CONTENT_40_WORDS)
finished.set()
worker = threading.Thread(target=_run)
worker.start()
released = False
try:
assert blocked.wait(timeout=1)
assert not finished.wait(timeout=0.05)
stuck.release()
released = True
assert finished.wait(timeout=1)
finally:
if not released:
stuck.release()
worker.join(timeout=1)
assert not worker.is_alive()
result = result_holder["result"]
assert result.compressed != CONTENT_40_WORDS
assert result.compressed_tokens == 20
def test_default_wait_allows_queued_batch(monkeypatch):
compressor = _make_compressor(monkeypatch, FakeModel())
monkeypatch.setattr(KompressCompressor, "_should_use_sequential_fallback", lambda self: False)
stuck, blocked = _make_block_tracking_semaphore(monkeypatch)
assert stuck.acquire(timeout=0)
finished = threading.Event()
result_holder = {}
contents = [CONTENT_40_WORDS, " ".join(f"x{i}" for i in range(30))]
def _run():
result_holder["results"] = compressor.compress_batch(contents)
finished.set()
worker = threading.Thread(target=_run)
worker.start()
released = False
try:
assert blocked.wait(timeout=1)
assert not finished.wait(timeout=0.05)
stuck.release()
released = True
assert finished.wait(timeout=1)
finally:
if not released:
stuck.release()
worker.join(timeout=1)
assert not worker.is_alive()
results = result_holder["results"]
assert [result.compressed_tokens for result in results] == [20, 15]
def test_default_max_concurrent():
assert kc._default_max_concurrent("onnx", "onnx") == 1
assert kc._default_max_concurrent("pytorch", "cpu") == 1
assert kc._default_max_concurrent("pytorch", "cuda") == 1
def test_execution_wait_budget(monkeypatch):
assert kc._execution_wait_budget_seconds() == 3.0
monkeypatch.setenv(KOMPRESS_EXECUTION_SEMAPHORE_WAIT_MS_ENV, "bogus")
assert kc._execution_wait_budget_seconds() == 3.0
monkeypatch.setenv(KOMPRESS_EXECUTION_SEMAPHORE_WAIT_MS_ENV, "-1")
assert kc._execution_wait_budget_seconds() == 0.0
def test_request_deadline_caps_default_wait_single(monkeypatch):
monkeypatch.setenv(KOMPRESS_REQUEST_DEADLINE_ENV, "10")
model = FakeModel()
compressor = _make_compressor(monkeypatch, model)
stuck = kc._execution_semaphore("onnx", "onnx")
assert stuck.acquire(timeout=0)
try:
started = time.monotonic()
result = compressor.compress(CONTENT_40_WORDS)
elapsed = time.monotonic() - started
finally:
stuck.release()
assert elapsed < 0.2
assert result.compressed == CONTENT_40_WORDS
assert model.calls == 0
def test_request_deadline_caps_default_wait_batch(monkeypatch):
monkeypatch.setenv(KOMPRESS_REQUEST_DEADLINE_ENV, "10")
model = FakeModel()
compressor = _make_compressor(monkeypatch, model)
monkeypatch.setattr(KompressCompressor, "_should_use_sequential_fallback", lambda self: False)
stuck = kc._execution_semaphore("onnx", "onnx")
assert stuck.acquire(timeout=0)
contents = [CONTENT_40_WORDS, " ".join(f"x{i}" for i in range(30))]
try:
started = time.monotonic()
results = compressor.compress_batch(contents)
elapsed = time.monotonic() - started
finally:
stuck.release()
assert elapsed < 0.2
assert [r.compressed for r in results] == contents
assert model.calls == 0
def test_carried_deadline_reaches_single_to_batch(monkeypatch):
monkeypatch.setenv(KOMPRESS_REQUEST_DEADLINE_ENV, "10")
model = FakeModel()
compressor = _make_compressor(monkeypatch, model)
load_state = {"calls": 0}
def fake_clock():
return 999.0 if load_state["calls"] >= 1 else 0.0
def fake_load(*_args, **_kwargs):
load_state["calls"] += 1
return model, FakeTokenizer(), "onnx"
monkeypatch.setattr(kc.time, "perf_counter", fake_clock)
monkeypatch.setattr(kc, "_load_kompress", fake_load)
monkeypatch.setattr(compressor, "_should_batch_single_content", lambda *_args, **_kwargs: True)
monkeypatch.setattr(compressor, "_should_use_sequential_fallback", lambda: False)
result = compressor.compress(CONTENT_40_WORDS)
assert result.compressed == CONTENT_40_WORDS
assert model.calls == 0
def test_carried_deadline_reaches_sequential_fallback(monkeypatch):
monkeypatch.setenv(KOMPRESS_REQUEST_DEADLINE_ENV, "10")
model = FakeModel()
compressor = _make_compressor(monkeypatch, model, chunk_words=40)
monkeypatch.setattr(kc.time, "perf_counter", lambda: 999.0 if model.calls >= 1 else 0.0)
monkeypatch.setattr(compressor, "_should_batch_single_content", lambda *_args, **_kwargs: False)
monkeypatch.setattr(compressor, "_should_use_sequential_fallback", lambda: True)
contents = [CONTENT_40_WORDS, " ".join(f"x{i}" for i in range(30))]
results = compressor.compress_batch(contents)
assert results[0].compressed != contents[0]
assert results[1].compressed == contents[1]
assert model.calls == 1
def test_acquire_bounded_unbounded_when_both_disabled():
semaphore = kc._execution_semaphore("onnx", "onnx")
assert kc._acquire_bounded(semaphore, None, None) is True