From b5900099af9127e0f86c085c7d08b5d947078ebe Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Thu, 22 Feb 2024 18:44:03 -0800 Subject: [PATCH 1/4] (feat) tpm/rpm limit by User --- litellm/proxy/_types.py | 4 + .../proxy/hooks/parallel_request_limiter.py | 181 ++++++++++++++---- .../tests/test_parallel_request_limiter.py | 50 +++++ 3 files changed, 203 insertions(+), 32 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index f0f3840947d..7f453980fbf 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -424,6 +424,10 @@ class LiteLLM_VerificationToken(LiteLLMBase): model_spend: Dict = {} model_max_budget: Dict = {} + # hidden params used for parallel request limiting, not required to create a token + user_id_rate_limits: Optional[dict] = None + team_id_rate_limits: Optional[dict] = None + class Config: protected_namespaces = () diff --git a/litellm/proxy/hooks/parallel_request_limiter.py b/litellm/proxy/hooks/parallel_request_limiter.py index 67f8d1ad2f4..021fbc5fb7f 100644 --- a/litellm/proxy/hooks/parallel_request_limiter.py +++ b/litellm/proxy/hooks/parallel_request_limiter.py @@ -24,46 +24,21 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): except: pass - async def async_pre_call_hook( + async def check_key_in_limits( self, user_api_key_dict: UserAPIKeyAuth, cache: DualCache, data: dict, call_type: str, + max_parallel_requests: int, + tpm_limit: int, + rpm_limit: int, + request_count_api_key: str, ): - self.print_verbose(f"Inside Max Parallel Request Pre-Call Hook") - api_key = user_api_key_dict.api_key - max_parallel_requests = user_api_key_dict.max_parallel_requests or sys.maxsize - tpm_limit = user_api_key_dict.tpm_limit or sys.maxsize - rpm_limit = user_api_key_dict.rpm_limit or sys.maxsize - - if api_key is None: - return - - if ( - max_parallel_requests == sys.maxsize - and tpm_limit == sys.maxsize - and rpm_limit == sys.maxsize - ): - return - - self.user_api_key_cache = cache # save the api key cache for updating the value - # ------------ - # 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}" - - request_count_api_key = f"{api_key}::{precise_minute}::request_count" - - # CHECK IF REQUEST ALLOWED current = cache.get_cache( key=request_count_api_key ) # {"current_requests": 1, "current_tpm": 1, "current_rpm": 10} - self.print_verbose(f"current: {current}") + # print(f"current: {current}") if current is None: new_val = { "current_requests": 1, @@ -88,10 +63,117 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): status_code=429, detail="Max parallel request 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 Parallel Request Pre-Call Hook") + api_key = user_api_key_dict.api_key + max_parallel_requests = user_api_key_dict.max_parallel_requests or sys.maxsize + tpm_limit = user_api_key_dict.tpm_limit or sys.maxsize + rpm_limit = user_api_key_dict.rpm_limit or sys.maxsize + + if api_key is None: + return + + self.user_api_key_cache = cache # save the api key cache for updating the value + # ------------ + # 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}" + + request_count_api_key = f"{api_key}::{precise_minute}::request_count" + + # CHECK IF REQUEST ALLOWED for key + current = cache.get_cache( + key=request_count_api_key + ) # {"current_requests": 1, "current_tpm": 1, "current_rpm": 10} + self.print_verbose(f"current: {current}") + if ( + max_parallel_requests == sys.maxsize + and tpm_limit == sys.maxsize + and rpm_limit == sys.maxsize + ): + pass + elif current is None: + new_val = { + "current_requests": 1, + "current_tpm": 0, + "current_rpm": 0, + } + cache.set_cache(request_count_api_key, new_val) + elif ( + int(current["current_requests"]) < max_parallel_requests + and current["current_tpm"] < tpm_limit + and current["current_rpm"] < rpm_limit + ): + # Increase count for this token + new_val = { + "current_requests": current["current_requests"] + 1, + "current_tpm": current["current_tpm"], + "current_rpm": current["current_rpm"], + } + cache.set_cache(request_count_api_key, new_val) + else: + raise HTTPException( + status_code=429, detail="Max parallel request limit reached." + ) + + # print("checking if user is in rate limits for user_id") + + # check if REQUEST ALLOWED for user_id + user_id = user_api_key_dict.user_id + _user_id_rate_limits = user_api_key_dict.user_id_rate_limits + + # print( + # f"USER ID RATE LIMITS: {_user_id_rate_limits}" + # ) + # get user tpm/rpm limits + + if _user_id_rate_limits is None: + return + user_tpm_limit = _user_id_rate_limits.get("tpm_limit") + user_rpm_limit = _user_id_rate_limits.get("rpm_limit") + 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 + 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}" + + 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, + cache=cache, + data=data, + call_type=call_type, + max_parallel_requests=max_parallel_requests, + request_count_api_key=request_count_api_key, + tpm_limit=user_tpm_limit, + rpm_limit=user_rpm_limit, + ) + return + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): try: self.print_verbose(f"INSIDE parallel request limiter ASYNC SUCCESS LOGGING") user_api_key = kwargs["litellm_params"]["metadata"]["user_api_key"] + user_api_key_user_id = kwargs["litellm_params"]["metadata"][ + "user_api_key_user_id" + ] if user_api_key is None: return @@ -121,7 +203,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): } # ------------ - # Update usage + # Update usage - API Key # ------------ new_val = { @@ -136,6 +218,41 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): self.user_api_key_cache.set_cache( request_count_api_key, new_val, ttl=60 ) # store in cache for 1 min. + + # ------------ + # Update usage - User + # ------------ + if user_api_key_user_id is None: + return + + total_tokens = 0 + + if isinstance(response_obj, ModelResponse): + total_tokens = response_obj.usage.total_tokens + + request_count_api_key = ( + f"{user_api_key_user_id}::{precise_minute}::request_count" + ) + + current = self.user_api_key_cache.get_cache(key=request_count_api_key) or { + "current_requests": 1, + "current_tpm": total_tokens, + "current_rpm": 1, + } + + new_val = { + "current_requests": max(current["current_requests"] - 1, 0), + "current_tpm": current["current_tpm"] + total_tokens, + "current_rpm": current["current_rpm"] + 1, + } + + self.print_verbose( + f"updated_value in success call: {new_val}, precise_minute: {precise_minute}" + ) + self.user_api_key_cache.set_cache( + request_count_api_key, new_val, ttl=60 + ) # store in cache for 1 min. + except Exception as e: self.print_verbose(e) # noqa diff --git a/litellm/tests/test_parallel_request_limiter.py b/litellm/tests/test_parallel_request_limiter.py index 17d79c36c9b..e402b617b78 100644 --- a/litellm/tests/test_parallel_request_limiter.py +++ b/litellm/tests/test_parallel_request_limiter.py @@ -139,6 +139,56 @@ async def test_pre_call_hook_tpm_limits(): assert e.status_code == 429 +@pytest.mark.asyncio +async def test_pre_call_hook_user_tpm_limits(): + """ + Test if error raised on hitting tpm limits + """ + # create user with tpm/rpm limits + + _api_key = "sk-12345" + user_api_key_dict = UserAPIKeyAuth( + api_key=_api_key, + user_id="ishaan", + user_id_rate_limits={"tpm_limit": 9, "rpm_limit": 10}, + ) + res = dict(user_api_key_dict) + print("dict user", res) + local_cache = DualCache() + parallel_request_handler = MaxParallelRequestsHandler() + + await parallel_request_handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, cache=local_cache, data={}, call_type="" + ) + + kwargs = { + "litellm_params": { + "metadata": {"user_api_key_user_id": "ishaan", "user_api_key": "gm"} + } + } + + await parallel_request_handler.async_log_success_event( + kwargs=kwargs, + response_obj=litellm.ModelResponse(usage=litellm.Usage(total_tokens=10)), + start_time="", + end_time="", + ) + + ## Expected cache val: {"current_requests": 0, "current_tpm": 0, "current_rpm": 1} + + try: + await parallel_request_handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={}, + call_type="", + ) + + pytest.fail(f"Expected call to fail") + except Exception as e: + assert e.status_code == 429 + + @pytest.mark.asyncio async def test_success_call_hook(): """ From 1fff8f81052fd935d7ea90b70d430f66b66935fb Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Thu, 22 Feb 2024 18:50:02 -0800 Subject: [PATCH 2/4] (fix) don't double check curr data and time --- litellm/proxy/hooks/parallel_request_limiter.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/litellm/proxy/hooks/parallel_request_limiter.py b/litellm/proxy/hooks/parallel_request_limiter.py index 021fbc5fb7f..df21b573b2b 100644 --- a/litellm/proxy/hooks/parallel_request_limiter.py +++ b/litellm/proxy/hooks/parallel_request_limiter.py @@ -147,11 +147,6 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): user_rpm_limit = sys.maxsize # now do the same tpm/rpm checks - 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}" - 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}") From a13243652f183d042ac4ceccce124bb68af263be Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Thu, 22 Feb 2024 19:16:22 -0800 Subject: [PATCH 3/4] (fix) failing parallel_Request_limiter test --- litellm/proxy/hooks/parallel_request_limiter.py | 17 ++++++----------- 1 file changed, 6 insertions(+), 11 deletions(-) diff --git a/litellm/proxy/hooks/parallel_request_limiter.py b/litellm/proxy/hooks/parallel_request_limiter.py index df21b573b2b..fb61fe3da6e 100644 --- a/litellm/proxy/hooks/parallel_request_limiter.py +++ b/litellm/proxy/hooks/parallel_request_limiter.py @@ -126,18 +126,12 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): status_code=429, detail="Max parallel request limit reached." ) - # print("checking if user is in rate limits for user_id") - # check if REQUEST ALLOWED for user_id user_id = user_api_key_dict.user_id _user_id_rate_limits = user_api_key_dict.user_id_rate_limits - # print( - # f"USER ID RATE LIMITS: {_user_id_rate_limits}" - # ) # get user tpm/rpm limits - - if _user_id_rate_limits is None: + if _user_id_rate_limits is None or _user_id_rate_limits == {}: return user_tpm_limit = _user_id_rate_limits.get("tpm_limit") user_rpm_limit = _user_id_rate_limits.get("rpm_limit") @@ -155,7 +149,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): cache=cache, data=data, call_type=call_type, - max_parallel_requests=max_parallel_requests, + max_parallel_requests=sys.maxsize, # TODO: Support max parallel requests for a user request_count_api_key=request_count_api_key, tpm_limit=user_tpm_limit, rpm_limit=user_rpm_limit, @@ -166,9 +160,10 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): try: self.print_verbose(f"INSIDE parallel request limiter ASYNC SUCCESS LOGGING") user_api_key = kwargs["litellm_params"]["metadata"]["user_api_key"] - user_api_key_user_id = kwargs["litellm_params"]["metadata"][ - "user_api_key_user_id" - ] + user_api_key_user_id = kwargs["litellm_params"]["metadata"].get( + "user_api_key_user_id", None + ) + if user_api_key is None: return From 2122a7cf1e343cc0f102c14e5ee206ce329fb3cd Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Thu, 22 Feb 2024 19:23:16 -0800 Subject: [PATCH 4/4] (docs) set user tpm/rpm limits --- docs/my-website/docs/proxy/users.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/my-website/docs/proxy/users.md b/docs/my-website/docs/proxy/users.md index 3eb0cb808b2..159b311a911 100644 --- a/docs/my-website/docs/proxy/users.md +++ b/docs/my-website/docs/proxy/users.md @@ -279,9 +279,9 @@ curl 'http://0.0.0.0:8000/key/generate' \ ## Set Rate Limits You can set: +- tpm limits (tokens per minute) +- rpm limits (requests per minute) - max parallel requests -- tpm limits -- rpm limits