AWSTemplateFormatVersion: 2010-09-09
Description: |
  MLOps SageMaker Project for multi-account ML model deployment. 
  This template creates a CI/CD pipeline to deploy a given inference image and pretrained Model to two separate account -- staging and production.

Parameters:
  SageMakerProjectName:
    Type: String
    Description: Name of the project
    MinLength: 1
    MaxLength: 32
    AllowedPattern: ^[a-zA-Z](-*[a-zA-Z0-9])*

  SageMakerProjectId:
    Type: String
    Description: Service generated Id of the project.

  ModelPackageGroupName:
    Type: String
    Description: Model package group name to monitor for model package state changes. Leave 'Auto' to auto generate.
    Default: 'Auto'
  
  MultiAccountDeployment:
    Type: String
    Description: Enable multi-account deployment of the model. The staging and production OU Ids must be set in the environment.
    AllowedValues:
      - 'YES'
      - 'NO'
    Default: 'NO'

Conditions:
  MLOpsArtifactBucketCondition: !Equals [ 'true', 'true' ]
  GenerateModelPackageNameCondition: !Equals [ !Ref ModelPackageGroupName, 'Auto' ]
  MultiAccountDeploymentCondition: !Equals [ !Ref MultiAccountDeployment, 'YES' ]

Resources:
  # Retrieve the environment variables
  GetEnvironmentConfiguration:
    Type: Custom::GetEnvironmentConfiguration
    Properties:
      ServiceToken: !ImportValue 'ds-get-environment-configuration-lambda-arn'
      SageMakerProjectName: !Ref SageMakerProjectName
      SSMParams:
        - 
          VariableName: 'DataBucketName'
          ParameterName: 'data-bucket-name'
        - 
          VariableName: 'ModelBucketName'
          ParameterName: 'model-bucket-name'
        - 
          VariableName: 'S3VPCEId'
          ParameterName: 's3-vpce-id'
        - 
          VariableName: 'S3KmsKeyId'
          ParameterName: 'kms-s3-key-arn'
        - 
          VariableName: 'PipelineExecutionRole'
          ParameterName: 'sm-pipeline-execution-role-arn'
        - 
          VariableName: 'OUStagingId'
          ParameterName: 'ou-staging-id'
        - 
          VariableName: 'OUProdId'
          ParameterName: 'ou-prod-id'
        - 
          VariableName: 'ProdAccountList'
          ParameterName: 'production-account-list'
        - 
          VariableName: 'StagingAccountList'
          ParameterName: 'staging-account-list'
        - 
          VariableName: 'ModelExecutionRole'
          ParameterName: 'sm-model-execution-role-name'
        - 
          VariableName: 'StackSetExecutionRole'
          ParameterName: 'stackset-execution-role-name'
        - 
          VariableName: 'StackSetAdministrationRole'
          ParameterName: 'stackset-administration-role-arn'
        - 
          VariableName: 'EbsKmsKeyArn'
          ParameterName: 'kms-ebs-key-arn'
        -
          VariableName: 'EnvTypeStagingName'
          ParameterName: 'env-type-staging-name'
        -
          VariableName: 'EnvTypeProdName'
          ParameterName: 'env-type-prod-name'
        - 
          VariableName: 'SeedCodeS3BucketName'
          ParameterName: 'seed-code-s3bucket-name'
        
  MlOpsArtifactsBucket:
    Type: AWS::S3::Bucket
    Condition: MLOpsArtifactBucketCondition
    DeletionPolicy: Retain
    UpdateReplacePolicy: Retain
    Properties:
      BucketName: !Sub sm-mlops-cp-${SageMakerProjectName}-${SageMakerProjectId} # 12+32+15=59 chars max/ 63 allowed
      AccessControl: Private
      PublicAccessBlockConfiguration:
        BlockPublicAcls: TRUE
        BlockPublicPolicy: TRUE
        IgnorePublicAcls: TRUE
        RestrictPublicBuckets: TRUE
      BucketEncryption:
        ServerSideEncryptionConfiguration:
          - ServerSideEncryptionByDefault:
              SSEAlgorithm: 'aws:kms'
              KMSMasterKeyID: !GetAtt GetEnvironmentConfiguration.S3KmsKeyId
      Tags:
        - Key: SageMakerProjectName
          Value: !Ref SageMakerProjectName
        - Key: SageMakerProjectId
          Value: !Ref SageMakerProjectId
        - Key: EnvironmentName
          Value: !GetAtt GetEnvironmentConfiguration.EnvironmentName
        - Key: EnvironmentType
          Value: !GetAtt GetEnvironmentConfiguration.EnvironmentType

  ModelDeployRegistryEventRule:
    Type: AWS::Events::Rule
    Properties:
      # Max length allowed: 64
      Name: !Sub sagemaker-${SageMakerProjectName}-${SageMakerProjectId}-model # max: 10+33+15+5=63 chars
      Description: "Rule to trigger a deployment when SageMaker Model registry is updated with a new model package. For example, a new model package is registered with Registry"
      EventPattern:
        source:
          - "aws.sagemaker"
        detail-type:
          - "SageMaker Model Package State Change"
        detail:
          ModelPackageGroupName: 
            - !If 
              - GenerateModelPackageNameCondition
              - !Sub '${SageMakerProjectName}-${SageMakerProjectId}'
              - !Ref ModelPackageGroupName
      State: "ENABLED"
      Targets:
        -
          Arn:
            !Sub 'arn:${AWS::Partition}:codepipeline:${AWS::Region}:${AWS::AccountId}:${ModelDeployPipeline}'
          RoleArn:
            !Sub 'arn:${AWS::Partition}:iam::${AWS::AccountId}:role/service-role/AmazonSageMakerServiceCatalogProductsUseRole'
          Id: !Sub sagemaker-${SageMakerProjectName}-trigger

  ModelDeployCodeCommitEventRule:
    Type: AWS::Events::Rule
    Properties:
      # Max length allowed: 64
      Name: !Sub sagemaker-${SageMakerProjectName}-${SageMakerProjectId}-code # max: 10+33+15+5=63 chars
      Description: "Rule to launch a pipeline run when ModelDeploy CodeCommit repository is updated"
      EventPattern:
        source:
          - "aws.codecommit"
        detail-type:
          - "CodeCommit Repository State Change"
        resources:
          - !GetAtt ModelDeployCodeCommitRepository.Arn
        detail:
          referenceType:
            - "branch"
          referenceName:
            - "main"
      State: "ENABLED"
      Targets:
        -
          Arn:
            !Sub 'arn:${AWS::Partition}:codepipeline:${AWS::Region}:${AWS::AccountId}:${ModelDeployPipeline}'
          RoleArn:
            !Sub 'arn:${AWS::Partition}:iam::${AWS::AccountId}:role/service-role/AmazonSageMakerServiceCatalogProductsUseRole'
          Id: !Sub codecommit-${SageMakerProjectName}-modeldeploy

  ModelDeployCodeCommitRepository:
    Type: AWS::CodeCommit::Repository
    Properties:
      # Max allowed length: 100 chars
      RepositoryName: !Sub sagemaker-${SageMakerProjectName}-${SageMakerProjectId}-model-deploy # max: 10+33+15+12=70
      RepositoryDescription: !Sub SageMaker Endpoint deployment infrastructure as code for the project ${SageMakerProjectName}
      Code:
        S3:
          Bucket: !GetAtt GetEnvironmentConfiguration.SeedCodeS3BucketName 
          Key: sagemaker-mlops/seed-code/mlops-model-deploy-v1.0.zip
        BranchName: main
      Tags:
        - Key: SageMakerProjectName
          Value: !Ref SageMakerProjectName
        - Key: SageMakerProjectId
          Value: !Ref SageMakerProjectId
        - Key: EnvironmentName
          Value: !GetAtt GetEnvironmentConfiguration.EnvironmentName
        - Key: EnvironmentType
          Value: !GetAtt GetEnvironmentConfiguration.EnvironmentType

  ModelDeployBuildProject:
    Type: AWS::CodeBuild::Project
    Properties:
      # Max length: 255 chars
      Name: !Sub sagemaker-${SageMakerProjectName}-${SageMakerProjectId}-modeldeploy # max: 10+33+15+10=68
      Description: Pulls the code from Model Deploy CodeCommit repository, builds the CloudFormation templates with the endpoint and deploys them
      ServiceRole: !Sub 'arn:${AWS::Partition}:iam::${AWS::AccountId}:role/service-role/AmazonSageMakerServiceCatalogProductsUseRole'
      Artifacts:
        Type: CODEPIPELINE
      Environment:
        Type: LINUX_CONTAINER
        ComputeType: BUILD_GENERAL1_SMALL
        Image: aws/codebuild/amazonlinux2-x86_64-standard:3.0
        EnvironmentVariables:
          - Name: SAGEMAKER_PROJECT_NAME
            Value: !Ref SageMakerProjectName
          - Name: SAGEMAKER_PROJECT_ID
            Value: !Ref SageMakerProjectId
          - Name: ENV_NAME
            Value: !GetAtt GetEnvironmentConfiguration.EnvironmentName
          - Name: ENV_TYPE
            Value: !GetAtt GetEnvironmentConfiguration.EnvironmentType
          - Name: ARTIFACT_BUCKET
            Value: !If 
              - MLOpsArtifactBucketCondition
              - !Ref MlOpsArtifactsBucket
              - !GetAtt GetEnvironmentConfiguration.DataBucketName
          - Name: DATA_BUCKET
            Value: !GetAtt GetEnvironmentConfiguration.DataBucketName
          - Name: MODEL_BUCKET
            Value: !GetAtt GetEnvironmentConfiguration.ModelBucketName
          - Name: SOURCE_MODEL_PACKAGE_GROUP_NAME
            Value: !If 
              - GenerateModelPackageNameCondition
              - !Sub '${SageMakerProjectName}-${SageMakerProjectId}'
              - !Ref ModelPackageGroupName
          - Name: AWS_REGION
            Value: !Ref AWS::Region
          - Name: MULTI_ACCOUNT_DEPLOYMENT
            Value: !Ref MultiAccountDeployment
          - Name: STAGING_ACCOUNT_LIST
            Value: !If 
                  - MultiAccountDeploymentCondition
                  - !GetAtt GetEnvironmentConfiguration.StagingAccountList
                  - !Ref 'AWS::AccountId'
          - Name: PROD_ACCOUNT_LIST
            Value: !If 
                  - MultiAccountDeploymentCondition
                  - !GetAtt GetEnvironmentConfiguration.ProdAccountList
                  - !Ref 'AWS::AccountId'
          - Name: STAGING_CONFIG_NAME
            Value: 'staging-config'
          - Name: PROD_CONFIG_NAME
            Value: 'prod-config'
          - Name: CFN_TEMPLATE_NAME
            Value: 'cfn-sm-endpoint'      
          - Name: SAGEMAKER_EXECUTION_ROLE_STAGING_NAME
            Value: !GetAtt GetEnvironmentConfiguration.ModelExecutionRole
          - Name: SAGEMAKER_EXECUTION_ROLE_PROD_NAME
            Value: !GetAtt GetEnvironmentConfiguration.ModelExecutionRole
          - Name: SAGEMAKER_EBS_KMS_KEY_ARN
            Value: !GetAtt GetEnvironmentConfiguration.EbsKmsKeyArn
          - Name: ENV_TYPE_STAGING_NAME
            Value: !GetAtt GetEnvironmentConfiguration.EnvTypeStagingName
          - Name: ENV_TYPE_PROD_NAME
            Value: !GetAtt GetEnvironmentConfiguration.EnvTypeProdName

      Source:
        Type: CODEPIPELINE
        BuildSpec: buildspec.yml
      TimeoutInMinutes: 15
      VpcConfig:
        SecurityGroupIds: !GetAtt GetEnvironmentConfiguration.SecurityGroups
        Subnets: !GetAtt GetEnvironmentConfiguration.SubnetIds
        VpcId: !GetAtt GetEnvironmentConfiguration.VpcId
      Tags:
        - Key: SageMakerProjectName
          Value: !Ref SageMakerProjectName
        - Key: SageMakerProjectId
          Value: !Ref SageMakerProjectId
        - Key: EnvironmentName
          Value: !GetAtt GetEnvironmentConfiguration.EnvironmentName
        - Key: EnvironmentType
          Value: !GetAtt GetEnvironmentConfiguration.EnvironmentType

  ModelDeployTestProject:
    Type: AWS::CodeBuild::Project
    Properties:
      # Max length: 255 chars
      Name: !Sub sagemaker-${SageMakerProjectName}-${SageMakerProjectId}-modeldeploy-test # max: 10+33+15+7=65
      Description: Test the deployment endpoint in the staging account
      ServiceRole: !Sub 'arn:${AWS::Partition}:iam::${AWS::AccountId}:role/service-role/AmazonSageMakerServiceCatalogProductsUseRole'
      Artifacts:
        Type: CODEPIPELINE
      Environment:
        Type: LINUX_CONTAINER
        ComputeType: BUILD_GENERAL1_SMALL
        Image: "aws/codebuild/amazonlinux2-x86_64-standard:3.0"
        EnvironmentVariables:
          - Name: AWS_REGION
            Value: !Ref "AWS::Region"
          - Name: BUILD_CONFIG
            Value: 'staging-config'
          - Name: TEST_RESULTS
            Value: 'test-results'

      Source:
        Type: CODEPIPELINE
        BuildSpec: test/buildspec.yml
      TimeoutInMinutes: 30
      VpcConfig:
        SecurityGroupIds: !GetAtt GetEnvironmentConfiguration.SecurityGroups
        Subnets: !GetAtt GetEnvironmentConfiguration.SubnetIds
        VpcId: !GetAtt GetEnvironmentConfiguration.VpcId
      Tags:
        - Key: SageMakerProjectName
          Value: !Ref SageMakerProjectName
        - Key: SageMakerProjectId
          Value: !Ref SageMakerProjectId
        - Key: EnvironmentName
          Value: !GetAtt GetEnvironmentConfiguration.EnvironmentName
        - Key: EnvironmentType
          Value: !GetAtt GetEnvironmentConfiguration.EnvironmentType

  ModelDeployPipeline:
    Type: AWS::CodePipeline::Pipeline
    Properties:
      # Max length: 100 chars
      Name: !Sub sagemaker-${SageMakerProjectName}-${SageMakerProjectId}-modeldeploy # max: 10+33+15+11=69
      RoleArn: !Sub 'arn:${AWS::Partition}:iam::${AWS::AccountId}:role/service-role/AmazonSageMakerServiceCatalogProductsUseRole'
      ArtifactStore:
        Type: S3
        Location: !If 
          - MLOpsArtifactBucketCondition
          - !Ref MlOpsArtifactsBucket
          - !GetAtt GetEnvironmentConfiguration.DataBucketName
      
      Tags:
        - Key: SageMakerProjectName
          Value: !Ref SageMakerProjectName
        - Key: SageMakerProjectId
          Value: !Ref SageMakerProjectId
        - Key: EnvironmentName
          Value: !GetAtt GetEnvironmentConfiguration.EnvironmentName
        - Key: EnvironmentType
          Value: !GetAtt GetEnvironmentConfiguration.EnvironmentType

      Stages:
        - Name: Source
          Actions:
            - Name: ModelDeploySource
              ActionTypeId:
                Category: Source
                Owner: AWS
                Provider: CodeCommit
                Version: '1'
              Configuration:
                PollForSourceChanges: false
                RepositoryName: !GetAtt ModelDeployCodeCommitRepository.Name
                BranchName: main
              OutputArtifacts:
                - Name: ModelDeploySourceArtifact

        # Stage: Build CFN Template for model deployment
        - Name: Build
          Actions:
            - Name: BuildEndpointDeploymentTemplate
              ActionTypeId:
                Category: Build
                Owner: AWS
                Provider: CodeBuild
                Version: '1'
              InputArtifacts:
                - Name: ModelDeploySourceArtifact
              OutputArtifacts:
                - Name: BuildArtifact
              Configuration:
                ProjectName: !Ref ModelDeployBuildProject
              RunOrder: 1

        # Stage: Deploy and test into the staging accounts
        - Name: DeployModelStaging
          Actions:

            # Action 1
            - Name: DeployStaging 
              InputArtifacts:
                - Name: BuildArtifact
              ActionTypeId:
                Category: Deploy
                Owner: AWS
                Version: '1'
                Provider: CloudFormationStackSet
              Configuration:
                StackSetName: !Sub 'sagemaker-${SageMakerProjectName}-${SageMakerProjectId}-deploy-${GetEnvironmentConfiguration.EnvTypeStagingName}'  #10+33+15+14=72 out of 128 max
                Description: 'SageMaker endpoint in the staging target OU'
                TemplatePath: BuildArtifact::cfn-sm-endpoint.yaml
                Parameters: BuildArtifact::staging-config.json
                Capabilities: CAPABILITY_NAMED_IAM
                PermissionModel: 'SELF_MANAGED'
                AdministrationRoleArn: !GetAtt GetEnvironmentConfiguration.StackSetAdministrationRole
                ExecutionRoleName: !GetAtt GetEnvironmentConfiguration.StackSetExecutionRole
                # For a self-managed model, targets can only be AWS accounts
                DeploymentTargets: !If 
                  - MultiAccountDeploymentCondition
                  - !GetAtt GetEnvironmentConfiguration.StagingAccountList
                  - !Ref 'AWS::AccountId'
                Regions: !Ref 'AWS::Region'
              RunOrder: 1

            # Action 2
            - Name: TestStaging
              ActionTypeId:
                Category: Build
                Owner: AWS
                Provider: CodeBuild
                Version: '1'
              InputArtifacts:
                - Name: ModelDeploySourceArtifact
                - Name: BuildArtifact
              OutputArtifacts:
                - Name: TestArtifact
              Configuration:
                ProjectName: !Ref ModelDeployTestProject
                PrimarySource: ModelDeploySourceArtifact
              RunOrder: 2

            # Action 3
            - Name: ApproveStagingDeployment
              ActionTypeId:
                Category: Approval
                Owner: AWS
                Version: '1'
                Provider: Manual
              Configuration:
                CustomData: !Sub 
                  - "Model ${ModelPackageGroupName} for the project ${SageMakerProjectName} is deployed to the staging accounts ${AccountList}"
                  - ModelPackageGroupName: !If 
                      - GenerateModelPackageNameCondition
                      - !Sub '${SageMakerProjectName}-${SageMakerProjectId}'
                      - !Ref ModelPackageGroupName
                    SageMakerProjectName: !Ref SageMakerProjectName
                    AccountList: !If 
                      - MultiAccountDeploymentCondition
                      - !GetAtt GetEnvironmentConfiguration.StagingAccountList
                      - !Ref 'AWS::AccountId'
                ExternalEntityLink: !Sub 'https://${AWS::Region}.console.aws.amazon.com/codesuite/codepipeline/pipelines/sagemaker-${SageMakerProjectName}-${SageMakerProjectId}-modeldeploy/view?region=${AWS::Region}'
              RunOrder: 3

        # Stage: Deploy into the production accounts
        - Name: DeployModelProd
          Actions:
            - Name: DeployProd
              InputArtifacts:
                - Name: BuildArtifact
              ActionTypeId:
                Category: Deploy
                Owner: AWS
                Version: '1'
                Provider: CloudFormationStackSet
              Configuration:
                StackSetName: !Sub 'sagemaker-${SageMakerProjectName}-${SageMakerProjectId}-deploy-${GetEnvironmentConfiguration.EnvTypeProdName}'  #10+33+15+14=72 out of 128 max
                Description: 'SageMaker endpoint in the production target OU'
                Parameters: BuildArtifact::prod-config.json
                TemplatePath: BuildArtifact::cfn-sm-endpoint.yaml
                Capabilities: CAPABILITY_NAMED_IAM
                PermissionModel: 'SELF_MANAGED'
                AdministrationRoleArn: !GetAtt GetEnvironmentConfiguration.StackSetAdministrationRole
                ExecutionRoleName: !GetAtt GetEnvironmentConfiguration.StackSetExecutionRole
                # For a self-managed model, targets can only be AWS accounts
                DeploymentTargets: !If 
                  - MultiAccountDeploymentCondition
                  - !GetAtt GetEnvironmentConfiguration.ProdAccountList
                  - !Ref 'AWS::AccountId'
                Regions: !Ref 'AWS::Region'
              RunOrder: 1