diff --git a/headroom/proxy/server.py b/headroom/proxy/server.py index b069335e0..59fe81cfb 100644 --- a/headroom/proxy/server.py +++ b/headroom/proxy/server.py @@ -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]]: diff --git a/tests/test_proxy_eager_preload_bind.py b/tests/test_proxy_eager_preload_bind.py index c168197a4..45b05c549 100644 --- a/tests/test_proxy_eager_preload_bind.py +++ b/tests/test_proxy_eager_preload_bind.py @@ -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() diff --git a/tests/test_proxy_health.py b/tests/test_proxy_health.py index af459535e..e2a2d89d1 100644 --- a/tests/test_proxy_health.py +++ b/tests/test_proxy_health.py @@ -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)