Description: Suggest baseline for training job
Parameters:
  ProjectPrefix:
    Type: String
    Description: |
      Makes resource privileges for services using this template scoped-limited.
      Changing the default must be done with care
    Default: PROJECT_PREFIX
  ModelName:
    Type: String
    Description: Name of the model
  TrainJobId:
    Type: String
    Description: Id of the Codepipeline + SagemakerJobs
  MLOpsRoleArn:
    Type: String
    Description: The role for executing the monitoring schedule
  BaselineInputUri:
    Type: String
    Description: The s3 uri for baseline input
  KmsKeyId:
    Description: AWS KMS key ID used to encrypt data at rest for S3 results.
    Type: String

Resources:
  SagemakerSuggestBaseline:
    Type: Custom::SuggestBaseline
    Properties:
      ServiceToken: !Sub arn:aws:lambda:${AWS::Region}:${AWS::AccountId}:function:${ProjectPrefix}-sagemaker-cfn-suggest-baseline
      ProcessingJobName: !Sub ${ProjectPrefix}-${ModelName}-pbl-${TrainJobId}
      BaselineInputUri: !Ref BaselineInputUri
      BaselineResultsUri: !Sub s3://sagemaker-${AWS::Region}-${AWS::AccountId}/${ModelName}/monitoring/baseline/${ProjectPrefix}-${ModelName}-pbl-${TrainJobId}
      KmsKeyId: !Ref KmsKeyId
      PassRoleArn: !Ref MLOpsRoleArn
      ExperimentName: !Ref ModelName
      TrialName: !Ref TrainJobId

Outputs:
  ProcessingJobName:
    Description: "The name of the baseline processing job"
    Value: !GetAtt SagemakerSuggestBaseline.ProcessingJobName