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

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

  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=
                  - "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"
                  - "IDLE_TIME=3600"
                  - "\n"
                  - 'echo "Fetching the autostop script"'
                  - "\n"
                  - "wget https://raw.githubusercontent.com/aws-samples/amazon-sagemaker-notebook-instance-lifecycle-config-samples/master/scripts/auto-stop-idle/autostop.py"
                  - "\n"
                  - 'echo "Starting the SageMaker autostop script in cron"'
                  - "\n"
                  - '(crontab -l 2>/dev/null; echo "*/5 * * * * /usr/bin/python $PWD/autostop.py --time $IDLE_TIME --ignore-connections") | crontab -'
                  - "\n"
                  - EOF
Outputs:
  NotebookName:
    Description: The name of the sagemaker notebook
    Value: !GetAtt "NeptuneNotebookInstance.NotebookInstanceName"