mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
fix(proxy): report deferred Kompress status and promote health from cache (#2564)
## Description When Kompress preload is deferred until first request, startup still logs "not installed" even if ML deps are present. After the model later loads into the module cache, /readyz and /health can keep reporting kompress as unhealthy because reconcile only inspected attached compressor instances. This PR reports deferred startup accurately and promotes health from the live module cache once the model is ready, without starting loads from health checks. Closes #2560 ## 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 - Treat eager-status `deferred` as installed-but-deferred at proxy startup and log that state instead of "not installed". - Promote `/readyz` and `/health` Kompress readiness from the module-level model cache when attached compressors are missing or not ready. - Keep health inspection free of lazy getters and download side effects. - Add regressions for deferred startup logging and cache-based health promotion. ## Testing - [x] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check .`) - [ ] Type checking passes (`mypy headroom`) - [x] New tests added for new functionality - [ ] Manual testing performed ### Test Output ```text $ PYTHONPATH=/tmp/headroom-2561 python -m pytest tests/test_proxy_health.py tests/test_proxy_eager_preload_bind.py -q -o addopts= 21 passed, 1 warning in 2.73s $ ruff format --check headroom/proxy/server.py tests/test_proxy_health.py tests/test_proxy_eager_preload_bind.py 3 files already formatted $ ruff check headroom/proxy/server.py tests/test_proxy_health.py tests/test_proxy_eager_preload_bind.py All checks passed! ``` ## Real Behavior Proof - Environment: Linux VPS, Python 3.11 venv with headroom-ai 0.32.1 wheel for `_core`, checked out main + this branch overlayed for source under test - Exact command / steps: `PYTHONPATH=/tmp/headroom-2561 python -m pytest tests/test_proxy_health.py tests/test_proxy_eager_preload_bind.py -q -o addopts=`; `ruff format --check` and `ruff check` on the three changed files - Observed result: 21 focused tests passed, including deferred startup log regression and module-cache health promotion; ruff format/check clean - Not tested: live multi-request proxy with real ONNX model download on this host; install-status follow-up mentioned in the issue comment ## Review Readiness - [x] I have performed a self-review - [x] This PR is ready for human review ## Checklist - [x] My code follows the project's style guidelines - [x] I have performed a self-review of my code - [ ] I have commented my code, particularly in hard-to-understand areas - [ ] I have made corresponding changes to the documentation - [x] My changes generate no new warnings - [x] I have added tests that prove my fix is effective or that my feature works - [x] New and existing unit tests pass locally with my changes - [x] I did **not** edit `CHANGELOG.md` — it is generated by release-please from my Conventional Commit PR title (a CI guard enforces this) ## Screenshots (if applicable) N/A ## Additional Notes - Scoped to Kompress status reporting only. The separate `headroom install status` ownership probe in the issue comment is left for a follow-up. Co-authored-by: axelray-dev <axelray-dev@users.noreply.github.com>
This commit is contained in:
parent
4bd121493d
commit
d50cfabedc
3 changed files with 68 additions and 2 deletions
|
|
@ -1708,14 +1708,16 @@ class HeadroomProxy(
|
|||
self.warmup.merge_transform_status(transform_status)
|
||||
|
||||
# Update internal status from eager loading results
|
||||
if eager_status.get("kompress") == "enabled":
|
||||
self._kompress_status = "enabled"
|
||||
if eager_status.get("kompress") in {"enabled", "deferred"}:
|
||||
self._kompress_status = eager_status["kompress"]
|
||||
if eager_status.get("code_aware") == "enabled":
|
||||
self._code_aware_status = "enabled"
|
||||
|
||||
# Log component status
|
||||
if self._kompress_status == "enabled":
|
||||
logger.info("Kompress: ENABLED (ModernBERT token compressor)")
|
||||
elif self._kompress_status == "deferred":
|
||||
logger.info("Kompress: DEFERRED (model loads on first request)")
|
||||
elif self.config.optimize:
|
||||
logger.info("Kompress: not installed (pip install headroom-ai[ml] for ML compression)")
|
||||
|
||||
|
|
@ -2630,6 +2632,21 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
|
|||
return True
|
||||
proxy.warmup.kompress.mark_loaded(handle=compressor, backend=backend)
|
||||
return True
|
||||
|
||||
try:
|
||||
from headroom.transforms.kompress_compressor import HF_MODEL_ID, _kompress_cache
|
||||
except ImportError:
|
||||
return True
|
||||
|
||||
cached = _kompress_cache.get(HF_MODEL_ID)
|
||||
if cached is not None:
|
||||
try:
|
||||
model, _tokenizer, backend = cached
|
||||
except (TypeError, ValueError):
|
||||
pass
|
||||
else:
|
||||
if backend and proxy.warmup.kompress.status != "loaded":
|
||||
proxy.warmup.kompress.mark_loaded(handle=model, backend=backend)
|
||||
return True
|
||||
|
||||
def _health_checks() -> dict[str, dict[str, Any]]:
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ bind still happens and transforms fall back to lazy loading.
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import threading
|
||||
import time
|
||||
|
||||
|
|
@ -118,3 +119,25 @@ async def test_startup_merges_warmup_for_normal_transforms(monkeypatch):
|
|||
assert proxy._kompress_status == "enabled"
|
||||
finally:
|
||||
await proxy.shutdown()
|
||||
|
||||
|
||||
async def test_startup_reports_deferred_kompress(caplog):
|
||||
proxy = _make_proxy(optimize=True)
|
||||
proxy.anthropic_pipeline = _FakePipeline([_FastTransform({"kompress": "deferred"})])
|
||||
proxy.openai_pipeline = _FakePipeline([])
|
||||
|
||||
try:
|
||||
# Proxy setup disables propagation on the ``headroom`` logger, so
|
||||
# attach caplog's handler directly to the logger that emits this line.
|
||||
server_mod.logger.addHandler(caplog.handler)
|
||||
try:
|
||||
with caplog.at_level(logging.INFO, logger=server_mod.logger.name):
|
||||
await proxy.startup()
|
||||
finally:
|
||||
server_mod.logger.removeHandler(caplog.handler)
|
||||
|
||||
assert proxy._kompress_status == "deferred"
|
||||
assert "Kompress: DEFERRED (model loads on first request)" in caplog.messages
|
||||
assert not any("Kompress: not installed" in message for message in caplog.messages)
|
||||
finally:
|
||||
await proxy.shutdown()
|
||||
|
|
|
|||
|
|
@ -3,6 +3,7 @@ from fastapi.testclient import TestClient
|
|||
|
||||
from headroom.proxy.models import ProxyConfig
|
||||
from headroom.proxy.server import create_app
|
||||
from headroom.transforms import kompress_compressor
|
||||
|
||||
|
||||
class _ReadyCompressor:
|
||||
|
|
@ -88,6 +89,31 @@ def test_readyz_promotes_deferred_kompress_after_runtime_load(monkeypatch):
|
|||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("attached", [False, True])
|
||||
def test_readyz_promotes_kompress_from_module_cache(monkeypatch, attached):
|
||||
model = object()
|
||||
monkeypatch.setattr(
|
||||
kompress_compressor,
|
||||
"_kompress_cache",
|
||||
{kompress_compressor.HF_MODEL_ID: (model, object(), "onnx")},
|
||||
)
|
||||
compressor = _ReadyCompressor(ready=False) if attached else None
|
||||
app, proxy = _health_app(monkeypatch, compressor)
|
||||
proxy.warmup.kompress.info["source_status"] = "deferred"
|
||||
|
||||
payload = TestClient(app).get("/readyz").json()
|
||||
|
||||
assert payload["checks"]["kompress"] == {
|
||||
"enabled": True,
|
||||
"ready": True,
|
||||
"status": "healthy",
|
||||
"backend": "onnx",
|
||||
}
|
||||
assert proxy.warmup.kompress.handle is model
|
||||
if compressor is not None:
|
||||
assert compressor.calls == ["is_ready"]
|
||||
|
||||
|
||||
def test_readyz_promotes_remote_kompress_backend(monkeypatch):
|
||||
compressor = _ReadyCompressor(backend="remote")
|
||||
app, proxy = _health_app(monkeypatch)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue