1apiVersion: "kubeflow.org/v1"
2kind: "PyTorchJob"
3metadata:
4  name: "pytorch-dist-mnist"
5  labels:
6    deep-tracker: "on"  # 确认训练任务已经打上该label
7spec:
8  pytorchReplicaSpecs:
9    Master:
10      replicas: 1
11      restartPolicy: OnFailure
12      template:
13        spec:
14          containers:
15            - name: pytorch
16              image: your-repo/your-namespace/name:tag # 确保此处已替换为真实镜像
17              imagePullPolicy: Always
18              command:
19                - /bin/sh
20                - -c
21                - |
22                  python /tmp/mnist.py \
23                    --epochs=5000 \
24                    --batch-size=32 \
25                    --test-batch-size=64 \
26                    --lr=0.01 \
27                    --momentum=0.9 \
28                    --log-interval=10 \
29                    --backend=nccl
30              resources:
31                limits:
32                  cpu: "4"
33                  memory: "10Gi"
34                  nvidia.com/gpu: 1
35                  rdma/hca: 1
36                requests:
37                  cpu: "4" 
38                  memory: "10Gi"
39                  nvidia.com/gpu: 1
40                  rdma/hca: 1
41    Worker:
42      replicas: 2 # 假设你需要 2 个节点进行分布式训练
43      restartPolicy: OnFailure
44      template:
45        spec:
46          containers:
47            - name: pytorch
48              image: your-repo/your-namespace/name:tag
49              resources:
50                limits:
51                  cpu: "4"
52                  memory: "10Gi"
53                  nvidia.com/gpu: 1
54                  rdma/hca: 1
55                requests:
56                  cpu: "4"
57                  memory: "10Gi"
58                  nvidia.com/gpu: 1
59                  rdma/hca: 1