fix(proxy): batch stale-response cleanup to avoid full-table-scan P2028 on large installs

enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py

@@ -9,6 +9,7 @@ from typing import TYPE_CHECKING
 import litellm
 from litellm._logging import verbose_proxy_logger
 from litellm.constants import (
+    MANAGED_OBJECT_CLEANUP_BATCH_SIZE,
     MANAGED_OBJECT_STALENESS_CUTOFF_DAYS,
     MAX_OBJECTS_PER_POLL_CYCLE,
 )
@@ -34,24 +35,38 @@ class CheckResponsesCost:
 
     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.
+        Mark up to MANAGED_OBJECT_CLEANUP_BATCH_SIZE stale response objects as
+        'stale_expired' per poll cycle.
+
+        We fetch IDs first then update by PK to avoid a full-table-scan UPDATE
+        that would lock 336K+ rows and trigger P2028 transaction timeouts.
+        The table has no index on (file_purpose, status, created_at), so an
+        unbounded update_many causes a sequential scan with massive row locking.
         """
         cutoff = datetime.now(timezone.utc) - timedelta(days=MANAGED_OBJECT_STALENESS_CUTOFF_DAYS)
-        result = await self.prisma_client.db.litellm_managedobjecttable.update_many(
+        terminal = ["completed", "complete", "failed", "expired", "cancelled", "stale_expired"]
+
+        stale_rows = await self.prisma_client.db.litellm_managedobjecttable.find_many(
             where={
                 "file_purpose": "response",
-                "status": {"not_in": ["completed", "complete", "failed", "expired", "cancelled", "stale_expired"]},
+                "status": {"not_in": terminal},
                 "created_at": {"lt": cutoff},
             },
+            select={"id": True},
+            take=MANAGED_OBJECT_CLEANUP_BATCH_SIZE,
+        )
+        if not stale_rows:
+            return
+
+        stale_ids = [row.id for row in stale_rows]
+        await self.prisma_client.db.litellm_managedobjecttable.update_many(
+            where={"id": {"in": stale_ids}},
             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"
-            )
+        verbose_proxy_logger.warning(
+            f"CheckResponsesCost: marked {len(stale_ids)} stale managed objects "
+            f"(older than {MANAGED_OBJECT_STALENESS_CUTOFF_DAYS} days) as stale_expired"
+        )
 
     async def check_responses_cost(self):
         """

litellm/constants.py

