fix(proxy): promote Kompress health after runtime load (#2402)

## Description

`/readyz` can keep reporting Kompress as `{"ready": false, "status":
"unhealthy", "backend": null}` after the live compressor has already
become ready. Startup intentionally records Kompress as `deferred`
without loading the model, `WarmupRegistry.merge_transform_status()`
stores that only as metadata, and the health check later serializes the
stale warmup slot instead of the live runtime compressor state. The
request path can already see the real readiness signal through
`KompressCompressor.is_ready()`, but nothing promotes the health surface
after startup.

This change keeps startup behavior untouched and reconciles Kompress
health from the live compressor right before `/readyz` serializes
component state. It adds side-effect-free runtime backend accessors for
local and remote Kompress implementations, promotes the warmup slot only
when the runtime compressor is ready, preserves loaded state on
transient inspection failures, and keeps Kompress excluded from
aggregate readiness.

Closes #2386

## 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

- `headroom/transforms/kompress_compressor.py`: add a side-effect-free
`ready_backend()` accessor that returns the cached backend for the
current model or `None`.
- `headroom/transforms/kompress_remote.py`: add `ready_backend()`
returning `"remote"` for the always-ready remote adapter.
- `headroom/proxy/server.py`: derive Kompress health from the live
enabled `ContentRouter` instances, promote the warmup slot only when
runtime readiness is real, respect per-provider re-enable overrides, and
preserve loaded state on transient inspection failures.
- `tests/test_proxy_health.py`: add focused regression, override,
pending, remote, no-instantiation, disabled, fail-open, and
aggregate-readiness coverage.
- `tests/test_kompress_preload_deferral.py`: keep startup-deferral proof
current if a helper needs the new accessor surface.

## Testing

- [x] Unit tests pass (`uv run pytest tests/test_proxy_health.py
tests/test_kompress_preload_deferral.py
tests/test_kompress_request_nonblocking.py -q`)
- [x] Linting passes (`uv run ruff check headroom/proxy/server.py
headroom/transforms/kompress_compressor.py
headroom/transforms/kompress_remote.py tests/test_proxy_health.py
tests/test_kompress_preload_deferral.py`)
- [x] Formatting passes (`uv run ruff format headroom/proxy/server.py
headroom/transforms/kompress_compressor.py
headroom/transforms/kompress_remote.py tests/test_proxy_health.py
tests/test_kompress_preload_deferral.py --check`)
- [x] New tests added for new functionality when applicable
- [x] Manual testing performed

### Test Output

```text
$ uv run pytest tests/test_proxy_health.py tests/test_kompress_preload_deferral.py tests/test_kompress_request_nonblocking.py -q
................................                                         [100%]
32 passed, 1 warning in 2.06s

$ uv run ruff check headroom/proxy/server.py headroom/transforms/kompress_compressor.py headroom/transforms/kompress_remote.py tests/test_proxy_health.py tests/test_kompress_preload_deferral.py
All checks passed!

$ uv run ruff format headroom/proxy/server.py headroom/transforms/kompress_compressor.py headroom/transforms/kompress_remote.py tests/test_proxy_health.py tests/test_kompress_preload_deferral.py --check
4 files already formatted
```

## Real Behavior Proof

- Environment: Windows host, local FastAPI test app with the same
`HeadroomProxy`, `WarmupRegistry`, and `/readyz` route used in
production
- Exact command / steps: run `uv run pytest tests/test_proxy_health.py
tests/test_kompress_preload_deferral.py
tests/test_kompress_request_nonblocking.py -q`, covering a deferred
startup slot, a pending resident compressor, a global-disable plus
`disable_kompress_anthropic=False` override, and a router whose lazy
getters would raise if health instantiated them
- Observed result: deferred runtime readiness promotes to `{"enabled":
true, "ready": true, "status": "healthy", "backend": "onnx"}`, a pending
resident compressor stays `{"ready": false, "backend": null}`, a
per-provider override re-enables health even when the global flag is
off, and the health path never instantiates Kompress
- Not tested: live remote Kompress endpoint behavior beyond the local
remote-adapter contract

## 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 made corresponding changes to the documentation if needed
- [x] My changes generate no new warnings
- [x] I have added tests that prove my fix is effective
- [x] New and existing unit tests pass locally with my changes

## Additional Notes

- `CHANGELOG.md` stays untouched because Headroom generates changelog
entries from conventional commits.
- Kompress remains a soft component already excluded from aggregate
readiness. This PR fixes only the per-component health report.
- The health path must remain read-only; it must not call `preload()`,
`ensure_background_load()`, `compress()`, or any network or model I/O.
This commit is contained in:
Rod Boev 2026-07-18 19:46:54 -04:00 committed by GitHub
parent fc9c63f18c
commit 54526bc858
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 268 additions and 1 deletions

View file

@ -2578,7 +2578,47 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
**details,
}
def _kompress_health_routers() -> list[ContentRouter]:
routers: list[ContentRouter] = []
for pipeline in (proxy.anthropic_pipeline, proxy.openai_pipeline):
for transform in getattr(pipeline, "transforms", ()):
if (
isinstance(transform, ContentRouter)
and transform.config.enable_kompress
and all(transform is not item for item in routers)
):
routers.append(transform)
return routers
def _reconcile_kompress_health() -> bool:
routers = _kompress_health_routers()
if not routers:
return False
compressors: list[Any] = []
for router in routers:
for name in ("_kompress", "_kompress_remote"):
compressor = getattr(router, name, None)
if compressor is not None and all(compressor is not item for item in compressors):
compressors.append(compressor)
for compressor in compressors:
try:
if not compressor.is_ready():
continue
backend = compressor.ready_backend()
if not backend:
continue
except Exception:
continue
if proxy.warmup.kompress.status == "loaded":
return True
proxy.warmup.kompress.mark_loaded(handle=compressor, backend=backend)
return True
return True
def _health_checks() -> dict[str, dict[str, Any]]:
kompress_enabled = _reconcile_kompress_health()
memory_status = (
proxy.memory_handler.health_status()
if proxy.memory_handler
@ -2625,7 +2665,7 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI:
error=_upstream_check_cache["error"],
),
"kompress": _component_health(
enabled=not config.disable_kompress,
enabled=kompress_enabled,
ready=proxy.warmup.kompress.status == "loaded",
backend=proxy.warmup.kompress.info.get("backend", None),
),

View file

@ -1177,6 +1177,11 @@ class KompressCompressor(Transform):
"""
return self.config.model_id in _kompress_cache
def ready_backend(self) -> str | None:
"""Return the cached backend without triggering a load."""
entry = _kompress_cache.get(self.config.model_id)
return entry[2] if entry is not None else None
def ensure_background_load(self) -> None:
"""Kick off a one-shot, non-blocking background download of the model.

View file

@ -61,6 +61,9 @@ class RemoteKompressCompressor:
def is_ready(self) -> bool:
return True
def ready_backend(self) -> str | None:
return "remote"
def preload(self, *, allow_download: bool = True) -> str:
return "remote"

View file

@ -1,9 +1,48 @@
import pytest
from fastapi.testclient import TestClient
from headroom.proxy.models import ProxyConfig
from headroom.proxy.server import create_app
class _ReadyCompressor:
def __init__(self, backend="onnx", ready=True, error=None):
self.backend = backend
self.ready = ready
self.error = error
self.calls = []
def is_ready(self):
self.calls.append("is_ready")
if self.error:
raise self.error
return self.ready
def ready_backend(self):
self.calls.append("ready_backend")
return self.backend
def _health_app(monkeypatch, compressor=None, *, disabled=False, **config_kwargs):
monkeypatch.setenv("HEADROOM_SKIP_UPSTREAM_CHECK", "1")
app = create_app(
ProxyConfig(
optimize=False,
cache_enabled=False,
rate_limit_enabled=False,
disable_kompress=disabled,
**config_kwargs,
)
)
app.state.ready = True
proxy = app.state.proxy
proxy.http_client = object()
router = proxy.anthropic_pipeline.transforms[-1]
if compressor is not None:
router._kompress = compressor
return app, proxy
def test_readyz_excludes_kompress_from_aggregate_readiness(monkeypatch):
monkeypatch.setenv("HEADROOM_SKIP_UPSTREAM_CHECK", "1")
@ -32,3 +71,183 @@ def test_readyz_excludes_kompress_from_aggregate_readiness(monkeypatch):
"status": "unhealthy",
"backend": None,
}
def test_readyz_promotes_deferred_kompress_after_runtime_load(monkeypatch):
compressor = _ReadyCompressor()
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",
}
def test_readyz_promotes_remote_kompress_backend(monkeypatch):
compressor = _ReadyCompressor(backend="remote")
app, proxy = _health_app(monkeypatch)
router = proxy.anthropic_pipeline.transforms[-1]
router._kompress = None
router._kompress_remote = compressor
payload = TestClient(app).get("/readyz").json()
assert payload["checks"]["kompress"]["backend"] == "remote"
assert payload["checks"]["kompress"]["ready"] is True
def test_readyz_keeps_pending_kompress_unloaded(monkeypatch):
compressor = _ReadyCompressor(backend="onnx", ready=False)
app, proxy = _health_app(monkeypatch, compressor)
router = proxy.anthropic_pipeline.transforms[-1]
router._kompress = compressor
payload = TestClient(app).get("/readyz").json()
assert payload["checks"]["kompress"] == {
"enabled": True,
"ready": False,
"status": "unhealthy",
"backend": None,
}
assert compressor.calls == ["is_ready"]
def test_readyz_never_starts_kompress_loading(monkeypatch):
compressor = _ReadyCompressor()
app, proxy = _health_app(monkeypatch, compressor)
router = proxy.anthropic_pipeline.transforms[-1]
router._kompress = compressor
TestClient(app).get("/readyz")
assert compressor.calls == ["is_ready", "ready_backend"]
def test_readyz_kompress_inspection_failure_fails_open(monkeypatch):
compressor = _ReadyCompressor(error=RuntimeError("inspection failed"))
app, proxy = _health_app(monkeypatch, compressor)
proxy.warmup.kompress.mark_loaded(handle=object(), backend="onnx")
payload = TestClient(app).get("/readyz").json()
assert payload["checks"]["kompress"]["ready"] is True
assert payload["checks"]["kompress"]["backend"] == "onnx"
def test_readyz_disabled_kompress_skips_inspection(monkeypatch):
compressor = _ReadyCompressor()
app, proxy = _health_app(monkeypatch, compressor, disabled=True)
router = proxy.anthropic_pipeline.transforms[-1]
router._kompress = compressor
payload = TestClient(app).get("/readyz").json()
assert payload["checks"]["kompress"] == {
"enabled": False,
"ready": True,
"status": "disabled",
"backend": None,
}
assert compressor.calls == []
def test_readyz_per_provider_kompress_override_reenables_health(monkeypatch):
compressor = _ReadyCompressor()
app, proxy = _health_app(
monkeypatch,
disabled=True,
disable_kompress_anthropic=False,
)
router = proxy.anthropic_pipeline.transforms[-1]
router._kompress = compressor
payload = TestClient(app).get("/readyz").json()
assert payload["checks"]["kompress"] == {
"enabled": True,
"ready": True,
"status": "healthy",
"backend": "onnx",
}
assert compressor.calls == ["is_ready", "ready_backend"]
def test_readyz_never_calls_lazy_kompress_getters(monkeypatch):
app, proxy = _health_app(monkeypatch)
router = proxy.anthropic_pipeline.transforms[-1]
def _boom():
raise AssertionError("health should not instantiate kompress")
router._get_kompress = _boom
router._get_remote_kompress = _boom
payload = TestClient(app).get("/readyz").json()
assert payload["checks"]["kompress"] == {
"enabled": True,
"ready": False,
"status": "unhealthy",
"backend": None,
}
@pytest.mark.parametrize(
("slot_status", "compressor", "disabled", "expected"),
[
(
"null",
None,
False,
{"enabled": True, "ready": False, "status": "unhealthy", "backend": None},
),
(
"null",
_ReadyCompressor(),
False,
{"enabled": True, "ready": True, "status": "healthy", "backend": "onnx"},
),
(
"null",
_ReadyCompressor(backend="remote"),
False,
{"enabled": True, "ready": True, "status": "healthy", "backend": "remote"},
),
(
"error",
_ReadyCompressor(),
False,
{"enabled": True, "ready": True, "status": "healthy", "backend": "onnx"},
),
(
"loaded",
_ReadyCompressor(),
False,
{"enabled": True, "ready": True, "status": "healthy", "backend": "existing"},
),
(
"null",
_ReadyCompressor(),
True,
{"enabled": False, "ready": True, "status": "disabled", "backend": None},
),
],
)
def test_readyz_kompress_state_matrix(monkeypatch, slot_status, compressor, disabled, expected):
app, proxy = _health_app(monkeypatch, compressor, disabled=disabled)
if slot_status == "error":
proxy.warmup.kompress.mark_error("not cached")
elif slot_status == "loaded":
proxy.warmup.kompress.mark_loaded(handle=object(), backend=expected["backend"])
if compressor is not None and expected["backend"] == "existing":
compressor.backend = "new"
payload = TestClient(app).get("/readyz").json()["checks"]["kompress"]
assert payload == expected