From cff3e0b75eeb836e9fafbcbf73fd4c968b001def Mon Sep 17 00:00:00 2001 From: harish-berri Date: Mon, 18 May 2026 23:21:04 -0700 Subject: [PATCH] =?UTF-8?q?refactor(bedrock/sagemaker):=20switch=20to=20la?= =?UTF-8?q?zy=20loading=20for=20response=20stre=E2=80=A6=20(#28189)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * refactor(bedrock/sagemaker): switch to lazy loading for response stream shapes - Replace eager loading of BEDROCK_RESPONSE_STREAM_SHAPE and SAGEMAKER_RESPONSE_STREAM_SHAPE with lazy loading via get_bedrock_response_stream_shape() and get_sagemaker_response_stream_shape() respectively. - This change optimizes performance by avoiding unnecessary imports and logging warnings unless the response stream shapes are actually needed. - Update relevant classes and tests to utilize the new lazy loading functions, ensuring consistent behavior across the codebase. * test(bedrock/sagemaker): add fixtures to clear response stream shape cache - Introduced `_reset_bedrock_response_stream_shape_cache` and `_reset_sagemaker_response_stream_shape_cache` fixtures to prevent lru_cache leakage between tests in their respective modules. - Updated tests to utilize these fixtures, ensuring that the response stream shape cache is cleared before and after each test run. - Added `pytest.importorskip("botocore")` to ensure that tests are skipped if the botocore library is not available. --- .../chat/invoke_agent/transformation.py | 4 +- litellm/llms/bedrock/chat/invoke_handler.py | 9 +-- litellm/llms/bedrock/common_utils.py | 25 +++--- litellm/llms/sagemaker/common_utils.py | 20 +++-- .../llms/bedrock/test_bedrock_common_utils.py | 79 ++++++++++++------ .../sagemaker/test_sagemaker_common_utils.py | 81 +++++++++++++------ 6 files changed, 146 insertions(+), 72 deletions(-) diff --git a/litellm/llms/bedrock/chat/invoke_agent/transformation.py b/litellm/llms/bedrock/chat/invoke_agent/transformation.py index e4072c2455..c88fa32b6a 100644 --- a/litellm/llms/bedrock/chat/invoke_agent/transformation.py +++ b/litellm/llms/bedrock/chat/invoke_agent/transformation.py @@ -299,9 +299,9 @@ class AmazonInvokeAgentConfig(BaseConfig, BaseAWSLLM): ) def _get_response_stream_shape(self): - from litellm.llms.bedrock.common_utils import BEDROCK_RESPONSE_STREAM_SHAPE + from litellm.llms.bedrock.common_utils import get_bedrock_response_stream_shape - return BEDROCK_RESPONSE_STREAM_SHAPE + return get_bedrock_response_stream_shape() def _extract_response_content(self, events: InvokeAgentEventList) -> str: """Extract the final response content from parsed events.""" diff --git a/litellm/llms/bedrock/chat/invoke_handler.py b/litellm/llms/bedrock/chat/invoke_handler.py index 92ca75db95..7a9916f1f3 100644 --- a/litellm/llms/bedrock/chat/invoke_handler.py +++ b/litellm/llms/bedrock/chat/invoke_handler.py @@ -68,9 +68,9 @@ from litellm.utils import CustomStreamWrapper, get_secret from ..base_aws_llm import BaseAWSLLM from ..common_utils import ( - BEDROCK_RESPONSE_STREAM_SHAPE, BedrockError, ModelResponseIterator, + get_bedrock_response_stream_shape, get_bedrock_tool_name, ) @@ -1828,7 +1828,8 @@ class AWSEventStreamDecoder: yield self._chunk_parser(chunk_data=_data) def _parse_message_from_event(self, event) -> Optional[str]: - if BEDROCK_RESPONSE_STREAM_SHAPE is None: + response_stream_shape = get_bedrock_response_stream_shape() + if response_stream_shape is None: raise BedrockError( status_code=500, message=( @@ -1837,9 +1838,7 @@ class AWSEventStreamDecoder: ), ) response_dict = event.to_response_dict() - parsed_response = self.parser.parse( - response_dict, BEDROCK_RESPONSE_STREAM_SHAPE - ) + parsed_response = self.parser.parse(response_dict, response_stream_shape) if response_dict["status_code"] != 200: decoded_body = response_dict["body"].decode() diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 0256d5d4b9..4f4729e401 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -4,6 +4,7 @@ from __future__ import annotations Common utilities used across bedrock chat/embedding/image generation """ +import functools import json import os from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union @@ -963,10 +964,8 @@ def _load_bedrock_response_stream_shape(): """ Load the ResponseStream shape from botocore's bundled bedrock-runtime schema. - Called once at module import time; the result is stored in - ``BEDROCK_RESPONSE_STREAM_SHAPE`` and reused for the process lifetime. Returns ``None`` if botocore is unavailable or the service model cannot be - loaded, so the module still imports cleanly. + loaded. """ try: from botocore.loaders import Loader @@ -977,15 +976,22 @@ def _load_bedrock_response_stream_shape(): return ServiceModel(service_dict).shape_for("ResponseStream") except Exception as e: verbose_logger.warning( - "litellm: could not pre-load bedrock-runtime response stream shape " + "litellm: could not load bedrock-runtime response stream shape " "— Bedrock event-stream decoding will be unavailable. Error: %s", e, ) return None -# Eagerly resolved once per process — avoids per-instance or per-request disk I/O. -BEDROCK_RESPONSE_STREAM_SHAPE = _load_bedrock_response_stream_shape() +@functools.lru_cache(maxsize=1) +def get_bedrock_response_stream_shape(): + """ + Lazily load and cache the bedrock-runtime ResponseStream shape for the process. + + Avoids importing botocore (and logging warnings) unless Bedrock event-stream + decoding is actually needed. + """ + return _load_bedrock_response_stream_shape() class BedrockEventStreamDecoderBase: @@ -999,7 +1005,8 @@ class BedrockEventStreamDecoderBase: self.parser = EventStreamJSONParser() def _parse_message_from_event(self, event) -> Optional[str]: - if BEDROCK_RESPONSE_STREAM_SHAPE is None: + response_stream_shape = get_bedrock_response_stream_shape() + if response_stream_shape is None: raise BedrockError( status_code=500, message=( @@ -1008,9 +1015,7 @@ class BedrockEventStreamDecoderBase: ), ) response_dict = event.to_response_dict() - parsed_response = self.parser.parse( - response_dict, BEDROCK_RESPONSE_STREAM_SHAPE - ) + parsed_response = self.parser.parse(response_dict, response_stream_shape) if response_dict["status_code"] != 200: decoded_body = response_dict["body"].decode() diff --git a/litellm/llms/sagemaker/common_utils.py b/litellm/llms/sagemaker/common_utils.py index 50c8ee4220..6c15d642f8 100644 --- a/litellm/llms/sagemaker/common_utils.py +++ b/litellm/llms/sagemaker/common_utils.py @@ -1,3 +1,4 @@ +import functools import json from typing import AsyncIterator, Iterator, List, Optional, Union @@ -22,14 +23,22 @@ def _load_sagemaker_response_stream_shape(): ) except Exception as e: verbose_logger.warning( - "litellm: could not pre-load sagemaker-runtime response stream shape " + "litellm: could not load sagemaker-runtime response stream shape " "— SageMaker event-stream decoding will be unavailable. Error: %s", e, ) return None -SAGEMAKER_RESPONSE_STREAM_SHAPE = _load_sagemaker_response_stream_shape() +@functools.lru_cache(maxsize=1) +def get_sagemaker_response_stream_shape(): + """ + Lazily load and cache the sagemaker-runtime stream shape for the process. + + Avoids importing botocore (and logging warnings) unless SageMaker event-stream + decoding is actually needed. + """ + return _load_sagemaker_response_stream_shape() class SagemakerError(BaseLLMException): @@ -207,7 +216,8 @@ class AWSEventStreamDecoder: verbose_logger.error(f"Final error parsing accumulated JSON: {e}") def _parse_message_from_event(self, event) -> Optional[str]: - if SAGEMAKER_RESPONSE_STREAM_SHAPE is None: + response_stream_shape = get_sagemaker_response_stream_shape() + if response_stream_shape is None: raise SagemakerError( status_code=500, message=( @@ -216,9 +226,7 @@ class AWSEventStreamDecoder: ), ) response_dict = event.to_response_dict() - parsed_response = self.parser.parse( - response_dict, SAGEMAKER_RESPONSE_STREAM_SHAPE - ) + parsed_response = self.parser.parse(response_dict, response_stream_shape) if response_dict["status_code"] != 200: raise ValueError(f"Bad response code, expected 200: {response_dict}") diff --git a/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py b/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py index 8fa9290d3d..c39fb427a0 100644 --- a/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py +++ b/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py @@ -14,18 +14,46 @@ from litellm.llms.bedrock.common_utils import BedrockModelInfo # --------------------------------------------------------------------------- # -# BEDROCK_RESPONSE_STREAM_SHAPE eager-load tests # +# get_bedrock_response_stream_shape lazy-load tests # # --------------------------------------------------------------------------- # -def test_bedrock_response_stream_shape_loaded_at_import(): +@pytest.fixture(autouse=True) +def _reset_bedrock_response_stream_shape_cache(): + """Prevent lru_cache leakage between tests in this module.""" + import litellm.llms.bedrock.common_utils as mod + + mod.get_bedrock_response_stream_shape.cache_clear() + yield + mod.get_bedrock_response_stream_shape.cache_clear() + + +def test_bedrock_response_stream_shape_lazy_loads_once(): """ - BEDROCK_RESPONSE_STREAM_SHAPE is resolved at module import time. + get_bedrock_response_stream_shape() loads from botocore at most once per process. + """ + from unittest.mock import MagicMock, patch + + import litellm.llms.bedrock.common_utils as mod + + sentinel = MagicMock() + with patch.object( + mod, "_load_bedrock_response_stream_shape", return_value=sentinel + ) as mock_load: + assert mod.get_bedrock_response_stream_shape() is sentinel + assert mod.get_bedrock_response_stream_shape() is sentinel + mock_load.assert_called_once() + + +def test_bedrock_response_stream_shape_loaded_on_first_access(): + """ + get_bedrock_response_stream_shape() loads once on first use. In a standard environment with botocore installed it must be non-None. """ - from litellm.llms.bedrock.common_utils import BEDROCK_RESPONSE_STREAM_SHAPE + pytest.importorskip("botocore") + from litellm.llms.bedrock.common_utils import get_bedrock_response_stream_shape - assert BEDROCK_RESPONSE_STREAM_SHAPE is not None + assert get_bedrock_response_stream_shape() is not None def test_bedrock_response_stream_shape_load_failure_returns_none(): @@ -38,6 +66,7 @@ def test_bedrock_response_stream_shape_load_failure_returns_none(): import litellm.llms.bedrock.common_utils as mod + pytest.importorskip("botocore") with patch( "botocore.loaders.Loader.load_service_model", side_effect=Exception("no data"), @@ -51,31 +80,29 @@ def test_bedrock_response_stream_shape_is_structure_shape(): The loaded shape should be the botocore StructureShape for ResponseStream, not a plain dict or any other type. """ + pytest.importorskip("botocore") from botocore.model import StructureShape - from litellm.llms.bedrock.common_utils import BEDROCK_RESPONSE_STREAM_SHAPE + from litellm.llms.bedrock.common_utils import get_bedrock_response_stream_shape - assert BEDROCK_RESPONSE_STREAM_SHAPE is not None, ( - "BEDROCK_RESPONSE_STREAM_SHAPE is None — botocore may not be installed" - ) - shape: StructureShape = BEDROCK_RESPONSE_STREAM_SHAPE # remove Optional + loaded_shape = get_bedrock_response_stream_shape() + assert ( + loaded_shape is not None + ), "get_bedrock_response_stream_shape() is None — botocore may not be installed" + shape: StructureShape = loaded_shape assert isinstance(shape, StructureShape) assert shape.name == "ResponseStream" -def test_bedrock_response_stream_shape_same_object_across_imports(): +def test_bedrock_response_stream_shape_same_object_across_calls(): """ - Both bedrock modules that use the shape must reference the identical object — - confirming the constant is not re-loaded per import. + Repeated calls must return the identical cached object. """ - from litellm.llms.bedrock.chat.invoke_handler import ( - BEDROCK_RESPONSE_STREAM_SHAPE as invoke_shape, - ) - from litellm.llms.bedrock.common_utils import ( - BEDROCK_RESPONSE_STREAM_SHAPE as common_shape, - ) + from litellm.llms.bedrock.common_utils import get_bedrock_response_stream_shape - assert common_shape is invoke_shape + first = get_bedrock_response_stream_shape() + second = get_bedrock_response_stream_shape() + assert first is second def test_bedrock_event_stream_decoder_base_uses_module_shape(): @@ -95,19 +122,23 @@ def test_bedrock_event_stream_decoder_base_uses_module_shape(): def test_bedrock_parse_message_from_event_raises_on_none_shape(): """ - When BEDROCK_RESPONSE_STREAM_SHAPE is None (botocore unavailable), + When get_bedrock_response_stream_shape() returns None (botocore unavailable), _parse_message_from_event must raise BedrockError before touching the botocore parser — not an opaque AttributeError from inside botocore. """ from unittest.mock import MagicMock, patch import litellm.llms.bedrock.common_utils as mod - from litellm.llms.bedrock.common_utils import BedrockError, BedrockEventStreamDecoderBase + from litellm.llms.bedrock.common_utils import ( + BedrockError, + BedrockEventStreamDecoderBase, + ) - decoder = BedrockEventStreamDecoderBase() + decoder = BedrockEventStreamDecoderBase.__new__(BedrockEventStreamDecoderBase) + decoder.parser = MagicMock() mock_event = MagicMock() - with patch.object(mod, "BEDROCK_RESPONSE_STREAM_SHAPE", None): + with patch.object(mod, "get_bedrock_response_stream_shape", return_value=None): with pytest.raises(BedrockError) as exc_info: decoder._parse_message_from_event(mock_event) diff --git a/tests/test_litellm/llms/sagemaker/test_sagemaker_common_utils.py b/tests/test_litellm/llms/sagemaker/test_sagemaker_common_utils.py index 9d7706557b..7e13459bca 100644 --- a/tests/test_litellm/llms/sagemaker/test_sagemaker_common_utils.py +++ b/tests/test_litellm/llms/sagemaker/test_sagemaker_common_utils.py @@ -12,18 +12,46 @@ from litellm.llms.sagemaker.completion.transformation import SagemakerConfig # --------------------------------------------------------------------------- # -# SAGEMAKER_RESPONSE_STREAM_SHAPE eager-load tests # +# get_sagemaker_response_stream_shape lazy-load tests # # --------------------------------------------------------------------------- # -def test_sagemaker_response_stream_shape_loaded_at_import(): +@pytest.fixture(autouse=True) +def _reset_sagemaker_response_stream_shape_cache(): + """Prevent lru_cache leakage between tests in this module.""" + import litellm.llms.sagemaker.common_utils as mod + + mod.get_sagemaker_response_stream_shape.cache_clear() + yield + mod.get_sagemaker_response_stream_shape.cache_clear() + + +def test_sagemaker_response_stream_shape_lazy_loads_once(): """ - SAGEMAKER_RESPONSE_STREAM_SHAPE is resolved at module import time. + get_sagemaker_response_stream_shape() loads from botocore at most once per process. + """ + from unittest.mock import MagicMock, patch + + import litellm.llms.sagemaker.common_utils as mod + + sentinel = MagicMock() + with patch.object( + mod, "_load_sagemaker_response_stream_shape", return_value=sentinel + ) as mock_load: + assert mod.get_sagemaker_response_stream_shape() is sentinel + assert mod.get_sagemaker_response_stream_shape() is sentinel + mock_load.assert_called_once() + + +def test_sagemaker_response_stream_shape_loaded_on_first_access(): + """ + get_sagemaker_response_stream_shape() loads once on first use. In a standard environment with botocore installed it must be non-None. """ - from litellm.llms.sagemaker.common_utils import SAGEMAKER_RESPONSE_STREAM_SHAPE + pytest.importorskip("botocore") + from litellm.llms.sagemaker.common_utils import get_sagemaker_response_stream_shape - assert SAGEMAKER_RESPONSE_STREAM_SHAPE is not None + assert get_sagemaker_response_stream_shape() is not None def test_sagemaker_response_stream_shape_load_failure_returns_none(): @@ -36,6 +64,7 @@ def test_sagemaker_response_stream_shape_load_failure_returns_none(): import litellm.llms.sagemaker.common_utils as mod + pytest.importorskip("botocore") with patch( "botocore.loaders.Loader.load_service_model", side_effect=Exception("no data"), @@ -49,14 +78,16 @@ def test_sagemaker_response_stream_shape_is_structure_shape(): The loaded shape should be the botocore StructureShape for InvokeEndpointWithResponseStreamOutput, not a plain dict or any other type. """ + pytest.importorskip("botocore") from botocore.model import StructureShape - from litellm.llms.sagemaker.common_utils import SAGEMAKER_RESPONSE_STREAM_SHAPE + from litellm.llms.sagemaker.common_utils import get_sagemaker_response_stream_shape - assert SAGEMAKER_RESPONSE_STREAM_SHAPE is not None, ( - "SAGEMAKER_RESPONSE_STREAM_SHAPE is None — botocore may not be installed" - ) - shape: StructureShape = SAGEMAKER_RESPONSE_STREAM_SHAPE # remove Optional + shape = get_sagemaker_response_stream_shape() + assert ( + shape is not None + ), "get_sagemaker_response_stream_shape() is None — botocore may not be installed" + shape: StructureShape = shape # remove Optional assert isinstance(shape, StructureShape) assert shape.name == "InvokeEndpointWithResponseStreamOutput" @@ -64,29 +95,25 @@ def test_sagemaker_response_stream_shape_is_structure_shape(): def test_sagemaker_response_stream_shape_not_reloaded_on_new_decoder(): """ Creating multiple AWSEventStreamDecoder instances must not trigger - additional botocore Loader calls — the shape is resolved once at import - time and reused. + additional botocore Loader calls — the shape is cached after first access. """ - from litellm.llms.sagemaker.common_utils import SAGEMAKER_RESPONSE_STREAM_SHAPE + from litellm.llms.sagemaker.common_utils import get_sagemaker_response_stream_shape - decoder_a = AWSEventStreamDecoder(model="test-model-a") - decoder_b = AWSEventStreamDecoder(model="test-model-b") + decoder_a = AWSEventStreamDecoder.__new__(AWSEventStreamDecoder) + decoder_b = AWSEventStreamDecoder.__new__(AWSEventStreamDecoder) - # Both decoders should use the same pre-loaded shape object (identity check) assert "_response_stream_shape_cache" not in decoder_a.__dict__ assert "_response_stream_shape_cache" not in decoder_b.__dict__ - # The module constant is still the same object - from litellm.llms.sagemaker.common_utils import ( - SAGEMAKER_RESPONSE_STREAM_SHAPE as shape_after, - ) - assert SAGEMAKER_RESPONSE_STREAM_SHAPE is shape_after + first = get_sagemaker_response_stream_shape() + second = get_sagemaker_response_stream_shape() + assert first is second def test_sagemaker_parse_message_from_event_raises_on_none_shape(): """ - When SAGEMAKER_RESPONSE_STREAM_SHAPE is None (botocore unavailable), - _parse_message_from_event must raise ValueError before touching the + When get_sagemaker_response_stream_shape() returns None (botocore unavailable), + _parse_message_from_event must raise SagemakerError before touching the botocore parser — not an opaque AttributeError from inside botocore. """ from unittest.mock import MagicMock, patch @@ -94,10 +121,14 @@ def test_sagemaker_parse_message_from_event_raises_on_none_shape(): import litellm.llms.sagemaker.common_utils as mod from litellm.llms.sagemaker.common_utils import SagemakerError - decoder = AWSEventStreamDecoder(model="test-model") + decoder = AWSEventStreamDecoder.__new__(AWSEventStreamDecoder) + decoder.model = "test-model" + decoder.parser = MagicMock() + decoder.content_blocks = [] + decoder.is_messages_api = None mock_event = MagicMock() - with patch.object(mod, "SAGEMAKER_RESPONSE_STREAM_SHAPE", None): + with patch.object(mod, "get_sagemaker_response_stream_shape", return_value=None): with pytest.raises(SagemakerError) as exc_info: decoder._parse_message_from_event(mock_event)