fix(proxy): cache litellm model resolution to stop repeated Provider List spam
## Description The proxy repeatedly prints LiteLLM's `Provider List: https://docs.litellm.ai/docs/providers` banner during normal operation, with no explanation or way to suppress it (#2851). Root cause: `_resolve_litellm_model()` in `headroom/proxy/savings_tracker.py` runs on every savings-tracking update (i.e. every request). For any model LiteLLM can't price (a custom/local/gateway model name — e.g. the reporter's local oMLX setup), the uncached fallback path calls `litellm.cost_per_token(...)` purely to probe resolvability. When that probe fails, LiteLLM prints the banner as an internal side effect before raising, and since the probe was never cached, it re-fires on every single request for the same unresolvable model. **Update:** review flagged that the first version of this fix cached into a plain, unbounded `dict` keyed by the (client-controlled) model name — a memory-retention path on a request-facing proxy, since a caller can grow it without limit by sending a new model string on every request. Replaced with a bounded `functools.lru_cache`; see Changes Made below. Closes #2851 ## 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/savings_tracker.py`: `_resolve_litellm_model()` is now decorated with `@lru_cache(maxsize=256)` instead of backing onto a hand-rolled unbounded `dict`. An evicted model name simply re-probes LiteLLM on next use — never a correctness issue, only whether the noisy failure banner reruns for that specific name. - `tests/conftest.py`: added a global `autouse` fixture, `_reset_litellm_model_resolution_cache`, that clears the cache before and after every test. It's process-lifetime and module-global, and several existing tests monkeypatch `savings_tracker.litellm` with different behavior per test while reusing common model names like `"gpt-4o"` — without a reset, whichever test resolves a name first silently wins that cache slot for the rest of the run and later tests stop exercising their own fake. - `tests/test_savings_tracker_litellm_resolution_cache.py` (new): regression tests for the three properties that actually matter — repeated resolution of one unknown model only probes LiteLLM once, resolving far more distinct names than the bound never grows the cache past it, and an evicted name is transparently re-probed rather than reusing a slot it no longer owns. - No behavior change for models LiteLLM can already price (fast path via `model_cost` lookup) — only the noisy uncached probe path is memoized, same as before. ## Testing - [x] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check .`) - [ ] Type checking passes (`mypy headroom`) — not run; `mypy` isn't installed in this environment - [x] New tests added for new functionality - [x] Manual testing performed ### Test Output ```text $ python3 -m pytest tests/test_proxy_savings_history.py tests/test_savings_tracker_zero_price.py \ tests/test_savings_tracker_litellm_resolution_cache.py -q tests/test_proxy_savings_history.py .................................... [ 73%] ... [ 79%] tests/test_savings_tracker_zero_price.py ....... [ 93%] tests/test_savings_tracker_litellm_resolution_cache.py ... [100%] 49 passed, 1 warning in 1.26s # Re-run in reversed file order to check for the exact order-dependence the # review flagged — same 49 passed, no failures either direction: $ python3 -m pytest tests/test_savings_tracker_litellm_resolution_cache.py \ tests/test_savings_tracker_zero_price.py tests/test_proxy_savings_history.py -q 49 passed, 1 warning in 1.11s $ python3 -m ruff check headroom/proxy/savings_tracker.py tests/conftest.py \ tests/test_savings_tracker_litellm_resolution_cache.py All checks passed! ``` ## Real Behavior Proof - Environment: macOS, Python 3.12.3, this repo checked out locally. - What changed since the last review pass: I got the compiled `headroom._core` Rust extension in hand (by installing the published `headroom-ai[all]` wheel into a separate venv and copying its `_core.abi3.so` next to this local source tree — same Python ABI, pure-Python edits in `savings_tracker.py` don't touch the compiled boundary). That unblocked the full test files this fix touches, including `tests/test_proxy_savings_history.py`, which was previously reported as untestable here. - Exact command / steps: three properties asserted directly against the real (now-bounded) cache in `tests/test_savings_tracker_litellm_resolution_cache.py`: 1. Resolve the same unresolvable model 5 times → assert the underlying `litellm.cost_per_token` probe fired exactly once. 2. Resolve `_MODEL_RESOLUTION_CACHE_MAXSIZE + 50` distinct model names → assert `_resolve_litellm_model.cache_info().currsize` stays at exactly `_MODEL_RESOLUTION_CACHE_MAXSIZE` (256), never higher — this is the actual memory-retention fix the review asked for. 3. Resolve one model, push exactly `maxsize` other distinct names through to evict it via LRU, then resolve it again → assert it re-probed (call count went 1 → 2), proving eviction is real and not just an untested cache_info number. - Observed result: all three pass; full affected-file suite (49 tests) passes in both forward and reversed run order, confirming the new `conftest.py` fixture actually fixes the cross-test leakage risk (verified by literally reordering the files, not just by inspection). - Not tested: a live HTTP request against a running `headroom proxy` process specifically re-exercising this bounded-cache commit — the earlier "20 simulated requests" proof against the previous (unbounded-dict) version of this fix was via a standalone script, not a real server; I have not repeated that specific live-server pass against this commit. The unit-level proof above exercises the exact same function (`_resolve_litellm_model`) the real proxy calls per-request from `headroom/proxy/server.py`, so I'm confident it generalizes, but flagging the gap rather than implying I re-ran it live. ## 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 — the bound/eviction rationale is commented above `_resolve_litellm_model`, and the cross-test leakage rationale is commented above the new `conftest.py` fixture - [ ] I have made corresponding changes to the documentation — N/A, internal implementation detail with no user-facing API/doc surface - [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` ## Additional Notes - `mypy` still hasn't been run — not installed in this sandbox, and I didn't want to widen the PR further by installing/configuring it just for this. Flagging rather than silently skipping. - The earlier "Additional Notes" gap about `test_proxy_savings_history.py` being untestable in this environment is resolved (see Real Behavior Proof) — it now runs and passes, including the pre-existing `test_litellm_resolution_and_savings_estimation_fallbacks` test that exercises `_resolve_litellm_model` with a mutated `model_cost` dict across several assertions in one test. - Deliberately did not also bound `headroom/pricing/litellm_pricing.py`'s sibling `_resolved_model_cache` — same shape of cache, arguably the same exposure — since it's outside this PR's diff and touching it wasn't asked for. Flagging in case a maintainer wants it as a fast follow-up rather than silently leaving it unmentioned. --------- Co-authored-by: connectsudhindra-gif <connectsudhindra-gif@users.noreply.github.com> Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com>
This commit is contained in:
@@ -16,6 +16,7 @@ import tempfile
|
||||
import threading
|
||||
from csv import DictWriter
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from functools import lru_cache
|
||||
from io import StringIO
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
@@ -164,6 +165,23 @@ def _normalize_model(value: Any) -> str:
|
||||
return cleaned or MODEL_UNKNOWN
|
||||
|
||||
|
||||
# `_resolve_litellm_model` is called on every savings-tracking update (i.e.
|
||||
# every request), and `model` is client-controlled — it comes straight off
|
||||
# the request body. For a model LiteLLM can't price (a custom / local /
|
||||
# gateway name), the uncached fallback below calls `litellm.cost_per_token`
|
||||
# purely to probe resolvability, which prints LiteLLM's noisy "Provider
|
||||
# List: https://docs.litellm.ai/docs/providers" banner on every failed probe
|
||||
# (#2851). Cache the resolution per model name so that probe runs at most
|
||||
# once per distinct model — bounded, not a plain dict: a request-facing
|
||||
# proxy must not let a caller grow an unbounded cache for free by sending a
|
||||
# fresh model string on every request. `maxsize` caps memory; LRU eviction
|
||||
# means a model that stops being sent eventually falls out and simply
|
||||
# re-probes if it's ever sent again — never a correctness issue, only
|
||||
# whether the probe (and its noisy failure banner) reruns.
|
||||
_MODEL_RESOLUTION_CACHE_MAXSIZE = 256
|
||||
|
||||
|
||||
@lru_cache(maxsize=_MODEL_RESOLUTION_CACHE_MAXSIZE)
|
||||
def _resolve_litellm_model(model: str) -> str:
|
||||
"""Resolve model name to one LiteLLM recognizes.
|
||||
|
||||
@@ -173,6 +191,12 @@ def _resolve_litellm_model(model: str) -> str:
|
||||
"claude-opus" identically to the live /stats path. Uses the shared result
|
||||
only when it maps to a priced model_cost key; otherwise falls through to the
|
||||
bare-prefix logic below. Fail-soft: pricing never breaks bookkeeping.
|
||||
|
||||
Bounded LRU cache, keyed by model name — see
|
||||
``_MODEL_RESOLUTION_CACHE_MAXSIZE`` above. Tests that mock the LiteLLM
|
||||
module across calls with the same model name must call
|
||||
``_resolve_litellm_model.cache_clear()`` between cases, or results from
|
||||
an earlier case leak in.
|
||||
"""
|
||||
litellm = _get_litellm_module()
|
||||
if litellm is None:
|
||||
|
||||
@@ -91,6 +91,29 @@ def _reset_copilot_routing_flag():
|
||||
reset_request_routed_to_copilot()
|
||||
|
||||
|
||||
# `savings_tracker._resolve_litellm_model` is an `lru_cache`d, module-global,
|
||||
# process-lifetime cache keyed by model name (bounded — see #2860). Many test
|
||||
# files monkeypatch `savings_tracker.litellm` to a fake with different
|
||||
# `model_cost`/`cost_per_token` behavior per test, but reuse common model
|
||||
# names like "gpt-4o" across them. Without a reset, whichever test resolves
|
||||
# "gpt-4o" first "wins" the cache entry for the rest of the run, and later
|
||||
# tests silently stop exercising their own fake — a real-not-hypothetical
|
||||
# order-dependence bug once the cache is process-lifetime instead of per-call.
|
||||
# Clear before AND after so a test's own within-test resolutions never leak
|
||||
# in from, or leak out to, a neighboring test either.
|
||||
@pytest.fixture(autouse=True)
|
||||
def _reset_litellm_model_resolution_cache():
|
||||
try:
|
||||
from headroom.proxy.savings_tracker import _resolve_litellm_model
|
||||
except ModuleNotFoundError:
|
||||
yield
|
||||
return
|
||||
|
||||
_resolve_litellm_model.cache_clear()
|
||||
yield
|
||||
_resolve_litellm_model.cache_clear()
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Global test hooks
|
||||
# =============================================================================
|
||||
|
||||
@@ -0,0 +1,88 @@
|
||||
"""Regression: `_resolve_litellm_model`'s cache must be bounded (PR #2860 review).
|
||||
|
||||
A plain unbounded dict cache keyed by a client-controlled model string is a
|
||||
memory-retention path on a request-facing proxy: a caller can grow it without
|
||||
limit by sending a new model name on every request. The fix uses a bounded
|
||||
`functools.lru_cache`. These tests pin the three properties that actually
|
||||
matter, independent of the litellm pricing behavior covered elsewhere:
|
||||
|
||||
- repeated resolution of the same unresolvable model only probes litellm once
|
||||
- the cache never grows past its bound, no matter how many distinct model
|
||||
names get resolved
|
||||
- an evicted name is transparently re-probed (never silently wrong or stuck)
|
||||
rather than growing the cache further
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import types
|
||||
|
||||
from headroom.proxy import savings_tracker as st
|
||||
|
||||
|
||||
def _fake_litellm_always_unresolvable(probe_calls: dict[str, int]) -> types.SimpleNamespace:
|
||||
"""A fake litellm where every model is unpriced and unresolvable.
|
||||
|
||||
`cost_per_token` always raises — exactly what a real custom/local model
|
||||
litellm has never heard of does — which is the call this cache exists to
|
||||
memoize (see the comment above `_resolve_litellm_model` in
|
||||
savings_tracker.py: that raise is also where real litellm prints its
|
||||
noisy "Provider List" banner, #2851).
|
||||
"""
|
||||
|
||||
def cost_per_token(*, model, prompt_tokens, completion_tokens):
|
||||
probe_calls[model] = probe_calls.get(model, 0) + 1
|
||||
raise RuntimeError("unknown model")
|
||||
|
||||
return types.SimpleNamespace(model_cost={}, cost_per_token=cost_per_token)
|
||||
|
||||
|
||||
def test_resolve_litellm_model_probes_unknown_model_once(monkeypatch):
|
||||
probe_calls: dict[str, int] = {}
|
||||
monkeypatch.setattr(
|
||||
st, "_get_litellm_module", lambda: _fake_litellm_always_unresolvable(probe_calls)
|
||||
)
|
||||
|
||||
for _ in range(5):
|
||||
resolved = st._resolve_litellm_model("widget-local-model")
|
||||
assert resolved == "widget-local-model"
|
||||
|
||||
assert probe_calls == {"widget-local-model": 1}
|
||||
|
||||
|
||||
def test_resolve_litellm_model_cache_is_bounded(monkeypatch):
|
||||
probe_calls: dict[str, int] = {}
|
||||
monkeypatch.setattr(
|
||||
st, "_get_litellm_module", lambda: _fake_litellm_always_unresolvable(probe_calls)
|
||||
)
|
||||
|
||||
extra_beyond_bound = 50
|
||||
for i in range(st._MODEL_RESOLUTION_CACHE_MAXSIZE + extra_beyond_bound):
|
||||
st._resolve_litellm_model(f"widget-local-model-{i}")
|
||||
|
||||
info = st._resolve_litellm_model.cache_info()
|
||||
assert info.maxsize == st._MODEL_RESOLUTION_CACHE_MAXSIZE
|
||||
# However many distinct names were resolved, the cache itself never
|
||||
# grows past its bound -- this is the actual memory-retention fix.
|
||||
assert info.currsize == st._MODEL_RESOLUTION_CACHE_MAXSIZE
|
||||
|
||||
|
||||
def test_resolve_litellm_model_evicted_name_reprobes(monkeypatch):
|
||||
probe_calls: dict[str, int] = {}
|
||||
monkeypatch.setattr(
|
||||
st, "_get_litellm_module", lambda: _fake_litellm_always_unresolvable(probe_calls)
|
||||
)
|
||||
|
||||
st._resolve_litellm_model("seed-model")
|
||||
assert probe_calls["seed-model"] == 1
|
||||
|
||||
# Push exactly `maxsize` new distinct names through without ever touching
|
||||
# "seed-model" again -- LRU eviction must push it out to make room.
|
||||
for i in range(st._MODEL_RESOLUTION_CACHE_MAXSIZE):
|
||||
st._resolve_litellm_model(f"filler-model-{i}")
|
||||
|
||||
# A resolvable name being evicted is not a correctness bug (it just
|
||||
# re-probes) -- the assertion that matters is that it *does* re-probe
|
||||
# rather than silently reusing a slot it no longer legitimately owns.
|
||||
st._resolve_litellm_model("seed-model")
|
||||
assert probe_calls["seed-model"] == 2
|
||||
Reference in New Issue
Block a user