Tensorflow训练中Worker如何动态伸缩

Tensorflow对于大部分算法用户来说能够满足使用需求 。但是在日常使用过程中也发现有一些bug导致日常生产过程中的训练异常中断，影响使用 。尤其是在离线集群进行大规模训练时，集群网络负载大，机器环境恶劣，导致故障频发，影响用户体验。 因此对于Tensorflow的稳定性改造，在大规模生产中是一项必要工作。

Worker 角色的 Failover 及动态扩缩容改造
为了提升用户使用体验提升集群资源有效利用率抵消用户超量配置训练资源的问题。我们主要对Tensorflow角色运行部分进行了如下改造：

Worker 角色的 Failover 。Worker是分布式训练中的计算角色，异常退出的高发户 。增加Failover机制，做到异常中断用户无感知，极大提升用户的使用体验。
Worker 角色 动态扩缩容。在线生产过程中数据量生产不是均匀产生的 。如国内业务白天数据量较大，晚上相对较少。国际业务 则正好相反。为了优化资源利用率，我们增加了Tensorflow动态扩缩容机制，在流程数据生产低谷时段 ，能够主动释放计算资源给其他计算任务。数据高峰时段则主动增加资源提高训练速度。

实现方案介绍
Tensorflow分布式训练采用的是PS模式，即参数服务器模式。训练过程中各个Worker （数据消费和梯度计算的执行者）将计算得到的梯度信息同步到PS上正向计算时Worker负责从PS上拉取参数 。PS与Worker之间的数据通信依赖于GRPC 。GRPC 连接在角色启动的过程中进行建立 。在原生Tensorflow中GRPC通信通道的建立依赖于每个角色启动时配置的 cluster_spec文件，是一种静态构建过程，为了实现Failover和动态扩缩容，我们需要打破这种静态机制，构建动态连接关系。

1.原生sparse配置打破worker之间关联

静态cluster_spec配置格式如下（以3Worker，2PS为例）

这种配置方式 Tensorflow 会 在启动时为PS&Worker和Worker&Worker之间都建立通信连接，当任意Worker失败时所有Worker 都会失败。原生Tensorflow还支持一种Sparse配置，可以打破Worker之间的依赖，这样任意Worker失败不会对其他Worker造成 影响。如Worker0启动配置方式如下：

通过sparse的方式可以切断Worker之间的相互影响。可以解决Worker缩容的问题。

2.通信通道注册机制打破PS，Worker角色依赖

在Tensorflow中通信通道建立是在角色启动时的find device阶段 。当任务已经运行时，如果worker发生重启，PS上的通信通道建立过程不会重复执行 ，新启动的角色无法被 PS发现。因此这种方式不能够解决 Failover 和 Worker扩容 的问题。为了打破这种 限制，需要在Tensorflow中增加通信参数注册机制。PS中的Worker通信信息不依赖于参数传入，改由GRPC服务动态注册 。实现通信信息注册更新，类调用关系如下图所示（依据Tensorflow1.8版本源码结构）：

注册机制实现机制如上所示，在实现过程中也遇到一些问题。Tensorflow在不同层次的通信对象做了缓存，如果在上层找到了通信对象则不会继续向下查找。这样在Worker Failover的时候会导致新通信地址无法得到更新因为上层缓存了旧 的通信对象（缓存 在GrpcWorkerCache中。我们采用的策略是将上层cache全部删除，只更新最下层的IP和Port信息 。这样在Worker Failover或者新增节点的时候会有一次通信对象rebuild 。解决了Worker能够任意重启和任意扩缩的问题剩下需要解决的问题就是调度策略的 问题 。

整个分布式任务的启动依赖于分布式调度框架，我们是基于yarn和docker实现。因此可以在 AppMaster的实现中结合在线数据源（Kafka）的数据生产速度进行调度具体实现不再赘述。这样一套资源自适应的Tensorflow分布式训练系统就落地了，实际上线效果也非常理想，不仅cover了因为集群环境恶劣和Tensorflow bug导致Worker abort问题，又有效提升了集群利用率 。