diff --git a/test/integration/scheduler_perf/option.go b/test/integration/scheduler_perf/option.go new file mode 100644 index 00000000000..6eaf7d6c778 --- /dev/null +++ b/test/integration/scheduler_perf/option.go @@ -0,0 +1,20 @@ +package benchmark + +import "k8s.io/kubernetes/test/utils/ktesting" + +type SchedulerPerfOption func(options *schedulerPerfOptions) + +// PrepareFn is a function that is called before the benchmarks run. +type PrepareFn func(tCtx ktesting.TContext) error + +type schedulerPerfOptions struct { + prepareFn PrepareFn +} + +// WithPrepareFn is the option to set a function that is called +// before the benchmarks run. (e.g. applying CRDs for custom plugins) +func WithPrepareFn(prepareFn PrepareFn) SchedulerPerfOption { + return func(s *schedulerPerfOptions) { + s.prepareFn = prepareFn + } +} diff --git a/test/integration/scheduler_perf/scheduler_perf.go b/test/integration/scheduler_perf/scheduler_perf.go index 19c1ddc352c..b6b8c030d38 100644 --- a/test/integration/scheduler_perf/scheduler_perf.go +++ b/test/integration/scheduler_perf/scheduler_perf.go @@ -1147,7 +1147,19 @@ func fixJSONOutput(b *testing.B) { // You can pass your own scheduler plugins via outOfTreePluginRegistry. // Also, you may want to put your plugins in PluginNames variable in this package // to collect metrics for them. -func RunBenchmarkPerfScheduling(b *testing.B, configFile string, topicName string, outOfTreePluginRegistry frameworkruntime.Registry) { +func RunBenchmarkPerfScheduling(b *testing.B, configFile string, topicName string, outOfTreePluginRegistry frameworkruntime.Registry, options ...SchedulerPerfOption) { + opts := &schedulerPerfOptions{ + prepareFn: nil, + } + + for _, option := range options { + option(opts) + } + + runBenchmarkPerSchedulingImpl(b, configFile, topicName, outOfTreePluginRegistry, opts) +} + +func runBenchmarkPerSchedulingImpl(b *testing.B, configFile string, topicName string, outOfTreePluginRegistry frameworkruntime.Registry, options *schedulerPerfOptions) { testCases, err := getTestCases(configFile) if err != nil { b.Fatal(err) @@ -1194,6 +1206,13 @@ func RunBenchmarkPerfScheduling(b *testing.B, configFile string, topicName strin b.Fatalf("workload %s is not valid: %v", w.Name, err) } + if options.prepareFn != nil { + err = options.prepareFn(tCtx) + if err != nil { + b.Fatalf("failed to run prepareFn: %v", err) + } + } + results, err := runWorkload(tCtx, tc, w, informerFactory) if err != nil { tCtx.Fatalf("Error running workload %s: %s", w.Name, err)