[Feature] Support Volcano for batch scheduling

ray-operator/controllers/ray/raycluster_controller.go

@@ -33,12 +33,14 @@ import (
 	"sigs.k8s.io/controller-runtime/pkg/manager"
 	"sigs.k8s.io/controller-runtime/pkg/reconcile"
 	"sigs.k8s.io/controller-runtime/pkg/source"
+	vcbetav1 "volcano.sh/apis/pkg/apis/scheduling/v1beta1"
 )
 
 var (
 	DefaultRequeueDuration    = 2 * time.Second
 	PrioritizeWorkersToDelete bool
 	ForcedClusterUpgrade      bool
+	EnableBatchScheduler      bool
 )
 
 // NewReconciler returns a new reconcile.Reconciler
@@ -336,6 +338,20 @@ func (r *RayClusterReconciler) reconcilePods(instance *rayiov1alpha1.RayCluster)
 	if err := r.List(context.TODO(), &headPods, client.InNamespace(instance.Namespace), filterLabels); err != nil {
 		return err
 	}
+	if EnableBatchScheduler {
+		var minMember int32
+		var totalResource corev1.ResourceList
+		if instance.Spec.EnableInTreeAutoscaling == nil || !*instance.Spec.EnableInTreeAutoscaling {
+			minMember = utils.CalculateDesiredReplicas(instance) + *instance.Spec.HeadGroupSpec.Replicas
+			totalResource = utils.CalculateDesiredResources(instance)
+		} else {
+			minMember = utils.CalculateMinReplicas(instance) + *instance.Spec.HeadGroupSpec.Replicas
+			totalResource = utils.CalculateMinResources(instance)
+		}
+		if err := r.syncPodGroup(instance, minMember, totalResource); err != nil {
+			return err
+		}
+	}
 	// Reconcile head Pod
 	if len(headPods.Items) == 1 {
 		headPod := headPods.Items[0]
@@ -654,6 +670,9 @@ func (r *RayClusterReconciler) createHeadPod(instance rayiov1alpha1.RayCluster)
 		Name:      pod.Name,
 		Namespace: pod.Namespace,
 	}
