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 4bb00408fc32..76e92538aaa7 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -966,6 +966,16 @@ async def check_batch_cost(self): ) elif response.status in PROVIDER_TERMINAL_BATCH_STATUSES: + from litellm.proxy.openai_files_endpoints.common_utils import ( + _completed_batch_safe_to_retire, + ) + + if response.status in ("completed", "complete") and not _completed_batch_safe_to_retire(response): + verbose_proxy_logger.info( + f"CheckBatchCost: batch {batch_id} is completed but its output file id " + f"has not appeared yet; leaving job {job.id} for the next poll cycle" + ) + continue await self._finalize_unbilled_terminal_job(job, response) # Record polling run metrics (always, even if nothing was processed) diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index d4cb8fe2374a..142aced4a38f 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -1334,6 +1334,25 @@ def batch_cost_poller_is_active() -> bool: return False +def _completed_batch_safe_to_retire(response: "LiteLLMBatch") -> bool: + """Whether a "completed" batch may be retired from cost recovery. + + ``batch_processed=True`` is the sole re-pickup gate for CheckBatchCost's + cost-recovery poller, so setting it retires the batch permanently. A batch can + reach ``status="completed"`` while ``output_file_id`` is still ``None`` (the + provider response briefly lags before the output id populates). Retiring in that + window loses the spend record forever. Retire only once we can prove there is + nothing left to recover: the output file has actually arrived, or the provider + reports no successful request lines. When counts are unknown, stay eligible so + the next poller pass revisits it. (#37713) + """ + if getattr(response, "output_file_id", None) is not None: + return True + request_counts = getattr(response, "request_counts", None) + completed = getattr(request_counts, "completed", None) + return completed == 0 + + async def update_batch_in_database( batch_id: str, unified_batch_id: str | Literal[False], @@ -1415,7 +1434,7 @@ async def update_batch_in_database( } poller_owns: Final = batch_cost_poller_is_active() if poller_owns_accounting is None else poller_owns_accounting - if db_status == "complete" and not poller_owns: + if db_status == "complete" and not poller_owns and _completed_batch_safe_to_retire(response): update_data["batch_processed"] = True try: diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py index 0065dbebc596..a1864c5e4809 100644 --- a/tests/proxy_unit_tests/test_check_batch_cost.py +++ b/tests/proxy_unit_tests/test_check_batch_cost.py @@ -1044,7 +1044,9 @@ async def test_completed_without_output_file_marked_processed_without_billing( Pre-fix it matched neither the completed-with-output branch nor the failed/expired/cancelled branch, so batch_processed stayed False and the row was re-selected on every poll cycle forever. It must now be marked terminal - exactly once, without being billed (no output means nothing to bill). + exactly once, without being billed: request_counts.completed == 0 proves the + missing output file means nothing to bill rather than a lagging output id + (#37713 keeps the lagging case eligible for the next cycle). """ import base64 from unittest.mock import patch @@ -1073,6 +1075,7 @@ async def test_completed_without_output_file_marked_processed_without_billing( mock_response.status = completed_status mock_response.output_file_id = None mock_response.error_file_id = "file-error-123" + mock_response.request_counts = MagicMock(completed=0, failed=3, total=3) mock_response.model_dump_json.return_value = ( f'{{"id":"batch-1","status":"{completed_status}"}}' ) @@ -1107,6 +1110,72 @@ async def test_completed_without_output_file_marked_processed_without_billing( mock_llm_router.get_deployment_credentials_with_provider.call_count == 0 ), "a batch with no output file must not enter the cost-tracking path" + @pytest.mark.asyncio + @pytest.mark.parametrize( + "request_counts", + [MagicMock(completed=7, failed=0, total=7), None], + ids=["lagging_output_id", "unknown_counts"], + ) + async def test_completed_with_lagging_output_file_left_for_next_cycle( + self, + check_batch_cost_instance, + mock_prisma_client, + mock_llm_router, + request_counts, + ): + """#37713 regression: a batch can report completed while its output_file_id is + still lagging behind at the provider. Retiring it in that window (or when the + request counts cannot prove there is nothing to bill) permanently loses the + spend record, so the poller must leave the row untouched and revisit it on the + next cycle once the output id has appeared. + """ + import base64 + from unittest.mock import patch + + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( + return_value=1 + ) + mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( + return_value=None + ) + + mock_job = MagicMock() + mock_job.id = "job-completed-lagging-output-1" + mock_job.unified_object_id = base64.urlsafe_b64encode( + b"litellm_proxy;model_id:model-123;llm_batch_id:batch-456" + ).decode() + mock_job.created_by = "user-1" + + assert check_batch_cost_instance._has_batch_processed_column is True + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( + return_value=[mock_job] + ) + + mock_response = MagicMock() + mock_response.status = "completed" + mock_response.output_file_id = None + mock_response.error_file_id = None + mock_response.request_counts = request_counts + + mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response) + mock_llm_router.get_deployment_credentials_with_provider = MagicMock( + return_value={"api_key": "sk-test"} + ) + + with patch( + "litellm.files.main.afile_content", + new_callable=AsyncMock, + ) as mock_afile_content: + await check_batch_cost_instance.check_batch_cost() + + assert ( + mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 0 + ), "a completed batch whose output id is still lagging must stay eligible for the next poll" + assert ( + mock_afile_content.await_count == 0 + ), "a batch with no output file must not be billed" + @pytest.mark.asyncio async def test_non_terminal_status_left_unprocessed( self, check_batch_cost_instance, mock_prisma_client, mock_llm_router diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py index 2161e345b409..7f84407f8b39 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py @@ -428,3 +428,47 @@ def test_add_internal_model_credentials_survives_a_failing_deployment_lookup(): add_internal_model_credentials(data=data, llm_router=router, model_id="deployment-gone") assert data == {"batch_id": "unified-batch-id"} + + +from litellm.proxy.openai_files_endpoints.common_utils import ( + _completed_batch_safe_to_retire, +) + + +def _completed_batch_for_retire( + output_file_id: str | None, completed: int | None = None +) -> LiteLLMBatch: + kwargs = dict( + id="batch-1", + completion_window="24h", + created_at=1234567890, + endpoint="/v1/chat/completions", + input_file_id="file-in", + object="batch", + status="completed", + output_file_id=output_file_id, + error_file_id=None, + ) + if completed is not None: + kwargs["request_counts"] = {"total": completed, "completed": completed, "failed": 0} + return LiteLLMBatch(**kwargs) + + +class TestCompletedBatchSafeToRetire: + """A completed batch is only safe to retire from cost recovery once its output + file has arrived or the provider proves no successful lines (#37713).""" + + def test_output_file_present_is_safe(self): + assert _completed_batch_safe_to_retire(_completed_batch_for_retire("file-out")) is True + + def test_no_output_and_no_successful_lines_is_safe(self): + # Every request line errored -> nothing left to recover. + assert _completed_batch_safe_to_retire(_completed_batch_for_retire(None, completed=0)) is True + + def test_no_output_but_successful_lines_is_not_safe(self): + # The bug: output_file_id is lagging; retiring here loses the spend record. + assert _completed_batch_safe_to_retire(_completed_batch_for_retire(None, completed=5)) is False + + def test_no_output_and_unknown_counts_is_not_safe(self): + # Counts unknown -> stay eligible so the next poller pass revisits it. + assert _completed_batch_safe_to_retire(_completed_batch_for_retire(None)) is False