mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
fix(bedrock): resolve the managed-batch output bucket on the inline accounting path too
A third path reads a completed batch's output file, and it could not resolve the bucket either. When cost is accounted from the retrieve itself rather than from the poller, the batch success handler calls _handle_completed_batch, which fetches the output file through _extract_file_access_credentials. That helper forwarded a whitelist covering Azure and Vertex, gcs_bucket_name included, but nothing for Bedrock, and retrieve_batch built its litellm_params through get_litellm_params, whose fixed signature drops the trusted credential snapshot. So the snapshot never reached the file read and it failed with "S3 bucket_name is required" for a bucket the deployment had configured, leaving the batch's cost unrecorded. Adding s3_bucket_name to that whitelist would not have worked. The Bedrock file config deliberately resolves the bucket only from the immutable server-side snapshot or the environment, never from a request param, because the bucket is what managed file ids are validated against. The snapshot is therefore what has to flow, exactly as it already does for the model-routed and cost-poller paths. retrieve_batch now re-adds the snapshot after get_litellm_params, the same way the file operations already do, the whitelist forwards it, and the proxy attaches it for router-routed managed batches from the deployment behind the unified id. Verified against a live proxy reading a real completed Bedrock batch: the cost row appears within seconds of the retrieve carrying the batch's real spend and usage, where before the read raised and no row was written. Resolving those credentials is best effort. A batch whose deployment no longer resolves, which happens when a model group is removed while batches are in flight, still serves its status instead of failing the request on the lookup. This matters for the OSS and polling-disabled configurations, where the retrieve path is the only thing that accounts for a batch at all.
This commit is contained in:
parent
460f0d29a9
commit
60fe4e464c
8 changed files with 154 additions and 1 deletions
|
|
@ -309,6 +309,7 @@ def _extract_file_access_credentials(litellm_params: dict | None) -> dict:
|
|||
"bucket_name",
|
||||
"timeout",
|
||||
"max_retries",
|
||||
"_litellm_internal_model_credentials",
|
||||
]
|
||||
for key in credential_keys:
|
||||
if key in litellm_params:
|
||||
|
|
|
|||
|
|
@ -22,6 +22,7 @@ from openai.types.batch import BatchRequestCounts
|
|||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.files.main import _add_trusted_model_credentials_to_litellm_params
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.llms.anthropic.batches.handler import AnthropicBatchesHandler
|
||||
from litellm.llms.azure.batches.handler import AzureBatchesAPI
|
||||
|
|
@ -527,6 +528,7 @@ def retrieve_batch(
|
|||
custom_llm_provider=custom_llm_provider,
|
||||
**kwargs,
|
||||
)
|
||||
_add_trusted_model_credentials_to_litellm_params(litellm_params, kwargs)
|
||||
if litellm_logging_obj is not None:
|
||||
litellm_logging_obj.update_from_kwargs(
|
||||
kwargs=kwargs,
|
||||
|
|
|
|||
|
|
@ -25,6 +25,7 @@ from litellm.proxy.openai_files_endpoints.common_utils import (
|
|||
_is_base64_encoded_unified_file_id,
|
||||
apply_team_provider_credentials,
|
||||
decode_model_from_file_id,
|
||||
add_internal_model_credentials_for_batch,
|
||||
encode_batch_response_ids,
|
||||
encode_file_id_with_model,
|
||||
ensure_batch_response_managed_file_ids,
|
||||
|
|
@ -537,6 +538,13 @@ async def retrieve_batch(
|
|||
detail={"error": "LLM Router not initialized. Ensure models added to proxy."},
|
||||
)
|
||||
|
||||
if unified_batch_id:
|
||||
add_internal_model_credentials_for_batch(
|
||||
data=data,
|
||||
llm_router=llm_router,
|
||||
model_id=get_model_id_from_unified_batch_id(unified_batch_id),
|
||||
)
|
||||
|
||||
response = await llm_router.aretrieve_batch(**data)
|
||||
response._hidden_params["unified_batch_id"] = unified_batch_id
|
||||
if unified_batch_id:
|
||||
|
|
|
|||
|
|
@ -465,6 +465,31 @@ def apply_team_provider_credentials(
|
|||
prepare_data_with_credentials(data=data, credentials=credentials)
|
||||
|
||||
|
||||
def add_internal_model_credentials_for_batch(
|
||||
data: dict,
|
||||
llm_router: "Router",
|
||||
model_id: str | None,
|
||||
) -> None:
|
||||
"""
|
||||
Attach the deployment's immutable server-side credential snapshot to a router-routed
|
||||
batch call (in-place).
|
||||
|
||||
Cost accounting for a completed batch reads the batch's output file, and the Bedrock
|
||||
file config resolves its bucket only from this snapshot, never from a request param,
|
||||
because the bucket is what managed file ids are validated against. Without it that
|
||||
read fails and the batch's cost is never recorded.
|
||||
"""
|
||||
if model_id is None:
|
||||
return
|
||||
try:
|
||||
credentials: Final = llm_router.get_deployment_credentials_with_provider(model_id=model_id)
|
||||
except Exception: # noqa: BLE001 # the snapshot only enables cost accounting; a batch whose deployment no longer resolves must still be retrievable
|
||||
return
|
||||
if credentials is None:
|
||||
return
|
||||
data["_litellm_internal_model_credentials"] = MappingProxyType(dict(credentials))
|
||||
|
||||
|
||||
def prepare_data_with_credentials(
|
||||
data: dict,
|
||||
credentials: dict,
|
||||
|
|
|
|||
|
|
@ -17,6 +17,7 @@ deterministic stand-ins so the arithmetic under test is the only variable.
|
|||
import json
|
||||
import os
|
||||
import sys
|
||||
from types import MappingProxyType
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
|
@ -1229,3 +1230,16 @@ async def test_calculate_batch_cost_and_usage_anthropic_end_to_end():
|
|||
assert cost == pytest.approx(1000 * 3e-6 / 2 + 8000 * 3e-7 / 2 + 2000 * 3.75e-6 / 2 + 200 * 15e-6 / 2)
|
||||
assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (11000, 200, 11200)
|
||||
assert models == ["claude-sonnet-4-5"]
|
||||
|
||||
|
||||
def test_extract_credentials_forwards_the_trusted_model_credential_snapshot():
|
||||
"""Bedrock resolves a batch's output bucket only from the immutable server-side
|
||||
snapshot, never from a request param, so cost accounting on the retrieve path cannot
|
||||
read the output file unless this key is forwarded. Without it the accounting raises
|
||||
"S3 bucket_name is required" for a bucket the deployment has configured, and the
|
||||
batch's cost is never recorded."""
|
||||
snapshot = MappingProxyType({"s3_bucket_name": "configured-bucket", "aws_region_name": "us-east-1"})
|
||||
|
||||
credentials = bu._extract_file_access_credentials({"_litellm_internal_model_credentials": snapshot})
|
||||
|
||||
assert credentials["_litellm_internal_model_credentials"] is snapshot
|
||||
|
|
|
|||
|
|
@ -28,6 +28,7 @@ import sys
|
|||
from contextlib import ExitStack
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict
|
||||
from types import MappingProxyType
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
|
@ -742,3 +743,33 @@ def test_resolve_timeout__httpx_timeout_returns_float_read():
|
|||
resolved = bm._resolve_timeout(_params(timeout=t), {}, "openai")
|
||||
assert isinstance(resolved, float)
|
||||
assert resolved == 99.0
|
||||
|
||||
|
||||
def test_retrieve__forwards_trusted_model_credentials_into_litellm_params(seams):
|
||||
"""The batch's cost is computed by reading its output file after the retrieve, and
|
||||
Bedrock resolves that bucket only from this immutable snapshot. get_litellm_params has
|
||||
a fixed signature that drops it, so without re-adding it here the snapshot never
|
||||
reaches the logging object and cost accounting fails on a bucket that is configured."""
|
||||
snapshot = MappingProxyType({"s3_bucket_name": "configured-bucket"})
|
||||
logging_obj = MagicMock()
|
||||
|
||||
bm.retrieve_batch(
|
||||
batch_id="batch-1",
|
||||
custom_llm_provider="openai",
|
||||
litellm_logging_obj=logging_obj,
|
||||
_litellm_internal_model_credentials=snapshot,
|
||||
)
|
||||
|
||||
litellm_params = logging_obj.update_from_kwargs.call_args.kwargs["litellm_params"]
|
||||
assert litellm_params["_litellm_internal_model_credentials"] is snapshot
|
||||
|
||||
|
||||
def test_retrieve__omits_trusted_model_credentials_when_not_supplied(seams):
|
||||
"""A retrieve with no snapshot must not invent an empty one, which would read as a
|
||||
configured bucket of nothing."""
|
||||
logging_obj = MagicMock()
|
||||
|
||||
bm.retrieve_batch(batch_id="batch-1", custom_llm_provider="openai", litellm_logging_obj=logging_obj)
|
||||
|
||||
litellm_params = logging_obj.update_from_kwargs.call_args.kwargs["litellm_params"]
|
||||
assert "_litellm_internal_model_credentials" not in litellm_params
|
||||
|
|
|
|||
|
|
@ -1138,7 +1138,11 @@ async def test_retrieve__unified_batch_id_routes_to_router(retrieve_harness):
|
|||
# DISPATCH - router fired, direct litellm did not.
|
||||
assert retrieve_harness.router_aretrieve.call_count == 1
|
||||
retrieve_harness.litellm_aretrieve.assert_not_called()
|
||||
retrieve_harness.creds_resolver.assert_not_called()
|
||||
|
||||
# Credentials are resolved for the deployment behind the unified id so the batch's
|
||||
# output file can be read for cost accounting. This id resolves to nothing here, and
|
||||
# the retrieve must still serve the batch rather than fail on the lookup.
|
||||
retrieve_harness.creds_resolver.assert_called_once_with(model_id="gpt-4o-mini")
|
||||
|
||||
# router receives the (still-encoded) batch id verbatim - this layer does
|
||||
# not decode it for the unified path.
|
||||
|
|
|
|||
|
|
@ -95,3 +95,71 @@ def test_apply_unified_file_ids_swaps_all_three_ids():
|
|||
"unified-out",
|
||||
"unified-err",
|
||||
)
|
||||
|
||||
|
||||
# =========================================================================== #
|
||||
# add_internal_model_credentials_for_batch - the snapshot that lets a completed
|
||||
# batch's output file be read, and therefore its cost be recorded
|
||||
# =========================================================================== #
|
||||
|
||||
|
||||
def test_add_internal_model_credentials_attaches_an_immutable_snapshot():
|
||||
"""Cost accounting for a completed batch reads its output file, and Bedrock resolves
|
||||
that bucket only from this snapshot. It must be immutable so nothing downstream can
|
||||
redirect the bucket that managed file ids are validated against."""
|
||||
from litellm.proxy.openai_files_endpoints.common_utils import (
|
||||
add_internal_model_credentials_for_batch,
|
||||
)
|
||||
|
||||
router = MagicMock()
|
||||
router.get_deployment_credentials_with_provider = MagicMock(
|
||||
return_value={"s3_bucket_name": "configured-bucket", "aws_region_name": "us-east-1"}
|
||||
)
|
||||
data = {"batch_id": "unified-batch-id"}
|
||||
|
||||
add_internal_model_credentials_for_batch(data=data, llm_router=router, model_id="deployment-1")
|
||||
|
||||
snapshot = data["_litellm_internal_model_credentials"]
|
||||
assert snapshot["s3_bucket_name"] == "configured-bucket"
|
||||
assert isinstance(snapshot, MappingProxyType)
|
||||
with pytest.raises(TypeError):
|
||||
snapshot["s3_bucket_name"] = "attacker-bucket"
|
||||
router.get_deployment_credentials_with_provider.assert_called_once_with(model_id="deployment-1")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model_id, credentials",
|
||||
[(None, {"s3_bucket_name": "b"}), ("deployment-1", None)],
|
||||
ids=["no-model-id", "deployment-has-no-credentials"],
|
||||
)
|
||||
def test_add_internal_model_credentials_is_a_noop_without_a_resolvable_deployment(model_id, credentials):
|
||||
"""An unroutable batch must be left alone rather than given an empty snapshot, which
|
||||
would look like a configured bucket of nothing."""
|
||||
from litellm.proxy.openai_files_endpoints.common_utils import (
|
||||
add_internal_model_credentials_for_batch,
|
||||
)
|
||||
|
||||
router = MagicMock()
|
||||
router.get_deployment_credentials_with_provider = MagicMock(return_value=credentials)
|
||||
data = {"batch_id": "unified-batch-id"}
|
||||
|
||||
add_internal_model_credentials_for_batch(data=data, llm_router=router, model_id=model_id)
|
||||
|
||||
assert "_litellm_internal_model_credentials" not in data
|
||||
|
||||
|
||||
def test_add_internal_model_credentials_survives_a_failing_deployment_lookup():
|
||||
"""The snapshot only enables cost accounting, so a batch whose deployment no longer
|
||||
resolves, which happens when a model group is removed while batches are in flight,
|
||||
must still be retrievable rather than failing the request on the lookup."""
|
||||
from litellm.proxy.openai_files_endpoints.common_utils import (
|
||||
add_internal_model_credentials_for_batch,
|
||||
)
|
||||
|
||||
router = MagicMock()
|
||||
router.get_deployment_credentials_with_provider = MagicMock(side_effect=KeyError("deployment-gone"))
|
||||
data = {"batch_id": "unified-batch-id"}
|
||||
|
||||
add_internal_model_credentials_for_batch(data=data, llm_router=router, model_id="deployment-gone")
|
||||
|
||||
assert data == {"batch_id": "unified-batch-id"}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue