# Copyright 2025 The Feast Authors

#

# Licensed under the Apache License, Version 2.0 (the "License");

# you may not use this file except in compliance with the License.

# You may obtain a copy of the License at

#

# https://www.apache.org/licenses/LICENSE-2.0

#

# Unless required by applicable law or agreed to in writing, software

# distributed under the License is distributed on an "AS IS" BASIS,

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

# See the License for the specific language governing permissions and

# limitations under the License.

"""

Centralized Ray Initialization Module for Feast.

This module combines configuration management and initialization logic for a

complete, self-contained Ray setup system.

"""

import logging

import os

from enum import Enum

from typing import Any, Dict, List, Optional, Union

import ray

from ray.data.context import DatasetContext

logger = logging.getLogger(__name__)

class RayExecutionMode(Enum):

"""Ray execution modes supported by Feast."""

LOCAL = "local"

REMOTE = "remote"

KUBERAY = "kuberay"

class RayConfigManager:

"""

Manages Ray configuration and execution mode determination.

Supports three main scenarios:

1. Local Ray: Single-machine development and testing

2. Remote Ray: Connect to existing Ray standalone cluster

3. KubeRay: Ray on Kubernetes with CodeFlare SDK

The manager determines execution mode based on configuration precedence:

1. Environment variable FEAST_RAY_EXECUTION_MODE (highest)

2. KubeRay mode (use_kuberay=True or cluster_name specified)

3. Remote mode (ray_address specified)

4. Local mode (default fallback)

"""

def __init__(self, config: Optional[Union[Dict[str, Any], object]] = None):

"""

Initialize Ray configuration manager.

Args:

config: Ray configuration (RayOfflineStoreConfig, RayComputeEngineConfig, or dict)

"""

self.config = config or {}

self._execution_mode: Optional[RayExecutionMode] = None

self._codeflare_config: Optional[Dict[str, Any]] = None

def determine_execution_mode(self) -> RayExecutionMode:

"""

Determine the appropriate Ray execution mode based on configuration.

Precedence (highest to lowest):

1. Environment variable FEAST_RAY_EXECUTION_MODE (explicit override)

2. KubeRay mode (use_kuberay=True or cluster_name specified)

3. Remote mode (ray_address specified)

4. Local mode (default fallback)

Returns:

RayExecutionMode enum value

"""

if self._execution_mode is not None:

return self._execution_mode

# 1. Check environment variable override first (highest precedence)

env_mode = os.getenv("FEAST_RAY_EXECUTION_MODE", "").lower()

if env_mode in ["local", "remote", "kuberay"]:

self._execution_mode = RayExecutionMode(env_mode)

logger.info(

f"Ray execution mode set via FEAST_RAY_EXECUTION_MODE: {env_mode}"

)

return self._execution_mode

# 2. Check for KubeRay configuration (second highest precedence)

use_kuberay = self._get_config_value("use_kuberay")

# Check for cluster_name in kuberay_conf

kuberay_conf = self._get_config_value("kuberay_conf", {}) or {}

cluster_name = kuberay_conf.get("cluster_name")

# Environment variables can enable KubeRay

if os.getenv("FEAST_USE_KUBERAY", "").lower() == "true":

use_kuberay = True

if os.getenv("FEAST_RAY_CLUSTER_NAME"):

cluster_name = os.getenv("FEAST_RAY_CLUSTER_NAME")

# KubeRay takes precedence over remote/local if configured

if use_kuberay or cluster_name:

self._execution_mode = RayExecutionMode.KUBERAY

reason = []

if use_kuberay:

reason.append("use_kuberay=True")

if cluster_name:

reason.append(f"cluster_name='{cluster_name}'")

logger.info(f"Ray execution mode: KubeRay ({', '.join(reason)})")

return self._execution_mode

# 3. Check for remote Ray configuration (third precedence)

ray_address = self._get_config_value("ray_address") or os.getenv("RAY_ADDRESS")

if ray_address:

self._execution_mode = RayExecutionMode.REMOTE

logger.info(f"Ray execution mode: Remote (ray_address='{ray_address}')")

