Merge pull request #20374 from ryan-crabbe/perf/cache-access-groups

perf: cache get_model_access_groups() no-args result on Router

litellm/router.py

@@ -471,6 +471,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:
@@ -6322,6 +6324,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:
@@ -6625,6 +6628,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:
@@ -6672,6 +6676,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
                         )
@@ -6705,6 +6710,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
                 )
@@ -7464,6 +7470,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", {})
@@ -7807,6 +7814,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,
@@ -7821,6 +7835,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)
@@ -7839,6 +7860,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(

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."""