AI计算教程，从单机到集群

在AI 计算：从单机到集群（上）中我们介绍了任务计算容器化的概念，解决单机计算的便利性问题。当我们进一步深入使用，尤其是有多台计算节点时，会面临一个新问题：多台 GPU 节点如何实现 GPU 资源的合理自动分配，达到共享使用 GPU 集群资源的目的。本篇介绍如何利用 Kubernetes 实现 GPU 资源的集群化管理，我们结合 Kubernetes 相关源码，详细分析 Kubernetes 框架下 GPU 设备管理逻辑和方式，以及如何实现对 GPU 资源的调度，帮助大家迈出从单机到集群的关键一步。

Kubernetes

容器技术简化了计算框架单机运行的问题，但是没法解决多机资源分配问题。虽然目前很多计算框架支持分布式运行，但缺少资源分配和调度的功能。这就是本篇我们所需要解决的问题，我们期望计算平台（集群）能够提供 GPU 任务调度、监控、失败重启等全生命周期管理的功能。当集群计算规模扩大时，如果没有这些功能，我们很难手工地去每一个计算节点上启动计算任务，也无法实时监控任务运行。除此之外，目前大多数分布式计算框架不支持生命周期管理，结束的训练进程并不会自动关闭，这也需要进行额外的处理。同时，当多用户共享使用计算资源时，如果依靠人工协调资源分配，会带来集群资源利用效率低下和使用繁琐等问题。总之，当存在多机集群或者多用户共享使用的情况时，我们需要一种平台帮助我们实现以下目标：

GPU 计算资源的自动调度与管理

计算任务的全生命周期管理

自动实现多用户任务共享使用计算资源

而我们本篇将要介绍的 Kubernetes 能够满足我们上述要求。Kubernetes (简称：k8s) 是 Google 开源的容器集群管理系统（内部代号：Borg）。Kubernetes 是一套完备的分布式系统平台，具有完备的集群管理能力，以容器技术为基础，为容器化的应用提供部署运行、资源调度、服务发现和动态伸缩等一系列完整功能，提高了大规模容器集群管理的便捷性，包括：容器自动化部署、可扩展资源自动调度机制以及多粒度的资源配额管理等。这是 Kubernetes 针对传统云服务和云计算所提供的一揽子功能。然而要实现我们所需要的目标，仅仅具备这些功能是不够的。我们需要能够对 GPU 资源进行管理和调度，包括更进一步的配额管理功能等。

Kubernetes 从 1.6 版本开始，逐渐开始支持 GPU 资源的调度，而且功能实现也在快速更新迭代，大致分为两个阶段：

在 1.8 版本之前，GPU 的管理在功能实现上比较简单粗暴，处于试验阶段，具备初步的基本功能，但还不能大规模的应用在生产环境。

在 1.8 版本之后，Kubernetes 重构了 GPU 的管理逻辑，引入 Device Manager 和 Device Plugin 组件，以更加规范完善的框架实现对扩展硬件的支持，不局限于只支持 GPU 硬件资源，还支持 FPGA，InfiniteBand 等。除此之外，还添加了设备健康检查等新功能。

本篇以 Kubernetes 1.10 版本为例，介绍 Kubernetes 如何实现 GPU 资源管理与调度，以及用户如何利用 Kubernetes 提交 GPU 任务。

Kubernetes 的 GPU 资源管理与调度

Kubernetes 1.8 版本之后，由 Device Manager 管理扩展硬件资源，包括 GPU，FPGA，InfiniteBand 等。不同硬件资源或者同一种硬件的不同厂商均可开发对应设备的 Device Plugin。Device Plugin 按照约定的 API 与 Kubernetes 中的 Device Manager 进行交互，实现对设备的管理、分配、回收以及运行状态监控和健康检查。

我们通过 Kubernetes 的代码可以发现，Device Manager 的功能逻辑集成在 Kubernetes 的核心组件 kubelet 代码中。Device Manager 与 Device Plugin 交互方式如下图所示。图中左侧是 Device Manager 的功能逻辑组件（绿色），右侧是 Device Plugin 的功能逻辑组件（红色），其中 Device Plugin 的红色色块代表功能执行，Device Manager 的绿色色块代表功能发起，空心矩形框为双方的 gRPC Server 逻辑功能。双方的 gPRC 交互过程分为四种，具体包括：

设备注册

设备监控

设备分配

设备回收

