{
  "Description": "A template to deploy Neptune Notebooks using CloudFormation resources",
  "Parameters": {
    "NotebookInstanceType": {
      "Description": "SageMaker Notebook instance type. Please refer to https://aws.amazon.com/sagemaker/pricing/ for up to date allowed instance types and pricing per AWS region.",
      "Type": "String",
      "Default": "ml.t3.medium",
      "AllowedValues": [
        "ml.t2.medium",
        "ml.t2.large",
        "ml.t2.xlarge",
        "ml.t2.2xlarge",
        "ml.t3.medium",
        "ml.t3.large",
        "ml.t3.xlarge",
        "ml.t3.2xlarge",
        "ml.m4.xlarge",
        "ml.m4.2xlarge",
        "ml.m4.4xlarge",
        "ml.m4.10xlarge",
        "ml.m4.16xlarge",
        "ml.m5.xlarge",
        "ml.m5.2xlarge",
        "ml.m5.4xlarge",
        "ml.m5.12xlarge",
        "ml.m5.24xlarge",
        "ml.c4.large",
        "ml.c4.xlarge",
        "ml.c4.2xlarge",
        "ml.c4.4xlarge",
        "ml.c4.8xlarge",
        "ml.c5.xlarge",
        "ml.c5.2xlarge",
        "ml.c5.4xlarge",
        "ml.c5.9xlarge",
        "ml.c5.18xlarge",
        "ml.c5d.xlarge",
        "ml.c5d.2xlarge",
        "ml.c5d.4xlarge",
        "ml.c5d.9xlarge",
        "ml.c5d.18xlarge",
        "ml.p3.2xlarge",
        "ml.p3.8xlarge",
        "ml.p3.16xlarge",
        "ml.p2.2xlarge",
        "ml.p2.8xlarge",
        "ml.p2.16xlarge"
      ],
      "ConstraintDescription": "Must be a valid SageMaker instance type."
    },
    "NotebookName": {
      "Description": "Name for the notebook instance. Ensure that this is unique.",
      "MaxLength": 38,
      "Type": "String"
    },
    "GraphEndpoint": {
      "Description": "Endpoint for Analytics graph.",
      "Type": "String"
    },
    "GraphPort": {
      "Description": "Port to access your Analytics Graph.",
      "Type": "String",
      "Default": "8182"
    },
    "GraphVPC": {
      "Description": "OPTIONAL: VPC Id of the VPC where your Analytics Graph is running. Leave empty for public endpoints.",
      "Type": "String"
    },
    "GraphSubnet": {
      "Description": "OPTIONAL: Id of a VPC subnet where your Analytics Graph is running. Leave empty for public endpoints.",
      "Type": "String"
    },
    "GraphSecurityGroup": {
      "Description": "OPTIONAL: Id of a VPC security group with access to your Analytics Graph. Leave empty for public endpoints.",
      "Type": "String"
    }
  },
  "Conditions": {
    "AddAnalyticsGraphVpc": {
      "Fn::Not": [
        {
          "Fn::Equals": [
            {
              "Ref": "GraphVPC"
            },
            ""
          ]
        }
      ]
    },
    "AddAnalyticsGraphSubnet": {
      "Fn::Not": [
        {
          "Fn::Equals": [
            {
              "Ref": "GraphSubnet"
            },
            ""
          ]
        }
      ]
    },
    "AddAnalyticsGraphSecurityGroup": {
      "Fn::Not": [
        {
          "Fn::Equals": [
            {
              "Ref": "GraphSecurityGroup"
            },
            ""
          ]
        }
      ]
    },
    "AddNetworkOptions": {
      "Fn::And": [
        {
          "Condition": "AddAnalyticsGraphVpc"
        },
        {
          "Condition": "AddAnalyticsGraphSubnet"
        }
      ]
    },
    "CreateSagemakerSecurityGroup": {
      "Fn::And": [
        {
          "Condition": "AddNetworkOptions"
        },
        {
          "Fn::Not": [
            {
              "Condition": "AddAnalyticsGraphSecurityGroup"
            }
          ]
        }
      ]
    },
    "IsIadRegion": {
      "Fn::Equals": [
        {
          "Ref": "AWS::Region"
        },
        "us-east-1"
      ]
    }
  },
  "Resources": {
    "SageMakerSecurityGroup": {
      "Type": "AWS::EC2::SecurityGroup",
      "Condition": "CreateSagemakerSecurityGroup",
      "Properties": {
        "GroupDescription": "Allow Access",
        "VpcId": {
          "Ref": "GraphVPC"
        },
        "Tags": [
          {
            "Key": "Name",
            "Value": {
              "Fn::Sub": "Neptune-Analytics"
            }
          },
          {
            "Key": "StackId",
            "Value": {
              "Fn::Sub": "${AWS::StackId}"
            }
          },
          {
            "Key": "Stack",
            "Value": {
              "Fn::Sub": "${AWS::Region}-${AWS::StackName}"
            }
          },
          {
            "Key": "Application",
            "Value": "NeptuneCloudformation"
          }
        ]
      }
    },
    "NeptuneAnalyticsNotebookInstance": {
      "Type": "AWS::SageMaker::NotebookInstance",
      "Properties": {
        "InstanceType": {
          "Ref": "NotebookInstanceType"
        },
        "PlatformIdentifier" : "notebook-al2-v3",
        "NotebookInstanceName": {
          "Fn::Join": [
            "",
            [
              "aws-neptune-analytics-",
              {
                "Ref": "NotebookName"
              }
            ]
          ]
        },
        "SubnetId": {
          "Fn::If": [
            "AddNetworkOptions",
            {
              "Ref": "GraphSubnet"
            },
            {
              "Ref": "AWS::NoValue"
            }
          ]
        },
        "SecurityGroupIds": {
          "Fn::If": [
            "AddNetworkOptions",
            {
              "Fn::If": [
                "AddAnalyticsGraphSecurityGroup",
                [
                  {
                    "Ref": "GraphSecurityGroup"
                  }
                ],
                [
                  {
                    "Fn::GetAtt": [
                      "SageMakerSecurityGroup",
                      "GroupId"
                    ]
                  }
                ]
              ]
            },
            {
              "Ref": "AWS::NoValue"
            }
          ]
        },
        "RoleArn": {
          "Fn::GetAtt": [
            "ExecutionRole",
            "Arn"
          ]
        },
        "LifecycleConfigName": {
          "Fn::GetAtt": [
            "NeptuneAnalyticsNotebookInstanceLifecycleConfig",
            "NotebookInstanceLifecycleConfigName"
          ]
        },
        "Tags": [
          {
            "Key": "StackId",
            "Value": {
              "Fn::Sub": "${AWS::StackId}"
            }
          },
          {
            "Key": "Stack",
            "Value": {
              "Fn::Sub": "${AWS::Region}-${AWS::StackName}"
            }
          },
          {
            "Key": "Application",
            "Value": "NeptuneCloudformation"
          },
          {
            "Key": "aws-neptune-analytics-graph-endpoint",
            "Value": {
              "Ref": "GraphEndpoint"
            }
          }
        ]
      }
    },
    "NeptuneAnalyticsNotebookInstanceLifecycleConfig": {
      "Type": "AWS::SageMaker::NotebookInstanceLifecycleConfig",
      "Properties": {
        "OnStart": [
          {
            "Content": {
              "Fn::Base64": {
                "Fn::Join": [
                  "",
                  [
                    "#!/bin/bash\n\n",
                    "sudo -u ec2-user -i <<'EOF'\n\n",
                    "echo \"export GRAPH_NOTEBOOK_AUTH_MODE=IAM\" >> ~/.bashrc\n",
                    "echo \"export GRAPH_NOTEBOOK_SSL=True\" >> ~/.bashrc\n",
                    "echo \"export GRAPH_NOTEBOOK_SERVICE=neptune-graph\" >> ~/.bashrc\n",
                    "echo \"export GRAPH_NOTEBOOK_HOST=",
                    {
                      "Ref": "GraphEndpoint"
                    },
                    "\" >> ~/.bashrc\n",
                    "echo \"export GRAPH_NOTEBOOK_PORT=",
                    {
                      "Ref": "GraphPort"
                    },
                    "\" >> ~/.bashrc\n",
                    "echo \"export NEPTUNE_LOAD_FROM_S3_ROLE_ARN=\" >> ~/.bashrc\n",
                    "echo \"export AWS_REGION=",
                    {
                      "Ref": "AWS::Region"
                    },
                    "\" >> ~/.bashrc\n\n",
                    "aws s3 cp s3://",
                    {
                      "Fn::If": [
                        "IsIadRegion",
                        "aws-neptune-notebook",
                        {
                          "Fn::Sub": "aws-neptune-notebook-${AWS::Region}"
                        }
                      ]
                    },
                    "/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_jl4x.sh\n\n",
                    "EOF"
                  ]
                ]
              }
            }
          }
        ]
      }
    },
    "ExecutionRole": {
      "Type": "AWS::IAM::Role",
      "Properties": {
        "AssumeRolePolicyDocument": {
          "Version": "2012-10-17",
          "Statement": [
            {
              "Effect": "Allow",
              "Principal": {
                "Service": [
                  "sagemaker.amazonaws.com"
                ]
              },
              "Action": [
                "sts:AssumeRole"
              ]
            }
          ]
        },
        "Path": "/",
        "Policies": [
          {
            "PolicyName": "SagemakerNotebookNeptuneAnalyticsPolicy",
            "PolicyDocument": {
              "Version": "2012-10-17",
              "Statement": [
                {
                  "Effect": "Allow",
                  "Action": [
                    "s3:GetObject",
                    "s3:ListBucket"
                  ],
                  "Resource": {
                    "Fn::Sub": "arn:${AWS::Partition}:s3:::*"
                  }
                },
                {
                  "Effect": "Allow",
                  "Action": [
                    "logs:CreateLogGroup",
                    "logs:CreateLogStream",
                    "logs:PutLogEvents"
                  ],
                  "Resource": {
                    "Fn::Sub": "arn:${AWS::Partition}:logs:${AWS::Region}:${AWS::AccountId}:log-group:/aws/sagemaker/*"
                  }
                },
                {
                  "Effect": "Allow",
                  "Action": "neptune-graph:*",
                  "Resource": {
                    "Fn::Sub": "arn:aws:neptune-graph:${AWS::Region}:${AWS::AccountId}:*/*"
                  }
                },
                {
                  "Effect": "Allow",
                  "Action": "sagemaker:DescribeNotebookInstance",
                  "Resource": {
                    "Fn::Sub": "arn:aws:sagemaker:${AWS::Region}:${AWS::AccountId}:notebook-instance/*"
                  }
                }
              ]
            }
          }
        ]
      }
    }
  },
  "Outputs": {
    "NeptuneAnalyticsNotebookInstanceId": {
      "Value": {
        "Ref": "NeptuneAnalyticsNotebookInstance"
      }
    },
    "NeptuneAnalyticsSagemakerNotebook": {
      "Value": {
        "Fn::Join": [
          "",
          [
            "https://",
            {
              "Fn::Select": [
                1,
                {
                  "Fn::Split": [
                    "/",
                    {
                      "Ref": "NeptuneAnalyticsNotebookInstance"
                    }
                  ]
                }
              ]
            },
            ".notebook.",
            {
              "Ref": "AWS::Region"
            },
            ".sagemaker.aws/"
          ]
        ]
      }
    },
    "NeptuneAnalyticsNotebookInstanceLifecycleConfigId": {
      "Value": {
        "Ref": "NeptuneAnalyticsNotebookInstanceLifecycleConfig"
      }
    }
  }
}