AWSTemplateFormatVersion: '2010-09-09'
Description: CloudFormation template to create a dedicated EFS directory for new
  SageMaker User Profiles in a SageMaker Studio Domain


## ==========================================================================
## =============================== PARAMETERS ===============================
Parameters:

  SageMakerDomainId:
    Description: SageMaker Studio Domain ID
    Type: String

  SageMakerStudioVpc:
    Description: SageMaker Studio VPC ID
    Type: AWS::EC2::VPC::Id

  SageMakerStudioSubnetId:
    Description: Subnet ID associated to SageMaker Studio Domain
    Type: List<AWS::EC2::Subnet::Id>

  SageMakerStudioSecurityGroupId:
    Description: SageMaker Studio Security group ID
    Type: AWS::EC2::SecurityGroup::Id


## ==========================================================================
## ================================ RESOURCES ===============================
Resources:


  ## ============ Lambda function: update Studio inbound SG rule ==============

  # IAM role for Lambda function
  LambdaSecurityGroupRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service:
                - lambda.amazonaws.com
            Action:
              - sts:AssumeRole
      Path: /
      Policies:
        - PolicyName: !Sub lambda-update-security-group
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: Allow
                Action:
                  - logs:CreateLogGroup
                Resource: !Sub arn:aws:logs:${AWS::Region}:${AWS::AccountId}:*
              - Effect: Allow
                Action:
                  - sagemaker:DescribeDomain
                  - elasticfilesystem:CreateMountTarget
                  - elasticfilesystem:DescribeMountTargets
                  - elasticfilesystem:DeleteMountTarget
                Resource: 
                  - !Sub arn:aws:elasticfilesystem:${AWS::Region}:${AWS::AccountId}:file-system/*
                  - !Sub arn:aws:elasticfilesystem:${AWS::Region}:${AWS::AccountId}:access-point/*
                  - !Sub arn:aws:sagemaker:${AWS::Region}:${AWS::AccountId}:domain/*
              - Effect: Allow
                Action:
                  - logs:CreateLogStream
                  - logs:PutLogEvents
                Resource: !Sub arn:aws:logs:${AWS::Region}:${AWS::AccountId}:log-group:/aws/lambda/${AWS::StackName}-lambda-security-group-function:*
              - Effect: Allow
                Action:
                  - ec2:DescribeSecurityGroups
                  - ec2:DescribeSubnets
                  - ec2:AuthorizeSecurityGroupIngress
                  - ec2:RevokeSecurityGroupIngress
                  - ec2:DescribeNetworkInterfaces
                  - ec2:DescribeVpcs

                Resource: '*'
              - Effect: Allow
                Action:
                  - ec2:CreateNetworkInterface
                Resource: 
                  - !Sub arn:aws:ec2:${AWS::Region}:${AWS::AccountId}:security-group/*
                  - !Sub arn:aws:ec2:${AWS::Region}:${AWS::AccountId}:network-interface/*
                  - !Sub arn:aws:ec2:${AWS::Region}:${AWS::AccountId}:subnet/*
              - Effect: Allow
                Action:
                  - ec2:DeleteNetworkInterface
                Resource: 
                  - !Sub arn:aws:ec2:${AWS::Region}:${AWS::AccountId}:*/*
      RoleName: !Sub ${AWS::StackName}-lambda-security-group-role

  # Lambda function
  LambdaUpdateSageMakerSecurityGroup:
    Type: AWS::Lambda::Function
    DependsOn:
      - EFS
      - EFSAccessPoint
    Properties:
      Description: Lambda function to update SageMaker inbound Security Group and create EFS mount targets in Sagemaker Studio Subnets
      FunctionName: !Sub ${AWS::StackName}-lambda-security-group-function
      VpcConfig:
        SecurityGroupIds:
          - !Ref SageMakerStudioSecurityGroupId
        SubnetIds: !Ref SageMakerStudioSubnetId
      Timeout: 900
      Handler: index.lambda_handler
      Role: !GetAtt LambdaSecurityGroupRole.Arn
      Runtime: python3.12
      Environment:
        Variables:
          sagemaker_domain_id: !Ref SageMakerDomainId
          lambda_security_group_id: !Ref SageMakerStudioSecurityGroupId
          efs_file_system: !Ref EFS
      Code:
        ZipFile: |
          import boto3
          import os
          import time
          import cfnresponse
          import logging
          logger = logging.getLogger()
          logger.setLevel(logging.INFO)
          ec2_client = boto3.client("ec2")
          sagemaker_client = boto3.client('sagemaker')
          efs_client = boto3.client('efs')
          def get_sagemaker_security_group_id(sagemaker_domain_id: str) -> str:
              # Get all Security Groups managed by SageMaker
              security_groups_managed_by_sagemaker = ec2_client.describe_security_groups(
                  Filters=[
                      {
                          "Name": "tag-key",
                          "Values": ["ManagedByAmazonSageMakerResource"]
                      }
                  ]
              )
              # Get name for the SageMaker Security Group allowing NFS inbound
              sagemaker_security_group_name = f"security-group-for-inbound-nfs-{sagemaker_domain_id}"
              # Get ID for sg_domain_name
              sagemaker_security_group_id = [sg["GroupId"] for sg in security_groups_managed_by_sagemaker["SecurityGroups"]
                                              if sg["GroupName"]==sagemaker_security_group_name][0]
              return sagemaker_security_group_id
          def describe_domain(domain_id):
              domain = sagemaker_client.describe_domain(
                DomainId=domain_id
              )
              subnets=domain['SubnetIds']
              return subnets
          def mount_efs(subnets, security_group,efs_system):
            for subnet in subnets:
              response = efs_client.create_mount_target(
                FileSystemId=efs_system,
                SubnetId=subnet,
                SecurityGroups=[
                    security_group
                  ]
              )
          def wait_mount_targets_available(list_mount_target):
            time.sleep(120)
            available = False
            for mount_target in list_mount_target["MountTargets"]:
                if mount_target["LifeCycleState"] == "available":
                  available = True
                else:
                  available = False
            return available

          def lambda_handler(event, context):
              request_type = event["RequestType"]
              sagemaker_domain_id = os.environ.get("sagemaker_domain_id")
              lambda_security_group_id = os.environ.get("lambda_security_group_id")
              efs_file_system= os.environ.get("efs_file_system")
              response_data = {}
              try:
                  sagemaker_security_group_id = get_sagemaker_security_group_id(sagemaker_domain_id)
                  sagemaker_subnets_ids=describe_domain(sagemaker_domain_id)
                  
                  if request_type == "Create":
                      # Add Security Group inbound rule to allow Lambda SG
                      ec2_client.authorize_security_group_ingress(
                          GroupId=sagemaker_security_group_id,
                          IpPermissions=[
                              {
                                  'IpProtocol': "tcp",
                                  "FromPort": 0,
                                  "ToPort": 65535,
                                  'UserIdGroupPairs': [{'GroupId': lambda_security_group_id}]
                              }
                          ]
                      )
                      # Create EFS Mount Target
                      mount_efs(sagemaker_subnets_ids, sagemaker_security_group_id,efs_file_system)
                      # Make sure Mount Targets LifeCycleState is available
                      list_mount_target = efs_client.describe_mount_targets(
                          FileSystemId=efs_file_system
                      )
                      mount_targets_available = wait_mount_targets_available(list_mount_target)
                      if not mount_targets_available:
                        logger.info("Mount targets still not available")
                        time.sleep(300)
                  elif request_type == "Delete":
                      # Remove Security Group inbound rule for Lambda SG
                      ec2_client.revoke_security_group_ingress(
                          GroupId=sagemaker_security_group_id,
                          IpPermissions=[
                              {
                                  'IpProtocol': "tcp",
                                  "FromPort": 0,
                                  "ToPort": 65535,
                                  'UserIdGroupPairs': [{'GroupId': lambda_security_group_id}]
                              }
                          ]
                      )

                      list_mount_target = efs_client.describe_mount_targets(
                          FileSystemId=efs_file_system
                      )
                      for mount_target in list_mount_target["MountTargets"]:
                        response = efs_client.delete_mount_target(
                            MountTargetId=mount_target["MountTargetId"]
                        )

                     
                  # Send SUCCESS response back to CFN custom Lambda-backed resource
                  response_data["sagemaker_security_group_id"] = sagemaker_security_group_id

                  cfnresponse.send(
                      event, context, cfnresponse.SUCCESS, response_data, event["LogicalResourceId"]
                  )
              except Exception as error:
                  # Send FAILED response back to CFN custom Lambda-backed resource
                  logger.info(f"Lambda custom resource failed with error: {error}")
                  cfnresponse.send(
                      event, context, cfnresponse.FAILED, response_data, event["LogicalResourceId"]
                  )
  # Lambda-backed custom resource to update NLB Target IP based on events
  CustomLambdaResource:
    Type: Custom::CustomLambda
    Properties:
      ServiceToken: !GetAtt LambdaUpdateSageMakerSecurityGroup.Arn

  ## ============================== EFS volume ================================

  EFS:
    Type: AWS::EFS::FileSystem
    Properties:
      PerformanceMode: generalPurpose
      FileSystemProtection:
        ReplicationOverwriteProtection: ENABLED
      LifecyclePolicies:
        - TransitionToIA: AFTER_30_DAYS
        - TransitionToArchive: AFTER_90_DAYS
      ThroughputMode: elastic
      BackupPolicy:
        Status: DISABLED
      Encrypted: true
      FileSystemTags:
        - Value: EFS-SageMakerStudio
          Key: Name

  EFSAccessPoint:
    Type: AWS::EFS::AccessPoint
    DependsOn: EFS
    Properties:
      FileSystemId: !Ref EFS
      RootDirectory:
        Path: /
      PosixUser:
        Uid: '0'
        Gid: '0'


  ## =========== Lambda function - Create EFS and update Studio users =========

  # Lambda IAM role
  LambdaEFSUserUpdateRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service:
                - lambda.amazonaws.com
            Action:
              - sts:AssumeRole
      Path: /
      Policies:
        - PolicyName: !Sub lambda-manage-efs-studio-users
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: Allow
                Action:
                  - ec2:DescribeNetworkInterfaces
                Resource: '*'
              - Effect: Allow
                Action:
                  - ec2:CreateNetworkInterface
                Resource: 
                  - !Sub arn:aws:ec2:${AWS::Region}:${AWS::AccountId}:security-group/*
                  - !Sub arn:aws:ec2:${AWS::Region}:${AWS::AccountId}:network-interface/*
                  - !Sub arn:aws:ec2:${AWS::Region}:${AWS::AccountId}:subnet/*
              - Effect: Allow
                Action:
                  - ec2:DeleteNetworkInterface
                Resource: 
                  - !Sub arn:aws:ec2:${AWS::Region}:${AWS::AccountId}:*/*
              - Effect: Allow
                Action:
                  - logs:CreateLogGroup
                Resource: !Sub arn:aws:logs:${AWS::Region}:${AWS::AccountId}:*
              - Effect: Allow
                Action:
                  - logs:CreateLogStream
                  - logs:PutLogEvents
                Resource: !Sub arn:aws:logs:${AWS::Region}:${AWS::AccountId}:log-group:/aws/lambda/${AWS::StackName}-lambda-efs-sagemaker-function:*
              - Effect: Allow
                Action:
                  - sagemaker:UpdateUserProfile
                  - sagemaker:UpdateDomain
                Resource: 
                  - !Sub arn:aws:sagemaker:${AWS::Region}:${AWS::AccountId}:user-profile/*
                  - !Sub arn:aws:sagemaker:${AWS::Region}:${AWS::AccountId}:domain/*
              - Effect: Allow
                Action:
                  - elasticfilesystem:ClientMount
                  - elasticfilesystem:ClientRootAccess
                  - elasticfilesystem:ClientWrite
                  - elasticfilesystem:DescribeMountTargets
                Resource: !GetAtt EFS.Arn
      RoleName: !Sub ${AWS::StackName}-lambda-efs-sagemaker-role

  # Lambda Function
  LambdaManageEFSUsers:
    Type: AWS::Lambda::Function
    DependsOn:
      - EFSAccessPoint
      - CustomLambdaResource
    Properties:
      Description: Lambda function to create an EFS directory for new SageMaker Studio users
      FunctionName: !Sub ${AWS::StackName}-lambda-efs-sagemaker-function
      VpcConfig:
        SecurityGroupIds:
          - !Ref SageMakerStudioSecurityGroupId
        SubnetIds: !Ref SageMakerStudioSubnetId
      Timeout: 60
      Handler: index.lambda_handler
      Role: !GetAtt LambdaEFSUserUpdateRole.Arn
      FileSystemConfigs:
        - Arn: !GetAtt EFSAccessPoint.Arn
          LocalMountPath: /mnt/efs
      Runtime: python3.12
      Environment:
        Variables:
          efs_id: !GetAtt EFS.FileSystemId
          domain_id: !Ref SageMakerDomainId
      Code:
        ZipFile: |
          import json
          import subprocess
          import boto3
          import os
          import logging
          logger = logging.getLogger()
          logger.setLevel(logging.INFO)
          sm_client = boto3.client('sagemaker')
          def lambda_handler(event, context):
              print(event)
              file_system=os.environ['efs_id']
              domain_id=os.environ['domain_id']
              event_name=event["detail"]["eventName"]
              event_domain_id=event["detail"]["requestParameters"]["domainId"]
              user=event["detail"]["requestParameters"]["userProfileName"]
              logger.info(f"Received eventName: {event_name}. Updating Studio User {user} for Domain {domain_id} and EFS {file_system}")
              
              if event_domain_id == domain_id:
                  # Permissions
                  repository=f'/mnt/efs/{user}'
                  subprocess.call(['mkdir', repository])
                  subprocess.call(['chown', '200001:1001', repository])
                  
                  #Update Sagemaker user
                  response = sm_client.update_user_profile(
                      DomainId=domain_id,
                      UserProfileName=user,
                      UserSettings={
                          'CustomFileSystemConfigs': [
                              {
                                  'EFSFileSystemConfig': {
                                      'FileSystemId': file_system,
                                      'FileSystemPath': f'/{user}'
                                  }
                              }
                          ]
                      }
                  )
                  logger.info(f"Updated Studio User {user} for Domain {domain_id} and EFS {file_system}")
                  return {
                      'statusCode': 200,
                      'body': json.dumps(f"Created dir and modified permissions for Studio User {user} and Domain {domain_id}")
                  }
              else:
                  return {
                      'statusCode': 400,
                      'body': json.dumps(f"Domain id {event_domain_id} is different than the domain id configured {domain_id}")
                  }

  ## ============================ EventBridge rule ============================

  EventBridgeRuleCreateUserProfile:
    Type: AWS::Events::Rule
    Properties:
      EventPattern:
        detail-type:
          - AWS API Call via CloudTrail
        source:
          - aws.sagemaker
        detail:
          eventSource:
            - sagemaker.amazonaws.com
          eventName:
            - CreateUserProfile
      Description: EventBridge Rule triggered when a new Studio User Profile is
        created. It will invoke a Lambda function that will create an EFS
        directory for such user.
      State: ENABLED
      Targets:
        - Arn: !GetAtt LambdaManageEFSUsers.Arn
          Id: TargetLambda
      Id: eb-rule-CreateUserProfile
      Name: eb-rule-CreateUserProfile

  PermissionForEventsToInvokeLambda:
    Type: AWS::Lambda::Permission
    Properties:
      FunctionName: !GetAtt LambdaManageEFSUsers.Arn
      Action: lambda:InvokeFunction
      SourceArn: !GetAtt EventBridgeRuleCreateUserProfile.Arn
      Principal: events.amazonaws.com


## ==========================================================================
## ================================ OUTPUTS =================================
Outputs:

  EFS:
    Description: Custom EFS for SageMaker Studio Domain
    Value: !Ref EFS

  LambdaFunction:
    Description: Lambda function to create EFS directories for new Studio users
    Value: !Ref LambdaManageEFSUsers