这是一个在 Kubernetes (K8s) 生态中至关重要的组件,用于在 K8s 集群中管理和暴露 NVIDIA GPU 资源。

我会从以下几个方面来为你剖析这个项目:

  1. 核心作用:它解决了什么问题?

  2. 工作原理:它如何与 Kubernetes 各组件协同工作?

  3. 关键组件和源码导读:项目的核心代码文件和它们的功能是什么?

  4. 部署与使用:如何在集群中部署并让 Pod 使用 GPU?


1. 核心作用 (What problem does it solve?)

默认情况下,Kubernetes 只知道如何管理和调度两种主要资源:CPU内存 (Memory)。它不认识 GPU、FPGA 或其他特殊硬件。

为了解决这个问题,Kubernetes 引入了 Device Plugin Framework (设备插件框架)。这个框架允许第三方厂商(如 NVIDIA、Intel、AMD 等)编写自己的“插件”,来向 Kubernetes 汇报节点上有哪些自定义硬件资源,并管理这些资源的分配。

NVIDIA/k8s-device-plugin 就是 NVIDIA官方提供的、实现了 K8s Device Plugin 框架的插件。它的核心作用就是:

  • 发现节点上的 GPU:在每个安装了 GPU 的节点上运行,并检测出所有可用的 NVIDIA GPU。

  • 向 Kubelet 注册并上报资源:告诉该节点的 Kubelet:“嗨,我这里有 N 个 NVIDIA GPU 资源,你可以把它们纳入调度了。”

  • 辅助分配 GPU:当一个 Pod 被调度到这个节点并请求 GPU 时,该插件会告诉 Kubelet 如何将一个具体的 GPU 分配给这个 Pod 的容器。

最终,用户就可以像申请 CPU 和内存一样,在 Pod 的 YAML 文件中简单地声明需要多少个 GPU,而 Kubernetes 调度器就能找到有可用 GPU 的节点来运行这个 Pod。


2. 工作原理 (How it Works?)

k8s-device-plugin 的工作流程完美地诠释了 Device Plugin 框架的机制。它主要通过 gRPC 与节点上的 Kubelet 组件进行通信。

整个生命周期可以分为以下几个关键步骤:

![alt text](https://d33wubrfki0l68.cloudfront.net/294a2039231263acc16a7536b940958742880f08/17397/images/blog/2017-12-08-introducing-device-plugins/device-plugin-interaction.png)

(图片来源:Kubernetes 官方博客)

第一步:启动与注册 (Start & Register)

  1. nvidia-device-plugin 通常以 DaemonSet 的形式部署,确保在集群中每个(或指定的)节点上都运行一个实例。

  2. 启动后,它首先会寻找 Kubelet 的 gRPC socket 文件,通常位于 /var/lib/kubelet/device-plugins/kubelet.sock。

  3. 找到后,它会作为 gRPC 客户端,调用 Kubelet 的 Register 接口,告诉 Kubelet:“你好,我是一个设备插件,我管理的资源类型是 nvidia.com/gpu,我的 gRPC 服务地址在 /var/lib/kubelet/device-plugins/nvidia.sock”。

第二步:设备发现与上报 (Discover & Report)

  1. 插件内部会使用 NVML (NVIDIA Management Library) 库来扫描节点上所有物理 GPU 的信息,比如它们的 UUID、健康状况等。

  2. Kubelet 在收到注册请求后,会作为 gRPC 客户端,连接到插件的 gRPC 服务 (nvidia.sock),并调用 ListAndWatch 接口。

  3. 插件通过 ListAndWatch 这个流式 gRPC 连接,将发现的 GPU 列表(以设备 ID 的形式)发送给 Kubelet。

  4. Kubelet 收到这个列表后,会更新该节点(Node)对象的状态,在 API Server 中将节点的 status.allocatable 字段增加 nvidia.com/gpu: N(N 是 GPU 数量)。

  5. 从此,Kubernetes 调度器 (Scheduler) 就知道这个节点有多少个可用的 GPU 资源了。

第三步:Pod 调度与设备分配 (Schedule & Allocate)

  1. 用户提交一个 Pod YAML,其中声明了需要 GPU 资源,例如:

    code Yaml

    downloadcontent_copy

    expand_less
        resources:
      limits:
        nvidia.com/gpu: 1
      
  2. 调度器看到这个请求,就会在集群中寻找一个 status.allocatable 字段中 nvidia.com/gpu 数量大于等于 1 的节点。

  3. 调度器成功将 Pod 调度到目标节点上。

  4. 目标节点上的 Kubelet 准备创建 Pod 的容器。在启动容器之前,它会再次调用设备插件的 gRPC 接口,这次调用的是 Allocate 方法。请求中会包含需要分配的 GPU 数量和设备 ID。

  5. nvidia-device-plugin 的 Allocate 方法被触发。它会选择一个可用的 GPU,并返回一组指令给 Kubelet。这些指令通常包括:

    • 要挂载到容器的设备文件:例如 /dev/nvidia0, /dev/nvidiactl。

    • 需要设置的环境变量:最关键的是 NVIDIA_VISIBLE_DEVICES。例如,NVIDIA_VISIBLE_DEVICES=GPU-f8065a31-4c5e-a6a3-285a-39401219b16e。这个环境变量会告诉容器内的 CUDA 程序只能看到和使用这一个特定的 GPU,从而实现资源隔离。

第四步:容器创建 (Container Creation)

  1. Kubelet 拿到 Allocate 返回的指令。

  2. 它将这些指令传递给底层的 容器运行时 (Container Runtime),如 Docker 或 containerd。

  3. 容器运行时使用这些指令来创建容器:挂载指定的设备文件、设置环境变量。

  4. 容器启动后,内部的应用程序(如 TensorFlow, PyTorch)通过 CUDA 库就能访问到被分配的那个 GPU 了。


3. 关键组件和源码导读

让我们来看一下项目的目录结构和核心文件。

code Code

downloadcontent_copy

expand_less

IGNORE_WHEN_COPYING_START

IGNORE_WHEN_COPYING_END

    .
├── cmd
│   └── nvidia-device-plugin
│       └── main.go         # 程序入口
├── deployments
│   └── static
│       └── nvidia-device-plugin.yml # DaemonSet 部署文件
├── pkg
│   ├── nvidia
│   │   ├── device.go       # 定义设备对象
│   │   ├── manager.go      # 设备管理的核心逻辑
│   │   ├── server.go       # gRPC 服务器的实现
│   │   └── server_v1.go    # 实现了 v1beta1 版本的 Device Plugin API
│   └── rm                  # 资源管理器相关
│       └── nvml.go         # 与 NVML 库交互的接口,用于发现和监控 GPU
└── ...
  

核心文件讲解:

  • cmd/nvidia-device-plugin/main.go:

    • 这是整个程序的入口

    • 它负责解析命令行参数(如 a ss-through-device-specs, fail-on-init-error 等)。

    • 最重要的是,它会启动设备插件的管理器 (manager)。

  • pkg/nvidia/manager.go:

    • 这是业务逻辑的核心。nvidiaDevicePlugin 结构体在这里定义。

    • Start() 方法是启动流程的开始,它会初始化 NVML,发现 GPU,然后启动 gRPC 服务器。

    • discover() 方法调用 rm.GetDevices() 来获取 GPU 列表。

    • 它还负责处理 OS 信号,实现优雅地关闭。

  • pkg/nvidia/server_v1.go:

    • 这里是 gRPC 服务的具体实现,完全遵循 K8s Device Plugin API 的 v1beta1 规范。

    • ListAndWatch 方法: 这是核心之一。它会创建一个流,并立即将当前发现的 GPU 列表发送给 Kubelet。之后它会保持连接,如果 GPU 状态发生变化(比如一个 GPU 变得不健康),它会发送更新。

    • Allocate 方法: 这是另一个核心。当 Kubelet 请求分配 GPU 时,这个方法会被调用。它会根据策略选择一个 GPU,并构建 ContainerAllocateResponse,其中包含了前面提到的设备路径和环境变量 NVIDIA_VISIBLE_DEVICES。

    • GetDevicePluginOptions 和 PreStartContainer 方法也是 API 的一部分,用于一些高级配置和准备工作。

  • pkg/rm/nvml.go:

    • 这是一个非常底层的模块,是与硬件交互的桥梁

    • 它通过 CGO 调用系统的 libnvidia-ml.so 动态链接库,也就是 NVML (NVIDIA Management Library)

    • 所有关于 GPU 的信息,如数量、UUID、健康状况、温度、功耗等,都是通过这个文件中的函数从驱动层获取的。它将底层的 C API 封装成了 Go 语言的接口。

代码阅读建议:

  1. 从 main.go 开始,看它是如何初始化和启动 manager.go 中的 nvidiaDevicePlugin 的。

  2. 接着看 manager.go 的 Start() 方法,理解插件启动、设备发现和 gRPC 服务器启动的流程。

  3. 然后重点阅读 server_v1.go 中的 ListAndWatch 和 Allocate 方法,这是理解与 Kubelet 交互的关键。

  4. 最后,如果你对底层硬件交互感兴趣,可以看 nvml.go,了解它是如何调用 NVML 库来获取 GPU 信息的。


4. 部署与使用

部署前提:

  1. 节点上已正确安装 NVIDIA 驱动。

  2. 节点上已安装 nvidia-container-toolkit,这样容器运行时才能感知到 GPU。

部署步骤:

通常使用官方提供的 DaemonSet YAML 文件进行部署:

code Bash

downloadcontent_copy

expand_less

IGNORE_WHEN_COPYING_START

IGNORE_WHEN_COPYING_END

    kubectl apply -f https://raw.githubusercontent.com/NVIDIA/k8s-device-plugin/v0.14.1/deployments/static/nvidia-device-plugin.yml
  

这个命令会创建一个 DaemonSet,在每个有 nvidia.com/gpu 标签(或没有特定污点)的节点上部署 nvidia-device-plugin Pod。

验证部署:

部署成功后,查看一个有 GPU 的节点,你会看到资源报告:

code Bash

downloadcontent_copy

expand_less

IGNORE_WHEN_COPYING_START

IGNORE_WHEN_COPYING_END

    kubectl describe node <your-gpu-node-name>
  

在输出的 Capacity 和 Allocatable 部分,你应该能看到:

code Code

downloadcontent_copy

expand_less

IGNORE_WHEN_COPYING_START

IGNORE_WHEN_COPYING_END

    Capacity:
  ...
  nvidia.com/gpu: 2
Allocatable:
  ...
  nvidia.com/gpu: 2
  

使用 GPU:

在 Pod 的 spec 中申请 GPU 资源:

code Yaml

downloadcontent_copy

expand_less

IGNORE_WHEN_COPYING_START

IGNORE_WHEN_COPYING_END

    apiVersion: v1
kind: Pod
metadata:
  name: cuda-vector-add
spec:
  restartPolicy: OnFailure
  containers:
    - name: cuda-vector-add
      image: "nvidia/cuda:11.6.2-base-ubuntu20.04"
      command: ["/bin/bash", "-c", "sleep 100000"] # 保持 Pod 运行
      resources:
        limits:
          nvidia.com/gpu: 1 # <-- 申请 1 个 GPU
  

当这个 Pod 启动后,你可以进入容器内部验证:

code Bash

downloadcontent_copy

expand_less

IGNORE_WHEN_COPYING_START

IGNORE_WHEN_COPYING_END

    kubectl exec -it cuda-vector-add -- nvidia-smi
  

你会看到 nvidia-smi 命令的输出,并且它应该只显示一个 GPU 设备,证明资源隔离成功。

总结

NVIDIA/k8s-device-plugin 是一个设计精良且遵循 Kubernetes 标准扩展模式的典范。它作为连接物理 GPU 硬件和 Kubernetes 集群管理系统的“翻译官”,通过标准的 gRPC 接口与 Kubelet 通信,巧妙地将复杂的 GPU 硬件资源无缝集成到 Kubernetes 的资源模型和调度体系中,极大地简化了在 Kubernetes 上运行 AI/ML 等 GPU 加速工作负载的复杂度。

Logo

有“AI”的1024 = 2048,欢迎大家加入2048 AI社区

更多推荐