remove the PreCycle plugin from scheduler (kubernetes-sigs#876)

* remove the PreCycle plugin from scheduler

Signed-off-by: Nir Rozenbaum <[email protected]>

* Apply suggestions from code review

Co-authored-by: Cong Liu <[email protected]>

---------

Signed-off-by: Nir Rozenbaum <[email protected]>
Co-authored-by: Cong Liu <[email protected]>

Author: 2 people

pkg/epp/metrics/metrics_test.go

@@ -571,11 +571,6 @@ func TestSchedulerPluginProcessingLatencies(t *testing.T) {
 		{
 			name: "multiple plugins",
 			latencies: []pluginLatency{
-				{
-					pluginType: "PreSchedule",
-					pluginName: "PluginA",
-					duration:   100 * time.Millisecond,
-				},
 				{
 					pluginType: "PostSchedule",
 					pluginName: "PluginB",

pkg/epp/metrics/testdata/scheduler_plugin_processing_latencies_metric

@@ -1,18 +1,5 @@
 # HELP inference_extension_scheduler_plugin_duration_seconds [ALPHA] Scheduler plugin processing latency distribution in seconds for each plugin type and plugin name.
 # TYPE inference_extension_scheduler_plugin_duration_seconds histogram
-inference_extension_scheduler_plugin_duration_seconds_bucket{plugin_name="PluginA",plugin_type="PreSchedule",le="0.0001"} 0
-inference_extension_scheduler_plugin_duration_seconds_bucket{plugin_name="PluginA",plugin_type="PreSchedule",le="0.0002"} 0
-inference_extension_scheduler_plugin_duration_seconds_bucket{plugin_name="PluginA",plugin_type="PreSchedule",le="0.0005"} 0
-inference_extension_scheduler_plugin_duration_seconds_bucket{plugin_name="PluginA",plugin_type="PreSchedule",le="0.001"} 0
-inference_extension_scheduler_plugin_duration_seconds_bucket{plugin_name="PluginA",plugin_type="PreSchedule",le="0.002"} 0
-inference_extension_scheduler_plugin_duration_seconds_bucket{plugin_name="PluginA",plugin_type="PreSchedule",le="0.005"} 0
-inference_extension_scheduler_plugin_duration_seconds_bucket{plugin_name="PluginA",plugin_type="PreSchedule",le="0.01"} 0
-inference_extension_scheduler_plugin_duration_seconds_bucket{plugin_name="PluginA",plugin_type="PreSchedule",le="0.02"} 0
-inference_extension_scheduler_plugin_duration_seconds_bucket{plugin_name="PluginA",plugin_type="PreSchedule",le="0.05"} 0
-inference_extension_scheduler_plugin_duration_seconds_bucket{plugin_name="PluginA",plugin_type="PreSchedule",le="0.1"} 1
-inference_extension_scheduler_plugin_duration_seconds_bucket{plugin_name="PluginA",plugin_type="PreSchedule",le="+Inf"} 1
-inference_extension_scheduler_plugin_duration_seconds_sum{plugin_name="PluginA",plugin_type="PreSchedule"} 0.1
-inference_extension_scheduler_plugin_duration_seconds_count{plugin_name="PluginA",plugin_type="PreSchedule"} 1
 inference_extension_scheduler_plugin_duration_seconds_bucket{plugin_name="PluginB",plugin_type="PostSchedule",le="0.0001"} 0
 inference_extension_scheduler_plugin_duration_seconds_bucket{plugin_name="PluginB",plugin_type="PostSchedule",le="0.0002"} 0
 inference_extension_scheduler_plugin_duration_seconds_bucket{plugin_name="PluginB",plugin_type="PostSchedule",le="0.0005"} 0

pkg/epp/scheduling/framework/plugins.go

@@ -22,7 +22,6 @@ import (
 
 const (
 	ProfilePickerType      = "ProfilePicker"
-	PreCyclePluginType     = "PreCycle"
 	FilterPluginType       = "Filter"
 	ScorerPluginType       = "Scorer"
 	PickerPluginType       = "Picker"
@@ -44,13 +43,6 @@ type ProfilePicker interface {
 	Pick(request *types.LLMRequest, profiles map[string]*SchedulerProfile, executionResults map[string]*types.Result) map[string]*SchedulerProfile
 }
 
-// PreCycle is called when the scheduler receives a new request and invokes a SchedulerProfile cycle.
-// It can be used for various initialization work.
-type PreCycle interface {
-	Plugin
-	PreCycle(ctx *types.SchedulingContext)
-}
-
 // Filter defines the interface for filtering a list of pods based on context.
 type Filter interface {
 	Plugin

pkg/epp/scheduling/framework/plugins/multi/prefix/plugin.go

@@ -114,7 +114,6 @@ func (s *schedulingContextState) Clone() types.StateData {
 }
 
 // compile-time type assertion
-var _ framework.PreCycle = &Plugin{}
 var _ framework.Scorer = &Plugin{}
 var _ framework.PostCycle = &Plugin{}
 
@@ -132,18 +131,6 @@ func (m *Plugin) Name() string {
 	return "prefix-cache"
 }
 
-// PreCycle initializes the prefix plugin state for the current scheduling cycle.
-func (m *Plugin) PreCycle(ctx *types.SchedulingContext) {
-	hashes := hashPrompt(ctx, m.HashBlockSize, m.MaxPrefixBlocksToMatch)
-	state := &schedulingContextState{
-		PrefixHashes:       hashes,
-		PrefixCacheServers: m.matchLongestPrefix(ctx, hashes, DefaultNumServersToMatch),
-	}
-
-	ctx.CycleState.Write(types.StateKey(m.Name()), state)
-	ctx.Logger.V(logutil.TRACE).Info(fmt.Sprintf("PreCycle, cached servers: %+v", state.PrefixCacheServers), "hashes", state.PrefixHashes)
-}
-
 // PostCycle records in the plugin cache the result of the scheduling selection.
 func (m *Plugin) PostCycle(ctx *types.SchedulingContext, res *types.Result) {
 	targetPod := res.TargetPod.GetPod()
@@ -160,13 +147,20 @@ func (m *Plugin) PostCycle(ctx *types.SchedulingContext, res *types.Result) {
 
 // Score returns the scoring result for the given list of pods based on context.
 func (m *Plugin) Score(ctx *types.SchedulingContext, pods []types.Pod) map[types.Pod]float64 {
-	scores := make(map[types.Pod]float64, len(pods))
-
-	state, err := m.getPrefixState(ctx.CycleState)
-	if err != nil {
-		ctx.Logger.Error(err, "failed to read prefix plugin cycle state")
-		return scores
+	// pre score step, hashing prompt and find longest prefix match.
+	hashes := hashPrompt(ctx, m.HashBlockSize, m.MaxPrefixBlocksToMatch)
+	numServers := DefaultNumServersToMatch
+	if numServers > len(pods) {
+		numServers = len(pods)
+	}
+	state := &schedulingContextState{
+		PrefixHashes:       hashes,
+		PrefixCacheServers: m.matchLongestPrefix(ctx, hashes, numServers),
 	}
+	ctx.CycleState.Write(types.StateKey(m.Name()), state)
+	ctx.Logger.V(logutil.TRACE).Info(fmt.Sprintf("cached servers: %+v", state.PrefixCacheServers), "hashes", state.PrefixHashes)
+	// calculate the scores of pods
+	scores := make(map[types.Pod]float64, len(pods))
 
 	total := len(state.PrefixHashes)
 	podScoreFunc := func(pod types.Pod) float64 {
@@ -185,9 +179,6 @@ func (m *Plugin) Score(ctx *types.SchedulingContext, pods []types.Pod) map[types
 
 // matchLongestPrefix returns a map of servers and length of prefix that each server caches.
 func (m *Plugin) matchLongestPrefix(ctx *types.SchedulingContext, hashes []BlockHash, numServers int) map[ServerID]int {
-	if numServers > len(ctx.PodsSnapshot) {
-		numServers = len(ctx.PodsSnapshot)
-	}
 	res := make(map[ServerID]int)
 	// Use a greedy strategy to search from the longest prefix.
 	// NOTE: It's possible to further optimize this with a binary search.

pkg/epp/scheduling/framework/plugins/multi/prefix/plugin_test.go

@@ -28,17 +28,14 @@ func TestPrefixPlugin(t *testing.T) {
 		Prompt:      "aaaaaa",
 	}
 	ctx := types.NewSchedulingContext(context.Background(), req1, nil, pods)
-	plugin.PreCycle(ctx)
+	scores := plugin.Score(ctx, pods)
 	state, err := plugin.getPrefixState(ctx.CycleState)
 	assert.NoError(t, err)
 	t.Logf("Hashes %+v, cached servers: %+v", state.PrefixHashes, state.PrefixCacheServers)
 	// Input size is 6, hash block size is 4, the last 2 characters are ignored.
 	// Total hashes = 2 (the first one is for the model)
 	assert.Equal(t, 2, len(state.PrefixHashes), "number of hashes is incorrect")
 	assert.Equal(t, 0, len(state.PrefixCacheServers), "there shouldn't be any cached servers")
-
-	// Updated to use the new Score method signature
-	scores := plugin.Score(ctx, pods)
 	assert.Equal(t, float64(0), scores[pod1], "score for pod1")
 	assert.Equal(t, float64(0), scores[pod2], "score for pod2")
 
@@ -52,17 +49,14 @@ func TestPrefixPlugin(t *testing.T) {
 		Prompt:      "bbbbbb",
 	}
 	ctx = types.NewSchedulingContext(context.Background(), req2, nil, pods)
-	plugin.PreCycle(ctx)
+	scores = plugin.Score(ctx, pods)
 	state, err = plugin.getPrefixState(ctx.CycleState)
 	assert.NoError(t, err)
 	t.Logf("Hashes %+v, cached servers: %+v", state.PrefixHashes, state.PrefixCacheServers)
 	// Input size is 6, hash block size is 4, the last 2 characters are ignored.
 	// Total hashes = 2 (the first one is for the model)
 	assert.Equal(t, 2, len(state.PrefixHashes), "number of hashes is incorrect")
 	assert.Equal(t, 0, len(state.PrefixCacheServers), "there shouldn't be any cached servers")
-
-	// Updated to use the new Score method signature
-	scores = plugin.Score(ctx, pods)
 	assert.Equal(t, float64(0), scores[pod1], "score for pod1")
 	assert.Equal(t, float64(0), scores[pod2], "score for pod2")
 
@@ -75,17 +69,14 @@ func TestPrefixPlugin(t *testing.T) {
 		Prompt:      "aaaabbbb",
 	}
 	ctx = types.NewSchedulingContext(context.Background(), req3, nil, pods)
-	plugin.PreCycle(ctx)
+	scores = plugin.Score(ctx, pods)
 	state, err = plugin.getPrefixState(ctx.CycleState)
 	assert.NoError(t, err)
 	t.Logf("Hashes %+v, cached servers: %+v", state.PrefixHashes, state.PrefixCacheServers)
 	// Input size is 8, hash block size is 4, so 2 hashes will be calculated.
 	// Total hashes = 3 (the first one is for the model)
 	assert.Equal(t, 3, len(state.PrefixHashes), "number of hashes is incorrect")
 	assert.Equal(t, 1, len(state.PrefixCacheServers), "pod1 should have cached the aaaa prefix")
-
-	// Updated to use the new Score method signature
-	scores = plugin.Score(ctx, pods)
 	assert.Equal(t, float64(2)/float64(3), scores[pod1], "score should be 2/3 - the model and the first prefix block match")
 	assert.Equal(t, float64(0), scores[pod2], "score for pod2")
 
@@ -97,17 +88,14 @@ func TestPrefixPlugin(t *testing.T) {
 		Prompt:      "aaaabbbb",
 	}
 	ctx = types.NewSchedulingContext(context.Background(), req4, nil, pods)
-	plugin.PreCycle(ctx)
+	scores = plugin.Score(ctx, pods)
 	state, err = plugin.getPrefixState(ctx.CycleState)
 	assert.NoError(t, err)
 	t.Logf("Hashes %+v, cached servers: %+v", state.PrefixHashes, state.PrefixCacheServers)
 	// Input size is 8, hash block size is 4, so 2 hashes will be calculated.
 	// Total hashes = 3 (the first one is for the model)
 	assert.Equal(t, 3, len(state.PrefixHashes), "number of hashes is incorrect")
 	assert.Equal(t, 0, len(state.PrefixCacheServers), "pod1 should have cached the aaaa prefix")
-
-	// Updated to use the new Score method signature
-	scores = plugin.Score(ctx, pods)
 	assert.Equal(t, float64(0), scores[pod1], "score for pod1")
 	assert.Equal(t, float64(0), scores[pod2], "score for pod2")
 
@@ -119,17 +107,14 @@ func TestPrefixPlugin(t *testing.T) {
 		Prompt:      "aaaabbbbcccc",
 	}
 	ctx = types.NewSchedulingContext(context.Background(), req5, nil, pods)
-	plugin.PreCycle(ctx)
+	scores = plugin.Score(ctx, pods)
 	state, err = plugin.getPrefixState(ctx.CycleState)
 	assert.NoError(t, err)
 	t.Logf("Hashes %+v, cached servers: %+v", state.PrefixHashes, state.PrefixCacheServers)
 	// Input size is 12, hash block size is 4, so 3 hashes will be calculated.
 	// Total hashes = 4 (the first one is for the model)
 	assert.Equal(t, 4, len(state.PrefixHashes), "number of hashes is incorrect")
 	assert.Equal(t, 1, len(state.PrefixCacheServers), "pod1 should have cached the aaaa prefix")
-
-	// Updated to use the new Score method signature
-	scores = plugin.Score(ctx, pods)
 	assert.Equal(t, 0.75, scores[pod1], "score should be 0.75 - the model and the first 2 prefix blocks match")
 	assert.Equal(t, float64(0), scores[pod2], "score for pod2")
 

pkg/epp/scheduling/framework/scheduler_profile.go

@@ -29,7 +29,6 @@ import (
 // NewSchedulerProfile creates a new SchedulerProfile object and returns its pointer.
 func NewSchedulerProfile() *SchedulerProfile {
 	return &SchedulerProfile{
-		preCyclePlugins:     []PreCycle{},
 		filters:             []Filter{},
 		scorers:             []*WeightedScorer{},
 		postCyclePlugins:    []PostCycle{},
@@ -40,21 +39,13 @@ func NewSchedulerProfile() *SchedulerProfile {
 
 // SchedulerProfile provides a profile configuration for the scheduler which influence routing decisions.
 type SchedulerProfile struct {
-	preCyclePlugins     []PreCycle
 	filters             []Filter
 	scorers             []*WeightedScorer
 	picker              Picker
 	postCyclePlugins    []PostCycle
 	PostResponsePlugins []PostResponse // TODO this field should get out of the scheduler
 }
 
-// WithPreCyclePlugins sets the given plugins as the PreCycle plugins.
-// If the SchedulerProfile has PreCycle plugins, this call replaces the existing plugins with the given ones.
-func (p *SchedulerProfile) WithPreCyclePlugins(plugins ...PreCycle) *SchedulerProfile {
-	p.preCyclePlugins = plugins
-	return p
-}
-
 // WithFilters sets the given filter plugins as the Filter plugins.
 // if the SchedulerProfile has Filter plugins, this call replaces the existing plugins with the given ones.
 func (p *SchedulerProfile) WithFilters(filters ...Filter) *SchedulerProfile {
@@ -96,9 +87,6 @@ func (p *SchedulerProfile) AddPlugins(pluginObjects ...Plugin) error {
 		} else if scorer, ok := plugin.(Scorer); ok { // if we got a Scorer instead of WeightedScorer that's an error.
 			return fmt.Errorf("failed to register scorer '%s' without a weight. follow function documentation to register a scorer", scorer.Name())
 		}
-		if preCyclePlugin, ok := plugin.(PreCycle); ok {
-			p.preCyclePlugins = append(p.preCyclePlugins, preCyclePlugin)
-		}
 		if filter, ok := plugin.(Filter); ok {
 			p.filters = append(p.filters, filter)
 		}
@@ -119,10 +107,8 @@ func (p *SchedulerProfile) AddPlugins(pluginObjects ...Plugin) error {
 }
 
 // RunCycle runs a SchedulerProfile cycle. In other words, it invokes all the SchedulerProfile plugins in this
-// order - PreCyclePlugins, Filters, Scorers, Picker, PostCyclePlugins. After completing all, it returns the result.
+// order - Filters, Scorers, Picker, PostCyclePlugins. After completing all, it returns the result.
 func (p *SchedulerProfile) RunCycle(ctx *types.SchedulingContext) (*types.Result, error) {
-	p.runPreCyclePlugins(ctx)
-
 	pods := p.runFilterPlugins(ctx)
 	if len(pods) == 0 {
 		return nil, errutil.Error{Code: errutil.Internal, Msg: "no pods available for the given request"}
@@ -137,15 +123,6 @@ func (p *SchedulerProfile) RunCycle(ctx *types.SchedulingContext) (*types.Result
 	return result, nil
 }
 
-func (p *SchedulerProfile) runPreCyclePlugins(ctx *types.SchedulingContext) {
-	for _, plugin := range p.preCyclePlugins {
-		ctx.Logger.V(logutil.DEBUG).Info("Running pre-cycle plugin", "plugin", plugin.Name())
-		before := time.Now()
-		plugin.PreCycle(ctx)
-		metrics.RecordSchedulerPluginProcessingLatency(PreCyclePluginType, plugin.Name(), time.Since(before))
-	}
-}
-
 func (p *SchedulerProfile) runFilterPlugins(ctx *types.SchedulingContext) []types.Pod {
 	loggerDebug := ctx.Logger.V(logutil.DEBUG)
 	filteredPods := ctx.PodsSnapshot

pkg/epp/scheduling/framework/scheduler_profile_test.go

@@ -61,7 +61,6 @@ func TestSchedulePlugins(t *testing.T) {
 		{
 			name: "all plugins executed successfully, all scorers with same weight",
 			profile: NewSchedulerProfile().
-				WithPreCyclePlugins(tp1, tp2).
 				WithFilters(tp1, tp2).
 				WithScorers(NewWeightedScorer(tp1, 1), NewWeightedScorer(tp2, 1)).
 				WithPicker(pickerPlugin).
@@ -79,7 +78,6 @@ func TestSchedulePlugins(t *testing.T) {
 		{
 			name: "all plugins executed successfully, different scorers weights",
 			profile: NewSchedulerProfile().
-				WithPreCyclePlugins(tp1, tp2).
 				WithFilters(tp1, tp2).
 				WithScorers(NewWeightedScorer(tp1, 60), NewWeightedScorer(tp2, 40)).
 				WithPicker(pickerPlugin).
@@ -97,7 +95,6 @@ func TestSchedulePlugins(t *testing.T) {
 		{
 			name: "filter all",
 			profile: NewSchedulerProfile().
-				WithPreCyclePlugins(tp1, tp2).
 				WithFilters(tp1, tp_filterAll).
 				WithScorers(NewWeightedScorer(tp1, 1), NewWeightedScorer(tp2, 1)).
 				WithPicker(pickerPlugin).
@@ -115,9 +112,6 @@ func TestSchedulePlugins(t *testing.T) {
 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
 			// Reset all plugins before each new test case.
-			for _, plugin := range test.profile.preCyclePlugins {
-				plugin.(*testPlugin).reset()
-			}
 			for _, plugin := range test.profile.filters {
 				plugin.(*testPlugin).reset()
 			}
@@ -159,12 +153,6 @@ func TestSchedulePlugins(t *testing.T) {
 				t.Errorf("Unexpected output (-want +got): %v", diff)
 			}
 			// Validate plugin execution counts dynamically
-			for _, plugin := range test.profile.preCyclePlugins {
-				tp, _ := plugin.(*testPlugin)
-				if tp.PreScheduleCallCount != 1 {
-					t.Errorf("Plugin %s PreSchedule() called %d times, expected 1", plugin.Name(), tp.PreScheduleCallCount)
-				}
-			}
 			for _, plugin := range test.profile.filters {
 				tp, _ := plugin.(*testPlugin)
 				if tp.FilterCallCount != 1 {
@@ -200,6 +188,12 @@ func TestSchedulePlugins(t *testing.T) {
 	}
 }
 
+// compile-time type assertion
+var _ Filter = &testPlugin{}
+var _ Scorer = &testPlugin{}
+var _ Picker = &testPlugin{}
+var _ PostCycle = &testPlugin{}
+
 // testPlugin is an implementation useful in unit tests.
 type testPlugin struct {
 	NameRes               string
@@ -208,7 +202,6 @@ type testPlugin struct {
 	ScoreRes              float64
 	FilterCallCount       int
 	FilterRes             []k8stypes.NamespacedName
-	PreScheduleCallCount  int
 	PostScheduleCallCount int
 	PickCallCount         int
 	NumOfPickerCandidates int
@@ -218,10 +211,6 @@ type testPlugin struct {
 
 func (tp *testPlugin) Name() string { return tp.NameRes }
 
-func (tp *testPlugin) PreCycle(ctx *types.SchedulingContext) {
-	tp.PreScheduleCallCount++
-}
-
 func (tp *testPlugin) Filter(ctx *types.SchedulingContext, pods []types.Pod) []types.Pod {
 	tp.FilterCallCount++
 	return findPods(ctx, tp.FilterRes...)
@@ -251,7 +240,6 @@ func (tp *testPlugin) PostCycle(ctx *types.SchedulingContext, res *types.Result)
 }
 
 func (tp *testPlugin) reset() {
-	tp.PreScheduleCallCount = 0
 	tp.FilterCallCount = 0
 	tp.ScoreCallCount = 0
 	tp.NumOfScoredPods = 0