Description: |
  CloudFormation template for creating SageMaker Notebook instance

Parameters:
  ProjectName:
    Type: String
    AllowedPattern: '[A-Za-z0-9\-]*'
    Description: Please specify your project name.  Used as a suffix for project resource names.

  EnvType:
    Description: System Environment
    Type: String
    Default: dev

  NotebookOwnerUsername:
    Description: Notebook owner username
    Type: String

  NotebookOwnerEmail:
    Description: Notebook owner email address
    Type: String

  InstanceType:
    Description: The instance type for the SageMaker notebook
    Type: String
    Default: ml.t3.medium
    AllowedValues:
      - ml.t3.medium
      - ml.t3.2xlarge
      - ml.m5.xlarge
      - ml.c5.2xlarge
      - ml.p3.2xlarge

  VolumeSize:
    Description: EBS Volume size in GB
    Type: Number
    Default: '10'

  NotebookSubnet:
    Description: The subnet ID for running the notebook instance
    Type: Number
    Default: 1
    AllowedValues:
      - 1
      - 2
      - 3

Outputs:
  NotebookUrl:
    Description: Link to open Notebook 
    Value: !Sub 'https://${AWS::Region}.console.aws.amazon.com/sagemaker/home?region=${AWS::Region}#/notebook-instances/openNotebook/${NotebookInstance.NotebookInstanceName}?view=lab'
  EbsKmsKeyId:
    Description: EBS KMS Key Id for the notebook encryption
    Value: !Ref SagemakerEbsKMS

