diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index ae6c042ab3a..e0ad14e3cb1 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -3003,6 +3003,7 @@ class LiteLLM_VerificationTokenView(LiteLLM_VerificationToken): team_alias: str | None = None team_tpm_limit: int | None = None team_rpm_limit: int | None = None + team_max_parallel_requests: int | None = None team_max_budget: float | None = None team_soft_budget: float | None = None team_models: list = [] @@ -3022,6 +3023,7 @@ class LiteLLM_VerificationTokenView(LiteLLM_VerificationToken): end_user_id: str | None = None end_user_tpm_limit: int | None = None end_user_rpm_limit: int | None = None + end_user_max_parallel_requests: int | None = None end_user_max_budget: float | None = None end_user_model_max_budget: dict | None = None @@ -3071,6 +3073,7 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob api_key: str | None = None user_role: LitellmUserRoles | None = None + user_max_parallel_requests: int | None = None allowed_model_region: AllowedModelRegion | None = None parent_otel_span: Span | None = None rpm_limit_per_model: dict[str, int] | None = None diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 9828311112e..4a28c2a1080 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -526,6 +526,9 @@ def _apply_budget_limits_to_end_user_params( if budget_info.rpm_limit is not None: end_user_params["end_user_rpm_limit"] = budget_info.rpm_limit + if budget_info.max_parallel_requests is not None: + end_user_params["end_user_max_parallel_requests"] = budget_info.max_parallel_requests + if budget_info.max_budget is not None: end_user_params["end_user_max_budget"] = budget_info.max_budget @@ -610,6 +613,8 @@ def update_valid_token_with_end_user_params(valid_token: UserAPIKeyAuth, end_use valid_token.end_user_tpm_limit = end_user_params["end_user_tpm_limit"] if end_user_params.get("end_user_rpm_limit") is not None: valid_token.end_user_rpm_limit = end_user_params["end_user_rpm_limit"] + if end_user_params.get("end_user_max_parallel_requests") is not None: + valid_token.end_user_max_parallel_requests = end_user_params["end_user_max_parallel_requests"] if end_user_params.get("allowed_model_region") is not None: valid_token.allowed_model_region = end_user_params["allowed_model_region"] if end_user_params.get("end_user_model_max_budget") is not None: @@ -1525,6 +1530,12 @@ async def _user_api_key_auth_builder( user_id=user_id, user_email=user_email, team_id=team_id, + team_alias=(team_object.team_alias if team_object is not None else None), + team_tpm_limit=(team_object.tpm_limit if team_object is not None else None), + team_rpm_limit=(team_object.rpm_limit if team_object is not None else None), + team_max_parallel_requests=(team_object.max_parallel_requests if team_object is not None else None), + team_models=(team_object.models if team_object is not None else []), + team_metadata=(team_object.metadata if team_object is not None else None), org_id=org_id, end_user_id=end_user_id, parent_otel_span=parent_otel_span, @@ -1535,11 +1546,17 @@ async def _user_api_key_auth_builder( valid_token = UserAPIKeyAuth( api_key=None, team_id=team_id, + team_alias=(team_object.team_alias if team_object is not None else None), + team_tpm_limit=(team_object.tpm_limit if team_object is not None else None), + team_rpm_limit=(team_object.rpm_limit if team_object is not None else None), + team_max_parallel_requests=(team_object.max_parallel_requests if team_object is not None else None), + team_models=(team_object.models if team_object is not None else []), user_role=( LitellmUserRoles(user_object.user_role) if user_object is not None and user_object.user_role is not None else LitellmUserRoles.INTERNAL_USER ), + user_max_parallel_requests=(user_object.max_parallel_requests if user_object is not None else None), user_id=user_id, user_email=user_email, org_id=org_id, diff --git a/litellm/proxy/hooks/parallel_request_limiter.py b/litellm/proxy/hooks/parallel_request_limiter.py index b313cb64c3f..15b74829032 100644 --- a/litellm/proxy/hooks/parallel_request_limiter.py +++ b/litellm/proxy/hooks/parallel_request_limiter.py @@ -407,6 +407,10 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): if user_rpm_limit is None: user_rpm_limit = sys.maxsize + user_max_parallel_requests = getattr(user_api_key_dict, "user_max_parallel_requests", sys.maxsize) + if user_max_parallel_requests is None: + user_max_parallel_requests = sys.maxsize + 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( @@ -414,7 +418,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): cache=cache, data=data, call_type=call_type, - max_parallel_requests=sys.maxsize, # TODO: Support max parallel requests for a user + max_parallel_requests=user_max_parallel_requests, #TODO:Support max parallel requests for a user(has been implemented) current=cache_objects["request_count_user_id"], request_count_api_key=request_count_api_key, tpm_limit=user_tpm_limit, @@ -435,6 +439,10 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): if team_rpm_limit is None: team_rpm_limit = sys.maxsize + team_max_parallel_requests = getattr(user_api_key_dict, "team_max_parallel_requests", sys.maxsize) + if team_max_parallel_requests is None: + team_max_parallel_requests = sys.maxsize + 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( @@ -442,7 +450,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): cache=cache, data=data, call_type=call_type, - max_parallel_requests=sys.maxsize, # TODO: Support max parallel requests for a team + max_parallel_requests=team_max_parallel_requests, # TODO: Support max parallel requests for a team(has been implemented) current=cache_objects["request_count_team_id"], request_count_api_key=request_count_api_key, tpm_limit=team_tpm_limit, @@ -462,6 +470,10 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): if end_user_rpm_limit is None: end_user_rpm_limit = sys.maxsize + end_user_max_parallel_requests = getattr(user_api_key_dict, "end_user_max_parallel_requests", sys.maxsize) + if end_user_max_parallel_requests is None: + end_user_max_parallel_requests = sys.maxsize + # now do the same tpm/rpm checks request_count_api_key = f"{user_api_key_dict.end_user_id}::{precise_minute}::request_count" @@ -471,7 +483,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): cache=cache, data=data, call_type=call_type, - max_parallel_requests=sys.maxsize, # TODO: Support max parallel requests for an End-User + max_parallel_requests=end_user_max_parallel_requests, # TODO: Support max parallel requests for an End-User(has been implemented) request_count_api_key=request_count_api_key, current=cache_objects["request_count_end_user_id"], tpm_limit=end_user_tpm_limit, @@ -762,6 +774,58 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): ttl=60, litellm_parent_otel_span=litellm_parent_otel_span, ) # save in cache for up to 1 min. + + user_api_key_user_id = _metadata.get("user_api_key_user_id", None) + if user_api_key_user_id is not None: + request_count_user = f"{user_api_key_user_id}::{precise_minute}::request_count" + current_user = await self.internal_usage_cache.async_get_cache( + key=request_count_user, + litellm_parent_otel_span=litellm_parent_otel_span, + ) or {"current_requests": 1, "current_tpm": 0, "current_rpm": 0} + new_val_user = { + "current_requests": max(current_user["current_requests"] - 1, 0), + "current_tpm": current_user["current_tpm"], + "current_rpm": current_user["current_rpm"], + } + await self.internal_usage_cache.async_set_cache( + request_count_user, new_val_user, ttl=60, + litellm_parent_otel_span=litellm_parent_otel_span, + ) + + user_api_key_team_id = _metadata.get("user_api_key_team_id", None) + if user_api_key_team_id is not None: + request_count_team = f"{user_api_key_team_id}::{precise_minute}::request_count" + current_team = await self.internal_usage_cache.async_get_cache( + key=request_count_team, + litellm_parent_otel_span=litellm_parent_otel_span, + ) or {"current_requests": 1, "current_tpm": 0, "current_rpm": 0} + new_val_team = { + "current_requests": max(current_team["current_requests"] - 1, 0), + "current_tpm": current_team["current_tpm"], + "current_rpm": current_team["current_rpm"], + } + await self.internal_usage_cache.async_set_cache( + request_count_team, new_val_team, ttl=60, + litellm_parent_otel_span=litellm_parent_otel_span, + ) + + user_api_key_end_user_id = _metadata.get("user_api_key_end_user_id", None) + if user_api_key_end_user_id is not None: + request_count_end_user = f"{user_api_key_end_user_id}::{precise_minute}::request_count" + current_end_user = await self.internal_usage_cache.async_get_cache( + key=request_count_end_user, + litellm_parent_otel_span=litellm_parent_otel_span, + ) or {"current_requests": 1, "current_tpm": 0, "current_rpm": 0} + new_val_end_user = { + "current_requests": max(current_end_user["current_requests"] - 1, 0), + "current_tpm": current_end_user["current_tpm"], + "current_rpm": current_end_user["current_rpm"], + } + await self.internal_usage_cache.async_set_cache( + request_count_end_user, new_val_end_user, ttl=60, + litellm_parent_otel_span=litellm_parent_otel_span, + ) + except Exception as e: verbose_proxy_logger.exception("Inside Parallel Request Limiter: An exception occurred - %s", e)