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"