[kubectl-plugin] Add head/worker node selector option (#3228)

* add node selector option for kubectl plugin create cluster

Signed-off-by: Troy Chiu <[email protected]>

* nit

Signed-off-by: Troy Chiu <[email protected]>

---------

Signed-off-by: Troy Chiu <[email protected]>

Author: troychiu

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

@@ -23,6 +23,8 @@ type CreateClusterOptions struct {
 	ioStreams              *genericclioptions.IOStreams
 	workerRayStartParams   map[string]string
 	headRayStartParams     map[string]string
+	headNodeSelectors      map[string]string
+	workerNodeSelectors    map[string]string
 	kubeContexter          util.KubeContexter
 	clusterName            string
 	rayVersion             string
@@ -106,6 +108,8 @@ func NewCreateClusterCommand(streams genericclioptions.IOStreams) *cobra.Command
 	cmd.Flags().BoolVar(&options.dryRun, "dry-run", false, "print the generated YAML instead of creating the cluster")
 	cmd.Flags().BoolVar(&options.wait, "wait", false, "wait for the cluster to be provisioned before returning. Returns an error if the cluster is not provisioned by the timeout specified")
 	cmd.Flags().DurationVar(&options.timeout, "timeout", defaultProvisionedTimeout, "the timeout for --wait")
+	cmd.Flags().StringToStringVar(&options.headNodeSelectors, "head-node-selectors", nil, "Node selectors to apply to all head pods in the cluster (e.g. --head-node-selector=cloud.google.com/gke-accelerator=nvidia-l4,cloud.google.com/gke-nodepool=my-node-pool)")
+	cmd.Flags().StringToStringVar(&options.workerNodeSelectors, "worker-node-selectors", nil, "Node selectors to apply to all worker pods in the cluster (e.g. --worker-node-selector=cloud.google.com/gke-accelerator=nvidia-l4,cloud.google.com/gke-nodepool=my-node-pool)")
 
 	options.configFlags.AddFlags(cmd.Flags())
 	return cmd
@@ -182,6 +186,8 @@ func (options *CreateClusterOptions) Run(ctx context.Context, k8sClient client.C
 			WorkerEphemeralStorage: options.workerEphemeralStorage,
 			WorkerGPU:              options.workerGPU,
 			WorkerRayStartParams:   options.workerRayStartParams,
+			HeadNodeSelectors:      options.headNodeSelectors,
+			WorkerNodeSelectors:    options.workerNodeSelectors,
 		},
 	}
 

kubectl-plugin/pkg/util/generation/generation.go

@@ -19,6 +19,8 @@ import (
 type RayClusterSpecObject struct {
 	HeadRayStartParams     map[string]string
 	WorkerRayStartParams   map[string]string
+	HeadNodeSelectors      map[string]string
+	WorkerNodeSelectors    map[string]string
 	RayVersion             string
 	Image                  string
 	HeadCPU                string
@@ -107,6 +109,7 @@ func (rayClusterSpecObject *RayClusterSpecObject) generateRayClusterSpec() *rayv
 			WithRayStartParams(headRayStartParams).
 			WithTemplate(corev1ac.PodTemplateSpec().
 				WithSpec(corev1ac.PodSpec().
+					WithNodeSelector(rayClusterSpecObject.HeadNodeSelectors).
 					WithContainers(corev1ac.Container().
 						WithName("ray-head").
 						WithImage(rayClusterSpecObject.Image).
@@ -122,6 +125,7 @@ func (rayClusterSpecObject *RayClusterSpecObject) generateRayClusterSpec() *rayv
 			WithReplicas(rayClusterSpecObject.WorkerReplicas).
 			WithTemplate(corev1ac.PodTemplateSpec().
 				WithSpec(corev1ac.PodSpec().
+					WithNodeSelector(rayClusterSpecObject.WorkerNodeSelectors).
 					WithContainers(corev1ac.Container().
 						WithName("ray-worker").
 						WithImage(rayClusterSpecObject.Image).

kubectl-plugin/pkg/util/generation/generation_test.go

@@ -257,6 +257,14 @@ func TestGenerateRayClusterSpec(t *testing.T) {
 		WorkerCPU:      "2",
 		WorkerMemory:   "10Gi",
 		WorkerGPU:      "0",
+		HeadNodeSelectors: map[string]string{
+			"head-selector1": "foo",
+			"head-selector2": "bar",
+		},
+		WorkerNodeSelectors: map[string]string{
+			"worker-selector1": "baz",
+			"worker-selector2": "qux",
+		},
 	}
 
 	expected := &rayv1ac.RayClusterSpecApplyConfiguration{
@@ -299,6 +307,10 @@ func TestGenerateRayClusterSpec(t *testing.T) {
 							},
 						},
 					},
+					NodeSelector: map[string]string{
+						"head-selector1": "foo",
+						"head-selector2": "bar",
+					},
 				},
 			},
 		},
@@ -325,6 +337,10 @@ func TestGenerateRayClusterSpec(t *testing.T) {
 								},
 							},
 						},
+						NodeSelector: map[string]string{
+							"worker-selector1": "baz",
+							"worker-selector2": "qux",
+						},
 					},
 				},
 			},