From 21984101e5ce0ac8ca9bbb6c8ff187df2dad8c89 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 17 Aug 2026 12:23:59 -0700 Subject: [PATCH 1/2] fix(proxy): bill cancelled and failed batches that still produced an output file --- .../proxy/common_utils/check_batch_cost.py | 16 ++--- .../proxy_unit_tests/test_check_batch_cost.py | 65 ++++++++----------- 2 files changed, 34 insertions(+), 47 deletions(-) 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 25b005973555..08456705925a 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -24,12 +24,16 @@ CHECK_BATCH_COST_USER_AGENT = "LiteLLM Proxy/CheckBatchCost" -TERMINAL_MANAGED_OBJECT_STATUSES: Final[Tuple[str, ...]] = ( +PROVIDER_TERMINAL_BATCH_STATUSES: Final[Tuple[str, ...]] = ( "completed", "complete", "failed", "expired", "cancelled", +) + +TERMINAL_MANAGED_OBJECT_STATUSES: Final[Tuple[str, ...]] = ( + *PROVIDER_TERMINAL_BATCH_STATUSES, "stale_expired", ) @@ -796,7 +800,7 @@ async def check_batch_cost(self): ## RETRIEVE THE BATCH JOB OUTPUT FILE if ( - response.status in ("completed", "complete", "expired") + response.status in PROVIDER_TERMINAL_BATCH_STATUSES and response.output_file_id is not None ): try: @@ -837,13 +841,7 @@ async def check_batch_cost(self): f"CheckBatchCost: failed to mark job {job.id} complete in DB: {db_err}" ) - elif response.status in ( - "completed", - "complete", - "failed", - "expired", - "cancelled", - ): + elif response.status in PROVIDER_TERMINAL_BATCH_STATUSES: try: from litellm.proxy.openai_files_endpoints.common_utils import ( _is_base64_encoded_unified_file_id, diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py index 72f8b87dd16d..fe97bd3ad5fb 100644 --- a/tests/proxy_unit_tests/test_check_batch_cost.py +++ b/tests/proxy_unit_tests/test_check_batch_cost.py @@ -817,12 +817,12 @@ async def test_terminal_status_persists_managed_output_file_ids( mock_llm_router, terminal_status, ): - """A cancelled/failed batch with provider output files must be persisted with - unified managed file IDs, never raw provider IDs. Raw IDs written here leak - to every later GET /batches/{id} and GET /batches because the terminal row is - final (batch_processed=True) and read paths only resolve, never mint. - (Expired with an output file is billed through the completed path instead, - covered by test_expired_with_output_file_is_billed.) + """A cancelled/failed batch with a provider error file (and no output file) must + be persisted with unified managed file IDs, never raw provider IDs. Raw IDs + written here leak to every later GET /batches/{id} and GET /batches because the + terminal row is final (batch_processed=True) and read paths only resolve, never + mint. (Any terminal status with an output file is billed through the completed + path instead, covered by test_terminal_status_with_output_file_is_billed.) """ import base64 import json @@ -832,15 +832,11 @@ async def test_terminal_status_persists_managed_output_file_ids( unified_batch_uid = base64.urlsafe_b64encode( b"litellm_proxy;model_id:model-123;llm_batch_id:batch-456" ).decode() - raw_output_file_id = "file-terminal-out-abc" raw_error_file_id = "file-terminal-err-xyz" raw_input_file_id = "file-terminal-in-123" unified_input_file_id = base64.urlsafe_b64encode( b"litellm_proxy:application/octet-stream;unified_id,in-1;target_model_names,gpt-5-batch" ).decode() - unified_output_file_id = base64.urlsafe_b64encode( - f"litellm_proxy:application/octet-stream;unified_id,u-1;llm_output_file_id,{raw_output_file_id}".encode() - ).decode() unified_error_file_id = base64.urlsafe_b64encode( f"litellm_proxy:application/octet-stream;unified_id,u-2;llm_output_file_id,{raw_error_file_id}".encode() ).decode() @@ -884,16 +880,13 @@ def find_managed_file(where): input_file_id=raw_input_file_id, object="batch", status=terminal_status, - output_file_id=raw_output_file_id, + output_file_id=None, error_file_id=raw_error_file_id, ) mock_llm_router.aretrieve_batch = AsyncMock(return_value=response) mock_hook = MagicMock() - mock_hook.get_unified_output_file_id.side_effect = [ - unified_output_file_id, - unified_error_file_id, - ] + mock_hook.get_unified_output_file_id.side_effect = [unified_error_file_id] mock_hook.store_unified_file_id = AsyncMock() check_batch_cost_instance.proxy_logging_obj.get_proxy_hook.return_value = ( mock_hook @@ -901,12 +894,7 @@ def find_managed_file(where): await check_batch_cost_instance.check_batch_cost() - mock_hook.get_unified_output_file_id.assert_any_call( - output_file_id=raw_output_file_id, - model_id="model-123", - model_name="gpt-5-batch", - ) - mock_hook.get_unified_output_file_id.assert_any_call( + mock_hook.get_unified_output_file_id.assert_called_once_with( output_file_id=raw_error_file_id, model_id="model-123", model_name="gpt-5-batch", @@ -915,10 +903,7 @@ def find_managed_file(where): next(iter(c.kwargs["model_mappings"].values())): c.kwargs["file_id"] for c in mock_hook.store_unified_file_id.call_args_list } - assert stored == { - raw_output_file_id: unified_output_file_id, - raw_error_file_id: unified_error_file_id, - } + assert stored == {raw_error_file_id: unified_error_file_id} for store_call in mock_hook.store_unified_file_id.call_args_list: assert store_call.kwargs["user_api_key_dict"].user_id == "user-1" assert store_call.kwargs["user_api_key_dict"].team_id == "team-1" @@ -932,9 +917,8 @@ def find_managed_file(where): persisted = json.loads(update_data["file_object"]) assert persisted["id"] == unified_batch_uid assert persisted["input_file_id"] == unified_input_file_id - assert persisted["output_file_id"] == unified_output_file_id + assert persisted["output_file_id"] is None assert persisted["error_file_id"] == unified_error_file_id - assert raw_output_file_id not in update_data["file_object"] assert raw_error_file_id not in update_data["file_object"] @pytest.mark.asyncio @@ -1067,12 +1051,17 @@ async def test_non_terminal_status_left_unprocessed( ), "a non-terminal batch must not be written back (would stop polling prematurely)" @pytest.mark.asyncio - async def test_expired_with_output_file_is_billed( - self, check_batch_cost_instance, mock_prisma_client, mock_llm_router + @pytest.mark.parametrize("terminal_status", ["expired", "cancelled", "failed"]) + async def test_terminal_status_with_output_file_is_billed( + self, + check_batch_cost_instance, + mock_prisma_client, + mock_llm_router, + terminal_status, ): - """An expired batch that still produced an output file served real request lines, - so it must be billed (cost tracked) and then marked processed, not silently - marked terminal without billing. + """A terminal (expired/cancelled/failed) batch that still produced an output file + served real request lines, so it must be billed (cost tracked) and then marked + processed, not silently marked terminal without billing. """ from unittest.mock import patch @@ -1085,7 +1074,7 @@ async def test_expired_with_output_file_is_billed( ) mock_job = MagicMock() - mock_job.id = "job-expired-with-output-1" + mock_job.id = "job-terminal-with-output-1" mock_job.unified_object_id = "dW5pZmllZF9iYXRjaF9pZA==" mock_job.created_by = "user-1" @@ -1095,10 +1084,10 @@ async def test_expired_with_output_file_is_billed( ) mock_response = MagicMock() - mock_response.status = "expired" + mock_response.status = terminal_status mock_response.output_file_id = "file-output-123" mock_response.model_dump_json.return_value = ( - '{"id":"batch-1","status":"expired"}' + f'{{"id":"batch-1","status":"{terminal_status}"}}' ) mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response) @@ -1164,7 +1153,7 @@ async def test_expired_with_output_file_is_billed( assert ( mock_afile_content.await_count == 1 - ), "expired batch with an output file must fetch results and be billed" + ), f"{terminal_status} batch with an output file must fetch results and be billed" mock_logging_obj.async_success_handler.assert_awaited_once() assert ( mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1 @@ -1174,8 +1163,8 @@ async def test_expired_with_output_file_is_billed( ]["data"] assert update_data["batch_processed"] is True assert ( - update_data["status"] == "expired" - ), "billed expired batch must keep its real terminal status in the DB" + 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_raw_output_file_id_converted_to_managed_id( From 2f9d331b4c6469a9577dd672f4eaf6c67c4e2b14 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 17 Aug 2026 12:36:39 -0700 Subject: [PATCH 2/2] fix(proxy): retire terminal batches whose advertised output file 404s instead of retrying --- .../proxy/common_utils/check_batch_cost.py | 93 ++++++++++++------- .../proxy_unit_tests/test_check_batch_cost.py | 77 +++++++++++++++ 2 files changed, 138 insertions(+), 32 deletions(-) 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 08456705925a..a05cbefd52e8 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 @@ def _batch_deployment_exists(self, model_id: str) -> bool: 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 @@ async def check_batch_cost(self): 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 @@ async def check_batch_cost(self): ) 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 fe97bd3ad5fb..2ac155028406 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 @@ async def test_terminal_status_with_output_file_is_billed( 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