chopratejas--headroom/tests/test_huggingface_tokenizer_timeout.py

"""HF tokenizer loading must be bounded (GH #1701): AutoTokenizer.from_pretrained
performs unbounded network downloads/retries; called lazily from the proxy's request
path it blocked the event loop for ~10 minutes and zombified the server. The fix
tries the local HF cache first (local_files_only=True), bounds the network attempt
with HEADROOM_HF_TOKENIZER_LOAD_TIMEOUT_SECS on a daemon thread, and fails open to
estimation — caching the failure so the hub is probed at most once per process.
"""

from __future__ import annotations

import sys
import time
import types
from typing import Any

import pytest

from headroom.tokenizers import huggingface as hf_mod
from headroom.tokenizers.huggingface import HuggingFaceTokenizer, _load_tokenizer


@pytest.fixture(autouse=True)
def _fresh_cache():
    _load_tokenizer.cache_clear()
    yield
    _load_tokenizer.cache_clear()


def _install_fake_transformers(monkeypatch: pytest.MonkeyPatch, from_pretrained) -> None:
    fake = types.ModuleType("transformers")
    fake.AutoTokenizer = type(
        "AutoTokenizer", (), {"from_pretrained": staticmethod(from_pretrained)}
    )
    monkeypatch.setitem(sys.modules, "transformers", fake)


def test_local_cache_tried_before_network(monkeypatch: pytest.MonkeyPatch) -> None:
    calls: list[dict[str, Any]] = []

    def fake_from_pretrained(name: str, **kwargs: Any):
        calls.append(kwargs)
        if kwargs.get("local_files_only"):
            raise OSError("not in cache")
        return "network-tokenizer"

    _install_fake_transformers(monkeypatch, fake_from_pretrained)
    monkeypatch.setenv("HEADROOM_HF_TOKENIZER_LOAD_TIMEOUT_SECS", "5")

    assert _load_tokenizer("some/model") == "network-tokenizer"
    assert calls[0].get("local_files_only") is True, "first attempt must be cache-only"
    assert not calls[1].get("local_files_only")


def test_cache_hit_never_touches_network(monkeypatch: pytest.MonkeyPatch) -> None:
    calls: list[dict[str, Any]] = []

    def fake_from_pretrained(name: str, **kwargs: Any):
        calls.append(kwargs)
        return "cached-tokenizer"

    _install_fake_transformers(monkeypatch, fake_from_pretrained)

    assert _load_tokenizer("some/model") == "cached-tokenizer"
    assert len(calls) == 1
    assert calls[0].get("local_files_only") is True


def test_slow_network_load_times_out_and_fails_open(monkeypatch: pytest.MonkeyPatch) -> None:
    def fake_from_pretrained(name: str, **kwargs: Any):
        if kwargs.get("local_files_only"):
            raise OSError("not in cache")
        time.sleep(60)  # simulates hung huggingface_hub download
        return "never"

    _install_fake_transformers(monkeypatch, fake_from_pretrained)
    monkeypatch.setenv("HEADROOM_HF_TOKENIZER_LOAD_TIMEOUT_SECS", "0.2")

    start = time.monotonic()
    assert _load_tokenizer("slow/model") is None
    assert time.monotonic() - start < 5, "load must unblock at the timeout, not the download"

    # Failure is cached (lru_cache) — the second call must not re-probe the hub.
    start = time.monotonic()
    assert _load_tokenizer("slow/model") is None
    assert time.monotonic() - start < 0.05


def test_timeout_zero_disables_network_loading(monkeypatch: pytest.MonkeyPatch) -> None:
    def fake_from_pretrained(name: str, **kwargs: Any):
        if kwargs.get("local_files_only"):
            raise OSError("not in cache")
        raise AssertionError("network load attempted despite timeout=0")

    _install_fake_transformers(monkeypatch, fake_from_pretrained)
    monkeypatch.setenv("HEADROOM_HF_TOKENIZER_LOAD_TIMEOUT_SECS", "0")

    assert _load_tokenizer("offline/model") is None


def test_count_messages_fails_open_to_estimation(monkeypatch: pytest.MonkeyPatch) -> None:
    def fake_from_pretrained(name: str, **kwargs: Any):
        raise OSError("unavailable")

    _install_fake_transformers(monkeypatch, fake_from_pretrained)
    monkeypatch.setenv("HEADROOM_HF_TOKENIZER_LOAD_TIMEOUT_SECS", "0.2")

    counter = HuggingFaceTokenizer("deepseek-chat")
    tokens = counter.count_messages([{"role": "user", "content": "hello world" * 50}])
    assert tokens > 0  # estimation fallback, no exception, no hang


def test_invalid_timeout_env_falls_back_to_default(monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.setenv("HEADROOM_HF_TOKENIZER_LOAD_TIMEOUT_SECS", "not-a-number")
    assert hf_mod._load_timeout_secs() == hf_mod._LOAD_TIMEOUT_DEFAULT