diff --git a/headroom/proxy/helpers.py b/headroom/proxy/helpers.py index 1e5cc62bc..e2914c289 100644 --- a/headroom/proxy/helpers.py +++ b/headroom/proxy/helpers.py @@ -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 diff --git a/headroom/proxy/server.py b/headroom/proxy/server.py index 76a1eacda..87e378297 100644 --- a/headroom/proxy/server.py +++ b/headroom/proxy/server.py @@ -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": diff --git a/tests/test_proxy_eager_preload_bind.py b/tests/test_proxy_eager_preload_bind.py new file mode 100644 index 000000000..c168197a4 --- /dev/null +++ b/tests/test_proxy_eager_preload_bind.py @@ -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()