Tensorflow分布式训练的调度方案

背景

随着机器学习和人工智能的迅猛发展，业界出现了许多开源的机器学习平台。由于机器学习与大数据天然的紧密结合，基于 Hadoop Yarn 的分布式任务调度仍是业界主流，但是随着容器化的发展，Docker + Kubernetes 的云原生组合，也展现出了很强的生命力。

痛点

在建设分布式训练平台的过程中，我们和机器学习的各个业务方，包括搜索推荐、图像算法、交易风控反作弊等，进行了深入沟通，调研他们的痛点。从中我们发现，算法业务方往往专注于模型和调参，而工程领域是他们相对薄弱的一个环节。建设一个强大的分布式平台，整合各个资源池，提供统一的机器学习框架，将能大大加快训练速度，提升效率，带来更多的可能性，此外还有助于提升资源利用率。

痛点一：对算力的需求越来越强烈

算力代表了生产力。深度学习在多个领域的出色表现，给业务带来了更多的可能性，同时对算力提出了越来越高的要求。深度学习模型参数的规模陡增，迭代的时间变的更长，从之前的小时级别，变成天级别，甚至月级别。以商品推荐为例，面对几亿的参数，近百亿的样本数量，即使采用 GPU 机器，也需要长达一个星期的训练时间；而图像业务拥有更多的参数和更复杂的模型，面对 TB 级的训练样本，单机场景下往往需要长达近一个月的训练时间。

再者，机器学习具有试错性非常强的特点，更快的训练速度可以带来更多的尝试，从而发掘更多的可能性。Tensorflow 从 0.8 版本开始支持分布式训练，至今为止，无论高阶还是低阶的 API，对分布式训练已经有了完善的支持。同时，Kubernetes 和 Hadoop 等具有完善的资源管理和调度功能，为 Tensorflow 分布式训练奠定资源层面的基础。

Tensorflow On Yarn 和 Tensorflow On Spark 是较早的解决方案，奇虎 360 的 Xlearning 也得到众人的青睐。而基于 Kubernetes 的 kubeflow ，也展现出旺盛的生命力。

痛点二：人肉管理的成本很高

业务方反馈的第二大问题是人肉管理的成本问题。人肉化的管理主要包含了部署和训练任务管理两大方面。

人肉部署
一个典型的场景是：团队内的成员共享一批机器，每次跑训练任务前，用户手动登陆机器，下载代码，安装对应的 python 包，过程繁琐且容易出现安装包版本的冲突。

由于不同的训练任务对 python 的版本和依赖完全不同，比如有些模型使用 python 2.7，有些使用 python 3.3，有些使用 tensorflow 1.8，有些使用 tensorflow 1.11 等等，非常容易出现依赖包冲突的问题。虽然沙箱能在一定程度上解决这问题，但是也带来了额外的管理负担。

现有已经开源的实现方案

Tensorflow on Yarn 方案

• Intel的Tensorflow on YARN: github地址
• 奇虎360开源的XLearning: github地址
• Hadoop 社区：github地址 ,相关讨论YARN-8220
• 美团 参考文档
• 我们公司采用的方案

Tensorflow on K8S 方案

• Google开源的Kubeflow: github

Tensorflow on Spark方案

• Yahoo开源的 Tensorflow on Spark github地址
• Databricks的Spark Deep Learning: github地址

三种方案架构

Tensorflow on Yarn 架构

架构描述：
Client：负责启动作业及获取作业执行状态；
ApplicationMaster（TFAM）：负责输入数据分片、启动及管理Container、执行日志保存等；
Container（Worker）：作业的实际执行者，负责启动Tensorflow的Worker或Parameter Server）进程，监控并向AM汇报进程状态，端口信息，上传作业的输出等。对于TensorFlow类型作业，还负责启动TensorBoard服务。

Tensorflow on Spark 架构

业界现在大部分spark 也是on Yarn，不排除spark standlone集群。spark 借助Executor 启动Tensorflow 的 ps 和 worker。

Tensorflow on K8S 架构

Tensorflow 和K8S 都是google 开源的项目，google 也在推动使用K8S 作为Tensorflow的调度框架。利用K8S的POD启动Tensorflow 的 ps 和 worker。主要基于 k8s 的 custom resource descriptor (CRD) and associated controller，很容易实现tf 的调度。

K8S vs Yarn

从Tensorflow 的调度方案开源实现来看，可以归结为在K8S 和 Yarn 之间的选型。

现在yarn 也支持docker ，在隔离性上和K8S的差距又小了一步。

No perfect match, only match perfect.

分布式机器学习训练需要大量的计算资源，现在大厂有大量计算资源的团队：Hadoop 平台团队、云服务团队。如果是Hadoop 团队，可能对Yarn最熟悉，Yarn机器规模最大，为了便于统一管理会采用Tensorflow On Yarn的方案；云服务团队更侧重与On K8S的方案。技术选型考虑历史因素、运维成本、开源项目社区活跃度、引入新的技术栈带来的收益折中。No perfect match, only match perfect.

Reference

• https://www.infoq.cn/article/k4uLIMaNccimGfuQk_cd