AWSTemplateFormatVersion: "2010-09-09"
Transform: AWS::Serverless-2016-10-31
Description: >
  A Sagemaker Algorithm Resource as a Cloudformation stack.

# Metadata associated with this stack.
Metadata:
  # Cloudformation interface for parameters.
  AWS::CloudFormation::Interface:
    ParameterGroups:
      # General parameters label.
      - Label:
          default: General Parameters
        Parameters:
          - Environment
          - AlgorithmName
          - TrainingImageURI
          - InferenceImageURI
    
    # Labels for the above parameters.
    ParameterLabels:
      Environment:
        default: Environment name
      AlgorithmName:
        default: Algorithm name
      TrainingImageURI:
        default: Training Image URI
      InferenceImageURI:
        default: Inference Image URI
     
# Parameters exposed by this template.
Parameters:
  # General parameters.
  Environment:
    Type: String
    Description: >
      The environment name on which you would like to deploy the project.
      This identifier will be used to tag created resources.
    Default: development
    MinLength: 1
    ConstraintDescription: The environment cannot be empty.
  AlgorithmName:
    Type: String
    Description: >
      The unique algorithm name on which you would like to create on Sagemaker.
    Default: h2o-gbm-algorithm
    MinLength: 1
    ConstraintDescription: The algorithm name cannot be empty.
  TrainingImageURI:
    Type: String
    Description: >
      URI of Training Image on Amazon ECR.
    MinLength: 1
    AllowedPattern: ".*.dkr.ecr..*.amazonaws.com/.*$"
    ConstraintDescription: The Training Image URI cannot be empty. Please set an URI in this account & region.
  InferenceImageURI:
    Type: String
    Description: >
      URI of Inference Image on Amazon ECR.
    MinLength: 1
    AllowedPattern: ".*.dkr.ecr..*.amazonaws.com/.*$"
    ConstraintDescription: The Inference Image URI cannot be empty. Please set an URI in this account & region.

    # Parameters exposed by this template.
Resources: 

  # Lambda Execution IAM Role for 'CreateAlgorithmLambdaIAMRole'
  CreateAlgorithmLambdaIAMRole:
    Type: 'AWS::IAM::Role'
    Properties:
      AssumeRolePolicyDocument:
        Version: 2012-10-17
        Statement:
          - Effect: Allow
            Principal:
              Service:
                - lambda.amazonaws.com
            Action:
              - 'sts:AssumeRole'
      Path: /
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole
      Policies:
        - PolicyName: AmazonSagemakerExecutionRolePolicy
          PolicyDocument:
            Version: "2012-10-17"
            Statement:
              - Effect: Allow
                Action:
                  - sagemaker:CreateAlgorithm
                  - sagemaker:DeleteAlgorithm
                Resource: 
                  - !Sub 'arn:aws:sagemaker:${AWS::Region}:${AWS::AccountId}:algorithm/*'
        - PolicyName: AmazonECRExecutionRolePolicy
          PolicyDocument:
            Version: "2012-10-17"
            Statement:
              - Effect: Allow
                Action:
                  - ecr:BatchCheckLayerAvailability
                  - ecr:BatchGetImage
                Resource:
                  - !Join
                    - ''
                    - - !Sub 'arn:aws:ecr:${AWS::Region}:${AWS::AccountId}:repository/'
                      - !Select [ "1", !Split ['/', !Ref TrainingImageURI ]]
                  - !Join
                    - ''
                    - - !Sub 'arn:aws:ecr:${AWS::Region}:${AWS::AccountId}:repository/'
                      - !Select [ "1", !Split ['/', !Ref InferenceImageURI ]]
  
  # Adds Create Algorithm Lambda
  CreateAlgorithmLambda:
    Type: 'AWS::Serverless::Function'
    Properties:
      CodeUri: lambdas/create-algorithm
      Handler: index.lambda_handler
      Role: !GetAtt CreateAlgorithmLambdaIAMRole.Arn
      Runtime: python3.8
      Timeout: 50
      Tags: 
        Name: !Ref "AWS::StackName"
        Environment: !Ref Environment
 
  # Executes Create Algorithm Custom Resource Lambda
  CreateAlgorithmCustomResource:
    Type: Custom::AppConfiguration
    Properties:
      ServiceToken: !GetAtt CreateAlgorithmLambda.Arn
      AlgorithmName: !Ref AlgorithmName
      TrainingImageURI: !Ref TrainingImageURI
      InferenceImageURI: !Ref InferenceImageURI

# The outputs to be generated by this template.
Outputs:
  Name:
    Description: >
      ML Parameter Provider Stack Name.
    Value: !Ref AWS::StackName
    Export:
      Name: !Sub ${AWS::StackName}-Name
  CreateAlgorithmLambdaArn:
    Description: >
      ARN for the Create Algorithm Custom Resource Lambda.
    Value: !GetAtt CreateAlgorithmLambda.Arn
    Export:
      Name: !Sub ${AWS::StackName}-CreateAlgorithmLambdaArn
  AlgorithmArn:
    Description: >
      ARN for the Algorithm Resource.
    Value: !Sub 'arn:aws:sagemaker:${AWS::Region}:${AWS::AccountId}:algorithm/${AlgorithmName}'
    Export:
      Name: !Sub ${AWS::StackName}-AlgorithmArn