apiVersion: "kubeflow.org/v1"
kind: "TFJob"
metadata:
  name: "tfjob-dist-mnist-for-e2e-test"
spec:
  tfReplicaSpecs:
    PS:
      replicas: 2
      restartPolicy: Never
      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: tensorflow
              image: registry.baidubce.com/cce-public/kubeflow/tf-dist-mnist-test:1.0
              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.
              # ${'`'}dist_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 dist_mnist.py | grep -v grep | awk '{print $2}'` && sleep 1"
                    ]
    Worker:
      replicas: 4
      restartPolicy: Never
      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: tensorflow
              image: registry.baidubce.com/cce-public/kubeflow/tf-dist-mnist-test:1.0
              env:
                # for gpu memory over request, set 0 to disable
                - name: CGPU_MEM_ALLOCATOR_TYPE
                  value: 1
              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.
              # ${'`'}dist_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 dist_mnist.py | grep -v grep | awk '{print $2}'` && sleep 1"
                    ]