From ed2d3f4bf86815584cf4e4561454848cef944101 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 22 Aug 2024 20:03:27 -0700 Subject: [PATCH] allow load testing sagemaker url --- litellm/llms/sagemaker.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/litellm/llms/sagemaker.py b/litellm/llms/sagemaker.py index 32146b9caed..76c3460f132 100644 --- a/litellm/llms/sagemaker.py +++ b/litellm/llms/sagemaker.py @@ -182,6 +182,10 @@ class SagemakerLLM(BaseAWSLLM): else: api_base = f"https://runtime.sagemaker.{aws_region_name}.amazonaws.com/endpoints/{model}/invocations" + sagemaker_base_url = optional_params.get("sagemaker_base_url", None) + if sagemaker_base_url is not None: + api_base = sagemaker_base_url + encoded_data = json.dumps(data).encode("utf-8") headers = {"Content-Type": "application/json"} if extra_headers is not None: