助力深度学习！阿里开源可插拔 GPU 共享调度工具

根据 Gartner 对全球 CIO 的调查结果显示，人工智能将成为 2019 年组织革命的颠覆性力量。对于人工智能来说，算力即正义，成本即能力，利用 Docker 和 Kubernetes 代表云原生技术为 AI 提供了一种新的工作模式，将 GPU 机器放到统一的资源池进行调度和管理，这避免了GPU 资源利用率低下和人工管理的成本。因此，全球主要的容器集群服务厂商 Kubernetes 都提供了 Nvidia GPU 容器集群调度能力，但是通常都是将一个 GPU 卡分配给一个容器。这虽然可以实现比较好的隔离性，确保使用 GPU 的应用不会被其他应用影响；对于深度学习模型训练的场景也非常适合，但是，针对模型开发和模型预测的场景还是会显得比较浪费。基于此，大家有了共享 GPU 的集群调度需求。

Kubernetes 共享 GPU 集群调度

共享 GPU 的集群调度就是能够让更多的模型开发和预测服务共享同一个 GPU 卡，进而提高集群中 Nvidia GPU 的利用率。而这就需要提供 GPU 资源的划分，而这里 GPU 资源划分的维度指的就是 GPU 显存和 Cuda Kernel 线程的划分。通常在集群级别谈支持共享 GPU 是以下两件事情：

1.调度

2.隔离，我们这里主要讨论的是调度，隔离的方案目前需要用户通过应用限制(比如使用 Tensorflow 的per_process_gpu_memory_fraction 来控制)，未来会提供基于 Nvidia 的 MPS 的可选项, 也会考虑 GPU 的方案。

而对于细粒度的 GPU 卡调度，目前 Kubernetes 社区并没有很好的方案，这是由于 Kubernetes 对于 GPU 这类扩展资源的定义仅仅支持整数粒度的加加减减，无法支持复杂资源的分配。比如用户希望使用 Pod A 占用半张 GPU卡，这在目前 Kubernetes 的架构设计中无法实现资源分配的记录和调用。这里挑战是多卡 GPU 共享是实际矢量资源问题，而 Extened Resource 是标量资源的描述。

针对此问题，我们设计了一个 Out Of Tree 的共享 GPU 调度方案，该方案依赖于 Kubernetes 的现有的工作机制:

Extended Resource 定义
Scheduler Extender 机制
Device Plugin 机制
Kubectl 的扩展机制

这个 GPU 共享调度扩展的好处是：利用 Kubernetes 的扩展和插件机制实现，对于 API Server，Scheduler，Controller Manager 以及 Kubelet 等核心组件没有侵入性。这就方便了使用者可以在不同 Kubernetes 版本上应用这个方案，无需 rebase 代码和重新构建 Kubernetes 二进制包。

用户场景

集群管理员：“我想提高集群的 GPU 使用率;在开发过程中，多个用户共享模型开发环境。”
应用开发人员：“我希望能够同时在 Volta GPU 上运行多个推理任务。”

目标

能够让使用者通过 API 描述对于一个可共享资源的申请，并能实现该种资源的调度

非目标

不支持该共享资源的隔离
不支持超卖

设计原则

明确问题简化设计，第一步只负责调度和部署，后续再实现运行时显存管控。
有很多的客户明确的诉求是首先可以支持多AI应用可以调度到同一个 GPU 上，他们可以接受从应用级别控制显存的大小，利用类似gpu_options.per_process_gpu_memory_fraction控制应用的显存使用量。那我们要解决的问题就先简化到以显存为调度标尺，并且把显存使用的大小以参数的方式传递给容器内部。
不做侵入式修改
本设计中不会修改 Kubernetes 核心的 Extended Resource 的设计， Scheduler 的实现，Device Plugin 的机制以及 Kubelet 的相关设计。重用 Extended Resource 描述共享资源的申请 API。这样的好处在于提供一个可以移植的方案，用户可以在原生 Kubernetes 上使用这个方案。
按显存和按卡调度的方式可以在集群内并存，但是同一个节点内是互斥的，不支持二者并存；要么是按卡数目，要么是按显存分配。

