基于Kubeflow的机器学习调度平台落地实战

作者：范德良 周佳煊 张振华

机器学习，特别是深度学习，在蘑菇街这样的电商平台有大量实际业务的落地场景，比如搜索推荐、图像算法、交易风控反作弊等等。随着业务的快速发展，之前已有的基于Yarn的调度平台已经无法满足大规模机器学习的计算需求，因此我们在2018年和算法工程团队一起建设了基于Kubeflow和Kubernetes的分布式机器学习平台，并深入到业务层面进行分布式改造，并且从Kubernetes、Tensorflow和业务等多个层面进行了一系列的性能调优。

背景

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

痛点

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

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

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

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

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

上述方案无一例外的将各个部门分散的机器纳入到统一的资源池，并提供资源管理和调度功能，让基于Excel表的资源管理和人肉调度成为过去，让用户更专注于算法模型，而非基础设施。在几十个worker下，无论是CPU还是GPU的分布式训练，训练速度都能得到近乎线性的提升，将单机的训练时间从月级别缩短到一天以内，提升效率的同时也大大提升了资源利用率。

蘑菇街早期的业务方往往独立维护各自团队的GPU机器“小池子”，机器的资源分配和管理存在盲区，缺乏统一管理和调度。GPU的资源存在不均衡和资源利用率低下的问题。事实上，大数据团队之前已经将CPU类型的训练任务接入了Xlearning，尝到了分布式训练的甜头，但也发现一些问题：

1．公司目前的Yarn不支持GPU资源管理，虽然近期版本已支持该特性，但存在稳定性风险。

2．缺乏资源隔离和限制，同节点的任务容易出现资源冲突。

3．监控信息不完善。在发生资源抢占时，往往无法定位根本原因。

4．缺少弹性能力，目前资源池是静态的，如果能借助公有云的弹性能力，在业务高峰期提供更大的算力，将能更快的满足业务需求。

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

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

人肉部署

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

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

此外，不同GPU机型依赖的Nvidia驱动也不同，较新的卡，比如V100所依赖的版本更高。人肉部署还需要管理和维护多个不同的驱动版本。

训练任务管理

人肉启动训练任务时，需要手动查看/评估资源的剩余可用情况，手动指定PS和Worker的数量，管理配置并进行服务发现。这些都给业务方带来了很大的负担。

解决的思路

Docker和Kubernetes很好的解决了人肉管理的问题：

1．部署：借助Docker良好的隔离性，各个容器的文件系统互不影响，将训练任务和驱动程序制作成镜像，避免了多次安装的繁琐。此外Kubernetes提供了服务发现功能，简化了分布式的部署。

2．训练任务生命周期管理：Kubernetes提供了生命周期管理的API，用户基于API即可一键式完成训练任务的增删改查，避免人工ssh的各种繁琐操作，可以大幅提升用户体验和效率。

痛点三：监控的缺失

监控也是分布式训练重要的一环，它是性能调优的重要依据。比如在PS-Worker的训练框架下，我们需要为每个PS/Worker配置合适的GPU/CPU/Memory，并设置合适的PS和Worker数量。如果某个参数配置不当，往往容易造成性能瓶颈，影响整体资源的利用率。比如当PS的网络很大时，我们需要增加PS节点，并对大参数进行partition；当worker CPU负载过高时，我们应该增加Worker的核数。

早期版本的Hadoop和Yarn并不支持GPU的资源可视化监控，而Kubernetes已拥有非常成熟监控方案Prometheus，它能周期采集CPU，内存，网络和GPU的监控数据，即每个PS/Worker的资源使用率都能得到详细的展示，为优化资源配置提供了依据。事实上，我们也是通过监控信息为不同的训练任务分配合适的资源配置，使得在训练速度和整体的吞吐率上达到一个较好的效果。

痛点四：资源隔离较弱

早期的机器学习平台基于Yarn的Angel和XLearning，由于Yarn缺乏对实例之间的资源隔离，我们在内存，网络，磁盘等均遇到不同程度的问题。

由于Yarn没有对任务的内存进行隔离，所以，业务方常常因对内存资源估计不准而导致worker的进程OOM。由于所有的进程都共用宿主机的IP，很容易造成端口冲突，此外磁盘IO和网络IO也很容易被打爆。

Kubeflow及核心组件

Kubeflow是由Google等公司于2017年推出的基于Kubernetes的开源项目，它支持常见的机器学习框架。

Kubeflow简介

The Kubeflow project is dedicated to making deployments of machine learning（ML）workflows on Kubernetes simple, portable and scalable.

