diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 4dbae6394f6..fb61a82e21a 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2922,6 +2922,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 = [] @@ -2941,6 +2942,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 @@ -2990,6 +2992,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 b39b1f330b3..5c1b3f59105 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -477,6 +477,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 @@ -561,6 +564,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: @@ -1479,6 +1484,7 @@ async def _user_api_key_auth_builder( 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, @@ -1493,12 +1499,14 @@ async def _user_api_key_auth_builder( 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 1ed009087cc..15b74829032 100644 --- a/litellm/proxy/hooks/parallel_request_limiter.py +++ b/litellm/proxy/hooks/parallel_request_limiter.py @@ -407,7 +407,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): if user_rpm_limit is None: user_rpm_limit = sys.maxsize - user_max_parallel_requests = user_api_key_dict.max_parallel_requests + 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 @@ -774,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)