diff --git a/litellm/__init__.py b/litellm/__init__.py index 09585bda7bf9..9ca52642d257 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -741,7 +741,7 @@ def is_openai_finetune_model(key: str) -> bool: def _populate_provider_model_sets(model_cost_map: Dict) -> None: - for key, value in model_cost_map.items(): + for key, value in tuple(model_cost_map.items()): if value.get("litellm_provider") == "openai" and not is_openai_finetune_model(key): open_ai_chat_completion_models.add(key) elif value.get("litellm_provider") == "text-completion-openai": diff --git a/litellm/utils.py b/litellm/utils.py index 2c9200fbad77..4d27c3a24f63 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -3036,13 +3036,23 @@ def _convert_stringified_numbers(value): "ap-northeast-1.", ) -_CACHE_PRICING_FIELDS: Final = ( - "cache_creation_input_token_cost", - "cache_creation_input_token_cost_above_1hr", - "cache_creation_input_token_cost_above_200k_tokens", - "cache_read_input_token_cost", - "cache_read_input_token_cost_above_200k_tokens", -) +_TOKEN_COST_FIELDS: Final = ("input_cost_per_token", "output_cost_per_token") + + +def _without_synthesized_costs( + model_info: Mapping[str, object], key: str, value: Mapping[str, object] +) -> Mapping[str, object]: + """``_get_model_info_helper`` synthesizes ``input_cost_per_token`` / ``output_cost_per_token`` = 0 when the raw + entry has neither (#30198). Persisting those zeros would turn a name-priced entry into a free one and make + ``_is_cost_explicitly_configured`` disable budget enforcement on the next re-registration.""" + raw_fields: Final = frozenset(litellm.model_cost.get(model_info["key"], litellm.model_cost.get(key, ()))) + return MappingProxyType( + { + name: field_value + for name, field_value in model_info.items() + if name not in _TOKEN_COST_FIELDS or name in raw_fields or name in value + } + ) def _resolve_builtin_model_cost_entry(key: str, provider: str) -> dict[str, object] | None: @@ -3051,30 +3061,40 @@ def _resolve_builtin_model_cost_entry(key: str, provider: str) -> dict[str, obje like ``bedrock/bedrock/bedrock/us.anthropic.claude-sonnet-4-6`` or region aliases). - Returns a copy of the matching entry so the caller can inherit its defaults - (most importantly cache pricing) without mutating the shared built-in. - Returns ``None`` when no safe match exists. + The key itself is tried first, so a catalog row whose lookup tripped + inherits from itself rather than from a region-stripped sibling priced at + the base rate. Returns a copy of the matching entry so the caller can + inherit its defaults without mutating the shared built-in. Returns ``None`` + when no candidate exists under a provider ``get_model_info`` would accept + for ``provider``, or, when the value names none, for the provider prefix + the key itself carries, which is what ``get_model_info`` derives when it + reads the key back. """ - candidates: Final[list[str]] = [] + candidates: Final[list[str]] = [key] segments: Final = key.split("/") + match_provider: Final = provider or (segments[0] if len(segments) > 1 and segments[0] in LlmProvidersSet else "") idx = 0 while idx < len(segments) - 1 and segments[idx] in LlmProvidersSet: idx += 1 candidates.append("/".join(segments[idx:])) - base: Final = candidates[-1] if candidates else key + base: Final = candidates[-1] for region_prefix in _BEDROCK_REGION_PREFIXES: if base.startswith(region_prefix): candidates.append(base[len(region_prefix) :]) - if provider: - stripped: Final = _strip_model_name(model=base, custom_llm_provider=provider) + if match_provider: + stripped: Final = _strip_model_name(model=base, custom_llm_provider=match_provider) if stripped != base: candidates.append(stripped) for candidate in candidates: entry = litellm.model_cost.get(candidate) - if entry is not None and entry.get("litellm_provider") is not None: + if ( + entry is not None + and entry.get("litellm_provider") is not None + and _check_provider_match(entry, match_provider) + ): return dict(entry) return None @@ -3213,20 +3233,22 @@ def register_model( else: builtin_model_info = _get_builtin_model_info_for_registration(model=_key_str) if builtin_model_info is not None: - existing_model = cast(dict, builtin_model_info) - model_cost_key = existing_model["key"] + existing_model = dict( # mutable-ok: merge target + _without_synthesized_costs(builtin_model_info, key=_key_str, value=value) + ) + model_cost_key = builtin_model_info["key"] else: # An exact entry ends the lookup ladder before the capability rules are # consulted, so seed from them: otherwise registering an unmapped model # shadows the very defaults it would have resolved to unregistered. - existing_model = dict(match_capability_generalizations(_key_str) or {}) # mutable-ok: merge target model_cost_key = key builtin_entry = _resolve_builtin_model_cost_entry(key=_key_str, provider=provider) - if builtin_entry is not None: - for field in _CACHE_PRICING_FIELDS: - if value.get(field) is None and builtin_entry.get(field) is not None: - existing_model[field] = builtin_entry[field] - elif ( + existing_model = ( + builtin_entry + if builtin_entry is not None + else dict(match_capability_generalizations(_key_str) or {}) # mutable-ok: merge target + ) + if builtin_entry is None and ( value.get("cache_creation_input_token_cost") is None and value.get("cache_read_input_token_cost") is None and value.get("tiered_pricing") is None @@ -3245,21 +3267,6 @@ def register_model( # custom pricing on subsequent cost lookups. if existing_model.get("litellm_provider") is None: existing_model.pop("litellm_provider", None) - # Same pattern for cost fields (#30198): ``_get_model_info_helper`` - # synthesizes ``input_cost_per_token`` / ``output_cost_per_token`` - # = 0 when they are absent from the raw entry. Writing those zeros - # back flips a sparse entry from "no cost keys" (priced via name) - # to "cost keys = 0" (free), which makes - # ``_is_cost_explicitly_configured`` return True and silently - # disables budget enforcement on the next re-registration. - _raw_entry = litellm.model_cost.get(model_cost_key) - if _raw_entry is None: - _raw_entry = litellm.model_cost.get(key) - if _raw_entry is None: - _raw_entry = {} - for _cost_field in ("input_cost_per_token", "output_cost_per_token"): - if _cost_field not in _raw_entry and _cost_field not in value: - existing_model.pop(_cost_field, None) ## override / add new keys to the existing model cost dictionary updated_dictionary = _update_dictionary(existing_model, value) litellm.model_cost.setdefault(model_cost_key, {}).update(updated_dictionary) @@ -5412,7 +5419,7 @@ def _rebuild_model_cost_lowercase_map() -> dict[str, str]: The rebuilt map (guaranteed to be not None). """ global _model_cost_lowercase_map - _model_cost_lowercase_map = {k.lower(): k for k in litellm.model_cost} + _model_cost_lowercase_map = {k.lower(): k for k in tuple(litellm.model_cost)} return _model_cost_lowercase_map @@ -5427,9 +5434,7 @@ def _handle_stale_map_entry_rebuild( Returns: The matched key if found after rebuild, None otherwise. """ - global _model_cost_lowercase_map - _model_cost_lowercase_map = _rebuild_model_cost_lowercase_map() - matched_key: Final = _model_cost_lowercase_map.get(potential_key_lower) + matched_key: Final = _rebuild_model_cost_lowercase_map().get(potential_key_lower) if matched_key is not None and matched_key in litellm.model_cost: return matched_key return None @@ -5469,18 +5474,19 @@ def _get_model_cost_key(potential_key: str) -> str | None: called and confirmed not to cause performance issues, add it to the allowed_helpers list in: tests/code_coverage_tests/check_get_model_cost_key_performance.py """ - global _model_cost_lowercase_map - # Exact match (O(1)) if potential_key in litellm.model_cost: return potential_key - # Case-insensitive lookup via map (O(1)) - if _model_cost_lowercase_map is None: - _model_cost_lowercase_map = _rebuild_model_cost_lowercase_map() + # Case-insensitive lookup via map (O(1)); bound locally because another thread's + # registration can invalidate the module-level map at any point + cached_lowercase_map: Final = _model_cost_lowercase_map + lowercase_map: Final = ( + cached_lowercase_map if cached_lowercase_map is not None else _rebuild_model_cost_lowercase_map() + ) potential_key_lower: Final = potential_key.lower() - matched_key = _model_cost_lowercase_map.get(potential_key_lower) + matched_key = lowercase_map.get(potential_key_lower) # Verify key exists (O(1) - handles model_cost.pop() case) if matched_key is not None and matched_key in litellm.model_cost: diff --git a/tests/test_litellm/test_register_model_custom_pricing.py b/tests/test_litellm/test_register_model_custom_pricing.py index 452a15334efc..a7160afb0460 100644 --- a/tests/test_litellm/test_register_model_custom_pricing.py +++ b/tests/test_litellm/test_register_model_custom_pricing.py @@ -11,6 +11,7 @@ import copy import os +from typing import Final, NoReturn import pytest @@ -795,6 +796,128 @@ def test_embedding_direct_sdk_custom_pricing_still_registers_shared_key(): _invalidate_model_cost_lowercase_map() +def test_register_model_prices_a_geo_alias_from_its_builtin_when_the_lookup_fails(monkeypatch): + """After a price data reload the deployment replay registers every bedrock + geo deployment (``bedrock/au.anthropic.``) with metadata only. When + ``get_model_info`` failed for a reason other than the model being unmapped + (a request thread registering pricing tripped its map rebuild), the + registration fell through to the unmapped path and wrote a raw-key entry + with no token pricing, which shadowed the canonical entry at $0 for every + request until the next reload happened to win the race. + """ + from litellm.types.utils import ModelResponse, Usage + + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + _invalidate_model_cost_lowercase_map() + canonical_key: Final = "au.anthropic.claude-sonnet-5" + alias: Final = f"bedrock/{canonical_key}" + canonical: Final = litellm.model_cost[canonical_key] + + def lookup_tripped_by_a_concurrent_registration(model: str, custom_llm_provider: str | None = None) -> NoReturn: + raise RuntimeError("dictionary changed size during iteration") + + with pytest.MonkeyPatch.context() as lookup_patch: + lookup_patch.setattr(litellm.utils, "get_model_info", lookup_tripped_by_a_concurrent_registration) + litellm.register_model( + {alias: {"litellm_provider": "bedrock", "mode": "chat", "supports_vision": True}}, + persist_across_reloads=False, + ) + + usage: Final = Usage(prompt_tokens=1000, completion_tokens=100, total_tokens=1100) + alias_cost: Final = litellm.completion_cost( + completion_response=ModelResponse(model=alias, usage=usage), model=alias, custom_llm_provider="bedrock" + ) + canonical_cost: Final = litellm.completion_cost( + completion_response=ModelResponse(model=canonical_key, usage=usage), + model=canonical_key, + custom_llm_provider="bedrock", + ) + try: + assert alias_cost == canonical_cost > 0 + assert litellm.get_model_info(alias)["input_cost_per_token"] == canonical["input_cost_per_token"] + assert litellm.get_model_info(alias)["supports_vision"] is True + finally: + _invalidate_model_cost_lowercase_map() + + +def test_register_model_keeps_a_catalog_rows_own_prices_when_the_lookup_fails(monkeypatch): + """A catalog row registered under its own key with no provider + (``au.anthropic.`` from a config ``model_info``) must inherit from + itself when ``get_model_info`` trips, never from its region-stripped + sibling, whose prices are the base rate rather than the geo rate. + """ + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + _invalidate_model_cost_lowercase_map() + geo_key: Final = "au.anthropic.claude-sonnet-5" + geo_row: Final = dict(litellm.model_cost[geo_key]) + sibling: Final = litellm.model_cost["anthropic.claude-sonnet-5"] + assert geo_row["input_cost_per_token"] != sibling["input_cost_per_token"] + + def lookup_tripped_by_a_concurrent_registration(model: str, custom_llm_provider: str | None = None) -> NoReturn: + raise RuntimeError("dictionary changed size during iteration") + + with pytest.MonkeyPatch.context() as lookup_patch: + lookup_patch.setattr(litellm.utils, "get_model_info", lookup_tripped_by_a_concurrent_registration) + litellm.register_model({geo_key: {"mode": "chat"}}, persist_across_reloads=False) + + try: + for field in ("input_cost_per_token", "output_cost_per_token", "cache_read_input_token_cost"): + assert litellm.model_cost[geo_key][field] == geo_row[field] + assert litellm.get_model_info(geo_key)[field] == geo_row[field] + assert litellm.model_cost[geo_key]["litellm_provider"] == geo_row["litellm_provider"] + finally: + _invalidate_model_cost_lowercase_map() + + +def test_register_model_inherits_builtin_token_pricing_for_unmapped_key(monkeypatch): + """A key shape ``get_model_info`` cannot resolve inherits the whole built-in + entry, not just its cache pricing, so ``get_model_info`` on the raw key + reports the built-in token prices instead of a synthesized $0.""" + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + _invalidate_model_cost_lowercase_map() + builtin_key: Final = "us.anthropic.claude-sonnet-4-6" + registered_key: Final = f"bedrock/bedrock/bedrock/{builtin_key}" + builtin: Final = litellm.model_cost[builtin_key] + + try: + litellm.register_model( + {registered_key: {"litellm_provider": "bedrock", "mode": "chat"}}, persist_across_reloads=False + ) + info: Final = litellm.get_model_info(registered_key) + assert info["input_cost_per_token"] == builtin["input_cost_per_token"] > 0 + assert info["output_cost_per_token"] == builtin["output_cost_per_token"] > 0 + assert info["cache_read_input_token_cost"] == builtin["cache_read_input_token_cost"] + finally: + _invalidate_model_cost_lowercase_map() + + +def test_register_model_never_inherits_pricing_across_providers(monkeypatch): + """The built-in entry an unmapped key inherits must belong to a provider + ``get_model_info`` would accept for the registered one, exactly as a direct + lookup would, so a bedrock key that happens to share a name with an OpenAI + model registers without OpenAI's prices.""" + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) + _invalidate_model_cost_lowercase_map() + registered_key: Final = "bedrock/gpt-4o" + assert litellm.model_cost["gpt-4o"]["litellm_provider"] == "openai" + assert litellm.model_cost["gpt-4o"]["cache_read_input_token_cost"] > 0 + + try: + litellm.register_model( + {registered_key: {"litellm_provider": "bedrock", "mode": "chat"}}, persist_across_reloads=False + ) + registered: Final = litellm.model_cost[registered_key] + assert "cache_read_input_token_cost" not in registered + assert "input_cost_per_token" not in registered + assert "output_cost_per_token" not in registered + finally: + _invalidate_model_cost_lowercase_map() + + def test_update_dictionary_merges_nested_dicts_without_aliasing(): """A nested dict must be merged copy-on-write: the pre-existing nested dict object stays untouched, and the caller's incoming nested dict is never diff --git a/tests/test_litellm/test_router_model_cost_isolation.py b/tests/test_litellm/test_router_model_cost_isolation.py index d73f5efa96bb..22264c7ceaf7 100644 --- a/tests/test_litellm/test_router_model_cost_isolation.py +++ b/tests/test_litellm/test_router_model_cost_isolation.py @@ -9,9 +9,14 @@ import asyncio import copy +import itertools import logging import os +import queue import re +import sys +import threading +from types import MappingProxyType from typing import Final from unittest.mock import Mock, patch @@ -2716,3 +2721,80 @@ def test_price_data_reload_refreshes_the_cached_model_group_and_deployment_info( assert router.cached_model_group_info("grp").input_cost_per_token == new_price assert router.cached_deployment_model_info("dep-a", "openai/gpt-4o")["input_cost_per_token"] == new_price + + +def test_price_data_reload_keeps_geo_deployments_priced_while_requests_register_pricing(monkeypatch): + """ + Request threads register per-request pricing into litellm.model_cost while a price + data reload walks it. The walk raised "dictionary changed size during iteration", + which aborted the reload after the catalog swap or failed the built-in lookup of the + bedrock geo deployment being replayed; that deployment then registered its own raw + key with metadata only and priced every request at $0 until the next reload. + """ + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + catalog: Final = litellm.get_model_cost_map(url="") + monkeypatch.setattr(litellm, "model_cost", dict(catalog)) + _invalidate_model_cost_lowercase_map() + backends: Final = tuple( + f"bedrock/au.anthropic.{model}" + for model in ("claude-sonnet-5", "claude-opus-5", "claude-haiku-4-5-20251001-v1:0", "claude-sonnet-4-6") + ) + canonical_prices: Final = MappingProxyType( + {backend: catalog[backend.removeprefix("bedrock/")]["input_cost_per_token"] for backend in backends} + ) + assert all(price > 0 for price in canonical_prices.values()) + router: Final = Router( + model_list=[ + { + "model_name": backend.rsplit(".", 1)[-1], + "litellm_params": {"model": backend, "aws_region_name": "ap-southeast-2"}, + "model_info": {"id": f"lit-5853-dep-{index}"}, + } + for index, backend in enumerate(backends) + ] + ) + stop: Final = threading.Event() + writer_outcome: Final[queue.SimpleQueue[int | Exception]] = queue.SimpleQueue() + + def register_until_stopped() -> None: + try: + for i in itertools.count(): + if stop.is_set(): + writer_outcome.put(i) + return + litellm.register_model( + { + f"lit-5853-writer-{i}": { + "litellm_provider": "lit-5853-writer", + "mode": "chat", + "input_cost_per_token": 1e-07, + "output_cost_per_token": 4e-07, + "cache_read_input_token_cost": 1e-08, + } + }, + persist_across_reloads=False, + ) + except Exception as failure: + writer_outcome.put(failure) + + writer: Final = threading.Thread(target=register_until_stopped, daemon=True) + saved_switch_interval: Final = sys.getswitchinterval() + sys.setswitchinterval(1e-4) + writer.start() + try: + for _ in range(10): + _simulate_price_data_reload_with_provider_sets(monkeypatch, dict(catalog)) + finally: + stop.set() + writer.join() + sys.setswitchinterval(saved_switch_interval) + + registrations_landed: Final = writer_outcome.get(timeout=5) + try: + assert isinstance(registrations_landed, int) and registrations_landed > 0, registrations_landed + assert tuple(backend for backend in backends if backend in litellm.model_cost) == () + for backend in backends: + assert litellm.get_model_info(backend)["input_cost_per_token"] == canonical_prices[backend] + assert tuple(router.get_model_names()) == tuple(backend.rsplit(".", 1)[-1] for backend in backends) + finally: + _invalidate_model_cost_lowercase_map() diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index dc6a6ad8cefa..5c85c940b734 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -4007,6 +4007,65 @@ def test_reapply_runtime_registrations_drops_request_scoped_registrations(monkey _invalidate_model_cost_lowercase_map() +def test_model_cost_walks_survive_a_registration_landing_mid_walk(monkeypatch): + """ + A price data reload rebuilds the lowercase lookup map and the provider model + sets by walking litellm.model_cost while request threads keep registering + per-request pricing into it. A key inserted mid-walk raised "dictionary + changed size during iteration", which aborted the reload between the + catalog swap and the deployment replay, or failed the built-in lookup of + the deployment being replayed. Both walks have to run over a snapshot. + """ + from litellm.utils import _invalidate_model_cost_lowercase_map, _rebuild_model_cost_lowercase_map + + class _KeyRegisteredMidWalk(str): + def lower(self) -> str: + litellm.model_cost.setdefault("lit-5853-registered-during-the-rebuild", {"mode": "chat"}) + return super().lower() + + class _EntryRegisteredMidWalk(dict): + def get(self, field, default=None): + litellm.model_cost.setdefault("lit-5853-registered-during-add-known-models", {"mode": "chat"}) + return super().get(field, default) + + monkeypatch.setattr(litellm, "model_cost", dict(litellm.model_cost)) + litellm.model_cost[_KeyRegisteredMidWalk("LIT-5853-WALKED-LAST")] = _EntryRegisteredMidWalk(mode="chat") + try: + assert _rebuild_model_cost_lowercase_map()["lit-5853-walked-last"] == "LIT-5853-WALKED-LAST" + litellm.add_known_models(model_cost_map=litellm.model_cost) + assert "lit-5853-registered-during-the-rebuild" in litellm.model_cost + assert "lit-5853-registered-during-add-known-models" in litellm.model_cost + finally: + _invalidate_model_cost_lowercase_map() + + +class _KeyWhoseLoweringLosesTheRace(str): + """Lowering the key is the one Python-level step between the lookup map's rebuild and + its read, so another thread's registration invalidating the map lands right there.""" + + def lower(self) -> str: + from litellm.utils import _invalidate_model_cost_lowercase_map + + _invalidate_model_cost_lowercase_map() + return super().lower() + + +def test_case_insensitive_model_cost_lookup_survives_an_invalidation_mid_lookup(): + """ + _get_model_cost_key read the module-level lowercase map after rebuilding it, so a + registration on another thread that invalidated the map in between turned the + lookup into an AttributeError on None, which the model info ladder reports as + "model isn't mapped yet" and register_model turns into a $0 shadow entry. + """ + from litellm.utils import _get_model_cost_key, _invalidate_model_cost_lowercase_map + + _invalidate_model_cost_lowercase_map() + try: + assert _get_model_cost_key(_KeyWhoseLoweringLosesTheRace("GPT-4O")) == "gpt-4o" + finally: + _invalidate_model_cost_lowercase_map() + + class _JsonCapture(logging.Handler): def __init__(self): super().__init__() @@ -5693,6 +5752,31 @@ def test_get_model_info_gemini(monkeypatch): assert info.get("rpm") is not None, f"{model} does not have rpm" +def test_register_model_without_a_provider_never_inherits_another_providers_entry(monkeypatch): + """ + The router registers a deployment's shared ``{provider}/{model}`` key without + ``litellm_provider``. Matching the built-in entry as a wildcard then handed + ``azure/claude-sonnet-4-6`` Anthropic's whole entry, provider included, and + reading the key back under the provider its prefix names raised "isn't mapped". + """ + from litellm.utils import _invalidate_model_cost_lowercase_map + + key: Final = "azure/claude-sonnet-4-6" + monkeypatch.setattr(litellm, "model_cost", dict(litellm.model_cost)) + litellm.model_cost.pop(key, None) + try: + litellm.register_model( + {key: {"input_cost_per_token": 1e-07, "output_cost_per_token": 4e-07}}, persist_across_reloads=False + ) + assert litellm.model_cost[key].get("litellm_provider") is None + info: Final = litellm.get_model_info(key) + assert info["litellm_provider"] == "azure" + assert info["input_cost_per_token"] == 1e-07 + assert info["output_cost_per_token"] == 4e-07 + finally: + _invalidate_model_cost_lowercase_map() + + @pytest.mark.parametrize( ("max_parallel_requests", "rpm", "tpm", "default_max_parallel_requests", "expected"), [