From d538cf489a67c1cbe3f4b4804922b9d26e6d6799 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 1 Oct 2025 18:35:34 -0700 Subject: [PATCH] [Feat] Fixes to dynamic rate limiter v3 - add saturatation detection (#15119) * test cases dynamic rate limits * fix _handle_generous_mode * docs add readme * use configs for vars * fix debug * add comment * test_dynamic_rate_limiter_v3.py * test_concurrent_pre_call_hooks_stress --- .../hooks/README.dynamic_rate_limiter_v3.md | 170 +++++ .../proxy/hooks/dynamic_rate_limiter_v3.py | 398 +++++++++-- litellm/types/utils.py | 10 + .../hooks/test_dynamic_rate_limiter_v3.py | 668 +++++++++++++++++- 4 files changed, 1184 insertions(+), 62 deletions(-) create mode 100644 litellm/proxy/hooks/README.dynamic_rate_limiter_v3.md diff --git a/litellm/proxy/hooks/README.dynamic_rate_limiter_v3.md b/litellm/proxy/hooks/README.dynamic_rate_limiter_v3.md new file mode 100644 index 0000000000..701f5e928c --- /dev/null +++ b/litellm/proxy/hooks/README.dynamic_rate_limiter_v3.md @@ -0,0 +1,170 @@ +# Dynamic Rate Limiter v3 - Saturation-Aware Priority-Based Rate Limiting + +## Overview + +The v3 dynamic rate limiter implements saturation-aware rate limiting with priority-based allocation. It balances resource efficiency (allowing unused capacity to be borrowed) with fairness guarantees (enforcing priorities during high load). + +**Key Behavior:** +- When system is under 80% capacity: Generous mode - allows priority borrowing +- When system is at/above 80% capacity: Strict mode - enforces normalized priority limits + +## How It Works + +### Flow Diagram + +``` +┌─────────────────────────────────────────────────────────────┐ +│ Incoming Request │ +└────────────────────────┬────────────────────────────────────┘ + │ + ▼ +┌─────────────────────────────────────────────────────────────┐ +│ 1. Check Model Saturation │ +│ - Query v3 limiter's Redis counters │ +│ - Calculate: current_usage / capacity │ +│ - Returns: 0.0 (empty) to 1.0+ (saturated) │ +└────────────────────────┬────────────────────────────────────┘ + │ + ▼ + ┌────────┴────────┐ + │ Saturation? │ + └────────┬────────┘ + │ + ┌───────────────┴───────────────┐ + │ │ + ▼ ▼ + < 80% (Generous) >= 80% (Strict) + │ │ + ▼ ▼ +┌─────────────────────┐ ┌─────────────────────┐ +│ Generous Mode │ │ Strict Mode │ +│ │ │ │ +│ - Enforce model- │ │ - Normalize │ +│ wide capacity │ │ priority weights │ +│ - No priority │ │ (if over 1.0) │ +│ restrictions │ │ │ +│ - Allows borrowing │ │ - Create priority- │ +│ │ │ specific │ +│ - First-come- │ │ descriptors │ +│ first-served │ │ │ +│ until capacity │ │ - Enforce strict │ +│ │ │ limits per │ +│ │ │ priority │ +└──────────┬──────────┘ └──────────┬──────────┘ + │ │ + │ ▼ + │ ┌──────────────────────┐ + │ │ Track model usage │ + │ │ for future │ + │ │ saturation checks │ + │ └──────────┬───────────┘ + │ │ + └───────────────┬───────────────┘ + │ + ▼ + ┌──────────────┐ + │ v3 Limiter │ + │ Check │ + └──────┬───────┘ + │ + ┌───────────────┴───────────────┐ + │ │ + ▼ ▼ + OVER_LIMIT OK + │ │ + ▼ ▼ + Return 429 Error Allow Request +``` + +## Configuration + +### Priority Reservation + +Set priority weights in your proxy configuration: + +```python +litellm.priority_reservation = { + "premium": 0.75, # 75% of capacity + "standard": 0.25 # 25% of capacity +} +``` + +### Priority Reservation Settings + +Configure saturation-aware behavior: + +```python +litellm.priority_reservation_settings = PriorityReservationSettings( + default_priority=0.5, # Default weight for users without explicit priority + saturation_threshold=0.80, # 80% - threshold for strict mode enforcement + tracking_multiplier=10 # 10x - multiplier for non-blocking tracking in strict mode +) +``` + +**Settings:** +- `default_priority` (default: 0.5) - Priority weight for users without explicit priority metadata +- `saturation_threshold` (default: 0.80) - Saturation level (0.0-1.0) at which strict priority enforcement begins +- `tracking_multiplier` (default: 10) - Multiplier for model-wide tracking limits in strict mode + +### User Priority Assignment + +Set priority in user metadata: + +```python +user_api_key_dict.metadata = {"priority": "premium"} +``` + +## Priority Weight Normalization + +If priorities sum to > 1.0, they are automatically normalized: + +``` +Input: {key_a: 0.60, key_b: 0.80} = 1.40 total +Output: {key_a: 0.43, key_b: 0.57} = 1.00 total +``` + +This ensures total allocation never exceeds model capacity. + +## Implementation Details + +### Saturation Detection + +- Queries v3 limiter's Redis counters for model-wide usage +- Checks both RPM and TPM, returns higher saturation value +- Non-blocking reads (doesn't increment counters) + +### Mode Selection + +**Generous Mode (< 80% saturation):** +- Creates single model-wide descriptor +- Enforces total capacity only +- Allows any priority to use available capacity +- Prevents over-subscription via model-wide limit + +**Strict Mode (>= 80% saturation):** +- Creates priority-specific descriptors with normalized weights +- Each priority gets its reserved allocation +- Tracks model-wide usage separately (non-blocking, 10x multiplier) +- Ensures fairness under load + +Test scenarios covered: +1. No rate limiting when under capacity +2. Priority queue behavior during saturation +3. Spillover capacity for default keys +4. Over-allocated priorities with normalization +5. Default priority value handling + + +### `_PROXY_DynamicRateLimitHandlerV3` + +Main handler class inheriting from `CustomLogger`. + +**Key Methods:** +- `async_pre_call_hook()` - Main entry point, routes to generous/strict mode +- `_check_model_saturation()` - Queries Redis for current usage +- `_handle_generous_mode()` - Enforces model-wide capacity only +- `_handle_strict_mode()` - Enforces normalized priority limits +- `_normalize_priority_weights()` - Handles over-allocation +- `_create_priority_based_descriptors()` - Creates rate limit descriptors + + diff --git a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py index 38e211dea5..5d0157f436 100644 --- a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py +++ b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py @@ -1,9 +1,9 @@ """ -Dynamic rate limiter v3 +Dynamic rate limiter v3 - Saturation-aware priority-based rate limiting """ import os -from typing import List, Literal, Optional, Union +from typing import Dict, List, Literal, Optional, Union from fastapi import HTTPException @@ -24,12 +24,18 @@ from litellm.types.router import ModelGroupInfo class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): """ - Simple validation version that uses v3 parallel request limiter for priority-based rate limiting. + Saturation-aware priority-based rate limiter using v3 infrastructure. - Key differences from original: - 1. Uses v3 limiter's sliding window approach instead of per-minute cache buckets - 2. Leverages Redis Lua scripts for atomic operations under high traffic - 3. Creates priority-specific rate limit descriptors + Key features: + 1. Reuses v3 limiter's Redis-based tracking (works across multiple instances) + 2. Only enforces priority limits when model is saturated (>80% usage) + 3. When under capacity, allows all requests (generous behavior) + 4. When saturated, enforces strict priority-based limits (fairness) + + How it works: + - Uses v3 limiter's counter keys to check model-wide saturation + - Saturation check reads existing counters without incrementing + - Priority enforcement reuses v3 limiter's atomic Lua scripts """ def __init__(self, internal_usage_cache: DualCache): self.internal_usage_cache = InternalUsageCache(dual_cache=internal_usage_cache) @@ -57,6 +63,107 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): weight = litellm.priority_reservation[priority] return weight + def _normalize_priority_weights(self) -> Dict[str, float]: + """ + Normalize priority weights if they sum to > 1.0 + + Handles over-allocation: {key_a: 0.60, key_b: 0.80} -> {key_a: 0.43, key_b: 0.57} + """ + if litellm.priority_reservation is None: + return {} + + weights = dict(litellm.priority_reservation) + total_weight = sum(weights.values()) + + if total_weight > 1.0: + normalized = {k: v / total_weight for k, v in weights.items()} + verbose_proxy_logger.debug( + f"Normalized over-allocated priorities: {weights} -> {normalized}" + ) + return normalized + + return weights + + async def _check_model_saturation( + self, + model: str, + model_group_info: ModelGroupInfo, + ) -> float: + """ + Check current saturation by directly querying v3 limiter's cache keys. + + Reuses v3 limiter's Redis-based tracking (works across multiple instances). + Reads counters WITHOUT incrementing them. + + Returns: + float: Saturation ratio (0.0 = empty, 1.0 = at capacity, >1.0 = over) + """ + try: + max_saturation = 0.0 + + # Query RPM saturation + if model_group_info.rpm is not None and model_group_info.rpm > 0: + # Use v3 limiter's key format: {key:value}:rate_limit_type + counter_key = self.v3_limiter.create_rate_limit_keys( + key="model_saturation_check", + value=model, + rate_limit_type="requests", + ) + + # Query cache for current counter value + counter_value = await self.internal_usage_cache.async_get_cache( + key=counter_key, + litellm_parent_otel_span=None, + local_only=False, # Check Redis too + ) + + if counter_value is not None: + current_requests = int(counter_value) + rpm_saturation = current_requests / model_group_info.rpm + max_saturation = max(max_saturation, rpm_saturation) + + verbose_proxy_logger.debug( + f"Model {model} RPM: {current_requests}/{model_group_info.rpm} " + f"({rpm_saturation:.1%})" + ) + + # Query TPM saturation + if model_group_info.tpm is not None and model_group_info.tpm > 0: + counter_key = self.v3_limiter.create_rate_limit_keys( + key="model_saturation_check", + value=model, + rate_limit_type="tokens", + ) + + counter_value = await self.internal_usage_cache.async_get_cache( + key=counter_key, + litellm_parent_otel_span=None, + local_only=False, + ) + + if counter_value is not None: + current_tokens = float(counter_value) + tpm_saturation = current_tokens / model_group_info.tpm + max_saturation = max(max_saturation, tpm_saturation) + + verbose_proxy_logger.debug( + f"Model {model} TPM: {current_tokens}/{model_group_info.tpm} " + f"({tpm_saturation:.1%})" + ) + + verbose_proxy_logger.debug( + f"Model {model} overall saturation: {max_saturation:.1%}" + ) + + return max_saturation + + except Exception as e: + verbose_proxy_logger.error( + f"Error checking saturation for {model}: {str(e)}" + ) + # Fail open: assume not saturated on error + return 0.0 + def _create_priority_based_descriptors( self, model: str, @@ -64,11 +171,10 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): priority: Optional[str], ) -> List[RateLimitDescriptor]: """ - Create rate limit descriptors based on priority and model group limits. + Create rate limit descriptors with normalized priority weights. - This is the key change: instead of calculating dynamic quotas based on active projects, - we create descriptors with priority-adjusted limits and let the v3 limiter handle - the actual rate limiting with its sliding window approach. + Uses normalized weights to handle over-allocation scenarios. + Only called when system is saturated. """ descriptors: List[RateLimitDescriptor] = [] @@ -79,8 +185,13 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): if model_group_info is None: return descriptors - # Get priority weight - priority_weight = self._get_priority_weight(priority) + # Get normalized priority weight (handles over-allocation) + normalized_weights = self._normalize_priority_weights() + priority_weight = normalized_weights.get(priority, None) if priority else None + if priority_weight is None: + # Fallback to non-normalized weight + priority_weight = self._get_priority_weight(priority) + # Create priority-specific rate limits # Use model:priority as the key to separate different priority levels @@ -88,16 +199,17 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): rate_limit_config: RateLimitDescriptorRateLimitObject = {} - # Apply priority weight to model limits + # Apply normalized priority weight to model limits if model_group_info.tpm is not None: - # Reserve portion of TPM based on priority + # Reserve portion of TPM based on normalized priority reserved_tpm = int(model_group_info.tpm * priority_weight) rate_limit_config["tokens_per_unit"] = reserved_tpm if model_group_info.rpm is not None: - # Reserve portion of RPM based on priority + # Reserve portion of RPM based on normalized priority reserved_rpm = int(model_group_info.rpm * priority_weight) rate_limit_config["requests_per_unit"] = reserved_rpm + if rate_limit_config: rate_limit_config["window_size"] = self.v3_limiter.window_size @@ -112,6 +224,171 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): return descriptors + def _create_model_tracking_descriptor( + self, + model: str, + model_group_info: ModelGroupInfo, + high_limit_multiplier: int = 1, + ) -> RateLimitDescriptor: + """ + Create a descriptor for tracking model-wide usage. + + Args: + model: Model name + model_group_info: Model configuration with RPM/TPM limits + high_limit_multiplier: Multiplier for limits (use >1 for tracking-only) + + Returns: + Rate limit descriptor for model-wide tracking + """ + return RateLimitDescriptor( + key="model_saturation_check", + value=model, + rate_limit={ + "requests_per_unit": ( + model_group_info.rpm * high_limit_multiplier + if model_group_info.rpm else None + ), + "tokens_per_unit": ( + model_group_info.tpm * high_limit_multiplier + if model_group_info.tpm else None + ), + "window_size": self.v3_limiter.window_size, + }, + ) + + async def _handle_generous_mode( + self, + model: str, + model_group_info: ModelGroupInfo, + user_api_key_dict: UserAPIKeyAuth, + key_priority: Optional[str], + ) -> None: + """ + Handle rate limiting in generous mode (under saturation threshold). + + In this mode, we enforce model-wide capacity but NOT priority-specific limits. + This allows lower-priority users to borrow unused capacity from higher-priority users. + + Args: + model: Model name + model_group_info: Model configuration + user_api_key_dict: User authentication info + key_priority: User's priority level + + Raises: + HTTPException: If model capacity is reached + """ + descriptor = self._create_model_tracking_descriptor( + model=model, + model_group_info=model_group_info, + high_limit_multiplier=1, # Enforce actual limits in generous mode + ) + + response = await self.v3_limiter.should_rate_limit( + descriptors=[descriptor], + parent_otel_span=user_api_key_dict.parent_otel_span, + ) + + if response["overall_code"] == "OVER_LIMIT": + for status in response["statuses"]: + if status["code"] == "OVER_LIMIT": + raise HTTPException( + status_code=429, + detail={ + "error": f"Model capacity reached for {model}. " + f"Priority: {key_priority}, " + f"Rate limit type: {status['rate_limit_type']}, " + f"Remaining: {status['limit_remaining']}" + }, + headers={ + "retry-after": str(self.v3_limiter.window_size), + "rate_limit_type": str(status["rate_limit_type"]), + "x-litellm-priority": key_priority or "default", + }, + ) + + async def _handle_strict_mode( + self, + model: str, + model_group_info: ModelGroupInfo, + user_api_key_dict: UserAPIKeyAuth, + key_priority: Optional[str], + saturation: float, + data: dict, + ) -> None: + """ + Handle rate limiting in strict mode (above saturation threshold). + + In this mode, we enforce priority-specific limits using normalized weights. + + Args: + model: Model name + model_group_info: Model configuration + user_api_key_dict: User authentication info + key_priority: User's priority level + saturation: Current saturation level + data: Request data dictionary + + Raises: + HTTPException: If priority-specific limit is exceeded + """ + # Create priority-based descriptors + descriptors = self._create_priority_based_descriptors( + model=model, + user_api_key_dict=user_api_key_dict, + priority=key_priority, + ) + + if not descriptors: + verbose_proxy_logger.debug("No rate limit descriptors created, allowing request") + return + + # Track model-wide usage for future saturation checks + # Why tracking_multiplier: v3_limiter.should_rate_limit() both increments AND checks limits. + # We need the increment (for saturation detection) but NOT the limit check (priority limits handle enforcement). + # Setting limit to 10x capacity ensures tracking never blocks while keeping accurate counters. + tracking_multiplier = litellm.priority_reservation_settings.tracking_multiplier + tracking_descriptor = self._create_model_tracking_descriptor( + model=model, + model_group_info=model_group_info, + high_limit_multiplier=tracking_multiplier, + ) + + await self.v3_limiter.should_rate_limit( + descriptors=[tracking_descriptor], + parent_otel_span=user_api_key_dict.parent_otel_span, + ) + + # Enforce priority-specific limits + response = await self.v3_limiter.should_rate_limit( + descriptors=descriptors, + parent_otel_span=user_api_key_dict.parent_otel_span, + ) + + if response["overall_code"] == "OVER_LIMIT": + for status in response["statuses"]: + if status["code"] == "OVER_LIMIT": + raise HTTPException( + status_code=429, + detail={ + "error": f"Priority-based rate limit exceeded for {status['descriptor_key']}. " + f"Priority: {key_priority}, " + f"Rate limit type: {status['rate_limit_type']}, " + f"Remaining: {status['limit_remaining']}, " + f"Model saturation: {saturation:.1%}" + }, + headers={ + "retry-after": str(self.v3_limiter.window_size), + "rate_limit_type": str(status["rate_limit_type"]), + "x-litellm-priority": key_priority or "default", + "x-litellm-saturation": f"{saturation:.2%}", + }, + ) + else: + # Store response for post-call hook + data["litellm_proxy_rate_limit_response"] = response + async def async_pre_call_hook( self, user_api_key_dict: UserAPIKeyAuth, @@ -130,60 +407,73 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): ], ) -> Optional[Union[Exception, str, dict]]: """ - Pre-call hook using v3 limiter for priority-based rate limiting. + Saturation-aware pre-call hook for priority-based rate limiting. + + This hook implements a two-mode rate limiting strategy: + - Generous mode (< 80% saturation): Enforces model capacity, allows priority borrowing + - Strict mode (>= 80% saturation): Enforces normalized priority-based limits + + Args: + user_api_key_dict: User authentication and metadata + cache: Dual cache instance + data: Request data containing model name + call_type: Type of API call being made + + Returns: + None if request is allowed, otherwise raises HTTPException """ if "model" not in data: return None + model = data["model"] key_priority: Optional[str] = user_api_key_dict.metadata.get("priority", None) - # Create priority-based descriptors - descriptors = self._create_priority_based_descriptors( - model=data["model"], - user_api_key_dict=user_api_key_dict, - priority=key_priority, + # Get model configuration + model_group_info: Optional[ModelGroupInfo] = self.llm_router.get_model_group_info( + model_group=model ) - - if not descriptors: - verbose_proxy_logger.debug("No rate limit descriptors created, allowing request") + if model_group_info is None: + verbose_proxy_logger.debug(f"No model group info for {model}, allowing request") return None + # Check current saturation level try: - # Use v3 limiter to check rate limits - response = await self.v3_limiter.should_rate_limit( - descriptors=descriptors, - parent_otel_span=user_api_key_dict.parent_otel_span, + saturation = await self._check_model_saturation(model, model_group_info) + + saturation_threshold = litellm.priority_reservation_settings.saturation_threshold + + verbose_proxy_logger.debug( + f"[Dynamic Rate Limiter] Model={model}, Saturation={saturation:.1%}, " + f"Threshold={saturation_threshold:.1%}, Priority={key_priority}" ) - - if response["overall_code"] == "OVER_LIMIT": - # Find which descriptor hit the limit - for status in response["statuses"]: - if status["code"] == "OVER_LIMIT": - raise HTTPException( - status_code=429, - detail={ - "error": f"Priority-based rate limit exceeded for {status['descriptor_key']}. " - f"Priority: {key_priority}, " - f"Rate limit type: {status['rate_limit_type']}, " - f"Remaining: {status['limit_remaining']}" - }, - headers={ - "retry-after": str(self.v3_limiter.window_size), - "rate_limit_type": str(status["rate_limit_type"]), - "x-litellm-priority": key_priority or "default", - }, - ) + + data["litellm_model_saturation"] = saturation + + # Route to appropriate mode based on saturation + if saturation < saturation_threshold: + await self._handle_generous_mode( + model=model, + model_group_info=model_group_info, + user_api_key_dict=user_api_key_dict, + key_priority=key_priority, + ) else: - # Store response for post-call hook - data["litellm_proxy_rate_limit_response"] = response - + await self._handle_strict_mode( + model=model, + model_group_info=model_group_info, + user_api_key_dict=user_api_key_dict, + key_priority=key_priority, + saturation=saturation, + data=data, + ) + except HTTPException: raise except Exception as e: - verbose_proxy_logger.exception( - f"Error in dynamic rate limiter v3 pre-call hook: {str(e)}" + verbose_proxy_logger.error( + f"Error in dynamic rate limiter: {str(e)}, allowing request" ) - # Allow request to proceed on unexpected errors + # Fail open on unexpected errors return None return None diff --git a/litellm/types/utils.py b/litellm/types/utils.py index bcf0fa1374..d303485b3d 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -2692,5 +2692,15 @@ class PriorityReservationSettings(BaseModel): default=0.5, description="Priority level to assign to API keys without explicit priority metadata. Should match a key in litellm.priority_reservation." ) + + saturation_threshold: float = Field( + default=0.80, + description="Saturation threshold (0.0-1.0) at which strict priority enforcement begins. Below this threshold, generous mode allows priority borrowing. Above this threshold, strict mode enforces normalized priority limits." + ) + + tracking_multiplier: int = Field( + default=10, + description="Multiplier for model-wide tracking limits in strict mode. Set to 10x because v3_limiter.should_rate_limit() both increments counters AND enforces limits - we need the counter increment (for saturation checks) but not the enforcement (priority limits handle that). High multiplier ensures tracking never blocks." + ) model_config = ConfigDict(protected_namespaces=()) diff --git a/tests/test_litellm/proxy/hooks/test_dynamic_rate_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_dynamic_rate_limiter_v3.py index 05e0d4287a..2b7c080e65 100644 --- a/tests/test_litellm/proxy/hooks/test_dynamic_rate_limiter_v3.py +++ b/tests/test_litellm/proxy/hooks/test_dynamic_rate_limiter_v3.py @@ -364,9 +364,12 @@ async def test_100_concurrent_priority_requests(): @pytest.mark.asyncio async def test_concurrent_pre_call_hooks_stress(): """ - Stress test: 50 concurrent pre-call hooks with priority enforcement. + Stress test: 50 concurrent pre-call hooks with saturation-aware priority enforcement. - This tests the actual rate limiting logic under concurrent load. + Tests priority-based rate limiting in strict mode (>80% saturation). + Mocks high saturation to force strict mode where priorities are enforced. + Premium users (80% allocation) should have >90% success rate. + Standard users (20% allocation) should have ~70% success rate with 30% random limiting. """ # Set up environment for premium feature os.environ["LITELLM_LICENSE"] = "test-license-key" @@ -398,10 +401,39 @@ async def test_concurrent_pre_call_hooks_stress(): successful_requests = [] rate_limited_requests = [] + # Mock saturation check to return high saturation (forces strict mode) + async def mock_get_cache(key, litellm_parent_otel_span=None, local_only=False): + """Mock cache to simulate high saturation.""" + # Return high usage to trigger strict mode (>80% saturation) + if ":requests" in key or ":tokens" in key: + return 1800 # 1800/2000 = 90% saturation + return None + async def mock_should_rate_limit(descriptors, parent_otel_span=None): - """Mock rate limiter that allows premium users, limits some standard users.""" + """Mock rate limiter that handles saturation-aware descriptors.""" descriptor = descriptors[0] - priority = descriptor["value"].split(":")[-1] + descriptor_key = descriptor["key"] + descriptor_value = descriptor["value"] + + # Handle model-wide tracking (for both generous and strict mode tracking) + if descriptor_key == "model_saturation_check": + # Always allow model-wide tracking (doesn't enforce in our mock) + return { + "overall_code": "OK", + "statuses": [ + { + "code": "OK", + "descriptor_key": descriptor_value, + "rate_limit_type": "tokens_per_unit", + "limit_remaining": 10000, + } + ], + } + + # Handle priority-specific enforcement in strict mode + if descriptor_key == "priority_model": + # Extract priority from value like "pre-call-stress-model:premium" + priority = descriptor_value.split(":")[-1] if priority == "premium": # Allow all premium requests @@ -410,7 +442,7 @@ async def test_concurrent_pre_call_hooks_stress(): "statuses": [ { "code": "OK", - "descriptor_key": descriptor["value"], + "descriptor_key": descriptor_value, "rate_limit_type": "tokens_per_unit", "limit_remaining": 1000, } @@ -426,7 +458,7 @@ async def test_concurrent_pre_call_hooks_stress(): "statuses": [ { "code": "OVER_LIMIT", - "descriptor_key": descriptor["value"], + "descriptor_key": descriptor_value, "rate_limit_type": "tokens_per_unit", "limit_remaining": 0, } @@ -438,9 +470,22 @@ async def test_concurrent_pre_call_hooks_stress(): "statuses": [ { "code": "OK", - "descriptor_key": descriptor["value"], + "descriptor_key": descriptor_value, "rate_limit_type": "tokens_per_unit", "limit_remaining": 100, + } + ], + } + + # Default: allow + return { + "overall_code": "OK", + "statuses": [ + { + "code": "OK", + "descriptor_key": descriptor_value, + "rate_limit_type": "tokens_per_unit", + "limit_remaining": 1000, } ], } @@ -466,6 +511,8 @@ async def test_concurrent_pre_call_hooks_stress(): with patch.object( handler.v3_limiter, "should_rate_limit", side_effect=mock_should_rate_limit + ), patch.object( + handler.internal_usage_cache, "async_get_cache", side_effect=mock_get_cache ): try: result = await handler.async_pre_call_hook( @@ -534,7 +581,7 @@ async def test_concurrent_pre_call_hooks_stress(): ), f"Premium success rate should be >= 90%, got {premium_success_rate:.2%}" assert ( standard_success_rate >= 0.5 - ), f"Standard success rate should be >= 50%, got {standard_success_rate:.2%}" + ), f"Standard success rate should be >= 50% (with 30% random limiting, allows for variance), got {standard_success_rate:.2%}" assert ( premium_success_rate > standard_success_rate ), "Premium should have higher success rate than standard" @@ -550,3 +597,608 @@ async def test_concurrent_pre_call_hooks_stress(): ) print(f" - Total successful: {successful_count}/50 ({successful_count/50:.1%})") print(f" - Priority system working: Premium > Standard success rates") + +# These tests make actual async_pre_call_hook calls to simulate real traffic + + +@pytest.mark.asyncio +async def test_fake_calls_case_1_no_rate_limiting_at_capacity(): + """ + Test Case 1: No Rate Limiting When At Capacity + + System: 100 RPM capacity + Key A: priority_reservation=0.75 (75 RPM reserved) + Key B: priority_reservation=0.25 (25 RPM reserved) + Traffic A: 50 RPM + Traffic B: 50 RPM + Expected A: 50 RPM (no limiting, under reserved capacity) + Expected B: 50 RPM (no limiting, under reserved capacity) + + When traffic is under individual reservations, no rate limiting should occur. + """ + os.environ["LITELLM_LICENSE"] = "test-license-key" + + # Set up priority reservations + litellm.priority_reservation = {"key_a": 0.75, "key_b": 0.25} + + dual_cache = DualCache() + handler = DynamicRateLimitHandler(internal_usage_cache=dual_cache) + + model = "fake-call-test-1" + total_rpm = 100 + + llm_router = Router( + model_list=[ + { + "model_name": model, + "litellm_params": { + "model": "gpt-3.5-turbo", + "api_key": "test-key", + "api_base": "test-base", + "rpm": total_rpm, + }, + } + ] + ) + handler.update_variables(llm_router=llm_router) + + # Create users + key_a_user = UserAPIKeyAuth() + key_a_user.metadata = {"priority": "key_a"} + key_a_user.user_id = "key_a_user" + + key_b_user = UserAPIKeyAuth() + key_b_user.metadata = {"priority": "key_b"} + key_b_user.user_id = "key_b_user" + + # Track results + successful_requests = {"key_a": 0, "key_b": 0} + rate_limited_requests = {"key_a": 0, "key_b": 0} + + async def make_request(user, priority_name, request_id): + """Make a single request and track the result.""" + try: + result = await handler.async_pre_call_hook( + user_api_key_dict=user, + cache=dual_cache, + data={"model": model}, + call_type="completion", + ) + + if result is None: + successful_requests[priority_name] += 1 + return {"status": "success", "priority": priority_name} + else: + rate_limited_requests[priority_name] += 1 + return {"status": "rate_limited", "priority": priority_name} + + except Exception as e: + rate_limited_requests[priority_name] += 1 + return {"status": "rate_limited", "priority": priority_name, "error": str(e)} + + # Send 50 requests from each priority (within capacity) + tasks = [] + + for i in range(50): + tasks.append(make_request(key_a_user, "key_a", f"key_a_{i}")) + + for i in range(50): + tasks.append(make_request(key_b_user, "key_b", f"key_b_{i}")) + + start_time = time.time() + results = await asyncio.gather(*tasks, return_exceptions=True) + end_time = time.time() + + # Analyze results + total_successful = successful_requests["key_a"] + successful_requests["key_b"] + total_rate_limited = rate_limited_requests["key_a"] + rate_limited_requests["key_b"] + + print(f"Test Case 1 - No Rate Limiting When At Capacity:") + print(f" - Duration: {end_time - start_time:.2f}s") + print(f" - Key A: {successful_requests['key_a']}/50 successful (reserved 75 RPM)") + print(f" - Key B: {successful_requests['key_b']}/50 successful (reserved 25 RPM)") + print(f" - Total successful: {total_successful}/100") + print(f" - Total rate limited: {total_rate_limited}/100") + + # Both keys should get all their requests since they're under capacity + assert successful_requests["key_a"] >= 45, f"Key A should get ≥45 requests, got {successful_requests['key_a']}" + assert successful_requests["key_b"] >= 45, f"Key B should get ≥45 requests, got {successful_requests['key_b']}" + + +@pytest.mark.asyncio +async def test_fake_calls_case_2_priority_queue_during_saturation(): + """ + Test Case 2: Priority Queue Behavior During Saturation + + System: 100 RPM capacity + Key A: priority_reservation=0.75 (75 RPM reserved) + Key B: priority_reservation=0.25 (25 RPM reserved) + Traffic A: 200 RPM + Traffic B: 200 RPM + Expected A: 75 RPM (75% of capacity) + Expected B: 25 RPM (25% of capacity) + + When total traffic exceeds capacity, rate limiting enforces priority reservations. + """ + os.environ["LITELLM_LICENSE"] = "test-license-key" + + litellm.priority_reservation = {"key_a": 0.75, "key_b": 0.25} + + dual_cache = DualCache() + handler = DynamicRateLimitHandler(internal_usage_cache=dual_cache) + + model = "fake-call-test-2" + total_rpm = 100 + + llm_router = Router( + model_list=[ + { + "model_name": model, + "litellm_params": { + "model": "gpt-3.5-turbo", + "api_key": "test-key", + "api_base": "test-base", + "rpm": total_rpm, + }, + } + ] + ) + handler.update_variables(llm_router=llm_router) + + # Create users + key_a_user = UserAPIKeyAuth() + key_a_user.metadata = {"priority": "key_a"} + key_a_user.user_id = "key_a_user" + + key_b_user = UserAPIKeyAuth() + key_b_user.metadata = {"priority": "key_b"} + key_b_user.user_id = "key_b_user" + + # Track results + successful_requests = {"key_a": 0, "key_b": 0} + rate_limited_requests = {"key_a": 0, "key_b": 0} + + async def make_request(user, priority_name, request_id): + """Make a single request and track the result.""" + try: + result = await handler.async_pre_call_hook( + user_api_key_dict=user, + cache=dual_cache, + data={"model": model}, + call_type="completion", + ) + + if result is None: + successful_requests[priority_name] += 1 + return {"status": "success", "priority": priority_name} + else: + rate_limited_requests[priority_name] += 1 + return {"status": "rate_limited", "priority": priority_name} + + except Exception as e: + rate_limited_requests[priority_name] += 1 + return {"status": "rate_limited", "priority": priority_name, "error": str(e)} + + # Send 200 requests from each priority (over capacity) + tasks = [] + + for i in range(200): + tasks.append(make_request(key_a_user, "key_a", f"key_a_{i}")) + + for i in range(200): + tasks.append(make_request(key_b_user, "key_b", f"key_b_{i}")) + + start_time = time.time() + results = await asyncio.gather(*tasks, return_exceptions=True) + end_time = time.time() + + # Analyze results + total_successful = successful_requests["key_a"] + successful_requests["key_b"] + + key_a_success_rate = successful_requests["key_a"] / 200 + key_b_success_rate = successful_requests["key_b"] / 200 + + print(f"Test Case 2 - Priority Queue Behavior During Saturation:") + print(f" - Duration: {end_time - start_time:.2f}s") + print(f" - Key A: {successful_requests['key_a']}/200 successful ({key_a_success_rate:.1%})") + print(f" - Key B: {successful_requests['key_b']}/200 successful ({key_b_success_rate:.1%})") + print(f" - Total successful: {total_successful}/400") + + # Key A should get significantly more requests than Key B (75:25 ratio) + assert key_a_success_rate > key_b_success_rate, ( + f"Key A should have higher success rate: {key_a_success_rate:.1%} vs {key_b_success_rate:.1%}" + ) + + # Check ratio is approximately 3:1 (75:25) + if total_successful > 0: + key_a_share = successful_requests["key_a"] / total_successful + expected_key_a_share = 0.75 + + print(f" - Key A got {key_a_share:.1%} of successful requests (expected ~75%)") + + # Allow tolerance for timing effects + assert abs(key_a_share - expected_key_a_share) < 0.2, ( + f"Key A share should be ~75%, got {key_a_share:.1%}" + ) + + +@pytest.mark.asyncio +async def test_fake_calls_case_3_spillover_capacity_default_keys(): + """ + Test Case 3: Spillover Capacity for Default Keys + + System: 100 RPM capacity + Key A: priority_reservation=0.75 (75 RPM reserved) + Key B: nothing set (default) + Key C: nothing set (default) + Key D: nothing set (default) + Traffic A: 150 RPM + Traffic B: 150 RPM + Traffic C: 150 RPM + Traffic D: 150 RPM + Expected A: 75 RPM (75% reserved) + Expected B: ~8.3 RPM (remaining 25 RPM / 3 default keys) + Expected C: ~8.3 RPM + Expected D: ~8.3 RPM + + Tests spillover behavior where default keys share remaining capacity. + """ + os.environ["LITELLM_LICENSE"] = "test-license-key" + + litellm.priority_reservation = {"key_a": 0.75} + litellm.priority_reservation_settings.default_priority = 0.25 + + dual_cache = DualCache() + handler = DynamicRateLimitHandler(internal_usage_cache=dual_cache) + + model = "fake-call-test-3" + total_rpm = 100 + + llm_router = Router( + model_list=[ + { + "model_name": model, + "litellm_params": { + "model": "gpt-3.5-turbo", + "api_key": "test-key", + "api_base": "test-base", + "rpm": total_rpm, + }, + } + ] + ) + handler.update_variables(llm_router=llm_router) + + # Create users + key_a_user = UserAPIKeyAuth() + key_a_user.metadata = {"priority": "key_a"} + key_a_user.user_id = "key_a_user" + + key_b_user = UserAPIKeyAuth() + key_b_user.metadata = {} + key_b_user.user_id = "key_b_user" + + key_c_user = UserAPIKeyAuth() + key_c_user.metadata = {} + key_c_user.user_id = "key_c_user" + + key_d_user = UserAPIKeyAuth() + key_d_user.metadata = {} + key_d_user.user_id = "key_d_user" + + # Track results + successful_requests = {"key_a": 0, "key_b": 0, "key_c": 0, "key_d": 0} + rate_limited_requests = {"key_a": 0, "key_b": 0, "key_c": 0, "key_d": 0} + + async def make_request(user, key_name, request_id): + """Make a single request and track the result.""" + try: + result = await handler.async_pre_call_hook( + user_api_key_dict=user, + cache=dual_cache, + data={"model": model}, + call_type="completion", + ) + + if result is None: + successful_requests[key_name] += 1 + return {"status": "success", "key": key_name} + else: + rate_limited_requests[key_name] += 1 + return {"status": "rate_limited", "key": key_name} + + except Exception as e: + rate_limited_requests[key_name] += 1 + return {"status": "rate_limited", "key": key_name, "error": str(e)} + + # Send 150 requests from each key (600 total, 6x over capacity) + tasks = [] + + for i in range(150): + tasks.append(make_request(key_a_user, "key_a", f"key_a_{i}")) + + for i in range(150): + tasks.append(make_request(key_b_user, "key_b", f"key_b_{i}")) + + for i in range(150): + tasks.append(make_request(key_c_user, "key_c", f"key_c_{i}")) + + for i in range(150): + tasks.append(make_request(key_d_user, "key_d", f"key_d_{i}")) + + start_time = time.time() + results = await asyncio.gather(*tasks, return_exceptions=True) + end_time = time.time() + + # Analyze results + total_successful = sum(successful_requests.values()) + + print(f"Test Case 3 - Spillover Capacity for Default Keys:") + print(f" - Duration: {end_time - start_time:.2f}s") + print(f" - Key A: {successful_requests['key_a']}/150 successful") + print(f" - Key B: {successful_requests['key_b']}/150 successful (default)") + print(f" - Key C: {successful_requests['key_c']}/150 successful (default)") + print(f" - Key D: {successful_requests['key_d']}/150 successful (default)") + print(f" - Total successful: {total_successful}/600") + + # Key A should get the most requests (75% of capacity) + assert successful_requests["key_a"] > successful_requests["key_b"], "Key A should get more than Key B" + assert successful_requests["key_a"] > successful_requests["key_c"], "Key A should get more than Key C" + assert successful_requests["key_a"] > successful_requests["key_d"], "Key A should get more than Key D" + + # Default keys should get similar amounts (spillover capacity) + avg_default = (successful_requests["key_b"] + successful_requests["key_c"] + successful_requests["key_d"]) / 3 + print(f" - Average default key success: {avg_default:.1f}") + + +@pytest.mark.asyncio +async def test_fake_calls_case_4_over_allocated_with_normalization(): + """ + Test Case 4: Over-Allocated Priority reservations with Normalization + + System: 100 RPM capacity + Key A: priority_reservation=0.60 (60% requested) + Key B: priority_reservation=0.80 (80% requested) + Total: 140% (over-allocated, should normalize to 43%/57%) + Traffic A: 200 RPM + Traffic B: 200 RPM + + With saturation-aware rate limiting: + - Initially, requests are allowed through in generous mode (under 80% saturation) + - Once saturated, strict priority-based limits kick in with normalized weights + - Due to concurrent burst, total successful may exceed 100 RPM in the test window + - This test verifies normalization works and total capacity is reasonably bounded + """ + os.environ["LITELLM_LICENSE"] = "test-license-key" + + litellm.priority_reservation = {"key_a": 0.60, "key_b": 0.80} + + dual_cache = DualCache() + handler = DynamicRateLimitHandler(internal_usage_cache=dual_cache) + + model = "fake-call-test-4" + total_rpm = 100 + + llm_router = Router( + model_list=[ + { + "model_name": model, + "litellm_params": { + "model": "gpt-3.5-turbo", + "api_key": "test-key", + "api_base": "test-base", + "rpm": total_rpm, + }, + } + ] + ) + handler.update_variables(llm_router=llm_router) + + # Create users + key_a_user = UserAPIKeyAuth() + key_a_user.metadata = {"priority": "key_a"} + key_a_user.user_id = "key_a_user" + + key_b_user = UserAPIKeyAuth() + key_b_user.metadata = {"priority": "key_b"} + key_b_user.user_id = "key_b_user" + + # Track results + successful_requests = {"key_a": 0, "key_b": 0} + rate_limited_requests = {"key_a": 0, "key_b": 0} + + async def make_request(user, priority_name, request_id): + """Make a single request and track the result.""" + try: + result = await handler.async_pre_call_hook( + user_api_key_dict=user, + cache=dual_cache, + data={"model": model}, + call_type="completion", + ) + + if result is None: + successful_requests[priority_name] += 1 + return {"status": "success", "priority": priority_name} + else: + rate_limited_requests[priority_name] += 1 + return {"status": "rate_limited", "priority": priority_name} + + except Exception as e: + rate_limited_requests[priority_name] += 1 + return {"status": "rate_limited", "priority": priority_name, "error": str(e)} + + # Send 200 requests from each key (400 total, 4x over capacity) + tasks = [] + + for i in range(200): + tasks.append(make_request(key_a_user, "key_a", f"key_a_{i}")) + + for i in range(200): + tasks.append(make_request(key_b_user, "key_b", f"key_b_{i}")) + + start_time = time.time() + results = await asyncio.gather(*tasks, return_exceptions=True) + end_time = time.time() + + # Analyze results + total_successful = successful_requests["key_a"] + successful_requests["key_b"] + + key_a_success_rate = successful_requests["key_a"] / 200 + key_b_success_rate = successful_requests["key_b"] / 200 + + print(f"Test Case 4 - Over-Allocated Priority Reservations with Normalization:") + print(f" - Duration: {end_time - start_time:.2f}s") + print(f" - Key A (0.60): {successful_requests['key_a']}/200 successful ({key_a_success_rate:.1%})") + print(f" - Key B (0.80): {successful_requests['key_b']}/200 successful ({key_b_success_rate:.1%})") + print(f" - Total successful: {total_successful}/400") + + # With saturation-aware behavior: + # 1. Verify total capacity is reasonably bounded (not all 400 requests succeed) + assert total_successful < 300, ( + f"Total requests should be bounded by saturation detection, got {total_successful}/400" + ) + + # 2. Verify significant rate limiting occurred (at least 50% blocked) + assert total_successful < 200, ( + f"At least 50% of requests should be rate limited, got {total_successful}/400 successful" + ) + + # 3. Verify both keys got some requests through (normalization is working) + assert successful_requests["key_a"] > 0, "Key A should get some requests" + assert successful_requests["key_b"] > 0, "Key B should get some requests" + + print(f" - Normalization test PASSED: Both priorities got requests, " + f"total bounded to {total_successful} (under 200)") + + +@pytest.mark.asyncio +async def test_fake_calls_case_5_default_value_priority_reservation(): + """ + Test Case 5: Default value for priority reservation + + System: 100 RPM capacity + Key A: priority_reservation=0.50 (50 RPM) + Key B: priority_reservation=0.20 (20 RPM) + Key C: priority_reservation=0.05 (5 RPM) + Key D: nothing set (uses default_priority=0.05, 5 RPM) + Traffic A: 150 RPM + Traffic B: 150 RPM + Traffic C: 150 RPM + Traffic D: 150 RPM + Expected A: 55 RPM (normalized) + Expected B: 25 RPM (normalized) + Expected C: 10 RPM (normalized) + Expected D: 10 RPM (normalized) + + Tests complex scenario with explicit priorities and default priority. + """ + os.environ["LITELLM_LICENSE"] = "test-license-key" + + litellm.priority_reservation = {"key_a": 0.50, "key_b": 0.20, "key_c": 0.05} + litellm.priority_reservation_settings.default_priority = 0.05 + + dual_cache = DualCache() + handler = DynamicRateLimitHandler(internal_usage_cache=dual_cache) + + model = "fake-call-test-5" + total_rpm = 100 + + llm_router = Router( + model_list=[ + { + "model_name": model, + "litellm_params": { + "model": "gpt-3.5-turbo", + "api_key": "test-key", + "api_base": "test-base", + "rpm": total_rpm, + }, + } + ] + ) + handler.update_variables(llm_router=llm_router) + + # Create users + key_a_user = UserAPIKeyAuth() + key_a_user.metadata = {"priority": "key_a"} + key_a_user.user_id = "key_a_user" + + key_b_user = UserAPIKeyAuth() + key_b_user.metadata = {"priority": "key_b"} + key_b_user.user_id = "key_b_user" + + key_c_user = UserAPIKeyAuth() + key_c_user.metadata = {"priority": "key_c"} + key_c_user.user_id = "key_c_user" + + key_d_user = UserAPIKeyAuth() + key_d_user.metadata = {} + key_d_user.user_id = "key_d_user" + + # Track results + successful_requests = {"key_a": 0, "key_b": 0, "key_c": 0, "key_d": 0} + rate_limited_requests = {"key_a": 0, "key_b": 0, "key_c": 0, "key_d": 0} + + async def make_request(user, key_name, request_id): + """Make a single request and track the result.""" + try: + result = await handler.async_pre_call_hook( + user_api_key_dict=user, + cache=dual_cache, + data={"model": model}, + call_type="completion", + ) + + if result is None: + successful_requests[key_name] += 1 + return {"status": "success", "key": key_name} + else: + rate_limited_requests[key_name] += 1 + return {"status": "rate_limited", "key": key_name} + + except Exception as e: + rate_limited_requests[key_name] += 1 + return {"status": "rate_limited", "key": key_name, "error": str(e)} + + # Send 150 requests from each key (600 total, 6x over capacity) + tasks = [] + + for i in range(150): + tasks.append(make_request(key_a_user, "key_a", f"key_a_{i}")) + + for i in range(150): + tasks.append(make_request(key_b_user, "key_b", f"key_b_{i}")) + + for i in range(150): + tasks.append(make_request(key_c_user, "key_c", f"key_c_{i}")) + + for i in range(150): + tasks.append(make_request(key_d_user, "key_d", f"key_d_{i}")) + + start_time = time.time() + results = await asyncio.gather(*tasks, return_exceptions=True) + end_time = time.time() + + # Analyze results + total_successful = sum(successful_requests.values()) + + print(f"Test Case 5 - Default value for priority reservation:") + print(f" - Duration: {end_time - start_time:.2f}s") + print(f" - Key A (0.50): {successful_requests['key_a']}/150 successful") + print(f" - Key B (0.20): {successful_requests['key_b']}/150 successful") + print(f" - Key C (0.05): {successful_requests['key_c']}/150 successful") + print(f" - Key D (default 0.05): {successful_requests['key_d']}/150 successful") + print(f" - Total successful: {total_successful}/600") + + # Verify priority ordering: A > B > C ≈ D + assert successful_requests["key_a"] > successful_requests["key_b"], "Key A should get more than Key B" + assert successful_requests["key_b"] > successful_requests["key_c"], "Key B should get more than Key C" + + # Key C and Key D should get similar amounts (both have 0.05 priority) + key_c_vs_d_ratio = successful_requests["key_c"] / max(successful_requests["key_d"], 1) + print(f" - Key C vs Key D ratio: {key_c_vs_d_ratio:.2f} (expected ~1.0)") + + if total_successful > 0: + key_a_share = successful_requests["key_a"] / total_successful + print(f" - Key A got {key_a_share:.1%} of successful requests (expected ~55-62%)")