update

Signed-off-by: You-Cheng Lin (Owen) <[email protected]>

Author: owenowenisme

ray-operator/controllers/ray/batchscheduler/kai-scheduler/kai_scheduler.go

@@ -9,7 +9,6 @@ package kaischeduler
 
 import (
 	"context"
-	"fmt"
 
 	"github.com/go-logr/logr"
 	corev1 "k8s.io/api/core/v1"
@@ -37,13 +36,7 @@ func GetPluginName() string { return "kai-scheduler" }
 
 func (k *KaiScheduler) Name() string { return GetPluginName() }
 
-func (k *KaiScheduler) DoBatchSchedulingOnSubmission(_ context.Context, object client.Object) error {
-	_, ok := object.(*rayv1.RayCluster)
-	if !ok {
-		return fmt.Errorf("currently only RayCluster is supported, got %T", object)
-	}
-	// yunikorn doesn't require any resources to be created upfront
-	// this is a no-opt for this implementation
+func (k *KaiScheduler) DoBatchSchedulingOnSubmission(_ context.Context, _ client.Object) error {
 	return nil
 }
 

ray-operator/controllers/ray/batchscheduler/yunikorn/yunikorn_scheduler.go

@@ -2,7 +2,6 @@ package yunikorn
 
 import (
 	"context"
-	"fmt"
 
 	"github.com/go-logr/logr"
 	corev1 "k8s.io/api/core/v1"
@@ -41,13 +40,7 @@ func (y *YuniKornScheduler) Name() string {
 	return GetPluginName()
 }
 
-func (y *YuniKornScheduler) DoBatchSchedulingOnSubmission(_ context.Context, object client.Object) error {
-	switch obj := object.(type) {
-	case *rayv1.RayCluster, *rayv1.RayJob:
-		// Supported types
-	default:
-		return fmt.Errorf("currently only RayCluster and RayJob are supported, got %T", obj)
-	}
+func (y *YuniKornScheduler) DoBatchSchedulingOnSubmission(_ context.Context, _ client.Object) error {
 	// yunikorn doesn't require any resources to be created upfront
 	// this is a no-opt for this implementation
 	return nil
@@ -59,9 +52,9 @@ func (y *YuniKornScheduler) DoBatchSchedulingOnSubmission(_ context.Context, obj
 // Currently we use this function to translate labels "yunikorn.apache.org/app-id" and "yunikorn.apache.org/queue"
 // to legacy labels "applicationId" and "queue", this is for the better compatibilities to support older yunikorn
 // versions.
-func (y *YuniKornScheduler) populatePodLabelsFromRayCluster(_ context.Context, app *rayv1.RayCluster, pod *corev1.Pod, sourceKey string, targetKey string) {
+func (y *YuniKornScheduler) populatePodLabelsFromRayCluster(_ context.Context, rayCluster *rayv1.RayCluster, pod *corev1.Pod, sourceKey string, targetKey string) {
 	// check labels
-	if value, exist := app.Labels[sourceKey]; exist {
+	if value, exist := rayCluster.Labels[sourceKey]; exist {
 		y.logger.Info("Updating pod label based on RayCluster labels",
 			"sourceKey", sourceKey, "targetKey", targetKey, "value", value)
 		pod.Labels[targetKey] = value
@@ -144,14 +137,14 @@ func (y *YuniKornScheduler) isGangSchedulingEnabled(obj client.Object) bool {
 	}
 }
 
-func (y *YuniKornScheduler) populateTaskGroupsAnnotationToPod(_ context.Context, app *rayv1.RayCluster, pod *corev1.Pod) {
+func (y *YuniKornScheduler) populateTaskGroupsAnnotationToPod(_ context.Context, rayCluster *rayv1.RayCluster, pod *corev1.Pod) {
 	var taskGroupsAnnotationValue string
 	var err error
-	if app.Annotations[YuniKornTaskGroupsAnnotationName] != "" {
-		taskGroupsAnnotationValue = app.Annotations[YuniKornTaskGroupsAnnotationName]
+	if rayCluster.Annotations[YuniKornTaskGroupsAnnotationName] != "" {
+		taskGroupsAnnotationValue = rayCluster.Annotations[YuniKornTaskGroupsAnnotationName]
 		y.logger.Info("using existing task groups annotation from RayCluster", "value", taskGroupsAnnotationValue)
 	} else {
-		taskGroups := newTaskGroupsFromApp(app)
+		taskGroups := newTaskGroupsFromRayCluster(rayCluster)
 		taskGroupsAnnotationValue, err = taskGroups.marshal()
 		if err != nil {
 			y.logger.Error(err, "failed to add gang scheduling related annotations to pod, "+

ray-operator/controllers/ray/batchscheduler/yunikorn/yunikorn_scheduler_test.go

@@ -391,8 +391,7 @@ func createRayJobWithLabels(name string, namespace string, rayClusterSpec *rayv1
 	return rayJob
 }
 
-func addHeadPodSpec(app *rayv1.RayCluster, resource corev1.ResourceList) {
-	// app.Spec.HeadGroupSpec.Template.Spec.Containers
+func addHeadPodSpec(rayCluster *rayv1.RayCluster, resource corev1.ResourceList) {
 	headContainers := []corev1.Container{
 		{
 			Name:  "head-pod",
@@ -404,10 +403,10 @@ func addHeadPodSpec(app *rayv1.RayCluster, resource corev1.ResourceList) {
 		},
 	}
 
-	app.Spec.HeadGroupSpec.Template.Spec.Containers = headContainers
+	rayCluster.Spec.HeadGroupSpec.Template.Spec.Containers = headContainers
 }
 
-func addWorkerPodSpec(app *rayv1.RayCluster, workerGroupName string,
+func addWorkerPodSpec(rayCluster *rayv1.RayCluster, workerGroupName string,
 	replicas int32, minReplicas int32, maxReplicas int32, resources corev1.ResourceList,
 ) {
 	workerContainers := []corev1.Container{
@@ -421,7 +420,7 @@ func addWorkerPodSpec(app *rayv1.RayCluster, workerGroupName string,
 		},
 	}
 
-	app.Spec.WorkerGroupSpecs = append(app.Spec.WorkerGroupSpecs, rayv1.WorkerGroupSpec{
+	rayCluster.Spec.WorkerGroupSpecs = append(rayCluster.Spec.WorkerGroupSpecs, rayv1.WorkerGroupSpec{
 		GroupName:   workerGroupName,
 		Replicas:    &replicas,
 		MinReplicas: &minReplicas,

ray-operator/controllers/ray/batchscheduler/yunikorn/yunikorn_task_groups.go

@@ -34,11 +34,11 @@ func newTaskGroups() *TaskGroups {
 	}
 }
 
-func newTaskGroupsFromApp(app *v1.RayCluster) *TaskGroups {
+func newTaskGroupsFromRayCluster(rayCluster *v1.RayCluster) *TaskGroups {
 	taskGroups := newTaskGroups()
 
 	// head group
-	headGroupSpec := app.Spec.HeadGroupSpec
+	headGroupSpec := rayCluster.Spec.HeadGroupSpec
 	headPodMinResource := utils.CalculatePodResource(headGroupSpec.Template.Spec)
 	taskGroups.addTaskGroup(
 		TaskGroup{
@@ -51,7 +51,7 @@ func newTaskGroupsFromApp(app *v1.RayCluster) *TaskGroups {
 		})
 
 	// worker groups
-	for _, workerGroupSpec := range app.Spec.WorkerGroupSpecs {
+	for _, workerGroupSpec := range rayCluster.Spec.WorkerGroupSpecs {
 		workerMinResource := utils.CalculatePodResource(workerGroupSpec.Template.Spec)
 		minWorkers := workerGroupSpec.MinReplicas
 		taskGroups.addTaskGroup(
@@ -68,11 +68,11 @@ func newTaskGroupsFromApp(app *v1.RayCluster) *TaskGroups {
 	return taskGroups
 }
 
-func newTaskGroupsFromRayJob(app *v1.RayJob, submitterTemplate *corev1.PodTemplateSpec) (*TaskGroups, error) {
+func newTaskGroupsFromRayJob(rayJob *v1.RayJob, submitterTemplate *corev1.PodTemplateSpec) (*TaskGroups, error) {
 	taskGroups := newTaskGroups()
 
 	// head group
-	headGroupSpec := app.Spec.RayClusterSpec.HeadGroupSpec
+	headGroupSpec := rayJob.Spec.RayClusterSpec.HeadGroupSpec
 	headPodMinResource := utils.CalculatePodResource(headGroupSpec.Template.Spec)
 	taskGroups.addTaskGroup(
 		TaskGroup{
@@ -85,7 +85,7 @@ func newTaskGroupsFromRayJob(app *v1.RayJob, submitterTemplate *corev1.PodTempla
 		})
 
 	// worker groups
-	for _, workerGroupSpec := range app.Spec.RayClusterSpec.WorkerGroupSpecs {
+	for _, workerGroupSpec := range rayJob.Spec.RayClusterSpec.WorkerGroupSpecs {
 		workerMinResource := utils.CalculatePodResource(workerGroupSpec.Template.Spec)
 		minWorkers := workerGroupSpec.MinReplicas
 		taskGroups.addTaskGroup(