mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-12 23:01:41 +00:00
fix(batches): handle provider cancellation and file cleanup gaps
This commit is contained in:
Yuneng Jiang
parent
a096dd615c
commit
7bff9bf9a2
6 changed files with 258 additions and 60 deletions
|
|
@ -7,7 +7,7 @@ from contextlib import suppress
|
|||
from functools import cache
|
||||
from itertools import chain
|
||||
from types import MappingProxyType
|
||||
from typing import Any, Final, TypeAlias, TypedDict
|
||||
from typing import Any, Final, Literal, TypeAlias, TypedDict
|
||||
from urllib.parse import unquote
|
||||
|
||||
import httpx
|
||||
|
|
@ -60,11 +60,8 @@ from litellm.utils import get_llm_provider
|
|||
from ..base_aws_llm import BaseAWSLLM
|
||||
from ..common_utils import BedrockError, merge_bedrock_aws_request_params, resolve_s3_encryption_key_id
|
||||
|
||||
# litellm_params key used to hand the SigV4-signed GET headers from
|
||||
# `transform_file_content_request` to `validate_environment` (the only hook
|
||||
# the shared file-content HTTP handler exposes for setting request headers).
|
||||
# Same pattern as the `upload_url` handoff in `transform_create_file_request`.
|
||||
S3_SIGNED_GET_HEADERS_PARAM: Final = "_s3_signed_get_headers"
|
||||
S3_SIGNED_REQUEST_HEADERS_PARAM: Final = "_s3_signed_request_headers"
|
||||
S3_DELETE_FILE_ID_PARAM: Final = "_s3_delete_file_id"
|
||||
|
||||
# litellm_params key carrying the size of the body uploaded to S3, handed from
|
||||
# `transform_create_file_request` to `transform_create_file_response`.
|
||||
|
|
@ -291,7 +288,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
|
|||
) -> dict:
|
||||
result: Final[dict[str, object]] = {}
|
||||
result.update(headers)
|
||||
signed_headers: Final = litellm_params.pop(S3_SIGNED_GET_HEADERS_PARAM, None)
|
||||
signed_headers: Final = litellm_params.pop(S3_SIGNED_REQUEST_HEADERS_PARAM, None)
|
||||
if isinstance(signed_headers, Mapping):
|
||||
result.update(signed_headers) # any-ok: untyped handoff headers
|
||||
# otherwise no extra headers - AWS credentials are handled by BaseAWSLLM
|
||||
|
|
@ -1187,18 +1184,31 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
|
|||
def transform_delete_file_request(
|
||||
self,
|
||||
file_id: str,
|
||||
optional_params: dict,
|
||||
litellm_params: dict,
|
||||
) -> tuple[str, dict]:
|
||||
raise NotImplementedError("BedrockFilesConfig does not support file deletion")
|
||||
optional_params: Mapping[str, object],
|
||||
litellm_params: MutableMapping[str, object],
|
||||
) -> tuple[str, dict[str, str]]:
|
||||
request: Final = self._transform_s3_file_request(
|
||||
file_id=file_id, method="DELETE", optional_params=optional_params, litellm_params=litellm_params
|
||||
)
|
||||
litellm_params[S3_DELETE_FILE_ID_PARAM] = file_id
|
||||
return request
|
||||
|
||||
def transform_delete_file_response(
|
||||
self,
|
||||
raw_response: httpx.Response,
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
litellm_params: dict,
|
||||
litellm_params: Mapping[str, object],
|
||||
) -> FileDeleted:
|
||||
raise NotImplementedError("BedrockFilesConfig does not support file deletion")
|
||||
if raw_response.status_code != 204:
|
||||
raise BedrockError(
|
||||
status_code=raw_response.status_code if raw_response.status_code >= 400 else 502,
|
||||
message=raw_response.text or f"S3 file deletion returned HTTP {raw_response.status_code}",
|
||||
headers=raw_response.headers,
|
||||
)
|
||||
file_id: Final = litellm_params.get(S3_DELETE_FILE_ID_PARAM)
|
||||
if not isinstance(file_id, str) or not file_id:
|
||||
raise ValueError("Missing file id for Bedrock file deletion response")
|
||||
return FileDeleted(id=file_id, deleted=True, object="file")
|
||||
|
||||
def transform_list_files_request(
|
||||
self,
|
||||
|
|
@ -1233,6 +1243,18 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
|
|||
if not file_id:
|
||||
raise ValueError("file_id is required for Bedrock file content retrieval")
|
||||
|
||||
return self._transform_s3_file_request(
|
||||
file_id=file_id, method="GET", optional_params=optional_params, litellm_params=litellm_params
|
||||
)
|
||||
|
||||
def _transform_s3_file_request(
|
||||
self,
|
||||
*,
|
||||
file_id: str,
|
||||
method: Literal["GET", "DELETE"],
|
||||
optional_params: Mapping[str, object],
|
||||
litellm_params: MutableMapping[str, object],
|
||||
) -> tuple[str, dict[str, str]]:
|
||||
s3_uri: Final = extract_s3_uri_from_file_id(file_id)
|
||||
bucket_name, object_key = _validate_file_id_against_configured_buckets(
|
||||
s3_uri=s3_uri,
|
||||
|
|
@ -1240,40 +1262,32 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
|
|||
allow_legacy_cloud_file_ids=should_allow_legacy_cloud_file_ids(litellm_params),
|
||||
)
|
||||
|
||||
# The shared file-content handler passes optional_params={}, so AWS
|
||||
# credentials/region arrive via litellm_params here (unlike the upload
|
||||
# path). s3_region_name wins over aws_region_name, same priority as
|
||||
# get_complete_file_url above.
|
||||
merged_params: Final[dict[str, object]] = {}
|
||||
merged_params.update(litellm_params)
|
||||
merged_params.update(optional_params)
|
||||
request_params: Final = _BedrockS3RequestParams.model_validate(merged_params)
|
||||
request_params: Final = _BedrockS3RequestParams.model_validate({**litellm_params, **optional_params})
|
||||
|
||||
region_preference: Final = request_params.s3_region_name or request_params.aws_region_name
|
||||
region_params: Final[dict[str, str | None]] = {"aws_region_name": region_preference}
|
||||
aws_region_name: Final = self._get_aws_region_name(optional_params=region_params, model="")
|
||||
|
||||
s3_endpoint_url = (
|
||||
s3_endpoint_url: Final = (
|
||||
request_params.s3_endpoint_url or f"https://s3.{aws_region_name}.{get_aws_dns_suffix(aws_region_name)}"
|
||||
).rstrip("/")
|
||||
url: Final = f"{s3_endpoint_url}/{bucket_name}/{encode_s3_object_key_for_url(object_key)}"
|
||||
|
||||
litellm_params[S3_SIGNED_GET_HEADERS_PARAM] = self._sign_s3_get_request(
|
||||
litellm_params[S3_SIGNED_REQUEST_HEADERS_PARAM] = self._sign_s3_request_without_body(
|
||||
api_base=url,
|
||||
aws_region_name=aws_region_name,
|
||||
request_params=request_params,
|
||||
method=method,
|
||||
)
|
||||
return url, {}
|
||||
|
||||
def _sign_s3_get_request(
|
||||
def _sign_s3_request_without_body(
|
||||
self,
|
||||
api_base: str,
|
||||
aws_region_name: str,
|
||||
request_params: _BedrockS3RequestParams,
|
||||
method: Literal["GET", "DELETE"] = "GET",
|
||||
) -> dict[str, str]:
|
||||
"""
|
||||
SigV4-sign an S3 GetObject request, mirroring `_sign_s3_request` (PUT).
|
||||
"""
|
||||
try:
|
||||
import hashlib
|
||||
|
||||
|
|
@ -1297,7 +1311,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
|
|||
|
||||
empty_body_hash: Final = hashlib.sha256(b"").hexdigest()
|
||||
aws_request: Final = AWSRequest( # any-ok: botocore AWSRequest is untyped
|
||||
method="GET",
|
||||
method=method,
|
||||
url=api_base,
|
||||
headers={"x-amz-content-sha256": empty_body_hash},
|
||||
)
|
||||
|
|
|
|||
|
|
@ -131,7 +131,12 @@ failures up to three times. Teardown attempts every registered cleanup before
|
|||
reporting failures as test errors. Already deleted files and batches that are
|
||||
terminal are safe to clean up again. Managed batch cancellation polls for up to eleven minutes
|
||||
before input deletion: the ten-minute provider window plus a propagation margin.
|
||||
Raw and model-encoded inputs can be deleted after cancellation is accepted
|
||||
Accepted cancellation may still report validating or in_progress while the provider
|
||||
updates its state. Raw and model-encoded batches are polled until cancelling or
|
||||
terminal before input deletion. OpenAI and Azure lifecycle cleanup also deletes
|
||||
output and error files returned by terminal batches. Bedrock deletion uses a signed S3 DELETE
|
||||
restricted to the configured storage buckets and managed file prefixes. The low-RPM
|
||||
test submits with its restricted key and cleans up with the test administrator key
|
||||
|
||||
Azure input uploads request `expires_after` anchored to `created_at` with
|
||||
`seconds=1209600`, and the lifecycle tests check the returned expiry. This is a
|
||||
|
|
|
|||
|
|
@ -1,16 +1,17 @@
|
|||
from builtins import ExceptionGroup
|
||||
from collections.abc import Callable
|
||||
from itertools import count
|
||||
from time import monotonic, sleep
|
||||
from typing import Final, Protocol
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
from batch_client import BatchObject, FileDeleteResponse
|
||||
from capabilities import is_managed_id
|
||||
from e2e_http import NetworkError, RateLimitedError, Result, Success, UnknownApiError
|
||||
from pydantic import BaseModel
|
||||
|
||||
CLEANUP_DELAYS: Final = (1.0, 2.0, 4.0)
|
||||
BATCH_TERMINAL_STATUSES: Final = frozenset({"completed", "failed", "expired", "cancelled"})
|
||||
BATCH_PENDING_STATUSES: Final = frozenset({"validating", "in_progress", "finalizing", "cancelling"})
|
||||
BATCH_CANCEL_TIMEOUT_SECONDS: Final = 660.0
|
||||
BATCH_CANCEL_POLL_SECONDS: Final = 10.0
|
||||
|
||||
|
|
@ -63,6 +64,7 @@ def cleanup_batch(
|
|||
*,
|
||||
key: str,
|
||||
provider: str | None = None,
|
||||
delete_output_files: bool = False,
|
||||
wait: Callable[[float], None] = sleep,
|
||||
clock: Callable[[], float] = monotonic,
|
||||
) -> None:
|
||||
|
|
@ -72,18 +74,28 @@ def cleanup_batch(
|
|||
f"Retrieve batch {batch_id} for cleanup",
|
||||
)
|
||||
if fetched.status in BATCH_TERMINAL_STATUSES:
|
||||
if delete_output_files:
|
||||
_cleanup_batch_outputs(client, fetched, key=key, provider=provider)
|
||||
return
|
||||
if fetched.status == "cancelling" and not needs_terminal_state:
|
||||
return
|
||||
if fetched.status != "cancelling":
|
||||
result: Final = cleanup_result(lambda: client.cancel_batch(batch_id, key=key, provider=provider))
|
||||
if not (isinstance(result, UnknownApiError) and result.status_code in {400, 409}):
|
||||
cancelled: Final = _require_cleanup_success(result, f"Cancel batch {batch_id}")
|
||||
assert cancelled.status in BATCH_TERMINAL_STATUSES | {"cancelling"}, (
|
||||
f"Cancel batch {batch_id} left status {cancelled.status}"
|
||||
)
|
||||
if not needs_terminal_state:
|
||||
return
|
||||
result: Final = (
|
||||
Success(status_code=200, data=fetched)
|
||||
if fetched.status == "cancelling"
|
||||
else cleanup_result(lambda: client.cancel_batch(batch_id, key=key, provider=provider))
|
||||
)
|
||||
conflicted: Final = isinstance(result, UnknownApiError) and result.status_code in {400, 409}
|
||||
if not conflicted:
|
||||
cancelled: Final = _require_cleanup_success(result, f"Cancel batch {batch_id}")
|
||||
assert cancelled.status in BATCH_TERMINAL_STATUSES | BATCH_PENDING_STATUSES, (
|
||||
f"Cancel batch {batch_id} left status {cancelled.status}"
|
||||
)
|
||||
if cancelled.status in BATCH_TERMINAL_STATUSES:
|
||||
if delete_output_files:
|
||||
_cleanup_batch_outputs(client, cancelled, key=key, provider=provider)
|
||||
return
|
||||
if cancelled.status == "cancelling" and not needs_terminal_state:
|
||||
return
|
||||
deadline: Final = clock() + BATCH_CANCEL_TIMEOUT_SECONDS
|
||||
for current in (
|
||||
_require_cleanup_success(
|
||||
|
|
@ -93,11 +105,36 @@ def cleanup_batch(
|
|||
for _ in count()
|
||||
):
|
||||
if current.status in BATCH_TERMINAL_STATUSES:
|
||||
if delete_output_files:
|
||||
_cleanup_batch_outputs(client, current, key=key, provider=provider)
|
||||
return
|
||||
assert current.status == "cancelling", f"Cancel batch {batch_id} left status {current.status}"
|
||||
if not needs_terminal_state:
|
||||
assert current.status in ({"cancelling"} if conflicted else BATCH_PENDING_STATUSES), (
|
||||
f"Cancel batch {batch_id} left status {current.status}"
|
||||
)
|
||||
if current.status == "cancelling" and not needs_terminal_state:
|
||||
return
|
||||
assert clock() < deadline, (
|
||||
f"Batch {batch_id} cancellation did not finish within {BATCH_CANCEL_TIMEOUT_SECONDS}s"
|
||||
)
|
||||
wait(BATCH_CANCEL_POLL_SECONDS)
|
||||
|
||||
|
||||
def _cleanup_batch_outputs(client: BatchCleanupClient, batch: BatchObject, *, key: str, provider: str | None) -> None:
|
||||
errors: Final = tuple(
|
||||
error
|
||||
for file_id in dict.fromkeys((batch.output_file_id, batch.error_file_id))
|
||||
if file_id is not None and file_id != batch.input_file_id
|
||||
if (error := _output_cleanup_error(client, file_id, key=key, provider=provider)) is not None
|
||||
)
|
||||
if errors:
|
||||
raise ExceptionGroup(f"Batch {batch.id} output cleanup failed", errors)
|
||||
|
||||
|
||||
def _output_cleanup_error(
|
||||
client: BatchCleanupClient, file_id: str, *, key: str, provider: str | None
|
||||
) -> Exception | None:
|
||||
try:
|
||||
cleanup_file(client, file_id, key=key, provider=provider)
|
||||
except Exception as error:
|
||||
return error
|
||||
return None
|
||||
|
|
|
|||
|
|
@ -4,7 +4,6 @@ from dataclasses import dataclass, field
|
|||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
from batch_cleanup import BATCH_CANCEL_TIMEOUT_SECONDS, CLEANUP_DELAYS, cleanup_batch, cleanup_file, cleanup_result
|
||||
from batch_client import AZURE_FILE_EXPIRY_SECONDS, BatchObject, FileDeleteResponse, batch_upload_form
|
||||
from capabilities import CAPABILITIES, Capability
|
||||
|
|
@ -219,6 +218,74 @@ class TestBatchCancellation:
|
|||
cleanup_batch(client, "batch-1", key="test-key", provider="azure")
|
||||
client.calls.assert_done()
|
||||
|
||||
@pytest.mark.parametrize("batch_id", ["batch-1", MANAGED_BATCH_ID])
|
||||
@pytest.mark.parametrize("pending_status", ["validating", "in_progress"])
|
||||
def test_accepted_cancellation_waits_through_stale_provider_status(
|
||||
self, batch_id: str, pending_status: str
|
||||
) -> None:
|
||||
client: Final = CleanupClient(
|
||||
calls=ExpectedCalls(
|
||||
iter(
|
||||
(
|
||||
f"retrieve vertex_ai {batch_id}",
|
||||
f"cancel vertex_ai {batch_id}",
|
||||
f"retrieve vertex_ai {batch_id}",
|
||||
f"retrieve vertex_ai {batch_id}",
|
||||
f"retrieve vertex_ai {batch_id}",
|
||||
"delete vertex_ai file-1",
|
||||
"delete key test-key",
|
||||
)
|
||||
)
|
||||
),
|
||||
batches=iter((batch("validating"), batch(pending_status), batch(pending_status), batch("cancelled"))),
|
||||
cancellations=iter((batch(pending_status),)),
|
||||
files=iter((deleted_file(),)),
|
||||
)
|
||||
delays: Final = ExpectedCalls(iter((10.0, 10.0)))
|
||||
manager: Final = ResourceManager(client=client, strict_cleanup=True)
|
||||
key: Final = manager.key()
|
||||
manager.defer(lambda: cleanup_file(client, "file-1", key=key, provider="vertex_ai"))
|
||||
manager.defer(lambda: cleanup_batch(client, batch_id, key=key, provider="vertex_ai", wait=delays))
|
||||
manager.teardown()
|
||||
client.calls.assert_done()
|
||||
delays.assert_done()
|
||||
|
||||
@pytest.mark.parametrize("output_delete_fails", [False, True])
|
||||
def test_batch_that_completed_before_cleanup_deletes_output_and_error_files(
|
||||
self, output_delete_fails: bool
|
||||
) -> None:
|
||||
client: Final = CleanupClient(
|
||||
calls=ExpectedCalls(
|
||||
iter(("retrieve openai batch-1", "delete openai file-output", "delete openai file-error"))
|
||||
),
|
||||
batches=iter(
|
||||
(
|
||||
Success(
|
||||
status_code=200,
|
||||
data=BatchObject(
|
||||
id="batch-1",
|
||||
status="completed",
|
||||
input_file_id="file-input",
|
||||
output_file_id="file-output",
|
||||
error_file_id="file-error",
|
||||
),
|
||||
),
|
||||
)
|
||||
),
|
||||
files=iter(
|
||||
(
|
||||
UnknownApiError(status_code=403, body="forbidden") if output_delete_fails else deleted_file(),
|
||||
deleted_file(),
|
||||
)
|
||||
),
|
||||
)
|
||||
if output_delete_fails:
|
||||
with pytest.raises(ExceptionGroup, match="output cleanup failed"):
|
||||
cleanup_batch(client, "batch-1", key="test-key", provider="openai", delete_output_files=True)
|
||||
else:
|
||||
cleanup_batch(client, "batch-1", key="test-key", provider="openai", delete_output_files=True)
|
||||
client.calls.assert_done()
|
||||
|
||||
@pytest.mark.parametrize("status", ["completed", "in_progress"])
|
||||
def test_cancellation_conflict_is_accepted_only_when_batch_became_inactive(self, status: str) -> None:
|
||||
client: Final = CleanupClient(
|
||||
|
|
|
|||
|
|
@ -25,7 +25,7 @@ from datetime import datetime, timedelta, timezone
|
|||
import pytest
|
||||
from pydantic import BaseModel
|
||||
|
||||
from e2e_config import PROXY_BASE_URL, unique_marker
|
||||
from e2e_config import MASTER_KEY, PROXY_BASE_URL, unique_marker
|
||||
|
||||
from batch_cleanup import cleanup_batch, cleanup_file
|
||||
from batch_client import (
|
||||
|
|
@ -257,7 +257,9 @@ def test_batch_lifecycle(
|
|||
require_successful_call(created)
|
||||
batch = BatchObject.model_validate_json(created.body)
|
||||
resources.defer(
|
||||
lambda: cleanup_batch(client, batch.id, key=key, provider=provider)
|
||||
lambda: cleanup_batch(
|
||||
client, batch.id, key=key, provider=provider, delete_output_files=cap.provider in {"openai", "azure"}
|
||||
)
|
||||
)
|
||||
|
||||
assert batch.id, f"create returned no batch id (body={created.body[:200]})"
|
||||
|
|
@ -801,7 +803,7 @@ class TestBatchEnqueuedTokenLimit:
|
|||
"""
|
||||
|
||||
def _upload_batch_file(
|
||||
self, client: BatchClient, resources: ResourceManager, key: str
|
||||
self, client: BatchClient, resources: ResourceManager, key: str, *, cleanup_key: str | None = None
|
||||
) -> FileObject:
|
||||
file = unwrap(
|
||||
client.upload_file(
|
||||
|
|
@ -811,7 +813,7 @@ class TestBatchEnqueuedTokenLimit:
|
|||
key=key,
|
||||
)
|
||||
)
|
||||
resources.defer(lambda: cleanup_file(client, file.id, key=key))
|
||||
resources.defer(lambda: cleanup_file(client, file.id, key=cleanup_key or key))
|
||||
return file
|
||||
|
||||
def _generate_enqueued_key(
|
||||
|
|
@ -848,7 +850,7 @@ class TestBatchEnqueuedTokenLimit:
|
|||
marker="rpm",
|
||||
rpm_limit=BATCH_RL_RPM_LIMIT,
|
||||
)
|
||||
file = self._upload_batch_file(client, resources, key)
|
||||
file = self._upload_batch_file(client, resources, key, cleanup_key=MASTER_KEY)
|
||||
|
||||
created = client.create_batch(body=BatchCreateBody(input_file_id=file.id), key=key)
|
||||
|
||||
|
|
@ -859,7 +861,7 @@ class TestBatchEnqueuedTokenLimit:
|
|||
)
|
||||
require_successful_call(created)
|
||||
batch = BatchObject.model_validate_json(created.body)
|
||||
resources.defer(lambda: cleanup_batch(client, batch.id, key=key))
|
||||
resources.defer(lambda: cleanup_batch(client, batch.id, key=MASTER_KEY, delete_output_files=True))
|
||||
|
||||
@pytest.mark.covers(
|
||||
"quota_management.ratelimit.batch_enqueued_tokens.blocks_when_exhausted",
|
||||
|
|
|
|||
|
|
@ -5,6 +5,8 @@ Test bedrock files transformation functionality
|
|||
import json
|
||||
import os
|
||||
from collections.abc import Mapping
|
||||
from contextlib import AsyncExitStack, closing
|
||||
from typing import Final
|
||||
from unittest.mock import MagicMock
|
||||
from urllib.parse import unquote, urlparse
|
||||
|
||||
|
|
@ -1855,6 +1857,77 @@ class TestBedrockBatchNonChatEndpointRecords:
|
|||
]
|
||||
|
||||
|
||||
class TestBedrockFileDeletion:
|
||||
S3_URI: Final = "s3://my-bucket/litellm-bedrock-files-model-abc.jsonl"
|
||||
URL: Final = "https://s3.us-west-2.amazonaws.com/my-bucket/litellm-bedrock-files-model-abc.jsonl"
|
||||
|
||||
def test_delete_file_sends_signed_delete_and_returns_matching_id(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
import httpx
|
||||
import respx
|
||||
|
||||
import litellm
|
||||
from litellm.llms.custom_httpx.http_handler import HTTPHandler
|
||||
|
||||
monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket")
|
||||
with respx.mock, closing(HTTPHandler()) as client:
|
||||
route: Final = respx.delete(self.URL).mock(return_value=httpx.Response(204))
|
||||
deleted: Final = litellm.file_delete(
|
||||
file_id=self.S3_URI, custom_llm_provider="bedrock", client=client,
|
||||
aws_access_key_id="AKIAEXAMPLE", aws_secret_access_key="test-secret", aws_region_name="us-west-2",
|
||||
)
|
||||
assert route.call_count == 1
|
||||
request: Final = route.calls[0].request
|
||||
assert request.content == b""
|
||||
signed: Final = AWSRequest(method="DELETE", url=self.URL, headers={
|
||||
"X-Amz-Date": request.headers["X-Amz-Date"],
|
||||
"X-Amz-Content-SHA256": request.headers["X-Amz-Content-SHA256"],
|
||||
})
|
||||
signed.context["timestamp"] = request.headers["X-Amz-Date"]
|
||||
auth: Final = S3SigV4Auth(Credentials("AKIAEXAMPLE", "test-secret"), "s3", "us-west-2")
|
||||
signature: Final = auth.signature(auth.string_to_sign(signed, auth.canonical_request(signed)), signed)
|
||||
assert request.headers["Authorization"].endswith(f"Signature={signature}")
|
||||
assert deleted.id == self.S3_URI and deleted.deleted is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_adelete_file_propagates_s3_errors(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
import httpx
|
||||
import respx
|
||||
|
||||
import litellm
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
|
||||
|
||||
monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket")
|
||||
monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True")
|
||||
async with AsyncExitStack() as stack:
|
||||
client: Final = AsyncHTTPHandler()
|
||||
stack.push_async_callback(client.close)
|
||||
with respx.mock:
|
||||
route: Final = respx.delete(self.URL).mock(
|
||||
return_value=httpx.Response(403, content=b"<Error><Code>AccessDenied</Code></Error>")
|
||||
)
|
||||
from litellm.llms.bedrock.common_utils import BedrockError
|
||||
|
||||
with pytest.raises(BedrockError, match="AccessDenied"):
|
||||
await litellm.afile_delete(
|
||||
file_id=self.S3_URI, custom_llm_provider="bedrock", client=client,
|
||||
aws_access_key_id="AKIAEXAMPLE", aws_secret_access_key="test-secret", aws_region_name="us-west-2",
|
||||
)
|
||||
assert route.call_count == 1
|
||||
|
||||
@pytest.mark.parametrize("file_id, message", [
|
||||
("s3://other-bucket/litellm-bedrock-files-model-abc.jsonl", "configured storage bucket"),
|
||||
("s3://my-bucket/private/data.jsonl", "LiteLLM-managed"),
|
||||
])
|
||||
def test_delete_rejects_untrusted_objects_before_signing(
|
||||
self, file_id: str, message: str, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
from litellm.llms.bedrock.files.transformation import BedrockFilesConfig
|
||||
|
||||
monkeypatch.setenv("AWS_S3_BUCKET_NAME", "my-bucket")
|
||||
with pytest.raises(ValueError, match=message):
|
||||
BedrockFilesConfig().transform_delete_file_request(file_id=file_id, optional_params={}, litellm_params={})
|
||||
|
||||
|
||||
class TestBedrockFileContentTransformation:
|
||||
"""SigV4-signed S3 GetObject retrieval of Bedrock batch output files."""
|
||||
|
||||
|
|
@ -1873,7 +1946,7 @@ class TestBedrockFileContentTransformation:
|
|||
import hashlib
|
||||
|
||||
from litellm.llms.bedrock.files.transformation import (
|
||||
S3_SIGNED_GET_HEADERS_PARAM,
|
||||
S3_SIGNED_REQUEST_HEADERS_PARAM,
|
||||
BedrockFilesConfig,
|
||||
)
|
||||
|
||||
|
|
@ -1889,7 +1962,7 @@ class TestBedrockFileContentTransformation:
|
|||
assert url == self.EXPECTED_URL
|
||||
assert params == {}
|
||||
|
||||
signed_headers = litellm_params[S3_SIGNED_GET_HEADERS_PARAM]
|
||||
signed_headers = litellm_params[S3_SIGNED_REQUEST_HEADERS_PARAM]
|
||||
content_hashes = {
|
||||
value
|
||||
for name, value in signed_headers.items()
|
||||
|
|
@ -2139,7 +2212,7 @@ class TestBedrockFileContentTransformation:
|
|||
def test_s3_region_name_wins_for_content_signing(self, monkeypatch):
|
||||
"""s3_region_name must override aws_region_name for both the URL and the signature."""
|
||||
from litellm.llms.bedrock.files.transformation import (
|
||||
S3_SIGNED_GET_HEADERS_PARAM,
|
||||
S3_SIGNED_REQUEST_HEADERS_PARAM,
|
||||
BedrockFilesConfig,
|
||||
)
|
||||
|
||||
|
|
@ -2154,17 +2227,17 @@ class TestBedrockFileContentTransformation:
|
|||
)
|
||||
|
||||
assert url.startswith("https://s3.eu-west-1.amazonaws.com/")
|
||||
authorization = litellm_params[S3_SIGNED_GET_HEADERS_PARAM]["Authorization"]
|
||||
authorization = litellm_params[S3_SIGNED_REQUEST_HEADERS_PARAM]["Authorization"]
|
||||
assert "/eu-west-1/s3/aws4_request" in authorization
|
||||
|
||||
def test_validate_environment_merges_and_pops_signed_get_headers(self):
|
||||
from litellm.llms.bedrock.files.transformation import (
|
||||
S3_SIGNED_GET_HEADERS_PARAM,
|
||||
S3_SIGNED_REQUEST_HEADERS_PARAM,
|
||||
BedrockFilesConfig,
|
||||
)
|
||||
|
||||
litellm_params = {
|
||||
S3_SIGNED_GET_HEADERS_PARAM: {"Authorization": "AWS4-HMAC-SHA256 test"}
|
||||
S3_SIGNED_REQUEST_HEADERS_PARAM: {"Authorization": "AWS4-HMAC-SHA256 test"}
|
||||
}
|
||||
|
||||
headers = BedrockFilesConfig().validate_environment(
|
||||
|
|
@ -2179,7 +2252,7 @@ class TestBedrockFileContentTransformation:
|
|||
"x-custom": "kept",
|
||||
"Authorization": "AWS4-HMAC-SHA256 test",
|
||||
}
|
||||
assert S3_SIGNED_GET_HEADERS_PARAM not in litellm_params
|
||||
assert S3_SIGNED_REQUEST_HEADERS_PARAM not in litellm_params
|
||||
|
||||
def test_transform_file_content_response_wraps_binary_content(self):
|
||||
import httpx
|
||||
|
|
@ -2379,7 +2452,7 @@ class TestBedrockFilesS3SignatureEncoding:
|
|||
self, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
from litellm.llms.bedrock.files.transformation import (
|
||||
S3_SIGNED_GET_HEADERS_PARAM,
|
||||
S3_SIGNED_REQUEST_HEADERS_PARAM,
|
||||
BedrockFilesConfig,
|
||||
)
|
||||
|
||||
|
|
@ -2402,7 +2475,7 @@ class TestBedrockFilesS3SignatureEncoding:
|
|||
method="GET",
|
||||
url=url,
|
||||
body=None,
|
||||
headers=litellm_params[S3_SIGNED_GET_HEADERS_PARAM],
|
||||
headers=litellm_params[S3_SIGNED_REQUEST_HEADERS_PARAM],
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -2457,7 +2530,7 @@ def test_sign_s3_request_assumes_role_with_external_id(monkeypatch):
|
|||
assert "ASIAFILESPUTROLE" in authorization
|
||||
|
||||
|
||||
def test_sign_s3_get_request_assumes_role_with_external_id(monkeypatch):
|
||||
def test_sign_s3_request_without_body_assumes_role_with_external_id(monkeypatch):
|
||||
"""A trust policy requiring sts:ExternalId must be satisfied when signing the S3 download request."""
|
||||
import datetime
|
||||
from unittest.mock import patch
|
||||
|
|
@ -2504,7 +2577,7 @@ def test_sign_s3_get_request_assumes_role_with_external_id(monkeypatch):
|
|||
assert request_params.aws_external_id == "external-id-files-get"
|
||||
|
||||
with patch.object(boto3, "client", return_value=FakeSTSClient()):
|
||||
signed_headers = BedrockFilesConfig()._sign_s3_get_request(
|
||||
signed_headers = BedrockFilesConfig()._sign_s3_request_without_body(
|
||||
api_base="https://s3.us-east-1.amazonaws.com/safe-bucket/litellm-bedrock-files-model-id-abc.jsonl",
|
||||
aws_region_name="us-east-1",
|
||||
request_params=request_params,
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue