From 80fb4ece9770e26d33b11bfd63d8bc146591eb74 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 2 Apr 2025 18:39:29 -0700 Subject: [PATCH] prom emit size of DB TX queues for observability --- litellm/integrations/prometheus_services.py | 10 +++--- .../db_transaction_queue/base_update_queue.py | 16 +++++++++ .../daily_spend_update_queue.py | 24 ++++++++++++-- .../db_transaction_queue/pod_lock_manager.py | 4 +-- .../redis_update_buffer.py | 33 +++++++++++++++++-- .../spend_update_queue.py | 24 ++++++++++++-- 6 files changed, 97 insertions(+), 14 deletions(-) diff --git a/litellm/integrations/prometheus_services.py b/litellm/integrations/prometheus_services.py index d14cbd7469..dddaa4d064 100644 --- a/litellm/integrations/prometheus_services.py +++ b/litellm/integrations/prometheus_services.py @@ -73,9 +73,7 @@ class PrometheusServicesLogger: service_metrics.append(counter_total_requests) if ServiceMetrics.GAUGE in metrics_to_initialize: - gauge = self.create_gauge( - service, type_of_request="pod_lock_manager" - ) + gauge = self.create_gauge(service, type_of_request="size") if gauge: service_metrics.append(gauge) @@ -95,12 +93,14 @@ class PrometheusServicesLogger: def _get_service_metrics_initialize( self, service: ServiceTypes ) -> List[ServiceMetrics]: + DEFAULT_METRICS = [ServiceMetrics.COUNTER, ServiceMetrics.GAUGE] if service not in DEFAULT_SERVICE_CONFIGS: - raise ValueError(f"Service {service} not found in DEFAULT_SERVICE_CONFIGS") + return DEFAULT_METRICS metrics = DEFAULT_SERVICE_CONFIGS.get(service, {}).get("metrics", []) if not metrics: - raise ValueError(f"No metrics found for service {service}") + verbose_logger.debug(f"No metrics found for service {service}") + return DEFAULT_METRICS return metrics def is_metric_registered(self, metric_name) -> bool: diff --git a/litellm/proxy/db/db_transaction_queue/base_update_queue.py b/litellm/proxy/db/db_transaction_queue/base_update_queue.py index b3c3c26c84..2bf2393127 100644 --- a/litellm/proxy/db/db_transaction_queue/base_update_queue.py +++ b/litellm/proxy/db/db_transaction_queue/base_update_queue.py @@ -2,8 +2,14 @@ Base class for in memory buffer for database transactions """ import asyncio +from typing import Optional from litellm._logging import verbose_proxy_logger +from litellm._service_logger import ServiceLogging + +service_logger_obj = ( + ServiceLogging() +) # used for tracking metrics for In memory buffer, redis buffer, pod lock manager class BaseUpdateQueue: @@ -16,6 +22,9 @@ class BaseUpdateQueue: """Enqueue an update.""" verbose_proxy_logger.debug("Adding update to queue: %s", update) await self.update_queue.put(update) + await self._emit_new_item_added_to_queue_event( + queue_size=self.update_queue.qsize() + ) async def flush_all_updates_from_in_memory_queue(self): """Get all updates from the queue.""" @@ -23,3 +32,10 @@ class BaseUpdateQueue: while not self.update_queue.empty(): updates.append(await self.update_queue.get()) return updates + + async def _emit_new_item_added_to_queue_event( + self, + queue_size: Optional[int] = None, + ): + """placeholder, emit event when a new item is added to the queue""" + pass diff --git a/litellm/proxy/db/db_transaction_queue/daily_spend_update_queue.py b/litellm/proxy/db/db_transaction_queue/daily_spend_update_queue.py index dedb8c8f8f..afae431370 100644 --- a/litellm/proxy/db/db_transaction_queue/daily_spend_update_queue.py +++ b/litellm/proxy/db/db_transaction_queue/daily_spend_update_queue.py @@ -1,9 +1,13 @@ import asyncio -from typing import Dict, List +from typing import Dict, List, Optional from litellm._logging import verbose_proxy_logger from litellm.proxy._types import DailyUserSpendTransaction -from litellm.proxy.db.db_transaction_queue.base_update_queue import BaseUpdateQueue +from litellm.proxy.db.db_transaction_queue.base_update_queue import ( + BaseUpdateQueue, + service_logger_obj, +) +from litellm.types.services import ServiceTypes class DailySpendUpdateQueue(BaseUpdateQueue): @@ -93,3 +97,19 @@ class DailySpendUpdateQueue(BaseUpdateQueue): else: aggregated_daily_spend_update_transactions[_key] = payload return aggregated_daily_spend_update_transactions + + async def _emit_new_item_added_to_queue_event( + self, + queue_size: Optional[int] = None, + ): + asyncio.create_task( + service_logger_obj.async_service_success_hook( + service=ServiceTypes.IN_MEMORY_DAILY_SPEND_UPDATE_QUEUE, + duration=0, + call_type="_emit_new_item_added_to_queue_event", + event_metadata={ + "gauge_labels": ServiceTypes.IN_MEMORY_DAILY_SPEND_UPDATE_QUEUE, + "gauge_value": queue_size, + }, + ) + ) diff --git a/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py b/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py index 3f63afe62a..cb4a43a802 100644 --- a/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py +++ b/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py @@ -3,9 +3,9 @@ import uuid from typing import TYPE_CHECKING, Any, Optional from litellm._logging import verbose_proxy_logger -from litellm._service_logger import ServiceLogging from litellm.caching.redis_cache import RedisCache from litellm.constants import DEFAULT_CRON_JOB_LOCK_TTL_SECONDS +from litellm.proxy.db.db_transaction_queue.base_update_queue import service_logger_obj from litellm.types.services import ServiceTypes if TYPE_CHECKING: @@ -13,8 +13,6 @@ if TYPE_CHECKING: else: ProxyLogging = Any -service_logger_obj = ServiceLogging() # used for tracking current pod lock status - class PodLockManager: """ diff --git a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py index ea1356159a..88741fbb18 100644 --- a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py +++ b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py @@ -4,6 +4,7 @@ Handles buffering database `UPDATE` transactions in Redis before committing them This is to prevent deadlocks and improve reliability """ +import asyncio import json from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union @@ -16,11 +17,13 @@ from litellm.constants import ( ) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy._types import DailyUserSpendTransaction, DBSpendUpdateTransactions +from litellm.proxy.db.db_transaction_queue.base_update_queue import service_logger_obj from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import ( DailySpendUpdateQueue, ) from litellm.proxy.db.db_transaction_queue.spend_update_queue import SpendUpdateQueue from litellm.secret_managers.main import str_to_bool +from litellm.types.services import ServiceTypes if TYPE_CHECKING: from litellm.proxy.utils import PrismaClient @@ -136,18 +139,27 @@ class RedisUpdateBuffer: return list_of_transactions = [safe_dumps(db_spend_update_transactions)] - await self.redis_cache.async_rpush( + current_redis_buffer_size = await self.redis_cache.async_rpush( key=REDIS_UPDATE_BUFFER_KEY, values=list_of_transactions, ) + await self._emit_new_item_added_to_redis_buffer_event( + queue_size=current_redis_buffer_size, + service=ServiceTypes.REDIS_SPEND_UPDATE_QUEUE, + ) list_of_daily_spend_update_transactions = [ safe_dumps(daily_spend_update_transactions) ] - await self.redis_cache.async_rpush( + + current_redis_buffer_size = await self.redis_cache.async_rpush( key=REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY, values=list_of_daily_spend_update_transactions, ) + await self._emit_new_item_added_to_redis_buffer_event( + queue_size=current_redis_buffer_size, + service=ServiceTypes.REDIS_DAILY_SPEND_UPDATE_QUEUE, + ) @staticmethod def _number_of_transactions_to_store_in_redis( @@ -300,3 +312,20 @@ class RedisUpdateBuffer: ) return combined_transaction + + async def _emit_new_item_added_to_redis_buffer_event( + self, + service: ServiceTypes, + queue_size: int, + ): + asyncio.create_task( + service_logger_obj.async_service_success_hook( + service=service, + duration=0, + call_type="_emit_new_item_added_to_queue_event", + event_metadata={ + "gauge_labels": service, + "gauge_value": queue_size, + }, + ) + ) diff --git a/litellm/proxy/db/db_transaction_queue/spend_update_queue.py b/litellm/proxy/db/db_transaction_queue/spend_update_queue.py index ce181d1478..60e9379751 100644 --- a/litellm/proxy/db/db_transaction_queue/spend_update_queue.py +++ b/litellm/proxy/db/db_transaction_queue/spend_update_queue.py @@ -1,5 +1,5 @@ import asyncio -from typing import List +from typing import List, Optional from litellm._logging import verbose_proxy_logger from litellm.proxy._types import ( @@ -7,7 +7,11 @@ from litellm.proxy._types import ( Litellm_EntityType, SpendUpdateQueueItem, ) -from litellm.proxy.db.db_transaction_queue.base_update_queue import BaseUpdateQueue +from litellm.proxy.db.db_transaction_queue.base_update_queue import ( + BaseUpdateQueue, + service_logger_obj, +) +from litellm.types.services import ServiceTypes class SpendUpdateQueue(BaseUpdateQueue): @@ -111,3 +115,19 @@ class SpendUpdateQueue(BaseUpdateQueue): transactions_dict[entity_id] += response_cost or 0 return db_spend_update_transactions + + async def _emit_new_item_added_to_queue_event( + self, + queue_size: Optional[int] = None, + ): + asyncio.create_task( + service_logger_obj.async_service_success_hook( + service=ServiceTypes.IN_MEMORY_SPEND_UPDATE_QUEUE, + duration=0, + call_type="_emit_new_item_added_to_queue_event", + event_metadata={ + "gauge_labels": ServiceTypes.IN_MEMORY_SPEND_UPDATE_QUEUE, + "gauge_value": queue_size, + }, + ) + )