fix(sagemaker.py): support model_id consistently. support dynamic args for async calls

litellm/llms/sagemaker.py

@@ -246,15 +246,28 @@ def completion(
                 model=model,
                 logging_obj=logging_obj,
                 data=data,
+                model_id=model_id,
+                aws_secret_access_key=aws_secret_access_key,
+                aws_access_key_id=aws_access_key_id,
+                aws_region_name=aws_region_name,
             )
             return response
-        response = client.invoke_endpoint_with_response_stream(
-            EndpointName=model,
-            ContentType="application/json",
-            Body=data,
-            CustomAttributes="accept_eula=true",
-        )
 
+        if model_id is not None:
+            response = client.invoke_endpoint_with_response_stream(
+                EndpointName=model,
+                InferenceComponentName=model_id,
+                ContentType="application/json",
+                Body=data,
+                CustomAttributes="accept_eula=true",
+            )
+        else:
+            response = client.invoke_endpoint_with_response_stream(
+                EndpointName=model,
+                ContentType="application/json",
+                Body=data,
+                CustomAttributes="accept_eula=true",
+            )
         return response["Body"]
     elif acompletion == True:
         _data = {"inputs": prompt, "parameters": inference_params}
@@ -265,31 +278,36 @@ def completion(
             model=model,
             logging_obj=logging_obj,
             data=_data,
+            model_id=model_id,
+            aws_secret_access_key=aws_secret_access_key,
+            aws_access_key_id=aws_access_key_id,
+            aws_region_name=aws_region_name,
         )
     data = json.dumps({"inputs": prompt, "parameters": inference_params}).encode(
         "utf-8"
     )
-    ## LOGGING
-    request_str = f"""
-    response = client.invoke_endpoint(
-        EndpointName={model},
-        ContentType="application/json",
-        Body={data},
-        CustomAttributes="accept_eula=true",
-    )
-    """  # type: ignore
-    logging_obj.pre_call(
-        input=prompt,
-        api_key="",
-        additional_args={
-            "complete_input_dict": data,
-            "request_str": request_str,
-            "hf_model_name": hf_model_name,
-        },
-    )
     ## COMPLETION CALL
     try:
         if model_id is not None:
