mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-08 03:08:45 +00:00
Merge pull request #24691 from Sameerlite/litellm_litellm_batch_metrics_4
feat(prometheus): add metrics for managed batch lifecycle
This commit is contained in:
commit
bda8c9d8ac
4 changed files with 384 additions and 1 deletions
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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 = ""
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue