mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-10 14:27:00 -04:00
806d2e468a
## Description The OpenAI and Gemini handlers resolved the tokenizer and counted the conversation inline on the event loop. When a model resolves to a HuggingFace tokenizer (the registry routes qwen, deepseek, llama, phi, falcon, and more there) a cold cache runs `AutoTokenizer.from_pretrained` behind a 10s `thread.join`, which freezes the whole server. That is the GH #1701 stall, now reachable from OpenAI and Gemini because `/v1/chat/completions` and `/v1/responses` are documented multi-provider passthroughs and receive those models. Anthropic already routed the same call through a fail-open `_count_tokens_offloaded` helper. This hoists that helper to the shared `HeadroomProxy` base and sends the OpenAI and Gemini sites through it too. No linked issue. This is the OpenAI and Gemini follow-on to #1738, which offloaded the Anthropic and batch paths. GH #1701 is the original freeze report. ## 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 - Hoisted `_count_tokens_offloaded` from `AnthropicHandlerMixin` to the shared `HeadroomProxy` base, next to `_run_compression_in_executor`. It resolves and counts on the bounded compression executor and fails open to estimation on timeout, error, or executor quarantine. - Routed 6 inline sites through it: `handle_openai_chat`, `handle_openai_responses`, `handle_gemini_generate_content`, `handle_google_cloudcode_stream`, `handle_gemini_count_tokens`, and `handle_gemini_stream_generate_content` (resolve only, keeps its per-part `count_text` loop). - Removed 6 now-dead local `get_tokenizer` imports. - Left batch's per-line counts inline on purpose. They run on an already-warm tokenizer, so offloading them adds executor churn without touching the cold load. Batch's `pipeline.apply` was already offloaded in #1738. - Extended the wiring guard to all 7 provider handlers, added a quarantine fail-open test and a `count_text` fail-open test, and stubbed the method on 2 mixin-only handler doubles. ## Testing - [x] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check .`) - [ ] Type checking passes (`mypy headroom`) - [x] New tests added for new functionality - [x] Manual testing performed ### Test Output ```text $ ruff check headroom/proxy/server.py headroom/proxy/handlers/{anthropic,openai,gemini}.py tests/test_tokenizer_count_offload.py All checks passed! $ pytest tests/test_tokenizer_count_offload.py 6 passed in 4.39s # offload suite + all 26 handler-calling test files + handler-helpers/batch/tokenizers $ pytest tests/test_tokenizer_count_offload.py tests/test_openai_codex_routing.py tests/test_gemini_nonjson_status.py ... tests/test_tokenizers.py 377 passed, 15 skipped, 15 warnings in 86.60s (0:01:26) ``` ## Real Behavior Proof - Environment: macOS (Darwin), Python 3.13, proxy built from this branch. A synthetic tokenizer that sleeps 0.5s on resolve+count stands in for a cold HuggingFace `from_pretrained` load, with a 10ms asyncio loop-canary running alongside. - Exact command / steps: monkeypatch `headroom.tokenizers.get_tokenizer` to the 0.5s-sleeping tokenizer, then time a concurrent canary across two counts, the offloaded `await proxy._count_tokens_offloaded("qwen2.5-coder", messages)` and the old inline `get_tokenizer(model).count_messages(messages)`. - Observed result: the offloaded path kept the loop live at 41 canary ticks during the 509ms count, the inline path froze it to 0 ticks over 502ms, and both returned the same token count. Full run was 377 passed, 15 skipped, 0 failed. The new quarantine test confirms an unrelated compression timeout downgrades counting to estimation instead of raising a 500. - Not tested: live HuggingFace downloads and real qwen/deepseek traffic. No API keys in this environment, so the Gemini and OpenAI integration tests skip on `GEMINI_API_KEY`/`OPENAI_API_KEY`. `mypy headroom` did not finish locally (cold-times-out past 10 minutes on this box), so type-checking is left to CI. ## 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) ## Additional Notes - No linked issue. Follow-on to #1738. - Batch per-line counts stay inline: they run on an already-warm tokenizer, so offloading them adds executor churn without addressing the cold load. - Found a 6th site mid-implementation. `handle_gemini_stream_generate_content` also resolved the tokenizer inline but counts via a `count_text` loop, so it takes the resolve-only path. Verified `EstimatingTokenCounter.count_text` exists, so its fail-open branch does not crash. - `mypy headroom` cold-times-out locally (server.py pulls the full graph). Deferred to CI's Linux shards, same as prior PRs on this file. `ruff` and `pytest` run clean. - Documentation checkbox left unchecked: this change ships no user-facing doc update. --------- Co-authored-by: JerrettDavis <mxjerrett@gmail.com>
92 lines
3.1 KiB
Python
92 lines
3.1 KiB
Python
"""A non-JSON Gemini upstream body must not be masked as a synthetic 502."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
|
|
from headroom.proxy.handlers.gemini import GeminiHandlerMixin
|
|
|
|
|
|
class _FakeRequest:
|
|
def __init__(self) -> None:
|
|
self.headers: dict[str, str] = {}
|
|
self.query_params: dict[str, str] = {}
|
|
self.url = SimpleNamespace(path="/v1beta/models/gemini-pro:generateContent", query="")
|
|
|
|
|
|
class _NonJsonResponse:
|
|
status_code = 503
|
|
content = b"<html>temporarily unavailable</html>"
|
|
headers = {"content-type": "text/html", "content-length": str(len(content))}
|
|
|
|
def json(self) -> object:
|
|
raise json.JSONDecodeError("not json", self.content.decode("utf-8"), 0)
|
|
|
|
|
|
class _FakeMetrics:
|
|
def __init__(self) -> None:
|
|
self.failed: list[str] = []
|
|
|
|
async def record_failed(self, *, provider: str, model: str = "") -> None:
|
|
self.failed.append(f"{provider}:{model}")
|
|
|
|
|
|
class _Handler(GeminiHandlerMixin):
|
|
GEMINI_API_URL = "https://gemini.example"
|
|
|
|
def __init__(self) -> None:
|
|
self.memory_handler = None
|
|
self.rate_limiter = None
|
|
self.usage_reporter = None
|
|
self.config = SimpleNamespace(
|
|
optimize=False,
|
|
anthropic_pre_upstream_memory_context_timeout_seconds=0.1,
|
|
)
|
|
self.metrics = _FakeMetrics()
|
|
self.outcomes = []
|
|
|
|
async def _next_request_id(self) -> str:
|
|
return "req-1"
|
|
|
|
async def _retry_request(self, method, url, headers, body): # noqa: ANN001, ANN201
|
|
return _NonJsonResponse()
|
|
|
|
async def _record_request_outcome(self, outcome) -> None: # noqa: ANN001
|
|
self.outcomes.append(outcome)
|
|
|
|
async def _count_tokens_offloaded(self, model, messages): # noqa: ANN001, ANN201
|
|
# Test stub for HeadroomProxy._count_tokens_offloaded: resolve the
|
|
# tokenizer and count inline (the real method offloads to the executor).
|
|
from headroom.tokenizers import get_tokenizer
|
|
|
|
tokenizer = get_tokenizer(model)
|
|
return tokenizer, tokenizer.count_messages(messages)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_generate_content_forwards_non_json_upstream_status(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
async def payload(request): # noqa: ANN001, ANN201
|
|
return {"contents": [{"role": "user", "parts": [{"text": "hello"}]}]}
|
|
|
|
class _Tokenizer:
|
|
def count_messages(self, messages): # noqa: ANN001, ANN201
|
|
return 7
|
|
|
|
monkeypatch.setattr("headroom.proxy.helpers._read_request_json", payload)
|
|
monkeypatch.setattr("headroom.tokenizers.get_tokenizer", lambda model: _Tokenizer())
|
|
|
|
handler = _Handler()
|
|
response = await handler.handle_gemini_generate_content(_FakeRequest(), "gemini-pro")
|
|
|
|
assert response.status_code == 503
|
|
assert response.body == _NonJsonResponse.content
|
|
assert response.headers["content-type"] == "text/html"
|
|
assert response.headers["x-headroom-tokens-before"] == "7"
|
|
assert response.headers["x-headroom-tokens-after"] == "7"
|
|
assert handler.metrics.failed == []
|
|
assert handler.outcomes[0].status_code == 503
|