#### What this does #### # identifies lowest tpm deployment import traceback from datetime import datetime from typing import Dict, List, Optional, Union from litellm import token_counter from litellm._logging import verbose_router_logger from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger from litellm.types.utils import LiteLLMPydanticObjectBase from litellm.utils import print_verbose class RoutingArgs(LiteLLMPydanticObjectBase): ttl: int = 1 * 60 # 1min (RPM/TPM expire key) class LowestTPMLoggingHandler(CustomLogger): test_flag: bool = False logged_success: int = 0 logged_failure: int = 0 default_cache_time_seconds: int = 1 * 60 * 60 # 1 hour def __init__( self, router_cache: DualCache, model_list: list, routing_args: dict = {} ): self.router_cache = router_cache self.model_list = model_list self.routing_args = RoutingArgs(**routing_args) def log_success_event(self, kwargs, response_obj, start_time, end_time): try: """ Update TPM/RPM usage on success """ if kwargs["litellm_params"].get("metadata") is None: pass else: model_group = kwargs["litellm_params"]["metadata"].get( "model_group", None ) id = kwargs["litellm_params"].get("model_info", {}).get("id", None) if model_group is None or id is None: return elif isinstance(id, int): id = str(id) total_tokens = response_obj["usage"]["total_tokens"] # ------------ # Setup values # ------------ current_minute = datetime.now().strftime("%H-%M") tpm_key = f"{model_group}:tpm:{current_minute}" rpm_key = f"{model_group}:rpm:{current_minute}" # ------------ # Update usage # ------------ ## TPM request_count_dict = self.router_cache.get_cache(key=tpm_key) or {} request_count_dict[id] = request_count_dict.get(id, 0) + total_tokens self.router_cache.set_cache( key=tpm_key, value=request_count_dict, ttl=self.routing_args.ttl ) ## RPM request_count_dict = self.router_cache.get_cache(key=rpm_key) or {} request_count_dict[id] = request_count_dict.get(id, 0) + 1 self.router_cache.set_cache( key=rpm_key, value=request_count_dict, ttl=self.routing_args.ttl ) ### TESTING ### if self.test_flag: self.logged_success += 1 except Exception as e: verbose_router_logger.error( "litellm.router_strategy.lowest_tpm_rpm.py::async_log_success_event(): Exception occured - {}".format( str(e) ) ) verbose_router_logger.debug(traceback.format_exc()) pass async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): try: """ Update TPM/RPM usage on success """ if kwargs["litellm_params"].get("metadata") is None: pass else: model_group = kwargs["litellm_params"]["metadata"].get( "model_group", None ) id = kwargs["litellm_params"].get("model_info", {}).get("id", None) if model_group is None or id is None: return elif isinstance(id, int): id = str(id) total_tokens = response_obj["usage"]["total_tokens"] # ------------ # Setup values # ------------ current_minute = datetime.now().strftime("%H-%M") tpm_key = f"{model_group}:tpm:{current_minute}" rpm_key = f"{model_group}:rpm:{current_minute}" # ------------ # Update usage # ------------ # update cache ## TPM request_count_dict = ( await self.router_cache.async_get_cache(key=tpm_key) or {} ) request_count_dict[id] = request_count_dict.get(id, 0) + total_tokens await self.router_cache.async_set_cache( key=tpm_key, value=request_count_dict, ttl=self.routing_args.ttl ) ## RPM request_count_dict = ( await self.router_cache.async_get_cache(key=rpm_key) or {} ) request_count_dict[id] = request_count_dict.get(id, 0) + 1 await self.router_cache.async_set_cache( key=rpm_key, value=request_count_dict, ttl=self.routing_args.ttl ) ### TESTING ### if self.test_flag: self.logged_success += 1 except Exception as e: verbose_router_logger.error( "litellm.router_strategy.lowest_tpm_rpm.py::async_log_success_event(): Exception occured - {}".format( str(e) ) ) verbose_router_logger.debug(traceback.format_exc()) pass def get_available_deployments( # noqa: PLR0915 self, model_group: str, healthy_deployments: list, messages: Optional[List[Dict[str, str]]] = None, input: Optional[Union[str, List]] = None, ): """ Returns a deployment with the lowest TPM/RPM usage. """ # get list of potential deployments verbose_router_logger.debug( f"get_available_deployments - Usage Based. model_group: {model_group}, healthy_deployments: {healthy_deployments}" ) current_minute = datetime.now().strftime("%H-%M") tpm_key = f"{model_group}:tpm:{current_minute}" rpm_key = f"{model_group}:rpm:{current_minute}" tpm_dict = self.router_cache.get_cache(key=tpm_key) rpm_dict = self.router_cache.get_cache(key=rpm_key) verbose_router_logger.debug( f"tpm_key={tpm_key}, tpm_dict: {tpm_dict}, rpm_dict: {rpm_dict}" ) try: input_tokens = token_counter(messages=messages, text=input) except Exception: input_tokens = 0 verbose_router_logger.debug(f"input_tokens={input_tokens}") # ----------------------- # Find lowest used model # ---------------------- lowest_tpm = float("inf") if tpm_dict is None: # base case - none of the deployments have been used # initialize a tpm dict with {model_id: 0} tpm_dict = {} for deployment in healthy_deployments: tpm_dict[deployment["model_info"]["id"]] = 0 else: for d in healthy_deployments: ## if healthy deployment not yet used if d["model_info"]["id"] not in tpm_dict: tpm_dict[d["model_info"]["id"]] = 0 all_deployments = tpm_dict deployment = None for item, item_tpm in all_deployments.items(): ## get the item from model list _deployment = None for m in healthy_deployments: if item == m["model_info"]["id"]: _deployment = m if _deployment is None: continue # skip to next one _deployment_tpm = None if _deployment_tpm is None: _deployment_tpm = _deployment.get("tpm") if _deployment_tpm is None: _deployment_tpm = _deployment.get("litellm_params", {}).get("tpm") if _deployment_tpm is None: _deployment_tpm = _deployment.get("model_info", {}).get("tpm") if _deployment_tpm is None: _deployment_tpm = float("inf") _deployment_rpm = None if _deployment_rpm is None: _deployment_rpm = _deployment.get("rpm") if _deployment_rpm is None: _deployment_rpm = _deployment.get("litellm_params", {}).get("rpm") if _deployment_rpm is None: _deployment_rpm = _deployment.get("model_info", {}).get("rpm") if _deployment_rpm is None: _deployment_rpm = float("inf") if item_tpm + input_tokens > _deployment_tpm: continue elif (rpm_dict is not None and item in rpm_dict) and ( rpm_dict[item] + 1 >= _deployment_rpm ): continue elif item_tpm < lowest_tpm: lowest_tpm = item_tpm deployment = _deployment print_verbose("returning picked lowest tpm/rpm deployment.") return deployment