Description: Template for sagemaker lambda and cloud formation custom resources
Transform: AWS::Serverless-2016-10-31
Resources:
  AddTransformHeaderFunction:
    Type: AWS::Serverless::Function
    Properties:
      FunctionName: mlops-add-transform-header
      CodeUri: .
      Handler: sagemaker_add_transform_header.lambda_handler
      Runtime: python3.7
      Role: !GetAtt SagemakerCustomResourceRole.Arn
      Description: "Prepend header to a batch transform job"

  CreateExperimentFunction:
    Type: AWS::Serverless::Function
    Properties:
      FunctionName: mlops-create-experiment
      CodeUri: .
      Handler: sagemaker_create_experiment.lambda_handler
      Runtime: python3.7
      Role: !GetAtt SagemakerCustomResourceRole.Arn
      Description: "Create a SageMaker experiment and trial"

  QueryDriftFunction:
    Type: AWS::Serverless::Function
    Properties:
      FunctionName: mlops-query-drift
      CodeUri: .
      Handler: sagemaker_query_drift.lambda_handler
      Runtime: python3.7
      Role: !GetAtt SagemakerCustomResourceRole.Arn
      Description: "Query processing job to return drift"

  QueryTrainingFunction:
    Type: AWS::Serverless::Function
    Properties:
      FunctionName: mlops-query-training
      CodeUri: .
      Handler: sagemaker_query_training.lambda_handler
      Runtime: python3.7
      Role: !GetAtt SagemakerCustomResourceRole.Arn
      Description: "Query training job to return results"

  SagemakerCustomResourceRole:
    Type: AWS::IAM::Role
    Properties:
      RoleName: sagemaker-cfn-custom-resource
      AssumeRolePolicyDocument:
        Statement:
          - Action:
              - sts:AssumeRole
            Effect: Allow
            Principal:
              Service:
                - lambda.amazonaws.com
        Version: "2012-10-17"
      ManagedPolicyArns:
        - "arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole"
      Policies:
        - PolicyDocument:
            Statement:
              - Sid: AllowSageMaker
                Effect: Allow
                Action:
                  - sagemaker:CreateExperiment
                  - sagemaker:CreateTrial
                  - sagemaker:CreateTrainingJob
                  - sagemaker:DescribeTrainingJob
                  - sagemaker:StopTrainingJob
                  - sagemaker:DescribeEndpoint
                  - sagemaker:UpdateEndpoint
                  - sagemaker:CreateEndpointConfig
                  - sagemaker:DescribeEndpointConfig
                  - sagemaker:DeleteEndpointConfig
                  - sagemaker:DescribeProcessingJob
                  - sagemaker:CreateProcessingJob
                  - sagemaker:StopProcessingJob
                  - kms:CreateGrant # Required if KmsKeyId specified
                Resource:
                  - !Sub arn:aws:sagemaker:${AWS::Region}:${AWS::AccountId}:*/*
              - Sid: S3Resources
                Effect: Allow
                Action:
                  - s3:GetObject*
                  - s3:PutObject
                Resource:
                  - !Sub arn:aws:s3:::sagemaker-${AWS::Region}-${AWS::AccountId}/*
                  - !Sub arn:aws:s3:::sagemaker-${AWS::Region}-${AWS::AccountId}
              - Sid: AllowLambda
                Effect: Allow
                Action:
                  - lambda:*
                Resource: "*"
              - Sid: AllowEvents
                Effect: Allow
                Action:
                  - events:* # Requires at least events:PutRule/events:RemoveTargets
                Resource: "*"
              - Sid: AllowPassRole
                Effect: Allow
                Action:
                  - iam:PassRole
                Resource: "*"
                Condition:
                  StringEquals:
                    iam:PassedToService: sagemaker.amazonaws.com
            Version: "2012-10-17"
          PolicyName: SagemakerCustomResource