add scheduling framework extender (#365)

Signed-off-by: saintube <saintube@foxmail.com>

Author: saintube

cmd/koord-scheduler/app/server.go

@@ -65,7 +65,7 @@ import (
 type Option func(frameworkext.ExtendedHandle, runtime.Registry) error
 
 // NewSchedulerCommand creates a *cobra.Command object with default parameters and registryOptions
-func NewSchedulerCommand(registryOptions ...Option) *cobra.Command {
+func NewSchedulerCommand(schedulingHooks []frameworkext.SchedulingPhaseHook, registryOptions ...Option) *cobra.Command {
 	opts := options.NewOptions()
 
 	cmd := &cobra.Command{
@@ -78,7 +78,7 @@ scenarios,ensuring the runtime quality of different workloads and users' demands
 for cost reduction and efficiency enhancement.
 `,
 		Run: func(cmd *cobra.Command, args []string) {
-			if err := runCommand(cmd, opts, registryOptions...); err != nil {
+			if err := runCommand(cmd, opts, schedulingHooks, registryOptions...); err != nil {
 				fmt.Fprintf(os.Stderr, "%v\n", err)
 				os.Exit(1)
 			}
@@ -110,7 +110,7 @@ for cost reduction and efficiency enhancement.
 }
 
 // runCommand runs the scheduler.
-func runCommand(cmd *cobra.Command, opts *options.Options, registryOptions ...Option) error {
+func runCommand(cmd *cobra.Command, opts *options.Options, schedulingHooks []frameworkext.SchedulingPhaseHook, registryOptions ...Option) error {
 	verflag.PrintAndExitIfRequested()
 	cliflag.PrintFlags(cmd.Flags())
 
@@ -122,7 +122,7 @@ func runCommand(cmd *cobra.Command, opts *options.Options, registryOptions ...Op
 		cancel()
 	}()
 
-	cc, sched, err := Setup(ctx, opts, registryOptions...)
+	cc, sched, err := Setup(ctx, opts, schedulingHooks, registryOptions...)
 	if err != nil {
 		return err
 	}
@@ -307,7 +307,7 @@ func WithPlugin(name string, factory runtime.PluginFactory) Option {
 }
 
 // Setup creates a completed config and a scheduler based on the command args and options
-func Setup(ctx context.Context, opts *options.Options, outOfTreeRegistryOptions ...Option) (*schedulerserverconfig.CompletedConfig, *scheduler.Scheduler, error) {
+func Setup(ctx context.Context, opts *options.Options, schedulingHooks []frameworkext.SchedulingPhaseHook, outOfTreeRegistryOptions ...Option) (*schedulerserverconfig.CompletedConfig, *scheduler.Scheduler, error) {
 	if cfg, err := latest.Default(); err != nil {
 		return nil, nil, err
 	} else {
@@ -373,5 +373,12 @@ func Setup(ctx context.Context, opts *options.Options, outOfTreeRegistryOptions
 	// TODO(joseph): Some extensions can also be made in the future,
 	//  such as replacing some interfaces in Scheduler to implement custom logic
 
+	// extend framework to hook run plugin functions
+	extendedFrameworkFactory := frameworkext.NewFrameworkExtenderFactory(extendedHandle, schedulingHooks...)
+	for k, v := range sched.Profiles {
+		sched.Profiles[k] = extendedFrameworkFactory.New(v)
+	}
+	// TODO: register event handlers for scheduler instance
+
 	return &cc, sched, nil
 }

cmd/koord-scheduler/main.go

@@ -24,6 +24,7 @@ import (
 	"k8s.io/component-base/logs"
 
 	"github.com/koordinator-sh/koordinator/cmd/koord-scheduler/app"
+	"github.com/koordinator-sh/koordinator/pkg/scheduler/frameworkext"
 	"github.com/koordinator-sh/koordinator/pkg/scheduler/plugins/compatibledefaultpreemption"
 	"github.com/koordinator-sh/koordinator/pkg/scheduler/plugins/loadaware"
 	"github.com/koordinator-sh/koordinator/pkg/scheduler/plugins/nodenumaresource"
@@ -35,10 +36,15 @@ import (
 func main() {
 	rand.Seed(time.Now().UnixNano())
 
+	// Register custom scheduling hooks for pre-process scheduling context before call plugins.
+	// e.g. change the nodeInfo and make a copy before calling filter plugins
+	var schedulingHooks []frameworkext.SchedulingPhaseHook
+
 	// Register custom plugins to the scheduler framework.
 	// Later they can consist of scheduler profile(s) and hence
 	// used by various kinds of workloads.
 	command := app.NewSchedulerCommand(
+		schedulingHooks,
 		app.WithPlugin(loadaware.Name, loadaware.New),
 		app.WithPlugin(nodenumaresource.Name, nodenumaresource.New),
 		app.WithPlugin(compatibledefaultpreemption.Name, compatibledefaultpreemption.New),

pkg/scheduler/frameworkext/framework_extender.go

@@ -17,10 +17,13 @@ limitations under the License.
 package frameworkext
 
 import (
+	"context"
 	"sync"
 
 	nrtinformers "github.com/k8stopologyawareschedwg/noderesourcetopology-api/pkg/generated/informers/externalversions"
+	corev1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/runtime"
+	"k8s.io/klog/v2"
 	"k8s.io/kubernetes/pkg/scheduler/framework"
 	frameworkruntime "k8s.io/kubernetes/pkg/scheduler/framework/runtime"
 
@@ -96,6 +99,96 @@ func (ext *frameworkExtendedHandleImpl) NodeResourceTopologySharedInformerFactor
 	return ext.nrtSharedInformerFactory
 }
 
+type FrameworkExtender interface {
+	framework.Framework
+}
+
+type FrameworkExtenderFactory interface {
+	New(f framework.Framework) FrameworkExtender
+}
+
+type SchedulingPhaseHook interface {
+	Name() string
+}
+
+type PreFilterPhaseHook interface {
+	SchedulingPhaseHook
+	PreFilterHook(handle ExtendedHandle, state *framework.CycleState, pod *corev1.Pod) (*corev1.Pod, bool)
+}
+
+type FilterPhaseHook interface {
+	SchedulingPhaseHook
+	FilterHook(handle ExtendedHandle, cycleState *framework.CycleState, pod *corev1.Pod, nodeInfo *framework.NodeInfo) (*corev1.Pod, *framework.NodeInfo, bool)
+}
+
+type frameworkExtenderFactoryImpl struct {
+	handle ExtendedHandle
+
+	// extend framework with SchedulingPhaseHook
+	preFilterHooks []PreFilterPhaseHook
+	filterHooks    []FilterPhaseHook
+}
+
+func NewFrameworkExtenderFactory(handle ExtendedHandle, hooks ...SchedulingPhaseHook) FrameworkExtenderFactory {
+	i := &frameworkExtenderFactoryImpl{
+		handle: handle,
+	}
+	for _, h := range hooks {
+		// a hook may register in multiple phases
+		preFilter, ok := h.(PreFilterPhaseHook)
+		if ok {
+			i.preFilterHooks = append(i.preFilterHooks, preFilter)
+		}
+		filter, ok := h.(FilterPhaseHook)
+		if ok {
+			i.filterHooks = append(i.filterHooks, filter)
+		}
+	}
+	return i
+}
+
+func (i *frameworkExtenderFactoryImpl) New(f framework.Framework) FrameworkExtender {
+	return &frameworkExtenderImpl{
+		Framework:      f,
+		handle:         i.handle,
+		preFilterHooks: i.preFilterHooks,
+		filterHooks:    i.filterHooks,
+	}
+}
+
+var _ framework.Framework = &frameworkExtenderImpl{}
+
+type frameworkExtenderImpl struct {
+	framework.Framework
+	handle ExtendedHandle
+
+	preFilterHooks []PreFilterPhaseHook
+	filterHooks    []FilterPhaseHook
+}
+
+func (ext *frameworkExtenderImpl) RunPreFilterPlugins(ctx context.Context, cycleState *framework.CycleState, pod *corev1.Pod) *framework.Status {
+	for _, hook := range ext.preFilterHooks {
+		newPod, hooked := hook.PreFilterHook(ext.handle, cycleState, pod)
+		if hooked {
+			klog.V(5).InfoS("RunPreFilterPlugins hooked", "meet PreFilterPhaseHook", "hook", hook.Name(), "pod", klog.KObj(pod))
+			return ext.Framework.RunPreFilterPlugins(ctx, cycleState, newPod)
+		}
+	}
+	return ext.Framework.RunPreFilterPlugins(ctx, cycleState, pod)
+}
+
+func (ext *frameworkExtenderImpl) RunFilterPlugins(ctx context.Context, cycleState *framework.CycleState, pod *corev1.Pod, nodeInfo *framework.NodeInfo) framework.PluginToStatus {
+	for _, hook := range ext.filterHooks {
+		// hook can change the args (cycleState, pod, nodeInfo) for filter plugins
+		newPod, newNodeInfo, hooked := hook.FilterHook(ext.handle, cycleState, pod, nodeInfo)
+		if hooked {
+			klog.V(5).InfoS("RunFilterPlugins hooked", "meet FilterPhaseHook", "hook", hook.Name(), "pod", klog.KObj(pod))
+			return ext.Framework.RunFilterPlugins(ctx, cycleState, newPod, newNodeInfo)
+		}
+	}
+	return ext.Framework.RunFilterPlugins(ctx, cycleState, pod, nodeInfo)
+}
+
 // PluginFactoryProxy is used to proxy the call to the PluginFactory function and pass in the ExtendedHandle for the custom plugin
 func PluginFactoryProxy(extendHandle ExtendedHandle, factoryFn frameworkruntime.PluginFactory) frameworkruntime.PluginFactory {
 	return func(args runtime.Object, handle framework.Handle) (framework.Plugin, error) {