Add test

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

Author: win5923

ray-operator/controllers/ray/raycluster_controller_unit_test.go

@@ -3643,3 +3643,125 @@ func TestReconcile_PodsWithAuthToken(t *testing.T) {
 		assert.True(t, authModeEnvFound, "Auth mode env vars not found")
 	}
 }
+
+func TestShouldRecreatePodsForUpgrade(t *testing.T) {
+	setupTest(t)
+	ctx := context.Background()
+
+	// Calculate template hashes for matching pods
+	headHash, err := common.GeneratePodTemplateHash(testRayCluster.Spec.HeadGroupSpec.Template)
+	require.NoError(t, err, "Failed to generate head template hash")
+	workerHash, err := common.GeneratePodTemplateHash(testRayCluster.Spec.WorkerGroupSpecs[0].Template)
+	require.NoError(t, err, "Failed to generate worker template hash")
+
+	// Helper function to create a pod with specific template hash
+	createPodWithHash := func(name string, nodeType rayv1.RayNodeType, groupName string, templateHash string) *corev1.Pod {
+		return &corev1.Pod{
+			ObjectMeta: metav1.ObjectMeta{
+				Name:      name,
+				Namespace: namespaceStr,
+				Labels: map[string]string{
+					utils.RayNodeLabelKey:      "yes",
+					utils.RayClusterLabelKey:   instanceName,
+					utils.RayNodeTypeLabelKey:  string(nodeType),
+					utils.RayNodeGroupLabelKey: groupName,
+				},
+				Annotations: map[string]string{
+					utils.PodTemplateHashKey: templateHash,
+				},
+			},
+			Spec: corev1.PodSpec{
+				Containers: []corev1.Container{
+					{Name: "ray-head", Image: "rayproject/ray:latest"},
+				},
+			},
+			Status: corev1.PodStatus{Phase: corev1.PodRunning},
+		}
+	}
+
+	tests := []struct {
+		name             string
+		upgradeStrategy  *rayv1.RayClusterUpgradeStrategy
+		pods             []runtime.Object
+		expectedRecreate bool
+	}{
+		{
+			name:             "No upgrade strategy",
+			upgradeStrategy:  nil,
+			pods:             testPods,
+			expectedRecreate: false,
+		},
+		{
+			name:             "Upgrade strategy type is nil",
+			upgradeStrategy:  &rayv1.RayClusterUpgradeStrategy{Type: nil},
+			pods:             testPods,
+			expectedRecreate: false,
+		},
+		{
+			name: "Upgrade strategy type is None",
+			upgradeStrategy: &rayv1.RayClusterUpgradeStrategy{
+				Type: ptr.To(rayv1.RayClusterUpgradeNone),
+			},
+			pods:             testPods,
+			expectedRecreate: false,
+		},
+		{
+			name: "Recreate strategy but no pods exist",
+			upgradeStrategy: &rayv1.RayClusterUpgradeStrategy{
+				Type: ptr.To(rayv1.Recreate),
+			},
+			pods:             []runtime.Object{},
+			expectedRecreate: false,
+		},
+		{
+			name: "Recreate strategy with matching template hash",
+			upgradeStrategy: &rayv1.RayClusterUpgradeStrategy{
+				Type: ptr.To(rayv1.Recreate),
+			},
+			pods: []runtime.Object{
+				createPodWithHash("head-pod", rayv1.HeadNode, headGroupNameStr, headHash),
+				createPodWithHash("worker-pod", rayv1.WorkerNode, groupNameStr, workerHash),
+			},
+			expectedRecreate: false,
+		},
+		{
+			name: "Recreate strategy with mismatched head template hash",
+			upgradeStrategy: &rayv1.RayClusterUpgradeStrategy{
+				Type: ptr.To(rayv1.Recreate),
+			},
+			pods: []runtime.Object{
+				createPodWithHash("head-pod", rayv1.HeadNode, headGroupNameStr, "old-head-hash"),
+				createPodWithHash("worker-pod", rayv1.WorkerNode, groupNameStr, workerHash),
+			},
+			expectedRecreate: true,
+		},
+		{
+			name: "Recreate strategy with mismatched worker template hash",
+			upgradeStrategy: &rayv1.RayClusterUpgradeStrategy{
+				Type: ptr.To(rayv1.Recreate),
+			},
+			pods: []runtime.Object{
+				createPodWithHash("head-pod", rayv1.HeadNode, headGroupNameStr, headHash),
+				createPodWithHash("worker-pod", rayv1.WorkerNode, groupNameStr, "old-worker-hash"),
+			},
+			expectedRecreate: true,
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			cluster := testRayCluster.DeepCopy()
+			cluster.Spec.UpgradeStrategy = tc.upgradeStrategy
+
+			fakeClient := clientFake.NewClientBuilder().WithRuntimeObjects(tc.pods...).Build()
+			testRayClusterReconciler := &RayClusterReconciler{
+				Client:   fakeClient,
+				Scheme:   scheme.Scheme,
+				Recorder: &record.FakeRecorder{},
+			}
+
+			result := testRayClusterReconciler.shouldRecreatePodsForUpgrade(ctx, cluster)
+			assert.Equal(t, tc.expectedRecreate, result)
+		})
+	}
+}

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

