mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-10 14:27:00 -04:00
## Description Three cold-start / robustness gaps found while debugging a user report of **0.12% savings across 722 requests** (49.8M input tokens, 60,920 saved). Closes # ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) ## Changes Made ### 1. The artifact fallback was unreachable for run-time failures `_create_onnx_session` tries `int8-wo` → `fp32` → `int8`, and its docstring describes exactly this scenario — but it only skipped a candidate when `InferenceSession(...)` **construction** threw. The int8 weight-only artifact carries `MatMulNBits` with `bits=8`. ORT's CPU kernel only handles 8-bit through the prepacked MLAS path, so a build or ISA without an 8-bit `SQNBitGemm` kernel falls into `ComputeBUnpacked`, which hard-asserts `nbits_ == 4`. That raises on `session.run()` **after** construction succeeded — so the fp32 candidate was never reached and ML compression was dead for the process lifetime. The reported log has 207 consecutive failures over three days. A two-token `_smoke_run` inside the existing candidate loop makes the fallback fire. `onnxruntime>=1.16.0` is unpinned, so which side of this an install lands on is a lottery. ### 2. A broken model cost an inference on every request, forever The per-request handler logged a `WARNING` and passed through with no latch — 207 identical lines that read as noise rather than "ML compression is dead". Now latches to passthrough after **3 consecutive** failures (any success resets the count) with one actionable `ERROR` naming the artifact override. ### 3. The model download began on the first request, not at startup #2001 was right to move Kompress off the startup path — on RHEL/CentOS 7-family hosts, entering cached native init before the port binds segfaults in `libarrow`/jemalloc with no Python traceback (#1908), which no `try/except` can catch. **This PR does not touch that.** But #2001 left the ~4-minute *download* on the first request, with every request in that window silently uncompressed behind one "model not ready" warning. Downloading is separable from loading. `prefetch_kompress_artifacts` resolves the files over plain `huggingface_hub` HTTP and never constructs an `InferenceSession` or imports `transformers`, so startup can prefetch bytes without touching the boundary #1908 crashes on. Native load stays deferred, status stays `deferred`, and a test asserts no session is constructed during prefetch. ## Testing - [x] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check .`) - [x] Type checking passes (`mypy headroom`) - [x] New tests added for new functionality - [x] Manual testing performed ### Test Output ```text $ .venv/bin/ruff check headroom/ tests/test_kompress_failsafe.py tests/test_kompress_preload_deferral.py --exclude headroom/dashboard/templates All checks passed! $ .venv/bin/mypy headroom/ Success: no issues found in 508 source files $ python -m pytest tests/test_kompress_failsafe.py tests/test_kompress_preload_deferral.py \ tests/test_kompress_request_nonblocking.py tests/test_force_kompress_all.py \ tests/test_kompress_must_keep.py tests/test_proxy_disable_kompress.py \ tests/test_proxy_per_provider_kompress.py tests/test_proxy_warmup.py \ tests/test_proxy_eager_preload_bind.py -q 95 passed in 10.51s ``` ## Real Behavior Proof - **Environment:** macOS 26.4 arm64, Python 3.12.6, onnxruntime 1.21.1, repo `.venv`. **(1) Fallback chain, against the real HF repo:** ```text WARNING ONNX artifact 'onnx/kompress-int8-wo.onnx' from chopratejas/kompress-v2-base is unusable (... nbits_ == 4 was false ...); trying next candidate SESSION OK -> ['input_ids', 'attention_mask'] SMOKE RUN OK on the selected artifact ``` Also confirmed the default artifact really is 8-bit, by loading the cached blob: `{'bits': [8], 'block_size': [128]}`. **(2) Files-only prefetch, with `InferenceSession` patched to raise:** ```text INFO Kompress: prefetching model artifacts for chopratejas/kompress-v2-base ... prefetch ok=True in 0.08s, no session constructed ``` - **Not tested / important caveat:** the user's exact failure **cannot be reproduced on this machine**. On ORT 1.21.1 arm64 the int8-wo artifact fails at *construction* (`matmul_nbits.cc:115`), which the pre-existing load-only fallback already caught. Their build fails at *execution* (`matmul_nbits.cc:442`, `ComputeBUnpacked`). So the run-time path is pinned with a fake ORT session that constructs fine and then rejects `run()` — a mechanism test, not a reproduction of their build. Confirming the fix on their host needs their `onnxruntime` version. - **Not tested:** no RHEL/CentOS 7 host available to re-verify #1908 non-regression; the argument is structural (prefetch never constructs a session) and asserted by `test_prefetch_never_constructs_a_session_or_imports_transformers`. ## 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 - [ ] 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 did **not** edit `CHANGELOG.md` 🤖 Generated with [Claude Code](https://claude.com/claude-code)
680 lines
24 KiB
Python
680 lines
24 KiB
Python
"""Fail-safe behavior for Kompress on degraded machines.
|
|
|
|
Reproduces the Windows incident where one pathologically slow ONNX inference
|
|
held the execution semaphore forever: every later compression blocked on an
|
|
unbounded acquire, every request hit the proxy's 30s stage timeout, and the
|
|
proxy delivered 0% savings plus +30s latency until restart. These tests pin
|
|
the three layers of defense (bounded acquire, wall-clock budget, preload
|
|
canary) and that the normal fast path is untouched.
|
|
|
|
No ML dependencies — the model/tokenizer are fakes injected via
|
|
``_load_kompress``.
|
|
"""
|
|
|
|
import threading
|
|
import time
|
|
|
|
import pytest
|
|
|
|
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,
|
|
)
|
|
|
|
|
|
class FakeEncoding:
|
|
"""Mimics a transformers BatchEncoding for is_split_into_words inputs.
|
|
|
|
One token per word, no special tokens — word_ids(i) is identity.
|
|
"""
|
|
|
|
def __init__(self, rows: list[list[str]]):
|
|
self._rows = rows
|
|
|
|
def __getitem__(self, key: str):
|
|
if key == "input_ids":
|
|
return [[0] * len(r) for r in self._rows]
|
|
if key == "attention_mask":
|
|
return [[1] * len(r) for r in self._rows]
|
|
raise KeyError(key)
|
|
|
|
def word_ids(self, batch_index: int = 0):
|
|
return list(range(len(self._rows[batch_index])))
|
|
|
|
|
|
class FakeTokenizer:
|
|
def __call__(self, words, **kwargs):
|
|
# is_split_into_words inputs: either one word list or a batch of them.
|
|
rows = words if words and isinstance(words[0], list) else [words]
|
|
return FakeEncoding(rows)
|
|
|
|
|
|
class FakeModel:
|
|
"""Keeps every other word; optional per-call delay to simulate slowness."""
|
|
|
|
def __init__(self, delay: float = 0.0):
|
|
self.delay = delay
|
|
self.calls = 0
|
|
|
|
def _tick(self):
|
|
self.calls += 1
|
|
if self.delay:
|
|
time.sleep(self.delay)
|
|
|
|
def get_keep_mask(self, input_ids, attention_mask):
|
|
self._tick()
|
|
return [[i % 2 == 0 for i in range(len(row))] for row in input_ids]
|
|
|
|
def get_scores(self, input_ids, attention_mask):
|
|
self._tick()
|
|
return [[1.0 if i % 2 == 0 else 0.0 for i in range(len(row))] for row in input_ids]
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _reset_module_state(monkeypatch):
|
|
kc._execution_semaphores.clear()
|
|
monkeypatch.setattr(kc, "_giveup_warned", False)
|
|
for env in (
|
|
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
|
|
kc._execution_semaphores.clear()
|
|
|
|
|
|
def _make_compressor(monkeypatch, model: FakeModel, **config_kwargs) -> KompressCompressor:
|
|
config_kwargs.setdefault("enable_ccr", False)
|
|
compressor = KompressCompressor(config=KompressConfig(**config_kwargs))
|
|
monkeypatch.setattr(
|
|
kc,
|
|
"_load_kompress",
|
|
lambda model_id, device="auto", **kwargs: (model, FakeTokenizer(), "onnx"),
|
|
)
|
|
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))
|
|
|
|
|
|
# ── Normal path: behavior and performance must be unchanged ───────────
|
|
|
|
|
|
def test_fast_model_compresses_normally(monkeypatch):
|
|
model = FakeModel()
|
|
compressor = _make_compressor(monkeypatch, model)
|
|
|
|
result = compressor.compress(CONTENT_40_WORDS)
|
|
|
|
assert result.compressed != CONTENT_40_WORDS
|
|
assert result.compressed_tokens == 20 # every other word kept
|
|
assert result.compression_ratio == 0.5
|
|
assert model.calls == 1
|
|
|
|
|
|
def test_fast_model_releases_semaphore(monkeypatch):
|
|
compressor = _make_compressor(monkeypatch, FakeModel())
|
|
compressor.compress(CONTENT_40_WORDS)
|
|
|
|
semaphore = kc._execution_semaphore("onnx", "onnx")
|
|
assert semaphore.acquire(timeout=0)
|
|
semaphore.release()
|
|
|
|
|
|
def test_semaphore_released_when_inference_raises(monkeypatch):
|
|
class ExplodingModel(FakeModel):
|
|
def get_keep_mask(self, input_ids, attention_mask):
|
|
raise RuntimeError("boom")
|
|
|
|
compressor = _make_compressor(monkeypatch, ExplodingModel())
|
|
result = compressor.compress(CONTENT_40_WORDS)
|
|
|
|
assert result.compressed == CONTENT_40_WORDS # passthrough, not an exception
|
|
semaphore = kc._execution_semaphore("onnx", "onnx")
|
|
assert semaphore.acquire(timeout=0)
|
|
semaphore.release()
|
|
|
|
|
|
# ── Bounded acquire: a stuck inference must not wedge other requests ──
|
|
|
|
|
|
def test_stuck_semaphore_passes_through_instead_of_blocking(monkeypatch):
|
|
monkeypatch.setenv(KOMPRESS_ACQUIRE_TIMEOUT_ENV, "0.1")
|
|
model = FakeModel()
|
|
compressor = _make_compressor(monkeypatch, model)
|
|
|
|
# Simulate the Windows incident: another thread holds the semaphore
|
|
# indefinitely (abandoned by its asyncio timeout but still running).
|
|
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 result.compressed == CONTENT_40_WORDS
|
|
assert model.calls == 0
|
|
assert elapsed < 3.0 # used to block forever
|
|
|
|
|
|
def test_stuck_semaphore_batch_passes_through(monkeypatch):
|
|
monkeypatch.setenv(KOMPRESS_ACQUIRE_TIMEOUT_ENV, "0.1")
|
|
compressor = _make_compressor(monkeypatch, FakeModel())
|
|
monkeypatch.setattr(KompressCompressor, "_should_use_sequential_fallback", lambda self: False)
|
|
|
|
stuck = kc._execution_semaphore("onnx", "onnx")
|
|
assert stuck.acquire(timeout=0)
|
|
try:
|
|
contents = [CONTENT_40_WORDS, " ".join(f"x{i}" for i in range(30))]
|
|
results = compressor.compress_batch(contents)
|
|
finally:
|
|
stuck.release()
|
|
|
|
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
|
|
semaphore.release()
|
|
|
|
|
|
def test_acquire_bounded_negative_remaining_does_not_raise():
|
|
semaphore = kc._execution_semaphore("onnx", "onnx")
|
|
assert semaphore.acquire(timeout=0)
|
|
try:
|
|
assert kc._acquire_bounded(semaphore, 5.0, -1.0) is False
|
|
finally:
|
|
semaphore.release()
|
|
|
|
|
|
# ── Wall-clock budget: give up before the proxy's stage timeout ───────
|
|
|
|
|
|
def test_time_budget_bails_to_passthrough(monkeypatch):
|
|
monkeypatch.setenv(KOMPRESS_TIME_BUDGET_ENV, "0.2")
|
|
model = FakeModel(delay=0.15)
|
|
compressor = _make_compressor(monkeypatch, model, chunk_words=10)
|
|
|
|
result = compressor.compress(CONTENT_40_WORDS) # 4 chunks at ~0.15s each
|
|
|
|
assert result.compressed == CONTENT_40_WORDS
|
|
assert model.calls < 4 # bailed before processing every chunk
|
|
|
|
|
|
def test_time_budget_disabled_processes_all_chunks(monkeypatch):
|
|
monkeypatch.setenv(KOMPRESS_TIME_BUDGET_ENV, "0")
|
|
model = FakeModel(delay=0.01)
|
|
compressor = _make_compressor(monkeypatch, model, chunk_words=10)
|
|
|
|
result = compressor.compress(CONTENT_40_WORDS)
|
|
|
|
assert model.calls == 4
|
|
assert result.compression_ratio == 0.5
|
|
|
|
|
|
def test_time_budget_batch_keeps_completed_texts(monkeypatch):
|
|
"""Mid-queue bail: fully processed texts stay compressed; any text with
|
|
an unprocessed chunk passes through whole (never partially dropped)."""
|
|
monkeypatch.setenv(KOMPRESS_TIME_BUDGET_ENV, "0.2")
|
|
model = FakeModel(delay=0.25) # one batch alone exhausts the budget
|
|
compressor = _make_compressor(monkeypatch, model)
|
|
monkeypatch.setattr(KompressCompressor, "_should_use_sequential_fallback", lambda self: False)
|
|
|
|
contents = [
|
|
" ".join(f"a{i}" for i in range(20)),
|
|
" ".join(f"b{i}" for i in range(20)),
|
|
" ".join(f"c{i}" for i in range(20)),
|
|
]
|
|
results = compressor.compress_batch(contents, batch_size=1)
|
|
|
|
assert len(results) == 3
|
|
# First batch ran; later ones bailed to passthrough.
|
|
assert results[0].compression_ratio == 0.5
|
|
assert results[1].compressed == contents[1]
|
|
assert results[2].compressed == contents[2]
|
|
# Every result preserves all information (compressed or original).
|
|
for r in results:
|
|
assert r.compressed
|
|
|
|
|
|
# ── Preload canary: detect degraded runtimes before live traffic ──────
|
|
|
|
|
|
def _join_canary(compressor: KompressCompressor) -> None:
|
|
assert compressor._canary_thread is not None
|
|
compressor._canary_thread.join(timeout=10)
|
|
assert not compressor._canary_thread.is_alive()
|
|
|
|
|
|
def test_canary_disables_kompress_on_slow_inference(monkeypatch, caplog):
|
|
monkeypatch.setenv(KOMPRESS_CANARY_THRESHOLD_ENV, "0.05")
|
|
model = FakeModel(delay=0.15)
|
|
compressor = _make_compressor(monkeypatch, model)
|
|
|
|
with caplog.at_level("WARNING"):
|
|
backend = compressor.preload()
|
|
_join_canary(compressor)
|
|
|
|
assert backend == "onnx"
|
|
assert compressor._degraded_reason is not None
|
|
assert model.calls == 2 # probe + one retry
|
|
assert "DISABLED" in caplog.text
|
|
|
|
result = compressor.compress(CONTENT_40_WORDS)
|
|
assert result.compressed == CONTENT_40_WORDS
|
|
assert model.calls == 2 # model never touched again
|
|
|
|
batch = compressor.compress_batch([CONTENT_40_WORDS])
|
|
assert batch[0].compressed == CONTENT_40_WORDS
|
|
assert model.calls == 2
|
|
|
|
|
|
def test_canary_fast_inference_stays_enabled(monkeypatch):
|
|
monkeypatch.setenv(KOMPRESS_CANARY_THRESHOLD_ENV, "5")
|
|
model = FakeModel()
|
|
compressor = _make_compressor(monkeypatch, model)
|
|
|
|
compressor.preload()
|
|
_join_canary(compressor)
|
|
|
|
assert compressor._degraded_reason is None
|
|
result = compressor.compress(CONTENT_40_WORDS)
|
|
assert result.compression_ratio == 0.5
|
|
|
|
|
|
def test_canary_retry_forgives_oneoff_warmup_slowness(monkeypatch):
|
|
"""First inference pays one-off warmup costs; only a slow retry condemns."""
|
|
monkeypatch.setenv(KOMPRESS_CANARY_THRESHOLD_ENV, "0.1")
|
|
|
|
class WarmupModel(FakeModel):
|
|
def get_keep_mask(self, input_ids, attention_mask):
|
|
self.calls += 1
|
|
if self.calls == 1:
|
|
time.sleep(0.2) # cold first run
|
|
return [[i % 2 == 0 for i in range(len(row))] for row in input_ids]
|
|
|
|
model = WarmupModel()
|
|
compressor = _make_compressor(monkeypatch, model)
|
|
compressor.preload()
|
|
_join_canary(compressor)
|
|
|
|
assert compressor._degraded_reason is None
|
|
assert model.calls == 2
|
|
|
|
|
|
def test_canary_disabled_via_env(monkeypatch):
|
|
monkeypatch.setenv(KOMPRESS_CANARY_THRESHOLD_ENV, "0")
|
|
model = FakeModel(delay=0.2)
|
|
compressor = _make_compressor(monkeypatch, model)
|
|
|
|
compressor.preload()
|
|
|
|
assert compressor._canary_thread is None # probe never scheduled
|
|
assert model.calls == 0
|
|
assert compressor._degraded_reason is None
|
|
|
|
|
|
def test_preload_does_not_block_on_slow_canary(monkeypatch):
|
|
"""The probe runs off the startup path: preload blocks proxy boot (the
|
|
HTTP server binds after it), and a slow probe once pushed the wrap-e2e
|
|
container past its 30s health-check timeout."""
|
|
monkeypatch.setenv(KOMPRESS_CANARY_THRESHOLD_ENV, "0.05")
|
|
model = FakeModel(delay=1.0)
|
|
compressor = _make_compressor(monkeypatch, model)
|
|
|
|
started = time.monotonic()
|
|
compressor.preload()
|
|
preload_elapsed = time.monotonic() - started
|
|
|
|
assert preload_elapsed < 0.5 # returns before the ~2s of probe inference
|
|
_join_canary(compressor)
|
|
assert compressor._degraded_reason is not None
|
|
|
|
|
|
def test_canary_probe_error_never_breaks_preload(monkeypatch):
|
|
class ExplodingModel(FakeModel):
|
|
def get_keep_mask(self, input_ids, attention_mask):
|
|
raise RuntimeError("probe boom")
|
|
|
|
compressor = _make_compressor(monkeypatch, ExplodingModel())
|
|
|
|
assert compressor.preload() == "onnx"
|
|
_join_canary(compressor)
|
|
assert compressor._degraded_reason is None
|
|
|
|
|
|
# ── Artifact selection: reject at LOAD what would fail at RUN ──────────────────
|
|
# Reported case: the int8 weight-only artifact carries MatMulNBits with bits=8.
|
|
# ORT's CPU kernel only handles 8-bit via the prepacked MLAS path, so a build
|
|
# without an 8-bit SQNBitGemm kernel falls into ComputeBUnpacked, which asserts
|
|
# nbits_ == 4. That raises on session.run() AFTER construction succeeded, so the
|
|
# load-only candidate loop never saw it and the fp32 fallback was unreachable:
|
|
# 207 consecutive per-request failures over three days, ML compression silently
|
|
# dead the whole time.
|
|
|
|
|
|
class _FakeOrtSession:
|
|
"""Constructs fine; optionally rejects execution the way ORT's CPU kernel does."""
|
|
|
|
def __init__(self, path: str, *, fails_at_run: bool):
|
|
self.path = path
|
|
self._fails_at_run = fails_at_run
|
|
self.runs = 0
|
|
|
|
def run(self, outputs, feeds):
|
|
self.runs += 1
|
|
if self._fails_at_run:
|
|
raise RuntimeError(
|
|
"[ONNXRuntimeError] : 6 : RUNTIME_EXCEPTION : Non-zero status code "
|
|
"returned while running MatMulNBits node ... nbits_ == 4 was false. "
|
|
"Only 4b quantization is supported for unpacked compute."
|
|
)
|
|
import numpy as np
|
|
|
|
return [np.zeros((1, 2), dtype=np.float32)]
|
|
|
|
|
|
def _install_fake_ort(monkeypatch, *, run_fails_for: set[str]):
|
|
"""Patch onnxruntime so InferenceSession succeeds but run() may not."""
|
|
created: list[_FakeOrtSession] = []
|
|
|
|
class _FakeOrt:
|
|
@staticmethod
|
|
def SessionOptions(): # noqa: N802 - mirrors the ORT API
|
|
return object()
|
|
|
|
@staticmethod
|
|
def InferenceSession(path, options=None, providers=None): # noqa: N802
|
|
session = _FakeOrtSession(path, fails_at_run=any(bad in path for bad in run_fails_for))
|
|
created.append(session)
|
|
return session
|
|
|
|
monkeypatch.setitem(__import__("sys").modules, "onnxruntime", _FakeOrt)
|
|
monkeypatch.setattr(kc, "_onnx_session_options", lambda _ort: object())
|
|
monkeypatch.setattr(kc, "hf_hub_download_local_first", lambda repo, fn, **kw: f"/cache/{fn}")
|
|
return created
|
|
|
|
|
|
def test_run_time_artifact_rejection_falls_through_to_next_candidate(monkeypatch, caplog):
|
|
"""A session that loads then fails at run must be skipped, not returned."""
|
|
created = _install_fake_ort(monkeypatch, run_fails_for={"int8-wo"})
|
|
|
|
with caplog.at_level("WARNING"):
|
|
session = kc._create_onnx_session("org/model", ["CPUExecutionProvider"])
|
|
|
|
# int8-wo was constructed, smoke-run, rejected; fp32 was selected instead.
|
|
assert "int8-wo" in created[0].path
|
|
assert created[0].runs == 1
|
|
assert "kompress-fp32.onnx" in session.path
|
|
assert "unusable" in caplog.text
|
|
|
|
|
|
def test_healthy_artifact_is_selected_after_one_smoke_run(monkeypatch):
|
|
created = _install_fake_ort(monkeypatch, run_fails_for=set())
|
|
|
|
session = kc._create_onnx_session("org/model", ["CPUExecutionProvider"])
|
|
|
|
# First candidate works, so no fallback and exactly one probe.
|
|
assert session is created[0]
|
|
assert len(created) == 1
|
|
assert session.runs == 1
|
|
|
|
|
|
def test_all_artifacts_failing_at_run_raises_rather_than_returning_a_dead_session(monkeypatch):
|
|
_install_fake_ort(monkeypatch, run_fails_for={"onnx/"})
|
|
|
|
with pytest.raises(FileNotFoundError, match="No loadable ONNX artifact"):
|
|
kc._create_onnx_session("org/model", ["CPUExecutionProvider"])
|
|
|
|
|
|
# ── Failure latch: a broken model stops costing us every request ───────────────
|
|
|
|
|
|
def test_repeated_inference_failures_latch_to_passthrough(monkeypatch, caplog):
|
|
class AlwaysFailingModel(FakeModel):
|
|
def get_keep_mask(self, input_ids, attention_mask):
|
|
self._tick()
|
|
raise RuntimeError("MatMulNBits nbits_ == 4 was false")
|
|
|
|
model = AlwaysFailingModel()
|
|
compressor = _make_compressor(monkeypatch, model)
|
|
monkeypatch.setenv(KOMPRESS_CANARY_THRESHOLD_ENV, "0") # no canary interference
|
|
|
|
with caplog.at_level("WARNING"):
|
|
for _ in range(kc._INFERENCE_FAILURE_LATCH):
|
|
assert compressor.compress(CONTENT_40_WORDS).compressed == CONTENT_40_WORDS
|
|
|
|
assert compressor._degraded_reason is not None
|
|
assert "DISABLED" in caplog.text
|
|
calls_at_latch = model.calls
|
|
|
|
# Latched: further calls short-circuit without touching the model again, so a
|
|
# broken artifact can't burn inference on every request for three days.
|
|
assert compressor.compress(CONTENT_40_WORDS).compressed == CONTENT_40_WORDS
|
|
assert model.calls == calls_at_latch
|
|
|
|
|
|
def test_a_success_resets_the_failure_count(monkeypatch):
|
|
class FlakyModel(FakeModel):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.fail_next = True
|
|
|
|
def get_keep_mask(self, input_ids, attention_mask):
|
|
if self.fail_next:
|
|
self._tick()
|
|
raise RuntimeError("transient")
|
|
return super().get_keep_mask(input_ids, attention_mask)
|
|
|
|
model = FlakyModel()
|
|
compressor = _make_compressor(monkeypatch, model)
|
|
monkeypatch.setenv(KOMPRESS_CANARY_THRESHOLD_ENV, "0")
|
|
|
|
# Two failures, then a success, then two more failures: never 3 in a row.
|
|
for _ in range(kc._INFERENCE_FAILURE_LATCH - 1):
|
|
compressor.compress(CONTENT_40_WORDS)
|
|
assert compressor._inference_failures == kc._INFERENCE_FAILURE_LATCH - 1
|
|
|
|
model.fail_next = False
|
|
compressor.compress(CONTENT_40_WORDS)
|
|
assert compressor._inference_failures == 0
|
|
assert compressor._degraded_reason is None
|
|
|
|
model.fail_next = True
|
|
for _ in range(kc._INFERENCE_FAILURE_LATCH - 1):
|
|
compressor.compress(CONTENT_40_WORDS)
|
|
assert compressor._degraded_reason is None
|