diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index d458d0f7c4a..91f0b9b2790 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -1040,22 +1040,6 @@ class ModelManagementAuthChecks: return True -def _deployment_name_and_model(deployment: Optional[Union[Deployment, Dict[str, object]]]) -> Tuple[Optional[str], str]: - """Return (model_name, litellm_params.model) for a deployment. - - delete_deployment is annotated to return a Deployment but hands back the raw - model_list dict at runtime, so both shapes are handled; the model defaults to "". - """ - if deployment is None: - return None, "" - if isinstance(deployment, dict): - name = deployment.get("model_name") - params = deployment.get("litellm_params") - model = params.get("model") if isinstance(params, dict) else None - return (name if isinstance(name, str) else None), (model if isinstance(model, str) else "") - return deployment.model_name, str(getattr(deployment.litellm_params, "model", "") or "") - - #### [BETA] - This is a beta endpoint, format might change based on user feedback. - https://github.com/BerriAI/litellm/issues/964 @router.post( "/model/delete", @@ -1127,19 +1111,7 @@ async def delete_model( ## DELETE FROM ROUTER ## if llm_router is not None: - deleted_deployment = llm_router.delete_deployment(id=model_info.id) - # delete_deployment only drops the deployment from model_list; the auto/ - # complexity router registries are keyed by model_name and would otherwise - # retain a stale (now unbacked) entry, so evict it here too. Guard on the - # auto_router/ prefix (as clear_cache does): a regular DB model that merely - # shares a model_name with a config-defined router must not evict that router, - # since add_deployment never restores config-defined routers. - deleted_name, deleted_model = _deployment_name_and_model(deleted_deployment) - if deleted_name is not None and deleted_model.startswith("auto_router/"): - llm_router.auto_routers.pop(deleted_name, None) - llm_router.complexity_routers.pop(deleted_name, None) - llm_router.adaptive_routers.pop(deleted_name, None) - llm_router.quality_routers.pop(deleted_name, None) + llm_router.delete_deployment(id=model_info.id) # Runs after the row delete so the sibling check sees post-delete state. if model_params.model_info.team_id is not None: diff --git a/litellm/router.py b/litellm/router.py index 78fe3ff025e..4aa2731466e 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -7734,6 +7734,51 @@ class Router: TaggedPreRoutingStrategy(tags=tags, strategy=strategy), ] + @staticmethod + def _unregister_pre_routing_strategy( + registry: dict[str, list[TaggedPreRoutingStrategy[_PreRoutingStrategyT]]], + model_name: str, + tags: tuple[str, ...], + ) -> bool: + """Drop the strategy registered for this exact (model_name, tags) pair, leaving + strategies registered under the same name with different tags in place. Returns + whether anything was actually dropped.""" + existing = registry.get(model_name, []) + remaining = [entry for entry in existing if entry.tags != tags] + if len(remaining) == len(existing): + return False + if remaining: + registry[model_name] = remaining + else: + registry.pop(model_name, None) + return True + + def _unregister_pre_routing_strategy_for_deployment(self, deployment: Deployment) -> None: + """ + Release the pre-routing strategy a deployment holds, so removing it from the + model_list also frees its (model_name, tags) slot. + + Without this, re-adding the deployment (an edit arriving via upsert_deployment, + or a router recreated under a name that was deleted earlier) hits the + "already exists" guard in `_register_pre_routing_strategy`, which + `ignore_invalid_deployments` swallows - the deployment then silently never + makes it back into the model_list. + + Released from every registry rather than the first match, because registration is + one-to-many: a complexity router configured with `adaptive` is also registered in + `adaptive_routers` under the same (model_name, tags) by the deferred finalize pass. + Guarded on the auto_router/ prefix so removing a *regular* deployment can't evict a + router that merely shares its model_name. + """ + if not deployment.litellm_params.model.startswith("auto_router/"): + return + model_name = deployment.model_name + tags = self._deployment_tags(deployment) + for registry in (self.auto_routers, self.complexity_routers, self.quality_routers): + self._unregister_pre_routing_strategy(registry, model_name, tags) + if self._unregister_pre_routing_strategy(self.adaptive_routers, model_name, tags): + self._sync_adaptive_router_hooks() + def _finalize_adaptive_router_if_configured(self) -> None: """Locate every adaptive-router deployment in the finalized model_list and build an AdaptiveRouter for each. Safe no-op when none are configured. @@ -7779,6 +7824,16 @@ class Router: TaggedPreRoutingStrategy(tags=tagged.tags, strategy=adaptive_router), ] + self._sync_adaptive_router_hooks() + + def _sync_adaptive_router_hooks(self) -> None: + """Rebuild the AdaptiveRouterPostCallHook set so it is exactly one hook per + currently registered adaptive router. Run at every point the adaptive registry + changes, otherwise a released router keeps recording turns through its hook.""" + from litellm.router_strategy.adaptive_router.hooks import ( + AdaptiveRouterPostCallHook, + ) + for callback in litellm.logging_callback_manager.get_custom_loggers_for_type(AdaptiveRouterPostCallHook): litellm.logging_callback_manager.remove_callback_from_all_lists(callback) for tagged_adaptive_routers in self.adaptive_routers.values(): @@ -8401,13 +8456,27 @@ class Router: self._invalidate_access_groups_cache() self._update_deployment_indices_after_removal(model_id=deployment_id, removal_idx=removal_idx) + # Free the outgoing deployment's pre-routing strategy slot (keyed by the + # OLD model_name/tags) before the re-add below re-registers it. + self._unregister_pre_routing_strategy_for_deployment(deployment=_deployment_on_router) + # if the model_id is not in router self.add_deployment(deployment=deployment) + # add_deployment() builds every strategy EXCEPT the adaptive one, which + # set_model_list() defers until the whole model_list is visible. Re-run that + # deferred pass so an adaptive router whose slot was just released above is + # rebuilt rather than left unregistered. + if self._is_adaptive_router_deployment(litellm_params=deployment.litellm_params) or ( + _deployment_on_router is not None + and self._is_adaptive_router_deployment(litellm_params=_deployment_on_router.litellm_params) + ): + self._finalize_adaptive_router_if_configured() return deployment except Exception as e: if self.ignore_invalid_deployments: - verbose_router_logger.debug( - f"Error upserting deployment: {e}, ignoring and continuing with other deployments." + verbose_router_logger.warning( + f"Error upserting deployment {deployment.model_name} (id={deployment.model_info.id}): {e}. " + "Dropping it and continuing with other deployments." ) return None else: @@ -8428,8 +8497,14 @@ class Router: try: if deployment_idx is not None: + try: + deployment_to_remove = self.get_deployment(model_id=id) + except Exception: + deployment_to_remove = None # Pop the item from the list first item = self.model_list.pop(deployment_idx) + if deployment_to_remove is not None: + self._unregister_pre_routing_strategy_for_deployment(deployment=deployment_to_remove) self._invalidate_model_group_info_cache() self._invalidate_access_groups_cache() self._update_deployment_indices_after_removal(model_id=id, removal_idx=deployment_idx) diff --git a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py index f3e5e2c9b71..bd5eda0197b 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py @@ -636,14 +636,33 @@ class TestDeleteModelClearsRouterRegistry: not just from model_list, or a stale (now unbacked) router entry lingers until restart. """ + @staticmethod + def _complexity_router_deployment(model_id: str, tags: list | None = None) -> dict: + return { + "model_name": "smart-router", + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": {"tiers": {"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o"}}, + "complexity_router_default_model": "gpt-4o", + **({"tags": tags} if tags else {}), + }, + "model_info": {"id": model_id, "db_model": True}, + } + @pytest.mark.asyncio - async def test_delete_model_pops_router_registries(self): + async def test_delete_model_releases_only_the_deleted_routers_slot(self): + """Deleting one tagged router must release its own slot and leave a sibling + sharing the model_name registered. A blanket pop(model_name) here would take + both down, and nothing reloads on the delete path to restore the survivor. + """ + import litellm + from litellm.proxy.management_endpoints.model_management_endpoints import ModelInfoDelete from litellm.proxy.management_endpoints.model_management_endpoints import ( delete_model as delete_model_endpoint, ) - from litellm.proxy.management_endpoints.model_management_endpoints import ModelInfoDelete model_id = "router-del-1" + surviving_id = "router-del-2" admin_user = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) db_row = LiteLLM_ProxyModelTable( model_id=model_id, @@ -660,16 +679,16 @@ class TestDeleteModelClearsRouterRegistry: mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=db_row) mock_prisma.db.litellm_proxymodeltable.delete = AsyncMock(return_value=db_row) - mock_router = MagicMock() - mock_router.delete_deployment = MagicMock( - return_value={ - "model_name": "smart-router", - "litellm_params": {"model": "auto_router/complexity_router"}, - "model_info": {"id": model_id}, - } + real_router = litellm.Router( + model_list=[ + {"model_name": "gpt-4o", "litellm_params": {"model": "gpt-4o"}}, + {"model_name": "gpt-4o-mini", "litellm_params": {"model": "gpt-4o-mini"}}, + self._complexity_router_deployment(model_id, tags=["team-a"]), + self._complexity_router_deployment(surviving_id, tags=["team-b"]), + ], + ignore_invalid_deployments=True, ) - mock_router.auto_routers = {"smart-router": MagicMock()} - mock_router.complexity_routers = {"smart-router": MagicMock()} + assert len(real_router.complexity_routers["smart-router"]) == 2 _PS = "litellm.proxy.proxy_server" with ( @@ -679,16 +698,17 @@ class TestDeleteModelClearsRouterRegistry: patch(f"{_PS}.proxy_logging_obj", MagicMock()), patch(f"{_PS}.general_settings", {}), patch(f"{_PS}.premium_user", True), - patch(f"{_PS}.llm_router", mock_router), + patch(f"{_PS}.llm_router", real_router), ): await delete_model_endpoint( model_info=ModelInfoDelete(id=model_id), user_api_key_dict=admin_user, ) - mock_router.delete_deployment.assert_called_once_with(id=model_id) - assert "smart-router" not in mock_router.auto_routers - assert "smart-router" not in mock_router.complexity_routers + assert model_id not in [m["model_info"]["id"] for m in real_router.model_list] + surviving = real_router.complexity_routers["smart-router"] + assert len(surviving) == 1 + assert surviving[0].tags == ("team-b",) @pytest.mark.asyncio async def test_delete_regular_model_preserves_config_router_sharing_name(self): diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index a9e5b3316e0..3b7bcad78b0 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -5936,3 +5936,274 @@ async def test_acreate_batch_request_bedrock_tags_override_deployment_tags(): bedrock_tags=request_tags, ) assert mock_sign.call_args.kwargs["data"]["tags"] == request_tags + + +class TestPreRoutingStrategyRegistryLifecycle: + """ + Regression tests: a deployment leaving the model_list must release the + pre-routing strategy slot it holds in `auto_routers` / `complexity_routers` / + `adaptive_routers` / `quality_routers`. + + Before this fix, editing an auto-router-family model (a UI save, which reaches + every other pod as an `upsert_deployment` from the periodic DB reload) popped + the deployment out of the model_list and then failed to re-add it: registration + raised "already exists" against the stale registry entry, and + `ignore_invalid_deployments=True` swallowed the error. The router vanished from + the Models page and stayed gone until a proxy restart, while the DB row and the + "saved successfully" response both looked fine. + """ + + @staticmethod + def _complexity_router_params(default_model: str, tags=None) -> dict: + return { + "model": "auto_router/complexity_router", + "complexity_router_config": { + "tiers": {"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o", "COMPLEX": "gpt-4o"} + }, + "complexity_router_default_model": default_model, + **({"tags": tags} if tags else {}), + } + + @classmethod + def _router_with_complexity_router(cls, default_model: str = "gpt-4o") -> "litellm.Router": + return litellm.Router( + model_list=[ + {"model_name": "gpt-4o", "litellm_params": {"model": "gpt-4o"}}, + {"model_name": "gpt-4o-mini", "litellm_params": {"model": "gpt-4o-mini"}}, + { + "model_name": "smart-router", + "litellm_params": cls._complexity_router_params(default_model), + "model_info": {"id": "router-1", "db_model": True}, + }, + ], + ignore_invalid_deployments=True, + ) + + @staticmethod + def _model_names(router: "litellm.Router") -> list: + return [model["model_name"] for model in router.model_list] + + def test_upsert_of_edited_router_keeps_it_routable(self): + from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo + + router = self._router_with_complexity_router() + + router.upsert_deployment( + deployment=Deployment( + model_name="smart-router", + litellm_params=LiteLLM_Params(**self._complexity_router_params("gpt-4o-mini")), + model_info=ModelInfo(id="router-1", db_model=True), + ) + ) + + assert "smart-router" in self._model_names(router) + registered = router.complexity_routers["smart-router"] + assert len(registered) == 1 + # the surviving strategy is the edited one, not the pre-edit leftover + assert registered[0].strategy.config.default_model == "gpt-4o-mini" + + def test_unchanged_upsert_leaves_router_untouched(self): + from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo + + router = self._router_with_complexity_router() + strategy_before = router.complexity_routers["smart-router"][0].strategy + + for _ in range(3): + router.upsert_deployment( + deployment=Deployment( + model_name="smart-router", + litellm_params=LiteLLM_Params(**self._complexity_router_params("gpt-4o")), + model_info=ModelInfo(id="router-1", db_model=True), + ) + ) + + assert "smart-router" in self._model_names(router) + assert router.complexity_routers["smart-router"][0].strategy is strategy_before + + def test_delete_frees_the_name_for_a_new_router(self): + from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo + + router = self._router_with_complexity_router() + + router.delete_deployment(id="router-1") + assert "smart-router" not in router.complexity_routers + + router.add_deployment( + deployment=Deployment( + model_name="smart-router", + litellm_params=LiteLLM_Params(**self._complexity_router_params("gpt-4o-mini")), + model_info=ModelInfo(id="router-2", db_model=True), + ) + ) + + assert "smart-router" in self._model_names(router) + assert router.complexity_routers["smart-router"][0].strategy.config.default_model == "gpt-4o-mini" + + def test_delete_only_frees_the_matching_tag_slot(self): + router = litellm.Router( + model_list=[ + {"model_name": "gpt-4o", "litellm_params": {"model": "gpt-4o"}}, + {"model_name": "gpt-4o-mini", "litellm_params": {"model": "gpt-4o-mini"}}, + { + "model_name": "shared-router", + "litellm_params": self._complexity_router_params("gpt-4o", tags=["team-a"]), + "model_info": {"id": "router-a"}, + }, + { + "model_name": "shared-router", + "litellm_params": self._complexity_router_params("gpt-4o-mini", tags=["team-b"]), + "model_info": {"id": "router-b"}, + }, + ], + ignore_invalid_deployments=True, + ) + assert len(router.complexity_routers["shared-router"]) == 2 + + router.delete_deployment(id="router-a") + + remaining = router.complexity_routers["shared-router"] + assert len(remaining) == 1 + assert remaining[0].tags == ("team-b",) + + def test_delete_of_regular_model_preserves_router_sharing_its_name(self): + router = litellm.Router( + model_list=[ + {"model_name": "gpt-4o", "litellm_params": {"model": "gpt-4o"}}, + {"model_name": "gpt-4o-mini", "litellm_params": {"model": "gpt-4o-mini"}}, + { + "model_name": "shared-name", + "litellm_params": self._complexity_router_params("gpt-4o"), + "model_info": {"id": "router-1"}, + }, + { + "model_name": "shared-name", + "litellm_params": {"model": "openai/gpt-4o"}, + "model_info": {"id": "regular-1"}, + }, + ], + ignore_invalid_deployments=True, + ) + strategy = router.complexity_routers["shared-name"][0].strategy + + router.delete_deployment(id="regular-1") + + assert router.complexity_routers["shared-name"][0].strategy is strategy + + def test_upsert_of_edited_adaptive_router_rebuilds_it(self): + """Adaptive routers are built by set_model_list()'s deferred pass, not by + add_deployment(), so releasing the slot on edit must be paired with a rebuild - + otherwise the edit silently turns adaptive routing off.""" + from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo + + def adaptive_params(available_models: list) -> dict: + return { + "model": "auto_router/adaptive_router", + "adaptive_router_config": {"available_models": available_models}, + } + + router = litellm.Router( + model_list=[ + {"model_name": "gpt-4o", "litellm_params": {"model": "openai/gpt-4o"}}, + {"model_name": "gpt-4o-mini", "litellm_params": {"model": "openai/gpt-4o-mini"}}, + { + "model_name": "adaptive-router", + "litellm_params": adaptive_params(["gpt-4o-mini"]), + "model_info": {"id": "router-1", "db_model": True}, + }, + ], + ignore_invalid_deployments=True, + ) + assert "adaptive-router" in router.adaptive_routers + + router.upsert_deployment( + deployment=Deployment( + model_name="adaptive-router", + litellm_params=LiteLLM_Params(**adaptive_params(["gpt-4o", "gpt-4o-mini"])), + model_info=ModelInfo(id="router-1", db_model=True), + ) + ) + + assert "adaptive-router" in self._model_names(router) + registered = router.adaptive_routers["adaptive-router"] + assert len(registered) == 1 + assert set(registered[0].strategy.config.available_models) == {"gpt-4o", "gpt-4o-mini"} + + def test_delete_of_adaptive_enabled_complexity_router_frees_both_registries(self): + """A complexity router with adaptive set is registered in BOTH complexity_routers + and adaptive_routers under the same (model_name, tags). Releasing only the first + match leaves the adaptive strategy live, so a deleted alias stays routable and its + post-call hook keeps recording.""" + import litellm as litellm_module + from litellm.router_strategy.adaptive_router.hooks import AdaptiveRouterPostCallHook + + params = { + "model": "auto_router/complexity_router", + "complexity_router_config": { + "tiers": {"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o"}, + "adaptive": True, + }, + "complexity_router_default_model": "gpt-4o", + } + router = litellm.Router( + model_list=[ + {"model_name": "gpt-4o", "litellm_params": {"model": "openai/gpt-4o"}}, + {"model_name": "gpt-4o-mini", "litellm_params": {"model": "openai/gpt-4o-mini"}}, + { + "model_name": "hybrid-router", + "litellm_params": params, + "model_info": {"id": "router-1", "db_model": True}, + }, + ], + ignore_invalid_deployments=True, + ) + assert "hybrid-router" in router.complexity_routers + assert "hybrid-router" in router.adaptive_routers + hooks = litellm_module.logging_callback_manager.get_custom_loggers_for_type(AdaptiveRouterPostCallHook) + assert len(hooks) == 1 + + router.delete_deployment(id="router-1") + + assert "hybrid-router" not in router.complexity_routers + assert "hybrid-router" not in router.adaptive_routers + remaining_hooks = litellm_module.logging_callback_manager.get_custom_loggers_for_type( + AdaptiveRouterPostCallHook + ) + assert remaining_hooks == [] + + def test_upsert_of_edited_quality_router_keeps_it_routable(self): + """_unregister_pre_routing_strategy_for_deployment dispatches on four prefixes; + quality_router is one of them and would otherwise go unexercised.""" + from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo + + def quality_params(default_model: str) -> dict: + return { + "model": "auto_router/quality_router", + "quality_router_default_model": default_model, + } + + router = litellm.Router( + model_list=[ + {"model_name": "gpt-4o", "litellm_params": {"model": "openai/gpt-4o"}}, + {"model_name": "gpt-4o-mini", "litellm_params": {"model": "openai/gpt-4o-mini"}}, + { + "model_name": "quality-router", + "litellm_params": quality_params("gpt-4o"), + "model_info": {"id": "router-1", "db_model": True}, + }, + ], + ignore_invalid_deployments=True, + ) + assert "quality-router" in router.quality_routers + + router.upsert_deployment( + deployment=Deployment( + model_name="quality-router", + litellm_params=LiteLLM_Params(**quality_params("gpt-4o-mini")), + model_info=ModelInfo(id="router-1", db_model=True), + ) + ) + + assert "quality-router" in self._model_names(router) + registered = router.quality_routers["quality-router"] + assert len(registered) == 1 + assert registered[0].strategy.config.default_model == "gpt-4o-mini"