1
package gpunode
2

3
import (
4
    "context"
5
    "fmt"
6

7
    "github.com/kubernetes-sigs/scheduler-plugins/pkg/apis/config/scheme"
8
    v1 "k8s.io/api/core/v1"
9
    "k8s.io/apimachinery/pkg/runtime"
10
    "k8s.io/kubernetes/pkg/scheduler/framework"
11
)
12

13
const Name = "NodeGPUFilter"
14

15
// FilterArgs 定义插件配置参数
16
type FilterArgs struct {
17
    // 需要 GPU 的 Pod 需要调度到哪些标签的节点
18
    GPULabelKey   string `json:"gpuLabelKey,omitempty"`
19
    GPULabelValue string `json:"gpuLabelValue,omitempty"`
20
}
21

22
// Filter 实现 Filter 扩展点
23
type Filter struct {
24
    args FilterArgs
25
}
26

27
var _ framework.FilterPlugin = &Filter{}
28

29
func (f *Filter) Name() string {
30
    return Name
31
}
32

33
// Filter 检查节点是否满足调度条件
34
func (f *Filter) Filter(ctx context.Context, _ *framework.CycleState, pod *v1.Pod, nodeInfo *framework.NodeInfo) *framework.Status {
35
    // 获取 Pod 是否需要 GPU
36
    needsGPU, ok := pod.Annotations["app.gpu"]
37
    if !ok || needsGPU != "true" {
38
        // Pod 不需要 GPU，直接通过
39
        return nil
40
    }
41

42
    // 检查节点是否有 GPU 标签
43
    node := nodeInfo.Node()
44
    if node == nil {
45
        return framework.NewStatus(framework.Unschedulable, "node is nil")
46
    }
47

48
    gpuValue, exists := node.Labels[f.args.GPULabelKey]
49
    if !exists || gpuValue != f.args.GPULabelValue {
50
        return framework.NewStatus(
51
            framework.Unschedulable,
52
            fmt.Sprintf("node %s does not have GPU label %s=%s",
53
                node.Name, f.args.GPULabelKey, f.args.GPULabelValue),
54
        )
55
    }
56

57
    return nil
58
}
59

60
// NewFilter 创建 Filter 插件实例
61
func NewFilter(_ context.Context, plArgs runtime.Object) (framework.Plugin, error) {
62
    args := &FilterArgs{}
63
    if err := scheme.Scheme.Convert(plArgs, args, nil); err != nil {
64
        return nil, fmt.Errorf("failed to convert args: %v", err)
65
    }
66

67
    if args.GPULabelKey == "" {
68
        args.GPULabelKey = "gpu.kubernetes.io"
69
    }
70
    if args.GPULabelValue == "" {
71
        args.GPULabelValue = "true"
72
    }
73

74
    return &Filter{args: *args}, nil
75
}

1
package gpunode
2

3
import (
4
    "context"
5
    "fmt"
6

7
    v1 "k8s.io/api/core/v1"
8
    "k8s.io/apimachinery/pkg/runtime"
9
    "k8s.io/kubernetes/pkg/scheduler/framework"
10
)
11

12
const (
13
    // GPU 节点的基础分数
14
    gpuNodeScore = 100
15
    // 普通节点的分数
16
    normalNodeScore = 50
17
)
18

19
// Score 实现 Score 扩展点
20
type Score struct {
21
    handle framework.Handle
22
}
23

24
var _ framework.ScorePlugin = &Score{}
25

26
func NewScore(_ context.Context, _ runtime.Object, h framework.Handle) (framework.Plugin, error) {
27
    return &Score{handle: h}, nil
28
}
29

30
func (s *Score) Name() string {
31
    return Name
32
}
33

34
// Score 对节点进行评分
35
// 注意：Score 函数接收的是节点名称，需要通过 handle 获取节点信息
36
func (s *Score) Score(ctx context.Context, state *framework.CycleState, pod *v1.Pod, nodeName string) (int64, *framework.Status) {
37
    // 通过 handle 获取节点信息
38
    nodeInfo, err := s.handle.SnapshotSharedLister().NodeInfos().Get(nodeName)
39
    if err != nil {
40
        return 0, framework.NewStatus(framework.Error, fmt.Sprintf("failed to get node %s: %v", nodeName, err))
41
    }
42

43
    node := nodeInfo.Node()
44
    if node == nil {
45
        return 0, framework.NewStatus(framework.Error, "node is nil")
46
    }
47

48
    // 检查是否为 GPU 节点
49
    needsGPU, _ := pod.Annotations["app.gpu"]
50
    if needsGPU == "true" {
51
        if gpuValue, ok := node.Labels["gpu.kubernetes.io"]; ok && gpuValue == "true" {
52
            return gpuNodeScore, nil
53
        }
54
        return 1, nil // GPU Pod 但节点无 GPU，给最低分
55
    }
56

57
    return normalNodeScore, nil
58
}
59

60
// ScoreExtensions 返回分数的扩展信息
61
func (s *Score) ScoreExtensions() framework.ScoreExtensions {
62
    return nil
63
}

1
apiVersion: kubescheduler.config.k8s.io/v1beta3
2
kind: KubeSchedulerConfiguration
3
profiles:
4
  - schedulerName: default-scheduler
5
    pluginConfig:
6
      - name: NodeGPUFilter
7
        args:
8
          gpuLabelKey: gpu.kubernetes.io
9
          gpuLabelValue: "true"