[RayJob] Add token authentication support for light weight job submitter (#4215)

* [RayJob] light weight job submitter auth token support

Signed-off-by: Future-Outlier <[email protected]>

* X-Ray-Authorization

Signed-off-by: Rueian <[email protected]>

---------

Signed-off-by: Future-Outlier <[email protected]>
Signed-off-by: Rueian <[email protected]>
Co-authored-by: Rueian <[email protected]>

Author: Future-Outlier

ray-operator/rayjob-submitter/cmd/main.go

@@ -64,7 +64,8 @@ func main() {
 	}
 	rayDashboardClient := &dashboardclient.RayDashboardClient{}
 	address = rayjobsubmitter.JobSubmissionURL(address)
-	rayDashboardClient.InitClient(&http.Client{Timeout: time.Second * 10}, address, "")
+	authToken := os.Getenv("RAY_AUTH_TOKEN")
+	rayDashboardClient.InitClient(&http.Client{Timeout: time.Second * 10}, address, authToken)
 	submissionId, err := rayDashboardClient.SubmitJobReq(context.Background(), &req)
 	if err != nil {
 		if strings.Contains(err.Error(), "Please use a different submission_id") {
@@ -76,7 +77,7 @@ func main() {
 		fmt.Fprintf(os.Stdout, "SUCC -- Job '%s' submitted successfully\n", submissionId)
 	}
 	fmt.Fprintf(os.Stdout, "INFO -- Tailing logs until the job finishes:\n")
-	err = rayjobsubmitter.TailJobLogs(address, submissionId, os.Stdout)
+	err = rayjobsubmitter.TailJobLogs(address, submissionId, authToken, os.Stdout)
 	exitOnError(err)
 }
 

ray-operator/rayjob-submitter/rayjob-submitter.go

@@ -31,12 +31,22 @@ func logTailingURL(address, submissionId string) (string, error) {
 	return address, nil
 }
 
-func TailJobLogs(address, submissionId string, out io.Writer) error {
+func TailJobLogs(address, submissionId string, authToken string, out io.Writer) error {
 	wsAddr, err := logTailingURL(address, submissionId)
 	if err != nil {
 		return err
 	}
-	c, _, err := websocket.Dial(context.Background(), wsAddr, nil)
+
+	var dialOptions *websocket.DialOptions
+	if authToken != "" {
+		dialOptions = &websocket.DialOptions{
+			HTTPHeader: map[string][]string{
+				"X-Ray-Authorization": {fmt.Sprintf("Bearer %s", authToken)},
+			},
+		}
+	}
+
+	c, _, err := websocket.Dial(context.Background(), wsAddr, dialOptions)
 	if err != nil {
 		return err
 	}