[kubectl-plugin][feat] support specifying number of head GPUs

when creating a RayCluster with `kubectl ray create cluster NAME --head-gpu N`. Similar to the `--worker-gpu` switch. Signed-off-by: David Xia <[email protected]>
ray-project · Feb 4, 2025 · dd07f38 · dd07f38
1 parent c8d34f4
commit dd07f38
Show file tree

Hide file tree

Showing 4 changed files with 32 additions and 5 deletions.
diff --git a/kubectl-plugin/pkg/cmd/create/create_cluster.go b/kubectl-plugin/pkg/cmd/create/create_cluster.go
@@ -22,6 +22,7 @@ type CreateClusterOptions struct {
 	image          string
 	headCPU        string
 	headMemory     string
+	headGPU        string
 	workerCPU      string
 	workerMemory   string
 	workerGPU      string
@@ -75,10 +76,11 @@ func NewCreateClusterCommand(streams genericclioptions.IOStreams) *cobra.Command
 	cmd.Flags().StringVar(&options.image, "image", options.image, "Ray image to use in the Ray Cluster yaml")
 	cmd.Flags().StringVar(&options.headCPU, "head-cpu", "2", "Number of CPU for the ray head. Default to 2")
 	cmd.Flags().StringVar(&options.headMemory, "head-memory", "4Gi", "Amount of memory to use for the ray head. Default to 4Gi")
+	cmd.Flags().StringVar(&options.headGPU, "head-gpu", "0", "Number of GPUs for the ray head. Default to 0")
 	cmd.Flags().Int32Var(&options.workerReplicas, "worker-replicas", 1, "Number of the worker group replicas. Default of 1")
 	cmd.Flags().StringVar(&options.workerCPU, "worker-cpu", "2", "Number of CPU for the ray worker. Default to 2")
 	cmd.Flags().StringVar(&options.workerMemory, "worker-memory", "4Gi", "Amount of memory to use for the ray worker. Default to 4Gi")
-	cmd.Flags().StringVar(&options.workerGPU, "worker-gpu", "0", "Number of GPU for the ray worker. Default to 0")
+	cmd.Flags().StringVar(&options.workerGPU, "worker-gpu", "0", "Number of GPUs for the ray worker. Default to 0")
 	cmd.Flags().BoolVar(&options.dryRun, "dry-run", false, "Will not apply the generated cluster and will print out the generated yaml")
 
 	options.configFlags.AddFlags(cmd.Flags())
@@ -129,6 +131,7 @@ func (options *CreateClusterOptions) Run(ctx context.Context, factory cmdutil.Fa
 			Image:          options.image,
 			HeadCPU:        options.headCPU,
 			HeadMemory:     options.headMemory,
+			HeadGPU:        options.headGPU,
 			WorkerReplicas: options.workerReplicas,
 			WorkerCPU:      options.workerCPU,
 			WorkerMemory:   options.workerMemory,

diff --git a/kubectl-plugin/pkg/cmd/create/create_cluster_test.go b/kubectl-plugin/pkg/cmd/create/create_cluster_test.go
@@ -92,6 +92,7 @@ func TestRayCreateClusterValidate(t *testing.T) {
 				rayVersion:     "ray-version",
 				image:          "ray-image",
 				headCPU:        "5",
+				headGPU:        "1",
 				headMemory:     "5Gi",
 				workerReplicas: 3,
 				workerCPU:      "4",

diff --git a/kubectl-plugin/pkg/util/generation/generation.go b/kubectl-plugin/pkg/util/generation/generation.go
@@ -12,10 +12,15 @@ import (
 	rayv1ac "github.com/ray-project/kuberay/ray-operator/pkg/client/applyconfiguration/ray/v1"
 )
 
+const (
+	resourceNvidiaGPU = "nvidia.com/gpu"
+)
+
 type RayClusterSpecObject struct {
 	RayVersion     string
 	Image          string
 	HeadCPU        string
+	HeadGPU        string
 	HeadMemory     string
 	WorkerCPU      string
 	WorkerGPU      string
@@ -96,13 +101,25 @@ func (rayClusterSpecObject *RayClusterSpecObject) generateRayClusterSpec() *rayv
 								corev1.ResourceMemory: resource.MustParse(rayClusterSpecObject.WorkerMemory),
 							}))))))
 
-	gpuResource := resource.MustParse(rayClusterSpecObject.WorkerGPU)
-	if !gpuResource.IsZero() {
+	headGPUResource := resource.MustParse(rayClusterSpecObject.HeadGPU)
+	if !headGPUResource.IsZero() {
+		var requests, limits corev1.ResourceList
+		requests = *rayClusterSpec.HeadGroupSpec.Template.Spec.Containers[0].Resources.Requests
+		limits = *rayClusterSpec.HeadGroupSpec.Template.Spec.Containers[0].Resources.Limits
+		requests[corev1.ResourceName(resourceNvidiaGPU)] = headGPUResource
+		limits[corev1.ResourceName(resourceNvidiaGPU)] = headGPUResource
+
+		rayClusterSpec.HeadGroupSpec.Template.Spec.Containers[0].Resources.Requests = &requests
+		rayClusterSpec.HeadGroupSpec.Template.Spec.Containers[0].Resources.Limits = &limits
+	}
+
+	workerGPUResource := resource.MustParse(rayClusterSpecObject.WorkerGPU)
+	if !workerGPUResource.IsZero() {
 		var requests, limits corev1.ResourceList
 		requests = *rayClusterSpec.WorkerGroupSpecs[0].Template.Spec.Containers[0].Resources.Requests
 		limits = *rayClusterSpec.WorkerGroupSpecs[0].Template.Spec.Containers[0].Resources.Limits
-		requests[corev1.ResourceName("nvidia.com/gpu")] = gpuResource
-		limits[corev1.ResourceName("nvidia.com/gpu")] = gpuResource
+		requests[corev1.ResourceName(resourceNvidiaGPU)] = workerGPUResource
+		limits[corev1.ResourceName(resourceNvidiaGPU)] = workerGPUResource
 
 		rayClusterSpec.WorkerGroupSpecs[0].Template.Spec.Containers[0].Resources.Requests = &requests
 		rayClusterSpec.WorkerGroupSpecs[0].Template.Spec.Containers[0].Resources.Limits = &limits

diff --git a/kubectl-plugin/pkg/util/generation/generation_test.go b/kubectl-plugin/pkg/util/generation/generation_test.go
@@ -20,6 +20,7 @@ func TestGenerateRayCluterApplyConfig(t *testing.T) {
 			Image:          "rayproject/ray:2.39.0",
 			HeadCPU:        "1",
 			HeadMemory:     "5Gi",
+			HeadGPU:        "1",
 			WorkerReplicas: 3,
 			WorkerCPU:      "2",
 			WorkerMemory:   "10Gi",
@@ -34,6 +35,7 @@ func TestGenerateRayCluterApplyConfig(t *testing.T) {
 	assert.Equal(t, testRayClusterYamlObject.RayVersion, *result.Spec.RayVersion)
 	assert.Equal(t, testRayClusterYamlObject.Image, *result.Spec.HeadGroupSpec.Template.Spec.Containers[0].Image)
 	assert.Equal(t, resource.MustParse(testRayClusterYamlObject.HeadCPU), *result.Spec.HeadGroupSpec.Template.Spec.Containers[0].Resources.Requests.Cpu())
+	assert.Equal(t, resource.MustParse(testRayClusterYamlObject.HeadGPU), *result.Spec.HeadGroupSpec.Template.Spec.Containers[0].Resources.Requests.Name(corev1.ResourceName("nvidia.com/gpu"), resource.DecimalSI))
 	assert.Equal(t, resource.MustParse(testRayClusterYamlObject.HeadMemory), *result.Spec.HeadGroupSpec.Template.Spec.Containers[0].Resources.Requests.Memory())
 	assert.Equal(t, "default-group", *result.Spec.WorkerGroupSpecs[0].GroupName)
 	assert.Equal(t, testRayClusterYamlObject.WorkerReplicas, *result.Spec.WorkerGroupSpecs[0].Replicas)
@@ -51,6 +53,7 @@ func TestGenerateRayJobApplyConfig(t *testing.T) {
 			RayVersion:     "2.39.0",
 			Image:          "rayproject/ray:2.39.0",
 			HeadCPU:        "1",
+			HeadGPU:        "1",
 			HeadMemory:     "5Gi",
 			WorkerReplicas: 3,
 			WorkerCPU:      "2",
@@ -83,6 +86,7 @@ func TestConvertRayClusterApplyConfigToYaml(t *testing.T) {
 			Image:          "rayproject/ray:2.39.0",
 			HeadCPU:        "1",
 			HeadMemory:     "5Gi",
+			HeadGPU:        "1",
 			WorkerReplicas: 3,
 			WorkerCPU:      "2",
 			WorkerMemory:   "10Gi",
@@ -119,9 +123,11 @@ spec:
             limits:
               cpu: "1"
               memory: 5Gi
+              nvidia.com/gpu: "1"
             requests:
               cpu: "1"
               memory: 5Gi
+              nvidia.com/gpu: "1"
   rayVersion: 2.39.0
   workerGroupSpecs:
   - groupName: default-group