scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-10 03:28:53 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 84

feat(proxy): implement parallel request limiter and related authentication modules

Browse source
This commit is contained in:
Nischalgouda 2026-09-12 12:45:31 +05:30
parent 1442667e95
commit bf941dc92b
3 changed files with 64 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
litellm/proxy/_types.py
View file

@ -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

8
litellm/proxy/auth/user_api_key_auth.py
View file

@ -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,

54
litellm/proxy/hooks/parallel_request_limiter.py
View file

@ -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)