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 42a9acbfd1e..cbe8d449b42 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -131,6 +131,10 @@ class CheckBatchCost: # every subsequent poll cycle. if self._has_batch_processed_column: try: + # Include "complete"/"completed" batches: the retrieve_batch + # endpoint may transition a batch to "complete" before + # CheckBatchCost runs. The batch_processed=False filter + # already prevents reprocessing finished batches. jobs = await self.prisma_client.db.litellm_managedobjecttable.find_many( where={ "file_purpose": "batch", @@ -140,8 +144,6 @@ class CheckBatchCost: "failed", "expired", "cancelled", - "complete", - "completed", "stale_expired", ] }, diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index 9fa2a51fa83..740e63b7f17 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -406,9 +406,11 @@ async def retrieve_batch( # noqa: PLR0915 verbose_proxy_logger=verbose_proxy_logger, ) - # If batch is in a terminal state, return immediately + # If batch is in a terminal state, return immediately. + # Include "complete" (DB-normalized form of "completed"). if response is not None and response.status in [ "completed", + "complete", "failed", "cancelled", "expired", diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index b75b2f4640f..49f17535333 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -831,14 +831,43 @@ async def update_batch_in_database( # Normalize status for database storage db_status = response.status if response.status != "completed" else "complete" - await prisma_client.db.litellm_managedobjecttable.update( - where={"unified_object_id": batch_id}, - data={ - "status": db_status, - "file_object": response.model_dump_json(), - "updated_at": litellm.utils.get_utc_datetime(), - }, - ) + update_data: dict = { + "status": db_status, + "file_object": response.model_dump_json(), + "updated_at": litellm.utils.get_utc_datetime(), + } + + # When a batch reaches completion, also mark batch_processed=True. + # The cost callback is enqueued asynchronously during the + # aretrieve_batch call that detected completion (via the @client + # decorator). It is not awaited, so there is a theoretical window + # where the callback hasn't executed yet. In practice the callback + # completes reliably. Setting the flag here unblocks file deletion + # which queries batch_processed=False. CheckBatchCost acts as a + # safety net for the rare case where the callback fails. + if db_status == "complete": + update_data["batch_processed"] = True + + try: + await prisma_client.db.litellm_managedobjecttable.update( + where={"unified_object_id": batch_id}, + data=update_data, + ) + except Exception as col_err: + # If the batch_processed column doesn't exist (old schema), + # retry without it so the status update still succeeds. + err_str = str(col_err).lower() + if "batch_processed" in err_str and update_data.get("batch_processed") is not None: + verbose_proxy_logger.warning( + f"batch_processed column not found, retrying update without it: {col_err}" + ) + update_data.pop("batch_processed", None) + await prisma_client.db.litellm_managedobjecttable.update( + where={"unified_object_id": batch_id}, + data=update_data, + ) + else: + raise except Exception as e: verbose_proxy_logger.error( f"Failed to update batch status in ManagedObjectTable: {e}" diff --git a/tests/openai_endpoints_tests/test_openai_batches_endpoint.py b/tests/openai_endpoints_tests/test_openai_batches_endpoint.py index 215ac0874f2..7e0c2771ad2 100644 --- a/tests/openai_endpoints_tests/test_openai_batches_endpoint.py +++ b/tests/openai_endpoints_tests/test_openai_batches_endpoint.py @@ -400,6 +400,8 @@ async def test_batch_status_sync_from_provider_to_database(): assert update_call_args.kwargs["data"]["status"] == "complete" # "completed" normalized to "complete" assert "file_object" in update_call_args.kwargs["data"] assert "updated_at" in update_call_args.kwargs["data"] + # batch_processed must be set to True when batch transitions to complete + assert update_call_args.kwargs["data"]["batch_processed"] is True # Verify logger was called with status change message mock_logger.info.assert_called() diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py index f6b8d567848..a84524f8244 100644 --- a/tests/proxy_unit_tests/test_check_batch_cost.py +++ b/tests/proxy_unit_tests/test_check_batch_cost.py @@ -84,8 +84,14 @@ class TestCheckBatchCost: assert find_call[1]["order"] == {"created_at": "asc"} not_in = find_call[1]["where"]["status"]["not_in"] assert "stale_expired" in not_in - assert "complete" in not_in - assert "completed" in not_in + # "complete"/"completed" are intentionally NOT excluded from the + # primary query — the batch_processed=False filter is sufficient. + # This allows CheckBatchCost to pick up batches that were + # transitioned to "complete" by the retrieve_batch endpoint + # before CheckBatchCost had a chance to process them. + assert "complete" not in not_in + assert "completed" not in not_in + assert find_call[1]["where"]["batch_processed"] is False @pytest.mark.asyncio async def test_fallback_query_used_when_batch_processed_missing(