scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-10 22:41:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 129

fix(cost-map): keep first fetch blocking, run retries in background

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
kerry 2026-09-09 02:05:08 +00:00
parent eeb7732fc1
commit 0086b62b45
4 changed files with 296 additions and 98 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
litellm/__init__.py
View file

@ -541,7 +541,7 @@ _key_management_system: Optional["KeyManagementSystem"] = None
#### PII MASKING ####
output_parse_pii: bool = False
#############################################
from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map
from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map, mark_litellm_import_complete
model_cost = get_model_cost_map(url=model_cost_map_url)
cost_discount_config: Dict[str, float] = {} # Provider-specific cost discounts {"vertex_ai": 0.05} = 5% discount
@ -2397,3 +2397,5 @@ def __getattr__(name: str) -> Any:
# ALL_LITELLM_RESPONSE_TYPES is lazy-loaded via __getattr__ to avoid loading utils at import time
mark_litellm_import_complete()

132
litellm/litellm_core_utils/get_model_cost_map.py
View file

@ -12,6 +12,7 @@ import asyncio
import json
import os
import random
import threading
import time
from collections.abc import Awaitable, Callable
from dataclasses import dataclass
@ -159,6 +160,15 @@ class GetModelCostMap:
RETRYABLE_FETCH_STATUS_CODES: Final = frozenset({429, 500, 502, 503, 504})
MODEL_COST_MAP_FETCH_MAX_ATTEMPTS: Final = 3
MODEL_COST_MAP_FETCH_MAX_WAIT_SECONDS: Final = 30.0
_litellm_import_complete = threading.Event()
def mark_litellm_import_complete() -> None:
_litellm_import_complete.set()
def _start_daemon_thread(fn: Callable[[], None]) -> None:
threading.Thread(target=fn, name="litellm-model-cost-map-retry", daemon=True).start()
@dataclass(frozen=True, slots=True)
@ -297,9 +307,15 @@ def _fetch_remote_model_cost_map_with_retry_sync(
sleep: Callable[[float], None],
rng: random.Random,
client: _SyncGetClient,
starting_attempt: int = 1,
initial_outcome: _FetchAttemptRetryable | None = None,
) -> ModelCostMapReloadResult:
for attempt in range(1, max_attempts + 1):
outcome = _attempt_fetch_sync(client=client, url=url, timeout=timeout)
for attempt in range(starting_attempt, max_attempts + 1):
outcome = (
initial_outcome
if initial_outcome is not None and attempt == starting_attempt
else _attempt_fetch_sync(client=client, url=url, timeout=timeout)
)
if not isinstance(outcome, _FetchAttemptRetryable):
return outcome
wait_seconds = _next_retry_wait(outcome=outcome, attempt=attempt, max_attempts=max_attempts, rng=rng)
@ -464,6 +480,70 @@ def _finalize_model_cost_map(model_cost: dict) -> dict:
return _expand_model_aliases(model_cost)
def adopt_model_cost_map(
new_model_cost_map: dict, # mutable-ok: public API preserves the mutable cost-map contract
) -> int:
import litellm
from litellm import utils
litellm.model_cost = new_model_cost_map
utils._invalidate_model_cost_lowercase_map() # pyright: ignore[reportPrivateUsage] # required cache invalidation
litellm.add_known_models(model_cost_map=new_model_cost_map)
fetched_model_count: Final = len(new_model_cost_map) if new_model_cost_map else 0
utils.reapply_runtime_model_cost_registrations()
return fetched_model_count
def _continue_remote_fetch_in_background(
url: str,
timeout: int,
max_attempts: int,
sleep: Callable[[float], None],
rng: random.Random,
client: _SyncGetClient,
first_outcome: _FetchAttemptRetryable,
apply: Callable[[dict], object], # mutable-ok: injected callback receives the mutable cost-map dict
) -> None:
try:
result: Final = _fetch_remote_model_cost_map_with_retry_sync(
url=url,
timeout=timeout,
max_attempts=max_attempts,
sleep=sleep,
rng=rng,
client=client,
initial_outcome=first_outcome,
)
if isinstance(result, ModelCostMapReloadUnavailable):
verbose_logger.warning(
"LiteLLM: Failed to fetch remote model cost map from %s after %d attempts; keeping local backup",
url,
max_attempts,
)
return
backup_model_count: Final = GetModelCostMap._get_backup_model_count() # pyright: ignore[reportPrivateUsage] # integrity cache
if not GetModelCostMap.validate_model_cost_map(
fetched_map=result.model_cost_map,
backup_model_count=backup_model_count,
):
verbose_logger.warning(
"LiteLLM: Fetched model cost map failed integrity check. Keeping local backup. url=%s",
url,
)
_cost_map_source_info.fallback_reason = "Remote data failed integrity validation"
return
finalized_map: Final = _finalize_model_cost_map(result.model_cost_map)
_litellm_import_complete.wait()
apply(finalized_map)
_cost_map_source_info.source = "remote"
_cost_map_source_info.url = url
_cost_map_source_info.is_env_forced = False
_cost_map_source_info.fallback_reason = None
_cost_map_source_info.loaded_at = datetime.now(timezone.utc)
except Exception as e:
verbose_logger.warning("LiteLLM: Background model cost map retry failed: %s", e)
def get_model_cost_map(
url: str,
timeout: int = 5,
@ -471,14 +551,19 @@ def get_model_cost_map(
sleep: Callable[[float], None] = time.sleep,
rng: random.Random | None = None,
client: "_SyncGetClient | None" = None,
start_background: Callable[[Callable[[], None]], None] = _start_daemon_thread,
apply: Callable[ # mutable-ok: injected callback receives the mutable cost-map dict
[dict],
object,
] = adopt_model_cost_map,
) -> dict:
"""
Public entry point — returns the model cost map dict.
1. If ``LITELLM_LOCAL_MODEL_COST_MAP`` is set, uses the local backup only.
2. Otherwise fetches from ``url``, retrying transient HTTP errors
(429/5xx/transport) with Retry-After-aware backoff, validates
integrity, and falls back to the local backup on any failure.
2. Otherwise fetches from ``url``, validates the first response, and falls
back to the local backup while retrying transient HTTP errors in the
background.
Only the backup model count is cached (a single int) for validation.
The full backup dict is only parsed when it must be *returned* as a
@ -497,14 +582,35 @@ def get_model_cost_map(
_cost_map_source_info.url = url
_cost_map_source_info.is_env_forced = False
result: Final = _fetch_remote_model_cost_map_with_retry_sync(
url=url,
timeout=timeout,
max_attempts=max_attempts,
sleep=sleep,
rng=rng if rng is not None else random.Random(),
client=client if client is not None else httpx,
)
fetch_client: Final = client if client is not None else httpx
fetch_rng: Final = rng if rng is not None else random.Random()
first_outcome: Final = _attempt_fetch_sync(client=fetch_client, url=url, timeout=timeout)
if isinstance(first_outcome, _FetchAttemptRetryable):
verbose_logger.warning(
"LiteLLM: model cost map fetch attempt 1/%d failed: %s; "
"using local backup while retrying in the background",
max_attempts,
first_outcome.reason,
)
_cost_map_source_info.source = "local"
_cost_map_source_info.fallback_reason = f"Remote fetch failed: {first_outcome.reason}"
local_map: Final = _finalize_model_cost_map(GetModelCostMap.load_local_model_cost_map())
if max_attempts > 1:
start_background(
lambda: _continue_remote_fetch_in_background(
url=url,
timeout=timeout,
max_attempts=max_attempts,
sleep=sleep,
rng=fetch_rng,
client=fetch_client,
first_outcome=first_outcome,
apply=apply,
)
)
return local_map
result: Final = first_outcome
if isinstance(result, ModelCostMapReloadUnavailable):
verbose_logger.warning(
"LiteLLM: Failed to fetch remote model cost map from %s: %s. Falling back to local backup.",

23
litellm/proxy/proxy_server.py
View file

@ -132,11 +132,7 @@ from litellm.types.utils import (
TextCompletionResponse,
TokenCountResponse,
)
from litellm.utils import (
_invalidate_model_cost_lowercase_map,
load_credentials_from_list,
reapply_runtime_model_cost_registrations,
)
from litellm.utils import load_credentials_from_list
if TYPE_CHECKING:
from aiohttp import ClientSession
@ -4411,20 +4407,9 @@ def resolve_classifier_plugin(
def _swap_in_model_cost_map(new_model_cost_map: dict) -> int:
"""Adopt a freshly fetched cost map into this process's litellm state, return the model count"""
litellm.model_cost = new_model_cost_map
# Invalidate case-insensitive lookup map since model_cost was replaced
_invalidate_model_cost_lowercase_map()
# Repopulate provider model sets (e.g. litellm.anthropic_models) so that
# wildcard patterns like "anthropic/*" include any newly added models.
litellm.add_known_models(model_cost_map=new_model_cost_map)
# Counted before the re-apply below, which writes into this same dict, so the
# number reported describes the fetched price data alone.
fetched_model_count: Final = len(new_model_cost_map) if new_model_cost_map else 0
# The swap discards everything registered at runtime (deployment model_info,
# register_model overrides), so put it back on top of the fresh catalog.
reapply_runtime_model_cost_registrations()
return fetched_model_count
from litellm.litellm_core_utils.get_model_cost_map import adopt_model_cost_map
return adopt_model_cost_map(new_model_cost_map)
def should_load_db_object(object_type: str | SupportedDBObjectType) -> bool:

235
tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py
View file

@ -20,21 +20,18 @@ from litellm.litellm_core_utils.get_model_cost_map import (
GetModelCostMap,
_count_model_entries,
_finalize_model_cost_map,
adopt_model_cost_map,
)
def _load_root_cost_map() -> dict:
path = os.path.join(
os.path.dirname(__file__), "../../../model_prices_and_context_window.json"
)
path = os.path.join(os.path.dirname(__file__), "../../../model_prices_and_context_window.json")
with open(path) as f:
return json.load(f)
def _make_models(n: int) -> dict:
return {
f"model-{i}": {"litellm_provider": "openai", "mode": "chat"} for i in range(n)
}
return {f"model-{i}": {"litellm_provider": "openai", "mode": "chat"} for i in range(n)}
def test_count_model_entries_excludes_reserved_keys():
@ -117,9 +114,7 @@ def test_finalize_pops_key_and_installs_rules():
def test_finalize_with_no_block_clears_rules():
previous = list(get_fallback_generalization_rules())
try:
set_fallback_generalizations(
[{"name": "stale", "pattern": r"^x", "model_info": {"a": 1}}]
)
set_fallback_generalizations([{"name": "stale", "pattern": r"^x", "model_info": {"a": 1}}])
_finalize_model_cost_map(_make_models(2))
assert match_capability_generalizations("x-1") is None
finally:
@ -306,9 +301,7 @@ def test_get_model_cost_map_stamps_loaded_at(monkeypatch):
from litellm.litellm_core_utils import get_model_cost_map as module
monkeypatch.setattr(module._cost_map_source_info, "loaded_at", None)
client, _calls = _mock_client(
[httpx.Response(200, content=_real_map_bytes())], client_cls=httpx.Client
)
client, _calls = _mock_client([httpx.Response(200, content=_real_map_bytes())], client_cls=httpx.Client)
before = datetime.now(timezone.utc)
module.get_model_cost_map(url="https://example.invalid/cost_map.json", client=client)
@ -317,6 +310,7 @@ def test_get_model_cost_map_stamps_loaded_at(monkeypatch):
assert loaded_at is not None
assert before <= loaded_at <= datetime.now(timezone.utc)
# ---------------------------------------------------------------------------
# refetch_model_cost_map: retry/backoff behavior for runtime reloads
# ---------------------------------------------------------------------------
@ -382,9 +376,7 @@ async def test_refetch_retries_429_honoring_retry_after():
]
)
sleeper = _SleepRecorder()
result = await refetch_model_cost_map(
url=_URL, sleep=sleeper, rng=random.Random(0), client=client
)
result = await refetch_model_cost_map(url=_URL, sleep=sleeper, rng=random.Random(0), client=client)
assert isinstance(result, ModelCostMapReloaded)
assert len(result.model_cost_map) > 100
assert calls["count"] == 3
@ -396,9 +388,7 @@ async def test_refetch_gives_up_after_max_attempts_with_exponential_backoff():
"""All 429 without Retry-After: exponential backoff waits, then a failure value."""
client, calls = _mock_client([httpx.Response(429)])
sleeper = _SleepRecorder()
result = await refetch_model_cost_map(
url=_URL, sleep=sleeper, rng=random.Random(0), client=client
)
result = await refetch_model_cost_map(url=_URL, sleep=sleeper, rng=random.Random(0), client=client)
assert isinstance(result, ModelCostMapReloadUnavailable)
assert "429" in result.reason
assert "after 3 attempts" in result.reason
@ -418,9 +408,7 @@ async def test_refetch_caps_retry_after_wait():
]
)
sleeper = _SleepRecorder()
result = await refetch_model_cost_map(
url=_URL, sleep=sleeper, rng=random.Random(0), client=client
)
result = await refetch_model_cost_map(url=_URL, sleep=sleeper, rng=random.Random(0), client=client)
assert isinstance(result, ModelCostMapReloaded)
assert sleeper.waits == [30.0]
@ -435,9 +423,7 @@ async def test_refetch_retries_transport_errors():
]
)
sleeper = _SleepRecorder()
result = await refetch_model_cost_map(
url=_URL, sleep=sleeper, rng=random.Random(0), client=client
)
result = await refetch_model_cost_map(url=_URL, sleep=sleeper, rng=random.Random(0), client=client)
assert isinstance(result, ModelCostMapReloaded)
assert calls["count"] == 2
assert len(sleeper.waits) == 1
@ -448,9 +434,7 @@ async def test_refetch_non_retryable_status_fails_immediately():
"""A 404 is permanent: one attempt, no sleeps, failure value."""
client, calls = _mock_client([httpx.Response(404)])
sleeper = _SleepRecorder()
result = await refetch_model_cost_map(
url=_URL, sleep=sleeper, rng=random.Random(0), client=client
)
result = await refetch_model_cost_map(url=_URL, sleep=sleeper, rng=random.Random(0), client=client)
assert isinstance(result, ModelCostMapReloadUnavailable)
assert "404" in result.reason
assert calls["count"] == 1
@ -461,9 +445,7 @@ async def test_refetch_non_retryable_status_fails_immediately():
async def test_refetch_invalid_json_fails_immediately():
client, calls = _mock_client([httpx.Response(200, content=b"not json")])
sleeper = _SleepRecorder()
result = await refetch_model_cost_map(
url=_URL, sleep=sleeper, rng=random.Random(0), client=client
)
result = await refetch_model_cost_map(url=_URL, sleep=sleeper, rng=random.Random(0), client=client)
assert isinstance(result, ModelCostMapReloadUnavailable)
assert "invalid JSON" in result.reason
assert calls["count"] == 1
@ -475,9 +457,7 @@ async def test_refetch_shrunk_map_fails_integrity_not_swapped_in():
"""A drastically shrunk upstream file is rejected instead of being adopted."""
tiny = json.dumps(_make_models(60)).encode()
client, _calls = _mock_client([httpx.Response(200, content=tiny)])
result = await refetch_model_cost_map(
url=_URL, sleep=_SleepRecorder(), rng=random.Random(0), client=client
)
result = await refetch_model_cost_map(url=_URL, sleep=_SleepRecorder(), rng=random.Random(0), client=client)
assert isinstance(result, ModelCostMapReloadUnavailable)
assert "integrity validation" in result.reason
@ -520,62 +500,187 @@ class _SyncSleepRecorder:
self.waits.append(seconds)
def test_boot_load_retries_transient_failures_instead_of_falling_back():
"""A refused connection then a 503 at pod boot used to pin the process to the bundled
backup for its lifetime; both are transient and must be retried before giving up."""
class _BackgroundRecorder:
def __init__(self):
self.callbacks = []
def __call__(self, callback):
self.callbacks.append(callback)
def test_boot_load_returns_local_map_and_schedules_transient_retry():
client, calls = _mock_client(
[
httpx.ConnectError("connection refused"),
],
client_cls=httpx.Client,
)
sleeper = _SyncSleepRecorder()
background = _BackgroundRecorder()
cost_map = get_model_cost_map(
url=_URL,
sleep=sleeper,
rng=random.Random(0),
client=client,
start_background=background,
)
assert calls["count"] == 1
assert sleeper.waits == []
assert len(background.callbacks) == 1
assert len(cost_map) > 100
source = get_model_cost_map_source_info()
assert source["source"] == "local"
assert source["fallback_reason"] is not None
def test_background_retry_adopts_valid_remote_map():
client, calls = _mock_client(
[
httpx.ConnectError("connection refused"),
httpx.Response(503),
httpx.Response(200, content=_real_map_bytes()),
],
client_cls=httpx.Client,
)
sleeper = _SyncSleepRecorder()
background = _BackgroundRecorder()
applied = []
cost_map = get_model_cost_map(url=_URL, sleep=sleeper, rng=random.Random(0), client=client)
get_model_cost_map(
url=_URL,
sleep=sleeper,
rng=random.Random(0),
client=client,
start_background=background,
apply=applied.append,
)
assert calls["count"] == 1
assert sleeper.waits == []
assert len(background.callbacks) == 1
background.callbacks[0]()
assert calls["count"] == 2
assert len(sleeper.waits) == 1
assert 2.0 <= sleeper.waits[0] < 3.0
assert len(applied) == 1
assert applied[0].keys() >= _load_root_cost_map().keys() - {"sample_spec", FALLBACK_GENERALIZATIONS_KEY}
source = get_model_cost_map_source_info()
assert source["source"] == "remote"
assert source["fallback_reason"] is None
def test_background_retry_keeps_local_map_after_remaining_failures():
client, calls = _mock_client(
[httpx.ConnectError("connection refused"), httpx.Response(503)],
client_cls=httpx.Client,
)
sleeper = _SyncSleepRecorder()
background = _BackgroundRecorder()
applied = []
get_model_cost_map(
url=_URL,
sleep=sleeper,
rng=random.Random(0),
client=client,
start_background=background,
apply=applied.append,
)
background.callbacks[0]()
assert calls["count"] == 3
assert len(sleeper.waits) == 2
assert 2.0 <= sleeper.waits[0] < 3.0
assert 4.0 <= sleeper.waits[1] < 5.0
assert applied == []
assert get_model_cost_map_source_info()["source"] == "local"
def test_boot_load_does_not_schedule_non_retryable_failure():
client, calls = _mock_client([httpx.Response(404)], client_cls=httpx.Client)
sleeper = _SyncSleepRecorder()
background = _BackgroundRecorder()
get_model_cost_map(
url=_URL,
sleep=sleeper,
rng=random.Random(0),
client=client,
start_background=background,
)
assert calls["count"] == 1
assert sleeper.waits == []
assert background.callbacks == []
source = get_model_cost_map_source_info()
assert source["source"] == "remote"
assert source["fallback_reason"] is None
assert source["source"] == "local"
assert source["fallback_reason"] is not None
def test_boot_load_success_does_not_schedule_background_retry():
client, calls = _mock_client([httpx.Response(200, content=_real_map_bytes())], client_cls=httpx.Client)
sleeper = _SyncSleepRecorder()
background = _BackgroundRecorder()
cost_map = get_model_cost_map(
url=_URL,
sleep=sleeper,
rng=random.Random(0),
client=client,
start_background=background,
)
assert calls["count"] == 1
assert sleeper.waits == []
assert background.callbacks == []
assert get_model_cost_map_source_info()["source"] == "remote"
assert cost_map.keys() >= _load_root_cost_map().keys() - {"sample_spec", FALLBACK_GENERALIZATIONS_KEY}
def test_boot_load_honors_retry_after_then_falls_back_after_max_attempts():
"""An outage longer than the retry budget still ends on the bundled backup, and the
recorded fallback reason says how many attempts were spent so operators can tell."""
client, calls = _mock_client(
[httpx.Response(429, headers={"Retry-After": "7"})], client_cls=httpx.Client
def test_boot_load_with_one_attempt_does_not_schedule_background_retry():
client, calls = _mock_client([httpx.ConnectError("connection refused")], client_cls=httpx.Client)
sleeper = _SyncSleepRecorder()
background = _BackgroundRecorder()
get_model_cost_map(
url=_URL,
max_attempts=1,
sleep=sleeper,
rng=random.Random(0),
client=client,
start_background=background,
)
sleeper = _SyncSleepRecorder()
cost_map = get_model_cost_map(url=_URL, sleep=sleeper, rng=random.Random(0), client=client)
assert calls["count"] == 3
assert sleeper.waits == [7.0, 7.0]
source = get_model_cost_map_source_info()
assert source["source"] == "local"
assert "after 3 attempts" in source["fallback_reason"]
assert len(cost_map) > 100
def test_boot_load_does_not_retry_permanent_failures():
"""A 404 or a malformed URL cannot heal by waiting: one attempt, no sleeps, backup."""
client, calls = _mock_client([httpx.Response(404)], client_cls=httpx.Client)
sleeper = _SyncSleepRecorder()
get_model_cost_map(url=_URL, sleep=sleeper, rng=random.Random(0), client=client)
assert calls["count"] == 1
assert sleeper.waits == []
assert background.callbacks == []
assert get_model_cost_map_source_info()["source"] == "local"
get_model_cost_map(url="not a url", sleep=sleeper, rng=random.Random(0))
assert sleeper.waits == []
assert get_model_cost_map_source_info()["source"] == "local"
def test_adopt_model_cost_map_replays_runtime_registration_and_provider_models():
import litellm
from litellm import utils as litellm_utils
original_model_cost = litellm.model_cost
original_registry = dict(litellm_utils._runtime_registered_model_cost)
original_anthropic_models = set(litellm.anthropic_models)
try:
litellm.register_model(
model_cost={"custom/deployment-model": {"litellm_provider": "custom", "max_input_tokens": 4321}}
)
models_count = adopt_model_cost_map({"anthropic/new-model": {"litellm_provider": "anthropic", "mode": "chat"}})
assert models_count == 1
assert "anthropic/new-model" in litellm.anthropic_models
assert litellm.model_cost["custom/deployment-model"]["max_input_tokens"] == 4321
finally:
litellm.model_cost = original_model_cost # test-quality-ok: restore the module state changed by adoption
litellm_utils._runtime_registered_model_cost.clear()
litellm_utils._runtime_registered_model_cost.update(original_registry)
litellm.anthropic_models.clear()
litellm.anthropic_models.update(original_anthropic_models)
litellm_utils._invalidate_model_cost_lowercase_map()
def test_boot_load_respects_local_env_override(monkeypatch):