@@ -38,11 +38,14 @@ func ValidateRayClusterMetadata(metadata metav1.ObjectMeta) error {
 }
 
 func ValidateRayClusterUpgradeOptions(instance *rayv1.RayCluster) error {
-	if instance.Spec.UpgradeStrategy != nil {
-		creatorCRDType := GetCRDType(instance.Labels[RayOriginatedFromCRDLabelKey])
-		if creatorCRDType == RayJobCRD || creatorCRDType == RayServiceCRD {
-			return fmt.Errorf("upgradeStrategy cannot be set when RayCluster is created by %s", creatorCRDType)
-		}
+	strategy := instance.Spec.UpgradeStrategy
+	if strategy == nil || strategy.Type == nil || *strategy.Type == rayv1.RayClusterUpgradeNone {
+		return nil
+	}
+
+	creatorCRDType := GetCRDType(instance.Labels[RayOriginatedFromCRDLabelKey])
+	if creatorCRDType == RayJobCRD || creatorCRDType == RayServiceCRD {
+		return fmt.Errorf("upgradeStrategy cannot be set when RayCluster is created by %s", creatorCRDType)
 	}
 	return nil
 }

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

@@ -1887,3 +1887,92 @@ func TestValidateClusterUpgradeOptions(t *testing.T) {
 		})
 	}
 }
