基于 Kubernetes 的 GPU 共享调度与迁移改造实践


背景

众所周知，当前对于深度学习的研究十分火热，如果想要取得好的效果，除了数据和算法两个要素外，强大的算力也是必不可少的。但由于目前主流的 NVIDIA GPU 比较昂贵，并且一般情况下独占卡的模式会对 GPU 这种宝贵的计算资源造成浪费，即不同用户对模型的理解深度不同，导致申请了独立的卡却没有把资源用满。因此，为了解决上述问题，通过 GPU 共享的虚拟化技术来提高资源利用率也逐渐成为当下的研究热点。

研究现状

从学术界和工业界的一些实现方式来看，GPU 共享的基本原理主要分为以下两部分：

  • 资源隔离：限制程序使用的 GPU 核数和显存，常见的思路就是劫持调用，具体可以分为用户态劫持或内核态劫持。
  • 并行模式：共享 GPU 情况下应用程序应该能够并行执行，有类似 CPU 的时间片模式和 MPS 模式。

从用户的使用角度来看，GPU 共享这种虚拟化技术主要体现在虚拟机层面和容器层面：

  • 虚拟机层面：将 GPU 硬件设备分割成很多虚拟 GPU 并映射到虚机里面，如 NVIDIA vGPU。
  • 容器层面：容器的本质还是进程，通过对驱动的某些关键接口进行封装劫持从而达到限制进程资源的目的，如 Gaia GPU 和 cGPU 等。

下面是对业内一些方案的介绍。

NVIDIA MPS

NVIDIA MPS (Multi-Process Service) 是 NVIDIA 公司推出的一套 GPU 共享方案。通过 NVIDIA 官方给出的文档我们可以看出，它是一种 C/S 架构模式，通过让多个 CUDA 程序共享同一个 GPU context，从而实现多个 CUDA 程序共享 GPU 的目的。用户可以通过设置 CUDA_MPS_ACTIVE_THREAD_PERCENTAGE 可用线程比例环境变量从而调整算力，但坏处是错误会互相影响。

NVIDIA MPS 系统架构

NVIDIA vGPU

NVIDIA vGPU 是 NVIDIA 的一款商业虚拟化产品，主要针对的是虚拟机平台，通过 vfio 提供了一个隔离环境，允许多个虚拟机能够同时访问宿主机的单个 NVIDIA GPU，每个 vGPU 类似于物理 GPU，有固定的显存大小。缺点是这种方案无法动态调整资源，并且需要购买 NVIDIA 的 license，比较昂贵。

NVIDIA vGPU 系统架构

GaiaGPU

GaiaGPU 是腾讯发表在 ISPA'18 上的一种 GPU 共享技术，同时代码也在 Github 开源，它是面向 K8S 中容器的一套 GPU 虚拟化方案。在工程设计上，该方案包括三个部分，CUDA 封装库 vCUDA、K8S Device Plugin 插件 gpu-manager daemonset 和 K8S 调度插件 gpu-admission。其中，vCUDA 通过劫持容器内用户程序对 CUDA Driver API 的调用来限制当前容器内进程对 GPU 和显存的使用。

vCUDA 调用示例

下表是其论文中描述的劫持的 CUDA Driver API。

CUDA Driver API 劫持列表

cGPU

今年 9 月阿里云在其商业化容器服务中推出了 cGPU 共享产品。和之前其开源的 GPU Sharing 工具不同，这次 cGPU 实现了对资源的隔离。从官方使用文档来看，它是通过一个内核驱动，劫持了对 driver 的调用来完成资源隔离，任务执行服从时间片模式。该方案目前没有开源和相关论文，只能在阿里云上使用。同样，在 10 月百度智能云也相继推出了 cGPU 共享技术，不过没有披露更多的技术细节。

cGPU容器架构图

rCUDA

这类方法和 GaiaGPU 类似，都是封装劫持 CUDA API 的调用，区别之处是这种方案可以实现 GPU 池化，即通过网络远程访问 GPU 资源。缺点是需要用户对静态链接的程序重新编译，CUDA 升级的时候需要进行修改来适配。类似思路的实现有驱动科技的商业化 OrionX 计算平台（猎户座计算平台），之前其有开源的 Orion vGPU，不过目前已不再维护。

rCUDA 架构

关于更多的研究进展，可以参考文末给出的参考文献，其中《针对深度学习的 GPU 共享》 一文对常见的 GPU 共享技术进行了综述，介绍了近年及 OSDI '20 上一些相关技术的最新进展。

下面，将对开源的 GaiaGPU 技术原理进行介绍。

GaiaGPU 原理

前置知识

为了更好地理解 GaiaGPU 的运行原理，这里先对一些相关的前置知识进行简单介绍，文中涉及 K8S 的部分基于 Kubernetes release-1.15。

NVIDIA Docker

之所以先介绍 NVIDIA Docker，是因为我们在容器环境下要想创建 GPU 应用程序必须要先安装 NVIDIA Docker。那么 Docker 和 NVIDIA Docker 又是什么关系呢？先看下面这幅图：

image.png

上图描述了创建一个 Docker 容器的基本流程，其中 runc 是一种符合 OCI 规范的容器运行时具体实现，是 Docker 的默认容器运行时。 NVIDIA Docker 实际上是通过在 runc 的基础上增加一个 prestart hook 来调用 libnvidia-container 挂载 GPU Device 和 CUDA Driver，从而让 NVIDIA GPU 能被容器识别并使用，以上过程可以用 NVIDIA 博客的一幅图来概括。

NVIDIA Docker 运行流程

这样，我们就可以用类似下面的命令来创建 GPU 容器了，其中 hook 会根据环境变量 NVIDIA_VISIBLE_DEVICES 来判断是否会分配 GPU 设备，以及挂载的设备 ID。

1
docker run -it -e NVIDIA_VISIBLE_DEVICES=xxx  ubuntu:18.04  /bin/bash

明白了 Docker 使用 NVIDIA GPU 的流程之后，再来看 K8S 创建 GPU 容器其实也就很简单了，只不过稍微繁琐了一些多了 dockershim，这也是为什么 Kubernetes 1.20 ChangeLog 里声称要“抛弃” Docker 吧。

K8S 创建容器过程

Kubernetes Scheduler

Kubernetes Scheduler 是 Kubernetes 的一个系统组件，负责监听 Kube-apiserver 中 PodSpec.NodeName 为空的 pod，为其选择合适的节点，然后由 Kubelet 在相应节点创建 pod。

目前 Kubernetes 调度器可以分为以下两类：

  • 基于谓词 (Predicates) 和优先级 (Priorities) 的调度器
  • 基于调度框架 (Scheduling Framework) 的调度器

由于基于 Scheduling Framework 的调度器对版本要求较高，在 1.19 达到稳定版（详见 Scheduling Framework 提案），所以这里暂不讨论。

基于谓词和优先级的调度器分为如下两个步骤，大体流程如下图所示：

  1. Predicates：过滤阶段，调度器会利用存储匹配相关、Pod 和 Node 匹配相关以及 Pod 和 Pod 匹配相关等来对节点进行过滤，如 NoDiskConfict、PodFitsHostPorts 和 MatchinterPodAffinity 算法
  2. Priorities：打分排序阶段，利用 Node 水位、Node/Pod 亲和性等优先级算法对过滤出来的节点计算得分，分数最高的节点即为要绑定的节点，如 LeastRequestedPriority、SelectorSpreadPriority 和 NodeAffinityPriority 算法
1
finalScoreNodeA = (weight1 * priorityFunc1) + (weight2 * priorityFunc2)

Kubernetes 调度流程

Scheduler Extender

Kubernetes 社区列举了 3 种扩展调度器的方法：

  • 在 Scheduler 源码中加入自己想要的规则然后重新编译
  • 实现自定义的 Scheduler 程序，其可以替代或与 K8S 默认调度器并行，实际使用时要在 Pod 的 spec.schedulerName 指定调度器名称
  • 调度器扩展程序 (Scheduler Extender)，实际上是一个可配置的 Webhook，包含了前面所述的 Predicates 和 Priorities 两个端点，负责对节点进行过滤和打分

