From 109cd93a397806f77feba0b20d92cbab885d4e51 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Fri, 29 Mar 2024 09:05:00 -0700 Subject: [PATCH] fix(sagemaker.py): support model_id consistently. support dynamic args for async calls --- litellm/llms/sagemaker.py | 201 +++++++++++++++++++++---------- litellm/tests/test_completion.py | 36 +++--- litellm/tests/test_streaming.py | 18 ++- 3 files changed, 165 insertions(+), 90 deletions(-) diff --git a/litellm/llms/sagemaker.py b/litellm/llms/sagemaker.py index 0aa987540e..2604875b47 100644 --- a/litellm/llms/sagemaker.py +++ b/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 diff --git a/litellm/tests/test_completion.py b/litellm/tests/test_completion.py index f23f59e854..c27dba4791 100644 --- a/litellm/tests/test_completion.py +++ b/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}") diff --git a/litellm/tests/test_streaming.py b/litellm/tests/test_streaming.py index 76bc610f2f..ababe7cd6e 100644 --- a/litellm/tests/test_streaming.py +++ b/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 = ""