diff --git a/litellm/router.py b/litellm/router.py index 95c95f86f87..2e256a8ce23 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -615,8 +615,9 @@ 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. delete_deployment gives a key back, so a -# later catalog refresh that starts serving that name is not treated as a deployment's own. +# capabilities for every other deployment of it. delete_deployment gives a key back once no +# live router still serves that id, so a later catalog refresh that starts serving the name +# is not treated as a deployment's own. _DEPLOYMENT_COST_MAP_KEYS: Final[set[str]] = set() # mutable-ok: ownership of shared cost-map keys @@ -9875,7 +9876,11 @@ class Router: _budget_limiter: Final = self._get_router_deployment_budget_limiter() if _budget_limiter is not None: _budget_limiter.unregister_deployment_budget(model_id=id) - _DEPLOYMENT_COST_MAP_KEYS.discard(id) + if not any( + router is not self and id in router.model_id_to_deployment_index_map + for router in tuple(_live_routers) + ): + _DEPLOYMENT_COST_MAP_KEYS.discard(id) try: self._unregister_pre_routing_strategy_for_deployment( deployment=item if isinstance(item, Deployment) else Deployment(**item) diff --git a/tests/test_litellm/test_router_model_cost_isolation.py b/tests/test_litellm/test_router_model_cost_isolation.py index 749af1da7c0..c7dca63feb1 100644 --- a/tests/test_litellm/test_router_model_cost_isolation.py +++ b/tests/test_litellm/test_router_model_cost_isolation.py @@ -348,6 +348,40 @@ def test_should_give_a_cost_map_key_back_when_the_deployment_is_deleted(): _restore_model_cost_entries(original) +def test_should_keep_the_cost_map_key_while_another_router_still_serves_it(): + """Two live routers can serve the same deployment id, and the claim is process-wide. + + Releasing it when only one of them drops the deployment would put the survivor back on + merging, so the price it just cleared would keep billing. + """ + from litellm.router import _DEPLOYMENT_COST_MAP_KEYS + + model_id = "deployment-served-twice" + original = {model_id: litellm.model_cost.get(model_id)} + entry = { + "model_name": "served-twice", + "litellm_params": {"model": "gpt-4o-mini", "mock_response": "ok"}, + "model_info": {"id": model_id, "input_cost_per_token": 0.005}, + } + first = Router(model_list=[entry]) + second = Router(model_list=[entry]) + + try: + assert model_id in _DEPLOYMENT_COST_MAP_KEYS + + assert first.delete_deployment(id=model_id) is not None + + assert model_id in _DEPLOYMENT_COST_MAP_KEYS, ( + "the claim was released while another router still served the deployment" + ) + + assert second.delete_deployment(id=model_id) is not None + assert model_id not in _DEPLOYMENT_COST_MAP_KEYS + finally: + _DEPLOYMENT_COST_MAP_KEYS.discard(model_id) + _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