return self._execution_mode

# 4. Default to local Ray (lowest precedence - fallback)

self._execution_mode = RayExecutionMode.LOCAL

logger.info(

"Ray execution mode: Local (default - no KubeRay or remote configuration found)"

)

return self._execution_mode

def get_kuberay_config(self) -> Dict[str, Any]:

"""

Get KubeRay/CodeFlare SDK configuration.

Returns:

Dictionary of KubeRay configuration with passthrough settings

"""

if self._codeflare_config is not None:

return self._codeflare_config

# Get passthrough configuration from kuberay_conf first

kuberay_conf = self._get_config_value("kuberay_conf", {}) or {}

config = {

"use_kuberay": (

os.getenv("FEAST_USE_KUBERAY", "").lower() == "true"

or self._get_config_value("use_kuberay", False)

),

# Get values from kuberay_conf or environment variables

"cluster_name": (

os.getenv("FEAST_RAY_CLUSTER_NAME") or kuberay_conf.get("cluster_name")

),

"namespace": (

os.getenv("FEAST_RAY_NAMESPACE")

or kuberay_conf.get("namespace", "default")

),

}

# Add authentication configuration from kuberay_conf or environment variables

auth_token = (

os.getenv("FEAST_RAY_AUTH_TOKEN")

or os.getenv("RAY_AUTH_TOKEN")

or kuberay_conf.get("auth_token")

)

if auth_token:

config["auth_token"] = auth_token

# Add authentication server URL

auth_server = (

os.getenv("FEAST_RAY_AUTH_SERVER")

or os.getenv("RAY_AUTH_SERVER")

or kuberay_conf.get("auth_server")

)

if auth_server:

config["auth_server"] = auth_server

# Add skip TLS verification setting

skip_tls = os.getenv(

"FEAST_RAY_SKIP_TLS", ""

).lower() == "true" or kuberay_conf.get("skip_tls", False)

config["skip_tls"] = skip_tls

# Add any additional configuration from kuberay_conf

for key, value in kuberay_conf.items():

if key not in config: # Don't override already processed keys

config[key] = value

self._codeflare_config = config

return config

def _get_config_value(self, key: str, default: Any = None) -> Any:

"""

Get configuration value from config object or dictionary.

Args:

key: Configuration key

default: Default value if key not found

Returns:

Configuration value

"""

if hasattr(self.config, key):

return getattr(self.config, key)

elif isinstance(self.config, dict):

return self.config.get(key, default)

else:

return default

class StandardRayWrapper:

"""Wrapper for Ray Native operations."""

def read_parquet(self, path: Union[str, List[str]], **kwargs) -> Any:

"""Read parquet files using standard Ray."""

return ray.data.read_parquet(path, **kwargs)

def read_csv(self, path: Union[str, List[str]], **kwargs) -> Any:

"""Read CSV files using standard Ray."""

return ray.data.read_csv(path, **kwargs)

def from_pandas(self, df: Any) -> Any:

"""Create dataset from pandas DataFrame using standard Ray."""

return ray.data.from_pandas(df)

def from_arrow(self, table: Any) -> Any:

"""Create dataset from Arrow table using standard Ray."""

return ray.data.from_arrow(table)

class CodeFlareRayWrapper:

"""Wrapper for Ray operations on KubeRay clusters using CodeFlare SDK."""

def __init__(

self,

cluster_name: str,

namespace: str,

auth_token: str,

auth_server: str,

skip_tls: bool = False,

enable_logging: bool = False,

):

"""Initialize CodeFlare Ray wrapper with cluster connection parameters."""

self.cluster_name = cluster_name

self.namespace = namespace

self.auth_token = auth_token

self.auth_server = auth_server

self.skip_tls = skip_tls

self.enable_logging = enable_logging

self.cluster = None

# Authenticate and setup Ray connection

self._authenticate_codeflare()

self._setup_ray_connection()

def _authenticate_codeflare(self):

"""Authenticate with CodeFlare SDK."""

try:

from codeflare_sdk import TokenAuthentication