我们以 Nvidia GPU 设备为例，详细介绍 Device Manager 和 Device Plugin 如何通过这四个交互过程实现对 GPU 的资源管理和调度。Nvidia 官方发布了k8s-device-plugin（github.com/NVIDIA/k8s-device-plugin），我们以此作为 Device Plugin 的示例，下文简称插件。

一、Device Plugin

准备工作

k8s-device-plugin 在运行之前，需要准备好 GPU 驱动，确保 GPU 能够正常工作。k8s-device-plugin 并不提供 GPU 驱动安装工作，需要管理员在每一个计算节点提前安装好 GPU 驱动。

k8s-device-plugin 以容器方式或者直接运行在计算节点上，推荐以 daemon set 方式，插件启动的 YAML 参考官方推荐配置文件（github.com/NVIDIA/k8s-device-plugin/blob/v1.10/nvidia-device-plugin.yml）。

k8s-device-plugin 需要和 Nvidia runtime 配合使用，在集群使用之前，每一个 GPU 计算节点上安装 Nvidia runtime，并修改 /etc/docker/daemon.json，设置 Nvidia runtime 为默认的 docker runtime，如下所示：

{    "default-runtime": "nvidia",    "runtimes": {        "nvidia": {            "path": "/usr/bin/nvidia-container-runtime",            "runtimeArgs": []        }    } }

Kubernetes 版本 1.8、1.9 需要添加 kubelet 参数 --feature-gates=DevicePlugins=true 开启 Device Plugin 功能，在 1.10 版本，此功能已经默认开启，无需再添加 kubelet 参数。

设备注册

Device Manager 内部维护一个供插件注册的 gRPC server，k8s-device-plugin 利用 gRPC 协议，并通过 /var/lib/kubelet/device-plugins/kubelet.sock 向 Device Manager 发起注册请求，对应的是图中红色色块的 Register 向左边的 Registry gRPC Server 发送注册信息，k8s-device-plugin 注册过程如下：

k8s-device-plugin 向 Device Manager 发起 RegisterRequest gRPC 请求，汇报 GPU 设备的信息：

设备名称 nvidia.com/gpu

API 版本号

k8s-device-plugin 的 gPRC server unix socket: /var/lib/kubelet/device-plugins/nvidia.sock

Device Manager 响应 RegisterRequest，利用 RegisterResponse 返回响应结果，如果设备注册成功后，更新节点扩展设备资源状态信息。

如果 Device Manager 响应成功，k8s-device-plugin 启动插件端的 gPRC Server，供 Device Manager 与 k8s-device-plugin 通信。

k8s-device-plugin 通过 kubelet.sock 持续监控 kubelet 的运行状态，当 kubelet 重启后，k8s-device-plugin 需要向 Device Manager 重新注册。由于双方需要通过 socket 通讯，当k8s-device-plugin 以容器的方式运行时，需要挂载主机的 /var/lib/kubelet/device-plugins/ 目录。

设备发现

当插件注册成功后，Device Manager 通过 ListAndWatch gRPC 请求获取当前设备的列表和健康状态，这个交互是双向的，如果设备状态发生改变，比如当 Device Plugin 检测到某个设备不健康的时候，就会主动通知 Device Manager。如果这个不健康的设备处于空闲状态，Device Manager 就会将其挪出可分配列表。如果该设备已经被某个任务使用，kubelet 中止此任务的使用。

设备发现的功能依赖每个插件根据不同设备的具体情况做出不同的处理。k8s-device-plugin 的设备发现过程调用如下：getDevices 函数获取 GPU UUID ，并传递给 Device Manager。GPU UUID 由 k8s-device-plugin 调用 Nvidia 的 NVML 库获取。NVIDIA Management Library (NVML)(developer.nvidia.com/nvidia-management-library-nvml)是 Nvidia 官方发布的基于 C 语言接口，用于监控和管理 Nvidia GPU 的工具库。其编译版随 GPU 驱动一起发布。Nvidia 的 nvidia-smi 和其他常用的第三方工具均使用 NVML 作为管理 GPU 的底层接口，k8s-device-plugin 由于是基于 go 语言的实现，所以直接使用了 nvidia-docker 1.0 中 NVML go 语言 Binding(github.com/NVIDIA/nvidia-docker/tree/1.0/src/nvml)。

func getDevices() []*pluginapi.Device {    n, err := nvml.GetDeviceCount()    check(err)     var devs []*pluginapi.Device    for i := uint(0); i < n; i++ {        d, err := nvml.NewDeviceLite(i)        check(err)        devs = append(devs, &pluginapi.Device{            ID:     d.UUID,            Health: pluginapi.Healthy,        })    }     return devs }

设备分配

设备分配由插件实现 Allocate 功能，负责配置硬件环境。kubelet 创建任务时，通过 gPRC 调用插件的 Allocate，完成设备的分配，以及确保设备在容器中能正常使用。常见的操作包括设置环境变量、挂载 volume、初始化容器所需的设备等。

k8s-device-plugin 是如何实现 GPU 设备分配呢？借助于AI 计算：从单机到集群（上）介绍的 Nvidia runtime，极大简化了 GPU 的分配逻辑，所以 k8s-device-plugin 的 Allocate 实现非常优雅，如下面代码所示，只需要设置容器的环境变量即可。其具体过程是：

Device Manager 传递待分配 GPU UUID 列表，调用 k8s-device-plugin 的 Allocate 模块。

k8s-device-plugin 的 Allocate 模块设置 Nvidia runtime 的环境变量 NVIDIA_VISIBLE_DEVICES。

response := pluginapi.ContainerAllocateResponse{            Envs: map[string]string{                "NVIDIA_VISIBLE_DEVICES": strings.Join(req.DevicesIDs, ","),            },        }

设备回收

在设备回收阶段，插件可以做一些比如驱动卸载等收尾工作。由于 GPU 不需要每次任务结束时卸载驱动，k8s-device-plugin 无需处理设备回收工作，任务资源的释放由 kubelet 控制。

二、Device Manager

Device Manager 的主要功能点包括：

提供注册 gRPC Server， 接受 Device Plugin 的注册请求，获取并维护节点上的设备列表。

与 Device Plugin 保持长连接通信，为 Device Plugin 提供回调函数，当节点设备状态改变时，Device Plugin 通知 Device Manager 根据设备列表信息，更新节点设备状态。

设备的分配与管理，维护当前运行任务与已使用设备的映射关系，提供待分配设备列表，并通过调用 Device Plugin 的 Allocate，协助 kubelet 完成任务设备分配相关的工作。

Device Manager 实现结构体如下，包括：

通信 socket 文件

负责注册的 gRPC Server

activePods 获取当前运行任务

sourceReady 用于从 checkpoint 移除不活动任务

callback 提供回调给 Device Plugin，用于设备监控状态

healthyDevices 健康设备列表，unhealthyDevices 问题设备列表，allocatedDevices 已使用设备列表

podDevices 记录运行任务和已使用的设备映射关系

// ManagerImpl is the structure in charge of managing Device Plugins.type ManagerImpl struct {    socketname string     socketdir  string  endpoints map[string]endpoint // Key is ResourceName     mutex     sync.Mutex     server *grpc.Server     // activePods is a method for listing active pods on the node    // so the amount of pluginResources requested by existing pods    // could be counted when updating allocated devices     activePods ActivePodsFunc     // sourcesReady provides the readiness of kubelet configuration sources such as apiserver update readiness.    // We use it to determine when we can purge inactive pods from checkpointed state.     sourcesReady config.SourcesReady     // callback is used for updating devices' states in one time call.    // e.g. a new device is advertised, two old devices are deleted and a running device fails.     callback monitorCallback // healthyDevices contains all of the registered healthy resourceNames and their exported device IDs.     healthyDevices map[string]sets.String          // unhealthyDevices contains all of the unhealthy devices and their exported device IDs.     unhealthyDevices map[string]sets.String // allocatedDevices contains allocated deviceIds, keyed by resourceName.     allocatedDevices map[string]sets.String // podDevices contains pod to allocated device mapping.     podDevices podDevices    store      utilstore.Store    pluginOpts map[string]*pluginapi.DevicePluginOptions }

由于篇幅限制，这里对 Device Manager 的实现细节不多做介绍，展开说一下 Device Manager 如何维护当前运行任务与已使用设备的映射关系。每个计算节点上的 Device Manager 创建保存当前运行任务与已使用设备映射关系的 checkpoint 文件 /var/lib/kubelet/device-plugins/kubelet_internal_checkpoint ，其内容如下，分为两部分：

PodDeviceEntries 保存节点正在运行任务的 PodID，ContainerName，ResourceName, 正在使用的 DeviceIDs 列表，以及 Device Plugin 返回的消息 AllocResp。其中，DeviceIDs 中多个 GPU UUID 代表多卡任务。

RegisteredDevices 保存节点处于健康状态的设备列表，以 ResourceName 为 key，其值为 DeviceIDs 列表。

