DRA P3：DRA 工作流程与源码分析

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11

DRA Driver（DaemonSet）启动
    │
    ├─ 1. 调用 NVML 扫描节点上的 GPU 设备
    │     → 收集每个 GPU 的属性：型号、架构、显存、CUDA 版本、驱动版本、UUID...
    │
    ├─ 2. 注册到 Kubelet（通过 plugins_registry socket）
    │     → Kubelet 发现 DRA Plugin，建立 gRPC 连接
    │
    └─ 3. 发布 ResourceSlice 到 API Server
          → 构建 ResourceSlice 对象并发布，包含所有 GPU 设备及其属性
          → 调度器、Kubelet、用户都可以通过 kubectl 查看

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22

// cmd/gpu-kubelet-plugin/driver.go:70
func NewDriver(ctx context.Context, config *Config) (*driver, error) {
    // 1. 初始化 NVML，枚举所有设备
    state, err := NewDeviceState(ctx, config)
    // ...
    driver := &driver{
        state:  state,
        pulock: flock.NewFlock(filepath.Join(config.DriverPluginPath(), DriverPrepUprepFlockFileName)),
    }
    // 2. 注册 DRA Plugin，暴露 Prepare/Unprepare 接口
    helper, err := kubeletplugin.Start(ctx, driver,
        kubeletplugin.DriverName(DriverName),
        kubeletplugin.Serialize(false), // 禁用自带序列化，改用 flock 控制
        // ...
    )
    driver.pluginhelper = helper
    // 3. 发布初始 ResourceSlice 到 API Server
    if err := driver.publishResources(ctx, config); err != nil {
        return nil, err
    }
    return driver, nil
}

1
2
3
4
5
6
7
8

