diff --git a/docs/my-website/docs/proxy/config_settings.md b/docs/my-website/docs/proxy/config_settings.md index 90d71fb95c8..90045755e9c 100644 --- a/docs/my-website/docs/proxy/config_settings.md +++ b/docs/my-website/docs/proxy/config_settings.md @@ -934,6 +934,9 @@ router_settings: | PROXY_BASE_URL | Base URL for proxy service | PROXY_BATCH_WRITE_AT | Time in seconds to wait before batch writing spend logs to the database. Default is 10 | PROXY_BATCH_POLLING_INTERVAL | Time in seconds to wait before polling a batch, to check if it's completed. Default is 6000s (1 hour) +| PROXY_BATCH_POLLING_ENABLED | Set to `false` to disable the `CheckBatchCost` and `CheckResponsesCost` background polling jobs entirely. Useful for emergency mitigation on installs with large numbers of stale managed objects. Default is `true` +| MAX_OBJECTS_PER_POLL_CYCLE | Maximum number of managed objects (batches / responses) fetched per polling cycle. Prevents OOM on installs with many stale rows. Default is `50` +| MANAGED_OBJECT_STALENESS_CUTOFF_DAYS | Managed objects older than this many days in a non-terminal state are marked `stale_expired` at the start of each poll cycle and skipped. Default is `7` | PROXY_BUDGET_RESCHEDULER_MAX_TIME | Maximum time in seconds to wait before checking database for budget resets. Default is 605 | PROXY_BUDGET_RESCHEDULER_MIN_TIME | Minimum time in seconds to wait before checking database for budget resets. Default is 597 | PYTHON_GC_THRESHOLD | GC thresholds ('gen0,gen1,gen2', e.g. '1000,50,50'); defaults to Python’s values. diff --git a/tests/proxy_unit_tests/test_check_responses_cost.py b/tests/proxy_unit_tests/test_check_responses_cost.py index 3bcacdfc05d..bbffd6e0e75 100644 --- a/tests/proxy_unit_tests/test_check_responses_cost.py +++ b/tests/proxy_unit_tests/test_check_responses_cost.py @@ -8,6 +8,7 @@ from unittest.mock import AsyncMock, MagicMock, Mock, patch import pytest +from litellm.constants import MAX_OBJECTS_PER_POLL_CYCLE from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse @@ -63,21 +64,45 @@ class TestCheckResponsesCost: self, check_responses_cost_instance, mock_prisma_client ): """Test check_responses_cost when there are no jobs to process""" - # Mock empty job list + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( + return_value=[] + ) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( + return_value=0 + ) + + await check_responses_cost_instance.check_responses_cost() + + # Verify find_many was called with pagination params + find_many_call = mock_prisma_client.db.litellm_managedobjecttable.find_many.call_args + assert find_many_call[1]["where"] == { + "status": {"in": ["queued", "in_progress"]}, + "file_purpose": "response", + } + assert find_many_call[1]["take"] == MAX_OBJECTS_PER_POLL_CYCLE + assert find_many_call[1]["order"] == {"created_at": "asc"} + + @pytest.mark.asyncio + async def test_cleanup_stale_managed_objects( + self, check_responses_cost_instance, mock_prisma_client + ): + """Stale rows (older than cutoff) are bulk-updated to stale_expired before polling.""" + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( + return_value=5 + ) mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( return_value=[] ) - # Should not raise any errors await check_responses_cost_instance.check_responses_cost() - # Verify find_many was called with correct parameters - mock_prisma_client.db.litellm_managedobjecttable.find_many.assert_called_once_with( - where={ - "status": {"in": ["queued", "in_progress"]}, - "file_purpose": "response", - } - ) + # The first update_many call should be the stale-row cleanup + calls = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list + stale_call = calls[0] + assert stale_call[1]["data"] == {"status": "stale_expired"} + where = stale_call[1]["where"] + assert "stale_expired" in where["status"]["not_in"] + assert "created_at" in where @pytest.mark.asyncio async def test_check_responses_cost_with_completed_response(