From 841663da16971b1e0d8e204fdf18e4bafedaf9e0 Mon Sep 17 00:00:00 2001 From: gglucass Date: Thu, 11 Jun 2026 19:53:03 +0200 Subject: [PATCH] fix(proxy): make Kompress eager preload cache-only so a cold cache can't block startup (#783) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Description `ContentRouter.eager_load_compressors()` runs a network `hf_hub_download` of the Kompress ONNX model on the **blocking startup/lifespan path**, before the proxy binds its port. On a cold cache this is unsafe: - the download can hang long enough to blow the supervisor's bind timeout, or - a native crash in the download/ML stack (an **uncatchable `Fatal Python error: Aborted` / SIGABRT**) kills the interpreter before it ever `listen()`s. Either way the supervisor sees "proxy never opened its port" and gives up. We observed this in the field from the desktop app (process aborted during `eager_load_compressors -> _load_kompress_onnx -> hf_hub_download` of `onnx/kompress-int8.onnx`, while the only Python thread was parked in the HuggingFace download file-lock; the abort came from a native thread, so `try/except` at the call site cannot catch it). The eager preload is a latency optimization and must never be able to block — or kill — startup. This change makes startup preload **cache-only**: if the model isn't already cached, we defer the download to first use (off the startup path) and bind the port normally. Warm starts are unchanged. ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) - [ ] New feature (non-breaking change that adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to change) - [ ] Documentation update - [ ] Performance improvement - [ ] Code refactoring (no functional changes) ## Changes Made - `onnx_runtime.hf_hub_download_local_first(...)`: added `allow_network` (default `True`). When `False`, a cache miss re-raises the local-lookup error instead of falling back to a network download. - `kompress_compressor`: added `allow_download` (default `True`) threaded through `preload()` -> `_load_kompress()` -> `_load_kompress_onnx()` / `_load_kompress_pytorch()` and the ModernBERT tokenizer load. Added `KompressModelNotCached`, raised when a cache-only load misses. Auto-mode no longer falls back to a PyTorch network download on a cache-only miss — it propagates so the caller can defer. - `content_router.eager_load_compressors()`: calls `preload(allow_download=False)`. On `KompressModelNotCached` it logs and reports the component as `"deferred"` (a status `warmup.merge_transform_status` already handles gracefully) instead of letting a cold download run on the startup path. Default (first-request) loading behavior and warm-start preload are unchanged. ## Testing - [x] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check .`) - [x] Type checking passes (`mypy headroom`) - [x] New tests added for new functionality - [ ] Manual testing performed New tests in `tests/test_kompress_preload_deferral.py` cover: cache-only `hf_hub_download_local_first` never hits the network; default still falls back; cache-only ONNX load raises `KompressModelNotCached`; auto-mode does **not** trigger a PyTorch download on a cache-only miss; and `eager_load_compressors` reports `deferred` (cold) / `enabled` (warm). Existing `_load_kompress` dispatch tests updated for the new keyword-only param. > Note on environment: I do not have a clean reproduction of the native SIGABRT itself (it depends on a specific machine's HF download/ML native stack), so the "Manual testing performed" box is left unchecked. The tests target the structural fix — that startup preload can no longer perform a network download — which is the precondition for the crash. ## Test Output ``` $ uv run pytest -v tests/test_kompress_preload_deferral.py tests/test_kompress_preload_deferral.py::test_local_first_no_network_when_disallowed PASSED tests/test_kompress_preload_deferral.py::test_local_first_falls_back_to_network_by_default PASSED tests/test_kompress_preload_deferral.py::test_load_kompress_onnx_cache_miss_raises_not_cached PASSED tests/test_kompress_preload_deferral.py::test_load_kompress_auto_does_not_pytorch_download_on_cache_miss PASSED tests/test_kompress_preload_deferral.py::test_eager_load_defers_when_model_not_cached PASSED tests/test_kompress_preload_deferral.py::test_eager_load_enabled_when_model_cached PASSED 6 passed in 4.82s $ uv run pytest tests/test_transforms/test_kompress_compressor.py tests/test_transforms_content_router.py tests/test_onnx_runtime.py tests/test_proxy_warmup.py 63 passed $ uv run ruff check # All checks passed! $ uv run mypy headroom/onnx_runtime.py headroom/transforms/kompress_compressor.py headroom/transforms/content_router.py Success: no issues found ``` ## 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 - [ ] I have updated the CHANGELOG.md if applicable (auto-generated from conventional commits) ## Additional Notes This contains the cold-start case. A native crash in onnxruntime *session init* (as opposed to the download) on first request would still be a separate issue; it is not what was observed here (the abort was during the HF download), and isolating it would be a larger, separate change. --------- Co-authored-by: Claude Opus 4.8 --- headroom/onnx_runtime.py | 12 +- headroom/transforms/content_router.py | 32 +++- headroom/transforms/kompress_compressor.py | 142 ++++++++++++++--- tests/test_kompress_preload_deferral.py | 143 ++++++++++++++++++ .../test_kompress_compressor.py | 10 +- 5 files changed, 305 insertions(+), 34 deletions(-) create mode 100644 tests/test_kompress_preload_deferral.py diff --git a/headroom/onnx_runtime.py b/headroom/onnx_runtime.py index 214f7efcd..4609528c1 100644 --- a/headroom/onnx_runtime.py +++ b/headroom/onnx_runtime.py @@ -7,7 +7,7 @@ import sys from typing import Any -def hf_hub_download_local_first(repo_id: str, filename: str) -> str: +def hf_hub_download_local_first(repo_id: str, filename: str, *, allow_network: bool = True) -> str: """Download a file from HuggingFace Hub, preferring the local cache. Tries ``local_files_only=True`` first to avoid a network HEAD request when @@ -17,12 +17,18 @@ def hf_hub_download_local_first(repo_id: str, filename: str) -> str: Args: repo_id: HuggingFace Hub repository identifier (e.g. ``"org/model"``). filename: Filename within the repository. + allow_network: When ``False``, never fall back to a network download — + a cache miss re-raises the local-lookup error. Used by startup + preload so a cold cache cannot block (or, via native crashes in the + download stack, kill) the process before it binds its port. Returns: Absolute path to the local cached file. Raises: - Any exception raised by ``hf_hub_download`` on a genuine download failure. + Any exception raised by ``hf_hub_download`` on a genuine download failure, + or the local-lookup error when ``allow_network`` is ``False`` and the + file is not cached. """ from huggingface_hub import hf_hub_download from huggingface_hub.errors import EntryNotFoundError, LocalEntryNotFoundError @@ -30,6 +36,8 @@ def hf_hub_download_local_first(repo_id: str, filename: str) -> str: try: return str(hf_hub_download(repo_id, filename, local_files_only=True)) except (LocalEntryNotFoundError, EntryNotFoundError, OSError): + if not allow_network: + raise return str(hf_hub_download(repo_id, filename)) diff --git a/headroom/transforms/content_router.py b/headroom/transforms/content_router.py index 6960d7692..18e6f612c 100644 --- a/headroom/transforms/content_router.py +++ b/headroom/transforms/content_router.py @@ -1622,14 +1622,36 @@ class ContentRouter(Transform): """ status: dict[str, str] = {} - # 1. ML text compressor: Kompress + # 1. ML text compressor: Kompress. + # + # Eager preload is cache-only (allow_download=False): on a cold cache we + # must NOT trigger a network download here, because this runs on the + # blocking startup/lifespan path before the proxy binds its port. A slow + # download stalls the bind, and a hard crash in the native download/ML + # stack (uncatchable SIGABRT) kills the interpreter before it ever + # listens — the proxy then "never opens its port" and the supervisor + # gives up. When the model isn't cached we defer to first use instead. if self.config.enable_kompress: + from .kompress_compressor import KompressModelNotCached + compressor = self._get_kompress() if compressor: - backend = compressor.preload() if hasattr(compressor, "preload") else "unknown" - logger.info("Kompress model pre-loaded at startup backend=%s", backend) - status["kompress"] = "enabled" - status["kompress_backend"] = str(backend) + if not hasattr(compressor, "preload"): + status["kompress"] = "enabled" + status["kompress_backend"] = "unknown" + else: + try: + backend = compressor.preload(allow_download=False) + except KompressModelNotCached: + logger.info( + "Kompress model not cached; deferring download to " + "first use to keep startup non-blocking" + ) + status["kompress"] = "deferred" + else: + logger.info("Kompress model pre-loaded at startup backend=%s", backend) + status["kompress"] = "enabled" + status["kompress_backend"] = str(backend) else: status["kompress"] = "unavailable" diff --git a/headroom/transforms/kompress_compressor.py b/headroom/transforms/kompress_compressor.py index 96fe98366..c888b2ae3 100644 --- a/headroom/transforms/kompress_compressor.py +++ b/headroom/transforms/kompress_compressor.py @@ -66,6 +66,29 @@ KOMPRESS_BATCH_SIZE_ENV = "HEADROOM_KOMPRESS_BATCH_SIZE" KompressBackend = Literal["auto", "onnx", "onnx_cpu", "onnx_coreml", "pytorch", "pytorch_mps"] +# HuggingFace local-lookup errors that mean "asset not in cache" rather than a +# genuine failure. Caught when loading cache-only so startup can defer instead. +try: + from huggingface_hub.errors import EntryNotFoundError, LocalEntryNotFoundError + + _NOT_CACHED_ERRORS: tuple[type[BaseException], ...] = ( + LocalEntryNotFoundError, + EntryNotFoundError, + OSError, + ) +except Exception: # pragma: no cover - huggingface_hub always present with [ml] + _NOT_CACHED_ERRORS = (OSError,) + + +class KompressModelNotCached(RuntimeError): + """Raised when a cache-only load is requested but the model is not cached. + + Used by startup eager-preload (``allow_download=False``) so the caller can + defer the download to first use instead of blocking the proxy startup path + on a network fetch. + """ + + # Model cache: model_id -> (model, tokenizer, backend) # Supports multiple models loaded simultaneously. _kompress_cache: dict[str, tuple[Any, Any, str]] = {} @@ -358,22 +381,38 @@ def _onnx_filename_candidates() -> tuple[str, ...]: return _DEFAULT_ONNX_FILENAMES -def _create_onnx_session(model_id: str, ort: Any, providers: list[Any]) -> Any: +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). + + When ``allow_download`` is ``False`` candidates are resolved from the local + cache only; if none is cached, :class:`KompressModelNotCached` is raised + instead of hitting the network. ``onnxruntime`` is imported only after a + candidate resolves, so a cache-only miss never requires it. """ last_err: Exception | None = None + cache_miss = False + ort: Any = None for filename in _onnx_filename_candidates(): try: - onnx_path = hf_hub_download_local_first(model_id, filename) + onnx_path = hf_hub_download_local_first( + model_id, filename, allow_network=allow_download + ) except Exception as exc: last_err = exc - logger.debug("ONNX artifact %r not in %s: %s", filename, model_id, exc) + cache_miss = cache_miss or isinstance(exc, _NOT_CACHED_ERRORS) + logger.debug("ONNX artifact %r unavailable for %s: %s", filename, model_id, exc) continue + if ort is None: + import onnxruntime + + ort = onnxruntime try: return ort.InferenceSession( onnx_path, @@ -388,6 +427,8 @@ def _create_onnx_session(model_id: str, ort: Any, providers: list[Any]) -> Any: model_id, exc, ) + if not allow_download and cache_miss: + raise KompressModelNotCached(model_id) from last_err raise FileNotFoundError( f"No loadable ONNX artifact in {model_id}; tried {_onnx_filename_candidates()}" ) from last_err @@ -397,11 +438,14 @@ def _load_kompress_onnx( model_id: str, *, use_coreml: bool = False, + allow_download: bool = True, ) -> tuple[Any, Any, str]: - """Download the ONNX model from HuggingFace and load with onnxruntime.""" - import onnxruntime as ort - from transformers import AutoTokenizer + """Download ONNX INT8 model from HuggingFace and load with onnxruntime. + When ``allow_download`` is ``False`` the model and tokenizer are loaded from + the local cache only; a cache miss raises :class:`KompressModelNotCached` + instead of hitting the network. + """ with _kompress_lock: if model_id in _kompress_cache: return _kompress_cache[model_id] @@ -435,17 +479,38 @@ def _load_kompress_onnx( else: providers = ["CPUExecutionProvider"] - session = _create_onnx_session(model_id, ort, providers) + session = _create_onnx_session(model_id, providers, allow_download=allow_download) model = _OnnxModel(session) - tokenizer = AutoTokenizer.from_pretrained("answerdotai/ModernBERT-base") + + from transformers import AutoTokenizer + + tokenizer = _load_modernbert_tokenizer(AutoTokenizer, allow_download=allow_download) _kompress_cache[model_id] = (model, tokenizer, backend) logger.info("Kompress ONNX loaded: %s backend=%s", model_id, backend) return model, tokenizer, backend -def _load_kompress_pytorch(model_id: str, device: str = "auto") -> tuple[Any, Any, str]: - """Download PyTorch model from HuggingFace and load with torch.""" +def _load_modernbert_tokenizer(auto_tokenizer: Any, *, allow_download: bool) -> Any: + """Load the ModernBERT tokenizer, cache-only when ``allow_download`` is False.""" + try: + return auto_tokenizer.from_pretrained( + "answerdotai/ModernBERT-base", local_files_only=not allow_download + ) + except _NOT_CACHED_ERRORS as exc: + if not allow_download: + raise KompressModelNotCached("answerdotai/ModernBERT-base") from exc + raise + + +def _load_kompress_pytorch( + model_id: str, device: str = "auto", *, allow_download: bool = True +) -> tuple[Any, Any, str]: + """Download PyTorch model from HuggingFace and load with torch. + + When ``allow_download`` is ``False`` weights and tokenizer are loaded from + the local cache only; a cache miss raises :class:`KompressModelNotCached`. + """ import torch from transformers import AutoTokenizer @@ -455,7 +520,14 @@ def _load_kompress_pytorch(model_id: str, device: str = "auto") -> tuple[Any, An logger.info("Downloading Kompress PyTorch model from %s ...", model_id) - weights_path = hf_hub_download_local_first(model_id, "model.safetensors") + try: + weights_path = hf_hub_download_local_first( + model_id, "model.safetensors", allow_network=allow_download + ) + except _NOT_CACHED_ERRORS as exc: + if not allow_download: + raise KompressModelNotCached(model_id) from exc + raise HeadroomCompressorModel = _get_model_class() model = HeadroomCompressorModel() @@ -476,7 +548,7 @@ def _load_kompress_pytorch(model_id: str, device: str = "auto") -> tuple[Any, An model.to(device) model.eval() - tokenizer = AutoTokenizer.from_pretrained("answerdotai/ModernBERT-base") + tokenizer = _load_modernbert_tokenizer(AutoTokenizer, allow_download=allow_download) _validate_pytorch_device(model, tokenizer, device) _kompress_cache[model_id] = (model, tokenizer, "pytorch") @@ -503,7 +575,9 @@ def _validate_pytorch_device(model: Any, tokenizer: Any, device: str) -> None: _ = scores[0].detach().cpu() -def _load_kompress(model_id: str = HF_MODEL_ID, device: str = "auto") -> tuple[Any, Any, str]: +def _load_kompress( + model_id: str = HF_MODEL_ID, device: str = "auto", *, allow_download: bool = True +) -> tuple[Any, Any, str]: """Load Kompress model, returns (model, tokenizer, backend). The default keeps the historic behavior: try ONNX CPU first @@ -516,6 +590,10 @@ def _load_kompress(model_id: str = HF_MODEL_ID, device: str = "auto") -> tuple[A - pytorch: force PyTorch with the configured device. - pytorch_mps: force PyTorch on Apple's MPS backend. + When ``allow_download`` is ``False`` the model is loaded from the local + cache only and a cache miss raises :class:`KompressModelNotCached` rather + than fetching from the network. + Models are cached by model_id — multiple models can coexist. """ if model_id in _kompress_cache: @@ -523,15 +601,17 @@ def _load_kompress(model_id: str = HF_MODEL_ID, device: str = "auto") -> tuple[A backend = _selected_backend() if backend in ("onnx", "onnx_cpu"): - return _load_kompress_onnx(model_id, use_coreml=False) + return _load_kompress_onnx(model_id, use_coreml=False, allow_download=allow_download) if backend == "onnx_coreml": - return _load_kompress_onnx(model_id, use_coreml=True) + return _load_kompress_onnx(model_id, use_coreml=True, allow_download=allow_download) if backend in ("pytorch", "pytorch_mps"): forced_device = "mps" if backend == "pytorch_mps" else device try: - return _load_kompress_pytorch(model_id, forced_device) + return _load_kompress_pytorch(model_id, forced_device, allow_download=allow_download) + except KompressModelNotCached: + raise except Exception as exc: if backend != "pytorch_mps": raise @@ -541,20 +621,27 @@ def _load_kompress(model_id: str = HF_MODEL_ID, device: str = "auto") -> tuple[A exc, ) if _is_onnx_available(): - return _load_kompress_onnx(model_id, use_coreml=False) - return _load_kompress_pytorch(model_id, "cpu") + return _load_kompress_onnx( + model_id, use_coreml=False, allow_download=allow_download + ) + return _load_kompress_pytorch(model_id, "cpu", allow_download=allow_download) # Auto mode: preserve stable default behavior. This avoids changing # compression quality/perf characteristics for existing installs while # allowing opt-in MPS/CoreML experiments via HEADROOM_KOMPRESS_BACKEND. if _is_onnx_available(): try: - return _load_kompress_onnx(model_id, use_coreml=False) + return _load_kompress_onnx(model_id, use_coreml=False, allow_download=allow_download) + except KompressModelNotCached: + # Cache-only miss: don't trigger a PyTorch network download as a + # fallback — propagate so the caller can defer. + if not allow_download: + raise except Exception as e: logger.warning("ONNX load failed for %s, trying PyTorch: %s", model_id, e) if _is_pytorch_available(): - return _load_kompress_pytorch(model_id, device) + return _load_kompress_pytorch(model_id, device, allow_download=allow_download) raise ImportError( "Kompress requires onnxruntime or torch. Install with: pip install headroom-ai[proxy]" @@ -653,10 +740,19 @@ class KompressCompressor(Transform): def __init__(self, config: KompressConfig | None = None): self.config = config or KompressConfig() - def preload(self) -> str: - """Load the backing model/tokenizer and return the selected backend.""" + def preload(self, *, allow_download: bool = True) -> str: + """Load the backing model/tokenizer and return the selected backend. - _model, _tokenizer, backend = _load_kompress(self.config.model_id, self.config.device) + When ``allow_download`` is ``False`` the model is loaded from the local + cache only; if it is not cached, :class:`KompressModelNotCached` is + raised so the caller can defer the download to first use. Startup eager + preload uses this so a cold cache cannot block the proxy from binding + its port. + """ + + _model, _tokenizer, backend = _load_kompress( + self.config.model_id, self.config.device, allow_download=allow_download + ) return backend def compress( diff --git a/tests/test_kompress_preload_deferral.py b/tests/test_kompress_preload_deferral.py new file mode 100644 index 000000000..ff32e6df4 --- /dev/null +++ b/tests/test_kompress_preload_deferral.py @@ -0,0 +1,143 @@ +"""Startup eager-preload must be cache-only so a cold cache cannot block or +crash the proxy before it binds its port. + +Regression for the production crash where ``eager_load_compressors`` ran a +network ``hf_hub_download`` of the Kompress ONNX model on the blocking +startup/lifespan path. On a cold cache that download could hang (300s bind +timeout) or hit a native ``SIGABRT`` in the download/ML stack, killing the +interpreter before it ever listened on its port. +""" + +from __future__ import annotations + +import pytest + +from headroom import onnx_runtime +from headroom.transforms import kompress_compressor as kc +from headroom.transforms.content_router import ContentRouter, ContentRouterConfig +from headroom.transforms.kompress_compressor import KompressModelNotCached + + +def test_local_first_no_network_when_disallowed(monkeypatch): + """allow_network=False must never fall back to a network download.""" + import huggingface_hub + from huggingface_hub.errors import LocalEntryNotFoundError + + calls: list[bool] = [] + + def fake_download(repo_id, filename, **kwargs): + local_only = kwargs.get("local_files_only", False) + calls.append(local_only) + if local_only: + raise LocalEntryNotFoundError("not cached") + return "/cache/networked" + + monkeypatch.setattr(huggingface_hub, "hf_hub_download", fake_download) + + with pytest.raises(LocalEntryNotFoundError): + onnx_runtime.hf_hub_download_local_first("org/model", "f.onnx", allow_network=False) + + # Only the local-only lookup ran; the network branch was never taken. + assert calls == [True] + + +def test_local_first_falls_back_to_network_by_default(monkeypatch): + """allow_network=True (default) keeps the historic cold-start behavior.""" + import huggingface_hub + from huggingface_hub.errors import LocalEntryNotFoundError + + calls: list[bool] = [] + + def fake_download(repo_id, filename, **kwargs): + local_only = kwargs.get("local_files_only", False) + calls.append(local_only) + if local_only: + raise LocalEntryNotFoundError("not cached") + return "/cache/networked" + + monkeypatch.setattr(huggingface_hub, "hf_hub_download", fake_download) + + path = onnx_runtime.hf_hub_download_local_first("org/model", "f.onnx") + assert path == "/cache/networked" + assert calls == [True, False] # local-only miss, then network download + + +def test_load_kompress_onnx_cache_miss_raises_not_cached(monkeypatch): + """A cache-only ONNX load surfaces KompressModelNotCached, not a network call.""" + from huggingface_hub.errors import LocalEntryNotFoundError + + monkeypatch.setattr(kc, "_kompress_cache", {}) + + def fake_local_first(repo_id, filename, *, allow_network=True): + assert allow_network is False # eager preload must request cache-only + raise LocalEntryNotFoundError("not cached") + + monkeypatch.setattr(kc, "hf_hub_download_local_first", fake_local_first) + + with pytest.raises(KompressModelNotCached): + kc._load_kompress_onnx("org/model", allow_download=False) + + +def test_load_kompress_auto_does_not_pytorch_download_on_cache_miss(monkeypatch): + """Auto mode must propagate the cache miss, not fall back to a PyTorch fetch.""" + monkeypatch.setattr(kc, "_kompress_cache", {}) + monkeypatch.setattr(kc, "_selected_backend", lambda: "auto") + monkeypatch.setattr(kc, "_is_onnx_available", lambda: True) + monkeypatch.setattr(kc, "_is_pytorch_available", lambda: True) + + def onnx_not_cached(model_id, *, use_coreml=False, allow_download=True): + raise KompressModelNotCached(model_id) + + def pytorch_should_not_run(*args, **kwargs): + raise AssertionError("PyTorch fallback must not download on a cache-only miss") + + monkeypatch.setattr(kc, "_load_kompress_onnx", onnx_not_cached) + monkeypatch.setattr(kc, "_load_kompress_pytorch", pytorch_should_not_run) + + with pytest.raises(KompressModelNotCached): + kc._load_kompress("org/model", allow_download=False) + + +class _StubCompressor: + def __init__(self, *, cached: bool): + self._cached = cached + self.preload_calls: list[bool] = [] + + def preload(self, *, allow_download: bool = True) -> str: + self.preload_calls.append(allow_download) + if self._cached: + return "onnx" + raise KompressModelNotCached("org/model") + + +def _router_kompress_only() -> ContentRouter: + return ContentRouter( + ContentRouterConfig( + enable_kompress=True, + enable_code_aware=False, + enable_smart_crusher=False, + ) + ) + + +def test_eager_load_defers_when_model_not_cached(monkeypatch): + router = _router_kompress_only() + stub = _StubCompressor(cached=False) + monkeypatch.setattr(router, "_get_kompress", lambda: stub) + + status = router.eager_load_compressors() + + assert status["kompress"] == "deferred" + assert stub.preload_calls == [False] # cache-only preload at startup + + +def test_eager_load_enabled_when_model_cached(monkeypatch): + router = _router_kompress_only() + stub = _StubCompressor(cached=True) + monkeypatch.setattr(router, "_get_kompress", lambda: stub) + + status = router.eager_load_compressors() + + assert status["kompress"] == "enabled" + assert status["kompress_backend"] == "onnx" + assert stub.preload_calls == [False] diff --git a/tests/test_transforms/test_kompress_compressor.py b/tests/test_transforms/test_kompress_compressor.py index f5847f36a..826a3014b 100644 --- a/tests/test_transforms/test_kompress_compressor.py +++ b/tests/test_transforms/test_kompress_compressor.py @@ -117,7 +117,7 @@ class TestKompressBackendSelection: monkeypatch.setattr( kmod, "_load_kompress_pytorch", - lambda model_id, device: ( + lambda model_id, device, *, allow_download=True: ( calls.append((model_id, device)) or ("model", "tokenizer", "pytorch") ), ) @@ -134,7 +134,7 @@ class TestKompressBackendSelection: monkeypatch.setattr( kmod, "_load_kompress_onnx", - lambda model_id, *, use_coreml=False: ( + lambda model_id, *, use_coreml=False, allow_download=True: ( calls.append((model_id, use_coreml)) or ("model", "tokenizer", "onnx_coreml") ), ) @@ -153,14 +153,16 @@ class TestKompressBackendSelection: monkeypatch.setattr( kmod, "_load_kompress_onnx", - lambda model_id, *, use_coreml=False: ( + lambda model_id, *, use_coreml=False, allow_download=True: ( calls.append("onnx") or ("model", "tokenizer", "onnx") ), ) monkeypatch.setattr( kmod, "_load_kompress_pytorch", - lambda model_id, device: calls.append("pytorch") or ("model", "tokenizer", "pytorch"), + lambda model_id, device, *, allow_download=True: ( + calls.append("pytorch") or ("model", "tokenizer", "pytorch") + ), ) assert kmod._load_kompress("model-c") == ("model", "tokenizer", "onnx")