NVIDIA GPU 设备实现深度解析

// pkg/device/nvidia/device.go
type NvidiaGPUDevices struct {
    config         NvidiaConfig
    ReportedGPUNum map[string]int64  // key: nodeName, value: 已上报 GPU 数量
    mu             sync.Mutex        // 保护 ReportedGPUNum 的并发访问
}

type NvidiaConfig struct {
    NodeDefaultConfig            `yaml:",inline"`
    ResourceCountName            string   // "nvidia.com/gpu"
    ResourceMemoryName           string   // "nvidia.com/gpumem"
    ResourceCoreName             string   // "nvidia.com/gpucores"
    ResourceMemoryPercentageName string   // "nvidia.com/gpumem-percentage"
    ResourcePriority             string   // "nvidia.com/priority"
    OverwriteEnv                 bool     // 是否覆盖 NVIDIA_VISIBLE_DEVICES
    DefaultMemory                int32    // 默认显存值 (MB)
    DefaultCores                 int32    // 默认算力百分比
    DefaultGPUNum                int32    // 默认 GPU 数量
    MemoryFactor                 int32    // 显存乘数因子
    DisableCoreLimit             bool     // 禁用算力限制
    MigGeometriesList            []AllowedMigGeometries  // MIG 几何模板
    GPUCorePolicy                GPUCoreUtilizationPolicy // 算力隔离策略
    RuntimeClassName             string   // Pod 运行时类名
}

type GPUCoreUtilizationPolicy string

const (
    DefaultCorePolicy GPUCoreUtilizationPolicy = "default"   // 默认行为
    ForceCorePolicy   GPUCoreUtilizationPolicy = "force"     // 强制算力隔离
    DisableCorePolicy GPUCoreUtilizationPolicy = "disable"   // 禁用算力隔离
)

// MIG 模板定义
type MigTemplate struct {
    Name   string  // 模板名称，如 "1g.5gb"
    Core   int32   // 算力份额
    Memory int32   // 显存大小 (MB)
    Count  int32   // 该模板在 Geometry 中的数量
}

// 一种 Geometry 是 MigTemplate 的组合
type Geometry []MigTemplate

// MIG 实例的运行时使用状态
type MigTemplateUsage struct {
    Name   string  // 模板名称
    Core   int32   // 算力份额
    Memory int32   // 显存大小
    InUse  bool    // 是否已被分配
}

// 设备的 MIG 使用状态
type MigInUse struct {
    Index     int32  // 当前使用的 Geometry 索引
    UsageList MIGS   // MIG 实例使用列表
}

knownMigGeometries:
  - models: ["A100", "A800"]
    allowedGeometries:
      - # Geometry 0: 7 个 1g.5gb 实例
        - name: "1g.5gb"
          core: 14
          memory: 5120
          count: 7
      - # Geometry 1: 3 个 2g.10gb + 1 个 1g.5gb
        - name: "2g.10gb"
          core: 28
          memory: 10240
          count: 3
        - name: "1g.5gb"
          core: 14
          memory: 5120
          count: 1
      - # Geometry 2: 2 个 3g.20gb + 1 个 1g.5gb
        - name: "3g.20gb"
          core: 42
          memory: 20480
          count: 2
        - name: "1g.5gb"
          core: 14
          memory: 5120
          count: 1

for _, val := range nodedevices {
    if val.Mode == MigMode {
        val.MIGTemplate = make([]device.Geometry, 0)
        for _, migTemplates := range dev.config.MigGeometriesList {
            found := false
            for _, migDevices := range migTemplates.Models {
                if strings.Contains(val.Type, migDevices) {
                    found = true
                    break
                }
            }
            if found {
                val.MIGTemplate = append(val.MIGTemplate, migTemplates.Geometries...)
                break
            }
        }
    }
}

func ExtractMigTemplatesFromUUID(uuid string) (templateIdx int, pos int, error)

type FilterDevice struct {
    UUID  []string  // 按 UUID 过滤
    Index []uint    // 按设备索引过滤
}

func assertNuma(annos map[string]string) bool {
    numabind, ok := annos[NumaBind]
    if ok {
        enforce, err := strconv.ParseBool(numabind)
        if err == nil && enforce {
            return true
        }
    }
    return false
}

if numa && prevnuma != dev.Numa {
    if k.Nums != originReq {
        // 记录 NUMA 不匹配原因
    }
    k.Nums = originReq       // 重置请求数量
    prevnuma = dev.Numa       // 更新当前 NUMA 节点
    tmpDevs = make(map[string]device.ContainerDevices)  // 清空已选设备
}

// pkg/device/nvidia/calculate_score.go
func calculateGPUPairScore(gpu0 *Device, gpu1 *Device) int {
    score := 0
    for _, link := range gpu0.Links[gpu1.Index] {
        switch link.Type {
        case P2PLinkCrossCPU:      score += 10
        case P2PLinkSameCPU:       score += 20
        case P2PLinkHostBridge:    score += 30
        case P2PLinkMultiSwitch:   score += 40
        case P2PLinkSingleSwitch:  score += 50
        case P2PLinkSameBoard:     score += 60
        case SingleNVLINKLink:     score += 100
        case TwoNVLINKLinks:       score += 200
        // ... 最高到 EighteenNVLINKLinks: 1800
        }
    }
    return score
}

func fitQuota(tmpDevs map[string]device.ContainerDevices, allocated *device.PodDevices,
              ns string, memreq int64, coresreq int64) bool {
    // 1. 累加当前分配方案中已使用的显存和算力
    // 2. 累加之前已分配的容器使用量
    // 3. 调用 QuotaManager.FitQuota() 检查是否超限
}