Modified to enable deployment of private network options with SageMaker Code Editor (#1269)

Co-authored-by: Taichiro Suzuki <taichirs@amazon.co.jp>

Author: tbrand

packages/cdk/lib/closed-network-stack.ts

@@ -14,6 +14,7 @@ import {
 
 export interface ClosedNetworkStackProps extends StackProps {
   readonly params: ProcessedStackInput;
+  readonly isSageMakerStudio: boolean;
 }
 
 export class ClosedNetworkStack extends Stack {
@@ -70,6 +71,7 @@ export class ClosedNetworkStack extends Stack {
       subnetIds: closedNetworkSubnetIds,
       hostedZone: closedVpc?.hostedZone,
       certificateArn: closedNetworkCertificateArn,
+      isSageMakerStudio: props.isSageMakerStudio,
     });
 
     const cognitoPrivateProxy = new CognitoPrivateProxy(

packages/cdk/lib/construct/closedNetwork/closed-web.ts

@@ -16,15 +16,16 @@ import {
   OperatingSystemFamily,
 } from 'aws-cdk-lib/aws-ecs';
 import { ApplicationLoadBalancedFargateService } from 'aws-cdk-lib/aws-ecs-patterns';
-import { Platform } from 'aws-cdk-lib/aws-ecr-assets';
+import { Platform, NetworkMode } from 'aws-cdk-lib/aws-ecr-assets';
 import { ApplicationLoadBalancer } from 'aws-cdk-lib/aws-elasticloadbalancingv2';
 
 export interface ClosedWebProps {
-  vpc: IVpc;
-  subnetIds?: string[] | null;
+  readonly vpc: IVpc;
+  readonly subnetIds?: string[] | null;
+  readonly isSageMakerStudio: boolean;
   // For HTTPS listener
-  hostedZone?: PrivateHostedZone;
-  certificateArn?: string | null;
+  readonly hostedZone?: PrivateHostedZone;
+  readonly certificateArn?: string | null;
 }
 
 export class ClosedWeb extends Construct {
@@ -77,6 +78,9 @@ export class ClosedWeb extends Construct {
       taskImageOptions: {
         image: ContainerImage.fromAsset('./fargate-s3-server', {
           platform: Platform.LINUX_AMD64,
+          networkMode: props.isSageMakerStudio
+            ? NetworkMode.custom('sagemaker')
+            : NetworkMode.DEFAULT,
         }),
         containerPort: 8080,
         environment: {

packages/cdk/lib/create-stacks.ts

@@ -86,6 +86,9 @@ export const createStacks = (app: cdk.App, params: ProcessedStackInput) => {
     inferenceProfileStacks
   );
 
+  // GenU Stack
+  const isSageMakerStudio = 'SAGEMAKER_APP_TYPE_LOWERCASE' in process.env;
+
   let closedNetworkStack: ClosedNetworkStack | undefined = undefined;
 
   if (params.closedNetworkMode) {
@@ -98,6 +101,7 @@ export const createStacks = (app: cdk.App, params: ProcessedStackInput) => {
           region: params.region,
         },
         params,
+        isSageMakerStudio,
       }
     );
   }
@@ -206,8 +210,6 @@ export const createStacks = (app: cdk.App, params: ProcessedStackInput) => {
     videoBucketRegionMap[region] = videoTmpBucketStack.bucketName;
   }
 
-  // GenU Stack
-  const isSageMakerStudio = 'SAGEMAKER_APP_TYPE_LOWERCASE' in process.env;
   const generativeAiUseCasesStack = new GenerativeAiUseCasesStack(
     app,
     `GenerativeAiUseCasesStack${updatedParams.env}`,

packages/cdk/test/__snapshots__/generative-ai-use-cases.test.ts.snap

@@ -1164,7 +1164,7 @@ exports[`GenerativeAiUseCases matches the snapshot (closed network mode) 1`] = `
             ],
             "Essential": true,
             "Image": {
-              "Fn::Sub": "123456890123.dkr.ecr.us-east-1.\${AWS::URLSuffix}/cdk-hnb659fds-container-assets-123456890123-us-east-1:7d1a9cbcd523fa1f30afeb1394c611b18778bbdc6343bc4d9ed3f9109ce669e9",
+              "Fn::Sub": "123456890123.dkr.ecr.us-east-1.\${AWS::URLSuffix}/cdk-hnb659fds-container-assets-123456890123-us-east-1:f42c9250931c270e4dbe045bdf1556d0f4363901fa75fc81605cf456a556646a",
             },
             "LogConfiguration": {
               "LogDriver": "awslogs",