AWSTemplateFormatVersion: "2010-09-09"
Description: CloudFormation Templates for provisioning sage maker studio and domain
Parameters:
  S3Bucket:
    Type: String
    Description: "S3 Bucket"
    Default: "test-studio-cf"
  UserProfileName:
    Type: String
    Description: "user profile name "
    Default: "DefaultUser"
  VPCId:
    Type: String
    Description: VPC Id
    Default: "vpc-xxxxxx"
  SubnetIds:
    Type: String
    Description: Subnet Ids
    Default: "subnet-xxxxx"

Resources:
  LambdaExecutionRole:
    Type: "AWS::IAM::Role"
    Properties:
      AssumeRolePolicyDocument:
        Version: 2012-10-17
        Statement:
          - Effect: Allow
            Principal:
              Service:
                - lambda.amazonaws.com
            Action:
              - "sts:AssumeRole"
      Path: /


  LambdaExecutionPolicy:
    Type: AWS::IAM::ManagedPolicy
    Properties:
      Path: /
      PolicyDocument:
        Version: 2012-10-17
        Statement:
          - Sid: CloudWatchLogsPermissions
            Effect: Allow
            Action:
              - logs:CreateLogGroup
              - logs:CreateLogStream
              - logs:PutLogEvents
            Resource: !Sub "arn:${AWS::Partition}:logs:*:*:*"
          - Sid: SageMakerDomainPermission
            Effect: Allow
            Action:
              - sagemaker:CreateDomain
              - sagemaker:DescribeDomain
              - sagemaker:DeleteDomain
              - sagemaker:UpdateDomain
              - sagemaker:CreateUserProfile
              - sagemaker:UpdateUserProfile
              - sagemaker:DeleteUserProfile
              - sagemaker:DescribeUserProfile
            Resource:
              - !Sub "arn:${AWS::Partition}:sagemaker:*:*:domain/*"
              - !Sub "arn:${AWS::Partition}:sagemaker:*:*:user-profile/*"
          - Sid: SageMakerExecPassRole
            Effect: Allow
            Action:
              - iam:PassRole
            Resource: !GetAtt  SageMakerExecutionRole.Arn
      Roles:
        - !Ref  LambdaExecutionRole

  SageMakerExecutionRole:
    Type: "AWS::IAM::Role"
    Properties:
      AssumeRolePolicyDocument:
        Version: 2012-10-17
        Statement:
          - Effect: Allow
            Principal:
              Service:
                - sagemaker.amazonaws.com
            Action:
              - "sts:AssumeRole"
      Path: /
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/AmazonSageMakerFullAccess

  StudioDomainFunction:
    Type: AWS::Lambda::Function
    DependsOn: LambdaExecutionPolicy
    Properties:
      Handler: domain_function.lambda_handler
      Role: !GetAtt LambdaExecutionRole.Arn
      Code:
        S3Bucket: !Ref S3Bucket
        S3Key: domain_function.zip
      Runtime: python3.8
      Timeout: 900
      Layers:
        - !Ref CfnResponseLayer

  CfnResponseLayer:
    Type: AWS::Lambda::LayerVersion
    Properties:
      CompatibleRuntimes:
        - python3.8
      Content:
        S3Bucket: !Ref S3Bucket
        S3Key: cfnResponse_layer.zip
      Description: cfn_response layer
      LayerName: cfn-response

  UserProfileFunction:
    Type: AWS::Lambda::Function
    Properties:
      Handler: UserProfile_function.lambda_handler
      Role: !GetAtt LambdaExecutionRole.Arn
      Code:
        S3Bucket: !Ref S3Bucket
        S3Key: UserProfile_function.zip
      Runtime: python3.8
      Timeout: 900
      Layers:
        - !Ref CfnResponseLayer

  StudioDomain:
    Type: Custom::StudioDomain
    Properties:
      ServiceToken: !GetAtt StudioDomainFunction.Arn
      VPC: !Ref VPCId
      SubnetIds: !Ref SubnetIds
      DomainName: "MyDomainName"
      DefaultUserSettings:
        ExecutionRole: !GetAtt SageMakerExecutionRole.Arn

  UserProfile:
    Type: Custom::UserProfile
    Properties:
      ServiceToken: !GetAtt UserProfileFunction.Arn
      DomainId: !GetAtt StudioDomain.DomainId
      UserProfileName: !Ref UserProfileName
      UserSettings:
        ExecutionRole: !GetAtt SageMakerExecutionRole.Arn