将设备使用映射关系通过 checkpoint 文件的方式保存到硬盘上，其目的是为了解决由于 kubelet 重启带来的设备映射关系信息丢失的问题。当 kubelet 重启时，自动读取硬盘上的 checkpoint 文件以获得重启前的设备使用映射关系，保证设备映射关系与实际任务使用的一致性，避免将已经分配的设备当做未使用设备重新分配使用的问题。

{  "PodDeviceEntries": [      {      "PodUID": "51a38fdb-3ef0-11e8-b8fe-0cc47ae55a2c",      "ContainerName": "task1",      "ResourceName": "nvidia.com/gpu",      "DeviceIDs": [          "GPU-77ccee89-7bbc-8838-a56e-f0ca79518232"       ],      "AllocResp": "CkIKFk5WSURJQV9WSVNJQkxFX0RFVklDRVMSKEdQVS03N2NjZWU4OS03YmJjLTg4MzgtYTU2ZS1mMGNhNzk1MTgyMzI="    },         {      "PodUID": "e7f290e2-3be0-11e8-b8fe-0cc47ae55a2c",      "ContainerName": "task2",      "ResourceName": "nvidia.com/gpu",      "DeviceIDs": [        "GPU-93d815e1-0bda-ea1f-08d9-0864e895553d",        "GPU-ffd50dfd-2578-342e-9a53-19b0f3d40852",        "GPU-68aed792-9550-6ef7-bd91-f8422efd7b5a",        "GPU-a6ec8254-2bd0-3237-142a-496fa2059d73"        ],      "AllocResp": "Cr4BChZOVklESUFfVklTSUJMRV9ERVZJQ0VTEqMBR1BVLTkzZDgxNWUxLTBiZGEtZWExZi0wOGQ5LTA4NjRlODk1NTUzZCxHUFUtZmZkNTBkZmQtMjU3OC0zNDJlLTlhNTMtMTliMGYzZDQwODUyLEdQVS02OGFlZDc5Mi05NTUwLTZlZjctYmQ5MS1mODQyMmVmZDdiNWEsR1BVLWE2ZWM4MjU0LTJiZDAtMzIzNy0xNDJhLTQ5NmZhMjA1OWQ3Mw=="      }  ],  "RegisteredDevices": {    "nvidia.com/gpu": [      "GPU-68aed792-9550-6ef7-bd91-f8422efd7b5a",      "GPU-02a18b6f-3098-7c10-33f9-ededd1b150b8",      "GPU-e4ce184a-d4a9-8b90-9ba1-9f40ec4cc2d7",      "GPU-68258d71-d70e-f8bd-9c9a-b7b5240c8b58",      "GPU-a6ec8254-2bd0-3237-142a-496fa2059d73",      "GPU-6d8322d9-b8b5-89ce-2804-5cbaacc7b6ef",      "GPU-93d815e1-0bda-ea1f-08d9-0864e895553d",      "GPU-ffd50dfd-2578-342e-9a53-19b0f3d40852",      "GPU-77ccee89-7bbc-8838-a56e-f0ca79518232",      "GPU-69bf1feb-66bc-67b9-c38e-c5a38ac93e20"         ]  } }

用户使用

我们从用户角度看一下，从任务提交开始的整个交互工作流程：

用户提交任务申请，通过在 YAML 文件里指定 nvidia.com/gpu 申请 X 个 GPU 卡

Scheduler 过滤满足条件的候选节点

任务 Pod 被分发到节点，该节点 Device Manager 决定待分配设备的 GPU UUID 列表

Device Manager 调用 gPRC Allocate，通知 Device Plugin 将 GPU UUID 列表中的设备映射到任务 Pod 中使用（比如上面介绍的：k8s-device-plugin 设置 Nvidia runtime 的环境变量 NVIDIA_VISIBLE_DEVICES）

任务完成创建

我们也给出 YAML 文件示例，如下面的 YAML 文件所示，用户提交 GPU 任务时，只需要通过 nvidia.com/gpu 指定 GPU 使用数量，Kubernetes 的 scheduler 查找合适的节点资源，并自动调度到满足要求的节点上，由节点上的 kubelet 完成任务的启动和运行。如果没有资源空余，则任务会处于 Pending 状态，等待任务需求的资源满足，则自动转入运行状态。

apiVersion: v1 kind: Pod metadata:  name: tensorflow spec:  containers:  - name: tensorflow    args: ["sleep 1d"]    command: ["/bin/sh", "-c"]    image: tensorflow/tensorflow:latest-gpu    resources:      limits:        nvidia.com/gpu: "1"

Kubernetes 支持对不同的 GPU 资源需求做筛选，比如，可以根据 GPU 型号做任务选择调度。如下所示，指定使用 Tesla P100 卡，则 Kubernetes 会自动调度任务到 P100 卡的节点上。

apiVersion: v1 kind: Pod metadata:  name: tensorflow spec:  containers:  - name: tensorflow    args: ["sleep 1d"]    command: ["/bin/sh", "-c"]    image: tensorflow/tensorflow:latest-gpu    resources:      limits:        nvidia.com/gpu: "1"  nodeSelector:    accelerator: nvidia-tesla-p100

这里需要进一步说明下使用 k8s-device-plugin 的一个小 bug，由于 GPU 计算节点上的 docker runtime 默认设置为 Nvidia runtime，而 Nvidia runtime 的 NVIDIA_VISIBLE_DEVICES 环境变量默认值为 all。所以，当用户提交非 GPU 任务时，如下所示，在 YAML 文件中没有指定 nvidia.com/gpu，则此时，在任务容器中能使用该节点上的所有 GPU 资源。这显然不是我们期望的，绕过了原有的资源分配逻辑。一种解决方式是在 YAML 文件中显式设置容器的环境变量 NVIDIA_VISIBLE_DEVICES，如下所示：

apiVersion: v1 kind: Pod metadata:  name: tensorflowspec:  containers:  - name: tensorflow    args: ["sleep 1d"]    command: ["/bin/sh", "-c"]    image: tensorflow/tensorflow:latest-gpu    env:      - name: NVIDIA_VISIBLE_DEVICES回顾与展望

本文开始提到 Kubernetes 1.8 之前版本的 GPU 管理比较简单，这里对 1.8 之前版本的实现逻辑简单介绍一下，并与本文介绍的 1.8 版本及之后的实现做对比，方便读者了解 Kubernetes GPU 资源管理和调度的演进历史，并能对当前实现方式的特点有直观的认识。

设备发现：1.8 之前版本是通过直接匹配计算节点上的设备文件 /dev/nvidia* 、/dev/nvidia-uvm 和 /dev/nvidiactl。这是一种折中做法，并不是真实的查询设备信息，存在不可靠的问题。资源的名称alpha.kubernetes.io/nvidia-gpu，和当前的nvidia.com/gpu也有区别。

设备分配：1.8 之前版本用户需要在 YAML 文件中显式挂载驱动相关动态库，1.8 之后版本用户不需要挂载。除此之外，在 Allocate 这部分，1.8 之前版本是通过 docker 的 API 实现，这其实不满足软件通用性的要求，也是一个临时的折中做法。1.8 之后版本是通过 Device Manager + Device Plugin 实现。

代码耦合：1.8 之前版本 GPU 资源管理的代码耦合在 Kubernetes 代码中，不利于社区贡献，而且增加了 Kubernetes 稳定运行的风险。

基于 Device Manager + Device Plugin 方式管理扩展硬件设备，不局限于 GPU，还可以管理 InfiniBand，高性能网卡，FPGA 等。

1.8 之后的版本新增了设备健康检查，此功能依赖于具体设备的 Device Plugin 实现，1.8 之前版本无健康监控功能。

目前，Kubernetes 社区对 GPU 的功能实现也在快速迭代，大致集中在两个方向：

当前 GPU 调度的数量均以整数为单位，考虑到利用 Kubernetes 做 GPU Inference，在后续功能改进上，将来有可能实现类似 0.5 这种非整数的调度单位，多容器之间共享 GPU 计算资源。

除了继续完善 GPU 管理调度的功能外，Kubernetes 与常用 AI 计算框架的结合也是社区工作的重点，两者的紧密配合将会带来更加便捷和高效的 AI 计算。

总结

在利用容器技术简化 AI 计算框架单机运行的基础上，本文详细介绍了利用 Kubernetes 实现集群的 GPU 资源调度和管理，重点介绍了当前的 Device Plugin + Device Manager 实现逻辑。并以 Nvidia 官方的 k8s-device-plugin 为例，分析了 k8s-device-plugin 与 Device Manager 的功能交互过程。在介绍 Device Manager 功能点之后，我们从用户的角度梳理了在 Kubernetes 上 GPU 任务提交的工作流程，并针对不同的应用场景，给出了三个示例。最后，我们回顾了 Kubernetes GPU 功能的演进历史，对比了 1.8 前后两个版本实现的优缺点，并展望了未来的发展方向。