fix: [notebook]: pin the source code in experimental notebook (#158)

* use s3_assets

* update template

* update srcCodeAsset

* rm DockerImage

* update Policy for notebook

* doc update

* add UT

* refine notebook role policy

* refine NotebookIamRole policy

* refine notebookRole

Co-authored-by: Yong Liu <yonmzn@amazon.com>

Author: amliuyong

docs/en/workshop/a-molecular-unfolding/evaluate-your-own-model.md

@@ -18,9 +18,7 @@ This solution is an open-source project under Apache License Version 2.0. You ca
 
 2. Clone the project to your own workspace.
 
-3. Change `githubRepo` path in the source/src/molecular-unfolding/cdk/construct-notebook.ts file to be your Git repository. 
-
-4. Make changes to source code. If you want to make changes for quantum algorithms, you can modify code under this directory **source/src/molecular-unfolding/utility**.
+3. Make changes to source code. If you want to make changes for quantum algorithms, you can modify code under this directory **source/src/molecular-unfolding/utility**.
 
 ## Deploy stack to your AWS account from CDK
 

docs/en/workshop/a-molecular-unfolding/notebook-experiment.md

@@ -41,7 +41,7 @@ Here we focus on the first phase of MD, ligand expansion, which aims to reduce s
 
     Figure 7: Notebook URL on the stack's Output tab
 
-4. Open the file under the directory **source/src/molecular-folding/molecular_unfolding.ipynb** and choose the kernel **qcenv**.
+4. Open the file under the directory **molecular-folding/molecular_unfolding.ipynb** and choose the kernel **qcenv**.
 
 5. Navigate through the notebook, which consists of four Steps:
 

docs/zh/workshop/a-molecular-unfolding/evaluate-your-own-model.md

@@ -20,9 +20,7 @@
 
 2. 将项目克隆到您自己的工作区。
 
-3. 修改文件`source/src/molecular-unfolding/cdk/construct-notebook.ts`中的`githubRepo`路径为您的Git仓库。
-
-4. 修改源代码。如果您想要修改量子算法，可以在目录`source/src/molecular-unfolding/utility`下修改代码。
+3. 修改源代码。如果您想要修改量子算法，可以在目录`source/src/molecular-unfolding/utility`下修改代码。
 
 ### 将堆栈从 CDK 部署到您的 AWS 账户
 

docs/zh/workshop/a-molecular-unfolding/notebook-experiment.md

@@ -43,7 +43,7 @@
     ![deployment output](../../images/deploy_output_notebook.png)
     图 7: 笔记本链接
 
-4. 打开笔记本文件**source/src/molecular-folding/molecular_unfolding.ipynbn**，并确保笔记本的核为**qcenv**。
+4. 打开笔记本文件**molecular-folding/molecular_unfolding.ipynb**，并确保笔记本的核为**qcenv**。
 
 5. 浏览整个笔记本。它由四个步骤组成：
 

source/package-lock.json

@@ -24,6 +24,7 @@
   "devDependencies": {
     "@types/aws-lambda": "^8.10.92",
     "@types/jest": "^27.0.3",
+    "@types/mustache": "^4.1.3",
     "@types/node": "14.17.7",
     "@typescript-eslint/eslint-plugin": "^4.29.0",
     "@typescript-eslint/parser": "^4.29.0",
@@ -42,10 +43,11 @@
   "dependencies": {
     "@aws-cdk/aws-batch-alpha": "^2.20.0-alpha.0",
     "@aws-sdk/client-cloudformation": "^3.48.0",
-    "aws-cdk-lib": "^2.20.0",
     "aws-cdk": "^2.20.0",
+    "aws-cdk-lib": "^2.20.0",
     "cdk-bootstrapless-synthesizer": "^2.2.0",
     "constructs": "^10.0.0",
+    "mustache": "^4.2.0",
     "source-map-support": "^0.5.16"
   },
   "bundledDependencies": [],

source/package.json

@@ -19,10 +19,14 @@ import {
   readFileSync,
 } from 'fs';
 
+import * as path from 'path';
+
 import {
   aws_s3 as s3,
   aws_kms as kms,
   aws_ec2 as ec2,
+  aws_s3_assets as s3_assets,
+  Fn,
 } from 'aws-cdk-lib';
 
 import {
@@ -35,6 +39,7 @@ import {
   Construct,
 } from 'constructs';
 
+import * as Mustache from 'mustache';
 import {
   RoleUtil,
 } from './utils/utils-role';
@@ -62,60 +67,46 @@ export class Notebook extends Construct {
     this.props = props;
     const INSTANCE_TYPE = 'ml.c5.xlarge';
 
-    this.roleUtil = RoleUtil.newInstance(this, props);
-
-    const githubRepo = 'https://github.com/awslabs/quantum-computing-exploration-for-drug-discovery-on-aws';
+    const srcCodeAsset = new s3_assets.Asset(this, 'srcCodeAsset', {
+      path: path.join(__dirname, '../../'),
+    });
 
-    const defaultCodeRepository = this.node.tryGetContext('default_code_repository') || githubRepo;
+    this.roleUtil = RoleUtil.newInstance(this, props);
 
     const notebookRole = this.roleUtil.createNotebookIamRole();
 
+    srcCodeAsset.grantRead(notebookRole);
+
     let onStartContent = readFileSync(`${__dirname}/resources/onStart.template`, 'utf-8');
 
-    const base64Encode = (str: string): string => Buffer.from(str, 'binary').toString('base64');
-    const onStartContentBase64 = base64Encode(onStartContent);
+    const rawOnStartContent = Mustache.render(onStartContent, {
+      s3_code_path: srcCodeAsset.s3ObjectUrl,
+    });
 
     const installBraketSdk = new CfnNotebookInstanceLifecycleConfig(this, 'install-braket-sdk', {
       onStart: [{
-        content: onStartContentBase64,
+        content: Fn.base64(ec2.UserData.custom(rawOnStartContent).render()),
       }],
     });
 
     const qcNotebookKey = new kms.Key(this, 'qcNotebookKey', {
       enableKeyRotation: true,
     });
 
-    let notebookInstance = null;
-    if (defaultCodeRepository.startsWith('https://')) {
-      notebookInstance = new CfnNotebookInstance(this, 'Notebook', {
-        instanceType: INSTANCE_TYPE,
-        roleArn: notebookRole.roleArn,
-        rootAccess: 'Enabled', // Lifecycle configurations need root access to be able to set up a notebook instance
-        lifecycleConfigName: installBraketSdk.attrNotebookInstanceLifecycleConfigName,
-        volumeSizeInGb: 50,
-        kmsKeyId: qcNotebookKey.keyId,
-        securityGroupIds: [this.props.notebookSg.securityGroupId],
-        subnetId: this.props.vpc.selectSubnets({
-          subnetType: ec2.SubnetType.PRIVATE_WITH_NAT,
-        }).subnetIds[0],
-        directInternetAccess: 'Disabled',
-        defaultCodeRepository: defaultCodeRepository,
-      });
-    } else {
-      notebookInstance = new CfnNotebookInstance(this, 'Notebook', {
-        instanceType: INSTANCE_TYPE,
-        roleArn: notebookRole.roleArn,
-        rootAccess: 'Enabled',
-        lifecycleConfigName: installBraketSdk.attrNotebookInstanceLifecycleConfigName,
-        volumeSizeInGb: 50,
-        kmsKeyId: qcNotebookKey.keyId,
-        securityGroupIds: [this.props.notebookSg.securityGroupId],
-        subnetId: this.props.vpc.selectSubnets({
-          subnetType: ec2.SubnetType.PRIVATE_WITH_NAT,
-        }).subnetIds[0],
-        directInternetAccess: 'Disabled',
-      });
-    }
+    const notebookInstance = new CfnNotebookInstance(this, 'Notebook', {
+      instanceType: INSTANCE_TYPE,
+      roleArn: notebookRole.roleArn,
+      rootAccess: 'Enabled',
+      lifecycleConfigName: installBraketSdk.attrNotebookInstanceLifecycleConfigName,
+      volumeSizeInGb: 50,
+      kmsKeyId: qcNotebookKey.keyId,
+      securityGroupIds: [this.props.notebookSg.securityGroupId],
+      subnetId: this.props.vpc.selectSubnets({
+        subnetType: ec2.SubnetType.PRIVATE_WITH_NAT,
+      }).subnetIds[0],
+      directInternetAccess: 'Disabled',
+    });
+
     this.notebookUrl = `https://console.aws.amazon.com/sagemaker/home?region=${this.props.region}#/notebook-instances/openNotebook/${notebookInstance.attrNotebookInstanceName}?view=classic`;
   }
 

source/src/molecular-unfolding/cdk/construct-notebook.ts

@@ -24,6 +24,12 @@ pip3 install jupyter ipykernel
 
 python3 -m ipykernel install --user --name=qcenv
 
+aws s3 cp {{&s3_code_path}} /home/ec2-user/SageMaker/code.zip
+
+cd /home/ec2-user/SageMaker/
+unzip code.zip -d ./
+rm code.zip
+
 EOS
 
 exit 0

source/src/molecular-unfolding/cdk/resources/onStart.template

@@ -23,8 +23,10 @@ import {
 
 import {
   Template,
+  Match,
 } from 'aws-cdk-lib/assertions';
 
+
 import {
   NotebookNestStack,
 } from '../src/molecular-unfolding/cdk/stack-notebook';
@@ -64,4 +66,28 @@ describe('Notebook', () => {
     template.resourceCountIs('AWS::SageMaker::NotebookInstanceLifecycleConfig', 1);
   });
 
+  test('NotebookInstanceLifecycleConfig is configured correctly', ()=>{
+    const template = initializeNestStackTemplate();
+    template.hasResourceProperties('AWS::SageMaker::NotebookInstanceLifecycleConfig', {
+      OnStart: [
+        {
+          Content: {
+            'Fn::Base64': {
+              'Fn::Join': [
+                '',
+                [
+                  Match.anyValue(),
+                  {
+                    'Fn::Sub': Match.anyValue(),
+                  },
+                  Match.anyValue(),
+                ],
+              ],
+            },
+          },
+        },
+      ],
+    });
+  });
+
 });