From ef008138a3df8368bd2203cc5d0e29ce8b887d8e Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Tue, 18 Mar 2025 19:44:04 -0700
Subject: [PATCH] feat(base_routing_strategy.py): handle updating in-memory
 keys

---
 litellm/proxy/_new_secret_config.yaml         |  4 ++-
 .../router_strategy/base_routing_strategy.py  | 26 ++++++++++++++-----
 litellm/router_strategy/lowest_tpm_rpm_v2.py  |  6 -----
 3 files changed, 22 insertions(+), 14 deletions(-)

diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml
index 00b3f325b0d..36ec0550481 100644
--- a/litellm/proxy/_new_secret_config.yaml
+++ b/litellm/proxy/_new_secret_config.yaml
@@ -4,12 +4,14 @@ model_list:
     model: azure/chatgpt-v-2
     api_key: os.environ/AZURE_API_KEY
     api_base: http://0.0.0.0:8090
+    rpm: 3
     
 litellm_settings:
+  num_retries: 0
   callbacks: ["otel"]
 
 router_settings:
-  # routing_strategy: usage-based-routing-v2 # 👈 KEY CHANGE
+  routing_strategy: usage-based-routing-v2 # 👈 KEY CHANGE
   redis_host: os.environ/REDIS_HOST
   redis_password: os.environ/REDIS_PASSWORD
   redis_port: os.environ/REDIS_PORT
diff --git a/litellm/router_strategy/base_routing_strategy.py b/litellm/router_strategy/base_routing_strategy.py
index b9a2f9cf04b..23b0c487a9a 100644
--- a/litellm/router_strategy/base_routing_strategy.py
+++ b/litellm/router_strategy/base_routing_strategy.py
@@ -3,8 +3,8 @@ Base class across routing strategies to abstract commmon functions like batch in
 """
 
 import asyncio
-from abc import ABC, abstractmethod
-from typing import List, Optional, Union
+from abc import ABC
+from typing import List, Optional, Set, Union
 
 from litellm._logging import verbose_router_logger
 from litellm.caching.caching import DualCache
@@ -28,6 +28,10 @@ class BaseRoutingStrategy(ABC):
                 )
             )
 
+        self.in_memory_keys_to_update: set[str] = (
+            set()
+        )  # Set with max size of 1000 keys
+
     async def _increment_value_in_current_window(
         self, key: str, value: Union[int, float], ttl: int
     ):
@@ -50,7 +54,7 @@ class BaseRoutingStrategy(ABC):
             ttl=ttl,
         )
         self.redis_increment_operation_queue.append(increment_op)
-
+        self.add_to_cache_keys(key=key)
         return result
 
     async def periodic_sync_in_memory_spend_with_redis(
@@ -104,9 +108,14 @@ class BaseRoutingStrategy(ABC):
                 f"Error syncing in-memory cache with Redis: {str(e)}"
             )
 
-    @abstractmethod
-    def get_cache_keys(self) -> List:
-        pass
+    def add_to_cache_keys(self, key: str):
+        self.in_memory_keys_to_update.add(key)
+
+    def get_cache_keys(self) -> Set[str]:
+        return self.in_memory_keys_to_update
+
+    def reset_cache_keys(self):
+        self.in_memory_keys_to_update = set()
 
     async def _sync_in_memory_spend_with_redis(self):
         """
@@ -132,9 +141,11 @@ class BaseRoutingStrategy(ABC):
             # 2. Fetch all current provider spend from Redis to update in-memory cache
             cache_keys = self.get_cache_keys()
 
+            cache_keys_list = list(cache_keys)
+
             # Batch fetch current spend values from Redis
             redis_values = await self.dual_cache.redis_cache.async_batch_get_cache(
-                key_list=cache_keys
+                key_list=cache_keys_list
             )
 
             # Update in-memory cache with Redis values
@@ -148,6 +159,7 @@ class BaseRoutingStrategy(ABC):
                             f"Updated in-memory cache for {key}: {value}"
                         )
 
+            self.reset_cache_keys()
         except Exception as e:
             verbose_router_logger.error(
                 f"Error syncing in-memory cache with Redis: {str(e)}"
diff --git a/litellm/router_strategy/lowest_tpm_rpm_v2.py b/litellm/router_strategy/lowest_tpm_rpm_v2.py
index 45bb14ee028..2b41033efe7 100644
--- a/litellm/router_strategy/lowest_tpm_rpm_v2.py
+++ b/litellm/router_strategy/lowest_tpm_rpm_v2.py
@@ -198,12 +198,6 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger):
                 result = await self._increment_value_in_current_window(
                     key=rpm_key, value=1, ttl=self.routing_args.ttl
                 )
-                # result = await self.router_cache.async_increment_cache(
-                #     key=rpm_key,
-                #     value=1,
-                #     ttl=self.routing_args.ttl,
-                #     parent_otel_span=parent_otel_span,
-                # )
                 if result is not None and result > deployment_rpm:
                     raise litellm.RateLimitError(
                         message="Deployment over defined rpm limit={}. current usage={}".format(