// cmd/gpu-kubelet-plugin/nvlib.go:174
func (l deviceLib) enumerateAllPossibleDevices() (*PerGPUAllocatableDevices, error) {
    perGPUAllocatable, err := l.GetPerGpuAllocatableDevices() // 遍历 NVML 获取所有物理 GPU
    if featuregates.Enabled(featuregates.PassthroughSupport) {
        err = l.enumerateGpuVfioDevices(perGPUAllocatable) // 额外枚举 VFIO PCI 设备
    }
    return perGPUAllocatable, nil
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20

// cmd/gpu-kubelet-plugin/driver.go:455
func (d *driver) publishResources(ctx context.Context, config *Config) error {
    if featuregates.Enabled(featuregates.DynamicMIG) {
        resources := d.GenerateDriverResources(config.flags.nodeName)
        return d.pluginhelper.PublishResources(ctx, resources)
    }
    // 普通模式：所有设备放进一个 Slice
    var resourceSlice resourceslice.Slice
    for _, devices := range d.state.perGPUAllocatable.allocatablesMap {
        for _, device := range devices {
            resourceSlice.Devices = append(resourceSlice.Devices, device.GetDevice())
        }
    }
    resources := resourceslice.DriverResources{
        Pools: map[string]resourceslice.Pool{
            config.flags.nodeName: {Slices: []resourceslice.Slice{resourceSlice}},
        },
    }
    return d.pluginhelper.PublishResources(ctx, resources)
}

1
2
3
4
5
6
7
8

apiVersion: resource.k8s.io/v1
kind: DeviceClass
metadata:
  name: gpu.nvidia.com
spec:
  selectors:
  - cel:
      expression: device.driver == 'gpu.nvidia.com' && device.attributes['gpu.nvidia.com'].type == 'gpu'

1
2
3
4
5
6
7

spec:
  selectors:
  - cel:
      expression: |
        device.driver == 'gpu.nvidia.com' &&
        device.attributes['gpu.nvidia.com'].type == 'gpu' &&
        device.attributes['gpu.nvidia.com'].productName == 'NVIDIA A100'

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29

apiVersion: resource.k8s.io/v1
kind: ResourceClaimTemplate
metadata:
  name: single-gpu
spec:
  spec:
    devices:
      requests:
      - name: gpu
        exactly:
          deviceClassName: gpu.nvidia.com
          allocationMode: ExactCount
          count: 1
---
apiVersion: v1
kind: Pod
metadata:
  name: gpu-test-pod
spec:
  containers:
  - name: cuda-container
    image: nvidia/cuda:12.1.0-base-ubuntu22.04
    command: ["nvidia-smi", "-L"]
    resources:
      claims:
      - name: gpu-claim
  resourceClaims:
  - name: gpu-claim
    resourceClaimTemplateName: single-gpu

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36

// pkg/scheduler/framework/plugins/dynamicresources/dynamicresources.go:444
func (pl *DynamicResources) PreFilter(ctx context.Context, state fwk.CycleState, pod *v1.Pod, nodes []fwk.NodeInfo) (*fwk.PreFilterResult, *fwk.Status) {
    // ...
    s := &stateData{}
    state.Write(stateKey, s)
    // 收集 Pod 引用的所有 ResourceClaim
    userClaims, err := pl.podResourceClaims(pod)
    // 处理扩展资源
    extendedResourceClaim, status := pl.preFilterExtendedResources(pod, logger, s)
    // 对已分配的 Claim，提取 NodeSelector
    for index, claim := range claims.all() {
        if claim.Status.Allocation != nil {
            if claim.Status.Allocation.NodeSelector != nil {
                nodeSelector, err := nodeaffinity.NewNodeSelector(claim.Status.Allocation.NodeSelector)
                s.informationsForClaim[index].availableOnNodes = nodeSelector
            }
        } else {
            numClaimsToAllocate++
            // 验证 DeviceClass 存在
            for _, request := range claim.Spec.Devices.Requests {
                // ...
            }
        }
    }
    // 收集全局已分配状态 + 列出所有 ResourceSlice + 创建结构化分配器
    if numClaimsToAllocate > 0 {
        // EnableDRAConsumableCapacity 启用时使用 GatherAllocatedState，否则使用 ListAllAllocatedDevices
        allocatedState, err := pl.draManager.ResourceClaims().GatherAllocatedState()
        slices, err := pl.draManager.ResourceSlices().ListWithDeviceTaintRules()
        features := AllocatorFeatures(pl.fts)
        allocator, err := structured.NewAllocator(ctx, features, *allocatedState, pl.draManager.DeviceClasses(), slices, pl.celCache)
        s.allocator = allocator
        s.nodeAllocations = make(map[string]nodeAllocation)
    }
    // ...
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26

// pkg/scheduler/framework/plugins/dynamicresources/dynamicresources.go:733
func (pl *DynamicResources) Filter(ctx context.Context, cs fwk.CycleState, pod *v1.Pod, nodeInfo fwk.NodeInfo) *fwk.Status {
    // ...
    // 构建待分配 Claim 列表
    claimsToAllocate := make([]*resourceapi.ResourceClaim, 0, state.claims.len())
    for index, claim := range state.claims.toAllocate() {
        if state.informationsForClaim[index].allocation != nil {
            pendingResult = append(pendingResult, *state.informationsForClaim[index].allocation)
            continue
        }
        claimsToAllocate = append(claimsToAllocate, claim)
    }
    // 调用分配器
    allocationResult, err := state.allocator.Allocate(allocCtx, node, claimsToAllocate)
    switch {
    case errors.Is(err, structured.ErrFailedAllocationOnNode):
        return statusUnschedulable(logger, err.Error(), ...)
    // ...
    }
    // 缓存分配结果
    state.nodeAllocations[node.Name] = nodeAllocation{
        allocationResults: allocations,
        // ...
    }
    // ...
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17

// pkg/scheduler/framework/plugins/dynamicresources/dynamicresources.go:1129
func (pl *DynamicResources) Reserve(ctx context.Context, cs fwk.CycleState, pod *v1.Pod, nodeName string) (status *fwk.Status) {
    // ...
    if state.allocator != nil {
        allocations, ok := state.nodeAllocations[nodeName]
        for index, claim := range state.claims.toAllocate() {
            allocation := &allocations.allocationResults[allocIndex]
            state.informationsForClaim[index].allocation = allocation
            claim = claim.DeepCopy()
            claim.Status.Allocation = allocation
            // 标记为"进行中分配"
            err := pl.draManager.ResourceClaims().SignalClaimPendingAllocation(claim.UID, claim)
            // ...
        }
    }
    // ...
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11

// pkg/scheduler/framework/plugins/dynamicresources/dynamicresources.go:1317
func (pl *DynamicResources) PreBind(ctx context.Context, cs fwk.CycleState, pod *v1.Pod, nodeName string) *fwk.Status {
    // ...
    for index, claim := range state.claims.all() {
        if !resourceclaim.IsReservedForPod(pod, claim, pl.fts.EnableDRAWorkloadResourceClaims) {
            claim, err := pl.bindClaim(ctx, state, podGroupState, index, pod, nodeName)
            // ...
        }
    }
    // ...
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12

allocateOne(remainingClaims, currentAllocation)
  │
  ├─ 基线情况：所有 Claim 已分配 → 返回成功
  ├─ FirstAvailable：按优先级顺序尝试子请求，第一个成功的胜出
  ├─ All 模式：按预计算列表顺序分配，每个设备必需
  └─ ExactCount 模式：遍历池/切片/设备，对每个候选：
     a. 跳过已使用的设备
     b. 检查 CEL 选择器
     c. 检查污点/容忍
     d. 检查约束（matchAttribute）
     e. 匹配则递归尝试下一个设备索引
     f. 失败时回溯（deallocate），尝试下一个候选

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18

Kubelet 收到绑定的 Pod
    │
    ├─ 1. PrepareResources（4 阶段）
    │     ├─ 阶段 1：验证
    │     │   → 获取每个 ResourceClaim，验证 Pod 在 ReservedFor 中
    │     │   → 解析每个需要的 DRA 驱动
    │     ├─ 阶段 2：缓存更新 + Checkpoint
    │     │   → 将 ClaimInfo 加入缓存
    │     │   → 如果已 Prepared → 跳过；否则构建 gRPC 批次
    │     ├─ 阶段 3：gRPC 调用
    │     │   → 调用 Driver 的 NodePrepareResources
    │     │   → Driver 返回 CDI 设备 ID
    │     └─ 阶段 4：标记已准备 + 最终 Checkpoint
    │
    ├─ 2. GetResources → 为容器运行时提供 CDI 设备列表
    │
    └─ 3. 容器运行时注入
          → 根据 CDI 描述文件挂载设备文件、驱动库，注入环境变量

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28

func (s *DeviceState) Prepare(ctx context.Context, claim *resourceapi.ResourceClaim) ([]kubeletplugin.Device, error) {
    s.Lock()
    defer s.Unlock()

    cp, err := s.getCheckpoint(ctx)
    // 如果 Claim 已经 PrepareCompleted → 直接返回缓存结果（幂等性）
    preparedClaim, exists := cp.V2.PreparedClaims[claimUID]
    if exists && preparedClaim.CheckpointState == ClaimCheckpointStatePrepareCompleted {
        return preparedClaim.PreparedDevices.GetDevices(), nil
    }
    // 如果处于 PrepareStarted（上次崩溃未完成），先回滚
    if exists && preparedClaim.CheckpointState == ClaimCheckpointStatePrepareStarted {
        s.unpreparePartiallyPrepairedClaim(claimUID, preparedClaim, cp)
    }
    // 更新 Checkpoint：标记为 PrepareStarted
    s.updateCheckpoint(ctx, func(cp *Checkpoint) {
        cp.V2.PreparedClaims[claimUID] = PreparedClaim{CheckpointState: ClaimCheckpointStatePrepareStarted, ...}
    })
    // 核心逻辑：按调度器的分配结果准备设备
    preparedDevices, err := s.prepareDevices(ctx, claim)
    // 生成 CDI 描述文件
    s.cdi.CreateClaimSpecFile(claimUID, preparedDevices)
    // 更新 Checkpoint：标记为 PrepareCompleted
    s.updateCheckpoint(ctx, func(cp *Checkpoint) {
        cp.V2.PreparedClaims[claimUID] = PreparedClaim{CheckpointState: ClaimCheckpointStatePrepareCompleted, PreparedDevices: preparedDevices}
    })
    return preparedDevices.GetDevices(), nil
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16

// cmd/gpu-kubelet-plugin/cdi.go:181
func (cdi *CDIHandler) CreateClaimSpecFile(claimUID string, preparedDevices PreparedDevices) error {
    commonEdits, err := cdi.GetCommonEditsCached() // 通用容器编辑（缓存 5 分钟）
    for _, group := range preparedDevices {
        for _, dev := range group.Devices {
            dname := fmt.Sprintf("%s-%s", claimUID, dev.CanonicalName()) // claim 专属名称
            if dev.Type() == GpuDeviceType {
                dspecsgpu, _ := cdi.GetDeviceSpecsByUUIDCached(dev.Gpu.Info.UUID) // 从缓存获取
                // ...
            }
            if dev.Type() == PreparedMigDeviceType { /* MIG：父 GPU 规格 + MIG 设备节点 */ }
            if dev.Type() == VfioDeviceType { /* VFIO：PCI 设备规格 */ }
        }
    }
    // ...
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11

Pod 终止
    │
    ├─ 1. Kubelet 清理
    │     → UnprepareResources：调用 Driver 的 NodeUnprepareResources
    │     → 移除 Pod 引用，当没有 Pod 引用时执行 Unprepare
    │
    ├─ 2. ResourceClaim Controller 清理
    │     → 移除 ReservedFor 条目
    │     → 当 ReservedFor 为空时：清除分配、移除 Finalizer、删除从模板生成的 Claim
    │
    └─ 3. 设备回到可分配状态

1
2
3
4
5
6
7
8
9

// cmd/gpu-kubelet-plugin/driver.go:420
func (d *driver) nodeUnprepareResource(ctx context.Context, claimRef kubeletplugin.NamespacedObject) error {
    release, err := d.pulock.Acquire(ctx, flock.WithTimeout(10*time.Second))
    if err != nil {
        return fmt.Errorf("error acquiring prep/unprep lock: %w", err)
    }
    defer release()
    return d.state.Unprepare(ctx, claimRef)
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20

// cmd/gpu-kubelet-plugin/device_state.go:426
func (s *DeviceState) Unprepare(ctx context.Context, claimRef kubeletplugin.NamespacedObject) error {
    s.Lock()
    defer s.Unlock()

    checkpoint, err := s.getCheckpoint(ctx)
    pc, exists := checkpoint.V2.PreparedClaims[claimUID]
    if !exists {
        return nil // 不在 Checkpoint → 直接返回（幂等性）
    }
    switch pc.CheckpointState {
    case ClaimCheckpointStatePrepareStarted:   // 上次崩溃未完成 → 回滚
        s.unpreparePartiallyPrepairedClaim(claimUID, pc, checkpoint)
    case ClaimCheckpointStatePrepareCompleted:  // 正常完成 → 清理
        s.unprepareDevices(ctx, claimUID, pc.PreparedDevices)
    }
    s.cdi.DeleteClaimSpecFile(claimUID)        // 删除 CDI 描述文件
    s.deleteClaimFromCheckpoint(ctx, claimRef)  // 从 Checkpoint 中删除该 Claim
    return nil
}