GPU 调度器实现

// Scheduling Framework 定义了调度器的扩展点
// 路径: kubernetes/pkg/scheduler/framework/interface.go

// PreFilterPlugin 在过滤之前执行，用于预处理或检查 Pod 信息
type PreFilterPlugin interface {
    Plugin
    // PreFilter 在调度周期开始时调用
    // 可用于计算 Pod 需要的资源、检查必要条件等
    PreFilter(ctx context.Context, state *CycleState, pod *v1.Pod) (*PreFilterResult, *Status)
    // PreFilterExtensions 返回扩展接口
    PreFilterExtensions() PreFilterExtensions
}

// FilterPlugin 过滤不满足条件的节点
type FilterPlugin interface {
    Plugin
    // Filter 对每个节点调用，返回节点是否适合运行 Pod
    // 这是最核心的过滤逻辑
    Filter(ctx context.Context, state *CycleState, pod *v1.Pod, nodeInfo *NodeInfo) *Status
}

// PostFilterPlugin 当所有节点都被过滤掉时执行
type PostFilterPlugin interface {
    Plugin
    // PostFilter 尝试抢占等操作使 Pod 可调度
    PostFilter(ctx context.Context, state *CycleState, pod *v1.Pod,
               filteredNodeStatusMap NodeToStatusMap) (*PostFilterResult, *Status)
}

// PreScorePlugin 在打分前执行预处理
type PreScorePlugin interface {
    Plugin
    PreScore(ctx context.Context, state *CycleState, pod *v1.Pod,
             nodes []*v1.Node) *Status
}

// ScorePlugin 为通过过滤的节点打分
type ScorePlugin interface {
    Plugin
    // Score 返回节点的得分，范围 [0, 100]
    Score(ctx context.Context, state *CycleState, pod *v1.Pod,
          nodeName string) (int64, *Status)
    // ScoreExtensions 返回归一化扩展
    ScoreExtensions() ScoreExtensions
}

// ReservePlugin 在绑定前预留资源
type ReservePlugin interface {
    Plugin
    // Reserve 预留节点上的资源
    Reserve(ctx context.Context, state *CycleState, pod *v1.Pod,
            nodeName string) *Status
    // Unreserve 取消预留（当后续步骤失败时调用）
    Unreserve(ctx context.Context, state *CycleState, pod *v1.Pod,
              nodeName string)
}

// PermitPlugin 在绑定前进行最终检查
type PermitPlugin interface {
    Plugin
    // Permit 可以返回：
    // - Success: 立即绑定
    // - Wait: 等待指定时间
    // - Deny: 拒绝调度
    Permit(ctx context.Context, state *CycleState, pod *v1.Pod,
           nodeName string) (*Status, time.Duration)
}

// BindPlugin 执行实际的绑定操作
type BindPlugin interface {
    Plugin
    // Bind 将 Pod 绑定到节点
    Bind(ctx context.Context, state *CycleState, pod *v1.Pod,
         nodeName string) *Status
}

// 路径: kubernetes/pkg/scheduler/framework/plugins/noderesources/fit.go

// NodeResourcesFit 检查节点是否有足够的资源
type NodeResourcesFit struct {
    // 忽略的资源列表
    ignoredResources      sets.String
    // 忽略的资源组
    ignoredResourceGroups sets.String
    // 是否启用 Pod overhead
    enablePodOverhead     bool
    // 评分策略
    scoreStrategy         *config.ScoringStrategy
}

// Filter 过滤资源不足的节点
func (f *NodeResourcesFit) Filter(ctx context.Context, cycleState *framework.CycleState,
    pod *v1.Pod, nodeInfo *framework.NodeInfo) *framework.Status {

    // 获取 Pod 请求的资源
    podRequest := f.getResourceRequest(pod)

    // 检查节点可用资源
    insufficientResources := fitsRequest(podRequest, nodeInfo,
        f.ignoredResources, f.ignoredResourceGroups)

    if len(insufficientResources) != 0 {
        // 资源不足，返回失败原因
        failureReasons := make([]string, 0, len(insufficientResources))
        for _, r := range insufficientResources {
            failureReasons = append(failureReasons, r.Reason)
        }
        return framework.NewStatus(framework.Unschedulable, failureReasons...)
    }
    return nil
}

// fitsRequest 检查节点资源是否满足请求
func fitsRequest(podRequest *preFilterState, nodeInfo *framework.NodeInfo,
    ignoredExtendedResources, ignoredResourceGroups sets.String) []InsufficientResource {

    insufficientResources := make([]InsufficientResource, 0, 4)

    // 检查 CPU
    if podRequest.MilliCPU > (nodeInfo.Allocatable.MilliCPU - nodeInfo.Requested.MilliCPU) {
        insufficientResources = append(insufficientResources, InsufficientResource{
            ResourceName: v1.ResourceCPU,
            Reason:       "Insufficient cpu",
            Requested:    podRequest.MilliCPU,
            Used:         nodeInfo.Requested.MilliCPU,
            Capacity:     nodeInfo.Allocatable.MilliCPU,
        })
    }

    // 检查内存
    if podRequest.Memory > (nodeInfo.Allocatable.Memory - nodeInfo.Requested.Memory) {
        insufficientResources = append(insufficientResources, InsufficientResource{
            ResourceName: v1.ResourceMemory,
            Reason:       "Insufficient memory",
            Requested:    podRequest.Memory,
            Used:         nodeInfo.Requested.Memory,
            Capacity:     nodeInfo.Allocatable.Memory,
        })
    }

    // 检查扩展资源（包括 GPU）
    for rName, rQuant := range podRequest.ScalarResources {
        // 跳过忽略的资源
        if ignoredExtendedResources.Has(string(rName)) {
            continue
        }

        // 检查资源组是否被忽略
        if ignoredResourceGroups.Len() > 0 {
            group := strings.Split(string(rName), "/")[0]
            if ignoredResourceGroups.Has(group) {
                continue
            }
        }

        // 关键：检查 GPU 等扩展资源
        if rQuant > (nodeInfo.Allocatable.ScalarResources[rName] -
                     nodeInfo.Requested.ScalarResources[rName]) {
            insufficientResources = append(insufficientResources, InsufficientResource{
                ResourceName: rName,
                Reason:       fmt.Sprintf("Insufficient %v", rName),
                Requested:    rQuant,
                Used:         nodeInfo.Requested.ScalarResources[rName],
                Capacity:     nodeInfo.Allocatable.ScalarResources[rName],
            })
        }
    }

    return insufficientResources
}

package gpuscheduler

import (
    "sync"
    "time"

    v1 "k8s.io/api/core/v1"
)

// GPUDevice 表示单个 GPU 设备
type GPUDevice struct {
    // 设备 UUID
    UUID string `json:"uuid"`
    // 设备索引
    Index int `json:"index"`
    // GPU 型号
    Model string `json:"model"`
    // 显存大小（MB）
    MemoryTotal int64 `json:"memoryTotal"`
    // 可用显存（MB）
    MemoryFree int64 `json:"memoryFree"`
    // 计算能力
    ComputeCapability string `json:"computeCapability"`
    // 是否支持 MIG
    MIGEnabled bool `json:"migEnabled"`
    // MIG 设备列表
    MIGDevices []MIGDevice `json:"migDevices,omitempty"`
    // 当前状态
    State DeviceState `json:"state"`
    // 分配给的 Pod
    AllocatedPod string `json:"allocatedPod,omitempty"`
    // 健康状态
    Health DeviceHealth `json:"health"`
    // NVLINK 连接
    NVLinkPeers []int `json:"nvlinkPeers,omitempty"`
    // PCIe 信息
    PCIeInfo PCIeInfo `json:"pcieInfo"`
}

// DeviceState GPU 设备状态
type DeviceState string

const (
    DeviceStateFree      DeviceState = "Free"
    DeviceStateAllocated DeviceState = "Allocated"
    DeviceStateReserved  DeviceState = "Reserved"
    DeviceStateFailed    DeviceState = "Failed"
)

// DeviceHealth GPU 健康状态
type DeviceHealth string

const (
    DeviceHealthy   DeviceHealth = "Healthy"
    DeviceUnhealthy DeviceHealth = "Unhealthy"
    DeviceUnknown   DeviceHealth = "Unknown"
)

// MIGDevice MIG 设备实例
type MIGDevice struct {
    // MIG UUID
    UUID string `json:"uuid"`
    // MIG profile（如 "1g.5gb"）
    Profile string `json:"profile"`
    // GPU 实例 ID
    GIIndex int `json:"giIndex"`
    // 计算实例 ID
    CIIndex int `json:"ciIndex"`
    // 状态
    State DeviceState `json:"state"`
    // 分配给的 Pod
    AllocatedPod string `json:"allocatedPod,omitempty"`
}

// PCIeInfo PCIe 拓扑信息
type PCIeInfo struct {
    // Bus ID
    BusID string `json:"busId"`
    // NUMA 节点
    NUMANode int `json:"numaNode"`
    // PCIe 代数
    Generation int `json:"generation"`
    // 链路宽度
    LinkWidth int `json:"linkWidth"`
}

// NodeGPUInfo 节点的 GPU 信息
type NodeGPUInfo struct {
    // 节点名称
    NodeName string `json:"nodeName"`
    // GPU 设备列表
    GPUs []GPUDevice `json:"gpus"`
    // GPU 拓扑
    Topology *GPUTopology `json:"topology"`
    // 驱动版本
    DriverVersion string `json:"driverVersion"`
    // CUDA 版本
    CUDAVersion string `json:"cudaVersion"`
    // 更新时间
    LastUpdated time.Time `json:"lastUpdated"`
}

// GPUTopology GPU 拓扑结构
type GPUTopology struct {
    // 节点名
    NodeName string `json:"nodeName"`
    // 连接矩阵，Links[i][j] 表示 GPU i 到 GPU j 的连接类型
    Links [][]LinkType `json:"links"`
    // NUMA 亲和性，NUMANodes[i] 表示 GPU i 所属的 NUMA 节点
    NUMANodes []int `json:"numaNodes"`
}

// LinkType GPU 间连接类型
type LinkType string

const (
    LinkTypeNone      LinkType = "None"
    LinkTypePCIe      LinkType = "PCIe"
    LinkTypeNVLink    LinkType = "NVLink"
    LinkTypeNVSwitch  LinkType = "NVSwitch"
    LinkTypeSameBoard LinkType = "SameBoard"
)

// GPUStateManager GPU 状态管理器
type GPUStateManager struct {
    mu sync.RWMutex
    // 节点 GPU 信息缓存
    nodeGPUInfo map[string]*NodeGPUInfo
    // 更新通道
    updateChan chan *NodeGPUInfo
    // 停止信号
    stopChan chan struct{}
}

// NewGPUStateManager 创建状态管理器
func NewGPUStateManager() *GPUStateManager {
    return &GPUStateManager{
        nodeGPUInfo: make(map[string]*NodeGPUInfo),
        updateChan:  make(chan *NodeGPUInfo, 100),
        stopChan:    make(chan struct{}),
    }
}

// GetNodeGPUInfo 获取节点 GPU 信息
func (m *GPUStateManager) GetNodeGPUInfo(nodeName string) (*NodeGPUInfo, bool) {
    m.mu.RLock()
    defer m.mu.RUnlock()
    info, ok := m.nodeGPUInfo[nodeName]
    return info, ok
}

// UpdateNodeGPUInfo 更新节点 GPU 信息
func (m *GPUStateManager) UpdateNodeGPUInfo(info *NodeGPUInfo) {
    m.mu.Lock()
    defer m.mu.Unlock()
    info.LastUpdated = time.Now()
    m.nodeGPUInfo[info.NodeName] = info
}

