From c42740a4b93f313488e45db7f404c2576d3cb6d4 Mon Sep 17 00:00:00 2001 From: Krish Dholakia Date: Sat, 7 Jun 2025 11:10:55 -0700 Subject: [PATCH] Simplify experimental multi-instance rate limiter - more accurate (#11424) * 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 --- docs/my-website/docs/proxy/config_settings.md | 1 + litellm/caching/dual_cache.py | 28 + litellm/caching/in_memory_cache.py | 18 +- litellm/caching/redis_cache.py | 39 +- litellm/proxy/hooks/__init__.py | 4 +- .../hooks/parallel_request_limiter_v2.py | 573 -------------- .../hooks/parallel_request_limiter_v3.py | 741 ++++++++++++++++++ .../router_strategy/base_routing_strategy.py | 43 +- litellm/types/caching.py | 10 + tests/local_testing/test_caching.py | 2 +- .../hooks/test_parallel_request_limiter_v2.py | 626 --------------- .../hooks/test_parallel_request_limiter_v3.py | 374 +++++++++ .../test_base_routing_strategy.py | 20 +- 13 files changed, 1246 insertions(+), 1233 deletions(-) delete mode 100644 litellm/proxy/hooks/parallel_request_limiter_v2.py create mode 100644 litellm/proxy/hooks/parallel_request_limiter_v3.py delete mode 100644 tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v2.py create mode 100644 tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py diff --git a/docs/my-website/docs/proxy/config_settings.md b/docs/my-website/docs/proxy/config_settings.md index 2afde3c4b1..e8db12e51f 100644 --- a/docs/my-website/docs/proxy/config_settings.md +++ b/docs/my-website/docs/proxy/config_settings.md @@ -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 diff --git a/litellm/caching/dual_cache.py b/litellm/caching/dual_cache.py index 8bef333758..ce07f7ce70 100644 --- a/litellm/caching/dual_cache.py +++ b/litellm/caching/dual_cache.py @@ -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: diff --git a/litellm/caching/in_memory_cache.py b/litellm/caching/in_memory_cache.py index f644fdfefd..47f911894a 100644 --- a/litellm/caching/in_memory_cache.py +++ b/litellm/caching/in_memory_cache.py @@ -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() diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index 5339719c01..b8091187bf 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -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 diff --git a/litellm/proxy/hooks/__init__.py b/litellm/proxy/hooks/__init__.py index aef3cc0fcc..83d7e17343 100644 --- a/litellm/proxy/hooks/__init__.py +++ b/litellm/proxy/hooks/__init__.py @@ -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 ### diff --git a/litellm/proxy/hooks/parallel_request_limiter_v2.py b/litellm/proxy/hooks/parallel_request_limiter_v2.py deleted file mode 100644 index 34b53a0c22..0000000000 --- a/litellm/proxy/hooks/parallel_request_limiter_v2.py +++ /dev/null @@ -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) - ) - ) diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py new file mode 100644 index 0000000000..246d57ee67 --- /dev/null +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -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)}" + ) diff --git a/litellm/router_strategy/base_routing_strategy.py b/litellm/router_strategy/base_routing_strategy.py index f37dccab85..6e410ef14a 100644 --- a/litellm/router_strategy/base_routing_strategy.py +++ b/litellm/router_strategy/base_routing_strategy.py @@ -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 ) diff --git a/litellm/types/caching.py b/litellm/types/caching.py index c15d4d1590..4893524998 100644 --- a/litellm/types/caching.py +++ b/litellm/types/caching.py @@ -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", { diff --git a/tests/local_testing/test_caching.py b/tests/local_testing/test_caching.py index a164894711..7560bccadb 100644 --- a/tests/local_testing/test_caching.py +++ b/tests/local_testing/test_caching.py @@ -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") diff --git a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v2.py b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v2.py deleted file mode 100644 index e1af7f441b..0000000000 --- a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v2.py +++ /dev/null @@ -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", - ) diff --git a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py new file mode 100644 index 0000000000..f84d108ba2 --- /dev/null +++ b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py @@ -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" diff --git a/tests/test_litellm/router_strategy/test_base_routing_strategy.py b/tests/test_litellm/router_strategy/test_base_routing_strategy.py index 62be5db306..2ce144e9a9 100644 --- a/tests/test_litellm/router_strategy/test_base_routing_strategy.py +++ b/tests/test_litellm/router_strategy/test_base_routing_strategy.py @@ -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