diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 13f134fce1d..e720428847d 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 run_aws_signing, sign_request_off_loop_if_aws +from litellm.llms.bedrock.base_aws_llm import SignsRequestsWithAWS, 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,12 @@ class BaseLLMHTTPHandler: headers=request_headers, ), ) - return await dispatch_async(*await run_aws_signing(sign_and_log, transformed)) + signed_request: Final = await ( + run_aws_signing(sign_and_log, transformed) + if isinstance(provider_config, SignsRequestsWithAWS) + else asyncio.to_thread(sign_and_log, transformed) + ) + return await dispatch_async(*signed_request) return transform_then_dispatch() 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 e24ea1ca611..2c7e476e9a8 100644 --- a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py +++ b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py @@ -3251,25 +3251,30 @@ async def test_sign_request_off_loop_if_aws_keeps_the_loop_serving_while_credent assert probe.served_during_refresh is True -async def test_run_aws_signing_leaves_the_default_executor_free_for_other_providers(): +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() + other provider's async entry point hops through that same executor. The scenario runs on its own loop + so the one-thread default executor it pins never leaks into the session loop.""" - def sign() -> str: - loop.call_soon_threadsafe(signing_parked.set) - refresh_done.wait() - return threading.current_thread().name + async def scenario() -> tuple[str, str]: + loop = asyncio.get_running_loop() + loop.set_default_executor(ThreadPoolExecutor(max_workers=1)) + signing_parked = asyncio.Event() + refresh_done = threading.Event() - 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() + def sign() -> str: + loop.call_soon_threadsafe(signing_parked.set) + refresh_done.wait() + return threading.current_thread().name - assert other_provider.name != await signing - assert (await signing).startswith("aws-signing") + 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() + return other_provider.name, await signing + + other_provider, signing_thread = asyncio.run(scenario()) + assert other_provider != signing_thread + assert signing_thread.startswith("aws-signing") diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py index 1140090f9cf..cea2d439198 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py @@ -20,6 +20,7 @@ from litellm.llms.base_llm.audio_transcription.transformation import ( BaseAudioTranscriptionConfig, ) from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException +from litellm.llms.bedrock.base_aws_llm import SignsRequestsWithAWS from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.custom_httpx.llm_http_handler import ( @@ -3431,6 +3432,26 @@ async def test_completion_signs_and_logs_off_the_event_loop_after_the_async_tran assert captured["body"] == {"transformed_by": "async"} assert config.sign_threads and all(thread is not loop_thread for thread in config.sign_threads) assert pre_call_threads and all(thread is not loop_thread for thread in pre_call_threads) + assert not any(thread.name.startswith("aws-signing") for thread in config.sign_threads + pre_call_threads) + + +class _AWSTransformRecordingConfig(SignsRequestsWithAWS, _TransformRecordingConfig): + pass + + +async def test_completion_signs_aws_configs_on_the_aws_signing_pool_after_the_async_transform(): + config = _AWSTransformRecordingConfig(transform_async=True) + pre_call_threads = [] + logging_obj = Mock(dynamic_success_callbacks=None, model_call_details={}) + logging_obj.pre_call.side_effect = lambda **kwargs: pre_call_threads.append(threading.current_thread()) + + pending, captured = _start_async_completion(config, logging_obj) + response = await pending + + assert response.choices[0].message.content == "async" + assert captured["body"] == {"transformed_by": "async"} + assert config.sign_threads and all(thread.name.startswith("aws-signing") for thread in config.sign_threads) + assert pre_call_threads and all(thread.name.startswith("aws-signing") for thread in pre_call_threads) async def test_completion_keeps_sync_transform_request_before_returning_by_default():