diff --git a/litellm/llms/bedrock/batches/handler.py b/litellm/llms/bedrock/batches/handler.py index 6efdd17f98d..801a11ca73b 100644 --- a/litellm/llms/bedrock/batches/handler.py +++ b/litellm/llms/bedrock/batches/handler.py @@ -1,3 +1,4 @@ +from collections.abc import Mapping from datetime import datetime from typing import TYPE_CHECKING, Any, Final, cast @@ -68,6 +69,19 @@ def _predict_output_file_uri(output_prefix: str, input_uri: str, job_id: str | N return f"{output_prefix}{job_id}/{input_basename}.out" +def _record_counts_from_response(response: Mapping[str, object]) -> BatchRequestCounts | None: + total_records: Final = response.get("totalRecordCount") + if not isinstance(total_records, int): + return None + success_records: Final = response.get("successRecordCount") + error_records: Final = response.get("errorRecordCount") + return BatchRequestCounts( + total=total_records, + completed=success_records if isinstance(success_records, int) else 0, + failed=error_records if isinstance(error_records, int) else 0, + ) + + def _to_epoch(value: Any) -> int | None: if value is None: return None @@ -271,11 +285,11 @@ class BedrockBatchesHandler: ``aws_external_id``). Unknown keys are ignored. Returns: - ``LiteLLMBatch`` shaped like an OpenAI Batch resource. Note that - ``request_counts`` is always ``(0, 0, 0)`` because - ``GetModelInvocationJob`` does not surface per-record counts; - callers that need accurate counts should parse - ``manifest.json.out`` from the output S3 prefix. + ``LiteLLMBatch`` shaped like an OpenAI Batch resource. + ``request_counts`` maps ``GetModelInvocationJob``'s + ``totalRecordCount`` / ``successRecordCount`` / ``errorRecordCount`` + when the provider reports them, and is ``None`` when it does not + (older botocore, or a status that omits counts). """ try: import boto3 @@ -386,7 +400,7 @@ class BedrockBatchesHandler: failed_at=completed_at if openai_status == "failed" else None, cancelled_at=completed_at if openai_status == "cancelled" else None, expired_at=completed_at if openai_status == "expired" else None, - request_counts=BatchRequestCounts(total=0, completed=0, failed=0), + request_counts=_record_counts_from_response(response), metadata=openai_batch_metadata, completion_window="24h", endpoint="/v1/chat/completions", diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index ddfdb56ac2c..992ed0d814d 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -1351,15 +1351,16 @@ def _completed_batch_safe_to_retire(response: "LiteLLMBatch") -> bool: provider response briefly lags before the output id populates). Retiring in that window loses the spend record forever. Retire only once we can prove there is nothing left to recover: the output file has actually arrived, or the provider - reports no successful request lines. When counts are unknown, stay eligible so - the next poller pass revisits it. (#37713) + reported a positive total with zero successful request lines, proving it + enumerated the batch and none succeeded. A zero or unknown total means counts + are unreported, so stay eligible and let the next poller pass revisit it. (#37713) """ if response.output_file_id is not None: return True request_counts = response.request_counts if request_counts is None: return False - return request_counts.completed == 0 + return request_counts.total > 0 and request_counts.completed == 0 async def update_batch_in_database( diff --git a/tests/test_litellm/llms/bedrock/batches/test_handler.py b/tests/test_litellm/llms/bedrock/batches/test_handler.py index d2dc89a7492..8c1357c345b 100644 --- a/tests/test_litellm/llms/bedrock/batches/test_handler.py +++ b/tests/test_litellm/llms/bedrock/batches/test_handler.py @@ -150,14 +150,38 @@ def test_handle_model_invocation_job_status_completed(patched_boto3): assert batch.completed_at == int(END_TIME.timestamp()) assert batch.failed_at is None assert batch.cancelled_at is None - # Per-record counts aren't reported by GetModelInvocationJob, so we leave - # them zeroed; consumers should parse manifest.json.out for accurate counts. - assert batch.request_counts.total == 0 + assert batch.request_counts is None assert batch.metadata["job_arn"] == JOB_ARN assert batch.metadata["output_file_uri"] == expected_out assert batch.metadata["output_s3_uri"] == OUTPUT_PREFIX +@pytest.mark.parametrize("success_count,error_count", [(100, 0), (86, 14)]) +def test_completed_job_maps_provider_record_counts(patched_boto3, success_count, error_count): + fake_client, _ = patched_boto3 + counted_response = _fake_boto3_response() + counted_response.update(totalRecordCount=100, successRecordCount=success_count, errorRecordCount=error_count) + fake_client.get_model_invocation_job.return_value = counted_response + + batch = BedrockBatchesHandler._handle_model_invocation_job_status(batch_id=JOB_ARN) + + assert batch.request_counts is not None + assert (batch.request_counts.total, batch.request_counts.completed, batch.request_counts.failed) == ( + 100, + success_count, + error_count, + ) + + +def test_missing_record_counts_leave_request_counts_none(patched_boto3): + fake_client, _ = patched_boto3 + fake_client.get_model_invocation_job.return_value = _fake_boto3_response() + + batch = BedrockBatchesHandler._handle_model_invocation_job_status(batch_id=JOB_ARN) + + assert batch.request_counts is None + + @pytest.mark.parametrize( "bedrock_status,openai_status", [ diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py index 7f84407f8b3..87cd2aaff1f 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py @@ -430,13 +430,15 @@ def test_add_internal_model_credentials_survives_a_failing_deployment_lookup(): assert data == {"batch_id": "unified-batch-id"} +from openai.types.batch import BatchRequestCounts + from litellm.proxy.openai_files_endpoints.common_utils import ( _completed_batch_safe_to_retire, ) def _completed_batch_for_retire( - output_file_id: str | None, completed: int | None = None + output_file_id: str | None, counts: BatchRequestCounts | None = None ) -> LiteLLMBatch: kwargs = dict( id="batch-1", @@ -449,26 +451,30 @@ def _completed_batch_for_retire( output_file_id=output_file_id, error_file_id=None, ) - if completed is not None: - kwargs["request_counts"] = {"total": completed, "completed": completed, "failed": 0} + if counts is not None: + kwargs["request_counts"] = counts return LiteLLMBatch(**kwargs) class TestCompletedBatchSafeToRetire: """A completed batch is only safe to retire from cost recovery once its output - file has arrived or the provider proves no successful lines (#37713).""" + file has arrived or the provider proves it enumerated a positive total of + request lines and none succeeded (#37713, LIT-6360).""" def test_output_file_present_is_safe(self): assert _completed_batch_safe_to_retire(_completed_batch_for_retire("file-out")) is True - def test_no_output_and_no_successful_lines_is_safe(self): - # Every request line errored -> nothing left to recover. - assert _completed_batch_safe_to_retire(_completed_batch_for_retire(None, completed=0)) is True + def test_no_output_and_synthesized_zero_counts_is_not_safe(self): + counts = BatchRequestCounts(total=0, completed=0, failed=0) + assert _completed_batch_safe_to_retire(_completed_batch_for_retire(None, counts)) is False def test_no_output_but_successful_lines_is_not_safe(self): - # The bug: output_file_id is lagging; retiring here loses the spend record. - assert _completed_batch_safe_to_retire(_completed_batch_for_retire(None, completed=5)) is False + counts = BatchRequestCounts(total=100, completed=100, failed=0) + assert _completed_batch_safe_to_retire(_completed_batch_for_retire(None, counts)) is False + + def test_no_output_and_all_lines_failed_is_safe(self): + counts = BatchRequestCounts(total=100, completed=0, failed=100) + assert _completed_batch_safe_to_retire(_completed_batch_for_retire(None, counts)) is True def test_no_output_and_unknown_counts_is_not_safe(self): - # Counts unknown -> stay eligible so the next poller pass revisits it. assert _completed_batch_safe_to_retire(_completed_batch_for_retire(None)) is False