scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-06 02:48:13 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

fix(sagemaker): Add role assumption support for embedding endpoint (#20435)

Browse source 
The SageMaker embedding handler was not using _load_credentials(),
which meant aws_role_name and aws_session_name parameters were
ignored. This prevented cross-account role assumption for embeddings
while it worked for completions.

Changes:
- Replace direct boto3 client creation with _load_credentials() call
- Create boto3.Session with assumed credentials
- Add comprehensive unit tests for role assumption

This aligns the embedding handler behavior with the completion handler,
which already supports role assumption via the BaseAWSLLM.get_credentials()
method.

Fixes cross-account SageMaker embedding access where users need to
assume a role in another account to invoke endpoints.
This commit is contained in:
jymmi 2026-03-09 20:58:14 -07:00 • committed by GitHub
parent 9314963697
commit 33d4b93bee
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 263 additions and 36 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

56
litellm/llms/sagemaker/completion/handler.py
View file

@ -583,35 +583,17 @@ class SagemakerLLM(BaseAWSLLM):
### BOTO3 INIT
import boto3
# pop aws_secret_access_key, aws_access_key_id, aws_region_name from kwargs, since completion calls fail with them
aws_secret_access_key = optional_params.pop("aws_secret_access_key", None)
aws_access_key_id = optional_params.pop("aws_access_key_id", None)
aws_region_name = optional_params.pop("aws_region_name", None)
# Use _load_credentials to support role assumption (aws_role_name, aws_session_name)
credentials, aws_region_name = self._load_credentials(optional_params)
if aws_access_key_id is not None:
# uses auth params passed to completion
# aws_access_key_id is not None, assume user is trying to auth using litellm.completion
client = boto3.client(
service_name="sagemaker-runtime",
aws_access_key_id=aws_access_key_id,
aws_secret_access_key=aws_secret_access_key,
region_name=aws_region_name,
)
else:
# aws_access_key_id is None, assume user is trying to auth using env variables
# boto3 automaticaly reads env variables
# we need to read region name from env
# I assume majority of users use .env for auth
region_name = (
get_secret("AWS_REGION_NAME")
or aws_region_name # get region from config file if specified
or "us-west-2" # default to us-west-2 if region not specified
)
client = boto3.client(
service_name="sagemaker-runtime",
region_name=region_name,
)
# Create boto3 session with the loaded credentials
session = boto3.Session(
aws_access_key_id=credentials.access_key,
aws_secret_access_key=credentials.secret_key,
aws_session_token=credentials.token,
region_name=aws_region_name,
)
client = session.client(service_name="sagemaker-runtime")
# pop streaming if it's in the optional params as 'stream' raises an error with sagemaker
inference_params = deepcopy(optional_params)
@ -628,7 +610,9 @@ class SagemakerLLM(BaseAWSLLM):
#### EMBEDDING LOGIC
# Transform request based on model type
provider_config = SagemakerEmbeddingConfig.get_model_config(model)
request_data = provider_config.transform_embedding_request(model, input, optional_params, {})
request_data = provider_config.transform_embedding_request(
model, input, optional_params, {}
)
data = json.dumps(request_data).encode("utf-8")
## LOGGING
@ -673,19 +657,19 @@ class SagemakerLLM(BaseAWSLLM):
)
print_verbose(f"raw model_response: {response}")
# Transform response based on model type
from httpx import Response as HttpxResponse
# Create a mock httpx Response object for the transformation
mock_response = HttpxResponse(
status_code=200,
content=json.dumps(response).encode('utf-8'),
headers={"content-type": "application/json"}
content=json.dumps(response).encode("utf-8"),
headers={"content-type": "application/json"},
)
model_response = EmbeddingResponse()
# Use the request_data that was already transformed above
return provider_config.transform_embedding_response(
model=model,
@ -695,5 +679,5 @@ class SagemakerLLM(BaseAWSLLM):
api_key=None,
request_data=request_data,
optional_params=optional_params,
litellm_params=litellm_params or {}
litellm_params=litellm_params or {},
)

