From d93baa3f2c1ca3c5123ecdee8510ee3c8d827f61 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 18:52:00 -0700 Subject: [PATCH] fix(bedrock): sign on a dedicated executor instead of the shared default one 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 --- .../providers/bedrock_agentcore/handler.py | 6 ++--- litellm/constants.py | 1 + litellm/integrations/s3_v2.py | 6 ++--- litellm/integrations/sqs.py | 4 +-- litellm/llms/bedrock/base_aws_llm.py | 20 +++++++++++++- litellm/llms/bedrock/chat/converse_handler.py | 7 +++-- litellm/llms/bedrock/count_tokens/handler.py | 4 +-- litellm/llms/bedrock/embed/embedding.py | 7 +++-- litellm/llms/bedrock/realtime/handler.py | 6 ++--- litellm/llms/custom_httpx/llm_http_handler.py | 4 +-- .../guardrail_hooks/bedrock_guardrails.py | 10 +++---- .../llm_passthrough_endpoints.py | 9 +++---- .../llms/bedrock/test_base_aws_llm.py | 26 +++++++++++++++++++ 13 files changed, 76 insertions(+), 34 deletions(-) diff --git a/litellm/a2a_protocol/providers/bedrock_agentcore/handler.py b/litellm/a2a_protocol/providers/bedrock_agentcore/handler.py index e01e215b5f0..306a8871b12 100644 --- a/litellm/a2a_protocol/providers/bedrock_agentcore/handler.py +++ b/litellm/a2a_protocol/providers/bedrock_agentcore/handler.py @@ -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, diff --git a/litellm/constants.py b/litellm/constants.py index f9389d22dea..060402cc08d 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -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" ) diff --git a/litellm/integrations/s3_v2.py b/litellm/integrations/s3_v2.py index e3297cfd27c..996ff9c75af 100644 --- a/litellm/integrations/s3_v2.py +++ b/litellm/integrations/s3_v2.py @@ -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()) diff --git a/litellm/integrations/sqs.py b/litellm/integrations/sqs.py index 31cc20f5995..d787375ca3c 100644 --- a/litellm/integrations/sqs.py +++ b/litellm/integrations/sqs.py @@ -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()) diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index eddd4d6f826..96804a1fa62 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -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) diff --git a/litellm/llms/bedrock/chat/converse_handler.py b/litellm/llms/bedrock/chat/converse_handler.py index 3e8e1226a55..6d48ff3f07c 100644 --- a/litellm/llms/bedrock/chat/converse_handler.py +++ b/litellm/llms/bedrock/chat/converse_handler.py @@ -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", diff --git a/litellm/llms/bedrock/count_tokens/handler.py b/litellm/llms/bedrock/count_tokens/handler.py index cae2315e744..d7fb510f057 100644 --- a/litellm/llms/bedrock/count_tokens/handler.py +++ b/litellm/llms/bedrock/count_tokens/handler.py @@ -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, diff --git a/litellm/llms/bedrock/embed/embedding.py b/litellm/llms/bedrock/embed/embedding.py index 4254d9b1d8b..be766eaedd0 100644 --- a/litellm/llms/bedrock/embed/embedding.py +++ b/litellm/llms/bedrock/embed/embedding.py @@ -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: diff --git a/litellm/llms/bedrock/realtime/handler.py b/litellm/llms/bedrock/realtime/handler.py index 792e3169a7a..ca2370303f2 100644 --- a/litellm/llms/bedrock/realtime/handler.py +++ b/litellm/llms/bedrock/realtime/handler.py @@ -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( diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 3106810b4b3..13f134fce1d 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -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() diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 0ae2c48a3b4..6a7ac4361b9 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -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, diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 037f663044f..bf6ee57b4fc 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -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", diff --git a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py index ca93f58f366..e24ea1ca611 100644 --- a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py +++ b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py @@ -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")