Improve CloudFormation template for Neptune DB Notebook (#571)

* Improve CloudFormation template for Neptune DB Notebook * update changelog * update authmode param
aws · Mar 20, 2024 · f835c0d · f835c0d
1 parent b6cd163
commit f835c0d
Show file tree

Hide file tree

Showing 3 changed files with 112 additions and 49 deletions.
diff --git a/ChangeLog.md b/ChangeLog.md
@@ -4,6 +4,7 @@ Starting with v1.31.6, this file will contain a record of major features and upd
 
 ## Upcoming
 - Updated `create-graph` CLI commands in Neptune Analytics samples ([Link to PR](https://github.com/aws/graph-notebook/pull/565))
+- Updated NeptuneDB Notebook CloudFormation template ([Link to PR](https://github.com/aws/graph-notebook/pull/571))
 - Added `@neptune_graph_only` magics decorator ([Link to PR](https://github.com/aws/graph-notebook/pull/569))
 - Added `%graph_pg_info` line magic ([Link to PR](https://github.com/aws/graph-notebook/pull/570))
 

diff --git a/additional-databases/sagemaker/neptune-notebook-cloudformation/README.md b/additional-databases/sagemaker/neptune-notebook-cloudformation/README.md
@@ -1,10 +1,10 @@
 ## Launching graph-notebook as Amazon Neptune Workbench via AWS CloudFormation
 
-The AWS CloudFormation template in this folder, [`neptune-workbench-stack.yaml`](neptune-workbench-stack.yaml), deploys Amazon Neptune workbench notebooks as resources, and includes the base 'Getting Started' notebooks. The workbench lets you work with your Amazon Neptune cluster using Jupyter notebooks hosted by Amazon SageMaker. You are billed for workbench resources through Amazon SageMaker, separately from your Neptune billing.
+The AWS CloudFormation template in this folder, [`neptune-workbench-stack.yaml`](neptune-workbench-stack.yaml), deploys Amazon Neptune workbench notebooks as resources, and includes the base 'Getting Started' notebooks. The workbench lets you work with your Amazon Neptune Database cluster using Jupyter notebooks hosted by Amazon SageMaker. You are billed for workbench resources through Amazon SageMaker, separately from your Neptune billing.
 
 ### Parameter details
 #### Minimum permissions for the SageMakerNotebookRole
-This is the ARN for the AWS IAM role that the notebook instance will assume. Make sure that this role has at least the following minimum permissions within its service role policy:
+You may opt to have your notebook instance assume an existing AWS IAM role, via the `SageMakerNotebookRoleArn` stack parameter. Make sure that this role has at least the following minimum permissions within its service role policy:
 
 ```json
 {
@@ -17,21 +17,36 @@ This is the ARN for the AWS IAM role that the notebook instance will assume. Mak
         "s3:ListBucket"
       ],
       "Resource": [
-        "arn:aws:s3:::aws-neptune-notebook",
-        "arn:aws:s3:::aws-neptune-notebook/*"
+        "arn:(AWS Partition):s3:::aws-neptune-notebook-(AWS Region)",
+        "arn:(AWS Partition):s3:::aws-neptune-notebook-(AWS Region)/*"
       ]
     },
     {
       "Effect": "Allow",
       "Action": "neptune-db:connect",
       "Resource": [
-        "your-cluster-arn/*"
+        "arn:(AWS Partition):neptune-db:(AWS Region):(AWS Account ID):(Cluster Resource ID)/*"
       ]
     }
   ]
 }
 ```
 
+If you would like to enable CloudWatch logging, also add:
+```json
+    {
+      "Effect": "Allow",
+      "Action": [
+        "logs:CreateLogGroup",
+        "logs:CreateLogStream",
+        "logs:PutLogEvents"
+      ],
+      "Resource": [
+        "arn:(AWS Partition):logs:(AWS Region):(AWS Account ID):log-group:/aws/sagemaker/*"
+      ]
+    }
+```
+
 The role should also establish the following trust relationship:
 
 ```json
@@ -48,11 +63,3 @@ The role should also establish the following trust relationship:
   ]
 }
 ```
-
-#### How to populate the 'Cluster' value within the AWS Console for Amazon Neptune Notebooks
-Add the following tags manually to the notebook instance.
-
-| Key | Value |
-| ------------- |-------------|
-| **aws-neptune-cluster-id** | Amazon Neptune database cluster ID (found under *DB cluster id* under *Configuration* of the selected cluster in the AWS console) |
-| **aws-neptune-resource-id** | Amazon Neptune cluster resource ID (found under *Resource id* under *Configuration* of the selected cluster in the AWS console) |
diff --git a/additional-databases/sagemaker/neptune-notebook-cloudformation/neptune-workbench-stack.yaml b/additional-databases/sagemaker/neptune-notebook-cloudformation/neptune-workbench-stack.yaml
@@ -52,11 +52,28 @@ Parameters:
     Description: The cluster endpoint of an existing Neptune cluster.
     Type: String
 
+  NeptuneClusterResourceId:
+    Description: The resource ID of the existing Neptune cluster.
+    Type: String
+
   NeptuneClusterPort:
-    Description: 'OPTIONAL: The Port of an existing Neptune cluster (default 8182).'
+    Description: 'OPTIONAL: The Port of the existing Neptune cluster (default 8182).'
     Type: String
     Default: '8182'
 
+  NeptuneClusterAuthMode:
+    Description: The IAM authentication setting on the existing Neptune cluster.
+    Type: String
+    Default: 'DEFAULT'
+    AllowedValues:
+      - 'IAM'
+      - 'DEFAULT'
+
+  NeptuneClusterLoadFromS3Arn:
+    Description: 'OPTIONAL: The ARN of the S3 bucket to load data from.'
+    Type: String
+    Default: ''
+
   NeptuneClusterSecurityGroups:
     Description: The VPC security group IDs. The security groups must be for the same VPC as specified in the subnet.
     Type: List<AWS::EC2::SecurityGroup::Id>
@@ -65,15 +82,21 @@ Parameters:
     Description: The ID of the subnet in a VPC to which you would like to have a connectivity from your ML compute instance.
     Type: AWS::EC2::Subnet::Id
 
-  SageMakerNotebookRole:
-    Description: The ARN for the IAM role that the notebook instance will assume.
+  SageMakerNotebookRoleArn:
+    Description: 'OPTIONAL: The ARN for the IAM role that the notebook instance will assume. If not provided, a role will be automatically created, based on the cluster information.'
     Type: String
-    AllowedPattern: ^arn:aws[a-z\-]*:iam::\d{12}:role/?[a-zA-Z_0-9+=,.@\-_/]+$
+    Default: ''
 
   SageMakerNotebookName:
     Description: The name of the Neptune notebook.
     Type: String
 
+Conditions:
+  UseSageMakerNotebookAutoRole:
+    Fn::Equals:
+      - !Ref SageMakerNotebookRoleArn
+      - ""
+
 Resources:
   NeptuneNotebookInstance:
     Type: AWS::SageMaker::NotebookInstance
@@ -89,42 +112,74 @@ Resources:
       SecurityGroupIds:
         Ref: NeptuneClusterSecurityGroups
       RoleArn:
-        Ref: SageMakerNotebookRole
+        Fn::If:
+          - UseSageMakerNotebookAutoRole
+          - Fn::GetAtt:
+              - SageMakerNotebookAutoRole
+              - Arn
+          - Ref: SageMakerNotebookRoleArn
       LifecycleConfigName:
         Fn::GetAtt:
-        - NeptuneNotebookInstanceLifecycleConfig
-        - NotebookInstanceLifecycleConfigName
+          - NeptuneNotebookInstanceLifecycleConfig
+          - NotebookInstanceLifecycleConfigName
+      Tags:
+        - Key: aws-neptune-cluster-id
+          Value: !Select [ 0, !Split [ ".", !Ref NeptuneClusterEndpoint ] ]
+        - Key: aws-neptune-resource-id
+          Value: !Ref NeptuneClusterResourceId
 
   NeptuneNotebookInstanceLifecycleConfig:
     Type: AWS::SageMaker::NotebookInstanceLifecycleConfig
     Properties:
       OnStart:
-      - Content:
-          Fn::Base64:
-            Fn::Join:
-            - ''
-            - - "#!/bin/bash\n"
-              - sudo -u ec2-user -i << 'EOF'
-              - "\n"
-              - echo 'export GRAPH_NOTEBOOK_AUTH_MODE=
-              - "DEFAULT' >> ~/.bashrc\n"
-              - echo 'export GRAPH_NOTEBOOK_HOST=
-              - !Ref NeptuneClusterEndpoint
-              - "' >> ~/.bashrc\n"
-              - echo 'export GRAPH_NOTEBOOK_PORT=
-              - !Ref NeptuneClusterPort
-              - "' >> ~/.bashrc\n"
-              - echo 'export NEPTUNE_LOAD_FROM_S3_ROLE_ARN=
-              - "' >> ~/.bashrc\n"
-              - echo 'export AWS_REGION=
-              - !Ref AWS::Region
-              - "' >> ~/.bashrc\n"
-              - aws s3 cp s3://aws-neptune-notebook/graph_notebook.tar.gz /tmp/graph_notebook.tar.gz
-              - "\n"
-              - rm -rf /tmp/graph_notebook
-              - "\n"
-              - tar -zxvf /tmp/graph_notebook.tar.gz -C /tmp
-              - "\n"
-              - /tmp/graph_notebook/install.sh
-              - "\n"
-              - EOF
+        - Content:
+            Fn::Base64: !Sub |
+              #!/bin/bash
+              sudo -u ec2-user -i << 'EOF'
+              echo 'export GRAPH_NOTEBOOK_AUTH_MODE=${NeptuneClusterAuthMode}' >> ~/.bashrc
+              echo 'export GRAPH_NOTEBOOK_SSL=True' >> ~/.bashrc
+              echo 'export GRAPH_NOTEBOOK_HOST=${NeptuneClusterEndpoint}' >> ~/.bashrc
+              echo 'export GRAPH_NOTEBOOK_PORT=${NeptuneClusterPort}' >> ~/.bashrc
+              echo "export GRAPH_NOTEBOOK_SERVICE=neptune-db" >> ~/.bashrc
+              echo 'export NEPTUNE_LOAD_FROM_S3_ROLE_ARN=${NeptuneClusterLoadFromS3Arn}' >> ~/.bashrc
+              echo 'export AWS_REGION=${AWS::Region}' >> ~/.bashrc
+              aws s3 cp s3://aws-neptune-notebook-${AWS::Region}/graph_notebook.tar.gz /tmp/graph_notebook.tar.gz
+              rm -rf /tmp/graph_notebook
+              tar -zxvf /tmp/graph_notebook.tar.gz -C /tmp
+              /tmp/graph_notebook/install.sh
+              EOF
+
+  SageMakerNotebookAutoRole:
+    Type: AWS::IAM::Role
+    Condition: UseSageMakerNotebookAutoRole
+    Properties:
+      AssumeRolePolicyDocument:
+        Version: "2012-10-17"
+        Statement:
+          - Effect: Allow
+            Principal:
+              Service: sagemaker.amazonaws.com
+            Action: sts:AssumeRole
+      Policies:
+        - PolicyName: SageMakerNotebookPolicy
+          PolicyDocument:
+            Version: "2012-10-17"
+            Statement:
+              - Effect: Allow
+                Action:
+                  - s3:GetObject
+                  - s3:ListBucket
+                Resource:
+                  - Fn::Sub: "arn:${AWS::Partition}:s3:::aws-neptune-notebook-${AWS::Region}"
+                  - Fn::Sub: "arn:${AWS::Partition}:s3:::aws-neptune-notebook-${AWS::Region}/*"
+              - Effect: Allow
+                Action: neptune-db:connect
+                Resource:
+                  - Fn::Sub: "arn:${AWS::Partition}:neptune-db:${AWS::Region}:${AWS::AccountId}:${NeptuneClusterResourceId}/*"
+              - Effect: Allow
+                Action:
+                  - logs:CreateLogGroup
+                  - logs:CreateLogStream
+                  - logs:PutLogEvents
+                Resource:
+                  - Fn::Sub: "arn:${AWS::Partition}:logs:${AWS::Region}:${AWS::AccountId}:log-group:/aws/sagemaker/*"