(testing) add router unit testing for send_llm_exception_alert , router_cooldown_event_callback , cooldown utils (#6258)

* add router unit testing for send_llm_exception_alert

* test router_cooldown_event_callback

* test test_router_cooldown_event_callback_no_prometheus

* test_router_cooldown_event_callback_no_deployment

* test_router_cooldown_event_callback_no_deployment

* add testing for test_should_run_cooldown_logic

* test_increment_deployment_successes_for_current_minute_does_not_write_to_redis

* test test_should_cooldown_deployment_allowed_fails_set_on_router

* use helper for _is_allowed_fails_set_on_router

* add complete testing for cooldown utils

* move router unit tests

* move router handle error

* fix test_send_llm_exception_alert_no_logger
This commit is contained in:
Ishaan Jaff 2024-10-16 23:19:51 +05:30 committed by GitHub
parent dee6de0105
commit 3ab2b86062
5 changed files with 564 additions and 26 deletions

View file

@ -12,8 +12,10 @@ if TYPE_CHECKING:
from litellm.router import Router as _Router
LitellmRouter = _Router
from litellm.integrations.prometheus import PrometheusLogger
else:
LitellmRouter = Any
PrometheusLogger = Any
async def router_cooldown_event_callback(
@ -56,34 +58,38 @@ async def router_cooldown_event_callback(
except Exception:
pass
# Trigger cooldown on Prometheus
from litellm.integrations.custom_logger import CustomLogger
from litellm.integrations.prometheus import PrometheusLogger
from litellm.litellm_core_utils.litellm_logging import (
get_custom_logger_compatible_class,
)
# get the prometheus logger from in memory loggers
prometheusLogger: Optional[CustomLogger] = get_custom_logger_compatible_class(
logging_integration="prometheus",
prometheusLogger: Optional[PrometheusLogger] = (
_get_prometheus_logger_from_callbacks()
)
if prometheusLogger is not None:
prometheusLogger.set_deployment_complete_outage(
litellm_model_name=_model_name,
model_id=model_id,
api_base=_api_base,
api_provider=llm_provider,
)
if isinstance(prometheusLogger, PrometheusLogger):
prometheusLogger.set_deployment_complete_outage(
litellm_model_name=_model_name,
model_id=model_id,
api_base=_api_base,
api_provider=llm_provider,
)
prometheusLogger.increment_deployment_cooled_down(
litellm_model_name=_model_name,
model_id=model_id,
api_base=_api_base,
api_provider=llm_provider,
exception_status=str(exception_status),
)
prometheusLogger.increment_deployment_cooled_down(
litellm_model_name=_model_name,
model_id=model_id,
api_base=_api_base,
api_provider=llm_provider,
exception_status=str(exception_status),
)
return
def _get_prometheus_logger_from_callbacks() -> Optional[PrometheusLogger]:
"""
Checks if prometheus is a initalized callback, if yes returns it
"""
from litellm.integrations.prometheus import PrometheusLogger
for _callback in litellm.callbacks:
if isinstance(_callback, PrometheusLogger):
return _callback
return None