[Bug][RayJob] Sidecar mode shouldn't restart head pod when head pod is deleted

ray-operator/controllers/ray/raycluster_controller.go

@@ -679,8 +679,8 @@ func (r *RayClusterReconciler) reconcilePods(ctx context.Context, instance *rayv
 			return errstd.New(reason)
 		}
 	} else if len(headPods.Items) == 0 {
-		originatedFrom := utils.GetCRDType(instance.Labels[utils.RayOriginatedFromCRDLabelKey])
-		if originatedFrom == utils.RayJobCRD {
+		if shouldSkipHeadPodRestart(instance) &&
+			meta.IsStatusConditionTrue(instance.Status.Conditions, string(rayv1.RayClusterProvisioned)) {
 			// Recreating the head Pod if the RayCluster created by RayJob is provisioned doesn't help RayJob.
 			//
 			// Case 1: GCS fault tolerance is disabled
@@ -692,13 +692,11 @@ func (r *RayClusterReconciler) reconcilePods(ctx context.Context, instance *rayv
 			//
 			// In this case, the worker Pods will not be killed by the new head Pod when it is created, but the submission ID has already been
 			// used by the old Ray job, so the new Ray job will fail.
-			if meta.IsStatusConditionTrue(instance.Status.Conditions, string(rayv1.RayClusterProvisioned)) {
-				logger.Info(
-					"reconcilePods: Found 0 head Pods for a RayJob-managed RayCluster; skipping head creation to let RayJob controller handle the failure",
-					"rayCluster", instance.Name,
-				)
-				return nil
-			}
+			logger.Info(
+				"reconcilePods: Found 0 head Pods for a sidecar-mode RayJob-managed RayCluster; skipping head creation to let RayJob controller handle the failure",
+				"rayCluster", instance.Name,
+			)
+			return nil
 		}
 		// Create head Pod if it does not exist.
 		logger.Info("reconcilePods: Found 0 head Pods; creating a head Pod for the RayCluster.")
@@ -830,7 +828,7 @@ func (r *RayClusterReconciler) reconcilePods(ctx context.Context, instance *rayv
 					}
 				}
 			}
-			logger.Info("reconcilePods", "found existing replica indices", "group", worker.GroupName, "indices", validReplicaIndices)
+			logger.Info("reconcilePods: found existing replica indices", "group", worker.GroupName, "indices", validReplicaIndices)
 		}
 		if diff > 0 {
 			// pods need to be added
@@ -1094,6 +1092,16 @@ func (r *RayClusterReconciler) reconcileMultiHostWorkerGroup(ctx context.Context
 	return nil
 }
 
+func shouldSkipHeadPodRestart(cluster *rayv1.RayCluster) bool {
+	if cluster == nil {
+		return false
+	}
+	if utils.GetCRDType(cluster.Labels[utils.RayOriginatedFromCRDLabelKey]) != utils.RayJobCRD {
+		return false
+	}
+	return cluster.Labels[utils.RayJobDisableHeadNodeRestartLabelKey] == "true"
+}
+
 // shouldDeletePod returns whether the Pod should be deleted and the reason
 //
 // @param pod: The Pod to be checked.

ray-operator/controllers/ray/rayjob_controller.go

@@ -109,6 +109,18 @@ func (r *RayJobReconciler) Reconcile(ctx context.Context, request ctrl.Request)
 		return ctrl.Result{}, nil
 	}
 
+	if rayJobInstance.Spec.SubmissionMode == rayv1.SidecarMode &&
+		rayJobInstance.ObjectMeta.DeletionTimestamp.IsZero() &&
+		addSidecarDisableHeadRestartLabel(rayJobInstance) {
+		logger.Info("Add label to disable head Pod recreation for sidecar-mode RayJob",
+			"label", utils.RayJobDisableHeadNodeRestartLabelKey)
+		if err := r.Update(ctx, rayJobInstance); err != nil {
+			logger.Error(err, "Failed to update RayJob labels")
+			return ctrl.Result{RequeueAfter: RayJobDefaultRequeueDuration}, err
+		}
+		return ctrl.Result{RequeueAfter: RayJobDefaultRequeueDuration}, 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
@@ -924,6 +936,15 @@ func (r *RayJobReconciler) getOrCreateRayClusterInstance(ctx context.Context, ra
 	}
 	logger.Info("Found the associated RayCluster for RayJob", "RayCluster", rayClusterNamespacedName)
 
+	if rayJobInstance.Spec.SubmissionMode == rayv1.SidecarMode &&
+		len(rayJobInstance.Spec.ClusterSelector) == 0 &&
+		ensureRayClusterHasSidecarDisableHeadRestartLabel(rayClusterInstance) {
+		if err := r.Update(ctx, rayClusterInstance); err != nil {
+			logger.Error(err, "Failed to update RayCluster with sidecar restart label", "RayCluster", rayClusterNamespacedName)
+			return nil, err
+		}
+	}
+
 	// Verify that RayJob is not in cluster selector mode first to avoid nil pointer dereference error during spec comparison.
 	// This is checked by ensuring len(rayJobInstance.Spec.ClusterSelector) equals 0.
 	if len(rayJobInstance.Spec.ClusterSelector) == 0 && !utils.CompareJsonStruct(rayClusterInstance.Spec, *rayJobInstance.Spec.RayClusterSpec) {
@@ -934,12 +955,17 @@ func (r *RayJobReconciler) getOrCreateRayClusterInstance(ctx context.Context, ra
 }
 
 func (r *RayJobReconciler) constructRayClusterForRayJob(rayJobInstance *rayv1.RayJob, rayClusterName string) (*rayv1.RayCluster, error) {
-	labels := make(map[string]string, len(rayJobInstance.Labels))
+	labels := map[string]string{}
 	for key, value := range rayJobInstance.Labels {
 		labels[key] = value
 	}
 	labels[utils.RayOriginatedFromCRNameLabelKey] = rayJobInstance.Name
 	labels[utils.RayOriginatedFromCRDLabelKey] = utils.RayOriginatedFromCRDLabelValue(utils.RayJobCRD)
+	labels[utils.RayJobSubmissionModeLabelKey] = string(rayJobInstance.Spec.SubmissionMode)
+	if rayJobInstance.Spec.SubmissionMode == rayv1.SidecarMode {
+		labels[utils.RayJobSubmissionModeLabelKey] = string(rayv1.SidecarMode)
+		labels[utils.RayJobDisableHeadNodeRestartLabelKey] = "true"
+	}
 	rayCluster := &rayv1.RayCluster{
 		ObjectMeta: metav1.ObjectMeta{
 			Labels:      labels,
@@ -973,6 +999,28 @@ func (r *RayJobReconciler) constructRayClusterForRayJob(rayJobInstance *rayv1.Ra
 	return rayCluster, nil
 }
 
+func addSidecarDisableHeadRestartLabel(rayJob *rayv1.RayJob) bool {
+	if rayJob.Labels == nil {
+		rayJob.Labels = map[string]string{}
+	}
+	if _, exists := rayJob.Labels[utils.RayJobDisableHeadNodeRestartLabelKey]; exists {
+		return false
+	}
+	rayJob.Labels[utils.RayJobDisableHeadNodeRestartLabelKey] = "true"
+	return true
+}
+
+func ensureRayClusterHasSidecarDisableHeadRestartLabel(rayCluster *rayv1.RayCluster) bool {
+	if rayCluster.Labels == nil {
+		rayCluster.Labels = map[string]string{}
+	}
+	if _, exists := rayCluster.Labels[utils.RayJobDisableHeadNodeRestartLabelKey]; exists {
+		return false
+	}
+	rayCluster.Labels[utils.RayJobDisableHeadNodeRestartLabelKey] = "true"
+	return true
+}
+
 func updateStatusToSuspendingIfNeeded(ctx context.Context, rayJob *rayv1.RayJob) bool {
 	logger := ctrl.LoggerFrom(ctx)
 	if !rayJob.Spec.Suspend {
@@ -1020,7 +1068,7 @@ func (r *RayJobReconciler) checkSubmitterAndUpdateStatusIfNeeded(ctx context.Con
 			// If head pod is deleted, mark the RayJob as failed
 			shouldUpdate = true
 			rayJob.Status.JobDeploymentStatus = rayv1.JobDeploymentStatusFailed
-			rayJob.Status.Reason = rayv1.AppFailed
+			rayJob.Status.Reason = rayv1.SubmissionFailed
 			rayJob.Status.Message = "Ray head pod not found."
 			return
 		}

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

@@ -26,6 +26,9 @@ const (
 	HashWithoutReplicasAndWorkersToDeleteKey = "ray.io/hash-without-replicas-and-workers-to-delete"
 	NumWorkerGroupsKey                       = "ray.io/num-worker-groups"
 	KubeRayVersion                           = "ray.io/kuberay-version"
+	RayJobSubmissionModeLabelKey             = "ray.io/job-submission-mode"
+	// RayJobDisableHeadNodeRestartLabelKey marks RayClusters created for sidecar-mode RayJobs to skip head Pod recreation after provisioning.
+	RayJobDisableHeadNodeRestartLabelKey = "ray.io/disable-provisioned-head-restart"
 
 	// Labels for feature RayMultihostIndexing
 	//

ray-operator/test/e2erayjob/rayjob_sidecar_mode_test.go

@@ -2,6 +2,7 @@ package e2erayjob
 
 import (
 	"testing"
+	"time"
 
 	. "github.com/onsi/gomega"
 	corev1 "k8s.io/api/core/v1"
@@ -10,6 +11,7 @@ import (
 	corev1ac "k8s.io/client-go/applyconfigurations/core/v1"
 
 	rayv1 "github.com/ray-project/kuberay/ray-operator/apis/ray/v1"
+	"github.com/ray-project/kuberay/ray-operator/controllers/ray/common"
 	"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"
@@ -174,17 +176,40 @@ env_vars:
 		g.Eventually(RayJob(test, rayJob.Namespace, rayJob.Name), TestTimeoutMedium).
 			Should(WithTransform(RayJobStatus, Equal(rayv1.JobStatusRunning)))
 
+		g.Eventually(RayJob(test, rayJob.Namespace, rayJob.Name), TestTimeoutMedium).
+			Should(WithTransform(func(job *rayv1.RayJob) string {
+				return job.Labels[utils.RayJobDisableHeadNodeRestartLabelKey]
+			}, Equal("true")))
+
 		// Fetch RayCluster and delete the head Pod
 		rayJob, err = GetRayJob(test, rayJob.Namespace, rayJob.Name)
 		g.Expect(err).NotTo(HaveOccurred())
 		rayCluster, err := GetRayCluster(test, rayJob.Namespace, rayJob.Status.RayClusterName)
 		g.Expect(err).NotTo(HaveOccurred())
+		g.Expect(rayCluster.Labels[utils.RayJobDisableHeadNodeRestartLabelKey]).To(Equal("true"))
 		headPod, err := GetHeadPod(test, rayCluster)
 		g.Expect(err).NotTo(HaveOccurred())
 		LogWithTimestamp(test.T(), "Deleting head Pod %s/%s for RayCluster %s", headPod.Namespace, headPod.Name, rayCluster.Name)
 		err = test.Client().Core().CoreV1().Pods(headPod.Namespace).Delete(test.Ctx(), headPod.Name, metav1.DeleteOptions{})
 		g.Expect(err).NotTo(HaveOccurred())
 
+		g.Eventually(func() int {
+			pods, listErr := test.Client().Core().CoreV1().Pods(rayCluster.Namespace).List(
+				test.Ctx(), common.RayClusterHeadPodsAssociationOptions(rayCluster).ToMetaV1ListOptions())
+			if listErr != nil {
+				return -1
+			}
+			return len(pods.Items)
+		}, TestTimeoutMedium, 2*time.Second).Should(Equal(0))
+		g.Consistently(func() int {
+			pods, listErr := test.Client().Core().CoreV1().Pods(rayCluster.Namespace).List(
+				test.Ctx(), common.RayClusterHeadPodsAssociationOptions(rayCluster).ToMetaV1ListOptions())
+			if listErr != nil {
+				return -1
+			}
+			return len(pods.Items)
+		}, TestTimeoutShort, 2*time.Second).Should(Equal(0))
+
 		// After head pod deletion, controller should mark RayJob as Failed with a specific message
 		g.Eventually(RayJob(test, rayJob.Namespace, rayJob.Name), TestTimeoutMedium).
 			Should(WithTransform(RayJobDeploymentStatus, Equal(rayv1.JobDeploymentStatusFailed)))

ray-operator/test/e2erayjob/rayjob_test.go

@@ -15,6 +15,7 @@ import (
 
 	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/common"
 	"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"
@@ -279,7 +280,7 @@ env_vars:
 			To(WithTransform(RayJobReason, Equal(rayv1.DeadlineExceeded)))
 	})
 
-	test.T().Run("RayJob fails when head Pod is deleted when job is running", func(_ *testing.T) {
+	test.T().Run("RayJob recreates head Pod when deleted while running", func(_ *testing.T) {
 		rayJobAC := rayv1ac.RayJob("delete-head-after-submit", namespace.Name).
 			WithSpec(rayv1ac.RayJobSpec().
 				WithRayClusterSpec(NewRayClusterSpec()).
@@ -289,6 +290,7 @@ env_vars:
 		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)
+		g.Expect(rayJob.Labels).To(Not(HaveKey(utils.RayJobDisableHeadNodeRestartLabelKey)))
 
 		// Wait until the RayJob's job status transitions to Running
 		LogWithTimestamp(test.T(), "Waiting for RayJob %s/%s to be 'Running'", rayJob.Namespace, rayJob.Name)
@@ -300,20 +302,28 @@ env_vars:
 		g.Expect(err).NotTo(HaveOccurred())
 		rayCluster, err := GetRayCluster(test, rayJob.Namespace, rayJob.Status.RayClusterName)
 		g.Expect(err).NotTo(HaveOccurred())
+		g.Expect(rayCluster.Labels).To(Not(HaveKeyWithValue(utils.RayJobDisableHeadNodeRestartLabelKey, "true")))
 		headPod, err := GetHeadPod(test, rayCluster)
 		g.Expect(err).NotTo(HaveOccurred())
 		LogWithTimestamp(test.T(), "Deleting head Pod %s/%s for RayCluster %s", headPod.Namespace, headPod.Name, rayCluster.Name)
 		err = test.Client().Core().CoreV1().Pods(headPod.Namespace).Delete(test.Ctx(), headPod.Name, metav1.DeleteOptions{})
 		g.Expect(err).NotTo(HaveOccurred())
 
-		// After head pod deletion, controller should mark RayJob as Failed with a specific message
-		g.Eventually(RayJob(test, rayJob.Namespace, rayJob.Name), TestTimeoutMedium).
-			Should(WithTransform(RayJobDeploymentStatus, Equal(rayv1.JobDeploymentStatusFailed)))
+		// Head pod should be recreated for non-sidecar modes and the RayJob should keep running/finish.
+		g.Eventually(func() int {
+			pods, listErr := test.Client().Core().CoreV1().Pods(rayCluster.Namespace).List(
+				test.Ctx(), common.RayClusterHeadPodsAssociationOptions(rayCluster).ToMetaV1ListOptions())
+			if listErr != nil {
+				return -1
+			}
+			return len(pods.Items)
+		}, TestTimeoutMedium, 2*time.Second).Should(Equal(1))
+		g.Consistently(RayJob(test, rayJob.Namespace, rayJob.Name), TestTimeoutShort).
+			ShouldNot(WithTransform(RayJobDeploymentStatus, Equal(rayv1.JobDeploymentStatusFailed)))
 		g.Eventually(RayJob(test, rayJob.Namespace, rayJob.Name), TestTimeoutMedium).
-			Should(WithTransform(RayJobReason, Or(
-				Equal(rayv1.JobDeploymentStatusTransitionGracePeriodExceeded),
-				Equal(rayv1.SubmissionFailed),
-			)))
+			Should(WithTransform(func(job *rayv1.RayJob) rayv1.JobDeploymentStatus {
+				return job.Status.JobDeploymentStatus
+			}, Or(Equal(rayv1.JobDeploymentStatusRunning), Equal(rayv1.JobDeploymentStatusComplete))))
 		// Cleanup
 		err = test.Client().Ray().RayV1().RayJobs(namespace.Name).Delete(test.Ctx(), rayJob.Name, metav1.DeleteOptions{})
 		g.Expect(err).NotTo(HaveOccurred())