fix(router): enforce model tpm limits against shared redis usage across replicas

The model tpm pre-call check read only the in-memory counter, so each proxy replica enforced the limit against its own traffic and the deployment admitted up to N times the configured tpm across N replicas. Read the shared Redis counter when the local counter is under the limit, keep the local counter authoritative when it is already at the limit, and fall back to local usage when Redis is unavailable

Supersedes #40854, Fixes #40291

Co-authored-by: Jahanzeb-git <jahanzebahmed2002@gmail.com>
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm/router_utils/pre_call_checks/model_rate_limit_check.py

@@ -18,6 +18,7 @@ import httpx
 import litellm
 from litellm._logging import verbose_router_logger
 from litellm.caching.dual_cache import DualCache
+from litellm.caching.redis_cache import RedisCircuitBreakerOpenError
 from litellm.integrations.custom_logger import CustomLogger
 from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import (
     ITPM_RESERVED_KEY,
@@ -136,6 +137,26 @@ class ModelRateLimitingCheck(CustomLogger):
 
         return tpm_key, rpm_key
 
+    def _get_current_tpm(self, tpm_key: str, tpm_limit: int) -> int | None:
+        local_tpm: Final = self.dual_cache.get_cache(key=tpm_key, local_only=True)
+        redis_cache: Final = self.dual_cache.redis_cache
+        if redis_cache is None or (local_tpm is not None and local_tpm >= tpm_limit):
+            return local_tpm
+        try:
+            return redis_cache.get_cache(key=tpm_key)
+        except RedisCircuitBreakerOpenError:
+            return local_tpm
+
+    async def _async_get_current_tpm(self, tpm_key: str, tpm_limit: int, parent_otel_span: Span | None) -> int | None:
+        local_tpm: Final = await self.dual_cache.async_get_cache(key=tpm_key, local_only=True)
+        redis_cache: Final = self.dual_cache.redis_cache
+        if redis_cache is None or (local_tpm is not None and local_tpm >= tpm_limit):
+            return local_tpm
+        try:
+            return await redis_cache.async_get_cache(key=tpm_key, parent_otel_span=parent_otel_span)
+        except RedisCircuitBreakerOpenError:
+            return local_tpm
+
     def pre_call_check(self, deployment: dict) -> dict | None:
         """
         Synchronous pre-call check for model rate limits.
@@ -168,8 +189,7 @@ class ModelRateLimitingCheck(CustomLogger):
 
             # Check TPM limit
             if tpm_limit is not None:
-                # First check local cache
-                current_tpm: Final = self.dual_cache.get_cache(key=tpm_key, local_only=True)
+                current_tpm: Final = self._get_current_tpm(tpm_key, tpm_limit)
                 if current_tpm is not None and current_tpm >= tpm_limit:
                     raise litellm.RateLimitError(
                         message=f"Model rate limit exceeded. TPM limit={tpm_limit}, current usage={current_tpm}",
@@ -249,8 +269,7 @@ class ModelRateLimitingCheck(CustomLogger):
 
             # Check TPM limit
             if tpm_limit is not None:
-                # First check local cache
-                current_tpm: Final = await self.dual_cache.async_get_cache(key=tpm_key, local_only=True)
+                current_tpm: Final = await self._async_get_current_tpm(tpm_key, tpm_limit, parent_otel_span)
                 if current_tpm is not None and current_tpm >= tpm_limit:
                     raise litellm.RateLimitError(
                         message=f"Model rate limit exceeded. TPM limit={tpm_limit}, current usage={current_tpm}",

tests/test_litellm/test_router/test_enforce_model_rate_limits.py

@@ -6,6 +6,7 @@ regardless of the routing strategy being used.
 """
 
 import asyncio
+from datetime import timedelta
 from unittest.mock import AsyncMock, MagicMock
 
 import pytest
@@ -13,10 +14,30 @@ import pytest
 import litellm
 from litellm import Router
 from litellm.caching.dual_cache import DualCache
+from litellm.caching.redis_cache import RedisCircuitBreakerOpenError
 from litellm.router_utils.pre_call_checks.model_rate_limit_check import (
     ModelRateLimitingCheck,
 )
 
+TPM_DEPLOYMENT = {
+    "tpm": 1000,
+    "litellm_params": {"model": "gpt-4"},
+    "model_info": {"id": "replica-test-id"},
+    "model_name": "test-model",
+}
+
+
+def _dual_cache_with_local_tpm(local_tpm: int, redis_cache: MagicMock | None) -> DualCache:
+    """In-memory tier holds ``local_tpm`` for this replica; the key is primed for this minute and the next
+    so a minute rollover between priming and the check cannot make the read miss."""
+    dual_cache = DualCache(redis_cache=redis_cache)
+    check = ModelRateLimitingCheck(dual_cache=dual_cache)
+    now = litellm.utils.get_utc_datetime()
+    for minute in (now, now + timedelta(minutes=1)):
+        tpm_key, _ = check._get_cache_keys(TPM_DEPLOYMENT, minute.strftime("%H-%M"))
+        dual_cache.set_cache(key=tpm_key, value=local_tpm, local_only=True)
+    return dual_cache
+
 
 class TestModelRateLimitingCheck:
     """Test the ModelRateLimitingCheck class directly."""
@@ -144,6 +165,52 @@ class TestModelRateLimitingCheck:
         assert "TPM limit=1000" in str(exc_info.value)
         assert "current usage=1000" in str(exc_info.value)
 
+    def test_pre_call_check_rejects_when_shared_tpm_is_over_limit_but_local_is_under(self):
+        """Another replica's usage in Redis must count even when this replica saw only a few tokens."""
+        redis_cache = MagicMock()
+        redis_cache.get_cache.return_value = 1000
+        check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(5, redis_cache))
+
+        with pytest.raises(litellm.RateLimitError) as exc_info:
+            check.pre_call_check(TPM_DEPLOYMENT)
+
+        assert "current usage=1000" in str(exc_info.value)
+
+    @pytest.mark.parametrize(
+        "redis_get", [MagicMock(return_value=None), MagicMock(side_effect=RedisCircuitBreakerOpenError())]
+    )
+    def test_pre_call_check_keeps_rejecting_on_local_usage_when_redis_read_fails(self, redis_get):
+        """A missing or failed Redis read must not admit traffic a replica already knows is over the limit."""
+        redis_cache = MagicMock()
+        redis_cache.get_cache = redis_get
+        check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(1000, redis_cache))
+
+        with pytest.raises(litellm.RateLimitError) as exc_info:
+            check.pre_call_check(TPM_DEPLOYMENT)
+
+        assert "current usage=1000" in str(exc_info.value)
+
+    def test_pre_call_check_falls_back_to_local_tpm_and_still_checks_rpm_when_redis_circuit_is_open(self):
+        redis_cache = MagicMock()
+        redis_cache.get_cache.side_effect = RedisCircuitBreakerOpenError()
+        redis_cache.increment_cache.return_value = 2
+        check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(5, redis_cache))
+
+        with pytest.raises(litellm.RateLimitError) as exc_info:
+            check.pre_call_check({**TPM_DEPLOYMENT, "rpm": 1})
+
+        assert "RPM limit=1" in str(exc_info.value)
+
+    def test_pre_call_check_without_redis_enforces_local_tpm_and_rpm(self):
+        check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(5, None))
+        deployment = {**TPM_DEPLOYMENT, "rpm": 1}
+
+        assert check.pre_call_check(deployment) == deployment
+        with pytest.raises(litellm.RateLimitError) as exc_info:
+            check.pre_call_check(deployment)
+
+        assert "RPM limit=1" in str(exc_info.value)
+
     def test_log_success_event_increments_cache(self):
         """Test that log_success_event correctly increments the cache."""
         mock_cache = MagicMock()
