Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat(AwsVllmComponent): Add AWS Cognito as auth service #32

Merged
merged 14 commits into from
Oct 10, 2024
Merged

feat(AwsVllmComponent): Add AWS Cognito as auth service #32

Show file tree
Hide file tree
Changes from 9 commits
Commits
Show all changes
14 commits
Select commit Hold shift + click to select a range
7173bb5
feat(AwsVllmComponent): Add AWS Cognito as authentication service for…
bramelfrink Oct 4, 2024
242689d
refactor(AwsVllmComponent): Require an API key when access is not for…
bramelfrink Oct 4, 2024
ff5af29
feat(vllm): Store API key in secrets manager, and create default usag…
bramelfrink Oct 7, 2024
83dfccf
feat(AwsVllmComponent): Add AWS Cognito as authentication service for…
bramelfrink Oct 4, 2024
f5a291a
feat(vllm): Allow for API key and corresponding usage plan when API i…
bramelfrink Oct 8, 2024
b8d2432
feat(vllm): Allow for API key and corresponding usage plan when API i…
bramelfrink Oct 8, 2024
8cef40e
Merge branch 'main' into feat/cognito-auth-service
bramelfrink Oct 9, 2024
bcd7729
fix(vllm): Ensure appropriate errors are raised
bramelfrink Oct 9, 2024
69c4423
chore: remove print
bramelfrink Oct 9, 2024
7a243d9
chore: reflect review feedback
bramelfrink Oct 9, 2024
0fbe6d8
chore: add tiered usage plans
bramelfrink Oct 9, 2024
6fe1bcf
fix: tests
bramelfrink Oct 9, 2024
7b0863f
fix: tests
bramelfrink Oct 10, 2024
914e370
Update src/damavand/cloud/aws/resources/vllm_component.py
bramelfrink Oct 10, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
246 changes: 143 additions & 103 deletions src/damavand/cloud/aws/resources/vllm_component.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -125,19 +125,29 @@ def __init__(
)

self.args = args

_ = self.model
_ = self.endpoint_config
_ = self.endpoint

if self.args.public_internet_access:
_ = self.api
_ = self.api_resource_completions
_ = self.api_method
_ = self.api_integration
_ = self.api_integration_response
_ = self.api_method_response
_ = self.api_deployment
_ = self.endpoint_ssm_parameter
_ = self.api
_ = self.api_resource_v1
_ = self.api_resource_chat
_ = self.api_resource_completions

# Only create API key if public internet access is set to False
if not self.args.public_internet_access:
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It looks like you are changing the behavior to always create the API Gateway which I do agree with but, change the arg name to something meaningful and related like this:

Suggested change
if not self.args.public_internet_access:
if not self.args.is_api_key_required:

_ = self.admin_api_key
_ = self.default_usage_plan
_ = self.api_key_usage_plan
_ = self.api_key_secret
_ = self.api_key_secret_version

_ = self.api_method
_ = self.api_integration
_ = self.api_integration_response
_ = self.api_method_response
_ = self.api_deployment

def get_service_assume_policy(self, service: str) -> dict[str, Any]:
"""Return the assume role policy for the requested service.
Expand Down Expand Up @@ -264,18 +274,8 @@ def endpoint(self) -> aws.sagemaker.Endpoint:
def api(self) -> aws.apigateway.RestApi:
"""
Return a public API for the SageMaker endpoint.

Raises
------
AttributeError
When public_internet_access is False.
"""

if not self.args.public_internet_access:
raise AttributeError(
"`api` is only available when public_internet_access is True"
)

return aws.apigateway.RestApi(
resource_name=f"{self._name}-api",
opts=ResourceOptions(parent=self),
Expand All @@ -290,17 +290,8 @@ def api_resource_v1(self) -> aws.apigateway.Resource:
"""
Return a resource for the API Gateway.

Raises
------
AttributeError
When public_internet_access is False.
"""

if not self.args.public_internet_access:
raise AttributeError(
"`api_resource`is only available when public_internet_access is True"
)

return aws.apigateway.Resource(
resource_name=f"{self._name}-api-resource-v1",
opts=ResourceOptions(parent=self),
Expand All @@ -321,11 +312,6 @@ def api_resource_chat(self) -> aws.apigateway.Resource:
When public_internet_access is False.
"""

