From c6be8326dbd0da26fcb2164f1a140929cd14fd92 Mon Sep 17 00:00:00 2001 From: Peter Muller Date: Mon, 1 Jul 2024 16:00:42 -0700 Subject: [PATCH] Allow calling SageMaker endpoints from different regions --- litellm/llms/sagemaker.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/litellm/llms/sagemaker.py b/litellm/llms/sagemaker.py index 8e75428bb7..079951b935 100644 --- a/litellm/llms/sagemaker.py +++ b/litellm/llms/sagemaker.py @@ -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",