From 1517a70e01bf18e495150a906ba5e0fed09ca646 Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Tue, 3 Feb 2026 16:00:24 -0800 Subject: [PATCH] perf: cache get_model_access_groups() no-args result on Router MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The no-args hot path (called on every proxy request) was rebuilding a defaultdict by iterating the full model list each time. Cache the result and invalidate at all 5 model_list mutation sites following the _invalidate_model_cost_lowercase_map() pattern. Line profile: 31.8µs/call → 1.1µs/call (29x improvement). --- litellm/router.py | 26 ++++++ tests/test_litellm/test_router.py | 126 ++++++++++++++++++++++++++++++ 2 files changed, 152 insertions(+) diff --git a/litellm/router.py b/litellm/router.py index d01c8443da..6fa6ecae93 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -470,6 +470,8 @@ class Router: [] ) # initialize an empty list - to allow _add_deployment and delete_deployment to work + self._access_groups_cache: Optional[Dict[str, List[str]]] = None + if allowed_fails is not None: self.allowed_fails = allowed_fails else: @@ -6055,6 +6057,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._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 for model in original_model_list: @@ -6358,6 +6361,7 @@ class Router: """ idx = len(self.model_list) self.model_list.append(model) + self._invalidate_access_groups_cache() # Update model_id index for O(1) lookup if model_id is not None: @@ -6405,6 +6409,7 @@ class Router: if removal_idx is not None: self.model_list.pop(removal_idx) + self._invalidate_access_groups_cache() self._update_deployment_indices_after_removal( model_id=deployment_id, removal_idx=removal_idx ) @@ -6438,6 +6443,7 @@ class Router: if deployment_idx is not None: # Pop the item from the list first item = self.model_list.pop(deployment_idx) + self._invalidate_access_groups_cache() self._update_deployment_indices_after_removal( model_id=id, removal_idx=deployment_idx ) @@ -7172,6 +7178,7 @@ class Router: """ # First populate the model_list self.model_list = [] + self._invalidate_access_groups_cache() for _, model in enumerate(model_list): # Extract model_info from the model dict model_info = model.get("model_info", {}) @@ -7508,6 +7515,13 @@ class Router: return returned_models + def _invalidate_access_groups_cache(self) -> None: + """Invalidate the cached access groups. + + Call this whenever self.model_list is modified to ensure the cache is rebuilt. + """ + self._access_groups_cache = None + def get_model_access_groups( self, model_name: Optional[str] = None, @@ -7522,6 +7536,13 @@ class Router: - model_access_group: Optional[str] - the received model access group from the user. If set, will only return models for that access group. - team_id: Optional[str] - the team id, to resolve team-specific models """ + # Check if this is the no-args hot path (cacheable) + _use_cache = model_name is None and model_access_group is None and team_id is None + + # Return cached result for the no-args hot path + if _use_cache and self._access_groups_cache is not None: + return self._access_groups_cache + from collections import defaultdict access_groups = defaultdict(list) @@ -7540,6 +7561,11 @@ class Router: model_name = m["model_name"] access_groups[group].append(model_name) + # Cache the result for the no-args hot path + if _use_cache: + self._access_groups_cache = dict(access_groups) + return self._access_groups_cache + return access_groups def _is_model_access_group_for_wildcard_route( diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 08ae804ea8..7f06b69995 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -925,6 +925,132 @@ def test_router_get_model_access_groups_team_only_models(): assert list(access_groups.keys()) == ["default-models"] +def test_get_model_access_groups_caching(): + """ + Test that get_model_access_groups caches the no-args result + and invalidates on deployment changes. + """ + from litellm.types.router import Deployment, LiteLLM_Params + + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4"}, + "model_info": {"access_groups": ["premium"]}, + }, + ] + ) + + # First call computes and populates cache + result1 = router.get_model_access_groups() + assert "premium" in result1 + + # All subsequent calls should return the same cached object (including first) + result2 = router.get_model_access_groups() + assert result1 is result2 + + # Calls with args should bypass cache + result_with_args = router.get_model_access_groups(model_name="gpt-4") + assert result_with_args is not result2 + + # Add a deployment — cache should be invalidated + router.add_deployment( + Deployment( + model_name="gpt-3.5", + litellm_params=LiteLLM_Params(model="gpt-3.5-turbo"), + model_info={"access_groups": ["default"]}, + ) + ) + result3 = router.get_model_access_groups() + assert result3 is not result2 + assert "premium" in result3 + assert "default" in result3 + + # Delete the deployment — cache should be invalidated again + deployment_id = None + for m in router.model_list: + if m.get("model_name") == "gpt-3.5": + deployment_id = m.get("model_info", {}).get("id") + break + assert deployment_id is not None + router.delete_deployment(id=deployment_id) + result4 = router.get_model_access_groups() + assert result4 is not result3 + assert "default" not in result4 + assert "premium" in result4 + + +def test_get_model_access_groups_cache_invalidation_set_model_list(): + """ + Test that set_model_list invalidates the access groups cache. + """ + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4"}, + "model_info": {"access_groups": ["premium"]}, + }, + ] + ) + + # Populate cache + result1 = router.get_model_access_groups() + assert "premium" in result1 + + # set_model_list should invalidate cache + router.set_model_list( + [ + { + "model_name": "claude-3", + "litellm_params": {"model": "anthropic/claude-3-opus-20240229"}, + "model_info": {"access_groups": ["research"]}, + }, + ] + ) + result2 = router.get_model_access_groups() + assert result2 is not result1 + assert "research" in result2 + assert "premium" not in result2 + + +def test_get_model_access_groups_cache_invalidation_upsert_deployment(): + """ + Test that upsert_deployment invalidates the access groups cache. + """ + from litellm.types.router import Deployment, LiteLLM_Params + + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4"}, + "model_info": {"access_groups": ["premium"]}, + }, + ] + ) + + # Populate cache + result1 = router.get_model_access_groups() + assert "premium" in result1 + + # Get the existing deployment's ID + existing_id = router.model_list[0]["model_info"]["id"] + + # Upsert with the same ID but different params — triggers pop + re-add + router.upsert_deployment( + Deployment( + model_name="gpt-4-updated", + litellm_params=LiteLLM_Params(model="gpt-4-turbo"), + model_info={"id": existing_id, "access_groups": ["updated-group"]}, + ) + ) + result2 = router.get_model_access_groups() + assert result2 is not result1 + assert "updated-group" in result2 + + @pytest.mark.asyncio async def test_acompletion_streaming_iterator(): """Test _acompletion_streaming_iterator for normal streaming and fallback behavior."""