MPI on Kubernetes

MPI(Message Passing Interface) 是一种可以支持点对点和广播的通信协议，具体实现的库有很多，使用比较流行的包括 Open Mpi, Intel MPI 等等，关于这些 MPI 库的介绍和使用，本文就不多赘述了，各位可以看看官方文档。

mpi-operator 是 Kubeflow 社区贡献的另一个关于深度/机器学习的一个 Operator，关于 mpi-operator 的 proposal，可以参考 mpi-operator-proposal。目前社区在 mpi-operator 主要用于 allreduce-style 的分布式训练，因为 mpi-operator 本质上就是给用户管理好多个进程之间的关系，所以天然支持的框架很多，包括 Horovod, TensorFlow, PyTorch, Apache MXNet 等等。而 mpi-operator 的基本架构是通过 Mpijob 的这种自定义资源对象来描述分布式机器学习的训练任务，同时实现了 Mpijob 的 Controller 来控制，其中分为 Launcher 和 Worker 这两种类型的工作负荷。

社区开源的 mpi-operator，开箱即用，但是在生产集群的应用，在某些方面，面对一些固定场景和业务的时候会有一定的限制。

1. 对于使用 GPU 资源的 Worker 有可能会调度到 单独的 GPU 集群，而 Launcher 会在其他集群上，所以跨集群 Launcher 和 Worker 的通信问题，需要额外的考虑
2. 希望通过 Pod IP 通信
3. Metrics 收集，目前社区版缺少 Mpijob 的基础指标
4. 需要支持更多的批调度组件
5. v1.8 和高版本集群的兼容，这里主要涉及到资源对象 status 这类的 subresource 的更新操作的兼容

对于用户，只要创建一个 Mpijob 的自定义资源对象，在 Template 配置好 Launcher 和 Worker 的相关信息，就相当于描述好一个分布式训练程序的执行过程了。

apiVersion: kubeflow.org/v1alpha2
kind: MPIJob
metadata:
  name: tensorflow-mnist
spec:
  slotsPerWorker: 1
  cleanPodPolicy: Running
  mpiReplicaSpecs:
    Launcher:
      replicas: 1
      template:
        spec:
          containers:
          - image: horovod-cpu:latest
            name: mpi-launcher
            command:
            - mpirun
            args:
            - -np
            - "2"
            - --allow-run-as-root
            - -bind-to
            - none
            - -map-by
            - slot
            - -x
            - LD_LIBRARY_PATH
            - -x
            - PATH
            - -mca
            - pml
            - ob1
            - -mca
            - btl
            - ^openib
            - python
            - /examples/tensorflow_mnist.py
            resources:
              limits:
                cpu: 1
                memory: 2Gi
    Worker:
      replicas: 2
      template:
        spec:
          containers:
          - command:
            - ""
            image: horovod-cpu:latest
            name: mpi-worker
            resources:
              limits:
                cpu: 2
                memory: 4Gi

Worker 本质上是 StatefulSet，在分布式训练的过程中，训练任务通常是有状态的，StatefulSet 正是管理这些的 Workload 的对象。通常，Launcher 会是一个比较轻量化的 Job，他主要完成几条命令的发送就可以了，通常是把命令通过 ssh/rsh 来发送接受命令，在 mpi-operator 里使用的是 kubectl 来给 Worker 发送命令，下图是其基础架构图。

Mpijob 启动的顺序是先启动 Worker 再启动 Launcher。

// pkg/controllers/v1alpha2/mpi_job_controller.go
// 先创建 Worker
worker, err = c.getOrCreateWorkerStatefulSet(mpiJob, workerReplicas)
if err != nil {
	return err
}

// 再创建 Launcher
if launcher == nil {
	launcher, err = c.kubeClient.BatchV1().Jobs(namespace).Create(c.newLauncher(mpiJob, c.kubectlDeliveryImage))
	if err != nil {
		return err
	}
}

其中 kubectl-delivery 的作用主要是将 kubectl 放入到 Launcher 容器内，之后可以通过 kubectl 来给 Worker 发送 mpirun 的命令，下图是其任务执行时候的时序图。

# Launcher 容器中执行的命令，就是给 Worker 下发 mpirun 的命令
/opt/kube/kubectl exec mpi-ea4304c32617ec5dvx89ht1et9-worker-0 -- /bin/sh -c PATH=/usr/local/bin:$PATH ; export PATH ; LD_LIBRARY_PATH=/usr/local/lib:$LD_LIBRARY_PATH ; export LD_LIBRARY_PATH ; DYLD_LIBRARY_PATH=/usr/local/lib:$DYLD_LIBRARY_PATH ; export DYLD_LIBRARY_PATH ; /usr/local/bin/orted -mca ess "env" -mca ess_base_jobid "2828599296" -mca ess_base_vpid 1 -mca ess_base_num_procs "2" -mca orte_node_regex "mpi-ea[4:4304]c32617ec5dvx89ht1et9-launcher-kljzn,mpi-ea[4:4304]c32617ec5dvx89ht1et9-worker-0@0(2)" -mca orte_hnp_uri "2828599296.0;tcp://6.16.105.7:36055" -mca plm "rsh" --tree-spawn -mca orte_parent_uri "2828599296.0;tcp://6.16.105.7:36055" -mca orte_default_hostfile "/etc/mpi/hostfile" -mca plm_rsh_agent "/etc/mpi/kubexec.sh" -mca hwloc_base_binding_policy "none" -mca rmaps_base_mapping_policy "slot" -mca pmix "^s1,s2,cray,isolated"