From 4c0c412bc5e670a020ec175090288dce20d859ee Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Fri, 13 Mar 2026 11:15:21 -0700 Subject: [PATCH] fix(proxy): prevent OOM/Prisma connection loss from unbounded managed-object poll (#23472) --- .../proxy/common_utils/check_batch_cost.py | 184 +++++++++++++++--- .../common_utils/check_responses_cost.py | 37 +++- litellm/constants.py | 9 + litellm/proxy/proxy_server.py | 7 +- 4 files changed, 209 insertions(+), 28 deletions(-) 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 bf8bc46f723..42a9acbfd1e 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -2,17 +2,24 @@ Polls LiteLLM_ManagedObjectTable to check if the batch job is complete, and if the cost has been tracked. """ -from litellm._uuid import uuid -from datetime import datetime +from datetime import datetime, timedelta, timezone from typing import TYPE_CHECKING, Optional from litellm._logging import verbose_proxy_logger +from litellm._uuid import uuid +from litellm.constants import ( + MANAGED_OBJECT_STALENESS_CUTOFF_DAYS, + MAX_OBJECTS_PER_POLL_CYCLE, +) if TYPE_CHECKING: from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.router import Router +CHECK_BATCH_COST_USER_AGENT = "LiteLLM Proxy/CheckBatchCost" + + class CheckBatchCost: def __init__( self, @@ -26,6 +33,69 @@ class CheckBatchCost: self.proxy_logging_obj: ProxyLogging = proxy_logging_obj self.prisma_client: PrismaClient = prisma_client self.llm_router: Router = llm_router + # Cached after the first poll cycle. Once we know the column is absent we skip + # the guaranteed-failing primary query on every subsequent cycle. + self._has_batch_processed_column: bool = True + + async def _get_user_info(self, batch_id, user_id) -> dict: + """ + Look up user email and key alias by user_id for enriching the S3 callback metadata. + Returns a dict with user_api_key_user_email and user_api_key_alias (both may be None). + """ + try: + user_row = await self.prisma_client.db.litellm_usertable.find_unique( + where={"user_id": user_id} + ) + if user_row is None: + return {} + return { + "user_api_key_user_email": getattr(user_row, "user_email", None), + "user_api_key_alias": getattr(user_row, "user_alias", None), + } + except Exception as e: + verbose_proxy_logger.error(f"CheckBatchCost: could not look up user {user_id} for batch {batch_id}: {e}") + return {} + + async def _cleanup_stale_managed_objects(self) -> None: + """ + Mark managed objects older than MANAGED_OBJECT_STALENESS_CUTOFF_DAYS days + in non-terminal states as 'stale_expired'. These will never complete and + should not be polled. + """ + cutoff = datetime.now(timezone.utc) - timedelta(days=MANAGED_OBJECT_STALENESS_CUTOFF_DAYS) + result = await self.prisma_client.db.litellm_managedobjecttable.update_many( + where={ + "file_purpose": "batch", + "status": {"not_in": ["completed", "complete", "failed", "expired", "cancelled", "stale_expired"]}, + "created_at": {"lt": cutoff}, + }, + data={"status": "stale_expired"}, + ) + if result > 0: + verbose_proxy_logger.warning( + f"CheckBatchCost: marked {result} stale managed objects " + f"(older than {MANAGED_OBJECT_STALENESS_CUTOFF_DAYS} days) as stale_expired" + ) + + async def _fallback_find_jobs(self) -> list: + """Query batch jobs without the batch_processed filter (for older schemas).""" + return await self.prisma_client.db.litellm_managedobjecttable.find_many( + where={ + "file_purpose": "batch", + "status": { + "not_in": [ + "failed", + "expired", + "cancelled", + "complete", + "completed", + "stale_expired", + ] + }, + }, + take=MAX_OBJECTS_PER_POLL_CYCLE, + order={"created_at": "asc"}, + ) async def check_batch_cost(self): """ @@ -48,14 +118,48 @@ class CheckBatchCost: get_model_id_from_unified_batch_id, ) - jobs = await self.prisma_client.db.litellm_managedobjecttable.find_many( - where={ - "status": {"in": ["validating", "in_progress", "finalizing"]}, - "file_purpose": "batch", - } - ) - completed_jobs = [] + try: + await self._cleanup_stale_managed_objects() + except Exception as cleanup_err: + verbose_proxy_logger.warning( + f"CheckBatchCost: stale cleanup failed (poll will continue): {cleanup_err}" + ) + # Look for all batches that have not yet been processed by CheckBatchCost. + # self._has_batch_processed_column is cached after the first probe so that + # older schemas don't pay a guaranteed-failing primary query + warning on + # every subsequent poll cycle. + if self._has_batch_processed_column: + try: + jobs = await self.prisma_client.db.litellm_managedobjecttable.find_many( + where={ + "file_purpose": "batch", + "batch_processed": False, + "status": { + "not_in": [ + "failed", + "expired", + "cancelled", + "complete", + "completed", + "stale_expired", + ] + }, + }, + take=MAX_OBJECTS_PER_POLL_CYCLE, + order={"created_at": "asc"}, + ) + except Exception as query_err: + if "batch_processed" not in str(query_err).lower() and "unknown column" not in str(query_err).lower() and "does not exist" not in str(query_err).lower(): + raise + # Permanent schema gap — cache the result so future cycles skip straight to fallback + self._has_batch_processed_column = False + verbose_proxy_logger.warning( + "CheckBatchCost: batch_processed column not found, querying without it" + ) + jobs = await self._fallback_find_jobs() + else: + jobs = await self._fallback_find_jobs() for job in jobs: # get the model from the job unified_object_id = job.unified_object_id @@ -107,6 +211,21 @@ class CheckBatchCost: f"Batch ID: {batch_id} is complete, tracking cost and usage" ) + # aretrieve_batch is called with the raw provider batch ID, so response.id + # is the raw provider value (e.g. "batch_20260223-0518.234"). We need the + # unified base64 ID in the S3 log so downstream consumers can correlate it + # back to the batch they submitted via the proxy. + # + # CheckBatchCost builds its own LiteLLMLogging object (logging_obj below) and + # calls async_success_handler(result=response) directly. That handler calls + # _build_standard_logging_payload(response, ...) which reads response.id at + # that point — so setting response.id here is sufficient. + # + # The HTTP endpoint does this substitution via the managed files hook + # (async_post_call_success_hook). CheckBatchCost bypasses that hook entirely, + # so we do it explicitly here. + response.id = job.unified_object_id + # This background job runs as default_user_id, so going through the HTTP endpoint # would trigger check_managed_file_id_access and get 403. Instead, extract the raw # provider file ID and call afile_content directly with deployment credentials. @@ -126,14 +245,14 @@ class CheckBatchCost: # Access content - handle both direct attribute and method call if hasattr(_file_content, 'content'): - content_bytes = _file_content.content + content_bytes = _file_content.content # type: ignore[union-attr] elif hasattr(_file_content, 'read'): - content_bytes = await _file_content.read() + content_bytes = await _file_content.read() # type: ignore[misc] else: - content_bytes = _file_content + content_bytes = _file_content # type: ignore[assignment] file_content_as_dict = _get_file_content_as_dictionary( - content_bytes + content_bytes # type: ignore[arg-type] ) deployment_info = self.llm_router.get_deployment(model_id=model_id) @@ -158,7 +277,7 @@ class CheckBatchCost: file_content_dictionary=file_content_as_dict, custom_llm_provider=llm_provider, # type: ignore model_name=model_name, - model_info=deployment_model_info, + model_info=deployment_model_info, # type: ignore[arg-type] ) ) logging_obj = LiteLLMLogging( @@ -171,11 +290,21 @@ class CheckBatchCost: function_id=str(uuid.uuid4()), ) + creator_user_id = job.created_by + user_info = await self._get_user_info(batch_id, job.created_by) + logging_obj.update_environment_variables( litellm_params={ + # set the user-agent header so that S3 callback consumers can easily identify CheckBatchCost callbacks + "proxy_server_request": { + "headers": { + "user-agent": CHECK_BATCH_COST_USER_AGENT, + } + }, "metadata": { - "user_api_key_user_id": job.created_by or "default-user-id", - } + "user_api_key_user_id": creator_user_id, + **user_info, + }, }, optional_params={}, ) @@ -188,11 +317,18 @@ class CheckBatchCost: ) # mark the job as complete - completed_jobs.append(job) - - if len(completed_jobs) > 0: - # mark the jobs as complete - await self.prisma_client.db.litellm_managedobjecttable.update_many( - where={"id": {"in": [job.id for job in completed_jobs]}}, - data={"status": "complete"}, - ) + try: + update_data: dict = { + "status": "complete", + "file_object": response.model_dump_json(), + } + if self._has_batch_processed_column: + update_data["batch_processed"] = True + await self.prisma_client.db.litellm_managedobjecttable.update( + where={"id": job.id}, + data=update_data, + ) + except Exception as db_err: + verbose_proxy_logger.error( + f"CheckBatchCost: failed to mark job {job.id} complete in DB: {db_err}" + ) diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py index 4ee6a89cc98..54fbc7abcc5 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py @@ -3,10 +3,15 @@ Polls LiteLLM_ManagedObjectTable to check if the response is complete. Cost tracking is handled automatically by litellm.aget_responses(). """ +from datetime import datetime, timedelta, timezone from typing import TYPE_CHECKING import litellm from litellm._logging import verbose_proxy_logger +from litellm.constants import ( + MANAGED_OBJECT_STALENESS_CUTOFF_DAYS, + MAX_OBJECTS_PER_POLL_CYCLE, +) if TYPE_CHECKING: from litellm.proxy.utils import PrismaClient, ProxyLogging @@ -27,6 +32,27 @@ class CheckResponsesCost: self.prisma_client: PrismaClient = prisma_client self.llm_router: Router = llm_router + async def _cleanup_stale_managed_objects(self) -> None: + """ + Mark managed objects older than MANAGED_OBJECT_STALENESS_CUTOFF_DAYS days + in non-terminal states as 'stale_expired'. These will never complete and + should not be polled. + """ + cutoff = datetime.now(timezone.utc) - timedelta(days=MANAGED_OBJECT_STALENESS_CUTOFF_DAYS) + result = await self.prisma_client.db.litellm_managedobjecttable.update_many( + where={ + "file_purpose": "response", + "status": {"not_in": ["completed", "complete", "failed", "expired", "cancelled", "stale_expired"]}, + "created_at": {"lt": cutoff}, + }, + data={"status": "stale_expired"}, + ) + if result > 0: + verbose_proxy_logger.warning( + f"CheckResponsesCost: marked {result} stale managed objects " + f"(older than {MANAGED_OBJECT_STALENESS_CUTOFF_DAYS} days) as stale_expired" + ) + async def check_responses_cost(self): """ Check if background responses are complete and track their cost. @@ -35,11 +61,20 @@ class CheckResponsesCost: - Cost is automatically tracked by litellm.aget_responses() - Mark completed/failed/cancelled responses as complete in the database """ + try: + await self._cleanup_stale_managed_objects() + except Exception as cleanup_err: + verbose_proxy_logger.warning( + f"CheckResponsesCost: stale cleanup failed (poll will continue): {cleanup_err}" + ) + jobs = await self.prisma_client.db.litellm_managedobjecttable.find_many( where={ "status": {"in": ["queued", "in_progress"]}, "file_purpose": "response", - } + }, + take=MAX_OBJECTS_PER_POLL_CYCLE, + order={"created_at": "asc"}, ) verbose_proxy_logger.debug(f"Found {len(jobs)} response jobs to check") diff --git a/litellm/constants.py b/litellm/constants.py index 89992b459c2..967626f387f 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1328,6 +1328,15 @@ PROXY_BUDGET_RESCHEDULER_MIN_TIME = int( os.getenv("PROXY_BUDGET_RESCHEDULER_MIN_TIME", 597) ) PROXY_BATCH_POLLING_INTERVAL = int(os.getenv("PROXY_BATCH_POLLING_INTERVAL", 3600)) +MAX_OBJECTS_PER_POLL_CYCLE = max(1, int(os.getenv("MAX_OBJECTS_PER_POLL_CYCLE", 50))) +MANAGED_OBJECT_STALENESS_CUTOFF_DAYS = max( + 1, int(os.getenv("MANAGED_OBJECT_STALENESS_CUTOFF_DAYS", 7)) +) +# Set PROXY_BATCH_POLLING_ENABLED=false to disable the CheckBatchCost and +# CheckResponsesCost background polling jobs entirely (e.g. to avoid DB load on +# installations with large numbers of stale managed objects). +_batch_polling_env = os.getenv("PROXY_BATCH_POLLING_ENABLED", "true").lower() +PROXY_BATCH_POLLING_ENABLED = _batch_polling_env == "true" PROXY_BUDGET_RESCHEDULER_MAX_TIME = int( os.getenv("PROXY_BUDGET_RESCHEDULER_MAX_TIME", 605) ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index aeb502599dc..a65890e44c1 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1,4 +1,3 @@ -import anyio import asyncio import copy import enum @@ -31,6 +30,7 @@ from typing import ( get_type_hints, ) +import anyio from pydantic import BaseModel, Json from litellm._uuid import uuid @@ -209,6 +209,7 @@ from litellm.constants import ( DEFAULT_MODEL_CREATED_AT_TIME, LITELLM_PROXY_ADMIN_NAME, PROMETHEUS_FALLBACK_STATS_SEND_TIME_HOURS, + PROXY_BATCH_POLLING_ENABLED, PROXY_BATCH_POLLING_INTERVAL, PROXY_BATCH_WRITE_AT, PROXY_BUDGET_RESCHEDULER_MAX_TIME, @@ -5606,7 +5607,7 @@ class ProxyStartupEvent: "Invalid maximum_spend_logs_retention_interval value" ) ### CHECK BATCH COST ### - if llm_router is not None: + if llm_router is not None and PROXY_BATCH_POLLING_ENABLED: try: from litellm_enterprise.proxy.common_utils.check_batch_cost import ( CheckBatchCost, @@ -5637,7 +5638,7 @@ class ProxyStartupEvent: pass ### CHECK RESPONSES COST ### - if llm_router is not None: + if llm_router is not None and PROXY_BATCH_POLLING_ENABLED: try: from litellm_enterprise.proxy.common_utils.check_responses_cost import ( CheckResponsesCost,