Scheduler Extender 实际使用时需要在 scheduler policy 配置（pkg/scheduler/api/v1/types.go:183）文件中指定访问方式，举例如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
{
    "kind":"Policy",
    "apiVersion":"v1",
    "predicates":[
        {
            "name":"CheckNodeUnschedulable"
        },
        {
            "name":"GeneralPredicates"
        },
        {
            "name":"HostName"
        }
        ...
    ],
    "priorities":[
        {
            "name":"EqualPriority",
            "weight":1
        },
        {
            "name":"MostRequestedPriority",
            "weight":1
        },
        {
            "name":"RequestedToCapacityRatioPriority",
            "weight":1
        }
        ...
    ],
    "extenders":[
        {
            "urlPrefix":"http://localhost:3456/scheduler",
            "filterVerb":"predicates",
            "enableHttps":false,
            "nodeCacheCapable":false
        }
    ]
}

这样，当 pod 被调度时，Scheduler Extender 会接收到调度器如下结构的参数（pkg/scheduler/api/v1/types.go:239），将节点根据自定义的过滤或打分排序逻辑处理后返回给调度器。

1
2
3
4
5
6
7
8
9
10
11
12
// ExtenderArgs represents the arguments needed by the extender to filter/prioritize
// nodes for a pod.
type ExtenderArgs struct {
    // Pod being scheduled
    Pod *apiv1.Pod `json:"pod"`
    // List of candidate nodes where the pod can be scheduled; to be populated
    // only if ExtenderConfig.NodeCacheCapable == false
    Nodes *apiv1.NodeList `json:"nodes,omitempty"`
    // List of candidate node names where the pod can be scheduled; to be
    // populated only if ExtenderConfig.NodeCacheCapable == true
    NodeNames *[]string `json:"nodenames,omitempty"`
}

Device Plugin

Device Plugin 是 Kubernetes 提供的一个设备插件框架，用来支持 GPU、FPGA 和高性能 NIC 等第三方设备，只要根据 Device Plugin 的接口实现一个特定设备的插件，就能实现 K8S 对设备的管理。实际上，Device Plugin 是作为 Server 端与 Kubelet (Client 端) 通过节点路径 /var/lib/kubelet/device-plugins/ 下的 Unix 本地套接字以 gRPC 的方式进行通信，比如 NVIDIA GPU Device Plugin 为 nvidia.sock，接口定义（pkg/kubelet/apis/deviceplugin/v1beta1/api.pb.go:567）如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
type DevicePluginServer interface {
    // GetDevicePluginOptions returns options to be communicated with Device
    // Manager
    GetDevicePluginOptions(context.Context, *Empty) (*DevicePluginOptions, error)
    // ListAndWatch returns a stream of List of Devices
    // Whenever a Device state change or a Device disappears, ListAndWatch
    // returns the new list
    ListAndWatch(*Empty, DevicePlugin_ListAndWatchServer) error
    // Allocate is called during container creation so that the Device
    // Plugin can run device specific operations and instruct Kubelet
    // of the steps to make the Device available in the container
    Allocate(context.Context, *AllocateRequest) (*AllocateResponse, error)
    // PreStartContainer is called, if indicated by Device Plugin during registeration phase,
    // before each container start. Device plugin can run device specific operations
    // such as reseting the device before making devices available to the container
    PreStartContainer(context.Context, *PreStartContainerRequest) (*PreStartContainerResponse, error)
}

其中， ListAndWatch 和 Allocate 是最主要的方法：

  • ListAndWatch：当 DevicePlugin 启动并向 Kubelet 注册后，Kubelet 会调用该 API 获取设备信息。值得注意的是，它是一个长连接，当设备健康状况发生变化时，DevicePlugin 会主动向 Kubelet 发送最新的设备信息
  • Allocate：当 Kubelet 要创建使用该设备的容器时， Kubelet 会调用该 API 为容器分配设备资源，得到使用该设备的必要信息，如设备列表、环境变量和挂载点等

Device Plugin 与 Kubelet 交互过程

GaiaGPU

了解了上述相关内容后，再来看 GaiaGPU 的实现会比较清晰。GaiaGPU 分为三个模块：

  • gpu-admission：Scheduler Extender，在宏观上实现对 GPU 资源的压缩调度，即优先占满一台节点的某张 GPU，减少碎片化
  • gpu-manager：Device Plugin，当 pod 绑定到节点后，该节点的 gpu-manager 会根据拓扑感知等算法选出最合适的 GPU 卡，提高通信性能的同时减少碎片化
  • vcuda：CUDA Wrapper，通过将部分 CUDA Driver API 封装拦截，构建为动态链接库挂载进容器，从而实现对 GPU 核和显存的限制

这里主要介绍一下 gpu-admission 和 gpu-manager 的工作原理。

gpu-admission

gpu-admission

gpu-admission 依照 Scheduler Extender 规范，实现了 Predicates 和 Prioritizes 过滤和打分两部分逻辑。其中，Predicates 部分定义如下：

1
2
3
4
5
6
7
type Predicate interface {
    // Name returns the name of this predictor
    Name() string
    // Filter returns the filter result of predictor, this will tell the suitable nodes to running
    // pod
    Filter(args schedulerapi.ExtenderArgs) *schedulerapi.ExtenderFilterResult
}

可以看到，Predicate 接口将 Scheduler 传过来的节点列表进行过滤，具体的过滤算子包括 quotaFilter 和 deviceFilter，前者会根据节点对指定 namespace 下的 GPU 限额进行过滤，后者则会根据 pod 的 GPU 申请量和 node 资源量匹配进行过滤。其中，deviceFilter 是我们要关注的主要部分。

首先，deviceFilter 会过滤出 GPU 节点，然后根据可分配核数、可分配显存和 GPU ID 来对节点进行从小到大排序。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
for i := range nodes {
    node := &nodes[i]
        //筛选出 GPU 节点
    if !util.IsGPUEnabledNode(node) {
        failedNodesMap[node.Name] = "no GPU device"
        continue
    }
    pods, err := gpuFilter.ListPodsOnNode(node)
    if err != nil {
        failedNodesMap[node.Name] = "failed to get pods on node"
        continue
    }
    nodeInfo := device.NewNodeInfo(node, pods)
    nodeInfoList = append(nodeInfoList, nodeInfo)
}
//根据可分配核数、可分配显存和 GPU ID 来对节点进行从小到大排序
sorter.Sort(nodeInfoList)

接着，在排序后的节点列表中为 pod 筛选节点，注意这里当找到符合要求的第一个节点时，后面的节点都会被跳过。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
for _, nodeInfo := range nodeInfoList {
    node := nodeInfo.GetNode()
    //如果找到一个节点满足条件且被成功打标记，则跳过后面节点
    if success {
        failedNodesMap[node.Name] = fmt.Sprintf("pod %s has already been matched to another node", pod.UID)
        continue
    }

    //通过传入一个 GPU 节点详细信息，创建一个Allocator，判断是否满足 pod GPU 申请要求
    alloc := algorithm.NewAllocator(nodeInfo)
    newPod, err := alloc.Allocate(pod)
   
    ...

    success = true  
}

为 pod 分配 GPU 时实际上是为 pod 下面每个需要 GPU 的容器进行分配，这里会根据容器申请的 GPU 数来判断是否可以共享 GPU，在选 GPU 卡时会像筛选节点一样，优先填满一张卡，具体逻辑如下，

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
func (alloc *allocator) AllocateOne(container *v1.Container) ([]*device.DeviceInfo, error) {
   
    ...

    node := alloc.nodeInfo.GetNode()
    //节点总的显存块数
    nodeTotalMemory := util.GetCapacityOfNode(node, util.VMemoryAnnotation)
    //GPU 卡数 = 节点 gpu 份数 / 100
    deviceCount := util.GetGPUDeviceCountOfNode(node)
    //每张 gpu 卡的显存量
    deviceTotalMemory := uint(nodeTotalMemory / deviceCount)
    //容器所需的 gpu 份数
    needCores := util.GetGPUResourceOfContainer(container, util.VCoreAnnotation)
    //容器所需显存块数
    needMemory := util.GetGPUResourceOfContainer(container, util.VMemoryAnnotation)
    switch {
    case needCores < util.HundredCore:
        //申请核数小于 100 为共享模式
        devs = NewShareMode(alloc.nodeInfo).Evaluate(needCores, needMemory)
        sharedMode = true
    default:
        devs = NewExclusiveMode(alloc.nodeInfo).Evaluate(needCores, needMemory)
    }
    if len(devs) == 0 {
        return nil, fmt.Errorf("failed to allocate for container %s", container.Name)
    }

    ...

    return devs, nil
}

共享模式下筛选 GPU 卡逻辑如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
func (al *shareMode) Evaluate(cores uint, memory uint) []*device.DeviceInfo {
    var (
        devs        []*device.DeviceInfo
        deviceCount = al.node.GetDeviceCount()
        tmpStore    = make([]*device.DeviceInfo, deviceCount)
        sorter      = shareModeSort(device.ByAllocatableCores, device.ByAllocatableMemory, device.ByID)
    )
    for i := 0; i < deviceCount; i++ {
        tmpStore[i] = al.node.GetDeviceMap()[i]
    }
    //根据可分配的 GPU 核数、显存数对 GPU 进行从小到大排序
    sorter.Sort(tmpStore)
    for _, dev := range tmpStore {
        //获取 GPU 核数和显存都大于给定值的 1 张 GPU 卡
        if dev.AllocatableCores() >= cores && dev.AllocatableMemory() >= memory {
            klog.V(4).Infof("Pick up %d , cores: %d, memory: %d", dev.GetID(), dev.AllocatableCores(), dev.AllocatableMemory())
            devs = append(devs, dev)
            break
        }
    }
    return devs
}

当分配 GPU 完成后，会对 pod 打一个 Annotation patch，标识这个 pod 已经被调度和具体的分配的信息，创建后的 pod yaml 如下所示。值得注意的是，patch 操作相当于使得 pod 发生了改变，会导致一次调度，不过在调度时判断一下就好了。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
apiVersion: v1
kind: Pod
metadata:
  annotations:
    # annotation patch
    tencent.com/gpu-assigned: "true"
    tencent.com/predicate-gpu-idx-0: "0"
  creationTimestamp: "2020-12-07T11:51:07Z"
  name: one-tencent-gpu
  namespace: danlu-efficiency
  ...
spec:
  containers:
  - image: hub.fuxi.netease.com/danlu-platform/dev/danlu-backend:gpu-test-app-time-cost
    name: one-tencent-gpu
    resources:
      limits:
        tencent.com/vcuda-core: "25"
        tencent.com/vcuda-memory: "10"
      requests:
        tencent.com/vcuda-core: "25"
        tencent.com/vcuda-memory: "10"
 ...

gpu-manager

这里仍然先用脑图展示一下 gpu-manager 包含的主要部分：

gpu-manager

  • VolumeManager：负责查找并复制与 NVIDIA GPU 相关可执行文件和动态链接库到 /etc/gpu-manager/vdriver 下，该路径下保存了一份原始库 (origin) 和一份经过 vCUDA 封装拦截的库 (nvidia)。在共享模式下，为容器分配 GPU 时会将被替换的库作为挂载点返回给 Kubelet，如下所示。
1
2
3
4
5
6
7
8
9
10
11
12
/etc/gpu-manager/
├── log
├── vdriver
│   ├── nvidia
│   │   ├── bin
│   │   ├── lib
│   │   └── lib64
│   └── origin
│       ├── bin
│       ├── lib
│       └── lib64
└── vm
  • ContainerRuntimeManager：根据 CgroupDriver （默认 cgroupfs）和容器运行时 Unix 本地套接字路径（默认 /var/run/dockershim.sock）建立 gRPC 连接，方便后续通过 CRI 获取容器内运行进程的相关信息等。
  • PodCache：watchdog 程序，通过 PodInformer 不断获取所在节点的 pod
  • VirtualManager：负责 GPU 分配后的管理工作
  • GPUTree：根据拓扑关系为所在节点的 GPU 卡构建一颗 GPU 树，每张 GPU 为树的叶子节点
  • NvidiaTopoAllocator：是 ListAndWatch 和 Allocate 的具体实现，负责为容器分配 GPU
  • Display：获取 Pod 的 GPU 使用信息
  • MetricsService：提供 Prometheus metric 接口

既然是 Device Plugin，我们自然比较关注 ListAndWatch 和 Allocate 方法，下面具体来讲一下 GPUTree 的构建和 NvidiaTopoAllocator 的工作原理。

GPUTree 的构建

gpu-manager 默认通过 nvml 库来构建 GPUTree，初始化代码如下，主要获取每张 GPU 的 minorID(可以当作卡的索引)、UUID、totalMem 等信息：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
for i := 0; i < int(num); i++ {
    dev, _ := nvml.DeviceGetHandleByIndex(uint(i))
    _, _, totalMem, _ := dev.DeviceGetMemoryInfo()
    pciInfo, _ := dev.DeviceGetPciInfo()
    minorID, _ := dev.DeviceGetMinorNumber()
    uuid, _ := dev.DeviceGetUUID()
        //叶节点掩码 Mask 为 1 << uint(index)
    n := t.allocateNode(i)
    //HundredCore 为常量 100，后面 NvidiaTopoAllocator 解释
    n.AllocatableMeta.Cores = HundredCore
    n.AllocatableMeta.Memory = int64(totalMem)
    n.Meta.TotalMemory = totalMem
    n.Meta.BusId = pciInfo.BusID
    n.Meta.MinorID = int(minorID)
    n.Meta.UUID = uuid

    t.addNode(n)
}

在得到了每张 GPU 的信息后，然后是获取 GPU 之间的拓扑关系并构造 GPUTree 的非叶子节点。在命令行下我们可以通过 nvidia-smi topo -i 得到 GPU 拓扑，如下表所示，其中，PIX、PHB 和 SYS 表示了不同的拓扑关系，通信消耗逐渐增大。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
        GPU0    GPU1    GPU2    GPU3    GPU4    GPU5    GPU6    GPU7    CPU Affinity
GPU0     X      PIX     PHB     PHB     SYS     SYS     SYS     SYS     0-13,28-41
GPU1    PIX      X      PHB     PHB     SYS     SYS     SYS     SYS     0-13,28-41
GPU2    PHB     PHB      X      PIX     SYS     SYS     SYS     SYS     0-13,28-41
GPU3    PHB     PHB     PIX      X      SYS     SYS     SYS     SYS     0-13,28-41
GPU4    SYS     SYS     SYS     SYS      X      PIX     PHB     PHB     14-27,42-55
GPU5    SYS     SYS     SYS     SYS     PIX      X      PHB     PHB     14-27,42-55
GPU6    SYS     SYS     SYS     SYS     PHB     PHB      X      PIX     14-27,42-55
GPU7    SYS     SYS     SYS     SYS     PHB     PHB     PIX      X      14-27,42-55

Legend:

  X    = Self
  SYS  = Connection traversing PCIe as well as the SMP interconnect between NUMA nodes (e.g., QPI/UPI)
  NODE = Connection traversing PCIe as well as the interconnect between PCIe Host Bridges within a NUMA node
  PHB  = Connection traversing PCIe as well as a PCIe Host Bridge (typically the CPU)
  PXB  = Connection traversing multiple PCIe switches (without traversing the PCIe Host Bridge)
  PIX  = Connection traversing a single PCIe switch
  NV#  = Connection traversing a bonded set of # NVLinks

类似的，在程序中可以由 nvml.DeviceGetTopologyCommonAncestor(devA, devB) 得到两张卡之间的拓扑类型，进而构造 GPUTree 的非叶子节点。

1
2
3
4
5
6
7
8
9
10
11
12
for cardA := uint(0); cardA < num; cardA++ {
    devA, _ := nvml.DeviceGetHandleByIndex(cardA)
    for cardB := cardA + 1; cardB < num; cardB++ {
        devB, _ := nvml.DeviceGetHandleByIndex(cardB)
        ntype, err := nvml.DeviceGetTopologyCommonAncestor(devA, devB)
        ...
        if newNode := t.join(nodes, ntype, int(cardA), int(cardB)); newNode != nil {
            klog.V(2).Infof("New node, type %d, mask %b", int(ntype), newNode.Mask)
            nodes[ntype] = append(nodes[ntype], newNode)
        }
    }
}

叶节点掩码 Mask = 1 << index // index 为卡的索引
非叶节点掩码 Mask = nodeA.Mask | nodeB.Mask

在得到了以上信息后，接下来就可以构建 GPUTree 了。构建树时，实际就是为每个节点寻找自己的父节点和子节点，下面代码展示了自底向上来设置节点的父节点和子节点。值得注意的是，这里使用了位运算的 trick，通过判断节点的掩码与运算是否为 0，来判断是否为父子节点。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
func (t *NvidiaTree) buildTree(nodes LevelMap) {
        //设置父节点
    for _, cur := range t.leaves {
        level := int(nvml.TOPOLOGY_SINGLE)
        self := cur
        for {
            for _, upperNode := range nodes[nvml.GpuTopologyLevel(level)] {
                //如果 mask 与操作后不为 0，则判断为其父节点，设置该节点的父节点字段
                if (upperNode.Mask & self.Mask) != 0 {
                    self.setParent(upperNode)
                    self = upperNode
                    break
                }
            }

            level += levelStep
            if level > int(nvml.TOPOLOGY_SYSTEM) {
                break
            }
        }
    }

    ...

    //自底向上将虚拟子节点设置为子节点
    // Transform vchildren to children
    for _, n := range t.leaves {
        cur := n.Parent
        for cur != nil {
            if len(cur.Children) == 0 {
                cur.Children = make([]*NvidiaNode, 0)

                for _, child := range cur.vchildren {
                    cur.Children = append(cur.Children, child)
                }
            }
            //自底向上设置子节点
            cur = cur.Parent
        }
    }
}

NvidiaTopoAllocator

ListAndWatch 和 Allocate 是为容器分配设备资源的主要方法。在此之前，我们再来回顾一下前面创建的 pod 的 yaml，你可能会困惑，这里 resources 字段表示什么意思，为什么是这种取值。

1
2
3
4
5
6
7
resources:
      limits:
        tencent.com/vcuda-core: "25"
        tencent.com/vcuda-memory: "10"
      requests:
        tencent.com/vcuda-core: "25"
        tencent.com/vcuda-memory: "10"

实际上，GaiaGPU 将 GPU 分成核数 (tencent.com/vcuda-core)和显存 (tencent.com/vcuda-memory)两种设备资源，加之 K8S 会将资源值转为整数，那么这里就用了一个 trick——在 GPU 核方面，将原先的 1 张卡即 nvidia.com/gpu: 1 划分为 100 份 tencent.com/vcuda-core，即 100 表示申请 1 张卡；显存 1 份表示 256Mi。那么，作为 Device Plugin 当被 Kubelet 调用 ListAndWatch 时，就会向 Kubelet 注册设备 (这里以8 张卡举例) 时，就会得到下面的“设备”：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
///var/lib/kubelet/device-plugins/kubelet_internal_checkpoint
{
    "Data":{
        ...
        "RegisteredDevices":{
            "nvidia.com/gpu":[
                "GPU-df1fa0de-0161-9bc1-eac7-825e9846fa53",
                "GPU-12c05a23-d8d7-aca7-d0c5-1df2e8d1f6a9",
                "GPU-cc1f9466-61bd-7c60-1f13-93bdcf5bd9c4",
                "GPU-38896278-98b0-645f-1fb8-b76bf875cce2",
                "GPU-f3d43c0f-d9eb-cce4-b8ec-8001bc0ee412",
                "GPU-03f50d8b-8228-247c-b186-860762e31dd1",
                "GPU-c8c2e7f6-b738-ed64-3c65-6ae3e489e4c3",
                "GPU-5572a570-a180-c4ff-298c-8dd4d4e66ffb"
            ],
            //长度为 800，即 8 张卡
            "tencent.com/vcuda-core":[
                "tencent.com/vcuda-core-381",
                "tencent.com/vcuda-core-389",
                "tencent.com/vcuda-core-463"
                ...
            ],
            //长度为 349，每个表示 256Mi
            "tencent.com/vcuda-memory":[
                "tencent.com/vcuda-memory-268435456-215",
                "tencent.com/vcuda-memory-268435456-263",
                "tencent.com/vcuda-memory-268435456-337"
                ...  
            ]
        }
    },
    "Checksum":1856670514
}

当创建 GPU 容器时我们最关注的就是 Allocate 方法，这里面根据 GPU 核的申请量分为三种情况：

  • shareMode：共享模式，即 GPU 核申请量小于 100，选出卡上剩余资源最少且满足容器需求的 GPU，减少碎片
  • linkMode：多卡模式，即 GPU 核申请量是大于 100 的整数，根据拓扑关系，选出通信性能最高的多张卡
  • fragmentMode：单卡模式，即 GPU 核申请量等于 100

代码的核心逻辑如下，即为每种情况设置了一种评估器，利用评估器在 GPUTree 中选出最合适的叶子节点 即 GPU：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
switch {
    case needCores > nvtree.HundredCore:
        eval, ok := ta.evaluators["link"]
        ...
        nodes = eval.Evaluate(needCores, 0)

    case needCores == nvtree.HundredCore:
        eval, ok := ta.evaluators["fragment"]
        ...
        nodes = eval.Evaluate(needCores, 0)

    default:        
        shareMode = true
        ...
        eval, ok := ta.evaluators["share"]
        nodes = eval.Evaluate(needCores, needMemory)
    ...

}

评估器 的逻辑基本类似，我们这里主要看下共享模式和多卡模式下的选卡实现：

  • 共享模式：直接在叶子节点，即实际 GPU 卡中选出符合要求的 1 张卡
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
// 共享模式
func (al *shareMode) Evaluate(cores int64, memory int64) []*nvidia.NvidiaNode {
    var (
        nodes    []*nvidia.NvidiaNode
        tmpStore = make([]*nvidia.NvidiaNode, al.tree.Total())
        sorter = shareModeSort(nvidia.ByAllocatableCores, nvidia.ByAllocatableMemory, nvidia.ByPids, nvidia.ByID)
    )

    for i := 0; i < al.tree.Total(); i++ {
        tmpStore[i] = al.tree.Leaves()[i]
    }

    //对所有叶子节点排序，直接找出可分配核数和显存数都大于等于请求量的 GPU 卡
    sorter.Sort(tmpStore)

    for _, node := range tmpStore {
        if node.AllocatableMeta.Cores >= cores && node.AllocatableMeta.Memory >= memory {
            nodes = append(nodes, node)
            break
        }
    }

    return nodes
}
  • 多卡模式：会根据拓扑关系、可分配节点数等选出多张卡，这点在排序函数 sorter 中体现，和前面联系起来就是自底向上建树，自顶向下选卡
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
func (al *linkMode) Evaluate(cores int64, memory int64) []*nvidia.NvidiaNode {
    var (
        sorter   = linkSort(nvidia.ByType, nvidia.ByAvailable, nvidia.ByMemory, nvidia.ByPids, nvidia.ByID)
        tmpStore = make(map[int]*nvidia.NvidiaNode)
        root     = al.tree.Root()
        nodes    = make([]*nvidia.NvidiaNode, 0)
        num      = int(cores / nvidia.HundredCore)
    )

    for _, node := range al.tree.Leaves() {
        //自底向上找到可用节点数量满足的节点，直到找到根节点，若还不满足则换下个叶子节点继续向上找
        for node != root {
            //自底向上找到可用节点数量满足的节点
            if node.Available() < num {
                node = node.Parent
                continue
            }
            //自底向上找到可用卡数满足的节点，记录进 tmpStore
            tmpStore[node.Meta.ID] = node
            break
        }
    }

    //将上面找到的节点放进候选列表
    candidates := make([]*nvidia.NvidiaNode, 0)
    for _, n := range tmpStore {
        candidates = append(candidates, n)
    }

    //对候选列表排序
    sorter.Sort(candidates)

    for _, n := range candidates[0].GetAvailableLeaves() {
        if num == 0 {
            break
        }
        nodes = append(nodes, n)
        num--
    }

    if num > 0 {
        return nil
    }

    return nodes
}

多卡模式

当分配完后，余下的工作就是将设备使用的相关信息返回给 Kubelet，从返回值的定义（pkg/kubelet/apis/deviceplugin/v1beta1/api.pb.go:273）我们可以看出 Kubelet 要启动容器会用到哪些参数：

1
2
3
4
5
6
7
8
9
10
type ContainerAllocateResponse struct {
    // List of environment variable to be set in the container to access one of more devices.
    Envs map[string]string `protobuf:"bytes,1,rep,name=envs" json:"envs,omitempty" protobuf_key:"bytes,1,opt,name=key,proto3" protobuf_val:"bytes,2,opt,name=value,proto3"`
    // Mounts for the container.
    Mounts []*Mount `protobuf:"bytes,2,rep,name=mounts" json:"mounts,omitempty"`
    // Devices for the container.
    Devices []*DeviceSpec `protobuf:"bytes,3,rep,name=devices" json:"devices,omitempty"`
    // Container annotations to pass to the container runtime
    Annotations map[string]string `protobuf:"bytes,4,rep,name=annotations" json:"annotations,omitempty" protobuf_key:"bytes,1,opt,name=key,proto3" protobuf_val:"bytes,2,opt,name=value,proto3"`
}

对应在 gpu-manager 中的主要工作如下所示：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
for _, n := range nodes {
    name := n.MinorName()
    ctntResp.Devices = append(ctntResp.Devices, &pluginapi.DeviceSpec{
        ContainerPath: name,
        HostPath:      name,
        Permissions:   "rwm",
    })
    //GPU 卡的 UUID，容器内 NVIDIA_VISIBLE_DEVICES 环境变量要用到
    deviceList = append(deviceList, n.Meta.UUID)
}

//更改相应的 Annotations tencent.com/vcuda-device:/dev/nvidia0,/dev/nvidia1...
ctntResp.Annotations[types.VDeviceAnnotation] = vDeviceAnnotationStr(nodes)
//在响应中为容器添加 /dev/nvidiactl 和 /dev/nvidia-uvm 挂载信息
// Append control device
ctntResp.Devices = append(ctntResp.Devices, &pluginapi.DeviceSpec{
    ContainerPath: types.NvidiaCtlDevice,
    HostPath:      types.NvidiaCtlDevice,
    Permissions:   "rwm",
})

ctntResp.Devices = append(ctntResp.Devices, &pluginapi.DeviceSpec{
    ContainerPath: types.NvidiaUVMDevice,
    HostPath:      types.NvidiaUVMDevice,
    Permissions:   "rwm",
})

// LD_LIBRARY_PATH
ctntResp.Envs["LD_LIBRARY_PATH"] = "/usr/local/nvidia/lib64"
for _, env := range container.Env {
    if env.Name == "compat32" && strings.ToLower(env.Value) == "true" {
        ctntResp.Envs["LD_LIBRARY_PATH"] = "/usr/local/nvidia/lib"
    }
}

// NVIDIA_VISIBLE_DEVICES
ctntResp.Envs["NVIDIA_VISIBLE_DEVICES"] = strings.Join(deviceList, ",")

//根据是否处于 shareMode，也就是单个 gpu 能否被共享来挂载不同的 host 目录
if shareMode {
    ctntResp.Mounts = append(ctntResp.Mounts, &pluginapi.Mount{
        ContainerPath: "/usr/local/nvidia",
        HostPath:      types.DriverLibraryPath,
        ReadOnly:      true,
    })
} else {
    ctntResp.Mounts = append(ctntResp.Mounts, &pluginapi.Mount{
        ContainerPath: "/usr/local/nvidia",
        HostPath:      types.DriverOriginLibraryPath,
        ReadOnly:      true,
    })
}

//将 host 上的 /etc/gpu-manager/vm/{podUID} 挂载到容器中，这个是为了容器内可以通过 vcuda.sock 和 virtual-manager 通信
ctntResp.Mounts = append(ctntResp.Mounts, &pluginapi.Mount{
    ContainerPath: types.VCUDA_MOUNTPOINT,
    HostPath:      filepath.Join(ta.config.VirtualManagerPath, string(pod.UID)),
    ReadOnly:      true,
})

伏羲丹炉 GPU 共享迁移改造实践

当理解了 GaiaGPU 的工作原理后，我们便要思考如何将其应用于我们的“丹炉”平台中来提高 GPU 资源利用率。不过摆在我们面前的问题是，目前节点上的老应用使用的都是 NVIDIA GPU，那么如何才能让节点上 NVIDIA GPU Pod 和 GaiaGPU Pod 共存，最终平滑迁移到 GaiaGPU 呢？仔细分析后，其实这里主要牵涉三个部分：

  • 调度器

在调度 GaiaGPU Pod 筛选节点和 GPU 时，需要将已经被 NVIDIA GPU Pod 占用的 GPU 排除，即避免出现一张卡上既有使用 NVIDIA GPU 也有使用 GaiaGPU 的容器；对于 NVIDIA GPU Pod，调度阶段无需改造，因为这里只会筛选节点，对于 GPU 的分配是在 NVIDIA Device Plugin 实现

  • Device Plugin

在调度到 GPU 节点后，接下来对 GPU 的具体分配工作是由 NVIDIA Device Plugin 和 gpu-manager 实现。那么为了达到两种模式的互相感知和隔离，在分配时 NVIDIA GPU Pod 就不能再占用被 GaiaGPU Pod 占用的卡，GaiaGPU Pod 也不能再分配已被 NVIDIA GPU Pod 占用的卡。但具体分配卡时，是 kubelet 分配的，所以这里需要对 NVIDIA Device Plugin 和 gpu-manager 进行改造

  • 监控

监控面临的问题是一样的，这里因为引入了 GaiaGPU，同一张卡上会有多个 Pod，因此这里需要将监控粒度缩小到容器，需要展示某张卡是被哪些 Pod 的哪些容器占用

以上操作确保了 NVIDIA GPU 和 GaiaGPU 两种模式的互相感知和隔离， 不会因为混用而对应用造成影响，最终实现平滑迁移。

Admission Webhook

Admission Webhook (动态准入控制器) 这里的工作主要是根据配置的项目 id，将其下的 NVIDIA GPU Pod 进行拦截替换为 GaiaGPU，将老的实例逐步平滑替换为 GaiaGPU，实现方式是通过对 nvidia.com/gpu 、schedulerName 、nodeSelector 和 tolerations 进行 JSON patch 来删除或替换，资源替换如下所示，其它同理：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
func replaceGPUResourcesAndSchedulerPatch(pod *corev1.Pod) (patch []byte, err error) {
    var patchOps []patchOperation
    var realVal float64
    //gpu patch
    for i, con := range pod.Spec.Containers {
        if val, ok := con.Resources.Limits[NvidiaGPUAnnotation]; ok {
            realVal, err = transformatQuantity(val)
            if err != nil {
                return nil, err
            }
            //先删除再添加
            patchOps = append(patchOps, patchOperation{
                Op:   "remove",
             // /spec/containers/%d/resources/limits/nvidia.com~1gpu
                Path: fmt.Sprintf(removeNvidiaGPULimitsJsonPatchPath, i),
            }, patchOperation{
                Op:    "add",
             // `/spec/containers/%d/resources/limits/tencent.com~1vcuda-core`
                Path:  fmt.Sprintf(addTencentGPUCoreLimitsJsonPatchPath, i),
                Value: realVal * 100,
            }, patchOperation{
                Op:    "add",
             //  `/spec/containers/%d/resources/limits/tencent.com~1vcuda-memory`
                Path:  fmt.Sprintf(addTencentGPUMemLimitsJsonPatchPath, i),
                Value: realVal * 42,
            })
        }
        if val, ok := con.Resources.Requests[NvidiaGPUAnnotation]; ok {
            realVal, err = transformatQuantity(val)
            if err != nil {
                return nil, err
            }
            //先删除再添加
            patchOps = append(patchOps, patchOperation{
                Op:   "remove",
              //  `/spec/containers/%d/resources/requests/nvidia.com~1gpu`
                Path: fmt.Sprintf(removeNvidiaGPURequestsJsonPatchPath, i),
            }, patchOperation{
                Op:    "add",
              //  `/spec/containers/%d/resources/requests/tencent.com~1vcuda-core`
                Path:  fmt.Sprintf(addTencentGPUCoreRequestsJsonPatchPath, i),
                Value: realVal * 100,
            }, patchOperation{
                Op:    "add",
             // `/spec/containers/%d/resources/requests/tencent.com~1vcuda-memory`
                Path:  fmt.Sprintf(addTencentGPUMemRequestsJsonPatchPath, i),
                Value: realVal * 42,
            })
        }
    }
    //scheduler patch
    patchOps = append(patchOps, patchOperation{
        Op:    "replace",
        // `/spec/schedulerName`
        Path:  addSchedulerNameJsonPatchPath,
        Value: schedulerName,
    })
    patch, err = json.Marshal(patchOps)
    klog.Infof("JSONPatch: %s", string(patch))
    return
}