// GetAvailableGPUs 获取节点可用的 GPU 列表
func (m *GPUStateManager) GetAvailableGPUs(nodeName string) []GPUDevice {
    m.mu.RLock()
    defer m.mu.RUnlock()

    info, ok := m.nodeGPUInfo[nodeName]
    if !ok {
        return nil
    }

    var available []GPUDevice
    for _, gpu := range info.GPUs {
        if gpu.State == DeviceStateFree && gpu.Health == DeviceHealthy {
            available = append(available, gpu)
        }
    }
    return available
}

// ReserveGPUs 预留 GPU 设备
func (m *GPUStateManager) ReserveGPUs(nodeName string, gpuIndices []int, podName string) error {
    m.mu.Lock()
    defer m.mu.Unlock()

    info, ok := m.nodeGPUInfo[nodeName]
    if !ok {
        return fmt.Errorf("node %s not found", nodeName)
    }

    // 检查所有 GPU 是否可用
    for _, idx := range gpuIndices {
        if idx >= len(info.GPUs) {
            return fmt.Errorf("GPU index %d out of range", idx)
        }
        if info.GPUs[idx].State != DeviceStateFree {
            return fmt.Errorf("GPU %d is not free", idx)
        }
    }

    // 预留 GPU
    for _, idx := range gpuIndices {
        info.GPUs[idx].State = DeviceStateReserved
        info.GPUs[idx].AllocatedPod = podName
    }

    return nil
}

// UnreserveGPUs 取消预留
func (m *GPUStateManager) UnreserveGPUs(nodeName string, gpuIndices []int) {
    m.mu.Lock()
    defer m.mu.Unlock()

    info, ok := m.nodeGPUInfo[nodeName]
    if !ok {
        return
    }

    for _, idx := range gpuIndices {
        if idx < len(info.GPUs) && info.GPUs[idx].State == DeviceStateReserved {
            info.GPUs[idx].State = DeviceStateFree
            info.GPUs[idx].AllocatedPod = ""
        }
    }
}

package gpuscheduler

import (
    "context"
    "fmt"
    "strconv"

    v1 "k8s.io/api/core/v1"
    "k8s.io/apimachinery/pkg/runtime"
    "k8s.io/kubernetes/pkg/scheduler/framework"
)

const (
    // PluginName 插件名称
    PluginName = "GPUScheduler"

    // GPU 资源名称
    ResourceNvidiaGPU = "nvidia.com/gpu"
    ResourceMIGPrefix = "nvidia.com/mig-"

    // 注解键
    AnnotationGPUModel       = "gpu.scheduler/model"
    AnnotationGPUMemory      = "gpu.scheduler/memory"
    AnnotationGPUTopology    = "gpu.scheduler/topology"
    AnnotationMIGProfile     = "gpu.scheduler/mig-profile"
    AnnotationPreferNVLink   = "gpu.scheduler/prefer-nvlink"
    AnnotationPreferSameNUMA = "gpu.scheduler/prefer-same-numa"

    // 状态键
    StateKeyGPURequest = "GPURequest"
)

// GPUSchedulerPlugin GPU 调度插件
type GPUSchedulerPlugin struct {
    handle       framework.Handle
    stateManager *GPUStateManager
}

// GPURequest GPU 请求信息
type GPURequest struct {
    // 请求的 GPU 数量
    Count int64
    // 指定的 GPU 型号（可选）
    Model string
    // 最小显存需求（MB）
    MinMemory int64
    // MIG profile（可选）
    MIGProfile string
    // 拓扑约束
    TopologyConstraint string
    // 是否偏好 NVLINK
    PreferNVLink bool
    // 是否偏好同一 NUMA
    PreferSameNUMA bool
}

// Name 返回插件名称
func (g *GPUSchedulerPlugin) Name() string {
    return PluginName
}

// New 创建插件实例
func New(obj runtime.Object, handle framework.Handle) (framework.Plugin, error) {
    return &GPUSchedulerPlugin{
        handle:       handle,
        stateManager: NewGPUStateManager(),
    }, nil
}

// PreFilter 预过滤阶段
func (g *GPUSchedulerPlugin) PreFilter(ctx context.Context, state *framework.CycleState,
    pod *v1.Pod) (*framework.PreFilterResult, *framework.Status) {

    // 解析 GPU 请求
    gpuRequest, err := g.parseGPURequest(pod)
    if err != nil {
        return nil, framework.NewStatus(framework.UnschedulableAndUnresolvable,
            fmt.Sprintf("failed to parse GPU request: %v", err))
    }

    // 如果不需要 GPU，跳过
    if gpuRequest.Count == 0 {
        return nil, framework.NewStatus(framework.Success)
    }

    // 验证请求的合法性
    if err := g.validateGPURequest(gpuRequest); err != nil {
        return nil, framework.NewStatus(framework.UnschedulableAndUnresolvable,
            fmt.Sprintf("invalid GPU request: %v", err))
    }

    // 将请求信息存入调度周期状态
    state.Write(StateKeyGPURequest, gpuRequest)

    return nil, framework.NewStatus(framework.Success)
}

// PreFilterExtensions 返回扩展接口
func (g *GPUSchedulerPlugin) PreFilterExtensions() framework.PreFilterExtensions {
    return nil
}

// parseGPURequest 解析 Pod 的 GPU 请求
func (g *GPUSchedulerPlugin) parseGPURequest(pod *v1.Pod) (*GPURequest, error) {
    request := &GPURequest{}

    // 从容器资源请求中获取 GPU 数量
    for _, container := range pod.Spec.Containers {
        if container.Resources.Requests != nil {
            // 检查标准 GPU 资源
            if gpuQty, ok := container.Resources.Requests[ResourceNvidiaGPU]; ok {
                request.Count += gpuQty.Value()
            }

            // 检查 MIG 资源
            for resName, qty := range container.Resources.Requests {
                if strings.HasPrefix(string(resName), ResourceMIGPrefix) {
                    request.Count += qty.Value()
                    // 提取 MIG profile
                    request.MIGProfile = strings.TrimPrefix(string(resName), ResourceMIGPrefix)
                }
            }
        }
    }

    // 从注解中获取附加需求
    if pod.Annotations != nil {
        if model, ok := pod.Annotations[AnnotationGPUModel]; ok {
            request.Model = model
        }

        if memory, ok := pod.Annotations[AnnotationGPUMemory]; ok {
            memMB, err := strconv.ParseInt(memory, 10, 64)
            if err != nil {
                return nil, fmt.Errorf("invalid memory annotation: %v", err)
            }
            request.MinMemory = memMB
        }

        if topology, ok := pod.Annotations[AnnotationGPUTopology]; ok {
            request.TopologyConstraint = topology
        }

        if mig, ok := pod.Annotations[AnnotationMIGProfile]; ok {
            request.MIGProfile = mig
        }

        if prefer, ok := pod.Annotations[AnnotationPreferNVLink]; ok {
            request.PreferNVLink = prefer == "true"
        }

        if prefer, ok := pod.Annotations[AnnotationPreferSameNUMA]; ok {
            request.PreferSameNUMA = prefer == "true"
        }
    }

    return request, nil
}

// validateGPURequest 验证 GPU 请求
func (g *GPUSchedulerPlugin) validateGPURequest(request *GPURequest) error {
    // 验证 GPU 数量
    if request.Count < 0 {
        return fmt.Errorf("negative GPU count: %d", request.Count)
    }

    if request.Count > 8 {
        // 通常单节点最多 8 GPU
        return fmt.Errorf("GPU count %d exceeds maximum", request.Count)
    }

    // 验证 MIG profile
    if request.MIGProfile != "" {
        validProfiles := []string{
            "1g.5gb", "1g.10gb", "1g.20gb",
            "2g.10gb", "2g.20gb",
            "3g.20gb", "3g.40gb",
            "4g.40gb",
            "7g.40gb", "7g.80gb",
        }
        valid := false
        for _, p := range validProfiles {
            if request.MIGProfile == p {
                valid = true
                break
            }
        }
        if !valid {
            return fmt.Errorf("invalid MIG profile: %s", request.MIGProfile)
        }
    }

    // 验证拓扑约束
    if request.TopologyConstraint != "" {
        validConstraints := []string{"single", "nvlink", "same-numa", "any"}
        valid := false
        for _, c := range validConstraints {
            if request.TopologyConstraint == c {
                valid = true
                break
            }
        }
        if !valid {
            return fmt.Errorf("invalid topology constraint: %s", request.TopologyConstraint)
        }
    }

    return nil
}

// Filter 过滤不满足 GPU 需求的节点
func (g *GPUSchedulerPlugin) Filter(ctx context.Context, state *framework.CycleState,
    pod *v1.Pod, nodeInfo *framework.NodeInfo) *framework.Status {

    // 获取 GPU 请求
    gpuRequestData, err := state.Read(StateKeyGPURequest)
    if err != nil {
        // 没有 GPU 请求，跳过
        return framework.NewStatus(framework.Success)
    }

    gpuRequest := gpuRequestData.(*GPURequest)
    if gpuRequest.Count == 0 {
        return framework.NewStatus(framework.Success)
    }

    nodeName := nodeInfo.Node().Name

    // 获取节点 GPU 信息
    nodeGPUInfo, ok := g.stateManager.GetNodeGPUInfo(nodeName)
    if !ok {
        return framework.NewStatus(framework.Unschedulable,
            fmt.Sprintf("node %s has no GPU info", nodeName))
    }

    // 检查 GPU 数量
    availableGPUs := g.stateManager.GetAvailableGPUs(nodeName)
    if int64(len(availableGPUs)) < gpuRequest.Count {
        return framework.NewStatus(framework.Unschedulable,
            fmt.Sprintf("insufficient GPUs: need %d, available %d",
                gpuRequest.Count, len(availableGPUs)))
    }

    // 检查 GPU 型号
    if gpuRequest.Model != "" {
        matchingGPUs := 0
        for _, gpu := range availableGPUs {
            if gpu.Model == gpuRequest.Model {
                matchingGPUs++
            }
        }
        if int64(matchingGPUs) < gpuRequest.Count {
            return framework.NewStatus(framework.Unschedulable,
                fmt.Sprintf("insufficient GPUs with model %s: need %d, available %d",
                    gpuRequest.Model, gpuRequest.Count, matchingGPUs))
        }
    }

    // 检查显存需求
    if gpuRequest.MinMemory > 0 {
        matchingGPUs := 0
        for _, gpu := range availableGPUs {
            if gpu.MemoryTotal >= gpuRequest.MinMemory {
                matchingGPUs++
            }
        }
        if int64(matchingGPUs) < gpuRequest.Count {
            return framework.NewStatus(framework.Unschedulable,
                fmt.Sprintf("insufficient GPUs with memory >= %dMB: need %d, available %d",
                    gpuRequest.MinMemory, gpuRequest.Count, matchingGPUs))
        }
    }

    // 检查 MIG 配置
    if gpuRequest.MIGProfile != "" {
        if !g.checkMIGAvailability(nodeGPUInfo, gpuRequest) {
            return framework.NewStatus(framework.Unschedulable,
                fmt.Sprintf("no available MIG profile %s", gpuRequest.MIGProfile))
        }
    }

    // 检查拓扑约束
    if gpuRequest.TopologyConstraint != "" {
        if !g.checkTopologyConstraint(nodeGPUInfo, gpuRequest) {
            return framework.NewStatus(framework.Unschedulable,
                fmt.Sprintf("topology constraint %s cannot be satisfied",
                    gpuRequest.TopologyConstraint))
        }
    }

    return framework.NewStatus(framework.Success)
}

// checkMIGAvailability 检查 MIG 可用性
func (g *GPUSchedulerPlugin) checkMIGAvailability(info *NodeGPUInfo,
    request *GPURequest) bool {

    availableCount := int64(0)

    for _, gpu := range info.GPUs {
        if !gpu.MIGEnabled {
            continue
        }
        for _, mig := range gpu.MIGDevices {
            if mig.Profile == request.MIGProfile && mig.State == DeviceStateFree {
                availableCount++
            }
        }
    }

    return availableCount >= request.Count
}

