[core] Cherry pick token authentication UX improvements (#58831)

Cherry pick: #58737

Signed-off-by: sampan <sampan@anyscale.com>
Signed-off-by: Edward Oakes <ed.nmi.oakes@gmail.com>
Co-authored-by: Sampan S Nayak <sampansnayak2@gmail.com>
Co-authored-by: sampan <sampan@anyscale.com>

Author: 3 people

doc/source/ray-core/api/exceptions.rst

@@ -40,3 +40,4 @@ Exceptions
     ray.exceptions.RaySystemError
     ray.exceptions.NodeDiedError
     ray.exceptions.UnserializableException
+    ray.exceptions.AuthenticationError

python/ray/_private/authentication/authentication_constants.py

@@ -1,27 +1,11 @@
-# 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
+# Authentication error messages
 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
+    "Token authentication is enabled but no authentication token was found."
 )
 
-# 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
-)
+TOKEN_INVALID_ERROR_MESSAGE = "Token authentication is enabled but the authentication token is invalid or incorrect."  # noqa: E501
 
+# HTTP header and cookie constants
 AUTHORIZATION_HEADER_NAME = "authorization"
 AUTHORIZATION_BEARER_PREFIX = "Bearer "
 RAY_AUTHORIZATION_HEADER_NAME = "x-ray-authorization"

python/ray/_private/authentication/authentication_token_setup.py

@@ -20,6 +20,7 @@
     AuthenticationTokenLoader,
     get_authentication_mode,
 )
+from ray.exceptions import AuthenticationError
 
 logger = logging.getLogger(__name__)
 
@@ -97,4 +98,6 @@ def ensure_token_if_auth_enabled(
             # Reload the cache so subsequent calls to token_loader read the new token.
             token_loader.reset_cache()
         else:
-            raise RuntimeError(TOKEN_AUTH_ENABLED_BUT_NO_TOKEN_FOUND_ERROR_MESSAGE)
+            raise AuthenticationError(
+                TOKEN_AUTH_ENABLED_BUT_NO_TOKEN_FOUND_ERROR_MESSAGE
+            )

python/ray/_private/authentication/http_token_authentication.py

@@ -117,13 +117,13 @@ def format_authentication_http_error(status: int, body: str) -> Optional[str]:
     if status == 401:
         return "Authentication required: {body}\n\n{details}".format(
             body=body,
-            details=authentication_constants.HTTP_REQUEST_MISSING_TOKEN_ERROR_MESSAGE,
+            details=authentication_constants.TOKEN_AUTH_ENABLED_BUT_NO_TOKEN_FOUND_ERROR_MESSAGE,
         )
 
     if status == 403:
         return "Authentication failed: {body}\n\n{details}".format(
             body=body,
-            details=authentication_constants.HTTP_REQUEST_INVALID_TOKEN_ERROR_MESSAGE,
+            details=authentication_constants.TOKEN_INVALID_ERROR_MESSAGE,
         )
 
     return None

python/ray/dashboard/client/src/authentication/TokenAuthenticationDialog.component.test.tsx

@@ -21,10 +21,10 @@ describe("TokenAuthenticationDialog", () => {
     );
 
     expect(
-      screen.getByText("Token Authentication Required"),
+      screen.getByText("Authentication Token Required"),
     ).toBeInTheDocument();
     expect(
-      screen.getByText(/token authentication is enabled for this cluster/i),
+      screen.getByText(/Token authentication is enabled/),
     ).toBeInTheDocument();
   });
 
@@ -38,10 +38,10 @@ describe("TokenAuthenticationDialog", () => {
     );
 
     expect(
-      screen.getByText("Token Authentication Required"),
+      screen.getByText("Authentication Token Required"),
     ).toBeInTheDocument();
     expect(
-      screen.getByText(/authentication token is invalid or has expired/i),
+      screen.getByText(/The existing authentication token is invalid/),
     ).toBeInTheDocument();
   });
 
@@ -71,7 +71,7 @@ describe("TokenAuthenticationDialog", () => {
       />,
     );
 
-    const input = screen.getByLabelText(/authentication token/i);
+    const input = screen.getByLabelText("Authentication Token");
     await user.type(input, "test-token-123");
 
     const submitButton = screen.getByRole("button", { name: /submit/i });
@@ -94,7 +94,7 @@ describe("TokenAuthenticationDialog", () => {
       />,
     );
 
-    const input = screen.getByLabelText(/authentication token/i);
+    const input = screen.getByLabelText("Authentication Token");
     await user.type(input, "test-token-123{Enter}");
 
     await waitFor(() => {
@@ -128,7 +128,7 @@ describe("TokenAuthenticationDialog", () => {
     const submitButton = screen.getByRole("button", { name: /submit/i });
     expect(submitButton).toBeDisabled();
 
-    const input = screen.getByLabelText(/authentication token/i);
+    const input = screen.getByLabelText("Authentication Token");
     await user.type(input, "test-token");
 
     expect(submitButton).not.toBeDisabled();
@@ -144,7 +144,7 @@ describe("TokenAuthenticationDialog", () => {
       />,
     );
 
-    const input = screen.getByLabelText(/authentication token/i);
+    const input = screen.getByLabelText("Authentication Token");
     await user.type(input, "secret-token");
 
     // Initially should be password type (hidden)
@@ -177,7 +177,7 @@ describe("TokenAuthenticationDialog", () => {
       />,
     );
 
-    const input = screen.getByLabelText(/authentication token/i);
+    const input = screen.getByLabelText("Authentication Token");
     await user.type(input, "test-token");
 
     const submitButton = screen.getByRole("button", { name: /submit/i });
@@ -200,7 +200,7 @@ describe("TokenAuthenticationDialog", () => {
 
     // Dialog should not be visible
     expect(
-      screen.queryByText("Token Authentication Required"),
+      screen.queryByText("Authentication Token Required"),
     ).not.toBeInTheDocument();
   });
 });

python/ray/dashboard/client/src/authentication/TokenAuthenticationDialog.tsx

@@ -68,11 +68,13 @@ export const TokenAuthenticationDialog: React.FC<TokenAuthenticationDialogProps>
       setShowToken(!showToken);
     };
 
-    // Different messages based on whether this is initial auth or re-auth
-    const title = "Token Authentication Required";
-    const message = hasExistingToken
-      ? "The authentication token is invalid or has expired. Please provide a valid authentication token."
-      : "Token authentication is enabled for this cluster. Please provide a valid authentication token.";
+    // Different messages based on whether this is initial auth or re-auth.
+    const title = "Authentication Token Required";
+    const message =
+      (hasExistingToken
+        ? "The existing authentication token is invalid."
+        : "Token authentication is enabled.") +
+      " Provide the matching authentication token for this cluster.\n- Local clusters: use `ray get-auth-token` to retrieve it.\n- Remote clusters: you must retrieve the token that was used when creating the cluster.\n\nSee: https://docs.ray.io/en/latest/ray-security/auth.html";
 
     return (
       <Dialog
@@ -88,7 +90,7 @@ export const TokenAuthenticationDialog: React.FC<TokenAuthenticationDialogProps>
         <DialogContent>
           <DialogContentText
             id="token-auth-dialog-description"
-            sx={{ marginBottom: 2 }}
+            sx={{ marginBottom: 2, whiteSpace: "pre-line" }}
           >
             {message}
           </DialogContentText>

python/ray/dashboard/modules/dashboard_sdk.py

@@ -26,6 +26,7 @@
 from ray._private.utils import split_address
 from ray.autoscaler._private.cli_logger import cli_logger
 from ray.dashboard.modules.job.common import uri_to_http_components
+from ray.exceptions import AuthenticationError
 from ray.util.annotations import DeveloperAPI, PublicAPI
 
 try:
@@ -323,7 +324,7 @@ def _do_request(
             response.status_code, response.text
         )
         if formatted_error:
-            raise RuntimeError(formatted_error)
+            raise AuthenticationError(formatted_error)
 
         return response
 

python/ray/exceptions.py

@@ -477,6 +477,30 @@ def __str__(self):
         return error_msg
 
 
+@PublicAPI
+class AuthenticationError(RayError):
+    """Indicates that an authentication error occurred.
+
+    Most commonly, this is caused by a missing or mismatching token set on the client
+    (e.g., a Ray CLI command interacting with a remote cluster).
+
+    Only applicable when `RAY_AUTH_MODE` is not set to `disabled`.
+    """
+
+    def __init__(self, message: str):
+        self.message = message
+
+        # Always hide traceback for cleaner output
+        self.__suppress_context__ = True
+        super().__init__(message)
+
+    def __str__(self) -> str:
+        return self.message + (
+            ". Ensure that you have `RAY_AUTH_MODE=token` set and the token for the cluster is available as the `RAY_AUTH_TOKEN` environment variable or a local file. "
+            "For more information, see: https://docs.ray.io/en/latest/ray-security/auth.html"
+        )
+
+
 @DeveloperAPI
 class UserCodeException(RayError):
     """Indicates that an exception occurred while executing user code.
@@ -963,4 +987,5 @@ def __str__(self):
     OufOfBandObjectRefSerializationException,
     RayCgraphCapacityExceeded,
     UnserializableException,
+    AuthenticationError,
 ]

python/ray/includes/common.pxd

@@ -133,6 +133,7 @@ cdef extern from "ray/common/status.h" namespace "ray" nogil:
         c_bool IsUnexpectedSystemExit()
         c_bool IsChannelError()
         c_bool IsChannelTimeoutError()
+        c_bool IsUnauthenticated()
 
         c_string ToString()
         c_string CodeAsString()

python/ray/includes/common.pxi

@@ -35,6 +35,7 @@ from ray.exceptions import (
     ActorDiedError,
     RayError,
     RaySystemError,
+    AuthenticationError,
     RayTaskError,
     ObjectStoreFullError,
     OutOfDiskError,
@@ -105,6 +106,8 @@ cdef int check_status(const CRayStatus& status) except -1 nogil:
         raise ValueError(message)
     elif status.IsIOError():
         raise IOError(message)
+    elif status.IsUnauthenticated():
+        raise AuthenticationError(message)
     elif status.IsRpcError():
         raise RpcError(message, rpc_code=status.rpc_code())
     elif status.IsIntentionalSystemExit():