从 Momenta 的 AI 训练平台谈起：为什么 Kubernetes 天然适合 AI 场景？

Momenta 作为国内首家自动驾驶独角兽公司，内部有近百名算法工程师，他们为深度学习技术提供了最基础的保障，但是随着算力的不断进步与演进，Momenta 亟需一套简单易用的线下算法训练实验平台供算法研发人员使用。为了解决线下训练中遇到的使用问题，Momenta 基于 Kubernetes 研发了一套高度可扩展的跨云 AI 训练方案，使内部算法工程师可以自由地在各个机房使用不同的硬件训练。

Momenta 基础架构技术负责人薛磊在 ArchSummit 全球架构师峰会（深圳站） 2019 发表了以《基于 Kubernetes 的跨云 AI 训练平台构建与展望》为主题的演讲，同时，InfoQ 记者在大会现场采访了薛磊老师，进一步挖掘 Momenta AI 训练平台的更多技术细节，探讨了 Kubernetes 与 AI 场景结合需要注意的问题，对于备受业界关注的 KubeFlow 社区近期进展也有了更多了解。

下文整理自 InfoQ 与薛磊老师的对话实录。

Momenta 的 AI 训练平台

InfoQ：Momenta 自研的 AI 训练平台，与各大公有云上提供的 AI 平台相比，有哪些不同之处？

薛磊：Momenta 主要定位在计算机视觉领域使用深度学习技术，因此不是像 Spark 那样偏传统通用的平台，而是会有比较多定制化的东西。深度学习场景有很多特点，我们会基于这些特点做深度优化。像国内公有云上的 AI 平台，比如 AWS 的 SageMaker、阿里云的 PAI 等，我们试用过一些，但是都无法完全满足 Momenta 的需求。Momenta 的 AI 训练场景涉及到内部大量不同级别的算法工程师，为了区分不同级别的工程师从而使资源得到最大化利用，需要实现用户隔离、排队机制、用户权限管理、资源管理、容量上限管理等诸多功能。我们从自研 AI 训练平台开始到现在，才逐一解决掉这些问题。国内云厂商的方案更加通用，但对于一个公司内部拥有众多算法研究人员的场景可能不是那么合适。

InfoQ：Momenta 自研的 AI 训练平台针对你们所需要的场景做了哪些定制化工作？

薛磊：定制化工作首先在于调度层面，我们需要做一些额外的设置。在深度学习场景中，很多框架都支持多机训练，因此我们在 K8s 中需要做到同时调度多台机器，即并行调度，这是其一。另外，我们内部有众多角色，比如实习生、全职员工、资深工程师等。虽然他们做的事情是相同的，都要用 GPU 资源去做实验或者训练，但是在资源比较紧张的情况下，我们需要做到按照优先级抢占排队等方式，才能让更高优先级的角色能够尽快完成对应的任务。

在权限管控层面，K8s 是状态无关的，我们在外面增加了一层用户相关的属性以及权限系统，让用户只能看到自己的资源，而无法看到个人之外的其他资源，这是从安全性角度出发所做的定制。

另外，深度学习场景的数据集往往非常大，针对这一特性，我们在存储层面也做了一些额外的优化。比如用高速对象存储来存储所有用户的训练集，这个训练集更多的时候是图片，然后在客户端或者计算节点做缓存，同时保证缓存的命中是一个恒定值，这样可以减少 CPU 利用率，最大限度地提升 GPU 的使用体验。除此之外，为了实现对混合云的支持，以及公有云、私有云混布方式的支持，我们也做了相应的定制化，让用户可以无感地在不同硬件设备资源之前切换。

整体来说，我们是基于 K8s 的架构以及即插即用组件的开发模式，定制了我们这一整套 AI 训练平台。

InfoQ：Momenta 的 AI 训练平台是只在内部使用还是会对外提供给用户？

薛磊：目前我们的平台主要还是在内部使用，今年第三季度应该会开始跟外部云厂商合作，把我们这一整套的解决方案基于他们的 PaaS 或者 IaaS 层做一个解决方案，对外提供我们的解决能力。

InfoQ：仅供内部使用的 AI 训练平台和可以供外部客户使用的 AI 训练平台是否会对技术提出了不同的要求？

薛磊：我们从整体的架构上来说还是比较通用的，针对这个场景是完全 OK 的，但是对于传统的并行计算场景还是需要做一些定制，但是这不是我们要 Cover 的点，我们还是重点关注深度学习这个垂直领域。像阿里云、AWS，他们有比较多的并行计算经验，这块就不用我们去操心了。

InfoQ：Momenta 为什么要将 AI 训练平台做成跨云平台？这是为了解决实际研发过程中的什么痛点？

薛磊：主要是为了满足产品迭代过程中产生的需求，2016 年年初 Momenta 刚刚成立的时候，很多云厂商还没有参与到 AI 训练平台这个场景里来，所以我们只能自建实现整套方案。随着自动驾驶和深度学习的发展，云厂商不断参与到各平台建设以及解决方案中，我们从去年开始接触各家云厂商，包括华为、阿里云、AWS、腾讯云等，他们提供了一些比较优惠的政策，所以我们开始在不同的云上做一些尝试。但我们主要是利用云的可扩展性，在任务比较紧急的情况下直接在线上扩展，提升交付能力，基本上只会用到云的对象存储和基础网络。

InfoQ：基于 K8s 实现跨云技术方案的好处是什么？难点又是什么？

薛磊：我们的跨云方案更多指的是用户可以无感地在不同的云之间切换，并不像某些线上业务那样对高可用有强需求，所以我们所说的跨云更多时候是指可以随意切换计算资源，并且能够支持数据集自动备份。跨云的难点在于，不同云之间资源的类型不太一样，如何将不同的硬件资源组织起来是其中一个挑战；另外，我们的数据集非常庞大，怎么在各个云之间自由地切换、调度、备份，也是一个挑战。

AI 训练平台和 K8s

InfoQ：现阶段业界对于 AI 训练平台存在哪些需求？对于自动驾驶领域的企业来说，对 AI 训练平台是否有行业相关的特定需求？

薛磊：比较通用的一个需求就是要能够支持多机调度，通过多机并行训练让模型训练越快越好。对于一些比较大的公司，可能还会需要多用户管理，小公司管理起来相对简单一些。除此之外 AI 训练平台对于存储的性能要求也比较高。我们在计算的时候，如果是一个人训练一个模型，那么其实简单的存储就 OK 了，但如果有一百个人而且每个人训练的模型、用的数据集都不一样，还需要跑多机任务，这个时候什么样的存储才能帮你解决这个问题？这是摆在我们眼前比较现实的一个问题。针对用户不断扩展的计算需求，存储怎么做适配，这是我们主要关注的问题，也是我在大会上演讲的重点内容。

InfoQ：K8s 有哪些对应的技术方案可以满足上面提到的需求？

薛磊：首先，K8s 支持 Device Plugin 即设备插件的机制，因为有了这种机制所以 GPU 的调度就变得可行了，同时高性能的网络，像 Infiniband 也变得可行。除此之外，它可以做到可定制化的调度器，像开源界比较火的 Kubebatch 就是一个可定制的调度器。另外，K8s 还支持 Operator、CRD 的机制，可以帮我们抽象用户任务。整体来说，我们主要得益于 K8s 的框架设计和插件化方式。

基于当然还是有一些美中不足的点，对于 AI 企业来说，会希望能够在调度层面做更多针对 AI 场景的专门的优化。我们比较关注基于硬件拓扑结构做调度，这块目前社区已经有一些尝试，但是还有很多工作需要做。

InfoQ：目前越来越多企业基于 K8s 搭建自己的 AI 训练平台，K8s 具备有哪些优势使之适合与 AI 场景结合？

薛磊： 我觉得 K8s 最大的优势在于它可以和任何场景结合 ，K8s 这种插件化、组件化的开发方式让它能够支持更多定制化的设计开发工作，它更多时候是一个通用的框架平台，但它又可以针对任何特定领域的需求做定制化的修改。

