From 90ee9e45871aa05afe80c7f7ea8ea4833745be3b Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 19 Sep 2025 16:04:45 -0700 Subject: [PATCH] [Feat] Dynamic Rate Limiter v3 - fixes to ensure priority routing works as expected (#14734) * fix: dynamic limiter v3 * fix: dynamic limiter v3 * feat: add dynamic limiter v3 * feat: add dynamic limiter v3 * feat: add dynamic limiter v3 in init litellm_logging * feat: add dynamic limiter v3 in init litellm_logging * fix: priority rate limiting * Potential fix for code scanning alert no. 3397: Clear-text logging of sensitive information Co-authored-by: Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com> * fix: priority rate limiting * fix: ruff * fix: mypy lint --------- Co-authored-by: Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com> --- litellm/__init__.py | 1 + .../custom_logger_registry.py | 2 + litellm/litellm_core_utils/litellm_logging.py | 32 ++ .../proxy/hooks/dynamic_rate_limiter_v3.py | 226 ++++++++ .../hooks/test_dynamic_rate_limiter_v3.py | 482 ++++++++++++++++++ 5 files changed, 743 insertions(+) create mode 100644 litellm/proxy/hooks/dynamic_rate_limiter_v3.py create mode 100644 tests/test_litellm/proxy/hooks/test_dynamic_rate_limiter_v3.py diff --git a/litellm/__init__.py b/litellm/__init__.py index 3c1d6e0696..038787f5cc 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -117,6 +117,7 @@ _custom_logger_compatible_callbacks_literal = Literal[ "logfire", "literalai", "dynamic_rate_limiter", + "dynamic_rate_limiter_v3", "langsmith", "prometheus", "otel", diff --git a/litellm/litellm_core_utils/custom_logger_registry.py b/litellm/litellm_core_utils/custom_logger_registry.py index 29e58b8101..bb8fa580e3 100644 --- a/litellm/litellm_core_utils/custom_logger_registry.py +++ b/litellm/litellm_core_utils/custom_logger_registry.py @@ -47,6 +47,7 @@ from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook i VectorStorePreCallHook, ) from litellm.proxy.hooks.dynamic_rate_limiter import _PROXY_DynamicRateLimitHandler +from litellm.proxy.hooks.dynamic_rate_limiter_v3 import _PROXY_DynamicRateLimitHandlerV3 class CustomLoggerRegistry: @@ -86,6 +87,7 @@ class CustomLoggerRegistry: "s3_v2": S3Logger, "aws_sqs": SQSLogger, "dynamic_rate_limiter": _PROXY_DynamicRateLimitHandler, + "dynamic_rate_limiter_v3": _PROXY_DynamicRateLimitHandlerV3, "vector_store_pre_call_hook": VectorStorePreCallHook, "dotprompt": DotpromptManager, "cloudzero": CloudZeroLogger, diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 0987f2799b..059fd9f1fc 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -3444,6 +3444,30 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915 dynamic_rate_limiter_obj.update_variables(llm_router=llm_router) _in_memory_loggers.append(dynamic_rate_limiter_obj) return dynamic_rate_limiter_obj # type: ignore + elif logging_integration == "dynamic_rate_limiter_v3": + from litellm.proxy.hooks.dynamic_rate_limiter_v3 import ( + _PROXY_DynamicRateLimitHandlerV3, + ) + + for callback in _in_memory_loggers: + if isinstance(callback, _PROXY_DynamicRateLimitHandlerV3): + return callback # type: ignore + + if internal_usage_cache is None: + raise Exception( + "Internal Error: Cache cannot be empty - internal_usage_cache={}".format( + internal_usage_cache + ) + ) + + dynamic_rate_limiter_obj_v3 = _PROXY_DynamicRateLimitHandlerV3( + internal_usage_cache=internal_usage_cache + ) + + if llm_router is not None and isinstance(llm_router, litellm.Router): + dynamic_rate_limiter_obj_v3.update_variables(llm_router=llm_router) + _in_memory_loggers.append(dynamic_rate_limiter_obj_v3) + return dynamic_rate_limiter_obj_v3 # type: ignore elif logging_integration == "langtrace": if "LANGTRACE_API_KEY" not in os.environ: raise ValueError("LANGTRACE_API_KEY not found in environment variables") @@ -3707,6 +3731,14 @@ def get_custom_logger_compatible_class( # noqa: PLR0915 for callback in _in_memory_loggers: if isinstance(callback, _PROXY_DynamicRateLimitHandler): return callback # type: ignore + elif logging_integration == "dynamic_rate_limiter_v3": + from litellm.proxy.hooks.dynamic_rate_limiter_v3 import ( + _PROXY_DynamicRateLimitHandlerV3, + ) + + for callback in _in_memory_loggers: + if isinstance(callback, _PROXY_DynamicRateLimitHandlerV3): + return callback # type: ignore elif logging_integration == "langtrace": from litellm.integrations.opentelemetry import OpenTelemetry diff --git a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py new file mode 100644 index 0000000000..fef03d5474 --- /dev/null +++ b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py @@ -0,0 +1,226 @@ +""" +Dynamic rate limiter v3 +""" + +import os +from typing import List, Literal, 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.utils import InternalUsageCache +from litellm.types.router import ModelGroupInfo + + +class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): + """ + Simple validation version that uses v3 parallel request limiter for priority-based rate limiting. + + 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 + """ + def __init__(self, internal_usage_cache: DualCache): + self.internal_usage_cache = InternalUsageCache(dual_cache=internal_usage_cache) + self.v3_limiter = _PROXY_MaxParallelRequestsHandler_v3(self.internal_usage_cache) + + def update_variables(self, llm_router: Router): + self.llm_router = llm_router + + def _get_priority_weight(self, priority: Optional[str]) -> float: + """Get the weight for a given priority from litellm.priority_reservation""" + weight: float = 1.0 + 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: + weight = litellm.priority_reservation[priority] + return weight + + def _create_priority_based_descriptors( + self, + model: str, + user_api_key_dict: UserAPIKeyAuth, + priority: Optional[str], + ) -> List[RateLimitDescriptor]: + """ + Create rate limit descriptors based on priority and model group limits. + + 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. + """ + descriptors: List[RateLimitDescriptor] = [] + + # 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 priority weight + priority_weight = self._get_priority_weight(priority) + + # Create priority-specific rate limits + # Use model:priority as the key to separate different priority levels + priority_key = f"{model}:{priority or 'default'}" + + rate_limit_config: RateLimitDescriptorRateLimitObject = {} + + # Apply priority weight to model limits + if model_group_info.tpm is not None: + # Reserve portion of TPM based on 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 + 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 + + async def async_pre_call_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + cache: DualCache, + data: dict, + call_type: Literal[ + "completion", + "text_completion", + "embeddings", + "image_generation", + "moderation", + "audio_transcription", + "pass_through_endpoint", + "rerank", + "mcp_call", + ], + ) -> Optional[Union[Exception, str, dict]]: + """ + Pre-call hook using v3 limiter for priority-based rate limiting. + """ + if "model" not in data: + return None + + 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, + ) + + if not descriptors: + verbose_proxy_logger.debug("No rate limit descriptors created, allowing request") + return None + + 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, + ) + + 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", + }, + ) + else: + # Store response for post-call hook + data["litellm_proxy_rate_limit_response"] = response + + except HTTPException: + raise + except Exception as e: + verbose_proxy_logger.exception( + f"Error in dynamic rate limiter v3 pre-call hook: {str(e)}" + ) + # Allow request to proceed 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): + key_priority: Optional[str] = user_api_key_dict.metadata.get("priority", None) + + # Get existing additional headers + additional_headers = getattr(response, "_hidden_params", {}).get("additional_headers", {}) or {} + + # Add priority information + additional_headers["x-litellm-priority"] = key_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 diff --git a/tests/test_litellm/proxy/hooks/test_dynamic_rate_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_dynamic_rate_limiter_v3.py new file mode 100644 index 0000000000..a7b5e4f1b1 --- /dev/null +++ b/tests/test_litellm/proxy/hooks/test_dynamic_rate_limiter_v3.py @@ -0,0 +1,482 @@ +""" +Test priority-based rate limiting for dynamic_rate_limiter_v3. + +Core tests to validate that priority weights are respected (0.9/0.1) instead of equal splitting (0.5/0.5). +""" +import asyncio +import os +import sys +import time +from unittest.mock import AsyncMock, patch + +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +import litellm +from litellm import DualCache, Router +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.hooks.dynamic_rate_limiter_v3 import ( + _PROXY_DynamicRateLimitHandlerV3 as DynamicRateLimitHandler, +) + + +@pytest.mark.asyncio +async def test_priority_weight_allocation(): + """ + Test that priority weights are correctly applied instead of equal splitting. + + With priority_reservation = {"high": 0.9, "low": 0.1}: + - High priority should get 90% of TPM (900 out of 1000) + - Low priority should get 10% of TPM (100 out of 1000) + + This validates the core fix where before it would split 50/50. + """ + # Set up priority reservations + litellm.priority_reservation = {"high": 0.9, "low": 0.1} + + dual_cache = DualCache() + handler = DynamicRateLimitHandler(internal_usage_cache=dual_cache) + + model = "test-model" + total_tpm = 1000 + + llm_router = Router( + model_list=[ + { + "model_name": model, + "litellm_params": { + "model": "gpt-3.5-turbo", + "api_key": "test-key", + "api_base": "test-base", + "tpm": total_tpm, + }, + } + ] + ) + handler.update_variables(llm_router=llm_router) + + # Test high priority allocation + high_priority_user = UserAPIKeyAuth() + high_priority_user.metadata = {"priority": "high"} + + high_descriptors = handler._create_priority_based_descriptors( + model=model, + user_api_key_dict=high_priority_user, + priority="high", + ) + + assert len(high_descriptors) == 1 + high_descriptor = high_descriptors[0] + expected_high_tpm = int(total_tpm * 0.9) # 900 + actual_high_tpm = high_descriptor["rate_limit"]["tokens_per_unit"] + + assert actual_high_tpm == expected_high_tpm, ( + f"High priority should get {expected_high_tpm} TPM (90%), got {actual_high_tpm}" + ) + assert high_descriptor["value"] == f"{model}:high" + + # Test low priority allocation + low_priority_user = UserAPIKeyAuth() + low_priority_user.metadata = {"priority": "low"} + + low_descriptors = handler._create_priority_based_descriptors( + model=model, + user_api_key_dict=low_priority_user, + priority="low", + ) + + assert len(low_descriptors) == 1 + low_descriptor = low_descriptors[0] + expected_low_tpm = int(total_tpm * 0.1) # 100 + actual_low_tpm = low_descriptor["rate_limit"]["tokens_per_unit"] + + assert actual_low_tpm == expected_low_tpm, ( + f"Low priority should get {expected_low_tpm} TPM (10%), got {actual_low_tpm}" + ) + assert low_descriptor["value"] == f"{model}:low" + + # Verify the ratio is 9:1, not 1:1 (equal splitting) + ratio = actual_high_tpm / actual_low_tpm + expected_ratio = 9.0 + assert abs(ratio - expected_ratio) < 0.1, ( + f"High:Low ratio should be {expected_ratio}:1, got {ratio}:1" + ) + + +@pytest.mark.asyncio +async def test_concurrent_priority_requests(): + """ + Test the core issue: 5 concurrent requests with different priorities should get + proper allocation based on priority weights, not equal splitting. + + This tests the exact scenario mentioned: priorities 0.9 and 0.1 should be 0.9/0.1, not 0.5/0.5. + """ + # Set up the exact scenario from the issue + litellm.priority_reservation = {"high": 0.9, "low": 0.1} + + dual_cache = DualCache() + handler = DynamicRateLimitHandler(internal_usage_cache=dual_cache) + + model = "test-model" + total_tpm = 1000 + + llm_router = Router( + model_list=[ + { + "model_name": model, + "litellm_params": { + "model": "gpt-3.5-turbo", + "api_key": "test-key", + "api_base": "test-base", + "tpm": total_tpm, + }, + } + ] + ) + handler.update_variables(llm_router=llm_router) + + # Create 5 concurrent users - 3 high priority, 2 low priority + high_priority_users = [] + low_priority_users = [] + + for i in range(3): # 3 high priority users + user = UserAPIKeyAuth() + user.metadata = {"priority": "high"} + user.user_id = f"high_user_{i}" + high_priority_users.append(user) + + for i in range(2): # 2 low priority users + user = UserAPIKeyAuth() + user.metadata = {"priority": "low"} + user.user_id = f"low_user_{i}" + low_priority_users.append(user) + + # Test all high priority users get the same allocation (not divided) + for user in high_priority_users: + descriptors = handler._create_priority_based_descriptors( + model=model, + user_api_key_dict=user, + priority="high", + ) + + assert len(descriptors) == 1 + descriptor = descriptors[0] + # Each high priority user should get 900 TPM, not divided by 3 + assert descriptor["rate_limit"]["tokens_per_unit"] == 900, ( + f"High priority user {user.user_id} should get 900 TPM, " + f"got {descriptor['rate_limit']['tokens_per_unit']}" + ) + assert descriptor["value"] == f"{model}:high" + + # Test all low priority users get the same allocation (not divided) + for user in low_priority_users: + descriptors = handler._create_priority_based_descriptors( + model=model, + user_api_key_dict=user, + priority="low", + ) + + assert len(descriptors) == 1 + descriptor = descriptors[0] + # Each low priority user should get 100 TPM, not divided by 2 + assert descriptor["rate_limit"]["tokens_per_unit"] == 100, ( + f"Low priority user {user.user_id} should get 100 TPM, " + f"got {descriptor['rate_limit']['tokens_per_unit']}" + ) + assert descriptor["value"] == f"{model}:low" + + +@pytest.mark.asyncio +async def test_100_concurrent_priority_requests(): + """ + Stress test: 100 concurrent requests with mixed priorities over 10 seconds. + + This validates that the priority system works correctly under high load: + - 70 high priority requests (should get 900 TPM each) + - 30 low priority requests (should get 100 TPM each) + - Spread across 10 seconds to simulate real-world load + """ + # Set up priority reservations + litellm.priority_reservation = {"high": 0.9, "low": 0.1} + + dual_cache = DualCache() + handler = DynamicRateLimitHandler(internal_usage_cache=dual_cache) + + model = "stress-test-model" + total_tpm = 1000 + + llm_router = Router( + model_list=[ + { + "model_name": model, + "litellm_params": { + "model": "gpt-3.5-turbo", + "api_key": "test-key", + "api_base": "test-base", + "tpm": total_tpm, + "rpm": 500, # Also test RPM limits + }, + } + ] + ) + handler.update_variables(llm_router=llm_router) + + # Create 100 users: 70 high priority, 30 low priority + all_users = [] + + # 70 high priority users + for i in range(70): + user = UserAPIKeyAuth() + user.metadata = {"priority": "high"} + user.user_id = f"high_stress_user_{i}" + all_users.append((user, "high", 900, 450)) # expected TPM, expected RPM + + # 30 low priority users + for i in range(30): + user = UserAPIKeyAuth() + user.metadata = {"priority": "low"} + user.user_id = f"low_stress_user_{i}" + all_users.append((user, "low", 100, 50)) # expected TPM, expected RPM + + async def test_user_descriptors(user_data): + """Test descriptor creation for a single user.""" + user, priority, expected_tpm, expected_rpm = user_data + + descriptors = handler._create_priority_based_descriptors( + model=model, + user_api_key_dict=user, + priority=priority, + ) + + assert len(descriptors) == 1, f"User {user.user_id} should have exactly 1 descriptor" + descriptor = descriptors[0] + + # Validate TPM allocation + actual_tpm = descriptor["rate_limit"]["tokens_per_unit"] + assert actual_tpm == expected_tpm, ( + f"User {user.user_id} ({priority}) should get {expected_tpm} TPM, got {actual_tpm}" + ) + + # Validate RPM allocation + actual_rpm = descriptor["rate_limit"]["requests_per_unit"] + assert actual_rpm == expected_rpm, ( + f"User {user.user_id} ({priority}) should get {expected_rpm} RPM, got {actual_rpm}" + ) + + # Validate descriptor key + assert descriptor["value"] == f"{model}:{priority}" + assert descriptor["key"] == "priority_model" + + return { + "user_id": user.user_id, + "priority": priority, + "tpm": actual_tpm, + "rpm": actual_rpm, + "success": True + } + + # Run all 100 requests concurrently to simulate high load + start_time = time.time() + + # Split into batches to simulate requests over 10 seconds + batch_size = 10 # 10 requests per batch + batches = [all_users[i:i + batch_size] for i in range(0, len(all_users), batch_size)] + + all_results = [] + + for batch_idx, batch in enumerate(batches): + # Process each batch concurrently + batch_tasks = [test_user_descriptors(user_data) for user_data in batch] + batch_results = await asyncio.gather(*batch_tasks, return_exceptions=True) + all_results.extend(batch_results) + + # Add small delay between batches to spread over ~10 seconds + if batch_idx < len(batches) - 1: # Don't sleep after last batch + await asyncio.sleep(1.0) # 1 second between batches + + end_time = time.time() + total_duration = end_time - start_time + + # Validate that the test ran over approximately 10 seconds + assert total_duration >= 9.0, f"Test should take ~10 seconds, took {total_duration:.2f}s" + assert total_duration <= 15.0, f"Test took too long: {total_duration:.2f}s" + + # Validate all requests were successful + successful_results = [r for r in all_results if isinstance(r, dict) and r.get("success")] + assert len(successful_results) == 100, f"Expected 100 successful results, got {len(successful_results)}" + + # Validate priority distribution + high_priority_results = [r for r in successful_results if r["priority"] == "high"] + low_priority_results = [r for r in successful_results if r["priority"] == "low"] + + assert len(high_priority_results) == 70, f"Expected 70 high priority results, got {len(high_priority_results)}" + assert len(low_priority_results) == 30, f"Expected 30 low priority results, got {len(low_priority_results)}" + + # Validate all high priority users got correct allocation + for result in high_priority_results: + assert result["tpm"] == 900, f"High priority user {result['user_id']} got {result['tpm']} TPM, expected 900" + assert result["rpm"] == 450, f"High priority user {result['user_id']} got {result['rpm']} RPM, expected 450" + + # Validate all low priority users got correct allocation + for result in low_priority_results: + assert result["tpm"] == 100, f"Low priority user {result['user_id']} got {result['tpm']} TPM, expected 100" + assert result["rpm"] == 50, f"Low priority user {result['user_id']} got {result['rpm']} RPM, expected 50" + + print(f"✅ Successfully processed 100 concurrent requests in {total_duration:.2f}s") + print(f" - 70 high priority users: 900 TPM, 450 RPM each") + print(f" - 30 low priority users: 100 TPM, 50 RPM each") + print(f" - Priority ratio maintained: 9:1 (TPM) and 9:1 (RPM)") + + +@pytest.mark.asyncio +async def test_concurrent_pre_call_hooks_stress(): + """ + Stress test: 50 concurrent pre-call hooks with priority enforcement. + + This tests the actual rate limiting logic under concurrent load. + """ + litellm.priority_reservation = {"premium": 0.8, "standard": 0.2} + + dual_cache = DualCache() + handler = DynamicRateLimitHandler(internal_usage_cache=dual_cache) + + model = "pre-call-stress-model" + total_tpm = 2000 + + llm_router = Router( + model_list=[ + { + "model_name": model, + "litellm_params": { + "model": "gpt-3.5-turbo", + "api_key": "test-key", + "api_base": "test-base", + "tpm": total_tpm, + }, + } + ] + ) + handler.update_variables(llm_router=llm_router) + + # Mock the v3 limiter to simulate different scenarios + successful_requests = [] + rate_limited_requests = [] + + async def mock_should_rate_limit(descriptors, parent_otel_span=None): + """Mock rate limiter that allows premium users, limits some standard users.""" + descriptor = descriptors[0] + priority = descriptor["value"].split(":")[-1] + + if priority == "premium": + # Allow all premium requests + return { + "overall_code": "OK", + "statuses": [{ + "code": "OK", + "descriptor_key": descriptor["value"], + "rate_limit_type": "tokens_per_unit", + "limit_remaining": 1000 + }] + } + else: + # Rate limit some standard requests (simulate load) + import random + if random.random() < 0.3: # 30% of standard requests get rate limited + return { + "overall_code": "OVER_LIMIT", + "statuses": [{ + "code": "OVER_LIMIT", + "descriptor_key": descriptor["value"], + "rate_limit_type": "tokens_per_unit", + "limit_remaining": 0 + }] + } + else: + return { + "overall_code": "OK", + "statuses": [{ + "code": "OK", + "descriptor_key": descriptor["value"], + "rate_limit_type": "tokens_per_unit", + "limit_remaining": 100 + }] + } + + # Create 50 users: 30 premium, 20 standard + users = [] + + for i in range(30): + user = UserAPIKeyAuth() + user.metadata = {"priority": "premium"} + user.user_id = f"premium_hook_user_{i}" + users.append((user, "premium")) + + for i in range(20): + user = UserAPIKeyAuth() + user.metadata = {"priority": "standard"} + user.user_id = f"standard_hook_user_{i}" + users.append((user, "standard")) + + async def make_request(user_data): + """Make a pre-call hook request.""" + user, priority = user_data + + with patch.object(handler.v3_limiter, 'should_rate_limit', side_effect=mock_should_rate_limit): + try: + result = await handler.async_pre_call_hook( + user_api_key_dict=user, + cache=DualCache(), + data={"model": model}, + call_type="completion", + ) + + # If no exception, request was allowed + successful_requests.append({ + "user_id": user.user_id, + "priority": priority, + "result": "allowed" + }) + return {"status": "success", "user_id": user.user_id, "priority": priority} + + except Exception as e: + # Request was rate limited + rate_limited_requests.append({ + "user_id": user.user_id, + "priority": priority, + "error": str(e) + }) + return {"status": "rate_limited", "user_id": user.user_id, "priority": priority} + + # Run all 50 requests concurrently + start_time = time.time() + tasks = [make_request(user_data) for user_data in users] + results = await asyncio.gather(*tasks, return_exceptions=True) + end_time = time.time() + + # Analyze results + successful_count = len([r for r in results if isinstance(r, dict) and r["status"] == "success"]) + rate_limited_count = len([r for r in results if isinstance(r, dict) and r["status"] == "rate_limited"]) + + # Validate that premium users were mostly successful (priority worked) + premium_results = [r for r in results if isinstance(r, dict) and r["priority"] == "premium"] + premium_success = len([r for r in premium_results if r["status"] == "success"]) + + standard_results = [r for r in results if isinstance(r, dict) and r["priority"] == "standard"] + standard_success = len([r for r in standard_results if r["status"] == "success"]) + + # Premium users should have higher success rate due to priority + premium_success_rate = premium_success / len(premium_results) if premium_results else 0 + standard_success_rate = standard_success / len(standard_results) if standard_results else 0 + + assert premium_success_rate >= 0.9, 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%}" + assert premium_success_rate > standard_success_rate, "Premium should have higher success rate than standard" + + total_duration = end_time - start_time + + print(f"✅ Processed 50 concurrent pre-call hooks in {total_duration:.2f}s") + print(f" - Premium users: {premium_success}/{len(premium_results)} success ({premium_success_rate:.1%})") + print(f" - Standard users: {standard_success}/{len(standard_results)} success ({standard_success_rate:.1%})") + print(f" - Total successful: {successful_count}/50 ({successful_count/50:.1%})") + print(f" - Priority system working: Premium > Standard success rates")