additional-databases/sagemaker/neptune-notebook-cloudformation/neptune-workbench-stack.yaml

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

  NeptuneClusterResourceId:
    Description: The resource ID of the existing Neptune cluster.
    Type: String

  NeptuneClusterPort:
    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>

  NeptuneClusterSubnetId:
    Description: The ID of the subnet in a VPC to which you would like to have connectivity from your ML compute instance.
    Type: AWS::EC2::Subnet::Id

  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
    Default: ''

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

  EncryptionKey:
    Description: 'OPTIONAL: The ARN of a KMS key to encrypt your notebook data with.'
    Type: String
    Default: ''

Conditions:
  UseSageMakerNotebookAutoRole:
    Fn::Equals:
      - !Ref SageMakerNotebookRoleArn
      - ""
  UseEncryptionKey:
    Fn::Not:
      - Fn::Equals:
        - Ref: EncryptionKey
        - ""

Mappings:
  RegionMap:
    cn-north-1:
      "bucket": "aws-neptune-notebook-cn-north-1"
    cn-northwest-1:
      "bucket": "aws-neptune-notebook-cn-northwest-1"
    us-east-2:
      "bucket": "aws-neptune-notebook-us-east-2"
    us-east-1:
      "bucket": "aws-neptune-notebook-us-east-1"
    us-west-1:
      "bucket": "aws-neptune-notebook-us-west-1"
    us-west-2:
      "bucket": "aws-neptune-notebook-us-west-2"
    ap-south-1:
      "bucket": "aws-neptune-notebook-ap-south-1"
    ap-northeast-2:
      "bucket": "aws-neptune-notebook-ap-northeast-2"
    ap-southeast-1:
      "bucket": "aws-neptune-notebook-ap-southeast-1"
    ap-southeast-2:
      "bucket": "aws-neptune-notebook-ap-southeast-2"
    ap-northeast-1:
      "bucket": "aws-neptune-notebook-ap-northeast-1"
    ca-central-1:
      "bucket": "aws-neptune-notebook-ca-central-1"
    eu-central-1:
      "bucket": "aws-neptune-notebook-eu-central-1"
    eu-west-1:
      "bucket": "aws-neptune-notebook-eu-west-1"
    eu-west-2:
      "bucket": "aws-neptune-notebook-eu-west-2"
    eu-west-3:
      "bucket": "aws-neptune-notebook-eu-west-3"
    eu-north-1:
      "bucket": "aws-neptune-notebook-eu-north-1"
    sa-east-1:
      "bucket": "aws-neptune-notebook-sa-east-1"
    af-south-1:
      "bucket": "aws-neptune-notebook-af-south-1"
    ap-east-1:
      "bucket": "aws-neptune-notebook-ap-east-1"
    me-south-1:
      "bucket": "aws-neptune-notebook-me-south-1"
    me-central-1:
      "bucket": "aws-neptune-notebook-me-central-1"
    ap-south-2:
      "bucket": "aws-neptune-notebook-ap-south-2"
    il-central-1:
      "bucket": "aws-neptune-notebook-il-central-1"
    us-gov-east-1:
      "bucket": "aws-neptune-notebook-us-gov-east-1"
    us-gov-west-1:
      "bucket": "aws-neptune-notebook-us-gov-west-1"

Resources:
  NeptuneNotebookInstance:
    Type: AWS::SageMaker::NotebookInstance
    Properties:
      NotebookInstanceName: !Join
                          - ''
                          - - 'aws-neptune-'
                            - !Ref SageMakerNotebookName
      InstanceType:
        Ref: NotebookInstanceType
      SubnetId:
        Ref: NeptuneClusterSubnetId
      SecurityGroupIds:
        Ref: NeptuneClusterSecurityGroups
      RoleArn:
        Fn::If:
          - UseSageMakerNotebookAutoRole
          - Fn::GetAtt:
              - SageMakerNotebookAutoRole
              - Arn
          - Ref: SageMakerNotebookRoleArn
      LifecycleConfigName:
        Fn::GetAtt:
          - NeptuneNotebookInstanceLifecycleConfig
          - NotebookInstanceLifecycleConfigName
      Tags:
        - Key: aws-neptune-cluster-id
          Value: !Select [ 0, !Split [ ".", !Ref NeptuneClusterEndpoint ] ]
        - Key: aws-neptune-resource-id
          Value: !Ref NeptuneClusterResourceId
      KmsKeyId:
        Fn::If:
          - UseEncryptionKey
          - Ref: EncryptionKey
          - Ref: AWS::NoValue

  NeptuneNotebookInstanceLifecycleConfig:
    Type: AWS::SageMaker::NotebookInstanceLifecycleConfig
    Properties:
      OnStart:
        - 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://${S3Bucket}/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
              - S3Bucket:
                  Fn::FindInMap:
                    - RegionMap
                    - Ref: AWS::Region
                    - bucket

  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:::${S3Bucket}"
                    - S3Bucket:
                        Fn::FindInMap:
                          - RegionMap
                          - Ref: AWS::Region
                          - bucket
                  - Fn::Sub:
                    - "arn:${AWS::Partition}:s3:::${S3Bucket}/*"
                    - S3Bucket:
                        Fn::FindInMap:
                          - RegionMap
                          - Ref: AWS::Region
                          - bucket
              - Effect: Allow
                Action: neptune-db:*
                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/*"