# 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, ShuffleStrategy

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 read_json(self, path: Union[str, List[str]], **kwargs) -> Any:

"""Read JSON/JSONL file(s)."""

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

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

"""Read plain-text file(s)."""

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

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

"""Read image files (PNG, JPEG, …) from a directory as numpy arrays."""

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

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

"""Read arbitrary binary files; each row has 'path' and 'bytes' columns."""

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

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

"""Read TFRecord file(s)."""

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

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

"""Read WebDataset tar shard(s)."""

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

def read_mongo(self, uri: str, database: str, collection: str, **kwargs) -> Any:

"""Read from a MongoDB collection."""

return ray.data.read_mongo(

uri=uri, database=database, collection=collection, **kwargs

)

def read_sql(self, sql: str, connection_url: str, **kwargs) -> Any:

"""Read from a SQL database. Builds the connection factory from connection_url."""

import sqlalchemy

def _connection_factory():

# raw_connection() returns a DB API2-compliant connection (with .cursor())

# from the underlying driver. engine.connect() returns a SQLAlchemy 2.0

# Connection object which does NOT expose .cursor(), failing Ray Data's

# DB API2 compliance check.

return sqlalchemy.create_engine(connection_url).raw_connection()

return ray.data.read_sql(sql, _connection_factory, **kwargs)

def from_huggingface(

self, dataset_name: str, split: str = "train", **kwargs

) -> Any:

"""Load a HuggingFace dataset and convert to a Ray Dataset."""

from datasets import load_dataset

hf_dataset = load_dataset(dataset_name, split=split, **kwargs)

return ray.data.from_huggingface(hf_dataset)

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,

num_gpus: float = 0,

worker_task_options: Optional[Dict[str, Any]] = None,

runtime_env: Optional[Dict[str, Any]] = None,

):

"""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.num_gpus = num_gpus

self.worker_task_options = worker_task_options or {}

self.extra_runtime_env = runtime_env or {}

self.cluster = None

# Authenticate and setup Ray connection

self._authenticate_codeflare()

self._setup_ray_connection()

def _authenticate_codeflare(self):

"""Authenticate with CodeFlare SDK. Skipped for in-cluster pods with no explicit token."""

if not self.auth_token or not self.auth_server:

logger.info(

"No auth_token/auth_server provided; assuming in-cluster auth via service account"

)

return

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"},

}

if self.extra_runtime_env:

extra_pip = self.extra_runtime_env.get("pip", [])

if extra_pip:

runtime_env["pip"] = list(

dict.fromkeys(runtime_env["pip"] + extra_pip)

)

extra_env_vars = self.extra_runtime_env.get("env_vars", {})

if extra_env_vars:

runtime_env["env_vars"].update(extra_env_vars)

for k, v in self.extra_runtime_env.items():

if k not in ("pip", "env_vars"):

runtime_env[k] = v

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

def _get_task_options(self, include_gpu: bool = False) -> Dict[str, Any]:

"""Build Ray .options() kwargs for a remote task dispatch.

Always includes the full worker_task_options passthrough dict.

num_gpus is merged in only when include_gpu=True, because I/O

operations (read_parquet, read_csv, from_pandas, from_arrow) are

pure data-movement and must not consume GPU slots. Pass

include_gpu=True only for compute/transformation tasks.

"""

opts: Dict[str, Any] = dict(self.worker_task_options)

if include_gpu and self.num_gpus:

opts["num_gpus"] = self.num_gpus

else:

# Guarantee no GPU slot is consumed for I/O tasks even if the

# caller placed num_gpus inside worker_task_options directly.

opts.pop("num_gpus", None)

return opts

# 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(file_path, read_kwargs):

import ray

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

opts = self._get_task_options()

remote_fn = _remote.options(**opts) if opts else _remote

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

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

"""Read CSV files - dispatched via @ray.remote to cluster workers."""

from feast.infra.ray_shared_utils import RemoteDatasetProxy

@ray.remote

def _remote(file_path, read_kwargs):

import ray

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

opts = self._get_task_options()

remote_fn = _remote.options(**opts) if opts else _remote

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

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

"""Read JSON/JSONL files - dispatched via @ray.remote to cluster workers."""

from feast.infra.ray_shared_utils import RemoteDatasetProxy

@ray.remote

def _remote(file_path, read_kwargs):

import ray

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

opts = self._get_task_options()

remote_fn = _remote.options(**opts) if opts else _remote

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

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

"""Read plain-text files - dispatched via @ray.remote to cluster workers."""

from feast.infra.ray_shared_utils import RemoteDatasetProxy

@ray.remote

def _remote(file_path, read_kwargs):

import ray

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

opts = self._get_task_options()

remote_fn = _remote.options(**opts) if opts else _remote

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

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

"""Read image directory - dispatched via @ray.remote to cluster workers."""

from feast.infra.ray_shared_utils import RemoteDatasetProxy

@ray.remote

def _remote(file_path, read_kwargs):

import ray

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

opts = self._get_task_options()

remote_fn = _remote.options(**opts) if opts else _remote

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

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

"""Read binary files - dispatched via @ray.remote to cluster workers."""

from feast.infra.ray_shared_utils import RemoteDatasetProxy

@ray.remote

def _remote(file_path, read_kwargs):

import ray

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

opts = self._get_task_options()

remote_fn = _remote.options(**opts) if opts else _remote

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

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

"""Read TFRecord files - dispatched via @ray.remote to cluster workers."""

from feast.infra.ray_shared_utils import RemoteDatasetProxy

@ray.remote

def _remote(file_path, read_kwargs):

import ray

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

opts = self._get_task_options()

remote_fn = _remote.options(**opts) if opts else _remote

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

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

"""Read WebDataset shards - dispatched via @ray.remote to cluster workers."""

from feast.infra.ray_shared_utils import RemoteDatasetProxy

@ray.remote

def _remote(file_path, read_kwargs):

import ray

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

opts = self._get_task_options()

remote_fn = _remote.options(**opts) if opts else _remote

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

def read_mongo(self, uri: str, database: str, collection: str, **kwargs) -> Any:

"""Read from MongoDB - dispatched via @ray.remote to cluster workers."""

from feast.infra.ray_shared_utils import RemoteDatasetProxy

@ray.remote

def _remote(mongo_uri, db, col, read_kwargs):

import ray

return ray.data.read_mongo(

uri=mongo_uri, database=db, collection=col, **read_kwargs

)

opts = self._get_task_options()

remote_fn = _remote.options(**opts) if opts else _remote

return RemoteDatasetProxy(remote_fn.remote(uri, database, collection, kwargs))

def read_sql(self, sql: str, connection_url: str, **kwargs) -> Any:

"""Read from SQL database - dispatched via @ray.remote to cluster workers.

