1/* 2Copyright 2019 The Kubernetes Authors. 3 4Licensed under the Apache License, Version 2.0 (the "License"); 5you may not use this file except in compliance with the License. 6You may obtain a copy of the License at 7 8 http://www.apache.org/licenses/LICENSE-2.0 9 10Unless required by applicable law or agreed to in writing, software 11distributed under the License is distributed on an "AS IS" BASIS, 12WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13See the License for the specific language governing permissions and 14limitations under the License. 15*/ 16 17package testing 18 19import ( 20 "k8s.io/apimachinery/pkg/runtime/schema" 21 "k8s.io/kube-scheduler/config/v1beta2" 22 schedulerapi "k8s.io/kubernetes/pkg/scheduler/apis/config" 23 "k8s.io/kubernetes/pkg/scheduler/apis/config/scheme" 24 "k8s.io/kubernetes/pkg/scheduler/framework" 25 "k8s.io/kubernetes/pkg/scheduler/framework/runtime" 26) 27 28var configDecoder = scheme.Codecs.UniversalDecoder() 29 30// NewFramework creates a Framework from the register functions and options. 31func NewFramework(fns []RegisterPluginFunc, profileName string, opts ...runtime.Option) (framework.Framework, error) { 32 registry := runtime.Registry{} 33 profile := &schedulerapi.KubeSchedulerProfile{ 34 SchedulerName: profileName, 35 Plugins: &schedulerapi.Plugins{}, 36 } 37 for _, f := range fns { 38 f(®istry, profile) 39 } 40 return runtime.NewFramework(registry, profile, opts...) 41} 42 43// RegisterPluginFunc is a function signature used in method RegisterFilterPlugin() 44// to register a Filter Plugin to a given registry. 45type RegisterPluginFunc func(reg *runtime.Registry, profile *schedulerapi.KubeSchedulerProfile) 46 47// RegisterQueueSortPlugin returns a function to register a QueueSort Plugin to a given registry. 48func RegisterQueueSortPlugin(pluginName string, pluginNewFunc runtime.PluginFactory) RegisterPluginFunc { 49 return RegisterPluginAsExtensions(pluginName, pluginNewFunc, "QueueSort") 50} 51 52// RegisterPreFilterPlugin returns a function to register a PreFilter Plugin to a given registry. 53func RegisterPreFilterPlugin(pluginName string, pluginNewFunc runtime.PluginFactory) RegisterPluginFunc { 54 return RegisterPluginAsExtensions(pluginName, pluginNewFunc, "PreFilter") 55} 56 57// RegisterFilterPlugin returns a function to register a Filter Plugin to a given registry. 58func RegisterFilterPlugin(pluginName string, pluginNewFunc runtime.PluginFactory) RegisterPluginFunc { 59 return RegisterPluginAsExtensions(pluginName, pluginNewFunc, "Filter") 60} 61 62// RegisterReservePlugin returns a function to register a Reserve Plugin to a given registry. 63func RegisterReservePlugin(pluginName string, pluginNewFunc runtime.PluginFactory) RegisterPluginFunc { 64 return RegisterPluginAsExtensions(pluginName, pluginNewFunc, "Reserve") 65} 66 67// RegisterPermitPlugin returns a function to register a Permit Plugin to a given registry. 68func RegisterPermitPlugin(pluginName string, pluginNewFunc runtime.PluginFactory) RegisterPluginFunc { 69 return RegisterPluginAsExtensions(pluginName, pluginNewFunc, "Permit") 70} 71 72// RegisterPreBindPlugin returns a function to register a PreBind Plugin to a given registry. 73func RegisterPreBindPlugin(pluginName string, pluginNewFunc runtime.PluginFactory) RegisterPluginFunc { 74 return RegisterPluginAsExtensions(pluginName, pluginNewFunc, "PreBind") 75} 76 77// RegisterScorePlugin returns a function to register a Score Plugin to a given registry. 78func RegisterScorePlugin(pluginName string, pluginNewFunc runtime.PluginFactory, weight int32) RegisterPluginFunc { 79 return RegisterPluginAsExtensionsWithWeight(pluginName, weight, pluginNewFunc, "Score") 80} 81 82// RegisterPreScorePlugin returns a function to register a Score Plugin to a given registry. 83func RegisterPreScorePlugin(pluginName string, pluginNewFunc runtime.PluginFactory) RegisterPluginFunc { 84 return RegisterPluginAsExtensions(pluginName, pluginNewFunc, "PreScore") 85} 86 87// RegisterBindPlugin returns a function to register a Bind Plugin to a given registry. 88func RegisterBindPlugin(pluginName string, pluginNewFunc runtime.PluginFactory) RegisterPluginFunc { 89 return RegisterPluginAsExtensions(pluginName, pluginNewFunc, "Bind") 90} 91 92// RegisterPluginAsExtensions returns a function to register a Plugin as given extensionPoints to a given registry. 93func RegisterPluginAsExtensions(pluginName string, pluginNewFunc runtime.PluginFactory, extensions ...string) RegisterPluginFunc { 94 return RegisterPluginAsExtensionsWithWeight(pluginName, 1, pluginNewFunc, extensions...) 95} 96 97// RegisterPluginAsExtensionsWithWeight returns a function to register a Plugin as given extensionPoints with weight to a given registry. 98func RegisterPluginAsExtensionsWithWeight(pluginName string, weight int32, pluginNewFunc runtime.PluginFactory, extensions ...string) RegisterPluginFunc { 99 return func(reg *runtime.Registry, profile *schedulerapi.KubeSchedulerProfile) { 100 reg.Register(pluginName, pluginNewFunc) 101 for _, extension := range extensions { 102 ps := getPluginSetByExtension(profile.Plugins, extension) 103 if ps == nil { 104 continue 105 } 106 ps.Enabled = append(ps.Enabled, schedulerapi.Plugin{Name: pluginName, Weight: weight}) 107 } 108 // Use defaults from latest config API version. 109 var gvk schema.GroupVersionKind 110 gvk = v1beta2.SchemeGroupVersion.WithKind(pluginName + "Args") 111 if args, _, err := configDecoder.Decode(nil, &gvk, nil); err == nil { 112 profile.PluginConfig = append(profile.PluginConfig, schedulerapi.PluginConfig{ 113 Name: pluginName, 114 Args: args, 115 }) 116 } 117 } 118} 119 120func getPluginSetByExtension(plugins *schedulerapi.Plugins, extension string) *schedulerapi.PluginSet { 121 switch extension { 122 case "QueueSort": 123 return &plugins.QueueSort 124 case "Filter": 125 return &plugins.Filter 126 case "PreFilter": 127 return &plugins.PreFilter 128 case "PreScore": 129 return &plugins.PreScore 130 case "Score": 131 return &plugins.Score 132 case "Bind": 133 return &plugins.Bind 134 case "Reserve": 135 return &plugins.Reserve 136 case "Permit": 137 return &plugins.Permit 138 case "PreBind": 139 return &plugins.PreBind 140 case "PostBind": 141 return &plugins.PostBind 142 default: 143 return nil 144 } 145} 146