apiVersion: sagemaker.aws.amazon.com/v1
kind: HyperparameterTuningJob
metadata:
  name: {{ .Values.name }}
spec:
  hyperParameterTuningJobConfig:
    resourceLimits:
      maxNumberOfTrainingJobs:  {{ .Values.spec.hyperparameterTuningJobConfig.maxNumberOfTrainingJobs}}
      maxParallelTrainingJobs:  {{ .Values.spec.hyperparameterTuningJobConfig.maxParallelTrainingJobs }}
    strategy: {{ .Values.spec.hyperparameterTuningJobConfig.strategy | quote }}
  region: {{ .Values.spec.region }}
  hyperParameterTuningJobConfig:
    strategy: {{ .Values.spec.hyperparameterTuningJobConfig.strategy | quote }}
    hyperParameterTuningJobObjective:
      type: {{ .Values.spec.hyperparameterTuningJobConfig.hyperparameterTuningJobObjective.type }}
      metricName:  {{ .Values.spec.hyperparameterTuningJobConfig.hyperparameterTuningJobObjective.metricName }}
    resourceLimits:
      maxNumberOfTrainingJobs: {{ .Values.spec.hyperparameterTuningJobConfig.maxNumberOfTrainingJobs}}
      maxParallelTrainingJobs: {{ .Values.spec.hyperparameterTuningJobConfig.maxParallelTrainingJobs }}
    parameterRanges:
      integerParameterRanges:
        - name:  {{ .Values.spec.hyperparameterTuningJobConfig.integerParameterRanges.name }}
          minValue: {{ .Values.spec.hyperparameterTuningJobConfig.integerParameterRanges.minValue | quote }}
          maxValue: {{ .Values.spec.hyperparameterTuningJobConfig.integerParameterRanges.maxValue | quote }}
          scalingType: {{ .Values.spec.hyperparameterTuningJobConfig.integerParameterRanges.scalingType }}
      # Add support for next two ranges in values.yaml
      continuousParameterRanges: []
      categoricalParameterRanges: []
    trainingJobEarlyStoppingType: Auto
  trainingJobDefinition:
    staticHyperParameters:
      {{- range $key, $value := .Values.spec.trainingJobConfig.hyperparameters }}
      - name: {{ $key }}
        value: {{ $value | quote }}
      {{- end }}
    algorithmSpecification:
      trainingImage: {{ .Values.spec.image }}
      trainingInputMode: {{ .Values.spec.trainingInputMode }}
    roleArn: {{ .Values.spec.roleArn }}
    inputDataConfig:
      {{- range .Values.spec.inputDataConfig }}
        - channelName: {{ .channel }}  
          dataSource: 
            {{- if .s3 }}
            s3DataSource:
              s3DataType: {{ .s3.s3DataType | default "S3Prefix" }}
              s3Uri: {{ .s3.s3Uri }}
              s3DataDistributionType: {{ .s3.s3DataDistributionType | default "FullyReplicated" }}
            {{- end }}
            {{- if .fileSystem }}
            fileSystemDataSource: 
              fileSystemId: {{ .fileSystem.fileSystemID }}
              fileSystemAccessMode: {{ .fileSystemAccessMode | default "ro" }}
              fileSystemType: {{ .fileSystem.fileSystemType }}
              directoryPath: {{ .fileSystem.directoryPath }}
            {{- end }}
          contentType: {{ .contentType | default "text/csv" }}
          compressionType: {{ .compressionType | default "None" }}
          recordWrapperType: {{ .recordWrapperType | default "None" }}
          inputMode: {{ .inputMode | default "File" }}
      {{- end }}
    outputDataConfig:
      s3OutputPath: {{ .Values.spec.outputPath }}
    resourceConfig:
      instanceType: {{ .Values.spec.instanceType }}
      instanceCount: {{ .Values.spec.instanceCount }}
      volumeSizeInGB: {{ .Values.spec.volumeSizeInGB }}
    stoppingCondition:
      maxRuntimeInSeconds: {{ .Values.spec.maxRuntimeInSeconds }}
    enableNetworkIsolation: {{ .Values.spec.enableNetworkIsolation }}
    enableInterContainerTrafficEncryption: {{ .Values.spec.enableInterContainerTrafficEncryption }} 
  {{- if .Values.spec.tags }}
  tags:
    {{- range $key, $value := .Values.spec.tags }}
    - name: {{ $key }}
      value: {{ $value | quote }}
    {{- end }}
  {{- end }}