Fix unit test

Signed-off-by: Andrey Velichkevich <andrey.velichkevich@gmail.com>
kubeflow · Feb 12, 2025 · 2bb5bea · 2bb5bea
1 parent e29ff98
commit 2bb5bea
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/pkg/runtime/framework/core/framework_test.go b/pkg/runtime/framework/core/framework_test.go
@@ -86,10 +86,12 @@ func TestNew(t *testing.T) {
 				watchExtensionPlugins: []framework.WatchExtensionPlugin{
 					&coscheduling.CoScheduling{},
 					&jobset.JobSet{},
+					&mpi.MPI{},
 				},
 				componentBuilderPlugins: []framework.ComponentBuilderPlugin{
 					&coscheduling.CoScheduling{},
 					&jobset.JobSet{},
+					&mpi.MPI{},
 				},
 				terminalConditionPlugins: []framework.TerminalConditionPlugin{
 					&jobset.JobSet{},
@@ -528,11 +530,12 @@ func TestWatchExtensionPlugins(t *testing.T) {
 		registry    fwkplugins.Registry
 		wantPlugins []framework.WatchExtensionPlugin
 	}{
-		"coscheding and jobset are performed": {
+		"coscheding, jobset, and mpi are performed": {
 			registry: fwkplugins.NewRegistry(),
 			wantPlugins: []framework.WatchExtensionPlugin{
 				&coscheduling.CoScheduling{},
 				&jobset.JobSet{},
+				&mpi.MPI{},
 			},
 		},
 		"an empty registry": {
@@ -541,7 +544,7 @@ func TestWatchExtensionPlugins(t *testing.T) {
 	}
 	cmpOpts := []cmp.Option{
 		cmpopts.SortSlices(func(a, b framework.Plugin) bool { return a.Name() < b.Name() }),
-		cmpopts.IgnoreUnexported(coscheduling.CoScheduling{}, jobset.JobSet{}),
+		cmpopts.IgnoreUnexported(coscheduling.CoScheduling{}, jobset.JobSet{}, mpi.MPI{}),
 	}
 	for name, tc := range cases {
 		t.Run(name, func(t *testing.T) {

diff --git a/pkg/runtime/framework/plugins/jobset/builder.go b/pkg/runtime/framework/plugins/jobset/builder.go
@@ -137,7 +137,7 @@ func (b *Builder) Initializer(trainJob *trainer.TrainJob) *Builder {
 }
 
 // Launcher updates JobSet values for the launcher Job.
-func (b *Builder) Launcher(info *runtime.Info, trainJob *kubeflowv2.TrainJob) *Builder {
+func (b *Builder) Launcher(info *runtime.Info, trainJob *trainer.TrainJob) *Builder {
 	for i, rJob := range b.Spec.ReplicatedJobs {
 		if rJob.Name == constants.JobLauncher {