// checkTopologyConstraint 检查拓扑约束
func (g *GPUSchedulerPlugin) checkTopologyConstraint(info *NodeGPUInfo,
    request *GPURequest) bool {

    availableGPUs := make([]int, 0)
    for i, gpu := range info.GPUs {
        if gpu.State == DeviceStateFree && gpu.Health == DeviceHealthy {
            availableGPUs = append(availableGPUs, i)
        }
    }

    if int64(len(availableGPUs)) < request.Count {
        return false
    }

    switch request.TopologyConstraint {
    case "single":
        // 单 GPU，任何节点都满足
        return request.Count == 1

    case "nvlink":
        // 需要 NVLINK 连接的 GPU
        return g.findNVLinkConnectedGPUs(info.Topology, availableGPUs, int(request.Count)) != nil

    case "same-numa":
        // 需要同一 NUMA 节点的 GPU
        return g.findSameNUMAGPUs(info.Topology, availableGPUs, int(request.Count)) != nil

    case "any":
        return true

    default:
        return true
    }
}

// findNVLinkConnectedGPUs 找到 NVLINK 连接的 GPU 组合
func (g *GPUSchedulerPlugin) findNVLinkConnectedGPUs(topo *GPUTopology,
    available []int, count int) []int {

    if topo == nil || len(topo.Links) == 0 {
        return nil
    }

    // 构建 NVLINK 连接图
    nvlinkGraph := make(map[int][]int)
    for _, i := range available {
        for _, j := range available {
            if i != j && topo.Links[i][j] == LinkTypeNVLink {
                nvlinkGraph[i] = append(nvlinkGraph[i], j)
            }
        }
    }

    // 使用 DFS 找到连通分量
    visited := make(map[int]bool)
    var result []int

    var dfs func(node int, group []int) []int
    dfs = func(node int, group []int) []int {
        if len(group) >= count {
            return group
        }

        visited[node] = true
        group = append(group, node)

        for _, neighbor := range nvlinkGraph[node] {
            if !visited[neighbor] {
                if found := dfs(neighbor, group); len(found) >= count {
                    return found
                }
            }
        }

        return group
    }

    for _, gpu := range available {
        if !visited[gpu] {
            visited = make(map[int]bool)
            if found := dfs(gpu, nil); len(found) >= count {
                return found[:count]
            }
        }
    }

    return nil
}

// findSameNUMAGPUs 找到同一 NUMA 节点的 GPU
func (g *GPUSchedulerPlugin) findSameNUMAGPUs(topo *GPUTopology,
    available []int, count int) []int {

    if topo == nil || len(topo.NUMANodes) == 0 {
        return nil
    }

    // 按 NUMA 节点分组
    numaGroups := make(map[int][]int)
    for _, gpu := range available {
        numa := topo.NUMANodes[gpu]
        numaGroups[numa] = append(numaGroups[numa], gpu)
    }

    // 找到满足数量要求的 NUMA 组
    for _, gpus := range numaGroups {
        if len(gpus) >= count {
            return gpus[:count]
        }
    }

    return nil
}

// Score 为节点打分
func (g *GPUSchedulerPlugin) Score(ctx context.Context, state *framework.CycleState,
    pod *v1.Pod, nodeName string) (int64, *framework.Status) {

    // 获取 GPU 请求
    gpuRequestData, err := state.Read(StateKeyGPURequest)
    if err != nil {
        return 0, framework.NewStatus(framework.Success)
    }

    gpuRequest := gpuRequestData.(*GPURequest)
    if gpuRequest.Count == 0 {
        return 0, framework.NewStatus(framework.Success)
    }

    nodeGPUInfo, ok := g.stateManager.GetNodeGPUInfo(nodeName)
    if !ok {
        return 0, framework.NewStatus(framework.Success)
    }

    score := int64(0)

    // 1. 拓扑得分（0-40分）
    topoScore := g.calculateTopologyScore(nodeGPUInfo, gpuRequest)
    score += topoScore

    // 2. 碎片化得分（0-30分）
    fragScore := g.calculateFragmentationScore(nodeGPUInfo, gpuRequest)
    score += fragScore

    // 3. 负载均衡得分（0-20分）
    balanceScore := g.calculateBalanceScore(nodeGPUInfo)
    score += balanceScore

    // 4. 型号匹配得分（0-10分）
    modelScore := g.calculateModelScore(nodeGPUInfo, gpuRequest)
    score += modelScore

    return score, framework.NewStatus(framework.Success)
}

// ScoreExtensions 返回得分扩展
func (g *GPUSchedulerPlugin) ScoreExtensions() framework.ScoreExtensions {
    return g
}

// NormalizeScore 归一化得分到 [0, 100]
func (g *GPUSchedulerPlugin) NormalizeScore(ctx context.Context, state *framework.CycleState,
    pod *v1.Pod, scores framework.NodeScoreList) *framework.Status {

    // 找到最大分数
    var maxScore int64 = 0
    for _, score := range scores {
        if score.Score > maxScore {
            maxScore = score.Score
        }
    }

    // 归一化
    if maxScore > 0 {
        for i := range scores {
            scores[i].Score = scores[i].Score * framework.MaxNodeScore / maxScore
        }
    }

    return framework.NewStatus(framework.Success)
}