+            ## LOGGING
+            request_str = f"""
+            response = client.invoke_endpoint(
+                EndpointName={model},
+                InferenceComponentName={model_id},
+                ContentType="application/json",
+                Body={data},
+                CustomAttributes="accept_eula=true",
+            )
+            """  # type: ignore
+            logging_obj.pre_call(
+                input=prompt,
+                api_key="",
+                additional_args={
+                    "complete_input_dict": data,
+                    "request_str": request_str,
+                    "hf_model_name": hf_model_name,
+                },
+            )
             response = client.invoke_endpoint(
                 EndpointName=model,
                 InferenceComponentName=model_id,
@@ -298,6 +316,24 @@ def completion(
                 CustomAttributes="accept_eula=true",
             )
         else:
+            ## LOGGING
+            request_str = f"""
+            response = client.invoke_endpoint(
+                EndpointName={model},
+                ContentType="application/json",
+                Body={data},
+                CustomAttributes="accept_eula=true",
+            )
+            """  # type: ignore
+            logging_obj.pre_call(
+                input=prompt,
+                api_key="",
+                additional_args={
+                    "complete_input_dict": data,
+                    "request_str": request_str,
+                    "hf_model_name": hf_model_name,
+                },
+            )
             response = client.invoke_endpoint(
                 EndpointName=model,
                 ContentType="application/json",
@@ -369,8 +405,12 @@ async def async_streaming(
     encoding,
     model_response: ModelResponse,
     model: str,
+    model_id: Optional[str],
     logging_obj: Any,
     data,
+    aws_secret_access_key: Optional[str],
+    aws_access_key_id: Optional[str],
+    aws_region_name: Optional[str],
 ):
     """
     Use aioboto3
@@ -379,11 +419,6 @@ async def async_streaming(
 
     session = aioboto3.Session()
 
-    # pop aws_secret_access_key, aws_access_key_id, aws_region_name from kwargs, since completion calls fail with them
-    aws_secret_access_key = optional_params.pop("aws_secret_access_key", None)
-    aws_access_key_id = optional_params.pop("aws_access_key_id", None)
-    aws_region_name = optional_params.pop("aws_region_name", None)
-
     if aws_access_key_id != None:
         # uses auth params passed to completion
         # aws_access_key_id is not None, assume user is trying to auth using litellm.completion
@@ -410,12 +445,21 @@ async def async_streaming(
 
     async with _client as client:
         try:
-            response = await client.invoke_endpoint_with_response_stream(
-                EndpointName=model,
-                ContentType="application/json",
-                Body=data,
-                CustomAttributes="accept_eula=true",
-            )
+            if model_id is not None:
+                response = await client.invoke_endpoint_with_response_stream(
+                    EndpointName=model,
+                    InferenceComponentName=model_id,
+                    ContentType="application/json",
+                    Body=data,
+                    CustomAttributes="accept_eula=true",
+                )
+            else:
+                response = await client.invoke_endpoint_with_response_stream(
+                    EndpointName=model,
+                    ContentType="application/json",
+                    Body=data,
+                    CustomAttributes="accept_eula=true",
+                )
         except Exception as e:
             raise SagemakerError(status_code=500, message=f"{str(e)}")
         response = response["Body"]
@@ -430,6 +474,10 @@ async def async_completion(
     model: str,
     logging_obj: Any,
     data: dict,
+    model_id: Optional[str],
+    aws_secret_access_key: Optional[str],
+    aws_access_key_id: Optional[str],
+    aws_region_name: Optional[str],
 ):
     """
     Use aioboto3
@@ -438,11 +486,6 @@ async def async_completion(
 
     session = aioboto3.Session()
 
-    # pop aws_secret_access_key, aws_access_key_id, aws_region_name from kwargs, since completion calls fail with them
-    aws_secret_access_key = optional_params.pop("aws_secret_access_key", None)
-    aws_access_key_id = optional_params.pop("aws_access_key_id", None)
-    aws_region_name = optional_params.pop("aws_region_name", None)
-
     if aws_access_key_id != None:
         # uses auth params passed to completion
         # aws_access_key_id is not None, assume user is trying to auth using litellm.completion
@@ -468,33 +511,63 @@ async def async_completion(
         )
 
     async with _client as client:
-        ## LOGGING
-        request_str = f"""
-        response = client.invoke_endpoint(
-            EndpointName={model},
-            ContentType="application/json",
-            Body={data},
-            CustomAttributes="accept_eula=true",
-        )
-        """  # type: ignore
-        logging_obj.pre_call(
-            input=data["inputs"],
-            api_key="",
-            additional_args={
-                "complete_input_dict": data,
-                "request_str": request_str,
-            },
-        )
         encoded_data = json.dumps(data).encode("utf-8")
         try:
-            response = await client.invoke_endpoint(
-                EndpointName=model,
-                ContentType="application/json",
-                Body=encoded_data,
-                CustomAttributes="accept_eula=true",
-            )
+            if model_id is not None:
+                ## LOGGING
+                request_str = f"""
+                response = client.invoke_endpoint(
+                    EndpointName={model},
+                    InferenceComponentName={model_id},
+                    ContentType="application/json",
+                    Body={data},
+                    CustomAttributes="accept_eula=true",
+                )
+                """  # type: ignore
+                logging_obj.pre_call(
+                    input=data["inputs"],
+                    api_key="",
+                    additional_args={
+                        "complete_input_dict": data,
+                        "request_str": request_str,
+                    },
+                )
+                response = await client.invoke_endpoint(
+                    EndpointName=model,
+                    InferenceComponentName=model_id,
+                    ContentType="application/json",
+                    Body=encoded_data,
+                    CustomAttributes="accept_eula=true",
+                )
+            else:
+                ## LOGGING
+                request_str = f"""
+                response = client.invoke_endpoint(
+                    EndpointName={model},
+                    ContentType="application/json",
+                    Body={data},
+                    CustomAttributes="accept_eula=true",
+                )
+                """  # type: ignore
+                logging_obj.pre_call(
+                    input=data["inputs"],
+                    api_key="",
+                    additional_args={
+                        "complete_input_dict": data,
+                        "request_str": request_str,
+                    },
+                )
+                response = await client.invoke_endpoint(
+                    EndpointName=model,
+                    ContentType="application/json",
+                    Body=encoded_data,
+                    CustomAttributes="accept_eula=true",
+                )
         except Exception as e:
-            raise SagemakerError(status_code=500, message=f"{str(e)}")
+            error_message = f"{str(e)}"
+            if "Inference Component Name header is required" in error_message:
+                error_message += "\n pass in via `litellm.completion(..., model_id={InferenceComponentName})`"
+            raise SagemakerError(status_code=500, message=error_message)
         response = await response["Body"].read()
         response = response.decode("utf8")
         ## LOGGING

litellm/tests/test_completion.py

@@ -1749,33 +1749,29 @@ def test_completion_sagemaker():
 # test_completion_sagemaker()
 
 
-@pytest.mark.skip(reason="AWS Suspended Account")
-def test_completion_sagemaker_stream():
+@pytest.mark.asyncio
+async def test_acompletion_sagemaker():
     try:
-        litellm.set_verbose = False
+        litellm.set_verbose = True
         print("testing sagemaker")
-        response = completion(
-            model="sagemaker/berri-benchmarking-Llama-2-70b-chat-hf-4",
+        response = await litellm.acompletion(
+            model="sagemaker/jumpstart-dft-hf-llm-mistral-7b-ins-20240329-150233",
+            model_id="huggingface-llm-mistral-7b-instruct-20240329-150233",
             messages=messages,
             temperature=0.2,
             max_tokens=80,
-            stream=True,
+            aws_region_name=os.getenv("AWS_REGION_NAME_2"),
+            aws_access_key_id=os.getenv("AWS_ACCESS_KEY_ID_2"),
+            aws_secret_access_key=os.getenv("AWS_SECRET_ACCESS_KEY_2"),
+            input_cost_per_second=0.000420,
         )
-
-        complete_streaming_response = ""
-        first_chunk_id, chunk_id = None, None
-        for i, chunk in enumerate(response):
-            print(chunk)
-            chunk_id = chunk.id
-            print(chunk_id)
-            if i == 0:
-                first_chunk_id = chunk_id
-            else:
-                assert chunk_id == first_chunk_id
-            complete_streaming_response += chunk.choices[0].delta.content or ""
         # Add any assertions here to check the response
-        # print(response)
-        assert len(complete_streaming_response) > 0
+        print(response)
+        cost = completion_cost(completion_response=response)
+        print("calculated cost", cost)
+        assert (
+            cost > 0.0 and cost < 1.0
+        )  # should never be > $1 for a single completion call
     except Exception as e:
         pytest.fail(f"Error occurred: {e}")
 

litellm/tests/test_streaming.py

@@ -1072,17 +1072,20 @@ def test_sagemaker_weird_response():
 # test_sagemaker_weird_response()
 
 
-@pytest.mark.skip(reason="AWS Suspended Account")
 @pytest.mark.asyncio
 async def test_sagemaker_streaming_async():
     try:
         messages = [{"role": "user", "content": "Hey, how's it going?"}]
         litellm.set_verbose = True
         response = await litellm.acompletion(
-            model="sagemaker/berri-benchmarking-Llama-2-70b-chat-hf-4",
+            model="sagemaker/jumpstart-dft-hf-llm-mistral-7b-ins-20240329-150233",
+            model_id="huggingface-llm-mistral-7b-instruct-20240329-150233",
             messages=messages,
-            max_tokens=100,
-            temperature=0.7,
+            temperature=0.2,
+            max_tokens=80,
+            aws_region_name=os.getenv("AWS_REGION_NAME_2"),
+            aws_access_key_id=os.getenv("AWS_ACCESS_KEY_ID_2"),
+            aws_secret_access_key=os.getenv("AWS_SECRET_ACCESS_KEY_2"),
             stream=True,
         )
         # Add any assertions here to check the response
@@ -1111,14 +1114,17 @@ async def test_sagemaker_streaming_async():
 # asyncio.run(test_sagemaker_streaming_async())
 
 
-@pytest.mark.skip(reason="AWS Suspended Account")
 def test_completion_sagemaker_stream():
     try:
         response = completion(
-            model="sagemaker/berri-benchmarking-Llama-2-70b-chat-hf-4",
+            model="sagemaker/jumpstart-dft-hf-llm-mistral-7b-ins-20240329-150233",
+            model_id="huggingface-llm-mistral-7b-instruct-20240329-150233",
             messages=messages,
             temperature=0.2,
             max_tokens=80,
+            aws_region_name=os.getenv("AWS_REGION_NAME_2"),
+            aws_access_key_id=os.getenv("AWS_ACCESS_KEY_ID_2"),
+            aws_secret_access_key=os.getenv("AWS_SECRET_ACCESS_KEY_2"),
             stream=True,
         )
         complete_response = ""