File size: 3,098 Bytes
e3278e4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 |
"""
Callbacks triggered on cooling down deployments
"""
import copy
from typing import TYPE_CHECKING, Any, Optional, Union
import litellm
from litellm._logging import verbose_logger
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(
litellm_router_instance: LitellmRouter,
deployment_id: str,
exception_status: Union[str, int],
cooldown_time: float,
):
"""
Callback triggered when a deployment is put into cooldown by litellm
- Updates deployment state on Prometheus
- Increments cooldown metric for deployment on Prometheus
"""
verbose_logger.debug("In router_cooldown_event_callback - updating prometheus")
_deployment = litellm_router_instance.get_deployment(model_id=deployment_id)
if _deployment is None:
verbose_logger.warning(
f"in router_cooldown_event_callback but _deployment is None for deployment_id={deployment_id}. Doing nothing"
)
return
_litellm_params = _deployment["litellm_params"]
temp_litellm_params = copy.deepcopy(_litellm_params)
temp_litellm_params = dict(temp_litellm_params)
_model_name = _deployment.get("model_name", None) or ""
_api_base = (
litellm.get_api_base(model=_model_name, optional_params=temp_litellm_params)
or ""
)
model_info = _deployment["model_info"]
model_id = model_info.id
litellm_model_name = temp_litellm_params.get("model") or ""
llm_provider = ""
try:
_, llm_provider, _, _ = litellm.get_llm_provider(
model=litellm_model_name,
custom_llm_provider=temp_litellm_params.get("custom_llm_provider"),
)
except Exception:
pass
# get the prometheus logger from in memory loggers
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,
)
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._async_success_callback:
if isinstance(_callback, PrometheusLogger):
return _callback
for global_callback in litellm.callbacks:
if isinstance(global_callback, PrometheusLogger):
return global_callback
return None
|