// calculateTopologyScore 计算拓扑得分
func (g *GPUSchedulerPlugin) calculateTopologyScore(info *NodeGPUInfo,
    request *GPURequest) int64 {

    if info.Topology == nil || request.Count <= 1 {
        return 20 // 单 GPU 或无拓扑信息，给中等分数
    }

    availableGPUs := make([]int, 0)
    for i, gpu := range info.GPUs {
        if gpu.State == DeviceStateFree && gpu.Health == DeviceHealthy {
            availableGPUs = append(availableGPUs, i)
        }
    }

    score := int64(0)

    // 检查 NVLINK 连接
    if request.PreferNVLink || request.TopologyConstraint == "nvlink" {
        nvlinkGPUs := g.findNVLinkConnectedGPUs(info.Topology, availableGPUs, int(request.Count))
        if nvlinkGPUs != nil {
            score += 40 // 满分
        } else {
            score += 10 // 无法满足 NVLINK，给低分
        }
    } else {
        // 即使没有要求，有 NVLINK 也加分
        nvlinkGPUs := g.findNVLinkConnectedGPUs(info.Topology, availableGPUs, int(request.Count))
        if nvlinkGPUs != nil {
            score += 30
        } else {
            score += 20
        }
    }

    // 检查 NUMA 亲和性
    if request.PreferSameNUMA || request.TopologyConstraint == "same-numa" {
        sameNUMAGPUs := g.findSameNUMAGPUs(info.Topology, availableGPUs, int(request.Count))
        if sameNUMAGPUs != nil {
            score += 10 // 额外加分
        }
    }

    return score
}

// calculateFragmentationScore 计算碎片化得分
// 目标：避免留下难以利用的 GPU 碎片
func (g *GPUSchedulerPlugin) calculateFragmentationScore(info *NodeGPUInfo,
    request *GPURequest) int64 {

    // 统计可用 GPU 数量
    availableCount := int64(0)
    for _, gpu := range info.GPUs {
        if gpu.State == DeviceStateFree && gpu.Health == DeviceHealthy {
            availableCount++
        }
    }

    totalGPUs := int64(len(info.GPUs))
    remaining := availableCount - request.Count

    // 策略：
    // 1. 优先选择能完全使用的节点（remaining = 0）
    // 2. 次选剩余 GPU 较多的节点（避免碎片）
    // 3. 避免留下 1 个 GPU 的情况（很难利用）

    if remaining == 0 {
        return 30 // 完美匹配
    }

    if remaining == 1 {
        return 5 // 留下 1 个 GPU，不太好
    }

    if remaining >= totalGPUs/2 {
        return 25 // 还有一半以上可用
    }

    return 15 // 中等
}

// calculateBalanceScore 计算负载均衡得分
func (g *GPUSchedulerPlugin) calculateBalanceScore(info *NodeGPUInfo) int64 {
    usedCount := 0
    totalCount := len(info.GPUs)

    for _, gpu := range info.GPUs {
        if gpu.State != DeviceStateFree {
            usedCount++
        }
    }

    // 计算使用率
    utilization := float64(usedCount) / float64(totalCount)

    // 偏好使用率较高的节点（但不是满载）
    // 这样可以让部分节点完全空闲，便于调度大任务
    if utilization >= 0.8 {
        return 10 // 接近满载，降低优先级
    }
    if utilization >= 0.5 {
        return 20 // 中等使用率，最优
    }
    if utilization >= 0.2 {
        return 15 // 使用率较低
    }
    return 10 // 几乎空闲
}

// calculateModelScore 计算 GPU 型号匹配得分
func (g *GPUSchedulerPlugin) calculateModelScore(info *NodeGPUInfo,
    request *GPURequest) int64 {

    if request.Model == "" {
        return 5 // 没有型号要求，给基础分
    }

    matchCount := 0
    for _, gpu := range info.GPUs {
        if gpu.Model == request.Model && gpu.State == DeviceStateFree {
            matchCount++
        }
    }

    if int64(matchCount) >= request.Count {
        return 10 // 满足型号要求
    }

    return 0
}

// Reserve 预留资源
func (g *GPUSchedulerPlugin) Reserve(ctx context.Context, state *framework.CycleState,
    pod *v1.Pod, nodeName string) *framework.Status {

    gpuRequestData, err := state.Read(StateKeyGPURequest)
    if err != nil {
        return framework.NewStatus(framework.Success)
    }

    gpuRequest := gpuRequestData.(*GPURequest)
    if gpuRequest.Count == 0 {
        return framework.NewStatus(framework.Success)
    }

    // 选择要分配的 GPU
    selectedGPUs, err := g.selectGPUs(nodeName, gpuRequest)
    if err != nil {
        return framework.NewStatus(framework.Unschedulable,
            fmt.Sprintf("failed to select GPUs: %v", err))
    }

    // 预留 GPU
    podName := fmt.Sprintf("%s/%s", pod.Namespace, pod.Name)
    if err := g.stateManager.ReserveGPUs(nodeName, selectedGPUs, podName); err != nil {
        return framework.NewStatus(framework.Unschedulable,
            fmt.Sprintf("failed to reserve GPUs: %v", err))
    }

    // 将选择的 GPU 存入状态
    state.Write("SelectedGPUs", selectedGPUs)

    return framework.NewStatus(framework.Success)
}

// Unreserve 取消预留
func (g *GPUSchedulerPlugin) Unreserve(ctx context.Context, state *framework.CycleState,
    pod *v1.Pod, nodeName string) {

    selectedGPUsData, err := state.Read("SelectedGPUs")
    if err != nil {
        return
    }

    selectedGPUs := selectedGPUsData.([]int)
    g.stateManager.UnreserveGPUs(nodeName, selectedGPUs)
}

