[RayJob][Status][5/n] Refactor getOrCreateK8sJob (#1750)

Author: kevin85421

ray-operator/controllers/ray/rayjob_controller.go

@@ -212,28 +212,7 @@ func (r *RayJobReconciler) Reconcile(ctx context.Context, request ctrl.Request)
 	rayDashboardClient.InitClient(rayJobInstance.Status.DashboardURL)
 
 	// Ensure k8s job has been created
-	jobName, wasJobCreated, err := r.getOrCreateK8sJob(ctx, rayJobInstance, rayClusterInstance)
-	if err != nil {
-		return ctrl.Result{RequeueAfter: RayJobDefaultRequeueDuration}, err
-	}
-
-	if wasJobCreated {
-		r.Log.Info("K8s job successfully created", "RayJob", rayJobInstance.Name, "jobId", jobName)
-		r.Recorder.Eventf(rayJobInstance, corev1.EventTypeNormal, "Created", "Created k8s job %s", jobName)
-	} else {
-		r.Log.Info("K8s job successfully retrieved", "RayJob", rayJobInstance.Name, "jobId", jobName)
-	}
-
-	// Check the status of the k8s job and update the RayJobInstance status accordingly.
-	// Get the k8s job
-	k8sJob := &batchv1.Job{}
-	err = r.Client.Get(ctx, types.NamespacedName{Name: jobName, Namespace: rayJobInstance.Namespace}, k8sJob)
-	if err != nil {
-		if errors.IsNotFound(err) {
-			r.Log.Info("Job not found", "RayJob", rayJobInstance.Name, "jobId", jobName)
-			return ctrl.Result{RequeueAfter: RayJobDefaultRequeueDuration}, err
-		}
-		r.Log.Error(err, "failed to get k8s job")
+	if err := r.createK8sJobIfNeed(ctx, rayJobInstance, rayClusterInstance); err != nil {
 		return ctrl.Result{RequeueAfter: RayJobDefaultRequeueDuration}, err
 	}
 
@@ -262,17 +241,13 @@ func (r *RayJobReconciler) Reconcile(ctx context.Context, request ctrl.Request)
 		// the RayJob is submitted against the RayCluster created by THIS job, then
 		// try to gracefully stop the Ray job and delete (suspend) the cluster
 		if rayJobInstance.Spec.Suspend && len(rayJobInstance.Spec.ClusterSelector) == 0 {
-			info, err := rayDashboardClient.GetJobInfo(ctx, rayJobInstance.Status.JobId)
-			if err != nil {
-				return ctrl.Result{RequeueAfter: RayJobDefaultRequeueDuration}, err
-			}
-			if !rayv1.IsJobTerminal(info.JobStatus) {
+			if !rayv1.IsJobTerminal(jobInfo.JobStatus) {
 				err := rayDashboardClient.StopJob(ctx, rayJobInstance.Status.JobId, &r.Log)
 				if err != nil {
 					return ctrl.Result{RequeueAfter: RayJobDefaultRequeueDuration}, err
 				}
 			}
-			if info.JobStatus != rayv1.JobStatusStopped {
+			if jobInfo.JobStatus != rayv1.JobStatusStopped {
 				return ctrl.Result{RequeueAfter: RayJobDefaultRequeueDuration}, nil
 			}
 
@@ -325,6 +300,7 @@ func (r *RayJobReconciler) Reconcile(ctx context.Context, request ctrl.Request)
 		}
 	}
 
+	// TODO (kevin85421): Use the source of truth `jobInfo.JobStatus` instead.
 	if isJobPendingOrRunning(rayJobInstance.Status.JobStatus) {
 		// Requeue the RayJob to poll its status from the running Ray job
 		return ctrl.Result{RequeueAfter: RayJobDefaultRequeueDuration}, nil
@@ -334,8 +310,8 @@ func (r *RayJobReconciler) Reconcile(ctx context.Context, request ctrl.Request)
 	return ctrl.Result{}, nil
 }
 
-// getOrCreateK8sJob creates a Kubernetes Job for the Ray Job if it doesn't exist, otherwise returns the existing one. It returns the Job name and a boolean indicating whether the Job was created.
-func (r *RayJobReconciler) getOrCreateK8sJob(ctx context.Context, rayJobInstance *rayv1.RayJob, rayClusterInstance *rayv1.RayCluster) (string, bool, error) {
+// createK8sJobIfNeed creates a Kubernetes Job for the RayJob if it doesn't exist.
+func (r *RayJobReconciler) createK8sJobIfNeed(ctx context.Context, rayJobInstance *rayv1.RayJob, rayClusterInstance *rayv1.RayCluster) error {
 	jobName := rayJobInstance.Name
 	jobNamespace := rayJobInstance.Namespace
 
@@ -346,18 +322,18 @@ func (r *RayJobReconciler) getOrCreateK8sJob(ctx context.Context, rayJobInstance
 			submitterTemplate, err := r.getSubmitterTemplate(rayJobInstance, rayClusterInstance)
 			if err != nil {
 				r.Log.Error(err, "failed to get submitter template")
-				return "", false, err
+				return err
 			}
 			return r.createNewK8sJob(ctx, rayJobInstance, submitterTemplate)
 		}
 
 		// Some other error occurred while trying to get the Job
-		r.Log.Error(err, "failed to get k8s Job")
-		return "", false, err
+		r.Log.Error(err, "failed to get Kubernetes Job")
+		return err
 	}
 
-	// Job already exists, instead of returning an error we return a "success"
-	return jobName, false, nil
+	r.Log.Info("Kubernetes Job already exists", "RayJob", rayJobInstance.Name, "Kubernetes Job", job.Name)
+	return nil
 }
 
 // getSubmitterTemplate builds the submitter pod template for the Ray job.
@@ -399,8 +375,8 @@ func (r *RayJobReconciler) getSubmitterTemplate(rayJobInstance *rayv1.RayJob, ra
 	return submitterTemplate, nil
 }
 
-// createNewK8sJob creates a new Kubernetes Job. It returns the Job's name and a boolean indicating whether a new Job was created.
-func (r *RayJobReconciler) createNewK8sJob(ctx context.Context, rayJobInstance *rayv1.RayJob, submitterTemplate corev1.PodTemplateSpec) (string, bool, error) {
+// createNewK8sJob creates a new Kubernetes Job. It returns an error.
+func (r *RayJobReconciler) createNewK8sJob(ctx context.Context, rayJobInstance *rayv1.RayJob, submitterTemplate corev1.PodTemplateSpec) error {
 	job := &batchv1.Job{
 		ObjectMeta: metav1.ObjectMeta{
 			Name:      rayJobInstance.Name,
@@ -422,17 +398,17 @@ func (r *RayJobReconciler) createNewK8sJob(ctx context.Context, rayJobInstance *
 	// Set the ownership in order to do the garbage collection by k8s.
 	if err := ctrl.SetControllerReference(rayJobInstance, job, r.Scheme); err != nil {
 		r.Log.Error(err, "failed to set controller reference")
-		return "", false, err
+		return err
 	}
 
 	// Create the Kubernetes Job
 	if err := r.Client.Create(ctx, job); err != nil {
 		r.Log.Error(err, "failed to create k8s Job")
-		return "", false, err
+		return err
 	}
-
-	// Return the Job's name and true indicating a new job was created
-	return job.Name, true, nil
+	r.Log.Info("Kubernetes Job created", "RayJob", rayJobInstance.Name, "Kubernetes Job", job.Name)
+	r.Recorder.Eventf(rayJobInstance, corev1.EventTypeNormal, "Created", "Created Kubernetes Job %s", job.Name)
+	return nil
 }
 
 func (r *RayJobReconciler) deleteCluster(ctx context.Context, rayJobInstance *rayv1.RayJob) (reconcile.Result, error) {

ray-operator/controllers/ray/rayjob_controller_unit_test.go

@@ -16,7 +16,7 @@ import (
 	clientFake "sigs.k8s.io/controller-runtime/pkg/client/fake"
 )
 
-func TestGetOrCreateK8sJob(t *testing.T) {
+func TestCreateK8sJobIfNeed(t *testing.T) {
 	newScheme := runtime.NewScheme()
 	_ = rayv1.AddToScheme(newScheme)
 	_ = batchv1.AddToScheme(newScheme)
@@ -67,21 +67,15 @@ func TestGetOrCreateK8sJob(t *testing.T) {
 		Recorder: &record.FakeRecorder{},
 	}
 
-	retrievedJobName, wasCreated, err := rayJobReconciler.getOrCreateK8sJob(ctx, rayJob, rayCluster)
-
+	err := rayJobReconciler.createK8sJobIfNeed(ctx, rayJob, rayCluster)
 	assert.NoError(t, err)
-	assert.False(t, wasCreated)
-	assert.Equal(t, "test-rayjob", retrievedJobName)
 
 	// Test 2: Create a new k8s job if it does not already exist
 	fakeClient = clientFake.NewClientBuilder().WithScheme(newScheme).WithRuntimeObjects(rayCluster, rayJob).Build()
 	rayJobReconciler.Client = fakeClient
 
-	retrievedJobName, wasCreated, err = rayJobReconciler.getOrCreateK8sJob(ctx, rayJob, rayCluster)
-
+	err = rayJobReconciler.createK8sJobIfNeed(ctx, rayJob, rayCluster)
 	assert.NoError(t, err)
-	assert.True(t, wasCreated)
-	assert.Equal(t, "test-rayjob", retrievedJobName)
 }
 
 func TestGetSubmitterTemplate(t *testing.T) {