Device Plugin 机制深度解析

// ==================== 主入口 ====================
// cmd/nvidia-device-plugin/main.go

package main

import (
    "github.com/NVIDIA/k8s-device-plugin/internal/plugin"
    "github.com/NVIDIA/go-nvml/pkg/nvml"
)

func main() {
    // 1. 初始化 NVML
    ret := nvml.Init()
    if ret != nvml.SUCCESS {
        log.Fatal("Failed to initialize NVML")
    }
    defer nvml.Shutdown()

    // 2. 加载配置
    config := loadConfig()

    // 3. 创建 Plugin Manager
    manager := plugin.NewManager(config)

    // 4. 启动并运行
    manager.Run()
}

// ==================== Plugin Manager ====================
// internal/plugin/manager.go

type Manager struct {
    config    *config.Config
    plugins   map[string]*NvidiaDevicePlugin
    nvml      nvml.Interface
}

func NewManager(config *config.Config) *Manager {
    return &Manager{
        config:  config,
        plugins: make(map[string]*NvidiaDevicePlugin),
        nvml:    nvml.New(),
    }
}

func (m *Manager) Run() error {
    // 1. 发现 GPU 设备
    devices, err := m.discoverDevices()
    if err != nil {
        return err
    }

    // 2. 根据配置决定暴露哪些资源
    // 可能是完整 GPU、MIG 设备、或时分复用份额
    resources := m.getResourcesToRegister(devices)

    // 3. 为每种资源创建 Device Plugin
    for _, resource := range resources {
        plugin := NewNvidiaDevicePlugin(resource, m.config)
        m.plugins[resource.Name] = plugin

        // 4. 启动 Plugin (在单独 goroutine)
        go plugin.Start()
    }

    // 5. 监听 kubelet 重启信号
    // kubelet 重启后需要重新注册
    m.watchKubeletRestart()

    // 6. 阻塞等待
    select {}
}

func (m *Manager) discoverDevices() ([]*Device, error) {
    // 使用 NVML 发现 GPU
    count, ret := m.nvml.DeviceGetCount()
    if ret != nvml.SUCCESS {
        return nil, fmt.Errorf("failed to get device count: %v", ret)
    }

    var devices []*Device
    for i := 0; i < count; i++ {
        device, ret := m.nvml.DeviceGetHandleByIndex(i)
        if ret != nvml.SUCCESS {
            continue
        }

        uuid, _ := device.GetUUID()
        name, _ := device.GetName()
        memory, _ := device.GetMemoryInfo()

        // 检查 MIG 模式
        migMode, _ := device.GetMigMode()

        devices = append(devices, &Device{
            Index:      i,
            UUID:       uuid,
            Name:       name,
            Memory:     memory.Total,
            MIGEnabled: migMode == nvml.DEVICE_MIG_ENABLE,
        })
    }

    return devices, nil
}

// ==================== Device Plugin 实现 ====================
// internal/plugin/server.go

type NvidiaDevicePlugin struct {
    resourceName string
    socket       string
    server       *grpc.Server
    devices      []*pluginapi.Device
    health       chan *pluginapi.Device
    stop         chan struct{}
}

func NewNvidiaDevicePlugin(resource *Resource, config *config.Config) *NvidiaDevicePlugin {
    return &NvidiaDevicePlugin{
        resourceName: resource.Name,
        socket:       fmt.Sprintf("%s-%s.sock", "nvidia", resource.Name),
        devices:      resource.Devices,
        health:       make(chan *pluginapi.Device),
        stop:         make(chan struct{}),
    }
}

func (p *NvidiaDevicePlugin) Start() error {
    // 1. 清理旧的 socket 文件
    socketPath := filepath.Join(pluginapi.DevicePluginPath, p.socket)
    os.Remove(socketPath)

    // 2. 创建 gRPC Server
    p.server = grpc.NewServer()
    pluginapi.RegisterDevicePluginServer(p.server, p)

    // 3. 监听 Unix Socket
    listener, err := net.Listen("unix", socketPath)
    if err != nil {
        return err
    }

    // 4. 启动 gRPC Server
    go p.server.Serve(listener)

    // 5. 向 kubelet 注册
    err = p.Register()
    if err != nil {
        return err
    }

    // 6. 启动健康检查
    go p.healthCheck()

    return nil
}

