From f58fefd58922e31825e71c7eac36db78900dc7bb Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 30 Mar 2024 20:01:36 -0700 Subject: [PATCH] fix(tpm_rpm_limiter.py): enable redis caching for tpm/rpm checks on keys/user/teams allows tpm/rpm checks to work across instances https://github.com/BerriAI/litellm/issues/2730 --- litellm/caching.py | 9 +- litellm/proxy/_new_secret_config.yaml | 5 +- litellm/proxy/hooks/tpm_rpm_limiter.py | 380 +++++++++++++++++++++++++ litellm/proxy/proxy_server.py | 23 +- litellm/proxy/utils.py | 17 +- 5 files changed, 423 insertions(+), 11 deletions(-) create mode 100644 litellm/proxy/hooks/tpm_rpm_limiter.py diff --git a/litellm/caching.py b/litellm/caching.py index 921ae1b21a..eea687edfe 100644 --- a/litellm/caching.py +++ b/litellm/caching.py @@ -796,6 +796,8 @@ class DualCache(BaseCache): self, in_memory_cache: Optional[InMemoryCache] = None, redis_cache: Optional[RedisCache] = None, + default_in_memory_ttl: Optional[float] = None, + default_redis_ttl: Optional[float] = None, ) -> None: super().__init__() # If in_memory_cache is not provided, use the default InMemoryCache @@ -803,11 +805,17 @@ class DualCache(BaseCache): # If redis_cache is not provided, use the default RedisCache self.redis_cache = redis_cache + self.default_in_memory_ttl = default_in_memory_ttl + self.default_redis_ttl = default_redis_ttl + def set_cache(self, key, value, local_only: bool = False, **kwargs): # Update both Redis and in-memory cache try: print_verbose(f"set cache: key: {key}; value: {value}") if self.in_memory_cache is not None: + if "ttl" not in kwargs and self.default_in_memory_ttl is not None: + kwargs["ttl"] = self.default_in_memory_ttl + self.in_memory_cache.set_cache(key, value, **kwargs) if self.redis_cache is not None and local_only == False: @@ -823,7 +831,6 @@ class DualCache(BaseCache): if self.in_memory_cache is not None: in_memory_result = self.in_memory_cache.get_cache(key, **kwargs) - print_verbose(f"in_memory_result: {in_memory_result}") if in_memory_result is not None: result = in_memory_result diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml index d218fddb19..dde07bb870 100644 --- a/litellm/proxy/_new_secret_config.yaml +++ b/litellm/proxy/_new_secret_config.yaml @@ -6,8 +6,9 @@ model_list: api_base: https://exampleopenaiendpoint-production.up.railway.app/ litellm_settings: - max_budget: 600020 - budget_duration: 30d + cache: true +# max_budget: 600020 +# budget_duration: 30d general_settings: master_key: sk-1234 diff --git a/litellm/proxy/hooks/tpm_rpm_limiter.py b/litellm/proxy/hooks/tpm_rpm_limiter.py new file mode 100644 index 0000000000..db1d1759ad --- /dev/null +++ b/litellm/proxy/hooks/tpm_rpm_limiter.py @@ -0,0 +1,380 @@ +# What is this? +## Checks TPM/RPM Limits for a key/user/team on the proxy +## Works with Redis - if given + +from typing import Optional, Literal +import litellm, traceback, sys +from litellm.caching import DualCache, RedisCache +from litellm.proxy._types import ( + UserAPIKeyAuth, + LiteLLM_VerificationTokenView, + LiteLLM_UserTable, + LiteLLM_TeamTable, +) +from litellm.integrations.custom_logger import CustomLogger +from fastapi import HTTPException +from litellm._logging import verbose_proxy_logger +from litellm import ModelResponse +from datetime import datetime + + +class _PROXY_MaxTPMRPMLimiter(CustomLogger): + user_api_key_cache = None + + # Class variables or attributes + def __init__(self, redis_usage_cache: Optional[RedisCache]): + self.redis_usage_cache = redis_usage_cache + self.internal_cache = DualCache( + redis_cache=redis_usage_cache, + default_in_memory_ttl=10, + default_redis_ttl=60, + ) + + def print_verbose(self, print_statement): + try: + verbose_proxy_logger.debug(print_statement) + if litellm.set_verbose: + print(print_statement) # noqa + except: + pass + + ## check if admin has set tpm/rpm limits for this key/user/team + + def _check_limits_set( + self, + user_api_key_cache: DualCache, + key: Optional[str], + user_id: Optional[str], + team_id: Optional[str], + ) -> bool: + ## key + if key is not None: + key_val = user_api_key_cache.get_cache(key=key) + if isinstance(key_val, dict): + key_val = LiteLLM_VerificationTokenView(**key_val) + + if isinstance(key_val, LiteLLM_VerificationTokenView): + user_api_key_tpm_limit = key_val.tpm_limit + + user_api_key_rpm_limit = key_val.rpm_limit + + if ( + user_api_key_tpm_limit is not None + or user_api_key_rpm_limit is not None + ): + return True + + ## team + if team_id is not None: + team_val = user_api_key_cache.get_cache(key=team_id) + if isinstance(team_val, dict): + team_val = LiteLLM_TeamTable(**team_val) + + if isinstance(team_val, LiteLLM_TeamTable): + team_tpm_limit = team_val.tpm_limit + + team_rpm_limit = team_val.rpm_limit + + if team_tpm_limit is not None or team_rpm_limit is not None: + return True + + ## user + if user_id is not None: + user_val = user_api_key_cache.get_cache(key=user_id) + if isinstance(user_val, dict): + user_val = LiteLLM_UserTable(**user_val) + + if isinstance(user_val, LiteLLM_UserTable): + user_tpm_limit = user_val.tpm_limit + + user_rpm_limit = user_val.rpm_limit + + if user_tpm_limit is not None or user_rpm_limit is not None: + return True + return False + + async def check_key_in_limits( + self, + user_api_key_dict: UserAPIKeyAuth, + current_minute_dict: dict, + tpm_limit: int, + rpm_limit: int, + request_count_api_key: str, + type: Literal["key", "user", "team"], + ): + if type == "key" and user_api_key_dict.api_key is not None: + current = current_minute_dict["key"].get(user_api_key_dict.api_key, None) + elif type == "user" and user_api_key_dict.user_id is not None: + current = current_minute_dict["user"].get(user_api_key_dict.user_id, None) + elif type == "team" and user_api_key_dict.team_id is not None: + current = current_minute_dict["team"].get(user_api_key_dict.team_id, None) + else: + return + + if current is None: + if tpm_limit == 0 or rpm_limit == 0: + # base case + raise HTTPException( + status_code=429, detail="Max tpm/rpm limit reached." + ) + elif current["current_tpm"] < tpm_limit and current["current_rpm"] < rpm_limit: + pass + else: + raise HTTPException(status_code=429, detail="Max tpm/rpm limit reached.") + + async def async_pre_call_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + cache: DualCache, + data: dict, + call_type: str, + ): + self.print_verbose( + f"Inside Max TPM/RPM Limiter Pre-Call Hook - {user_api_key_dict}" + ) + api_key = user_api_key_dict.api_key + # check if REQUEST ALLOWED for user_id + user_id = user_api_key_dict.user_id + ## get team tpm/rpm limits + team_id = user_api_key_dict.team_id + + _set_limits = self._check_limits_set( + user_api_key_cache=cache, key=api_key, user_id=user_id, team_id=team_id + ) + + if _set_limits == False: + return + + # ------------ + # Setup values + # ------------ + + self.user_api_key_cache = cache + + current_date = datetime.now().strftime("%Y-%m-%d") + current_hour = datetime.now().strftime("%H") + current_minute = datetime.now().strftime("%M") + precise_minute = f"{current_date}-{current_hour}-{current_minute}" + cache_key = "usage:{}".format(precise_minute) + current_minute_dict = await self.internal_cache.async_get_cache( + key=cache_key + ) # {"usage:{curr_minute}": {"key": {: {"current_requests": 1, "current_tpm": 1, "current_rpm": 10}}}} + + if current_minute_dict is None: + current_minute_dict = {"key": {}, "user": {}, "team": {}} + + if api_key is not None: + tpm_limit = getattr(user_api_key_dict, "tpm_limit", sys.maxsize) + if tpm_limit is None: + tpm_limit = sys.maxsize + rpm_limit = getattr(user_api_key_dict, "rpm_limit", sys.maxsize) + if rpm_limit is None: + rpm_limit = sys.maxsize + request_count_api_key = f"{api_key}::{precise_minute}::request_count" + await self.check_key_in_limits( + user_api_key_dict=user_api_key_dict, + current_minute_dict=current_minute_dict, + request_count_api_key=request_count_api_key, + tpm_limit=tpm_limit, + rpm_limit=rpm_limit, + type="key", + ) + + if user_id is not None: + _user_id_rate_limits = user_api_key_dict.user_id_rate_limits + + # get user tpm/rpm limits + if _user_id_rate_limits is not None and isinstance( + _user_id_rate_limits, dict + ): + user_tpm_limit = _user_id_rate_limits.get("tpm_limit", None) + user_rpm_limit = _user_id_rate_limits.get("rpm_limit", None) + if user_tpm_limit is None: + user_tpm_limit = sys.maxsize + if user_rpm_limit is None: + user_rpm_limit = sys.maxsize + + # now do the same tpm/rpm checks + request_count_api_key = f"{user_id}::{precise_minute}::request_count" + # print(f"Checking if {request_count_api_key} is allowed to make request for minute {precise_minute}") + await self.check_key_in_limits( + user_api_key_dict=user_api_key_dict, + current_minute_dict=current_minute_dict, + request_count_api_key=request_count_api_key, + tpm_limit=user_tpm_limit, + rpm_limit=user_rpm_limit, + type="user", + ) + + # TEAM RATE LIMITS + if team_id is not None: + team_tpm_limit = getattr(user_api_key_dict, "team_tpm_limit", sys.maxsize) + + if team_tpm_limit is None: + team_tpm_limit = sys.maxsize + team_rpm_limit = getattr(user_api_key_dict, "team_rpm_limit", sys.maxsize) + if team_rpm_limit is None: + team_rpm_limit = sys.maxsize + + if team_tpm_limit is None: + team_tpm_limit = sys.maxsize + if team_rpm_limit is None: + team_rpm_limit = sys.maxsize + + # now do the same tpm/rpm checks + request_count_api_key = f"{team_id}::{precise_minute}::request_count" + + # print(f"Checking if {request_count_api_key} is allowed to make request for minute {precise_minute}") + await self.check_key_in_limits( + user_api_key_dict=user_api_key_dict, + current_minute_dict=current_minute_dict, + request_count_api_key=request_count_api_key, + tpm_limit=team_tpm_limit, + rpm_limit=team_rpm_limit, + type="team", + ) + + return + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + try: + self.print_verbose(f"INSIDE TPM RPM Limiter ASYNC SUCCESS LOGGING") + + user_api_key = kwargs["litellm_params"]["metadata"]["user_api_key"] + user_api_key_user_id = kwargs["litellm_params"]["metadata"].get( + "user_api_key_user_id", None + ) + user_api_key_team_id = kwargs["litellm_params"]["metadata"].get( + "user_api_key_team_id", None + ) + + _limits_set = self._check_limits_set( + user_api_key_cache=self.user_api_key_cache, + key=user_api_key, + user_id=user_api_key_user_id, + team_id=user_api_key_team_id, + ) + + if _limits_set == False: # don't waste cache calls if no tpm/rpm limits set + return + + # ------------ + # Setup values + # ------------ + + current_date = datetime.now().strftime("%Y-%m-%d") + current_hour = datetime.now().strftime("%H") + current_minute = datetime.now().strftime("%M") + precise_minute = f"{current_date}-{current_hour}-{current_minute}" + + total_tokens = 0 + + if isinstance(response_obj, ModelResponse): + total_tokens = response_obj.usage.total_tokens + + """ + - get value from redis + - increment requests + 1 + - increment tpm + 1 + - increment rpm + 1 + - update value in-memory + redis + """ + cache_key = "usage:{}".format(precise_minute) + if ( + self.internal_cache.redis_cache is not None + ): # get straight from redis if possible + current_minute_dict = ( + await self.internal_cache.redis_cache.async_get_cache( + key=cache_key, + ) + ) # {"usage:{current_minute}": {"key": {}, "team": {}, "user": {}}} + else: + current_minute_dict = await self.internal_cache.async_get_cache( + key=cache_key, + ) + + if current_minute_dict is None: + current_minute_dict = {"key": {}, "user": {}, "team": {}} + + _cache_updated = False # check if a cache update is required. prevent unnecessary rewrites. + + # ------------ + # Update usage - API Key + # ------------ + + if user_api_key is not None: + _cache_updated = True + ## API KEY ## + if user_api_key in current_minute_dict["key"]: + current_key_usage = current_minute_dict["key"][user_api_key] + new_val = { + "current_tpm": current_key_usage["current_tpm"] + total_tokens, + "current_rpm": current_key_usage["current_rpm"] + 1, + } + else: + new_val = { + "current_tpm": total_tokens, + "current_rpm": 1, + } + + current_minute_dict["key"][user_api_key] = new_val + + self.print_verbose( + f"updated_value in success call: {new_val}, precise_minute: {precise_minute}" + ) + + # ------------ + # Update usage - User + # ------------ + if user_api_key_user_id is not None: + _cache_updated = True + total_tokens = 0 + + if isinstance(response_obj, ModelResponse): + total_tokens = response_obj.usage.total_tokens + + if user_api_key_user_id in current_minute_dict["key"]: + current_key_usage = current_minute_dict["key"][user_api_key_user_id] + new_val = { + "current_tpm": current_key_usage["current_tpm"] + total_tokens, + "current_rpm": current_key_usage["current_rpm"] + 1, + } + else: + new_val = { + "current_tpm": total_tokens, + "current_rpm": 1, + } + + current_minute_dict["user"][user_api_key_user_id] = new_val + + # ------------ + # Update usage - Team + # ------------ + if user_api_key_team_id is not None: + _cache_updated = True + total_tokens = 0 + + if isinstance(response_obj, ModelResponse): + total_tokens = response_obj.usage.total_tokens + + if user_api_key_team_id in current_minute_dict["key"]: + current_key_usage = current_minute_dict["key"][user_api_key_team_id] + new_val = { + "current_tpm": current_key_usage["current_tpm"] + total_tokens, + "current_rpm": current_key_usage["current_rpm"] + 1, + } + else: + new_val = { + "current_tpm": total_tokens, + "current_rpm": 1, + } + + current_minute_dict["team"][user_api_key_team_id] = new_val + + if _cache_updated == True: + await self.internal_cache.async_set_cache( + key=cache_key, value=current_minute_dict + ) + + except Exception as e: + self.print_verbose(e) # noqa diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index e440581d28..6c570eb59a 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -102,7 +102,7 @@ from litellm.proxy.secret_managers.google_kms import load_google_kms from litellm.proxy.secret_managers.aws_secret_manager import load_aws_secret_manager import pydantic from litellm.proxy._types import * -from litellm.caching import DualCache +from litellm.caching import DualCache, RedisCache from litellm.proxy.health_check import perform_health_check from litellm._logging import verbose_router_logger, verbose_proxy_logger from litellm.proxy.auth.handle_jwt import JWTHandler @@ -281,6 +281,9 @@ otel_logging = False prisma_client: Optional[PrismaClient] = None custom_db_client: Optional[DBClient] = None user_api_key_cache = DualCache() +redis_usage_cache: Optional[RedisCache] = ( + None # redis cache used for tracking spend, tpm/rpm limits +) user_custom_auth = None user_custom_key_generate = None use_background_health_checks = None @@ -299,7 +302,9 @@ disable_spend_logs = False jwt_handler = JWTHandler() prompt_injection_detection_obj: Optional[_OPTIONAL_PromptInjectionDetection] = None ### INITIALIZE GLOBAL LOGGING OBJECT ### -proxy_logging_obj = ProxyLogging(user_api_key_cache=user_api_key_cache) +proxy_logging_obj = ProxyLogging( + user_api_key_cache=user_api_key_cache, redis_usage_cache=redis_usage_cache +) ### REDIS QUEUE ### async_result = None celery_app_conn = None @@ -909,6 +914,10 @@ async def user_api_key_auth( models=valid_token.team_models, ) + user_api_key_cache.set_cache( + key=valid_token.team_id, value=_team_obj + ) # save team table in cache - used for tpm/rpm limiting - tpm_rpm_limiter.py + _end_user_object = None if "user" in request_data: _id = "end_user_id:{}".format(request_data["user"]) @@ -1905,7 +1914,7 @@ class ProxyConfig: """ Load config values into proxy global state """ - global master_key, user_config_file_path, otel_logging, user_custom_auth, user_custom_auth_path, user_custom_key_generate, use_background_health_checks, health_check_interval, use_queue, custom_db_client, proxy_budget_rescheduler_max_time, proxy_budget_rescheduler_min_time, ui_access_mode, litellm_master_key_hash, proxy_batch_write_at, disable_spend_logs, prompt_injection_detection_obj + global master_key, user_config_file_path, otel_logging, user_custom_auth, user_custom_auth_path, user_custom_key_generate, use_background_health_checks, health_check_interval, use_queue, custom_db_client, proxy_budget_rescheduler_max_time, proxy_budget_rescheduler_min_time, ui_access_mode, litellm_master_key_hash, proxy_batch_write_at, disable_spend_logs, prompt_injection_detection_obj, redis_usage_cache # Load existing config config = await self.get_config(config_file_path=config_file_path) @@ -1967,6 +1976,7 @@ class ProxyConfig: "password": cache_password, } ) + # Assuming cache_type, cache_host, cache_port, and cache_password are strings print( # noqa f"{blue_color_code}Cache Type:{reset_color_code} {cache_type}" @@ -1991,7 +2001,14 @@ class ProxyConfig: cache_params[key] = litellm.get_secret(value) ## to pass a complete url, or set ssl=True, etc. just set it as `os.environ[REDIS_URL] = `, _redis.py checks for REDIS specific environment variables + litellm.cache = Cache(**cache_params) + + if litellm.cache is not None and isinstance( + litellm.cache.cache, RedisCache + ): + ## INIT PROXY REDIS USAGE CLIENT ## + redis_usage_cache = litellm.cache.cache print( # noqa f"{blue_color_code}Set Cache on LiteLLM Proxy: {vars(litellm.cache.cache)}{reset_color_code}" ) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index f70d67aac0..7d3e30f865 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -12,13 +12,14 @@ from litellm.proxy._types import ( LiteLLM_TeamTable, Member, ) -from litellm.caching import DualCache +from litellm.caching import DualCache, RedisCache from litellm.llms.custom_httpx.httpx_handler import HTTPHandler from litellm.proxy.hooks.parallel_request_limiter import ( _PROXY_MaxParallelRequestsHandler, ) from litellm import ModelResponse, EmbeddingResponse, ImageResponse from litellm.proxy.hooks.max_budget_limiter import _PROXY_MaxBudgetLimiter +from litellm.proxy.hooks.tpm_rpm_limiter import _PROXY_MaxTPMRPMLimiter from litellm.proxy.hooks.cache_control_check import _PROXY_CacheControlCheck from litellm.integrations.custom_logger import CustomLogger from litellm.proxy.db.base_client import CustomDB @@ -46,16 +47,21 @@ class ProxyLogging: - support the max parallel request integration """ - def __init__(self, user_api_key_cache: DualCache): + def __init__( + self, user_api_key_cache: DualCache, redis_usage_cache: Optional[RedisCache] + ): ## INITIALIZE LITELLM CALLBACKS ## self.call_details: dict = {} self.call_details["user_api_key_cache"] = user_api_key_cache - self.max_parallel_request_limiter = _PROXY_MaxParallelRequestsHandler() + # self.max_parallel_request_limiter = _PROXY_MaxParallelRequestsHandler() + self.max_tpm_rpm_limiter = _PROXY_MaxTPMRPMLimiter( + redis_usage_cache=redis_usage_cache + ) self.max_budget_limiter = _PROXY_MaxBudgetLimiter() self.cache_control_check = _PROXY_CacheControlCheck() self.alerting: Optional[List] = None self.alerting_threshold: float = 300 # default to 5 min. threshold - pass + self.redis_usage_cache = redis_usage_cache def update_values( self, alerting: Optional[List], alerting_threshold: Optional[float] @@ -66,7 +72,8 @@ class ProxyLogging: def _init_litellm_callbacks(self): print_verbose(f"INITIALIZING LITELLM CALLBACKS!") - litellm.callbacks.append(self.max_parallel_request_limiter) + # litellm.callbacks.append(self.max_parallel_request_limiter) + litellm.callbacks.append(self.max_tpm_rpm_limiter) litellm.callbacks.append(self.max_budget_limiter) litellm.callbacks.append(self.cache_control_check) litellm.success_callback.append(self.response_taking_too_long_callback)