[RayCluster][Fix] DesiredReplicas, MinReplicas and MaxReplicas should respect workerGroupSpec.Suspend (#2728)

Author: rueian

ray-operator/controllers/ray/raycluster_controller_unit_test.go

@@ -1864,6 +1864,73 @@ func TestCalculateStatusWithoutDesiredReplicas(t *testing.T) {
 	assert.Nil(t, newInstance.Status.StateTransitionTimes)
 }
 
+// TestCalculateStatusWithSuspendedWorkerGroups tests that the cluster CR should be marked as Ready without workers
+// and all desired resources are not counted with suspended workers
+func TestCalculateStatusWithSuspendedWorkerGroups(t *testing.T) {
+	setupTest(t)
+
+	testRayCluster.Spec.WorkerGroupSpecs[0].Suspend = ptr.To[bool](true)
+	testRayCluster.Spec.WorkerGroupSpecs[0].MinReplicas = ptr.To[int32](100)
+	testRayCluster.Spec.WorkerGroupSpecs[0].MaxReplicas = ptr.To[int32](100)
+	testRayCluster.Spec.WorkerGroupSpecs[0].Template.Spec.Containers[0].Resources.Requests = corev1.ResourceList{
+		corev1.ResourceCPU:    resource.MustParse("100m"),
+		corev1.ResourceMemory: resource.MustParse("100Mi"),
+	}
+
+	// Create a new scheme with CRDs, Pod, Service schemes.
+	newScheme := runtime.NewScheme()
+	_ = rayv1.AddToScheme(newScheme)
+	_ = corev1.AddToScheme(newScheme)
+
+	// Mock data
+	headServiceIP := "aaa.bbb.ccc.ddd"
+	headService, err := common.BuildServiceForHeadPod(context.Background(), *testRayCluster, nil, nil)
+	assert.Nil(t, err, "Failed to build head service.")
+	headService.Spec.ClusterIP = headServiceIP
+	headPod := &corev1.Pod{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      "headNode",
+			Namespace: namespaceStr,
+			Labels: map[string]string{
+				utils.RayClusterLabelKey:  instanceName,
+				utils.RayNodeTypeLabelKey: string(rayv1.HeadNode),
+			},
+		},
+		Status: corev1.PodStatus{
+			PodIP: headNodeIP,
+			Phase: corev1.PodRunning,
+			Conditions: []corev1.PodCondition{
+				{
+					Type:   corev1.PodReady,
+					Status: corev1.ConditionTrue,
+				},
+			},
+		},
+	}
+	runtimeObjects := []runtime.Object{headPod, headService}
+
+	// Initialize a fake client with newScheme and runtimeObjects.
+	fakeClient := clientFake.NewClientBuilder().WithScheme(newScheme).WithRuntimeObjects(runtimeObjects...).Build()
+	ctx := context.Background()
+
+	// Initialize a RayCluster reconciler.
+	r := &RayClusterReconciler{
+		Client:   fakeClient,
+		Recorder: &record.FakeRecorder{},
+		Scheme:   scheme.Scheme,
+	}
+
+	newInstance, err := r.calculateStatus(ctx, testRayCluster, nil)
+	assert.Nil(t, err)
+	assert.Equal(t, newInstance.Status.DesiredWorkerReplicas, int32(0))
+	assert.Equal(t, newInstance.Status.MinWorkerReplicas, int32(0))
+	assert.Equal(t, newInstance.Status.MaxWorkerReplicas, int32(0))
+	assert.Equal(t, newInstance.Status.DesiredCPU, resource.Quantity{})
+	assert.Equal(t, newInstance.Status.DesiredMemory, resource.Quantity{})
+	assert.Equal(t, newInstance.Status.State, rayv1.Ready) //nolint:staticcheck // https://github.com/ray-project/kuberay/pull/2288
+	assert.NotNil(t, newInstance.Status.StateTransitionTimes)
+}
+
 // TestCalculateStatusWithReconcileErrorBackAndForth tests that the cluster CR should not be marked as Ready if reconcileErr != nil
 // and the Ready state should not be removed after being Ready even if reconcileErr != nil
 func TestCalculateStatusWithReconcileErrorBackAndForth(t *testing.T) {

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

@@ -322,6 +322,9 @@ func GetWorkerGroupDesiredReplicas(ctx context.Context, workerGroupSpec rayv1.Wo
 	log := ctrl.LoggerFrom(ctx)
 	// Always adhere to min/max replicas constraints.
 	var workerReplicas int32
+	if workerGroupSpec.Suspend != nil && *workerGroupSpec.Suspend {
+		return 0
+	}
 	if *workerGroupSpec.MinReplicas > *workerGroupSpec.MaxReplicas {
 		log.Info("minReplicas is greater than maxReplicas, using maxReplicas as desired replicas. "+
 			"Please fix this to avoid any unexpected behaviors.", "minReplicas", *workerGroupSpec.MinReplicas, "maxReplicas", *workerGroupSpec.MaxReplicas)
@@ -352,6 +355,9 @@ func CalculateDesiredReplicas(ctx context.Context, cluster *rayv1.RayCluster) in
 func CalculateMinReplicas(cluster *rayv1.RayCluster) int32 {
 	count := int32(0)
 	for _, nodeGroup := range cluster.Spec.WorkerGroupSpecs {
+		if nodeGroup.Suspend != nil && *nodeGroup.Suspend {
+			continue
+		}
 		count += *nodeGroup.MinReplicas
 	}
 
@@ -362,6 +368,9 @@ func CalculateMinReplicas(cluster *rayv1.RayCluster) int32 {
 func CalculateMaxReplicas(cluster *rayv1.RayCluster) int32 {
 	count := int32(0)
 	for _, nodeGroup := range cluster.Spec.WorkerGroupSpecs {
+		if nodeGroup.Suspend != nil && *nodeGroup.Suspend {
+			continue
+		}
 		count += *nodeGroup.MaxReplicas
 	}
 
@@ -405,6 +414,9 @@ func CalculateDesiredResources(cluster *rayv1.RayCluster) corev1.ResourceList {
 	headPodResource := CalculatePodResource(cluster.Spec.HeadGroupSpec.Template.Spec)
 	desiredResourcesList = append(desiredResourcesList, headPodResource)
 	for _, nodeGroup := range cluster.Spec.WorkerGroupSpecs {
+		if nodeGroup.Suspend != nil && *nodeGroup.Suspend {
+			continue
+		}
 		podResource := CalculatePodResource(nodeGroup.Template.Spec)
 		for i := int32(0); i < *nodeGroup.Replicas; i++ {
 			desiredResourcesList = append(desiredResourcesList, podResource)

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

@@ -513,6 +513,57 @@ func TestGetWorkerGroupDesiredReplicas(t *testing.T) {
 	workerGroupSpec.MinReplicas = &maxReplicas
 	workerGroupSpec.MaxReplicas = &minReplicas
 	assert.Equal(t, GetWorkerGroupDesiredReplicas(ctx, workerGroupSpec), *workerGroupSpec.MaxReplicas)
+
+	// Test 6: `WorkerGroupSpec.Suspend` is true.
+	suspend := true
+	workerGroupSpec.MinReplicas = &maxReplicas
+	workerGroupSpec.MaxReplicas = &minReplicas
+	workerGroupSpec.Suspend = &suspend
+	assert.Equal(t, GetWorkerGroupDesiredReplicas(ctx, workerGroupSpec), int32(0))
+}
+
+func TestCalculateMinReplicas(t *testing.T) {
+	// Test 1
+	minReplicas := int32(1)
+	rayCluster := &rayv1.RayCluster{
+		Spec: rayv1.RayClusterSpec{
+			WorkerGroupSpecs: []rayv1.WorkerGroupSpec{
+				{
+					MinReplicas: &minReplicas,
+				},
+			},
+		},
+	}
+	assert.Equal(t, CalculateMinReplicas(rayCluster), minReplicas)
+
+	// Test 2
+	suspend := true
+	for i := range rayCluster.Spec.WorkerGroupSpecs {
+		rayCluster.Spec.WorkerGroupSpecs[i].Suspend = &suspend
+	}
+	assert.Equal(t, CalculateMinReplicas(rayCluster), int32(0))
+}
+
+func TestCalculateMaxReplicas(t *testing.T) {
+	// Test 1
+	maxReplicas := int32(1)
+	rayCluster := &rayv1.RayCluster{
+		Spec: rayv1.RayClusterSpec{
+			WorkerGroupSpecs: []rayv1.WorkerGroupSpec{
+				{
+					MaxReplicas: &maxReplicas,
+				},
+			},
+		},
+	}
+	assert.Equal(t, CalculateMaxReplicas(rayCluster), maxReplicas)
+
+	// Test 2
+	suspend := true
+	for i := range rayCluster.Spec.WorkerGroupSpecs {
+		rayCluster.Spec.WorkerGroupSpecs[i].Suspend = &suspend
+	}
+	assert.Equal(t, CalculateMaxReplicas(rayCluster), int32(0))
 }
 
 func TestCalculateDesiredReplicas(t *testing.T) {