apiVersion: v1
kind: ConfigMap
metadata:
  name: trial-template
  namespace: kubeflow
  labels:
    app: katib-trial-templates
data:
  defaultTrialTemplate.yaml: |-
    apiVersion: batch/v1
    kind: Job
    spec:
      template:
        spec:
          containers:
            - name: training-container
              image: docker.io/kubeflowkatib/mxnet-mnist:v1beta1-45c5727
              command:
                - "python3"
                - "/opt/mxnet-mnist/mnist.py"
                - "--batch-size=64"
                - "--lr=${trialParameters.learningRate}"
                - "--num-layers=${trialParameters.numberLayers}"
                - "--optimizer=${trialParameters.optimizer}"
          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: docker.io/kubeflowkatib/enas-cnn-cifar10-cpu:v1beta1-45c5727
              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: docker.io/kubeflowkatib/pytorch-mnist:v1beta1-45c5727
                  imagePullPolicy: Always
                  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: docker.io/kubeflowkatib/pytorch-mnist:v1beta1-45c5727
                  imagePullPolicy: Always
                  command:
                    - "python3"
                    - "/opt/pytorch-mnist/mnist.py"
                    - "--epochs=1"
                    - "--lr=${trialParameters.learningRate}"
                    - "--momentum=${trialParameters.momentum}"