From 35849bf7fb8193e3e3c67f45f48d0fc1223bc635 Mon Sep 17 00:00:00 2001
From: Arpit Singh <arpsingh@nvidia.com>
Date: Wed, 6 Jul 2022 00:36:22 -0700
Subject: [PATCH] KEP-3327: Add CPUManager policy option to align CPUs by
 Socket instead of by NUMA node

---
 pkg/kubelet/cm/cpumanager/policy_options.go | 11 ++++++
 pkg/kubelet/cm/cpumanager/policy_static.go  | 39 ++++++++++++++++++---
 2 files changed, 46 insertions(+), 4 deletions(-)

diff --git a/pkg/kubelet/cm/cpumanager/policy_options.go b/pkg/kubelet/cm/cpumanager/policy_options.go
index 64b14833e58..a61e14ac856 100644
--- a/pkg/kubelet/cm/cpumanager/policy_options.go
+++ b/pkg/kubelet/cm/cpumanager/policy_options.go
@@ -28,11 +28,13 @@ import (
 const (
 	FullPCPUsOnlyOption            string = "full-pcpus-only"
 	DistributeCPUsAcrossNUMAOption string = "distribute-cpus-across-numa"
+	AlignBySocketOption            string = "align-by-socket"
 )
 
 var (
 	alphaOptions = sets.NewString(
 		DistributeCPUsAcrossNUMAOption,
+		AlignBySocketOption,
 	)
 	betaOptions = sets.NewString(
 		FullPCPUsOnlyOption,
@@ -69,6 +71,9 @@ type StaticPolicyOptions struct {
 	// Flag to evenly distribute CPUs across NUMA nodes in cases where more
 	// than one NUMA node is required to satisfy the allocation.
 	DistributeCPUsAcrossNUMA bool
+	// Flag to ensure CPU's are considered aligned at socket boundary rather than
+	// NUMA boundary
+	AlignBySocket bool
 }
 
 func NewStaticPolicyOptions(policyOptions map[string]string) (StaticPolicyOptions, error) {
@@ -91,6 +96,12 @@ func NewStaticPolicyOptions(policyOptions map[string]string) (StaticPolicyOption
 				return opts, fmt.Errorf("bad value for option %q: %w", name, err)
 			}
 			opts.DistributeCPUsAcrossNUMA = optValue
+		case AlignBySocketOption:
+			optValue, err := strconv.ParseBool(value)
+			if err != nil {
+				return opts, fmt.Errorf("bad value for option %q: %w", name, err)
+			}
+			opts.AlignBySocket = optValue
 		default:
 			// this should never be reached, we already detect unknown options,
 			// but we keep it as further safety.
diff --git a/pkg/kubelet/cm/cpumanager/policy_static.go b/pkg/kubelet/cm/cpumanager/policy_static.go
index c2aca10640f..0c3c241d360 100644
--- a/pkg/kubelet/cm/cpumanager/policy_static.go
+++ b/pkg/kubelet/cm/cpumanager/policy_static.go
@@ -325,10 +325,7 @@ func (p *staticPolicy) allocateCPUs(s state.State, numCPUs int, numaAffinity bit
 	// If there are aligned CPUs in numaAffinity, attempt to take those first.
 	result := cpuset.NewCPUSet()
 	if numaAffinity != nil {
-		alignedCPUs := cpuset.NewCPUSet()
-		for _, numaNodeID := range numaAffinity.GetBits() {
-			alignedCPUs = alignedCPUs.Union(allocatableCPUs.Intersection(p.topology.CPUDetails.CPUsInNUMANodes(numaNodeID)))
-		}
+		alignedCPUs := p.getAlignedCPUs(numaAffinity, allocatableCPUs)
 
 		numAlignedToAlloc := alignedCPUs.Size()
 		if numCPUs < numAlignedToAlloc {
@@ -571,6 +568,10 @@ func (p *staticPolicy) generateCPUTopologyHints(availableCPUs cpuset.CPUSet, reu
 	// to the minAffinitySize. Only those with an equal number of bits set (and
 	// with a minimal set of numa nodes) will be considered preferred.
 	for i := range hints {
+		if p.options.AlignBySocket && p.isHintSocketAligned(hints[i].NUMANodeAffinity) {
+			hints[i].Preferred = true
+			continue
+		}
 		if hints[i].NUMANodeAffinity.Count() == minAffinitySize {
 			hints[i].Preferred = true
 		}
@@ -578,3 +579,33 @@ func (p *staticPolicy) generateCPUTopologyHints(availableCPUs cpuset.CPUSet, reu
 
 	return hints
 }
+
+func (p *staticPolicy) isHintSocketAligned(hint bitmask.BitMask) bool {
+	numaNodes := hint.GetBits()
+	if p.topology.CPUDetails.SocketsInNUMANodes(numaNodes[:]...).Size() == 1 {
+		return true
+	}
+	return false
+}
+
+// getAlignedCPUs return set of aligned CPUs based on numa affinity mask and configured policy options.
+func (p *staticPolicy) getAlignedCPUs(numaAffinity bitmask.BitMask, allocatableCPUs cpuset.CPUSet) cpuset.CPUSet {
+	alignedCPUs := cpuset.NewCPUSet()
+	numaBits := numaAffinity.GetBits()
+	// If align-by-socket policy option is enabled, NUMA based hint is expanded to
+	// socket aligned hint. It will ensure that first socket aligned available CPUs are
+	// allocated before we try to find CPUs across socket to satisfy allocation request.
+	if p.options.AlignBySocket {
+		socketBits := p.topology.CPUDetails.SocketsInNUMANodes(numaBits...).ToSliceNoSort()
+		for _, socketID := range socketBits {
+			alignedCPUs = alignedCPUs.Union(allocatableCPUs.Intersection(p.topology.CPUDetails.CPUsInSockets(socketID)))
+		}
+		return alignedCPUs
+	}
+
+	for _, numaNodeID := range numaBits {
+		alignedCPUs = alignedCPUs.Union(allocatableCPUs.Intersection(p.topology.CPUDetails.CPUsInNUMANodes(numaNodeID)))
+	}
+
+	return alignedCPUs
+}