mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-22 18:25:39 +00:00
[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:
@@ -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
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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%)")
|
||||
|
||||
Reference in New Issue
Block a user