func (p *NvidiaDevicePlugin) Register() error {
    // 连接 kubelet 的注册服务
    conn, err := grpc.Dial(
        pluginapi.KubeletSocket,
        grpc.WithInsecure(),
        grpc.WithDialer(func(addr string, timeout time.Duration) (net.Conn, error) {
            return net.DialTimeout("unix", addr, timeout)
        }),
    )
    if err != nil {
        return err
    }
    defer conn.Close()

    client := pluginapi.NewRegistrationClient(conn)

    // 发送注册请求
    req := &pluginapi.RegisterRequest{
        Version:      pluginapi.Version,
        Endpoint:     p.socket,
        ResourceName: p.resourceName,
        Options: &pluginapi.DevicePluginOptions{
            PreStartRequired:                true,
            GetPreferredAllocationAvailable: true,
        },
    }

    _, err = client.Register(context.Background(), req)
    return err
}

// ListAndWatch 实现 - 流式上报设备列表
func (p *NvidiaDevicePlugin) ListAndWatch(
    e *pluginapi.Empty,
    s pluginapi.DevicePlugin_ListAndWatchServer,
) error {
    // 首次发送完整设备列表
    s.Send(&pluginapi.ListAndWatchResponse{Devices: p.devices})

    // 持续监听设备状态变化
    for {
        select {
        case <-p.stop:
            return nil
        case device := <-p.health:
            // 设备健康状态变化，重新发送完整列表
            p.updateDeviceHealth(device)
            s.Send(&pluginapi.ListAndWatchResponse{Devices: p.devices})
        }
    }
}

// Allocate 实现 - 分配设备给容器
func (p *NvidiaDevicePlugin) Allocate(
    ctx context.Context,
    req *pluginapi.AllocateRequest,
) (*pluginapi.AllocateResponse, error) {
    response := &pluginapi.AllocateResponse{}

    for _, containerReq := range req.ContainerRequests {
        // 获取请求的设备 ID
        deviceIDs := containerReq.DevicesIDs

        // 验证设备是否存在且健康
        for _, id := range deviceIDs {
            if !p.deviceExists(id) {
                return nil, fmt.Errorf("device %s not found", id)
            }
            if !p.deviceHealthy(id) {
                return nil, fmt.Errorf("device %s is unhealthy", id)
            }
        }

        // 构建容器响应
        containerResp := &pluginapi.ContainerAllocateResponse{
            // 设置环境变量，告诉容器运行时使用哪些 GPU
            Envs: map[string]string{
                "NVIDIA_VISIBLE_DEVICES": strings.Join(deviceIDs, ","),
            },
            // 设备节点挂载 (实际上 nvidia-container-runtime 会处理)
            Devices: []*pluginapi.DeviceSpec{
                {
                    ContainerPath: "/dev/nvidiactl",
                    HostPath:      "/dev/nvidiactl",
                    Permissions:   "rw",
                },
                {
                    ContainerPath: "/dev/nvidia-uvm",
                    HostPath:      "/dev/nvidia-uvm",
                    Permissions:   "rw",
                },
            },
        }

        // 添加每个 GPU 的设备节点
        for i, id := range deviceIDs {
            containerResp.Devices = append(containerResp.Devices, &pluginapi.DeviceSpec{
                ContainerPath: fmt.Sprintf("/dev/nvidia%d", i),
                HostPath:      fmt.Sprintf("/dev/nvidia%d", p.getDeviceIndex(id)),
                Permissions:   "rw",
            })
        }

        response.ContainerResponses = append(response.ContainerResponses, containerResp)
    }

    return response, nil
}

