feat: add grace period after sumitter finished (#4091)

* feat: add grace period after sumitter finished

* feat: generate new field in rayjob CRD

* fix: fix terminate logic + add default timeout value

* refactor: remove submitter finish timeout in rayjob status & use default timeout

* feat: get finish time from job and container

* test: for submitter finished timeout

* fix: set finishedAt to nil if submitter not finished

* test: ensure timeout close to the set timeout

* refactor: make status msg more readable

* fix: simplify finishedAt time format

* refactor: remove isSubmitterFinished and use finishedAt only

* fix: add LastTransitionTime to JobCondition in test

* Trigger CI

Signed-off-by: machichima <[email protected]>

* refactor: fix lint and nit

* fix: fix test

---------

Signed-off-by: machichima <[email protected]>

Author: machichima

ray-operator/controllers/ray/rayjob_controller.go

@@ -34,8 +34,9 @@ import (
 )
 
 const (
-	RayJobDefaultRequeueDuration = 3 * time.Second
-	PythonUnbufferedEnvVarName   = "PYTHONUNBUFFERED"
+	RayJobDefaultRequeueDuration    = 3 * time.Second
+	PythonUnbufferedEnvVarName      = "PYTHONUNBUFFERED"
+	DefaultSubmitterFinishedTimeout = 30 * time.Second
 )
 
 // RayJobReconciler reconciles a RayJob object
@@ -258,13 +259,18 @@ func (r *RayJobReconciler) Reconcile(ctx context.Context, request ctrl.Request)
 			return ctrl.Result{RequeueAfter: RayJobDefaultRequeueDuration}, err
 		}
 
-		var isSubmitterFinished bool
+		var finishedAt *time.Time
 		if rayJobInstance.Spec.SubmissionMode == rayv1.K8sJobMode || rayJobInstance.Spec.SubmissionMode == rayv1.SidecarMode {
 			var shouldUpdate bool
-			shouldUpdate, isSubmitterFinished, err = r.checkSubmitterAndUpdateStatusIfNeeded(ctx, rayJobInstance)
+			shouldUpdate, finishedAt, err = r.checkSubmitterAndUpdateStatusIfNeeded(ctx, rayJobInstance)
 			if err != nil {
 				return ctrl.Result{RequeueAfter: RayJobDefaultRequeueDuration}, err
 			}
+
+			if checkSubmitterFinishedTimeoutAndUpdateStatusIfNeeded(ctx, rayJobInstance, finishedAt) {
+				break
+			}
+
 			if shouldUpdate {
 				break
 			}
@@ -288,7 +294,8 @@ func (r *RayJobReconciler) Reconcile(ctx context.Context, request ctrl.Request)
 					}
 					return ctrl.Result{RequeueAfter: RayJobDefaultRequeueDuration}, nil
 				}
-				if isSubmitterFinished {
+				// finishedAt will only be set if submitter finished
+				if finishedAt != nil {
 					rayJobInstance.Status.JobDeploymentStatus = rayv1.JobDeploymentStatusFailed
 					rayJobInstance.Status.Reason = rayv1.AppFailed
 					rayJobInstance.Status.Message = "Submitter completed but Ray job not found in RayCluster."
@@ -306,10 +313,10 @@ func (r *RayJobReconciler) Reconcile(ctx context.Context, request ctrl.Request)
 		jobDeploymentStatus := rayv1.JobDeploymentStatusRunning
 		reason := rayv1.JobFailedReason("")
 		isJobTerminal := rayv1.IsJobTerminal(jobInfo.JobStatus)
-		// If in K8sJobMode, further refine the terminal condition by checking if the submitter Job has finished.
+		// If in K8sJobMode or SidecarMode, further refine the terminal condition by checking if the submitter has finished.
 		// See https://github.com/ray-project/kuberay/pull/1919 for reasons.
 		if utils.HasSubmitter(rayJobInstance) {
-			isJobTerminal = isJobTerminal && isSubmitterFinished
+			isJobTerminal = isJobTerminal && finishedAt != nil
 		}
 
 		if isJobTerminal {
@@ -973,10 +980,10 @@ func updateStatusToSuspendingIfNeeded(ctx context.Context, rayJob *rayv1.RayJob)
 	return true
 }
 
-func (r *RayJobReconciler) checkSubmitterAndUpdateStatusIfNeeded(ctx context.Context, rayJob *rayv1.RayJob) (shouldUpdate, isSubmitterFinished bool, err error) {
+func (r *RayJobReconciler) checkSubmitterAndUpdateStatusIfNeeded(ctx context.Context, rayJob *rayv1.RayJob) (shouldUpdate bool, finishedAt *time.Time, err error) {
 	logger := ctrl.LoggerFrom(ctx)
 	shouldUpdate = false
-	isSubmitterFinished = false
+	finishedAt = nil
 	var submitterContainerStatus *corev1.ContainerStatus
 	var condition *batchv1.JobCondition
 
@@ -1023,7 +1030,7 @@ func (r *RayJobReconciler) checkSubmitterAndUpdateStatusIfNeeded(ctx context.Con
 			}
 		}
 
-		isSubmitterFinished = isSubmitterContainerFinished(headPod)
+		finishedAt = getSubmitterContainerFinishedTime(headPod)
 		return
 	case rayv1.K8sJobMode:
 		job := &batchv1.Job{}
@@ -1037,6 +1044,7 @@ func (r *RayJobReconciler) checkSubmitterAndUpdateStatusIfNeeded(ctx context.Con
 			logger.Error(err, "Failed to get the submitter Kubernetes Job for RayJob", "NamespacedName", namespacedName)
 			return
 		}
+
 		shouldUpdate, condition = checkK8sJobStatus(job)
 		if shouldUpdate {
 			logger.Info("The submitter Kubernetes Job has failed. Attempting to transition the status to `Failed`.",
@@ -1052,7 +1060,14 @@ func (r *RayJobReconciler) checkSubmitterAndUpdateStatusIfNeeded(ctx context.Con
 				rayJob.Status.Message = fmt.Sprintf("Job submission has failed. Reason: %s. Message: %s", condition.Reason, condition.Message)
 			}
 		}
-		_, isSubmitterFinished = utils.IsJobFinished(job)
+
+		var jobFinishedCondition *batchv1.JobCondition
+		// Find the terminal condition to get its LastTransitionTime
+		jobFinishedCondition = getJobFinishedCondition(job)
+		if jobFinishedCondition != nil && !jobFinishedCondition.LastTransitionTime.IsZero() {
+			finishedAt = &jobFinishedCondition.LastTransitionTime.Time
+		}
+
 		return
 	default:
 		// This means that the KubeRay logic is wrong, and it's better to panic as a system error than to allow the operator to
@@ -1071,6 +1086,15 @@ func checkK8sJobStatus(job *batchv1.Job) (bool, *batchv1.JobCondition) {
 	return false, nil
 }
 
+func getJobFinishedCondition(job *batchv1.Job) *batchv1.JobCondition {
+	for _, condition := range job.Status.Conditions {
+		if (condition.Type == batchv1.JobComplete || condition.Type == batchv1.JobFailed) && condition.Status == corev1.ConditionTrue {
+			return &condition
+		}
+	}
+	return nil
+}
+
 func checkSidecarContainerStatus(headPod *corev1.Pod) (bool, *corev1.ContainerStatus) {
 	for _, containerStatus := range headPod.Status.ContainerStatuses {
 		if containerStatus.Name == utils.SubmitterContainerName {
@@ -1099,6 +1123,31 @@ func checkActiveDeadlineAndUpdateStatusIfNeeded(ctx context.Context, rayJob *ray
 	return true
 }
 
+func checkSubmitterFinishedTimeoutAndUpdateStatusIfNeeded(ctx context.Context, rayJob *rayv1.RayJob, finishedAt *time.Time) bool {
+	logger := ctrl.LoggerFrom(ctx)
+
+	// Check if timeout is configured and submitter has finished
+	if finishedAt == nil {
+		return false
+	}
+
+	// Check if timeout has been exceeded
+	if time.Now().Before(finishedAt.Add(DefaultSubmitterFinishedTimeout)) {
+		return false
+	}
+
+	logger.Info("The RayJob has passed the submitterFinishedTimeoutSeconds. Transition the status to terminal.",
+		"SubmitterFinishedTime", finishedAt,
+		"SubmitterFinishedTimeoutSeconds", DefaultSubmitterFinishedTimeout.String())
+
+	rayJob.Status.JobStatus = rayv1.JobStatusFailed
+	rayJob.Status.JobDeploymentStatus = rayv1.JobDeploymentStatusFailed
+	rayJob.Status.Reason = rayv1.JobDeploymentStatusTransitionGracePeriodExceeded
+	rayJob.Status.Message = fmt.Sprintf("The RayJob submitter finished at %v but the ray job did not reach terminal state within %v",
+		finishedAt.Format(time.DateTime), DefaultSubmitterFinishedTimeout)
+	return true
+}
+
 func checkTransitionGracePeriodAndUpdateStatusIfNeeded(ctx context.Context, rayJob *rayv1.RayJob) bool {
 	logger := ctrl.LoggerFrom(ctx)
 	if rayv1.IsJobTerminal(rayJob.Status.JobStatus) && rayJob.Status.JobDeploymentStatus == rayv1.JobDeploymentStatusRunning {
@@ -1129,16 +1178,16 @@ func checkTransitionGracePeriodAndUpdateStatusIfNeeded(ctx context.Context, rayJ
 	return false
 }
 
-func isSubmitterContainerFinished(pod *corev1.Pod) bool {
+func getSubmitterContainerFinishedTime(pod *corev1.Pod) *time.Time {
 	for _, containerStatus := range pod.Status.ContainerStatuses {
 		if containerStatus.Name == utils.SubmitterContainerName {
 			if containerStatus.State.Terminated != nil {
-				return true
+				return &containerStatus.State.Terminated.FinishedAt.Time
 			}
 			break
 		}
 	}
-	return false
+	return nil
 }
 
 // handleDeletionRules processes the DeletionRules with a impact-aware strategy.

ray-operator/controllers/ray/rayjob_controller_test.go

@@ -290,7 +290,7 @@ var _ = Context("RayJob with different submission modes", func() {
 
 				// Update the submitter Kubernetes Job to Complete.
 				conditions := []batchv1.JobCondition{
-					{Type: batchv1.JobComplete, Status: corev1.ConditionTrue},
+					{Type: batchv1.JobComplete, Status: corev1.ConditionTrue, LastTransitionTime: metav1.Now()},
 				}
 				job.Status.Conditions = conditions
 				Expect(k8sClient.Status().Update(ctx, job)).Should(Succeed())
@@ -515,7 +515,7 @@ var _ = Context("RayJob with different submission modes", func() {
 
 				// Update the submitter Kubernetes Job to Complete.
 				conditions := []batchv1.JobCondition{
-					{Type: batchv1.JobComplete, Status: corev1.ConditionTrue},
+					{Type: batchv1.JobComplete, Status: corev1.ConditionTrue, LastTransitionTime: metav1.Now()},
 				}
 				job.Status.Conditions = conditions
 				Expect(k8sClient.Status().Update(ctx, job)).Should(Succeed())
@@ -680,7 +680,7 @@ var _ = Context("RayJob with different submission modes", func() {
 
 				// Update the submitter Kubernetes Job to Complete.
 				conditions := []batchv1.JobCondition{
-					{Type: batchv1.JobComplete, Status: corev1.ConditionTrue},
+					{Type: batchv1.JobComplete, Status: corev1.ConditionTrue, LastTransitionTime: metav1.Now()},
 				}
 				job.Status.Conditions = conditions
 				Expect(k8sClient.Status().Update(ctx, job)).Should(Succeed())
@@ -775,7 +775,7 @@ var _ = Context("RayJob with different submission modes", func() {
 
 				// Update the submitter Kubernetes Job to Complete.
 				conditions := []batchv1.JobCondition{
-					{Type: batchv1.JobComplete, Status: corev1.ConditionTrue},
+					{Type: batchv1.JobComplete, Status: corev1.ConditionTrue, LastTransitionTime: metav1.Now()},
 				}
 				job.Status.Conditions = conditions
 				Expect(k8sClient.Status().Update(ctx, job)).Should(Succeed())
@@ -1002,7 +1002,7 @@ var _ = Context("RayJob with different submission modes", func() {
 
 				// Update the submitter Kubernetes Job to Complete.
 				conditions := []batchv1.JobCondition{
-					{Type: batchv1.JobComplete, Status: corev1.ConditionTrue},
+					{Type: batchv1.JobComplete, Status: corev1.ConditionTrue, LastTransitionTime: metav1.Now()},
 				}
 				job.Status.Conditions = conditions
 				Expect(k8sClient.Status().Update(ctx, job)).Should(Succeed())
@@ -1141,7 +1141,7 @@ var _ = Context("RayJob with different submission modes", func() {
 
 				// Update the submitter Kubernetes Job to Complete.
 				conditions := []batchv1.JobCondition{
-					{Type: batchv1.JobComplete, Status: corev1.ConditionTrue},
+					{Type: batchv1.JobComplete, Status: corev1.ConditionTrue, LastTransitionTime: metav1.Now()},
 				}
 				job.Status.Conditions = conditions
 				Expect(k8sClient.Status().Update(ctx, job)).Should(Succeed())
@@ -1280,7 +1280,7 @@ var _ = Context("RayJob with different submission modes", func() {
 
 				// Update the submitter Kubernetes Job to Complete.
 				conditions := []batchv1.JobCondition{
-					{Type: batchv1.JobComplete, Status: corev1.ConditionTrue},
+					{Type: batchv1.JobComplete, Status: corev1.ConditionTrue, LastTransitionTime: metav1.Now()},
 				}
 				job.Status.Conditions = conditions
 				Expect(k8sClient.Status().Update(ctx, job)).Should(Succeed())
@@ -1436,7 +1436,7 @@ var _ = Context("RayJob with different submission modes", func() {
 
 				// Update the submitter Kubernetes Job to Complete.
 				conditions := []batchv1.JobCondition{
-					{Type: batchv1.JobComplete, Status: corev1.ConditionTrue},
+					{Type: batchv1.JobComplete, Status: corev1.ConditionTrue, LastTransitionTime: metav1.Now()},
 				}
 				job.Status.Conditions = conditions
 				Expect(k8sClient.Status().Update(ctx, job)).Should(Succeed())
@@ -1580,7 +1580,7 @@ var _ = Context("RayJob with different submission modes", func() {
 
 				// Update the submitter Kubernetes Job to Complete.
 				conditions := []batchv1.JobCondition{
-					{Type: batchv1.JobComplete, Status: corev1.ConditionTrue},
+					{Type: batchv1.JobComplete, Status: corev1.ConditionTrue, LastTransitionTime: metav1.Now()},
 				}
 				job.Status.Conditions = conditions
 				Expect(k8sClient.Status().Update(ctx, job)).Should(Succeed())
@@ -1696,7 +1696,7 @@ var _ = Context("RayJob with different submission modes", func() {
 
 				// Update the submitter Kubernetes Job to Complete.
 				conditions := []batchv1.JobCondition{
-					{Type: batchv1.JobComplete, Status: corev1.ConditionTrue},
+					{Type: batchv1.JobComplete, Status: corev1.ConditionTrue, LastTransitionTime: metav1.Now()},
 				}
 				job.Status.Conditions = conditions
 				Expect(k8sClient.Status().Update(ctx, job)).Should(Succeed())
@@ -1824,7 +1824,7 @@ var _ = Context("RayJob with different submission modes", func() {
 
 				// Update the submitter Kubernetes Job to Complete.
 				conditions := []batchv1.JobCondition{
-					{Type: batchv1.JobComplete, Status: corev1.ConditionTrue},
+					{Type: batchv1.JobComplete, Status: corev1.ConditionTrue, LastTransitionTime: metav1.Now()},
 				}
 				job.Status.Conditions = conditions
 				Expect(k8sClient.Status().Update(ctx, job)).Should(Succeed())
@@ -1985,7 +1985,7 @@ var _ = Context("RayJob with different submission modes", func() {
 
 				// Update the submitter Kubernetes Job to Complete.
 				conditions := []batchv1.JobCondition{
-					{Type: batchv1.JobComplete, Status: corev1.ConditionTrue},
+					{Type: batchv1.JobComplete, Status: corev1.ConditionTrue, LastTransitionTime: metav1.Now()},
 				}
 				job.Status.Conditions = conditions
 				Expect(k8sClient.Status().Update(ctx, job)).Should(Succeed())

ray-operator/test/e2erayjob/rayjob_test.go

@@ -3,13 +3,18 @@ package e2erayjob
 import (
 	"strings"
 	"testing"
+	"time"
 
 	. "github.com/onsi/gomega"
+	"github.com/stretchr/testify/assert"
+	batchv1 "k8s.io/api/batch/v1"
+	corev1 "k8s.io/api/core/v1"
 	k8serrors "k8s.io/apimachinery/pkg/api/errors"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/utils/ptr"
 
 	rayv1 "github.com/ray-project/kuberay/ray-operator/apis/ray/v1"
+	"github.com/ray-project/kuberay/ray-operator/controllers/ray"
 	"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"
@@ -360,4 +365,73 @@ env_vars:
 		g.Expect(err).NotTo(HaveOccurred())
 		LogWithTimestamp(test.T(), "Deleted RayJob %s/%s successfully", *rayJobAC.Namespace, *rayJobAC.Name)
 	})
+
+	test.T().Run("RayJob has exceed SubmitterFinishedTimeout", func(_ *testing.T) {
+		rayJobAC := rayv1ac.RayJob("submitter-timeout", namespace.Name).
+			WithSpec(rayv1ac.RayJobSpec().
+				WithRayClusterSpec(NewRayClusterSpec(MountConfigMap[rayv1ac.RayClusterSpecApplyConfiguration](jobs, "/home/ray/jobs"))).
+				WithEntrypoint("python /home/ray/jobs/long_running.py").
+				WithShutdownAfterJobFinishes(true).
+				WithSubmitterPodTemplate(JobSubmitterPodTemplateApplyConfiguration()))
+
+		rayJob, err := test.Client().Ray().RayV1().RayJobs(namespace.Name).Apply(test.Ctx(), rayJobAC, TestApplyOptions)
+		g.Expect(err).NotTo(HaveOccurred())
+		LogWithTimestamp(test.T(), "Created RayJob %s/%s successfully", rayJob.Namespace, rayJob.Name)
+
+		// Wait until RayJob status and deployment status is running
+		LogWithTimestamp(test.T(), "Waiting for RayJob %s/%s deployment status to be 'Running'", rayJob.Namespace, rayJob.Name)
+		g.Eventually(RayJob(test, rayJob.Namespace, rayJob.Name), TestTimeoutMedium).
+			Should(WithTransform(RayJobDeploymentStatus, Equal(rayv1.JobDeploymentStatusRunning)))
+		LogWithTimestamp(test.T(), "Waiting for Ray job %s/%s to be actually running in Ray cluster", rayJob.Namespace, rayJob.Name)
+		g.Eventually(RayJob(test, rayJob.Namespace, rayJob.Name), TestTimeoutMedium).
+			Should(WithTransform(RayJobStatus, Equal(rayv1.JobStatusRunning)))
+
+		// Wait for the submitter job to be created
+		LogWithTimestamp(test.T(), "Waiting for submitter job to be created")
+		g.Eventually(Jobs(test, namespace.Name)).ShouldNot(BeEmpty())
+
+		// Force the submitter job to complete by updating its status
+		// This simulates the submitter finishing but the Ray job still running
+		LogWithTimestamp(test.T(), "Updating submitter job status to complete")
+		job, err := test.Client().Core().BatchV1().Jobs(namespace.Name).Get(test.Ctx(), rayJob.Name, metav1.GetOptions{})
+		g.Expect(err).NotTo(HaveOccurred())
+		now := metav1.Now()
+		job.Status.Conditions = append(job.Status.Conditions, batchv1.JobCondition{
+			Type:               batchv1.JobComplete,
+			Status:             corev1.ConditionTrue,
+			LastProbeTime:      now,
+			LastTransitionTime: now,
+			Reason:             "Completed",
+			Message:            "Job completed successfully for timeout test",
+		})
+		job.Status.CompletionTime = &now
+		job.Status.Succeeded = 1
+
+		_, err = test.Client().Core().BatchV1().Jobs(namespace.Name).UpdateStatus(test.Ctx(), job, metav1.UpdateOptions{})
+		g.Expect(err).NotTo(HaveOccurred())
+		LogWithTimestamp(test.T(), "Successfully marked submitter job as completed at %v", now.Time)
+
+		// Record the start time for timeout measurement
+		timeoutStartTime := time.Now()
+
+		// Wait for the timeout to trigger
+		LogWithTimestamp(test.T(), "Waiting for RayJob %s/%s to exceed SubmitterFinishedTimeout", rayJob.Namespace, rayJob.Name)
+		g.Eventually(RayJob(test, rayJob.Namespace, rayJob.Name), TestTimeoutShort).
+			Should(WithTransform(RayJobDeploymentStatus, Equal(rayv1.JobDeploymentStatusFailed)))
+
+		// Measure the actual timeout duration and verify the timeout duration is close to DefaultSubmitterFinishedTimeout
+		actualTimeoutDuration := time.Since(timeoutStartTime)
+		expectedTimeout := ray.DefaultSubmitterFinishedTimeout
+		assert.InDelta(test.T(), expectedTimeout.Seconds(), actualTimeoutDuration.Seconds(), 5.0,
+			"Actual timeout duration should be close to DefaultSubmitterFinishedTimeout")
+
+		// Get the updated rayJob
+		rayJob, err = GetRayJob(test, rayJob.Namespace, rayJob.Name)
+		g.Expect(err).NotTo(HaveOccurred())
+
+		reason := rayJob.Status.Reason
+		message := rayJob.Status.Message
+		g.Expect(reason).To(Equal(rayv1.JobDeploymentStatusTransitionGracePeriodExceeded))
+		g.Expect(message).To(MatchRegexp(`The RayJob submitter finished at .* but the ray job did not reach terminal state within .*`))
+	})
 }