+	if EnableBatchScheduler {
+		pod.Annotations[vcbetav1.KubeGroupNameAnnotationKey] = instance.Name
+	}
 
 	r.Log.Info("createHeadPod", "head pod with name", pod.GenerateName)
 	if err := r.Create(context.TODO(), &pod); err != nil {
@@ -682,6 +701,9 @@ func (r *RayClusterReconciler) createWorkerPod(instance rayiov1alpha1.RayCluster
 		Name:      pod.Name,
 		Namespace: pod.Namespace,
 	}
+	if EnableBatchScheduler {
+		pod.Annotations[vcbetav1.KubeGroupNameAnnotationKey] = instance.Name
+	}
 	replica := pod
 	if err := r.Create(context.TODO(), &replica); err != nil {
 		if errors.IsAlreadyExists(err) {
@@ -771,6 +793,10 @@ func (r *RayClusterReconciler) SetupWithManager(mgr ctrl.Manager, reconcileConcu
 			IsController: true,
 			OwnerType:    &rayiov1alpha1.RayCluster{},
 		}).
+		Watches(&source.Kind{Type: &vcbetav1.PodGroup{}}, &handler.EnqueueRequestForOwner{
+			IsController: true,
+			OwnerType:    &rayiov1alpha1.RayCluster{},
+		}).
 		WithOptions(controller.Options{MaxConcurrentReconciles: reconcileConcurrency}).
 		Complete(r)
 }
@@ -1049,3 +1075,53 @@ func (r *RayClusterReconciler) updateClusterState(instance *rayiov1alpha1.RayClu
 	instance.Status.State = clusterState
 	return r.Status().Update(context.Background(), instance)
 }
+
+func (r *RayClusterReconciler) syncPodGroup(instance *rayiov1alpha1.RayCluster, size int32, totalResource corev1.ResourceList) error {
+	// use cluster name as pod group name
+	podGroupName := instance.Name
+	podGroupIdentifier := types.NamespacedName{
+		Name:      podGroupName,
+		Namespace: instance.Namespace,
+	}
+	podGroup := &vcbetav1.PodGroup{}
+	if err := r.Get(context.TODO(), podGroupIdentifier, podGroup); err != nil {
+		if !errors.IsNotFound(err) {
+			return err
+		}
+
+		// use the same name as ray cluster
+		podGroup.Name = instance.Name
+		podGroup.Namespace = instance.Namespace
+		podGroup.OwnerReferences = []metav1.OwnerReference{
+			*metav1.NewControllerRef(instance, rayiov1alpha1.SchemeGroupVersion.WithKind("RayCluster")),
+		}
+		podGroup.Spec.MinMember = size
+		podGroup.Spec.MinResources = &totalResource
+		if queue, ok := instance.Spec.HeadGroupSpec.Template.Labels["volcano.sh/queue-name"]; ok {
+			podGroup.Spec.Queue = queue
+		}
+		podGroup.Spec.PriorityClassName = instance.Spec.HeadGroupSpec.Template.Spec.PriorityClassName
+		podGroup.Status = vcbetav1.PodGroupStatus{
+			Phase: vcbetav1.PodGroupPending,
+		}
+
+		if err := r.Create(context.TODO(), podGroup); err != nil {
+			if errors.IsAlreadyExists(err) {
+				r.Log.Info("pod group already exist, no need to create")
+				return nil
+			}
+			r.Log.Error(err, "Pod group create error!", "PodGroup.Error", err)
+			return err
+		}
+	} else {
+		if podGroup.Spec.MinMember != size || !utils.Equals(*podGroup.Spec.MinResources, totalResource) {
+			podGroup.Spec.MinMember = size
+			podGroup.Spec.MinResources = &totalResource
+			if err := r.Update(context.TODO(), podGroup); err != nil {
+				r.Log.Error(err, "Fail to update pod group!", "podGroup", podGroupIdentifier)
+				return err
+			}
+		}
+	}
+	return nil
+}

ray-operator/controllers/ray/utils/util.go

@@ -213,6 +213,73 @@ func CalculateAvailableReplicas(pods corev1.PodList) int32 {
 	return count
 }
 
+func CalculateDesiredResources(cluster *rayiov1alpha1.RayCluster) corev1.ResourceList {
+	desiredResourcesList := []corev1.ResourceList{{}}
+	headPodResource := calculatePodResource(cluster.Spec.HeadGroupSpec.Template.Spec)
+	for i := int32(0); i < *cluster.Spec.HeadGroupSpec.Replicas; i++ {
+		desiredResourcesList = append(desiredResourcesList, headPodResource)
+	}
+	for _, nodeGroup := range cluster.Spec.WorkerGroupSpecs {
+		podResource := calculatePodResource(nodeGroup.Template.Spec)
+		for i := int32(0); i < *nodeGroup.Replicas; i++ {
+			desiredResourcesList = append(desiredResourcesList, podResource)
+		}
+	}
+	return sumResourceList(desiredResourcesList)
+}
+
+func CalculateMinResources(cluster *rayiov1alpha1.RayCluster) corev1.ResourceList {
+	minResourcesList := []corev1.ResourceList{{}}
+	headPodResource := calculatePodResource(cluster.Spec.HeadGroupSpec.Template.Spec)
+	for i := int32(0); i < *cluster.Spec.HeadGroupSpec.Replicas; i++ {
+		minResourcesList = append(minResourcesList, headPodResource)
+	}
+	for _, nodeGroup := range cluster.Spec.WorkerGroupSpecs {
+		podResource := calculatePodResource(nodeGroup.Template.Spec)
+		for i := int32(0); i < *nodeGroup.MinReplicas; i++ {
+			minResourcesList = append(minResourcesList, podResource)
+		}
+	}
+	return sumResourceList(minResourcesList)
+}
+
+func calculatePodResource(podSpec corev1.PodSpec) corev1.ResourceList {
+	podResource := corev1.ResourceList{}
+	for _, container := range podSpec.Containers {
+		containerResource := container.Resources.Requests
+		for name, quantity := range container.Resources.Limits {
+			if _, ok := containerResource[name]; !ok {
+				containerResource[name] = quantity
+			}
+		}
+		for name, quantity := range containerResource {
+			if totalQuantity, ok := podResource[name]; ok {
+				totalQuantity.Add(quantity)
+				podResource[name] = totalQuantity
+			} else {
+				podResource[name] = quantity
+			}
+		}
+	}
+	return podResource
+}
+
+func sumResourceList(list []corev1.ResourceList) corev1.ResourceList {
+	totalResource := corev1.ResourceList{}
+	for _, l := range list {
+		for name, quantity := range l {
+
+			if value, ok := totalResource[name]; !ok {
+				totalResource[name] = quantity.DeepCopy()
+			} else {
+				value.Add(quantity)
+				totalResource[name] = value
+			}
+		}
+	}
+	return totalResource
+}
+
 func Contains(elems []string, searchTerm string) bool {
 	for _, s := range elems {
 		if searchTerm == s {
@@ -359,3 +426,22 @@ func GenerateJsonHash(obj interface{}) (string, error) {
 
 	return hashStr, nil
 }
+
+// Equals returns true if the two lists are equivalent
+func Equals(a corev1.ResourceList, b corev1.ResourceList) bool {
+	if len(a) != len(b) {
+		return false
+	}
+
+	for key, value1 := range a {
+		value2, found := b[key]
+		if !found {
+			return false
+		}
+		if value1.Cmp(value2) != 0 {
+			return false
+		}
+	}
+
+	return true
+}

ray-operator/go.mod

@@ -21,6 +21,7 @@ require (
 	k8s.io/code-generator v0.23.0
 	k8s.io/utils v0.0.0-20210930125809-cb0fa318a74b
 	sigs.k8s.io/controller-runtime v0.11.1
+	volcano.sh/apis v1.6.0-alpha.0.0.20221012070524-685db38b4fae
 )
 
 require (

ray-operator/go.sum

@@ -948,6 +948,8 @@ k8s.io/component-base v0.23.0 h1:UAnyzjvVZ2ZR1lF35YwtNY6VMN94WtOnArcXBu34es8=
 k8s.io/component-base v0.23.0/go.mod h1:DHH5uiFvLC1edCpvcTDV++NKULdYYU6pR9Tt3HIKMKI=
 k8s.io/gengo v0.0.0-20210813121822-485abfe95c7c h1:GohjlNKauSai7gN4wsJkeZ3WAJx4Sh+oT/b5IYn5suA=
 k8s.io/gengo v0.0.0-20210813121822-485abfe95c7c/go.mod h1:FiNAH4ZV3gBg2Kwh89tzAEV2be7d5xI0vBa/VySYy3E=
+k8s.io/klog v1.0.0 h1:Pt+yjF5aB1xDSVbau4VsWe+dQNzA0qv1LlXdC2dF6Q8=
+k8s.io/klog v1.0.0/go.mod h1:4Bi6QPql/J/LkTDqv7R/cd3hPo4k2DG6Ptcz060Ez5I=
 k8s.io/klog/v2 v2.0.0/go.mod h1:PBfzABfn139FHAV07az/IF9Wp1bkk3vpT2XSJ76fSDE=
 k8s.io/klog/v2 v2.2.0/go.mod h1:Od+F08eJP+W3HUb4pSrPpgp9DGU4GzlpG/TmITuYh/Y=
 k8s.io/klog/v2 v2.30.0 h1:bUO6drIvCIsvZ/XFgfxoGFQU/a4Qkh0iAlvUR7vlHJw=
@@ -972,3 +974,5 @@ sigs.k8s.io/structured-merge-diff/v4 v4.2.0/go.mod h1:j/nl6xW8vLS49O8YvXW1ocPhZa
 sigs.k8s.io/yaml v1.2.0/go.mod h1:yfXDCHCao9+ENCvLSE62v9VSji2MKu5jeNfTrofGhJc=
 sigs.k8s.io/yaml v1.3.0 h1:a2VclLzOGrwOHDiV8EfBGhvjHvP46CtW5j6POvhYGGo=
 sigs.k8s.io/yaml v1.3.0/go.mod h1:GeOyir5tyXNByN85N/dRIT9es5UQNerPYEKK56eTBm8=
+volcano.sh/apis v1.6.0-alpha.0.0.20221012070524-685db38b4fae h1:H7yidKnIq/Y7KmjFP5xFSmE7xL674226D8pEoA/RfG8=
+volcano.sh/apis v1.6.0-alpha.0.0.20221012070524-685db38b4fae/go.mod h1:drNMGuHPn1ew7oBSDQb5KRey6tXOQksbUtw3gPxF3Vo=

ray-operator/main.go

@@ -21,6 +21,7 @@ import (
 	k8szap "sigs.k8s.io/controller-runtime/pkg/log/zap"
 
 	rayv1alpha1 "github.com/ray-project/kuberay/ray-operator/apis/ray/v1alpha1"
+	vcbetav1 "volcano.sh/apis/pkg/apis/scheduling/v1beta1"
 	// +kubebuilder:scaffold:imports
 )
 
@@ -35,6 +36,7 @@ var (
 func init() {
 	utilruntime.Must(clientgoscheme.AddToScheme(scheme))
 	utilruntime.Must(rayv1alpha1.AddToScheme(scheme))
+	utilruntime.Must(vcbetav1.AddToScheme(scheme))
 	// +kubebuilder:scaffold:scheme
 }
 
@@ -63,6 +65,8 @@ func main() {
 		"Forced cluster upgrade flag")
 	flag.StringVar(&logFile, "log-file-path", "",
 		"Synchronize logs to local file")
+	flag.BoolVar(&ray.EnableBatchScheduler, "enable-batch-scheduler", false,
+		"Enable batch scheduler. Currently is volcano, which supports gang scheduler policy.")
 
 	opts := k8szap.Options{
 		Development: true,