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:
Tejas Chopra 2026-08-03 10:42:48 -07:00 committed by GitHub
parent 8262a4a321
commit 224578e80b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 393 additions and 11 deletions

View file

@ -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"

View file

@ -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],

View file

@ -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

View file

@ -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")