Update in filter and decode filters implementation, both are based on the ByRoleLabel + add related test in scheduler test

Signed-off-by: Maya Barnea <mayab@il.ibm.com>

Author: mayabar

pkg/plugins/filter/pd_role_filter.go

@@ -8,11 +8,8 @@ import (
 )
 
 const (
-	// DecodeFilterType is the type of the DecodeFilter
-	DecodeFilterType = "decode-filter"
-
-	// PrefillFilterType is the type of the PrefillFilter
-	PrefillFilterType = "prefill-filter"
+	// ByRoleLabelFilterType is the type of the ByLabelsFilter
+	ByRoleLabelFilterType = "role-label"
 
 	// RoleLabel name
 	RoleLabel = "llm-d.ai/role"
@@ -24,90 +21,64 @@ const (
 	RoleBoth = "both"
 )
 
-// compile-time type assertion
-var _ framework.Filter = &PrefillFilter{}
-
-// NewPrefillFilter creates a new instance of the DecodeFilter
-func NewPrefillFilter() *PrefillFilter {
-	return &PrefillFilter{
-		name: PrefillFilterType,
-	}
-}
-
-// PrefillFilter - filters out pods that are not marked with role Prefill
-type PrefillFilter struct {
+// ByRoleLabel - filters out pods based on the role defined by RoleLabel label
+type ByRoleLabel struct {
+	// name defines the filter name
 	name string
+	// validRoles defines list of valid role header values
+	validRoles map[string]bool
+	// allowsNoRolesLabel - if true pods without role label will be considered as valid (not filtered out)
+	allowsNoRolesLabel bool
 }
 
-// Type returns the type of the filter
-func (pf *PrefillFilter) Type() string {
-	return PrefillFilterType
-}
-
-// Name returns the name of the instance of the filter.
-func (pf *PrefillFilter) Name() string {
-	return pf.name
-}
-
-// WithName sets the name of the filter.
-func (pf *PrefillFilter) WithName(name string) *PrefillFilter {
-	pf.name = name
-	return pf
-}
+var _ framework.Filter = &ByRoleLabel{} // validate interface conformance
 
-// Filter filters out all pods that are not marked as "prefill"
-func (pf *PrefillFilter) Filter(_ context.Context, _ *types.CycleState, _ *types.LLMRequest, pods []types.Pod) []types.Pod {
-	filteredPods := []types.Pod{}
+// NewByRoleLabel creates and returns an instance of the RoleBasedFilter based on the input parameters
+// name - the filter name
+// rolesArr - list of valid roles
+func NewByRoleLabel(name string, allowsNoRolesLabel bool, rolesArr ...string) *ByRoleLabel {
+	roles := map[string]bool{}
 
-	for _, pod := range pods {
-		role := pod.GetPod().Labels[RoleLabel]
-		if role == RolePrefill { // TODO: doesn't RoleBoth also imply Prefill?
-			filteredPods = append(filteredPods, pod)
-		}
+	for _, role := range rolesArr {
+		roles[role] = true
 	}
-	return filteredPods
-}
-
-// compile-time type assertion
-var _ framework.Filter = &DecodeFilter{}
 
-// NewDecodeFilter creates a new instance of the DecodeFilter
-func NewDecodeFilter() *DecodeFilter {
-	return &DecodeFilter{
-		name: DecodeFilterType,
-	}
+	return &ByRoleLabel{name: name, allowsNoRolesLabel: allowsNoRolesLabel, validRoles: roles}
 }
 
-// DecodeFilter - filters out pods that are not marked with role Decode or Both
-type DecodeFilter struct {
-	name string
+// NewPrefillFilter creates and returns an instance of the Filter configured for prefill role
+func NewPrefillFilter() framework.Filter {
+	return NewByRoleLabel("prefill-filter", false, RolePrefill)
 }
 
-// Type returns the type of the filter
-func (df *DecodeFilter) Type() string {
-	return DecodeFilterType
+// NewDecodeFilter creates and returns an instance of the Filter configured for decode role
+func NewDecodeFilter() framework.Filter {
+	return NewByRoleLabel("decode-filter", true, RoleDecode, RoleBoth)
 }
 
-// Name returns the name of the instance of the filter.
-func (df *DecodeFilter) Name() string {
-	return df.name
+// Type returns the type of the filter
+func (f *ByRoleLabel) Type() string {
+	return ByRoleLabelFilterType
 }
 
-// WithName sets the name of the filter.
-func (df *DecodeFilter) WithName(name string) *DecodeFilter {
-	df.name = name
-	return df
+// Name returns the name of the filter
+func (f *ByRoleLabel) Name() string {
+	return f.name
 }
 
-// Filter removes all pods that are not marked as "decode" or "both"
-func (df *DecodeFilter) Filter(_ context.Context, _ *types.CycleState, _ *types.LLMRequest, pods []types.Pod) []types.Pod {
+// Filter filters out all pods that are not marked with one of roles from the validRoles collection
+// or has no role label in case allowsNoRolesLabel is true
+func (f *ByRoleLabel) Filter(_ context.Context, _ *types.CycleState, _ *types.LLMRequest, pods []types.Pod) []types.Pod {
 	filteredPods := []types.Pod{}
 
 	for _, pod := range pods {
-		role, defined := pod.GetPod().Labels[RoleLabel]
-		if !defined || role == RoleDecode || role == RoleBoth {
+		role, labelDefined := pod.GetPod().Labels[RoleLabel]
+		_, roleExists := f.validRoles[role]
+
+		if (!labelDefined && f.allowsNoRolesLabel) || roleExists {
 			filteredPods = append(filteredPods, pod)
 		}
 	}
+
 	return filteredPods
 }

pkg/scheduling/pd/scheduler.go

@@ -31,21 +31,13 @@ func CreatePDSchedulerConfig(ctx context.Context, pdConfig *config.Config, prefi
 	// otherwise, PD is enabled.
 
 	// create decode scheduling profile.
-	decodeFilter, err := filter.NewDecodeFilter()
-	if err != nil {
-		return nil, err
-	}
-	decodeProfile, err := createSchedulerProfile(ctx, decodeFilter, picker.NewMaxScorePicker(), pdConfig.DecodeSchedulerPlugins, prefixScorer, true)
+	decodeProfile, err := createSchedulerProfile(ctx, filter.NewDecodeFilter(), picker.NewMaxScorePicker(), pdConfig.DecodeSchedulerPlugins, prefixScorer, true)
 	if err != nil {
 		return nil, fmt.Errorf("falied to create decode scheduling profile - %w", err)
 	}
 
 	// create prefil scheduling profile.
-	prefillFilter, err := filter.NewPrefillFilter()
-	if err != nil {
-		return nil, err
-	}
-	prefilProfile, err := createSchedulerProfile(ctx, prefillFilter, picker.NewMaxScorePicker(), pdConfig.PrefillSchedulerPlugins, prefixScorer, true)
+	prefilProfile, err := createSchedulerProfile(ctx, filter.NewPrefillFilter(), picker.NewMaxScorePicker(), pdConfig.PrefillSchedulerPlugins, prefixScorer, true)
 	if err != nil {
 		return nil, fmt.Errorf("falied to create prefill scheduling profile - %w", err)
 	}

pkg/scheduling/pd/scheduler_test.go

@@ -60,13 +60,30 @@ func TestPDSchedule(t *testing.T) {
 			WaitingModels: map[string]int{},
 		},
 	}
+	noRolePod1 := &backendmetrics.FakePodMetrics{
+		Pod: &backend.Pod{
+			NamespacedName: k8stypes.NamespacedName{Name: "noRolePod1"},
+			Address:        "1.1.1.1",
+		},
+		Metrics: &backendmetrics.MetricsState{},
+	}
+	noRolePod2 := &backendmetrics.FakePodMetrics{
+		Pod: &backend.Pod{
+			NamespacedName: k8stypes.NamespacedName{Name: "noRolePod2"},
+			Address:        "2.2.2.2",
+		},
+		Metrics: &backendmetrics.MetricsState{},
+	}
 
 	tests := []struct {
-		name    string
-		req     *types.LLMRequest
-		input   []backendmetrics.PodMetrics
-		wantRes *types.SchedulingResult
-		err     bool
+		name            string
+		req             *types.LLMRequest
+		input           []backendmetrics.PodMetrics
+		wantRes         *types.SchedulingResult
+		wantHeaders     map[string]string
+		unwantedHeaders []string
+		unwantedPodIDs  []string
+		err             bool
 	}{
 		{
 			name: "no pods in datastore",
@@ -153,6 +170,16 @@ func TestPDSchedule(t *testing.T) {
 				PrimaryProfileName: "decode",
 			},
 		},
+		{
+			name: "TestRoles",
+			req: &types.LLMRequest{
+				TargetModel: "critical",
+				Prompt:      "12345678901",
+			},
+			input:          []backendmetrics.PodMetrics{pod1, noRolePod1, noRolePod2},
+			wantRes:        nil, // doesn't mater which pod was selected
+			unwantedPodIDs: []string{pod1.GetPod().NamespacedName.String()},
+		},
 	}
 
 	ctx := context.Background()
@@ -189,8 +216,37 @@ func TestPDSchedule(t *testing.T) {
 				t.Errorf("Unexpected error, got %v, want %v", err, test.err)
 			}
 
-			if diff := cmp.Diff(test.wantRes, got); diff != "" {
-				t.Errorf("Unexpected output (-want +got): %v", diff)
+			if test.wantRes != nil {
+				if diff := cmp.Diff(test.wantRes, got); diff != "" {
+					t.Errorf("Unexpected output (-want +got): %v", diff)
+				}
+
+				for header, value := range test.wantHeaders {
+					gotValue, ok := test.req.Headers[header]
+					if !ok {
+						t.Errorf("Missing header: %s", header)
+					} else if gotValue != value {
+						t.Errorf("Wrong header value for %s: want %s got %s)", header, value, gotValue)
+					}
+				}
+
+				for _, header := range test.unwantedHeaders {
+					if _, exists := test.req.Headers[header]; exists {
+						t.Errorf("Unwanted header %s exists", header)
+					}
+				}
+			}
+
+			if len(test.unwantedPodIDs) > 0 {
+				// ensure that target pod is not one of the unwanted
+				profileRes, found := got.ProfileResults[got.PrimaryProfileName]
+				if found {
+					for _, podID := range test.unwantedPodIDs {
+						if podID == profileRes.TargetPod.GetPod().NamespacedName.String() {
+							t.Errorf("Unwanted pod was selected: %s", podID)
+						}
+					}
+				}
 			}
 		})
 	}