From 24f5d321a2ca199a4e28d50b1048ecf47ef44fe7 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Mon, 23 Mar 2026 10:04:02 -0700 Subject: [PATCH] fix(proxy): batch stale-response cleanup to avoid full-table-scan P2028 on large installs --- .../common_utils/check_responses_cost.py | 35 +++-- litellm/constants.py | 5 + .../test_responses_background_cost.py | 132 ++++++++++++++---- 3 files changed, 137 insertions(+), 35 deletions(-) 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 54fbc7abcc5..d7a29688a94 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py +++ b/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): """ diff --git a/litellm/constants.py b/litellm/constants.py index c0dd115210c..1c76ba244ab 100644 --- a/litellm/constants.py +++ b/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). diff --git a/tests/test_litellm/integrations/test_responses_background_cost.py b/tests/test_litellm/integrations/test_responses_background_cost.py index 5cb42704181..2dcc1848743 100644 --- a/tests/test_litellm/integrations/test_responses_background_cost.py +++ b/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()