chore: import upstream snapshot with attribution

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"

python/ray/_private/ray_experimental_perf.py

@@ -0,0 +1,336 @@
+"""This is the script for `ray microbenchmark`."""
+
+import asyncio
+import logging
+import multiprocessing
+
+import ray
+import ray.experimental.channel as ray_channel
+from ray._common.utils import (
+    get_or_create_event_loop,
+)
+from ray._private.ray_microbenchmark_helpers import asyncio_timeit, timeit
+from ray._private.test_utils import get_actor_node_id
+from ray.dag import InputNode, MultiOutputNode
+from ray.dag.compiled_dag_node import CompiledDAG
+
+logger = logging.getLogger(__name__)
+
+
+@ray.remote
+class DAGActor:
+    def echo(self, x):
+        return x
+
+    def echo_multiple(self, *x):
+        return x
+
+
+def check_optimized_build():
+    if not ray._raylet.OPTIMIZED:
+        msg = (
+            "WARNING: Unoptimized build! "
+            "To benchmark an optimized build, try:\n"
+            "\tbazel run -c opt //:gen_ray_pkg\n"
+            "You can also make this permanent by adding\n"
+            "\tbuild --compilation_mode=opt\n"
+            "to your user-wide ~/.bazelrc file. "
+            "(Do not add this to the project-level .bazelrc file.)"
+        )
+        logger.warning(msg)
+
+
+def create_driver_actor():
+    return CompiledDAG.DAGDriverProxyActor.options(
+        label_selector={
+            ray._raylet.RAY_NODE_ID_KEY: ray.get_runtime_context().get_node_id()
+        }
+    ).remote()
+
+
+def main(results=None):
+    results = results or []
+    loop = get_or_create_event_loop()
+
+    check_optimized_build()
+
+    print("Tip: set TESTS_TO_RUN='pattern' to run a subset of benchmarks")
+
+    #################################################
+    # Perf tests for channels, used in compiled DAGs.
+    #################################################
+    ray.init()
+
+    def put_channel_small(chans, do_get=False):
+        for chan in chans:
+            chan.write(b"0")
+            if do_get:
+                chan.read()
+
+    @ray.remote
+    class ChannelReader:
+        def ready(self):
+            return
+
+        def read(self, chans):
+            while True:
+                for chan in chans:
+                    chan.read()
+
+    driver_actor = create_driver_actor()
+    driver_node = get_actor_node_id(driver_actor)
+    chans = [ray_channel.Channel(None, [(driver_actor, driver_node)], 1000)]
+    results += timeit(
+        "[unstable] local put:local get, single channel calls",
+        lambda: put_channel_small(chans, do_get=True),
+    )
+
+    reader = ChannelReader.remote()
+    reader_node = get_actor_node_id(reader)
+    chans = [ray_channel.Channel(None, [(reader, reader_node)], 1000)]
+    ray.get(reader.ready.remote())
+    reader.read.remote(chans)
+    results += timeit(
+        "[unstable] local put:1 remote get, single channel calls",
+        lambda: put_channel_small(chans),
+    )
+    ray.kill(reader)
+
+    n_cpu = multiprocessing.cpu_count() // 2
+    print(f"Testing multiple readers/channels, n={n_cpu}")
+
+    reader_and_node_list = []
+    for _ in range(n_cpu):
+        reader = ChannelReader.remote()
+        reader_node = get_actor_node_id(reader)
+        reader_and_node_list.append((reader, reader_node))
+    chans = [ray_channel.Channel(None, reader_and_node_list, 1000)]
+    ray.get([reader.ready.remote() for reader, _ in reader_and_node_list])
+    for reader, _ in reader_and_node_list:
+        reader.read.remote(chans)
+    results += timeit(
+        "[unstable] local put:n remote get, single channel calls",
+        lambda: put_channel_small(chans),
+    )
+    for reader, _ in reader_and_node_list:
+        ray.kill(reader)
+
+    reader = ChannelReader.remote()
+    reader_node = get_actor_node_id(reader)
+    chans = [
+        ray_channel.Channel(None, [(reader, reader_node)], 1000) for _ in range(n_cpu)
+    ]
+    ray.get(reader.ready.remote())
+    reader.read.remote(chans)
+    results += timeit(
+        "[unstable] local put:1 remote get, n channels calls",
+        lambda: put_channel_small(chans),
+    )
+    ray.kill(reader)
+
+    reader_and_node_list = []
+    for _ in range(n_cpu):
+        reader = ChannelReader.remote()
+        reader_node = get_actor_node_id(reader)
+        reader_and_node_list.append((reader, reader_node))
+    chans = [
+        ray_channel.Channel(None, [reader_and_node_list[i]], 1000) for i in range(n_cpu)
+    ]
+    ray.get([reader.ready.remote() for reader, _ in reader_and_node_list])
+    for chan, reader_node_tuple in zip(chans, reader_and_node_list):
+        reader = reader_node_tuple[0]
+        reader.read.remote([chan])
+    results += timeit(
+        "[unstable] local put:n remote get, n channels calls",
+        lambda: put_channel_small(chans),
+    )
+    for reader, _ in reader_and_node_list:
+        ray.kill(reader)
+
+    # Tests for compiled DAGs.
+
+    def _exec(dag, num_args=1, payload_size=1):
+        output_ref = dag.execute(*[b"x" * payload_size for _ in range(num_args)])
+        ray.get(output_ref)
+
+    async def exec_async(tag):
+        async def _exec_async():
+            fut = await compiled_dag.execute_async(b"x")
+            if not isinstance(fut, list):
+                await fut
+            else:
+                await asyncio.gather(*fut)
+
+        return await asyncio_timeit(
+            tag,
+            _exec_async,
+        )
+
+    # Single-actor DAG calls
+
+    a = DAGActor.remote()
+    with InputNode() as inp:
+        dag = a.echo.bind(inp)
+
+    results += timeit(
+        "[unstable] single-actor DAG calls", lambda: ray.get(dag.execute(b"x"))
+    )
+    compiled_dag = dag.experimental_compile()
+    results += timeit(
+        "[unstable] compiled single-actor DAG calls", lambda: _exec(compiled_dag)
+    )
+    del a
+
+    # Single-actor asyncio DAG calls
+
+    a = DAGActor.remote()
+    with InputNode() as inp:
+        dag = a.echo.bind(inp)
+    compiled_dag = dag.experimental_compile(enable_asyncio=True)
+    results += loop.run_until_complete(
+        exec_async(
+            "[unstable] compiled single-actor asyncio DAG calls",
+        )
+    )
+    del a
+
+    # Scatter-gather DAG calls
+
+    n_cpu = multiprocessing.cpu_count() // 2
+    actors = [DAGActor.remote() for _ in range(n_cpu)]
+    with InputNode() as inp:
+        dag = MultiOutputNode([a.echo.bind(inp) for a in actors])
+    results += timeit(
+        f"[unstable] scatter-gather DAG calls, n={n_cpu} actors",
+        lambda: ray.get(dag.execute(b"x")),
+    )
+    compiled_dag = dag.experimental_compile()
+    results += timeit(
+        f"[unstable] compiled scatter-gather DAG calls, n={n_cpu} actors",
+        lambda: _exec(compiled_dag),
+    )
+
+    # Scatter-gather asyncio DAG calls
+
+    actors = [DAGActor.remote() for _ in range(n_cpu)]
+    with InputNode() as inp:
+        dag = MultiOutputNode([a.echo.bind(inp) for a in actors])
+    compiled_dag = dag.experimental_compile(enable_asyncio=True)
+    results += loop.run_until_complete(
+        exec_async(
+            f"[unstable] compiled scatter-gather asyncio DAG calls, n={n_cpu} actors",
+        )
+    )
+
+    # Chain DAG calls
+
+    actors = [DAGActor.remote() for _ in range(n_cpu)]
+    with InputNode() as inp:
+        dag = inp
+        for a in actors:
+            dag = a.echo.bind(dag)
+    results += timeit(
+        f"[unstable] chain DAG calls, n={n_cpu} actors",
+        lambda: ray.get(dag.execute(b"x")),
+    )
+    compiled_dag = dag.experimental_compile()
+    results += timeit(
+        f"[unstable] compiled chain DAG calls, n={n_cpu} actors",
+        lambda: _exec(compiled_dag),
+    )
+
+    # Chain asyncio DAG calls
+
+    actors = [DAGActor.remote() for _ in range(n_cpu)]
+    with InputNode() as inp:
+        dag = inp
+        for a in actors:
+            dag = a.echo.bind(dag)
+    compiled_dag = dag.experimental_compile(enable_asyncio=True)
+    results += loop.run_until_complete(
+        exec_async(f"[unstable] compiled chain asyncio DAG calls, n={n_cpu} actors")
+    )
+
+    # Multiple args with small payloads
+
+    n_actors = 8
+    assert (
+        n_cpu > n_actors
+    ), f"n_cpu ({n_cpu}) must be greater than n_actors ({n_actors})"
+
+    actors = [DAGActor.remote() for _ in range(n_actors)]
+    with InputNode() as inp:
+        dag = MultiOutputNode([actors[i].echo.bind(inp[i]) for i in range(n_actors)])
+    payload_size = 1
+    results += timeit(
+        f"[unstable] multiple args with small payloads DAG calls, n={n_actors} actors",
+        lambda: ray.get(dag.execute(*[b"x" * payload_size for _ in range(n_actors)])),
+    )
+    compiled_dag = dag.experimental_compile()
+    results += timeit(
+        f"[unstable] compiled multiple args with small payloads DAG calls, "
+        f"n={n_actors} actors",
+        lambda: _exec(compiled_dag, num_args=n_actors, payload_size=payload_size),
+    )
+
+    # Multiple args with medium payloads
+
+    actors = [DAGActor.remote() for _ in range(n_actors)]
+    with InputNode() as inp:
+        dag = MultiOutputNode([actors[i].echo.bind(inp[i]) for i in range(n_actors)])
+    payload_size = 1024 * 1024
+    results += timeit(
+        f"[unstable] multiple args with medium payloads DAG calls, n={n_actors} actors",
+        lambda: ray.get(dag.execute(*[b"x" * payload_size for _ in range(n_actors)])),
+    )
+    compiled_dag = dag.experimental_compile()
+    results += timeit(
+        "[unstable] compiled multiple args with medium payloads DAG calls, "
+        f"n={n_actors} actors",
+        lambda: _exec(compiled_dag, num_args=n_actors, payload_size=payload_size),
+    )
+
+    # Multiple args with large payloads
+
+    actors = [DAGActor.remote() for _ in range(n_actors)]
+    with InputNode() as inp:
+        dag = MultiOutputNode([actors[i].echo.bind(inp[i]) for i in range(n_actors)])
+    payload_size = 10 * 1024 * 1024
+    results += timeit(
+        f"[unstable] multiple args with large payloads DAG calls, n={n_actors} actors",
+        lambda: ray.get(dag.execute(*[b"x" * payload_size for _ in range(n_actors)])),
+    )
+    compiled_dag = dag.experimental_compile()
+    results += timeit(
+        "[unstable] compiled multiple args with large payloads DAG calls, "
+        f"n={n_actors} actors",
+        lambda: _exec(compiled_dag, num_args=n_actors, payload_size=payload_size),
+    )
+
+    # Worst case for multiple arguments: a single actor takes all the arguments
+    # with small payloads.
+
+    actor = DAGActor.remote()
+    n_args = 8
+    with InputNode() as inp:
+        dag = actor.echo_multiple.bind(*[inp[i] for i in range(n_args)])
+    payload_size = 1
+    results += timeit(
+        "[unstable] single-actor with all args with small payloads DAG calls, "
+        "n=1 actors",
+        lambda: ray.get(dag.execute(*[b"x" * payload_size for _ in range(n_args)])),
+    )
+    compiled_dag = dag.experimental_compile()
+    results += timeit(
+        "[unstable] single-actor with all args with small payloads DAG calls, "
+        "n=1 actors",
+        lambda: _exec(compiled_dag, num_args=n_args, payload_size=payload_size),
+    )
+
+    ray.shutdown()
+
+    return results
+
+
+if __name__ == "__main__":
+    main()

python/ray/_private/ray_logging/__init__.py

