TrainingOperator--PyTorchJob实现机制分析

前言

由 Pytorch分布式训练（一）_chenxy02的博客-CSDN博客 可知Pytorch分布式训练实现进程间寻址，主要依靠以下 四个参数:

• MASTER_ADDR
• MASTER_PORT
• WORLD_SIZE
• RANK

MASTER_PORT和MASTER_ADDR的目的是告诉进程组中负责进程通信协调的核心进程的IP地址和端口。 RANK参数是该进程的id，WORLD_SIZE是说明进程组中进程的个数。
 

从上一篇博客我们可以想到 手工启动分布式训练 有以下缺点：

1、得手工在多个节点上启动多个 python脚本，配置不同的 rank等参数

2、如果是在K8S上起多个 pod 进行分布式训练 ，还得解决IP不确定的问题

3、难以满足用户想在特定情景下动态伸缩节点数的要求

带着这三个问题，我们来解析一下 PyTorchJob (training-operator实现的CRD之一) 的实现机制。

Training-operator 介绍

代码仓库：GitHub - kubeflow/training-operator: Training operators on Kubernetes.

如下，Training-operator 提供各位了 各种 K8S 的 custom resources。 使得在 K8S 上进行 TensorFlow/PyTorch/Apache MXNet/XGBoost/MPI/Paddle 框架的分布式训练变得容易。

Training-operator 安装

由 training-operator 在 v1.6 之后才支持 PaddlePaddle，这里我们安装 v1.6.0 版本。

# 在线安装
kubectl apply -k "github.com/kubeflow/training-operator/manifests/overlays/standalone?ref=v1.6.0"

# 离线安装则提前下好training-operator的项目代码，以及镜像 kubeflow/training-operator:v1-5a5f92d

检查安装结果

PytorchJob 入门使用

apply 一个 kind 为 PyTorchJob 的 yaml ，如下：

kubectl apply -n kubeflow -f <<EOF
apiVersion: "kubeflow.org/v1"
kind: PyTorchJob
metadata:
  name: pytorch-simple-001
  namespace: kubeflow
spec:
  pytorchReplicaSpecs:
    Master:
      replicas: 1
      restartPolicy: OnFailure
      template:
        spec:
          containers:
            - name: pytorch
              image: kubeflowkatib/pytorch-mnist:v1beta1-45c5727
              imagePullPolicy: Always
              command:
                - "python3"
                - "/opt/pytorch-mnist/mnist.py"
                - "--epochs=1"
    Worker:
      replicas: 2
      restartPolicy: OnFailure
      template:
        spec:
          containers:
            - name: pytorch
              image: kubeflowkatib/pytorch-mnist:v1beta1-45c5727
              imagePullPolicy: Always
              command:
                - "python3"
                - "/opt/pytorch-mnist/mnist.py"
                - "--epochs=1"
EOF

正常情况下，我们会看到三个running的pod，如下：

 训练完成后，Pod 的状态变成 Completed，如下：

PytorchJob 源码分析

关于 traning-operator 所使用的脚手架工具 kubebuilder 的代码结构可参见 Kubeflow--TFJob实现机制分析_chenxy02的博客-CSDN博客 ， 这里我们直接看PyTorchJob的调谐方法 Reconcile()

代码地址：pkg/controller.v1/pytorch/pytorchjob_controller.go

其中，r.ReconcileHPA(pytorchjob) 会根据 pytorchjob.spec.elasticPolicy 字段，对训练任务进行动态扩缩容。利用此，也便可以解决 上述第三个问题——“在特定情景下动态伸缩节点数”。

接下来 我们进到 r.ReconcileJobs(pytorchJob, ……)方法，看一下 PyTorchJob 怎么解决上述前两个问题。在 ReconcileJobs() 方法中调用了 ReconcilePods()方法负责调谐出Pod，pod的数量由 pytorchjob.spec.pytorchReplicaSpecs 决定。

代码地址：github.com/kubeflow/common/pkg/controller.v1/common/pod.go

进入到 jc.createNewPod() 可以发现，training-operator是通过 jc.Controller.SetClusterSpec(）方法，根据不同分布式训练框架的需要，为训练容器注入相应的配置。

代码地址：github.com/kubeflow/common/pkg/controller.v1/common/pod.go

代码地址：pkg/controller.v1/pytorch/pytorchjob_controller.go

进入在setPodEnv() 便可以看到 training-operator是怎么为pytorch分布式训练的容器注入所需的 MASTER_PORT、MASTER_ADDR、WORLD_SIZE、RANK 等环境变量。

至此，便解决了上述第一个问题—— “启动多个 python脚本，配置不同的 rank值”

至于上述第二个问题——"多个 pod 进行分布式训练 ，IP不确定的问题"。则是通过为每一个训练节点创建一个Services，通过记录集群内域名来解决。详见：jc.Controller.ReconcileServices方法（代码地址：github.com/kubeflow/common/pkg/controller.v1/common/job.go）

概念对齐

最后我们再对齐一下 WORLD_SIZE、RANK 这两个PyTorchJob自动生成的环境变量该怎么理解。

在多机单进程训练的时候， RANK确实可以是PyTorch官方的对rank的解释（代表该进程在 分布式训练中的序号），WORLD_SIZE 也确定是代表 rank的总数量。

在多机多进程训练则有些不一样，我们先回顾一下，手工做多机多卡（多训练进程）训练是怎么做的。

 在这里，RANK 则对应的是 node_rank (训练节点的序号)，WORLD_SIZE则对应的是 nnode (节点的总数)。