fix(proxy): reset sibling tpm/rpm counters when shared window rolls over

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

litellm/proxy/hooks/parallel_request_limiter_v3.py

@@ -91,10 +91,16 @@ 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])
 local window_size = tonumber(ARGV[2])
+local reset_windows = {}
 
 -- Process each window/counter pair
 for i = 1, #KEYS, 2 do
@@ -106,6 +112,11 @@ 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
         redis.call('SET', window_key, tostring(now))
         redis.call('SET', counter_key, increment_value)
         redis.call('EXPIRE', window_key, window_size)
@@ -151,6 +162,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 +213,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)
@@ -1019,6 +1036,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
         This follows the same logic as the Redis Lua script but uses async cache operations.
         """
         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):
@@ -1036,6 +1054,16 @@ 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)
                 await self.internal_usage_cache.async_set_cache(
                     key=window_key,
                     value=str(now_int),
@@ -2049,9 +2077,20 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
 
         # Pass 2: apply increments.
         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

@@ -5602,6 +5602,75 @@ 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
 @pytest.mark.parametrize(
     "key_metadata, team_metadata, expected_output_estimate, tier",