diff --git a/docs/my-website/docs/proxy/config_settings.md b/docs/my-website/docs/proxy/config_settings.md index 528a5c109035..11cec01fdee4 100644 --- a/docs/my-website/docs/proxy/config_settings.md +++ b/docs/my-website/docs/proxy/config_settings.md @@ -1032,6 +1032,7 @@ router_settings: | SENDGRID_SENDER_EMAIL | Email address used as the sender in SendGrid email transactions | SPEND_LOGS_URL | URL for retrieving spend logs | SPEND_LOG_CLEANUP_BATCH_SIZE | Number of logs deleted per batch during cleanup. Default is 1000 +| STALE_OBJECT_CLEANUP_BATCH_SIZE | Max number of stale managed objects updated per cleanup cycle. Default is 1000 | SSL_CERTIFICATE | Path to the SSL certificate file | SSL_ECDH_CURVE | ECDH curve for SSL/TLS key exchange (e.g., 'X25519' to disable PQC). | SSL_SECURITY_LEVEL | [BETA] Security level for SSL/TLS connections. E.g. `DEFAULT@SECLEVEL=1` 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 54fbc7abcc59..dc0168683c80 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py @@ -11,6 +11,7 @@ from litellm.constants import ( MANAGED_OBJECT_STALENESS_CUTOFF_DAYS, MAX_OBJECTS_PER_POLL_CYCLE, + STALE_OBJECT_CLEANUP_BATCH_SIZE, ) if TYPE_CHECKING: @@ -32,21 +33,49 @@ def __init__( self.prisma_client: PrismaClient = prisma_client self.llm_router: Router = llm_router + async def _expire_stale_rows( + self, cutoff: datetime, batch_size: int + ) -> int: + """Execute the bounded UPDATE that marks stale rows as 'stale_expired'. + + Isolated so it can be swapped / mocked in tests without touching the + orchestration logic in ``_cleanup_stale_managed_objects``. + + Uses PostgreSQL syntax (``$1::timestamptz``, ``LIMIT``, double-quoted + identifiers) which is the only dialect the proxy supports — every + ``schema.prisma`` in the repo sets ``provider = "postgresql"``. + Same pattern as ``spend_log_cleanup.py``. + """ + return await self.prisma_client.db.execute_raw( + """ + UPDATE "LiteLLM_ManagedObjectTable" + SET "status" = 'stale_expired' + WHERE "id" IN ( + SELECT "id" FROM "LiteLLM_ManagedObjectTable" + WHERE "file_purpose" = 'response' + AND "status" NOT IN ('completed', 'complete', 'failed', 'expired', 'cancelled', 'stale_expired') + AND "created_at" < $1::timestamptz + ORDER BY "created_at" ASC + LIMIT $2 + ) + """, + cutoff, + batch_size, + ) + 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. + + Runs as a single DB query with a subquery LIMIT so no rows are loaded + into Python memory. Processes at most STALE_OBJECT_CLEANUP_BATCH_SIZE + rows per invocation to avoid overwhelming the DB when there is a large + backlog. """ 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"}, - ) + result = await self._expire_stale_rows(cutoff, STALE_OBJECT_CLEANUP_BATCH_SIZE) if result > 0: verbose_proxy_logger.warning( f"CheckResponsesCost: marked {result} stale managed objects " diff --git a/litellm/constants.py b/litellm/constants.py index 1af53b2dae08..28c6c0cc0e39 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1367,6 +1367,9 @@ MANAGED_OBJECT_STALENESS_CUTOFF_DAYS = max( 1, int(os.getenv("MANAGED_OBJECT_STALENESS_CUTOFF_DAYS", 7)) ) +STALE_OBJECT_CLEANUP_BATCH_SIZE = max( + 1, int(os.getenv("STALE_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/local_testing/test_embedding.py b/tests/local_testing/test_embedding.py index c43cad78b1e7..3f1a397ebcbe 100644 --- a/tests/local_testing/test_embedding.py +++ b/tests/local_testing/test_embedding.py @@ -373,35 +373,6 @@ def test_openai_azure_embedding_optional_arg(): # test_openai_embedding() -@pytest.mark.parametrize( - "model, api_base", - [ - ("embed-english-v2.0", None), - ], -) -@pytest.mark.parametrize("sync_mode", [True, False]) -@pytest.mark.asyncio -async def test_cohere_embedding(sync_mode, model, api_base): - try: - # litellm.set_verbose=True - data = { - "model": model, - "input": ["good morning from litellm", "this is another item"], - "input_type": "search_query", - "api_base": api_base, - } - if sync_mode: - response = embedding(**data) - else: - response = await litellm.aembedding(**data) - - print(f"response:", response) - - assert isinstance(response.usage, litellm.Usage) - except Exception as e: - pytest.fail(f"Error occurred: {e}") - - # test_cohere_embedding()