diff --git a/headroom/proxy/server.py b/headroom/proxy/server.py index f7c34f4f3..a6dff74f3 100644 --- a/headroom/proxy/server.py +++ b/headroom/proxy/server.py @@ -2646,7 +2646,9 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI: continue if proxy.warmup.kompress.status == "loaded": return True - proxy.warmup.kompress.mark_loaded(handle=compressor, backend=backend) + proxy.warmup.kompress.mark_loaded( + handle=compressor, backend=backend, source_status="runtime" + ) return True try: @@ -2662,7 +2664,9 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI: pass else: if backend and proxy.warmup.kompress.status != "loaded": - proxy.warmup.kompress.mark_loaded(handle=model, backend=backend) + proxy.warmup.kompress.mark_loaded( + handle=model, backend=backend, source_status="runtime" + ) return True def _health_checks() -> dict[str, dict[str, Any]]: @@ -3292,6 +3296,13 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI: @app.get("/debug/warmup", dependencies=[Depends(_require_loopback)]) async def debug_warmup(): + # Promote a deferred Kompress slot from live runtime state before we + # serialize. Without this the registry keeps reporting the startup + # snapshot (``status: null`` / ``source_status: deferred``) forever, + # even while the model is loaded and compressing, unless somebody + # happens to hit /health or /readyz first. Same read-only + # reconciliation those endpoints run (issue #2624). + _reconcile_kompress_health() warmup_registry = getattr(proxy, "warmup", None) payload = warmup_registry.to_dict() if warmup_registry is not None else {} payload["runtime"] = _runtime_payload() diff --git a/tests/test_proxy_debug_endpoints.py b/tests/test_proxy_debug_endpoints.py index e5ab4529c..6cb8e4017 100644 --- a/tests/test_proxy_debug_endpoints.py +++ b/tests/test_proxy_debug_endpoints.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio +from contextlib import contextmanager import pytest @@ -421,6 +422,100 @@ def test_debug_warmup_reports_registry_slots(client): assert data["runtime"]["websocket_sessions"]["active_relay_tasks"] == 0 +class _KompressStub: + """Read-only stand-in exposing the accessors the health reconciler uses. + + ``preload`` / ``ensure_background_load`` raise so the tests fail loudly if + ``/debug/warmup`` ever triggers a model load instead of just observing. + """ + + def __init__(self, *, backend="onnx", ready=True): + self.backend = backend + self.ready = ready + self.calls: list[str] = [] + + def is_ready(self): + self.calls.append("is_ready") + return self.ready + + def ready_backend(self): + self.calls.append("ready_backend") + return self.backend + + def preload(self): + raise AssertionError("/debug/warmup must never preload kompress") + + def ensure_background_load(self): + raise AssertionError("/debug/warmup must never start a background load") + + +@contextmanager +def _deferred_kompress_client(compressor): + """Client whose kompress slot still carries the startup ``deferred`` mark. + + Mirrors the real cold-start shape: ``eager_load_compressors`` reported + ``deferred``, the model then loaded on the request path, and nothing wrote + the promotion back to the registry. + """ + config = ProxyConfig( + optimize=False, + cache_enabled=False, + rate_limit_enabled=False, + cost_tracking_enabled=False, + ) + app = create_app(config) + with TestClient( + app, + base_url="http://127.0.0.1", + client=("127.0.0.1", 12345), + ) as test_client: + proxy = app.state.proxy + router = proxy.anthropic_pipeline.transforms[-1] + router._kompress = compressor + proxy.warmup.kompress.mark_null() + proxy.warmup.kompress.info["source_status"] = "deferred" + yield proxy, test_client + + +def _clear_kompress_cache(monkeypatch): + """Neutralize the process-global ONNX cache the reconciler falls back to.""" + try: + from headroom.transforms import kompress_compressor + except ImportError: + return + monkeypatch.setattr(kompress_compressor, "_kompress_cache", {}, raising=False) + + +def test_debug_warmup_promotes_deferred_kompress_after_runtime_load(): + compressor = _KompressStub() + with _deferred_kompress_client(compressor) as (_proxy, client): + slot = client.get("/debug/warmup").json()["kompress"] + + assert slot["status"] == "loaded" + assert slot["info"]["backend"] == "onnx" + assert slot["info"]["source_status"] == "runtime" + + +def test_debug_warmup_keeps_pending_kompress_null(monkeypatch): + _clear_kompress_cache(monkeypatch) + compressor = _KompressStub(ready=False) + with _deferred_kompress_client(compressor) as (_proxy, client): + slot = client.get("/debug/warmup").json()["kompress"] + + assert slot["status"] == "null" + assert slot["info"]["source_status"] == "deferred" + assert compressor.calls == ["is_ready"] + + +def test_debug_warmup_never_starts_kompress_loading(): + compressor = _KompressStub() + with _deferred_kompress_client(compressor) as (_proxy, client): + client.get("/debug/warmup") + + # Observation only: no preload(), no ensure_background_load(), no compress(). + assert compressor.calls == ["is_ready", "ready_backend"] + + def test_debug_ws_sessions_reports_live_session(app_and_client): app, client = app_and_client proxy = app.state.proxy diff --git a/tests/test_proxy_health.py b/tests/test_proxy_health.py index e2a2d89d1..efdad4cc1 100644 --- a/tests/test_proxy_health.py +++ b/tests/test_proxy_health.py @@ -87,6 +87,9 @@ def test_readyz_promotes_deferred_kompress_after_runtime_load(monkeypatch): "status": "healthy", "backend": "onnx", } + # The promotion must also clear the startup marker, otherwise the slot + # serializes as loaded-but-deferred in /debug/warmup. + assert proxy.warmup.kompress.info["source_status"] == "runtime" @pytest.mark.parametrize("attached", [False, True])