vinez/headroom
1
0
Fork 0
mirror of https://github.com/headroomlabs-ai/headroom.git synced 2026-08-10 14:27:00 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix(kompress): never block the request path on the cold-cache model download (#1161)

Browse source 
Closes #1146.

## Problem

On a cold cache, the first request that reaches the Kompress deep
compressor triggers an inline `hf_hub_download` of the 274 MB
`chopratejas/kompress-v2-base` ONNX model **on the request thread**.
That download races the proxy's compression budget
(`HEADROOM_COMPRESSION_TIMEOUT_SECONDS`, default 30s — the
`compression_first_stage` timeout): the fetch is cancelled mid-transfer,
**nothing finalizes in the HF cache**, and the request fails open
(uncompressed). Because the partial blob never lands, every subsequent
request repeats the same ~30s hang + fail-open, so the deep compressor
never actually becomes available through the proxy.

This is a **distinct root cause from #946** (which concerns the timeout
itself). Here the model must simply never be fetched synchronously on a
latency-sensitive request.

## Fix

Make the request path cache-only and move the one-time download
off-thread.

**`kompress_compressor.py`**
- `compress(..., allow_download=False)` — new keyword (default `True`,
so the direct API and `compress_batch` are unchanged) that resolves the
model cache-only; on a cold cache it raises `KompressModelNotCached` and
passes through instead of blocking on the network.
- `is_ready()` — lockless cache-membership check, safe to call on the
hot path.
- `ensure_background_download(model_id, device)` — starts at most one
daemon thread per model to pull the artifact down out of band (a
finished/failed thread is replaced, so a transient failure can be
retried by a later request). The compression timeout does not bound this
thread.

**`content_router.py`** — gate the deep path on readiness:
- not ready → return passthrough immediately and kick off the background
download;
- ready → `compress(allow_download=False)` (cache-only, no network on
the request thread).

Net effect: the cold-cache deep path returns in ~0 ms (passthrough)
instead of hanging ~30 s; the model downloads once in the background;
subsequent requests transparently use the deep compressor once it is
cached.

## Verification

Clean install of `headroom-ai==0.26.0` (main `@9f7f3ad` + this patch),
HF cache empty:

- **Cold-cache router request**: returns in **56 ms**, passthrough
(output == input), fetch handed to a background daemon thread — vs. the
pre-fix inline hang.
- **A/B on the same `compress()` call** with a slow-fetch stand-in:
pre-fix (`allow_download=True`) blocked **24.00 s** on the request
thread; post-fix (`allow_download=False`) returned **59 ms**.
- **4 new regression tests** in
`tests/test_kompress_request_nonblocking.py` (cache-only passthrough;
one-thread-per-model background download; router skips deep path when
not ready; router stays cache-only when ready) — `4 passed`.

## Files
- `headroom/transforms/kompress_compressor.py` — non-blocking load +
cache-only `compress`
- `headroom/transforms/content_router.py` — gate deep path on
`is_ready()`; background fetch when cold
- `tests/test_kompress_request_nonblocking.py` — regression coverage

---------

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
xspawnnn 2026-06-19 09:33:05 -07:00 • committed by GitHub
parent 5b84691770
commit 3fc2a78a5e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 243 additions and 13 deletions
Download patch file Download diff file Expand all files Collapse all files

31
headroom/transforms/content_router.py
View file

@ -1681,21 +1681,28 @@ class ContentRouter(Transform):
compressed: str | None = None
compressed_tokens: int | None = None
# Primary: Kompress — downloads from chopratejas/kompress-v2-base on first use
# Primary: Kompress. On a cold cache the model is fetched once in the
# background (ensure_background_load) instead of blocking this request
# thread on a 274MB download that races the compression timeout and
# fails open. Until it is cached, route around the deep path.
if self.config.enable_kompress:
compressor = self._get_kompress()
if compressor:
try:
result = compressor.compress(
text_to_compress,
context=context,
question=question,
target_ratio=getattr(self, "_runtime_target_ratio", None),
)
compressed = result.compressed
compressed_tokens = result.compressed_tokens
except Exception as e:
logger.warning("Kompress failed: %s", e)
if not compressor.is_ready():
compressor.ensure_background_load()
else:
try:
result = compressor.compress(
text_to_compress,
context=context,
question=question,
target_ratio=getattr(self, "_runtime_target_ratio", None),
allow_download=False,
)
compressed = result.compressed
compressed_tokens = result.compressed_tokens
except Exception as e:
logger.warning("Kompress failed: %s", e)
if compressed is None:
return content, len(content.split())

83
headroom/transforms/kompress_compressor.py
View file

