Merge pull request #20376 from ryan-crabbe/fix/model-group-info-cache-invalidation

fix: add cache invalidation for _cached_get_model_group_info

litellm/router.py

@@ -6423,6 +6423,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_model_group_info_cache()
         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
 
@@ -6733,6 +6734,7 @@ class Router:
         """
         idx = len(self.model_list)
         self.model_list.append(model)
+        self._invalidate_model_group_info_cache()
         self._invalidate_access_groups_cache()
 
         # Update model_id index for O(1) lookup
@@ -6781,6 +6783,7 @@ class Router:
 
                     if removal_idx is not None:
                         self.model_list.pop(removal_idx)
+                        self._invalidate_model_group_info_cache()
                         self._invalidate_access_groups_cache()
                         self._update_deployment_indices_after_removal(
                             model_id=deployment_id, removal_idx=removal_idx
@@ -6815,6 +6818,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_model_group_info_cache()
                 self._invalidate_access_groups_cache()
                 self._update_deployment_indices_after_removal(
                     model_id=id, removal_idx=deployment_idx
@@ -7576,6 +7580,7 @@ class Router:
         """
         # First populate the model_list
         self.model_list = []
+        self._invalidate_model_group_info_cache()
         self._invalidate_access_groups_cache()
         for _, model in enumerate(model_list):
             # Extract model_info from the model dict
@@ -7920,6 +7925,13 @@ class Router:
 
         return returned_models
 
+    def _invalidate_model_group_info_cache(self) -> None:
+        """Invalidate the cached model group info.
+
+        Call this whenever self.model_list is modified to ensure the cache is rebuilt.
+        """
+        self._cached_get_model_group_info.cache_clear()
+
     def _invalidate_access_groups_cache(self) -> None:
         """Invalidate the cached access groups.
 

tests/test_litellm/test_router.py

@@ -925,6 +925,73 @@ def test_router_get_model_access_groups_team_only_models():
     assert list(access_groups.keys()) == ["default-models"]
 
 
+def test_cached_get_model_group_info():
+    """
+    Test that _cached_get_model_group_info caches results 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", "api_key": "fake"},
+                "model_info": {"tpm": 1000, "rpm": 100},
+            },
+        ]
+    )
+
+    # First call should compute and cache
+    result1 = router._cached_get_model_group_info("gpt-4")
+    assert result1 is not None
+    assert result1.tpm == 1000
+
+    # Second call should hit cache (same object)
+    result2 = router._cached_get_model_group_info("gpt-4")
+    assert result1 is result2
+
+    # Add a deployment — cache should be invalidated
+    router.add_deployment(
+        Deployment(
+            model_name="gpt-4",
+            litellm_params=LiteLLM_Params(model="gpt-4", api_key="fake2"),
+            model_info={"tpm": 2000, "rpm": 200},
+        )
+    )
+    result3 = router._cached_get_model_group_info("gpt-4")
+    assert result3 is not result2
+    assert result3 is not None
+    assert result3.tpm == 3000  # 1000 + 2000
+
+    # Delete a deployment — cache should be invalidated
+    deployment_id = router.model_list[-1]["model_info"]["id"]
+    router.delete_deployment(id=deployment_id)
+    result4 = router._cached_get_model_group_info("gpt-4")
+    assert result4 is not result3
+    assert result4 is not None
+    assert result4.tpm == 1000
+
+    # set_model_list — cache should be invalidated
+    router.set_model_list(
+        [
+            {
+                "model_name": "gpt-4",
+                "litellm_params": {"model": "gpt-4", "api_key": "fake"},
+                "model_info": {"tpm": 5000},
+            },
+        ]
+    )
+    result5 = router._cached_get_model_group_info("gpt-4")
+    assert result5 is not result4
+    assert result5 is not None
+    assert result5.tpm == 5000
+
+    # Verify cache still works after invalidation
+    result6 = router._cached_get_model_group_info("gpt-4")
+    assert result5 is result6
+
+
 def test_get_model_access_groups_caching():
     """
     Test that get_model_access_groups caches the no-args result