InfoQ：除了 K8s 以外，还有其他比较好的技术方案吗？

薛磊：并行计算场景，一般 Yarn 用的比较多，它也是比较老牌的调度平台。另外就是 Mesos，还有 K8s。我们为什么选 K8s？部分原因跟 Momenta 的业务背景有关，另外，K8s 支持高度定制化，同时它的社区发展也非常火爆。虽然 Yarn 做起来会更简单一点，比如它可能天然支持批量调度，但是从定制化开发的角度来说可能不是那么好做。而且在 17 年的时候，K8s 已经大行其道了，所以我们基本上就直接选择了 K8s。

KubeFlow

InfoQ：聊到 K8s 和 AI 就不得不提 KubeFlow，您也是 KubeFlow 的贡献者之一，能跟我们分享一下过去半年 KubeFlow 社区重点在做哪些工作吗？

薛磊：KuberFlow 社区成立于 17 年年底，成立以来做了非常多工作，一开始专注于做一个套解决方案，攘括了非常多开源组件，帮助算法工程师和做技术架构的工程师将各自负责的工作融合起来；后来慢慢加入了更多新设计的产品，比如说像后面 KubeFlow 的 pipeline，简化了用户部署的难度，操作起来会更加简单，也引入了一些任务管理、用户管理的功能，还有阿里云贡献出来的 Arena（目前已经合并回 KubeFlow 社区），也能帮用户更方便地管理整套流程。这半年，据我了解，KubeFlow 社区投入了很多精力在做 AutoML，即自动化调参，还有就是将周边生态做得更好，帮助用户降低部署成本，引入了一些做实验的场景。总体来说，发展还是非常迅速的。

InfoQ：那么 KubeFlow 接下来还有什么比较重要的规划吗？

薛磊：下一步比较重要的应该是面向 1.0 发布的准备工作，正式 Release 前要做一些稳定性、易用性、文档等相关工作。

InfoQ：谈谈 KubeFlow 的优势和不足之处。

薛磊：优势是对于用户来说会简化很多自己去调研的成本，KubeFlow 集成了很多东西，用户在训练完之后，到 Serving，再到后面的测试，所有的东西都可以通过 KubeFlow 来管理，另外，包括模型管理、数据集管理等功能都做得比较完善。但是缺点就在于它有点囊括太多东西了，如果我们只需要训练，不需要 Serving，而且我们有自己的模型管理组件，KubeFlow 要怎么跟我们现有的组件或者定制化的东西进行组合，可能是比较麻烦的一个点。

Info Q：KubeFlow 项目现在是否足够成熟？什么样的企业适合选用 KubeFlow？为什么？

薛磊：就目前来说，国内用起来可能稍微不是那么友好，因为它是 Google 发起的，基本上都是英文文档，从使用到最终部署的时候，可能都会有一些网络传输的问题。它的容器镜像全都在 GCR 上面，所以天然对于国内用户来说可能不是那么友好。如果是想体验 KubeFlow 功能的话，我推荐用阿里的 Arena 可能会更简单一些，而且他们也会针对性地做一些国内的适配。AS 第一天下午来自华为云的马达也分享了他们基于 K8s 的一套解决方案，我感觉业界大家都在向前去推动。KubeFlow 是一个很好的选择，但还是那句话，因为每家公司都有不太一样的地方，你可能需要做一些深入的定制化的东西，这其中可能需要做更多沟通和调研的工作。

采访嘉宾介绍：

薛磊，Momenta 基础架构技术负责人。目前在 Momenta 从事 AI 基础架构研发的工作，带领团队开发了跨云多区域异构计算 GPU 平台，该平台基于 Kubernetes 框架，结合 AI 训练的特别场景设计开发出了一套针对算法研发人员方便易用的训练平台。与此同时也是 KubeFlow 的贡献者，其中 caffe2-operator 的作者。 在此之前，作为 Hypercontainer 的早期员工，参与了大多数项目，是容器技术早期践行者。