From 5cf6bf2cabddd7e103ff6c9a756a13feb107920d Mon Sep 17 00:00:00 2001 From: David Xia Date: Tue, 18 Feb 2025 16:19:50 -0500 Subject: [PATCH] [refactor][kubectl-plugin] make `--node-type` an enum flag (#3018) Signed-off-by: David Xia --- kubectl-plugin/pkg/cmd/log/log.go | 53 ++++++++++++++++++++++---- kubectl-plugin/pkg/cmd/log/log_test.go | 25 +++++------- 2 files changed, 55 insertions(+), 23 deletions(-) diff --git a/kubectl-plugin/pkg/cmd/log/log.go b/kubectl-plugin/pkg/cmd/log/log.go index 232f6ce9cc7..e156b16ac52 100644 --- a/kubectl-plugin/pkg/cmd/log/log.go +++ b/kubectl-plugin/pkg/cmd/log/log.go @@ -32,12 +32,51 @@ import ( const filePathInPod = "/tmp/ray/session_latest/logs/" +type nodeTypeEnum string + +const ( + allNodeType nodeTypeEnum = "all" + headNodeType nodeTypeEnum = "head" + workerNodeType nodeTypeEnum = "worker" +) + +// String is used both by fmt.Print and by Cobra in help text +func (e *nodeTypeEnum) String() string { + return string(*e) +} + +// Set must have pointer receiver so it doesn't change the value of a copy +func (e *nodeTypeEnum) Set(v string) error { + val := strings.ToLower(v) + + switch val { + case string(allNodeType), string(headNodeType), string(workerNodeType): + *e = nodeTypeEnum(val) + return nil + default: + return fmt.Errorf("must be one of %q, %q, or %q", allNodeType, headNodeType, workerNodeType) + } +} + +// Type is only used in help text +func (e *nodeTypeEnum) Type() string { + return "enum" +} + +func nodeTypeCompletion(_ *cobra.Command, _ []string, _ string) ([]string, cobra.ShellCompDirective) { + return []string{ + fmt.Sprintf("%s\tdownload logs from %s Ray nodes", allNodeType, allNodeType), + fmt.Sprintf("%s\tdownload logs from %s Ray nodes", headNodeType, headNodeType), + fmt.Sprintf("%s\tdownload logs from %s Ray nodes", workerNodeType, workerNodeType), + }, cobra.ShellCompDirectiveDefault +} + type ClusterLogOptions struct { configFlags *genericclioptions.ConfigFlags ioStreams *genericclioptions.IOStreams Executor RemoteExecutor outputDir string - nodeType string + nodeType nodeTypeEnum ResourceName string ResourceType util.ResourceType } @@ -73,6 +112,7 @@ func NewClusterLogOptions(streams genericclioptions.IOStreams) *ClusterLogOption configFlags: genericclioptions.NewConfigFlags(true), ioStreams: &streams, Executor: &DefaultRemoteExecutor{}, + nodeType: allNodeType, } } @@ -106,7 +146,10 @@ func NewClusterLogCommand(streams genericclioptions.IOStreams) *cobra.Command { }, } cmd.Flags().StringVar(&options.outputDir, "out-dir", options.outputDir, "directory to save the logs to") - cmd.Flags().StringVar(&options.nodeType, "node-type", options.nodeType, "type of Ray node from which to download log, supports 'worker', 'head', or 'all'") + cmd.Flags().Var(&options.nodeType, "node-type", "type of Ray node from which to download logs, supports: 'worker', 'head', or 'all'") + + cobra.CheckErr(cmd.RegisterFlagCompletionFunc("node-type", nodeTypeCompletion)) + options.configFlags.AddFlags(cmd.Flags()) return cmd } @@ -139,12 +182,6 @@ func (options *ClusterLogOptions) Complete(cmd *cobra.Command, args []string) er options.ResourceName = typeAndName[1] } - if options.nodeType == "" { - options.nodeType = "all" - } else { - options.nodeType = strings.ToLower(options.nodeType) - } - return nil } diff --git a/kubectl-plugin/pkg/cmd/log/log_test.go b/kubectl-plugin/pkg/cmd/log/log_test.go index c6b5611bae2..482a5da9334 100644 --- a/kubectl-plugin/pkg/cmd/log/log_test.go +++ b/kubectl-plugin/pkg/cmd/log/log_test.go @@ -117,18 +117,16 @@ func TestRayClusterLogComplete(t *testing.T) { tests := []struct { name string - nodeType string + nodeType nodeTypeEnum expectedResourceType util.ResourceType expectedResourceName string - expectedNodeType string args []string hasErr bool }{ { - name: "valid request with RayCluster with empty nodetype input", + name: "valid request with RayCluster", expectedResourceType: util.RayCluster, expectedResourceName: "test-raycluster", - expectedNodeType: "all", args: []string{"test-raycluster"}, hasErr: false, }, @@ -137,7 +135,6 @@ func TestRayClusterLogComplete(t *testing.T) { expectedResourceType: util.RayCluster, expectedResourceName: "test-raycluster", args: []string{"rayCluster/test-raycluster"}, - expectedNodeType: "all", hasErr: false, }, { @@ -145,7 +142,6 @@ func TestRayClusterLogComplete(t *testing.T) { expectedResourceType: util.RayService, expectedResourceName: "test-rayservice", args: []string{"rayservice/test-rayservice"}, - expectedNodeType: "all", hasErr: false, }, { @@ -153,7 +149,6 @@ func TestRayClusterLogComplete(t *testing.T) { expectedResourceType: util.RayJob, expectedResourceName: "test-rayJob", args: []string{"rayJob/test-rayJob"}, - expectedNodeType: "all", hasErr: false, }, { @@ -177,6 +172,7 @@ func TestRayClusterLogComplete(t *testing.T) { t.Run(tc.name, func(t *testing.T) { testStreams, _, _, _ := genericclioptions.NewTestIOStreams() fakeClusterLogOptions := NewClusterLogOptions(testStreams) + fakeClusterLogOptions.nodeType = tc.nodeType err := fakeClusterLogOptions.Complete(cmd, tc.args) if tc.hasErr { require.Error(t, err) @@ -184,7 +180,6 @@ func TestRayClusterLogComplete(t *testing.T) { require.NoError(t, err) assert.Equal(t, tc.expectedResourceType, fakeClusterLogOptions.ResourceType) assert.Equal(t, tc.expectedResourceName, fakeClusterLogOptions.ResourceName) - assert.Equal(t, tc.expectedNodeType, fakeClusterLogOptions.nodeType) } }) } @@ -227,7 +222,7 @@ func TestRayClusterLogValidate(t *testing.T) { }, outputDir: fakeDir, ResourceName: "fake-cluster", - nodeType: "head", + nodeType: headNodeType, ioStreams: &testStreams, }, expectError: "no context is currently set, use \"--context\" or \"kubectl config use-context \" to select a new one", @@ -241,7 +236,7 @@ func TestRayClusterLogValidate(t *testing.T) { }, outputDir: fakeDir, ResourceName: "fake-cluster", - nodeType: "head", + nodeType: workerNodeType, ioStreams: &testStreams, }, }, @@ -254,7 +249,7 @@ func TestRayClusterLogValidate(t *testing.T) { }, outputDir: fakeDir, ResourceName: "fake-cluster", - nodeType: "head", + nodeType: allNodeType, ioStreams: &testStreams, }, }, @@ -277,7 +272,7 @@ func TestRayClusterLogValidate(t *testing.T) { configFlags: fakeConfigFlags, outputDir: fakeDir, ResourceName: "fake-cluster", - nodeType: "head", + nodeType: headNodeType, ioStreams: &testStreams, }, expectError: "", @@ -289,7 +284,7 @@ func TestRayClusterLogValidate(t *testing.T) { configFlags: fakeConfigFlags, outputDir: "", ResourceName: "fake-cluster", - nodeType: "head", + nodeType: headNodeType, ioStreams: &testStreams, }, expectError: "", @@ -301,7 +296,7 @@ func TestRayClusterLogValidate(t *testing.T) { configFlags: fakeConfigFlags, outputDir: "randomPath-here", ResourceName: "fake-cluster", - nodeType: "head", + nodeType: headNodeType, ioStreams: &testStreams, }, expectError: "Directory does not exist. Failed with: stat randomPath-here: no such file or directory", @@ -313,7 +308,7 @@ func TestRayClusterLogValidate(t *testing.T) { configFlags: fakeConfigFlags, outputDir: kubeConfigWithCurrentContext, ResourceName: "fake-cluster", - nodeType: "head", + nodeType: headNodeType, ioStreams: &testStreams, }, expectError: "Path is not a directory. Please input a directory and try again",