Benchmark test for PodAffinity

2025-09-21 01:50:55 +00:00 · 2019-09-17 20:23:20 -07:00
parent ef69b488bd
commit e85473c2df
5 changed files with 235 additions and 7 deletions
--- a/pkg/scheduler/testing/workload_prep.go
+++ b/pkg/scheduler/testing/workload_prep.go
@@ -22,15 +22,15 @@ import (
 	"k8s.io/api/core/v1"
 )

-// MakeNodesAndPods serves as a testing helper for EvenPodsSpread feature.
+// MakeNodesAndPodsForEvenPodsSpread serves as a testing helper for EvenPodsSpread feature.
 // It builds a fake cluster containing running Pods and Nodes.
 // The size of Pods and Nodes are determined by input arguments.
 // The specs of Pods and Nodes are generated with the following rules:
 // - If `pod` has "node" as a topologyKey, each generated node is applied with a unique label: "node: node<i>".
 // - If `pod` has "zone" as a topologyKey, each generated node is applied with a rotating label: "zone: zone[0-9]".
-// - Depending on "lableSelector.MatchExpressions[0].Key" the `pod` has in each topologySpreadConstraint,
+// - Depending on "labelSelector.MatchExpressions[0].Key" the `pod` has in each topologySpreadConstraint,
 //   each generated pod will be applied with label "key1", "key1,key2", ..., "key1,key2,...,keyN" in a rotating manner.
