Device Plugin 框架与 HAMi 实现

// cmd/device-plugin/nvidia/main.go
func start(c *cli.Context, o *options) error {
    util.NodeName = os.Getenv(util.NodeNameEnvName)
    client.InitGlobalClient()
    // 监听 kubelet socket 目录变化
    watcher, err := watch.Files(kubeletSocketDir)
    // 启动插件
    plugins, restartPlugins, err := startPlugins(c, o)
    // 事件循环：监控 kubelet 重启、信号处理
    for {
        select {
        case <-restartTimeout:
            goto restart
        case event := <-watcher.Events:
            // kubelet socket 重建 -> 重启插件
        case s := <-sigs:
            // SIGHUP -> 重启，其他信号 -> 退出
        }
    }
}

func (plugin *NvidiaDevicePlugin) getAPIDevices() *[]*device.DeviceInfo {
    devs := plugin.Devices()
    nvml.Init()
    for UUID := range devs {
        ndev, _ := nvml.DeviceGetHandleByUUID(UUID)
        idx, _ := ndev.GetIndex()
        memory, _ := ndev.GetMemoryInfo()
        Model, _ := ndev.GetName()
        numa, _ := GetNumaNode(ndev)
        // 应用 DeviceMemoryScaling / DeviceCoreScaling / DeviceSplitCount
        res = append(res, &device.DeviceInfo{
            ID: UUID, Index: idx, Count: splitCount,
            Devmem: scaledMemory, Devcore: scaledCore,
            Type: Model, Numa: numa, Health: health,
        })
    }
}

func (plugin *NvidiaDevicePlugin) RegisterInAnnotation() error {
    devices := plugin.getAPIDevices()
    encodeddevices := device.MarshalNodeDevices(*devices)
    if encodeddevices == plugin.deviceCache {
        return nil  // 设备信息无变化，跳过
    }
    plugin.deviceCache = encodeddevices
    annos[nvidia.RegisterAnnos] = encodeddevices
    return util.PatchNodeAnnotations(node, annos)
}

// pkg/device-plugin/nvidiadevice/nvinternal/plugin/server.go
func (plugin *NvidiaDevicePlugin) Register(kubeletSocket string) error {
    conn, _ := plugin.dial(kubeletSocket, 5*time.Second)
    client := kubeletdevicepluginv1beta1.NewRegistrationClient(conn)
    reqt := &kubeletdevicepluginv1beta1.RegisterRequest{
        Version:      kubeletdevicepluginv1beta1.Version,
        Endpoint:     path.Base(plugin.socket),
        ResourceName: string(plugin.rm.Resource()),
    }
    client.Register(context.Background(), reqt)
}

// pkg/device/devices.go
func CheckHealth(devType string, node *corev1.Node) (bool, bool) {
    handshake := node.Annotations[util.HandshakeAnnos[devType]]
    if strings.Contains(handshake, "Requesting") {
        formertime, _ := time.Parse(time.DateTime,
            strings.Split(handshake, "_")[1])
        // 如果 Requesting 超过 60 秒未更新，则认为不健康
        return time.Now().Before(formertime.Add(time.Second * 60)), false
    } else if strings.Contains(handshake, "Deleted") {
        return true, false
    } else {
        // 初次握手或已 Reported，发起新的 Requesting
        return true, true
    }
}

func (r *nvmlResourceManager) checkHealth(stop <-chan interface{},
    devices Devices, unhealthy chan<- *Device, ...) error {

    // 获取要关注的 XID 错误列表
    xids := getHealthCheckXids()

    // 为每个设备注册事件监听
    eventMask := uint64(nvml.EventTypeXidCriticalError |
                        nvml.EventTypeDoubleBitEccError |
                        nvml.EventTypeSingleBitEccError)
    for _, d := range devices {
        gpu.RegisterEvents(eventMask & supportedEvents, eventSet)
    }

    // 事件循环
    for {
        e, ret := eventSet.Wait(5000)  // 5 秒超时轮询
        if e.EventType == nvml.EventTypeXidCriticalError {
            if !xids.IsDisabled(e.EventData) {
                unhealthy <- d  // 将设备标记为不健康
            }
        }
    }
}

func (nv *NvidiaDevicePlugin) DisableOtherNVMLOperation() {
    CreateMigApplyLock()
    nv.disableHealthChecks <- true       // 暂停健康检查
    nv.disableWatchAndRegister <- true   // 暂停设备注册
    // 等待两个协程确认已暂停
}

func (nv *NvidiaDevicePlugin) EnableOtherNVMLOperation() {
    RemoveMigApplyLock()
    nv.disableHealthChecks <- false      // 恢复健康检查
    nv.disableWatchAndRegister <- false  // 恢复设备注册
}