perf(scheduler): avoid pod-affinity pre-score allocations (#1983)

Signed-off-by: Erez Freiberger <enoodle@gmail.com>

Author: enoodle

.changes/unreleased/changed-20260727-142159.yaml

@@ -0,0 +1,3 @@
+kind: Fixed
+body: |-
+  Removed unnecessary pod-affinity pre-score allocation churn

Makefile

@@ -154,7 +154,7 @@ BENCH_OUTPUT ?= benchmark-results.txt
 # because some reclaim benchmarks require -benchtime=1x and only a curated subset
 # should run in CI.
 BENCH_SPECIAL_PACKAGES := ./pkg/scheduler/actions/reclaim
-BENCH_SPECIAL_REGEX := '^BenchmarkReclaim(WithMissingPVCJobs|UnschedulableDistributedJob_(10|50|100)Node)$$'
+BENCH_SPECIAL_REGEX := '^BenchmarkReclaim(WithMissingPVCJobs|UnschedulableDistributedJob_((10|50|100)Node|AntiAffinity100Node))$$'
 
 .PHONY: benchstat
 benchstat: $(BENCHSTAT)

pkg/scheduler/actions/reclaim/reclaim_benchmark_test.go

@@ -18,6 +18,8 @@ import (
 	"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/api/pod_status"
 	"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/api/podgroup_info/subgroup_info"
 	"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/api/topology_info"
+	"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/conf"
+	"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/conf_util"
 	"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/constants"
 	"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/test_utils"
 	"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/test_utils/jobs_fake"
@@ -239,6 +241,10 @@ func buildUnschedulableDistributedReclaimBenchmarkTopology(
 	}
 
 	jobs = append(jobs, distributedJob)
+	var schedulerConfig *conf.SchedulerConfiguration
+	if params.FailureMode == unschedulableDistributedReclaimAntiAffinity {
+		schedulerConfig = schedulerConfigWithPodAffinity()
+	}
 
 	return test_utils.TestTopologyBasic{
 		Name:       "unschedulable distributed reclaim benchmark",
@@ -259,10 +265,21 @@ func buildUnschedulableDistributedReclaimBenchmarkTopology(
 		},
 		Mocks: &test_utils.TestMock{
 			CacheRequirements: &test_utils.CacheMocking{},
+			SchedulerConf:     schedulerConfig,
 		},
 	}
 }
 
+func schedulerConfigWithPodAffinity() *conf.SchedulerConfiguration {
+	config, err := conf_util.GetDefaultSchedulerConf()
+	if err != nil {
+		panic(err)
+	}
+	config.ScenarioSearchBudgets = nil
+	config.Tiers[0].Plugins = append(config.Tiers[0].Plugins, conf.PluginOption{Name: "podaffinity"})
+	return config
+}
+
 func unschedulableDistributedRackCount(params unschedulableDistributedReclaimBenchmarkParams) int {
 	switch params.FailureMode {
 	case unschedulableDistributedReclaimAntiAffinity:

pkg/scheduler/plugins/podaffinity/podaffinity.go

@@ -4,9 +4,6 @@
 package podaffinity
 
 import (
-	ksf "k8s.io/kube-scheduler/framework"
-	k8sframework "k8s.io/kubernetes/pkg/scheduler/framework"
-
 	"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/api"
 	"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/api/common_info"
 	"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/api/node_info"
@@ -50,13 +47,8 @@ func (pp *podAffinityPlugin) OnSessionOpen(ssn *framework.Session) {
 }
 
 func (pp *podAffinityPlugin) nodePreOrderFn(k8sPlugins *k8s_internal.SessionScoreFns) api.NodePreOrderFn {
-	return func(task *pod_info.PodInfo, fittingNodes []*node_info.NodeInfo) error {
-		var nodes []ksf.NodeInfo
-		for range fittingNodes {
-			nodes = append(nodes, &k8sframework.NodeInfo{})
-		}
-
-		status := k8sPlugins.PrePodAffinity(task.Pod, nodes)
+	return func(task *pod_info.PodInfo, _ []*node_info.NodeInfo) error {
+		status := k8sPlugins.PrePodAffinity(task.Pod, nil)
 		if status.IsSkip() {
 			pp.skipOrderFn.add(task.UID)
 		}

pkg/scheduler/plugins/podaffinity/podaffinity_test.go

@@ -59,6 +59,21 @@ var _ = Describe("Preferred Pod Affinity", func() {
 				},
 			},
 		}
+		exampleAntiAffinity := &v1.Affinity{
+			PodAntiAffinity: &v1.PodAntiAffinity{
+				PreferredDuringSchedulingIgnoredDuringExecution: []v1.WeightedPodAffinityTerm{
+					{
+						Weight: 10,
+						PodAffinityTerm: v1.PodAffinityTerm{
+							LabelSelector: &metav1.LabelSelector{
+								MatchLabels: exampleLabels,
+							},
+							TopologyKey: nodeLabelName,
+						},
+					},
+				},
+			},
+		}
 
 		for testName, testData := range map[string]testInput{
 			"No Pod Affinity": {
@@ -156,6 +171,42 @@ var _ = Describe("Preferred Pod Affinity", func() {
 				}, resource_info.NewResourceVectorMap()),
 				expectedScore: 2 * 10 * scores.K8sPlugins,
 			},
+			"Matching Pod Anti-Affinity": {
+				nodes: []*v1.Node{
+					{
+						ObjectMeta: metav1.ObjectMeta{
+							Name: "node-1",
+							Labels: map[string]string{
+								nodeLabelName: "node-1",
+							},
+						},
+					},
+				},
+				pods: []*v1.Pod{
+					{
+						ObjectMeta: metav1.ObjectMeta{
+							Name:      "existing-pod-1",
+							Namespace: "test",
+							Labels:    exampleLabels,
+						},
+						Spec: v1.PodSpec{
+							NodeName: "node-1",
+							Affinity: exampleAntiAffinity.DeepCopy(),
+						},
+					},
+				},
+				task: pod_info.NewTaskInfo(&v1.Pod{
+					ObjectMeta: metav1.ObjectMeta{
+						Name:      "test-pod",
+						Namespace: "test",
+						Labels:    exampleLabels,
+					},
+					Spec: v1.PodSpec{
+						Affinity: exampleAntiAffinity.DeepCopy(),
+					},
+				}, resource_info.NewResourceVectorMap()),
+				expectedScore: -2 * 10 * scores.K8sPlugins,
+			},
 		} {
 			testName := testName
 			testData := testData