[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
This commit is contained in:
Ishaan Jaff
2025-10-01 18:35:34 -07:00
committed by GitHub
parent 2bc5d93f23
commit d538cf489a
4 changed files with 1184 additions and 62 deletions
@@ -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
+344 -54
View File
@@ -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
+10
View File
@@ -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=())
@@ -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%)")