Resources:
  SageMakerExecRole:
    Type: 'AWS::IAM::Role'
    Properties:
      AssumeRolePolicyDocument:
        Version: 2012-10-17
        Statement:
          - Effect: Allow
            Principal:
              Service:
                - sagemaker.amazonaws.com
            Action:
              - 'sts:AssumeRole'
      Path: /service-role/
      RoleName: !Sub "ds-notebook-role-${ProjectName}-${EnvType}-${NotebookOwnerUsername}"
      ManagedPolicyArns:
        - !Ref SageMakerExecPolicy
        - 'arn:aws:iam::aws:policy/AmazonS3ReadOnlyAccess'
      Tags:
        - Key: ProjectName
          Value: !Ref ProjectName
        - Key: EnvironmentType
          Value: !Ref EnvType

  SageMakerExecPolicy:
    Type: 'AWS::IAM::ManagedPolicy'
    Properties:
      ManagedPolicyName: !Sub "ds-notebook-policy-${ProjectName}-${EnvType}-${NotebookOwnerUsername}"
      PolicyDocument:
        Version: 2012-10-17
        Statement:
          - Action:
              - 'ssm:GetParameters'
              - 'ssm:GetParameter'
            Resource: !Sub 'arn:aws:ssm:${AWS::Region}:${AWS::AccountId}:parameter/ds-*'
            Effect: Allow
          - Action:
              - 'sagemaker:*'
            Resource: '*'
            Effect: Allow
          - Effect: Allow
            Action:
              - 'application-autoscaling:DeleteScalingPolicy'
              - 'application-autoscaling:DeleteScheduledAction'
              - 'application-autoscaling:DeregisterScalableTarget'
              - 'application-autoscaling:DescribeScalableTargets'
              - 'application-autoscaling:DescribeScalingActivities'
              - 'application-autoscaling:DescribeScalingPolicies'
              - 'application-autoscaling:DescribeScheduledActions'
              - 'application-autoscaling:PutScalingPolicy'
              - 'application-autoscaling:PutScheduledAction'
              - 'application-autoscaling:RegisterScalableTarget'
              - 'cloudwatch:DeleteAlarms'
              - 'cloudwatch:DescribeAlarms'
              - 'cloudwatch:GetMetricData'
              - 'cloudwatch:GetMetricStatistics'
              - 'cloudwatch:ListMetrics'
              - 'cloudwatch:PutMetricAlarm'
              - 'cloudwatch:PutMetricData'
              - 'ec2:CreateNetworkInterface'
              - 'ec2:CreateNetworkInterfacePermission'
              - 'ec2:DeleteNetworkInterface'
              - 'ec2:DeleteNetworkInterfacePermission'
              - 'ec2:DescribeDhcpOptions'
              - 'ec2:DescribeNetworkInterfaces'
              - 'ec2:DescribeRouteTables'
              - 'ec2:DescribeSecurityGroups'
              - 'ec2:DescribeSubnets'
              - 'ec2:DescribeVpcEndpoints'
              - 'ec2:DescribeVpcs'
              - 'ecr:BatchCheckLayerAvailability'
              - 'ecr:BatchGetImage'
              - 'ecr:CreateRepository'
              - 'ecr:GetAuthorizationToken'
              - 'ecr:GetDownloadUrlForLayer'
              - 'ecr:Describe*'
              - 'elastic-inference:Connect'
              - 'iam:ListRoles'
              - 'kms:CreateGrant'
              - 'kms:Decrypt'
              - 'kms:DescribeKey'
              - 'kms:Encrypt'
              - 'kms:GenerateDataKey'
              - 'kms:ListAliases'
              - 'lambda:ListFunctions'
              - 'logs:CreateLogGroup'
              - 'logs:CreateLogStream'
              - 'logs:DescribeLogStreams'
              - 'logs:GetLogEvents'
              - 'logs:PutLogEvents'
              - 'sns:ListTopics'
              - 'codecommit:BatchGetRepositories'
              - 'codecommit:GitPull'
              - 'codecommit:GitPush'
              - 'codecommit:CreateBranch'
              - 'codecommit:DeleteBranch'
              - 'codecommit:GetBranch'
              - 'codecommit:ListBranches'
              - 'codecommit:CreatePullRequest'
              - 'codecommit:GetPullRequest'
              - 'codecommit:CreateCommit'
              - 'codecommit:GetCommit'
              - 'codecommit:GetCommitHistory'
              - 'codecommit:GetDifferences'
              - 'codecommit:GetReferences'
              - 'codecommit:CreateRepository'
              - 'codecommit:GetRepository'
              - 'codecommit:ListRepositories'
            Resource: '*'
          - Effect: Allow
            Action:
              - 'ecr:SetRepositoryPolicy'
              - 'ecr:CompleteLayerUpload'
              - 'ecr:BatchDeleteImage'
              - 'ecr:UploadLayerPart'
              - 'ecr:DeleteRepositoryPolicy'
              - 'ecr:InitiateLayerUpload'
              - 'ecr:DeleteRepository'
              - 'ecr:PutImage'
            Resource: 'arn:aws:ecr:*:*:repository/*sagemaker*'
          - Effect: Allow
            Action:
              - 's3:GetObject'
              - 's3:PutObject'
              - 's3:DeleteObject'
            Resource:
              - !Sub 'arn:aws:s3:::ds-data-bucket-${ProjectName}-${EnvType}-*'
              - !Sub 'arn:aws:s3:::ds-data-bucket-${ProjectName}-${EnvType}-*/*'
              - !Sub 'arn:aws:s3:::ds-model-bucket-${ProjectName}-${EnvType}-*'
              - !Sub 'arn:aws:s3:::ds-model-bucket-${ProjectName}-${EnvType}-*/*'
          - Effect: Allow
            Action:
              - 's3:GetBucketLocation'
              - 's3:ListBucket'
              - 's3:ListAllMyBuckets'
            Resource: '*'
          - Effect: Allow
            Action:
              - 's3:GetObject'
            Resource: '*'
            Condition:
              StringEqualsIgnoreCase:
                's3:ExistingObjectTag/SageMaker': 'true'
          - Effect: Allow
            Action:
              - 'lambda:InvokeFunction'
            Resource:
              - 'arn:aws:lambda:*:*:function:*SageMaker*'
              - 'arn:aws:lambda:*:*:function:*sagemaker*'
              - 'arn:aws:lambda:*:*:function:*Sagemaker*'
              - 'arn:aws:lambda:*:*:function:*LabelingFunction*'
          - Action: 'iam:CreateServiceLinkedRole'
            Effect: Allow
            Resource: >-
              arn:aws:iam::*:role/aws-service-role/sagemaker.application-autoscaling.amazonaws.com/AWSServiceRoleForApplicationAutoScaling_SageMakerEndpoint
            Condition:
              StringLike:
                'iam:AWSServiceName': sagemaker.application-autoscaling.amazonaws.com
          - Effect: Allow
            Action:
              - 'sns:Subscribe'
              - 'sns:CreateTopic'
            Resource:
              - 'arn:aws:sns:*:*:*SageMaker*'
              - 'arn:aws:sns:*:*:*Sagemaker*'
              - 'arn:aws:sns:*:*:*sagemaker*'
          - Effect: Allow
            Action:
              - 'iam:PassRole'
            Resource: '*'
            Condition:
              StringEquals:
                'iam:PassedToService':
                  - sagemaker.amazonaws.com

  SagemakerEbsKMS:
    Type: 'AWS::KMS::Key'
    Properties:
      Description: Generated KMS Key for sagemaker Notebook's EBS encryption
      EnableKeyRotation: true
      Enabled: true
      KeyPolicy:
        Version: 2012-10-17
        Id: KmsKey-EbsKmsSagemakerKey
        Statement:
          - Action:
              - 'kms:DeleteAlias'
              - 'kms:DescribeKey'
              - 'kms:EnableKey'
              - 'kms:GetKeyPolicy'
              - 'kms:UpdateAlias'
              - 'kms:CreateAlias'
              - 'kms:GetKeyPolicy'
              - 'kms:CreateGrant'
              - 'kms:DisableKey'
              - 'kms:Revoke*'
              - 'kms:Disable*'
              - 'kms:CancelKeyDeletion'
              - 'kms:ScheduleKeyDeletion'
              - 'kms:PutKeyPolicy'
              - 'kms:RevokeGrant'
              - 'kms:TagResource'
              - 'kms:UnTagResource'
              - 'kms:EnableKeyRotation'
              - 'kms:ListResourceTags'
            Effect: Allow
            Principal:
              AWS: !Sub "arn:aws:iam::${AWS::AccountId}:root"
            Resource:
              - '*'
      Tags:
        - Key: ProjectName
          Value: !Ref ProjectName
        - Key: EnvironmentType
          Value: !Ref EnvType

  NotebookInstance:
    Type: 'AWS::SageMaker::NotebookInstance'
    Properties:
      DirectInternetAccess: Disabled
      InstanceType: !Ref InstanceType
      KmsKeyId: !Ref SagemakerEbsKMS
      NotebookInstanceName: !Sub "ds-notebook-${ProjectName}-${EnvType}-${NotebookOwnerUsername}"
      DefaultCodeRepository: 
        Fn::ImportValue:
          !Sub "ds-source-${ProjectName}-${EnvType}-url"
      LifecycleConfigName: 
        Fn::ImportValue:
          !Sub 'ds-notebook-lc-${ProjectName}-${EnvType}'
      RoleArn: !GetAtt SageMakerExecRole.Arn
      SecurityGroupIds:
        - Fn::ImportValue:
            !Sub 'ds-sagemaker-sg-${ProjectName}-${EnvType}'
      SubnetId: 
        "Fn::ImportValue":
          !Sub 'ds-subnet${NotebookSubnet}-${ProjectName}-${EnvType}'
      VolumeSizeInGB: !Ref VolumeSize
      Tags:
        - Key: NotebookOwnerUsername
          Value: !Ref NotebookOwnerUsername
        - Key: NotebookOwnerEmail
          Value: !Ref NotebookOwnerEmail
        - Key: ProjectName
          Value: !Ref ProjectName
        - Key: EnvironmentType
          Value: !Ref EnvType