scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-29 06:33:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 200

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

Browse source 
Backport of #47184 to stable/1.105.x.
Cherry-picked from 3fb95c8fe3 (main).
Adapted to the line: keeps the line's _ManagedObjectRow name and adds its status property, ports postgres_storable_json into managed_files for the claim write (the rest of managed_files keeps the line's code), and drops the S3 batch line-item call that #41691 added on main. Tests converted from MagicMock responses to LiteLLMBatch where the claim now copies the response.
This commit is contained in:
joshua-berri 2026-10-20 16:04:27 -07:00 • committed by Mrinal Chanshetty
parent 74af50c64a
commit d74b58224a
6 changed files with 243 additions and 124 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -2,6 +2,7 @@
Polls LiteLLM_ManagedObjectTable to check if the batch job is complete, and if the cost has been tracked.
"""
import json
from collections.abc import Sequence
from dataclasses import replace as dataclasses_replace
from datetime import datetime, timedelta, timezone
@ -53,6 +54,9 @@ class _ManagedObjectRow(Protocol):
@property
def created_by(self) -> str | None: ...
@property
def status(self) -> str | None: ...
@property
def file_object(self) -> object: ...
@ -336,7 +340,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
@ -350,10 +354,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(
@ -362,7 +372,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
@ -374,7 +386,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(
@ -671,8 +687,6 @@ class CheckBatchCost:
@staticmethod
def _get_input_file_id(job: "_ManagedObjectRow") -> Optional[str]:
import json
from litellm.types.utils import LiteLLMBatch
file_object = job.file_object
@ -829,8 +843,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
@ -899,7 +912,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"
@ -918,7 +933,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)
@ -1027,10 +1042,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,
@ -1061,22 +1073,20 @@ 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:
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
elif response.status in PROVIDER_TERMINAL_BATCH_STATUSES:
from litellm.proxy.openai_files_endpoints.common_utils import (

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

@ -22,7 +22,7 @@ from uuid import NAMESPACE_URL, uuid5
import httpx
from fastapi import HTTPException
from pydantic import ValidationError
from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError
import litellm
from litellm import Router, verbose_logger
@ -205,6 +205,31 @@ def _managed_object_table(prisma_client: PrismaClient) -> _ManagedObjectTableAct
return prisma_client.db.litellm_managedobjecttable
_JSON_VALUE_ADAPTER: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue)
_POSTGRES_UNSTORABLE_CHARACTER: Final = "\x00"
def _storable_text(text: str) -> str:
return text.replace(_POSTGRES_UNSTORABLE_CHARACTER, "")
def _storable_json_value(value: JsonValue) -> JsonValue:
if isinstance(value, str):
return _storable_text(value)
if isinstance(value, list):
return [_storable_json_value(item) for item in value]
if isinstance(value, dict):
return {_storable_text(key): _storable_json_value(item) for key, item in value.items()}
return value
def postgres_storable_json(file_object: BaseModel) -> str:
dumped: Final = file_object.model_dump_json()
if "\\u0000" not in dumped:
return dumped
return _JSON_VALUE_ADAPTER.dump_json(_storable_json_value(_JSON_VALUE_ADAPTER.validate_json(dumped))).decode()
def _storage_metadata_of(file_object: OpenAIFileObject | None) -> Mapping[str, str]:
hidden_params: Final = cast( # cast-ok: _hidden_params is an untyped attribute the upload path sets
"Mapping[str, object]", getattr(file_object, "_hidden_params", None) or {}

18
litellm/batches/batch_utils.py
View file

@ -3,7 +3,7 @@ from collections.abc import Iterable, Iterator, Mapping
from dataclasses import dataclass
from dataclasses import replace as dataclasses_replace
from enum import Enum
from typing import Any, Final, Literal
from typing import Any, Final, Literal, cast
import litellm
from litellm._logging import verbose_logger
@ -748,9 +748,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] = batch_job_output_file.get("response", None) or {}
return _response.get("status_code", None) == 200
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
)

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 json
import logging
from types import MappingProxyType
from typing import Final, Literal
import httpx
import pytest
@ -2451,3 +2452,30 @@ async def test_flag_sends_every_vertex_row_down_the_native_path_when_a_model_is_
assert calls == []
assert (result.successful_requests, result.failed_requests) == (0, 1)
@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

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

@ -15,6 +15,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
@ -512,10 +514,16 @@ class TestCheckBatchCost:
mock_job.created_by = "user-1"
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(
@ -696,11 +704,15 @@ class TestCheckBatchCost:
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)
@ -764,15 +776,13 @@ class TestCheckBatchCost:
await check_batch_cost_instance.check_batch_cost()
assert (
mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1
), "Expected update() to be called exactly once for the completed job"
update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[
1
]["data"]
assert (
update_data["batch_processed"] is True
), "update() must include batch_processed=True when column is present"
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_many.call_args[1]["data"]
assert update_data["batch_processed"] is True, (
"update() must include batch_processed=True when column is present"
)
assert update_data["status"] == "complete"
@pytest.mark.asyncio
@ -899,7 +909,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"
)
@ -1374,11 +1384,15 @@ class TestCheckBatchCost:
return_value=[mock_job]
)
mock_response = MagicMock()
mock_response.status = terminal_status
mock_response.output_file_id = "file-output-123"
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)
@ -1446,12 +1460,8 @@ class TestCheckBatchCost:
mock_afile_content.await_count == 1
), 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
@ -1970,14 +1980,15 @@ 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(
@ -2050,8 +2061,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"
@ -2202,14 +2213,15 @@ 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(
@ -2280,8 +2292,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"
@ -3149,13 +3161,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"
@ -3237,8 +3256,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
@ -3273,7 +3293,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)
@ -3365,6 +3385,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
@ -3376,13 +3420,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
@ -3402,18 +3442,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)
@ -3422,10 +3461,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):
@ -3459,7 +3504,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
@ -3492,7 +3537,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)