diff --git a/headroom/transforms/content_router.py b/headroom/transforms/content_router.py index 6196cf8fb..a47773f7d 100644 --- a/headroom/transforms/content_router.py +++ b/headroom/transforms/content_router.py @@ -4047,6 +4047,22 @@ class ContentRouter(Transform): logger.debug("HTMLExtractor not available (install trafilatura)") 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]: """Pre-load compressors at startup to avoid first-request latency. @@ -4061,7 +4077,18 @@ class ContentRouter(Transform): # 1. ML text compressor: Kompress. # # 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: compressor = self._get_kompress() if compressor: @@ -4069,8 +4096,10 @@ class ContentRouter(Transform): status["kompress"] = "enabled" status["kompress_backend"] = "unknown" else: - logger.info("Kompress model preload deferred until first request") 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: status["kompress"] = "unavailable" diff --git a/headroom/transforms/kompress_compressor.py b/headroom/transforms/kompress_compressor.py index c2a39870b..37c4260cf 100644 --- a/headroom/transforms/kompress_compressor.py +++ b/headroom/transforms/kompress_compressor.py @@ -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_COREML_CACHE_DIR_ENV = "HEADROOM_KOMPRESS_COREML_CACHE_DIR" 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_DEFAULT = 3000 KOMPRESS_BATCH_SIZE_ENV = "HEADROOM_KOMPRESS_BATCH_SIZE" @@ -600,15 +606,42 @@ def _onnx_filename_candidates() -> tuple[str, ...]: 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( model_id: str, providers: list[Any], *, allow_download: bool = True ) -> Any: """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 - session-load failure (e.g. the weight-only int8 artifact uses the - MatMulNBits contrib op, which old onnxruntime builds can't run — those - installs fall through to the fp32 artifact instead of losing Kompress). + A candidate is skipped on download miss (file not in the repo), on + session-load failure, or on smoke-run failure (e.g. the weight-only int8 + artifact uses the MatMulNBits contrib op, which some onnxruntime builds + 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 cache only; if none is cached, :class:`KompressModelNotCached` is raised @@ -633,15 +666,17 @@ def _create_onnx_session( ort = onnxruntime try: - return ort.InferenceSession( + session = ort.InferenceSession( onnx_path, _onnx_session_options(ort), providers=providers, ) + _smoke_run(session) + return session except Exception as exc: last_err = exc 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, model_id, exc, @@ -1041,6 +1076,70 @@ def ensure_background_download(model_id: str = HF_MODEL_ID, device: str = "auto" 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( model_id: str = HF_MODEL_ID, device: str = "cpu", @@ -1170,10 +1269,14 @@ class KompressCompressor(Transform): def __init__(self, config: KompressConfig | None = None): self.config = config or KompressConfig() - # Set by the preload canary when inference is too slow to be useful; - # compress()/compress_batch() then pass content through untouched. + # Set by the preload canary when inference is too slow to be useful, or by + # the failure latch when inference raises repeatedly; compress()/ + # compress_batch() then pass content through untouched. self._degraded_reason: str | 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: """Load the backing model/tokenizer and return the selected backend. @@ -1563,6 +1666,9 @@ class KompressCompressor(Transform): 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 except KompressModelNotCached: @@ -1572,9 +1678,41 @@ class KompressCompressor(Transform): ) return self._passthrough(content, n_words) except Exception as e: - logger.warning("Kompress compression failed: %s", e) + self._record_inference_failure(e) 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( self, contents: list[str], diff --git a/tests/test_kompress_failsafe.py b/tests/test_kompress_failsafe.py index 138ba9d29..3fb305c93 100644 --- a/tests/test_kompress_failsafe.py +++ b/tests/test_kompress_failsafe.py @@ -535,3 +535,146 @@ def test_canary_probe_error_never_breaks_preload(monkeypatch): 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 diff --git a/tests/test_kompress_preload_deferral.py b/tests/test_kompress_preload_deferral.py index 143c61153..2dff7a99d 100644 --- a/tests/test_kompress_preload_deferral.py +++ b/tests/test_kompress_preload_deferral.py @@ -127,6 +127,9 @@ def test_eager_load_defers_kompress_regardless_of_cache_state(monkeypatch, cache router = _router_kompress_only() stub = _StubCompressor(cached=cache_state == "cached") 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() @@ -164,6 +167,7 @@ def test_non_kompress_warmups_continue_when_kompress_is_deferred(monkeypatch): router = _router_kompress_only() stub = _StubCompressor(cached=True) 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._get_magika", lambda: object()) @@ -174,6 +178,74 @@ def test_non_kompress_warmups_continue_when_kompress_is_deferred(monkeypatch): 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 async def test_proxy_startup_does_not_enter_cached_kompress_native_loader(monkeypatch): pytest.importorskip("httpx")