use router_cooldown_handler

litellm/router_utils/cooldown_callbacks.py

@@ -0,0 +1,51 @@
+"""
+Callbacks triggered on cooling down deployments
+"""
+
+import copy
+from typing import TYPE_CHECKING, Any, Union
+
+import litellm
+from litellm._logging import verbose_logger
+
+if TYPE_CHECKING:
+    from litellm.router import Router as _Router
+
+    LitellmRouter = _Router
+else:
+    LitellmRouter = Any
+
+
+async def router_cooldown_handler(
+    litellm_router_instance: LitellmRouter,
+    deployment_id: str,
+    exception_status: Union[str, int],
+    cooldown_time: float,
+):
+    _deployment = litellm_router_instance.get_deployment(model_id=deployment_id)
+    if _deployment is None:
+        verbose_logger.warning(
+            f"in router_cooldown_handler 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)
+    _api_base = litellm.get_api_base(
+        model=_model_name, optional_params=temp_litellm_params
+    )
+    model_info = _deployment["model_info"]
+    model_id = model_info.id
+
+    # Trigger cooldown on Prometheus
+    from litellm.litellm_core_utils.litellm_logging import prometheusLogger
+
+    if prometheusLogger is not None:
+        prometheusLogger.set_llm_outage_metric(
+            litellm_model_name=_model_name,
+            model_id=model_id,
+            api_base="",
+            api_provider="",
+        )
+    pass