mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
fix(proxy): bind before eager preload so a hung compressor load can't block startup (#1500)
## Description On Windows, `headroom proxy` with optimization enabled sometimes never opens its listening port. `HeadroomProxy.startup()` runs inside the ASGI lifespan, which completes **before** uvicorn binds the socket, and the eager compressor/parser/detector preload ran synchronously there. The per-transform loop already swallows exceptions, so the only thing that can still block the bind is a **hang or an uncatchable native stall** during a model load. That matches the report exactly, including that `--no-optimize` (which skips the preload) binds fine. This decouples the preload from the bind by running it off the event loop under a timeout, so startup always returns and the port binds. Closes #790 ## 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 - `proxy/server.py`: - Extracted the eager-preload loop into a pure sync helper `_eager_preload_transforms()` that returns `(eager_status, transform_statuses)` and does **not** mutate `self.warmup` (so it is safe to run off-thread). - `startup()` now runs it via `asyncio.wait_for(asyncio.to_thread(self._eager_preload_transforms), timeout=EAGER_PRELOAD_TIMEOUT_SECONDS)`. On timeout/exception it logs a warning and continues with empty status, so startup returns and uvicorn binds; transforms fall back to lazy loading on first use. Warmup status is merged on the main thread after the await. - `proxy/helpers.py`: added `EAGER_PRELOAD_TIMEOUT_SECONDS` (default 120s, override via `HEADROOM_EAGER_PRELOAD_TIMEOUT_SECONDS`). The preload is cache-only (`allow_download=False`), so the cap only ever fires on a true hang, never on normal load. - Tests: `tests/test_proxy_eager_preload_bind.py` — helper dedup/exception-swallow, and (via a real `startup()`) that a hung preload no longer blocks startup from returning while a normal transform still merges its warmup status. The happy path is unchanged: a fast preload still completes before `startup()` returns and still populates `self.warmup`. ## Testing - [x] Unit tests pass (`pytest tests/test_proxy_eager_preload_bind.py`) - [x] Linting passes (`ruff check`) - [x] Type checking passes (`mypy headroom`) - [x] New tests added for new functionality - [x] Manual testing performed (live Windows proxy smoke — see proof) ### Test Output ```text $ pytest tests/test_proxy_eager_preload_bind.py -q tests\test_proxy_eager_preload_bind.py ... [100%] 3 passed in 7.66s $ ruff check headroom/proxy/server.py headroom/proxy/helpers.py tests/test_proxy_eager_preload_bind.py All checks passed! $ mypy headroom --ignore-missing-imports # changed files: no new errors ``` ## Real Behavior Proof - Environment: Windows 11, Python 3.13.11, `headroom` 0.28.0, Rust `_core` loaded. - Exact command / steps: start the proxy with optimization enabled (which runs the preload), then curl `/health`. ```text headroom proxy --port 8799 --no-telemetry # optimization ENABLED (runs the preload) curl http://127.0.0.1:8799/health ``` - Observed result: the port binds and `/health` returns HTTP 200 with the preload-bearing startup reported healthy: ```text HTTP_STATUS=200 {"service":"headroom-proxy","status":"healthy","ready":true, "checks":{"startup":{"enabled":true,"ready":true,"status":"healthy","error":null}, ...}, "config":{"optimize":true, ...}, "rust_core":"loaded"} ``` Startup completed and the socket bound with `optimize:true` on a Windows host — the path that previously could hang before binding. - Not tested: a real native model-load hang on Windows (no reliable way to induce the uncatchable native stall on demand). The regression test proves the timeout/bind decoupling deterministically by injecting a transform that blocks past the timeout and asserting `startup()` still returns promptly. ## Review Readiness - [x] I have performed a self-review - [x] This PR is ready for human review ## Additional Notes - Linux CI cannot reproduce the native Windows hang; the regression test proves the decoupling (startup returns despite a blocking preload), not the native root cause. Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
parent
27e010e38f
commit
d5ac07fc45
3 changed files with 196 additions and 26 deletions
|
|
@ -781,6 +781,18 @@ try:
|
|||
except ValueError:
|
||||
COMPRESSION_TIMEOUT_SECONDS = 30.0
|
||||
|
||||
# Eager startup preload timeout in seconds. The preload (compressor/parser models,
|
||||
# cache-only, allow_download=False) runs off the event loop during startup; this
|
||||
# bound only fires on a true hang or an uncatchable native stall so the proxy still
|
||||
# binds its port instead of never opening (GH #790). Override via
|
||||
# HEADROOM_EAGER_PRELOAD_TIMEOUT_SECONDS. Falls back to 120 on an unparseable value.
|
||||
try:
|
||||
EAGER_PRELOAD_TIMEOUT_SECONDS = float(
|
||||
os.environ.get("HEADROOM_EAGER_PRELOAD_TIMEOUT_SECONDS", "120")
|
||||
)
|
||||
except ValueError:
|
||||
EAGER_PRELOAD_TIMEOUT_SECONDS = 120.0
|
||||
|
||||
# Maximum compression cache sessions (prevents unbounded memory growth)
|
||||
MAX_COMPRESSION_CACHE_SESSIONS = 500
|
||||
|
||||
|
|
|
|||
|
|
@ -123,6 +123,7 @@ from headroom.proxy.cost import (
|
|||
)
|
||||
from headroom.proxy.helpers import (
|
||||
COMPRESSION_TIMEOUT_SECONDS, # noqa: F401
|
||||
EAGER_PRELOAD_TIMEOUT_SECONDS,
|
||||
MAX_COMPRESSION_CACHE_SESSIONS, # noqa: F401
|
||||
MAX_MESSAGE_ARRAY_LENGTH, # noqa: F401
|
||||
MAX_REQUEST_BODY_SIZE, # noqa: F401
|
||||
|
|
@ -1293,6 +1294,44 @@ class HeadroomProxy(
|
|||
return "available" # Available but not enabled
|
||||
return "disabled"
|
||||
|
||||
def _eager_preload_transforms(self) -> tuple[dict[str, str], list[dict[str, str]]]:
|
||||
"""Eagerly load every compressor/parser/detector once (dedup by ``id()``).
|
||||
|
||||
Pure load: returns the merged ``eager_status`` plus the per-transform
|
||||
status dicts for the caller to merge into ``self.warmup`` on the main
|
||||
thread (``WarmupRegistry`` is not written off-thread). This runs via
|
||||
``asyncio.to_thread`` so a slow or hung native model load cannot keep
|
||||
startup from binding the port (#790).
|
||||
"""
|
||||
eager_status: dict[str, str] = {}
|
||||
transform_statuses: list[dict[str, str]] = []
|
||||
seen_transform_ids: set[int] = set()
|
||||
for pipeline in (self.anthropic_pipeline, self.openai_pipeline):
|
||||
for transform in pipeline.transforms:
|
||||
if id(transform) in seen_transform_ids:
|
||||
continue
|
||||
seen_transform_ids.add(id(transform))
|
||||
if not hasattr(transform, "eager_load_compressors"):
|
||||
continue
|
||||
try:
|
||||
transform_status = transform.eager_load_compressors()
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"Eager preload failed for %s: %s",
|
||||
type(transform).__name__,
|
||||
exc,
|
||||
)
|
||||
continue
|
||||
if not isinstance(transform_status, dict):
|
||||
continue
|
||||
# Merge: later writers win only if the key wasn't set. Preload a
|
||||
# transform ONCE — if another pipeline also has
|
||||
# ``eager_load_compressors`` it contributes only new keys.
|
||||
for key, value in transform_status.items():
|
||||
eager_status.setdefault(key, value)
|
||||
transform_statuses.append(transform_status)
|
||||
return eager_status, transform_statuses
|
||||
|
||||
async def startup(self):
|
||||
"""Initialize async resources."""
|
||||
self.pipeline_extensions.emit(
|
||||
|
|
@ -1379,32 +1418,31 @@ class HeadroomProxy(
|
|||
|
||||
if self.config.optimize:
|
||||
logger.info("Pre-loading compressors and parsers...")
|
||||
seen_transform_ids: set[int] = set()
|
||||
pipelines = (self.anthropic_pipeline, self.openai_pipeline)
|
||||
for pipeline in pipelines:
|
||||
for transform in pipeline.transforms:
|
||||
if id(transform) in seen_transform_ids:
|
||||
continue
|
||||
seen_transform_ids.add(id(transform))
|
||||
if not hasattr(transform, "eager_load_compressors"):
|
||||
continue
|
||||
try:
|
||||
transform_status = transform.eager_load_compressors()
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"Eager preload failed for %s: %s",
|
||||
type(transform).__name__,
|
||||
exc,
|
||||
)
|
||||
continue
|
||||
if not isinstance(transform_status, dict):
|
||||
continue
|
||||
# Merge: later writers win only if the key wasn't set.
|
||||
# Preload a transform ONCE — if another pipeline also has
|
||||
# ``eager_load_compressors`` it contributes only new keys.
|
||||
for key, value in transform_status.items():
|
||||
eager_status.setdefault(key, value)
|
||||
self.warmup.merge_transform_status(transform_status)
|
||||
# Run the preload OFF the event loop with a bound. The loop body
|
||||
# already swallows per-transform Exceptions, so the only thing that
|
||||
# can still block ASGI lifespan startup (and therefore the socket
|
||||
# bind) is a hang or an uncatchable native stall during a model load
|
||||
# on Windows — the "never opens its port" failure in #790. Capping it
|
||||
# means startup always returns and uvicorn binds; on timeout the
|
||||
# transforms simply fall back to lazy loading on first use.
|
||||
transform_statuses: list[dict[str, str]] = []
|
||||
try:
|
||||
eager_status, transform_statuses = await asyncio.wait_for(
|
||||
asyncio.to_thread(self._eager_preload_transforms),
|
||||
timeout=EAGER_PRELOAD_TIMEOUT_SECONDS,
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"Eager preload exceeded %.0fs or failed (%s); continuing so "
|
||||
"the proxy still binds — transforms load lazily on first use.",
|
||||
EAGER_PRELOAD_TIMEOUT_SECONDS,
|
||||
exc,
|
||||
)
|
||||
eager_status, transform_statuses = {}, []
|
||||
# Merge warmup status on the main thread (WarmupRegistry is not
|
||||
# written off-thread).
|
||||
for transform_status in transform_statuses:
|
||||
self.warmup.merge_transform_status(transform_status)
|
||||
|
||||
# Update internal status from eager loading results
|
||||
if eager_status.get("kompress") == "enabled":
|
||||
|
|
|
|||
120
tests/test_proxy_eager_preload_bind.py
Normal file
120
tests/test_proxy_eager_preload_bind.py
Normal file
|
|
@ -0,0 +1,120 @@
|
|||
"""Startup must bind its port even when eager preload hangs (#790).
|
||||
|
||||
``HeadroomProxy.startup()`` runs inside the ASGI lifespan, which completes
|
||||
*before* uvicorn binds the socket. The eager compressor/parser preload used to
|
||||
run synchronously there, so a hang or an uncatchable native stall during a model
|
||||
load (observed on Windows) left the proxy "never opening its port". The preload
|
||||
now runs off the event loop under ``asyncio.wait_for`` with
|
||||
``EAGER_PRELOAD_TIMEOUT_SECONDS``; on timeout startup logs and continues so the
|
||||
bind still happens and transforms fall back to lazy loading.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import threading
|
||||
import time
|
||||
|
||||
import pytest
|
||||
|
||||
pytest.importorskip("fastapi")
|
||||
|
||||
import headroom.proxy.server as server_mod
|
||||
from headroom.proxy.server import ProxyConfig, create_app
|
||||
|
||||
|
||||
def _make_proxy(*, optimize: bool):
|
||||
config = ProxyConfig(
|
||||
optimize=optimize,
|
||||
cache_enabled=False,
|
||||
rate_limit_enabled=False,
|
||||
cost_tracking_enabled=False,
|
||||
log_requests=False,
|
||||
ccr_inject_tool=False,
|
||||
ccr_handle_responses=False,
|
||||
ccr_context_tracking=False,
|
||||
image_optimize=False,
|
||||
subscription_tracking_enabled=False,
|
||||
)
|
||||
return create_app(config).state.proxy
|
||||
|
||||
|
||||
class _FastTransform:
|
||||
def __init__(self, status):
|
||||
self._status = status
|
||||
|
||||
def eager_load_compressors(self):
|
||||
return self._status
|
||||
|
||||
|
||||
class _RaisingTransform:
|
||||
def eager_load_compressors(self):
|
||||
raise RuntimeError("boom")
|
||||
|
||||
|
||||
class _NonDictTransform:
|
||||
def eager_load_compressors(self):
|
||||
return "not-a-dict"
|
||||
|
||||
|
||||
class _HangingTransform:
|
||||
"""Simulates a model load that hangs forever (released via the event)."""
|
||||
|
||||
def __init__(self, release: threading.Event):
|
||||
self._release = release
|
||||
|
||||
def eager_load_compressors(self):
|
||||
# Safety cap so a misbehaving test can never wedge the suite.
|
||||
self._release.wait(timeout=30)
|
||||
return {"hang": "done"}
|
||||
|
||||
|
||||
class _FakePipeline:
|
||||
def __init__(self, transforms):
|
||||
self.transforms = transforms
|
||||
|
||||
|
||||
def test_eager_preload_dedupes_and_swallows_failures():
|
||||
proxy = _make_proxy(optimize=False)
|
||||
shared = _FastTransform({"shared": "enabled"})
|
||||
proxy.anthropic_pipeline = _FakePipeline([shared, _FastTransform({"kompress": "enabled"})])
|
||||
# ``shared`` appears in both pipelines and must load exactly once; the
|
||||
# raising and non-dict transforms must be skipped without aborting.
|
||||
proxy.openai_pipeline = _FakePipeline([shared, _RaisingTransform(), _NonDictTransform()])
|
||||
|
||||
eager_status, statuses = proxy._eager_preload_transforms()
|
||||
|
||||
assert eager_status == {"shared": "enabled", "kompress": "enabled"}
|
||||
assert statuses == [{"shared": "enabled"}, {"kompress": "enabled"}]
|
||||
|
||||
|
||||
async def test_startup_binds_despite_hung_preload(monkeypatch):
|
||||
monkeypatch.setattr(server_mod, "EAGER_PRELOAD_TIMEOUT_SECONDS", 0.3)
|
||||
proxy = _make_proxy(optimize=True)
|
||||
release = threading.Event()
|
||||
proxy.anthropic_pipeline = _FakePipeline([_HangingTransform(release)])
|
||||
proxy.openai_pipeline = _FakePipeline([])
|
||||
|
||||
try:
|
||||
start = time.monotonic()
|
||||
await proxy.startup() # must NOT wait on the hung load
|
||||
elapsed = time.monotonic() - start
|
||||
# Returns shortly after the 0.3s preload timeout, far below the 30s hang.
|
||||
assert elapsed < 10
|
||||
finally:
|
||||
release.set()
|
||||
await proxy.shutdown()
|
||||
|
||||
|
||||
async def test_startup_merges_warmup_for_normal_transforms(monkeypatch):
|
||||
proxy = _make_proxy(optimize=True)
|
||||
captured: list[dict] = []
|
||||
monkeypatch.setattr(proxy.warmup, "merge_transform_status", captured.append)
|
||||
proxy.anthropic_pipeline = _FakePipeline([_FastTransform({"kompress": "enabled"})])
|
||||
proxy.openai_pipeline = _FakePipeline([])
|
||||
|
||||
try:
|
||||
await proxy.startup()
|
||||
assert {"kompress": "enabled"} in captured
|
||||
assert proxy._kompress_status == "enabled"
|
||||
finally:
|
||||
await proxy.shutdown()
|
||||
Loading…
Add table
Add a link
Reference in a new issue