tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-07-17 18:17:52 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor(bedrock/sagemaker): switch to lazy loading for response stre… (#28189)

Browse Source 
* 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.
This commit is contained in:
harish-berri
2026-05-18 23:21:04 -07:00
committed by GitHub
parent 0290c7bc00
commit cff3e0b75e
6 changed files with 146 additions and 72 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/llms/bedrock/chat/invoke_agent/transformation.py
+2 -2
View File
@@ -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."""
litellm/llms/bedrock/chat/invoke_handler.py
+4 -5
View File
@@ -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()
litellm/llms/bedrock/common_utils.py
+15 -10
View File
@@ -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()
litellm/llms/sagemaker/common_utils.py
+14 -6
View File
@@ -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}")
tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py
+55 -24
View File
@@ -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)
tests/test_litellm/llms/sagemaker/test_sagemaker_common_utils.py
+56 -25
View File
@@ -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)