Description: Toolchain template which provides the resources needed to represent infrastructure as code.
  This template specifically creates a CI/CD pipeline to deploy a given inference image and pretrained Model to two stages in CD -- 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.

  SageMakerExecutionRoleStagingName:
    Type: String
    Description: Name of the role created in the Staging account used by SageMaker to deploy the model
    MinLength: 1
    MaxLength: 64
    AllowedPattern: ^[a-zA-Z](-*[a-zA-Z0-9])*
  SageMakerExecutionRoleProdName:
    Type: String
    Description: Name of the role created in the Prod account used by SageMaker to deploy the model
    MinLength: 1
    MaxLength: 64
    AllowedPattern: ^[a-zA-Z](-*[a-zA-Z0-9])*
        
  OrganizationalUnitStagingId:
    Type: String
    Description: Id of the organizational unit that holds the staging account
  OrganizationalUnitProdId:
    Type: String
    Description: Id of the organizational unit that holds the prod account
    

Resources:
  MlOpsArtifactsBucket:
    Type: AWS::S3::Bucket
    DeletionPolicy: Retain
    Properties:
      BucketName: !Sub sagemaker-project-${SageMakerProjectId} # 58 chars max/ 64 allowed

  ModelDeploySageMakerEventRule:
    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:
            - !Ref SageMakerProjectName
      State: "ENABLED"
      Targets:
        -
          Arn:
            !Join [ ':', [ 'arn', !Ref 'AWS::Partition', 'codepipeline', !Ref 'AWS::Region', !Ref 'AWS::AccountId', !Ref ModelDeployPipeline ] ]
          RoleArn:
            !Join [ ':', [ 'arn', !Ref 'AWS::Partition', 'iam:', !Ref '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+4=62 chars
      Description: "Rule to trigger a deployment when CodeCommit is updated with a commit"
      EventPattern:
        source:
          - "aws.codecommit"
        detail-type:
          - "CodeCommit Repository State Change"
        resources:
          - !GetAtt ModelDeployCodeCommitRepository.Arn
        detail:
          referenceType:
            - "branch"
          referenceName:
            - "main"
      State: "ENABLED"
      Targets:
        -
          Arn:
            !Join [ ':', [ 'arn', !Ref 'AWS::Partition', 'codepipeline', !Ref 'AWS::Region', !Ref 'AWS::AccountId', !Ref ModelDeployPipeline ] ]
          RoleArn:
            !Join [ ':', [ 'arn', !Ref 'AWS::Partition', 'iam:', !Ref 'AWS::AccountId', 'role/service-role/AmazonSageMakerServiceCatalogProductsUseRole'] ]
          Id: !Sub codecommit-${SageMakerProjectName}-trigger

  ModelDeployCodeCommitRepository:
    Type: AWS::CodeCommit::Repository
    Properties:
      # Max allowed length: 100 chars
      RepositoryName: !Sub sagemaker-${SageMakerProjectName}-${SageMakerProjectId}-modeldeploy # max: 10+33+15+11=69
      RepositoryDescription: !Sub SageMaker Endpoint deployment infrastructure as code for the Project ${SageMakerProjectName}
      Code:
        S3:
          Bucket: aws-ml-blog 
          Key: artifacts/sagemaker-pipeline-blog-resources/multi-account-deployment-v1.0.zip
        BranchName: main

  ModelDeployBuildProject:
    Type: AWS::CodeBuild::Project
    Properties:
      # Max length: 255 chars
      Name: !Sub sagemaker-${SageMakerProjectName}-${SageMakerProjectId}-modeldeploy # max: 10+33+15+11=69
      Description: Builds the Cfn template which defines the Endpoint with specified configuration
      ServiceRole:
        !Join [ ':', [ 'arn', !Ref 'AWS::Partition', 'iam:', !Ref '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: ARTIFACT_BUCKET
           Value: !Ref MlOpsArtifactsBucket
         - Name: SOURCE_MODEL_PACKAGE_GROUP_NAME
           Value: !Ref SageMakerProjectName
         - Name: AWS_REGION
           Value: !Ref AWS::Region
         # these values are used by the build system to output to the output artifacts.
         # further down, we use these names in the Cfn deployment steps
         - Name: EXPORT_TEMPLATE_NAME
           Value: template-export.yml
         - Name: EXPORT_TEMPLATE_STAGING_CONFIG
           Value: staging-config-export.json
         - Name: EXPORT_TEMPLATE_PROD_CONFIG
           Value: prod-config-export.json
         - Name: SAGEMAKER_EXECUTION_ROLE_STAGING_NAME
           Value: !Ref SageMakerExecutionRoleStagingName
         - Name: SAGEMAKER_EXECUTION_ROLE_PROD_NAME
           Value: !Ref SageMakerExecutionRoleProdName
         - Name: ORGANIZATIONAL_UNIT_STAGING_ID
           Value: !Ref OrganizationalUnitStagingId
         - Name: ORGANIZATIONAL_UNIT_PROD_ID
           Value: !Ref OrganizationalUnitProdId

      Source:
        Type: CODEPIPELINE
        BuildSpec: buildspec.yml
      TimeoutInMinutes: 30

  ModelDeployTestProject:
    Type: AWS::CodeBuild::Project
    Properties:
      # Max length: 255 chars
      Name: !Sub sagemaker-${SageMakerProjectName}-${SageMakerProjectId}-testing # max: 10+33+15+7=65
      Description: Test the deployment endpoint
      ServiceRole:
        !Join [ ':', [ 'arn', !Ref 'AWS::Partition', 'iam:', !Ref '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: AWS_REGION
            Value: !Ref "AWS::Region"
          - Name: BUILD_CONFIG
            Value: staging-config-export.json
          - Name: EXPORT_TEST_RESULTS
            Value: test-results.json
          - Name: SAGEMAKER_EXECUTION_ROLE_NAME
            Value: !Ref SageMakerExecutionRoleStagingName
          - Name: ORGANIZATIONAL_UNIT_ID
            Value: !Ref OrganizationalUnitStagingId
      Source:
        Type: CODEPIPELINE
        BuildSpec: test/buildspec.yml
      TimeoutInMinutes: 30

  ModelDeployPipeline:
    Type: AWS::CodePipeline::Pipeline
    DependsOn: MlOpsArtifactsBucket
    Properties:
      # Max length: 100 chars
      Name: !Sub sagemaker-${SageMakerProjectName}-${SageMakerProjectId}-modeldeploy # max: 10+33+15+11=69
      RoleArn:
        !Join [ ':', [ 'arn', !Ref 'AWS::Partition', 'iam:', !Ref 'AWS::AccountId', 'role/service-role/AmazonSageMakerServiceCatalogProductsUseRole'] ]
      ArtifactStore:
        Type: S3
        Location:
          !Ref MlOpsArtifactsBucket
      Stages:
        - Name: Source
          Actions:
            - Name: ModelDeployInfraCode
              ActionTypeId:
                Category: Source
                Owner: AWS
                Provider: CodeCommit
                Version: 1
              Configuration:
                # need to explicitly set this to false per https://docs.aws.amazon.com/codepipeline/latest/userguide/update-change-detection.html
                PollForSourceChanges: false
                RepositoryName: !GetAtt ModelDeployCodeCommitRepository.Name
                BranchName: main
              OutputArtifacts:
                - Name: SourceArtifact

        - Name: Build
          Actions:
            - Name: BuildDeploymentTemplates
              ActionTypeId:
                Category: Build
                Owner: AWS
                Provider: CodeBuild
                Version: 1
              InputArtifacts:
                - Name: SourceArtifact
              OutputArtifacts:
                - Name: BuildArtifact
              Configuration:
                ProjectName: !Ref ModelDeployBuildProject
              RunOrder: 1

        - Name: DeployResourcesStaging
          Actions:
            - Name: DeployStaging 
              InputArtifacts:
                - Name: BuildArtifact
              ActionTypeId:
                Category: Deploy
                Owner: AWS
                Version: 1
                Provider: CloudFormationStackSet
              Configuration:
                Capabilities: CAPABILITY_NAMED_IAM
                PermissionModel: SERVICE_MANAGED
                OrganizationsAutoDeployment: Disabled
                StackSetName: !Sub sagemaker-${SageMakerProjectName}-${SageMakerProjectId}-deploy-staging  #10+33+15+14=72 out of 128 max
                Parameters: BuildArtifact::staging-config-export.json
                # The buildspec.yml in the application stack uses this file name,
                TemplatePath: BuildArtifact::template-export.yml
                DeploymentTargets: !Ref OrganizationalUnitStagingId
                Regions: !Ref 'AWS::Region'
              RunOrder: 1
            - Name: TestStaging
              ActionTypeId:
                Category: Build
                Owner: AWS
                Provider: CodeBuild
                Version: 1
              InputArtifacts:
                - Name: SourceArtifact
                - Name: BuildArtifact
              OutputArtifacts:
                - Name: TestArtifact
              Configuration:
                ProjectName: !Ref ModelDeployTestProject
                PrimarySource: SourceArtifact
              RunOrder: 2
            - Name: ApproveDeployment
              ActionTypeId:
                Category: Approval
                Owner: AWS
                Version: 1
                Provider: Manual
              Configuration:
                CustomData: "Approve this model for Production"
              RunOrder: 3
        - Name: DeployResourcesProd
          Actions:
            - Name: DeployProd
              InputArtifacts:
                - Name: BuildArtifact
              ActionTypeId:
                Category: Deploy
                Owner: AWS
                Version: 1
                Provider: CloudFormationStackSet
              Configuration:
                Capabilities: CAPABILITY_NAMED_IAM
                PermissionModel: SERVICE_MANAGED
                OrganizationsAutoDeployment: Disabled
                StackSetName: !Sub sagemaker-${SageMakerProjectName}-${SageMakerProjectId}-deploy-prod  #10+33+15+14=72 out of 128 max
                Parameters: BuildArtifact::prod-config-export.json
                TemplatePath: BuildArtifact::template-export.yml
                DeploymentTargets: !Ref OrganizationalUnitProdId
                Regions: !Ref 'AWS::Region'
              RunOrder: 1