auth = TokenAuthentication(

token=self.auth_token,

server=self.auth_server,

skip_tls=self.skip_tls,

)

auth.login()

except Exception as e:

logger.error(f"CodeFlare authentication failed: {e}")

raise

def _setup_ray_connection(self):

"""Setup Ray connection to KubeRay cluster using TLS certificates."""

try:

from codeflare_sdk import generate_cert, get_cluster

self.cluster = get_cluster(

cluster_name=self.cluster_name, namespace=self.namespace

)

if self.cluster is None:

raise RuntimeError(

f"Failed to find KubeRay cluster '{self.cluster_name}' in namespace '{self.namespace}'"

)

generate_cert.generate_tls_cert(self.cluster_name, self.namespace)

generate_cert.export_env(self.cluster_name, self.namespace)

cluster_uri = self.cluster.cluster_uri()

runtime_env = {

"pip": ["feast"],

"env_vars": {"RAY_DISABLE_IMPORT_WARNING": "1"},

}

ray.shutdown()

logging_level = "INFO" if self.enable_logging else "ERROR"

ray.init(

address=cluster_uri,

ignore_reinit_error=True,

logging_level=logging_level,

log_to_driver=self.enable_logging,

runtime_env=runtime_env,

)

logger.info(f"Ray connected successfully to cluster: {self.cluster_name}")

except Exception as e:

logger.error(f"Ray connection failed: {e}")

raise

# Ray Data API methods - wrapped in @ray.remote to execute on cluster workers

def read_parquet(self, path: Union[str, List[str]], **kwargs) -> Any:

"""Read parquet files - runs remotely on KubeRay cluster workers."""

from feast.infra.ray_shared_utils import RemoteDatasetProxy

@ray.remote

def _remote_read_parquet(file_path, read_kwargs):

import ray

return ray.data.read_parquet(file_path, **read_kwargs)

return RemoteDatasetProxy(_remote_read_parquet.remote(path, kwargs))

def read_csv(self, path: Union[str, List[str]], **kwargs) -> Any:

"""Read CSV files - runs remotely on KubeRay cluster workers."""

from feast.infra.ray_shared_utils import RemoteDatasetProxy

@ray.remote

def _remote_read_csv(file_path, read_kwargs):

import ray

return ray.data.read_csv(file_path, **read_kwargs)

return RemoteDatasetProxy(_remote_read_csv.remote(path, kwargs))

def from_pandas(self, df: Any) -> Any:

"""Create dataset from pandas DataFrame - runs remotely on KubeRay cluster workers."""

from feast.infra.ray_shared_utils import RemoteDatasetProxy

@ray.remote

def _remote_from_pandas(dataframe):

import ray

return ray.data.from_pandas(dataframe)

return RemoteDatasetProxy(_remote_from_pandas.remote(df))

def from_arrow(self, table: Any) -> Any:

"""Create dataset from Arrow table - runs remotely on KubeRay cluster workers."""

from feast.infra.ray_shared_utils import RemoteDatasetProxy

@ray.remote

def _remote_from_arrow(arrow_table):

import ray

return ray.data.from_arrow(arrow_table)

return RemoteDatasetProxy(_remote_from_arrow.remote(table))

# Global state tracking

_ray_initialized = False

_ray_wrapper: Optional[Union[StandardRayWrapper, CodeFlareRayWrapper]] = None

def _suppress_ray_logging() -> None:

"""Suppress Ray and Ray Data logging completely."""

import warnings

# Suppress Ray warnings

warnings.filterwarnings("ignore", category=DeprecationWarning, module="ray")

warnings.filterwarnings("ignore", category=UserWarning, module="ray")

# Set environment variables to suppress Ray output

os.environ["RAY_DISABLE_IMPORT_WARNING"] = "1"

os.environ["RAY_SUPPRESS_UNVERIFIED_TLS_WARNING"] = "1"

os.environ["RAY_LOG_LEVEL"] = "ERROR"

os.environ["RAY_DATA_LOG_LEVEL"] = "ERROR"

os.environ["RAY_DISABLE_PROGRESS_BARS"] = "1"

# Suppress all Ray-related loggers

