## 打造企业级云深度学习平台
文/陈迪豪
### 深度学习服务介绍
机器学习与人工智能,相信大家已经耳熟能详,随着大规模标记数据的积累、神经网络算法的成熟以及高性能通用 GPU 的推广,深度学习逐渐成为计算机专家以及大数据科学家的研究重点。近年来,无论是图像的分类、识别和检测,还是语音生成、自然语言处理,甚至是 AI 下围棋或者打游戏都基于深度学习有了很大的突破。而随着 TensorFlow、Caffe 等开源框架的发展,深度学习的门槛变得越来越低,甚至初中生都可以轻易实现一个图像分类或者自动驾驶的神经网络模型,但目前最前沿的成果主要还是出自 Google、微软等巨头企业。
Google 不仅拥有优秀的人才储备和大数据资源,其得天独厚的基础架构也极大推动了 AI 业务的发展,得益于内部的大规模集群调度系统 Borg,开发者可以快速申请大量GPU资源进行模型训练和上线模型服务,并且通过资源共享和自动调度保证整体资源利用率也很高。Google 开源了 TensorFlow 深度学习框架,让开发者可以在本地轻易地组合 MLP、CNN 和 RNN 等模块实现复杂的神经网络模型,但 TensorFlow 只是一个数值计算库,并不能解决资源隔离、任务调度等问题,将深度学习框架集成到基于云计算的基础架构上将是下一个关键任务。
除了 Google、微软,国内的百度也开源了 PaddlePaddle 分布式计算框架,并且官方集成了
Kubernetes 等容器调度系统,用户可以基于 PaddlePaddle 框架实现神经网络模型,同时利用容器的隔离性和 Kubernetes 的资源共享、自动调度、故障恢复等特性,但平台不能支持更多深度学习框架接口。而亚马逊和腾讯云相继推出了面向开发者的公有云服务,可以同时支持多种主流的开源深度学习框架,阿里、金山和小米也即将推出基于 GPU 的云深度学习服务,还有无数企业在默默地研发内部的机器学习平台和大数据服务。
面对如此眼花缭乱的云服务和开源技术,架构师该如何考虑其中的技术细节,从用户的角度又该如何选择这些平台或者服务呢。我将介绍小米云深度学习平台的架构设计与实现细节,希望能给 AI 领域的研发人员提供一些思考和启示。
### 云深度学习平台设计
云深度学习平台,我定义为 Cloud Machine Learning,就是基于云计算的机器学习和深度学习平台。首先 TensorFlow、MXNet 是深度学习框架或者深度学习平台,但并不是云深度学习平台,它们虽然可以组成一个分布式计算集群进行模型训练,但需要用户在计算服务器上手动启动和管理进程,并没有云计算中任务隔离、资源共享、自动调度、故障恢复以及按需计费等功能。因此我们需要区分深度学习类库以及深度学习平台之间的关系,而这些类库实现的随机梯度下降和反向传播等算法却是深度学习应用所必须的,这是一种全新的编程范式,需要我们已有的基础架构去支持。
云计算和大数据发展超过了整整十年,在业界催生非常多优秀的开源工具,如实现了类似 AWS IaaS 功能的 OpenStack 项目,还有 Hadoop、Spark、Hive 等大数据存储和处理框架,以及近年很火的 Docker、Kubernetes 等容器项目,这些都是构建现代云计算服务的基石。这些云服务有共同的特点,例如我们使用 HDFS 进行数据存储,用户不需要手动申请物理资源就可以做到开箱即用,用户数据保存在几乎无限制的公共资源池中,并且通过租户隔离保证数据安全,集群在节点故障或者水平扩容时自动触发 Failover 且不会影响用户业务。虽然 Spark 通过 MLib 接口提供部分机器学习算法功能,但绝不能替代 TensorFlow、Caffe 等深度学习框架的作用,因此我们仍需要实现 Cloud Machine Learning
服务,并且确保实现云服务的基本特性——我将其总结为下面几条:
- 屏蔽硬件资源保证开箱即用
- 缩短业务环境部署和启动时间
- 提供“无限”的存储和计算能力
- 实现多租户隔离保证数据安全
- 实现错误容忍和自动故障迁移
- 提高集群利用率和降低性能损耗
相比于 MapReduce 或者 Spark 任务,深度学习的模型训练时间周期长,而且需要调优的超参数更多,平台设计还需要考虑以下几点:
- 支持通用 GPU 等异构化硬件
- 支持主流的深度学习框架接口
- 支持无人值守的超参数自动调优
- 支持从模型训练到上线的工作流
这是我个人对云深度学习平台的需求理解,也是小米在实现 cloud-ml 服务时的基本设计原则。虽然涉及到高可用、分布式等颇具实现难度的问题,但借助目前比较成熟的云计算框架和开源技术,我们的架构和实现基本满足了前面所有的需求,当然如果有更多需求和想法欢迎随时交流。
### 云深度学习平台架构
遵循前面的平台设计原则,我们的系统架构也愈加清晰明了,为了满足小米内部的所有深度学习和机器学习需求,需要有一个多租户、任务隔离、资源共享、支持多框架和
GPU 的通用服务平台。通过实现经典的 MLP、CNN 或 RNN 算法并不能满足业务快速发展的需求,因此我们需要支持 TensorFlow 等用户自定义的模型结构,并且支持高性能 GPU 和分布式训练是这个云深度学习平台的必须功能,不仅仅是模型训练,我们还希望集成模型服务等功能来最大化用户的使用效益。
计算机领域有句名言“任何计算机问题都可以通过增加一个中间层来解决”。无论是 AWS、OpenStack、Hadoop、Spark 还是 TCP/IP 都是这样做的,通过增加一个抽象层来屏蔽底层资源,对上层提供更易用或者更可靠的访问接口。小米的 cloud-ml 平台也需要实现对底层物理资源的屏蔽,尤其是对 GPU 资源的抽象和调度,但我们不需要重新实现,因为社区已经有了很多成熟的分布式解决方案,如 OpenStack、Yarn和Kubernetes。目前 OpenStack 和 Yarn 对 GPU 调度支持有所欠缺,虚拟机也存在启动速度慢、性能 overhead 较大等问题,而容器方案中的 Kubernetes 和 Mesos 发展迅速,支持
GPU 调度等功能,是目前最值得推荐的架构选型之一。
目前小米 cloud-ml 平台的任务调度和物理机管理基于多节点的分布式 Kubernetes 集群,对于 OpenStack、Yarn 和 Mesos 我们也保留了实现接口,可以通过实现 Mesos 后端让用户的任务调度到 Mesos 集群进行训练,最终返回给用户一致的使用接口。目前 Kubernetes 最新稳定版是1.6,已经支持 Nvidia GPU 的调度和访问,对于其他厂商 GPU 暂不支持但基本能满足企业内部的需求,而且 Pod、Deployment、Job、StatefulSet 等功能日趋稳定,加上 Docker、Prometheus、Harbor 等生态项目的成熟,已经在大量生产环境验证过,可以满足通用 PaaS 或者 Cloud Machine learning 等定制服务平台的需求。
使用 Kubernetes 管理用户的 Docker 容器,还解决了资源隔离的问题,保证不同深度学习训练任务间的环境不会冲突,并且可以针对训练任务和模型服务使用 Job 和 Deployment 等不同的接口,充分利用分布式容器编排系统的重调度和负载均衡功能。但是,Kubernetes 并没有完善的多租户和 Quota 管理功能,难以与企业内部的权限管理系统对接,这要求我们对 Kubernetes API 进行再一次“抽象”。我们通过 API Server 实现了内部的 AKSK 签名和认证授权机制,在处理用户请求时加入多租户和 Quota 配额功能,并且对外提供简单易用的
RESTful API,进一步简化了整个云深度学习平台的使用流程,整体架构设计如图1。
图1 云深度学习平台整体架构
通过实现 API Server,我们对外提供了 API、SDK、命令行以及 Web 控制台多种访问方式,最大程度上满足了用户复杂多变的使用环境。集群内置了 Docker 镜像仓库服务,托管了我们支持的17个深度学习框架的容器镜像,让用户不需要任何初始化命令就可以一键创建各框架的开发环境、训练任务以及模型服务。多副本的 API Server
和 Etcd 集群,保证了整个集群所有组件的高可用,和 Hadoop 或者 Spark 一样,我们的 cloud-ml 服务在任意一台服务器经历断网、宕机、磁盘故障等暴力测试下都能自动 Failover 保证业务不受任何影响。
前面提到,我们通过抽象层定义了云深度学习平台的接口,无论后端使用 Kubernetes、Mesos、Yarn 甚至是 OpenStack、AWS 都可以支持。通过容器的抽象可以定义任务的运行环境,目前已经支持17个主流的深度学习框架,用户甚至可以在不改任何一行代码的情况下定义自己的运行环境或者使用自己实现的深度学习框架。在灵活的架构下,我们还实现了分布式训练、超参数自动调优、前置命令、NodeSelector、Bring Your Own Image 和
FUSE 集成等功能,将在下面逐一介绍。
### 云深度学习平台实现
前面提到我们后端使用 Kubernetes 编排系统,通过 API Server 实现授权认证和 Quota 配额功能。由于云深度学习服务是一个计算服务,和我以前做过的分布式存储服务有着本质的区别,计算服务离线运算时间较长,客户端请求延时要求较低而且吞吐很小,因此我们的 API 服务在易用性和高性能上可以选择前者,目前主流的 Web 服务器都可以满足需求。基于 Web 服务器我们可以实现集成内部权限管理系统的业务逻辑,小米生态云提供了类似 AWS 的 AKSK 签名认证机制,用户注册登录后可以自行创建 Access key 和 Secret key,请求时在客户端进行 AKSK 的签名后发送,这样用户不需要把账号密码或密钥加到请求中,即使密钥泄露也可以由用户来禁用,请求时即使签名被嗅探也只能重放当前的请求内容,是非常可靠的安全机制。除此之外,我们参考 OpenStack 项目的体系架构,实现了多租户和 Quota 功能,通过认证和授权的请求需要经过 Quota 配额检查,在高可用数据库中持久化相应的数据,这样平台管理员就可以动态修改每个租户的 Quota,而且用户可以随时查看自身的审计信息。
小米 cloud-ml 服务实现了深度学习模型的开发、训练、调优、测试、部署和预测等完整功能,都是通过提交到后端的 Kubernetes 集群来实现,完整的功能介绍可以查看官方[文档](http://docs.api.xiaomi.com/cloud-ml/)。Kubernetes 对外提供了 RESTful API 访问接口,通过 YAML 或者 JSON 来描述不同的任务类型,不同编程语言实现的系统也可以使用社区开发的 SDK 来访问。对于我们支持的多个深度学习框架,还有开发环境、训练任务、模型服务等功能,都需要定制 Docker 镜像,提交到 Kubernetes 时指定使用的容器镜像、启动命令等参数。通过对 Kubernetes API 的封装,我们可以简化 Kubernetes 的使用细节,保证了对 Mesos、Yarn 等后端支持的兼容性,同时避免了直接暴露 Kubernetes API 带来的授权问题以及安全隐患。
除了可以启动单个容器执行用户的训练代码,小米 cloud-ml 平台也支持 TensorFlow 的分布式训练,使用时只需要传入 ps 和 worker 个数即可。考虑到对 TensorFlow 原生 API 的兼容性,我们并没有定制修改 TensorFlow 代码,用户甚至可以在本地安装开源的 TensorFlow 测试后再提交,同样可以运行在云平台上。但本地运行分布式 TensorFlow 需要在多台服务器上手动起进程,同时要避免进程使用的端口与其他服务冲突,而且要考虑系统环境、内存不足、磁盘空间等问题,代码更新和运维压力成倍增加,Cloud Machine Learning
下的分布式 TensorFlow 只需要在提交任务时多加两个参数即可。有人觉得手动启动分布式 TensorFlow 非常繁琐,在云端实现逻辑是否更加复杂?其实并不是,通过云服务的控制节点,我们在启动任务前就可以分配不会冲突的端口资源,启动时通过容器隔离环境资源,而用户不需要传入 Cluster spec 等繁琐的参数,我们遵循 Google CloudML 标准,会自动生成 Cluster spec 等信息通过环境变量加入到容器的启动任务中。这样无论是单机版训练任务,还是几个节点的分布式任务,甚至是上百节点的分布式训练任务,cloud-ml 平台都可以通过相同的镜像和代码来运行,只是启动时传入的环境变量不同,在不改变任何外部依赖的情况下优雅地实现了看似复杂的分布式训练功能。
图2 云深度学习平台分布式训练
看到这里大家可能认为,小米的 cloud-ml 平台和 Google 的 CloudML 服务,都有点类似之前很火的
PaaS(Platform as a Service)或者 CaaS(Container as a Service)服务。确实如此,基于 Kubernetes 或者 Mesos 我们可以很容易实现一个通用的 CaaS,用户上传应用代码和 Docker 镜像,由平台调度和运行,但不同的是 Cloud Machine Learning 简化了与机器学习无关的功能。我们不需要用户了解 PaaS 的所有功能,也不需要支持所有编程语言的运行环境,暴露提交任务、查看任务、删除任务等更简单的使用接口即可,而要支持不同规模的 TensorFlow 应用代码,用户需要以标准的 Python 打包方式上传代码。
Python 的标准打包方式独立于 TensorFlow 或者小米 cloud-ml 平台,幸运的是目前 Google CloudML 也支持 Python 的标准打包方式,通过这种标准接口,我们甚至发现 Google CloudML 打包好的 samples 代码甚至可以直接提交到小米 cloud-ml 平台上训练。这是非常有意思的尝试,意味着用户可以使用原生的 TensorFlow 接口来实现自己的模型,在本地计算资源不充足的情况下可以提交到 Google CloudML 服务上训练,同时可以一行代码不用改直接提交到小米或者其他云服务厂商中的云平台上训练。如果大家在实现内部的云深度学习平台,不妨也参考下标准的 Python 打包方式,这样用户同一份代码就可以兼容所有云平台,避免厂商绑定。
除了训练任务,Cloud Machine Learning 平台最好也能集成模型服务、开发环境等功能。对于模型服务,TensorFlow 社区开源了 TensorFlow Serving 项目,可以加载任意 TensorFlow 模型并且提供统一的访问接口,而 Caffe 社区也提供了 Web demo 项目方便用户使用。目前 Kubernetes 和 Mesos 都实现了类似 Deployment 的功能,通过制作 TensorFlow Serving 等服务的容器镜像,我们可以很方便地为用户快速启动对应的模型服务。通过对 Kubernetes API 的封装,我们在暴露给用户 API 时也提供了 replicas 等参数,这样用户就可以直接通过 Kubernetes API 来创建多副本的 Deployment 实例,并且由 Kubernetes 来实现负载均衡等功能。除此之外,TensorFlow Serving 本身还支持在线模型升级和同时加载多个模型版本等功能,我们在保证 TensorFlow Serving 容器正常运行的情况下,允许用户更新分布式对象存储中的模型文件就可以轻易地支持在线模型升级的功能。
对于比较小众但有特定使用场景的深度学习框架,Cloud Macine Learning 的开发环境、训练任务和模型服务都支持 Bring Your Own Image 功能,也就是说用户可以定制自己的 Docker 镜像并在提交任务时指定使用。这种灵活的设置极大地降低了平台管理者的维护成本,我们不需要根据每个用户的需求定制通用的 Golden image,事实上也不可能有完美的镜像可以满足所有需求,用户不同的模型可能有任意的 Python 或者非 Python 依赖,甚至是自己实现的私有深度学习框架也可以直接提交到 Cloud Machine Learning 平台上训练。内测 BYOI 功能时,我们还惊喜地发现这个功能对于我们开发新的深度学习框架支持,以及提前测试镜像升级有非常大的帮助,同时用户自己实现的 Caffe 模型服务和 XGBoost
模型服务也可以完美支持。
当然 Cloud Machine Learning 平台还可以实现很多有意思的功能,例如通过对线上不同 GPU 机型打 label,通过 NodeSelector 功能可以允许用户选择具体的 GPU 型号进行更细粒度的调度,这需要我们暴露更底层 Kubernetes API 实现,这在集群测试中也是非常有用的功能。而无论是基于 GPU 的训练任务还是模型服务,我们都制作了对应的 CUDA 容器镜像,通过 Kubernetes 调度到对应的 GPU 计算节点就可以访问本地图像处理硬件进行高性能运算了。小米 cloud-ml 还开放了前置命令和后置命令功能,允许用户在启动训练任务前和训练任务结束后执行自定义命令,对于不支持分布式存储的深度学习框架,可以在前置命令中挂载 S3 fuse和 FDS fuse 到本地目录,或者初始化 HDFS 的 Kerberos 账号,灵活的接口可以实现更多用户自定义的功能。还有超参数自动调优功能,与 Google CloudML 类似,用户可以提交时指定多组超参数配置,云平台可以自动分配资源起多实例并行计算,为了支持读取用户自定义的指标数据,我们实现了类似 TensorBoard 的 Python 接口直接访问 TensorFlow event file 数据,并通过命令行返回给用户最优的超参数组合。最后还有 TensorFlow Application Template 功能,在 Cloud Machine Learning 平台上用户可以将自己的模型代码公开或者使用官方维护的开源 TensorFlow 应用,用户提交任务时可以直接指定这些开源模板进行训练,模板已经实现了 MLP、CNN、RNN 和 LR 等经典神经网络结构,还可以通过超参数来配置神经网络每一层的节点数和层数,而且可以支持任意稠密和稀疏的数据集,这样不需要编写代码就可以在云平台上训练自己的数据快速生成 AI 模型了。
在前面的平台设计和平台架构后,要实现完整的云深度学习服务并不困难,尤其是集成了Docker、Etcd、Kubernetes、TensorFlow 等优秀开源项目,组件间通过 API 松耦合地交互,需要的重复工作主要是打通企业内部权限系统和将用户请求转化成 Kubernetes 等后端请求而已,而支持标准的打包方式还可以让业务代码在任意云平台上无缝迁移。
### 云深度学习平台实践
目前小米云深度学习平台已经在内部各业务部门推广使用,相比于直接使用物理机,云服务拥有超高的资源利用率、快速的启动时间、近乎“无限”的计算资源、自动的故障迁移、支持分布式训练和超参数自动调优等优点,相信可以得到更好的推广和应用。
除了完成上述的功能,我们在实践时也听取了用户反馈进行改进。例如有内部用户反馈,在云端训练的 TensorFlow 应用把 event file 也导出到分布式存储中,使用 TensorBoard 需要先下载文件再从本地起服务相对麻烦,因此我们在原有基础架构实现了 TensorboardService 功能,可以一键启动
TensorBoard 服务,用户只需要用浏览器就可以打开使用。
管理 GPU 资源和排查 GPU 调度问题也是相当繁琐的,尤其是需要管理不同 GPU 设备和不同 CUDA 版本的异构集群,我们统一规范了 CUDA 的安装方式,保证 Kubernetes 调度的容器可以正常访问宿主机的 GPU 设备。当然对于 GPU 资源的调度和释放,我们有完善的测试文档可以保证每一个 GPU 都可以正常使用,根据测试需求实现的 NodeSelector 功能也帮忙我们更快地定位问题。
由于已经支持几十个功能和十几个深度学习框架,每次升级都可能影响已有服务的功能,因此我们会在多节点的分布式 staging 集群进行上线演习和测试,并且实现 smoke test 脚本进行完整的功能性测试。服务升级需要更新代码,但是为了保证不影响线上业务,无论是
Kubernetes 还是我们实现的 API Server 都有多副本提供服务,通过高可用技术先迁移服务进行滚动升级,对于一些单机运行的脚本也通过 Etcd 实现了高可用的抢主机制,保证所有组件没有单点故障。
大家可以通过前面提到的文档地址和 cloud-ml-sdk 项目了解到更多细节,或者关注我微博(@tobe-陈迪豪)与我交流。
### 总结
本文介绍了实现企业级云深度学习平台需要的概念和知识,基于小米 cloud-ml 服务探讨了云平台的设计、架构、实现以及实践这四方面的内容,希望大家看完有所收获。