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 08456705925..a05cbefd52e 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -290,6 +290,57 @@ class CheckBatchCost: 404 must not retire the row; the staleness sweep bounds it instead.""" return self.llm_router.get_deployment(model_id=model_id) is not None + @staticmethod + def _is_output_file_gone_at_provider(error: Exception, output_file_id: Optional[str]) -> bool: + """A 404 naming the output file means there is nothing to fetch on this or any + later poll: providers like Vertex AI advertise an output path for every batch, + including terminal ones that never wrote it. Any other failure may be + transient, so it keeps retrying until the staleness sweep bounds it.""" + import openai + + from litellm.exceptions import NotFoundError + + if not output_file_id: + return False + return isinstance(error, (NotFoundError, openai.NotFoundError)) and output_file_id in str(error) + + async def _finalize_unbilled_terminal_job( + self, job: "LiteLLM_ManagedObjectTable", response: "LiteLLMBatch" + ) -> None: + """Persist a terminal batch that has nothing billable, converting any raw + provider file ids to managed ids, and take it out of the poll page.""" + try: + from litellm.proxy.openai_files_endpoints.common_utils import ( + _is_base64_encoded_unified_file_id, + ensure_batch_response_managed_file_ids, + ) + + response.id = job.unified_object_id + await ensure_batch_response_managed_file_ids( + response=response, + managed_files_obj=self.proxy_logging_obj.get_proxy_hook("managed_files"), + prisma_client=self.prisma_client, + verbose_proxy_logger=verbose_proxy_logger, + db_batch_object=job, + unified_batch_id=_is_base64_encoded_unified_file_id(job.unified_object_id), + ) + update_data: Final[dict] = { + "status": response.status, + "file_object": response.model_dump_json(), + **({"batch_processed": True} if self._has_batch_processed_column else {}), + } + await self.prisma_client.db.litellm_managedobjecttable.update( + where={"id": job.id}, + data=update_data, + ) + verbose_proxy_logger.info( + f"CheckBatchCost: marked job {job.id} as {response.status} in DB" + ) + except Exception as db_err: + verbose_proxy_logger.error( + f"CheckBatchCost: failed to mark job {job.id} as {response.status} in DB: {db_err}" + ) + @staticmethod def _record_error( prom_logger: Optional["PrometheusLogger"], error_type: str @@ -812,6 +863,15 @@ class CheckBatchCost: prom_logger=prom_logger, ) except Exception as tracking_err: + if self._is_output_file_gone_at_provider( + tracking_err, response.output_file_id + ) and self._batch_deployment_exists(model_id): + verbose_proxy_logger.warning( + f"CheckBatchCost: output file {response.output_file_id} of batch {batch_id} " + f"does not exist at the provider; retiring job {job.id} unbilled" + ) + await self._finalize_unbilled_terminal_job(job, response) + continue verbose_proxy_logger.error( f"CheckBatchCost: failed to track cost for batch {batch_id} " f"(job {job.id}); leaving it unprocessed so the next poll retries: {tracking_err}" @@ -842,38 +902,7 @@ class CheckBatchCost: ) elif response.status in PROVIDER_TERMINAL_BATCH_STATUSES: - try: - from litellm.proxy.openai_files_endpoints.common_utils import ( - _is_base64_encoded_unified_file_id, - ensure_batch_response_managed_file_ids, - ) - - response.id = job.unified_object_id - await ensure_batch_response_managed_file_ids( - response=response, - managed_files_obj=self.proxy_logging_obj.get_proxy_hook("managed_files"), - prisma_client=self.prisma_client, - verbose_proxy_logger=verbose_proxy_logger, - db_batch_object=job, - unified_batch_id=_is_base64_encoded_unified_file_id(job.unified_object_id), - ) - update_data = { - "status": response.status, - "file_object": response.model_dump_json(), - } - if self._has_batch_processed_column: - update_data["batch_processed"] = True - await self.prisma_client.db.litellm_managedobjecttable.update( - where={"id": job.id}, - data=update_data, - ) - verbose_proxy_logger.info( - f"CheckBatchCost: marked job {job.id} as {response.status} in DB" - ) - except Exception as db_err: - verbose_proxy_logger.error( - f"CheckBatchCost: failed to mark job {job.id} as {response.status} in DB: {db_err}" - ) + await self._finalize_unbilled_terminal_job(job, response) # Record polling run metrics (always, even if nothing was processed) if prom_logger: diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py index fe97bd3ad5f..2ac15502840 100644 --- a/tests/proxy_unit_tests/test_check_batch_cost.py +++ b/tests/proxy_unit_tests/test_check_batch_cost.py @@ -1166,6 +1166,83 @@ class TestCheckBatchCost: update_data["status"] == terminal_status ), f"billed {terminal_status} batch must keep its real terminal status in the DB" + @pytest.mark.asyncio + async def test_terminal_batch_with_missing_output_file_is_retired_unbilled( + self, check_batch_cost_instance, mock_prisma_client, mock_llm_router + ): + """A terminal batch whose advertised output file 404s at the provider has + nothing to fetch on this or any later poll (Vertex AI advertises an output + path for every batch, even ones that never wrote it), so the job must be + retired as terminal on the first cycle instead of retrying until the + staleness sweep gives up on it. + """ + import base64 + from unittest.mock import patch + + from litellm.exceptions import NotFoundError + + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( + return_value=0 + ) + 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-output-gone-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] + ) + + missing_output_file_id = "gs://batch-out/job-1/predictions.jsonl" + mock_response = MagicMock() + mock_response.status = "failed" + mock_response.output_file_id = missing_output_file_id + mock_response.error_file_id = None + mock_response.model_dump_json.return_value = ( + '{"id":"batch-1","status":"failed"}' + ) + + 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, + side_effect=NotFoundError( + message=f"404: output file {missing_output_file_id} does not exist", + model="gemini-2.5-pro", + llm_provider="vertex_ai", + ), + ) as mock_afile_content, + patch( + "litellm.batches.batch_utils.calculate_batch_cost_and_usage", + new_callable=AsyncMock, + ) as mock_calculate, + ): + await check_batch_cost_instance.check_batch_cost() + + assert mock_afile_content.await_count == 1 + mock_calculate.assert_not_awaited() + assert ( + mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1 + ), "a terminal batch with a 404ing output file must be retired, not retried forever" + update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[ + 1 + ]["data"] + assert update_data["status"] == "failed" + assert update_data["batch_processed"] is True + @pytest.mark.asyncio async def test_raw_output_file_id_converted_to_managed_id( self, check_batch_cost_instance, mock_prisma_client, mock_llm_router