scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 8

fix(batches): support AWS Bedrock batch cancellation via StopModelInvocationJob (#33986)

Browse source
This commit is contained in:
Arjun Pakhan 2026-07-21 06:16:02 +00:00
parent 3f5186f9af
commit fc36825dfd
3 changed files with 124 additions and 102 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

9
litellm/batches/main.py
View file

@ -1087,9 +1087,16 @@ def cancel_batch(
timeout=timeout,
max_retries=optional_params.max_retries,
)
elif custom_llm_provider == "bedrock":
from litellm.llms.bedrock.batches.handler import BedrockBatchesHandler
response = BedrockBatchesHandler.cancel_batch(
batch_id=batch_id,
**kwargs,
)
else:
raise litellm.exceptions.BadRequestError(
message="LiteLLM doesn't support {} for 'cancel_batch'. Only 'openai', 'azure', and 'vertex_ai' are supported.".format(
message="LiteLLM doesn't support {} for 'cancel_batch'. Only 'openai', 'azure', 'vertex_ai', and 'bedrock' are supported.".format(
custom_llm_provider
),
model="n/a",

162
litellm/llms/bedrock/batches/handler.py
View file

@ -6,9 +6,7 @@ from openai.types.batch import Metadata as OpenAIBatchMetadata
from litellm.types.utils import LiteLLMBatch
# AWS Bedrock model-invocation-job statuses → OpenAI Batch statuses.
# Mirrors the mapping used by `BedrockBatchesConfig.transform_create_batch_response`
# so create / retrieve return consistent statuses.
# AWS Bedrock model-invocation-job statuses -> OpenAI Batch statuses.
_BEDROCK_MIJ_STATUS_TO_OPENAI = {
"Submitted": "validating",
"Validating": "validating",
@ -44,17 +42,6 @@ def _extract_job_id_from_arn(arn: str) -> Optional[str]:
def _predict_output_file_uri(
output_prefix: str, input_uri: str, job_id: Optional[str]
) -> Optional[str]:
"""
Compute the deterministic per-job result file URI Bedrock writes to.
Bedrock lays results out as::
<output_prefix>/<job-id>/<basename(input_uri)>.out
We compute it client-side so OpenAI-style ``client.files.content(output_file_id)``
works without an extra S3 ``ListObjectsV2`` round-trip. Returns ``None`` if we
don't have enough info; callers should fall back to the bare prefix.
"""
if not output_prefix or not input_uri or not job_id:
return None
if not output_prefix.endswith("/"):
@ -76,40 +63,76 @@ def _to_epoch(value: Any) -> Optional[int]:
class BedrockBatchesHandler:
"""
Handler for Bedrock Batches.
"""Handler for Bedrock Batches."""
Specific providers/models needed some special handling.
@staticmethod
def cancel_batch(
batch_id: str,
aws_region_name: Optional[str] = None,
logging_obj=None,
**kwargs,
) -> "LiteLLMBatch":
"""
Cancel an AWS Bedrock batch model invocation job using StopModelInvocationJob.
"""
try:
import boto3
from botocore.exceptions import ClientError
except ImportError as exc:
raise ImportError(
"Missing boto3/botocore to call bedrock. Run 'pip install boto3'."
) from exc
E.g. Twelve Labs Embedding Async Invoke
"""
region = (
aws_region_name or _extract_region_from_bedrock_arn(batch_id) or "us-east-1"
)
from litellm.llms.bedrock.batches.transformation import BedrockBatchesConfig
creds = BedrockBatchesConfig().get_credentials(
aws_access_key_id=kwargs.get("aws_access_key_id"),
aws_secret_access_key=kwargs.get("aws_secret_access_key"),
aws_session_token=kwargs.get("aws_session_token"),
aws_region_name=region,
aws_session_name=kwargs.get("aws_session_name"),
aws_profile_name=kwargs.get("aws_profile_name"),
aws_role_name=kwargs.get("aws_role_name"),
aws_web_identity_token=kwargs.get("aws_web_identity_token"),
aws_sts_endpoint=kwargs.get("aws_sts_endpoint"),
aws_external_id=kwargs.get("aws_external_id"),
)
client = boto3.client(
"bedrock",
region_name=region,
aws_access_key_id=creds.access_key,
aws_secret_access_key=creds.secret_key,
aws_session_token=creds.token,
)
try:
client.stop_model_invocation_job(jobIdentifier=batch_id)
except ClientError as e:
# Idempotency: if job is already Stopping/Stopped/Completed, swallow ValidationException
if e.response.get("Error", {}).get("Code") != "ValidationException":
raise e
return BedrockBatchesHandler._handle_model_invocation_job_status(
batch_id=batch_id,
aws_region_name=region,
logging_obj=logging_obj,
**kwargs,
)
@staticmethod
def _handle_async_invoke_status(
batch_id: str, aws_region_name: str, logging_obj=None, **kwargs
) -> "LiteLLMBatch":
"""
Handle async invoke status check for AWS Bedrock.
This is for Twelve Labs Embedding Async Invoke.
Args:
batch_id: The async invoke ARN
aws_region_name: AWS region name
**kwargs: Additional parameters
Returns:
dict: Status information including status, output_file_id (S3 URL), etc.
"""
import asyncio
from litellm.llms.bedrock.embed.embedding import BedrockEmbedding
async def _async_get_status():
# Create embedding handler instance
embedding_handler = BedrockEmbedding()
# Get the status of the async invoke job
status_response = await embedding_handler._get_async_invoke_status(
invocation_arn=batch_id,
aws_region_name=aws_region_name,
@ -117,18 +140,13 @@ class BedrockBatchesHandler:
**kwargs,
)
# Transform response to a LiteLLMBatch object
from litellm.types.utils import LiteLLMBatch
openai_batch_metadata: OpenAIBatchMetadata = {
"output_file_id": status_response["outputDataConfig"][
"s3OutputDataConfig"
]["s3Uri"],
"output_file_id": status_response["outputDataConfig"]["s3OutputDataConfig"]["s3Uri"],
"failure_message": status_response.get("failureMessage") or "",
"model_arn": status_response["modelArn"],
}
result = LiteLLMBatch(
return LiteLLMBatch(
id=status_response["invocationArn"],
object="batch",
status=status_response["status"],
@ -151,10 +169,6 @@ class BedrockBatchesHandler:
input_file_id="",
)
return result
# Since this function is called from within an async context via run_in_executor,
# we need to create a new event loop in a thread to avoid conflicts
import concurrent.futures
def run_in_thread():
@ -176,37 +190,6 @@ class BedrockBatchesHandler:
logging_obj=None,
**kwargs,
) -> "LiteLLMBatch":
"""
Handle ``GetModelInvocationJob`` status check for AWS Bedrock bulk batch
inference jobs (the ARN type returned by ``CreateModelInvocationJob``).
``CreateModelInvocationJob`` lives on the Bedrock **control plane**
(``bedrock.<region>.amazonaws.com``), distinct from the data-plane
``bedrock-runtime`` endpoint that serves Twelve Labs async-invoke ARNs.
The two ARN families therefore can't share a handler — see
``litellm/batches/main.py`` for the dispatch.
Args:
batch_id: A ``arn:aws:bedrock:<region>:<acct>:model-invocation-job/<id>``
ARN (or just the trailing job id; both are accepted by
``GetModelInvocationJob``).
aws_region_name: Region for the boto3 ``bedrock`` client. If omitted,
we fall back to parsing the region out of ``batch_id`` itself.
logging_obj: Optional litellm logging object.
**kwargs: Optional AWS credential overrides
(``aws_access_key_id``, ``aws_secret_access_key``,
``aws_session_token``, ``aws_profile_name``,
``aws_role_name``, ``aws_session_name``,
``aws_web_identity_token``, ``aws_sts_endpoint``,
``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.
"""
try:
import boto3
except ImportError as exc:
@ -214,15 +197,10 @@ class BedrockBatchesHandler:
"Missing boto3 to call bedrock. Run 'pip install boto3'."
) from exc
# Resolve region: explicit > parsed-from-ARN > us-east-1 (boto3 default).
region = (
aws_region_name or _extract_region_from_bedrock_arn(batch_id) or "us-east-1"
)
# Resolve credentials through the same path the rest of the bedrock
# provider uses, so model_list / env / role-assumption configs are
# honored. We instantiate BedrockBatchesConfig (which extends
# BaseAWSLLM) lazily to avoid a circular import at module load.
from litellm.llms.bedrock.batches.transformation import BedrockBatchesConfig
creds = BedrockBatchesConfig().get_credentials(
@ -247,10 +225,6 @@ class BedrockBatchesHandler:
)
if logging_obj is not None:
# Use the bare job id in the logged URL so we don't double up the
# `model-invocation-job/` segment when `batch_id` is a full ARN.
# `GetModelInvocationJob` accepts either form, but only the bare id
# produces a sensible-looking URL in logs.
url_path_id = _extract_job_id_from_arn(batch_id) or batch_id
logging_obj.pre_call(
input=batch_id,
@ -291,26 +265,12 @@ class BedrockBatchesHandler:
.get("s3Uri", "")
)
# Bedrock returns the output *prefix* the user supplied at job creation.
# Actual results land at <prefix>/<job-id>/<basename(input)>.out — we
# surface that single-file URI as `output_file_id` so the OpenAI-style
# download flow works without an extra S3 listing call. We deliberately
# do NOT fall back to the bare prefix when prediction fails: a prefix
# is not a downloadable object, so handing it back as `output_file_id`
# would reproduce the very NoSuchKey bug this handler exists to fix.
# The bare prefix is preserved in metadata for callers that want the
# `manifest.json.out` or want to do their own listing.
job_arn = response.get("jobArn", batch_id)
job_id = _extract_job_id_from_arn(job_arn)
output_file_uri = _predict_output_file_uri(output_prefix, input_uri, job_id)
completed_at = _to_epoch(response.get("endTime"))
# Note: metadata uses "" (not None) for unknown URIs to satisfy the
# OpenAI Batch metadata schema, which is `dict[str, str]`. The
# `output_file_id` field on the LiteLLMBatch itself does carry None
# correctly (see below), so callers should branch on that, not on
# `metadata["output_file_uri"]`.
openai_batch_metadata: OpenAIBatchMetadata = {
"model_arn": response.get("modelId", ""),
"job_arn": job_arn,

55
tests/test_bedrock_cancel_batch.py Normal file
View file

@ -0,0 +1,55 @@
from unittest.mock import MagicMock, patch
import pytest
import litellm
from litellm.llms.bedrock.batches.handler import BedrockBatchesHandler
@patch("boto3.client")
def test_bedrock_cancel_batch_handler(mock_boto_client):
mock_client_instance = MagicMock()
mock_boto_client.return_value = mock_client_instance
mock_client_instance.stop_model_invocation_job.return_value = {}
mock_client_instance.get_model_invocation_job.return_value = {
"jobArn": "arn:aws:bedrock:us-east-1:123456789012:model-invocation-job/test-job-id",
"status": "Stopping",
"submitTime": 1700000000,
"lastModifiedTime": 1700000100,
}
res = BedrockBatchesHandler.cancel_batch(
batch_id="arn:aws:bedrock:us-east-1:123456789012:model-invocation-job/test-job-id",
aws_region_name="us-east-1",
aws_access_key_id="test",
aws_secret_access_key="test",
)
mock_client_instance.stop_model_invocation_job.assert_called_once_with(
jobIdentifier="arn:aws:bedrock:us-east-1:123456789012:model-invocation-job/test-job-id"
)
assert res.status == "cancelling"
@patch("boto3.client")
def test_litellm_cancel_batch_bedrock_dispatcher(mock_boto_client):
mock_client_instance = MagicMock()
mock_boto_client.return_value = mock_client_instance
mock_client_instance.stop_model_invocation_job.return_value = {}
mock_client_instance.get_model_invocation_job.return_value = {
"jobArn": "arn:aws:bedrock:us-east-1:123456789012:model-invocation-job/test-job-id",
"status": "Stopped",
"submitTime": 1700000000,
"lastModifiedTime": 1700000100,
}
res = litellm.cancel_batch(
batch_id="arn:aws:bedrock:us-east-1:123456789012:model-invocation-job/test-job-id",
custom_llm_provider="bedrock",
aws_region_name="us-east-1",
aws_access_key_id="test",
aws_secret_access_key="test",
)
assert res.status == "cancelled"