From 224460d4c9950f3a8ada9624080a44d03cc5f976 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Fri, 4 Oct 2024 16:56:22 +0530
Subject: [PATCH] fix prometheus track cooldown events on custom logger (#6060)

---
 litellm/proxy/proxy_config.yaml            | 12 +++++++++---
 litellm/router_utils/cooldown_callbacks.py | 14 +++++++++-----
 2 files changed, 18 insertions(+), 8 deletions(-)

diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml
index d9b275ae8..d611aa87b 100644
--- a/litellm/proxy/proxy_config.yaml
+++ b/litellm/proxy/proxy_config.yaml
@@ -3,8 +3,14 @@ model_list:
     litellm_params:
       model: openai/gpt-5
       api_key: fake-key
-      api_base: https://exampleopenaiendpoint-production.up.railway.app/
-
-
+      api_base: https://exampleopenaiendpoint-production.up.railwaz.app/
+  - model_name: db-openai-endpoint
+    litellm_params:
+      model: openai/gpt-5
+      api_key: fake-key
+      api_base: https://exampleopenaiendpoint-production.up.railwxaz.app/
+
+litellm_settings:
+  callbacks: ["prometheus"]
 
 
diff --git a/litellm/router_utils/cooldown_callbacks.py b/litellm/router_utils/cooldown_callbacks.py
index a1bbf929f..d1c2431fd 100644
--- a/litellm/router_utils/cooldown_callbacks.py
+++ b/litellm/router_utils/cooldown_callbacks.py
@@ -3,7 +3,7 @@ Callbacks triggered on cooling down deployments
 """
 
 import copy
-from typing import TYPE_CHECKING, Any, Union
+from typing import TYPE_CHECKING, Any, Optional, Union
 
 import litellm
 from litellm._logging import verbose_logger
@@ -57,12 +57,16 @@ async def router_cooldown_event_callback(
         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,
+    )
 
-    prometheusLogger = None
-    for callback in litellm.callbacks:
-        if isinstance(callback, PrometheusLogger):
-            prometheusLogger = callback
+    # get the prometheus logger from in memory loggers
+    prometheusLogger: Optional[CustomLogger] = get_custom_logger_compatible_class(
+        logging_integration="prometheus",
+    )
 
     if prometheusLogger is not None: