mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
3 commits
| Author | SHA1 | Message | Date | |
|---|---|---|---|---|
JD Davis
|
039cd2431a
|
fix(proxy): preserve merged session and quarantine contracts (#2943)
## Description Forward-fixes two integration contracts exposed while auditing the large August 12 merge batch on `main`. The Codex WebSocket request-ID hardening correctly gave every emitted dashboard/feed row a unique ID, but it also changed the human-readable `PERF` prefix from the stable WebSocket session ID to that per-emission ID. That broke operator correlation and the contract documented by the original merge. This PR separates storage identity from log correlation: rows remain unique, while `PERF` lines remain grouped under the session ID. The same audit found two tokenizer quarantine tests still modeling the pre-time-cap behavior. Timeout debt no longer activates quarantine after its deadline expires. The tests now establish a live deadline and therefore continue to exercise the intended fail-open branch without weakening the production guard. ## Changes - Add an optional `RequestOutcome.perf_request_id` correlation field, defaulting to the existing `request_id` behavior for all current callers. - Set that field to the stable session ID for both per-turn and residual Codex WebSocket emissions. - Strengthen the lifecycle regression test to prove the unique feed-row ID is not used as the `PERF` prefix. - Update tokenizer quarantine tests to model an active, time-capped quarantine. This is a forward fix; it does not revert the unique WebSocket request IDs or the time-capped quarantine behavior. ## Merge-batch audit context - Audited 48 squash merges from ` |
||
|
inix
|
806d2e468a
|
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> |
||
|
Parideboy
|
46d5d685d9
|
fix(proxy): bound HF tokenizer load and offload token counting off event loop (#1738)
## Description Fixes #1701. On Windows, `headroom proxy --anthropic-api-url https://api.deepseek.com/anthropic` froze: the first `/v1/messages` request took ~610s (`optimization_latency_ms=609972`) with only router/lifecycle markers, and afterwards the whole server was a zombie — `/livez`, `/readyz` and `/health` hung until the process was killed. `HEADROOM_DETECT_BACKEND=python` was already set, so this was not the #575/#845 native-detect deadlock. Root cause: DeepSeek model names route to the HuggingFace tokenizer backend (`MODEL_PATTERNS` in `headroom/tokenizers/registry.py`). `HuggingFaceTokenizer` loads lazily, so the registry's construction-time fallback never fires; the first `count_messages` calls `AutoTokenizer.from_pretrained(..., trust_remote_code=True)` — unbounded network downloads/retries — and this ran **synchronously inside the async Anthropic messages handler** (`get_tokenizer(model)` + `tokenizer.count_messages(messages)`), outside the 30s `_run_compression_in_executor` bound. huggingface_hub retry chains on a restricted network easily reach ~10 minutes, blocking the entire asyncio event loop; subsequent on-loop counting kept it pinned. tiktoken got a bounded eager load for the same bug class long ago (#956); the HF backend never did. ## Type of Change - [x] Bug fix (non-breaking change which fixes an issue) - [ ] New feature (non-breaking change which adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to change) - [ ] Documentation update - [ ] Refactoring (no functional changes) ## Changes Made - `headroom/tokenizers/huggingface.py`: `_load_tokenizer` now tries the local HF cache first (`local_files_only=True`, no network), then bounds the network load with `HEADROOM_HF_TOKENIZER_LOAD_TIMEOUT_SECS` (default 10s; `0` disables network loads) on a daemon thread. Timeouts/failures return `None` (cached by `lru_cache`, so the hub is probed at most once per process per tokenizer) and `count_messages` fails open to char-based estimation via the existing `_use_fallback()` path. - `headroom/proxy/handlers/anthropic.py`: new `AnthropicHandlerMixin._count_tokens_offloaded(model, messages)` runs `get_tokenizer` + `count_messages` on the compression executor bounded by `COMPRESSION_TIMEOUT_SECONDS`, failing open to `EstimatingTokenCounter` (downgrade logged once per model). Used in `handle_anthropic_messages` (the issue's hot path, both count sites) and `handle_anthropic_batch_create`; the batch path's inline `anthropic_pipeline.apply()` is now offloaded via `_run_compression_in_executor` (mirrors the #1612 image-compression offload). - `headroom/proxy/handlers/batch.py`: the two remaining inline `openai_pipeline.apply()` calls (`handle_google_batch_create`, `_compress_batch_jsonl`) are offloaded the same way; existing `except` blocks keep the pass-through fail-open semantics. - Tests: `tests/test_huggingface_tokenizer_timeout.py` (cache-first, bounded timeout, failure caching, timeout=0, fail-open estimation), `tests/test_tokenizer_count_offload.py` (wiring guards, runs on `headroom-compress` worker, event loop stays responsive during slow tokenizer work, fail-open), plus `_run_compression_in_executor` stub on the batch test double. ## Testing - [x] All existing tests pass - [x] Added new tests for the changes - [ ] Manual testing performed ``` $ python -m pytest tests/test_huggingface_tokenizer_timeout.py tests/test_tokenizer_count_offload.py tests/test_image_compression_offload.py tests/test_gemini_compression_offload.py tests/test_tokenizers tests/test_proxy_handlers_batch.py -q 50 passed $ ruff check . # No issues found $ ruff format --check . # 1043 files already formatted $ mypy headroom --ignore-missing-imports # 0 errors ``` ## Real Behavior Proof - Environment: Windows 11 Pro (10.0.26200), Python 3.13, local checkout of this branch with the Rust core built. - Exact command / steps: `python -m pytest tests/test_tokenizer_count_offload.py -q` — includes `test_count_tokens_offloaded_keeps_loop_responsive`, which reproduces the issue's mechanism: a tokenizer whose `count_messages` blocks (stand-in for the unbounded `AutoTokenizer.from_pretrained` network load) while an asyncio ticker measures event-loop liveness. Also `python -m pytest tests/test_huggingface_tokenizer_timeout.py -q` with a `from_pretrained` stub that sleeps 60s and `HEADROOM_HF_TOKENIZER_LOAD_TIMEOUT_SECS=0.2`. - Observed result: with the fix, the slow count runs on a `headroom-compress` worker thread and the loop keeps ticking (`ticks >= 5`; inline it yields ~0 — the zombie). The 60s-hung HF load unblocks at the 0.2s timeout, falls back to estimation, and the second call returns instantly (failure cached, no re-probe). All 10 new tests pass. - Not tested: live reproduction against `api.deepseek.com` from a network where HF hub downloads stall (the reporter's exact environment); actual HF vocab download timing on a healthy network. ## Review Readiness - [x] I have performed a self-review - [x] This PR is ready for human review Co-authored-by: Claude Fable 5 <noreply@anthropic.com> |