From 167b8bb1f53b308b83f9c3a4530d1af958c005d8 Mon Sep 17 00:00:00 2001 From: michelligabriele Date: Mon, 16 Mar 2026 18:45:56 +0100 Subject: [PATCH] fix(batch): sync all status changes to DB in CheckBatchCost polling job The CheckBatchCost polling job only updated the database when a batch reached "completed" status. Intermediate states (in_progress, finalizing) and terminal failure states (failed, expired, cancelled) were never written back. Since list_user_batches reads exclusively from the DB, enterprise users saw permanently stale "validating" status in the Swagger UI and GET /v1/batches responses. This adds a status synchronization step after aretrieve_batch that updates the DB for all status changes. Terminal failure states also get batch_processed=True to stop futile re-polling. --- .../proxy/common_utils/check_batch_cost.py | 48 +++++++ .../proxy/test_managed_files_access_check.py | 124 ++++++++++++++++++ 2 files changed, 172 insertions(+) diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py index cbe8d449b42..92527e605d6 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -204,6 +204,54 @@ class CheckBatchCost: ) continue + # Sync batch status to DB for all status changes (not just completion). + # This ensures list_user_batches returns fresh status from the provider. + _current_db_status = job.status + _provider_status = response.status + _normalized_status = ( + _provider_status if _provider_status != "completed" else "complete" + ) + + if _normalized_status != _current_db_status: + _is_terminal_failure = _provider_status in ( + "failed", + "expired", + "cancelled", + ) + + if _provider_status != "completed": + # For non-completed status changes, update DB now. + # Completed batches are handled below with cost tracking. + try: + update_data: dict = { + "status": _normalized_status, + "file_object": response.model_dump_json(), + } + if _is_terminal_failure: + if self._has_batch_processed_column: + update_data["batch_processed"] = True + verbose_proxy_logger.info( + f"Batch {batch_id} reached terminal state " + f"'{_provider_status}', marking as processed" + ) + else: + verbose_proxy_logger.info( + f"Batch {batch_id} status changed: " + f"{_current_db_status} -> {_provider_status}" + ) + 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 sync status for " + f"batch {batch_id}: {db_err}" + ) + + if _is_terminal_failure: + continue # No cost tracking needed for failed/expired/cancelled + ## RETRIEVE THE BATCH JOB OUTPUT FILE if ( response.status == "completed" diff --git a/tests/test_litellm/enterprise/proxy/test_managed_files_access_check.py b/tests/test_litellm/enterprise/proxy/test_managed_files_access_check.py index 8cb642b7a44..8593a0f10ee 100644 --- a/tests/test_litellm/enterprise/proxy/test_managed_files_access_check.py +++ b/tests/test_litellm/enterprise/proxy/test_managed_files_access_check.py @@ -207,3 +207,127 @@ async def test_check_batch_cost_should_call_afile_content_directly_with_credenti assert "file_object" in update_call_kwargs["data"], ( "file_object must be written to DB so list_batches reads updated status" ) + + +# --- Status sync tests: polling job updates DB for all status changes --- + + +def _make_checker_for_status_test(mock_job, batch_response): + """Helper to create a CheckBatchCost instance for status sync tests.""" + from litellm_enterprise.proxy.common_utils.check_batch_cost import CheckBatchCost + + mock_prisma = MagicMock() + mock_prisma.db.litellm_managedobjecttable.find_many = AsyncMock( + return_value=[mock_job] + ) + mock_prisma.db.litellm_managedobjecttable.update = AsyncMock() + mock_prisma.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) + + mock_router = MagicMock() + mock_router.aretrieve_batch = AsyncMock(return_value=batch_response) + + checker = CheckBatchCost( + proxy_logging_obj=MagicMock(), + prisma_client=mock_prisma, + llm_router=mock_router, + ) + return checker, mock_prisma + + +def _make_mock_job(status: str) -> MagicMock: + """Create a mock DB job with a specific status.""" + unified_raw = "litellm_proxy;model_id:model-deploy-xyz;llm_batch_id:batch-123;llm_output_file_id:file-raw-output" + mock_job = MagicMock() + mock_job.unified_object_id = base64.b64encode(unified_raw.encode()).decode() + mock_job.created_by = "user-A" + mock_job.id = "job-1" + mock_job.status = status + return mock_job + + +def _make_batch_response(status: str): + """Create a LiteLLMBatch with a specific status.""" + from litellm.types.utils import LiteLLMBatch + + return LiteLLMBatch( + id="batch-123", + completion_window="24h", + created_at=1700000000, + endpoint="/v1/chat/completions", + input_file_id="file-input", + object="batch", + status=status, + ) + + +@pytest.mark.asyncio +async def test_check_batch_cost_updates_db_for_in_progress_status(): + """ + When provider returns in_progress, the polling job should update the DB + status and file_object so list_batches shows fresh data. + """ + mock_job = _make_mock_job(status="validating") + batch_response = _make_batch_response(status="in_progress") + checker, mock_prisma = _make_checker_for_status_test(mock_job, batch_response) + + await checker.check_batch_cost() + + # DB should be updated with in_progress status + mock_prisma.db.litellm_managedobjecttable.update.assert_called_once() + update_kwargs = mock_prisma.db.litellm_managedobjecttable.update.call_args.kwargs + assert update_kwargs["data"]["status"] == "in_progress" + assert "file_object" in update_kwargs["data"] + # batch_processed should NOT be set for intermediate states + assert "batch_processed" not in update_kwargs["data"] + + +@pytest.mark.asyncio +async def test_check_batch_cost_marks_failed_batch_as_processed(): + """ + When provider returns failed, the polling job should update DB status + AND set batch_processed=True to stop re-polling. + """ + mock_job = _make_mock_job(status="validating") + batch_response = _make_batch_response(status="failed") + checker, mock_prisma = _make_checker_for_status_test(mock_job, batch_response) + + await checker.check_batch_cost() + + mock_prisma.db.litellm_managedobjecttable.update.assert_called_once() + update_kwargs = mock_prisma.db.litellm_managedobjecttable.update.call_args.kwargs + assert update_kwargs["data"]["status"] == "failed" + assert update_kwargs["data"]["batch_processed"] is True + assert "file_object" in update_kwargs["data"] + + +@pytest.mark.asyncio +async def test_check_batch_cost_marks_expired_batch_as_processed(): + """ + When provider returns expired, the polling job should update DB status + AND set batch_processed=True to stop re-polling. + """ + mock_job = _make_mock_job(status="in_progress") + batch_response = _make_batch_response(status="expired") + checker, mock_prisma = _make_checker_for_status_test(mock_job, batch_response) + + await checker.check_batch_cost() + + mock_prisma.db.litellm_managedobjecttable.update.assert_called_once() + update_kwargs = mock_prisma.db.litellm_managedobjecttable.update.call_args.kwargs + assert update_kwargs["data"]["status"] == "expired" + assert update_kwargs["data"]["batch_processed"] is True + + +@pytest.mark.asyncio +async def test_check_batch_cost_skips_update_when_status_unchanged(): + """ + When provider returns the same status as DB, no update should be made. + """ + mock_job = _make_mock_job(status="validating") + batch_response = _make_batch_response(status="validating") + checker, mock_prisma = _make_checker_for_status_test(mock_job, batch_response) + + await checker.check_batch_cost() + + # No DB update should have been made + mock_prisma.db.litellm_managedobjecttable.update.assert_not_called()