1apiVersion: "kubeflow.org/v1"
2kind: "MXJob"
3metadata:
4  name: "mxnet-job"
5spec:
6  jobMode: MXTrain
7  mxReplicaSpecs:
8    Scheduler:
9      replicas: 1
10      restartPolicy: Never
11      template:
12        metadata:
13          annotations:
14            sidecar.istio.io/inject: "false"
15            # if your libcuda.so.1 is in custom path, set the correct path with the following annotation
16            # kubernetes.io/baidu-cgpu.nvidia-driver-lib: /usr/lib64
17        spec:
18          schedulerName: volcano
19          containers:
20            - name: mxnet
21              image: registry.baidubce.com/cce-public/kubeflow/cce-public/mxjob/mxnet:gpu
22              resources:
23                limits:
24                  baidu.com/v100_32g_cgpu: "1"
25                  # for gpu core/memory isolation
26                  baidu.com/v100_32g_cgpu_core: 5
27                  baidu.com/v100_32g_cgpu_memory: "1"
28              # if gpu core isolation is enabled, set the following preStop hook for graceful shutdown.
29              # ${'`'}train_mnist.py${'`'} needs to be replaced with the name of your gpu process.
30              lifecycle:
31                preStop:
32                  exec:
33                    command: [
34                      "/bin/sh", "-c",
35                      "kill -10 ${'`'}ps -ef | grep train_mnist.py | grep -v grep | awk '{print $2}'${'`'} && sleep 1"
36                    ]
37    Server:
38      replicas: 1
39      restartPolicy: Never
40      template:
41        metadata:
42          annotations:
43            sidecar.istio.io/inject: "false"
44            # if your libcuda.so.1 is in custom path, set the correct path with the following annotation
45            # kubernetes.io/baidu-cgpu.nvidia-driver-lib: /usr/lib64
46        spec:
47          schedulerName: volcano
48          containers:
49            - name: mxnet
50              image: registry.baidubce.com/cce-public/kubeflow/cce-public/mxjob/mxnet:gpu
51              resources:
52                limits:
53                  baidu.com/v100_32g_cgpu: "1"
54                  # for gpu core/memory isolation
55                  baidu.com/v100_32g_cgpu_core: 5
56                  baidu.com/v100_32g_cgpu_memory: "1"
57              # if gpu core isolation is enabled, set the following preStop hook for graceful shutdown.
58              # ${'`'}train_mnist.py${'`'} needs to be replaced with the name of your gpu process.
59              lifecycle:
60                preStop:
61                  exec:
62                    command: [
63                      "/bin/sh", "-c",
64                      "kill -10 ${'`'}ps -ef | grep train_mnist.py | grep -v grep | awk '{print $2}'${'`'} && sleep 1"
65                    ]
66    Worker:
67      replicas: 1
68      restartPolicy: Never
69      template:
70        metadata:
71          annotations:
72            sidecar.istio.io/inject: "false"
73            # if your libcuda.so.1 is in custom path, set the correct path with the following annotation
74            # kubernetes.io/baidu-cgpu.nvidia-driver-lib: /usr/lib64
75        spec:
76          schedulerName: volcano
77          containers:
78            - name: mxnet
79              image: registry.baidubce.com/cce-public/kubeflow/cce-public/mxjob/mxnet:gpu
80              env:
81                # for gpu memory over request, set 0 to disable
82                - name: CGPU_MEM_ALLOCATOR_TYPE
83                  value: “1”
84              command: ["python"]
85              args: [
86                "/incubator-mxnet/example/image-classification/train_mnist.py",
87                "--num-epochs","10","--num-layers","2","--kv-store","dist_device_sync","--gpus","0"
88              ]
89              resources:
90                requests:
91                  cpu: 1
92                  memory: 1Gi
93                limits:
94                  baidu.com/v100_32g_cgpu: "1"
95                  # for gpu core/memory isolation
96                  baidu.com/v100_32g_cgpu_core: 20
97                  baidu.com/v100_32g_cgpu_memory: "4"
98              # if gpu core isolation is enabled, set the following preStop hook for graceful shutdown.
99              # ${'`'}train_mnist.py${'`'} needs to be replaced with the name of your gpu process.
100              lifecycle:
101                preStop:
102                  exec:
103                    command: [
104                      "/bin/sh", "-c",
105                      "kill -10 ${'`'}ps -ef | grep train_mnist.py | grep -v grep | awk '{print $2}'${'`'} && sleep 1"
106                    ]