Integration: KAI Scheduler

helm-chart/kuberay-operator/values.yaml

@@ -70,12 +70,16 @@ logging:
 #  4. Use PodGroup
 #       batchScheduler:
 #         name: scheduler-plugins
-#
+
+#  5. Use Kai Scheduler
+#       batchScheduler:
+#         name: kai-scheduler
+
 batchScheduler:
   # Deprecated. This option will be removed in the future.
   # Note, for backwards compatibility. When it sets to true, it enables volcano scheduler integration.
   enabled: false
-  # Set the customized scheduler name, supported values are "volcano", "yunikorn" or "scheduler-plugins", do not set
+  # Set the customized scheduler name, supported values are "volcano", "yunikorn", "kai-scheduler" or "scheduler-plugins", do not set
   # "batchScheduler.enabled=true" at the same time as it will override this option.
   name: ""
 

ray-operator/apis/config/v1alpha1/config_utils.go

@@ -5,6 +5,7 @@ import (
 
 	"github.com/go-logr/logr"
 
+	kaischeduler "github.com/ray-project/kuberay/ray-operator/controllers/ray/batchscheduler/kai-scheduler"
 	schedulerplugins "github.com/ray-project/kuberay/ray-operator/controllers/ray/batchscheduler/scheduler-plugins"
 	"github.com/ray-project/kuberay/ray-operator/controllers/ray/batchscheduler/volcano"
 	"github.com/ray-project/kuberay/ray-operator/controllers/ray/batchscheduler/yunikorn"
@@ -23,7 +24,7 @@ func ValidateBatchSchedulerConfig(logger logr.Logger, config Configuration) erro
 
 	if len(config.BatchScheduler) > 0 {
 		// if a customized scheduler is configured, check it is supported
-		if config.BatchScheduler == volcano.GetPluginName() || config.BatchScheduler == yunikorn.GetPluginName() || config.BatchScheduler == schedulerplugins.GetPluginName() {
+		if config.BatchScheduler == volcano.GetPluginName() || config.BatchScheduler == yunikorn.GetPluginName() || config.BatchScheduler == schedulerplugins.GetPluginName() || config.BatchScheduler == kaischeduler.GetPluginName() {
 			logger.Info("Feature flag batch-scheduler is enabled",
 				"scheduler name", config.BatchScheduler)
 		} else {

ray-operator/apis/config/v1alpha1/config_utils_test.go

@@ -6,6 +6,7 @@ import (
 	"github.com/go-logr/logr"
 	"github.com/go-logr/logr/testr"
 
+	kaischeduler "github.com/ray-project/kuberay/ray-operator/controllers/ray/batchscheduler/kai-scheduler"
 	schedulerPlugins "github.com/ray-project/kuberay/ray-operator/controllers/ray/batchscheduler/scheduler-plugins"
 	"github.com/ray-project/kuberay/ray-operator/controllers/ray/batchscheduler/volcano"
 	"github.com/ray-project/kuberay/ray-operator/controllers/ray/batchscheduler/yunikorn"
@@ -71,6 +72,16 @@ func TestValidateBatchSchedulerConfig(t *testing.T) {
 			},
 			wantErr: false,
 		},
+		{
+			name: "valid option, batch-scheduler=kai-scheduler",
+			args: args{
+				logger: testr.New(t),
+				config: Configuration{
+					BatchScheduler: kaischeduler.GetPluginName(),
+				},
+			},
+			wantErr: false,
+		},
 		{
 			name: "invalid option, invalid scheduler name",
 			args: args{

ray-operator/apis/config/v1alpha1/configuration_types.go

@@ -44,7 +44,7 @@ type Configuration struct {
 	LogStdoutEncoder string `json:"logStdoutEncoder,omitempty"`
 
 	// BatchScheduler enables the batch scheduler integration with a specific scheduler
-	// based on the given name, currently, supported values are volcano and yunikorn.
+	// based on the given name, currently, supported values are volcano, yunikorn, kai-scheduler.
 	BatchScheduler string `json:"batchScheduler,omitempty"`
 
 	// HeadSidecarContainers includes specification for a sidecar container

ray-operator/config/samples/ray-cluster.kai-gpu-sharing.yaml

@@ -0,0 +1,35 @@
+apiVersion: ray.io/v1
+kind: RayCluster
+metadata:
+  name: raycluster-half-gpu
+  labels:
+    kai.scheduler/queue: team-a
+spec:
+  headGroupSpec:
+    template:
+      spec:
+        containers:
+        - name: head
+          image: rayproject/ray:2.46.0
+          resources:
+            limits:
+              cpu: "1"
+              memory: "2Gi"
+
+  # ---- Two workers share one GPU (0.5 each) ----
+  workerGroupSpecs:
+  - groupName: shared-gpu
+    replicas: 2
+    minReplicas: 2
+    template:
+      metadata:
+        annotations:
+          gpu-fraction: "0.5"
+      spec:
+        containers:
+        - name: worker
+          image: rayproject/ray:2.46.0
+          resources:
+            limits:
+              cpu: "1"
+              memory: "2Gi"

ray-operator/config/samples/ray-cluster.kai-scheduler-queues.yaml

@@ -0,0 +1,38 @@
+apiVersion: scheduling.run.ai/v2
+kind: Queue
+metadata:
+  name: department-1
+spec:
+  resources:
+    cpu:
+      quota: -1
+      limit: -1
+      overQuotaWeight: 1
+    gpu:
+      quota: -1
+      limit: -1
+      overQuotaWeight: 1
+    memory:
+      quota: -1
+      limit: -1
+      overQuotaWeight: 1
+---
+apiVersion: scheduling.run.ai/v2
+kind: Queue
+metadata:
+  name: team-a
+spec:
+  parentQueue: department-1
+  resources:
+    cpu:
+      quota: -1
+      limit: -1
+      overQuotaWeight: 1
+    gpu:
+      quota: -1
+      limit: -1
+      overQuotaWeight: 1
+    memory:
+      quota: -1
+      limit: -1
+      overQuotaWeight: 1

ray-operator/config/samples/ray-cluster.kai-scheduler.yaml

@@ -0,0 +1,31 @@
+#A simple example raycluster with KAI
+apiVersion: ray.io/v1
+kind: RayCluster
+metadata:
+  name: raycluster-sample
+  labels:
+    kai.scheduler/queue: team-a
+spec:
+  headGroupSpec:
+    template:
+      spec:
+        containers:
+        - name: ray-head
+          image: rayproject/ray:2.46.0
+          resources:
+            requests:
+              cpu: "1"
+              memory: "2Gi"
+  workerGroupSpecs:
+  - groupName: worker
+    replicas: 2
+    minReplicas: 2
+    template:
+      spec:
+        containers:
+        - name: ray-worker
+          image: rayproject/ray:2.46.0
+          resources:
+            requests:
+              cpu: "1"
+              memory: "1Gi"

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

@@ -0,0 +1,66 @@
+package kaischeduler
+
+// This KAI plugin relies on KAI-Scheduler's
+// built-in PodGrouper to create PodGroups at
+// runtime, so the plugin itself only needs to:
+//   1. expose the scheduler name,
+//   2. stamp pods with schedulerName + queue label.
+// No PodGroup create/patch logic is included.
+
+import (
+	"context"
+
+	corev1 "k8s.io/api/core/v1"
+	"k8s.io/apimachinery/pkg/runtime"
+	"k8s.io/client-go/rest"
+	ctrl "sigs.k8s.io/controller-runtime"
+	"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"
+)
+
+const (
+	QueueLabelName = "kai.scheduler/queue"
+)
+
+type KaiScheduler struct{}
+
+type KaiSchedulerFactory struct{}
+
+func GetPluginName() string { return "kai-scheduler" }
+
+func (k *KaiScheduler) Name() string { return GetPluginName() }
+
+func (k *KaiScheduler) DoBatchSchedulingOnSubmission(_ context.Context, _ *rayv1.RayCluster) error {
+	return nil
+}
+
+func (k *KaiScheduler) AddMetadataToPod(ctx context.Context, app *rayv1.RayCluster, _ string, pod *corev1.Pod) {
+	pod.Spec.SchedulerName = k.Name()
+
+	queue, ok := app.Labels[QueueLabelName]
+	if !ok || queue == "" {
+		logger := ctrl.LoggerFrom(ctx).WithName("kai-scheduler")
+		logger.Info("Queue label missing from RayCluster; pods will remain pending",
+			"requiredLabel", QueueLabelName,
+			"rayCluster", app.Name)
+		return
+	}
+	if pod.Labels == nil {
+		pod.Labels = make(map[string]string)
+	}
+	pod.Labels[QueueLabelName] = queue
+}
+
+func (kf *KaiSchedulerFactory) New(_ context.Context, _ *rest.Config, _ client.Client) (schedulerinterface.BatchScheduler, error) {
+	return &KaiScheduler{}, nil
+}
+
+func (kf *KaiSchedulerFactory) AddToScheme(_ *runtime.Scheme) {
+}
+
+func (kf *KaiSchedulerFactory) ConfigureReconciler(b *builder.Builder) *builder.Builder {
+	return b
+}

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

@@ -0,0 +1,141 @@
+package kaischeduler
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	corev1 "k8s.io/api/core/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+
+	rayv1 "github.com/ray-project/kuberay/ray-operator/apis/ray/v1"
+)
+
+func createTestRayCluster(labels map[string]string) *rayv1.RayCluster {
+	return &rayv1.RayCluster{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      "test-cluster",
+			Namespace: "default",
+			Labels:    labels,
+		},
+	}
+}
+
+func createTestPod() *corev1.Pod {
+	return &corev1.Pod{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      "test-pod",
+			Namespace: "default",
+			Labels: map[string]string{
+				"ray.io/cluster":   "test-cluster",
+				"ray.io/node-type": "worker",
+				"app":              "ray",
+			},
+		},
+		Spec: corev1.PodSpec{
+			Containers: []corev1.Container{{
+				Name:  "ray-worker",
+				Image: "rayproject/ray:latest",
+			}},
+		},
+	}
+}
+
+func TestAddMetadataToPod_WithQueueLabel(t *testing.T) {
+	a := assert.New(t)
+	scheduler := &KaiScheduler{}
+	ctx := context.Background()
+
+	// Create RayCluster with queue label
+	rayCluster := createTestRayCluster(map[string]string{
+		QueueLabelName: "test-queue",
+	})
+	pod := createTestPod()
+
+	// Call AddMetadataToPod
+	scheduler.AddMetadataToPod(ctx, rayCluster, "test-group", pod)
+
+	// Assert scheduler name is set to kai-scheduler
+	a.Equal("kai-scheduler", pod.Spec.SchedulerName)
+
+	// Assert queue label is propagated to pod
+	a.NotNil(pod.Labels)
+	a.Equal("test-queue", pod.Labels[QueueLabelName])
+}
+
+func TestAddMetadataToPod_WithoutQueueLabel(t *testing.T) {
+	a := assert.New(t)
+	scheduler := &KaiScheduler{}
+	ctx := context.Background()
+
+	// Create RayCluster without queue label
+	rayCluster := createTestRayCluster(map[string]string{})
+	pod := createTestPod()
+
+	// Call AddMetadataToPod
+	scheduler.AddMetadataToPod(ctx, rayCluster, "test-group", pod)
+
+	// Assert scheduler name is still set (always required)
+	a.Equal("kai-scheduler", pod.Spec.SchedulerName)
+
+	// Assert queue label is not added to pod when missing from RayCluster
+	if pod.Labels != nil {
+		_, exists := pod.Labels[QueueLabelName]
+		a.False(exists)
+	}
+}
+
+func TestAddMetadataToPod_WithEmptyQueueLabel(t *testing.T) {
+	a := assert.New(t)
+	scheduler := &KaiScheduler{}
+	ctx := context.Background()
+
+	// Create RayCluster with empty queue label
+	rayCluster := createTestRayCluster(map[string]string{
+		QueueLabelName: "",
+	})
+	pod := createTestPod()
+
+	// Call AddMetadataToPod
+	scheduler.AddMetadataToPod(ctx, rayCluster, "test-group", pod)
+
+	// Assert scheduler name is still set
+	a.Equal("kai-scheduler", pod.Spec.SchedulerName)
+
+	// Assert empty queue label is treated as missing
+	if pod.Labels != nil {
+		_, exists := pod.Labels[QueueLabelName]
+		a.False(exists)
+	}
+}
+
+func TestAddMetadataToPod_PreservesExistingPodLabels(t *testing.T) {
+	a := assert.New(t)
+	scheduler := &KaiScheduler{}
+	ctx := context.Background()
+
+	// Create RayCluster with queue label
+	rayCluster := createTestRayCluster(map[string]string{
+		QueueLabelName: "test-queue",
+	})
+
+	// Create pod with existing labels
+	pod := createTestPod()
+	pod.Labels = map[string]string{
+		"existing-label": "existing-value",
+		"app":            "ray",
+	}
+
+	// Call AddMetadataToPod
+	scheduler.AddMetadataToPod(ctx, rayCluster, "test-group", pod)
+
+	// Assert scheduler name is set
+	a.Equal("kai-scheduler", pod.Spec.SchedulerName)
+
+	// Assert queue label is added
+	a.Equal("test-queue", pod.Labels[QueueLabelName])
+
+	// Assert existing labels are preserved
+	a.Equal("existing-value", pod.Labels["existing-label"])
+	a.Equal("ray", pod.Labels["app"])
+}