tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-06 02:23:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: reduce get_deployment cost to O(1) (#14967)

Browse Source 
* fix: reduce get_deployment cost to O(1)

* fix: add unit test

* fix: cleaner

* fix: reference errors

* fix: add missing unit tests
This commit is contained in:
Alexsander Hamir
2025-09-27 13:44:10 -07:00
committed by GitHub
parent e270e0a797
commit a4eec173bc
3 changed files with 199 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/proxy/proxy_server.py
+6 -5
View File
@@ -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()
litellm/router.py
+88 -19
View File
@@ -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]:
tests/router_unit_tests/test_router_index_management.py
+105
View File
@@ -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