Kubeflow旨在支持多种机器学习框架运行在Kubernetes之上，比如Tensorflow, Pytorch, Caffe等常见框架。它包含了operator、pipeline、超参数调优、serving等诸多模块。它通过提供对应的operator，基于Kubernetes的Pod/headless Service等基础资源为框架提供与之相配的更高层次的资源。比如tf-operator为Tensorflow提供了job维度的生命周期管理能力，以满足Tensorflow分布式训练的资源和拓扑需求，达到了一键式部署Tensorflow训练任务的效果。

Kubeflow包含如下operator，分别对应主流的分布式计算框架。蘑菇街主要采用了kubeflow中的tf-operator，来实现对机器学习任务的统一管理。

Distributed Tensorflow

蘑菇街业务方使用的计算框架主要是Tensorflow，因此有必要先介绍一下Tensorflow的分布式训练，Tensorflow支持如下三种分布式策略：

1．MirroredStrategy：适用于单机多卡的训练场景，功能有限，不在本文讨论范围内。

2．ParameterServerStrategy：用于多机多卡场景，主要分为worker节点和PS节点，其中模型参数全部存储在PS节点，worker在每个step计算完梯度后向PS更新梯度，蘑菇街当前使用这种方案。

3．CollectiveAllReduceStrategy：用于多机多卡场景，通过all-reduce的方式融合梯度，只需要worker节点，不需要PS节点，从另外一个角度说，该节点既充当worker角色，又充当PS角色。该方案是带宽优化的，具有性能好，可扩展性强的特点，是Tensorflow推荐的未来方向。

以ParameterServerStrategy为例，一个分布式训练集群至少需要两种类型的节点：PS和worker。由于在训练中需要一个worker节点来评估效果和保存checkpoint，因此单独把该节点作为chief（或者叫master）节点。通常情况下，一个集群需要多个worker节点，多个PS节点，一个chief节点。所有worker节点的CPU/内存/GPU等资源配置完全相同，所有PS节点的CPU/内存等资源配置也相同。从资源拓扑角度出发，如果能够提供一种Kubernetes资源，用户可以基于该资源定义PS/worker/chief的数量和规格，用户就可以一键式创建分布式集群，大大简化了分布式集群的部署和配置。tf-operator定义了TFJob资源，用户可以借助tf-operator在Kubernetes上一键拉起分布式训练集群。

从Tensorflow分布式训练的使用方式出发，用户在启动每个节点的任务时，需要传入集群所有节点的网络信息。这意味着分布式训练集群的每个节点需要预先知道所有其它节点的网络地址信息，即要求服务发现功能。tf-operator基于Kubernetes headless service，完美的提供了服务发现功能，无需用户再手工指定PS/Worker的IP信息，极大的降低了用户的部署成本。

落地实践

主要包含了以下几部分。

Tensorflow生命周期管理（tf-operator）

在部署tf-operator之后，首先需要在Kubernetes中创建对应的TFJob CRD，之后就可以创建TFJob资源了。

在如下的样例中，我们定义了一个具有10个worker,4个ps,1个chief的分布式训练集群。从TFJob参数不难发现，它对ParameterServerStrategy和CollectiveAllReduceStrategy这两种策略方式都支持，只是在CollectiveAllReduceStrategy场景下，无需要设置PS资源。

        apiVersion: "kubeflow.org/v1alpha1"
        kind: "TFJob"
        metadata:
          name: "example-job"
        spec:
          replicaSpecs:
            - replicas: 1
              tfReplicaType: CHIEF
              template:
                spec:
                  containers:
                    - image: gcr.io/tf-on-k8s-dogfood/chief_sample:latest
                      name: tensorflow
                  restartPolicy: OnFailure
            - replicas: 10
              tfReplicaType: WORKER
              template:
                spec:
                  containers:
                    - image: gcr.io/tf-on-k8s-dogfood/worker_sample:latest
                      name: tensorflow
                  restartPolicy: OnFailure
            - replicas: 4
              tfReplicaType: PS
              template:
                spec:
                  containers:
                    - image: gcr.io/tf-on-k8s-dogfood/ps_sample:latest
                      name: tensorflow

Tf-operator启动后，通过list-watch不断的监听TFJob资源相关事件，当收到创建TFJob事件时，tf-operator依次创建PS/Worker/Chief（Master）Replica资源。以PS Replica为例，根据replicas数量依次创建等同数量的pod，并为每个pod创建headless service。此外，它还生成TF_CONFIG环境变量，这个环境变量记录了所有pod的域名和端口，最终在创建pod时候注入到容器中。

任务调度（kube-batch）

