AWSTemplateFormatVersion: '2010-09-09'

Description: A template to deploy Neptune Notebooks using CloudFormation resources.

Parameters:
  NotebookInstanceType:
    Description: The notebook instance type.
    Type: String
    Default: ml.t2.medium
    AllowedValues:
    - ml.t2.medium
    - ml.t2.large
    - ml.t2.xlarge
    - ml.t2.2xlarge
    - ml.t3.2xlarge
    - ml.t3.large
    - ml.t3.medium
    - ml.t3.xlarge
    - ml.m4.xlarge
    - ml.m4.2xlarge
    - ml.m4.4xlarge
    - ml.m4.10xlarge
    - ml.m4.16xlarge
    - ml.m5.12xlarge
    - ml.m5.24xlarge
    - ml.m5.2xlarge
    - ml.m5.4xlarge
    - ml.m5.xlarge
    - ml.p2.16xlarge
    - ml.p2.8xlarge
    - ml.p2.xlarge
    - ml.p3.16xlarge
    - ml.p3.2xlarge
    - ml.p3.8xlarge
    - ml.c4.2xlarge
    - ml.c4.4xlarge
    - ml.c4.8xlarge
    - ml.c4.xlarge
    - ml.c5.18xlarge
    - ml.c5.2xlarge
    - ml.c5.4xlarge
    - ml.c5.9xlarge
    - ml.c5.xlarge
    - ml.c5d.18xlarge
    - ml.c5d.2xlarge
    - ml.c5d.4xlarge
    - ml.c5d.9xlarge
    - ml.c5d.xlarge
    ConstraintDescription: Must be a valid SageMaker instance type.

  NeptuneClusterEndpoint:
    Description: The cluster endpoint of an existing Neptune cluster.
    Type: String

  NeptuneClusterPort:
    Description: 'OPTIONAL: The Port of an existing Neptune cluster (default 8182).'
    Type: String
    Default: '8182'

  NeptuneClusterAuthMode:
    Description: 'The IAM authentication mode on the specified Neptune cluster.'
    Type: String
    Default: DEFAULT
    AllowedValues:
      - DEFAULT
      - IAM

  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>

  NeptuneClusterSubnetId:
    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.
    Type: String
    AllowedPattern: ^arn:aws[a-z\-]*:iam::\d{12}:role/?[a-zA-Z_0-9+=,.@\-_/]+$

  SageMakerNotebookName:
    Description: The name of the Neptune notebook.
    Type: String

Resources:
  NeptuneNotebookInstance:
    Type: AWS::SageMaker::NotebookInstance
    Properties:
      NotebookInstanceName: !Join
                          - ''
                          - - 'aws-neptune-'
                            - !Ref SageMakerNotebookName
      InstanceType:
        Ref: NotebookInstanceType
      SubnetId:
        Ref: NeptuneClusterSubnetId
      SecurityGroupIds:
        Ref: NeptuneClusterSecurityGroups
      RoleArn:
        Ref: SageMakerNotebookRole
      LifecycleConfigName:
        Fn::GetAtt:
        - NeptuneNotebookInstanceLifecycleConfig
        - NotebookInstanceLifecycleConfigName

  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=
              - !Ref NeptuneClusterAuthMode
              - "' >> ~/.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