详细设计

前提：

依旧延用 Kubernetes Extended Resource 定义，但是衡量维度最小单位从 1 个 GPU 卡变为 GPU 显存的 MiB。如果所节点使用的 GPU 为单卡 16GiB 显存，它对应的资源就是 16276MiB；
由于用户对于共享GPU的诉求在于模型开发和模型预测场景，在此场景下，用户申请的GPU资源上限不会超过一张卡，也就是申请的资源上限为单卡。

而我们的工作首先是定义了两个新的 Extended Resource：第一个是 gpu-mem，对应的是 GPU 显存;第二个是 gpu-count，对应的是 GPU 卡数。通过两个标量资源描述矢量资源, 并且结合这一资源，提供支持共享 GPU 的工作机制。下面是基本的架构图：

助力深度学习！阿里开源可插拔 GPU 共享调度工具

核心功能模块:

GPU Share Scheduler Extender: 利用 Kubernetes 的调度器扩展机制，负责在全局调度器 Filter 和 Bind 的时候判断节点上单个 GPU 卡是否能够提供足够的 GPU Mem，并且在 Bind 的时刻将 GPU 的分配结果通过 annotation 记录到 Pod Spec 以供后续 Filter 检查分配结果。
GPU Share Device Plugin: 利用 Device Plugin 机制，在节点上被 Kubelet 调用负责 GPU 卡的分配，依赖 scheduler Extender 分配结果执行。

具体流程:

1. 资源上报

GPU Share Device Plugin 利用 nvml 库查询到 GPU 卡的数量和每张 GPU 卡的显存，通过ListAndWatch()将节点的 GPU 总显存（数量显存）作为另外 Extended Resource 汇报给 Kubelet； Kubelet 进一步汇报给 Kubernetes API Server。举例说明，如果节点含有两块 GPU 卡，并且每块卡包含 16276MiB，从用户的角度来看：该节点的 GPU 资源为 16276 2 = 32552; 同时也会将节点上的 GPU 卡数量 2 作为另外一个 Extended Resource 上报。

2. 扩展调度

GPU Share Scheduler Extender 可以在分配 gpu-mem 给 Pod 的同时将分配信息以 annotation 的形式保留在 Pod spec 中，并且在过滤时刻根据此信息判断每张卡是否包含足够可用的 gpu-mem 分配。

2.1 Kubernetes 默认调度器在进行完所有过滤(filter)行为后会通过 http 方式调用 GPU Share Scheduler Extender的filter 方法, 这是由于默认调度器计算 Extended Resource 时，只能判断资源总量是否有满足需求的空闲资源，无法具体判断单张卡上是否满足需求；所以就需要由 GPU Share Scheduler Extender 检查单张卡上是否含有可用资源。

以下图为例，在由 3 个包含两块 GPU 卡的节点组成的 Kubernetes 集群中，当用户申请gpu-mem=8138时，默认调度器会扫描所有节点，发现 N1 所剩的资源为 (16276 * 2 - 16276 -12207 = 4069 )不满足资源需求，N1 节点被过滤掉。

而 N2 和 N3 节点所剩资源都为 8138MiB，从整体调度的角度看，都符合默认调度器的条件；此时默认调度器会委托 GPU Share Scheduler Extender 进行二次过滤，在二次过滤中，GPU Share Scheduler Extender 需要判断单张卡是否满足调度需求，在查看 N2 节点时发现该节点虽然有 8138MiB 可用资源，但是落到每张卡上看，GPU0 和分别 GPU1 只有 4069MiB 的可用资源，无法满足单卡 8138MiB 的诉求。而 N3 节点虽然也是总共有 8138MiB 可用资源，但是这些可用资源都属于 GPU0，满足单卡可调度的需求。由此，通过 GPU Share Scheduler Extender 的筛选就可以实现精准的条件筛选。

助力深度学习！阿里开源可插拔 GPU 共享调度工具

2.2 当调度器找到满足条件的节点，就会委托 GPU Share Scheduler Extender 的 bind 方法进行节点和 Pod 的绑定，这里 Extender 需要做的是两件事情

以 binpack 的规则找到节点中最优选择的 GPU 卡 id，此处的最优含义是对于同一个节点不同的 GPU 卡，以 binpack 的原则作为判断条件，优先选择空闲资源满足条件但同时又是所剩资源最少的 GPU 卡，并且将其作为ALIYUN_COM_GPU_MEM_IDX保存到 Pod 的 annotation 中；同时也保存该 Pod 申请的 GPU Memory 作为ALIYUN_COM_GPU_MEM_POD和ALIYUN_COM_GPU_MEM_ASSUME_TIME保存至 Pod 的 annotation 中，并且在此时进行 Pod 和所选节点的绑定。

注意：这时还会保存ALIYUN_COM_GPU_MEM_ASSIGNED的 Pod annotation，它被初始化为“false”。它表示该 Pod 在调度时刻被指定到了某块 GPU 卡，但是并没有真正在节点上创建该 Pod。ALIYUN_COM_GPU_MEM_ASSUME_TIME代表了指定时间。

如果此时发现分配节点上没有 GPU 资源符合条件，此时不进行绑定，直接不报错退出，默认调度器会在 assume 超时后重新调度。

调用 Kubernetes API 执行节点和 Pod 的绑定

以下图为例，当 GPU Share Scheduler Extender 要把 gpu-mem：8138 的 Pod 和经过筛选出来的节点 N1 绑定，首先会比较不同 GPU 的可用资源，分别为 GPU0(12207),GPU1(8138),GPU2(4069),GPU3(16276),其中 GPU2 所剩资源不满足需求，被舍弃掉；而另外三个满足条件的 GPU 中, GPU1 恰恰是符合空闲资源满足条件但同时又是所剩资源最少的 GPU 卡，因此 GPU1 被选出。

助力深度学习！阿里开源可插拔 GPU 共享调度工具

3. 节点上运行

当 Pod 和节点绑定的事件被 Kubelet 接收到后，Kubelet 就会在节点上创建真正的 Pod 实体，在这个过程中, Kubelet 会调用 GPU Share Device Plugin 的Allocate方法, Allocate方法的参数是 Pod 申请的 gpu-mem。而在Allocate方法中，会根据 GPU Share Scheduler Extender 的调度决策运行对应的 Pod

会列出该节点中所有状态为 Pending 并且ALIYUN_COM_GPU_MEM_ASSIGNED为false的 GPU Share Pod
选择出其中 Pod Annotation 的ALIYUN_COM_GPU_MEM_POD的数量与 Allocate 申请数量一致的 Pod。如果有多个符合这种条件的 Pod，就会选择其中ALIYUN_COM_GPU_MEM_ASSUME_TIME最早的 Pod。
将该 Pod 的 annotation ALIYUN_COM_GPU_MEM_ASSIGNED设置为true，并且将 Pod annotation 中的 GPU 信息转化为环境变量返回给 Kubelet 用以真正的创建 Pod。

助力深度学习！阿里开源可插拔 GPU 共享调度工具

部署

请参照部署文档

测试样例

首先创建一个使用aliyun.com/gpu-mem的应用

apiVersion: apps/v1
kind: Deployment
metadata:
  name: binpack-1
  labels:
    app: binpack-1
spec:
  replicas: 1
  selector: # define how the deployment finds the pods it manages
    matchLabels:
      app: binpack-1
  template: # define the pods specifications
    metadata:
      labels:
        app: binpack-1
    spec:
      containers:
      - name: binpack-1
        image: cheyang/gpu-player:v2
        resources:
          limits:
            # MiB
            aliyun.com/gpu-mem: 1024