mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
fix(proxy): offload OpenAI and Gemini tokenizer counting off the event loop (#2498)
## 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>
This commit is contained in:
inix
GitHub
parent
5d23a0aec2
commit
806d2e468a
7 changed files with 291 additions and 73 deletions
|
|
@ -72,46 +72,9 @@ class AnthropicHandlerMixin:
|
|||
"""Mixin providing Anthropic API handler methods for HeadroomProxy."""
|
||||
|
||||
async def _count_tokens_offloaded(self, model, messages): # noqa: ANN001, ANN201
|
||||
"""Resolve a tokenizer and count messages off the event loop.
|
||||
from headroom.proxy.token_counting import count_tokens_offloaded
|
||||
|
||||
Tokenizer resolution can be expensive on first use (HuggingFace
|
||||
backends may download vocab files) and counting a full Claude Code
|
||||
conversation is CPU-bound, so both run on the compression executor
|
||||
bounded by ``COMPRESSION_TIMEOUT_SECONDS`` (GH #1701: an unbounded
|
||||
on-loop load froze the whole server). On timeout or error this
|
||||
fails open to character-based estimation.
|
||||
|
||||
Returns:
|
||||
Tuple of ``(tokenizer, token_count)``. The tokenizer is fully
|
||||
initialized, so later ``count_messages`` calls on it are pure
|
||||
CPU work.
|
||||
"""
|
||||
from headroom.proxy.helpers import COMPRESSION_TIMEOUT_SECONDS
|
||||
from headroom.tokenizers import EstimatingTokenCounter, get_tokenizer
|
||||
|
||||
def _resolve_and_count(): # noqa: ANN202
|
||||
tokenizer = get_tokenizer(model)
|
||||
return tokenizer, tokenizer.count_messages(messages)
|
||||
|
||||
try:
|
||||
return await self._run_compression_in_executor(
|
||||
_resolve_and_count,
|
||||
timeout=float(COMPRESSION_TIMEOUT_SECONDS),
|
||||
)
|
||||
except Exception as e: # fail open — includes asyncio.TimeoutError
|
||||
# Log the downgrade once per model, not per request.
|
||||
fallback_models = getattr(self, "_token_count_fallback_models", None)
|
||||
if fallback_models is None:
|
||||
fallback_models = set()
|
||||
self._token_count_fallback_models = fallback_models
|
||||
if model not in fallback_models:
|
||||
fallback_models.add(model)
|
||||
logger.warning(
|
||||
f"Token counting for model {model} failed or timed out "
|
||||
f"({e.__class__.__name__}); falling back to estimation"
|
||||
)
|
||||
estimator = EstimatingTokenCounter()
|
||||
return estimator, estimator.count_messages(messages)
|
||||
return await count_tokens_offloaded(self, model, messages)
|
||||
|
||||
@staticmethod
|
||||
def _resolve_ccr_workspace(
|
||||
|
|
|
|||
|
|
@ -38,6 +38,16 @@ def _usage_int(value: Any, default: int = 0) -> int:
|
|||
class GeminiHandlerMixin:
|
||||
"""Mixin providing Gemini API handler methods for HeadroomProxy."""
|
||||
|
||||
async def _count_tokens_offloaded(self, model, messages): # noqa: ANN001, ANN201
|
||||
from headroom.proxy.token_counting import count_tokens_offloaded
|
||||
|
||||
return await count_tokens_offloaded(self, model, messages)
|
||||
|
||||
async def _count_texts_offloaded(self, model, texts): # noqa: ANN001, ANN201
|
||||
from headroom.proxy.token_counting import count_texts_offloaded
|
||||
|
||||
return await count_texts_offloaded(self, model, texts)
|
||||
|
||||
def _is_cloudcode_antigravity_request(
|
||||
self, body: dict[str, Any], headers: dict[str, str]
|
||||
) -> bool:
|
||||
|
|
@ -259,7 +269,6 @@ class GeminiHandlerMixin:
|
|||
from fastapi.responses import JSONResponse, Response
|
||||
|
||||
from headroom.proxy.helpers import MAX_REQUEST_BODY_SIZE, _read_request_json
|
||||
from headroom.tokenizers import get_tokenizer
|
||||
from headroom.utils import extract_user_query
|
||||
|
||||
start_time = time.time()
|
||||
|
|
@ -491,9 +500,8 @@ class GeminiHandlerMixin:
|
|||
headers=response_headers,
|
||||
)
|
||||
|
||||
# Token counting
|
||||
tokenizer = get_tokenizer(model)
|
||||
original_tokens = tokenizer.count_messages(messages)
|
||||
# Token counting (offloaded off the event loop — GH #1701)
|
||||
tokenizer, original_tokens = await self._count_tokens_offloaded(model, messages)
|
||||
|
||||
# Optimization
|
||||
transforms_applied: list[str] = []
|
||||
|
|
@ -816,7 +824,6 @@ class GeminiHandlerMixin:
|
|||
from fastapi.responses import JSONResponse
|
||||
|
||||
from headroom.proxy.helpers import _read_request_json
|
||||
from headroom.tokenizers import get_tokenizer
|
||||
from headroom.utils import extract_user_query
|
||||
|
||||
start_time = time.time()
|
||||
|
|
@ -880,8 +887,10 @@ class GeminiHandlerMixin:
|
|||
if isinstance(contents, list) and idx < len(contents)
|
||||
}
|
||||
|
||||
tokenizer = get_tokenizer(model)
|
||||
original_tokens = tokenizer.count_messages(messages) if messages else 0
|
||||
# Token counting (offloaded off the event loop — GH #1701)
|
||||
tokenizer, original_tokens = await self._count_tokens_offloaded(model, messages)
|
||||
if not messages:
|
||||
original_tokens = 0
|
||||
optimized_messages = messages
|
||||
optimized_tokens = original_tokens
|
||||
transforms_applied: list[str] = []
|
||||
|
|
@ -981,7 +990,6 @@ class GeminiHandlerMixin:
|
|||
from fastapi.responses import JSONResponse
|
||||
|
||||
from headroom.proxy.helpers import _read_request_json
|
||||
from headroom.tokenizers import get_tokenizer
|
||||
|
||||
start_time = time.time()
|
||||
request_id = await self._next_request_id()
|
||||
|
|
@ -1019,14 +1027,17 @@ class GeminiHandlerMixin:
|
|||
request_id=request_id,
|
||||
)
|
||||
|
||||
# Token counting
|
||||
tokenizer = get_tokenizer(model)
|
||||
original_tokens = 0
|
||||
for content in contents:
|
||||
parts = content.get("parts", [])
|
||||
for part in parts:
|
||||
if "text" in part:
|
||||
original_tokens += tokenizer.count_text(part["text"])
|
||||
# Token counting (offloaded off the event loop — GH #1701). Reuse the
|
||||
# shared _dict_parts coercion and keep only str text values: count_text
|
||||
# raises on a non-str part value and the fail-open path re-runs the same
|
||||
# input, so a malformed part would otherwise 500 the streaming request.
|
||||
text_parts = [
|
||||
part["text"]
|
||||
for content in (contents if isinstance(contents, list) else [])
|
||||
for part in self._dict_parts(content)
|
||||
if isinstance(part.get("text"), str)
|
||||
]
|
||||
_, original_tokens = await self._count_texts_offloaded(model, text_parts)
|
||||
|
||||
optimization_latency = (time.time() - start_time) * 1000
|
||||
|
||||
|
|
@ -1069,7 +1080,6 @@ class GeminiHandlerMixin:
|
|||
from fastapi.responses import JSONResponse, Response
|
||||
|
||||
from headroom.proxy.helpers import _read_request_json
|
||||
from headroom.tokenizers import get_tokenizer
|
||||
from headroom.utils import extract_user_query
|
||||
|
||||
start_time = time.time()
|
||||
|
|
@ -1140,9 +1150,8 @@ class GeminiHandlerMixin:
|
|||
headers=response_headers,
|
||||
)
|
||||
|
||||
# Token counting (original)
|
||||
tokenizer = get_tokenizer(model)
|
||||
original_tokens = tokenizer.count_messages(messages)
|
||||
# Token counting (original, offloaded off the event loop — GH #1701)
|
||||
tokenizer, original_tokens = await self._count_tokens_offloaded(model, messages)
|
||||
|
||||
# Apply compression using the same pipeline as generateContent
|
||||
transforms_applied: list[str] = []
|
||||
|
|
|
|||
|
|
@ -1316,6 +1316,11 @@ def _prefers_http1_passthrough(base_url: str) -> bool:
|
|||
class OpenAIHandlerMixin:
|
||||
"""Mixin providing OpenAI API handler methods for HeadroomProxy."""
|
||||
|
||||
async def _count_tokens_offloaded(self, model, messages): # noqa: ANN001, ANN201
|
||||
from headroom.proxy.token_counting import count_tokens_offloaded
|
||||
|
||||
return await count_tokens_offloaded(self, model, messages)
|
||||
|
||||
OPENAI_RESPONSES_ROUTER_MIN_BYTES = 512
|
||||
OPENAI_RESPONSES_OUTPUT_TYPES = _RESPONSES_OUTPUT_ITEM_TYPES
|
||||
|
||||
|
|
@ -2576,7 +2581,6 @@ class OpenAIHandlerMixin:
|
|||
_read_request_json,
|
||||
)
|
||||
from headroom.proxy.modes import is_cache_mode, is_token_mode
|
||||
from headroom.tokenizers import get_tokenizer
|
||||
from headroom.utils import extract_user_query
|
||||
|
||||
start_time = time.time()
|
||||
|
|
@ -2905,9 +2909,8 @@ class OpenAIHandlerMixin:
|
|||
|
||||
return Response(content=cached.response_body, headers=response_headers)
|
||||
|
||||
# Token counting
|
||||
tokenizer = get_tokenizer(model)
|
||||
original_tokens = tokenizer.count_messages(messages)
|
||||
# Token counting (offloaded off the event loop — GH #1701)
|
||||
tokenizer, original_tokens = await self._count_tokens_offloaded(model, messages)
|
||||
|
||||
# Hook: pre_compress
|
||||
_hook_biases = None
|
||||
|
|
@ -4257,7 +4260,6 @@ class OpenAIHandlerMixin:
|
|||
MAX_REQUEST_BODY_SIZE,
|
||||
read_request_json_with_bytes,
|
||||
)
|
||||
from headroom.tokenizers import get_tokenizer
|
||||
from headroom.utils import extract_user_query
|
||||
|
||||
start_time = time.time()
|
||||
|
|
@ -4474,9 +4476,8 @@ class OpenAIHandlerMixin:
|
|||
detail=f"Rate limited. Retry after {wait_seconds:.1f}s",
|
||||
)
|
||||
|
||||
# Token counting on converted messages
|
||||
tokenizer = get_tokenizer(model)
|
||||
original_tokens = tokenizer.count_messages(messages)
|
||||
# Token counting on converted messages (offloaded off the event loop — GH #1701)
|
||||
tokenizer, original_tokens = await self._count_tokens_offloaded(model, messages)
|
||||
|
||||
# Defaults below feed downstream telemetry and memory injection.
|
||||
# If optimization remains enabled, the Responses payload is compressed
|
||||
|
|
|
|||
79
headroom/proxy/token_counting.py
Normal file
79
headroom/proxy/token_counting.py
Normal file
|
|
@ -0,0 +1,79 @@
|
|||
"""Offloaded token-count helpers shared by proxy handlers.
|
||||
|
||||
Tokenizer resolution can be expensive on first use (HuggingFace backends may
|
||||
download vocab files) and counting a full Claude Code conversation is CPU-bound,
|
||||
so both run on the caller's compression executor bounded by
|
||||
``COMPRESSION_TIMEOUT_SECONDS`` (GH #1701: an unbounded on-loop load froze the
|
||||
whole server). On timeout, error, or a missing executor this fails open to
|
||||
character-based estimation.
|
||||
|
||||
Shared by every provider handler mixin (Anthropic, OpenAI, Gemini): the OpenAI
|
||||
``/v1/chat/completions`` and ``/v1/responses`` endpoints are multi-provider
|
||||
passthroughs, so an HF-routed model (qwen, deepseek, llama, ...) can reach them
|
||||
and trigger the same cold load.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from collections.abc import Callable
|
||||
from typing import Any, cast
|
||||
|
||||
logger = logging.getLogger("headroom.proxy")
|
||||
|
||||
|
||||
def _record_fallback_model(owner: Any, model: Any, message: str) -> None:
|
||||
fallback_models = getattr(owner, "_token_count_fallback_models", None)
|
||||
if fallback_models is None:
|
||||
fallback_models = set()
|
||||
owner._token_count_fallback_models = fallback_models
|
||||
if model not in fallback_models:
|
||||
fallback_models.add(model)
|
||||
logger.warning(message)
|
||||
|
||||
|
||||
async def _count_offloaded(owner: Any, model: Any, count: Callable[[Any], int]) -> tuple[Any, int]:
|
||||
"""Resolve a tokenizer and apply ``count`` off the event loop when possible.
|
||||
|
||||
``count`` maps a resolved tokenizer to a token total. Returns
|
||||
``(tokenizer, total)``; the returned tokenizer is fully initialized, so later
|
||||
counts on it are pure CPU work. Fails open to ``EstimatingTokenCounter`` when
|
||||
the owner has no compression executor, or on timeout/error.
|
||||
"""
|
||||
from headroom.proxy.helpers import COMPRESSION_TIMEOUT_SECONDS
|
||||
from headroom.tokenizers import EstimatingTokenCounter, get_tokenizer
|
||||
|
||||
runner = getattr(owner, "_run_compression_in_executor", None)
|
||||
if runner is None:
|
||||
estimator = EstimatingTokenCounter()
|
||||
return estimator, count(estimator)
|
||||
|
||||
def _resolve_and_count() -> tuple[Any, int]:
|
||||
tokenizer = get_tokenizer(model)
|
||||
return tokenizer, count(tokenizer)
|
||||
|
||||
try:
|
||||
result = await runner(_resolve_and_count, timeout=float(COMPRESSION_TIMEOUT_SECONDS))
|
||||
return cast(tuple[Any, int], result)
|
||||
except Exception as e: # fail open — includes asyncio.TimeoutError
|
||||
_record_fallback_model(
|
||||
owner,
|
||||
model,
|
||||
f"Token counting for model {model} failed or timed out "
|
||||
f"({e.__class__.__name__}); falling back to estimation",
|
||||
)
|
||||
estimator = EstimatingTokenCounter()
|
||||
return estimator, count(estimator)
|
||||
|
||||
|
||||
async def count_tokens_offloaded(owner: Any, model: Any, messages: Any) -> tuple[Any, int]:
|
||||
"""Resolve a tokenizer and count ``messages`` off the event loop when possible."""
|
||||
return await _count_offloaded(owner, model, lambda counter: counter.count_messages(messages))
|
||||
|
||||
|
||||
async def count_texts_offloaded(owner: Any, model: Any, texts: Any) -> tuple[Any, int]:
|
||||
"""Resolve a tokenizer and count text fragments off the event loop when possible."""
|
||||
text_list = list(texts)
|
||||
return await _count_offloaded(
|
||||
owner, model, lambda counter: sum(counter.count_text(text) for text in text_list)
|
||||
)
|
||||
|
|
@ -57,6 +57,14 @@ class _Handler(GeminiHandlerMixin):
|
|||
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(
|
||||
|
|
|
|||
|
|
@ -205,6 +205,14 @@ class _DummyOpenAIHandler(OpenAIHandlerMixin):
|
|||
# synchronously so MagicMock call_count assertions fire.
|
||||
return fn()
|
||||
|
||||
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)
|
||||
|
||||
async def _record_request_outcome(self, outcome) -> None:
|
||||
# Test stub: delegates to the production funnel so wire shape
|
||||
# matches HeadroomProxy._record_request_outcome.
|
||||
|
|
|
|||
|
|
@ -4,8 +4,10 @@ handler. For HF-backed models (e.g. deepseek-*) first use triggers an unbounded
|
|||
network download, freezing the whole server (610s request, then /livez, /readyz
|
||||
and /health hang until kill). The fix routes resolution + counting through
|
||||
HeadroomProxy._count_tokens_offloaded (compression executor, bounded by
|
||||
COMPRESSION_TIMEOUT_SECONDS, fail-open to estimation), and offloads the inline
|
||||
batch pipeline.apply() calls the same way.
|
||||
COMPRESSION_TIMEOUT_SECONDS, fail-open to estimation) — shared by every provider
|
||||
handler (Anthropic, OpenAI, Gemini), since the OpenAI passthrough endpoints
|
||||
receive the same HF-backed models — and offloads the inline batch
|
||||
pipeline.apply() calls the same way.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
|
@ -17,7 +19,18 @@ import time
|
|||
|
||||
from headroom.proxy.handlers.anthropic import AnthropicHandlerMixin
|
||||
from headroom.proxy.handlers.batch import BatchHandlerMixin
|
||||
from headroom.proxy.server import ProxyConfig, create_app
|
||||
from headroom.proxy.handlers.gemini import GeminiHandlerMixin
|
||||
from headroom.proxy.handlers.openai import OpenAIHandlerMixin
|
||||
from headroom.proxy.server import (
|
||||
CompressionQuarantinedError,
|
||||
ProxyConfig,
|
||||
create_app,
|
||||
)
|
||||
from headroom.proxy.token_counting import (
|
||||
_count_offloaded,
|
||||
count_texts_offloaded,
|
||||
count_tokens_offloaded,
|
||||
)
|
||||
from headroom.tokenizers import EstimatingTokenCounter
|
||||
|
||||
|
||||
|
|
@ -36,11 +49,36 @@ def _make_proxy(): # noqa: ANN202 — returns the internal HeadroomProxy
|
|||
def test_handlers_offload_token_counting_and_batch_apply() -> None:
|
||||
"""Wiring guard: the request paths must use the offloaded helpers, not inline
|
||||
get_tokenizer/count_messages or pipeline.apply on the event loop."""
|
||||
fn = AnthropicHandlerMixin.handle_anthropic_messages
|
||||
# Every provider handler that counts the original conversation must route
|
||||
# resolution + counting through the shared fail-open helper, never inline on
|
||||
# the loop. OpenAI /chat + /responses are multi-provider passthroughs, so an
|
||||
# HF-routed model (qwen, deepseek, llama, ...) can reach them and cold-load.
|
||||
for mixin, method in (
|
||||
(AnthropicHandlerMixin, "handle_anthropic_messages"),
|
||||
(OpenAIHandlerMixin, "handle_openai_chat"),
|
||||
(OpenAIHandlerMixin, "handle_openai_responses"),
|
||||
(GeminiHandlerMixin, "handle_gemini_generate_content"),
|
||||
(GeminiHandlerMixin, "handle_google_cloudcode_stream"),
|
||||
(GeminiHandlerMixin, "handle_gemini_count_tokens"),
|
||||
):
|
||||
fn = getattr(mixin, method)
|
||||
assert inspect.iscoroutinefunction(fn), f"{method} must be async"
|
||||
src = inspect.getsource(fn)
|
||||
assert "_count_tokens_offloaded(" in src, f"{method}: token counting not offloaded"
|
||||
assert "tokenizer = get_tokenizer(" not in src, (
|
||||
f"{method}: tokenizer resolved inline on the loop"
|
||||
)
|
||||
|
||||
fn = GeminiHandlerMixin.handle_gemini_stream_generate_content
|
||||
assert inspect.iscoroutinefunction(fn)
|
||||
src = inspect.getsource(fn)
|
||||
assert "_count_tokens_offloaded(" in src, "token counting not offloaded"
|
||||
assert "_count_texts_offloaded(" in src, "streaming Gemini text counting not offloaded"
|
||||
assert "tokenizer = get_tokenizer(" not in src, "tokenizer resolved inline on the loop"
|
||||
assert "count_text(" not in src, "streaming Gemini count_text still runs on the loop"
|
||||
assert "_dict_parts(" in src, "streaming Gemini must reuse the shared _dict_parts coercion"
|
||||
assert 'isinstance(part.get("text"), str)' in src, (
|
||||
"streaming Gemini must skip non-str text so count_text can't 500"
|
||||
)
|
||||
|
||||
for mixin, method in (
|
||||
(AnthropicHandlerMixin, "handle_anthropic_batch_create"),
|
||||
|
|
@ -54,7 +92,7 @@ def test_handlers_offload_token_counting_and_batch_apply() -> None:
|
|||
assert "_run_compression_in_executor(" in src, f"{method}: apply() not offloaded"
|
||||
assert "COMPRESSION_TIMEOUT_SECONDS" in src, f"{method}: offload missing timeout"
|
||||
|
||||
helper_src = inspect.getsource(AnthropicHandlerMixin._count_tokens_offloaded)
|
||||
helper_src = inspect.getsource(_count_offloaded)
|
||||
assert "COMPRESSION_TIMEOUT_SECONDS" in helper_src
|
||||
assert "EstimatingTokenCounter" in helper_src, "helper must fail open to estimation"
|
||||
|
||||
|
|
@ -124,3 +162,115 @@ async def test_count_tokens_offloaded_fails_open(monkeypatch) -> None: # noqa:
|
|||
assert tokens > 0
|
||||
# Logged-once bookkeeping records the downgraded model.
|
||||
assert "deepseek-chat" in proxy._token_count_fallback_models
|
||||
|
||||
|
||||
async def test_count_tokens_offloaded_fails_open_on_executor_quarantine() -> None:
|
||||
"""Now that OpenAI/Gemini counting shares the compression executor, an
|
||||
unrelated request's compression timeout can quarantine it — the next
|
||||
``_run_compression_in_executor`` call raises ``CompressionQuarantinedError``
|
||||
immediately (process-wide state). A request that is only counting tokens
|
||||
must not 500 on that; it fails open to estimation like any other error."""
|
||||
# The executor's ``except Exception`` fail-open only catches the quarantine
|
||||
# error because it subclasses Exception — pin that contract.
|
||||
assert issubclass(CompressionQuarantinedError, Exception)
|
||||
|
||||
proxy = _make_proxy()
|
||||
# Record a concurrent compression as timed out so the real executor guard
|
||||
# quarantines the next call — no mock of the helper itself.
|
||||
proxy._compression_timed_out_in_flight = 1
|
||||
|
||||
tokenizer, tokens = await proxy._count_tokens_offloaded(
|
||||
"qwen2.5-coder", [{"role": "user", "content": "hello world"}]
|
||||
)
|
||||
|
||||
assert isinstance(tokenizer, EstimatingTokenCounter)
|
||||
assert tokens > 0
|
||||
assert "qwen2.5-coder" in proxy._token_count_fallback_models
|
||||
|
||||
|
||||
async def test_count_tokens_offloaded_returns_count_text_capable_tokenizer() -> None:
|
||||
"""The fail-open tokenizer should still support text counting for callers
|
||||
that need per-fragment accounting."""
|
||||
proxy = _make_proxy()
|
||||
# Quarantine forces the fail-open branch (an EstimatingTokenCounter).
|
||||
proxy._compression_timed_out_in_flight = 1
|
||||
|
||||
# The empty-messages count is intentionally discarded by that handler
|
||||
# (it sums text parts itself), so only the tokenizer matters here.
|
||||
tokenizer, _ = await proxy._count_tokens_offloaded("qwen2.5-coder", [])
|
||||
|
||||
assert isinstance(tokenizer, EstimatingTokenCounter)
|
||||
# The streaming handler's per-part loop must not raise on the fallback.
|
||||
assert tokenizer.count_text("hello world") > 0
|
||||
|
||||
|
||||
async def test_count_texts_offloaded_runs_on_worker_thread(monkeypatch) -> None: # noqa: ANN001
|
||||
proxy = _make_proxy()
|
||||
loop_thread = threading.current_thread().name
|
||||
seen: dict[str, str] = {}
|
||||
|
||||
class _SpyTokenizer(EstimatingTokenCounter):
|
||||
def count_text(self, text): # noqa: ANN001, ANN201
|
||||
seen["thread"] = threading.current_thread().name
|
||||
return super().count_text(text)
|
||||
|
||||
monkeypatch.setattr("headroom.tokenizers.get_tokenizer", lambda *a, **k: _SpyTokenizer())
|
||||
|
||||
_, tokens = await proxy._count_texts_offloaded("gemini-pro", ["hello", "world"])
|
||||
|
||||
assert tokens > 0
|
||||
assert seen["thread"].startswith("headroom-compress")
|
||||
assert seen["thread"] != loop_thread
|
||||
|
||||
|
||||
async def test_count_texts_offloaded_fails_open(monkeypatch) -> None: # noqa: ANN001
|
||||
"""The texts variant downgrades to estimation on a resolution error, the same
|
||||
as the messages variant (its fail-open branch was previously uncovered)."""
|
||||
proxy = _make_proxy()
|
||||
|
||||
def _boom(*a, **k): # noqa: ANN002, ANN003, ANN202
|
||||
raise RuntimeError("tokenizer backend exploded")
|
||||
|
||||
monkeypatch.setattr("headroom.tokenizers.get_tokenizer", _boom)
|
||||
|
||||
tokenizer, tokens = await proxy._count_texts_offloaded("deepseek-chat", ["hello", "world"])
|
||||
|
||||
assert isinstance(tokenizer, EstimatingTokenCounter)
|
||||
assert tokens > 0
|
||||
assert "deepseek-chat" in proxy._token_count_fallback_models
|
||||
|
||||
|
||||
async def test_count_offloaded_without_executor_estimates() -> None:
|
||||
"""An owner with no compression executor (a lightweight caller or test double)
|
||||
fails open to estimation inline instead of crashing on the missing runner."""
|
||||
|
||||
class _NoExecutorOwner:
|
||||
pass
|
||||
|
||||
owner = _NoExecutorOwner()
|
||||
|
||||
tok, n_msg = await count_tokens_offloaded(
|
||||
owner, "gpt-4", [{"role": "user", "content": "hello world"}]
|
||||
)
|
||||
assert isinstance(tok, EstimatingTokenCounter)
|
||||
assert n_msg > 0
|
||||
|
||||
tok2, n_txt = await count_texts_offloaded(owner, "gemini-pro", ["hello", "world"])
|
||||
assert isinstance(tok2, EstimatingTokenCounter)
|
||||
assert n_txt > 0
|
||||
|
||||
|
||||
async def test_count_texts_offloaded_sums_fragments(monkeypatch) -> None: # noqa: ANN001
|
||||
"""The streaming rewrite sums per-fragment counts, matching the old per-part
|
||||
count_text loop it replaced."""
|
||||
proxy = _make_proxy()
|
||||
monkeypatch.setattr(
|
||||
"headroom.tokenizers.get_tokenizer", lambda *a, **k: EstimatingTokenCounter()
|
||||
)
|
||||
fragments = ["hello", "world", "foo"]
|
||||
|
||||
_, total = await proxy._count_texts_offloaded("gemini-pro", fragments)
|
||||
|
||||
est = EstimatingTokenCounter()
|
||||
assert total == sum(est.count_text(f) for f in fragments)
|
||||
assert total > 0
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue