mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-10 14:27:00 -04:00
3a27c4dacb
## Description
`/debug/warmup` serialized the warmup registry verbatim, so a Kompress
slot left at the startup snapshot kept reporting `{"status": "null",
"info": {"source_status": "deferred"}}` forever — even while the ONNX
model was loaded and actively compressing.
`/health` and `/readyz` already fix this: #2402 added
`_reconcile_kompress_health()`, which promotes the slot from live
runtime state. The debug route never called it, so its answer depended
on whether a health probe happened to run first. That is the half of
#2624 still reproducing on `main`.
Second defect: `WarmupSlot.mark_loaded()` only *updates* `info`, so the
startup-planted `source_status: "deferred"` survived promotion and the
slot serialized as the self-contradictory `{"status": "loaded", "info":
{"source_status": "deferred", "backend": "onnx"}}`.
Closes #2624
## 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/proxy/server.py`: call the existing
`_reconcile_kompress_health()` in the `/debug/warmup` route before
serializing the registry. The reconciler never instantiates a compressor
and never calls `preload()` / `ensure_background_load()` / `compress()`
— it only reads `is_ready()` / `ready_backend()` on an already resident
instance, or falls back to the module-level ONNX cache — so the endpoint
stays side-effect free and idempotent.
- `headroom/proxy/server.py`: stamp `source_status="runtime"` at both
`mark_loaded()` promotion sites in `_reconcile_kompress_health()` (the
resident-compressor path and the `_kompress_cache` fallback),
overwriting the stale startup marker.
- `tests/test_proxy_debug_endpoints.py`: three regression tests plus a
read-only compressor stub whose `preload` / `ensure_background_load`
raise, so a future change that makes the debug route trigger a load
fails loudly.
- `tests/test_proxy_health.py`: assert the promoted slot's
`info["source_status"] == "runtime"`.
## Testing
- [x] Unit tests pass (`pytest`)
- [x] Linting passes (`ruff check .`)
- [x] Type checking passes (`mypy headroom`)
- [x] New tests added for new functionality
- [ ] Manual testing performed
### Test Output
```text
$ pytest tests/test_proxy_debug_endpoints.py tests/test_proxy_health.py tests/test_proxy_warmup.py -q
tests\test_proxy_debug_endpoints.py ............................. [ 52%]
tests\test_proxy_health.py ................. [ 83%]
tests\test_proxy_warmup.py ......... [100%]
============================= 55 passed in 36.56s =============================
$ ruff check headroom/proxy/server.py tests/test_proxy_debug_endpoints.py tests/test_proxy_health.py
All checks passed!
$ ruff format --check headroom/proxy/server.py tests/test_proxy_debug_endpoints.py tests/test_proxy_health.py
3 files already formatted
$ mypy headroom --ignore-missing-imports
Success: no issues found in 506 source files
```
The three new tests were confirmed to be genuine regression tests: with
the `server.py` change reverted and the tests kept, all three fail.
```text
$ git stash push -- headroom/proxy/server.py && pytest tests/test_proxy_debug_endpoints.py -q -k kompress
FAILED tests/test_proxy_debug_endpoints.py::test_debug_warmup_promotes_deferred_kompress_after_runtime_load
FAILED tests/test_proxy_debug_endpoints.py::test_debug_warmup_keeps_pending_kompress_null
FAILED tests/test_proxy_debug_endpoints.py::test_debug_warmup_never_starts_kompress_loading
====================== 3 failed, 26 deselected in 3.98s =======================
```
## Real Behavior Proof
- Environment: Windows 11, Python 3.13.11, pytest 9.1.1, ruff 0.15.x,
mypy 1.20.2, branch based on `main` at 6d5516dc
- Exact command / steps: `pytest tests/test_proxy_debug_endpoints.py
tests/test_proxy_health.py tests/test_proxy_warmup.py -q`, then `git
stash push -- headroom/proxy/server.py` and re-run `pytest
tests/test_proxy_debug_endpoints.py -q -k kompress` to confirm the new
tests fail without the fix
- Observed result: 55 passed with the fix. Without the fix the three new
`/debug/warmup` tests fail — the slot stays `status: "null"` with
`info.source_status: "deferred"` and the stub records zero calls, i.e.
the endpoint never looked at live runtime state. With the fix the same
slot serializes as `{"status": "loaded", "info": {"source_status":
"runtime", "backend": "onnx"}}` and the stub records exactly
`["is_ready", "ready_backend"]` — no load triggered.
- Not tested: the live end-to-end proxy path (cold start, real ONNX
download, real request traffic). This machine has no `onnxruntime` /
`transformers` installed, so a real Kompress load cannot run here; the
tests substitute a stub at the same seam `_reconcile_kompress_health()`
reads. Unrelated to this change, that missing-dependency environment
also makes the pre-existing
`tests/test_kompress_preload_deferral.py::test_proxy_startup_does_not_enter_cached_kompress_native_loader`
fail locally (it reports `source_status: "unavailable"` instead of
`"deferred"`); it fails identically on unmodified `main`.
## 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
- [x] 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)
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
282 lines
8.5 KiB
Python
282 lines
8.5 KiB
Python
import pytest
|
|
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:
|
|
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")
|
|
|
|
app = create_app(
|
|
ProxyConfig(
|
|
optimize=False,
|
|
cache_enabled=False,
|
|
rate_limit_enabled=False,
|
|
)
|
|
)
|
|
app.state.ready = True
|
|
proxy = app.state.proxy
|
|
proxy.http_client = object()
|
|
proxy.warmup.kompress.mark_error("model not cached")
|
|
|
|
client = TestClient(app)
|
|
response = client.get("/readyz")
|
|
|
|
assert response.status_code == 200
|
|
payload = response.json()
|
|
assert payload["ready"] is True
|
|
assert payload["status"] == "healthy"
|
|
assert payload["checks"]["kompress"] == {
|
|
"enabled": True,
|
|
"ready": False,
|
|
"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",
|
|
}
|
|
# 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])
|
|
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)
|
|
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
|