@@ -245,6 +312,58 @@ class TestModelRateLimitingCheckAsync:
 
         assert "TPM limit=1000" in str(exc_info.value)
 
+    @pytest.mark.asyncio
+    async def test_async_pre_call_check_rejects_when_shared_tpm_is_over_limit_but_local_is_under(self):
+        """Another replica's usage in Redis must count even when this replica saw only a few tokens."""
+        redis_cache = MagicMock()
+        redis_cache.async_get_cache = AsyncMock(return_value=1000)
+        check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(5, redis_cache))
+
+        with pytest.raises(litellm.RateLimitError) as exc_info:
+            await check.async_pre_call_check(TPM_DEPLOYMENT)
+
+        assert "current usage=1000" in str(exc_info.value)
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize(
+        "redis_get", [AsyncMock(return_value=None), AsyncMock(side_effect=RedisCircuitBreakerOpenError())]
+    )
+    async def test_async_pre_call_check_keeps_rejecting_on_local_usage_when_redis_read_fails(self, redis_get):
+        """A missing or failed Redis read must not admit traffic a replica already knows is over the limit."""
+        redis_cache = MagicMock()
+        redis_cache.async_get_cache = redis_get
+        check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(1000, redis_cache))
+
+        with pytest.raises(litellm.RateLimitError) as exc_info:
+            await check.async_pre_call_check(TPM_DEPLOYMENT)
+
+        assert "current usage=1000" in str(exc_info.value)
+
+    @pytest.mark.asyncio
+    async def test_async_pre_call_check_falls_back_to_local_tpm_and_still_checks_rpm_when_redis_circuit_is_open(
+        self,
+    ):
+        redis_cache = MagicMock()
+        redis_cache.async_get_cache = AsyncMock(side_effect=RedisCircuitBreakerOpenError())
+        redis_cache.async_increment = AsyncMock(return_value=2)
+        check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(5, redis_cache))
+
+        with pytest.raises(litellm.RateLimitError) as exc_info:
+            await check.async_pre_call_check({**TPM_DEPLOYMENT, "rpm": 1})
+
+        assert "RPM limit=1" in str(exc_info.value)
+
+    @pytest.mark.asyncio
+    async def test_async_pre_call_check_without_redis_enforces_local_tpm_and_rpm(self):
+        check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(5, None))
+        deployment = {**TPM_DEPLOYMENT, "rpm": 1}
+
+        assert await check.async_pre_call_check(deployment) == deployment
+        with pytest.raises(litellm.RateLimitError) as exc_info:
+            await check.async_pre_call_check(deployment)
+
+        assert "RPM limit=1" in str(exc_info.value)
+
     @pytest.mark.asyncio
     async def test_async_log_success_event_increments_cache(self):
         """Test that async_log_success_event correctly increments the cache."""