feat(qrm): support getting specified pool name in QoS conf

kubewharf · Dec 12, 2023 · bff65a5 · bff65a5
1 parent 0a516df
commit bff65a5
Show file tree

Hide file tree

Showing 4 changed files with 90 additions and 14 deletions.
diff --git a/pkg/agent/qrm-plugins/cpu/dynamicpolicy/state/state.go b/pkg/agent/qrm-plugins/cpu/dynamicpolicy/state/state.go
@@ -161,20 +161,7 @@ func (ai *AllocationInfo) GetSpecifiedPoolName() string {
 		return cpuadvisor.EmptyOwnerPoolName
 	}
 
-	switch ai.QoSLevel {
-	case consts.PodAnnotationQoSLevelSharedCores:
-		specifiedPoolName := ai.Annotations[consts.PodAnnotationCPUEnhancementCPUSet]
-		if specifiedPoolName != cpuadvisor.EmptyOwnerPoolName {
-			return specifiedPoolName
-		}
-		return PoolNameShare
-	case consts.PodAnnotationQoSLevelReclaimedCores:
-		return PoolNameReclaim
-	case consts.PodAnnotationQoSLevelDedicatedCores:
-		return PoolNameDedicated
-	default:
-		return cpuadvisor.EmptyOwnerPoolName
-	}
+	return GetSpecifiedPoolName(ai.QoSLevel, ai.Annotations[consts.PodAnnotationCPUEnhancementCPUSet])
 }
 
 // CheckMainContainer returns true if the AllocationInfo is for main container

diff --git a/pkg/agent/qrm-plugins/cpu/dynamicpolicy/state/util.go b/pkg/agent/qrm-plugins/cpu/dynamicpolicy/state/util.go
@@ -24,7 +24,9 @@ import (
 	"k8s.io/apimachinery/pkg/util/sets"
 	"k8s.io/klog/v2"
 
+	apiconsts "github.com/kubewharf/katalyst-api/pkg/consts"
 	"github.com/kubewharf/katalyst-core/pkg/agent/qrm-plugins/cpu/consts"
+	"github.com/kubewharf/katalyst-core/pkg/agent/qrm-plugins/cpu/dynamicpolicy/cpuadvisor"
 	advisorapi "github.com/kubewharf/katalyst-core/pkg/agent/qrm-plugins/cpu/dynamicpolicy/cpuadvisor"
 	"github.com/kubewharf/katalyst-core/pkg/util/machine"
 )
@@ -223,3 +225,19 @@ func GetPoolType(poolName string) string {
 		return PoolNameShare
 	}
 }
+
+func GetSpecifiedPoolName(qosLevel, cpusetEnhancementValue string) string {
+	switch qosLevel {
+	case apiconsts.PodAnnotationQoSLevelSharedCores:
+		if cpusetEnhancementValue != cpuadvisor.EmptyOwnerPoolName {
+			return cpusetEnhancementValue
+		}
+		return PoolNameShare
+	case apiconsts.PodAnnotationQoSLevelReclaimedCores:
+		return PoolNameReclaim
+	case apiconsts.PodAnnotationQoSLevelDedicatedCores:
+		return PoolNameDedicated
+	default:
+		return cpuadvisor.EmptyOwnerPoolName
+	}
+}
diff --git a/pkg/agent/qrm-plugins/cpu/dynamicpolicy/state/util_test.go b/pkg/agent/qrm-plugins/cpu/dynamicpolicy/state/util_test.go
@@ -509,3 +509,52 @@ func TestGenerateCPUMachineStateByPodEntries(t *testing.T) {
 		os.RemoveAll(tmpDir)
 	}
 }
+
+func TestGetSpecifiedPoolName(t *testing.T) {
+	type args struct {
+		qosLevel               string
+		cpusetEnhancementValue string
+	}
+	tests := []struct {
+		name string
+		args args
+		want string
+	}{
+		{
+			name: "shared_cores with empty cpusetEnhancementValue",
+			args: args{
+				qosLevel: consts.PodAnnotationQoSLevelSharedCores,
+			},
+			want: PoolNameShare,
+		},
+		{
+			name: "shared_cores with non-empty cpusetEnhancementValue",
+			args: args{
+				qosLevel:               consts.PodAnnotationQoSLevelSharedCores,
+				cpusetEnhancementValue: "offline",
+			},
+			want: "offline",
+		},
+		{
+			name: "dedicated_cores with empty cpusetEnhancementValue",
+			args: args{
+				qosLevel: consts.PodAnnotationQoSLevelDedicatedCores,
+			},
+			want: PoolNameDedicated,
+		},
+		{
+			name: "reclaimed_cores with empty cpusetEnhancementValue",
+			args: args{
+				qosLevel: consts.PodAnnotationQoSLevelReclaimedCores,
+			},
+			want: PoolNameReclaim,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if got := GetSpecifiedPoolName(tt.args.qosLevel, tt.args.cpusetEnhancementValue); got != tt.want {
+				t.Errorf("GetSpecifiedPoolName() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}
diff --git a/pkg/config/generic/qos.go b/pkg/config/generic/qos.go
@@ -24,6 +24,7 @@ import (
 	"k8s.io/apimachinery/pkg/util/sets"
 
 	apiconsts "github.com/kubewharf/katalyst-api/pkg/consts"
+	"github.com/kubewharf/katalyst-core/pkg/agent/qrm-plugins/cpu/dynamicpolicy/state"
 	"github.com/kubewharf/katalyst-core/pkg/util/general"
 	"github.com/kubewharf/katalyst-core/pkg/util/qos/helper"
 )
@@ -354,6 +355,27 @@ func (c *QoSConfiguration) checkQosMatched(annotations map[string]string, qosVal
 	return true, false, nil
 }
 
+func (c *QoSConfiguration) GetSpecifiedPoolNameForPod(pod *v1.Pod) (string, error) {
+	if pod == nil {
+		return "", fmt.Errorf("nil pod")
+	}
+	return c.GetSpecifiedPoolName(c.GetQoSEnhancementsForPod(pod), pod.Annotations)
+}
+
+// GetSpecifiedPoolName returns the specified cpuset pool name for given enhancements and annotations;
+func (c *QoSConfiguration) GetSpecifiedPoolName(enhancements, annotations map[string]string) (string, error) {
+	qosLevel, err := c.GetQoSLevel(annotations)
+
+	if err != nil {
+		return "", fmt.Errorf("GetQoSLevel failed with error: %v", err)
+	}
+
+	enhancementKVs := helper.ParseKatalystQOSEnhancement(enhancements, annotations,
+		apiconsts.PodAnnotationCPUEnhancementKey)
+
+	return state.GetSpecifiedPoolName(qosLevel, enhancementKVs[apiconsts.PodAnnotationCPUEnhancementCPUSet]), nil
+}
+
 // checkKeyValueMatched checks whether the given key-value pair exists in the map
 // if the returns value equals 1, it represents
 // - key not exists