allow load testing sagemaker url

This commit is contained in:
Ishaan Jaff 2024-08-22 20:03:27 -07:00
parent 3f96a8c667
commit ed2d3f4bf8

View file

@ -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: