scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-11 03:38:38 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 82

[Fix] - Router: add model_name index for O(1) deployment lookups (#15113)

Browse source 
* perf(router): add model_name index for O(1) deployment lookups

Add model_name_to_deployment_indices mapping to optimize _get_all_deployments()
from O(n) to O(1) + O(k) lookups.

- Add model_name_to_deployment_indices: Dict[str, List[int]]
- Add _build_model_name_index() to build/maintain the index
- Update _add_model_to_list_and_index_map() to maintain both indices
- Refactor to use idx = len(self.model_list) before append (cleaner)
- Optimize _get_all_deployments() to use index instead of linear scan

* test(router): add test coverage for _build_model_name_index

Add single comprehensive test for _build_model_name_index() function to fix
code coverage CI failure.

The test verifies:
- Index correctly maps model_name to deployment indices
- Handles multiple deployments per model_name
- Clears and rebuilds index correctly

Fixes: CI code coverage error for _build_model_name_index
This commit is contained in:
Alexsander Hamir 2025-10-06 08:14:11 -07:00 • committed by GitHub
parent 9383ffccd6
commit ddb90c9ad7
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 106 additions and 17 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

71
litellm/router.py
View file

@ -416,6 +416,9 @@ class Router:
# Initialize model ID to deployment index mapping for O(1) lookups
self.model_id_to_deployment_index_map: Dict[str, int] = {}
# 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]] = {}
if model_list is not None:
# Build model index immediately to enable O(1) lookups from the start
@ -5097,6 +5100,7 @@ class Router:
original_model_list = copy.deepcopy(model_list)
self.model_list = []
self.model_id_to_deployment_index_map = {} # Reset the index
self.model_name_to_deployment_indices = {} # Reset the model_name 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:
@ -5138,6 +5142,9 @@ class Router:
f"\nInitialized Model List {self.get_model_names()}"
)
self.model_names = [m["model_name"] for m in model_list]
# Build model_name index for O(1) lookups
self._build_model_name_index(self.model_list)
def _add_deployment(self, deployment: Deployment) -> Deployment:
import os
@ -5365,20 +5372,27 @@ class Router:
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.
Helper method to add a model to the model_list and update both indices.
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"]
"""
idx = len(self.model_list)
self.model_list.append(model)
# Update model index for O(1) lookup
# Update model_id 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
self.model_id_to_deployment_index_map[model_id] = idx
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
)
self.model_id_to_deployment_index_map[model["model_info"]["id"]] = idx
# Update model_name index for O(1) lookup
model_name = model.get("model_name")
if model_name:
if model_name not in self.model_name_to_deployment_indices:
self.model_name_to_deployment_indices[model_name] = []
self.model_name_to_deployment_indices[model_name].append(idx)
def upsert_deployment(self, deployment: Deployment) -> Optional[Deployment]:
"""
@ -6094,6 +6108,22 @@ class Router:
additional_headers[header] = value
return response
def _build_model_name_index(self, model_list: list) -> None:
"""
Build model_name -> deployment indices mapping for O(1) lookups.
This index allows us to find all deployments for a given model_name in O(1) time
instead of O(n) linear scan through the entire model_list.
"""
self.model_name_to_deployment_indices.clear()
for idx, model in enumerate(model_list):
model_name = model.get("model_name")
if model_name:
if model_name not in self.model_name_to_deployment_indices:
self.model_name_to_deployment_indices[model_name] = []
self.model_name_to_deployment_indices[model_name].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.
@ -6198,18 +6228,27 @@ class Router:
Used for accurate 'get_model_list'.
if team_id specified, only return team-specific models
Optimized with O(1) index lookup instead of O(n) linear scan.
"""
returned_models: List[DeploymentTypedDict] = []
for model in self.model_list:
if self.should_include_deployment(
model_name=model_name, model=model, team_id=team_id
):
if model_alias is not None:
alias_model = copy.deepcopy(model)
alias_model["model_name"] = model_alias
returned_models.append(alias_model)
else:
returned_models.append(model)
# 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]
# O(k) where k = deployments for this model_name (typically 1-10)
for idx in indices:
model = self.model_list[idx]
if self.should_include_deployment(
model_name=model_name, model=model, team_id=team_id
):
if model_alias is not None:
alias_model = copy.deepcopy(model)
alias_model["model_name"] = model_alias
returned_models.append(alias_model)
else:
returned_models.append(model)
return returned_models

52
tests/router_unit_tests/test_router_index_management.py
View file

@ -77,7 +77,6 @@ class TestRouterIndexManagement:
# 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
@ -127,3 +126,54 @@ class TestRouterIndexManagement:
# Test: Empty router
empty_router = Router(model_list=[])
assert empty_router.has_model_id("any-id") == False
def test_build_model_name_index(self, router):
"""Test _build_model_name_index 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"},
},
{
"model_name": "gpt-4", # Duplicate model_name, different deployment
"litellm_params": {"model": "gpt-4"},
"model_info": {"id": "model-3"},
},
]
# Test: Build index from model list
router._build_model_name_index(model_list)
# Verify: model_name_to_deployment_indices is correctly built
assert "gpt-3.5-turbo" in router.model_name_to_deployment_indices
assert "gpt-4" in router.model_name_to_deployment_indices
# Verify: gpt-3.5-turbo has single deployment
assert router.model_name_to_deployment_indices["gpt-3.5-turbo"] == [0]
# Verify: gpt-4 has multiple deployments
assert router.model_name_to_deployment_indices["gpt-4"] == [1, 2]
# Test: Rebuild index (should clear and rebuild)
new_model_list = [
{
"model_name": "claude-3",
"litellm_params": {"model": "claude-3"},
"model_info": {"id": "model-4"},
},
]
router._build_model_name_index(new_model_list)
# Verify: Old entries are cleared
assert "gpt-3.5-turbo" not in router.model_name_to_deployment_indices
assert "gpt-4" not in router.model_name_to_deployment_indices
# Verify: New entry is added
assert "claude-3" in router.model_name_to_deployment_indices
assert router.model_name_to_deployment_indices["claude-3"] == [0]