chore: initial public snapshot for github upload
This commit is contained in:
Your Name
@@ -0,0 +1,249 @@
|
||||
#### 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, routing_args: dict = {}):
|
||||
self.router_cache = router_cache
|
||||
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 "litellm_params" not in kwargs or kwargs["litellm_params"] is None:
|
||||
return
|
||||
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 "litellm_params" not in kwargs or kwargs["litellm_params"] is None:
|
||||
return
|
||||
if kwargs["litellm_params"].get("metadata") is None:
|
||||
pass
|
||||
else:
|
||||
model_group = kwargs["litellm_params"]["metadata"].get(
|
||||
"model_group", None
|
||||
)
|
||||
|
||||
model_info = kwargs["litellm_params"].get("model_info")
|
||||
id = None
|
||||
if model_info is not None and isinstance(model_info, dict):
|
||||
id = model_info.get("id", None)
|
||||
if model_group is None or id is None:
|
||||
return
|
||||
elif isinstance(id, int):
|
||||
id = str(id)
|
||||
|
||||
if "usage" not in response_obj:
|
||||
return
|
||||
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.exception(
|
||||
"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
|
||||
Reference in New Issue
Block a user