Merge pull request #24691 from Sameerlite/litellm_litellm_batch_metrics_4

feat(prometheus): add metrics for managed batch lifecycle
This commit is contained in:
yuneng-jiang 2026-03-27 09:52:59 -07:00 • committed by GitHub
commit bda8c9d8ac
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 384 additions and 1 deletions

View file

@ -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,
)

View file

@ -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:

View file

@ -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 = ""

View file

@ -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)