From 3fb95c8fe385ebe422d82d8cb10c9e52fac266fb Mon Sep 17 00:00:00 2001 From: joshua-berri Date: Sat, 17 Oct 2026 19:10:22 -0500 Subject: [PATCH] fix(batches): preserve spend across malformed output and metadata failures (#47184) * fix(batches): preserve recoverable spend and terminal reconciliation state * fix(batches): narrow validated envelope types * fix(batches): serialize stored metadata when releasing claims * fix(batches): sanitize metadata before claiming completed batches * refactor(batches): expose existing metadata serializer for shared use --------- Co-authored-by: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> --- .../proxy/common_utils/check_batch_cost.py | 68 ++++--- .../proxy/hooks/managed_files.py | 8 +- litellm/batches/batch_utils.py | 16 +- tests/unit/batches/test_batch_utils.py | 28 +++ .../proxy/test_managed_files_access_check.py | 21 +- .../common_utils/test_check_batch_cost.py | 187 +++++++++++------- 6 files changed, 209 insertions(+), 119 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 94c3674945d..2d2bc92dbd2 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -3,6 +3,7 @@ Polls LiteLLM_ManagedObjectTable to check if the batch job is complete, and if t """ import asyncio +import json from collections.abc import Mapping, Sequence from dataclasses import replace as dataclasses_replace from datetime import datetime, timedelta, timezone @@ -227,7 +228,7 @@ class CheckBatchCost: "so it will no longer be polled" ) - async def _claim_job_for_costing(self, job: "ManagedObjectRow") -> bool: + async def _claim_job_for_costing(self, job: "ManagedObjectRow", response: "LiteLLMBatch") -> bool: """ Atomically flip batch_processed from false to true, returning whether this pod won the row. Every pod and uvicorn worker schedules its own poller against the shared @@ -241,10 +242,16 @@ class CheckBatchCost: """ if not self._has_batch_processed_column: return True + from litellm_enterprise.proxy.hooks.managed_files import postgres_storable_json + try: claimed: Final = await _managed_object_table(self.prisma_client).update_many( where={"id": job.id, "batch_processed": False}, - data={"batch_processed": True}, + data={ + "batch_processed": True, + "status": response.status if response.status != "completed" else "complete", + "file_object": postgres_storable_json(response), + }, ) except Exception as db_err: verbose_proxy_logger.error( @@ -253,7 +260,9 @@ class CheckBatchCost: return False return claimed > 0 - async def _release_job_claim(self, job: "ManagedObjectRow") -> None: + async def _release_job_claim( + self, job: "ManagedObjectRow", original_status: str | None, original_file_object: object + ) -> None: """Give a claimed row back once billing it failed, so a later poll cycle retries it. Safe to match on batch_processed=True: while this poller is active the retrieve @@ -265,7 +274,11 @@ class CheckBatchCost: try: await _managed_object_table(self.prisma_client).update_many( where={"id": job.id, "batch_processed": True}, - data={"batch_processed": False}, + data={ + "batch_processed": False, + "status": original_status, + "file_object": json.dumps(original_file_object), + }, ) except Exception as db_err: verbose_proxy_logger.error( @@ -564,8 +577,6 @@ class CheckBatchCost: @staticmethod def _get_input_file_id(job: "ManagedObjectRow") -> Optional[str]: - import json - from litellm.types.utils import LiteLLMBatch file_object: Final = job.file_object @@ -718,8 +729,7 @@ class CheckBatchCost: ) except Exception as _e: verbose_proxy_logger.warning( - f"CheckBatchCost: failed to create managed file ID for " - f"{_file_attr}={_raw_file_id!r}: {_e}" + f"CheckBatchCost: failed to create managed file ID for {_file_attr}={_raw_file_id!r}: {_e}" ) # Pass the deployment's router-registered pricing (litellm_params custom @@ -788,7 +798,9 @@ class CheckBatchCost: custom_llm_provider=str(llm_provider) if llm_provider else None, ) - if not await self._claim_job_for_costing(job): + original_status: Final = job.status + original_file_object: Final = job.file_object + if not await self._claim_job_for_costing(job, response.model_copy(update={"usage": batch_result.usage})): verbose_proxy_logger.info( f"CheckBatchCost: batch {batch_id} (job {job.id}) was claimed by another pod " "in this window, so its cost is already being tracked there" @@ -807,7 +819,7 @@ class CheckBatchCost: batch_completion_cost=batch_result.completion_cost, ) except Exception: - await self._release_job_claim(job) + await self._release_job_claim(job, original_status, original_file_object) raise # Record batch duration (completed_at - created_at) @@ -1100,10 +1112,7 @@ class CheckBatchCost: continue ## RETRIEVE THE BATCH JOB OUTPUT FILE - if ( - response.status in PROVIDER_TERMINAL_BATCH_STATUSES - and response.output_file_id is not None - ): + if response.status in PROVIDER_TERMINAL_BATCH_STATUSES and response.output_file_id is not None: try: tracked = await self._track_completed_batch_cost( job=job, @@ -1134,24 +1143,19 @@ class CheckBatchCost: # Track this job for the final metrics summary processed_models.append(tracked) - # mark the job as complete - try: - update_data: dict = { - "status": response.status if response.status != "completed" else "complete", - "file_object": response.model_dump_json(), - } - if self._has_batch_processed_column: - update_data["batch_processed"] = True - await _managed_object_table(self.prisma_client).update( - where={"id": job.id}, - data=update_data, - ) - except Exception as db_err: - verbose_proxy_logger.error( - f"CheckBatchCost: failed to mark job {job.id} complete in DB: {db_err}" - ) - if not self._has_batch_processed_column: - # nothing claimed the row, so it stays in the poll page and its lines wait for that retry + if not self._has_batch_processed_column: + try: + await _managed_object_table(self.prisma_client).update( + where={"id": job.id}, + data={ + "status": response.status if response.status != "completed" else "complete", + "file_object": response.model_dump_json(), + }, + ) + except Exception as db_err: + verbose_proxy_logger.error( + f"CheckBatchCost: failed to mark job {job.id} complete in DB: {db_err}" + ) continue self._store_batch_line_items_in_background(job, response, model_id, batch_id, billed=True) diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index a33b33662ab..439770390f8 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -416,7 +416,7 @@ def _storable_json_value(value: JsonValue) -> JsonValue: return value -def _postgres_storable_json(file_object: BaseModel) -> str: +def postgres_storable_json(file_object: BaseModel) -> str: dumped: Final = file_object.model_dump_json() if "\\u0000" not in dumped: return dumped @@ -650,7 +650,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): ) # FIX: Update status and file_object on every operation to keep state in sync update_columns: Final = { - "file_object": _postgres_storable_json(stored_object), + "file_object": postgres_storable_json(stored_object), "status": file_object.status, "updated_by": user_api_key_dict.user_id, } @@ -666,7 +666,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): data={ "create": { "unified_object_id": unified_object_id, - "file_object": _postgres_storable_json(stored_object), + "file_object": postgres_storable_json(stored_object), "model_object_id": model_object_id, "file_purpose": file_purpose, "created_by": resolve_resource_owner_id(user_api_key_dict), @@ -702,7 +702,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): "unified_object_id": response_id, "model_object_id": decrypted_response_id(response_id), "file_purpose": "response", - "file_object": _postgres_storable_json(_without_model_output(response)), + "file_object": postgres_storable_json(_without_model_output(response)), "status": PENDING_RESPONSE_STATUSES[0], "created_by": resolve_resource_owner_id(user_api_key_dict), "team_id": user_api_key_dict.team_id, diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index a50cdf75d33..45b0eaa0b57 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -794,11 +794,17 @@ def _batch_response_was_successful( message batch results lines report ``result.type == "succeeded"``; Bedrock batch output lines report ``modelOutput`` (and no ``error``). """ - if custom_llm_provider == "anthropic": - return _get_anthropic_result_from_batch_results_line(batch_job_output_file).get("type") == "succeeded" if custom_llm_provider == "bedrock": return batch_job_output_file.get("modelOutput") is not None and batch_job_output_file.get("error") is None - _response: Final[dict[str, object]] = cast( # cast-ok: batch response bodies come from provider JSON - dict[str, object], batch_job_output_file.get("response", None) or {} + envelope_key: Final = "result" if custom_llm_provider == "anthropic" else "response" + envelope: Final = batch_job_output_file.get(envelope_key) + if not isinstance(envelope, Mapping): + if envelope is not None: + verbose_logger.warning("batch output %s must be an object, treating the line as unsuccessful", envelope_key) + return False + response: Final = cast(Mapping[object, object], envelope) # cast-ok: mapping shape validated above + return ( + response.get("type") == "succeeded" + if custom_llm_provider == "anthropic" + else response.get("status_code") == 200 ) - return _response.get("status_code", None) == 200 diff --git a/tests/unit/batches/test_batch_utils.py b/tests/unit/batches/test_batch_utils.py index 912d4cfe84c..630b00e172b 100644 --- a/tests/unit/batches/test_batch_utils.py +++ b/tests/unit/batches/test_batch_utils.py @@ -17,6 +17,7 @@ deterministic stand-ins so the arithmetic under test is the only variable. import asyncio, json, time import logging from types import MappingProxyType +from typing import Final, Literal import httpx import pytest @@ -3347,3 +3348,30 @@ async def test_batch_retrieve_cost_tracking_with_partial_explicit_data(): assert mock_batch._hidden_params["batch_successful_requests"] == 8 assert mock_batch._hidden_params["batch_failed_requests"] == 0 assert mock_batch.usage == expected_usage + + +@pytest.mark.parametrize("provider", ("openai", "anthropic")) +@pytest.mark.parametrize("envelope", (1, "invalid", [1], None, {})) +@pytest.mark.parametrize("include_valid", (True, False)) +def test_invalid_batch_envelopes_preserve_recoverable_cost( + provider: Literal["openai", "anthropic"], envelope: object, include_valid: bool +) -> None: + invalid: Final[dict[str, object]] = {"response" if provider == "openai" else "result": envelope} + valid: Final[dict[str, object]] = ( + _success_row(model="batch-accounting-control", usage=_usage(100, 20)) + if provider == "openai" + else { + "result": { + "type": "succeeded", + "message": {"model": "batch-accounting-control", "usage": {"input_tokens": 100, "output_tokens": 20}}, + } + } + ) + entries: Final = (valid, invalid, valid) if include_valid else (invalid,) + result: Final = bu._aggregate_batch_cost_usage_models( + entries, provider, model_info={"input_cost_per_token_batches": 0.001, "output_cost_per_token_batches": 0.002} + ) + assert result.cost == pytest.approx(0.28 if include_valid else 0) + assert result.usage.prompt_tokens == (200 if include_valid else 0) + assert result.usage.completion_tokens == (40 if include_valid else 0) + assert (result.successful_requests, result.failed_requests) == (2 if include_valid else 0, 1) diff --git a/tests/unit/enterprise/proxy/test_managed_files_access_check.py b/tests/unit/enterprise/proxy/test_managed_files_access_check.py index 62302446a29..0cdc212502c 100644 --- a/tests/unit/enterprise/proxy/test_managed_files_access_check.py +++ b/tests/unit/enterprise/proxy/test_managed_files_access_check.py @@ -9,6 +9,8 @@ with deployment credentials, bypassing the managed files access-control hooks. """ import base64 +import json +from typing import Final import pytest from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, patch @@ -367,13 +369,14 @@ async def test_check_batch_cost_should_call_afile_content_directly_with_credenti # managed_files_obj.afile_content should NOT have been called mock_managed_files_hook.afile_content.assert_not_called() - # Verify the DB update writes batch_processed, status, and file_object - mock_prisma.db.litellm_managedobjecttable.update.assert_called_once() - update_call_kwargs = ( - mock_prisma.db.litellm_managedobjecttable.update.call_args.kwargs + claim_data: Final = next( + call.kwargs["data"] + for call in mock_prisma.db.litellm_managedobjecttable.update_many.call_args_list + if call.kwargs["where"].get("id") == mock_job.id ) - assert update_call_kwargs["data"]["batch_processed"] is True - assert update_call_kwargs["data"]["status"] == "complete" - assert ( - "file_object" in update_call_kwargs["data"] - ), "file_object must be written to DB so list_batches reads updated status" + assert claim_data["batch_processed"] is True + assert claim_data["status"] == "complete" + persisted_batch: Final = json.loads(claim_data["file_object"]) + assert persisted_batch["status"] == "completed" + assert persisted_batch["usage"]["prompt_tokens"] == 10 + assert persisted_batch["usage"]["completion_tokens"] == 5 diff --git a/tests/unit/proxy/common_utils/test_check_batch_cost.py b/tests/unit/proxy/common_utils/test_check_batch_cost.py index c974e55c0c9..eb2a5b7f66c 100644 --- a/tests/unit/proxy/common_utils/test_check_batch_cost.py +++ b/tests/unit/proxy/common_utils/test_check_batch_cost.py @@ -17,6 +17,8 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi import HTTPException +from litellm.types.utils import LiteLLMBatch + if TYPE_CHECKING: from litellm.batches.batch_utils import BatchCostUsageResult @@ -668,10 +670,16 @@ class TestCheckBatchCost: 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 = "file-output-123" - mock_response.model_dump_json.return_value = '{"id":"batch-1","status":"completed"}' + mock_response: Final = LiteLLMBatch( + id="batch-1", + status="completed", + output_file_id="file-output-123", + created_at=1, + completion_window="24h", + endpoint="/v1/chat/completions", + input_file_id="file-input-1", + object="batch", + ) 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"}) @@ -730,10 +738,10 @@ class TestCheckBatchCost: await check_batch_cost_instance.check_batch_cost() - assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1, ( + assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count >= 3, ( "Expected update() to be called exactly once for the completed job" ) - update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[1]["data"] + update_data = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args[1]["data"] assert update_data["batch_processed"] is True, ( "update() must include batch_processed=True when column is present" ) @@ -863,7 +871,7 @@ class TestCheckBatchCost: mock_update_database.assert_awaited_once() assert mock_update_database.call_args.kwargs["response_cost"] == 0.01 - assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1, ( + assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count >= 3, ( "the job must still be marked processed once cost tracking succeeds" ) @@ -1285,11 +1293,16 @@ class TestCheckBatchCost: 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 = terminal_status - mock_response.output_file_id = "file-output-123" - mock_response.error_file_id = None - mock_response.model_dump_json.return_value = f'{{"id":"batch-1","status":"{terminal_status}"}}' + mock_response: Final = LiteLLMBatch( + id="batch-1", + status=terminal_status, + output_file_id="file-output-123", + created_at=1, + completion_window="24h", + endpoint="/v1/chat/completions", + input_file_id="file-input-1", + object="batch", + ) 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"}) @@ -1352,8 +1365,8 @@ class TestCheckBatchCost: 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 - update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[1]["data"] + assert mock_prisma_client.db.litellm_managedobjecttable.update_many.call_count >= 3 + update_data = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args[1]["data"] assert update_data["batch_processed"] is True assert update_data["status"] == terminal_status, ( f"billed {terminal_status} batch must keep its real terminal status in the DB" @@ -1847,13 +1860,16 @@ class TestUnmanagedVertexRouting: router.resolve_model_name_from_model_id.return_value = "gemini-2.5-flash" router.get_model_ids.return_value = ["deploy-1"] - mock_response = MagicMock() - mock_response.status = "completed" - mock_response.output_file_id = "gs://bucket/out/predictions.jsonl" - mock_response.error_file_id = None - mock_response.completed_at = None - mock_response.created_at = None - mock_response.model_dump_json.return_value = '{"id":"8823717160934178816","status":"completed"}' + mock_response: Final = LiteLLMBatch( + id="8823717160934178816", + status="completed", + output_file_id="gs://bucket/out/predictions.jsonl", + created_at=1, + completion_window="24h", + endpoint="/v1/chat/completions", + input_file_id="file-input-1", + object="batch", + ) router.aretrieve_batch = AsyncMock(return_value=mock_response) router.get_deployment_credentials_with_provider = MagicMock( return_value={"vertex_project": "p", "vertex_location": "us-central1"} @@ -1921,8 +1937,8 @@ class TestUnmanagedVertexRouting: mock_logging_obj.async_success_handler.assert_awaited_once() assert mock_logging_obj.async_success_handler.call_args[1]["batch_cost"] == 0.01 - assert prisma.db.litellm_managedobjecttable.update.call_count == 1 - update_data = prisma.db.litellm_managedobjecttable.update.call_args[1]["data"] + assert prisma.db.litellm_managedobjecttable.update_many.call_count >= 3 + update_data = prisma.db.litellm_managedobjecttable.update_many.call_args[1]["data"] assert update_data["batch_processed"] is True assert update_data["status"] == "complete" @@ -2065,13 +2081,16 @@ class TestUnmanagedBedrockRouting: router.resolve_model_name_from_model_id.return_value = "claude-sonnet-4" router.get_model_ids.return_value = ["deploy-1"] - mock_response = MagicMock() - mock_response.status = "completed" - mock_response.output_file_id = "s3://bucket/out/predictions.jsonl" - mock_response.error_file_id = None - mock_response.completed_at = None - mock_response.created_at = None - mock_response.model_dump_json.return_value = f'{{"id":"{self._ARN}","status":"completed"}}' + mock_response: Final = LiteLLMBatch( + id=self._ARN, + status="completed", + output_file_id="s3://bucket/out/predictions.jsonl", + created_at=1, + completion_window="24h", + endpoint="/v1/chat/completions", + input_file_id="file-input-1", + object="batch", + ) router.aretrieve_batch = AsyncMock(return_value=mock_response) router.get_deployment_credentials_with_provider = MagicMock(return_value={"aws_region_name": "us-east-1"}) @@ -2135,8 +2154,8 @@ class TestUnmanagedBedrockRouting: mock_logging_obj.async_success_handler.assert_awaited_once() assert mock_logging_obj.async_success_handler.call_args[1]["batch_cost"] == 0.02 - assert prisma.db.litellm_managedobjecttable.update.call_count == 1 - update_data = prisma.db.litellm_managedobjecttable.update.call_args[1]["data"] + assert prisma.db.litellm_managedobjecttable.update_many.call_count >= 3 + update_data = prisma.db.litellm_managedobjecttable.update_many.call_args[1]["data"] assert update_data["batch_processed"] is True assert update_data["status"] == "complete" @@ -3003,13 +3022,20 @@ class TestMultiPodBatchCostClaim: @staticmethod def _router(): - response = MagicMock() - response.status = "completed" - response.output_file_id = _CLAIM_OUTPUT_FILE_ID - response.error_file_id = None - response.created_at = 1 - response.completed_at = 2 - response.model_dump_json.return_value = '{"id":"batch-456","status":"completed"}' + from litellm.types.utils import LiteLLMBatch + + response: Final = LiteLLMBatch( + id="batch-456", + status="completed", + output_file_id=_CLAIM_OUTPUT_FILE_ID, + error_file_id=None, + created_at=1, + completed_at=2, + completion_window="24h", + endpoint="/v1/chat/completions", + input_file_id="file-input-1", + object="batch", + ) deployment = MagicMock() deployment.litellm_params.custom_llm_provider = "openai" @@ -3089,8 +3115,9 @@ class TestMultiPodBatchCostClaim: ("openai/batch-test-model", 2, 0, "batch-test-model"), ), ) + @pytest.mark.parametrize("failure_envelope", (None, 1, "invalid", [1])) async def test_batch_accounting_preserves_models_and_mixed_results( - self, deployment_model: str, successes: int, failures: int, expected_model: str + self, deployment_model: str, successes: int, failures: int, expected_model: str, failure_envelope: object ) -> None: import base64 from datetime import datetime @@ -3125,7 +3152,7 @@ class TestMultiPodBatchCostClaim: }, "error": None, } - failure: Final = {"response": None, "error": {"code": "batch_expired", "message": "expired"}} + failure: Final = {"response": failure_envelope, "error": {"code": "batch_expired", "message": "expired"}} rows: Final = (success,) * successes + (failure,) * failures content: Final = b"\n".join(json.dumps(row).encode() for row in rows) @@ -3217,6 +3244,30 @@ class TestMultiPodBatchCostClaim: with patch("litellm.proxy.proxy_server.scheduler", scheduler): await guard._check_file_deletion_allowed(file_id) + @pytest.mark.asyncio + @pytest.mark.parametrize("status", ("completed", "failed", "expired", "cancelled")) + @pytest.mark.parametrize("metadata_text", ("before\x00after", r"\u0000")) + async def test_cost_claim_persists_terminal_metadata_without_a_later_write( + self, status: str, metadata_text: str + ) -> None: + row: Final = _FakeManagedObjectRow() + journal: Final = [] + prisma: Final = self._prisma(row, journal) + router: Final = self._router() + router.aretrieve_batch.return_value = router.aretrieve_batch.return_value.model_copy( + update={"status": status, "metadata": {"annotation": metadata_text}} + ) + prisma.db.litellm_managedobjecttable.update.side_effect = RuntimeError("final write unavailable") + worker: Final = self._instance(prisma, router) + with self._billing_patches(journal) as logging_obj: + await worker.check_batch_cost() + await worker.check_batch_cost() + assert row.batch_processed is True + assert row.status == ("complete" if status == "completed" else status) + assert json.loads(row.file_object)["status"] == status + assert json.loads(row.file_object)["metadata"] == {"annotation": metadata_text.replace("\x00", "")} + assert logging_obj.async_success_handler.await_count == 1 + @pytest.mark.asyncio async def test_winning_pod_claims_the_row_between_fetching_and_billing(self): """The claim flips batch_processed false -> true after the results are in hand and @@ -3228,13 +3279,9 @@ class TestMultiPodBatchCostClaim: with self._billing_patches(journal) as logging_obj: await self._instance(prisma, self._router()).check_batch_cost() - assert journal == ["fetch", "claim", "bill", "finalize"] - assert self._claim_calls(prisma) == [ - { - "where": {"id": "job-claim-1", "batch_processed": False}, - "data": {"batch_processed": True}, - } - ] + assert journal == ["fetch", "claim", "bill"] + assert row.status == "complete" + assert json.loads(row.file_object)["usage"] == {"prompt_tokens": 10, "completion_tokens": 5} logging_obj.async_success_handler.assert_awaited_once() assert row.batch_processed is True @@ -3254,18 +3301,17 @@ class TestMultiPodBatchCostClaim: assert journal == ["fetch"] logging_obj.async_success_handler.assert_not_awaited() - assert self._claim_calls(prisma) == [ - { - "where": {"id": "job-claim-1", "batch_processed": False}, - "data": {"batch_processed": True}, - } - ] + assert row.status == "in_progress" + assert json.loads(row.file_object)["status"] == "in_progress" @pytest.mark.asyncio - async def test_a_failed_spend_log_write_releases_the_claim(self): + @pytest.mark.parametrize("stored_as_json_string", (True, False)) + async def test_a_failed_spend_log_write_releases_the_claim(self, stored_as_json_string: bool) -> None: """A transient failure while billing a claimed batch must hand the row back, or its spend is silently lost instead of being retried on the next cycle.""" row = _FakeManagedObjectRow() + original_file_object: Final = row.file_object if stored_as_json_string else json.loads(row.file_object) + row.file_object = original_file_object journal = [] prisma = self._prisma(row, journal) @@ -3274,10 +3320,16 @@ class TestMultiPodBatchCostClaim: assert journal == ["fetch", "claim", "bill", "release"] assert row.batch_processed is False - assert self._claim_calls(prisma)[-1] == { - "where": {"id": "job-claim-1", "batch_processed": True}, - "data": {"batch_processed": False}, - } + assert row.status == "in_progress" + assert json.loads(row.file_object) == original_file_object + + with self._billing_patches(journal) as logging_obj: + await self._instance(prisma, self._router()).check_batch_cost() + await self._instance(prisma, self._router()).check_batch_cost() + + assert row.batch_processed is True + assert row.status == "complete" + logging_obj.async_success_handler.assert_awaited_once() @pytest.mark.asyncio async def test_a_worker_interrupted_mid_costing_leaves_the_batch_billable(self): @@ -3309,7 +3361,7 @@ class TestMultiPodBatchCostClaim: with self._billing_patches(survivor_journal) as survivor_logging: await self._instance(survivor_prisma, self._router()).check_batch_cost() - assert survivor_journal == ["fetch", "claim", "bill", "finalize"] + assert survivor_journal == ["fetch", "claim", "bill"] survivor_logging.async_success_handler.assert_awaited_once() assert row.batch_processed is True @@ -3340,7 +3392,7 @@ class TestMultiPodBatchCostClaim: finish_fetch.set() await asyncio.wait_for(costing, timeout=5) - assert journal == ["fetch", "claim", "bill", "finalize"] + assert journal == ["fetch", "claim", "bill"] assert row.batch_processed is True await self._run_deletion_guard(prisma, _CLAIM_OUTPUT_FILE_ID) @@ -3465,7 +3517,6 @@ class TestBatchLineItemsFromThePoller: f"read {_CLAIM_OUTPUT_FILE_ID}", "claim", "bill", - "finalize", f"read {_CLAIM_OUTPUT_FILE_ID}", "read file-input-1", "lines", @@ -3539,7 +3590,7 @@ class TestBatchLineItemsFromThePoller: await instance.check_batch_cost() await self._stored(instance) - assert journal[-4:] == ["finalize", f"read {_CLAIM_OUTPUT_FILE_ID}", "read file-input-1", "lines"] + assert journal[-4:] == ["bill", f"read {_CLAIM_OUTPUT_FILE_ID}", "read file-input-1", "lines"] ((line,),) = [payload["lines"] for payload in logger.payloads] assert line["request"] == json.loads(self._INPUT) assert line["response"] == json.loads(self._OUTPUT) @@ -3558,7 +3609,7 @@ class TestBatchLineItemsFromThePoller: await instance.check_batch_cost() await self._stored(instance) - assert [step for step in journal if not step.startswith("read")] == ["claim", "bill", "finalize", "lines"] + assert [step for step in journal if not step.startswith("read")] == ["claim", "bill", "lines"] assert row.batch_processed is True assert row.status == "complete" @@ -3574,7 +3625,7 @@ class TestBatchLineItemsFromThePoller: await instance.check_batch_cost() await self._stored(instance) - assert journal == [f"read {_CLAIM_OUTPUT_FILE_ID}", "claim", "bill", "finalize"] + assert journal == [f"read {_CLAIM_OUTPUT_FILE_ID}", "claim", "bill"] @pytest.mark.asyncio async def test_a_batch_where_every_line_failed_sends_its_error_lines_after_it_is_retired(self, monkeypatch): @@ -3751,7 +3802,6 @@ class TestBatchLineItemsFromThePoller: f"read {_CLAIM_OUTPUT_FILE_ID}", "claim", "bill", - "finalize", f"read {_CLAIM_OUTPUT_FILE_ID}", "read file-input-1", "lines started", @@ -3834,7 +3884,6 @@ class TestBatchLineItemsFromThePoller: f"read {_CLAIM_OUTPUT_FILE_ID}", "claim", "bill", - "finalize", "store batch-earlier", "store batch-456", ] @@ -3863,7 +3912,7 @@ class TestBatchLineItemsFromThePoller: await instance.check_batch_cost() await self._stored(instance) - assert journal == [f"read {_CLAIM_OUTPUT_FILE_ID}", "claim", "bill", "finalize"] + assert journal == [f"read {_CLAIM_OUTPUT_FILE_ID}", "claim", "bill"] assert logger.payloads == [] assert (row.batch_processed, row.status) == (True, "complete") assert f"the line items of batch batch-456 (job {row.id}) were not stored" in caplog.text @@ -3887,7 +3936,7 @@ class TestBatchLineItemsFromThePoller: await instance.check_batch_cost() await self._stored(instance) - assert journal == [f"read {_CLAIM_OUTPUT_FILE_ID}", "claim", "bill", "finalize"] + assert journal == [f"read {_CLAIM_OUTPUT_FILE_ID}", "claim", "bill"] assert logger.payloads == [] assert (row.batch_processed, row.status) == (True, "complete") assert f"the line items of batch batch-456 (job {row.id}) were not stored" in caplog.text