ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Allow calling SageMaker endpoints from different regions

Browse Source
This commit is contained in:
Peter Muller 2024-07-01 16:00:42 -07:00
parent bfab76a810
commit c6be8326db
No known key found for this signature in database
GPG Key ID: A6094745C95B1613
1 changed files with 8 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
litellm/llms/sagemaker.py
View File

@ -185,7 +185,8 @@ def completion(
# I assume majority of users use .env for auth
region_name = (
get_secret("AWS_REGION_NAME")
or "us-west-2" # default to us-west-2 if user not specified
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",
@ -439,7 +440,8 @@ async def async_streaming(
# I assume majority of users use .env for auth
region_name = (
get_secret("AWS_REGION_NAME")
or "us-west-2" # default to us-west-2 if user not specified
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 = session.client(
service_name="sagemaker-runtime",
@ -506,7 +508,8 @@ async def async_completion(
# I assume majority of users use .env for auth
region_name = (
get_secret("AWS_REGION_NAME")
or "us-west-2" # default to us-west-2 if user not specified
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 = session.client(
service_name="sagemaker-runtime",
@ -661,7 +664,8 @@ def embedding(
# I assume majority of users use .env for auth
region_name = (
get_secret("AWS_REGION_NAME")
or "us-west-2" # default to us-west-2 if user not specified
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",