From 46358db912d15003ad9ddd1daf231a31ffa847ca Mon Sep 17 00:00:00 2001
From: Kevin <kpostlet@redhat.com>
Date: Mon, 22 Jul 2024 15:43:00 -0400
Subject: [PATCH] automatically inject default queue if not provided

Signed-off-by: Kevin <kpostlet@redhat.com>
---
 pkg/controllers/raycluster_webhook.go      | 28 ++++++++++++++++++++++
 pkg/controllers/raycluster_webhook_test.go |  2 ++
 2 files changed, 30 insertions(+)

diff --git a/pkg/controllers/raycluster_webhook.go b/pkg/controllers/raycluster_webhook.go
index 4e4b259f6..f2fa552ab 100644
--- a/pkg/controllers/raycluster_webhook.go
+++ b/pkg/controllers/raycluster_webhook.go
@@ -21,12 +21,14 @@ import (
 	"strconv"
 
 	rayv1 "github.com/ray-project/kuberay/ray-operator/apis/ray/v1"
+	kueuev1beta1 "sigs.k8s.io/kueue/apis/kueue/v1beta1"
 
 	corev1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/runtime"
 	"k8s.io/apimachinery/pkg/util/validation/field"
 	"k8s.io/utils/ptr"
 	ctrl "sigs.k8s.io/controller-runtime"
+	"sigs.k8s.io/controller-runtime/pkg/client"
 	logf "sigs.k8s.io/controller-runtime/pkg/log"
 	"sigs.k8s.io/controller-runtime/pkg/webhook"
 	"sigs.k8s.io/controller-runtime/pkg/webhook/admission"
@@ -45,6 +47,7 @@ var rayclusterlog = logf.Log.WithName("raycluster-resource")
 
 func SetupRayClusterWebhookWithManager(mgr ctrl.Manager, cfg *config.KubeRayConfiguration) error {
 	rayClusterWebhookInstance := &rayClusterWebhook{
+		Client: mgr.GetClient(),
 		Config: cfg,
 	}
 	return ctrl.NewWebhookManagedBy(mgr).
@@ -58,6 +61,7 @@ func SetupRayClusterWebhookWithManager(mgr ctrl.Manager, cfg *config.KubeRayConf
 // +kubebuilder:webhook:path=/validate-ray-io-v1-raycluster,mutating=false,failurePolicy=fail,sideEffects=None,groups=ray.io,resources=rayclusters,verbs=create;update,versions=v1,name=vraycluster.ray.openshift.ai,admissionReviewVersions=v1
 
 type rayClusterWebhook struct {
+	client.Client
 	Config *config.KubeRayConfiguration
 }
 
@@ -77,6 +81,30 @@ func (w *rayClusterWebhook) Default(ctx context.Context, obj runtime.Object) err
 		rayCluster.Spec.HeadGroupSpec.Template.Spec.ServiceAccountName = rayCluster.Name + "-oauth-proxy"
 	}
 
+	// add default queue label if not present
+	if rayCluster.GetLabels() == nil {
+		rayCluster.Labels = make(map[string]string)
+	}
+	err := w.Client.List(ctx, &kueuev1beta1.LocalQueueList{})
+	if err != nil {
+		rayclusterlog.Error(err, "Failed to list LocalQueues, Kueue CRD might not be installed")
+		_, ok := rayCluster.Labels["kueue.x-k8s.io/queue-name"]
+		if !ok {
+			// check if CRD Kueue LocalQueue exists
+			localQueues := &kueuev1beta1.LocalQueueList{}
+			err := w.Client.List(ctx, localQueues)
+			if err == nil {
+				for _, localQueue := range localQueues.Items {
+					is_default, ok := localQueue.Labels["kueue.x-k8s.io/default-queue"]
+					if ok && is_default == "true" {
+						rayCluster.Labels["kueue.x-k8s.io/queue-name"] = localQueue.Name
+						break
+					}
+				}
+			}
+		}
+	}
+
 	if ptr.Deref(w.Config.MTLSEnabled, true) {
 		rayclusterlog.V(2).Info("Adding create-cert Init Containers")
 		// HeadGroupSpec
diff --git a/pkg/controllers/raycluster_webhook_test.go b/pkg/controllers/raycluster_webhook_test.go
index d8e4f8c43..2997725d5 100644
--- a/pkg/controllers/raycluster_webhook_test.go
+++ b/pkg/controllers/raycluster_webhook_test.go
@@ -22,6 +22,7 @@ import (
 	. "github.com/onsi/gomega"
 	"github.com/project-codeflare/codeflare-common/support"
 	rayv1 "github.com/ray-project/kuberay/ray-operator/apis/ray/v1"
+	"sigs.k8s.io/controller-runtime/pkg/client/fake"
 
 	corev1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
@@ -35,6 +36,7 @@ var (
 	rayClusterName = "test-raycluster"
 
 	rcWebhook = &rayClusterWebhook{
+		Client: fake.NewFakeClient(),
 		Config: &config.KubeRayConfiguration{},
 	}
 )