template.yml

AWSTemplateFormatVersion: 2010-09-09
Description: SageMaker Studio Module
Mappings:
  RegionMap:
    us-east-1:
      datascience: 'arn:aws:sagemaker:us-east-1:081325390199:image/datascience-1.0'
      datawrangler: >-
        arn:aws:sagemaker:us-east-1:663277389841:image/sagemaker-data-wrangler-1.0
    us-east-2:
      datascience: 'arn:aws:sagemaker:us-east-2:429704687514:image/datascience-1.0'
      datawrangler: >-
        arn:aws:sagemaker:us-east-2:415577184552:image/sagemaker-data-wrangler-1.0
    us-west-1:
      datascience: 'arn:aws:sagemaker:us-west-1:742091327244:image/datascience-1.0'
      datawrangler: >-
        arn:aws:sagemaker:us-west-1:926135532090:image/sagemaker-data-wrangler-1.0
    us-west-2:
      datascience: 'arn:aws:sagemaker:us-west-2:236514542706:image/datascience-1.0'
      datawrangler: >-
        arn:aws:sagemaker:us-west-2:174368400705:image/sagemaker-data-wrangler-1.0
    af-south-1:
      datascience: 'arn:aws:sagemaker:af-south-1:559312083959:image/datascience-1.0'
      datawrangler: >-
        arn:aws:sagemaker:af-south-1:143210264188:image/sagemaker-data-wrangler-1.0
    ap-east-1:
      datascience: 'arn:aws:sagemaker:ap-east-1:493642496378:image/datascience-1.0'
      datawrangler: >-
        arn:aws:sagemaker:ap-east-1:707077482487:image/sagemaker-data-wrangler-1.0
    ap-south-1:
      datascience: 'arn:aws:sagemaker:ap-south-1:394103062818:image/datascience-1.0'
      datawrangler: >-
        arn:aws:sagemaker:ap-south-1:089933028263:image/sagemaker-data-wrangler-1.0
    ap-northeast-2:
      datascience: 'arn:aws:sagemaker:ap-northeast-2:806072073708:image/datascience-1.0'
      datawrangler: >-
        arn:aws:sagemaker:ap-northeast-2:131546521161:image/sagemaker-data-wrangler-1.0
    ap-southeast-1:
      datascience: 'arn:aws:sagemaker:ap-southeast-1:492261229750:image/datascience-1.0'
      datawrangler: >-
        arn:aws:sagemaker:ap-southeast-1:119527597002:image/sagemaker-data-wrangler-1.0
    ap-southeast-2:
      datascience: 'arn:aws:sagemaker:ap-southeast-2:452832661640:image/datascience-1.0'
      datawrangler: >-
        arn:aws:sagemaker:ap-southeast-2:422173101802:image/sagemaker-data-wrangler-1.0
    ap-northeast-1:
      datascience: 'arn:aws:sagemaker:ap-northeast-1:102112518831:image/datascience-1.0'
      datawrangler: >-
        arn:aws:sagemaker:ap-northeast-1:649008135260:image/sagemaker-data-wrangler-1.0
    ca-central-1:
      datascience: 'arn:aws:sagemaker:ca-central-1:310906938811:image/datascience-1.0'
      datawrangler: >-
        arn:aws:sagemaker:ca-central-1:557239378090:image/sagemaker-data-wrangler-1.0
    eu-central-1:
      datascience: 'arn:aws:sagemaker:eu-central-1:936697816551:image/datascience-1.0'
      datawrangler: >-
        arn:aws:sagemaker:eu-central-1:024640144536:image/sagemaker-data-wrangler-1.0
    eu-west-1:
      datascience: 'arn:aws:sagemaker:eu-west-1:470317259841:image/datascience-1.0'
      datawrangler: >-
        arn:aws:sagemaker:eu-west-1:245179582081:image/sagemaker-data-wrangler-1.0
    eu-west-2:
      datascience: 'arn:aws:sagemaker:eu-west-2:712779665605:image/datascience-1.0'
      datawrangler: >-
        arn:aws:sagemaker:eu-west-2:894491911112:image/sagemaker-data-wrangler-1.0
    eu-west-3:
      datascience: 'arn:aws:sagemaker:eu-west-3:615547856133:image/datascience-1.0'
      datawrangler: >-
        arn:aws:sagemaker:eu-west-3:807237891255:image/sagemaker-data-wrangler-1.0
    eu-north-1:
      datascience: 'arn:aws:sagemaker:eu-north-1:243637512696:image/datascience-1.0'
      datawrangler: >-
        arn:aws:sagemaker:eu-north-1:054986407534:image/sagemaker-data-wrangler-1.0
    eu-south-1:
      datascience: >-
        arn:aws:sagemaker:eu-south-1:488287956546:image/sagemaker-data-wrangler-1.0
      datawrangler: 'arn:aws:sagemaker:eu-south-1:592751261982:image/datascience-1.0'
    sa-east-1:
      datascience: 'arn:aws:sagemaker:sa-east-1:782484402741:image/datascience-1.0'
      datawrangler: >-
        arn:aws:sagemaker:sa-east-1:424196993095:image/sagemaker-data-wrangler-1.0