@@ -1356,6 +1356,11 @@ MAX_OBJECTS_PER_POLL_CYCLE = max(1, int(os.getenv("MAX_OBJECTS_PER_POLL_CYCLE",
 MANAGED_OBJECT_STALENESS_CUTOFF_DAYS = max(
     1, int(os.getenv("MANAGED_OBJECT_STALENESS_CUTOFF_DAYS", 7))
 )
+# Max rows to mark stale per poll cycle. One batch per cycle keeps DB load bounded
+# and avoids a full-table-scan UPDATE that would lock 336K+ rows at once.
+MANAGED_OBJECT_CLEANUP_BATCH_SIZE = max(
+    1, int(os.getenv("MANAGED_OBJECT_CLEANUP_BATCH_SIZE", 1000))
+)
 # 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).

tests/test_litellm/integrations/test_responses_background_cost.py

@@ -9,7 +9,10 @@ from unittest.mock import AsyncMock, MagicMock, Mock, patch
 
 import pytest
 
-from litellm.constants import MAX_OBJECTS_PER_POLL_CYCLE
+from litellm.constants import (
+    MANAGED_OBJECT_CLEANUP_BATCH_SIZE,
+    MAX_OBJECTS_PER_POLL_CYCLE,
+)
 from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse
 
 
@@ -50,11 +53,9 @@ class TestResponsesBackgroundCostTracking:
             output=[],
             usage=None,
         )
-        
+
         # Add hidden params with model_id (simulating what base_process_llm_request does)
-        response._hidden_params = {
-            "model_id": "model-deployment-id-123"
-        }
+        response._hidden_params = {"model_id": "model-deployment-id-123"}
 
         # Mock request data
         data = {
@@ -73,7 +74,7 @@ class TestResponsesBackgroundCostTracking:
                 # Get model_id from hidden params
                 hidden_params = getattr(response, "_hidden_params", {}) or {}
                 model_id = hidden_params.get("model_id", None)
-                
+
                 if model_id:
                     # Store in managed objects table using response.id directly
                     await mock_managed_files_obj.store_unified_object_id(
@@ -192,7 +193,7 @@ class TestResponsesBackgroundCostTracking:
             if response.status in ["queued", "in_progress"]:
                 hidden_params = getattr(response, "_hidden_params", {}) or {}
                 model_id = hidden_params.get("model_id", None)
-                
+
                 if model_id:  # This will be False
                     await mock_managed_files_obj.store_unified_object_id(
                         unified_object_id=response.id,
@@ -241,7 +242,7 @@ class TestResponsesBackgroundCostTracking:
                 if response.status in ["queued", "in_progress"]:
                     hidden_params = getattr(response, "_hidden_params", {}) or {}
                     model_id = hidden_params.get("model_id", None)
-                    
+
                     if model_id:
                         await mock_managed_files_obj.store_unified_object_id(
                             unified_object_id=response.id,
@@ -265,6 +266,7 @@ def _check_responses_cost_module_available():
         from litellm_enterprise.proxy.common_utils.check_responses_cost import (  # noqa: F401
             CheckResponsesCost,
         )
+
         return True
     except ImportError:
         return False
@@ -272,7 +274,7 @@ def _check_responses_cost_module_available():
 
 @pytest.mark.skipif(
     not _check_responses_cost_module_available(),
-    reason="litellm_enterprise.proxy.common_utils.check_responses_cost module not available (enterprise-only feature)"
+    reason="litellm_enterprise.proxy.common_utils.check_responses_cost module not available (enterprise-only feature)",
 )
 class TestCheckResponsesCost:
     """Tests for the CheckResponsesCost polling class"""
@@ -337,15 +339,19 @@ class TestCheckResponsesCost:
         # Should not raise any errors
         await checker.check_responses_cost()
 
-        # Verify find_many was called with correct parameters (includes pagination)
-        mock_prisma_client.db.litellm_managedobjecttable.find_many.assert_called_once_with(
-            where={
-                "status": {"in": ["queued", "in_progress"]},
-                "file_purpose": "response",
-            },
-            take=MAX_OBJECTS_PER_POLL_CYCLE,
-            order={"created_at": "asc"},
+        # find_many is called at least twice: once by cleanup (returns []), once for jobs.
+        # Verify the jobs query was among the calls.
+        calls = (
+            mock_prisma_client.db.litellm_managedobjecttable.find_many.call_args_list
         )
+        jobs_calls = [
+            c
+            for c in calls
+            if c.kwargs.get("where", {}).get("status")
+            == {"in": ["queued", "in_progress"]}
+        ]
+        assert len(jobs_calls) == 1
+        assert jobs_calls[0].kwargs["take"] == MAX_OBJECTS_PER_POLL_CYCLE
 
     @pytest.mark.asyncio
     async def test_check_responses_cost_with_completed_job(
@@ -398,9 +404,12 @@ class TestCheckResponsesCost:
 
             # Verify update_many was called to mark job as completed
             # (stale cleanup also calls update_many, so check the specific completion call)
-            update_many_calls = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list
+            update_many_calls = (
+                mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list
+            )
             completion_calls = [
-                c for c in update_many_calls
+                c
+                for c in update_many_calls
                 if c.kwargs.get("where", {}).get("id") is not None
             ]
             assert len(completion_calls) == 1
@@ -450,9 +459,12 @@ class TestCheckResponsesCost:
 
             # Verify job was marked as completed even though it failed
             # (stale cleanup also calls update_many, so check the specific completion call)
-            update_many_calls = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list
+            update_many_calls = (
+                mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list
+            )
             completion_calls = [
-                c for c in update_many_calls
+                c
+                for c in update_many_calls
                 if c.kwargs.get("where", {}).get("id") is not None
             ]
             assert len(completion_calls) == 1
@@ -500,9 +512,12 @@ class TestCheckResponsesCost:
 
             # Verify no completion update_many was called (job still in progress)
             # (stale cleanup may still call update_many, so filter for completion calls)
-            update_many_calls = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list
+            update_many_calls = (
+                mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list
+            )
             completion_calls = [
-                c for c in update_many_calls
+                c
+                for c in update_many_calls
                 if c.kwargs.get("where", {}).get("id") is not None
             ]
             assert len(completion_calls) == 0
@@ -544,9 +559,76 @@ class TestCheckResponsesCost:
 
             # Verify no completion update_many was called (error occurred)
             # (stale cleanup may still call update_many, so filter for completion calls)
-            update_many_calls = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list
+            update_many_calls = (
+                mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list
+            )
             completion_calls = [
-                c for c in update_many_calls
+                c
+                for c in update_many_calls
                 if c.kwargs.get("where", {}).get("id") is not None
             ]
             assert len(completion_calls) == 0
+
+    @pytest.mark.asyncio
+    async def test_cleanup_stale_managed_objects_single_batch(
+        self, mock_proxy_logging_obj, mock_prisma_client, mock_llm_router
+    ):
+        """
+        _cleanup_stale_managed_objects must fetch IDs first then update by PK — not
+        issue an unbounded update_many with a WHERE on unindexed columns.
+
+        With 336K rows and no index on (file_purpose, status, created_at), the old
+        approach locked the entire table in one transaction and triggered P2028.
+        """
+        from litellm_enterprise.proxy.common_utils.check_responses_cost import (
+            CheckResponsesCost,
+        )
+
+        stale_rows = [MagicMock(id=f"stale-{i}") for i in range(5)]
+        mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
+            return_value=stale_rows
+        )
+        mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock()
+
+        checker = CheckResponsesCost(
+            proxy_logging_obj=mock_proxy_logging_obj,
+            prisma_client=mock_prisma_client,
+            llm_router=mock_llm_router,
+        )
+        await checker._cleanup_stale_managed_objects()
+
+        # find_many must use select={"id": True} and take= to bound the scan
+        find_call = mock_prisma_client.db.litellm_managedobjecttable.find_many.call_args
+        assert find_call.kwargs["select"] == {"id": True}
+        assert find_call.kwargs["take"] == MANAGED_OBJECT_CLEANUP_BATCH_SIZE
+
+        # update_many must be scoped to the fetched IDs (uses PK index, not full scan)
+        update_call = (
+            mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args
+        )
+        assert update_call.kwargs["where"] == {"id": {"in": [r.id for r in stale_rows]}}
+        assert update_call.kwargs["data"] == {"status": "stale_expired"}
+
+    @pytest.mark.asyncio
+    async def test_cleanup_stale_managed_objects_no_stale_rows(
+        self, mock_proxy_logging_obj, mock_prisma_client, mock_llm_router
+    ):
+        """When there are no stale rows, cleanup does one find_many and no update_many."""
+        from litellm_enterprise.proxy.common_utils.check_responses_cost import (
+            CheckResponsesCost,
+        )
+
+        mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
+            return_value=[]
+        )
+        mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock()
+
+        checker = CheckResponsesCost(
+            proxy_logging_obj=mock_proxy_logging_obj,
+            prisma_client=mock_prisma_client,
+            llm_router=mock_llm_router,
+        )
+        await checker._cleanup_stale_managed_objects()
+
+        mock_prisma_client.db.litellm_managedobjecttable.find_many.assert_called_once()
+        mock_prisma_client.db.litellm_managedobjecttable.update_many.assert_not_called()