@ -678,6 +678,57 @@ def unload_kompress_model(model_id: str | None = None) -> bool:
return True
# ── Background model download ─────────────────────────────────────────
#
# The proxy request path must never block on a cold model download. A first
# deep-path request would otherwise resolve the 274MB ONNX artifact via an
# inline hf_hub_download on the request thread, where it races the proxy's
# compression timeout (HEADROOM_COMPRESSION_TIMEOUT_SECONDS, default 30s). The
# fetch is cancelled mid-transfer, the blob never finalizes in the HF cache,
# and every subsequent request re-hangs and fails open. Instead the request
# path resolves the model cache-only (allow_download=False) and pulls it down
# once here, in a daemon thread that the compression timeout does not bound.
_download_threads: dict[str, threading.Thread] = {}
_download_threads_lock = threading.Lock()
def _background_download(model_id: str, device: str) -> None:
try:
logger.info("Kompress: downloading model %s in the background ...", model_id)
_load_kompress(model_id, device, allow_download=True)
logger.info("Kompress: background model download complete for %s", model_id)
except Exception as exc:
logger.warning("Kompress: background model download failed for %s: %s", model_id, exc)
def ensure_background_download(model_id: str = HF_MODEL_ID, device: str = "auto") -> None:
"""Start a one-shot background download of the model if it isn't cached.
Idempotent and non-blocking: at most one download thread runs per model_id,
and a finished or failed thread is replaced on the next call so a transient
network failure can be retried by a later request. Once the download
completes the deep path activates on subsequent requests without ever
blocking one on the network.
"""
if model_id in _kompress_cache:
return
with _download_threads_lock:
if model_id in _kompress_cache:
return
existing = _download_threads.get(model_id)
if existing is not None and existing.is_alive():
return
thread = threading.Thread(
target=_background_download,
args=(model_id, device),
name=f"kompress-download-{model_id.replace('/', '-')}",
daemon=True,
)
_download_threads[model_id] = thread
thread.start()
# ── Compressor ────────────────────────────────────────────────────────
@ -755,6 +806,21 @@ class KompressCompressor(Transform):
)
return backend
def is_ready(self) -> bool:
"""True if the model is loaded so :meth:`compress` won't touch the network.
A plain cache-membership check — no lock, no I/O — safe to call on the
hot request path to decide whether to run the deep compressor or skip it.
"""
return self.config.model_id in _kompress_cache
def ensure_background_load(self) -> None:
"""Kick off a one-shot, non-blocking background download of the model.
No-op when the model is already cached or a download is already running.
"""
ensure_background_download(self.config.model_id, self.config.device)
def compress(
self,
content: str,
@ -762,6 +828,8 @@ class KompressCompressor(Transform):
content_type: str | None = None,
question: str | None = None,
target_ratio: float | None = None,
*,
allow_download: bool = True,
) -> KompressResult:
"""Compress content using Kompress model.
@ -773,6 +841,11 @@ class KompressCompressor(Transform):
target_ratio: If None (default), model decides how much to keep using
score threshold. If set (e.g. 0.3), forces that keep ratio.
The proxy never sets this — only user-facing API does.
allow_download: When False, load the model from the local cache only;
a cache miss passes through instead of fetching from the network.
The proxy sets this False so a cold model never blocks the request
thread (see ``ensure_background_download``); direct callers keep
the historic auto-download-on-first-use behavior.
Returns:
KompressResult with compressed text.
@ -784,7 +857,9 @@ class KompressCompressor(Transform):
return self._passthrough(content, n_words)
try:
model, tokenizer, backend = _load_kompress(self.config.model_id, self.config.device)
model, tokenizer, backend = _load_kompress(
self.config.model_id, self.config.device, allow_download=allow_download
)
is_onnx = backend == "onnx"
device_type = _model_device_type(model, backend)
@ -919,6 +994,12 @@ class KompressCompressor(Transform):
return result
except KompressModelNotCached:
logger.debug(
"Kompress model %s not cached; passing through without compression",
self.config.model_id,
)
return self._passthrough(content, n_words)
except Exception as e:
logger.warning("Kompress compression failed: %s", e)
return self._passthrough(content, n_words)

136
tests/test_kompress_request_nonblocking.py Normal file
View file

@ -0,0 +1,136 @@
"""The proxy request path must never block on a cold Kompress model download.
Counterpart to ``test_kompress_preload_deferral.py`` (which covers the startup
path). A first deep-path request used to resolve the 274MB ONNX artifact via an
inline ``hf_hub_download`` on the request thread, where it raced the proxy's
``HEADROOM_COMPRESSION_TIMEOUT_SECONDS`` budget (GH #946 / #1146): the fetch was
cancelled mid-transfer, nothing cached, and every request re-hung and failed
open. The request path now resolves the model cache-only and pulls it down once
in a background daemon thread instead.
"""
from __future__ import annotations
from headroom.transforms import kompress_compressor as kc
from headroom.transforms.content_router import ContentRouter, ContentRouterConfig
from headroom.transforms.kompress_compressor import KompressCompressor
def test_compress_cache_only_passes_through_without_network(monkeypatch):
"""compress(allow_download=False) on a cold cache must not hit the network."""
from huggingface_hub.errors import LocalEntryNotFoundError
monkeypatch.setattr(kc, "_kompress_cache", {})
monkeypatch.setattr(kc, "_selected_backend", lambda: "onnx")
def fake_local_first(repo_id, filename, *, allow_network=True):
assert allow_network is False, "request path must resolve the model cache-only"
raise LocalEntryNotFoundError("not cached")
monkeypatch.setattr(kc, "hf_hub_download_local_first", fake_local_first)
text = " ".join(["token"] * 50) # >= 10 words: not the short-content passthrough
result = KompressCompressor().compress(text, allow_download=False)
assert result.compressed == text
assert result.compression_ratio == 1.0
def test_ensure_background_download_runs_one_thread_per_model(monkeypatch):
"""At most one download thread per model; retried after it dies; skipped once cached."""
monkeypatch.setattr(kc, "_kompress_cache", {})
monkeypatch.setattr(kc, "_download_threads", {})
created: list[object] = []
class FakeThread:
def __init__(self, *, target, args, name, daemon):
self.target, self.args, self.name, self.daemon = target, args, name, daemon
self._alive = True
created.append(self)
def start(self): # do not actually run — simulate a live download
pass
def is_alive(self):
return self._alive
monkeypatch.setattr(kc.threading, "Thread", FakeThread)
kc.ensure_background_download("org/model", "cpu")
kc.ensure_background_download("org/model", "cpu") # thread alive -> no second start
assert len(created) == 1
assert created[0].daemon is True
created[0]._alive = False # simulate the download finishing/failing
kc.ensure_background_download("org/model", "cpu") # dead -> retry
assert len(created) == 2
kc._kompress_cache["org/model"] = ("model", "tokenizer", "onnx")
kc.ensure_background_download("org/model", "cpu") # cached -> no-op
assert len(created) == 2
def _kompress_router() -> ContentRouter:
return ContentRouter(
ContentRouterConfig(
enable_kompress=True,
enable_code_aware=False,
enable_smart_crusher=False,
)
)
def test_router_skips_deep_path_and_fetches_in_background_when_not_ready(monkeypatch):
router = _kompress_router()
calls = {"ensure": 0, "compress": 0}
class NotReadyKompress:
def is_ready(self) -> bool:
return False
def ensure_background_load(self) -> None:
calls["ensure"] += 1
def compress(self, *args, **kwargs):
calls["compress"] += 1
raise AssertionError("must not run the deep path before the model is cached")
monkeypatch.setattr(router, "_get_kompress", lambda: NotReadyKompress())
text = " ".join(["content"] * 40)
out, tokens = router._try_ml_compressor(text, context="")
assert out == text # passthrough, unchanged
assert calls["ensure"] == 1 # background fetch kicked off
assert calls["compress"] == 0 # deep path skipped, no inline download
def test_router_compresses_cache_only_when_ready(monkeypatch):
router = _kompress_router()
seen: dict[str, object] = {}
class ReadyResult:
compressed = "kept words"
compressed_tokens = 2
class ReadyKompress:
def is_ready(self) -> bool:
return True
def ensure_background_load(self) -> None:
raise AssertionError("must not fetch when the model is already cached")
def compress(
self, content, *, context="", question=None, target_ratio=None, allow_download=True
):
seen["allow_download"] = allow_download
return ReadyResult()
monkeypatch.setattr(router, "_get_kompress", lambda: ReadyKompress())
text = " ".join(["content"] * 40)
out, tokens = router._try_ml_compressor(text, context="")
assert seen["allow_download"] is False # request path stays cache-only even when ready
assert out == "kept words"

6
tests/test_transforms/test_content_router.py
View file

@ -150,6 +150,12 @@ def test_force_kompress_routes_anthropic_tool_result_to_targeted_kompress(
captured: dict[str, object] = {}
class FakeKompress:
def is_ready(self) -> bool:
return True
def ensure_background_load(self) -> None:
pass
def compress(self, content, **kwargs):
captured.update(kwargs)
compressed = " ".join(content.split()[:20])