fix: add validation for load-aware scorer to handle invalid queue thresholds (#240)

Signed-off-by: Kfir Toledo <[email protected]>

Author: kfirtoledo

pkg/plugins/scorer/load_aware_scorer.go

@@ -5,9 +5,11 @@ import (
 	"encoding/json"
 	"fmt"
 
+	"sigs.k8s.io/controller-runtime/pkg/log"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/plugins"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/scheduling/framework"
 	"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/scheduling/types"
+	logutil "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/util/logging"
 )
 
 const (
@@ -26,19 +28,24 @@ type loadAwareScorerParameters struct {
 var _ framework.Scorer = &LoadAwareScorer{}
 
 // LoadAwareScorerFactory defines the factory function for the LoadAwareScorer
-func LoadAwareScorerFactory(name string, rawParameters json.RawMessage, _ plugins.Handle) (plugins.Plugin, error) {
+func LoadAwareScorerFactory(name string, rawParameters json.RawMessage, handle plugins.Handle) (plugins.Plugin, error) {
 	parameters := loadAwareScorerParameters{Threshold: QueueThresholdDefault}
 	if rawParameters != nil {
 		if err := json.Unmarshal(rawParameters, &parameters); err != nil {
 			return nil, fmt.Errorf("failed to parse the parameters of the '%s' scorer - %w", LoadAwareScorerType, err)
 		}
 	}
 
-	return NewLoadAwareScorer(parameters.Threshold).WithName(name), nil
+	return NewLoadAwareScorer(handle.Context(), parameters.Threshold).WithName(name), nil
 }
 
 // NewLoadAwareScorer creates a new load based scorer
-func NewLoadAwareScorer(queueThreshold int) *LoadAwareScorer {
+func NewLoadAwareScorer(ctx context.Context, queueThreshold int) *LoadAwareScorer {
+	if queueThreshold <= 0 {
+		queueThreshold = QueueThresholdDefault
+		log.FromContext(ctx).V(logutil.DEFAULT).Info(fmt.Sprintf("queueThreshold %d should be positive, using default queue threshold %d", queueThreshold, QueueThresholdDefault))
+	}
+
 	return &LoadAwareScorer{
 		typedName:      plugins.TypedName{Type: LoadAwareScorerType},
 		queueThreshold: float64(queueThreshold),

pkg/plugins/scorer/load_aware_scorer_test.go

@@ -27,7 +27,7 @@ func TestLoadBasedScorer(t *testing.T) {
 	}{
 		{
 			name:   "load based scorer",
-			scorer: scorer.NewLoadAwareScorer(0),
+			scorer: scorer.NewLoadAwareScorer(context.Background(), 10),
 			req: &types.LLMRequest{
 				TargetModel: "critical",
 			},

pkg/scheduling/pd/scheduler_test.go

@@ -195,7 +195,7 @@ func TestPDSchedule(t *testing.T) {
 
 			decodeSchedulerProfile := framework.NewSchedulerProfile().
 				WithFilters(filter.NewDecodeFilter()).
-				WithScorers(framework.NewWeightedScorer(scorer.NewLoadAwareScorer(scorer.QueueThresholdDefault), 1)).
+				WithScorers(framework.NewWeightedScorer(scorer.NewLoadAwareScorer(ctx, scorer.QueueThresholdDefault), 1)).
 				WithPicker(picker.NewMaxScorePicker())
 			err = decodeSchedulerProfile.AddPlugins(framework.NewWeightedScorer(prefixScorer, 0))
 			assert.NoError(t, err, "SchedulerProfile AddPlugins returned unexpected error")