+
+func TestValidateRayClusterUpgradeOptions(t *testing.T) {
+	tests := []struct {
+		upgradeStrategy   *rayv1.RayClusterUpgradeStrategy
+		name              string
+		originatedFromCRD string
+		errorMessage      string
+		expectError       bool
+	}{
+		{
+			name:            "No upgrade strategy",
+			upgradeStrategy: nil,
+			expectError:     false,
+		},
+		{
+			name: "Upgrade strategy set None for RayCluster",
+			upgradeStrategy: &rayv1.RayClusterUpgradeStrategy{
+				Type: ptr.To(rayv1.RayClusterUpgradeNone),
+			},
+			originatedFromCRD: string(RayClusterCRD),
+			expectError:       false,
+		},
+		{
+			name: "Upgrade strategy set Recreate for RayCluster",
+			upgradeStrategy: &rayv1.RayClusterUpgradeStrategy{
+				Type: ptr.To(rayv1.Recreate),
+			},
+			originatedFromCRD: string(RayClusterCRD),
+			expectError:       false,
+		},
+		{
+			name: "Upgrade strategy set Recreate for RayCluster created by RayJob",
+			upgradeStrategy: &rayv1.RayClusterUpgradeStrategy{
+				Type: ptr.To(rayv1.Recreate),
+			},
+			originatedFromCRD: string(RayJobCRD),
+			expectError:       true,
+			errorMessage:      "upgradeStrategy cannot be set when RayCluster is created by RayJob",
+		},
+		{
+			name: "Upgrade strategy set Recreate for RayCluster created by RayService",
+			upgradeStrategy: &rayv1.RayClusterUpgradeStrategy{
+				Type: ptr.To(rayv1.Recreate),
+			},
+			originatedFromCRD: string(RayServiceCRD),
+			expectError:       true,
+			errorMessage:      "upgradeStrategy cannot be set when RayCluster is created by RayService",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			cluster := &rayv1.RayCluster{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      "test-cluster",
+					Namespace: "default",
+					Labels:    map[string]string{},
+				},
+				Spec: rayv1.RayClusterSpec{
+					UpgradeStrategy: tt.upgradeStrategy,
+					HeadGroupSpec: rayv1.HeadGroupSpec{
+						Template: corev1.PodTemplateSpec{
+							Spec: corev1.PodSpec{
+								Containers: []corev1.Container{
+									{Name: "ray-head", Image: "rayproject/ray:latest"},
+								},
+							},
+						},
+					},
+				},
+			}
+
+			if tt.originatedFromCRD != "" {
+				cluster.Labels[RayOriginatedFromCRDLabelKey] = tt.originatedFromCRD
+			}
+
+			err := ValidateRayClusterUpgradeOptions(cluster)
+
+			if tt.expectError {
+				require.Error(t, err, "Expected error for test case: %s", tt.name)
+				if tt.errorMessage != "" {
+					assert.Contains(t, err.Error(), tt.errorMessage, "Error message mismatch for test case: %s", tt.name)
+				}
+			} else {
+				require.NoError(t, err, "Unexpected error for test case: %s", tt.name)
+			}
+		})
+	}
+}

ray-operator/test/e2e/raycluster_test.go

@@ -9,6 +9,7 @@ import (
 	"k8s.io/apimachinery/pkg/api/errors"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/types"
+	corev1ac "k8s.io/client-go/applyconfigurations/core/v1"
 	"k8s.io/utils/ptr"
 
 	rayv1 "github.com/ray-project/kuberay/ray-operator/apis/ray/v1"
@@ -208,3 +209,91 @@ func TestRayClusterScalingDown(t *testing.T) {
 		g.Expect(err).NotTo(HaveOccurred(), "Failed to remove finalizer from pod %s/%s", namespace.Name, pod.Name)
 	}
 }
+
+func TestRayClusterUpgradeStrategy(t *testing.T) {
+	test := With(t)
+	g := NewWithT(t)
+
+	// Create a namespace
+	namespace := test.NewTestNamespace()
+
+	// Create RayCluster with Recreate upgrade strategy
+	rayClusterAC := rayv1ac.RayCluster("raycluster-upgrade", namespace.Name).
+		WithSpec(rayv1ac.RayClusterSpec().
+			WithRayVersion(GetRayVersion()).
+			WithUpgradeStrategy(rayv1ac.RayClusterUpgradeStrategy().
+				WithType(rayv1.Recreate)).
+			WithHeadGroupSpec(rayv1ac.HeadGroupSpec().
+				WithRayStartParams(map[string]string{"dashboard-host": "0.0.0.0"}).
+				WithTemplate(HeadPodTemplateApplyConfiguration())).
+			WithWorkerGroupSpecs(rayv1ac.WorkerGroupSpec().
+				WithReplicas(2).
+				WithMinReplicas(1).
+				WithMaxReplicas(5).
+				WithGroupName("small-group").
+				WithRayStartParams(map[string]string{"num-cpus": "1"}).
+				WithTemplate(WorkerPodTemplateApplyConfiguration())))
+
+	rayCluster, err := test.Client().Ray().RayV1().RayClusters(namespace.Name).Apply(test.Ctx(), rayClusterAC, TestApplyOptions)
+	g.Expect(err).NotTo(HaveOccurred())
+	LogWithTimestamp(test.T(), "Created RayCluster %s/%s with Recreate upgrade strategy", rayCluster.Namespace, rayCluster.Name)
+
+	LogWithTimestamp(test.T(), "Waiting for RayCluster %s/%s to become ready", rayCluster.Namespace, rayCluster.Name)
+	g.Eventually(RayCluster(test, rayCluster.Namespace, rayCluster.Name), TestTimeoutMedium).
+		Should(WithTransform(RayClusterState, Equal(rayv1.Ready)))
+
+	// Get initial pods
+	headPod, err := GetHeadPod(test, rayCluster)
+	g.Expect(err).NotTo(HaveOccurred())
+	initialHeadPodName := headPod.Name
+	LogWithTimestamp(test.T(), "Initial head pod name: %s", initialHeadPodName)
+
+	workerPods, err := GetWorkerPods(test, rayCluster)
+	g.Expect(err).NotTo(HaveOccurred())
+	g.Expect(workerPods).To(HaveLen(2), "Should have 2 worker pods initially")
+
+	LogWithTimestamp(test.T(), "Updating RayCluster %s/%s pod template to trigger upgrade", rayCluster.Namespace, rayCluster.Name)
+	// Update head pod template with a new environment variable to trigger recreation
+	rayClusterAC.Spec.HeadGroupSpec.Template.Spec.Containers[0].
+		WithEnv(corev1ac.EnvVar().
+			WithName("TEST_UPGRADE_ENV").
+			WithValue("upgraded"))
+
+	rayCluster, err = test.Client().Ray().RayV1().RayClusters(namespace.Name).Apply(test.Ctx(), rayClusterAC, TestApplyOptions)
+	g.Expect(err).NotTo(HaveOccurred())
+	LogWithTimestamp(test.T(), "Applied updated RayCluster spec")
+
+	LogWithTimestamp(test.T(), "Waiting for all pods to be deleted due to Recreate strategy")
+	g.Eventually(func(gg Gomega) {
+		pods, err := GetAllPods(test, rayCluster)
+		gg.Expect(err).NotTo(HaveOccurred())
+		gg.Expect(pods).To(BeEmpty(), "All pods should be deleted")
+	}, TestTimeoutMedium).Should(Succeed())
+
+	LogWithTimestamp(test.T(), "Waiting for RayCluster %s/%s to become ready again", rayCluster.Namespace, rayCluster.Name)
+	g.Eventually(RayCluster(test, rayCluster.Namespace, rayCluster.Name), TestTimeoutMedium).
+		Should(WithTransform(RayClusterState, Equal(rayv1.Ready)))
+
+	// Verify new pods are created
+	newHeadPod, err := GetHeadPod(test, rayCluster)
+	g.Expect(err).NotTo(HaveOccurred())
+	g.Expect(newHeadPod.Name).NotTo(Equal(initialHeadPodName), "Head pod should be recreated with a different name")
+	LogWithTimestamp(test.T(), "New head pod name: %s", newHeadPod.Name)
+
+	// Verify the new environment variable is applied
+	found := false
+	for _, container := range newHeadPod.Spec.Containers {
+		for _, env := range container.Env {
+			if env.Name == "TEST_UPGRADE_ENV" && env.Value == "upgraded" {
+				found = true
+				break
+			}
+		}
+	}
+	g.Expect(found).To(BeTrue(), "New head pod should have the updated environment variable")
+
+	newWorkerPods, err := GetWorkerPods(test, rayCluster)
+	g.Expect(err).NotTo(HaveOccurred())
+	g.Expect(newWorkerPods).To(HaveLen(2), "Should have 2 worker pods after upgrade")
+	LogWithTimestamp(test.T(), "Successfully verified RayCluster upgrade with Recreate strategy")
+}