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 36

Merge pull request #34087 from ArjunPakhan/fix/bedrock-cancel-batch

Browse source 
fix(batches): support AWS Bedrock batch cancellation via `StopModelInvocationJob`
This commit is contained in:
Mateo Wang 2026-08-17 10:00:57 -07:00 • committed by GitHub
commit 41c3133d0e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 186 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
litellm/batches/main.py
View file

@ -826,7 +826,7 @@ def list_batches(
async def acancel_batch(
batch_id: str,
model: str | None = None,
custom_llm_provider: Literal["openai", "azure", "vertex_ai"] = "openai",
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock"] = "openai",
metadata: dict[str, str] | None = None,
extra_headers: dict[str, str] | None = None,
extra_body: dict[str, str] | None = None,
@ -872,7 +872,7 @@ async def acancel_batch(
def cancel_batch(
batch_id: str,
model: str | None = None,
custom_llm_provider: Literal["openai", "azure", "vertex_ai"] | str = "openai",
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock"] | str = "openai",
metadata: dict[str, str] | None = None,
extra_headers: dict[str, str] | None = None,
extra_body: dict[str, str] | None = None,
@ -993,9 +993,14 @@ def cancel_batch(
timeout=timeout,
max_retries=optional_params.max_retries,
)
elif custom_llm_provider == "bedrock":
response = BedrockBatchesHandler.cancel_batch(
batch_id=batch_id,
**kwargs,
)
else:
raise litellm.exceptions.BadRequestError(
message=f"LiteLLM doesn't support {custom_llm_provider} for 'cancel_batch'. Only 'openai', 'azure', and 'vertex_ai' are supported.",
message=f"LiteLLM doesn't support {custom_llm_provider} for 'cancel_batch'. Only 'openai', 'azure', 'vertex_ai', and 'bedrock' are supported.",
model="n/a",
llm_provider=custom_llm_provider,
response=httpx.Response(

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

@ -1,11 +1,14 @@
from datetime import datetime
from typing import Any, Final, cast
from typing import TYPE_CHECKING, Any, Final, cast
from openai.types.batch import BatchRequestCounts
from openai.types.batch import Metadata as OpenAIBatchMetadata
from litellm.types.utils import LiteLLMBatch
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
# 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.
@ -22,6 +25,8 @@ _BEDROCK_MIJ_STATUS_TO_OPENAI: Final = {
"Expired": "expired",
}
_CANCEL_IDEMPOTENT_STATUSES: Final = frozenset({"cancelling", "cancelled", "completed", "failed", "expired"})
def _extract_region_from_bedrock_arn(arn: str) -> str | None:
"""ARN shape: ``arn:aws:bedrock:<region>:<account>:<type>/<id>``"""
@ -82,6 +87,81 @@ class BedrockBatchesHandler:
E.g. Twelve Labs Embedding Async Invoke
"""
@staticmethod
def cancel_batch(
batch_id: str,
aws_region_name: str | None = None,
logging_obj: "LiteLLMLoggingObj | None" = None,
aws_access_key_id: str | None = None,
aws_secret_access_key: str | None = None,
aws_session_token: str | None = None,
aws_session_name: str | None = None,
aws_profile_name: str | None = None,
aws_role_name: str | None = None,
aws_web_identity_token: str | None = None,
aws_sts_endpoint: str | None = None,
aws_external_id: str | None = None,
**kwargs: object, # kwargs-ok: litellm.cancel_batch forwards arbitrary user kwargs verbatim
) -> "LiteLLMBatch":
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
region: Final = aws_region_name or _extract_region_from_bedrock_arn(batch_id) or "us-east-1"
from litellm.llms.bedrock.batches.transformation import BedrockBatchesConfig
creds: Final = BedrockBatchesConfig().get_credentials(
aws_access_key_id=aws_access_key_id,
aws_secret_access_key=aws_secret_access_key,
aws_session_token=aws_session_token,
aws_region_name=region,
aws_session_name=aws_session_name,
aws_profile_name=aws_profile_name,
aws_role_name=aws_role_name,
aws_web_identity_token=aws_web_identity_token,
aws_sts_endpoint=aws_sts_endpoint,
aws_external_id=aws_external_id,
)
client: Final = 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,
)
def job_status() -> "LiteLLMBatch":
return BedrockBatchesHandler._handle_model_invocation_job_status(
batch_id=batch_id,
aws_region_name=region,
logging_obj=logging_obj,
aws_access_key_id=aws_access_key_id,
aws_secret_access_key=aws_secret_access_key,
aws_session_token=aws_session_token,
aws_session_name=aws_session_name,
aws_profile_name=aws_profile_name,
aws_role_name=aws_role_name,
aws_web_identity_token=aws_web_identity_token,
aws_sts_endpoint=aws_sts_endpoint,
aws_external_id=aws_external_id,
)
try:
client.stop_model_invocation_job(jobIdentifier=batch_id)
except ClientError as e:
if e.response.get("Error", {}).get("Code") not in ("ValidationException", "ConflictException"):
raise
current_batch: Final = job_status()
if current_batch.status not in _CANCEL_IDEMPOTENT_STATUSES:
raise
return current_batch
return job_status()
@staticmethod
def _handle_async_invoke_status(batch_id: str, aws_region_name: str, logging_obj=None, **kwargs) -> "LiteLLMBatch":
"""

97
tests/test_litellm/llms/bedrock/batches/test_handler.py
View file

@ -336,3 +336,100 @@ def test_logging_url_uses_bare_id_when_only_id_passed(patched_boto3):
assert pre_kwargs["additional_args"]["api_base"] == (
f"https://bedrock.us-west-2.amazonaws.com/model-invocation-job/{JOB_ID}"
)
def test_cancel_batch_stops_job_and_returns_mapped_status(patched_boto3):
fake_client, boto_client_factory = patched_boto3
fake_client.get_model_invocation_job.return_value = _fake_boto3_response(status="Stopping")
batch = BedrockBatchesHandler.cancel_batch(batch_id=JOB_ARN)
fake_client.stop_model_invocation_job.assert_called_once_with(jobIdentifier=JOB_ARN)
_, kwargs = boto_client_factory.call_args
assert kwargs["region_name"] == "us-west-2"
assert batch.status == "cancelling"
def test_cancel_batch_tolerates_already_terminal_job(patched_boto3):
from botocore.exceptions import ClientError
fake_client, _ = patched_boto3
fake_client.stop_model_invocation_job.side_effect = ClientError(
{"Error": {"Code": "ValidationException", "Message": "Job is already in a terminal state"}},
"StopModelInvocationJob",
)
fake_client.get_model_invocation_job.return_value = _fake_boto3_response(status="Stopped")
batch = BedrockBatchesHandler.cancel_batch(batch_id=JOB_ARN)
assert batch.status == "cancelled"
def test_cancel_batch_tolerates_conflict_on_already_stopped_job(patched_boto3):
from botocore.exceptions import ClientError
fake_client, _ = patched_boto3
fake_client.stop_model_invocation_job.side_effect = ClientError(
{"Error": {"Code": "ConflictException", "Message": "Job cannot be stopped in its current state"}},
"StopModelInvocationJob",
)
fake_client.get_model_invocation_job.return_value = _fake_boto3_response(status="Stopped")
batch = BedrockBatchesHandler.cancel_batch(batch_id=JOB_ARN)
assert batch.status == "cancelled"
def test_cancel_batch_reraises_conflict_when_job_not_terminal(patched_boto3):
from botocore.exceptions import ClientError
fake_client, _ = patched_boto3
fake_client.stop_model_invocation_job.side_effect = ClientError(
{"Error": {"Code": "ConflictException", "Message": "Operation conflicts with current job state"}},
"StopModelInvocationJob",
)
fake_client.get_model_invocation_job.return_value = _fake_boto3_response(status="InProgress")
with pytest.raises(ClientError):
BedrockBatchesHandler.cancel_batch(batch_id=JOB_ARN)
def test_cancel_batch_reraises_validation_error_when_job_not_terminal(patched_boto3):
from botocore.exceptions import ClientError
fake_client, _ = patched_boto3
fake_client.stop_model_invocation_job.side_effect = ClientError(
{"Error": {"Code": "ValidationException", "Message": "Cannot stop job in current state"}},
"StopModelInvocationJob",
)
fake_client.get_model_invocation_job.return_value = _fake_boto3_response(status="InProgress")
with pytest.raises(ClientError):
BedrockBatchesHandler.cancel_batch(batch_id=JOB_ARN)
def test_cancel_batch_reraises_other_client_errors(patched_boto3):
from botocore.exceptions import ClientError
fake_client, _ = patched_boto3
fake_client.stop_model_invocation_job.side_effect = ClientError(
{"Error": {"Code": "AccessDeniedException", "Message": "not authorized"}},
"StopModelInvocationJob",
)
with pytest.raises(ClientError):
BedrockBatchesHandler.cancel_batch(batch_id=JOB_ARN)
fake_client.get_model_invocation_job.assert_not_called()
def test_litellm_cancel_batch_dispatches_to_bedrock(patched_boto3):
import litellm
fake_client, _ = patched_boto3
fake_client.get_model_invocation_job.return_value = _fake_boto3_response(status="Stopped")
batch = litellm.cancel_batch(batch_id=JOB_ARN, custom_llm_provider="bedrock")
fake_client.stop_model_invocation_job.assert_called_once_with(jobIdentifier=JOB_ARN)
assert batch.status == "cancelled"