phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-26 11:14:04 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

refactor(router.py): code cleanup

Browse source
This commit is contained in:
Krrish Dholakia 2023-11-17 17:05:38 -08:00
parent 7d70bf84a7
commit 452946b2f8
1 changed files with 102 additions and 102 deletions
Download patch file Download diff file Expand all files Collapse all files

204
litellm/router.py
View file

@ -301,119 +301,119 @@ class Router:
data[k] = v data[k] = v
return await litellm.aembedding(**{**data, "input": input, "caching": self.cache_responses, **kwargs}) return await litellm.aembedding(**{**data, "input": input, "caching": self.cache_responses, **kwargs})
def deployment_callback( # def deployment_callback(
self, # self,
kwargs, # kwargs to completion # kwargs, # kwargs to completion
completion_response, # response from completion # completion_response, # response from completion
start_time, end_time # start/end time # start_time, end_time # start/end time
): # ):
""" # """
Function LiteLLM submits a callback to after a successful # Function LiteLLM submits a callback to after a successful
completion. Purpose of this is ti update TPM/RPM usage per model # completion. Purpose of this is ti update TPM/RPM usage per model
""" # """
model_name = kwargs.get('model', None) # i.e. gpt35turbo # model_name = kwargs.get('model', None) # i.e. gpt35turbo
custom_llm_provider = kwargs.get("litellm_params", {}).get('custom_llm_provider', None) # i.e. azure # custom_llm_provider = kwargs.get("litellm_params", {}).get('custom_llm_provider', None) # i.e. azure
if custom_llm_provider: # if custom_llm_provider:
model_name = f"{custom_llm_provider}/{model_name}" # model_name = f"{custom_llm_provider}/{model_name}"
total_tokens = completion_response['usage']['total_tokens'] # total_tokens = completion_response['usage']['total_tokens']
self._set_deployment_usage(model_name, total_tokens) # self._set_deployment_usage(model_name, total_tokens)
def get_available_deployment(self, # def get_available_deployment(self,
model: str, # model: str,
messages: Optional[List[Dict[str, str]]] = None, # messages: Optional[List[Dict[str, str]]] = None,
input: Optional[Union[str, List]] = None): # input: Optional[Union[str, List]] = None):
""" # """
Returns a deployment with the lowest TPM/RPM usage. # Returns a deployment with the lowest TPM/RPM usage.
""" # """
# get list of potential deployments # # get list of potential deployments
potential_deployments = [] # potential_deployments = []
for item in self.model_list: # for item in self.model_list:
if item["model_name"] == model: # if item["model_name"] == model:
potential_deployments.append(item) # potential_deployments.append(item)
# set first model as current model to calculate token count # # set first model as current model to calculate token count
deployment = potential_deployments[0] # deployment = potential_deployments[0]
# get encoding # # get encoding
token_count = 0 # token_count = 0
if messages is not None: # if messages is not None:
token_count = litellm.token_counter(model=deployment["model_name"], messages=messages) # token_count = litellm.token_counter(model=deployment["model_name"], messages=messages)
elif input is not None: # elif input is not None:
if isinstance(input, List): # if isinstance(input, List):
input_text = "".join(text for text in input) # input_text = "".join(text for text in input)
else: # else:
input_text = input # input_text = input
token_count = litellm.token_counter(model=deployment["model_name"], text=input_text) # token_count = litellm.token_counter(model=deployment["model_name"], text=input_text)
# ----------------------- # # -----------------------
# Find lowest used model # # Find lowest used model
# ---------------------- # # ----------------------
lowest_tpm = float("inf") # lowest_tpm = float("inf")
deployment = None # deployment = None
# Go through all the models to get tpm, rpm # # Go through all the models to get tpm, rpm
for item in potential_deployments: # for item in potential_deployments:
item_tpm, item_rpm = self._get_deployment_usage(deployment_name=item["litellm_params"]["model"]) # item_tpm, item_rpm = self._get_deployment_usage(deployment_name=item["litellm_params"]["model"])
if item_tpm == 0: # if item_tpm == 0:
return item # return item
elif item_tpm + token_count > item["tpm"] or item_rpm + 1 >= item["rpm"]: # elif item_tpm + token_count > item["tpm"] or item_rpm + 1 >= item["rpm"]:
continue # continue
elif item_tpm < lowest_tpm: # elif item_tpm < lowest_tpm:
lowest_tpm = item_tpm # lowest_tpm = item_tpm
deployment = item # deployment = item
# if none, raise exception # # if none, raise exception
if deployment is None: # if deployment is None:
raise ValueError("No models available.") # raise ValueError("No models available.")
# return model # # return model
return deployment # return deployment
def _get_deployment_usage( # def _get_deployment_usage(
self, # self,
deployment_name: str # deployment_name: str
): # ):
# ------------ # # ------------
# Setup values # # Setup values
# ------------ # # ------------
current_minute = datetime.now().strftime("%H-%M") # current_minute = datetime.now().strftime("%H-%M")
tpm_key = f'{deployment_name}:tpm:{current_minute}' # tpm_key = f'{deployment_name}:tpm:{current_minute}'
rpm_key = f'{deployment_name}:rpm:{current_minute}' # rpm_key = f'{deployment_name}:rpm:{current_minute}'
# ------------ # # ------------
# Return usage # # Return usage
# ------------ # # ------------
tpm = self.cache.get_cache(cache_key=tpm_key) or 0 # tpm = self.cache.get_cache(cache_key=tpm_key) or 0
rpm = self.cache.get_cache(cache_key=rpm_key) or 0 # rpm = self.cache.get_cache(cache_key=rpm_key) or 0
return int(tpm), int(rpm) # return int(tpm), int(rpm)
def increment(self, key: str, increment_value: int): # def increment(self, key: str, increment_value: int):
# get value # # get value
cached_value = self.cache.get_cache(cache_key=key) # cached_value = self.cache.get_cache(cache_key=key)
# update value # # update value
try: # try:
cached_value = cached_value + increment_value # cached_value = cached_value + increment_value
except: # except:
cached_value = increment_value # cached_value = increment_value
# save updated value # # save updated value
self.cache.add_cache(result=cached_value, cache_key=key, ttl=self.default_cache_time_seconds) # self.cache.add_cache(result=cached_value, cache_key=key, ttl=self.default_cache_time_seconds)
def _set_deployment_usage( # def _set_deployment_usage(
self, # self,
model_name: str, # model_name: str,
total_tokens: int # total_tokens: int
): # ):
# ------------ # # ------------
# Setup values # # Setup values
# ------------ # # ------------
current_minute = datetime.now().strftime("%H-%M") # current_minute = datetime.now().strftime("%H-%M")
tpm_key = f'{model_name}:tpm:{current_minute}' # tpm_key = f'{model_name}:tpm:{current_minute}'
rpm_key = f'{model_name}:rpm:{current_minute}' # rpm_key = f'{model_name}:rpm:{current_minute}'
# ------------ # # ------------
# Update usage # # Update usage
# ------------ # # ------------
self.increment(tpm_key, total_tokens) # self.increment(tpm_key, total_tokens)
self.increment(rpm_key, 1) # self.increment(rpm_key, 1)