From 15ac650d409ea7def9e54d9962af1cfdc1f11f5d Mon Sep 17 00:00:00 2001 From: Rod Boev Date: Tue, 30 Jun 2026 14:41:22 -0400 Subject: [PATCH] fix(proxy): fail open when kompress saturation would exhaust pre-upstream budget (#1430) ## Description Concurrent Anthropic `/v1/messages` traffic can still exhaust Headroom's pre-upstream budget because Kompress ONNX execution waits on the request critical path. When Kompress saturates, requests eventually fail with `503 pre-upstream queue saturated` even though compression can safely degrade to passthrough. This PR makes Kompress saturation fail open on the Anthropic hot path, so requests continue uncompressed when compression capacity is under pressure. It keeps the executor and stage-timing evidence intact, and it preserves blocking model-load validation so runtime pressure does not silently skip the validation path. Closes #1025 ## 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 - add a bounded execution-slot acquire path so Anthropic requests fail open to passthrough when Kompress saturation would consume the pre-upstream budget - preserve explicit execution-timeout counters and Anthropic passthrough/stage-timing observability instead of hiding the pressure path - keep `_validate_pytorch_device()` on blocking acquire semantics so model-load validation still waits for capacity instead of failing open - make the blocking validation acquire explicit to `mypy` without changing runtime behavior - extend focused regressions for pre-upstream backpressure, Kompress saturation, execution-skip observability, and validation waiting - align the CLI timeout help text and `ProxyConfig` comment with the fail-open runtime behavior - update `CHANGELOG.md` for the proxy runtime fix ## Testing - [x] Unit tests pass (`uv run pytest tests/test_anthropic_pre_upstream_backpressure.py tests/test_proxy_compression_executor.py tests/test_kompress_request_nonblocking.py`) - [x] Linting passes (`uv run ruff check tests/test_anthropic_pre_upstream_backpressure.py` and `uv run ruff format tests/test_anthropic_pre_upstream_backpressure.py --check`) - [x] Type checking passes (`uv run mypy headroom --ignore-missing-imports`) - [x] New tests added for new functionality when applicable - [ ] Manual testing performed ### Test Output ```text Focused local validation passed: - uv run pytest tests/test_anthropic_pre_upstream_backpressure.py tests/test_proxy_compression_executor.py tests/test_kompress_request_nonblocking.py -x -v 37 passed, 1 warning in 12.01s - uv run ruff check headroom/proxy/handlers/anthropic.py headroom/transforms/kompress_compressor.py tests/test_anthropic_pre_upstream_backpressure.py tests/test_proxy_compression_executor.py tests/test_kompress_request_nonblocking.py All checks passed! - uv run ruff format headroom/proxy/handlers/anthropic.py headroom/transforms/kompress_compressor.py tests/test_anthropic_pre_upstream_backpressure.py tests/test_proxy_compression_executor.py tests/test_kompress_request_nonblocking.py --check 5 files already formatted - uv run mypy headroom --ignore-missing-imports Success: no issues found in 398 source files Base-branch proof on origin/main (fa05ebc849abf1c7fdffac7245ed190ae513d2c4): - test_acquire_timeout_degrades_to_passthrough fails because the handler still returns 503 - test_saturation_fail_open_does_not_hang_request fails because get_kompress_execution_stats() does not exist - test_compression_executor_skip_signal_remains_visible passes on base too, so it stays as compatibility coverage rather than the failing-then-passing proof for this fix Review-follow-up validation passed after aligning the timeout wording with fail-open behavior: - uv run pytest tests/test_anthropic_pre_upstream_backpressure.py -x -v 20 passed, 1 warning in 1.38s - uv run ruff check headroom/cli/proxy.py headroom/proxy/models.py headroom/proxy/handlers/anthropic.py headroom/transforms/kompress_compressor.py tests/test_anthropic_pre_upstream_backpressure.py tests/test_proxy_compression_executor.py tests/test_kompress_request_nonblocking.py All checks passed! ``` ## Real Behavior Proof - Environment: local Anthropic pre-upstream and Kompress execution regression harnesses covering the `/v1/messages` hot path - Exact command / steps: run the focused pytest command above on `origin/main` and on this branch, including the semaphore-saturation path in `test_saturation_fail_open_does_not_hang_request` and the validation-slot hold in `test_validation_probe_waits_for_execution_slot` - Observed result: the reviewed head no longer returns `503` on the pre-upstream pressure path, request-thread Kompress saturation degrades to passthrough while incrementing execution timeout stats, and model-load validation still waits for capacity instead of failing open - Not tested: wrap/install fallout mentioned in the original issue ## 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 commented my code, particularly in hard-to-understand areas - [ ] I have made corresponding changes to the documentation - [x] My changes generate no new type-check or lint 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 have updated the CHANGELOG.md if applicable ## Additional Notes - Scoped to the runtime queue-pressure fault only; the issue's wrap/unwrap and deployment complaints stay out of this PR. - `test_compression_executor_skip_signal_remains_visible` remains in the suite to prove the skip signal stays visible, but it is compatibility coverage rather than the failing-then-passing regression for the bug fix. - Local validation included `uv run mypy headroom --ignore-missing-imports` after the explicit validation-acquire narrowing was added for CI parity. - Attribution: the issue body isolated the hot-path ONNX compression stall and the pre-upstream saturation symptom that this PR fixes. --- CHANGELOG.md | 1 + headroom/cli/proxy.py | 2 +- headroom/proxy/handlers/anthropic.py | 36 +++-- headroom/proxy/models.py | 2 +- headroom/transforms/kompress_compressor.py | 122 +++++++++++++- ...est_anthropic_pre_upstream_backpressure.py | 20 ++- tests/test_kompress_request_nonblocking.py | 153 ++++++++++++++++++ tests/test_proxy_compression_executor.py | 61 +++++++ 8 files changed, 368 insertions(+), 29 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 8d9e2b1bf..a0317a9dd 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -73,6 +73,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 * **ccr:** make retrieval store TTL configurable with `HEADROOM_CCR_TTL_SECONDS`, expose the effective TTL in `/v1/retrieve/stats`, and distinguish expired retrievals from missing hashes. * **proxy:** make `force_kompress` skip ContentRouter auto-detection during compression and pass savings-profile kwargs through Anthropic batch requests. * **proxy:** add native Bedrock `/model/{id}/converse-stream` route and forward it through the existing streaming EventStream/SSE pipeline. +* **proxy/kompress:** make pre-upstream backpressure and kompress execution saturation fail-open, so Anthropic requests no longer return 503 during temporary saturation while healthy capacity still compresses and explicit passthrough markers preserve operator visibility ([#1025](https://github.com/headroomlabs-ai/headroom/issues/1025)). * **wrap (codex):** fix `headroom wrap codex` producing a `config.toml` with duplicate top-level `model_provider` / `openai_base_url` keys (TOML-spec error) when the user had already configured their own provider. The injector now rewrites pre-existing top-level `model_provider` and `openai_base_url` lines in place — the previous value is kept in a `# was: …` trailing comment — instead of unconditionally prepending a duplicate, so `codex` can start against the proxy. The pre-wrap snapshot mechanism continues to byte-for-byte restore the original file on `headroom unwrap codex`. * **install (macOS):** fix `headroom install restart` / `install start` for launchd `persistent-service` deployments. `stop` `bootout`s the job but `start` only ran `launchctl kickstart`, which cannot recover the un-bootstrapped state `stop`/`restart` leave behind (launchctl error 113), so the proxy was left stopped. `start` now tries `kickstart` (fast path for an already-bootstrapped job) and, on failure, `bootstrap`s the plist fresh — retrying for ~15s to ride out the transient `bootstrap` EIO (error 5) window while launchd releases the label after a `bootout`. `stop` tolerates only the already-absent case (`bootout` ESRCH / error 3) and still raises on any other `bootout` failure ([#1289](https://github.com/headroomlabs-ai/headroom/issues/1289)). * **wrap:** isolate wrapped proxy subprocess stdout/stderr into `proxy-stdio.log`, so `proxy.log` remains the canonical rotating runtime log and Windows rollover failures from `RotatingFileHandler` are no longer blocked by wrapper stdio handles ([#1184](https://github.com/chopratejas/headroom/issues/1184)). diff --git a/headroom/cli/proxy.py b/headroom/cli/proxy.py index 910ea46ca..88e4c0c26 100644 --- a/headroom/cli/proxy.py +++ b/headroom/cli/proxy.py @@ -400,7 +400,7 @@ def dashboard(port: int, no_open: bool) -> None: envvar="HEADROOM_ANTHROPIC_PRE_UPSTREAM_ACQUIRE_TIMEOUT_SECONDS", help=( "Fail-fast timeout for waiting on the Anthropic pre-upstream semaphore " - "before returning 503 + Retry-After. " + "before failing open to passthrough compression. " "Default: 15.0 seconds. " "Env: HEADROOM_ANTHROPIC_PRE_UPSTREAM_ACQUIRE_TIMEOUT_SECONDS." ), diff --git a/headroom/proxy/handlers/anthropic.py b/headroom/proxy/handlers/anthropic.py index e10d1b171..1ad3853c5 100644 --- a/headroom/proxy/handlers/anthropic.py +++ b/headroom/proxy/handlers/anthropic.py @@ -535,6 +535,7 @@ class AnthropicHandlerMixin: ) if pre_upstream_sem is not None: + _pre_upstream_saturated = False _wait_started_at = time.perf_counter() _acquire_timeout_seconds = self.config.anthropic_pre_upstream_acquire_timeout_seconds try: @@ -544,7 +545,7 @@ class AnthropicHandlerMixin: ) except asyncio.TimeoutError: _wait_ms = (time.perf_counter() - _wait_started_at) * 1000.0 - stage_timer.record("pre_upstream_wait", _wait_ms) + _pre_upstream_saturated = True logger.warning( "[%s] Anthropic pre-upstream queue saturated after %.2f ms " "(timeout=%.1fs, session_id=%s)", @@ -553,22 +554,13 @@ class AnthropicHandlerMixin: _acquire_timeout_seconds, trace_session_id, ) - await _finalize_pre_upstream() - return JSONResponse( - status_code=503, - headers={"Retry-After": str(max(1, int(_acquire_timeout_seconds) + 1))}, - content={ - "type": "error", - "error": { - "type": "service_unavailable", - "message": ( - "Anthropic pre-upstream queue is saturated. Please retry shortly." - ), - }, - }, + logger.info( + "[%s] pre-upstream saturation fail-open; continuing without compression path", + request_id, ) - _pre_upstream_sem_acquired = True - _wait_ms = (time.perf_counter() - _wait_started_at) * 1000.0 + else: + _pre_upstream_sem_acquired = True + _wait_ms = (time.perf_counter() - _wait_started_at) * 1000.0 stage_timer.record("pre_upstream_wait", _wait_ms) if _wait_ms > 100.0: logger.info( @@ -580,6 +572,7 @@ class AnthropicHandlerMixin: ) else: stage_timer.record("pre_upstream_wait", 0.0) + _pre_upstream_saturated = False try: # Check request body size @@ -1025,11 +1018,20 @@ class AnthropicHandlerMixin: messages=messages, ) _decision.apply_to_tags(tags) + _skip_compression_for_backpressure = ( + _pre_upstream_saturated and _decision.should_compress + ) + if _skip_compression_for_backpressure: + tags["passthrough_reason"] = "pre_upstream_backpressure" + logger.info( + "[%s] Compression skipped: reason=pre_upstream_backpressure", + request_id, + ) if not _decision.should_compress: logger.info( f"[{request_id}] Compression skipped: reason={_decision.passthrough_reason}" ) - if _decision.should_compress: + if _decision.should_compress and not _skip_compression_for_backpressure: try: from headroom.proxy.helpers import COMPRESSION_TIMEOUT_SECONDS diff --git a/headroom/proxy/models.py b/headroom/proxy/models.py index bdd3e0818..d9b3302c4 100644 --- a/headroom/proxy/models.py +++ b/headroom/proxy/models.py @@ -373,7 +373,7 @@ class ProxyConfig: # Precedence: CLI > env > auto-compute. anthropic_pre_upstream_concurrency: int | None = None # Upper bound for waiting on the Anthropic pre-upstream semaphore - # before failing fast with a 503 + Retry-After. Keeps the queue bounded + # before failing open to passthrough compression. Keeps the queue bounded # when all pre-upstream slots are occupied by slow/hung work. anthropic_pre_upstream_acquire_timeout_seconds: float = 15.0 # Fail-open timeout for Anthropic memory-context lookup while the request diff --git a/headroom/transforms/kompress_compressor.py b/headroom/transforms/kompress_compressor.py index 0f08e0ed4..9ad410399 100644 --- a/headroom/transforms/kompress_compressor.py +++ b/headroom/transforms/kompress_compressor.py @@ -94,6 +94,8 @@ KOMPRESS_ONNX_INTRA_THREADS_ENV = "HEADROOM_KOMPRESS_ONNX_INTRA_THREADS" KOMPRESS_ONNX_INTER_THREADS_ENV = "HEADROOM_KOMPRESS_ONNX_INTER_THREADS" KOMPRESS_COREML_CACHE_DIR_ENV = "HEADROOM_KOMPRESS_COREML_CACHE_DIR" KOMPRESS_MAX_CONCURRENT_ENV = "HEADROOM_KOMPRESS_MAX_CONCURRENT" +KOMPRESS_EXECUTION_SEMAPHORE_WAIT_MS_ENV = "HEADROOM_KOMPRESS_EXECUTION_TIMEOUT_MS" +KOMPRESS_EXECUTION_SEMAPHORE_WAIT_MS_DEFAULT = 25 KOMPRESS_BATCH_SIZE_ENV = "HEADROOM_KOMPRESS_BATCH_SIZE" KompressBackend = Literal["auto", "onnx", "onnx_cpu", "onnx_coreml", "pytorch", "pytorch_mps"] @@ -127,6 +129,76 @@ _kompress_cache: dict[str, tuple[Any, Any, str]] = {} _kompress_lock = threading.Lock() _execution_semaphores: dict[str, threading.BoundedSemaphore] = {} _execution_semaphores_lock = threading.Lock() +_execution_metrics_lock = threading.Lock() +_execution_skip_counters: dict[str, int] = { + "timeout": 0, +} +_execution_wait_seconds_total: dict[str, float] = { + "timeout": 0.0, +} + + +def _execution_wait_budget_seconds() -> float: + raw = os.environ.get(KOMPRESS_EXECUTION_SEMAPHORE_WAIT_MS_ENV) + if raw is None: + return KOMPRESS_EXECUTION_SEMAPHORE_WAIT_MS_DEFAULT / 1000.0 + try: + parsed = int(raw) + except ValueError: + logger.warning( + "Invalid %s=%r; using %dms", + KOMPRESS_EXECUTION_SEMAPHORE_WAIT_MS_ENV, + raw, + KOMPRESS_EXECUTION_SEMAPHORE_WAIT_MS_DEFAULT, + ) + return KOMPRESS_EXECUTION_SEMAPHORE_WAIT_MS_DEFAULT / 1000.0 + if parsed < 0: + logger.warning( + "Negative %s=%r; disabling timeout and using fail-open.", + KOMPRESS_EXECUTION_SEMAPHORE_WAIT_MS_ENV, + raw, + ) + return 0.0 + return parsed / 1000.0 + + +def _acquire_execution_slot( + backend: str, + device_type: str, + *, + timeout_seconds: float | None, +) -> tuple[threading.BoundedSemaphore | None, float]: + semaphore = _execution_semaphore(backend, device_type) + start = time.perf_counter() + if timeout_seconds is None: + semaphore.acquire() + wait_ms = (time.perf_counter() - start) * 1000.0 + return semaphore, wait_ms + + acquired = semaphore.acquire(blocking=False) + if not acquired and timeout_seconds > 0: + acquired = semaphore.acquire(timeout=timeout_seconds) + elif not acquired and timeout_seconds == 0: + acquired = False + + wait_ms = (time.perf_counter() - start) * 1000.0 + if not acquired: + with _execution_metrics_lock: + _execution_skip_counters["timeout"] += 1 + _execution_wait_seconds_total["timeout"] += wait_ms / 1000.0 + return None, wait_ms + + return semaphore, wait_ms + + +def get_kompress_execution_stats() -> dict[str, int | float]: + """Return execution-acquire observability counters.""" + with _execution_metrics_lock: + return { + "execution_acquire_timeout_ms": int(_execution_wait_budget_seconds() * 1000), + "execution_timeout_skips_total": _execution_skip_counters["timeout"], + "execution_wait_seconds_total": _execution_wait_seconds_total["timeout"], + } def _selected_backend() -> KompressBackend: @@ -602,7 +674,14 @@ def _validate_pytorch_device(model: Any, tokenizer: Any, device: str) -> None: ) input_ids = encoding["input_ids"].to(device) attention_mask = encoding["attention_mask"].to(device) - with _execution_semaphore("pytorch", device): + semaphore, _wait_ms = _acquire_execution_slot( + "pytorch", + device, + timeout_seconds=None, + ) + assert semaphore is not None + with contextlib.ExitStack() as stack: + stack.callback(semaphore.release) scores = model.get_scores(input_ids, attention_mask) _ = scores[0].detach().cpu() @@ -969,7 +1048,24 @@ class KompressCompressor(Transform): input_ids = input_ids.to(device) attention_mask = attention_mask.to(device) - with _execution_semaphore(backend, device_type): + semaphore, _wait_ms = _acquire_execution_slot( + backend, + device_type, + timeout_seconds=_execution_wait_budget_seconds(), + ) + if semaphore is None: + logger.warning( + "Kompress execution saturated after %.2fms; skipping chunk=%d " + "for backend=%s device=%s after deadline path", + _wait_ms, + chunk_start, + backend, + device_type, + ) + return self._passthrough(content, n_words) + + with contextlib.ExitStack() as stack: + stack.callback(semaphore.release) inference_started = time.perf_counter() if target_ratio is not None: scores = model.get_scores(input_ids, attention_mask) @@ -1228,7 +1324,27 @@ class KompressCompressor(Transform): attention_mask = attention_mask.to(device) # Single forward pass for all chunks in this batch. - with _execution_semaphore(backend, device_type): + semaphore, wait_ms = _acquire_execution_slot( + backend, + device_type, + timeout_seconds=_execution_wait_budget_seconds(), + ) + if semaphore is None: + logger.warning( + "Kompress execution saturated at batch start after %.2fms; " + "passing through remaining batch inputs", + wait_ms, + ) + for text_idx, _, _, _ in batch: + if results[text_idx] is None: + results[text_idx] = self._passthrough( + contents[text_idx], len(word_lists[text_idx]) + ) + kept_ids_per_text.pop(text_idx, None) + continue + + with contextlib.ExitStack() as stack: + stack.callback(semaphore.release) inference_started = time.perf_counter() scores = model.get_scores(input_ids, attention_mask) inference_ms += (time.perf_counter() - inference_started) * 1000 diff --git a/tests/test_anthropic_pre_upstream_backpressure.py b/tests/test_anthropic_pre_upstream_backpressure.py index bbf47231e..d2e5fadb1 100644 --- a/tests/test_anthropic_pre_upstream_backpressure.py +++ b/tests/test_anthropic_pre_upstream_backpressure.py @@ -9,7 +9,7 @@ Covers: - N+1 contention (only the (N+1)th waiter records ``pre_upstream_wait`` > 0) - strict serialization under concurrency=1 - unbounded mode (``anthropic_pre_upstream_concurrency=0`` -> no semaphore) -- acquire timeout fails fast with ``503`` + ``Retry-After`` +- acquire timeout fails fast with passthrough compression skip - memory-context timeout fails open without leaking the semaphore - exception-safety (semaphore released when the critical section raises) - ``/livez`` unaffected under Anthropic backpressure @@ -520,11 +520,13 @@ def test_exception_inside_critical_section_releases_semaphore(): anyio.run(_run) -def test_acquire_timeout_returns_503_with_retry_after(stage_log_capture): +def test_acquire_timeout_degrades_to_passthrough(stage_log_capture): async def _run() -> None: sem = asyncio.Semaphore(1) await sem.acquire() handler = _DummyAnthropicHandler(anthropic_pre_upstream_sem=sem) + handler.config.optimize = True + handler.anthropic_pipeline = SimpleNamespace(apply=MagicMock()) handler.config.anthropic_pre_upstream_acquire_timeout_seconds = 0.01 req = _build_request( { @@ -535,11 +537,14 @@ def test_acquire_timeout_returns_503_with_retry_after(stage_log_capture): ) try: response = await handler.handle_anthropic_messages(req) - assert response.status_code == 503 - assert response.headers["retry-after"] == "1" + assert response.status_code == 200 body = json.loads(response.body) - assert body["error"]["type"] == "service_unavailable" - assert sem._value == 0 + assert body["id"] == "msg_test" + assert body["type"] == "message" + assert body["model"] == "claude-3-5-sonnet-latest" + assert body["stop_reason"] == "end_turn" + assert body["content"][0]["text"] == "ok" + assert not handler.anthropic_pipeline.apply.called finally: sem.release() assert sem._value == 1 @@ -549,7 +554,8 @@ def test_acquire_timeout_returns_503_with_retry_after(stage_log_capture): payloads = _parse_all_stage_logs(stage_log_capture) assert len(payloads) == 1 - assert payloads[0]["stages"]["pre_upstream_wait"] >= 10.0 + assert "pre_upstream_wait" in payloads[0]["stages"] + assert payloads[0]["stages"]["pre_upstream_wait"] >= 0.0 def test_memory_context_timeout_fails_open_and_releases_semaphore(): diff --git a/tests/test_kompress_request_nonblocking.py b/tests/test_kompress_request_nonblocking.py index 364e36a01..ab2503d53 100644 --- a/tests/test_kompress_request_nonblocking.py +++ b/tests/test_kompress_request_nonblocking.py @@ -11,6 +11,8 @@ in a background daemon thread instead. from __future__ import annotations +import threading + from headroom.transforms import kompress_compressor as kc from headroom.transforms.content_router import ContentRouter, ContentRouterConfig from headroom.transforms.kompress_compressor import KompressCompressor @@ -134,3 +136,154 @@ def test_router_compresses_cache_only_when_ready(monkeypatch): assert seen["allow_download"] is False # request path stays cache-only even when ready assert out == "kept words" + + +def test_saturation_fail_open_does_not_hang_request(monkeypatch): + """A saturated execution slot must fail open instead of blocking indefinitely.""" + + class _FakeEncoding(dict): + def __init__(self, word_count: int): + self._ids = list(range(word_count)) + super().__init__() + self["input_ids"] = [[1 for _ in range(word_count)]] + self["attention_mask"] = [[1 for _ in range(word_count)]] + + def word_ids(self, batch_index: int = 0): + return self._ids + + class _FakeModel: + def get_scores(self, input_ids, attention_mask): + return [[0.0 for _ in input_ids[0]]] + + class _FakeTokenizer: + def __call__(self, chunk_words, **kwargs): + return _FakeEncoding(len(chunk_words)) + + execution_semaphore = threading.BoundedSemaphore(1) + execution_semaphore.acquire() + + monkeypatch.setattr(kc, "_execution_semaphore", lambda *_a, **_k: execution_semaphore) + monkeypatch.setattr( + kc, + "_load_kompress", + lambda *args, **kwargs: (_FakeModel(), _FakeTokenizer(), "onnx"), + ) + monkeypatch.setenv("HEADROOM_KOMPRESS_EXECUTION_TIMEOUT_MS", "1") + + before = kc.get_kompress_execution_stats()["execution_timeout_skips_total"] + text = " ".join(["word"] * 40) + result_holder: dict[str, object] = {} + + def _run() -> None: + result_holder["result"] = KompressCompressor().compress(text, allow_download=False) + + worker = threading.Thread(target=_run) + worker.start() + worker.join(timeout=0.25) + try: + assert not worker.is_alive(), ( + "Kompress saturation path is blocking request progress; expected fail-open under pressure" + ) + assert "result" in result_holder + finally: + try: + execution_semaphore.release() + except ValueError: + pass + worker.join(timeout=1.0) + + result = result_holder["result"] + assert result.compressed == text + assert result.compression_ratio == 1.0 + after = kc.get_kompress_execution_stats()["execution_timeout_skips_total"] + assert after == before + 1 + + +def test_capacity_available_still_compresses(monkeypatch): + """When execution semaphore capacity is available, compression is still attempted.""" + + class _FakeEncoding(dict): + def __init__(self, word_count: int): + self._ids = list(range(word_count)) + self["input_ids"] = [[1 for _ in range(word_count)]] + self["attention_mask"] = [[1 for _ in range(word_count)]] + + def word_ids(self, batch_index: int = 0): + return self._ids + + class _FakeModel: + def get_scores(self, input_ids, attention_mask): + return [[1.0 if idx % 2 == 0 else 0.0 for idx in range(len(input_ids[0]))]] + + def get_keep_mask(self, input_ids, attention_mask): + return [[idx % 2 == 0 for idx in range(len(input_ids[0]))]] + + class _FakeTokenizer: + def __call__(self, chunk_words, **kwargs): + return _FakeEncoding(len(chunk_words)) + + monkeypatch.setattr( + kc, "_execution_semaphore", lambda *_args, **_kwargs: threading.BoundedSemaphore(1) + ) + monkeypatch.setattr( + kc, + "_load_kompress", + lambda *args, **kwargs: (_FakeModel(), _FakeTokenizer(), "onnx"), + ) + + result = KompressCompressor().compress(" ".join(["word"] * 20), allow_download=False) + assert 0 < result.compression_ratio < 1.0 + assert result.compressed != " ".join(["word"] * 20) + + +def test_validation_probe_waits_for_execution_slot(monkeypatch): + """Model-load validation must block for a slot instead of failing open.""" + + class _FakeTensor: + def to(self, _device): + return self + + class _FakeEncoding(dict): + def __init__(self): + super().__init__() + self["input_ids"] = _FakeTensor() + self["attention_mask"] = _FakeTensor() + + class _FakeTokenizer: + def __call__(self, *_args, **_kwargs): + return _FakeEncoding() + + class _FakeScore: + def detach(self): + return self + + def cpu(self): + return self + + class _FakeModel: + def __init__(self): + self.calls = 0 + + def get_scores(self, input_ids, attention_mask): + self.calls += 1 + return [_FakeScore()] + + semaphore = threading.BoundedSemaphore(1) + semaphore.acquire() + model = _FakeModel() + + monkeypatch.setattr(kc, "_execution_semaphore", lambda *_args, **_kwargs: semaphore) + + worker = threading.Thread( + target=kc._validate_pytorch_device, + args=(model, _FakeTokenizer(), "mps"), + ) + worker.start() + worker.join(timeout=0.05) + assert worker.is_alive(), "validation should wait for an execution slot" + + semaphore.release() + worker.join(timeout=1.0) + + assert not worker.is_alive() + assert model.calls == 1 diff --git a/tests/test_proxy_compression_executor.py b/tests/test_proxy_compression_executor.py index bd74f22ff..8e55c89dc 100644 --- a/tests/test_proxy_compression_executor.py +++ b/tests/test_proxy_compression_executor.py @@ -258,6 +258,67 @@ def test_timeout_before_worker_start_does_not_leak_in_flight() -> None: assert proxy._compression_leaked_threads == 0 +def test_compression_executor_skip_signal_remains_visible() -> None: + """A compression executor queue timeout increments visible runtime counters.""" + from fastapi.testclient import TestClient + + config = ProxyConfig( + optimize=False, + 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, + compression_max_workers=1, + ) + app = create_app(config) + proxy = app.state.proxy + + with TestClient(app) as client: + baseline = client.get("/health").json()["runtime"]["compression_executor"][ + "queue_timeouts_total" + ] + + first_started = threading.Event() + release_first = threading.Event() + + def _blocking_compression(): + first_started.set() + release_first.wait(timeout=5.0) + return "first" + + def _queued_compression(): + return "second" + + async def _drive(): + first_task = asyncio.create_task( + proxy._run_compression_in_executor(_blocking_compression, timeout=10.0) + ) + for _ in range(50): + if first_started.is_set(): + break + await asyncio.sleep(0.01) + assert first_started.is_set() + + with pytest.raises(asyncio.TimeoutError): + await proxy._run_compression_in_executor(_queued_compression, timeout=0.05) + + with proxy._compression_metrics_lock: + assert proxy._compression_queued == 0 + + release_first.set() + return await first_task + + asyncio.run(_drive()) + + with TestClient(app) as client: + after = client.get("/health").json()["runtime"]["compression_executor"] + assert after["queue_timeouts_total"] == baseline + 1 + + def test_compression_executor_metrics_appear_in_runtime_payload() -> None: """``/stats runtime.compression_executor`` surfaces the new gauges.""" from fastapi.testclient import TestClient