Set schedulerName to Yunikorn (kubeflow#2153)

Signed-off-by: Jacob Salway <jacob.salway@gmail.com>
jbhalodia-slack · Sep 8, 2024 · 62b4ca6 · 62b4ca6
1 parent c810ece
commit 62b4ca6
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 0 deletions.
diff --git a/internal/scheduler/yunikorn/scheduler.go b/internal/scheduler/yunikorn/scheduler.go
@@ -109,6 +109,11 @@ func (s *Scheduler) Schedule(app *v1beta2.SparkApplication) error {
 		})
 	}
 
+	// Ensure that the driver and executors pods are scheduled by Yunikorn
+	// if it is installed with the admissions controller disabled
+	app.Spec.Driver.SchedulerName = util.StringPtr(SchedulerName)
+	app.Spec.Executor.SchedulerName = util.StringPtr(SchedulerName)
+
 	// Yunikorn re-uses the application ID set by the driver under the label "spark-app-selector",
 	// so there is no need to set an application ID
 	// https://github.com/apache/yunikorn-k8shim/blob/2278b3217c702ccb796e4d623bc7837625e5a4ec/pkg/common/utils/utils.go#L168-L171

diff --git a/internal/scheduler/yunikorn/scheduler_test.go b/internal/scheduler/yunikorn/scheduler_test.go
@@ -242,6 +242,9 @@ func TestSchedule(t *testing.T) {
 				assert.Equal(t, *options.Queue, tc.app.Spec.Driver.Labels[queueLabel])
 				assert.Equal(t, *options.Queue, tc.app.Spec.Executor.Labels[queueLabel])
 			}
+
+			assert.Equal(t, "yunikorn", *tc.app.Spec.Driver.SchedulerName)
+			assert.Equal(t, "yunikorn", *tc.app.Spec.Executor.SchedulerName)
 		})
 	}
 }