mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
3 commits
| Author | SHA1 | Message | Date | |
|---|---|---|---|---|
Tejas Chopra
|
f9807fd69e
|
feat(proxy): let extensions report cost savings and their own latency (#3051)
## What
Two changes that let a proxy extension report **what it saved** and
**what it cost**, so both show up under `/stats`, the dashboard, and
Prometheus.
`record_scope_savings` already existed and already accepted `usd` — the
one channel in the proxy that can express savings *without* tokens. Two
things stopped it working end to end.
### 1. Savings were silently dropped on Gemini traffic (bug)
`bind_scope` shares one attribution ledger between ASGI middleware and
the request handler. Anthropic and OpenAI call it; **Gemini never did**,
so anything an extension recorded into the request scope was discarded
for Gemini traffic only — silently, because an empty ledger and an
unbound one are indistinguishable at the outcome funnel. Now bound at
all four Gemini tag sites.
### 2. An extension's own latency was invisible (gap)
`overhead_ms` is measured *inside* the handler, and an ASGI extension
**wraps** that handler — so every millisecond it spends reaches the
client while every timing surface stays flat. An extension that halves
the bill and adds 200 ms per request is a trade the operator has to see
both halves of, and only one half was reaching the dashboard.
`record_scope_timing(scope, stage, ms)` is the symmetric counterpart to
`record_scope_savings`, carried on the same bound ledger and merged into
`RequestOutcome.pipeline_timing` at the outcome funnel — one place, so
every provider picks it up at once.
## API surface
```python
from headroom.proxy.savings_attribution import record_scope_savings, record_scope_timing
record_scope_savings(scope, "my_extension", tokens=0, usd=0.004) # money without tokens
record_scope_timing(scope, "my_extension", elapsed_ms)
```
Both take the ASGI `scope`, because middleware has no other way in.
Documented in `extensions.py` — the module extension authors actually
read, and the stability contract for this interface.
- Savings → `/stats` `savings.by_source`, dashboard card,
`headroom_savings_attributed_usd_total{source=...}`
- Timing → `/stats` `pipeline_timing`, dashboard Performance panel,
`headroom_transform_timing_ms_*`
**Attribution only.** These rows explain the headline total; they are
never added to it.
## Changes to existing behavior
- `public_tags` now strips `_headroom_stage_timing` as well as
`_headroom_savings_attribution`. Both ride on `tags` because that is the
one dict reaching the outcome funnel from every handler, and a list and
a dict must not land in a string-keyed label store.
- `pipeline_timing` passed to `metrics.record_request` is merged rather
than passed through **only when an extension contributed timings**; with
no extension the handler's own dict is passed through unchanged
(asserted by identity in the tests).
- Stage names are extension-supplied, so they are capped at 16 and
namespaced `ext:` — `deep_copy` reported by a plugin must never
accumulate into the same series as `deep_copy` measured by the pipeline.
A handler's own timing wins a collision (unreachable while the prefix
stands; the safe way round if it ever goes).
## Failure modes
Both calls are bounded (32 sources, 16 stages), never raise, and never
change a response — telemetry from a plugin must not be able to break
the request it is describing. Non-positive and non-numeric durations are
ignored: a zero is a clock artifact, not an observation, and averaging
it in would drag the mean down exactly where the stage is cheapest to
skip. `timings_from_tags` tolerates junk on the tag.
## Test-double fix
Three Gemini test fakes (`FakeRequest`, `_FakeRequest`,
`_VertexGeminiImageRequest`) had no `.scope`, which every real Starlette
`Request` has. They now do. This is a double that had drifted from the
type it stands in for; the alternative was weakening the handler to
tolerate a request shape that cannot occur in production.
---
## Real behavior proof
**Setup:** macOS 15.4 (darwin 25.4.0), Python 3.12.13, this branch at
`
|
||
|
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> |
||
|
Abhay Singh
|
f723925be7
|
fix(proxy/gemini): forward a non-JSON upstream body with its real status (#2174)
## Description
`handle_gemini_generate_content` turns a non-JSON upstream error
response into a generic 502, hiding the real status and body.
After the upstream call it extracts usage from `response.json()`:
```python
try:
resp_json = response.json()
usage = resp_json.get("usageMetadata", {})
...
cache_read_tokens = usage.get("cachedContentTokenCount", 0)
except (KeyError, TypeError, AttributeError) as e: # <-- missing JSONDecodeError / ValueError
...
```
`response.json()` raises `json.JSONDecodeError` (a `ValueError`
subclass) for a non-JSON body. That isn't in the tuple, so it escapes to
the function's outer `except Exception`, which returns a synthetic 502
and discards the real `response.status_code` / `response.content` (which
the success path forwards verbatim). An overloaded Google/Vertex/Copilot
frontend commonly returns a 503/500/429 with an HTML or empty body, so
the client sees a generic 502 instead of the true status — defeating
retry/backoff and dropping the diagnostic.
The all-non-text early-exit branch in the same handler already handles
this correctly with the full tuple (`except (json.JSONDecodeError,
ValueError, KeyError, TypeError, AttributeError)`) and then forwards the
real status/content.
## Fix
Add `json.JSONDecodeError, ValueError` to the token-extraction `except`,
matching that sibling. On a non-JSON body the extraction is skipped
(token metrics keep their fallbacks) and the handler falls through to
`return Response(content=response.content,
status_code=response.status_code, ...)` — the real status and body.
Closes #
## 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
- `headroom/proxy/handlers/gemini.py`: broaden the token-extraction
`except` in `handle_gemini_generate_content` to include
`json.JSONDecodeError, ValueError`.
- `tests/test_gemini_nonjson_status.py`: new test asserting that except
clause catches the JSON/ValueError family (guards against the
regression).
- `CHANGELOG.md`: Bug Fixes entry.
## Testing
- [ ] Unit tests pass (`pytest`)
- [x] Linting passes (`ruff check .`)
- [ ] Type checking passes (`mypy headroom`)
- [x] New tests added for new functionality
- [ ] Manual testing performed
### Test Output
```text
$ uvx ruff@0.15.17 check headroom/proxy/handlers/gemini.py tests/test_gemini_nonjson_status.py
All checks passed!
$ python -m py_compile headroom/proxy/handlers/gemini.py tests/test_gemini_nonjson_status.py
OK
```
## Real Behavior Proof
- Environment: Windows 11, Python 3.12, `uvx ruff@0.15.17`. A full
`pytest` OOM-kills this box (ML stack import), so I verified the
exception handling with a dependency-free script that models the
token-extraction try/except plus the handler's verbatim-forward return,
and left the full pytest to CI.
- Exact command / steps: sent a 503 response whose `.json()` raises
`JSONDecodeError` (non-JSON body) through the old tuple and the new
tuple, plus a normal JSON 200 as a control.
- Observed result: old lets `JSONDecodeError` escape (→ the outer
handler's synthetic 502); new catches it and forwards the real 503; the
JSON 200 still extracts tokens under both. The new test asserts the real
handler's except clause includes the JSON/ValueError family.
- Not tested: a live overloaded Gemini upstream; full local `pytest`
deferred to CI (OOM).
## 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
- [ ] New and existing unit tests pass locally with my changes
- [x] I have updated the CHANGELOG.md if applicable
## Additional Notes
The "unit tests pass locally" / "type checking" boxes are unchecked
because the full suite imports the ML stack, which I can't run here; the
change adds two exception types matching an existing, tested sibling
branch, verified by the standalone proof and a source-level regression
guard (a full handler-integration harness for Gemini doesn't exist
in-tree, and the existing gemini integration tests hit a live API).
---------
Co-authored-by: JerrettDavis <mxjerrett@gmail.com>
Co-authored-by: Tejas Chopra <chopratejas@gmail.com>
|