gpu-admission 改造

gpu-admission 的默认可分配 GPU 视图是各节点上的所有 GPU，那么当节点上同时存在 NVIDIA GPU Pod 和 GaiaGPU Pod 时，我们理所应当要将已经被占用的 GPU 过滤掉，即只在那些空闲的 GPU 上进行分配。那么在调度 GPU Pod 时需要做一步处理——统计节点上 NVIDIA GPU Pod 和 GaiaGPU Pod 的使用量，为 GPU Pod 筛选有空闲 GPU 的节点，同时为 GaiaGPU Pod 分配空闲 GPU。

首先，在 deviceFilter 中先筛选出有空闲 GPU 的节点：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
nodeDeviceCapacity := util.GetNvidiaGPUDeviceCountOfNode(node)
var nodeDeviceAllocated int
nodeDeviceAllocated += util.GetTencentGpuCountOfNode(pods)
nodeDeviceAllocated += util.GetNvidiaGpuCountOfNode(pods)
if util.MoreOneGpu(pod) {
    nodeDeviceAllocated += util.GetGpuCountOfPod(pod)
    klog.Infof("Node:%s, NodeGpuCapacity:%v,  NodeGpuAllocated+PodGpuRequest:%v", node.Name, nodeDeviceCapacity, nodeDeviceAllocated)
} else {
    klog.Infof("Node:%s, NodeGpuCapacity:%v,  NodeGpuAllocated:%v", node.Name, nodeDeviceCapacity, nodeDeviceAllocated)
}
if (nodeDeviceCapacity - nodeDeviceAllocated) < 0 {
    failedNodesMap[node.Name] = "no GPU device available"
    continue
}
nodeInfo := device.NewNodeInfo(node, pods)
nodeInfoList = append(nodeInfoList, nodeInfo)

然后在构造每个节点的具体信息时，要同时考虑 NVIDIA GPU 和 GaiaGPU 的使用量。需要注意的是，在前面讲解 gpu-admission 创建 GaiaGPU Pod 时，最后会打一个如下的 Annotation patch ，标识了具体的分配的信息，再结合 Pod 信息计算节点上 GaiaGPU 使用量就比较容易：

1
2
3
# annotation patch
tencent.com/gpu-assigned: "true"
tencent.com/predicate-gpu-idx-0: "0"

而由于 NVIDIA GPU Pod 本身并没有标识分配了哪张 GPU，并且可能存在一些 Pod 申请了 NVIDIA GPU 但处于 ErrImagePull 或 ImagePullBackOff 等异常状态，因此这里在获取 NVIDIA GPU 使用量时我们有两种途径：

  • 进入 Pod 容器内根据 /dev/nvidia[0-9] 来获取容器具体占用的哪张卡
  • 上面方法执行失败的话，进入同节点 nvidia-device-plugin 容器内，通过 /var/lib/kubelet/device-plugins/ 下的 kubelet_internal_checkpoint （kubelet 分配设备时会写入该文件）和 nvidia_checkpoint （nvidia-device-plugin 改造后会写该文件，后文会讲）获取：
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
for i, c := range pod.Spec.Containers {
    gpu := util.GetGPUResourceOfContainer(&c, "nvidia.com/gpu")
    if gpu > 0 {
        var indexs []int
        if devices, ok := Cache.Get(pod.Namespace, pod.Name, c.Name); ok {
            indexs = devices
        } else {
            //方法一
            stdout, _, err := util.ExecInPod(pod.Namespace, pod.Name, c.Name, "ls /dev/")
            if err != nil || len(stdout) == 0 {
                if indexs, err = util.GetNvidiaPredicateIdxOfContainer(pod, i); err != nil || len(indexs) == 0 {
                    //方法二
                    ret, err := util.GetGpuIdxOfContainerFromCheckpoint(pods, pod, c.Name)
                    if err != nil && len(ret) == 0 {
                        errFound = true
                        unrecognizedGpu += gpu
                    }
                    indexs = ret
                }
            } else {
                devList := strings.Split(stdout, "\n")
                for _, dev := range devList {
                    if match, _ := regexp.MatchString("nvidia[0-9]", dev); match {
                        indexStr := strings.Split(dev, "nvidia")[1]
                        index, err := strconv.Atoi(indexStr)
                        indexs = append(indexs, index)
                    }
                }
            }
                ...
        }
        ...
        continue
    }
    ...
}

在获取了 NVIDIA GPU Pod 和 GaiaGPU Pod 的使用量后，就可以计算该节点的剩余 GPU 量从而对 Pod 进行分配 GPU 了。

nvidia device plugin 改造

当创建 NVIDIA GPU Pod 时，我们需要将 GaiaGPU Pod 占用的 GPU 过滤掉，但由 Allocate 的参数可以看到，GPU 的分配实际是由 kubelet 指定的：

1
2
3
4
5
6
7
8
Allocate(context.Context, *AllocateRequest) (*AllocateResponse, error)

type AllocateRequest struct {
    ContainerRequests []*ContainerAllocateRequest `protobuf:"bytes,1,rep,name=container_requests,json=containerRequests" json:"container_requests,omitempty"`
}
type ContainerAllocateRequest struct {
    DevicesIDs []string `protobuf:"bytes,1,rep,name=devicesIDs" json:"devicesIDs,omitempty"`
}

然而 kubelet 实际并不知道哪些 GPU 被 NVIDIA GPU Pod 占用，哪些被 GaiaGPU 占用，它只会根据自己之前已被分配的 GPU 来得到剩余 “可分配” 的 GPU，将其作为参数传给 nvidia device plugin，但实际上，这部分 GPU 可能已被 GaiaGPU 占用，所以我们不能直接将 kubelet 指定的 GPU 分配给容器。那么当 kubelet 调用 Allocate 分配 GPU 时，我们应该先判断 GPU 是否被占用，若空闲则可正常分配，否则应选择其它剩余空闲 GPU。

首先，先获取待分配 GPU 的 Pod、容器以及其在 Pod 中的索引：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
//首先获取待分配 GPU Pod 的 容器，这里从我们程序里设置的 cache 中获取
candidatePods, err := na.cache.GetCandidatePods()
for _, pod := range candidatePods {
    if found {
        break
    }
    for i, c := range pod.Spec.Containers {
        gpuReq := GetNvidiaGPUResourceOfContainer(&c)
        if gpuReq == 0 || reqCount != gpuReq {
            continue
        }
       
        if _, err := GetNvidiaPredicateIdxOfContainer(pod, i); err == nil {
            continue
        }

        candidatePod = pod
        candidateContainer = &candidatePod.Spec.Containers[i]
        candidateContainerIdx = i
        found = true
        break
    }
}