// GetPreferredAllocation 实现 - 返回首选的设备分配方案
// 用于拓扑感知调度
func (p *NvidiaDevicePlugin) GetPreferredAllocation(
    ctx context.Context,
    req *pluginapi.PreferredAllocationRequest,
) (*pluginapi.PreferredAllocationResponse, error) {
    response := &pluginapi.PreferredAllocationResponse{}

    for _, containerReq := range req.ContainerRequests {
        available := containerReq.AvailableDeviceIDs
        mustInclude := containerReq.MustIncludeDeviceIDs
        size := int(containerReq.AllocationSize)

        // 获取首选分配
        // 考虑 NVLink 连接、NUMA 亲和性等
        preferred := p.getPreferredDevices(available, mustInclude, size)

        response.ContainerResponses = append(response.ContainerResponses,
            &pluginapi.ContainerPreferredAllocationResponse{
                DeviceIDs: preferred,
            })
    }

    return response, nil
}

// 获取首选设备 - 考虑拓扑
func (p *NvidiaDevicePlugin) getPreferredDevices(
    available []string,
    mustInclude []string,
    size int,
) []string {
    // 简化实现：优先选择同一 NVLink 域的 GPU

    // 1. 首先包含必须包含的设备
    selected := make(map[string]bool)
    for _, id := range mustInclude {
        selected[id] = true
    }

    // 2. 获取设备拓扑信息
    topology := p.getDeviceTopology()

    // 3. 按 NVLink 连接度排序
    // 优先选择与已选设备有 NVLink 连接的设备
    candidates := p.sortByNVLinkAffinity(available, selected, topology)

    // 4. 选择剩余设备
    for _, id := range candidates {
        if len(selected) >= size {
            break
        }
        if !selected[id] {
            selected[id] = true
        }
    }

    // 转换为列表
    result := make([]string, 0, len(selected))
    for id := range selected {
        result = append(result, id)
    }
    return result
}

// 健康检查
func (p *NvidiaDevicePlugin) healthCheck() {
    for {
        select {
        case <-p.stop:
            return
        case <-time.After(30 * time.Second):
            for _, device := range p.devices {
                health := p.checkDeviceHealth(device.ID)
                if health != device.Health {
                    device.Health = health
                    p.health <- device
                }
            }
        }
    }
}

func (p *NvidiaDevicePlugin) checkDeviceHealth(id string) string {
    // 使用 NVML 检查设备健康状态
    handle, ret := nvml.DeviceGetHandleByUUID(id)
    if ret != nvml.SUCCESS {
        return pluginapi.Unhealthy
    }

    // 检查 ECC 错误
    eccErrors, ret := handle.GetTotalEccErrors(
        nvml.MEMORY_ERROR_TYPE_UNCORRECTED,
        nvml.VOLATILE_ECC,
    )
    if ret == nvml.SUCCESS && eccErrors > 0 {
        return pluginapi.Unhealthy
    }

    // 检查温度
    temp, ret := handle.GetTemperature(nvml.TEMPERATURE_GPU)
    if ret == nvml.SUCCESS && temp > 90 {
        return pluginapi.Unhealthy
    }

    return pluginapi.Healthy
}

// internal/rm/mig.go

// MIG 资源管理器
type MIGResourceManager struct {
    nvml    nvml.Interface
    devices map[string]*MIGDevice
}

type MIGDevice struct {
    ParentUUID       string
    GIProfileID      int
    CIProfileID      int
    UUID             string
    Memory           uint64
    Name             string  // 如 "3g.40gb"
}

