Support --address flag for kubectl ray job submit (#3922)

Signed-off-by: JosefNagelschmidt <info@jn.contact>
Co-authored-by: Jun-Hao Wan <ken89@kimo.com>

Author: JosefNagelschmidt

kubectl-plugin/pkg/cmd/job/job_submit.go

@@ -35,7 +35,7 @@ import (
 const (
 	dashboardAddr      = "http://localhost:8265"
 	clusterTimeout     = 120.0
-	portforwardtimeout = 60.0
+	portForwardTimeout = 60.0
 )
 
 type SubmitJobOptions struct {
@@ -45,6 +45,7 @@ type SubmitJobOptions struct {
 	workerNodeSelectors      map[string]string
 	headNodeSelectors        map[string]string
 	logColor                 string
+	address                  string
 	image                    string
 	fileName                 string
 	workingDir               string
@@ -86,8 +87,9 @@ type JobInfo struct {
 
 var (
 	jobSubmitLong = templates.LongDesc(`
-		Submit Ray job to Ray cluster as one would using Ray CLI e.g. 'ray job submit ENTRYPOINT'. Command supports all options that 'ray job submit' supports, except '--address'.
-		If Ray cluster is already setup, use 'kubectl ray session' instead.
+		Submit Ray job to Ray cluster as one would using Ray CLI e.g. 'ray job submit ENTRYPOINT'.
+		If Ray cluster is already setup, use 'kubectl ray session' instead. If '--address' is set, we connect directly
+		without port-forwarding; if empty, we port-forward to localhost:8265.
 
 		If no RayJob YAML file is specified, the command will create a default RayJob for the user.
 
@@ -149,6 +151,7 @@ func NewJobSubmitCommand(cmdFactory cmdutil.Factory, streams genericclioptions.I
 		},
 	}
 	cmd.Flags().StringVarP(&options.fileName, "filename", "f", "", "Path and name of the Ray Job YAML file")
+	cmd.Flags().StringVar(&options.address, "address", "", "Ray Dashboard base URL (e.g., https://ray.example.com). If set, skips port-forwarding.")
 	cmd.Flags().StringVar(&options.submissionID, "submission-id", "", "ID to specify for the Ray job. If not provided, one will be generated")
 	cmd.Flags().StringVar(&options.runtimeEnv, "runtime-env", "", "Path and name to the runtime env YAML file.")
 	cmd.Flags().StringVar(&options.workingDir, "working-dir", "", "Directory containing files that your job will run in")
@@ -307,6 +310,11 @@ func (options *SubmitJobOptions) Validate(cmd *cobra.Command) error {
 		}
 	}
 
+	if cmd.Flags().Changed("address") {
+		if strings.TrimSpace(options.address) == "" {
+			return fmt.Errorf("--address was provided but is empty")
+		}
+	}
 	return nil
 }
 
@@ -421,55 +429,62 @@ func (options *SubmitJobOptions) Run(ctx context.Context, factory cmdutil.Factor
 		return fmt.Errorf("Timed out waiting for cluster")
 	}
 
-	svcName, err := k8sClients.GetRayHeadSvcName(ctx, options.namespace, util.RayCluster, options.cluster)
-	if err != nil {
-		return fmt.Errorf("Failed to find service name: %w", err)
-	}
-
-	// start port forward section
-	portForwardCmd := portforward.NewCmdPortForward(factory, *options.ioStreams)
-	portForwardCmd.SetArgs([]string{"service/" + svcName, fmt.Sprintf("%d:%d", 8265, 8265)})
-
-	// create new context for port-forwarding so we can cancel the context to stop the port forwarding only
-	portforwardctx, cancel := context.WithCancel(ctx)
-	defer cancel()
-	go func() {
-		fmt.Printf("Port Forwarding service %s\n", svcName)
-		if err := portForwardCmd.ExecuteContext(portforwardctx); err != nil {
-			log.Fatalf("Error occurred while port-forwarding Ray dashboard: %v", err)
+	if options.address == "" {
+		svcName, err := k8sClients.GetRayHeadSvcName(ctx, options.namespace, util.RayCluster, options.cluster)
+		if err != nil {
+			return fmt.Errorf("Failed to find service name: %w", err)
 		}
-	}()
 
-	// Wait for port forward to be ready
-	var portforwardReady bool
-	portforwardWaitStartTime := time.Now()
-	currTime = portforwardWaitStartTime
+		// start port forward section
+		portForwardCmd := portforward.NewCmdPortForward(factory, *options.ioStreams)
+		portForwardCmd.SetArgs([]string{"service/" + svcName, fmt.Sprintf("%d:%d", 8265, 8265)})
+
+		// create new context for port-forwarding so we can cancel the context to stop the port forwarding only
+		portForwardCtx, cancel := context.WithCancel(ctx)
+		defer cancel()
+		go func() {
+			fmt.Printf("Port forwarding service %s\n", svcName)
+			if err := portForwardCmd.ExecuteContext(portForwardCtx); err != nil {
+				log.Fatalf("Error occurred while port-forwarding Ray dashboard: %v", err)
+			}
+		}()
 
-	portforwardCheckRequest, err := http.NewRequestWithContext(ctx, http.MethodGet, dashboardAddr, nil)
-	if err != nil {
-		return fmt.Errorf("Error occurred when trying to create request to probe cluster endpoint: %w", err)
-	}
-	httpClient := http.Client{
-		Timeout: 5 * time.Second,
-	}
-	fmt.Printf("Waiting for portforwarding...")
-	for !portforwardReady && currTime.Sub(portforwardWaitStartTime).Seconds() <= portforwardtimeout {
-		time.Sleep(2 * time.Second)
-		rayDashboardResponse, err := httpClient.Do(portforwardCheckRequest)
+		// Wait for port forward to be ready
+		var portForwardReady bool
+		portForwardWaitStartTime := time.Now()
+		currTime = portForwardWaitStartTime
+
+		portforwardCheckRequest, err := http.NewRequestWithContext(ctx, http.MethodGet, dashboardAddr, nil)
 		if err != nil {
-			err = fmt.Errorf("Error occurred when waiting for portforwarding: %w", err)
-			fmt.Println(err)
+			return fmt.Errorf("Error occurred when trying to create request to probe cluster endpoint: %w", err)
 		}
-		if rayDashboardResponse.StatusCode >= 200 && rayDashboardResponse.StatusCode < 300 {
-			portforwardReady = true
+		httpClient := http.Client{
+			Timeout: 5 * time.Second,
 		}
-		rayDashboardResponse.Body.Close()
-		currTime = time.Now()
-	}
-	if !portforwardReady {
-		return fmt.Errorf("Timed out waiting for port forwarding")
+		fmt.Printf("Waiting for port forwarding...")
+		for !portForwardReady && currTime.Sub(portForwardWaitStartTime).Seconds() <= portForwardTimeout {
+			time.Sleep(2 * time.Second)
+			rayDashboardResponse, err := httpClient.Do(portforwardCheckRequest)
+			if err != nil {
+				err = fmt.Errorf("Error occurred when waiting for port forwarding: %w", err)
+				fmt.Println(err)
+				currTime = time.Now()
+				continue
+			}
+			if rayDashboardResponse.StatusCode >= 200 && rayDashboardResponse.StatusCode < 300 {
+				portForwardReady = true
+			}
+			rayDashboardResponse.Body.Close()
+			currTime = time.Now()
+		}
+		if !portForwardReady {
+			return fmt.Errorf("Timed out waiting for port forwarding")
+		}
+		options.address = dashboardAddr
+		fmt.Printf("Port forwarding started on %s\n", options.address)
+	} else {
+		fmt.Printf("Using address %s (no port-forwarding)\n", options.address)
 	}
-	fmt.Printf("Portforwarding started on %s\n", dashboardAddr)
 
 	// If submission ID is not provided by the user, generate one.
 	if options.submissionID == "" {
@@ -608,7 +623,11 @@ func (options *SubmitJobOptions) Run(ctx context.Context, factory cmdutil.Factor
 }
 
 func (options *SubmitJobOptions) raySubmitCmd() ([]string, error) {
-	raySubmitCmd := []string{"ray", "job", "submit", "--address", dashboardAddr}
+	addr := options.address
+	if addr == "" {
+		addr = dashboardAddr
+	}
+	raySubmitCmd := []string{"ray", "job", "submit", "--address", addr}
 
 	if len(options.runtimeEnv) > 0 {
 		raySubmitCmd = append(raySubmitCmd, "--runtime-env", options.runtimeEnv)

kubectl-plugin/pkg/cmd/job/job_submit_test.go

@@ -195,6 +195,66 @@ func TestRayJobSubmitWithoutYamlValidate(t *testing.T) {
 	}
 }
 
+func TestRayJobSubmit_AddressValidation(t *testing.T) {
+	testStreams, _, _, _ := genericclioptions.NewTestIOStreams()
+	cmdFactory := cmdutil.NewFactory(genericclioptions.NewConfigFlags(true))
+
+	test := []struct {
+		name             string
+		address          string
+		expectError      string
+		expectNormalized string
+		flagChanged      bool
+	}{
+		{
+			name:        "address flag not set: port-forward mode",
+			address:     "",
+			flagChanged: false,
+		},
+		{
+			name:        "address flag set but empty: error",
+			address:     "",
+			flagChanged: true,
+			expectError: "--address was provided but is empty",
+		},
+		{
+			name:        "valid https address: OK",
+			address:     "https://ingress.example.com",
+			flagChanged: true,
+		},
+	}
+
+	for _, tc := range test {
+		t.Run(tc.name, func(t *testing.T) {
+			opts := &SubmitJobOptions{
+				cmdFactory: cmdFactory,
+				ioStreams:  &testStreams,
+				rayjobName: "rayjob-sample",
+				workingDir: "Fake/File/Path",
+			}
+
+			cmd := &cobra.Command{}
+			cmd.Flags().StringVar(&opts.address, "address", "", "Ray Dashboard base URL")
+
+			if tc.flagChanged {
+				require.NoError(t, cmd.Flags().Set("address", tc.address))
+			} else {
+				opts.address = tc.address
+			}
+
+			err := opts.Validate(cmd)
+			if tc.expectError != "" {
+				require.EqualError(t, err, tc.expectError)
+			} else {
+				require.NoError(t, err)
+				if tc.expectNormalized != "" {
+					require.Equal(t, tc.expectNormalized, opts.address)
+				}
+			}
+		})
+	}
+}
+
 func TestRayJobSubmitCmdFlagsOverrideYaml(t *testing.T) {
 	testStreams, _, _, _ := genericclioptions.NewTestIOStreams()
 	cmdFactory := cmdutil.NewFactory(genericclioptions.NewConfigFlags(true))
@@ -438,6 +498,7 @@ func TestRaySubmitCmd(t *testing.T) {
 	fakeSubmitJobOptions.verify = "True"
 	fakeSubmitJobOptions.workingDir = "/fake/working/dir"
 	fakeSubmitJobOptions.entryPoint = "python fake_python_script.py"
+	fakeSubmitJobOptions.address = dashboardAddr
 
 	actualCmd, err := fakeSubmitJobOptions.raySubmitCmd()
 	require.NoError(t, err)
@@ -489,3 +550,51 @@ func TestRayJobSubmit_FlagsHaveDefaults(t *testing.T) {
 	assert.Equal(t, 0, opts.entryPointMemory, "default entrypoint-memory should be 0")
 	assert.False(t, opts.noWait, "default no-wait should be false")
 }
+
+func TestRaySubmitCmd_AddressSelection(t *testing.T) {
+	streams, _, _, _ := genericclioptions.NewTestIOStreams()
+	factory := cmdutil.NewFactory(genericclioptions.NewConfigFlags(true))
+
+	makeCmd := func(addr string) ([]string, error) {
+		opts := NewJobSubmitOptions(factory, streams)
+		opts.workingDir = "/fake/working/dir"
+		opts.entryPoint = "python fake.py"
+		opts.address = addr
+		return opts.raySubmitCmd()
+	}
+
+	tests := []struct {
+		name         string
+		address      string
+		expectedAddr string
+	}{
+		{
+			name:         "no address provided: falls back to dashboardAddr",
+			address:      "",
+			expectedAddr: dashboardAddr,
+		},
+		{
+			name:         "custom address provided: uses custom",
+			address:      "https://ingress.example.com",
+			expectedAddr: "https://ingress.example.com",
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			cmd, err := makeCmd(tc.address)
+			require.NoError(t, err)
+
+			require.GreaterOrEqual(t, len(cmd), 5, "command too short")
+			assert.Equal(t, "ray", cmd[0])
+			assert.Equal(t, "job", cmd[1])
+			assert.Equal(t, "submit", cmd[2])
+
+			assert.Equal(t, "--address", cmd[3])
+			assert.Equal(t, tc.expectedAddr, cmd[4])
+
+			require.Contains(t, cmd, "--working-dir")
+			require.Contains(t, cmd, "--")
+		})
+	}
+}