通过引入kube-batch，满足了业务方对批量任务调度的需求。没有采用Kubernetes默认的调度器，主要是基于以下两点考虑：

GANG scheduling:Kubernetes默认的调度器是以pod为粒度的，并不支持GANG scheduling。机器学习的训练任务要求集群保持一个整体，要么所有的pod都能成功创建，要么没有一个pod能被创建。试想资源处于临界状态时，如果采用默认的调度器调度一个分布式训练集群，则会导致部分节点因资源不足而无法成功调度，那些成功创建的pod空占用资源，但是无法进行训练。

任务排队：Kubernetes默认的调度器无法提供队列调度的功能，而不会像Hadoop任务那样进行排队。而面向机器学习/大数据/HPC的批量任务调度，往往要求系统能维护一个或多个队列，进行任务排队。当任务结束并且释放资源后，后续的任务可以继续执行。

Kube-batch目前是Kubernetes SIGs旗下的一个孵化项目，是一个运行在Kubernetes上面向机器学习/大数据/HPC的批调度器（batch scheduler），它支持以Pod Group为单位进行资源调度，并支持preempt和priority。对于暂时无法满足资源条件的Pod，在Kubernetes中会处于pending状态，直到资源释放从而继续执行。

Kube-batch的基本流程如下图，它通过list-watch监听Pod, Queue, PodGroup和Node等资源，在本地维护一份集群资源的全局缓存，依次通过如下的策略（reclaim, allocate, preemption, predict）完成资源的调度。

ikube-batch旨在通过配置策略，提供gang scheduling的调度能力，由于是基于queue的调度思想，kube-batch对机器学习及大数据任务的调度具有很大的潜力。

由于kube-batch对任务的actions有四种，可以根据自身的业务特点，指定只使用其中的一种或者多种：如果简单的分配，使用allocation即可；如果资源存在碎片，让调度器能够感知并重新分配，可以将allocation与backfill结合起来使用。如此既可满足业务的调度需求，又可以在一定程度上提升调度性能。

认证和授权（Dex）

我们引入了由CoreOS研发的dex组件，并嵌入到我们的SDK中，实现了与公司LDAP的对接。我们将Kubernetes中的namespace映射成不同的应用组，应用的负责人会自动授予（Rolebinding）管理员的角色（Role），同时针对一些特权用户还会创建clusterRole及clusterRolebinding，方便这些特权用户访问某些需要高级权限的API（如查询node信息等）。

多租户的隔离性

平台需要接入多个不同的业务方，自然需要考虑多租户间的隔离性。我们利用Kubernetes提供的resource quota，以namespace为单位对资源进行分配。

性能调优

这里介绍一下蘑菇街在分布式机器学习调优的一些经验，主要分为

Kubernetes层面、Tensorflow层面和业务层面的一些调优。

Kubernetes层面调优

CPUShare VS CPUSet

以商品推荐常用的wide and deep模型为例，该模型采用CPU资源进行训练，宿主机的规格为80C256G。

从直觉上说，训练本质上是大量的矩阵运算，而矩阵运算在内存地址空间具有良好连续性的特点。若充分利用CPU cache，将有助于提升cache的命中率，最终加快训练速度。Kubernetes从1.8开始支持CPU Manager特性，该特性支持对于Guaranteed类型的pod，采用CpuSet为Pod分配独占的CPU core。在相同训练任务下，对比CpuSet和CpuShare对训练速度的影响，发现在worker CPU核数较少的情况下，CpuSet的性能远远超过CpuShare。

虽然Guaranteed类型的Pod牺牲了资源弹性，但是CpuSet带来的性能收益要高于弹性带来的收益。即使对CpuShare容器不设置cpu limits，当跑满容器的CPU资源时，相同cpu requests下，CpuSet容器的性能依旧比CpuShare的性能高出20%以上。

在kubelet中开启CPU Manaer特性的配置参数如下

        --feature-gates=CPUManager=true
        --cpu-manager-policy=static
        --cpu-manager-reconcile-period=5s
        --kube-reserved=cpu=500m

Pod Affinity

PS/Worker训练框架下，PS节点作为中心节点，网络流量往往非常大，容易把带宽跑满。通过Pod AntiAffinity将所有PS节点尽可能打散到不同的宿主机上，从而分摊网络流量。

            podAntiAffinity:
              preferredDuringSchedulingIgnoredDuringExecution:
              - weight: 100
                podAffinityTerm:
                  labelSelector:
                    matchExpressions:
                    - key: tf-replica-type
                      operator: NotIn
                      values:
                      - ps
                  topologyKey: kubernetes.io/hostname

设置合适的资源配比

