mirror of
https://github.com/BerriAI/litellm.git
synced 2025-04-25 18:54:30 +00:00
d57be47b0f
* ci(config.yml): add a 'check_code_quality' step Addresses https://github.com/BerriAI/litellm/issues/5991 * ci(config.yml): check why circle ci doesn't pick up this test * ci(config.yml): fix to run 'check_code_quality' tests * fix(__init__.py): fix unprotected import * fix(__init__.py): don't remove unused imports * build(ruff.toml): update ruff.toml to ignore unused imports * fix: fix: ruff + pyright - fix linting + type-checking errors * fix: fix linting errors * fix(lago.py): fix module init error * fix: fix linting errors * ci(config.yml): cd into correct dir for checks * fix(proxy_server.py): fix linting error * fix(utils.py): fix bare except causes ruff linting errors * fix: ruff - fix remaining linting errors * fix(clickhouse.py): use standard logging object * fix(__init__.py): fix unprotected import * fix: ruff - fix linting errors * fix: fix linting errors * ci(config.yml): cleanup code qa step (formatting handled in local_testing) * fix(_health_endpoints.py): fix ruff linting errors * ci(config.yml): just use ruff in check_code_quality pipeline for now * build(custom_guardrail.py): include missing file * style(embedding_handler.py): fix ruff check
221 lines
8.5 KiB
Python
221 lines
8.5 KiB
Python
#### What this does ####
|
|
# identifies least busy deployment
|
|
# How is this achieved?
|
|
# - Before each call, have the router print the state of requests {"deployment": "requests_in_flight"}
|
|
# - use litellm.input_callbacks to log when a request is just about to be made to a model - {"deployment-id": traffic}
|
|
# - use litellm.success + failure callbacks to log when a request completed
|
|
# - in get_available_deployment, for a given model group name -> pick based on traffic
|
|
|
|
import os
|
|
import random
|
|
import traceback
|
|
from typing import Optional
|
|
|
|
import dotenv # type: ignore
|
|
import requests
|
|
|
|
from litellm.caching import DualCache
|
|
from litellm.integrations.custom_logger import CustomLogger
|
|
|
|
|
|
class LeastBusyLoggingHandler(CustomLogger):
|
|
test_flag: bool = False
|
|
logged_success: int = 0
|
|
logged_failure: int = 0
|
|
|
|
def __init__(self, router_cache: DualCache, model_list: list):
|
|
self.router_cache = router_cache
|
|
self.mapping_deployment_to_id: dict = {}
|
|
self.model_list = model_list
|
|
|
|
def log_pre_api_call(self, model, messages, kwargs):
|
|
"""
|
|
Log when a model is being used.
|
|
|
|
Caching based on model group.
|
|
"""
|
|
try:
|
|
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)
|
|
|
|
request_count_api_key = f"{model_group}_request_count"
|
|
# update cache
|
|
request_count_dict = (
|
|
self.router_cache.get_cache(key=request_count_api_key) or {}
|
|
)
|
|
request_count_dict[id] = request_count_dict.get(id, 0) + 1
|
|
|
|
self.router_cache.set_cache(
|
|
key=request_count_api_key, value=request_count_dict
|
|
)
|
|
except Exception:
|
|
pass
|
|
|
|
def log_success_event(self, kwargs, response_obj, start_time, end_time):
|
|
try:
|
|
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)
|
|
|
|
request_count_api_key = f"{model_group}_request_count"
|
|
# decrement count in cache
|
|
request_count_dict = (
|
|
self.router_cache.get_cache(key=request_count_api_key) or {}
|
|
)
|
|
request_count_value: Optional[int] = request_count_dict.get(id, 0)
|
|
if request_count_value is None:
|
|
return
|
|
request_count_dict[id] = request_count_value - 1
|
|
self.router_cache.set_cache(
|
|
key=request_count_api_key, value=request_count_dict
|
|
)
|
|
|
|
### TESTING ###
|
|
if self.test_flag:
|
|
self.logged_success += 1
|
|
except Exception:
|
|
pass
|
|
|
|
def log_failure_event(self, kwargs, response_obj, start_time, end_time):
|
|
try:
|
|
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)
|
|
|
|
request_count_api_key = f"{model_group}_request_count"
|
|
# decrement count in cache
|
|
request_count_dict = (
|
|
self.router_cache.get_cache(key=request_count_api_key) or {}
|
|
)
|
|
request_count_value: Optional[int] = request_count_dict.get(id, 0)
|
|
if request_count_value is None:
|
|
return
|
|
request_count_dict[id] = request_count_value - 1
|
|
self.router_cache.set_cache(
|
|
key=request_count_api_key, value=request_count_dict
|
|
)
|
|
|
|
### TESTING ###
|
|
if self.test_flag:
|
|
self.logged_failure += 1
|
|
except Exception:
|
|
pass
|
|
|
|
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
|
|
try:
|
|
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)
|
|
|
|
request_count_api_key = f"{model_group}_request_count"
|
|
# decrement count in cache
|
|
request_count_dict = (
|
|
self.router_cache.get_cache(key=request_count_api_key) or {}
|
|
)
|
|
request_count_value: Optional[int] = request_count_dict.get(id, 0)
|
|
if request_count_value is None:
|
|
return
|
|
request_count_dict[id] = request_count_value - 1
|
|
self.router_cache.set_cache(
|
|
key=request_count_api_key, value=request_count_dict
|
|
)
|
|
|
|
### TESTING ###
|
|
if self.test_flag:
|
|
self.logged_success += 1
|
|
except Exception:
|
|
pass
|
|
|
|
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
|
|
try:
|
|
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)
|
|
|
|
request_count_api_key = f"{model_group}_request_count"
|
|
# decrement count in cache
|
|
request_count_dict = (
|
|
self.router_cache.get_cache(key=request_count_api_key) or {}
|
|
)
|
|
request_count_value: Optional[int] = request_count_dict.get(id, 0)
|
|
if request_count_value is None:
|
|
return
|
|
request_count_dict[id] = request_count_value - 1
|
|
self.router_cache.set_cache(
|
|
key=request_count_api_key, value=request_count_dict
|
|
)
|
|
|
|
### TESTING ###
|
|
if self.test_flag:
|
|
self.logged_failure += 1
|
|
except Exception:
|
|
pass
|
|
|
|
def get_available_deployments(self, model_group: str, healthy_deployments: list):
|
|
request_count_api_key = f"{model_group}_request_count"
|
|
deployments = self.router_cache.get_cache(key=request_count_api_key) or {}
|
|
all_deployments = deployments
|
|
for d in healthy_deployments:
|
|
## if healthy deployment not yet used
|
|
if d["model_info"]["id"] not in all_deployments:
|
|
all_deployments[d["model_info"]["id"]] = 0
|
|
# map deployment to id
|
|
# pick least busy deployment
|
|
min_traffic = float("inf")
|
|
min_deployment = None
|
|
for k, v in all_deployments.items():
|
|
if v < min_traffic:
|
|
min_traffic = v
|
|
min_deployment = k
|
|
if min_deployment is not None:
|
|
## check if min deployment is a string, if so, cast it to int
|
|
for m in healthy_deployments:
|
|
if m["model_info"]["id"] == min_deployment:
|
|
return m
|
|
min_deployment = random.choice(healthy_deployments)
|
|
else:
|
|
min_deployment = random.choice(healthy_deployments)
|
|
return min_deployment
|