From a4eec173bccc8067a3844900f45b6e7a3bf3d8f4 Mon Sep 17 00:00:00 2001 From: Alexsander Hamir Date: Sat, 27 Sep 2025 13:44:10 -0700 Subject: [PATCH] fix: reduce get_deployment cost to O(1) (#14967) * fix: reduce get_deployment cost to O(1) * fix: add unit test * fix: cleaner * fix: reference errors * fix: add missing unit tests --- litellm/proxy/proxy_server.py | 11 +- litellm/router.py | 107 ++++++++++++++---- .../test_router_index_management.py | 105 +++++++++++++++++ 3 files changed, 199 insertions(+), 24 deletions(-) create mode 100644 tests/router_unit_tests/test_router_index_management.py diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 6d2e56c444..f031b9ffbd 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -637,11 +637,6 @@ async def proxy_startup_event(app: FastAPI): user_api_key_cache=user_api_key_cache, ) - if use_background_health_checks: - asyncio.create_task( - _run_background_health_check() - ) # start the background health check coroutine. - if prompt_injection_detection_obj is not None: # [TODO] - REFACTOR THIS prompt_injection_detection_obj.update_environment(router=llm_router) @@ -664,6 +659,12 @@ async def proxy_startup_event(app: FastAPI): await ProxyStartupEvent._update_default_team_member_budget() + # Start background health checks AFTER models are loaded and index is built + if use_background_health_checks: + asyncio.create_task( + _run_background_health_check() + ) # start the background health check coroutine. + ## [Optional] Initialize dd tracer ProxyStartupEvent._init_dd_tracer() diff --git a/litellm/router.py b/litellm/router.py index f093096f84..a1cfc22e7d 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -409,7 +409,12 @@ class Router: ) # {"TEAM_ID": PatternMatchRouter} self.auto_routers: Dict[str, "AutoRouter"] = {} + # Initialize model ID to deployment index mapping for O(1) lookups + self.model_id_to_deployment_index_map: Dict[str, int] = {} + if model_list is not None: + # Build model index immediately to enable O(1) lookups from the start + self._build_model_id_to_deployment_index_map(model_list) model_list = copy.deepcopy(model_list) self.set_model_list(model_list) self.healthy_deployments: List = self.model_list # type: ignore @@ -4974,7 +4979,7 @@ class Router: model = deployment.to_json(exclude_none=True) - self.model_list.append(model) + self._add_model_to_list_and_index_map(model=model, model_id=deployment.model_info.id) return deployment except Exception as e: if self.ignore_invalid_deployments: @@ -5085,6 +5090,7 @@ class Router: def set_model_list(self, model_list: list): original_model_list = copy.deepcopy(model_list) self.model_list = [] + self.model_id_to_deployment_index_map = {} # Reset the index # we add api_base/api_key each model so load balancing between azure/gpt on api_base1 and api_base2 works for model in original_model_list: @@ -5323,10 +5329,42 @@ class Router: self._add_deployment(deployment=deployment) # add to model names - self.model_list.append(_deployment) + self._add_model_to_list_and_index_map(model=_deployment, model_id=deployment.model_info.id) self.model_names.append(deployment.model_name) return deployment + def _update_deployment_indices_after_removal(self, model_id: str, removal_idx: int) -> None: + """ + Helper method to update deployment indices after a deployment has been removed from model_list. + + Parameters: + - model_id: str - the id of the deployment that was removed + - removal_idx: int - the index where the deployment was removed from model_list + """ + # Update indices for all models after the removed one + for deployment_id, idx in self.model_id_to_deployment_index_map.items(): + if idx > removal_idx: + self.model_id_to_deployment_index_map[deployment_id] = idx - 1 + # Remove the deleted model from index + if model_id in self.model_id_to_deployment_index_map: + del self.model_id_to_deployment_index_map[model_id] + + + def _add_model_to_list_and_index_map(self, model: dict, model_id: Optional[str] = None) -> None: + """ + Helper method to add a model to the model_list and update the model_id_to_deployment_index_map. + + Parameters: + - model: dict - the model to add to the list + - model_id: Optional[str] - the model ID to use for indexing. If None, will try to get from model["model_info"]["id"] + """ + self.model_list.append(model) + # Update model index for O(1) lookup + if model_id is not None: + self.model_id_to_deployment_index_map[model_id] = len(self.model_list) - 1 + elif model.get("model_info", {}).get("id") is not None: + self.model_id_to_deployment_index_map[model["model_info"]["id"]] = len(self.model_list) - 1 + def upsert_deployment(self, deployment: Deployment) -> Optional[Deployment]: """ Add or update deployment @@ -5352,12 +5390,15 @@ class Router: # if there is a new litellm param -> then update the deployment # remove the previous deployment removal_idx: Optional[int] = None - for idx, model in enumerate(self.model_list): - if model["model_info"]["id"] == deployment.model_info.id: - removal_idx = idx + deployment_id = deployment.model_info.id + deployment_fast_mapping = self.model_id_to_deployment_index_map + + if deployment_id in deployment_fast_mapping: + removal_idx = deployment_fast_mapping[deployment_id] - if removal_idx is not None: - self.model_list.pop(removal_idx) + if removal_idx is not None: + self.model_list.pop(removal_idx) + self._update_deployment_indices_after_removal(model_id=deployment_id, removal_idx=removal_idx) # if the model_id is not in router self.add_deployment(deployment=deployment) @@ -5381,13 +5422,14 @@ class Router: - OR None (if deleted deployment not found) """ deployment_idx = None - for idx, m in enumerate(self.model_list): - if m["model_info"]["id"] == id: - deployment_idx = idx + if id in self.model_id_to_deployment_index_map: + deployment_idx = self.model_id_to_deployment_index_map[id] try: if deployment_idx is not None: + # Pop the item from the list first item = self.model_list.pop(deployment_idx) + self._update_deployment_indices_after_removal(model_id=id, removal_idx=deployment_idx) return item else: return None @@ -5400,15 +5442,17 @@ class Router: Raise Exception -> if model found in invalid format """ - for model in self.model_list: - if "model_info" in model and "id" in model["model_info"]: - if model_id == model["model_info"]["id"]: - if isinstance(model, dict): - return Deployment(**model) - elif isinstance(model, Deployment): - return model - else: - raise Exception("Model invalid format - {}".format(type(model))) + # Use O(1) lookup via model_id_to_deployment_index_map only + if model_id in self.model_id_to_deployment_index_map: + idx = self.model_id_to_deployment_index_map[model_id] + model = self.model_list[idx] + if isinstance(model, dict): + return Deployment(**model) + elif isinstance(model, Deployment): + return model + else: + raise Exception("Model invalid format - {}".format(type(model))) + return None def get_deployment_credentials(self, model_id: str) -> Optional[dict]: @@ -6026,6 +6070,31 @@ class Router: additional_headers[header] = value return response + def _build_model_id_to_deployment_index_map(self, model_list: list): + """ + Build model index from model list to enable O(1) lookups immediately. + This is called during initialization to avoid the race condition where + requests arrive before model_id_to_deployment_index_map is populated. + """ + # First populate the model_list + self.model_list = [] + for _, model in enumerate(model_list): + # Extract model_info from the model dict + model_info = model.get("model_info", {}) + model_id = model_info.get("id") + + # If no ID exists, generate one using the same logic as set_model_list + if model_id is None: + model_name = model.get("model_name", "") + litellm_params = model.get("litellm_params", {}) + model_id = self._generate_model_id(model_name, litellm_params) + # Update the model_info in the original list + if "model_info" not in model: + model["model_info"] = {} + model["model_info"]["id"] = model_id + + self._add_model_to_list_and_index_map(model=model, model_id=model_id) + def get_model_ids( self, model_name: Optional[str] = None, exclude_team_models: bool = False ) -> List[str]: diff --git a/tests/router_unit_tests/test_router_index_management.py b/tests/router_unit_tests/test_router_index_management.py new file mode 100644 index 0000000000..bdd52f66ba --- /dev/null +++ b/tests/router_unit_tests/test_router_index_management.py @@ -0,0 +1,105 @@ +import sys +import os +import pytest + +sys.path.insert( + 0, os.path.abspath("../..") +) # Adds the parent directory to the system path +from litellm import Router + + +class TestRouterIndexManagement: + """Test cases for router index management functions""" + + @pytest.fixture + def router(self): + """Create a router instance for testing""" + return Router(model_list=[]) + + def test_update_deployment_indices_after_removal(self, router): + """Test _update_deployment_indices_after_removal function""" + # Setup: Add models to router with proper structure + router.model_list = [ + {"model": "test1", "model_info": {"id": "model-1"}}, + {"model": "test2", "model_info": {"id": "model-2"}}, + {"model": "test3", "model_info": {"id": "model-3"}} + ] + router.model_id_to_deployment_index_map = {"model-1": 0, "model-2": 1, "model-3": 2} + + # Test: Remove model-2 (index 1) + router._update_deployment_indices_after_removal(model_id="model-2", removal_idx=1) + + # Verify: model-2 is removed from index + assert "model-2" not in router.model_id_to_deployment_index_map + # Verify: model-3 index is updated (2 -> 1) + assert router.model_id_to_deployment_index_map["model-3"] == 1 + # Verify: model-1 index remains unchanged + assert router.model_id_to_deployment_index_map["model-1"] == 0 + + def test_build_model_id_to_deployment_index_map(self, router): + """Test _build_model_id_to_deployment_index_map function""" + model_list = [ + { + "model_name": "gpt-3.5-turbo", + "litellm_params": {"model": "gpt-3.5-turbo"}, + "model_info": {"id": "model-1"}, + }, + { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "model-2"}, + }, + ] + + # Test: Build index from model list + router._build_model_id_to_deployment_index_map(model_list) + + # Verify: model_list is populated + assert len(router.model_list) == 2 + # Verify: model_id_to_deployment_index_map is correctly built + assert router.model_id_to_deployment_index_map["model-1"] == 0 + assert router.model_id_to_deployment_index_map["model-2"] == 1 + + def test_add_model_to_list_and_index_map_from_model_info(self, router): + """Test _add_model_to_list_and_index_map extracting model_id from model_info""" + # Setup: Empty router + router.model_list = [] + router.model_id_to_deployment_index_map = {} + + # Test: Add model without explicit model_id + model = {"model": "test-model", "model_info": {"id": "model-info-id"}} + router._add_model_to_list_and_index_map(model=model) + + # Verify: Model added to list + assert len(router.model_list) == 1 + assert router.model_list[0] == model + + # Verify: Index map uses model_info.id + assert router.model_id_to_deployment_index_map["model-info-id"] == 0 + + + def test_add_model_to_list_and_index_map_multiple_models(self, router): + """Test _add_model_to_list_and_index_map with multiple models to verify indexing""" + # Setup: Empty router + router.model_list = [] + router.model_id_to_deployment_index_map = {} + + # Test: Add multiple models + model1 = {"model": "model1", "model_info": {"id": "id-1"}} + model2 = {"model": "model2", "model_info": {"id": "id-2"}} + model3 = {"model": "model3", "model_info": {"id": "id-3"}} + + router._add_model_to_list_and_index_map(model=model1, model_id="id-1") + router._add_model_to_list_and_index_map(model=model2, model_id="id-2") + router._add_model_to_list_and_index_map(model=model3, model_id="id-3") + + # Verify: All models added to list + assert len(router.model_list) == 3 + assert router.model_list[0] == model1 + assert router.model_list[1] == model2 + assert router.model_list[2] == model3 + + # Verify: Correct indices in map + assert router.model_id_to_deployment_index_map["id-1"] == 0 + assert router.model_id_to_deployment_index_map["id-2"] == 1 + assert router.model_id_to_deployment_index_map["id-3"] == 2