Parameters:
  SageMakerDomainName:
    Type: String
    Description: SageMaker Studio domain id
    Default: studio-domain
  SecurityGroupId:
    Type: String
    Description: >-
      List of security group ids separated by comma. Can be empty if it is
      stored in SSM
    Default: ''
  SubnetId:
    Type: String
    Description: List of subnet ids separated by comma. Can be empty if it is stored in SSM
    Default: ''
  VpcId:
    Type: String
    Description: VPC ID used for the SageMaker domain. Can be empty if it is stored in SSM
    Default: ''
Conditions:
  SecurityGroupIdEmpty: !Equals
    - !Ref SecurityGroupId
    - ''
  SubnetIdEmpty: !Equals
    - !Ref SubnetId
    - ''
  VpcIdEmpty: !Equals
    - !Ref VpcId
    - ''
Resources:
  S3BucketStudio:
    Type: 'AWS::S3::Bucket'
    Properties:
      VersioningConfiguration:
        Status: Enabled
      PublicAccessBlockConfiguration:
        BlockPublicAcls: true
        IgnorePublicAcls: true
        BlockPublicPolicy: true
        RestrictPublicBuckets: true
  DefaultPolicy:
    Type: 'AWS::IAM::ManagedPolicy'
    Properties:
      PolicyDocument:
        Version: 2012-10-17
        Statement:
          - Sid: SageMakerBasePolicy
            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'
              - 'elastic-inference:Connect'
              - 'iam:GetRole'
              - 'iam:ListRoles'
              - 'iam:PassRole'
              - 'lambda:InvokeFunction'
              - 'lambda:ListFunctions'
              - 'logs:CreateLogGroup'
              - 'logs:CreateLogStream'
              - 'logs:DescribeLogStreams'
              - 'logs:GetLogEvents'
              - 'logs:PutLogEvents'
              - 'sts:AssumeRole'
            Resource: '*'
  KMSPolicy:
    Type: AWS::IAM::ManagedPolicy
    Properties:
      PolicyDocument:
        Version: "2012-10-17"
        Statement:
          - Effect: Allow
            Action:
              - "kms:CreateGrant"
              - "kms:Decrypt"
              - "kms:DescribeKey"
              - "kms:Encrypt"
              - "kms:ReEncrypt*"
              - "kms:GenerateDataKey"
              - "kms:ListKeys"
              - "kms:ListAliases"
            Resource:
              - "*"
  S3Policy:
    Type: 'AWS::IAM::ManagedPolicy'
    Properties:
      PolicyDocument:
        Version: 2012-10-17
        Statement:
          - Action:
              - 's3:Create*'
            Effect: Allow
            Resource: '*'
          - Effect: Allow
            Action:
              - 's3:AbortMultipartUpload'
              - 's3:DeleteObject'
              - 's3:Describe*'
              - 's3:GetObject'
              - 's3:PutBucket*'
              - 's3:PutObject*'
              - 's3:GetBucket*'
              - 's3:ListBucket'
            Resource:
              - '*'
  ServicesPolicy:
    Type: 'AWS::IAM::ManagedPolicy'
    Properties:
      PolicyDocument:
        Version: 2012-10-17
        Statement:
          - Effect: Allow
            Action:
              - 'cloudformation:*'
              - 'cloudwatch:*'
              - 'codecommit:*'
              - 'ecr:*'
              - 'lambda:*'
              - 'logs:*'
              - 'sagemaker:*'
              - 'servicecatalog:*'
              - 'ssm:*'
              - 'sts:GetCallerIdentity'
            Resource: '*'
          - Effect: Allow
            Action:
              - 'iam:PassRole'
            Resource:
              - !Sub 'arn:aws:iam::${AWS::AccountId}:role/*'
  SageMakerExecutionRole:
    Type: 'AWS::IAM::Role'
    Properties:
      RoleName: mlops-sagemaker-execution-role
      AssumeRolePolicyDocument:
        Version: 2012-10-17
        Statement:
          - Effect: Allow
            Action: 'sts:AssumeRole'
            Principal:
              Service:
                - cloudformation.amazonaws.com
                - lambda.amazonaws.com
                - sagemaker.amazonaws.com
      Path: /
      ManagedPolicyArns:
        - !Ref DefaultPolicy
        - !Ref KMSPolicy
        - !Ref S3Policy
        - !Ref ServicesPolicy
        - arn:aws:iam::aws:policy/ComprehendFullAccess
        - arn:aws:iam::aws:policy/TranslateFullAccess
  SageMakerStudioDomain:
    Type: 'AWS::SageMaker::Domain'
    Properties:
      AppNetworkAccessType: VpcOnly
      AuthMode: IAM
      DefaultUserSettings:
        ExecutionRole: !GetAtt
          - SageMakerExecutionRole
          - Arn
        SecurityGroups: !If
          - SecurityGroupIdEmpty
          - - '{{resolve:ssm:/networking/sg}}'
          - - !Ref SecurityGroupId
        SharingSettings:
          NotebookOutputOption: Allowed
          S3OutputPath: !Sub 's3://${S3BucketStudio}/studio'
      DomainName: !Ref SageMakerDomainName
      SubnetIds: !If
        - SubnetIdEmpty
        - - '{{resolve:ssm:/networking/subnet/private}}'
        - - !Ref SubnetId
      VpcId: !If 
        - VpcIdEmpty
        - '{{resolve:ssm:/networking/vpc}}'
        - !Ref VpcId
  EnableSageMakerProjects:
    Type: 'Custom::ResourceForEnablingSageMakerProjects'
    Properties:
      ServiceToken: !GetAtt
        - EnableSageMakerProjectsLambda
        - Arn
      ExecutionRole: !GetAtt
        - SageMakerExecutionRole
        - Arn
  EnableSageMakerProjectsLambda:
    Type: 'AWS::Lambda::Function'
    DependsOn: SageMakerStudioDomain
    Properties:
      Code:
        ZipFile: |
          # Function: EnableSagemakerProjects
          # Purpose:  Enables Sagemaker Projects
          import json
          import boto3
          import cfnresponse
          from botocore.exceptions import ClientError

          client = boto3.client("sagemaker")
          sc_client = boto3.client("servicecatalog")
          def lambda_handler(event, context):
              try:
                  response_status = cfnresponse.SUCCESS
                  if "RequestType" in event and event["RequestType"] == "Create":
                      enable_sm_projects(event["ResourceProperties"]["ExecutionRole"])
                  cfnresponse.send(event, context, response_status, {}, "")
              except ClientError as exception:
                  print(exception)
                  cfnresponse.send(event, context, cfnresponse.FAILED, {}, physicalResourceId=event.get("PhysicalResourceId"), reason=str(exception))

          def enable_sm_projects(studio_role_arn):
              # enable Project on account level (accepts portfolio share)
              response = client.enable_sagemaker_servicecatalog_portfolio()
              # associate studio role with portfolio
              response = sc_client.list_accepted_portfolio_shares()
              portfolio_id = ""
              for portfolio in response["PortfolioDetails"]:
                  if portfolio["ProviderName"] == "Amazon SageMaker":
                      portfolio_id = portfolio["Id"]
              response = sc_client.associate_principal_with_portfolio(
                  PortfolioId=portfolio_id,
                  PrincipalARN=studio_role_arn,
                  PrincipalType="IAM"
              )
      Description: Enable Sagemaker Projects
      Handler: index.lambda_handler
      MemorySize: 128
      Role: !GetAtt
        - SageMakerExecutionRole
        - Arn
      Runtime: python3.8
      Timeout: 120
  SageMakerStudioUserProfile:
    Type: 'AWS::SageMaker::UserProfile'
    Properties:
      DomainId: !Ref SageMakerStudioDomain
      UserProfileName: bpistone
  JupyterApp:
    Type: 'AWS::SageMaker::App'
    DependsOn: SageMakerStudioUserProfile
    Properties:
      AppName: default
      AppType: JupyterServer
      DomainId: !Ref SageMakerStudioDomain
      UserProfileName: bpistone
  DataScienceApp:
    Type: 'AWS::SageMaker::App'
    DependsOn: SageMakerStudioUserProfile
    Properties:
      AppName: sm-mlops-datascience-ml-t3-medium
      AppType: KernelGateway
      DomainId: !Ref SageMakerStudioDomain
      ResourceSpec:
        InstanceType: ml.t3.medium
        SageMakerImageArn: !FindInMap
          - RegionMap
          - !Ref 'AWS::Region'
          - datascience
      UserProfileName: bpistone
  DeleteKernelGatewayAppsNoKGCondition:
    Type: 'Custom::DeleteKernelGatewayApps'
    DependsOn: SageMakerStudioUserProfile
    Properties:
      ServiceToken: !GetAtt
        - DeleteKernelGatewayAppsLambda
        - Arn
      DomainId: !Ref SageMakerStudioDomain
  DeleteKernelGatewayAppsKGCondition:
    Type: 'Custom::DeleteKernelGatewayApps'
    DependsOn:
      - DataScienceApp
    Properties:
      ServiceToken: !GetAtt
        - DeleteKernelGatewayAppsLambda
        - Arn
      DomainId: !Ref SageMakerStudioDomain
  DeleteKernelGatewayAppsLambda:
    Type: 'AWS::Lambda::Function'
    Properties:
      Code:
        ZipFile: |
          # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
          # SPDX-License-Identifier: MIT-0
          import time
          import boto3
          import logging
          import cfnresponse
          from botocore.exceptions import ClientError
          sm_client = boto3.client('sagemaker')
          logger = logging.getLogger(__name__)
          def delete_apps(domain_id):
              logging.info(f'Start deleting apps for domain id: {domain_id}')
              try:
                  sm_client.describe_domain(DomainId=domain_id)
              except:
                  logging.info(f'Cannot retrieve {domain_id}')
                  return
              for p in sm_client.get_paginator('list_apps').paginate(DomainIdEquals=domain_id):
                  for a in p['Apps']:
                      if a['AppType'] == 'KernelGateway' and a['Status'] != 'Deleted':
                          sm_client.delete_app(DomainId=a['DomainId'], UserProfileName=a['UserProfileName'], AppType=a['AppType'], AppName=a['AppName'])

              apps = 1
              while apps:
                  apps = 0
                  for p in sm_client.get_paginator('list_apps').paginate(DomainIdEquals=domain_id):
                      apps += len([a['AppName'] for a in p['Apps'] if a['AppType'] == 'KernelGateway' and a['Status'] != 'Deleted'])
                  logging.info(f'Number of active KernelGateway apps: {str(apps)}')
                  time.sleep(5)
              logger.info(f'KernelGateway apps for {domain_id} deleted')
              return
          def lambda_handler(event, context):
              response_data = {}
              physicalResourceId = event.get('PhysicalResourceId')
              try:
                  if event['RequestType'] in ['Create', 'Update']:
                      physicalResourceId = event.get('ResourceProperties')['DomainId']

                  elif event['RequestType'] == 'Delete':
                      delete_apps(physicalResourceId)
                  cfnresponse.send(event, context, cfnresponse.SUCCESS, response_data, physicalResourceId=physicalResourceId)
              except ClientError as exception:
                  logging.error(exception)
                  cfnresponse.send(event, context, cfnresponse.FAILED, response_data, physicalResourceId=physicalResourceId, reason=str(exception))
      Description: Delete KernelGateway apps to clean up
      Handler: index.lambda_handler
      MemorySize: 128
      Role: !GetAtt
        - SageMakerExecutionRole
        - Arn
      Runtime: python3.8
      Timeout: 900