// selectGPUs 选择要分配的 GPU
func (g *GPUSchedulerPlugin) selectGPUs(nodeName string, request *GPURequest) ([]int, error) {
    info, ok := g.stateManager.GetNodeGPUInfo(nodeName)
    if !ok {
        return nil, fmt.Errorf("node %s not found", nodeName)
    }

    availableGPUs := make([]int, 0)
    for i, gpu := range info.GPUs {
        if gpu.State == DeviceStateFree && gpu.Health == DeviceHealthy {
            // 检查型号
            if request.Model != "" && gpu.Model != request.Model {
                continue
            }
            // 检查显存
            if request.MinMemory > 0 && gpu.MemoryTotal < request.MinMemory {
                continue
            }
            availableGPUs = append(availableGPUs, i)
        }
    }

    count := int(request.Count)

    // 根据拓扑约束选择
    switch request.TopologyConstraint {
    case "nvlink":
        selected := g.findNVLinkConnectedGPUs(info.Topology, availableGPUs, count)
        if selected != nil {
            return selected, nil
        }
        return nil, fmt.Errorf("cannot find NVLINK connected GPUs")

    case "same-numa":
        selected := g.findSameNUMAGPUs(info.Topology, availableGPUs, count)
        if selected != nil {
            return selected, nil
        }
        return nil, fmt.Errorf("cannot find GPUs in same NUMA")

    default:
        // 优先选择 NVLINK 连接的 GPU
        if request.PreferNVLink && info.Topology != nil {
            selected := g.findNVLinkConnectedGPUs(info.Topology, availableGPUs, count)
            if selected != nil {
                return selected, nil
            }
        }

        // 优先选择同一 NUMA 的 GPU
        if request.PreferSameNUMA && info.Topology != nil {
            selected := g.findSameNUMAGPUs(info.Topology, availableGPUs, count)
            if selected != nil {
                return selected, nil
            }
        }

        // 默认：选择前 N 个
        if len(availableGPUs) >= count {
            return availableGPUs[:count], nil
        }

        return nil, fmt.Errorf("insufficient GPUs")
    }
}

// cmd/gpu-collector/main.go
package main

import (
    "context"
    "encoding/json"
    "fmt"
    "os"
    "time"

    "github.com/NVIDIA/go-nvml/pkg/nvml"
    metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
    "k8s.io/client-go/kubernetes"
    "k8s.io/client-go/rest"
)

type GPUCollector struct {
    nodeName   string
    kubeClient kubernetes.Interface
    interval   time.Duration
}

func NewGPUCollector() (*GPUCollector, error) {
    // 获取节点名
    nodeName := os.Getenv("NODE_NAME")
    if nodeName == "" {
        return nil, fmt.Errorf("NODE_NAME env not set")
    }

    // 创建 k8s 客户端
    config, err := rest.InClusterConfig()
    if err != nil {
        return nil, fmt.Errorf("failed to get k8s config: %v", err)
    }

    client, err := kubernetes.NewForConfig(config)
    if err != nil {
        return nil, fmt.Errorf("failed to create k8s client: %v", err)
    }

    return &GPUCollector{
        nodeName:   nodeName,
        kubeClient: client,
        interval:   30 * time.Second,
    }, nil
}

func (c *GPUCollector) Run(ctx context.Context) error {
    // 初始化 NVML
    ret := nvml.Init()
    if ret != nvml.SUCCESS {
        return fmt.Errorf("failed to initialize NVML: %v", nvml.ErrorString(ret))
    }
    defer nvml.Shutdown()

    ticker := time.NewTicker(c.interval)
    defer ticker.Stop()

    for {
        select {
        case <-ctx.Done():
            return ctx.Err()
        case <-ticker.C:
            if err := c.collectAndUpdate(); err != nil {
                fmt.Printf("Failed to collect GPU info: %v\n", err)
            }
        }
    }
}

func (c *GPUCollector) collectAndUpdate() error {
    info, err := c.collectGPUInfo()
    if err != nil {
        return err
    }

    return c.updateNodeGPUInfo(info)
}

func (c *GPUCollector) collectGPUInfo() (*NodeGPUInfo, error) {
    // 获取驱动版本
    driverVersion, ret := nvml.SystemGetDriverVersion()
    if ret != nvml.SUCCESS {
        return nil, fmt.Errorf("failed to get driver version: %v", nvml.ErrorString(ret))
    }

    // 获取 CUDA 版本
    cudaVersion, ret := nvml.SystemGetCudaDriverVersion_v2()
    if ret != nvml.SUCCESS {
        return nil, fmt.Errorf("failed to get CUDA version: %v", nvml.ErrorString(ret))
    }

    // 获取 GPU 数量
    count, ret := nvml.DeviceGetCount()
    if ret != nvml.SUCCESS {
        return nil, fmt.Errorf("failed to get device count: %v", nvml.ErrorString(ret))
    }

    info := &NodeGPUInfo{
        NodeName:      c.nodeName,
        DriverVersion: driverVersion,
        CUDAVersion:   fmt.Sprintf("%d.%d", cudaVersion/1000, (cudaVersion%1000)/10),
        GPUs:          make([]GPUDevice, count),
    }

    // 收集每个 GPU 的信息
    for i := 0; i < count; i++ {
        device, ret := nvml.DeviceGetHandleByIndex(i)
        if ret != nvml.SUCCESS {
            continue
        }

        info.GPUs[i] = c.collectDeviceInfo(device, i)
    }

    // 收集拓扑信息
    info.Topology = c.collectTopology(count)

    return info, nil
}

func (c *GPUCollector) collectDeviceInfo(device nvml.Device, index int) GPUDevice {
    gpu := GPUDevice{Index: index}

    // UUID
    uuid, ret := device.GetUUID()
    if ret == nvml.SUCCESS {
        gpu.UUID = uuid
    }

    // 型号
    name, ret := device.GetName()
    if ret == nvml.SUCCESS {
        gpu.Model = name
    }

    // 显存
    memory, ret := device.GetMemoryInfo()
    if ret == nvml.SUCCESS {
        gpu.MemoryTotal = int64(memory.Total / 1024 / 1024) // MB
        gpu.MemoryFree = int64(memory.Free / 1024 / 1024)
    }

    // 计算能力
    major, minor, ret := device.GetCudaComputeCapability()
    if ret == nvml.SUCCESS {
        gpu.ComputeCapability = fmt.Sprintf("%d.%d", major, minor)
    }

    // MIG 状态
    mode, _, ret := device.GetMigMode()
    if ret == nvml.SUCCESS {
        gpu.MIGEnabled = mode == nvml.DEVICE_MIG_ENABLE
    }

    // 如果 MIG 启用，收集 MIG 设备
    if gpu.MIGEnabled {
        gpu.MIGDevices = c.collectMIGDevices(device)
    }

    // PCIe 信息
    pciInfo, ret := device.GetPciInfo()
    if ret == nvml.SUCCESS {
        gpu.PCIeInfo = PCIeInfo{
            BusID: fmt.Sprintf("%04x:%02x:%02x.0",
                pciInfo.Domain, pciInfo.Bus, pciInfo.Device),
        }
    }

    // NUMA 节点
    // 注意：需要通过系统文件获取
    numaNode, err := c.getNUMANode(gpu.PCIeInfo.BusID)
    if err == nil {
        gpu.PCIeInfo.NUMANode = numaNode
    }

    // 健康状态
    gpu.Health = DeviceHealthy

    return gpu
}

