Merge branch 'master' into KUBERAY-2237

Author: cchen777

helm-chart/kuberay-operator/crds/ray.io_rayclusters.yaml

@@ -182,6 +182,7 @@ const (
 type HeadInfo struct {
 	PodIP       string `json:"podIP,omitempty"`
 	ServiceIP   string `json:"serviceIP,omitempty"`
+	PodName     string `json:"podName,omitempty"`
 	ServiceName string `json:"serviceName,omitempty"`
 }
 

helm-chart/kuberay-operator/crds/ray.io_rayjobs.yaml

@@ -11,12 +11,14 @@ import (
 	rayv1 "github.com/ray-project/kuberay/ray-operator/apis/ray/v1"
 	schedulerinterface "github.com/ray-project/kuberay/ray-operator/controllers/ray/batchscheduler/interface"
 	"github.com/ray-project/kuberay/ray-operator/controllers/ray/batchscheduler/volcano"
+	"github.com/ray-project/kuberay/ray-operator/controllers/ray/batchscheduler/yunikorn"
 	"github.com/ray-project/kuberay/ray-operator/controllers/ray/utils"
 )
 
 var schedulerContainers = map[string]schedulerinterface.BatchSchedulerFactory{
 	schedulerinterface.GetDefaultPluginName(): &schedulerinterface.DefaultBatchSchedulerFactory{},
 	volcano.GetPluginName():                   &volcano.VolcanoBatchSchedulerFactory{},
+	yunikorn.GetPluginName():                  &yunikorn.YuniKornSchedulerFactory{},
 }
 
 func GetRegisteredNames() []string {

helm-chart/kuberay-operator/crds/ray.io_rayservices.yaml

@@ -0,0 +1,72 @@
+package yunikorn
+
+import (
+	"context"
+
+	"github.com/go-logr/logr"
+	corev1 "k8s.io/api/core/v1"
+	"k8s.io/apimachinery/pkg/runtime"
+	"k8s.io/client-go/rest"
+	"sigs.k8s.io/controller-runtime/pkg/builder"
+	logf "sigs.k8s.io/controller-runtime/pkg/log"
+
+	rayv1 "github.com/ray-project/kuberay/ray-operator/apis/ray/v1"
+	schedulerinterface "github.com/ray-project/kuberay/ray-operator/controllers/ray/batchscheduler/interface"
+)
+
+const (
+	SchedulerName                     string = "yunikorn"
+	YuniKornPodApplicationIDLabelName string = "applicationId"
+	YuniKornPodQueueLabelName         string = "queue"
+	RayClusterApplicationIDLabelName  string = "yunikorn.apache.org/application-id"
+	RayClusterQueueLabelName          string = "yunikorn.apache.org/queue-name"
+)
+
+type YuniKornScheduler struct {
+	log logr.Logger
+}
+
+type YuniKornSchedulerFactory struct{}
+
+func GetPluginName() string {
+	return SchedulerName
+}
+
+func (y *YuniKornScheduler) Name() string {
+	return GetPluginName()
+}
+
+func (y *YuniKornScheduler) DoBatchSchedulingOnSubmission(_ context.Context, _ *rayv1.RayCluster) error {
+	// yunikorn doesn't require any resources to be created upfront
+	// this is a no-opt for this implementation
+	return nil
+}
+
+func (y *YuniKornScheduler) populatePodLabels(app *rayv1.RayCluster, pod *corev1.Pod, sourceKey string, targetKey string) {
+	// check labels
+	if value, exist := app.Labels[sourceKey]; exist {
+		y.log.Info("Updating pod label based on RayCluster annotations",
+			"sourceKey", sourceKey, "targetKey", targetKey, "value", value)
+		pod.Labels[targetKey] = value
+	}
+}
+
+func (y *YuniKornScheduler) AddMetadataToPod(app *rayv1.RayCluster, _ string, pod *corev1.Pod) {
+	y.populatePodLabels(app, pod, RayClusterApplicationIDLabelName, YuniKornPodApplicationIDLabelName)
+	y.populatePodLabels(app, pod, RayClusterQueueLabelName, YuniKornPodQueueLabelName)
+	pod.Spec.SchedulerName = y.Name()
+}
+
+func (yf *YuniKornSchedulerFactory) New(_ *rest.Config) (schedulerinterface.BatchScheduler, error) {
+	return &YuniKornScheduler{
+		log: logf.Log.WithName(SchedulerName),
+	}, nil
+}
+
+func (yf *YuniKornSchedulerFactory) AddToScheme(_ *runtime.Scheme) {
+	// No extra scheme needs to be registered
+}
+
+func (yf *YuniKornSchedulerFactory) ConfigureReconciler(b *builder.Builder) *builder.Builder {
+	return b
+}

ray-operator/apis/ray/v1/raycluster_types.go

@@ -0,0 +1,92 @@
+package yunikorn
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	v1 "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 TestPopulatePodLabels(t *testing.T) {
+	yk := &YuniKornScheduler{}
+
+	// --- case 1
+	// Ray Cluster CR has labels defined
+	job1 := "job-1-01234"
+	queue1 := "root.default"
+
+	rayCluster1 := createRayClusterWithLabels(
+		"ray-cluster-with-labels",
+		"test",
+		map[string]string{
+			RayClusterApplicationIDLabelName: job1,
+			RayClusterQueueLabelName:         queue1,
+		},
+	)
+
+	rayPod := createPod("my-pod-1", "test")
+	yk.populatePodLabels(rayCluster1, rayPod, RayClusterApplicationIDLabelName, YuniKornPodApplicationIDLabelName)
+	yk.populatePodLabels(rayCluster1, rayPod, RayClusterQueueLabelName, YuniKornPodQueueLabelName)
+	assert.Equal(t, podLabelsContains(rayPod, YuniKornPodApplicationIDLabelName, job1), true)
+	assert.Equal(t, podLabelsContains(rayPod, YuniKornPodQueueLabelName, queue1), true)
+
+	// --- case 2
+	// Ray Cluster CR has nothing
+	// In this case, the pod will not be populated with the required labels
+	job2 := "job-2-01234"
+	queue2 := "root.default"
+
+	rayCluster2 := createRayClusterWithLabels(
+		"ray-cluster-without-labels",
+		"test",
+		nil, // empty labels
+	)
+	rayPod3 := createPod("my-pod-2", "test")
+	yk.populatePodLabels(rayCluster2, rayPod3, RayClusterApplicationIDLabelName, YuniKornPodApplicationIDLabelName)
+	yk.populatePodLabels(rayCluster2, rayPod3, RayClusterQueueLabelName, YuniKornPodQueueLabelName)
+	assert.Equal(t, podLabelsContains(rayPod3, YuniKornPodApplicationIDLabelName, job2), false)
+	assert.Equal(t, podLabelsContains(rayPod3, YuniKornPodQueueLabelName, queue2), false)
+}
+
+func createRayClusterWithLabels(name string, namespace string, labels map[string]string) *rayv1.RayCluster {
+	rayCluster := &rayv1.RayCluster{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      name,
+			Namespace: namespace,
+			Labels:    labels,
+		},
+	}
+
+	return rayCluster
+}
+
+func createPod(name string, namespace string) *v1.Pod {
+	return &v1.Pod{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:        name,
+			Namespace:   namespace,
+			Labels:      make(map[string]string),
+			Annotations: make(map[string]string),
+		},
+	}
+}
+
+func podLabelsContains(pod *v1.Pod, key string, value string) bool {
+	if pod == nil {
+		return false
+	}
+
+	if len(pod.Labels) > 0 {
+		labelValue, exist := pod.Labels[key]
+		if exist {
+			if labelValue == value {
+				return true
+			}
+		}
+	}
+
+	return false
+}