[Core] handling auth_mode=token in python ray.init() calls (ray-project#57835)

## Description
builds atop of ray-project#58047, this pr
ensures the following when `auth_mode` is `token`:
calling `ray.init() `(without passing an existing cluster address) ->
check if token is present, generate and store in default path if not
present
calling `ray.init(address="xyz")` (connecting to an existing cluster) ->
check if token is present, raise exception if one is not present

---------

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

Author: 3 people

python/ray/_private/authentication/__init__.py

@@ -0,0 +1,6 @@
+import uuid
+
+
+# TODO: this is a placeholder for the actual authentication token generator. Will be replaced with a proper implementation.
+def generate_new_authentication_token() -> str:
+    return uuid.uuid4().hex

python/ray/_private/authentication/authentication_token_generator.py

@@ -0,0 +1,104 @@
+"""Authentication token setup for Ray.
+
+This module provides functions to generate and save authentication tokens
+for Ray's token-based authentication system. Token loading and caching is
+handled by the C++ AuthenticationTokenLoader.
+"""
+
+import logging
+from pathlib import Path
+from typing import Any, Dict, Optional
+
+from ray._private.authentication.authentication_token_generator import (
+    generate_new_authentication_token,
+)
+from ray._raylet import (
+    AuthenticationMode,
+    AuthenticationTokenLoader,
+    get_authentication_mode,
+)
+
+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.
+
+    Returns:
+        The newly generated authentication token.
+    """
+    # Generate a UUID-based token
+    token = generate_new_authentication_token()
+
+    token_path = _get_default_token_path()
+    try:
+        # Create directory if it doesn't exist
+        token_path.parent.mkdir(parents=True, exist_ok=True)
+
+        # Write token to file with explicit flush and fsync
+        with open(token_path, "w") as f:
+            f.write(token)
+
+        logger.info(f"Generated new authentication token and saved to {token_path}")
+    except Exception:
+        raise
+
+
+def _get_default_token_path() -> Path:
+    """Get the default token file path (~/.ray/auth_token).
+
+    Returns:
+        Path object pointing to ~/.ray/auth_token
+    """
+    return Path.home() / ".ray" / "auth_token"
+
+
+def ensure_token_if_auth_enabled(
+    system_config: Optional[Dict[str, Any]] = None, create_token_if_missing: bool = True
+) -> None:
+    """Check authentication settings and set up token resources if authentication is enabled.
+
+    Ray calls this early during ray.init() to do the following for token-based authentication:
+    1. Check whether you enabled token-based authentication.
+    2. Make sure a token is available if authentication is enabled.
+    3. Generate and save a default token for new local clusters if one doesn't already exist.
+
+    Args:
+        system_config: Ray raises an error if you set auth_mode in system_config instead of the environment.
+        create_token_if_missing: Generate a new token if one doesn't already exist.
+
+    Raises:
+        RuntimeError: Ray raises this error if authentication is enabled but no token is found when connecting
+            to an existing cluster.
+    """
+
+    # Check if you enabled token authentication.
+    if get_authentication_mode() != AuthenticationMode.TOKEN:
+        if (
+            system_config
+            and "auth_mode" in system_config
+            and system_config["auth_mode"] != "disabled"
+        ):
+            raise RuntimeError(
+                "Set authentication mode can only be set with the `RAY_auth_mode` environment variable, not using the system_config."
+            )
+        return
+
+    token_loader = AuthenticationTokenLoader.instance()
+
+    if not token_loader.has_token():
+        if create_token_if_missing:
+            # Generate a new token.
+            generate_and_save_token()
+
+            # 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)

python/ray/_private/authentication/authentication_token_setup.py

@@ -62,6 +62,9 @@
 from ray._common import ray_option_utils
 from ray._common.constants import RAY_WARN_BLOCKING_GET_INSIDE_ASYNC_ENV_VAR
 from ray._common.utils import load_class
+from ray._private.authentication.authentication_token_setup import (
+    ensure_token_if_auth_enabled,
+)
 from ray._private.client_mode_hook import client_mode_hook
 from ray._private.custom_types import TensorTransportEnum
 from ray._private.function_manager import FunctionActorManager
@@ -1865,6 +1868,9 @@ def sigterm_handler(signum, frame):
     if bootstrap_address is None:
         # In this case, we need to start a new cluster.
 
+        # Setup and verify authentication for new cluster
+        ensure_token_if_auth_enabled(_system_config, create_token_if_missing=True)
+
         # Don't collect usage stats in ray.init() unless it's a nightly wheel.
         from ray._common.usage import usage_lib
 
@@ -1952,6 +1958,9 @@ def sigterm_handler(signum, frame):
                 "an existing cluster."
             )
 
+        # Setup and verify authentication for connecting to existing cluster
+        ensure_token_if_auth_enabled(_system_config, create_token_if_missing=False)
+
         # In this case, we only need to connect the node.
         ray_params = ray._private.parameter.RayParams(
             node_ip_address=_node_ip_address,

python/ray/_private/worker.py

@@ -201,6 +201,7 @@ include "includes/metric.pxi"
 include "includes/setproctitle.pxi"
 include "includes/raylet_client.pxi"
 include "includes/gcs_subscriber.pxi"
+include "includes/rpc_token_authentication.pxi"
 
 import ray
 from ray.exceptions import (

python/ray/_raylet.pyx

@@ -0,0 +1,15 @@
+from libcpp cimport bool as c_bool
+
+cdef extern from "ray/rpc/authentication/authentication_mode.h" namespace "ray::rpc" nogil:
+    cdef enum CAuthenticationMode "ray::rpc::AuthenticationMode":
+        DISABLED "ray::rpc::AuthenticationMode::DISABLED"
+        TOKEN "ray::rpc::AuthenticationMode::TOKEN"
+
+    CAuthenticationMode GetAuthenticationMode()
+
+cdef extern from "ray/rpc/authentication/authentication_token_loader.h" namespace "ray::rpc" nogil:
+    cdef cppclass CAuthenticationTokenLoader "ray::rpc::AuthenticationTokenLoader":
+        @staticmethod
+        CAuthenticationTokenLoader& instance()
+        c_bool HasToken()
+        void ResetCache()

python/ray/includes/rpc_token_authentication.pxd

@@ -0,0 +1,45 @@
+from ray.includes.rpc_token_authentication cimport (
+    CAuthenticationMode,
+    GetAuthenticationMode,
+    CAuthenticationTokenLoader,
+)
+
+
+# Authentication mode enum exposed to Python
+class AuthenticationMode:
+    DISABLED = CAuthenticationMode.DISABLED
+    TOKEN = CAuthenticationMode.TOKEN
+
+
+def get_authentication_mode():
+    """Get the current authentication mode.
+
+    Returns:
+        AuthenticationMode enum value (DISABLED or TOKEN)
+    """
+    return GetAuthenticationMode()
+
+
+class AuthenticationTokenLoader:
+    """Python wrapper for C++ AuthenticationTokenLoader singleton."""
+
+    @staticmethod
+    def instance():
+        """Get the singleton instance (returns a wrapper for convenience)."""
+        return AuthenticationTokenLoader()
+
+    def has_token(self):
+        """Check if an authentication token exists without crashing.
+
+        Returns:
+            bool: True if a token exists, False otherwise
+        """
+        return CAuthenticationTokenLoader.instance().HasToken()
+
+    def reset_cache(self):
+        """Reset the C++ authentication token cache.
+
+        This forces the token loader to reload the token from environment
+        variables or files on the next request.
+        """
+        CAuthenticationTokenLoader.instance().ResetCache()

python/ray/includes/rpc_token_authentication.pxi

@@ -368,6 +368,7 @@ py_test_module_list(
         "test_task_metrics.py",
         "test_tempdir.py",
         "test_tls_auth.py",
+        "test_token_auth_integration.py",
         "test_traceback.py",
         "test_worker_capping.py",
         "test_worker_state.py",