Accepts connection_url (str) instead of a connection_factory callable so

the argument is serialisable across the Ray object store boundary.

"""

from feast.infra.ray_shared_utils import RemoteDatasetProxy

@ray.remote

def _remote(sql_query, conn_url, read_kwargs):

import ray

import sqlalchemy

def _factory():

return sqlalchemy.create_engine(conn_url).raw_connection()

return ray.data.read_sql(sql_query, _factory, **read_kwargs)

opts = self._get_task_options()

remote_fn = _remote.options(**opts) if opts else _remote

return RemoteDatasetProxy(remote_fn.remote(sql, connection_url, kwargs))

def from_huggingface(

self, dataset_name: str, split: str = "train", **kwargs

) -> Any:

"""Load a HuggingFace dataset on a remote Ray worker and return a Ray Dataset.

The dataset is loaded directly on the worker to avoid serializing

memory-mapped HF Dataset objects across the network and to keep

driver memory usage near zero.

"""

from feast.infra.ray_shared_utils import RemoteDatasetProxy

@ray.remote

def _remote(ds_name, ds_split, extra_kwargs):

import ray

from datasets import load_dataset

hf_dataset = load_dataset(ds_name, split=ds_split, **extra_kwargs)

return ray.data.from_huggingface(hf_dataset)

opts = self._get_task_options()

remote_fn = _remote.options(**opts) if opts else _remote

return RemoteDatasetProxy(remote_fn.remote(dataset_name, split, kwargs))

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

"""Create dataset from pandas DataFrame - dispatched via @ray.remote."""

from feast.infra.ray_shared_utils import RemoteDatasetProxy

@ray.remote

def _remote(dataframe):

import ray

return ray.data.from_pandas(dataframe)

opts = self._get_task_options()

remote_fn = _remote.options(**opts) if opts else _remote

return RemoteDatasetProxy(remote_fn.remote(df))

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

"""Create dataset from Arrow table - dispatched via @ray.remote."""

from feast.infra.ray_shared_utils import RemoteDatasetProxy

@ray.remote

def _remote(arrow_table):

import ray

return ray.data.from_arrow(arrow_table)

opts = self._get_task_options()

remote_fn = _remote.options(**opts) if opts else _remote

return RemoteDatasetProxy(remote_fn.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 = ShuffleStrategy.SORT_SHUFFLE_PULL_BASED

ctx.enable_tensor_extension_casting = False

if hasattr(ctx, "enable_arrow_backed_pandas_conversion"):

ctx.enable_arrow_backed_pandas_conversion = 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 = ShuffleStrategy.SORT_SHUFFLE_PULL_BASED

ctx.enable_tensor_extension_casting = False

if hasattr(ctx, "enable_arrow_backed_pandas_conversion"):

ctx.enable_arrow_backed_pandas_conversion = 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.get("auth_token", ""),

auth_server=kuberay_config.get("auth_server", ""),

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

enable_logging=enable_logging,

num_gpus=getattr(config, "num_gpus", 0) or 0,

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

runtime_env=kuberay_config.get("runtime_env"),

)

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 = ShuffleStrategy.SORT_SHUFFLE_PULL_BASED

ctx.enable_tensor_extension_casting = False

if hasattr(ctx, "enable_arrow_backed_pandas_conversion"):

ctx.enable_arrow_backed_pandas_conversion = 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")