ray_loggers = [

"ray",

"ray.data",

"ray.data.dataset",

"ray.data.context",

"ray.data._internal.streaming_executor",

"ray.data._internal.execution",

"ray.data._internal",

"ray.tune",

"ray.serve",

"ray.util",

"ray._private",

]

for logger_name in ray_loggers:

logging.getLogger(logger_name).setLevel(logging.ERROR)

# Configure DatasetContext to disable progress bars

try:

ctx = DatasetContext.get_current()

ctx.enable_progress_bars = False

if hasattr(ctx, "verbose_progress"):

ctx.verbose_progress = False

except Exception:

pass # Ignore if Ray Data is not available

def _initialize_local_ray(config: Any, enable_logging: bool = False) -> None:

"""

Initialize Ray in local mode.

Args:

config: Configuration object (RayOfflineStoreConfig or RayComputeEngineConfig)

enable_logging: Whether to enable Ray logging

"""

logger.info("Initializing Ray in LOCAL mode")

ray_init_kwargs: Dict[str, Any] = {

"ignore_reinit_error": True,

"include_dashboard": False,

}

if enable_logging:

ray_init_kwargs.update(

{

"log_to_driver": True,

"logging_level": "INFO",

}

)

else:

ray_init_kwargs.update(

{

"log_to_driver": False,

"logging_level": "ERROR",

}

)

_suppress_ray_logging()

# Add local configuration

ray_init_kwargs.update(

{

"_node_ip_address": os.getenv("RAY_NODE_IP", "127.0.0.1"),

"num_cpus": os.cpu_count() or 4,

}

)

# Merge with user-provided ray_conf if available

if hasattr(config, "ray_conf") and config.ray_conf:

ray_init_kwargs.update(config.ray_conf)

# Initialize Ray

ray.init(**ray_init_kwargs)

# Configure DatasetContext

ctx = DatasetContext.get_current()

ctx.shuffle_strategy = "sort" # type: ignore

ctx.enable_tensor_extension_casting = False

# Log cluster info

if enable_logging:

cluster_resources = ray.cluster_resources()

logger.info(

f"Ray local cluster initialized with {cluster_resources.get('CPU', 0)} CPUs, "

f"{cluster_resources.get('memory', 0) / (1024**3):.1f}GB memory"

)

def _initialize_remote_ray(config: Any, enable_logging: bool = False) -> None:

"""

Initialize Ray in remote mode (connect to existing Ray cluster).

Args:

config: Configuration object with ray_address

enable_logging: Whether to enable Ray logging

"""

ray_address = getattr(config, "ray_address", None)

if not ray_address:

ray_address = os.getenv("RAY_ADDRESS")

if not ray_address:

raise ValueError("ray_address must be specified for remote Ray mode")

logger.info(f"Initializing Ray in REMOTE mode, connecting to: {ray_address}")

ray_init_kwargs: Dict[str, Any] = {

"address": ray_address,

"ignore_reinit_error": True,

"include_dashboard": False,

}

if enable_logging:

ray_init_kwargs.update(

{

"log_to_driver": True,

"logging_level": "INFO",

}

)

else:

ray_init_kwargs.update(

{

"log_to_driver": False,

"logging_level": "ERROR",

}

)

_suppress_ray_logging()

# Merge with user-provided ray_conf if available

if hasattr(config, "ray_conf") and config.ray_conf:

ray_init_kwargs.update(config.ray_conf)

# Initialize Ray

ray.init(**ray_init_kwargs)

# Configure DatasetContext

ctx = DatasetContext.get_current()

ctx.shuffle_strategy = "sort" # type: ignore

ctx.enable_tensor_extension_casting = False

# Log cluster info

if enable_logging:

cluster_resources = ray.cluster_resources()

logger.info(

f"Ray remote cluster initialized with {cluster_resources.get('CPU', 0)} CPUs, "

f"{cluster_resources.get('memory', 0) / (1024**3):.1f}GB memory"

)

def _initialize_kuberay(config: Any, enable_logging: bool = False) -> None:

"""

Initialize Ray in KubeRay mode using CodeFlare SDK.

Args:

config: Configuration object with KubeRay settings

enable_logging: Whether to enable Ray logging

"""

global _ray_wrapper

logger.info("Initializing Ray in KUBERAY mode using CodeFlare SDK")

if not enable_logging:

_suppress_ray_logging()

# Get KubeRay configuration

config_manager = RayConfigManager(config)

kuberay_config = config_manager.get_kuberay_config()

# Initialize CodeFlare Ray wrapper - this connects to the cluster

_ray_wrapper = CodeFlareRayWrapper(

cluster_name=kuberay_config["cluster_name"],

namespace=kuberay_config["namespace"],

auth_token=kuberay_config["auth_token"],

auth_server=kuberay_config["auth_server"],

skip_tls=kuberay_config.get("skip_tls", False),

enable_logging=enable_logging,

)

logger.info("KubeRay cluster connection established via CodeFlare SDK")

def ensure_ray_initialized(

config: Optional[Any] = None, force_reinit: bool = False

) -> None:

"""

Ensure Ray is initialized with appropriate configuration.

This is the main entry point for Ray initialization across all Feast components.

It automatically detects the execution mode and initializes Ray accordingly.

Args:

config: Configuration object (RayOfflineStoreConfig, RayComputeEngineConfig, or RepoConfig)

force_reinit: If True, reinitialize Ray even if already initialized

Raises:

ValueError: If configuration is invalid or required parameters are missing

"""

global _ray_initialized

# Check if already initialized

if _ray_initialized and not force_reinit:

logger.debug("Ray already initialized, skipping initialization")

return

# Extract Ray-specific config if RepoConfig is provided

ray_config = config

if config and hasattr(config, "offline_store"):

ray_config = config.offline_store

elif config and hasattr(config, "batch_engine"):

ray_config = config.batch_engine

# Determine enable_logging setting

enable_logging = (

getattr(ray_config, "enable_ray_logging", False) if ray_config else False

)

# Use RayConfigManager to determine execution mode

config_manager = RayConfigManager(ray_config)

execution_mode = config_manager.determine_execution_mode()

logger.info(f"Ray execution mode detected: {execution_mode.value}")

# Check if Ray is already initialized (from external source)

if ray.is_initialized() and not force_reinit:

logger.info("Ray is already initialized externally, using existing cluster")

# Configure DatasetContext even if Ray is already initialized

ctx = DatasetContext.get_current()

ctx.shuffle_strategy = "sort" # type: ignore

ctx.enable_tensor_extension_casting = False

if not enable_logging:

_suppress_ray_logging()

_ray_initialized = True

return

# Initialize based on execution mode

try:

if execution_mode == RayExecutionMode.KUBERAY:

_initialize_kuberay(ray_config, enable_logging)

elif execution_mode == RayExecutionMode.REMOTE:

_initialize_remote_ray(ray_config, enable_logging)

else: # LOCAL

_initialize_local_ray(ray_config, enable_logging)

_ray_initialized = True

logger.info(f"Ray initialized successfully in {execution_mode.value} mode")

except Exception as e:

logger.error(f"Failed to initialize Ray in {execution_mode.value} mode: {e}")

raise

def get_ray_wrapper() -> Union[StandardRayWrapper, CodeFlareRayWrapper]:

"""

Get the appropriate Ray wrapper based on current initialization mode.

Returns:

StandardRayWrapper for local/remote modes, CodeFlareRayWrapper for KubeRay mode

"""

global _ray_wrapper

if _ray_wrapper is None:

# Return a standard Ray wrapper for local/remote modes

_ray_wrapper = StandardRayWrapper()

return _ray_wrapper

def is_ray_initialized() -> bool:

"""Check if Ray has been initialized via this module."""

return _ray_initialized

def shutdown_ray() -> None:

"""Shutdown Ray and reset initialization state."""

global _ray_initialized, _ray_wrapper

if ray.is_initialized():

logger.info("Shutting down Ray")

ray.shutdown()

_ray_initialized = False

_ray_wrapper = None

logger.info("Ray shutdown complete")