diff --git a/litellm/router.py b/litellm/router.py index 46d35352c3..3a6c514989 100644 --- a/litellm/router.py +++ b/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. diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 5732deda6f..4bb9685f5e 100644 --- a/tests/test_litellm/test_router.py +++ b/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