Description: Wait for a 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
  TrainJobRequest:
    Type: String
    Description: Training job request json
  KmsKeyId:
    Description: AWS KMS key ID used to encrypt data at rest on the ML storage volume attached to training job.
    Type: String

Resources:
  SagemakerTrainingJob:
    Type: Custom::TrainingJob
    Properties:
      ServiceToken: !Sub "arn:aws:lambda:${AWS::Region}:${AWS::AccountId}:function:${ProjectPrefix}-sagemaker-cfn-training-job"
      TrainingJobName: !Sub ${ProjectPrefix}-${ModelName}-${TrainJobId}
      TrainingJobRequest: !Ref TrainJobRequest
      ExperimentName: !Ref ModelName
      TrialName: !Ref TrainJobId
      KmsKeyId: !Ref KmsKeyId