chore: import upstream snapshot with attribution

python/ray/__init__.py

@@ -0,0 +1,299 @@
+# isort: skip_file
+from ray._private import log  # isort: skip # noqa: F401
+import logging
+import os
+import sys
+from typing import TYPE_CHECKING
+
+log.generate_logging_config()
+logger = logging.getLogger(__name__)
+
+
+def _configure_system():
+    import os
+    import platform
+    import sys
+
+    """Wraps system configuration to avoid 'leaking' variables into ray."""
+
+    # Sanity check pickle5 if it has been installed.
+    if "pickle5" in sys.modules:
+        if sys.version_info >= (3, 8):
+            logger.warning(
+                "Package pickle5 becomes unnecessary in Python 3.8 and above. "
+                "Its presence may confuse libraries including Ray. "
+                "Please uninstall the package."
+            )
+
+        import importlib.metadata
+
+        try:
+            version_str = importlib.metadata.version("pickle5")
+            version = tuple(int(n) for n in version_str.split("."))
+            if version < (0, 0, 10):
+                logger.warning(
+                    "Although not used by Ray, a version of pickle5 that leaks memory "
+                    "is found in the environment. Please run 'pip install pickle5 -U' "
+                    "to upgrade."
+                )
+        except importlib.metadata.PackageNotFoundError:
+            logger.warning(
+                "You are using the 'pickle5' module, but "
+                "the exact version is unknown (possibly carried as "
+                "an internal component by another module). Please "
+                "make sure you are using pickle5 >= 0.0.10 because "
+                "previous versions may leak memory."
+            )
+
+    # Importing psutil. Must be before ray._raylet is
+    # initialized.
+    thirdparty_files = os.path.join(
+        os.path.abspath(os.path.dirname(__file__)), "thirdparty_files"
+    )
+    sys.path.insert(0, thirdparty_files)
+
+    if (
+        platform.system() == "Linux"
+        and "Microsoft".lower() in platform.release().lower()
+    ):
+        from ray._private import compat  # noqa: E402
+
+        compat.patch_psutil()
+
+    # Expose ray ABI symbols which may be dependent by other shared
+    # libraries such as _streaming.so. See BUILD.bazel:_raylet
+    python_shared_lib_suffix = ".so" if sys.platform != "win32" else ".pyd"
+    so_path = os.path.join(
+        os.path.dirname(__file__), "_raylet" + python_shared_lib_suffix
+    )
+    if os.path.exists(so_path):
+        import ctypes
+        from ctypes import CDLL
+
+        CDLL(so_path, ctypes.RTLD_GLOBAL)
+
+
+_configure_system()
+# Delete configuration function.
+del _configure_system
+
+from ray import _version  # noqa: E402
+
+__commit__ = _version.commit
+__version__ = _version.version
+
+import ray._raylet  # noqa: E402
+
+from ray._raylet import (  # noqa: E402,F401
+    ActorClassID,
+    ActorID,
+    NodeID,
+    Config as _Config,
+    JobID,
+    WorkerID,
+    FunctionID,
+    ObjectID,
+    ObjectRef,
+    ObjectRefGenerator,
+    DynamicObjectRefGenerator,
+    TaskID,
+    UniqueID,
+    Language,
+    PlacementGroupID,
+    ClusterID,
+)
+
+_config = _Config()
+
+from ray._private.state import (  # noqa: E402,F401
+    nodes,
+    timeline,
+    cluster_resources,
+    available_resources,
+)
+from ray._private.worker import (  # noqa: E402,F401
+    SCRIPT_MODE,
+    WORKER_MODE,
+    RESTORE_WORKER_MODE,
+    SPILL_WORKER_MODE,
+    cancel,
+    get,
+    get_actor,
+    get_gpu_ids,
+    init,
+    is_initialized,
+    put,
+    kill,
+    remote,
+    shutdown,
+    wait,
+)
+
+from ray._private.ray_logging.logging_config import LoggingConfig  # noqa: E402
+
+# We import ray.actor because some code is run in actor.py which initializes
+# some functions in the worker.
+import ray.actor  # noqa: E402,F401
+from ray.actor import method  # noqa: E402,F401
+
+# TODO(qwang): We should remove this exporting in Ray2.0.
+from ray.cross_language import java_function, java_actor_class  # noqa: E402,F401
+from ray.runtime_context import get_runtime_context  # noqa: E402,F401
+from ray import internal  # noqa: E402,F401
+from ray import util  # noqa: E402,F401
+from ray import _private  # noqa: E402,F401
+
+# We import ClientBuilder so that modules can inherit from `ray.ClientBuilder`.
+from ray.client_builder import client, ClientBuilder  # noqa: E402,F401
+
+
+class _DeprecationWrapper:
+    def __init__(self, name, real_worker):
+        self._name = name
+        self._real_worker = real_worker
+        self._warned = set()
+
+    def __getattr__(self, attr):
+        value = getattr(self._real_worker, attr)
+        if attr not in self._warned:
+            self._warned.add(attr)
+            logger.warning(
+                f"DeprecationWarning: `ray.{self._name}.{attr}` is a private "
+                "attribute and access will be removed in a future Ray version."
+            )
+        return value
+
+
+# TODO(ekl) remove this entirely after 3rd party libraries are all migrated.
+worker = _DeprecationWrapper("worker", ray._private.worker)
+ray_constants = _DeprecationWrapper("ray_constants", ray._private.ray_constants)
+serialization = _DeprecationWrapper("serialization", ray._private.serialization)
+state = _DeprecationWrapper("state", ray._private.state)
+
+
+# Pulic Ray APIs
+__all__ = [
+    "__version__",
+    "_config",
+    "get_runtime_context",
+    "autoscaler",
+    "available_resources",
+    "cancel",
+    "client",
+    "ClientBuilder",
+    "cluster_resources",
+    "get",
+    "get_actor",
+    "get_gpu_ids",
+    "init",
+    "is_initialized",
+    "java_actor_class",
+    "java_function",
+    "cpp_function",
+    "kill",
+    "Language",
+    "method",
+    "nodes",
+    "put",
+    "remote",
+    "shutdown",
+    "timeline",
+    "wait",
+    "SCRIPT_MODE",
+    "WORKER_MODE",
+    "LoggingConfig",
+]
+
+# Public APIs that should automatically trigger ray.init().
+AUTO_INIT_APIS = {
+    "cancel",
+    "get",
+    "get_actor",
+    "get_gpu_ids",
+    "kill",
+    "put",
+    "wait",
+    "get_runtime_context",
+}
+
+# Public APIs that should not automatically trigger ray.init().
+NON_AUTO_INIT_APIS = {
+    "ClientBuilder",
+    "Language",
+    "SCRIPT_MODE",
+    "WORKER_MODE",
+    "__version__",
+    "_config",
+    "autoscaler",
+    "available_resources",
+    "client",
+    "cluster_resources",
+    "cpp_function",
+    "init",
+    "is_initialized",
+    "java_actor_class",
+    "java_function",
+    "method",
+    "nodes",
+    "remote",
+    "shutdown",
+    "timeline",
+    "LoggingConfig",
+}
+
+assert set(__all__) == AUTO_INIT_APIS | NON_AUTO_INIT_APIS
+from ray._private.auto_init_hook import wrap_auto_init_for_all_apis  # noqa: E402
+
+wrap_auto_init_for_all_apis(AUTO_INIT_APIS)
+del wrap_auto_init_for_all_apis
+
+# Subpackages
+__all__ += [
+    "actor",
+    "autoscaler",
+    "data",
+    "internal",
+    "util",
+    "widgets",
+    "workflow",
+]
+
+# ID types
+__all__ += [
+    "ActorClassID",
+    "ActorID",
+    "NodeID",
+    "JobID",
+    "WorkerID",
+    "FunctionID",
+    "ObjectID",
+    "ObjectRef",
+    "ObjectRefGenerator",
+    "DynamicObjectRefGenerator",
+    "TaskID",
+    "UniqueID",
+    "PlacementGroupID",
+]
+
+
+# Delay importing of expensive, isolated subpackages. Note that for proper type
+# checking support these imports must be kept in sync between type checking and
+# runtime behavior.
+if TYPE_CHECKING:
+    from ray import autoscaler
+    from ray import data
+    from ray import workflow
+else:
+
+    def __getattr__(name: str):
+        import importlib
+
+        if name in ["data", "workflow", "autoscaler"]:
+            return importlib.import_module("." + name, __name__)
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+del os
+del logging
+del sys
+del TYPE_CHECKING

python/ray/_common/README.md

@@ -0,0 +1,7 @@
+# Common Utilities Shared Across the Libraries
+
+This directory contains logic shared across Ray Core and the native libraries.
+
+- All dependencies by the libraries on non-public APIs in the repo should live here. Libraries should _not_ depend on `ray._private`.
+- Interfaces exposed in this directory should be treated similarly to a "developer API."
+- End users and external libraries not inside the Ray repo should not depend on any code in `ray._common` (the same as `ray._private`).

python/ray/_common/constants.py

@@ -0,0 +1,7 @@
+# Prefix for the node id resource that is automatically added to each node.
+# For example, a node may have id `node:172.23.42.1`.
+NODE_ID_PREFIX = "node:"
+# The system resource that head node has.
+HEAD_NODE_RESOURCE_NAME = NODE_ID_PREFIX + "__internal_head__"
+
+RAY_WARN_BLOCKING_GET_INSIDE_ASYNC_ENV_VAR = "RAY_WARN_BLOCKING_GET_INSIDE_ASYNC"

python/ray/_common/deprecation.py

@@ -0,0 +1,181 @@
+import inspect
+import logging
+from functools import wraps
+from typing import Any, Callable, Optional, TypeVar, Union, cast, overload
+
+from ray.util import log_once
+from ray.util.annotations import _mark_annotated
+
+# TypeVar for preserving function/class signatures through decorators.
+# Note: These decorators also accept properties, but we use Callable for the
+# common case. Properties work at runtime but won't get full type inference.
+F = TypeVar("F", bound=Callable[..., Any])
+
+logger = logging.getLogger(__name__)
+
+# A constant to use for any configuration that should be deprecated
+# (to check, whether this config has actually been assigned a proper value or
+# not).
+DEPRECATED_VALUE = -1
+
+
+def deprecation_warning(
+    old: str,
+    new: Optional[str] = None,
+    *,
+    help: Optional[str] = None,
+    error: Optional[Union[bool, type[Exception]]] = None,
+    stacklevel: int = 2,
+) -> None:
+    """Warns (via the `logger` object) or throws a deprecation warning/error.
+
+    Args:
+        old: A description of the "thing" that is to be deprecated.
+        new: A description of the new "thing" that replaces it.
+        help: An optional help text to tell the user, what to
+            do instead of using `old`.
+        error: Whether or which exception to raise. If True, raise ValueError.
+            If False, just warn. If `error` is-a subclass of Exception,
+            raise that Exception.
+        stacklevel: The stacklevel to use for the warning message.
+            Use 2 to point to where this function is called, 3+ to point
+            further up the stack.
+
+    Raises:
+        ValueError: If `error=True`.
+        Exception: Of type `error`, iff `error` is a sub-class of `Exception`.
+    """
+    msg = "`{}` has been deprecated.{}".format(
+        old, (" Use `{}` instead.".format(new) if new else f" {help}" if help else "")
+    )
+
+    if error:
+        if not isinstance(error, bool) and issubclass(error, Exception):
+            # error is an Exception
+            raise error(msg)
+        else:
+            # error is a boolean, construct ValueError ourselves
+            raise ValueError(msg)
+    else:
+        logger.warning(
+            "DeprecationWarning: " + msg + " This will raise an error in the future!",
+            stacklevel=stacklevel,
+        )
+
+
+@overload
+def Deprecated(
+    old: None = None,
+    *,
+    new: Optional[str] = None,
+    help: Optional[str] = None,
+    error: Union[bool, type[Exception]],
+) -> Callable[[F], F]:
+    ...
+
+
+@overload
+def Deprecated(
+    old: str,
+    *,
+    new: Optional[str] = None,
+    help: Optional[str] = None,
+    error: Union[bool, type[Exception]],
+) -> Callable[[F], F]:
+    ...
+
+
+def Deprecated(
+    old: Optional[str] = None,
+    *,
+    new: Optional[str] = None,
+    help: Optional[str] = None,
+    error: Union[bool, type[Exception]],
+) -> Callable[[F], F]:
+    """Decorator for documenting a deprecated class, method, or function.
+
+    Automatically adds a `deprecation.deprecation_warning(old=...,
+    error=False)` to not break existing code at this point to the decorated
+    class' constructor, method, or function.
+
+    In a next major release, this warning should then be made an error
+    (by setting error=True), which means at this point that the
+    class/method/function is no longer supported, but will still inform
+    the user about the deprecation event.
+
+    In a further major release, the class, method, function should be erased
+    entirely from the codebase.
+
+
+    .. testcode::
+        :skipif: True
+
+        from ray._common.deprecation import Deprecated
+        # Deprecated class: Patches the constructor to warn if the class is
+        # used.
+        @Deprecated(new="NewAndMuchCoolerClass", error=False)
+        class OldAndUncoolClass:
+            ...
+
+        # Deprecated class method: Patches the method to warn if called.
+        class StillCoolClass:
+            ...
+            @Deprecated(new="StillCoolClass.new_and_much_cooler_method()",
+                        error=False)
+            def old_and_uncool_method(self, uncool_arg):
+                ...
+
+        # Deprecated function: Patches the function to warn if called.
+        @Deprecated(new="new_and_much_cooler_function", error=False)
+        def old_and_uncool_function(*uncool_args):
+            ...
+    """
+
+    def _inner(obj: F) -> F:
+        # A deprecated class.
+        if inspect.isclass(obj):
+            # Patch the class' init method to raise the warning/error.
+            obj_init = obj.__init__
+
+            def patched_init(*args, **kwargs):
+                if log_once(old or obj.__name__):
+                    deprecation_warning(
+                        old=old or obj.__name__,
+                        new=new,
+                        help=help,
+                        error=error,
+                        stacklevel=3,
+                    )
+                return obj_init(*args, **kwargs)
+
+            obj.__init__ = patched_init
+            _mark_annotated(obj)
+            # Return the patched class (with the warning/error when
+            # instantiated).
+            return obj
+
+        # A deprecated class method or function.
+        # Patch with the warning/error at the beginning.
+        def _ctor(*args, **kwargs):
+            if log_once(old or obj.__name__):
+                deprecation_warning(
+                    old=old or obj.__name__,
+                    new=new,
+                    help=help,
+                    error=error,
+                    stacklevel=3,
+                )
+            # Call the deprecated method/function.
+            return obj(*args, **kwargs)
+
+        # Only apply @wraps for actual callables, not properties/descriptors.
+        # Setting __wrapped__ on a property causes inspect.unwrap() to return
+        # the property, which breaks inspect.signature() in the tracing helper.
+        if callable(obj):
+            _ctor = wraps(obj)(_ctor)
+
+        # Return the patched class method/function.
+        return cast(F, _ctor)
+
+    # Return the prepared decorator.
+    return _inner

python/ray/_common/filters.py

@@ -0,0 +1,55 @@
+import logging
+from typing import Any, Dict
+
+import ray
+from ray._common.logging_constants import LogKey
+
+
+class CoreContextFilter(logging.Filter):
+    TASK_LEVEL_LOG_KEYS = [
+        LogKey.TASK_ID.value,
+        LogKey.TASK_NAME.value,
+        LogKey.TASK_FUNCTION_NAME.value,
+    ]
+
+    @classmethod
+    def get_ray_core_logging_context(cls) -> Dict[str, Any]:
+        """
+        Get the ray core logging context as a dict.
+        Only use this function if you need include the attributes to the log record
+        yourself by bypassing the filter.
+        """
+        if not ray.is_initialized():
+            # There is no additional context if ray is not initialized
+            return {}
+
+        runtime_context = ray.get_runtime_context()
+        ray_core_logging_context = {
+            LogKey.JOB_ID.value: runtime_context.get_job_id(),
+            LogKey.WORKER_ID.value: runtime_context.get_worker_id(),
+            LogKey.NODE_ID.value: runtime_context.get_node_id(),
+        }
+        if runtime_context.worker.mode == ray.WORKER_MODE:
+            ray_core_logging_context[
+                LogKey.ACTOR_ID.value
+            ] = runtime_context.get_actor_id()
+            ray_core_logging_context[
+                LogKey.TASK_ID.value
+            ] = runtime_context.get_task_id()
+            ray_core_logging_context[
+                LogKey.TASK_NAME.value
+            ] = runtime_context.get_task_name()
+            ray_core_logging_context[
+                LogKey.TASK_FUNCTION_NAME.value
+            ] = runtime_context.get_task_function_name()
+            ray_core_logging_context[
+                LogKey.ACTOR_NAME.value
+            ] = runtime_context.get_actor_name()
+        return ray_core_logging_context
+
+    def filter(self, record):
+        context = self.get_ray_core_logging_context()
+        for key, value in context.items():
+            if value is not None and value != "":
+                setattr(record, key, value)
+        return True

python/ray/_common/formatters.py

@@ -0,0 +1,120 @@
+import json
+import logging
+from abc import ABC, abstractmethod
+from typing import Any, Dict, List
+
+from ray._common.logging_constants import (
+    LOGGER_FLATTEN_KEYS,
+    LOGRECORD_STANDARD_ATTRS,
+    LogKey,
+)
+from ray._private.log import INTERNAL_TIMESTAMP_LOG_KEY
+from ray._private.ray_constants import LOGGER_FORMAT
+
+
+def _append_flatten_attributes(formatted_attrs: Dict[str, Any], key: str, value: Any):
+    """Flatten the dictionary values for special keys and append the values in place.
+
+    If the key is in `LOGGER_FLATTEN_KEYS`, the value will be flattened and appended
+    to the `formatted_attrs` dictionary. Otherwise, the key-value pair will be appended
+    directly.
+    """
+    if key in LOGGER_FLATTEN_KEYS:
+        if not isinstance(value, dict):
+            raise ValueError(
+                f"Expected a dictionary passing into {key}, but got {type(value)}"
+            )
+        for k, v in value.items():
+            if k in formatted_attrs:
+                raise KeyError(f"Found duplicated key in the log record: {k}")
+            formatted_attrs[k] = v
+    else:
+        formatted_attrs[key] = value
+
+
+class AbstractFormatter(logging.Formatter, ABC):
+    def __init__(self, fmt=None, datefmt=None, style="%", validate=True) -> None:
+        super().__init__(fmt, datefmt, style, validate)
+        self._additional_log_standard_attrs = []
+
+    def set_additional_log_standard_attrs(
+        self, additional_log_standard_attrs: List[str]
+    ) -> None:
+        self._additional_log_standard_attrs = additional_log_standard_attrs
+
+    @property
+    def additional_log_standard_attrs(self) -> List[str]:
+        return self._additional_log_standard_attrs
+
+    def generate_record_format_attrs(
+        self,
+        record: logging.LogRecord,
+        exclude_default_standard_attrs,
+    ) -> dict:
+        record_format_attrs = {}
+
+        # If `exclude_default_standard_attrs` is False, include the standard attributes.
+        # Otherwise, include only Ray and user-provided context.
+        if not exclude_default_standard_attrs:
+            record_format_attrs.update(
+                {
+                    LogKey.ASCTIME.value: self.formatTime(record),
+                    LogKey.LEVELNAME.value: record.levelname,
+                    LogKey.MESSAGE.value: record.getMessage(),
+                    LogKey.FILENAME.value: record.filename,
+                    LogKey.LINENO.value: record.lineno,
+                    LogKey.PROCESS.value: record.process,
+                }
+            )
+            if record.exc_info:
+                if not record.exc_text:
+                    record.exc_text = self.formatException(record.exc_info)
+                record_format_attrs[LogKey.EXC_TEXT.value] = record.exc_text
+
+        # Add the user specified additional standard attributes.
+        for key in self._additional_log_standard_attrs:
+            _append_flatten_attributes(
+                record_format_attrs, key, getattr(record, key, None)
+            )
+
+        for key, value in record.__dict__.items():
+            # Both Ray and user-provided context are stored in `record_format`.
+            if key not in LOGRECORD_STANDARD_ATTRS:
+                _append_flatten_attributes(record_format_attrs, key, value)
+
+        # Format the internal timestamp to the standardized `timestamp_ns` key.
+        if INTERNAL_TIMESTAMP_LOG_KEY in record_format_attrs:
+            record_format_attrs[LogKey.TIMESTAMP_NS.value] = record_format_attrs.pop(
+                INTERNAL_TIMESTAMP_LOG_KEY
+            )
+
+        return record_format_attrs
+
+    @abstractmethod
+    def format(self, record: logging.LogRecord) -> str:
+        pass
+
+
+class JSONFormatter(AbstractFormatter):
+    def format(self, record: logging.LogRecord) -> str:
+        record_format_attrs = self.generate_record_format_attrs(
+            record, exclude_default_standard_attrs=False
+        )
+        return json.dumps(record_format_attrs)
+
+
+class TextFormatter(AbstractFormatter):
+    def __init__(self, fmt=None, datefmt=None, style="%", validate=True) -> None:
+        super().__init__(fmt, datefmt, style, validate)
+        self._inner_formatter = logging.Formatter(LOGGER_FORMAT)
+
+    def format(self, record: logging.LogRecord) -> str:
+        s = self._inner_formatter.format(record)
+        record_format_attrs = self.generate_record_format_attrs(
+            record, exclude_default_standard_attrs=True
+        )
+
+        additional_attrs = " ".join(
+            [f"{key}={value}" for key, value in record_format_attrs.items()]
+        )
+        return f"{s} {additional_attrs}"

python/ray/_common/logging_constants.py

@@ -0,0 +1,63 @@
+"""Logging-related constants shared across Ray libraries.
+
+Used to distinguish standard Python logging attributes from Ray or user context.
+See: https://docs.python.org/3/library/logging.html#logrecord-attributes
+"""
+
+from enum import Enum
+
+LOGRECORD_STANDARD_ATTRS = frozenset(
+    {
+        "args",
+        "asctime",
+        "created",
+        "exc_info",
+        "exc_text",
+        "filename",
+        "funcName",
+        "levelname",
+        "levelno",
+        "lineno",
+        "message",
+        "module",
+        "msecs",
+        "msg",
+        "name",
+        "pathname",
+        "process",
+        "processName",
+        "relativeCreated",
+        "stack_info",
+        "thread",
+        "threadName",
+        "taskName",
+    }
+)
+
+LOGGER_FLATTEN_KEYS = {
+    "ray_serve_extra_fields",
+}
+
+
+class LogKey(str, Enum):
+    # Core context
+    JOB_ID = "job_id"
+    WORKER_ID = "worker_id"
+    NODE_ID = "node_id"
+    ACTOR_ID = "actor_id"
+    TASK_ID = "task_id"
+    ACTOR_NAME = "actor_name"
+    TASK_NAME = "task_name"
+    TASK_FUNCTION_NAME = "task_func_name"
+
+    # Logger built-in context
+    ASCTIME = "asctime"
+    LEVELNAME = "levelname"
+    MESSAGE = "message"
+    FILENAME = "filename"
+    LINENO = "lineno"
+    EXC_TEXT = "exc_text"
+    PROCESS = "process"
+
+    # Ray logging context
+    TIMESTAMP_NS = "timestamp_ns"

python/ray/_common/network_utils.py

@@ -0,0 +1,27 @@
+import socket
+from contextlib import closing
+
+from ray._raylet import (  # noqa: F401
+    build_address,
+    get_all_interfaces_ip,
+    get_localhost_ip,
+    is_ipv6,
+    is_localhost,
+    node_ip_address_from_perspective,
+    parse_address,
+)
+
+
+def find_free_port(family: socket.AddressFamily = socket.AF_INET) -> int:
+    """Find a free port on the local machine.
+
+    Args:
+        family: The socket address family (AF_INET for IPv4, AF_INET6 for IPv6).
+            Defaults to AF_INET.
+
+    Returns:
+        An available port number.
+    """
+    with closing(socket.socket(family, socket.SOCK_STREAM)) as s:
+        s.bind(("", 0))
+        return s.getsockname()[1]

python/ray/_common/observability/__init__.py

@@ -0,0 +1,12 @@
+"""This module provides the Python API for emitting internal Ray events
+via the ONE-Event framework. Events are buffered and exported through
+the C++ RayEventRecorder.
+"""
+
+from ray._common.observability.internal_event import InternalEventBuilder
+from ray._common.observability.platform_events import PlatformEventBuilder
+
+__all__ = [
+    "InternalEventBuilder",
+    "PlatformEventBuilder",
+]

python/ray/_common/observability/internal_event.py

@@ -0,0 +1,111 @@
+"""Base class for building internal Ray events.
+
+This module provides the base class for creating event builders that can
+emit events to dashboard-agents aggregator agent service.
+"""
+
+from abc import ABC, abstractmethod
+from typing import Optional
+
+from ray._raylet import RayEvent
+from ray.core.generated.events_base_event_pb2 import RayEvent as RayEventProto
+from ray.util.annotations import DeveloperAPI
+
+
+@DeveloperAPI
+class InternalEventBuilder(ABC):
+    """Abstract base class for building internal Ray events.
+
+    Subclasses implement specific event types
+    and must implement get_entity_id() and serialize_event_data().
+    """
+
+    def __init__(
+        self,
+        source_type: int,
+        event_type: int,
+        nested_event_field_number: int,
+        severity: int = RayEventProto.Severity.INFO,
+        message: str = "",
+        session_name: str = "",
+    ):
+        """Initialize the event builder.
+
+        Args:
+            source_type: RayEvent.SourceType enum value. Use
+                RayEventProto.SourceType.<NAME> constants (e.g.,
+                RayEventProto.SourceType.JOBS,
+                RayEventProto.SourceType.GCS).
+            event_type: RayEvent.EventType enum value. Use
+                RayEventProto.EventType.<NAME> constants (e.g.,
+                RayEventProto.EventType.SUBMISSION_JOB_DEFINITION_EVENT,
+                RayEventProto.EventType.DRIVER_JOB_LIFECYCLE_EVENT).
+            nested_event_field_number: The field number in RayEvent proto for the
+                nested event message. Use RayEventProto.<FIELD>_FIELD_NUMBER
+                constants (e.g.,
+                RayEventProto.SUBMISSION_JOB_DEFINITION_EVENT_FIELD_NUMBER).
+            severity: RayEvent.Severity enum value (default INFO).
+            message: Optional message associated with the event.
+            session_name: The Ray session name.
+        """
+        self._source_type = source_type
+        self._event_type = event_type
+        self._nested_event_field_number = nested_event_field_number
+        self._severity = severity
+        self._message = message
+        self._session_name = session_name
+
+    @abstractmethod
+    def get_entity_id(self) -> str:
+        """Return the unique entity ID for this event.
+
+        The entity ID is used to associate related events (e.g., definition
+        and lifecycle events for the same job).
+
+        Returns:
+            A string identifier unique to this entity (e.g., node_id/task_id/actor_id/job_id).
+        """
+        pass
+
+    @abstractmethod
+    def serialize_event_data(self) -> bytes:
+        """Serialize the event-specific protobuf data to bytes.
+
+        Returns:
+            Serialized protobuf bytes of the nested event message
+            (e.g., SubmissionJobDefinitionEvent.SerializeToString()).
+        """
+        pass
+
+    def build(
+        self,
+        event_id: Optional[bytes] = None,
+        timestamp_ns: Optional[int] = None,
+    ) -> RayEvent:
+        """Build the Cython RayEvent object for submission.
+
+        Args:
+            event_id: Optional explicit event id bytes. When omitted, the C++ layer
+                generates a random id (matching the convention used by the other
+                RayEventInterface subclasses). Provide an explicit value to reuse
+                an id from an upstream source (e.g., a Kubernetes event uid).
+            timestamp_ns: Optional explicit event timestamp in nanoseconds since the
+                unix epoch. When omitted, the C++ layer captures the current time
+                at construction. Provide an explicit value for platform events that
+                carry their own source timestamp.
+
+        Returns:
+            A RayEvent object.
+        """
+        return RayEvent(
+            source_type=self._source_type,
+            event_type=self._event_type,
+            severity=self._severity,
+            entity_id=self.get_entity_id(),
+            message=self._message,
+            session_name=self._session_name,
+            serialized_data=self.serialize_event_data(),
+            nested_event_field_number=self._nested_event_field_number,
+            event_id=event_id if event_id is not None else b"",
+            timestamp_ns=int(timestamp_ns) if timestamp_ns is not None else 0,
+        )

python/ray/_common/observability/platform_events.py

@@ -0,0 +1,67 @@
+from typing import Dict, Optional
+
+from ray._common.observability.internal_event import InternalEventBuilder
+from ray.core.generated.events_base_event_pb2 import RayEvent as RayEventProto
+from ray.core.generated.platform_event_pb2 import PlatformEvent, Source
+from ray.util.annotations import DeveloperAPI
+
+
+@DeveloperAPI
+class PlatformEventBuilder(InternalEventBuilder):
+    """Builder for creating infrastructure PlatformEvents (e.g., from Kubernetes)."""
+
+    def __init__(
+        self,
+        event_uid: str,
+        platform: int = Source.Platform.PLATFORM_UNSPECIFIED,  # Source.Platform enum
+        object_kind: str = "",
+        object_name: str = "",
+        reason: str = "",
+        message: str = "",
+        severity: int = RayEventProto.Severity.INFO,
+        component: str = "",
+        source_metadata: Optional[Dict[str, str]] = None,
+        custom_fields: Optional[Dict[str, str]] = None,
+        session_name: str = "",
+    ):
+        super().__init__(
+            source_type=RayEventProto.SourceType.CLUSTER_LIFECYCLE,
+            event_type=RayEventProto.EventType.PLATFORM_EVENT,
+            nested_event_field_number=RayEventProto.PLATFORM_EVENT_FIELD_NUMBER,
+            severity=severity,
+            message=message,
+            session_name=session_name,
+        )
+        self._event_uid = event_uid
+        self._platform = platform
+        self._object_kind = object_kind
+        self._object_name = object_name
+        self._reason = reason
+        self._component = component
+        self._source_metadata = source_metadata
+        self._custom_fields = custom_fields
+
+    def get_entity_id(self) -> str:
+        return self._event_uid
+
+    def serialize_event_data(self) -> bytes:
+        source_proto = Source(
+            platform=self._platform,
+            component=self._component,
+        )
+        if self._source_metadata:
+            for k, v in self._source_metadata.items():
+                source_proto.metadata[k] = v
+
+        event = PlatformEvent(
+            source=source_proto,
+            object_kind=self._object_kind,
+            object_name=self._object_name,
+            message=self._message,
+            reason=self._reason,
+        )
+        if self._custom_fields:
+            for k, v in self._custom_fields.items():
+                event.custom_fields[k] = v
+
+        return event.SerializeToString()

python/ray/_common/pydantic_compat.py

@@ -0,0 +1,113 @@
+# ruff: noqa
+import packaging.version
+
+# Pydantic is a dependency of `ray["default"]` but not the minimal installation,
+# so handle the case where it isn't installed.
+try:
+    import pydantic
+
+    PYDANTIC_INSTALLED = True
+except ImportError:
+    pydantic = None
+    PYDANTIC_INSTALLED = False
+
+
+if not PYDANTIC_INSTALLED:
+    IS_PYDANTIC_2 = False
+    BaseModel = None
+    Extra = None
+    Field = None
+    NonNegativeFloat = None
+    NonNegativeInt = None
+    PositiveFloat = None
+    PositiveInt = None
+    PrivateAttr = None
+    StrictInt = None
+    ValidationError = None
+    root_validator = None
+    validator = None
+
+    def is_subclass_of_base_model(obj):
+        return False
+
+elif not hasattr(pydantic, "__version__") or packaging.version.parse(
+    pydantic.__version__
+) < packaging.version.parse("2.0"):
+    raise ImportError(
+        "Pydantic v1 is no longer supported in Ray. " "Please upgrade to `pydantic>=2`."
+    )
+else:
+    IS_PYDANTIC_2 = True
+    from pydantic import (
+        BaseModel,
+        Extra,
+        Field,
+        NonNegativeFloat,
+        NonNegativeInt,
+        PositiveFloat,
+        PositiveInt,
+        PrivateAttr,
+        StrictInt,
+        ValidationError,
+        root_validator,
+        validator,
+    )
+
+    def is_subclass_of_base_model(obj):
+        return issubclass(obj, BaseModel)
+
+
+def _iter_model_field_types():
+    model_field_types = []
+
+    try:
+        from pydantic.fields import ModelField as model_field_type
+    except ImportError:
+        pass
+    else:
+        model_field_types.append(model_field_type)
+
+    try:
+        from pydantic.v1.fields import ModelField as compat_model_field_type
+    except ImportError:
+        pass
+    else:
+        if compat_model_field_type not in model_field_types:
+            model_field_types.append(compat_model_field_type)
+
+    return model_field_types
+
+
+def register_pydantic_serializers(serialization_context):
+    if not PYDANTIC_INSTALLED:
+        return
+
+    # Pydantic's Cython validators are not serializable.
+    # https://github.com/cloudpipe/cloudpickle/issues/408
+    #
+    # FastAPI can still surface Pydantic's v1 compatibility ModelField under
+    # Pydantic v2, so we need to register serializers for both types until that
+    # compatibility path is no longer used upstream.
+    for model_field_type in _iter_model_field_types():
+        serialization_context._register_cloudpickle_serializer(
+            model_field_type,
+            custom_serializer=lambda o: {
+                "name": o.name,
+                # outer_type_ is the original type for ModelFields,
+                # while type_ can be updated later with the nested type
+                # like int for List[int].
+                "type_": o.outer_type_,
+                "class_validators": o.class_validators,
+                "model_config": o.model_config,
+                "default": o.default,
+                "default_factory": o.default_factory,
+                "required": o.required,
+                "alias": o.alias,
+                "field_info": o.field_info,
+            },
+            custom_deserializer=(
+                lambda kwargs, model_field_type=model_field_type: model_field_type(
+                    **kwargs
+                )
+            ),
+        )

python/ray/_common/ray_constants.py

@@ -0,0 +1,5 @@
+# Default max_concurrency option in @ray.remote for async actors.
+DEFAULT_MAX_CONCURRENCY_ASYNC = 1000
+
+LOGGING_ROTATE_BYTES = 512 * 1024 * 1024  # 512MB.
+LOGGING_ROTATE_BACKUP_COUNT = 5  # 5 Backup files at max.

python/ray/_common/ray_option_utils.py

@@ -0,0 +1,454 @@
+"""Manage, parse and validate options for Ray tasks, actors and actor methods."""
+import warnings
+from dataclasses import dataclass
+from typing import Any, Callable, Dict, Optional, Tuple, Union
+
+import ray
+from ray._private import ray_constants
+from ray._private.label_utils import (
+    validate_fallback_strategy,
+    validate_label_selector,
+)
+from ray._private.utils import get_ray_doc_version
+from ray.util.placement_group import PlacementGroup
+from ray.util.scheduling_strategies import (
+    NodeAffinitySchedulingStrategy,
+    NodeLabelSchedulingStrategy,
+    PlacementGroupSchedulingStrategy,
+)
+
+
+@dataclass
+class Option:
+    # Type constraint of an option.
+    type_constraint: Optional[Union[type, Tuple[type]]] = None
+    # Value constraint of an option.
+    # The callable should return None if there is no error.
+    # Otherwise, return the error message.
+    value_constraint: Optional[Callable[[Any], Optional[str]]] = None
+    # Default value.
+    default_value: Any = None
+
+    def validate(self, keyword: str, value: Any):
+        """Validate the option."""
+        if self.type_constraint is not None:
+            if not isinstance(value, self.type_constraint):
+                raise TypeError(
+                    f"The type of keyword '{keyword}' must be {self.type_constraint}, "
+                    f"but received type {type(value)}"
+                )
+        if self.value_constraint is not None:
+            possible_error_message = self.value_constraint(value)
+            if possible_error_message:
+                raise ValueError(possible_error_message)
+
+
+def _counting_option(name: str, infinite: bool = True, default_value: Any = None):
+    """This is used for positive and discrete options.
+
+    Args:
+        name: The name of the option keyword.
+        infinite: If True, user could use -1 to represent infinity.
+        default_value: The default value for this option.
+
+    Returns:
+        An Option object.
+    """
+    if infinite:
+        return Option(
+            (int, type(None)),
+            lambda x: None
+            if (x is None or x >= -1)
+            else f"The keyword '{name}' only accepts None, 0, -1"
+            " or a positive integer, where -1 represents infinity.",
+            default_value=default_value,
+        )
+    return Option(
+        (int, type(None)),
+        lambda x: None
+        if (x is None or x >= 0)
+        else f"The keyword '{name}' only accepts None, 0 or a positive integer.",
+        default_value=default_value,
+    )
+
+
+def _validate_resource_quantity(name, quantity):
+    if quantity < 0:
+        return f"The quantity of resource {name} cannot be negative"
+    if (
+        isinstance(quantity, float)
+        and quantity != 0.0
+        and int(quantity * ray._raylet.RESOURCE_UNIT_SCALING) == 0
+    ):
+        return (
+            f"The precision of the fractional quantity of resource {name}"
+            " cannot go beyond 0.0001"
+        )
+    resource_name = "GPU" if name == "num_gpus" else name
+    if resource_name in ray._private.accelerators.get_all_accelerator_resource_names():
+        (
+            valid,
+            error_message,
+        ) = ray._private.accelerators.get_accelerator_manager_for_resource(
+            resource_name
+        ).validate_resource_request_quantity(
+            quantity
+        )
+        if not valid:
+            return error_message
+    return None
+
+
+def _resource_option(name: str, default_value: Any = None):
+    """This is used for resource related options."""
+    return Option(
+        (float, int, type(None)),
+        lambda x: None if (x is None) else _validate_resource_quantity(name, x),
+        default_value=default_value,
+    )
+
+
+def _validate_resources(resources: Optional[Dict[str, float]]) -> Optional[str]:
+    if resources is None:
+        return None
+
+    if "CPU" in resources or "GPU" in resources:
+        return (
+            "Use the 'num_cpus' and 'num_gpus' keyword instead of 'CPU' and 'GPU' "
+            "in 'resources' keyword"
+        )
+
+    for name, quantity in resources.items():
+        possible_error_message = _validate_resource_quantity(name, quantity)
+        if possible_error_message:
+            return possible_error_message
+
+    return None
+
+
+_common_options = {
+    "label_selector": Option((dict, type(None)), lambda x: validate_label_selector(x)),
+    "fallback_strategy": Option(
+        (list, type(None)), lambda x: validate_fallback_strategy(x)
+    ),
+    "accelerator_type": Option((str, type(None))),
+    "memory": _resource_option("memory"),
+    "name": Option((str, type(None))),
+    "num_cpus": _resource_option("num_cpus"),
+    "num_gpus": _resource_option("num_gpus"),
+    "object_store_memory": _counting_option("object_store_memory", False),
+    # TODO(suquark): "placement_group", "placement_group_bundle_index"
+    # and "placement_group_capture_child_tasks" are deprecated,
+    # use "scheduling_strategy" instead.
+    "placement_group": Option(
+        (type(None), str, PlacementGroup), default_value="default"
+    ),
+    "placement_group_bundle_index": Option(int, default_value=-1),
+    "placement_group_capture_child_tasks": Option((bool, type(None))),
+    "resources": Option((dict, type(None)), lambda x: _validate_resources(x)),
+    "runtime_env": Option((dict, type(None))),
+    "scheduling_strategy": Option(
+        (
+            type(None),
+            str,
+            PlacementGroupSchedulingStrategy,
+            NodeAffinitySchedulingStrategy,
+            NodeLabelSchedulingStrategy,
+        )
+    ),
+    "enable_task_events": Option(bool, default_value=True),
+    "_labels": Option((dict, type(None))),
+}
+
+
+def issubclass_safe(obj: Any, cls_: type) -> bool:
+    try:
+        return issubclass(obj, cls_)
+    except TypeError:
+        return False
+
+
+_task_only_options = {
+    "max_calls": _counting_option("max_calls", False, default_value=0),
+    # Normal tasks may be retried on failure this many times.
+    # TODO(swang): Allow this to be set globally for an application.
+    "max_retries": _counting_option(
+        "max_retries", default_value=ray_constants.DEFAULT_TASK_MAX_RETRIES
+    ),
+    # override "_common_options"
+    "num_cpus": _resource_option("num_cpus", default_value=1),
+    "num_returns": Option(
+        (int, str, type(None)),
+        lambda x: None
+        if (x is None or x == "dynamic" or x == "streaming" or x >= 0)
+        else "Default None. When None is passed, "
+        "The default value is 1 for a task and actor task, and "
+        "'streaming' for generator tasks and generator actor tasks. "
+        "The keyword 'num_returns' only accepts None, "
+        "a non-negative integer, "
+        "'streaming' (for generators), or 'dynamic'. 'dynamic' flag "
+        "will be deprecated in the future, and it is recommended to use "
+        "'streaming' instead.",
+        default_value=None,
+    ),
+    "object_store_memory": Option(  # override "_common_options"
+        (int, type(None)),
+        lambda x: None
+        if (x is None)
+        else "Setting 'object_store_memory' is not implemented for tasks",
+    ),
+    "retry_exceptions": Option(
+        (bool, list, tuple),
+        lambda x: None
+        if (
+            isinstance(x, bool)
+            or (
+                isinstance(x, (list, tuple))
+                and all(issubclass_safe(x_, Exception) for x_ in x)
+            )
+        )
+        else "retry_exceptions must be either a boolean or a list of exceptions",
+        default_value=False,
+    ),
+    "_generator_backpressure_num_objects": Option(
+        (int, type(None)),
+        lambda x: None
+        if x != 0
+        else (
+            "_generator_backpressure_num_objects=0 is not allowed. "
+            "Use a value > 0. If the value is equal to 1, the behavior "
+            "is identical to Python generator (generator 1 object "
+            "whenever `next` is called). Use -1 to disable this feature. "
+        ),
+    ),
+    "_num_objects_per_yield": Option(
+        (int, type(None)),
+        lambda x: None
+        if (x is None or x > 0)
+        else (
+            "_num_objects_per_yield is a private streaming generator option "
+            "that must be set to a positive integer."
+        ),
+        default_value=1,
+    ),
+}
+
+_actor_only_options = {
+    "concurrency_groups": Option((list, dict, type(None))),
+    "enable_tensor_transport": Option((bool, type(None)), default_value=None),
+    "lifetime": Option(
+        (str, type(None)),
+        lambda x: None
+        if x in (None, "detached", "non_detached")
+        else "actor `lifetime` argument must be one of 'detached', "
+        "'non_detached' and 'None'.",
+    ),
+    "max_concurrency": _counting_option("max_concurrency", False),
+    "max_restarts": _counting_option("max_restarts", default_value=0),
+    "max_task_retries": _counting_option("max_task_retries", default_value=0),
+    "max_pending_calls": _counting_option("max_pending_calls", default_value=-1),
+    "namespace": Option((str, type(None))),
+    "get_if_exists": Option(bool, default_value=False),
+    "allow_out_of_order_execution": Option((bool, type(None))),
+    # Actor-wide cap on the number of unconsumed streaming-generator
+    # objects across all generator tasks running on the actor. Coexists
+    # with the per-method `_generator_backpressure_num_objects`: both
+    # apply, and the producer blocks on whichever is tighter. -1 (or
+    # None / unset) disables the actor-wide cap.
+    "_actor_generator_backpressure_num_objects": Option(
+        (int, type(None)),
+        lambda x: None
+        if (x is None or x > 0 or x == -1)
+        else (
+            "_actor_generator_backpressure_num_objects must be > 0 to cap the "
+            "actor's total unconsumed generator objects, or -1 to disable. "
+            f"Got {x}."
+        ),
+    ),
+}
+
+# Priority is important here because during dictionary update, same key with higher
+# priority overrides the same key with lower priority. We make use of priority
+# to set the correct default value for tasks / actors.
+
+# priority: _common_options > _actor_only_options > _task_only_options
+valid_options: Dict[str, Option] = {
+    **_task_only_options,
+    **_actor_only_options,
+    **_common_options,
+}
+# priority: _task_only_options > _common_options
+task_options: Dict[str, Option] = {**_common_options, **_task_only_options}
+# priority: _actor_only_options > _common_options
+actor_options: Dict[str, Option] = {**_common_options, **_actor_only_options}
+
+remote_args_error_string = (
+    "The @ray.remote decorator must be applied either with no arguments and no "
+    "parentheses, for example '@ray.remote', or it must be applied using some of "
+    f"the arguments in the list {list(valid_options.keys())}, for example "
+    "'@ray.remote(num_returns=2, resources={\"CustomResource\": 1})'."
+)
+
+
+def _check_deprecate_placement_group(options: Dict[str, Any]):
+    """Check if deprecated placement group option exists."""
+    placement_group = options.get("placement_group", "default")
+    scheduling_strategy = options.get("scheduling_strategy")
+    # TODO(suquark): @ray.remote(placement_group=None) is used in
+    # "python/ray.data._internal/remote_fn.py" and many other places,
+    # while "ray.data.read_api.read_datasource" set "scheduling_strategy=SPREAD".
+    # This might be a bug, but it is also ok to allow them co-exist.
+    if (placement_group not in ("default", None)) and (scheduling_strategy is not None):
+        raise ValueError(
+            "Placement groups should be specified via the "
+            "scheduling_strategy option. "
+            "The placement_group option is deprecated."
+        )
+
+
+def _warn_if_using_deprecated_placement_group(
+    options: Dict[str, Any], caller_stacklevel: int
+):
+    placement_group = options["placement_group"]
+    placement_group_bundle_index = options["placement_group_bundle_index"]
+    placement_group_capture_child_tasks = options["placement_group_capture_child_tasks"]
+    if placement_group != "default":
+        warnings.warn(
+            "placement_group parameter is deprecated. Use "
+            "scheduling_strategy=PlacementGroupSchedulingStrategy(...) "
+            "instead, see the usage at "
+            f"https://docs.ray.io/en/{get_ray_doc_version()}/ray-core/package-ref.html#ray-remote.",  # noqa: E501
+            DeprecationWarning,
+            stacklevel=caller_stacklevel + 1,
+        )
+    if placement_group_bundle_index != -1:
+        warnings.warn(
+            "placement_group_bundle_index parameter is deprecated. Use "
+            "scheduling_strategy=PlacementGroupSchedulingStrategy(...) "
+            "instead, see the usage at "
+            f"https://docs.ray.io/en/{get_ray_doc_version()}/ray-core/package-ref.html#ray-remote.",  # noqa: E501
+            DeprecationWarning,
+            stacklevel=caller_stacklevel + 1,
+        )
+    if placement_group_capture_child_tasks:
+        warnings.warn(
+            "placement_group_capture_child_tasks parameter is deprecated. Use "
+            "scheduling_strategy=PlacementGroupSchedulingStrategy(...) "
+            "instead, see the usage at "
+            f"https://docs.ray.io/en/{get_ray_doc_version()}/ray-core/package-ref.html#ray-remote.",  # noqa: E501
+            DeprecationWarning,
+            stacklevel=caller_stacklevel + 1,
+        )
+
+
+def validate_task_options(
+    options: Dict[str, Any],
+    in_options: bool,
+    is_generator_callable: Optional[bool] = None,
+):
+    """Options check for Ray tasks.
+
+    Args:
+        options: Options for Ray tasks.
+        in_options: If True, we are checking the options under the context of
+            ".options()".
+        is_generator_callable: Optional bool indicating whether the callable is a
+            generator function. If provided and num_returns is 'streaming' or
+            'dynamic', validates that the callable is a generator.
+    """
+    for k, v in options.items():
+        if k not in task_options:
+            raise ValueError(
+                f"Invalid option keyword {k} for remote functions. "
+                f"Valid ones are {list(task_options)}."
+            )
+        task_options[k].validate(k, v)
+    if in_options and "max_calls" in options:
+        raise ValueError("Setting 'max_calls' is not supported in '.options()'.")
+    _check_deprecate_placement_group(options)
+
+    if is_generator_callable is not None:
+        num_returns = options.get("num_returns")
+        if num_returns is not None:
+            validate_num_returns(is_generator_callable, num_returns)
+
+
+def validate_actor_options(options: Dict[str, Any], in_options: bool):
+    """Options check for Ray actors.
+
+    Args:
+        options: Options for Ray actors.
+        in_options: If True, we are checking the options under the context of
+            ".options()".
+    """
+    for k, v in options.items():
+        if k not in actor_options:
+            raise ValueError(
+                f"Invalid option keyword {k} for actors. "
+                f"Valid ones are {list(actor_options)}."
+            )
+        actor_options[k].validate(k, v)
+
+    if in_options and "concurrency_groups" in options:
+        raise ValueError(
+            "Setting 'concurrency_groups' is not supported in '.options()'."
+        )
+
+    if options.get("get_if_exists") and not options.get("name"):
+        raise ValueError("The actor name must be specified to use `get_if_exists`.")
+
+    if "object_store_memory" in options:
+        warnings.warn(
+            "Setting 'object_store_memory'"
+            " for actors is deprecated since it doesn't actually"
+            " reserve the required object store memory."
+            f" Use object spilling that's enabled by default (https://docs.ray.io/en/{get_ray_doc_version()}/ray-core/objects/object-spilling.html) "  # noqa: E501
+            "instead to bypass the object store memory size limitation.",
+            DeprecationWarning,
+            stacklevel=1,
+        )
+
+    _check_deprecate_placement_group(options)
+
+
+def validate_num_returns(is_generator_callable: bool, num_returns: Any) -> None:
+    """Validate num_returns for @ray.remote and @ray.method decorators.
+
+    This function validates:
+    1. If num_returns is an integer < 0, it should fail fast.
+    2. If num_returns='streaming' or 'dynamic' is used with a non-generator
+       function, it should fail fast.
+
+    Args:
+        is_generator_callable: Whether the callable is a generator function or
+            async generator function.
+        num_returns: The num_returns value to validate.
+
+    Raises:
+        ValueError: If num_returns < 0, or if num_returns is 'streaming' or 'dynamic'
+            but the callable is not a generator function or async generator function.
+    """
+    if num_returns is None:
+        return
+
+    # Validate num_returns < 0
+    if isinstance(num_returns, int) and num_returns < 0:
+        raise ValueError(f"num_returns must be >= 0, but got {num_returns}.")
+
+    # Validate num_returns='streaming' or 'dynamic' for generator functions
+    if num_returns in ("streaming", "dynamic") and not is_generator_callable:
+        raise ValueError(
+            f"num_returns='{num_returns}' can only be used with generator functions "
+            f"(functions that use 'yield'). "
+            f"The decorated function is not a generator function."
+        )
+
+
+def update_options(
+    original_options: Dict[str, Any], new_options: Dict[str, Any]
+) -> Dict[str, Any]:
+    """Update original options with new options and return.
+    The returned updated options contain shallow copy of original options.
+    """
+
+    return {**original_options, **new_options}

python/ray/_common/retry.py

@@ -0,0 +1,162 @@
+import functools
+import logging
+import random
+import re
+import time
+import traceback
+from collections.abc import Sequence
+from typing import Callable, Optional, TypeVar
+
+try:
+    from typing import ParamSpec
+except ImportError:
+    from typing_extensions import ParamSpec
+
+logger = logging.getLogger(__name__)
+
+R = TypeVar("R")
+P = ParamSpec("P")
+
+
+def format_exception(exc: BaseException, include_cause: bool = False) -> str:
+    """Format ``exc`` as ``"ClassName: message"`` for substring/regex matching.
+
+    Uses `traceback.format_exception_only` so the class name is preserved
+    and callers can match on either the class name or the message.
+
+    Args:
+        exc: The exception to format.
+        include_cause: If True and ``exc.__cause__`` is set (``raise X from Y``),
+            append the cause exception after the base exception. This is useful when
+            we want to match on an exception encountered in the UDF (e.g. ``RateLimitError``)
+            which is wrapped in a ``UserCodeException`` by Ray Data.
+
+    Returns:
+        A single-string representation of ``exc`` in the form
+        ``"ClassName: message"``. When ``include_cause`` is True and
+        ``exc.__cause__`` is set, the cause's formatted form is appended
+        after a single space. See the example below.
+
+    Example:
+        For a ``UserCodeException`` wrapping a ``RateLimitError``, calling ``format_exception(e, include_cause=True)``
+        returns::
+
+            ray.exceptions.UserCodeException: UDF failed to process a data block. RateLimitError: Error code: 429 - rate limited
+    """
+    s = "".join(traceback.format_exception_only(type(exc), exc)).rstrip("\n")
+    if include_cause and exc.__cause__:
+        cause = exc.__cause__
+        s += " " + "".join(traceback.format_exception_only(type(cause), cause)).rstrip(
+            "\n"
+        )
+    return s
+
+
+def matches_error(pattern: str, error_str: str) -> bool:
+    """True if ``pattern`` matches ``error_str`` as a substring or as a regex.
+
+    Substring is tried first so literal patterns are not interpreted as regex.
+    Invalid regex patterns return False instead of raising.
+
+    Args:
+        pattern: Pattern to match, tried first as a substring then as a regex.
+        error_str: Formatted exception string.
+
+    Returns:
+        True if ``pattern`` matches ``error_str`` as a substring or as a regex.
+    """
+    if pattern in error_str:
+        return True
+    try:
+        return bool(re.search(pattern, error_str))
+    except re.error:
+        return False
+
+
+def call_with_retry(
+    f: Callable[P, R],
+    description: str,
+    match: Optional[Sequence[str]] = None,
+    max_attempts: int = 10,
+    max_backoff_s: int = 32,
+    *args: P.args,
+    **kwargs: P.kwargs,
+) -> R:
+    """Retry a function with exponential backoff.
+
+    Args:
+        f: The function to retry.
+        description: An imperative description of the function being retried. For
+            example, "open the file".
+        match: A sequence of patterns to match in the exception message. Each
+            pattern is first checked as a substring, then as a regex. If
+            ``None``, any error is retried.
+        max_attempts: The maximum number of attempts to retry.
+        max_backoff_s: The maximum number of seconds to backoff.
+        *args: Arguments to pass to the function.
+        **kwargs: Keyword arguments to pass to the function.
+
+    Returns:
+        The result of the function.
+    """
+    # TODO: consider inverse match and matching exception type
+    assert max_attempts >= 1, f"`max_attempts` must be positive. Got {max_attempts}."
+
+    for i in range(max_attempts):
+        try:
+            return f(*args, **kwargs)
+        except Exception as e:
+            exception_str = format_exception(e)
+            is_retryable = match is None or any(
+                matches_error(pattern, exception_str) for pattern in match
+            )
+            if is_retryable and i + 1 < max_attempts:
+                # Retry with binary exponential backoff with 20% random jitter.
+                backoff = min(2**i, max_backoff_s) * (random.uniform(0.8, 1.2))
+                logger.debug(
+                    f"Retrying {i+1} attempts to {description} after {backoff} seconds."
+                )
+                time.sleep(backoff)
+            else:
+                if is_retryable:
+                    logger.debug(
+                        f"Failed to {description} after {max_attempts} attempts. Raising."
+                    )
+                else:
+                    logger.debug(
+                        f"Did not find a match for {exception_str}. Raising after {i+1} attempts."
+                    )
+                raise e from None
+
+
+def retry(
+    description: str,
+    match: Optional[Sequence[str]] = None,
+    max_attempts: int = 10,
+    max_backoff_s: int = 32,
+) -> Callable[[Callable[P, R]], Callable[P, R]]:
+    """Decorator-based version of call_with_retry.
+
+    Args:
+        description: An imperative description of the function being retried. For
+            example, "open the file".
+        match: A sequence of patterns to match in the exception message. Each
+            pattern is first checked as a substring, then as a regex. If
+            ``None``, any error is retried.
+        max_attempts: The maximum number of attempts to retry.
+        max_backoff_s: The maximum number of seconds to backoff.
+
+    Returns:
+        A Callable that can be applied in a normal decorator fashion.
+    """
+
+    def decorator(func: Callable[P, R]) -> Callable[P, R]:
+        @functools.wraps(func)
+        def inner(*args: P.args, **kwargs: P.kwargs) -> R:
+            return call_with_retry(
+                func, description, match, max_attempts, max_backoff_s, *args, **kwargs
+            )
+
+        return inner
+
+    return decorator

python/ray/_common/serialization.py

@@ -0,0 +1,34 @@
+import io
+import logging
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    pass
+
+
+import ray._private.utils
+import ray.cloudpickle as pickle
+import ray.exceptions
+from ray._private import ray_constants
+from ray.util import inspect_serializability
+
+logger = logging.getLogger(__name__)
+ALLOW_OUT_OF_BAND_OBJECT_REF_SERIALIZATION = ray_constants.env_bool(
+    "RAY_allow_out_of_band_object_ref_serialization", True
+)
+
+
+def pickle_dumps(obj: Any, error_msg: str):
+    """Wrap cloudpickle.dumps to provide better error message
+    when the object is not serializable.
+    """
+    try:
+        return pickle.dumps(obj)
+    except (TypeError, ray.exceptions.OufOfBandObjectRefSerializationException) as e:
+        sio = io.StringIO()
+        inspect_serializability(obj, print_file=sio)
+        msg = f"{error_msg}:\n{sio.getvalue()}"
+        if isinstance(e, TypeError):
+            raise TypeError(msg) from e
+        else:
+            raise ray.exceptions.OufOfBandObjectRefSerializationException(msg)

python/ray/_common/signature.py

@@ -0,0 +1,163 @@
+import inspect
+import logging
+from inspect import Parameter
+from typing import Any, Dict, List, Tuple
+
+from ray._private.inspect_util import is_cython
+
+# Logger for this module. It should be configured at the entry point
+# into the program using Ray. Ray provides a default configuration at
+# entry/init points.
+logger = logging.getLogger(__name__)
+
+# This dummy type is also defined in ArgumentsBuilder.java. Please keep it
+# synced.
+DUMMY_TYPE = b"__RAY_DUMMY__"
+
+
+def get_signature(func: Any) -> inspect.Signature:
+    """Get signature parameters.
+
+    Support Cython functions by grabbing relevant attributes from the Cython
+    function and attaching to a no-op function. This is somewhat brittle, since
+    inspect may change, but given that inspect is written to a PEP, we hope
+    it is relatively stable. Future versions of Python may allow overloading
+    the inspect 'isfunction' and 'ismethod' functions / create ABC for Python
+    functions. Until then, it appears that Cython won't do anything about
+    compatability with the inspect module.
+
+    Args:
+        func: The function whose signature should be checked.
+
+    Returns:
+        A function signature object, which includes the names of the keyword
+            arguments as well as their default values.
+
+    Raises:
+        TypeError: A type error if the signature is not supported
+    """
+    # The first condition for Cython functions, the latter for Cython instance
+    # methods
+    if is_cython(func):
+        attrs = ["__code__", "__annotations__", "__defaults__", "__kwdefaults__"]
+
+        if all(hasattr(func, attr) for attr in attrs):
+            original_func = func
+
+            def func():
+                return
+
+            for attr in attrs:
+                setattr(func, attr, getattr(original_func, attr))
+        else:
+            raise TypeError(f"{func!r} is not a Python function we can process")
+
+    return inspect.signature(func)
+
+
+def extract_signature(func: Any, ignore_first: bool = False) -> List[Parameter]:
+    """Extract the function signature from the function.
+
+    Args:
+        func: The function whose signature should be extracted.
+        ignore_first: True if the first argument should be ignored. This should
+            be used when func is a method of a class.
+
+    Returns:
+        List of Parameter objects representing the function signature.
+    """
+    signature_parameters = list(get_signature(func).parameters.values())
+
+    if ignore_first:
+        if len(signature_parameters) == 0:
+            raise ValueError(
+                "Methods must take a 'self' argument, but the "
+                f"method '{func.__name__}' does not have one."
+            )
+        signature_parameters = signature_parameters[1:]
+
+    return signature_parameters
+
+
+def validate_args(
+    signature_parameters: List[Parameter], args: Tuple[Any, ...], kwargs: Dict[str, Any]
+) -> None:
+    """Validates the arguments against the signature.
+
+    Args:
+        signature_parameters: The list of Parameter objects
+            representing the function signature, obtained from
+            `extract_signature`.
+        args: The positional arguments passed into the function.
+        kwargs: The keyword arguments passed into the function.
+
+    Raises:
+        TypeError: Raised if arguments do not fit in the function signature.
+    """
+    reconstructed_signature = inspect.Signature(parameters=signature_parameters)
+    try:
+        reconstructed_signature.bind(*args, **kwargs)
+    except TypeError as exc:  # capture a friendlier stacktrace
+        raise TypeError(str(exc)) from None
+
+
+def flatten_args(
+    signature_parameters: List[Parameter], args: Tuple[Any, ...], kwargs: Dict[str, Any]
+) -> List[Any]:
+    """Validates the arguments against the signature and flattens them.
+
+    The flat list representation is a serializable format for arguments.
+    Since the flatbuffer representation of function arguments is a list, we
+    combine both keyword arguments and positional arguments. We represent
+    this with two entries per argument value - [DUMMY_TYPE, x] for positional
+    arguments and [KEY, VALUE] for keyword arguments. See the below example.
+    See `recover_args` for logic restoring the flat list back to args/kwargs.
+
+    Args:
+        signature_parameters: The list of Parameter objects
+            representing the function signature, obtained from
+            `extract_signature`.
+        args: The positional arguments passed into the function.
+        kwargs: The keyword arguments passed into the function.
+
+    Returns:
+        List of args and kwargs. Non-keyword arguments are prefixed
+            by internal enum DUMMY_TYPE.
+
+    Raises:
+        TypeError: Raised if arguments do not fit in the function signature.
+    """
+    validate_args(signature_parameters, args, kwargs)
+    list_args = []
+    for arg in args:
+        list_args += [DUMMY_TYPE, arg]
+
+    for keyword, arg in kwargs.items():
+        list_args += [keyword, arg]
+    return list_args
+
+
+def recover_args(flattened_args: List[Any]) -> Tuple[List[Any], Dict[str, Any]]:
+    """Recreates `args` and `kwargs` from the flattened arg list.
+
+    Args:
+        flattened_args: List of args and kwargs. This should be the output of
+            `flatten_args`.
+
+    Returns:
+        args: The non-keyword arguments passed into the function.
+        kwargs: The keyword arguments passed into the function.
+    """
+    assert (
+        len(flattened_args) % 2 == 0
+    ), "Flattened arguments need to be even-numbered. See `flatten_args`."
+    args = []
+    kwargs = {}
+    for name_index in range(0, len(flattened_args), 2):
+        name, arg = flattened_args[name_index], flattened_args[name_index + 1]
+        if name == DUMMY_TYPE:
+            args.append(arg)
+        else:
+            kwargs[name] = arg
+
+    return args, kwargs

python/ray/_common/test_utils.py

@@ -0,0 +1,530 @@
+"""Test utilities for Ray.
+
+This module contains test utility classes that are distributed with the Ray package
+and can be used by external libraries and tests. These utilities must remain in
+_common/ (not in tests/) to be accessible in the Ray package distribution.
+"""
+
+import asyncio
+import inspect
+import logging
+import os
+import subprocess
+import sys
+import threading
+import time
+import traceback
+import uuid
+from collections import defaultdict
+from collections.abc import Awaitable
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any, Callable, Dict, Iterator, List, Optional, Set
+
+import ray
+import ray._common.usage.usage_lib as ray_usage_lib
+import ray._private.utils
+from ray._common.network_utils import build_address
+from ray._common.utils import decode
+
+logger = logging.getLogger(__name__)
+
+try:
+    from prometheus_client.core import Metric
+    from prometheus_client.parser import Sample, text_string_to_metric_families
+except (ImportError, ModuleNotFoundError):
+    Metric = None
+    Sample = None
+
+    def text_string_to_metric_families(*args, **kwargs):
+        raise ModuleNotFoundError("`prometheus_client` not found")
+
+
+@ray.remote(num_cpus=0)
+class SignalActor:
+    """A Ray actor for coordinating test execution through signals.
+
+    Useful for testing async coordination, waiting for specific states,
+    and synchronizing multiple actors or tasks in tests.
+    """
+
+    def __init__(self):
+        self.ready_event = asyncio.Event()
+        self.num_waiters = 0
+
+    def send(self, clear: bool = False):
+        self.ready_event.set()
+        if clear:
+            self.ready_event.clear()
+
+    async def wait(self, should_wait: bool = True):
+        if should_wait:
+            self.num_waiters += 1
+            await self.ready_event.wait()
+            self.num_waiters -= 1
+
+    async def cur_num_waiters(self) -> int:
+        return self.num_waiters
+
+
+@ray.remote(num_cpus=0)
+class Semaphore:
+    """A Ray actor implementing a semaphore for test coordination.
+
+    Useful for testing resource limiting, concurrency control,
+    and coordination between multiple actors or tasks.
+    """
+
+    def __init__(self, value: int = 1):
+        self._sema = asyncio.Semaphore(value=value)
+
+    async def acquire(self):
+        await self._sema.acquire()
+
+    async def release(self):
+        self._sema.release()
+
+    async def locked(self) -> bool:
+        return self._sema.locked()
+
+
+__all__ = ["SignalActor", "Semaphore"]
+
+
+def wait_for_condition(
+    condition_predictor: Callable[..., bool],
+    timeout: float = 10,
+    retry_interval_ms: float = 100,
+    raise_exceptions: bool = False,
+    **kwargs: Any,
+):
+    """Wait until a condition is met or time out with an exception.
+
+    Args:
+        condition_predictor: A function that predicts the condition.
+        timeout: Maximum timeout in seconds.
+        retry_interval_ms: Retry interval in milliseconds.
+        raise_exceptions: If true, exceptions that occur while executing
+            condition_predictor won't be caught and instead will be raised.
+        **kwargs: Arguments to pass to the condition_predictor.
+
+    Returns:
+        None: Returns when the condition is met.
+
+    Raises:
+        RuntimeError: If the condition is not met before the timeout expires.
+    """
+    start = time.monotonic()
+    last_ex = None
+    while time.monotonic() - start <= timeout:
+        try:
+            if condition_predictor(**kwargs):
+                return
+        except Exception:
+            if raise_exceptions:
+                raise
+            last_ex = ray._private.utils.format_error_message(traceback.format_exc())
+        time.sleep(retry_interval_ms / 1000.0)
+    message = "The condition wasn't met before the timeout expired."
+    if last_ex is not None:
+        message += f" Last exception: {last_ex}"
+    raise RuntimeError(message)
+
+
+async def async_wait_for_condition(
+    condition_predictor: Callable[..., Awaitable[bool]],
+    timeout: float = 10,
+    retry_interval_ms: float = 100,
+    **kwargs: Any,
+):
+    """Wait until a condition is met or time out with an exception.
+
+    Args:
+        condition_predictor: A function that predicts the condition.
+        timeout: Maximum timeout in seconds.
+        retry_interval_ms: Retry interval in milliseconds.
+        **kwargs: Arguments to pass to the condition_predictor.
+
+    Returns:
+        None: Returns when the condition is met.
+
+    Raises:
+        RuntimeError: If the condition is not met before the timeout expires.
+    """
+    start = time.monotonic()
+    last_ex = None
+    while time.monotonic() - start <= timeout:
+        try:
+            if inspect.iscoroutinefunction(condition_predictor):
+                if await condition_predictor(**kwargs):
+                    return
+            else:
+                if condition_predictor(**kwargs):
+                    return
+        except Exception as ex:
+            last_ex = ex
+        await asyncio.sleep(retry_interval_ms / 1000.0)
+    message = "The condition wasn't met before the timeout expired."
+    if last_ex is not None:
+        message += f" Last exception: {last_ex}"
+    raise RuntimeError(message)
+
+
+@contextmanager
+def simulate_s3_bucket(
+    port: int = 5002,
+    region: str = "us-west-2",
+) -> Iterator[str]:
+    """Context manager that simulates an S3 bucket and yields the URI.
+
+    Args:
+        port: The port of the localhost endpoint where S3 is being served.
+        region: The S3 region.
+
+    Yields:
+        str: URI for the simulated S3 bucket.
+    """
+    from moto.server import ThreadedMotoServer
+
+    old_env = os.environ
+    os.environ["AWS_ACCESS_KEY_ID"] = "testing"
+    os.environ["AWS_SECRET_ACCESS_KEY"] = "testing"
+    os.environ["AWS_SECURITY_TOKEN"] = "testing"
+    os.environ["AWS_SESSION_TOKEN"] = "testing"
+
+    s3_server = f"http://{build_address('localhost', port)}"
+    server = ThreadedMotoServer(port=port)
+    server.start()
+    url = f"s3://{uuid.uuid4().hex}?region={region}&endpoint_override={s3_server}"
+    yield url
+    server.stop()
+    os.environ = old_env
+
+
+class TelemetryCallsite(Enum):
+    DRIVER = "driver"
+    ACTOR = "actor"
+    TASK = "task"
+
+
+def _get_library_usages() -> Set[str]:
+    return set(
+        ray_usage_lib.get_library_usages_to_report(
+            ray.experimental.internal_kv.internal_kv_get_gcs_client()
+        )
+    )
+
+
+def _get_extra_usage_tags() -> Dict[str, str]:
+    return ray_usage_lib.get_extra_usage_tags_to_report(
+        ray.experimental.internal_kv.internal_kv_get_gcs_client()
+    )
+
+
+def check_library_usage_telemetry(
+    use_lib_fn: Callable[[], None],
+    *,
+    callsite: TelemetryCallsite,
+    expected_library_usages: List[Set[str]],
+    expected_extra_usage_tags: Optional[Dict[str, str]] = None,
+):
+    """Helper for writing tests to validate library usage telemetry.
+
+    `use_lib_fn` is a callable that will be called from the provided callsite.
+    After calling it, the telemetry data to export will be validated against
+    expected_library_usages and expected_extra_usage_tags.
+    """
+    assert len(_get_library_usages()) == 0, _get_library_usages()
+
+    if callsite == TelemetryCallsite.DRIVER:
+        use_lib_fn()
+    elif callsite == TelemetryCallsite.ACTOR:
+
+        @ray.remote
+        class A:
+            def __init__(self):
+                use_lib_fn()
+
+        a = A.remote()
+        ray.get(a.__ray_ready__.remote())
+    elif callsite == TelemetryCallsite.TASK:
+
+        @ray.remote
+        def f():
+            use_lib_fn()
+
+        ray.get(f.remote())
+    else:
+        assert False, f"Unrecognized callsite: {callsite}"
+
+    library_usages = _get_library_usages()
+    extra_usage_tags = _get_extra_usage_tags()
+
+    assert library_usages in expected_library_usages, library_usages
+    if expected_extra_usage_tags:
+        assert all(
+            [extra_usage_tags[k] == v for k, v in expected_extra_usage_tags.items()]
+        ), extra_usage_tags
+
+
+class FakeTimer:
+    def __init__(self, start_time: Optional[float] = None):
+        self._lock = threading.Lock()
+        self.reset(start_time=start_time)
+
+    def reset(self, start_time: Optional[float] = None):
+        with self._lock:
+            if start_time is None:
+                start_time = time.time()
+            self._curr = start_time
+
+    def time(self) -> float:
+        return self._curr
+
+    def advance(self, by: float):
+        with self._lock:
+            self._curr += by
+
+    def realistic_sleep(self, amt: float):
+        with self._lock:
+            self._curr += amt + 0.001
+
+
+def is_named_tuple(cls):
+    """Return True if cls is a namedtuple and False otherwise."""
+    b = cls.__bases__
+    if len(b) != 1 or b[0] is not tuple:
+        return False
+    f = getattr(cls, "_fields", None)
+    if not isinstance(f, tuple):
+        return False
+    return all(type(n) is str for n in f)
+
+
+def assert_tensors_equivalent(obj1, obj2):
+    """
+    Recursively compare objects with special handling for torch.Tensor.
+
+    Tensors are considered equivalent if:
+      - Same dtype and shape
+      - Same device type (e.g., both 'cpu' or both 'cuda'), index ignored
+      - Values are equal (or close for floats)
+    """
+    import torch
+
+    if isinstance(obj1, torch.Tensor) and isinstance(obj2, torch.Tensor):
+        # 1. dtype
+        assert obj1.dtype == obj2.dtype, f"dtype mismatch: {obj1.dtype} vs {obj2.dtype}"
+        # 2. shape
+        assert obj1.shape == obj2.shape, f"shape mismatch: {obj1.shape} vs {obj2.shape}"
+        # 3. device type must match (cpu/cpu or cuda/cuda), ignore index
+        assert (
+            obj1.device.type == obj2.device.type
+        ), f"Device type mismatch: {obj1.device} vs {obj2.device}"
+
+        # 4. Compare values safely on CPU
+        t1_cpu = obj1.cpu()
+        t2_cpu = obj2.cpu()
+        if obj1.dtype.is_floating_point or obj1.dtype.is_complex:
+            assert torch.allclose(
+                t1_cpu, t2_cpu, atol=1e-6, rtol=1e-5
+            ), "Floating-point tensors not close"
+        else:
+            assert torch.equal(t1_cpu, t2_cpu), "Integer/bool tensors not equal"
+        return
+
+    # Type must match
+    if type(obj1) is not type(obj2):
+        raise AssertionError(f"Type mismatch: {type(obj1)} vs {type(obj2)}")
+
+    # Handle namedtuples
+    if is_named_tuple(type(obj1)):
+        assert len(obj1) == len(obj2)
+        for a, b in zip(obj1, obj2):
+            assert_tensors_equivalent(a, b)
+    elif isinstance(obj1, dict):
+        assert obj1.keys() == obj2.keys()
+        for k in obj1:
+            assert_tensors_equivalent(obj1[k], obj2[k])
+    elif isinstance(obj1, (list, tuple)):
+        assert len(obj1) == len(obj2)
+        for a, b in zip(obj1, obj2):
+            assert_tensors_equivalent(a, b)
+    elif hasattr(obj1, "__dict__") and hasattr(obj2, "__dict__"):
+        # Compare user-defined objects by their public attributes
+        keys1 = {
+            k
+            for k in obj1.__dict__.keys()
+            if not k.startswith("_ray_") and k != "_pytype_"
+        }
+        keys2 = {
+            k
+            for k in obj2.__dict__.keys()
+            if not k.startswith("_ray_") and k != "_pytype_"
+        }
+        assert keys1 == keys2, f"Object attribute keys differ: {keys1} vs {keys2}"
+        for k in keys1:
+            assert_tensors_equivalent(obj1.__dict__[k], obj2.__dict__[k])
+    else:
+        # Fallback for primitives: int, float, str, bool, etc.
+        assert obj1 == obj2, f"Non-tensor values differ: {obj1} vs {obj2}"
+
+
+def run_string_as_driver(
+    driver_script: str, env: Dict = None, encode: str = "utf-8"
+) -> str:
+    """Run a driver as a separate process.
+
+    Args:
+        driver_script: A string to run as a Python script.
+        env: The environment variables for the driver.
+        encode: The encoding to use for the driver script.
+
+    Returns:
+        The script's output.
+    """
+    proc = subprocess.Popen(
+        [sys.executable, "-"],
+        stdin=subprocess.PIPE,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.STDOUT,
+        env=env,
+    )
+    with proc:
+        output = proc.communicate(driver_script.encode(encoding=encode))[0]
+        if proc.returncode:
+            print(decode(output, encode_type=encode))
+            logger.error(proc.stderr)
+            raise subprocess.CalledProcessError(
+                proc.returncode, proc.args, output, proc.stderr
+            )
+        out = decode(output, encode_type=encode)
+    return out
+
+
+@dataclass
+class MetricSamplePattern:
+    name: Optional[str] = None
+    value: Optional[str] = None
+    partial_label_match: Optional[Dict[str, str]] = None
+
+    def matches(self, sample: "Sample"):
+        if self.name is not None:
+            if self.name != sample.name:
+                return False
+
+        if self.value is not None:
+            if self.value != sample.value:
+                return False
+
+        if self.partial_label_match is not None:
+            for label, value in self.partial_label_match.items():
+                if sample.labels.get(label) != value:
+                    return False
+
+        return True
+
+
+@dataclass
+class PrometheusTimeseries:
+    """A collection of timeseries from multiple addresses. Each timeseries is a
+    collection of samples with the same metric name and labels. Concretely:
+    - components_dict: a dictionary of addresses to the Component labels
+    - metric_descriptors: a dictionary of metric names to the Metric object
+    - metric_samples: the latest value of each label
+    """
+
+    components_dict: Dict[str, Set[str]] = field(default_factory=dict)
+    metric_descriptors: Dict[str, "Metric"] = field(default_factory=dict)
+    metric_samples: Dict[frozenset, "Sample"] = field(default_factory=dict)
+
+    def flush(self):
+        self.components_dict.clear()
+        self.metric_descriptors.clear()
+        self.metric_samples.clear()
+
+
+def fetch_raw_prometheus(prom_addresses, timeout=None):
+    # Local import so minimal dependency tests can run without requests
+    import requests
+
+    for address in prom_addresses:
+        try:
+            kwargs = {} if timeout is None else {"timeout": timeout}
+            response = requests.get(f"http://{address}/metrics", **kwargs)
+            yield address, response.text
+        except requests.exceptions.ConnectionError:
+            continue
+        except requests.exceptions.Timeout:
+            continue
+
+
+def fetch_prometheus(prom_addresses, timeout=None):
+    components_dict = {}
+    metric_descriptors = {}
+    metric_samples = []
+
+    for address in prom_addresses:
+        if address not in components_dict:
+            components_dict[address] = set()
+
+    for address, response in fetch_raw_prometheus(prom_addresses, timeout=timeout):
+        for metric in text_string_to_metric_families(response):
+            for sample in metric.samples:
+                metric_descriptors[sample.name] = metric
+                metric_samples.append(sample)
+                if "Component" in sample.labels:
+                    components_dict[address].add(sample.labels["Component"])
+    return components_dict, metric_descriptors, metric_samples
+
+
+def fetch_prometheus_timeseries(
+    prom_addreses: List[str],
+    result: PrometheusTimeseries,
+    timeout=None,
+) -> PrometheusTimeseries:
+    components_dict, metric_descriptors, metric_samples = fetch_prometheus(
+        prom_addreses, timeout=timeout
+    )
+    for address, components in components_dict.items():
+        if address not in result.components_dict:
+            result.components_dict[address] = set()
+        result.components_dict[address].update(components)
+    result.metric_descriptors.update(metric_descriptors)
+    for sample in metric_samples:
+        # udpate sample to the latest value
+        result.metric_samples[
+            frozenset(list(sample.labels.items()) + [("_metric_name_", sample.name)])
+        ] = sample
+    return result
+
+
+def fetch_prometheus_metrics(prom_addresses: List[str]) -> Dict[str, List[Any]]:
+    """Return prometheus metrics from the given addresses.
+
+    Args:
+        prom_addresses: List of metrics_agent addresses to collect metrics from.
+
+    Returns:
+        Dict mapping from metric name to list of samples for the metric.
+    """
+    _, _, samples = fetch_prometheus(prom_addresses)
+    samples_by_name = defaultdict(list)
+    for sample in samples:
+        samples_by_name[sample.name].append(sample)
+    return samples_by_name
+
+
+def fetch_prometheus_metric_timeseries(
+    prom_addresses: List[str],
+    result: PrometheusTimeseries,
+    timeout=None,
+) -> Dict[str, List[Any]]:
+    samples = fetch_prometheus_timeseries(
+        prom_addresses, result, timeout=timeout
+    ).metric_samples.values()
+    samples_by_name = defaultdict(list)
+    for sample in samples:
+        samples_by_name[sample.name].append(sample)
+    return samples_by_name

python/ray/_common/tests/BUILD.bazel

@@ -0,0 +1,53 @@
+load("@rules_python//python:defs.bzl", "py_library")
+load("//bazel:python.bzl", "py_test_module_list")
+
+py_library(
+    name = "conftest",
+    srcs = glob(["**/conftest.py"]),
+    visibility = [
+        "//python/ray/_common/tests:__subpackages__",
+    ],
+    deps = ["//python/ray/tests:conftest"],
+)
+
+# Small tests.
+py_test_module_list(
+    size = "small",
+    files = [
+        "test_deprecation.py",
+        "test_filters.py",
+        "test_formatters.py",
+        "test_logging_constants.py",
+        "test_network_utils.py",
+        "test_ray_option_utils.py",
+        "test_retry.py",
+        "test_signal_semaphore_utils.py",
+        "test_signature.py",
+        "test_tls_utils.py",
+        "test_utils.py",
+        "test_wait_for_condition.py",
+    ],
+    tags = [
+        "exclusive",
+        "team:core",
+    ],
+    deps = [
+        ":conftest",
+        "//:ray_lib",
+    ],
+)
+
+py_test_module_list(
+    size = "large",
+    files = [
+        "test_usage_stats.py",
+    ],
+    tags = [
+        "exclusive",
+        "team:core",
+    ],
+    deps = [
+        ":conftest",
+        "//:ray_lib",
+    ],
+)

python/ray/_common/tests/conftest.py

@@ -0,0 +1,2 @@
+# Imports for filters and formatters tests
+pytest_plugins = ["ray.tests.conftest"]

python/ray/_common/tests/test_deprecation.py

@@ -0,0 +1,97 @@
+import sys
+from unittest.mock import patch
+
+import pytest
+
+from ray._common.deprecation import (
+    DEPRECATED_VALUE,
+    Deprecated,
+    deprecation_warning,
+)
+
+
+def test_deprecation_warning_warn():
+    with patch("ray._common.deprecation.logger.warning") as mock_warning:
+        deprecation_warning("old_feature", "new_feature")
+
+        mock_warning.assert_called_once()
+        args, _ = mock_warning.call_args
+        assert (
+            "DeprecationWarning: `old_feature` has been deprecated. Use `new_feature` instead."
+            in args[0]
+        )
+
+
+def test_deprecation_warning_error():
+    with pytest.raises(ValueError) as excinfo:
+        deprecation_warning("old_feature", error=True)
+    assert "`old_feature` has been deprecated." in str(excinfo.value)
+
+
+def test_deprecated_decorator_function():
+    with patch("ray._common.deprecation.logger.warning") as mock_warning, patch(
+        "ray._common.deprecation.log_once"
+    ) as mock_log_once:
+        mock_log_once.return_value = True
+
+        @Deprecated(old="old_func", new="new_func", error=False)
+        def old_func():
+            return "result"
+
+        result = old_func()
+        assert result == "result"
+        mock_warning.assert_called_once()
+
+
+def test_deprecated_decorator_class():
+    with patch("ray._common.deprecation.logger.warning") as mock_warning, patch(
+        "ray._common.deprecation.log_once"
+    ) as mock_log_once:
+        mock_log_once.return_value = True
+
+        @Deprecated(old="OldClass", new="NewClass", error=False)
+        class OldClass:
+            pass
+
+        instance = OldClass()
+        assert isinstance(instance, OldClass)
+        mock_warning.assert_called_once()
+
+
+def test_deprecated_decorator_method():
+    with patch("ray._common.deprecation.logger.warning") as mock_warning, patch(
+        "ray._common.deprecation.log_once"
+    ) as mock_log_once:
+        mock_log_once.return_value = True
+
+        class MyClass:
+            @Deprecated(old="old_method", new="new_method", error=False)
+            def old_method(self):
+                return "method_result"
+
+        instance = MyClass()
+        result = instance.old_method()
+        assert result == "method_result"
+        mock_warning.assert_called_once()
+
+
+def test_deprecated_decorator_error():
+    with patch("ray._common.deprecation.log_once") as mock_log_once:
+        mock_log_once.return_value = True
+
+        @Deprecated(old="old_func", error=True)
+        def old_func():
+            pass
+
+        with pytest.raises(ValueError):
+            old_func()
+
+
+def test_deprecated_value_constant():
+    assert (
+        DEPRECATED_VALUE == -1
+    ), f"DEPRECATED_VALUE should be -1, but got {DEPRECATED_VALUE}"
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-sv", __file__]))

python/ray/_common/tests/test_filters.py

@@ -0,0 +1,160 @@
+import logging
+import logging.config
+import sys
+from concurrent.futures import ThreadPoolExecutor
+
+import pytest
+
+import ray
+from ray._common.filters import CoreContextFilter
+
+
+class TestCoreContextFilter:
+    def test_driver_process(self, shutdown_only):
+        log_context = ["job_id", "worker_id", "node_id"]
+        filter = CoreContextFilter()
+        record = logging.makeLogRecord({})
+        assert filter.filter(record)
+        # Ray is not initialized so no context except PID which should be available
+        for attr in log_context:
+            assert not hasattr(record, attr)
+        # PID should be available even when Ray is not initialized
+        assert hasattr(record, "process")
+        assert hasattr(record, "_ray_timestamp_ns")
+
+        ray.init()
+        record = logging.makeLogRecord({})
+        assert filter.filter(record)
+        runtime_context = ray.get_runtime_context()
+        expected_values = {
+            "job_id": runtime_context.get_job_id(),
+            "worker_id": runtime_context.get_worker_id(),
+            "node_id": runtime_context.get_node_id(),
+            "process": record.process,
+        }
+        for attr in log_context:
+            assert hasattr(record, attr)
+            assert getattr(record, attr) == expected_values[attr]
+        # This is not a worker process, so actor_id and task_id should not exist.
+        for attr in ["actor_id", "task_id"]:
+            assert not hasattr(record, attr)
+        assert hasattr(record, "_ray_timestamp_ns")
+
+    def test_task_process(self, shutdown_only):
+        @ray.remote
+        def f():
+            filter = CoreContextFilter()
+            record = logging.makeLogRecord({})
+            assert filter.filter(record)
+            should_exist = ["job_id", "worker_id", "node_id", "task_id", "process"]
+            runtime_context = ray.get_runtime_context()
+            expected_values = {
+                "job_id": runtime_context.get_job_id(),
+                "worker_id": runtime_context.get_worker_id(),
+                "node_id": runtime_context.get_node_id(),
+                "task_id": runtime_context.get_task_id(),
+                "task_name": runtime_context.get_task_name(),
+                "task_func_name": runtime_context.get_task_function_name(),
+                "process": record.process,
+            }
+            for attr in should_exist:
+                assert hasattr(record, attr)
+                assert getattr(record, attr) == expected_values[attr]
+            assert not hasattr(record, "actor_id")
+            assert not hasattr(record, "actor_name")
+            assert hasattr(record, "_ray_timestamp_ns")
+
+        obj_ref = f.remote()
+        ray.get(obj_ref)
+
+    def test_actor_process(self, shutdown_only):
+        @ray.remote
+        class A:
+            def f(self):
+                filter = CoreContextFilter()
+                record = logging.makeLogRecord({})
+                assert filter.filter(record)
+                should_exist = [
+                    "job_id",
+                    "worker_id",
+                    "node_id",
+                    "actor_id",
+                    "task_id",
+                    "process",
+                ]
+                runtime_context = ray.get_runtime_context()
+                expected_values = {
+                    "job_id": runtime_context.get_job_id(),
+                    "worker_id": runtime_context.get_worker_id(),
+                    "node_id": runtime_context.get_node_id(),
+                    "actor_id": runtime_context.get_actor_id(),
+                    "actor_name": runtime_context.get_actor_name(),
+                    "task_id": runtime_context.get_task_id(),
+                    "task_name": runtime_context.get_task_name(),
+                    "task_func_name": runtime_context.get_task_function_name(),
+                    "process": record.process,
+                }
+                for attr in should_exist:
+                    assert hasattr(record, attr)
+                    assert getattr(record, attr) == expected_values[attr]
+                assert hasattr(record, "_ray_timestamp_ns")
+
+                # Record should not have the attribute with a value of an empty string.
+                assert runtime_context.get_actor_name() == ""
+                assert not hasattr(record, "actor_name")
+
+        actor = A.remote()
+        ray.get(actor.f.remote())
+
+    def test_actor_process_with_thread(self, shutdown_only):
+        @ray.remote
+        class MockedRayDataWorker:
+            def _check_log_record_in_thread(self):
+                filter = CoreContextFilter()
+                record = logging.makeLogRecord({})
+
+                assert filter.filter(record)
+                should_exist = [
+                    "job_id",
+                    "worker_id",
+                    "node_id",
+                    "actor_id",
+                    "task_id",
+                    "process",
+                ]
+                runtime_context = ray.get_runtime_context()
+                expected_values = {
+                    "job_id": runtime_context.get_job_id(),
+                    "worker_id": runtime_context.get_worker_id(),
+                    "node_id": runtime_context.get_node_id(),
+                    "actor_id": runtime_context.get_actor_id(),
+                    "task_id": runtime_context.get_task_id(),
+                    "process": record.process,
+                }
+                for attr in should_exist:
+                    assert hasattr(record, attr)
+                    assert getattr(record, attr) == expected_values[attr]
+                assert hasattr(record, "_ray_timestamp_ns")
+
+                # Record should not have the attribute with a value of an empty string.
+                assert runtime_context.get_actor_name() == ""
+                assert not hasattr(record, "actor_name")
+
+                assert runtime_context.get_task_name() == ""
+                assert not hasattr(record, "task_name")
+
+                assert runtime_context.get_task_function_name() == ""
+                assert not hasattr(record, "task_function_name")
+
+                return record
+
+            def map(self):
+                with ThreadPoolExecutor(max_workers=1) as executor:
+                    executor.submit(self._check_log_record_in_thread).result()
+
+        actor = MockedRayDataWorker.remote()
+        ray.get(actor.map.remote())
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-sv", __file__]))

python/ray/_common/tests/test_formatters.py

@@ -0,0 +1,162 @@
+import json
+import logging
+import logging.config
+import sys
+
+import pytest
+
+from ray._common.formatters import JSONFormatter, TextFormatter
+
+
+class TestJSONFormatter:
+    def test_empty_record(self, shutdown_only):
+        formatter = JSONFormatter()
+        record = logging.makeLogRecord({})
+        formatted = formatter.format(record)
+
+        record_dict = json.loads(formatted)
+        should_exist = [
+            "process",
+            "asctime",
+            "levelname",
+            "message",
+            "filename",
+            "lineno",
+            "timestamp_ns",
+        ]
+        for key in should_exist:
+            assert key in record_dict
+        assert len(record_dict) == len(should_exist)
+        assert "exc_text" not in record_dict
+
+    def test_record_with_exception(self, shutdown_only):
+        formatter = JSONFormatter()
+        record = logging.makeLogRecord({})
+        try:
+            raise ValueError("test")
+        except ValueError:
+            record.exc_info = sys.exc_info()
+        formatted = formatter.format(record)
+        record_dict = json.loads(formatted)
+        should_exist = [
+            "process",
+            "asctime",
+            "levelname",
+            "message",
+            "filename",
+            "lineno",
+            "exc_text",
+            "timestamp_ns",
+        ]
+        for key in should_exist:
+            assert key in record_dict
+        assert "Traceback (most recent call last):" in record_dict["exc_text"]
+        assert len(record_dict) == len(should_exist)
+
+    def test_record_with_user_provided_context(self, shutdown_only):
+        formatter = JSONFormatter()
+        record = logging.makeLogRecord({"user": "ray"})
+        formatted = formatter.format(record)
+        record_dict = json.loads(formatted)
+        should_exist = [
+            "process",
+            "asctime",
+            "levelname",
+            "message",
+            "filename",
+            "lineno",
+            "user",
+            "timestamp_ns",
+        ]
+        for key in should_exist:
+            assert key in record_dict
+        assert record_dict["user"] == "ray"
+        assert len(record_dict) == len(should_exist)
+        assert "exc_text" not in record_dict
+
+    def test_record_with_flatten_keys_invalid_value(self, shutdown_only):
+        formatter = JSONFormatter()
+        record = logging.makeLogRecord({"ray_serve_extra_fields": "not_a_dict"})
+        with pytest.raises(ValueError):
+            formatter.format(record)
+
+    def test_record_with_flatten_keys_valid_dict(self, shutdown_only):
+        formatter = JSONFormatter()
+        record = logging.makeLogRecord(
+            {"ray_serve_extra_fields": {"key1": "value1", "key2": 2}}
+        )
+        formatted = formatter.format(record)
+        record_dict = json.loads(formatted)
+        should_exist = [
+            "process",
+            "asctime",
+            "levelname",
+            "message",
+            "filename",
+            "lineno",
+            "key1",
+            "key2",
+            "timestamp_ns",
+        ]
+        for key in should_exist:
+            assert key in record_dict
+        assert record_dict["key1"] == "value1", record_dict
+        assert record_dict["key2"] == 2
+        assert "ray_serve_extra_fields" not in record_dict
+        assert len(record_dict) == len(should_exist)
+        assert "exc_text" not in record_dict
+
+    def test_record_with_valid_additional_log_standard_attrs(self, shutdown_only):
+        formatter = JSONFormatter()
+        formatter.set_additional_log_standard_attrs(["name"])
+        record = logging.makeLogRecord({})
+        formatted = formatter.format(record)
+
+        record_dict = json.loads(formatted)
+        should_exist = [
+            "process",
+            "asctime",
+            "levelname",
+            "message",
+            "filename",
+            "lineno",
+            "timestamp_ns",
+            "name",
+        ]
+        for key in should_exist:
+            assert key in record_dict
+        assert len(record_dict) == len(should_exist)
+
+
+class TestTextFormatter:
+    def test_record_with_user_provided_context(self):
+        formatter = TextFormatter()
+        record = logging.makeLogRecord({"user": "ray"})
+        formatted = formatter.format(record)
+        assert "user=ray" in formatted
+
+    def test_record_with_exception(self):
+        formatter = TextFormatter()
+        record = logging.LogRecord(
+            name="test_logger",
+            level=logging.INFO,
+            pathname="test.py",
+            lineno=1000,
+            msg="Test message",
+            args=None,
+            exc_info=None,
+        )
+        formatted = formatter.format(record)
+        for s in ["INFO", "Test message", "test.py:1000", "--"]:
+            assert s in formatted
+
+    def test_record_with_valid_additional_log_standard_attrs(self, shutdown_only):
+        formatter = TextFormatter()
+        formatter.set_additional_log_standard_attrs(["name"])
+        record = logging.makeLogRecord({})
+        formatted = formatter.format(record)
+        assert "name=" in formatted
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-sv", __file__]))

python/ray/_common/tests/test_logging_constants.py

@@ -0,0 +1,83 @@
+import sys
+
+import pytest
+
+from ray._common.logging_constants import (
+    LOGGER_FLATTEN_KEYS,
+    LOGRECORD_STANDARD_ATTRS,
+    LogKey,
+)
+
+
+def test_logrecord_standard_attrs_is_frozenset():
+    assert isinstance(LOGRECORD_STANDARD_ATTRS, frozenset)
+
+
+def test_logrecord_standard_attrs_contains_standard_names():
+    expected = frozenset(
+        {
+            "args",
+            "asctime",
+            "created",
+            "exc_info",
+            "exc_text",
+            "filename",
+            "funcName",
+            "levelname",
+            "levelno",
+            "lineno",
+            "message",
+            "module",
+            "msecs",
+            "msg",
+            "name",
+            "pathname",
+            "process",
+            "processName",
+            "relativeCreated",
+            "stack_info",
+            "thread",
+            "threadName",
+            "taskName",
+        }
+    )
+    assert LOGRECORD_STANDARD_ATTRS == expected
+
+
+def test_logrecord_standard_attrs_has_expected_size():
+    assert len(LOGRECORD_STANDARD_ATTRS) == 23
+
+
+def test_logger_flatten_keys_is_set():
+    assert isinstance(LOGGER_FLATTEN_KEYS, set)
+    assert "ray_serve_extra_fields" in LOGGER_FLATTEN_KEYS
+
+
+def test_logkey_is_enum():
+    from enum import Enum
+
+    assert issubclass(LogKey, Enum)
+    expected = {
+        "JOB_ID": "job_id",
+        "WORKER_ID": "worker_id",
+        "NODE_ID": "node_id",
+        "ACTOR_ID": "actor_id",
+        "TASK_ID": "task_id",
+        "ACTOR_NAME": "actor_name",
+        "TASK_NAME": "task_name",
+        "TASK_FUNCTION_NAME": "task_func_name",
+        "ASCTIME": "asctime",
+        "LEVELNAME": "levelname",
+        "MESSAGE": "message",
+        "FILENAME": "filename",
+        "LINENO": "lineno",
+        "EXC_TEXT": "exc_text",
+        "PROCESS": "process",
+        "TIMESTAMP_NS": "timestamp_ns",
+    }
+    actual = {member.name: member.value for member in LogKey}
+    assert actual == expected
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/_common/tests/test_network_utils.py

@@ -0,0 +1,17 @@
+import sys
+
+import pytest
+
+from ray._common.network_utils import is_localhost
+
+
+def test_is_localhost():
+    assert is_localhost("localhost")
+    assert is_localhost("127.0.0.1")
+    assert is_localhost("::1")
+    assert not is_localhost("8.8.8.8")
+    assert not is_localhost("2001:db8::1")
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/_common/tests/test_ray_option_utils.py

@@ -0,0 +1,197 @@
+import re
+import sys
+from unittest.mock import patch
+
+import pytest
+
+from ray._common.ray_option_utils import (
+    Option,
+    _check_deprecate_placement_group,
+    _counting_option,
+    _resource_option,
+    _validate_resource_quantity,
+    _validate_resources,
+    update_options,
+    validate_actor_options,
+    validate_task_options,
+)
+from ray.util.placement_group import PlacementGroup
+
+
+class TestOptionValidation:
+    def test_option_validate(self):
+        opt = Option(
+            type_constraint=int, value_constraint=lambda v: "error" if v < 0 else None
+        )
+        opt.validate("test", 1)
+        with pytest.raises(TypeError):
+            opt.validate("test", "a")
+        with pytest.raises(ValueError, match="error"):
+            opt.validate("test", -1)
+
+    def test_counting_option(self):
+        # Test infinite counting option
+        opt_inf = _counting_option("test_inf", infinite=True)
+        opt_inf.validate("test_inf", 5)
+        opt_inf.validate("test_inf", 0)
+        opt_inf.validate("test_inf", -1)  # Represents infinity
+        opt_inf.validate("test_inf", None)
+        with pytest.raises(ValueError):
+            opt_inf.validate("test_inf", -2)
+        with pytest.raises(TypeError):
+            opt_inf.validate("test_inf", 1.5)
+
+        # Test non-infinite counting option
+        opt_non_inf = _counting_option("test_non_inf", infinite=False)
+        opt_non_inf.validate("test_non_inf", 5)
+        opt_non_inf.validate("test_non_inf", 0)
+        opt_non_inf.validate("test_non_inf", None)
+        with pytest.raises(ValueError):
+            opt_non_inf.validate("test_non_inf", -1)
+
+    @patch("ray._raylet.RESOURCE_UNIT_SCALING", 10000)
+    @patch(
+        "ray._private.accelerators.get_all_accelerator_resource_names",
+        return_value={"GPU", "TPU"},
+    )
+    @patch("ray._private.accelerators.get_accelerator_manager_for_resource")
+    def test_validate_resource_quantity(self, mock_get_manager, mock_get_all_names):
+        # Valid cases
+        assert _validate_resource_quantity("CPU", 1) is None
+        assert _validate_resource_quantity("memory", 0) is None
+        assert _validate_resource_quantity("custom", 0.5) is None
+
+        # Invalid cases
+        err = _validate_resource_quantity("CPU", -1)
+        assert isinstance(err, str)
+        assert "cannot be negative" in err
+        err = _validate_resource_quantity("CPU", 0.00001)
+        assert isinstance(err, str)
+        assert "cannot go beyond 0.0001" in err
+
+        # Accelerator validation
+        mock_manager_instance = mock_get_manager.return_value
+        mock_manager_instance.validate_resource_request_quantity.return_value = (
+            False,
+            "mock error",
+        )
+        err = _validate_resource_quantity("GPU", 1.5)
+        assert isinstance(err, str)
+        assert "mock error" in err
+        mock_get_manager.assert_called_with("GPU")
+        mock_manager_instance.validate_resource_request_quantity.assert_called_with(1.5)
+
+        mock_manager_instance.validate_resource_request_quantity.return_value = (
+            True,
+            "",
+        )
+        assert _validate_resource_quantity("TPU", 1) is None
+
+    def test_resource_option(self):
+        opt = _resource_option("CPU")
+        opt.validate("CPU", 1)
+        opt.validate("CPU", 0.5)
+        opt.validate("CPU", None)
+        with pytest.raises(TypeError):
+            opt.validate("CPU", "1")
+        with pytest.raises(ValueError):
+            opt.validate("CPU", -1.0)
+
+    def test_validate_resources(self):
+        assert _validate_resources(None) is None
+        assert _validate_resources({"custom": 1}) is None
+        err = _validate_resources({"CPU": 1, "GPU": 1})
+        assert isinstance(err, str)
+        assert "Use the 'num_cpus' and 'num_gpus' keyword" in err
+        err = _validate_resources({"custom": -1})
+        assert isinstance(err, str)
+        assert "cannot be negative" in err
+
+
+class TestTaskActorOptionValidation:
+    def test_validate_task_options_valid(self):
+        validate_task_options({"num_cpus": 2, "max_retries": 3}, in_options=False)
+
+    def test_validate_task_options_invalid_keyword(self):
+        with pytest.raises(ValueError, match="Invalid option keyword"):
+            validate_task_options({"invalid_option": 1}, in_options=False)
+
+    def test_validate_task_options_in_options_invalid(self):
+        with pytest.raises(
+            ValueError,
+            match=re.escape("Setting 'max_calls' is not supported in '.options()'."),
+        ):
+            validate_task_options({"max_calls": 5}, in_options=True)
+
+    def test_validate_actor_options_valid(self):
+        validate_actor_options({"max_concurrency": 2, "name": "abc"}, in_options=False)
+
+    def test_validate_actor_options_invalid_keyword(self):
+        with pytest.raises(ValueError, match="Invalid option keyword"):
+            validate_actor_options({"invalid_option": 1}, in_options=False)
+
+    def test_validate_actor_options_in_options_invalid(self):
+        with pytest.raises(
+            ValueError,
+            match=re.escape(
+                "Setting 'concurrency_groups' is not supported in '.options()'."
+            ),
+        ):
+            validate_actor_options({"concurrency_groups": {}}, in_options=True)
+
+    def test_validate_actor_get_if_exists_no_name(self):
+        with pytest.raises(
+            ValueError, match="must be specified to use `get_if_exists`"
+        ):
+            validate_actor_options({"get_if_exists": True}, in_options=False)
+
+    def test_validate_actor_object_store_memory_warning(self):
+        with pytest.warns(
+            DeprecationWarning,
+            match="Setting 'object_store_memory' for actors is deprecated",
+        ):
+            validate_actor_options({"object_store_memory": 100}, in_options=False)
+
+    def test_check_deprecate_placement_group(self):
+        pg = PlacementGroup.empty()
+        # No error if only one is specified
+        _check_deprecate_placement_group({"placement_group": pg})
+        _check_deprecate_placement_group({"scheduling_strategy": "SPREAD"})
+
+        # Error if both are specified
+        with pytest.raises(
+            ValueError, match="Placement groups should be specified via"
+        ):
+            _check_deprecate_placement_group(
+                {"placement_group": pg, "scheduling_strategy": "SPREAD"}
+            )
+
+        # Check no error with default or None placement_group
+        _check_deprecate_placement_group(
+            {"placement_group": "default", "scheduling_strategy": "SPREAD"}
+        )
+        _check_deprecate_placement_group(
+            {"placement_group": None, "scheduling_strategy": "SPREAD"}
+        )
+
+
+class TestUpdateOptions:
+    def test_simple_update(self):
+        original = {"num_cpus": 1, "name": "a"}
+        new = {"num_cpus": 2, "num_gpus": 1}
+        updated = update_options(original, new)
+        assert updated == {"num_cpus": 2, "name": "a", "num_gpus": 1}
+
+    def test_update_with_empty_new(self):
+        original = {"num_cpus": 1}
+        updated = update_options(original, {})
+        assert updated == original
+
+    def test_update_empty_original(self):
+        new = {"num_cpus": 1}
+        updated = update_options({}, new)
+        assert updated == new
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-sv", __file__]))

python/ray/_common/tests/test_retry.py

@@ -0,0 +1,136 @@
+import sys
+
+import pytest
+
+from ray._common.retry import (
+    call_with_retry,
+    retry,
+)
+
+
+def test_call_with_retry_immediate_success_with_args():
+    def func(a, b):
+        return [a, b]
+
+    assert call_with_retry(func, "func", [], 1, 0, "a", "b") == ["a", "b"]
+
+
+def test_retry_immediate_success_with_object_args():
+    class MyClass:
+        @retry("func", [], 1, 0)
+        def func(self, a, b):
+            return [a, b]
+
+    assert MyClass().func("a", "b") == ["a", "b"]
+
+
+@pytest.mark.parametrize("use_decorator", [True, False])
+def test_retry_last_attempt_successful_with_appropriate_wait_time(
+    monkeypatch, use_decorator
+):
+    sleep_total = 0
+
+    def sleep(x):
+        nonlocal sleep_total
+        sleep_total += x
+
+    monkeypatch.setattr("time.sleep", sleep)
+    monkeypatch.setattr("random.uniform", lambda a, b: 1)
+
+    pattern = "have not reached 4th attempt"
+    call_count = 0
+
+    def func():
+        nonlocal call_count
+        call_count += 1
+        if call_count == 4:
+            return "success"
+        raise ValueError(pattern)
+
+    args = ["func", [pattern], 4, 3]
+    if use_decorator:
+        assert retry(*args)(func)() == "success"
+    else:
+        assert call_with_retry(func, *args) == "success"
+    assert sleep_total == 6  # 1 + 2 + 3
+
+
+@pytest.mark.parametrize("use_decorator", [True, False])
+def test_retry_unretryable_error(use_decorator):
+    call_count = 0
+
+    def func():
+        nonlocal call_count
+        call_count += 1
+        raise ValueError("unretryable error")
+
+    args = ["func", ["only retryable error"], 10, 0]
+    with pytest.raises(ValueError, match="unretryable error"):
+        if use_decorator:
+            retry(*args)(func)()
+        else:
+            call_with_retry(func, *args)
+    assert call_count == 1
+
+
+@pytest.mark.parametrize("use_decorator", [True, False])
+def test_retry_fail_all_attempts_retry_all_errors(use_decorator):
+    call_count = 0
+
+    def func():
+        nonlocal call_count
+        call_count += 1
+        raise ValueError(str(call_count))
+
+    args = ["func", None, 3, 0]
+    with pytest.raises(ValueError):
+        if use_decorator:
+            retry(*args)(func)()
+        else:
+            call_with_retry(func, *args)
+    assert call_count == 3
+
+
+def test_call_with_retry_matches_class_name():
+    """Patterns can match the exception class name (e.g., 'RateLimit')."""
+
+    class RateLimitError(Exception):
+        pass
+
+    call_count = 0
+
+    def func():
+        nonlocal call_count
+        call_count += 1
+        raise RateLimitError("Error code: 429")
+
+    with pytest.raises(RateLimitError):
+        call_with_retry(func, "func", ["RateLimit"], 3, 0)
+    assert call_count == 3
+
+
+@pytest.mark.parametrize(
+    "pattern,should_retry",
+    [
+        # Valid regex that is not a literal substring, that matches via regex search
+        (r"\d{3}", True),
+        # Invalid regex, re.error is handled by returning False and the error is not retried.
+        (r"[unclosed", False),
+    ],
+)
+def test_call_with_retry_regex_matching(pattern, should_retry):
+    call_count = 0
+
+    def func():
+        nonlocal call_count
+        call_count += 1
+        raise ValueError("Error code: 429")
+
+    with pytest.raises(ValueError):
+        call_with_retry(func, "func", [pattern], 3, 0)
+
+    assert call_count == (3 if should_retry else 1)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-sv", __file__]))

python/ray/_common/tests/test_signal_semaphore_utils.py

@@ -0,0 +1,103 @@
+"""Tests for Ray test utility classes.
+
+This module contains pytest-based tests for SignalActor and Semaphore classes
+from ray._common.test_utils. These test utility classes are used for coordination
+and synchronization in Ray tests.
+"""
+
+import sys
+import time
+
+import pytest
+
+import ray
+from ray._common.test_utils import Semaphore, SignalActor, wait_for_condition
+
+
+@pytest.fixture(scope="module")
+def ray_init():
+    """Initialize Ray for the test module."""
+    ray.init(num_cpus=4)
+    yield
+    ray.shutdown()
+
+
+def test_signal_actor_basic(ray_init):
+    """Test basic SignalActor functionality - send and wait operations."""
+    signal = SignalActor.remote()
+
+    # Test initial state
+    assert ray.get(signal.cur_num_waiters.remote()) == 0
+
+    # Test send and wait
+    ray.get(signal.send.remote())
+    signal.wait.remote()
+    assert ray.get(signal.cur_num_waiters.remote()) == 0
+
+
+def test_signal_actor_multiple_waiters(ray_init):
+    """Test SignalActor with multiple waiters and signal clearing."""
+    signal = SignalActor.remote()
+
+    # Create multiple waiters
+    for _ in range(3):
+        signal.wait.remote()
+
+    # Check number of waiters
+    wait_for_condition(lambda: ray.get(signal.cur_num_waiters.remote()) == 3)
+
+    # Send signal and wait for all waiters
+    ray.get(signal.send.remote())
+
+    # Verify all waiters are done
+    wait_for_condition(lambda: ray.get(signal.cur_num_waiters.remote()) == 0)
+
+    # check that .wait() doesn't block if the signal is already sent
+    ray.get(signal.wait.remote())
+
+    assert ray.get(signal.cur_num_waiters.remote()) == 0
+
+    # clear the signal
+    ray.get(signal.send.remote(clear=True))
+    signal.wait.remote()
+    # Verify all waiters are done
+    wait_for_condition(lambda: ray.get(signal.cur_num_waiters.remote()) == 1)
+
+    ray.get(signal.send.remote())
+
+
+def test_semaphore_basic(ray_init):
+    """Test basic Semaphore functionality - acquire, release, and lock status."""
+    sema = Semaphore.remote(value=2)
+
+    # Test initial state
+    wait_for_condition(lambda: ray.get(sema.locked.remote()) is False)
+
+    # Test acquire and release
+    ray.get(sema.acquire.remote())
+    ray.get(sema.acquire.remote())
+    wait_for_condition(lambda: ray.get(sema.locked.remote()) is True)
+
+    ray.get(sema.release.remote())
+    ray.get(sema.release.remote())
+    wait_for_condition(lambda: ray.get(sema.locked.remote()) is False)
+
+
+def test_semaphore_concurrent(ray_init):
+    """Test Semaphore with concurrent workers to verify resource limiting."""
+    sema = Semaphore.remote(value=2)
+
+    def worker():
+        ray.get(sema.acquire.remote())
+        time.sleep(0.1)
+        ray.get(sema.release.remote())
+
+    # Create multiple workers
+    _ = [worker() for _ in range(4)]
+
+    # Verify semaphore is not locked
+    wait_for_condition(lambda: ray.get(sema.locked.remote()) is False)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-sv", __file__]))

python/ray/_common/tests/test_signature.py

@@ -0,0 +1,515 @@
+"""Tests for Ray signature utility functions.
+
+This module contains pytest-based tests for signature-related functions in
+ray._common.signature. These functions are used for extracting, validating,
+and flattening function signatures for serialization.
+"""
+
+import inspect
+import sys
+from typing import Any, Optional
+from unittest.mock import Mock, patch
+
+import pytest
+
+from ray._common.signature import (
+    DUMMY_TYPE,
+    extract_signature,
+    flatten_args,
+    get_signature,
+    recover_args,
+    validate_args,
+)
+
+
+class TestGetSignature:
+    """Tests for the get_signature utility function."""
+
+    def test_regular_function(self):
+        """Test getting signature from a regular Python function."""
+
+        def test_func(a, b=10, *args, **kwargs):
+            return a + b
+
+        sig = get_signature(test_func)
+        assert sig is not None
+        assert len(sig.parameters) == 4
+        assert "a" in sig.parameters
+        assert "b" in sig.parameters
+        assert sig.parameters["b"].default == 10
+
+    def test_function_with_annotations(self):
+        """Test getting signature from a function with type annotations."""
+
+        def test_func(a: int, b: str = "default") -> str:
+            return f"{a}{b}"
+
+        sig = get_signature(test_func)
+        assert sig is not None
+        assert len(sig.parameters) == 2
+        assert sig.parameters["a"].annotation is int
+        assert sig.parameters["b"].annotation is str
+        assert sig.parameters["b"].default == "default"
+
+    def test_function_no_parameters(self):
+        """Test getting signature from a function with no parameters."""
+
+        def test_func():
+            return "hello"
+
+        sig = get_signature(test_func)
+        assert sig is not None
+        assert len(sig.parameters) == 0
+
+    def test_lambda_function(self):
+        """Test getting signature from a lambda function."""
+        sig = get_signature(lambda x, y=5: x + y)
+        assert sig is not None
+        assert len(sig.parameters) == 2  # x, y
+        assert sig.parameters["y"].default == 5
+
+    @patch("ray._common.signature.is_cython")
+    def test_cython_function_with_attributes(self, mock_is_cython):
+        """Test getting signature from a Cython function with required attributes."""
+        mock_is_cython.return_value = True
+
+        def original_func(x=10):
+            return x
+
+        mock_func = Mock()
+        mock_func.__code__ = original_func.__code__
+        mock_func.__annotations__ = original_func.__annotations__
+        mock_func.__defaults__ = original_func.__defaults__
+        mock_func.__kwdefaults__ = original_func.__kwdefaults__
+
+        sig = get_signature(mock_func)
+        assert sig is not None
+        assert len(sig.parameters) == 1
+        assert "x" in sig.parameters
+
+    @patch("ray._common.signature.is_cython")
+    def test_cython_function_missing_attributes(self, mock_is_cython):
+        """Test error handling for Cython function missing required attributes."""
+        mock_is_cython.return_value = True
+
+        # Create a mock Cython function missing required attributes
+        mock_func = Mock()
+        del mock_func.__code__  # Remove required attribute
+
+        with pytest.raises(TypeError, match="is not a Python function we can process"):
+            get_signature(mock_func)
+
+    def test_method_signature(self):
+        """Test getting signature from a class method."""
+
+        class TestClass:
+            def test_method(self, a, b=20):
+                return a + b
+
+        sig = get_signature(TestClass.test_method)
+        assert sig is not None
+        assert len(sig.parameters) == 3  # self, a, b
+        assert "self" in sig.parameters
+        assert "a" in sig.parameters
+        assert "b" in sig.parameters
+        assert sig.parameters["b"].default == 20
+
+
+class TestExtractSignature:
+    """Tests for the extract_signature utility function."""
+
+    def test_function_without_ignore_first(self):
+        """Test extracting signature from function without ignoring first parameter."""
+
+        def test_func(a, b=10, c=None):
+            return a + b
+
+        params = extract_signature(test_func, ignore_first=False)
+        assert len(params) == 3
+        assert params[0].name == "a"
+        assert params[1].name == "b"
+        assert params[1].default == 10
+        assert params[2].name == "c"
+        assert params[2].default is None
+
+    def test_method_with_ignore_first(self):
+        """Test extracting signature from method ignoring 'self' parameter."""
+
+        class TestClass:
+            def test_method(self, a, b=20):
+                return a + b
+
+        params = extract_signature(TestClass.test_method, ignore_first=True)
+        assert len(params) == 2
+        assert params[0].name == "a"
+        assert params[1].name == "b"
+        assert params[1].default == 20
+
+    def test_function_with_ignore_first(self):
+        """Test extracting signature from regular function with ignore_first=True."""
+
+        def test_func(x, y, z=30):
+            return x + y + z
+
+        params = extract_signature(test_func, ignore_first=True)
+        assert len(params) == 2
+        assert params[0].name == "y"
+        assert params[1].name == "z"
+        assert params[1].default == 30
+
+    def test_empty_parameters_with_ignore_first(self):
+        """Test error handling when method has no parameters but ignore_first=True."""
+
+        def test_func():
+            return "hello"
+
+        with pytest.raises(ValueError, match="Methods must take a 'self' argument"):
+            extract_signature(test_func, ignore_first=True)
+
+    def test_single_parameter_with_ignore_first(self):
+        """Test extracting signature from method with only 'self' parameter."""
+
+        class TestClass:
+            def test_method(self):
+                return "hello"
+
+        params = extract_signature(TestClass.test_method, ignore_first=True)
+        assert len(params) == 0
+
+    def test_varargs_and_kwargs(self):
+        """Test extracting signature with *args and **kwargs."""
+
+        def test_func(a, b=10, *args, **kwargs):
+            return a + b
+
+        params = extract_signature(test_func, ignore_first=False)
+        assert len(params) == 4
+        assert params[0].name == "a"
+        assert params[1].name == "b"
+        assert params[2].name == "args"
+        assert params[2].kind == inspect.Parameter.VAR_POSITIONAL
+        assert params[3].name == "kwargs"
+        assert params[3].kind == inspect.Parameter.VAR_KEYWORD
+
+
+class TestValidateArgs:
+    """Tests for the validate_args utility function."""
+
+    def test_valid_positional_args(self):
+        """Test validation with valid positional arguments."""
+
+        def test_func(a, b, c=30):
+            return a + b + c
+
+        params = extract_signature(test_func)
+        # Should not raise an exception
+        validate_args(params, (1, 2), {})
+        validate_args(params, (1, 2, 3), {})
+
+    def test_valid_keyword_args(self):
+        """Test validation with valid keyword arguments."""
+
+        def test_func(a, b=20, c=30):
+            return a + b + c
+
+        params = extract_signature(test_func)
+        # Should not raise an exception
+        validate_args(params, (1,), {"b": 2})
+        validate_args(params, (1,), {"b": 2, "c": 3})
+        validate_args(params, (), {"a": 1, "b": 2, "c": 3})
+
+    def test_valid_mixed_args(self):
+        """Test validation with mixed positional and keyword arguments."""
+
+        def test_func(a, b, c=30):
+            return a + b + c
+
+        params = extract_signature(test_func)
+        # Should not raise an exception
+        validate_args(params, (1,), {"b": 2})
+        validate_args(params, (1, 2), {"c": 3})
+
+    def test_too_many_positional_args(self):
+        """Test error handling for too many positional arguments."""
+
+        def test_func(a, b):
+            return a + b
+
+        params = extract_signature(test_func)
+        with pytest.raises(TypeError):
+            validate_args(params, (1, 2, 3), {})
+
+    def test_missing_required_args(self):
+        """Test error handling for missing required arguments."""
+
+        def test_func(a, b, c=30):
+            return a + b + c
+
+        params = extract_signature(test_func)
+        with pytest.raises(TypeError):
+            validate_args(params, (1,), {})  # Missing 'b'
+
+    def test_unexpected_keyword_args(self):
+        """Test error handling for unexpected keyword arguments."""
+
+        def test_func(a, b):
+            return a + b
+
+        params = extract_signature(test_func)
+        with pytest.raises(TypeError):
+            validate_args(params, (1, 2), {"c": 3})
+
+    def test_duplicate_args(self):
+        """Test error handling for duplicate arguments (positional and keyword)."""
+
+        def test_func(a, b, c=30):
+            return a + b + c
+
+        params = extract_signature(test_func)
+        with pytest.raises(TypeError):
+            validate_args(params, (1, 2), {"b": 3})  # 'b' specified twice
+
+    def test_varargs_validation(self):
+        """Test validation with *args and **kwargs."""
+
+        def test_func(a, b=20, *args, **kwargs):
+            return a + b
+
+        params = extract_signature(test_func)
+        # Should not raise an exception
+        validate_args(params, (1, 2, 3, 4), {"extra": 5})
+        validate_args(params, (1,), {"b": 2, "extra": 3})
+
+
+class TestFlattenArgs:
+    """Tests for the flatten_args utility function."""
+
+    def test_only_positional_args(self):
+        """Test flattening with only positional arguments."""
+
+        def test_func(a, b, c):
+            return a + b + c
+
+        params = extract_signature(test_func)
+        flattened = flatten_args(params, (1, 2, 3), {})
+
+        expected = [DUMMY_TYPE, 1, DUMMY_TYPE, 2, DUMMY_TYPE, 3]
+        assert flattened == expected
+
+    def test_only_keyword_args(self):
+        """Test flattening with only keyword arguments."""
+
+        def test_func(a=1, b=2, c=3):
+            return a + b + c
+
+        params = extract_signature(test_func)
+        flattened = flatten_args(params, (), {"a": 10, "b": 20, "c": 30})
+
+        expected = ["a", 10, "b", 20, "c", 30]
+        assert flattened == expected
+
+    def test_mixed_args(self):
+        """Test flattening with mixed positional and keyword arguments."""
+
+        def test_func(a, b, c=30):
+            return a + b + c
+
+        params = extract_signature(test_func)
+        flattened = flatten_args(params, (1, 2), {"c": 3})
+
+        expected = [DUMMY_TYPE, 1, DUMMY_TYPE, 2, "c", 3]
+        assert flattened == expected
+
+    def test_empty_args(self):
+        """Test flattening with no arguments."""
+
+        def test_func():
+            return "hello"
+
+        params = extract_signature(test_func)
+        flattened = flatten_args(params, (), {})
+
+        assert flattened == []
+
+    def test_complex_types(self):
+        """Test flattening with complex argument types."""
+
+        def test_func(a, b, c=None):
+            return a + b
+
+        params = extract_signature(test_func)
+        complex_args = ([1, 2, 3], {"key": "value"})
+        complex_kwargs = {"c": {"nested": "dict"}}
+
+        flattened = flatten_args(params, complex_args, complex_kwargs)
+
+        expected = [
+            DUMMY_TYPE,
+            [1, 2, 3],
+            DUMMY_TYPE,
+            {"key": "value"},
+            "c",
+            {"nested": "dict"},
+        ]
+        assert flattened == expected
+
+    def test_invalid_args_raises_error(self):
+        """Test that invalid arguments raise TypeError during flattening."""
+
+        def test_func(a, b):
+            return a + b
+
+        params = extract_signature(test_func)
+        with pytest.raises(TypeError):
+            flatten_args(params, (1, 2, 3), {})  # Too many args
+
+
+class TestRecoverArgs:
+    """Tests for the recover_args utility function."""
+
+    def test_only_positional_args(self):
+        """Test recovering only positional arguments."""
+        flattened = [DUMMY_TYPE, 1, DUMMY_TYPE, 2, DUMMY_TYPE, 3]
+        args, kwargs = recover_args(flattened)
+
+        assert args == [1, 2, 3]
+        assert kwargs == {}
+
+    def test_only_keyword_args(self):
+        """Test recovering only keyword arguments."""
+        flattened = ["a", 10, "b", 20, "c", 30]
+        args, kwargs = recover_args(flattened)
+
+        assert args == []
+        assert kwargs == {"a": 10, "b": 20, "c": 30}
+
+    def test_mixed_args(self):
+        """Test recovering mixed positional and keyword arguments."""
+        flattened = [DUMMY_TYPE, 1, DUMMY_TYPE, 2, "c", 3]
+        args, kwargs = recover_args(flattened)
+
+        assert args == [1, 2]
+        assert kwargs == {"c": 3}
+
+    def test_empty_flattened(self):
+        """Test recovering from empty flattened list."""
+        flattened = []
+        args, kwargs = recover_args(flattened)
+
+        assert args == []
+        assert kwargs == {}
+
+    def test_complex_types(self):
+        """Test recovering complex argument types."""
+        flattened = [
+            DUMMY_TYPE,
+            [1, 2, 3],
+            DUMMY_TYPE,
+            {"key": "value"},
+            "c",
+            {"nested": "dict"},
+        ]
+        args, kwargs = recover_args(flattened)
+
+        assert args == [[1, 2, 3], {"key": "value"}]
+        assert kwargs == {"c": {"nested": "dict"}}
+
+    def test_invalid_odd_length(self):
+        """Test error handling for odd-length flattened list."""
+        flattened = [DUMMY_TYPE, 1, "key"]  # Odd length
+        with pytest.raises(
+            AssertionError, match="Flattened arguments need to be even-numbered"
+        ):
+            recover_args(flattened)
+
+    def test_preserve_order(self):
+        """Test that argument order is preserved during flatten/recover."""
+
+        def test_func(a, b, c, d, e):
+            return a + b + c + d + e
+
+        params = extract_signature(test_func)
+        original_args = (1, 2, 3)
+        original_kwargs = {"d": 4, "e": 5}
+
+        flattened = flatten_args(params, original_args, original_kwargs)
+        recovered_args, recovered_kwargs = recover_args(flattened)
+
+        assert recovered_args == [1, 2, 3]
+        assert recovered_kwargs == {"d": 4, "e": 5}
+
+
+class TestIntegration:
+    """Integration tests for signature utilities working together."""
+
+    def test_complete_workflow(self):
+        """Test complete workflow from function to flatten/recover."""
+
+        def test_func(x: int, y: str = "default", z: Optional[Any] = None):
+            return f"{x}_{y}_{z}"
+
+        # Extract signature
+        params = extract_signature(test_func)
+        assert len(params) == 3
+
+        # Validate arguments
+        args = (42, "hello")
+        kwargs = {"z": [1, 2, 3]}
+        validate_args(params, args, kwargs)
+
+        # Flatten arguments
+        flattened = flatten_args(params, args, kwargs)
+        expected = [DUMMY_TYPE, 42, DUMMY_TYPE, "hello", "z", [1, 2, 3]]
+        assert flattened == expected
+
+        # Recover arguments
+        recovered_args, recovered_kwargs = recover_args(flattened)
+        assert recovered_args == list(args)
+        assert recovered_kwargs == kwargs
+
+    def test_method_workflow_with_ignore_first(self):
+        """Test complete workflow for class methods with ignore_first=True."""
+
+        class TestClass:
+            def test_method(self, a: int, b: str = "test"):
+                return f"{a}_{b}"
+
+        # Extract signature ignoring 'self'
+        params = extract_signature(TestClass.test_method, ignore_first=True)
+        assert len(params) == 2
+        assert params[0].name == "a"
+        assert params[1].name == "b"
+
+        # Validate and flatten
+        args = (100,)
+        kwargs = {"b": "custom"}
+        validate_args(params, args, kwargs)
+        flattened = flatten_args(params, args, kwargs)
+
+        # Recover and verify
+        recovered_args, recovered_kwargs = recover_args(flattened)
+        assert recovered_args == list(args)
+        assert recovered_kwargs == kwargs
+
+    def test_varargs_kwargs_workflow(self):
+        """Test workflow with functions that have *args and **kwargs."""
+
+        def test_func(a, b=10, *args, **kwargs):
+            return a + b + sum(args) + sum(kwargs.values())
+
+        params = extract_signature(test_func)
+
+        # Test with extra positional and keyword arguments
+        args = (1, 2, 3, 4, 5)
+        kwargs = {"extra1": 10, "extra2": 20}
+
+        validate_args(params, args, kwargs)
+        flattened = flatten_args(params, args, kwargs)
+        recovered_args, recovered_kwargs = recover_args(flattened)
+
+        assert recovered_args == list(args)
+        assert recovered_kwargs == kwargs
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", "-s", __file__]))

python/ray/_common/tests/test_tls_utils.py

@@ -0,0 +1,41 @@
+import sys
+
+import pytest
+
+from ray._common.tls_utils import generate_self_signed_tls_certs
+
+
+def test_generate_self_signed_tls_certs_returns_tuple():
+    cert_contents, key_contents = generate_self_signed_tls_certs()
+    assert isinstance(cert_contents, str)
+    assert isinstance(key_contents, str)
+
+
+def test_generate_self_signed_tls_certs_pem_format():
+    cert_contents, key_contents = generate_self_signed_tls_certs()
+    assert cert_contents.strip().startswith("-----BEGIN CERTIFICATE-----")
+    assert cert_contents.strip().endswith("-----END CERTIFICATE-----")
+    assert key_contents.strip().startswith("-----BEGIN")
+    assert "PRIVATE KEY" in key_contents
+
+
+def test_generate_self_signed_tls_certs_usable_for_ssl():
+    import ssl
+    import tempfile
+
+    cert_contents, key_contents = generate_self_signed_tls_certs()
+    with (
+        tempfile.NamedTemporaryFile(mode="w", suffix=".crt") as cf,
+        tempfile.NamedTemporaryFile(mode="w", suffix=".key") as kf,
+    ):
+        cf.write(cert_contents)
+        cf.flush()
+        kf.write(key_contents)
+        kf.flush()
+
+        ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
+        ctx.load_cert_chain(cf.name, kf.name)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))

python/ray/_common/tests/test_utils.py

@@ -0,0 +1,291 @@
+"""Tests for Ray utility functions.
+
+This module contains pytest-based tests for utility functions in ray._common.utils.
+Test utility classes (SignalActor, Semaphore) are in ray._common.test_utils to
+ensure they're included in the Ray package distribution.
+"""
+
+import asyncio
+import os
+import sys
+import tempfile
+import warnings
+
+import pytest
+
+from ray._common.utils import (
+    _BACKGROUND_TASKS,
+    env_bool,
+    get_or_create_event_loop,
+    get_system_memory,
+    load_class,
+    run_background_task,
+    try_to_create_directory,
+)
+
+# Optional imports for testing
+try:
+    import psutil
+
+    PSUTIL_AVAILABLE = True
+except ImportError:
+    PSUTIL_AVAILABLE = False
+
+
+class TestGetOrCreateEventLoop:
+    """Tests for the get_or_create_event_loop utility function."""
+
+    def test_existing_event_loop(self):
+        # With running event loop
+        expect_loop = asyncio.new_event_loop()
+        expect_loop.set_debug(True)
+        asyncio.set_event_loop(expect_loop)
+        with warnings.catch_warnings():
+            # Assert no deprecating warnings raised for python>=3.10.
+            warnings.simplefilter("error")
+            actual_loop = get_or_create_event_loop()
+
+            assert actual_loop == expect_loop, "Loop should not be recreated."
+
+    def test_new_event_loop(self):
+        with warnings.catch_warnings():
+            # Assert no deprecating warnings raised for python>=3.10.
+            warnings.simplefilter("error")
+            loop = get_or_create_event_loop()
+            assert loop is not None, "new event loop should be created."
+
+
+class TestEnvBool:
+    """Tests for the env_bool utility function."""
+
+    _KEY = "_RAY_TEST_ENV_BOOL"
+
+    @pytest.mark.parametrize(
+        "env_value, expected",
+        [
+            ("1", True),
+            ("true", True),
+            ("True", True),
+            ("TRUE", True),
+            ("0", False),
+            ("false", False),
+            ("False", False),
+            ("FALSE", False),
+            ("yes", False),
+            ("no", False),
+            ("", False),
+        ],
+    )
+    def test_env_bool_values(self, env_value, expected, monkeypatch):
+        monkeypatch.setenv(self._KEY, env_value)
+        assert env_bool(self._KEY, False) is expected
+
+    def test_env_bool_default_when_unset(self, monkeypatch):
+        monkeypatch.delenv(self._KEY, raising=False)
+        assert env_bool(self._KEY, False) is False
+        assert env_bool(self._KEY, True) is True
+
+
+@pytest.mark.asyncio
+async def test_run_background_task():
+    """Test the run_background_task utility function."""
+    result = {}
+
+    async def co():
+        result["start"] = 1
+        await asyncio.sleep(0)
+        result["end"] = 1
+
+    run_background_task(co())
+
+    # Background task is running.
+    assert len(_BACKGROUND_TASKS) == 1
+    # co executed.
+    await asyncio.sleep(0)
+    # await asyncio.sleep(0) from co is reached.
+    await asyncio.sleep(0)
+    # co finished and callback called.
+    await asyncio.sleep(0)
+    # The task should be removed from the set once it finishes.
+    assert len(_BACKGROUND_TASKS) == 0
+
+    assert result.get("start") == 1
+    assert result.get("end") == 1
+
+
+class TestTryToCreateDirectory:
+    """Tests for the try_to_create_directory utility function."""
+
+    def test_create_new_directory(self):
+        """Test creating a new directory."""
+        with tempfile.TemporaryDirectory() as temp_dir:
+            test_dir = os.path.join(temp_dir, "test_dir")
+            try_to_create_directory(test_dir)
+            assert os.path.exists(test_dir), "Directory should be created"
+            assert os.path.isdir(test_dir), "Path should be a directory"
+
+    def test_existing_directory(self):
+        """Test creating a directory that already exists."""
+        with tempfile.TemporaryDirectory() as temp_dir:
+            test_dir = os.path.join(temp_dir, "existing_dir")
+            # Create directory first
+            os.makedirs(test_dir)
+            # Should work without error
+            try_to_create_directory(test_dir)
+            assert os.path.exists(test_dir), "Directory should still exist"
+
+    def test_nested_directory_creation(self):
+        """Test creating nested directory structure."""
+        with tempfile.TemporaryDirectory() as temp_dir:
+            nested_dir = os.path.join(temp_dir, "nested", "deep", "structure")
+            try_to_create_directory(nested_dir)
+            assert os.path.exists(nested_dir), "Nested directory should be created"
+
+    def test_tilde_expansion(self):
+        """Test directory creation with tilde expansion."""
+        with tempfile.TemporaryDirectory() as temp_dir:
+            fake_home = os.path.join(temp_dir, "fake_home")
+            os.makedirs(fake_home, exist_ok=True)
+
+            # Mock the expanduser for this test
+            original_expanduser = os.path.expanduser
+            os.path.expanduser = (
+                lambda path: path.replace("~", fake_home)
+                if path.startswith("~")
+                else path
+            )
+
+            try:
+                tilde_dir = "~/test_tilde_dir"
+                try_to_create_directory(tilde_dir)
+                expected_path = os.path.join(fake_home, "test_tilde_dir")
+                assert os.path.exists(
+                    expected_path
+                ), "Tilde-expanded directory should be created"
+            finally:
+                # Restore original expanduser
+                os.path.expanduser = original_expanduser
+
+
+class TestLoadClass:
+    """Tests for the load_class utility function."""
+
+    def test_load_builtin_class(self):
+        """Test loading a builtin class."""
+        list_class = load_class("builtins.list")
+        assert list_class is list, "Should load the builtin list class"
+
+    def test_load_module(self):
+        """Test loading a module."""
+        path_module = load_class("os.path")
+        import os.path
+
+        assert path_module is os.path, "Should load os.path module"
+
+    def test_load_function(self):
+        """Test loading a function from a module."""
+        makedirs_func = load_class("os.makedirs")
+        assert makedirs_func is os.makedirs, "Should load os.makedirs function"
+
+    def test_load_standard_library_class(self):
+        """Test loading a standard library class."""
+        temp_dir_class = load_class("tempfile.TemporaryDirectory")
+        assert (
+            temp_dir_class is tempfile.TemporaryDirectory
+        ), "Should load TemporaryDirectory class"
+
+    def test_load_nested_module_class(self):
+        """Test loading a class from a nested module."""
+        datetime_class = load_class("datetime.datetime")
+        import datetime
+
+        assert (
+            datetime_class is datetime.datetime
+        ), "Should load datetime.datetime class"
+
+    def test_invalid_path_error(self):
+        """Test error handling for invalid paths."""
+        with pytest.raises(ValueError, match="valid path like mymodule.provider_class"):
+            load_class("invalid")
+
+    def test_nonexistent_module_error(self):
+        """Test error handling for nonexistent modules."""
+        with pytest.raises((ImportError, ModuleNotFoundError)):
+            load_class("nonexistent_module.SomeClass")
+
+    def test_nonexistent_attribute_error(self):
+        """Test error handling for nonexistent attributes."""
+        with pytest.raises(AttributeError):
+            load_class("os.NonexistentClass")
+
+
+class TestGetSystemMemory:
+    """Tests for the get_system_memory utility function."""
+
+    @pytest.mark.skipif(not PSUTIL_AVAILABLE, reason="psutil not available")
+    def test_cgroups_v1_with_low_limit(self):
+        """Test cgroups v1 with low memory limit."""
+        with tempfile.NamedTemporaryFile("w") as memory_limit_file:
+            memory_limit_file.write("1073741824")  # 1GB
+            memory_limit_file.flush()
+            memory = get_system_memory(
+                memory_limit_filename=memory_limit_file.name,
+                memory_limit_filename_v2="__does_not_exist__",
+            )
+            assert memory == 1073741824, "Should return cgroup limit when low"
+
+    @pytest.mark.skipif(not PSUTIL_AVAILABLE, reason="psutil not available")
+    def test_cgroups_v1_with_high_limit(self):
+        """Test cgroups v1 with high memory limit (should fallback to psutil)."""
+        with tempfile.NamedTemporaryFile("w") as memory_limit_file:
+            memory_limit_file.write(str(2**63))  # Very high limit
+            memory_limit_file.flush()
+            psutil_memory = psutil.virtual_memory().total
+            memory = get_system_memory(
+                memory_limit_filename=memory_limit_file.name,
+                memory_limit_filename_v2="__does_not_exist__",
+            )
+            assert (
+                memory == psutil_memory
+            ), "Should fallback to psutil when cgroup limit is very high"
+
+    @pytest.mark.skipif(not PSUTIL_AVAILABLE, reason="psutil not available")
+    def test_cgroups_v2_with_limit(self):
+        """Test cgroups v2 with memory limit set."""
+        with tempfile.NamedTemporaryFile("w") as memory_max_file:
+            memory_max_file.write("2147483648\n")  # 2GB with newline
+            memory_max_file.flush()
+            memory = get_system_memory(
+                memory_limit_filename="__does_not_exist__",
+                memory_limit_filename_v2=memory_max_file.name,
+            )
+            assert memory == 2147483648, "Should return cgroups v2 limit"
+
+    @pytest.mark.skipif(not PSUTIL_AVAILABLE, reason="psutil not available")
+    def test_cgroups_v2_unlimited(self):
+        """Test cgroups v2 with unlimited memory (max)."""
+        with tempfile.NamedTemporaryFile("w") as memory_max_file:
+            memory_max_file.write("max")
+            memory_max_file.flush()
+            psutil_memory = psutil.virtual_memory().total
+            memory = get_system_memory(
+                memory_limit_filename="__does_not_exist__",
+                memory_limit_filename_v2=memory_max_file.name,
+            )
+            assert (
+                memory == psutil_memory
+            ), "Should fallback to psutil when cgroups v2 is unlimited"
+
+    @pytest.mark.skipif(not PSUTIL_AVAILABLE, reason="psutil not available")
+    def test_no_cgroup_files(self):
+        """Test fallback to psutil when no cgroup files exist."""
+        psutil_memory = psutil.virtual_memory().total
+        memory = get_system_memory(
+            memory_limit_filename="__does_not_exist__",
+            memory_limit_filename_v2="__also_does_not_exist__",
+        )
+        assert memory == psutil_memory, "Should use psutil when no cgroup files exist"
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-sv", __file__]))

python/ray/_common/tests/test_wait_for_condition.py

@@ -0,0 +1,320 @@
+import asyncio
+import sys
+import time
+
+import pytest
+
+from ray._common.test_utils import async_wait_for_condition, wait_for_condition
+
+
+class TestWaitForCondition:
+    """Tests for the synchronous wait_for_condition function."""
+
+    def test_immediate_true_condition(self):
+        """Test that function returns immediately when condition is already true."""
+
+        def always_true():
+            return True
+
+        wait_for_condition(always_true, timeout=5)
+
+    def test_condition_becomes_true(self):
+        """Test waiting for a condition that becomes true after some time."""
+        counter = {"value": 0}
+
+        def condition():
+            counter["value"] += 1
+            return counter["value"] >= 3
+
+        wait_for_condition(condition, timeout=5, retry_interval_ms=50)
+
+        assert counter["value"] >= 3
+
+    def test_timeout_raises_runtime_error(self):
+        """Test that timeout raises RuntimeError with appropriate message."""
+
+        def always_false():
+            return False
+
+        with pytest.raises(RuntimeError) as exc_info:
+            wait_for_condition(always_false, timeout=0.2, retry_interval_ms=50)
+
+        assert "condition wasn't met before the timeout expired" in str(exc_info.value)
+
+    def test_condition_with_kwargs(self):
+        """Test passing kwargs to the condition predictor."""
+
+        def condition_with_args(target, current=0):
+            return current >= target
+
+        wait_for_condition(condition_with_args, timeout=1, target=5, current=10)
+
+        # Should not raise an exception since current >= target
+
+    def test_exception_handling_default(self):
+        """Test that exceptions are caught by default and timeout occurs."""
+
+        def failing_condition():
+            raise ValueError("Test exception")
+
+        with pytest.raises(RuntimeError) as exc_info:
+            wait_for_condition(failing_condition, timeout=0.2, retry_interval_ms=50)
+
+        error_msg = str(exc_info.value)
+        assert "condition wasn't met before the timeout expired" in error_msg
+        assert "Last exception:" in error_msg
+        assert "ValueError: Test exception" in error_msg
+
+    def test_exception_handling_raise_true(self):
+        """Test that exceptions are raised when raise_exceptions=True."""
+
+        def failing_condition():
+            raise ValueError("Test exception")
+
+        with pytest.raises(ValueError) as exc_info:
+            wait_for_condition(failing_condition, timeout=1, raise_exceptions=True)
+
+        assert "Test exception" in str(exc_info.value)
+
+    def test_custom_retry_interval(self):
+        """Test that custom retry intervals are respected."""
+        call_times = []
+
+        def condition():
+            call_times.append(time.time())
+            return len(call_times) >= 3
+
+        wait_for_condition(condition, timeout=5, retry_interval_ms=200)
+
+        # Verify that calls were spaced approximately 200ms apart
+        if len(call_times) >= 2:
+            interval = call_times[1] - call_times[0]
+            assert 0.15 <= interval <= 0.25  # Allow some tolerance
+
+    def test_condition_with_mixed_results(self):
+        """Test condition that fails initially then succeeds."""
+        attempts = {"count": 0}
+
+        def intermittent_condition():
+            attempts["count"] += 1
+            # Succeed on the 4th attempt
+            return attempts["count"] >= 4
+
+        wait_for_condition(intermittent_condition, timeout=2, retry_interval_ms=100)
+        assert attempts["count"] >= 4
+
+
+class TestAsyncWaitForCondition:
+    """Tests for the asynchronous async_wait_for_condition function."""
+
+    @pytest.mark.asyncio
+    async def test_immediate_true_condition(self):
+        """Test that function returns immediately when condition is already true."""
+
+        def always_true():
+            return True
+
+        await async_wait_for_condition(always_true, timeout=5)
+
+    @pytest.mark.asyncio
+    async def test_async_condition_becomes_true(self):
+        """Test waiting for an async condition that becomes true after some time."""
+        counter = {"value": 0}
+
+        async def async_condition():
+            counter["value"] += 1
+            await asyncio.sleep(0.01)  # Small async operation
+            return counter["value"] >= 3
+
+        await async_wait_for_condition(async_condition, timeout=5, retry_interval_ms=50)
+
+        assert counter["value"] >= 3
+
+    @pytest.mark.asyncio
+    async def test_sync_condition_becomes_true(self):
+        """Test waiting for a sync condition in async context."""
+        counter = {"value": 0}
+
+        def sync_condition():
+            counter["value"] += 1
+            return counter["value"] >= 3
+
+        await async_wait_for_condition(sync_condition, timeout=5, retry_interval_ms=50)
+        assert counter["value"] >= 3
+
+    @pytest.mark.asyncio
+    async def test_timeout_raises_runtime_error(self):
+        """Test that timeout raises RuntimeError with appropriate message."""
+
+        def always_false():
+            return False
+
+        with pytest.raises(RuntimeError) as exc_info:
+            await async_wait_for_condition(
+                always_false, timeout=0.2, retry_interval_ms=50
+            )
+
+        assert "condition wasn't met before the timeout expired" in str(exc_info.value)
+
+    @pytest.mark.asyncio
+    async def test_condition_with_kwargs(self):
+        """Test passing kwargs to the condition predictor."""
+
+        def condition_with_args(target, current=0):
+            return current >= target
+
+        await async_wait_for_condition(
+            condition_with_args, timeout=1, target=5, current=10
+        )
+
+        # Should not raise an exception since current >= target
+
+    @pytest.mark.asyncio
+    async def test_async_condition_with_kwargs(self):
+        """Test passing kwargs to an async condition predictor."""
+
+        async def async_condition_with_args(target, current=0):
+            await asyncio.sleep(0.01)
+            return current >= target
+
+        await async_wait_for_condition(
+            async_condition_with_args, timeout=1, target=5, current=10
+        )
+
+        # Should not raise an exception since current >= target
+
+    @pytest.mark.asyncio
+    async def test_exception_handling(self):
+        """Test that exceptions are caught and timeout occurs."""
+
+        def failing_condition():
+            raise ValueError("Test exception")
+
+        with pytest.raises(RuntimeError) as exc_info:
+            await async_wait_for_condition(
+                failing_condition, timeout=0.2, retry_interval_ms=50
+            )
+
+        error_msg = str(exc_info.value)
+        assert "condition wasn't met before the timeout expired" in error_msg
+        assert "Last exception:" in error_msg
+
+    @pytest.mark.asyncio
+    async def test_async_exception_handling(self):
+        """Test that exceptions from async conditions are caught."""
+
+        async def async_failing_condition():
+            await asyncio.sleep(0.01)
+            raise ValueError("Async test exception")
+
+        with pytest.raises(RuntimeError) as exc_info:
+            await async_wait_for_condition(
+                async_failing_condition, timeout=0.2, retry_interval_ms=50
+            )
+
+        error_msg = str(exc_info.value)
+        assert "condition wasn't met before the timeout expired" in error_msg
+        assert "Last exception:" in error_msg
+
+    @pytest.mark.asyncio
+    async def test_custom_retry_interval(self):
+        """Test that custom retry intervals are respected."""
+        call_times = []
+
+        def condition():
+            call_times.append(time.time())
+            return len(call_times) >= 3
+
+        await async_wait_for_condition(condition, timeout=5, retry_interval_ms=200)
+
+        # Verify that calls were spaced approximately 200ms apart
+        if len(call_times) >= 2:
+            interval = call_times[1] - call_times[0]
+            assert 0.15 <= interval <= 0.25  # Allow some tolerance
+
+    @pytest.mark.asyncio
+    async def test_mixed_sync_async_conditions(self):
+        """Test that both sync and async conditions work in the same test."""
+        sync_counter = {"value": 0}
+        async_counter = {"value": 0}
+
+        def sync_condition():
+            sync_counter["value"] += 1
+            return sync_counter["value"] >= 2
+
+        async def async_condition():
+            async_counter["value"] += 1
+            await asyncio.sleep(0.01)
+            return async_counter["value"] >= 2
+
+        # Test sync condition
+        await async_wait_for_condition(sync_condition, timeout=2, retry_interval_ms=50)
+        assert sync_counter["value"] >= 2
+
+        # Test async condition
+        await async_wait_for_condition(async_condition, timeout=2, retry_interval_ms=50)
+        assert async_counter["value"] >= 2
+
+
+class TestEdgeCases:
+    """Tests for edge cases and boundary conditions."""
+
+    def test_zero_timeout(self):
+        """Test behavior with zero timeout."""
+
+        def slow_condition():
+            time.sleep(0.1)
+            return True
+
+        with pytest.raises(RuntimeError):
+            wait_for_condition(slow_condition, timeout=0, retry_interval_ms=50)
+
+    @pytest.mark.asyncio
+    async def test_async_zero_timeout(self):
+        """Test async behavior with zero timeout."""
+
+        async def slow_condition():
+            await asyncio.sleep(0.1)
+            return True
+
+        with pytest.raises(RuntimeError):
+            await async_wait_for_condition(
+                slow_condition, timeout=0, retry_interval_ms=50
+            )
+
+    def test_very_small_retry_interval(self):
+        """Test with very small retry interval."""
+        counter = {"value": 0}
+
+        def condition():
+            counter["value"] += 1
+            return counter["value"] >= 5
+
+        start_time = time.time()
+        wait_for_condition(condition, timeout=1, retry_interval_ms=1)
+        elapsed = time.time() - start_time
+
+        # Should complete quickly due to small retry interval
+        assert elapsed < 0.5
+        assert counter["value"] >= 5
+
+    @pytest.mark.asyncio
+    async def test_async_very_small_retry_interval(self):
+        """Test async version with very small retry interval."""
+        counter = {"value": 0}
+
+        def condition():
+            counter["value"] += 1
+            return counter["value"] >= 5
+
+        start_time = time.time()
+        await async_wait_for_condition(condition, timeout=1, retry_interval_ms=1)
+        elapsed = time.time() - start_time
+
+        # Should complete quickly due to small retry interval
+        assert elapsed < 0.5
+        assert counter["value"] >= 5
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-sv", __file__]))

python/ray/_common/tls_utils.py

@@ -0,0 +1,101 @@
+"""TLS utilities shared across Ray libraries (e.g. Serve)."""
+
+import datetime
+import os
+import socket
+from typing import Tuple
+
+from ray._common.network_utils import (
+    get_localhost_ip,
+    node_ip_address_from_perspective,
+)
+
+
+def generate_self_signed_tls_certs() -> Tuple[str, str]:
+    """Create self-signed key/cert pair for testing.
+
+    Returns:
+        Tuple of (cert_pem_contents, key_pem_contents).
+
+    Raises:
+        ImportError: If the ``cryptography`` library is not installed.
+    """
+    try:
+        from cryptography import x509
+        from cryptography.hazmat.backends import default_backend
+        from cryptography.hazmat.primitives import hashes, serialization
+        from cryptography.hazmat.primitives.asymmetric import rsa
+        from cryptography.x509.oid import NameOID
+    except ImportError as e:
+        raise ImportError(
+            "Using self-signed TLS certs requires `cryptography`. "
+            "Install it with: pip install cryptography"
+        ) from e
+
+    key = rsa.generate_private_key(
+        public_exponent=65537, key_size=2048, backend=default_backend()
+    )
+    key_contents = key.private_bytes(
+        encoding=serialization.Encoding.PEM,
+        format=serialization.PrivateFormat.PKCS8,
+        encryption_algorithm=serialization.NoEncryption(),
+    ).decode()
+
+    subject = x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, "ray-internal")])
+    altnames = x509.SubjectAlternativeName(
+        [
+            x509.DNSName(socket.gethostbyname(socket.gethostname())),
+            x509.DNSName(get_localhost_ip()),
+            x509.DNSName(node_ip_address_from_perspective()),
+            x509.DNSName("localhost"),
+        ]
+    )
+    now = datetime.datetime.utcnow()
+    cert = (
+        x509.CertificateBuilder()
+        .subject_name(subject)
+        .issuer_name(subject)
+        .add_extension(altnames, critical=False)
+        .public_key(key.public_key())
+        .serial_number(x509.random_serial_number())
+        .not_valid_before(now)
+        .not_valid_after(now + datetime.timedelta(days=365))
+        .sign(key, hashes.SHA256(), default_backend())
+    )
+
+    cert_contents = cert.public_bytes(serialization.Encoding.PEM).decode()
+    return cert_contents, key_contents
+
+
+def add_port_to_grpc_server(server, address):
+    import grpc
+
+    if os.environ.get("RAY_USE_TLS", "0").lower() in ("1", "true"):
+        server_cert_chain, private_key, ca_cert = load_certs_from_env()
+        credentials = grpc.ssl_server_credentials(
+            [(private_key, server_cert_chain)],
+            root_certificates=ca_cert,
+            require_client_auth=ca_cert is not None,
+        )
+        return server.add_secure_port(address, credentials)
+    else:
+        return server.add_insecure_port(address)
+
+
+def load_certs_from_env():
+    tls_env_vars = ["RAY_TLS_SERVER_CERT", "RAY_TLS_SERVER_KEY", "RAY_TLS_CA_CERT"]
+    if any(v not in os.environ for v in tls_env_vars):
+        raise RuntimeError(
+            "If the environment variable RAY_USE_TLS is set to true "
+            "then RAY_TLS_SERVER_CERT, RAY_TLS_SERVER_KEY and "
+            "RAY_TLS_CA_CERT must also be set."
+        )
+
+    with open(os.environ["RAY_TLS_SERVER_CERT"], "rb") as f:
+        server_cert_chain = f.read()
+    with open(os.environ["RAY_TLS_SERVER_KEY"], "rb") as f:
+        private_key = f.read()
+    with open(os.environ["RAY_TLS_CA_CERT"], "rb") as f:
+        ca_cert = f.read()
+
+    return server_cert_chain, private_key, ca_cert

python/ray/_common/usage/usage_constants.py

@@ -0,0 +1,63 @@
+SCHEMA_VERSION = "0.1"
+
+# The key to store / obtain cluster metadata.
+CLUSTER_METADATA_KEY = b"CLUSTER_METADATA"
+
+# The name of a json file where usage stats will be written.
+USAGE_STATS_FILE = "usage_stats.json"
+
+USAGE_STATS_ENABLED_ENV_VAR = "RAY_USAGE_STATS_ENABLED"
+
+USAGE_STATS_SOURCE_ENV_VAR = "RAY_USAGE_STATS_SOURCE"
+
+USAGE_STATS_SOURCE_OSS = "OSS"
+
+USAGE_STATS_ENABLED_FOR_CLI_MESSAGE = (
+    "Usage stats collection is enabled. To disable this, add `--disable-usage-stats` "
+    "to the command that starts the cluster, or run the following command:"
+    " `ray disable-usage-stats` before starting the cluster. "
+    "See https://docs.ray.io/en/master/cluster/usage-stats.html for more details."
+)
+
+USAGE_STATS_ENABLED_FOR_RAY_INIT_MESSAGE = (
+    "Usage stats collection is enabled. To disable this, run the following command:"
+    " `ray disable-usage-stats` before starting Ray. "
+    "See https://docs.ray.io/en/master/cluster/usage-stats.html for more details."
+)
+
+USAGE_STATS_DISABLED_MESSAGE = "Usage stats collection is disabled."
+
+USAGE_STATS_ENABLED_BY_DEFAULT_FOR_CLI_MESSAGE = (
+    "Usage stats collection is enabled by default without user confirmation "
+    "because this terminal is detected to be non-interactive. "
+    "To disable this, add `--disable-usage-stats` to the command that starts "
+    "the cluster, or run the following command:"
+    " `ray disable-usage-stats` before starting the cluster. "
+    "See https://docs.ray.io/en/master/cluster/usage-stats.html for more details."
+)
+
+USAGE_STATS_ENABLED_BY_DEFAULT_FOR_RAY_INIT_MESSAGE = (
+    "Usage stats collection is enabled by default for nightly wheels. "
+    "To disable this, run the following command:"
+    " `ray disable-usage-stats` before starting Ray. "
+    "See https://docs.ray.io/en/master/cluster/usage-stats.html for more details."
+)
+
+USAGE_STATS_CONFIRMATION_MESSAGE = (
+    "Enable usage stats collection? "
+    "This prompt will auto-proceed in 10 seconds to avoid blocking cluster startup."
+)
+
+LIBRARY_USAGE_SET_NAME = "library_usage_"
+
+HARDWARE_USAGE_SET_NAME = "hardware_usage_"
+
+# Keep in-sync with the same constants defined in usage_stats_client.h
+EXTRA_USAGE_TAG_PREFIX = "extra_usage_tag_"
+USAGE_STATS_NAMESPACE = "usage_stats"
+
+KUBERNETES_SERVICE_HOST_ENV = "KUBERNETES_SERVICE_HOST"
+KUBERAY_ENV = "RAY_USAGE_STATS_KUBERAY_IN_USE"
+
+PROVIDER_KUBERNETES_GENERIC = "kubernetes"
+PROVIDER_KUBERAY = "kuberay"

python/ray/_common/utils.py

@@ -0,0 +1,491 @@
+import asyncio
+import binascii
+import errno
+import importlib
+import inspect
+import logging
+import os
+import random
+import string
+import sys
+import tempfile
+import time
+from abc import ABC, abstractmethod
+from inspect import signature
+from types import ModuleType
+from typing import Any, Coroutine, Dict, Optional, Tuple
+
+import ray
+from ray._raylet import GcsClient, NodeID
+from ray.core.generated.gcs_pb2 import GcsNodeInfo
+from ray.core.generated.gcs_service_pb2 import GetAllNodeInfoRequest
+
+import psutil
+
+logger = logging.getLogger(__name__)
+
+
+def env_integer(key, default):
+    if key in os.environ:
+        value = os.environ[key]
+        try:
+            return int(value)
+        except ValueError:
+            logger.debug(
+                f"Found {key} in environment, but value must "
+                f"be an integer. Got: {value}. Returning "
+                f"provided default {default}."
+            )
+            return default
+    return default
+
+
+def env_float(key, default):
+    if key in os.environ:
+        value = os.environ[key]
+        try:
+            return float(value)
+        except ValueError:
+            logger.debug(
+                f"Found {key} in environment, but value must "
+                f"be a float. Got: {value}. Returning "
+                f"provided default {default}."
+            )
+            return default
+    return default
+
+
+def env_bool(key, default):
+    if key in os.environ:
+        val = os.environ[key].lower()
+        return val == "true" or val == "1"
+    return default
+
+
+def import_module_and_attr(
+    full_path: str, *, reload_module: bool = False
+) -> Tuple[ModuleType, Any]:
+    """Given a full import path to a module attr, return the imported module and attr.
+
+    If `reload_module` is set, the module will be reloaded using `importlib.reload`.
+
+    Args:
+        full_path: The full import path to the module and attr.
+        reload_module: Whether to reload the module.
+
+    Returns:
+        A tuple of the imported module and attr.
+    """
+    if ":" in full_path:
+        if full_path.count(":") > 1:
+            raise ValueError(
+                f'Got invalid import path "{full_path}". An '
+                "import path may have at most one colon."
+            )
+        module_name, attr_name = full_path.split(":")
+    else:
+        last_period_idx = full_path.rfind(".")
+        module_name = full_path[:last_period_idx]
+        attr_name = full_path[last_period_idx + 1 :]
+    module = importlib.import_module(module_name)
+    if reload_module:
+        importlib.reload(module)
+    return module, getattr(module, attr_name)
+
+
+def import_attr(full_path: str, *, reload_module: bool = False) -> Any:
+    """Given a full import path to a module attr, return the imported attr.
+
+    If `reload_module` is set, the module will be reloaded using `importlib.reload`.
+
+    For example, the following are equivalent:
+        MyClass = import_attr("module.submodule:MyClass")
+        MyClass = import_attr("module.submodule.MyClass")
+        from module.submodule import MyClass
+
+    Args:
+        full_path: The full import path to the module and attr.
+        reload_module: Whether to reload the module.
+
+    Returns:
+        Imported attr
+    """
+    return import_module_and_attr(full_path, reload_module=reload_module)[1]
+
+
+def get_or_create_event_loop() -> asyncio.AbstractEventLoop:
+    """Get a running async event loop if one exists, otherwise create one.
+
+    This function serves as a proxy for the deprecating get_event_loop().
+    It tries to get the running loop first, and if no running loop
+    could be retrieved:
+    - For python version <3.10: it falls back to the get_event_loop
+        call.
+    - For python version >= 3.10: it uses the same python implementation
+        of _get_event_loop() at asyncio/events.py.
+
+    Ideally, one should use high level APIs like asyncio.run() with python
+    version >= 3.7, if not possible, one should create and manage the event
+    loops explicitly.
+    """
+    vers_info = sys.version_info
+    if vers_info.major >= 3 and vers_info.minor >= 10:
+        # This follows the implementation of the deprecating `get_event_loop`
+        # in python3.10's asyncio. See python3.10/asyncio/events.py
+        # _get_event_loop()
+        try:
+            loop = asyncio.get_running_loop()
+            assert loop is not None
+            return loop
+        except RuntimeError as e:
+            # No running loop, relying on the error message as for now to
+            # differentiate runtime errors.
+            assert "no running event loop" in str(e)
+            try:
+                loop = asyncio.get_event_loop_policy().get_event_loop()
+                return loop
+            except RuntimeError:
+                # Python 3.14+: get_event_loop() no longer creates a loop automatically
+                # See: https://docs.python.org/3.14/library/asyncio-eventloop.html
+                loop = asyncio.new_event_loop()
+                asyncio.set_event_loop(loop)
+                return loop
+
+    return asyncio.get_event_loop()
+
+
+_BACKGROUND_TASKS = set()
+
+
+def run_background_task(coroutine: Coroutine) -> asyncio.Task:
+    """Schedule a task reliably to the event loop.
+
+    This API is used when you don't want to cache the reference of `asyncio.Task`.
+    For example,
+
+    ```
+    get_event_loop().create_task(coroutine(*args))
+    ```
+
+    The above code doesn't guarantee to schedule the coroutine to the event loops
+
+    When using create_task in a  "fire and forget" way, we should keep the references
+    alive for the reliable execution. This API is used to fire and forget
+    asynchronous execution.
+
+    https://docs.python.org/3/library/asyncio-task.html#creating-tasks
+    """
+    task = get_or_create_event_loop().create_task(coroutine)
+    # Add task to the set. This creates a strong reference.
+    _BACKGROUND_TASKS.add(task)
+
+    # To prevent keeping references to finished tasks forever,
+    # make each task remove its own reference from the set after
+    # completion:
+    task.add_done_callback(_BACKGROUND_TASKS.discard)
+    return task
+
+
+# Used in gpu detection
+RESOURCE_CONSTRAINT_PREFIX = "accelerator_type:"
+PLACEMENT_GROUP_BUNDLE_RESOURCE_NAME = "bundle"
+
+
+def resources_from_ray_options(options_dict: Dict[str, Any]) -> Dict[str, Any]:
+    """Determine a task's resource requirements.
+
+    Args:
+        options_dict: The dictionary that contains resources requirements.
+
+    Returns:
+        A dictionary of the resource requirements for the task.
+    """
+    resources = (options_dict.get("resources") or {}).copy()
+
+    if "CPU" in resources or "GPU" in resources:
+        raise ValueError(
+            "The resources dictionary must not contain the key 'CPU' or 'GPU'"
+        )
+    elif "memory" in resources or "object_store_memory" in resources:
+        raise ValueError(
+            "The resources dictionary must not "
+            "contain the key 'memory' or 'object_store_memory'"
+        )
+    elif PLACEMENT_GROUP_BUNDLE_RESOURCE_NAME in resources:
+        raise ValueError(
+            "The resource should not include `bundle` which "
+            f"is reserved for Ray. resources: {resources}"
+        )
+
+    num_cpus = options_dict.get("num_cpus")
+    num_gpus = options_dict.get("num_gpus")
+    memory = options_dict.get("memory")
+    object_store_memory = options_dict.get("object_store_memory")
+    accelerator_type = options_dict.get("accelerator_type")
+
+    if num_cpus is not None:
+        resources["CPU"] = num_cpus
+    if num_gpus is not None:
+        resources["GPU"] = num_gpus
+    if memory is not None:
+        resources["memory"] = int(memory)
+    if object_store_memory is not None:
+        resources["object_store_memory"] = object_store_memory
+    if accelerator_type is not None:
+        resources[f"{RESOURCE_CONSTRAINT_PREFIX}{accelerator_type}"] = 0.001
+
+    return resources
+
+
+# Match the standard alphabet used for UUIDs.
+RANDOM_STRING_ALPHABET = string.ascii_lowercase + string.digits
+
+
+def get_random_alphanumeric_string(length: int):
+    """Generates random string of length consisting exclusively of
+    - Lower-case ASCII chars
+    - Digits
+    """
+    return "".join(random.choices(RANDOM_STRING_ALPHABET, k=length))
+
+
+_PRINTED_WARNING = set()
+
+
+def get_call_location(back: int = 1):
+    """
+    Get the location (filename and line number) of a function caller, `back`
+    frames up the stack.
+
+    Args:
+        back: The number of frames to go up the stack, not including this
+            function.
+
+    Returns:
+        A string with the filename and line number of the caller.
+        For example, "myfile.py:123".
+    """
+    stack = inspect.stack()
+    try:
+        frame = stack[back + 1]
+        return f"{frame.filename}:{frame.lineno}"
+    except IndexError:
+        return "UNKNOWN"
+
+
+def resolve_user_ray_temp_dir(gcs_client: GcsClient, node_id: str):
+    """
+    Get the ray temp directory.
+
+    If a temp dir was specified for this node, this function will
+    retrieve the information from GCS. Otherwise, it will fallback to the
+    default ray temp directory.
+
+    Args:
+        gcs_client: The GCS client.
+        node_id: The ID of the node to fetch the temp dir for.
+                 E.g.: "1a9904d8aa3de65367830e2aef6313a5b2e9d4b0e3725e0dceeacb1b"
+                        (hex string representation of the node ID)
+
+    Returns:
+        The path to the ray temp directory.
+    """
+    # check if temp dir is available from runtime context
+    if ray.is_initialized() and ray.get_runtime_context().get_node_id() == node_id:
+        return ray.get_runtime_context().get_temp_dir()
+
+    # Fetch temp dir as specified by --temp-dir at creation time.
+    try:
+        # Create node selector for node_id filter
+        node_selector = GetAllNodeInfoRequest.NodeSelector()
+        node_selector.node_id = NodeID.from_hex(node_id).binary()
+
+        node_infos = gcs_client.get_all_node_info(
+            node_selectors=[node_selector],
+            state_filter=GcsNodeInfo.GcsNodeState.ALIVE,
+        ).values()
+    except Exception as e:
+        raise Exception(
+            f"Failed to get node info from GCS when fetching tempdir for node {node_id}: {e}"
+        )
+    if not node_infos:
+        raise Exception(
+            f"No node info associated with ALIVE state found for node {node_id} in GCS"
+        )
+
+    node_info = next(iter(node_infos))
+    if node_info is not None:
+        temp_dir = getattr(node_info, "temp_dir", None)
+        if temp_dir is not None:
+            return temp_dir
+        else:
+            raise Exception(
+                "Node temp_dir was not found in NodeInfo. did the node's raylet start successfully?"
+            )
+
+
+def get_default_system_temp_dir():
+    if "RAY_TMPDIR" in os.environ:
+        return os.environ["RAY_TMPDIR"]
+    elif sys.platform.startswith("linux") and "TMPDIR" in os.environ:
+        return os.environ["TMPDIR"]
+    elif sys.platform.startswith("darwin") or sys.platform.startswith("linux"):
+        # Ideally we wouldn't need this fallback, but keep it for now for
+        # for compatibility
+        tempdir = os.path.join(os.sep, "tmp")
+    else:
+        tempdir = tempfile.gettempdir()
+
+    return tempdir
+
+
+def get_default_ray_temp_dir():
+    return os.path.join(get_default_system_temp_dir(), "ray")
+
+
+def get_ray_address_file(temp_dir: Optional[str]):
+    if temp_dir is None:
+        temp_dir = get_default_ray_temp_dir()
+    return os.path.join(temp_dir, "ray_current_cluster")
+
+
+def reset_ray_address(temp_dir: Optional[str] = None):
+    address_file = get_ray_address_file(temp_dir)
+    if os.path.exists(address_file):
+        try:
+            os.remove(address_file)
+        except OSError:
+            pass
+
+
+def load_class(path):
+    """Load a class at runtime given a full path.
+
+    Example of the path: mypkg.mysubpkg.myclass
+    """
+    class_data = path.split(".")
+    if len(class_data) < 2:
+        raise ValueError("You need to pass a valid path like mymodule.provider_class")
+    module_path = ".".join(class_data[:-1])
+    class_str = class_data[-1]
+    module = importlib.import_module(module_path)
+    return getattr(module, class_str)
+
+
+def get_system_memory(
+    # For cgroups v1:
+    memory_limit_filename: str = "/sys/fs/cgroup/memory/memory.limit_in_bytes",
+    # For cgroups v2:
+    memory_limit_filename_v2: str = "/sys/fs/cgroup/memory.max",
+):
+    """Return the total amount of system memory in bytes.
+
+    Args:
+        memory_limit_filename: The path to the file that contains the memory
+            limit for the Docker container. Defaults to
+            /sys/fs/cgroup/memory/memory.limit_in_bytes.
+        memory_limit_filename_v2: The path to the file that contains the memory
+            limit for the Docker container in cgroups v2. Defaults to
+            /sys/fs/cgroup/memory.max.
+
+    Returns:
+        The total amount of system memory in bytes.
+    """
+    # Try to accurately figure out the memory limit if we are in a docker
+    # container. Note that this file is not specific to Docker and its value is
+    # often much larger than the actual amount of memory.
+    docker_limit = None
+    if os.path.exists(memory_limit_filename):
+        with open(memory_limit_filename, "r") as f:
+            docker_limit = int(f.read().strip())
+    elif os.path.exists(memory_limit_filename_v2):
+        with open(memory_limit_filename_v2, "r") as f:
+            # Don't forget to strip() the newline:
+            max_file = f.read().strip()
+            if max_file.isnumeric():
+                docker_limit = int(max_file)
+            else:
+                # max_file is "max", i.e. is unset.
+                docker_limit = None
+
+    # Use psutil if it is available.
+    psutil_memory_in_bytes = psutil.virtual_memory().total
+
+    if docker_limit is not None:
+        # We take the min because the cgroup limit is very large if we aren't
+        # in Docker.
+        return min(docker_limit, psutil_memory_in_bytes)
+
+    return psutil_memory_in_bytes
+
+
+def binary_to_hex(identifier):
+    hex_identifier = binascii.hexlify(identifier)
+    hex_identifier = hex_identifier.decode()
+    return hex_identifier
+
+
+def hex_to_binary(hex_identifier):
+    return binascii.unhexlify(hex_identifier)
+
+
+def try_make_directory_shared(directory_path):
+    try:
+        os.chmod(directory_path, 0o0777)
+    except OSError as e:
+        # Silently suppress the PermissionError that is thrown by the chmod.
+        # This is done because the user attempting to change the permissions
+        # on a directory may not own it. The chmod is attempted whether the
+        # directory is new or not to avoid race conditions.
+        # ray-project/ray/#3591
+        if e.errno in [errno.EACCES, errno.EPERM]:
+            pass
+        else:
+            raise
+
+
+def try_to_create_directory(directory_path):
+    # Attempt to create a directory that is globally readable/writable.
+    directory_path = os.path.expanduser(directory_path)
+    os.makedirs(directory_path, exist_ok=True)
+    # Change the log directory permissions so others can use it. This is
+    # important when multiple people are using the same machine.
+    try_make_directory_shared(directory_path)
+
+
+def get_function_args(callable):
+    all_parameters = frozenset(signature(callable).parameters)
+    return list(all_parameters)
+
+
+def decode(byte_str: str, allow_none: bool = False, encode_type: str = "utf-8"):
+    """Make this unicode in Python 3, otherwise leave it as bytes.
+
+    Args:
+        byte_str: The byte string to decode.
+        allow_none: If true, then we will allow byte_str to be None in which
+            case we will return an empty string. TODO(rkn): Remove this flag.
+            This is only here to simplify upgrading to flatbuffers 1.10.0.
+        encode_type: The encoding type to use for decoding. Defaults to "utf-8".
+
+    Returns:
+        A byte string in Python 2 and a unicode string in Python 3.
+    """
+    if byte_str is None and allow_none:
+        return ""
+
+    if not isinstance(byte_str, bytes):
+        raise ValueError(f"The argument {byte_str} must be a bytes object.")
+    return byte_str.decode(encode_type)
+
+
+class TimerBase(ABC):
+    @abstractmethod
+    def time(self) -> float:
+        """Return the current time."""
+        raise NotImplementedError
+
+
+class Timer(TimerBase):
+    def time(self) -> float:
+        return time.time()

python/ray/_private/BUILD.bazel

@@ -0,0 +1,18 @@
+load("//bazel:python.bzl", "doctest")
+
+doctest(
+    files = glob(
+        ["**/*.py"],
+        exclude = ["**/thirdparty_files/**"],
+    ),
+    tags = ["team:core"],
+)
+
+filegroup(
+    name = "src_files",
+    srcs = glob(
+        ["**/*.py"],
+        exclude = ["**/thirdparty_files/**"],
+    ),
+    visibility = ["//:__pkg__"],
+)

python/ray/_private/accelerators/__init__.py

@@ -0,0 +1,92 @@
+from typing import Optional, Set
+
+from ray._private.accelerators.accelerator import (
+    RAY_ACCEL_ENV_VAR_OVERRIDE_ON_ZERO_ENV_VAR,
+    AcceleratorManager,
+)
+from ray._private.accelerators.amd_gpu import AMDGPUAcceleratorManager
+from ray._private.accelerators.furiosa import FuriosaAcceleratorManager
+from ray._private.accelerators.hpu import HPUAcceleratorManager
+from ray._private.accelerators.intel_gpu import IntelGPUAcceleratorManager
+from ray._private.accelerators.metax_gpu import MetaxGPUAcceleratorManager
+from ray._private.accelerators.neuron import NeuronAcceleratorManager
+from ray._private.accelerators.npu import NPUAcceleratorManager
+from ray._private.accelerators.nvidia_gpu import NvidiaGPUAcceleratorManager
+from ray._private.accelerators.rbln import RBLNAcceleratorManager
+from ray._private.accelerators.tpu import TPUAcceleratorManager
+
+
+def get_all_accelerator_managers() -> Set[AcceleratorManager]:
+    """Get all accelerator managers supported by Ray."""
+    return {
+        NvidiaGPUAcceleratorManager,
+        IntelGPUAcceleratorManager,
+        AMDGPUAcceleratorManager,
+        TPUAcceleratorManager,
+        NeuronAcceleratorManager,
+        HPUAcceleratorManager,
+        NPUAcceleratorManager,
+        RBLNAcceleratorManager,
+        MetaxGPUAcceleratorManager,
+        FuriosaAcceleratorManager,
+    }
+
+
+def get_all_accelerator_resource_names() -> Set[str]:
+    """Get all resource names for accelerators."""
+    return {
+        accelerator_manager.get_resource_name()
+        for accelerator_manager in get_all_accelerator_managers()
+    }
+
+
+def get_accelerator_manager_for_resource(
+    resource_name: str,
+) -> Optional[AcceleratorManager]:
+    """Get the corresponding accelerator manager for the given
+    accelerator resource name
+
+    E.g., TPUAcceleratorManager is returned if resource name is "TPU"
+    """
+    try:
+        return get_accelerator_manager_for_resource._resource_name_to_accelerator_manager.get(  # noqa: E501
+            resource_name, None
+        )
+    except AttributeError:
+        # Lazy initialization.
+        resource_name_to_accelerator_manager = {
+            accelerator_manager.get_resource_name(): accelerator_manager
+            for accelerator_manager in get_all_accelerator_managers()
+        }
+        # Special handling for GPU resource name since multiple accelerator managers
+        # have the same GPU resource name.
+        if AMDGPUAcceleratorManager.get_current_node_num_accelerators() > 0:
+            resource_name_to_accelerator_manager["GPU"] = AMDGPUAcceleratorManager
+        elif IntelGPUAcceleratorManager.get_current_node_num_accelerators() > 0:
+            resource_name_to_accelerator_manager["GPU"] = IntelGPUAcceleratorManager
+        elif MetaxGPUAcceleratorManager.get_current_node_num_accelerators() > 0:
+            resource_name_to_accelerator_manager["GPU"] = MetaxGPUAcceleratorManager
+        else:
+            resource_name_to_accelerator_manager["GPU"] = NvidiaGPUAcceleratorManager
+        get_accelerator_manager_for_resource._resource_name_to_accelerator_manager = (
+            resource_name_to_accelerator_manager
+        )
+        return resource_name_to_accelerator_manager.get(resource_name, None)
+
+
+__all__ = [
+    "NvidiaGPUAcceleratorManager",
+    "IntelGPUAcceleratorManager",
+    "AMDGPUAcceleratorManager",
+    "TPUAcceleratorManager",
+    "NeuronAcceleratorManager",
+    "HPUAcceleratorManager",
+    "NPUAcceleratorManager",
+    "RBLNAcceleratorManager",
+    "MetaxGPUAcceleratorManager",
+    "FuriosaAcceleratorManager",
+    "get_all_accelerator_managers",
+    "get_all_accelerator_resource_names",
+    "get_accelerator_manager_for_resource",
+    "RAY_ACCEL_ENV_VAR_OVERRIDE_ON_ZERO_ENV_VAR",
+]

python/ray/_private/accelerators/accelerator.py

@@ -0,0 +1,158 @@
+from abc import ABC, abstractmethod
+from typing import Dict, List, Optional, Tuple
+
+# https://github.com/ray-project/ray/issues/54868
+# Ray no longer overrides accelerator ids environment variables when the number
+# of accelerators is zero. For example, if a user sets `num_gpus=0` in
+# `ray.init()`, the environment variable `CUDA_VISIBLE_DEVICES` will not be
+# set to an empty string.
+#
+# Set RAY_ACCEL_ENV_VAR_OVERRIDE_ON_ZERO=1 to restore the old behavior where
+# Ray would override accelerator env vars even when zero accelerators are assigned.
+#
+RAY_ACCEL_ENV_VAR_OVERRIDE_ON_ZERO_ENV_VAR = "RAY_ACCEL_ENV_VAR_OVERRIDE_ON_ZERO"
+
+
+class AcceleratorManager(ABC):
+    """This class contains all the functions needed for supporting
+    an accelerator family in Ray."""
+
+    @staticmethod
+    @abstractmethod
+    def get_resource_name() -> str:
+        """Get the name of the resource representing this accelerator family.
+
+        Returns:
+            The resource name: e.g., the resource name for NVIDIA GPUs is "GPU"
+        """
+
+    @staticmethod
+    @abstractmethod
+    def get_visible_accelerator_ids_env_var() -> str:
+        """Get the env var that sets the ids of visible accelerators of this family.
+
+        Returns:
+            The env var for setting visible accelerator ids: e.g.,
+                CUDA_VISIBLE_DEVICES for NVIDIA GPUs.
+        """
+
+    @staticmethod
+    @abstractmethod
+    def get_current_node_num_accelerators() -> int:
+        """Get the total number of accelerators of this family on the current node.
+
+        Returns:
+            The detected total number of accelerators of this family.
+            Return 0 if the current node doesn't contain accelerators of this family.
+        """
+
+    @staticmethod
+    @abstractmethod
+    def get_current_node_accelerator_type() -> Optional[str]:
+        """Get the type of the accelerator of this family on the current node.
+
+        Currently Ray only supports single accelerator type of
+        an accelerator family on each node.
+
+        The result should only be used when get_current_node_num_accelerators() > 0.
+
+        Returns:
+            The detected accelerator type of this family: e.g., H100 for NVIDIA GPU.
+            Return None if it's unknown or the node doesn't have
+            accelerators of this family.
+        """
+
+    @staticmethod
+    @abstractmethod
+    def get_current_node_additional_resources() -> Optional[Dict[str, float]]:
+        """Get any additional resources required for the current node.
+
+        In case a particular accelerator type requires considerations for
+        additional resources (e.g. for TPUs, providing the TPU pod type and
+        TPU name), this function can be used to provide the
+        additional logical resources.
+
+        Returns:
+            A dictionary representing additional resources that may be
+            necessary for a particular accelerator type.
+        """
+
+    @staticmethod
+    @abstractmethod
+    def validate_resource_request_quantity(
+        quantity: float,
+    ) -> Tuple[bool, Optional[str]]:
+        """Validate the resource request quantity of this accelerator resource.
+
+        Args:
+            quantity: The resource request quantity to be validated.
+
+        Returns:
+            (valid, error_message) tuple: the first element of the tuple
+                indicates whether the given quantity is valid or not,
+                the second element is the error message
+                if the given quantity is invalid.
+        """
+
+    @staticmethod
+    @abstractmethod
+    def get_current_process_visible_accelerator_ids() -> Optional[List[str]]:
+        """Get the ids of accelerators of this family that are visible to the current process.
+
+        Returns:
+            The list of visiable accelerator ids.
+            Return None if all accelerators are visible.
+        """
+
+    @staticmethod
+    @abstractmethod
+    def set_current_process_visible_accelerator_ids(ids: List[str]) -> None:
+        """Set the ids of accelerators of this family that are visible to the current process.
+
+        Args:
+            ids: The ids of visible accelerators of this family.
+        """
+
+    @staticmethod
+    def get_ec2_instance_num_accelerators(
+        instance_type: str, instances: dict
+    ) -> Optional[int]:
+        """Get the number of accelerators of this family on ec2 instance with given type.
+
+        Args:
+            instance_type: The ec2 instance type.
+            instances: Map from ec2 instance type to instance metadata returned by
+                ec2 `describe-instance-types`.
+
+        Returns:
+            The number of accelerators of this family on the ec2 instance
+            with given type.
+            Return None if it's unknown.
+        """
+        return None
+
+    @staticmethod
+    def get_ec2_instance_accelerator_type(
+        instance_type: str, instances: dict
+    ) -> Optional[str]:
+        """Get the accelerator type of this family on ec2 instance with given type.
+
+        Args:
+            instance_type: The ec2 instance type.
+            instances: Map from ec2 instance type to instance metadata returned by
+                ec2 `describe-instance-types`.
+
+        Returns:
+            The accelerator type of this family on the ec2 instance with given type.
+            Return None if it's unknown.
+        """
+        return None
+
+    @staticmethod
+    def get_current_node_accelerator_labels() -> Optional[Dict[str, str]]:
+        """Get accelerator related Ray node labels of the curent node.
+
+        Returns:
+            A dictionary mapping accelerator related label keys to values.
+        """
+        return None

python/ray/_private/accelerators/amd_gpu.py

@@ -0,0 +1,153 @@
+import logging
+import os
+from typing import List, Optional, Tuple
+
+from ray._private.accelerators.accelerator import AcceleratorManager
+from ray._private.ray_constants import env_bool
+
+logger = logging.getLogger(__name__)
+
+HIP_VISIBLE_DEVICES_ENV_VAR = "HIP_VISIBLE_DEVICES"
+NOSET_HIP_VISIBLE_DEVICES_ENV_VAR = "RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES"
+
+amd_product_dict = {
+    "0x66a1": "AMD-Instinct-MI50",
+    "0x738c": "AMD-Instinct-MI100",
+    "0x7408": "AMD-Instinct-MI250X",
+    "0x740c": "AMD-Instinct-MI250X-MI250",
+    "0x740f": "AMD-Instinct-MI210",
+    "0x74a0": "AMD-Instinct-MI300A",
+    "0x74a1": "AMD-Instinct-MI300X-OAM",
+    "0x74a2": "AMD-Instinct-MI308X-OAM",
+    "0x74a9": "AMD-Instinct-MI300X-HF",
+    "0x74a5": "AMD-Instinct-MI325X-OAM",
+    "0x75a0": "AMD-Instinct-MI350X-OAM",
+    "0x75a3": "AMD-Instinct-MI355X-OAM",
+    "0x6798": "AMD-Radeon-R9-200-HD-7900",
+    "0x6799": "AMD-Radeon-HD-7900",
+    "0x679A": "AMD-Radeon-HD-7900",
+    "0x679B": "AMD-Radeon-HD-7900",
+}
+
+
+class AMDGPUAcceleratorManager(AcceleratorManager):
+    """AMD GPU accelerators."""
+
+    @staticmethod
+    def get_resource_name() -> str:
+        return "GPU"
+
+    @staticmethod
+    def get_visible_accelerator_ids_env_var() -> str:
+        if (
+            HIP_VISIBLE_DEVICES_ENV_VAR not in os.environ
+            and "ROCR_VISIBLE_DEVICES" in os.environ
+        ):
+            raise RuntimeError(
+                f"Please use {HIP_VISIBLE_DEVICES_ENV_VAR} instead of ROCR_VISIBLE_DEVICES"
+            )
+
+        return HIP_VISIBLE_DEVICES_ENV_VAR
+
+    @staticmethod
+    def get_current_process_visible_accelerator_ids() -> Optional[List[str]]:
+        amd_visible_devices = os.environ.get(
+            AMDGPUAcceleratorManager.get_visible_accelerator_ids_env_var(), None
+        )
+
+        if amd_visible_devices is None:
+            return None
+
+        if amd_visible_devices == "":
+            return []
+
+        if amd_visible_devices == "NoDevFiles":
+            return []
+
+        return list(amd_visible_devices.split(","))
+
+    @staticmethod
+    def get_current_node_num_accelerators() -> int:
+        import ray._private.thirdparty.pyamdsmi as pyamdsmi
+
+        num_gpus = 0
+
+        try:
+            pyamdsmi.smi_initialize()
+            num_gpus = pyamdsmi.smi_get_device_count()
+        except Exception:
+            pass
+        finally:
+            try:
+                pyamdsmi.smi_shutdown()
+            except Exception:
+                pass
+
+        return num_gpus
+
+    @staticmethod
+    def get_current_node_accelerator_type() -> Optional[str]:
+        try:
+            device_ids = AMDGPUAcceleratorManager._get_amd_device_ids()
+            if device_ids is None:
+                return None
+            return AMDGPUAcceleratorManager._gpu_name_to_accelerator_type(device_ids[0])
+        except Exception:
+            return None
+
+    @staticmethod
+    def _gpu_name_to_accelerator_type(name):
+        if name is None:
+            return None
+        try:
+            match = amd_product_dict[name]
+            return match
+        except Exception:
+            return None
+
+    @staticmethod
+    def validate_resource_request_quantity(
+        quantity: float,
+    ) -> Tuple[bool, Optional[str]]:
+        return (True, None)
+
+    @staticmethod
+    def set_current_process_visible_accelerator_ids(
+        visible_amd_devices: List[str],
+    ) -> None:
+        if env_bool(NOSET_HIP_VISIBLE_DEVICES_ENV_VAR, False):
+            return
+
+        os.environ[
+            AMDGPUAcceleratorManager.get_visible_accelerator_ids_env_var()
+        ] = ",".join([str(i) for i in visible_amd_devices])
+
+    @staticmethod
+    def _get_amd_device_ids() -> List[str]:
+        """Get the list of GPUs IDs
+        Example:
+            On a node with 2x MI210 GPUs
+            pyamdsmi library python bindings
+            return: ['0x740f', '0x740f']
+        Returns:
+            A list of strings containing GPU IDs
+        """
+        import ray._private.thirdparty.pyamdsmi as pyamdsmi
+
+        device_ids = []
+        try:
+            pyamdsmi.smi_initialize()
+            num_devices = pyamdsmi.smi_get_device_count()
+            for i in range(num_devices):
+                did = pyamdsmi.smi_get_device_id(i)
+                if did >= 0:
+                    device_ids.append(hex(did))
+        except Exception:
+            return None
+        finally:
+            try:
+                pyamdsmi.pyamdsmi_shutdown()
+            except Exception:
+                pass
+
+        return device_ids

python/ray/_private/accelerators/furiosa.py

@@ -0,0 +1,209 @@
+import logging
+import os
+from functools import lru_cache
+from typing import List, Optional, Tuple
+
+from ray._private.accelerators.accelerator import AcceleratorManager
+from ray._private.ray_constants import env_bool
+
+logger = logging.getLogger(__name__)
+
+# Ray uses ``FURIOSA_DEVICES`` to track which Furiosa NPUs are assigned to a
+# worker/actor process. The value uses ``npu:<index>`` notation. Ray's
+# scheduler operates at the device level, so the value Ray writes is always
+# the device-level form (e.g. ``npu:0,npu:3``). Bare integer IDs are also
+# accepted on read for convenience.
+#
+# Note that ``furiosa-llm``'s Python API does not honor ``FURIOSA_DEVICES``
+# automatically; callers must pass ``devices=os.environ["FURIOSA_DEVICES"]``
+# (or an equivalent list) explicitly to ``furiosa_llm.LLM(...)``. The
+# ``furiosa-llm`` CLI does read the value but accepts a richer
+# ``npu:X:Y`` (PE-level) form that Ray does not currently preserve through
+# worker scheduling; see ``_strip_npu_prefix`` below.
+FURIOSA_VISIBLE_DEVICES_ENV_VAR = "FURIOSA_DEVICES"
+NOSET_FURIOSA_VISIBLE_DEVICES_ENV_VAR = "RAY_EXPERIMENTAL_NOSET_FURIOSA_DEVICES"
+
+_FURIOSA_DEVICE_PREFIX = "npu:"
+
+
+@lru_cache(maxsize=None)
+def _ensure_furiosa_initialized() -> bool:
+    """Run ``furiosa_smi_py.init()`` exactly once per process."""
+    from furiosa_smi_py import init
+
+    init()
+    return True
+
+
+def _get_furiosa_list_devices():
+    """Lazy import + one-shot init of ``furiosa_smi_py``.
+
+    Returns the current ``list_devices`` callable on success, or ``None`` if
+    the SDK is unavailable or initialization fails. ``list_devices`` is
+    re-imported on each call so test monkeypatches on the module attribute
+    take effect.
+    """
+    try:
+        _ensure_furiosa_initialized()
+        from furiosa_smi_py import list_devices
+    except Exception as e:
+        logger.debug("furiosa_smi_py is unavailable: %s", e)
+        return None
+    return list_devices
+
+
+def _strip_npu_prefix(token: str) -> str:
+    """Return the numeric device index from an ``npu:<id>`` token.
+
+    Accepts bare integers (``"3"``) as well as the prefixed form
+    (``"npu:3"``) so that values written by other tooling round-trip
+    cleanly.
+    """
+    token = token.strip()
+    if token.startswith(_FURIOSA_DEVICE_PREFIX):
+        token = token[len(_FURIOSA_DEVICE_PREFIX) :]
+    # ``furiosa-llm`` accepts both ``npu:X`` (whole NPU) and ``npu:X:Y``
+    # (PE-level, e.g. ``npu:0:0-3`` for fused PE 0-3 of NPU 0). Ray's
+    # scheduler currently operates at the device level only, so we keep
+    # the device index and drop any trailing PE selector. Round-tripping
+    # PE-level partitioning through worker scheduling is tracked as a
+    # follow-up enhancement.
+    return token.split(":", 1)[0]
+
+
+class FuriosaAcceleratorManager(AcceleratorManager):
+    """FuriosaAI NPU accelerators.
+
+    Resource name is ``FURIOSA``. The accelerator type is reported as
+    ``FURIOSA_<ARCH>`` where ``<ARCH>`` is the architecture identifier
+    that the Furiosa SMI SDK exposes via its ``Arch`` enum. The current
+    SDK variants are ``Rngd``, ``RngdS``, ``RngdMax`` and ``RngdPlus``,
+    which surface here as ``FURIOSA_RNGD``, ``FURIOSA_RNGDS``,
+    ``FURIOSA_RNGDMAX`` and ``FURIOSA_RNGDPLUS`` respectively.
+    Supporting any architecture the SDK reports keeps this manager
+    forward-compatible with new SKUs as Furiosa adds them to
+    ``furiosa_smi_py``.
+
+    Device visibility is tracked through the ``FURIOSA_DEVICES``
+    environment variable, formatted as ``npu:<id>`` tokens. The value
+    can be passed to the ``furiosa-llm`` CLI (e.g.,
+    ``furiosa-llm serve --devices "$FURIOSA_DEVICES" ...``). When
+    invoking the ``furiosa_llm.LLM`` Python API directly, the assigned
+    devices must be passed explicitly, e.g.
+    ``LLM(model_path, devices=os.environ["FURIOSA_DEVICES"])``;
+    ``LLM(devices=None)`` allocates all visible NPUs and would bypass
+    Ray's per-worker isolation.
+    """
+
+    @staticmethod
+    def get_resource_name() -> str:
+        return "FURIOSA"
+
+    @staticmethod
+    def get_visible_accelerator_ids_env_var() -> str:
+        return FURIOSA_VISIBLE_DEVICES_ENV_VAR
+
+    @staticmethod
+    def get_current_process_visible_accelerator_ids() -> Optional[List[str]]:
+        visible_devices = os.environ.get(
+            FuriosaAcceleratorManager.get_visible_accelerator_ids_env_var()
+        )
+        if visible_devices is None:
+            return None
+        if visible_devices == "":
+            return []
+        return [
+            _strip_npu_prefix(token)
+            for token in visible_devices.split(",")
+            if token.strip()
+        ]
+
+    @staticmethod
+    def get_current_node_num_accelerators() -> int:
+        """Detects the number of Furiosa NPU devices on the current machine."""
+        list_devices = _get_furiosa_list_devices()
+        if list_devices is None:
+            return 0
+        try:
+            return len(list_devices())
+        except Exception as e:
+            logger.debug("Could not list Furiosa NPU devices: %s", e)
+            return 0
+
+    @staticmethod
+    def get_current_node_accelerator_type() -> Optional[str]:
+        """Gets the architecture of the Furiosa NPU on the current node.
+
+        Returns a string like ``FURIOSA_RNGD``, ``FURIOSA_RNGDMAX``,
+        ``FURIOSA_RNGDS``, or ``FURIOSA_RNGDPLUS``. Ray assumes a single
+        accelerator type per node, so the architecture of the first detected
+        device is used.
+
+        The architecture is read from
+        ``device.device_info().arch()`` to mirror the upstream Furiosa SMI
+        interface. The Arch enum string is normalized into an
+        accelerator-type label: ``+`` is mapped to ``plus`` so distinct
+        SKUs do not collide (``rngd+`` becomes ``FURIOSA_RNGDPLUS``,
+        matching the PyO3 enum form ``RngdPlus``), and any remaining
+        non-alphanumeric characters are stripped.
+        """
+        list_devices = _get_furiosa_list_devices()
+        if list_devices is None:
+            return None
+        try:
+            devices = list_devices()
+            if not devices:
+                return None
+
+            arch_obj = devices[0].device_info().arch()
+            if arch_obj is None:
+                return None
+
+            # PyO3 enums typically stringify as "<EnumName.Variant>",
+            # "EnumName.Variant", or just "Variant". Take the trailing
+            # component.
+            raw = str(arch_obj).split(".")[-1].strip()
+            if not raw:
+                return None
+            # Map special suffixes to their alphabetic equivalents so that the
+            # ``Arch::ToString`` form ("rngd+") and the PyO3 enum form
+            # ("RngdPlus") produce the same Ray accelerator type label.
+            # Without this, stripping "+" would collapse "rngd+" into
+            # "rngd", colliding with the distinct ``Rngd`` SKU.
+            raw = raw.replace("+", "plus")
+            normalized = "".join(ch for ch in raw if ch.isalnum()).upper()
+            if not normalized:
+                return None
+            return f"FURIOSA_{normalized}"
+        except Exception as e:
+            logger.debug("Failed to detect Furiosa NPU type: %s", e)
+            return None
+
+    @staticmethod
+    def validate_resource_request_quantity(
+        quantity: float,
+    ) -> Tuple[bool, Optional[str]]:
+        if isinstance(quantity, float) and not quantity.is_integer():
+            return (
+                False,
+                f"{FuriosaAcceleratorManager.get_resource_name()} resource quantity"
+                " must be a whole number. Furiosa NPUs do not support"
+                " fractional resource sharing."
+                f" The specified quantity {quantity} is invalid.",
+            )
+        return (True, None)
+
+    @staticmethod
+    def set_current_process_visible_accelerator_ids(
+        visible_furiosa_devices: List[str],
+    ) -> None:
+        if env_bool(NOSET_FURIOSA_VISIBLE_DEVICES_ENV_VAR, False):
+            return
+
+        formatted = ",".join(
+            f"{_FURIOSA_DEVICE_PREFIX}{_strip_npu_prefix(str(d))}"
+            for d in visible_furiosa_devices
+        )
+        os.environ[
+            FuriosaAcceleratorManager.get_visible_accelerator_ids_env_var()
+        ] = formatted

python/ray/_private/accelerators/hpu.py

@@ -0,0 +1,122 @@
+import logging
+import os
+from functools import lru_cache
+from importlib.util import find_spec
+from typing import List, Optional, Tuple
+
+from ray._private.accelerators.accelerator import AcceleratorManager
+from ray._private.ray_constants import env_bool
+
+logger = logging.getLogger(__name__)
+
+HABANA_VISIBLE_DEVICES_ENV_VAR = "HABANA_VISIBLE_MODULES"
+NOSET_HABANA_VISIBLE_MODULES_ENV_VAR = "RAY_EXPERIMENTAL_NOSET_HABANA_VISIBLE_MODULES"
+
+
+@lru_cache()
+def is_package_present(package_name: str) -> bool:
+    try:
+        return find_spec(package_name) is not None
+    except ModuleNotFoundError:
+        return False
+
+
+HPU_PACKAGE_AVAILABLE = is_package_present("habana_frameworks")
+
+
+class HPUAcceleratorManager(AcceleratorManager):
+    """Intel Habana(HPU) accelerators."""
+
+    @staticmethod
+    def get_resource_name() -> str:
+        return "HPU"
+
+    @staticmethod
+    def get_visible_accelerator_ids_env_var() -> str:
+        return HABANA_VISIBLE_DEVICES_ENV_VAR
+
+    @staticmethod
+    def get_current_process_visible_accelerator_ids() -> Optional[List[str]]:
+        hpu_visible_devices = os.environ.get(
+            HPUAcceleratorManager.get_visible_accelerator_ids_env_var(), None
+        )
+
+        if hpu_visible_devices is None:
+            return None
+
+        if hpu_visible_devices == "":
+            return []
+
+        return list(hpu_visible_devices.split(","))
+
+    @staticmethod
+    def get_current_node_num_accelerators() -> int:
+        """Attempt to detect the number of HPUs on this machine.
+        Returns:
+            The number of HPUs if any were detected, otherwise 0.
+        """
+        if HPU_PACKAGE_AVAILABLE:
+            import habana_frameworks.torch.hpu as torch_hpu
+
+            if torch_hpu.is_available():
+                return torch_hpu.device_count()
+            else:
+                logging.info("HPU devices not available")
+                return 0
+        else:
+            return 0
+
+    @staticmethod
+    def is_initialized() -> bool:
+        """Attempt to check if HPU backend is initialized.
+        Returns:
+            True if backend initialized else False.
+        """
+        if HPU_PACKAGE_AVAILABLE:
+            import habana_frameworks.torch.hpu as torch_hpu
+
+            if torch_hpu.is_available() and torch_hpu.is_initialized():
+                return True
+            else:
+                return False
+        else:
+            return False
+
+    @staticmethod
+    def get_current_node_accelerator_type() -> Optional[str]:
+        """Attempt to detect the HPU family type.
+        Returns:
+            The device name (GAUDI, GAUDI2) if detected else None.
+        """
+        if HPUAcceleratorManager.is_initialized():
+            import habana_frameworks.torch.hpu as torch_hpu
+
+            return f"Intel-{torch_hpu.get_device_name()}"
+        else:
+            logging.info("HPU type cannot be detected")
+            return None
+
+    @staticmethod
+    def validate_resource_request_quantity(
+        quantity: float,
+    ) -> Tuple[bool, Optional[str]]:
+        if isinstance(quantity, float) and not quantity.is_integer():
+            return (
+                False,
+                f"{HPUAcceleratorManager.get_resource_name()} resource quantity"
+                " must be whole numbers. "
+                f"The specified quantity {quantity} is invalid.",
+            )
+        else:
+            return (True, None)
+
+    @staticmethod
+    def set_current_process_visible_accelerator_ids(
+        visible_hpu_devices: List[str],
+    ) -> None:
+        if env_bool(NOSET_HABANA_VISIBLE_MODULES_ENV_VAR, False):
+            return
+
+        os.environ[
+            HPUAcceleratorManager.get_visible_accelerator_ids_env_var()
+        ] = ",".join([str(i) for i in visible_hpu_devices])

python/ray/_private/accelerators/intel_gpu.py

@@ -0,0 +1,104 @@
+import logging
+import os
+from typing import List, Optional, Tuple
+
+from ray._private.accelerators.accelerator import AcceleratorManager
+from ray._private.ray_constants import env_bool
+
+logger = logging.getLogger(__name__)
+
+ONEAPI_DEVICE_SELECTOR_ENV_VAR = "ONEAPI_DEVICE_SELECTOR"
+NOSET_ONEAPI_DEVICE_SELECTOR_ENV_VAR = "RAY_EXPERIMENTAL_NOSET_ONEAPI_DEVICE_SELECTOR"
+ONEAPI_DEVICE_BACKEND_TYPE = "level_zero"
+ONEAPI_DEVICE_TYPE = "gpu"
+
+
+class IntelGPUAcceleratorManager(AcceleratorManager):
+    """Intel GPU accelerators."""
+
+    @staticmethod
+    def get_resource_name() -> str:
+        return "GPU"
+
+    @staticmethod
+    def get_visible_accelerator_ids_env_var() -> str:
+        return ONEAPI_DEVICE_SELECTOR_ENV_VAR
+
+    @staticmethod
+    def get_current_process_visible_accelerator_ids() -> Optional[List[str]]:
+        oneapi_visible_devices = os.environ.get(
+            IntelGPUAcceleratorManager.get_visible_accelerator_ids_env_var(), None
+        )
+        if oneapi_visible_devices is None:
+            return None
+
+        if oneapi_visible_devices == "":
+            return []
+
+        if oneapi_visible_devices == "NoDevFiles":
+            return []
+
+        prefix = ONEAPI_DEVICE_BACKEND_TYPE + ":"
+
+        return list(oneapi_visible_devices.split(prefix)[1].split(","))
+
+    @staticmethod
+    def get_current_node_num_accelerators() -> int:
+        try:
+            import dpctl
+        except ImportError:
+            dpctl = None
+        if dpctl is None:
+            return 0
+
+        num_gpus = 0
+        try:
+            dev_info = ONEAPI_DEVICE_BACKEND_TYPE + ":" + ONEAPI_DEVICE_TYPE
+            context = dpctl.SyclContext(dev_info)
+            num_gpus = context.device_count
+        except Exception:
+            num_gpus = 0
+        return num_gpus
+
+    @staticmethod
+    def get_current_node_accelerator_type() -> Optional[str]:
+        """Get the name of first Intel GPU. (supposed only one GPU type on a node)
+        Example:
+            name: 'Intel(R) Data Center GPU Max 1550'
+            return name: 'Intel-GPU-Max-1550'
+        Returns:
+            A string representing the name of Intel GPU type.
+        """
+        try:
+            import dpctl
+        except ImportError:
+            dpctl = None
+        if dpctl is None:
+            return None
+
+        accelerator_type = None
+        try:
+            dev_info = ONEAPI_DEVICE_BACKEND_TYPE + ":" + ONEAPI_DEVICE_TYPE + ":0"
+            dev = dpctl.SyclDevice(dev_info)
+            accelerator_type = "Intel-GPU-" + "-".join(dev.name.split(" ")[-2:])
+        except Exception:
+            accelerator_type = None
+        return accelerator_type
+
+    @staticmethod
+    def validate_resource_request_quantity(
+        quantity: float,
+    ) -> Tuple[bool, Optional[str]]:
+        return (True, None)
+
+    @staticmethod
+    def set_current_process_visible_accelerator_ids(
+        visible_xpu_devices: List[str],
+    ) -> None:
+        if env_bool(NOSET_ONEAPI_DEVICE_SELECTOR_ENV_VAR, False):
+            return
+
+        prefix = ONEAPI_DEVICE_BACKEND_TYPE + ":"
+        os.environ[
+            IntelGPUAcceleratorManager.get_visible_accelerator_ids_env_var()
+        ] = prefix + ",".join([str(i) for i in visible_xpu_devices])

python/ray/_private/accelerators/metax_gpu.py

@@ -0,0 +1,90 @@
+import logging
+import os
+from typing import List, Optional, Tuple
+
+from ray._private.accelerators.accelerator import AcceleratorManager
+
+logger = logging.getLogger(__name__)
+
+CUDA_VISIBLE_DEVICES_ENV_VAR = "CUDA_VISIBLE_DEVICES"
+NOSET_CUDA_VISIBLE_DEVICES_ENV_VAR = "RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES"
+
+
+class MetaxGPUAcceleratorManager(AcceleratorManager):
+    """Metax GPU accelerators."""
+
+    @staticmethod
+    def get_resource_name() -> str:
+        return "GPU"
+
+    @staticmethod
+    def get_visible_accelerator_ids_env_var() -> str:
+        return CUDA_VISIBLE_DEVICES_ENV_VAR
+
+    @staticmethod
+    def get_current_process_visible_accelerator_ids() -> Optional[List[str]]:
+        cuda_visible_devices = os.environ.get(
+            MetaxGPUAcceleratorManager.get_visible_accelerator_ids_env_var(), None
+        )
+        if cuda_visible_devices is None:
+            return None
+
+        if cuda_visible_devices == "":
+            return []
+
+        return list(cuda_visible_devices.split(","))
+
+    @staticmethod
+    def get_current_node_num_accelerators() -> int:
+        try:
+            import pymxsml.mxsml_extension as pymxsml
+
+            try:
+                pymxsml.mxSmlExInit()
+            except pymxsml.MXSMLEXError:
+                return 0
+            device_count = pymxsml.mxSmlExDeviceGetCount()
+            pymxsml.mxSmlExShutdown()
+            return device_count
+        except Exception as e:
+            logger.debug("Could not import pymxsml: %s", e)
+        return 0
+
+    @staticmethod
+    def get_current_node_accelerator_type() -> Optional[str]:
+        try:
+            import pymxsml.mxsml_extension as pymxsml
+
+            try:
+                pymxsml.mxSmlExInit()
+            except pymxsml.MXSMLEXError:
+                return None
+            device_name = None
+            device_count = pymxsml.mxSmlExDeviceGetCount()
+            if device_count > 0:
+                handle = pymxsml.mxSmlExDeviceGetHandleByIndex(0)
+                device_name = pymxsml.mxSmlExDeviceGetName(handle)
+                if isinstance(device_name, bytes):
+                    device_name = device_name.decode("utf-8")
+            pymxsml.mxSmlExShutdown()
+            return device_name
+        except Exception:
+            logger.warning("Failed to detect GPU type.", exc_info=True)
+        return None
+
+    @staticmethod
+    def validate_resource_request_quantity(
+        quantity: float,
+    ) -> Tuple[bool, Optional[str]]:
+        return (True, None)
+
+    @staticmethod
+    def set_current_process_visible_accelerator_ids(
+        visible_cuda_devices: List[str],
+    ) -> None:
+        if os.environ.get(NOSET_CUDA_VISIBLE_DEVICES_ENV_VAR):
+            return
+
+        os.environ[
+            MetaxGPUAcceleratorManager.get_visible_accelerator_ids_env_var()
+        ] = ",".join(visible_cuda_devices)

python/ray/_private/accelerators/neuron.py

@@ -0,0 +1,133 @@
+import json
+import logging
+import os
+import subprocess
+import sys
+from typing import List, Optional, Tuple
+
+from ray._private.accelerators.accelerator import AcceleratorManager
+from ray._private.ray_constants import env_bool
+
+logger = logging.getLogger(__name__)
+
+NEURON_RT_VISIBLE_CORES_ENV_VAR = "NEURON_RT_VISIBLE_CORES"
+NOSET_AWS_NEURON_RT_VISIBLE_CORES_ENV_VAR = (
+    "RAY_EXPERIMENTAL_NOSET_NEURON_RT_VISIBLE_CORES"
+)
+
+# https://awsdocs-neuron.readthedocs-hosted.com/en/latest/general/arch/neuron-hardware/inf2-arch.html#aws-inf2-arch
+# https://awsdocs-neuron.readthedocs-hosted.com/en/latest/general/arch/neuron-hardware/trn1-arch.html#aws-trn1-arch
+# Subject to removal after the information is available via public API
+AWS_NEURON_INSTANCE_MAP = {
+    "trn1.2xlarge": 2,
+    "trn1.32xlarge": 32,
+    "trn1n.32xlarge": 32,
+    "inf2.xlarge": 2,
+    "inf2.8xlarge": 2,
+    "inf2.24xlarge": 12,
+    "inf2.48xlarge": 24,
+}
+
+
+class NeuronAcceleratorManager(AcceleratorManager):
+    """AWS Inferentia and Trainium accelerators."""
+
+    @staticmethod
+    def get_resource_name() -> str:
+        return "neuron_cores"
+
+    @staticmethod
+    def get_visible_accelerator_ids_env_var() -> str:
+        return NEURON_RT_VISIBLE_CORES_ENV_VAR
+
+    @staticmethod
+    def get_current_process_visible_accelerator_ids() -> Optional[List[str]]:
+        neuron_visible_cores = os.environ.get(
+            NeuronAcceleratorManager.get_visible_accelerator_ids_env_var(), None
+        )
+
+        if neuron_visible_cores is None:
+            return None
+
+        if neuron_visible_cores == "":
+            return []
+
+        return list(neuron_visible_cores.split(","))
+
+    @staticmethod
+    def get_current_node_num_accelerators() -> int:
+        """
+        Attempt to detect the number of Neuron cores on this machine.
+
+        Returns:
+            The number of Neuron cores if any were detected, otherwise 0.
+        """
+        nc_count: int = 0
+        neuron_path = "/opt/aws/neuron/bin/"
+        if sys.platform.startswith("linux") and os.path.isdir(neuron_path):
+            result = subprocess.run(
+                [os.path.join(neuron_path, "neuron-ls"), "--json-output"],
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+            )
+            if result.returncode == 0 and result.stdout:
+                neuron_devices = json.loads(result.stdout)
+                for neuron_device in neuron_devices:
+                    nc_count += neuron_device.get("nc_count", 0)
+        return nc_count
+
+    @staticmethod
+    def get_current_node_accelerator_type() -> Optional[str]:
+        from ray.util.accelerators import AWS_NEURON_CORE
+
+        return AWS_NEURON_CORE
+
+    @staticmethod
+    def validate_resource_request_quantity(
+        quantity: float,
+    ) -> Tuple[bool, Optional[str]]:
+        if isinstance(quantity, float) and not quantity.is_integer():
+            return (
+                False,
+                f"{NeuronAcceleratorManager.get_resource_name()} resource quantity"
+                " must be whole numbers. "
+                f"The specified quantity {quantity} is invalid.",
+            )
+        else:
+            return (True, None)
+
+    @staticmethod
+    def set_current_process_visible_accelerator_ids(
+        visible_neuron_core_ids: List[str],
+    ) -> None:
+        """Set the NEURON_RT_VISIBLE_CORES environment variable based on
+        given visible_neuron_core_ids.
+
+        Args:
+            visible_neuron_core_ids: List of str representing core IDs.
+        """
+        if env_bool(NOSET_AWS_NEURON_RT_VISIBLE_CORES_ENV_VAR, False):
+            return
+
+        os.environ[
+            NeuronAcceleratorManager.get_visible_accelerator_ids_env_var()
+        ] = ",".join([str(i) for i in visible_neuron_core_ids])
+
+    @staticmethod
+    def get_ec2_instance_num_accelerators(
+        instance_type: str, instances: dict
+    ) -> Optional[int]:
+        # TODO: AWS SDK (public API) doesn't yet expose the NeuronCore
+        # information. It will be available (work-in-progress)
+        # as xxAcceleratorInfo in InstanceTypeInfo.
+        # https://docs.aws.amazon.com/AWSEC2/latest/APIReference/API_InstanceTypeInfo.html
+        # See https://github.com/ray-project/ray/issues/38473
+        return AWS_NEURON_INSTANCE_MAP.get(instance_type.lower(), None)
+
+    @staticmethod
+    def get_ec2_instance_accelerator_type(
+        instance_type: str, instances: dict
+    ) -> Optional[str]:
+        from ray.util.accelerators import AWS_NEURON_CORE
+
+        return AWS_NEURON_CORE

python/ray/_private/accelerators/npu.py

@@ -0,0 +1,100 @@
+import glob
+import logging
+import os
+from typing import List, Optional, Tuple
+
+from ray._private.accelerators.accelerator import AcceleratorManager
+from ray._private.ray_constants import env_bool
+
+logger = logging.getLogger(__name__)
+
+ASCEND_RT_VISIBLE_DEVICES_ENV_VAR = "ASCEND_RT_VISIBLE_DEVICES"
+NOSET_ASCEND_RT_VISIBLE_DEVICES_ENV_VAR = (
+    "RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES"
+)
+
+
+class NPUAcceleratorManager(AcceleratorManager):
+    """Ascend NPU accelerators."""
+
+    @staticmethod
+    def get_resource_name() -> str:
+        return "NPU"
+
+    @staticmethod
+    def get_visible_accelerator_ids_env_var() -> str:
+        return ASCEND_RT_VISIBLE_DEVICES_ENV_VAR
+
+    @staticmethod
+    def get_current_process_visible_accelerator_ids() -> Optional[List[str]]:
+        ascend_visible_devices = os.environ.get(
+            NPUAcceleratorManager.get_visible_accelerator_ids_env_var(), None
+        )
+
+        if ascend_visible_devices is None:
+            return None
+
+        if ascend_visible_devices == "":
+            return []
+
+        if ascend_visible_devices == "NoDevFiles":
+            return []
+
+        return list(ascend_visible_devices.split(","))
+
+    @staticmethod
+    def get_current_node_num_accelerators() -> int:
+        """Attempt to detect the number of NPUs on this machine.
+
+        NPU chips are represented as devices within `/dev/`, either as `/dev/davinci?`.
+
+        Returns:
+            The number of NPUs if any were detected, otherwise 0.
+        """
+        try:
+            import acl
+
+            device_count, ret = acl.rt.get_device_count()
+            if ret == 0:
+                return device_count
+        except Exception as e:
+            logger.debug("Could not import AscendCL: %s", e)
+
+        try:
+            npu_files = glob.glob("/dev/davinci[0-9]*")
+            return len(npu_files)
+        except Exception as e:
+            logger.debug("Failed to detect number of NPUs: %s", e)
+        return 0
+
+    @staticmethod
+    def get_current_node_accelerator_type() -> Optional[str]:
+        """Get the type of the Ascend NPU on the current node.
+
+        Returns:
+            A string of the type, such as "Ascend910A", "Ascend910B", "Ascend310P1".
+        """
+        try:
+            import acl
+
+            return acl.get_soc_name()
+        except Exception:
+            logger.exception("Failed to detect NPU type.")
+        return None
+
+    @staticmethod
+    def validate_resource_request_quantity(
+        quantity: float,
+    ) -> Tuple[bool, Optional[str]]:
+        return (True, None)
+
+    @staticmethod
+    def set_current_process_visible_accelerator_ids(
+        visible_npu_devices: List[str],
+    ) -> None:
+        if env_bool(NOSET_ASCEND_RT_VISIBLE_DEVICES_ENV_VAR, False):
+            return
+
+        os.environ[
+            NPUAcceleratorManager.get_visible_accelerator_ids_env_var()
+        ] = ",".join([str(i) for i in visible_npu_devices])

python/ray/_private/accelerators/nvidia_gpu.py

@@ -0,0 +1,145 @@
+import logging
+import os
+import re
+from typing import List, Optional, Tuple
+
+from ray._private.accelerators.accelerator import AcceleratorManager
+from ray._private.ray_constants import env_bool
+
+logger = logging.getLogger(__name__)
+
+CUDA_VISIBLE_DEVICES_ENV_VAR = "CUDA_VISIBLE_DEVICES"
+NOSET_CUDA_VISIBLE_DEVICES_ENV_VAR = "RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES"
+
+# Capture the accelerator model from the NVML device name: the run of leading
+# all-caps tokens (e.g. "RTX", "PRO") up to and including the first token that
+# contains a digit. This keeps datacenter cards stable ("Tesla V100-SXM2-16GB"
+# -> "V100", "NVIDIA A100-SXM4-40GB" -> "A100") while disambiguating the RTX
+# line, whose first token is only a brand prefix ("NVIDIA RTX PRO 6000 Blackwell
+# Server Edition" -> "RTX PRO 6000"). A trailing SKU suffix after a hyphen is
+# dropped. Mixed-case consumer names ("NVIDIA GeForce RTX 5090") don't match and
+# fall back to a hyphen-joined product name in _gpu_name_to_accelerator_type.
+NVIDIA_GPU_NAME_PATTERN = re.compile(r"\w+\s+((?:[A-Z]+\s+)*[A-Z0-9]*\d[A-Z0-9]*)")
+
+
+class NvidiaGPUAcceleratorManager(AcceleratorManager):
+    """NVIDIA GPU accelerators."""
+
+    @staticmethod
+    def get_resource_name() -> str:
+        return "GPU"
+
+    @staticmethod
+    def get_visible_accelerator_ids_env_var() -> str:
+        return CUDA_VISIBLE_DEVICES_ENV_VAR
+
+    @staticmethod
+    def get_current_process_visible_accelerator_ids() -> Optional[List[str]]:
+        cuda_visible_devices = os.environ.get(
+            NvidiaGPUAcceleratorManager.get_visible_accelerator_ids_env_var(), None
+        )
+        if cuda_visible_devices is None:
+            return None
+
+        if cuda_visible_devices == "":
+            return []
+
+        if cuda_visible_devices == "NoDevFiles":
+            return []
+
+        return list(cuda_visible_devices.split(","))
+
+    @staticmethod
+    def get_current_node_num_accelerators() -> int:
+        import ray._private.thirdparty.pynvml as pynvml
+
+        try:
+            pynvml.nvmlInit()
+        except pynvml.NVMLError:
+            return 0  # pynvml init failed
+        device_count = pynvml.nvmlDeviceGetCount()
+        pynvml.nvmlShutdown()
+        return device_count
+
+    @staticmethod
+    def get_current_node_accelerator_type() -> Optional[str]:
+        import ray._private.thirdparty.pynvml as pynvml
+
+        try:
+            pynvml.nvmlInit()
+        except pynvml.NVMLError:
+            return None  # pynvml init failed
+        device_count = pynvml.nvmlDeviceGetCount()
+        cuda_device_type = None
+        if device_count > 0:
+            handle = pynvml.nvmlDeviceGetHandleByIndex(0)
+            device_name = pynvml.nvmlDeviceGetName(handle)
+            if isinstance(device_name, bytes):
+                device_name = device_name.decode("utf-8")
+            cuda_device_type = (
+                NvidiaGPUAcceleratorManager._gpu_name_to_accelerator_type(device_name)
+            )
+        pynvml.nvmlShutdown()
+        return cuda_device_type
+
+    @staticmethod
+    def _gpu_name_to_accelerator_type(name):
+        if name is None:
+            return None
+        match = NVIDIA_GPU_NAME_PATTERN.match(name)
+        result = match.group(1).replace(" ", "-") if match else None
+        if result and len(result) > 1:
+            return result
+        # The pattern above requires an all-uppercase/numeric model token, which
+        # works for datacenter cards ("Tesla V100-SXM2-16GB" -> "V100",
+        # "NVIDIA RTX PRO 6000 ..." -> "RTX-PRO-6000") but not for consumer
+        # cards whose product line is mixed case ("NVIDIA GeForce RTX 5090").
+        # Fall back to a hyphen-joined product name so callers get a useful
+        # accelerator_type label like "GeForce-RTX-5090".
+        cleaned = re.sub(r"^NVIDIA\s+", "", name).strip()
+        return cleaned.replace(" ", "-") if cleaned else None
+
+    @staticmethod
+    def validate_resource_request_quantity(
+        quantity: float,
+    ) -> Tuple[bool, Optional[str]]:
+        return (True, None)
+
+    @staticmethod
+    def set_current_process_visible_accelerator_ids(
+        visible_cuda_devices: List[str],
+    ) -> None:
+        if env_bool(NOSET_CUDA_VISIBLE_DEVICES_ENV_VAR, False):
+            return
+
+        os.environ[
+            NvidiaGPUAcceleratorManager.get_visible_accelerator_ids_env_var()
+        ] = ",".join([str(i) for i in visible_cuda_devices])
+
+    @staticmethod
+    def get_ec2_instance_num_accelerators(
+        instance_type: str, instances: dict
+    ) -> Optional[int]:
+        if instance_type not in instances:
+            return None
+
+        gpus = instances[instance_type].get("GpuInfo", {}).get("Gpus")
+        if gpus is not None:
+            # TODO(ameer): currently we support one gpu type per node.
+            assert len(gpus) == 1
+            return gpus[0]["Count"]
+        return None
+
+    @staticmethod
+    def get_ec2_instance_accelerator_type(
+        instance_type: str, instances: dict
+    ) -> Optional[str]:
+        if instance_type not in instances:
+            return None
+
+        gpus = instances[instance_type].get("GpuInfo", {}).get("Gpus")
+        if gpus is not None:
+            # TODO(ameer): currently we support one gpu type per node.
+            assert len(gpus) == 1
+            return gpus[0]["Name"]
+        return None

python/ray/_private/accelerators/rbln.py

@@ -0,0 +1,81 @@
+import logging
+import os
+from typing import List, Optional, Tuple
+
+from ray._private.accelerators.accelerator import AcceleratorManager
+from ray._private.ray_constants import env_bool
+
+logger = logging.getLogger(__name__)
+
+RBLN_RT_VISIBLE_DEVICES_ENV_VAR = "RBLN_DEVICES"
+NOSET_RBLN_RT_VISIBLE_DEVICES_ENV_VAR = "RAY_EXPERIMENTAL_NOSET_RBLN_RT_VISIBLE_DEVICES"
+
+
+class RBLNAcceleratorManager(AcceleratorManager):
+    """Rebellions RBLN accelerators."""
+
+    @staticmethod
+    def get_resource_name() -> str:
+        return "RBLN"
+
+    @staticmethod
+    def get_visible_accelerator_ids_env_var() -> str:
+        return RBLN_RT_VISIBLE_DEVICES_ENV_VAR
+
+    @staticmethod
+    def get_current_process_visible_accelerator_ids() -> Optional[List[str]]:
+        visible_devices = os.environ.get(
+            RBLNAcceleratorManager.get_visible_accelerator_ids_env_var()
+        )
+        if visible_devices is None:
+            return None
+        if visible_devices == "":
+            return []
+        return visible_devices.split(",")
+
+    @staticmethod
+    def get_current_node_num_accelerators() -> int:
+        """Detects the number of RBLN devices on the current machine."""
+        try:
+            from rebel import device_count
+
+            return device_count()
+        except Exception as e:
+            logger.debug("Could not detect RBLN devices: %s", e)
+            return 0
+
+    @staticmethod
+    def get_current_node_accelerator_type() -> Optional[str]:
+        """Gets the type of RBLN NPU on the current node."""
+        try:
+            from rebel import get_npu_name
+
+            return get_npu_name()
+        except Exception as e:
+            logger.exception("Failed to detect RBLN NPU type: %s", e)
+            return None
+
+    @staticmethod
+    def validate_resource_request_quantity(
+        quantity: float,
+    ) -> Tuple[bool, Optional[str]]:
+        if isinstance(quantity, float) and not quantity.is_integer():
+            return (
+                False,
+                f"{RBLNAcceleratorManager.get_resource_name()} resource quantity"
+                " must be whole numbers. "
+                f"The specified quantity {quantity} is invalid.",
+            )
+        else:
+            return (True, None)
+
+    @staticmethod
+    def set_current_process_visible_accelerator_ids(
+        visible_rbln_devices: List[str],
+    ) -> None:
+        if env_bool(NOSET_RBLN_RT_VISIBLE_DEVICES_ENV_VAR, False):
+            return
+
+        os.environ[
+            RBLNAcceleratorManager.get_visible_accelerator_ids_env_var()
+        ] = ",".join(map(str, visible_rbln_devices))

python/ray/_private/accelerators/tpu.py

@@ -0,0 +1,766 @@
+import glob
+import logging
+import os
+import re
+from functools import lru_cache
+from typing import Dict, List, Optional, Set, Tuple
+
+import requests
+
+import ray
+from ray._private.accelerators.accelerator import AcceleratorManager
+from ray._private.ray_constants import env_bool
+from ray.util.placement_group import (
+    PlacementGroup,
+    placement_group,
+    remove_placement_group,
+)
+from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
+
+logger = logging.getLogger(__name__)
+
+
+TPU_VALID_CHIP_OPTIONS = (1, 2, 4, 8)
+GKE_TPU_ACCELERATOR_TYPE_ENV_VAR = "TPU_ACCELERATOR_TYPE"
+GKE_TPU_TOPOLOGY_ENV_VAR = "TPU_TOPOLOGY"
+GKE_TPU_WORKER_ID_ENV_VAR = "TPU_WORKER_ID"
+GKE_TPU_NAME_ENV_VAR = "TPU_NAME"
+
+# Constants for accessing the `accelerator-type` from TPU VM
+# instance metadata.
+# See https://cloud.google.com/compute/docs/metadata/overview
+# for more details about VM instance metadata.
+GCE_TPU_ACCELERATOR_ENDPOINT = (
+    "http://metadata.google.internal/computeMetadata/v1/instance/attributes/"
+)
+GCE_TPU_HEADERS = {"Metadata-Flavor": "Google"}
+GCE_TPU_ACCELERATOR_KEY = "accelerator-type"
+GCE_TPU_ENV_KEY = "tpu-env"
+GCE_TPU_INSTANCE_ID_KEY = "instance-id"
+GCE_TPU_WORKER_ID_KEY = "agent-worker-number"
+
+TPU_VISIBLE_CHIPS_ENV_VAR = "TPU_VISIBLE_CHIPS"
+
+NOSET_TPU_VISIBLE_CHIPS_ENV_VAR = "RAY_EXPERIMENTAL_NOSET_TPU_VISIBLE_CHIPS"
+
+# The following defines environment variables that allow
+# us to access a subset of TPU visible chips.
+#
+# See: https://github.com/google/jax/issues/14977 for an example/more details.
+TPU_CHIPS_PER_HOST_BOUNDS_ENV_VAR = "TPU_CHIPS_PER_HOST_BOUNDS"
+TPU_CHIPS_PER_HOST_BOUNDS_1_CHIP_CONFIG = "1,1,1"
+TPU_CHIPS_PER_HOST_BOUNDS_2_CHIP_CONFIG = "1,2,1"
+
+TPU_HOST_BOUNDS_ENV_VAR = "TPU_HOST_BOUNDS"
+TPU_SINGLE_HOST_BOUNDS = "1,1,1"
+
+# By default TPU VMs come with 4 chips per host and 2 tensorcores per chip.
+# For more details: https://cloud.google.com/tpu/docs/system-architecture-tpu-vm
+DEFAULT_TPU_NUM_CHIPS_PER_HOST = 4
+DEFAULT_TPU_NUM_CORES_PER_CHIP = 2
+
+# Accelerators that support up to 8 chips per host for single-host topologies: v5e, v6e
+TPU_8_CHIPS_PER_HOST_TYPES = ("v5litepod", "v6e")
+
+# Topologies that are always sub-host or single-host
+TPU_SINGLE_HOST_TOPOLOGIES = ("1x1", "2x2", "2x4")
+
+# Accelerators that are 2 cores per chip: v2, v3, v4, v5p, v7x
+# Accelerators that are 1 core per chip: v5e, v6e
+SINGLE_CORE_TPU_TYPES = ("v5litepod", "v6e")
+
+# The valid TPU types.
+VALID_TPU_TYPES = ("v2", "v3", "v4", "v5p", "v5litepod", "v6e", "v7x")
+
+# This is only used to construct TPU 3D topologies
+def _get_larger_3d_topologies(max_x: int, max_y: int, max_z: int) -> Set[str]:
+    """Returns a set of larger 3D TPU topologies given the max x,y,z value. Using DEFAULT_TPU_NUM_CHIPS_PER_HOST as increment"""
+    topologies = set()
+    for x in range(
+        DEFAULT_TPU_NUM_CHIPS_PER_HOST, max_x + 1, DEFAULT_TPU_NUM_CHIPS_PER_HOST
+    ):
+        for y in range(
+            DEFAULT_TPU_NUM_CHIPS_PER_HOST, max_y + 1, DEFAULT_TPU_NUM_CHIPS_PER_HOST
+        ):
+            for z in range(
+                DEFAULT_TPU_NUM_CHIPS_PER_HOST,
+                max_z + 1,
+                DEFAULT_TPU_NUM_CHIPS_PER_HOST,
+            ):
+                topologies.add(f"{x}x{y}x{z}")
+
+    return topologies
+
+
+# The valid TPU topologies for each of the TPU types.
+VALID_TPU_TOPOLOGY = {
+    "v2": {"4x4", "4x8", "8x8", "8x16", "16x16"},
+    "v3": {"4x4", "4x8", "8x8", "8x16", "16x16", "16x32", "32x32"},
+    "v4": {"2x2x1", "2x2x2", "2x2x4", "2x4x4"}.union(
+        _get_larger_3d_topologies(12, 12, 16)
+    ),
+    "v5p": {
+        "2x2x1",
+        "2x2x2",
+        "2x2x4",
+        "2x4x4",
+    }.union(_get_larger_3d_topologies(16, 16, 24)),
+    "v5litepod": {"1x1", "2x2", "2x4", "2x8", "4x4", "4x8", "8x8", "8x16", "16x16"},
+    "v6e": {"1x1", "2x2", "2x4", "2x8", "4x4", "4x8", "8x8", "8x16", "16x16"},
+    "v7x": {
+        "2x2x1",
+        "2x2x2",
+        "2x2x4",
+        "2x4x4",
+        "4x4x4",
+        "4x4x8",
+        "4x8x8",
+        "8x8x8",
+        "8x8x16",
+        "8x16x16",
+    },
+}
+
+
+def _get_tpu_metadata(key: str) -> Optional[str]:
+    """Poll and get TPU metadata."""
+    try:
+        accelerator_type_request = requests.get(
+            os.path.join(GCE_TPU_ACCELERATOR_ENDPOINT, key),
+            headers=GCE_TPU_HEADERS,
+        )
+        if (
+            accelerator_type_request.status_code == 200
+            and accelerator_type_request.text
+        ):
+            return accelerator_type_request.text
+        else:
+            logging.debug(
+                "Unable to poll TPU GCE Metadata. Got "
+                f"status code: {accelerator_type_request.status_code} and "
+                f"content: {accelerator_type_request.text}"
+            )
+    except requests.RequestException as e:
+        logging.debug("Unable to poll the TPU GCE Metadata: %s", e)
+    return None
+
+
+def _accelerator_type_check(accelerator_type: str):
+    if not accelerator_type.startswith(VALID_TPU_TYPES):
+        raise ValueError(
+            f"Invalid accelerator type: {accelerator_type}. Must start with one of: {VALID_TPU_TYPES}"
+        )
+
+
+def get_total_chips_from_accelerator_type(accelerator_type: str) -> int:
+    """Calculates total chips from a GCP accelerator ("pod") type string (e.g. "v6e-16")."""
+    _accelerator_type_check(accelerator_type)
+
+    parts = accelerator_type.split("-")
+    if len(parts) < 2:
+        raise ValueError(
+            f"Accelerator type must include size (e.g. 'v6e-8'), got: {accelerator_type}"
+        )
+
+    num_cores = int(parts[1])
+    cores_per_chip = get_tpu_cores_per_chip(accelerator_type)
+
+    return num_cores // cores_per_chip
+
+
+def get_num_tpu_visible_chips_per_host(accelerator_type: str) -> int:
+    _accelerator_type_check(accelerator_type)
+
+    if accelerator_type.startswith(TPU_8_CHIPS_PER_HOST_TYPES):
+        total_chips = get_total_chips_from_accelerator_type(accelerator_type)
+
+        # Sub/single-host topologies return their exact chip count
+        if total_chips <= 8:
+            return total_chips
+
+    # Multi-host topologies default to 4 visible chips per host
+    return DEFAULT_TPU_NUM_CHIPS_PER_HOST
+
+
+def get_tpu_cores_per_chip(accelerator_type: str) -> int:
+    _accelerator_type_check(accelerator_type)
+    if accelerator_type.startswith(SINGLE_CORE_TPU_TYPES):
+        return 1
+
+    return DEFAULT_TPU_NUM_CORES_PER_CHIP
+
+
+def get_num_chips_from_topology(topology: str) -> int:
+    """
+    Calculates the total number of chips in a TPU topology.
+    Ex: "2x2x2" -> 8
+    """
+    total_chips = 1
+    for dim in topology.strip().lower().split("x"):
+        total_chips *= int(dim)
+    return total_chips
+
+
+def infer_tpu_pod_type_from_topology(
+    topology: str, accelerator_type: str
+) -> Optional[str]:
+    """Infer the TPU pod type (e.g. v4-32) from topology and accelerator type."""
+    if not topology or not accelerator_type:
+        return None
+    try:
+        num_chips = get_num_chips_from_topology(topology)
+        generation = accelerator_type.lower().replace("tpu-", "")
+        num_cores = num_chips * get_tpu_cores_per_chip(generation)
+
+        return f"{generation}-{num_cores}"
+    except Exception as e:
+        raise ValueError(
+            f"Failed to infer pod type from topology '{topology}' "
+            f"and type '{accelerator_type}'"
+        ) from e
+
+
+def fetch_tpu_slice_name_from_pg(pg):
+    @ray.remote(num_cpus=0)
+    def _get_tpu_slice_name():
+        return TPUAcceleratorManager.get_current_node_tpu_name()
+
+    tpu_name_ref = _get_tpu_slice_name.options(
+        scheduling_strategy=PlacementGroupSchedulingStrategy(
+            placement_group=pg, placement_group_bundle_index=0
+        )
+    ).remote()
+
+    return ray.get(tpu_name_ref)
+
+
+def get_chips_per_host(topology: str, accelerator_version: str) -> int:
+    """Get the number of chips per host based on topology and accelerator version.
+
+    Rules for determining the default number of chips per host:
+        - Default for most TPU generations (v4, v5p, v7x, etc.) is 4 chips per host.
+        - For v5e and v6e:
+            - Topologies with <= 8 chips use the exact chip count (e.g. 1x1 -> 1).
+              These topologies are always sub or single-host.
+            - Multi-host topologies (> 8 chips) default to 4-chip hosts.
+
+    Args:
+        topology: The TPU topology string (e.g. "2x2x2", "2x4").
+        accelerator_version: The accelerator version string (e.g. "v4", "v6e").
+
+    Returns:
+        The default number of chips per host for the given configuration.
+    """
+    total_chips = get_num_chips_from_topology(topology)
+
+    # Check for 8-chip host types (v5litepod, v6e) for single host setups
+    if (
+        accelerator_version.strip().lower() in TPU_8_CHIPS_PER_HOST_TYPES
+        and topology.strip().lower() in TPU_SINGLE_HOST_TOPOLOGIES
+    ):
+        return total_chips
+
+    return DEFAULT_TPU_NUM_CHIPS_PER_HOST
+
+
+DEFAULT_TPU_HEAD_RESERVATION_TIMEOUT_S: float = 100.0
+
+
+def reserve_tpu_slice(
+    topology: str,
+    accelerator_type: str,
+    timeout_s: Optional[float] = DEFAULT_TPU_HEAD_RESERVATION_TIMEOUT_S,
+) -> Optional[Tuple[str, PlacementGroup]]:
+    """Reserves a TPU slice using its head resource and returns the slice name.
+    This enables gang scheduling of training workers with multi-host TPUs.
+    This is used by JaxTrainer with TPUs in Ray Train.
+
+    Args:
+        topology: The TPU topology string (e.g. "2x2x2").
+        accelerator_type: The accelerator type of the node (e.g. "TPU-V4").
+        timeout_s: The maximum time in seconds to wait for the TPU head
+            placement group to become ready. The head reservation must succeed
+            before the slice name can be retrieved, so this call is necessarily
+            blocking. Defaults to ``DEFAULT_TPU_HEAD_RESERVATION_TIMEOUT_S``.
+            Pass ``None`` to wait indefinitely.
+
+    Returns:
+        A tuple of a string representing a unique TPU slice name and the placement
+        group handle reserving the TPU head.
+
+    Raises:
+        TimeoutError: If the TPU head placement group does not become ready
+            within ``timeout_s`` seconds.
+    """
+    pod_type = infer_tpu_pod_type_from_topology(topology, accelerator_type)
+    if pod_type is None:
+        return None
+
+    # Reserve a slice by creating a placement group on the TPU head.
+    head_label_selector = {
+        "ray.io/tpu-worker-id": "0",
+        "ray.io/tpu-pod-type": pod_type,
+    }
+    head_placement_group = placement_group(
+        bundles=[{f"TPU-{pod_type}-head": 1}],
+        bundle_label_selector=[head_label_selector],
+    )
+
+    logger.debug(
+        "Waiting up to %s seconds to reserve multi-host slice head.", timeout_s
+    )
+    ready, _ = ray.wait([head_placement_group.ready()], timeout=timeout_s)
+
+    if not ready:
+        # Clean up the pending head reservation so that resources are not
+        # held while the caller decides whether to retry.
+        try:
+            remove_placement_group(head_placement_group)
+        except Exception:
+            logger.exception(
+                "Failed to clean up pending TPU head placement group after timeout."
+            )
+        raise TimeoutError(
+            "Failed to reserve TPU head for slice with shape: {} after {} "
+            "seconds. Ensure your cluster has sufficient resources. Requesting "
+            "TPU head node with labels: {}. Current resources: {}".format(
+                pod_type,
+                timeout_s,
+                head_label_selector,
+                ray.available_resources(),
+            )
+        )
+
+    # Retrieve the unique slice ID.
+    slice_name = fetch_tpu_slice_name_from_pg(head_placement_group)
+    if slice_name is None:
+        raise RuntimeError(
+            "Failed to retrieve TPU slice name after reserving head placement group. "
+            "Ensure that TPU slice metadata is available and correctly configured on multi-host nodes."
+        )
+
+    return (slice_name, head_placement_group)
+
+
+class TPUAcceleratorManager(AcceleratorManager):
+    """Google TPU accelerators."""
+
+    @staticmethod
+    def get_resource_name() -> str:
+        return "TPU"
+
+    @staticmethod
+    def get_visible_accelerator_ids_env_var() -> str:
+        return TPU_VISIBLE_CHIPS_ENV_VAR
+
+    @staticmethod
+    def get_current_process_visible_accelerator_ids() -> Optional[List[str]]:
+        tpu_visible_chips = os.environ.get(
+            TPUAcceleratorManager.get_visible_accelerator_ids_env_var(), None
+        )
+
+        if tpu_visible_chips is None:
+            return None
+
+        if tpu_visible_chips == "":
+            return []
+
+        return list(tpu_visible_chips.split(","))
+
+    @staticmethod
+    @lru_cache()
+    def get_current_node_num_accelerators() -> int:
+        """Attempt to detect the number of TPUs on this machine.
+
+        TPU chips are represented as devices within `/dev/`, either as
+        `/dev/accel*` or `/dev/vfio/*`.
+
+        Returns:
+            The number of TPUs if any were detected, otherwise 0.
+        """
+        # Real TPU chips are exposed as character devices at /dev/accel0,
+        # /dev/accel1, etc. NVIDIA drivers 570.x and later (Blackwell-class
+        # GPUs such as the RTX 5090) instead create /dev/accel as a *directory*
+        # containing /dev/accel/accel0, which the non-recursive glob below
+        # would otherwise miscount as a TPU chip. Filter directory entries out
+        # so both GKE and GCE TPU detection keep working while rejecting the
+        # NVIDIA false positive.
+        accel_chips = [p for p in glob.glob("/dev/accel*") if not os.path.isdir(p)]
+        if accel_chips:
+            return len(accel_chips)
+
+        try:
+            vfio_entries = os.listdir("/dev/vfio")
+            numeric_entries = [int(entry) for entry in vfio_entries if entry.isdigit()]
+            return len(numeric_entries)
+        except FileNotFoundError as e:
+            logger.debug("Failed to detect number of TPUs: %s", e)
+            return 0
+
+    @staticmethod
+    def is_valid_tpu_accelerator_type(tpu_accelerator_type: str) -> bool:
+        """Check whether the tpu accelerator_type is formatted correctly.
+
+        The accelerator_type field typically follows a form of v{generation}-{cores/chips},
+        but newer generations like 7x may follow tpu{generation}-{cores/chips}.
+
+        See the following for more information:
+        https://cloud.google.com/sdk/gcloud/reference/compute/tpus/tpu-vm/accelerator-types/describe
+
+        Args:
+            tpu_accelerator_type: The string representation of the accelerator type
+                to be checked for validity.
+
+        Returns:
+            True if it's valid, false otherwise.
+        """
+        # 1. Legacy format: v2-8, v3-32.
+        # 2. Newer format with letters in generation: v5litepod-16, v6e-4.
+        # 3. Ironwood TPU format which contains a tpu prefix: tpu7x-16.
+        expected_pattern = re.compile(r"^(v|tpu)\d+[a-zA-Z]*-\d+$")
+        if not expected_pattern.match(tpu_accelerator_type):
+            return False
+        return True
+
+    @staticmethod
+    def is_valid_tpu_accelerator_topology(
+        tpu_accelerator_version: str, tpu_topology: str
+    ) -> bool:
+        """Check whether the tpu topology is valid.
+
+        The accelerator_type field follows a form of v{generation}.
+        The accelerator_topology field follows either the form {A}x{B} or {A}x{B}x{C} depending on the v{generation}
+
+        Args:
+            tpu_accelerator_version: The string representation of the accelerator version. (e.g. v6e, V5P)
+            tpu_topology: The string representation of the accelerator topology
+                to be checked for validity
+
+        Returns:
+            True if it's a valid topology, False otherwise.
+        """
+        tpu_version_formatted = tpu_accelerator_version.strip().lower().split("-")[0]
+        if tpu_version_formatted.startswith("tpu"):
+            tpu_version_formatted = "v" + tpu_version_formatted[3:]
+        if (
+            tpu_version_formatted.lower() not in VALID_TPU_TOPOLOGY
+            or tpu_topology.strip().lower()
+            not in VALID_TPU_TOPOLOGY[tpu_version_formatted]
+        ):
+            return False
+        return True
+
+    @staticmethod
+    def validate_resource_request_quantity(
+        quantity: float,
+    ) -> Tuple[bool, Optional[str]]:
+        if quantity not in TPU_VALID_CHIP_OPTIONS:
+            return (
+                False,
+                f"The number of requested 'TPU' was set to {quantity} which "
+                "is not a supported chip configuration. Supported configs: "
+                f"{TPU_VALID_CHIP_OPTIONS}",
+            )
+        else:
+            return (True, None)
+
+    @staticmethod
+    def set_current_process_visible_accelerator_ids(
+        visible_tpu_chips: List[str],
+    ) -> None:
+        """Set TPU environment variables based on the provided visible_tpu_chips.
+
+        To access a subset of the TPU visible chips, we must use a combination of
+        environment variables that tells the compiler (via ML framework) the:
+        - Visible chips
+        - The physical bounds of chips per host
+        - The host bounds within the context of a TPU pod.
+
+        See: https://github.com/google/jax/issues/14977 for an example/more details.
+
+        Args:
+            visible_tpu_chips: List of str representing TPU chips.
+        """
+        if env_bool(NOSET_TPU_VISIBLE_CHIPS_ENV_VAR, False):
+            return
+
+        num_visible_tpu_chips = len(visible_tpu_chips)
+        num_accelerators_on_node = (
+            TPUAcceleratorManager.get_current_node_num_accelerators()
+        )
+        if num_visible_tpu_chips == num_accelerators_on_node:
+            # Let the ML framework use the defaults
+            os.environ.pop(TPU_CHIPS_PER_HOST_BOUNDS_ENV_VAR, None)
+            os.environ.pop(TPU_HOST_BOUNDS_ENV_VAR, None)
+            return
+        os.environ[
+            TPUAcceleratorManager.get_visible_accelerator_ids_env_var()
+        ] = ",".join([str(i) for i in visible_tpu_chips])
+        if num_visible_tpu_chips == 1:
+            os.environ[
+                TPU_CHIPS_PER_HOST_BOUNDS_ENV_VAR
+            ] = TPU_CHIPS_PER_HOST_BOUNDS_1_CHIP_CONFIG
+            os.environ[TPU_HOST_BOUNDS_ENV_VAR] = TPU_SINGLE_HOST_BOUNDS
+        elif num_visible_tpu_chips == 2:
+            os.environ[
+                TPU_CHIPS_PER_HOST_BOUNDS_ENV_VAR
+            ] = TPU_CHIPS_PER_HOST_BOUNDS_2_CHIP_CONFIG
+            os.environ[TPU_HOST_BOUNDS_ENV_VAR] = TPU_SINGLE_HOST_BOUNDS
+
+    @staticmethod
+    def get_current_node_tpu_pod_type() -> Optional[str]:
+        """Get the TPU pod type of the current node if applicable.
+
+        Individual TPU VMs within a TPU pod must know what type
+        of pod it is a part of. This is necessary for the
+        ML framework to work properly.
+
+        The logic is different if the TPU was provisioned via:
+        ```
+        gcloud tpus tpu-vm create ...
+        ```
+        (i.e. a GCE VM), vs through GKE:
+        - GCE VMs will always have a metadata server to poll this info
+        - GKE VMS will have environment variables preset.
+
+        Returns:
+            A string representing the current TPU pod type, e.g.
+            v4-16.
+
+        """
+        # Start with GKE-based check
+        accelerator_type = os.getenv(GKE_TPU_ACCELERATOR_TYPE_ENV_VAR, "")
+        if not accelerator_type:
+            # GCE-based VM check
+            accelerator_type = _get_tpu_metadata(key=GCE_TPU_ACCELERATOR_KEY)
+        if accelerator_type and TPUAcceleratorManager.is_valid_tpu_accelerator_type(
+            tpu_accelerator_type=accelerator_type
+        ):
+            if accelerator_type.lower().startswith("tpu"):
+                return "v" + accelerator_type.lower()[3:]
+
+            return accelerator_type
+        logging.debug("Failed to get a valid accelerator type.")
+        return None
+
+    @staticmethod
+    def get_current_node_tpu_name() -> Optional[str]:
+        """Return the name of the TPU pod that this worker node is a part of.
+
+        For instance, if the TPU was created with name "my-tpu", this function
+        will return "my-tpu".
+
+        If created through the Ray cluster launcher, the
+        name will typically be something like "ray-my-tpu-cluster-worker-aa946781-tpu".
+
+        In case the TPU was created through KubeRay, we currently expect that the
+        environment variable TPU_NAME is set per TPU pod slice, in which case
+        this function will return the value of that environment variable.
+
+        """
+        try:
+            # Start with GKE-based check
+            tpu_name = os.getenv(GKE_TPU_NAME_ENV_VAR, None)
+            if not tpu_name:
+                # GCE-based VM check
+                tpu_name = _get_tpu_metadata(key=GCE_TPU_INSTANCE_ID_KEY)
+            return tpu_name
+        except ValueError as e:
+            logging.debug("Could not get TPU name: %s", e)
+            return None
+
+    @staticmethod
+    def get_current_node_tpu_worker_id() -> Optional[int]:
+        """Return the worker index of the TPU pod."""
+        try:
+            # Start with GKE-based check
+            worker_id = os.getenv(GKE_TPU_WORKER_ID_ENV_VAR, None)
+            if not worker_id:
+                # GCE-based VM check
+                worker_id = _get_tpu_metadata(key=GCE_TPU_WORKER_ID_KEY)
+            if worker_id:
+                return int(worker_id)
+            else:
+                return None
+        except ValueError as e:
+            logging.debug("Could not get TPU worker id: %s", e)
+            return None
+
+    @staticmethod
+    def get_num_workers_in_current_tpu_pod() -> Optional[int]:
+        """Return the total number of workers in a TPU pod."""
+        tpu_pod_type = TPUAcceleratorManager.get_current_node_tpu_pod_type()
+        chips_per_host = TPUAcceleratorManager.get_current_node_num_accelerators()
+        cores_per_chip = get_tpu_cores_per_chip(tpu_pod_type)  # Hard-coded map.
+        cores_per_host = chips_per_host * cores_per_chip
+        if tpu_pod_type and cores_per_host > 0:
+            num_cores = int(tpu_pod_type.split("-")[1])
+            num_workers = num_cores // cores_per_host
+            # If the chip count doesn't fill a full host, a sub-host is still treated as a host.
+            if num_cores % cores_per_host != 0:
+                num_workers += 1
+            return num_workers
+        else:
+            logging.debug("Could not get num workers in TPU pod.")
+            return None
+
+    @staticmethod
+    def get_current_node_tpu_topology() -> Optional[str]:
+        try:
+            # Attempt GKE based lookup first
+            if topology := os.environ.get(GKE_TPU_TOPOLOGY_ENV_VAR):
+                return topology.strip().lower()
+            # GCE-based VM check using TPU env string.
+            tpu_env = _get_tpu_metadata(key=GCE_TPU_ENV_KEY)
+            if tpu_env:
+                topology = re.search(r"TOPOLOGY:\s*'([^']+)'", tpu_env)
+                if topology:
+                    return topology.group(1).strip().lower()
+        except ValueError as e:
+            logging.debug("Could not get TPU topology: %s", e)
+            return None
+
+    @staticmethod
+    def get_current_node_accelerator_type() -> Optional[str]:
+        """Attempt to detect the TPU accelerator type.
+
+        The output of this function will return the "ray accelerator type"
+        resource (e.g. TPU-V4) that indicates the TPU version.
+
+        We also expect that our TPU nodes contain a "TPU pod type"
+        resource, which indicates information about the topology of
+        the TPU pod slice.
+
+        We expect that the "TPU pod type" resource to be used when
+        running multi host workers, i.e. when TPU units are pod slices.
+
+        We expect that the "ray accelerator type" resource to be used when
+        running single host workers, i.e. when TPU units are single hosts.
+
+        Returns:
+            A string representing the TPU accelerator type,
+            e.g. "TPU-V2", "TPU-V3", "TPU-V4" if applicable, else None.
+
+        """
+
+        def tpu_pod_type_to_ray_accelerator_type(
+            tpu_pod_type: str,
+        ) -> Optional[str]:
+            return "TPU-" + str(tpu_pod_type.split("-")[0].upper())
+
+        ray_accelerator_type = None
+        tpu_pod_type = TPUAcceleratorManager.get_current_node_tpu_pod_type()
+
+        if tpu_pod_type is not None:
+            ray_accelerator_type = tpu_pod_type_to_ray_accelerator_type(
+                tpu_pod_type=tpu_pod_type
+            )
+            if ray_accelerator_type is None:
+                logger.info(
+                    "While trying to autodetect a TPU type, "
+                    f"received malformed accelerator_type: {tpu_pod_type}"
+                )
+
+        if ray_accelerator_type is None:
+            logging.info("Failed to auto-detect TPU type.")
+
+        return ray_accelerator_type
+
+    @staticmethod
+    def get_current_node_additional_resources() -> Optional[Dict[str, float]]:
+        """Get additional resources required for TPU nodes.
+
+        This will populate the TPU pod type and the TPU name which
+        is used for TPU pod execution.
+
+        When running workloads on a TPU pod, we need a way to run
+        the same binary on every worker in the TPU pod.
+
+        See https://jax.readthedocs.io/en/latest/multi_process.html
+        for more information.
+
+        To do this in ray, we take advantage of custom resources. We
+        mark worker 0 of the TPU pod as a "coordinator" that identifies
+        the other workers in the TPU pod. We therefore need:
+        - worker 0 to be targetable.
+        - all workers in the TPU pod to have a unique identifier consistent
+        within a TPU pod.
+
+        So assuming we want to run the following workload:
+
+        @ray.remote
+        def my_jax_fn():
+            import jax
+            return jax.device_count()
+
+        We could broadcast this on a TPU pod (e.g. a v4-16) as follows:
+
+        @ray.remote(resources={"TPU-v4-16-head"})
+        def run_jax_fn(executable):
+            # Note this will execute on worker 0
+            tpu_name = ray.util.tpu.get_current_pod_name()
+            num_hosts = ray.util.tpu.get_current_pod_worker_count()
+            tpu_executable = executable.options(resources={"TPU": 4, tpu_name: 1})
+            return [tpu_executable.remote() for _ in range(num_hosts)]
+
+        Returns:
+            A dictionary representing additional resources that may be
+            necessary for a particular accelerator type.
+
+        """
+        resources = {}
+        tpu_name = TPUAcceleratorManager.get_current_node_tpu_name()
+        worker_id = TPUAcceleratorManager.get_current_node_tpu_worker_id()
+        tpu_pod_type = TPUAcceleratorManager.get_current_node_tpu_pod_type()
+
+        if tpu_name and worker_id is not None and tpu_pod_type:
+            pod_head_resource_name = f"TPU-{tpu_pod_type}-head"
+            # Add the name of the TPU to the resource.
+            resources[tpu_name] = 1
+            # Only add in the TPU pod type resource to worker 0.
+            if worker_id == 0:
+                resources[pod_head_resource_name] = 1
+        else:
+            logging.info(
+                "Failed to configure TPU pod. Got: "
+                "tpu_name: %s, worker_id: %s, accelerator_type: %s",
+                tpu_name,
+                worker_id,
+                tpu_pod_type,
+            )
+        if resources:
+            return resources
+        return None
+
+    @staticmethod
+    def get_current_node_accelerator_labels() -> Dict[str, str]:
+        """Get default TPU-specific Ray node labels for the current node.
+
+        For TPUs, these labels include:
+        - ray.io/tpu-slice-name: the name of the TPU Pod or slice
+        - ray.io/tpu-worker-id: the integer worker ID within the slice
+        - ray.io/tpu-topology: the TPU topology (e.g. 4x4)
+        - ray.io/tpu-pod-type: the TPU pod type (e.g. v4-8)
+
+        Returns:
+            A dictionary of TPU label keys and resolved values.
+        """
+        tpu_labels = {}
+
+        tpu_name = TPUAcceleratorManager.get_current_node_tpu_name()
+        if tpu_name:
+            tpu_labels[ray._raylet.RAY_NODE_TPU_SLICE_NAME_KEY] = tpu_name
+
+        worker_id = TPUAcceleratorManager.get_current_node_tpu_worker_id()
+        if worker_id is not None:
+            tpu_labels[ray._raylet.RAY_NODE_TPU_WORKER_ID_KEY] = str(worker_id)
+
+        tpu_topology = TPUAcceleratorManager.get_current_node_tpu_topology()
+        if tpu_topology:
+            tpu_labels[ray._raylet.RAY_NODE_TPU_TOPOLOGY_KEY] = tpu_topology
+
+        pod_type = TPUAcceleratorManager.get_current_node_tpu_pod_type()
+        if pod_type:
+            tpu_labels[ray._raylet.RAY_NODE_TPU_POD_TYPE_KEY] = pod_type
+
+        return tpu_labels

python/ray/_private/arrow_serialization.py

@@ -0,0 +1,798 @@
+# arrow_serialization.py must resides outside of ray.data, otherwise
+# it causes circular dependency issues for AsyncActors due to
+# ray.data's lazy import.
+# see https://github.com/ray-project/ray/issues/30498 for more context.
+import logging
+import os
+import sys
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Callable, List, Optional, Tuple
+
+from ray._private.utils import is_in_test
+
+if TYPE_CHECKING:
+    import pyarrow
+
+RAY_DISABLE_CUSTOM_ARROW_JSON_OPTIONS_SERIALIZATION = (
+    "RAY_DISABLE_CUSTOM_ARROW_JSON_OPTIONS_SERIALIZATION"
+)
+RAY_DISABLE_CUSTOM_ARROW_DATA_SERIALIZATION = (
+    "RAY_DISABLE_CUSTOM_ARROW_DATA_SERIALIZATION"
+)
+
+logger = logging.getLogger(__name__)
+
+# Whether we have already warned the user about bloated fallback serialization.
+_serialization_fallback_set = set()
+
+
+def _register_custom_datasets_serializers(serialization_context):
+    try:
+        import pyarrow as pa  # noqa: F401
+    except ModuleNotFoundError:
+        # No pyarrow installed so not using Arrow, so no need for custom serializers.
+        return
+
+    # Register all custom serializers required by Datasets.
+    _register_arrow_data_serializer(serialization_context)
+    _register_arrow_json_readoptions_serializer(serialization_context)
+    _register_arrow_json_parseoptions_serializer(serialization_context)
+
+
+# Register custom Arrow JSON ReadOptions serializer to workaround it not being picklable
+# in Arrow < 8.0.0.
+def _register_arrow_json_readoptions_serializer(serialization_context):
+    if (
+        os.environ.get(
+            RAY_DISABLE_CUSTOM_ARROW_JSON_OPTIONS_SERIALIZATION,
+            "0",
+        )
+        == "1"
+    ):
+        return
+
+    import pyarrow.json as pajson
+
+    serialization_context._register_cloudpickle_serializer(
+        pajson.ReadOptions,
+        custom_serializer=lambda opts: (opts.use_threads, opts.block_size),
+        custom_deserializer=lambda args: pajson.ReadOptions(*args),
+    )
+
+
+def _register_arrow_json_parseoptions_serializer(serialization_context):
+    if (
+        os.environ.get(
+            RAY_DISABLE_CUSTOM_ARROW_JSON_OPTIONS_SERIALIZATION,
+            "0",
+        )
+        == "1"
+    ):
+        return
+
+    import pyarrow.json as pajson
+
+    serialization_context._register_cloudpickle_serializer(
+        pajson.ParseOptions,
+        custom_serializer=lambda opts: (
+            opts.explicit_schema,
+            opts.newlines_in_values,
+            opts.unexpected_field_behavior,
+        ),
+        custom_deserializer=lambda args: pajson.ParseOptions(*args),
+    )
+
+
+# Register custom Arrow data serializer to work around zero-copy slice pickling bug.
+# See https://issues.apache.org/jira/browse/ARROW-10739.
+def _register_arrow_data_serializer(serialization_context):
+    """Custom reducer for Arrow data that works around a zero-copy slicing pickling
+    bug by using the Arrow IPC format for the underlying serialization.
+
+    Background:
+        Arrow has both array-level slicing and buffer-level slicing; both are zero-copy,
+        but the former has a serialization bug where the entire buffer is serialized
+        instead of just the slice, while the latter's serialization works as expected
+        and only serializes the slice of the buffer. I.e., array-level slicing doesn't
+        propagate the slice down to the buffer when serializing the array.
+
+        We work around this by registering a custom cloudpickle reducers for Arrow
+        Tables that delegates serialization to the Arrow IPC format; thankfully, Arrow's
+        IPC serialization has fixed this buffer truncation bug.
+
+    See https://issues.apache.org/jira/browse/ARROW-10739.
+    """
+    if os.environ.get(RAY_DISABLE_CUSTOM_ARROW_DATA_SERIALIZATION, "0") == "1":
+        return
+
+    import pyarrow as pa
+
+    serialization_context._register_cloudpickle_reducer(pa.Table, _arrow_table_reduce)
+    serialization_context._register_cloudpickle_reducer(pa.Schema, _arrow_schema_reduce)
+
+
+def _arrow_schema_reduce(
+    schema: "pyarrow.Schema",
+) -> Tuple[Callable[["bytes"], "pyarrow.Schema"], Tuple[bytes]]:
+    """Custom reducer for Arrow Schema that uses IPC serialization for performance.
+
+    Arrow's native IPC serialization for schemas is significantly faster than
+    cloudpickle (10-20x for serialization, 2-3x for deserialization), making
+    this optimization particularly valuable for workloads with large schemas.
+    """
+    # Use Arrow's native IPC serialization which is much faster than cloudpickle
+    return _restore_schema_from_ipc, (schema.serialize().to_pybytes(),)
+
+
+def _restore_schema_from_ipc(buf: bytes) -> "pyarrow.Schema":
+    """Restore an Arrow Schema serialized to Arrow IPC format."""
+    import pyarrow as pa
+
+    return pa.ipc.read_schema(pa.BufferReader(buf))
+
+
+def _arrow_table_reduce(t: "pyarrow.Table"):
+    """Custom reducer for Arrow Tables that works around a zero-copy slice pickling bug.
+    Background:
+        Arrow has both array-level slicing and buffer-level slicing; both are zero-copy,
+        but the former has a serialization bug where the entire buffer is serialized
+        instead of just the slice, while the latter's serialization works as expected
+        and only serializes the slice of the buffer. I.e., array-level slicing doesn't
+        propagate the slice down to the buffer when serializing the array.
+        All that these copy methods do is, at serialization time, take the array-level
+        slicing and translate them to buffer-level slicing, so only the buffer slice is
+        sent over the wire instead of the entire buffer.
+    See https://issues.apache.org/jira/browse/ARROW-10739.
+    """
+    global _serialization_fallback_set
+
+    # Reduce the ChunkedArray columns.
+    reduced_columns = []
+    for column_name in t.column_names:
+        column = t[column_name]
+        try:
+            # Delegate to ChunkedArray reducer.
+            reduced_column = _arrow_chunked_array_reduce(column)
+        except Exception as e:
+            if not _is_dense_union(column.type) and is_in_test():
+                # If running in a test and the column is not a dense union array
+                # (which we expect to need a fallback), we want to raise the error,
+                # not fall back.
+                raise e from None
+            if type(column.type) not in _serialization_fallback_set:
+                logger.warning(
+                    "Failed to complete optimized serialization of Arrow Table, "
+                    f"serialization of column '{column_name}' of type {column.type} "
+                    "failed, so we're falling back to Arrow IPC serialization for the "
+                    "table. Note that this may result in slower serialization and more "
+                    "worker memory utilization. Serialization error:",
+                    exc_info=True,
+                )
+                _serialization_fallback_set.add(type(column.type))
+            # Fall back to Arrow IPC-based workaround for the entire table.
+            return _arrow_table_ipc_reduce(t)
+        else:
+            # Column reducer succeeded, add reduced column to list.
+            reduced_columns.append(reduced_column)
+    return _reconstruct_table, (reduced_columns, t.schema)
+
+
+def _reconstruct_table(
+    reduced_columns: List[Tuple[List["pyarrow.Array"], "pyarrow.DataType"]],
+    schema: "pyarrow.Schema",
+) -> "pyarrow.Table":
+    """Restore a serialized Arrow Table, reconstructing each reduced column."""
+    import pyarrow as pa
+
+    # Reconstruct each reduced column.
+    columns = []
+    for chunks_payload, type_ in reduced_columns:
+        columns.append(_reconstruct_chunked_array(chunks_payload, type_))
+
+    return pa.Table.from_arrays(columns, schema=schema)
+
+
+def _arrow_chunked_array_reduce(
+    ca: "pyarrow.ChunkedArray",
+) -> Tuple[List["PicklableArrayPayload"], "pyarrow.DataType"]:
+    """Custom reducer for Arrow ChunkedArrays that works around a zero-copy slice
+    pickling bug. This reducer does not return a reconstruction function, since it's
+    expected to be reconstructed by the Arrow Table reconstructor.
+    """
+    # Convert chunks to serialization payloads.
+    chunk_payloads = []
+    for chunk in ca.chunks:
+        chunk_payload = PicklableArrayPayload.from_array(chunk)
+        chunk_payloads.append(chunk_payload)
+    return chunk_payloads, ca.type
+
+
+def _reconstruct_chunked_array(
+    chunks: List["PicklableArrayPayload"], type_: "pyarrow.DataType"
+) -> "pyarrow.ChunkedArray":
+    """Restore a serialized Arrow ChunkedArray from chunks and type."""
+    import pyarrow as pa
+
+    # Reconstruct chunks from serialization payloads.
+    chunks = [chunk.to_array() for chunk in chunks]
+    return pa.chunked_array(chunks, type_)
+
+
+@dataclass
+class PicklableArrayPayload:
+    """Picklable array payload, holding data buffers and array metadata.
+
+    This is a helper container for pickling and reconstructing nested Arrow Arrays while
+    ensuring that the buffers that underly zero-copy slice views are properly truncated.
+    """
+
+    # Array type.
+    type: "pyarrow.DataType"
+    # Length of array.
+    length: int
+    # Underlying data buffers.
+    buffers: List["pyarrow.Buffer"]
+    # Cached null count.
+    null_count: int
+    # Slice offset into base array.
+    offset: int
+    # Serialized array payloads for nested (child) arrays.
+    children: List["PicklableArrayPayload"]
+
+    @classmethod
+    def from_array(self, a: "pyarrow.Array") -> "PicklableArrayPayload":
+        """Create a picklable array payload from an Arrow Array.
+
+        This will recursively accumulate data buffer and metadata payloads that are
+        ready for pickling; namely, the data buffers underlying zero-copy slice views
+        will be properly truncated.
+        """
+        return _array_to_array_payload(a)
+
+    def to_array(self) -> "pyarrow.Array":
+        """Reconstruct an Arrow Array from this picklable payload."""
+        return _array_payload_to_array(self)
+
+
+def _array_payload_to_array(payload: "PicklableArrayPayload") -> "pyarrow.Array":
+    """Reconstruct an Arrow Array from a possibly nested PicklableArrayPayload."""
+    import pyarrow as pa
+
+    children = [child_payload.to_array() for child_payload in payload.children]
+
+    if pa.types.is_dictionary(payload.type):
+        # Dedicated path for reconstructing a DictionaryArray, since
+        # Array.from_buffers() doesn't work for DictionaryArrays.
+        assert len(children) == 2, len(children)
+        indices, dictionary = children
+        return pa.DictionaryArray.from_arrays(
+            indices,
+            dictionary,
+            ordered=payload.type.ordered,  # Explicitly pass the ordered flag to from_arrays() to prevent dropping it as ordered=False by default
+        )
+    elif pa.types.is_map(payload.type) and len(children) > 1:
+        # In pyarrow<7.0.0, the underlying map child array is not exposed, so we work
+        # with the key and item arrays.
+        assert len(children) == 3, len(children)
+        offsets, keys, items = children
+        return pa.MapArray.from_arrays(offsets, keys, items)
+    elif isinstance(payload.type, pa.BaseExtensionType):
+        assert len(children) == 1, len(children)
+        storage = children[0]
+        return payload.type.wrap_array(storage)
+    else:
+        # Common case: use Array.from_buffers() to construct an array of a certain type.
+        return pa.Array.from_buffers(
+            type=payload.type,
+            length=payload.length,
+            buffers=payload.buffers,
+            null_count=payload.null_count,
+            offset=payload.offset,
+            children=children,
+        )
+
+
+def _array_to_array_payload(a: "pyarrow.Array") -> "PicklableArrayPayload":
+    """Serialize an Arrow Array to an PicklableArrayPayload for later pickling.
+
+    This function's primary purpose is to dispatch to the handler for the input array
+    type.
+    """
+    import pyarrow as pa
+
+    if _is_dense_union(a.type):
+        # Dense unions are not supported.
+        # TODO(Clark): Support dense unions.
+        raise NotImplementedError(
+            "Custom slice view serialization of dense union arrays is not yet "
+            "supported."
+        )
+
+    # Dispatch to handler for array type.
+    if pa.types.is_null(a.type):
+        return _null_array_to_array_payload(a)
+    elif _is_primitive(a.type):
+        return _primitive_array_to_array_payload(a)
+    elif _is_binary(a.type):
+        return _binary_array_to_array_payload(a)
+    elif pa.types.is_list(a.type) or pa.types.is_large_list(a.type):
+        return _list_array_to_array_payload(a)
+    elif pa.types.is_fixed_size_list(a.type):
+        return _fixed_size_list_array_to_array_payload(a)
+    elif pa.types.is_struct(a.type):
+        return _struct_array_to_array_payload(a)
+    elif pa.types.is_union(a.type):
+        return _union_array_to_array_payload(a)
+    elif pa.types.is_dictionary(a.type):
+        return _dictionary_array_to_array_payload(a)
+    elif pa.types.is_map(a.type):
+        return _map_array_to_array_payload(a)
+    elif isinstance(a.type, pa.BaseExtensionType):
+        return _extension_array_to_array_payload(a)
+    else:
+        raise ValueError("Unhandled Arrow array type:", a.type)
+
+
+def _is_primitive(type_: "pyarrow.DataType") -> bool:
+    """Whether the provided Array type is primitive (boolean, numeric, temporal or
+    fixed-size binary)."""
+    import pyarrow as pa
+
+    return (
+        pa.types.is_integer(type_)
+        or pa.types.is_floating(type_)
+        or pa.types.is_decimal(type_)
+        or pa.types.is_boolean(type_)
+        or pa.types.is_temporal(type_)
+        or pa.types.is_fixed_size_binary(type_)
+    )
+
+
+def _is_binary(type_: "pyarrow.DataType") -> bool:
+    """Whether the provided Array type is a variable-sized binary type."""
+    import pyarrow as pa
+
+    return (
+        pa.types.is_string(type_)
+        or pa.types.is_large_string(type_)
+        or pa.types.is_binary(type_)
+        or pa.types.is_large_binary(type_)
+    )
+
+
+def _null_array_to_array_payload(a: "pyarrow.NullArray") -> "PicklableArrayPayload":
+    """Serialize null array to PicklableArrayPayload."""
+    # Buffer scheme: [None]
+    return PicklableArrayPayload(
+        type=a.type,
+        length=len(a),
+        buffers=[None],  # Single null buffer is expected.
+        null_count=a.null_count,
+        offset=0,
+        children=[],
+    )
+
+
+def _primitive_array_to_array_payload(a: "pyarrow.Array") -> "PicklableArrayPayload":
+    """Serialize primitive (numeric, temporal, boolean) arrays to
+    PicklableArrayPayload.
+    """
+    assert _is_primitive(a.type), a.type
+    # Buffer scheme: [bitmap, data]
+    buffers = a.buffers()
+    assert len(buffers) == 2, len(buffers)
+
+    # Copy bitmap buffer, if needed.
+    bitmap_buf = buffers[0]
+    if a.null_count > 0:
+        bitmap_buf = _copy_bitpacked_buffer_if_needed(bitmap_buf, a.offset, len(a))
+    else:
+        bitmap_buf = None
+
+    # Copy data buffer, if needed.
+    data_buf = buffers[1]
+    if data_buf is not None:
+        data_buf = _copy_buffer_if_needed(buffers[1], a.type, a.offset, len(a))
+
+    return PicklableArrayPayload(
+        type=a.type,
+        length=len(a),
+        buffers=[bitmap_buf, data_buf],
+        null_count=a.null_count,
+        offset=0,
+        children=[],
+    )
+
+
+def _binary_array_to_array_payload(a: "pyarrow.Array") -> "PicklableArrayPayload":
+    """Serialize binary (variable-sized binary, string) arrays to
+    PicklableArrayPayload.
+    """
+    assert _is_binary(a.type), a.type
+    # Buffer scheme: [bitmap, value_offsets, data]
+    buffers = a.buffers()
+    assert len(buffers) == 3, len(buffers)
+
+    # Copy bitmap buffer, if needed.
+    if a.null_count > 0:
+        bitmap_buf = _copy_bitpacked_buffer_if_needed(buffers[0], a.offset, len(a))
+    else:
+        bitmap_buf = None
+
+    # Copy offset buffer, if needed.
+    offset_buf = buffers[1]
+    offset_buf, data_offset, data_length = _copy_offsets_buffer_if_needed(
+        offset_buf, a.type, a.offset, len(a)
+    )
+    data_buf = buffers[2]
+    data_buf = _copy_buffer_if_needed(data_buf, None, data_offset, data_length)
+    return PicklableArrayPayload(
+        type=a.type,
+        length=len(a),
+        buffers=[bitmap_buf, offset_buf, data_buf],
+        null_count=a.null_count,
+        offset=0,
+        children=[],
+    )
+
+
+def _list_array_to_array_payload(a: "pyarrow.Array") -> "PicklableArrayPayload":
+    """Serialize list (regular and large) arrays to PicklableArrayPayload."""
+    # Dedicated path for ListArrays. These arrays have a nested set of bitmap and
+    # offset buffers, eventually bottoming out on a data buffer.
+    # Buffer scheme:
+    # [bitmap, offsets, bitmap, offsets, ..., bitmap, data]
+    buffers = a.buffers()
+    assert len(buffers) > 1, len(buffers)
+
+    # Copy bitmap buffer, if needed.
+    if a.null_count > 0:
+        bitmap_buf = _copy_bitpacked_buffer_if_needed(buffers[0], a.offset, len(a))
+    else:
+        bitmap_buf = None
+
+    # Copy offset buffer, if needed.
+    offset_buf = buffers[1]
+    offset_buf, child_offset, child_length = _copy_offsets_buffer_if_needed(
+        offset_buf, a.type, a.offset, len(a)
+    )
+
+    # Propagate slice to child.
+    child = a.values.slice(child_offset, child_length)
+
+    return PicklableArrayPayload(
+        type=a.type,
+        length=len(a),
+        buffers=[bitmap_buf, offset_buf],
+        null_count=a.null_count,
+        offset=0,
+        children=[_array_to_array_payload(child)],
+    )
+
+
+def _fixed_size_list_array_to_array_payload(
+    a: "pyarrow.FixedSizeListArray",
+) -> "PicklableArrayPayload":
+    """Serialize fixed size list arrays to PicklableArrayPayload."""
+    # Dedicated path for fixed-size lists.
+    # Buffer scheme:
+    # [bitmap, values_bitmap, values_data, values_subbuffers...]
+    buffers = a.buffers()
+    assert len(buffers) >= 1, len(buffers)
+
+    # Copy bitmap buffer, if needed.
+    if a.null_count > 0:
+        bitmap_buf = _copy_bitpacked_buffer_if_needed(buffers[0], a.offset, len(a))
+    else:
+        bitmap_buf = None
+
+    # Propagate slice to child.
+    child_offset = a.type.list_size * a.offset
+    child_length = a.type.list_size * len(a)
+    child = a.values.slice(child_offset, child_length)
+
+    return PicklableArrayPayload(
+        type=a.type,
+        length=len(a),
+        buffers=[bitmap_buf],
+        null_count=a.null_count,
+        offset=0,
+        children=[_array_to_array_payload(child)],
+    )
+
+
+def _struct_array_to_array_payload(a: "pyarrow.StructArray") -> "PicklableArrayPayload":
+    """Serialize struct arrays to PicklableArrayPayload."""
+    # Dedicated path for StructArrays.
+    # StructArrays have a top-level bitmap buffer and one or more children arrays.
+    # Buffer scheme: [bitmap, None, child_bitmap, child_data, ...]
+    buffers = a.buffers()
+    assert len(buffers) >= 1, len(buffers)
+
+    # Copy bitmap buffer, if needed.
+    if a.null_count > 0:
+        bitmap_buf = _copy_bitpacked_buffer_if_needed(buffers[0], a.offset, len(a))
+    else:
+        bitmap_buf = None
+
+    # Get field children payload.
+    # Offsets and truncations are already propagated to the field arrays, so we can
+    # serialize them as-is.
+    children = [_array_to_array_payload(a.field(i)) for i in range(a.type.num_fields)]
+    return PicklableArrayPayload(
+        type=a.type,
+        length=len(a),
+        buffers=[bitmap_buf],
+        null_count=a.null_count,
+        offset=0,
+        children=children,
+    )
+
+
+def _union_array_to_array_payload(a: "pyarrow.UnionArray") -> "PicklableArrayPayload":
+    """Serialize union arrays to PicklableArrayPayload."""
+    import pyarrow as pa
+
+    # Dedicated path for UnionArrays.
+    # UnionArrays have a top-level bitmap buffer and type code buffer, and one or
+    # more children arrays.
+    # Buffer scheme: [None, typecodes, child_bitmap, child_data, ...]
+    assert not _is_dense_union(a.type)
+    buffers = a.buffers()
+    assert len(buffers) > 1, len(buffers)
+
+    bitmap_buf = buffers[0]
+    assert bitmap_buf is None, bitmap_buf
+
+    # Copy type code buffer, if needed.
+    type_code_buf = buffers[1]
+    type_code_buf = _copy_buffer_if_needed(type_code_buf, pa.int8(), a.offset, len(a))
+
+    # Get field children payload.
+    # Offsets and truncations are already propagated to the field arrays, so we can
+    # serialize them as-is.
+    children = [_array_to_array_payload(a.field(i)) for i in range(a.type.num_fields)]
+    return PicklableArrayPayload(
+        type=a.type,
+        length=len(a),
+        buffers=[bitmap_buf, type_code_buf],
+        null_count=a.null_count,
+        offset=0,
+        children=children,
+    )
+
+
+def _dictionary_array_to_array_payload(
+    a: "pyarrow.DictionaryArray",
+) -> "PicklableArrayPayload":
+    """Serialize dictionary arrays to PicklableArrayPayload."""
+    # Dedicated path for DictionaryArrays.
+    # Buffer scheme: [indices_bitmap, indices_data] (dictionary stored separately)
+    indices_payload = _array_to_array_payload(a.indices)
+    dictionary_payload = _array_to_array_payload(a.dictionary)
+    return PicklableArrayPayload(
+        type=a.type,
+        length=len(a),
+        buffers=[],
+        null_count=a.null_count,
+        offset=0,
+        children=[indices_payload, dictionary_payload],
+    )
+
+
+def _map_array_to_array_payload(a: "pyarrow.MapArray") -> "PicklableArrayPayload":
+    """Serialize map arrays to PicklableArrayPayload."""
+    import pyarrow as pa
+
+    # Dedicated path for MapArrays.
+    # Buffer scheme: [bitmap, offsets, child_struct_array_buffers, ...]
+    buffers = a.buffers()
+    assert len(buffers) > 0, len(buffers)
+
+    # Copy bitmap buffer, if needed.
+    if a.null_count > 0:
+        bitmap_buf = _copy_bitpacked_buffer_if_needed(buffers[0], a.offset, len(a))
+    else:
+        bitmap_buf = None
+
+    new_buffers = [bitmap_buf]
+
+    # Copy offsets buffer, if needed.
+    offset_buf = buffers[1]
+    offset_buf, data_offset, data_length = _copy_offsets_buffer_if_needed(
+        offset_buf, a.type, a.offset, len(a)
+    )
+
+    if isinstance(a, pa.lib.ListArray):
+        # Map arrays directly expose the one child struct array in pyarrow>=7.0.0, which
+        # is easier to work with than the raw buffers.
+        new_buffers.append(offset_buf)
+        children = [_array_to_array_payload(a.values.slice(data_offset, data_length))]
+    else:
+        # In pyarrow<7.0.0, the child struct array is not exposed, so we work with the
+        # key and item arrays.
+        buffers = a.buffers()
+        assert len(buffers) > 2, len(buffers)
+        # Reconstruct offsets array.
+        offsets = pa.Array.from_buffers(
+            pa.int32(), len(a) + 1, [bitmap_buf, offset_buf]
+        )
+        # Propagate slice to keys.
+        keys = a.keys.slice(data_offset, data_length)
+        # Propagate slice to items.
+        items = a.items.slice(data_offset, data_length)
+        children = [
+            _array_to_array_payload(offsets),
+            _array_to_array_payload(keys),
+            _array_to_array_payload(items),
+        ]
+    return PicklableArrayPayload(
+        type=a.type,
+        length=len(a),
+        buffers=new_buffers,
+        null_count=a.null_count,
+        offset=0,
+        children=children,
+    )
+
+
+def _extension_array_to_array_payload(
+    a: "pyarrow.ExtensionArray",
+) -> "PicklableArrayPayload":
+    storage_payload = _array_to_array_payload(a.storage)
+    return PicklableArrayPayload(
+        type=a.type,
+        length=len(a),
+        buffers=[],
+        null_count=a.null_count,
+        offset=0,
+        children=[storage_payload],
+    )
+
+
+def _copy_buffer_if_needed(
+    buf: "pyarrow.Buffer",
+    type_: Optional["pyarrow.DataType"],
+    offset: int,
+    length: int,
+) -> "pyarrow.Buffer":
+    """Copy buffer, if needed."""
+    import pyarrow as pa
+
+    if type_ is not None and pa.types.is_boolean(type_):
+        # Arrow boolean array buffers are bit-packed, with 8 entries per byte,
+        # and are accessed via bit offsets.
+        buf = _copy_bitpacked_buffer_if_needed(buf, offset, length)
+    else:
+        type_bytewidth = type_.bit_width // 8 if type_ is not None else 1
+        buf = _copy_normal_buffer_if_needed(buf, type_bytewidth, offset, length)
+    return buf
+
+
+def _copy_normal_buffer_if_needed(
+    buf: "pyarrow.Buffer",
+    byte_width: int,
+    offset: int,
+    length: int,
+) -> "pyarrow.Buffer":
+    """Copy buffer, if needed."""
+    byte_offset = offset * byte_width
+    byte_length = length * byte_width
+    if offset > 0 or byte_length < buf.size:
+        # Array is a zero-copy slice, so we need to copy to a new buffer before
+        # serializing; this slice of the underlying buffer (not the array) will ensure
+        # that the buffer is properly copied at pickle-time.
+        buf = buf.slice(byte_offset, byte_length)
+    return buf
+
+
+def _copy_bitpacked_buffer_if_needed(
+    buf: "pyarrow.Buffer",
+    offset: int,
+    length: int,
+) -> "pyarrow.Buffer":
+    """Copy bit-packed binary buffer, if needed."""
+    bit_offset = offset % 8
+    byte_offset = offset // 8
+    byte_length = _bytes_for_bits(bit_offset + length) // 8
+    if offset > 0 or byte_length < buf.size:
+        buf = buf.slice(byte_offset, byte_length)
+        if bit_offset != 0:
+            # Need to manually shift the buffer to eliminate the bit offset.
+            buf = _align_bit_offset(buf, bit_offset, byte_length)
+    return buf
+
+
+def _copy_offsets_buffer_if_needed(
+    buf: "pyarrow.Buffer",
+    arr_type: "pyarrow.DataType",
+    offset: int,
+    length: int,
+) -> Tuple["pyarrow.Buffer", int, int]:
+    """Copy the provided offsets buffer, returning the copied buffer and the
+    offset + length of the underlying data.
+    """
+    import pyarrow as pa
+    import pyarrow.compute as pac
+
+    if (
+        pa.types.is_large_list(arr_type)
+        or pa.types.is_large_string(arr_type)
+        or pa.types.is_large_binary(arr_type)
+        or pa.types.is_large_unicode(arr_type)
+    ):
+        offset_type = pa.int64()
+    else:
+        offset_type = pa.int32()
+    # Copy offset buffer, if needed.
+    buf = _copy_buffer_if_needed(buf, offset_type, offset, length + 1)
+    # Reconstruct the offset array so we can determine the offset and length
+    # of the child array.
+    offsets = pa.Array.from_buffers(offset_type, length + 1, [None, buf])
+    child_offset = offsets[0].as_py()
+    child_length = offsets[-1].as_py() - child_offset
+    # Create new offsets aligned to 0 for the copied data buffer slice.
+    offsets = pac.subtract(offsets, child_offset)
+    if pa.types.is_int32(offset_type):
+        # We need to cast the resulting Int64Array back down to an Int32Array.
+        offsets = offsets.cast(offset_type, safe=False)
+    buf = offsets.buffers()[1]
+    return buf, child_offset, child_length
+
+
+def _bytes_for_bits(n: int) -> int:
+    """Round up n to the nearest multiple of 8.
+    This is used to get the byte-padded number of bits for n bits.
+    """
+    return (n + 7) & (-8)
+
+
+def _align_bit_offset(
+    buf: "pyarrow.Buffer",
+    bit_offset: int,
+    byte_length: int,
+) -> "pyarrow.Buffer":
+    """Align the bit offset into the buffer with the front of the buffer by shifting
+    the buffer and eliminating the offset.
+    """
+    import pyarrow as pa
+
+    bytes_ = buf.to_pybytes()
+    bytes_as_int = int.from_bytes(bytes_, sys.byteorder)
+    bytes_as_int >>= bit_offset
+    bytes_ = bytes_as_int.to_bytes(byte_length, sys.byteorder)
+    return pa.py_buffer(bytes_)
+
+
+def _arrow_table_ipc_reduce(table: "pyarrow.Table"):
+    """Custom reducer for Arrow Table that works around a zero-copy slicing pickling
+    bug by using the Arrow IPC format for the underlying serialization.
+
+    This is currently used as a fallback for unsupported types (or unknown bugs) for
+    the manual buffer truncation workaround, e.g. for dense unions.
+    """
+    from pyarrow.ipc import RecordBatchStreamWriter
+    from pyarrow.lib import BufferOutputStream
+
+    output_stream = BufferOutputStream()
+    with RecordBatchStreamWriter(output_stream, schema=table.schema) as wr:
+        wr.write_table(table)
+    # NOTE: output_stream.getvalue() materializes the serialized table to a single
+    # contiguous bytestring, resulting in a few copy. This adds 1-2 extra copies on the
+    # serialization side, and 1 extra copy on the deserialization side.
+    return _restore_table_from_ipc, (output_stream.getvalue(),)
+
+
+def _restore_table_from_ipc(buf: bytes) -> "pyarrow.Table":
+    """Restore an Arrow Table serialized to Arrow IPC format."""
+    from pyarrow.ipc import RecordBatchStreamReader
+
+    with RecordBatchStreamReader(buf) as reader:
+        return reader.read_all()
+
+
+def _is_dense_union(type_: "pyarrow.DataType") -> bool:
+    """Whether the provided Arrow type is a dense union."""
+    import pyarrow as pa
+
+    return pa.types.is_union(type_) and type_.mode == "dense"

python/ray/_private/async_compat.py

@@ -0,0 +1,54 @@
+"""
+This file should only be imported from Python 3.
+It will raise SyntaxError when importing from Python 2.
+"""
+import asyncio
+import inspect
+from functools import lru_cache
+
+from ray._private.ray_constants import env_bool
+
+try:
+    import uvloop
+except ImportError:
+    uvloop = None
+
+
+def get_new_event_loop():
+    """Construct a new event loop. Ray will use uvloop if it exists and is enabled"""
+    if uvloop and env_bool("RAY_USE_UVLOOP", True):
+        return uvloop.new_event_loop()
+    else:
+        return asyncio.new_event_loop()
+
+
+def try_install_uvloop():
+    """Installs uvloop as event-loop implementation for asyncio (if available and enabled)"""
+    if uvloop and env_bool("RAY_USE_UVLOOP", True):
+        uvloop.install()
+    else:
+        pass
+
+
+def is_async_func(func) -> bool:
+    """Return True if the function is an async or async generator method."""
+    return inspect.iscoroutinefunction(func) or inspect.isasyncgenfunction(func)
+
+
+@lru_cache(maxsize=2**10)
+def has_async_methods(cls: object) -> bool:
+    """Return True if the class has any async methods."""
+    return len(inspect.getmembers(cls, predicate=is_async_func)) > 0
+
+
+@lru_cache(maxsize=2**10)
+def sync_to_async(func):
+    """Wrap a blocking function in an async function"""
+
+    if is_async_func(func):
+        return func
+
+    async def wrapper(*args, **kwargs):
+        return func(*args, **kwargs)
+
+    return wrapper

python/ray/_private/async_utils.py

@@ -0,0 +1,65 @@
+# Adapted from [aiodebug](https://gitlab.com/quantlane/libs/aiodebug)
+
+# Copyright 2016-2022 Quantlane s.r.o.
+
+#    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
+
+#        http://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.
+
+# Modifications:
+# - Removed the dependency to `logwood`.
+# - Renamed `monitor_loop_lag.enable()` to just `enable_monitor_loop_lag()`.
+# - Miscellaneous changes to make it work with Ray.
+
+import asyncio
+import asyncio.events
+from typing import Callable, Optional, Set
+
+# Strong references to background tasks to prevent them from being garbage
+# collected mid-execution. The asyncio event loop only keeps weak references
+# to tasks, so a task with no other strong references can be collected at
+# any time. See https://docs.python.org/3/library/asyncio-task.html#asyncio.create_task
+_BACKGROUND_TASKS: Set[asyncio.Task] = set()
+
+
+def enable_monitor_loop_lag(
+    callback: Callable[[float], None],
+    interval_s: float = 0.25,
+    loop: Optional[asyncio.AbstractEventLoop] = None,
+) -> None:
+    """Start logging event loop lags to the callback.
+
+    In ideal circumstances they should be very close to zero. Lags may increase
+    if event loop callbacks block for too long.
+
+    Note: this works for all event loops, including uvloop.
+
+    Args:
+        callback: Callback to call with the lag in seconds.
+        interval_s: How often to measure the lag, in seconds.
+        loop: The event loop to monitor. If None, uses the running loop.
+    """
+    if loop is None:
+        loop = asyncio.get_running_loop()
+    if loop is None:
+        raise ValueError("No provided loop, nor running loop found.")
+
+    async def monitor():
+        while loop.is_running():
+            t0 = loop.time()
+            await asyncio.sleep(interval_s)
+            lag = loop.time() - t0 - interval_s  # Should be close to zero.
+            callback(lag)
+
+    task = loop.create_task(monitor(), name="async_utils.monitor_loop_lag")
+    # Keep a strong reference so the task isn't garbage collected mid-execution.
+    _BACKGROUND_TASKS.add(task)
+    task.add_done_callback(_BACKGROUND_TASKS.discard)

python/ray/_private/authentication/authentication_constants.py

@@ -0,0 +1,14 @@
+# Authentication error messages
+TOKEN_AUTH_ENABLED_BUT_NO_TOKEN_FOUND_ERROR_MESSAGE = (
+    "Token authentication is enabled but no authentication token was found"
+)
+
+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"
+
+AUTHENTICATION_TOKEN_COOKIE_NAME = "ray-authentication-token"
+AUTHENTICATION_TOKEN_COOKIE_MAX_AGE = 30 * 24 * 60 * 60  # 30 days

python/ray/_private/authentication/authentication_token_generator.py

@@ -0,0 +1,9 @@
+import secrets
+
+
+def generate_new_authentication_token() -> str:
+    """Generate an authentication token for the cluster.
+
+    256 bits of entropy is considered sufficient to be durable to brute force attacks.
+    """
+    return secrets.token_hex(32)

python/ray/_private/authentication/authentication_token_setup.py

@@ -0,0 +1,103 @@
+"""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_constants import (
+    TOKEN_AUTH_ENABLED_BUT_NO_TOKEN_FOUND_ERROR_MESSAGE,
+)
+from ray._private.authentication.authentication_token_generator import (
+    generate_new_authentication_token,
+)
+from ray._raylet import (
+    AuthenticationMode,
+    AuthenticationTokenLoader,
+    get_authentication_mode,
+)
+from ray.exceptions import AuthenticationError
+
+logger = logging.getLogger(__name__)
+
+
+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(ignore_auth_mode=True):
+        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 AuthenticationError(
+                TOKEN_AUTH_ENABLED_BUT_NO_TOKEN_FOUND_ERROR_MESSAGE
+            )

python/ray/_private/authentication/authentication_utils.py

@@ -0,0 +1,58 @@
+try:
+    from ray._raylet import (
+        AuthenticationMode,
+        get_authentication_mode,
+        validate_authentication_token,
+    )
+
+    _RAYLET_AVAILABLE = True
+except ImportError:
+    # ray._raylet not available during doc builds
+    _RAYLET_AVAILABLE = False
+
+
+def is_token_auth_enabled() -> bool:
+    """Check if token authentication is enabled.
+
+    Returns:
+        bool: True if AUTH_MODE is set to "token".
+    """
+    if not _RAYLET_AVAILABLE:
+        return False
+
+    return get_authentication_mode() == AuthenticationMode.TOKEN
+
+
+def validate_request_token(auth_header: str) -> bool:
+    """Validate the Authorization header from an HTTP request.
+
+    Args:
+        auth_header: The Authorization header value (e.g., "Bearer <token>")
+
+    Returns:
+        bool: True if token is valid, False otherwise
+    """
+    if not _RAYLET_AVAILABLE or not auth_header:
+        return False
+
+    # validate_authentication_token expects full "Bearer <token>" format
+    # and performs equality comparison via C++ layer
+    return validate_authentication_token(auth_header)
+
+
+def get_authentication_mode_name(mode: AuthenticationMode) -> str:
+    """Convert AuthenticationMode enum value to string name.
+
+    Args:
+        mode: AuthenticationMode enum value from ray._raylet
+
+    Returns:
+        String name: "disabled", "token"
+    """
+    from ray._raylet import AuthenticationMode
+
+    _MODE_NAMES = {
+        AuthenticationMode.DISABLED: "disabled",
+        AuthenticationMode.TOKEN: "token",
+    }
+    return _MODE_NAMES.get(mode, "unknown")

python/ray/_private/authentication/grpc_authentication_client_interceptor.py

@@ -0,0 +1,155 @@
+"""gRPC client interceptor for token-based authentication."""
+
+import logging
+from collections import namedtuple
+from typing import Tuple
+
+import grpc
+from grpc import aio as aiogrpc
+
+from ray._raylet import AuthenticationTokenLoader
+
+logger = logging.getLogger(__name__)
+
+
+# Named tuple to hold client call details
+_ClientCallDetails = namedtuple(
+    "_ClientCallDetails",
+    ("method", "timeout", "metadata", "credentials", "wait_for_ready", "compression"),
+)
+
+
+def _get_authentication_metadata_tuple() -> Tuple[Tuple[str, str], ...]:
+    """Get gRPC metadata tuple for authentication. Currently only supported for token authentication.
+
+    Returns:
+        tuple: Empty tuple or ((AUTHORIZATION_HEADER_NAME, "Bearer <token>"),)
+    """
+    token_loader = AuthenticationTokenLoader.instance()
+    if not token_loader.has_token():
+        return ()
+
+    headers = token_loader.get_token_for_http_header()
+
+    # Convert HTTP header dict to gRPC metadata tuple
+    # gRPC expects: (("key", "value"), ...)
+    return tuple((k, v) for k, v in headers.items())
+
+
+class SyncAuthenticationMetadataClientInterceptor(
+    grpc.UnaryUnaryClientInterceptor,
+    grpc.UnaryStreamClientInterceptor,
+    grpc.StreamUnaryClientInterceptor,
+    grpc.StreamStreamClientInterceptor,
+):
+    """Synchronous gRPC client interceptor that adds authentication metadata."""
+
+    def _intercept_call_details(self, client_call_details):
+        """Helper method to add authentication metadata to client call details."""
+        metadata = list(client_call_details.metadata or [])
+        metadata.extend(_get_authentication_metadata_tuple())
+
+        return _ClientCallDetails(
+            method=client_call_details.method,
+            timeout=client_call_details.timeout,
+            metadata=metadata,
+            credentials=client_call_details.credentials,
+            wait_for_ready=getattr(client_call_details, "wait_for_ready", None),
+            compression=getattr(client_call_details, "compression", None),
+        )
+
+    def intercept_unary_unary(self, continuation, client_call_details, request):
+        new_details = self._intercept_call_details(client_call_details)
+        return continuation(new_details, request)
+
+    def intercept_unary_stream(self, continuation, client_call_details, request):
+        new_details = self._intercept_call_details(client_call_details)
+        return continuation(new_details, request)
+
+    def intercept_stream_unary(
+        self, continuation, client_call_details, request_iterator
+    ):
+        new_details = self._intercept_call_details(client_call_details)
+        return continuation(new_details, request_iterator)
+
+    def intercept_stream_stream(
+        self, continuation, client_call_details, request_iterator
+    ):
+        new_details = self._intercept_call_details(client_call_details)
+        return continuation(new_details, request_iterator)
+
+
+def _intercept_call_details_async(client_call_details):
+    """Helper to add authentication metadata to client call details (async version)."""
+    metadata = list(client_call_details.metadata or [])
+    metadata.extend(_get_authentication_metadata_tuple())
+
+    return _ClientCallDetails(
+        method=client_call_details.method,
+        timeout=client_call_details.timeout,
+        metadata=metadata,
+        credentials=client_call_details.credentials,
+        wait_for_ready=getattr(client_call_details, "wait_for_ready", None),
+        compression=getattr(client_call_details, "compression", None),
+    )
+
+
+# NOTE: gRPC aio's Channel.__init__ uses if-elif chains to categorize interceptors,
+# so a single class inheriting from multiple interceptor types will only be registered
+# for the first matching type. We must use separate classes for each RPC type.
+# See: https://github.com/grpc/grpc/blob/master/src/python/grpcio/grpc/aio/_channel.py
+
+
+class _AsyncUnaryUnaryAuthInterceptor(aiogrpc.UnaryUnaryClientInterceptor):
+    """Async unary-unary interceptor that adds authentication metadata."""
+
+    async def intercept_unary_unary(self, continuation, client_call_details, request):
+        new_details = _intercept_call_details_async(client_call_details)
+        return await continuation(new_details, request)
+
+
+class _AsyncUnaryStreamAuthInterceptor(aiogrpc.UnaryStreamClientInterceptor):
+    """Async unary-stream interceptor that adds authentication metadata."""
+
+    async def intercept_unary_stream(self, continuation, client_call_details, request):
+        new_details = _intercept_call_details_async(client_call_details)
+        return await continuation(new_details, request)
+
+
+class _AsyncStreamUnaryAuthInterceptor(aiogrpc.StreamUnaryClientInterceptor):
+    """Async stream-unary interceptor that adds authentication metadata."""
+
+    async def intercept_stream_unary(
+        self, continuation, client_call_details, request_iterator
+    ):
+        new_details = _intercept_call_details_async(client_call_details)
+        return await continuation(new_details, request_iterator)
+
+
+class _AsyncStreamStreamAuthInterceptor(aiogrpc.StreamStreamClientInterceptor):
+    """Async stream-stream interceptor that adds authentication metadata."""
+
+    async def intercept_stream_stream(
+        self, continuation, client_call_details, request_iterator
+    ):
+        new_details = _intercept_call_details_async(client_call_details)
+        return await continuation(new_details, request_iterator)
+
+
+def get_async_auth_interceptors():
+    """Get a list of async authentication interceptors for all RPC types.
+
+    Returns a list of separate interceptor instances, one for each RPC type,
+    because gRPC aio channels only register multi-inheritance interceptors
+    for the first matching type.
+
+    Returns:
+        List of interceptor instances for unary-unary, unary-stream,
+        stream-unary, and stream-stream RPCs.
+    """
+    return [
+        _AsyncUnaryUnaryAuthInterceptor(),
+        _AsyncUnaryStreamAuthInterceptor(),
+        _AsyncStreamUnaryAuthInterceptor(),
+        _AsyncStreamStreamAuthInterceptor(),
+    ]

python/ray/_private/authentication/grpc_authentication_server_interceptor.py

@@ -0,0 +1,232 @@
+"""gRPC server interceptor for token-based authentication."""
+
+import logging
+from typing import Awaitable, Callable
+
+import grpc
+from grpc import aio as aiogrpc
+
+from ray._private.authentication.authentication_constants import (
+    AUTHORIZATION_HEADER_NAME,
+)
+from ray._private.authentication.authentication_utils import (
+    is_token_auth_enabled,
+    validate_request_token,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def _authenticate_request(metadata: tuple) -> bool:
+    """Authenticate incoming request. currently only supports token authentication.
+
+    Args:
+        metadata: gRPC metadata tuple of (key, value) pairs
+    Returns:
+        True if authentication succeeds or is not required, False otherwise
+    """
+    if not is_token_auth_enabled():
+        return True
+
+    # Extract authorization header from metadata
+    auth_header = None
+    for key, value in metadata:
+        if key.lower() == AUTHORIZATION_HEADER_NAME:
+            auth_header = value
+            break
+
+    if not auth_header:
+        logger.warning("Authentication required but no authorization header provided")
+        return False
+
+    # Validate the token format and value
+    # validate_request_token returns bool (True if valid, False otherwise)
+    return validate_request_token(auth_header)
+
+
+class AsyncAuthenticationServerInterceptor(aiogrpc.ServerInterceptor):
+    """Async gRPC server interceptor that validates authentication tokens.
+
+    This interceptor checks the "authorization" metadata header for a valid
+    Bearer token when token authentication is enabled via RAY_AUTH_MODE=token.
+
+    If the token is missing or invalid, the request is rejected with UNAUTHENTICATED status.
+    """
+
+    async def intercept_service(
+        self,
+        continuation: Callable[
+            [grpc.HandlerCallDetails], Awaitable[grpc.RpcMethodHandler]
+        ],
+        handler_call_details: grpc.HandlerCallDetails,
+    ) -> grpc.RpcMethodHandler:
+        """Intercept service calls to validate authentication.
+
+        This method is called once per RPC to get the handler. We wrap the handler
+        to validate authentication before executing the actual RPC method.
+        """
+        # Get the actual handler
+        handler = await continuation(handler_call_details)
+
+        if handler is None:
+            return None
+
+        async def _abort_if_unauthenticated(context):
+            """Abort the RPC if authentication fails."""
+            if not _authenticate_request(context.invocation_metadata()):
+                await context.abort(
+                    grpc.StatusCode.UNAUTHENTICATED,
+                    "Invalid or missing authentication token",
+                )
+
+        # Wrap the RPC behavior with authentication check
+        def wrap_unary_response(behavior):
+            """Wrap a unary response RPC method to validate authentication first."""
+            if behavior is None:
+                return None
+
+            async def wrapped(request_or_iterator, context):
+                await _abort_if_unauthenticated(context)
+                return await behavior(request_or_iterator, context)
+
+            return wrapped
+
+        def wrap_stream_response(behavior):
+            """Wrap a streaming response RPC method to validate authentication first."""
+            if behavior is None:
+                return None
+
+            async def wrapped(request_or_iterator, context):
+                await _abort_if_unauthenticated(context)
+                async for response in behavior(request_or_iterator, context):
+                    yield response
+
+            return wrapped
+
+        # Create a wrapper class that implements RpcMethodHandler interface
+        class AuthenticatedHandler:
+            """Wrapper handler that validates authentication."""
+
+            def __init__(
+                self, original_handler, unary_wrapper_func, stream_wrapper_func
+            ):
+                self._original = original_handler
+                self._wrap_unary = unary_wrapper_func
+                self._wrap_stream = stream_wrapper_func
+
+            @property
+            def request_streaming(self):
+                return self._original.request_streaming
+
+            @property
+            def response_streaming(self):
+                return self._original.response_streaming
+
+            @property
+            def request_deserializer(self):
+                return self._original.request_deserializer
+
+            @property
+            def response_serializer(self):
+                return self._original.response_serializer
+
+            @property
+            def unary_unary(self):
+                return self._wrap_unary(self._original.unary_unary)
+
+            @property
+            def unary_stream(self):
+                return self._wrap_stream(self._original.unary_stream)
+
+            @property
+            def stream_unary(self):
+                return self._wrap_unary(self._original.stream_unary)
+
+            @property
+            def stream_stream(self):
+                return self._wrap_stream(self._original.stream_stream)
+
+        return AuthenticatedHandler(handler, wrap_unary_response, wrap_stream_response)
+
+
+class SyncAuthenticationServerInterceptor(grpc.ServerInterceptor):
+    """Synchronous gRPC server interceptor that validates authentication tokens.
+
+    This interceptor checks the "authorization" metadata header for a valid
+    Bearer token when token authentication is enabled via RAY_AUTH_MODE=token.
+    If the token is missing or invalid, the request is rejected with UNAUTHENTICATED status.
+    """
+
+    def intercept_service(
+        self,
+        continuation: Callable[[grpc.HandlerCallDetails], grpc.RpcMethodHandler],
+        handler_call_details: grpc.HandlerCallDetails,
+    ) -> grpc.RpcMethodHandler:
+        """Intercept service calls to validate authentication.
+
+        This method is called once per RPC to get the handler. We wrap the handler
+        to validate authentication before executing the actual RPC method.
+        """
+        # Get the actual handler
+        handler = continuation(handler_call_details)
+
+        if handler is None:
+            return None
+
+        # Wrap the RPC behavior with authentication check
+        def wrap_rpc_behavior(behavior):
+            """Wrap an RPC method to validate authentication first."""
+            if behavior is None:
+                return None
+
+            def wrapped(request_or_iterator, context):
+                if not _authenticate_request(context.invocation_metadata()):
+                    context.abort(
+                        grpc.StatusCode.UNAUTHENTICATED,
+                        "Invalid or missing authentication token",
+                    )
+                return behavior(request_or_iterator, context)
+
+            return wrapped
+
+        # Create a wrapper class that implements RpcMethodHandler interface
+        class AuthenticatedHandler:
+            """Wrapper handler that validates authentication."""
+
+            def __init__(self, original_handler, wrapper_func):
+                self._original = original_handler
+                self._wrap = wrapper_func
+
+            @property
+            def request_streaming(self):
+                return self._original.request_streaming
+
+            @property
+            def response_streaming(self):
+                return self._original.response_streaming
+
+            @property
+            def request_deserializer(self):
+                return self._original.request_deserializer
+
+            @property
+            def response_serializer(self):
+                return self._original.response_serializer
+
+            @property
+            def unary_unary(self):
+                return self._wrap(self._original.unary_unary)
+
+            @property
+            def unary_stream(self):
+                return self._wrap(self._original.unary_stream)
+
+            @property
+            def stream_unary(self):
+                return self._wrap(self._original.stream_unary)
+
+            @property
+            def stream_stream(self):
+                return self._wrap(self._original.stream_stream)
+
+        return AuthenticatedHandler(handler, wrap_rpc_behavior)

python/ray/_private/authentication/http_token_authentication.py

@@ -0,0 +1,129 @@
+import logging
+from types import ModuleType
+from typing import Dict, List, Optional
+
+from ray._private.authentication import (
+    authentication_constants,
+    authentication_utils as auth_utils,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def get_token_auth_middleware(
+    aiohttp_module: ModuleType,
+    whitelisted_exact_paths: Optional[List[str]] = None,
+    whitelisted_path_prefixes: Optional[List[str]] = None,
+):
+    """Internal helper to create token auth middleware with provided modules.
+
+    Args:
+        aiohttp_module: The aiohttp module to use
+        whitelisted_exact_paths: List of exact paths that don't require authentication
+        whitelisted_path_prefixes: List of path prefixes that don't require authentication
+    Returns:
+        An aiohttp middleware function
+    """
+
+    @aiohttp_module.web.middleware
+    async def token_auth_middleware(request, handler):
+        """Middleware to validate bearer tokens when token authentication is enabled.
+
+        In minimal Ray installations (without ray._raylet), this middleware is a no-op
+        and passes all requests through without authentication.
+        """
+        # No-op if token auth is not enabled or raylet is not available
+        if not auth_utils.is_token_auth_enabled():
+            return await handler(request)
+
+        # skip authentication for whitelisted paths
+        if (whitelisted_exact_paths and request.path in whitelisted_exact_paths) or (
+            whitelisted_path_prefixes
+            and request.path.startswith(tuple(whitelisted_path_prefixes))
+        ):
+            return await handler(request)
+
+        # Try to get authentication token from multiple sources (in priority order):
+        # 1. Standard "Authorization" header (for API clients, SDKs)
+        # 2. Fallback "X-Ray-Authorization" header (for proxies and KubeRay)
+        # 3. Cookie (for web dashboard sessions)
+
+        auth_header = request.headers.get(
+            authentication_constants.AUTHORIZATION_HEADER_NAME, ""
+        )
+
+        if not auth_header:
+            auth_header = request.headers.get(
+                authentication_constants.RAY_AUTHORIZATION_HEADER_NAME, ""
+            )
+
+        if not auth_header:
+            token = request.cookies.get(
+                authentication_constants.AUTHENTICATION_TOKEN_COOKIE_NAME
+            )
+            if token:
+                # Format as Bearer token for validation
+                auth_header = (
+                    authentication_constants.AUTHORIZATION_BEARER_PREFIX + token
+                )
+
+        if not auth_header:
+            return aiohttp_module.web.Response(
+                status=401, text="Unauthorized: Missing authentication token"
+            )
+
+        if not auth_utils.validate_request_token(auth_header):
+            return aiohttp_module.web.Response(
+                status=403, text="Forbidden: Invalid authentication token"
+            )
+
+        return await handler(request)
+
+    return token_auth_middleware
+
+
+def get_auth_headers_if_auth_enabled(user_headers: Dict[str, str]) -> Dict[str, str]:
+
+    if not auth_utils.is_token_auth_enabled():
+        return {}
+
+    from ray._raylet import AuthenticationTokenLoader
+
+    # Check if user provided their own Authorization header (case-insensitive)
+    has_user_auth = any(
+        key.lower() == authentication_constants.AUTHORIZATION_HEADER_NAME
+        for key in user_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 format_authentication_http_error(status: int, body: str) -> Optional[str]:
+    """Return a user-friendly authentication error message, if applicable."""
+
+    if status == 401:
+        return "Authentication required: {body}\n\n{details}".format(
+            body=body,
+            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.TOKEN_INVALID_ERROR_MESSAGE,
+        )
+
+    return None

python/ray/_private/authentication_test_utils.py

@@ -0,0 +1,157 @@
+import os
+import shutil
+import tempfile
+from contextlib import contextmanager
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Dict, Optional
+
+from ray._raylet import AuthenticationTokenLoader, Config
+
+_AUTH_ENV_VARS = ("RAY_AUTH_MODE", "RAY_AUTH_TOKEN", "RAY_AUTH_TOKEN_PATH")
+_DEFAULT_AUTH_TOKEN_RELATIVE_PATH = Path(".ray") / "auth_token"
+
+
+def reset_auth_token_state() -> None:
+    """Reset authentication token and AUTH_MODE ray config."""
+
+    AuthenticationTokenLoader.instance().reset_cache()
+    Config.initialize("")
+
+
+def set_auth_mode(mode: str) -> None:
+    """Set the authentication mode environment variable."""
+
+    os.environ["RAY_AUTH_MODE"] = mode
+
+
+def set_env_auth_token(token: str) -> None:
+    """Configure the authentication token via environment variable."""
+
+    os.environ["RAY_AUTH_TOKEN"] = token
+    os.environ.pop("RAY_AUTH_TOKEN_PATH", None)
+
+
+def set_auth_token_path(token: str, path: Path) -> None:
+    """Write the authentication token to a specific path and point the loader to it."""
+
+    token_path = Path(path)
+    if token is not None:
+        token_path.parent.mkdir(parents=True, exist_ok=True)
+        token_path.write_text(token)
+    os.environ["RAY_AUTH_TOKEN_PATH"] = str(token_path)
+    os.environ.pop("RAY_AUTH_TOKEN", None)
+
+
+def set_default_auth_token(token: str) -> Path:
+    """Write the authentication token to the default ~/.ray/auth_token location."""
+
+    default_path = Path.home() / _DEFAULT_AUTH_TOKEN_RELATIVE_PATH
+    default_path.parent.mkdir(parents=True, exist_ok=True)
+    default_path.write_text(token)
+    return default_path
+
+
+def clear_auth_token_sources(remove_default: bool = False) -> None:
+    """Clear authentication-related environment variables and optional default token file."""
+
+    for var in ("RAY_AUTH_TOKEN", "RAY_AUTH_TOKEN_PATH"):
+        os.environ.pop(var, None)
+
+    if remove_default:
+        default_path = Path.home() / _DEFAULT_AUTH_TOKEN_RELATIVE_PATH
+        default_path.unlink(missing_ok=True)
+
+
+@dataclass
+class AuthenticationEnvSnapshot:
+    original_env: Dict[str, Optional[str]]
+    original_home: Optional[str]
+    home_was_set: bool
+    temp_home: Optional[Path]
+    default_token_path: Path
+    default_token_exists: bool
+    default_token_contents: Optional[str]
+
+    @classmethod
+    def capture(cls) -> "AuthenticationEnvSnapshot":
+        """Capture current authentication-related environment state."""
+
+        original_env = {var: os.environ.get(var) for var in _AUTH_ENV_VARS}
+        home_was_set = "HOME" in os.environ
+        original_home = os.environ.get("HOME")
+        temp_home: Optional[Path] = None
+
+        if not home_was_set:
+            # in CI $HOME may not be set which can cause issues with tests related to default auth token file.
+            test_tmpdir = os.environ.get("TEST_TMPDIR")
+            base_dir = Path(test_tmpdir) if test_tmpdir else Path(tempfile.gettempdir())
+            temp_home = base_dir / "ray_test_home"
+            temp_home.mkdir(parents=True, exist_ok=True)
+            os.environ["HOME"] = str(temp_home)
+
+        default_token_path = Path.home() / _DEFAULT_AUTH_TOKEN_RELATIVE_PATH
+        default_token_exists = default_token_path.exists()
+        default_token_contents = (
+            default_token_path.read_text() if default_token_exists else None
+        )
+
+        return cls(
+            original_env=original_env,
+            original_home=original_home,
+            home_was_set=home_was_set,
+            temp_home=temp_home,
+            default_token_path=default_token_path,
+            default_token_exists=default_token_exists,
+            default_token_contents=default_token_contents,
+        )
+
+    def clear_default_token(self) -> None:
+        """Remove the default token file for the current HOME."""
+
+        self.default_token_path.unlink(missing_ok=True)
+
+    def restore(self) -> None:
+        """Restore the captured environment, HOME, and default token file state."""
+        # delete any custom token files that may have been created during the test
+        custom_token_path = os.environ.get("RAY_AUTH_TOKEN_PATH")
+        if custom_token_path is not None:
+            custom_token_path = Path(custom_token_path)
+            if custom_token_path.exists():
+                custom_token_path.unlink(missing_ok=True)
+
+        for var, value in self.original_env.items():
+            if value is None:
+                os.environ.pop(var, None)
+            else:
+                os.environ[var] = value
+
+        if self.home_was_set:
+            if self.original_home is None:
+                os.environ.pop("HOME", None)
+            else:
+                os.environ["HOME"] = self.original_home
+
+        if self.default_token_exists:
+            self.default_token_path.parent.mkdir(parents=True, exist_ok=True)
+            self.default_token_path.write_text(self.default_token_contents or "")
+        else:
+            self.default_token_path.unlink(missing_ok=True)
+
+        if not self.home_was_set:
+            current_home = os.environ.get("HOME")
+            if self.temp_home is not None and current_home == str(self.temp_home):
+                os.environ.pop("HOME", None)
+            if self.temp_home is not None and self.temp_home.exists():
+                shutil.rmtree(self.temp_home, ignore_errors=True)
+
+
+@contextmanager
+def authentication_env_guard():
+    """Context manager that restores authentication environment state on exit."""
+
+    snapshot = AuthenticationEnvSnapshot.capture()
+    try:
+        yield snapshot
+    finally:
+        snapshot.restore()

python/ray/_private/auto_init_hook.py

@@ -0,0 +1,32 @@
+import os
+import threading
+from functools import wraps
+
+import ray
+
+auto_init_lock = threading.Lock()
+enable_auto_connect = os.environ.get("RAY_ENABLE_AUTO_CONNECT", "") != "0"
+
+
+def auto_init_ray():
+    if enable_auto_connect and not ray.is_initialized():
+        with auto_init_lock:
+            if not ray.is_initialized():
+                ray.init()
+
+
+def wrap_auto_init(fn):
+    @wraps(fn)
+    def auto_init_wrapper(*args, **kwargs):
+        auto_init_ray()
+        return fn(*args, **kwargs)
+
+    return auto_init_wrapper
+
+
+def wrap_auto_init_for_all_apis(api_names):
+    """Wrap public APIs with automatic ray.init."""
+    for api_name in api_names:
+        api = getattr(ray, api_name, None)
+        assert api is not None, api_name
+        setattr(ray, api_name, wrap_auto_init(api))

python/ray/_private/client_mode_hook.py

@@ -0,0 +1,221 @@
+import os
+import threading
+from contextlib import contextmanager
+from functools import wraps
+from typing import Any, Callable, Optional, Tuple, TypeVar, cast, overload
+
+from ray._private.auto_init_hook import auto_init_ray
+
+F = TypeVar("F", bound=Callable[..., Any])
+
+# Attr set on func defs to mark they have been converted to client mode.
+RAY_CLIENT_MODE_ATTR = "__ray_client_mode_key__"
+
+# Global setting of whether client mode is enabled. This default to OFF,
+# but is enabled upon ray.client(...).connect() or in tests.
+is_client_mode_enabled = os.environ.get("RAY_CLIENT_MODE", "0") == "1"
+
+# When RAY_CLIENT_MODE == 1, we treat it as default enabled client mode
+# This is useful for testing
+is_client_mode_enabled_by_default = is_client_mode_enabled
+os.environ.update({"RAY_CLIENT_MODE": "0"})
+
+is_init_called = False
+
+# Local setting of whether to ignore client hook conversion. This defaults
+# to TRUE and is disabled when the underlying 'real' Ray function is needed.
+_client_hook_status_on_thread = threading.local()
+_client_hook_status_on_thread.status = True
+
+
+def _get_client_hook_status_on_thread():
+    """Get's the value of `_client_hook_status_on_thread`.
+    Since `_client_hook_status_on_thread` is a thread-local variable, we may
+    need to add and set the 'status' attribute.
+    """
+    global _client_hook_status_on_thread
+    if not hasattr(_client_hook_status_on_thread, "status"):
+        _client_hook_status_on_thread.status = True
+    return _client_hook_status_on_thread.status
+
+
+def _set_client_hook_status(val: bool):
+    global _client_hook_status_on_thread
+    _client_hook_status_on_thread.status = val
+
+
+def _disable_client_hook():
+    global _client_hook_status_on_thread
+    out = _get_client_hook_status_on_thread()
+    _client_hook_status_on_thread.status = False
+    return out
+
+
+def _explicitly_enable_client_mode():
+    """Force client mode to be enabled.
+    NOTE: This should not be used in tests, use `enable_client_mode`.
+    """
+    global is_client_mode_enabled
+    is_client_mode_enabled = True
+
+
+def _explicitly_disable_client_mode():
+    global is_client_mode_enabled
+    is_client_mode_enabled = False
+
+
+@contextmanager
+def disable_client_hook():
+    val = _disable_client_hook()
+    try:
+        yield None
+    finally:
+        _set_client_hook_status(val)
+
+
+@contextmanager
+def enable_client_mode():
+    _explicitly_enable_client_mode()
+    try:
+        yield None
+    finally:
+        _explicitly_disable_client_mode()
+
+
+@overload
+def client_mode_hook(func: F) -> F:
+    ...
+
+
+@overload
+def client_mode_hook(*, local_only_kwargs: Tuple[str, ...] = ()) -> Callable[[F], F]:
+    ...
+
+
+def client_mode_hook(
+    func: Optional[F] = None, *, local_only_kwargs: Tuple[str, ...] = ()
+):
+    """Decorator for whether to use the 'regular' ray version of a function,
+    or the Ray Client version of that function.
+
+    Args:
+        func: This function. This is set when this function is used
+            as a bare decorator.
+        local_only_kwargs: Names of keyword arguments that apply only to the
+            local (non-client) implementation. They are stripped before the
+            call is redirected to the Ray Client, so the client API does not
+            need to accept them. On the local path they pass through unchanged.
+            as a decorator.
+
+    Returns:
+        The wrapped function that dispatches to the regular or client version.
+    """
+
+    def decorator(func: F) -> F:
+        from ray.util.client import ray
+
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            # NOTE(hchen): DO NOT use "import" inside this function.
+            # Because when it's called within a `__del__` method, this error
+            # will be raised (see #35114):
+            # ImportError: sys.meta_path is None, Python is likely shutting down.
+            if client_mode_should_convert():
+                # Legacy code
+                # we only convert init function if RAY_CLIENT_MODE=1
+                if func.__name__ != "init" or is_client_mode_enabled_by_default:
+                    if local_only_kwargs:
+                        kwargs = {
+                            k: v
+                            for k, v in kwargs.items()
+                            if k not in local_only_kwargs
+                        }
+                    return getattr(ray, func.__name__)(*args, **kwargs)
+            return func(*args, **kwargs)
+
+        return cast(F, wrapper)
+
+    # Support both `@client_mode_hook` and
+    # `@client_mode_hook(local_only_kwargs=...)` usage.
+    if func is not None:
+        return decorator(func)
+    return decorator
+
+
+def client_mode_should_convert():
+    """Determines if functions should be converted to client mode."""
+
+    # `is_client_mode_enabled_by_default` is used for testing with
+    # `RAY_CLIENT_MODE=1`. This flag means all tests run with client mode.
+    return (
+        is_client_mode_enabled or is_client_mode_enabled_by_default
+    ) and _get_client_hook_status_on_thread()
+
+
+def client_mode_wrap(func):
+    """Wraps a function called during client mode for execution as a remote
+    task.
+
+    Can be used to implement public features of ray client which do not
+    belong in the main ray API (`ray.*`), yet require server-side execution.
+    An example is the creation of placement groups:
+    `ray.util.placement_group.placement_group()`. When called on the client
+    side, this function is wrapped in a task to facilitate interaction with
+    the GCS.
+    """
+
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        from ray.util.client import ray
+
+        auto_init_ray()
+        # Directly pass this through since `client_mode_wrap` is for
+        # Placement Group APIs
+        if client_mode_should_convert():
+            f = ray.remote(num_cpus=0)(func)
+            ref = f.remote(*args, **kwargs)
+            return ray.get(ref)
+        return func(*args, **kwargs)
+
+    return wrapper
+
+
+def client_mode_convert_function(func_cls, in_args, in_kwargs, **kwargs):
+    """Runs a preregistered ray RemoteFunction through the ray client.
+
+    The common case for this is to transparently convert that RemoteFunction
+    to a ClientRemoteFunction. This happens in circumstances where the
+    RemoteFunction is declared early, in a library and only then is Ray used in
+    client mode -- necessitating a conversion.
+    """
+    from ray.util.client import ray
+
+    key = getattr(func_cls, RAY_CLIENT_MODE_ATTR, None)
+
+    # Second part of "or" is needed in case func_cls is reused between Ray
+    # client sessions in one Python interpreter session.
+    if (key is None) or (not ray._converted_key_exists(key)):
+        key = ray._convert_function(func_cls)
+        setattr(func_cls, RAY_CLIENT_MODE_ATTR, key)
+    client_func = ray._get_converted(key)
+    return client_func._remote(in_args, in_kwargs, **kwargs)
+
+
+def client_mode_convert_actor(actor_cls, in_args, in_kwargs, **kwargs):
+    """Runs a preregistered actor class on the ray client
+
+    The common case for this decorator is for instantiating an ActorClass
+    transparently as a ClientActorClass. This happens in circumstances where
+    the ActorClass is declared early, in a library and only then is Ray used in
+    client mode -- necessitating a conversion.
+    """
+    from ray.util.client import ray
+
+    key = getattr(actor_cls, RAY_CLIENT_MODE_ATTR, None)
+    # Second part of "or" is needed in case actor_cls is reused between Ray
+    # client sessions in one Python interpreter session.
+    if (key is None) or (not ray._converted_key_exists(key)):
+        key = ray._convert_actor(actor_cls)
+        setattr(actor_cls, RAY_CLIENT_MODE_ATTR, key)
+    client_actor = ray._get_converted(key)
+    return client_actor._remote(in_args, in_kwargs, **kwargs)

python/ray/_private/collections_utils.py

@@ -0,0 +1,10 @@
+from typing import Any, List
+
+
+def split(items: List[Any], chunk_size: int):
+    """Splits provided list into chunks of given size"""
+
+    assert chunk_size > 0, "Chunk size has to be > 0"
+
+    for i in range(0, len(items), chunk_size):
+        yield items[i : i + chunk_size]

python/ray/_private/compat.py

@@ -0,0 +1,40 @@
+import io
+import platform
+
+
+def patch_psutil():
+    """WSL's /proc/meminfo has an inconsistency where it
+    nondeterministically omits a space after colons (after "SwapFree:"
+    in my case).
+    psutil then splits on spaces and then parses the wrong field,
+    crashing on the 'int(fields[1])' expression in
+    psutil._pslinux.virtual_memory().
+    Workaround: We ensure there is a space following each colon.
+    """
+    assert (
+        platform.system() == "Linux"
+        and "Microsoft".lower() in platform.release().lower()
+    )
+
+    try:
+        import psutil._pslinux
+    except ImportError:
+        psutil = None
+    psutil_open_binary = None
+    if psutil:
+        try:
+            psutil_open_binary = psutil._pslinux.open_binary
+        except AttributeError:
+            pass
+    # Only patch it if it doesn't seem to have been patched already
+    if psutil_open_binary and psutil_open_binary.__name__ == "open_binary":
+
+        def psutil_open_binary_patched(fname, *args, **kwargs):
+            f = psutil_open_binary(fname, *args, **kwargs)
+            if fname == "/proc/meminfo":
+                with f:
+                    # Make sure there's a space after colons
+                    return io.BytesIO(f.read().replace(b":", b": "))
+            return f
+
+        psutil._pslinux.open_binary = psutil_open_binary_patched

python/ray/_private/conftest_utils.py

@@ -0,0 +1,15 @@
+import pytest
+
+import ray._private.ray_constants as ray_constants
+
+
+@pytest.fixture
+def set_override_dashboard_url(monkeypatch, request):
+    override_url = getattr(request, "param", "https://external_dashboard_url")
+    with monkeypatch.context() as m:
+        if override_url:
+            m.setenv(
+                ray_constants.RAY_OVERRIDE_DASHBOARD_URL,
+                override_url,
+            )
+        yield

python/ray/_private/custom_types.py

@@ -0,0 +1,154 @@
+from typing import Literal
+
+from ray.core.generated.common_pb2 import (
+    ErrorType,
+    Language,
+    TaskStatus,
+    TaskType,
+    WorkerExitType,
+    WorkerType,
+)
+from ray.core.generated.gcs_pb2 import (
+    ActorTableData,
+    GcsNodeInfo,
+    PlacementGroupTableData,
+)
+
+ACTOR_STATUS = [
+    "DEPENDENCIES_UNREADY",
+    "PENDING_CREATION",
+    "ALIVE",
+    "RESTARTING",
+    "DEAD",
+]
+TypeActorStatus = Literal[tuple(ACTOR_STATUS)]
+PLACEMENT_GROUP_STATUS = [
+    "PENDING",
+    "PREPARED",
+    "CREATED",
+    "REMOVED",
+    "RESCHEDULING",
+]
+TypePlacementGroupStatus = Literal[tuple(PLACEMENT_GROUP_STATUS)]
+TASK_STATUS = [
+    "NIL",
+    "PENDING_ARGS_AVAIL",
+    "PENDING_NODE_ASSIGNMENT",
+    "PENDING_OBJ_STORE_MEM_AVAIL",
+    "PENDING_ARGS_FETCH",
+    "SUBMITTED_TO_WORKER",
+    "PENDING_ACTOR_TASK_ARGS_FETCH",
+    "PENDING_ACTOR_TASK_ORDERING_OR_CONCURRENCY",
+    "RUNNING",
+    "RUNNING_IN_RAY_GET",
+    "RUNNING_IN_RAY_WAIT",
+    "FINISHED",
+    "FAILED",
+    "GETTING_AND_PINNING_ARGS",
+]
+TypeTaskStatus = Literal[tuple(TASK_STATUS)]
+NODE_STATUS = ["ALIVE", "DEAD"]
+TypeNodeStatus = Literal[tuple(NODE_STATUS)]
+WORKER_TYPE = [
+    "WORKER",
+    "DRIVER",
+    "SPILL_WORKER",
+    "RESTORE_WORKER",
+]
+TypeWorkerType = Literal[tuple(WORKER_TYPE)]
+WORKER_EXIT_TYPE = [
+    "SYSTEM_ERROR",
+    "INTENDED_SYSTEM_EXIT",
+    "USER_ERROR",
+    "INTENDED_USER_EXIT",
+    "NODE_OUT_OF_MEMORY",
+]
+TypeWorkerExitType = Literal[tuple(WORKER_EXIT_TYPE)]
+TASK_TYPE = [
+    "NORMAL_TASK",
+    "ACTOR_CREATION_TASK",
+    "ACTOR_TASK",
+    "DRIVER_TASK",
+]
+TypeTaskType = Literal[tuple(TASK_TYPE)]
+# TODO(kevin85421): `class ReferenceType(Enum)` is defined in
+# `dashboard/memory_utils.py` to avoid complex dependencies. I redefined
+# it here. Eventually, we should remove the one in `dashboard/memory_utils.py`
+# and define it under `ray/_private`.
+REFERENCE_TYPE = [
+    "ACTOR_HANDLE",
+    "PINNED_IN_MEMORY",
+    "LOCAL_REFERENCE",
+    "USED_BY_PENDING_TASK",
+    "CAPTURED_IN_OBJECT",
+    "UNKNOWN_STATUS",
+]
+TypeReferenceType = Literal[tuple(REFERENCE_TYPE)]
+# The ErrorType enum is used in the export API so it is public
+# and any modifications must be backward compatible.
+ERROR_TYPE = [
+    "WORKER_DIED",
+    "ACTOR_DIED",
+    "TASK_EXECUTION_EXCEPTION",
+    "OBJECT_IN_PLASMA",
+    "TASK_CANCELLED",
+    "ACTOR_CREATION_FAILED",
+    "RUNTIME_ENV_SETUP_FAILED",
+    "OBJECT_LOST",
+    "OWNER_DIED",
+    "OBJECT_DELETED",
+    "DEPENDENCY_RESOLUTION_FAILED",
+    "OBJECT_UNRECONSTRUCTABLE_MAX_ATTEMPTS_EXCEEDED",
+    "OBJECT_UNRECONSTRUCTABLE_LINEAGE_EVICTED",
+    "OBJECT_FETCH_TIMED_OUT",
+    "LOCAL_RAYLET_DIED",
+    "TASK_PLACEMENT_GROUP_REMOVED",
+    "ACTOR_PLACEMENT_GROUP_REMOVED",
+    "TASK_UNSCHEDULABLE_ERROR",
+    "ACTOR_UNSCHEDULABLE_ERROR",
+    "OUT_OF_DISK_ERROR",
+    "OBJECT_FREED",
+    "OUT_OF_MEMORY",
+    "NODE_DIED",
+    "END_OF_STREAMING_GENERATOR",
+    "ACTOR_UNAVAILABLE",
+    "GENERATOR_TASK_FAILED_FOR_OBJECT_RECONSTRUCTION",
+    "OBJECT_UNRECONSTRUCTABLE_PUT",
+    "OBJECT_UNRECONSTRUCTABLE_RETRIES_DISABLED",
+    "OBJECT_UNRECONSTRUCTABLE_BORROWED",
+    "OBJECT_UNRECONSTRUCTABLE_REF_NOT_FOUND",
+    "OBJECT_UNRECONSTRUCTABLE_TASK_CANCELLED",
+    "OBJECT_UNRECONSTRUCTABLE_LINEAGE_DISABLED",
+    "WORKER_STARTUP_FAILED",
+]
+# The Language enum is used in the export API so it is public
+# and any modifications must be backward compatible.
+LANGUAGE = ["PYTHON", "JAVA", "CPP"]
+
+
+def validate_protobuf_enum(grpc_enum, custom_enum):
+    """Validate the literal contains the correct enum values from protobuf"""
+    enum_vals = set(grpc_enum.DESCRIPTOR.values_by_name.keys())
+    # Sometimes, the grpc enum is mocked, and it
+    # doesn't include any values in that case.
+    if len(enum_vals) > 0:
+        assert enum_vals == set(
+            custom_enum
+        ), """Literals in `custom_types.py` and `.proto` files are out of sync. \
+Consider building //:install_py_proto with Bazel or updating `custom_types.py`."""
+
+
+# Do the enum validation here.
+# It is necessary to avoid regression. Alternatively, we can auto generate this
+# directly by protobuf.
+validate_protobuf_enum(ActorTableData.ActorState, ACTOR_STATUS)
+validate_protobuf_enum(
+    PlacementGroupTableData.PlacementGroupState, PLACEMENT_GROUP_STATUS
+)
+validate_protobuf_enum(TaskStatus, TASK_STATUS)
+validate_protobuf_enum(GcsNodeInfo.GcsNodeState, NODE_STATUS)
+validate_protobuf_enum(WorkerType, WORKER_TYPE)
+validate_protobuf_enum(WorkerExitType, WORKER_EXIT_TYPE)
+validate_protobuf_enum(TaskType, TASK_TYPE)
+validate_protobuf_enum(ErrorType, ERROR_TYPE)
+validate_protobuf_enum(Language, LANGUAGE)

python/ray/_private/dict.py

@@ -0,0 +1,253 @@
+import copy
+from collections import deque
+from collections.abc import Mapping, Sequence
+from typing import Dict, List, Optional, TypeVar, Union
+
+from ray.util.annotations import Deprecated
+
+T = TypeVar("T")
+
+
+@Deprecated
+def merge_dicts(d1: dict, d2: dict) -> dict:
+    """
+    Args:
+        d1: Dict 1.
+        d2: Dict 2.
+
+    Returns:
+         A new dict that is d1 and d2 deep merged.
+    """
+    merged = copy.deepcopy(d1)
+    deep_update(merged, d2, True, [])
+    return merged
+
+
+@Deprecated
+def deep_update(
+    original: dict,
+    new_dict: dict,
+    new_keys_allowed: bool = False,
+    allow_new_subkey_list: Optional[List[str]] = None,
+    override_all_if_type_changes: Optional[List[str]] = None,
+    override_all_key_list: Optional[List[str]] = None,
+) -> dict:
+    """Updates original dict with values from new_dict recursively.
+
+    If new key is introduced in new_dict, then if new_keys_allowed is not
+    True, an error will be thrown. Further, for sub-dicts, if the key is
+    in the allow_new_subkey_list, then new subkeys can be introduced.
+
+    Args:
+        original: Dictionary with default values.
+        new_dict: Dictionary with values to be updated
+        new_keys_allowed: Whether new keys are allowed.
+        allow_new_subkey_list: List of keys that
+            correspond to dict values where new subkeys can be introduced.
+            This is only at the top level.
+        override_all_if_type_changes: List of top level
+            keys with value=dict, for which we always simply override the
+            entire value (dict), iff the "type" key in that value dict changes.
+        override_all_key_list: List of top level keys
+            for which we override the entire value if the key is in the new_dict.
+
+    Returns:
+        The updated original dict.
+    """
+    allow_new_subkey_list = allow_new_subkey_list or []
+    override_all_if_type_changes = override_all_if_type_changes or []
+    override_all_key_list = override_all_key_list or []
+
+    for k, value in new_dict.items():
+        if k not in original and not new_keys_allowed:
+            raise Exception("Unknown config parameter `{}` ".format(k))
+
+        # Both orginal value and new one are dicts.
+        if (
+            isinstance(original.get(k), dict)
+            and isinstance(value, dict)
+            and k not in override_all_key_list
+        ):
+            # Check old type vs old one. If different, override entire value.
+            if (
+                k in override_all_if_type_changes
+                and "type" in value
+                and "type" in original[k]
+                and value["type"] != original[k]["type"]
+            ):
+                original[k] = value
+            # Allowed key -> ok to add new subkeys.
+            elif k in allow_new_subkey_list:
+                deep_update(
+                    original[k],
+                    value,
+                    True,
+                    override_all_key_list=override_all_key_list,
+                )
+            # Non-allowed key.
+            else:
+                deep_update(
+                    original[k],
+                    value,
+                    new_keys_allowed,
+                    override_all_key_list=override_all_key_list,
+                )
+        # Original value not a dict OR new value not a dict:
+        # Override entire value.
+        else:
+            original[k] = value
+    return original
+
+
+@Deprecated
+def flatten_dict(
+    dt: Dict,
+    delimiter: str = "/",
+    prevent_delimiter: bool = False,
+    flatten_list: bool = False,
+):
+    """Flatten dict.
+
+    Output and input are of the same dict type.
+    Input dict remains the same after the operation.
+    """
+
+    def _raise_delimiter_exception():
+        raise ValueError(
+            f"Found delimiter `{delimiter}` in key when trying to flatten "
+            f"array. Please avoid using the delimiter in your specification."
+        )
+
+    dt = copy.copy(dt)
+    if prevent_delimiter and any(delimiter in key for key in dt):
+        # Raise if delimiter is any of the keys
+        _raise_delimiter_exception()
+
+    while_check = (dict, list) if flatten_list else dict
+
+    while any(isinstance(v, while_check) for v in dt.values()):
+        remove = []
+        add = {}
+        for key, value in dt.items():
+            if isinstance(value, dict):
+                for subkey, v in value.items():
+                    if prevent_delimiter and delimiter in subkey:
+                        # Raise if delimiter is in any of the subkeys
+                        _raise_delimiter_exception()
+
+                    add[delimiter.join([key, str(subkey)])] = v
+                remove.append(key)
+            elif flatten_list and isinstance(value, list):
+                for i, v in enumerate(value):
+                    if prevent_delimiter and delimiter in subkey:
+                        # Raise if delimiter is in any of the subkeys
+                        _raise_delimiter_exception()
+
+                    add[delimiter.join([key, str(i)])] = v
+                remove.append(key)
+
+        dt.update(add)
+        for k in remove:
+            del dt[k]
+    return dt
+
+
+@Deprecated
+def unflatten_dict(dt: Dict[str, T], delimiter: str = "/") -> Dict[str, T]:
+    """Unflatten dict. Does not support unflattening lists."""
+    dict_type = type(dt)
+    out = dict_type()
+    for key, val in dt.items():
+        path = key.split(delimiter)
+        item = out
+        for k in path[:-1]:
+            item = item.setdefault(k, dict_type())
+            if not isinstance(item, dict_type):
+                raise TypeError(
+                    f"Cannot unflatten dict due the key '{key}' "
+                    f"having a parent key '{k}', which value is not "
+                    f"of type {dict_type} (got {type(item)}). "
+                    "Change the key names to resolve the conflict."
+                )
+        item[path[-1]] = val
+    return out
+
+
+@Deprecated
+def unflatten_list_dict(dt: Dict[str, T], delimiter: str = "/") -> Dict[str, T]:
+    """Unflatten nested dict and list.
+
+    This function now has some limitations:
+    (1) The keys of dt must be str.
+    (2) If unflattened dt (the result) contains list, the index order must be
+        ascending when accessing dt. Otherwise, this function will throw
+        AssertionError.
+    (3) The unflattened dt (the result) shouldn't contain dict with number
+        keys.
+
+    Be careful to use this function. If you want to improve this function,
+    please also improve the unit test. See #14487 for more details.
+
+    Args:
+        dt: Flattened dictionary that is originally nested by multiple
+            list and dict.
+        delimiter: Delimiter of keys.
+
+    Returns:
+        The unflattened nested dict/list.
+
+    Example:
+        >>> dt = {"aaa/0/bb": 12, "aaa/1/cc": 56, "aaa/1/dd": 92}
+        >>> unflatten_list_dict(dt)
+        {'aaa': [{'bb': 12}, {'cc': 56, 'dd': 92}]}
+    """
+    out_type = list if list(dt)[0].split(delimiter, 1)[0].isdigit() else type(dt)
+    out = out_type()
+    for key, val in dt.items():
+        path = key.split(delimiter)
+
+        item = out
+        for i, k in enumerate(path[:-1]):
+            next_type = list if path[i + 1].isdigit() else dict
+            if isinstance(item, dict):
+                item = item.setdefault(k, next_type())
+            elif isinstance(item, list):
+                if int(k) >= len(item):
+                    item.append(next_type())
+                    assert int(k) == len(item) - 1
+                item = item[int(k)]
+
+        if isinstance(item, dict):
+            item[path[-1]] = val
+        elif isinstance(item, list):
+            item.append(val)
+            assert int(path[-1]) == len(item) - 1
+    return out
+
+
+@Deprecated
+def unflattened_lookup(
+    flat_key: str, lookup: Union[Mapping, Sequence], delimiter: str = "/", **kwargs
+) -> Union[Mapping, Sequence]:
+    """
+    Unflatten `flat_key` and iteratively look up in `lookup`. E.g.
+    `flat_key="a/0/b"` will try to return `lookup["a"][0]["b"]`.
+    """
+    if flat_key in lookup:
+        return lookup[flat_key]
+    keys = deque(flat_key.split(delimiter))
+    base = lookup
+    while keys:
+        key = keys.popleft()
+        try:
+            if isinstance(base, Mapping):
+                base = base[key]
+            elif isinstance(base, Sequence):
+                base = base[int(key)]
+            else:
+                raise KeyError()
+        except KeyError as e:
+            if "default" in kwargs:
+                return kwargs["default"]
+            raise e
+    return base

python/ray/_private/event/event_logger.py

@@ -0,0 +1,202 @@
+import json
+import logging
+import os
+import pathlib
+import random
+import socket
+import string
+import threading
+from datetime import datetime
+from typing import Dict, Optional
+
+from google.protobuf.json_format import Parse
+
+from ray._private.protobuf_compat import message_to_dict
+from ray.core.generated.event_pb2 import Event
+
+global_logger = logging.getLogger(__name__)
+
+
+def get_event_id():
+    return "".join([random.choice(string.hexdigits) for _ in range(36)])
+
+
+class EventLoggerAdapter:
+    def __init__(self, source: Event.SourceType, logger: logging.Logger):
+        """Adapter for the Python logger that's used to emit events.
+
+        When events are emitted, they are aggregated and available via
+        state API and dashboard.
+
+        This class is thread-safe.
+        """
+        self.logger = logger
+        # Aligned with `event.proto`'s `message Event``
+        self.source = source
+        self.source_hostname = socket.gethostname()
+        self.source_pid = os.getpid()
+
+        # The below fields must be protected by this lock.
+        self.lock = threading.Lock()
+        # {str -> str} typed dict
+        self.global_context = {}
+
+    def set_global_context(self, global_context: Dict[str, str] = None):
+        """Set the global metadata.
+
+        This method overwrites the global metadata if it is called more than once.
+        """
+        with self.lock:
+            self.global_context = {} if not global_context else global_context
+
+    def trace(self, message: str, **kwargs):
+        self._emit(Event.Severity.TRACE, message, **kwargs)
+
+    def debug(self, message: str, **kwargs):
+        self._emit(Event.Severity.DEBUG, message, **kwargs)
+
+    def info(self, message: str, **kwargs):
+        self._emit(Event.Severity.INFO, message, **kwargs)
+
+    def warning(self, message: str, **kwargs):
+        self._emit(Event.Severity.WARNING, message, **kwargs)
+
+    def error(self, message: str, **kwargs):
+        self._emit(Event.Severity.ERROR, message, **kwargs)
+
+    def fatal(self, message: str, **kwargs):
+        self._emit(Event.Severity.FATAL, message, **kwargs)
+
+    def _emit(self, severity: Event.Severity, message: str, **kwargs):
+        # NOTE: Python logger is thread-safe,
+        # so we don't need to protect it using locks.
+        event = Event()
+        event.event_id = get_event_id()
+        event.timestamp = int(datetime.now().timestamp())
+        event.message = message
+        event.severity = severity
+        # TODO(sang): Support event type & schema.
+        event.label = ""
+        event.source_type = self.source
+        event.source_hostname = self.source_hostname
+        event.source_pid = self.source_pid
+        custom_fields = event.custom_fields
+        with self.lock:
+            for k, v in self.global_context.items():
+                if v is not None and k is not None:
+                    custom_fields[k] = v
+        for k, v in kwargs.items():
+            if v is not None and k is not None:
+                custom_fields[k] = v
+
+        self.logger.info(
+            json.dumps(
+                message_to_dict(
+                    event,
+                    always_print_fields_with_no_presence=True,
+                    preserving_proto_field_name=True,
+                    use_integers_for_enums=False,
+                )
+            )
+        )
+
+        # Force flush all handlers so that we won't lose events.
+        for handler in self.logger.handlers[:]:
+            try:
+                handler.flush()
+            except Exception:
+                global_logger.exception("Failed to flush event logger handler.")
+
+
+def _build_event_file_logger(source: Event.SourceType, sink_dir: str):
+    logger = logging.getLogger("_ray_event_logger")
+    logger.setLevel(logging.INFO)
+    dir_path = pathlib.Path(sink_dir) / "events"
+    filepath = dir_path / f"event_{source}.log"
+    dir_path.mkdir(exist_ok=True)
+    filepath.touch(exist_ok=True)
+    # Configure the logger.
+    handler = logging.FileHandler(filepath)
+    formatter = logging.Formatter("%(message)s")
+    handler.setFormatter(formatter)
+    logger.addHandler(handler)
+    logger.propagate = False
+    return logger
+
+
+# This lock must be used when accessing or updating global event logger dict.
+_event_logger_lock = threading.Lock()
+_event_logger = {}
+
+
+def get_event_logger(source: Event.SourceType, sink_dir: str):
+    """Get the event logger of the current process.
+
+    There's only 1 event logger per (process, source).
+
+    TODO(sang): Support more impl than file-based logging.
+                Currently, the interface also ties to the
+                file-based logging impl.
+
+    Args:
+        source: The source of the event.
+        sink_dir: The directory to sink event logs.
+
+    Returns:
+        The event logger adapter for the given source.
+    """
+    with _event_logger_lock:
+        global _event_logger
+        source_name = Event.SourceType.Name(source)
+        if source_name not in _event_logger:
+            logger = _build_event_file_logger(source_name, sink_dir)
+            _event_logger[source_name] = EventLoggerAdapter(source, logger)
+
+        return _event_logger[source_name]
+
+
+def parse_event(event_str: str) -> Optional[Event]:
+    """Parse an event from a string.
+
+    Args:
+        event_str: The string to parse. Expect to be a JSON serialized
+            Event protobuf.
+
+    Returns:
+        The parsed event if parsable, else None
+    """
+    try:
+        return Parse(event_str, Event())
+    except Exception:
+        global_logger.exception(f"Failed to parse event: {event_str}")
+        return None
+
+
+def filter_event_by_level(event: Event, filter_event_level: str) -> bool:
+    """Filter an event based on event level.
+
+    Args:
+        event: The event to filter.
+        filter_event_level: The event level string to filter by. Any events
+            that are lower than this level will be filtered.
+
+    Returns:
+        True if the event should be filtered, else False.
+    """
+
+    event_levels = {
+        Event.Severity.TRACE: 0,
+        Event.Severity.DEBUG: 1,
+        Event.Severity.INFO: 2,
+        Event.Severity.WARNING: 3,
+        Event.Severity.ERROR: 4,
+        Event.Severity.FATAL: 5,
+    }
+
+    filter_event_level = filter_event_level.upper()
+    filter_event_level = Event.Severity.Value(filter_event_level)
+
+    if event_levels[event.severity] < event_levels[filter_event_level]:
+        return True
+
+    return False

python/ray/_private/event/export_event_logger.py

@@ -0,0 +1,261 @@
+import json
+import logging
+import pathlib
+import random
+import string
+import threading
+from datetime import datetime
+from enum import Enum
+from typing import Union
+
+from ray._private import ray_constants
+from ray._private.protobuf_compat import message_to_dict
+from ray.core.generated.export_dataset_metadata_pb2 import (
+    ExportDatasetMetadata,
+)
+from ray.core.generated.export_dataset_operator_event_pb2 import (
+    ExportDatasetOperatorEventData,
+)
+from ray.core.generated.export_dataset_operator_schema_pb2 import (
+    ExportDatasetOperatorSchema,
+)
+from ray.core.generated.export_event_pb2 import ExportEvent
+from ray.core.generated.export_submission_job_event_pb2 import (
+    ExportSubmissionJobEventData,
+)
+from ray.core.generated.export_train_state_pb2 import (
+    ExportTrainRunAttemptEventData,
+    ExportTrainRunEventData,
+)
+
+global_logger = logging.getLogger(__name__)
+
+# This contains the union of export event data types which emit events
+# using the python ExportEventLoggerAdapter
+ExportEventDataType = Union[
+    ExportSubmissionJobEventData,
+    ExportTrainRunEventData,
+    ExportTrainRunAttemptEventData,
+    ExportDatasetMetadata,
+    ExportDatasetOperatorEventData,
+    ExportDatasetOperatorSchema,
+]
+
+
+class EventLogType(Enum):
+    """Enum class representing different types of export event logs.
+
+    Each enum value contains a log type name and a set of supported event data types.
+
+    Attributes:
+        TRAIN_STATE: Export events related to training state, supporting train run and attempt events.
+        SUBMISSION_JOB: Export events related to job submissions.
+        DATASET_METADATA: Export events related to dataset metadata.
+        DATASET_OPERATOR_EVENT: Export events related to Ray Data operator.
+        DATASET_OPERATOR_SCHEMA: Export schema related to Ray Data operator.
+    """
+
+    TRAIN_STATE = (
+        "EXPORT_TRAIN_STATE",
+        {ExportTrainRunEventData, ExportTrainRunAttemptEventData},
+    )
+    SUBMISSION_JOB = ("EXPORT_SUBMISSION_JOB", {ExportSubmissionJobEventData})
+    DATASET_METADATA = ("EXPORT_DATASET_METADATA", {ExportDatasetMetadata})
+    DATASET_OPERATOR_EVENT = (
+        "EXPORT_DATASET_OPERATOR_EVENT",
+        {ExportDatasetOperatorEventData},
+    )
+    DATASET_OPERATOR_SCHEMA = (
+        "EXPORT_DATASET_OPERATOR_SCHEMA",
+        {ExportDatasetOperatorSchema},
+    )
+
+    def __init__(self, log_type_name: str, event_types: set[ExportEventDataType]):
+        """Initialize an EventLogType enum value.
+
+        Args:
+            log_type_name: String identifier for the log type. This name is used to construct the log file name.
+                See `_build_export_event_file_logger` for more details.
+            event_types: Set of event data types that this log type supports.
+        """
+        self.log_type_name = log_type_name
+        self.event_types = event_types
+
+    def supports_event_type(self, event_type: ExportEventDataType) -> bool:
+        """Check if this log type supports the given event data type.
+
+        Args:
+            event_type: The event data type to check for support.
+
+        Returns:
+            bool: True if the event type is supported, False otherwise.
+        """
+        return type(event_type) in self.event_types
+
+
+def generate_event_id():
+    return "".join([random.choice(string.hexdigits) for _ in range(18)])
+
+
+class ExportEventLoggerAdapter:
+    def __init__(self, log_type: EventLogType, logger: logging.Logger):
+        """Adapter for the Python logger that's used to emit export events."""
+        self.logger = logger
+        self.log_type = log_type
+
+    def send_event(self, event_data: ExportEventDataType):
+        # NOTE: Python logger is thread-safe,
+        # so we don't need to protect it using locks.
+        try:
+            event = self._create_export_event(event_data)
+        except TypeError:
+            global_logger.exception(
+                "Failed to create ExportEvent from event_data so no "
+                "event will be written to file."
+            )
+            return
+
+        event_as_str = self._export_event_to_string(event)
+
+        self.logger.info(event_as_str)
+        # Force flush all handlers so that we won't lose events.
+        for handler in self.logger.handlers[:]:
+            try:
+                handler.flush()
+            except Exception:
+                global_logger.exception("Failed to flush export event logger handler.")
+
+    def _create_export_event(self, event_data: ExportEventDataType) -> ExportEvent:
+        event = ExportEvent()
+        event.event_id = generate_event_id()
+        event.timestamp = int(datetime.now().timestamp())
+        if isinstance(event_data, ExportSubmissionJobEventData):
+            event.submission_job_event_data.CopyFrom(event_data)
+            event.source_type = ExportEvent.SourceType.EXPORT_SUBMISSION_JOB
+        elif isinstance(event_data, ExportTrainRunEventData):
+            event.train_run_event_data.CopyFrom(event_data)
+            event.source_type = ExportEvent.SourceType.EXPORT_TRAIN_RUN
+        elif isinstance(event_data, ExportTrainRunAttemptEventData):
+            event.train_run_attempt_event_data.CopyFrom(event_data)
+            event.source_type = ExportEvent.SourceType.EXPORT_TRAIN_RUN_ATTEMPT
+        elif isinstance(event_data, ExportDatasetMetadata):
+            event.dataset_metadata.CopyFrom(event_data)
+            event.source_type = ExportEvent.SourceType.EXPORT_DATASET_METADATA
+        elif isinstance(event_data, ExportDatasetOperatorEventData):
+            event.dataset_operator_event_data.CopyFrom(event_data)
+            event.source_type = ExportEvent.SourceType.EXPORT_DATASET_OPERATOR_EVENT
+        elif isinstance(event_data, ExportDatasetOperatorSchema):
+            event.dataset_operator_schema.CopyFrom(event_data)
+            event.source_type = ExportEvent.SourceType.EXPORT_DATASET_OPERATOR_SCHEMA
+        else:
+            raise TypeError(f"Invalid event_data type: {type(event_data)}")
+        if not self.log_type.supports_event_type(event_data):
+            global_logger.error(
+                f"event_data has source type {event.source_type}, however "
+                f"the event was sent to a logger with log type {self.log_type.log_type_name}. "
+                f"The event will still be written to the file of {self.log_type.log_type_name} "
+                "but this indicates a bug in the code."
+            )
+            pass
+        return event
+
+    def _export_event_to_string(self, event: ExportEvent) -> str:
+        event_data_json = {}
+        proto_to_dict_options = {
+            "always_print_fields_with_no_presence": True,
+            "preserving_proto_field_name": True,
+            "use_integers_for_enums": False,
+        }
+        event_data_field_set = event.WhichOneof("event_data")
+        if event_data_field_set:
+            event_data_json = message_to_dict(
+                getattr(event, event_data_field_set),
+                **proto_to_dict_options,
+            )
+        else:
+            global_logger.error(
+                f"event_data missing from export event with id {event.event_id} "
+                f"and type {event.source_type}. An empty event will be written, "
+                "but this indicates a bug in the code. "
+            )
+            pass
+        event_json = {
+            "event_id": event.event_id,
+            "timestamp": event.timestamp,
+            "source_type": ExportEvent.SourceType.Name(event.source_type),
+            "event_data": event_data_json,
+        }
+        return json.dumps(event_json)
+
+
+def _build_export_event_file_logger(
+    log_type_name: str, sink_dir: str
+) -> logging.Logger:
+    logger = logging.getLogger("_ray_export_event_logger_" + log_type_name)
+    logger.setLevel(logging.INFO)
+    dir_path = pathlib.Path(sink_dir) / "export_events"
+    filepath = dir_path / f"event_{log_type_name}.log"
+    dir_path.mkdir(exist_ok=True)
+    filepath.touch(exist_ok=True)
+    # Configure the logger.
+    # Default is 100 MB max file size
+    handler = logging.handlers.RotatingFileHandler(
+        filepath,
+        maxBytes=(ray_constants.RAY_EXPORT_EVENT_MAX_FILE_SIZE_BYTES),
+        backupCount=ray_constants.RAY_EXPORT_EVENT_MAX_BACKUP_COUNT,
+    )
+    logger.addHandler(handler)
+    logger.propagate = False
+    return logger
+
+
+# This lock must be used when accessing or updating global event logger dict.
+_export_event_logger_lock = threading.Lock()
+_export_event_logger = {}
+
+
+def get_export_event_logger(log_type: EventLogType, sink_dir: str) -> logging.Logger:
+    """Get the export event logger of the current process.
+
+    There's only one logger per export event source.
+
+    Args:
+        log_type: The type of the export event.
+        sink_dir: The directory to sink event logs.
+
+    Returns:
+        The export event logger adapter for the given log type.
+    """
+    with _export_event_logger_lock:
+        global _export_event_logger
+        log_type_name = log_type.log_type_name
+        if log_type_name not in _export_event_logger:
+            logger = _build_export_event_file_logger(log_type.log_type_name, sink_dir)
+            _export_event_logger[log_type_name] = ExportEventLoggerAdapter(
+                log_type, logger
+            )
+
+        return _export_event_logger[log_type_name]
+
+
+def check_export_api_enabled(
+    source: ExportEvent.SourceType,
+) -> bool:
+    """
+    Check RAY_ENABLE_EXPORT_API_WRITE and RAY_ENABLE_EXPORT_API_WRITE_CONFIG environment
+    variables to verify if export events should be written for the given source type.
+
+    Args:
+        source: The source of the export event.
+
+    Returns:
+        True if the export API is enabled for the given source, else False.
+    """
+    if ray_constants.RAY_ENABLE_EXPORT_API_WRITE:
+        return True
+    source_name = ExportEvent.SourceType.Name(source)
+    return (
+        source_name in ray_constants.RAY_ENABLE_EXPORT_API_WRITE_CONFIG
+        if ray_constants.RAY_ENABLE_EXPORT_API_WRITE_CONFIG
+        else False
+    )

python/ray/_private/external_storage.py

@@ -0,0 +1,676 @@
+import abc
+import logging
+import os
+import random
+import shutil
+import time
+import urllib
+import uuid
+from collections import namedtuple
+from typing import IO, List, Optional, Tuple, Union
+
+import ray
+from ray._private.ray_constants import DEFAULT_OBJECT_PREFIX
+from ray._raylet import ObjectRef
+
+ParsedURL = namedtuple("ParsedURL", "base_url, offset, size")
+logger = logging.getLogger(__name__)
+
+
+def create_url_with_offset(*, url: str, offset: int, size: int) -> str:
+    """Methods to create a URL with offset.
+
+    When ray spills objects, it fuses multiple objects
+    into one file to optimize the performance. That says, each object
+    needs to keep tracking of its own special url to store metadata.
+
+    This method creates an url_with_offset, which is used internally
+    by Ray.
+
+    Created url_with_offset can be passed to the self._get_base_url method
+    to parse the filename used to store files.
+
+    Example) file://path/to/file?offset=""&size=""
+
+    Args:
+        url: url to the object stored in the external storage.
+        offset: Offset from the beginning of the file to
+            the first bytes of this object.
+        size: Size of the object that is stored in the url.
+            It is used to calculate the last offset.
+
+    Returns:
+        url_with_offset stored internally to find
+        objects from external storage.
+    """
+    return f"{url}?offset={offset}&size={size}"
+
+
+def parse_url_with_offset(url_with_offset: str) -> Tuple[str, int, int]:
+    """Parse url_with_offset to retrieve information.
+
+    base_url is the url where the object ref
+    is stored in the external storage.
+
+    Args:
+        url_with_offset: url created by create_url_with_offset.
+
+    Returns:
+        named tuple of base_url, offset, and size.
+    """
+    parsed_result = urllib.parse.urlparse(url_with_offset)
+    query_dict = urllib.parse.parse_qs(parsed_result.query)
+    # Split by ? to remove the query from the url.
+    base_url = parsed_result.geturl().split("?")[0]
+    if "offset" not in query_dict or "size" not in query_dict:
+        raise ValueError(f"Failed to parse URL: {url_with_offset}")
+    offset = int(query_dict["offset"][0])
+    size = int(query_dict["size"][0])
+    return ParsedURL(base_url=base_url, offset=offset, size=size)
+
+
+class ExternalStorage(metaclass=abc.ABCMeta):
+    """The base class for external storage.
+
+    This class provides some useful functions for zero-copy object
+    put/get from plasma store. Also it specifies the interface for
+    object spilling.
+
+    When inheriting this class, please make sure to implement validation
+    logic inside __init__ method. When ray instance starts, it will
+    instantiating external storage to validate the config.
+    """
+
+    HEADER_LENGTH = 24
+    CORE_WORKER_INIT_GRACE_PERIOD_S = 1
+
+    def __init__(self):
+        # NOTE(edoakes): do not access this field directly. Use the `core_worker`
+        # property instead to handle initialization race conditions.
+        self._core_worker: Optional["ray._raylet.CoreWorker"] = None
+
+    @property
+    def core_worker(self) -> "ray._raylet.CoreWorker":
+        """Get the core_worker initialized in this process.
+
+        In rare cases, the core worker may not be fully initialized by the time an I/O
+        worker begins to execute an operation because there is no explicit flag set to
+        indicate that the Python layer is ready to execute tasks.
+        """
+        if self._core_worker is None:
+            worker = ray._private.worker.global_worker
+            start = time.time()
+            while not worker.connected:
+                time.sleep(0.001)
+                if time.time() - start > self.CORE_WORKER_INIT_GRACE_PERIOD_S:
+                    raise RuntimeError(
+                        "CoreWorker didn't initialize within grace period of "
+                        f"{self.CORE_WORKER_INIT_GRACE_PERIOD_S}s."
+                    )
+
+            self._core_worker = worker.core_worker
+
+        return self._core_worker
+
+    def _get_objects_from_store(self, object_refs):
+        # Since the object should always exist in the plasma store before
+        # spilling, it can directly get the object from the local plasma
+        # store.
+        # issue: https://github.com/ray-project/ray/pull/13831
+        return self.core_worker.get_if_local(object_refs)
+
+    def _put_object_to_store(
+        self, metadata, data_size, file_like, object_ref, owner_address
+    ):
+        self.core_worker.put_file_like_object(
+            metadata, data_size, file_like, object_ref, owner_address
+        )
+
+    def _write_multiple_objects(
+        self, f: IO, object_refs: List[ObjectRef], owner_addresses: List[str], url: str
+    ) -> List[str]:
+        """Fuse all given objects into a given file handle.
+
+        Args:
+            f: File handle to fusion all given object refs.
+            object_refs: Object references to fusion to a single file.
+            owner_addresses: Owner addresses for the provided objects.
+            url: url where the object ref is stored
+                in the external storage.
+
+        Returns:
+            List of urls_with_offset of fused objects.
+            The order of returned keys are equivalent to the one
+            with given object_refs.
+
+        Raises:
+            ValueError: when given configuration for
+                the external storage is invalid.
+        """
+        keys = []
+        offset = 0
+        ray_object_pairs = self._get_objects_from_store(object_refs)
+        for ref, (buf, metadata, _), owner_address in zip(
+            object_refs, ray_object_pairs, owner_addresses
+        ):
+            address_len = len(owner_address)
+            metadata_len = len(metadata)
+            if buf is None and len(metadata) == 0:
+                error = f"Object {ref.hex()} does not exist."
+                raise ValueError(error)
+            buf_len = 0 if buf is None else len(buf)
+            header = (
+                address_len.to_bytes(8, byteorder="little")
+                + metadata_len.to_bytes(8, byteorder="little")
+                + buf_len.to_bytes(8, byteorder="little")
+                + owner_address
+                + metadata
+            )
+            # 24 bytes to store owner address, metadata, and buffer lengths.
+            payload_len = self.HEADER_LENGTH + address_len + metadata_len + buf_len
+            written_bytes = f.write(header)
+            if buf_len:
+                written_bytes += f.write(memoryview(buf))
+            assert written_bytes == payload_len
+            url_with_offset = create_url_with_offset(
+                url=url, offset=offset, size=written_bytes
+            )
+            keys.append(url_with_offset.encode())
+            offset += written_bytes
+        # Necessary because pyarrow.io.NativeFile does not flush() on close().
+        f.flush()
+        return keys
+
+    def _size_check(
+        self,
+        address_len: int,
+        metadata_len: int,
+        buffer_len: int,
+        obtained_data_size: int,
+    ):
+        """Check whether or not the obtained_data_size is as expected.
+
+        Args:
+            address_len: Length of the address.
+            metadata_len: Actual metadata length of the object.
+            buffer_len: Actual buffer length of the object.
+            obtained_data_size: Data size specified in the url_with_offset.
+
+        Raises:
+            ValueError: If obtained_data_size is different from
+                address_len + metadata_len + buffer_len + 24 (first 8 bytes to store length).
+        """
+        data_size_in_bytes = (
+            address_len + metadata_len + buffer_len + self.HEADER_LENGTH
+        )
+        if data_size_in_bytes != obtained_data_size:
+            raise ValueError(
+                f"Obtained data has a size of {data_size_in_bytes}, "
+                "although it is supposed to have the "
+                f"size of {obtained_data_size}."
+            )
+
+    @abc.abstractmethod
+    def spill_objects(
+        self, object_refs: List[ObjectRef], owner_addresses: List[str]
+    ) -> List[str]:
+        """Spill objects to the external storage. Objects are specified
+        by their object refs.
+
+        Args:
+            object_refs: The list of the refs of the objects to be spilled.
+            owner_addresses: Owner addresses for the provided objects.
+
+        Returns:
+            A list of internal URLs with object offset.
+        """
+
+    @abc.abstractmethod
+    def restore_spilled_objects(
+        self, object_refs: List[ObjectRef], url_with_offset_list: List[str]
+    ) -> int:
+        """Restore objects from the external storage.
+
+        Args:
+            object_refs: List of object IDs (note that it is not ref).
+            url_with_offset_list: List of url_with_offset.
+
+        Returns:
+            The total number of bytes restored.
+        """
+
+    @abc.abstractmethod
+    def delete_spilled_objects(self, urls: List[str]):
+        """Delete objects that are spilled to the external storage.
+
+        Args:
+            urls: URLs that store spilled object files.
+
+        NOTE: This function should not fail if some of the urls
+        do not exist.
+        """
+
+    @abc.abstractmethod
+    def destroy_external_storage(self):
+        """Destroy external storage when a head node is down.
+
+        NOTE: This is currently working when the cluster is
+        started by ray.init
+        """
+
+
+class NullStorage(ExternalStorage):
+    """The class that represents an uninitialized external storage."""
+
+    def spill_objects(self, object_refs, owner_addresses) -> List[str]:
+        raise NotImplementedError("External storage is not initialized")
+
+    def restore_spilled_objects(self, object_refs, url_with_offset_list):
+        raise NotImplementedError("External storage is not initialized")
+
+    def delete_spilled_objects(self, urls: List[str]):
+        raise NotImplementedError("External storage is not initialized")
+
+    def destroy_external_storage(self):
+        raise NotImplementedError("External storage is not initialized")
+
+
+class FileSystemStorage(ExternalStorage):
+    """The class for filesystem-like external storage."""
+
+    def __init__(
+        self,
+        node_id: str,
+        directory_path: Union[str, List[str]],
+        buffer_size: Optional[int] = None,
+    ):
+        """Initialize FileSystemStorage.
+
+        Args:
+            node_id: The ID of the node this storage is associated with.
+            directory_path: A path or list of paths to spill objects to.
+            buffer_size: File buffer size used for writes.
+
+        Raises:
+            ValueError: Raises directory path to spill objects doesn't exist.
+        """
+        super().__init__()
+
+        # -- A list of directory paths to spill objects --
+        self._directory_paths = []
+        # -- Current directory to spill objects --
+        self._current_directory_index = 0
+        # -- File buffer size to spill objects --
+        self._buffer_size = -1
+
+        # Validation.
+        assert (
+            directory_path is not None
+        ), "directory_path should be provided to use object spilling."
+        if isinstance(directory_path, str):
+            directory_path = [directory_path]
+        assert isinstance(
+            directory_path, list
+        ), "Directory_path must be either a single string or a list of strings"
+        if buffer_size is not None:
+            assert isinstance(buffer_size, int), "buffer_size must be an integer."
+            self._buffer_size = buffer_size
+
+        # Create directories.
+        for path in directory_path:
+            full_dir_path = os.path.join(path, f"{DEFAULT_OBJECT_PREFIX}_{node_id}")
+            os.makedirs(full_dir_path, exist_ok=True)
+            if not os.path.exists(full_dir_path):
+                raise ValueError(
+                    "The given directory path to store objects, "
+                    f"{full_dir_path}, could not be created."
+                )
+            self._directory_paths.append(full_dir_path)
+        assert len(self._directory_paths) == len(directory_path)
+        # Choose the current directory.
+        # It chooses a random index to maximize multiple directories that are
+        # mounted at different point.
+        self._current_directory_index = random.randrange(0, len(self._directory_paths))
+
+    def spill_objects(self, object_refs, owner_addresses) -> List[str]:
+        if len(object_refs) == 0:
+            return []
+        # Choose the current directory path by round robin order.
+        self._current_directory_index = (self._current_directory_index + 1) % len(
+            self._directory_paths
+        )
+        directory_path = self._directory_paths[self._current_directory_index]
+
+        filename = _get_unique_spill_filename(object_refs)
+        url = f"{os.path.join(directory_path, filename)}"
+        with open(url, "wb", buffering=self._buffer_size) as f:
+            return self._write_multiple_objects(f, object_refs, owner_addresses, url)
+
+    def restore_spilled_objects(
+        self, object_refs: List[ObjectRef], url_with_offset_list: List[str]
+    ):
+        total = 0
+        for i in range(len(object_refs)):
+            object_ref = object_refs[i]
+            url_with_offset = url_with_offset_list[i].decode()
+            # Retrieve the information needed.
+            parsed_result = parse_url_with_offset(url_with_offset)
+            base_url = parsed_result.base_url
+            offset = parsed_result.offset
+            # Read a part of the file and recover the object.
+            with open(base_url, "rb") as f:
+                f.seek(offset)
+                address_len = int.from_bytes(f.read(8), byteorder="little")
+                metadata_len = int.from_bytes(f.read(8), byteorder="little")
+                buf_len = int.from_bytes(f.read(8), byteorder="little")
+                self._size_check(address_len, metadata_len, buf_len, parsed_result.size)
+                total += buf_len
+                owner_address = f.read(address_len)
+                metadata = f.read(metadata_len)
+                # read remaining data to our buffer
+                self._put_object_to_store(
+                    metadata, buf_len, f, object_ref, owner_address
+                )
+        return total
+
+    def delete_spilled_objects(self, urls: List[str]):
+        for url in urls:
+            path = parse_url_with_offset(url.decode()).base_url
+            try:
+                os.remove(path)
+            except FileNotFoundError:
+                # Occurs when the urls are retried during worker crash/failure.
+                pass
+
+    def destroy_external_storage(self):
+        for directory_path in self._directory_paths:
+            self._destroy_external_storage(directory_path)
+
+    def _destroy_external_storage(self, directory_path):
+        # There's a race condition where IO workers are still
+        # deleting each objects while we try deleting the
+        # whole directory. So we should keep trying it until
+        # The directory is actually deleted.
+        while os.path.isdir(directory_path):
+            try:
+                shutil.rmtree(directory_path)
+            except (FileNotFoundError):
+                # If exception occurs when other IO workers are
+                # deleting the file at the same time.
+                pass
+            except Exception:
+                logger.exception(
+                    "Error cleaning up spill files. "
+                    "You might still have remaining spilled "
+                    "objects inside `ray_spilled_objects` directory."
+                )
+                break
+
+
+class ExternalStorageSmartOpenImpl(ExternalStorage):
+    """The external storage class implemented by smart_open.
+    (https://github.com/RaRe-Technologies/smart_open)
+
+    Smart open supports multiple backend with the same APIs.
+
+    To use this implementation, you should pre-create the given uri.
+    For example, if your uri is a local file path, you should pre-create
+    the directory.
+    """
+
+    def __init__(
+        self,
+        node_id: str,
+        uri: Union[str, list],
+        override_transport_params: dict = None,
+        buffer_size: int = 1024
+        * 1024,  # For remote spilling, at least 1MB is recommended.
+    ):
+        """Initialize ExternalStorageSmartOpenImpl.
+
+        Args:
+            node_id: The ID of the node this storage is associated with.
+            uri: Storage URI used for smart open.
+            override_transport_params: Overriding the default value of
+                transport_params for smart-open library.
+            buffer_size: File buffer size used for writes.
+
+        Raises:
+            ModuleNotFoundError: If it fails to setup. For example, if smart
+                open library is not downloaded, this will fail.
+        """
+        super().__init__()
+
+        try:
+            from smart_open import open  # noqa
+        except ModuleNotFoundError as e:
+            raise ModuleNotFoundError(
+                "Smart open is chosen to be a object spilling "
+                "external storage, but smart_open and boto3 "
+                f"is not downloaded. Original error: {e}"
+            )
+
+        # Validation
+        assert uri is not None, "uri should be provided to use object spilling."
+        if isinstance(uri, str):
+            uri = [uri]
+        assert isinstance(uri, list), "uri must be a single string or list of strings."
+        assert isinstance(buffer_size, int), "buffer_size must be an integer."
+
+        uri_is_s3 = [u.startswith("s3://") for u in uri]
+        self.is_for_s3 = all(uri_is_s3)
+        if not self.is_for_s3:
+            assert not any(uri_is_s3), "all uri's must be s3 or none can be s3."
+            self._uris = uri
+        else:
+            self._uris = [u.strip("/") for u in uri]
+        assert len(self._uris) == len(uri)
+
+        self._current_uri_index = random.randrange(0, len(self._uris))
+        self.prefix = f"{DEFAULT_OBJECT_PREFIX}_{node_id}"
+        self.override_transport_params = override_transport_params or {}
+
+        if self.is_for_s3:
+            import boto3  # noqa
+
+            # Setup boto3. It is essential because if we don't create boto
+            # session, smart_open will create a new session for every
+            # open call.
+            self.s3 = boto3.resource(service_name="s3")
+
+            # smart_open always seek to 0 if we don't set this argument.
+            # This will lead us to call a Object.get when it is not necessary,
+            # so defer seek and call seek before reading objects instead.
+            self.transport_params = {
+                "defer_seek": True,
+                "resource": self.s3,
+                "buffer_size": buffer_size,
+            }
+        else:
+            self.transport_params = {}
+
+        self.transport_params.update(self.override_transport_params)
+
+    def spill_objects(self, object_refs, owner_addresses) -> List[str]:
+        if len(object_refs) == 0:
+            return []
+        from smart_open import open
+
+        # Choose the current uri by round robin order.
+        self._current_uri_index = (self._current_uri_index + 1) % len(self._uris)
+        uri = self._uris[self._current_uri_index]
+
+        key = f"{self.prefix}-{_get_unique_spill_filename(object_refs)}"
+        url = f"{uri}/{key}"
+
+        with open(
+            url,
+            mode="wb",
+            transport_params=self.transport_params,
+        ) as file_like:
+            return self._write_multiple_objects(
+                file_like, object_refs, owner_addresses, url
+            )
+
+    def restore_spilled_objects(
+        self, object_refs: List[ObjectRef], url_with_offset_list: List[str]
+    ):
+        from smart_open import open
+
+        total = 0
+        for i in range(len(object_refs)):
+            object_ref = object_refs[i]
+            url_with_offset = url_with_offset_list[i].decode()
+
+            # Retrieve the information needed.
+            parsed_result = parse_url_with_offset(url_with_offset)
+            base_url = parsed_result.base_url
+            offset = parsed_result.offset
+
+            with open(base_url, "rb", transport_params=self.transport_params) as f:
+                # smart open seek reads the file from offset-end_of_the_file
+                # when the seek is called.
+                f.seek(offset)
+                address_len = int.from_bytes(f.read(8), byteorder="little")
+                metadata_len = int.from_bytes(f.read(8), byteorder="little")
+                buf_len = int.from_bytes(f.read(8), byteorder="little")
+                self._size_check(address_len, metadata_len, buf_len, parsed_result.size)
+                owner_address = f.read(address_len)
+                total += buf_len
+                metadata = f.read(metadata_len)
+                # read remaining data to our buffer
+                self._put_object_to_store(
+                    metadata, buf_len, f, object_ref, owner_address
+                )
+        return total
+
+    def delete_spilled_objects(self, urls: List[str]):
+        pass
+
+    def destroy_external_storage(self):
+        pass
+
+
+_external_storage = NullStorage()
+
+
+class UnstableFileStorage(FileSystemStorage):
+    """This class is for testing with writing failure."""
+
+    def __init__(self, node_id: str, **kwargs):
+        super().__init__(node_id, **kwargs)
+        self._failure_rate = 0.1
+        self._partial_failure_ratio = 0.2
+
+    def spill_objects(self, object_refs, owner_addresses) -> List[str]:
+        r = random.random() < self._failure_rate
+        failed = r < self._failure_rate
+        partial_failed = r < self._partial_failure_ratio
+        if failed:
+            raise IOError("Spilling object failed intentionally for testing.")
+        elif partial_failed:
+            i = random.choice(range(len(object_refs)))
+            return super().spill_objects(object_refs[:i], owner_addresses)
+        else:
+            return super().spill_objects(object_refs, owner_addresses)
+
+
+class SlowFileStorage(FileSystemStorage):
+    """This class is for testing slow object spilling."""
+
+    def __init__(self, node_id: str, **kwargs):
+        super().__init__(node_id, **kwargs)
+        self._min_delay = 1
+        self._max_delay = 2
+
+    def spill_objects(self, object_refs, owner_addresses) -> List[str]:
+        delay = random.random() * (self._max_delay - self._min_delay) + self._min_delay
+        time.sleep(delay)
+        return super().spill_objects(object_refs, owner_addresses)
+
+
+def setup_external_storage(config, node_id, session_name):
+    """Setup the external storage according to the config."""
+    assert node_id is not None, "node_id should be provided."
+    global _external_storage
+    if config:
+        storage_type = config["type"]
+        if storage_type == "filesystem":
+            _external_storage = FileSystemStorage(node_id, **config["params"])
+        elif storage_type == "smart_open":
+            _external_storage = ExternalStorageSmartOpenImpl(
+                node_id, **config["params"]
+            )
+        elif storage_type == "mock_distributed_fs":
+            # This storage is used to unit test distributed external storages.
+            # TODO(sang): Delete it after introducing the mock S3 test.
+            _external_storage = FileSystemStorage(node_id, **config["params"])
+        elif storage_type == "unstable_fs":
+            # This storage is used to unit test unstable file system for fault
+            # tolerance.
+            _external_storage = UnstableFileStorage(node_id, **config["params"])
+        elif storage_type == "slow_fs":
+            # This storage is used to unit test slow filesystems.
+            _external_storage = SlowFileStorage(node_id, **config["params"])
+        else:
+            raise ValueError(f"Unknown external storage type: {storage_type}")
+    else:
+        _external_storage = NullStorage()
+    return _external_storage
+
+
+def reset_external_storage():
+    global _external_storage
+    _external_storage = NullStorage()
+
+
+def spill_objects(
+    object_refs: List[ObjectRef], owner_addresses: List[str]
+) -> List[str]:
+    """Spill objects to the external storage. Objects are specified
+    by their object refs.
+
+    Args:
+        object_refs: The list of the refs of the objects to be spilled.
+        owner_addresses: The owner addresses of the provided object refs.
+
+    Returns:
+        A list of keys corresponding to the input object refs.
+    """
+    return _external_storage.spill_objects(object_refs, owner_addresses)
+
+
+def restore_spilled_objects(
+    object_refs: List[ObjectRef], url_with_offset_list: List[str]
+):
+    """Restore objects from the external storage.
+
+    Args:
+        object_refs: List of object IDs (note that it is not ref).
+        url_with_offset_list: List of url_with_offset.
+
+    Returns:
+        The total number of bytes restored.
+    """
+    return _external_storage.restore_spilled_objects(object_refs, url_with_offset_list)
+
+
+def delete_spilled_objects(urls: List[str]):
+    """Delete objects that are spilled to the external storage.
+
+    Args:
+        urls: URLs that store spilled object files.
+    """
+    _external_storage.delete_spilled_objects(urls)
+
+
+def _get_unique_spill_filename(object_refs: List[ObjectRef]) -> str:
+    """Generate a unique spill file name.
+
+    Args:
+        object_refs: objects to be spilled in this file.
+
+    Returns:
+        A unique filename for the spilled object batch.
+    """
+    return f"{uuid.uuid4().hex}-multi-{len(object_refs)}"

python/ray/_private/function_manager.py

@@ -0,0 +1,750 @@
+import dis
+import hashlib
+import importlib
+import inspect
+import json
+import logging
+import os
+import sys
+import threading
+import time
+import traceback
+from collections import defaultdict, namedtuple
+from typing import Callable, Optional
+
+import ray
+import ray._private.profiling as profiling
+from ray import cloudpickle as pickle
+from ray._common.serialization import pickle_dumps
+from ray._private import ray_constants
+from ray._private.inspect_util import (
+    is_class_method,
+    is_function_or_method,
+    is_static_method,
+)
+from ray._private.ray_constants import KV_NAMESPACE_FUNCTION_TABLE
+from ray._private.utils import (
+    check_oversized_function,
+    ensure_str,
+    format_error_message,
+)
+from ray._raylet import (
+    WORKER_PROCESS_SETUP_HOOK_KEY_NAME_GCS,
+    JobID,
+    PythonFunctionDescriptor,
+)
+from ray.remote_function import RemoteFunction
+from ray.util.tracing.tracing_helper import _inject_tracing_into_class
+
+FunctionExecutionInfo = namedtuple(
+    "FunctionExecutionInfo", ["function", "function_name", "max_calls"]
+)
+ImportedFunctionInfo = namedtuple(
+    "ImportedFunctionInfo",
+    ["job_id", "function_id", "function_name", "function", "module", "max_calls"],
+)
+
+"""FunctionExecutionInfo: A named tuple storing remote function information."""
+
+logger = logging.getLogger(__name__)
+
+
+def make_function_table_key(key_type: bytes, job_id: JobID, key: Optional[bytes]):
+    if key is None:
+        return b":".join([key_type, job_id.hex().encode()])
+    else:
+        return b":".join([key_type, job_id.hex().encode(), key])
+
+
+def build_setup_hook_export_entry(
+    setup_func: Callable, job_id: JobID
+) -> tuple[bytes, bytes, bytes]:
+    """Compute the exported payload and GCS key for a setup hook callable.
+
+    Args:
+        setup_func: The setup hook function to export.
+        job_id: The job ID to export the setup hook for.
+
+    Returns:
+        A tuple of (pickled_function, function_id, key).
+    """
+    pickled_function = pickle_dumps(
+        setup_func,
+        "Cannot serialize the worker_process_setup_hook " f"{setup_func.__name__}",
+    )
+    function_to_run_id = hashlib.shake_128(pickled_function).digest(
+        ray_constants.ID_SIZE
+    )
+    key = make_function_table_key(
+        # This value should match with gcs_function_manager.h.
+        # Otherwise, it won't be GC'ed.
+        WORKER_PROCESS_SETUP_HOOK_KEY_NAME_GCS.encode(),
+        # b"FunctionsToRun",
+        job_id,
+        function_to_run_id,
+    )
+    return pickled_function, function_to_run_id, key
+
+
+class FunctionActorManager:
+    """A class used to export/load remote functions and actors.
+    Attributes:
+        _worker: The associated worker that this manager related.
+        _functions_to_export: The remote functions to export when
+            the worker gets connected.
+        _actors_to_export: The actors to export when the worker gets
+            connected.
+        _function_execution_info: The function_id
+            and execution_info.
+        _num_task_executions: The function
+            execution times.
+        imported_actor_classes: The set of actor classes keys (format:
+            ActorClass:function_id) that are already in GCS.
+    """
+
+    def __init__(self, worker: "ray._private.worker.Worker"):
+        """Initialize FunctionActorManager.
+
+        Args:
+            worker: The worker this manager belongs to.
+        """
+        self._worker = worker
+        self._functions_to_export = []
+        self._actors_to_export = []
+        # This field is a dictionary that maps function IDs
+        # to a FunctionExecutionInfo object. This should only be used on
+        # workers that execute remote functions.
+        self._function_execution_info = defaultdict(lambda: {})
+        self._num_task_executions = defaultdict(lambda: {})
+        # A set of all of the actor class keys that have been imported by the
+        # import thread. It is safe to convert this worker into an actor of
+        # these types.
+        self.imported_actor_classes = set()
+        self._loaded_actor_classes = {}
+        # Deserialize an ActorHandle will call load_actor_class(). If a
+        # function closure captured an ActorHandle, the deserialization of the
+        # function will be:
+        #         -> fetch_and_register_remote_function (acquire lock)
+        #         -> _load_actor_class_from_gcs (acquire lock, too)
+        # So, the lock should be a reentrant lock.
+        self.lock = threading.RLock()
+
+        self.execution_infos = {}
+        # This is the counter to keep track of how many keys have already
+        # been exported so that we can find next key quicker.
+        self._num_exported = 0
+        # This is to protect self._num_exported when doing exporting
+        self._export_lock = threading.Lock()
+
+    def increase_task_counter(self, function_descriptor):
+        function_id = function_descriptor.function_id
+        self._num_task_executions[function_id] += 1
+
+    def get_task_counter(self, function_descriptor):
+        function_id = function_descriptor.function_id
+        return self._num_task_executions[function_id]
+
+    def compute_collision_identifier(self, function_or_class: Callable) -> bytes:
+        """The identifier is used to detect excessive duplicate exports.
+        The identifier is used to determine when the same function or class is
+        exported many times. This can yield false positives.
+        Args:
+            function_or_class: The function or class to compute an identifier
+                for.
+        Returns:
+            The identifier. Note that different functions or classes can give
+                rise to same identifier. However, the same function should
+                hopefully always give rise to the same identifier. TODO(rkn):
+                verify if this is actually the case. Note that if the
+                identifier is incorrect in any way, then we may give warnings
+                unnecessarily or fail to give warnings, but the application's
+                behavior won't change.
+        """
+        import io
+
+        string_file = io.StringIO()
+        dis.dis(function_or_class, file=string_file, depth=2)
+        collision_identifier = function_or_class.__name__ + ":" + string_file.getvalue()
+
+        # Return a hash of the identifier in case it is too large.
+        return hashlib.sha256(collision_identifier.encode("utf-8")).digest()
+
+    def load_function_or_class_from_local(self, module_name, function_or_class_name):
+        """Try to load a function or class in the module from local."""
+        module = importlib.import_module(module_name)
+        parts = [part for part in function_or_class_name.split(".") if part]
+        object = module
+        try:
+            for part in parts:
+                object = getattr(object, part)
+            return object
+        except Exception:
+            return None
+
+    def export_setup_func(
+        self, setup_func: Callable, timeout: Optional[int] = None
+    ) -> bytes:
+        """Export the setup hook function and return the key."""
+        pickled_function, function_to_run_id, key = build_setup_hook_export_entry(
+            setup_func, self._worker.current_job_id.binary()
+        )
+
+        check_oversized_function(
+            pickled_function, setup_func.__name__, "function", self._worker
+        )
+
+        try:
+            self._worker.gcs_client.internal_kv_put(
+                key,
+                pickle.dumps(
+                    {
+                        "job_id": self._worker.current_job_id.binary(),
+                        "function_id": function_to_run_id,
+                        "function": pickled_function,
+                    }
+                ),
+                # overwrite
+                True,
+                ray_constants.KV_NAMESPACE_FUNCTION_TABLE,
+                timeout=timeout,
+            )
+        except Exception as e:
+            logger.exception(
+                "Failed to export the setup hook " f"{setup_func.__name__}."
+            )
+            raise e
+
+        return key
+
+    def export(self, remote_function: RemoteFunction) -> None:
+        """Pickle a remote function and export it to redis.
+
+        Args:
+            remote_function: the RemoteFunction object.
+        """
+        if self._worker.load_code_from_local:
+            function_descriptor = remote_function._function_descriptor
+            module_name, function_name = (
+                function_descriptor.module_name,
+                function_descriptor.function_name,
+            )
+            # If the function is dynamic, we still export it to GCS
+            # even if load_code_from_local is set True.
+            if (
+                self.load_function_or_class_from_local(module_name, function_name)
+                is not None
+            ):
+                return
+        function = remote_function._function
+        pickled_function = remote_function._pickled_function
+
+        check_oversized_function(
+            pickled_function,
+            remote_function._function_name,
+            "remote function",
+            self._worker,
+        )
+        key = make_function_table_key(
+            b"RemoteFunction",
+            self._worker.current_job_id,
+            remote_function._function_descriptor.function_id.binary(),
+        )
+        if self._worker.gcs_client.internal_kv_exists(key, KV_NAMESPACE_FUNCTION_TABLE):
+            return
+        val = pickle.dumps(
+            {
+                "job_id": self._worker.current_job_id.binary(),
+                "function_id": remote_function._function_descriptor.function_id.binary(),  # noqa: E501
+                "function_name": remote_function._function_name,
+                "module": function.__module__,
+                "function": pickled_function,
+                "collision_identifier": self.compute_collision_identifier(function),
+                "max_calls": remote_function._max_calls,
+            }
+        )
+        self._worker.gcs_client.internal_kv_put(
+            key, val, True, KV_NAMESPACE_FUNCTION_TABLE
+        )
+
+    def fetch_registered_method(
+        self, key: str, timeout: Optional[int] = None
+    ) -> Optional[ImportedFunctionInfo]:
+        vals = self._worker.gcs_client.internal_kv_get(
+            key, KV_NAMESPACE_FUNCTION_TABLE, timeout=timeout
+        )
+        if vals is None:
+            return None
+        else:
+            vals = pickle.loads(vals)
+            fields = [
+                "job_id",
+                "function_id",
+                "function_name",
+                "function",
+                "module",
+                "max_calls",
+            ]
+            return ImportedFunctionInfo._make(vals.get(field) for field in fields)
+
+    def fetch_and_register_remote_function(self, key):
+        """Import a remote function."""
+        remote_function_info = self.fetch_registered_method(key)
+        if not remote_function_info:
+            return False
+        (
+            job_id_str,
+            function_id_str,
+            function_name,
+            serialized_function,
+            module,
+            max_calls,
+        ) = remote_function_info
+
+        function_id = ray.FunctionID(function_id_str)
+        job_id = ray.JobID(job_id_str)
+        max_calls = int(max_calls)
+
+        # This function is called by ImportThread. This operation needs to be
+        # atomic. Otherwise, there is race condition. Another thread may use
+        # the temporary function above before the real function is ready.
+        with self.lock:
+            self._num_task_executions[function_id] = 0
+
+            try:
+                function = pickle.loads(serialized_function)
+            except Exception:
+                # If an exception was thrown when the remote function was
+                # imported, we record the traceback and notify the scheduler
+                # of the failure.
+                traceback_str = format_error_message(traceback.format_exc())
+
+                def f(*args, **kwargs):
+                    raise RuntimeError(
+                        "The remote function failed to import on the "
+                        "worker. This may be because needed library "
+                        "dependencies are not installed in the worker "
+                        "environment or cannot be found from sys.path "
+                        f"{sys.path}:\n\n{traceback_str}"
+                    )
+
+                # Use a placeholder method when function pickled failed
+                self._function_execution_info[function_id] = FunctionExecutionInfo(
+                    function=f, function_name=function_name, max_calls=max_calls
+                )
+
+                # Log the error message. Log at DEBUG level to avoid overly
+                # spamming the log on import failure. The user gets the error
+                # via the RuntimeError message above.
+                logger.debug(
+                    "Failed to unpickle the remote function "
+                    f"'{function_name}' with "
+                    f"function ID {function_id.hex()}. "
+                    f"Job ID:{job_id}."
+                    f"Traceback:\n{traceback_str}. "
+                )
+            else:
+                # The below line is necessary. Because in the driver process,
+                # if the function is defined in the file where the python
+                # script was started from, its module is `__main__`.
+                # However in the worker process, the `__main__` module is a
+                # different module, which is `default_worker.py`
+                function.__module__ = module
+                self._function_execution_info[function_id] = FunctionExecutionInfo(
+                    function=function, function_name=function_name, max_calls=max_calls
+                )
+        return True
+
+    def get_execution_info(
+        self, job_id: JobID, function_descriptor: PythonFunctionDescriptor
+    ) -> FunctionExecutionInfo:
+        """Get the FunctionExecutionInfo of a remote function.
+        Args:
+            job_id: ID of the job that the function belongs to.
+            function_descriptor: The FunctionDescriptor of the function to get.
+        Returns:
+            A FunctionExecutionInfo object.
+        """
+        function_id = function_descriptor.function_id
+        # If the function has already been loaded,
+        # There's no need to load again
+        if function_id in self._function_execution_info:
+            return self._function_execution_info[function_id]
+        if self._worker.load_code_from_local:
+            # Load function from local code.
+            if not function_descriptor.is_actor_method():
+                # If the function is not able to be loaded,
+                # try to load it from GCS,
+                # even if load_code_from_local is set True
+                if self._load_function_from_local(function_descriptor) is True:
+                    return self._function_execution_info[function_id]
+        # Load function from GCS.
+        # Wait until the function to be executed has actually been
+        # registered on this worker. We will push warnings to the user if
+        # we spend too long in this loop.
+        # The driver function may not be found in sys.path. Try to load
+        # the function from GCS.
+        with profiling.profile("wait_for_function"):
+            self._wait_for_function(function_descriptor, job_id)
+        try:
+            function_id = function_descriptor.function_id
+            info = self._function_execution_info[function_id]
+        except KeyError as e:
+            message = (
+                "Error occurs in get_execution_info: "
+                "job_id: %s, function_descriptor: %s. Message: %s"
+                % (job_id, function_descriptor, e)
+            )
+            raise KeyError(message)
+        return info
+
+    def _load_function_from_local(self, function_descriptor):
+        assert not function_descriptor.is_actor_method()
+        function_id = function_descriptor.function_id
+
+        module_name, function_name = (
+            function_descriptor.module_name,
+            function_descriptor.function_name,
+        )
+
+        object = self.load_function_or_class_from_local(module_name, function_name)
+        if object is not None:
+            # Directly importing from local may break function with dynamic ray.remote,
+            # such as the _start_controller function utilized for the Ray service.
+            if isinstance(object, RemoteFunction):
+                function = object._function
+            else:
+                function = object
+            self._function_execution_info[function_id] = FunctionExecutionInfo(
+                function=function,
+                function_name=function_name,
+                max_calls=0,
+            )
+            self._num_task_executions[function_id] = 0
+            return True
+        else:
+            return False
+
+    def _wait_for_function(
+        self,
+        function_descriptor: PythonFunctionDescriptor,
+        job_id: str,
+        timeout: float = 10,
+    ):
+        """Wait until the function to be executed is present on this worker.
+        This method will simply loop until the import thread has imported the
+        relevant function. If we spend too long in this loop, that may indicate
+        a problem somewhere and we will push an error message to the user.
+        If this worker is an actor, then this will wait until the actor has
+        been defined.
+
+        Args:
+            function_descriptor: The FunctionDescriptor of the function that
+                we want to execute.
+            job_id: The ID of the job to push the error message to
+                if this times out.
+            timeout: Seconds to wait before pushing a warning to the user.
+        """
+        start_time = time.time()
+        # Only send the warning once.
+        warning_sent = False
+        while True:
+            with self.lock:
+                if self._worker.actor_id.is_nil():
+                    if function_descriptor.function_id in self._function_execution_info:
+                        break
+                    else:
+                        key = make_function_table_key(
+                            b"RemoteFunction",
+                            job_id,
+                            function_descriptor.function_id.binary(),
+                        )
+                        if self.fetch_and_register_remote_function(key) is True:
+                            break
+                else:
+                    assert not self._worker.actor_id.is_nil()
+                    # Actor loading will happen when execute_task is called.
+                    assert self._worker.actor_id in self._worker.actors
+                    break
+
+            if time.time() - start_time > timeout:
+                warning_message = (
+                    "This worker was asked to execute a function "
+                    f"that has not been registered ({function_descriptor}, "
+                    f"node={self._worker.node_ip_address}, "
+                    f"worker_id={self._worker.worker_id.hex()}, "
+                    f"pid={os.getpid()}). You may have to restart Ray."
+                )
+                if not warning_sent:
+                    logger.error(warning_message)
+                    ray._private.utils.push_error_to_driver(
+                        self._worker,
+                        ray_constants.WAIT_FOR_FUNCTION_PUSH_ERROR,
+                        warning_message,
+                        job_id=job_id,
+                    )
+                warning_sent = True
+            time.sleep(0.001)
+
+    def export_actor_class(
+        self, Class, actor_creation_function_descriptor, actor_method_names
+    ):
+        if self._worker.load_code_from_local:
+            module_name, class_name = (
+                actor_creation_function_descriptor.module_name,
+                actor_creation_function_descriptor.class_name,
+            )
+            # If the class is dynamic, we still export it to GCS
+            # even if load_code_from_local is set True.
+            if (
+                self.load_function_or_class_from_local(module_name, class_name)
+                is not None
+            ):
+                return
+
+        # `current_job_id` shouldn't be NIL, unless:
+        # 1) This worker isn't an actor;
+        # 2) And a previous task started a background thread, which didn't
+        #    finish before the task finished, and still uses Ray API
+        #    after that.
+        assert not self._worker.current_job_id.is_nil(), (
+            "You might have started a background thread in a non-actor "
+            "task, please make sure the thread finishes before the "
+            "task finishes."
+        )
+        job_id = self._worker.current_job_id
+        key = make_function_table_key(
+            b"ActorClass",
+            job_id,
+            actor_creation_function_descriptor.function_id.binary(),
+        )
+        serialized_actor_class = pickle_dumps(
+            Class,
+            f"Could not serialize the actor class "
+            f"{actor_creation_function_descriptor.repr}",
+        )
+        actor_class_info = {
+            "class_name": actor_creation_function_descriptor.class_name.split(".")[-1],
+            "module": actor_creation_function_descriptor.module_name,
+            "class": serialized_actor_class,
+            "job_id": job_id.binary(),
+            "collision_identifier": self.compute_collision_identifier(Class),
+            "actor_method_names": json.dumps(list(actor_method_names)),
+        }
+
+        check_oversized_function(
+            actor_class_info["class"],
+            actor_class_info["class_name"],
+            "actor",
+            self._worker,
+        )
+
+        self._worker.gcs_client.internal_kv_put(
+            key, pickle.dumps(actor_class_info), True, KV_NAMESPACE_FUNCTION_TABLE
+        )
+        # TODO(rkn): Currently we allow actor classes to be defined
+        # within tasks. I tried to disable this, but it may be necessary
+        # because of https://github.com/ray-project/ray/issues/1146.
+
+    def load_actor_class(
+        self,
+        job_id: JobID,
+        actor_creation_function_descriptor: PythonFunctionDescriptor,
+    ) -> type:
+        """Load the actor class.
+        Args:
+            job_id: job ID of the actor.
+            actor_creation_function_descriptor: Function descriptor of
+                the actor constructor.
+        Returns:
+            The actor class.
+        """
+        function_id = actor_creation_function_descriptor.function_id
+        # Check if the actor class already exists in the cache.
+        actor_class = self._loaded_actor_classes.get(function_id, None)
+        if actor_class is None:
+            # Load actor class.
+            if self._worker.load_code_from_local:
+                # Load actor class from local code first.
+                actor_class = self._load_actor_class_from_local(
+                    actor_creation_function_descriptor
+                )
+                # If the actor is unable to be loaded
+                # from local, try to load it
+                # from GCS even if load_code_from_local is set True
+                if actor_class is None:
+                    actor_class = self._load_actor_class_from_gcs(
+                        job_id, actor_creation_function_descriptor
+                    )
+
+            else:
+                # Load actor class from GCS.
+                actor_class = self._load_actor_class_from_gcs(
+                    job_id, actor_creation_function_descriptor
+                )
+
+            # Re-inject tracing into the loaded class. This is necessary because
+            # cloudpickle doesn't preserve __signature__ attributes on module-level
+            # functions. When a class is pickled and unpickled, user-defined methods
+            # are looked up from the module, losing the __signature__ that was set by
+            # _inject_tracing_into_class during actor creation. Re-injecting tracing
+            # ensures the method signatures include _ray_trace_ctx when tracing is
+            # enabled, matching the behavior expected by _tracing_actor_method_invocation.
+            _inject_tracing_into_class(actor_class)
+
+            # Save the loaded actor class in cache.
+            self._loaded_actor_classes[function_id] = actor_class
+
+            # Generate execution info for the methods of this actor class.
+            module_name = actor_creation_function_descriptor.module_name
+            actor_class_name = actor_creation_function_descriptor.class_name
+            actor_methods = inspect.getmembers(
+                actor_class, predicate=is_function_or_method
+            )
+            for actor_method_name, actor_method in actor_methods:
+                # Actor creation function descriptor use a unique function
+                # hash to solve actor name conflict. When constructing an
+                # actor, the actor creation function descriptor will be the
+                # key to find __init__ method execution info. So, here we
+                # use actor creation function descriptor as method descriptor
+                # for generating __init__ method execution info.
+                if actor_method_name == "__init__":
+                    method_descriptor = actor_creation_function_descriptor
+                else:
+                    method_descriptor = PythonFunctionDescriptor(
+                        module_name, actor_method_name, actor_class_name
+                    )
+                method_id = method_descriptor.function_id
+                executor = self._make_actor_method_executor(
+                    actor_method_name, actor_method
+                )
+                self._function_execution_info[method_id] = FunctionExecutionInfo(
+                    function=executor,
+                    function_name=actor_method_name,
+                    max_calls=0,
+                )
+                self._num_task_executions[method_id] = 0
+            self._num_task_executions[function_id] = 0
+        return actor_class
+
+    def _load_actor_class_from_local(self, actor_creation_function_descriptor):
+        """Load actor class from local code."""
+        module_name, class_name = (
+            actor_creation_function_descriptor.module_name,
+            actor_creation_function_descriptor.class_name,
+        )
+
+        object = self.load_function_or_class_from_local(module_name, class_name)
+
+        if object is not None:
+            if isinstance(object, ray.actor.ActorClass):
+                return object.__ray_metadata__.modified_class
+            else:
+                return ray.actor._modify_class(object)
+        else:
+            return None
+
+    def _create_fake_actor_class(
+        self, actor_class_name, actor_method_names, traceback_str
+    ):
+        class TemporaryActor:
+            async def __dummy_method(self):
+                """Dummy method for this fake actor class to work for async actors.
+                Without this method, this temporary actor class fails to initialize
+                if the original actor class was async."""
+                pass
+
+        def temporary_actor_method(*args, **kwargs):
+            raise RuntimeError(
+                f"The actor with name {actor_class_name} "
+                "failed to import on the worker. This may be because "
+                "needed library dependencies are not installed in the "
+                f"worker environment:\n\n{traceback_str}"
+            )
+
+        for method in actor_method_names:
+            setattr(TemporaryActor, method, temporary_actor_method)
+
+        return TemporaryActor
+
+    def _load_actor_class_from_gcs(self, job_id, actor_creation_function_descriptor):
+        """Load actor class from GCS."""
+        key = make_function_table_key(
+            b"ActorClass",
+            job_id,
+            actor_creation_function_descriptor.function_id.binary(),
+        )
+
+        # Fetch raw data from GCS.
+        vals = self._worker.gcs_client.internal_kv_get(key, KV_NAMESPACE_FUNCTION_TABLE)
+        fields = ["job_id", "class_name", "module", "class", "actor_method_names"]
+        if vals is None:
+            vals = {}
+        else:
+            vals = pickle.loads(vals)
+        (job_id_str, class_name, module, pickled_class, actor_method_names) = (
+            vals.get(field) for field in fields
+        )
+
+        class_name = ensure_str(class_name)
+        module_name = ensure_str(module)
+        job_id = ray.JobID(job_id_str)
+        actor_method_names = json.loads(ensure_str(actor_method_names))
+
+        actor_class = None
+        try:
+            with self.lock:
+                actor_class = pickle.loads(pickled_class)
+        except Exception:
+            logger.debug("Failed to load actor class %s.", class_name)
+            # If an exception was thrown when the actor was imported, we record
+            # the traceback and notify the scheduler of the failure.
+            traceback_str = format_error_message(traceback.format_exc())
+            # The actor class failed to be unpickled, create a fake actor
+            # class instead (just to produce error messages and to prevent
+            # the driver from hanging).
+            actor_class = self._create_fake_actor_class(
+                class_name, actor_method_names, traceback_str
+            )
+
+        # The below line is necessary. Because in the driver process,
+        # if the function is defined in the file where the python script
+        # was started from, its module is `__main__`.
+        # However in the worker process, the `__main__` module is a
+        # different module, which is `default_worker.py`
+        actor_class.__module__ = module_name
+        return actor_class
+
+    def _make_actor_method_executor(self, method_name: str, method: Callable):
+        """Make an executor that wraps a user-defined actor method.
+        The wrapped method updates the worker's internal state and performs any
+        necessary checkpointing operations.
+        Args:
+            method_name: The name of the actor method.
+            method: The actor method to wrap. This should be a
+                method defined on the actor class and should therefore take an
+                instance of the actor as the first argument.
+        Returns:
+            A function that executes the given actor method on the worker's
+                stored instance of the actor. The function also updates the
+                worker's internal state to record the executed method.
+        """
+
+        def actor_method_executor(__ray_actor, *args, **kwargs):
+            # Execute the assigned method.
+            is_bound = is_class_method(method) or is_static_method(
+                type(__ray_actor), method_name
+            )
+            if is_bound:
+                return method(*args, **kwargs)
+            else:
+                return method(__ray_actor, *args, **kwargs)
+
+        # Set method_name and method as attributes to the executor closure
+        # so we can make decision based on these attributes in task executor.
+        # Precisely, asyncio support requires to know whether:
+        # - the method is a ray internal method: starts with __ray
+        # - the method is a coroutine function: defined by async def
+        actor_method_executor.name = method_name
+        actor_method_executor.method = method
+
+        return actor_method_executor

python/ray/_private/gc_collect_manager.py

@@ -0,0 +1,51 @@
+import gc
+import logging
+import threading
+import time
+from typing import Callable, Optional
+
+logger = logging.getLogger(__name__)
+
+
+class PythonGCThread(threading.Thread):
+    """A background thread that triggers Python garbage collection.
+
+    This thread waits for GC events from CoreWorker and triggers `gc.collect()` when
+    when requested."""
+
+    def __init__(self, *, gc_collect_func: Optional[Callable] = None):
+        logger.debug("Starting Python GC thread")
+        super().__init__(name="PythonGCThread", daemon=True)
+        self._should_exit = False
+        self._gc_event = threading.Event()
+        # Sets the gc_collect_func (only for testing), defaults to gc.collect
+        self._gc_collect_func = gc_collect_func or gc.collect
+
+    def trigger_gc(self) -> None:
+        self._gc_event.set()
+
+    def run(self):
+        while not self._should_exit:
+            self._gc_event.wait()
+            self._gc_event.clear()
+
+            if self._should_exit:
+                break
+
+            try:
+                start = time.monotonic()
+                num_freed = self._gc_collect_func()
+                if num_freed > 0:
+                    logger.debug(
+                        "gc.collect() freed {} refs in {} seconds".format(
+                            num_freed, time.monotonic() - start
+                        )
+                    )
+            except Exception as e:
+                logger.error(f"Error during GC: {e}")
+
+    def stop(self):
+        logger.debug("Stopping Python GC thread")
+        self._should_exit = True
+        self._gc_event.set()
+        self.join()

python/ray/_private/gcs_pubsub.py

@@ -0,0 +1,293 @@
+import asyncio
+import logging
+import random
+from collections import deque
+from typing import List, Optional, Tuple
+
+import grpc
+from grpc import aio as aiogrpc
+
+import ray._private.gcs_utils as gcs_utils
+from ray._common.utils import get_or_create_event_loop
+from ray.core.generated import (
+    gcs_pb2,
+    gcs_service_pb2,
+    gcs_service_pb2_grpc,
+    pubsub_pb2,
+)
+
+logger = logging.getLogger(__name__)
+
+_OBSERVABILITY_PUBSUB_CHANNELS = (
+    pubsub_pb2.RAY_ERROR_INFO_CHANNEL,
+    pubsub_pb2.RAY_LOG_CHANNEL,
+    pubsub_pb2.RAY_NODE_RESOURCE_USAGE_CHANNEL,
+)
+
+
+class _SubscriberBase:
+    def __init__(self, worker_id: bytes = None):
+        self._worker_id = worker_id
+        # self._subscriber_id needs to match the binary format of a random
+        # SubscriberID / UniqueID, which is 28 (kUniqueIDSize) random bytes.
+        self._subscriber_id = bytes(bytearray(random.getrandbits(8) for _ in range(28)))
+        self._last_batch_size = 0
+        self._max_processed_sequence_id = 0
+        self._publisher_id = b""
+
+    # Batch size of the result from last poll. Used to indicate whether the
+    # subscriber can keep up.
+    @property
+    def last_batch_size(self):
+        return self._last_batch_size
+
+    def _subscribe_request(self, channel):
+        cmd = pubsub_pb2.Command(channel_type=channel, subscribe_message={})
+        req = gcs_service_pb2.GcsSubscriberCommandBatchRequest(
+            subscriber_id=self._subscriber_id, sender_id=self._worker_id, commands=[cmd]
+        )
+        return req
+
+    def _poll_request(self):
+        return gcs_service_pb2.GcsSubscriberPollRequest(
+            subscriber_id=self._subscriber_id,
+            max_processed_sequence_id=self._max_processed_sequence_id,
+            publisher_id=self._publisher_id,
+        )
+
+    def _unsubscribe_request(self, channels):
+        req = gcs_service_pb2.GcsSubscriberCommandBatchRequest(
+            subscriber_id=self._subscriber_id, sender_id=self._worker_id, commands=[]
+        )
+        for channel in channels:
+            req.commands.append(
+                pubsub_pb2.Command(channel_type=channel, unsubscribe_message={})
+            )
+        return req
+
+    @staticmethod
+    def _should_terminate_polling(e: grpc.RpcError) -> None:
+        # Caller only expects polling to be terminated after deadline exceeded.
+        if e.code() == grpc.StatusCode.DEADLINE_EXCEEDED:
+            return True
+        # Could be a temporary connection issue. Suppress error.
+        # TODO: reconnect GRPC channel?
+        if e.code() == grpc.StatusCode.UNAVAILABLE:
+            return True
+        return False
+
+
+class _AioSubscriber(_SubscriberBase):
+    """Async io subscriber to GCS.
+
+    Usage example common to Aio subscribers:
+        subscriber = GcsAioXxxSubscriber(address="...")
+        await subscriber.subscribe()
+        while running:
+            ...... = await subscriber.poll()
+            ......
+        await subscriber.close()
+    """
+
+    def __init__(
+        self,
+        pubsub_channel_type,
+        worker_id: bytes = None,
+        address: str = None,
+        channel: aiogrpc.Channel = None,
+    ):
+        super().__init__(worker_id)
+
+        if address:
+            assert channel is None, "address and channel cannot both be specified"
+            channel = gcs_utils.create_gcs_channel(address, aio=True)
+        else:
+            assert channel is not None, "One of address and channel must be specified"
+        if pubsub_channel_type in _OBSERVABILITY_PUBSUB_CHANNELS:
+            self._stub = gcs_service_pb2_grpc.ObservabilityPubSubServiceStub(channel)
+        else:
+            self._stub = gcs_service_pb2_grpc.ControlPlanePubSubGcsServiceStub(channel)
+
+        # Type of the channel.
+        self._channel = pubsub_channel_type
+        # A queue of received PubMessage.
+        self._queue = deque()
+        # Indicates whether the subscriber has closed.
+        self._close = asyncio.Event()
+
+    async def subscribe(self) -> None:
+        """Registers a subscription for the subscriber's channel type.
+
+        Before the registration, published messages in the channel will not be
+        saved for the subscriber.
+        """
+        if self._close.is_set():
+            return
+        req = self._subscribe_request(self._channel)
+        await self._stub.GcsSubscriberCommandBatch(req, timeout=30)
+
+    async def _poll_call(self, req, timeout=None):
+        # Wrap GRPC _AioCall as a coroutine.
+        return await self._stub.GcsSubscriberPoll(req, timeout=timeout)
+
+    async def _poll(self, timeout=None) -> None:
+        while len(self._queue) == 0:
+            req = self._poll_request()
+            poll = get_or_create_event_loop().create_task(
+                self._poll_call(req, timeout=timeout)
+            )
+            close = get_or_create_event_loop().create_task(self._close.wait())
+            done, others = await asyncio.wait(
+                [poll, close], timeout=timeout, return_when=asyncio.FIRST_COMPLETED
+            )
+            # Cancel the other task if needed to prevent memory leak.
+            other_task = others.pop()
+            if not other_task.done():
+                other_task.cancel()
+            if poll not in done or close in done:
+                # Request timed out or subscriber closed.
+                break
+            try:
+                self._last_batch_size = len(poll.result().pub_messages)
+                if poll.result().publisher_id != self._publisher_id:
+                    if self._publisher_id != b"":
+                        logger.debug(
+                            f"replied publisher_id {poll.result().publisher_id} "
+                            f"different from {self._publisher_id}, this should "
+                            "only happen during gcs failover."
+                        )
+                    self._publisher_id = poll.result().publisher_id
+                    self._max_processed_sequence_id = 0
+                for msg in poll.result().pub_messages:
+                    if msg.sequence_id <= self._max_processed_sequence_id:
+                        logger.warning(f"Ignoring out of order message {msg}")
+                        continue
+                    self._max_processed_sequence_id = msg.sequence_id
+                    self._queue.append(msg)
+            except grpc.RpcError as e:
+                if self._should_terminate_polling(e):
+                    return
+                raise
+
+    async def close(self) -> None:
+        """Closes the subscriber and its active subscription."""
+
+        # Mark close to terminate inflight polling and prevent future requests.
+        if self._close.is_set():
+            return
+        self._close.set()
+        req = self._unsubscribe_request(channels=[self._channel])
+        try:
+            await self._stub.GcsSubscriberCommandBatch(req, timeout=5)
+        except Exception:
+            pass
+        self._stub = None
+
+
+class GcsAioResourceUsageSubscriber(_AioSubscriber):
+    def __init__(
+        self,
+        worker_id: bytes = None,
+        address: str = None,
+        channel: grpc.Channel = None,
+    ):
+        super().__init__(
+            pubsub_pb2.RAY_NODE_RESOURCE_USAGE_CHANNEL, worker_id, address, channel
+        )
+
+    async def poll(self, timeout: Optional[float] = None) -> Tuple[bytes, str]:
+        """Polls for new resource usage message.
+
+        Args:
+            timeout: Optional timeout in seconds for the poll request.
+
+        Returns:
+            A tuple of string reporter ID and resource usage json string.
+        """
+        await self._poll(timeout=timeout)
+        return self._pop_resource_usage(self._queue)
+
+    @staticmethod
+    def _pop_resource_usage(queue):
+        if len(queue) == 0:
+            return None, None
+        msg = queue.popleft()
+        return msg.key_id.decode(), msg.node_resource_usage_message.json
+
+
+class GcsAioActorSubscriber(_AioSubscriber):
+    def __init__(
+        self,
+        worker_id: bytes = None,
+        address: str = None,
+        channel: grpc.Channel = None,
+    ):
+        super().__init__(pubsub_pb2.GCS_ACTOR_CHANNEL, worker_id, address, channel)
+
+    @property
+    def queue_size(self):
+        return len(self._queue)
+
+    async def poll(
+        self, batch_size: int, timeout: Optional[float] = None
+    ) -> List[Tuple[bytes, gcs_pb2.ActorTableData]]:
+        """Polls for new actor message.
+
+        Args:
+            batch_size: Maximum number of messages to return.
+            timeout: Optional timeout in seconds for the poll request.
+
+        Returns:
+            A list of tuples of binary actor ID and actor table data.
+        """
+        await self._poll(timeout=timeout)
+        return self._pop_actors(self._queue, batch_size=batch_size)
+
+    @staticmethod
+    def _pop_actors(queue, batch_size):
+        if len(queue) == 0:
+            return []
+        popped = 0
+        msgs = []
+        while len(queue) > 0 and popped < batch_size:
+            msg = queue.popleft()
+            msgs.append((msg.key_id, msg.actor_message))
+            popped += 1
+        return msgs
+
+
+class GcsAioNodeInfoSubscriber(_AioSubscriber):
+    def __init__(
+        self,
+        worker_id: bytes = None,
+        address: str = None,
+        channel: grpc.Channel = None,
+    ):
+        super().__init__(pubsub_pb2.GCS_NODE_INFO_CHANNEL, worker_id, address, channel)
+
+    async def poll(
+        self, batch_size: int, timeout: Optional[float] = None
+    ) -> List[Tuple[bytes, gcs_pb2.GcsNodeInfo]]:
+        """Polls for new node info message.
+
+        Args:
+            batch_size: Maximum number of messages to return.
+            timeout: Optional timeout in seconds for the poll request.
+
+        Returns:
+            A list of tuples of (node_id, GcsNodeInfo).
+        """
+        await self._poll(timeout=timeout)
+        return self._pop_node_infos(self._queue, batch_size=batch_size)
+
+    @staticmethod
+    def _pop_node_infos(queue, batch_size):
+        if len(queue) == 0:
+            return []
+        popped = 0
+        msgs = []
+        while len(queue) > 0 and popped < batch_size:
+            msg = queue.popleft()
+            msgs.append((msg.key_id, msg.node_info_message))
+            popped += 1
+        return msgs

python/ray/_private/gcs_utils.py

@@ -0,0 +1,159 @@
+import logging
+from typing import Optional
+
+from ray._private import ray_constants
+from ray.core.generated.common_pb2 import ErrorType, JobConfig
+from ray.core.generated.gcs_pb2 import (
+    ActorTableData,
+    AvailableResources,
+    ErrorTableData,
+    GcsEntry,
+    GcsNodeInfo,
+    JobTableData,
+    PlacementGroupTableData,
+    PubSubMessage,
+    ResourceDemand,
+    ResourceLoad,
+    ResourcesData,
+    ResourceUsageBatchData,
+    TablePrefix,
+    TablePubsub,
+    TaskEvents,
+    TotalResources,
+    WorkerTableData,
+)
+
+logger = logging.getLogger(__name__)
+
+__all__ = [
+    "ActorTableData",
+    "GcsNodeInfo",
+    "AvailableResources",
+    "TotalResources",
+    "JobTableData",
+    "JobConfig",
+    "ErrorTableData",
+    "ErrorType",
+    "GcsEntry",
+    "ResourceUsageBatchData",
+    "ResourcesData",
+    "TablePrefix",
+    "TablePubsub",
+    "TaskEvents",
+    "ResourceDemand",
+    "ResourceLoad",
+    "PubSubMessage",
+    "WorkerTableData",
+    "PlacementGroupTableData",
+]
+
+
+WORKER = 0
+DRIVER = 1
+
+# Cap messages at 512MB
+_MAX_MESSAGE_LENGTH = 512 * 1024 * 1024
+# Send keepalive every 60s
+_GRPC_KEEPALIVE_TIME_MS = 60 * 1000
+# Keepalive should be replied < 60s
+_GRPC_KEEPALIVE_TIMEOUT_MS = 60 * 1000
+
+# Also relying on these defaults:
+# grpc.keepalive_permit_without_calls=0: No keepalive without inflight calls.
+# grpc.use_local_subchannel_pool=0: Subchannels are shared.
+_GRPC_OPTIONS = [
+    *ray_constants.GLOBAL_GRPC_OPTIONS,
+    ("grpc.max_send_message_length", _MAX_MESSAGE_LENGTH),
+    ("grpc.max_receive_message_length", _MAX_MESSAGE_LENGTH),
+    ("grpc.keepalive_time_ms", _GRPC_KEEPALIVE_TIME_MS),
+    ("grpc.keepalive_timeout_ms", _GRPC_KEEPALIVE_TIMEOUT_MS),
+]
+
+
+def create_gcs_channel(address: str, aio: bool = False):
+    """Returns a GRPC channel to GCS.
+
+    Args:
+        address: GCS address string, e.g. ip:port
+        aio: Whether using grpc.aio
+
+    Returns:
+        grpc.Channel or grpc.aio.Channel to GCS
+    """
+    from ray._private.grpc_utils import init_grpc_channel
+
+    return init_grpc_channel(address, options=_GRPC_OPTIONS, asynchronous=aio)
+
+
+class GcsChannel:
+    def __init__(self, gcs_address: Optional[str] = None, aio: bool = False):
+        self._gcs_address = gcs_address
+        self._aio = aio
+
+    @property
+    def address(self):
+        return self._gcs_address
+
+    def connect(self):
+        # GCS server uses a cached port, so it should use the same port after
+        # restarting. This means GCS address should stay the same for the
+        # lifetime of the Ray cluster.
+        self._channel = create_gcs_channel(self._gcs_address, self._aio)
+
+    def channel(self):
+        return self._channel
+
+
+def cleanup_redis_storage(
+    host: str,
+    port: int,
+    password: str,
+    use_ssl: bool,
+    storage_namespace: str,
+    username: Optional[str] = None,
+):
+    """This function is used to cleanup the GCS storage in Redis.
+    It supports Redis in cluster and non-cluster modes.
+
+    Args:
+       host: The Redis host address.
+       port: The Redis port.
+       password: The Redis password.
+       use_ssl: Whether to encrypt the connection.
+       storage_namespace: The namespace of the storage to be deleted.
+       username: The Redis username.
+
+    Returns:
+        The result of deleting the GCS key prefix from the Redis storage.
+    """
+
+    from ray._raylet import del_key_prefix_from_storage  # type: ignore
+
+    if not isinstance(host, str):
+        raise ValueError("Host must be a string")
+
+    if username is None:
+        username = ""
+
+    if not isinstance(username, str):
+        raise ValueError("Username must be a string")
+
+    if not isinstance(password, str):
+        raise ValueError("Password must be a string")
+
+    if port < 0:
+        raise ValueError(f"Invalid port: {port}")
+
+    if not isinstance(use_ssl, bool):
+        raise TypeError("use_ssl must be a boolean")
+
+    if not isinstance(storage_namespace, str):
+        raise ValueError("storage namespace must be a string")
+
+    # Right now, GCS stores all data in multiple hashes with keys prefixed by
+    # storage_namespace. So we only need to delete the specific key prefix to cleanup
+    # the cluster's data.
+    # Note this deletes all keys with prefix `RAY{key_prefix}@`, not `{key_prefix}`.
+    return del_key_prefix_from_storage(
+        host, port, username, password, use_ssl, storage_namespace
+    )

python/ray/_private/grpc_utils.py

@@ -0,0 +1,155 @@
+import os
+from concurrent import futures
+from typing import Any, Optional, Sequence, Tuple
+
+import grpc
+from grpc import aio as aiogrpc
+
+import ray
+from ray._common.tls_utils import load_certs_from_env
+from ray._private.authentication import authentication_utils
+
+
+def init_grpc_channel(
+    address: str,
+    options: Optional[Sequence[Tuple[str, Any]]] = None,
+    asynchronous: bool = False,
+    credentials: Optional[grpc.ChannelCredentials] = None,
+):
+    """Create a gRPC channel with authentication interceptors if token auth is enabled.
+
+    This function handles:
+    - TLS configuration via RAY_USE_TLS environment variable or custom credentials
+    - Authentication interceptors when token auth is enabled
+    - Keepalive settings from Ray config
+    - Both synchronous and asynchronous channels
+
+    Args:
+        address: The gRPC server address (host:port)
+        options: Optional gRPC channel options as sequence of (key, value) tuples
+        asynchronous: If True, create async channel; otherwise sync
+        credentials: Optional custom gRPC credentials for TLS. If provided, takes
+            precedence over RAY_USE_TLS environment variable.
+
+    Returns:
+        grpc.Channel or grpc.aio.Channel: Configured gRPC channel with interceptors
+    """
+    grpc_module = aiogrpc if asynchronous else grpc
+
+    options = options or []
+    options_dict = dict(options)
+    options_dict["grpc.keepalive_time_ms"] = options_dict.get(
+        "grpc.keepalive_time_ms", ray._config.grpc_client_keepalive_time_ms()
+    )
+    options_dict["grpc.keepalive_timeout_ms"] = options_dict.get(
+        "grpc.keepalive_timeout_ms", ray._config.grpc_client_keepalive_timeout_ms()
+    )
+    options = options_dict.items()
+
+    # Build interceptors list
+    interceptors = []
+    if authentication_utils.is_token_auth_enabled():
+        from ray._private.authentication.grpc_authentication_client_interceptor import (
+            SyncAuthenticationMetadataClientInterceptor,
+            get_async_auth_interceptors,
+        )
+
+        if asynchronous:
+            interceptors.extend(get_async_auth_interceptors())
+        else:
+            interceptors.append(SyncAuthenticationMetadataClientInterceptor())
+
+    # Determine channel type and credentials
+    if credentials is not None:
+        # Use provided custom credentials (takes precedence)
+        channel_creator = grpc_module.secure_channel
+        base_args = (address, credentials)
+    elif os.environ.get("RAY_USE_TLS", "0").lower() in ("1", "true"):
+        # Use TLS from environment variables
+        server_cert_chain, private_key, ca_cert = load_certs_from_env()
+        tls_credentials = grpc.ssl_channel_credentials(
+            certificate_chain=server_cert_chain,
+            private_key=private_key,
+            root_certificates=ca_cert,
+        )
+        channel_creator = grpc_module.secure_channel
+        base_args = (address, tls_credentials)
+    else:
+        # Insecure channel
+        channel_creator = grpc_module.insecure_channel
+        base_args = (address,)
+
+    # Create channel (async channels get interceptors in constructor, sync via intercept_channel)
+    if asynchronous:
+        channel = channel_creator(
+            *base_args, options=options, interceptors=interceptors
+        )
+    else:
+        channel = channel_creator(*base_args, options=options)
+        if interceptors:
+            channel = grpc.intercept_channel(channel, *interceptors)
+
+    return channel
+
+
+def create_grpc_server_with_interceptors(
+    max_workers: Optional[int] = None,
+    thread_name_prefix: str = "grpc_server",
+    options: Optional[Sequence[Tuple[str, Any]]] = None,
+    asynchronous: bool = False,
+):
+    """Create a gRPC server with authentication interceptors if token auth is enabled.
+
+    This function handles:
+    - Authentication interceptors when token auth is enabled
+    - Both synchronous and asynchronous servers
+    - Thread pool configuration for sync servers
+
+    Args:
+        max_workers: Max thread pool workers (required for sync, ignored for async)
+        thread_name_prefix: Thread name prefix for sync thread pool
+        options: Optional gRPC server options as sequence of (key, value) tuples
+        asynchronous: If True, create async server; otherwise sync
+
+    Returns:
+        grpc.Server or grpc.aio.Server: Configured gRPC server with interceptors
+    """
+    grpc_module = aiogrpc if asynchronous else grpc
+
+    # Build interceptors list
+    interceptors = []
+    if authentication_utils.is_token_auth_enabled():
+        if asynchronous:
+            from ray._private.authentication.grpc_authentication_server_interceptor import (
+                AsyncAuthenticationServerInterceptor,
+            )
+
+            interceptors.append(AsyncAuthenticationServerInterceptor())
+        else:
+            from ray._private.authentication.grpc_authentication_server_interceptor import (
+                SyncAuthenticationServerInterceptor,
+            )
+
+            interceptors.append(SyncAuthenticationServerInterceptor())
+
+    # Create server
+    if asynchronous:
+        server = grpc_module.server(
+            interceptors=interceptors if interceptors else None,
+            options=options,
+        )
+    else:
+        if max_workers is None:
+            raise ValueError("max_workers is required for synchronous gRPC servers")
+
+        executor = futures.ThreadPoolExecutor(
+            max_workers=max_workers,
+            thread_name_prefix=thread_name_prefix,
+        )
+        server = grpc_module.server(
+            executor,
+            interceptors=interceptors if interceptors else None,
+            options=options,
+        )
+
+    return server

python/ray/_private/inspect_util.py

@@ -0,0 +1,52 @@
+import inspect
+
+
+def is_cython(obj):
+    """Check if an object is a Cython function or method"""
+
+    # TODO(suo): We could split these into two functions, one for Cython
+    # functions and another for Cython methods.
+    # TODO(suo): There doesn't appear to be a Cython function 'type' we can
+    # check against via isinstance. Please correct me if I'm wrong.
+    def check_cython(x):
+        return type(x).__name__ == "cython_function_or_method"
+
+    # Check if function or method, respectively
+    return check_cython(obj) or (
+        hasattr(obj, "__func__") and check_cython(obj.__func__)
+    )
+
+
+def is_function_or_method(obj: object) -> bool:
+    """Check if an object is a function or method.
+
+    Args:
+        obj: The Python object in question.
+
+    Returns:
+        True if the object is an function or method.
+    """
+    return inspect.isfunction(obj) or inspect.ismethod(obj) or is_cython(obj)
+
+
+def is_class_method(f):
+    """Returns whether the given method is a class_method."""
+    return hasattr(f, "__self__") and f.__self__ is not None
+
+
+def is_static_method(cls: type, f_name: str) -> bool:
+    """Returns whether the class has a static method with the given name.
+
+    Args:
+        cls: The Python class (i.e. object of type `type`) to
+            search for the method in.
+        f_name: The name of the method to look up in this class
+            and check whether or not it is static.
+
+    Returns:
+        True if ``cls`` defines a static method named ``f_name``.
+    """
+    for base_cls in inspect.getmro(cls):
+        if f_name in base_cls.__dict__:
+            return isinstance(base_cls.__dict__[f_name], staticmethod)
+    return False

python/ray/_private/internal_api.py

@@ -0,0 +1,269 @@
+import warnings
+from typing import List, Tuple
+
+import ray
+import ray._private.profiling as profiling
+import ray._private.services as services
+import ray._private.worker
+from ray._common.network_utils import build_address
+from ray._private.state import GlobalState
+from ray._raylet import GcsClientOptions
+from ray.core.generated import common_pb2
+
+__all__ = ["free", "global_gc"]
+MAX_MESSAGE_LENGTH = ray._config.max_grpc_message_size()
+
+
+def global_gc():
+    """Trigger gc.collect() on all workers in the cluster."""
+
+    worker = ray._private.worker.global_worker
+    worker.core_worker.global_gc()
+
+
+def get_state_from_address(address=None):
+    address = services.canonicalize_bootstrap_address_or_die(address)
+
+    state = GlobalState()
+    options = GcsClientOptions.create(
+        address, None, allow_cluster_id_nil=True, fetch_cluster_id_if_nil=False
+    )
+    state._initialize_global_state(options)
+    return state
+
+
+def memory_summary(
+    address=None,
+    group_by="NODE_ADDRESS",
+    sort_by="OBJECT_SIZE",
+    units="B",
+    line_wrap=True,
+    stats_only=False,
+    num_entries=None,
+):
+    from ray.dashboard.memory_utils import memory_summary
+
+    state = get_state_from_address(address)
+    reply = get_memory_info_reply(state)
+
+    if stats_only:
+        return store_stats_summary(reply)
+    return memory_summary(
+        state, group_by, sort_by, line_wrap, units, num_entries
+    ) + store_stats_summary(reply)
+
+
+def get_memory_info_reply(
+    state, node_manager_address=None, node_manager_port=None, timeout_seconds=60.0
+):
+    """Returns global memory info."""
+
+    from ray._private.grpc_utils import init_grpc_channel
+    from ray.core.generated import node_manager_pb2, node_manager_pb2_grpc
+
+    # We can ask any Raylet for the global memory info, that Raylet internally
+    # asks all nodes in the cluster for memory stats.
+    if node_manager_address is None or node_manager_port is None:
+        # We should ask for a raylet that is alive.
+        raylet = None
+        for node in state.node_table():
+            if node["Alive"]:
+                raylet = node
+                break
+        assert raylet is not None, "Every raylet is dead"
+        raylet_address = build_address(
+            raylet["NodeManagerAddress"], raylet["NodeManagerPort"]
+        )
+    else:
+        raylet_address = build_address(node_manager_address, node_manager_port)
+
+    channel = init_grpc_channel(
+        raylet_address,
+        options=[
+            ("grpc.max_send_message_length", MAX_MESSAGE_LENGTH),
+            ("grpc.max_receive_message_length", MAX_MESSAGE_LENGTH),
+        ],
+    )
+
+    stub = node_manager_pb2_grpc.NodeManagerServiceStub(channel)
+    reply = stub.FormatGlobalMemoryInfo(
+        node_manager_pb2.FormatGlobalMemoryInfoRequest(include_memory_info=False),
+        timeout=timeout_seconds,
+    )
+    return reply
+
+
+def node_stats(
+    node_manager_address=None, node_manager_port=None, include_memory_info=True
+):
+    """Returns NodeStats object describing memory usage in the cluster."""
+
+    from ray._private.grpc_utils import init_grpc_channel
+    from ray.core.generated import node_manager_pb2, node_manager_pb2_grpc
+
+    # We can ask any Raylet for the global memory info.
+    assert node_manager_address is not None and node_manager_port is not None
+    raylet_address = build_address(node_manager_address, node_manager_port)
+    channel = init_grpc_channel(
+        raylet_address,
+        options=[
+            ("grpc.max_send_message_length", MAX_MESSAGE_LENGTH),
+            ("grpc.max_receive_message_length", MAX_MESSAGE_LENGTH),
+        ],
+    )
+
+    stub = node_manager_pb2_grpc.NodeManagerServiceStub(channel)
+    node_stats = stub.GetNodeStats(
+        node_manager_pb2.GetNodeStatsRequest(include_memory_info=include_memory_info),
+        timeout=30.0,
+    )
+    return node_stats
+
+
+def store_stats_summary(reply):
+    """Returns formatted string describing object store stats in all nodes."""
+    store_summary = "--- Aggregate object store stats across all nodes ---\n"
+    # TODO(ekl) it would be nice if we could provide a full memory usage
+    # breakdown by type (e.g., pinned by worker, primary, etc.)
+    store_summary += (
+        "Plasma memory usage {} MiB, {} objects, {}% full, {}% "
+        "needed\n".format(
+            int(reply.store_stats.object_store_bytes_used / (1024 * 1024)),
+            reply.store_stats.num_local_objects,
+            round(
+                100
+                * reply.store_stats.object_store_bytes_used
+                / reply.store_stats.object_store_bytes_avail,
+                2,
+            ),
+            round(
+                100
+                * reply.store_stats.object_store_bytes_primary_copy
+                / reply.store_stats.object_store_bytes_avail,
+                2,
+            ),
+        )
+    )
+    if reply.store_stats.object_store_bytes_fallback > 0:
+        store_summary += "Plasma filesystem mmap usage: {} MiB\n".format(
+            int(reply.store_stats.object_store_bytes_fallback / (1024 * 1024))
+        )
+    if reply.store_stats.spill_time_total_s > 0:
+        store_summary += (
+            "Spilled {} MiB, {} objects, avg write throughput {} MiB/s\n".format(
+                int(reply.store_stats.spilled_bytes_total / (1024 * 1024)),
+                reply.store_stats.spilled_objects_total,
+                int(
+                    reply.store_stats.spilled_bytes_total
+                    / (1024 * 1024)
+                    / reply.store_stats.spill_time_total_s
+                ),
+            )
+        )
+    if reply.store_stats.restore_time_total_s > 0:
+        store_summary += (
+            "Restored {} MiB, {} objects, avg read throughput {} MiB/s\n".format(
+                int(reply.store_stats.restored_bytes_total / (1024 * 1024)),
+                reply.store_stats.restored_objects_total,
+                int(
+                    reply.store_stats.restored_bytes_total
+                    / (1024 * 1024)
+                    / reply.store_stats.restore_time_total_s
+                ),
+            )
+        )
+    if reply.store_stats.object_pulls_queued:
+        store_summary += "Object fetches queued, waiting for available memory."
+
+    return store_summary
+
+
+def free(object_refs: list, local_only: bool = False):
+    """
+    DeprecationWarning: `free` is a deprecated API and will be
+    removed in a future version of Ray. If you have a use case
+    for this API, please open an issue on GitHub.
+
+    Free a list of IDs from the in-process and plasma object stores.
+
+    This function is a low-level API which should be used in restricted
+    scenarios.
+
+    If local_only is false, the request will be send to all object stores.
+
+    This method will not return any value to indicate whether the deletion is
+    successful or not. This function is an instruction to the object store. If
+    some of the objects are in use, the object stores will delete them later
+    when the ref count is down to 0.
+
+    Examples:
+
+        .. testcode::
+
+            import ray
+
+            @ray.remote
+            def f():
+                return 0
+
+            obj_ref = f.remote()
+            ray.get(obj_ref)  # wait for object to be created first
+            free([obj_ref])  # unpin & delete object globally
+
+    Args:
+        object_refs: List of object refs to delete.
+        local_only: Whether only deleting the list of objects in local
+            object store or all object stores.
+
+    Returns:
+        None.
+    """
+    warnings.warn(
+        "`free` is a deprecated API and will be removed in a future version of Ray. "
+        "If you have a use case for this API, please open an issue on GitHub.",
+        DeprecationWarning,
+    )
+    worker = ray._private.worker.global_worker
+
+    if isinstance(object_refs, ray.ObjectRef):
+        object_refs = [object_refs]
+
+    if not isinstance(object_refs, list):
+        raise TypeError(
+            "free() expects a list of ObjectRef, got {}".format(type(object_refs))
+        )
+
+    # Make sure that the values are object refs.
+    for object_ref in object_refs:
+        if not isinstance(object_ref, ray.ObjectRef):
+            raise TypeError(
+                "Attempting to call `free` on the value {}, "
+                "which is not an ray.ObjectRef.".format(object_ref)
+            )
+
+    worker.check_connected()
+    with profiling.profile("ray.free"):
+        if len(object_refs) == 0:
+            return
+
+        worker.core_worker.free_objects(object_refs, local_only)
+
+
+def get_local_ongoing_lineage_reconstruction_tasks() -> List[
+    Tuple[common_pb2.LineageReconstructionTask, int]
+]:
+    """Return the locally submitted ongoing retry tasks
+       triggered by lineage reconstruction.
+
+    NOTE: for the lineage reconstruction task status,
+    this method only returns the status known to the submitter
+    (i.e. it returns SUBMITTED_TO_WORKER instead of RUNNING).
+
+    The return type is a list of pairs where pair.first is the
+    lineage reconstruction task info and pair.second is the number
+    of ongoing lineage reconstruction tasks of this type.
+    """
+
+    worker = ray._private.worker.global_worker
+    worker.check_connected()
+    return worker.core_worker.get_local_ongoing_lineage_reconstruction_tasks()

python/ray/_private/label_utils.py

@@ -0,0 +1,230 @@
+import json
+import re
+from typing import (
+    Any,
+    Dict,
+    List,
+    Optional,
+)
+
+import yaml
+
+import ray._private.ray_constants as ray_constants
+
+# Regex patterns used to validate that labels conform to Kubernetes label syntax rules.
+# https://kubernetes.io/docs/concepts/overview/working-with-objects/labels/#syntax-and-character-set
+
+# Regex for mandatory name (DNS label) or value
+# Examples:
+#   Valid matches: "a", "label-name", "a-._b", "123", "this_is_a_valid_label"
+#   Invalid matches: "-abc", "abc-", "my@label"
+LABEL_REGEX = re.compile(r"([a-zA-Z0-9]([a-zA-Z0-9_.-]{0,61}[a-zA-Z0-9])?)")
+
+# Regex for optional prefix (DNS subdomain)
+# Examples:
+#   Valid matches: "abc", "sub.domain.example", "my-label", "123.456.789"
+#   Invalid matches: "-abc", "prefix_", "sub..domain", sub.$$.example
+LABEL_PREFIX_REGEX = rf"^({LABEL_REGEX.pattern}?(\.{LABEL_REGEX.pattern}?)*)$"
+
+# Supported operators for label selector conditions. Not (!) conditions are handled separately.
+LABEL_OPERATORS = {"in"}
+# Create a pattern string dynamically based on the LABEL_OPERATORS
+OPERATOR_PATTERN = "|".join([re.escape(operator) for operator in LABEL_OPERATORS])
+
+# Regex to match valid label selector operators and values
+# Examples:
+#   Valid matches: "spot", "!GPU", "213521", "in(A123, B456, C789)", "!in(spot, on-demand)", "valid-value"
+#   Invalid matches: "-spot", "spot_", "in()", "in(spot,", "in(H100, TPU!GPU)", "!!!in(H100, TPU)"
+LABEL_SELECTOR_REGEX = re.compile(
+    rf"^!?(?:{OPERATOR_PATTERN})?\({LABEL_REGEX.pattern}(?:, ?{LABEL_REGEX.pattern})*\)$|^!?{LABEL_REGEX.pattern}$"
+)
+
+
+def parse_node_labels_json(labels_json: str) -> Dict[str, str]:
+    labels = json.loads(labels_json)
+    if not isinstance(labels, dict):
+        raise ValueError("The format after deserialization is not a key-value pair map")
+    for key, value in labels.items():
+        if not isinstance(key, str):
+            raise ValueError("The key is not string type.")
+        if not isinstance(value, str):
+            raise ValueError(f'The value of the "{key}" is not string type')
+
+    # Validate parsed custom node labels don't begin with ray.io prefix
+    validate_node_labels(labels)
+
+    return labels
+
+
+def parse_node_labels_string(labels_str: str) -> Dict[str, str]:
+    labels = {}
+
+    # Remove surrounding quotes if they exist
+    if len(labels_str) > 1 and labels_str.startswith('"') and labels_str.endswith('"'):
+        labels_str = labels_str[1:-1]
+
+    if labels_str == "":
+        return labels
+
+    # Labels argument should consist of a string of key=value pairs
+    # separated by commas. Labels follow Kubernetes label syntax.
+    label_pairs = labels_str.split(",")
+    for pair in label_pairs:
+        # Split each pair by `=`
+        key_value = pair.split("=")
+        if len(key_value) != 2:
+            raise ValueError("Label string is not a key-value pair.")
+        key = key_value[0].strip()
+        value = key_value[1].strip()
+        labels[key] = value
+
+    # Validate parsed node labels follow expected Kubernetes label syntax
+    validate_node_label_syntax(labels)
+
+    return labels
+
+
+def parse_node_labels_from_yaml_file(path: str) -> Dict[str, str]:
+    if path == "":
+        return {}
+    with open(path, "r") as file:
+        # Expects valid YAML content
+        labels = yaml.safe_load(file)
+        if not isinstance(labels, dict):
+            raise ValueError(
+                "The format after deserialization is not a key-value pair map."
+            )
+        for key, value in labels.items():
+            if not isinstance(key, str):
+                raise ValueError("The key is not string type.")
+            if not isinstance(value, str):
+                raise ValueError(f'The value of "{key}" is not string type.')
+
+    # Validate parsed node labels follow expected Kubernetes label syntax
+    validate_node_label_syntax(labels)
+
+    return labels
+
+
+# TODO (ryanaoleary@): This function will be removed after the migration to the label
+# selector API from NodeLabelSchedulingPolicy is complete.
+def validate_node_labels(labels: Dict[str, str]):
+    if labels is None:
+        return
+    for key in labels.keys():
+        if key.startswith(ray_constants.RAY_DEFAULT_LABEL_KEYS_PREFIX):
+            raise ValueError(
+                f"Custom label keys `{key}` cannot start with the prefix "
+                f"`{ray_constants.RAY_DEFAULT_LABEL_KEYS_PREFIX}`. "
+                f"This is reserved for Ray defined labels."
+            )
+
+
+def validate_label_key(key: str) -> Optional[str]:
+    if "/" in key:
+        prefix, name = key.rsplit("/", 1)
+        if len(prefix) > 253 or not re.fullmatch(LABEL_PREFIX_REGEX, prefix):
+            return str(
+                f"Invalid label key prefix `{prefix}`. Prefix must be a series of DNS labels "
+                f"separated by dots (.), not longer than 253 characters in total."
+            )
+    else:
+        name = key
+    if len(name) > 63 or not re.fullmatch(LABEL_REGEX, name):
+        return str(
+            f"Invalid label key name `{name}`. Name must be 63 chars or less beginning and ending "
+            f"with an alphanumeric character ([a-z0-9A-Z]) with dashes (-), underscores (_),"
+            f"dots (.), and alphanumerics between."
+        )
+    return None
+
+
+def validate_label_value(value: str):
+    if value == "":
+        return
+    if len(value) > 63 or not re.fullmatch(LABEL_REGEX, value):
+        raise ValueError(
+            f"Invalid label key value `{value}`. Value must be 63 chars or less beginning and ending "
+            f"with an alphanumeric character ([a-z0-9A-Z]) with dashes (-), underscores (_),"
+            f"dots (.), and alphanumerics between."
+        )
+
+
+def validate_label_selector(label_selector: Optional[Dict[str, str]]) -> Optional[str]:
+    if label_selector is None:
+        return None
+
+    for key, value in label_selector.items():
+        possible_error_message = validate_label_key(key)
+        if possible_error_message:
+            return possible_error_message
+        if value is not None:
+            possible_error_message = validate_label_selector_value(value)
+            if possible_error_message:
+                return possible_error_message
+
+    return None
+
+
+def validate_label_selector_value(selector: str) -> Optional[str]:
+    if selector == "":
+        return None
+    if not re.fullmatch(LABEL_SELECTOR_REGEX, selector):
+        return str(
+            f"Invalid label selector value `{selector}`. The label selector value should contain optional operators and a label value. Supported operators are: ! and {LABEL_OPERATORS}. "
+            f"Value must be 63 chars or less beginning and ending "
+            f"with an alphanumeric character ([a-z0-9A-Z]) with dashes (-), underscores (_),"
+            f"dots (.), and alphanumerics between."
+        )
+
+    return None
+
+
+# TODO (ryanaoleary@): This function will replace `validate_node_labels` after
+# the migration from NodeLabelSchedulingPolicy to the Label Selector API is complete.
+def validate_node_label_syntax(labels: Dict[str, str]):
+    if labels is None:
+        return
+    for key, value in labels.items():
+        possible_error_message = validate_label_key(key)
+        if possible_error_message:
+            raise ValueError(possible_error_message)
+        if value is not None:
+            validate_label_value(value)
+
+
+def validate_fallback_strategy(
+    fallback_strategy: Optional[List[Dict[str, Any]]]
+) -> Optional[str]:
+    if fallback_strategy is None:
+        return None
+
+    # Supported options for `fallback_strategy` scheduling.
+    supported_options = {"label_selector"}
+
+    for strategy in fallback_strategy:
+        if not isinstance(strategy, dict):
+            return "Each element in fallback_strategy must be a dictionary."
+
+        if not strategy:
+            return "Empty dictionary found in `fallback_strategy`."
+
+        # Validate `fallback_strategy` only contains supported options.
+        for option in strategy:
+            if option not in supported_options:
+                return (
+                    f"Unsupported option found: '{option}'. "
+                    f"Only {list(supported_options)} is currently supported."
+                )
+
+        # Validate the 'label_selector' dictionary.
+        label_selector = strategy.get("label_selector")
+        if label_selector:
+            if not isinstance(label_selector, dict):
+                return 'The value of "label_selector" must be a dictionary.'
+
+            error_message = validate_label_selector(label_selector)
+            if error_message:
+                return error_message
+
+    return None

python/ray/_private/log.py

@@ -0,0 +1,150 @@
+import logging
+import threading
+import time
+from typing import Optional, Union
+
+INTERNAL_TIMESTAMP_LOG_KEY = "_ray_timestamp_ns"
+
+
+def _print_loggers():
+    """Print a formatted list of loggers and their handlers for debugging."""
+    loggers = {logging.root.name: logging.root}
+    loggers.update(dict(sorted(logging.root.manager.loggerDict.items())))
+    for name, logger in loggers.items():
+        if isinstance(logger, logging.Logger):
+            print(f"  {name}: disabled={logger.disabled}, propagate={logger.propagate}")
+            for handler in logger.handlers:
+                print(f"    {handler}")
+
+
+def clear_logger(logger: Union[str, logging.Logger]):
+    """Reset a logger, clearing its handlers and enabling propagation.
+
+    Args:
+        logger: Logger to be cleared
+    """
+    if isinstance(logger, str):
+        logger = logging.getLogger(logger)
+    logger.propagate = True
+    logger.handlers.clear()
+
+
+class PlainRayHandler(logging.StreamHandler):
+    """A plain log handler.
+
+    This handler writes to whatever sys.stderr points to at emit-time,
+    not at instantiation time. See docs for logging._StderrHandler.
+    """
+
+    def __init__(self):
+        super().__init__()
+        self.plain_handler = logging._StderrHandler()
+        self.plain_handler.level = self.level
+        self.plain_handler.formatter = logging.Formatter(fmt="%(message)s")
+
+    def emit(self, record: logging.LogRecord):
+        """Emit the log message.
+
+        If this is a worker, bypass fancy logging and just emit the log record.
+        If this is the driver, emit the message using the appropriate console handler.
+
+        Args:
+            record: Log record to be emitted
+        """
+        import ray
+
+        if (
+            hasattr(ray, "_private")
+            and hasattr(ray._private, "worker")
+            and ray._private.worker.global_worker.mode
+            == ray._private.worker.WORKER_MODE
+        ):
+            self.plain_handler.emit(record)
+        else:
+            logging._StderrHandler.emit(self, record)
+
+
+logger_initialized = False
+logging_config_lock = threading.Lock()
+
+
+def _setup_log_record_factory():
+    """Setup log record factory to add _ray_timestamp_ns to LogRecord."""
+    old_factory = logging.getLogRecordFactory()
+
+    def record_factory(*args, **kwargs):
+        record = old_factory(*args, **kwargs)
+        # Python logging module starts to use `time.time_ns()` to generate `created`
+        # from Python 3.13 to avoid the precision loss caused by the float type.
+        # Here, we generate the `created` for the LogRecord to support older Python
+        # versions.
+        ct = time.time_ns()
+        record.created = ct / 1e9
+
+        record.__dict__[INTERNAL_TIMESTAMP_LOG_KEY] = ct
+
+        return record
+
+    logging.setLogRecordFactory(record_factory)
+
+
+def generate_logging_config():
+    """Generate the default Ray logging configuration."""
+    with logging_config_lock:
+        global logger_initialized
+        if logger_initialized:
+            return
+        logger_initialized = True
+
+        plain_formatter = logging.Formatter(
+            "%(asctime)s\t%(levelname)s %(filename)s:%(lineno)s -- %(message)s"
+        )
+
+        default_handler = PlainRayHandler()
+        default_handler.setFormatter(plain_formatter)
+
+        ray_logger = logging.getLogger("ray")
+        ray_logger.setLevel(logging.INFO)
+        ray_logger.addHandler(default_handler)
+        ray_logger.propagate = False
+
+        # Special handling for ray.rllib: only warning-level messages passed through
+        # See https://github.com/ray-project/ray/pull/31858 for related PR
+        rllib_logger = logging.getLogger("ray.rllib")
+        rllib_logger.setLevel(logging.WARN)
+
+        # Set up the LogRecord factory.
+        _setup_log_record_factory()
+
+
+def setup_process_exit_logger(
+    process_exit_log_path: str,
+    level: int = logging.INFO,
+    formatter: Optional[logging.Formatter] = None,
+) -> logging.Logger:
+    """Configure and return the 'ray.process_exit' logger with a FileHandler."""
+    logger = logging.getLogger("ray.process_exit")
+    logger.setLevel(level)
+    logger.propagate = False
+
+    fh = logging.FileHandler(process_exit_log_path, encoding="utf-8")
+    if formatter is None:
+        formatter = logging.Formatter(
+            "%(asctime)s\t%(levelname)s %(filename)s:%(lineno)s -- %(message)s"
+        )
+    fh.setFormatter(formatter)
+    logger.addHandler(fh)
+    return logger
+
+
+def format_returncode(rc: Optional[int]) -> str:
+    """Return a consistent string for process return code."""
+    if rc is None:
+        return "None"
+    try:
+        rc_int = int(rc)
+    except Exception:
+        return str(rc)
+    if rc_int < 0:
+        return f"{rc_int} (signal {-rc_int})"
+    return f"{rc_int}"

python/ray/_private/log_monitor.py

@@ -0,0 +1,642 @@
+import argparse
+import errno
+import glob
+import logging
+import logging.handlers
+import os
+import platform
+import re
+import shutil
+import sys
+import time
+import traceback
+from typing import Callable, List, Optional, Set
+
+import ray._private.ray_constants as ray_constants
+import ray._private.utils
+from ray._private import logging_utils
+from ray._private.ray_logging import setup_component_logger
+from ray._raylet import GcsClient
+
+# Logger for this module. It should be configured at the entry point
+# into the program using Ray. Ray provides a default configuration at
+# entry/init points.
+logger = logging.getLogger(__name__)
+
+# The groups are job id, and pid.
+WORKER_LOG_PATTERN = re.compile(r".*worker.*-([0-9a-f]+)-(\d+)")
+# The groups are job id.
+RUNTIME_ENV_SETUP_PATTERN = re.compile(r".*runtime_env_setup-(\d+).log")
+# Log name update interval under pressure.
+# We need it because log name update is CPU intensive and uses 100%
+# of cpu when there are many log files.
+LOG_NAME_UPDATE_INTERVAL_S = float(os.getenv("LOG_NAME_UPDATE_INTERVAL_S", 0.5))
+# Once there are more files than this threshold,
+# log monitor start giving backpressure to lower cpu usages.
+RAY_LOG_MONITOR_MANY_FILES_THRESHOLD = int(
+    os.getenv("RAY_LOG_MONITOR_MANY_FILES_THRESHOLD", 1000)
+)
+RAY_RUNTIME_ENV_LOG_TO_DRIVER_ENABLED = int(
+    os.getenv("RAY_RUNTIME_ENV_LOG_TO_DRIVER_ENABLED", 0)
+)
+
+
+class LogFileInfo:
+    def __init__(
+        self,
+        filename=None,
+        size_when_last_opened=None,
+        file_position=None,
+        file_handle=None,
+        is_err_file=False,
+        job_id=None,
+        worker_pid=None,
+    ):
+        assert (
+            filename is not None
+            and size_when_last_opened is not None
+            and file_position is not None
+        )
+        self.filename = filename
+        self.size_when_last_opened = size_when_last_opened
+        self.file_position = file_position
+        self.file_handle = file_handle
+        self.is_err_file = is_err_file
+        self.job_id = job_id
+        self.worker_pid = worker_pid
+        self.actor_name = None
+        self.task_name = None
+
+    def reopen_if_necessary(self):
+        """Check if the file's inode has changed and reopen it if necessary.
+
+        There are a variety of reasons what we would logically consider a file
+        would have different inodes, such as log rotation or file syncing
+        semantics.
+
+        If the file is smaller than our recorded file position, we assume it has been
+        rotated and start reading it from the beginning.
+        """
+        try:
+            open_inode = None
+            if self.file_handle and not self.file_handle.closed:
+                open_inode = os.fstat(self.file_handle.fileno()).st_ino
+
+            new_statinfo = os.stat(self.filename)
+            if new_statinfo.st_ino != open_inode:
+                self.file_handle = open(self.filename, "rb")
+
+                # If the new file is smaller than the last read position, assume that
+                # the file has been rotated and read from the beginning. Else, continue
+                # from the existing file position.
+                if new_statinfo.st_size < self.file_position:
+                    self.file_position = 0
+
+                self.file_handle.seek(self.file_position)
+                self.size_when_last_opened = new_statinfo.st_size
+            else:
+                # Inode unchanged, but the file may have been truncated and
+                # rewritten in place. Compare against both the last read
+                # position and the last observed file size so we can detect
+                # rewrites even when the new content grows beyond the old
+                # position before the next poll.
+                if new_statinfo.st_size < max(
+                    self.file_position, self.size_when_last_opened
+                ):
+                    reopened_file = open(self.filename, "rb")
+                    self.file_handle.close()
+                    self.file_handle = reopened_file
+                    self.file_position = 0
+                    self.size_when_last_opened = new_statinfo.st_size
+        except Exception:
+            logger.debug(f"file no longer exists, skip re-opening of {self.filename}")
+
+    def __repr__(self):
+        return (
+            "FileInfo(\n"
+            f"\tfilename: {self.filename}\n"
+            f"\tsize_when_last_opened: {self.size_when_last_opened}\n"
+            f"\tfile_position: {self.file_position}\n"
+            f"\tfile_handle: {self.file_handle}\n"
+            f"\tis_err_file: {self.is_err_file}\n"
+            f"\tjob_id: {self.job_id}\n"
+            f"\tworker_pid: {self.worker_pid}\n"
+            f"\tactor_name: {self.actor_name}\n"
+            f"\ttask_name: {self.task_name}\n"
+            ")"
+        )
+
+
+class LogMonitor:
+    """A monitor process for monitoring Ray log files.
+
+    This class maintains a list of open files and a list of closed log files. We
+    can't simply leave all files open because we'll run out of file
+    descriptors.
+
+    The "run" method of this class will cycle between doing several things:
+    1. First, it will check if any new files have appeared in the log
+       directory. If so, they will be added to the list of closed files.
+    2. Then, if we are unable to open any new files, we will close all of the
+       files.
+    3. Then, we will open as many closed files as we can that may have new
+       lines (judged by an increase in file size since the last time the file
+       was opened).
+    4. Then we will loop through the open files and see if there are any new
+       lines in the file. If so, we will publish them to Ray pubsub.
+
+    Attributes:
+        ip: The hostname of this machine, for grouping log messages.
+        logs_dir: The directory that the log files are in.
+        log_filenames: This is the set of filenames of all files in
+            open_file_infos and closed_file_infos.
+        open_file_infos (list[LogFileInfo]): Info for all of the open files.
+        closed_file_infos (list[LogFileInfo]): Info for all of the closed
+            files.
+        can_open_more_files: True if we can still open more files and
+            false otherwise.
+        max_files_open: The maximum number of files that can be open.
+    """
+
+    def __init__(
+        self,
+        node_ip_address: str,
+        logs_dir: str,
+        gcs_client: GcsClient,
+        is_proc_alive_fn: Callable[[int], bool],
+        max_files_open: int = ray_constants.LOG_MONITOR_MAX_OPEN_FILES,
+        gcs_address: Optional[str] = None,
+    ):
+        """Initialize the log monitor object."""
+        self.ip: str = node_ip_address
+        self.logs_dir: str = logs_dir
+        self.gcs_client = gcs_client
+        self.log_filenames: Set[str] = set()
+        self.open_file_infos: List[LogFileInfo] = []
+        self.closed_file_infos: List[LogFileInfo] = []
+        self.can_open_more_files: bool = True
+        self.max_files_open: int = max_files_open
+        self.is_proc_alive_fn: Callable[[int], bool] = is_proc_alive_fn
+        self.is_autoscaler_v2: bool = self.get_is_autoscaler_v2(gcs_address)
+
+        logger.info(
+            f"Starting log monitor with [max open files={max_files_open}],"
+            f" [is_autoscaler_v2={self.is_autoscaler_v2}]"
+        )
+
+    def get_is_autoscaler_v2(self, gcs_address: Optional[str]) -> bool:
+        """Check if autoscaler v2 is enabled."""
+        if gcs_address is None:
+            return False
+
+        if not ray.experimental.internal_kv._internal_kv_initialized():
+            ray.experimental.internal_kv._initialize_internal_kv(self.gcs_client)
+        from ray.autoscaler.v2.utils import is_autoscaler_v2
+
+        return is_autoscaler_v2()
+
+    def _close_all_files(self):
+        """Close all open files (so that we can open more)."""
+        while len(self.open_file_infos) > 0:
+            file_info = self.open_file_infos.pop(0)
+            file_info.file_handle.close()
+            file_info.file_handle = None
+
+            proc_alive = True
+            # Test if the worker process that generated the log file
+            # is still alive. Only applies to worker processes.
+            # For all other system components, we always assume they are alive.
+            if (
+                file_info.worker_pid != "raylet"
+                and file_info.worker_pid != "gcs_server"
+                and file_info.worker_pid != "autoscaler"
+                and file_info.worker_pid != "runtime_env"
+                and file_info.worker_pid is not None
+            ):
+                assert not isinstance(file_info.worker_pid, str), (
+                    "PID should be an int type. " f"Given PID: {file_info.worker_pid}."
+                )
+                proc_alive = self.is_proc_alive_fn(file_info.worker_pid)
+                if not proc_alive:
+                    # The process is not alive any more, so move the log file
+                    # out of the log directory so glob.glob will not be slowed
+                    # by it.
+                    target = os.path.join(
+                        self.logs_dir, "old", os.path.basename(file_info.filename)
+                    )
+                    try:
+                        shutil.move(file_info.filename, target)
+                    except (IOError, OSError) as e:
+                        if e.errno == errno.ENOENT:
+                            logger.warning(
+                                f"Warning: The file {file_info.filename} was not found."
+                            )
+                        else:
+                            raise e
+
+            if proc_alive:
+                self.closed_file_infos.append(file_info)
+
+        self.can_open_more_files = True
+
+    def update_log_filenames(self):
+        """Update the list of log files to monitor."""
+        monitor_log_paths = []
+        # output of user code is written here
+        monitor_log_paths += glob.glob(
+            f"{self.logs_dir}/worker*[.out|.err]"
+        ) + glob.glob(f"{self.logs_dir}/java-worker*.log")
+        # segfaults and other serious errors are logged here
+        monitor_log_paths += glob.glob(f"{self.logs_dir}/raylet*.err")
+        # monitor logs are needed to report autoscaler events
+        # TODO(rickyx): remove this after migration.
+        if not self.is_autoscaler_v2:
+            # We publish monitor logs in autoscaler v1
+            monitor_log_paths += glob.glob(f"{self.logs_dir}/monitor.log")
+        else:
+            # We publish autoscaler events directly in autoscaler v2
+            monitor_log_paths += glob.glob(
+                f"{self.logs_dir}/events/event_AUTOSCALER.log"
+            )
+
+        # If gcs server restarts, there can be multiple log files.
+        monitor_log_paths += glob.glob(f"{self.logs_dir}/gcs_server*.err")
+
+        # Add libtpu logs if they exist in the Ray container.
+        tpu_log_dir = f"{self.logs_dir}/tpu_logs"
+        if os.path.isdir(tpu_log_dir):
+            monitor_log_paths += glob.glob(f"{self.logs_dir}/tpu_logs/**")
+
+        # runtime_env setup process is logged here
+        if RAY_RUNTIME_ENV_LOG_TO_DRIVER_ENABLED:
+            monitor_log_paths += glob.glob(f"{self.logs_dir}/runtime_env*.log")
+        for file_path in monitor_log_paths:
+            if os.path.isfile(file_path) and file_path not in self.log_filenames:
+                worker_match = WORKER_LOG_PATTERN.match(file_path)
+                if worker_match:
+                    worker_pid = int(worker_match.group(2))
+                else:
+                    worker_pid = None
+                job_id = None
+
+                # Perform existence check first because most file will not be
+                # including runtime_env. This saves some cpu cycle.
+                if "runtime_env" in file_path:
+                    runtime_env_job_match = RUNTIME_ENV_SETUP_PATTERN.match(file_path)
+                    if runtime_env_job_match:
+                        job_id = runtime_env_job_match.group(1)
+
+                is_err_file = file_path.endswith("err")
+
+                self.log_filenames.add(file_path)
+                self.closed_file_infos.append(
+                    LogFileInfo(
+                        filename=file_path,
+                        size_when_last_opened=0,
+                        file_position=0,
+                        file_handle=None,
+                        is_err_file=is_err_file,
+                        job_id=job_id,
+                        worker_pid=worker_pid,
+                    )
+                )
+                log_filename = os.path.basename(file_path)
+                logger.info(f"Beginning to track file {log_filename}")
+
+    def open_closed_files(self):
+        """Open some closed files if they may have new lines.
+
+        Opening more files may require us to close some of the already open
+        files.
+        """
+        if not self.can_open_more_files:
+            # If we can't open any more files. Close all of the files.
+            self._close_all_files()
+
+        files_with_no_updates = []
+        while len(self.closed_file_infos) > 0:
+            if len(self.open_file_infos) >= self.max_files_open:
+                self.can_open_more_files = False
+                break
+
+            file_info = self.closed_file_infos.pop(0)
+            assert file_info.file_handle is None
+            # Get the file size to see if it has gotten bigger since we last
+            # opened it.
+            try:
+                file_size = os.path.getsize(file_info.filename)
+            except (IOError, OSError) as e:
+                # Catch "file not found" errors.
+                if e.errno == errno.ENOENT:
+                    logger.warning(
+                        f"Warning: The file {file_info.filename} was not found."
+                    )
+                    self.log_filenames.remove(file_info.filename)
+                    continue
+                raise e
+
+            # If some new lines have been added to this file, try to reopen the
+            # file.
+            if file_size > file_info.size_when_last_opened:
+                try:
+                    f = open(file_info.filename, "rb")
+                except (IOError, OSError) as e:
+                    if e.errno == errno.ENOENT:
+                        logger.warning(
+                            f"Warning: The file {file_info.filename} was not found."
+                        )
+                        self.log_filenames.remove(file_info.filename)
+                        continue
+                    else:
+                        raise e
+
+                f.seek(file_info.file_position)
+                file_info.size_when_last_opened = file_size
+                file_info.file_handle = f
+                self.open_file_infos.append(file_info)
+            else:
+                files_with_no_updates.append(file_info)
+
+        if len(self.open_file_infos) >= self.max_files_open:
+            self.can_open_more_files = False
+        # Add the files with no changes back to the list of closed files.
+        self.closed_file_infos += files_with_no_updates
+
+    def check_log_files_and_publish_updates(self):
+        """Gets updates to the log files and publishes them.
+
+        Returns:
+            True if anything was published and false otherwise.
+        """
+        anything_published = False
+        lines_to_publish = []
+
+        def flush():
+            nonlocal lines_to_publish
+            nonlocal anything_published
+            if len(lines_to_publish) > 0:
+                data = {
+                    "ip": self.ip,
+                    "pid": file_info.worker_pid,
+                    "job": file_info.job_id,
+                    "is_err": file_info.is_err_file,
+                    "lines": lines_to_publish,
+                    "actor_name": file_info.actor_name,
+                    "task_name": file_info.task_name,
+                }
+                try:
+                    self.gcs_client.publish_logs(data)
+                except Exception:
+                    logger.exception(f"Failed to publish log messages {data}")
+                anything_published = True
+                lines_to_publish = []
+
+        for file_info in self.open_file_infos:
+            assert not file_info.file_handle.closed
+            file_info.reopen_if_necessary()
+
+            max_num_lines_to_read = ray_constants.LOG_MONITOR_NUM_LINES_TO_READ
+            for _ in range(max_num_lines_to_read):
+                try:
+                    next_line = file_info.file_handle.readline()
+                    # Replace any characters not in UTF-8 with
+                    # a replacement character, see
+                    # https://stackoverflow.com/a/38565489/10891801
+                    next_line = next_line.decode("utf-8", "replace")
+                    if next_line == "":
+                        break
+                    next_line = next_line.rstrip("\r\n")
+
+                    if next_line.startswith(ray_constants.LOG_PREFIX_ACTOR_NAME):
+                        flush()  # Possible change of task/actor name.
+                        file_info.actor_name = next_line.split(
+                            ray_constants.LOG_PREFIX_ACTOR_NAME, 1
+                        )[1]
+                        file_info.task_name = None
+                    elif next_line.startswith(ray_constants.LOG_PREFIX_TASK_NAME):
+                        flush()  # Possible change of task/actor name.
+                        file_info.task_name = next_line.split(
+                            ray_constants.LOG_PREFIX_TASK_NAME, 1
+                        )[1]
+                    elif next_line.startswith(ray_constants.LOG_PREFIX_JOB_ID):
+                        file_info.job_id = next_line.split(
+                            ray_constants.LOG_PREFIX_JOB_ID, 1
+                        )[1]
+                    elif next_line.startswith(
+                        "Windows fatal exception: access violation"
+                    ):
+                        # We are suppressing the
+                        # 'Windows fatal exception: access violation'
+                        # message on workers on Windows here.
+                        # As far as we know it is harmless,
+                        # but is frequently popping up if Python
+                        # functions are run inside the core
+                        # worker C extension. See the investigation in
+                        # github.com/ray-project/ray/issues/18944
+                        # Also skip the following line, which is an
+                        # empty line.
+                        file_info.file_handle.readline()
+                    else:
+                        lines_to_publish.append(next_line)
+                except Exception:
+                    logger.error(
+                        f"Error: Reading file: {file_info.filename}, "
+                        f"position: {file_info.file_info.file_handle.tell()} "
+                        "failed."
+                    )
+                    raise
+
+            if file_info.file_position == 0:
+                # make filename windows-agnostic
+                filename = file_info.filename.replace("\\", "/")
+                if "/raylet" in filename:
+                    file_info.worker_pid = "raylet"
+                elif "/gcs_server" in filename:
+                    file_info.worker_pid = "gcs_server"
+                elif "/monitor" in filename or "event_AUTOSCALER" in filename:
+                    file_info.worker_pid = "autoscaler"
+                elif "/runtime_env" in filename:
+                    file_info.worker_pid = "runtime_env"
+
+            # Record the current position in the file.
+            file_info.file_position = file_info.file_handle.tell()
+            flush()
+
+        return anything_published
+
+    def should_update_filenames(self, last_file_updated_time: float) -> bool:
+        """Return true if filenames should be updated.
+
+        This method is used to apply the backpressure on file updates because
+        that requires heavy glob operations which use lots of CPUs.
+
+        Args:
+            last_file_updated_time: The last time filenames are updated.
+
+        Returns:
+            True if filenames should be updated. False otherwise.
+        """
+        elapsed_seconds = float(time.time() - last_file_updated_time)
+        return (
+            len(self.log_filenames) < RAY_LOG_MONITOR_MANY_FILES_THRESHOLD
+            or elapsed_seconds > LOG_NAME_UPDATE_INTERVAL_S
+        )
+
+    def run(self):
+        """Run the log monitor.
+
+        This will scan the file system once every LOG_NAME_UPDATE_INTERVAL_S to
+        check if there are new log files to monitor. It will also publish new
+        log lines.
+        """
+        last_updated = time.time()
+        while True:
+            if self.should_update_filenames(last_updated):
+                self.update_log_filenames()
+                last_updated = time.time()
+
+            self.open_closed_files()
+            anything_published = self.check_log_files_and_publish_updates()
+            # If nothing was published, then wait a little bit before checking
+            # for logs to avoid using too much CPU.
+            if not anything_published:
+                time.sleep(0.1)
+
+
+def is_proc_alive(pid):
+    # Import locally to make sure the bundled version is used if needed
+    import psutil
+
+    try:
+        return psutil.Process(pid).is_running()
+    except psutil.NoSuchProcess:
+        # The process does not exist.
+        return False
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description=("Parse GCS server address for the log monitor to connect to.")
+    )
+    parser.add_argument(
+        "--gcs-address", required=False, type=str, help="The address (ip:port) of GCS."
+    )
+    parser.add_argument(
+        "--node-ip-address",
+        required=False,
+        type=str,
+        help="The IP address of the node.",
+    )
+    parser.add_argument(
+        "--logging-level",
+        required=False,
+        type=str,
+        default=ray_constants.LOGGER_LEVEL,
+        choices=ray_constants.LOGGER_LEVEL_CHOICES,
+        help=ray_constants.LOGGER_LEVEL_HELP,
+    )
+    parser.add_argument(
+        "--logging-format",
+        required=False,
+        type=str,
+        default=ray_constants.LOGGER_FORMAT,
+        help=ray_constants.LOGGER_FORMAT_HELP,
+    )
+    parser.add_argument(
+        "--logging-filename",
+        required=False,
+        type=str,
+        default=ray_constants.LOG_MONITOR_LOG_FILE_NAME,
+        help="Specify the name of log file, "
+        "log to stderr if set empty, default is "
+        f'"{ray_constants.LOG_MONITOR_LOG_FILE_NAME}"',
+    )
+    parser.add_argument(
+        "--session-dir",
+        required=True,
+        type=str,
+        help="Specify the path of the session directory used by Ray processes.",
+    )
+    parser.add_argument(
+        "--logs-dir",
+        required=True,
+        type=str,
+        help="Specify the path of the log directory used by Ray processes.",
+    )
+    parser.add_argument(
+        "--logging-rotate-bytes",
+        required=True,
+        type=int,
+        help="Specify the max bytes for rotating log file.",
+    )
+    parser.add_argument(
+        "--logging-rotate-backup-count",
+        required=True,
+        type=int,
+        help="Specify the backup count of rotated log file.",
+    )
+    parser.add_argument(
+        "--stdout-filepath",
+        required=False,
+        default="",
+        type=str,
+        help="The filepath to dump log monitor stdout.",
+    )
+    parser.add_argument(
+        "--stderr-filepath",
+        required=False,
+        default="",
+        type=str,
+        help="The filepath to dump log monitor stderr.",
+    )
+
+    args = parser.parse_args()
+
+    # Disable log rotation for windows platform.
+    logging_rotation_bytes = args.logging_rotate_bytes if sys.platform != "win32" else 0
+    logging_rotation_backup_count = (
+        args.logging_rotate_backup_count if sys.platform != "win32" else 1
+    )
+    logging_params = dict(
+        logging_level=args.logging_level,
+        logging_format=args.logging_format,
+        log_dir=args.logs_dir,
+        filename=args.logging_filename,
+        max_bytes=logging_rotation_bytes,
+        backup_count=logging_rotation_backup_count,
+    )
+    logger = setup_component_logger(**logging_params)
+
+    # Setup stdout/stderr redirect files if redirection enabled
+    logging_utils.redirect_stdout_stderr_if_needed(
+        args.stdout_filepath,
+        args.stderr_filepath,
+        logging_rotation_bytes,
+        logging_rotation_backup_count,
+    )
+
+    gcs_client = GcsClient(address=args.gcs_address)
+    log_monitor = LogMonitor(
+        args.node_ip_address,
+        args.logs_dir,
+        gcs_client,
+        is_proc_alive,
+        gcs_address=args.gcs_address,
+    )
+
+    try:
+        log_monitor.run()
+    except Exception as e:
+        # Something went wrong, so push an error to all drivers.
+        traceback_str = ray._private.utils.format_error_message(traceback.format_exc())
+        message = (
+            f"The log monitor on node {platform.node()} "
+            f"failed with the following error:\n{traceback_str}"
+        )
+        ray._private.utils.publish_error_to_driver(
+            ray_constants.LOG_MONITOR_DIED_ERROR,
+            message,
+            gcs_client=gcs_client,
+        )
+        logger.error(message)
+        raise e

python/ray/_private/logging_utils.py

@@ -0,0 +1,49 @@
+import sys
+
+from ray._private.utils import open_log
+from ray._raylet import StreamRedirector
+
+
+def redirect_stdout_stderr_if_needed(
+    stdout_filepath: str,
+    stderr_filepath: str,
+    rotation_bytes: int,
+    rotation_backup_count: int,
+):
+    """This function sets up redirection for stdout and stderr if needed, based on the given rotation parameters.
+
+    params:
+    stdout_filepath: the filepath stdout will be redirected to; if empty, stdout will not be redirected.
+    stderr_filepath: the filepath stderr will be redirected to; if empty, stderr will not be redirected.
+    rotation_bytes: number of bytes which triggers file rotation.
+    rotation_backup_count: the max size of rotation files.
+    """
+
+    # Setup redirection for stdout and stderr.
+    if stdout_filepath:
+        StreamRedirector.redirect_stdout(
+            stdout_filepath,
+            rotation_bytes,
+            rotation_backup_count,
+            False,  # tee_to_stdout
+            False,  # tee_to_stderr
+        )
+    if stderr_filepath:
+        StreamRedirector.redirect_stderr(
+            stderr_filepath,
+            rotation_bytes,
+            rotation_backup_count,
+            False,  # tee_to_stdout
+            False,  # tee_to_stderr
+        )
+
+    # Setup python system stdout/stderr.
+    stdout_fileno = sys.stdout.fileno()
+    stderr_fileno = sys.stderr.fileno()
+    # We also manually set sys.stdout and sys.stderr because that seems to
+    # have an effect on the output buffering. Without doing this, stdout
+    # and stderr are heavily buffered resulting in seemingly lost logging
+    # statements. We never want to close the stdout file descriptor, dup2 will
+    # close it when necessary and we don't want python's GC to close it.
+    sys.stdout = open_log(stdout_fileno, unbuffered=True, closefd=False)
+    sys.stderr = open_log(stderr_fileno, unbuffered=True, closefd=False)

python/ray/_private/memory_monitor.py

@@ -0,0 +1,165 @@
+import logging
+import os
+import platform
+import sys
+import time
+
+import ray  # noqa F401
+
+# Import ray before psutil will make sure we use psutil's bundled version
+from ray._common.utils import get_system_memory
+
+import psutil  # noqa E402
+
+logger = logging.getLogger(__name__)
+
+
+def get_rss(memory_info):
+    """Get the estimated non-shared memory usage from psutil memory_info."""
+    mem = memory_info.rss
+    # OSX doesn't have the shared attribute
+    if hasattr(memory_info, "shared"):
+        mem -= memory_info.shared
+    return mem
+
+
+def get_shared(virtual_memory):
+    """Get the estimated shared memory usage from psutil virtual mem info."""
+    # OSX doesn't have the shared attribute
+    if hasattr(virtual_memory, "shared"):
+        return virtual_memory.shared
+    else:
+        return 0
+
+
+def get_top_n_memory_usage(n: int = 10):
+    """Get the top n memory usage of the process
+
+    Params:
+        n: Number of top n process memory usage to return.
+    Returns:
+        (str) The formatted string of top n process memory usage.
+    """
+    proc_stats = []
+    for proc in psutil.process_iter(["memory_info", "cmdline"]):
+        try:
+            proc_stats.append(
+                (get_rss(proc.info["memory_info"]), proc.pid, proc.info["cmdline"])
+            )
+        except psutil.NoSuchProcess:
+            # We should skip the process that has exited. Refer this
+            # issue for more detail:
+            # https://github.com/ray-project/ray/issues/14929
+            continue
+        except psutil.AccessDenied:
+            # On MacOS, the proc_pidinfo call (used to get per-process
+            # memory info) fails with a permission denied error when used
+            # on a process that isn’t owned by the same user. For now, we
+            # drop the memory info of any such process, assuming that
+            # processes owned by other users (e.g. root) aren't Ray
+            # processes and will be of less interest when an OOM happens
+            # on a Ray node.
+            # See issue for more detail:
+            # https://github.com/ray-project/ray/issues/11845#issuecomment-849904019  # noqa: E501
+            continue
+    proc_str = "PID\tMEM\tCOMMAND"
+    for rss, pid, cmdline in sorted(proc_stats, reverse=True)[:n]:
+        proc_str += "\n{}\t{}GiB\t{}".format(
+            pid, round(rss / (1024**3), 2), " ".join(cmdline)[:100].strip()
+        )
+    return proc_str
+
+
+class RayOutOfMemoryError(Exception):
+    def __init__(self, msg):
+        Exception.__init__(self, msg)
+
+    @staticmethod
+    def get_message(used_gb, total_gb, threshold):
+        proc_str = get_top_n_memory_usage(n=10)
+        return (
+            "More than {}% of the memory on ".format(int(100 * threshold))
+            + "node {} is used ({} / {} GB). ".format(
+                platform.node(), round(used_gb, 2), round(total_gb, 2)
+            )
+            + f"The top 10 memory consumers are:\n\n{proc_str}"
+            + "\n\nIn addition, up to {} GiB of shared memory is ".format(
+                round(get_shared(psutil.virtual_memory()) / (1024**3), 2)
+            )
+            + "currently being used by the Ray object store.\n---\n"
+            "--- Tip: Use the `ray memory` command to list active "
+            "objects in the cluster.\n"
+            "--- To disable OOM exceptions, set "
+            "RAY_DISABLE_MEMORY_MONITOR=1.\n---\n"
+        )
+
+
+class MemoryMonitor:
+    """Helper class for raising errors on low memory.
+
+    This presents a much cleaner error message to users than what would happen
+    if we actually ran out of memory.
+
+    The monitor tries to use the cgroup memory limit and usage if it is set
+    and available so that it is more reasonable inside containers. Otherwise,
+    it uses `psutil` to check the memory usage.
+
+    The environment variable `RAY_MEMORY_MONITOR_ERROR_THRESHOLD` can be used
+    to overwrite the default error_threshold setting.
+
+    Used by test only. For production code use memory_monitor_interface.cc
+    """
+
+    def __init__(self, error_threshold=0.95, check_interval=1):
+        # Note: it takes ~50us to check the memory usage through psutil, so
+        # throttle this check at most once a second or so.
+        self.check_interval = check_interval
+        self.last_checked = 0
+        try:
+            self.error_threshold = float(
+                os.getenv("RAY_MEMORY_MONITOR_ERROR_THRESHOLD")
+            )
+        except (ValueError, TypeError):
+            self.error_threshold = error_threshold
+        # Try to read the cgroup memory limit if it is available.
+        try:
+            with open("/sys/fs/cgroup/memory/memory.limit_in_bytes", "rb") as f:
+                self.cgroup_memory_limit_gb = int(f.read()) / (1024**3)
+        except IOError:
+            self.cgroup_memory_limit_gb = sys.maxsize / (1024**3)
+        if not psutil:
+            logger.warning(
+                "WARNING: Not monitoring node memory since `psutil` "
+                "is not installed. Install this with "
+                "`pip install psutil` to enable "
+                "debugging of memory-related crashes."
+            )
+        self.disabled = (
+            "RAY_DEBUG_DISABLE_MEMORY_MONITOR" in os.environ
+            or "RAY_DISABLE_MEMORY_MONITOR" in os.environ
+        )
+
+    def get_memory_usage(self):
+        from ray._private.utils import get_used_memory
+
+        total_gb = get_system_memory() / (1024**3)
+        used_gb = get_used_memory() / (1024**3)
+
+        return used_gb, total_gb
+
+    def raise_if_low_memory(self):
+        if self.disabled:
+            return
+
+        if time.time() - self.last_checked > self.check_interval:
+            self.last_checked = time.time()
+            used_gb, total_gb = self.get_memory_usage()
+
+            if used_gb > total_gb * self.error_threshold:
+                raise RayOutOfMemoryError(
+                    RayOutOfMemoryError.get_message(
+                        used_gb, total_gb, self.error_threshold
+                    )
+                )
+            else:
+                logger.debug(f"Memory usage is {used_gb} / {total_gb}")

python/ray/_private/metrics_agent.py

@@ -0,0 +1,911 @@
+import json
+import logging
+import os
+import re
+import threading
+import time
+import traceback
+from collections import defaultdict, namedtuple
+from typing import Any, Dict, List, Set, Tuple, Union
+
+from opencensus.metrics.export.metric_descriptor import MetricDescriptorType
+from opencensus.metrics.export.value import ValueDouble
+from opencensus.stats import aggregation, measure as measure_module
+from opencensus.stats.aggregation_data import (
+    CountAggregationData,
+    DistributionAggregationData,
+    LastValueAggregationData,
+    SumAggregationData,
+)
+from opencensus.stats.base_exporter import StatsExporter
+from opencensus.stats.stats_recorder import StatsRecorder
+from opencensus.stats.view import View
+from opencensus.stats.view_manager import ViewManager
+from opencensus.tags import (
+    tag_key as tag_key_module,
+    tag_map as tag_map_module,
+    tag_value as tag_value_module,
+)
+from prometheus_client.core import (
+    CounterMetricFamily,
+    GaugeMetricFamily,
+    HistogramMetricFamily,
+    Metric as PrometheusMetric,
+)
+
+import ray
+from ray._common.network_utils import build_address
+from ray._private.ray_constants import env_bool
+from ray._private.telemetry.metric_cardinality import (
+    WORKER_ID_TAG_KEY,
+    MetricCardinality,
+)
+from ray._raylet import GcsClient
+from ray.core.generated.metrics_pb2 import Metric
+from ray.util.metrics import _is_invalid_metric_name
+
+logger = logging.getLogger(__name__)
+
+# Env var key to decide worker timeout.
+# If the worker doesn't report for more than
+# this time, we treat workers as dead.
+RAY_WORKER_TIMEOUT_S = "RAY_WORKER_TIMEOUT_S"
+GLOBAL_COMPONENT_KEY = "CORE"
+RE_NON_ALPHANUMS = re.compile(r"[^a-zA-Z0-9]")
+
+
+class Gauge(View):
+    """Gauge representation of opencensus view.
+
+    This class is used to collect process metrics from the reporter agent.
+    Cpp metrics should be collected in a different way.
+    """
+
+    def __init__(self, name, description, unit, tags: List[str]):
+        if _is_invalid_metric_name(name):
+            raise ValueError(
+                f"Invalid metric name: {name}. Metric will be discarded "
+                "and data will not be collected or published. "
+                "Metric names can only contain letters, numbers, _, and :. "
+                "Metric names cannot start with numbers."
+            )
+        self._measure = measure_module.MeasureInt(name, description, unit)
+        self._description = description
+        tags = [tag_key_module.TagKey(tag) for tag in tags]
+        self._view = View(
+            name, description, tags, self.measure, aggregation.LastValueAggregation()
+        )
+
+    @property
+    def measure(self):
+        return self._measure
+
+    @property
+    def view(self):
+        return self._view
+
+    @property
+    def name(self):
+        return self.measure.name
+
+    @property
+    def description(self):
+        return self._description
+
+
+Record = namedtuple("Record", ["gauge", "value", "tags"])
+
+
+def fix_grpc_metric(metric: Metric):
+    """
+    Fix the inbound `opencensus.proto.metrics.v1.Metric` protos to make it acceptable
+    by opencensus.stats.DistributionAggregationData.
+
+    - metric name: gRPC OpenCensus metrics have names with slashes and dots, e.g.
+    `grpc.io/client/server_latency`[1]. However Prometheus metric names only take
+    alphanums,underscores and colons[2]. We santinize the name by replacing non-alphanum
+    chars to underscore, like the official opencensus prometheus exporter[3].
+    - distribution bucket bounds: The Metric proto asks distribution bucket bounds to
+    be > 0 [4]. However, gRPC OpenCensus metrics have their first bucket bound == 0 [1].
+    This makes the `DistributionAggregationData` constructor to raise Exceptions. This
+    applies to all bytes and milliseconds (latencies). The fix: we update the initial 0
+    bounds to be 0.000_000_1. This will not affect the precision of the metrics, since
+    we don't expect any less-than-1 bytes, or less-than-1-nanosecond times.
+
+    [1] https://github.com/census-instrumentation/opencensus-specs/blob/master/stats/gRPC.md#units  # noqa: E501
+    [2] https://prometheus.io/docs/concepts/data_model/#metric-names-and-labels
+    [3] https://github.com/census-instrumentation/opencensus-cpp/blob/50eb5de762e5f87e206c011a4f930adb1a1775b1/opencensus/exporters/stats/prometheus/internal/prometheus_utils.cc#L39 # noqa: E501
+    [4] https://github.com/census-instrumentation/opencensus-proto/blob/master/src/opencensus/proto/metrics/v1/metrics.proto#L218 # noqa: E501
+    """
+
+    if not metric.metric_descriptor.name.startswith("grpc.io/"):
+        return
+
+    metric.metric_descriptor.name = RE_NON_ALPHANUMS.sub(
+        "_", metric.metric_descriptor.name
+    )
+
+    for series in metric.timeseries:
+        for point in series.points:
+            if point.HasField("distribution_value"):
+                dist_value = point.distribution_value
+                bucket_bounds = dist_value.bucket_options.explicit.bounds
+                if len(bucket_bounds) > 0 and bucket_bounds[0] == 0:
+                    bucket_bounds[0] = 0.000_000_1
+
+
+class OpencensusProxyMetric:
+    def __init__(self, name: str, desc: str, unit: str, label_keys: List[str]):
+        """Represents the OpenCensus metrics that will be proxy exported."""
+        self._name = name
+        self._desc = desc
+        self._unit = unit
+        # -- The label keys of the metric --
+        self._label_keys = label_keys
+        # -- The data that needs to be proxy exported --
+        # tuple of label values -> data (OpenCesnsus Aggregation data)
+        self._data = {}
+
+    @property
+    def name(self):
+        return self._name
+
+    @property
+    def desc(self):
+        return self._desc
+
+    @property
+    def unit(self):
+        return self._unit
+
+    @property
+    def label_keys(self):
+        return self._label_keys
+
+    @property
+    def data(self):
+        return self._data
+
+    def is_distribution_aggregation_data(self):
+        """Check if the metric is a distribution aggreation metric."""
+        return len(self._data) > 0 and isinstance(
+            next(iter(self._data.values())), DistributionAggregationData
+        )
+
+    def add_data(self, label_values: Tuple, data: Any):
+        """Add the data to the metric.
+
+        Args:
+            label_values: The label values of the metric.
+            data: The data to be added.
+        """
+        self._data[label_values] = data
+
+    def record(self, metric: Metric):
+        """Parse the Opencensus Protobuf and store the data.
+
+        The data can be accessed via `data` API once recorded.
+        """
+        timeseries = metric.timeseries
+
+        if len(timeseries) == 0:
+            return
+
+        # Create the aggregation and fill it in the our stats
+        for series in timeseries:
+            labels = tuple(val.value for val in series.label_values)
+
+            # Aggregate points.
+            for point in series.points:
+                if (
+                    metric.metric_descriptor.type
+                    == MetricDescriptorType.CUMULATIVE_INT64
+                ):
+                    data = CountAggregationData(point.int64_value)
+                elif (
+                    metric.metric_descriptor.type
+                    == MetricDescriptorType.CUMULATIVE_DOUBLE
+                ):
+                    data = SumAggregationData(ValueDouble, point.double_value)
+                elif metric.metric_descriptor.type == MetricDescriptorType.GAUGE_DOUBLE:
+                    data = LastValueAggregationData(ValueDouble, point.double_value)
+                elif (
+                    metric.metric_descriptor.type
+                    == MetricDescriptorType.CUMULATIVE_DISTRIBUTION
+                ):
+                    dist_value = point.distribution_value
+                    counts_per_bucket = [bucket.count for bucket in dist_value.buckets]
+                    bucket_bounds = dist_value.bucket_options.explicit.bounds
+                    data = DistributionAggregationData(
+                        dist_value.sum / dist_value.count,
+                        dist_value.count,
+                        dist_value.sum_of_squared_deviation,
+                        counts_per_bucket,
+                        bucket_bounds,
+                    )
+                else:
+                    raise ValueError("Summary is not supported")
+                self._data[labels] = data
+
+
+class Component:
+    def __init__(self, id: str):
+        """Represent a component that requests to proxy export metrics
+
+        Args:
+            id: Id of this component.
+        """
+        self.id = id
+        # -- The time this component reported its metrics last time --
+        # It is used to figure out if this component is stale.
+        self._last_reported_time = time.monotonic()
+        # -- Metrics requested to proxy export from this component --
+        # metrics_name (str) -> metric (OpencensusProxyMetric)
+        self._metrics = {}
+
+    @property
+    def metrics(self) -> Dict[str, OpencensusProxyMetric]:
+        """Return the metrics requested to proxy export from this component."""
+        return self._metrics
+
+    @property
+    def last_reported_time(self):
+        return self._last_reported_time
+
+    def record(self, metrics: List[Metric]):
+        """Parse the Opencensus protobuf and store metrics.
+
+        Metrics can be accessed via `metrics` API for proxy export.
+
+        Args:
+            metrics: A list of Opencensus protobuf for proxy export.
+        """
+        self._last_reported_time = time.monotonic()
+        for metric in metrics:
+            fix_grpc_metric(metric)
+            descriptor = metric.metric_descriptor
+            name = descriptor.name
+            label_keys = [label_key.key for label_key in descriptor.label_keys]
+
+            if name not in self._metrics:
+                self._metrics[name] = OpencensusProxyMetric(
+                    name, descriptor.description, descriptor.unit, label_keys
+                )
+            self._metrics[name].record(metric)
+
+
+class OpenCensusProxyCollector:
+    def __init__(self, namespace: str, component_timeout_s: int = 60):
+        """Prometheus collector implementation for opencensus proxy export.
+
+        Prometheus collector requires to implement `collect` which is
+        invoked whenever Prometheus queries the endpoint.
+
+        The class is thread-safe.
+
+        Args:
+            namespace: Prometheus namespace.
+            component_timeout_s: Number of seconds after which a component
+                without new reports is considered stale and its metrics are
+                no longer exported.
+        """
+        # -- Protect `self._components` --
+        self._components_lock = threading.Lock()
+        # -- Timeout until the component is marked as stale --
+        # Once the component is considered as stale,
+        # the metrics from that worker won't be exported.
+        self._component_timeout_s = component_timeout_s
+        # -- Prometheus namespace --
+        self._namespace = namespace
+        # -- Component that requests to proxy export metrics --
+        # Component means core worker, raylet, and GCS.
+        # component_id -> Components
+        # For workers, they contain worker ids.
+        # For other components (raylet, GCS),
+        # they contain the global key `GLOBAL_COMPONENT_KEY`.
+        self._components = {}
+        # Whether we want to export counter as gauge.
+        # This is for bug compatibility.
+        # See https://github.com/ray-project/ray/pull/43795.
+        self._export_counter_as_gauge = env_bool("RAY_EXPORT_COUNTER_AS_GAUGE", True)
+
+    def record(self, metrics: List[Metric], worker_id_hex: str = None):
+        """Record the metrics reported from the component that reports it.
+
+        Args:
+            metrics: A list of opencensus protobuf to proxy export metrics.
+            worker_id_hex: A worker id that reports these metrics.
+                If None, it means they are reported from Raylet or GCS.
+        """
+        key = GLOBAL_COMPONENT_KEY if not worker_id_hex else worker_id_hex
+        with self._components_lock:
+            if key not in self._components:
+                self._components[key] = Component(key)
+            self._components[key].record(metrics)
+
+    def clean_stale_components(self):
+        """Clean up stale components.
+
+        Stale means the component is dead or unresponsive.
+
+        Stale components won't be reported to Prometheus anymore.
+        """
+        with self._components_lock:
+            stale_components = []
+            stale_component_ids = []
+            for id, component in self._components.items():
+                elapsed = time.monotonic() - component.last_reported_time
+                if elapsed > self._component_timeout_s:
+                    stale_component_ids.append(id)
+                    logger.info(
+                        "Metrics from a worker ({}) is cleaned up due to "
+                        "timeout. Time since last report {}s".format(id, elapsed)
+                    )
+            for id in stale_component_ids:
+                stale_components.append(self._components.pop(id))
+            return stale_components
+
+    # TODO(sang): add start and end timestamp
+    def to_prometheus_metrics(
+        self,
+        metric_name: str,
+        metric_description: str,
+        label_keys: List[str],
+        metric_units: str,
+        label_values: Tuple[tag_value_module.TagValue],
+        agg_data: Any,
+        metrics_map: Dict[str, List[PrometheusMetric]],
+    ) -> None:
+        """to_metric translate the data that OpenCensus create
+        to Prometheus format, using Prometheus Metric object.
+
+        This method is from Opencensus Prometheus Exporter.
+
+        Args:
+            metric_name: Name of the metric.
+            metric_description: Description of the metric.
+            label_keys: The fixed label keys of the metric.
+            metric_units: Units of the metric.
+            label_values: The values of `label_keys`.
+            agg_data: `opencensus.stats.aggregation_data.AggregationData` object.
+                Aggregated data that needs to be converted as Prometheus samples
+            metrics_map: The converted metric is added to this map.
+
+        """
+        assert self._components_lock.locked()
+        metric_name = f"{self._namespace}_{metric_name}"
+        assert len(label_values) == len(label_keys), (label_values, label_keys)
+        # Prometheus requires that all tag values be strings hence
+        # the need to cast none to the empty string before exporting. See
+        # https://github.com/census-instrumentation/opencensus-python/issues/480
+        label_values = [tv if tv else "" for tv in label_values]
+
+        if isinstance(agg_data, CountAggregationData):
+            metrics = metrics_map.get(metric_name)
+            if not metrics:
+                metric = CounterMetricFamily(
+                    name=metric_name,
+                    documentation=metric_description,
+                    unit=metric_units,
+                    labels=label_keys,
+                )
+                metrics = [metric]
+                metrics_map[metric_name] = metrics
+            metrics[0].add_metric(labels=label_values, value=agg_data.count_data)
+            return
+
+        if isinstance(agg_data, SumAggregationData):
+            # This should be emitted as prometheus counter
+            # but we used to emit it as prometheus gauge.
+            # To keep the backward compatibility
+            # (changing from counter to gauge changes the metric name
+            # since prometheus client will add "_total" suffix to counter
+            # per OpenMetrics specification),
+            # we now emit both counter and gauge and in the
+            # next major Ray release (3.0) we can stop emitting gauge.
+            # This leaves people enough time to migrate their dashboards.
+            # See https://github.com/ray-project/ray/pull/43795.
+            metrics = metrics_map.get(metric_name)
+            if not metrics:
+                metric = CounterMetricFamily(
+                    name=metric_name,
+                    documentation=metric_description,
+                    labels=label_keys,
+                )
+                metrics = [metric]
+                metrics_map[metric_name] = metrics
+            metrics[0].add_metric(labels=label_values, value=agg_data.sum_data)
+
+            if not self._export_counter_as_gauge:
+                pass
+            elif metric_name.endswith("_total"):
+                # In this case, we only need to emit prometheus counter
+                # since for metric name already ends with _total suffix
+                # prometheus client won't change it
+                # so there is no backward compatibility issue.
+                # See https://prometheus.github.io/client_python/instrumenting/counter/
+                pass
+            else:
+                if len(metrics) == 1:
+                    metric = GaugeMetricFamily(
+                        name=metric_name,
+                        documentation=(
+                            f"(DEPRECATED, use {metric_name}_total metric instead) "
+                            f"{metric_description}"
+                        ),
+                        labels=label_keys,
+                    )
+                    metrics.append(metric)
+                assert len(metrics) == 2
+                metrics[1].add_metric(labels=label_values, value=agg_data.sum_data)
+            return
+
+        elif isinstance(agg_data, DistributionAggregationData):
+            assert agg_data.bounds == sorted(agg_data.bounds)
+            # buckets are a list of buckets. Each bucket is another list with
+            # a pair of bucket name and value, or a triple of bucket name,
+            # value, and exemplar. buckets need to be in order.
+            buckets = []
+            cum_count = 0  # Prometheus buckets expect cumulative count.
+            for ii, bound in enumerate(agg_data.bounds):
+                cum_count += agg_data.counts_per_bucket[ii]
+                bucket = [str(bound), cum_count]
+                buckets.append(bucket)
+            # Prometheus requires buckets to be sorted, and +Inf present.
+            # In OpenCensus we don't have +Inf in the bucket bonds so need to
+            # append it here.
+            buckets.append(["+Inf", agg_data.count_data])
+            metrics = metrics_map.get(metric_name)
+            if not metrics:
+                metric = HistogramMetricFamily(
+                    name=metric_name,
+                    documentation=metric_description,
+                    labels=label_keys,
+                )
+                metrics = [metric]
+                metrics_map[metric_name] = metrics
+            metrics[0].add_metric(
+                labels=label_values,
+                buckets=buckets,
+                sum_value=agg_data.sum,
+            )
+            return
+
+        elif isinstance(agg_data, LastValueAggregationData):
+            metrics = metrics_map.get(metric_name)
+            if not metrics:
+                metric = GaugeMetricFamily(
+                    name=metric_name,
+                    documentation=metric_description,
+                    labels=label_keys,
+                )
+                metrics = [metric]
+                metrics_map[metric_name] = metrics
+            metrics[0].add_metric(labels=label_values, value=agg_data.value)
+            return
+
+        else:
+            raise ValueError(f"unsupported aggregation type {type(agg_data)}")
+
+    def _aggregate_metric_data(
+        self,
+        datas: List[
+            Union[LastValueAggregationData, CountAggregationData, SumAggregationData]
+        ],
+    ) -> Union[LastValueAggregationData, CountAggregationData, SumAggregationData]:
+        assert len(datas) > 0
+        sample = datas[0]
+        if isinstance(sample, LastValueAggregationData):
+            return LastValueAggregationData(
+                ValueDouble, sum([data.value for data in datas])
+            )
+        if isinstance(sample, CountAggregationData):
+            return CountAggregationData(sum([data.count_data for data in datas]))
+        if isinstance(sample, SumAggregationData):
+            return SumAggregationData(
+                ValueDouble, sum([data.sum_data for data in datas])
+            )
+
+        raise ValueError(
+            f"Unsupported aggregation type {type(sample)}. "
+            "Supported types are "
+            f"{CountAggregationData}, {LastValueAggregationData}, {SumAggregationData}."
+            f"Got {datas}."
+        )
+
+    def _aggregate_with_recommended_cardinality(
+        self,
+        per_worker_metrics: List[OpencensusProxyMetric],
+    ) -> List[OpencensusProxyMetric]:
+        """Collect per-worker metrics, aggregate them into per-node metrics and convert
+        them to Prometheus format.
+
+        Args:
+            per_worker_metrics: A list of per-worker metrics for the same metric name.
+        Returns:
+            A list of per-node metrics for the same metric name, with the high
+            cardinality labels removed and the values aggregated.
+        """
+        metric = next(iter(per_worker_metrics), None)
+        if not metric or WORKER_ID_TAG_KEY not in metric.label_keys:
+            # No high cardinality labels, return the original metrics.
+            return per_worker_metrics
+
+        worker_id_label_index = metric.label_keys.index(WORKER_ID_TAG_KEY)
+        # map from the tuple of label values without worker_id to the list of per worker
+        # task metrics
+        label_value_to_data: Dict[
+            Tuple,
+            List[
+                Union[
+                    LastValueAggregationData,
+                    CountAggregationData,
+                    SumAggregationData,
+                ]
+            ],
+        ] = defaultdict(list)
+        for metric in per_worker_metrics:
+            for label_values, data in metric.data.items():
+                # remove the worker_id from the label values
+                label_value_to_data[
+                    label_values[:worker_id_label_index]
+                    + label_values[worker_id_label_index + 1 :]
+                ].append(data)
+
+        aggregated_metric = OpencensusProxyMetric(
+            name=metric.name,
+            desc=metric.desc,
+            unit=metric.unit,
+            # remove the worker_id from the label keys
+            label_keys=metric.label_keys[:worker_id_label_index]
+            + metric.label_keys[worker_id_label_index + 1 :],
+        )
+        for label_values, datas in label_value_to_data.items():
+            aggregated_metric.add_data(
+                label_values,
+                self._aggregate_metric_data(datas),
+            )
+
+        return [aggregated_metric]
+
+    def collect(self):  # pragma: NO COVER
+        """Collect fetches the statistics from OpenCensus
+        and delivers them as Prometheus Metrics.
+        Collect is invoked every time a prometheus.Gatherer is run
+        for example when the HTTP endpoint is invoked by Prometheus.
+
+        This method is required as a Prometheus Collector.
+        """
+        with self._components_lock:
+            # First construct the list of opencensus metrics to be converted to
+            # prometheus metrics.  For LEGACY cardinality level, this comprises all
+            # metrics from all components.  For RECOMMENDED cardinality level, we need
+            # to remove the high cardinality labels and aggreate the component metrics.
+            open_cencus_metrics: List[OpencensusProxyMetric] = []
+            # The metrics that need to be aggregated with recommended cardinality. Key
+            # is the metric name and value is the list of per-worker metrics.
+            to_lower_cardinality: Dict[str, List[OpencensusProxyMetric]] = defaultdict(
+                list
+            )
+            cardinality_level = MetricCardinality.get_cardinality_level()
+            for component in self._components.values():
+                for metric in component.metrics.values():
+                    if (
+                        cardinality_level == MetricCardinality.RECOMMENDED
+                        and not metric.is_distribution_aggregation_data()
+                    ):
+                        # We reduce the cardinality for all metrics except for histogram
+                        # metrics. The aggregation of histogram metrics from worker
+                        # level to node level is not well defined. In addition, we
+                        # currently have very few histogram metrics in Ray
+                        # so the impact of them is negligible.
+                        to_lower_cardinality[metric.name].append(metric)
+                    else:
+                        open_cencus_metrics.append(metric)
+            for per_worker_metrics in to_lower_cardinality.values():
+                open_cencus_metrics.extend(
+                    self._aggregate_with_recommended_cardinality(
+                        per_worker_metrics,
+                    )
+                )
+
+            prometheus_metrics_map = {}
+            for metric in open_cencus_metrics:
+                for label_values, data in metric.data.items():
+                    self.to_prometheus_metrics(
+                        metric.name,
+                        metric.desc,
+                        metric.label_keys,
+                        metric.unit,
+                        label_values,
+                        data,
+                        prometheus_metrics_map,
+                    )
+
+        for metrics in prometheus_metrics_map.values():
+            for metric in metrics:
+                yield metric
+
+
+class MetricsAgent:
+    def __init__(
+        self,
+        view_manager: ViewManager,
+        stats_recorder: StatsRecorder,
+        stats_exporter: StatsExporter = None,
+    ):
+        """A class to record and export metrics.
+
+        The class exports metrics in 2 different ways.
+        - Directly record and export metrics using OpenCensus.
+        - Proxy metrics from other core components
+            (e.g., raylet, GCS, core workers).
+
+        This class is thread-safe.
+        """
+        # Lock required because gRPC server uses
+        # multiple threads to process requests.
+        self._lock = threading.Lock()
+
+        #
+        # Opencensus components to record metrics.
+        #
+
+        # Managing views to export metrics
+        # If the stats_exporter is None, we disable all metrics export.
+        self.view_manager = view_manager
+        # A class that's used to record metrics
+        # emitted from the current process.
+        self.stats_recorder = stats_recorder
+        # A class to export metrics.
+        self.stats_exporter = stats_exporter
+        # -- A Prometheus custom collector to proxy export metrics --
+        # `None` if the prometheus server is not started.
+        self.proxy_exporter_collector = None
+
+        if self.stats_exporter is None:
+            # If the exporter is not given,
+            # we disable metrics collection.
+            self.view_manager = None
+        else:
+            self.view_manager.register_exporter(stats_exporter)
+            self.proxy_exporter_collector = OpenCensusProxyCollector(
+                self.stats_exporter.options.namespace,
+                component_timeout_s=int(os.getenv(RAY_WORKER_TIMEOUT_S, 120)),
+            )
+
+        # Registered view names.
+        self._registered_views: Set[str] = set()
+
+    def record_and_export(self, records: List[Record], global_tags=None):
+        """Directly record and export stats from the same process."""
+        global_tags = global_tags or {}
+        with self._lock:
+            if not self.view_manager:
+                return
+
+            for record in records:
+                gauge = record.gauge
+                value = record.value
+                tags = record.tags
+                try:
+                    self._record_gauge(gauge, value, {**tags, **global_tags})
+                except Exception as e:
+                    logger.error(
+                        f"Failed to record metric {gauge.name} with value {value} with tags {tags!r} and global tags {global_tags!r} due to: {e!r}"
+                    )
+
+    def _record_gauge(self, gauge: Gauge, value: float, tags: dict):
+        if gauge.name not in self._registered_views:
+            self.view_manager.register_view(gauge.view)
+            self._registered_views.add(gauge.name)
+        measurement_map = self.stats_recorder.new_measurement_map()
+        tag_map = tag_map_module.TagMap()
+        for key, tag_val in tags.items():
+            try:
+                tag_key = tag_key_module.TagKey(key)
+            except ValueError as e:
+                logger.error(
+                    f"Failed to create tag key {key} for metric {gauge.name} due to: {e!r}"
+                )
+                raise e
+            try:
+                tag_value = tag_value_module.TagValue(tag_val)
+            except ValueError as e:
+                logger.error(
+                    f"Failed to create tag value {tag_val} for key {key} for metric {gauge.name} due to: {e!r}"
+                )
+                raise e
+            tag_map.insert(tag_key, tag_value)
+        measurement_map.measure_float_put(gauge.measure, value)
+        # NOTE: When we record this metric, timestamp will be renewed.
+        measurement_map.record(tag_map)
+
+    def proxy_export_metrics(self, metrics: List[Metric], worker_id_hex: str = None):
+        """Proxy export metrics specified by a Opencensus Protobuf.
+
+        This API is used to export metrics emitted from
+        core components.
+
+        Args:
+            metrics: A list of protobuf Metric defined from OpenCensus.
+            worker_id_hex: The worker ID it proxies metrics export. None
+                if the metric is not from a worker (i.e., raylet, GCS).
+
+        Returns:
+            None.
+        """
+        with self._lock:
+            if not self.view_manager:
+                return
+
+        self._proxy_export_metrics(metrics, worker_id_hex)
+
+    def _proxy_export_metrics(self, metrics: List[Metric], worker_id_hex: str = None):
+        self.proxy_exporter_collector.record(metrics, worker_id_hex)
+
+    def clean_all_dead_worker_metrics(self):
+        """Clean dead worker's metrics.
+
+        Worker metrics are cleaned up and won't be exported once
+        it is considered as dead.
+
+        This method has to be periodically called by a caller.
+        """
+        with self._lock:
+            if not self.view_manager:
+                return
+
+        self.proxy_exporter_collector.clean_stale_components()
+
+
+class PrometheusServiceDiscoveryWriter(threading.Thread):
+    """A class to support Prometheus service discovery.
+
+    It supports file-based service discovery. Checkout
+    https://prometheus.io/docs/guides/file-sd/ for more details.
+
+    Args:
+        gcs_address: Gcs address for this cluster.
+        temp_dir: Temporary directory used by
+            Ray to store logs and metadata.
+        session_dir: Session-specific directory for this Ray session.
+            If provided, the discovery file is written here instead of
+            temp_dir, and a backward-compatible symlink is created at
+            the old temp_dir location.
+    """
+
+    def __init__(self, gcs_address: str, temp_dir: str, session_dir: str = None):
+        gcs_client_options = ray._raylet.GcsClientOptions.create(
+            gcs_address, None, allow_cluster_id_nil=True, fetch_cluster_id_if_nil=False
+        )
+        self.gcs_address = gcs_address
+
+        ray._private.state.state._initialize_global_state(gcs_client_options)
+        self.temp_dir = temp_dir
+        self.session_dir = session_dir if session_dir else temp_dir
+        # Tracks whether the backward-compatible symlink has been successfully created.
+        # This prevents recreating the symlink on every periodic write, avoiding
+        # unnecessary disk I/O, race conditions, and log flooding.
+        self._symlink_created = False
+        # If symlink creation fails (e.g., due to lack of permissions on Windows
+        # without developer mode, or restricted filesystems), this fallback flag is set
+        # to True. When True, the writer copies the file directly instead of symlinking.
+        self._use_fallback_copy = False
+        self.default_service_discovery_flush_period = 5
+
+        # The last service discovery content that PrometheusServiceDiscoveryWriter has seen
+        self.latest_service_discovery_content = []
+        self._content_lock = threading.RLock()
+
+        super().__init__()
+
+    def get_latest_service_discovery_content(self):
+        """Return the latest stored service discovery content."""
+        with self._content_lock:
+            return self.latest_service_discovery_content
+
+    def get_file_discovery_content(self):
+        """Return the content for Prometheus service discovery."""
+        nodes = ray.nodes()
+        metrics_export_addresses = [
+            build_address(node["NodeManagerAddress"], node["MetricsExportPort"])
+            for node in nodes
+            if node["alive"] is True
+        ]
+        gcs_client = GcsClient(address=self.gcs_address)
+        autoscaler_addr = gcs_client.internal_kv_get(b"AutoscalerMetricsAddress", None)
+        if autoscaler_addr:
+            metrics_export_addresses.append(autoscaler_addr.decode("utf-8"))
+        dashboard_addr = gcs_client.internal_kv_get(b"DashboardMetricsAddress", None)
+        if dashboard_addr:
+            metrics_export_addresses.append(dashboard_addr.decode("utf-8"))
+        content = [{"labels": {"job": "ray"}, "targets": metrics_export_addresses}]
+        with self._content_lock:
+            self.latest_service_discovery_content = content
+        return json.dumps(content)
+
+    def write(self):
+        # Write a file based on https://prometheus.io/docs/guides/file-sd/
+        # Write should be atomic. Otherwise, Prometheus raises an error that
+        # json file format is invalid because it reads a file when
+        # file is re-written. Note that Prometheus still works although we
+        # have this error.
+        temp_file_name = self.get_temp_file_name()
+        with open(temp_file_name, "w") as json_file:
+            json_file.write(self.get_file_discovery_content())
+        # NOTE: os.replace is atomic on both Linux and Windows, so we won't
+        # have race condition reading this file.
+        os.replace(temp_file_name, self.get_target_file_name())
+        # Create a backward-compatible symlink at the old temp_dir location
+        # so that existing Prometheus configurations that reference the old
+        # path continue to work. Verify if the symlink is still valid and
+        # pointing to the correct target, repairing it if it has been deleted or modified.
+        if self.session_dir != self.temp_dir:
+            legacy_path = os.path.join(
+                self.temp_dir,
+                ray._private.ray_constants.PROMETHEUS_SERVICE_DISCOVERY_FILE,
+            )
+            if self._symlink_created and not self._use_fallback_copy:
+                try:
+                    if not (
+                        os.path.islink(legacy_path)
+                        and os.readlink(legacy_path) == self.get_target_file_name()
+                    ):
+                        self._symlink_created = False
+                except OSError:
+                    self._symlink_created = False
+
+            if not self._symlink_created and not self._use_fallback_copy:
+                try:
+                    if os.path.islink(legacy_path) or os.path.exists(legacy_path):
+                        os.remove(legacy_path)
+                    os.symlink(self.get_target_file_name(), legacy_path)
+                    self._symlink_created = True
+                except OSError:
+                    logger.warning(
+                        f"Failed to create backward-compatible symlink at "
+                        f"{legacy_path}. Falling back to copying the service discovery file."
+                    )
+                    self._use_fallback_copy = True
+
+            if self._use_fallback_copy:
+                try:
+                    import shutil
+
+                    temp_legacy_path = legacy_path + ".tmp"
+                    shutil.copy(self.get_target_file_name(), temp_legacy_path)
+                    os.replace(temp_legacy_path, legacy_path)
+                except OSError as e:
+                    logger.warning(
+                        f"Failed to copy service discovery file to legacy path {legacy_path}: {e}"
+                    )
+
+    def get_target_file_name(self):
+        return os.path.join(
+            self.session_dir,
+            ray._private.ray_constants.PROMETHEUS_SERVICE_DISCOVERY_FILE,
+        )
+
+    def get_temp_file_name(self):
+        return os.path.join(
+            self.session_dir,
+            "{}_{}".format(
+                "tmp", ray._private.ray_constants.PROMETHEUS_SERVICE_DISCOVERY_FILE
+            ),
+        )
+
+    def run(self):
+        while True:
+            # This thread won't be broken by exceptions.
+            try:
+                self.write()
+            except Exception as e:
+                logger.warning(
+                    "Writing a service discovery file, {},failed.".format(
+                        self.get_target_file_name()
+                    )
+                )
+                logger.warning(traceback.format_exc())
+                logger.warning(f"Error message: {e}")
+            time.sleep(self.default_service_discovery_flush_period)

python/ray/_private/object_ref_generator.py

@@ -0,0 +1,385 @@
+from __future__ import annotations
+
+import asyncio
+import collections
+from typing import TYPE_CHECKING, Deque, Iterator, Optional
+
+import ray
+from ray.exceptions import GetTimeoutError, ObjectRefStreamEndOfStreamError
+from ray.util.annotations import DeveloperAPI, PublicAPI
+
+if TYPE_CHECKING:
+    from ray._private.worker import Worker
+
+
+@DeveloperAPI
+class DynamicObjectRefGenerator:
+    def __init__(self, refs: Deque["ray.ObjectRef"]):
+        # TODO(swang): As an optimization, can also store the generator
+        # ObjectID so that we don't need to keep individual ref counts for the
+        # inner ObjectRefs.
+        self._refs: Deque["ray.ObjectRef"] = collections.deque(refs)
+
+    def __iter__(self) -> Iterator("ray.ObjectRef"):
+        while self._refs:
+            yield self._refs.popleft()
+
+    def __len__(self) -> int:
+        return len(self._refs)
+
+
+@PublicAPI
+class ObjectRefGenerator:
+    """A generator to obtain object references from a task in a streaming manner.
+
+    The class is compatible with the Python generator and async generator interfaces.
+
+    The class is not thread-safe.
+
+    Do not initialize the class and create an instance directly.
+    The instance should be created by `.remote`.
+
+    .. testcode::
+
+        import ray
+        from typing import Generator
+
+        @ray.remote(num_returns="streaming")
+        def gen() -> Generator[int, None, None]:
+            for i in range(5):
+                yield i
+
+        obj_ref_gen: ray.ObjectRefGenerator = gen.remote()
+        for obj_ref in obj_ref_gen:
+            print("Got:", ray.get(obj_ref))
+    """
+
+    def __init__(self, generator_ref: "ray.ObjectRef", worker: "Worker"):
+        # The reference to a generator task.
+        self._generator_ref = generator_ref
+        # True if an exception has been raised from the generator task.
+        self._generator_task_raised = False
+        # Ray's worker class. ray._private.worker.global_worker
+        self.worker = worker
+        self.worker.check_connected()
+        assert hasattr(worker, "core_worker")
+
+    # Public APIs
+
+    def __iter__(self) -> "ObjectRefGenerator":
+        return self
+
+    def __next__(self) -> "ray.ObjectRef":
+        """Waits until a next ref is available and returns the object ref.
+
+        Raises StopIteration if there's no more objects
+        to generate.
+
+        The object ref will contain an exception if the task fails.
+        When the generator task returns N objects, it can return
+        up to N + 1 objects (if there's a system failure, the
+        last object will contain a system level exception).
+        """
+        return self._next_sync()
+
+    def send(self, value):
+        raise NotImplementedError("`gen.send` is not supported.")
+
+    def throw(self, value):
+        raise NotImplementedError("`gen.throw` is not supported.")
+
+    def close(self):
+        raise NotImplementedError("`gen.close` is not supported.")
+
+    def __aiter__(self) -> "ObjectRefGenerator":
+        return self
+
+    async def __anext__(self):
+        return await self._next_async()
+
+    async def asend(self, value):
+        raise NotImplementedError("`gen.asend` is not supported.")
+
+    async def athrow(self, value):
+        raise NotImplementedError("`gen.athrow` is not supported.")
+
+    async def aclose(self):
+        raise NotImplementedError("`gen.aclose` is not supported.")
+
+    def completed(self) -> "ray.ObjectRef":
+        """Returns an object ref that is ready when
+        a generator task completes.
+
+        If the task is failed unexpectedly (e.g., worker failure),
+        the `ray.get(gen.completed())` raises an exception.
+
+        The function returns immediately.
+        """
+        return self._generator_ref
+
+    def next_ready(self) -> bool:
+        """If True, it means the output of next(gen) is ready and
+        ray.get(next(gen)) returns immediately. False otherwise.
+
+        It returns False when next(gen) raises a StopIteration
+        (this condition should be checked using is_finished).
+
+        The function returns immediately.
+        """
+        self.worker.check_connected()
+        core_worker = self.worker.core_worker
+
+        if self.is_finished():
+            return False
+
+        expected_ref, is_ready = core_worker.peek_object_ref_stream(self._generator_ref)
+
+        if is_ready:
+            return True
+
+        ready, _ = ray.wait([expected_ref], timeout=0, fetch_local=False)
+        return len(ready) > 0
+
+    def is_finished(self) -> bool:
+        """If True, it means the generator is finished
+        and all output is taken. False otherwise.
+
+        When True, if next(gen) is called, it will raise StopIteration
+        or StopAsyncIteration
+
+        The function returns immediately.
+        """
+        self.worker.check_connected()
+        core_worker = self.worker.core_worker
+
+        finished = core_worker.is_object_ref_stream_finished(self._generator_ref)
+
+        if finished:
+            if self._generator_task_raised:
+                return True
+            else:
+                # We should try ray.get on a generator ref.
+                # If it raises an exception and
+                # _generator_task_raised is not set,
+                # this means the last ref is not taken yet.
+                try:
+                    ray.get(self._generator_ref)
+                except Exception:
+                    # The exception from _generator_ref
+                    # hasn't been taken yet.
+                    return False
+                else:
+                    return True
+        else:
+            return False
+
+    # Private APIs
+
+    def _get_next_object_id_binary(self) -> bytes:
+        """Return the binary id of the next object in the stream."""
+        self.worker.check_connected()
+        return self.worker.core_worker.peek_next_object_id_binary(self._generator_ref)
+
+    def _stream_exhausted(self) -> bool:
+        """Whether the stream's end-of-stream marker has been reached and all
+        yielded refs consumed.
+
+        Non-blocking, in-memory check (unlike ``is_finished``, this does not
+        ``ray.get`` the generator return object). When True, the only thing
+        left is the end-of-stream ``ray.get`` of the return object that
+        ``_next_sync`` performs to surface ``StopIteration`` / task errors.
+        """
+        self.worker.check_connected()
+        return self.worker.core_worker.is_object_ref_stream_finished(
+            self._generator_ref
+        )
+
+    def _get_next_ref_n(self, num_refs: int) -> list["ray.ObjectRef"]:
+        """Return the next num_refs references from a generator without consuming them.
+
+        The returned refs are not consumed; wait for the last one to become ready
+        before calling ``_consume_next_ref_n`` to advance the stream.
+
+        Args:
+            num_refs: The number of references to return, starting from the
+                current head of the stream. Must be positive.
+
+        Returns:
+            A list of exactly num_refs ObjectRefs corresponding to the next
+            results in the stream, starting from the current head.
+        """
+        if num_refs <= 0:
+            raise ValueError("num_refs must be positive")
+        self.worker.check_connected()
+        core_worker = self.worker.core_worker
+        return [
+            ref
+            for ref, _ in core_worker.peek_object_ref_stream_n(
+                self._generator_ref, num_refs
+            )
+        ]
+
+    def _consume_next_ref_n(self, num_refs: int) -> None:
+        """Consume (advance) the next num_refs references from a generator.
+
+        The caller must have waited for the last requested ref to become ready
+        (see ``_get_next_ref_n``); otherwise this raises ``ValueError`` instead
+        of silently advancing past unwritten objects.
+
+        If fewer than num_refs references remain before the end of the stream,
+        only the remaining references are consumed and the call returns
+        without raising.
+
+        Args:
+            num_refs: The number of references to consume, starting from the
+                current head of the stream. Must be positive.
+        """
+        if num_refs <= 0:
+            raise ValueError("num_refs must be positive")
+        self.worker.check_connected()
+        core_worker = self.worker.core_worker
+        try:
+            core_worker.try_read_next_object_ref_stream_n(self._generator_ref, num_refs)
+        except ObjectRefStreamEndOfStreamError:
+            return
+
+    def _next_sync(self, timeout_s: Optional[int | float] = None) -> "ray.ObjectRef":
+        """Waits for timeout_s and returns the object ref if available.
+
+        If an object is not available within the given timeout, it
+        returns a nil object reference.
+
+        If -1 timeout is provided, it means it waits infinitely.
+
+        Waiting is implemented as busy waiting.
+
+        Raises StopIteration if there's no more objects
+        to generate.
+
+        The object ref will contain an exception if the task fails.
+        When the generator task returns N objects, it can return
+        up to N + 1 objects (if there's a system failure, the
+        last object will contain a system level exception).
+
+        Args:
+            timeout_s: If the next object is not ready within
+                this timeout, it returns the nil object ref.
+
+        Returns:
+            ObjectRef corresponding to the next result in the stream.
+        """
+        core_worker = self.worker.core_worker
+
+        # Wait for the next ObjectRef to become ready.
+        expected_ref, is_ready = core_worker.peek_object_ref_stream(self._generator_ref)
+
+        if not is_ready:
+            _, unready = ray.wait([expected_ref], timeout=timeout_s, fetch_local=False)
+            if len(unready) > 0:
+                return ray.ObjectRef.nil()
+
+        try:
+            ref = core_worker.try_read_next_object_ref_stream(self._generator_ref)
+            assert not ref.is_nil()
+        except ObjectRefStreamEndOfStreamError:
+            if self._generator_task_raised:
+                # Exception has been returned.
+                raise StopIteration from None
+
+            try:
+                # The generator ref contains an exception
+                # if there's any failure. It contains nothing otherwise.
+                # In that case, it should raise StopIteration.
+                #
+                # Bound this get by the caller's timeout: the return object
+                # can be remote — or lost to a failed node and pending
+                # reconstruction — and an unbounded get would block the
+                # caller until it is restored (e.g. the Ray Data scheduling
+                # thread; a saturated cluster can then deadlock, since the
+                # blocked consumer is what releases backpressured CPUs).
+                # Per this method's contract, a timeout is reported as "no
+                # object ready yet" (nil ref) so the caller retries.
+                ray.get(
+                    self._generator_ref,
+                    timeout=(None if timeout_s is None or timeout_s < 0 else timeout_s),
+                )
+            except GetTimeoutError:
+                return ray.ObjectRef.nil()
+            except Exception:
+                self._generator_task_raised = True
+                return self._generator_ref
+            else:
+                # The task finished without an exception.
+                raise StopIteration from None
+        return ref
+
+    async def _suppress_exceptions(self, ref: "ray.ObjectRef") -> None:
+        # Wrap a streamed ref to avoid asyncio warnings about not retrieving
+        # the exception when we are just waiting for the ref to become ready.
+        # The exception will get returned (or warned) to the user once they
+        # actually await the ref.
+        try:
+            await ref
+        except Exception:
+            pass
+
+    async def _next_async(self, timeout_s: Optional[int | float] = None):
+        """Same API as _next_sync, but it is for async context."""
+        core_worker = self.worker.core_worker
+        ref, is_ready = core_worker.peek_object_ref_stream(self._generator_ref)
+
+        if not is_ready:
+            # TODO(swang): Avoid fetching the value.
+            _, unready = await asyncio.wait(
+                [asyncio.create_task(self._suppress_exceptions(ref))], timeout=timeout_s
+            )
+            if len(unready) > 0:
+                return ray.ObjectRef.nil()
+
+        try:
+            ref = core_worker.try_read_next_object_ref_stream(self._generator_ref)
+            assert not ref.is_nil()
+        except ObjectRefStreamEndOfStreamError:
+            if self._generator_task_raised:
+                # Exception has been returned.
+                raise StopAsyncIteration from None
+
+            try:
+                # The generator ref contains an exception
+                # if there's any failure. It contains nothing otherwise.
+                # In that case, it should raise StopAsyncIteration.
+                #
+                # Bound this await by the caller's timeout, mirroring
+                # _next_sync: the return object can be remote — or lost to a
+                # failed node and pending reconstruction — and an unbounded
+                # await would block the caller until it is restored. Per this
+                # method's contract, a timeout is reported as "no object
+                # ready yet" (nil ref) so the caller retries.
+                if timeout_s is None or timeout_s < 0:
+                    await self._generator_ref
+                else:
+                    await asyncio.wait_for(self._generator_ref, timeout=timeout_s)
+            except asyncio.TimeoutError:
+                return ray.ObjectRef.nil()
+            except Exception:
+                self._generator_task_raised = True
+                return self._generator_ref
+            else:
+                # Meaning the task succeed without failure raise StopAsyncIteration.
+                raise StopAsyncIteration from None
+
+        return ref
+
+    def __del__(self):
+        if hasattr(self.worker, "core_worker"):
+            # The stream is created when a task is first submitted.
+            # NOTE: This can be called multiple times
+            # because python doesn't guarantee __del__ is called
+            # only once.
+            self.worker.core_worker.async_delete_object_ref_stream(self._generator_ref)
+
+    def __getstate__(self):
+        raise TypeError(
+            "You cannot return or pass a generator to other task. "
+            "Serializing a ObjectRefGenerator is not allowed."
+        )

python/ray/_private/parameter.py

@@ -0,0 +1,482 @@
+import logging
+import os
+import pathlib
+from typing import Dict, List, Optional
+
+import ray._private.ray_constants as ray_constants
+from ray._common.network_utils import get_localhost_ip
+from ray._private.resource_isolation_config import ResourceIsolationConfig
+from ray._private.utils import get_ray_client_dependency_error
+
+logger = logging.getLogger(__name__)
+
+
+class RayParams:
+    """A class used to store the parameters used by Ray.
+
+    Attributes:
+        redis_address: The address of the Redis server to connect to. If
+            this address is not provided, then this command will start Redis, a
+            raylet, a plasma store, a plasma manager, and some workers.
+            It will also kill these processes when Python exits.
+        redis_port: The port that the primary Redis shard should listen
+            to. If None, then it will fall back to
+            ray._private.ray_constants.DEFAULT_PORT, or a random port if the default is
+            not available.
+        redis_shard_ports: A list of the ports to use for the non-primary Redis
+            shards. If None, then it will fall back to the ports right after
+            redis_port, or random ports if those are not available.
+        num_cpus: Number of CPUs to configure the raylet with.
+        num_gpus: Number of GPUs to configure the raylet with.
+        resources: A dictionary mapping the name of a resource to the quantity
+            of that resource available.
+        labels: The key-value labels of the node.
+        memory: Total available memory for workers requesting memory.
+        object_store_memory: The amount of memory (in bytes) to start the
+            object store with.
+        object_manager_port int: The port to use for the object manager.
+        node_manager_port: The port to use for the node manager.
+        gcs_server_port: The port to use for the GCS server.
+        node_ip_address: The IP address of the node that we are on.
+        min_worker_port: The lowest port number that workers will bind
+            on. If not set or set to 0, random ports will be chosen.
+        max_worker_port: The highest port number that workers will bind
+            on. If set, min_worker_port must also be set.
+        worker_port_list: An explicit list of ports to be used for
+            workers (comma-separated). Overrides min_worker_port and
+            max_worker_port.
+        ray_client_server_port: The port number the ray client server
+            will bind on. If not set, the ray client server will not
+            be started.
+        redirect_output: True if stdout and stderr for non-worker
+            processes should be redirected to files and false otherwise.
+        log_to_stderr: If set, controls whether non-worker stdout/stderr should be
+            written to stderr (True) or redirected to log files (False). This is the
+            preferred replacement for the deprecated `redirect_output` field.
+        external_addresses: The address of external Redis server to
+            connect to, in format of "ip1:port1,ip2:port2,...".  If this
+            address is provided, then ray won't start Redis instances in the
+            head node but use external Redis server(s) instead.
+        num_redis_shards: The number of Redis shards to start in addition to
+            the primary Redis shard.
+        redis_max_clients: If provided, attempt to configure Redis with this
+            maxclients number.
+        redis_username: Prevents external clients without the username
+            from connecting to Redis if provided.
+        redis_password: Prevents external clients without the password
+            from connecting to Redis if provided.
+        plasma_directory: A directory where the Plasma memory mapped files will
+            be created.
+        object_spilling_directory: The path to spill objects to. The same path will
+            be used as the object store fallback directory as well.
+        worker_path: The path of the source code that will be run by the
+            worker.
+        setup_worker_path: The path of the Python file that will set up
+            the environment for the worker process.
+        huge_pages: Boolean flag indicating whether to start the Object
+            Store with hugetlbfs support. Requires plasma_directory.
+        include_dashboard: Boolean flag indicating whether to start the web
+            UI, which displays the status of the Ray cluster. If this value is
+            None, then the UI will be started if the relevant dependencies are
+            present.
+        dashboard_host: The host to bind the dashboard server to. Use localhost
+            (127.0.0.1/::1) for local access only, or 0.0.0.0/:: for all
+            interfaces. Defaults to localhost.
+        dashboard_port: The port to bind the dashboard server to.
+            Defaults to 8265.
+        dashboard_agent_listen_port: The port for dashboard agents to listen on
+            for HTTP requests.
+            Defaults to 52365.
+        runtime_env_agent_port: The port at which the runtime env agent
+            listens to for HTTP.
+            Defaults to random available port.
+        plasma_store_socket_name: If provided, it specifies the socket
+            name used by the plasma store.
+        raylet_socket_name: If provided, it specifies the socket path
+            used by the raylet process.
+        temp_dir: If provided, it will specify the root temporary
+            directory for the Ray process. Must be an absolute path.
+        runtime_env_dir_name: If provided, specifies the directory that
+            will be created in the session dir to hold runtime_env files.
+        include_log_monitor: If True, then start a log monitor to
+            monitor the log files for all processes on this node and push their
+            contents to Redis.
+        autoscaling_config: path to autoscaling config file.
+        metrics_agent_port: The port to bind metrics agent.
+        metrics_export_port: The port at which metrics are exposed
+            through a Prometheus endpoint.
+        no_monitor: If True, the ray autoscaler monitor for this cluster
+            will not be started.
+        _system_config: Configuration for overriding RayConfig
+            defaults. Used to set system configuration and for experimental Ray
+            core feature flags.
+        enable_object_reconstruction: Enable plasma reconstruction on
+            failure.
+        ray_debugger_external: If true, make the Ray debugger for a
+            worker available externally to the node it is running on. This will
+            bind on 0.0.0.0 instead of localhost.
+        env_vars: Override environment variables for the raylet.
+        session_name: The current Ray session name.
+        webui: The url of the UI.
+        cluster_id: The cluster ID in hex string.
+        resource_isolation_config: settings for cgroupv2 based isolation of ray
+            system processes (defaults to no isolation if config not provided)
+        proxy_server_url: The proxy url to redirect dashboard backend request to.
+            By default, the dashboard requests will be directed to the Ray api server.
+            Ex: http://historyserver:8080
+    """
+
+    def __init__(
+        self,
+        redis_address: Optional[str] = None,
+        gcs_address: Optional[str] = None,
+        num_cpus: Optional[int] = None,
+        num_gpus: Optional[int] = None,
+        resources: Optional[Dict[str, float]] = None,
+        labels: Optional[Dict[str, str]] = None,
+        memory: Optional[float] = None,
+        object_store_memory: Optional[float] = None,
+        redis_port: Optional[int] = None,
+        redis_shard_ports: Optional[List[int]] = None,
+        object_manager_port: Optional[int] = None,
+        node_manager_port: int = 0,
+        gcs_server_port: Optional[int] = None,
+        node_ip_address: Optional[str] = None,
+        node_name: Optional[str] = None,
+        min_worker_port: Optional[int] = None,
+        max_worker_port: Optional[int] = None,
+        worker_port_list: Optional[List[int]] = None,
+        ray_client_server_port: Optional[int] = None,
+        redirect_output: Optional[bool] = None,
+        log_to_stderr: Optional[bool] = None,
+        external_addresses: Optional[List[str]] = None,
+        num_redis_shards: Optional[int] = None,
+        redis_max_clients: Optional[int] = None,
+        redis_username: Optional[str] = ray_constants.REDIS_DEFAULT_USERNAME,
+        redis_password: Optional[str] = ray_constants.REDIS_DEFAULT_PASSWORD,
+        plasma_directory: Optional[str] = None,
+        object_spilling_directory: Optional[str] = None,
+        worker_path: Optional[str] = None,
+        setup_worker_path: Optional[str] = None,
+        huge_pages: Optional[bool] = False,
+        include_dashboard: Optional[bool] = None,
+        dashboard_host: Optional[str] = get_localhost_ip(),
+        dashboard_port: Optional[bool] = ray_constants.DEFAULT_DASHBOARD_PORT,
+        dashboard_agent_listen_port: Optional[
+            int
+        ] = ray_constants.DEFAULT_DASHBOARD_AGENT_LISTEN_PORT,
+        runtime_env_agent_port: Optional[int] = None,
+        plasma_store_socket_name: Optional[str] = None,
+        raylet_socket_name: Optional[str] = None,
+        temp_dir: Optional[str] = None,
+        runtime_env_dir_name: Optional[str] = None,
+        include_log_monitor: Optional[str] = None,
+        autoscaling_config: Optional[str] = None,
+        ray_debugger_external: bool = False,
+        _system_config: Optional[Dict[str, str]] = None,
+        enable_object_reconstruction: Optional[bool] = False,
+        metrics_agent_port: Optional[int] = None,
+        metrics_export_port: Optional[int] = None,
+        tracing_startup_hook=None,
+        no_monitor: Optional[bool] = False,
+        env_vars: Optional[Dict[str, str]] = None,
+        session_name: Optional[str] = None,
+        webui: Optional[str] = None,
+        cluster_id: Optional[str] = None,
+        node_id: Optional[str] = None,
+        resource_isolation_config: Optional[ResourceIsolationConfig] = None,
+        proxy_server_url: Optional[str] = None,
+    ):
+        self.redis_address = redis_address
+        self.gcs_address = gcs_address
+        self.num_cpus = num_cpus
+        self.num_gpus = num_gpus
+        self.memory = memory
+        self.object_store_memory = object_store_memory
+        self.resources = resources
+        self.redis_port = redis_port
+        self.redis_shard_ports = redis_shard_ports
+        self.object_manager_port = object_manager_port
+        self.node_manager_port = node_manager_port
+        self.gcs_server_port = gcs_server_port
+        self.node_ip_address = node_ip_address
+        self.node_name = node_name
+        self.min_worker_port = min_worker_port
+        self.max_worker_port = max_worker_port
+        self.worker_port_list = worker_port_list
+        self.ray_client_server_port = ray_client_server_port
+        self.redirect_output = redirect_output
+        self.log_to_stderr = log_to_stderr
+        self.external_addresses = external_addresses
+        self.num_redis_shards = num_redis_shards
+        self.redis_max_clients = redis_max_clients
+        self.redis_username = redis_username
+        self.redis_password = redis_password
+        self.plasma_directory = plasma_directory
+        self.object_spilling_directory = object_spilling_directory
+        self.worker_path = worker_path
+        self.setup_worker_path = setup_worker_path
+        self.huge_pages = huge_pages
+        self.include_dashboard = include_dashboard
+        self.dashboard_host = dashboard_host
+        self.dashboard_port = dashboard_port
+        self.dashboard_agent_listen_port = dashboard_agent_listen_port
+        self.runtime_env_agent_port = runtime_env_agent_port
+        self.plasma_store_socket_name = plasma_store_socket_name
+        self.raylet_socket_name = raylet_socket_name
+        self.temp_dir = temp_dir
+        self.runtime_env_dir_name = (
+            runtime_env_dir_name or ray_constants.DEFAULT_RUNTIME_ENV_DIR_NAME
+        )
+        self.include_log_monitor = include_log_monitor
+        self.autoscaling_config = autoscaling_config
+        self.metrics_agent_port = metrics_agent_port
+        self.metrics_export_port = metrics_export_port
+        self.tracing_startup_hook = tracing_startup_hook
+        self.no_monitor = no_monitor
+        self.ray_debugger_external = ray_debugger_external
+        self.env_vars = env_vars
+        self.session_name = session_name
+        self.webui = webui
+        self._system_config = _system_config or {}
+        self._enable_object_reconstruction = enable_object_reconstruction
+        self.labels = labels
+        self._check_usage()
+        self.cluster_id = cluster_id
+        self.node_id = node_id
+        self.proxy_server_url = proxy_server_url
+
+        self.resource_isolation_config = resource_isolation_config
+        if not self.resource_isolation_config:
+            self.resource_isolation_config = ResourceIsolationConfig(
+                enable_resource_isolation=False
+            )
+
+        # Set the internal config options for object reconstruction.
+        if enable_object_reconstruction:
+            # Turn off object pinning.
+            if self._system_config is None:
+                self._system_config = dict()
+            print(self._system_config)
+            self._system_config["lineage_pinning_enabled"] = True
+
+    def update(self, **kwargs):
+        """Update the settings according to the keyword arguments.
+
+        Args:
+            kwargs: The keyword arguments to set corresponding fields.
+        """
+        for arg in kwargs:
+            if hasattr(self, arg):
+                setattr(self, arg, kwargs[arg])
+            else:
+                raise ValueError(f"Invalid RayParams parameter in update: {arg}")
+
+        self._check_usage()
+
+    def update_if_absent(self, **kwargs):
+        """Update the settings when the target fields are None.
+
+        Args:
+            kwargs: The keyword arguments to set corresponding fields.
+        """
+        for arg in kwargs:
+            if hasattr(self, arg):
+                if getattr(self, arg) is None:
+                    setattr(self, arg, kwargs[arg])
+            else:
+                raise ValueError(
+                    f"Invalid RayParams parameter in update_if_absent: {arg}"
+                )
+
+        self._check_usage()
+
+    def update_pre_selected_port(self):
+        """Update the pre-selected port information
+
+        Returns:
+            The dictionary mapping of component -> ports.
+        """
+
+        def wrap_port(port):
+            # 0 port means select a random port for the grpc server.
+            if port is None or port == 0:
+                return []
+            else:
+                return [port]
+
+        # Create a dictionary of the component -> port mapping.
+        pre_selected_ports = {
+            "gcs": wrap_port(self.redis_port),
+            "object_manager": wrap_port(self.object_manager_port),
+            "node_manager": wrap_port(self.node_manager_port),
+            "gcs_server": wrap_port(self.gcs_server_port),
+            "client_server": wrap_port(self.ray_client_server_port),
+            "dashboard": wrap_port(self.dashboard_port),
+            "dashboard_agent_grpc": wrap_port(self.metrics_agent_port),
+            "dashboard_agent_http": wrap_port(self.dashboard_agent_listen_port),
+            "runtime_env_agent": wrap_port(self.runtime_env_agent_port),
+            "metrics_export": wrap_port(self.metrics_export_port),
+        }
+        redis_shard_ports = self.redis_shard_ports
+        if redis_shard_ports is None:
+            redis_shard_ports = []
+        pre_selected_ports["redis_shards"] = redis_shard_ports
+        if self.worker_port_list is None:
+            if self.min_worker_port is not None and self.max_worker_port is not None:
+                pre_selected_ports["worker_ports"] = list(
+                    range(self.min_worker_port, self.max_worker_port + 1)
+                )
+            else:
+                # The dict is not updated when it requires random ports.
+                pre_selected_ports["worker_ports"] = []
+        else:
+            pre_selected_ports["worker_ports"] = [
+                int(port) for port in self.worker_port_list.split(",")
+            ]
+
+        # Update the pre selected port set.
+        self.reserved_ports = set()
+        for comp, port_list in pre_selected_ports.items():
+            for port in port_list:
+                if port in self.reserved_ports:
+                    raise ValueError(
+                        f"Ray component {comp} is trying to use "
+                        f"a port number {port} that is used by other components.\n"
+                        f"Port information: {self._format_ports(pre_selected_ports)}\n"
+                        "If you allocate ports, please make sure the same port "
+                        "is not used by multiple components."
+                    )
+                self.reserved_ports.add(port)
+
+    def _check_usage(self):
+        if self.worker_port_list is not None:
+            for port_str in self.worker_port_list.split(","):
+                try:
+                    port = int(port_str)
+                except ValueError as e:
+                    raise ValueError(
+                        "worker_port_list must be a comma-separated "
+                        f"list of integers: {e}"
+                    ) from None
+
+                if port < 1024 or port > 65535:
+                    raise ValueError(
+                        "Ports in worker_port_list must be "
+                        f"between 1024 and 65535. Got: {port}"
+                    )
+
+        # Used primarily for testing.
+        if os.environ.get("RAY_USE_RANDOM_PORTS", False):
+            if self.min_worker_port is None and self.max_worker_port is None:
+                self.min_worker_port = 0
+                self.max_worker_port = 0
+
+        if self.min_worker_port is not None:
+            if self.min_worker_port != 0 and (
+                self.min_worker_port < 1024 or self.min_worker_port > 65535
+            ):
+                raise ValueError(
+                    "min_worker_port must be 0 or an integer between 1024 and 65535."
+                )
+
+        if self.max_worker_port is not None:
+            if self.min_worker_port is None:
+                raise ValueError(
+                    "If max_worker_port is set, min_worker_port must also be set."
+                )
+            elif self.max_worker_port != 0:
+                if self.max_worker_port < 1024 or self.max_worker_port > 65535:
+                    raise ValueError(
+                        "max_worker_port must be 0 or an integer between "
+                        "1024 and 65535."
+                    )
+                elif self.max_worker_port <= self.min_worker_port:
+                    raise ValueError(
+                        "max_worker_port must be higher than min_worker_port."
+                    )
+        if self.ray_client_server_port is not None:
+            if get_ray_client_dependency_error() is not None:
+                raise ValueError(
+                    "Ray Client requires pip package `ray[client]`. "
+                    "If you installed the minimal Ray (e.g. `pip install ray`), "
+                    "please reinstall by executing `pip install ray[client]`."
+                )
+            if (
+                self.ray_client_server_port < 1024
+                or self.ray_client_server_port > 65535
+            ):
+                raise ValueError(
+                    "ray_client_server_port must be an integer "
+                    "between 1024 and 65535."
+                )
+        if self.runtime_env_agent_port is not None:
+            if self.runtime_env_agent_port != 0 and (
+                self.runtime_env_agent_port < 1024
+                or self.runtime_env_agent_port > 65535
+            ):
+                raise ValueError(
+                    "runtime_env_agent_port must be 0 (auto-assign) or an integer "
+                    "between 1024 and 65535."
+                )
+
+        if self.resources is not None:
+
+            def build_error(resource, alternative):
+                return (
+                    f"{self.resources} -> `{resource}` cannot be a "
+                    "custom resource because it is one of the default resources "
+                    f"({ray_constants.DEFAULT_RESOURCES}). "
+                    f"Use `{alternative}` instead. For example, use `ray start "
+                    f"--{alternative.replace('_', '-')}=1` instead of "
+                    f"`ray start --resources={{'{resource}': 1}}`"
+                )
+
+            assert "CPU" not in self.resources, build_error("CPU", "num_cpus")
+            assert "GPU" not in self.resources, build_error("GPU", "num_gpus")
+            assert "memory" not in self.resources, build_error("memory", "memory")
+            assert "object_store_memory" not in self.resources, build_error(
+                "object_store_memory", "object_store_memory"
+            )
+
+        if self.redirect_output is not None:
+            raise DeprecationWarning("The redirect_output argument is deprecated.")
+
+        if self.temp_dir is not None and not os.path.isabs(self.temp_dir):
+            raise ValueError("temp_dir must be absolute path or None.")
+
+        if self.temp_dir is not None and os.getenv("VIRTUAL_ENV"):
+            is_relative = True
+            try:
+                (
+                    pathlib.Path(self.temp_dir)
+                    .resolve()
+                    .relative_to(pathlib.Path(os.getenv("VIRTUAL_ENV")).resolve())
+                )
+            except ValueError:
+                is_relative = False
+
+            if is_relative:
+                raise ValueError(
+                    "temp_dir must not be child directory of virtualenv root"
+                )
+
+    def _format_ports(self, pre_selected_ports):
+        """Format the pre-selected ports information to be more human-readable."""
+        ports = pre_selected_ports.copy()
+
+        for comp, port_list in ports.items():
+            if len(port_list) == 1:
+                ports[comp] = port_list[0]
+            elif len(port_list) == 0:
+                # Nothing is selected, meaning it will be randomly selected.
+                ports[comp] = "random"
+            elif comp == "worker_ports":
+                min_port = port_list[0]
+                max_port = port_list[len(port_list) - 1]
+                if len(port_list) < 50:
+                    port_range_str = str(port_list)
+                else:
+                    port_range_str = f"from {min_port} to {max_port}"
+                ports[comp] = f"{len(port_list)} ports {port_range_str}"
+        return ports

python/ray/_private/path_utils.py

@@ -0,0 +1,32 @@
+import pathlib
+import urllib
+
+"""Cross-platform utilities for manipulating paths and URIs.
+
+NOTE: All functions in this file must support POSIX and Windows.
+"""
+
+
+def is_path(path_or_uri: str) -> bool:
+    """Returns True if uri_or_path is a path and False otherwise.
+
+    Windows paths start with a drive name which can be interpreted as
+    a URI scheme by urlparse and thus needs to be treated differently
+    form POSIX paths.
+
+    E.g. Creating a directory returns the path 'C:\\Users\\mp5n6ul72w\\working_dir'
+    will have the scheme 'C:'.
+    """
+    if not isinstance(path_or_uri, str):
+        raise TypeError(f" path_or_uri must be a string, got {type(path_or_uri)}.")
+
+    parsed_path = pathlib.Path(path_or_uri)
+    parsed_uri = urllib.parse.urlparse(path_or_uri)
+
+    if isinstance(parsed_path, pathlib.PurePosixPath):
+        return not parsed_uri.scheme
+    elif isinstance(parsed_path, pathlib.PureWindowsPath):
+        return parsed_uri.scheme == parsed_path.drive.strip(":").lower()
+    else:
+        # this should never happen.
+        raise TypeError(f"Unsupported path type: {type(parsed_path).__name__}")

python/ray/_private/process_watcher.py

@@ -0,0 +1,197 @@
+import asyncio
+import io
+import logging
+import os
+import sys
+from concurrent.futures import ThreadPoolExecutor
+
+import ray
+import ray._private.ray_constants as ray_constants
+import ray.dashboard.consts as dashboard_consts
+from ray._common.utils import run_background_task
+from ray._raylet import GcsClient
+from ray.dashboard.consts import _PARENT_DEATH_THREASHOLD
+
+# Import psutil after ray so the packaged version is used.
+import psutil
+
+logger = logging.getLogger(__name__)
+
+# TODO: move all consts from dashboard_consts to ray_constants and rename to remove
+# DASHBOARD_ prefixes.
+
+# Publishes at most this number of lines of Raylet logs, when the Raylet dies
+# unexpectedly.
+_RAYLET_LOG_MAX_PUBLISH_LINES = 20
+
+# Reads at most this amount of Raylet logs from the tail, for publishing and
+# checking if the Raylet was terminated gracefully.
+_RAYLET_LOG_MAX_TAIL_SIZE = 1 * 1024**2
+
+try:
+    create_task = asyncio.create_task
+except AttributeError:
+    create_task = asyncio.ensure_future
+
+
+def get_raylet_pid():
+    # TODO(edoakes): RAY_RAYLET_PID isn't properly set on Windows. This is
+    # only used for fate-sharing with the raylet and we need a different
+    # fate-sharing mechanism for Windows anyways.
+    if sys.platform in ["win32", "cygwin"]:
+        return None
+    raylet_pid = int(os.environ["RAY_RAYLET_PID"])
+    assert raylet_pid > 0
+    logger.info("raylet pid is %s", raylet_pid)
+    return raylet_pid
+
+
+def create_check_raylet_task(log_dir, gcs_client, parent_dead_callback, loop):
+    """
+    Creates an asyncio task to periodically check if the raylet process is still
+    running. If raylet is dead for _PARENT_DEATH_THREASHOLD (5) times, prepare to exit
+    as follows:
+
+    - Write logs about whether the raylet exit is graceful, by looking into the raylet
+    log and search for term "SIGTERM",
+    - Flush the logs via GcsClient,
+    - Exit.
+    """
+    if sys.platform in ["win32", "cygwin"]:
+        raise RuntimeError("can't check raylet process in Windows.")
+    raylet_pid = get_raylet_pid()
+
+    if dashboard_consts.PARENT_HEALTH_CHECK_BY_PIPE:
+        logger.info("check_parent_via_pipe")
+        check_parent_task = _check_parent_via_pipe(
+            log_dir, gcs_client, loop, parent_dead_callback
+        )
+    else:
+        logger.info("_check_parent")
+        check_parent_task = _check_parent(
+            raylet_pid, log_dir, gcs_client, parent_dead_callback
+        )
+
+    return run_background_task(check_parent_task)
+
+
+def report_raylet_error_logs(log_dir: str, gcs_client: GcsClient):
+    log_path = os.path.join(log_dir, "raylet.out")
+    error = False
+    msg = "Raylet is terminated. "
+    try:
+        with open(log_path, "r", encoding="utf-8") as f:
+            # Seek to _RAYLET_LOG_MAX_TAIL_SIZE from the end if the
+            # file is larger than that.
+            f.seek(0, io.SEEK_END)
+            pos = max(0, f.tell() - _RAYLET_LOG_MAX_TAIL_SIZE)
+            f.seek(pos, io.SEEK_SET)
+            # Read remaining logs by lines.
+            raylet_logs = f.readlines()
+            # Assume the SIGTERM message must exist within the last
+            # _RAYLET_LOG_MAX_TAIL_SIZE of the log file.
+            if any("Raylet received SIGTERM" in line for line in raylet_logs):
+                msg += "Termination is graceful."
+                logger.info(msg)
+            else:
+                msg += (
+                    "Termination is unexpected. Possible reasons "
+                    "include: (1) SIGKILL by the user or system "
+                    "OOM killer, (2) Invalid memory access from "
+                    "Raylet causing SIGSEGV or SIGBUS, "
+                    "(3) Other termination signals. "
+                    f"Last {_RAYLET_LOG_MAX_PUBLISH_LINES} lines "
+                    "of the Raylet logs:\n"
+                )
+                msg += "    " + "    ".join(
+                    raylet_logs[-_RAYLET_LOG_MAX_PUBLISH_LINES:]
+                )
+                error = True
+    except Exception as e:
+        msg += f"Failed to read Raylet logs at {log_path}: {e}!"
+        logger.exception(msg)
+        error = True
+    if error:
+        logger.error(msg)
+        # TODO: switch to async if necessary.
+        ray._private.utils.publish_error_to_driver(
+            ray_constants.RAYLET_DIED_ERROR,
+            msg,
+            gcs_client=gcs_client,
+        )
+    else:
+        logger.info(msg)
+
+
+async def _check_parent_via_pipe(
+    log_dir: str, gcs_client: GcsClient, loop, parent_dead_callback
+):
+    while True:
+        try:
+            # Read input asynchronously.
+            # The parent (raylet) should have redirected its pipe
+            # to stdin. If we read 0 bytes from stdin, it means
+            # the process is dead.
+            with ThreadPoolExecutor(max_workers=1) as executor:
+                input_data = await loop.run_in_executor(
+                    executor, lambda: sys.stdin.readline()
+                )
+            if len(input_data) == 0:
+                # cannot read bytes from parent == parent is dead.
+                parent_dead_callback("_check_parent_via_pipe: The parent is dead.")
+                report_raylet_error_logs(log_dir, gcs_client)
+                sys.exit(0)
+        except Exception as e:
+            logger.exception(
+                "raylet health checking is failed. "
+                f"The agent process may leak. Exception: {e}"
+            )
+
+
+async def _check_parent(raylet_pid, log_dir, gcs_client, parent_dead_callback):
+    """Check if raylet is dead and fate-share if it is."""
+    try:
+        curr_proc = psutil.Process()
+        parent_death_cnt = 0
+        while True:
+            parent = curr_proc.parent()
+            # If the parent is dead, it is None.
+            parent_gone = parent is None
+            init_assigned_for_parent = False
+            parent_changed = False
+
+            if parent:
+                # Sometimes, the parent is changed to the `init` process.
+                # In this case, the parent.pid is 1.
+                init_assigned_for_parent = parent.pid == 1
+                # Sometimes, the parent is dead, and the pid is reused
+                # by other processes. In this case, this condition is triggered.
+                parent_changed = raylet_pid != parent.pid
+
+            if parent_gone or init_assigned_for_parent or parent_changed:
+                parent_death_cnt += 1
+                logger.warning(
+                    f"Raylet is considered dead {parent_death_cnt} X. "
+                    f"If it reaches to {_PARENT_DEATH_THREASHOLD}, the agent "
+                    f"will kill itself. Parent: {parent}, "
+                    f"parent_gone: {parent_gone}, "
+                    f"init_assigned_for_parent: {init_assigned_for_parent}, "
+                    f"parent_changed: {parent_changed}."
+                )
+                if parent_death_cnt < _PARENT_DEATH_THREASHOLD:
+                    await asyncio.sleep(
+                        dashboard_consts.DASHBOARD_AGENT_CHECK_PARENT_INTERVAL_S
+                    )
+                    continue
+
+                parent_dead_callback("_check_parent: The parent is dead.")
+                report_raylet_error_logs(log_dir, gcs_client)
+                sys.exit(0)
+            else:
+                parent_death_cnt = 0
+            await asyncio.sleep(
+                dashboard_consts.DASHBOARD_AGENT_CHECK_PARENT_INTERVAL_S
+            )
+    except Exception:
+        logger.exception("Failed to check parent PID, exiting.")
+        sys.exit(1)

python/ray/_private/profiling.py

@@ -0,0 +1,238 @@
+import json
+import os
+from collections import defaultdict
+from dataclasses import asdict, dataclass
+from typing import Dict, List, Optional, Union
+
+import ray
+
+
+class _NullLogSpan:
+    """A log span context manager that does nothing"""
+
+    def __enter__(self):
+        pass
+
+    def __exit__(self, type, value, tb):
+        pass
+
+
+PROFILING_ENABLED = "RAY_PROFILING" in os.environ
+NULL_LOG_SPAN = _NullLogSpan()
+
+# Colors are specified at
+# https://github.com/catapult-project/catapult/blob/master/tracing/tracing/base/color_scheme.html.  # noqa: E501
+_default_color_mapping = defaultdict(
+    lambda: "generic_work",
+    {
+        "worker_idle": "cq_build_abandoned",
+        "task": "rail_response",
+        "task:deserialize_arguments": "rail_load",
+        "task:execute": "rail_animation",
+        "task:store_outputs": "rail_idle",
+        "wait_for_function": "detailed_memory_dump",
+        "ray.get": "good",
+        "ray.put": "terrible",
+        "ray.wait": "vsync_highlight_color",
+        "submit_task": "background_memory_dump",
+        "fetch_and_run_function": "detailed_memory_dump",
+        "register_remote_function": "detailed_memory_dump",
+    },
+)
+
+
+@dataclass(init=True)
+class ChromeTracingCompleteEvent:
+    # https://docs.google.com/document/d/1CvAClvFfyA5R-PhYUmn5OOQtYMH4h6I0nSsKchNAySU/preview#heading=h.lpfof2aylapb # noqa
+    # The event categories. This is a comma separated list of categories
+    # for the event. The categories can be used to hide events in
+    # the Trace Viewer UI.
+    cat: str
+    # The string displayed on the event.
+    name: str
+    # The identifier for the group of rows that the event
+    # appears in.
+    pid: int
+    # The identifier for the row that the event appears in.
+    tid: int
+    # The start time in microseconds.
+    ts: int
+    # The duration in microseconds.
+    dur: int
+    # This is the name of the color to display the box in.
+    cname: str
+    # The extra user-defined data.
+    args: Dict[str, Union[str, int]]
+    # The event type (X means the complete event).
+    ph: str = "X"
+
+
+@dataclass(init=True)
+class ChromeTracingMetadataEvent:
+    # https://docs.google.com/document/d/1CvAClvFfyA5R-PhYUmn5OOQtYMH4h6I0nSsKchNAySU/preview#bookmark=id.iycbnb4z7i9g # noqa
+    name: str
+    # Metadata arguments. E.g., name: <metadata_name>
+    args: Dict[str, str]
+    # The process id of this event. In Ray, pid indicates the node.
+    pid: int
+    # The thread id of this event. In Ray, tid indicates each worker.
+    tid: int = None
+    # M means the metadata event.
+    ph: str = "M"
+
+
+def profile(event_type: str, extra_data: Optional[Dict[str, str]] = None):
+    """Profile a span of time so that it appears in the timeline visualization.
+
+    Note that this only works in the raylet code path.
+
+    This function can be used as follows (both on the driver or within a task).
+
+    .. testcode::
+        import ray._private.profiling as profiling
+
+        with profiling.profile("custom event", extra_data={'key': 'val'}):
+            # Do some computation here.
+            x = 1 * 2
+
+    Optionally, a dictionary can be passed as the "extra_data" argument, and
+    it can have keys "name" and "cname" if you want to override the default
+    timeline display text and box color. Other values will appear at the bottom
+    of the chrome tracing GUI when you click on the box corresponding to this
+    profile span.
+
+    Args:
+        event_type: A string describing the type of the event.
+        extra_data: This must be a dictionary mapping strings to strings. This
+            data will be added to the json objects that are used to populate
+            the timeline, so if you want to set a particular color, you can
+            simply set the "cname" attribute to an appropriate color.
+            Similarly, if you set the "name" attribute, then that will set the
+            text displayed on the box in the timeline.
+
+    Returns:
+        An object that can profile a span of time via a "with" statement.
+    """
+    if not PROFILING_ENABLED:
+        return NULL_LOG_SPAN
+    worker = ray._private.worker.global_worker
+    return worker.core_worker.profile_event(event_type.encode("ascii"), extra_data)
+
+
+def chrome_tracing_dump(
+    tasks: List[dict],
+) -> str:
+    """Generate a chrome/perfetto tracing dump using task events.
+
+    Args:
+        tasks: List of tasks generated by a state API list_tasks(detail=True).
+
+    Returns:
+        Json serialized dump to create a chrome/perfetto tracing.
+    """
+    # All events from given tasks.
+    all_events = []
+
+    # Chrome tracing doesn't have a concept of "node". Instead, we use
+    # chrome tracing's pid == ray's node.
+    # chrome tracing's tid == ray's process.
+    # Note that pid or tid is usually integer, but ray's node/process has
+    # ids in string.
+    # Unfortunately, perfetto doesn't allow to have string as a value of pid/tid.
+    # To workaround it, we use Metadata event from chrome tracing schema
+    # (https://docs.google.com/document/d/1CvAClvFfyA5R-PhYUmn5OOQtYMH4h6I0nSsKchNAySU/preview#heading=h.xqopa5m0e28f) # noqa
+    # which allows pid/tid -> name mapping. In order to use this schema
+    # we build node_ip/(node_ip, worker_id) -> arbitrary index mapping.
+
+    # node ip address -> node idx.
+    node_to_index = {}
+    # Arbitrary index mapped to the ip address.
+    node_idx = 0
+    # (node index, worker id) -> worker idx
+    worker_to_index = {}
+    # Arbitrary index mapped to the (node index, worker id).
+    worker_idx = 0
+
+    for task in tasks:
+        profiling_data = task.get("profiling_data", [])
+        if profiling_data:
+            node_ip_address = profiling_data["node_ip_address"]
+            component_events = profiling_data["events"]
+            component_type = profiling_data["component_type"]
+            component_id = component_type + ":" + profiling_data["component_id"]
+
+            if component_type not in ["worker", "driver"]:
+                continue
+
+            for event in component_events:
+                extra_data = event["extra_data"]
+                # Propagate extra data.
+                extra_data["task_id"] = task["task_id"]
+                extra_data["job_id"] = task["job_id"]
+                extra_data["attempt_number"] = task["attempt_number"]
+                extra_data["func_or_class_name"] = task["func_or_class_name"]
+                extra_data["actor_id"] = task["actor_id"]
+                event_name = event["event_name"]
+
+                # build a id -> arbitrary index mapping
+                if node_ip_address not in node_to_index:
+                    node_to_index[node_ip_address] = node_idx
+                    # Whenever new node ip is introduced, we increment the index.
+                    node_idx += 1
+
+                if (
+                    node_to_index[node_ip_address],
+                    component_id,
+                ) not in worker_to_index:  # noqa
+                    worker_to_index[
+                        (node_to_index[node_ip_address], component_id)
+                    ] = worker_idx  # noqa
+                    worker_idx += 1
+
+                # Modify the name with the additional user-defined extra data.
+                cname = _default_color_mapping[event["event_name"]]
+                name = event_name
+
+                if "cname" in extra_data:
+                    cname = _default_color_mapping[event["extra_data"]["cname"]]
+                if "name" in extra_data:
+                    name = extra_data["name"]
+
+                new_event = ChromeTracingCompleteEvent(
+                    cat=event_name,
+                    name=name,
+                    pid=node_to_index[node_ip_address],
+                    tid=worker_to_index[(node_to_index[node_ip_address], component_id)],
+                    ts=event["start_time"] * 1e3,
+                    dur=(event["end_time"] * 1e3) - (event["start_time"] * 1e3),
+                    cname=cname,
+                    args=extra_data,
+                )
+                all_events.append(asdict(new_event))
+
+    for node, i in node_to_index.items():
+        all_events.append(
+            asdict(
+                ChromeTracingMetadataEvent(
+                    name="process_name",
+                    pid=i,
+                    args={"name": f"Node {node}"},
+                )
+            )
+        )
+
+    for worker, i in worker_to_index.items():
+        all_events.append(
+            asdict(
+                ChromeTracingMetadataEvent(
+                    name="thread_name",
+                    ph="M",
+                    tid=i,
+                    pid=worker[0],
+                    args={"name": worker[1]},
+                )
+            )
+        )
+
+    # Handle task event disabled.
+    return json.dumps(all_events)

python/ray/_private/prometheus_exporter.py

@@ -0,0 +1,375 @@
+# NOTE: This file has been copied from OpenCensus Python exporter.
+# It is because OpenCensus Prometheus exporter hasn't released for a while
+# and the latest version has a compatibility issue with the latest OpenCensus
+# library.
+
+import logging
+import re
+import threading
+from typing import Any
+from wsgiref.simple_server import make_server
+
+from opencensus.common.transports import sync
+from opencensus.stats import aggregation_data as aggregation_data_module, base_exporter
+from prometheus_client import make_wsgi_app
+from prometheus_client.core import (
+    REGISTRY,
+    CollectorRegistry,
+    CounterMetricFamily,
+    GaugeMetricFamily,
+    HistogramMetricFamily,
+    UnknownMetricFamily,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class Options(object):
+    """Options contains options for configuring the exporter.
+
+    The address can be empty as the prometheus client will
+    assume it's localhost.
+
+    Args:
+        namespace: The prometheus namespace to be used. Defaults to ''.
+        port: The Prometheus port to be used. Defaults to 8000.
+        address: The Prometheus address to be used. Defaults to ''.
+        registry: A Prometheus collector registry instance.
+    """
+
+    def __init__(
+        self,
+        namespace: str = "",
+        port: int = 8000,
+        address: str = "",
+        registry: CollectorRegistry = REGISTRY,
+    ):
+        self._namespace = namespace
+        self._registry = registry
+        self._port = int(port)
+        self._address = address
+
+    @property
+    def registry(self):
+        """Prometheus Collector Registry instance"""
+        return self._registry
+
+    @property
+    def namespace(self):
+        """Prefix to be used with view name"""
+        return self._namespace
+
+    @property
+    def port(self):
+        """Port number to listen"""
+        return self._port
+
+    @property
+    def address(self):
+        """Endpoint address (default is localhost)"""
+        return self._address
+
+
+class Collector(object):
+    """Collector represents the Prometheus Collector object"""
+
+    def __init__(self, options=Options(), view_name_to_data_map=None):
+        if view_name_to_data_map is None:
+            view_name_to_data_map = {}
+        self._options = options
+        self._registry = options.registry
+        self._view_name_to_data_map = view_name_to_data_map
+        self._registered_views = {}
+
+    @property
+    def options(self):
+        """Options to be used to configure the exporter"""
+        return self._options
+
+    @property
+    def registry(self):
+        """Prometheus Collector Registry instance"""
+        return self._registry
+
+    @property
+    def view_name_to_data_map(self):
+        """Map with all view data objects
+        that will be sent to Prometheus
+        """
+        return self._view_name_to_data_map
+
+    @property
+    def registered_views(self):
+        """Map with all registered views"""
+        return self._registered_views
+
+    def register_view(self, view):
+        """register_view will create the needed structure
+        in order to be able to sent all data to Prometheus
+        """
+        v_name = get_view_name(self.options.namespace, view)
+
+        if v_name not in self.registered_views:
+            desc = {
+                "name": v_name,
+                "documentation": view.description,
+                "labels": list(map(sanitize, view.columns)),
+                "units": view.measure.unit,
+            }
+            self.registered_views[v_name] = desc
+
+    def add_view_data(self, view_data):
+        """Add view data object to be sent to server"""
+        self.register_view(view_data.view)
+        v_name = get_view_name(self.options.namespace, view_data.view)
+        self.view_name_to_data_map[v_name] = view_data
+
+    # TODO: add start and end timestamp
+    def to_metric(
+        self,
+        desc: dict,
+        tag_values: tuple,
+        agg_data: Any,
+        metrics_map: dict,
+    ) -> None:
+        """Translate the data that OpenCensus creates to Prometheus format.
+
+        Args:
+            desc: The map that describes view definition.
+            tag_values: TagValue object used as label values.
+            agg_data: Aggregated data that needs to be converted as Prometheus samples.
+            metrics_map: A map of metric name to Prometheus Metric object that is
+                populated by this method.
+        """
+        metric_name = desc["name"]
+        metric_description = desc["documentation"]
+        label_keys = desc["labels"]
+        metric_units = desc["units"]
+        assert len(tag_values) == len(label_keys), (tag_values, label_keys)
+        # Prometheus requires that all tag values be strings hence
+        # the need to cast none to the empty string before exporting. See
+        # https://github.com/census-instrumentation/opencensus-python/issues/480
+        tag_values = [tv if tv else "" for tv in tag_values]
+
+        if isinstance(agg_data, aggregation_data_module.CountAggregationData):
+            metric = metrics_map.get(metric_name)
+            if not metric:
+                metric = CounterMetricFamily(
+                    name=metric_name,
+                    documentation=metric_description,
+                    unit=metric_units,
+                    labels=label_keys,
+                )
+                metrics_map[metric_name] = metric
+            metric.add_metric(labels=tag_values, value=agg_data.count_data)
+            return
+
+        elif isinstance(agg_data, aggregation_data_module.DistributionAggregationData):
+
+            assert agg_data.bounds == sorted(agg_data.bounds)
+            # buckets are a list of buckets. Each bucket is another list with
+            # a pair of bucket name and value, or a triple of bucket name,
+            # value, and exemplar. buckets need to be in order.
+            buckets = []
+            cum_count = 0  # Prometheus buckets expect cumulative count.
+            for ii, bound in enumerate(agg_data.bounds):
+                cum_count += agg_data.counts_per_bucket[ii]
+                bucket = [str(bound), cum_count]
+                buckets.append(bucket)
+            # Prometheus requires buckets to be sorted, and +Inf present.
+            # In OpenCensus we don't have +Inf in the bucket bonds so need to
+            # append it here.
+            buckets.append(["+Inf", agg_data.count_data])
+            metric = metrics_map.get(metric_name)
+            if not metric:
+                metric = HistogramMetricFamily(
+                    name=metric_name,
+                    documentation=metric_description,
+                    labels=label_keys,
+                )
+                metrics_map[metric_name] = metric
+            metric.add_metric(
+                labels=tag_values,
+                buckets=buckets,
+                sum_value=agg_data.sum,
+            )
+            return
+
+        elif isinstance(agg_data, aggregation_data_module.SumAggregationData):
+            metric = metrics_map.get(metric_name)
+            if not metric:
+                metric = UnknownMetricFamily(
+                    name=metric_name,
+                    documentation=metric_description,
+                    labels=label_keys,
+                )
+                metrics_map[metric_name] = metric
+            metric.add_metric(labels=tag_values, value=agg_data.sum_data)
+            return
+
+        elif isinstance(agg_data, aggregation_data_module.LastValueAggregationData):
+            metric = metrics_map.get(metric_name)
+            if not metric:
+                metric = GaugeMetricFamily(
+                    name=metric_name,
+                    documentation=metric_description,
+                    labels=label_keys,
+                )
+                metrics_map[metric_name] = metric
+            metric.add_metric(labels=tag_values, value=agg_data.value)
+            return
+
+        else:
+            raise ValueError(f"unsupported aggregation type {type(agg_data)}")
+
+    def collect(self):  # pragma: NO COVER
+        """Collect fetches the statistics from OpenCensus
+        and delivers them as Prometheus Metrics.
+        Collect is invoked every time a prometheus.Gatherer is run
+        for example when the HTTP endpoint is invoked by Prometheus.
+        """
+        # Make a shallow copy of self._view_name_to_data_map, to avoid seeing
+        # concurrent modifications when iterating through the dictionary.
+        metrics_map = {}
+        for v_name, view_data in self._view_name_to_data_map.copy().items():
+            if v_name not in self.registered_views:
+                continue
+            desc = self.registered_views[v_name]
+            for tag_values in view_data.tag_value_aggregation_data_map:
+                agg_data = view_data.tag_value_aggregation_data_map[tag_values]
+                self.to_metric(desc, tag_values, agg_data, metrics_map)
+
+        for metric in metrics_map.values():
+            yield metric
+
+
+class PrometheusStatsExporter(base_exporter.StatsExporter):
+    """Exporter exports stats to Prometheus.
+
+    Users need to register the exporter as an HTTP Handler to be able to export.
+
+    Args:
+        options: An options object with the parameters to instantiate the
+            prometheus exporter.
+        gatherer: A Prometheus collector registry instance.
+        transport: An instance of a Transport to send data with.
+        collector: An instance of the Prometheus Collector object.
+    """
+
+    def __init__(
+        self,
+        options: Options,
+        gatherer: CollectorRegistry,
+        transport: Any = sync.SyncTransport,
+        collector: "Collector" = Collector(),
+    ):
+        self._options = options
+        self._gatherer = gatherer
+        self._collector = collector
+        self._transport = transport(self)
+        self._port = self.serve_http()
+        REGISTRY.register(self._collector)
+
+    @property
+    def transport(self):
+        """The transport way to be sent data to server
+        (default is sync).
+        """
+        return self._transport
+
+    @property
+    def collector(self):
+        """Collector class instance to be used
+        to communicate with Prometheus
+        """
+        return self._collector
+
+    @property
+    def gatherer(self):
+        """Prometheus Collector Registry instance"""
+        return self._gatherer
+
+    @property
+    def options(self):
+        """Options to be used to configure the exporter"""
+        return self._options
+
+    @property
+    def port(self):
+        """The port the HTTP server is listening on."""
+        return self._port
+
+    def export(self, view_data):
+        """export send the data to the transport class
+        in order to be sent to Prometheus in a sync or async way.
+        """
+        if view_data is not None:  # pragma: NO COVER
+            self.transport.export(view_data)
+
+    def on_register_view(self, view):
+        return NotImplementedError("Not supported by Prometheus")
+
+    def emit(self, view_data):  # pragma: NO COVER
+        """Emit exports to the Prometheus if view data has one or more rows.
+        Each OpenCensus AggregationData will be converted to
+        corresponding Prometheus Metric: SumData will be converted
+        to Untyped Metric, CountData will be a Counter Metric
+        DistributionData will be a Histogram Metric.
+        """
+
+        for v_data in view_data:
+            if v_data.tag_value_aggregation_data_map is None:
+                v_data.tag_value_aggregation_data_map = {}
+
+            self.collector.add_view_data(v_data)
+
+    def serve_http(self):
+        """serve_http serves the Prometheus endpoint."""
+        address = self.options.address or ""
+        httpd = make_server(address, self.options.port, make_wsgi_app())
+        t = threading.Thread(target=httpd.serve_forever)
+        t.daemon = True
+        t.start()
+        # Return the actual port (in case port=0 was specified)
+        return httpd.server_address[1]
+
+
+def new_stats_exporter(option):
+    """new_stats_exporter returns an exporter
+    that exports stats to Prometheus.
+    """
+    if option.namespace == "":
+        raise ValueError("Namespace can not be empty string.")
+
+    collector = new_collector(option)
+
+    exporter = PrometheusStatsExporter(
+        options=option, gatherer=option.registry, collector=collector
+    )
+    return exporter
+
+
+def new_collector(options):
+    """new_collector should be used
+    to create instance of Collector class in order to
+    prevent the usage of constructor directly
+    """
+    return Collector(options=options)
+
+
+def get_view_name(namespace, view):
+    """create the name for the view"""
+    name = ""
+    if namespace != "":
+        name = namespace + "_"
+    return sanitize(name + view.name)
+
+
+_NON_LETTERS_NOR_DIGITS_RE = re.compile(r"[^\w]", re.UNICODE | re.IGNORECASE)
+
+
+def sanitize(key):
+    """sanitize the given metric name or label according to Prometheus rule.
+    Replace all characters other than [A-Za-z0-9_] with '_'.
+    """
+    return _NON_LETTERS_NOR_DIGITS_RE.sub("_", key)

python/ray/_private/protobuf_compat.py

@@ -0,0 +1,52 @@
+import inspect
+
+from google.protobuf.json_format import MessageToDict, MessageToJson
+
+"""
+This module provides a compatibility layer for different versions of the protobuf
+library.
+"""
+
+_protobuf_has_old_arg_name_cached = None
+
+
+def _protobuf_has_old_arg_name():
+    """Cache the inspect result to avoid doing it for every single message."""
+    global _protobuf_has_old_arg_name_cached
+    if _protobuf_has_old_arg_name_cached is None:
+        params = inspect.signature(MessageToDict).parameters
+        _protobuf_has_old_arg_name_cached = "including_default_value_fields" in params
+    return _protobuf_has_old_arg_name_cached
+
+
+def rename_always_print_fields_with_no_presence(kwargs):
+    """
+    Protobuf version 5.26.0rc2 renamed argument for `MessageToDict` and `MessageToJson`:
+    `including_default_value_fields` -> `always_print_fields_with_no_presence`.
+    See https://github.com/protocolbuffers/protobuf/commit/06e7caba58ede0220b110b89d08f329e5f8a7537#diff-8de817c14d6a087981503c9aea38730b1b3e98f4e306db5ff9d525c7c304f234L129  # noqa: E501
+
+    We choose to always use the new argument name. If user used the old arg, we raise an
+    error.
+
+    If protobuf does not have the new arg name but have the old arg name, we rename our
+    arg to the old one.
+    """
+    old_arg_name = "including_default_value_fields"
+    new_arg_name = "always_print_fields_with_no_presence"
+    if old_arg_name in kwargs:
+        raise ValueError(f"{old_arg_name} is deprecated, please use {new_arg_name}")
+
+    if new_arg_name in kwargs and _protobuf_has_old_arg_name():
+        kwargs[old_arg_name] = kwargs.pop(new_arg_name)
+
+    return kwargs
+
+
+def message_to_dict(*args, **kwargs):
+    kwargs = rename_always_print_fields_with_no_presence(kwargs)
+    return MessageToDict(*args, **kwargs)
+
+
+def message_to_json(*args, **kwargs):
+    kwargs = rename_always_print_fields_with_no_presence(kwargs)
+    return MessageToJson(*args, **kwargs)

python/ray/_private/ray_client_microbenchmark.py

@@ -0,0 +1,117 @@
+import inspect
+import logging
+import sys
+
+import numpy as np
+
+from ray._private.ray_microbenchmark_helpers import timeit
+from ray.util.client.ray_client_helpers import ray_start_client_server
+
+
+def benchmark_get_calls(ray, results):
+    value = ray.put(0)
+
+    def get_small():
+        ray.get(value)
+
+    results += timeit("client: get calls", get_small)
+
+
+def benchmark_tasks_and_get_batch(ray, results):
+    @ray.remote
+    def small_value():
+        return b"ok"
+
+    def small_value_batch():
+        submitted = [small_value.remote() for _ in range(1000)]
+        ray.get(submitted)
+        return 0
+
+    results += timeit("client: tasks and get batch", small_value_batch)
+
+
+def benchmark_put_calls(ray, results):
+    def put_small():
+        ray.put(0)
+
+    results += timeit("client: put calls", put_small)
+
+
+def benchmark_remote_put_calls(ray, results):
+    @ray.remote
+    def do_put_small():
+        for _ in range(100):
+            ray.put(0)
+
+    def put_multi_small():
+        ray.get([do_put_small.remote() for _ in range(10)])
+
+    results += timeit("client: tasks and put batch", put_multi_small, 1000)
+
+
+def benchmark_put_large(ray, results):
+    arr = np.zeros(100 * 1024 * 1024, dtype=np.int64)
+
+    def put_large():
+        ray.put(arr)
+
+    results += timeit("client: put gigabytes", put_large, 8 * 0.1)
+
+
+def benchmark_simple_actor(ray, results):
+    @ray.remote(num_cpus=0)
+    class Actor:
+        def small_value(self):
+            return b"ok"
+
+        def small_value_arg(self, x):
+            return b"ok"
+
+        def small_value_batch(self, n):
+            ray.get([self.small_value.remote() for _ in range(n)])
+
+    a = Actor.remote()
+
+    def actor_sync():
+        ray.get(a.small_value.remote())
+
+    results += timeit("client: 1:1 actor calls sync", actor_sync)
+
+    def actor_async():
+        ray.get([a.small_value.remote() for _ in range(1000)])
+
+    results += timeit("client: 1:1 actor calls async", actor_async, 1000)
+
+    a = Actor.options(max_concurrency=16).remote()
+
+    def actor_concurrent():
+        ray.get([a.small_value.remote() for _ in range(1000)])
+
+    results += timeit("client: 1:1 actor calls concurrent", actor_concurrent, 1000)
+
+
+def main(results=None):
+    results = results or []
+
+    ray_config = {"logging_level": logging.WARNING}
+
+    def ray_connect_handler(job_config=None, **ray_init_kwargs):
+        from ray._private.client_mode_hook import disable_client_hook
+
+        with disable_client_hook():
+            import ray as real_ray
+
+            if not real_ray.is_initialized():
+                real_ray.init(**ray_config)
+
+    for name, obj in inspect.getmembers(sys.modules[__name__]):
+        if not name.startswith("benchmark_"):
+            continue
+        with ray_start_client_server(ray_connect_handler=ray_connect_handler) as ray:
+            obj(ray, results)
+
+    return results
+
+
+if __name__ == "__main__":
+    main()

python/ray/_private/ray_cluster_perf.py

@@ -0,0 +1,51 @@
+"""This is the script for `ray clusterbenchmark`."""
+
+import time
+
+import numpy as np
+
+import ray
+from ray.cluster_utils import Cluster
+
+
+def main():
+    cluster = Cluster(
+        initialize_head=True,
+        connect=True,
+        head_node_args={"object_store_memory": 20 * 1024 * 1024 * 1024, "num_cpus": 16},
+    )
+    cluster.add_node(
+        object_store_memory=20 * 1024 * 1024 * 1024, num_gpus=1, num_cpus=16
+    )
+
+    object_ref_list = []
+    for i in range(0, 10):
+        object_ref = ray.put(np.random.rand(1024 * 128, 1024))
+        object_ref_list.append(object_ref)
+
+    @ray.remote(num_gpus=1)
+    def f(object_ref_list):
+        diffs = []
+        for object_ref in object_ref_list:
+            before = time.time()
+            ray.get(object_ref)
+            after = time.time()
+            diffs.append(after - before)
+            time.sleep(1)
+        return np.mean(diffs), np.std(diffs)
+
+    time_diff, time_diff_std = ray.get(f.remote(object_ref_list))
+
+    print(
+        "latency to get an 1G object over network",
+        round(time_diff, 2),
+        "+-",
+        round(time_diff_std, 2),
+    )
+
+    ray.shutdown()
+    cluster.shutdown()
+
+
+if __name__ == "__main__":
+    main()

python/ray/_private/ray_constants.py

@@ -0,0 +1,628 @@
+"""Ray constants used in the Python code."""
+
+import json
+import logging
+import os
+import sys
+
+from ray._common.utils import env_bool, env_float, env_integer  # noqa: F401
+
+logger = logging.getLogger(__name__)
+
+
+def env_set_by_user(key):
+    return key in os.environ
+
+
+# Whether event logging to driver is enabled. Set to 0 to disable.
+AUTOSCALER_EVENTS = env_integer("RAY_SCHEDULER_EVENTS", 1)
+
+# Whether to disable the C++ failure signal handler that provides stack traces
+# on crashes. Disabling this is necessary when using Java libraries
+# because Ray's signal handler conflicts with the JVM's signal handling.
+RAY_DISABLE_FAILURE_SIGNAL_HANDLER = env_bool(
+    "RAY_DISABLE_FAILURE_SIGNAL_HANDLER", False
+)
+
+RAY_LOG_TO_DRIVER = env_bool("RAY_LOG_TO_DRIVER", True)
+
+# Filter level under which events will be filtered out, i.e. not printing to driver
+RAY_LOG_TO_DRIVER_EVENT_LEVEL = os.environ.get("RAY_LOG_TO_DRIVER_EVENT_LEVEL", "INFO")
+
+# Internal kv keys for storing monitor debug status.
+DEBUG_AUTOSCALING_ERROR = "__autoscaling_error"
+DEBUG_AUTOSCALING_STATUS = "__autoscaling_status"
+DEBUG_AUTOSCALING_STATUS_LEGACY = "__autoscaling_status_legacy"
+
+ID_SIZE = 28
+
+# The following constants are used to create default values for
+# resource isolation when it is enabled.
+# TODO(54703): Link to OSS documentation about the feature once it's available.
+DEFAULT_CGROUP_PATH = "/sys/fs/cgroup"
+# The default proportion of cpu cores to reserve for ray system processes.
+DEFAULT_SYSTEM_RESERVED_CPU_PROPORTION = env_float(
+    "RAY_DEFAULT_SYSTEM_RESERVED_CPU_PROPORTION", 0.05
+)
+# The default minimum number of cpu cores to reserve for ray system processes.
+# This value is used if the available_cores * DEFAULT_SYSTEM_RESERVED_CPU_PROPORTION < this value.
+DEFAULT_MIN_SYSTEM_RESERVED_CPU_CORES = env_float(
+    "RAY_DEFAULT_MIN_SYSTEM_RESERVED_CPU_CORES", 1.0
+)
+# The default maximum number of cpu cores to reserve for ray system processes.
+# This value is used if the available_cores * DEFAULT_SYSTEM_RESERVED_CPU_PROPORTION > this value.
+DEFAULT_MAX_SYSTEM_RESERVED_CPU_CORES = env_float(
+    "RAY_DEFAULT_MAX_SYSTEM_RESERVED_CPU_CORES", 3.0
+)
+# The values for SYSTEM_RESERVED_MEMORY do not include the memory reserveed
+# for the object store.
+# The default proportion available memory to reserve for ray system processes.
+DEFAULT_SYSTEM_RESERVED_MEMORY_PROPORTION = env_float(
+    "RAY_DEFAULT_SYSTEM_RESERVED_MEMORY_PROPORTION", 0.10
+)
+# The default minimum number of bytes to reserve for ray system processes.
+# This value is used if the available_memory * DEFAULT_SYSTEM_RESERVED_MEMORY_PROPORTION < this value.
+DEFAULT_MIN_SYSTEM_RESERVED_MEMORY_BYTES = env_integer(
+    "RAY_DEFAULT_MIN_SYSTEM_RESERVED_MEMORY_BYTES", 500 * (1024**2)  # 500MB
+)
+# The default maximum number of bytes to reserve for ray system processes.
+# This value is used if the available_memory * DEFAULT_SYSTEM_RESERVED_MEMORY_PROPORTION > this value.
+DEFAULT_MAX_SYSTEM_RESERVED_MEMORY_BYTES = env_integer(
+    "RAY_DEFAULT_MAX_SYSTEM_RESERVED_MEMORY_BYTES", (10) * (1024**3)
+)
+# The default buffer size between the physical memory limit enforced by resource isolation
+# and the logical memory limit available for scheduling user tasks. This buffer can be tuned
+# to allocate more or less memory room for tolerating passing in the wrong logical memory
+# estimate at the cost of lower memory utilization.
+DEFAULT_USER_PHYSICAL_LOGICAL_MEMORY_LIMIT_BUFFER_BYTES = env_integer(
+    "RAY_DEFAULT_USER_PHYSICAL_LOGICAL_MEMORY_LIMIT_BUFFER_BYTES",
+    500 * (1024**2),  # 500MiB
+)
+
+# The default maximum number of bytes to allocate to the object store unless
+# overridden by the user.
+DEFAULT_OBJECT_STORE_MAX_MEMORY_BYTES = env_integer(
+    "RAY_DEFAULT_OBJECT_STORE_MAX_MEMORY_BYTES", (200) * (10**9)  # 200 GB
+)
+# The default proportion of available memory allocated to the object store
+DEFAULT_OBJECT_STORE_MEMORY_PROPORTION = env_float(
+    "RAY_DEFAULT_OBJECT_STORE_MEMORY_PROPORTION",
+    0.3,
+)
+# The smallest cap on the memory used by the object store that we allow.
+# This must be greater than MEMORY_RESOURCE_UNIT_BYTES
+OBJECT_STORE_MINIMUM_MEMORY_BYTES = 75 * 1024 * 1024
+# Each ObjectRef currently uses about 3KB of caller memory.
+CALLER_MEMORY_USAGE_PER_OBJECT_REF = 3000
+# Above this number of bytes, raise an error by default unless the user sets
+# RAY_ALLOW_SLOW_STORAGE=1. This avoids swapping with large object stores.
+REQUIRE_SHM_SIZE_THRESHOLD = 10**10
+# Mac with 16GB memory has degraded performance when the object store size is
+# greater than 2GB.
+# (see https://github.com/ray-project/ray/issues/20388 for details)
+# The workaround here is to limit capacity to 2GB for Mac by default,
+# and raise error if the capacity is overwritten by user.
+MAC_DEGRADED_PERF_MMAP_SIZE_LIMIT = (2) * (2**30)
+# If a user does not specify a port for the primary Ray service,
+# we attempt to start the service running at this port.
+DEFAULT_PORT = 6379
+
+RAY_ADDRESS_ENVIRONMENT_VARIABLE = "RAY_ADDRESS"
+RAY_API_SERVER_ADDRESS_ENVIRONMENT_VARIABLE = "RAY_API_SERVER_ADDRESS"
+RAY_NAMESPACE_ENVIRONMENT_VARIABLE = "RAY_NAMESPACE"
+RAY_RUNTIME_ENV_ENVIRONMENT_VARIABLE = "RAY_RUNTIME_ENV"
+RAY_RUNTIME_ENV_URI_PIN_EXPIRATION_S_ENV_VAR = (
+    "RAY_RUNTIME_ENV_TEMPORARY_REFERENCE_EXPIRATION_S"
+)
+# Ray populates this env var to the working dir in the creation of a runtime env.
+# For example, `pip` and `conda` users can use this environment variable to locate the
+# `requirements.txt` file.
+RAY_RUNTIME_ENV_CREATE_WORKING_DIR_ENV_VAR = "RAY_RUNTIME_ENV_CREATE_WORKING_DIR"
+# Defaults to 10 minutes. This should be longer than the total time it takes for
+# the local working_dir and py_modules to be uploaded, or these files might get
+# garbage collected before the job starts.
+RAY_RUNTIME_ENV_URI_PIN_EXPIRATION_S_DEFAULT = 10 * 60
+# If set to 1, then `.gitignore` files will not be parsed and loaded into "excludes"
+# when using a local working_dir or py_modules.
+RAY_RUNTIME_ENV_IGNORE_GITIGNORE = "RAY_RUNTIME_ENV_IGNORE_GITIGNORE"
+# Default directories to exclude when packaging working_dir.
+# Override by setting the RAY_OVERRIDE_RUNTIME_ENV_DEFAULT_EXCLUDES
+# (comma-separated) environment variable. Set to an empty string to disable.
+# `.git` is necessary since it is never in .gitignore.
+RAY_RUNTIME_ENV_DEFAULT_EXCLUDES = ".git,.venv,venv,__pycache__"
+
+
+def get_runtime_env_default_excludes() -> list[str]:
+    """Get default excludes for working_dir, overridable via RAY_OVERRIDE_RUNTIME_ENV_DEFAULT_EXCLUDES environment variable."""
+    val = os.environ.get(
+        "RAY_OVERRIDE_RUNTIME_ENV_DEFAULT_EXCLUDES", RAY_RUNTIME_ENV_DEFAULT_EXCLUDES
+    )
+    return [x.strip() for x in val.split(",") if x.strip()]
+
+
+# Hook for running a user-specified runtime-env hook. This hook will be called
+# unconditionally given the runtime_env dict passed for ray.init. It must return
+# a rewritten runtime_env dict. Example: "your.module.runtime_env_hook".
+RAY_RUNTIME_ENV_HOOK = "RAY_RUNTIME_ENV_HOOK"
+# Hook that is invoked on `ray start`. It will be given the cluster parameters and
+# whether we are the head node as arguments. The function can modify the params class,
+# but otherwise returns void. Example: "your.module.ray_start_hook".
+RAY_START_HOOK = "RAY_START_HOOK"
+# Hook that is invoked on `ray job submit`. It will be given all the same args as the
+# job.cli.submit() function gets, passed as kwargs to this function.
+RAY_JOB_SUBMIT_HOOK = "RAY_JOB_SUBMIT_HOOK"
+# Headers to pass when using the Job CLI. It will be given to
+# instantiate a Job SubmissionClient.
+RAY_JOB_HEADERS = "RAY_JOB_HEADERS"
+
+# Timeout waiting for the dashboard to come alive during node startup.
+RAY_DASHBOARD_STARTUP_TIMEOUT_S = env_integer("RAY_DASHBOARD_STARTUP_TIMEOUT_S", 60)
+
+# Enable profiling endpoints in the dashboard.
+RAY_DASHBOARD_ENABLE_PROFILING = env_bool("RAY_DASHBOARD_ENABLE_PROFILING", False)
+
+DEFAULT_DASHBOARD_PORT = 8265
+DASHBOARD_ADDRESS = "dashboard"
+DASHBOARD_CLIENT_MAX_SIZE = 100 * 1024**2
+PROMETHEUS_SERVICE_DISCOVERY_FILE = "prom_metrics_service_discovery.json"
+DEFAULT_DASHBOARD_AGENT_LISTEN_PORT = 52365
+# Default resource requirements for actors when no resource requirements are
+# specified.
+DEFAULT_ACTOR_METHOD_CPU_SIMPLE = 1
+DEFAULT_ACTOR_CREATION_CPU_SIMPLE = 0
+# Default resource requirements for actors when some resource requirements are
+# specified in .
+DEFAULT_ACTOR_METHOD_CPU_SPECIFIED = 0
+DEFAULT_ACTOR_CREATION_CPU_SPECIFIED = 1
+# Default number of return values for each actor method.
+DEFAULT_ACTOR_METHOD_NUM_RETURN_VALS = 1
+
+# Wait 30 seconds for client to reconnect after unexpected disconnection
+DEFAULT_CLIENT_RECONNECT_GRACE_PERIOD = 30
+
+# If a remote function or actor (or some other export) has serialized size
+# greater than this quantity, print an warning.
+FUNCTION_SIZE_WARN_THRESHOLD = 10**7
+FUNCTION_SIZE_ERROR_THRESHOLD = env_integer("FUNCTION_SIZE_ERROR_THRESHOLD", (10**8))
+
+# If remote functions with the same source are imported this many times, then
+# print a warning.
+DUPLICATE_REMOTE_FUNCTION_THRESHOLD = 100
+
+# The maximum resource quantity that is allowed. TODO(rkn): This could be
+# relaxed, but the current implementation of the node manager will be slower
+# for large resource quantities due to bookkeeping of specific resource IDs.
+MAX_RESOURCE_QUANTITY = 100e12
+
+# Number of units 1 resource can be subdivided into.
+MIN_RESOURCE_GRANULARITY = 0.0001
+
+# Set this environment variable to populate the dashboard URL with
+# an external hosted Ray dashboard URL (e.g. because the
+# dashboard is behind a proxy or load balancer). This only overrides
+# the dashboard URL when returning or printing to a user through a public
+# API, but not in the internal KV store.
+RAY_OVERRIDE_DASHBOARD_URL = "RAY_OVERRIDE_DASHBOARD_URL"
+
+
+# Different types of Ray errors that can be pushed to the driver.
+# TODO(rkn): These should be defined in flatbuffers and must be synced with
+# the existing C++ definitions.
+PICKLING_LARGE_OBJECT_PUSH_ERROR = "pickling_large_object"
+WAIT_FOR_FUNCTION_PUSH_ERROR = "wait_for_function"
+VERSION_MISMATCH_PUSH_ERROR = "version_mismatch"
+WORKER_CRASH_PUSH_ERROR = "worker_crash"
+WORKER_DIED_PUSH_ERROR = "worker_died"
+WORKER_POOL_LARGE_ERROR = "worker_pool_large"
+PUT_RECONSTRUCTION_PUSH_ERROR = "put_reconstruction"
+RESOURCE_DEADLOCK_ERROR = "resource_deadlock"
+REMOVED_NODE_ERROR = "node_removed"
+MONITOR_DIED_ERROR = "monitor_died"
+LOG_MONITOR_DIED_ERROR = "log_monitor_died"
+DASHBOARD_AGENT_DIED_ERROR = "dashboard_agent_died"
+DASHBOARD_DIED_ERROR = "dashboard_died"
+RAYLET_DIED_ERROR = "raylet_died"
+DETACHED_ACTOR_ANONYMOUS_NAMESPACE_ERROR = "detached_actor_anonymous_namespace"
+EXCESS_QUEUEING_WARNING = "excess_queueing_warning"
+
+# Used by autoscaler to set the node custom resources and labels
+# from cluster.yaml.
+RESOURCES_ENVIRONMENT_VARIABLE = "RAY_OVERRIDE_RESOURCES"
+LABELS_ENVIRONMENT_VARIABLE = "RAY_OVERRIDE_LABELS"
+
+# Temporary flag to disable log processing in the dashboard.  This is useful
+# if the dashboard is overloaded by logs and failing to process other
+# dashboard API requests (e.g. Job Submission).
+DISABLE_DASHBOARD_LOG_INFO = env_integer("RAY_DISABLE_DASHBOARD_LOG_INFO", 0)
+
+LOGGER_FORMAT = "%(asctime)s\t%(levelname)s %(filename)s:%(lineno)s -- %(message)s"
+LOGGER_FORMAT_ESCAPE = json.dumps(LOGGER_FORMAT.replace("%", "%%"))
+LOGGER_FORMAT_HELP = f"The logging format. default={LOGGER_FORMAT_ESCAPE}"
+# Configure the default logging levels for various Ray components.
+# TODO (kevin85421): Currently, I don't encourage Ray users to configure
+# `RAY_LOGGER_LEVEL` until its scope and expected behavior are clear and
+# easy to understand. Now, only Ray developers should use it.
+LOGGER_LEVEL = os.environ.get("RAY_LOGGER_LEVEL", "info")
+LOGGER_LEVEL_CHOICES = ["debug", "info", "warning", "error", "critical"]
+LOGGER_LEVEL_HELP = (
+    "The logging level threshold, choices=['debug', 'info',"
+    " 'warning', 'error', 'critical'], default='info'"
+)
+
+LOGGING_REDIRECT_STDERR_ENVIRONMENT_VARIABLE = "RAY_LOG_TO_STDERR"
+# Logging format when logging stderr. This should be formatted with the
+# component before setting the formatter, e.g. via
+#   format = LOGGER_FORMAT_STDERR.format(component="dashboard")
+#   handler.setFormatter(logging.Formatter(format))
+LOGGER_FORMAT_STDERR = (
+    "%(asctime)s\t%(levelname)s ({component}) %(filename)s:%(lineno)s -- %(message)s"
+)
+
+# Constants used to define the different process types.
+PROCESS_TYPE_REAPER = "reaper"
+PROCESS_TYPE_MONITOR = "monitor"
+PROCESS_TYPE_RAY_CLIENT_SERVER = "ray_client_server"
+PROCESS_TYPE_LOG_MONITOR = "log_monitor"
+PROCESS_TYPE_DASHBOARD = "dashboard"
+PROCESS_TYPE_DASHBOARD_AGENT = "dashboard_agent"
+PROCESS_TYPE_RUNTIME_ENV_AGENT = "runtime_env_agent"
+PROCESS_TYPE_WORKER = "worker"
+PROCESS_TYPE_RAYLET = "raylet"
+PROCESS_TYPE_REDIS_SERVER = "redis_server"
+PROCESS_TYPE_GCS_SERVER = "gcs_server"
+PROCESS_TYPE_PYTHON_CORE_WORKER_DRIVER = "python-core-driver"
+PROCESS_TYPE_PYTHON_CORE_WORKER = "python-core-worker"
+
+# Log file names
+MONITOR_LOG_FILE_NAME = f"{PROCESS_TYPE_MONITOR}.log"
+LOG_MONITOR_LOG_FILE_NAME = f"{PROCESS_TYPE_LOG_MONITOR}.log"
+
+# Enable log deduplication.
+RAY_DEDUP_LOGS = env_bool("RAY_DEDUP_LOGS", True)
+RAY_FLUSH_DRIVER_LOGS = env_bool("RAY_FLUSH_DRIVER_LOGS", False)
+# How many seconds of messages to buffer for log deduplication.
+RAY_DEDUP_LOGS_AGG_WINDOW_S = env_integer("RAY_DEDUP_LOGS_AGG_WINDOW_S", 5)
+
+# Regex for log messages to never deduplicate, or None. This takes precedence over
+# the skip regex below. A default pattern is set for testing.
+TESTING_NEVER_DEDUP_TOKEN = "__ray_testing_never_deduplicate__"
+RAY_DEDUP_LOGS_ALLOW_REGEX = os.environ.get(
+    "RAY_DEDUP_LOGS_ALLOW_REGEX", TESTING_NEVER_DEDUP_TOKEN
+)
+
+# Regex for log messages to always skip / suppress, or None.
+RAY_DEDUP_LOGS_SKIP_REGEX = os.environ.get("RAY_DEDUP_LOGS_SKIP_REGEX")
+
+AGENT_PROCESS_TYPE_DASHBOARD_AGENT = "ray::DashboardAgent"
+AGENT_PROCESS_TYPE_RUNTIME_ENV_AGENT = "ray::RuntimeEnvAgent"
+
+AGENT_PROCESS_LIST = [
+    AGENT_PROCESS_TYPE_DASHBOARD_AGENT,
+    AGENT_PROCESS_TYPE_RUNTIME_ENV_AGENT,
+]
+
+WORKER_PROCESS_TYPE_IDLE_WORKER = "ray::IDLE"
+WORKER_PROCESS_TYPE_SPILL_WORKER_NAME = "SpillWorker"
+WORKER_PROCESS_TYPE_RESTORE_WORKER_NAME = "RestoreWorker"
+WORKER_PROCESS_TYPE_SPILL_WORKER_IDLE = (
+    f"ray::IDLE_{WORKER_PROCESS_TYPE_SPILL_WORKER_NAME}"
+)
+WORKER_PROCESS_TYPE_RESTORE_WORKER_IDLE = (
+    f"ray::IDLE_{WORKER_PROCESS_TYPE_RESTORE_WORKER_NAME}"
+)
+WORKER_PROCESS_TYPE_SPILL_WORKER = f"ray::SPILL_{WORKER_PROCESS_TYPE_SPILL_WORKER_NAME}"
+WORKER_PROCESS_TYPE_RESTORE_WORKER = (
+    f"ray::RESTORE_{WORKER_PROCESS_TYPE_RESTORE_WORKER_NAME}"
+)
+WORKER_PROCESS_TYPE_SPILL_WORKER_DELETE = (
+    f"ray::DELETE_{WORKER_PROCESS_TYPE_SPILL_WORKER_NAME}"
+)
+WORKER_PROCESS_TYPE_RESTORE_WORKER_DELETE = (
+    f"ray::DELETE_{WORKER_PROCESS_TYPE_RESTORE_WORKER_NAME}"
+)
+
+# The number of files the log monitor will open. If more files exist, they will
+# be ignored.
+LOG_MONITOR_MAX_OPEN_FILES = int(
+    os.environ.get("RAY_LOG_MONITOR_MAX_OPEN_FILES", "200")
+)
+
+# The maximum batch of lines to be read in a single iteration. We _always_ try
+# to read this number of lines even if there aren't any new lines.
+LOG_MONITOR_NUM_LINES_TO_READ = int(
+    os.environ.get("RAY_LOG_MONITOR_NUM_LINES_TO_READ", "1000")
+)
+
+# Autoscaler events are denoted by the ":event_summary:" magic token.
+LOG_PREFIX_EVENT_SUMMARY = ":event_summary:"
+# Cluster-level info events are denoted by the ":info_message:" magic token. These may
+# be emitted in the stderr of Ray components.
+LOG_PREFIX_INFO_MESSAGE = ":info_message:"
+# Actor names are recorded in the logs with this magic token as a prefix.
+LOG_PREFIX_ACTOR_NAME = ":actor_name:"
+# Task names are recorded in the logs with this magic token as a prefix.
+LOG_PREFIX_TASK_NAME = ":task_name:"
+# Job ids are recorded in the logs with this magic token as a prefix.
+LOG_PREFIX_JOB_ID = ":job_id:"
+
+# The object metadata field uses the following format: It is a comma
+# separated list of fields. The first field is mandatory and is the
+# type of the object (see types below) or an integer, which is interpreted
+# as an error value. The second part is optional and if present has the
+# form DEBUG:<breakpoint_id>, it is used for implementing the debugger.
+
+# A constant used as object metadata to indicate the object is cross language.
+OBJECT_METADATA_TYPE_CROSS_LANGUAGE = b"XLANG"
+# A constant used as object metadata to indicate the object is python specific.
+OBJECT_METADATA_TYPE_PYTHON = b"PYTHON"
+# A constant used as object metadata to indicate the object is raw bytes.
+OBJECT_METADATA_TYPE_RAW = b"RAW"
+
+# A constant used as object metadata to indicate the object is an actor handle.
+# This value should be synchronized with the Java definition in
+# ObjectSerializer.java
+# TODO(fyrestone): Serialize the ActorHandle via the custom type feature
+# of XLANG.
+OBJECT_METADATA_TYPE_ACTOR_HANDLE = b"ACTOR_HANDLE"
+
+# A constant indicating the debugging part of the metadata (see above).
+OBJECT_METADATA_DEBUG_PREFIX = b"DEBUG:"
+
+AUTOSCALER_RESOURCE_REQUEST_CHANNEL = b"autoscaler_resource_request"
+
+REDIS_DEFAULT_USERNAME = ""
+
+REDIS_DEFAULT_PASSWORD = ""
+
+# The Mach kernel page size in bytes.
+MACH_PAGE_SIZE_BYTES = 4096
+
+# The max number of bytes for task execution error message.
+MAX_APPLICATION_ERROR_LENGTH = env_integer("RAY_MAX_APPLICATION_ERROR_LENGTH", 500)
+
+# Max 64 bit integer value, which is needed to ensure against overflow
+# in C++ when passing integer values cross-language.
+MAX_INT64_VALUE = 9223372036854775807
+
+# Object Spilling related constants
+DEFAULT_OBJECT_PREFIX = "ray_spilled_objects"
+
+GCS_PORT_ENVIRONMENT_VARIABLE = "RAY_GCS_SERVER_PORT"
+
+HEALTHCHECK_EXPIRATION_S = os.environ.get("RAY_HEALTHCHECK_EXPIRATION_S", 10)
+
+# Filename of "shim process" that sets up Python worker environment.
+# Should be kept in sync with kSetupWorkerFilename in
+# src/ray/common/constants.h.
+SETUP_WORKER_FILENAME = "setup_worker.py"
+
+# Directory name where runtime_env resources will be created & cached.
+DEFAULT_RUNTIME_ENV_DIR_NAME = "runtime_resources"
+
+# The timeout seconds for the creation of runtime env,
+# dafault timeout is 10 minutes
+DEFAULT_RUNTIME_ENV_TIMEOUT_SECONDS = 600
+
+# The timeout seconds for the GCS server request.
+# Try fetching from the cpp environment variable first.
+GCS_SERVER_REQUEST_TIMEOUT_SECONDS = int(
+    os.environ.get("RAY_gcs_server_request_timeout_seconds", "60")
+)
+
+# Used to separate lines when formatting the call stack where an ObjectRef was
+# created.
+CALL_STACK_LINE_DELIMITER = " | "
+
+# The default gRPC max message size is 4 MiB, we use a larger number of 512 MiB
+# NOTE: This is equal to the C++ limit of (RAY_CONFIG::max_grpc_message_size)
+GRPC_CPP_MAX_MESSAGE_SIZE = 512 * 1024 * 1024
+
+# The gRPC send & receive max length for "dashboard agent" server.
+# NOTE: This is equal to the C++ limit of RayConfig::max_grpc_message_size
+#       and HAVE TO STAY IN SYNC with it (ie, meaning that both of these values
+#       have to be set at the same time)
+AGENT_GRPC_MAX_MESSAGE_LENGTH = env_integer(
+    "AGENT_GRPC_MAX_MESSAGE_LENGTH", 20 * 1024 * 1024  # 20MB
+)
+
+
+# GRPC options
+GRPC_ENABLE_HTTP_PROXY = (
+    1
+    if os.environ.get("RAY_grpc_enable_http_proxy", "0").lower() in ("1", "true")
+    else 0
+)
+GLOBAL_GRPC_OPTIONS = (("grpc.enable_http_proxy", GRPC_ENABLE_HTTP_PROXY),)
+
+# Internal kv namespaces
+KV_NAMESPACE_DASHBOARD = b"dashboard"
+KV_NAMESPACE_SESSION = b"session"
+KV_NAMESPACE_TRACING = b"tracing"
+KV_NAMESPACE_PDB = b"ray_pdb"
+KV_NAMESPACE_HEALTHCHECK = b"healthcheck"
+KV_NAMESPACE_JOB = b"job"
+KV_NAMESPACE_CLUSTER = b"cluster"
+KV_HEAD_NODE_ID_KEY = b"head_node_id"
+# TODO: Set package for runtime env
+# We need to update ray client for this since runtime env use ray client
+# This might introduce some compatibility issues so leave it here for now.
+KV_NAMESPACE_PACKAGE = None
+KV_NAMESPACE_FUNCTION_TABLE = b"fun"
+
+LANGUAGE_WORKER_TYPES = ["python", "java", "cpp"]
+
+NEURON_CORES = "neuron_cores"
+GPU = "GPU"
+TPU = "TPU"
+NPU = "NPU"
+HPU = "HPU"
+
+
+RAY_WORKER_NICENESS = "RAY_worker_niceness"
+
+# Default max_retries option in @ray.remote for non-actor
+# tasks.
+DEFAULT_TASK_MAX_RETRIES = 3
+
+# Default max_concurrency option in @ray.remote for threaded actors.
+DEFAULT_MAX_CONCURRENCY_THREADED = 1
+
+# Ray internal flags. These flags should not be set by users, and we strip them on job
+# submission.
+# This should be consistent with src/ray/common/ray_internal_flag_def.h
+RAY_INTERNAL_FLAGS = [
+    "RAY_JOB_ID",
+    "RAY_RAYLET_PID",
+    "RAY_OVERRIDE_NODE_ID_FOR_TESTING",
+]
+
+DEFAULT_RESOURCES = {"CPU", "GPU", "memory", "object_store_memory"}
+
+# Supported Python versions for runtime env's "conda" field. Ray downloads
+# Ray wheels into the conda environment, so the Ray wheels for these Python
+# versions must be available online.
+RUNTIME_ENV_CONDA_PY_VERSIONS = [(3, 9), (3, 10), (3, 11), (3, 12)]
+
+# Whether to enable Ray clusters (in addition to local Ray).
+# Ray clusters are not explicitly supported for Windows and OSX.
+IS_WINDOWS_OR_OSX = sys.platform == "darwin" or sys.platform == "win32"
+ENABLE_RAY_CLUSTERS_ENV_VAR = "RAY_ENABLE_WINDOWS_OR_OSX_CLUSTER"
+ENABLE_RAY_CLUSTER = env_bool(
+    ENABLE_RAY_CLUSTERS_ENV_VAR,
+    not IS_WINDOWS_OR_OSX,
+)
+
+SESSION_LATEST = "session_latest"
+NUM_PORT_RETRIES = 40
+NUM_REDIS_GET_RETRIES = int(os.environ.get("RAY_NUM_REDIS_GET_RETRIES", "20"))
+
+# Turn this on if actor task log's offsets are expected to be recorded.
+# With this enabled, actor tasks' log could be queried with task id.
+RAY_ENABLE_RECORD_ACTOR_TASK_LOGGING = env_bool(
+    "RAY_ENABLE_RECORD_ACTOR_TASK_LOGGING", False
+)
+
+# RuntimeEnv env var to indicate it exports a function
+WORKER_PROCESS_SETUP_HOOK_ENV_VAR = "__RAY_WORKER_PROCESS_SETUP_HOOK_ENV_VAR"
+RAY_WORKER_PROCESS_SETUP_HOOK_LOAD_TIMEOUT_ENV_VAR = (
+    "RAY_WORKER_PROCESS_SETUP_HOOK_LOAD_TIMEOUT"  # noqa
+)
+
+RAY_DEFAULT_LABEL_KEYS_PREFIX = "ray.io/"
+
+RAY_TPU_MAX_CONCURRENT_CONNECTIONS_ENV_VAR = "RAY_TPU_MAX_CONCURRENT_ACTIVE_CONNECTIONS"
+
+RAY_NODE_IP_FILENAME = "node_ip_address.json"
+
+RAY_LOGGING_CONFIG_ENCODING = os.environ.get("RAY_LOGGING_CONFIG_ENCODING")
+
+RAY_BACKEND_LOG_JSON_ENV_VAR = "RAY_BACKEND_LOG_JSON"
+
+# Write export API event of all resource types to file if enabled.
+# RAY_enable_export_api_write_config will not be considered if
+# this is enabled.
+RAY_ENABLE_EXPORT_API_WRITE = env_bool("RAY_enable_export_api_write", False)
+
+# Comma separated string containing individual resource
+# to write export API events for. This configuration is only used if
+# RAY_enable_export_api_write is not enabled. Full list of valid
+# resource types in ExportEvent.SourceType enum in
+# src/ray/protobuf/export_api/export_event.proto
+# Example config:
+# `export RAY_enable_export_api_write_config='EXPORT_SUBMISSION_JOB,EXPORT_ACTOR'`
+RAY_ENABLE_EXPORT_API_WRITE_CONFIG_STR = os.environ.get(
+    "RAY_enable_export_api_write_config", ""
+)
+RAY_ENABLE_EXPORT_API_WRITE_CONFIG = RAY_ENABLE_EXPORT_API_WRITE_CONFIG_STR.split(",")
+
+RAY_EXPORT_EVENT_MAX_FILE_SIZE_BYTES = env_bool(
+    "RAY_EXPORT_EVENT_MAX_FILE_SIZE_BYTES", 100 * 1e6
+)
+
+RAY_EXPORT_EVENT_MAX_BACKUP_COUNT = env_bool("RAY_EXPORT_EVENT_MAX_BACKUP_COUNT", 20)
+
+# Comma-separated list of event types that are emitted through the Python
+# EventRecorder (One-Event Framework) to the AggregatorAgent.
+# Valid values are the names of EventType entries defined in
+# src/ray/protobuf/public/events_base_event.proto
+# Defaults to PLATFORM_EVENTS if not set.
+RAY_ENABLE_PYTHON_RAY_EVENT_TYPES = frozenset(
+    {
+        t.strip()
+        for t in os.environ.get(
+            "RAY_ENABLE_PYTHON_RAY_EVENT_TYPES", "PLATFORM_EVENT"
+        ).split(",")
+        if t.strip()
+    }
+)
+
+
+# If this flag is set and you run the driver with `uv run`, Ray propagates the `uv run`
+# environment to all workers. Ray does this by setting the `py_executable` to the
+# `uv run`` command line and by propagating the working directory
+# via the `working_dir` plugin so uv finds the pyproject.toml.
+# If you enable RAY_ENABLE_UV_RUN_RUNTIME_ENV AND you run the driver
+# with `uv run`, Ray deactivates the regular RAY_RUNTIME_ENV_HOOK
+# because in most cases the hooks wouldn't work unless you specifically make the code
+# for the runtime env hook available in your uv environment and make sure your hook
+# is compatible with your uv runtime environment. If you want to combine a custom
+# RAY_RUNTIME_ENV_HOOK with `uv run`, you should flag off RAY_ENABLE_UV_RUN_RUNTIME_ENV
+# and call ray._private.runtime_env.uv_runtime_env_hook.hook manually in your hook or
+# manually set the py_executable in your runtime environment hook.
+RAY_ENABLE_UV_RUN_RUNTIME_ENV = env_bool("RAY_ENABLE_UV_RUN_RUNTIME_ENV", True)
+
+# Prometheus metric cardinality level setting, either "legacy" or "recommended".
+#
+# Legacy: report all metrics to prometheus with the set of labels that are reported by
+#   the component, including WorkerId, (task or actor) Name, etc. This is the default.
+# Recommended: report only the node level metrics to prometheus. This means that the
+#   WorkerId will be removed from all metrics.
+# Low: Same as recommended, but also drop the Name label for tasks and actors.
+RAY_METRIC_CARDINALITY_LEVEL = os.environ.get(
+    "RAY_metric_cardinality_level", "recommended"
+)
+
+# Whether enable OpenTelemetry as the metrics collection backend. The default is
+# using OpenCensus.
+RAY_ENABLE_OPEN_TELEMETRY = env_bool("RAY_enable_open_telemetry", True)
+
+# How long to wait for a fetch for an RDT object to complete during ray.get before timing out and raising an exception to the user.
+#
+# NOTE: This is a tenth of `RayConfig::fetch_fail_timeout_milliseconds` by default as RDT transfers are expected to be much faster.
+RDT_FETCH_FAIL_TIMEOUT_SECONDS = (
+    env_integer("RAY_rdt_fetch_fail_timeout_milliseconds", 60000) / 1000
+)
+
+# Whether to enable zero-copy serialization for PyTorch tensors.
+# When enabled, Ray serializes PyTorch tensors by converting them to NumPy arrays
+# and leveraging pickle5's zero-copy buffer sharing. This avoids copying the
+# underlying tensor data, which can improve performance when passing large tensors
+# across tasks or actors. Note that this is experimental and should be used with caution
+# as we won't copy and allow a write to shared memory. One process changing a tensor
+# after ray.get could be reflected in another process.
+#
+# This feature is experimental and works best under the following conditions:
+# - The tensor has `requires_grad=False` (i.e., is detached from the autograd graph).
+# - The tensor is contiguous in memory
+# - Performance benefits from this are larger if the tensor resides in CPU memory
+# - You are not using Ray Direct Transport
+#
+# Tensors on GPU or non-contiguous tensors are still supported: Ray will
+# automatically move them to CPU and/or make them contiguous as needed.
+# While this incurs an initial copy, subsequent serialization may still benefit
+# from reduced overhead compared to the default path.
+#
+# Use with caution and ensure tensors meet the above criteria before enabling.
+# Default: False.
+RAY_ENABLE_ZERO_COPY_TORCH_TENSORS = env_bool(
+    "RAY_ENABLE_ZERO_COPY_TORCH_TENSORS", False
+)
+
+# Max number of cached NIXL remote agents. When exceeded, the least recently used
+# remote agent is evicted. When set to 0, there will be no remote agent reuse.
+NIXL_REMOTE_AGENT_CACHE_MAXSIZE = env_integer(
+    "RAY_NIXL_REMOTE_AGENT_CACHE_MAXSIZE", 1000
+)
+
+# Name of the environment variable for the Redis password.
+RAY_REDIS_PASSWORD_ENV = "RAY_REDIS_PASSWORD"