AWSTemplateFormatVersion: 2010-09-09
Description: |
  MLOps SageMaker Project for CI/CD pipeline for model build, train, validate. 
  This template creates a CI/CD pipeline to build, train, and validate an ML model using SageMaker pipelines

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.

Conditions:
  MLOpsArtifactBucketCondition: !Equals [ 'true', 'true' ]

Resources:

  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: 'S3KmsKeyId'
          ParameterName: 'kms-s3-key-arn'
        - 
          VariableName: 'PipelineExecutionRole'
          ParameterName: 'sm-pipeline-execution-role-arn'
        - 
          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

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

  ModelBuildCodeCommitRepository:
    Type: AWS::CodeCommit::Repository
    Properties:
      # Max allowed length: 100 chars
      RepositoryName: !Sub sagemaker-${SageMakerProjectName}-${SageMakerProjectId}-model-build-train # max: 10+33+15+18=76
      RepositoryDescription: !Sub SageMaker Model building infrastructure as code for the project ${SageMakerProjectName}
      Code:
        S3:
          Bucket: !GetAtt GetEnvironmentConfiguration.SeedCodeS3BucketName 
          Key: sagemaker-mlops/seed-code/mlops-model-build-train-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

  ModelPipelineBuildProject:
    Type: AWS::CodeBuild::Project
    Properties:
      # Max length: 255 chars
      Name: !Sub sagemaker-${SageMakerProjectName}-${SageMakerProjectId}-modelbuild # max: 10+33+15+10=68
      Description: Pulls the code from Model Build CodeCommit repository, creates the SageMaker Pipeline, and executes it
      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: DATA_BUCKET
            Value: !GetAtt GetEnvironmentConfiguration.DataBucketName
          - Name: SAGEMAKER_PIPELINE_NAME
            Value: !Sub sagemaker-${SageMakerProjectName}
          - Name: SAGEMAKER_PIPELINE_ROLE_ARN
            Value: !GetAtt GetEnvironmentConfiguration.PipelineExecutionRole
          - Name: AWS_REGION
            Value: !Ref AWS::Region
      Source:
        Type: CODEPIPELINE
        BuildSpec: buildspec.yml
      TimeoutInMinutes: 480
      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

  ModelBuildPipeline:
    Type: AWS::CodePipeline::Pipeline
    Properties:
      # Max length: 100 chars
      Name: !Sub sagemaker-${SageMakerProjectName}-${SageMakerProjectId}-modelbuild # max: 10+33+15+10=68
      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: ModelBuildSource
              ActionTypeId:
                Category: Source
                Owner: AWS
                Provider: CodeCommit
                Version: '1'
              Configuration:
                PollForSourceChanges: 'false'
                RepositoryName: !GetAtt ModelBuildCodeCommitRepository.Name
                BranchName: main
              OutputArtifacts:
                - Name: ModelBuildSourceArtifact

        - Name: Build
          Actions:
            - Name: BuildAndExecuteSageMakerPipeline
              ActionTypeId:
                Category: Build
                Owner: AWS
                Provider: CodeBuild
                Version: '1'
              InputArtifacts:
                - Name: ModelBuildSourceArtifact
              OutputArtifacts:
                - Name: ModelBuildBuildArtifact
              Configuration:
                ProjectName: !Ref ModelPipelineBuildProject
              RunOrder: 1