Add role overrides (#172)

Co-authored-by: Evan Stohlmann <[email protected]>
awslabs · Dec 5, 2024 · 208d6e7 · 208d6e7
1 parent 23fb773
commit 208d6e7
Show file tree

Hide file tree

Showing 34 changed files with 973 additions and 741 deletions.
diff --git a/Makefile b/Makefile
@@ -217,6 +217,7 @@ cleanTypeScript:
 	@find . -type d -name "build" -exec rm -rf {} +
 	@find . -type d -name ".tscache" -exec rm -rf {} +
 	@find . -type d -name ".jest_cache" -exec rm -rf {} +
+	@find . -type d -name "node_modules" -exec rm -rf {} +
 
 
 ## Delete CloudFormation outputs
@@ -280,11 +281,16 @@ endif
 ## Tear down all infrastructure
 destroy: cleanMisc
 	$(call print_config)
+ifneq (,$(findstring true, $(HEADLESS)))
+	npx cdk destroy ${STACK} --force $(if $(PROFILE),--profile ${PROFILE});
+else
 	@printf "Is the configuration correct? [y/N]  "\
 	&& read confirm_config &&\
 	if [ $${confirm_config:-'N'} = 'y' ]; then \
 		npx cdk destroy ${STACK} --force $(if $(PROFILE),--profile ${PROFILE}); \
 	fi;
+endif
+
 
 
 #################################################################################

diff --git a/ecs_model_deployer/src/lib/ecsCluster.ts b/ecs_model_deployer/src/lib/ecsCluster.ts
@@ -12,7 +12,7 @@
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
-*/
+ */
 
 // ECS Cluster Construct.
 import { CfnOutput, Duration, RemovalPolicy } from 'aws-cdk-lib';
@@ -38,25 +38,29 @@ import {
     Volume,
 } from 'aws-cdk-lib/aws-ecs';
 import { ApplicationLoadBalancer, BaseApplicationListenerProps } from 'aws-cdk-lib/aws-elasticloadbalancingv2';
-import { IRole, ManagedPolicy, ServicePrincipal, Role } from 'aws-cdk-lib/aws-iam';
+import { IRole, ManagedPolicy, Role, ServicePrincipal } from 'aws-cdk-lib/aws-iam';
 import { StringParameter } from 'aws-cdk-lib/aws-ssm';
 import { Construct } from 'constructs';
 
 import { createCdkId } from './utils';
-import { BaseProps, ECSConfig, Ec2Metadata, EcsSourceType } from './ecs-schema';
+import { BaseProps, Ec2Metadata, ECSConfig, EcsSourceType } from './ecs-schema';
 
 /**
  * Properties for the ECSCluster Construct.
  *
  * @property {IVpc} vpc - The virtual private cloud (VPC).
- * @property {SecurityGroups} securityGroups - The security group that the ECS cluster should use.
+ * @property {ISecurityGroup} securityGroup - The security group that the ECS cluster should use.
  * @property {ECSConfig} ecsConfig - The configuration for the cluster.
+ * @property {string} taskRoleName? - The role applied to the task
+ * @property {string} executionRoleName? - The role used for executing the task
  */
 type ECSClusterProps = {
     ecsConfig: ECSConfig;
     securityGroup: ISecurityGroup;
     vpc: IVpc;
     subnetSelection?: SubnetSelection;
+    taskRoleName?: string;
+    executionRoleName?: string;
 } & BaseProps;
 
 /**
@@ -79,7 +83,7 @@ export class ECSCluster extends Construct {
    */
     constructor (scope: Construct, id: string, props: ECSClusterProps) {
         super(scope, id);
-        const { config, vpc, securityGroup, ecsConfig, subnetSelection } = props;
+        const { config, vpc, securityGroup, ecsConfig, subnetSelection, taskRoleName, executionRoleName } = props;
 
         // Create ECS cluster
         const cluster = new Cluster(this, createCdkId([ecsConfig.identifier, 'Cl']), {
@@ -182,30 +186,17 @@ export class ECSCluster extends Construct {
             environment.SSL_CERT_FILE = config.certificateAuthorityBundle;
         }
 
-        const taskPolicyId = createCdkId([config.deploymentName, 'ECSPolicy']);
-        const taskPolicyStringParam = StringParameter.fromStringParameterName(this, 'taskPolicyStringParam',
-            `${config.deploymentPrefix}/policies/${taskPolicyId}`
-        );
-        const taskPolicy = ManagedPolicy.fromManagedPolicyArn(this, taskPolicyId, taskPolicyStringParam.stringValue);
-        const role_id = ecsConfig.identifier;
-        const roleName = createCdkId([config.deploymentName, role_id, 'Role']);
-        const taskRole = new Role(this, createCdkId([role_id, 'Role']), {
-            assumedBy: new ServicePrincipal('ecs-tasks.amazonaws.com'),
-            roleName,
-            description: `Allow ${role_id} ${role_id} ECS task access to AWS resources`,
-            managedPolicies: [taskPolicy],
-        });
-        new StringParameter(this, createCdkId([config.deploymentName, role_id, 'SP']), {
-            parameterName: `${config.deploymentPrefix}/roles/${role_id}`,
-            stringValue: taskRole.roleArn,
-            description: `Role ARN for LISA ${role_id} ${role_id} ECS Task`,
-        });
+        const roleId = ecsConfig.identifier;
+        const taskRole = taskRoleName ?
+            Role.fromRoleName(this, createCdkId([config.deploymentName, roleId]), taskRoleName) :
+            this.createTaskRole(config.deploymentName, config.deploymentPrefix, roleId);
 
         // Create ECS task definition
-        const taskDefinition = new Ec2TaskDefinition(this, createCdkId([ecsConfig.identifier, 'Ec2TaskDefinition']), {
-            family: createCdkId([config.deploymentName, ecsConfig.identifier], 32, 2),
-            taskRole: taskRole,
+        const taskDefinition = new Ec2TaskDefinition(this, createCdkId([roleId, 'Ec2TaskDefinition']), {
+            family: createCdkId([config.deploymentName, roleId], 32, 2),
             volumes: volumes,
+            taskRole,
+            ...(executionRoleName && { executionRole: Role.fromRoleName(this, createCdkId([config.deploymentName, roleId, 'EX']), executionRoleName) }),
         });
 
         // Add container to task definition
@@ -351,4 +342,28 @@ export class ECSCluster extends Construct {
         this.container = container;
         this.taskRole = taskRole;
     }
+
+    createTaskRole (deploymentName: string, deploymentPrefix: string | undefined, roleId: string): IRole {
+        const taskPolicyId = createCdkId([deploymentName, 'ECSPolicy']);
+        const taskPolicyStringParam = StringParameter.fromStringParameterName(this, 'taskPolicyStringParam',
+            `${deploymentPrefix}/policies/${taskPolicyId}`,
+        );
+
+        const taskPolicy = ManagedPolicy.fromManagedPolicyArn(this, taskPolicyId, taskPolicyStringParam.stringValue);
+        const roleName = createCdkId([roleId, 'Role']);
+        const role = new Role(this, roleName, {
+            assumedBy: new ServicePrincipal('ecs-tasks.amazonaws.com'),
+            roleName,
+            description: `Allow ${roleId} ECS task access to AWS resources`,
+            managedPolicies: [taskPolicy],
+        });
+
+        new StringParameter(this, createCdkId([deploymentName, roleId, 'SP']), {
+            parameterName: `${deploymentPrefix}/roles/${roleId}`,
+            stringValue: role.roleArn,
+            description: `Role ARN for LISA ${roleId} ECS Task`,
+        });
+
+        return role;
+    }
 }
diff --git a/lib/api-base/ecsCluster.ts b/lib/api-base/ecsCluster.ts
@@ -177,6 +177,13 @@ export class ECSCluster extends Construct {
             environment.SSL_CERT_FILE = config.certificateAuthorityBundle;
         }
 
+        // Retrieve execution role if it has been overridden
+        const executionRole = config.roles ? Role.fromRoleArn(
+            this,
+            createCdkId([ecsConfig.identifier, 'ER']),
+            StringParameter.valueForStringParameter(this, `${config.deploymentPrefix}/roles/${ecsConfig.identifier}EX`),
+        ) : undefined;
+
         // Create ECS task definition
         const taskRole = Role.fromRoleArn(
             this,
@@ -185,8 +192,9 @@ export class ECSCluster extends Construct {
         );
         const taskDefinition = new Ec2TaskDefinition(this, createCdkId([ecsConfig.identifier, 'Ec2TaskDefinition']), {
             family: createCdkId([config.deploymentName, ecsConfig.identifier], 32, 2),
-            taskRole: taskRole,
-            volumes: volumes,
+            volumes,
+            ...(taskRole && {taskRole}),
+            ...(executionRole && {executionRole}),
         });
 
         // Add container to task definition

diff --git a/lib/chat/api/configuration.ts b/lib/chat/api/configuration.ts
@@ -17,7 +17,6 @@
 import { IAuthorizer, RestApi } from 'aws-cdk-lib/aws-apigateway';
 import * as dynamodb from 'aws-cdk-lib/aws-dynamodb';
 import { ISecurityGroup } from 'aws-cdk-lib/aws-ec2';
-import { Role } from 'aws-cdk-lib/aws-iam';
 import { LayerVersion, Runtime } from 'aws-cdk-lib/aws-lambda';
 import { StringParameter } from 'aws-cdk-lib/aws-ssm';
 import { Construct } from 'constructs';
@@ -27,6 +26,7 @@ import { BaseProps } from '../../schema';
 import { createLambdaRole } from '../../core/utils';
 import { Vpc } from '../../networking/vpc';
 import { AwsCustomResource, PhysicalResourceId } from 'aws-cdk-lib/custom-resources';
+import { IRole } from 'aws-cdk-lib/aws-iam';
 
 /**
  * Properties for ConfigurationApi Construct.
@@ -78,7 +78,7 @@ export class ConfigurationApi extends Construct {
             removalPolicy: config.removalPolicy,
         });
 
-        const lambdaRole: Role = createLambdaRole(this, config.deploymentName, 'ConfigurationApi', configTable.tableArn);
+        const lambdaRole: IRole = createLambdaRole(this, config.deploymentName, 'ConfigurationApi', configTable.tableArn, config.roles?.LambdaConfigurationApiExecutionRole);
 
         // Populate the App Config table with default config
         const date = new Date();

diff --git a/lib/chat/api/session.ts b/lib/chat/api/session.ts
@@ -16,8 +16,8 @@
 
 import { IAuthorizer, RestApi } from 'aws-cdk-lib/aws-apigateway';
 import * as dynamodb from 'aws-cdk-lib/aws-dynamodb';
+import { IRole } from 'aws-cdk-lib/aws-iam';
 import { ISecurityGroup } from 'aws-cdk-lib/aws-ec2';
-import { Role } from 'aws-cdk-lib/aws-iam';
 import { LayerVersion, Runtime } from 'aws-cdk-lib/aws-lambda';
 import { StringParameter } from 'aws-cdk-lib/aws-ssm';
 import { Construct } from 'constructs';
@@ -146,7 +146,7 @@ export class SessionApi extends Construct {
             },
         ];
 
-        const lambdaRole: Role = createLambdaRole(this, config.deploymentName, 'SessionApi', sessionTable.tableArn);
+        const lambdaRole: IRole = createLambdaRole(this, config.deploymentName, 'SessionApi', sessionTable.tableArn, config.roles?.LambdaExecutionRole);
 
         apis.forEach((f) => {
             const lambdaFunction = registerAPIEndpoint(

diff --git a/lib/core/api_base.ts b/lib/core/api_base.ts
@@ -21,6 +21,7 @@ import { Construct } from 'constructs';
 import { CustomAuthorizer } from '../api-base/authorizer';
 import { BaseProps } from '../schema';
 import { Vpc } from '../networking/vpc';
+import { Role } from 'aws-cdk-lib/aws-iam';
 
 type LisaApiBaseStackProps = {
     vpc: Vpc;
@@ -63,6 +64,10 @@ export class LisaApiBaseStack extends Stack {
             config: config,
             securityGroups: [vpc.securityGroups.lambdaSg],
             vpc,
+            ...(config.roles &&
+            {
+                role: Role.fromRoleName(this, 'AuthorizerRole', config.roles.RestApiAuthorizerRole),
+            })
         });
 
         this.restApi = restApi;

diff --git a/lib/core/api_deployment.ts b/lib/core/api_deployment.ts
@@ -14,7 +14,7 @@
   limitations under the License.
 */
 
-import { Stack, StackProps } from 'aws-cdk-lib';
+import { CfnOutput, Stack, StackProps, Aws } from 'aws-cdk-lib';
 import { Deployment, RestApi } from 'aws-cdk-lib/aws-apigateway';
 import { Construct } from 'constructs';
 
@@ -42,5 +42,10 @@ export class LisaApiDeploymentStack extends Stack {
         // Hack to allow deploying to an existing stage
         // https://github.com/aws/aws-cdk/issues/25582
         (deployment as any).resource.stageName = config.deploymentStage;
+
+        new CfnOutput(this, 'ApiUrl', {
+            value: `https://${restApiId}.execute-api.${this.region}.${Aws.URL_SUFFIX}/${config.deploymentStage}`,
+            description: 'API Gateway URL'
+        });
     }
 }
diff --git a/lib/core/iam/roles.ts b/lib/core/iam/roles.ts
@@ -0,0 +1,73 @@
+/**
+ Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+
+ Licensed under the Apache License, Version 2.0 (the "License").
+ You may not use this file except in compliance with the License.
+ You may obtain a copy of the License at
+
+ http://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing, software
+ distributed under the License is distributed on an "AS IS" BASIS,
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ See the License for the specific language governing permissions and
+ limitations under the License.
+ */
+
+/**
+ * List of all roles used for overrides
+ */
+export enum Roles {
+    DOCKER_IMAGE_BUILDER_DEPLOYMENT_ROLE = 'DockerImageBuilderDeploymentRole',
+    DOCKER_IMAGE_BUILDER_EC2_ROLE = 'DockerImageBuilderEC2Role',
+    DOCKER_IMAGE_BUILDER_ROLE = 'DockerImageBuilderRole',
+    DOCS_ROLE = 'DocsRole',
+    DOCS_DEPLOYER_ROLE = 'DocsDeployerRole',
+    ECS_MODEL_DEPLOYER_ROLE = 'ECSModelDeployerRole',
+    ECS_MODEL_TASK_ROLE = 'ECSModelTaskRole',
+    // eslint-disable-next-line no-unused-vars
+    ECS_REST_API_ROLE = 'ECSRestApiRole',
+    // eslint-disable-next-line no-unused-vars
+    ECS_REST_API_EX_ROLE = 'ECSRestApiExRole',
+    LAMBDA_EXECUTION_ROLE = 'LambdaExecutionRole',
+    LAMBDA_CONFIGURATION_API_EXECUTION_ROLE = 'LambdaConfigurationApiExecutionRole',
+    MODEL_API_ROLE = 'ModelApiRole',
+    MODEL_SFN_LAMBDA_ROLE = 'ModelsSfnLambdaRole',
+    MODEL_SFN_ROLE = 'ModelSfnRole',
+    RAG_LAMBDA_EXECUTION_ROLE = 'LisaRagLambdaExecutionRole',
+    REST_API_AUTHORIZER_ROLE = 'RestApiAuthorizerRole',
+    S3_READER_ROLE = 'S3ReaderRole',
+    UI_DEPLOYMENT_ROLE = 'UIDeploymentRole',
+}
+
+export const RoleNames: Record<string, string> = {
+    'DockerImageBuilderDeploymentRole': 'DockerImageBuilderDeploymentRole',
+    'DockerImageBuilderEC2Role': 'DockerImageBuilderEC2Role',
+    'DockerImageBuilderRole': 'DockerImageBuilderRole',
+    'DocsRole': 'DocsRole',
+    'DocsDeployerRole': 'DocsDeployerRole',
+    'ECSModelDeployerRole': 'ECSModelDeployerRole',
+    'ECSModelTaskRole': 'ECSModelTaskRole',
+    'ECSRestApiRole': 'ECSRestApiRole',
+    'ECSRestApiExRole': 'ECSRestApiExRole',
+    'LambdaExecutionRole': 'LambdaExecutionRole',
+    'LambdaConfigurationApiExecutionRole': 'LambdaConfigurationApiExecutionRole',
+    'ModelApiRole': 'ModelApiRole',
+    'ModelsSfnLambdaRole': 'ModelsSfnLambdaRole',
+    'ModelSfnRole': 'ModelSfnRole',
+    'RagLambdaExecutionRole': 'RAGRole',
+    'RestApiAuthorizerRole': 'RestApiAuthorizerRole',
+    'S3ReaderRole': 'S3ReaderRole',
+    'UIDeploymentRole': 'UIDeploymentRole',
+};
+
+export function of (key: string): Roles {
+    const keys = Object.keys(Roles).filter((x) => x === key);
+    if (keys.length > 0)
+        return Roles[keys[0] as keyof typeof Roles] as Roles;
+    else {
+        throw Error(`No Roles entry exists for ${key}`);
+    }
+}
+
+export const ROLE = 'Role';