mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-23 18:25:22 +00:00
[Fixes] Dynamic Rate Limiter - Dynamic rate limiting token count increases/decreases by 1 instead of actual count + Redis TTL (#17558)
* fix async_log_success_event for _PROXY_DynamicRateLimitHandlerV3 * test_async_log_success_event_increments_by_actual_tokens * fix redis TTL * Potential fix for code scanning alert no. 3873: Clear-text logging of sensitive information Co-authored-by: Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com> --------- Co-authored-by: Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com>
This commit is contained in:
Ishaan Jaff
co-authored by
Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com>
Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com>
parent
4d39a1a18f
commit
a78f40f75a
@@ -614,3 +614,121 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
|
||||
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)}"
|
||||
)
|
||||
|
||||
@@ -65,6 +65,12 @@ for i = 1, #KEYS, 2 do
|
||||
table.insert(results, increment_value) -- counter
|
||||
else
|
||||
local counter = redis.call('INCR', counter_key)
|
||||
-- This happens when window_key exists but counter_key doesn't (e.g., tokens key
|
||||
-- created after requests key when both share the same window_key)
|
||||
local current_ttl = redis.call('TTL', counter_key)
|
||||
if current_ttl == -1 then
|
||||
redis.call('EXPIRE', counter_key, window_size)
|
||||
end
|
||||
table.insert(results, window_start) -- window_start
|
||||
table.insert(results, counter) -- counter
|
||||
end
|
||||
|
||||
@@ -1,67 +1,16 @@
|
||||
model_list:
|
||||
- model_name: qwen-25vl-72b
|
||||
- model_name: openai/gpt-4o-mini
|
||||
litellm_params:
|
||||
model: bedrock/openai/arn:aws:bedrock:us-east-1:046319184608:imported-model/0m2lasirsp6z
|
||||
model: openai/gpt-4o-mini
|
||||
tpm: 1000
|
||||
|
||||
guardrails:
|
||||
- guardrail_name: "bedrock-pre-guard"
|
||||
litellm_params:
|
||||
guardrail: bedrock
|
||||
mode: "pre_call"
|
||||
guardrailIdentifier: ff6ujrregl1q
|
||||
guardrailVersion: "DRAFT"
|
||||
|
||||
|
||||
# like MCPs/vector stores
|
||||
search_tools:
|
||||
- search_tool_name: litellm-search
|
||||
litellm_params:
|
||||
search_provider: perplexity
|
||||
api_key: os.environ/PERPLEXITYAI_API_KEY
|
||||
- search_tool_name: firecrawl-search
|
||||
litellm_params:
|
||||
search_provider: firecrawl
|
||||
api_key: os.environ/FIRECRAWL_API_KEY
|
||||
|
||||
|
||||
litellm_settings:
|
||||
max_end_user_budget_id: "2f6634cd-c631-4d3b-96c7-ad510ea06eaf"
|
||||
# Comprehensive logging settings
|
||||
store_audit_logs: true
|
||||
verbose: true
|
||||
log_level: "DEBUG" # Options: DEBUG, INFO, WARNING, ERROR
|
||||
callbacks: ["s3_v2", "smtp_email"]
|
||||
s3_callback_params:
|
||||
s3_endpoint_url: "https://localhost:443" # Replace with your Minio server URL and port
|
||||
s3_aws_access_key_id: "minioadmin"
|
||||
s3_aws_secret_access_key: "minioadmin"
|
||||
s3_region_name: "minio" # This can be any value for Minio
|
||||
s3_bucket_name: "litellm-test" # Replace with your bucket name
|
||||
s3_use_ssl: False
|
||||
s3_verify: False
|
||||
cache: True
|
||||
cache_params:
|
||||
type: local
|
||||
drop_params: True
|
||||
callbacks: ["dynamic_rate_limiter_v3"]
|
||||
priority_reservation:
|
||||
"prod": 0.9 # 90% reserved for production
|
||||
"dev": 0.1 # 10% reserved for development
|
||||
|
||||
|
||||
general_settings:
|
||||
store_prompts_in_spend_logs: True
|
||||
pass_through_endpoints:
|
||||
- path: "/special/rerank"
|
||||
target: "https://api.cohere.com/v1/rerank"
|
||||
headers:
|
||||
Authorization: "Bearer os.environ/COHERE_API_KEY"
|
||||
guardrails:
|
||||
bedrock-pre-guard:
|
||||
request_fields: ["documents[*].text"]
|
||||
|
||||
|
||||
vector_store_registry:
|
||||
- vector_store_name: "bedrock-litellm-website-knowledgebase"
|
||||
litellm_params:
|
||||
vector_store_id: "T37J8R4WTM"
|
||||
custom_llm_provider: "bedrock"
|
||||
vector_store_description: "Bedrock vector store for the Litellm website knowledgebase"
|
||||
vector_store_metadata:
|
||||
source: "https://www.litellm.com/docs"
|
||||
|
||||
@@ -1323,3 +1323,187 @@ async def test_default_priority_shared_pool():
|
||||
print(f" - 3 keys without priority share ONE pool: {desc_a[0]['value']}")
|
||||
print(f" - Shared pool limit: {desc_a[0]['rate_limit']['requests_per_unit']} RPM")
|
||||
print(f" - Explicit priority 'prod' uses separate pool: {desc_prod[0]['value']}")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_log_success_event_increments_by_actual_tokens():
|
||||
"""
|
||||
Test that async_log_success_event increments token counters by actual token usage.
|
||||
|
||||
This validates the fix for Bug 1: Token count was incrementing by 1 instead of actual usage.
|
||||
The async_log_success_event should increment both model_saturation_check and priority_model
|
||||
counters by the actual completion_tokens (when rate_limit_type=output).
|
||||
"""
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from litellm.types.utils import ModelResponse, Usage
|
||||
|
||||
os.environ["LITELLM_LICENSE"] = "test-license-key"
|
||||
litellm.priority_reservation = {"dev": 0.1, "prod": 0.9}
|
||||
|
||||
dual_cache = DualCache()
|
||||
handler = DynamicRateLimitHandler(internal_usage_cache=dual_cache)
|
||||
|
||||
model = "test-token-increment"
|
||||
llm_router = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": model,
|
||||
"litellm_params": {
|
||||
"model": "gpt-3.5-turbo",
|
||||
"api_key": "test-key",
|
||||
"api_base": "test-base",
|
||||
"tpm": 1000,
|
||||
},
|
||||
}
|
||||
]
|
||||
)
|
||||
handler.update_variables(llm_router=llm_router)
|
||||
|
||||
# Track what gets incremented
|
||||
increment_calls = []
|
||||
|
||||
async def mock_increment(pipeline_operations, parent_otel_span=None):
|
||||
for op in pipeline_operations:
|
||||
increment_calls.append({
|
||||
"key": op["key"],
|
||||
"increment_value": op["increment_value"],
|
||||
})
|
||||
|
||||
handler.v3_limiter.async_increment_tokens_with_ttl_preservation = mock_increment
|
||||
|
||||
# Create mock response with 50 completion tokens
|
||||
mock_response = MagicMock(spec=ModelResponse)
|
||||
mock_response.usage = MagicMock(spec=Usage)
|
||||
mock_response.usage.prompt_tokens = 10
|
||||
mock_response.usage.completion_tokens = 50
|
||||
mock_response.usage.total_tokens = 60
|
||||
|
||||
# Create kwargs with priority in user_api_key_auth_metadata
|
||||
kwargs = {
|
||||
"standard_logging_object": {
|
||||
"metadata": {
|
||||
"user_api_key_auth_metadata": {"priority": "dev"},
|
||||
},
|
||||
"model_group": model,
|
||||
},
|
||||
"litellm_params": {
|
||||
"metadata": {"model_group": model},
|
||||
},
|
||||
}
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.common_utils.callback_utils.get_model_group_from_litellm_kwargs",
|
||||
return_value=model,
|
||||
):
|
||||
await handler.async_log_success_event(
|
||||
kwargs=kwargs,
|
||||
response_obj=mock_response,
|
||||
start_time=None,
|
||||
end_time=None,
|
||||
)
|
||||
|
||||
# Verify increments happened with actual token count (50 completion tokens)
|
||||
assert len(increment_calls) == 2, f"Expected 2 increment calls, got {len(increment_calls)}"
|
||||
|
||||
# Both should increment by 50 (completion_tokens, since rate_limit_type defaults to 'output')
|
||||
for call in increment_calls:
|
||||
assert call["increment_value"] == 50, (
|
||||
f"Expected increment of 50 tokens, got {call['increment_value']} for key {call['key']}"
|
||||
)
|
||||
|
||||
# Verify correct keys were used
|
||||
keys = [call["key"] for call in increment_calls]
|
||||
assert any("model_saturation_check" in k for k in keys), "Should increment model_saturation_check"
|
||||
assert any("priority_model" in k and "dev" in k for k in keys), "Should increment priority_model with 'dev' priority"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_log_success_event_uses_team_priority_from_auth_metadata():
|
||||
"""
|
||||
Test that async_log_success_event correctly retrieves priority from user_api_key_auth_metadata.
|
||||
|
||||
This validates the fix where priority is retrieved from standard_logging_metadata.user_api_key_auth_metadata
|
||||
instead of just standard_logging_metadata.priority. This is important for team-based priority inheritance.
|
||||
"""
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from litellm.types.utils import ModelResponse, Usage
|
||||
|
||||
os.environ["LITELLM_LICENSE"] = "test-license-key"
|
||||
litellm.priority_reservation = {"team_priority": 0.8, "default": 0.2}
|
||||
|
||||
dual_cache = DualCache()
|
||||
handler = DynamicRateLimitHandler(internal_usage_cache=dual_cache)
|
||||
|
||||
model = "test-team-priority"
|
||||
llm_router = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": model,
|
||||
"litellm_params": {
|
||||
"model": "gpt-3.5-turbo",
|
||||
"api_key": "test-key",
|
||||
"api_base": "test-base",
|
||||
"tpm": 1000,
|
||||
},
|
||||
}
|
||||
]
|
||||
)
|
||||
handler.update_variables(llm_router=llm_router)
|
||||
|
||||
# Track incremented keys to verify priority is used correctly
|
||||
incremented_keys = []
|
||||
|
||||
async def mock_increment(pipeline_operations, parent_otel_span=None):
|
||||
for op in pipeline_operations:
|
||||
incremented_keys.append(op["key"])
|
||||
|
||||
handler.v3_limiter.async_increment_tokens_with_ttl_preservation = mock_increment
|
||||
|
||||
# Create mock response
|
||||
mock_response = MagicMock(spec=ModelResponse)
|
||||
mock_response.usage = MagicMock(spec=Usage)
|
||||
mock_response.usage.prompt_tokens = 10
|
||||
mock_response.usage.completion_tokens = 20
|
||||
mock_response.usage.total_tokens = 30
|
||||
|
||||
# Simulate team metadata inheritance: priority is in user_api_key_auth_metadata
|
||||
# This is how the proxy passes team metadata to the callback
|
||||
kwargs = {
|
||||
"standard_logging_object": {
|
||||
"metadata": {
|
||||
# Priority NOT at top level (this would fail before the fix)
|
||||
# Priority IS in user_api_key_auth_metadata (team inheritance)
|
||||
"user_api_key_auth_metadata": {"priority": "team_priority"},
|
||||
},
|
||||
"model_group": model,
|
||||
},
|
||||
"litellm_params": {
|
||||
"metadata": {"model_group": model},
|
||||
},
|
||||
}
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.common_utils.callback_utils.get_model_group_from_litellm_kwargs",
|
||||
return_value=model,
|
||||
):
|
||||
await handler.async_log_success_event(
|
||||
kwargs=kwargs,
|
||||
response_obj=mock_response,
|
||||
start_time=None,
|
||||
end_time=None,
|
||||
)
|
||||
|
||||
# Verify the priority_model key uses 'team_priority' (not 'default_pool')
|
||||
priority_keys = [k for k in incremented_keys if "priority_model" in k]
|
||||
assert len(priority_keys) == 1, f"Expected 1 priority_model key, got {len(priority_keys)}"
|
||||
|
||||
# The key should contain 'team_priority', not 'default_pool'
|
||||
assert "team_priority" in priority_keys[0], (
|
||||
f"Expected priority key to use 'team_priority' from user_api_key_auth_metadata, "
|
||||
f"got key: {priority_keys[0]}"
|
||||
)
|
||||
assert "default_pool" not in priority_keys[0], (
|
||||
f"Priority key should NOT use 'default_pool', should use team's priority. Got: {priority_keys[0]}"
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user