scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-23 05:34:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 24

fix(batches): preserve spend across malformed output and metadata failures (#47184)

Browse source 
* 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>
This commit is contained in:
joshua-berri 2026-10-17 19:10:22 -05:00 • committed by GitHub
parent 6c68c2cbef
commit 3fb95c8fe3
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 209 additions and 119 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

68
enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
View file

@ -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)

8
enterprise/litellm_enterprise/proxy/hooks/managed_files.py
View file

@ -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,

16
litellm/batches/batch_utils.py
View file

@ -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

28
tests/unit/batches/test_batch_utils.py
View file

@ -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)

21
tests/unit/enterprise/proxy/test_managed_files_access_check.py
View file

@ -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

187
tests/unit/proxy/common_utils/test_check_batch_cost.py
View file

@ -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