不同模型的计算量和参数各不相同，因此针对每个模型都应该设置一个合适的PS/Worker的规格和数量。在监控完善的条件下，可以根据监控数据分析瓶颈，调整实例的规格和数量，使得整体的训练速度和平台吞吐量能达到较好的水平。

· Kube-batch调度

由于kube-batch默认开启“reclaim, allocate, backfill, preempt”四种actions，导致每次调度时轮询周期较长，通过配置actions为allocate一种可以提高30%的调度效率。

Tensorflow层面调优

Tensorflow的调优主要参考了官网文档。

· 线程数

由于sysconf系统调用等隔离性的问题，容器观察到的CPU信息往往是宿主机的。Tensorflow默认的线程数为CPU核数，如此情况下，Tensorflow创建的线程数远远超过实际分配到的CPU核数。同样以wide and deep模型为例，通过保持与cpu limit

一致的线程数，上下文切换降低约40%，训练速度提升约5%。

        config = tf.ConfigProto()

        config.intra_op_parallelism_threads = cpu_limit_cores
        config.inter_op_parallelism_threads = cpu_limit_cores
        tf.Session(config=config)

业务层面调优

· Partition

在某次训练中发现PS流量节点的分布不均匀，其中某个PS节点的流量非常大，而其它PS节点的流量相对较低。通过分析timeline.json发现某个embedding向量非常大，所以通过partition，将该tensor分散到不同的PS节点上，从而避免了某个PS节点成为瓶颈。

        partitioner = tf.fixed_size_partitioner(ps_number, axis=0)
        with tf.variable_scope("emb_layer", partitioner= partitioner) as
        scope:

· Adam优化器

由于Adam优化器会更新所有参数的梯度，所以在大embedding向量下，如果采用adam优化器，会大大增加计算量，严重影响训练速度。因此建议采用Lazy_Adam_Optimizer或者Adadelta优化器。

总结和体会

目前上述基于Kubernetes的机器学习平台已经在生产环境为多个业务方提供了服务。建设这套平台的过程，也是我们探索的过程。以下我们总结了一些还不尽如人意的地方，以及我们对未来的展望。

tf-operator

从落地的情况来看，tf-operator的功能能满足基本的要求，稳定性较高。但是在故障恢复稍有欠缺，对于pod级别的故障，依赖kubelet来恢复；对于node级别的故障，目前还不支持故障恢复。分布训练下，故障概率随着worker数量和训练时间的增加而增加。worker作为无状态节点，故障恢复既是可行的，也是非常有必要的。

kube-batch

目前kube-batch功能薄弱，成熟度有待商榷。比如kube-batch只有predict功能，没有priority功能。并且predict功能也非常薄弱，仅支持部分基础的filter，比如PodMatchNodeSelector, PodFitsHostPorts以及基本的资源调度等。特别是PodAffinity特性，对于PS-Worker架构非常有用，因为PS节点的网络流量非常大，所以需要PS节点之间反亲和，将各个PS节点分散。此外，kube-batch也不支持多任务之间依赖关系。

Kube-batch的落地效果差强人意，社区的维护力度较低。除了功能薄弱以外，我们也碰上了诸多的问题，处于勉强可用状态。我们建议可将pod群维度的资源判断功能放到上层，只有当空闲资源满足创建整个分布式训练集群时，再将请求发送给Kubernetes，由Kubernetes默认的调度器完成调度，当然，这种方式也存在一些缺点。

GPU虚拟化

目前GPU相关的监控、隔离性以及更细粒度的调度，仍然是业界的痛点问题。Nvidia提供的GPU虚拟化方案需要收取高额的Lincense费用，目前业界还没有免费的GPU虚拟化方案。

在实际的业务场景中，一些业务的GPU使用率并不高，但以虚拟机或者容器方式运行时往往会独占一块GPU卡，虽然可以通过设置CUDA_VISIBLE_DEVICES来实现多个容器共享一块GPU卡，但任务之间的隔离性是无法保证的。

另外，Kubernetes进行GPU资源分配时，默认还不感知GPU的Topology，而不同分配的策略，对训练的性能也会产生很大的影响。YARN和Kubernetes开源社区都在努力增强这块的调度能力。

Kubeflow展望

Kubeflow是目前基于Kubernetes的主流机器学习解决方案，它抽象了和机器学习相关的PS-Worker模型，实现了一套pipeline的工作流，支持超参数训练和Jupyter notebooks集成等能力。

由于Kubeflow解决了基于Yarn的机器学习平台的痛点，同时容器化越来越普及。基于Kubernetes这样大规模的企业级容器平台，将会是未来的方向，相信Kubeflow在2019年将会有更好的发展。