@@ -0,0 +1,380 @@
+import logging
+import logging.handlers
+import os
+import re
+import sys
+import threading
+import time
+from dataclasses import dataclass
+from typing import Any, Callable, Dict, Iterable, List, Optional, Set, Tuple, Union
+
+import colorama
+
+import ray
+from ray._private.ray_constants import (
+    RAY_DEDUP_LOGS,
+    RAY_DEDUP_LOGS_AGG_WINDOW_S,
+    RAY_DEDUP_LOGS_ALLOW_REGEX,
+    RAY_DEDUP_LOGS_SKIP_REGEX,
+)
+from ray.experimental.tqdm_ray import RAY_TQDM_MAGIC
+from ray.util.debug import log_once
+
+
+def setup_logger(
+    logging_level: int,
+    logging_format: str,
+):
+    """Setup default logging for ray."""
+    logger = logging.getLogger("ray")
+    if logging_format:
+        # Overwrite the formatters for all default handlers.
+        formatter = logging.Formatter(logging_format)
+        for handler in logger.handlers:
+            handler.setFormatter(formatter)
+    if isinstance(logging_level, str):
+        logging_level = logging.getLevelName(logging_level.upper())
+    logger.setLevel(logging_level)
+
+
+def setup_component_logger(
+    *,
+    logging_level: Union[int, str],
+    logging_format: str,
+    log_dir: str,
+    filename: Union[str, Iterable[str]],
+    max_bytes: int,
+    backup_count: int,
+    logger_name: Optional[str] = None,
+    propagate: bool = True,
+):
+    """Configure the logger that is used for Ray's python components.
+
+    For example, it should be used for monitor, dashboard, and log monitor.
+    The only exception is workers. They use the different logging config.
+
+    Ray's python components generally should not write to stdout/stderr, because
+    messages written there will be redirected to the head node. For deployments where
+    there may be thousands of workers, this would create unacceptable levels of log
+    spam. For this reason, we disable the "ray" logger's handlers, and enable
+    propagation so that log messages that actually do need to be sent to the head node
+    can reach it.
+
+    Args:
+        logging_level: Logging level in string or logging enum.
+        logging_format: Logging format string.
+        log_dir: Log directory path. If empty, logs will go to
+            stderr.
+        filename: A single filename or an iterable of filenames to write logs to.
+            If empty, logs will go to stderr.
+        max_bytes: Same argument as RotatingFileHandler's maxBytes.
+        backup_count: Same argument as RotatingFileHandler's backupCount.
+        logger_name: Used to create or get the corresponding
+            logger in getLogger call. It will get the root logger by default.
+        propagate: Whether to propagate the log to the parent logger.
+    Returns:
+        the created or modified logger.
+    """
+    ray._private.log.clear_logger("ray")
+
+    logger = logging.getLogger(logger_name)
+    if isinstance(logging_level, str):
+        logging_level = logging.getLevelName(logging_level.upper())
+    logger.setLevel(logging_level)
+
+    filenames = [filename] if isinstance(filename, str) else filename
+
+    for filename in filenames:
+        if not filename or not log_dir:
+            handler = logging.StreamHandler()
+        else:
+            handler = logging.handlers.RotatingFileHandler(
+                os.path.join(log_dir, filename),
+                maxBytes=max_bytes,
+                backupCount=backup_count,
+            )
+        handler.setLevel(logging_level)
+        handler.setFormatter(logging.Formatter(logging_format))
+        logger.addHandler(handler)
+
+    logger.propagate = propagate
+    return logger
+
+
+def run_callback_on_events_in_ipython(event: str, cb: Callable):
+    """
+    Register a callback to be run after each cell completes in IPython.
+    E.g.:
+        This is used to flush the logs after each cell completes.
+
+    If IPython is not installed, this function does nothing.
+
+    Args:
+        event: The IPython event to subscribe to (e.g. ``post_run_cell``).
+        cb: The callback to run.
+    """
+    if "IPython" in sys.modules:
+        from IPython import get_ipython
+
+        ipython = get_ipython()
+        # Register a callback on cell completion.
+        if ipython is not None:
+            ipython.events.register(event, cb)
+
+
+"""
+All components underneath here is used specifically for the default_worker.py.
+"""
+
+# It's worth noticing that filepath format should be kept in sync with function
+# `GetWorkerOutputFilepath` under file "src/ray/core_worker/core_worker_process.cc".
+def get_worker_log_file_name(worker_type, job_id=None):
+    if job_id is None:
+        job_id = os.environ.get("RAY_JOB_ID")
+    if worker_type == "WORKER":
+        if job_id is None:
+            job_id = ""
+        worker_name = "worker"
+    else:
+        job_id = ""
+        worker_name = "io_worker"
+
+    # Make sure these values are set already.
+    assert ray._private.worker._global_node is not None
+    assert ray._private.worker.global_worker is not None
+    filename = f"{worker_name}-{ray.get_runtime_context().get_worker_id()}-"
+    if job_id:
+        filename += f"{job_id}-"
+    filename += f"{os.getpid()}"
+    return filename
+
+
+def configure_log_file(out_file, err_file):
+    # If either of the file handles are None, there are no log files to
+    # configure since we're redirecting all output to stdout and stderr.
+    if out_file is None or err_file is None:
+        return
+    stdout_fileno = sys.stdout.fileno()
+    stderr_fileno = sys.stderr.fileno()
+    # C++ logging requires redirecting the stdout file descriptor. Note that
+    # dup2 will automatically close the old file descriptor before overriding
+    # it.
+    os.dup2(out_file.fileno(), stdout_fileno)
+    os.dup2(err_file.fileno(), 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 = ray._private.utils.open_log(
+        stdout_fileno, unbuffered=True, closefd=False
+    )
+    sys.stderr = ray._private.utils.open_log(
+        stderr_fileno, unbuffered=True, closefd=False
+    )
+
+
+class WorkerStandardStreamDispatcher:
+    def __init__(self):
+        self.handlers = []
+        self._lock = threading.Lock()
+
+    def add_handler(self, name: str, handler: Callable) -> None:
+        with self._lock:
+            self.handlers.append((name, handler))
+
+    def remove_handler(self, name: str) -> None:
+        with self._lock:
+            new_handlers = [pair for pair in self.handlers if pair[0] != name]
+            self.handlers = new_handlers
+
+    def emit(self, data):
+        with self._lock:
+            for pair in self.handlers:
+                _, handle = pair
+                handle(data)
+
+
+global_worker_stdstream_dispatcher = WorkerStandardStreamDispatcher()
+
+
+# Regex for canonicalizing log lines.
+NUMBERS = re.compile(r"(\d+|0x[0-9a-fA-F]+)")
+
+# Batch of log lines including ip, pid, lines, etc.
+LogBatch = Dict[str, Any]
+
+
+def _canonicalise_log_line(line):
+    # Remove words containing numbers or hex, since those tend to differ between
+    # workers.
+    return " ".join(x for x in line.split() if not NUMBERS.search(x))
+
+
+@dataclass
+class DedupState:
+    # Timestamp of the earliest log message seen of this pattern.
+    timestamp: int
+
+    # The number of un-printed occurrences for this pattern.
+    count: int
+
+    # Latest instance of this log pattern.
+    line: int
+
+    # Latest metadata dict for this log pattern, not including the lines field.
+    metadata: LogBatch
+
+    # Set of (ip, pid) sources which have emitted this pattern.
+    sources: Set[Tuple[str, int]]
+
+    # The string that should be printed to stdout.
+    def formatted(self) -> str:
+        return self.line + _color(
+            f" [repeated {self.count}x across cluster]" + _warn_once()
+        )
+
+
+class LogDeduplicator:
+    def __init__(
+        self,
+        agg_window_s: int,
+        allow_re: Optional[str],
+        skip_re: Optional[str],
+        *,
+        _timesource=None,
+    ):
+        self.agg_window_s = agg_window_s
+        if allow_re:
+            self.allow_re = re.compile(allow_re)
+        else:
+            self.allow_re = None
+        if skip_re:
+            self.skip_re = re.compile(skip_re)
+        else:
+            self.skip_re = None
+        # Buffer of up to RAY_DEDUP_LOGS_AGG_WINDOW_S recent log patterns.
+        # This buffer is cleared if the pattern isn't seen within the window.
+        self.recent: Dict[str, DedupState] = {}
+        self.timesource = _timesource or (lambda: time.time())
+
+        run_callback_on_events_in_ipython("post_execute", self.flush)
+
+    def deduplicate(self, batch: LogBatch) -> List[LogBatch]:
+        """Rewrite a batch of lines to reduce duplicate log messages.
+
+        Args:
+            batch: The batch of lines from a single source.
+
+        Returns:
+            List of batches from this and possibly other previous sources to print.
+        """
+        if not RAY_DEDUP_LOGS:
+            return [batch]
+
+        now = self.timesource()
+        metadata = batch.copy()
+        del metadata["lines"]
+        source = (metadata.get("ip"), metadata.get("pid"))
+        output: List[LogBatch] = [dict(**metadata, lines=[])]
+
+        # Decide which lines to emit from the input batch. Put the outputs in the
+        # first output log batch (output[0]).
+        for line in batch["lines"]:
+            if RAY_TQDM_MAGIC in line or (self.allow_re and self.allow_re.search(line)):
+                output[0]["lines"].append(line)
+                continue
+            elif self.skip_re and self.skip_re.search(line):
+                continue
+            dedup_key = _canonicalise_log_line(line)
+
+            if dedup_key == "":
+                # Don't dedup messages that are empty after canonicalization.
+                # Because that's all the information users want to see.
+                output[0]["lines"].append(line)
+                continue
+
+            if dedup_key in self.recent:
+                sources = self.recent[dedup_key].sources
+                sources.add(source)
+                # We deduplicate the warnings/error messages from raylet by default.
+                if len(sources) > 1 or batch["pid"] == "raylet":
+                    state = self.recent[dedup_key]
+                    self.recent[dedup_key] = DedupState(
+                        state.timestamp,
+                        state.count + 1,
+                        line,
+                        metadata,
+                        sources,
+                    )
+                else:
+                    # Don't dedup messages from the same source, just print.
+                    output[0]["lines"].append(line)
+            else:
+                self.recent[dedup_key] = DedupState(now, 0, line, metadata, {source})
+                output[0]["lines"].append(line)
+
+        # Flush patterns from the buffer that are older than the aggregation window.
+        while self.recent:
+            if now - next(iter(self.recent.values())).timestamp < self.agg_window_s:
+                break
+            dedup_key = next(iter(self.recent))
+            state = self.recent.pop(dedup_key)
+            # we already logged an instance of this line immediately when received,
+            # so don't log for count == 0
+            if state.count > 1:
+                # (Actor pid=xxxx) [repeated 2x across cluster] ...
+                output.append(dict(**state.metadata, lines=[state.formatted()]))
+                # Continue aggregating for this key but reset timestamp and count.
+                state.timestamp = now
+                state.count = 0
+                self.recent[dedup_key] = state
+            elif state.count > 0:
+                # Aggregation wasn't fruitful, print the line and stop aggregating.
+                output.append(dict(state.metadata, lines=[state.line]))
+
+        return output
+
+    def flush(self) -> List[dict]:
+        """Return all buffered log messages and clear the buffer.
+
+        Returns:
+            List of log batches to print.
+        """
+        output = []
+        for state in self.recent.values():
+            if state.count > 1:
+                output.append(
+                    dict(
+                        state.metadata,
+                        lines=[state.formatted()],
+                    )
+                )
+            elif state.count > 0:
+                output.append(dict(state.metadata, **{"lines": [state.line]}))
+        self.recent.clear()
+        return output
+
+
+def _warn_once() -> str:
+    if log_once("log_dedup_warning"):
+        return (
+            " (Ray deduplicates logs by default. Set RAY_DEDUP_LOGS=0 to "
+            "disable log deduplication, or see https://docs.ray.io/en/master/"
+            "ray-observability/user-guides/configure-logging.html#log-deduplication "
+            "for more options.)"
+        )
+    else:
+        return ""
+
+
+def _color(msg: str) -> str:
+    return "{}{}{}".format(colorama.Fore.GREEN, msg, colorama.Style.RESET_ALL)
+
+
+stdout_deduplicator = LogDeduplicator(
+    RAY_DEDUP_LOGS_AGG_WINDOW_S, RAY_DEDUP_LOGS_ALLOW_REGEX, RAY_DEDUP_LOGS_SKIP_REGEX
+)
+stderr_deduplicator = LogDeduplicator(
+    RAY_DEDUP_LOGS_AGG_WINDOW_S, RAY_DEDUP_LOGS_ALLOW_REGEX, RAY_DEDUP_LOGS_SKIP_REGEX
+)

python/ray/_private/ray_logging/default_impl.py

@@ -0,0 +1,4 @@
+def get_logging_configurator():
+    from ray._private.ray_logging.logging_config import DefaultLoggingConfigurator
+
+    return DefaultLoggingConfigurator()

python/ray/_private/ray_logging/logging_config.py

@@ -0,0 +1,171 @@
+import logging
+from abc import ABC, abstractmethod
+from dataclasses import asdict, dataclass, field, fields
+from typing import Dict, Set
+
+from ray._common.filters import CoreContextFilter
+from ray._common.formatters import JSONFormatter, TextFormatter
+from ray._common.logging_constants import LOGRECORD_STANDARD_ATTRS
+from ray._private.ray_logging import default_impl
+from ray.util.annotations import PublicAPI
+
+
+class LoggingConfigurator(ABC):
+    @abstractmethod
+    def get_supported_encodings(self) -> Set[str]:
+        raise NotImplementedError
+
+    @abstractmethod
+    def configure(self, logging_config: "LoggingConfig"):
+        raise NotImplementedError
+
+
+class DefaultLoggingConfigurator(LoggingConfigurator):
+    def __init__(self):
+        self._encoding_to_formatter = {
+            "TEXT": TextFormatter(),
+            "JSON": JSONFormatter(),
+        }
+
+    def get_supported_encodings(self) -> Set[str]:
+        return self._encoding_to_formatter.keys()
+
+    def configure(self, logging_config: "LoggingConfig"):
+        formatter = self._encoding_to_formatter[logging_config.encoding]
+        formatter.set_additional_log_standard_attrs(
+            logging_config.additional_log_standard_attrs
+        )
+
+        core_context_filter = CoreContextFilter()
+        handler = logging.StreamHandler()
+        handler.setLevel(logging_config.log_level)
+        handler.setFormatter(formatter)
+        handler.addFilter(core_context_filter)
+
+        root_logger = logging.getLogger()
+        root_logger.setLevel(logging_config.log_level)
+        root_logger.addHandler(handler)
+
+        ray_logger = logging.getLogger("ray")
+        ray_logger.setLevel(logging_config.log_level)
+        # Remove all existing handlers added by `ray/__init__.py`.
+        for h in ray_logger.handlers[:]:
+            ray_logger.removeHandler(h)
+        ray_logger.addHandler(handler)
+        ray_logger.propagate = False
+
+
+_logging_configurator: LoggingConfigurator = default_impl.get_logging_configurator()
+
+
+# Class defines the logging configurations for a Ray job.
+# To add a new logging configuration: (1) add a new field to this class; (2) Update the
+# logic in the __post_init__ method in this class to add the validation logic;
+# (3) Update the configure method in the DefaultLoggingConfigurator
+# class to use the new field.
+@PublicAPI(stability="alpha")
+@dataclass
+class LoggingConfig:
+    encoding: str = "TEXT"
+    log_level: str = "INFO"
+    # The list of valid attributes are defined as LOGRECORD_STANDARD_ATTRS in
+    # constants.py.
+    additional_log_standard_attrs: list = field(default_factory=list)
+
+    def __post_init__(self):
+        if self.encoding not in _logging_configurator.get_supported_encodings():
+            raise ValueError(
+                f"Invalid encoding type: {self.encoding}. "
+                "Valid encoding types are: "
+                f"{list(_logging_configurator.get_supported_encodings())}"
+            )
+
+        for attr in self.additional_log_standard_attrs:
+            if attr not in LOGRECORD_STANDARD_ATTRS:
+                raise ValueError(
+                    f"Unknown python logging standard attribute: {attr}. "
+                    "The valid attributes are: "
+                    f"{set(LOGRECORD_STANDARD_ATTRS)}"
+                )
+
+    def to_dict(self) -> Dict[str, object]:
+        """Serialize to a plain dict suitable for JSON transport."""
+        return asdict(self)
+
+    @classmethod
+    def from_dict(cls, d: Dict[str, object]) -> "LoggingConfig":
+        """Create a LoggingConfig from a dict, ignoring unknown keys."""
+        known = {f.name for f in fields(cls)}
+        return cls(**{k: v for k, v in d.items() if k in known})
+
+    def _configure_logging(self):
+        """Set up the logging configuration for the current process."""
+        _logging_configurator.configure(self)
+
+    def _apply(self):
+        """Set up the logging configuration."""
+        self._configure_logging()
+
+
+LoggingConfig.__doc__ = """
+    Logging configuration for a Ray job. These configurations are used to set up the
+    root logger of the driver process and all Ray tasks and actor processes that belong
+    to the job.
+
+    Examples: 1. Configure the logging to use TEXT encoding.
+        .. testcode::
+
+            import ray
+            import logging
+
+            ray.init(
+                logging_config=ray.LoggingConfig(encoding="TEXT", log_level="INFO", additional_log_standard_attrs=['name'])
+            )
+
+            @ray.remote
+            def f():
+                logger = logging.getLogger(__name__)
+                logger.info("This is a Ray task")
+
+            ray.get(f.remote())
+            ray.shutdown()
+
+        .. testoutput::
+            :options: +MOCK
+
+            2025-02-12 12:25:16,836 INFO test-log-config.py:11 -- This is a Ray task name=__main__ job_id=01000000 worker_id=51188d9448be4664bf2ea26ac410b67acaaa970c4f31c5ad3ae776a5 node_id=f683dfbffe2c69984859bc19c26b77eaf3866c458884c49d115fdcd4 task_id=c8ef45ccd0112571ffffffffffffffffffffffff01000000 task_name=f task_func_name=test-log-config.f timestamp_ns=1739391916836884000
+
+    2. Configure the logging to use JSON encoding.
+        .. testcode::
+
+            import ray
+            import logging
+
+            ray.init(
+                logging_config=ray.LoggingConfig(encoding="JSON", log_level="INFO", additional_log_standard_attrs=['name'])
+            )
+
+            @ray.remote
+            def f():
+                logger = logging.getLogger(__name__)
+                logger.info("This is a Ray task")
+
+            ray.get(f.remote())
+            ray.shutdown()
+
+        .. testoutput::
+            :options: +MOCK
+
+            {"asctime": "2025-02-12 12:25:48,766", "levelname": "INFO", "message": "This is a Ray task", "filename": "test-log-config.py", "lineno": 11, "name": "__main__", "job_id": "01000000", "worker_id": "6d307578014873fcdada0fa22ea6d49e0fb1f78960e69d61dfe41f5a", "node_id": "69e3a5e68bdc7eb8ac9abb3155326ee3cc9fc63ea1be04d11c0d93c7", "task_id": "c8ef45ccd0112571ffffffffffffffffffffffff01000000", "task_name": "f", "task_func_name": "test-log-config.f", "timestamp_ns": 1739391948766949000}
+
+    Args:
+        encoding: Encoding type for the logs. The valid values are
+            {list(_logging_configurator.get_supported_encodings())}
+        log_level: Log level for the logs. Defaults to 'INFO'. You can set
+            it to 'DEBUG' to receive more detailed debug logs.
+        additional_log_standard_attrs: List of additional standard python logger attributes to
+            include in the log. Defaults to an empty list. The list of already
+            included standard attributes are: "asctime", "levelname", "message",
+            "filename", "lineno", "exc_text". The list of valid attributes are specified
+            here: http://docs.python.org/library/logging.html#logrecord-attributes
+    """  # noqa: E501

python/ray/_private/ray_microbenchmark_helpers.py

@@ -0,0 +1,92 @@
+import os
+import time
+from contextlib import contextmanager
+from typing import List, Optional, Tuple
+
+import numpy as np
+
+import ray
+
+# Only run tests matching this filter pattern.
+
+filter_pattern = os.environ.get("TESTS_TO_RUN", "")
+skip_pattern = os.environ.get("TESTS_TO_SKIP", "")
+
+
+def timeit(
+    name, fn, multiplier=1, warmup_time_sec=10
+) -> List[Optional[Tuple[str, float, float]]]:
+    if filter_pattern and filter_pattern not in name:
+        return [None]
+    if skip_pattern and skip_pattern in name:
+        return [None]
+    # sleep for a while to avoid noisy neigbhors.
+    # related issue: https://github.com/ray-project/ray/issues/22045
+    time.sleep(warmup_time_sec)
+    # warmup
+    start = time.perf_counter()
+    count = 0
+    while time.perf_counter() - start < 1:
+        fn()
+        count += 1
+    # real run
+    step = count // 10 + 1
+    stats = []
+    for _ in range(4):
+        start = time.perf_counter()
+        count = 0
+        while time.perf_counter() - start < 2:
+            for _ in range(step):
+                fn()
+            count += step
+        end = time.perf_counter()
+        stats.append(multiplier * count / (end - start))
+
+    mean = np.mean(stats)
+    sd = np.std(stats)
+    print(name, "per second", round(mean, 2), "+-", round(sd, 2))
+    return [(name, mean, sd)]
+
+
+async def asyncio_timeit(
+    name, async_fn, multiplier=1, warmup_time_sec=10
+) -> List[Optional[Tuple[str, float, float]]]:
+    if filter_pattern and filter_pattern not in name:
+        return [None]
+    if skip_pattern and skip_pattern in name:
+        return [None]
+    # sleep for a while to avoid noisy neigbhors.
+    # related issue: https://github.com/ray-project/ray/issues/22045
+    time.sleep(warmup_time_sec)
+    # warmup
+    start = time.perf_counter()
+    count = 0
+    while time.perf_counter() - start < 1:
+        await async_fn()
+        count += 1
+    # real run
+    step = count // 10 + 1
+    stats = []
+    for _ in range(4):
+        start = time.perf_counter()
+        count = 0
+        while time.perf_counter() - start < 2:
+            for _ in range(step):
+                await async_fn()
+            count += step
+        end = time.perf_counter()
+        stats.append(multiplier * count / (end - start))
+
+    mean = np.mean(stats)
+    sd = np.std(stats)
+    print(name, "per second", round(mean, 2), "+-", round(sd, 2))
+    return [(name, mean, sd)]
+
+
+@contextmanager
+def ray_setup_and_teardown(**init_args):
+    ray.init(**init_args)
+    try:
+        yield None
+    finally:
+        ray.shutdown()

python/ray/_private/ray_perf.py

@@ -0,0 +1,330 @@
+"""This is the script for `ray microbenchmark`."""
+
+import asyncio
+import logging
+import multiprocessing
+
+import numpy as np
+
+import ray
+from ray._private.ray_client_microbenchmark import main as client_microbenchmark_main
+from ray._private.ray_microbenchmark_helpers import timeit
+
+logger = logging.getLogger(__name__)
+
+
+@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([small_value.remote() for _ in range(n)])
+
+
+@ray.remote
+class AsyncActor:
+    async def small_value(self):
+        return b"ok"
+
+    async def small_value_with_arg(self, x):
+        return b"ok"
+
+    async def small_value_batch(self, n):
+        await asyncio.wait([small_value.remote() for _ in range(n)])
+
+
+@ray.remote(num_cpus=0)
+class Client:
+    def __init__(self, servers):
+        if not isinstance(servers, list):
+            servers = [servers]
+        self.servers = servers
+
+    def small_value_batch(self, n):
+        results = []
+        for s in self.servers:
+            results.extend([s.small_value.remote() for _ in range(n)])
+        ray.get(results)
+
+    def small_value_batch_arg(self, n):
+        x = ray.put(0)
+        results = []
+        for s in self.servers:
+            results.extend([s.small_value_arg.remote(x) for _ in range(n)])
+        ray.get(results)
+
+
+@ray.remote
+def small_value():
+    return b"ok"
+
+
+@ray.remote
+def small_value_batch(n):
+    submitted = [small_value.remote() for _ in range(n)]
+    ray.get(submitted)
+    return 0
+
+
+@ray.remote
+def create_object_containing_ref():
+    obj_refs = []
+    for _ in range(10000):
+        obj_refs.append(ray.put(1))
+    return obj_refs
+
+
+def check_optimized_build():
+    if not ray._raylet.OPTIMIZED:
+        msg = (
+            "WARNING: Unoptimized build! "
+            "To benchmark an optimized build, try:\n"
+            "\tbazel run -c opt //:gen_ray_pkg\n"
+            "You can also make this permanent by adding\n"
+            "\tbuild --compilation_mode=opt\n"
+            "to your user-wide ~/.bazelrc file. "
+            "(Do not add this to the project-level .bazelrc file.)"
+        )
+        logger.warning(msg)
+
+
+def main(results=None):
+    results = results or []
+
+    check_optimized_build()
+
+    print("Tip: set TESTS_TO_RUN='pattern' to run a subset of benchmarks")
+
+    ray.init()
+
+    value = ray.put(0)
+
+    def get_small():
+        ray.get(value)
+
+    def put_small():
+        ray.put(0)
+
+    @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)])
+
+    arr = np.zeros(100 * 1024 * 1024, dtype=np.int64)
+
+    results += timeit("single client get calls (Plasma Store)", get_small)
+
+    results += timeit("single client put calls (Plasma Store)", put_small)
+
+    results += timeit("multi client put calls (Plasma Store)", put_multi_small, 1000)
+
+    def put_large():
+        ray.put(arr)
+
+    results += timeit("single client put gigabytes", put_large, 8 * 0.1)
+
+    def small_value_batch():
+        submitted = [small_value.remote() for _ in range(1000)]
+        ray.get(submitted)
+        return 0
+
+    results += timeit("single client tasks and get batch", small_value_batch)
+
+    @ray.remote
+    def do_put():
+        for _ in range(10):
+            ray.put(np.zeros(10 * 1024 * 1024, dtype=np.int64))
+
+    def put_multi():
+        ray.get([do_put.remote() for _ in range(10)])
+
+    results += timeit("multi client put gigabytes", put_multi, 10 * 8 * 0.1)
+
+    obj_containing_ref = create_object_containing_ref.remote()
+
+    def get_containing_object_ref():
+        ray.get(obj_containing_ref)
+
+    results += timeit(
+        "single client get object containing 10k refs", get_containing_object_ref
+    )
+
+    def wait_multiple_refs():
+        num_objs = 1000
+        not_ready = [small_value.remote() for _ in range(num_objs)]
+        # We only need to trigger the fetch_local once for each object,
+        # raylet will persist these fetch requests even after ray.wait returns.
+        # See https://github.com/ray-project/ray/issues/30375.
+        fetch_local = True
+        for _ in range(num_objs):
+            _ready, not_ready = ray.wait(not_ready, fetch_local=fetch_local)
+            if fetch_local:
+                fetch_local = False
+
+    results += timeit("single client wait 1k refs", wait_multiple_refs)
+
+    def small_task():
+        ray.get(small_value.remote())
+
+    results += timeit("single client tasks sync", small_task)
+
+    def small_task_async():
+        ray.get([small_value.remote() for _ in range(1000)])
+
+    results += timeit("single client tasks async", small_task_async, 1000)
+
+    n = 10000
+    m = 4
+    actors = [Actor.remote() for _ in range(m)]
+
+    def multi_task():
+        submitted = [a.small_value_batch.remote(n) for a in actors]
+        ray.get(submitted)
+
+    results += timeit("multi client tasks async", multi_task, n * m)
+
+    a = Actor.remote()
+
+    def actor_sync():
+        ray.get(a.small_value.remote())
+
+    results += timeit("1:1 actor calls sync", actor_sync)
+
+    a = Actor.remote()
+
+    def actor_async():
+        ray.get([a.small_value.remote() for _ in range(1000)])
+
+    results += timeit("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("1:1 actor calls concurrent", actor_concurrent, 1000)
+
+    n = 5000
+    n_cpu = multiprocessing.cpu_count() // 2
+    actors = [Actor._remote() for _ in range(n_cpu)]
+    client = Client.remote(actors)
+
+    def actor_async_direct():
+        ray.get(client.small_value_batch.remote(n))
+
+    results += timeit("1:n actor calls async", actor_async_direct, n * len(actors))
+
+    n_cpu = multiprocessing.cpu_count() // 2
+    a = [Actor.remote() for _ in range(n_cpu)]
+
+    @ray.remote
+    def work(actors):
+        ray.get([actors[i % n_cpu].small_value.remote() for i in range(n)])
+
+    def actor_multi2():
+        ray.get([work.remote(a) for _ in range(m)])
+
+    results += timeit("n:n actor calls async", actor_multi2, m * n)
+
+    n = 1000
+    actors = [Actor._remote() for _ in range(n_cpu)]
+    clients = [Client.remote(a) for a in actors]
+
+    def actor_multi2_direct_arg():
+        ray.get([c.small_value_batch_arg.remote(n) for c in clients])
+
+    results += timeit(
+        "n:n actor calls with arg async", actor_multi2_direct_arg, n * len(clients)
+    )
+
+    a = AsyncActor.remote()
+
+    def actor_sync():
+        ray.get(a.small_value.remote())
+
+    results += timeit("1:1 async-actor calls sync", actor_sync)
+
+    a = AsyncActor.remote()
+
+    def async_actor():
+        ray.get([a.small_value.remote() for _ in range(1000)])
+
+    results += timeit("1:1 async-actor calls async", async_actor, 1000)
+
+    a = AsyncActor.remote()
+
+    def async_actor():
+        ray.get([a.small_value_with_arg.remote(i) for i in range(1000)])
+
+    results += timeit("1:1 async-actor calls with args async", async_actor, 1000)
+
+    n = 5000
+    n_cpu = multiprocessing.cpu_count() // 2
+    actors = [AsyncActor.remote() for _ in range(n_cpu)]
+    client = Client.remote(actors)
+
+    def async_actor_async():
+        ray.get(client.small_value_batch.remote(n))
+
+    results += timeit("1:n async-actor calls async", async_actor_async, n * len(actors))
+
+    n = 5000
+    m = 4
+    n_cpu = multiprocessing.cpu_count() // 2
+    a = [AsyncActor.remote() for _ in range(n_cpu)]
+
+    @ray.remote
+    def async_actor_work(actors):
+        ray.get([actors[i % n_cpu].small_value.remote() for i in range(n)])
+
+    def async_actor_multi():
+        ray.get([async_actor_work.remote(a) for _ in range(m)])
+
+    results += timeit("n:n async-actor calls async", async_actor_multi, m * n)
+    ray.shutdown()
+
+    ############################
+    # End of channel perf tests.
+    ############################
+
+    NUM_PGS = 100
+    NUM_BUNDLES = 1
+    ray.init(resources={"custom": 100})
+
+    def placement_group_create_removal(num_pgs):
+        pgs = [
+            ray.util.placement_group(
+                bundles=[{"custom": 0.001} for _ in range(NUM_BUNDLES)]
+            )
+            for _ in range(num_pgs)
+        ]
+        [pg.wait(timeout_seconds=30) for pg in pgs]
+        # Include placement group removal here to clean up.
+        # If we don't clean up placement groups, the whole performance
+        # gets slower as it runs more.
+        # Since timeit function runs multiple times without
+        # the cleaning logic, we should have this method here.
+        for pg in pgs:
+            ray.util.remove_placement_group(pg)
+
+    results += timeit(
+        "placement group create/removal",
+        lambda: placement_group_create_removal(NUM_PGS),
+        NUM_PGS,
+    )
+    ray.shutdown()
+
+    client_microbenchmark_main(results)
+
+    return results
+
+
+if __name__ == "__main__":
+    main()

python/ray/_private/ray_process_reaper.py

@@ -0,0 +1,60 @@
+import atexit
+import os
+import signal
+import sys
+import time
+
+"""
+This is a lightweight "reaper" process used to ensure that ray processes are
+cleaned up properly when the main ray process dies unexpectedly (e.g.,
+segfaults or gets SIGKILLed). Note that processes may not be cleaned up
+properly if this process is SIGTERMed or SIGKILLed.
+
+It detects that its parent has died by reading from stdin, which must be
+inherited from the parent process so that the OS will deliver an EOF if the
+parent dies. When this happens, the reaper process kills the rest of its
+process group (first attempting graceful shutdown with SIGTERM, then escalating
+to SIGKILL).
+"""
+
+SIGTERM_GRACE_PERIOD_SECONDS = 1
+
+
+def reap_process_group(*args):
+    def sigterm_handler(*args):
+        # Give a one-second grace period for other processes to clean up.
+        time.sleep(SIGTERM_GRACE_PERIOD_SECONDS)
+        # SIGKILL the pgroup (including ourselves) as a last-resort.
+        if sys.platform == "win32":
+            atexit.unregister(sigterm_handler)
+            os.kill(0, signal.CTRL_BREAK_EVENT)
+        else:
+            os.killpg(0, signal.SIGKILL)
+
+    # Set a SIGTERM handler to handle SIGTERMing ourselves with the group.
+    if sys.platform == "win32":
+        atexit.register(sigterm_handler)
+    else:
+        signal.signal(signal.SIGTERM, sigterm_handler)
+
+    # Our parent must have died, SIGTERM the group (including ourselves).
+    if sys.platform == "win32":
+        os.kill(0, signal.CTRL_C_EVENT)
+    else:
+        os.killpg(0, signal.SIGTERM)
+
+
+def main():
+    # Read from stdin forever. Because stdin is a file descriptor
+    # inherited from our parent process, we will get an EOF if the parent
+    # dies, which is signaled by an empty return from read().
+    # We intentionally don't set any signal handlers here, so a SIGTERM from
+    # the parent can be used to kill this process gracefully without it killing
+    # the rest of the process group.
+    while len(sys.stdin.read()) != 0:
+        pass
+    reap_process_group()
+
+
+if __name__ == "__main__":
+    main()

python/ray/_private/resource_and_label_spec.py

@@ -0,0 +1,473 @@
+import json
+import logging
+import os
+from typing import Dict, Optional, Tuple
+
+import ray
+import ray._private.ray_constants as ray_constants
+from ray._common.constants import HEAD_NODE_RESOURCE_NAME, NODE_ID_PREFIX
+from ray._common.utils import RESOURCE_CONSTRAINT_PREFIX
+from ray._private import accelerators
+from ray._private.accelerators import AcceleratorManager
+from ray._private.resource_isolation_config import ResourceIsolationConfig
+
+logger = logging.getLogger(__name__)
+
+
+class ResourceAndLabelSpec:
+    """Represents the resource and label configuration passed to a raylet.
+
+    All fields can be None. Before starting services, resolve() should be
+    called to return a ResourceAndLabelSpec with unknown values filled in with
+    merged values based on the local machine and user specifications.
+    """
+
+    def __init__(
+        self,
+        num_cpus: Optional[int] = None,
+        num_gpus: Optional[int] = None,
+        memory: Optional[float] = None,
+        object_store_memory: Optional[float] = None,
+        resources: Optional[Dict[str, float]] = None,
+        labels: Optional[Dict[str, str]] = None,
+    ):
+        """
+        Initialize a ResourceAndLabelSpec
+
+        Args:
+            num_cpus: The CPUs allocated for this raylet.
+            num_gpus: The GPUs allocated for this raylet.
+            memory: The memory allocated for this raylet.
+            object_store_memory: The object store memory allocated for this raylet.
+            resources: The custom resources allocated for this raylet.
+            labels: The labels associated with this node. Labels can be used along
+                with resources for scheduling.
+        """
+        self.num_cpus = num_cpus
+        self.num_gpus = num_gpus
+        self.memory = memory
+        self.object_store_memory = object_store_memory
+        self.resources = resources
+        self.labels = labels
+        self._is_resolved = False
+
+    def resolved(self) -> bool:
+        """Returns if resolve() has been called for this ResourceAndLabelSpec
+        and default values are filled out."""
+        return self._is_resolved
+
+    def _all_fields_set(self) -> bool:
+        """Returns whether all fields in this ResourceAndLabelSpec are not None."""
+        return all(
+            v is not None
+            for v in (
+                self.num_cpus,
+                self.num_gpus,
+                self.memory,
+                self.object_store_memory,
+                self.resources,
+                self.labels,
+            )
+        )
+
+    def to_resource_dict(self):
+        """Returns a dict suitable to pass to raylet initialization.
+
+        This renames num_cpus / num_gpus to "CPU" / "GPU",
+        and check types and values.
+        """
+        assert self.resolved()
+
+        resources = dict(
+            self.resources,
+            CPU=self.num_cpus,
+            GPU=self.num_gpus,
+            memory=int(self.memory),
+            object_store_memory=int(self.object_store_memory),
+        )
+
+        resources = {
+            resource_label: resource_quantity
+            for resource_label, resource_quantity in resources.items()
+            if resource_quantity != 0
+        }
+
+        # Check types.
+        for resource_label, resource_quantity in resources.items():
+            assert isinstance(resource_quantity, int) or isinstance(
+                resource_quantity, float
+            ), (
+                f"{resource_label} ({type(resource_quantity)}): " f"{resource_quantity}"
+            )
+            if (
+                isinstance(resource_quantity, float)
+                and not resource_quantity.is_integer()
+            ):
+                raise ValueError(
+                    "Resource quantities must all be whole numbers. "
+                    "Violated by resource '{}' in {}.".format(resource_label, resources)
+                )
+            if resource_quantity < 0:
+                raise ValueError(
+                    "Resource quantities must be nonnegative. "
+                    "Violated by resource '{}' in {}.".format(resource_label, resources)
+                )
+            if resource_quantity > ray_constants.MAX_RESOURCE_QUANTITY:
+                raise ValueError(
+                    "Resource quantities must be at most {}. "
+                    "Violated by resource '{}' in {}.".format(
+                        ray_constants.MAX_RESOURCE_QUANTITY, resource_label, resources
+                    )
+                )
+
+        return resources
+
+    def resolve(
+        self,
+        is_head: bool,
+        node_ip_address: Optional[str] = None,
+        resource_isolation_config: Optional[ResourceIsolationConfig] = None,
+    ) -> "ResourceAndLabelSpec":
+        """Fills out this ResourceAndLabelSpec instance with merged values from system defaults and user specification.
+
+        Args:
+            is_head: Whether this is the head node.
+            node_ip_address: The IP address of the node that we are on.
+                This is used to automatically create a node id resource.
+            resource_isolation_config: Optional resource isolation config. When
+                enabled and memory is not explicitly set, the system reserved
+                memory for resource isolation is subtracted from available user memory.
+
+        Returns:
+            ResourceAndLabelSpec: This instance with all fields resolved.
+        """
+
+        self._resolve_resources(is_head=is_head, node_ip_address=node_ip_address)
+
+        # Resolve accelerator-specific resources
+        (
+            accelerator_manager,
+            num_accelerators,
+        ) = ResourceAndLabelSpec._get_current_node_accelerator(
+            self.num_gpus, self.resources
+        )
+        self._resolve_accelerator_resources(accelerator_manager, num_accelerators)
+
+        # Default num_gpus value if unset by user and unable to auto-detect.
+        if self.num_gpus is None:
+            self.num_gpus = 0
+
+        # Resolve and merge node labels from all sources (params, env, and default).
+        self._resolve_labels(accelerator_manager)
+
+        # Resolve memory resources
+        self._resolve_memory_resources(resource_isolation_config)
+
+        self._is_resolved = True
+        assert self._all_fields_set()
+        return self
+
+    @staticmethod
+    def _load_env_resources() -> Dict[str, float]:
+        """Load resource overrides from the environment, if present."""
+        env_resources = {}
+        env_string = os.getenv(ray_constants.RESOURCES_ENVIRONMENT_VARIABLE)
+        if env_string:
+            try:
+                env_resources = json.loads(env_string)
+            except Exception:
+                logger.exception(f"Failed to load {env_string}")
+                raise
+            logger.debug(f"Autoscaler overriding resources: {env_resources}.")
+        return env_resources
+
+    @staticmethod
+    def _merge_resources(env_dict: Dict[str, float], params_dict: Dict[str, float]):
+        """Merge environment and Ray param-provided resources, with env values taking precedence.
+        Returns separated special case params (CPU/GPU/memory) and the merged resource dict.
+        """
+        num_cpus = env_dict.pop("CPU", None)
+        num_gpus = env_dict.pop("GPU", None)
+        memory = env_dict.pop("memory", None)
+        object_store_memory = env_dict.pop("object_store_memory", None)
+
+        result = params_dict.copy()
+        result.update(env_dict)
+
+        for key in set(env_dict.keys()).intersection(params_dict or {}):
+            if params_dict[key] != env_dict[key]:
+                logger.warning(
+                    f"Autoscaler is overriding your resource: {key}: "
+                    f"{params_dict[key]} with {env_dict[key]}."
+                )
+
+        return num_cpus, num_gpus, memory, object_store_memory, result
+
+    def _resolve_resources(
+        self, is_head: bool, node_ip_address: Optional[str] = None
+    ) -> None:
+        """Resolve CPU, GPU, and custom resources. Merges resources from environment,
+        Ray params, and defaults in that order of precedence."""
+
+        # Load environment override resources and merge with resources passed
+        # in from Ray Params. Separates special case params if found in env.
+        env_resources = ResourceAndLabelSpec._load_env_resources()
+        (
+            num_cpus,
+            num_gpus,
+            memory,
+            object_store_memory,
+            merged_resources,
+        ) = ResourceAndLabelSpec._merge_resources(env_resources, self.resources or {})
+
+        self.num_cpus = self.num_cpus if num_cpus is None else num_cpus
+        self.num_gpus = self.num_gpus if num_gpus is None else num_gpus
+        self.memory = self.memory if memory is None else memory
+        self.object_store_memory = (
+            self.object_store_memory
+            if object_store_memory is None
+            else object_store_memory
+        )
+        self.resources = merged_resources
+
+        if node_ip_address is None:
+            node_ip_address = ray.util.get_node_ip_address()
+
+        # Automatically create a node id resource on each node. This is
+        # queryable with ray._private.state.node_ids() and
+        # ray._private.state.current_node_id().
+        self.resources[NODE_ID_PREFIX + node_ip_address] = 1.0
+
+        # Automatically create a head node resource.
+        if HEAD_NODE_RESOURCE_NAME in self.resources:
+            raise ValueError(
+                f"{HEAD_NODE_RESOURCE_NAME}"
+                " is a reserved resource name, use another name instead."
+            )
+        if is_head:
+            self.resources[HEAD_NODE_RESOURCE_NAME] = 1.0
+
+        # Auto-detect CPU count if not explicitly set
+        if self.num_cpus is None:
+            self.num_cpus = ray._private.utils.get_num_cpus()
+
+    @staticmethod
+    def _load_env_labels() -> Dict[str, str]:
+        env_override_labels = {}
+        env_override_labels_string = os.getenv(
+            ray_constants.LABELS_ENVIRONMENT_VARIABLE
+        )
+        if env_override_labels_string:
+            try:
+                env_override_labels = json.loads(env_override_labels_string)
+            except Exception:
+                logger.exception(f"Failed to load {env_override_labels_string}")
+                raise
+            logger.info(f"Autoscaler overriding labels: {env_override_labels}.")
+
+        return env_override_labels
+
+    @staticmethod
+    def _get_default_labels(
+        accelerator_manager: Optional[AcceleratorManager],
+    ) -> Dict[str, str]:
+        default_labels = {}
+
+        # Get environment variables populated from K8s Pod Spec
+        node_group = os.environ.get(ray._raylet.NODE_TYPE_NAME_ENV, "")
+        market_type = os.environ.get(ray._raylet.NODE_MARKET_TYPE_ENV, "")
+        availability_region = os.environ.get(ray._raylet.NODE_REGION_ENV, "")
+        availability_zone = os.environ.get(ray._raylet.NODE_ZONE_ENV, "")
+
+        # Map environment variables to default ray node labels
+        if market_type:
+            default_labels[ray._raylet.RAY_NODE_MARKET_TYPE_KEY] = market_type
+        if node_group:
+            default_labels[ray._raylet.RAY_NODE_GROUP_KEY] = node_group
+        if availability_zone:
+            default_labels[ray._raylet.RAY_NODE_ZONE_KEY] = availability_zone
+        if availability_region:
+            default_labels[ray._raylet.RAY_NODE_REGION_KEY] = availability_region
+
+        # Get accelerator type from AcceleratorManager
+        if accelerator_manager:
+            accelerator_type = accelerator_manager.get_current_node_accelerator_type()
+            if accelerator_type:
+                default_labels[
+                    ray._raylet.RAY_NODE_ACCELERATOR_TYPE_KEY
+                ] = accelerator_type
+
+            # Set TPU specific default labels to enable multi-host scheduling.
+            if accelerator_manager.get_resource_name() == "TPU":
+                tpu_labels = accelerator_manager.get_current_node_accelerator_labels()
+                if tpu_labels:
+                    default_labels.update(tpu_labels)
+
+        return default_labels
+
+    def _resolve_labels(
+        self, accelerator_manager: Optional[AcceleratorManager]
+    ) -> None:
+        """Resolve and merge environment override, user-input from params, and Ray default
+        labels in that order of precedence."""
+
+        # Start with a dictionary filled out with Ray default labels
+        merged = ResourceAndLabelSpec._get_default_labels(accelerator_manager)
+
+        # Merge user-specified labels from Ray params
+        for key, val in (self.labels or {}).items():
+            if key in merged and merged[key] != val:
+                logger.warning(
+                    f"User label is overriding Ray default label: {key}: "
+                    f"{key}: {merged[key]} to "
+                    f"{key}: {self.labels[key]}."
+                )
+            merged[key] = val
+
+        # Merge autoscaler override labels from environment
+        env_labels = ResourceAndLabelSpec._load_env_labels()
+        for key, val in (env_labels or {}).items():
+            if key in merged and merged[key] != val:
+                logger.warning(
+                    "Autoscaler is overriding your label:"
+                    f"{key}: {merged[key]} to "
+                    f"{key}: {env_labels[key]}."
+                )
+            merged[key] = val
+
+        self.labels = merged
+
+    def _resolve_accelerator_resources(self, accelerator_manager, num_accelerators):
+        """Detect and update accelerator resources on a node."""
+        if not accelerator_manager:
+            return
+
+        accelerator_resource_name = accelerator_manager.get_resource_name()
+        visible_accelerator_ids = (
+            accelerator_manager.get_current_process_visible_accelerator_ids()
+        )
+
+        # Check that the number of accelerators that the raylet wants doesn't
+        # exceed the amount allowed by visible accelerator ids.
+        if (
+            num_accelerators is not None
+            and visible_accelerator_ids is not None
+            and num_accelerators > len(visible_accelerator_ids)
+        ):
+            raise ValueError(
+                f"Attempting to start raylet with {num_accelerators} "
+                f"{accelerator_resource_name}, "
+                f"but {accelerator_manager.get_visible_accelerator_ids_env_var()} "
+                f"contains {visible_accelerator_ids}."
+            )
+
+        if accelerator_resource_name == "GPU":
+            self.num_gpus = num_accelerators
+        else:
+            self.resources[accelerator_resource_name] = num_accelerators
+
+        accelerator_type = accelerator_manager.get_current_node_accelerator_type()
+        if accelerator_type:
+            self.resources[f"{RESOURCE_CONSTRAINT_PREFIX}{accelerator_type}"] = 1
+        additional_resources = (
+            accelerator_manager.get_current_node_additional_resources()
+        )
+        if additional_resources:
+            self.resources.update(additional_resources)
+
+    def _resolve_memory_resources(
+        self,
+        resource_isolation_config: Optional[ResourceIsolationConfig] = None,
+    ):
+        """
+        Resolves logical and object store memory resources if not
+        explicitly set.
+
+        Args:
+            resource_isolation_config: Optional resource isolation config. When
+                enabled and memory is not explicitly set, the system reserved
+                memory for resource isolation is subtracted from available user memory.
+        """
+        # Choose a default object store size.
+        system_memory = ray._common.utils.get_system_memory()
+        if (
+            resource_isolation_config is not None
+            and resource_isolation_config.is_enabled()
+        ):
+            available_memory_bytes = (
+                system_memory
+                - resource_isolation_config.system_reserved_memory
+                - ray_constants.DEFAULT_USER_PHYSICAL_LOGICAL_MEMORY_LIMIT_BUFFER_BYTES
+            )
+        else:
+            available_memory_bytes = ray._private.utils.estimate_available_memory()
+        if self.object_store_memory is None:
+            self.object_store_memory = ray._private.utils.resolve_object_store_memory(
+                available_memory_bytes
+            )
+
+        memory = self.memory
+        if memory is None:
+            memory = available_memory_bytes - self.object_store_memory
+            if memory < 100e6 and memory < 0.05 * system_memory:
+                raise ValueError(
+                    "After taking into account object store and redis memory "
+                    "usage, the amount of memory on this node available for "
+                    "tasks and actors ({} GB) is less than {}% of total. "
+                    "You can adjust these settings with "
+                    "ray.init(memory=<bytes>, "
+                    "object_store_memory=<bytes>).".format(
+                        round(memory / 1e9, 2), int(100 * (memory / system_memory))
+                    )
+                )
+
+        self.memory = memory
+
+    @staticmethod
+    def _get_current_node_accelerator(
+        num_gpus: Optional[int], resources: Dict[str, float]
+    ) -> Tuple[AcceleratorManager, int]:
+        """
+        Returns the AcceleratorManager and accelerator count for the accelerator
+        associated with this node. This assumes each node has at most one accelerator type.
+        If no accelerators are present, returns None.
+
+        The resolved accelerator count uses num_gpus (for GPUs) or resources if set, and
+        otherwise falls back to the count auto-detected by the AcceleratorManager. The
+        resolved accelerator count is capped by the number of visible accelerators.
+
+        Args:
+            num_gpus: GPU count (if provided by user).
+            resources: Resource dictionary containing custom resource keys.
+
+        Returns:
+            Tuple[Optional[AcceleratorManager], int]: A tuple containing the accelerator
+            manager (or None) the final resolved accelerator count.
+        """
+        for resource_name in accelerators.get_all_accelerator_resource_names():
+            accelerator_manager = accelerators.get_accelerator_manager_for_resource(
+                resource_name
+            )
+            if accelerator_manager is None:
+                continue
+            # Respect configured value for GPUs if set
+            if resource_name == "GPU":
+                num_accelerators = num_gpus
+            else:
+                num_accelerators = resources.get(resource_name)
+            if num_accelerators is None:
+                num_accelerators = (
+                    accelerator_manager.get_current_node_num_accelerators()
+                )
+                visible_accelerator_ids = (
+                    accelerator_manager.get_current_process_visible_accelerator_ids()
+                )
+                if visible_accelerator_ids is not None:
+                    num_accelerators = min(
+                        num_accelerators, len(visible_accelerator_ids)
+                    )
+
+            if num_accelerators > 0:
+                return accelerator_manager, num_accelerators
+
+        return None, 0

python/ray/_private/resource_isolation_config.py

@@ -0,0 +1,282 @@
+import logging
+from typing import Optional
+
+import ray._common.utils
+import ray._private.ray_constants as ray_constants
+import ray._private.utils as utils
+
+logger = logging.getLogger(__name__)
+
+# See https://docs.kernel.org/admin-guide/cgroup-v2.html#weights
+# for information about cpu weights
+_CGROUP_CPU_MAX_WEIGHT: int = 10000
+
+
+class ResourceIsolationConfig:
+    """Configuration for enabling resource isolation by reserving memory and cpu for ray system processes through cgroupv2.
+
+    Validates configuration for resource isolation by enforcing types, correct combinations of values, applying default values,
+    and sanity checking cpu and memory reservations. Also, converts system_reserved_cpu into cpu.weights for cgroupv2.
+
+    Attributes:
+        enable_resource_isolation: True if cgroupv2 based isolation of ray
+            system processes is enabled.
+        cgroup_path: The path for the cgroup the raylet should use to enforce
+            resource isolation.
+        system_reserved_cpu: The amount of cores reserved for ray system
+            processes. Must be >= ray_constants.MINIMUM_SYSTEM_RESERVED_CPU_CORES
+            and < the total number of cores available.
+        system_reserved_memory: The amount of memory in bytes reserved
+            for ray system processes. Must be >= ray_constants.MINIMUM_SYSTEM_RESERVED_MEMORY_BYTES
+            and < the total memory available.
+
+    TODO(54703): Link documentation when it's available.
+    """
+
+    def __init__(
+        self,
+        enable_resource_isolation: bool = False,
+        cgroup_path: Optional[str] = None,
+        system_reserved_cpu: Optional[float] = None,
+        system_reserved_memory: Optional[int] = None,
+    ):
+        """
+        Raises:
+            ValueError: On invalid inputs.
+
+        Args:
+            enable_resource_isolation: True if cgroupv2 based isolation of ray
+                system processes is enabled.
+            cgroup_path: The path for the cgroup the raylet should use to enforce
+                resource isolation.
+            system_reserved_cpu: The amount of cores reserved for ray system
+                processes. Must be >= ray_constants.MINIMUM_SYSTEM_RESERVED_CPU_CORES
+                and < the total number of cores available.
+            system_reserved_memory: The amount of memory in bytes reserved
+                for ray system processes. Must be >= ray_constants.MINIMUM_SYSTEM_RESERVED_MEMORY_BYTES
+                and < the total memory available.
+        """
+        self._resource_isolation_enabled = enable_resource_isolation
+        self.cgroup_path = cgroup_path
+        self.system_reserved_memory = system_reserved_memory
+        self.system_pids = ""
+
+        # cgroupv2 cpu.weight calculated from system_reserved_cpu assumes ray uses all available cores.
+        self.system_reserved_cpu_weight: int = None
+
+        if not enable_resource_isolation:
+            if self.cgroup_path:
+                raise ValueError(
+                    "cgroup_path cannot be set when resource isolation is not enabled. "
+                    "Set enable_resource_isolation to True if you're using ray.init or use the "
+                    "--enable-resource-isolation flag if you're using the ray cli."
+                )
+            if system_reserved_cpu is not None:
+                raise ValueError(
+                    "system_reserved_cpu cannot be set when resource isolation is not enabled. "
+                    "Set enable_resource_isolation to True if you're using ray.init or use the "
+                    "--enable-resource-isolation flag if you're using the ray cli."
+                )
+
+            if system_reserved_memory is not None:
+                raise ValueError(
+                    "system_reserved_memory cannot be set when resource isolation is not enabled. "
+                    "Set enable_resource_isolation to True if you're using ray.init or use the "
+                    "--enable-resource-isolation flag if you're using the ray cli."
+                )
+            return
+
+        self.system_reserved_cpu_weight = self._validate_and_get_system_reserved_cpu(
+            system_reserved_cpu
+        )
+
+        self.system_reserved_memory = self._validate_and_get_system_reserved_memory(
+            system_reserved_memory
+        )
+
+        self.cgroup_path = self._validate_and_get_cgroup_path(cgroup_path)
+
+    def is_enabled(self) -> bool:
+        return self._resource_isolation_enabled
+
+    def add_system_pids(self, system_pids: str):
+        """A comma-separated list of pids to move into the system cgroup."""
+        self.system_pids = system_pids
+
+    @staticmethod
+    def _validate_and_get_cgroup_path(cgroup_path: Optional[str]) -> str:
+        """Returns the ray_constants.DEFAULT_CGROUP_PATH if cgroup_path is not specified.
+
+        Args:
+            cgroup_path: The path for the cgroup the raylet should use to enforce
+                resource isolation.
+
+        Returns:
+            str: The validated cgroup path.
+
+        Raises:
+            ValueError: If cgroup_path is not a string.
+        """
+        if not cgroup_path:
+            cgroup_path = ray_constants.DEFAULT_CGROUP_PATH
+
+        if not isinstance(cgroup_path, str):
+            raise ValueError(
+                f"Invalid value={cgroup_path} for cgroup_path. "
+                "Use a string to represent the path for the cgroup that the raylet should use "
+                "to enable resource isolation."
+            )
+
+        return cgroup_path
+
+    @staticmethod
+    def _validate_and_get_system_reserved_cpu(
+        system_reserved_cpu: Optional[float],
+    ) -> int:
+        """If system_reserved_cpu is specified, validates it, otherwise returns the default value.
+
+        Validation entails checking the type, ensuring that the value is in range, and converts it
+        into cpu.weights for cgroupv2. See https://docs.kernel.org/admin-guide/cgroup-v2.html#weights
+        for more information.
+
+        If system_reserved_cpu is not specified, returns a default value between
+        [DEFAULT_MIN_SYSTEM_RESERVED_CPU_CORES, DEFAULT_MAX_SYSTEM_RESERVED_CPU_CORES].
+
+        # TODO(54703): The errors from this method are user-facing and thus need
+        to be linked the user-facing documentation once it's available.
+
+        Args:
+            system_reserved_cpu: The amount of cores reserved for ray system
+                processes. Must be >= ray_constants.MINIMUM_SYSTEM_RESERVED_CPU_CORES
+                and < the total number of cores available.
+
+        Raises:
+            ValueError: If system_reserved_cpu is specified, but invalid or if the system
+                does not have enough available cpus.
+
+        Returns:
+            The cgroup v2 cpu.weight value derived from the reserved cpu cores.
+        """
+        available_system_cpus = utils.get_num_cpus(truncate=False)
+
+        if available_system_cpus < ray_constants.DEFAULT_MIN_SYSTEM_RESERVED_CPU_CORES:
+            raise ValueError(
+                f"The available number of cpu cores on this system {available_system_cpus} is less than "
+                f"the minimum amount that is required for ray's system processes. "
+                f"Pick a number of cpu cores greater than or equal to {ray_constants.DEFAULT_MIN_SYSTEM_RESERVED_CPU_CORES}"
+            )
+
+        if system_reserved_cpu is None:
+            system_reserved_cpu = float(
+                min(
+                    max(
+                        ray_constants.DEFAULT_MIN_SYSTEM_RESERVED_CPU_CORES,
+                        ray_constants.DEFAULT_SYSTEM_RESERVED_CPU_PROPORTION
+                        * available_system_cpus,
+                    ),
+                    ray_constants.DEFAULT_MAX_SYSTEM_RESERVED_CPU_CORES,
+                )
+            )
+
+        if not (
+            isinstance(system_reserved_cpu, float)
+            or isinstance(system_reserved_cpu, int)
+        ):
+            raise ValueError(
+                f"Invalid value={system_reserved_cpu} for system_reserved_cpu. "
+                "Use a float to represent the number of cores that need to be reserved for "
+                "ray system processes to enable resource isolation."
+            )
+
+        system_reserved_cpu = float(system_reserved_cpu)
+
+        if system_reserved_cpu < ray_constants.DEFAULT_MIN_SYSTEM_RESERVED_CPU_CORES:
+            raise ValueError(
+                f"The requested system_reserved_cpu={system_reserved_cpu} is less than "
+                f"the minimum number of cpus that can be used for resource isolation. "
+                "Pick a number of cpu cores to reserve for ray system processes "
+                f"greater than or equal to {ray_constants.DEFAULT_MIN_SYSTEM_RESERVED_CPU_CORES}"
+            )
+
+        if system_reserved_cpu >= available_system_cpus:
+            raise ValueError(
+                f"The requested system_reserved_cpu={system_reserved_cpu} is greater than or equal to "
+                f"the number of cpus available={available_system_cpus}. "
+                "Pick a smaller number of cpu cores to reserve for ray system processes."
+            )
+
+        # Converting the number of cores the user defined into cpu.weights
+        # This assumes that ray is allowed to use all available CPU
+        # cores and distribute them between system, worker and
+        # user processes
+        return int(
+            (system_reserved_cpu / float(available_system_cpus))
+            * _CGROUP_CPU_MAX_WEIGHT
+        )
+
+    @staticmethod
+    def _validate_and_get_system_reserved_memory(
+        system_reserved_memory: Optional[int],
+    ) -> int:
+        """If system_reserved_memory is not specified, returns the default value. Otherwise,
+        checks the type, makes sure that the value is in range.
+
+        Args:
+            system_reserved_memory: The amount of memory in bytes reserved
+                for ray system processes. Must be >= ray_constants.MINIMUM_SYSTEM_RESERVED_MEMORY_BYTES
+                and < the total memory available.
+
+        Returns:
+            int: The validated system reserved memory in bytes.
+
+        Raises:
+            ValueError: If system_reserved_memory is specified, but invalid.
+        """
+        available_system_memory = ray._common.utils.get_system_memory()
+
+        if (
+            available_system_memory
+            < ray_constants.DEFAULT_MIN_SYSTEM_RESERVED_MEMORY_BYTES
+        ):
+            raise ValueError(
+                f"The available memory on this system {available_system_memory} is less than "
+                f"the minimum amount that is required for ray's system processes. "
+                f"Pick a number of bytes greater than or equal to {ray_constants.DEFAULT_MIN_SYSTEM_RESERVED_MEMORY_BYTES}"
+            )
+
+        if system_reserved_memory is None:
+            system_reserved_memory = int(
+                min(
+                    max(
+                        ray_constants.DEFAULT_MIN_SYSTEM_RESERVED_MEMORY_BYTES,
+                        ray_constants.DEFAULT_SYSTEM_RESERVED_MEMORY_PROPORTION
+                        * available_system_memory,
+                    ),
+                    ray_constants.DEFAULT_MAX_SYSTEM_RESERVED_MEMORY_BYTES,
+                )
+            )
+
+        if not isinstance(system_reserved_memory, int):
+            raise ValueError(
+                f"Invalid value {system_reserved_memory} for system_reserved_memory. "
+                "Use an integer to represent the number bytes that need to be reserved for "
+                "ray system processes to enable resource isolation."
+            )
+
+        if (
+            system_reserved_memory
+            < ray_constants.DEFAULT_MIN_SYSTEM_RESERVED_MEMORY_BYTES
+        ):
+            raise ValueError(
+                f"The requested system_reserved_memory {system_reserved_memory} is less than "
+                f"the minimum number of bytes that can be used for resource isolation. "
+                "Pick a number of bytes to reserve for ray system processes "
+                f"greater than or equal to {ray_constants.DEFAULT_MIN_SYSTEM_RESERVED_MEMORY_BYTES}"
+            )
+
+        if system_reserved_memory > available_system_memory:
+            raise ValueError(
+                f"The total requested system_reserved_memory={system_reserved_memory} "
+                f"is greater than the amount of memory available={available_system_memory}."
+            )
+        return system_reserved_memory

python/ray/_private/runtime_env/BUILD.bazel

@@ -0,0 +1,42 @@
+# TODO(hjiang): All existing pythons are not using bazel as build system, which leads to missing BUILD file and targets.
+# Revisit if we decide to support bazel build in the future.
+
+load("@rules_python//python:defs.bzl", "py_library")
+
+package(default_visibility = ["//visibility:public"])
+
+py_library(
+    name = "validation",
+    srcs = ["validation.py"],
+)
+
+py_library(
+    name = "utils",
+    srcs = ["utils.py"],
+)
+
+py_library(
+    name = "virtualenv_utils",
+    srcs = ["virtualenv_utils.py"],
+    deps = [
+        ":utils",
+    ],
+)
+
+py_library(
+    name = "dependency_utils",
+    srcs = ["dependency_utils.py"],
+    deps = [
+        ":utils",
+    ],
+)
+
+py_library(
+    name = "uv",
+    srcs = ["uv.py"],
+    deps = [
+        ":dependency_utils",
+        ":utils",
+        ":virtualenv_utils",
+    ],
+)

python/ray/_private/runtime_env/__init__.py

@@ -0,0 +1,3 @@
+# List of files to exclude from the Ray directory when using runtime_env for
+# Ray development. These are not necessary in the Ray workers.
+RAY_WORKER_DEV_EXCLUDES = ["raylet", "gcs_server", "cpp/", "tests/", "core/src"]

python/ray/_private/runtime_env/_clonevirtualenv.py

@@ -0,0 +1,334 @@
+#!/usr/bin/env python
+
+from __future__ import with_statement
+
+import logging
+import optparse
+import os
+import os.path
+import re
+import shutil
+import subprocess
+import sys
+import itertools
+
+__version__ = "0.5.7"
+
+
+logger = logging.getLogger()
+
+
+env_bin_dir = "bin"
+if sys.platform == "win32":
+    env_bin_dir = "Scripts"
+    _WIN32 = True
+else:
+    _WIN32 = False
+
+
+class UserError(Exception):
+    pass
+
+
+def _dirmatch(path, matchwith):
+    """Check if path is within matchwith's tree.
+    >>> _dirmatch('/home/foo/bar', '/home/foo/bar')
+    True
+    >>> _dirmatch('/home/foo/bar/', '/home/foo/bar')
+    True
+    >>> _dirmatch('/home/foo/bar/etc', '/home/foo/bar')
+    True
+    >>> _dirmatch('/home/foo/bar2', '/home/foo/bar')
+    False
+    >>> _dirmatch('/home/foo/bar2/etc', '/home/foo/bar')
+    False
+    """
+    matchlen = len(matchwith)
+    if path.startswith(matchwith) and path[matchlen : matchlen + 1] in [os.sep, ""]:
+        return True
+    return False
+
+
+def _virtualenv_sys(venv_path):
+    """obtain version and path info from a virtualenv."""
+    executable = os.path.join(venv_path, env_bin_dir, "python")
+    if _WIN32:
+        env = os.environ.copy()
+    else:
+        env = {}
+    # Must use "executable" as the first argument rather than as the
+    # keyword argument "executable" to get correct value from sys.path
+    p = subprocess.Popen(
+        [
+            executable,
+            "-c",
+            "import sys;"
+            'print ("%d.%d" % (sys.version_info.major, sys.version_info.minor));'
+            'print ("\\n".join(sys.path));',
+        ],
+        env=env,
+        stdout=subprocess.PIPE,
+    )
+    stdout, err = p.communicate()
+    assert not p.returncode and stdout
+    lines = stdout.decode("utf-8").splitlines()
+    return lines[0], list(filter(bool, lines[1:]))
+
+
+def clone_virtualenv(src_dir, dst_dir):
+    if not os.path.exists(src_dir):
+        raise UserError("src dir %r does not exist" % src_dir)
+    if os.path.exists(dst_dir):
+        raise UserError("dest dir %r exists" % dst_dir)
+    # sys_path = _virtualenv_syspath(src_dir)
+    logger.info("cloning virtualenv '%s' => '%s'..." % (src_dir, dst_dir))
+    shutil.copytree(
+        src_dir, dst_dir, symlinks=True, ignore=shutil.ignore_patterns("*.pyc")
+    )
+    version, sys_path = _virtualenv_sys(dst_dir)
+    logger.info("fixing scripts in bin...")
+    fixup_scripts(src_dir, dst_dir, version)
+
+    has_old = lambda s: any(i for i in s if _dirmatch(i, src_dir))  # noqa: E731
+
+    if has_old(sys_path):
+        # only need to fix stuff in sys.path if we have old
+        # paths in the sys.path of new python env. right?
+        logger.info("fixing paths in sys.path...")
+        fixup_syspath_items(sys_path, src_dir, dst_dir)
+    v_sys = _virtualenv_sys(dst_dir)
+    remaining = has_old(v_sys[1])
+    assert not remaining, v_sys
+    fix_symlink_if_necessary(src_dir, dst_dir)
+
+
+def fix_symlink_if_necessary(src_dir, dst_dir):
+    # sometimes the source virtual environment has symlinks that point to itself
+    # one example is $OLD_VIRTUAL_ENV/local/lib points to $OLD_VIRTUAL_ENV/lib
+    # this function makes sure
+    # $NEW_VIRTUAL_ENV/local/lib will point to $NEW_VIRTUAL_ENV/lib
+    # usually this goes unnoticed unless one tries to upgrade a package though pip,
+    # so this bug is hard to find.
+    logger.info("scanning for internal symlinks that point to the original virtual env")
+    for dirpath, dirnames, filenames in os.walk(dst_dir):
+        for a_file in itertools.chain(filenames, dirnames):
+            full_file_path = os.path.join(dirpath, a_file)
+            if os.path.islink(full_file_path):
+                target = os.path.realpath(full_file_path)
+                if target.startswith(src_dir):
+                    new_target = target.replace(src_dir, dst_dir)
+                    logger.debug("fixing symlink in %s" % (full_file_path,))
+                    os.remove(full_file_path)
+                    os.symlink(new_target, full_file_path)
+
+
+def fixup_scripts(old_dir, new_dir, version, rewrite_env_python=False):
+    bin_dir = os.path.join(new_dir, env_bin_dir)
+    root, dirs, files = next(os.walk(bin_dir))
+    pybinre = re.compile(r"pythonw?([0-9]+(\.[0-9]+(\.[0-9]+)?)?)?$")
+    for file_ in files:
+        filename = os.path.join(root, file_)
+        if file_ in ["python", "python%s" % version, "activate_this.py"]:
+            continue
+        elif file_.startswith("python") and pybinre.match(file_):
+            # ignore other possible python binaries
+            continue
+        elif file_.endswith(".pyc"):
+            # ignore compiled files
+            continue
+        elif file_ == "activate" or file_.startswith("activate."):
+            fixup_activate(os.path.join(root, file_), old_dir, new_dir)
+        elif os.path.islink(filename):
+            fixup_link(filename, old_dir, new_dir)
+        elif os.path.isfile(filename):
+            fixup_script_(
+                root,
+                file_,
+                old_dir,
+                new_dir,
+                version,
+                rewrite_env_python=rewrite_env_python,
+            )
+
+
+def fixup_script_(root, file_, old_dir, new_dir, version, rewrite_env_python=False):
+    old_shebang = "#!%s/bin/python" % os.path.normcase(os.path.abspath(old_dir))
+    new_shebang = "#!%s/bin/python" % os.path.normcase(os.path.abspath(new_dir))
+    env_shebang = "#!/usr/bin/env python"
+
+    filename = os.path.join(root, file_)
+    with open(filename, "rb") as f:
+        if f.read(2) != b"#!":
+            # no shebang
+            return
+        f.seek(0)
+        lines = f.readlines()
+
+    if not lines:
+        # warn: empty script
+        return
+
+    def rewrite_shebang(version=None):
+        logger.debug("fixing %s" % filename)
+        shebang = new_shebang
+        if version:
+            shebang = shebang + version
+        shebang = (shebang + "\n").encode("utf-8")
+        with open(filename, "wb") as f:
+            f.write(shebang)
+            f.writelines(lines[1:])
+
+    try:
+        bang = lines[0].decode("utf-8").strip()
+    except UnicodeDecodeError:
+        # binary file
+        return
+
+    # This takes care of the scheme in which shebang is of type
+    # '#!/venv/bin/python3' while the version of system python
+    # is of type 3.x e.g. 3.5.
+    short_version = bang[len(old_shebang) :]
+
+    if not bang.startswith("#!"):
+        return
+    elif bang == old_shebang:
+        rewrite_shebang()
+    elif bang.startswith(old_shebang) and bang[len(old_shebang) :] == version:
+        rewrite_shebang(version)
+    elif (
+        bang.startswith(old_shebang)
+        and short_version
+        and bang[len(old_shebang) :] == short_version
+    ):
+        rewrite_shebang(short_version)
+    elif rewrite_env_python and bang.startswith(env_shebang):
+        if bang == env_shebang:
+            rewrite_shebang()
+        elif bang[len(env_shebang) :] == version:
+            rewrite_shebang(version)
+    else:
+        # can't do anything
+        return
+
+
+def fixup_activate(filename, old_dir, new_dir):
+    logger.debug("fixing %s" % filename)
+    with open(filename, "rb") as f:
+        data = f.read().decode("utf-8")
+
+    data = data.replace(old_dir, new_dir)
+    with open(filename, "wb") as f:
+        f.write(data.encode("utf-8"))
+
+
+def fixup_link(filename, old_dir, new_dir, target=None):
+    logger.debug("fixing %s" % filename)
+    if target is None:
+        target = os.readlink(filename)
+
+    origdir = os.path.dirname(os.path.abspath(filename)).replace(new_dir, old_dir)
+    if not os.path.isabs(target):
+        target = os.path.abspath(os.path.join(origdir, target))
+        rellink = True
+    else:
+        rellink = False
+
+    if _dirmatch(target, old_dir):
+        if rellink:
+            # keep relative links, but don't keep original in case it
+            # traversed up out of, then back into the venv.
+            # so, recreate a relative link from absolute.
+            target = target[len(origdir) :].lstrip(os.sep)
+        else:
+            target = target.replace(old_dir, new_dir, 1)
+
+    # else: links outside the venv, replaced with absolute path to target.
+    _replace_symlink(filename, target)
+
+
+def _replace_symlink(filename, newtarget):
+    tmpfn = "%s.new" % filename
+    os.symlink(newtarget, tmpfn)
+    os.rename(tmpfn, filename)
+
+
+def fixup_syspath_items(syspath, old_dir, new_dir):
+    for path in syspath:
+        if not os.path.isdir(path):
+            continue
+        path = os.path.normcase(os.path.abspath(path))
+        if _dirmatch(path, old_dir):
+            path = path.replace(old_dir, new_dir, 1)
+            if not os.path.exists(path):
+                continue
+        elif not _dirmatch(path, new_dir):
+            continue
+        root, dirs, files = next(os.walk(path))
+        for file_ in files:
+            filename = os.path.join(root, file_)
+            if filename.endswith(".pth"):
+                fixup_pth_file(filename, old_dir, new_dir)
+            elif filename.endswith(".egg-link"):
+                fixup_egglink_file(filename, old_dir, new_dir)
+
+
+def fixup_pth_file(filename, old_dir, new_dir):
+    logger.debug("fixup_pth_file %s" % filename)
+
+    with open(filename, "r") as f:
+        lines = f.readlines()
+
+    has_change = False
+
+    for num, line in enumerate(lines):
+        line = (line.decode("utf-8") if hasattr(line, "decode") else line).strip()
+
+        if not line or line.startswith("#") or line.startswith("import "):
+            continue
+        elif _dirmatch(line, old_dir):
+            lines[num] = line.replace(old_dir, new_dir, 1)
+            has_change = True
+
+    if has_change:
+        with open(filename, "w") as f:
+            payload = os.linesep.join([line.strip() for line in lines]) + os.linesep
+            f.write(payload)
+
+
+def fixup_egglink_file(filename, old_dir, new_dir):
+    logger.debug("fixing %s" % filename)
+    with open(filename, "rb") as f:
+        link = f.read().decode("utf-8").strip()
+    if _dirmatch(link, old_dir):
+        link = link.replace(old_dir, new_dir, 1)
+        with open(filename, "wb") as f:
+            link = (link + "\n").encode("utf-8")
+            f.write(link)
+
+
+def main():
+    parser = optparse.OptionParser(
+        "usage: %prog [options] /path/to/existing/venv /path/to/cloned/venv"
+    )
+    parser.add_option(
+        "-v", action="count", dest="verbose", default=False, help="verbosity"
+    )
+    options, args = parser.parse_args()
+    try:
+        old_dir, new_dir = args
+    except ValueError:
+        print("virtualenv-clone %s" % (__version__,))
+        parser.error("not enough arguments given.")
+    old_dir = os.path.realpath(old_dir)
+    new_dir = os.path.realpath(new_dir)
+    loglevel = (logging.WARNING, logging.INFO, logging.DEBUG)[min(2, options.verbose)]
+    logging.basicConfig(level=loglevel, format="%(message)s")
+    try:
+        clone_virtualenv(old_dir, new_dir)
+    except UserError:
+        e = sys.exc_info()[1]
+        parser.error(str(e))
+
+
+if __name__ == "__main__":
+    main()

python/ray/_private/runtime_env/agent/main.py

@@ -0,0 +1,266 @@
+import argparse
+import logging
+import os
+import socket
+import sys
+
+import ray
+import ray._private.ray_constants as ray_constants
+from ray._common.utils import (
+    get_or_create_event_loop,
+)
+from ray._private import logging_utils
+from ray._private.authentication.http_token_authentication import (
+    get_token_auth_middleware,
+)
+from ray._private.process_watcher import create_check_raylet_task
+from ray._raylet import RUNTIME_ENV_AGENT_PORT_NAME, GcsClient, persist_port
+from ray.core.generated import (
+    runtime_env_agent_pb2,
+)
+
+
+def import_libs():
+    my_dir = os.path.abspath(os.path.dirname(__file__))
+    sys.path.insert(0, os.path.join(my_dir, "thirdparty_files"))  # for aiohttp
+    sys.path.insert(0, my_dir)  # for runtime_env_agent and runtime_env_consts
+
+
+import_libs()
+
+import aiohttp  # noqa: E402
+import runtime_env_consts  # noqa: E402
+from aiohttp import web  # noqa: E402
+from runtime_env_agent import RuntimeEnvAgent  # noqa: E402
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Runtime env agent.")
+    parser.add_argument(
+        "--node-id",
+        required=True,
+        type=str,
+        help="the unique ID of this node.",
+    )
+    parser.add_argument(
+        "--node-ip-address",
+        required=True,
+        type=str,
+        help="the IP address of this node.",
+    )
+    parser.add_argument(
+        "--runtime-env-agent-port",
+        required=True,
+        type=int,
+        default=None,
+        help="The port on which the runtime env agent will receive HTTP requests.",
+    )
+    parser.add_argument(
+        "--session-dir",
+        required=True,
+        type=str,
+        default=None,
+        help="The path of this ray session directory.",
+    )
+
+    parser.add_argument(
+        "--gcs-address", required=True, type=str, help="The address (ip:port) of GCS."
+    )
+    parser.add_argument(
+        "--cluster-id-hex", required=True, type=str, help="The cluster id in hex."
+    )
+    parser.add_argument(
+        "--runtime-env-dir",
+        required=True,
+        type=str,
+        default=None,
+        help="Specify the path of the resource directory used by runtime_env.",
+    )
+
+    parser.add_argument(
+        "--logging-level",
+        required=False,
+        type=lambda s: logging.getLevelName(s.upper()),
+        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=runtime_env_consts.RUNTIME_ENV_AGENT_LOG_FILENAME,
+        help="Specify the name of log file, "
+        'log to stdout if set empty, default is "{}".'.format(
+            runtime_env_consts.RUNTIME_ENV_AGENT_LOG_FILENAME
+        ),
+    )
+    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(
+        "--log-dir",
+        required=True,
+        type=str,
+        default=None,
+        help="Specify the path of log directory.",
+    )
+    parser.add_argument(
+        "--temp-dir",
+        required=True,
+        type=str,
+        default=None,
+        help="Specify the path of the temporary directory use by Ray process.",
+    )
+    parser.add_argument(
+        "--stdout-filepath",
+        required=False,
+        type=str,
+        default="",
+        help="The filepath to dump runtime env agent stdout.",
+    )
+    parser.add_argument(
+        "--stderr-filepath",
+        required=False,
+        type=str,
+        default="",
+        help="The filepath to dump runtime env agent 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.log_dir,
+        filename=args.logging_filename,
+        max_bytes=logging_rotation_bytes,
+        backup_count=logging_rotation_backup_count,
+    )
+
+    # 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, cluster_id=args.cluster_id_hex)
+    agent = RuntimeEnvAgent(
+        runtime_env_dir=args.runtime_env_dir,
+        logging_params=logging_params,
+        gcs_client=gcs_client,
+        temp_dir=args.temp_dir,
+        address=args.node_ip_address,
+        runtime_env_agent_port=args.runtime_env_agent_port,
+    )
+
+    ray._raylet.setproctitle(ray_constants.AGENT_PROCESS_TYPE_RUNTIME_ENV_AGENT)
+
+    # POST /get_or_create_runtime_env
+    # body is serialzied protobuf GetOrCreateRuntimeEnvRequest
+    # reply is serialzied protobuf GetOrCreateRuntimeEnvReply
+    async def get_or_create_runtime_env(request: web.Request) -> web.Response:
+        data = await request.read()
+        request = runtime_env_agent_pb2.GetOrCreateRuntimeEnvRequest()
+        request.ParseFromString(data)
+        reply = await agent.GetOrCreateRuntimeEnv(request)
+        return web.Response(
+            body=reply.SerializeToString(), content_type="application/octet-stream"
+        )
+
+    # POST /delete_runtime_env_if_possible
+    # body is serialzied protobuf DeleteRuntimeEnvIfPossibleRequest
+    # reply is serialzied protobuf DeleteRuntimeEnvIfPossibleReply
+    async def delete_runtime_env_if_possible(request: web.Request) -> web.Response:
+        data = await request.read()
+        request = runtime_env_agent_pb2.DeleteRuntimeEnvIfPossibleRequest()
+        request.ParseFromString(data)
+        reply = await agent.DeleteRuntimeEnvIfPossible(request)
+        return web.Response(
+            body=reply.SerializeToString(), content_type="application/octet-stream"
+        )
+
+    # POST /get_runtime_envs_info
+    # body is serialzied protobuf GetRuntimeEnvsInfoRequest
+    # reply is serialzied protobuf GetRuntimeEnvsInfoReply
+    async def get_runtime_envs_info(request: web.Request) -> web.Response:
+        data = await request.read()
+        request = runtime_env_agent_pb2.GetRuntimeEnvsInfoRequest()
+        request.ParseFromString(data)
+        reply = await agent.GetRuntimeEnvsInfo(request)
+        return web.Response(
+            body=reply.SerializeToString(), content_type="application/octet-stream"
+        )
+
+    app = web.Application(middlewares=[get_token_auth_middleware(aiohttp)])
+
+    app.router.add_post("/get_or_create_runtime_env", get_or_create_runtime_env)
+    app.router.add_post(
+        "/delete_runtime_env_if_possible", delete_runtime_env_if_possible
+    )
+    app.router.add_post("/get_runtime_envs_info", get_runtime_envs_info)
+
+    loop = get_or_create_event_loop()
+    check_raylet_task = None
+    if sys.platform not in ["win32", "cygwin"]:
+
+        def parent_dead_callback(msg):
+            agent._logger.info(
+                "Raylet is dead! Exiting Runtime Env Agent. "
+                f"addr: {args.node_ip_address}, "
+                f"port: {args.runtime_env_agent_port}\n"
+                f"{msg}"
+            )
+
+        # No need to await this task.
+        check_raylet_task = create_check_raylet_task(
+            args.log_dir, gcs_client, parent_dead_callback, loop
+        )
+
+    port = args.runtime_env_agent_port or 0
+    infos = socket.getaddrinfo(args.node_ip_address, port, type=socket.SOCK_STREAM)
+    family, socktype, proto, _, sockaddr = infos[0]
+    sock = socket.socket(family, socktype, proto)
+    sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+    sock.bind(sockaddr)
+
+    bound_port = sock.getsockname()[1]
+    persist_port(
+        args.session_dir,
+        args.node_id,
+        RUNTIME_ENV_AGENT_PORT_NAME,
+        bound_port,
+    )
+
+    try:
+        web.run_app(app, sock=sock, loop=loop)
+    except SystemExit as e:
+        agent._logger.info(f"SystemExit! {e}")
+        # We have to poke the task exception, or there's an error message
+        # "task exception was never retrieved".
+        if check_raylet_task is not None:
+            check_raylet_task.exception()
+        sys.exit(e.code)

python/ray/_private/runtime_env/agent/runtime_env_agent.py

@@ -0,0 +1,620 @@
+import asyncio
+import logging
+import os
+import time
+import traceback
+from collections import defaultdict
+from dataclasses import dataclass
+from typing import Callable, Dict, List, Set, Tuple
+
+import ray
+import ray._private.runtime_env.agent.runtime_env_consts as runtime_env_consts
+from ray._common.utils import get_or_create_event_loop
+from ray._private.ray_constants import (
+    DEFAULT_RUNTIME_ENV_TIMEOUT_SECONDS,
+)
+from ray._private.ray_logging import setup_component_logger
+from ray._private.runtime_env.conda import CondaPlugin
+from ray._private.runtime_env.context import RuntimeEnvContext
+from ray._private.runtime_env.default_impl import get_image_uri_plugin_cls
+from ray._private.runtime_env.image_uri import ContainerPlugin
+from ray._private.runtime_env.java_jars import JavaJarsPlugin
+from ray._private.runtime_env.nsight import NsightPlugin
+from ray._private.runtime_env.pip import PipPlugin
+from ray._private.runtime_env.plugin import (
+    RuntimeEnvPlugin,
+    RuntimeEnvPluginManager,
+    create_for_plugin_if_needed,
+)
+from ray._private.runtime_env.py_executable import PyExecutablePlugin
+from ray._private.runtime_env.py_modules import PyModulesPlugin
+from ray._private.runtime_env.rocprof_sys import RocProfSysPlugin
+from ray._private.runtime_env.uv import UvPlugin
+from ray._private.runtime_env.working_dir import WorkingDirPlugin
+from ray._raylet import GcsClient
+from ray.core.generated import runtime_env_agent_pb2
+from ray.core.generated.runtime_env_common_pb2 import (
+    RuntimeEnvState as ProtoRuntimeEnvState,
+)
+from ray.runtime_env import RuntimeEnv, RuntimeEnvConfig
+
+default_logger = logging.getLogger(__name__)
+
+# TODO(edoakes): this is used for unit tests. We should replace it with a
+# better pluggability mechanism once available.
+SLEEP_FOR_TESTING_S = os.environ.get("RAY_RUNTIME_ENV_SLEEP_FOR_TESTING_S")
+
+
+@dataclass
+class CreatedEnvResult:
+    # Whether or not the env was installed correctly.
+    success: bool
+    # If success is True, will be a serialized RuntimeEnvContext
+    # If success is False, will be an error message.
+    result: str
+    # The time to create a runtime env in ms.
+    creation_time_ms: int
+
+
+# e.g., "working_dir"
+UriType = str
+
+
+class ReferenceTable:
+    """
+    The URI reference table which is used for GC.
+    When the reference count is decreased to zero,
+    the URI should be removed from this table and
+    added to cache if needed.
+    """
+
+    def __init__(
+        self,
+        uris_parser: Callable[[RuntimeEnv], Tuple[str, UriType]],
+        unused_uris_callback: Callable[[List[Tuple[str, UriType]]], None],
+        unused_runtime_env_callback: Callable[[str], None],
+    ):
+        # Runtime Environment reference table. The key is serialized runtime env and
+        # the value is reference count.
+        self._runtime_env_reference: Dict[str, int] = defaultdict(int)
+        # URI reference table. The key is URI parsed from runtime env and the value
+        # is reference count.
+        self._uri_reference: Dict[str, int] = defaultdict(int)
+        self._uris_parser = uris_parser
+        self._unused_uris_callback = unused_uris_callback
+        self._unused_runtime_env_callback = unused_runtime_env_callback
+        # send the `DeleteRuntimeEnvIfPossible` RPC when the client exits. The URI won't
+        # be leaked now because the reference count will be reset to zero when the job
+        # finished.
+        self._reference_exclude_sources: Set[str] = {
+            "client_server",
+        }
+
+    def _increase_reference_for_uris(self, uris):
+        default_logger.debug(f"Increase reference for uris {uris}.")
+        for uri, _ in uris:
+            self._uri_reference[uri] += 1
+
+    def _decrease_reference_for_uris(self, uris):
+        default_logger.debug(f"Decrease reference for uris {uris}.")
+        unused_uris = list()
+        for uri, uri_type in uris:
+            if self._uri_reference[uri] > 0:
+                self._uri_reference[uri] -= 1
+                if self._uri_reference[uri] == 0:
+                    unused_uris.append((uri, uri_type))
+                    del self._uri_reference[uri]
+            else:
+                default_logger.warning(f"URI {uri} does not exist.")
+        if unused_uris:
+            default_logger.info(f"Unused uris {unused_uris}.")
+            self._unused_uris_callback(unused_uris)
+        return unused_uris
+
+    def _increase_reference_for_runtime_env(self, serialized_env: str):
+        default_logger.debug(f"Increase reference for runtime env {serialized_env}.")
+        self._runtime_env_reference[serialized_env] += 1
+
+    def _decrease_reference_for_runtime_env(self, serialized_env: str):
+        """Decrease reference count for the given [serialized_env]. Throw exception if we cannot decrement reference."""
+        default_logger.debug(f"Decrease reference for runtime env {serialized_env}.")
+        unused = False
+        if self._runtime_env_reference[serialized_env] > 0:
+            self._runtime_env_reference[serialized_env] -= 1
+            if self._runtime_env_reference[serialized_env] == 0:
+                unused = True
+                del self._runtime_env_reference[serialized_env]
+        else:
+            default_logger.warning(f"Runtime env {serialized_env} does not exist.")
+            raise ValueError(
+                f"{serialized_env} cannot decrement reference since the reference count is 0"
+            )
+        if unused:
+            default_logger.info(f"Unused runtime env {serialized_env}.")
+            self._unused_runtime_env_callback(serialized_env)
+
+    def increase_reference(
+        self, runtime_env: RuntimeEnv, serialized_env: str, source_process: str
+    ) -> None:
+        if source_process in self._reference_exclude_sources:
+            return
+        self._increase_reference_for_runtime_env(serialized_env)
+        uris = self._uris_parser(runtime_env)
+        self._increase_reference_for_uris(uris)
+
+    def decrease_reference(
+        self, runtime_env: RuntimeEnv, serialized_env: str, source_process: str
+    ) -> None:
+        """Decrease reference count for runtime env and uri. Throw exception if decrement reference count fails."""
+        if source_process in self._reference_exclude_sources:
+            return
+        self._decrease_reference_for_runtime_env(serialized_env)
+        uris = self._uris_parser(runtime_env)
+        self._decrease_reference_for_uris(uris)
+
+    @property
+    def runtime_env_refs(self) -> Dict[str, int]:
+        """Return the runtime_env -> ref count mapping.
+
+        Returns:
+            The mapping of serialized runtime env -> ref count.
+        """
+        return self._runtime_env_reference
+
+
+class RuntimeEnvAgent:
+    """An RPC server to create and delete runtime envs.
+
+    Attributes:
+        dashboard_agent: The DashboardAgent object contains global config.
+    """
+
+    def __init__(
+        self,
+        runtime_env_dir: str,
+        logging_params: dict,
+        gcs_client: GcsClient,
+        temp_dir: str,
+        address: str,
+        runtime_env_agent_port: int,
+    ):
+        """Initialize the runtime env agent.
+
+        Args:
+            runtime_env_dir: Directory used to store runtime env resources.
+            logging_params: Keyword arguments forwarded to
+                :func:`setup_component_logger` to configure the agent logger.
+            gcs_client: GCS client used to fetch package data.
+            temp_dir: Temporary directory used by plugins (e.g. container plugin).
+            address: IP address that the agent is listening on, used for logging.
+            runtime_env_agent_port: Port that the agent is listening on, used for
+                logging.
+        """
+        super().__init__()
+
+        self._logger = default_logger
+        self._logging_params = logging_params
+        self._logger = setup_component_logger(
+            logger_name=default_logger.name, **self._logging_params
+        )
+        # Don't propagate logs to the root logger, because these logs
+        # might contain sensitive information. Instead, these logs should
+        # be confined to the runtime env agent log file `self.LOG_FILENAME`.
+        self._logger.propagate = False
+
+        self._logger.info("Starting runtime env agent at pid %s", os.getpid())
+        self._logger.info(f"Parent raylet pid is {os.environ.get('RAY_RAYLET_PID')}")
+
+        self._runtime_env_dir = runtime_env_dir
+        self._per_job_logger_cache = dict()
+        # Cache the results of creating envs to avoid repeatedly calling into
+        # conda and other slow calls.
+        self._env_cache: Dict[str, CreatedEnvResult] = dict()
+        # Maps a serialized runtime env to a lock that is used
+        # to prevent multiple concurrent installs of the same env.
+        self._env_locks: Dict[str, asyncio.Lock] = dict()
+        self._gcs_client = gcs_client
+
+        self._pip_plugin = PipPlugin(self._runtime_env_dir)
+        self._uv_plugin = UvPlugin(self._runtime_env_dir)
+        self._conda_plugin = CondaPlugin(self._runtime_env_dir)
+        self._py_modules_plugin = PyModulesPlugin(
+            self._runtime_env_dir, self._gcs_client
+        )
+        self._py_executable_plugin = PyExecutablePlugin()
+        self._java_jars_plugin = JavaJarsPlugin(self._runtime_env_dir, self._gcs_client)
+        self._working_dir_plugin = WorkingDirPlugin(
+            self._runtime_env_dir, self._gcs_client
+        )
+        self._container_plugin = ContainerPlugin(temp_dir)
+        # TODO(jonathan-anyscale): change the plugin to ProfilerPlugin
+        # and unify with nsight and other profilers.
+        self._nsight_plugin = NsightPlugin(self._runtime_env_dir)
+        self._rocprof_sys_plugin = RocProfSysPlugin(self._runtime_env_dir)
+        self._image_uri_plugin = get_image_uri_plugin_cls()(temp_dir)
+
+        # TODO(architkulkarni): "base plugins" and third-party plugins should all go
+        # through the same code path.  We should never need to refer to
+        # self._xxx_plugin, we should just iterate through self._plugins.
+        self._base_plugins: List[RuntimeEnvPlugin] = [
+            self._working_dir_plugin,
+            self._uv_plugin,
+            self._pip_plugin,
+            self._conda_plugin,
+            self._py_modules_plugin,
+            self._py_executable_plugin,
+            self._java_jars_plugin,
+            self._container_plugin,
+            self._nsight_plugin,
+            self._rocprof_sys_plugin,
+            self._image_uri_plugin,
+        ]
+        self._plugin_manager = RuntimeEnvPluginManager()
+        for plugin in self._base_plugins:
+            self._plugin_manager.add_plugin(plugin)
+
+        self._reference_table = ReferenceTable(
+            self.uris_parser,
+            self.unused_uris_processor,
+            self.unused_runtime_env_processor,
+        )
+
+        self._logger.info(
+            "Listening to address %s, port %d", address, runtime_env_agent_port
+        )
+
+        try:
+            self._node_ip = ray.util.get_node_ip_address()
+            self._node_prefix = f"[Node {self._node_ip}] "
+        except Exception as e:
+            self._logger.warning(f"Failed to get node IP address, using fallback: {e}")
+            self._node_prefix = "[Node unknown] "
+
+    def uris_parser(self, runtime_env: RuntimeEnv):
+        result = list()
+        for name, plugin_setup_context in self._plugin_manager.plugins.items():
+            plugin = plugin_setup_context.class_instance
+            uris = plugin.get_uris(runtime_env)
+            for uri in uris:
+                result.append((uri, UriType(name)))
+        return result
+
+    def unused_uris_processor(self, unused_uris: List[Tuple[str, UriType]]) -> None:
+        for uri, uri_type in unused_uris:
+            self._plugin_manager.plugins[str(uri_type)].uri_cache.mark_unused(uri)
+
+    def unused_runtime_env_processor(self, unused_runtime_env: str) -> None:
+        def delete_runtime_env():
+            del self._env_cache[unused_runtime_env]
+            self._logger.info(
+                "Runtime env %s removed from env-level cache.", unused_runtime_env
+            )
+
+        if unused_runtime_env in self._env_cache:
+            if not self._env_cache[unused_runtime_env].success:
+                loop = get_or_create_event_loop()
+                # Cache the bad runtime env result by ttl seconds.
+                loop.call_later(
+                    runtime_env_consts.BAD_RUNTIME_ENV_CACHE_TTL_SECONDS,
+                    delete_runtime_env,
+                )
+            else:
+                delete_runtime_env()
+
+    def get_or_create_logger(self, job_id: bytes, log_files: List[str]):
+        job_id = job_id.decode()
+        if job_id not in self._per_job_logger_cache:
+            params = self._logging_params.copy()
+            params["filename"] = [f"runtime_env_setup-{job_id}.log", *log_files]
+            params["logger_name"] = f"runtime_env_{job_id}"
+            params["propagate"] = False
+            per_job_logger = setup_component_logger(**params)
+            self._per_job_logger_cache[job_id] = per_job_logger
+        return self._per_job_logger_cache[job_id]
+
+    async def GetOrCreateRuntimeEnv(self, request):
+        self._logger.debug(
+            f"Got request from {request.source_process} to increase "
+            "reference for runtime env: "
+            f"{request.serialized_runtime_env}."
+        )
+
+        async def _setup_runtime_env(
+            runtime_env: RuntimeEnv,
+            runtime_env_config: RuntimeEnvConfig,
+        ):
+            log_files = runtime_env_config.get("log_files", [])
+            # Use a separate logger for each job.
+            per_job_logger = self.get_or_create_logger(request.job_id, log_files)
+            context = RuntimeEnvContext(env_vars=runtime_env.env_vars())
+
+            # Warn about unrecognized fields in the runtime env.
+            for name, _ in runtime_env.plugins():
+                if name not in self._plugin_manager.plugins:
+                    per_job_logger.warning(
+                        f"runtime_env field {name} is not recognized by "
+                        "Ray and will be ignored.  In the future, unrecognized "
+                        "fields in the runtime_env will raise an exception."
+                    )
+
+            # Creates each runtime env URI by their priority. `working_dir` is special
+            # because it needs to be created before other plugins. All other plugins are
+            # created in the priority order (smaller priority value -> earlier to
+            # create), with a special environment variable being set to the working dir.
+            # ${RAY_RUNTIME_ENV_CREATE_WORKING_DIR}
+
+            # First create working dir...
+            working_dir_ctx = self._plugin_manager.plugins[WorkingDirPlugin.name]
+            await create_for_plugin_if_needed(
+                runtime_env,
+                working_dir_ctx.class_instance,
+                working_dir_ctx.uri_cache,
+                context,
+                per_job_logger,
+            )
+
+            # Then within the working dir, create the other plugins.
+            working_dir_uri_or_none = runtime_env.working_dir_uri()
+            with self._working_dir_plugin.with_working_dir_env(working_dir_uri_or_none):
+                """Run setup for each plugin unless it has already been cached."""
+                for (
+                    plugin_setup_context
+                ) in self._plugin_manager.sorted_plugin_setup_contexts():
+                    plugin = plugin_setup_context.class_instance
+                    if plugin.name != WorkingDirPlugin.name:
+                        uri_cache = plugin_setup_context.uri_cache
+                        await create_for_plugin_if_needed(
+                            runtime_env, plugin, uri_cache, context, per_job_logger
+                        )
+            return context
+
+        async def _create_runtime_env_with_retry(
+            runtime_env: RuntimeEnv,
+            setup_timeout_seconds: int,
+            runtime_env_config: RuntimeEnvConfig,
+        ) -> Tuple[bool, str, str]:
+            """Create runtime env with retry times. This function won't raise exceptions.
+
+            Args:
+                runtime_env: The instance of RuntimeEnv class.
+                setup_timeout_seconds: The timeout of runtime environment creation for
+                    each attempt.
+                runtime_env_config: The configuration for the runtime environment.
+
+            Returns:
+                Tuple[bool, str, str]: A tuple containing:
+                    - result (bool): Whether the creation was successful
+                    - runtime_env_context (str): The serialized context if successful, None otherwise
+                    - error_message (str): Error message if failed, None otherwise
+            """
+            self._logger.info(
+                f"Creating runtime env: {serialized_env} with timeout "
+                f"{setup_timeout_seconds} seconds."
+            )
+            num_retries = runtime_env_consts.RUNTIME_ENV_RETRY_TIMES
+            error_message = None
+            serialized_context = None
+            for i in range(num_retries):
+                # Only sleep when retrying.
+                if i != 0:
+                    await asyncio.sleep(
+                        runtime_env_consts.RUNTIME_ENV_RETRY_INTERVAL_MS / 1000
+                    )
+
+                try:
+                    runtime_env_setup_task = _setup_runtime_env(
+                        runtime_env, runtime_env_config
+                    )
+                    runtime_env_context = await asyncio.wait_for(
+                        runtime_env_setup_task, timeout=setup_timeout_seconds
+                    )
+                    serialized_context = runtime_env_context.serialize()
+                    error_message = None
+                    break
+                except Exception as e:
+                    err_msg = f"Failed to create runtime env {serialized_env}."
+                    self._logger.exception(err_msg)
+                    error_message = "".join(
+                        traceback.format_exception(type(e), e, e.__traceback__)
+                    )
+                    if isinstance(e, asyncio.TimeoutError):
+                        hint = (
+                            f"Failed to install runtime_env within the "
+                            f"timeout of {setup_timeout_seconds} seconds. Consider "
+                            "increasing the timeout in the runtime_env config. "
+                            "For example: \n"
+                            '    runtime_env={"config": {"setup_timeout_seconds":'
+                            " 1800}, ...}\n"
+                            "If not provided, the default timeout is "
+                            f"{DEFAULT_RUNTIME_ENV_TIMEOUT_SECONDS} seconds. "
+                        )
+                        error_message = hint + error_message
+
+            if error_message:
+                self._logger.error(
+                    "runtime_env creation failed %d times, giving up.",
+                    num_retries,
+                )
+                return False, None, error_message
+            else:
+                self._logger.info(
+                    "Successfully created runtime env: %s, context: %s",
+                    serialized_env,
+                    serialized_context,
+                )
+                return True, serialized_context, None
+
+        try:
+            serialized_env = request.serialized_runtime_env
+            runtime_env = RuntimeEnv.deserialize(serialized_env)
+        except Exception as e:
+            self._logger.exception(
+                "[Increase] Failed to parse runtime env: " f"{serialized_env}"
+            )
+
+            error_message = "".join(
+                traceback.format_exception(type(e), e, e.__traceback__)
+            )
+
+            return runtime_env_agent_pb2.GetOrCreateRuntimeEnvReply(
+                status=runtime_env_agent_pb2.AGENT_RPC_STATUS_FAILED,
+                error_message=f"{self._node_prefix}{error_message}",
+            )
+
+        # Increase reference
+        self._reference_table.increase_reference(
+            runtime_env, serialized_env, request.source_process
+        )
+
+        if serialized_env not in self._env_locks:
+            # async lock to prevent the same env being concurrently installed
+            self._env_locks[serialized_env] = asyncio.Lock()
+
+        async with self._env_locks[serialized_env]:
+            if serialized_env in self._env_cache:
+                serialized_context = self._env_cache[serialized_env]
+                result = self._env_cache[serialized_env]
+                if result.success:
+                    context = result.result
+                    self._logger.info(
+                        "Runtime env already created "
+                        f"successfully. Env: {serialized_env}, "
+                        f"context: {context}"
+                    )
+                    return runtime_env_agent_pb2.GetOrCreateRuntimeEnvReply(
+                        status=runtime_env_agent_pb2.AGENT_RPC_STATUS_OK,
+                        serialized_runtime_env_context=context,
+                    )
+                else:
+                    error_message = result.result
+                    self._logger.info(
+                        "Runtime env already failed. "
+                        f"Env: {serialized_env}, "
+                        f"err: {error_message}"
+                    )
+                    # Recover the reference.
+                    self._reference_table.decrease_reference(
+                        runtime_env, serialized_env, request.source_process
+                    )
+                    return runtime_env_agent_pb2.GetOrCreateRuntimeEnvReply(
+                        status=runtime_env_agent_pb2.AGENT_RPC_STATUS_FAILED,
+                        error_message=f"{self._node_prefix}{error_message}",
+                    )
+
+            if SLEEP_FOR_TESTING_S:
+                self._logger.info(f"Sleeping for {SLEEP_FOR_TESTING_S}s.")
+                time.sleep(int(SLEEP_FOR_TESTING_S))
+
+            runtime_env_config = RuntimeEnvConfig.from_proto(request.runtime_env_config)
+
+            # accroding to the document of `asyncio.wait_for`,
+            # None means disable timeout logic
+            setup_timeout_seconds = (
+                None
+                if runtime_env_config["setup_timeout_seconds"] == -1
+                else runtime_env_config["setup_timeout_seconds"]
+            )
+
+            start = time.perf_counter()
+            (
+                successful,
+                serialized_context,
+                error_message,
+            ) = await _create_runtime_env_with_retry(
+                runtime_env,
+                setup_timeout_seconds,
+                runtime_env_config,
+            )
+            creation_time_ms = int(round((time.perf_counter() - start) * 1000, 0))
+            if not successful:
+                # Recover the reference.
+                self._reference_table.decrease_reference(
+                    runtime_env, serialized_env, request.source_process
+                )
+            # Add the result to env cache.
+            self._env_cache[serialized_env] = CreatedEnvResult(
+                successful,
+                serialized_context if successful else error_message,
+                creation_time_ms,
+            )
+            # Reply the RPC
+            return runtime_env_agent_pb2.GetOrCreateRuntimeEnvReply(
+                status=runtime_env_agent_pb2.AGENT_RPC_STATUS_OK
+                if successful
+                else runtime_env_agent_pb2.AGENT_RPC_STATUS_FAILED,
+                serialized_runtime_env_context=serialized_context,
+                error_message=f"{self._node_prefix}{error_message}"
+                if not successful
+                else "",
+            )
+
+    async def DeleteRuntimeEnvIfPossible(self, request):
+        self._logger.info(
+            f"Got request from {request.source_process} to decrease "
+            "reference for runtime env: "
+            f"{request.serialized_runtime_env}."
+        )
+
+        try:
+            runtime_env = RuntimeEnv.deserialize(request.serialized_runtime_env)
+        except Exception as e:
+            self._logger.exception(
+                "[Decrease] Failed to parse runtime env: "
+                f"{request.serialized_runtime_env}"
+            )
+
+            error_message = "".join(
+                traceback.format_exception(type(e), e, e.__traceback__)
+            )
+
+            return runtime_env_agent_pb2.GetOrCreateRuntimeEnvReply(
+                status=runtime_env_agent_pb2.AGENT_RPC_STATUS_FAILED,
+                error_message=f"{self._node_prefix}{error_message}",
+            )
+
+        try:
+            self._reference_table.decrease_reference(
+                runtime_env, request.serialized_runtime_env, request.source_process
+            )
+        except Exception as e:
+            return runtime_env_agent_pb2.DeleteRuntimeEnvIfPossibleReply(
+                status=runtime_env_agent_pb2.AGENT_RPC_STATUS_FAILED,
+                error_message=f"{self._node_prefix}Failed to decrement reference for runtime env for {str(e)}",
+            )
+
+        return runtime_env_agent_pb2.DeleteRuntimeEnvIfPossibleReply(
+            status=runtime_env_agent_pb2.AGENT_RPC_STATUS_OK
+        )
+
+    async def GetRuntimeEnvsInfo(self, request):
+        """Return the runtime env information of the node."""
+        # TODO(sang): Currently, it only includes runtime_env information.
+        # We should include the URI information which includes,
+        # URIs
+        # Caller
+        # Ref counts
+        # Cache information
+        # Metrics (creation time & success)
+        # Deleted URIs
+        limit = request.limit if request.HasField("limit") else -1
+        runtime_env_states = defaultdict(ProtoRuntimeEnvState)
+        runtime_env_refs = self._reference_table.runtime_env_refs
+        for runtime_env, ref_cnt in runtime_env_refs.items():
+            runtime_env_states[runtime_env].runtime_env = runtime_env
+            runtime_env_states[runtime_env].ref_cnt = ref_cnt
+        for runtime_env, result in self._env_cache.items():
+            runtime_env_states[runtime_env].runtime_env = runtime_env
+            runtime_env_states[runtime_env].success = result.success
+            if not result.success:
+                runtime_env_states[runtime_env].error = result.result
+            runtime_env_states[runtime_env].creation_time_ms = result.creation_time_ms
+
+        reply = runtime_env_agent_pb2.GetRuntimeEnvsInfoReply()
+        count = 0
+        for runtime_env_state in runtime_env_states.values():
+            if limit != -1 and count >= limit:
+                break
+            count += 1
+            reply.runtime_env_states.append(runtime_env_state)
+        reply.total = len(runtime_env_states)
+        return reply

python/ray/_private/runtime_env/agent/runtime_env_consts.py

@@ -0,0 +1,20 @@
+import ray._private.ray_constants as ray_constants
+
+RUNTIME_ENV_RETRY_TIMES = ray_constants.env_integer("RUNTIME_ENV_RETRY_TIMES", 3)
+
+RUNTIME_ENV_RETRY_INTERVAL_MS = ray_constants.env_integer(
+    "RUNTIME_ENV_RETRY_INTERVAL_MS", 1000
+)
+
+# Cache TTL for bad runtime env. After this time, delete the cache and retry to create
+# runtime env if needed.
+BAD_RUNTIME_ENV_CACHE_TTL_SECONDS = ray_constants.env_integer(
+    "BAD_RUNTIME_ENV_CACHE_TTL_SECONDS", 60 * 10
+)
+
+RUNTIME_ENV_LOG_FILENAME = "runtime_env.log"
+RUNTIME_ENV_AGENT_PORT_PREFIX = "RUNTIME_ENV_AGENT_PORT_PREFIX:"
+RUNTIME_ENV_AGENT_LOG_FILENAME = "runtime_env_agent.log"
+RUNTIME_ENV_AGENT_CHECK_PARENT_INTERVAL_S_ENV_NAME = (
+    "RAY_RUNTIME_ENV_AGENT_CHECK_PARENT_INTERVAL_S"  # noqa
+)

python/ray/_private/runtime_env/conda.py

@@ -0,0 +1,400 @@
+import hashlib
+import json
+import logging
+import os
+import runpy
+import shutil
+import subprocess
+import sys
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+
+import yaml
+from filelock import FileLock
+
+import ray
+from ray._common.utils import (
+    get_or_create_event_loop,
+    try_to_create_directory,
+)
+from ray._private.runtime_env.conda_utils import (
+    create_conda_env_if_needed,
+    delete_conda_env,
+    get_conda_activate_commands,
+    get_conda_envs,
+    get_conda_info_json,
+)
+from ray._private.runtime_env.context import RuntimeEnvContext
+from ray._private.runtime_env.packaging import Protocol, parse_uri
+from ray._private.runtime_env.plugin import RuntimeEnvPlugin
+from ray._private.runtime_env.validation import parse_and_validate_conda
+from ray._private.utils import (
+    get_directory_size_bytes,
+    get_master_wheel_url,
+    get_release_wheel_url,
+    get_wheel_filename,
+)
+
+default_logger = logging.getLogger(__name__)
+
+_WIN32 = os.name == "nt"
+
+
+def _resolve_current_ray_path() -> str:
+    # When ray is built from source with pip install -e,
+    # ray.__file__ returns .../python/ray/__init__.py and this function returns
+    # ".../python".
+    # When ray is installed from a prebuilt binary, ray.__file__ returns
+    # .../site-packages/ray/__init__.py and this function returns
+    # ".../site-packages".
+    return os.path.split(os.path.split(ray.__file__)[0])[0]
+
+
+def _get_ray_setup_spec():
+    """Find the Ray setup_spec from the currently running Ray.
+
+    This function works even when Ray is built from source with pip install -e.
+    """
+    ray_source_python_path = _resolve_current_ray_path()
+    setup_py_path = os.path.join(ray_source_python_path, "setup.py")
+    return runpy.run_path(setup_py_path)["setup_spec"]
+
+
+def _resolve_install_from_source_ray_dependencies():
+    """Find the Ray dependencies when Ray is installed from source."""
+    deps = (
+        _get_ray_setup_spec().install_requires + _get_ray_setup_spec().extras["default"]
+    )
+    # Remove duplicates
+    return list(set(deps))
+
+
+def _inject_ray_to_conda_site(
+    conda_path, logger: Optional[logging.Logger] = default_logger
+):
+    """Write the current Ray site package directory to a new site"""
+    if _WIN32:
+        python_binary = os.path.join(conda_path, "python")
+    else:
+        python_binary = os.path.join(conda_path, "bin/python")
+    site_packages_path = (
+        subprocess.check_output(
+            [
+                python_binary,
+                "-c",
+                "import sysconfig; print(sysconfig.get_paths()['purelib'])",
+            ]
+        )
+        .decode()
+        .strip()
+    )
+
+    ray_path = _resolve_current_ray_path()
+    logger.warning(
+        f"Injecting {ray_path} to environment site-packages {site_packages_path} "
+        "because _inject_current_ray flag is on."
+    )
+
+    maybe_ray_dir = os.path.join(site_packages_path, "ray")
+    if os.path.isdir(maybe_ray_dir):
+        logger.warning(f"Replacing existing ray installation with {ray_path}")
+        shutil.rmtree(maybe_ray_dir)
+
+    # See usage of *.pth file at
+    # https://docs.python.org/3/library/site.html
+    with open(os.path.join(site_packages_path, "ray_shared.pth"), "w") as f:
+        f.write(ray_path)
+
+
+def _current_py_version():
+    return ".".join(map(str, sys.version_info[:3]))  # like 3.6.10
+
+
+def current_ray_pip_specifier(
+    logger: Optional[logging.Logger] = default_logger,
+) -> Optional[str]:
+    """The pip requirement specifier for the running version of Ray.
+
+    Args:
+        logger: Logger used to warn when the running Ray version cannot be
+            detected (e.g. when running a source build).
+
+    Returns:
+        A string which can be passed to `pip install` to install the
+        currently running Ray version, or None if running on a version
+        built from source locally (likely if you are developing Ray).
+
+    Examples:
+        Returns "https://s3-us-west-2.amazonaws.com/ray-wheels/[..].whl"
+            if running a stable release, a nightly or a specific commit
+    """
+    if os.environ.get("RAY_CI_POST_WHEEL_TESTS"):
+        # Running in Buildkite CI after the wheel has been built.
+        # Wheels are at in the ray/.whl directory, but use relative path to
+        # allow for testing locally if needed.
+        return os.path.join(
+            Path(ray.__file__).resolve().parents[2], ".whl", get_wheel_filename()
+        )
+    elif ray.__commit__ == "{{RAY_COMMIT_SHA}}":
+        # Running on a version built from source locally.
+        if os.environ.get("RAY_RUNTIME_ENV_LOCAL_DEV_MODE") != "1":
+            logger.warning(
+                "Current Ray version could not be detected, most likely "
+                "because you have manually built Ray from source.  To use "
+                "runtime_env in this case, set the environment variable "
+                "RAY_RUNTIME_ENV_LOCAL_DEV_MODE=1."
+            )
+        return None
+    elif "dev" in ray.__version__:
+        # Running on a nightly wheel.
+        return get_master_wheel_url()
+    else:
+        return get_release_wheel_url()
+
+
+def inject_dependencies(
+    conda_dict: Dict[Any, Any],
+    py_version: str,
+    pip_dependencies: Optional[List[str]] = None,
+) -> Dict[Any, Any]:
+    """Add Ray, Python and (optionally) extra pip dependencies to a conda dict.
+
+    Args:
+        conda_dict: A dict representing the JSON-serialized conda
+            environment YAML file.  This dict will be modified and returned.
+        py_version: A string representing a Python version to inject
+            into the conda dependencies, e.g. "3.7.7"
+        pip_dependencies: A list of pip dependencies that
+            will be prepended to the list of pip dependencies in
+            the conda dict.  If the conda dict does not already have a "pip"
+            field, one will be created.
+    Returns:
+        The modified dict.  (Note: the input argument conda_dict is modified
+        and returned.)
+    """
+    if pip_dependencies is None:
+        pip_dependencies = []
+    if conda_dict.get("dependencies") is None:
+        conda_dict["dependencies"] = []
+
+    # Inject Python dependency.
+    deps = conda_dict["dependencies"]
+
+    # Add current python dependency.  If the user has already included a
+    # python version dependency, conda will raise a readable error if the two
+    # are incompatible, e.g:
+    #   ResolvePackageNotFound: - python[version='3.5.*,>=3.6']
+    deps.append(f"python={py_version}")
+
+    if "pip" not in deps:
+        deps.append("pip")
+
+    # Insert pip dependencies.
+    found_pip_dict = False
+    for dep in deps:
+        if isinstance(dep, dict) and dep.get("pip") and isinstance(dep["pip"], list):
+            dep["pip"] = pip_dependencies + dep["pip"]
+            found_pip_dict = True
+            break
+    if not found_pip_dict:
+        deps.append({"pip": pip_dependencies})
+
+    return conda_dict
+
+
+def _get_conda_env_hash(conda_dict: Dict) -> str:
+    # Set `sort_keys=True` so that different orderings yield the same hash.
+    serialized_conda_spec = json.dumps(conda_dict, sort_keys=True)
+    hash = hashlib.sha1(serialized_conda_spec.encode("utf-8")).hexdigest()
+    return hash
+
+
+def get_uri(runtime_env: Dict) -> Optional[str]:
+    """Return `"conda://<hashed_dependencies>"`, or None if no GC required."""
+    conda = runtime_env.get("conda")
+    if conda is not None:
+        if isinstance(conda, str):
+            # User-preinstalled conda env.  We don't garbage collect these, so
+            # we don't track them with URIs.
+            uri = None
+        elif isinstance(conda, dict):
+            uri = f"conda://{_get_conda_env_hash(conda_dict=conda)}"
+        else:
+            raise TypeError(
+                "conda field received by RuntimeEnvAgent must be "
+                f"str or dict, not {type(conda).__name__}."
+            )
+    else:
+        uri = None
+    return uri
+
+
+def _get_conda_dict_with_ray_inserted(
+    runtime_env: "RuntimeEnv",  # noqa: F821
+    logger: Optional[logging.Logger] = default_logger,
+) -> Dict[str, Any]:
+    """Returns the conda spec with the Ray and `python` dependency inserted."""
+    conda_dict = json.loads(runtime_env.conda_config())
+    assert conda_dict is not None
+
+    ray_pip = current_ray_pip_specifier(logger=logger)
+    if ray_pip:
+        extra_pip_dependencies = [ray_pip, "ray[default]"]
+    elif runtime_env.get_extension("_inject_current_ray"):
+        extra_pip_dependencies = _resolve_install_from_source_ray_dependencies()
+    else:
+        extra_pip_dependencies = []
+    conda_dict = inject_dependencies(
+        conda_dict, _current_py_version(), extra_pip_dependencies
+    )
+    return conda_dict
+
+
+class CondaPlugin(RuntimeEnvPlugin):
+
+    name = "conda"
+
+    def __init__(self, resources_dir: str):
+        self._resources_dir = os.path.join(resources_dir, "conda")
+        try_to_create_directory(self._resources_dir)
+
+        # It is not safe for multiple processes to install conda envs
+        # concurrently, even if the envs are different, so use a global
+        # lock for all conda installs and deletions.
+        # See https://github.com/ray-project/ray/issues/17086
+        self._installs_and_deletions_file_lock = os.path.join(
+            self._resources_dir, "ray-conda-installs-and-deletions.lock"
+        )
+        # A set of named conda environments (instead of yaml or dict)
+        # that are validated to exist.
+        # NOTE: It has to be only used within the same thread, which
+        # is an event loop.
+        # Also, we don't need to GC this field because it is pretty small.
+        self._validated_named_conda_env = set()
+
+    def _get_path_from_hash(self, hash: str) -> str:
+        """Generate a path from the hash of a conda or pip spec.
+
+        The output path also functions as the name of the conda environment
+        when using the `--prefix` option to `conda create` and `conda remove`.
+
+        Example output:
+            /tmp/ray/session_2021-11-03_16-33-59_356303_41018/runtime_resources
+                /conda/ray-9a7972c3a75f55e976e620484f58410c920db091
+        """
+        return os.path.join(self._resources_dir, hash)
+
+    def get_uris(self, runtime_env: "RuntimeEnv") -> List[str]:  # noqa: F821
+        """Return the conda URI from the RuntimeEnv if it exists, else return []."""
+        conda_uri = runtime_env.conda_uri()
+        if conda_uri:
+            return [conda_uri]
+        return []
+
+    def delete_uri(
+        self, uri: str, logger: Optional[logging.Logger] = default_logger
+    ) -> int:
+        """Delete URI and return the number of bytes deleted."""
+        logger.info(f"Got request to delete URI {uri}")
+        protocol, hash = parse_uri(uri)
+        if protocol != Protocol.CONDA:
+            raise ValueError(
+                "CondaPlugin can only delete URIs with protocol "
+                f"conda.  Received protocol {protocol}, URI {uri}"
+            )
+
+        conda_env_path = self._get_path_from_hash(hash)
+        local_dir_size = get_directory_size_bytes(conda_env_path)
+
+        with FileLock(self._installs_and_deletions_file_lock):
+            successful = delete_conda_env(prefix=conda_env_path, logger=logger)
+        if not successful:
+            logger.warning(f"Error when deleting conda env {conda_env_path}. ")
+            return 0
+
+        return local_dir_size
+
+    async def create(
+        self,
+        uri: Optional[str],
+        runtime_env: "RuntimeEnv",  # noqa: F821
+        context: RuntimeEnvContext,
+        logger: logging.Logger = default_logger,
+    ) -> int:
+        if not runtime_env.has_conda():
+            return 0
+
+        def _create():
+            result = parse_and_validate_conda(runtime_env.get("conda"))
+
+            if isinstance(result, str):
+                # The conda env name is given.
+                # In this case, we only verify if the given
+                # conda env exists.
+
+                # If the env is already validated, do nothing.
+                if result in self._validated_named_conda_env:
+                    return 0
+
+                conda_info = get_conda_info_json()
+                envs = get_conda_envs(conda_info)
+
+                # We accept `result` as a conda name or full path.
+                if not any(result == env[0] or result == env[1] for env in envs):
+                    raise ValueError(
+                        f"The given conda environment '{result}' "
+                        f"from the runtime env {runtime_env} doesn't "
+                        "exist from the output of `conda info --json`. "
+                        "You can only specify an env that already exists. "
+                        f"Please make sure to create an env {result} "
+                    )
+                self._validated_named_conda_env.add(result)
+                return 0
+
+            logger.debug(
+                "Setting up conda for runtime_env: " f"{runtime_env.serialize()}"
+            )
+            protocol, hash = parse_uri(uri)
+            conda_env_name = self._get_path_from_hash(hash)
+
+            conda_dict = _get_conda_dict_with_ray_inserted(runtime_env, logger=logger)
+
+            logger.info(f"Setting up conda environment with {runtime_env}")
+            with FileLock(self._installs_and_deletions_file_lock):
+                try:
+                    conda_yaml_file = os.path.join(
+                        self._resources_dir, "environment.yml"
+                    )
+                    with open(conda_yaml_file, "w") as file:
+                        yaml.dump(conda_dict, file)
+                    create_conda_env_if_needed(
+                        conda_yaml_file, prefix=conda_env_name, logger=logger
+                    )
+                finally:
+                    os.remove(conda_yaml_file)
+
+                if runtime_env.get_extension("_inject_current_ray"):
+                    _inject_ray_to_conda_site(conda_path=conda_env_name, logger=logger)
+            logger.info(f"Finished creating conda environment at {conda_env_name}")
+            return get_directory_size_bytes(conda_env_name)
+
+        loop = get_or_create_event_loop()
+        return await loop.run_in_executor(None, _create)
+
+    def modify_context(
+        self,
+        uris: List[str],
+        runtime_env: "RuntimeEnv",  # noqa: F821
+        context: RuntimeEnvContext,
+        logger: Optional[logging.Logger] = default_logger,
+    ):
+        if not runtime_env.has_conda():
+            return
+
+        if runtime_env.conda_env_name():
+            conda_env_name = runtime_env.conda_env_name()
+        else:
+            protocol, hash = parse_uri(runtime_env.conda_uri())
+            conda_env_name = self._get_path_from_hash(hash)
+        context.py_executable = "python"
+        context.command_prefix += get_conda_activate_commands(conda_env_name)

python/ray/_private/runtime_env/conda_utils.py

@@ -0,0 +1,285 @@
+import hashlib
+import json
+import logging
+import os
+import shutil
+import subprocess
+from typing import List, Optional, Tuple, Union
+
+"""Utilities for conda.  Adapted from https://github.com/mlflow/mlflow."""
+
+# Name of environment variable indicating a path to a conda installation. Ray
+# will default to running "conda" if unset.
+RAY_CONDA_HOME = "RAY_CONDA_HOME"
+
+_WIN32 = os.name == "nt"
+
+
+def get_conda_activate_commands(conda_env_name: str) -> List[str]:
+    """
+    Get a list of commands to run to silently activate the given conda env.
+    """
+    #  Checking for newer conda versions
+    if not _WIN32 and ("CONDA_EXE" in os.environ or RAY_CONDA_HOME in os.environ):
+        conda_path = get_conda_bin_executable("conda")
+        activate_conda_env = [
+            ".",
+            f"{os.path.dirname(conda_path)}/../etc/profile.d/conda.sh",
+            "&&",
+        ]
+        activate_conda_env += ["conda", "activate", conda_env_name]
+
+    else:
+        activate_path = get_conda_bin_executable("activate")
+        if not _WIN32:
+            # Use bash command syntax
+            activate_conda_env = ["source", activate_path, conda_env_name]
+        else:
+            conda_path = get_conda_bin_executable("conda")
+            activate_conda_env = [conda_path, "activate", conda_env_name]
+    return activate_conda_env + ["1>&2", "&&"]
+
+
+def get_conda_bin_executable(executable_name: str) -> str:
+    """
+    Return path to the specified executable, assumed to be discoverable within
+    a conda installation.
+
+    The conda home directory (expected to contain a 'bin' subdirectory on
+    linux) is configurable via the ``RAY_CONDA_HOME`` environment variable. If
+    ``RAY_CONDA_HOME`` is unspecified, try the ``CONDA_EXE`` environment
+    variable set by activating conda. If neither is specified, this method
+    returns `executable_name`.
+    """
+    conda_home = os.environ.get(RAY_CONDA_HOME)
+    if conda_home:
+        if _WIN32:
+            candidate = os.path.join(conda_home, "%s.exe" % executable_name)
+            if os.path.exists(candidate):
+                return candidate
+            candidate = os.path.join(conda_home, "%s.bat" % executable_name)
+            if os.path.exists(candidate):
+                return candidate
+        else:
+            return os.path.join(conda_home, "bin/%s" % executable_name)
+    else:
+        conda_home = "."
+    # Use CONDA_EXE as per https://github.com/conda/conda/issues/7126
+    if "CONDA_EXE" in os.environ:
+        conda_bin_dir = os.path.dirname(os.environ["CONDA_EXE"])
+        if _WIN32:
+            candidate = os.path.join(conda_home, "%s.exe" % executable_name)
+            if os.path.exists(candidate):
+                return candidate
+            candidate = os.path.join(conda_home, "%s.bat" % executable_name)
+            if os.path.exists(candidate):
+                return candidate
+        else:
+            return os.path.join(conda_bin_dir, executable_name)
+    if _WIN32:
+        return executable_name + ".bat"
+    return executable_name
+
+
+def _get_conda_env_name(conda_env_path: str) -> str:
+    conda_env_contents = open(conda_env_path).read()
+    return "ray-%s" % hashlib.sha1(conda_env_contents.encode("utf-8")).hexdigest()
+
+
+def create_conda_env_if_needed(
+    conda_yaml_file: str, prefix: str, logger: Optional[logging.Logger] = None
+) -> None:
+    """
+    Given a conda YAML, creates a conda environment containing the required
+    dependencies if such a conda environment doesn't already exist.
+    Args:
+        conda_yaml_file: The path to a conda `environment.yml` file.
+        prefix: Directory to install the environment into via
+            the `--prefix` option to conda create.  This also becomes the name
+            of the conda env; i.e. it can be passed into `conda activate` and
+            `conda remove`
+        logger: Logger used to surface progress and errors; defaults to the
+            module logger when not provided.
+    """
+    if logger is None:
+        logger = logging.getLogger(__name__)
+
+    conda_path = get_conda_bin_executable("conda")
+    try:
+        exec_cmd([conda_path, "--help"], throw_on_error=False)
+    except (EnvironmentError, FileNotFoundError):
+        raise ValueError(
+            f"Could not find Conda executable at '{conda_path}'. "
+            "Ensure Conda is installed as per the instructions at "
+            "https://conda.io/projects/conda/en/latest/"
+            "user-guide/install/index.html. "
+            "You can also configure Ray to look for a specific "
+            f"Conda executable by setting the {RAY_CONDA_HOME} "
+            "environment variable to the path of the Conda executable."
+        )
+
+    _, stdout, _ = exec_cmd([conda_path, "env", "list", "--json"])
+    envs = json.loads(stdout[stdout.index("{") :])["envs"]
+
+    if prefix in envs:
+        logger.info(f"Conda environment {prefix} already exists.")
+        return
+
+    create_cmd = [
+        conda_path,
+        "env",
+        "create",
+        "--file",
+        conda_yaml_file,
+        "--prefix",
+        prefix,
+    ]
+
+    logger.info(f"Creating conda environment {prefix}")
+    exit_code, output = exec_cmd_stream_to_logger(create_cmd, logger)
+    if exit_code != 0:
+        if os.path.exists(prefix):
+            shutil.rmtree(prefix)
+        raise RuntimeError(
+            f"Failed to install conda environment {prefix}:\nOutput:\n{output}"
+        )
+
+
+def delete_conda_env(prefix: str, logger: Optional[logging.Logger] = None) -> bool:
+    if logger is None:
+        logger = logging.getLogger(__name__)
+
+    logger.info(f"Deleting conda environment {prefix}")
+
+    conda_path = get_conda_bin_executable("conda")
+    delete_cmd = [conda_path, "remove", "-p", prefix, "--all", "-y"]
+    exit_code, output = exec_cmd_stream_to_logger(delete_cmd, logger)
+
+    if exit_code != 0:
+        logger.debug(f"Failed to delete conda environment {prefix}:\n{output}")
+        return False
+
+    return True
+
+
+def get_conda_env_list() -> list:
+    """
+    Get conda env list in full paths.
+    """
+    conda_path = get_conda_bin_executable("conda")
+    try:
+        exec_cmd([conda_path, "--help"], throw_on_error=False)
+    except EnvironmentError:
+        raise ValueError(f"Could not find Conda executable at {conda_path}.")
+    _, stdout, _ = exec_cmd([conda_path, "env", "list", "--json"])
+    envs = json.loads(stdout)["envs"]
+    return envs
+
+
+def get_conda_info_json() -> dict:
+    """
+    Get `conda info --json` output.
+
+    Returns dict of conda info. See [1] for more details. We mostly care about these
+    keys:
+
+    - `conda_prefix`: str The path to the conda installation.
+    - `envs`: List[str] absolute paths to conda environments.
+
+    [1] https://github.com/conda/conda/blob/main/conda/cli/main_info.py
+    """
+    conda_path = get_conda_bin_executable("conda")
+    try:
+        exec_cmd([conda_path, "--help"], throw_on_error=False)
+    except EnvironmentError:
+        raise ValueError(f"Could not find Conda executable at {conda_path}.")
+    _, stdout, _ = exec_cmd([conda_path, "info", "--json"])
+    return json.loads(stdout)
+
+
+def get_conda_envs(conda_info: dict) -> List[Tuple[str, str]]:
+    """
+    Gets the conda environments, as a list of (name, path) tuples.
+    """
+    prefix = conda_info["conda_prefix"]
+    ret = []
+    for env in conda_info["envs"]:
+        if env == prefix:
+            ret.append(("base", env))
+        else:
+            ret.append((os.path.basename(env), env))
+    return ret
+
+
+class ShellCommandException(Exception):
+    pass
+
+
+def exec_cmd(
+    cmd: List[str], throw_on_error: bool = True, logger: Optional[logging.Logger] = None
+) -> Union[int, Tuple[int, str, str]]:
+    """
+    Runs a command as a child process.
+
+    A convenience wrapper for running a command from a Python script.
+
+    Note on the return value: A tuple of the exit code,
+    standard output and standard error is returned.
+
+    Args:
+        cmd: the command to run, as a list of strings
+        throw_on_error: if true, raises an Exception if the exit code of the
+            program is nonzero
+        logger: Unused; retained for API compatibility.
+
+    Returns:
+        A tuple of (exit_code, stdout, stderr) from the child process.
+    """
+    child = subprocess.Popen(
+        cmd,
+        stdout=subprocess.PIPE,
+        stdin=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+        universal_newlines=True,
+    )
+    (stdout, stderr) = child.communicate()
+    exit_code = child.wait()
+    if throw_on_error and exit_code != 0:
+        raise ShellCommandException(
+            "Non-zero exit code: %s\n\nSTDOUT:\n%s\n\nSTDERR:%s"
+            % (exit_code, stdout, stderr)
+        )
+    return exit_code, stdout, stderr
+
+
+def exec_cmd_stream_to_logger(
+    cmd: List[str], logger: logging.Logger, n_lines: int = 50, **kwargs
+) -> Tuple[int, str]:
+    """Runs a command as a child process, streaming output to the logger.
+
+    The last n_lines lines of output are also returned (stdout and stderr).
+    """
+    if "env" in kwargs and _WIN32 and "PATH" not in [x.upper() for x in kwargs.keys]:
+        raise ValueError("On windows, Popen requires 'PATH' in 'env'")
+    child = subprocess.Popen(
+        cmd,
+        universal_newlines=True,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.STDOUT,
+        **kwargs,
+    )
+    last_n_lines = []
+    with child.stdout:
+        for line in iter(child.stdout.readline, b""):
+            exit_code = child.poll()
+            if exit_code is not None:
+                break
+            line = line.strip()
+            if not line:
+                continue
+            last_n_lines.append(line.strip())
+            last_n_lines = last_n_lines[-n_lines:]
+            logger.info(line.strip())
+
+    exit_code = child.wait()
+    return exit_code, "\n".join(last_n_lines)

python/ray/_private/runtime_env/constants.py

@@ -0,0 +1,28 @@
+# Env var set by job manager to pass runtime env and metadata to subprocess
+RAY_JOB_CONFIG_JSON_ENV_VAR = "RAY_JOB_CONFIG_JSON_ENV_VAR"
+
+# The plugin config which should be loaded when ray cluster starts.
+# It is a json formatted config,
+# e.g. [{"class": "xxx.xxx.xxx_plugin", "priority": 10}].
+RAY_RUNTIME_ENV_PLUGINS_ENV_VAR = "RAY_RUNTIME_ENV_PLUGINS"
+
+# The field name of plugin class in the plugin config.
+RAY_RUNTIME_ENV_CLASS_FIELD_NAME = "class"
+
+# The field name of priority in the plugin config.
+RAY_RUNTIME_ENV_PRIORITY_FIELD_NAME = "priority"
+
+# The default priority of runtime env plugin.
+RAY_RUNTIME_ENV_PLUGIN_DEFAULT_PRIORITY = 10
+
+# The minimum priority of runtime env plugin.
+RAY_RUNTIME_ENV_PLUGIN_MIN_PRIORITY = 0
+
+# The maximum priority of runtime env plugin.
+RAY_RUNTIME_ENV_PLUGIN_MAX_PRIORITY = 100
+
+# The schema files or directories of plugins which should be loaded in workers.
+RAY_RUNTIME_ENV_PLUGIN_SCHEMAS_ENV_VAR = "RAY_RUNTIME_ENV_PLUGIN_SCHEMAS"
+
+# The file suffix of runtime env plugin schemas.
+RAY_RUNTIME_ENV_PLUGIN_SCHEMA_SUFFIX = ".json"

python/ray/_private/runtime_env/context.py

@@ -0,0 +1,108 @@
+import json
+import logging
+import os
+import shlex
+import subprocess
+import sys
+from typing import Dict, List, Optional
+
+from ray._private.services import get_ray_jars_dir
+from ray._private.utils import update_envs
+from ray.core.generated.common_pb2 import Language
+from ray.util.annotations import DeveloperAPI
+
+logger = logging.getLogger(__name__)
+
+
+@DeveloperAPI
+class RuntimeEnvContext:
+    """A context used to describe the created runtime env."""
+
+    def __init__(
+        self,
+        command_prefix: List[str] = None,
+        env_vars: Dict[str, str] = None,
+        py_executable: Optional[str] = None,
+        override_worker_entrypoint: Optional[str] = None,
+        java_jars: List[str] = None,
+    ):
+        self.command_prefix = command_prefix or []
+        self.env_vars = env_vars or {}
+        self.py_executable = py_executable or sys.executable
+        self.override_worker_entrypoint: Optional[str] = override_worker_entrypoint
+        self.java_jars = java_jars or []
+
+    def serialize(self) -> str:
+        return json.dumps(self.__dict__)
+
+    @staticmethod
+    def deserialize(json_string):
+        return RuntimeEnvContext(**json.loads(json_string))
+
+    def exec_worker(self, passthrough_args: List[str], language: Language):
+        update_envs(self.env_vars)
+
+        if language == Language.PYTHON and sys.platform == "win32":
+            executable = [self.py_executable]
+        elif language == Language.PYTHON:
+            executable = ["exec", self.py_executable]
+        elif language == Language.JAVA:
+            executable = ["java"]
+            ray_jars = os.path.join(get_ray_jars_dir(), "*")
+
+            local_java_jars = []
+            for java_jar in self.java_jars:
+                local_java_jars.append(f"{java_jar}/*")
+                local_java_jars.append(java_jar)
+
+            class_path_args = ["-cp", ray_jars + ":" + str(":".join(local_java_jars))]
+            passthrough_args = class_path_args + passthrough_args
+        elif sys.platform == "win32":
+            executable = []
+        else:
+            executable = ["exec"]
+
+        # By default, raylet uses the path to default_worker.py on host.
+        # However, the path to default_worker.py inside the container
+        # can be different. We need the user to specify the path to
+        # default_worker.py inside the container.
+        if self.override_worker_entrypoint:
+            logger.debug(
+                f"Changing the worker entrypoint from {passthrough_args[0]} to "
+                f"{self.override_worker_entrypoint}."
+            )
+            passthrough_args[0] = self.override_worker_entrypoint
+
+        if sys.platform == "win32":
+
+            def quote(s):
+                s = s.replace("&", "%26")
+                return s
+
+            passthrough_args = [quote(s) for s in passthrough_args]
+
+            cmd = [*self.command_prefix, *executable, *passthrough_args]
+            logger.debug(f"Exec'ing worker with command: {cmd}")
+            subprocess.Popen(cmd, shell=True).wait()
+        else:
+            # We use shlex to do the necessary shell escape
+            # of special characters in passthrough_args.
+            passthrough_args = [shlex.quote(s) for s in passthrough_args]
+            cmd = [*self.command_prefix, *executable, *passthrough_args]
+            # TODO(SongGuyang): We add this env to command for macOS because it doesn't
+            # work for the C++ process of `os.execvp`. We should find a better way to
+            # fix it.
+            MACOS_LIBRARY_PATH_ENV_NAME = "DYLD_LIBRARY_PATH"
+            if MACOS_LIBRARY_PATH_ENV_NAME in os.environ:
+                cmd.insert(
+                    0,
+                    f"{MACOS_LIBRARY_PATH_ENV_NAME}="
+                    f"{os.environ[MACOS_LIBRARY_PATH_ENV_NAME]}",
+                )
+            logger.debug(f"Exec'ing worker with command: {cmd}")
+            # PyCharm will monkey patch the os.execvp at
+            # .pycharm_helpers/pydev/_pydev_bundle/pydev_monkey.py
+            # The monkey patched os.execvp function has a different
+            # signature. So, we use os.execvp("executable", args=[])
+            # instead of os.execvp(file="executable", args=[])
+            os.execvp("bash", args=["bash", "-c", " ".join(cmd)])

python/ray/_private/runtime_env/default_impl.py

@@ -0,0 +1,5 @@
+from ray._private.runtime_env.image_uri import ImageURIPlugin
+
+
+def get_image_uri_plugin_cls():
+    return ImageURIPlugin

python/ray/_private/runtime_env/dependency_utils.py

@@ -0,0 +1,118 @@
+"""Util functions to manage dependency requirements."""
+
+import logging
+import os
+import tempfile
+from contextlib import asynccontextmanager
+from typing import List, Optional, Tuple
+
+from ray._private.runtime_env import virtualenv_utils
+from ray._private.runtime_env.utils import check_output_cmd
+
+INTERNAL_PIP_FILENAME = "ray_runtime_env_internal_pip_requirements.txt"
+MAX_INTERNAL_PIP_FILENAME_TRIES = 100
+
+
+def gen_requirements_txt(requirements_file: str, pip_packages: List[str]):
+    """Dump [pip_packages] to the given [requirements_file] for later env setup."""
+    with open(requirements_file, "w") as file:
+        for line in pip_packages:
+            file.write(line + "\n")
+
+
+@asynccontextmanager
+async def check_ray(python: str, cwd: str, logger: logging.Logger):
+    """A context manager to check ray is not overwritten.
+
+    Currently, we only check ray version and path. It works for virtualenv,
+        - ray is in Python's site-packages.
+        - ray is overwritten during yield.
+        - ray is in virtualenv's site-packages.
+    """
+
+    async def _get_ray_version_and_path() -> Tuple[str, str]:
+        with tempfile.TemporaryDirectory(
+            prefix="check_ray_version_tempfile"
+        ) as tmp_dir:
+            ray_version_path = os.path.join(tmp_dir, "ray_version.txt")
+            check_ray_cmd = [
+                python,
+                "-c",
+                """
+import ray
+with open(r"{ray_version_path}", "wt") as f:
+    f.write(ray.__version__)
+    f.write(" ")
+    f.write(ray.__path__[0])
+                """.format(
+                    ray_version_path=ray_version_path
+                ),
+            ]
+            if virtualenv_utils._WIN32:
+                env = os.environ.copy()
+            else:
+                env = {}
+            output = await check_output_cmd(
+                check_ray_cmd, logger=logger, cwd=cwd, env=env
+            )
+            logger.info(f"try to write ray version information in: {ray_version_path}")
+            with open(ray_version_path, "rt") as f:
+                output = f.read()
+            # print after import ray may have  endings, so we strip them by *_
+            ray_version, ray_path, *_ = [s.strip() for s in output.split()]
+        return ray_version, ray_path
+
+    version, path = await _get_ray_version_and_path()
+    yield
+    actual_version, actual_path = await _get_ray_version_and_path()
+    if actual_version != version:
+        raise RuntimeError(
+            "Changing the ray version is not allowed: \n"
+            f"  current version: {actual_version}, "
+            f"  expect version: {version}, "
+            f"  current path: {actual_path}, "
+            f"  expect path: {path}, "
+            "Please ensure the dependencies in the runtime_env pip field "
+            "do not install a different version of Ray."
+        )
+    if actual_path != path:
+        logger.info(
+            f"Detected new Ray package with the same version at {actual_path} (vs system {path})."
+        )
+
+
+def get_requirements_file(target_dir: str, pip_list: Optional[List[str]]) -> str:
+    """Returns the path to the requirements file to use for this runtime env.
+
+    If pip_list is not None, we will check if the internal pip filename is in any of
+    the entries of pip_list. If so, we will append numbers to the end of the
+    filename until we find one that doesn't conflict. This prevents infinite
+    recursion if the user specifies the internal pip filename in their pip list.
+
+    Args:
+        target_dir: The directory to store the requirements file in.
+        pip_list: A list of pip requirements specified by the user.
+
+    Returns:
+        The path to the requirements file to use for this runtime env.
+    """
+
+    def filename_in_pip_list(filename: str) -> bool:
+        for pip_entry in pip_list:
+            if filename in pip_entry:
+                return True
+        return False
+
+    filename = INTERNAL_PIP_FILENAME
+    if pip_list is not None:
+        i = 1
+        while filename_in_pip_list(filename) and i < MAX_INTERNAL_PIP_FILENAME_TRIES:
+            filename = f"{INTERNAL_PIP_FILENAME}.{i}"
+            i += 1
+        if i == MAX_INTERNAL_PIP_FILENAME_TRIES:
+            raise RuntimeError(
+                "Could not find a valid filename for the internal "
+                "pip requirements file. Please specify a different "
+                "pip list in your runtime env."
+            )
+    return os.path.join(target_dir, filename)

python/ray/_private/runtime_env/image_uri.py

@@ -0,0 +1,229 @@
+import asyncio
+import logging
+import os
+import tempfile
+from typing import List, Optional
+
+from ray._private.runtime_env.context import RuntimeEnvContext
+from ray._private.runtime_env.plugin import RuntimeEnvPlugin
+
+default_logger = logging.getLogger(__name__)
+
+
+async def _create_impl(image_uri: str, logger: logging.Logger):
+    # Pull image if it doesn't exist
+    # Also get path to `default_worker.py` inside the image.
+    with tempfile.TemporaryDirectory() as tmpdir:
+        os.chmod(tmpdir, 0o777)
+        result_file = os.path.join(tmpdir, "worker_path.txt")
+        get_worker_path_script = """
+import ray._private.workers.default_worker as dw
+with open('/shared/worker_path.txt', 'w') as f:
+    f.write(dw.__file__)
+"""
+        cmd = [
+            "podman",
+            "run",
+            "--rm",
+            "-v",
+            f"{tmpdir}:/shared:Z",
+            image_uri,
+            "python",
+            "-c",
+            get_worker_path_script,
+        ]
+
+        logger.info("Pulling image %s", image_uri)
+
+        process = await asyncio.create_subprocess_exec(
+            *cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE
+        )
+
+        stdout, stderr = await process.communicate()
+
+        if process.returncode != 0:
+            raise RuntimeError(
+                f"Podman command failed: cmd={cmd}, returncode={process.returncode}, stdout={stdout.decode()}, stderr={stderr.decode()}"
+            )
+
+        if not os.path.exists(result_file):
+            raise FileNotFoundError(
+                f"Worker path file not created when getting worker path for image {image_uri}"
+            )
+
+        with open(result_file, "r") as f:
+            worker_path = f.read().strip()
+
+        if not worker_path.endswith(".py"):
+            raise ValueError(
+                f"Invalid worker path inferred in image {image_uri}: {worker_path}"
+            )
+
+        logger.info(f"Inferred worker path in image {image_uri}: {worker_path}")
+        return worker_path
+
+
+def _modify_context_impl(
+    image_uri: str,
+    worker_path: str,
+    run_options: Optional[List[str]],
+    context: RuntimeEnvContext,
+    logger: logging.Logger,
+    ray_tmp_dir: str,
+):
+    context.override_worker_entrypoint = worker_path
+
+    container_driver = "podman"
+    container_command = [
+        container_driver,
+        "run",
+        "-v",
+        ray_tmp_dir + ":" + ray_tmp_dir,
+        "--cgroup-manager=cgroupfs",
+        "--network=host",
+        "--pid=host",
+        "--ipc=host",
+        # NOTE(zcin): Mounted volumes in rootless containers are
+        # owned by the user `root`. The user on host (which will
+        # usually be `ray` if this is being run in a ray docker
+        # image) who started the container is mapped using user
+        # namespaces to the user `root` in a rootless container. In
+        # order for the Ray Python worker to access the mounted ray
+        # tmp dir, we need to use keep-id mode which maps the user
+        # as itself (instead of as `root`) into the container.
+        # https://www.redhat.com/sysadmin/rootless-podman-user-namespace-modes
+        "--userns=keep-id",
+    ]
+
+    # Environment variables to set in container
+    env_vars = dict()
+
+    # Propagate all host environment variables that have the prefix "RAY_"
+    # This should include RAY_RAYLET_PID
+    for env_var_name, env_var_value in os.environ.items():
+        if env_var_name.startswith("RAY_"):
+            env_vars[env_var_name] = env_var_value
+
+    # Support for runtime_env['env_vars']
+    env_vars.update(context.env_vars)
+
+    # Set environment variables
+    for env_var_name, env_var_value in env_vars.items():
+        container_command.append("--env")
+        container_command.append(f"{env_var_name}='{env_var_value}'")
+
+    # The RAY_JOB_ID environment variable is needed for the default worker.
+    # It won't be set at the time setup() is called, but it will be set
+    # when worker command is executed, so we use RAY_JOB_ID=$RAY_JOB_ID
+    # for the container start command
+    container_command.append("--env")
+    container_command.append("RAY_JOB_ID=$RAY_JOB_ID")
+
+    if run_options:
+        container_command.extend(run_options)
+    # TODO(chenk008): add resource limit
+    container_command.append("--entrypoint")
+    container_command.append("python")
+    container_command.append(image_uri)
+
+    # Example:
+    # podman run -v /tmp/ray:/tmp/ray
+    # --cgroup-manager=cgroupfs --network=host --pid=host --ipc=host
+    # --userns=keep-id --env RAY_RAYLET_PID=23478 --env RAY_JOB_ID=$RAY_JOB_ID
+    # --entrypoint python rayproject/ray:nightly-py39
+    container_command_str = " ".join(container_command)
+    logger.info(f"Starting worker in container with prefix {container_command_str}")
+
+    context.py_executable = container_command_str
+
+
+class ImageURIPlugin(RuntimeEnvPlugin):
+    """Starts worker in a container of a custom image."""
+
+    name = "image_uri"
+
+    @staticmethod
+    def get_compatible_keys():
+        return {"image_uri", "config", "env_vars"}
+
+    def __init__(self, ray_tmp_dir: str):
+        self._ray_tmp_dir = ray_tmp_dir
+
+    async def create(
+        self,
+        uri: Optional[str],
+        runtime_env: "RuntimeEnv",  # noqa: F821
+        context: RuntimeEnvContext,
+        logger: logging.Logger,
+    ) -> float:
+        if not runtime_env.image_uri():
+            return
+
+        self.worker_path = await _create_impl(runtime_env.image_uri(), logger)
+
+    def modify_context(
+        self,
+        uris: List[str],
+        runtime_env: "RuntimeEnv",  # noqa: F821
+        context: RuntimeEnvContext,
+        logger: Optional[logging.Logger] = default_logger,
+    ):
+        if not runtime_env.image_uri():
+            return
+
+        _modify_context_impl(
+            runtime_env.image_uri(),
+            self.worker_path,
+            [],
+            context,
+            logger,
+            self._ray_tmp_dir,
+        )
+
+
+class ContainerPlugin(RuntimeEnvPlugin):
+    """Starts worker in container."""
+
+    name = "container"
+
+    def __init__(self, ray_tmp_dir: str):
+        self._ray_tmp_dir = ray_tmp_dir
+
+    async def create(
+        self,
+        uri: Optional[str],
+        runtime_env: "RuntimeEnv",  # noqa: F821
+        context: RuntimeEnvContext,
+        logger: logging.Logger,
+    ) -> float:
+        if not runtime_env.has_py_container() or not runtime_env.py_container_image():
+            return
+
+        self.worker_path = await _create_impl(runtime_env.py_container_image(), logger)
+
+    def modify_context(
+        self,
+        uris: List[str],
+        runtime_env: "RuntimeEnv",  # noqa: F821
+        context: RuntimeEnvContext,
+        logger: Optional[logging.Logger] = default_logger,
+    ):
+        if not runtime_env.has_py_container() or not runtime_env.py_container_image():
+            return
+
+        if runtime_env.py_container_worker_path():
+            logger.warning(
+                "You are using `container.worker_path`, but the path to "
+                "`default_worker.py` is now automatically detected from the image. "
+                "`container.worker_path` is deprecated and will be removed in future "
+                "versions."
+            )
+
+        _modify_context_impl(
+            runtime_env.py_container_image(),
+            runtime_env.py_container_worker_path() or self.worker_path,
+            runtime_env.py_container_run_options(),
+            context,
+            logger,
+            self._ray_tmp_dir,
+        )

python/ray/_private/runtime_env/java_jars.py

@@ -0,0 +1,104 @@
+import logging
+import os
+from typing import Dict, List, Optional
+
+from ray._common.utils import try_to_create_directory
+from ray._private.runtime_env.context import RuntimeEnvContext
+from ray._private.runtime_env.packaging import (
+    delete_package,
+    download_and_unpack_package,
+    get_local_dir_from_uri,
+    is_jar_uri,
+)
+from ray._private.runtime_env.plugin import RuntimeEnvPlugin
+from ray._private.utils import get_directory_size_bytes
+from ray._raylet import GcsClient
+from ray.exceptions import RuntimeEnvSetupError
+
+default_logger = logging.getLogger(__name__)
+
+
+class JavaJarsPlugin(RuntimeEnvPlugin):
+
+    name = "java_jars"
+
+    def __init__(self, resources_dir: str, gcs_client: GcsClient):
+        self._resources_dir = os.path.join(resources_dir, "java_jars_files")
+        self._gcs_client = gcs_client
+        try_to_create_directory(self._resources_dir)
+
+    def _get_local_dir_from_uri(self, uri: str):
+        return get_local_dir_from_uri(uri, self._resources_dir)
+
+    def delete_uri(
+        self, uri: str, logger: Optional[logging.Logger] = default_logger
+    ) -> int:
+        """Delete URI and return the number of bytes deleted."""
+        local_dir = get_local_dir_from_uri(uri, self._resources_dir)
+        local_dir_size = get_directory_size_bytes(local_dir)
+
+        deleted = delete_package(uri, self._resources_dir)
+        if not deleted:
+            logger.warning(f"Tried to delete nonexistent URI: {uri}.")
+            return 0
+
+        return local_dir_size
+
+    def get_uris(self, runtime_env: dict) -> List[str]:
+        return runtime_env.java_jars()
+
+    async def _download_jars(
+        self, uri: str, logger: Optional[logging.Logger] = default_logger
+    ):
+        """Download a jar URI."""
+        try:
+            jar_file = await download_and_unpack_package(
+                uri, self._resources_dir, self._gcs_client, logger=logger
+            )
+        except Exception as e:
+            raise RuntimeEnvSetupError(
+                "Failed to download jar file: {}".format(e)
+            ) from e
+        module_dir = self._get_local_dir_from_uri(uri)
+        logger.debug(f"Succeeded to download jar file {jar_file} .")
+        return module_dir
+
+    async def create(
+        self,
+        uri: str,
+        runtime_env: "RuntimeEnv",  # noqa: F821
+        context: RuntimeEnvContext,
+        logger: Optional[logging.Logger] = default_logger,
+    ) -> int:
+        if not uri:
+            return 0
+        if is_jar_uri(uri):
+            module_dir = await self._download_jars(uri=uri, logger=logger)
+        else:
+            try:
+                module_dir = await download_and_unpack_package(
+                    uri, self._resources_dir, self._gcs_client, logger=logger
+                )
+            except Exception as e:
+                raise RuntimeEnvSetupError(
+                    "Failed to download jar file: {}".format(e)
+                ) from e
+
+        return get_directory_size_bytes(module_dir)
+
+    def modify_context(
+        self,
+        uris: List[str],
+        runtime_env_dict: Dict,
+        context: RuntimeEnvContext,
+        logger: Optional[logging.Logger] = default_logger,
+    ):
+        for uri in uris:
+            module_dir = self._get_local_dir_from_uri(uri)
+            if not module_dir.exists():
+                raise ValueError(
+                    f"Local directory {module_dir} for URI {uri} does "
+                    "not exist on the cluster. Something may have gone wrong while "
+                    "downloading, unpacking or installing the java jar files."
+                )
+            context.java_jars.append(str(module_dir))

python/ray/_private/runtime_env/nsight.py

@@ -0,0 +1,149 @@
+import asyncio
+import copy
+import logging
+import os
+import subprocess
+import sys
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple
+
+from ray._common.utils import (
+    try_to_create_directory,
+)
+from ray._private.runtime_env.context import RuntimeEnvContext
+from ray._private.runtime_env.plugin import RuntimeEnvPlugin
+from ray.exceptions import RuntimeEnvSetupError
+
+default_logger = logging.getLogger(__name__)
+
+# Nsight options used when runtime_env={"_nsight": "default"}
+NSIGHT_DEFAULT_CONFIG = {
+    "t": "cuda,cudnn,cublas,nvtx",
+    "o": "'worker_process_%p'",
+    "stop-on-exit": "true",
+}
+
+
+def parse_nsight_config(nsight_config: Dict[str, str]) -> List[str]:
+    """
+    Function to convert dictionary of nsight options into
+    nsight command line
+
+    The function returns:
+    - List[str]: nsys profile cmd line split into list of str
+    """
+    nsight_cmd = ["nsys", "profile"]
+    for option, option_val in nsight_config.items():
+        # option standard based on
+        # https://www.gnu.org/software/libc/manual/html_node/Argument-Syntax.html
+        if len(option) > 1:
+            nsight_cmd.append(f"--{option}={option_val}")
+        else:
+            nsight_cmd += [f"-{option}", option_val]
+    return nsight_cmd
+
+
+class NsightPlugin(RuntimeEnvPlugin):
+    name = "_nsight"
+
+    def __init__(self, resources_dir: str):
+        self.nsight_cmd = []
+
+        # replace this with better way to get logs dir
+        session_dir, runtime_dir = os.path.split(resources_dir)
+        self._nsight_dir = Path(session_dir) / "logs" / "nsight"
+        try_to_create_directory(self._nsight_dir)
+
+    async def _check_nsight_script(
+        self, nsight_config: Dict[str, str]
+    ) -> Tuple[bool, str]:
+        """
+        Function to validate if nsight_config is a valid nsight profile options
+        Args:
+            nsight_config: dictionary mapping nsight option to it's value
+        Returns:
+            a tuple consists of a boolean indicating if the nsight_config
+            is valid option and an error message if the nsight_config is invalid
+        """
+
+        # use empty as nsight report test filename
+        nsight_config_copy = copy.deepcopy(nsight_config)
+        nsight_config_copy["o"] = str(Path(self._nsight_dir) / "empty")
+        nsight_cmd = parse_nsight_config(nsight_config_copy)
+        try:
+            nsight_cmd = nsight_cmd + [sys.executable, "-c", '""']
+            process = await asyncio.create_subprocess_exec(
+                *nsight_cmd,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+            )
+            stdout, stderr = await process.communicate()
+            error_msg = stderr.strip() if stderr.strip() != "" else stdout.strip()
+
+            # cleanup test.nsys-rep file
+            clean_up_cmd = ["rm", f"{nsight_config_copy['o']}.nsys-rep"]
+            cleanup_process = await asyncio.create_subprocess_exec(
+                *clean_up_cmd,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+            )
+            _, _ = await cleanup_process.communicate()
+            if process.returncode == 0:
+                return True, None
+            else:
+                return False, error_msg
+        except FileNotFoundError:
+            return False, ("nsight is not installed")
+
+    async def create(
+        self,
+        uri: Optional[str],
+        runtime_env: "RuntimeEnv",  # noqa: F821
+        context: RuntimeEnvContext,
+        logger: logging.Logger = default_logger,
+    ) -> int:
+        nsight_config = runtime_env.nsight()
+        if not nsight_config:
+            return 0
+
+        if nsight_config and sys.platform != "linux":
+            raise RuntimeEnvSetupError(
+                "Nsight CLI is only available in Linux.\n"
+                "More information can be found in "
+                "https://docs.nvidia.com/nsight-compute/NsightComputeCli/index.html"
+            )
+
+        if isinstance(nsight_config, str):
+            if nsight_config == "default":
+                nsight_config = NSIGHT_DEFAULT_CONFIG
+            else:
+                raise RuntimeEnvSetupError(
+                    f"Unsupported nsight config: {nsight_config}. "
+                    "The supported config is 'default' or "
+                    "Dictionary of nsight options"
+                )
+
+        is_valid_nsight_cmd, error_msg = await self._check_nsight_script(nsight_config)
+        if not is_valid_nsight_cmd:
+            logger.warning(error_msg)
+            raise RuntimeEnvSetupError(
+                "nsight profile failed to run with the following "
+                f"error message:\n {error_msg}"
+            )
+        # add set output path to logs dir
+        nsight_config["o"] = str(
+            Path(self._nsight_dir) / nsight_config.get("o", NSIGHT_DEFAULT_CONFIG["o"])
+        )
+
+        self.nsight_cmd = parse_nsight_config(nsight_config)
+        return 0
+
+    def modify_context(
+        self,
+        uris: List[str],
+        runtime_env: "RuntimeEnv",  # noqa: F821
+        context: RuntimeEnvContext,
+        logger: Optional[logging.Logger] = default_logger,
+    ):
+        logger.info("Running nsight profiler")
+        context.py_executable = " ".join(self.nsight_cmd) + " python"

python/ray/_private/runtime_env/pip.py

@@ -0,0 +1,422 @@
+import asyncio
+import hashlib
+import json
+import logging
+import os
+import re
+import shutil
+import sys
+from asyncio import create_task, get_running_loop
+from typing import Dict, List, Optional
+
+from ray._common.utils import try_to_create_directory
+from ray._private.runtime_env import dependency_utils, virtualenv_utils
+from ray._private.runtime_env.packaging import Protocol, parse_uri
+from ray._private.runtime_env.plugin import RuntimeEnvPlugin
+from ray._private.runtime_env.utils import check_output_cmd
+from ray._private.utils import get_directory_size_bytes
+
+default_logger = logging.getLogger(__name__)
+
+# Matches unresolved environment variable placeholders such as
+# ${RAY_RUNTIME_ENV_CREATE_WORKING_DIR} or $VAR. Such placeholders are
+# expanded by the runtime env agent only after the driver-side hash is
+# computed, so any path containing one is not a real path on the driver
+# and must not be opened during hash computation.
+_ENV_VAR_PATTERN = re.compile(r"\$\{[^}]+\}|\$[A-Za-z_][A-Za-z0-9_]*")
+
+
+def _parse_requirements_file(file_path: str) -> List[str]:
+    packages = []
+    with open(file_path, "r", encoding="utf-8") as f:
+        for line in f:
+            # Strip whitespace and remove inline comments (preceded by a space)
+            line = line.split(" #")[0].strip()
+            if not line or line.startswith("#"):
+                continue
+            packages.append(line)
+    return packages
+
+
+def _get_pip_hash(pip_dict: Dict) -> str:
+    pip_dict_copy = pip_dict.copy()
+    # Using a list as a stack for iterative processing to handle nested requirements.
+    # Each item is a tuple (package_spec, parent_dir), where parent_dir is the directory
+    # of the file that contained this package spec (None for top-level packages)
+    packages_to_process = [
+        (pkg, None) for pkg in reversed(pip_dict_copy.get("packages", []))
+    ]
+    expanded_packages = []
+    # Track visited files using absolute paths to prevent circular references
+    visited_files = set()
+
+    while packages_to_process:
+        pkg, parent_dir = packages_to_process.pop()
+
+        file_path = None
+        if pkg.startswith("-r"):
+            file_path = pkg[2:].lstrip()
+        elif pkg.startswith("--requirement"):
+            file_path = pkg[len("--requirement") :].lstrip()
+            if file_path.startswith("="):
+                file_path = file_path[1:].lstrip()
+        else:
+            expanded_packages.append(pkg)
+            continue
+
+        if file_path is not None:
+            # If the path contains an unresolved environment variable
+            # placeholder (e.g. ${RAY_RUNTIME_ENV_CREATE_WORKING_DIR}),
+            # we cannot open it on the driver. Keep the original spec
+            # string in the hash input so the URI still reflects the user
+            # input, and let the runtime env agent expand and read the
+            # file on the worker side.
+            if _ENV_VAR_PATTERN.search(file_path):
+                expanded_packages.append(pkg)
+                continue
+
+            if parent_dir and not os.path.isabs(file_path):
+                file_path = os.path.join(parent_dir, file_path)
+
+            try:
+                abs_file_path = os.path.abspath(file_path)
+            except Exception:
+                default_logger.warning(f"Invalid path: {file_path}")
+                continue
+
+            if abs_file_path in visited_files:
+                default_logger.warning(
+                    f"Skipping circular reference to {abs_file_path}"
+                )
+                continue
+            visited_files.add(abs_file_path)
+
+            file_dir = os.path.dirname(abs_file_path)
+            packages_from_file = _parse_requirements_file(abs_file_path)
+            packages_to_process.extend(
+                [(p, file_dir) for p in reversed(packages_from_file)]
+            )
+    pip_dict_copy["packages"] = expanded_packages
+    serialized_pip_spec = json.dumps(pip_dict_copy, sort_keys=True)
+    hash_val = hashlib.sha1(serialized_pip_spec.encode("utf-8")).hexdigest()
+    return hash_val
+
+
+def get_uri(runtime_env: Dict) -> Optional[str]:
+    """Return `"pip://<hashed_dependencies>"`, or None if no GC required."""
+    pip = runtime_env.get("pip")
+    if pip is not None:
+        if isinstance(pip, dict):
+            uri = "pip://" + _get_pip_hash(pip_dict=pip)
+        elif isinstance(pip, list):
+            uri = "pip://" + _get_pip_hash(pip_dict=dict(packages=pip))
+        else:
+            raise TypeError(
+                "pip field received by RuntimeEnvAgent must be "
+                f"list or dict, not {type(pip).__name__}."
+            )
+    else:
+        uri = None
+    return uri
+
+
+class PipProcessor:
+    def __init__(
+        self,
+        target_dir: str,
+        runtime_env: "RuntimeEnv",  # noqa: F821
+        logger: Optional[logging.Logger] = default_logger,
+    ):
+        try:
+            import virtualenv  # noqa: F401 ensure virtualenv exists.
+        except ImportError:
+            raise RuntimeError(
+                f"Please install virtualenv "
+                f"`{sys.executable} -m pip install virtualenv`"
+                f"to enable pip runtime env."
+            )
+        logger.debug("Setting up pip for runtime_env: %s", runtime_env)
+        self._target_dir = target_dir
+        self._runtime_env = runtime_env
+        self._logger = logger
+
+        self._pip_config = self._runtime_env.pip_config()
+        self._pip_env = os.environ.copy()
+        self._pip_env.update(self._runtime_env.env_vars())
+
+    @classmethod
+    async def _ensure_pip_version(
+        cls,
+        path: str,
+        pip_version: Optional[str],
+        cwd: str,
+        pip_env: Dict,
+        logger: logging.Logger,
+    ):
+        """Run the pip command to reinstall pip to the specified version."""
+        if not pip_version:
+            return
+
+        python = virtualenv_utils.get_virtualenv_python(path)
+        # Ensure pip version.
+        pip_reinstall_cmd = [
+            python,
+            "-m",
+            "pip",
+            "install",
+            "--disable-pip-version-check",
+            f"pip{pip_version}",
+        ]
+        logger.info("Installing pip with version %s", pip_version)
+
+        await check_output_cmd(pip_reinstall_cmd, logger=logger, cwd=cwd, env=pip_env)
+
+    async def _pip_check(
+        self,
+        path: str,
+        pip_check: bool,
+        cwd: str,
+        pip_env: Dict,
+        logger: logging.Logger,
+    ):
+        """Run the pip check command to check python dependency conflicts.
+        If exists conflicts, the exit code of pip check command will be non-zero.
+        """
+        if not pip_check:
+            logger.info("Skip pip check.")
+            return
+        python = virtualenv_utils.get_virtualenv_python(path)
+
+        await check_output_cmd(
+            [python, "-m", "pip", "check", "--disable-pip-version-check"],
+            logger=logger,
+            cwd=cwd,
+            env=pip_env,
+        )
+
+        logger.info("Pip check on %s successfully.", path)
+
+    async def _install_pip_packages(
+        self,
+        path: str,
+        pip_packages: List[str],
+        cwd: str,
+        pip_env: Dict,
+        logger: logging.Logger,
+    ):
+        virtualenv_path = virtualenv_utils.get_virtualenv_path(path)
+        python = virtualenv_utils.get_virtualenv_python(path)
+        # TODO(fyrestone): Support -i, --no-deps, --no-cache-dir, ...
+        pip_requirements_file = dependency_utils.get_requirements_file(
+            path, pip_packages
+        )
+
+        # Avoid blocking the event loop.
+        loop = get_running_loop()
+        await loop.run_in_executor(
+            None,
+            dependency_utils.gen_requirements_txt,
+            pip_requirements_file,
+            pip_packages,
+        )
+
+        # Install all dependencies
+        # The default options for pip install are
+        #
+        # --disable-pip-version-check
+        #   Don't periodically check PyPI to determine whether a new version
+        #   of pip is available for download.
+        #
+        # --no-cache-dir
+        #   Disable the cache, the pip runtime env is a one-time installation,
+        #   and we don't need to handle the pip cache broken.
+        #
+        # Allow users to specify their own options to install packages via `pip`.
+        pip_install_cmd = [
+            python,
+            "-m",
+            "pip",
+            "install",
+            "-r",
+            pip_requirements_file,
+        ]
+
+        pip_opt_list = self._pip_config.get(
+            "pip_install_options", ["--disable-pip-version-check", "--no-cache-dir"]
+        )
+        pip_install_cmd.extend(pip_opt_list)
+
+        logger.info("Installing python requirements to %s", virtualenv_path)
+
+        await check_output_cmd(pip_install_cmd, logger=logger, cwd=cwd, env=pip_env)
+
+    async def _run(self):
+        path = self._target_dir
+        logger = self._logger
+        pip_packages = self._pip_config["packages"]
+        # We create an empty directory for exec cmd so that the cmd will
+        # run more stable. e.g. if cwd has ray, then checking ray will
+        # look up ray in cwd instead of site packages.
+        exec_cwd = os.path.join(path, "exec_cwd")
+        os.makedirs(exec_cwd, exist_ok=True)
+        try:
+            await virtualenv_utils.create_or_get_virtualenv(path, exec_cwd, logger)
+            python = virtualenv_utils.get_virtualenv_python(path)
+            async with dependency_utils.check_ray(python, exec_cwd, logger):
+                # Ensure pip version.
+                await self._ensure_pip_version(
+                    path,
+                    self._pip_config.get("pip_version", None),
+                    exec_cwd,
+                    self._pip_env,
+                    logger,
+                )
+                # Install pip packages.
+                await self._install_pip_packages(
+                    path,
+                    pip_packages,
+                    exec_cwd,
+                    self._pip_env,
+                    logger,
+                )
+                # Check python environment for conflicts.
+                await self._pip_check(
+                    path,
+                    self._pip_config.get("pip_check", False),
+                    exec_cwd,
+                    self._pip_env,
+                    logger,
+                )
+        except Exception:
+            logger.info("Delete incomplete virtualenv: %s", path)
+            shutil.rmtree(path, ignore_errors=True)
+            logger.exception("Failed to install pip packages.")
+            raise
+
+    def __await__(self):
+        return self._run().__await__()
+
+
+class PipPlugin(RuntimeEnvPlugin):
+    name = "pip"
+
+    def __init__(self, resources_dir: str):
+        self._pip_resources_dir = os.path.join(resources_dir, "pip")
+        self._creating_task = {}
+        # Maps a URI to a lock that is used to prevent multiple concurrent
+        # installs of the same virtualenv, see #24513
+        self._create_locks: Dict[str, asyncio.Lock] = {}
+        # Key: created hashes. Value: size of the pip dir.
+        self._created_hash_bytes: Dict[str, int] = {}
+        try_to_create_directory(self._pip_resources_dir)
+
+    def _get_path_from_hash(self, hash_val: str) -> str:
+        """Generate a path from the hash of a pip spec.
+
+        Example output:
+            /tmp/ray/session_2021-11-03_16-33-59_356303_41018/runtime_resources
+                /pip/ray-9a7972c3a75f55e976e620484f58410c920db091
+        """
+        return os.path.join(self._pip_resources_dir, hash_val)
+
+    def get_uris(self, runtime_env: "RuntimeEnv") -> List[str]:  # noqa: F821
+        """Return the pip URI from the RuntimeEnv if it exists, else return []."""
+        pip_uri = runtime_env.pip_uri()
+        if pip_uri:
+            return [pip_uri]
+        return []
+
+    def delete_uri(
+        self, uri: str, logger: Optional[logging.Logger] = default_logger
+    ) -> int:
+        """Delete URI and return the number of bytes deleted."""
+        logger.info("Got request to delete pip URI %s", uri)
+        protocol, hash_val = parse_uri(uri)
+        if protocol != Protocol.PIP:
+            raise ValueError(
+                "PipPlugin can only delete URIs with protocol "
+                f"pip. Received protocol {protocol}, URI {uri}"
+            )
+
+        # Cancel running create task.
+        task = self._creating_task.pop(hash_val, None)
+        if task is not None:
+            task.cancel()
+
+        del self._created_hash_bytes[hash_val]
+
+        pip_env_path = self._get_path_from_hash(hash_val)
+        local_dir_size = get_directory_size_bytes(pip_env_path)
+        del self._create_locks[uri]
+        try:
+            shutil.rmtree(pip_env_path)
+        except OSError as e:
+            logger.warning(f"Error when deleting pip env {pip_env_path}: {str(e)}")
+            return 0
+
+        return local_dir_size
+
+    async def create(
+        self,
+        uri: str,
+        runtime_env: "RuntimeEnv",  # noqa: F821
+        context: "RuntimeEnvContext",  # noqa: F821
+        logger: Optional[logging.Logger] = default_logger,
+    ) -> int:
+        if not runtime_env.has_pip():
+            return 0
+
+        protocol, hash_val = parse_uri(uri)
+        target_dir = self._get_path_from_hash(hash_val)
+
+        async def _create_for_hash():
+            await PipProcessor(
+                target_dir,
+                runtime_env,
+                logger,
+            )
+
+            loop = get_running_loop()
+            return await loop.run_in_executor(
+                None, get_directory_size_bytes, target_dir
+            )
+
+        if uri not in self._create_locks:
+            # async lock to prevent the same virtualenv being concurrently installed
+            self._create_locks[uri] = asyncio.Lock()
+
+        async with self._create_locks[uri]:
+            if hash_val in self._created_hash_bytes:
+                return self._created_hash_bytes[hash_val]
+            self._creating_task[hash_val] = task = create_task(_create_for_hash())
+            task.add_done_callback(lambda _: self._creating_task.pop(hash_val, None))
+            pip_dir_bytes = await task
+            self._created_hash_bytes[hash_val] = pip_dir_bytes
+            return pip_dir_bytes
+
+    def modify_context(
+        self,
+        uris: List[str],
+        runtime_env: "RuntimeEnv",  # noqa: F821
+        context: "RuntimeEnvContext",  # noqa: F821
+        logger: logging.Logger = default_logger,
+    ):
+        if not runtime_env.has_pip():
+            return
+        # PipPlugin only uses a single URI.
+        uri = uris[0]
+        # Update py_executable.
+        protocol, hash_val = parse_uri(uri)
+        target_dir = self._get_path_from_hash(hash_val)
+        virtualenv_python = virtualenv_utils.get_virtualenv_python(target_dir)
+
+        if not os.path.exists(virtualenv_python):
+            raise ValueError(
+                f"Local directory {target_dir} for URI {uri} does "
+                "not exist on the cluster. Something may have gone wrong while "
+                "installing the runtime_env `pip` packages."
+            )
+        context.py_executable = virtualenv_python
+        context.command_prefix += virtualenv_utils.get_virtualenv_activate_command(
+            target_dir
+        )

python/ray/_private/runtime_env/plugin.py

@@ -0,0 +1,265 @@
+import json
+import logging
+import os
+from abc import ABC
+from typing import Any, Dict, List, Optional, Type
+
+from ray._common.utils import import_attr
+from ray._private.runtime_env.constants import (
+    RAY_RUNTIME_ENV_CLASS_FIELD_NAME,
+    RAY_RUNTIME_ENV_PLUGIN_DEFAULT_PRIORITY,
+    RAY_RUNTIME_ENV_PLUGIN_MAX_PRIORITY,
+    RAY_RUNTIME_ENV_PLUGIN_MIN_PRIORITY,
+    RAY_RUNTIME_ENV_PLUGINS_ENV_VAR,
+    RAY_RUNTIME_ENV_PRIORITY_FIELD_NAME,
+)
+from ray._private.runtime_env.context import RuntimeEnvContext
+from ray._private.runtime_env.uri_cache import URICache
+from ray.util.annotations import DeveloperAPI
+
+default_logger = logging.getLogger(__name__)
+
+
+@DeveloperAPI
+class RuntimeEnvPlugin(ABC):
+    """Abstract base class for runtime environment plugins."""
+
+    name: str = None
+    priority: int = RAY_RUNTIME_ENV_PLUGIN_DEFAULT_PRIORITY
+
+    @staticmethod
+    def validate(runtime_env_dict: dict) -> None:
+        """Validate user entry for this plugin.
+
+        The method is invoked upon installation of runtime env.
+
+        Args:
+            runtime_env_dict: The user-supplied runtime environment dict.
+
+        Raises:
+            ValueError: If the validation fails.
+        """
+        pass
+
+    def get_uris(self, runtime_env: "RuntimeEnv") -> List[str]:  # noqa: F821
+        return []
+
+    async def create(
+        self,
+        uri: Optional[str],
+        runtime_env: "RuntimeEnv",  # noqa: F821
+        context: RuntimeEnvContext,
+        logger: logging.Logger,
+    ) -> float:
+        """Create and install the runtime environment.
+
+        Gets called in the runtime env agent at install time. The URI can be
+        used as a caching mechanism.
+
+        Args:
+            uri: A URI uniquely describing this resource.
+            runtime_env: The RuntimeEnv object.
+            context: Auxiliary information supplied by Ray.
+            logger: A logger to log messages during the context modification.
+
+        Returns:
+            float: The disk space taken up by this plugin installation for this
+                environment. e.g. for working_dir, this downloads the files to the
+                local node.
+        """
+        return 0
+
+    def modify_context(
+        self,
+        uris: List[str],
+        runtime_env: "RuntimeEnv",  # noqa: F821
+        context: RuntimeEnvContext,
+        logger: logging.Logger,
+    ) -> None:
+        """Modify context to change worker startup behavior.
+
+        For example, you can use this to prepend "cd <dir>" command to worker
+        startup, or add new environment variables.
+
+        Args:
+            uris: The URIs used by this resource.
+            runtime_env: The RuntimeEnv object.
+            context: Auxiliary information supplied by Ray.
+            logger: A logger to log messages during the context modification.
+        """
+        return
+
+    def delete_uri(self, uri: str, logger: logging.Logger) -> float:
+        """Delete the runtime environment given uri.
+
+        Args:
+            uri: A URI uniquely describing this resource.
+            logger: The logger used to log messages during the deletion.
+
+        Returns:
+            float: The amount of space reclaimed by the deletion.
+        """
+        return 0
+
+
+class PluginSetupContext:
+    def __init__(
+        self,
+        name: str,
+        class_instance: RuntimeEnvPlugin,
+        priority: int,
+        uri_cache: URICache,
+    ):
+        self.name = name
+        self.class_instance = class_instance
+        self.priority = priority
+        self.uri_cache = uri_cache
+
+
+class RuntimeEnvPluginManager:
+    """This manager is used to load plugins in runtime env agent."""
+
+    def __init__(self):
+        self.plugins: Dict[str, PluginSetupContext] = {}
+        plugin_config_str = os.environ.get(RAY_RUNTIME_ENV_PLUGINS_ENV_VAR)
+        if plugin_config_str:
+            plugin_configs = json.loads(plugin_config_str)
+            self.load_plugins(plugin_configs)
+
+    def validate_plugin_class(self, plugin_class: Type[RuntimeEnvPlugin]) -> None:
+        if not issubclass(plugin_class, RuntimeEnvPlugin):
+            raise RuntimeError(
+                f"Invalid runtime env plugin class {plugin_class}. "
+                "The plugin class must inherit "
+                "ray._private.runtime_env.plugin.RuntimeEnvPlugin."
+            )
+        if not plugin_class.name:
+            raise RuntimeError(f"No valid name in runtime env plugin {plugin_class}.")
+        if plugin_class.name in self.plugins:
+            raise RuntimeError(
+                f"The name of runtime env plugin {plugin_class} conflicts "
+                f"with {self.plugins[plugin_class.name]}.",
+            )
+
+    def validate_priority(self, priority: Any) -> None:
+        if (
+            not isinstance(priority, int)
+            or priority < RAY_RUNTIME_ENV_PLUGIN_MIN_PRIORITY
+            or priority > RAY_RUNTIME_ENV_PLUGIN_MAX_PRIORITY
+        ):
+            raise RuntimeError(
+                f"Invalid runtime env priority {priority}, "
+                "it should be an integer between "
+                f"{RAY_RUNTIME_ENV_PLUGIN_MIN_PRIORITY} "
+                f"and {RAY_RUNTIME_ENV_PLUGIN_MAX_PRIORITY}."
+            )
+
+    def load_plugins(self, plugin_configs: List[Dict]) -> None:
+        """Load runtime env plugins and create URI caches for them."""
+        for plugin_config in plugin_configs:
+            if (
+                not isinstance(plugin_config, dict)
+                or RAY_RUNTIME_ENV_CLASS_FIELD_NAME not in plugin_config
+            ):
+                raise RuntimeError(
+                    f"Invalid runtime env plugin config {plugin_config}, "
+                    "it should be a object which contains the "
+                    f"{RAY_RUNTIME_ENV_CLASS_FIELD_NAME} field."
+                )
+            plugin_class = import_attr(plugin_config[RAY_RUNTIME_ENV_CLASS_FIELD_NAME])
+            self.validate_plugin_class(plugin_class)
+
+            # The priority should be an integer between 0 and 100.
+            # The default priority is 10. A smaller number indicates a
+            # higher priority and the plugin will be set up first.
+            if RAY_RUNTIME_ENV_PRIORITY_FIELD_NAME in plugin_config:
+                priority = plugin_config[RAY_RUNTIME_ENV_PRIORITY_FIELD_NAME]
+            else:
+                priority = plugin_class.priority
+            self.validate_priority(priority)
+
+            class_instance = plugin_class()
+            self.plugins[plugin_class.name] = PluginSetupContext(
+                plugin_class.name,
+                class_instance,
+                priority,
+                self.create_uri_cache_for_plugin(class_instance),
+            )
+
+    def add_plugin(self, plugin: RuntimeEnvPlugin) -> None:
+        """Add a plugin to the manager and create a URI cache for it.
+
+        Args:
+            plugin: The class instance of the plugin.
+        """
+        plugin_class = type(plugin)
+        self.validate_plugin_class(plugin_class)
+        self.validate_priority(plugin_class.priority)
+        self.plugins[plugin_class.name] = PluginSetupContext(
+            plugin_class.name,
+            plugin,
+            plugin_class.priority,
+            self.create_uri_cache_for_plugin(plugin),
+        )
+
+    def create_uri_cache_for_plugin(self, plugin: RuntimeEnvPlugin) -> URICache:
+        """Create a URI cache for a plugin.
+
+        Args:
+            plugin: The plugin instance whose URIs the cache will manage.
+
+        Returns:
+            The created URI cache for the plugin.
+        """
+        # Set the max size for the cache.  Defaults to 10 GB.
+        cache_size_env_var = f"RAY_RUNTIME_ENV_{plugin.name}_CACHE_SIZE_GB".upper()
+        cache_size_bytes = int(
+            (1024**3) * float(os.environ.get(cache_size_env_var, 10))
+        )
+        return URICache(plugin.delete_uri, cache_size_bytes)
+
+    def sorted_plugin_setup_contexts(self) -> List[PluginSetupContext]:
+        """Get the sorted plugin setup contexts, sorted by increasing priority.
+
+        Returns:
+            The sorted plugin setup contexts.
+        """
+        return sorted(self.plugins.values(), key=lambda x: x.priority)
+
+
+async def create_for_plugin_if_needed(
+    runtime_env: "RuntimeEnv",  # noqa: F821
+    plugin: RuntimeEnvPlugin,
+    uri_cache: URICache,
+    context: RuntimeEnvContext,
+    logger: logging.Logger = default_logger,
+):
+    """Set up the environment using the plugin if not already set up and cached."""
+    if plugin.name not in runtime_env or runtime_env[plugin.name] is None:
+        return
+
+    plugin.validate(runtime_env)
+
+    uris = plugin.get_uris(runtime_env)
+
+    if not uris:
+        logger.debug(
+            f"No URIs for runtime env plugin {plugin.name}; "
+            "create always without checking the cache."
+        )
+        await plugin.create(None, runtime_env, context, logger=logger)
+
+    for uri in uris:
+        if uri not in uri_cache:
+            logger.debug(f"Cache miss for URI {uri}.")
+            size_bytes = await plugin.create(uri, runtime_env, context, logger=logger)
+            uri_cache.add(uri, size_bytes, logger=logger)
+        else:
+            logger.info(
+                f"Runtime env {plugin.name} {uri} is already installed "
+                "and will be reused. Search "
+                "all runtime_env_setup-*.log to find the corresponding setup log."
+            )
+            uri_cache.mark_used(uri, logger=logger)
+
+    plugin.modify_context(uris, runtime_env, context, logger)

python/ray/_private/runtime_env/plugin_schema_manager.py

@@ -0,0 +1,97 @@
+import json
+import logging
+import os
+from typing import List
+
+import jsonschema
+
+from ray._private.runtime_env.constants import (
+    RAY_RUNTIME_ENV_PLUGIN_SCHEMA_SUFFIX,
+    RAY_RUNTIME_ENV_PLUGIN_SCHEMAS_ENV_VAR,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class RuntimeEnvPluginSchemaManager:
+    """This manager is used to load plugin json schemas."""
+
+    default_schema_path = os.path.join(
+        os.path.dirname(__file__), "../../runtime_env/schemas"
+    )
+    schemas = {}
+    loaded = False
+
+    @classmethod
+    def _load_schemas(cls, schema_paths: List[str]):
+        for schema_path in schema_paths:
+            try:
+                with open(schema_path) as f:
+                    schema = json.load(f)
+            except json.decoder.JSONDecodeError:
+                logger.error("Invalid runtime env schema %s, skip it.", schema_path)
+                continue
+            except OSError:
+                logger.error("Cannot open runtime env schema %s, skip it.", schema_path)
+                continue
+            if "title" not in schema:
+                logger.error(
+                    "No valid title in runtime env schema %s, skip it.", schema_path
+                )
+                continue
+            if schema["title"] in cls.schemas:
+                logger.error(
+                    "The 'title' of runtime env schema %s conflicts with %s, skip it.",
+                    schema_path,
+                    cls.schemas[schema["title"]],
+                )
+                continue
+            cls.schemas[schema["title"]] = schema
+
+    @classmethod
+    def _load_default_schemas(cls):
+        schema_json_files = list()
+        for root, _, files in os.walk(cls.default_schema_path):
+            for f in files:
+                if f.endswith(RAY_RUNTIME_ENV_PLUGIN_SCHEMA_SUFFIX):
+                    schema_json_files.append(os.path.join(root, f))
+            logger.debug(
+                f"Loading the default runtime env schemas: {schema_json_files}."
+            )
+            cls._load_schemas(schema_json_files)
+
+    @classmethod
+    def _load_schemas_from_env_var(cls):
+        # The format of env var:
+        # "/path/to/env_1_schema.json,/path/to/env_2_schema.json,/path/to/schemas_dir/"
+        schema_paths = os.environ.get(RAY_RUNTIME_ENV_PLUGIN_SCHEMAS_ENV_VAR)
+        if schema_paths:
+            schema_json_files = list()
+            for path in schema_paths.split(","):
+                if path.endswith(RAY_RUNTIME_ENV_PLUGIN_SCHEMA_SUFFIX):
+                    schema_json_files.append(path)
+                elif os.path.isdir(path):
+                    for root, _, files in os.walk(path):
+                        for f in files:
+                            if f.endswith(RAY_RUNTIME_ENV_PLUGIN_SCHEMA_SUFFIX):
+                                schema_json_files.append(os.path.join(root, f))
+            logger.info(
+                f"Loading the runtime env schemas from env var: {schema_json_files}."
+            )
+            cls._load_schemas(schema_json_files)
+
+    @classmethod
+    def validate(cls, name, instance):
+        if not cls.loaded:
+            # Load the schemas lazily.
+            cls._load_default_schemas()
+            cls._load_schemas_from_env_var()
+            cls.loaded = True
+        # if no schema matches, skip the validation.
+        if name in cls.schemas:
+            jsonschema.validate(instance=instance, schema=cls.schemas[name])
+
+    @classmethod
+    def clear(cls):
+        cls.schemas.clear()
+        cls.loaded = False

python/ray/_private/runtime_env/protocol.py

@@ -0,0 +1,315 @@
+import enum
+import os
+from urllib.parse import urlparse
+
+RAY_RUNTIME_ENV_HTTP_USER_AGENT_ENV_VAR = "RAY_RUNTIME_ENV_HTTP_USER_AGENT"
+RAY_RUNTIME_ENV_BEARER_TOKEN_ENV_VAR = "RAY_RUNTIME_ENV_BEARER_TOKEN"
+_DEFAULT_HTTP_USER_AGENT = "ray-runtime-env-curl/1.0"
+
+
+class ProtocolsProvider:
+    _MISSING_DEPENDENCIES_WARNING = (
+        "Note that these must be preinstalled "
+        "on all nodes in the Ray cluster; it is not "
+        "sufficient to install them in the runtime_env."
+    )
+
+    @classmethod
+    def get_protocols(cls):
+        return {
+            # For packages dynamically uploaded and managed by the GCS.
+            "gcs",
+            # For conda environments installed locally on each node.
+            "conda",
+            # For pip environments installed locally on each node.
+            "pip",
+            # For uv environments install locally on each node.
+            "uv",
+            # Remote http path, assumes everything packed in one zip file.
+            "http",
+            # Remote https path, assumes everything packed in one zip file.
+            "https",
+            # Remote s3 path, assumes everything packed in one zip file.
+            "s3",
+            # Remote google storage path, assumes everything packed in one zip file.
+            "gs",
+            # Remote azure blob storage path, assumes everything packed in one zip file.
+            "azure",
+            # Remote Azure Blob File System Secure path, assumes everything packed in one zip file.
+            "abfss",
+            # File storage path, assumes everything packed in one zip file.
+            "file",
+        }
+
+    @classmethod
+    def get_remote_protocols(cls):
+        return {"http", "https", "s3", "gs", "azure", "abfss", "file"}
+
+    @classmethod
+    def _handle_s3_protocol(cls):
+        """Set up S3 protocol handling.
+
+        Returns:
+            tuple: (open_file function, transport_params)
+
+        Raises:
+            ImportError: If required dependencies are not installed.
+        """
+        try:
+            import boto3
+            from smart_open import open as open_file
+        except ImportError:
+            raise ImportError(
+                "You must `pip install smart_open[s3]` "
+                "to fetch URIs in s3 bucket. " + cls._MISSING_DEPENDENCIES_WARNING
+            )
+
+        # Create S3 client, falling back to unsigned for public buckets
+        session = boto3.Session()
+        # session.get_credentials() will return None if no credentials can be found.
+        if session.get_credentials():
+            # If credentials are found, use a standard signed client.
+            s3_client = session.client("s3")
+        else:
+            # No credentials found, fall back to an unsigned client for public buckets.
+            from botocore import UNSIGNED
+            from botocore.config import Config
+
+            s3_client = boto3.client("s3", config=Config(signature_version=UNSIGNED))
+
+        transport_params = {"client": s3_client}
+        return open_file, transport_params
+
+    @classmethod
+    def _handle_gs_protocol(cls):
+        """Set up Google Cloud Storage protocol handling.
+
+        Returns:
+            tuple: (open_file function, transport_params)
+
+        Raises:
+            ImportError: If required dependencies are not installed.
+        """
+        try:
+            from google.cloud import storage  # noqa: F401
+            from smart_open import open as open_file
+        except ImportError:
+            raise ImportError(
+                "You must `pip install smart_open[gcs]` "
+                "to fetch URIs in Google Cloud Storage bucket."
+                + cls._MISSING_DEPENDENCIES_WARNING
+            )
+
+        return open_file, None
+
+    @classmethod
+    def _handle_azure_protocol(cls):
+        """Set up Azure blob storage protocol handling.
+
+        Returns:
+            tuple: (open_file function, transport_params)
+
+        Raises:
+            ImportError: If required dependencies are not installed.
+            ValueError: If required environment variables are not set.
+        """
+        try:
+            from azure.identity import DefaultAzureCredential
+            from azure.storage.blob import BlobServiceClient  # noqa: F401
+            from smart_open import open as open_file
+        except ImportError:
+            raise ImportError(
+                "You must `pip install azure-storage-blob azure-identity smart_open[azure]` "
+                "to fetch URIs in Azure Blob Storage. "
+                + cls._MISSING_DEPENDENCIES_WARNING
+            )
+
+        # Define authentication variable
+        azure_storage_account_name = os.getenv("AZURE_STORAGE_ACCOUNT")
+
+        if not azure_storage_account_name:
+            raise ValueError(
+                "Azure Blob Storage authentication requires "
+                "AZURE_STORAGE_ACCOUNT environment variable to be set."
+            )
+
+        account_url = f"https://{azure_storage_account_name}.blob.core.windows.net/"
+        transport_params = {
+            "client": BlobServiceClient(
+                account_url=account_url, credential=DefaultAzureCredential()
+            )
+        }
+
+        return open_file, transport_params
+
+    @classmethod
+    def _handle_abfss_protocol(cls):
+        """Set up Azure Blob File System Secure (ABFSS) protocol handling.
+
+        Returns:
+            tuple: (open_file function, transport_params)
+
+        Raises:
+            ImportError: If required dependencies are not installed.
+            ValueError: If the ABFSS URI format is invalid.
+        """
+        try:
+            import adlfs
+            from azure.identity import DefaultAzureCredential
+        except ImportError:
+            raise ImportError(
+                "You must `pip install adlfs azure-identity` "
+                "to fetch URIs in Azure Blob File System Secure. "
+                + cls._MISSING_DEPENDENCIES_WARNING
+            )
+
+        def open_file(uri, mode, *, transport_params=None):
+            # Parse and validate the ABFSS URI
+            parsed = urlparse(uri)
+
+            # Validate ABFSS URI format: abfss://container@account.dfs.core.windows.net/path
+            if not parsed.netloc or "@" not in parsed.netloc:
+                raise ValueError(
+                    f"Invalid ABFSS URI format - missing container@account: {uri}"
+                )
+
+            container_part, hostname_part = parsed.netloc.split("@", 1)
+
+            # Validate container name (must be non-empty)
+            if not container_part:
+                raise ValueError(
+                    f"Invalid ABFSS URI format - empty container name: {uri}"
+                )
+
+            # Validate hostname format
+            if not hostname_part or not hostname_part.endswith(".dfs.core.windows.net"):
+                raise ValueError(
+                    f"Invalid ABFSS URI format - invalid hostname (must end with .dfs.core.windows.net): {uri}"
+                )
+
+            # Extract and validate account name
+            azure_storage_account_name = hostname_part.split(".")[0]
+            if not azure_storage_account_name:
+                raise ValueError(
+                    f"Invalid ABFSS URI format - empty account name: {uri}"
+                )
+
+            # Handle ABFSS URI with adlfs
+            filesystem = adlfs.AzureBlobFileSystem(
+                account_name=azure_storage_account_name,
+                credential=DefaultAzureCredential(),
+            )
+            return filesystem.open(uri, mode)
+
+        return open_file, None
+
+    @classmethod
+    def _http_headers(cls) -> dict:
+        headers = {
+            "User-Agent": os.environ.get(
+                RAY_RUNTIME_ENV_HTTP_USER_AGENT_ENV_VAR, _DEFAULT_HTTP_USER_AGENT
+            ),
+            "Accept": "*/*",
+        }
+
+        bearer_token = os.environ.get(RAY_RUNTIME_ENV_BEARER_TOKEN_ENV_VAR)
+        if bearer_token:
+            headers["Authorization"] = f"Bearer {bearer_token}"
+
+        return headers
+
+    @classmethod
+    def _handle_http_protocol(cls):
+        """Set up HTTP/HTTPS protocol handling with curl-like headers."""
+
+        try:
+            from smart_open import open as smart_open_open
+        except ImportError:
+            raise ImportError(
+                "You must `pip install smart_open` to fetch HTTP/HTTPS URIs. "
+                + cls._MISSING_DEPENDENCIES_WARNING
+            )
+
+        def open_file(uri, mode, *, transport_params=None):
+            params = {
+                "headers": cls._http_headers(),
+                "timeout": 60,
+            }
+            if transport_params:
+                params.update(transport_params)
+            return smart_open_open(uri, mode, transport_params=params)
+
+        return open_file, None
+
+    @classmethod
+    def download_remote_uri(cls, protocol: str, source_uri: str, dest_file: str):
+        """Download file from remote URI to destination file.
+
+        Args:
+            protocol: The protocol to use for downloading (e.g., 's3', 'https').
+            source_uri: The source URI to download from.
+            dest_file: The destination file path to save to.
+
+        Raises:
+            ImportError: If required dependencies for the protocol are not installed.
+        """
+        assert protocol in cls.get_remote_protocols()
+
+        tp = None
+        open_file = None
+
+        if protocol == "file":
+            source_uri = source_uri[len("file://") :]
+
+            def open_file(uri, mode, *, transport_params=None):
+                return open(uri, mode)
+
+        elif protocol in ("http", "https"):
+            open_file, tp = cls._handle_http_protocol()
+        elif protocol == "s3":
+            open_file, tp = cls._handle_s3_protocol()
+        elif protocol == "gs":
+            open_file, tp = cls._handle_gs_protocol()
+        elif protocol == "azure":
+            open_file, tp = cls._handle_azure_protocol()
+        elif protocol == "abfss":
+            open_file, tp = cls._handle_abfss_protocol()
+        else:
+            try:
+                from smart_open import open as open_file
+            except ImportError:
+                raise ImportError(
+                    "You must `pip install smart_open` "
+                    f"to fetch {protocol.upper()} URIs. "
+                    + cls._MISSING_DEPENDENCIES_WARNING
+                )
+
+        with open_file(source_uri, "rb", transport_params=tp) as fin:
+            with open(dest_file, "wb") as fout:
+                fout.write(fin.read())
+
+
+Protocol = enum.Enum(
+    "Protocol",
+    {protocol.upper(): protocol for protocol in ProtocolsProvider.get_protocols()},
+)
+
+
+@classmethod
+def _remote_protocols(cls):
+    # Returns a list of protocols that support remote storage
+    # These protocols should only be used with paths that end in
+    # ".zip", ".whl", ".tar.gz", or ".tgz"
+    return [
+        cls[protocol.upper()] for protocol in ProtocolsProvider.get_remote_protocols()
+    ]
+
+
+Protocol.remote_protocols = _remote_protocols
+
+
+def _download_remote_uri(self, source_uri, dest_file):
+    return ProtocolsProvider.download_remote_uri(self.value, source_uri, dest_file)
+
+
+Protocol.download_remote_uri = _download_remote_uri

python/ray/_private/runtime_env/py_executable.py

@@ -0,0 +1,45 @@
+import logging
+from typing import List, Optional
+
+from ray._private.runtime_env.context import RuntimeEnvContext
+from ray._private.runtime_env.plugin import RuntimeEnvPlugin
+
+default_logger = logging.getLogger(__name__)
+
+
+class PyExecutablePlugin(RuntimeEnvPlugin):
+    """This plugin allows running Ray workers with a custom Python executable.
+
+    You can use it with
+    `ray.init(runtime_env={"py_executable": "<command> <args>"})`. If you specify
+    a `working_dir` in the runtime environment, the executable will have access
+    to the working directory, for example, to a requirements.txt for a package manager,
+    a script for a debugger, or the executable could be a shell script in the
+    working directory. You can also use this plugin to run worker processes
+    in a custom profiler or use a custom Python interpreter or `python` with
+    custom arguments.
+    """
+
+    name = "py_executable"
+
+    def __init__(self):
+        pass
+
+    async def create(
+        self,
+        uri: Optional[str],
+        runtime_env: "RuntimeEnv",  # noqa: F821
+        context: RuntimeEnvContext,
+        logger: logging.Logger = default_logger,
+    ) -> int:
+        return 0
+
+    def modify_context(
+        self,
+        uris: List[str],
+        runtime_env: "RuntimeEnv",  # noqa: F821
+        context: RuntimeEnvContext,
+        logger: Optional[logging.Logger] = default_logger,
+    ):
+        logger.info("Running py_executable plugin")
+        context.py_executable = runtime_env.py_executable()

python/ray/_private/runtime_env/py_modules.py

@@ -0,0 +1,242 @@
+import logging
+import os
+from pathlib import Path
+from types import ModuleType
+from typing import Any, Dict, List, Optional
+
+from ray._common.utils import try_to_create_directory
+from ray._private.runtime_env.context import RuntimeEnvContext
+from ray._private.runtime_env.packaging import (
+    Protocol,
+    delete_package,
+    download_and_unpack_package,
+    get_local_dir_from_uri,
+    get_uri_for_directory,
+    get_uri_for_file,
+    get_uri_for_package,
+    install_wheel_package,
+    is_whl_uri,
+    package_exists,
+    parse_uri,
+    upload_package_if_needed,
+    upload_package_to_gcs,
+)
+from ray._private.runtime_env.plugin import RuntimeEnvPlugin
+from ray._private.runtime_env.working_dir import set_pythonpath_in_context
+from ray._private.utils import get_directory_size_bytes
+from ray._raylet import GcsClient
+from ray.exceptions import RuntimeEnvSetupError
+
+default_logger = logging.getLogger(__name__)
+
+
+def _check_is_uri(s: str) -> bool:
+    try:
+        protocol, path = parse_uri(s)
+    except ValueError:
+        protocol, path = None, None
+
+    supported_extensions = (".zip", ".whl", ".tar.gz", ".tgz")
+    if protocol in Protocol.remote_protocols() and not any(
+        path.endswith(ext) for ext in supported_extensions
+    ):
+        raise ValueError(
+            "Only .zip, .whl, .tar.gz, and .tgz files supported for remote URIs."
+        )
+
+    return protocol is not None
+
+
+def upload_py_modules_if_needed(
+    runtime_env: Dict[str, Any],
+    include_gitignore: bool,
+    scratch_dir: Optional[str] = None,
+    logger: Optional[logging.Logger] = default_logger,
+    upload_fn=None,
+) -> Dict[str, Any]:
+    """Uploads the entries in py_modules and replaces them with a list of URIs.
+
+    For each entry that is already a URI, this is a no-op.
+    """
+    py_modules = runtime_env.get("py_modules")
+    if py_modules is None:
+        return runtime_env
+
+    if not isinstance(py_modules, list):
+        raise TypeError(
+            "py_modules must be a List of local paths, imported modules, or "
+            f"URIs, got {type(py_modules)}."
+        )
+
+    py_modules_uris = []
+    for module in py_modules:
+        if isinstance(module, str):
+            # module_path is a local path or a URI.
+            module_path = module
+        elif isinstance(module, Path):
+            module_path = str(module)
+        elif isinstance(module, ModuleType):
+            if not hasattr(module, "__path__"):
+                # This is a single-file module.
+                module_path = module.__file__
+            else:
+                # NOTE(edoakes): Python allows some installed Python packages to
+                # be split into multiple directories. We could probably handle
+                # this, but it seems tricky & uncommon. If it's a problem for
+                # users, we can add this support on demand.
+                if len(module.__path__) > 1:
+                    raise ValueError(
+                        "py_modules only supports modules whose __path__"
+                        " has length 1 or those who are single-file."
+                    )
+                [module_path] = module.__path__
+        else:
+            raise TypeError(
+                "py_modules must be a list of file paths, URIs, "
+                f"or imported modules, got {type(module)}."
+            )
+
+        if _check_is_uri(module_path):
+            module_uri = module_path
+        else:
+            # module_path is a local path.
+            if Path(module_path).is_dir() or Path(module_path).suffix == ".py":
+                is_dir = Path(module_path).is_dir()
+                excludes = runtime_env.get("excludes", None)
+                if is_dir:
+                    module_uri = get_uri_for_directory(
+                        module_path,
+                        include_gitignore=include_gitignore,
+                        excludes=excludes,
+                    )
+                else:
+                    module_uri = get_uri_for_file(module_path)
+                if upload_fn is None:
+                    if scratch_dir is None:
+                        scratch_dir = os.getcwd()
+                    try:
+                        upload_package_if_needed(
+                            module_uri,
+                            scratch_dir,
+                            module_path,
+                            include_gitignore=include_gitignore,
+                            include_parent_dir=is_dir,
+                            excludes=excludes,
+                            logger=logger,
+                        )
+                    except Exception as e:
+                        from ray.util.spark.utils import is_in_databricks_runtime
+
+                        if is_in_databricks_runtime():
+                            raise RuntimeEnvSetupError(
+                                f"Failed to upload module {module_path} to the Ray "
+                                f"cluster, please ensure there are only files under "
+                                f"the module path, notebooks under the path are "
+                                f"not allowed, original exception: {e}"
+                            ) from e
+                        raise RuntimeEnvSetupError(
+                            f"Failed to upload module {module_path} to the Ray "
+                            f"cluster: {e}"
+                        ) from e
+                else:
+                    upload_fn(module_path, excludes=excludes)
+            elif Path(module_path).suffix == ".whl":
+                module_uri = get_uri_for_package(Path(module_path))
+                if upload_fn is None:
+                    if not package_exists(module_uri):
+                        try:
+                            upload_package_to_gcs(
+                                module_uri, Path(module_path).read_bytes()
+                            )
+                        except Exception as e:
+                            raise RuntimeEnvSetupError(
+                                f"Failed to upload {module_path} to the Ray "
+                                f"cluster: {e}"
+                            ) from e
+                else:
+                    upload_fn(module_path, excludes=None, is_file=True)
+            else:
+                raise ValueError(
+                    "py_modules entry must be a .py file, "
+                    "a directory, or a .whl file; "
+                    f"got {module_path}"
+                )
+
+        py_modules_uris.append(module_uri)
+
+    # TODO(architkulkarni): Expose a single URI for py_modules.  This plugin
+    # should internally handle the "sub-URIs", the individual modules.
+
+    runtime_env["py_modules"] = py_modules_uris
+    return runtime_env
+
+
+class PyModulesPlugin(RuntimeEnvPlugin):
+
+    name = "py_modules"
+
+    def __init__(self, resources_dir: str, gcs_client: GcsClient):
+        self._resources_dir = os.path.join(resources_dir, "py_modules_files")
+        self._gcs_client = gcs_client
+        try_to_create_directory(self._resources_dir)
+
+    def _get_local_dir_from_uri(self, uri: str):
+        return get_local_dir_from_uri(uri, self._resources_dir)
+
+    def delete_uri(
+        self, uri: str, logger: Optional[logging.Logger] = default_logger
+    ) -> int:
+        """Delete URI and return the number of bytes deleted."""
+        logger.info("Got request to delete pymodule URI %s", uri)
+        local_dir = get_local_dir_from_uri(uri, self._resources_dir)
+        local_dir_size = get_directory_size_bytes(local_dir)
+
+        deleted = delete_package(uri, self._resources_dir)
+        if not deleted:
+            logger.warning(f"Tried to delete nonexistent URI: {uri}.")
+            return 0
+
+        return local_dir_size
+
+    def get_uris(self, runtime_env) -> List[str]:
+        return runtime_env.py_modules()
+
+    async def create(
+        self,
+        uri: str,
+        runtime_env: "RuntimeEnv",  # noqa: F821
+        context: RuntimeEnvContext,
+        logger: Optional[logging.Logger] = default_logger,
+    ) -> int:
+
+        module_dir = await download_and_unpack_package(
+            uri, self._resources_dir, self._gcs_client, logger=logger
+        )
+
+        if is_whl_uri(uri):
+            wheel_uri = module_dir
+            module_dir = self._get_local_dir_from_uri(uri)
+            await install_wheel_package(
+                wheel_uri=wheel_uri, target_dir=module_dir, logger=logger
+            )
+
+        return get_directory_size_bytes(module_dir)
+
+    def modify_context(
+        self,
+        uris: List[str],
+        runtime_env_dict: Dict,
+        context: RuntimeEnvContext,
+        logger: Optional[logging.Logger] = default_logger,
+    ):
+        module_dirs = []
+        for uri in uris:
+            module_dir = self._get_local_dir_from_uri(uri)
+            if not module_dir.exists():
+                raise ValueError(
+                    f"Local directory {module_dir} for URI {uri} does "
+                    "not exist on the cluster. Something may have gone wrong while "
+                    "downloading, unpacking or installing the py_modules files."
+                )
+            module_dirs.append(str(module_dir))
+        set_pythonpath_in_context(os.pathsep.join(module_dirs), context)

python/ray/_private/runtime_env/rocprof_sys.py

@@ -0,0 +1,173 @@
+import asyncio
+import copy
+import logging
+import os
+import subprocess
+import sys
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple
+
+from ray._common.utils import try_to_create_directory
+from ray._private.runtime_env.context import RuntimeEnvContext
+from ray._private.runtime_env.plugin import RuntimeEnvPlugin
+from ray.exceptions import RuntimeEnvSetupError
+
+default_logger = logging.getLogger(__name__)
+
+# rocprof-sys config used when runtime_env={"_rocprof_sys": "default"}
+# Refer to the following link for more information on rocprof-sys options
+# https://rocm.docs.amd.com/projects/rocprofiler-systems/en/docs-6.4.0/how-to/understanding-rocprof-sys-output.html
+ROCPROFSYS_DEFAULT_CONFIG = {
+    "env": {
+        "ROCPROFSYS_TIME_OUTPUT": "false",
+        "ROCPROFSYS_OUTPUT_PREFIX": "worker_process_%p",
+    },
+    "args": {
+        "F": "true",
+    },
+}
+
+
+def parse_rocprof_sys_config(
+    rocprof_sys_config: Dict[str, str]
+) -> Tuple[List[str], List[str]]:
+    """
+    Function to convert dictionary of rocprof-sys options into
+    rocprof-sys-python command line
+
+    The function returns:
+    - List[str]: rocprof-sys-python cmd line split into list of str
+    """
+    rocprof_sys_cmd = ["rocprof-sys-python"]
+    rocprof_sys_env = {}
+    if "args" in rocprof_sys_config:
+        # Parse rocprof-sys arg options
+        for option, option_val in rocprof_sys_config["args"].items():
+            # option standard based on
+            # https://www.gnu.org/software/libc/manual/html_node/Argument-Syntax.html
+            if len(option) > 1:
+                rocprof_sys_cmd.append(f"--{option}={option_val}")
+            else:
+                rocprof_sys_cmd += [f"-{option}", option_val]
+    if "env" in rocprof_sys_config:
+        rocprof_sys_env = rocprof_sys_config["env"]
+    rocprof_sys_cmd.append("--")
+    return rocprof_sys_cmd, rocprof_sys_env
+
+
+class RocProfSysPlugin(RuntimeEnvPlugin):
+    name = "_rocprof_sys"
+
+    def __init__(self, resources_dir: str):
+        self.rocprof_sys_cmd = []
+        self.rocprof_sys_env = {}
+
+        # replace this with better way to get logs dir
+        session_dir, runtime_dir = os.path.split(resources_dir)
+        self._rocprof_sys_dir = Path(session_dir) / "logs" / "rocprof_sys"
+        try_to_create_directory(self._rocprof_sys_dir)
+
+    async def _check_rocprof_sys_script(
+        self, rocprof_sys_config: Dict[str, str]
+    ) -> Tuple[bool, str]:
+        """
+        Function to validate if rocprof_sys_config is a valid rocprof_sys profile options
+        Args:
+            rocprof_sys_config: dictionary mapping rocprof_sys option to it's value
+        Returns:
+            a tuple consists of a boolean indicating if the rocprof_sys_config
+            is valid option and an error message if the rocprof_sys_config is invalid
+        """
+
+        # use empty as rocprof_sys report test filename
+        test_folder = str(Path(self._rocprof_sys_dir) / "test")
+        rocprof_sys_cmd, rocprof_sys_env = parse_rocprof_sys_config(rocprof_sys_config)
+        rocprof_sys_env_copy = copy.deepcopy(rocprof_sys_env)
+        rocprof_sys_env_copy["ROCPROFSYS_OUTPUT_PATH"] = test_folder
+        rocprof_sys_env_copy.update(os.environ)
+        try_to_create_directory(test_folder)
+
+        # Create a test python file to run rocprof_sys
+        with open(f"{test_folder}/test.py", "w") as f:
+            f.write("import time\n")
+        try:
+            rocprof_sys_cmd = rocprof_sys_cmd + [f"{test_folder}/test.py"]
+            process = await asyncio.create_subprocess_exec(
+                *rocprof_sys_cmd,
+                env=rocprof_sys_env_copy,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+            )
+            stdout, stderr = await process.communicate()
+            error_msg = stderr.strip() if stderr.strip() != "" else stdout.strip()
+
+            # cleanup temp file
+            clean_up_cmd = ["rm", "-r", test_folder]
+            cleanup_process = await asyncio.create_subprocess_exec(
+                *clean_up_cmd,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+            )
+            _, _ = await cleanup_process.communicate()
+            if process.returncode == 0:
+                return True, None
+            else:
+                return False, error_msg
+        except FileNotFoundError:
+            return False, ("rocprof_sys is not installed")
+
+    async def create(
+        self,
+        uri: Optional[str],
+        runtime_env: "RuntimeEnv",  # noqa: F821
+        context: RuntimeEnvContext,
+        logger: logging.Logger = default_logger,
+    ) -> int:
+        rocprof_sys_config = runtime_env.rocprof_sys()
+        if not rocprof_sys_config:
+            return 0
+
+        if rocprof_sys_config and sys.platform != "linux":
+            raise RuntimeEnvSetupError("rocprof-sys CLI is only available in Linux.\n")
+
+        if isinstance(rocprof_sys_config, str):
+            if rocprof_sys_config == "default":
+                rocprof_sys_config = ROCPROFSYS_DEFAULT_CONFIG
+            else:
+                raise RuntimeEnvSetupError(
+                    f"Unsupported rocprof_sys config: {rocprof_sys_config}. "
+                    "The supported config is 'default' or "
+                    "Dictionary of rocprof_sys options"
+                )
+
+        is_valid_rocprof_sys_config, error_msg = await self._check_rocprof_sys_script(
+            rocprof_sys_config
+        )
+        if not is_valid_rocprof_sys_config:
+            logger.warning(error_msg)
+            raise RuntimeEnvSetupError(
+                "rocprof-sys profile failed to run with the following "
+                f"error message:\n {error_msg}"
+            )
+        # add set output path to logs dir
+        if "env" not in rocprof_sys_config:
+            rocprof_sys_config["env"] = {}
+        rocprof_sys_config["env"]["ROCPROFSYS_OUTPUT_PATH"] = str(
+            Path(self._rocprof_sys_dir)
+        )
+
+        self.rocprof_sys_cmd, self.rocprof_sys_env = parse_rocprof_sys_config(
+            rocprof_sys_config
+        )
+        return 0
+
+    def modify_context(
+        self,
+        uris: List[str],
+        runtime_env: "RuntimeEnv",  # noqa: F821
+        context: RuntimeEnvContext,
+        logger: Optional[logging.Logger] = default_logger,
+    ):
+        logger.info("Running rocprof-sys profiler")
+        context.py_executable = " ".join(self.rocprof_sys_cmd)
+        context.env_vars.update(self.rocprof_sys_env)

python/ray/_private/runtime_env/setup_hook.py

@@ -0,0 +1,261 @@
+import base64
+import logging
+import os
+import traceback
+from typing import Any, Callable, Dict, Optional, Union
+
+import ray
+import ray._private.ray_constants as ray_constants
+import ray.cloudpickle as pickle
+from ray._common.utils import load_class
+from ray._private.function_manager import build_setup_hook_export_entry
+from ray.runtime_env import RuntimeEnv
+
+logger = logging.getLogger(__name__)
+
+RUNTIME_ENV_FUNC_IDENTIFIER = "ray_runtime_env_func::"
+
+
+def get_import_export_timeout():
+    return int(
+        os.environ.get(
+            ray_constants.RAY_WORKER_PROCESS_SETUP_HOOK_LOAD_TIMEOUT_ENV_VAR, "60"
+        )
+    )
+
+
+def decode_function_key(key: bytes) -> str:
+    # b64encode only includes A-Z, a-z, 0-9, + and / characters
+    return RUNTIME_ENV_FUNC_IDENTIFIER + base64.b64encode(key).decode()
+
+
+def _encode_function_key(key: str) -> bytes:
+    assert key.startswith(RUNTIME_ENV_FUNC_IDENTIFIER)
+    return base64.b64decode(key[len(RUNTIME_ENV_FUNC_IDENTIFIER) :])
+
+
+def _raise_setup_hook_conflict(existing_hook_value: str, setup_hook_desc: str) -> None:
+    raise RuntimeError(
+        "Conflicting worker_process_setup_hook: the setup hook env "
+        f"var is already set to '{existing_hook_value}', but "
+        f"runtime_env specifies {setup_hook_desc}."
+    )
+
+
+def export_setup_func_callable(
+    runtime_env: Union[Dict[str, Any], RuntimeEnv],
+    setup_func: Callable,
+    worker: "ray.Worker",
+) -> Union[Dict[str, Any], RuntimeEnv]:
+    assert isinstance(setup_func, Callable)
+    try:
+        key = worker.function_actor_manager.export_setup_func(
+            setup_func, timeout=get_import_export_timeout()
+        )
+    except Exception as e:
+        raise ray.exceptions.RuntimeEnvSetupError(
+            "Failed to export the setup function."
+        ) from e
+    env_vars = runtime_env.get("env_vars", {})
+    assert ray_constants.WORKER_PROCESS_SETUP_HOOK_ENV_VAR not in env_vars, (
+        f"The env var, {ray_constants.WORKER_PROCESS_SETUP_HOOK_ENV_VAR}, "
+        "is not permitted because it is reserved for the internal use."
+    )
+    env_vars[ray_constants.WORKER_PROCESS_SETUP_HOOK_ENV_VAR] = decode_function_key(key)
+    runtime_env["env_vars"] = env_vars
+    # Note: This field is no-op. We don't have a plugin for the setup hook
+    # because we can implement it simply using an env var.
+    # This field is just for the observability purpose, so we store
+    # the name of the method.
+    runtime_env["worker_process_setup_hook"] = setup_func.__name__
+    return runtime_env
+
+
+def export_setup_func_module(
+    runtime_env: Union[Dict[str, Any], RuntimeEnv],
+    setup_func_module: str,
+) -> Union[Dict[str, Any], RuntimeEnv]:
+    assert isinstance(setup_func_module, str)
+    env_vars = runtime_env.get("env_vars", {})
+    assert ray_constants.WORKER_PROCESS_SETUP_HOOK_ENV_VAR not in env_vars, (
+        f"The env var, {ray_constants.WORKER_PROCESS_SETUP_HOOK_ENV_VAR}, "
+        "is not permitted because it is reserved for the internal use."
+    )
+    env_vars[ray_constants.WORKER_PROCESS_SETUP_HOOK_ENV_VAR] = setup_func_module
+    runtime_env["env_vars"] = env_vars
+    return runtime_env
+
+
+def _check_setup_hook_consistency(
+    existing_hook_value: str,
+    setup_func: Union[Callable, str],
+    worker: "ray.Worker",
+) -> None:
+    """Validate that an already-set hook env var is consistent with setup_func.
+
+    When the env var is already populated (e.g. inherited from a job supervisor),
+    we compare it against the `worker_process_setup_hook` field in the runtime_env
+    to detect silent mismatches.
+
+    Args:
+        existing_hook_value: The value of the existing hook env var.
+        setup_func: The setup function or module path.
+        worker: The worker instance.
+
+    Raises:
+        RuntimeError: If a conflict between the existing env var and setup_func is detected.
+    """
+    if isinstance(setup_func, Callable):
+        try:
+            _encode_function_key(existing_hook_value)
+        except Exception:
+            _raise_setup_hook_conflict(
+                existing_hook_value, f"callable '{setup_func.__name__}'"
+            )
+        _check_callable_hooks_match(existing_hook_value, setup_func, worker)
+    elif isinstance(setup_func, str):
+        try:
+            _encode_function_key(existing_hook_value)
+            existing_is_callable_ref = True
+        except Exception:
+            existing_is_callable_ref = False
+        if existing_is_callable_ref or existing_hook_value != setup_func:
+            _raise_setup_hook_conflict(existing_hook_value, f"'{setup_func}'")
+
+
+def _check_callable_hooks_match(
+    existing_hook_value: str,
+    setup_func: Callable,
+    worker: "ray.Worker",
+) -> None:
+    """Verify a callable produces the same GCS key as the existing env var."""
+    _, _, expected_key = build_setup_hook_export_entry(
+        setup_func, worker.current_job_id.binary()
+    )
+    expected_env_value = decode_function_key(expected_key)
+
+    if existing_hook_value != expected_env_value:
+        _raise_setup_hook_conflict(
+            existing_hook_value, f"callable '{setup_func.__name__}'"
+        )
+
+
+def upload_worker_process_setup_hook_if_needed(
+    runtime_env: Union[Dict[str, Any], RuntimeEnv],
+    worker: "ray.Worker",
+) -> Union[Dict[str, Any], RuntimeEnv]:
+    """Uploads the worker_process_setup_hook to GCS with a key.
+
+    runtime_env["worker_process_setup_hook"] is converted to a decoded key
+    that can load the worker setup hook function from GCS.
+    i.e., you can use internalKV.Get(runtime_env["worker_process_setup_hook])
+    to access the worker setup hook from GCS.
+
+    Args:
+        runtime_env: The runtime_env. The value will be modified
+            when returned.
+        worker: ray.worker instance.
+
+    Returns:
+        The modified runtime_env with the setup hook processed into an env var.
+    """
+    setup_func = runtime_env.get("worker_process_setup_hook")
+
+    if setup_func is None:
+        return runtime_env
+
+    env_vars = runtime_env.get("env_vars", {})
+    existing_hook = env_vars.get(ray_constants.WORKER_PROCESS_SETUP_HOOK_ENV_VAR)
+
+    if existing_hook is not None:
+        # A setup hook is already populated (e.g. inherited from job supervisor).
+        # Validate that it is consistent with the current worker_process_setup_hook.
+        _check_setup_hook_consistency(existing_hook, setup_func, worker)
+        return runtime_env
+
+    if isinstance(setup_func, Callable):
+        return export_setup_func_callable(runtime_env, setup_func, worker)
+    elif isinstance(setup_func, str):
+        return export_setup_func_module(runtime_env, setup_func)
+    else:
+        raise TypeError(
+            "worker_process_setup_hook must be a function, " f"got {type(setup_func)}."
+        )
+
+
+def load_and_execute_setup_hook(
+    worker_process_setup_hook_key: str,
+) -> Optional[str]:
+    """Load the setup hook from a given key and execute.
+
+    Args:
+        worker_process_setup_hook_key: The key to import the setup hook
+            from GCS.
+    Returns:
+        An error message if it fails. None if it succeeds.
+    """
+    assert worker_process_setup_hook_key is not None
+    if not worker_process_setup_hook_key.startswith(RUNTIME_ENV_FUNC_IDENTIFIER):
+        return load_and_execute_setup_hook_module(worker_process_setup_hook_key)
+    else:
+        return load_and_execute_setup_hook_func(worker_process_setup_hook_key)
+
+
+def load_and_execute_setup_hook_module(
+    worker_process_setup_hook_key: str,
+) -> Optional[str]:
+    try:
+        setup_func = load_class(worker_process_setup_hook_key)
+        setup_func()
+        return None
+    except Exception:
+        error_message = (
+            "Failed to execute the setup hook method, "
+            f"{worker_process_setup_hook_key} "
+            "from ``ray.init(runtime_env="
+            f"{{'worker_process_setup_hook': {worker_process_setup_hook_key}}})``. "
+            "Please make sure the given module exists and is available "
+            "from ray workers. For more details, see the error trace below.\n"
+            f"{traceback.format_exc()}"
+        )
+        return error_message
+
+
+def load_and_execute_setup_hook_func(
+    worker_process_setup_hook_key: str,
+) -> Optional[str]:
+    worker = ray._private.worker.global_worker
+    assert worker.connected
+    func_manager = worker.function_actor_manager
+    try:
+        worker_setup_func_info = func_manager.fetch_registered_method(
+            _encode_function_key(worker_process_setup_hook_key),
+            timeout=get_import_export_timeout(),
+        )
+    except Exception:
+        error_message = (
+            "Failed to import setup hook within "
+            f"{get_import_export_timeout()} seconds.\n"
+            f"{traceback.format_exc()}"
+        )
+        return error_message
+
+    try:
+        setup_func = pickle.loads(worker_setup_func_info.function)
+    except Exception:
+        error_message = (
+            "Failed to deserialize the setup hook method.\n" f"{traceback.format_exc()}"
+        )
+        return error_message
+
+    try:
+        setup_func()
+    except Exception:
+        error_message = (
+            f"Failed to execute the setup hook method. Function name:"
+            f"{worker_setup_func_info.function_name}\n"
+            f"{traceback.format_exc()}"
+        )
+        return error_message
+
+    return None

python/ray/_private/runtime_env/uri_cache.py

@@ -0,0 +1,115 @@
+import logging
+from typing import Callable, Optional, Set
+
+default_logger = logging.getLogger(__name__)
+
+DEFAULT_MAX_URI_CACHE_SIZE_BYTES = (1024**3) * 10  # 10 GB
+
+
+class URICache:
+    """Caches URIs up to a specified total size limit.
+
+    URIs are represented by strings. Each URI has an associated size on disk.
+    When a URI is added to the URICache, it is marked as "in use".
+    When a URI is no longer in use, the user of this class should call
+    `mark_unused` to signal that the URI is safe for deletion.
+
+    URIs in the cache can be marked as "in use" by calling `mark_used`.
+
+    Deletion of URIs on disk does not occur until the size limit is exceeded.
+    When this happens, URIs that are not in use are deleted randomly until the
+    size limit is satisfied, or there are no more URIs that are not in use.
+
+    It is possible for the total size on disk to exceed the size limit if all
+    the URIs are in use.
+
+    """
+
+    def __init__(
+        self,
+        delete_fn: Optional[Callable[[str, logging.Logger], int]] = None,
+        max_total_size_bytes: int = DEFAULT_MAX_URI_CACHE_SIZE_BYTES,
+        debug_mode: bool = False,
+    ):
+        # Maps URIs to the size in bytes of their corresponding disk contents.
+        self._used_uris: Set[str] = set()
+        self._unused_uris: Set[str] = set()
+
+        if delete_fn is None:
+            self._delete_fn = lambda uri, logger: 0
+        else:
+            self._delete_fn = delete_fn
+
+        # Total size of both used and unused URIs in the cache.
+        self._total_size_bytes = 0
+        self.max_total_size_bytes = max_total_size_bytes
+
+        # Used in `self._check_valid()` for testing.
+        self._debug_mode = debug_mode
+
+    def mark_unused(self, uri: str, logger: logging.Logger = default_logger):
+        """Mark a URI as unused and okay to be deleted."""
+        if uri not in self._used_uris:
+            logger.info(f"URI {uri} is already unused.")
+        else:
+            self._unused_uris.add(uri)
+            self._used_uris.remove(uri)
+        logger.info(f"Marked URI {uri} unused.")
+        self._evict_if_needed(logger)
+        self._check_valid()
+
+    def mark_used(self, uri: str, logger: logging.Logger = default_logger):
+        """Mark a URI as in use. URIs in use will not be deleted."""
+        if uri in self._used_uris:
+            return
+        elif uri in self._unused_uris:
+            self._used_uris.add(uri)
+            self._unused_uris.remove(uri)
+        else:
+            raise ValueError(
+                f"Got request to mark URI {uri} used, but this "
+                "URI is not present in the cache."
+            )
+        logger.info(f"Marked URI {uri} used.")
+        self._check_valid()
+
+    def add(self, uri: str, size_bytes: int, logger: logging.Logger = default_logger):
+        """Add a URI to the cache and mark it as in use."""
+        if uri in self._unused_uris:
+            self._unused_uris.remove(uri)
+
+        self._used_uris.add(uri)
+        self._total_size_bytes += size_bytes
+
+        self._evict_if_needed(logger)
+        self._check_valid()
+        logger.info(f"Added URI {uri} with size {size_bytes}")
+
+    def get_total_size_bytes(self) -> int:
+        return self._total_size_bytes
+
+    def _evict_if_needed(self, logger: logging.Logger = default_logger):
+        """Evict unused URIs (if they exist) until total size <= max size."""
+        while (
+            self._unused_uris
+            and self.get_total_size_bytes() > self.max_total_size_bytes
+        ):
+            # TODO(architkulkarni): Evict least recently used URI instead
+            arbitrary_unused_uri = next(iter(self._unused_uris))
+            self._unused_uris.remove(arbitrary_unused_uri)
+            num_bytes_deleted = self._delete_fn(arbitrary_unused_uri, logger)
+            self._total_size_bytes -= num_bytes_deleted
+            logger.info(
+                f"Deleted URI {arbitrary_unused_uri} with size " f"{num_bytes_deleted}."
+            )
+
+    def _check_valid(self):
+        """(Debug mode only) Check "used" and "unused" sets are disjoint."""
+        if self._debug_mode:
+            assert self._used_uris & self._unused_uris == set()
+
+    def __contains__(self, uri):
+        return uri in self._used_uris or uri in self._unused_uris
+
+    def __repr__(self):
+        return str(self.__dict__)

python/ray/_private/runtime_env/utils.py

@@ -0,0 +1,117 @@
+import asyncio
+import itertools
+import logging
+import subprocess
+import textwrap
+import types
+from typing import List
+
+
+class SubprocessCalledProcessError(subprocess.CalledProcessError):
+    """The subprocess.CalledProcessError with stripped stdout."""
+
+    LAST_N_LINES = 50
+
+    def __init__(self, *args, cmd_index=None, **kwargs):
+        self.cmd_index = cmd_index
+        super().__init__(*args, **kwargs)
+
+    @staticmethod
+    def _get_last_n_line(str_data: str, last_n_lines: int) -> str:
+        if last_n_lines < 0:
+            return str_data
+        lines = str_data.strip().split("\n")
+        return "\n".join(lines[-last_n_lines:])
+
+    def __str__(self):
+        str_list = (
+            []
+            if self.cmd_index is None
+            else [f"Run cmd[{self.cmd_index}] failed with the following details."]
+        )
+        str_list.append(super().__str__())
+        out = {
+            "stdout": self.stdout,
+            "stderr": self.stderr,
+        }
+        for name, s in out.items():
+            if s:
+                subtitle = f"Last {self.LAST_N_LINES} lines of {name}:"
+                last_n_line_str = self._get_last_n_line(s, self.LAST_N_LINES).strip()
+                str_list.append(
+                    f"{subtitle}\n{textwrap.indent(last_n_line_str, ' ' * 4)}"
+                )
+        return "\n".join(str_list)
+
+
+async def check_output_cmd(
+    cmd: List[str],
+    *,
+    logger: logging.Logger,
+    cmd_index_gen: types.GeneratorType = itertools.count(1),
+    **kwargs,
+) -> str:
+    """Run command with arguments and return its output.
+
+    If the return code was non-zero it raises a CalledProcessError. The
+    CalledProcessError object will have the return code in the returncode
+    attribute and any output in the output attribute.
+
+    Args:
+        cmd: The cmdline should be a sequence of program arguments or else
+            a single string or path-like object. The program to execute is
+            the first item in cmd.
+        logger: The logger instance.
+        cmd_index_gen: The cmd index generator, default is itertools.count(1).
+        **kwargs: All arguments are passed to the create_subprocess_exec.
+
+    Returns:
+        The stdout of cmd.
+
+    Raises:
+        CalledProcessError: If the return code of cmd is not 0.
+    """
+
+    cmd_index = next(cmd_index_gen)
+    logger.info("Run cmd[%s] %s", cmd_index, repr(cmd))
+
+    proc = None
+    try:
+        proc = await asyncio.create_subprocess_exec(
+            *cmd,
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.STDOUT,
+            **kwargs,
+        )
+        # Use communicate instead of polling stdout:
+        #   * Avoid deadlocks due to streams pausing reading or writing and blocking the
+        #     child process. Please refer to:
+        #     https://docs.python.org/3/library/asyncio-subprocess.html#asyncio.asyncio.subprocess.Process.stderr
+        #   * Avoid mixing multiple outputs of concurrent cmds.
+        stdout, _ = await proc.communicate()
+    except asyncio.exceptions.CancelledError as e:
+        # since Python 3.9, when cancelled, the inner process needs to throw as it is
+        # for asyncio to timeout properly https://bugs.python.org/issue40607
+        raise e
+    except BaseException as e:
+        raise RuntimeError(f"Run cmd[{cmd_index}] got exception.") from e
+    else:
+        stdout = stdout.decode("utf-8")
+        if stdout:
+            logger.info("Output of cmd[%s]: %s", cmd_index, stdout)
+        else:
+            logger.info("No output for cmd[%s]", cmd_index)
+        if proc.returncode != 0:
+            raise SubprocessCalledProcessError(
+                proc.returncode, cmd, output=stdout, cmd_index=cmd_index
+            )
+        return stdout
+    finally:
+        if proc is not None:
+            # Kill process.
+            try:
+                proc.kill()
+            except ProcessLookupError:
+                pass
+            # Wait process exit.
+            await proc.wait()

python/ray/_private/runtime_env/uv.py

@@ -0,0 +1,344 @@
+"""Util class to install packages via uv."""
+
+import asyncio
+import hashlib
+import json
+import logging
+import os
+import shutil
+import sys
+from asyncio import create_task, get_running_loop
+from typing import Dict, List, Optional
+
+from ray._common.utils import try_to_create_directory
+from ray._private.runtime_env import dependency_utils, virtualenv_utils
+from ray._private.runtime_env.packaging import Protocol, parse_uri
+from ray._private.runtime_env.plugin import RuntimeEnvPlugin
+from ray._private.runtime_env.utils import check_output_cmd
+from ray._private.utils import get_directory_size_bytes
+
+default_logger = logging.getLogger(__name__)
+
+
+def _get_uv_hash(uv_dict: Dict) -> str:
+    """Get a deterministic hash value for `uv` related runtime envs."""
+    serialized_uv_spec = json.dumps(uv_dict, sort_keys=True)
+    hash_val = hashlib.sha1(serialized_uv_spec.encode("utf-8")).hexdigest()
+    return hash_val
+
+
+def get_uri(runtime_env: Dict) -> Optional[str]:
+    """Return `"uv://<hashed_dependencies>"`, or None if no GC required."""
+    uv = runtime_env.get("uv")
+    if uv is not None:
+        if isinstance(uv, dict):
+            uri = "uv://" + _get_uv_hash(uv_dict=uv)
+        elif isinstance(uv, list):
+            uri = "uv://" + _get_uv_hash(uv_dict=dict(packages=uv))
+        else:
+            raise TypeError(
+                "uv field received by RuntimeEnvAgent must be "
+                f"list or dict, not {type(uv).__name__}."
+            )
+    else:
+        uri = None
+    return uri
+
+
+class UvProcessor:
+    def __init__(
+        self,
+        target_dir: str,
+        runtime_env: "RuntimeEnv",  # noqa: F821
+        logger: Optional[logging.Logger] = default_logger,
+    ):
+        try:
+            import virtualenv  # noqa: F401 ensure virtualenv exists.
+        except ImportError:
+            raise RuntimeError(
+                f"Please install virtualenv "
+                f"`{sys.executable} -m pip install virtualenv`"
+                f"to enable uv runtime env."
+            )
+
+        logger.debug("Setting up uv for runtime_env: %s", runtime_env)
+        self._target_dir = target_dir
+        # An empty directory is created to execute cmd.
+        self._exec_cwd = os.path.join(self._target_dir, "exec_cwd")
+        self._runtime_env = runtime_env
+        self._logger = logger
+
+        self._uv_config = self._runtime_env.uv_config()
+        self._uv_env = os.environ.copy()
+        self._uv_env.update(self._runtime_env.env_vars())
+
+    async def _install_uv(
+        self, path: str, cwd: str, pip_env: dict, logger: logging.Logger
+    ):
+        """Before package install, make sure the required version `uv` (if specifieds)
+        is installed.
+        """
+        virtualenv_path = virtualenv_utils.get_virtualenv_path(path)
+        python = virtualenv_utils.get_virtualenv_python(path)
+
+        def _get_uv_exec_to_install() -> str:
+            """Get `uv` executable with version to install."""
+            uv_version = self._uv_config.get("uv_version", None)
+            if uv_version:
+                return f"uv{uv_version}"
+            # Use default version.
+            return "uv"
+
+        uv_install_cmd = [
+            python,
+            "-m",
+            "pip",
+            "install",
+            "--disable-pip-version-check",
+            "--no-cache-dir",
+            _get_uv_exec_to_install(),
+        ]
+        logger.info("Installing package uv to %s", virtualenv_path)
+        await check_output_cmd(uv_install_cmd, logger=logger, cwd=cwd, env=pip_env)
+
+    async def _check_uv_existence(
+        self, path: str, cwd: str, env: dict, logger: logging.Logger
+    ) -> bool:
+        """Check and return the existence of `uv` in virtual env."""
+        python = virtualenv_utils.get_virtualenv_python(path)
+
+        check_existence_cmd = [
+            python,
+            "-m",
+            "uv",
+            "--version",
+        ]
+
+        try:
+            # If `uv` doesn't exist, exception will be thrown.
+            await check_output_cmd(check_existence_cmd, logger=logger, cwd=cwd, env=env)
+            return True
+        except Exception:
+            return False
+
+    async def _uv_check(sef, python: str, cwd: str, logger: logging.Logger) -> None:
+        """Check virtual env dependency compatibility.
+        If any incompatibility detected, exception will be thrown.
+
+        param:
+            python: the path for python executable within virtual environment.
+        """
+        cmd = [python, "-m", "uv", "pip", "check"]
+        await check_output_cmd(
+            cmd,
+            logger=logger,
+            cwd=cwd,
+        )
+
+    async def _install_uv_packages(
+        self,
+        path: str,
+        uv_packages: List[str],
+        cwd: str,
+        pip_env: Dict,
+        logger: logging.Logger,
+    ):
+        """Install required python packages via `uv`."""
+        virtualenv_path = virtualenv_utils.get_virtualenv_path(path)
+        python = virtualenv_utils.get_virtualenv_python(path)
+        # TODO(fyrestone): Support -i, --no-deps, --no-cache-dir, ...
+        requirements_file = dependency_utils.get_requirements_file(path, uv_packages)
+
+        # Check existence for `uv` and see if we could skip `uv` installation.
+        uv_exists = await self._check_uv_existence(path, cwd, pip_env, logger)
+
+        # Install uv, which acts as the default package manager.
+        if (not uv_exists) or (self._uv_config.get("uv_version", None) is not None):
+            await self._install_uv(path, cwd, pip_env, logger)
+
+        # Avoid blocking the event loop.
+        loop = get_running_loop()
+        await loop.run_in_executor(
+            None, dependency_utils.gen_requirements_txt, requirements_file, uv_packages
+        )
+
+        # Install all dependencies.
+        #
+        # Difference with pip:
+        # 1. `--disable-pip-version-check` has no effect for uv.
+        uv_install_cmd = [
+            python,
+            "-m",
+            "uv",
+            "pip",
+            "install",
+            "-r",
+            requirements_file,
+        ]
+
+        uv_opt_list = self._uv_config.get("uv_pip_install_options", ["--no-cache"])
+        if uv_opt_list:
+            uv_install_cmd += uv_opt_list
+
+        logger.info("Installing python requirements to %s", virtualenv_path)
+        await check_output_cmd(uv_install_cmd, logger=logger, cwd=cwd, env=pip_env)
+
+        # Check python environment for conflicts.
+        if self._uv_config.get("uv_check", False):
+            await self._uv_check(python, cwd, logger)
+
+    async def _run(self):
+        path = self._target_dir
+        logger = self._logger
+        uv_packages = self._uv_config["packages"]
+        # We create an empty directory for exec cmd so that the cmd will
+        # run more stable. e.g. if cwd has ray, then checking ray will
+        # look up ray in cwd instead of site packages.
+        os.makedirs(self._exec_cwd, exist_ok=True)
+        try:
+            await virtualenv_utils.create_or_get_virtualenv(
+                path, self._exec_cwd, logger
+            )
+            python = virtualenv_utils.get_virtualenv_python(path)
+            async with dependency_utils.check_ray(python, self._exec_cwd, logger):
+                # Install packages with uv.
+                await self._install_uv_packages(
+                    path,
+                    uv_packages,
+                    self._exec_cwd,
+                    self._uv_env,
+                    logger,
+                )
+        except Exception:
+            logger.info("Delete incomplete virtualenv: %s", path)
+            shutil.rmtree(path, ignore_errors=True)
+            logger.exception("Failed to install uv packages.")
+            raise
+
+    def __await__(self):
+        return self._run().__await__()
+
+
+class UvPlugin(RuntimeEnvPlugin):
+    name = "uv"
+
+    def __init__(self, resources_dir: str):
+        self._uv_resource_dir = os.path.join(resources_dir, "uv")
+        self._creating_task = {}
+        # Maps a URI to a lock that is used to prevent multiple concurrent
+        # installs of the same virtualenv, see #24513
+        self._create_locks: Dict[str, asyncio.Lock] = {}
+        # Key: created hashes. Value: size of the uv dir.
+        self._created_hash_bytes: Dict[str, int] = {}
+        try_to_create_directory(self._uv_resource_dir)
+
+    def _get_path_from_hash(self, hash_val: str) -> str:
+        """Generate a path from the hash of a uv spec.
+
+        Example output:
+            /tmp/ray/session_2021-11-03_16-33-59_356303_41018/runtime_resources
+                /uv/ray-9a7972c3a75f55e976e620484f58410c920db091
+        """
+        return os.path.join(self._uv_resource_dir, hash_val)
+
+    def get_uris(self, runtime_env: "RuntimeEnv") -> List[str]:  # noqa: F821
+        """Return the uv URI from the RuntimeEnv if it exists, else return []."""
+        uv_uri = runtime_env.uv_uri()
+        if uv_uri:
+            return [uv_uri]
+        return []
+
+    def delete_uri(
+        self, uri: str, logger: Optional[logging.Logger] = default_logger
+    ) -> int:
+        """Delete URI and return the number of bytes deleted."""
+        logger.info("Got request to delete uv URI %s", uri)
+        protocol, hash_val = parse_uri(uri)
+        if protocol != Protocol.UV:
+            raise ValueError(
+                "UvPlugin can only delete URIs with protocol "
+                f"uv. Received protocol {protocol}, URI {uri}"
+            )
+
+        # Cancel running create task.
+        task = self._creating_task.pop(hash_val, None)
+        if task is not None:
+            task.cancel()
+
+        del self._created_hash_bytes[hash_val]
+
+        uv_env_path = self._get_path_from_hash(hash_val)
+        local_dir_size = get_directory_size_bytes(uv_env_path)
+        del self._create_locks[uri]
+        try:
+            shutil.rmtree(uv_env_path)
+        except OSError as e:
+            logger.warning(f"Error when deleting uv env {uv_env_path}: {str(e)}")
+            return 0
+
+        return local_dir_size
+
+    async def create(
+        self,
+        uri: str,
+        runtime_env: "RuntimeEnv",  # noqa: F821
+        context: "RuntimeEnvContext",  # noqa: F821
+        logger: Optional[logging.Logger] = default_logger,
+    ) -> int:
+        if not runtime_env.has_uv():
+            return 0
+
+        protocol, hash_val = parse_uri(uri)
+        target_dir = self._get_path_from_hash(hash_val)
+
+        async def _create_for_hash():
+            await UvProcessor(
+                target_dir,
+                runtime_env,
+                logger,
+            )
+
+            loop = get_running_loop()
+            return await loop.run_in_executor(
+                None, get_directory_size_bytes, target_dir
+            )
+
+        if uri not in self._create_locks:
+            # async lock to prevent the same virtualenv being concurrently installed
+            self._create_locks[uri] = asyncio.Lock()
+
+        async with self._create_locks[uri]:
+            if hash_val in self._created_hash_bytes:
+                return self._created_hash_bytes[hash_val]
+            self._creating_task[hash_val] = task = create_task(_create_for_hash())
+            task.add_done_callback(lambda _: self._creating_task.pop(hash_val, None))
+            uv_dir_bytes = await task
+            self._created_hash_bytes[hash_val] = uv_dir_bytes
+            return uv_dir_bytes
+
+    def modify_context(
+        self,
+        uris: List[str],
+        runtime_env: "RuntimeEnv",  # noqa: F821
+        context: "RuntimeEnvContext",  # noqa: F821
+        logger: logging.Logger = default_logger,
+    ):
+        if not runtime_env.has_uv():
+            return
+        # UvPlugin only uses a single URI.
+        uri = uris[0]
+        # Update py_executable.
+        protocol, hash_val = parse_uri(uri)
+        target_dir = self._get_path_from_hash(hash_val)
+        virtualenv_python = virtualenv_utils.get_virtualenv_python(target_dir)
+
+        if not os.path.exists(virtualenv_python):
+            raise ValueError(
+                f"Local directory {target_dir} for URI {uri} does "
+                "not exist on the cluster. Something may have gone wrong while "
+                "installing the runtime_env `uv` packages."
+            )
+        context.py_executable = virtualenv_python
+        context.command_prefix += virtualenv_utils.get_virtualenv_activate_command(
+            target_dir
+        )

python/ray/_private/runtime_env/uv_runtime_env_hook.py

@@ -0,0 +1,452 @@
+import argparse
+import copy
+import optparse
+import os
+import pathlib
+import platform
+import sys
+import urllib.parse
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Tuple
+
+import psutil
+
+
+def _is_path(path_or_uri: str) -> bool:
+    """Returns True if uri_or_path is a path and False otherwise.
+
+    IMPORTANT: This is a duplicate of ray._private.path_utils.is_path().
+
+    Why we can't import from path_utils:
+    - This hook runs via `uv run --no-project uv_runtime_env_hook.py` in test scenarios
+    - UV creates a minimal environment without dependencies installed yet
+    - Importing from ray._private.path_utils triggers the full Ray import chain:
+      ray._private.path_utils → ray/__init__.py → ray._private.worker →
+      ray.widgets → ray.widgets.util → packaging.version
+    - The 'packaging' module is not available in the minimal UV environment,
+      causing: ModuleNotFoundError: No module named 'packaging.version'
+
+    This duplicate implementation uses only stdlib (pathlib, urllib.parse)
+    to avoid the dependency issue. If you modify this function, ensure you
+    also update ray._private.path_utils.is_path() to keep them in sync.
+    """
+    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__}")
+
+
+def _create_uv_run_parser():
+    """Create and return the argument parser for 'uv run' command."""
+
+    parser = optparse.OptionParser(prog="uv run", add_help_option=False)
+
+    # Disable interspersed args to stop parsing when we hit the first
+    # argument that is not recognized by the parser.
+    parser.disable_interspersed_args()
+
+    # Main options group
+    main_group = optparse.OptionGroup(parser, "Main options")
+    main_group.add_option("--extra", action="append", dest="extras")
+    main_group.add_option("--all-extras", action="store_true")
+    main_group.add_option("--no-extra", action="append", dest="no_extras")
+    main_group.add_option("--no-dev", action="store_true")
+    main_group.add_option("--group", action="append", dest="groups")
+    main_group.add_option("--no-group", action="append", dest="no_groups")
+    main_group.add_option("--no-default-groups", action="store_true")
+    main_group.add_option("--only-group", action="append", dest="only_groups")
+    main_group.add_option("--all-groups", action="store_true")
+    main_group.add_option("-m", "--module")
+    main_group.add_option("--only-dev", action="store_true")
+    main_group.add_option("--no-editable", action="store_true")
+    main_group.add_option("--exact", action="store_true")
+    main_group.add_option("--env-file", action="append", dest="env_files")
+    main_group.add_option("--no-env-file", action="store_true")
+    parser.add_option_group(main_group)
+
+    # With options
+    with_group = optparse.OptionGroup(parser, "With options")
+    with_group.add_option("--with", action="append", dest="with_packages")
+    with_group.add_option("--with-editable", action="append", dest="with_editable")
+    with_group.add_option(
+        "--with-requirements", action="append", dest="with_requirements"
+    )
+    parser.add_option_group(with_group)
+
+    # Environment options
+    env_group = optparse.OptionGroup(parser, "Environment options")
+    env_group.add_option("--isolated", action="store_true")
+    env_group.add_option("--active", action="store_true")
+    env_group.add_option("--no-sync", action="store_true")
+    env_group.add_option("--locked", action="store_true")
+    env_group.add_option("--frozen", action="store_true")
+    parser.add_option_group(env_group)
+
+    # Script options
+    script_group = optparse.OptionGroup(parser, "Script options")
+    script_group.add_option("-s", "--script", action="store_true")
+    script_group.add_option("--gui-script", action="store_true")
+    parser.add_option_group(script_group)
+
+    # Workspace options
+    workspace_group = optparse.OptionGroup(parser, "Workspace options")
+    workspace_group.add_option("--all-packages", action="store_true")
+    workspace_group.add_option("--package")
+    workspace_group.add_option("--no-project", action="store_true")
+    parser.add_option_group(workspace_group)
+
+    # Index options
+    index_group = optparse.OptionGroup(parser, "Index options")
+    index_group.add_option("--index", action="append", dest="indexes")
+    index_group.add_option("--default-index")
+    index_group.add_option("-i", "--index-url")
+    index_group.add_option(
+        "--extra-index-url", action="append", dest="extra_index_urls"
+    )
+    index_group.add_option("-f", "--find-links", action="append", dest="find_links")
+    index_group.add_option("--no-index", action="store_true")
+    index_group.add_option(
+        "--index-strategy",
+        type="choice",
+        choices=["first-index", "unsafe-first-match", "unsafe-best-match"],
+    )
+    index_group.add_option(
+        "--keyring-provider", type="choice", choices=["disabled", "subprocess"]
+    )
+    parser.add_option_group(index_group)
+
+    # Resolver options
+    resolver_group = optparse.OptionGroup(parser, "Resolver options")
+    resolver_group.add_option("-U", "--upgrade", action="store_true")
+    resolver_group.add_option(
+        "-P", "--upgrade-package", action="append", dest="upgrade_packages"
+    )
+    resolver_group.add_option(
+        "--resolution", type="choice", choices=["highest", "lowest", "lowest-direct"]
+    )
+    resolver_group.add_option(
+        "--prerelease",
+        type="choice",
+        choices=[
+            "disallow",
+            "allow",
+            "if-necessary",
+            "explicit",
+            "if-necessary-or-explicit",
+        ],
+    )
+    resolver_group.add_option(
+        "--fork-strategy", type="choice", choices=["fewest", "requires-python"]
+    )
+    resolver_group.add_option("--exclude-newer")
+    resolver_group.add_option("--no-sources", action="store_true")
+    parser.add_option_group(resolver_group)
+
+    # Installer options
+    installer_group = optparse.OptionGroup(parser, "Installer options")
+    installer_group.add_option("--reinstall", action="store_true")
+    installer_group.add_option(
+        "--reinstall-package", action="append", dest="reinstall_packages"
+    )
+    installer_group.add_option(
+        "--link-mode", type="choice", choices=["clone", "copy", "hardlink", "symlink"]
+    )
+    installer_group.add_option("--compile-bytecode", action="store_true")
+    parser.add_option_group(installer_group)
+
+    # Build options
+    build_group = optparse.OptionGroup(parser, "Build options")
+    build_group.add_option(
+        "-C", "--config-setting", action="append", dest="config_settings"
+    )
+    build_group.add_option("--no-build-isolation", action="store_true")
+    build_group.add_option(
+        "--no-build-isolation-package",
+        action="append",
+        dest="no_build_isolation_packages",
+    )
+    build_group.add_option("--no-build", action="store_true")
+    build_group.add_option(
+        "--no-build-package", action="append", dest="no_build_packages"
+    )
+    build_group.add_option("--no-binary", action="store_true")
+    build_group.add_option(
+        "--no-binary-package", action="append", dest="no_binary_packages"
+    )
+    parser.add_option_group(build_group)
+
+    # Cache options
+    cache_group = optparse.OptionGroup(parser, "Cache options")
+    cache_group.add_option("-n", "--no-cache", action="store_true")
+    cache_group.add_option("--cache-dir")
+    cache_group.add_option("--refresh", action="store_true")
+    cache_group.add_option(
+        "--refresh-package", action="append", dest="refresh_packages"
+    )
+    parser.add_option_group(cache_group)
+
+    # Python options
+    python_group = optparse.OptionGroup(parser, "Python options")
+    python_group.add_option("-p", "--python")
+    python_group.add_option("--managed-python", action="store_true")
+    python_group.add_option("--no-managed-python", action="store_true")
+    python_group.add_option("--no-python-downloads", action="store_true")
+    # note: the following is a legacy option and will be removed at some point
+    # https://github.com/astral-sh/uv/pull/12246
+    python_group.add_option(
+        "--python-preference",
+        type="choice",
+        choices=["only-managed", "managed", "system", "only-system"],
+    )
+    parser.add_option_group(python_group)
+
+    # Global options
+    global_group = optparse.OptionGroup(parser, "Global options")
+    global_group.add_option("-q", "--quiet", action="count", default=0)
+    global_group.add_option("-v", "--verbose", action="count", default=0)
+    global_group.add_option(
+        "--color", type="choice", choices=["auto", "always", "never"]
+    )
+    global_group.add_option("--native-tls", action="store_true")
+    global_group.add_option("--offline", action="store_true")
+    global_group.add_option(
+        "--allow-insecure-host", action="append", dest="insecure_hosts"
+    )
+    global_group.add_option("--no-progress", action="store_true")
+    global_group.add_option("--directory")
+    global_group.add_option("--project")
+    global_group.add_option("--config-file")
+    global_group.add_option("--no-config", action="store_true")
+    parser.add_option_group(global_group)
+
+    return parser
+
+
+def _parse_args(
+    parser: optparse.OptionParser, args: List[str]
+) -> Tuple[optparse.Values, List[str]]:
+    """
+    Parse the command-line options found in 'args'.
+
+    Replacement for parser.parse_args that handles unknown arguments
+    by keeping them in the command list instead of erroring and
+    discarding them.
+    """
+    parser.rargs = args
+    parser.largs = []
+    options = parser.get_default_values()
+    try:
+        parser._process_args(parser.largs, parser.rargs, options)
+    except optparse.BadOptionError as err:
+        # If we hit an argument that is not recognized, we put it
+        # back into the unconsumed arguments
+        parser.rargs = [err.opt_str] + parser.rargs
+    return options, parser.rargs
+
+
+def _check_working_dir_files(
+    uv_run_args: optparse.Values, runtime_env: Dict[str, Any]
+) -> None:
+    """
+    Check that the files required by uv are local to the working_dir. This catches
+    the most common cases of how things are different in Ray, i.e. not the whole file
+    system will be available on the workers, only the working_dir.
+
+    The function won't return anything, it just raises a RuntimeError if there is an error.
+    """
+    working_dir = Path(runtime_env["working_dir"]).resolve()
+
+    # Check if the requirements.txt file is in the working_dir
+    if uv_run_args.with_requirements:
+        for requirements_file in uv_run_args.with_requirements:
+            if not Path(requirements_file).resolve().is_relative_to(working_dir):
+                raise RuntimeError(
+                    f"You specified --with-requirements={uv_run_args.with_requirements} but "
+                    f"the requirements file is not in the working_dir {runtime_env['working_dir']}, "
+                    "so the workers will not have access to the file. Make sure "
+                    "the requirements file is in the working directory. "
+                    "You can do so by specifying --directory in 'uv run', by changing the current "
+                    "working directory before running 'uv run', or by using the 'working_dir' "
+                    "parameter of the runtime_environment."
+                )
+
+    # Check if the pyproject.toml file is in the working_dir
+    pyproject = None
+    if uv_run_args.no_project:
+        pyproject = None
+    elif uv_run_args.project:
+        pyproject = Path(uv_run_args.project)
+    else:
+        # Walk up the directory tree until pyproject.toml is found
+        current_path = Path.cwd().resolve()
+        while current_path != current_path.parent:
+            if (current_path / "pyproject.toml").exists():
+                pyproject = Path(current_path / "pyproject.toml")
+                break
+            current_path = current_path.parent
+
+    if pyproject and not pyproject.resolve().is_relative_to(working_dir):
+        raise RuntimeError(
+            f"Your {pyproject.resolve()} is not in the working_dir {runtime_env['working_dir']}, "
+            "so the workers will not have access to the file. Make sure "
+            "the pyproject.toml file is in the working directory. "
+            "You can do so by specifying --directory in 'uv run', by changing the current "
+            "working directory before running 'uv run', or by using the 'working_dir' "
+            "parameter of the runtime_environment."
+        )
+
+
+def _get_uv_run_cmdline() -> Optional[List[str]]:
+    """
+    Return the command line of the first ancestor process that was run with
+    "uv run" and None if there is no such ancestor.
+
+    uv spawns the python process as a child process, so we first check the
+    parent process command line. We also check our parent's parents since
+    the Ray driver might be run as a subprocess of the 'uv run' process.
+    """
+    parents = psutil.Process().parents()
+    for parent in parents:
+        try:
+            cmdline = parent.cmdline()
+            if (
+                len(cmdline) > 1
+                and os.path.basename(cmdline[0]) == "uv"
+                and cmdline[1] == "run"
+            ):
+                return cmdline
+        except psutil.NoSuchProcess:
+            continue
+        except psutil.AccessDenied:
+            continue
+    return None
+
+
+def hook(runtime_env: Optional[Dict[str, Any]]) -> Dict[str, Any]:
+    """Hook that detects if the driver is run in 'uv run' and sets the runtime environment accordingly."""
+
+    runtime_env = copy.deepcopy(runtime_env) or {}
+
+    cmdline = _get_uv_run_cmdline()
+    if not cmdline:
+        # This means the driver was not run in a 'uv run' environment -- in this case
+        # we leave the runtime environment unchanged
+        return runtime_env
+
+    # First check that the "uv" and "pip" runtime environments are not used.
+    if "uv" in runtime_env or "pip" in runtime_env:
+        raise RuntimeError(
+            "You are using the 'pip' or 'uv' runtime environments together with "
+            "'uv run'. These are not compatible since 'uv run' will run the workers "
+            "in an isolated environment -- please add the 'pip' or 'uv' dependencies to your "
+            "'uv run' environment e.g. by including them in your pyproject.toml."
+        )
+
+    # Extract the arguments uv_run_args of 'uv run' that are not part of the command.
+    args_to_parse = cmdline[2:]  # Remove 'uv run' prefix
+    original_length = len(
+        args_to_parse
+    )  # Save before parsing (parser modifies in-place)
+
+    parser = _create_uv_run_parser()
+    (options, command) = _parse_args(parser, args_to_parse)
+
+    # Calculate how many arguments were consumed by the parser.
+    # Since disable_interspersed_args() is set, parsing stops at the first
+    # unrecognized argument (the command), so all consumed args are uv options.
+    args_consumed = original_length - len(command)
+    uv_run_args = cmdline[: 2 + args_consumed]
+
+    # Remove the "--directory" argument since it has already been taken into
+    # account when setting the current working directory of the current process.
+    # Also remove the "--module" argument, since the default_worker.py is
+    # invoked as a script and not as a module.
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--directory")
+    parser.add_argument("-m", "--module")
+    _, remaining_uv_run_args = parser.parse_known_args(uv_run_args)
+
+    # Pin the worker Python for `uv run` unless the driver already specified one.
+    #
+    # Without this, uv may resolve a different Python than the Ray driver (e.g. 3.12 instead of 3.11),
+    # which causes Ray to fail with a Python version mismatch. (See https://github.com/ray-project/ray/issues/59639)
+    # order of precedence:
+    # 1. options.python (driver specified)
+    # 2. env_vars["UV_PYTHON"]
+    # 3. current os.environ["UV_PYTHON"]
+    # 4. platform.python_version() (since uv run uses the same Python as the driver)
+    if not options.python:
+        env_vars = runtime_env.get("env_vars") or {}
+        uv_python = (
+            env_vars.get("UV_PYTHON")
+            or os.environ.get("UV_PYTHON")
+            or platform.python_version()
+        )
+        remaining_uv_run_args = remaining_uv_run_args + ["--python", uv_python]
+
+    # Append "python" to the end so that when Ray adds "-m default_worker.py",
+    # it becomes "uv run --python X.Y.Z python -m default_worker.py"
+    remaining_uv_run_args = remaining_uv_run_args + ["python"]
+
+    runtime_env["py_executable"] = " ".join(remaining_uv_run_args)
+
+    # If the user specified a working_dir, we always honor it, otherwise
+    # use the same working_dir that uv run would use
+    if "working_dir" not in runtime_env:
+        runtime_env["working_dir"] = os.getcwd()
+
+    # Validate that pyproject.toml and requirements files are within working_dir
+    # This prevents runtime errors on workers when files are not accessible
+    # Only validate for local paths - remote URIs will be downloaded by Ray
+    working_dir = runtime_env["working_dir"]
+    if _is_path(working_dir):
+        _check_working_dir_files(options, runtime_env)
+
+    return runtime_env
+
+
+# This __main__ is used for unit testing if the runtime_env_hook picks up the
+# right settings.
+if __name__ == "__main__":
+    import json
+
+    test_parser = argparse.ArgumentParser()
+    test_parser.add_argument("--extra-args", action="store_true")
+    test_parser.add_argument("runtime_env")
+    args = test_parser.parse_args()
+
+    # If the env variable is set, add one more level of subprocess indirection
+    if os.environ.get("RAY_TEST_UV_ADD_SUBPROCESS_INDIRECTION") == "1":
+        import subprocess
+
+        env = os.environ.copy()
+        env.pop("RAY_TEST_UV_ADD_SUBPROCESS_INDIRECTION")
+        subprocess.check_call([sys.executable] + sys.argv, env=env)
+        sys.exit(0)
+
+    # If the following env variable is set, we use multiprocessing
+    # spawn to start the subprocess, since it uses a different way to
+    # modify the command line than subprocess.check_call
+    if os.environ.get("RAY_TEST_UV_MULTIPROCESSING_SPAWN") == "1":
+        import multiprocessing
+
+        multiprocessing.set_start_method("spawn")
+        pool = multiprocessing.Pool(processes=1)
+        runtime_env = json.loads(args.runtime_env)
+        print(json.dumps(pool.apply(hook, (runtime_env,))))
+        sys.exit(0)
+
+    # We purposefully modify sys.argv here to make sure the hook is robust
+    # against such modification.
+    sys.argv.pop(1)
+    runtime_env = json.loads(args.runtime_env)
+    print(json.dumps(hook(runtime_env)))

python/ray/_private/runtime_env/validation.py

@@ -0,0 +1,466 @@
+import logging
+import sys
+from collections import OrderedDict
+from pathlib import Path
+from typing import Dict, List, Optional, Union
+
+import yaml
+
+from ray._private.path_utils import is_path
+from ray._private.runtime_env.packaging import parse_path
+
+logger = logging.getLogger(__name__)
+
+
+def validate_path(path: str) -> None:
+    """Parse the path to ensure it is well-formed and exists."""
+    parse_path(path)
+
+
+def validate_uri(uri: str):
+    try:
+        from ray._private.runtime_env.packaging import Protocol, parse_uri
+
+        protocol, path = parse_uri(uri)
+    except ValueError:
+        raise ValueError(
+            f"{uri} is not a valid URI. Passing directories or modules to "
+            "be dynamically uploaded is only supported at the job level "
+            "(i.e., passed to `ray.init`)."
+        )
+
+    supported_extensions = (".zip", ".whl", ".tar.gz", ".tgz")
+    if protocol in Protocol.remote_protocols() and not any(
+        path.endswith(ext) for ext in supported_extensions
+    ):
+        raise ValueError(
+            "Only .zip, .whl, .tar.gz, and .tgz files supported for remote URIs."
+        )
+
+
+def _handle_local_deps_requirement_file(requirements_file: str):
+    """Read the given [requirements_file], and return all required dependencies."""
+    requirements_path = Path(requirements_file)
+    if not requirements_path.is_file():
+        raise ValueError(f"{requirements_path} is not a valid file")
+    return requirements_path.read_text().strip().split("\n")
+
+
+def validate_py_modules_uris(py_modules_uris: List[str]) -> List[str]:
+    """Parses and validates a 'py_modules' option.
+
+    Expects py_modules to be a list of URIs.
+    """
+    if not isinstance(py_modules_uris, list):
+        raise TypeError(
+            "`py_modules` must be a list of strings, got " f"{type(py_modules_uris)}."
+        )
+
+    for module in py_modules_uris:
+
+        if not isinstance(module, str):
+            raise TypeError("`py_module` must be a string, got " f"{type(module)}.")
+
+        validate_uri(module)
+
+
+def parse_and_validate_py_modules(py_modules: List[str]) -> List[str]:
+    """Parses and validates a 'py_modules' option.
+
+    Expects py_modules to be a list of local paths or URIs.
+    """
+    if not isinstance(py_modules, list):
+        raise TypeError(
+            "`py_modules` must be a list of strings, got " f"{type(py_modules)}."
+        )
+
+    for module in py_modules:
+
+        if not isinstance(module, str):
+            raise TypeError("`py_module` must be a string, got " f"{type(module)}.")
+
+        if is_path(module):
+            validate_path(module)
+        else:
+            validate_uri(module)
+
+    return py_modules
+
+
+def validate_working_dir_uri(working_dir_uri: str) -> str:
+    """Parses and validates a 'working_dir' option."""
+    if not isinstance(working_dir_uri, str):
+        raise TypeError(
+            "`working_dir` must be a string, got " f"{type(working_dir_uri)}."
+        )
+
+    validate_uri(working_dir_uri)
+
+
+def parse_and_validate_working_dir(working_dir: str) -> str:
+    """Parses and validates a 'working_dir' option.
+
+    This can be a URI or a path.
+    """
+    assert working_dir is not None
+
+    if not isinstance(working_dir, str):
+        raise TypeError("`working_dir` must be a string, got " f"{type(working_dir)}.")
+
+    if is_path(working_dir):
+        validate_path(working_dir)
+    else:
+        validate_uri(working_dir)
+
+    return working_dir
+
+
+def parse_and_validate_conda(conda: Union[str, dict]) -> Union[str, dict]:
+    """Parses and validates a user-provided 'conda' option.
+
+    Conda can be one of three cases:
+        1) A dictionary describing the env. This is passed through directly.
+        2) A string referring to the name of a preinstalled conda env.
+        3) A string pointing to a local conda YAML file. This is detected
+           by looking for a '.yaml' or '.yml' suffix. In this case, the file
+           will be read as YAML and passed through as a dictionary.
+    """
+    assert conda is not None
+
+    if sys.platform == "win32":
+        logger.warning(
+            "runtime environment support is experimental on Windows. "
+            "If you run into issues please file a report at "
+            "https://github.com/ray-project/ray/issues."
+        )
+
+    result = conda
+    if isinstance(conda, str):
+        file_path = Path(conda)
+        if file_path.suffix in (".yaml", ".yml"):
+            if not file_path.is_file():
+                raise ValueError(f"Can't find conda YAML file {file_path}.")
+            try:
+                result = yaml.safe_load(file_path.read_text())
+            except Exception as e:
+                raise ValueError(f"Failed to read conda file {file_path}: {e}.")
+        elif file_path.is_absolute():
+            if not file_path.is_dir():
+                raise ValueError(f"Can't find conda env directory {file_path}.")
+            result = str(file_path)
+    elif isinstance(conda, dict):
+        result = conda
+    else:
+        raise TypeError(
+            "runtime_env['conda'] must be of type str or " f"dict, got {type(conda)}."
+        )
+
+    return result
+
+
+def parse_and_validate_uv(uv: Union[str, List[str], Dict]) -> Optional[Dict]:
+    """Parses and validates a user-provided 'uv' option.
+
+    The value of the input 'uv' field can be one of two cases:
+        1) A List[str] describing the requirements. This is passed through.
+           Example usage: ["tensorflow", "requests"]
+        2) a string containing the path to a local pip “requirements.txt” file.
+        3) A python dictionary that has one field:
+            a) packages (required, List[str]): a list of uv packages, it same as 1).
+            b) uv_check (optional, bool): whether to enable pip check at the end of uv
+               install, default to False.
+            c) uv_version (optional, str): user provides a specific uv to use; if
+               unspecified, default version of uv will be used.
+            d) uv_pip_install_options (optional, List[str]): user-provided options for
+              `uv pip install` command, default to ["--no-cache"].
+
+    The returned parsed value will be a list of packages. If a Ray library
+    (e.g. "ray[serve]") is specified, it will be deleted and replaced by its
+    dependencies (e.g. "uvicorn", "requests").
+    """
+    assert uv is not None
+    if sys.platform == "win32":
+        logger.warning(
+            "runtime environment support is experimental on Windows. "
+            "If you run into issues please file a report at "
+            "https://github.com/ray-project/ray/issues."
+        )
+
+    result: str = ""
+    if isinstance(uv, str):
+        uv_list = _handle_local_deps_requirement_file(uv)
+        result = dict(packages=uv_list, uv_check=False)
+    elif isinstance(uv, list) and all(isinstance(dep, str) for dep in uv):
+        result = dict(packages=uv, uv_check=False)
+    elif isinstance(uv, dict):
+        if set(uv.keys()) - {
+            "packages",
+            "uv_check",
+            "uv_version",
+            "uv_pip_install_options",
+        }:
+            raise ValueError(
+                "runtime_env['uv'] can only have these fields: "
+                "packages, uv_check, uv_version and uv_pip_install_options, but got: "
+                f"{list(uv.keys())}"
+            )
+        if "packages" not in uv:
+            raise ValueError(
+                f"runtime_env['uv'] must include field 'packages', but got {uv}"
+            )
+        if "uv_check" in uv and not isinstance(uv["uv_check"], bool):
+            raise TypeError(
+                "runtime_env['uv']['uv_check'] must be of type bool, "
+                f"got {type(uv['uv_check'])}"
+            )
+        if "uv_version" in uv and not isinstance(uv["uv_version"], str):
+            raise TypeError(
+                "runtime_env['uv']['uv_version'] must be of type str, "
+                f"got {type(uv['uv_version'])}"
+            )
+        if "uv_pip_install_options" in uv:
+            if not isinstance(uv["uv_pip_install_options"], list):
+                raise TypeError(
+                    "runtime_env['uv']['uv_pip_install_options'] must be of type "
+                    f"list[str] got {type(uv['uv_pip_install_options'])}"
+                )
+            # Check each item in installation option.
+            for idx, cur_opt in enumerate(uv["uv_pip_install_options"]):
+                if not isinstance(cur_opt, str):
+                    raise TypeError(
+                        "runtime_env['uv']['uv_pip_install_options'] must be of type "
+                        f"list[str] got {type(cur_opt)} for {idx}-th item."
+                    )
+
+        result = uv.copy()
+        result["uv_check"] = uv.get("uv_check", False)
+        result["uv_pip_install_options"] = uv.get(
+            "uv_pip_install_options", ["--no-cache"]
+        )
+        if not isinstance(uv["packages"], list):
+            raise ValueError(
+                "runtime_env['uv']['packages'] must be of type list, "
+                f"got: {type(uv['packages'])}"
+            )
+    else:
+        raise TypeError(
+            "runtime_env['uv'] must be of type " f"List[str], or dict, got {type(uv)}"
+        )
+
+    # Deduplicate packages for package lists.
+    result["packages"] = list(OrderedDict.fromkeys(result["packages"]))
+
+    if len(result["packages"]) == 0:
+        result = None
+    logger.debug(f"Rewrote runtime_env `uv` field from {uv} to {result}.")
+    return result
+
+
+def parse_and_validate_pip(pip: Union[str, List[str], Dict]) -> Optional[Dict]:
+    """Parses and validates a user-provided 'pip' option.
+
+    The value of the input 'pip' field can be one of two cases:
+        1) A List[str] describing the requirements. This is passed through.
+        2) A string pointing to a local requirements file. In this case, the
+           file contents will be read split into a list.
+        3) A python dictionary that has three fields:
+            a) packages (required, List[str]): a list of pip packages, it same as 1).
+            b) pip_check (optional, bool): whether to enable pip check at the end of pip
+               install, default to False.
+            c) pip_version (optional, str): the version of pip, ray will spell
+               the package name 'pip' in front of the `pip_version` to form the final
+               requirement string, the syntax of a requirement specifier is defined in
+               full in PEP 508.
+            d) pip_install_options (optional, List[str]): user-provided options for
+              `pip install` command, defaults to ["--disable-pip-version-check", "--no-cache-dir"].
+
+    The returned parsed value will be a list of pip packages. If a Ray library
+    (e.g. "ray[serve]") is specified, it will be deleted and replaced by its
+    dependencies (e.g. "uvicorn", "requests").
+    """
+    assert pip is not None
+    result = None
+
+    if sys.platform == "win32":
+        logger.warning(
+            "runtime environment support is experimental on Windows. "
+            "If you run into issues please file a report at "
+            "https://github.com/ray-project/ray/issues."
+        )
+    if isinstance(pip, str):
+        # We have been given a path to a requirements.txt file.
+        pip_list = _handle_local_deps_requirement_file(pip)
+        result = dict(
+            packages=pip_list,
+            pip_check=False,
+        )
+    elif isinstance(pip, list) and all(isinstance(dep, str) for dep in pip):
+        result = dict(packages=pip, pip_check=False)
+    elif isinstance(pip, dict):
+        if set(pip.keys()) - {
+            "packages",
+            "pip_check",
+            "pip_install_options",
+            "pip_version",
+        }:
+            raise ValueError(
+                "runtime_env['pip'] can only have these fields: "
+                "packages, pip_check, pip_install_options and pip_version, but got: "
+                f"{list(pip.keys())}"
+            )
+
+        if "pip_check" in pip and not isinstance(pip["pip_check"], bool):
+            raise TypeError(
+                "runtime_env['pip']['pip_check'] must be of type bool, "
+                f"got {type(pip['pip_check'])}"
+            )
+        if "pip_version" in pip:
+            if not isinstance(pip["pip_version"], str):
+                raise TypeError(
+                    "runtime_env['pip']['pip_version'] must be of type str, "
+                    f"got {type(pip['pip_version'])}"
+                )
+        if "pip_install_options" in pip:
+            if not isinstance(pip["pip_install_options"], list):
+                raise TypeError(
+                    "runtime_env['pip']['pip_install_options'] must be of type "
+                    f"list[str] got {type(pip['pip_install_options'])}"
+                )
+            # Check each item in installation option.
+            for idx, cur_opt in enumerate(pip["pip_install_options"]):
+                if not isinstance(cur_opt, str):
+                    raise TypeError(
+                        "runtime_env['pip']['pip_install_options'] must be of type "
+                        f"list[str] got {type(cur_opt)} for {idx}-th item."
+                    )
+
+        result = pip.copy()
+        # Contrary to pip_check, we do not insert the default value of pip_install_options.
+        # This is to maintain backwards compatibility with ray==2.0.1
+        result["pip_check"] = pip.get("pip_check", False)
+
+        if "packages" not in pip:
+            raise ValueError(
+                f"runtime_env['pip'] must include field 'packages', but got {pip}"
+            )
+        elif isinstance(pip["packages"], str):
+            result["packages"] = _handle_local_deps_requirement_file(pip["packages"])
+        elif not isinstance(pip["packages"], list):
+            raise ValueError(
+                "runtime_env['pip']['packages'] must be of type str of list, "
+                f"got: {type(pip['packages'])}"
+            )
+    else:
+        raise TypeError(
+            "runtime_env['pip'] must be of type str or " f"List[str], got {type(pip)}"
+        )
+
+    # Eliminate duplicates to prevent `pip install` from erroring. Use
+    # OrderedDict to preserve the order of the list.  This makes the output
+    # deterministic and easier to debug, because pip install can have
+    # different behavior depending on the order of the input.
+    result["packages"] = list(OrderedDict.fromkeys(result["packages"]))
+
+    if len(result["packages"]) == 0:
+        result = None
+
+    logger.debug(f"Rewrote runtime_env `pip` field from {pip} to {result}.")
+
+    return result
+
+
+def parse_and_validate_container(container: List[str]) -> List[str]:
+    """Parses and validates a user-provided 'container' option.
+
+    This is passed through without validation (for now).
+    """
+    assert container is not None
+    return container
+
+
+def parse_and_validate_excludes(excludes: List[str]) -> List[str]:
+    """Parses and validates a user-provided 'excludes' option.
+
+    This is validated to verify that it is of type List[str].
+
+    If an empty list is passed, we return `None` for consistency.
+    """
+    assert excludes is not None
+
+    if isinstance(excludes, list) and len(excludes) == 0:
+        return None
+
+    if isinstance(excludes, list) and all(isinstance(path, str) for path in excludes):
+        return excludes
+    else:
+        raise TypeError(
+            "runtime_env['excludes'] must be of type "
+            f"List[str], got {type(excludes)}"
+        )
+
+
+def parse_and_validate_env_vars(env_vars: Dict[str, str]) -> Optional[Dict[str, str]]:
+    """Parses and validates a user-provided 'env_vars' option.
+
+    This is validated to verify that all keys and vals are strings.
+
+    If an empty dictionary is passed, we return `None` for consistency.
+
+    Args:
+        env_vars: A dictionary of environment variables to set in the
+            runtime environment.
+
+    Returns:
+        The validated env_vars dictionary, or None if it was empty.
+
+    Raises:
+        TypeError: If the env_vars is not a dictionary of strings. The error message
+            will include the type of the invalid value.
+    """
+    assert env_vars is not None
+    if len(env_vars) == 0:
+        return None
+
+    if not isinstance(env_vars, dict):
+        raise TypeError(
+            "runtime_env['env_vars'] must be of type "
+            f"Dict[str, str], got {type(env_vars)}"
+        )
+
+    for key, val in env_vars.items():
+        if not isinstance(key, str):
+            raise TypeError(
+                "runtime_env['env_vars'] must be of type "
+                f"Dict[str, str], but the key {key} is of type {type(key)}"
+            )
+        if not isinstance(val, str):
+            raise TypeError(
+                "runtime_env['env_vars'] must be of type "
+                f"Dict[str, str], but the value {val} is of type {type(val)}"
+            )
+
+    return env_vars
+
+
+# Dictionary mapping runtime_env options with the function to parse and
+# validate them.
+OPTION_TO_VALIDATION_FN = {
+    "py_modules": parse_and_validate_py_modules,
+    "working_dir": parse_and_validate_working_dir,
+    "excludes": parse_and_validate_excludes,
+    "conda": parse_and_validate_conda,
+    "pip": parse_and_validate_pip,
+    "uv": parse_and_validate_uv,
+    "env_vars": parse_and_validate_env_vars,
+    "container": parse_and_validate_container,
+}
+
+# RuntimeEnv can be created with local paths
+# for these options. However, after the packages
+# for these options have been uploaded to GCS,
+# they must be URIs. These functions provide the ability
+# to validate that these options only contain well-formed URIs.
+OPTION_TO_NO_PATH_VALIDATION_FN = {
+    "working_dir": validate_working_dir_uri,
+    "py_modules": validate_py_modules_uris,
+}