mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
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:
parent
a09ba6c087
commit
5bd2266f16
3 changed files with 306 additions and 10 deletions
|
|
@ -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
|
||||
]
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue