NVIDIA k8s-device-plugin

简介

一个专门为 Kubernetes 设计的组件，用于将 NVIDIA GPU 集成到 Kubernetes 集群中，使其作为可调度资源供容器化工作负载使用。它是 NVIDIA GPU 支持在 Kubernetes 中运行的关键部分，特别是在需要高性能计算（HPC）、机器学习或图形处理的场景下。

---

流程图

启动插件
   ↓
初始化 NVML → 发现 GPU
   ↓
启动 gRPC 服务器 → 注册到 kubelet
   ↓
ListAndWatch 上报设备列表 ← 健康检查（循环）
   ↓
Pod 请求 GPU → Allocate 配置环境变量
   ↓
持续运行并更新状态

工作流程概述

1. 插件启动: 插件以 DaemonSet 的形式运行在每个有 NVIDIA GPU 的节点上。
2. 设备发现: 通过 NVIDIA Management Library (NVML) 检测节点上的 GPU。
3. 注册到 kubelet: 向节点的 kubelet 注册插件，声明其管理的资源（例如 nvidia.com/gpu）。
4. 设备上报: 通过 gRPC 接口向 kubelet 上报可用 GPU 的列表及其状态。
5. 健康检查: 持续监控 GPU 的健康状态，并更新上报信息。
6. 资源分配: 当 Pod 请求 GPU 时，插件为容器配置访问 GPU 的环境。
7. 持续运行: 插件保持运行，定期更新设备状态并响应 kubelet 的请求。

---

1. 插件启动

• 运行方式: NVIDIA Device Plugin 通常以 DaemonSet 部署在 Kubernetes 集群中，确保每个有 GPU 的节点上都有一个插件实例。
• 初始化: 插件启动时会加载配置（如资源名称 nvidia.com/gpu），并初始化 NVML 以访问 GPU 硬件。
• 代码入口: 在 main.go 中

func main() {
    plugin, err := NewNvidiaDevicePlugin(/* 配置参数 */)
    if err != nil {
        log.Fatal(err)
    }
    plugin.Serve()
}

---

2. 设备发现

• 目的: 识别节点上所有可用的 NVIDIA GPU。
• 实现: 使用 NVML 的 API 获取 GPU 数量和详细信息。
• 关键代码: 在 nvml.go 中：

func (n *nvmlDeviceLister) GetDevices() ([]*pluginapi.Device, error) {
    count, ret := nvml.DeviceGetCount()
    if ret != nvml.SUCCESS {
        return nil, fmt.Errorf("nvml.DeviceGetCount() failed: %v", ret)
    }
    var devs []*pluginapi.Device
    for i := 0; i < count; i++ {
        dev, ret := nvml.DeviceGetHandleByIndex(i)
        if ret != nvml.SUCCESS {
            continue
        }
        uuid, ret := nvml.DeviceGetUUID(dev)
        if ret != nvml.SUCCESS {
            continue
        }
        devs = append(devs, &pluginapi.Device{
            ID:     uuid,
            Health: pluginapi.Healthy,
        })
    }
    return devs, nil
}

输出: 插件生成一个 pluginapi.Device 列表，每个设备包含唯一的 UUID 和初始健康状态。

---

3. 注册到 kubelet

• 目的: 通知 kubelet 插件的存在及其管理的资源。
• 实现: 插件启动一个 gRPC 服务器，并通过 kubelet 的设备插件注册接口进行注册。
• 关键代码: 在 server.go 中：

func (m *NvidiaDevicePlugin) Serve() error {
    m.server = grpc.NewServer()
    pluginapi.RegisterDevicePluginServer(m.server, m)
    lis, err := net.Listen("unix", m.socket) // e.g., /var/lib/kubelet/device-plugins/nvidia.sock
    if err != nil {
        return err
    }
    go m.server.Serve(lis)
    return m.registerWithKubelet()
}

func (m *NvidiaDevicePlugin) registerWithKubelet() error {
    conn, err := grpc.Dial(kubeletSocket, grpc.WithInsecure(), grpc.WithDialer(...))
    if err != nil {
        return err
    }
    defer conn.Close()
    client := pluginapi.NewRegistrationClient(conn)
    req := &pluginapi.RegisterRequest{
        Version:      pluginapi.Version,
        Endpoint:     path.Base(m.socket),
        ResourceName: m.resourceName, // "nvidia.com/gpu"
    }
    _, err = client.Register(context.Background(), req)
    return err
}

