[Core] pass auth token in dashboard head python client sdk (#58281)

Supports token based authentication in dashboard head sdk, all clients
which build on top of the submission_client will now support token auth
out of the box. so this covers all cli commands like job submit, state
api, serve related cli commands etc.

---------

Signed-off-by: sampan <[email protected]>
Signed-off-by: Sampan S Nayak <[email protected]>
Signed-off-by: Edward Oakes <[email protected]>
Co-authored-by: sampan <[email protected]>
Co-authored-by: Edward Oakes <[email protected]>

Author: 3 people

python/ray/_private/authentication/authentication_constants.py

@@ -0,0 +1,23 @@
+# Token setup instructions (used in multiple contexts)
+TOKEN_SETUP_INSTRUCTIONS = """Please provide an authentication token using one of these methods:
+  1. Set the RAY_AUTH_TOKEN environment variable
+  2. Set the RAY_AUTH_TOKEN_PATH environment variable (pointing to a token file)
+  3. Create a token file at the default location: ~/.ray/auth_token"""
+
+# When token auth is enabled but no token is found anywhere
+TOKEN_AUTH_ENABLED_BUT_NO_TOKEN_FOUND_ERROR_MESSAGE = (
+    "Token authentication is enabled but no authentication token was found. "
+    + TOKEN_SETUP_INSTRUCTIONS
+)
+
+# When HTTP request fails with 401 (Unauthorized - missing token)
+HTTP_REQUEST_MISSING_TOKEN_ERROR_MESSAGE = (
+    "The Ray cluster requires authentication, but no token was provided.\n\n"
+    + TOKEN_SETUP_INSTRUCTIONS
+)
+
+# When HTTP request fails with 403 (Forbidden - invalid token)
+HTTP_REQUEST_INVALID_TOKEN_ERROR_MESSAGE = (
+    "The authentication token you provided is invalid or incorrect.\n\n"
+    + TOKEN_SETUP_INSTRUCTIONS
+)

python/ray/_private/authentication/authentication_token_setup.py

@@ -9,6 +9,9 @@
 from pathlib import Path
 from typing import Any, Dict, Optional
 
+from ray._private.authentication.authentication_constants import (
+    TOKEN_AUTH_ENABLED_BUT_NO_TOKEN_FOUND_ERROR_MESSAGE,
+)
 from ray._private.authentication.authentication_token_generator import (
     generate_new_authentication_token,
 )
@@ -20,13 +23,6 @@
 
 logger = logging.getLogger(__name__)
 
-TOKEN_AUTH_ENABLED_BUT_NO_TOKEN_FOUND_ERROR_MESSAGE = (
-    "Token authentication is enabled but no authentication token was found. Please provide a token with one of these options:\n"
-    + "  1. RAY_AUTH_TOKEN environment variable\n"
-    + "  2. RAY_AUTH_TOKEN_PATH environment variable (path to token file)\n"
-    + "  3. Default token file: ~/.ray/auth_token"
-)
-
 
 def generate_and_save_token() -> None:
     """Generate a new random token and save it in the default token path.

python/ray/dashboard/modules/dashboard_sdk.py

@@ -12,6 +12,7 @@
 import yaml
 
 import ray
+from ray._private.authentication import authentication_constants
 from ray._private.runtime_env.packaging import (
     create_package,
     get_uri_for_directory,
@@ -20,7 +21,9 @@
 from ray._private.runtime_env.py_modules import upload_py_modules_if_needed
 from ray._private.runtime_env.working_dir import upload_working_dir_if_needed
 from ray._private.utils import split_address
+from ray._raylet import AuthenticationTokenLoader
 from ray.autoscaler._private.cli_logger import cli_logger
+from ray.dashboard.authentication_utils import is_token_auth_enabled
 from ray.dashboard.modules.job.common import uri_to_http_components
 from ray.util.annotations import DeveloperAPI, PublicAPI
 
@@ -222,7 +225,9 @@ def __init__(
         self._default_metadata = cluster_info.metadata or {}
         # Headers used for all requests sent to job server, optional and only
         # needed for cases like authentication to remote cluster.
-        self._headers = cluster_info.headers
+        self._headers = cluster_info.headers or {}
+        self._headers.update(**self._get_auth_headers())
+
         # Set SSL verify parameter for the requests library and create an ssl_context
         # object when needed for the aiohttp library.
         self._verify = verify
@@ -242,6 +247,36 @@ def __init__(
             else:
                 self._ssl_context = None
 
+    def _get_auth_headers(self) -> Dict[str, str]:
+        """Get authentication headers if token auth is enabled.
+
+        Returns:
+            dict: Authentication headers to merge with request headers.
+                  Empty dict if no auth needed or token unavailable.
+        """
+        if not is_token_auth_enabled():
+            return {}
+
+        # Check if user provided their own Authorization header (case-insensitive)
+        has_user_auth = any(
+            key.lower() == "authorization" for key in self._headers.keys()
+        )
+        if has_user_auth:
+            # User has provided their own auth header, don't override
+            return {}
+
+        token_loader = AuthenticationTokenLoader.instance()
+        auth_headers = token_loader.get_token_for_http_header()
+
+        if not auth_headers:
+            # Token auth enabled but no token found
+            logger.warning(
+                "Token authentication is enabled but no token was found. "
+                "Requests to authenticated clusters will fail."
+            )
+
+        return auth_headers
+
     def _check_connection_and_version(
         self, min_version: str = "1.9", version_error_message: str = None
     ):
@@ -293,14 +328,15 @@ def _do_request(
         json_data: Optional[dict] = None,
         **kwargs,
     ) -> "requests.Response":
-        """Perform the actual HTTP request
+        """Perform the actual HTTP request with authentication error handling.
 
         Keyword arguments other than "cookies", "headers" are forwarded to the
         `requests.request()`.
         """
         url = self._address + endpoint
         logger.debug(f"Sending request to {url} with json data: {json_data or {}}.")
-        return requests.request(
+
+        response = requests.request(
             method,
             url,
             cookies=self._cookies,
@@ -311,6 +347,22 @@ def _do_request(
             **kwargs,
         )
 
+        # Check for authentication errors and provide helpful messages
+        if response.status_code == 401:
+            # Unauthorized - missing or no token provided
+            raise RuntimeError(
+                f"Authentication required: {response.text}\n\n"
+                + authentication_constants.HTTP_REQUEST_MISSING_TOKEN_ERROR_MESSAGE
+            )
+        elif response.status_code == 403:
+            # Forbidden - invalid token
+            raise RuntimeError(
+                f"Authentication failed: {response.text}\n\n"
+                + authentication_constants.HTTP_REQUEST_INVALID_TOKEN_ERROR_MESSAGE
+            )
+
+        return response
+
     def _package_exists(
         self,
         package_uri: str,

python/ray/dashboard/tests/test_dashboard_auth.py

@@ -1,117 +1,68 @@
 """Tests for dashboard token authentication."""
-import os
+
 import sys
 
 import pytest
 import requests
 
-import ray
-from ray._raylet import Config
-
-
-@pytest.fixture
-def start_ray_with_env_vars(request):
-    """Clean up environment variables after each test."""
-    env_vars = getattr(request, "param", {}).pop("env_vars", {})
-    os.environ.update(**env_vars)
-    Config.initialize("")
-
-    yield ray.init()
-
-    ray.shutdown()
-    for k in env_vars.keys():
-        del os.environ[k]
 
-
-TEST_TOKEN = "test_token_12345678901234567890123456789012"
-
-
-@pytest.mark.parametrize(
-    "start_ray_with_env_vars",
-    [
-        {
-            "env_vars": {"RAY_auth_mode": "token", "RAY_AUTH_TOKEN": TEST_TOKEN},
-        },
-    ],
-    indirect=True,
-)
-def test_auth_enabled_valid_token(start_ray_with_env_vars):
+def test_dashboard_request_requires_auth_with_valid_token(
+    setup_cluster_with_token_auth,
+):
     """Test that requests succeed with valid token when auth is enabled."""
-    dashboard_url = start_ray_with_env_vars.address_info["webui_url"]
 
-    # Request with valid auth should succeed
-    headers = {"Authorization": f"Bearer {TEST_TOKEN}"}
+    cluster_info = setup_cluster_with_token_auth
+    headers = {"Authorization": f"Bearer {cluster_info['token']}"}
+
     response = requests.get(
-        f"http://{dashboard_url}/api/component_activities",
+        f"{cluster_info['dashboard_url']}/api/component_activities",
         headers=headers,
     )
+
     assert response.status_code == 200
 
 
-@pytest.mark.parametrize(
-    "start_ray_with_env_vars",
-    [
-        {
-            "env_vars": {"RAY_auth_mode": "token", "RAY_AUTH_TOKEN": TEST_TOKEN},
-        },
-    ],
-    indirect=True,
-)
-def test_auth_enabled_missing_token(start_ray_with_env_vars):
+def test_dashboard_request_requires_auth_missing_token(setup_cluster_with_token_auth):
     """Test that requests fail without token when auth is enabled."""
-    dashboard_url = start_ray_with_env_vars.address_info["webui_url"]
 
-    # GET without auth should fail with 401
+    cluster_info = setup_cluster_with_token_auth
+
     response = requests.get(
-        f"http://{dashboard_url}/api/component_activities",
+        f"{cluster_info['dashboard_url']}/api/component_activities",
         json={"test": "data"},
     )
+
     assert response.status_code == 401
 
 
-@pytest.mark.parametrize(
-    "start_ray_with_env_vars",
-    [
-        {
-            "env_vars": {"RAY_auth_mode": "token", "RAY_AUTH_TOKEN": TEST_TOKEN},
-        },
-    ],
-    indirect=True,
-)
-def test_auth_enabled_invalid_token(start_ray_with_env_vars):
+def test_dashboard_request_requires_auth_invalid_token(setup_cluster_with_token_auth):
     """Test that requests fail with invalid token when auth is enabled."""
-    dashboard_url = start_ray_with_env_vars.address_info["webui_url"]
 
-    # Request with wrong token should fail with 403
-    headers = {"Authorization": "Bearer INCORRECT_TOKEN"}
+    cluster_info = setup_cluster_with_token_auth
+    headers = {"Authorization": "Bearer wrong_token_00000000000000000000000000000000"}
+
     response = requests.get(
-        f"http://{dashboard_url}/api/component_activities",
+        f"{cluster_info['dashboard_url']}/api/component_activities",
         json={"test": "data"},
         headers=headers,
     )
+
     assert response.status_code == 403
 
 
-@pytest.mark.parametrize(
-    "start_ray_with_env_vars",
-    [
-        {
-            "env_vars": {"RAY_auth_mode": "disabled"},
-        },
-    ],
-    indirect=True,
-)
-def test_auth_disabled(start_ray_with_env_vars):
+def test_dashboard_auth_disabled(setup_cluster_without_token_auth):
     """Test that auth is not enforced when auth_mode is disabled."""
-    dashboard_url = start_ray_with_env_vars.address_info["webui_url"]
 
-    # GET without auth should succeed when auth is disabled
+    cluster_info = setup_cluster_without_token_auth
+
     response = requests.get(
-        f"http://{dashboard_url}/api/component_activities", json={"test": "data"}
+        f"{cluster_info['dashboard_url']}/api/component_activities",
+        json={"test": "data"},
     )
-    # Should not return 401 or 403
+
     assert response.status_code == 200
 
 
 if __name__ == "__main__":
+
     sys.exit(pytest.main(["-vv", __file__]))

python/ray/includes/rpc_token_authentication.pxd

@@ -16,6 +16,7 @@ cdef extern from "ray/rpc/authentication/authentication_token.h" namespace "ray:
         CAuthenticationToken(string value)
         c_bool empty()
         c_bool Equals(const CAuthenticationToken& other)
+        string ToAuthorizationHeaderValue()
         @staticmethod
         CAuthenticationToken FromMetadata(string metadata_value)
 

python/ray/includes/rpc_token_authentication.pxi

@@ -11,6 +11,7 @@ class AuthenticationMode:
     DISABLED = CAuthenticationMode.DISABLED
     TOKEN = CAuthenticationMode.TOKEN
 
+_AUTHORIZATION_HEADER_NAME = "authorization"
 
 def get_authentication_mode():
     """Get the current authentication mode.
@@ -69,3 +70,26 @@ class AuthenticationTokenLoader:
         variables or files on the next request.
         """
         CAuthenticationTokenLoader.instance().ResetCache()
+
+    def get_token_for_http_header(self) -> dict:
+        """Get authentication token as a dictionary for HTTP headers.
+
+        This method loads the token from C++ AuthenticationTokenLoader and returns it
+        as a dictionary that can be merged with existing headers. It returns an empty
+        dictionary if:
+        - A token does not exist
+        - The token is empty
+
+        Returns:
+            dict: Empty dict or {"authorization": "Bearer <token>"}
+        """
+        if not self.has_token():
+            return {}
+
+        # Get the token from C++ layer
+        cdef optional[CAuthenticationToken] token_opt = CAuthenticationTokenLoader.instance().GetToken()
+
+        if not token_opt.has_value() or token_opt.value().empty():
+            return {}
+
+        return {_AUTHORIZATION_HEADER_NAME: token_opt.value().ToAuthorizationHeaderValue().decode('utf-8')}

python/ray/tests/BUILD.bazel

@@ -578,6 +578,7 @@ py_test_module_list(
         "test_runtime_env_py_executable.py",
         "test_state_api_summary.py",
         "test_streaming_generator_regression.py",
+        "test_submission_client_auth.py",
         "test_system_metrics.py",
         "test_task_events_3.py",
         "test_task_metrics_reconstruction.py",