fix(proxy/debug): reconcile Kompress warmup state in /debug/warmup (#2711)

## 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>
This commit is contained in:
Parideboy 2026-08-02 22:15:01 +02:00 committed by GitHub
parent 9ce5af02b1
commit 3a27c4dacb
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 111 additions and 2 deletions

View file

@ -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()

View file

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

View file

@ -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])