然后，为该容器分配真实的 GPU。在这里，我们通过构造 devsUsage，记录了节点上 GPU 的 UUID 以及是否被占用，那么在遍历 kubelet 发来的 GPU 列表时，我们就可以判断其是否能被分配，若不能分配则跳过，从剩余空闲 GPU 中分配：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
if found {
    gpuReq := GetNvidiaGPUResourceOfContainer(candidateContainer)
    var devAlloc []string
    for i, devUsage := range devsUsage {
        for _, devId := range req.DevicesIDs {
            if devUsage.ID == devId && !devUsage.InUse {
                devAlloc = append(devAlloc, devId)
                devsUsage[i].InUse = true
                break
            }
        }
    }

    num := int(gpuReq) - len(devAlloc)
    for _, devUsage := range devsUsage {
        if num <= 0 {
            break
        }
        if devUsage.InUse {
            continue
        }
        devAlloc = append(devAlloc, devUsage.ID)
        num--
    }
}

最后，当为这个 Pod 的当前容器分配完 GPU 后，在其 Annotation 打 patch，避免下次重复获取 candidateContainer：

1
2
3
4
5
6
7
annotationMap := make(map[string]string)
devIdx := getDevIdx(devAlloc, devs)
annotationMap[NvidiaGPUPredicateGPUIndexPrefix+strconv.Itoa(candidateContainerIdx)] = strings.Join(devIdx, ",")
if err := na.patchPodWithAnnotations(candidatePod, annotationMap); err != nil {
    return nil, err
}
resps.ContainerResponses = append(resps.ContainerResponses, &response)

当 kubelet 收到上述返回结果后，就会将其写入 /var/lib/kubelet/device-plugins/kubelet_internal_checkpoint 。kubelet_internal_checkpoint 记录了所有注册的设备，并且 Kubelet 为每个容器分配完设备后都会将分配信息写入该文件，其中 AllocResp 是经 base64 编码后的分配信息。值得注意的是，下面为容器 xiaoxi-test-rew 实际分配的 GPU (AllocResp base64解码得到) 就与 kubelet 指定的 DeviceIDs 不同，而 kubelet 实际上也是根据 AllocResp 的值来作为创建容器的参数：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
// /var/lib/kubelet/device-plugins/kubelet_internal_checkpoint
{
    "Data":{
        "PodDeviceEntries":[
            {
                "PodUID":"f7c7d715-39c6-11eb-8b01-6c92bf66acae",
                "ContainerName":"test-app-container",
                "ResourceName":"tencent.com/vcuda-core",
                "DeviceIDs":[
                    "tencent.com/vcuda-core-253",
                    "tencent.com/vcuda-core-637",
                    ...
                ],
             "AllocResp":"CkIKFk5WSURJQV9WSVNJQkxFX0RFVklDRVMSKEdQVS1jZDI3YTI4Yi01ZmM3LTIzMTAtYzEyYi04NWYzNmQ5MWU5NTIKKgoPTERfTElCUkFSWV9QQVRIEhcvdXNyL2xvY2FsL252aWRpYS9saWI2NBI2ChEvdXNyL2xvY2FsL252aWRpYRIfL2V0Yy9ncHUtbWFuYWdlci92ZHJpdmVyL252aWRpYRgBEkgKCi9ldGMvdmN1ZGESOC9ldGMvZ3B1LW1hbmFnZXIvdm0vZjdjN2Q3MTUtMzljNi0xMWViLThiMDEtNmM5MmJmNjZhY2FlGAEaIQoML2Rldi9udmlkaWEwEgwvZGV2L252aWRpYTAaA3J3bRolCg4vZGV2L252aWRpYWN0bBIOL2Rldi9udmlkaWFjdGwaA3J3bRonCg8vZGV2L252aWRpYS11dm0SDy9kZXYvbnZpZGlhLXV2bRoDcndtIigKGHRlbmNlbnQuY29tL3ZjdWRhLWRldmljZRIML2Rldi9udmlkaWEwIiYKGHRlbmNlbnQuY29tL3ZjdWRhLW1lbW9yeRIKODMyMTQ5OTEzNiIcChZ0ZW5jZW50LmNvbS92Y3VkYS1jb3JlEgI3NQ=="
            },
            {
                "PodUID":"f7c7d715-39c6-11eb-8b01-6c92bf66acae",
                "ContainerName":"test-app-container",
                "ResourceName":"tencent.com/vcuda-memory",
                "DeviceIDs":[
                    "tencent.com/vcuda-memory-268435456-123"
                    ...
                ],
                "AllocResp":""
            },
            {
                "PodUID":"fc17fb5e-394e-11eb-8b01-6c92bf66acae",
                "ContainerName":"xiaoxi-test-rew",
                "ResourceName":"nvidia.com/gpu",
                "DeviceIDs":[
                    "GPU-00f2b915-9faa-053e-fa1a-6ff9472e3f39"
                ],                          
                //base64 解码后 GPU-3c4e132d-e8e9-c6b4-26e3-cf57261d2288
               "AllocResp":"CkIKFk5WSURJQV9WSVNJQkxFX0RFVklDRVMSKEdQVS0zYzRlMTMyZC1lOGU5LWM2YjQtMjZlMy1jZjU3MjYxZDIyODg="
            }
        ],
        "RegisteredDevices":{
            "nvidia.com/gpu":[
                "GPU-1dd8b6f1-6240-2577-667c-1c446f6ffc6d"
                ...
            ]
        }
    }
}

同样地，nvidia device plugin 分配完成之后会将本次的真实分配信息写入 /var/lib/kubelet/device-plugins/checkpoint/nvidia_checkpoint 中，两者数据结构保持一致。可以看到， nvidia_checkpoint 其实是 kubelet_internal_checkpoint 的子集，只不过这里为容器分配的 DeviceIDs 对应 kubelet_internal_checkpoint 中的 AllocResp：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
{
    "Data":{
        "PodDeviceEntries":[
            {
                "PodUID":"fc17fb5e-394e-11eb-8b01-6c92bf66acae",
                "ContainerName":"xiaoxi-test-rew",
                "DeviceIDs":[
                    "GPU-3c4e132d-e8e9-c6b4-26e3-cf57261d2288"
                ],
                "ResourceName":"nvidia.com/gpu"
            }
        ],
        "RegisteredDevices":null
    }
}

gpu-manager 改造

gpu-manager 同 nvidia device plugin 一样，在为 GaiaGPU 容器分配 GPU 时，需要将已经被 NVIDIA GPU Pod 占用的 GPU 排除掉，那么这里就需要获取当前哪些卡被 NVIDIA GPU Pod 占用，并将其标记，同时还要将已经不被 NVIDIA GPU Pod 占用的 GPU 给释放掉，即标记和回收两个步骤。

要想获取哪些 GPU 被 NVIDIA GPU Pod 占用，我们可以结合前面提到的 /var/lib/kubelet/device-plugins/kubelet_internal_checkpoint 和 /var/lib/kubelet/device-plugins/checkpoint/nvidia_checkpoint 文件中获取。

kubelet_internal_checkpoint 记录了所有 Kubelet 为容器分配的设备信息，对于 nvidia_checkpoint，正如前面介绍 nvidia device plugin 改造时讲过的，这里保存了创建新的 NVIDIA GPU 容器时分配的真实 GPU UUID，那么结合两个文件便可得到哪些 Pod 的哪些容器占用了哪些 GPU，主要代码逻辑如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
func (ta *NvidiaTopoAllocator) getNvidiaGPUUsedInfoFromCheckpoint() (map[string][]string, []string, error) {
    all, err := utils.GetPodID2GPUIDMapFromCheckpoint(ta.config.DevicePluginPath, types.CheckPointFileName)
    part, err := utils.GetPodID2GPUIDMapFromCheckpoint(ta.config.CheckpointPath, types.NvidiaCheckPointFileName)
   
    l := len(all)
    usedMap := make(map[string][]string, l)
    podIds := make([]string, l)
    for podId := range all {
        if v, ok := part[podId]; ok {
            //新实例按 nvidia_checkpoint
            usedMap[podId] = v
        } else {
            //旧实例按 kubelet_internal_checkpoint
            usedMap[podId] = all[podId]
        }
        podIds = append(podIds, podId)
    }
    return usedMap, podIds, nil
}

但是，kubelet_internal_checkpoint 会存在延迟更新的问题，即当一个 NVIDIA GPU Pod 是 Failed、停止或被删除时，这个文件里依然保存其 GPU 占用信息。

通过分析 Kubelet 的代码 (pkg/kubelet/cm/devicemanager/manager.go:521) 可以看到，writeCheckpoint 负责 kubelet_internal_checkpoint 的写入，代码中 podDevices 是 kubelet 在内存中维护的为 pod 分配设备信息的 cache，对应前面看到的 kubelet_internal_checkpoint 内容，具体实现如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
// Checkpoints device to container allocation information to disk.
func (m *ManagerImpl) writeCheckpoint() error {
    m.mutex.Lock()
    registeredDevs := make(map[string][]string)
    for resource, devices := range m.healthyDevices {
        registeredDevs[resource] = devices.UnsortedList()
    }
    data := checkpoint.New(m.podDevices.toCheckpointData(),
        registeredDevs)
    m.mutex.Unlock()
    err := m.checkpointManager.CreateCheckpoint(kubeletDeviceManagerCheckpoint, data)
    if err != nil {
        return fmt.Errorf("failed to write checkpoint file %q: %v", kubeletDeviceManagerCheckpoint, err)
    }
    return nil
}

writeCheckpoint 只有三处会被调用。其中，只有 allocateContainerResources (pkg/kubelet/cm/devicemanager/manager.go:651) 即为容器分配设备资源时才会将该 cache 刷新落盘，而后面两处调用并未更新 podDevices，因此会导致 kubelet_internal_checkpoint 更新延迟。同理，nvidia_checkpoint 也只在分配 GPU 时更新，不过因其为 kubelet_internal_checkpoint 的子集所以可以不用考虑：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
func (m *ManagerImpl) allocateContainerResources(pod *v1.Pod, container *v1.Container, devicesToReuse map[string]sets.String) error {
    allocatedDevicesUpdated := false
    for k, v := range container.Resources.Limits {
        resource := string(k)
        needed := int(v.Value())
        //更新 podDevices
        if !allocatedDevicesUpdated {
            m.updateAllocatedDevices(m.activePods())
            allocatedDevicesUpdated = true
        }
        allocDevices, err := m.devicesToAllocate(podUID, contName, resource, needed, devicesToReuse[resource])
        if err != nil {
            return err
        }
        if allocDevices == nil || len(allocDevices) <= 0 {
            continue
        }
        m.mutex.Lock()
        //获取 device plugin gRPC 服务
        eI, ok := m.endpoints[resource]
        m.mutex.Unlock()
        devs := allocDevices.UnsortedList()
        //调用 device plugin allocate 方法
        resp, err := eI.e.allocate(devs)
        m.mutex.Lock()
        m.podDevices.insert(podUID, contName, resource, allocDevices, resp.ContainerResponses[0])
        m.mutex.Unlock()
    }

    //写 kubelet_internal_checkpoint 文件
    //由于前面 podDevices 被更新，所以这里文件也会被刷新
    return m.writeCheckpoint()
}

writeCheckpoint

综上所述，我们需要通过 Pod Informer 找出当前存在的 Pod，两者取 差集 将那些不存在的 Pod 给过滤掉，同时将其占有的 GPU 释放，所以总的逻辑如下：每次将 NVIDIA GPU Pod 占用的 GPU 标记为 已占用 即不能分配给 GaiaGPU。另外，每次获取占用信息时都将其保存在程序的 nvidiaGPUUsedCache 中，每次分配时进行对比，将之前存在但此时不存在的 NVIDIA GPU Pod 所占用的 GPU 给释放掉：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
func (ta *NvidiaTopoAllocator) updateNvidiaGPUUsedInfo() {
    //获取当前活跃的 pod
    activeNvidiaPods := watchdog.GetActiveNvidiaPods()
   
    //获取 checkpoint 文件的 GPU 占用信息
    checkpointNvidiaUsedInfo, _, err := ta.getNvidiaGPUUsedInfoFromCheckpoint()
   
    for podUid, gpuUUIDs := range checkpointNvidiaUsedInfo {
        ta.nvidiaGPUUsedCache[podUid] = gpuUUIDs
    }
   
    lastActivePodUids := sets.NewString()
    activePodUids := sets.NewString()
    for podUid := range ta.nvidiaGPUUsedCache {
        lastActivePodUids.Insert(podUid)
    }
    for uid := range activeNvidiaPods {
        activePodUids.Insert(uid)
    }

    //回收
    podsToBeRemoved := lastActivePodUids.Difference(activePodUids).List()
    for _, podUid := range podsToBeRemoved {
        //将checkpoint中非活跃的pod的NVIDIA gpu 回收
        nvidiaGPUUUIDs := ta.nvidiaGPUUsedCache[podUid]
        for _, uuid := range nvidiaGPUUUIDs {
            node := ta.tree.QueryByUUID(uuid)
            ta.tree.MarkFree(node, nvtree.HundredCore, int64(node.Meta.TotalMemory))
        }
        //删除 cache 多余的
        delete(ta.nvidiaGPUUsedCache, podUid)
    }
    ta.markOccupyByNvidiaGPU(ta.nvidiaGPUUsedCache)
}

//标记
func (ta *NvidiaTopoAllocator) markOccupyByNvidiaGPU(nvidiaGPUUsedInfo map[string][]string) {
    for podUid, gpuIds := range nvidiaGPUUsedInfo {
        for _, uuid := range gpuIds {
            node := ta.tree.QueryByUUID(uuid)
            ta.tree.MarkOccupied(node, nvtree.HundredCore, int64(node.Meta.TotalMemory))
        }
    }
}

对于 GPU 的标记和回收涉及对树的更新，这里同样利用位运算，在树上层层传播：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
//标记
func (t *NvidiaTree) occupyNode(n *NvidiaNode) {
    for p := n.Parent; p != nil; p = p.Parent {
        if p.Mask&n.Mask == n.Mask {
            p.Mask ^= n.Mask
        }
    }
}

//回收
func (t *NvidiaTree) freeNode(n *NvidiaNode) {
    for p := n.Parent; p != nil; p = p.Parent {
        p.Mask |= n.Mask
    }
}

通过上述改造，NVIDIA GPU Pod 和 GaiaGPU Pod 就可以互相感知并同时运行在同一节点上了，最终可以确保将老实例逐步平滑迁移为新的方案。

总结

以上是这段时间对 GPU 共享这块工作的总结和思考，由于时间和水平有限，文中可能会存在错误，欢迎大家批评指正～

参考

  • NVIDIA MPS
  • NVIDIA vGPU
  • 在 Kubernetes 中使用 vGPU 实现机器学习任务共享 GPU
  • GaiaGPU: Sharing GPUs in Container Clouds
  • Gaia Scheduler: A Kubernetes-Based Scheduler Framework
  • vcuda-controller
  • gpu-manager
  • gpu-admission
  • GaiaStack上 的 GPU 虚拟化技术
  • GPUManager 虚拟化方案
  • [Gaia Scheduler] gpu-manager 启动流程分析
  • [Gaia Scheduler] gpu-manager 的虚拟化 gpu 分配流程
  • 开源工具 GPU Sharing：支持Kubernetes集群细粒度
  • KubeCon 2019: Minimizing GPU Cost For Your Deep Learning Workload On Kubernetes
  • gpushare-scheduler-extender
  • gpushare-device-plugin
  • 用尽每一寸GPU，阿里云 cGPU 容器技术白皮书重磅发布
  • 百度智能云天合 2.0 更轻量级的 GPU 资源共享技术
  • 针对深度学习的 GPU 共享
  • nvidia-docker
  • Nvidia GPU 如何在 Kubernetes 里工作
  • docker，containerd，runc，docker-shim
  • 深入理解 Nvidia-docker2.0
  • Docker and OCI Runtimes
  • 白话 Kubernetes Runtime
  • kubelet scheduler 源码分析：调度器的工作原理
  • 从零开始入门 K8s：调度器的调度流程和算法介绍
  • 自定义 Kubernetes 调度器
  • Kubernetes的Device Plugin设计解读
  • Dynamic Admission Control
  • 深入理解 Kubernetes Admission Webhook
  • [k8s源码分析][kubelet] devicemanager 之重启 kubelet 和 device-plugin