From ffbe1845ef8327d9d57adaff2f1fd48d5aa75597 Mon Sep 17 00:00:00 2001 From: David Manouchehri Date: Wed, 31 Jul 2024 16:23:06 +0000 Subject: [PATCH 1/2] (bedrock_httpx.py) - Add support for custom STS endpoints, e.g. for FIPS. --- litellm/llms/bedrock_httpx.py | 31 ++++++++++++++++++++++++++-- litellm/tests/test_secret_manager.py | 24 ++++++++++++++++++++- 2 files changed, 52 insertions(+), 3 deletions(-) diff --git a/litellm/llms/bedrock_httpx.py b/litellm/llms/bedrock_httpx.py index a6cf32c77e..015223ecbe 100644 --- a/litellm/llms/bedrock_httpx.py +++ b/litellm/llms/bedrock_httpx.py @@ -383,6 +383,7 @@ class BedrockLLM(BaseLLM): aws_profile_name: Optional[str] = None, aws_role_name: Optional[str] = None, aws_web_identity_token: Optional[str] = None, + aws_sts_endpoint: Optional[str] = None, ): """ Return a boto3.Credentials object @@ -403,6 +404,7 @@ class BedrockLLM(BaseLLM): aws_profile_name, aws_role_name, aws_web_identity_token, + aws_sts_endpoint, ] # Iterate over parameters and update if needed @@ -421,6 +423,7 @@ class BedrockLLM(BaseLLM): aws_profile_name, aws_role_name, aws_web_identity_token, + aws_sts_endpoint, ) = params_to_check ### CHECK STS ### @@ -432,12 +435,19 @@ class BedrockLLM(BaseLLM): print_verbose( f"IN Web Identity Token: {aws_web_identity_token} | Role Name: {aws_role_name} | Session Name: {aws_session_name}" ) + + if aws_sts_endpoint is None: + sts_endpoint = f"https://sts.{aws_region_name}.amazonaws.com" + else: + sts_endpoint = aws_sts_endpoint + iam_creds_cache_key = json.dumps( { "aws_web_identity_token": aws_web_identity_token, "aws_role_name": aws_role_name, "aws_session_name": aws_session_name, "aws_region_name": aws_region_name, + "aws_sts_endpoint": sts_endpoint, } ) @@ -454,7 +464,7 @@ class BedrockLLM(BaseLLM): sts_client = boto3.client( "sts", region_name=aws_region_name, - endpoint_url=f"https://sts.{aws_region_name}.amazonaws.com", + endpoint_url=sts_endpoint, ) # https://docs.aws.amazon.com/STS/latest/APIReference/API_AssumeRoleWithWebIdentity.html @@ -849,6 +859,7 @@ class BedrockLLM(BaseLLM): "aws_bedrock_runtime_endpoint", None ) # https://bedrock-runtime.{region_name}.amazonaws.com aws_web_identity_token = optional_params.pop("aws_web_identity_token", None) + aws_sts_endpoint = optional_params.pop("aws_sts_endpoint", None) ### SET REGION NAME ### if aws_region_name is None: @@ -878,6 +889,7 @@ class BedrockLLM(BaseLLM): aws_profile_name=aws_profile_name, aws_role_name=aws_role_name, aws_web_identity_token=aws_web_identity_token, + aws_sts_endpoint=aws_sts_endpoint, ) ### SET RUNTIME ENDPOINT ### @@ -1536,6 +1548,7 @@ class BedrockConverseLLM(BaseLLM): aws_profile_name: Optional[str] = None, aws_role_name: Optional[str] = None, aws_web_identity_token: Optional[str] = None, + aws_sts_endpoint: Optional[str] = None, ): """ Return a boto3.Credentials object @@ -1552,6 +1565,7 @@ class BedrockConverseLLM(BaseLLM): aws_profile_name, aws_role_name, aws_web_identity_token, + aws_sts_endpoint, ] # Iterate over parameters and update if needed @@ -1570,6 +1584,7 @@ class BedrockConverseLLM(BaseLLM): aws_profile_name, aws_role_name, aws_web_identity_token, + aws_sts_endpoint, ) = params_to_check ### CHECK STS ### @@ -1578,12 +1593,22 @@ class BedrockConverseLLM(BaseLLM): and aws_role_name is not None and aws_session_name is not None ): + print_verbose( + f"IN Web Identity Token: {aws_web_identity_token} | Role Name: {aws_role_name} | Session Name: {aws_session_name}" + ) + + if aws_sts_endpoint is None: + sts_endpoint = f"https://sts.{aws_region_name}.amazonaws.com" + else: + sts_endpoint = aws_sts_endpoint + iam_creds_cache_key = json.dumps( { "aws_web_identity_token": aws_web_identity_token, "aws_role_name": aws_role_name, "aws_session_name": aws_session_name, "aws_region_name": aws_region_name, + "aws_sts_endpoint": sts_endpoint, } ) @@ -1600,7 +1625,7 @@ class BedrockConverseLLM(BaseLLM): sts_client = boto3.client( "sts", region_name=aws_region_name, - endpoint_url=f"https://sts.{aws_region_name}.amazonaws.com", + endpoint_url=sts_endpoint, ) # https://docs.aws.amazon.com/STS/latest/APIReference/API_AssumeRoleWithWebIdentity.html @@ -1815,6 +1840,7 @@ class BedrockConverseLLM(BaseLLM): "aws_bedrock_runtime_endpoint", None ) # https://bedrock-runtime.{region_name}.amazonaws.com aws_web_identity_token = optional_params.pop("aws_web_identity_token", None) + aws_sts_endpoint = optional_params.pop("aws_sts_endpoint", None) ### SET REGION NAME ### if aws_region_name is None: @@ -1844,6 +1870,7 @@ class BedrockConverseLLM(BaseLLM): aws_profile_name=aws_profile_name, aws_role_name=aws_role_name, aws_web_identity_token=aws_web_identity_token, + aws_sts_endpoint=aws_sts_endpoint, ) ### SET RUNTIME ENDPOINT ### diff --git a/litellm/tests/test_secret_manager.py b/litellm/tests/test_secret_manager.py index cd2f2731fc..904b291bfe 100644 --- a/litellm/tests/test_secret_manager.py +++ b/litellm/tests/test_secret_manager.py @@ -13,7 +13,7 @@ import pytest from litellm import get_secret from litellm.proxy.secret_managers.aws_secret_manager import load_aws_secret_manager from litellm.llms.azure import get_azure_ad_token_from_oidc -from litellm.llms.bedrock_httpx import BedrockLLM +from litellm.llms.bedrock_httpx import BedrockLLM, BedrockConverseLLM @pytest.mark.skip(reason="AWS Suspended Account") @@ -113,3 +113,25 @@ def test_oidc_circle_v1_with_amazon(): aws_role_name=aws_role_name, aws_session_name="assume-v1-session", ) + +@pytest.mark.skipif( + os.environ.get("CIRCLE_OIDC_TOKEN") is None, + reason="Cannot run without being in CircleCI Runner", +) +def test_oidc_circle_v1_with_amazon_fips(): + # The purpose of this test is to validate that we can assume a role in a FIPS region + + # TODO: This is using ai.moda's IAM role, we should use LiteLLM's IAM role eventually + aws_role_name = ( + "arn:aws:iam::335785316107:role/litellm-github-unit-tests-circleci-v1-assume-only" + ) + aws_web_identity_token = "oidc/circleci/" + + bllm = BedrockConverseLLM() + creds = bllm.get_credentials( + aws_region_name="us-west-1", + aws_web_identity_token=aws_web_identity_token, + aws_role_name=aws_role_name, + aws_session_name="assume-v1-session-fips", + aws_sts_endpoint="https://sts-fips.us-west-1.amazonaws.com", + ) From 507529e8df86a8e268312a575f5446efc7355f0c Mon Sep 17 00:00:00 2001 From: David Manouchehri Date: Wed, 31 Jul 2024 16:51:29 +0000 Subject: [PATCH 2/2] (test_bedrock_completion.py) - Use FIPS endpoints for testing. --- litellm/tests/test_bedrock_completion.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/litellm/tests/test_bedrock_completion.py b/litellm/tests/test_bedrock_completion.py index 498cfb5d89..4b1e544c7f 100644 --- a/litellm/tests/test_bedrock_completion.py +++ b/litellm/tests/test_bedrock_completion.py @@ -558,7 +558,7 @@ def test_completion_bedrock_httpx_command_r_sts_oidc_auth(): import os aws_web_identity_token = "oidc/circleci_v2/" - aws_region_name = os.environ["AWS_REGION_NAME"] + aws_region_name = "us-west-2" # aws_role_name = os.environ["AWS_TEMP_ROLE_NAME"] # TODO: This is using ai.moda's IAM role, we should use LiteLLM's IAM role eventually aws_role_name = "arn:aws:iam::335785316107:role/litellm-github-unit-tests-circleci" @@ -575,6 +575,8 @@ def test_completion_bedrock_httpx_command_r_sts_oidc_auth(): aws_web_identity_token=aws_web_identity_token, aws_role_name=aws_role_name, aws_session_name="my-test-session", + aws_sts_endpoint="https://sts-fips.us-west-2.amazonaws.com", + aws_bedrock_runtime_endpoint="https://bedrock-runtime-fips.us-west-2.amazonaws.com", ) # Add any assertions here to check the response print(response)