# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # SPDX-License-Identifier: MIT-0 AWSTemplateFormatVersion: 2010-09-09 Description: | Create a SageMaker Studio domain Metadata: AWS::CloudFormation::Interface: ParameterGroups: - Label: default: Data Science Environment Parameters: - EnvironmentName - Label: default: Amazon SageMaker Studio Parameters: - DomainName - AuthMode - Label: default: Network and Storage Configuration Parameters: - VPCId - SageMakerDomainSubnetIds - SageMakerDomainSecurityGroupIds - SageMakerDomainStorageKMSKeyId - Label: default: Permissions Parameters: - SageMakerDefaultExecutionRoleArn - SetupLambdaExecutionRoleArn ParameterLabels: EnvironmentName: default: Environment name DomainName: default: Domain name AuthMode: default: Authentication mode VPCId: default: VPC SageMakerDomainSubnetIds: default: Subnet(s) SageMakerDomainSecurityGroupIds: default: Security group(s) SageMakerDomainStorageKMSKeyId: default: Storage encryption key NetworkAccessType: default: Network access for SageMaker Studio SageMakerDefaultExecutionRoleArn: default: SageMaker default execution role SetupLambdaExecutionRoleArn: default: Execution role for setup Lambda function Outputs: SageMakerDomainId: Description: SageMaker Domain Id Value: !GetAtt SageMakerStudioDomain.DomainId Export: Name: 'ds-sagemaker-domain-id' Parameters: EnvironmentName: Type: String AllowedPattern: '[a-z0-9\-]*' Description: Please specify your SageMaker environment name. Default: 'sm-environment' DomainName: Type: String Description: SageMaker Studio domain name. Leave empty to auto generate. Default: '' VPCId: Type: AWS::EC2::VPC::Id Description: Choose a VPC for SageMaker Studio and SageMaker workloads SageMakerDomainSubnetIds: Type: List Description: Choose subnets or provide a comma-delimited list of subnet ids SageMakerDomainSecurityGroupIds: Type: List Description: Choose security groups for SageMaker Studio and SageMaker workloads SageMakerDomainStorageKMSKeyId: Type: String Description: SageMaker uses an AWS managed CMK to encrypt your EFS and EBS file systems by default. To use a customer managed CMK, enter its key Id. Default: '' NetworkAccessType: Type: String AllowedValues: - 'PublicInternetOnly' - 'VpcOnly' Description: Choose how SageMaker Studio accesses resources over the Network Default: 'VpcOnly' AuthMode: Type: String AllowedValues: - 'IAM' Description: The mode of authentication that members use to access the domain. Only IAM is supported for this solution. Default: 'IAM' SageMakerDefaultExecutionRoleArn: Type: String Description: The ARN of the SageMaker execution role SetupLambdaExecutionRoleArn: Type: String Description: The ARN of the execution role for the Lambda function for SageMaker Studio setup Conditions: GenerateDomainNameCondition: !Equals [ !Ref DomainName, '' ] SageMakerEFSKMSKeyCondition: !Not [ !Equals [ !Ref SageMakerDomainStorageKMSKeyId, ''] ] Resources: SageMakerStudioDomain: Type: AWS::SageMaker::Domain Properties: AppNetworkAccessType: !Ref NetworkAccessType AuthMode: !Ref AuthMode DefaultUserSettings: ExecutionRole: !Ref SageMakerDefaultExecutionRoleArn SecurityGroups: !Ref SageMakerDomainSecurityGroupIds DomainName: !If - GenerateDomainNameCondition - !Sub '${EnvironmentName}-${AWS::Region}-sagemaker-domain' - !Ref DomainName KmsKeyId: !If [ SageMakerEFSKMSKeyCondition, !Ref SageMakerDomainStorageKMSKeyId, !Ref 'AWS::NoValue' ] SubnetIds: !Ref SageMakerDomainSubnetIds VpcId: !Ref VPCId Tags: - Key: EnvironmentName Value: !Ref EnvironmentName EnableSageMakerProjects: Type: Custom::ResourceForEnablingSageMakerProjects DependsOn: SageMakerStudioDomain Properties: ServiceToken: !GetAtt EnableSageMakerProjectsLambda.Arn ExecutionRole: !Ref SageMakerDefaultExecutionRoleArn DeleteDomainApps: Type: Custom::DeleteDomainApps Properties: ServiceToken: !GetAtt DeleteDomainAppsLambda.Arn DomainId: !GetAtt SageMakerStudioDomain.DomainId DeleteDomainAppsLambda: Type: AWS::Lambda::Function Properties: ReservedConcurrentExecutions: 1 Code: ZipFile: | # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # SPDX-License-Identifier: MIT-0 import time import boto3 import logging import json import cfnresponse from botocore.exceptions import ClientError sm_client = boto3.client('sagemaker') logger = logging.getLogger(__name__) def delete_user_profiles(domain_id): logger.info(f'Start deleting user profiles for domain id: {domain_id}') for p in sm_client.get_paginator('list_user_profiles').paginate(DomainIdEquals=domain_id): for up in p['UserProfiles']: if up['Status'] not in ('Deleting', 'Pending'): sm_client.delete_user_profile(DomainId=up['DomainId'], UserProfileName=up['UserProfileName']) up = 1 while up: up = 0 for p in sm_client.get_paginator('list_user_profiles').paginate(DomainIdEquals=domain_id): up += len([u['UserProfileName'] for u in p['UserProfiles'] if u['Status'] != 'Deleted']) logger.info(f'Number of active user profiles: {str(up)}') time.sleep(5) def delete_apps(domain_id): logger.info(f'Start deleting apps for domain id: {domain_id}') try: sm_client.describe_domain(DomainId=domain_id) except: logger.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['Status'] != 'Deleted': logger.info(f"Deleting {a['AppType']}:{a['AppName']}") 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['Status'] != 'Deleted']) logger.info(f'Number of active apps: {str(apps)}') time.sleep(5) logger.info(f'Apps for {domain_id} deleted') return def lambda_handler(event, context): response_data = {} try: physicalResourceId = event.get('PhysicalResourceId') logger.info(json.dumps(event)) if event['RequestType'] in ['Create', 'Update']: physicalResourceId = event.get('ResourceProperties')['DomainId'] elif event['RequestType'] == 'Delete': delete_apps(physicalResourceId) delete_user_profiles(physicalResourceId) cfnresponse.send(event, context, cfnresponse.SUCCESS, response_data, physicalResourceId=physicalResourceId) except (Exception, ClientError) as exception: logger.error(exception) cfnresponse.send(event, context, cfnresponse.FAILED, response_data, physicalResourceId=physicalResourceId, reason=str(exception)) Description: Delete SageMaker domain apps to clean up Handler: index.lambda_handler MemorySize: 128 Role: !Ref SetupLambdaExecutionRoleArn Runtime: python3.8 Timeout: 900 Tags: - Key: EnvironmentName Value: !Ref EnvironmentName EnableSageMakerProjectsLambda: Type: AWS::Lambda::Function DependsOn: SageMakerStudioDomain Properties: ReservedConcurrentExecutions: 1 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 (Exception, 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: !Ref SetupLambdaExecutionRoleArn Runtime: python3.8 Timeout: 60 Tags: - Key: EnvironmentName Value: !Ref EnvironmentName # SSM parameter SageMakerDomainIdSSM: Type: 'AWS::SSM::Parameter' Properties: Name: !Sub "${EnvironmentName}-sagemaker-domain-id" Type: String Value: !GetAtt SageMakerStudioDomain.DomainId Description: !Sub 'SageMaker Studio domain id for ${SageMakerStudioDomain.DomainArn}'