243
tests/test_litellm/llms/sagemaker/test_sagemaker_embedding_role_assumption.py Normal file
View file

@ -0,0 +1,243 @@
"""
Test cases for SageMaker embedding role assumption support
This module tests that the SageMaker embedding handler properly supports
AWS IAM role assumption via aws_role_name and aws_session_name parameters,
matching the behavior of the completion handler.
"""
import json
import os
import sys
from datetime import timezone
from unittest.mock import MagicMock, call, patch
sys.path.insert(0, os.path.abspath("../../../../.."))
from botocore.credentials import Credentials
from litellm.llms.sagemaker.completion.handler import SagemakerLLM
from litellm.types.utils import EmbeddingResponse
class TestSagemakerEmbeddingRoleAssumption:
"""Test that SageMaker embedding supports role assumption like completion does"""
def setup_method(self):
self.sagemaker_llm = SagemakerLLM()
def test_embedding_uses_load_credentials(self):
"""
Test that embedding() calls _load_credentials() to support role assumption.
This ensures aws_role_name and aws_session_name parameters are properly handled.
"""
# Mock credentials that would be returned after role assumption
mock_credentials = Credentials(
access_key="assumed-access-key",
secret_key="assumed-secret-key",
token="assumed-session-token",
)
# Mock the SageMaker client response
mock_sagemaker_client = MagicMock()
mock_sagemaker_client.invoke_endpoint.return_value = {
"Body": MagicMock(
read=MagicMock(return_value=json.dumps({"embedding": [[0.1, 0.2, 0.3]]}).encode())
)
}
# Mock boto3.Session to return our mock client
mock_session = MagicMock()
mock_session.client.return_value = mock_sagemaker_client
with patch.object(
self.sagemaker_llm, "_load_credentials", return_value=(mock_credentials, "us-east-1")
) as mock_load_creds, patch("boto3.Session", return_value=mock_session):
# Create mock logging object
mock_logging = MagicMock()
optional_params = {
"aws_role_name": "arn:aws:iam::123456789012:role/TestRole",
"aws_session_name": "test-session",
}
self.sagemaker_llm.embedding(
model="test-endpoint",
input=["hello world"],
model_response=EmbeddingResponse(),
print_verbose=print,
encoding=None,
logging_obj=mock_logging,
optional_params=optional_params,
)
# Verify _load_credentials was called with the optional_params
mock_load_creds.assert_called_once()
# Verify boto3.Session was created with the assumed credentials
mock_session_calls = mock_session.client.call_args_list
assert len(mock_session_calls) == 1
assert mock_session_calls[0] == call(service_name="sagemaker-runtime")
def test_embedding_role_assumption_with_sts(self):
"""
Test the full role assumption flow for embeddings, similar to completion.
Verifies that STS assume_role is called when aws_role_name is provided.
"""
# Mock the STS client for role assumption
mock_sts_client = MagicMock()
# Mock the STS response with proper expiration handling
mock_expiry = MagicMock()
mock_expiry.tzinfo = timezone.utc
time_diff = MagicMock()
time_diff.total_seconds.return_value = 3600
mock_expiry.__sub__ = MagicMock(return_value=time_diff)
mock_sts_response = {
"Credentials": {
"AccessKeyId": "assumed-access-key",
"SecretAccessKey": "assumed-secret-key",
"SessionToken": "assumed-session-token",
"Expiration": mock_expiry,
}
}
mock_sts_client.assume_role.return_value = mock_sts_response
# Mock the SageMaker client response
mock_sagemaker_client = MagicMock()
mock_sagemaker_client.invoke_endpoint.return_value = {
"Body": MagicMock(
read=MagicMock(return_value=json.dumps({"embedding": [[0.1, 0.2, 0.3]]}).encode())
)
}
# Mock boto3.Session for SageMaker client creation
mock_session = MagicMock()
mock_session.client.return_value = mock_sagemaker_client
def mock_boto3_client(service_name, **kwargs):
if service_name == "sts":
return mock_sts_client
return mock_sagemaker_client
with patch("boto3.client", side_effect=mock_boto3_client), \
patch("boto3.Session", return_value=mock_session):
mock_logging = MagicMock()
optional_params = {
"aws_role_name": "arn:aws:iam::123456789012:role/CrossAccountRole",
"aws_session_name": "litellm-embedding-session",
"aws_region_name": "us-east-1",
}
self.sagemaker_llm.embedding(
model="test-endpoint",
input=["hello world"],
model_response=EmbeddingResponse(),
print_verbose=print,
encoding=None,
logging_obj=mock_logging,
optional_params=optional_params,
)
# Verify STS assume_role was called with correct parameters
mock_sts_client.assume_role.assert_called_once()
call_args = mock_sts_client.assume_role.call_args
assert call_args[1]["RoleArn"] == "arn:aws:iam::123456789012:role/CrossAccountRole"
assert call_args[1]["RoleSessionName"] == "litellm-embedding-session"
def test_embedding_without_role_assumption(self):
"""
Test that embedding works without role assumption when aws_role_name is not provided.
Should use default credentials from environment/instance profile.
"""
# Mock the SageMaker client response
mock_sagemaker_client = MagicMock()
mock_sagemaker_client.invoke_endpoint.return_value = {
"Body": MagicMock(
read=MagicMock(return_value=json.dumps({"embedding": [[0.1, 0.2, 0.3]]}).encode())
)
}
mock_session = MagicMock()
mock_session.client.return_value = mock_sagemaker_client
# Mock credentials returned from environment
mock_credentials = Credentials(
access_key="env-access-key",
secret_key="env-secret-key",
token=None,
)
with patch.object(
self.sagemaker_llm, "_load_credentials", return_value=(mock_credentials, "us-west-2")
), patch("boto3.Session", return_value=mock_session):
mock_logging = MagicMock()
# No aws_role_name provided
optional_params = {
"aws_region_name": "us-west-2",
}
result = self.sagemaker_llm.embedding(
model="test-endpoint",
input=["hello world"],
model_response=EmbeddingResponse(),
print_verbose=print,
encoding=None,
logging_obj=mock_logging,
optional_params=optional_params,
)
# Should still work and return embeddings
assert result is not None
def test_embedding_session_created_with_assumed_credentials(self):
"""
Test that boto3.Session is created with the credentials from role assumption.
This verifies the credentials flow from _load_credentials to the SageMaker client.
"""
mock_credentials = Credentials(
access_key="assumed-key",
secret_key="assumed-secret",
token="assumed-token",
)
mock_sagemaker_client = MagicMock()
mock_sagemaker_client.invoke_endpoint.return_value = {
"Body": MagicMock(
read=MagicMock(return_value=json.dumps({"embedding": [[0.1, 0.2, 0.3]]}).encode())
)
}
with patch.object(
self.sagemaker_llm, "_load_credentials", return_value=(mock_credentials, "us-east-1")
), patch("boto3.Session") as mock_session_class:
mock_session = MagicMock()
mock_session.client.return_value = mock_sagemaker_client
mock_session_class.return_value = mock_session
mock_logging = MagicMock()
self.sagemaker_llm.embedding(
model="test-endpoint",
input=["hello world"],
model_response=EmbeddingResponse(),
print_verbose=print,
encoding=None,
logging_obj=mock_logging,
optional_params={},
)
# Verify Session was created with the assumed credentials
mock_session_class.assert_called_once_with(
aws_access_key_id="assumed-key",
aws_secret_access_key="assumed-secret",
aws_session_token="assumed-token",
region_name="us-east-1",
)