Skip to content

Commit ab4fa93

Browse files
EkinKarabulut金嘉淇10353800
EkinKarabulut
authored and
金嘉淇10353800
committed
Integration: KAI Scheduler (ray-project#3886)
1 parent 52e027a commit ab4fa93

File tree

11 files changed

+444
-5
lines changed

11 files changed

+444
-5
lines changed

‎helm-chart/kuberay-operator/values.yaml

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -70,12 +70,16 @@ logging:
7070
# 4. Use PodGroup
7171
# batchScheduler:
7272
# name: scheduler-plugins
73-
#
73+
74+
# 5. Use Kai Scheduler
75+
# batchScheduler:
76+
# name: kai-scheduler
77+
7478
batchScheduler:
7579
# Deprecated. This option will be removed in the future.
7680
# Note, for backwards compatibility. When it sets to true, it enables volcano scheduler integration.
7781
enabled: false
78-
# Set the customized scheduler name, supported values are "volcano", "yunikorn" or "scheduler-plugins", do not set
82+
# Set the customized scheduler name, supported values are "volcano", "yunikorn", "kai-scheduler" or "scheduler-plugins", do not set
7983
# "batchScheduler.enabled=true" at the same time as it will override this option.
8084
name: ""
8185

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ import (
55

66
"github.com/go-logr/logr"
77

8+
kaischeduler "github.com/ray-project/kuberay/ray-operator/controllers/ray/batchscheduler/kai-scheduler"
89
schedulerplugins "github.com/ray-project/kuberay/ray-operator/controllers/ray/batchscheduler/scheduler-plugins"
910
"github.com/ray-project/kuberay/ray-operator/controllers/ray/batchscheduler/volcano"
1011
"github.com/ray-project/kuberay/ray-operator/controllers/ray/batchscheduler/yunikorn"
@@ -23,7 +24,7 @@ func ValidateBatchSchedulerConfig(logger logr.Logger, config Configuration) erro
2324

2425
if len(config.BatchScheduler) > 0 {
2526
// if a customized scheduler is configured, check it is supported
26-
if config.BatchScheduler == volcano.GetPluginName() || config.BatchScheduler == yunikorn.GetPluginName() || config.BatchScheduler == schedulerplugins.GetPluginName() {
27+
if config.BatchScheduler == volcano.GetPluginName() || config.BatchScheduler == yunikorn.GetPluginName() || config.BatchScheduler == schedulerplugins.GetPluginName() || config.BatchScheduler == kaischeduler.GetPluginName() {
2728
logger.Info("Feature flag batch-scheduler is enabled",
2829
"scheduler name", config.BatchScheduler)
2930
} else {

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

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import (
66
"github.com/go-logr/logr"
77
"github.com/go-logr/logr/testr"
88

9+
kaischeduler "github.com/ray-project/kuberay/ray-operator/controllers/ray/batchscheduler/kai-scheduler"
910
schedulerPlugins "github.com/ray-project/kuberay/ray-operator/controllers/ray/batchscheduler/scheduler-plugins"
1011
"github.com/ray-project/kuberay/ray-operator/controllers/ray/batchscheduler/volcano"
1112
"github.com/ray-project/kuberay/ray-operator/controllers/ray/batchscheduler/yunikorn"
@@ -71,6 +72,16 @@ func TestValidateBatchSchedulerConfig(t *testing.T) {
7172
},
7273
wantErr: false,
7374
},
75+
{
76+
name: "valid option, batch-scheduler=kai-scheduler",
77+
args: args{
78+
logger: testr.New(t),
79+
config: Configuration{
80+
BatchScheduler: kaischeduler.GetPluginName(),
81+
},
82+
},
83+
wantErr: false,
84+
},
7485
{
7586
name: "invalid option, invalid scheduler name",
7687
args: args{

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,7 @@ type Configuration struct {
4444
LogStdoutEncoder string `json:"logStdoutEncoder,omitempty"`
4545

4646
// BatchScheduler enables the batch scheduler integration with a specific scheduler
47-
// based on the given name, currently, supported values are volcano and yunikorn.
47+
// based on the given name, currently, supported values are volcano, yunikorn, kai-scheduler.
4848
BatchScheduler string `json:"batchScheduler,omitempty"`
4949

5050
// HeadSidecarContainers includes specification for a sidecar container

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

Lines changed: 98 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,98 @@
1+
# KAI Scheduler Example - GPU Sharing with RayCluster
2+
# KAI Scheduler uses a hierarchical queue system for resource management and fair sharing.
3+
# These queues must be created before any RayCluster can be scheduled by KAI.
4+
5+
# NOTE: This is a DEMO configuration with unlimited quotas (-1) for easy testing.
6+
# In real-world deployments, you should set appropriate CPU/GPU/memory quotas and limits
7+
# based on your cluster's actual resources and organizational needs.
8+
9+
# GPU Sharing Note: This example utilizes time slicing GPU sharing.
10+
# KAI Scheduler also supports MPS and other GPU sharing methods.
11+
# For more information, check the KAI Scheduler documentation.
12+
13+
# Parent queue: Represents a department or high-level organizational unit
14+
apiVersion: scheduling.run.ai/v2
15+
kind: Queue
16+
metadata:
17+
name: department-1
18+
spec:
19+
# priority: 100 # Optional: Higher priority queues get surplus resources first
20+
resources:
21+
# quota: Guaranteed resources for this queue
22+
# limit: Maximum resources this queue can use
23+
# overQuotaWeight: How surplus resources are shared among queues
24+
# Note: Using -1 (unlimited) for demo purposes
25+
cpu:
26+
quota: -1
27+
limit: -1
28+
overQuotaWeight: 1
29+
gpu:
30+
quota: -1
31+
limit: -1
32+
overQuotaWeight: 1
33+
memory:
34+
quota: -1
35+
limit: -1
36+
overQuotaWeight: 1
37+
---
38+
# Child queue: Represents a team within the department-1
39+
apiVersion: scheduling.run.ai/v2
40+
kind: Queue
41+
metadata:
42+
name: team-a
43+
spec:
44+
parentQueue: department-1 # Inherits from parent queue
45+
# priority: 50 # Optional: Team priority within department
46+
resources:
47+
# quota: Guaranteed resources for this queue
48+
# limit: Maximum resources this queue can use
49+
# overQuotaWeight: How surplus resources are shared among queues
50+
# Note: Using -1 (unlimited) for demo purposes
51+
cpu:
52+
quota: -1
53+
limit: -1
54+
overQuotaWeight: 1
55+
gpu:
56+
quota: -1
57+
limit: -1
58+
overQuotaWeight: 1
59+
memory:
60+
quota: -1
61+
limit: -1
62+
overQuotaWeight: 1
63+
---
64+
# RayCluster with KAI Scheduler and GPU Sharing
65+
apiVersion: ray.io/v1
66+
kind: RayCluster
67+
metadata:
68+
name: raycluster-half-gpu
69+
labels:
70+
kai.scheduler/queue: team-a # REQUIRED: Queue assignment for scheduling
71+
spec:
72+
headGroupSpec:
73+
template:
74+
spec:
75+
containers:
76+
- name: head
77+
image: rayproject/ray:2.46.0
78+
resources:
79+
limits:
80+
cpu: "1"
81+
memory: "2Gi"
82+
# ---- Two workers share one GPU (0.5 each) ----
83+
workerGroupSpecs:
84+
- groupName: shared-gpu
85+
replicas: 2
86+
minReplicas: 2
87+
template:
88+
metadata:
89+
annotations:
90+
gpu-fraction: "0.5" # Request 0.5 GPU per pod (two pods share one GPU)
91+
spec:
92+
containers:
93+
- name: worker
94+
image: rayproject/ray:2.46.0
95+
resources:
96+
limits:
97+
cpu: "1"
98+
memory: "2Gi"

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

Lines changed: 90 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,90 @@
1+
# KAI Scheduler Example - Basic RayCluster
2+
# KAI Scheduler uses a hierarchical queue system for resource management and fair sharing.
3+
# These queues must be created before any RayCluster can be scheduled by KAI.
4+
5+
# NOTE: This is a DEMO configuration with unlimited quotas (-1) for easy testing.
6+
# In real-world deployments, you should set appropriate CPU/GPU/memory quotas and limits
7+
# based on your cluster's actual resources and organizational needs.
8+
9+
# Parent queue: Represents a department or high-level organizational unit
10+
apiVersion: scheduling.run.ai/v2
11+
kind: Queue
12+
metadata:
13+
name: department-1
14+
spec:
15+
# priority: 100 # Optional: Higher priority queues get surplus resources first
16+
resources:
17+
# quota: Guaranteed resources for this queue
18+
# limit: Maximum resources this queue can use
19+
# overQuotaWeight: How surplus resources are shared among queues
20+
# Note: Using -1 (unlimited) for demo purposes
21+
cpu:
22+
quota: -1
23+
limit: -1
24+
overQuotaWeight: 1
25+
gpu:
26+
quota: -1
27+
limit: -1
28+
overQuotaWeight: 1
29+
memory:
30+
quota: -1
31+
limit: -1
32+
overQuotaWeight: 1
33+
---
34+
# Child queue: Represents a team within the department-1
35+
apiVersion: scheduling.run.ai/v2
36+
kind: Queue
37+
metadata:
38+
name: team-a
39+
spec:
40+
parentQueue: department-1 # Inherits from parent queue
41+
# priority: 50 # Optional: Team priority within department
42+
resources:
43+
# quota: Guaranteed resources for this queue
44+
# limit: Maximum resources this queue can use
45+
# overQuotaWeight: How surplus resources are shared among queues
46+
# Note: Using -1 (unlimited) for demo purposes
47+
cpu:
48+
quota: -1
49+
limit: -1
50+
overQuotaWeight: 1
51+
gpu:
52+
quota: -1
53+
limit: -1
54+
overQuotaWeight: 1
55+
memory:
56+
quota: -1
57+
limit: -1
58+
overQuotaWeight: 1
59+
---
60+
# RayCluster with KAI Scheduler
61+
apiVersion: ray.io/v1
62+
kind: RayCluster
63+
metadata:
64+
name: raycluster-sample
65+
labels:
66+
kai.scheduler/queue: team-a # REQUIRED: Queue assignment for scheduling
67+
spec:
68+
headGroupSpec:
69+
template:
70+
spec:
71+
containers:
72+
- name: ray-head
73+
image: rayproject/ray:2.46.0
74+
resources:
75+
requests:
76+
cpu: "1"
77+
memory: "2Gi"
78+
workerGroupSpecs:
79+
- groupName: worker
80+
replicas: 2
81+
minReplicas: 2
82+
template:
83+
spec:
84+
containers:
85+
- name: ray-worker
86+
image: rayproject/ray:2.46.0
87+
resources:
88+
requests:
89+
cpu: "1"
90+
memory: "1Gi"

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

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,70 @@
1+
package kaischeduler
2+
3+
// This KAI plugin relies on KAI-Scheduler's
4+
// built-in PodGrouper to create PodGroups at
5+
// runtime, so the plugin itself only needs to:
6+
// 1. expose the scheduler name,
7+
// 2. stamp pods with schedulerName + queue label.
8+
// No PodGroup create/patch logic is included.
9+
10+
import (
11+
"context"
12+
13+
"github.com/go-logr/logr"
14+
corev1 "k8s.io/api/core/v1"
15+
"k8s.io/apimachinery/pkg/runtime"
16+
"k8s.io/client-go/rest"
17+
"sigs.k8s.io/controller-runtime/pkg/builder"
18+
"sigs.k8s.io/controller-runtime/pkg/client"
19+
logf "sigs.k8s.io/controller-runtime/pkg/log"
20+
21+
rayv1 "github.com/ray-project/kuberay/ray-operator/apis/ray/v1"
22+
schedulerinterface "github.com/ray-project/kuberay/ray-operator/controllers/ray/batchscheduler/interface"
23+
)
24+
25+
const (
26+
QueueLabelName = "kai.scheduler/queue"
27+
)
28+
29+
type KaiScheduler struct {
30+
log logr.Logger
31+
}
32+
33+
type KaiSchedulerFactory struct{}
34+
35+
func GetPluginName() string { return "kai-scheduler" }
36+
37+
func (k *KaiScheduler) Name() string { return GetPluginName() }
38+
39+
func (k *KaiScheduler) DoBatchSchedulingOnSubmission(_ context.Context, _ *rayv1.RayCluster) error {
40+
return nil
41+
}
42+
43+
func (k *KaiScheduler) AddMetadataToPod(_ context.Context, app *rayv1.RayCluster, _ string, pod *corev1.Pod) {
44+
pod.Spec.SchedulerName = k.Name()
45+
46+
queue, ok := app.Labels[QueueLabelName]
47+
if !ok || queue == "" {
48+
k.log.Info("Queue label missing from RayCluster; pods will remain pending",
49+
"requiredLabel", QueueLabelName,
50+
"rayCluster", app.Name)
51+
return
52+
}
53+
if pod.Labels == nil {
54+
pod.Labels = make(map[string]string)
55+
}
56+
pod.Labels[QueueLabelName] = queue
57+
}
58+
59+
func (kf *KaiSchedulerFactory) New(_ context.Context, _ *rest.Config, _ client.Client) (schedulerinterface.BatchScheduler, error) {
60+
return &KaiScheduler{
61+
log: logf.Log.WithName("kai-scheduler"),
62+
}, nil
63+
}
64+
65+
func (kf *KaiSchedulerFactory) AddToScheme(_ *runtime.Scheme) {
66+
}
67+
68+
func (kf *KaiSchedulerFactory) ConfigureReconciler(b *builder.Builder) *builder.Builder {
69+
return b
70+
}

0 commit comments

Comments
 (0)