add ray.io/kuberay-version annotations for head pod and worker pods

Signed-off-by: win5923 <ken89@kimo.com>

Author: win5923

ray-operator/controllers/ray/common/pod.go

@@ -188,10 +188,12 @@ func DefaultHeadPodTemplate(ctx context.Context, instance rayv1.RayCluster, head
 	// This ensures privilege of KubeRay users are contained within the namespace of the RayCluster.
 	podTemplate.ObjectMeta.Namespace = instance.Namespace
 
+	if podTemplate.Annotations == nil {
+		podTemplate.Annotations = make(map[string]string)
+	}
+	// Set the KubeRay version used to create the pod
+	podTemplate.Annotations[utils.KubeRayVersion] = utils.KUBERAY_VERSION
 	if templateHash != "" {
-		if podTemplate.Annotations == nil {
-			podTemplate.Annotations = make(map[string]string)
-		}
 		podTemplate.Annotations[utils.PodTemplateHashKey] = templateHash
 	}
 
@@ -335,15 +337,16 @@ func DefaultWorkerPodTemplate(ctx context.Context, instance rayv1.RayCluster, wo
 
 	podTemplate := workerSpec.Template
 	podTemplate.GenerateName = podName
-
 	// Pods created by RayCluster should be restricted to the namespace of the RayCluster.
 	// This ensures privilege of KubeRay users are contained within the namespace of the RayCluster.
 	podTemplate.ObjectMeta.Namespace = instance.Namespace
 
+	if podTemplate.Annotations == nil {
+		podTemplate.Annotations = make(map[string]string)
+	}
+	// Set the KubeRay version used to create the pod
+	podTemplate.Annotations[utils.KubeRayVersion] = utils.KUBERAY_VERSION
 	if templateHash != "" {
-		if podTemplate.Annotations == nil {
-			podTemplate.Annotations = make(map[string]string)
-		}
 		podTemplate.Annotations[utils.PodTemplateHashKey] = templateHash
 	}
 	// The Ray worker should only start once the GCS server is ready.

ray-operator/controllers/ray/raycluster_controller.go

@@ -1137,6 +1137,21 @@ func (r *RayClusterReconciler) shouldRecreatePodsForUpgrade(ctx context.Context,
 		return false
 	}
 
+	// Case 1: If the KubeRay version has changed, update annotations then check in the next reconciliation
+	for _, pod := range allPods.Items {
+		podVersion := pod.Annotations[utils.KubeRayVersion]
+		if podVersion != "" && podVersion != utils.KUBERAY_VERSION {
+			logger.Info("Pods have different KubeRay version, updating pod annotations",
+				"pod", pod.Name,
+				"podVersion", podVersion,
+				"currentVersion", utils.KUBERAY_VERSION)
+			if err := r.updatePodsAnnotations(ctx, instance, &allPods); err != nil {
+				logger.Error(err, "Failed to update pod annotations for KubeRay version change")
+			}
+			return false
+		}
+	}
+
 	headHash, err := common.GeneratePodTemplateHash(instance.Spec.HeadGroupSpec.Template)
 	if err != nil {
 		logger.Error(err, "Failed to generate head template hash")
@@ -1153,7 +1168,7 @@ func (r *RayClusterReconciler) shouldRecreatePodsForUpgrade(ctx context.Context,
 		workerHashMap[workerGroup.GroupName] = hash
 	}
 
-	// Check each pod to see if its template hash matches the current spec
+	// Case 2: If the pod template hash has changed, recreate all pods
 	for _, pod := range allPods.Items {
 		nodeType := pod.Labels[utils.RayNodeTypeLabelKey]
 		actualHash := pod.Annotations[utils.PodTemplateHashKey]
@@ -1183,6 +1198,61 @@ func (r *RayClusterReconciler) shouldRecreatePodsForUpgrade(ctx context.Context,
 	return false
 }
 
+// updatePodsAnnotations updates pod annotations to match the current KubeRay version and PodTemplateHashKey
+func (r *RayClusterReconciler) updatePodsAnnotations(ctx context.Context, instance *rayv1.RayCluster, allPods *corev1.PodList) error {
+	logger := ctrl.LoggerFrom(ctx)
+
+	for i := range allPods.Items {
+		pod := &allPods.Items[i]
+		podVersion := pod.Annotations[utils.KubeRayVersion]
+
+		if podVersion == utils.KUBERAY_VERSION || podVersion == "" {
+			continue
+		}
+
+		newHash, err := r.calculatePodTemplateHash(instance, pod)
+		if err != nil {
+			return err
+		}
+
+		if pod.Annotations == nil {
+			pod.Annotations = make(map[string]string)
+		}
+		pod.Annotations[utils.KubeRayVersion] = utils.KUBERAY_VERSION
+		pod.Annotations[utils.PodTemplateHashKey] = newHash
+
+		if err := r.Update(ctx, pod); err != nil {
+			return err
+		}
+
+		logger.Info("Updated pod annotations", "pod", pod.Name, "version", utils.KUBERAY_VERSION)
+	}
+
+	return nil
+}
+
+// calculatePodTemplateHash calculates the hash for a pod's template based on its node type and group
+func (r *RayClusterReconciler) calculatePodTemplateHash(instance *rayv1.RayCluster, pod *corev1.Pod) (string, error) {
+	nodeType := pod.Labels[utils.RayNodeTypeLabelKey]
+
+	switch rayv1.RayNodeType(nodeType) {
+	case rayv1.HeadNode:
+		return common.GeneratePodTemplateHash(instance.Spec.HeadGroupSpec.Template)
+
+	case rayv1.WorkerNode:
+		groupName := pod.Labels[utils.RayNodeGroupLabelKey]
+		for _, workerGroup := range instance.Spec.WorkerGroupSpecs {
+			if workerGroup.GroupName == groupName {
+				return common.GeneratePodTemplateHash(workerGroup.Template)
+			}
+		}
+		return "", fmt.Errorf("worker group %s not found in RayCluster spec", groupName)
+
+	default:
+		return "", fmt.Errorf("unknown node type: %s", nodeType)
+	}
+}
+
 // shouldDeletePod returns whether the Pod should be deleted and the reason
 //
 // @param pod: The Pod to be checked.

ray-operator/controllers/ray/raycluster_controller_unit_test.go

@@ -3668,6 +3668,7 @@ func TestShouldRecreatePodsForUpgrade(t *testing.T) {
 				},
 				Annotations: map[string]string{
 					utils.PodTemplateHashKey: templateHash,
+					utils.KubeRayVersion:     utils.KUBERAY_VERSION,
 				},
 			},
 			Spec: corev1.PodSpec{
@@ -3679,6 +3680,13 @@ func TestShouldRecreatePodsForUpgrade(t *testing.T) {
 		}
 	}
 
+	// Helper function to create a pod with specific template hash and KubeRay version
+	createPodWithHashAndVersion := func(name string, nodeType rayv1.RayNodeType, groupName string, templateHash string, kuberayVersion string) *corev1.Pod {
+		pod := createPodWithHash(name, nodeType, groupName, templateHash)
+		pod.Annotations[utils.KubeRayVersion] = kuberayVersion
+		return pod
+	}
+
 	tests := []struct {
 		name             string
 		upgradeStrategy  *rayv1.RayClusterUpgradeStrategy
@@ -3746,6 +3754,17 @@ func TestShouldRecreatePodsForUpgrade(t *testing.T) {
 			},
 			expectedRecreate: true,
 		},
+		{
+			name: "Recreate strategy with different KubeRay version - should update annotations and not recreate",
+			upgradeStrategy: &rayv1.RayClusterUpgradeStrategy{
+				Type: ptr.To(rayv1.RayClusterRecreate),
+			},
+			pods: []runtime.Object{
+				createPodWithHashAndVersion("head-pod", rayv1.HeadNode, headGroupNameStr, "old-head-hash", "v1.0.0"),
+				createPodWithHashAndVersion("worker-pod", rayv1.WorkerNode, groupNameStr, "old-worker-hash", "v1.0.0"),
+			},
+			expectedRecreate: false,
+		},
 	}
 
 	for _, tc := range tests {