-func MakeNodesAndPods(pod *v1.Pod, existingPodsNum, allNodesNum, filteredNodesNum int) (existingPods []*v1.Pod, allNodes []*v1.Node, filteredNodes []*v1.Node) {
+func MakeNodesAndPodsForEvenPodsSpread(pod *v1.Pod, existingPodsNum, allNodesNum, filteredNodesNum int) (existingPods []*v1.Pod, allNodes []*v1.Node, filteredNodes []*v1.Node) {
 	var topologyKeys []string
 	var labels []string
 	zones := 10
@@ -65,3 +65,78 @@ func MakeNodesAndPods(pod *v1.Pod, existingPodsNum, allNodesNum, filteredNodesNu
 	}
 	return
 }
+
+// MakeNodesAndPodsForPodAffinity serves as a testing helper for Pod(Anti)Affinity feature.
+// It builds a fake cluster containing running Pods and Nodes.
+// For simplicity, the Nodes will be labelled with "region", "zone" and "node". Nodes[i] will be applied with:
+// - "region": "region" + i%3
+// - "zone": "zone" + i%10
+// - "node": "node" + i
+// The Pods will be applied with various combinations of PodAffinity and PodAntiAffinity terms.
+func MakeNodesAndPodsForPodAffinity(existingPodsNum, allNodesNum int) (existingPods []*v1.Pod, allNodes []*v1.Node) {
+	tpKeyToSizeMap := map[string]int{
+		"region": 3,
+		"zone":   10,
+		"node":   allNodesNum,
+	}
+	// build nodes to spread across all topology domains
+	for i := 0; i < allNodesNum; i++ {
+		nodeName := fmt.Sprintf("node%d", i)
+		nodeWrapper := MakeNode().Name(nodeName)
+		for tpKey, size := range tpKeyToSizeMap {
+			nodeWrapper = nodeWrapper.Label(tpKey, fmt.Sprintf("%s%d", tpKey, i%size))
+		}
+		allNodes = append(allNodes, nodeWrapper.Obj())
+	}
+
+	labels := []string{"foo", "bar", "baz"}
+	tpKeys := []string{"region", "zone", "node"}
+
+	// Build pods.
+	// Each pod will be created with one affinity and one anti-affinity terms using all combinations of
+	// affinity and anti-affinity kinds listed below
+	// e.g., the first pod will have {affinity, anti-affinity} terms of kinds {NilPodAffinity, NilPodAffinity};
+	// the second will be {NilPodAffinity, PodAntiAffinityWithRequiredReq}, etc.
+	affinityKinds := []PodAffinityKind{
+		NilPodAffinity,
+		PodAffinityWithRequiredReq,
+		PodAffinityWithPreferredReq,
+		PodAffinityWithRequiredPreferredReq,
+	}
+	antiAffinityKinds := []PodAffinityKind{
+		NilPodAffinity,
+		PodAntiAffinityWithRequiredReq,
+		PodAntiAffinityWithPreferredReq,
+		PodAntiAffinityWithRequiredPreferredReq,
+	}
+
+	totalSize := len(affinityKinds) * len(antiAffinityKinds)
+	for i := 0; i < existingPodsNum; i++ {
+		podWrapper := MakePod().Name(fmt.Sprintf("pod%d", i)).Node(fmt.Sprintf("node%d", i%allNodesNum))
+		label, tpKey := labels[i%len(labels)], tpKeys[i%len(tpKeys)]
+
+		affinityIdx := i % totalSize
+		// len(affinityKinds) is equal to len(antiAffinityKinds)
+		leftIdx, rightIdx := affinityIdx/len(affinityKinds), affinityIdx%len(affinityKinds)
+		podWrapper = podWrapper.PodAffinityExists(label, tpKey, affinityKinds[leftIdx])
+		podWrapper = podWrapper.PodAntiAffinityExists(label, tpKey, antiAffinityKinds[rightIdx])
+		existingPods = append(existingPods, podWrapper.Obj())
+	}
+
+	return
+}
+
+// MakeNodesAndPods serves as a testing helper to generate regular Nodes and Pods
+// that don't use any advanced scheduling features.
+func MakeNodesAndPods(existingPodsNum, allNodesNum int) (existingPods []*v1.Pod, allNodes []*v1.Node) {
+	// build nodes
+	for i := 0; i < allNodesNum; i++ {
+		allNodes = append(allNodes, MakeNode().Name(fmt.Sprintf("node%d", i)).Obj())
+	}
+	// build pods
+	for i := 0; i < existingPodsNum; i++ {
+		podWrapper := MakePod().Name(fmt.Sprintf("pod%d", i)).Node(fmt.Sprintf("node%d", i%allNodesNum))
+		existingPods = append(existingPods, podWrapper.Obj())
+	}
+	return
+}
--- a/pkg/scheduler/testing/wrappers.go
+++ b/pkg/scheduler/testing/wrappers.go
@@ -190,7 +190,7 @@ func (p *PodWrapper) NodeSelector(m map[string]string) *PodWrapper {
 }

 // NodeAffinityIn creates a HARD node affinity (with the operator In)
-// and injects into the innner pod.
+// and injects into the inner pod.
 func (p *PodWrapper) NodeAffinityIn(key string, vals []string) *PodWrapper {
 	if p.Spec.Affinity == nil {
 		p.Spec.Affinity = &v1.Affinity{}
@@ -204,7 +204,7 @@ func (p *PodWrapper) NodeAffinityIn(key string, vals []string) *PodWrapper {
 }

 // NodeAffinityNotIn creates a HARD node affinity (with the operator NotIn)
-// and injects into the innner pod.
+// and injects into the inner pod.
 func (p *PodWrapper) NodeAffinityNotIn(key string, vals []string) *PodWrapper {
 	if p.Spec.Affinity == nil {
 		p.Spec.Affinity = &v1.Affinity{}
@@ -217,6 +217,104 @@ func (p *PodWrapper) NodeAffinityNotIn(key string, vals []string) *PodWrapper {
 	return p
 }

+// PodAffinityKind represents different kinds of PodAffinity.
+type PodAffinityKind int
+
+const (
+	// NilPodAffinity is a no-op which doesn't apply any PodAffinity.
+	NilPodAffinity PodAffinityKind = iota
+	// PodAffinityWithRequiredReq applies a HARD requirement to pod.spec.affinity.PodAffinity.
+	PodAffinityWithRequiredReq
+	// PodAffinityWithPreferredReq applies a SOFT requirement to pod.spec.affinity.PodAffinity.
+	PodAffinityWithPreferredReq
+	// PodAffinityWithRequiredPreferredReq applies HARD and SOFT requirements to pod.spec.affinity.PodAffinity.
+	PodAffinityWithRequiredPreferredReq
+	// PodAntiAffinityWithRequiredReq applies a HARD requirement to pod.spec.affinity.PodAntiAffinity.
+	PodAntiAffinityWithRequiredReq
+	// PodAntiAffinityWithPreferredReq applies a SOFT requirement to pod.spec.affinity.PodAntiAffinity.
+	PodAntiAffinityWithPreferredReq
+	// PodAntiAffinityWithRequiredPreferredReq applies HARD and SOFT requirements to pod.spec.affinity.PodAntiAffinity.
+	PodAntiAffinityWithRequiredPreferredReq
+)
+
+// PodAffinityExists creates an PodAffinity with the operator "Exists"
+// and injects into the inner pod.
+func (p *PodWrapper) PodAffinityExists(labelKey, topologyKey string, kind PodAffinityKind) *PodWrapper {
+	if kind == NilPodAffinity {
+		return p
+	}
+
+	if p.Spec.Affinity == nil {
+		p.Spec.Affinity = &v1.Affinity{}
+	}
+	if p.Spec.Affinity.PodAffinity == nil {
+		p.Spec.Affinity.PodAffinity = &v1.PodAffinity{}
+	}
+	labelSelector := MakeLabelSelector().Exists(labelKey).Obj()
+	term := v1.PodAffinityTerm{LabelSelector: labelSelector, TopologyKey: topologyKey}
+	switch kind {
+	case PodAffinityWithRequiredReq:
+		p.Spec.Affinity.PodAffinity.RequiredDuringSchedulingIgnoredDuringExecution = append(
+			p.Spec.Affinity.PodAffinity.RequiredDuringSchedulingIgnoredDuringExecution,
+			term,
+		)
+	case PodAffinityWithPreferredReq:
+		p.Spec.Affinity.PodAffinity.PreferredDuringSchedulingIgnoredDuringExecution = append(
+			p.Spec.Affinity.PodAffinity.PreferredDuringSchedulingIgnoredDuringExecution,
+			v1.WeightedPodAffinityTerm{Weight: 1, PodAffinityTerm: term},
+		)
+	case PodAffinityWithRequiredPreferredReq:
+		p.Spec.Affinity.PodAffinity.RequiredDuringSchedulingIgnoredDuringExecution = append(
+			p.Spec.Affinity.PodAffinity.RequiredDuringSchedulingIgnoredDuringExecution,
+			term,
+		)
+		p.Spec.Affinity.PodAffinity.PreferredDuringSchedulingIgnoredDuringExecution = append(
+			p.Spec.Affinity.PodAffinity.PreferredDuringSchedulingIgnoredDuringExecution,
+			v1.WeightedPodAffinityTerm{Weight: 1, PodAffinityTerm: term},
+		)
+	}
+	return p
+}
+
+// PodAntiAffinityExists creates an PodAntiAffinity with the operator "Exists"
+// and injects into the inner pod.
+func (p *PodWrapper) PodAntiAffinityExists(labelKey, topologyKey string, kind PodAffinityKind) *PodWrapper {
+	if kind == NilPodAffinity {
+		return p
+	}
+
+	if p.Spec.Affinity == nil {
+		p.Spec.Affinity = &v1.Affinity{}
+	}
+	if p.Spec.Affinity.PodAntiAffinity == nil {
+		p.Spec.Affinity.PodAntiAffinity = &v1.PodAntiAffinity{}
+	}
+	labelSelector := MakeLabelSelector().Exists(labelKey).Obj()
+	term := v1.PodAffinityTerm{LabelSelector: labelSelector, TopologyKey: topologyKey}
+	switch kind {
+	case PodAntiAffinityWithRequiredReq:
+		p.Spec.Affinity.PodAntiAffinity.RequiredDuringSchedulingIgnoredDuringExecution = append(
+			p.Spec.Affinity.PodAntiAffinity.RequiredDuringSchedulingIgnoredDuringExecution,
+			term,
+		)
+	case PodAntiAffinityWithPreferredReq:
+		p.Spec.Affinity.PodAntiAffinity.PreferredDuringSchedulingIgnoredDuringExecution = append(
+			p.Spec.Affinity.PodAntiAffinity.PreferredDuringSchedulingIgnoredDuringExecution,
+			v1.WeightedPodAffinityTerm{Weight: 1, PodAffinityTerm: term},
+		)
+	case PodAntiAffinityWithRequiredPreferredReq:
+		p.Spec.Affinity.PodAntiAffinity.RequiredDuringSchedulingIgnoredDuringExecution = append(
+			p.Spec.Affinity.PodAntiAffinity.RequiredDuringSchedulingIgnoredDuringExecution,
+			term,
+		)
+		p.Spec.Affinity.PodAntiAffinity.PreferredDuringSchedulingIgnoredDuringExecution = append(
+			p.Spec.Affinity.PodAntiAffinity.PreferredDuringSchedulingIgnoredDuringExecution,
+			v1.WeightedPodAffinityTerm{Weight: 1, PodAffinityTerm: term},
+		)
+	}
+	return p
+}
+
 // SpreadConstraint constructs a TopologySpreadConstraint object and injects
 // into the inner pod.
 func (p *PodWrapper) SpreadConstraint(maxSkew int, tpKey string, mode v1.UnsatisfiableConstraintAction, selector *metav1.LabelSelector) *PodWrapper {