From 315b098a1e51da320ef356285c857c7aaa224391 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Sun, 6 Sep 2026 10:46:56 +0000 Subject: [PATCH] fix(router): decide cost-map eviction by what this registrar created The previous guard read a catalog entry off `litellm_provider`, so a deployment that declares its own provider in model_info was treated as one and kept billing at a price it no longer carried. It also only held for a single registration: a second one under a colliding id saw the id the first merge left behind and evicted the catalog entry anyway. Track the cost-map keys this registrar creates instead. A key it created is evicted before re-registration; one it did not is left to merge, which is what a deployment id colliding with a catalog model name needs. Also folds the required-fields comment into the docstring that already gives the reason. --- .../model_management_endpoints.py | 9 ++-- litellm/router.py | 25 +++++++---- .../test_router_model_cost_isolation.py | 45 ++++++++++++++++--- 3 files changed, 58 insertions(+), 21 deletions(-) diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 113477dbbdd..742b9d9817f 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -743,8 +743,6 @@ _OWNERSHIP_FIELDS: Final = frozenset( } ) -# Clearing a required field writes a row no reload can rebuild: both blobs load through -# LiteLLM_Params / ModelInfo, which reject it. _STORED_REQUIRED_FIELDS: Final = frozenset( name for model in (LiteLLM_Params, ModelInfo) for name, field in model.model_fields.items() if field.is_required() ) @@ -754,9 +752,10 @@ _NULL_CLEAR_IGNORED_FIELDS: Final = _OWNERSHIP_FIELDS | _STORED_REQUIRED_FIELDS def _explicitly_cleared_fields(patch: BaseModel | None) -> frozenset[str]: """The keys a patch sends as an explicit null, which update_db_model removes from the - stored blob (JSON Merge Patch). Ownership keys and the keys the stored models require - are left alone, and the PTU keys are handled by _explicitly_cleared_ptu_fields, whose - clear is gated on the feature flag. Applied after both blobs merge, so a model_info blob + stored blob (JSON Merge Patch). Ownership keys are left alone, as are the keys the stored + models require, since clearing one writes a row no reload can rebuild through + LiteLLM_Params / ModelInfo. The PTU keys are handled by _explicitly_cleared_ptu_fields, + whose clear is gated on the feature flag. Applied after both blobs merge, so a model_info blob the UI echoes back cannot resurrect a pricing key the litellm_params patch clears. """ if patch is None: diff --git a/litellm/router.py b/litellm/router.py index 306fee83f2d..0dd4af038b9 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -612,6 +612,13 @@ RETRY_BREADCRUMB_EXCLUDED_KWARGS: Final = frozenset( RETRY_BREADCRUMB_LIMIT: Final = 4 +# Cost-map keys created by _register_deployment_in_model_cost, which shares one flat +# namespace with the built-in model catalog. Only a key it created may be evicted, or a +# deployment whose id names a real model would strip that model's pricing and +# capabilities for every other deployment of it. +_DEPLOYMENT_COST_MAP_KEYS: Final[set[str]] = set() # mutable-ok: ownership of shared cost-map keys + + class Router: model_names: set = set() cache_responses: bool | None = False @@ -9745,10 +9752,12 @@ class Router: """Write a deployment's metadata into ``litellm.model_cost``. Runs when a deployment is added and again after a price data reload, so - the entries a refresh rebuilds are the ones a fresh boot would produce. The - deployment's own ``model_id`` entry is replaced rather than merged, so a - price cleared from the deployment does not linger from an earlier - registration and keep billing at the old rate. + the entries a refresh rebuilds are the ones a fresh boot would produce. An + entry this function created is replaced rather than merged, so a price cleared + from the deployment does not linger from an earlier registration and keep + billing at the old rate. An entry it did not create is left to merge, because + a deployment id that collides with a catalog model name shares that model's + entry with every other deployment of it. Nothing is recorded for replay: a refresh walks the live routers instead, so a deleted, repointed or never-added deployment, and a discarded router, drop out of the rebuild on their own. @@ -9765,12 +9774,10 @@ class Router: } if model_id is not None: - # Deployments key into the same cost map as the built-in catalog, so evict only an - # entry this registration owns; a deployment id that names a real model keeps the - # old merge rather than stripping what every other deployment of it reads. - registered: Final = litellm.model_cost.get(model_id) - if registered is not None and registered.get("litellm_provider") is None: + if model_id in _DEPLOYMENT_COST_MAP_KEYS: litellm.model_cost.pop(model_id, None) + elif model_id not in litellm.model_cost: + _DEPLOYMENT_COST_MAP_KEYS.add(model_id) litellm.register_model( model_cost={model_id: model_info}, persist_across_reloads=False, diff --git a/tests/test_litellm/test_router_model_cost_isolation.py b/tests/test_litellm/test_router_model_cost_isolation.py index 28cd70425e9..b9cbcbdfb08 100644 --- a/tests/test_litellm/test_router_model_cost_isolation.py +++ b/tests/test_litellm/test_router_model_cost_isolation.py @@ -260,19 +260,21 @@ def test_should_not_strip_a_builtin_entry_when_a_deployment_id_collides_with_it( whose id happens to name a real model must not evict that model's entry. Stripping it would take the pricing and capability flags every other deployment of that - model reads, process-wide, until the next price-map reload. + model reads, process-wide, until the next price-map reload. Registering twice, because + the first registration is what would mark the entry as this deployment's own. """ colliding_id = "gpt-4o" original = {colliding_id: litellm.model_cost.get(colliding_id)} builtin_max_tokens = litellm.model_cost[colliding_id]["max_tokens"] try: - Router._register_deployment_in_model_cost( - model_id=colliding_id, - model_info={"id": colliding_id, "db_model": True, "mode": "chat"}, - model="gpt-4o-mini", - custom_llm_provider="openai", - ) + for _ in range(2): + Router._register_deployment_in_model_cost( + model_id=colliding_id, + model_info={"id": colliding_id, "db_model": True, "mode": "chat"}, + model="gpt-4o-mini", + custom_llm_provider="openai", + ) entry = litellm.model_cost[colliding_id] assert entry["max_tokens"] == builtin_max_tokens, ( @@ -284,6 +286,35 @@ def test_should_not_strip_a_builtin_entry_when_a_deployment_id_collides_with_it( _restore_model_cost_entries(original) +def test_should_drop_a_stale_price_even_when_the_deployment_declares_a_provider(): + """A deployment may carry `litellm_provider` in its own model_info, which must not be + read as "this is a catalog entry" and stop the stale price from being dropped.""" + model_id = "deployment-provider-tagged" + original = {model_id: litellm.model_cost.get(model_id)} + + try: + Router._register_deployment_in_model_cost( + model_id=model_id, + model_info={"id": model_id, "litellm_provider": "openai", "input_cost_per_token": 0.005}, + model="gpt-4o-mini", + custom_llm_provider="openai", + ) + assert litellm.model_cost[model_id]["input_cost_per_token"] == 0.005 + + Router._register_deployment_in_model_cost( + model_id=model_id, + model_info={"id": model_id, "litellm_provider": "openai", "mode": "chat"}, + model="gpt-4o-mini", + custom_llm_provider="openai", + ) + + assert litellm.model_cost[model_id].get("input_cost_per_token") != 0.005, ( + "a deployment that declares its provider kept billing at the price it no longer carries" + ) + finally: + _restore_model_cost_entries(original) + + def test_should_preserve_builtin_pricing_regardless_of_deployment_order(): """ The built-in pricing should be preserved no matter which deployment