diff --git a/litellm/router_utils/pre_call_checks/model_rate_limit_check.py b/litellm/router_utils/pre_call_checks/model_rate_limit_check.py index af3d7ddfac7..79ea6dc36ec 100644 --- a/litellm/router_utils/pre_call_checks/model_rate_limit_check.py +++ b/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}", diff --git a/tests/test_litellm/test_router/test_enforce_model_rate_limits.py b/tests/test_litellm/test_router/test_enforce_model_rate_limits.py index 1def253ac93..ee665051106 100644 --- a/tests/test_litellm/test_router/test_enforce_model_rate_limits.py +++ b/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."""