wip: feat: cpu numa affinity

go.mod

@@ -175,6 +175,7 @@ require (
 )
 
 replace (
+	github.com/kubewharf/katalyst-api => github.com/yehlemias/katalyst-api v0.0.0-20260123040154-7734ccb2ac3f
 	k8s.io/api => k8s.io/api v0.24.6
 	k8s.io/apiextensions-apiserver => k8s.io/apiextensions-apiserver v0.24.6
 	k8s.io/apimachinery => k8s.io/apimachinery v0.24.6

go.sum

@@ -574,8 +574,6 @@ github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ=
 github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI=
 github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
 github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE=
-github.com/kubewharf/katalyst-api v0.5.9-0.20260108125536-85e136f5902c h1:ohKHA5TOlW9487menKnKH2M14LeIq1xQ1yW4xp8x9o8=
-github.com/kubewharf/katalyst-api v0.5.9-0.20260108125536-85e136f5902c/go.mod h1:BZMVGVl3EP0eCn5xsDgV41/gjYkoh43abIYxrB10e3k=
 github.com/kubewharf/kubelet v1.24.6-kubewharf.9 h1:jOTYZt7h/J7I8xQMKMUcJjKf5UFBv37jHWvNp5VRFGc=
 github.com/kubewharf/kubelet v1.24.6-kubewharf.9/go.mod h1:MxbSZUx3wXztFneeelwWWlX7NAAStJ6expqq7gY2J3c=
 github.com/kyoh86/exportloopref v0.1.7/go.mod h1:h1rDl2Kdj97+Kwh4gdz3ujE7XHmH51Q0lUiZ1z4NLj8=
@@ -937,6 +935,8 @@ github.com/xiang90/probing v0.0.0-20190116061207-43a291ad63a2 h1:eY9dn8+vbi4tKz5
 github.com/xiang90/probing v0.0.0-20190116061207-43a291ad63a2/go.mod h1:UETIi67q53MR2AWcXfiuqkDkRtnGDLqkBTpCHuJHxtU=
 github.com/xlab/treeprint v0.0.0-20181112141820-a009c3971eca/go.mod h1:ce1O1j6UtZfjr22oyGxGLbauSBp2YVXpARAosm7dHBg=
 github.com/xordataexchange/crypt v0.0.3-0.20170626215501-b2862e3d0a77/go.mod h1:aYKd//L2LvnjZzWKhF00oedf4jCCReLcmhLdhm1A27Q=
+github.com/yehlemias/katalyst-api v0.0.0-20260123040154-7734ccb2ac3f h1:W8u1tV9FcEUze0qAA0IfMAzUOWPHoGoHv1yqAfy/Qs0=
+github.com/yehlemias/katalyst-api v0.0.0-20260123040154-7734ccb2ac3f/go.mod h1:BZMVGVl3EP0eCn5xsDgV41/gjYkoh43abIYxrB10e3k=
 github.com/yuin/goldmark v1.1.25/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
 github.com/yuin/goldmark v1.1.27/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
 github.com/yuin/goldmark v1.1.32/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=

pkg/agent/qrm-plugins/commonstate/pool.go

@@ -23,6 +23,7 @@ import (
 	apiconsts "github.com/kubewharf/katalyst-api/pkg/consts"
 	cpuconsts "github.com/kubewharf/katalyst-core/pkg/agent/qrm-plugins/cpu/consts"
 	"github.com/kubewharf/katalyst-core/pkg/util/machine"
+	qosutil "github.com/kubewharf/katalyst-core/pkg/util/qos"
 )
 
 // notice that pool-name may not have direct mapping relations with qos-level, for instance
@@ -46,10 +47,11 @@ const (
 // FakedContainerName represents a placeholder since pool entry has no container-level
 // FakedNUMAID represents a placeholder since pools like shared/reclaimed will not contain a specific numa
 const (
-	FakedContainerName = ""
-	FakedNUMAID        = -1
-	NameSeparator      = "#"
-	NUMAPoolInfix      = "-NUMA"
+	FakedContainerName    = ""
+	FakedNUMAID           = -1
+	NameSeparator         = "#"
+	NUMAPoolInfix         = "-NUMA"
+	NUMAAffinityPoolInfix = "-Affinity"
 )
 
 const (
@@ -68,6 +70,11 @@ func IsShareNUMABindingPool(poolName string) bool {
 	return strings.Contains(poolName, NUMAPoolInfix)
 }
 
+// IsShareNUMAAffinityPool checks whether the pool is numa affinity pool
+func IsShareNUMAAffinityPool(poolName string) bool {
+	return strings.Contains(poolName, NUMAAffinityPoolInfix) || IsShareNUMABindingPool(poolName)
+}
+
 func GetPoolType(poolName string) string {
 	if IsIsolationPool(poolName) {
 		return PoolNamePrefixIsolation
@@ -101,8 +108,8 @@ func GetSpecifiedPoolName(qosLevel, cpusetEnhancementValue string) string {
 	}
 }
 
-// GetSpecifiedNUMABindingNUMAID parses the numa id for AllocationInfo
-func GetSpecifiedNUMABindingNUMAID(annotations map[string]string) (int, error) {
+// GetSpecifiedNUMAID parses the numa id for AllocationInfo
+func GetSpecifiedNUMAID(annotations map[string]string) (int, error) {
 	if _, ok := annotations[cpuconsts.CPUStateAnnotationKeyNUMAHint]; !ok {
 		return FakedNUMAID, nil
 	}
@@ -118,15 +125,17 @@ func GetSpecifiedNUMABindingNUMAID(annotations map[string]string) (int, error) {
 	return numaSet.ToSliceNoSortInt()[0], nil
 }
 
-// GetSpecifiedNUMABindingPoolName get numa_binding pool name
-// for numa_binding shared_cores according to enhancements and NUMA hint
-func GetSpecifiedNUMABindingPoolName(qosLevel string, annotations map[string]string) (string, error) {
+// GetSpecifiedNUMAPoolName get numa pool name:
+// - NUMAPoolInfix for numa_binding shared_cores, e.g. share-NUMA0
+// - NUMAAffinityPoolInfix for numa_affinity shared_cores, e.g. share-Affinity0
+func GetSpecifiedNUMAPoolName(qosLevel string, annotations map[string]string) (string, error) {
 	if qosLevel != apiconsts.PodAnnotationQoSLevelSharedCores ||
-		annotations[apiconsts.PodAnnotationMemoryEnhancementNumaBinding] != apiconsts.PodAnnotationMemoryEnhancementNumaBindingEnable {
-		return EmptyOwnerPoolName, fmt.Errorf("GetSpecifiedNUMABindingPoolName is only for numa_binding shared_cores")
+		(annotations[apiconsts.PodAnnotationMemoryEnhancementNumaBinding] != apiconsts.PodAnnotationMemoryEnhancementNumaBindingEnable &&
+			annotations[apiconsts.PodAnnotationCPUEnhancementNumaAffinity] != apiconsts.PodAnnotationCPUEnhancementNumaAffinityEnable) {
+		return EmptyOwnerPoolName, fmt.Errorf("GetSpecifiedNUMAPoolName is only for numa_binding or numa_affinity shared_cores")
 	}
 
-	numaID, err := GetSpecifiedNUMABindingNUMAID(annotations)
+	numaID, err := GetSpecifiedNUMAID(annotations)
 	if err != nil {
 		return EmptyOwnerPoolName, err
 	}
@@ -141,5 +150,10 @@ func GetSpecifiedNUMABindingPoolName(qosLevel string, annotations map[string]str
 		return EmptyOwnerPoolName, fmt.Errorf("empty specifiedPoolName")
 	}
 
-	return GetNUMAPoolName(specifiedPoolName, numaID), nil
+	numaPoolInfix := NUMAAffinityPoolInfix
+	if qosutil.AnnotationsIndicateNUMABinding(annotations) {
+		numaPoolInfix = NUMAPoolInfix
+	}
+
+	return GetNUMAPoolName(specifiedPoolName, numaPoolInfix, numaID), nil
 }

pkg/agent/qrm-plugins/commonstate/state.go

@@ -136,17 +136,17 @@ func (am *AllocationMeta) GetSpecifiedPoolName() string {
 	return GetSpecifiedPoolName(am.QoSLevel, am.Annotations[consts.PodAnnotationCPUEnhancementCPUSet])
 }
 
-// GetSpecifiedNUMABindingNUMAID parses the numa id for AllocationInfo
-func (am *AllocationMeta) GetSpecifiedNUMABindingNUMAID() (int, error) {
+// GetSpecifiedNUMAID parses the numa id for AllocationInfo
+func (am *AllocationMeta) GetSpecifiedNUMAID() (int, error) {
 	if am == nil {
 		return FakedNUMAID, fmt.Errorf("empty am")
 	}
 
-	return GetSpecifiedNUMABindingNUMAID(am.Annotations)
+	return GetSpecifiedNUMAID(am.Annotations)
 }
 
-// SetSpecifiedNUMABindingNUMAID set the numa id for AllocationInfo
-func (am *AllocationMeta) SetSpecifiedNUMABindingNUMAID(numaID uint64) {
+// SetSpecifiedNUMAID set the numa id for AllocationInfo
+func (am *AllocationMeta) SetSpecifiedNUMAID(numaID uint64) {
 	if am == nil {
 		return
 	}
@@ -158,14 +158,14 @@ func (am *AllocationMeta) SetSpecifiedNUMABindingNUMAID(numaID uint64) {
 	am.Annotations[cpuconsts.CPUStateAnnotationKeyNUMAHint] = machine.NewCPUSet(int(numaID)).String()
 }
 
-// GetSpecifiedNUMABindingPoolName get numa_binding pool name
-// for numa_binding shared_cores according to enhancements and NUMA hint
-func (am *AllocationMeta) GetSpecifiedNUMABindingPoolName() (string, error) {
-	return GetSpecifiedNUMABindingPoolName(am.QoSLevel, am.Annotations)
+// GetSpecifiedNUMAPoolName get numa pool name
+// for numa_binding or numa_affinity shared_cores according to enhancements and NUMA hint
+func (am *AllocationMeta) GetSpecifiedNUMAPoolName() (string, error) {
+	return GetSpecifiedNUMAPoolName(am.QoSLevel, am.Annotations)
 }
 
-func GetNUMAPoolName(candidateSpecifiedPoolName string, targetNUMANode int) string {
-	return fmt.Sprintf("%s%s%d", candidateSpecifiedPoolName, NUMAPoolInfix, targetNUMANode)
+func GetNUMAPoolName(candidateSpecifiedPoolName, poolInfix string, targetNUMANode int) string {
+	return fmt.Sprintf("%s%s%d", candidateSpecifiedPoolName, poolInfix, targetNUMANode)
 }
 
 // CheckMainContainer returns true if the AllocationInfo is for main container
@@ -316,3 +316,84 @@ func (am *AllocationMeta) CheckDedicatedPool() bool {
 	}
 	return am.OwnerPoolName == PoolNameDedicated
 }
+
+// CheckNUMAAffinity returns true if the AllocationInfo is for pod with numa-affinity enhancement
+func (am *AllocationMeta) CheckNUMAAffinity() bool {
+	if am == nil {
+		return false
+	}
+	return am.Annotations[consts.PodAnnotationCPUEnhancementNumaAffinity] ==
+		consts.PodAnnotationCPUEnhancementNumaAffinityEnable || am.CheckNUMABinding()
+}
+
+// CheckDedicatedNUMAAffinity returns true if the AllocationInfo is for pod with
+// dedicated-qos and numa-affinity enhancement
+func (am *AllocationMeta) CheckDedicatedNUMAAffinity() bool {
+	return am.CheckDedicated() && am.CheckNUMAAffinity()
+}
+
+// CheckSharedNUMAAffinity returns true if the AllocationInfo is for pod with
+// shared-qos and numa-affinity enhancement
+func (am *AllocationMeta) CheckSharedNUMAAffinity() bool {
+	return am.CheckShared() && am.CheckNUMAAffinity()
+}
+
+// CheckSharedOrDedicatedNUMAAffinity returns true if the AllocationInfo is for pod with
+// shared-qos or dedicated-qos and numa-affinity enhancement
+func (am *AllocationMeta) CheckSharedOrDedicatedNUMAAffinity() bool {
+	return am.CheckSharedNUMAAffinity() || am.CheckDedicatedNUMAAffinity()
+}
+
+// CheckNonBindingNUMAAffinity returns true if the AllocationInfo is for pod with
+// numa-affinity enhancement and without numa-binding enhancement
+func (am *AllocationMeta) CheckNonBindingNUMAAffinity() bool {
+	if am == nil {
+		return false
+	}
+	return am.Annotations[consts.PodAnnotationCPUEnhancementNumaAffinity] ==
+		consts.PodAnnotationCPUEnhancementNumaAffinityEnable && !am.CheckNUMABinding()
+}
+
+// CheckNonBindingSharedNUMAAffinity returns true if the AllocationInfo is for pod with
+// shared-qos and numa-affinity enhancement and without numa-binding enhancement
+func (am *AllocationMeta) CheckNonBindingSharedNUMAAffinity() bool {
+	return am.CheckShared() && am.CheckNonBindingNUMAAffinity()
+}
+
+// CheckNonBindingDedicatedNUMAAffinity returns true if the AllocationInfo is for pod with
+// dedicated-qos and numa-affinity enhancement and without numa-binding enhancement
+func (am *AllocationMeta) CheckNonBindingDedicatedNUMAAffinity() bool {
+	return am.CheckDedicated() && am.CheckNonBindingNUMAAffinity()
+}
+
+// CheckNonBindingSharedOrDedicatedNUMAAffinity returns true if the AllocationInfo is for pod with
+// shared-qos or dedicated-qos and numa-affinity enhancement and without numa-binding enhancement
+func (am *AllocationMeta) CheckNonBindingSharedOrDedicatedNUMAAffinity() bool {
+	return am.CheckNonBindingSharedNUMAAffinity() || am.CheckNonBindingDedicatedNUMAAffinity()
+}
+
+// CheckSharedNumaBindingOrDedicatedNUMAAffinity returns true if the AllocationInfo is for pod with
+// shared numa-bind or dedicated numa-affinity.
+func (am *AllocationMeta) CheckSharedNumaBindingOrDedicatedNUMAAffinity() bool {
+	return am.CheckSharedNUMABinding() || am.CheckDedicatedNUMAAffinity()
+}
+
+// CheckReclaimedNUMAAffinity returns true if the AllocationInfo is for pod with
+// reclaimed numa-bind or dedicated numa-affinity.
+func (am *AllocationMeta) CheckReclaimedNUMAAffinity() bool {
+	return am.CheckReclaimed() && am.CheckNUMAAffinity()
+}
+
+// CheckReclaimedActualNUMAAffinity returns true if the AllocationInfo is for pod with
+// reclaimed-qos and numa-affinity enhancement and numa hint is not empty, which means
+// the container is allocated on a specific NUMA node
+func (am *AllocationMeta) CheckReclaimedActualNUMAAffinity() bool {
+	return am.CheckReclaimedNUMAAffinity() && am.CheckActualNUMABinding()
+}
+
+// CheckReclaimedNonActualNUMAAffinity returns true if the AllocationInfo is for pod with
+// reclaimed-qos and numa-affinity enhancement and numa hint is empty, which means
+// the container can be allocated on multi NUMA nodes
+func (am *AllocationMeta) CheckReclaimedNonActualNUMAAffinity() bool {
+	return am.CheckReclaimed() && !am.CheckActualNUMABinding()
+}

pkg/agent/qrm-plugins/commonstate/util.go

@@ -60,6 +60,58 @@ func CheckNUMABindingAntiAffinity(meta *AllocationMeta, annotations map[string]s
 	return false
 }
 
+func CheckNUMABindingWithAffinity(meta *AllocationMeta, annotations map[string]string) bool {
+	if meta == nil {
+		return false
+	} else if len(annotations) == 0 {
+		return false
+	}
+
+	// if qos level is same as candidate, filter out this numa
+	if meta.GetQoSLevel() != annotations[consts.PodAnnotationQoSLevelKey] {
+		return false
+	}
+
+	if meta.CheckNUMANotShare() ||
+		qos.AnnotationsIndicateNUMANotShare(annotations) {
+		return false
+	}
+
+	if meta.CheckNUMABinding() {
+		// considering isolation, use specified pool instead of actual pool name here
+		candidateSpecifiedPoolName := GetSpecifiedPoolName(consts.PodAnnotationQoSLevelSharedCores,
+			annotations[consts.PodAnnotationCPUEnhancementCPUSet])
+		aiSpecifiedPoolName := meta.GetSpecifiedPoolName()
+
+		// shared_cores with numa binding doesn't support two share type pools with same specified name existing at same NUMA
+		return candidateSpecifiedPoolName == aiSpecifiedPoolName
+	}
+	return false
+}
+
+// CheckNonCPUAffinityNUMA returns true, if the AllocationMeta indicates that the meta is numa affinity candidate.
+func CheckNonCPUAffinityNUMA(meta *AllocationMeta) bool {
+	if meta == nil {
+		return false
+	}
+
+	return meta.CheckNUMAAffinity()
+}
+
+// CheckNonBindingCPUAffinityNUMA returns true, if the AllocationMeta indicates that the meta is numa affinity candidate.
+// Now we consume that different numa affinity cpu pool can share the same numa node.
+func CheckNonBindingCPUAffinityNUMA(meta *AllocationMeta, annotations map[string]string) bool {
+	if meta == nil {
+		return false
+	}
+
+	if meta.CheckNUMANotShare() || qos.AnnotationsIndicateNUMANotShare(annotations) {
+		return false
+	}
+
+	return meta.CheckNonBindingNUMAAffinity()
+}
+
 // GenerateGenericContainerAllocationMeta generates a generic container's allocation metadata.
 // This function populates the AllocationMeta struct using data from the resource request and other parameters.
 // Parameters: