Merge pull request #40350 from BerriAI/litellm_cost_map_background_retries

fix(cost-map): keep first fetch blocking, run retries in background
This commit is contained in:
kerry-berri 2026-09-09 09:59:07 -07:00 committed by GitHub
commit bfcc6404d3
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 240 additions and 114 deletions

View file

@ -546,7 +546,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
@ -2405,3 +2405,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()

View file

@ -13,6 +13,7 @@ import hashlib
import json
import os
import random
import threading
import time
from collections.abc import Awaitable, Callable
from dataclasses import dataclass, replace
@ -176,6 +177,11 @@ 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()
@dataclass(frozen=True, slots=True)
@ -314,12 +320,13 @@ async def _fetch_remote_model_cost_map_with_retry(
def _fetch_remote_model_cost_map_with_retry_sync(
url: str,
timeout: int,
max_attempts: int,
attempts: range,
sleep: Callable[[float], None],
rng: random.Random,
client: _SyncGetClient,
) -> ModelCostMapReloadResult:
for attempt in range(1, max_attempts + 1):
max_attempts: Final = attempts.stop - 1
for attempt in attempts:
outcome = _attempt_fetch_sync(client=client, url=url, timeout=timeout)
if not isinstance(outcome, _FetchAttemptRetryable):
return outcome
@ -520,6 +527,68 @@ def _finalize_loaded_model_cost_map(loaded: ModelCostMapReloaded) -> ModelCostMa
return replace(loaded, model_cost_map=_finalize_model_cost_map(loaded.model_cost_map))
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 _retry_remote_fetch_in_background(
url: str,
timeout: int,
max_attempts: int,
sleep: Callable[[float], None],
rng: random.Random,
client: _SyncGetClient,
first_outcome: _FetchAttemptRetryable,
) -> None:
try:
first_wait: Final = _next_retry_wait(outcome=first_outcome, attempt=1, max_attempts=max_attempts, rng=rng)
if isinstance(first_wait, ModelCostMapReloadUnavailable):
return
sleep(first_wait)
result: Final = _fetch_remote_model_cost_map_with_retry_sync(
url=url,
timeout=timeout,
attempts=range(2, max_attempts + 1),
sleep=sleep,
rng=rng,
client=client,
)
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
_litellm_import_complete.wait()
if not GetModelCostMap.validate_model_cost_map(
fetched_map=result.model_cost_map,
backup_model_count=GetModelCostMap._get_backup_model_count(), # pyright: ignore[reportPrivateUsage] # integrity cache
):
verbose_logger.warning(
"LiteLLM: Fetched model cost map failed integrity check. Using local backup instead. url=%s",
url,
)
return
finalized: Final = _finalize_loaded_model_cost_map(result).model_cost_map
_cost_map_source_info.source = "remote"
_cost_map_source_info.fallback_reason = None
_cost_map_source_info.loaded_at = datetime.now(timezone.utc)
adopt_model_cost_map(finalized)
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,
@ -532,9 +601,7 @@ def get_model_cost_map(
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``, retrying transient errors in a background thread.
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
@ -553,24 +620,34 @@ 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,
)
if isinstance(result, ModelCostMapReloadUnavailable):
fetch_client: Final = client if client is not None else httpx
fetch_rng: Final = rng if rng is not None else random.Random()
outcome: Final = _attempt_fetch_sync(client=fetch_client, url=url, timeout=timeout)
if isinstance(outcome, _FetchAttemptRetryable) and max_attempts > 1:
threading.Thread(
target=_retry_remote_fetch_in_background,
kwargs={ # mutable-ok: threading requires a mutable keyword-arguments mapping
"url": url,
"timeout": timeout,
"max_attempts": max_attempts,
"sleep": sleep,
"rng": fetch_rng,
"client": fetch_client,
"first_outcome": outcome,
},
name="litellm-model-cost-map-retry",
daemon=True,
).start()
if not isinstance(outcome, ModelCostMapReloaded):
verbose_logger.warning(
"LiteLLM: Failed to fetch remote model cost map from %s: %s. Falling back to local backup.",
url,
result.reason,
outcome.reason,
)
_cost_map_source_info.source = "local"
_cost_map_source_info.fallback_reason = f"Remote fetch failed: {result.reason}"
_cost_map_source_info.fallback_reason = f"Remote fetch failed: {outcome.reason}"
return _finalize_loaded_model_cost_map(GetModelCostMap.load_local_model_cost_map_with_revision()).model_cost_map
content: Final = result.model_cost_map
content: Final = outcome.model_cost_map
# Validate using cached count (cheap int comparison, no file I/O)
if not GetModelCostMap.validate_model_cost_map(
@ -587,4 +664,4 @@ def get_model_cost_map(
_cost_map_source_info.source = "remote"
_cost_map_source_info.fallback_reason = None
return _finalize_loaded_model_cost_map(result).model_cost_map
return _finalize_loaded_model_cost_map(outcome).model_cost_map

View file

@ -138,11 +138,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
@ -4436,20 +4432,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:

View file

@ -3079,7 +3079,7 @@ def register_model(
# Convert stringified numbers to appropriate numeric types
loaded_model_cost = model_cost
elif isinstance(model_cost, str):
loaded_model_cost = litellm.get_model_cost_map(url=model_cost)
loaded_model_cost = litellm.get_model_cost_map(url=model_cost, max_attempts=1)
if persist_across_reloads:
_registrations: Final[Mapping[str, Mapping[str, object]]] = loaded_model_cost

View file

@ -6,6 +6,7 @@ count actual model entries, not reserved meta keys) and the extraction of the
import json
import os
import threading
import pytest
@ -26,9 +27,7 @@ from litellm.litellm_core_utils.get_model_cost_map import (
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)
@ -44,9 +43,7 @@ def test_git_blob_id_is_what_git_hash_object_prints():
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():
@ -129,9 +126,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:
@ -317,9 +312,7 @@ def test_get_model_cost_map_stamps_loaded_at():
from litellm.litellm_core_utils import get_model_cost_map as module
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)
@ -328,6 +321,7 @@ def test_get_model_cost_map_stamps_loaded_at():
assert loaded_at is not None
assert before <= loaded_at <= datetime.now(timezone.utc)
# ---------------------------------------------------------------------------
# refetch_model_cost_map: retry/backoff behavior for runtime reloads
# ---------------------------------------------------------------------------
@ -394,9 +388,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
@ -408,9 +400,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
@ -430,9 +420,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]
@ -447,9 +435,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
@ -460,9 +446,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
@ -473,9 +457,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
@ -487,9 +469,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
@ -586,69 +566,125 @@ from litellm.litellm_core_utils.get_model_cost_map import (
class _SyncSleepRecorder:
"""Injected in place of time.sleep so the boot path's waits are asserted without delay."""
def __init__(self):
def __init__(self, block=False):
self.waits = []
self.block = block
self.started = threading.Event()
self.release = threading.Event()
def __call__(self, seconds: float) -> None:
if self.block:
self.started.set()
self.release.wait(timeout=10)
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."""
def _retry_threads():
return [thread for thread in threading.enumerate() if thread.name == "litellm-model-cost-map-retry"]
def test_boot_load_success_does_not_start_background_retry():
client, calls = _mock_client([httpx.Response(200, content=_real_map_bytes())], client_cls=httpx.Client)
sleeper = _SyncSleepRecorder()
cost_map = get_model_cost_map(
url=_URL,
sleep=sleeper,
rng=random.Random(0),
client=client,
)
assert calls["count"] == 1
assert sleeper.waits == []
assert _retry_threads() == []
assert cost_map.keys() >= _load_root_cost_map().keys() - {"sample_spec", FALLBACK_GENERALIZATIONS_KEY}
assert get_model_cost_map_source_info()["source"] == "remote"
def test_boot_load_transient_failure_returns_local_then_background_retry_adopts_remote(monkeypatch):
import litellm
from litellm import utils as litellm_utils
from litellm.litellm_core_utils import get_model_cost_map as module
original_model_cost = litellm.model_cost
monkeypatch.setattr(litellm, "model_cost", dict(original_model_cost))
for name, provider_models in tuple(vars(litellm).items()):
if name.endswith("_models") and isinstance(provider_models, set):
monkeypatch.setattr(litellm, name, set(provider_models))
monkeypatch.setattr(litellm, "models_by_provider", dict(litellm.models_by_provider))
monkeypatch.setattr(
litellm_utils,
"_runtime_registered_model_cost",
dict(litellm_utils._runtime_registered_model_cost),
)
source_info = module._cost_map_source_info
for name in ("source", "url", "is_env_forced", "fallback_reason", "loaded_at", "source_revision", "etag"):
monkeypatch.setattr(source_info, name, getattr(source_info, name))
remote_map = _load_root_cost_map()
remote_map["claude-remote-only-test"] = {"litellm_provider": "anthropic", "mode": "chat"}
client, calls = _mock_client(
[
httpx.ConnectError("connection refused"),
httpx.Response(503),
httpx.Response(200, content=_real_map_bytes()),
httpx.Response(200, content=json.dumps(remote_map).encode()),
],
client_cls=httpx.Client,
)
sleeper = _SyncSleepRecorder()
sleeper = _SyncSleepRecorder(block=True)
litellm.register_model({"my-runtime-model": {"litellm_provider": "custom", "max_input_tokens": 4321}})
cost_map = get_model_cost_map(url=_URL, sleep=sleeper, rng=random.Random(0), client=client)
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
source = get_model_cost_map_source_info()
assert source["source"] == "remote"
assert source["fallback_reason"] is None
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
cost_map = get_model_cost_map(
url=_URL,
max_attempts=3,
sleep=sleeper,
rng=random.Random(0),
client=client,
)
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
assert calls["count"] == 1
assert sleeper.waits == []
assert sleeper.started.wait(timeout=10)
threads = _retry_threads()
try:
assert len(threads) == 1
assert "claude-remote-only-test" not in cost_map
assert cost_map.keys() == _finalize_model_cost_map(GetModelCostMap.load_local_model_cost_map()).keys()
source = get_model_cost_map_source_info()
assert source["source"] == "local"
assert source["fallback_reason"].startswith("Remote fetch failed:")
sleeper.release.set()
for thread in threads:
thread.join(timeout=10)
assert all(not thread.is_alive() for thread in threads)
assert sleeper.waits and 2.0 <= sleeper.waits[0] < 3.0
assert calls["count"] == 2
assert "claude-remote-only-test" in litellm.model_cost
assert "claude-remote-only-test" in litellm.anthropic_models
assert "my-runtime-model" in litellm.model_cost
source = get_model_cost_map_source_info()
assert source["source"] == "remote"
assert source["fallback_reason"] is None
finally:
sleeper.release.set()
for thread in _retry_threads():
thread.join(timeout=10)
def test_boot_load_does_not_retry_permanent_failures():
"""A 404 or a malformed URL cannot heal by waiting: one attempt, no sleeps, backup."""
def test_boot_load_does_not_retry_non_retryable_failure():
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)
get_model_cost_map(
url=_URL,
sleep=sleeper,
rng=random.Random(0),
client=client,
)
assert calls["count"] == 1
assert sleeper.waits == []
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"
assert _retry_threads() == []
source = get_model_cost_map_source_info()
assert source["source"] == "local"
assert source["fallback_reason"] is not None
def test_boot_load_respects_local_env_override(monkeypatch):
@ -701,7 +737,9 @@ def test_boot_load_that_fails_the_integrity_check_reports_the_backup_not_the_rej
)
get_model_cost_map(url=_URL, sleep=_SyncSleepRecorder(), rng=random.Random(0), client=remote)
shrunk_body = b'{"gpt-5.4-mini": {"mode": "chat", "input_cost_per_token": 1e-06, "output_cost_per_token": 2e-06}}'
shrunk, _ = _mock_client([httpx.Response(200, headers={"ETag": 'W/"shrunk"'}, content=shrunk_body)], client_cls=httpx.Client)
shrunk, _ = _mock_client(
[httpx.Response(200, headers={"ETag": 'W/"shrunk"'}, content=shrunk_body)], client_cls=httpx.Client
)
get_model_cost_map(url=_URL, sleep=_SyncSleepRecorder(), rng=random.Random(0), client=shrunk)

View file

@ -2,10 +2,12 @@ import asyncio
import json
import logging
import os
import threading
from datetime import datetime, timedelta, timezone
from typing import Final
from unittest.mock import AsyncMock, MagicMock, patch
import httpx
import pytest
import respx
from jsonschema import validate
@ -2382,6 +2384,28 @@ def test_register_model_with_scientific_notation():
_invalidate_model_cost_lowercase_map()
@respx.mock
def test_register_model_url_fetch_uses_single_attempt(monkeypatch):
monkeypatch.delenv("LITELLM_LOCAL_MODEL_COST_MAP", raising=False)
monkeypatch.setattr(litellm, "model_cost", dict(litellm.model_cost))
before = dict(litellm.model_cost)
threads_before = {thread.name for thread in threading.enumerate()}
route = respx.get("https://example.invalid/custom_pricing.json").mock(
return_value=httpx.Response(503)
)
litellm.register_model(model_cost="https://example.invalid/custom_pricing.json")
threads_after = {thread.name for thread in threading.enumerate()}
assert route.call_count == 1
assert not (threads_after - threads_before) & {"litellm-model-cost-map-retry"}
assert not any(
thread.name == "litellm-model-cost-map-retry" and thread.is_alive()
for thread in threading.enumerate()
)
assert litellm.model_cost.keys() >= before.keys()
def test_register_model_openrouter_without_slash():
"""
Test that register_model handles openrouter models without '/' in the name.