mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-10 14:27:00 -04:00
fix(proxy): make Kompress eager preload cache-only so a cold cache can't block startup (#783)
## 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 <changed files> # 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 <noreply@anthropic.com>
This commit is contained in:
gglucass
GitHub
parent
9bff5752bb
commit
841663da16
5 changed files with 305 additions and 34 deletions
|
|
@ -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))
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
143
tests/test_kompress_preload_deferral.py
Normal file
143
tests/test_kompress_preload_deferral.py
Normal file
|
|
@ -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]
|
||||
|
|
@ -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")
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue