From 9b39893980703ef3e840c5acbe249ab63ca22f8f Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Wed, 24 Jan 2024 10:29:28 -0800 Subject: [PATCH 1/7] v0 set arn number on dynamoDB /key/gen --- docs/my-website/docs/proxy/virtual_keys.md | 4 ++- litellm/proxy/_types.py | 2 ++ litellm/proxy/db/dynamo_db.py | 33 ++++++++++++++++++++++ litellm/proxy/proxy_config.yaml | 14 +++++---- 4 files changed, 46 insertions(+), 7 deletions(-) diff --git a/docs/my-website/docs/proxy/virtual_keys.md b/docs/my-website/docs/proxy/virtual_keys.md index e1c89bbc216..999db6055d0 100644 --- a/docs/my-website/docs/proxy/virtual_keys.md +++ b/docs/my-website/docs/proxy/virtual_keys.md @@ -554,7 +554,9 @@ general_settings: "region_name": "us-west-2" "user_table_name": "your-user-table", "key_table_name": "your-token-table", - "config_table_name": "your-config-table" + "config_table_name": "your-config-table", + "aws_role_name": "your-aws_role_name", + "aws_session_name": "your-aws_session_name", } ``` diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index d5dc841cb26..88ee88bcd0f 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -208,6 +208,8 @@ class DynamoDBArgs(LiteLLMBase): key_table_name: str = "LiteLLM_VerificationToken" config_table_name: str = "LiteLLM_Config" spend_table_name: str = "LiteLLM_SpendLogs" + aws_role_name: Optional[str] = None + aws_session_name: Optional[str] = None class ConfigGeneralSettings(LiteLLMBase): diff --git a/litellm/proxy/db/dynamo_db.py b/litellm/proxy/db/dynamo_db.py index 83cf6b15724..50330ccf3bf 100644 --- a/litellm/proxy/db/dynamo_db.py +++ b/litellm/proxy/db/dynamo_db.py @@ -52,6 +52,32 @@ class DynamoDBWrapper(CustomDB): self.database_arguments = database_arguments self.region_name = database_arguments.region_name + def set_env_vars_based_on_arn(self): + if self.database_arguments.aws_role_name is None: + return + verbose_proxy_logger.debug( + f"DynamoDB: setting env vars based on arn={self.database_arguments.aws_role_name}" + ) + import boto3, os + + sts_client = boto3.client("sts") + assumed_role = sts_client.assume_role( + RoleArn=self.database_arguments.aws_role_name, + RoleSessionName=self.database_arguments.aws_session_name, + ) + + aws_access_key_id = assumed_role["Credentials"]["AccessKeyId"] + aws_secret_access_key = assumed_role["Credentials"]["SecretAccessKey"] + aws_session_token = assumed_role["Credentials"]["SessionToken"] + + verbose_proxy_logger.debug( + f"Got STS assumed Role, aws_access_key_id={aws_access_key_id}" + ) + # set these in the env so aiodynamo can use them + os.environ["AWS_ACCESS_KEY_ID"] = aws_access_key_id + os.environ["AWS_SECRET_ACCESS_KEY"] = aws_secret_access_key + os.environ["AWS_SESSION_TOKEN"] = aws_session_token + async def connect(self): """ Connect to DB, and creating / updating any tables @@ -74,6 +100,7 @@ class DynamoDBWrapper(CustomDB): import aiohttp verbose_proxy_logger.debug("DynamoDB Wrapper - Attempting to connect") + self.set_env_vars_based_on_arn() # before making ClientSession check if ssl_verify=False if self.database_arguments.ssl_verify == False: client_session = ClientSession(connector=aiohttp.TCPConnector(ssl=False)) @@ -170,6 +197,8 @@ class DynamoDBWrapper(CustomDB): from aiohttp import ClientSession import aiohttp + self.set_env_vars_based_on_arn() + if self.database_arguments.ssl_verify == False: client_session = ClientSession(connector=aiohttp.TCPConnector(ssl=False)) else: @@ -210,6 +239,8 @@ class DynamoDBWrapper(CustomDB): from aiohttp import ClientSession import aiohttp + self.set_env_vars_based_on_arn() + if self.database_arguments.ssl_verify == False: client_session = ClientSession(connector=aiohttp.TCPConnector(ssl=False)) else: @@ -253,6 +284,7 @@ class DynamoDBWrapper(CustomDB): async def update_data( self, key: str, value: dict, table_name: Literal["user", "key", "config"] ): + self.set_env_vars_based_on_arn() from aiodynamo.client import Client from aiodynamo.credentials import Credentials, StaticCredentials from aiodynamo.http.httpx import HTTPX @@ -323,4 +355,5 @@ class DynamoDBWrapper(CustomDB): """ Not Implemented yet. """ + self.set_env_vars_based_on_arn() return super().delete_data(keys, table_name) diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml index 97168b19f9a..2bed73bcb8b 100644 --- a/litellm/proxy/proxy_config.yaml +++ b/litellm/proxy/proxy_config.yaml @@ -67,12 +67,14 @@ litellm_settings: general_settings: master_key: sk-1234 - # database_type: "dynamo_db" - # database_args: { # 👈 all args - https://github.com/BerriAI/litellm/blob/befbcbb7ac8f59835ce47415c128decf37aac328/litellm/proxy/_types.py#L190 - # "billing_mode": "PAY_PER_REQUEST", - # "region_name": "us-west-2", - # "ssl_verify": False - # } + database_type: "dynamo_db" + database_args: { # 👈 all args - https://github.com/BerriAI/litellm/blob/befbcbb7ac8f59835ce47415c128decf37aac328/litellm/proxy/_types.py#L190 + "billing_mode": "PAY_PER_REQUEST", + "region_name": "us-west-2", + "ssl_verify": False, + "aws_role_name": "", + "aws_session_name": "", + } From d2c1e1554fc59f4cb66cd5b4fa9015aa9733499c Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Fri, 2 Feb 2024 10:49:09 -0800 Subject: [PATCH 2/7] (fix) gitignore --- .gitignore | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/.gitignore b/.gitignore index 77ee0fbefa9..863e9c6416c 100644 --- a/.gitignore +++ b/.gitignore @@ -35,3 +35,8 @@ hosted_config.yaml litellm/proxy/tests/node_modules litellm/proxy/tests/package.json litellm/proxy/tests/package-lock.json +ui/litellm-dashboard/.next +ui/litellm-dashboard/node_modules +ui/litellm-dashboard/next-env.d.ts +ui/litellm-dashboard/package.json +ui/litellm-dashboard/package-lock.json \ No newline at end of file From 9c13c65ce24b27aa8a6c24ed0355d5c5a094af17 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Fri, 2 Feb 2024 10:57:06 -0800 Subject: [PATCH 3/7] (dynamo) use assume_role_with_web_identity --- litellm/proxy/_types.py | 5 +++++ litellm/proxy/db/dynamo_db.py | 7 ++++++- 2 files changed, 11 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 88ee88bcd0f..10290fc3c11 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -210,6 +210,11 @@ class DynamoDBArgs(LiteLLMBase): spend_table_name: str = "LiteLLM_SpendLogs" aws_role_name: Optional[str] = None aws_session_name: Optional[str] = None + aws_web_identity_token: Optional[str] = None + aws_provider_id: Optional[str] = None + aws_policy_arns: Optional[List[str]] = None + aws_policy: Optional[str] = None + aws_duration_seconds: Optional[int] = None class ConfigGeneralSettings(LiteLLMBase): diff --git a/litellm/proxy/db/dynamo_db.py b/litellm/proxy/db/dynamo_db.py index 50330ccf3bf..3ea2b285ae3 100644 --- a/litellm/proxy/db/dynamo_db.py +++ b/litellm/proxy/db/dynamo_db.py @@ -61,9 +61,14 @@ class DynamoDBWrapper(CustomDB): import boto3, os sts_client = boto3.client("sts") - assumed_role = sts_client.assume_role( + assumed_role = sts_client.assume_role_with_web_identity( RoleArn=self.database_arguments.aws_role_name, RoleSessionName=self.database_arguments.aws_session_name, + WebIdentityToken=self.database_arguments.aws_web_identity_token, + ProviderId=self.database_arguments.aws_provider_id, + PolicyArns=self.database_arguments.aws_policy_arns, + Policy=self.database_arguments.aws_policy, + DurationSeconds=self.database_arguments.aws_duration_seconds, ) aws_access_key_id = assumed_role["Credentials"]["AccessKeyId"] From e59aca925acd3fe8d269261f33c75e5e02203648 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Tue, 6 Feb 2024 08:18:17 -0800 Subject: [PATCH 4/7] fix(proxy_server.py): update dynamodb load config --- litellm/proxy/proxy_server.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index cc4222f8e12..4fc9222cee1 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1076,6 +1076,19 @@ class ProxyConfig: database_type == "dynamo_db" or database_type == "dynamodb" ): database_args = general_settings.get("database_args", None) + ### LOAD FROM os.environ/ ### + for k, v in database_args.items(): + if isinstance(v, str) and v.startswith("os.environ/"): + database_args[k] = litellm.get_secret(v) + if isinstance(k, str) and k == "aws_web_identity_token": + if os.path.exists(v): + with open(v, "r") as file: + token_content = file.read() + database_args[k] = json.loads(token_content) + else: + verbose_proxy_logger.info( + f"DynamoDB Loading - {v} is not a valid file path" + ) custom_db_client = DBClient( custom_db_args=database_args, custom_db_type=database_type ) From 551a4c60a7bf2b25e56e5436c57326b1e3eddfa9 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 7 Feb 2024 22:52:15 -0800 Subject: [PATCH 5/7] fix(dynamo_db.py): eliminating unneccessary params --- litellm/proxy/db/dynamo_db.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/litellm/proxy/db/dynamo_db.py b/litellm/proxy/db/dynamo_db.py index 3ea2b285ae3..f315a74be55 100644 --- a/litellm/proxy/db/dynamo_db.py +++ b/litellm/proxy/db/dynamo_db.py @@ -65,10 +65,6 @@ class DynamoDBWrapper(CustomDB): RoleArn=self.database_arguments.aws_role_name, RoleSessionName=self.database_arguments.aws_session_name, WebIdentityToken=self.database_arguments.aws_web_identity_token, - ProviderId=self.database_arguments.aws_provider_id, - PolicyArns=self.database_arguments.aws_policy_arns, - Policy=self.database_arguments.aws_policy, - DurationSeconds=self.database_arguments.aws_duration_seconds, ) aws_access_key_id = assumed_role["Credentials"]["AccessKeyId"] From c4038e38e5d0116a1df610a968f3e3b52685b62a Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Mon, 12 Feb 2024 19:54:39 -0800 Subject: [PATCH 6/7] (fix) dyanmo fixes --- litellm/proxy/_types.py | 2 ++ litellm/proxy/db/dynamo_db.py | 10 +++++++++- 2 files changed, 11 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 10290fc3c11..028d6f590c4 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -215,6 +215,8 @@ class DynamoDBArgs(LiteLLMBase): aws_policy_arns: Optional[List[str]] = None aws_policy: Optional[str] = None aws_duration_seconds: Optional[int] = None + assume_role_aws_role_name: Optional[str] = None + assume_role_aws_session_name: Optional[str] = None class ConfigGeneralSettings(LiteLLMBase): diff --git a/litellm/proxy/db/dynamo_db.py b/litellm/proxy/db/dynamo_db.py index f315a74be55..78e81c1542b 100644 --- a/litellm/proxy/db/dynamo_db.py +++ b/litellm/proxy/db/dynamo_db.py @@ -61,12 +61,20 @@ class DynamoDBWrapper(CustomDB): import boto3, os sts_client = boto3.client("sts") - assumed_role = sts_client.assume_role_with_web_identity( + + # call 1 + non_used_assumed_role = sts_client.assume_role_with_web_identity( RoleArn=self.database_arguments.aws_role_name, RoleSessionName=self.database_arguments.aws_session_name, WebIdentityToken=self.database_arguments.aws_web_identity_token, ) + # call 2 + assumed_role = sts_client.assume_role( + RoleArn=self.database_arguments.assume_role_aws_role_name, + RoleSessionName=self.database_arguments.assume_role_aws_session_name, + ) + aws_access_key_id = assumed_role["Credentials"]["AccessKeyId"] aws_secret_access_key = assumed_role["Credentials"]["SecretAccessKey"] aws_session_token = assumed_role["Credentials"]["SessionToken"] From 7d27690203a07dea3c9e50c7cda61819d58c5529 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Mon, 12 Feb 2024 21:41:13 -0800 Subject: [PATCH 7/7] (fix) use correct value for os.environ dynamo --- litellm/proxy/proxy_config.yaml | 1 + litellm/proxy/proxy_server.py | 13 +++++++++---- 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml index 2bed73bcb8b..de4719d6edd 100644 --- a/litellm/proxy/proxy_config.yaml +++ b/litellm/proxy/proxy_config.yaml @@ -74,6 +74,7 @@ general_settings: "ssl_verify": False, "aws_role_name": "", "aws_session_name": "", + "aws_web_identity_token": "os.environ/AWS_WEB_IDENTITY_TOKEN" } diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 4fc9222cee1..85142be6fe2 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1081,14 +1081,19 @@ class ProxyConfig: if isinstance(v, str) and v.startswith("os.environ/"): database_args[k] = litellm.get_secret(v) if isinstance(k, str) and k == "aws_web_identity_token": - if os.path.exists(v): - with open(v, "r") as file: + value = database_args[k] + verbose_proxy_logger.debug( + f"Loading AWS Web Identity Token from file: {value}" + ) + if os.path.exists(value): + with open(value, "r") as file: token_content = file.read() - database_args[k] = json.loads(token_content) + database_args[k] = token_content else: verbose_proxy_logger.info( - f"DynamoDB Loading - {v} is not a valid file path" + f"DynamoDB Loading - {value} is not a valid file path" ) + verbose_proxy_logger.debug(f"database_args: {database_args}") custom_db_client = DBClient( custom_db_args=database_args, custom_db_type=database_type )