diff --git a/headroom/providers/anthropic.py b/headroom/providers/anthropic.py index b34878f2d..3e341bd5d 100644 --- a/headroom/providers/anthropic.py +++ b/headroom/providers/anthropic.py @@ -274,11 +274,18 @@ class AnthropicTokenCounter(TokenCounter): ) _FALLBACK_WARNING_SHOWN = True - # Load tiktoken as fallback + # Load tiktoken as fallback — bounded, so a stalled vocab download can't + # hang token counting inside a request (tiktoken's downloader has no + # network timeout); on timeout we estimate by characters instead (GH #956). try: - import tiktoken + from headroom.tokenizers.tiktoken_counter import ( + TiktokenLoadError, + load_encoding, + ) - self._encoding = tiktoken.get_encoding("cl100k_base") + self._encoding = load_encoding("cl100k_base") + except TiktokenLoadError: + self._encoding = None # count_text() falls back to a character estimate except ImportError: if not self._use_api: warnings.warn( diff --git a/headroom/tokenizers/registry.py b/headroom/tokenizers/registry.py index 797f0b642..a2579c0a0 100644 --- a/headroom/tokenizers/registry.py +++ b/headroom/tokenizers/registry.py @@ -287,10 +287,24 @@ class TokenizerRegistry: return "estimation" def _create_tiktoken(self, model: str) -> TokenCounter: - """Create tiktoken-based tokenizer.""" - try: - from .tiktoken_counter import TiktokenCounter + """Create tiktoken-based tokenizer. + Forces the (bounded) encoding load up front so a stalled vocab download + falls back to estimation instead of hanging later inside a request (GH #956). + """ + try: + from .tiktoken_counter import ( + TiktokenCounter, + TiktokenLoadError, + get_encoding_for_model, + load_encoding, + ) + + try: + load_encoding(get_encoding_for_model(model)) + except TiktokenLoadError as exc: + logger.warning("tiktoken unavailable (%s); using estimation.", exc) + return EstimatingTokenCounter() return TiktokenCounter(model) except ImportError: logger.warning("tiktoken not installed. Install with: pip install tiktoken") diff --git a/headroom/tokenizers/tiktoken_counter.py b/headroom/tokenizers/tiktoken_counter.py index a3ccab025..a22fd9f9d 100644 --- a/headroom/tokenizers/tiktoken_counter.py +++ b/headroom/tokenizers/tiktoken_counter.py @@ -10,11 +10,38 @@ It supports multiple encodings: from __future__ import annotations +import logging +import os +import threading from functools import lru_cache from typing import Any from .base import BaseTokenizer +logger = logging.getLogger(__name__) + + +class TiktokenLoadError(RuntimeError): + """Raised when a tiktoken encoding can't be loaded in time. + + tiktoken downloads its BPE vocab on first use via ``requests.get`` with no + timeout, so a stalled/firewalled connection can block indefinitely. We bound + that load and raise this instead, so callers fall back to estimation rather + than hanging the request (see GH #956). + """ + + +# Encoding names whose bounded load already timed out — don't block on them again. +_load_failed: set[str] = set() + + +def _load_timeout_seconds() -> float: + try: + return float(os.environ.get("HEADROOM_TIKTOKEN_LOAD_TIMEOUT_SECONDS", "10")) + except (TypeError, ValueError): + return 10.0 + + # Model to encoding mapping MODEL_TO_ENCODING = { # GPT-4o family (o200k_base) @@ -78,10 +105,54 @@ DEFAULT_ENCODING = "cl100k_base" @lru_cache(maxsize=8) def _get_encoding(encoding_name: str): - """Get tiktoken encoding, cached for performance.""" + """Get a tiktoken encoding, cached for performance. + + Bounded by ``HEADROOM_TIKTOKEN_LOAD_TIMEOUT_SECONDS`` (default 10s): tiktoken's + vocab download has no network timeout, so we run the load on a worker thread + and raise :class:`TiktokenLoadError` if it doesn't finish in time, letting + callers fall back to estimation rather than hang the request (GH #956). The + first timed-out encoding is remembered so later calls fail fast instead of + re-blocking on every request. + """ import tiktoken - return tiktoken.get_encoding(encoding_name) + if encoding_name in _load_failed: + raise TiktokenLoadError(f"tiktoken encoding {encoding_name!r} previously failed to load") + + box: dict[str, Any] = {} + + def _load() -> None: + try: + box["enc"] = tiktoken.get_encoding(encoding_name) + except BaseException as exc: # noqa: BLE001 - re-raised in the calling thread + box["err"] = exc + + worker = threading.Thread(target=_load, name=f"tiktoken-load-{encoding_name}", daemon=True) + worker.start() + worker.join(_load_timeout_seconds()) + + if worker.is_alive(): + _load_failed.add(encoding_name) + logger.warning( + "tiktoken encoding %r did not load within %.1fs (likely a stalled vocab " + "download); falling back to token estimation. Pre-populate TIKTOKEN_CACHE_DIR " + "or tune HEADROOM_TIKTOKEN_LOAD_TIMEOUT_SECONDS.", + encoding_name, + _load_timeout_seconds(), + ) + raise TiktokenLoadError(f"tiktoken encoding {encoding_name!r} load timed out") + if "err" in box: + raise box["err"] + return box["enc"] + + +def load_encoding(encoding_name: str) -> Any: + """Public, bounded tiktoken-encoding loader. + + Returns the tiktoken encoding, or raises :class:`TiktokenLoadError` if the + vocab can't be loaded within the timeout (see :func:`_get_encoding`, GH #956). + """ + return _get_encoding(encoding_name) def get_encoding_for_model(model: str) -> str: diff --git a/tests/test_tokenizers/test_tiktoken_load_timeout.py b/tests/test_tokenizers/test_tiktoken_load_timeout.py new file mode 100644 index 000000000..4e1970f31 --- /dev/null +++ b/tests/test_tokenizers/test_tiktoken_load_timeout.py @@ -0,0 +1,81 @@ +"""tiktoken vocab loading must be bounded (GH #956). + +tiktoken downloads its BPE vocab via ``requests.get`` with no timeout, so a +stalled/firewalled connection blocks indefinitely. The proxy calls this lazily +inside a request worker, so the only bound was the 30s compression timeout — +yielding "every request times out, 0 compression". The bounded loader caps the +wait and falls back to estimation instead. +""" + +from __future__ import annotations + +import time + +import pytest + +from headroom.tokenizers import tiktoken_counter as tc +from headroom.tokenizers.estimator import EstimatingTokenCounter +from headroom.tokenizers.registry import TokenizerRegistry + + +@pytest.fixture(autouse=True) +def _reset_encoding_state(): + tc._get_encoding.cache_clear() + tc._load_failed.clear() + yield + tc._get_encoding.cache_clear() + tc._load_failed.clear() + + +def _stalled_get_encoding(_name: str): + # Simulates tiktoken's unbounded network download stalling. + time.sleep(2.0) + return object() + + +def test_load_encoding_is_bounded_on_stall(monkeypatch: pytest.MonkeyPatch) -> None: + import tiktoken + + monkeypatch.setattr(tiktoken, "get_encoding", _stalled_get_encoding) + monkeypatch.setenv("HEADROOM_TIKTOKEN_LOAD_TIMEOUT_SECONDS", "0.2") + + start = time.perf_counter() + with pytest.raises(tc.TiktokenLoadError): + tc.load_encoding("stall-enc") + elapsed = time.perf_counter() - start + assert elapsed < 1.5, f"load was not bounded (took {elapsed:.2f}s vs the 2s stall)" + + +def test_failed_encoding_short_circuits(monkeypatch: pytest.MonkeyPatch) -> None: + import tiktoken + + monkeypatch.setattr(tiktoken, "get_encoding", _stalled_get_encoding) + monkeypatch.setenv("HEADROOM_TIKTOKEN_LOAD_TIMEOUT_SECONDS", "0.2") + + with pytest.raises(tc.TiktokenLoadError): + tc.load_encoding("stall-enc-2") + + # A second request must fail instantly via the _load_failed short-circuit, + # not wait out the timeout again (this is what makes it not "every request"). + start = time.perf_counter() + with pytest.raises(tc.TiktokenLoadError): + tc.load_encoding("stall-enc-2") + assert time.perf_counter() - start < 0.1 + + +def test_fast_load_returns_encoding(monkeypatch: pytest.MonkeyPatch) -> None: + import tiktoken + + sentinel = object() + monkeypatch.setattr(tiktoken, "get_encoding", lambda _name: sentinel) + assert tc.load_encoding("fast-enc") is sentinel + + +def test_registry_falls_back_to_estimator_on_stall(monkeypatch: pytest.MonkeyPatch) -> None: + import tiktoken + + monkeypatch.setattr(tiktoken, "get_encoding", _stalled_get_encoding) + monkeypatch.setenv("HEADROOM_TIKTOKEN_LOAD_TIMEOUT_SECONDS", "0.2") + + counter = TokenizerRegistry()._create_tiktoken("gpt-4") + assert isinstance(counter, EstimatingTokenCounter)