---
AWSTemplateFormatVersion: "2010-09-09"
Description: This stack deploys a Sagemaker Notebook to use with Amazon DocumentDB.

Parameters:

  DocDBSubnet:
    Type: AWS::EC2::Subnet::Id
    Description: Subnet of VPC with access to DocumentDB cluster
    
  DocDBSecurityGroup:
    Type: AWS::EC2::SecurityGroup::Id
    Description: Security Group with access to DocumentDB cluster

  SageMakerInstanceType:
    Type: String
    Description: Instance type for SageMaker
    Default: ml.m5.xlarge
    AllowedValues:
      - 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
      - ml.m5.12xlarge
      - ml.m5.24xlarge
      - ml.m5.2xlarge
      - ml.m5.4xlarge
      - ml.m5.xlarge
      - ml.t3.2xlarge
      - ml.t3.large
      - ml.t3.medium
      - ml.t3.xlarge

Metadata:
  AWS::CloudFormation::Interface: 
    ParameterGroups: 
      - 
        Label: 
          default: "DocumentDB Configuration"
        Parameters: 
          - DocDBSubnet
          - DocDBSecurityGroup
      -
        Label: 
          default: "SageMaker Configuration"
        Parameters:
          - SageMakerInstanceType

Resources:

# SageMaker
  SageMakerRole:
    Type: AWS::IAM::Role
    Properties: 
      AssumeRolePolicyDocument: 
        Version: "2012-10-17"
        Statement: 
          - 
            Effect: "Allow"
            Principal: 
              Service: 
                - "sagemaker.amazonaws.com"
            Action: 
              - "sts:AssumeRole"
      Description: SageMaker notebook instance role granting read access to Secrets Manager
      RoleName: !Sub ${AWS::StackName}-SageMakerRoleName
      Tags: 
        - Key: Name
          Value: !Sub ${AWS::StackName}-SageMakerRoleName

  SageMakerNotebookInstance:
    Type: AWS::SageMaker::NotebookInstance
    Properties: 
      InstanceType: !Ref SageMakerInstanceType
      NotebookInstanceName: !Sub ${AWS::StackName}-SageMakerInstance
      RoleArn: !GetAtt SageMakerRole.Arn
      DirectInternetAccess: Enabled
      LifecycleConfigName: !GetAtt SageMakerNotebookInstanceLifecycleConfig.NotebookInstanceLifecycleConfigName
      SecurityGroupIds: 
        - !Ref DocDBSecurityGroup
      SubnetId: !Ref DocDBSubnet
      Tags:
        - Key: Name
          Value: !Sub ${AWS::StackName}-SageMakerInstance
      VolumeSizeInGB: 5

  SageMakerNotebookInstanceLifecycleConfig:
    Type: "AWS::SageMaker::NotebookInstanceLifecycleConfig"
    Properties:
      OnStart:
        - Content:
            Fn::Base64: 
              !Sub |
                #!/bin/bash
                sudo -u ec2-user -i <<'EOF'
                source /home/ec2-user/anaconda3/bin/activate python3
                pip install --upgrade pymongo
                source /home/ec2-user/anaconda3/bin/deactivate
                cd /home/ec2-user/SageMaker
                wget https://s3.amazonaws.com/rds-downloads/rds-combined-ca-bundle.pem
                EOF

Outputs:
  StackName:
    Value: !Sub ${AWS::StackName}
  SageMakerRole:
    Value: !Ref SageMakerRole
  SageMakerNotebookInstance:
    Value: !Ref SageMakerNotebookInstance
  SageMakerNotebookInstanceLifecycleConfig:
    Value: !Ref SageMakerNotebookInstanceLifecycleConfig