Sagemaker endpoint capability to inject boto3 client for cross account scenarios (#10728)

- **Description: Allow to inject boto3 client for Cross account access
type of scenarios in using Sagemaker Endpoint **
  - **Issue:#10634 #10184** 
  - **Dependencies: None** 
  - **Tag maintainer:** 
  - **Twitter handle:lethargicoder**

Co-authored-by: Vikram(VS) <vssht@amazon.com>
pull/10785/head
Vikram Shitole 12 months ago committed by GitHub
parent c8f386db97
commit a4e858b111
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

@ -80,6 +80,22 @@ class SagemakerEndpoint(LLM):
""" """
""" """
Args:
region_name: The aws region e.g., `us-west-2`.
Fallsback to AWS_DEFAULT_REGION env variable
or region specified in ~/.aws/config.
credentials_profile_name: The name of the profile in the ~/.aws/credentials
or ~/.aws/config files, which has either access keys or role information
specified. If not specified, the default credential profile or, if on an
EC2 instance, credentials from IMDS will be used.
client: boto3 client for Sagemaker Endpoint
content_handler: Implementation for model specific LLMContentHandler
Example: Example:
.. code-block:: python .. code-block:: python
@ -98,8 +114,21 @@ class SagemakerEndpoint(LLM):
region_name=region_name, region_name=region_name,
credentials_profile_name=credentials_profile_name credentials_profile_name=credentials_profile_name
) )
#Use with boto3 client
client = boto3.client(
"sagemaker-runtime",
region_name=region_name
)
se = SagemakerEndpoint(
endpoint_name=endpoint_name,
client=client
)
""" """
client: Any #: :meta private: client: Any = None
"""Boto3 client for sagemaker runtime"""
endpoint_name: str = "" endpoint_name: str = ""
"""The name of the endpoint from the deployed Sagemaker model. """The name of the endpoint from the deployed Sagemaker model.
@ -157,6 +186,10 @@ class SagemakerEndpoint(LLM):
@root_validator() @root_validator()
def validate_environment(cls, values: Dict) -> Dict: def validate_environment(cls, values: Dict) -> Dict:
"""Dont do anything if client provided externally"""
if values.get("client") is not None:
return values
"""Validate that AWS credentials to and python package exists in environment.""" """Validate that AWS credentials to and python package exists in environment."""
try: try:
import boto3 import boto3

Loading…
Cancel
Save