fix(proxy): prevent OOM/Prisma connection loss from unbounded managed-object poll (#23472)

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}"
+                    )

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

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

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,