func (m *MIGResourceManager) DiscoverMIGDevices() ([]*MIGDevice, error) {
    var migDevices []*MIGDevice

    // 获取所有 GPU
    count, _ := m.nvml.DeviceGetCount()

    for i := 0; i < count; i++ {
        device, _ := m.nvml.DeviceGetHandleByIndex(i)

        // 检查是否启用 MIG
        migMode, _ := device.GetMigMode()
        if migMode != nvml.DEVICE_MIG_ENABLE {
            continue
        }

        parentUUID, _ := device.GetUUID()

        // 获取所有 GPU Instance
        giCount, _ := device.GetMaxMigDeviceCount()

        for gi := 0; gi < giCount; gi++ {
            giHandle, ret := device.GetMigDeviceHandleByIndex(gi)
            if ret != nvml.SUCCESS {
                continue
            }

            // 获取 MIG 设备信息
            uuid, _ := giHandle.GetUUID()
            memory, _ := giHandle.GetMemoryInfo()

            // 获取 Profile 信息
            giInfo, _ := giHandle.GetGpuInstanceProfileInfo()
            ciInfo, _ := giHandle.GetComputeInstanceProfileInfo()

            migDevice := &MIGDevice{
                ParentUUID:  parentUUID,
                GIProfileID: giInfo.Id,
                CIProfileID: ciInfo.Id,
                UUID:        uuid,
                Memory:      memory.Total,
                Name:        m.getProfileName(giInfo, ciInfo),
            }

            migDevices = append(migDevices, migDevice)
        }
    }

    return migDevices, nil
}

// 根据 MIG Profile 生成资源名称
func (m *MIGResourceManager) GetResourceName(device *MIGDevice) string {
    // 格式: nvidia.com/mig-<profile>
    // 例如: nvidia.com/mig-3g.40gb
    return fmt.Sprintf("nvidia.com/mig-%s", device.Name)
}

// MIG 策略配置
type MIGStrategy string

const (
    // Single: 假设所有 GPU 都是相同的 MIG 配置
    // 适合所有 GPU 配置相同的集群
    MIGStrategySingle MIGStrategy = "single"

    // Mixed: 支持 MIG 和非 MIG GPU 混合
    // 会暴露 nvidia.com/gpu 和 nvidia.com/mig-* 两种资源
    MIGStrategyMixed MIGStrategy = "mixed"

    // None: 忽略 MIG 设备，只暴露完整 GPU
    MIGStrategyNone MIGStrategy = "none"
)

// custom-device-plugin/main.go
// 自定义 Device Plugin 开发模板

package main

import (
    "context"
    "fmt"
    "net"
    "os"
    "path/filepath"
    "time"

    "google.golang.org/grpc"
    pluginapi "k8s.io/kubelet/pkg/apis/deviceplugin/v1beta1"
)

const (
    resourceName = "example.com/custom-device"
    socketName   = "custom-device.sock"
)

// CustomDevicePlugin 自定义设备插件
type CustomDevicePlugin struct {
    devices     []*pluginapi.Device
    socket      string
    server      *grpc.Server
    healthCheck chan *pluginapi.Device
    stop        chan struct{}
}

// NewCustomDevicePlugin 创建新的插件实例
func NewCustomDevicePlugin() *CustomDevicePlugin {
    return &CustomDevicePlugin{
        socket:      filepath.Join(pluginapi.DevicePluginPath, socketName),
        healthCheck: make(chan *pluginapi.Device),
        stop:        make(chan struct{}),
    }
}

// discoverDevices 发现设备
// 这里需要实现具体的设备发现逻辑
func (p *CustomDevicePlugin) discoverDevices() []*pluginapi.Device {
    // 示例：发现 4 个自定义设备
    devices := make([]*pluginapi.Device, 4)
    for i := 0; i < 4; i++ {
        devices[i] = &pluginapi.Device{
            ID:     fmt.Sprintf("device-%d", i),
            Health: pluginapi.Healthy,
            // 可选：添加拓扑信息
            Topology: &pluginapi.TopologyInfo{
                Nodes: []*pluginapi.NUMANode{
                    {ID: int64(i % 2)}, // 假设有 2 个 NUMA 节点
                },
            },
        }
    }
    return devices
}

