feat(parallel_request_limiter_v2.py): add sliding window logic (#11283)

* feat(parallel_request_limiter_v2.py): add sliding window logic

allows rate limiting to work across minutes

* fix(parallel_request_limiter_v2.py): decrement usage on rate limit error

* fix(base_routing_strategy.py): fix merge from redis - preserve values in in-memory cache during gap b/w push to redis and read from redis

* fix(base_routing_strategy.py): catch the delta change during redis sync

ensures values are kept in sync

* fix(parallel_request_limiter_v2.py): update tpm tracking to use slot key logic

* fix: fix linting error

* test: update testing

* test: update tests

* test: skip on rate limit or internal server errors

* test: use pytest fixture instead

* test: bump mistral model

Author: krrishdholakia

litellm/proxy/hooks/parallel_request_limiter_v2.py

@@ -132,42 +132,116 @@ async def check_key_in_limits_v2(
     ):
         ## INCREMENT CURRENT USAGE
         increment_list: List[Tuple[str, int]] = []
+        decrement_list: List[Tuple[str, int]] = []
+        slots_to_check: List[str] = []
         increment_value_by_group = {
             "request_count": 1,
             "tpm": 0,
             "rpm": 1,
         }
-        for group in ["request_count", "rpm", "tpm"]:
-            key = self._get_current_usage_key(
-                user_api_key_dict=user_api_key_dict,
-                precise_minute=precise_minute,
-                model=data.get("model", None),
-                rate_limit_type=rate_limit_type,
-                group=cast(RateLimitGroups, group),
-            )
-            if key is None:
-                continue
-            increment_list.append((key, increment_value_by_group[group]))
+
+        # Get current time and calculate the last 4 15s slots
+        current_time = datetime.now()
+        current_slot = (
+            current_time.second // 15
+        )  # This gives us 0-3 for the current 15s slot
+        slots_to_check = []
+        slot_cache_keys = []
+        # Calculate the last 4 slots, handling minute boundaries
+        for i in range(4):
+            slot_number = (current_slot - i) % 4  # This ensures we wrap around properly
+            minute = current_time.minute
+            hour = current_time.hour
+
+            # If we need to look at previous minute
+            if current_slot - i < 0:
+                if minute == 0:
+                    # If we're at minute 0, go to previous hour
+                    hour = (current_time.hour - 1) % 24
+                    minute = 59
+                else:
+                    minute = current_time.minute - 1
+
+            slot_key = f"{current_time.strftime('%Y-%m-%d')}-{hour:02d}-{minute:02d}-{slot_number}"
+            slots_to_check.append(slot_key)
+
+        # For each slot, create keys for all rate limit groups
+        for slot_key in slots_to_check:
+            for group in ["request_count", "rpm", "tpm"]:
+                key = self._get_current_usage_key(
+                    user_api_key_dict=user_api_key_dict,
+                    precise_minute=slot_key,
+                    model=data.get("model", None),
+                    rate_limit_type=rate_limit_type,
+                    group=cast(RateLimitGroups, group),
+                )
+                if key is None:
+                    continue
+                # Only increment the current slot
+                if slot_key == slots_to_check[0]:
+                    increment_list.append((key, increment_value_by_group[group]))
+                    decrement_list.append(
+                        (key, -1 if increment_value_by_group[group] == 1 else 0)
+                    )
+                slot_cache_keys.append(key)
 
         if (
             not max_parallel_requests and not rpm_limit and not tpm_limit
         ):  # no rate limits
             return
 
-        results = await self._increment_value_list_in_current_window(
+        # Use the existing atomic increment-and-check functionality
+        await self._increment_value_list_in_current_window(
             increment_list=increment_list,
             ttl=60,
         )
+
+        # Get the current values for all slots to check limits
+        current_values = await self.internal_usage_cache.async_batch_get_cache(
+            slot_cache_keys
+        )
+        if current_values is None:
+            current_values = [None] * len(slot_cache_keys)
+
+        # Calculate totals across all slots, handling None values
+        # Group values by type (request_count, rpm, tpm)
+        request_counts = []
+        rpm_counts = []
+        tpm_counts = []
+
+        for i in range(0, len(current_values), 3):
+            request_counts.append(
+                current_values[i] if current_values[i] is not None else 0
+            )
+            rpm_counts.append(
+                current_values[i + 1] if current_values[i + 1] is not None else 0
+            )
+            tpm_counts.append(
+                current_values[i + 2] if current_values[i + 2] is not None else 0
+            )
+
+        # Calculate totals across all slots
+        total_requests = sum(request_counts)
+        total_rpm = sum(rpm_counts)
+        total_tpm = sum(tpm_counts)
+
         should_raise_error = False
         if max_parallel_requests is not None:
-            should_raise_error = results[0] > max_parallel_requests
+            should_raise_error = total_requests > max_parallel_requests
         if rpm_limit is not None:
-            should_raise_error = should_raise_error or results[1] > rpm_limit
+            should_raise_error = should_raise_error or total_rpm > rpm_limit
         if tpm_limit is not None:
-            should_raise_error = should_raise_error or results[2] > tpm_limit
+            should_raise_error = should_raise_error or total_tpm > tpm_limit
+
         if should_raise_error:
+            ## DECREMENT CURRENT USAGE - so we don't keep failing subsequent requests
+            await self._increment_value_list_in_current_window(
+                increment_list=decrement_list,
+                ttl=60,
+            )
+
             raise self.raise_rate_limit_error(
-                additional_details=f"{CommonProxyErrors.max_parallel_request_limit_reached.value}. Hit limit for {rate_limit_type}. Current usage: max_parallel_requests: {results[0]}, current_rpm: {results[1]}, current_tpm: {results[2]}. Current limits: max_parallel_requests: {max_parallel_requests}, rpm_limit: {rpm_limit}, tpm_limit: {tpm_limit}."
+                additional_details=f"{CommonProxyErrors.max_parallel_request_limit_reached.value}. Hit limit for {rate_limit_type}. Current usage: max_parallel_requests: {total_requests}, current_rpm: {total_rpm}, current_tpm: {total_tpm}. Current limits: max_parallel_requests: {max_parallel_requests}, rpm_limit: {rpm_limit}, tpm_limit: {tpm_limit}."
             )
 
     def time_to_next_minute(self) -> float:
@@ -356,11 +430,18 @@ async def _update_usage_in_cache_post_call(
         }
 
         rate_limit_types = ["key", "user", "customer", "team", "model_per_key"]
+        current_time = datetime.now()
+        current_hour = current_time.hour
+        current_minute = current_time.minute
+        current_slot = (
+            current_time.second // 15
+        )  # This gives us 0-3 for the current 15s slot
+        slot_key = f"{current_time.strftime('%Y-%m-%d')}-{current_hour:02d}-{current_minute:02d}-{current_slot}"
         for rate_limit_type in rate_limit_types:
             for group in ["request_count", "rpm", "tpm"]:
                 key = self._get_current_usage_key(
                     user_api_key_dict=user_api_key_dict,
-                    precise_minute=precise_minute,
+                    precise_minute=slot_key,
                     model=model,
                     rate_limit_type=cast(RateLimitTypes, rate_limit_type),
                     group=cast(RateLimitGroups, group),

litellm/proxy/proxy_server.py

@@ -2729,7 +2729,7 @@ async def _init_non_llm_objects_in_db(self, prisma_client: PrismaClient):
         """
         await self._init_guardrails_in_db(prisma_client=prisma_client)
         await self._init_vector_stores_in_db(prisma_client=prisma_client)
-        await self._init_mcp_servers_in_db()
+        # await self._init_mcp_servers_in_db()
 
     async def _init_guardrails_in_db(self, prisma_client: PrismaClient):
         from litellm.proxy.guardrails.guardrail_registry import (

litellm/router_strategy/base_routing_strategy.py

@@ -178,38 +178,54 @@ async def _sync_in_memory_spend_with_redis(self):
             await self._push_in_memory_increments_to_redis()
 
             # 2. Fetch all current provider spend from Redis to update in-memory cache
-            pattern = self.get_key_pattern_to_sync()
-            cache_keys: Optional[Union[Set[str], List[str]]] = None
-            if pattern:
-                cache_keys = await self.dual_cache.redis_cache.async_scan_iter(
-                    pattern=pattern
-                )
-
-            if cache_keys is None:
-                cache_keys = (
-                    self.get_in_memory_keys_to_update()
-                )  # if no pattern OR redis cache does not support scan_iter, use in-memory keys
+            cache_keys = (
+                self.get_in_memory_keys_to_update()
+            )  # if no pattern OR redis cache does not support scan_iter, use in-memory keys
 
             if isinstance(cache_keys, set):
                 cache_keys_list = list(cache_keys)
             else:
                 cache_keys_list = cache_keys
 
-            # Batch fetch current spend values from Redis
+            # 1. Snapshot in-memory before
+            in_memory_before_dict = {}
+            in_memory_before = (
+                await self.dual_cache.in_memory_cache.async_batch_get_cache(
+                    keys=cache_keys_list
+                )
+            )
+            for k, v in zip(cache_keys_list, in_memory_before):
+                in_memory_before_dict[k] = v
+
+            # 2. Fetch from Redis
             redis_values = await self.dual_cache.redis_cache.async_batch_get_cache(
                 key_list=cache_keys_list
             )
 
-            # Update in-memory cache with Redis values
-            if isinstance(redis_values, dict):  # Check if redis_values is a dictionary
-                for key, value in redis_values.items():
-                    if value is not None:
-                        await self.dual_cache.in_memory_cache.async_set_cache(
-                            key=key, value=float(value)
-                        )
-                        # verbose_router_logger.debug(
-                        #     f"Updated in-memory cache for {key}: {value}"
-                        # )
+            # 3. Snapshot in-memory after
+            in_memory_after = (
+                await self.dual_cache.in_memory_cache.async_batch_get_cache(
+                    keys=cache_keys_list
+                )
+            )
+            in_memory_after_dict = {}
+            for k, v in zip(cache_keys_list, in_memory_after):
+                in_memory_after_dict[k] = v
+
+            # 4. Merge
+            for key in cache_keys_list:
+                redis_val = float(redis_values.get(key, 0) or 0)
+                before = float(in_memory_before_dict.get(key, 0) or 0)
+                after = float(in_memory_after_dict.get(key, 0) or 0)
+                delta = after - before
+                if delta > 0:
+                    await self._increment_value_in_current_window(
+                        key=key, value=delta, ttl=60
+                    )
+                merged = redis_val + delta
+                await self.dual_cache.in_memory_cache.async_set_cache(
+                    key=key, value=merged
+                )
 
             self.reset_in_memory_keys_to_update()
         except Exception as e:

tests/llm_translation/base_llm_unit_tests.py

@@ -8,6 +8,7 @@
 import uuid
 import time
 import base64
+import inspect
 
 sys.path.insert(
     0, os.path.abspath("../..")
@@ -76,11 +77,20 @@ def get_base_completion_call_args(self) -> dict:
         """Must return the base completion call args"""
         pass
 
-
     def get_base_completion_call_args_with_reasoning_model(self) -> dict:
         """Must return the base completion call args with reasoning_effort"""
         return {}
 
+    @pytest.fixture(autouse=True)
+    def _handle_rate_limits(self):
+        """Fixture to handle rate limit errors for all test methods"""
+        try:
+            yield
+        except litellm.RateLimitError:
+            pytest.skip("Rate limit exceeded")
+        except litellm.InternalServerError:
+            pytest.skip("Model is overloaded")
+
     def test_developer_role_translation(self):
         """
         Test that the developer role is translated correctly for non-OpenAI providers.

tests/llm_translation/test_cohere.py

@@ -164,6 +164,7 @@ def test_completion_cohere():
 # FYI - cohere_chat looks quite unstable, even when testing locally
 @pytest.mark.asyncio
 @pytest.mark.parametrize("sync_mode", [True, False])
+@pytest.mark.flaky(retries=3, delay=1)
 async def test_chat_completion_cohere(sync_mode):
     try:
         litellm.set_verbose = True

tests/llm_translation/test_mistral_api.py

@@ -31,7 +31,7 @@
 class TestMistralCompletion(BaseLLMChatTest):
     def get_base_completion_call_args(self) -> dict:
         litellm.set_verbose = True
-        return {"model": "mistral/mistral-small-latest"}
+        return {"model": "mistral/mistral-medium-latest"}
 
     def test_tool_call_no_arguments(self, tool_call_no_arguments):
         """Test that tool calls with no arguments is translated correctly. Relevant issue: https://github.com/BerriAI/litellm/issues/6833"""