mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
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:
parent
fc9c63f18c
commit
54526bc858
4 changed files with 268 additions and 1 deletions
|
|
@ -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),
|
||||
),
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue