Merge pull request #135650 from pravk03/benchmark-test

Add node declared features scheduler perf tests

test/integration/scheduler_perf/dra.go

@@ -82,7 +82,7 @@ func (op *createResourceClaimsOp) isValid(allowParameterization bool) error {
 func (op *createResourceClaimsOp) collectsMetrics() bool {
 	return false
 }
-func (op *createResourceClaimsOp) patchParams(w *workload) (realOp, error) {
+func (op *createResourceClaimsOp) patchParams(w *Workload) (realOp, error) {
 	if op.CountParam != "" {
 		var err error
 		op.Count, err = w.Params.get(op.CountParam[1:])
@@ -159,7 +159,7 @@ func (op *createResourceDriverOp) isValid(allowParameterization bool) error {
 func (op *createResourceDriverOp) collectsMetrics() bool {
 	return false
 }
-func (op *createResourceDriverOp) patchParams(w *workload) (realOp, error) {
+func (op *createResourceDriverOp) patchParams(w *Workload) (realOp, error) {
 	if op.MaxClaimsPerNodeParam != "" {
 		var err error
 		op.MaxClaimsPerNode, err = w.Params.get(op.MaxClaimsPerNodeParam[1:])
@@ -267,7 +267,7 @@ func (op *allocResourceClaimsOp) isValid(allowParameterization bool) error {
 func (op *allocResourceClaimsOp) collectsMetrics() bool {
 	return false
 }
-func (op *allocResourceClaimsOp) patchParams(w *workload) (realOp, error) {
+func (op *allocResourceClaimsOp) patchParams(w *Workload) (realOp, error) {
 	return op, op.isValid(false)
 }
 

test/integration/scheduler_perf/dra/dra_test.go

@@ -62,7 +62,7 @@ func TestSchedulerPerf(t *testing.T) {
 			// - "default": don't change features
 			var options []perf.SchedulerPerfOption
 			if allocatorName == "stable" {
-				options = append(options, perf.WithPreRunFn(func(tCtx ktesting.TContext) error {
+				options = append(options, perf.WithPreRunFn(func(tCtx ktesting.TContext, _ *perf.Workload) (func(), error) {
 					gate := utilfeature.DefaultFeatureGate.(featuregate.MutableVersionedFeatureGate)
 					overrides := featuregatetesting.FeatureOverrides{
 						features.DRAPrioritizedList: false,
@@ -74,7 +74,7 @@ func TestSchedulerPerf(t *testing.T) {
 						overrides[features.DRAConsumableCapacity] = false
 					}
 					featuregatetesting.SetFeatureGatesDuringTest(tCtx, utilfeature.DefaultFeatureGate, overrides)
-					return nil
+					return nil, nil
 				}))
 			}
 

test/integration/scheduler_perf/executor.go

@@ -62,9 +62,10 @@ type WorkloadExecutor struct {
 	podInformer                  coreinformers.PodInformer
 	throughputErrorMargin        float64
 	testCase                     *testCase
-	workload                     *workload
+	workload                     *Workload
 	topicName                    string
 	nextNodeIndex                int
+	opts                         *schedulerPerfOptions
 }
 
 func (e *WorkloadExecutor) wait() {
@@ -109,6 +110,15 @@ func (e *WorkloadExecutor) runCreateNodesOp(tCtx ktesting.TContext, opIndex int,
 		return err
 	}
 	e.nextNodeIndex += op.Count
+	if e.opts != nil && e.opts.nodeUpdateFn != nil {
+		nodes, err := waitListAllNodes(tCtx, tCtx.Client())
+		if err != nil {
+			return fmt.Errorf("failed to list nodes for nodeUpdateFn: %w", err)
+		}
+		if err := e.opts.nodeUpdateFn(tCtx, e.scheduler, e.workload, nodes); err != nil {
+			return fmt.Errorf("nodeUpdateFn failed: %w", err)
+		}
+	}
 	return nil
 }
 

test/integration/scheduler_perf/executor_test.go

@@ -629,7 +629,7 @@ func TestMetricThreshold(t *testing.T) {
 				return []testDataCollector{&mockDataCollector{dataItems: tc.dataItems}}
 			}
 
-			workload := &workload{
+			workload := &Workload{
 				Name: "some/workload",
 				Threshold: thresholds{
 					valuesByTopic: map[string]float64{"example": tc.thresholdValue},

test/integration/scheduler_perf/misc/performance-config.yaml

@@ -100,7 +100,14 @@
       initNodes: 5000
       initPods: 5000
       measurePods: 50000
-
+  - name: 5000Nodes_50000Pods_NodeDeclaredFeaturesDisabled
+    labels: [performance]
+    featureGates:
+      NodeDeclaredFeatures: false
+    params:
+      initNodes: 5000
+      initPods: 5000
+      measurePods: 50000
 # This test case simulates the scheduling of daemonset.
 # https://github.com/kubernetes/kubernetes/issues/124709
 - name: SchedulingDaemonset

test/integration/scheduler_perf/nodedeclaredfeatures/node_declared_features_test.go

@@ -0,0 +1,189 @@
+/*
+Copyright The Kubernetes Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package nodedeclaredfeatures
+
+import (
+	"context"
+	"fmt"
+	"os"
+	"slices"
+	"sort"
+	"testing"
+	"time"
+
+	v1 "k8s.io/api/core/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/util/version"
+	"k8s.io/apimachinery/pkg/util/wait"
+	_ "k8s.io/component-base/logs/json/register"
+	ndf "k8s.io/component-helpers/nodedeclaredfeatures"
+	ndffeatures "k8s.io/component-helpers/nodedeclaredfeatures/features"
+	"k8s.io/kubernetes/pkg/features"
+	"k8s.io/kubernetes/pkg/scheduler"
+	perf "k8s.io/kubernetes/test/integration/scheduler_perf"
+	"k8s.io/kubernetes/test/utils/ktesting"
+)
+
+// This test benchmarks the NodeDeclaredFeatures features performance impact.
+//
+// It mocks a variable number of node features registered (via 'numNodeDeclaredFeatures' param) but the
+// inference function of all the registered features look for pod level resouces ('spec.resources').
+// The tests creates a pod with pod level resources to ensure that all features are inferred
+// during scheduling.
+
+func TestMain(m *testing.M) {
+	if err := perf.InitTests(); err != nil {
+		fmt.Fprintf(os.Stderr, "%v\n", err)
+		os.Exit(1)
+	}
+
+	m.Run()
+}
+
+// mockFeature is a mock implementation of the Feature interface for testing.
+type mockFeature struct {
+	name       string
+	maxVersion *version.Version
+}
+
+func (f *mockFeature) Name() string {
+	return f.name
+}
+
+func (f *mockFeature) Discover(cfg *ndf.NodeConfiguration) bool {
+	// This function is called by kubelet to discover features on the node and report them in the node status.
+	// In the test context, feature discovery is bypassed and we update node status directly after
+	// the node is created (see updateNodesWithDeclaredFeatures()), so this function is a no-op.
+	return true
+}
+
+func (f *mockFeature) InferForScheduling(podInfo *ndf.PodInfo) bool {
+	// Check if the pod is using pod level resources.
+	return podInfo.Spec != nil && podInfo.Spec.Resources != nil
+}
+
+func (f *mockFeature) InferForUpdate(oldPodInfo, newPodInfo *ndf.PodInfo) bool {
+	return false
+}
+
+func (f *mockFeature) Requirements() *ndf.FeatureRequirements {
+	return nil
+}
+
+func (f *mockFeature) MaxVersion() *version.Version {
+	return f.maxVersion
+}
+
+func createMockFeature(name string, maxVersionStr string) ndf.Feature {
+	var v *version.Version
+	if maxVersionStr != "" {
+		v = version.MustParseSemantic(maxVersionStr)
+	}
+	return &mockFeature{
+		name:       name,
+		maxVersion: v,
+	}
+}
+
+func setupFeatures(numNodeDeclaredFeatures int) func() {
+	nodeFeatures := make([]ndf.Feature, 0, numNodeDeclaredFeatures)
+	for i := 1; i <= numNodeDeclaredFeatures; i++ {
+		featureName := fmt.Sprintf("Feature%d", i)
+		nodeFeatures = append(nodeFeatures, createMockFeature(featureName, ""))
+	}
+	originalAllFeatures := ndffeatures.AllFeatures
+	ndffeatures.AllFeatures = nodeFeatures
+	return func() {
+		ndffeatures.AllFeatures = originalAllFeatures
+	}
+}
+
+func preInitNodeDeclaredFeatures(t ktesting.TContext, w *perf.Workload) (func(), error) {
+	if !w.FeatureGates[features.NodeDeclaredFeatures] {
+		t.Logf("Skipping NodeDeclaredFeatures pre-init as the feature gate is disabled")
+		return func() {}, nil
+	}
+
+	numNodeDeclaredFeatures, err := w.GetParam("numNodeDeclaredFeatures")
+	if err != nil {
+		t.Logf("numNodeDeclaredFeatures param not specified in workload config")
+		numNodeDeclaredFeatures = 0
+	}
+
+	t.Logf("PreInit: Setting up %d mock features for %s", numNodeDeclaredFeatures, w.Name)
+	// Setup mock features to be registered in the NDF library.
+	cleanupMockFeatures := setupFeatures(numNodeDeclaredFeatures)
+
+	return cleanupMockFeatures, nil
+}
+
+func updateNodesWithDeclaredFeatures(tCtx ktesting.TContext, scheduler *scheduler.Scheduler, w *perf.Workload, nodes *v1.NodeList) error {
+	if !w.FeatureGates[features.NodeDeclaredFeatures] {
+		return nil
+	}
+	numNodeDeclaredFeatures, err := w.GetParam("numNodeDeclaredFeatures")
+	if err != nil {
+		tCtx.Logf("numNodeDeclaredFeatures param not specified in workload config")
+		numNodeDeclaredFeatures = 0
+	}
+
+	var featureNames []string
+	for i := 1; i <= numNodeDeclaredFeatures; i++ {
+		featureNames = append(featureNames, fmt.Sprintf("Feature%d", i))
+	}
+	sort.Strings(featureNames)
+
+	for _, node := range nodes.Items {
+		nodeToUpdate := node.DeepCopy()
+		nodeToUpdate.Status.DeclaredFeatures = featureNames
+		_, err := tCtx.Client().CoreV1().Nodes().UpdateStatus(tCtx.Context, nodeToUpdate, metav1.UpdateOptions{})
+		if err != nil {
+			return fmt.Errorf("failed to update node %s status: %w", node.Name, err)
+		}
+		tCtx.Logf("Updated node %s status with %d features", node.Name, len(featureNames))
+	}
+
+	schedulerCache := scheduler.Cache
+	err = wait.PollUntilContextTimeout(tCtx.Context, 100*time.Millisecond, 60*time.Second, true, func(ctx context.Context) (bool, error) {
+		for _, n := range nodes.Items {
+			nodeInfo, err := schedulerCache.GetNode(n.Name)
+			if err != nil {
+				return false, nil
+
+			}
+			cachedNode := nodeInfo.Node()
+			if !slices.Equal(cachedNode.Status.DeclaredFeatures, featureNames) {
+				return false, nil
+			}
+		}
+		return true, nil
+	})
+
+	if err != nil {
+		return fmt.Errorf("timeout waiting for all nodes to reflect status update: %w", err)
+	}
+
+	return nil
+}
+
+func TestSchedulerPerf(t *testing.T) {
+	perf.RunIntegrationPerfScheduling(t, "performance-config.yaml", perf.WithPreRunFn(preInitNodeDeclaredFeatures), perf.WithNodeUpdateFn(updateNodesWithDeclaredFeatures))
+}
+
+func BenchmarkPerfScheduling(b *testing.B) {
+	perf.RunBenchmarkPerfScheduling(b, "performance-config.yaml", "nodedeclaredfeatures", nil, perf.WithPreRunFn(preInitNodeDeclaredFeatures), perf.WithNodeUpdateFn(updateNodesWithDeclaredFeatures))
+}

test/integration/scheduler_perf/nodedeclaredfeatures/performance-config.yaml

@@ -0,0 +1,69 @@
+# The following labels are used in this file. (listed in ascending order of the number of covered test cases)
+#
+# - integration-test: test cases to run as the integration test, usually to spot some issues in the scheduler implementation or scheduler-perf itself.
+# - performance: test cases to run in the performance test.
+# - short: supplemental label for the above two labels (must not used alone), which literally means short execution time test cases.
+#
+# Specifically, the CIs use labels like the following:
+# - `ci-kubernetes-integration-master` (`integration-test`): Test cases are chosen based on a tradeoff between code coverage and overall runtime. 
+# It basically covers all test cases but with their smallest workload. 
+# - `pull-kubernetes-integration` (`integration-test`,`short`): Test cases are chosen so that they should take less than total 5 min to complete.
+# - `ci-benchmark-scheduler-perf` (`performance`): Long enough test cases are chosen (ideally, longer than 10 seconds) 
+# to provide meaningful samples for the pod scheduling rate.
+#
+# Also, `performance`+`short` isn't used in the CIs, but it's used to test the performance test locally.
+# (Sometimes, the test cases with `integration-test` are too small to spot issues.)
+#
+# Combining `performance` and `short` selects suitable workloads for a local
+# before/after comparisons with benchstat.
+
+- name: NodeDeclaredFeaturesEnabled
+  defaultPodTemplatePath: templates/pod-with-pod-level-resources.yaml
+  workloadTemplate:
+    - opcode: createNodes
+      countParam: $nodes
+      nodeTemplatePath: templates/node-base.yaml
+    - opcode: createPods
+      namespace: init
+      countParam: $initPods
+    - opcode: createPods
+      namespace: test
+      countParam: $measurePods
+      collectMetrics: true
+  workloads:
+    - name: 5Nodes5DeclaredFeatures
+      labels: [integration-test, short]
+      featureGates:
+        NodeDeclaredFeatures: true
+      params:
+        nodes: 5
+        initPods: 0
+        measurePods: 10
+        numNodeDeclaredFeatures: 5
+    - name: 5000Nodes20DeclaredFeatures
+      labels: [performance]
+      featureGates:
+        NodeDeclaredFeatures: true
+      params:
+        nodes: 5000
+        initPods: 5000
+        measurePods: 5000
+        numNodeDeclaredFeatures: 20
+    - name: 5000Nodes40DeclaredFeatures
+      labels: [performance]
+      featureGates:
+        NodeDeclaredFeatures: true
+      params:
+        nodes: 5000
+        initPods: 5000
+        measurePods: 5000
+        numNodeDeclaredFeatures: 40
+    - name: 5000Nodes100DeclaredFeatures
+      labels: [performance]
+      featureGates:
+        NodeDeclaredFeatures: true
+      params:
+        nodes: 5000
+        initPods: 5000
+        measurePods: 5000
+        numNodeDeclaredFeatures: 100

test/integration/scheduler_perf/nodedeclaredfeatures/templates/node-base.yaml

@@ -0,0 +1,17 @@
+apiVersion: v1
+kind: Node
+metadata:
+  generateName: node-
+spec: {}
+status:
+  capacity:
+    cpu: "128"
+    memory: "16Gi"
+    pods: "110"
+  allocatable:
+    cpu: "128"
+    memory: "16Gi"
+    pods: "110"
+  phase: Running
+  # DeclaredFeatures will be added dynamically by the test based on numNodeDeclaredFeatures param in the config.
+

test/integration/scheduler_perf/nodedeclaredfeatures/templates/pod-with-pod-level-resources.yaml

@@ -0,0 +1,12 @@
+apiVersion: v1
+kind: Pod
+metadata:
+  generateName: pod-
+spec:
+  resources:
+    requests:
+      cpu: "100m"
+      memory: "200Mi"
+  containers:
+  - name: container-a
+    image: registry.k8s.io/pause:3.10.1

test/integration/scheduler_perf/operations.go

@@ -71,7 +71,7 @@ func (c *createAny) collectsMetrics() bool {
 	return false
 }
 
-func (c createAny) patchParams(w *workload) (realOp, error) {
+func (c createAny) patchParams(w *Workload) (realOp, error) {
 	if c.CountParam != "" {
 		count, err := w.Params.get(c.CountParam[1:])
 		if err != nil {
@@ -215,7 +215,7 @@ func (*createNodesOp) collectsMetrics() bool {
 	return false
 }
 
-func (cno createNodesOp) patchParams(w *workload) (realOp, error) {
+func (cno createNodesOp) patchParams(w *Workload) (realOp, error) {
 	if cno.CountParam != "" {
 		var err error
 		cno.Count, err = w.Params.get(cno.CountParam[1:])
@@ -253,7 +253,7 @@ func (*createNamespacesOp) collectsMetrics() bool {
 	return false
 }
 
-func (cmo createNamespacesOp) patchParams(w *workload) (realOp, error) {
+func (cmo createNamespacesOp) patchParams(w *Workload) (realOp, error) {
 	if cmo.CountParam != "" {
 		var err error
 		cmo.Count, err = w.Params.get(cmo.CountParam[1:])
@@ -342,7 +342,7 @@ func (cpo *createPodsOp) collectsMetrics() bool {
 	return cpo.CollectMetrics
 }
 
-func (cpo createPodsOp) patchParams(w *workload) (realOp, error) {
+func (cpo createPodsOp) patchParams(w *Workload) (realOp, error) {
 	if cpo.CountParam != "" {
 		var err error
 		cpo.Count, err = w.Params.get(cpo.CountParam[1:])
@@ -395,7 +395,7 @@ func (cpso *createPodSetsOp) collectsMetrics() bool {
 	return cpso.CreatePodsOp.CollectMetrics
 }
 
-func (cpso createPodSetsOp) patchParams(w *workload) (realOp, error) {
+func (cpso createPodSetsOp) patchParams(w *Workload) (realOp, error) {
 	if cpso.CountParam != "" {
 		var err error
 		cpso.Count, err = w.Params.get(cpso.CountParam[1:])
@@ -442,7 +442,7 @@ func (dpo *deletePodsOp) collectsMetrics() bool {
 	return false
 }
 
-func (dpo deletePodsOp) patchParams(w *workload) (realOp, error) {
+func (dpo deletePodsOp) patchParams(w *Workload) (realOp, error) {
 	return &dpo, nil
 }
 
@@ -489,7 +489,7 @@ func (*churnOp) collectsMetrics() bool {
 	return false
 }
 
-func (co churnOp) patchParams(w *workload) (realOp, error) {
+func (co churnOp) patchParams(w *Workload) (realOp, error) {
 	return &co, nil
 }
 
@@ -530,7 +530,7 @@ func (*barrierOp) collectsMetrics() bool {
 	return false
 }
 
-func (bo barrierOp) patchParams(w *workload) (realOp, error) {
+func (bo barrierOp) patchParams(w *Workload) (realOp, error) {
 	if bo.StageRequirement == "" {
 		bo.StageRequirement = Scheduled
 	}
@@ -556,7 +556,7 @@ func (so *sleepOp) collectsMetrics() bool {
 	return false
 }
 
-func (so sleepOp) patchParams(w *workload) (realOp, error) {
+func (so sleepOp) patchParams(w *Workload) (realOp, error) {
 	if so.DurationParam != "" {
 		durationStr, err := getParam[string](w.Params, so.DurationParam[1:])
 		if err != nil {
@@ -595,7 +595,7 @@ func (*startCollectingMetricsOp) collectsMetrics() bool {
 	return false
 }
 
-func (scm startCollectingMetricsOp) patchParams(_ *workload) (realOp, error) {
+func (scm startCollectingMetricsOp) patchParams(_ *Workload) (realOp, error) {
 	return &scm, nil
 }
 
@@ -615,7 +615,7 @@ func (*stopCollectingMetricsOp) collectsMetrics() bool {
 	return true
 }
 
-func (scm stopCollectingMetricsOp) patchParams(_ *workload) (realOp, error) {
+func (scm stopCollectingMetricsOp) patchParams(_ *Workload) (realOp, error) {
 	return &scm, nil
 }
 

test/integration/scheduler_perf/options.go

@@ -17,6 +17,8 @@ limitations under the License.
 package benchmark
 
 import (
+	v1 "k8s.io/api/core/v1"
+	"k8s.io/kubernetes/pkg/scheduler"
 	frameworkruntime "k8s.io/kubernetes/pkg/scheduler/framework/runtime"
 	"k8s.io/kubernetes/test/utils/ktesting"
 )
@@ -28,10 +30,18 @@ type SchedulerPerfOption func(options *schedulerPerfOptions)
 // Alternatively, it may also return a non-nil error.
 type HookFn func(tCtx ktesting.TContext) error
 
+// PreRunFn hook function is called for each workload after feature gates are set,
+// but before he scheduler is started. It returns an optional cleanup function and an error.
+type PreRunFn func(tCtx ktesting.TContext, w *Workload) (func(), error)
+
+// NodeUpdateFn is a function called after nodes are created in a workload using createNodesOp.
+type NodeUpdateFn func(tCtx ktesting.TContext, scheduler *scheduler.Scheduler, w *Workload, nodes *v1.NodeList) error
+
 type schedulerPerfOptions struct {
 	outOfTreePluginRegistry frameworkruntime.Registry
-	preRunFn                HookFn
+	preRunFn                PreRunFn
 	prepareFn               HookFn
+	nodeUpdateFn            NodeUpdateFn
 }
 
 // WithPrepareFn is the option to set a function that is called
@@ -43,10 +53,18 @@ func WithPrepareFn(prepareFn HookFn) SchedulerPerfOption {
 	}
 }
 
+// WithNodeUpdateFn is the option to set a function that is called
+// after nodes are created by createNodesOp within a workload execution.
+func WithNodeUpdateFn(fn NodeUpdateFn) SchedulerPerfOption {
+	return func(s *schedulerPerfOptions) {
+		s.nodeUpdateFn = fn
+	}
+}
+
 // WithPreRunFn is the option to set a function that is called
 // after configuring the process (logging, feature gates) and
 // before running any code (etcd, scheduler).
-func WithPreRunFn(preRunFn HookFn) SchedulerPerfOption {
+func WithPreRunFn(preRunFn PreRunFn) SchedulerPerfOption {
 	return func(s *schedulerPerfOptions) {
 		s.preRunFn = preRunFn
 	}

test/integration/scheduler_perf/scheduler_perf.go

@@ -261,7 +261,7 @@ type testCase struct {
 	// be executed serially one after another.
 	WorkloadTemplate []op
 	// List of workloads to run under this testCase.
-	Workloads []*workload
+	Workloads []*Workload
 	// SchedulerConfigPath is the path of scheduler configuration
 	// Optional
 	SchedulerConfigPath string
@@ -298,10 +298,10 @@ func (tc *testCase) workloadNamesUnique() error {
 	return nil
 }
 
-// workload is a subtest under a testCase that tests the scheduler performance
+// Workload is a subtest under a testCase that tests the scheduler performance
 // for a certain ordering of ops. The set of nodes created and pods scheduled
-// in a workload may be heterogeneous.
-type workload struct {
+// in a Workload may be heterogeneous.
+type Workload struct {
 	// Name of the workload.
 	Name string
 	// Values of parameters used in the workloadTemplate.
@@ -331,7 +331,12 @@ type workload struct {
 	FeatureGates map[featuregate.Feature]bool
 }
 
-func (w *workload) isValid(mcc *metricsCollectorConfig) error {
+// GetParam retrieves a parameter from the Workload's parameters as an integer.
+func (w *Workload) GetParam(key string) (int, error) {
+	return w.Params.get(key)
+}
+
+func (w *Workload) isValid(mcc *metricsCollectorConfig) error {
 	if w.Threshold.value < 0 {
 		return fmt.Errorf("invalid Threshold=%f; should be non-negative", w.Threshold.value)
 	}
@@ -344,7 +349,7 @@ func (w *workload) isValid(mcc *metricsCollectorConfig) error {
 	return w.ThresholdMetricSelector.isValid(mcc)
 }
 
-func (w *workload) setDefaults(testCaseThresholdMetricSelector *thresholdMetricSelector) {
+func (w *Workload) setDefaults(testCaseThresholdMetricSelector *thresholdMetricSelector) {
 	if w.ThresholdMetricSelector != nil {
 		return
 	}
@@ -484,7 +489,7 @@ func getParam[T float64 | string | bool](p params, key string) (T, error) {
 }
 
 // unusedParams returns the names of unusedParams
-func (w workload) unusedParams() []string {
+func (w Workload) unusedParams() []string {
 	var ret []string
 	for name := range w.Params.params {
 		if !w.Params.isUsed[name] {
@@ -559,7 +564,7 @@ type realOp interface {
 	// type, even though calls will be made from with a *realOp. This is because
 	// callers don't want the receiver to inadvertently modify the realOp
 	// (instead, it's returned as a return value).
-	patchParams(w *workload) (realOp, error)
+	patchParams(w *Workload) (realOp, error)
 }
 
 // runnableOp is an interface implemented by some operations. It makes it possible
@@ -619,7 +624,7 @@ func initTestOutput(tb testing.TB) io.Writer {
 
 var specialFilenameChars = regexp.MustCompile(`[^a-zA-Z0-9-_]`)
 
-func setupTestCase(t testing.TB, tc *testCase, featureGates map[featuregate.Feature]bool, opts *schedulerPerfOptions) (*scheduler.Scheduler, informers.SharedInformerFactory, ktesting.TContext) {
+func setupTestCase(t testing.TB, tc *testCase, featureGates map[featuregate.Feature]bool, workload *Workload, opts *schedulerPerfOptions) (*scheduler.Scheduler, informers.SharedInformerFactory, ktesting.TContext) {
 	tCtx := ktesting.Init(t, initoption.PerTestOutput(UseTestingLog))
 	artifacts, doArtifacts := os.LookupEnv("ARTIFACTS")
 	if !UseTestingLog && doArtifacts {
@@ -700,8 +705,12 @@ func setupTestCase(t testing.TB, tc *testCase, featureGates map[featuregate.Feat
 	featuregatetesting.SetFeatureGatesDuringTest(t, utilfeature.DefaultFeatureGate, featureGates)
 
 	if opts.preRunFn != nil {
-		if err := opts.preRunFn(tCtx); err != nil {
-			t.Fatalf("pre-run: %v", err)
+		cleanup, err := opts.preRunFn(tCtx, workload)
+		if err != nil {
+			t.Fatalf("failed to run preRunFn for workload %s: %v", workload.Name, err)
+		}
+		if cleanup != nil {
+			t.Cleanup(cleanup)
 		}
 	}
 
@@ -811,7 +820,7 @@ func RunBenchmarkPerfScheduling(b *testing.B, configFile string, topicName strin
 					fixJSONOutput(b)
 
 					featureGates := featureGatesMerge(tc.FeatureGates, w.FeatureGates)
-					scheduler, informerFactory, tCtx := setupTestCase(b, tc, featureGates, opts)
+					scheduler, informerFactory, tCtx := setupTestCase(b, tc, featureGates, w, opts)
 
 					err := w.isValid(tc.MetricsCollectorConfig)
 					if err != nil {
@@ -825,7 +834,7 @@ func RunBenchmarkPerfScheduling(b *testing.B, configFile string, topicName strin
 						}
 					}
 
-					results, err := runWorkload(tCtx, tc, w, topicName, scheduler, informerFactory)
+					results, err := runWorkload(tCtx, tc, w, topicName, scheduler, informerFactory, opts)
 					if err != nil {
 						tCtx.Fatalf("Error running workload %s: %s", w.Name, err)
 					}
@@ -933,13 +942,13 @@ func RunIntegrationPerfScheduling(t *testing.T, configFile string, options ...Sc
 						t.Skipf("disabled by label filter %q", TestSchedulingLabelFilter)
 					}
 					featureGates := featureGatesMerge(tc.FeatureGates, w.FeatureGates)
-					scheduler, informerFactory, tCtx := setupTestCase(t, tc, featureGates, opts)
+					scheduler, informerFactory, tCtx := setupTestCase(t, tc, featureGates, w, opts)
 					err := w.isValid(tc.MetricsCollectorConfig)
 					if err != nil {
 						t.Fatalf("workload %s is not valid: %v", w.Name, err)
 					}
 
-					_, err = runWorkload(tCtx, tc, w, "" /* topic name not relevant */, scheduler, informerFactory)
+					_, err = runWorkload(tCtx, tc, w, "" /* topic name not relevant */, scheduler, informerFactory, opts)
 					if err != nil {
 						tCtx.Fatalf("Error running workload %s: %s", w.Name, err)
 					}
@@ -976,7 +985,7 @@ func loadSchedulerConfig(file string) (*config.KubeSchedulerConfiguration, error
 	return nil, fmt.Errorf("couldn't decode as KubeSchedulerConfiguration, got %s: ", gvk)
 }
 
-func unrollWorkloadTemplate(tb ktesting.TB, wt []op, w *workload) []op {
+func unrollWorkloadTemplate(tb ktesting.TB, wt []op, w *Workload) []op {
 	var unrolled []op
 	for opIndex, o := range wt {
 		realOp, err := o.realOp.patchParams(w)
@@ -1075,7 +1084,7 @@ func checkEmptyInFlightEvents() error {
 	return nil
 }
 
-func runWorkload(tCtx ktesting.TContext, tc *testCase, w *workload, topicName string, scheduler *scheduler.Scheduler, informerFactory informers.SharedInformerFactory) ([]DataItem, error) {
+func runWorkload(tCtx ktesting.TContext, tc *testCase, w *Workload, topicName string, scheduler *scheduler.Scheduler, informerFactory informers.SharedInformerFactory, opts *schedulerPerfOptions) ([]DataItem, error) {
 	b, benchmarking := tCtx.TB().(*testing.B)
 	if benchmarking {
 		start := time.Now()
@@ -1115,6 +1124,7 @@ func runWorkload(tCtx ktesting.TContext, tc *testCase, w *workload, topicName st
 		testCase:                     tc,
 		workload:                     w,
 		topicName:                    topicName,
+		opts:                         opts,
 	}
 
 	tCtx.TB().Cleanup(func() {

test/integration/scheduler_perf/update.go

@@ -73,7 +73,7 @@ func (c *updateAny) collectsMetrics() bool {
 	return false
 }
 
-func (c updateAny) patchParams(w *workload) (realOp, error) {
+func (c updateAny) patchParams(w *Workload) (realOp, error) {
 	if c.CountParam != "" {
 		count, err := w.Params.get(c.CountParam[1:])
 		if err != nil {