From 1328afe612b7532f90f9c2f20360d5460dee5297 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Tue, 18 Mar 2025 19:09:53 -0700 Subject: [PATCH] fix(lowest_tpm_rpm_v2.py): support batch writing increments to redis --- litellm/constants.py | 1 + .../proxy/_experimental/out/onboarding.html | 1 - litellm/proxy/_new_secret_config.yaml | 6 +- .../router_strategy/base_routing_strategy.py | 154 ++++++++++++++++++ litellm/router_strategy/lowest_tpm_rpm_v2.py | 24 ++- 5 files changed, 176 insertions(+), 10 deletions(-) delete mode 100644 litellm/proxy/_experimental/out/onboarding.html create mode 100644 litellm/router_strategy/base_routing_strategy.py diff --git a/litellm/constants.py b/litellm/constants.py index b4551a78f5f..eb59858d433 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -7,6 +7,7 @@ DEFAULT_MAX_RETRIES = 2 DEFAULT_FAILURE_THRESHOLD_PERCENT = ( 0.5 # default cooldown a deployment if 50% of requests fail in a given minute ) +DEFAULT_REDIS_SYNC_INTERVAL = 1 DEFAULT_COOLDOWN_TIME_SECONDS = 5 DEFAULT_REPLICATE_POLLING_RETRIES = 5 DEFAULT_REPLICATE_POLLING_DELAY_SECONDS = 1 diff --git a/litellm/proxy/_experimental/out/onboarding.html b/litellm/proxy/_experimental/out/onboarding.html deleted file mode 100644 index 82f43619df3..00000000000 --- a/litellm/proxy/_experimental/out/onboarding.html +++ /dev/null @@ -1 +0,0 @@ -LiteLLM Dashboard \ No newline at end of file diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml index 64100277a80..00b3f325b0d 100644 --- a/litellm/proxy/_new_secret_config.yaml +++ b/litellm/proxy/_new_secret_config.yaml @@ -3,13 +3,13 @@ model_list: litellm_params: model: azure/chatgpt-v-2 api_key: os.environ/AZURE_API_KEY - api_base: os.environ/AZURE_API_BASE + api_base: http://0.0.0.0:8090 litellm_settings: - callbacks: ["prometheus"] + 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 new file mode 100644 index 00000000000..b9a2f9cf04b --- /dev/null +++ b/litellm/router_strategy/base_routing_strategy.py @@ -0,0 +1,154 @@ +""" +Base class across routing strategies to abstract commmon functions like batch incrementing redis +""" + +import asyncio +from abc import ABC, abstractmethod +from typing import List, Optional, Union + +from litellm._logging import verbose_router_logger +from litellm.caching.caching import DualCache +from litellm.caching.redis_cache import RedisPipelineIncrementOperation +from litellm.constants import DEFAULT_REDIS_SYNC_INTERVAL + + +class BaseRoutingStrategy(ABC): + def __init__( + self, + dual_cache: DualCache, + should_batch_redis_writes: bool, + default_sync_interval: Optional[Union[int, float]], + ): + self.dual_cache = dual_cache + self.redis_increment_operation_queue: List[RedisPipelineIncrementOperation] = [] + if should_batch_redis_writes: + asyncio.create_task( + self.periodic_sync_in_memory_spend_with_redis( + default_sync_interval=default_sync_interval + ) + ) + + async def _increment_value_in_current_window( + self, key: str, value: Union[int, float], ttl: int + ): + """ + Increment spend within existing budget window + + Runs once the budget start time exists in Redis Cache (on the 2nd and subsequent requests to the same provider) + + - Increments the spend in memory cache (so spend instantly updated in memory) + - Queues the increment operation to Redis Pipeline (using batched pipeline to optimize performance. Using Redis for multi instance environment of LiteLLM) + """ + result = await self.dual_cache.in_memory_cache.async_increment( + key=key, + value=value, + ttl=ttl, + ) + increment_op = RedisPipelineIncrementOperation( + key=key, + increment_value=value, + ttl=ttl, + ) + self.redis_increment_operation_queue.append(increment_op) + + return result + + async def periodic_sync_in_memory_spend_with_redis( + self, default_sync_interval: Optional[Union[int, float]] + ): + """ + Handler that triggers sync_in_memory_spend_with_redis every DEFAULT_REDIS_SYNC_INTERVAL seconds + + Required for multi-instance environment usage of provider budgets + """ + default_sync_interval = default_sync_interval or DEFAULT_REDIS_SYNC_INTERVAL + while True: + try: + await self._sync_in_memory_spend_with_redis() + await asyncio.sleep( + default_sync_interval + ) # Wait for DEFAULT_REDIS_SYNC_INTERVAL seconds before next sync + except Exception as e: + verbose_router_logger.error(f"Error in periodic sync task: {str(e)}") + await asyncio.sleep( + default_sync_interval + ) # Still wait DEFAULT_REDIS_SYNC_INTERVAL seconds on error before retrying + + async def _push_in_memory_increments_to_redis(self): + """ + How this works: + - async_log_success_event collects all provider spend increments in `redis_increment_operation_queue` + - This function pushes all increments to Redis in a batched pipeline to optimize performance + + Only runs if Redis is initialized + """ + try: + if not self.dual_cache.redis_cache: + return # Redis is not initialized + + verbose_router_logger.debug( + "Pushing Redis Increment Pipeline for queue: %s", + self.redis_increment_operation_queue, + ) + if len(self.redis_increment_operation_queue) > 0: + asyncio.create_task( + self.dual_cache.redis_cache.async_increment_pipeline( + increment_list=self.redis_increment_operation_queue, + ) + ) + + self.redis_increment_operation_queue = [] + + except Exception as e: + verbose_router_logger.error( + f"Error syncing in-memory cache with Redis: {str(e)}" + ) + + @abstractmethod + def get_cache_keys(self) -> List: + pass + + async def _sync_in_memory_spend_with_redis(self): + """ + Ensures in-memory cache is updated with latest Redis values for all provider spends. + + Why Do we need this? + - Optimization to hit sub 100ms latency. Performance was impacted when redis was used for read/write per request + - Use provider budgets in multi-instance environment, we use Redis to sync spend across all instances + + What this does: + 1. Push all provider spend increments to Redis + 2. Fetch all current provider spend from Redis to update in-memory cache + """ + + try: + # No need to sync if Redis cache is not initialized + if self.dual_cache.redis_cache is None: + return + + # 1. Push all provider spend increments to Redis + await self._push_in_memory_increments_to_redis() + + # 2. Fetch all current provider spend from Redis to update in-memory cache + cache_keys = self.get_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 + ) + + # 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}" + ) + + 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 64f086036be..45bb14ee028 100644 --- a/litellm/router_strategy/lowest_tpm_rpm_v2.py +++ b/litellm/router_strategy/lowest_tpm_rpm_v2.py @@ -15,6 +15,8 @@ from litellm.types.router import RouterErrors from litellm.types.utils import LiteLLMPydanticObjectBase, StandardLoggingPayload from litellm.utils import get_utc_datetime, print_verbose +from .base_routing_strategy import BaseRoutingStrategy + if TYPE_CHECKING: from opentelemetry.trace import Span as _Span @@ -27,7 +29,7 @@ class RoutingArgs(LiteLLMPydanticObjectBase): ttl: int = 1 * 60 # 1min (RPM/TPM expire key) -class LowestTPMLoggingHandler_v2(CustomLogger): +class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): """ Updated version of TPM/RPM Logging. @@ -51,6 +53,12 @@ class LowestTPMLoggingHandler_v2(CustomLogger): self.router_cache = router_cache self.model_list = model_list self.routing_args = RoutingArgs(**routing_args) + BaseRoutingStrategy.__init__( + self, + dual_cache=router_cache, + should_batch_redis_writes=True, + default_sync_interval=0.1, + ) def pre_call_check(self, deployment: Dict) -> Optional[Dict]: """ @@ -104,6 +112,7 @@ class LowestTPMLoggingHandler_v2(CustomLogger): ) else: # if local result below limit, check redis ## prevent unnecessary redis checks + result = self.router_cache.increment_cache( key=rpm_key, value=1, ttl=self.routing_args.ttl ) @@ -186,12 +195,15 @@ class LowestTPMLoggingHandler_v2(CustomLogger): ) else: # if local result below limit, check redis ## prevent unnecessary redis checks - result = await self.router_cache.async_increment_cache( - key=rpm_key, - value=1, - ttl=self.routing_args.ttl, - parent_otel_span=parent_otel_span, + 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(