func (c *GPUCollector) collectMIGDevices(device nvml.Device) []MIGDevice {
    var migDevices []MIGDevice

    // 获取最大 GPU 实例数
    maxGI, ret := device.GetMaxMigDeviceCount()
    if ret != nvml.SUCCESS || maxGI == 0 {
        return migDevices
    }

    // 枚举 MIG 设备
    for i := 0; i < maxGI; i++ {
        migDevice, ret := device.GetMigDeviceHandleByIndex(i)
        if ret != nvml.SUCCESS {
            continue
        }

        mig := MIGDevice{
            GIIndex: i,
        }

        uuid, ret := migDevice.GetUUID()
        if ret == nvml.SUCCESS {
            mig.UUID = uuid
        }

        // 获取 profile 信息
        attr, ret := migDevice.GetAttributes()
        if ret == nvml.SUCCESS {
            mig.Profile = fmt.Sprintf("%dg.%dgb",
                attr.GpuInstanceSliceCount,
                attr.MemorySizeMB/1024)
        }

        mig.State = DeviceStateFree
        migDevices = append(migDevices, mig)
    }

    return migDevices
}

func (c *GPUCollector) collectTopology(gpuCount int) *GPUTopology {
    topo := &GPUTopology{
        NodeName:  c.nodeName,
        Links:     make([][]LinkType, gpuCount),
        NUMANodes: make([]int, gpuCount),
    }

    for i := 0; i < gpuCount; i++ {
        topo.Links[i] = make([]LinkType, gpuCount)

        device1, ret := nvml.DeviceGetHandleByIndex(i)
        if ret != nvml.SUCCESS {
            continue
        }

        for j := 0; j < gpuCount; j++ {
            if i == j {
                topo.Links[i][j] = LinkTypeSameBoard
                continue
            }

            device2, ret := nvml.DeviceGetHandleByIndex(j)
            if ret != nvml.SUCCESS {
                topo.Links[i][j] = LinkTypeNone
                continue
            }

            // 获取 GPU 间的拓扑关系
            pathInfo, ret := nvml.DeviceGetTopologyCommonAncestor(device1, device2)
            if ret != nvml.SUCCESS {
                topo.Links[i][j] = LinkTypeNone
                continue
            }

            switch pathInfo {
            case nvml.TOPOLOGY_INTERNAL:
                topo.Links[i][j] = LinkTypeSameBoard
            case nvml.TOPOLOGY_SINGLE:
                // 检查是否有 NVLINK
                nvlinkCap, ret := device1.GetNvLinkCapability(uint32(j), nvml.NVLINK_CAP_P2P_SUPPORTED)
                if ret == nvml.SUCCESS && nvlinkCap > 0 {
                    topo.Links[i][j] = LinkTypeNVLink
                } else {
                    topo.Links[i][j] = LinkTypePCIe
                }
            case nvml.TOPOLOGY_MULTIPLE:
                topo.Links[i][j] = LinkTypePCIe
            case nvml.TOPOLOGY_HOSTBRIDGE:
                topo.Links[i][j] = LinkTypePCIe
            case nvml.TOPOLOGY_NODE:
                topo.Links[i][j] = LinkTypePCIe
            case nvml.TOPOLOGY_SYSTEM:
                topo.Links[i][j] = LinkTypePCIe
            default:
                topo.Links[i][j] = LinkTypeNone
            }
        }
    }

    return topo
}

func (c *GPUCollector) getNUMANode(busID string) (int, error) {
    // 读取 /sys/bus/pci/devices/{busID}/numa_node
    path := fmt.Sprintf("/sys/bus/pci/devices/%s/numa_node", busID)
    data, err := os.ReadFile(path)
    if err != nil {
        return 0, err
    }

    var numa int
    fmt.Sscanf(string(data), "%d", &numa)
    return numa, nil
}

func (c *GPUCollector) updateNodeGPUInfo(info *NodeGPUInfo) error {
    // 这里应该更新 CRD，简化示例使用 ConfigMap
    data, err := json.Marshal(info)
    if err != nil {
        return err
    }

    cm := &corev1.ConfigMap{
        ObjectMeta: metav1.ObjectMeta{
            Name:      fmt.Sprintf("gpu-info-%s", c.nodeName),
            Namespace: "kube-system",
            Labels: map[string]string{
                "app":  "gpu-collector",
                "node": c.nodeName,
            },
        },
        Data: map[string]string{
            "gpu-info": string(data),
        },
    }

    _, err = c.kubeClient.CoreV1().ConfigMaps("kube-system").Update(
        context.Background(), cm, metav1.UpdateOptions{})
    if err != nil {
        // 如果不存在则创建
        _, err = c.kubeClient.CoreV1().ConfigMaps("kube-system").Create(
            context.Background(), cm, metav1.CreateOptions{})
    }

    return err
}

func main() {
    collector, err := NewGPUCollector()
    if err != nil {
        fmt.Printf("Failed to create collector: %v\n", err)
        os.Exit(1)
    }

    ctx := context.Background()
    if err := collector.Run(ctx); err != nil {
        fmt.Printf("Collector error: %v\n", err)
        os.Exit(1)
    }
}