如何利用Kubernetes Device Plugin管理异构计算设备（FPGA/NPU）？

在现代AI基础设施中，除了标准的GPU，FPGA（现场可编程门阵列）和NPU（神经网络处理器）等异构计算设备扮演着越来越重要的角色。这些设备提供了更高的能效比和定制化的计算能力。然而，Kubernetes原生只支持基本的CPU和内存调度。要让Kubelet识别、分配和管理这些非标准设备，我们必须依赖Kubernetes Device Plugin机制。

本文将深入讲解Device Plugin的工作原理，并提供一个基于Device Plugin API管理异构资源的实操框架。

1. Device Plugin 架构概述

Kubernetes Device Plugin 是一种 Sidecar 模式的扩展机制，允许第三方供应商（如FPGA/NPU厂商）通过实现特定的gRPC服务，向Kubelet通告和管理其设备。

核心工作流程：

1. 发现 (Discovery): Device Plugin 启动后，在主机文件系统（通常是 /var/lib/kubelet/device-plugins/）上创建一个 Unix Socket。
2. 注册 (Registration): Plugin 通过该 Socket 连接到 Kubelet，并向 Kubelet 注册自己提供的资源类型（例如：company.com/fpga_accel）。
3. 监控 (ListAndWatch): Plugin 持续监控设备的健康状态，并通过 gRPC 流实时报告给 Kubelet。
4. 分配 (Allocation): 当 Pod 请求该资源时，Kubelet 调用 Plugin 的 Allocate 接口。Plugin 负责准备设备（例如，进行必要的驱动初始化，设置cgroup规则），并将分配信息返回给 Kubelet，Kubelet随后将这些信息注入到Pod的运行时环境中。

2. Device Plugin gRPC 接口实现要点

一个 Device Plugin 必须实现 Kubernetes 定义的 device.proto 中的 gRPC 服务。以下是实现一个假定 NPU Plugin 的关键步骤。

2.1 定义 gRPC 服务（device.proto 简化）

Device Plugin 至少需要实现 Registration 服务和 DevicePlugin 服务。

// Registration.proto
service Registration {
    rpc Register(RegisterRequest) returns (Empty) {}
}

// DevicePlugin.proto
service DevicePlugin {
    rpc ListAndWatch(Empty) returns (stream ListAndWatchResponse) {}
    rpc Allocate(AllocateRequest) returns (AllocateResponse) {}
}

2.2 插件核心逻辑 (Python 伪代码)

为了演示注册和报告设备的过程，我们使用 Python 结构来展示关键逻辑。一个实际的 NPU 插件需要通过特定的 vendor SDK 来获取设备信息和健康状态。

import grpc
import time
import os

# 假设我们已经生成了 device_plugin_pb2.py 和 device_plugin_pb2_grpc.py

SOCKET_PATH = '/var/lib/kubelet/device-plugins/my-npu.sock'
RESOURCE_NAME = 'mycompany.com/npu_accelerator'

class NPUDevicePlugin(device_plugin_pb2_grpc.DevicePluginServicer):
    def __init__(self):
        # 模拟两个 NPU 设备 ID
        self.devices = {
            'npu-0': device_plugin_pb2.Device.HEALTHY,
            'npu-1': device_plugin_pb2.Device.HEALTHY
        }

    def ListAndWatch(self, request, context):
        # 持续发送设备列表和健康状态
        while True:
            devices_list = []
            for id, health in self.devices.items():
                devices_list.append(
                    device_plugin_pb2.Device(id=id, health=health)
                )

            response = device_plugin_pb2.ListAndWatchResponse(devices=devices_list)
            yield response
            time.sleep(10) # 每10秒报告一次状态

    def Allocate(self, request, context):
        # Kubelet 请求分配设备时调用
        response = device_plugin_pb2.AllocateResponse()
        container_response = device_plugin_pb2.ContainerAllocateResponse()

        # 实际操作：获取设备驱动信息、挂载路径、设置cgroup
        for id in request.container_requests[0].device_ids:
            print(f"Allocating device: {id}")

            # 假设 NPU 需要一个特殊的设备文件挂载
            container_response.devices.append(
                device_plugin_pb2.DeviceSpec(
                    host_path=f'/dev/{id}', 
                    container_path=f'/dev/npu_accel/{id}', 
                    permissions='rwm'
                )
            )

        response.container_responses.append(container_response)
        return response

# 注册逻辑
def register_to_kubelet():
    # Kubelet gRPC Server 监听在 /var/lib/kubelet/device-plugins/kubelet.sock
    channel = grpc.insecure_channel(
        'unix:///var/lib/kubelet/device-plugins/kubelet.sock'
    )
    # ... (连接和发送 RegisterRequest，包含 RESOURCE_NAME 和 SOCKET_PATH)
    print(f"Plugin registered for {RESOURCE_NAME}")

3. 部署 Device Plugin (使用 DaemonSet)

Device Plugin 必须作为 DaemonSet 部署，确保它运行在集群中的每个节点上（或通过 Node Selector 运行在拥有异构设备的节点上）。

关键配置点：

1. HostPath Volume: 必须挂载 Kubelet 的 Device Plugin 目录 (/var/lib/kubelet/device-plugins)，用于通信。
2. Privileged Mode: 通常需要以特权模式运行，以便访问和配置底层设备驱动和/dev目录。

Device Plugin DaemonSet YAML 示例

apiVersion: apps/v1
kind: DaemonSet
metadata:
  name: npu-device-plugin
  namespace: kube-system
spec:
  selector:
    matchLabels:
      name: npu-plugin
  template:
    metadata:
      labels:
        name: npu-plugin
    spec:
      hostNetwork: true
      tolerations:
      - key: "node-role.kubernetes.io/master"
        operator: "Exists"
        effect: "NoSchedule"
      containers:
      - name: npu-plugin-container
        image: your-registry.com/npu-plugin:v1.0
        securityContext:
          privileged: true  # 允许访问底层设备
        volumeMounts:
        - name: device-plugin-socket
          mountPath: /var/lib/kubelet/device-plugins
        - name: dev-dir
          mountPath: /dev  # 如果需要在Allocate中修改设备文件
      volumes:
      - name: device-plugin-socket
        hostPath:
          path: /var/lib/kubelet/device-plugins
      - name: dev-dir
        hostPath:
          path: /dev

4. 验证和调度异构资源

部署 DaemonSet 成功后，Kubelet 会识别新的资源类型。您可以通过查看节点状态来验证资源是否被注册。

4.1 检查 Node Status

kubectl describe node <node-name>

您应该在 Capacity 和 Allocatable 部分看到您的自定义资源：

Capacity:
  cpu:                16
  memory:             64Gi
  mycompany.com/npu_accelerator: 2  # 成功注册了2个NPU
Allocatable:
  cpu:                16
  memory:             64Gi
  mycompany.com/npu_accelerator: 2

4.2 Pod 请求资源

现在，AI 推理或训练工作负载可以通过标准的 resources.limits 字段请求这些 NPU 设备。

apiVersion: v1
kind: Pod
metadata:
  name: npu-worker-pod
spec:
  containers:
  - name: ai-inference-app
    image: deeplearning/npu-image:latest
    resources:
      limits:
        mycompany.com/npu_accelerator: 1

当 Pod 调度到具有 NPU 的节点上时，Kubelet 会调用 Device Plugin 的 Allocate 方法，确保所需的设备文件和挂载点正确配置到容器中，从而实现对异构计算资源的精细化管理和使用。