Merge pull request #134 from aws-samples/wangzt-dev

wzt1001 · web-flow · commit af9123bfa00a · 2024-07-01T09:08:10.000-07:00
cdk robustness update
diff --git a/source/resources/bin/main.ts b/source/resources/bin/main.ts
@@ -1,11 +1,13 @@
-#!/usr/bin/env node
 import * as cdk from 'aws-cdk-lib';
 import { MainStack } from '../lib/main-stack';
+
 const devEnv = {
     account: process.env.CDK_DEFAULT_ACCOUNT,
     region: process.env.CDK_DEFAULT_REGION,
 };
 
 const app = new cdk.App();
-new MainStack(app, 'GenBiMainStack', { env: devEnv });
+const deployRds = process.argv.includes('--deploy-rds'); // Check if --deploy-rds flag is present
+
+new MainStack(app, 'GenBiMainStack', { env: devEnv, deployRds }); // Pass deployRDS flag to MainStack constructor
 app.synth();
diff --git a/source/resources/lib/aos/aos-stack.ts b/source/resources/lib/aos/aos-stack.ts
@@ -7,24 +7,17 @@ import { AnyPrincipal, Effect, PolicyStatement } from "aws-cdk-lib/aws-iam";
 import * as crypto from 'crypto';
 
 export class AOSStack extends cdk.Stack {
-  _vpc;
   _securityGroup;
   public readonly endpoint: string;
   public readonly OSMasterUserSecretName: string;
   public readonly OSHostSecretName: string;
 
-  constructor(scope: Construct, id: string, props: cdk.StackProps) {
+  constructor(scope: Construct, id: string, props: cdk.StackProps & {vpc: ec2.Vpc} & { subnets: cdk.aws_ec2.ISubnet[] }) {
     super(scope, id, props);
 
-    this._vpc = ec2.Vpc.fromLookup(this, "VPC", {
-      isDefault: true,
-    });
-    // Lookup a VPC
-    // this._vpc = props.vpc;
-
     // Create a Security Group for OpenSearch
     this._securityGroup = new ec2.SecurityGroup(this, 'GenBIOpenSearchSG', {
-      vpc: this._vpc,
+      vpc: props.vpc,
       description: 'Allow access to OpenSearch',
       allowAllOutbound: true
     });
@@ -52,14 +45,15 @@ export class AOSStack extends cdk.Stack {
     this._securityGroup.addIngressRule(ec2.Peer.anyIpv4(), ec2.Port.tcp(443), 'Allow HTTPS access');
 
     // Find subnets in different availability zones
-    const subnets = this._vpc.selectSubnets({
-      subnetType: ec2.SubnetType.PUBLIC,
+    const subnets = props.vpc.selectSubnets({
+      subnetType: ec2.SubnetType.PRIVATE_WITH_EGRESS,
     }).subnets;
+    // const subnets = this._vpc.selectSubnets().subnets;
 
     // Create the OpenSearch domain
     const domain = new opensearch.Domain(this, 'GenBiOpenSearchDomain', {
       version: opensearch.EngineVersion.OPENSEARCH_2_9,
-      vpc: this._vpc,
+      vpc: props.vpc,
       securityGroups: [this._securityGroup],
       accessPolicies: [new PolicyStatement({
           effect: Effect.ALLOW,
@@ -68,21 +62,26 @@ export class AOSStack extends cdk.Stack {
           resources: [`arn:${this.partition}:es:${cdk.Aws.REGION}:${cdk.Aws.ACCOUNT_ID}:domain/*`]
       })]
       ,
-      vpcSubnets: [
-        { subnets: [subnets[0]] },
-      ],      
-      // vpcSubnets: SubnetSelection(one_per_az=True, subnet_type=aws_ec2.SubnetType.PUBLIC),
+      vpcSubnets: [ {subnets:  subnets.slice(0, 2)}],
+      // vpcSubnets: [
+      //   { subnets: [subnets[0]] },
+      // ],      
       capacity: {
-        dataNodes: 1,
+        dataNodes: 2,
         dataNodeInstanceType: 'm5.large.search',
         multiAzWithStandbyEnabled: false
       },
+      // capacity: {
+      //   dataNodes: 1,
+      //   dataNodeInstanceType: 'm5.large.search',
+      //   multiAzWithStandbyEnabled: false
+      // },
       ebs: {
         volumeType: ec2.EbsDeviceVolumeType.GP3,
         volumeSize: 20,
       },
       zoneAwareness: {
-        enabled: false
+        availabilityZoneCount: 2
       },
       nodeToNodeEncryption: true,
       encryptionAtRest: {
diff --git a/source/resources/lib/ecs/ecs-stack.ts b/source/resources/lib/ecs/ecs-stack.ts
@@ -9,16 +9,23 @@ import * as ecs_patterns from 'aws-cdk-lib/aws-ecs-patterns';
 import * as path from 'path';
 
 export class ECSStack extends cdk.Stack {
-  _vpc;
   public readonly streamlitEndpoint: string;
   public readonly frontendEndpoint: string;
   public readonly apiEndpoint: string;
-constructor(scope: Construct, id: string, props: cdk.StackProps & { cognitoUserPoolId: string} & { cognitoUserPoolClientId: string} & {OSMasterUserSecretName: string} & {OSHostSecretName: string}) {
+constructor(scope: Construct, id: string, props: cdk.StackProps 
+  & { vpc: ec2.Vpc}
+  & { subnets: cdk.aws_ec2.ISubnet[] } & { cognitoUserPoolId: string} 
+  & { cognitoUserPoolClientId: string} & {OSMasterUserSecretName: string} 
+  & {OSHostSecretName: string}) {
     super(scope, id, props);
-    // Create a VPC
-    this._vpc = ec2.Vpc.fromLookup(this, "VPC", {
-        isDefault: true,
-    });
+
+    // 选择所有的 isolated 和 private with egress 子网
+    // const isolatedSubnets = this._vpc.selectSubnets({ subnetType: ec2.SubnetType.PRIVATE_ISOLATED }).subnets;
+    // const privateSubnets = this._vpc.selectSubnets({ subnetType: ec2.SubnetType.PRIVATE_WITH_EGRESS }).subnets;
+
+    // 合并所有非公共子网
+    // const nonPublicSubnets = [...isolatedSubnets, ...privateSubnets];
+    // const subnets = this._vpc.selectSubnets().subnets;
 
     // Create ECR repositories and Docker image assets
     const services = [
@@ -47,7 +54,7 @@ constructor(scope: Construct, id: string, props: cdk.StackProps & { cognitoUserP
     
     //  Create an ECS cluster
     const cluster = new ecs.Cluster(this, 'GenBiCluster', {
-      vpc: this._vpc,
+      vpc: props.vpc,
     });
 
     const taskExecutionRole = new iam.Role(this, 'TaskExecutionRole', {
@@ -164,8 +171,8 @@ constructor(scope: Construct, id: string, props: cdk.StackProps & { cognitoUserP
       cluster: cluster,
       taskDefinition: taskDefinitionStreamlit,
       publicLoadBalancer: true,
-      // taskSubnets: { subnetType: ec2.SubnetType.PRIVATE_WITH_EGRESS },
-      assignPublicIp: true
+      taskSubnets: { subnets: props.subnets },
+      assignPublicIp: false
     });
 
     // ======= 2. API Service =======
@@ -204,8 +211,8 @@ constructor(scope: Construct, id: string, props: cdk.StackProps & { cognitoUserP
       cluster: cluster,
       taskDefinition: taskDefinitionAPI,
       publicLoadBalancer: true,
-      // taskSubnets: { subnetType: ec2.SubnetType.PUBLIC },
-      assignPublicIp: true
+      taskSubnets: { subnets: props.subnets },
+      assignPublicIp: false
     });
 
     // ======= 3. Frontend Service =======
@@ -251,7 +258,8 @@ constructor(scope: Construct, id: string, props: cdk.StackProps & { cognitoUserP
       taskDefinition: taskDefinitionFrontend,
       publicLoadBalancer: true,
       // taskSubnets: { subnetType: ec2.SubnetType.PUBLIC },
-      assignPublicIp: true
+      taskSubnets: { subnets: props.subnets },
+      assignPublicIp: false
     });
 
     this.streamlitEndpoint = fargateServiceStreamlit.loadBalancer.loadBalancerDnsName;
diff --git a/source/resources/lib/main-stack.ts b/source/resources/lib/main-stack.ts
@@ -1,4 +1,4 @@
-import { Duration, Stack, StackProps, CfnParameter, CfnOutput } from 'aws-cdk-lib';
+import { StackProps, CfnParameter, CfnOutput } from 'aws-cdk-lib';
 import * as cdk from 'aws-cdk-lib';
 import { Construct } from 'constructs';
 import * as ec2 from 'aws-cdk-lib/aws-ec2';
@@ -7,14 +7,21 @@ import { LLMStack } from './model/llm-stack';
 import { ECSStack } from './ecs/ecs-stack';
 import { CognitoStack } from './cognito/cognito-stack';
 import { RDSStack } from './rds/rds-stack';
+import { VPCStack } from './vpc/vpc-stack';
+
+interface MainStackProps extends StackProps {
+  deployRds?: boolean;
+}
 
 export class MainStack extends cdk.Stack {
-  constructor(scope: Construct, id: string, props: StackProps={}) {
+  constructor(scope: Construct, id: string, props: MainStackProps={ deployRds: false }) {
     super(scope, id, props);
 
-    // Looking for the default VPC
-    const vpc = ec2.Vpc.fromLookup(this, "VPC", {
-      isDefault: true,
+    const _deployRds = props.deployRds || false;
+
+    // ======== Step 0. Define the VPC =========
+    const _VpcStack = new VPCStack(this, 'vpc-Stack', {
+      env: props.env,
     });
 
     // ======== Step 1. Define the LLMStack =========
@@ -26,7 +33,9 @@ export class MainStack extends cdk.Stack {
     
     // ======== Step 2. Define the AOSStack ========= 
     const _AosStack = new AOSStack(this, 'aos-Stack', {
-      env: props.env
+      env: props.env,
+      vpc: _VpcStack.vpc,
+      subnets: _VpcStack.vpc.selectSubnets({ subnetType: ec2.SubnetType.PRIVATE_WITH_EGRESS }).subnets,
     });
 
     const aosEndpoint = _AosStack.endpoint;
@@ -45,24 +54,32 @@ export class MainStack extends cdk.Stack {
     // pass the aosEndpoint and aosPassword to the ecs stack
     const _EcsStack = new ECSStack(this, 'ecs-Stack', {
       env: props.env,
+      vpc: _VpcStack.vpc,
+      subnets: _VpcStack.vpc.selectSubnets({ subnetType: ec2.SubnetType.PRIVATE_WITH_EGRESS }).subnets,
       cognitoUserPoolId: _CognitoStack.userPoolId,
       cognitoUserPoolClientId: _CognitoStack.userPoolClientId,
       OSMasterUserSecretName: _AosStack.OSMasterUserSecretName,
       OSHostSecretName: _AosStack.OSHostSecretName,
     });
-
+    _AosStack.addDependency(_VpcStack);
     _EcsStack.addDependency(_AosStack);
     _EcsStack.addDependency(_CognitoStack);
+    _EcsStack.addDependency(_VpcStack);
 
     new cdk.CfnOutput(this, 'AOSDomainEndpoint', {
       value: aosEndpoint,
       description: 'The endpoint of the OpenSearch domain'
     });
-
-    // new cdk.CfnOutput(this, 'RDSEndpoint', {
-    //   value: _RdsStack.endpoint,
-    //   description: 'The endpoint of the RDS instance'
-    // });
+    
+    if (_deployRds) {
+      const _RdsStack = new RDSStack(this, 'rds-Stack', {
+        env: props.env,
+      });
+      new cdk.CfnOutput(this, 'RDSEndpoint', {
+        value: _RdsStack.endpoint,
+        description: 'The endpoint of the RDS instance',
+      });
+    }
     
     new cdk.CfnOutput(this, 'StreamlitEndpoint', {
       value: _EcsStack.streamlitEndpoint,
