Merge pull request #15574 from BerriAI/litellm_router_index_change

perf(router): optimize model lookups with O(1) index maps

litellm/router.py

@@ -5522,9 +5522,15 @@ class Router:
         Returns -> Deployment or None
 
         Raise Exception -> if model found in invalid format
+        
+        Optimized with O(1) index lookup instead of O(n) linear scan.
         """
-        for model in self.model_list:
-            if model["model_name"] == model_group_name:
+        # O(1) lookup in model_name index
+        if model_group_name in self.model_name_to_deployment_indices:
+            indices = self.model_name_to_deployment_indices[model_group_name]
+            if indices:
+                # Return first deployment for this model_name
+                model = self.model_list[indices[0]]
                 if isinstance(model, dict):
                     return Deployment(**model)
                 elif isinstance(model, Deployment):
@@ -5634,11 +5640,13 @@ class Router:
         Returns
         - dict: the model in list with 'model_name', 'litellm_params', Optional['model_info']
         - None: could not find deployment in list
+        
+        Optimized with O(1) index lookup instead of O(n) linear scan.
         """
-        for model in self.model_list:
-            if "model_info" in model and "id" in model["model_info"]:
-                if id == model["model_info"]["id"]:
-                    return model
+        # O(1) lookup via model_id_to_deployment_index_map
+        if id in self.model_id_to_deployment_index_map:
+            idx = self.model_id_to_deployment_index_map[id]
+            return self.model_list[idx]
         return None
 
     def get_model_group(self, id: str) -> Optional[List]:
@@ -6172,17 +6180,33 @@ class Router:
         if 'model_name' is none, returns all.
 
         Returns list of model id's.
+        
+        Optimized with O(1) or O(k) index lookup when model_name provided,
+        instead of O(n) linear scan.
         """        
         ids = []
-        for model in self.model_list:
-            if "model_info" in model and "id" in model["model_info"]:
-                id = model["model_info"]["id"]
-                if exclude_team_models and model["model_info"].get("team_id"):
-                    continue
-                if model_name is not None and model["model_name"] == model_name:
-                    ids.append(id)
-                elif model_name is None:
-                    ids.append(id)
+        
+        if model_name is not None:
+            # O(1) lookup in model_name index, then O(k) iteration where k = deployments for this model_name
+            if model_name in self.model_name_to_deployment_indices:
+                indices = self.model_name_to_deployment_indices[model_name]
+                for idx in indices:
+                    model = self.model_list[idx]
+                    if "model_info" in model and "id" in model["model_info"]:
+                        if exclude_team_models and model["model_info"].get("team_id"):
+                            continue
+                        ids.append(model["model_info"]["id"])
+        else:
+            # When model_name is None, return all model IDs
+            # Use the index map keys for O(n) where n = total deployments
+            for model_id in self.model_id_to_deployment_index_map.keys():
+                idx = self.model_id_to_deployment_index_map[model_id]
+                model = self.model_list[idx]
+                if "model_info" in model and "id" in model["model_info"]:
+                    if exclude_team_models and model["model_info"].get("team_id"):
+                        continue
+                    ids.append(model_id)
+        
         return ids
 
     def has_model_id(self, candidate_id: str) -> bool:

tests/router_unit_tests/test_router_index_management.py

@@ -1,6 +1,7 @@
 import sys
 import os
 import pytest
+import ast
 
 sys.path.insert(
     0, os.path.abspath("../..")
@@ -178,6 +179,92 @@ class TestRouterIndexManagement:
         assert "claude-3" in router.model_name_to_deployment_indices
         assert router.model_name_to_deployment_indices["claude-3"] == [0]
 
+    def test_no_linear_scans_in_router(self):
+        """
+        Static analysis test to ensure Router doesn't use O(n) linear scans.
+        
+        Scans router.py for 'in self.model_list' pattern which indicates
+        inefficient O(n) iteration instead of using index-based O(1) lookups.
+        
+        Methods should use:
+        - model_id_to_deployment_index_map for O(1) model_id lookups
+        - model_name_to_deployment_indices for O(1) + O(k) model_name lookups
+        """
+        # Methods that are allowed to iterate through self.model_list
+        ALLOWED_METHODS = [
+            "_get_deployment_by_litellm_model",  # Edge case: lookup by litellm_params.model (not indexed)
+        ]
+        
+        # Get path to router.py
+        router_file = os.path.join(
+            os.path.dirname(os.path.dirname(os.path.dirname(__file__))),
+            "litellm",
+            "router.py"
+        )
+        
+        # Read the file
+        with open(router_file, 'r') as f:
+            content = f.read()
+        
+        # Parse with AST
+        tree = ast.parse(content)
+        
+        # Find violations
+        violations = []
+        ignore_methods = set(ALLOWED_METHODS)
+        
+        for node in ast.walk(tree):
+            if isinstance(node, ast.FunctionDef):
+                method_name = node.name
+                
+                # Skip ignored methods
+                if method_name in ignore_methods:
+                    continue
+                
+                # Get source for this method
+                try:
+                    method_source = ast.get_source_segment(content, node)
+                    if not method_source:
+                        continue
+                    
+                    # Check for the anti-pattern: "in self.model_list"
+                    # This catches: for x in self.model_list, if x in self.model_list, etc.
+                    if "in self.model_list" in method_source:
+                        # Extract the specific line for better error reporting
+                        lines = method_source.split('\n')
+                        pattern_line = None
+                        for line in lines:
+                            if "in self.model_list" in line:
+                                pattern_line = line.strip()
+                                break
+                        
+                        violations.append({
+                            "method": method_name,
+                            "line": node.lineno,
+                            "pattern": pattern_line or "in self.model_list"
+                        })
+                except Exception:
+                    # Skip if we can't get source segment
+                    pass
+        
+        # Assert no violations
+        if violations:
+            error_msg = "\n".join([
+                f"  - {v['method']}() at line {v['line']}: {v['pattern']}"
+                for v in violations
+            ])
+            
+            pytest.fail(
+                f"\n{'='*70}\n"
+                f"Found O(n) linear scan pattern in router.py:\n\n"
+                f"{error_msg}\n\n"
+                f"These methods should use index maps instead:\n"
+                f"  - model_id_to_deployment_index_map (for model_id lookups)\n"
+                f"  - model_name_to_deployment_indices (for model_name lookups)\n\n"
+                f"If a method legitimately needs O(n) iteration, add it to\n"
+                f"ALLOWED_METHODS in this test method.\n"
+                f"{'='*70}\n"
+            )
     def test_model_names_is_set(self):
         """Verify that model_names uses a set for O(1) lookups, not a list (O(n))"""
         router = Router(model_list=[])