mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
fix(kompress): reject artifacts that fail at run, and prefetch model files at startup (#2740)
## 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)
This commit is contained in:
parent
8262a4a321
commit
224578e80b
4 changed files with 393 additions and 11 deletions
|
|
@ -4047,6 +4047,22 @@ class ContentRouter(Transform):
|
||||||
logger.debug("HTMLExtractor not available (install trafilatura)")
|
logger.debug("HTMLExtractor not available (install trafilatura)")
|
||||||
return self._html_extractor
|
return self._html_extractor
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _prefetch_kompress_artifacts_async(kompress_config: Any) -> bool:
|
||||||
|
"""Start a background download of the Kompress model files, if needed.
|
||||||
|
|
||||||
|
Files only — see ``prefetch_kompress_artifacts`` for why startup must not
|
||||||
|
build the model. Returns ``True`` when a prefetch is running.
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
from .kompress_compressor import HF_MODEL_ID, ensure_background_prefetch
|
||||||
|
|
||||||
|
model_id = getattr(kompress_config, "model_id", None) or HF_MODEL_ID
|
||||||
|
return ensure_background_prefetch(str(model_id))
|
||||||
|
except Exception as e: # pragma: no cover - defensive; never break startup
|
||||||
|
logger.debug("Kompress artifact prefetch skipped: %s", e)
|
||||||
|
return False
|
||||||
|
|
||||||
def eager_load_compressors(self) -> dict[str, str]:
|
def eager_load_compressors(self) -> dict[str, str]:
|
||||||
"""Pre-load compressors at startup to avoid first-request latency.
|
"""Pre-load compressors at startup to avoid first-request latency.
|
||||||
|
|
||||||
|
|
@ -4061,7 +4077,18 @@ class ContentRouter(Transform):
|
||||||
# 1. ML text compressor: Kompress.
|
# 1. ML text compressor: Kompress.
|
||||||
#
|
#
|
||||||
# Native model initialization stays out of the blocking startup/lifespan
|
# Native model initialization stays out of the blocking startup/lifespan
|
||||||
# path. The existing lazy request path loads Kompress on first use.
|
# path. The existing lazy request path loads Kompress on first use. This is
|
||||||
|
# load-bearing, NOT laziness: on RHEL/CentOS 7-family hosts entering cached
|
||||||
|
# Kompress native init before the port binds segfaults in libarrow/jemalloc
|
||||||
|
# with no Python traceback (#1908, fixed by #2001) — a crash no try/except
|
||||||
|
# can catch. Do not call `preload()` here.
|
||||||
|
#
|
||||||
|
# What we CAN do at startup is prefetch the model FILES. Downloading is
|
||||||
|
# pure huggingface_hub HTTP — no ONNX session, no transformers import, so it
|
||||||
|
# never touches the native path that #1908 crashes on. That removes the real
|
||||||
|
# cold-start cost: previously the ~4-minute download began on the FIRST
|
||||||
|
# REQUEST, and every request in that window went silently uncompressed
|
||||||
|
# behind a single "model not ready" warning.
|
||||||
if self.config.enable_kompress:
|
if self.config.enable_kompress:
|
||||||
compressor = self._get_kompress()
|
compressor = self._get_kompress()
|
||||||
if compressor:
|
if compressor:
|
||||||
|
|
@ -4069,8 +4096,10 @@ class ContentRouter(Transform):
|
||||||
status["kompress"] = "enabled"
|
status["kompress"] = "enabled"
|
||||||
status["kompress_backend"] = "unknown"
|
status["kompress_backend"] = "unknown"
|
||||||
else:
|
else:
|
||||||
logger.info("Kompress model preload deferred until first request")
|
|
||||||
status["kompress"] = "deferred"
|
status["kompress"] = "deferred"
|
||||||
|
if self._prefetch_kompress_artifacts_async(getattr(compressor, "config", None)):
|
||||||
|
status["kompress_artifacts"] = "prefetching"
|
||||||
|
logger.info("Kompress model preload deferred until first request")
|
||||||
else:
|
else:
|
||||||
status["kompress"] = "unavailable"
|
status["kompress"] = "unavailable"
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -96,6 +96,12 @@ KOMPRESS_ONNX_INTRA_THREADS_ENV = "HEADROOM_KOMPRESS_ONNX_INTRA_THREADS"
|
||||||
KOMPRESS_ONNX_INTER_THREADS_ENV = "HEADROOM_KOMPRESS_ONNX_INTER_THREADS"
|
KOMPRESS_ONNX_INTER_THREADS_ENV = "HEADROOM_KOMPRESS_ONNX_INTER_THREADS"
|
||||||
KOMPRESS_COREML_CACHE_DIR_ENV = "HEADROOM_KOMPRESS_COREML_CACHE_DIR"
|
KOMPRESS_COREML_CACHE_DIR_ENV = "HEADROOM_KOMPRESS_COREML_CACHE_DIR"
|
||||||
KOMPRESS_MAX_CONCURRENT_ENV = "HEADROOM_KOMPRESS_MAX_CONCURRENT"
|
KOMPRESS_MAX_CONCURRENT_ENV = "HEADROOM_KOMPRESS_MAX_CONCURRENT"
|
||||||
|
# Consecutive inference failures before Kompress latches to passthrough for the
|
||||||
|
# rest of the process. 3 rides out a transient error while still catching a model
|
||||||
|
# that is broken for this install on the first few requests rather than the 200th.
|
||||||
|
# ponytail: fixed count, not a rate window — a broken artifact fails every call,
|
||||||
|
# so there is nothing a window would tell us that three strikes doesn't.
|
||||||
|
_INFERENCE_FAILURE_LATCH = 3
|
||||||
KOMPRESS_EXECUTION_SEMAPHORE_WAIT_MS_ENV = "HEADROOM_KOMPRESS_EXECUTION_TIMEOUT_MS"
|
KOMPRESS_EXECUTION_SEMAPHORE_WAIT_MS_ENV = "HEADROOM_KOMPRESS_EXECUTION_TIMEOUT_MS"
|
||||||
KOMPRESS_EXECUTION_SEMAPHORE_WAIT_MS_DEFAULT = 3000
|
KOMPRESS_EXECUTION_SEMAPHORE_WAIT_MS_DEFAULT = 3000
|
||||||
KOMPRESS_BATCH_SIZE_ENV = "HEADROOM_KOMPRESS_BATCH_SIZE"
|
KOMPRESS_BATCH_SIZE_ENV = "HEADROOM_KOMPRESS_BATCH_SIZE"
|
||||||
|
|
@ -600,15 +606,42 @@ def _onnx_filename_candidates() -> tuple[str, ...]:
|
||||||
return _DEFAULT_ONNX_FILENAMES
|
return _DEFAULT_ONNX_FILENAMES
|
||||||
|
|
||||||
|
|
||||||
|
def _smoke_run(session: Any) -> None:
|
||||||
|
"""Run one tiny forward pass so a broken artifact fails HERE, not per request.
|
||||||
|
|
||||||
|
Some ONNX Runtime builds accept a session and then reject it at execution.
|
||||||
|
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 a load-only check never sees it and the fp32
|
||||||
|
fallback below is unreachable. Observed in the wild as 207 consecutive
|
||||||
|
per-request failures over three days with ML compression silently dead.
|
||||||
|
|
||||||
|
Two tokens through the real graph, so it costs milliseconds rather than the
|
||||||
|
seconds the timed canary takes (kernel dispatch is what fails, not compute).
|
||||||
|
"""
|
||||||
|
import numpy as np
|
||||||
|
|
||||||
|
session.run(
|
||||||
|
["final_scores"],
|
||||||
|
{
|
||||||
|
"input_ids": np.zeros((1, 2), dtype=np.int64),
|
||||||
|
"attention_mask": np.ones((1, 2), dtype=np.int64),
|
||||||
|
},
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
def _create_onnx_session(
|
def _create_onnx_session(
|
||||||
model_id: str, providers: list[Any], *, allow_download: bool = True
|
model_id: str, providers: list[Any], *, allow_download: bool = True
|
||||||
) -> Any:
|
) -> Any:
|
||||||
"""Resolve and load the model's ONNX artifact, trying candidates in order.
|
"""Resolve and load the model's ONNX artifact, trying candidates in order.
|
||||||
|
|
||||||
A candidate is skipped on download miss (file not in the repo) or on
|
A candidate is skipped on download miss (file not in the repo), on
|
||||||
session-load failure (e.g. the weight-only int8 artifact uses the
|
session-load failure, or on smoke-run failure (e.g. the weight-only int8
|
||||||
MatMulNBits contrib op, which old onnxruntime builds can't run — those
|
artifact uses the MatMulNBits contrib op, which some onnxruntime builds
|
||||||
installs fall through to the fp32 artifact instead of losing Kompress).
|
accept at load and then reject at execution — those installs fall through to
|
||||||
|
the fp32 artifact instead of losing Kompress). See :func:`_smoke_run`.
|
||||||
|
|
||||||
When ``allow_download`` is ``False`` candidates are resolved from the local
|
When ``allow_download`` is ``False`` candidates are resolved from the local
|
||||||
cache only; if none is cached, :class:`KompressModelNotCached` is raised
|
cache only; if none is cached, :class:`KompressModelNotCached` is raised
|
||||||
|
|
@ -633,15 +666,17 @@ def _create_onnx_session(
|
||||||
|
|
||||||
ort = onnxruntime
|
ort = onnxruntime
|
||||||
try:
|
try:
|
||||||
return ort.InferenceSession(
|
session = ort.InferenceSession(
|
||||||
onnx_path,
|
onnx_path,
|
||||||
_onnx_session_options(ort),
|
_onnx_session_options(ort),
|
||||||
providers=providers,
|
providers=providers,
|
||||||
)
|
)
|
||||||
|
_smoke_run(session)
|
||||||
|
return session
|
||||||
except Exception as exc:
|
except Exception as exc:
|
||||||
last_err = exc
|
last_err = exc
|
||||||
logger.warning(
|
logger.warning(
|
||||||
"ONNX artifact %r from %s failed to load (%s); trying next candidate",
|
"ONNX artifact %r from %s is unusable (%s); trying next candidate",
|
||||||
filename,
|
filename,
|
||||||
model_id,
|
model_id,
|
||||||
exc,
|
exc,
|
||||||
|
|
@ -1041,6 +1076,70 @@ def ensure_background_download(model_id: str = HF_MODEL_ID, device: str = "auto"
|
||||||
thread.start()
|
thread.start()
|
||||||
|
|
||||||
|
|
||||||
|
def prefetch_kompress_artifacts(model_id: str = HF_MODEL_ID) -> bool:
|
||||||
|
"""Download the model's ONNX artifact to the local cache. No native init.
|
||||||
|
|
||||||
|
Deliberately weaker than :func:`warm_kompress_model`: it resolves files over
|
||||||
|
plain huggingface_hub HTTP and never constructs an ``InferenceSession`` or
|
||||||
|
imports ``transformers``. That distinction is the whole point — entering
|
||||||
|
Kompress *native* init on the proxy's startup path segfaults in
|
||||||
|
libarrow/jemalloc on RHEL/CentOS 7-family hosts (#1908, fixed by #2001), so
|
||||||
|
startup may prefetch bytes but must not build the model.
|
||||||
|
|
||||||
|
Stops at the first candidate that resolves: the loader tries them in the same
|
||||||
|
order, so fetching the rest would be wasted bandwidth.
|
||||||
|
|
||||||
|
Returns ``True`` if an artifact is now cached locally.
|
||||||
|
"""
|
||||||
|
if model_id in _kompress_cache:
|
||||||
|
return True
|
||||||
|
for filename in _onnx_filename_candidates():
|
||||||
|
try:
|
||||||
|
hf_hub_download_local_first(model_id, filename, allow_network=True)
|
||||||
|
return True
|
||||||
|
except Exception as exc:
|
||||||
|
logger.debug("Kompress prefetch: %r unavailable for %s: %s", filename, model_id, exc)
|
||||||
|
return False
|
||||||
|
|
||||||
|
|
||||||
|
def ensure_background_prefetch(model_id: str = HF_MODEL_ID) -> bool:
|
||||||
|
"""Start a one-shot background artifact prefetch. Non-blocking, idempotent.
|
||||||
|
|
||||||
|
Returns ``True`` when a prefetch is running or was started, ``False`` when the
|
||||||
|
model is already cached (nothing to do) or Kompress isn't installed. Shares the
|
||||||
|
per-model thread registry with :func:`ensure_background_download` so the two
|
||||||
|
can't race to fetch the same files.
|
||||||
|
"""
|
||||||
|
if not is_kompress_available() or model_id in _kompress_cache:
|
||||||
|
return False
|
||||||
|
with _download_threads_lock:
|
||||||
|
if model_id in _kompress_cache:
|
||||||
|
return False
|
||||||
|
existing = _download_threads.get(model_id)
|
||||||
|
if existing is not None and existing.is_alive():
|
||||||
|
return True
|
||||||
|
|
||||||
|
def _run() -> None:
|
||||||
|
logger.info("Kompress: prefetching model artifacts for %s ...", model_id)
|
||||||
|
if prefetch_kompress_artifacts(model_id):
|
||||||
|
logger.info(
|
||||||
|
"Kompress: artifact prefetch complete for %s; the model loads on "
|
||||||
|
"first use without a download stall.",
|
||||||
|
model_id,
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
logger.warning("Kompress: artifact prefetch found no usable file for %s", model_id)
|
||||||
|
|
||||||
|
thread = threading.Thread(
|
||||||
|
target=_run,
|
||||||
|
name=f"kompress-prefetch-{model_id.replace('/', '-')}",
|
||||||
|
daemon=True,
|
||||||
|
)
|
||||||
|
_download_threads[model_id] = thread
|
||||||
|
thread.start()
|
||||||
|
return True
|
||||||
|
|
||||||
|
|
||||||
def warm_kompress_model(
|
def warm_kompress_model(
|
||||||
model_id: str = HF_MODEL_ID,
|
model_id: str = HF_MODEL_ID,
|
||||||
device: str = "cpu",
|
device: str = "cpu",
|
||||||
|
|
@ -1170,10 +1269,14 @@ class KompressCompressor(Transform):
|
||||||
|
|
||||||
def __init__(self, config: KompressConfig | None = None):
|
def __init__(self, config: KompressConfig | None = None):
|
||||||
self.config = config or KompressConfig()
|
self.config = config or KompressConfig()
|
||||||
# Set by the preload canary when inference is too slow to be useful;
|
# Set by the preload canary when inference is too slow to be useful, or by
|
||||||
# compress()/compress_batch() then pass content through untouched.
|
# the failure latch when inference raises repeatedly; compress()/
|
||||||
|
# compress_batch() then pass content through untouched.
|
||||||
self._degraded_reason: str | None = None
|
self._degraded_reason: str | None = None
|
||||||
self._canary_thread: threading.Thread | None = None
|
self._canary_thread: threading.Thread | None = None
|
||||||
|
# Consecutive inference failures — reset by any success, so a transient
|
||||||
|
# error can't accumulate toward the latch across a healthy run.
|
||||||
|
self._inference_failures: int = 0
|
||||||
|
|
||||||
def preload(self, *, allow_download: bool = True) -> str:
|
def preload(self, *, allow_download: bool = True) -> str:
|
||||||
"""Load the backing model/tokenizer and return the selected backend.
|
"""Load the backing model/tokenizer and return the selected backend.
|
||||||
|
|
@ -1563,6 +1666,9 @@ class KompressCompressor(Transform):
|
||||||
result.tokens_saved,
|
result.tokens_saved,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
# A real inference landed — clear the strike count so only CONSECUTIVE
|
||||||
|
# failures can reach the latch.
|
||||||
|
self._inference_failures = 0
|
||||||
return result
|
return result
|
||||||
|
|
||||||
except KompressModelNotCached:
|
except KompressModelNotCached:
|
||||||
|
|
@ -1572,9 +1678,41 @@ class KompressCompressor(Transform):
|
||||||
)
|
)
|
||||||
return self._passthrough(content, n_words)
|
return self._passthrough(content, n_words)
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.warning("Kompress compression failed: %s", e)
|
self._record_inference_failure(e)
|
||||||
return self._passthrough(content, n_words)
|
return self._passthrough(content, n_words)
|
||||||
|
|
||||||
|
def _record_inference_failure(self, exc: BaseException) -> None:
|
||||||
|
"""Log a failed inference, and latch to degraded after repeated failures.
|
||||||
|
|
||||||
|
A model that fails once may be transient; one that fails every call is
|
||||||
|
broken for this process and will never recover on its own. Without a latch
|
||||||
|
that state is a per-request WARNING forever — the reported case logged 207
|
||||||
|
identical lines across three days while every request silently went
|
||||||
|
uncompressed, which read as noise rather than "ML compression is dead".
|
||||||
|
Latching converts it into one actionable line plus a `/debug/warmup`
|
||||||
|
signal, and stops paying for a call that cannot succeed.
|
||||||
|
"""
|
||||||
|
self._inference_failures += 1
|
||||||
|
if self._degraded_reason is not None:
|
||||||
|
return
|
||||||
|
if self._inference_failures < _INFERENCE_FAILURE_LATCH:
|
||||||
|
logger.warning(
|
||||||
|
"Kompress compression failed (%d/%d before disabling): %s",
|
||||||
|
self._inference_failures,
|
||||||
|
_INFERENCE_FAILURE_LATCH,
|
||||||
|
exc,
|
||||||
|
)
|
||||||
|
return
|
||||||
|
self._degraded_reason = f"{self._inference_failures} consecutive inference failures: {exc}"
|
||||||
|
logger.error(
|
||||||
|
"Kompress inference failed %d times consecutively (%s) — ML compression "
|
||||||
|
"DISABLED for this run; content passes through uncompressed. Pin a working "
|
||||||
|
"ONNX artifact via %s=onnx/kompress-fp32.onnx, or report the error above.",
|
||||||
|
self._inference_failures,
|
||||||
|
exc,
|
||||||
|
KOMPRESS_ONNX_FILENAME_ENV,
|
||||||
|
)
|
||||||
|
|
||||||
def compress_batch(
|
def compress_batch(
|
||||||
self,
|
self,
|
||||||
contents: list[str],
|
contents: list[str],
|
||||||
|
|
|
||||||
|
|
@ -535,3 +535,146 @@ def test_canary_probe_error_never_breaks_preload(monkeypatch):
|
||||||
assert compressor.preload() == "onnx"
|
assert compressor.preload() == "onnx"
|
||||||
_join_canary(compressor)
|
_join_canary(compressor)
|
||||||
assert compressor._degraded_reason is None
|
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
|
||||||
|
|
|
||||||
|
|
@ -127,6 +127,9 @@ def test_eager_load_defers_kompress_regardless_of_cache_state(monkeypatch, cache
|
||||||
router = _router_kompress_only()
|
router = _router_kompress_only()
|
||||||
stub = _StubCompressor(cached=cache_state == "cached")
|
stub = _StubCompressor(cached=cache_state == "cached")
|
||||||
monkeypatch.setattr(router, "_get_kompress", lambda: stub)
|
monkeypatch.setattr(router, "_get_kompress", lambda: stub)
|
||||||
|
# Artifact prefetch is files-only, but it still reaches the network — keep it
|
||||||
|
# out of this assertion so the test stays about the native-preload boundary.
|
||||||
|
monkeypatch.setattr(router, "_prefetch_kompress_artifacts_async", lambda _cfg: False)
|
||||||
|
|
||||||
status = router.eager_load_compressors()
|
status = router.eager_load_compressors()
|
||||||
|
|
||||||
|
|
@ -164,6 +167,7 @@ def test_non_kompress_warmups_continue_when_kompress_is_deferred(monkeypatch):
|
||||||
router = _router_kompress_only()
|
router = _router_kompress_only()
|
||||||
stub = _StubCompressor(cached=True)
|
stub = _StubCompressor(cached=True)
|
||||||
monkeypatch.setattr(router, "_get_kompress", lambda: stub)
|
monkeypatch.setattr(router, "_get_kompress", lambda: stub)
|
||||||
|
monkeypatch.setattr(router, "_prefetch_kompress_artifacts_async", lambda _cfg: False)
|
||||||
monkeypatch.setattr("headroom.compression.detector._magika_available", lambda: True)
|
monkeypatch.setattr("headroom.compression.detector._magika_available", lambda: True)
|
||||||
monkeypatch.setattr("headroom.compression.detector._get_magika", lambda: object())
|
monkeypatch.setattr("headroom.compression.detector._get_magika", lambda: object())
|
||||||
|
|
||||||
|
|
@ -174,6 +178,74 @@ def test_non_kompress_warmups_continue_when_kompress_is_deferred(monkeypatch):
|
||||||
assert stub.preload_calls == []
|
assert stub.preload_calls == []
|
||||||
|
|
||||||
|
|
||||||
|
# ── Startup artifact prefetch (files only, never native init) ──────────────────
|
||||||
|
# The cold-start cost #2001 left behind: the ~4-minute model download began on the
|
||||||
|
# FIRST REQUEST, so every request in that window went silently uncompressed.
|
||||||
|
# Prefetching FILES at startup is safe because it is plain huggingface_hub HTTP —
|
||||||
|
# it never constructs an InferenceSession, which is the boundary that segfaults in
|
||||||
|
# libarrow/jemalloc on RHEL/CentOS 7-family hosts (#1908).
|
||||||
|
|
||||||
|
|
||||||
|
def test_eager_load_starts_artifact_prefetch_without_native_preload(monkeypatch):
|
||||||
|
router = _router_kompress_only()
|
||||||
|
stub = _StubCompressor(cached=False)
|
||||||
|
monkeypatch.setattr(router, "_get_kompress", lambda: stub)
|
||||||
|
prefetch_calls: list[object] = []
|
||||||
|
monkeypatch.setattr(
|
||||||
|
router,
|
||||||
|
"_prefetch_kompress_artifacts_async",
|
||||||
|
lambda cfg: (prefetch_calls.append(cfg), True)[1],
|
||||||
|
)
|
||||||
|
|
||||||
|
status = router.eager_load_compressors()
|
||||||
|
|
||||||
|
assert status["kompress_artifacts"] == "prefetching"
|
||||||
|
# The #2001 invariant still holds: no native preload on the startup path.
|
||||||
|
assert status["kompress"] == "deferred"
|
||||||
|
assert stub.preload_calls == []
|
||||||
|
assert len(prefetch_calls) == 1
|
||||||
|
|
||||||
|
|
||||||
|
def test_prefetch_never_constructs_a_session_or_imports_transformers(monkeypatch):
|
||||||
|
"""The safety property that makes startup prefetch legal at all."""
|
||||||
|
monkeypatch.setattr(kc, "_kompress_cache", {})
|
||||||
|
requested: list[str] = []
|
||||||
|
|
||||||
|
def fake_local_first(repo_id, filename, *, allow_network=True):
|
||||||
|
requested.append(filename)
|
||||||
|
return f"/cache/{filename}"
|
||||||
|
|
||||||
|
monkeypatch.setattr(kc, "hf_hub_download_local_first", fake_local_first)
|
||||||
|
|
||||||
|
def explode(*args, **kwargs):
|
||||||
|
raise AssertionError("prefetch must not build the model")
|
||||||
|
|
||||||
|
monkeypatch.setattr(kc, "_load_kompress", explode)
|
||||||
|
monkeypatch.setattr(kc, "_load_kompress_onnx", explode)
|
||||||
|
|
||||||
|
assert kc.prefetch_kompress_artifacts("org/model") is True
|
||||||
|
# Stops at the first candidate that resolves — the loader tries the same order.
|
||||||
|
assert requested == [kc._onnx_filename_candidates()[0]]
|
||||||
|
|
||||||
|
|
||||||
|
def test_prefetch_reports_false_when_no_artifact_resolves(monkeypatch):
|
||||||
|
monkeypatch.setattr(kc, "_kompress_cache", {})
|
||||||
|
|
||||||
|
def always_missing(repo_id, filename, *, allow_network=True):
|
||||||
|
raise OSError("not found")
|
||||||
|
|
||||||
|
monkeypatch.setattr(kc, "hf_hub_download_local_first", always_missing)
|
||||||
|
|
||||||
|
assert kc.prefetch_kompress_artifacts("org/model") is False
|
||||||
|
|
||||||
|
|
||||||
|
def test_background_prefetch_is_noop_when_model_already_cached(monkeypatch):
|
||||||
|
monkeypatch.setattr(kc, "_kompress_cache", {"org/model": object()})
|
||||||
|
monkeypatch.setattr(kc, "is_kompress_available", lambda: True)
|
||||||
|
|
||||||
|
assert kc.ensure_background_prefetch("org/model") is False
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_proxy_startup_does_not_enter_cached_kompress_native_loader(monkeypatch):
|
async def test_proxy_startup_does_not_enter_cached_kompress_native_loader(monkeypatch):
|
||||||
pytest.importorskip("httpx")
|
pytest.importorskip("httpx")
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue