Skip to content

Commit ae3c527

Browse files
win5923win5923
committed
Extract AddSchedulerNameToObject to batchscheduler utils package
Signed-off-by: win5923 <[email protected]>
1 parent 40bd7f0 commit ae3c527

File tree

5 files changed

+76
-34
lines changed

5 files changed

+76
-34
lines changed

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

Lines changed: 9 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -9,16 +9,18 @@ package kaischeduler
99

1010
import (
1111
"context"
12+
"fmt"
1213

13-
corev1 "k8s.io/api/core/v1"
1414
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
1515
"k8s.io/apimachinery/pkg/runtime"
1616
"k8s.io/client-go/rest"
1717
ctrl "sigs.k8s.io/controller-runtime"
1818
"sigs.k8s.io/controller-runtime/pkg/builder"
1919
"sigs.k8s.io/controller-runtime/pkg/client"
2020

21+
rayv1 "github.com/ray-project/kuberay/ray-operator/apis/ray/v1"
2122
schedulerinterface "github.com/ray-project/kuberay/ray-operator/controllers/ray/batchscheduler/interface"
23+
"github.com/ray-project/kuberay/ray-operator/controllers/ray/batchscheduler/utils"
2224
)
2325

2426
const (
@@ -33,13 +35,17 @@ func GetPluginName() string { return "kai-scheduler" }
3335

3436
func (k *KaiScheduler) Name() string { return GetPluginName() }
3537

36-
func (k *KaiScheduler) DoBatchSchedulingOnSubmission(_ context.Context, _ metav1.Object) error {
38+
func (k *KaiScheduler) DoBatchSchedulingOnSubmission(_ context.Context, object metav1.Object) error {
39+
_, ok := object.(*rayv1.RayCluster)
40+
if !ok {
41+
return fmt.Errorf("currently only RayCluster is supported, got %T", object)
42+
}
3743
return nil
3844
}
3945

4046
func (k *KaiScheduler) AddMetadataToChildResource(ctx context.Context, parent metav1.Object, child metav1.Object, _ string) {
4147
logger := ctrl.LoggerFrom(ctx).WithName("kai-scheduler")
42-
addSchedulerNameToObject(child, k.Name())
48+
utils.AddSchedulerNameToObject(child, k.Name())
4349

4450
parentLabel := parent.GetLabels()
4551
queue, ok := parentLabel[QueueLabelName]
@@ -57,15 +63,6 @@ func (k *KaiScheduler) AddMetadataToChildResource(ctx context.Context, parent me
5763
child.SetLabels(childLabels)
5864
}
5965

60-
func addSchedulerNameToObject(obj metav1.Object, schedulerName string) {
61-
switch obj := obj.(type) {
62-
case *corev1.Pod:
63-
obj.Spec.SchedulerName = schedulerName
64-
case *corev1.PodTemplateSpec:
65-
obj.Spec.SchedulerName = schedulerName
66-
}
67-
}
68-
6966
func (kf *KaiSchedulerFactory) New(_ context.Context, _ *rest.Config, _ client.Client) (schedulerinterface.BatchScheduler, error) {
7067
return &KaiScheduler{}, nil
7168
}

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

Lines changed: 2 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,6 @@ import (
44
"context"
55
"fmt"
66

7-
corev1 "k8s.io/api/core/v1"
87
"k8s.io/apimachinery/pkg/api/errors"
98
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
109
"k8s.io/apimachinery/pkg/runtime"
@@ -17,6 +16,7 @@ import (
1716

1817
rayv1 "github.com/ray-project/kuberay/ray-operator/apis/ray/v1"
1918
schedulerinterface "github.com/ray-project/kuberay/ray-operator/controllers/ray/batchscheduler/interface"
19+
batchschedulerutils "github.com/ray-project/kuberay/ray-operator/controllers/ray/batchscheduler/utils"
2020
"github.com/ray-project/kuberay/ray-operator/controllers/ray/utils"
2121
)
2222

@@ -105,16 +105,7 @@ func (k *KubeScheduler) AddMetadataToChildResource(_ context.Context, parent met
105105
labels[kubeSchedulerPodGroupLabelKey] = parent.GetName()
106106
child.SetLabels(labels)
107107
}
108-
addSchedulerNameToObject(child, k.Name())
109-
}
110-
111-
func addSchedulerNameToObject(obj metav1.Object, schedulerName string) {
112-
switch obj := obj.(type) {
113-
case *corev1.Pod:
114-
obj.Spec.SchedulerName = schedulerName
115-
case *corev1.PodTemplateSpec:
116-
obj.Spec.SchedulerName = schedulerName
117-
}
108+
batchschedulerutils.AddSchedulerNameToObject(child, k.Name())
118109
}
119110

120111
func (k *KubeScheduler) isGangSchedulingEnabled(obj metav1.Object) bool {

‎ray-operator/controllers/ray/batchscheduler/utils/utils.go‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
package utils
2+
3+
import (
4+
corev1 "k8s.io/api/core/v1"
5+
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
6+
)
7+
8+
// AddSchedulerNameToObject sets the schedulerName field on Pod and PodTemplateSpec resources.
9+
// Used to assign batch scheduler names to:
10+
// - Head pod and worker pod in RayCluster
11+
// - Job in RayJob
12+
func AddSchedulerNameToObject(obj metav1.Object, schedulerName string) {
13+
switch obj := obj.(type) {
14+
case *corev1.Pod:
15+
obj.Spec.SchedulerName = schedulerName
16+
case *corev1.PodTemplateSpec:
17+
obj.Spec.SchedulerName = schedulerName
18+
}
19+
}

‎ray-operator/controllers/ray/batchscheduler/utils/utils_test.go‎

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
1+
package utils
2+
3+
import (
4+
"testing"
5+
6+
corev1 "k8s.io/api/core/v1"
7+
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
8+
)
9+
10+
func TestAddSchedulerNameToObject(t *testing.T) {
11+
schedulerName := "test-scheduler"
12+
13+
t.Run("Pod object should have schedulerName set", func(t *testing.T) {
14+
pod := &corev1.Pod{
15+
ObjectMeta: metav1.ObjectMeta{
16+
Name: "test-pod",
17+
Namespace: "default",
18+
},
19+
Spec: corev1.PodSpec{},
20+
}
21+
22+
AddSchedulerNameToObject(pod, schedulerName)
23+
24+
if pod.Spec.SchedulerName != schedulerName {
25+
t.Errorf("expected schedulerName to be %q, got %q", schedulerName, pod.Spec.SchedulerName)
26+
}
27+
})
28+
29+
t.Run("PodTemplateSpec object should have schedulerName set", func(t *testing.T) {
30+
podTemplate := &corev1.PodTemplateSpec{
31+
ObjectMeta: metav1.ObjectMeta{
32+
Name: "test-template",
33+
Namespace: "default",
34+
},
35+
Spec: corev1.PodSpec{},
36+
}
37+
38+
AddSchedulerNameToObject(podTemplate, schedulerName)
39+
40+
if podTemplate.Spec.SchedulerName != schedulerName {
41+
t.Errorf("expected schedulerName to be %q, got %q", schedulerName, podTemplate.Spec.SchedulerName)
42+
}
43+
})
44+
}

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

Lines changed: 2 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,6 @@ package yunikorn
33
import (
44
"context"
55

6-
corev1 "k8s.io/api/core/v1"
76
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
87
"k8s.io/apimachinery/pkg/runtime"
98
"k8s.io/client-go/rest"
@@ -13,6 +12,7 @@ import (
1312

1413
rayv1 "github.com/ray-project/kuberay/ray-operator/apis/ray/v1"
1514
schedulerinterface "github.com/ray-project/kuberay/ray-operator/controllers/ray/batchscheduler/interface"
15+
batchschedulerutils "github.com/ray-project/kuberay/ray-operator/controllers/ray/batchscheduler/utils"
1616
"github.com/ray-project/kuberay/ray-operator/controllers/ray/utils"
1717
)
1818

@@ -84,15 +84,6 @@ func populateLabelsFromObject(parent metav1.Object, child metav1.Object, sourceK
8484
child.SetLabels(labels)
8585
}
8686

87-
func addSchedulerNameToObject(obj metav1.Object, schedulerName string) {
88-
switch obj := obj.(type) {
89-
case *corev1.Pod:
90-
obj.Spec.SchedulerName = schedulerName
91-
case *corev1.PodTemplateSpec:
92-
obj.Spec.SchedulerName = schedulerName
93-
}
94-
}
95-
9687
func getTaskGroupsAnnotationValue(obj metav1.Object) (string, error) {
9788
taskGroups := newTaskGroups()
9889
switch obj := obj.(type) {
@@ -127,7 +118,7 @@ func (y *YuniKornScheduler) AddMetadataToChildResource(ctx context.Context, pare
127118

128119
populateLabelsFromObject(parent, child, RayApplicationIDLabelName, YuniKornPodApplicationIDLabelName)
129120
populateLabelsFromObject(parent, child, RayApplicationQueueLabelName, YuniKornPodQueueLabelName)
130-
addSchedulerNameToObject(child, y.Name())
121+
batchschedulerutils.AddSchedulerNameToObject(child, y.Name())
131122

132123
if y.isGangSchedulingEnabled(parent) {
133124
logger.Info("gang scheduling is enabled, propagating task groups annotation to child", "name", child.GetName(), "namespace", child.GetNamespace())

0 commit comments

Comments
 (0)