if not self.args.public_internet_access:
raise AttributeError(
"`api_resource`is only available when public_internet_access is True"
)

return aws.apigateway.Resource(
resource_name=f"{self._name}-api-resource-chat",
opts=ResourceOptions(parent=self),
Expand All @@ -339,18 +325,8 @@ def api_resource_chat(self) -> aws.apigateway.Resource:
def api_resource_completions(self) -> aws.apigateway.Resource:
"""
Return a resource for the API Gateway.

Raises
------
AttributeError
When public_internet_access is False.
"""

if not self.args.public_internet_access:
raise AttributeError(
"`api_resource`is only available when public_internet_access is True"
)

return aws.apigateway.Resource(
resource_name=f"{self._name}-api-resource-completions",
opts=ResourceOptions(parent=self),
Expand All @@ -359,41 +335,150 @@ def api_resource_completions(self) -> aws.apigateway.Resource:
path_part="completions",
)


@property
@cache
def api_method(self) -> aws.apigateway.Method:
"""
Return a method for the API Gateway.
"""

if self.args.public_internet_access:
return aws.apigateway.Method(
resource_name=f"{self._name}-api-method",
opts=ResourceOptions(parent=self),
rest_api=self.api.id,
resource_id=self.api_resource_completions.id,
http_method="POST",
authorization="NONE",
)
else:
return aws.apigateway.Method(
resource_name=f"{self._name}-api-method",
opts=ResourceOptions(parent=self),
rest_api=self.api.id,
resource_id=self.api_resource_completions.id,
http_method="POST",
authorization="NONE",
api_key_required=True,
)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You can remove the if condition just by passing the public_internet_access to the api_key_required argument:

Suggested change
if self.args.public_internet_access:
return aws.apigateway.Method(
resource_name=f"{self._name}-api-method",
opts=ResourceOptions(parent=self),
rest_api=self.api.id,
resource_id=self.api_resource_completions.id,
http_method="POST",
authorization="NONE",
)
else:
return aws.apigateway.Method(
resource_name=f"{self._name}-api-method",
opts=ResourceOptions(parent=self),
rest_api=self.api.id,
resource_id=self.api_resource_completions.id,
http_method="POST",
authorization="NONE",
api_key_required=True,
)
return aws.apigateway.Method(
resource_name=f"{self._name}-api-method",
opts=ResourceOptions(parent=self),
rest_api=self.api.id,
resource_id=self.api_resource_completions.id,
http_method="POST",
authorization="NONE",
api_key_required=self.args.public_internet_access,
)


@property
@cache
def admin_api_key(self) -> aws.apigateway.ApiKey:
"""
Return the admin API key for the API Gateway


Raises
------
AttributeError
When public_internet_access is False.
When public_internet_access is True.
"""
if self.args.public_internet_access:
raise AttributeError("`admin_api_key` is only available when public_internet_access is False")

return aws.apigateway.ApiKey(
resource_name=f"{self._name}-api-key",
opts=ResourceOptions(parent=self),
)

@property
@cache
def api_key_secret(self) -> aws.secretsmanager.Secret:
"""
Return the secret for the API key

if not self.args.public_internet_access:
raise AttributeError(
"`api_method`is only available when public_internet_access is True"
)

return aws.apigateway.Method(
resource_name=f"{self._name}-api-method",
Raises
------
AttributeError
When public_internet_access is True.
"""
if self.args.public_internet_access:
raise AttributeError("`admin_api_secret` is only available when public_internet_access is False")

return aws.secretsmanager.Secret(
resource_name=f"{self._name}-api-key-secret",
opts=ResourceOptions(parent=self),
rest_api=self.api.id,
resource_id=self.api_resource_completions.id,
http_method="POST",
authorization="NONE",
)

@property
def api_sagemaker_integration_uri(self) -> pulumi.Output[str]:
@cache
def api_key_secret_version(self) -> aws.secretsmanager.SecretVersion:
"""
Return the SageMaker model integration URI for the API Gateway
Return the secret version for the API key

Raises
------
AttributeError
When public_internet_access is False.
When public_internet_access is True.
"""
if self.args.public_internet_access:
raise AttributeError("`api_key_secret_version` is only available when public_internet_access is False")

return aws.secretsmanager.SecretVersion(
resource_name=f"{self._name}-api-key-secret-version",
opts=ResourceOptions(parent=self, depends_on=[self.api_key_secret]),
secret_id=self.api_key_secret.id,
secret_string=self.admin_api_key.id,
)


@property
@cache
def default_usage_plan(self) -> aws.apigateway.UsagePlan:
"""
Return a default usage plan for the API Gateway, that does not limit the usage.

Raises
------
AttributeError
When public_internet_access is True.
"""
if self.args.public_internet_access:
raise AttributeError("`default_usage_plan` is only available when public_internet_access is False")

return aws.apigateway.UsagePlan(
resource_name=f"{self._name}-api-usage-plan",
opts=ResourceOptions(parent=self),
api_stages=[
aws.apigateway.UsagePlanApiStageArgs(
api_id=self.api.id,
# NOTE: How do we want to deal with API stages vs. AWS environments?
stage=self.args.api_env_name,
)
],
)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can we have a two or three other default usage plan?



@property
@cache
def api_key_usage_plan(self) -> aws.apigateway.UsagePlanKey:
"""
Return the usage plan key for the API Gateway

Raises
------
AttributeError
When public_internet_access is True.
"""
if self.args.public_internet_access:
raise AttributeError("`api_key_usage_plan` is only available when public_internet_access is False")

return aws.apigateway.UsagePlanKey(
resource_name=f"{self._name}-api-usage-plan-key",
opts=ResourceOptions(parent=self),
key_id=self.admin_api_key.id,
key_type="API_KEY",
usage_plan_id=self.default_usage_plan.id,
)

@property
def api_sagemaker_integration_uri(self) -> pulumi.Output[str]:
"""
Return the SageMaker model integration URI for the API Gateway

"""

return self.endpoint.name.apply(
Expand All @@ -414,17 +499,8 @@ def api_access_sagemaker_role(self) -> aws.iam.Role:
"""
Return an execution role for APIGateway to access SageMaker endpoints.

Raises
------
AttributeError
When public_internet_access is False.
"""

if not self.args.public_internet_access:
raise AttributeError(
"`api_access_sagemaker_rol`is only available when public_internet_access is True"
)

return aws.iam.Role(
resource_name=f"{self._name}-api-sagemaker-access-role",
opts=ResourceOptions(parent=self),
Expand All @@ -440,17 +516,8 @@ def api_integration(self) -> aws.apigateway.Integration:
"""
Return a sagemaker integration for the API Gateway.

Raises
------
AttributeError
When public_internet_access is False.
"""

if not self.args.public_internet_access:
raise AttributeError(
"`api_integration`is only available when public_internet_access is True"
)

return aws.apigateway.Integration(
resource_name=f"{self._name}-api-integration",
opts=ResourceOptions(parent=self),
Expand All @@ -469,17 +536,8 @@ def api_integration_response(self) -> aws.apigateway.IntegrationResponse:
"""
Return a sagemaker integration response for the API Gateway.

Raises
------
AttributeError
When public_internet_access is False.
"""

if not self.args.public_internet_access:
raise AttributeError(
"`api_integration_response`is only available when public_internet_access is True"
)

return aws.apigateway.IntegrationResponse(
resource_name=f"{self._name}-api-integration-response",
opts=ResourceOptions(parent=self, depends_on=[self.api_integration]),
Expand All @@ -495,23 +553,14 @@ def api_method_response(self) -> aws.apigateway.MethodResponse:
"""
Return a sagemaker method response for the API Gateway.

Raises
------
AttributeError
When public_internet_access is False.
"""

if not self.args.public_internet_access:
raise AttributeError(
"`api_method_response`is only available when public_internet_access is True"
)

return aws.apigateway.MethodResponse(
resource_name=f"{self._name}-api-method-response",
opts=ResourceOptions(parent=self),
rest_api=self.api.id,
resource_id=self.api_resource_completions.id,
http_method="POST",
http_method=self.api_method.http_method,
status_code="200",
)

Expand All @@ -521,17 +570,8 @@ def api_deployment(self) -> aws.apigateway.Deployment:
"""
Return an API deployment for the API Gateway.

Raises
------
AttributeError
When public_internet_access is False.
"""

if not self.args.public_internet_access:
raise AttributeError(
"`api_deploy`is only available when public_internet_access is True"
)

return aws.apigateway.Deployment(
resource_name=f"{self._name}-api-deploy",
opts=ResourceOptions(
Expand Down
Loading
Loading