sagemaker-domain.yaml

AWSTemplateFormatVersion: '2010-09-09'
Description: >
  This CFN template creates a SageMaker domain, a user profile, a Jupyter App Space

Parameters:
  UserProfileNamePrefix:
    Type: String
    Description: The user profile name prefix for the SageMaker workshop
    Default: 'studio-user-ts'
  DomainNamePrefix:
    Type: String
    Description: The domain name prefix of the Sagemaker domain
    Default: 'timeseries-domain'
  JupyterLabAppInstance:
    Type: String
    Description: Name of an ML instance to run JupyterLab App space
    Default: 'ml.g5.4xlarge'
  OwnerID:
    Type: String
    Description: Owner identification to be written as a resource tag
    Default: 'not-set'
  CreateMLflowServer:
    Type: String
    Description: Create an MLflow trackings server as part of the domain
    Default: 'NO'
    AllowedValues:
      - 'YES'
      - 'NO'

Conditions:
  CreateMLflowServerCondition: !Equals [ !Ref CreateMLflowServer, 'YES' ]
  OwnerIDCondition: !Equals [ !Ref OwnerID, 'not-set' ]

Resources:
  LambdaExecutionRole:
    Type: AWS::IAM::Role
    Properties:
      Policies: 
        - PolicyName: SetupLambdaPermissions
          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:ListDomains
                  - sagemaker:CreateDomain
                  - sagemaker:DescribeDomain
                  - sagemaker:DeleteDomain
                  - sagemaker:UpdateDomain
                  - sagemaker:ListUserProfiles
                  - sagemaker:CreateUserProfile
                  - sagemaker:UpdateUserProfile
                  - sagemaker:DeleteUserProfile
                  - sagemaker:DescribeUserProfile
                  - sagemaker:ListApps
                  - sagemaker:CreateApp
                  - sagemaker:DescribeApp
                  - sagemaker:DeleteApp
                  - sagemaker:UpdateApp
                Resource:
                  - !Sub "arn:${AWS::Partition}:sagemaker:*:*:domain/*"
                  - !Sub "arn:${AWS::Partition}:sagemaker:*:*:user-profile/*"
                  - !Sub "arn:${AWS::Partition}:sagemaker:*:*:app/*"
              - Sid: SageMakerProjectsPermission
                Effect: Allow
                Action:
                  - servicecatalog:AcceptPortfolioShare
                  - sagemaker:EnableSagemakerServicecatalogPortfolio
                  - sagemaker:DisableSagemakerServicecatalogPortfolio
                Resource: '*'
              - Sid: ServiceCatalogPermission
                Effect: Allow
                Action:
                  - servicecatalog:ListAcceptedPortfolioShares
                  - servicecatalog:AssociatePrincipalWithPortfolio
                Resource: '*'
              - Sid: SageMakerExecPassRole
                Effect: Allow
                Action:
                  - iam:PassRole
                Resource: !GetAtt SageMakerExecutionRole.Arn
              - Sid: S3BucketCORSPermission
                Effect: Allow
                Action:
                  - s3:GetBucketCORS
                  - s3:PutBucketCORS
                  - s3:CreateBucket
                  - s3:ListBucket
                  - s3:HeadBucket
                Resource: !Sub "arn:${AWS::Partition}:s3:::${StudioS3Bucket}" 
              - Sid: S3BucketDeletePermission
                Effect: Allow
                Action:
                  - s3:ListBucket
                  - s3:ListBucketVersions
                  - s3:DeleteBucket
                  - s3:DeleteObject
                  - s3:DeleteObjectVersion
                Resource: 
                  - !Sub "arn:${AWS::Partition}:s3:::${StudioS3Bucket}" 
                  - !Sub "arn:${AWS::Partition}:s3:::${StudioS3Bucket}/*" 
      AssumeRolePolicyDocument:
        Version: 2012-10-17
        Statement:
          - Effect: Allow
            Principal:
              Service:
                - lambda.amazonaws.com
            Action:
              - "sts:AssumeRole"
      Path: /
      ManagedPolicyArns:
        - 'arn:aws:iam::aws:policy/AmazonEC2ReadOnlyAccess'
        - 'arn:aws:iam::aws:policy/IAMReadOnlyAccess'
        - 'arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole'

  SageMakerExecutionRole:
    Type: AWS::IAM::Role
    Properties:
      Policies: 
        - PolicyName: logs
          PolicyDocument:
            Version: 2012-10-17
            Statement:
              - 
                Effect: Allow
                Action:
                  - logs:FilterLogEvents
                Resource: !Sub "arn:${AWS::Partition}:logs:*:${AWS::AccountId}:*"
        - PolicyName: iam-access
          PolicyDocument:
            Version: 2012-10-17
            Statement:
              - 
                Effect: Allow
                Action:
                  - iam:GetRole
                  - iam:GetRolePolicy
                Resource: '*'
        - PolicyName: pass-role
          PolicyDocument:
            Version: 2012-10-17
            Statement:
              - 
                Effect: Allow
                Action:
                  - iam:PassRole
                Resource: 'arn:aws:iam::*:role/*'
                Condition:
                  StringLike:
                    iam:PassedToService:
                      - sagemaker.amazonaws.com
                      - events.amazonaws.com
        - PolicyName: kms-access
          PolicyDocument:
            Version: 2012-10-17
            Statement:
              - 
                Effect: Allow
                Action:
                  - kms:CreateKey
                  - kms:Get*
                  - kms:List*
                Resource: '*'
        - PolicyName: list-tags
          PolicyDocument: 
            Version: 2012-10-17
            Statement:
              -
                Effect: Allow
                Action:
                  - sagemaker:ListTags
                Resource: '*'
        - PolicyName: get-service-quotas
          PolicyDocument:
            Version: 2012-10-17
            Statement:
            - 
              Effect: Allow
              Action:
                - servicequotas:GetServiceQuota
              Resource: !Sub 'arn:${AWS::Partition}:servicequotas:${AWS::Region}:${AWS::AccountId}:sagemaker/*'
        - PolicyName: mlflow-permissions
          PolicyDocument:
            Version: 2012-10-17
            Statement:
            -
              Effect: Allow
              Action:
                - sagemaker-mlflow:*
                - sagemaker:CreateMlflowTrackingServer
                - sagemaker:UpdateMlflowTrackingServer
                - sagemaker:DeleteMlflowTrackingServer
                - sagemaker:StartMlflowTrackingServer
                - sagemaker:StopMlflowTrackingServer
                - sagemaker:CreatePresignedMlflowTrackingServerUrl
              Resource: "*"
        - PolicyName: q-developer-access
          PolicyDocument: 
            Version: 2012-10-17
            Statement:
              -
                Effect: Allow
                Action:
                  - q:SendMessage
                Resource: '*'
        - PolicyName: bedrock-models-access
          PolicyDocument:
            Version: 2012-10-17
            Statement:
            - 
              Effect: Allow
              Action:
                - bedrock:GetFoundationModelAvailability
                - bedrock:InvokeModel
                - bedrock:InvokeModelWithResponseStream
              Resource:
                - 'arn:aws:bedrock:*::foundation-model/*'
        - PolicyName: headless-execution-permissions
          PolicyDocument:
            Version: 2012-10-17
            Statement:
            - 
              Effect: Allow
              Action:
                - events:TagResource
                - events:DeleteRule
                - events:PutTargets
                - events:DescribeRule
                - events:PutRule
                - events:RemoveTargets
                - events:DisableRule
                - events:EnableRule
              Resource: '*'
              Condition:
                StringEquals:
                  'aws:ResourceTag/sagemaker:is-scheduling-notebook-job': 'true'
            -
              Effect: Allow
              Action: 
                - s3:CreateBucket
                - s3:PutBucketVersioning
                - s3:PutEncryptionConfiguration
              Resource: 'arn:aws:s3:::sagemaker-automated-execution-*'
            - 
              Effect: Allow
              Action: 
                - sagemaker:AddTags
              Resource: 
                - 'arn:aws:sagemaker:*:*:training-job/*'
                - 'arn:aws:sagemaker:*:*:pipeline/*'
            - 
              Effect: Allow
              Action: 
                - ec2:CreateNetworkInterface
                - ec2:CreateNetworkInterfacePermission
                - ec2:CreateVpcEndpoint
                - ec2:DeleteNetworkInterface
                - ec2:DeleteNetworkInterfacePermission
                - ec2:DescribeDhcpOptions
                - ec2:DescribeNetworkInterfaces
                - ec2:DescribeRouteTables
                - ec2:DescribeSecurityGroups
                - ec2:DescribeSubnets
                - ec2:DescribeVpcEndpoints
                - ec2:DescribeVpcs
                - ecr:BatchCheckLayerAvailability
                - ecr:BatchGetImage
                - ecr:GetDownloadUrlForLayer
                - ecr:GetAuthorizationToken
                - s3:ListBucket
                - s3:GetBucketLocation
                - s3:GetEncryptionConfiguration
                - s3:PutObject
                - s3:DeleteObject
                - s3:GetObject
                - s3:PutObjectTagging
                - sagemaker:DescribeDomain
                - sagemaker:DescribeUserProfile
                - sagemaker:DescribeSpace
                - sagemaker:DescribeStudioLifecycleConfig
                - sagemaker:DescribeImageVersion
                - sagemaker:DescribeAppImageConfig
                - sagemaker:CreateTrainingJob
                - sagemaker:DescribeTrainingJob
                - sagemaker:StopTrainingJob
                - sagemaker:Search
                - sagemaker:CreatePipeline
                - sagemaker:DescribePipeline
                - sagemaker:DeletePipeline
                - sagemaker:StartPipelineExecution
                - glue:GetConnections
              Resource: "*"
      AssumeRolePolicyDocument: 
        Version: 2012-10-17
        Statement:
          - 
            Effect: Allow
            Principal: 
              Service: 
                - sagemaker.amazonaws.com
                - events.amazonaws.com
                - forecast.amazonaws.com
                - lambda.amazonaws.com
                - bedrock.amazonaws.com
                - redshift.amazonaws.com
            Action: 
              - sts:AssumeRole
              
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/AmazonSageMakerFullAccess
        - arn:aws:iam::aws:policy/AWSCloudFormationFullAccess
        - arn:aws:iam::aws:policy/AmazonSageMakerCanvasFullAccess
        - arn:aws:iam::aws:policy/AmazonSageMakerCanvasDataPrepFullAccess
        - arn:aws:iam::aws:policy/service-role/AmazonSageMakerCanvasDirectDeployAccess
        - arn:aws:iam::aws:policy/AmazonSageMakerCanvasAIServicesAccess
        - arn:aws:iam::aws:policy/AmazonSageMakerCanvasBedrockAccess
        - arn:aws:iam::aws:policy/AWSPriceListServiceFullAccess
        - !Sub "arn:${AWS::Partition}:iam::aws:policy/AmazonSageMakerPipelinesIntegrations"
  
  EnableCanvasSettingsLambda:
    Type: AWS::Lambda::Function
    Properties:
      Code:
        ZipFile: |
          # Function: CFEnableSagemakerCanvasSettings
          # Purpose:  Enables Sagemaker Canvas Settings
          import json
          import boto3
          import cfnresponse

          client = boto3.client('sagemaker')

          def lambda_handler(event, context):
              response_status = cfnresponse.SUCCESS
              sagemaker_domain_id = event['ResourceProperties']['SageMakerDomainId']
              sagemaker_execution_role = event['ResourceProperties']['SageMakerExecutionRoleArn']
              canvas_bucket_artifacts = event['ResourceProperties']['CanvasBucketName']

              try:
                if 'RequestType' in event and event['RequestType'] == 'Create':
                    client.update_domain(
                        DomainId=sagemaker_domain_id,
                        DefaultUserSettings={
                          'CanvasAppSettings': {
                            'WorkspaceSettings': {'S3ArtifactPath': f's3://{canvas_bucket_artifacts}/'},
                            'TimeSeriesForecastingSettings': {'Status': 'ENABLED'},
                            'ModelRegisterSettings': {'Status': 'ENABLED'},
                            'DirectDeploySettings': {'Status': 'ENABLED'},
                            'KendraSettings': {'Status': 'DISABLED'}, # Change to ENABLED when you want to use Kendra for RAG
                            'GenerativeAiSettings': {'AmazonBedrockRoleArn':sagemaker_execution_role},
                            # Uncomment and modify the below if you need to add OAuth for Salesforce or Snowflake
                            # 'IdentityProviderOAuthSettings': [
                            #     {
                            #         'DataSourceName': 'SalesforceGenie'|'Snowflake',
                            #         'Status': 'ENABLED'|'DISABLED',
                            #         'SecretArn': 'string'
                            #     },
                            # ],
                          }
                        }
                    )
                cfnresponse.send(event, context, response_status, {}, '')
              except Exception as e:
                cfnresponse.send(event, context, cfnresponse.FAILED, {'Error': str(e)}, context.log_stream_name)

      Description: Enable SageMaker Canvas Settings
      Handler: index.lambda_handler
      MemorySize: 128
      Role: !GetAtt LambdaExecutionRole.Arn
      Runtime: python3.12
      Timeout: 30
      Tags:
        - Key: owner
          Value: !If [ OwnerIDCondition, !Ref 'AWS::StackId', !Ref OwnerID ]

  EnableCanvasSettings:
    Type: Custom::EnableCanvasSettings
    Properties:
      ServiceToken: !GetAtt EnableCanvasSettingsLambda.Arn
      SageMakerDomainId: !GetAtt StudioDomain.DomainId
      SageMakerExecutionRoleArn: !GetAtt SageMakerExecutionRole.Arn
      CanvasBucketName: !Ref StudioS3Bucket

  SetS3BucketCORSPolicyLambda:
    Type: AWS::Lambda::Function
    Properties:
      Role: !GetAtt LambdaExecutionRole.Arn
      Handler: index.handler
      Runtime: python3.12
      Timeout: 30
      MemorySize: 128
      Code:
        ZipFile: |
          import json
          import boto3
          import botocore
          import cfnresponse
          import os

          s3 = boto3.client('s3')

          def handler(event, context):
              bucket_name = event['ResourceProperties']['BucketName']
              cors_configuration = {
                  'CORSRules': [{
                      'AllowedHeaders': ['*'],
                      'AllowedMethods': ['POST', 'PUT', 'GET', 'HEAD', 'DELETE'],
                      'AllowedOrigins': ['https://*.sagemaker.aws'],
                      'ExposeHeaders': [
                          'ETag', 'x-amz-delete-marker', 'x-amz-id-2',
                          'x-amz-request-id', 'x-amz-server-side-encryption',
                          'x-amz-version-id'
                      ]
                  }]
              }

              response_status = cfnresponse.SUCCESS

              if 'RequestType' in event and event['RequestType'] != 'Create':
                cfnresponse.send(event, context, response_status, {}, '')

              try:
                  # Check if the bucket exists
                  s3.head_bucket(Bucket=bucket_name)
                  s3.put_bucket_cors(Bucket=bucket_name, CORSConfiguration=cors_configuration)
              except botocore.exceptions.ClientError as e:
                  # If a 404 error occurs, then the bucket does not exist
                  if e.response['Error']['Code'] == '404':
                      try:
                          # Create the bucket
                          s3.create_bucket(
                              Bucket=bucket_name,
                              CreateBucketConfiguration={'LocationConstraint': os.getenv("AWS_REGION")}
                          )
                          # Set the CORS configuration
                          s3.put_bucket_cors(Bucket=bucket_name, CORSConfiguration=cors_configuration)
                      except botocore.exceptions.ClientError as e:
                          # If an error occurs during bucket creation or CORS setting, return an error
                          print(f"Error creating or setting CORS for bucket {bucket_name}: {e}")
                          cfnresponse.send(event, context, cfnresponse.FAILED, {'Error': str(e)}, bucket_name)
                          return
                  else:
                      # If any other error occurs, return an error
                      print(f"Error accessing bucket {bucket_name}: {e}")
                      cfnresponse.send(event, context, cfnresponse.FAILED, {'Error': str(e)}, bucket_name)
                      return
              except Exception as e:
                  # If any other unexpected error occurs, return an error
                  print(f"Unexpected error: {e}")
                  cfnresponse.send(event, context, cfnresponse.FAILED, {'Error': str(e)}, bucket_name)
                  return

              cfnresponse.send(event, context, cfnresponse.SUCCESS, {'BucketName': bucket_name}, bucket_name)
      Tags:
        - Key: owner
          Value: !If [ OwnerIDCondition, !Ref 'AWS::StackId', !Ref OwnerID ]

  SetS3BucketCORSPolicy:
    Type: Custom::SetS3BucketCORSPolicy
    Properties:
      ServiceToken: !GetAtt SetS3BucketCORSPolicyLambda.Arn
      BucketName: !Ref StudioS3Bucket

  DefaultVpcLambda:
    Type: AWS::Lambda::Function
    Properties:
      Code:
        ZipFile: |
          import json
          import boto3
          import cfnresponse
          import logging
          import traceback

          logger = logging.getLogger()
          logger.setLevel(logging.INFO)
          ec2 = boto3.client('ec2')

          def lambda_handler(event, context):     
            try:         
              if 'RequestType' in event and event['RequestType'] == 'Create':
                  vpc_id = get_default_vpc_id()
                  subnets =  get_subnets_for_vpc(vpc_id)
                  cfnresponse.send(event, context, cfnresponse.SUCCESS, {'VpcId': vpc_id , "Subnets" : subnets}, '')
              else:
                  cfnresponse.send(event, context, cfnresponse.SUCCESS, {},'')
            except:
              logger.exception(f"CFGetDefaultVpcIdTut:failed :{traceback.format_exc()}")
              cfnresponse.send(event, context, cfnresponse.FAILED, {})

          def get_default_vpc_id():
              vpcs = ec2.describe_vpcs(Filters=[{'Name': 'is-default', 'Values': ['true']}])
              vpcs = vpcs['Vpcs']
              vpc_id = vpcs[0]['VpcId']
              return vpc_id


          def get_subnets_for_vpc(vpcId):
              response = ec2.describe_subnets(
                  Filters=[
                      {
                          'Name': 'vpc-id',
                          'Values': [vpcId]
                      }
                  ]
              )
              subnet_ids = []
              for subnet in response['Subnets']:
                  subnet_ids.append(subnet['SubnetId'])
              return subnet_ids 
      Description: Return default VPC ID and Subnets
      Handler: index.lambda_handler
      MemorySize: 128
      Role: !GetAtt LambdaExecutionRole.Arn
      Runtime: python3.10
      Timeout: 60
      Tags:
        - Key: owner
          Value: !If [ OwnerIDCondition, !Ref 'AWS::StackId', !Ref OwnerID ]

  DefaultVpcFinder:
    Type: Custom::ResourceForFindingDefaultVpc
    Properties:
      ServiceToken: !GetAtt DefaultVpcLambda.Arn

  DeleteS3BucketContentLambda:
    Type: AWS::Lambda::Function
    Properties:
      Code:
        ZipFile: |
          import json, boto3
          import cfnresponse
          import logging
          
          logger = logging.getLogger()
          logger.setLevel(logging.INFO)

          def lambda_handler(event, context):
              logger.info("event: {}".format(event))
              try:
                  bucket = event['ResourceProperties']['BucketName']
                  logger.info("bucket: {}, event['RequestType']: {}".format(bucket,event['RequestType']))
                  if event['RequestType'] == 'Delete':
                      s3 = boto3.resource('s3')
                      bucket = s3.Bucket(bucket)
                      for obj in bucket.objects.filter():
                          logger.info("delete obj: {}".format(obj))
                          s3.Object(bucket.name, obj.key).delete()

                  sendResponseCfn(event, context, cfnresponse.SUCCESS)
              except Exception as e:
                  logger.info("Exception: {}".format(e))
                  sendResponseCfn(event, context, cfnresponse.FAILED)

          def sendResponseCfn(event, context, responseStatus):
              responseData = {}
              responseData['Data'] = {}
              cfnresponse.send(event, context, responseStatus, responseData, "CustomResourcePhysicalID") 

      Handler: index.lambda_handler
      Description: "Delete all objects in S3 bucket"
      Role: !GetAtt LambdaExecutionRole.Arn
      Runtime: python3.10
      Timeout: 60
      Tags:
        - Key: owner
          Value: !If [ OwnerIDCondition, !Ref 'AWS::StackId', !Ref OwnerID ]
      
  DeleteS3BucketContentOnDelete:
    Type: Custom::DeleteS3BucketContent
    Properties:
      ServiceToken: !GetAtt DeleteS3BucketContentLambda.Arn
      BucketName: !Ref StudioS3Bucket

  DelayLambda:
    Type: AWS::Lambda::Function
    Properties:
      Handler: index.lambda_handler
      Role: !GetAtt LambdaExecutionRole.Arn
      Runtime: python3.10
      Timeout: 900
      Code:
        ZipFile: |
          import time
          import cfnresponse

          def lambda_handler(event, context):
              try:
                  delay_sec = int(event.get('ResourceProperties', {}).get('DelayInSec', 60))
                  time.sleep(delay_sec)
                  cfnresponse.send(event, context, cfnresponse.SUCCESS, {
                      'Data': f'Delayed for {delay_sec} seconds'
                  }, event.get('PhysicalResourceId', 'DelayCustomResource'))
              except Exception as e:
                  cfnresponse.send(event, context, cfnresponse.FAILED, {
                      'Error': str(e)
                  }, event.get('PhysicalResourceId', 'DelayCustomResource'))
      Tags:
        - Key: owner
          Value: !If [ OwnerIDCondition, !Ref 'AWS::StackId', !Ref OwnerID ]
      
  Delay:
    DependsOn:
      - StudioDomain
    Type: Custom::DelayLambda
    Properties:
      ServiceToken: !GetAtt DelayLambda.Arn
      DelayInSec: 60

  StudioDomain:
    Type: AWS::SageMaker::Domain
    Properties: 
      AppNetworkAccessType: PublicInternetOnly
      AuthMode: IAM
      DomainSettings:
        DockerSettings:
          EnableDockerAccess: ENABLED
      DefaultUserSettings: 
        ExecutionRole: !GetAtt SageMakerExecutionRole.Arn
        StudioWebPortal: ENABLED
        DefaultLandingUri: 'studio::'
      DomainName: !Join ["-", [!Ref DomainNamePrefix, !Select [0, !Split [ "-", !Select [2, !Split ["/", !Ref AWS::StackId]]]]]]
      SubnetIds: !GetAtt DefaultVpcFinder.Subnets
      VpcId: !GetAtt DefaultVpcFinder.VpcId
      Tags: 
        - Key: owner
          Value: !If [ OwnerIDCondition, !Ref 'AWS::StackId', !Ref OwnerID ]

  StudioUserProfile:
    DependsOn:
      - Delay
    Type: AWS::SageMaker::UserProfile
    Properties: 
      DomainId: !GetAtt StudioDomain.DomainId
      UserProfileName: !Join ["-", [!Ref UserProfileNamePrefix, !Select [0, !Split [ "-", !Select [2, !Split ["/", !Ref AWS::StackId]]]]]]
      UserSettings:
        DefaultLandingUri: 'studio::'
        ExecutionRole: !GetAtt SageMakerExecutionRole.Arn
      Tags: 
        - Key: domain-id
          Value: !Ref StudioDomain
        - Key: owner
          Value: !If [ OwnerIDCondition, !Ref 'AWS::StackId', !Ref OwnerID ]

### S3 Bucket similar to the one created by the create domain action in the UI
  StudioS3Bucket:
    Type: AWS::S3::Bucket
    Properties:
      BucketName: !Join
        - "-"
        - - "sagemaker-studio"
          - !Ref AWS::AccountId
          - !Select [0, !Split [ "-", !Select [2, !Split ["/", !Ref AWS::StackId]]]]
      Tags:
        - Key: owner
          Value: !If [ OwnerIDCondition, !Ref 'AWS::StackId', !Ref OwnerID ]

  LifeCycleConfigLambdaRole:
    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 'LifeCycleConfigLambdaPolicy-${AWS::StackName}'
          PolicyDocument:
            Version: 2012-10-17
            Statement:
              - Effect: Allow
                Action:
                  - 'sagemaker:CreateStudioLifecycleConfig'
                  - 'sagemaker:DeleteStudioLifecycleConfig'
                Resource: !Sub 'arn:aws:sagemaker:${AWS::Region}:${AWS::AccountId}:studio-lifecycle-config/*'
              - Effect: Allow
                Action:
                  - 'sagemaker:UpdateUserProfile'
                  - 'sagemaker:DeleteUserProfile'
                Resource: !Sub 'arn:aws:sagemaker:${AWS::Region}:${AWS::AccountId}:user-profile/*'
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole

  LifeCycleConfigLambda:
    DependsOn:
      - StudioUserProfile
    Type: 'AWS::Lambda::Function'
    Properties:
      Description: Add LifeCycle Configuration to copy workshop files to Studio
      Handler: index.lambda_handler
      Role: !GetAtt LifeCycleConfigLambdaRole.Arn
      Runtime: python3.10
      Timeout: 60
      Code:
        ZipFile: !Join
          - |+

          - - 'import boto3'
            - 'import base64'
            - 'import cfnresponse'
            - ''
            - 'client = boto3.client(''sagemaker'')'
            - 'lcc_up1 = ''\n''.join(('
            - '    ''#!/bin/bash'','
            - '    '''','
            - '    ''set -ex'','
            - '    '''','
            - '    ''if [ ! -z "${SM_JOB_DEF_VERSION}" ]'','
            - '    ''then'','
            - '    ''   echo "Running in job mode, skip lcc"'','
            - '    ''else'','
            - '    ''   git clone https://github.com/aws-samples/modern-time-series-forecasting-on-aws.git || { echo "Error: Failed to clone repository"; exit 0; }'','
            - '    ''   echo "Files cloned from GitHub repo"'','
            - '    ''fi'','
            - '    '''','
            - '))'
            - ''
            - !Sub 'lcc_name_up1 = "${AWS::StackName}-clone-repo"'
            - !Sub 'up1 = "${StudioUserProfile}"'
            - ''
            - 'def get_lcc_base64_string(lcc_string):'
            - '    lcc_bytes = lcc_string.encode("ascii")'
            - '    base64_lcc_bytes = base64.b64encode(lcc_bytes)'
            - '    base64_lcc_string = base64_lcc_bytes.decode("ascii")'
            - '    return base64_lcc_string'
            - ''
            - ''
            - 'def apply_lcc_to_user_profile(base64_lcc_string, lcc_config_name, profile):'
            - '    response = client.create_studio_lifecycle_config('
            - '        StudioLifecycleConfigName=lcc_config_name,'
            - '        StudioLifecycleConfigContent=base64_lcc_string,'
            - '        StudioLifecycleConfigAppType="JupyterLab",'
            - '   )'
            - ''
            - '    lcc_arn = response["StudioLifecycleConfigArn"]'
            - '    update_up = client.update_user_profile('
            - '        DomainId=profile.split("|")[1],'
            - '        UserProfileName=profile.split("|")[0],'
            - '        UserSettings={'
            - '            "JupyterLabAppSettings": {'
            - '                "DefaultResourceSpec": {"LifecycleConfigArn": lcc_arn},'
            - '                "LifecycleConfigArns": [lcc_arn]'
            - '            }'
            - '        }'
            - '    )'
            - '    return update_up'
            - ''
            - ''
            - 'def lambda_handler(event, context):'
            - '    print(event)'
            - '    try:'
            - '        base64_lcc_up1_string = get_lcc_base64_string(lcc_up1)'
            - '        updated_up1 = apply_lcc_to_user_profile('
            - '            base64_lcc_up1_string,'
            - '            lcc_name_up1,'
            - '            up1'
            - '        )'
            - '        print("Response User Profile LCC update for UP1")'
            - '        print(updated_up1)'
            - ''
            - '        response_value = 120'
            - '        response_data = {"Data": response_value}'
            - '        cfnresponse.send(event, context, cfnresponse.SUCCESS, response_data)'
            - '    except Exception as e:'
            - '        if "RequestType" in event:'
            - '            if event["RequestType"] == "Delete":'
            - '                try:'
            - '                    response1 = client.delete_studio_lifecycle_config('
            - '                        StudioLifecycleConfigName=lcc_name_up1'
            - '                    )'
            - '                    print(response1)'
            - '                    response_data = {}'
            - '                    cfnresponse.send(event, context, cfnresponse.SUCCESS, response_data)'
            - '                    return'
            - '                except Exception as e2:'
            - '                    print(e2)'
            - '                    response_data = e2'
            - '                    cfnresponse.send(event, context, cfnresponse.SUCCESS, response_data)'
            - '                    return'
            - '        print(e)'
            - '        response_data = {"Data": str(e)}'
            - '        cfnresponse.send(event, context, cfnresponse.FAILED, response_data)'
      Tags:
        - Key: owner
          Value: !If [ OwnerIDCondition, !Ref 'AWS::StackId', !Ref OwnerID ]

  LifeCycleConfigLambdaInvoke:
    Type: AWS::CloudFormation::CustomResource
    Version: "1.0"
    Properties:
      ServiceToken: !GetAtt LifeCycleConfigLambda.Arn

  JupyterAppSpace:
    Type: AWS::SageMaker::Space
    Properties:
      DomainId: !GetAtt StudioDomain.DomainId
      OwnershipSettings:
        OwnerUserProfileName: !Select [2, !Split ["/", !GetAtt StudioUserProfile.UserProfileArn ]]
      SpaceDisplayName: ts-space
      SpaceName: ts-space
      SpaceSettings:
        AppType: JupyterLab
        JupyterLabAppSettings:
          CodeRepositories: 
            - RepositoryUrl: https://github.com/aws-samples/modern-time-series-forecasting-on-aws.git
          DefaultResourceSpec:
            InstanceType: !Ref JupyterLabAppInstance
        SpaceStorageSettings:
          EbsStorageSettings:
            EbsVolumeSizeInGb: 100
      SpaceSharingSettings:
        SharingType: Private

  # MLflow tracking server
  MlFlowServer:
    Type: AWS::SageMaker::MlflowTrackingServer
    Condition: CreateMLflowServerCondition
    Properties:
      ArtifactStoreUri: !Sub 's3://${StudioS3Bucket}/mlflow/${StudioDomain.DomainId}'
      AutomaticModelRegistration: True
      RoleArn: !GetAtt SageMakerExecutionRole.Arn
      TrackingServerName: !Sub 'mlflow-${StudioDomain.DomainId}'
      TrackingServerSize: Small
      Tags:
        - Key: domain-id
          Value: !Ref StudioDomain
        - Key: owner
          Value: !If [ OwnerIDCondition, !Ref 'AWS::StackId', !Ref OwnerID ]