From bbd8ca3b3dbaf4d5ad536501477984811929ce5d Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Wed, 25 Mar 2026 16:17:57 +0530 Subject: [PATCH] feat(prometheus): add metrics for managed batch lifecycle - Add Prometheus metrics for managed batch and file operations - Track batch creation, file size, duration, and deletion events - Add CheckBatchCost polling metrics (jobs polled/processed, errors) - Record metrics in managed_files hook and check_batch_cost utility - Metrics include labels for model, provider, user, and status Made-with: Cursor --- .../proxy/common_utils/check_batch_cost.py | 51 ++++- .../proxy/hooks/managed_files.py | 68 ++++++ litellm/integrations/prometheus.py | 202 ++++++++++++++++++ litellm/types/integrations/prometheus.py | 64 ++++++ 4 files changed, 384 insertions(+), 1 deletion(-) diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py index cbe8d449b42..356f6ecd4b5 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -3,7 +3,7 @@ Polls LiteLLM_ManagedObjectTable to check if the batch job is complete, and if t """ from datetime import datetime, timedelta, timezone -from typing import TYPE_CHECKING, Optional +from typing import TYPE_CHECKING, List, Optional, Tuple from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid @@ -118,6 +118,15 @@ class CheckBatchCost: get_model_id_from_unified_batch_id, ) + try: + from litellm.integrations.prometheus import PrometheusLogger + prom_logger = PrometheusLogger.get_instance() + except Exception as e: + verbose_proxy_logger.error(f"CheckBatchCost: could not get Prometheus logger: {e}") + prom_logger = None + + processed_models: List[Tuple[Optional[str], Optional[str]]] = [] + try: await self._cleanup_stale_managed_objects() except Exception as cleanup_err: @@ -172,6 +181,8 @@ class CheckBatchCost: verbose_proxy_logger.info( f"Skipping job {unified_object_id} because it is not a valid unified object id" ) + if prom_logger: + prom_logger.record_check_batch_cost_error("invalid_unified_id") continue else: unified_object_id = decoded_unified_object_id @@ -183,6 +194,8 @@ class CheckBatchCost: verbose_proxy_logger.info( f"Skipping job {unified_object_id} because it is not a valid model id" ) + if prom_logger: + prom_logger.record_check_batch_cost_error("invalid_model_id") continue verbose_proxy_logger.info( @@ -202,6 +215,8 @@ class CheckBatchCost: verbose_proxy_logger.info( f"Skipping job {unified_object_id} because of error querying model ID: {model_id} for cost and usage of batch ID: {batch_id}: {e}" ) + if prom_logger: + prom_logger.record_check_batch_cost_error("provider_retrieval_error") continue ## RETRIEVE THE BATCH JOB OUTPUT FILE @@ -257,11 +272,25 @@ class CheckBatchCost: content_bytes # type: ignore[arg-type] ) + # Record output file size + if prom_logger and content_bytes: + try: + prom_logger.record_managed_file_size( + size_bytes=len(content_bytes), # type: ignore + purpose="batch", + file_type="output", + model=model_id, + ) + except Exception: + pass + deployment_info = self.llm_router.get_deployment(model_id=model_id) if deployment_info is None: verbose_proxy_logger.info( f"Skipping job {unified_object_id} because it is not a valid deployment info" ) + if prom_logger: + prom_logger.record_check_batch_cost_error("deployment_not_found") continue custom_llm_provider = deployment_info.litellm_params.custom_llm_provider litellm_model_name = deployment_info.litellm_params.model @@ -318,6 +347,19 @@ class CheckBatchCost: batch_models=batch_models, ) + # Record batch duration (completed_at - created_at) + if prom_logger and response.completed_at and response.created_at: + duration_seconds = float(response.completed_at - response.created_at) + if duration_seconds >= 0: + prom_logger.record_managed_batch_duration( + duration_seconds=duration_seconds, + model=model_name, + api_provider=str(llm_provider) if llm_provider else None, + ) + + # Track this job for the final metrics summary + processed_models.append((model_name, str(llm_provider) if llm_provider else None)) + # mark the job as complete try: update_data: dict = { @@ -334,3 +376,10 @@ class CheckBatchCost: verbose_proxy_logger.error( f"CheckBatchCost: failed to mark job {job.id} complete in DB: {db_err}" ) + + # Record polling run metrics (always, even if nothing was processed) + if prom_logger: + prom_logger.record_check_batch_cost_run( + jobs_polled=len(jobs), + processed_models=processed_models if processed_models else None, + ) diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index dc14937d46b..60c564072a0 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -74,6 +74,13 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): self.internal_usage_cache = internal_usage_cache self.prisma_client = prisma_client + @staticmethod + def _get_prometheus_logger(): + """Find PrometheusLogger from litellm.callbacks, if registered.""" + from litellm.integrations.prometheus import PrometheusLogger + + return PrometheusLogger.get_instance() + async def store_unified_file_id( self, file_id: str, @@ -905,6 +912,31 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): model_mappings=model_mappings, user_api_key_dict=user_api_key_dict, ) + + # Emit Prometheus metrics for managed file creation + prom_logger = self._get_prometheus_logger() + if prom_logger: + first_model = target_model_names_list[0] if target_model_names_list else None + first_provider = "" + if responses: + first_provider = getattr(responses[0], "_hidden_params", {}).get("custom_llm_provider") or "" + prom_logger.record_managed_file_created( + model=first_model or "", + api_provider=first_provider, + user=user_api_key_dict.user_id or "", + user_email=getattr(user_api_key_dict, "user_email", None) or "", + api_key_alias=user_api_key_dict.key_alias or "", + ) + if response.bytes and response.bytes > 0: + prom_logger.record_managed_file_size( + size_bytes=response.bytes, + purpose=response.purpose or "batch", + file_type="input", + model=first_model, + api_provider=first_provider, + user=user_api_key_dict.user_id, + ) + return response @staticmethod @@ -1083,6 +1115,31 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): file_purpose="batch", user_api_key_dict=user_api_key_dict, ) + + # Only record batch creation metric on actual create (not retrieve/cancel). + # unified_file_id in _hidden_params is only set by the create_batch endpoint. + original_unified_file_id = response._hidden_params.get("unified_file_id") + if original_unified_file_id: + prom_logger = self._get_prometheus_logger() + if prom_logger: + batch_provider = "" + if model_name: + try: + from litellm.litellm_core_utils.get_llm_provider_logic import ( + get_llm_provider, + ) + _, batch_provider, _, _ = get_llm_provider(model=model_name) + except Exception: + if "/" in model_name: + batch_provider = model_name.split("/")[0] + prom_logger.record_managed_batch_created( + model=model_name or "", + api_provider=batch_provider, + user=user_api_key_dict.user_id or "", + user_email=getattr(user_api_key_dict, "user_email", None) or "", + api_key_alias=user_api_key_dict.key_alias or "", + ) + elif isinstance(response, LiteLLMFineTuningJob): ## Check if unified_file_id is in the response unified_file_id = response._hidden_params.get( @@ -1332,6 +1389,11 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): f"Alternatively, wait for all batches to complete and for cost to be computed (batch_processed=true)." ) + # Record blocked deletion metric + prom_logger = self._get_prometheus_logger() + if prom_logger: + prom_logger.record_managed_file_deleted(result="blocked") + raise HTTPException( status_code=400, detail=error_message, @@ -1365,6 +1427,12 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): file_id, litellm_parent_otel_span ) + # Record successful deletion metric only on actual success + if stored_file_object or delete_response: + prom_logger = self._get_prometheus_logger() + if prom_logger: + prom_logger.record_managed_file_deleted(result="success") + if stored_file_object: return stored_file_object elif delete_response: diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 90306c11a42..8de01aac60b 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -65,6 +65,17 @@ def _get_cached_end_user_id_for_cost_tracking(): class PrometheusLogger(CustomLogger): # Class variables or attributes + + @staticmethod + def get_instance() -> Optional["PrometheusLogger"]: + """Find the PrometheusLogger instance from litellm.callbacks, if registered.""" + import litellm + + for cb in litellm.callbacks: + if isinstance(cb, PrometheusLogger): + return cb + return None + def __init__( # noqa: PLR0915 self, **kwargs, @@ -440,6 +451,76 @@ class PrometheusLogger(CustomLogger): labelnames=[], ) + ######################################## + # Managed Batch Metrics + ######################################## + self.litellm_managed_batch_created_total = self._counter_factory( + name="litellm_managed_batch_created_total", + documentation="Total number of managed batches created", + labelnames=[ + "model", + "api_provider", + "user", + "user_email", + "api_key_alias", + ], + ) + + self.litellm_managed_file_size_bytes = self._gauge_factory( + "litellm_managed_file_size_bytes", + "Size of the most recent managed batch file in bytes (last-seen value per label combination)", + labelnames=["purpose", "file_type", "model", "api_provider", "user"], + ) + + self.litellm_managed_batch_duration_seconds = self._histogram_factory( + "litellm_managed_batch_duration_seconds", + "Duration of completed managed batches in seconds (completed_at - created_at)", + labelnames=["model", "api_provider"], + buckets=BATCH_DURATION_BUCKETS, + ) + + self.litellm_managed_file_created_total = self._counter_factory( + name="litellm_managed_file_created_total", + documentation="Total number of managed files created", + labelnames=[ + "model", + "api_provider", + "user", + "user_email", + "api_key_alias", + ], + ) + + self.litellm_managed_file_deleted_total = self._counter_factory( + name="litellm_managed_file_deleted_total", + documentation="Total number of managed file deletions (success or blocked)", + labelnames=["result"], + ) + + self.litellm_check_batch_cost_jobs_polled = self._gauge_factory( + "litellm_check_batch_cost_jobs_polled", + "Number of unprocessed batches found by the last CheckBatchCost poll", + labelnames=[], + ) + + self.litellm_check_batch_cost_jobs_processed_total = self._counter_factory( + name="litellm_check_batch_cost_jobs_processed_total", + documentation="Total number of batches successfully cost-tracked by CheckBatchCost", + labelnames=["model", "api_provider"], + ) + + self.litellm_check_batch_cost_errors_total = self._counter_factory( + name="litellm_check_batch_cost_errors_total", + documentation="Total number of errors in CheckBatchCost by error type", + labelnames=["error_type"], + ) + + self.litellm_check_batch_cost_last_run_timestamp = self._gauge_factory( + "litellm_check_batch_cost_last_run_timestamp", + "Unix timestamp of the last CheckBatchCost job run", + labelnames=[], + ) + except Exception as e: print_verbose(f"Got exception on init prometheus client {str(e)}") raise e @@ -2158,6 +2239,127 @@ class PrometheusLogger(CustomLogger): except Exception as e: verbose_logger.debug(f"Error recording guardrail metrics: {str(e)}") + ######################################## + # Managed Batch Metric Recording Methods + ######################################## + + def record_managed_batch_created( + self, + model: Optional[str], + api_provider: Optional[str], + user: Optional[str], + user_email: Optional[str], + api_key_alias: Optional[str], + ): + try: + self.litellm_managed_batch_created_total.labels( + model=model, + api_provider=api_provider, + user=user, + user_email=user_email, + api_key_alias=api_key_alias, + ).inc() + except Exception as e: + verbose_logger.warning(f"Error recording batch created metric: {e}") + + def record_managed_file_size( + self, + size_bytes: int, + purpose: str, + file_type: str, + model: Optional[str] = None, + api_provider: Optional[str] = None, + user: Optional[str] = None, + ): + """Record the size of a managed file. Uses a gauge (last-seen value per label combination).""" + try: + self.litellm_managed_file_size_bytes.labels( + purpose=purpose, + file_type=file_type, + model=model or "", + api_provider=api_provider or "", + user=user or "", + ).set(size_bytes) + except Exception as e: + verbose_logger.warning(f"Error recording file size metric: {e}") + + def record_managed_batch_duration( + self, + duration_seconds: float, + model: Optional[str] = None, + api_provider: Optional[str] = None, + ): + try: + self.litellm_managed_batch_duration_seconds.labels( + model=model or "", + api_provider=api_provider or "", + ).observe(duration_seconds) + except Exception as e: + verbose_logger.warning(f"Error recording batch duration metric: {e}") + + def record_managed_file_created( + self, + model: Optional[str], + api_provider: Optional[str], + user: Optional[str], + user_email: Optional[str], + api_key_alias: Optional[str], + ): + try: + self.litellm_managed_file_created_total.labels( + model=model, + api_provider=api_provider, + user=user, + user_email=user_email, + api_key_alias=api_key_alias, + ).inc() + except Exception as e: + verbose_logger.warning(f"Error recording file created metric: {e}") + + def record_managed_file_deleted(self, result: str): + """Record a managed file deletion attempt. result is 'success' or 'blocked'.""" + try: + self.litellm_managed_file_deleted_total.labels(result=result).inc() + except Exception as e: + verbose_logger.warning(f"Error recording file deleted metric: {e}") + + def record_check_batch_cost_run( + self, + jobs_polled: int, + processed_models: Optional[List[Tuple[Optional[str], Optional[str]]]] = None, + ): + """ + Record CheckBatchCost polling metrics. + + Args: + jobs_polled: Number of unprocessed batches found + processed_models: List of (model, api_provider) tuples for processed jobs + """ + import time + + try: + self.litellm_check_batch_cost_last_run_timestamp.set(time.time()) + self.litellm_check_batch_cost_jobs_polled.set(jobs_polled) + + if processed_models: + for model, api_provider in processed_models: + self.litellm_check_batch_cost_jobs_processed_total.labels( + model=model or "", + api_provider=api_provider or "", + ).inc() + except Exception as e: + verbose_logger.warning(f"Error recording check batch cost metrics: {e}") + + def record_check_batch_cost_error(self, error_type: str): + try: + self.litellm_check_batch_cost_errors_total.labels( + error_type=error_type, + ).inc() + except Exception as e: + verbose_logger.warning( + f"Error recording check batch cost error metric: {e}" + ) + @staticmethod def _get_exception_class_name(exception: Exception) -> str: exception_class_name = "" diff --git a/litellm/types/integrations/prometheus.py b/litellm/types/integrations/prometheus.py index 0856d8a6f9b..05913e609f6 100644 --- a/litellm/types/integrations/prometheus.py +++ b/litellm/types/integrations/prometheus.py @@ -159,6 +159,23 @@ LATENCY_BUCKETS = ( float("inf"), ) +# Batch jobs can run for minutes to hours; buckets span 1 min → 24 h. +BATCH_DURATION_BUCKETS = ( + 60.0, + 120.0, + 300.0, + 600.0, + 900.0, + 1800.0, + 3600.0, + 7200.0, + 14400.0, + 28800.0, + 43200.0, + 86400.0, + float("inf"), +) + class UserAPIKeyLabelNames(Enum): END_USER = "end_user" @@ -238,6 +255,16 @@ DEFINED_PROMETHEUS_METRICS = Literal[ "litellm_llm_api_failed_requests_metric", "litellm_callback_logging_failures_metric", "litellm_in_flight_requests", + # Managed batch metrics + "litellm_managed_batch_created_total", + "litellm_managed_file_size_bytes", + "litellm_managed_batch_duration_seconds", + "litellm_managed_file_created_total", + "litellm_managed_file_deleted_total", + "litellm_check_batch_cost_jobs_polled", + "litellm_check_batch_cost_jobs_processed_total", + "litellm_check_batch_cost_errors_total", + "litellm_check_batch_cost_last_run_timestamp", ] @@ -618,6 +645,43 @@ class PrometheusMetricLabels: litellm_cache_misses_metric = _cache_metric_labels litellm_cached_tokens_metric = _cache_metric_labels + # Managed batch metrics + _batch_user_labels = [ + UserAPIKeyLabelNames.v1_LITELLM_MODEL_NAME.value, + UserAPIKeyLabelNames.API_PROVIDER.value, + UserAPIKeyLabelNames.USER.value, + UserAPIKeyLabelNames.USER_EMAIL.value, + UserAPIKeyLabelNames.API_KEY_ALIAS.value, + ] + + litellm_managed_batch_created_total = _batch_user_labels + + litellm_managed_file_size_bytes: List[ + str + ] = [] # labels: purpose, file_type, model, api_provider, user (custom) + + litellm_managed_batch_duration_seconds = [ + UserAPIKeyLabelNames.v1_LITELLM_MODEL_NAME.value, + UserAPIKeyLabelNames.API_PROVIDER.value, + ] + + litellm_managed_file_created_total = _batch_user_labels + + litellm_managed_file_deleted_total: List[ + str + ] = [] # only "result" label, added at metric creation + + litellm_check_batch_cost_jobs_polled: List[str] = [] + + litellm_check_batch_cost_jobs_processed_total = [ + UserAPIKeyLabelNames.v1_LITELLM_MODEL_NAME.value, + UserAPIKeyLabelNames.API_PROVIDER.value, + ] + + litellm_check_batch_cost_errors_total: List[str] = [] # label: error_type (custom) + + litellm_check_batch_cost_last_run_timestamp: List[str] = [] + @staticmethod def get_labels(label_name: DEFINED_PROMETHEUS_METRICS) -> List[str]: default_labels = getattr(PrometheusMetricLabels, label_name)