[Feature] Add ManagedBy field to RayJob (#2589)

mszadkow · web-flow · commit e6af2cc607af · 2024-12-09T13:27:31.000-05:00
* Add managedBy field to the spec

* Make use of managedBy field in RayJob reconciliation

* Update after code review
diff --git a/docs/reference/api.md b/docs/reference/api.md
@@ -159,6 +159,7 @@ _Appears in:_
 | `metadata` _object (keys:string, values:string)_ | Refer to Kubernetes API documentation for fields of `metadata`. |  |  |
 | `clusterSelector` _object (keys:string, values:string)_ | clusterSelector is used to select running rayclusters by labels |  |  |
 | `submitterConfig` _[SubmitterConfig](#submitterconfig)_ | Configurations of submitter k8s job. |  |  |
+| `managedBy` _string_ | ManagedBy is used to indicate the controller or entity that manages a RayJob.<br />The value must be either empty, 'ray.io/kuberay-operator' or 'kueue.x-k8s.io/multikueue'.<br />The kuberay-operator reconciles a RayJob which doesn't have this field at all or<br />the field value is the reserved string 'ray.io/kuberay-operator',<br />but delegates reconciling the RayJob with 'kueue.x-k8s.io/multikueue' to the Kueue.<br />The field is immutable. |  |  |
 | `entrypoint` _string_ | INSERT ADDITIONAL SPEC FIELDS - desired state of cluster<br />Important: Run "make" to regenerate code after modifying this file |  |  |
 | `runtimeEnvYAML` _string_ | RuntimeEnvYAML represents the runtime environment configuration<br />provided as a multi-line YAML string. |  |  |
 | `jobId` _string_ | If jobId is not set, a new jobId will be auto-generated. |  |  |
diff --git a/helm-chart/kuberay-operator/crds/ray.io_rayjobs.yaml b/helm-chart/kuberay-operator/crds/ray.io_rayjobs.yaml
diff --git a/ray-operator/apis/ray/v1/rayjob_types.go b/ray-operator/apis/ray/v1/rayjob_types.go
@@ -86,6 +86,15 @@ type RayJobSpec struct {
 	ClusterSelector map[string]string `json:"clusterSelector,omitempty"`
 	// Configurations of submitter k8s job.
 	SubmitterConfig *SubmitterConfig `json:"submitterConfig,omitempty"`
+	// ManagedBy is used to indicate the controller or entity that manages a RayJob.
+	// The value must be either empty, 'ray.io/kuberay-operator' or 'kueue.x-k8s.io/multikueue'.
+	// The kuberay-operator reconciles a RayJob which doesn't have this field at all or
+	// the field value is the reserved string 'ray.io/kuberay-operator',
+	// but delegates reconciling the RayJob with 'kueue.x-k8s.io/multikueue' to the Kueue.
+	// The field is immutable.
+	// +kubebuilder:validation:XValidation:rule="self == oldSelf",message="the managedBy field is immutable"
+	// +kubebuilder:validation:XValidation:rule="self in ['', 'ray.io/kuberay-operator', 'kueue.x-k8s.io/multikueue']",message="the managedBy field value must be either empty, 'ray.io/kuberay-operator' or 'kueue.x-k8s.io/multikueue'"
+	ManagedBy *string `json:"managedBy,omitempty"`
 	// INSERT ADDITIONAL SPEC FIELDS - desired state of cluster
 	// Important: Run "make" to regenerate code after modifying this file
 	Entrypoint string `json:"entrypoint,omitempty"`
diff --git a/ray-operator/apis/ray/v1/zz_generated.deepcopy.go b/ray-operator/apis/ray/v1/zz_generated.deepcopy.go
diff --git a/ray-operator/config/crd/bases/ray.io_rayjobs.yaml b/ray-operator/config/crd/bases/ray.io_rayjobs.yaml
diff --git a/ray-operator/controllers/ray/rayjob_controller.go b/ray-operator/controllers/ray/rayjob_controller.go
@@ -92,6 +92,11 @@ func (r *RayJobReconciler) Reconcile(ctx context.Context, request ctrl.Request)
 		return ctrl.Result{RequeueAfter: RayJobDefaultRequeueDuration}, err
 	}
 
+	if manager := utils.ManagedByExternalController(rayJobInstance.Spec.ManagedBy); manager != nil {
+		logger.Info("Skipping RayJob managed by a custom controller", "managed-by", manager)
+		return ctrl.Result{}, nil
+	}
+
 	if !rayJobInstance.ObjectMeta.DeletionTimestamp.IsZero() {
 		logger.Info("RayJob is being deleted", "DeletionTimestamp", rayJobInstance.ObjectMeta.DeletionTimestamp)
 		// If the JobStatus is not terminal, it is possible that the Ray job is still running. This includes
diff --git a/ray-operator/controllers/ray/utils/constant.go b/ray-operator/controllers/ray/utils/constant.go
@@ -188,6 +188,9 @@ const (
 	// as well as the user-agent. This constant is updated before release.
 	// TODO: Update KUBERAY_VERSION to be a build-time variable.
 	KUBERAY_VERSION = "nightly"
+
+	// KubeRayController represents the value of the default job controller
+	KubeRayController = "ray.io/kuberay-operator"
 )
 
 type ServiceType string
diff --git a/ray-operator/controllers/ray/utils/util.go b/ray-operator/controllers/ray/utils/util.go
@@ -601,3 +601,10 @@ type ClientProvider interface {
 	GetDashboardClient(mgr manager.Manager) func() RayDashboardClientInterface
 	GetHttpProxyClient(mgr manager.Manager) func() RayHttpProxyClientInterface
 }
+
+func ManagedByExternalController(controllerName *string) *string {
+	if controllerName != nil && *controllerName != KubeRayController {
+		return controllerName
+	}
+	return nil
+}
diff --git a/ray-operator/pkg/client/applyconfiguration/ray/v1/rayjobspec.go b/ray-operator/pkg/client/applyconfiguration/ray/v1/rayjobspec.go
diff --git a/ray-operator/test/e2e/rayjob_cluster_selector_test.go b/ray-operator/test/e2e/rayjob_cluster_selector_test.go
@@ -2,8 +2,10 @@ package e2e
 
 import (
 	"testing"
+	"time"
 
 	. "github.com/onsi/gomega"
+	"k8s.io/apimachinery/pkg/api/errors"
 
 	rayv1 "github.com/ray-project/kuberay/ray-operator/apis/ray/v1"
 	"github.com/ray-project/kuberay/ray-operator/controllers/ray/utils"
@@ -87,4 +89,48 @@ env_vars:
 		g.Expect(GetRayJob(test, rayJob.Namespace, rayJob.Name)).
 			To(WithTransform(RayJobStatus, Equal(rayv1.JobStatusFailed)))
 	})
+
+	test.T().Run("RayJob should be created but not to be updated when managed externally", func(_ *testing.T) {
+		t.Parallel()
+
+		// RayJob
+		rayJobAC := rayv1ac.RayJob("managed-externally", namespace.Name).
+			WithSpec(rayv1ac.RayJobSpec().
+				WithClusterSelector(map[string]string{utils.RayClusterLabelKey: rayCluster.Name}).
+				WithEntrypoint("python /home/ray/jobs/counter.py").
+				WithRuntimeEnvYAML(`
+env_vars:
+  counter_name: test_counter
+`).
+				WithSubmitterPodTemplate(jobSubmitterPodTemplateApplyConfiguration()).
+				WithManagedBy(MultiKueueController))
+
+		rayJob, err := test.Client().Ray().RayV1().RayJobs(namespace.Name).Apply(test.Ctx(), rayJobAC, TestApplyOptions)
+		g.Expect(err).NotTo(HaveOccurred())
+		test.T().Logf("Created RayJob %s/%s successfully", rayJob.Namespace, rayJob.Name)
+
+		// Assert the Ray job status has not been updated
+		g.Consistently(func(gg Gomega) {
+			rayJob, err = GetRayJob(test, rayJob.Namespace, rayJob.Name)
+			gg.Expect(err).ToNot(HaveOccurred())
+			gg.Expect(rayJob.Status.JobDeploymentStatus).To(Equal(rayv1.JobDeploymentStatusNew))
+		}, time.Second*3, time.Millisecond*500).Should(Succeed())
+	})
+
+	test.T().Run("RayJob should not be created due to managedBy invalid value", func(_ *testing.T) {
+		// RayJob
+		rayJobAC := rayv1ac.RayJob("managed-externally-invalid", namespace.Name).
+			WithSpec(rayv1ac.RayJobSpec().
+				WithClusterSelector(map[string]string{utils.RayClusterLabelKey: rayCluster.Name}).
+				WithEntrypoint("python /home/ray/jobs/counter.py").
+				WithRuntimeEnvYAML(`
+env_vars:
+  counter_name: test_counter
+`).
+				WithSubmitterPodTemplate(jobSubmitterPodTemplateApplyConfiguration()).
+				WithManagedBy("invalid.com/controller"))
+
+		_, err := test.Client().Ray().RayV1().RayJobs(namespace.Name).Apply(test.Ctx(), rayJobAC, TestApplyOptions)
+		g.Expect(errors.IsInvalid(err)).To(BeTrue(), "error: %v", err)
+	})
 }
diff --git a/ray-operator/test/e2e/rayjob_test.go b/ray-operator/test/e2e/rayjob_test.go
@@ -6,10 +6,12 @@ import (
 
 	. "github.com/onsi/gomega"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/utils/ptr"
 
 	k8serrors "k8s.io/apimachinery/pkg/api/errors"
 
 	rayv1 "github.com/ray-project/kuberay/ray-operator/apis/ray/v1"
+	"github.com/ray-project/kuberay/ray-operator/controllers/ray/utils"
 	rayv1ac "github.com/ray-project/kuberay/ray-operator/pkg/client/applyconfiguration/ray/v1"
 	. "github.com/ray-project/kuberay/ray-operator/test/support"
 )
@@ -245,4 +247,49 @@ env_vars:
 		g.Expect(GetRayJob(test, rayJob.Namespace, rayJob.Name)).
 			To(WithTransform(RayJobReason, Equal(rayv1.DeadlineExceeded)))
 	})
+
+	test.T().Run("RayJob should be created, but not updated when managed externally", func(_ *testing.T) {
+		// RayJob
+		rayJobAC := rayv1ac.RayJob("managed-externally", namespace.Name).
+			WithSpec(rayv1ac.RayJobSpec().
+				WithRayClusterSpec(newRayClusterSpec(mountConfigMap[rayv1ac.RayClusterSpecApplyConfiguration](jobs, "/home/ray/jobs"))).
+				WithEntrypoint("python /home/ray/jobs/counter.py").
+				WithRuntimeEnvYAML(`
+env_vars:
+  counter_name: test_counter
+`).
+				WithShutdownAfterJobFinishes(true).
+				WithSubmitterPodTemplate(jobSubmitterPodTemplateApplyConfiguration()).
+				WithManagedBy(MultiKueueController))
+
+		rayJob, err := test.Client().Ray().RayV1().RayJobs(namespace.Name).Apply(test.Ctx(), rayJobAC, TestApplyOptions)
+		g.Expect(err).NotTo(HaveOccurred())
+		test.T().Logf("Created RayJob %s/%s successfully", rayJob.Namespace, rayJob.Name)
+
+		// Should not to be able to change managedBy field as it's immutable
+		rayJobAC.Spec.WithManagedBy(utils.KubeRayController)
+		_, err = test.Client().Ray().RayV1().RayJobs(namespace.Name).Apply(test.Ctx(), rayJobAC, TestApplyOptions)
+		g.Expect(err).To(HaveOccurred())
+		g.Eventually(RayJob(test, *rayJobAC.Namespace, *rayJobAC.Name)).
+			Should(WithTransform(RayJobManagedBy, Equal(ptr.To(MultiKueueController))))
+
+		// Refresh the RayJob status and assert it has not been updated
+		g.Eventually(RayJob(test, rayJob.Namespace, rayJob.Name)).
+			Should(WithTransform(RayJobDeploymentStatus, Equal(rayv1.JobDeploymentStatusNew)))
+
+		// Assert the associated RayCluster has not beed created
+		rcList, err := test.Client().Ray().RayV1().RayClusters(rayJob.Namespace).List(test.Ctx(), metav1.ListOptions{})
+		g.Expect(err).NotTo(HaveOccurred())
+		for _, rc := range rcList.Items {
+			g.Expect(rc.Name).NotTo(HaveSuffix(*rayJobAC.Name))
+		}
+
+		// Assert the submitter Job has not been created
+		g.Eventually(Jobs(test, namespace.Name)).Should(BeEmpty())
+
+		// Delete the RayJob
+		err = test.Client().Ray().RayV1().RayJobs(namespace.Name).Delete(test.Ctx(), *rayJobAC.Name, metav1.DeleteOptions{})
+		g.Expect(err).NotTo(HaveOccurred())
+		test.T().Logf("Deleted RayJob %s/%s successfully", *rayJobAC.Namespace, *rayJobAC.Name)
+	})
 }
diff --git a/ray-operator/test/support/ray.go b/ray-operator/test/support/ray.go
@@ -30,6 +30,10 @@ func RayJobDeploymentStatus(job *rayv1.RayJob) rayv1.JobDeploymentStatus {
 	return job.Status.JobDeploymentStatus
 }
 
+func RayJobManagedBy(job *rayv1.RayJob) *string {
+	return job.Spec.ManagedBy
+}
+
 func RayJobReason(job *rayv1.RayJob) rayv1.JobFailedReason {
 	return job.Status.Reason
 }
diff --git a/ray-operator/test/support/utils.go b/ray-operator/test/support/utils.go
@@ -16,6 +16,8 @@ type OutputType string
 
 const (
 	Log OutputType = "log"
+	// MultiKueueController represents the vaue of the MultiKueue controller
+	MultiKueueController = "kueue.x-k8s.io/multikueue"
 )
 
 func WriteToOutputDir(t Test, fileName string, fileType OutputType, data []byte) {

Original file line number	Diff line number	Diff line change
`@@ -188,6 +188,9 @@ const (`
`188`	`188`	`// as well as the user-agent. This constant is updated before release.`
`189`	`189`	`// TODO: Update KUBERAY_VERSION to be a build-time variable.`
`190`	`190`	`KUBERAY_VERSION = "nightly"`
	`191`	`+`
	`192`	`+ // KubeRayController represents the value of the default job controller`
	`193`	`+ KubeRayController = "ray.io/kuberay-operator"`
`191`	`194`	`)`
`192`	`195`
`193`	`196`	`type ServiceType string`