fix(proxy): serialize in-memory window rollover and reset siblings once per expired window

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm/proxy/hooks/parallel_request_limiter_v3.py

@@ -100,7 +100,6 @@ BATCH_RATE_LIMITER_SCRIPT: Final = """
 local results = {}
 local now = tonumber(ARGV[1])
 local window_size = tonumber(ARGV[2])
-local reset_windows = {}
 
 -- Process each window/counter pair
 for i = 1, #KEYS, 2 do
@@ -112,11 +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
-        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
+        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)
@@ -1035,8 +1031,16 @@ 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]] = []
-        reset_windows: Final[set[str]] = set()  # mutable-ok: tracks windows reset during this call
 
         # Process each window/counter pair
         for i in range(0, len(keys), 2):
@@ -1054,16 +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
-                if window_key not in reset_windows:
-                    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,
-                        )
-                    reset_windows.add(window_key)
+                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),
@@ -2076,21 +2078,24 @@ 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]] = []
-        reset_windows: Final[set[str]] = set()  # mutable-ok: tracks windows reset during this call
         for meta, state in zip(per_counter_meta, descriptor_state):
             new_counter = meta["increment"] if state["window_expired"] else state["current"] + meta["increment"]
             if state["window_expired"]:
-                if meta["window_key"] not in reset_windows:
-                    for sibling_counter_key in _sibling_counter_keys(meta["window_key"]):
-                        await self.internal_usage_cache.async_set_cache(
-                            key=sibling_counter_key,
-                            value=0,
-                            ttl=meta["window_size"],
-                            litellm_parent_otel_span=parent_otel_span,
-                            local_only=True,
-                        )
-                    reset_windows.add(meta["window_key"])
                 await self.internal_usage_cache.async_set_cache(
                     key=meta["window_key"],
                     value=str(now_int),

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,
@@ -5671,6 +5673,86 @@ async def test_tpm_reservation_resets_sibling_tokens_with_request_window(monkeyp
     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",