apiVersion: "kubeflow.org/v1"
kind: "PyTorchJob"
metadata:
  name: "pytorch-dist-mnist-gloo"
spec:
  pytorchReplicaSpecs:
    Master:
      replicas: 1
      restartPolicy: OnFailure
      template:
        metadata:
          annotations:
            sidecar.istio.io/inject: "false"
            # if your libcuda.so.1 is in custom path, set the correct path with the following annotation
            # kubernetes.io/baidu-cgpu.nvidia-driver-lib: /usr/lib64
        spec:
          schedulerName: volcano
          containers:
            - name: pytorch
              image: registry.baidubce.com/cce-public/kubeflow/pytorch-dist-mnist-test-with-data:1.0
              args: ["--backend", "gloo"]
              # Comment out the below resources to use the CPU.
              resources:
                requests:
                  cpu: 1
                  memory: 1Gi
                limits:
                  baidu.com/v100_32g_cgpu: "1"
                  # for gpu core/memory isolation
                  baidu.com/v100_32g_cgpu_core: 10
                  baidu.com/v100_32g_cgpu_memory: "2"
              # if gpu core isolation is enabled, set the following preStop hook for graceful shutdown.
              # `mnist.py` needs to be replaced with the name of your gpu process.
              lifecycle:
                preStop:
                  exec:
                    command: [
                      "/bin/sh", "-c",
                      "kill -10 `ps -ef | grep mnist.py | grep -v grep | awk '{print $2}'` && sleep 1"
                    ]
    Worker:
      replicas: 1
      restartPolicy: OnFailure
      template:
        metadata:
          annotations:
            sidecar.istio.io/inject: "false"
        spec:
          schedulerName: volcano
          containers:
            - name: pytorch
              image: registry.baidubce.com/cce-public/kubeflow/pytorch-dist-mnist-test-with-data:1.0
              env:
                # for gpu memory over request, set 0 to disable
                - name: CGPU_MEM_ALLOCATOR_TYPE
                  value: “1”
              args: ["--backend", "gloo"]
              resources:
                requests:
                  cpu: 1
                  memory: 1Gi
                limits:
                  baidu.com/v100_32g_cgpu: "1"
                  # for gpu core/memory isolation
                  baidu.com/v100_32g_cgpu_core: 20
                  baidu.com/v100_32g_cgpu_memory: "4"
              # if gpu core isolation is enabled, set the following preStop hook for graceful shutdown.
              # `mnist.py` needs to be replaced with the name of your gpu process.
              lifecycle:
                preStop:
                  exec:
                    command: [
                      "/bin/sh", "-c",
                      "kill -10 `ps -ef | grep mnist.py | grep -v grep | awk '{print $2}'` && sleep 1"
                    ]