tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-18 04:28:19 +00:00
Code Issues Packages Projects Releases Wiki Activity

Simplify experimental multi-instance rate limiter - more accurate (#11424)

Browse Source 
* refactor: comment out circuit breaker

causes incorrect rate limiting in high traffic

* fix(base_routing_strategy.py): don't reset value if redis val is lower than current in-memory value

Fixes issue where redis might be trailing in-memory value

* fix(parallel_request_limiter_v2.py): if in-memory higher than redis, don't reset value; add previous slot keys to redis increment to correctly 'get' them

* fix(parallel_request_limiter_v3.py): v3 implementation of parallel request limiter

does not use background redis syncing - increments redis in call

 simplify rate limiting logic, to improve accuracy

* fix: fix ruff errors

* fix(parallel_request_limiter_v3.py): don't decrement limit on post call success - causes double decrements

* fix(parallel_request_limiter_v3.py): working accurate multi-instance logic

ensured just 100 requests allowed on 100 users, 10 ramp up, 100 rpm limit key, 2 instances

* fix(parallel_request_limiter_v3.py): working accurate rate limiting with time window resets

allows rate limiting to work across multiple windows

* test: add unit tests for v3 rate limiter

* fix(parallel_request_limiter_v3.py): return window value into in-memory cache

allows in-memory cache checks to be used correctly

* refactor(parallel_request_limiter_v3.py): refactor rate limiting to work for multiple window/counter key pairs

enables using for user/team/model rate limiting

* feat(parallel_request_limiter_v3.py): working rate limiting, across key/user/team/end-user

* fix(parallel_request_limiter_v3.py): add model specific rate limiting

* fix(parallel_request_limiter_v3.py): ignore if no rate limits set

skip unecessary rate limit checks - if no limits set

* fix(parallel_request_limiter_v3.py): initial commit bringing token rate limits back

* fix(parallel_request_limiter_v3.py): increment by value in list + update assertions to handle tokens + max parallel requests

* test(parallel_request_limiter_v3.py): more testing

* fix(parallel_request_limiter.py): working in-memory cache limiter

* fix(redis_cache.py): ignore linting error - use safe hasattr

* fix(parallel_request_limiter_v3.py): fix linting error

* refactor: remove redundant parallel_Request_limiter_v2.py

old / inaccurate implementation

* test: update tests

* style: cleanup

* test: update test

* docs(config_settings.md): document new env var

* test(test_base_routing_strategy.py): update test
This commit is contained in:
Krish Dholakia
2025-06-07 11:10:55 -07:00
committed by GitHub
parent 1d86fc84fe
commit c42740a4b9
13 changed files with 1246 additions and 1233 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/my-website/docs/proxy/config_settings.md
+1
View File
@@ -521,6 +521,7 @@ router_settings:
| LITELLM_LOCAL_MODEL_COST_MAP | Local configuration for model cost mapping in LiteLLM
| LITELLM_LOG | Enable detailed logging for LiteLLM
| LITELLM_MODE | Operating mode for LiteLLM (e.g., production, development)
| LITELLM_RATE_LIMIT_WINDOW_SIZE | Rate limit window size for LiteLLM. Default is 60
| LITELLM_SALT_KEY | Salt key for encryption in LiteLLM
| LITELLM_SECRET_AWS_KMS_LITELLM_LICENSE | AWS KMS encrypted license for LiteLLM
| LITELLM_TOKEN | Access token for LiteLLM integration
litellm/caching/dual_cache.py
+28
View File
@@ -14,6 +14,9 @@ import traceback
from concurrent.futures import ThreadPoolExecutor
from typing import TYPE_CHECKING, Any, List, Optional, Union
if TYPE_CHECKING:
from litellm.types.caching import RedisPipelineIncrementOperation
import litellm
from litellm._logging import print_verbose, verbose_logger
@@ -373,6 +376,31 @@ class DualCache(BaseCache):
except Exception as e:
raise e # don't log if exception is raised
async def async_increment_cache_pipeline(
self,
increment_list: List["RedisPipelineIncrementOperation"],
local_only: bool = False,
parent_otel_span: Optional[Span] = None,
**kwargs,
) -> Optional[List[float]]:
try:
result: Optional[List[float]] = None
if self.in_memory_cache is not None:
result = await self.in_memory_cache.async_increment_pipeline(
increment_list=increment_list,
parent_otel_span=parent_otel_span,
)
if self.redis_cache is not None and local_only is False:
result = await self.redis_cache.async_increment_pipeline(
increment_list=increment_list,
parent_otel_span=parent_otel_span,
)
return result
except Exception as e:
raise e # don't log if exception is raised
async def async_set_cache_sadd(
self, key, value: List, local_only: bool = False, **kwargs
) -> None:
litellm/caching/in_memory_cache.py
+16 -2
View File
@@ -11,7 +11,10 @@ Has 4 methods:
import json
import sys
import time
from typing import Any, List, Optional
from typing import TYPE_CHECKING, Any, List, Optional
if TYPE_CHECKING:
from litellm.types.caching import RedisPipelineIncrementOperation
from pydantic import BaseModel
@@ -140,7 +143,7 @@ class InMemoryCache(BaseCache):
self.cache_dict[key] = value
if self.allow_ttl_override(key): # if ttl is not set, set it to default ttl
if "ttl" in kwargs and kwargs["ttl"] is not None:
self.ttl_dict[key] = time.time() + kwargs["ttl"]
self.ttl_dict[key] = time.time() + float(kwargs["ttl"])
else:
self.ttl_dict[key] = time.time() + self.default_ttl
@@ -219,6 +222,17 @@ class InMemoryCache(BaseCache):
await self.async_set_cache(key, value, **kwargs)
return value
async def async_increment_pipeline(
self, increment_list: List["RedisPipelineIncrementOperation"], **kwargs
) -> Optional[List[float]]:
results = []
for increment in increment_list:
result = await self.async_increment(
increment["key"], increment["increment_value"], **kwargs
)
results.append(result)
return results
def flush_cache(self):
self.cache_dict.clear()
self.ttl_dict.clear()
litellm/caching/redis_cache.py
+35 -4
View File
@@ -294,6 +294,36 @@ class RedisCache(BaseCache):
)
raise e
def async_register_script(self, script: str) -> Any:
"""
Register a Lua script with Redis asynchronously.
Works with both standalone Redis and Redis Cluster.
Args:
script (str): The Lua script to register
Returns:
Any: A script object that can be called with keys and args
"""
try:
_redis_client = self.init_async_client()
# For standalone Redis
if hasattr(_redis_client, "register_script"):
return _redis_client.register_script(script) # type: ignore
# For Redis Cluster
elif hasattr(_redis_client, "script_load"):
# Load the script and get its SHA
script_sha = _redis_client.script_load(script) # type: ignore
# Return a callable that uses evalsha
async def script_callable(keys: List[str], args: List[Any]) -> Any:
return _redis_client.evalsha(script_sha, len(keys), *keys, *args) # type: ignore
return script_callable
except Exception as e:
verbose_logger.error(f"Error registering Redis script: {str(e)}")
raise e
async def async_set_cache(self, key, value, **kwargs):
from redis.asyncio import Redis
@@ -980,8 +1010,11 @@ class RedisCache(BaseCache):
pipe.expire(cache_key, _td)
# Execute the pipeline and return results
results = await pipe.execute()
print_verbose(f"Increment ASYNC Redis Cache PIPELINE: results: {results}")
return results
# only return float values
verbose_logger.debug(
f"Increment ASYNC Redis Cache PIPELINE: results: {results}"
)
return [r for r in results if isinstance(r, float)]
async def async_increment_pipeline(
self, increment_list: List[RedisPipelineIncrementOperation], **kwargs
@@ -1011,8 +1044,6 @@ class RedisCache(BaseCache):
async with _redis_client.pipeline(transaction=False) as pipe:
results = await self._pipeline_increment_helper(pipe, increment_list)
print_verbose(f"pipeline increment results: {results}")
## LOGGING ##
end_time = time.time()
_duration = end_time - start_time
litellm/proxy/hooks/__init__.py
+2 -2
View File
@@ -5,7 +5,7 @@ from . import *
from .cache_control_check import _PROXY_CacheControlCheck
from .max_budget_limiter import _PROXY_MaxBudgetLimiter
from .parallel_request_limiter import _PROXY_MaxParallelRequestsHandler
from .parallel_request_limiter_v2 import _PROXY_MaxParallelRequestsHandler_v2
from .parallel_request_limiter_v3 import _PROXY_MaxParallelRequestsHandler_v3
### CHECK IF ENTERPRISE HOOKS ARE AVAILABLE ###
@@ -23,7 +23,7 @@ PROXY_HOOKS = {
## FEATURE FLAG HOOKS ##
if os.getenv("EXPERIMENTAL_MULTI_INSTANCE_RATE_LIMITING", "false").lower() == "true":
PROXY_HOOKS["parallel_request_limiter"] = _PROXY_MaxParallelRequestsHandler_v2
PROXY_HOOKS["parallel_request_limiter"] = _PROXY_MaxParallelRequestsHandler_v3
### update PROXY_HOOKS with ENTERPRISE_PROXY_HOOKS ###
litellm/proxy/hooks/parallel_request_limiter_v2.py
-573
View File
@@ -1,573 +0,0 @@
"""
V2 Implementation of Parallel Requests, TPM, RPM Limiting on the proxy
Designed to work on a multi-instance setup, where multiple instances are writing to redis simultaneously
"""
import asyncio
import sys
from datetime import datetime, timedelta
from typing import (
TYPE_CHECKING,
Any,
List,
Literal,
Optional,
Tuple,
TypedDict,
Union,
cast,
)
from fastapi import HTTPException
import litellm
from litellm import DualCache, ModelResponse
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:
from opentelemetry.trace import Span as _Span
from litellm.proxy.utils import InternalUsageCache as _InternalUsageCache
Span = Union[_Span, Any]
InternalUsageCache = _InternalUsageCache
else:
Span = Any
InternalUsageCache = Any
class CacheObject(TypedDict):
current_global_requests: Optional[dict]
request_count_api_key: Optional[int]
request_count_api_key_model: Optional[dict]
request_count_user_id: Optional[dict]
request_count_team_id: Optional[dict]
request_count_end_user_id: Optional[dict]
rpm_api_key: Optional[int]
tpm_api_key: Optional[int]
RateLimitGroups = Literal["request_count", "tpm", "rpm"]
RateLimitTypes = Literal["key", "model_per_key", "user", "customer", "team"]
class _PROXY_MaxParallelRequestsHandler_v2(BaseRoutingStrategy, CustomLogger):
# Class variables or attributes
def __init__(self, internal_usage_cache: InternalUsageCache):
self.internal_usage_cache = internal_usage_cache
BaseRoutingStrategy.__init__(
self,
dual_cache=internal_usage_cache.dual_cache,
should_batch_redis_writes=True,
default_sync_interval=1,
)
def print_verbose(self, print_statement):
try:
verbose_proxy_logger.debug(print_statement)
if litellm.set_verbose:
print(print_statement) # noqa
except Exception:
pass
@property
def prefix(self) -> str:
return "parallel_request_limiter_v2"
def _get_current_usage_key(
self,
user_api_key_dict: UserAPIKeyAuth,
precise_minute: str,
model: Optional[str],
rate_limit_type: Literal["key", "model_per_key", "user", "customer", "team"],
group: RateLimitGroups,
) -> 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
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" 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" 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" 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:
return None
def get_key_pattern_to_sync(self) -> Optional[str]:
return self.prefix + "::"
def _get_slots_to_check(self, current_slot: int) -> List[str]:
slots_to_check = []
current_time = datetime.now()
for i in range(4):
slot_number = (current_slot - i) % 4 # This ensures we wrap around properly
minute = current_time.minute
hour = current_time.hour
# If we need to look at previous minute
if current_slot - i < 0:
if minute == 0:
# If we're at minute 0, go to previous hour
hour = (current_time.hour - 1) % 24
minute = 59
else:
minute = current_time.minute - 1
slot_key = f"{current_time.strftime('%Y-%m-%d')}-{hour:02d}-{minute:02d}-{slot_number}"
slots_to_check.append(slot_key)
return slots_to_check
async def check_key_in_limits_v2(
self,
user_api_key_dict: UserAPIKeyAuth,
data: dict,
max_parallel_requests: Optional[int],
precise_minute: str,
tpm_limit: Optional[int],
rpm_limit: Optional[int],
rate_limit_type: Literal["key", "model_per_key", "user", "customer", "team"],
):
## INCREMENT CURRENT USAGE
increment_list: List[Tuple[str, int]] = []
decrement_list: List[Tuple[str, int]] = []
slots_to_check: List[str] = []
increment_value_by_group = {
"request_count": 1,
"tpm": 0,
"rpm": 1,
}
# Get current time and calculate the last 4 15s slots
current_time = datetime.now()
current_slot = (
current_time.second // 15
) # This gives us 0-3 for the current 15s slot
slots_to_check = self._get_slots_to_check(current_slot)
slot_cache_keys = []
# Calculate the last 4 slots, handling minute boundaries
# For each slot, create keys for all rate limit groups
for slot_key in slots_to_check:
for group in ["request_count", "rpm", "tpm"]:
key = self._get_current_usage_key(
user_api_key_dict=user_api_key_dict,
precise_minute=slot_key,
model=data.get("model", None),
rate_limit_type=rate_limit_type,
group=cast(RateLimitGroups, group),
)
if key is None:
continue
# Only increment the current slot
if slot_key == slots_to_check[0]:
increment_list.append((key, increment_value_by_group[group]))
decrement_list.append(
(key, -1 if increment_value_by_group[group] == 1 else 0)
)
else:
self.add_to_in_memory_keys_to_update(key=key)
slot_cache_keys.append(key)
if (
not max_parallel_requests and not rpm_limit and not tpm_limit
): # no rate limits
return
# Use the existing atomic increment-and-check functionality
await self._increment_value_list_in_current_window(
increment_list=increment_list,
ttl=60,
)
# Get the current values for all slots to check limits
current_values = await self.internal_usage_cache.async_batch_get_cache(
slot_cache_keys
)
if current_values is None:
current_values = [None] * len(slot_cache_keys)
# Calculate totals across all slots, handling None values
# Group values by type (request_count, rpm, tpm)
request_counts = []
rpm_counts = []
tpm_counts = []
for i in range(0, len(current_values), 3):
request_counts.append(
current_values[i] if current_values[i] is not None else 0
)
rpm_counts.append(
current_values[i + 1] if current_values[i + 1] is not None else 0
)
tpm_counts.append(
current_values[i + 2] if current_values[i + 2] is not None else 0
)
# Calculate totals across all slots
total_requests = sum(request_counts)
total_rpm = sum(rpm_counts)
total_tpm = sum(tpm_counts)
should_raise_error = False
if max_parallel_requests is not None:
should_raise_error = total_requests > max_parallel_requests
if rpm_limit is not None:
should_raise_error = should_raise_error or total_rpm > rpm_limit
if tpm_limit is not None:
should_raise_error = should_raise_error or total_tpm > tpm_limit
if should_raise_error:
## DECREMENT CURRENT USAGE - so we don't keep failing subsequent requests
await self._increment_value_list_in_current_window(
increment_list=decrement_list,
ttl=60,
)
raise self.raise_rate_limit_error(
additional_details=f"{CommonProxyErrors.max_parallel_request_limit_reached.value}. Hit limit for {rate_limit_type}. Current usage: max_parallel_requests: {total_requests}, current_rpm: {total_rpm}, current_tpm: {total_tpm}. Current limits: max_parallel_requests: {max_parallel_requests}, rpm_limit: {rpm_limit}, tpm_limit: {tpm_limit}."
)
def time_to_next_minute(self) -> float:
# Get the current time
now = datetime.now()
# Calculate the next minute
next_minute = (now + timedelta(minutes=1)).replace(second=0, microsecond=0)
# Calculate the difference in seconds
seconds_to_next_minute = (next_minute - now).total_seconds()
return seconds_to_next_minute
def raise_rate_limit_error(
self, additional_details: Optional[str] = None
) -> HTTPException:
"""
Raise an HTTPException with a 429 status code and a retry-after header
"""
error_message = "Max parallel request limit reached"
if additional_details is not None:
error_message = error_message + " " + additional_details
raise HTTPException(
status_code=429,
detail=f"Max parallel request limit reached {additional_details}",
headers={"retry-after": str(self.time_to_next_minute())},
)
async def async_pre_call_hook( # noqa: PLR0915
self,
user_api_key_dict: UserAPIKeyAuth,
cache: DualCache,
data: dict,
call_type: str,
):
self.print_verbose("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
if max_parallel_requests is None:
max_parallel_requests = sys.maxsize
if data is None:
data = {}
global_max_parallel_requests = data.get("metadata", {}).get(
"global_max_parallel_requests", None
)
tpm_limit = getattr(user_api_key_dict, "tpm_limit", sys.maxsize)
if tpm_limit is None:
tpm_limit = sys.maxsize
rpm_limit = getattr(user_api_key_dict, "rpm_limit", sys.maxsize)
if rpm_limit is None:
rpm_limit = sys.maxsize
# ------------
# Setup values
# ------------
if global_max_parallel_requests is not None:
# get value from cache
_key = "global_max_parallel_requests"
current_global_requests = await self.internal_usage_cache.async_get_cache(
key=_key,
local_only=True,
litellm_parent_otel_span=user_api_key_dict.parent_otel_span,
)
# check if below limit
if current_global_requests is None:
current_global_requests = 1
# if above -> raise error
if current_global_requests >= global_max_parallel_requests:
return self.raise_rate_limit_error(
additional_details=f"Hit Global Limit: Limit={global_max_parallel_requests}, current: {current_global_requests}"
)
# if below -> increment
else:
await self.internal_usage_cache.async_increment_cache(
key=_key,
value=1,
local_only=True,
litellm_parent_otel_span=user_api_key_dict.parent_otel_span,
)
requested_model = data.get("model", None)
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}"
tasks = []
if api_key is not None:
# CHECK IF REQUEST ALLOWED for key
tasks.append(
self.check_key_in_limits_v2(
user_api_key_dict=user_api_key_dict,
data=data,
max_parallel_requests=max_parallel_requests,
precise_minute=precise_minute,
tpm_limit=tpm_limit,
rpm_limit=rpm_limit,
rate_limit_type="key",
)
)
if user_api_key_dict.user_id is not None:
# CHECK IF REQUEST ALLOWED for key
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.user_tpm_limit,
rpm_limit=user_api_key_dict.user_rpm_limit,
rate_limit_type="user",
)
)
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,
data=data,
max_parallel_requests=None,
precise_minute=precise_minute,
tpm_limit=user_api_key_dict.team_tpm_limit,
rpm_limit=user_api_key_dict.team_rpm_limit,
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
async def _update_usage_in_cache_post_call(
self,
user_api_key_dict: UserAPIKeyAuth,
precise_minute: str,
model: Optional[str],
total_tokens: int,
litellm_parent_otel_span: Union[Span, None] = None,
):
increment_list: List[Tuple[str, int]] = []
increment_value_by_group = {
"request_count": -1,
"tpm": total_tokens,
"rpm": 0,
}
rate_limit_types = ["key", "user", "customer", "team", "model_per_key"]
current_time = datetime.now()
current_hour = current_time.hour
current_minute = current_time.minute
current_slot = (
current_time.second // 15
) # This gives us 0-3 for the current 15s slot
slot_key = f"{current_time.strftime('%Y-%m-%d')}-{current_hour:02d}-{current_minute:02d}-{current_slot}"
for rate_limit_type in rate_limit_types:
for group in ["request_count", "rpm", "tpm"]:
key = self._get_current_usage_key(
user_api_key_dict=user_api_key_dict,
precise_minute=slot_key,
model=model,
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
await self._increment_value_list_in_current_window(
increment_list=increment_list,
ttl=60,
)
async def async_log_success_event( # noqa: PLR0915
self, kwargs, response_obj, start_time, end_time
):
from litellm.proxy.common_utils.callback_utils import (
get_model_group_from_litellm_kwargs,
)
litellm_parent_otel_span: Union[Span, None] = _get_parent_otel_span_from_kwargs(
kwargs=kwargs
)
try:
self.print_verbose("INSIDE parallel request limiter ASYNC SUCCESS LOGGING")
# ------------
# Setup values
# ------------
global_max_parallel_requests = kwargs["litellm_params"]["metadata"].get(
"global_max_parallel_requests", None
)
user_api_key = kwargs["litellm_params"]["metadata"]["user_api_key"]
user_api_key_user_id = kwargs["litellm_params"]["metadata"].get(
"user_api_key_user_id", None
)
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") or kwargs["litellm_params"][
"metadata"
].get("user_api_key_end_user_id", None)
# ------------
# Setup values
# ------------
if global_max_parallel_requests is not None:
# get value from cache
_key = "global_max_parallel_requests"
# decrement
await self.internal_usage_cache.async_increment_cache(
key=_key,
value=-1,
local_only=True,
litellm_parent_otel_span=litellm_parent_otel_span,
)
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}"
model_group = get_model_group_from_litellm_kwargs(kwargs)
total_tokens = 0
if isinstance(response_obj, ModelResponse):
total_tokens = response_obj.usage.total_tokens # type: ignore
# ------------
# Update usage - API Key
# ------------
await self._update_usage_in_cache_post_call(
user_api_key_dict=UserAPIKeyAuth(
api_key=user_api_key,
user_id=user_api_key_user_id,
team_id=user_api_key_team_id,
end_user_id=user_api_key_end_user_id,
),
precise_minute=precise_minute,
model=model_group,
total_tokens=total_tokens,
)
except Exception as e:
verbose_proxy_logger.exception(
"Inside Parallel Request Limiter: An exception occurred - {}".format(
str(e)
)
)
async def async_post_call_failure_hook(
self,
request_data: dict,
original_exception: Exception,
user_api_key_dict: UserAPIKeyAuth,
traceback_str: Optional[str] = None,
):
try:
self.print_verbose("Inside Max Parallel Request Failure Hook")
model_group = request_data.get("model", None)
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}"
## decrement call count if call failed
await self._update_usage_in_cache_post_call(
user_api_key_dict=user_api_key_dict,
precise_minute=precise_minute,
model=model_group,
total_tokens=0,
)
except Exception as e:
verbose_proxy_logger.exception(
"Inside Parallel Request Limiter: An exception occurred - {}".format(
str(e)
)
)
litellm/proxy/hooks/parallel_request_limiter_v3.py
+741
View File
@@ -0,0 +1,741 @@
"""
This is a rate limiter implementation based on a similar one by Envoy proxy.
This is currently in development and not yet ready for production.
"""
import os
import sys
from datetime import datetime
from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, TypedDict, Union
from fastapi import HTTPException
import litellm
from litellm import DualCache
from litellm._logging import verbose_proxy_logger
from litellm.integrations.custom_logger import CustomLogger
from litellm.proxy._types import UserAPIKeyAuth
if TYPE_CHECKING:
from opentelemetry.trace import Span as _Span
from litellm.proxy.utils import InternalUsageCache as _InternalUsageCache
from litellm.types.caching import RedisPipelineIncrementOperation
Span = Union[_Span, Any]
InternalUsageCache = _InternalUsageCache
else:
Span = Any
InternalUsageCache = Any
RATE_LIMITER_SCRIPT = """
local window_key = KEYS[1]
local counter_key = KEYS[2]
local now = ARGV[1]
local window_size = ARGV[2]
-- Check if window exists and is valid
local window_start = redis.call('GET', window_key)
if not window_start or (tonumber(now) - tonumber(window_start)) >= tonumber(window_size) then
-- Reset window and counter
redis.call('SET', window_key, now)
redis.call('SET', counter_key, 1)
redis.call('EXPIRE', window_key, window_size)
redis.call('EXPIRE', counter_key, window_size)
return {1, now}
end
-- Increment counter
local counter = redis.call('INCR', counter_key)
return {counter, window_start}
"""
BATCH_RATE_LIMITER_SCRIPT = """
local results = {}
local now = tonumber(ARGV[1])
local window_size = tonumber(ARGV[2])
-- Process each window/counter pair
for i = 1, #KEYS, 2 do
local window_key = KEYS[i]
local counter_key = KEYS[i + 1]
local increment_value = tonumber(KEYS[i + 2]) or 1
-- Check if window exists and is valid
local window_start = redis.call('GET', window_key)
if not window_start or (now - tonumber(window_start)) >= window_size then
-- Reset window and counter
redis.call('SET', window_key, tostring(now))
redis.call('SET', counter_key, increment_value)
redis.call('EXPIRE', window_key, window_size)
redis.call('EXPIRE', counter_key, window_size)
table.insert(results, tostring(now)) -- window_start
table.insert(results, increment_value) -- counter
else
local counter = redis.call('INCR', counter_key)
table.insert(results, window_start) -- window_start
table.insert(results, counter) -- counter
end
end
return results
"""
class RateLimitDescriptorRateLimitObject(TypedDict, total=False):
requests_per_unit: Optional[int]
tokens_per_unit: Optional[int]
max_parallel_requests: Optional[int]
window_size: Optional[int]
class RateLimitDescriptor(TypedDict):
key: str
value: str
rate_limit: Optional[RateLimitDescriptorRateLimitObject]
class RateLimitResponse(TypedDict):
overall_code: str
statuses: List[Dict[str, Any]]
class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
def __init__(self, internal_usage_cache: InternalUsageCache):
self.internal_usage_cache = internal_usage_cache
if self.internal_usage_cache.dual_cache.redis_cache is not None:
self.rate_limiter_script = (
self.internal_usage_cache.dual_cache.redis_cache.async_register_script(
RATE_LIMITER_SCRIPT
)
)
self.batch_rate_limiter_script = (
self.internal_usage_cache.dual_cache.redis_cache.async_register_script(
BATCH_RATE_LIMITER_SCRIPT
)
)
else:
self.rate_limiter_script = None
self.batch_rate_limiter_script = None
self.window_size = int(os.getenv("LITELLM_RATE_LIMIT_WINDOW_SIZE", 60))
def print_verbose(self, print_statement):
try:
verbose_proxy_logger.debug(print_statement)
if litellm.set_verbose:
print(print_statement) # noqa
except Exception:
pass
async def rate_limiter_script_handler(
self,
window_key: str,
counter_key: str,
now: float,
window_size: float,
parent_otel_span: Optional[Span] = None,
) -> int:
"""
Update Redis
Update in-memory cache
Return the new count and window value
"""
if self.rate_limiter_script is not None:
result = await self.rate_limiter_script(
keys=[window_key, counter_key], args=[now, window_size]
)
counter_value, window_value = result
# Update in-memory cache
await self.internal_usage_cache.async_set_cache(
key=counter_key,
value=counter_value,
ttl=window_size,
litellm_parent_otel_span=parent_otel_span,
local_only=True,
)
await self.internal_usage_cache.async_set_cache(
key=window_key,
value=window_value,
ttl=window_size,
litellm_parent_otel_span=parent_otel_span,
local_only=True,
)
return counter_value
else: # in-memory only implementation
current_window = await self.internal_usage_cache.async_get_cache(
key=window_key,
litellm_parent_otel_span=parent_otel_span,
)
if current_window is None or (now - current_window) >= window_size:
# Set new window start time
await self.internal_usage_cache.async_set_cache(
key=window_key,
value=now,
ttl=window_size,
litellm_parent_otel_span=parent_otel_span,
)
# Reset counter
result = await self.internal_usage_cache.async_increment_cache(
key=counter_key,
value=1,
ttl=window_size,
litellm_parent_otel_span=parent_otel_span,
)
else:
# Get current count
result = (
await self.internal_usage_cache.async_increment_cache(
key=counter_key,
value=1,
ttl=window_size,
litellm_parent_otel_span=parent_otel_span,
)
or 1
)
return int(result)
def create_rate_limit_keys(
self,
key: str,
value: str,
rate_limit_type: Literal["requests", "tokens", "max_parallel_requests"],
) -> str:
"""
Create the rate limit keys for the given key and value.
"""
counter_key = f"{{{key}:{value}}}:{rate_limit_type}"
return counter_key
def is_cache_list_over_limit(
self,
keys_to_fetch: List[str],
cache_values: List[Any],
key_metadata: Dict[str, Any],
) -> RateLimitResponse:
"""
Check if the cache values are over the limit.
"""
statuses = []
overall_code = "OK"
for i in range(0, len(cache_values), 2):
item_code = "OK"
window_key = keys_to_fetch[i]
counter_key = keys_to_fetch[i + 1]
counter_value = cache_values[i + 1]
requests_limit = key_metadata[window_key]["requests_limit"]
max_parallel_requests_limit = key_metadata[window_key][
"max_parallel_requests_limit"
]
tokens_limit = key_metadata[window_key]["tokens_limit"]
# Determine which limit to use for current_limit and limit_remaining
if counter_key.endswith(":requests"):
current_limit = requests_limit
elif counter_key.endswith(":max_parallel_requests"):
current_limit = max_parallel_requests_limit
elif counter_key.endswith(":tokens"):
current_limit = tokens_limit
else:
current_limit = None
if (
counter_key.endswith(":requests")
and requests_limit is not None
and counter_value is not None
and int(counter_value) + 1 > requests_limit
):
overall_code = "OVER_LIMIT"
item_code = "OVER_LIMIT"
elif (
counter_key.endswith(":max_parallel_requests")
and max_parallel_requests_limit is not None
and counter_value is not None
and int(counter_value) + 1 > max_parallel_requests_limit
):
overall_code = "OVER_LIMIT"
item_code = "OVER_LIMIT"
elif (
counter_key.endswith(":tokens")
and tokens_limit is not None
and counter_value is not None
and int(counter_value) + 1 > tokens_limit
):
overall_code = "OVER_LIMIT"
item_code = "OVER_LIMIT"
# Only compute limit_remaining if current_limit is not None
if current_limit is None:
limit_remaining = None
elif counter_value is None:
limit_remaining = current_limit
else:
limit_remaining = current_limit - counter_value
statuses.append(
{
"code": item_code,
"current_limit": current_limit,
"limit_remaining": limit_remaining,
}
)
return RateLimitResponse(overall_code=overall_code, statuses=statuses)
async def should_rate_limit(
self,
descriptors: List[RateLimitDescriptor],
parent_otel_span: Optional[Span] = None,
read_only: bool = False,
) -> RateLimitResponse:
"""
Check if any of the rate limit descriptors should be rate limited.
Returns a RateLimitResponse with the overall code and status for each descriptor.
Uses batch operations for Redis to improve performance.
"""
now = datetime.now().timestamp()
now_int = int(now) # Convert to integer for Redis Lua script
# Collect all keys and their metadata upfront
keys_to_fetch: List[str] = []
key_metadata = {} # Store metadata for each key
for descriptor in descriptors:
key = descriptor["key"]
value = descriptor["value"]
rate_limit = descriptor.get("rate_limit", {}) or {}
requests_limit = rate_limit.get("requests_per_unit")
tokens_limit = rate_limit.get("tokens_per_unit")
max_parallel_requests_limit = rate_limit.get("max_parallel_requests")
window_size = rate_limit.get("window_size") or self.window_size
window_key = f"{{{key}:{value}}}:window"
if requests_limit is not None:
key = self.create_rate_limit_keys(key, value, "requests")
keys_to_fetch.extend([window_key, key])
elif tokens_limit is not None:
key = self.create_rate_limit_keys(key, value, "tokens")
keys_to_fetch.extend([window_key, key])
elif max_parallel_requests_limit is not None:
key = self.create_rate_limit_keys(key, value, "max_parallel_requests")
keys_to_fetch.extend([window_key, key])
else:
continue
key_metadata[window_key] = {
"requests_limit": int(requests_limit)
if requests_limit is not None
else None,
"tokens_limit": int(tokens_limit) if tokens_limit is not None else None,
"max_parallel_requests_limit": int(max_parallel_requests_limit)
if max_parallel_requests_limit is not None
else None,
"window_size": int(window_size),
}
## CHECK IN-MEMORY CACHE
cache_values = await self.internal_usage_cache.async_batch_get_cache(
keys=keys_to_fetch,
parent_otel_span=parent_otel_span,
local_only=True,
)
if cache_values is not None:
rate_limit_response = self.is_cache_list_over_limit(
keys_to_fetch, cache_values, key_metadata
)
if rate_limit_response["overall_code"] == "OVER_LIMIT":
return rate_limit_response
## IF under limit, check Redis
if self.batch_rate_limiter_script is not None:
cache_values = await self.batch_rate_limiter_script(
keys=keys_to_fetch,
args=[now_int, self.window_size], # Use integer timestamp
)
# update in-memory cache with new values
for i in range(0, len(cache_values), 2):
window_key = keys_to_fetch[i]
counter_key = keys_to_fetch[i + 1]
window_value = cache_values[i]
counter_value = cache_values[i + 1]
await self.internal_usage_cache.async_set_cache(
key=counter_key,
value=counter_value,
ttl=self.window_size,
litellm_parent_otel_span=parent_otel_span,
local_only=True,
)
await self.internal_usage_cache.async_set_cache(
key=window_key,
value=window_value,
ttl=self.window_size,
litellm_parent_otel_span=parent_otel_span,
local_only=True,
)
else:
raise ValueError("Batch rate limiter script is not initialized")
rate_limit_response = self.is_cache_list_over_limit(
keys_to_fetch, cache_values, key_metadata
)
return rate_limit_response
async def async_pre_call_hook(
self,
user_api_key_dict: UserAPIKeyAuth,
cache: DualCache,
data: dict,
call_type: str,
):
"""
Pre-call hook to check rate limits before making the API call.
"""
from litellm.proxy.auth.auth_utils import (
get_key_model_rpm_limit,
get_key_model_tpm_limit,
)
verbose_proxy_logger.debug("Inside Rate Limit Pre-Call Hook")
# Create rate limit descriptors
descriptors = []
# API Key rate limits
if user_api_key_dict.api_key:
descriptors.append(
RateLimitDescriptor(
key="api_key",
value=user_api_key_dict.api_key,
rate_limit={
"requests_per_unit": user_api_key_dict.rpm_limit,
"tokens_per_unit": user_api_key_dict.tpm_limit,
"max_parallel_requests": user_api_key_dict.max_parallel_requests,
"window_size": self.window_size, # 1 minute window
},
)
)
# User rate limits
if user_api_key_dict.user_id:
descriptors.append(
RateLimitDescriptor(
key="user",
value=user_api_key_dict.user_id,
rate_limit={
"requests_per_unit": user_api_key_dict.user_rpm_limit,
"tokens_per_unit": user_api_key_dict.user_tpm_limit,
"window_size": self.window_size,
},
)
)
# Team rate limits
if user_api_key_dict.team_id:
descriptors.append(
RateLimitDescriptor(
key="team",
value=user_api_key_dict.team_id,
rate_limit={
"requests_per_unit": user_api_key_dict.team_rpm_limit,
"tokens_per_unit": user_api_key_dict.team_tpm_limit,
"window_size": self.window_size,
},
)
)
# End user rate limits
if user_api_key_dict.end_user_id:
descriptors.append(
RateLimitDescriptor(
key="end_user",
value=user_api_key_dict.end_user_id,
rate_limit={
"requests_per_unit": user_api_key_dict.end_user_rpm_limit,
"tokens_per_unit": user_api_key_dict.end_user_tpm_limit,
"window_size": self.window_size,
},
)
)
# Model rate limits
requested_model = data.get("model", None)
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]
descriptors.append(
RateLimitDescriptor(
key="model_per_key",
value=f"{user_api_key_dict.api_key}:{requested_model}",
rate_limit={
"requests_per_unit": model_specific_rpm_limit,
"tokens_per_unit": model_specific_tpm_limit,
"window_size": self.window_size,
},
)
)
# Check rate limits
response = await self.should_rate_limit(
descriptors=descriptors,
parent_otel_span=user_api_key_dict.parent_otel_span,
)
if response["overall_code"] == "OVER_LIMIT":
# Find which descriptor hit the limit
for i, status in enumerate(response["statuses"]):
if status["code"] == "OVER_LIMIT":
descriptor = descriptors[i]
raise HTTPException(
status_code=429,
detail=f"Rate limit exceeded for {descriptor['key']}: {descriptor['value']}. Remaining: {status['limit_remaining']}",
headers={
"retry-after": str(self.window_size)
}, # Retry after 1 minute
)
def _create_pipeline_operations(
self,
key: str,
value: str,
rate_limit_type: Literal["requests", "tokens", "max_parallel_requests"],
total_tokens: int,
) -> List["RedisPipelineIncrementOperation"]:
"""
Create pipeline operations for TPM increments
"""
from litellm.types.caching import RedisPipelineIncrementOperation
pipeline_operations: List[RedisPipelineIncrementOperation] = []
counter_key = self.create_rate_limit_keys(
key=key,
value=value,
rate_limit_type="tokens",
)
pipeline_operations.append(
RedisPipelineIncrementOperation(
key=counter_key,
increment_value=total_tokens,
ttl=self.window_size,
)
)
return pipeline_operations
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
"""
Update TPM usage on successful API calls by incrementing counters using pipeline
"""
from litellm.litellm_core_utils.core_helpers import (
_get_parent_otel_span_from_kwargs,
)
from litellm.proxy.common_utils.callback_utils import (
get_model_group_from_litellm_kwargs,
)
from litellm.types.caching import RedisPipelineIncrementOperation
from litellm.types.utils import ModelResponse, Usage
litellm_parent_otel_span: Union[Span, None] = _get_parent_otel_span_from_kwargs(
kwargs
)
try:
self.print_verbose("INSIDE parallel request limiter ASYNC SUCCESS LOGGING")
# Get metadata from kwargs
user_api_key = kwargs["litellm_params"]["metadata"]["user_api_key"]
user_api_key_user_id = kwargs["litellm_params"]["metadata"].get(
"user_api_key_user_id", None
)
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") or kwargs["litellm_params"][
"metadata"
].get("user_api_key_end_user_id", None)
model_group = get_model_group_from_litellm_kwargs(kwargs)
# Get total tokens from response
total_tokens = 0
if isinstance(response_obj, ModelResponse):
_usage = getattr(response_obj, "usage", None)
if _usage and isinstance(_usage, Usage):
total_tokens = _usage.total_tokens
# Create pipeline operations for TPM increments
pipeline_operations: List[RedisPipelineIncrementOperation] = []
# API Key TPM
if user_api_key:
# MAX PARALLEL REQUESTS - only support for API Key, just decrement the counter
counter_key = self.create_rate_limit_keys(
key="api_key",
value=user_api_key,
rate_limit_type="max_parallel_requests",
)
pipeline_operations.append(
RedisPipelineIncrementOperation(
key=counter_key,
increment_value=-1,
ttl=self.window_size,
)
)
pipeline_operations.extend(
self._create_pipeline_operations(
key="api_key",
value=user_api_key,
rate_limit_type="tokens",
total_tokens=total_tokens,
)
)
# User TPM
if user_api_key_user_id:
# TPM
pipeline_operations.extend(
self._create_pipeline_operations(
key="user",
value=user_api_key_user_id,
rate_limit_type="tokens",
total_tokens=total_tokens,
)
)
# Team TPM
if user_api_key_team_id:
pipeline_operations.extend(
self._create_pipeline_operations(
key="team",
value=user_api_key_team_id,
rate_limit_type="tokens",
total_tokens=total_tokens,
)
)
# End User TPM
if user_api_key_end_user_id:
pipeline_operations.extend(
self._create_pipeline_operations(
key="end_user",
value=user_api_key_end_user_id,
rate_limit_type="tokens",
total_tokens=total_tokens,
)
)
# Model-specific TPM
if model_group and user_api_key:
pipeline_operations.extend(
self._create_pipeline_operations(
key="model_per_key",
value=f"{user_api_key}:{model_group}",
rate_limit_type="tokens",
total_tokens=total_tokens,
)
)
# Execute all increments in a single pipeline
if pipeline_operations:
await self.internal_usage_cache.dual_cache.async_increment_cache_pipeline(
increment_list=pipeline_operations,
litellm_parent_otel_span=litellm_parent_otel_span,
)
except Exception as e:
verbose_proxy_logger.exception(
f"Error in rate limit success event: {str(e)}"
)
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
"""
No-op for failure event since we handle increments in should_rate_limit
"""
pass
async def async_post_call_success_hook(
self, data: dict, user_api_key_dict: UserAPIKeyAuth, response
):
"""
Post-call hook to update rate limit headers in the response.
"""
try:
descriptors = []
# API Key
if user_api_key_dict.api_key:
descriptors.append(
RateLimitDescriptor(
key="api_key",
value=user_api_key_dict.api_key,
rate_limit={
"requests_per_unit": user_api_key_dict.rpm_limit
or sys.maxsize,
"window_size": self.window_size,
},
)
)
# Check rate limits
# rate_limit_response = await self.should_rate_limit(
# descriptors=descriptors,
# parent_otel_span=user_api_key_dict.parent_otel_span,
# read_only=True,
# )
# # Update response headers
# if hasattr(response, "_hidden_params"):
# _hidden_params = getattr(response, "_hidden_params")
# else:
# _hidden_params = None
# if _hidden_params is not None and (
# isinstance(_hidden_params, BaseModel)
# or isinstance(_hidden_params, dict)
# ):
# if isinstance(_hidden_params, BaseModel):
# _hidden_params = _hidden_params.model_dump()
# _additional_headers = _hidden_params.get("additional_headers", {}) or {}
# # Add rate limit headers
# for i, status in enumerate(rate_limit_response["statuses"]):
# descriptor = descriptors[i]
# prefix = f"x-ratelimit-{descriptor['key']}"
# _additional_headers[f"{prefix}-remaining-requests"] = status[
# "limit_remaining"
# ]
# _additional_headers[f"{prefix}-limit-requests"] = status[
# "current_limit"
# ]
# setattr(
# response,
# "_hidden_params",
# {**_hidden_params, "additional_headers": _additional_headers},
# )
except Exception as e:
verbose_proxy_logger.exception(
f"Error in rate limit post-call hook: {str(e)}"
)
litellm/router_strategy/base_routing_strategy.py
+30 -13
View File
@@ -140,18 +140,31 @@ class BaseRoutingStrategy(ABC):
compressed_ops[op["key"]] = op
ops_to_remove.append(idx)
# Convert back to list
compressed_queue = list(compressed_ops.values())
await self.dual_cache.redis_cache.async_increment_pipeline(
increment_list=compressed_queue,
increment_result = (
await self.dual_cache.redis_cache.async_increment_pipeline(
increment_list=compressed_queue,
)
)
self.redis_increment_operation_queue = [
op
for idx, op in enumerate(self.redis_increment_operation_queue)
if idx not in ops_to_remove
]
if increment_result is not None:
return_result = {
key["key"]: op
for key, op in zip(compressed_queue, increment_result)
}
else:
return_result = {}
return return_result
except Exception as e:
verbose_router_logger.error(
f"Error syncing in-memory cache with Redis: {str(e)}"
@@ -202,10 +215,7 @@ class BaseRoutingStrategy(ABC):
self.get_in_memory_keys_to_update()
) # if no pattern OR redis cache does not support scan_iter, use in-memory keys
if isinstance(cache_keys, set):
cache_keys_list = list(cache_keys)
else:
cache_keys_list = cache_keys
cache_keys_list = list(cache_keys)
# 1. Snapshot in-memory before
in_memory_before_dict = {}
@@ -218,12 +228,9 @@ class BaseRoutingStrategy(ABC):
in_memory_before_dict[k] = float(v or 0)
# 1. Push all provider spend increments to Redis
await self._push_in_memory_increments_to_redis()
# 2. Fetch from Redis
redis_values = await self.dual_cache.redis_cache.async_batch_get_cache(
key_list=cache_keys_list
)
redis_values = await self._push_in_memory_increments_to_redis()
if redis_values is None:
return
# 4. Merge
for key in cache_keys_list:
@@ -233,7 +240,17 @@ class BaseRoutingStrategy(ABC):
await self.dual_cache.in_memory_cache.async_get_cache(key=key) or 0
)
delta = after - before
merged = redis_val + delta
if after <= redis_val:
merged = redis_val + delta
else:
continue
# elif "rpm" in key: # redis is behind in-memory cache
# # shut down the proxy
# print(f"self.redis_increment_operation_queue: {self.redis_increment_operation_queue}")
# print(f"Redis_val={redis_val} is behind in-memory cache_val={after} for key: {key}. This should not happen, since we should be updating redis with in-memory cache.")
# import os
# os._exit(1)
# raise Exception(f"Redis is behind in-memory cache for key: {key}. This should not happen, since we should be updating redis with in-memory cache.")
await self.dual_cache.in_memory_cache.async_set_cache(
key=key, value=merged
)
litellm/types/caching.py
+10
View File
@@ -37,6 +37,16 @@ class RedisPipelineIncrementOperation(TypedDict):
ttl: Optional[int]
class RedisPipelineSetOperation(TypedDict):
"""
TypeDict for 1 Redis Pipeline Set Operation
"""
key: str
value: Any
ttl: Optional[int]
DynamicCacheControl = TypedDict(
"DynamicCacheControl",
{
tests/local_testing/test_caching.py
+1 -1
View File
@@ -2524,7 +2524,7 @@ async def test_redis_increment_pipeline():
results = await redis_cache.async_increment_pipeline(increment_list)
# Verify results
assert len(results) == 8 # 4 increment operations + 4 expire operations
assert len(results) == 4
# Verify the values were actually set in Redis
value1 = await redis_cache.async_get_cache("test_key1")
tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v2.py
-626
View File
@@ -1,626 +0,0 @@
"""
Unit Tests for the max parallel request limiter v2 for the proxy
"""
import asyncio
import os
import sys
from datetime import datetime
import pytest
from fastapi import HTTPException
import litellm
from litellm import Router
from litellm.caching.caching import DualCache
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.hooks.parallel_request_limiter_v2 import (
_PROXY_MaxParallelRequestsHandler_v2 as _PROXY_MaxParallelRequestsHandler,
)
from litellm.proxy.utils import InternalUsageCache, ProxyLogging, hash_token
@pytest.mark.flaky(reruns=3)
@pytest.mark.asyncio
async def test_normal_router_call_v2(monkeypatch):
"""
Test normal router call with parallel request limiter v2
"""
model_list = [
{
"model_name": "azure-model",
"litellm_params": {
"model": "azure/gpt-turbo",
"api_key": "os.environ/AZURE_FRANCE_API_KEY",
"api_base": "https://openai-france-1234.openai.azure.com",
"rpm": 1440,
},
"model_info": {"id": 1},
},
{
"model_name": "azure-model",
"litellm_params": {
"model": "azure/gpt-35-turbo",
"api_key": "os.environ/AZURE_EUROPE_API_KEY",
"api_base": "https://my-endpoint-europe-berri-992.openai.azure.com",
"rpm": 6,
},
"model_info": {"id": 2},
},
]
router = Router(
model_list=model_list,
set_verbose=False,
num_retries=3,
) # type: ignore
_api_key = "sk-12345"
_api_key = hash_token(_api_key)
user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, max_parallel_requests=1)
local_cache = DualCache()
parallel_request_handler = _PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(local_cache)
)
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=""
)
current_time = datetime.now()
current_hour = current_time.hour
current_minute = current_time.minute
current_slot = (
current_time.second // 15
) # This gives us 0-3 for the current 15s slot
slot_key = f"{current_time.strftime('%Y-%m-%d')}-{current_hour:02d}-{current_minute:02d}-{current_slot}"
print(f"slot_key: {slot_key}")
request_count_api_key = parallel_request_handler._get_current_usage_key(
user_api_key_dict=user_api_key_dict,
precise_minute=slot_key,
model=None,
rate_limit_type="key",
group="request_count",
)
await asyncio.sleep(1)
assert (
parallel_request_handler.internal_usage_cache.get_cache(
key=request_count_api_key
)
== 1
)
# normal call
response = await router.acompletion(
model="azure-model",
messages=[{"role": "user", "content": "Hey, how's it going?"}],
metadata={"user_api_key": _api_key},
mock_response="hello",
)
await asyncio.sleep(1) # success is done in a separate thread
print(f"local_cache in normal call: {local_cache}")
assert (
parallel_request_handler.internal_usage_cache.get_cache(
key=request_count_api_key
)
== 0
)
@pytest.mark.parametrize(
"rate_limit_object",
[
"key",
"model_per_key",
"user",
"customer",
"team",
],
)
@pytest.mark.flaky(reruns=3)
@pytest.mark.asyncio
async def test_normal_router_call_tpm(monkeypatch, rate_limit_object):
"""
Test normal router call with parallel request limiter v2
"""
model_list = [
{
"model_name": "azure-model",
"litellm_params": {
"model": "azure/gpt-turbo",
"api_key": "os.environ/AZURE_FRANCE_API_KEY",
"api_base": "https://openai-france-1234.openai.azure.com",
"rpm": 1440,
},
"model_info": {"id": 1},
},
{
"model_name": "azure-model",
"litellm_params": {
"model": "azure/gpt-35-turbo",
"api_key": "os.environ/AZURE_EUROPE_API_KEY",
"api_base": "https://my-endpoint-europe-berri-992.openai.azure.com",
"rpm": 6,
},
"model_info": {"id": 2},
},
]
router = Router(
model_list=model_list,
set_verbose=False,
num_retries=3,
) # type: ignore
_api_key = "sk-12345"
_api_key = hash_token(_api_key)
if rate_limit_object == "key":
user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, tpm_limit=10)
elif rate_limit_object == "user":
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)
)
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={"model": "azure-model"},
call_type="",
)
current_time = datetime.now()
current_hour = current_time.hour
current_minute = current_time.minute
current_slot = (
current_time.second // 15
) # This gives us 0-3 for the current 15s slot
slot_key = f"{current_time.strftime('%Y-%m-%d')}-{current_hour:02d}-{current_minute:02d}-{current_slot}"
print(f"slot_key: {slot_key}")
request_count_api_key = parallel_request_handler._get_current_usage_key(
user_api_key_dict=user_api_key_dict,
precise_minute=slot_key,
model="azure-model",
rate_limit_type=rate_limit_object,
group="tpm",
)
print(f"request_count_api_key: {request_count_api_key}")
await asyncio.sleep(1)
assert (
parallel_request_handler.internal_usage_cache.get_cache(
key=request_count_api_key
)
== 0
)
# normal call
response = await router.acompletion(
model="azure-model",
messages=[{"role": "user", "content": "Hey, how's it going?"}],
metadata={
"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",
)
await asyncio.sleep(1) # success is done in a separate thread
print(f"request_count_api_key: {request_count_api_key}")
next_slot_key = f"{current_time.strftime('%Y-%m-%d')}-{current_hour:02d}-{current_minute:02d}-{current_slot + 1 if current_slot < 3 else 0}"
request_count_api_key_next_slot = parallel_request_handler._get_current_usage_key(
user_api_key_dict=user_api_key_dict,
precise_minute=next_slot_key,
model="azure-model",
rate_limit_type=rate_limit_object,
group="tpm",
)
## check if current slot matches response.usage.total_tokens else next slot
current_slot_get_cache = parallel_request_handler.internal_usage_cache.get_cache(
key=request_count_api_key
)
next_slot_get_cache = parallel_request_handler.internal_usage_cache.get_cache(
key=request_count_api_key_next_slot
)
assert (
current_slot_get_cache == response.usage.total_tokens
or next_slot_get_cache == response.usage.total_tokens
)
@pytest.mark.parametrize(
"rate_limit_object",
[
"key",
"model_per_key",
"user",
"customer",
"team",
],
)
@pytest.mark.flaky(reruns=3)
@pytest.mark.asyncio
async def test_normal_router_call_rpm(monkeypatch, rate_limit_object):
"""
Test normal router call with parallel request limiter v2
"""
model_list = [
{
"model_name": "azure-model",
"litellm_params": {
"model": "azure/gpt-turbo",
"api_key": "os.environ/AZURE_FRANCE_API_KEY",
"api_base": "https://openai-france-1234.openai.azure.com",
"rpm": 1440,
},
"model_info": {"id": 1},
},
{
"model_name": "azure-model",
"litellm_params": {
"model": "azure/gpt-35-turbo",
"api_key": "os.environ/AZURE_EUROPE_API_KEY",
"api_base": "https://my-endpoint-europe-berri-992.openai.azure.com",
"rpm": 6,
},
"model_info": {"id": 2},
},
]
router = Router(
model_list=model_list,
set_verbose=False,
num_retries=3,
) # type: ignore
_api_key = "sk-12345"
_api_key = hash_token(_api_key)
if rate_limit_object == "key":
user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, rpm_limit=1)
elif rate_limit_object == "user":
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)
)
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={"model": "azure-model"},
call_type="",
)
current_time = datetime.now()
current_hour = current_time.hour
current_minute = current_time.minute
current_slot = (
current_time.second // 15
) # This gives us 0-3 for the current 15s slot
slot_key = f"{current_time.strftime('%Y-%m-%d')}-{current_hour:02d}-{current_minute:02d}-{current_slot}"
request_count_api_key = parallel_request_handler._get_current_usage_key(
user_api_key_dict=user_api_key_dict,
precise_minute=slot_key,
model="azure-model",
rate_limit_type=rate_limit_object,
group="rpm",
)
await asyncio.sleep(1)
assert (
parallel_request_handler.internal_usage_cache.get_cache(
key=request_count_api_key
)
== 1
)
# normal call
response = await router.acompletion(
model="azure-model",
messages=[{"role": "user", "content": "Hey, how's it going?"}],
metadata={
"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",
)
await asyncio.sleep(1) # success is done in a separate thread
print(f"request_count_api_key: {request_count_api_key}")
assert (
parallel_request_handler.internal_usage_cache.get_cache(
key=request_count_api_key
)
== 1
)
with pytest.raises(HTTPException):
await parallel_request_handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=local_cache,
data={"model": "azure-model"},
call_type="",
)
@pytest.mark.flaky(reruns=3)
@pytest.mark.asyncio
async def test_streaming_router_call_v2(monkeypatch):
"""
Test streaming router call with parallel request limiter v2
"""
model_list = [
{
"model_name": "azure-model",
"litellm_params": {
"model": "azure/gpt-turbo",
"api_key": "os.environ/AZURE_FRANCE_API_KEY",
"api_base": "https://openai-france-1234.openai.azure.com",
"rpm": 1440,
},
"model_info": {"id": 1},
},
{
"model_name": "azure-model",
"litellm_params": {
"model": "azure/gpt-35-turbo",
"api_key": "os.environ/AZURE_EUROPE_API_KEY",
"api_base": "https://my-endpoint-europe-berri-992.openai.azure.com",
"rpm": 6,
},
"model_info": {"id": 2},
},
]
router = Router(
model_list=model_list,
set_verbose=False,
num_retries=3,
) # type: ignore
_api_key = "sk-12345"
_api_key = hash_token(_api_key)
user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, max_parallel_requests=1)
local_cache = DualCache()
print(f"litellm callbacks pre-set: {litellm.callbacks}")
parallel_request_handler = _PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(local_cache)
)
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=""
)
current_time = datetime.now()
current_hour = current_time.hour
current_minute = current_time.minute
current_slot = (
current_time.second // 15
) # This gives us 0-3 for the current 15s slot
slot_key = f"{current_time.strftime('%Y-%m-%d')}-{current_hour:02d}-{current_minute:02d}-{current_slot}"
request_count_api_key = parallel_request_handler._get_current_usage_key(
user_api_key_dict=user_api_key_dict,
precise_minute=slot_key,
model=None,
rate_limit_type="key",
group="request_count",
)
await asyncio.sleep(1)
assert (
parallel_request_handler.internal_usage_cache.get_cache(
key=request_count_api_key
)
== 1
)
# streaming call
print(f"litellm callbacks: {litellm.callbacks}")
response = await router.acompletion(
model="azure-model",
messages=[{"role": "user", "content": "Hey, how's it going?"}],
stream=True,
metadata={"user_api_key": _api_key},
mock_response="hello",
)
async for chunk in response:
continue
await asyncio.sleep(3) # success is done in a separate thread
print(f"local_cache in streaming call: {local_cache}")
assert (
parallel_request_handler.internal_usage_cache.get_cache(
key=request_count_api_key
)
== 0
)
@pytest.mark.parametrize(
"rate_limit_object",
[
"key",
# "model_per_key",
"user",
# "customer",
"team",
],
)
@pytest.mark.flaky(reruns=3)
@pytest.mark.asyncio
async def test_bad_router_call_v2(monkeypatch, rate_limit_object):
"""
Test bad router call with parallel request limiter v2
"""
model_list = [
{
"model_name": "azure-model",
"litellm_params": {
"model": "azure/gpt-turbo",
"api_key": "os.environ/AZURE_FRANCE_API_KEY",
"api_base": "https://openai-france-1234.openai.azure.com",
"rpm": 1440,
},
"model_info": {"id": 1},
},
{
"model_name": "azure-model",
"litellm_params": {
"model": "azure/gpt-35-turbo",
"api_key": "os.environ/AZURE_EUROPE_API_KEY",
"api_base": "https://my-endpoint-europe-berri-992.openai.azure.com",
"rpm": 6,
},
"model_info": {"id": 2},
},
]
router = Router(
model_list=model_list,
set_verbose=False,
num_retries=3,
) # type: ignore
_api_key = "sk-12345"
_api_key = hash_token(_api_key)
if rate_limit_object == "key":
user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, rpm_limit=1)
elif rate_limit_object == "user":
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)
local_cache = DualCache()
parallel_request_handler = _PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(local_cache)
)
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=""
)
current_time = datetime.now()
current_hour = current_time.hour
current_minute = current_time.minute
current_slot = (
current_time.second // 15
) # This gives us 0-3 for the current 15s slot
slot_key = f"{current_time.strftime('%Y-%m-%d')}-{current_hour:02d}-{current_minute:02d}-{current_slot}"
request_count_api_key = parallel_request_handler._get_current_usage_key(
user_api_key_dict=user_api_key_dict,
precise_minute=slot_key,
model=None,
rate_limit_type=rate_limit_object,
group="rpm",
)
await asyncio.sleep(1)
assert (
parallel_request_handler.internal_usage_cache.get_cache(
key=request_count_api_key
)
== 1
)
# bad streaming call
await parallel_request_handler.async_post_call_failure_hook(
request_data={},
original_exception=Exception("test"),
user_api_key_dict=user_api_key_dict,
)
assert (
parallel_request_handler.internal_usage_cache.get_cache(
key=request_count_api_key
)
== 1
)
@pytest.mark.asyncio
async def test_check_key_in_limits_v2_sliding_window():
"""
Test the check_key_in_limits_v2 function with sliding window logic
"""
print("Starting test")
_api_key = "sk-12345"
_api_key = hash_token(_api_key)
user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, rpm_limit=2)
local_cache = DualCache()
parallel_request_handler = _PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(local_cache)
)
print("Created handler")
# Get current time and calculate slots
current_time = datetime.now()
current_slot = (current_time.minute * 60 + current_time.second) // 15
current_slot_key = (
f"{current_time.strftime('%Y-%m-%d')}-{current_time.hour:02d}-{current_slot}"
)
print(f"Current slot key: {current_slot_key}")
print("Making first request")
# Test 1: First request should succeed
await parallel_request_handler.check_key_in_limits_v2(
user_api_key_dict=user_api_key_dict,
data={},
max_parallel_requests=None,
precise_minute=current_slot_key,
tpm_limit=None,
rpm_limit=3,
rate_limit_type="key",
)
print("First request completed")
print("Making second request")
# Test 2: Second request should succeed
await parallel_request_handler.check_key_in_limits_v2(
user_api_key_dict=user_api_key_dict,
data={},
max_parallel_requests=None,
precise_minute=current_slot_key,
tpm_limit=None,
rpm_limit=3,
rate_limit_type="key",
)
print("Second request completed")
print("Verifying cache")
# Make third request - should fail
with pytest.raises(HTTPException):
await parallel_request_handler.check_key_in_limits_v2(
user_api_key_dict=user_api_key_dict,
data={},
max_parallel_requests=None,
precise_minute=current_slot_key,
tpm_limit=None,
rpm_limit=2,
rate_limit_type="key",
)
tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py
+374
View File
@@ -0,0 +1,374 @@
"""
Unit Tests for the max parallel request limiter v3 for the proxy
"""
import asyncio
import os
import sys
from datetime import datetime
import pytest
from fastapi import HTTPException
import litellm
from litellm import Router
from litellm.caching.caching import DualCache
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.hooks.parallel_request_limiter_v3 import (
_PROXY_MaxParallelRequestsHandler_v3 as _PROXY_MaxParallelRequestsHandler,
)
from litellm.proxy.utils import InternalUsageCache, ProxyLogging, hash_token
@pytest.mark.flaky(reruns=3)
@pytest.mark.asyncio
async def test_sliding_window_rate_limit_v3(monkeypatch):
"""
Test the sliding window rate limiting functionality
"""
monkeypatch.setenv("LITELLM_RATE_LIMIT_WINDOW_SIZE", "2")
_api_key = "sk-12345"
_api_key = hash_token(_api_key)
user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, rpm_limit=3)
local_cache = DualCache()
parallel_request_handler = _PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(local_cache)
)
# Mock the batch_rate_limiter_script to simulate window expiry and use correct key construction
window_starts = {}
async def mock_batch_rate_limiter(*args, **kwargs):
keys = kwargs.get("keys") if "keys" in kwargs else args[0]
now = kwargs.get("args")[0] if "args" in kwargs else args[1][0]
window_size = kwargs.get("args")[1] if "args" in kwargs else args[1][1]
results = []
for i in range(0, len(keys), 3):
window_key = keys[i]
counter_key = keys[i + 1]
# Simulate window expiry
prev_window = window_starts.get(window_key)
prev_counter = await local_cache.async_get_cache(key=counter_key) or 0
if prev_window is None or (now - prev_window) >= window_size:
# Window expired, reset
window_starts[window_key] = now
new_counter = 1
await local_cache.async_set_cache(
key=window_key, value=now, ttl=window_size
)
await local_cache.async_set_cache(
key=counter_key, value=new_counter, ttl=window_size
)
else:
new_counter = prev_counter + 1
await local_cache.async_set_cache(
key=counter_key, value=new_counter, ttl=window_size
)
results.append(now)
results.append(new_counter)
return results
parallel_request_handler.batch_rate_limiter_script = mock_batch_rate_limiter
# First request should succeed
await parallel_request_handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict, cache=local_cache, data={}, call_type=""
)
# Second request should succeed
await parallel_request_handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict, cache=local_cache, data={}, call_type=""
)
# Third request should fail
with pytest.raises(HTTPException) as exc_info:
await parallel_request_handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=local_cache,
data={},
call_type="",
)
assert exc_info.value.status_code == 429
assert "Rate limit exceeded" in str(exc_info.value.detail)
# Wait for window to expire (2 seconds)
await asyncio.sleep(3)
print("WAITED 3 seconds")
print(f"local_cache: {local_cache.in_memory_cache.cache_dict}")
# After window expires, should be able to make requests again
await parallel_request_handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict, cache=local_cache, data={}, call_type=""
)
@pytest.mark.asyncio
async def test_rate_limiter_script_return_values_v3(monkeypatch):
"""
Test that the rate limiter script returns both counter and window values correctly
"""
monkeypatch.setenv("LITELLM_RATE_LIMIT_WINDOW_SIZE", "2")
_api_key = "sk-12345"
_api_key = hash_token(_api_key)
user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, rpm_limit=3)
local_cache = DualCache()
parallel_request_handler = _PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(local_cache)
)
# Mock the batch_rate_limiter_script to simulate window expiry and use correct key construction
window_starts = {}
async def mock_batch_rate_limiter(*args, **kwargs):
keys = kwargs.get("keys") if "keys" in kwargs else args[0]
now = kwargs.get("args")[0] if "args" in kwargs else args[1][0]
window_size = kwargs.get("args")[1] if "args" in kwargs else args[1][1]
results = []
for i in range(0, len(keys), 3):
window_key = keys[i]
counter_key = keys[i + 1]
# Simulate window expiry
prev_window = window_starts.get(window_key)
prev_counter = await local_cache.async_get_cache(key=counter_key) or 0
if prev_window is None or (now - prev_window) >= window_size:
# Window expired, reset
window_starts[window_key] = now
new_counter = 1
await local_cache.async_set_cache(
key=window_key, value=now, ttl=window_size
)
await local_cache.async_set_cache(
key=counter_key, value=new_counter, ttl=window_size
)
else:
new_counter = prev_counter + 1
await local_cache.async_set_cache(
key=counter_key, value=new_counter, ttl=window_size
)
results.append(now)
results.append(new_counter)
return results
parallel_request_handler.batch_rate_limiter_script = mock_batch_rate_limiter
# Make first request
await parallel_request_handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict, cache=local_cache, data={}, call_type=""
)
# Verify both counter and window values are stored in cache
window_key = f"{{api_key:{_api_key}}}:window"
counter_key = f"{{api_key:{_api_key}}}:requests"
window_value = await local_cache.async_get_cache(key=window_key)
counter_value = await local_cache.async_get_cache(key=counter_key)
assert window_value is not None, "Window value should be stored in cache"
assert counter_value is not None, "Counter value should be stored in cache"
assert counter_value == 1, "Counter should be 1 after first request"
# Make second request
await parallel_request_handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict, cache=local_cache, data={}, call_type=""
)
# Verify counter increased but window stayed same
new_window_value = await local_cache.async_get_cache(key=window_key)
new_counter_value = await local_cache.async_get_cache(key=counter_key)
assert (
new_window_value == window_value
), "Window value should not change within window"
assert new_counter_value == 2, "Counter should be 2 after second request"
# Wait for window to expire
await asyncio.sleep(3)
# Make request after window expiry
await parallel_request_handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict, cache=local_cache, data={}, call_type=""
)
# Verify new window and reset counter
final_window_value = await local_cache.async_get_cache(key=window_key)
final_counter_value = await local_cache.async_get_cache(key=counter_key)
assert final_window_value != window_value, "Window value should change after expiry"
assert final_counter_value == 1, "Counter should reset to 1 after window expiry"
@pytest.mark.parametrize(
"rate_limit_object",
[
"api_key",
"model_per_key",
"user",
"end_user",
"team",
],
)
@pytest.mark.flaky(reruns=3)
@pytest.mark.asyncio
async def test_normal_router_call_tpm_v3(monkeypatch, rate_limit_object):
"""
Test normal router call with parallel request limiter v3 for TPM rate limiting
"""
monkeypatch.setenv("LITELLM_RATE_LIMIT_WINDOW_SIZE", "2")
model_list = [
{
"model_name": "azure-model",
"litellm_params": {
"model": "azure/gpt-turbo",
"api_key": "os.environ/AZURE_FRANCE_API_KEY",
"api_base": "https://openai-france-1234.openai.azure.com",
"rpm": 1440,
},
"model_info": {"id": 1},
},
{
"model_name": "azure-model",
"litellm_params": {
"model": "azure/gpt-35-turbo",
"api_key": "os.environ/AZURE_EUROPE_API_KEY",
"api_base": "https://my-endpoint-europe-berri-992.openai.azure.com",
"rpm": 6,
},
"model_info": {"id": 2},
},
]
router = Router(
model_list=model_list,
set_verbose=False,
num_retries=3,
) # type: ignore
_api_key = "sk-12345"
_api_key = hash_token(_api_key)
if rate_limit_object == "api_key":
user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, tpm_limit=10)
elif rate_limit_object == "user":
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 == "end_user":
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)
)
# Mock the batch_rate_limiter_script to simulate window expiry and use correct key construction
window_starts = {}
async def mock_batch_rate_limiter(*args, **kwargs):
keys = kwargs.get("keys") if "keys" in kwargs else args[0]
now = kwargs.get("args")[0] if "args" in kwargs else args[1][0]
window_size = kwargs.get("args")[1] if "args" in kwargs else args[1][1]
results = []
for i in range(0, len(keys), 3):
window_key = keys[i]
counter_key = keys[i + 1]
# Simulate window expiry
prev_window = window_starts.get(window_key)
prev_counter = await local_cache.async_get_cache(key=counter_key) or 0
if prev_window is None or (now - prev_window) >= window_size:
# Window expired, reset
window_starts[window_key] = now
new_counter = 1
await local_cache.async_set_cache(
key=window_key, value=now, ttl=window_size
)
await local_cache.async_set_cache(
key=counter_key, value=new_counter, ttl=window_size
)
else:
new_counter = prev_counter + 1
await local_cache.async_set_cache(
key=counter_key, value=new_counter, ttl=window_size
)
results.append(now)
results.append(new_counter)
return results
parallel_request_handler.batch_rate_limiter_script = mock_batch_rate_limiter
monkeypatch.setattr(litellm, "callbacks", [parallel_request_handler])
# Helper to get the correct value for key construction
def get_value_for_key(rate_limit_object, user_api_key_dict, model_name):
if rate_limit_object == "api_key":
return user_api_key_dict.api_key
elif rate_limit_object == "user":
return user_api_key_dict.user_id
elif rate_limit_object == "team":
return user_api_key_dict.team_id
elif rate_limit_object == "end_user":
return user_api_key_dict.end_user_id
elif rate_limit_object == "model_per_key":
return f"{user_api_key_dict.api_key}:{model_name}"
return None
value = get_value_for_key(rate_limit_object, user_api_key_dict, "azure-model")
counter_key = parallel_request_handler.create_rate_limit_keys(
rate_limit_object, value, "tokens"
)
# First request should succeed
await parallel_request_handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=local_cache,
data={"model": "azure-model"},
call_type="",
)
# normal call
response = await router.acompletion(
model="azure-model",
messages=[{"role": "user", "content": "Hey, how's it going?"}],
metadata={
"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",
)
await asyncio.sleep(1) # success is done in a separate thread
# Verify the token count is tracked
counter_value = await local_cache.async_get_cache(key=counter_key)
print(f"local_cache: {local_cache.in_memory_cache.cache_dict}")
assert (
counter_value is not None
), f"Counter value should be stored in cache for {counter_key}"
# Make another request to test rate limiting
with pytest.raises(HTTPException) as exc_info:
await parallel_request_handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=local_cache,
data={"model": "azure-model"},
call_type="",
)
# Wait for window to expire
await asyncio.sleep(3)
# Make request after window expiry
await parallel_request_handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=local_cache,
data={"model": "azure-model"},
call_type="",
)
# Verify new window and reset counter
final_counter_value = await local_cache.async_get_cache(key=counter_key)
assert final_counter_value == 1, "Counter should reset to 1 after window expiry"
tests/test_litellm/router_strategy/test_base_routing_strategy.py
+8 -12
View File
@@ -96,8 +96,13 @@ async def test_push_in_memory_increments_to_redis(base_strategy, mock_dual_cache
@pytest.mark.asyncio
async def test_sync_in_memory_spend_with_redis(base_strategy, mock_dual_cache):
from litellm.types.caching import RedisPipelineIncrementOperation
# Setup test data
base_strategy.in_memory_keys_to_update = {"key1"}
base_strategy.redis_increment_operation_queue = [
RedisPipelineIncrementOperation(key="key1", increment_value=10, ttl=3600),
]
# Mock the in-memory cache batch get responses for before snapshot
in_memory_before_future: asyncio.Future[List[str]] = asyncio.Future()
@@ -108,8 +113,8 @@ async def test_sync_in_memory_spend_with_redis(base_strategy, mock_dual_cache):
# Mock Redis batch get response
redis_future: asyncio.Future[Dict[str, str]] = asyncio.Future()
redis_future.set_result({"key1": "15.0"}) # Redis values
mock_dual_cache.redis_cache.async_batch_get_cache.return_value = redis_future
redis_future.set_result([15.0]) # Redis values
mock_dual_cache.redis_cache.async_increment_pipeline.return_value = redis_future
# Mock in-memory get for after snapshot
in_memory_after_future: asyncio.Future[Optional[str]] = asyncio.Future()
@@ -120,18 +125,9 @@ async def test_sync_in_memory_spend_with_redis(base_strategy, mock_dual_cache):
await base_strategy._sync_in_memory_spend_with_redis()
# Verify Redis batch get was called with correct keys
key_list = mock_dual_cache.redis_cache.async_batch_get_cache.call_args.kwargs[
"key_list"
]
assert sorted(key_list) == sorted(["key1"])
# Verify in-memory cache was updated with merged values
# For key1: redis_val(15.0) + delta(8.0 - 5.0) = 18.0
assert mock_dual_cache.in_memory_cache.async_set_cache.call_count == 1
# Verify the final merged values
set_cache_calls = mock_dual_cache.in_memory_cache.async_set_cache.call_args_list
print(f"set_cache_calls: {set_cache_calls}")
assert any(
call.kwargs["key"] == "key1" and float(call.kwargs["value"]) == 18.0
for call in set_cache_calls