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,开箱即用,但是在生产集群的应用,在某些方面,面对一些固定场景和业务的时候会有一定的限制。
- 对于使用 GPU 资源的
Worker
有可能会调度到 单独的 GPU 集群,而Launcher
会在其他集群上,所以跨集群Launcher
和Worker
的通信问题,需要额外的考虑 - 希望通过 Pod IP 通信
- Metrics 收集,目前社区版缺少
Mpijob
的基础指标 - 需要支持更多的批调度组件
- 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"