// Start 启动插件
func (p *CustomDevicePlugin) Start() error {
    // 1. 发现设备
    p.devices = p.discoverDevices()
    if len(p.devices) == 0 {
        return fmt.Errorf("no devices found")
    }

    // 2. 清理旧 socket
    os.Remove(p.socket)

    // 3. 创建 gRPC server
    p.server = grpc.NewServer()
    pluginapi.RegisterDevicePluginServer(p.server, p)

    // 4. 监听 socket
    listener, err := net.Listen("unix", p.socket)
    if err != nil {
        return err
    }

    // 5. 启动 server
    go func() {
        p.server.Serve(listener)
    }()

    // 6. 等待 server 启动
    conn, err := p.dial(p.socket, 5*time.Second)
    if err != nil {
        return err
    }
    conn.Close()

    // 7. 注册到 kubelet
    if err := p.register(); err != nil {
        return err
    }

    // 8. 启动健康检查
    go p.runHealthCheck()

    return nil
}

// register 向 kubelet 注册
func (p *CustomDevicePlugin) register() error {
    conn, err := p.dial(pluginapi.KubeletSocket, 5*time.Second)
    if err != nil {
        return err
    }
    defer conn.Close()

    client := pluginapi.NewRegistrationClient(conn)

    req := &pluginapi.RegisterRequest{
        Version:      pluginapi.Version,
        Endpoint:     socketName,
        ResourceName: resourceName,
        Options: &pluginapi.DevicePluginOptions{
            PreStartRequired:                false,
            GetPreferredAllocationAvailable: true,
        },
    }

    _, err = client.Register(context.Background(), req)
    return err
}

// GetDevicePluginOptions 返回插件选项
func (p *CustomDevicePlugin) GetDevicePluginOptions(
    ctx context.Context,
    e *pluginapi.Empty,
) (*pluginapi.DevicePluginOptions, error) {
    return &pluginapi.DevicePluginOptions{
        PreStartRequired:                false,
        GetPreferredAllocationAvailable: true,
    }, nil
}

// ListAndWatch 列出并监控设备
func (p *CustomDevicePlugin) ListAndWatch(
    e *pluginapi.Empty,
    s pluginapi.DevicePlugin_ListAndWatchServer,
) error {
    // 发送初始设备列表
    s.Send(&pluginapi.ListAndWatchResponse{Devices: p.devices})

    // 持续监控
    for {
        select {
        case <-p.stop:
            return nil
        case device := <-p.healthCheck:
            // 更新设备健康状态
            for i, d := range p.devices {
                if d.ID == device.ID {
                    p.devices[i].Health = device.Health
                    break
                }
            }
            // 发送更新后的设备列表
            s.Send(&pluginapi.ListAndWatchResponse{Devices: p.devices})
        }
    }
}

// Allocate 分配设备
func (p *CustomDevicePlugin) Allocate(
    ctx context.Context,
    req *pluginapi.AllocateRequest,
) (*pluginapi.AllocateResponse, error) {
    response := &pluginapi.AllocateResponse{}

    for _, containerReq := range req.ContainerRequests {
        containerResp := &pluginapi.ContainerAllocateResponse{
            // 设置环境变量
            Envs: map[string]string{
                "CUSTOM_DEVICES": fmt.Sprintf("%v", containerReq.DevicesIDs),
            },
            // 挂载设备
            Mounts: []*pluginapi.Mount{
                {
                    ContainerPath: "/dev/custom",
                    HostPath:      "/dev/custom",
                    ReadOnly:      false,
                },
            },
            // 设备节点
            Devices: []*pluginapi.DeviceSpec{},
        }

        // 为每个请求的设备添加设备节点
        for _, id := range containerReq.DevicesIDs {
            containerResp.Devices = append(containerResp.Devices, &pluginapi.DeviceSpec{
                ContainerPath: fmt.Sprintf("/dev/custom/%s", id),
                HostPath:      fmt.Sprintf("/dev/custom/%s", id),
                Permissions:   "rw",
            })
        }

        response.ContainerResponses = append(response.ContainerResponses, containerResp)
    }

    return response, nil
}

