"""
Dynamic rate limiter v3 - Saturation-aware priority-based rate limiting
"""
import os
from datetime import datetime
from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Union
from fastapi import HTTPException
import litellm
from litellm import ModelResponse, Router
from litellm._logging import verbose_proxy_logger
from litellm.caching.caching import DualCache
from litellm.integrations.custom_logger import CustomLogger
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.hooks.parallel_request_limiter_v3 import (
RateLimitDescriptor,
RateLimitDescriptorRateLimitObject,
_PROXY_MaxParallelRequestsHandler_v3,
)
from litellm.proxy.hooks.rate_limiter_utils import convert_priority_to_percent
from litellm.proxy.utils import InternalUsageCache
from litellm.types.router import ModelGroupInfo
from litellm.types.utils import CallTypesLiteral
if TYPE_CHECKING:
from litellm.types.utils import PriorityReservationSettings
def _get_priority_settings() -> "PriorityReservationSettings":
"""
Get the priority reservation settings, guaranteed to be non-None.
The settings are lazy-loaded in litellm.__init__ and always return an instance.
This helper provides proper type narrowing for mypy.
"""
settings = litellm.priority_reservation_settings
if settings is None:
# This should never happen due to lazy loading, but satisfy mypy
from litellm.types.utils import PriorityReservationSettings
return PriorityReservationSettings()
return settings
class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
"""
Saturation-aware priority-based rate limiter using v3 infrastructure.
Key features:
1. Model capacity ALWAYS enforced at 100% (prevents over-allocation)
2. Priority usage tracked from first request (accurate accounting)
3. Priority limits only enforced when saturated >= threshold
4. Three-phase checking prevents partial counter increments
5. Reuses v3 limiter's Redis-based tracking (multi-instance safe)
How it works:
- Phase 1: Read-only check of ALL limits (no increments)
- Phase 2: Decide enforcement based on saturation
- Phase 3: Increment counters only if request allowed
- When under-saturated: priorities can borrow unused capacity (generous)
- When saturated: strict priority-based limits enforced (fair)
- Uses v3 limiter's atomic Lua scripts for race-free increments
"""
def __init__(
self,
internal_usage_cache: DualCache,
time_provider: Optional[Callable[[], datetime]] = None,
):
self.internal_usage_cache = InternalUsageCache(dual_cache=internal_usage_cache)
self.v3_limiter = _PROXY_MaxParallelRequestsHandler_v3(
self.internal_usage_cache, time_provider=time_provider
)
def update_variables(self, llm_router: Router):
self.llm_router = llm_router
def _get_saturation_check_cache_ttl(self) -> int:
"""Get the configurable TTL for local cache when reading saturation values."""
return _get_priority_settings().saturation_check_cache_ttl
async def _get_saturation_value_from_cache(
self,
counter_key: str,
) -> Optional[str]:
"""
Get saturation value with configurable local cache TTL.
Uses DualCache with configurable TTL for local cache storage.
TTL is configurable via litellm.priority_reservation_settings.saturation_check_cache_ttl
Args:
counter_key: The cache key for the saturation counter
Returns:
Counter value as string, or None if not found
"""
local_cache_ttl = self._get_saturation_check_cache_ttl()
return await self.internal_usage_cache.async_get_cache(
key=counter_key,
litellm_parent_otel_span=None,
local_only=False,
ttl=local_cache_ttl,
)
def _get_priority_weight(
self, priority: Optional[str], model_info: Optional[ModelGroupInfo] = None
) -> float:
"""Get the weight for a given priority from litellm.priority_reservation"""
weight: float = _get_priority_settings().default_priority
if (
litellm.priority_reservation is None
or priority not in litellm.priority_reservation
):
verbose_proxy_logger.debug(
"Priority Reservation not set for the given priority."
)
elif priority is not None and litellm.priority_reservation is not None:
if os.getenv("LITELLM_LICENSE", None) is None:
verbose_proxy_logger.error(
"PREMIUM FEATURE: Reserving tpm/rpm by priority is a premium feature. Please add a 'LITELLM_LICENSE' to your .env to enable this.\nGet a license: https://docs.litellm.ai/docs/proxy/enterprise."
)
else:
value = litellm.priority_reservation[priority]
weight = convert_priority_to_percent(value, model_info)
return weight
def _get_priority_from_user_api_key_dict(
self, user_api_key_dict: UserAPIKeyAuth
) -> Optional[str]:
"""
Get priority from user_api_key_dict.
Checks team metadata first (takes precedence), then falls back to key metadata.
Args:
user_api_key_dict: User authentication info
Returns:
Priority string if found, None otherwise
"""
priority: Optional[str] = None
# Check team metadata first (takes precedence)
if user_api_key_dict.team_metadata is not None:
priority = user_api_key_dict.team_metadata.get("priority", None)
# Fall back to key metadata
if priority is None:
priority = user_api_key_dict.metadata.get("priority", None)
return priority
def _normalize_priority_weights(
self, model_info: ModelGroupInfo
) -> 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}
Converts absolute rpm/tpm values to percentages based on model capacity.
"""
if litellm.priority_reservation is None:
return {}
# Convert all values to percentages first
weights: Dict[str, float] = {}
for k, v in litellm.priority_reservation.items():
weights[k] = convert_priority_to_percent(v, model_info)
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
def _get_priority_allocation(
self,
model: str,
priority: Optional[str],
normalized_weights: Dict[str, float],
model_info: Optional[ModelGroupInfo] = None,
) -> tuple[float, str]:
"""
Get priority weight and pool key for a given priority.
For explicit priorities: returns specific allocation and unique pool key
For default priority: returns default allocation and shared pool key
Args:
model: Model name
priority: Priority level (None for default)
normalized_weights: Pre-computed normalized weights
model_info: Model configuration (optional, for fallback conversion)
Returns:
tuple: (priority_weight, priority_key)
"""
# Check if this key has an explicit priority in litellm.priority_reservation
has_explicit_priority = (
priority is not None
and litellm.priority_reservation is not None
and priority in litellm.priority_reservation
)
if has_explicit_priority and priority is not None:
# Explicit priority: get its specific allocation
priority_weight = normalized_weights.get(
priority, self._get_priority_weight(priority, model_info)
)
# Use unique key per priority level
priority_key = f"{model}:{priority}"
else:
# No explicit priority: share the default_priority pool with ALL other default keys
priority_weight = _get_priority_settings().default_priority
# Use shared key for all default-priority requests
priority_key = f"{model}:default_pool"
return priority_weight, priority_key
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 - always read from Redis for multi-node consistency
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 Redis directly for current counter value (skip local cache for consistency)
counter_value = await self._get_saturation_value_from_cache(
counter_key=counter_key
)
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._get_saturation_value_from_cache(
counter_key=counter_key
)
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,
user_api_key_dict: UserAPIKeyAuth,
priority: Optional[str],
) -> List[RateLimitDescriptor]:
"""
Create rate limit descriptors with normalized priority weights.
Uses normalized weights to handle over-allocation scenarios.
For explicit priorities: each priority gets its own pool (e.g., prod gets 75%)
For default priority: ALL keys without explicit priority share ONE pool (e.g., all share 25%)
"""
descriptors: List[RateLimitDescriptor] = []
if litellm.priority_reservation is None:
return descriptors
# Get model group info
model_group_info: Optional[
ModelGroupInfo
] = self.llm_router.get_model_group_info(model_group=model)
if model_group_info is None:
return descriptors
# Get normalized priority weight and pool key
normalized_weights = self._normalize_priority_weights(model_group_info)
priority_weight, priority_key = self._get_priority_allocation(
model=model,
priority=priority,
normalized_weights=normalized_weights,
model_info=model_group_info,
)
rate_limit_config: RateLimitDescriptorRateLimitObject = {}
# Apply priority weight to model limits
if model_group_info.tpm is not None:
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:
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
descriptors.append(
RateLimitDescriptor(
key="priority_model",
value=priority_key,
rate_limit=rate_limit_config,
)
)
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 _check_rate_limits(
self,
model: str,
model_group_info: ModelGroupInfo,
user_api_key_dict: UserAPIKeyAuth,
priority: Optional[str],
saturation: float,
data: dict,
) -> None:
"""
Check rate limits using THREE-PHASE approach to prevent partial increments.
Phase 1: Read-only check of ALL limits (no increments)
Phase 2: Decide which limits to enforce based on saturation
Phase 3: Increment ALL counters atomically (model + priority)
This prevents the bug where:
- Model counter increments in stage 1
- Priority check fails in stage 2
- Request blocked but model counter already incremented
Key behaviors:
- All checks performed first (read-only)
- Only increment counters if request will be allowed
- Model capacity: Always enforced at 100%
- Priority limits: Only enforced when saturated >= threshold
- Both counters tracked from first request (accurate accounting)
Args:
model: Model name
model_group_info: Model configuration
user_api_key_dict: User authentication info
priority: User's priority level
saturation: Current saturation level
data: Request data dictionary
Raises:
HTTPException: If any limit is exceeded
"""
import json
saturation_threshold = _get_priority_settings().saturation_threshold
should_enforce_priority = saturation >= saturation_threshold
# Build ALL descriptors upfront
descriptors_to_check: List[RateLimitDescriptor] = []
# Model-wide descriptor (always enforce)
model_wide_descriptor = self._create_model_tracking_descriptor(
model=model,
model_group_info=model_group_info,
high_limit_multiplier=1,
)
descriptors_to_check.append(model_wide_descriptor)
# Priority descriptors (always track, conditionally enforce)
priority_descriptors = self._create_priority_based_descriptors(
model=model,
user_api_key_dict=user_api_key_dict,
priority=priority,
)
if priority_descriptors:
descriptors_to_check.extend(priority_descriptors)
# PHASE 1: Read-only check of ALL limits (no increments)
check_response = await self.v3_limiter.should_rate_limit(
descriptors=descriptors_to_check,
parent_otel_span=user_api_key_dict.parent_otel_span,
read_only=True, # CRITICAL: Don't increment counters yet
)
verbose_proxy_logger.debug(
f"Read-only check: {json.dumps(check_response, indent=2)}"
)
# PHASE 2: Decide which limits to enforce
if check_response["overall_code"] == "OVER_LIMIT":
for status in check_response["statuses"]:
if status["code"] == "OVER_LIMIT":
descriptor_key = status["descriptor_key"]
# Model-wide limit exceeded (ALWAYS enforce)
if descriptor_key == "model_saturation_check":
raise HTTPException(
status_code=429,
detail={
"error": f"Model capacity reached for {model}. "
f"Priority: {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": priority or "default",
},
)
# Priority limit exceeded (ONLY enforce when saturated)
elif descriptor_key == "priority_model" and should_enforce_priority:
verbose_proxy_logger.debug(
f"Enforcing priority limits for {model}, saturation: {saturation:.1%}, "
f"priority: {priority}"
)
raise HTTPException(
status_code=429,
detail={
"error": f"Priority-based rate limit exceeded. "
f"Priority: {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": priority or "default",
"x-litellm-saturation": f"{saturation:.2%}",
},
)
# PHASE 3: Increment counters separately to avoid early-exit issues
# Model counter must ALWAYS increment, but priority counter might be over limit
# If we increment them together, v3_limiter's in-memory check will exit early
# and skip incrementing the model counter
# Step 3a: Increment model-wide counter (always)
model_increment_response = await self.v3_limiter.should_rate_limit(
descriptors=[model_wide_descriptor],
parent_otel_span=user_api_key_dict.parent_otel_span,
read_only=False,
)
# Step 3b: Increment priority counter (may be over limit, but we still track it)
if priority_descriptors:
priority_increment_response = await self.v3_limiter.should_rate_limit(
descriptors=priority_descriptors,
parent_otel_span=user_api_key_dict.parent_otel_span,
read_only=False,
)
# Combine responses for post-call hook
combined_response = {
"overall_code": model_increment_response["overall_code"],
"statuses": model_increment_response["statuses"]
+ priority_increment_response["statuses"],
}
data["litellm_proxy_rate_limit_response"] = combined_response
else:
data["litellm_proxy_rate_limit_response"] = model_increment_response
async def async_pre_call_hook(
self,
user_api_key_dict: UserAPIKeyAuth,
cache: DualCache,
data: dict,
call_type: CallTypesLiteral,
) -> Optional[Union[Exception, str, dict]]:
"""
Saturation-aware pre-call hook for priority-based rate limiting.
Flow:
1. Check current saturation level
2. THREE-PHASE rate limit check:
- PHASE 1: Read-only check of ALL limits (no increments)
- PHASE 2: Decide which limits to enforce based on saturation
- PHASE 3: Increment ALL counters atomically if request allowed
This three-phase approach ensures:
- Model capacity is NEVER exceeded (always enforced at 100%)
- Priority usage tracked from first request (accurate metrics)
- Counters only increment when request will be allowed (prevents phantom usage)
- When under-saturated: priorities can borrow unused capacity (generous)
- When saturated: fair allocation based on normalized priority weights (strict)
Example with 100 RPM model, 60% priority allocation, 80% threshold:
- Saturation < 80%: Priority can use up to 100 RPM (model limit enforced only)
- Saturation >= 80%: Priority limited to 60 RPM (both limits enforced)
Prevents bugs where:
- Model counter increments but priority check fails → model over-capacity
- Priority counter increments but not enforced → inaccurate metrics
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"]
priority = self._get_priority_from_user_api_key_dict(
user_api_key_dict=user_api_key_dict
)
# Get model configuration
model_group_info: Optional[
ModelGroupInfo
] = self.llm_router.get_model_group_info(model_group=model)
if model_group_info is None:
verbose_proxy_logger.debug(
f"No model group info for {model}, allowing request"
)
return None
try:
# STEP 1: Check current saturation level
saturation = await self._check_model_saturation(model, model_group_info)
saturation_threshold = _get_priority_settings().saturation_threshold
verbose_proxy_logger.debug(
f"[Dynamic Rate Limiter] Model={model}, Saturation={saturation:.1%}, "
f"Threshold={saturation_threshold:.1%}, Priority={priority}"
)
# STEP 2: Check rate limits in THREE phases
# Phase 1: Read-only check of ALL limits (no increments)
# Phase 2: Decide which limits to enforce (based on saturation)
# Phase 3: Increment ALL counters only if request will be allowed
# This prevents partial increments and ensures accurate tracking
await self._check_rate_limits(
model=model,
model_group_info=model_group_info,
user_api_key_dict=user_api_key_dict,
priority=priority,
saturation=saturation,
data=data,
)
except HTTPException:
raise
except Exception as e:
verbose_proxy_logger.error(
f"Error in dynamic rate limiter: {str(e)}, allowing request"
)
# Fail open on unexpected errors
return None
return None
async def async_post_call_success_hook(
self, data: dict, user_api_key_dict: UserAPIKeyAuth, response
):
"""
Post-call hook to add rate limit headers to response.
Leverages v3 limiter's post-call hook functionality.
"""
try:
# Call v3 limiter's post-call hook to add standard rate limit headers
await self.v3_limiter.async_post_call_success_hook(
data=data, user_api_key_dict=user_api_key_dict, response=response
)
# Add additional priority-specific headers
if isinstance(response, ModelResponse):
priority = self._get_priority_from_user_api_key_dict(
user_api_key_dict=user_api_key_dict
)
# Get existing additional headers
additional_headers = (
getattr(response, "_hidden_params", {}).get(
"additional_headers", {}
)
or {}
)
# Add priority information
additional_headers["x-litellm-priority"] = priority or "default"
additional_headers["x-litellm-rate-limiter-version"] = "v3"
# Update response
if not hasattr(response, "_hidden_params"):
response._hidden_params = {}
response._hidden_params["additional_headers"] = additional_headers
return response
except Exception as e:
verbose_proxy_logger.exception(
f"Error in dynamic rate limiter v3 post-call hook: {str(e)}"
)
return response
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
"""
Update token usage for priority-based rate limiting after successful API calls.
Increments token counters for:
- model_saturation_check: Model-wide token tracking
- priority_model: Priority-specific token tracking
"""
from litellm.litellm_core_utils.core_helpers import (
_get_parent_otel_span_from_kwargs,
)
from litellm.proxy.common_utils.callback_utils import (
get_model_group_from_litellm_kwargs,
)
from litellm.types.caching import RedisPipelineIncrementOperation
from litellm.types.utils import Usage
try:
verbose_proxy_logger.debug(
"INSIDE dynamic rate limiter ASYNC SUCCESS LOGGING"
)
litellm_parent_otel_span = _get_parent_otel_span_from_kwargs(kwargs)
# Get metadata from standard_logging_object
standard_logging_object = kwargs.get("standard_logging_object") or {}
standard_logging_metadata = standard_logging_object.get("metadata") or {}
# Get model and priority
model_group = get_model_group_from_litellm_kwargs(kwargs)
if not model_group:
return
# Get priority from user_api_key_auth_metadata in standard_logging_metadata
# This is where user_api_key_dict.metadata is stored during pre-call
user_api_key_auth_metadata = (
standard_logging_metadata.get("user_api_key_auth_metadata") or {}
)
key_priority: Optional[str] = user_api_key_auth_metadata.get("priority")
# Get total tokens from response
total_tokens = 0
rate_limit_type = self.v3_limiter.get_rate_limit_type()
if isinstance(response_obj, ModelResponse):
_usage = getattr(response_obj, "usage", None)
if _usage and isinstance(_usage, Usage):
if rate_limit_type == "output":
total_tokens = _usage.completion_tokens
elif rate_limit_type == "input":
total_tokens = _usage.prompt_tokens
elif rate_limit_type == "total":
total_tokens = _usage.total_tokens
if total_tokens == 0:
return
# Create pipeline operations for token increments
pipeline_operations: List[RedisPipelineIncrementOperation] = []
# Model-wide token tracking (model_saturation_check)
model_token_key = self.v3_limiter.create_rate_limit_keys(
key="model_saturation_check",
value=model_group,
rate_limit_type="tokens",
)
pipeline_operations.append(
RedisPipelineIncrementOperation(
key=model_token_key,
increment_value=total_tokens,
ttl=self.v3_limiter.window_size,
)
)
# Priority-specific token tracking (priority_model)
# Determine priority key (same logic as _get_priority_allocation)
has_explicit_priority = (
key_priority is not None
and litellm.priority_reservation is not None
and key_priority in litellm.priority_reservation
)
if has_explicit_priority and key_priority is not None:
priority_key = f"{model_group}:{key_priority}"
else:
priority_key = f"{model_group}:default_pool"
priority_token_key = self.v3_limiter.create_rate_limit_keys(
key="priority_model",
value=priority_key,
rate_limit_type="tokens",
)
pipeline_operations.append(
RedisPipelineIncrementOperation(
key=priority_token_key,
increment_value=total_tokens,
ttl=self.v3_limiter.window_size,
)
)
# Execute token increments with TTL preservation
if pipeline_operations:
await self.v3_limiter.async_increment_tokens_with_ttl_preservation(
pipeline_operations=pipeline_operations,
parent_otel_span=litellm_parent_otel_span,
)
# Only log 'priority' if it's known safe; otherwise, redact.
SAFE_PRIORITIES = {"low", "medium", "high", "default"}
logged_priority = (
key_priority if key_priority in SAFE_PRIORITIES else "REDACTED"
)
verbose_proxy_logger.debug(
f"[Dynamic Rate Limiter] Incremented tokens by {total_tokens} for "
f"model={model_group}, priority={logged_priority}"
)
except Exception as e:
verbose_proxy_logger.exception(
f"Error in dynamic rate limiter success event: {str(e)}"
)