From ef9ea1f8f20b53cae0ff8050dee04591b6f053f8 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Mon, 23 Mar 2026 16:44:30 +0530 Subject: [PATCH] fix(router): address Greptile P1/P2 performance issues - Guard against llm_router=None to prevent silent deletion - Add O(1) team_model index to avoid O(n) scan on every team request Made-with: Cursor --- .../model_management_endpoints.py | 26 ++++---- litellm/router.py | 61 +++++++++++++++++++ 2 files changed, 76 insertions(+), 11 deletions(-) diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index eb85112a66..bfd67ea4e5 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -474,11 +474,15 @@ async def _update_existing_team_model_assignment( if old_public_name and public_model_name != old_public_name: from litellm.proxy.proxy_server import llm_router - other_deployments_with_old_name = [] - if llm_router: + if llm_router is None: + verbose_proxy_logger.warning( + "llm_router not initialized; skipping old public name cleanup to preserve sibling deployments" + ) + else: all_deployments = llm_router.get_model_list( model_name=old_public_name, team_id=team_id ) + other_deployments_with_old_name = [] if all_deployments: other_deployments_with_old_name = [ d @@ -488,15 +492,15 @@ async def _update_existing_team_model_assignment( == old_public_name ] - if not other_deployments_with_old_name: - await team_model_delete( - data=TeamModelDeleteRequest( - team_id=team_id, - models=[old_public_name], - ), - http_request=Request(scope={"type": "http"}), - user_api_key_dict=user_api_key_dict, - ) + if not other_deployments_with_old_name: + await team_model_delete( + data=TeamModelDeleteRequest( + team_id=team_id, + models=[old_public_name], + ), + http_request=Request(scope={"type": "http"}), + user_api_key_dict=user_api_key_dict, + ) await team_model_add( data=TeamModelAddRequest( diff --git a/litellm/router.py b/litellm/router.py index 767938d11f..c9d79beacc 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -467,6 +467,8 @@ class Router: # Initialize model name to deployment indices mapping for O(1) lookups # Maps model_name -> list of indices in model_list self.model_name_to_deployment_indices: Dict[str, List[int]] = {} + # Maps (team_id, team_public_model_name) -> list of indices in model_list + self.team_model_to_deployment_indices: Dict[Tuple[str, str], List[int]] = {} if model_list is not None: # set_model_list will build indices automatically @@ -6835,6 +6837,7 @@ class Router: self.model_list = [] self.model_id_to_deployment_index_map = {} # Reset the index self.model_name_to_deployment_indices = {} # Reset the model_name index + self.team_model_to_deployment_indices = {} # Reset the team_model index self._invalidate_model_group_info_cache() self._invalidate_access_groups_cache() # we add api_base/api_key each model so load balancing between azure/gpt on api_base1 and api_base2 works @@ -7150,6 +7153,26 @@ class Router: else: del self.model_name_to_deployment_indices[model_name] + # Update team_model_to_deployment_indices + for key, indices in list(self.team_model_to_deployment_indices.items()): + # Remove the deleted index + if removal_idx in indices: + indices.remove(removal_idx) + + # Decrement all indices greater than removal_idx + updated_indices = [] + for idx in indices: + if idx > removal_idx: + updated_indices.append(idx - 1) + else: + updated_indices.append(idx) + + # Update or remove the entry + if len(updated_indices) > 0: + self.team_model_to_deployment_indices[key] = updated_indices + else: + del self.team_model_to_deployment_indices[key] + def _add_model_to_list_and_index_map( self, model: dict, model_id: Optional[str] = None ) -> None: @@ -7178,6 +7201,17 @@ class Router: self.model_name_to_deployment_indices[model_name] = [] self.model_name_to_deployment_indices[model_name].append(idx) + # Update team_model index for O(1) team-scoped lookup + team_id = model.get("model_info", {}).get("team_id") + team_public_model_name = model.get("model_info", {}).get( + "team_public_model_name" + ) + if team_id and team_public_model_name: + key = (team_id, team_public_model_name) + if key not in self.team_model_to_deployment_indices: + self.team_model_to_deployment_indices[key] = [] + self.team_model_to_deployment_indices[key].append(idx) + def upsert_deployment(self, deployment: Deployment) -> Optional[Deployment]: """ Add or update deployment @@ -8008,6 +8042,7 @@ class Router: instead of O(n) linear scan through the entire model_list. """ self.model_name_to_deployment_indices.clear() + self.team_model_to_deployment_indices.clear() for idx, model in enumerate(model_list): model_name = model.get("model_name") @@ -8016,6 +8051,16 @@ class Router: self.model_name_to_deployment_indices[model_name] = [] self.model_name_to_deployment_indices[model_name].append(idx) + team_id = model.get("model_info", {}).get("team_id") + team_public_model_name = model.get("model_info", {}).get( + "team_public_model_name" + ) + if team_id and team_public_model_name: + key = (team_id, team_public_model_name) + if key not in self.team_model_to_deployment_indices: + self.team_model_to_deployment_indices[key] = [] + self.team_model_to_deployment_indices[key].append(idx) + def _build_model_id_to_deployment_index_map(self, model_list: list): """ Build model index from model list to enable O(1) lookups immediately. @@ -8200,6 +8245,22 @@ class Router: """ returned_models: List[DeploymentTypedDict] = [] + # O(1) lookup in team_model index when team_id is provided + if team_id is not None: + key = (team_id, model_name) + if key in self.team_model_to_deployment_indices: + indices = self.team_model_to_deployment_indices[key] + # O(k) where k = team deployments for this model_name (typically 1-10) + for idx in indices: + model = self.model_list[idx] + if model_alias is not None: + alias_model = model.copy() + alias_model["model_name"] = model_alias + returned_models.append(alias_model) + else: + returned_models.append(model) + return returned_models + # O(1) lookup in model_name index if model_name in self.model_name_to_deployment_indices: indices = self.model_name_to_deployment_indices[model_name]