From bcdb53920af27fb98602426901a35be1231bb3b8 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 17 Jun 2025 14:44:24 -0700 Subject: [PATCH] [Fix] Bug Fix for using prom metrics config (#11779) * fixes for using config for prom metrics * test_set_llm_deployment_success_metrics_with_label_filtering * fixes for deployment failure * fix code qa checks * test_async_post_call_success_hook --- litellm/integrations/prometheus.py | 177 ++++++++---------- litellm/types/integrations/prometheus.py | 30 +++ .../test_prometheus_unit_tests.py | 173 +++++++++++++++-- 3 files changed, 263 insertions(+), 117 deletions(-) diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 6e096ce5b36..7da154d31e5 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -40,7 +40,7 @@ class PrometheusLogger(CustomLogger): from prometheus_client import Counter, Gauge, Histogram from litellm.proxy.proxy_server import CommonProxyErrors, premium_user - + # Always initialize label_filters, even for non-premium users self.label_filters = self._parse_prometheus_config() @@ -1070,6 +1070,31 @@ class PrometheusLogger(CustomLogger): llm_provider = _litellm_params.get("custom_llm_provider", None) + # Create enum_values for the label factory (always create for use in different metrics) + enum_values = UserAPIKeyLabelValues( + litellm_model_name=litellm_model_name, + model_id=model_id, + api_base=api_base, + api_provider=llm_provider, + exception_status=( + str(getattr(exception, "status_code", None)) if exception else None + ), + exception_class=( + self._get_exception_class_name(exception) if exception else None + ), + requested_model=model_group, + hashed_api_key=standard_logging_payload["metadata"][ + "user_api_key_hash" + ], + api_key_alias=standard_logging_payload["metadata"][ + "user_api_key_alias" + ], + team=standard_logging_payload["metadata"]["user_api_key_team_id"], + team_alias=standard_logging_payload["metadata"][ + "user_api_key_team_alias" + ], + ) + """ log these labels ["litellm_model_name", "model_id", "api_base", "api_provider"] @@ -1081,25 +1106,14 @@ class PrometheusLogger(CustomLogger): api_provider=llm_provider or "", ) if exception is not None: - self.litellm_deployment_failure_responses.labels( - litellm_model_name=litellm_model_name, - model_id=model_id, - api_base=api_base, - api_provider=llm_provider, - exception_status=str(getattr(exception, "status_code", None)), - exception_class=self._get_exception_class_name(exception), - requested_model=model_group, - hashed_api_key=standard_logging_payload["metadata"][ - "user_api_key_hash" - ], - api_key_alias=standard_logging_payload["metadata"][ - "user_api_key_alias" - ], - team=standard_logging_payload["metadata"]["user_api_key_team_id"], - team_alias=standard_logging_payload["metadata"][ - "user_api_key_team_alias" - ], - ).inc() + + _labels = prometheus_label_factory( + supported_enum_labels=self.get_labels_for_metric( + metric_name="litellm_deployment_failure_responses" + ), + enum_values=enum_values, + ) + self.litellm_deployment_failure_responses.labels(**_labels).inc() # tag based tracking if standard_logging_payload is not None and isinstance( @@ -1122,23 +1136,13 @@ class PrometheusLogger(CustomLogger): } ).inc() - self.litellm_deployment_total_requests.labels( - litellm_model_name=litellm_model_name, - model_id=model_id, - api_base=api_base, - api_provider=llm_provider, - requested_model=model_group, - hashed_api_key=standard_logging_payload["metadata"][ - "user_api_key_hash" - ], - api_key_alias=standard_logging_payload["metadata"][ - "user_api_key_alias" - ], - team=standard_logging_payload["metadata"]["user_api_key_team_id"], - team_alias=standard_logging_payload["metadata"][ - "user_api_key_team_alias" - ], - ).inc() + _labels = prometheus_label_factory( + supported_enum_labels=self.get_labels_for_metric( + metric_name="litellm_deployment_total_requests" + ), + enum_values=enum_values, + ) + self.litellm_deployment_total_requests.labels(**_labels).inc() pass except Exception as e: @@ -1165,9 +1169,7 @@ class PrometheusLogger(CustomLogger): if standard_logging_payload is None: return - model_group = standard_logging_payload["model_group"] api_base = standard_logging_payload["api_base"] - _response_headers = request_kwargs.get("response_headers") _litellm_params = request_kwargs.get("litellm_params", {}) or {} _metadata = _litellm_params.get("metadata", {}) litellm_model_name = request_kwargs.get("model", None) @@ -1191,14 +1193,13 @@ class PrometheusLogger(CustomLogger): if litellm_overhead_time_ms := standard_logging_payload[ "hidden_params" ].get("litellm_overhead_time_ms"): - self.litellm_overhead_latency_metric.labels( - model_group, - llm_provider, - api_base, - litellm_model_name, - standard_logging_payload["metadata"]["user_api_key_hash"], - standard_logging_payload["metadata"]["user_api_key_alias"], - ).observe( + _labels = prometheus_label_factory( + supported_enum_labels=self.get_labels_for_metric( + metric_name="litellm_overhead_latency_metric" + ), + enum_values=enum_values, + ) + self.litellm_overhead_latency_metric.labels(**_labels).observe( litellm_overhead_time_ms / 1000 ) # set as seconds @@ -1209,24 +1210,26 @@ class PrometheusLogger(CustomLogger): "api_base", "litellm_model_name" """ - self.litellm_remaining_requests_metric.labels( - model_group, - llm_provider, - api_base, - litellm_model_name, - standard_logging_payload["metadata"]["user_api_key_hash"], - standard_logging_payload["metadata"]["user_api_key_alias"], - ).set(remaining_requests) + _labels = prometheus_label_factory( + supported_enum_labels=self.get_labels_for_metric( + metric_name="litellm_remaining_requests_metric" + ), + enum_values=enum_values, + ) + self.litellm_remaining_requests_metric.labels(**_labels).set( + remaining_requests + ) if remaining_tokens: - self.litellm_remaining_tokens_metric.labels( - model_group, - llm_provider, - api_base, - litellm_model_name, - standard_logging_payload["metadata"]["user_api_key_hash"], - standard_logging_payload["metadata"]["user_api_key_alias"], - ).set(remaining_tokens) + _labels = prometheus_label_factory( + supported_enum_labels=self.get_labels_for_metric( + metric_name="litellm_remaining_tokens_metric" + ), + enum_values=enum_values, + ) + self.litellm_remaining_tokens_metric.labels(**_labels).set( + remaining_tokens + ) """ log these labels @@ -1239,41 +1242,21 @@ class PrometheusLogger(CustomLogger): api_provider=llm_provider or "", ) - self.litellm_deployment_success_responses.labels( - litellm_model_name=litellm_model_name, - model_id=model_id, - api_base=api_base, - api_provider=llm_provider, - requested_model=model_group, - hashed_api_key=standard_logging_payload["metadata"][ - "user_api_key_hash" - ], - api_key_alias=standard_logging_payload["metadata"][ - "user_api_key_alias" - ], - team=standard_logging_payload["metadata"]["user_api_key_team_id"], - team_alias=standard_logging_payload["metadata"][ - "user_api_key_team_alias" - ], - ).inc() + _labels = prometheus_label_factory( + supported_enum_labels=self.get_labels_for_metric( + metric_name="litellm_deployment_success_responses" + ), + enum_values=enum_values, + ) + self.litellm_deployment_success_responses.labels(**_labels).inc() - self.litellm_deployment_total_requests.labels( - litellm_model_name=litellm_model_name, - model_id=model_id, - api_base=api_base, - api_provider=llm_provider, - requested_model=model_group, - hashed_api_key=standard_logging_payload["metadata"][ - "user_api_key_hash" - ], - api_key_alias=standard_logging_payload["metadata"][ - "user_api_key_alias" - ], - team=standard_logging_payload["metadata"]["user_api_key_team_id"], - team_alias=standard_logging_payload["metadata"][ - "user_api_key_team_alias" - ], - ).inc() + _labels = prometheus_label_factory( + supported_enum_labels=self.get_labels_for_metric( + metric_name="litellm_deployment_total_requests" + ), + enum_values=enum_values, + ) + self.litellm_deployment_total_requests.labels(**_labels).inc() # Track deployment Latency response_ms: timedelta = end_time - start_time diff --git a/litellm/types/integrations/prometheus.py b/litellm/types/integrations/prometheus.py index 13714871b86..358ff3d50bd 100644 --- a/litellm/types/integrations/prometheus.py +++ b/litellm/types/integrations/prometheus.py @@ -76,6 +76,9 @@ class UserAPIKeyLabelNames(Enum): DEFINED_PROMETHEUS_METRICS = Literal[ "litellm_llm_api_latency_metric", "litellm_request_total_latency_metric", + "litellm_overhead_latency_metric", + "litellm_remaining_requests_metric", + "litellm_remaining_tokens_metric", "litellm_proxy_total_requests_metric", "litellm_proxy_failed_requests_metric", "litellm_deployment_latency_per_output_token", @@ -157,6 +160,33 @@ class PrometheusMetricLabels: UserAPIKeyLabelNames.TEAM_ALIAS.value, ] + litellm_overhead_latency_metric = [ + UserAPIKeyLabelNames.REQUESTED_MODEL.value, + UserAPIKeyLabelNames.API_PROVIDER.value, + UserAPIKeyLabelNames.API_BASE.value, + UserAPIKeyLabelNames.v2_LITELLM_MODEL_NAME.value, + UserAPIKeyLabelNames.API_KEY_HASH.value, + UserAPIKeyLabelNames.API_KEY_ALIAS.value, + ] + + litellm_remaining_requests_metric = [ + UserAPIKeyLabelNames.REQUESTED_MODEL.value, + UserAPIKeyLabelNames.API_PROVIDER.value, + UserAPIKeyLabelNames.API_BASE.value, + UserAPIKeyLabelNames.v2_LITELLM_MODEL_NAME.value, + UserAPIKeyLabelNames.API_KEY_HASH.value, + UserAPIKeyLabelNames.API_KEY_ALIAS.value, + ] + + litellm_remaining_tokens_metric = [ + UserAPIKeyLabelNames.REQUESTED_MODEL.value, + UserAPIKeyLabelNames.API_PROVIDER.value, + UserAPIKeyLabelNames.API_BASE.value, + UserAPIKeyLabelNames.v2_LITELLM_MODEL_NAME.value, + UserAPIKeyLabelNames.API_KEY_HASH.value, + UserAPIKeyLabelNames.API_KEY_ALIAS.value, + ] + litellm_requests_metric = [ UserAPIKeyLabelNames.END_USER.value, UserAPIKeyLabelNames.API_KEY_HASH.value, diff --git a/tests/logging_callback_tests/test_prometheus_unit_tests.py b/tests/logging_callback_tests/test_prometheus_unit_tests.py index 5b40166286f..119fe3cd46a 100644 --- a/tests/logging_callback_tests/test_prometheus_unit_tests.py +++ b/tests/logging_callback_tests/test_prometheus_unit_tests.py @@ -797,6 +797,7 @@ def test_set_llm_deployment_success_metrics(prometheus_logger): "x_ratelimit_remaining_requests": 123, "x_ratelimit_remaining_tokens": 4321, } + standard_logging_payload["model_group"] = "my_custom_model_group" standard_logging_payload["hidden_params"]["litellm_overhead_time_ms"] = 100 # Create test data @@ -810,6 +811,7 @@ def test_set_llm_deployment_success_metrics(prometheus_logger): } enum_values = UserAPIKeyLabelValues( + requested_model=standard_logging_payload["model_group"], litellm_model_name=standard_logging_payload["model"], api_provider=standard_logging_payload["custom_llm_provider"], hashed_api_key=standard_logging_payload["metadata"]["user_api_key_hash"], @@ -834,26 +836,28 @@ def test_set_llm_deployment_success_metrics(prometheus_logger): # Verify remaining requests metric prometheus_logger.litellm_remaining_requests_metric.labels.assert_called_once_with( - "openai-gpt", # model_group / requested model from create_standard_logging_payload() - "openai", # llm provider - "https://api.openai.com", # api base - "gpt-3.5-turbo", # actual model used - litellm model name - standard_logging_payload["metadata"]["user_api_key_hash"], - standard_logging_payload["metadata"]["user_api_key_alias"], + api_base="https://api.openai.com", + api_key_alias=standard_logging_payload["metadata"]["user_api_key_alias"], + api_provider="openai", + hashed_api_key=standard_logging_payload["metadata"]["user_api_key_hash"], + litellm_model_name="gpt-3.5-turbo", + requested_model="my_custom_model_group", ) + prometheus_logger.litellm_remaining_requests_metric.labels().set.assert_called_once_with( 123 ) # Verify remaining tokens metric prometheus_logger.litellm_remaining_tokens_metric.labels.assert_called_once_with( - "openai-gpt", # model_group / requested model from create_standard_logging_payload() - "openai", # llm provider - "https://api.openai.com", # api base - "gpt-3.5-turbo", # actual model used - litellm model name - standard_logging_payload["metadata"]["user_api_key_hash"], - standard_logging_payload["metadata"]["user_api_key_alias"], + api_base="https://api.openai.com", + api_key_alias=standard_logging_payload["metadata"]["user_api_key_alias"], + api_provider="openai", + hashed_api_key=standard_logging_payload["metadata"]["user_api_key_hash"], + litellm_model_name="gpt-3.5-turbo", + requested_model="my_custom_model_group", ) + prometheus_logger.litellm_remaining_tokens_metric.labels().set.assert_called_once_with( 4321 ) @@ -872,7 +876,7 @@ def test_set_llm_deployment_success_metrics(prometheus_logger): model_id="model-123", api_base="https://api.openai.com", api_provider="openai", - requested_model="openai-gpt", # requested model from create_standard_logging_payload() + requested_model="my_custom_model_group", hashed_api_key=standard_logging_payload["metadata"]["user_api_key_hash"], api_key_alias=standard_logging_payload["metadata"]["user_api_key_alias"], team=standard_logging_payload["metadata"]["user_api_key_team_id"], @@ -886,7 +890,7 @@ def test_set_llm_deployment_success_metrics(prometheus_logger): model_id="model-123", api_base="https://api.openai.com", api_provider="openai", - requested_model="openai-gpt", # requested model from create_standard_logging_payload() + requested_model="my_custom_model_group", hashed_api_key=standard_logging_payload["metadata"]["user_api_key_hash"], api_key_alias=standard_logging_payload["metadata"]["user_api_key_alias"], team=standard_logging_payload["metadata"]["user_api_key_team_id"], @@ -906,12 +910,12 @@ def test_set_llm_deployment_success_metrics(prometheus_logger): team_alias=standard_logging_payload["metadata"]["user_api_key_team_alias"], ) prometheus_logger.litellm_overhead_latency_metric.labels.assert_called_once_with( - "openai-gpt", # model_group / requested model from create_standard_logging_payload() - "openai", # llm provider - "https://api.openai.com", # api base - "gpt-3.5-turbo", # actual model used - litellm model name - standard_logging_payload["metadata"]["user_api_key_hash"], - standard_logging_payload["metadata"]["user_api_key_alias"], + api_base="https://api.openai.com", + api_key_alias=standard_logging_payload["metadata"]["user_api_key_alias"], + api_provider="openai", + hashed_api_key=standard_logging_payload["metadata"]["user_api_key_hash"], + litellm_model_name="gpt-3.5-turbo", + requested_model="my_custom_model_group", ) # Calculate expected latency per token (1 second / 10 tokens = 0.1 seconds per token) @@ -1517,3 +1521,132 @@ def test_get_exception_class_name(prometheus_logger): ) assert prometheus_logger._get_exception_class_name(context_window_error) == "ContextWindowExceededError" + +def test_set_llm_deployment_success_metrics_with_label_filtering(): + """ + Test that set_llm_deployment_success_metrics correctly uses prometheus_label_factory + and respects label filtering configuration to prevent "Incorrect label names" errors. + """ + from litellm.types.integrations.prometheus import PrometheusMetricsConfig + + # Create a prometheus logger with label filtering configuration + config = [ + PrometheusMetricsConfig( + group="test_group", + metrics=[ + "litellm_overhead_latency_metric", + "litellm_remaining_requests_metric", + "litellm_remaining_tokens_metric", + "litellm_deployment_success_responses", + "litellm_deployment_total_requests" + ], + include_labels=["requested_model", "api_provider", "hashed_api_key"] # Limited labels + ) + ] + + # Mock litellm.prometheus_metrics_config + with patch('litellm.prometheus_metrics_config', config): + # Clear registry before creating new logger + collectors = list(REGISTRY._collector_to_names.keys()) + for collector in collectors: + REGISTRY.unregister(collector) + + prometheus_logger = PrometheusLogger() + + # Mock all the metrics used in the method + prometheus_logger.litellm_overhead_latency_metric = MagicMock() + prometheus_logger.litellm_remaining_requests_metric = MagicMock() + prometheus_logger.litellm_remaining_tokens_metric = MagicMock() + prometheus_logger.litellm_deployment_success_responses = MagicMock() + prometheus_logger.litellm_deployment_total_requests = MagicMock() + prometheus_logger.set_deployment_healthy = MagicMock() + + # Create standard logging payload + standard_logging_payload = create_standard_logging_payload() + standard_logging_payload["hidden_params"]["additional_headers"] = { + "x_ratelimit_remaining_requests": 123, + "x_ratelimit_remaining_tokens": 4321, + } + standard_logging_payload["hidden_params"]["litellm_overhead_time_ms"] = 100 + + # Create test data + request_kwargs = { + "model": "gpt-3.5-turbo", + "litellm_params": { + "custom_llm_provider": "openai", + "metadata": {"model_info": {"id": "model-123"}}, + }, + "standard_logging_object": standard_logging_payload, + } + + enum_values = UserAPIKeyLabelValues( + litellm_model_name=standard_logging_payload["model"], + api_provider=standard_logging_payload["custom_llm_provider"], + hashed_api_key=standard_logging_payload["metadata"]["user_api_key_hash"], + api_key_alias=standard_logging_payload["metadata"]["user_api_key_alias"], + team=standard_logging_payload["metadata"]["user_api_key_team_id"], + team_alias=standard_logging_payload["metadata"]["user_api_key_team_alias"], + requested_model=standard_logging_payload["model_group"], + model=standard_logging_payload["model"], + model_id=standard_logging_payload["model_id"], + api_base=standard_logging_payload["api_base"], + ) + + start_time = datetime.now() + end_time = start_time + timedelta(seconds=1) + output_tokens = 10 + + # Call the function - this should not raise "Incorrect label names" error + prometheus_logger.set_llm_deployment_success_metrics( + request_kwargs=request_kwargs, + start_time=start_time, + end_time=end_time, + output_tokens=output_tokens, + enum_values=enum_values, + ) + + # Verify that metrics were called with filtered labels (only the configured ones) + # The exact labels depend on what get_labels_for_metric returns for each metric + + # Verify overhead latency metric was called with filtered labels + prometheus_logger.litellm_overhead_latency_metric.labels.assert_called_once() + overhead_labels = prometheus_logger.litellm_overhead_latency_metric.labels.call_args[1] + + # Should only contain the filtered labels that are supported for this metric + expected_filtered_labels = {"requested_model", "api_provider", "hashed_api_key"} + actual_labels = set(k for k in overhead_labels.keys() if k is not None) + + # Verify that only expected labels are present (subset of configured labels) + assert actual_labels <= expected_filtered_labels + + # Verify remaining requests metric was called with filtered labels + prometheus_logger.litellm_remaining_requests_metric.labels.assert_called_once() + requests_labels = prometheus_logger.litellm_remaining_requests_metric.labels.call_args[1] + actual_labels = set(k for k in requests_labels.keys() if k is not None) + assert actual_labels <= expected_filtered_labels + + # Verify remaining tokens metric was called with filtered labels + prometheus_logger.litellm_remaining_tokens_metric.labels.assert_called_once() + tokens_labels = prometheus_logger.litellm_remaining_tokens_metric.labels.call_args[1] + actual_labels = set(k for k in tokens_labels.keys() if k is not None) + assert actual_labels <= expected_filtered_labels + + # Verify deployment success responses metric was called with filtered labels + prometheus_logger.litellm_deployment_success_responses.labels.assert_called_once() + success_labels = prometheus_logger.litellm_deployment_success_responses.labels.call_args[1] + actual_labels = set(k for k in success_labels.keys() if k is not None) + assert actual_labels <= expected_filtered_labels + + # Verify deployment total requests metric was called with filtered labels + prometheus_logger.litellm_deployment_total_requests.labels.assert_called_once() + total_labels = prometheus_logger.litellm_deployment_total_requests.labels.call_args[1] + actual_labels = set(total_labels.keys()) + assert actual_labels.issubset(expected_filtered_labels.union({None})) + + # Verify all metrics were actually called (no exceptions were raised) + prometheus_logger.litellm_overhead_latency_metric.labels().observe.assert_called_once() + prometheus_logger.litellm_remaining_requests_metric.labels().set.assert_called_once_with(123) + prometheus_logger.litellm_remaining_tokens_metric.labels().set.assert_called_once_with(4321) + prometheus_logger.litellm_deployment_success_responses.labels().inc.assert_called_once() + prometheus_logger.litellm_deployment_total_requests.labels().inc.assert_called_once() +