mxnet分布式形式代码阅读小结

mxnet的分布式形式研究小结

// 未经博主本人允许，谢绝转载，谢谢

目前mxnet支持多种分布式形式，根据官方文档，支持的形式有ssh、mpirun、sge集群、yarn集群等几种形式。ssh和mpirun是需要指定机器列表的，ssh需要提交机对其他机器有信任关系。如官方文档所述：

『 Launching a distributed job is a bit different from running on a single machine. MXNet provides  tools/launch.py  to start a job by using  ssh ,  mpi ,  sge , or  yarn .』

最近主要看了基于yarn的mxnet深度学习作业，下面做一个小结，记录一些思考

可以使用mxnet/tools/launch.py向yarn集群提交一个mxnet，launch.py的参数--launcher指定为yarn即可

launch.py会调用mxnet/dmlc-core/tracker/dmlc_tracker/yarn.py中定义的submit，主要作用是export一些yarn作业需要的环境变量

最终都会调用tracker.py的submit函数，当提交时指定的server数量大于0时，会在当前提交机用一个线程启动kvstore中的scheduler角色。

def submit(nworker, nserver, fun_submit, hostIP='auto', pscmd=None):
    if nserver == 0:
        pscmd = None

    envs = {'DMLC_NUM_WORKER' : nworker,
            'DMLC_NUM_SERVER' : nserver}
    hostIP = get_host_ip(hostIP)

    if nserver == 0:
        rabit = RabitTracker(hostIP=hostIP, nslave=nworker)
        envs.update(rabit.slave_envs())
        rabit.start(nworker)
    else:
	// 这里会调用c++的库在本地启动scheduler，接收后续启动的来自worker和server的心跳
        pserver = PSTracker(hostIP=hostIP, cmd=pscmd, envs=envs)
        envs.update(pserver.slave_envs())
    fun_submit(nworker, nserver, envs)

    if nserver == 0:
        rabit.join()
    else:
        pserver.join()

然后向yarn集群提交一个作业，这个作业的ApplicationMaster在dmlc-core中已经实现，会根据worker和server的数量向Yarn集群的RM申请资源，拿到资源后，启动worker和server，在环境变量中export scheduler的地址、自身的角色（自己是worker还是server）等必要信息。这个AM会在worker或者server异常的时候重新启动它们。

在机器上启动了worker或者server后，单机也需要做计算的工作，如果角色是worker，那么计算后负责参数推送和拉取到server。

这里没有输入数据的切分和分配，每个计算节点读取什么样的输入，是由计算节点自身的rank号来决定自己读取总输入的哪一部份。

个人的想法：

跟想象中的分布式作业相比，这里scheduler这样具有『协调』功能的角色，和ApplicationMaster同样具有『协调和管理』功能的角色是割裂的，同时对作业的总体进度也没有把控，不具备分布式作业中常见的处理慢节点、控制进度、缩扩容等等的决策和实施能力，分布式的实现相对来说是比较简单的。