// GetPreferredAllocation 返回首选分配方案
func (p *CustomDevicePlugin) GetPreferredAllocation(
    ctx context.Context,
    req *pluginapi.PreferredAllocationRequest,
) (*pluginapi.PreferredAllocationResponse, error) {
    response := &pluginapi.PreferredAllocationResponse{}

    for _, containerReq := range req.ContainerRequests {
        // 简单实现：按顺序选择
        var preferred []string
        for _, id := range containerReq.AvailableDeviceIDs {
            preferred = append(preferred, id)
            if len(preferred) >= int(containerReq.AllocationSize) {
                break
            }
        }

        response.ContainerResponses = append(response.ContainerResponses,
            &pluginapi.ContainerPreferredAllocationResponse{
                DeviceIDs: preferred,
            })
    }

    return response, nil
}

// PreStartContainer 容器启动前钩子
func (p *CustomDevicePlugin) PreStartContainer(
    ctx context.Context,
    req *pluginapi.PreStartContainerRequest,
) (*pluginapi.PreStartContainerResponse, error) {
    // 可以在这里做一些容器启动前的准备工作
    return &pluginapi.PreStartContainerResponse{}, nil
}

// runHealthCheck 运行健康检查
func (p *CustomDevicePlugin) runHealthCheck() {
    ticker := time.NewTicker(30 * time.Second)
    defer ticker.Stop()

    for {
        select {
        case <-p.stop:
            return
        case <-ticker.C:
            for _, device := range p.devices {
                health := p.checkHealth(device.ID)
                if health != device.Health {
                    p.healthCheck <- &pluginapi.Device{
                        ID:     device.ID,
                        Health: health,
                    }
                }
            }
        }
    }
}

// checkHealth 检查设备健康状态
func (p *CustomDevicePlugin) checkHealth(id string) string {
    // 实现具体的健康检查逻辑
    // 例如：检查设备文件是否存在
    devicePath := fmt.Sprintf("/dev/custom/%s", id)
    if _, err := os.Stat(devicePath); os.IsNotExist(err) {
        return pluginapi.Unhealthy
    }
    return pluginapi.Healthy
}

// Stop 停止插件
func (p *CustomDevicePlugin) Stop() {
    close(p.stop)
    if p.server != nil {
        p.server.Stop()
    }
}

// dial 连接到 socket
func (p *CustomDevicePlugin) dial(socket string, timeout time.Duration) (*grpc.ClientConn, error) {
    ctx, cancel := context.WithTimeout(context.Background(), timeout)
    defer cancel()

    return grpc.DialContext(ctx, socket,
        grpc.WithInsecure(),
        grpc.WithBlock(),
        grpc.WithContextDialer(func(ctx context.Context, addr string) (net.Conn, error) {
            return net.DialTimeout("unix", addr, timeout)
        }),
    )
}

func main() {
    plugin := NewCustomDevicePlugin()

    // 启动插件
    if err := plugin.Start(); err != nil {
        fmt.Printf("Failed to start plugin: %v\n", err)
        os.Exit(1)
    }

    fmt.Println("Custom Device Plugin started")

    // 监听 kubelet 重启
    // kubelet 重启后需要重新注册
    watcher, err := fsnotify.NewWatcher()
    if err != nil {
        fmt.Printf("Failed to create watcher: %v\n", err)
        os.Exit(1)
    }
    defer watcher.Close()

    watcher.Add(pluginapi.DevicePluginPath)

    for {
        select {
        case event := <-watcher.Events:
            if event.Name == pluginapi.KubeletSocket && event.Op&fsnotify.Create == fsnotify.Create {
                fmt.Println("Kubelet restarted, re-registering...")
                plugin.Stop()
                plugin = NewCustomDevicePlugin()
                if err := plugin.Start(); err != nil {
                    fmt.Printf("Failed to restart plugin: %v\n", err)
                }
            }
        case err := <-watcher.Errors:
            fmt.Printf("Watcher error: %v\n", err)
        }
    }
}