Kubeflow提供TFJob、PytorchJob、MPIJob等CRD（Custom Resource Definition），基于这些CRD，可以在Kubernetes集群上运行分布式训练。

介绍

Kubeflow中的training operator，包括有tf-operator、pytorch-operator、mpi-operator等，这些operators服务于各种深度学习的分布式训练，主要工作有：

• 在Kubernetes集群上创建Pod，将各个训练进程拉起。
• 配置服务发现的信息，创建相关Kubernetes资源。
• 监控并更新整个任务的状态（如：jobCreateCount，jobDeleteCount，jobSuccessfulCount，jobFailedCount，jobRestartCount）。

因此，可以将Kubeflow中的operators用作下发创建分布式训练任务。

前提工作

• k8s集群，集群版本 >= 1.23
• 训练模型的代码

安装Training Operator

在kubeflow/training-operator项目中找到对应的YAML文件，安装Training Operator选择步骤1或步骤2即可。

• Master分支

kubectl apply -k "kubeflow/training-operator/manifests/overlays/standalone"

• Stable版本

kubectl apply -k "kubeflow/training-operator/manifests/overlays/standalone?ref=v1.5.0"

• 验证

kubectl get pods -nkubeflow # 查看training operator是否在运行
kubectl get crd # 查看tfjob、pytorchjob、mpijob等crd

创建Job

TFJob

TFJob是Kubernetes的API对象资源，用于在集群上运行TensorFlow训练任务。官方提供了多个训练案例，本文以 tf-dist-mnist 训练为例。

• 制作镜像

镜像基于Tensorflow 1.5.0 的官方镜像，并将代码复制到镜像内，并配制好entrypoint即可。镜像文件Dockerfile示例如下：

FROM tensorflow/tensorflow:1.5.0

ADD ./ /var/tf_dist_mnist
ENTRYPOINT ["python", "/var/tf_dist_mnist/dist_mnist.py"]

根据Dockerfile构建Docker镜像，指定Dockerfile和构建上下文路径。

docker build -f Dockerfile -t kubeflow/tf-dist-mnist-test:1.0 ./ # 构建的镜像tag是 kubeflow/tf-dist-mnist-test:1.0，./是将当前目录作为构建上下文路径

额外操作：将镜像上传到仓库。

docker login <registry-url> # 需要输入用户名和密码
docker tag source_image[:tag] <registry-url>/<repository>[:tag]
docker push <registry-url>/<repository>[:tag]

• 准备 TFJob的YAML文件，定义了2个PS和4个Worker。镜像地址可以更改为上传后在镜像仓库内的地址。

apiVersion: "kubeflow.org/v1"
kind: "TFJob"
metadata:
  name: "tf-dist-mnist-for-e2e-test"
  namespace: "kubeflow"
spec:
  tfReplicaSpecs:
    PS:
      replicas: 2
      restartPolicy: Never
      template:
        spec:
          containers:
            - name: tensorflow
              image: kubeflow/tf-dist-mnist-test:1.0
    Worker:
      replicas: 4
      restartPolicy: Never
      template:
        spec:
          containers:
            - name: tensorflow
              image: kubeflow/tf-dist-mnist-test:1.0

• 提交任务

kubectl create -f ./tf-dist-mnist.yaml

• 查看任务状态

kubectl -n kubeflow get -o yaml tfjobs tf-dist-mnist-for-e2e-test # 查看job的状态，tf-dist-mnist-for-e2e-test是${JOBNAME}
kubectl -n kubeflow describe tfjobs tf-dist-mnist-for-e2e-test # 查看job的最近的events，k8s默认保留一小时内的events
kubectl logs ${PODNAME} # 查看TFJob创建的pod的日志
kubectl -n kubeflow describe pods ${PODNAME} # 查看pod的事件，确保它们被调度

具体可参考 kubeflow官方文档。

PyTorchJob

PyTorchJob是Kubernetes的API对象资源，用于在集群上运行PyTorch训练任务。官方提供了多个训练案例，本文以分布式MNIST示例训练为例。

• 创建PyTorchJob，定义了一个Master和一个Worker。

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

• 提交任务

kubectl create -f pytorch-simple.yaml 

• 查看任务状态

kubectl get pods -l job-name=pytorch-simple -n kubeflow
kubectl get -o yaml pytorchjobs pytorch-simple -n kubeflow

MPIJob

MPI Operator对接Horovod的Elastic模式，可以动态调整Worker数量，在集群中运行弹性训练任务，充分利用集群闲置资源。

• 构建镜像。

在本实例中直接使用官方提供的镜像 horovod/horovod:0.20.0-tf2.3.0-torch1.6.0-mxnet1.5.0-py3.7-cpu，镜像中已经包含训练代码 tensorflow2_mnist_elastic。

注意点：

• • 镜像比较大，拉取会比较费时。如果集群没有配置代理，可以考虑先将镜像本地拉取后，传输到集群中。

docker pull <image_name>
docker save -o image.tar <image_name> # 将本地镜像保存为tar归档文件，可以传输到集群中
docker load -i image.tar # 加载镜像

• • 训练代码中需要从网络中加载MNIST数据集，如果没有配置代理，可以考虑基于现有的镜像，把数据集拷贝到镜像中，并将相应代码改为加载本地数据集。

docker run -it --rm image_name /bin/bash # 启动一个临时容器查看镜像内部代码
docker cp /path/dataset <container_id>:/data # 将数据集复制到容器内部，docker ps 可以找到容器id，同时修改相应代码加载本地数据集
docker commit <container_id> new_image # 将修改后的容器保存为新镜像

• 创建Job

准备 MPIJob的YAML文件。

apiVersion: kubeflow.org/v1
kind: MPIJob
metadata:
  name: tensorflow-mnist-elastic
spec:
  slotsPerWorker: 1
  cleanPodPolicy: Running
  mpiReplicaSpecs:
    Launcher:
      replicas: 1
      template:
        spec:
          containers:
          - image: horovod/horovod:0.20.0-tf2.3.0-torch1.6.0-mxnet1.5.0-py3.7-cpu
            name: mpi-launcher
            command:
            - horovodrun
            args:
            - -np
            - "2"
            - --min-np
            - "1"
            - --max-np
            - "3"
            - --host-discovery-script
            - /etc/mpi/discover_hosts.sh
            - python
            - /examples/elastic/tensorflow2_mnist_elastic.py
            resources:
              limits:
                cpu: 1
                memory: 2Gi
    Worker:
      replicas: 2
      template:
        spec:
          containers:
          - image: horovod/horovod:0.20.0-tf2.3.0-torch1.6.0-mxnet1.5.0-py3.7-cpu
            name: mpi-worker
            resources:
              limits:
                cpu: 2
                memory: 4Gi

• 提交任务

kubectl create -f ./tensorflow-mnist-elastic.yaml

• 查看任务状态

kubectl get -o yaml mpijob tensorflow-mnist-elastic # 查看mpijob的状态