mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-29 06:33:28 +00:00
fix(batches): preserve spend across malformed output and metadata failures (#47184)
Backport of #47184 to stable/1.104.x.
Cherry-picked from 3fb95c8fe3 (main).
Ported by hand onto the line's poller, whose surrounding code predates several main-only batch features: the claim now writes the terminal status and NUL-stripped file_object, a failed spend-log write restores both, and the managed-files postgres_storable_json sanitizer (from #45703 on main) is added to the line since this fix needs it.
This commit is contained in:
parent
bb1e40c3e5
commit
f7b0b09e75
6 changed files with 230 additions and 117 deletions
|
|
@ -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
|
||||
|
|
@ -336,7 +337,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 +351,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 +369,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 +383,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 +684,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
|
||||
|
|
@ -899,7 +910,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 +931,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)
|
||||
|
|
@ -1061,22 +1074,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:
|
||||
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}"
|
||||
)
|
||||
|
||||
elif response.status in PROVIDER_TERMINAL_BATCH_STATUSES:
|
||||
from litellm.proxy.openai_files_endpoints.common_utils import (
|
||||
|
|
|
|||
|
|
@ -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 {}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -744,9 +744,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
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -2417,3 +2418,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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -16,6 +16,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
|
||||
|
||||
|
|
@ -697,11 +699,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)
|
||||
|
|
@ -765,15 +771,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
|
||||
|
|
@ -900,7 +904,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"
|
||||
)
|
||||
|
||||
|
|
@ -1375,11 +1379,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)
|
||||
|
|
@ -1447,16 +1455,12 @@ 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
|
||||
), f"billed {terminal_status} batch must keep its real terminal status in the DB"
|
||||
assert 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_error_file_failures_add_to_failed_request_count(
|
||||
|
|
@ -1971,14 +1975,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(
|
||||
|
|
@ -2051,8 +2056,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"
|
||||
|
||||
|
|
@ -2203,14 +2208,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(
|
||||
|
|
@ -2281,8 +2287,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"
|
||||
|
||||
|
|
@ -3150,13 +3156,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"
|
||||
|
|
@ -3238,8 +3251,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
|
||||
|
|
@ -3274,7 +3288,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)
|
||||
|
||||
|
|
@ -3366,6 +3380,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
|
||||
|
|
@ -3377,13 +3415,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
|
||||
|
||||
|
|
@ -3403,18 +3437,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)
|
||||
|
||||
|
|
@ -3423,10 +3456,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):
|
||||
|
|
@ -3460,7 +3499,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
|
||||
|
||||
|
|
@ -3493,7 +3532,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)
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue