From 9724f06365e9b39d7b1a94d15bca79b33ed46106 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Sun, 6 Sep 2026 11:04:32 +0000 Subject: [PATCH] fix(router): hold a cost-map key while another live router still serves it The claim is process-wide but the release was per-deletion, so with two routers serving one deployment id, the first deletion put the survivor back on merging and the price it had just cleared would keep billing. Release the key only once no live router still serves that id. --- litellm/router.py | 11 ++++-- .../test_router_model_cost_isolation.py | 34 +++++++++++++++++++ 2 files changed, 42 insertions(+), 3 deletions(-) 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