fix prometheus track cooldown events on custom logger (#6060)

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"]
 
 

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: