Modify kai scheduler and sheduler plugins

Signed-off-by: win5923 <[email protected]>

Author: win5923

ray-operator/config/samples/ray-job.volcano-scheduler-queue.yaml

@@ -11,7 +11,7 @@ spec:
 apiVersion: ray.io/v1
 kind: RayJob
 metadata:
-  name: rayjob-sample-2
+  name: rayjob-sample-0
   labels:
     ray.io/scheduler-name: volcano
     volcano.sh/queue-name: kuberay-test-queue
@@ -23,16 +23,10 @@ spec:
       - pendulum==2.1.2
     env_vars:
       counter_name: "test_counter"
-  # rayClusterSpec specifies the RayCluster instance to be created by the RayJob controller.
   rayClusterSpec:
-    rayVersion: '2.46.0' # should match the Ray version in the image of the containers
-    # Ray head pod template
+    rayVersion: '2.46.0'
     headGroupSpec:
-      # The `rayStartParams` are used to configure the `ray start` command.
-      # See https://github.com/ray-project/kuberay/blob/master/docs/guidance/rayStartParams.md for the default settings of `rayStartParams` in KubeRay.
-      # See https://docs.ray.io/en/latest/cluster/cli.html#ray-start for all available options in `rayStartParams`.
       rayStartParams: {}
-      #pod template
       template:
         spec:
           containers:
@@ -41,7 +35,7 @@ spec:
             ports:
             - containerPort: 6379
               name: gcs-server
-            - containerPort: 8265 # Ray dashboard
+            - containerPort: 8265
               name: dashboard
             - containerPort: 10001
               name: client
@@ -56,31 +50,22 @@ spec:
             - mountPath: /home/ray/samples
               name: code-sample
           volumes:
-          # You set volumes at the Pod level, then mount them into containers inside that Pod
           - name: code-sample
             configMap:
-              # Provide the name of the ConfigMap you want to mount.
               name: ray-job-code-sample
-              # An array of keys from the ConfigMap to create as files
               items:
               - key: sample_code.py
                 path: sample_code.py
     workerGroupSpecs:
-    # the pod replicas in this group typed worker
     - replicas: 2
       minReplicas: 2
       maxReplicas: 2
-      # logical group name, for this called small-group, also can be functional
       groupName: small-group
-      # The `rayStartParams` are used to configure the `ray start` command.
-      # See https://github.com/ray-project/kuberay/blob/master/docs/guidance/rayStartParams.md for the default settings of `rayStartParams` in KubeRay.
-      # See https://docs.ray.io/en/latest/cluster/cli.html#ray-start for all available options in `rayStartParams`.
       rayStartParams: {}
-      #pod template
       template:
         spec:
           containers:
-          - name: ray-worker # must consist of lower case alphanumeric characters or '-', and must start and end with an alphanumeric character (e.g. 'my-name',  or '123-abc'
+          - name: ray-worker
             image: rayproject/ray:2.46.0
             resources:
               limits:
@@ -89,23 +74,6 @@ spec:
               requests:
                 cpu: "1"
                 memory: "1Gi"
-
-  # SubmitterPodTemplate is the template for the pod that will run the `ray job submit` command against the RayCluster.
-  # If SubmitterPodTemplate is specified, the first container is assumed to be the submitter container.
-  # submitterPodTemplate:
-  #   spec:
-  #     restartPolicy: Never
-  #     containers:
-  #     - name: my-custom-rayjob-submitter-pod
-  #       image: rayproject/ray:2.46.0
-  #       # If Command is not specified, the correct command will be supplied at runtime using the RayJob spec `entrypoint` field.
-  #       # Specifying Command is not recommended.
-  #       # command: ["sh", "-c", "ray job submit --address=http://$RAY_DASHBOARD_ADDRESS --submission-id=$RAY_JOB_SUBMISSION_ID -- echo hello world"]
-
-
-######################Ray code sample#################################
-# this sample is from https://docs.ray.io/en/latest/cluster/job-submission.html#quick-start-example
-# it is mounted into the container and executed to show the Ray job at work
 ---
 apiVersion: v1
 kind: ConfigMap

ray-operator/controllers/ray/batchscheduler/interface/interface.go

@@ -3,14 +3,11 @@ package schedulerinterface
 import (
 	"context"
 
-	corev1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/runtime"
 	"k8s.io/client-go/rest"
 	"sigs.k8s.io/controller-runtime/pkg/builder"
 	"sigs.k8s.io/controller-runtime/pkg/client"
-
-	rayv1 "github.com/ray-project/kuberay/ray-operator/apis/ray/v1"
 )
 
 // BatchScheduler manages submitting RayCluster pods to a third-party scheduler.
@@ -23,11 +20,6 @@ type BatchScheduler interface {
 	// For most batch schedulers, this results in the creation of a PodGroup.
 	DoBatchSchedulingOnSubmission(ctx context.Context, object metav1.Object) error
 
-	// AddMetadataToPod enriches the pod with metadata necessary to tie it to the scheduler.
-	// For example, setting labels for queues / priority, and setting schedulerName.
-	// This function will be removed once Rayjob Volcano scheduler integration is completed.
-	AddMetadataToPod(ctx context.Context, rayCluster *rayv1.RayCluster, groupName string, pod *corev1.Pod)
-
 	// AddMetadataToChildResource enriches the child resource (batchv1.Job, rayv1.RayCluster) with metadata necessary to tie it to the scheduler.
 	// For example, setting labels for queues / priority, and setting schedulerName.
 	AddMetadataToChildResource(ctx context.Context, parent metav1.Object, child metav1.Object, groupName string)
@@ -63,9 +55,6 @@ func (d *DefaultBatchScheduler) DoBatchSchedulingOnSubmission(_ context.Context,
 	return nil
 }
 
-func (d *DefaultBatchScheduler) AddMetadataToPod(_ context.Context, _ *rayv1.RayCluster, _ string, _ *corev1.Pod) {
-}
-
 func (d *DefaultBatchScheduler) AddMetadataToChildResource(_ context.Context, _ metav1.Object, _ metav1.Object, _ string) {
 }
 

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

@@ -18,7 +18,6 @@ import (
 	"sigs.k8s.io/controller-runtime/pkg/builder"
 	"sigs.k8s.io/controller-runtime/pkg/client"
 
-	rayv1 "github.com/ray-project/kuberay/ray-operator/apis/ray/v1"
 	schedulerinterface "github.com/ray-project/kuberay/ray-operator/controllers/ray/batchscheduler/interface"
 )
 
@@ -38,23 +37,33 @@ func (k *KaiScheduler) DoBatchSchedulingOnSubmission(_ context.Context, _ metav1
 	return nil
 }
 
-func (k *KaiScheduler) AddMetadataToPod(ctx context.Context, app *rayv1.RayCluster, _ string, pod *corev1.Pod) {
+func (k *KaiScheduler) AddMetadataToChildResource(ctx context.Context, parent metav1.Object, child metav1.Object, _ string) {
 	logger := ctrl.LoggerFrom(ctx).WithName("kai-scheduler")
-	pod.Spec.SchedulerName = k.Name()
+	addSchedulerNameToObject(child, k.Name())
 
-	queue, ok := app.Labels[QueueLabelName]
+	parentLabel := parent.GetLabels()
+	queue, ok := parentLabel[QueueLabelName]
 	if !ok || queue == "" {
-		logger.Info("Queue label missing from RayCluster; pods will remain pending",
+		logger.Info("Queue label missing from parent; child will remain pending",
 			"requiredLabel", QueueLabelName)
 		return
 	}
-	if pod.Labels == nil {
-		pod.Labels = make(map[string]string)
+
+	childLabels := child.GetLabels()
+	if childLabels == nil {
+		childLabels = make(map[string]string)
 	}
-	pod.Labels[QueueLabelName] = queue
+	childLabels[QueueLabelName] = queue
+	child.SetLabels(childLabels)
 }
 
-func (k *KaiScheduler) AddMetadataToChildResource(_ context.Context, _ metav1.Object, _ metav1.Object, _ string) {
+func addSchedulerNameToObject(obj metav1.Object, schedulerName string) {
+	switch obj := obj.(type) {
+	case *corev1.Pod:
+		obj.Spec.SchedulerName = schedulerName
+	case *corev1.PodTemplateSpec:
+		obj.Spec.SchedulerName = schedulerName
+	}
 }
 
 func (kf *KaiSchedulerFactory) New(_ context.Context, _ *rest.Config, _ client.Client) (schedulerinterface.BatchScheduler, error) {

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

@@ -41,7 +41,7 @@ func createTestPod() *corev1.Pod {
 	}
 }
 
-func TestAddMetadataToPod_WithQueueLabel(t *testing.T) {
+func TestAddMetadataToChildResource_WithQueueLabel(t *testing.T) {
 	a := assert.New(t)
 	scheduler := &KaiScheduler{}
 	ctx := context.Background()
@@ -52,8 +52,8 @@ func TestAddMetadataToPod_WithQueueLabel(t *testing.T) {
 	})
 	pod := createTestPod()
 
-	// Call AddMetadataToPod
-	scheduler.AddMetadataToPod(ctx, rayCluster, "test-group", pod)
+	// Call AddMetadataToChildResource
+	scheduler.AddMetadataToChildResource(ctx, rayCluster, pod, "test-group")
 
 	// Assert scheduler name is set to kai-scheduler
 	a.Equal("kai-scheduler", pod.Spec.SchedulerName)
@@ -63,7 +63,7 @@ func TestAddMetadataToPod_WithQueueLabel(t *testing.T) {
 	a.Equal("test-queue", pod.Labels[QueueLabelName])
 }
 
-func TestAddMetadataToPod_WithoutQueueLabel(t *testing.T) {
+func TestAddMetadataToChildResource_WithoutQueueLabel(t *testing.T) {
 	a := assert.New(t)
 	scheduler := &KaiScheduler{}
 	ctx := context.Background()
@@ -72,8 +72,8 @@ func TestAddMetadataToPod_WithoutQueueLabel(t *testing.T) {
 	rayCluster := createTestRayCluster(map[string]string{})
 	pod := createTestPod()
 
-	// Call AddMetadataToPod
-	scheduler.AddMetadataToPod(ctx, rayCluster, "test-group", pod)
+	// Call AddMetadataToChildResource
+	scheduler.AddMetadataToChildResource(ctx, rayCluster, pod, "test-group")
 
 	// Assert scheduler name is still set (always required)
 	a.Equal("kai-scheduler", pod.Spec.SchedulerName)
@@ -85,7 +85,7 @@ func TestAddMetadataToPod_WithoutQueueLabel(t *testing.T) {
 	}
 }
 
-func TestAddMetadataToPod_WithEmptyQueueLabel(t *testing.T) {
+func TestAddMetadataToChildResource_WithEmptyQueueLabel(t *testing.T) {
 	a := assert.New(t)
 	scheduler := &KaiScheduler{}
 	ctx := context.Background()
@@ -96,8 +96,8 @@ func TestAddMetadataToPod_WithEmptyQueueLabel(t *testing.T) {
 	})
 	pod := createTestPod()
 
-	// Call AddMetadataToPod
-	scheduler.AddMetadataToPod(ctx, rayCluster, "test-group", pod)
+	// Call AddMetadataToChildResource
+	scheduler.AddMetadataToChildResource(ctx, rayCluster, pod, "test-group")
 
 	// Assert scheduler name is still set
 	a.Equal("kai-scheduler", pod.Spec.SchedulerName)
@@ -109,7 +109,7 @@ func TestAddMetadataToPod_WithEmptyQueueLabel(t *testing.T) {
 	}
 }
 
-func TestAddMetadataToPod_PreservesExistingPodLabels(t *testing.T) {
+func TestAddMetadataToChildResource_PreservesExistingPodLabels(t *testing.T) {
 	a := assert.New(t)
 	scheduler := &KaiScheduler{}
 	ctx := context.Background()
@@ -126,8 +126,8 @@ func TestAddMetadataToPod_PreservesExistingPodLabels(t *testing.T) {
 		"app":            "ray",
 	}
 
-	// Call AddMetadataToPod
-	scheduler.AddMetadataToPod(ctx, rayCluster, "test-group", pod)
+	// Call AddMetadataToChildResource
+	scheduler.AddMetadataToChildResource(ctx, rayCluster, pod, "test-group")
 
 	// Assert scheduler name is set
 	a.Equal("kai-scheduler", pod.Spec.SchedulerName)

ray-operator/controllers/ray/batchscheduler/scheduler-plugins/scheduler_plugins.go

@@ -93,21 +93,30 @@ func (k *KubeScheduler) DoBatchSchedulingOnSubmission(ctx context.Context, objec
 	return nil
 }
 
-// AddMetadataToPod adds essential labels and annotations to the Ray pod
-// the scheduler needs these labels and annotations in order to do the scheduling properly
-func (k *KubeScheduler) AddMetadataToPod(_ context.Context, rayCluster *rayv1.RayCluster, _ string, pod *corev1.Pod) {
-	// when gang scheduling is enabled, extra labels need to be added to all pods
-	if k.isGangSchedulingEnabled(rayCluster) {
-		pod.Labels[kubeSchedulerPodGroupLabelKey] = rayCluster.Name
+func (k *KubeScheduler) AddMetadataToChildResource(_ context.Context, parent metav1.Object, child metav1.Object, _ string) {
+	// when gang scheduling is enabled, extra labels need to be added to all child resources
+	if k.isGangSchedulingEnabled(parent) {
+		labels := child.GetLabels()
+		if labels == nil {
+			labels = make(map[string]string)
+		}
+		labels[kubeSchedulerPodGroupLabelKey] = parent.GetName()
+		child.SetLabels(labels)
 	}
-	pod.Spec.SchedulerName = k.Name()
+	addSchedulerNameToObject(child, k.Name())
 }
 
-func (k *KubeScheduler) AddMetadataToChildResource(_ context.Context, _ metav1.Object, _ metav1.Object, _ string) {
+func addSchedulerNameToObject(obj metav1.Object, schedulerName string) {
+	switch obj := obj.(type) {
+	case *corev1.Pod:
+		obj.Spec.SchedulerName = schedulerName
+	case *corev1.PodTemplateSpec:
+		obj.Spec.SchedulerName = schedulerName
+	}
 }
 
-func (k *KubeScheduler) isGangSchedulingEnabled(app *rayv1.RayCluster) bool {
-	_, exist := app.Labels[utils.RayGangSchedulingEnabled]
+func (k *KubeScheduler) isGangSchedulingEnabled(obj metav1.Object) bool {
+	_, exist := obj.GetLabels()[utils.RayGangSchedulingEnabled]
 	return exist
 }
 

ray-operator/controllers/ray/batchscheduler/scheduler-plugins/scheduler_plugins_test.go

@@ -117,7 +117,7 @@ func TestCreatePodGroupWithMultipleHosts(t *testing.T) {
 	a.Equal(int32(5), podGroup.Spec.MinMember)
 }
 
-func TestAddMetadataToPod(t *testing.T) {
+func TestAddMetadataToChildResource(t *testing.T) {
 	tests := []struct {
 		name         string
 		enableGang   bool
@@ -150,7 +150,7 @@ func TestAddMetadataToPod(t *testing.T) {
 			}
 
 			scheduler := &KubeScheduler{}
-			scheduler.AddMetadataToPod(context.TODO(), &cluster, "worker", pod)
+			scheduler.AddMetadataToChildResource(context.TODO(), &cluster, pod, "worker")
 
 			if tt.enableGang {
 				a.Equal(cluster.Name, pod.Labels[kubeSchedulerPodGroupLabelKey])