diff --git a/litellm/__init__.py b/litellm/__init__.py index 42c0ea881fd..4e3754399da 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -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() diff --git a/litellm/litellm_core_utils/get_model_cost_map.py b/litellm/litellm_core_utils/get_model_cost_map.py index 9cba5db8ab7..4118cd3420e 100644 --- a/litellm/litellm_core_utils/get_model_cost_map.py +++ b/litellm/litellm_core_utils/get_model_cost_map.py @@ -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.", diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 1cd08fe27c0..efe38376353 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -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: diff --git a/tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py b/tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py index 18185126775..dce1a431f13 100644 --- a/tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py +++ b/tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py @@ -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):