mxjob-byteps.yaml

---
apiVersion: kubeflow.org/v1beta1
kind: Experiment
metadata:
  namespace: kubeflow
  name: mxjob-byteps
spec:
  objective:
    type: maximize
    goal: 0.99
    objectiveMetricName: Train-accuracy
  algorithm:
    algorithmName: random
  parallelTrialCount: 1
  maxTrialCount: 4
  maxFailedTrialCount: 3
  parameters:
    - name: lr
      parameterType: double
      feasibleSpace:
        min: "0.1"
        max: "0.11"
  trialTemplate:
    primaryContainerName: mxnet
    # In this example we can collect metrics only from the Worker pods.
    primaryPodLabels:
      training.kubeflow.org/replica-type: worker
    trialParameters:
      - name: learningRate
        description: Learning rate for the training model
        reference: lr
    trialSpec:
      apiVersion: kubeflow.org/v1
      kind: MXJob
      spec:
        jobMode: MXTrain
        runPolicy:
          cleanPodPolicy: None
        mxReplicaSpecs:
          Scheduler:
            replicas: 1
            restartPolicy: Never
            template:
              spec:
                containers:
                  - name: mxnet
                    image: docker.io/bytepsimage/mxnet
                    command: ["bpslaunch"]
          Server:
            replicas: 1
            restartPolicy: Never
            template:
              spec:
                containers:
                  - name: mxnet
                    image: docker.io/bytepsimage/mxnet
                    command: ["bpslaunch"]
          Worker:
            replicas: 1
            restartPolicy: Never
            template:
              spec:
                containers:
                  - name: mxnet
                    image: docker.io/bytepsimage/mxnet
                    command: ["bpslaunch"]
                    args:
                      [
                        "python3",
                        "/usr/local/byteps/example/mxnet/train_imagenet_byteps.py",
                        "--benchmark",
                        "1",
                        "--lr=${trialParameters.learningRate}",
                        "--num-examples=1000",
                        "--num-epochs=4",
                      ]
                    volumeMounts:
                      - mountPath: /dev/shm
                        name: dshm
                    resources:
                      limits:
                        nvidia.com/gpu: 1
                volumes:
                  - name: dshm
                    emptyDir:
                      medium: Memory