fix(proxy): retire terminal batches whose advertised output file 404s instead of retrying

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:

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