[fix][kubectl-plugin] set worker group CPU limit

kubectl-plugin/pkg/cmd/create/create_workergroup.go

@@ -20,6 +20,10 @@ import (
 	"k8s.io/kubectl/pkg/util/templates"
 )
 
+const (
+	resourceNvidiaGPU = "nvidia.com/gpu"
+)
+
 type CreateWorkerGroupOptions struct {
 	configFlags       *genericclioptions.ConfigFlags
 	ioStreams         *genericclioptions.IOStreams
@@ -132,6 +136,19 @@ func (options *CreateWorkerGroupOptions) Run(ctx context.Context, factory cmduti
 	}
 
 	newRayCluster := rayCluster.DeepCopy()
+
+	newRayCluster.Spec.WorkerGroupSpecs = append(newRayCluster.Spec.WorkerGroupSpecs, createWorkerGroupSpec(options))
+
+	newRayCluster, err = k8sClient.RayClient().RayV1().RayClusters(*options.configFlags.Namespace).Update(ctx, newRayCluster, metav1.UpdateOptions{})
+	if err != nil {
+		return fmt.Errorf("error updating Ray cluster with new worker group: %w", err)
+	}
+
+	fmt.Printf("Updated Ray cluster %s/%s with new worker group\n", newRayCluster.Namespace, newRayCluster.Name)
+	return nil
+}
+
+func createWorkerGroupSpec(options *CreateWorkerGroupOptions) rayv1.WorkerGroupSpec {
 	podTemplate := corev1.PodTemplateSpec{
 		Spec: corev1.PodSpec{
 			Containers: []corev1.Container{
@@ -155,25 +172,16 @@ func (options *CreateWorkerGroupOptions) Run(ctx context.Context, factory cmduti
 
 	gpuResource := resource.MustParse(options.workerGPU)
 	if !gpuResource.IsZero() {
-		podTemplate.Spec.Containers[0].Resources.Requests[corev1.ResourceName("nvidia.com/gpu")] = gpuResource
-		podTemplate.Spec.Containers[0].Resources.Limits[corev1.ResourceName("nvidia.com/gpu")] = gpuResource
+		podTemplate.Spec.Containers[0].Resources.Requests[corev1.ResourceName(resourceNvidiaGPU)] = gpuResource
+		podTemplate.Spec.Containers[0].Resources.Limits[corev1.ResourceName(resourceNvidiaGPU)] = gpuResource
 	}
 
-	workerGroup := rayv1.WorkerGroupSpec{
+	return rayv1.WorkerGroupSpec{
 		GroupName:      options.groupName,
 		Replicas:       &options.workerReplicas,
 		MinReplicas:    &options.workerMinReplicas,
 		MaxReplicas:    &options.workerMaxReplicas,
 		RayStartParams: map[string]string{},
 		Template:       podTemplate,
 	}
-	newRayCluster.Spec.WorkerGroupSpecs = append(newRayCluster.Spec.WorkerGroupSpecs, workerGroup)
-
-	newRayCluster, err = k8sClient.RayClient().RayV1().RayClusters(*options.configFlags.Namespace).Update(ctx, newRayCluster, metav1.UpdateOptions{})
-	if err != nil {
-		return fmt.Errorf("error updating Ray cluster with new worker group: %w", err)
-	}
-
-	fmt.Printf("Updated Ray cluster %s/%s with new worker group\n", newRayCluster.Namespace, newRayCluster.Name)
-	return nil
 }

kubectl-plugin/pkg/cmd/create/create_workergroup_test.go

@@ -0,0 +1,57 @@
+package create
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	corev1 "k8s.io/api/core/v1"
+	"k8s.io/apimachinery/pkg/api/resource"
+	"k8s.io/utils/ptr"
+
+	rayv1 "github.com/ray-project/kuberay/ray-operator/apis/ray/v1"
+)
+
+func TestCreateWorkerGroupSpec(t *testing.T) {
+	options := &CreateWorkerGroupOptions{
+		groupName:         "example-group",
+		image:             "DEADBEEF",
+		workerReplicas:    3,
+		workerMinReplicas: 1,
+		workerMaxReplicas: 5,
+		workerCPU:         "2",
+		workerMemory:      "5Gi",
+		workerGPU:         "1",
+	}
+
+	expected := rayv1.WorkerGroupSpec{
+		RayStartParams: map[string]string{},
+		GroupName:      "example-group",
+		Template: corev1.PodTemplateSpec{
+			Spec: corev1.PodSpec{
+				Containers: []corev1.Container{
+					{
+						Name:  "ray-worker",
+						Image: "DEADBEEF",
+						Resources: corev1.ResourceRequirements{
+							Requests: corev1.ResourceList{
+								corev1.ResourceCPU:    resource.MustParse("2"),
+								corev1.ResourceMemory: resource.MustParse("5Gi"),
+								resourceNvidiaGPU:     resource.MustParse("1"),
+							},
+							Limits: corev1.ResourceList{
+								corev1.ResourceCPU:    resource.MustParse("2"),
+								corev1.ResourceMemory: resource.MustParse("5Gi"),
+								resourceNvidiaGPU:     resource.MustParse("1"),
+							},
+						},
+					},
+				},
+			},
+		},
+		Replicas:    ptr.To[int32](3),
+		MinReplicas: ptr.To[int32](1),
+		MaxReplicas: ptr.To[int32](5),
+	}
+
+	assert.Equal(t, expected, createWorkerGroupSpec(options))
+}