scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-11 22:51:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 123

fix(bedrock): sign on a dedicated executor instead of the shared default one

Browse source 
asyncio.to_thread puts every Bedrock signing on the loop's default executor, the same pool every provider's async entry point hops through, so signings parked on botocore's refresh lock queued unrelated providers behind Bedrock. run_aws_signing runs them on a 16-thread pool only AWS signing uses
This commit is contained in:
mateo-berri 2026-09-09 18:52:00 -07:00
parent e27a549aa9
commit d93baa3f2c
13 changed files with 76 additions and 34 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
litellm/a2a_protocol/providers/bedrock_agentcore/handler.py
View file

@ -5,7 +5,6 @@ Sends JSON-RPC envelopes directly to AgentCore endpoints, bypassing the
completion bridge that would otherwise strip the envelope.
"""
import asyncio
import json
from collections.abc import AsyncIterator, Mapping
from typing import Any, Final
@ -14,6 +13,7 @@ from litellm._logging import verbose_logger
from litellm.a2a_protocol.providers.bedrock_agentcore.transformation import (
BedrockAgentCoreA2ATransformation,
)
from litellm.llms.bedrock.base_aws_llm import run_aws_signing
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
from litellm.types.llms.custom_http import httpxSpecialProvider
@ -46,7 +46,7 @@ class BedrockAgentCoreA2AHandler:
Returns:
A2A JSON-RPC response dict from the AgentCore agent
"""
url, headers, body = await asyncio.to_thread(
url, headers, body = await run_aws_signing(
BedrockAgentCoreA2ATransformation.get_url_and_signed_request,
request_id=request_id,
params=params,
@ -93,7 +93,7 @@ class BedrockAgentCoreA2AHandler:
Yields:
A2A streaming response events from the AgentCore agent
"""
url, headers, body = await asyncio.to_thread(
url, headers, body = await run_aws_signing(
BedrockAgentCoreA2ATransformation.get_url_and_signed_request,
request_id=request_id,
params=params,

1
litellm/constants.py
View file

@ -570,6 +570,7 @@ LOGGING_WORKER_AGGRESSIVE_CLEAR_COOLDOWN_SECONDS: Final = float(
LOGGING_EXECUTOR_MAX_THREADS: Final = get_env_int("LOGGING_EXECUTOR_MAX_THREADS", 100)
LOGGING_EXECUTOR_MAX_PENDING_TASKS: Final = get_env_int("LOGGING_EXECUTOR_MAX_PENDING_TASKS", 10_000)
LOGGING_EXECUTOR_DROPPED_TASK_LOG_INTERVAL_SECONDS: Final = 30.0
AWS_SIGNING_MAX_THREADS: Final = 16
DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE: Final = os.getenv(
"DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE", "streaming.chunk.yield"
)

6
litellm/integrations/s3_v2.py
View file

@ -24,7 +24,7 @@ from litellm.integrations.s3 import (
from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, run_aws_signing
from litellm.llms.custom_httpx.http_handler import (
_get_httpx_client,
get_async_httpx_client,
@ -366,7 +366,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
# Sign the request
aws_request: Final = AWSRequest(method="PUT", url=url, data=json_string, headers=headers)
aws_region_name: Final = self.get_aws_region_name_for_non_llm_api_calls(aws_region_name=self.s3_region_name)
await asyncio.to_thread(S3SigV4Auth(credentials, "s3", aws_region_name).add_auth, aws_request)
await run_aws_signing(S3SigV4Auth(credentials, "s3", aws_region_name).add_auth, aws_request)
# Prepare the signed headers
signed_headers: Final = dict(aws_request.headers.items())
@ -597,7 +597,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
# Sign the request
aws_request: Final = AWSRequest(method="GET", url=url, headers=headers)
await asyncio.to_thread(S3SigV4Auth(credentials, "s3", self.s3_region_name).add_auth, aws_request)
await run_aws_signing(S3SigV4Auth(credentials, "s3", self.s3_region_name).add_auth, aws_request)
# Prepare the signed headers
signed_headers: Final = dict(aws_request.headers.items())

4
litellm/integrations/sqs.py
View file

@ -22,7 +22,7 @@ from litellm.constants import (
SQS_SEND_MESSAGE_ACTION,
)
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, run_aws_signing
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
httpxSpecialProvider,
@ -295,7 +295,7 @@ class SQSLogger(CustomBatchLogger, BaseAWSLLM):
data=prepped.body,
headers=prepped.headers,
)
await asyncio.to_thread(SigV4Auth(credentials, "sqs", self.sqs_region_name).add_auth, aws_request)
await run_aws_signing(SigV4Auth(credentials, "sqs", self.sqs_region_name).add_auth, aws_request)
signed_headers: Final = dict(aws_request.headers.items())

20
litellm/llms/bedrock/base_aws_llm.py
View file

@ -1,12 +1,15 @@
import asyncio
import base64
import contextvars
import hashlib
import json
import os
import re
import urllib.parse
from collections.abc import Callable, Mapping
from concurrent.futures import ThreadPoolExecutor
from datetime import datetime
from functools import partial
from threading import Lock
from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, ParamSpec, TypeVar, cast, get_args, overload
@ -17,6 +20,7 @@ from litellm._logging import verbose_logger
from litellm.caching.caching import DualCache
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.constants import (
AWS_SIGNING_MAX_THREADS,
BEDROCK_EMBEDDING_PROVIDERS_LITERAL,
BEDROCK_IAM_CACHE_FETCH_LOCK_STRIPES,
BEDROCK_IAM_CACHE_MAX_ENTRIES,
@ -1697,6 +1701,20 @@ def sign_aws_json_post(
_SignParams = ParamSpec("_SignParams")
_SignedRequest = TypeVar("_SignedRequest")
AWS_SIGNING_EXECUTOR: Final = ThreadPoolExecutor(max_workers=AWS_SIGNING_MAX_THREADS, thread_name_prefix="aws-signing")
async def run_aws_signing(
sign: Callable[_SignParams, _SignedRequest],
/,
*args: _SignParams.args,
**kwargs: _SignParams.kwargs, # kwargs-ok: ParamSpec forwarding keeps the wrapped signing signature
) -> _SignedRequest:
context: Final = contextvars.copy_context()
return await asyncio.get_running_loop().run_in_executor(
AWS_SIGNING_EXECUTOR, partial(context.run, sign, *args, **kwargs)
)
async def sign_request_off_loop_if_aws(
provider_config: object,
@ -1706,5 +1724,5 @@ async def sign_request_off_loop_if_aws(
**kwargs: _SignParams.kwargs, # kwargs-ok: ParamSpec forwarding keeps the wrapped sign_request signature
) -> _SignedRequest:
if isinstance(provider_config, SignsRequestsWithAWS):
return await asyncio.to_thread(sign_request, *args, **kwargs)
return await run_aws_signing(sign_request, *args, **kwargs)
return sign_request(*args, **kwargs)

7
litellm/llms/bedrock/chat/converse_handler.py
View file

@ -1,4 +1,3 @@
import asyncio
import json
from collections.abc import Mapping
from types import MappingProxyType
@ -22,7 +21,7 @@ from litellm.rust_bridge.chat_completions import rust_chat_completions_accepts
from litellm.types.utils import ModelResponse
from litellm.utils import CustomStreamWrapper
from ..base_aws_llm import BaseAWSLLM, Credentials, bedrock_bearer_token
from ..base_aws_llm import BaseAWSLLM, Credentials, bedrock_bearer_token, run_aws_signing
from ..common_utils import BedrockError, _get_all_bedrock_regions, error_response_text
from .invoke_handler import AWSEventStreamDecoder, MockResponseIterator, make_call
@ -137,7 +136,7 @@ class BedrockConverseLLM(BaseAWSLLM):
)
data: Final = json.dumps(request_data)
prepped: Final = await asyncio.to_thread(
prepped: Final = await run_aws_signing(
self.get_request_headers,
credentials=credentials,
aws_region_name=litellm_params.get("aws_region_name") or "us-west-2",
@ -208,7 +207,7 @@ class BedrockConverseLLM(BaseAWSLLM):
)
data: Final = json.dumps(request_data)
prepped: Final = await asyncio.to_thread(
prepped: Final = await run_aws_signing(
self.get_request_headers,
credentials=credentials,
aws_region_name=litellm_params.get("aws_region_name") or "us-west-2",

4
litellm/llms/bedrock/count_tokens/handler.py
View file

@ -4,13 +4,13 @@ AWS Bedrock CountTokens API handler.
Simplified handler leveraging existing LiteLLM Bedrock infrastructure.
"""
import asyncio
from typing import Any, Final
import httpx
import litellm
from litellm._logging import verbose_logger
from litellm.llms.bedrock.base_aws_llm import run_aws_signing
from litellm.llms.bedrock.common_utils import BedrockError
from litellm.llms.bedrock.count_tokens.transformation import BedrockCountTokensConfig
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, get_async_httpx_client
@ -77,7 +77,7 @@ class BedrockCountTokensHandler(BedrockCountTokensConfig):
# Extract api_key for bearer token auth if provided
api_key: Final = litellm_params.get("api_key", None)
headers: Final = {"Content-Type": "application/json"}
signed_headers, signed_body = await asyncio.to_thread(
signed_headers, signed_body = await run_aws_signing(
self._sign_request,
service_name="bedrock",
headers=headers,

7
litellm/llms/bedrock/embed/embedding.py
View file

@ -2,7 +2,6 @@
Handles embedding calls to Bedrock's `/invoke` endpoint
"""
import asyncio
import copy
import json
import urllib.parse
@ -27,7 +26,7 @@ from litellm.types.llms.bedrock import (
)
from litellm.types.utils import EmbeddingResponse, LlmProviders
from ..base_aws_llm import AWSPreparedRequest, BaseAWSLLM, Credentials, bedrock_bearer_token
from ..base_aws_llm import AWSPreparedRequest, BaseAWSLLM, Credentials, bedrock_bearer_token, run_aws_signing
from ..common_utils import BedrockError
from .amazon_nova_transformation import AmazonNovaEmbeddingConfig
from .amazon_titan_g1_transformation import AmazonTitanG1Config
@ -357,7 +356,7 @@ class BedrockEmbedding(BaseAWSLLM):
if extra_headers is not None:
headers = {"Content-Type": "application/json", **extra_headers}
prepped = await asyncio.to_thread(
prepped = await run_aws_signing(
self.get_request_headers,
credentials=credentials,
aws_region_name=aws_region_name,
@ -638,7 +637,7 @@ class BedrockEmbedding(BaseAWSLLM):
credentials=credentials, url=status_url, headers=headers, aws_region_name=aws_region_name
)
prepped: Final = await asyncio.to_thread(sign_status_request)
prepped: Final = await run_aws_signing(sign_status_request)
# LOGGING
if logging_obj is not None:

6
litellm/llms/bedrock/realtime/handler.py
View file

@ -21,7 +21,7 @@ from litellm.litellm_core_utils.realtime_streaming import DefaultLoggedRealTimeE
from litellm.types.llms.openai import OpenAIRealtimeEvents
from litellm.types.realtime import RealtimeResponseTransformInput
from ..base_aws_llm import BaseAWSLLM
from ..base_aws_llm import BaseAWSLLM, run_aws_signing
from ..common_utils import BedrockError
from .transformation import BedrockRealtimeConfig
@ -149,7 +149,7 @@ class BedrockRealtime(BaseAWSLLM):
verbose_proxy_logger.debug("Bedrock Realtime: Connecting to %s with model %s", endpoint_uri, model)
credentials: Final = await asyncio.to_thread(
credentials: Final = await run_aws_signing(
self.get_credentials,
aws_access_key_id=aws_access_key_id,
aws_secret_access_key=aws_secret_access_key,
@ -170,7 +170,7 @@ class BedrockRealtime(BaseAWSLLM):
"or configure credentials in the environment"
),
)
frozen_credentials: Final = await asyncio.to_thread(credentials.get_frozen_credentials)
frozen_credentials: Final = await run_aws_signing(credentials.get_frozen_credentials)
# Initialize Bedrock client with aws_sdk_bedrock_runtime
config: Final = Config(

4
litellm/llms/custom_httpx/llm_http_handler.py
View file

@ -77,7 +77,7 @@ from litellm.llms.base_llm.vector_store_files.transformation import (
BaseVectorStoreFilesConfig,
)
from litellm.llms.base_llm.videos.transformation import BaseVideoConfig
from litellm.llms.bedrock.base_aws_llm import sign_request_off_loop_if_aws
from litellm.llms.bedrock.base_aws_llm import run_aws_signing, sign_request_off_loop_if_aws
from litellm.llms.custom_httpx.container_handler import raise_for_error_status
from litellm.llms.custom_httpx.http_handler import (
AsyncHTTPHandler,
@ -638,7 +638,7 @@ class BaseLLMHTTPHandler:
headers=request_headers,
),
)
return await dispatch_async(*await asyncio.to_thread(sign_and_log, transformed))
return await dispatch_async(*await run_aws_signing(sign_and_log, transformed))
return transform_then_dispatch()

10
litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py
View file

@ -44,7 +44,7 @@ from litellm.llms.anthropic.chat.guardrail_translation.handler import AnthropicM
from litellm.llms.base_llm.guardrail_translation.utils import (
effective_scan_only_tool_results_for_guardrail,
)
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, bedrock_bearer_token
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, bedrock_bearer_token, run_aws_signing
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
httpxSpecialProvider,
@ -917,7 +917,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
source,
)
return BedrockGuardrailResponse()
credentials, aws_region_name = await asyncio.to_thread(
credentials, aws_region_name = await run_aws_signing(
self._load_credentials, bearer_token=bedrock_bearer_token(api_key)
)
allow_chunking: Final = not self._content_uses_contextual_grounding(content)
@ -1180,7 +1180,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
**base_request_data,
"content": content,
} # mutable-ok: outbound JSON request body
prepared_request: Final = await asyncio.to_thread(
prepared_request: Final = await run_aws_signing(
self._prepare_request,
credentials=credentials,
data=bedrock_request_data,
@ -1878,12 +1878,12 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
return BedrockGuardrailResponse()
api_key: Final[str | None] = request_data.get("api_key") if request_data else None
credentials, aws_region_name = await asyncio.to_thread(
credentials, aws_region_name = await run_aws_signing(
self._load_credentials, bearer_token=bedrock_bearer_token(api_key)
)
body: Final[dict[str, object]] = {"messages": checks_messages, "checks": self.checks}
prepared_request: Final = await asyncio.to_thread(
prepared_request: Final = await run_aws_signing(
self._prepare_request,
credentials=credentials,
data=body,

9
litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py
View file

@ -8,7 +8,6 @@ Use litellm with Anthropic SDK, Vertex AI SDK, Cohere SDK, etc.
from __future__ import annotations
import asyncio
import hmac
import inspect
import json
@ -1131,7 +1130,7 @@ async def bedrock_proxy_route(
)
# Add or update query parameters
from litellm.llms.bedrock.base_aws_llm import sign_aws_json_post
from litellm.llms.bedrock.base_aws_llm import run_aws_signing, sign_aws_json_post
from litellm.llms.bedrock.chat import BedrockConverseLLM
bedrock_llm: Final = BedrockConverseLLM()
@ -1140,7 +1139,7 @@ async def bedrock_proxy_route(
data: Final = await _json_request_body(request)
except Exception as e:
raise HTTPException(status_code=400, detail={"error": e})
prepped: Final = await asyncio.to_thread(
prepped: Final = await run_aws_signing(
sign_aws_json_post,
get_credentials=bedrock_llm.get_credentials,
service_name="bedrock",
@ -1236,10 +1235,10 @@ async def comprehend_medical_proxy_route(
if "stream" in data:
raise HTTPException(status_code=400, detail="'stream' is not a Comprehend Medical request member")
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, sign_aws_json_post
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, run_aws_signing, sign_aws_json_post
target_url: Final = f"https://comprehendmedical.{aws_region_name}.{get_aws_dns_suffix(aws_region_name)}/"
prepped: Final = await asyncio.to_thread(
prepped: Final = await run_aws_signing(
sign_aws_json_post,
get_credentials=partial(BaseAWSLLM().get_credentials, aws_region_name=aws_region_name),
service_name="comprehendmedical",

26
tests/test_litellm/llms/bedrock/test_base_aws_llm.py
View file

@ -1,5 +1,6 @@
import asyncio
import json
from concurrent.futures import ThreadPoolExecutor
import os
import threading
import time
@ -23,6 +24,7 @@ from litellm.llms.bedrock.base_aws_llm import (
AwsAuthError,
BaseAWSLLM,
Boto3CredentialsInfo,
run_aws_signing,
sign_request_off_loop_if_aws,
)
from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe
@ -3247,3 +3249,27 @@ async def test_sign_request_off_loop_if_aws_keeps_the_loop_serving_while_credent
assert "Authorization" in signed
assert probe.served_during_refresh is True
async def test_run_aws_signing_leaves_the_default_executor_free_for_other_providers():
"""A signing parked on botocore's refresh lock must not hold a default-executor thread, since every
other provider's async entry point hops through that same executor."""
loop = asyncio.get_running_loop()
loop.set_default_executor(ThreadPoolExecutor(max_workers=1))
signing_parked = asyncio.Event()
refresh_done = threading.Event()
def sign() -> str:
loop.call_soon_threadsafe(signing_parked.set)
refresh_done.wait()
return threading.current_thread().name
signing = asyncio.create_task(run_aws_signing(sign))
try:
await asyncio.wait_for(signing_parked.wait(), timeout=5)
other_provider = await asyncio.wait_for(loop.run_in_executor(None, threading.current_thread), timeout=5)
finally:
refresh_done.set()
assert other_provider.name != await signing
assert (await signing).startswith("aws-signing")