结果: kubelet 知道插件管理 nvidia.com/gpu 资源，并通过 gRPC 端点与其通信

---

4. 设备上报

• 目的: 将 GPU 的可用性和状态实时上报给 kubelet。
• 实现: 插件实现 Kubernetes Device Plugin API 的 ListAndWatch 方法，通过 gRPC 流定期发送设备列表。
• 关键代码: 在 server.go 中：

func (m *NvidiaDevicePlugin) ListAndWatch(e *pluginapi.Empty, s pluginapi.DevicePlugin_ListAndWatchServer) error {
    prevDevices := make(map[string]*pluginapi.Device)
    for {
        currentDevices, err := m.getDevices()
        if err != nil {
            return err
        }
        if !devicesEqual(prevDevices, currentDevices) {
            resp := &pluginapi.ListAndWatchResponse{Devices: currentDevices}
            if err := s.Send(resp); err != nil {
                return err
            }
            prevDevices = make(map[string]*pluginapi.Device)
            for _, dev := range currentDevices {
                prevDevices[dev.ID] = dev
            }
        }
        time.Sleep(10 * time.Second) // 定期更新
    }
}

结果: kubelet 收到设备列表，了解当前可用的 GPU 及其健康状态（如 Healthy 或 Unhealthy）。

---

5. 健康检查

• 目的: 监控 GPU 的运行状态，检测故障或异常。

• 实现: 插件通过 NVML 检查 GPU 的电源状态、温度、XID 错误等，并在状态变化时更新设备列表。

• 关键代码: 在 server.go 中：

  func (m *NvidiaDevicePlugin) healthCheck() {
      for {
          devices, _ := m.getDevices()
          for _, dev := range devices {
              healthy := m.checkHealth(dev.ID)
              if dev.Health != healthy {
                  dev.Health = healthy
                  m.updateDevices(dev)
              }
          }
          time.Sleep(5 * time.Second)
      }
  }
  
  func (m *NvidiaDevicePlugin) checkHealth(uuid string) string {
      dev, ret := nvml.DeviceGetHandleByUUID(uuid)
      if ret != nvml.SUCCESS {
          return pluginapi.Unhealthy
      }
      xidErrors := m.getXidErrors(dev)
      if len(xidErrors) > 0 {
          return pluginapi.Unhealthy
      }
      return pluginapi.Healthy
  }

结果: 不健康的 GPU 被标记为 Unhealthy，不会被分配给 Pod。

---

6. 资源分配

• 目的: 当 Pod 请求 GPU 时，插件为容器配置访问权限。

• 实现: 实现 Allocate 方法，返回环境变量（如 NVIDIA_VISIBLE_DEVICES）给容器。

• 关键代码: 在 server.go 中：

  func (m *NvidiaDevicePlugin) Allocate(ctx context.Context, req *pluginapi.AllocateRequest) (*pluginapi.AllocateResponse, error) {
      var responses []*pluginapi.ContainerAllocateResponse
      for _, r := range req.ContainerRequests {
          resp := &pluginapi.ContainerAllocateResponse{
              Envs: map[string]string{
                  "NVIDIA_VISIBLE_DEVICES": strings.Join(r.DevicesIDs, ","),
              },
          }
          responses = append(responses, resp)
      }
      return &pluginapi.AllocateResponse{ContainerResponses: responses}, nil
  }

结果: 容器启动时通过 NVIDIA 容器运行时访问指定的 GPU。

---

7. 持续运行

• 行为: 插件保持 gRPC 服务器运行，持续执行 ListAndWatch 和健康检查。
• 容错: 如果插件崩溃或节点重启，DaemonSet 会重新部署插件，重复上述流程。

---

结束

NVIDIA Kubernetes Device Plugin 是连接 NVIDIA GPU 硬件和 Kubernetes 调度系统的桥梁。它通过发现 GPU、注册资源、上报状态和分配设备，使得 GPU 在 Kubernetes 中成为一种可管理的资源类型。它的设计充分利用了 NVML 的硬件访问能力和 Kubernetes 的插件机制，为 GPU 密集型工作负载提供了高效、云原生的支持。

简单来说，它让 Kubernetes "认识"并"使用" NVIDIA GPU，就像管理 CPU 和内存一样自然。