1
16
17 package framework
18
19 import (
20 "context"
21
22 v1 "k8s.io/api/core/v1"
23 "k8s.io/apimachinery/pkg/runtime/schema"
24 kubeschedulerconfigv1 "k8s.io/kube-scheduler/config/v1"
25 schedulerapi "k8s.io/kubernetes/pkg/scheduler/apis/config"
26 "k8s.io/kubernetes/pkg/scheduler/apis/config/scheme"
27 "k8s.io/kubernetes/pkg/scheduler/framework"
28 "k8s.io/kubernetes/pkg/scheduler/framework/runtime"
29 )
30
31 var configDecoder = scheme.Codecs.UniversalDecoder()
32
33
34 func NewFramework(ctx context.Context, fns []RegisterPluginFunc, profileName string, opts ...runtime.Option) (framework.Framework, error) {
35 registry := runtime.Registry{}
36 profile := &schedulerapi.KubeSchedulerProfile{
37 SchedulerName: profileName,
38 Plugins: &schedulerapi.Plugins{},
39 }
40 for _, f := range fns {
41 f(®istry, profile)
42 }
43 return runtime.NewFramework(ctx, registry, profile, opts...)
44 }
45
46
47
48 type RegisterPluginFunc func(reg *runtime.Registry, profile *schedulerapi.KubeSchedulerProfile)
49
50
51 func RegisterQueueSortPlugin(pluginName string, pluginNewFunc runtime.PluginFactory) RegisterPluginFunc {
52 return RegisterPluginAsExtensions(pluginName, pluginNewFunc, "QueueSort")
53 }
54
55
56 func RegisterPreFilterPlugin(pluginName string, pluginNewFunc runtime.PluginFactory) RegisterPluginFunc {
57 return RegisterPluginAsExtensions(pluginName, pluginNewFunc, "PreFilter")
58 }
59
60
61 func RegisterFilterPlugin(pluginName string, pluginNewFunc runtime.PluginFactory) RegisterPluginFunc {
62 return RegisterPluginAsExtensions(pluginName, pluginNewFunc, "Filter")
63 }
64
65
66 func RegisterReservePlugin(pluginName string, pluginNewFunc runtime.PluginFactory) RegisterPluginFunc {
67 return RegisterPluginAsExtensions(pluginName, pluginNewFunc, "Reserve")
68 }
69
70
71 func RegisterPermitPlugin(pluginName string, pluginNewFunc runtime.PluginFactory) RegisterPluginFunc {
72 return RegisterPluginAsExtensions(pluginName, pluginNewFunc, "Permit")
73 }
74
75
76 func RegisterPreBindPlugin(pluginName string, pluginNewFunc runtime.PluginFactory) RegisterPluginFunc {
77 return RegisterPluginAsExtensions(pluginName, pluginNewFunc, "PreBind")
78 }
79
80
81 func RegisterScorePlugin(pluginName string, pluginNewFunc runtime.PluginFactory, weight int32) RegisterPluginFunc {
82 return RegisterPluginAsExtensionsWithWeight(pluginName, weight, pluginNewFunc, "Score")
83 }
84
85
86 func RegisterPreScorePlugin(pluginName string, pluginNewFunc runtime.PluginFactory) RegisterPluginFunc {
87 return RegisterPluginAsExtensions(pluginName, pluginNewFunc, "PreScore")
88 }
89
90
91 func RegisterBindPlugin(pluginName string, pluginNewFunc runtime.PluginFactory) RegisterPluginFunc {
92 return RegisterPluginAsExtensions(pluginName, pluginNewFunc, "Bind")
93 }
94
95
96 func RegisterPluginAsExtensions(pluginName string, pluginNewFunc runtime.PluginFactory, extensions ...string) RegisterPluginFunc {
97 return RegisterPluginAsExtensionsWithWeight(pluginName, 1, pluginNewFunc, extensions...)
98 }
99
100
101 func RegisterPluginAsExtensionsWithWeight(pluginName string, weight int32, pluginNewFunc runtime.PluginFactory, extensions ...string) RegisterPluginFunc {
102 return func(reg *runtime.Registry, profile *schedulerapi.KubeSchedulerProfile) {
103 reg.Register(pluginName, pluginNewFunc)
104 for _, extension := range extensions {
105 ps := getPluginSetByExtension(profile.Plugins, extension)
106 if ps == nil {
107 continue
108 }
109 ps.Enabled = append(ps.Enabled, schedulerapi.Plugin{Name: pluginName, Weight: weight})
110 }
111
112 var gvk schema.GroupVersionKind
113 gvk = kubeschedulerconfigv1.SchemeGroupVersion.WithKind(pluginName + "Args")
114 if args, _, err := configDecoder.Decode(nil, &gvk, nil); err == nil {
115 profile.PluginConfig = append(profile.PluginConfig, schedulerapi.PluginConfig{
116 Name: pluginName,
117 Args: args,
118 })
119 }
120 }
121 }
122
123 func getPluginSetByExtension(plugins *schedulerapi.Plugins, extension string) *schedulerapi.PluginSet {
124 switch extension {
125 case "QueueSort":
126 return &plugins.QueueSort
127 case "Filter":
128 return &plugins.Filter
129 case "PreFilter":
130 return &plugins.PreFilter
131 case "PreScore":
132 return &plugins.PreScore
133 case "Score":
134 return &plugins.Score
135 case "Bind":
136 return &plugins.Bind
137 case "Reserve":
138 return &plugins.Reserve
139 case "Permit":
140 return &plugins.Permit
141 case "PreBind":
142 return &plugins.PreBind
143 case "PostBind":
144 return &plugins.PostBind
145 default:
146 return nil
147 }
148 }
149
150
151 func BuildNodeInfos(nodes []*v1.Node) []*framework.NodeInfo {
152 res := make([]*framework.NodeInfo, len(nodes))
153 for i := 0; i < len(nodes); i++ {
154 res[i] = framework.NewNodeInfo()
155 res[i].SetNode(nodes[i])
156 }
157 return res
158 }
159
View as plain text