diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 89826694bb6..cdec5922fff 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -91,6 +91,11 @@ else: _REQUEST_RATE_LIMIT_DATA: Final = TypeAdapter(Mapping[str, object]) +def _sibling_counter_keys(window_key: str) -> tuple[str, str]: + prefix: Final = window_key.removesuffix(":window") + return f"{prefix}:requests", f"{prefix}:tokens" + + BATCH_RATE_LIMITER_SCRIPT: Final = """ local results = {} local now = tonumber(ARGV[1]) @@ -106,6 +111,8 @@ for i = 1, #KEYS, 2 do local window_start = redis.call('GET', window_key) if not window_start or (now - tonumber(window_start)) >= window_size then -- Reset window and counter + local prefix = string.sub(window_key, 1, -(#':window') - 1) + redis.call('DEL', prefix .. ':requests', prefix .. ':tokens') redis.call('SET', window_key, tostring(now)) redis.call('SET', counter_key, increment_value) redis.call('EXPIRE', window_key, window_size) @@ -151,6 +158,7 @@ CHECK_AND_INCREMENT_BY_N_SCRIPT: Final = """ local time_reply = redis.call('TIME') local now = tonumber(time_reply[1]) local descriptor_count = #KEYS / 2 +local reset_windows = {} -- Pass 1: read state, validate. Abort without writing if any over limit. local descriptor_state = {} @@ -201,6 +209,11 @@ for i = 1, descriptor_count do if window_expired then active_window_start = now + if not reset_windows[window_key] then + local prefix = string.sub(window_key, 1, -(#':window') - 1) + redis.call('DEL', prefix .. ':requests', prefix .. ':tokens') + reset_windows[window_key] = true + end redis.call('SET', window_key, tostring(now)) redis.call('SET', counter_key, increment) redis.call('EXPIRE', window_key, window_size) @@ -1018,6 +1031,15 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): Implement sliding window rate limiting logic using in-memory cache operations. This follows the same logic as the Redis Lua script but uses async cache operations. """ + async with self._check_and_increment_lock: + return await self._in_memory_cache_sliding_window(keys=keys, now_int=now_int, window_size=window_size) + + async def _in_memory_cache_sliding_window( + self, + keys: list[str], + now_int: int, + window_size: int, + ) -> CacheCounterValues: results: Final[list[CacheCounterValue | None]] = [] # Process each window/counter pair @@ -1036,6 +1058,14 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): # Check if window exists and is valid if window_start is None or (now_int - int(window_start)) >= window_size: # Reset window and counter + for sibling_counter_key in _sibling_counter_keys(window_key): + await self.internal_usage_cache.async_set_cache( + key=sibling_counter_key, + value=0, + ttl=window_size, + litellm_parent_otel_span=None, + local_only=True, + ) await self.internal_usage_cache.async_set_cache( key=window_key, value=str(now_int), @@ -2048,6 +2078,20 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) # Pass 2: apply increments. + expired_windows: Final[Mapping[str, int]] = { + meta["window_key"]: meta["window_size"] + for meta, state in zip(per_counter_meta, descriptor_state) + if state["window_expired"] + } + for window_key, window_size in expired_windows.items(): + for sibling_counter_key in _sibling_counter_keys(window_key): + await self.internal_usage_cache.async_set_cache( + key=sibling_counter_key, + value=0, + ttl=window_size, + litellm_parent_otel_span=parent_otel_span, + local_only=True, + ) statuses: Final[list[RateLimitStatus]] = [] for meta, state in zip(per_counter_meta, descriptor_state): new_counter = meta["increment"] if state["window_expired"] else state["current"] + meta["increment"] diff --git a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py index 85023b94207..5889b1b513f 100644 --- a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py @@ -18,11 +18,13 @@ from fastapi import HTTPException import litellm from litellm import Router from litellm.caching.caching import DualCache +from litellm.caching.in_memory_cache import InMemoryCache from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.hooks.parallel_request_limiter_v3 import ( PARALLEL_REQUEST_SLOT_TTL_SECONDS, ParallelSlotAcquisition, + RateLimitDescriptor, RequestRateLimiterStash, _request_stash, get_or_create_request_stash, @@ -5602,6 +5604,155 @@ async def _reserved_tokens_for( return int(await local_cache.async_get_cache(key=tokens_key) or 0) +@pytest.mark.asyncio +async def test_tpm_reservation_resets_sibling_tokens_with_request_window(monkeypatch): + monkeypatch.setenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", "true") + time_controller = TimeController() + local_cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache), + time_provider=time_controller.now, + ) + user_api_key_dict = UserAPIKeyAuth( + api_key=hash_token("sk-window-reset-siblings"), + tpm_limit=1000, + rpm_limit=1000, + ) + + async def request(call_id): + data = { + "model": "gpt-4o", + "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 200, + "litellm_call_id": call_id, + "metadata": { + "user_api_key": user_api_key_dict.api_key, + "user_api_key_user_id": user_api_key_dict.user_id, + }, + } + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data=data, + call_type="completion", + ) + await handler.async_log_success_event( + kwargs={ + "litellm_call_id": call_id, + "litellm_params": { + "metadata": { + "user_api_key": user_api_key_dict.api_key, + "user_api_key_user_id": user_api_key_dict.user_id, + "model_group": "gpt-4o", + } + }, + "standard_logging_object": { + "metadata": { + "user_api_key_hash": user_api_key_dict.api_key, + "user_api_key_user_id": user_api_key_dict.user_id, + } + }, + }, + response_obj=ModelResponse( + model="gpt-4o", + usage=Usage(prompt_tokens=100, completion_tokens=200, total_tokens=300), + ), + start_time=datetime.now(), + end_time=datetime.now(), + ) + + tokens_key = handler.create_rate_limit_keys( + key="api_key", value=user_api_key_dict.api_key, rate_limit_type="tokens" + ) + for index in range(3): + await request(f"call-{index}") + assert await local_cache.async_get_cache(key=tokens_key) == (index + 1) * 300 + + time_controller.advance(61) + await request("call-after-window-reset") + assert await local_cache.async_get_cache(key=tokens_key) == 300 + + +@pytest.mark.asyncio +async def test_atomic_tpm_reservation_rollover_resets_sibling_requests_counter(): + local_cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + window_size = 60 + now_int = int(time.time()) + window_key = "{api_key:atomic-rollover}:window" + requests_key = handler.create_rate_limit_keys("api_key", "atomic-rollover", "requests") + tokens_key = handler.create_rate_limit_keys("api_key", "atomic-rollover", "tokens") + for key, value in ((window_key, str(now_int - window_size - 1)), (requests_key, 3), (tokens_key, 900)): + await local_cache.async_set_cache(key=key, value=value, ttl=window_size) + + tpm_pass = await handler.atomic_check_and_increment_by_n( + descriptors=[ + RateLimitDescriptor( + key="api_key", + value="atomic-rollover", + rate_limit={"tokens_per_unit": 1000, "window_size": window_size}, + ) + ], + increments=[{"tokens": 200}], + ) + assert tpm_pass["overall_code"] == "OK" + assert await local_cache.async_get_cache(key=tokens_key) == 200 + + rpm_pass = await handler.should_rate_limit( + descriptors=[ + RateLimitDescriptor( + key="api_key", + value="atomic-rollover", + rate_limit={"requests_per_unit": 5, "window_size": window_size}, + ) + ], + skip_tpm_check=True, + ) + assert rpm_pass["overall_code"] == "OK" + assert [status["limit_remaining"] for status in rpm_pass["statuses"]] == [4] + assert await local_cache.async_get_cache(key=requests_key) == 1 + + +class _YieldingInMemoryCache(InMemoryCache): + async def async_get_cache(self, key: str, **kwargs: object) -> object: + value = await super().async_get_cache(key, **kwargs) + await asyncio.sleep(0) + return value + + +@pytest.mark.asyncio +async def test_window_rollover_reset_does_not_erase_concurrent_sibling_increment(): + local_cache = DualCache(in_memory_cache=_YieldingInMemoryCache()) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + window_size = 60 + now_int = int(time.time()) + window_key = "{api_key:concurrent-rollover}:window" + requests_key = handler.create_rate_limit_keys("api_key", "concurrent-rollover", "requests") + tokens_key = handler.create_rate_limit_keys("api_key", "concurrent-rollover", "tokens") + for key, value in ((window_key, str(now_int - window_size - 1)), (requests_key, 3), (tokens_key, 900)): + await local_cache.async_set_cache(key=key, value=value, ttl=window_size) + + tpm_descriptor = RateLimitDescriptor( + key="api_key", + value="concurrent-rollover", + rate_limit={"tokens_per_unit": 1000, "window_size": window_size}, + ) + rpm_pass, tpm_pass = await asyncio.gather( + handler.in_memory_cache_sliding_window( + keys=[window_key, requests_key], now_int=now_int, window_size=window_size + ), + handler.atomic_check_and_increment_by_n( + descriptors=[tpm_descriptor], + increments=[{"requests": 0, "tokens": 200}], + ), + ) + + assert rpm_pass == [str(now_int), 1] + assert tpm_pass["overall_code"] == "OK" + assert await local_cache.async_get_cache(key=requests_key) == 1 + assert await local_cache.async_get_cache(key=tokens_key) == 200 + + @pytest.mark.asyncio @pytest.mark.parametrize( "key_metadata, team_metadata, expected_output_estimate, tier",