diff --git a/enterprise/enterprise_hooks/parallel_request_limiter_v2.py b/enterprise/enterprise_hooks/parallel_request_limiter_v2.py index 0ca5d03f4d..3e77831371 100644 --- a/enterprise/enterprise_hooks/parallel_request_limiter_v2.py +++ b/enterprise/enterprise_hooks/parallel_request_limiter_v2.py @@ -26,6 +26,10 @@ from litellm._logging import verbose_proxy_logger from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.core_helpers import _get_parent_otel_span_from_kwargs from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth +from litellm.proxy.auth.auth_utils import ( + get_key_model_rpm_limit, + get_key_model_tpm_limit, +) from litellm.router_strategy.base_routing_strategy import BaseRoutingStrategy if TYPE_CHECKING: @@ -85,25 +89,33 @@ class _PROXY_MaxParallelRequestsHandler(BaseRoutingStrategy, CustomLogger): model: Optional[str], rate_limit_type: Literal["key", "model_per_key", "user", "customer", "team"], group: RateLimitGroups, - ) -> str: - if rate_limit_type == "key": + ) -> Optional[str]: + if rate_limit_type == "key" and user_api_key_dict.api_key is not None: return ( f"{self.prefix}::{user_api_key_dict.api_key}::{precise_minute}::{group}" ) - elif rate_limit_type == "model_per_key" and model is not None: + elif ( + rate_limit_type == "model_per_key" + and model is not None + and user_api_key_dict.api_key is not None + ): return f"{self.prefix}::{user_api_key_dict.api_key}::{model}::{precise_minute}::{group}" - elif rate_limit_type == "user": + elif rate_limit_type == "user" and user_api_key_dict.user_id is not None: return ( f"{self.prefix}::{user_api_key_dict.user_id}::{precise_minute}::{group}" ) - elif rate_limit_type == "customer": + elif ( + rate_limit_type == "customer" and user_api_key_dict.end_user_id is not None + ): return f"{self.prefix}::{user_api_key_dict.end_user_id}::{precise_minute}::{group}" - elif rate_limit_type == "team": + elif rate_limit_type == "team" and user_api_key_dict.team_id is not None: return ( f"{self.prefix}::{user_api_key_dict.team_id}::{precise_minute}::{group}" ) + elif rate_limit_type == "model_per_key" and model is not None: + return f"{self.prefix}::{user_api_key_dict.api_key}::{model}::{precise_minute}::{group}" else: - raise ValueError(f"Invalid rate limit type: {rate_limit_type}") + return None def get_key_pattern_to_sync(self) -> Optional[str]: return self.prefix + "::" @@ -133,6 +145,8 @@ class _PROXY_MaxParallelRequestsHandler(BaseRoutingStrategy, CustomLogger): rate_limit_type=rate_limit_type, group=cast(RateLimitGroups, group), ) + if key is None: + continue increment_list.append((key, increment_value_by_group[group])) if ( @@ -153,7 +167,7 @@ class _PROXY_MaxParallelRequestsHandler(BaseRoutingStrategy, CustomLogger): should_raise_error = should_raise_error or results[2] > tpm_limit if should_raise_error: raise self.raise_rate_limit_error( - additional_details=f"{CommonProxyErrors.max_parallel_request_limit_reached.value}. Hit limit for {rate_limit_type}. Current limits: max_parallel_requests: {max_parallel_requests}, tpm_limit: {tpm_limit}, rpm_limit: {rpm_limit}" + additional_details=f"{CommonProxyErrors.max_parallel_request_limit_reached.value}. Hit limit for {rate_limit_type}. Current usage: max_parallel_requests: {results[0]}, current_rpm: {results[1]}, current_tpm: {results[2]}. Current limits: max_parallel_requests: {max_parallel_requests}, rpm_limit: {rpm_limit}, tpm_limit: {tpm_limit}." ) def time_to_next_minute(self) -> float: @@ -233,7 +247,7 @@ class _PROXY_MaxParallelRequestsHandler(BaseRoutingStrategy, CustomLogger): local_only=True, litellm_parent_otel_span=user_api_key_dict.parent_otel_span, ) - _model = data.get("model", None) + requested_model = data.get("model", None) current_date = datetime.now().strftime("%Y-%m-%d") current_hour = datetime.now().strftime("%H") @@ -254,7 +268,7 @@ class _PROXY_MaxParallelRequestsHandler(BaseRoutingStrategy, CustomLogger): rate_limit_type="key", ) ) - elif user_api_key_dict.user_id is not None: + if user_api_key_dict.user_id is not None: # CHECK IF REQUEST ALLOWED for key tasks.append( self.check_key_in_limits_v2( @@ -267,7 +281,7 @@ class _PROXY_MaxParallelRequestsHandler(BaseRoutingStrategy, CustomLogger): rate_limit_type="user", ) ) - elif user_api_key_dict.team_id is not None: + if user_api_key_dict.team_id is not None: tasks.append( self.check_key_in_limits_v2( user_api_key_dict=user_api_key_dict, @@ -279,6 +293,49 @@ class _PROXY_MaxParallelRequestsHandler(BaseRoutingStrategy, CustomLogger): rate_limit_type="team", ) ) + if user_api_key_dict.end_user_id is not None: + tasks.append( + self.check_key_in_limits_v2( + user_api_key_dict=user_api_key_dict, + data=data, + max_parallel_requests=None, + precise_minute=precise_minute, + tpm_limit=user_api_key_dict.end_user_tpm_limit, + rpm_limit=user_api_key_dict.end_user_rpm_limit, + rate_limit_type="customer", + ) + ) + if requested_model and ( + get_key_model_tpm_limit(user_api_key_dict) is not None + or get_key_model_rpm_limit(user_api_key_dict) is not None + ): + _tpm_limit_for_key_model = get_key_model_tpm_limit(user_api_key_dict) or {} + _rpm_limit_for_key_model = get_key_model_rpm_limit(user_api_key_dict) or {} + + should_check_rate_limit = False + if requested_model in _tpm_limit_for_key_model: + should_check_rate_limit = True + elif requested_model in _rpm_limit_for_key_model: + should_check_rate_limit = True + + if should_check_rate_limit: + model_specific_tpm_limit: Optional[int] = None + model_specific_rpm_limit: Optional[int] = None + if requested_model in _tpm_limit_for_key_model: + model_specific_tpm_limit = _tpm_limit_for_key_model[requested_model] + if requested_model in _rpm_limit_for_key_model: + model_specific_rpm_limit = _rpm_limit_for_key_model[requested_model] + tasks.append( + self.check_key_in_limits_v2( + user_api_key_dict=user_api_key_dict, + data=data, + max_parallel_requests=None, + precise_minute=precise_minute, + tpm_limit=model_specific_tpm_limit, + rpm_limit=model_specific_rpm_limit, + rate_limit_type="model_per_key", + ) + ) await asyncio.gather(*tasks) return @@ -298,7 +355,7 @@ class _PROXY_MaxParallelRequestsHandler(BaseRoutingStrategy, CustomLogger): "rpm": 0, } - rate_limit_types = ["key", "user", "customer", "team"] + rate_limit_types = ["key", "user", "customer", "team", "model_per_key"] for rate_limit_type in rate_limit_types: for group in ["request_count", "rpm", "tpm"]: key = self._get_current_usage_key( @@ -308,7 +365,8 @@ class _PROXY_MaxParallelRequestsHandler(BaseRoutingStrategy, CustomLogger): rate_limit_type=cast(RateLimitTypes, rate_limit_type), group=cast(RateLimitGroups, group), ) - + if key is None: + continue increment_list.append((key, increment_value_by_group[group])) if increment_list: # Only call if we have values to increment @@ -344,7 +402,9 @@ class _PROXY_MaxParallelRequestsHandler(BaseRoutingStrategy, CustomLogger): user_api_key_team_id = kwargs["litellm_params"]["metadata"].get( "user_api_key_team_id", None ) - user_api_key_end_user_id = kwargs.get("user") + user_api_key_end_user_id = kwargs.get("user") or kwargs["litellm_params"][ + "metadata" + ].get("user_api_key_end_user_id", None) # ------------ # Setup values diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index c59ed19805..6bb5801f9a 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -14,7 +14,7 @@ import inspect import json import time from datetime import timedelta -from typing import TYPE_CHECKING, Any, List, Optional, Tuple, Union +from typing import TYPE_CHECKING, Any, List, Optional, Tuple, Union, cast import litellm from litellm._logging import print_verbose, verbose_logger @@ -83,7 +83,9 @@ class RedisCache(BaseCache): redis_kwargs.update(kwargs) self.redis_client = get_redis_client(**redis_kwargs) - self.redis_async_client: Optional[async_redis_client] = None + self.redis_async_client: Optional[ + Union[async_redis_client, async_redis_cluster_client] + ] = None self.redis_kwargs = redis_kwargs self.async_redis_conn_pool = get_redis_connection_pool(**redis_kwargs) @@ -134,13 +136,27 @@ class RedisCache(BaseCache): def init_async_client( self, ) -> Union[async_redis_client, async_redis_cluster_client]: - from .._redis import get_redis_async_client + from litellm import in_memory_llm_clients_cache - if self.redis_async_client is None: - self.redis_async_client = get_redis_async_client( + from .._redis import get_redis_async_client, get_redis_connection_pool + + cached_client = in_memory_llm_clients_cache.get_cache(key="async-redis-client") + if cached_client is not None: + redis_async_client = cast( + Union[async_redis_client, async_redis_cluster_client], cached_client + ) + else: + # Create new connection pool and client for current event loop + self.async_redis_conn_pool = get_redis_connection_pool(**self.redis_kwargs) + redis_async_client = get_redis_async_client( connection_pool=self.async_redis_conn_pool, **self.redis_kwargs ) - return self.redis_async_client + in_memory_llm_clients_cache.set_cache( + key="async-redis-client", value=self.redis_async_client + ) + + self.redis_async_client = redis_async_client # type: ignore + return redis_async_client def check_and_fix_namespace(self, key: str) -> str: """ diff --git a/tests/enterprise/enterprise_hooks/test_parallel_request_limiter_v2.py b/tests/enterprise/enterprise_hooks/test_parallel_request_limiter_v2.py index 0a472db638..769b082295 100644 --- a/tests/enterprise/enterprise_hooks/test_parallel_request_limiter_v2.py +++ b/tests/enterprise/enterprise_hooks/test_parallel_request_limiter_v2.py @@ -106,9 +106,9 @@ async def test_normal_router_call_v2(monkeypatch): "rate_limit_object", [ "key", - # "model_per_key", + "model_per_key", "user", - # "customer", + "customer", "team", ], ) @@ -154,6 +154,13 @@ async def test_normal_router_call_tpm(monkeypatch, rate_limit_object): user_api_key_dict = UserAPIKeyAuth(user_id="12345", user_tpm_limit=10) elif rate_limit_object == "team": user_api_key_dict = UserAPIKeyAuth(team_id="12345", team_tpm_limit=10) + elif rate_limit_object == "customer": + user_api_key_dict = UserAPIKeyAuth(end_user_id="12345", end_user_tpm_limit=10) + elif rate_limit_object == "model_per_key": + user_api_key_dict = UserAPIKeyAuth( + api_key=_api_key, + metadata={"model_tpm_limit": {"azure-model": 10}}, + ) local_cache = DualCache() parallel_request_handler = _PROXY_MaxParallelRequestsHandler( internal_usage_cache=InternalUsageCache(local_cache) @@ -161,7 +168,10 @@ async def test_normal_router_call_tpm(monkeypatch, rate_limit_object): monkeypatch.setattr(litellm, "callbacks", [parallel_request_handler]) await parallel_request_handler.async_pre_call_hook( - user_api_key_dict=user_api_key_dict, cache=local_cache, data={}, call_type="" + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"model": "azure-model"}, + call_type="", ) current_date = datetime.now().strftime("%Y-%m-%d") @@ -171,7 +181,7 @@ async def test_normal_router_call_tpm(monkeypatch, rate_limit_object): request_count_api_key = parallel_request_handler._get_current_usage_key( user_api_key_dict=user_api_key_dict, precise_minute=precise_minute, - model=None, + model="azure-model", rate_limit_type=rate_limit_object, group="tpm", ) @@ -191,6 +201,7 @@ async def test_normal_router_call_tpm(monkeypatch, rate_limit_object): "user_api_key": _api_key, "user_api_key_user_id": user_api_key_dict.user_id, "user_api_key_team_id": user_api_key_dict.team_id, + "user_api_key_end_user_id": user_api_key_dict.end_user_id, }, mock_response="hello", ) @@ -210,9 +221,9 @@ async def test_normal_router_call_tpm(monkeypatch, rate_limit_object): "rate_limit_object", [ "key", - # "model_per_key", + "model_per_key", "user", - # "customer", + "customer", "team", ], ) @@ -258,6 +269,13 @@ async def test_normal_router_call_rpm(monkeypatch, rate_limit_object): user_api_key_dict = UserAPIKeyAuth(user_id="12345", user_rpm_limit=1) elif rate_limit_object == "team": user_api_key_dict = UserAPIKeyAuth(team_id="12345", team_rpm_limit=1) + elif rate_limit_object == "customer": + user_api_key_dict = UserAPIKeyAuth(end_user_id="12345", end_user_rpm_limit=1) + elif rate_limit_object == "model_per_key": + user_api_key_dict = UserAPIKeyAuth( + api_key=_api_key, + metadata={"model_rpm_limit": {"azure-model": 1}}, + ) local_cache = DualCache() parallel_request_handler = _PROXY_MaxParallelRequestsHandler( internal_usage_cache=InternalUsageCache(local_cache) @@ -265,7 +283,10 @@ async def test_normal_router_call_rpm(monkeypatch, rate_limit_object): monkeypatch.setattr(litellm, "callbacks", [parallel_request_handler]) await parallel_request_handler.async_pre_call_hook( - user_api_key_dict=user_api_key_dict, cache=local_cache, data={}, call_type="" + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"model": "azure-model"}, + call_type="", ) current_date = datetime.now().strftime("%Y-%m-%d") @@ -275,7 +296,7 @@ async def test_normal_router_call_rpm(monkeypatch, rate_limit_object): request_count_api_key = parallel_request_handler._get_current_usage_key( user_api_key_dict=user_api_key_dict, precise_minute=precise_minute, - model=None, + model="azure-model", rate_limit_type=rate_limit_object, group="rpm", ) @@ -295,6 +316,7 @@ async def test_normal_router_call_rpm(monkeypatch, rate_limit_object): "user_api_key": _api_key, "user_api_key_user_id": user_api_key_dict.user_id, "user_api_key_team_id": user_api_key_dict.team_id, + "user_api_key_end_user_id": user_api_key_dict.end_user_id, }, mock_response="hello", ) @@ -313,7 +335,7 @@ async def test_normal_router_call_rpm(monkeypatch, rate_limit_object): await parallel_request_handler.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=local_cache, - data={}, + data={"model": "azure-model"}, call_type="", )