Merge pull request #21957 from jquinter/fix/flaky-rpm-limit-test

fix: atomic RPM rate limiting in model rate limit check

litellm/router_utils/pre_call_checks/model_rate_limit_check.py

@@ -129,27 +129,8 @@ class ModelRateLimitingCheck(CustomLogger):
                         ),
                     )
 
-            # Check RPM limit
+            # Check RPM limit (atomic increment-first to avoid race conditions)
             if rpm_limit is not None:
-                # First check local cache
-                current_rpm = self.dual_cache.get_cache(key=rpm_key, local_only=True)
-                if current_rpm >= rpm_limit:
-                    raise litellm.RateLimitError(
-                        message=f"Model rate limit exceeded. RPM limit={rpm_limit}, current usage={current_rpm}",
-                        llm_provider="",
-                        model=model_name,
-                        response=httpx.Response(
-                            status_code=429,
-                            content=f"{RouterErrors.user_defined_ratelimit_error.value} rpm limit={rpm_limit}. current usage={current_rpm}. id={model_id}, model_group={model_group}",
-                            headers={"retry-after": str(60)},
-                            request=httpx.Request(
-                                method="model_rate_limit_check",
-                                url="https://github.com/BerriAI/litellm",
-                            ),
-                        ),
-                    )
-
-                # Check redis cache and increment
                 current_rpm = self.dual_cache.increment_cache(
                     key=rpm_key, value=1, ttl=RoutingArgs.ttl
                 )
@@ -226,30 +207,8 @@ class ModelRateLimitingCheck(CustomLogger):
                         num_retries=0,  # Don't retry - return 429 immediately
                     )
 
-            # Check RPM limit
+            # Check RPM limit (atomic increment-first to avoid race conditions)
             if rpm_limit is not None:
-                # First check local cache
-                current_rpm = await self.dual_cache.async_get_cache(
-                    key=rpm_key, local_only=True
-                )
-                if current_rpm is not None and current_rpm >= rpm_limit:
-                    raise litellm.RateLimitError(
-                        message=f"Model rate limit exceeded. RPM limit={rpm_limit}, current usage={current_rpm}",
-                        llm_provider="",
-                        model=model_name,
-                        response=httpx.Response(
-                            status_code=429,
-                            content=f"{RouterErrors.user_defined_ratelimit_error.value} rpm limit={rpm_limit}. current usage={current_rpm}. id={model_id}, model_group={model_group}",
-                            headers={"retry-after": str(60)},
-                            request=httpx.Request(
-                                method="model_rate_limit_check",
-                                url="https://github.com/BerriAI/litellm",
-                            ),
-                        ),
-                        num_retries=0,  # Don't retry - return 429 immediately
-                    )
-
-                # Check redis cache and increment
                 current_rpm = await self.dual_cache.async_increment_cache(
                     key=rpm_key,
                     value=1,

tests/test_litellm/test_router/test_enforce_model_rate_limits.py

@@ -5,12 +5,14 @@ This feature allows users to enforce TPM/RPM limits set on model deployments
 regardless of the routing strategy being used.
 """
 
+import asyncio
 from unittest.mock import AsyncMock, MagicMock
 
 import pytest
 
 import litellm
 from litellm import Router
+from litellm.caching.dual_cache import DualCache
 from litellm.router_utils.pre_call_checks.model_rate_limit_check import (
     ModelRateLimitingCheck,
 )
@@ -88,7 +90,7 @@ class TestModelRateLimitingCheck:
     def test_pre_call_check_raises_rate_limit_error_when_over_rpm(self):
         """Test that RateLimitError is raised when RPM limit is exceeded."""
         mock_cache = MagicMock()
-        mock_cache.get_cache.return_value = 10  # Already at limit
+        mock_cache.increment_cache.return_value = 11  # Over limit after increment
 
         check = ModelRateLimitingCheck(dual_cache=mock_cache)
 
@@ -103,12 +105,11 @@ class TestModelRateLimitingCheck:
             check.pre_call_check(deployment)
 
         assert "RPM limit=10" in str(exc_info.value)
-        assert "current usage=10" in str(exc_info.value)
+        assert "current usage=11" in str(exc_info.value)
 
     def test_pre_call_check_allows_request_under_limit(self):
         """Test that requests are allowed when under the limit."""
         mock_cache = MagicMock()
-        mock_cache.get_cache.return_value = 5
         mock_cache.increment_cache.return_value = 6
 
         check = ModelRateLimitingCheck(dual_cache=mock_cache)
@@ -188,7 +189,8 @@ class TestModelRateLimitingCheckAsync:
     async def test_async_pre_call_check_raises_rate_limit_error_when_over_rpm(self):
         """Test that RateLimitError is raised when RPM limit is exceeded (async)."""
         mock_cache = MagicMock()
-        mock_cache.async_get_cache = AsyncMock(return_value=10)  # Already at limit
+        mock_cache.async_get_cache = AsyncMock(return_value=None)
+        mock_cache.async_increment_cache = AsyncMock(return_value=11)  # Over limit
 
         check = ModelRateLimitingCheck(dual_cache=mock_cache)
 
@@ -208,7 +210,7 @@ class TestModelRateLimitingCheckAsync:
     async def test_async_pre_call_check_allows_request_under_limit(self):
         """Test that requests are allowed when under the limit (async)."""
         mock_cache = MagicMock()
-        mock_cache.async_get_cache = AsyncMock(return_value=5)
+        mock_cache.async_get_cache = AsyncMock(return_value=None)
         mock_cache.async_increment_cache = AsyncMock(return_value=6)
 
         check = ModelRateLimitingCheck(dual_cache=mock_cache)
@@ -313,3 +315,41 @@ class TestRouterWithEnforceModelRateLimits:
                 break
 
         assert found, "ModelRateLimitingCheck should be in litellm.callbacks"
+
+
+class TestModelRateLimitConcurrency:
+    """Test that RPM rate limiting is atomic under concurrent requests."""
+
+    @pytest.mark.asyncio
+    async def test_concurrent_requests_respect_rpm_limit(self):
+        """
+        Fire 4 concurrent async requests with RPM limit of 2.
+        Exactly 2 should succeed and 2 should raise RateLimitError.
+
+        This test validates the atomic increment-first pattern:
+        the old check-then-increment pattern would let 3+ through
+        due to a race condition on the local cache read.
+        """
+        dual_cache = DualCache()
+        check = ModelRateLimitingCheck(dual_cache=dual_cache)
+
+        deployment = {
+            "rpm": 2,
+            "litellm_params": {"model": "gpt-4"},
+            "model_info": {"id": "concurrent-test-id"},
+            "model_name": "test-model",
+        }
+
+        async def attempt_request():
+            return await check.async_pre_call_check(deployment)
+
+        results = await asyncio.gather(
+            *[attempt_request() for _ in range(4)],
+            return_exceptions=True,
+        )
+
+        successes = [r for r in results if not isinstance(r, Exception)]
+        failures = [r for r in results if isinstance(r, litellm.RateLimitError)]
+
+        assert len(successes) == 2, f"Expected 2 successes, got {len(successes)}"
+        assert len(failures) == 2, f"Expected 2 rate limit errors, got {len(failures)}"