tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-02 10:21:52 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add customer + model per key level multi-instance tpm/rpm limiting (#10518)

Browse Source 
* fix(redis_cache.py): handle multiple event loops

* fix(parallel_request_limiter_v2.py): add customer tpm limiting

* fix(parallel_request_limiter.py): add customer rpm limiting

* fix(parallel_request_limiter_v2.py): add model per key + customer tpm/rpm limiting

* fix(parallel_request_limiter_v2.py): make error more informative

* fix: fix ruff error

* fix: generate new poetry lock
This commit is contained in:
Krish Dholakia
2025-05-03 10:28:55 -07:00
committed by GitHub
parent 2138170861
commit 552b7e4013
3 changed files with 127 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files
enterprise/enterprise_hooks/parallel_request_limiter_v2.py
+74 -14
View File
@@ -26,6 +26,10 @@ from litellm._logging import verbose_proxy_logger
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.core_helpers import _get_parent_otel_span_from_kwargs
from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth
from litellm.proxy.auth.auth_utils import (
get_key_model_rpm_limit,
get_key_model_tpm_limit,
)
from litellm.router_strategy.base_routing_strategy import BaseRoutingStrategy
if TYPE_CHECKING:
@@ -85,25 +89,33 @@ class _PROXY_MaxParallelRequestsHandler(BaseRoutingStrategy, CustomLogger):
model: Optional[str],
rate_limit_type: Literal["key", "model_per_key", "user", "customer", "team"],
group: RateLimitGroups,
) -> str:
if rate_limit_type == "key":
) -> Optional[str]:
if rate_limit_type == "key" and user_api_key_dict.api_key is not None:
return (
f"{self.prefix}::{user_api_key_dict.api_key}::{precise_minute}::{group}"
)
elif rate_limit_type == "model_per_key" and model is not None:
elif (
rate_limit_type == "model_per_key"
and model is not None
and user_api_key_dict.api_key is not None
):
return f"{self.prefix}::{user_api_key_dict.api_key}::{model}::{precise_minute}::{group}"
elif rate_limit_type == "user":
elif rate_limit_type == "user" and user_api_key_dict.user_id is not None:
return (
f"{self.prefix}::{user_api_key_dict.user_id}::{precise_minute}::{group}"
)
elif rate_limit_type == "customer":
elif (
rate_limit_type == "customer" and user_api_key_dict.end_user_id is not None
):
return f"{self.prefix}::{user_api_key_dict.end_user_id}::{precise_minute}::{group}"
elif rate_limit_type == "team":
elif rate_limit_type == "team" and user_api_key_dict.team_id is not None:
return (
f"{self.prefix}::{user_api_key_dict.team_id}::{precise_minute}::{group}"
)
elif rate_limit_type == "model_per_key" and model is not None:
return f"{self.prefix}::{user_api_key_dict.api_key}::{model}::{precise_minute}::{group}"
else:
raise ValueError(f"Invalid rate limit type: {rate_limit_type}")
return None
def get_key_pattern_to_sync(self) -> Optional[str]:
return self.prefix + "::"
@@ -133,6 +145,8 @@ class _PROXY_MaxParallelRequestsHandler(BaseRoutingStrategy, CustomLogger):
rate_limit_type=rate_limit_type,
group=cast(RateLimitGroups, group),
)
if key is None:
continue
increment_list.append((key, increment_value_by_group[group]))
if (
@@ -153,7 +167,7 @@ class _PROXY_MaxParallelRequestsHandler(BaseRoutingStrategy, CustomLogger):
should_raise_error = should_raise_error or results[2] > tpm_limit
if should_raise_error:
raise self.raise_rate_limit_error(
additional_details=f"{CommonProxyErrors.max_parallel_request_limit_reached.value}. Hit limit for {rate_limit_type}. Current limits: max_parallel_requests: {max_parallel_requests}, tpm_limit: {tpm_limit}, rpm_limit: {rpm_limit}"
additional_details=f"{CommonProxyErrors.max_parallel_request_limit_reached.value}. Hit limit for {rate_limit_type}. Current usage: max_parallel_requests: {results[0]}, current_rpm: {results[1]}, current_tpm: {results[2]}. Current limits: max_parallel_requests: {max_parallel_requests}, rpm_limit: {rpm_limit}, tpm_limit: {tpm_limit}."
)
def time_to_next_minute(self) -> float:
@@ -233,7 +247,7 @@ class _PROXY_MaxParallelRequestsHandler(BaseRoutingStrategy, CustomLogger):
local_only=True,
litellm_parent_otel_span=user_api_key_dict.parent_otel_span,
)
_model = data.get("model", None)
requested_model = data.get("model", None)
current_date = datetime.now().strftime("%Y-%m-%d")
current_hour = datetime.now().strftime("%H")
@@ -254,7 +268,7 @@ class _PROXY_MaxParallelRequestsHandler(BaseRoutingStrategy, CustomLogger):
rate_limit_type="key",
)
)
elif user_api_key_dict.user_id is not None:
if user_api_key_dict.user_id is not None:
# CHECK IF REQUEST ALLOWED for key
tasks.append(
self.check_key_in_limits_v2(
@@ -267,7 +281,7 @@ class _PROXY_MaxParallelRequestsHandler(BaseRoutingStrategy, CustomLogger):
rate_limit_type="user",
)
)
elif user_api_key_dict.team_id is not None:
if user_api_key_dict.team_id is not None:
tasks.append(
self.check_key_in_limits_v2(
user_api_key_dict=user_api_key_dict,
@@ -279,6 +293,49 @@ class _PROXY_MaxParallelRequestsHandler(BaseRoutingStrategy, CustomLogger):
rate_limit_type="team",
)
)
if user_api_key_dict.end_user_id is not None:
tasks.append(
self.check_key_in_limits_v2(
user_api_key_dict=user_api_key_dict,
data=data,
max_parallel_requests=None,
precise_minute=precise_minute,
tpm_limit=user_api_key_dict.end_user_tpm_limit,
rpm_limit=user_api_key_dict.end_user_rpm_limit,
rate_limit_type="customer",
)
)
if requested_model and (
get_key_model_tpm_limit(user_api_key_dict) is not None
or get_key_model_rpm_limit(user_api_key_dict) is not None
):
_tpm_limit_for_key_model = get_key_model_tpm_limit(user_api_key_dict) or {}
_rpm_limit_for_key_model = get_key_model_rpm_limit(user_api_key_dict) or {}
should_check_rate_limit = False
if requested_model in _tpm_limit_for_key_model:
should_check_rate_limit = True
elif requested_model in _rpm_limit_for_key_model:
should_check_rate_limit = True
if should_check_rate_limit:
model_specific_tpm_limit: Optional[int] = None
model_specific_rpm_limit: Optional[int] = None
if requested_model in _tpm_limit_for_key_model:
model_specific_tpm_limit = _tpm_limit_for_key_model[requested_model]
if requested_model in _rpm_limit_for_key_model:
model_specific_rpm_limit = _rpm_limit_for_key_model[requested_model]
tasks.append(
self.check_key_in_limits_v2(
user_api_key_dict=user_api_key_dict,
data=data,
max_parallel_requests=None,
precise_minute=precise_minute,
tpm_limit=model_specific_tpm_limit,
rpm_limit=model_specific_rpm_limit,
rate_limit_type="model_per_key",
)
)
await asyncio.gather(*tasks)
return
@@ -298,7 +355,7 @@ class _PROXY_MaxParallelRequestsHandler(BaseRoutingStrategy, CustomLogger):
"rpm": 0,
}
rate_limit_types = ["key", "user", "customer", "team"]
rate_limit_types = ["key", "user", "customer", "team", "model_per_key"]
for rate_limit_type in rate_limit_types:
for group in ["request_count", "rpm", "tpm"]:
key = self._get_current_usage_key(
@@ -308,7 +365,8 @@ class _PROXY_MaxParallelRequestsHandler(BaseRoutingStrategy, CustomLogger):
rate_limit_type=cast(RateLimitTypes, rate_limit_type),
group=cast(RateLimitGroups, group),
)
if key is None:
continue
increment_list.append((key, increment_value_by_group[group]))
if increment_list: # Only call if we have values to increment
@@ -344,7 +402,9 @@ class _PROXY_MaxParallelRequestsHandler(BaseRoutingStrategy, CustomLogger):
user_api_key_team_id = kwargs["litellm_params"]["metadata"].get(
"user_api_key_team_id", None
)
user_api_key_end_user_id = kwargs.get("user")
user_api_key_end_user_id = kwargs.get("user") or kwargs["litellm_params"][
"metadata"
].get("user_api_key_end_user_id", None)
# ------------
# Setup values
litellm/caching/redis_cache.py
+22 -6
View File
@@ -14,7 +14,7 @@ import inspect
import json
import time
from datetime import timedelta
from typing import TYPE_CHECKING, Any, List, Optional, Tuple, Union
from typing import TYPE_CHECKING, Any, List, Optional, Tuple, Union, cast
import litellm
from litellm._logging import print_verbose, verbose_logger
@@ -83,7 +83,9 @@ class RedisCache(BaseCache):
redis_kwargs.update(kwargs)
self.redis_client = get_redis_client(**redis_kwargs)
self.redis_async_client: Optional[async_redis_client] = None
self.redis_async_client: Optional[
Union[async_redis_client, async_redis_cluster_client]
] = None
self.redis_kwargs = redis_kwargs
self.async_redis_conn_pool = get_redis_connection_pool(**redis_kwargs)
@@ -134,13 +136,27 @@ class RedisCache(BaseCache):
def init_async_client(
self,
) -> Union[async_redis_client, async_redis_cluster_client]:
from .._redis import get_redis_async_client
from litellm import in_memory_llm_clients_cache
if self.redis_async_client is None:
self.redis_async_client = get_redis_async_client(
from .._redis import get_redis_async_client, get_redis_connection_pool
cached_client = in_memory_llm_clients_cache.get_cache(key="async-redis-client")
if cached_client is not None:
redis_async_client = cast(
Union[async_redis_client, async_redis_cluster_client], cached_client
)
else:
# Create new connection pool and client for current event loop
self.async_redis_conn_pool = get_redis_connection_pool(**self.redis_kwargs)
redis_async_client = get_redis_async_client(
connection_pool=self.async_redis_conn_pool, **self.redis_kwargs
)
return self.redis_async_client
in_memory_llm_clients_cache.set_cache(
key="async-redis-client", value=self.redis_async_client
)
self.redis_async_client = redis_async_client # type: ignore
return redis_async_client
def check_and_fix_namespace(self, key: str) -> str:
"""
tests/enterprise/enterprise_hooks/test_parallel_request_limiter_v2.py
+31 -9
View File
@@ -106,9 +106,9 @@ async def test_normal_router_call_v2(monkeypatch):
"rate_limit_object",
[
"key",
# "model_per_key",
"model_per_key",
"user",
# "customer",
"customer",
"team",
],
)
@@ -154,6 +154,13 @@ async def test_normal_router_call_tpm(monkeypatch, rate_limit_object):
user_api_key_dict = UserAPIKeyAuth(user_id="12345", user_tpm_limit=10)
elif rate_limit_object == "team":
user_api_key_dict = UserAPIKeyAuth(team_id="12345", team_tpm_limit=10)
elif rate_limit_object == "customer":
user_api_key_dict = UserAPIKeyAuth(end_user_id="12345", end_user_tpm_limit=10)
elif rate_limit_object == "model_per_key":
user_api_key_dict = UserAPIKeyAuth(
api_key=_api_key,
metadata={"model_tpm_limit": {"azure-model": 10}},
)
local_cache = DualCache()
parallel_request_handler = _PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(local_cache)
@@ -161,7 +168,10 @@ async def test_normal_router_call_tpm(monkeypatch, rate_limit_object):
monkeypatch.setattr(litellm, "callbacks", [parallel_request_handler])
await parallel_request_handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict, cache=local_cache, data={}, call_type=""
user_api_key_dict=user_api_key_dict,
cache=local_cache,
data={"model": "azure-model"},
call_type="",
)
current_date = datetime.now().strftime("%Y-%m-%d")
@@ -171,7 +181,7 @@ async def test_normal_router_call_tpm(monkeypatch, rate_limit_object):
request_count_api_key = parallel_request_handler._get_current_usage_key(
user_api_key_dict=user_api_key_dict,
precise_minute=precise_minute,
model=None,
model="azure-model",
rate_limit_type=rate_limit_object,
group="tpm",
)
@@ -191,6 +201,7 @@ async def test_normal_router_call_tpm(monkeypatch, rate_limit_object):
"user_api_key": _api_key,
"user_api_key_user_id": user_api_key_dict.user_id,
"user_api_key_team_id": user_api_key_dict.team_id,
"user_api_key_end_user_id": user_api_key_dict.end_user_id,
},
mock_response="hello",
)
@@ -210,9 +221,9 @@ async def test_normal_router_call_tpm(monkeypatch, rate_limit_object):
"rate_limit_object",
[
"key",
# "model_per_key",
"model_per_key",
"user",
# "customer",
"customer",
"team",
],
)
@@ -258,6 +269,13 @@ async def test_normal_router_call_rpm(monkeypatch, rate_limit_object):
user_api_key_dict = UserAPIKeyAuth(user_id="12345", user_rpm_limit=1)
elif rate_limit_object == "team":
user_api_key_dict = UserAPIKeyAuth(team_id="12345", team_rpm_limit=1)
elif rate_limit_object == "customer":
user_api_key_dict = UserAPIKeyAuth(end_user_id="12345", end_user_rpm_limit=1)
elif rate_limit_object == "model_per_key":
user_api_key_dict = UserAPIKeyAuth(
api_key=_api_key,
metadata={"model_rpm_limit": {"azure-model": 1}},
)
local_cache = DualCache()
parallel_request_handler = _PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(local_cache)
@@ -265,7 +283,10 @@ async def test_normal_router_call_rpm(monkeypatch, rate_limit_object):
monkeypatch.setattr(litellm, "callbacks", [parallel_request_handler])
await parallel_request_handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict, cache=local_cache, data={}, call_type=""
user_api_key_dict=user_api_key_dict,
cache=local_cache,
data={"model": "azure-model"},
call_type="",
)
current_date = datetime.now().strftime("%Y-%m-%d")
@@ -275,7 +296,7 @@ async def test_normal_router_call_rpm(monkeypatch, rate_limit_object):
request_count_api_key = parallel_request_handler._get_current_usage_key(
user_api_key_dict=user_api_key_dict,
precise_minute=precise_minute,
model=None,
model="azure-model",
rate_limit_type=rate_limit_object,
group="rpm",
)
@@ -295,6 +316,7 @@ async def test_normal_router_call_rpm(monkeypatch, rate_limit_object):
"user_api_key": _api_key,
"user_api_key_user_id": user_api_key_dict.user_id,
"user_api_key_team_id": user_api_key_dict.team_id,
"user_api_key_end_user_id": user_api_key_dict.end_user_id,
},
mock_response="hello",
)
@@ -313,7 +335,7 @@ async def test_normal_router_call_rpm(monkeypatch, rate_limit_object):
await parallel_request_handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=local_cache,
data={},
data={"model": "azure-model"},
call_type="",
)