plugin server - adding option for plugins to expose http api (#12)

Author: enoodle

cmd/scheduler/app/options/options.go

@@ -54,6 +54,7 @@ type ServerOption struct {
 	AllowConsolidatingReclaim         bool
 	NumOfStatusRecordingWorkers       int
 	GlobalDefaultStalenessGracePeriod time.Duration
+	PluginServerPort                  int
 
 	QPS   int
 	Burst int
@@ -97,6 +98,7 @@ func (s *ServerOption) AddFlags(fs *pflag.FlagSet) {
 	fs.BoolVar(&s.AllowConsolidatingReclaim, "allow-consolidating-reclaim", true, "Do not count pipelined pods towards 'reclaimed' resources")
 	fs.IntVar(&s.NumOfStatusRecordingWorkers, "num-of-status-recording-workers", defaultNumOfStatusRecordingWorkers, "specifies the max number of go routines spawned to update pod and podgroups conditions and events. Defaults to 5")
 	fs.DurationVar(&s.GlobalDefaultStalenessGracePeriod, "default-staleness-grace-period", defaultStalenessGracePeriod, "Global default staleness grace period duration. Negative values means infinite. Defaults to 60s")
+	fs.IntVar(&s.PluginServerPort, "plugin-server-port", 8081, "The port to bind for plugin server requests")
 
 	utilfeature.DefaultMutableFeatureGate.AddFlag(fs)
 }

cmd/scheduler/app/options/options_test.go

@@ -46,6 +46,7 @@ func TestAddFlags(t *testing.T) {
 		GlobalDefaultStalenessGracePeriod: defaultStalenessGracePeriod,
 		NumOfStatusRecordingWorkers:       defaultNumOfStatusRecordingWorkers,
 		NodePoolLabelKey:                  defaultNodePoolLabelKey,
+		PluginServerPort:                  8081,
 	}
 
 	if !reflect.DeepEqual(expected, s) {

cmd/scheduler/app/server.go

@@ -85,6 +85,11 @@ func RunApp() error {
 		os.Exit(1)
 	}
 
+	mux := http.NewServeMux()
+	go func() {
+		_ = http.ListenAndServe(fmt.Sprintf(":%d", so.PluginServerPort), mux)
+	}()
+
 	setupProfiling(so)
 	if err := setupLogging(so); err != nil {
 		fmt.Printf("Failed to initialize loggers: %v", err)
@@ -96,7 +101,7 @@ func RunApp() error {
 	config.QPS = float32(so.QPS)
 	config.Burst = so.Burst
 
-	return Run(so, config)
+	return Run(so, config, mux)
 }
 
 func setupProfiling(so *options.ServerOption) {
@@ -122,14 +127,15 @@ func setupLogging(so *options.ServerOption) error {
 	return nil
 }
 
-func Run(opt *options.ServerOption, config *restclient.Config) error {
+func Run(opt *options.ServerOption, config *restclient.Config, mux *http.ServeMux) error {
 	if opt.PrintVersion {
 		version.PrintVersion()
 	}
 
 	scheduler, err := scheduler.NewScheduler(config,
 		opt.SchedulerConf,
 		BuildSchedulerParams(opt),
+		mux,
 	)
 	if err != nil {
 		return err

pkg/env-tests/scheduler/scheduler.go

@@ -39,7 +39,7 @@ func RunScheduler(cfg *rest.Config, stopCh chan struct{}) error {
 
 	params := app.BuildSchedulerParams(opt)
 
-	s, err := scheduler.NewScheduler(cfg, "", params)
+	s, err := scheduler.NewScheduler(cfg, "", params, nil)
 	if err != nil {
 		return err
 	}

pkg/scheduler/framework/framework.go

@@ -4,6 +4,7 @@
 package framework
 
 import (
+	"net/http"
 	"time"
 
 	"k8s.io/apimachinery/pkg/types"
@@ -15,11 +16,15 @@ import (
 )
 
 func OpenSession(cache cache.Cache, config *conf.SchedulerConfiguration,
-	schedulerParams *conf.SchedulerParams, sessionId types.UID) (*Session, error) {
+	schedulerParams *conf.SchedulerParams, sessionId types.UID, mux *http.ServeMux) (*Session, error) {
 	openSessionStart := time.Now()
 	defer metrics.UpdateOpenSessionDuration(openSessionStart)
 
-	ssn, err := openSession(cache, sessionId, *schedulerParams)
+	if server == nil {
+		server = newPluginServer(mux)
+	}
+
+	ssn, err := openSession(cache, sessionId, *schedulerParams, mux)
 	if err != nil {
 		return nil, err
 	}

pkg/scheduler/framework/server.go

@@ -0,0 +1,45 @@
+// Copyright 2025 NVIDIA CORPORATION
+// SPDX-License-Identifier: Apache-2.0
+
+package framework
+
+import (
+	"fmt"
+	"net/http"
+)
+
+type PluginServer struct {
+	registeredPlugins map[string]func(http.ResponseWriter, *http.Request)
+	mux               *http.ServeMux
+}
+
+func newPluginServer(mux *http.ServeMux) *PluginServer {
+	ps := &PluginServer{
+		registeredPlugins: map[string]func(http.ResponseWriter, *http.Request){},
+		mux:               mux,
+	}
+	if ps.mux != nil {
+		ps.mux.HandleFunc("/", ps.handlePlugin)
+	}
+
+	return ps
+}
+
+func (ps *PluginServer) registerPlugin(path string, handler func(http.ResponseWriter, *http.Request)) error {
+	if ps.mux == nil {
+		return fmt.Errorf("server not initialized")
+	}
+
+	ps.registeredPlugins[path] = handler
+	return nil
+}
+
+func (ps *PluginServer) handlePlugin(w http.ResponseWriter, r *http.Request) {
+	path := r.URL.Path
+	handler, ok := ps.registeredPlugins[path]
+	if !ok {
+		http.Error(w, "plugin not found", http.StatusNotFound)
+		return
+	}
+	handler(w, r)
+}

pkg/scheduler/framework/session.go

@@ -5,6 +5,7 @@ package framework
 
 import (
 	"fmt"
+	"net/http"
 	"sort"
 	"sync"
 	"time"
@@ -29,6 +30,8 @@ import (
 	"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/scheduler_util"
 )
 
+var server *PluginServer
+
 type Session struct {
 	UID   types.UID
 	Cache cache.Cache
@@ -60,6 +63,7 @@ type Session struct {
 	plugins         map[string]Plugin
 	eventHandlers   []*EventHandler
 	schedulerParams conf.SchedulerParams
+	mux             *http.ServeMux
 
 	k8sPodState map[types.UID]k8s_internal.SessionState
 }
@@ -307,7 +311,7 @@ func (ssn *Session) clear() {
 	ssn.JobOrderFns = nil
 }
 
-func openSession(cache cache.Cache, sessionId types.UID, schedulerParams conf.SchedulerParams) (*Session, error) {
+func openSession(cache cache.Cache, sessionId types.UID, schedulerParams conf.SchedulerParams, mux *http.ServeMux) (*Session, error) {
 	ssn := &Session{
 		UID:   sessionId,
 		Cache: cache,
@@ -318,8 +322,8 @@ func openSession(cache cache.Cache, sessionId types.UID, schedulerParams conf.Sc
 
 		plugins:         map[string]Plugin{},
 		schedulerParams: schedulerParams,
-
-		k8sPodState: map[types.UID]k8s_internal.SessionState{},
+		mux:             mux,
+		k8sPodState:     map[types.UID]k8s_internal.SessionState{},
 	}
 
 	log.InfraLogger.V(2).Infof("Taking cluster snapshot ...")

pkg/scheduler/framework/session_plugins.go

@@ -4,6 +4,8 @@
 package framework
 
 import (
+	"net/http"
+
 	"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/api"
 	"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/api/common_info"
 	"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/api/node_info"
@@ -63,6 +65,13 @@ func (ssn *Session) AddGetQueueAllocatedResourcesFn(of api.QueueResource) {
 	ssn.GetQueueAllocatedResourcesFns = append(ssn.GetQueueAllocatedResourcesFns, of)
 }
 
+func (ssn *Session) AddHttpHandler(path string, handler func(http.ResponseWriter, *http.Request)) {
+	if server == nil {
+		return
+	}
+	server.registerPlugin(path, handler)
+}
+
 func (ssn *Session) CanReclaimResources(reclaimer *reclaimer_info.ReclaimerInfo) bool {
 	for _, canReclaimFn := range ssn.CanReclaimResourcesFns {
 		return canReclaimFn(reclaimer)

pkg/scheduler/plugins/podaffinity/podaffinity_test.go

@@ -203,7 +203,9 @@ func testPodPreferredAffinity(testData testInput, clusterAffinityInfo pod_affini
 		mockCache,
 		&conf.SchedulerConfiguration{Tiers: []conf.Tier{}},
 		&conf.SchedulerParams{},
-		sessionId)
+		sessionId,
+		nil,
+	)
 	Expect(err).To(Succeed())
 
 	Expect(ssn.NodePreOrderFns).To(HaveLen(0), "Expected no node pre-order functions")

pkg/scheduler/plugins/proportion/proportion_test.go

@@ -765,7 +765,7 @@ var _ = Describe("Set Fair Share in Proportion", func() {
 						RestrictSchedulingNodes: testData.isRestrictNode,
 						SchedulerName:           schedulerName,
 					},
-					"1")
+					"1", nil)
 				if got := getNodeResources(session, testData.node); !reflect.DeepEqual(got, testData.want) {
 					Fail(fmt.Sprintf("getNodeResources() = %v, want %v", got, testData.want))
 				}