---
apiVersion: v1
kind: ConfigMap
metadata:
  name: trial-templates
  namespace: kubeflow
  labels:
    katib.kubeflow.org/component: trial-templates
data:
  defaultTrialTemplate.yaml: |-
    apiVersion: batch/v1
    kind: Job
    spec:
      template:
        spec:
          containers:
            - name: training-container
              image: ghcr.io/kubeflow/katib/pytorch-mnist-cpu:v0.18.0
              command:
                - "python3"
                - "/opt/pytorch-mnist/mnist.py"
                - "--epochs=1"
                - "--batch-size=16"
                - "--lr=${trialParameters.learningRate}"
                - "--momentum=${trialParameters.momentum}"
          restartPolicy: Never
  # For ConfigMap templates double quotes must set in commands to correct parse JSON parameters in Trial Template (e.g nn_config, architecture)
  enasCPUTemplate: |-
    apiVersion: batch/v1
    kind: Job
    spec:
      template:
        spec:
          containers:
            - name: training-container
              image: ghcr.io/kubeflow/katib/enas-cnn-cifar10-cpu:v0.18.0
              command:
                - python3
                - -u
                - RunTrial.py
                - --num_epochs=1
                - "--architecture=\"${trialParameters.neuralNetworkArchitecture}\""
                - "--nn_config=\"${trialParameters.neuralNetworkConfig}\""
          restartPolicy: Never
  pytorchJobTemplate: |-
    apiVersion: kubeflow.org/v1
    kind: PyTorchJob
    spec:
      pytorchReplicaSpecs:
        Master:
          replicas: 1
          restartPolicy: OnFailure
          template:
            spec:
              containers:
                - name: pytorch
                  image: ghcr.io/kubeflow/katib/pytorch-mnist-cpu:v0.18.0
                  command:
                    - "python3"
                    - "/opt/pytorch-mnist/mnist.py"
                    - "--epochs=1"
                    - "--lr=${trialParameters.learningRate}"
                    - "--momentum=${trialParameters.momentum}"
        Worker:
          replicas: 2
          restartPolicy: OnFailure
          template:
            spec:
              containers:
                - name: pytorch
                  image: ghcr.io/kubeflow/katib/pytorch-mnist-cpu:v0.18.0
                  command:
                    - "python3"
                    - "/opt/pytorch-mnist/mnist.py"
                    - "--epochs=1"
                    - "--lr=${trialParameters.learningRate}"
                    - "--momentum=${trialParameters.momentum}"