【凡益堂】如何在Kubernetes实现GPU调度及共享
概论
近年来AI技术的繁荣和深化，尤其是深度学习的崛起，离不开海量数据和计算力的提升。尤其是对Nvidia的GPU的利用，让深度学习获得几十倍的性能提升，才彻底打开了AI想象空间。虽然智慧芯片近年来有着百花齐放的风景，最典型的例如Google的TPU，但是平心而论，从普惠意义和生态上，Nvidia的GPU仍然占有主导位置。
不过，Nvidia的GPU无疑是昂贵的，所以如何最大化利用好GPU的硬件资源，是每一个算力平台型产品都要考虑的问题。比如，有多个用户使用GPU服务器进行训练时，如何保证资源合理的分配非常重要。得益于Nvidia公司为Docker写的Runtime，也就是Nvidia-Docker，使得在Docker里使用GPU成为可能。从容器粒度来管理和使用GPU要比从主机角度容易很多，因为运行GPU的AI任务通常配置非常复杂，这种复杂包括管理员从管理GPU卡的分配和使用者切换不同的训练环境，而容器可以封装不同的训练环境，很大程度上降低复杂性。此外，借助Kubernetes来管理Nvidia-Docker，使得GPU任务的分配更加简单和合理，目前已成为几乎所有主流的AI算力平台的方案。
Kubernetes支持通过Device-Plugin的方式来增加对默认资源(CPU,Memory等)之外的设备支持，而第三方可以通过编写相应的Device-Plugin来增加对设备的支持。目前Nvidia也是通过这样的方式对GPU进行支持。
K8s + Nvidia-Device-Plugin的方式两个限制：每一块GPU同时最多只能被一个容器使用；没有考虑GPU卡之间的通道亲和性。 这样的方式已经满足了部分AI算力需求场景，但是某些场景下还是有缺点和限制的: 每一块GPU同时最多只能被一个容器使用，这在训练模式下没有任何问题，但是在开发调试模式下会造成巨大的资源浪费。因为开发调试模式下用户大部分时间并没有实质运行GPU的资源，但却排他的独占了昂贵的GPU。此外在多卡主机架构下，GPU卡直接的连接通常是不一样的，有的通过Nvlink相连，有的通过PCIe，而不同的连接方式性能差别非常大。而没有考虑同个主机里GPU卡直接的通道亲和性时，也会给多卡计算时需要发生数据传输时(如all_reduce操作)带来过高的通信开销。那么自然而然，在同一个容器挂载多张GPU卡时，我们当然更希望会挂载通道亲和性更好的卡。
本文会介绍K8s进行GPU调度的通用流程和我们的一些改造方案。包括用于支持容器GPU挂载的Nvidia-Docker、K8s中将GPU作为拓展资源调度的Device-Plugin机制，以及针对原生Nvidia-Device-Plugin存在的问题的改造方案。
Nvidia-Docker的简单介绍
Nvidia-Docker是Nvidia官方对容器做的拓展，用以让容器能够支持Nvidia的GPU设备。据官方统计的数据标明，目前Nvidia-Docker的下载量已经超过200万次，可以得知目前使用Nvidia-Docker来做AI系统环境已经是非常主流的做法。
这里不详细介绍Nvidia-Docker的原理了，详细的原理可以阅读其官方的设计文档，这里只作简单的介绍。从2015年开始，Docker容器的诞生了一套容器运行时标准OCI(Open Containers Initiative)，它包含容器运行时标准(Runtime-Spec)和 容器镜像标准(Image-Spec)。而著名的Runc则是这套标准的一个默认实现，然而任何满足该标准的实现都可以注册为容器的运行时拓展。Containerd则包装了Runc和其它功能如生命周期管理等，以Daemon的形式运行在主机。Nvidia-Docker正是基于这一套拓展标准增加Nvidia GPU的容器支持。Nvidia-Docker主要原理是将对GPU的支持放入一个兼容OCI标准的运行时库拓展libnvidia-container中，并在Runtime的API中进行调用，在libnvidia-container中通过共享和调用主机侧的nvidia-driver实现对GPU的支持。在容器启动时，Runc会调用一个叫做nvidia-container-runtime-hook的hook，这个hook会去检查相应的环境是否具有GPU的支持和一些环境检查，完成之后容器启动，在运行时容器内进程也是通过libnvidia-container暴露的接口进行交互，从而实现容器对GPU的透传和运行时支持。
（图片来源：https://devblogs.nvidia.com/gpu-containers-runtime）
值得注意的是，Nvidia-Docker容器会使用主机侧的Nvidia-Driver，再上层的软件栈如cuda/cudnn，AI框架等，则在容器里面提供。此外，多个Nvidia-Docker可以挂载同一个GPU，只要通过环境变量指定就好，并没有数量上的限制。
为了方便理解后面Device-Plugin的机制，这里简单介绍一下Nvidia-Docker挂载不同GPU设备的方式。Nvidia-Docker的使用非常简单，它通过指定一个环境变量来指定将要挂载的GPU设备，不过要在Docker的配置文件中指定Docker的Runtime为Nvidia-Docker，或者通过命令行显式指定也可以：
nvidia-docker run -e NVIDIA_VISIBLE_DEVICES=0,1 --runtime=nvidia -it tensorflow/tensorflow-gpu:v1.13 bash滑动查看
如果在Docker配置中已经做过相关配置，那么就可以简化为：
docker run -e NVIDIA_VISIBLE_DEVICES=0,1 -it tensorflow/tensorflow-gpu:v1.13 bash滑动查看
这里NVIDIA_VISIBLE_DEVICES这个环境变量用来指定需要绑定的GPU卡的逻辑ID，就可以实现容器中对该卡的绑定，使用上非常简单。
K8s的Device-Plugin机制
K8s通过Device-Plugin的机制对非默认的资源设备进行支持，例如RDMA设备、AMD GPU等，当然也包括本文最关心的Nvidia GPU。通过编写相应的Device-Plugin，第三方资源设备上可以在K8s中添加对相应设备的支持，让用户获得和原生资源(CPU,Memory等)近乎一样的使用体验。
Device-Plugin机制本质上是一个RPC服务。K8s定义了一个RPC调用接口，第三方资源设备方可以通过实现该接口来让该设备得以在K8s侧得以支持，并且在使用方式上和默认资源没有太大区别。Device-Plugin以Daemonset的方式在主机侧运行，并且通过一个Socket文件与Kubelet进行通信，从而通过Kubelet给K8s上报相关信息。部署了该Daemonset的主机节点在k8s看来会包含由Device-Plugin注册的硬件资源。Device-Plugin总的原理如下：
（图片来源：https://medium.com/@Alibaba_Cloud）
首先，Device-Plugin需要向K8s注册该资源，注册机制通过实现以下RPC接口得以实现：
service Registration { rpc Register(RegisterRequest) returns (Empty) {}}滑动查看
在详细的rpc调用中，该接口会上报socket名称、Device-Plugin的Api Version等信息，当然，更重要的是它会上报ResourceName,该ResourceName会被K8s登记为该自定义设备的名称，而名称的规范是vendor-domain/resource，例如，Nvidia的GPU就被定义为nvidia.com/gpu，在用户申请资源时，就需要使用该名称。例如，在创建POD的资源配置里，需要这样指定该资源：
apiVersion: v1kind: Podmetadata: name: demo-podspec: containers: - name: demo-container-1 image: k8s.gcr.io/pause:2.0 resources: limits: nvidia.com/gpu: 2 滑动查看
注册之后，Device-Plugin还需要上报主机侧的设备的数量和状态，例如，如果主机节点上有8块GPU卡，Device-Plugin会将该数量的资源数量和资源id列表告知K8s。此外，当有Pod向K8s申请该资源时，K8s会从上报的id列表中按照一定策略返回满足用户需求的资源数量的id序列，当该id列表返回给Device-Plugin，再由Device-Plugin根据一定策略映射到真正的资源设备。以上的过程主要由以下的RPC调用实现的：
service DevicePlugin { rpc ListAndWatch(Empty) returns (stream ListAndWatchResponse) {} rpc Allocate(AllocateRequest) returns (AllocateResponse) {} }滑动查看
这里的ListAndWatch将由Device-Plugin调用NVML库获取主机侧的GPU设备和状态，并返回给k8s相应的设备列表。而Allocate将在容器创建被调用，用来返回一些能够使用主机上该资源的特殊配置，比如一些环境变量，再将这些信息给到Kubelet，并在容器启动的时候传给容器。对于Nvidia GPU而言，主要是前面的提到的环境变量NVIDIA_VISIBLE_DEVICES，这样容器启动的时候就能够挂载相应的GPU设备了。
值得注意的是，容器侧并没有任何gpu虚拟化或者显存分配的策略，所以Nvidia-Device-Plugin分配的粒度是单张卡，并且绝对的一一映射，即上报的GPU数量就是绝对的GPU数量，k8s侧负责资源的分配，再由Device-Plugin受K8s的返回的需要挂载GPU数量和ID，并将容器映射到实际的GPU设备上。
这样的方式在某些场景下是合理的，比如在强计算量任务的模式下，能够避免不同进程对GPU卡资源的争抢以致发生显存OOM等现象。但是在某些场景下缺会造成巨大的资源浪费。比如有某些容器环境仅仅是给到用户进行算法的开发和调试，这样的任务在绝大部分时间里并不会实际使用GPU，这种情况下能够让容器可以适当的共享GPU是有很大价值的，毕竟GPU的价格非常昂贵，我们需要提供共享机制最大化资源的使用。此外Nvidia-Device-Plugin并没有考虑GPU的亲和性，这有可能会让单容器多卡的容器遭遇较差的计算性能。这里会介绍我们的实现思路。
如何让不同容器共享GPU？
前面介绍过，Device-Plugin通过ListAndWatch接口上报GPU资源列表，那么自然而然，我们会想，如果能够伪造出更多虚拟的GPU ID给K8s，K8s在分配POD资源的时候返回虚拟的id，再由Device-Plugin映射回真实的id从而实现GPU卡的复用，就像虚拟内存地址到物理内存地址映射一样，而虚拟内存可以比物理内存大很多。是的，主要的思路就是这样，构造出虚拟的GPU设备id，在容器真正启动的时候再映射到真实的GPU设备。但是有一个非常重要的问题需要解决：怎么样保证GPU负载的大体上的平衡，而不会出现某些卡上绑定了太多的容器，而另一些卡上则没有任何容器？这个现象是有可能出现的，容器的寿命不一，不断的创建和删除容器极有可能会造成GPU资源分配的不均匀。所以虚拟id到真实id的映射不能够是通过一个简单的线性映射关系来决定，而是需要通过考虑在任务创建时的真实的GPU负载来动态的决定挂载GPU的id。
解决这个问题我们评估过两个方案：调用NVML获取GPU卡的实时状态，并选择负载较少的卡进行分配；借助外部数据库存放每个节点的GPU负载信息，Device-Plugin在Allocate的时候调用Redis的信息查看负载情况。我们最终采用的是第二种方法，因为NVML只能够查询进程占用、资源占用的情况，如果一个容器绑定了一块GPU卡，但容器里面并没有任何进程使用GPU，那么用NVML并不能查看真实的容器绑定关系。我们现阶段还是希望把负载局限在卡上绑定的容器数，而不是真实的GPU使用率。
当决定采用借助于一个外部的Redis数据库的方案，用来存放每个节点的实时动态的状态，我们会在redis中为每一个节点维护一个单独的map，记录每个GPU id上分配的容器数量和任务名称。当Device-Plugin每次决定Allocate分配时，会去查询Redis中查询该节点下各个GPU id的容器负载，选择最低的部分进行绑定，并且将相应的GPU id上的容器负载增加1。当资源释放时，Device-Plugin并不能知道释放的消息，我们通过K8s的调度服务的Informer机制，在自定义的Informer中捕捉到释放的POD的节点信息和任务名称，并以此并将Redis中相应的GPU id的资源数减去1。通过这种方式维护和监控着GPU资源的分配信息。这里仅仅介绍解决的思路，具体细节不再展开。
如何让多卡任务绑定亲和性高的卡？
GPU的通道亲和性在多卡训练任务中非常重要，因为不同的连接介质对卡与卡之间的数据传输速度影响非常大。以下是一个典型的8卡GPU的卡间通道拓扑图。可以看出有的卡之间是通过Nvlink(NV1等)相连，有的是通过PCIe(PIX)相连。
而不同的GPU通道会导致完全不一样的数据传输性能，通常它们之间的速度传输能相差很多倍，例如，Nvlink可以达到几十GB/s，而PCIe通常只有10 GB/s左右的吞吐性能。下图是Nvidia Tesla P100的系列的直观的连通拓扑图和通道传输性能：
（图片来源：https://www.nvidia.com）
正如前面所说，Nvidia-Device-Plugin并没有考虑通道的亲和性，也就是说在一个单容器双卡的容器中，通过K8s的调度极有可能将两个通过PCIe连接的卡绑定到一起，即便有在同一个Nvlink通道的卡的存在，这显然是不合理的。高亲和性和容器负载均衡有时会是相互矛盾的需求，如果追求绝对的亲和性，那么可能要牺牲容器任务的负载均衡。我们采用的算法策略是尽量让这两个指标取得一个均衡。
如果不考虑真实的GPU任务负载，单纯的让高亲和性的卡绑定到一起是比较容易实现的。类似于共享GPU实现的思路，在多卡绑定的任务重，我们可以在Device-Plugin里面调用NVML，获得GPU卡的连接拓扑，从而知道它们之间的亲和性关系。然后当Allocate的时候，选择让高亲和性通道间的GPU卡分配到一起即可。但是，如果考虑到高亲和性的卡中间有部分的容器任务负载非常高，那么这个时候可能要高负载的影响。比较合理的方式是使用评分函数，按照一定策略给不同的可选组合评分，选择得分最高的组合。但是我们采用较简单和直接的策略：首先选出负载最小的一个GPU id，再选择跟这个id在同一高亲和性通道的GPU卡，并挑选其中任务负载最小的卡进行绑定。具体的细节不再展开了，主要是对NVML库的调用，可以拿到主机的通道拓扑，剩下的工作是顺理成章的。
总结
本文简单介绍了Docker对Nvidia GPU的支持方案以及K8s的Device-Plugin机制。并且针对现有的Nvidia-Device-Plugin的某些场景缺陷提供解决思路。主要是针对GPU卡的任务间共享和通道亲和性的优化。然而，这些改造都是要看场景的，有的场景值得和迫切需要这么做，但是有的场景却非常不适合。这样的改造会增加外部依赖并且让Nvidia-Device-Plugin的GPU卡的绑定策略变得更加复杂，所以我个人强烈建议只有在必要的时候进行这些改造。而平台的交互式测试、验证场景，正是这样改造的场景和动力。