Add files using upload-large-folder tool

lib/python3.12/site-packages/ray/air/__init__.py

@@ -0,0 +1,21 @@
+from ray.air.config import (
+    CheckpointConfig,
+    FailureConfig,
+    RunConfig,
+    ScalingConfig,
+)
+from ray.air.data_batch_type import DataBatchType
+from ray.air.execution.resources.request import AcquiredResources, ResourceRequest
+from ray.air.result import Result
+import ray.data  # noqa: F401  # TODO: This is a hack to avoid circular import
+
+__all__ = [
+    "DataBatchType",
+    "RunConfig",
+    "Result",
+    "ScalingConfig",
+    "FailureConfig",
+    "CheckpointConfig",
+    "AcquiredResources",
+    "ResourceRequest",
+]

lib/python3.12/site-packages/ray/air/_internal/config.py

@@ -0,0 +1,47 @@
+import dataclasses
+from typing import Iterable
+
+
+def ensure_only_allowed_dataclass_keys_updated(
+    dataclass: dataclasses.dataclass,
+    allowed_keys: Iterable[str],
+):
+    """
+    Validate dataclass by raising an exception if any key not included in
+    ``allowed_keys`` differs from the default value.
+
+    A ``ValueError`` will also be raised if any of the ``allowed_keys``
+    is not present in ``dataclass.__dict__``.
+
+    Args:
+        dataclass: Dict or dataclass to check.
+        allowed_keys: dataclass attribute keys that can have a value different than
+            the default one.
+    """
+    default_data = dataclass.__class__()
+
+    allowed_keys = set(allowed_keys)
+
+    # TODO: split keys_not_in_dict validation to a separate function.
+    keys_not_in_dict = [key for key in allowed_keys if key not in default_data.__dict__]
+    if keys_not_in_dict:
+        raise ValueError(
+            f"Key(s) {keys_not_in_dict} are not present in "
+            f"{dataclass.__class__.__name__}. "
+            "Remove them from `allowed_keys`. "
+            f"Valid keys: {list(default_data.__dict__.keys())}"
+        )
+
+    # These keys should not have been updated in the `dataclass` object
+    prohibited_keys = set(default_data.__dict__) - allowed_keys
+
+    bad_keys = [
+        key
+        for key in prohibited_keys
+        if dataclass.__dict__[key] != default_data.__dict__[key]
+    ]
+    if bad_keys:
+        raise ValueError(
+            f"Key(s) {bad_keys} are not allowed to be updated in the current context. "
+            "Remove them from the dataclass."
+        )

lib/python3.12/site-packages/ray/air/_internal/device_manager/__init__.py

@@ -0,0 +1,92 @@
+import logging
+import threading
+from typing import Optional
+
+import ray
+import ray._private.ray_constants as ray_constants
+from ray.air._internal.device_manager.cpu import CPUTorchDeviceManager
+from ray.air._internal.device_manager.hpu import HPUTorchDeviceManager
+from ray.air._internal.device_manager.npu import NPUTorchDeviceManager
+from ray.air._internal.device_manager.nvidia_gpu import CUDATorchDeviceManager
+from ray.air._internal.device_manager.torch_device_manager import TorchDeviceManager
+
+logger = logging.getLogger(__name__)
+
+
+DEFAULT_TORCH_DEVICE_MANAGER_CLS = CPUTorchDeviceManager
+
+
+SUPPORTED_ACCELERATOR_TORCH_DEVICE_MANAGER = {
+    ray_constants.GPU: CUDATorchDeviceManager,
+    ray_constants.HPU: HPUTorchDeviceManager,
+    ray_constants.NPU: NPUTorchDeviceManager,
+}
+
+
+def register_custom_torch_dist_backend(backend: Optional[str] = None) -> None:
+    if backend == "hccl":
+        # The name for the communication backend of Habana and torch-npu is the same.
+        HPUTorchDeviceManager.register_custom_torch_dist_backend()
+
+        NPUTorchDeviceManager.register_custom_torch_dist_backend()
+
+
+_torch_device_manager = None
+_torch_device_manager_lock = threading.Lock()
+
+
+def get_torch_device_manager_by_context() -> TorchDeviceManager:
+    global _torch_device_manager
+
+    with _torch_device_manager_lock:
+        if not _torch_device_manager:
+            existing_device_manager_cls = None
+            resources = ray.get_runtime_context().get_accelerator_ids()
+
+            # select correct accelerator type from resources
+            for resource_type, resource_value in resources.items():
+                device_manager_cls = SUPPORTED_ACCELERATOR_TORCH_DEVICE_MANAGER.get(
+                    resource_type, None
+                )
+                if resource_value and device_manager_cls:
+                    # An error will raise when multiple accelerators are specified.
+                    if existing_device_manager_cls:
+                        raise RuntimeError(
+                            "Unable to determine the appropriate DeviceManager "
+                            f"for the specified resources {resources}."
+                        )
+                    else:
+                        existing_device_manager_cls = device_manager_cls
+
+            device_manager_cls = (
+                existing_device_manager_cls or DEFAULT_TORCH_DEVICE_MANAGER_CLS
+            )
+
+            _torch_device_manager = device_manager_cls()
+
+    return _torch_device_manager
+
+
+def get_torch_device_manager_by_device_type(device_type: str):
+    if device_type.lower() == ray_constants.GPU.lower() or device_type == "cuda":
+        return CUDATorchDeviceManager()
+    elif device_type.lower() == ray_constants.NPU.lower():
+        return NPUTorchDeviceManager()
+    elif device_type.lower() == ray_constants.HPU.lower():
+        return HPUTorchDeviceManager()
+    elif device_type.lower() == "cpu":
+        return CPUTorchDeviceManager()
+
+    raise RuntimeError(f"Device type {device_type} cannot be recognized.")
+
+
+__all__ = [
+    TorchDeviceManager,
+    CPUTorchDeviceManager,
+    CUDATorchDeviceManager,
+    HPUTorchDeviceManager,
+    NPUTorchDeviceManager,
+    register_custom_torch_dist_backend,
+    get_torch_device_manager_by_context,
+    get_torch_device_manager_by_device_type,
+]

lib/python3.12/site-packages/ray/air/_internal/device_manager/cpu.py

@@ -0,0 +1,30 @@
+from contextlib import contextmanager
+from typing import List
+
+import torch
+
+from ray.air._internal.device_manager.torch_device_manager import TorchDeviceManager
+
+
+class CPUTorchDeviceManager(TorchDeviceManager):
+    """CPU device manager"""
+
+    def is_available(self) -> bool():
+        return True
+
+    def get_devices(self) -> List[torch.device]:
+        """Gets the correct torch device list configured for this process."""
+        return [torch.device("cpu")]
+
+    def supports_stream(self) -> bool:
+        """Validate if the device type support create a stream"""
+        return False
+
+    def get_stream_context(self, stream):
+        """Return empty context mananger for CPU."""
+
+        @contextmanager
+        def default_context_manager():
+            yield
+
+        return default_context_manager()

lib/python3.12/site-packages/ray/air/_internal/device_manager/hpu.py

@@ -0,0 +1,50 @@
+from contextlib import contextmanager
+from typing import List, Union
+
+import torch
+
+from ray._private.accelerators.hpu import HPU_PACKAGE_AVAILABLE
+from ray.air._internal.device_manager.torch_device_manager import TorchDeviceManager
+
+if HPU_PACKAGE_AVAILABLE:
+    import habana_frameworks.torch.hpu as torch_hpu
+
+
+class HPUTorchDeviceManager(TorchDeviceManager):
+    """HPU device manager"""
+
+    @staticmethod
+    def register_custom_torch_dist_backend():
+        if HPU_PACKAGE_AVAILABLE:
+            import habana_frameworks.torch.core  # noqa: F401
+            import habana_frameworks.torch.distributed.hccl  # noqa: F401
+
+    def is_available(self) -> bool():
+        if not HPU_PACKAGE_AVAILABLE:
+            return False
+
+        return torch_hpu.is_available()
+
+    def get_devices(self) -> List[torch.device]:
+        if not self.is_available():
+            raise RuntimeError(
+                "Using HPUTorchDeviceManager but torch hpu is not available."
+            )
+
+        return [torch.device("hpu")]
+
+    def set_device(self, device: Union[torch.device, int, str, None]):
+        torch_hpu.set_device(device)
+
+    def supports_stream(self) -> bool:
+        """Validate if the device type support create a stream"""
+        return False
+
+    def get_stream_context(self, stream):
+        """Get HPU stream context manager, empty so far."""
+
+        @contextmanager
+        def default_context_manager():
+            yield
+
+        return default_context_manager()

lib/python3.12/site-packages/ray/air/_internal/device_manager/npu.py

@@ -0,0 +1,104 @@
+import os
+from importlib.util import find_spec
+from typing import List, Union
+
+import torch
+
+import ray
+import ray._private.ray_constants as ray_constants
+from ray._private.accelerators.npu import ASCEND_RT_VISIBLE_DEVICES_ENV_VAR
+from ray.air._internal.device_manager.torch_device_manager import TorchDeviceManager
+
+
+def is_package_present(package_name: str) -> bool:
+    try:
+        return find_spec(package_name) is not None
+    except ModuleNotFoundError:
+        return False
+
+
+NPU_TORCH_PACKAGE_AVAILABLE = is_package_present("torch_npu")
+
+
+if NPU_TORCH_PACKAGE_AVAILABLE:
+    import torch_npu  # noqa: F401
+
+
+class NPUTorchDeviceManager(TorchDeviceManager):
+    """Ascend NPU device manager"""
+
+    @staticmethod
+    def register_custom_torch_dist_backend():
+        if NPU_TORCH_PACKAGE_AVAILABLE:
+            import torch_npu  # noqa: F401, F811
+
+    def is_available(self) -> bool:
+        if not NPU_TORCH_PACKAGE_AVAILABLE:
+            return False
+
+        return torch.npu.is_available()
+
+    def get_devices(self) -> List[torch.device]:
+        """Gets the correct torch device list configured for this process.
+
+        Returns a list of torch NPU devices allocated for the current worker.
+        If no NPUs are assigned, then it returns a list with a single CPU device.
+        """
+        if NPU_TORCH_PACKAGE_AVAILABLE and torch.npu.is_available():
+            npu_ids = [
+                str(id)
+                for id in ray.get_runtime_context().get_accelerator_ids()[
+                    ray_constants.NPU
+                ]
+            ]
+
+            device_ids = []
+
+            if len(npu_ids) > 0:
+                npu_visible_str = os.environ.get(ASCEND_RT_VISIBLE_DEVICES_ENV_VAR, "")
+                if npu_visible_str and npu_visible_str != "NoDevFiles":
+                    npu_visible_list = npu_visible_str.split(",")
+                else:
+                    npu_visible_list = []
+
+                for npu_id in npu_ids:
+                    try:
+                        device_ids.append(npu_visible_list.index(npu_id))
+                    except IndexError:
+                        raise RuntimeError(
+                            "ASCEND_RT_VISIBLE_DEVICES set incorrectly. "
+                            f"Got {npu_visible_str}, expected to include {npu_id}. "
+                            "Did you override the `ASCEND_RT_VISIBLE_DEVICES` "
+                            "environment variable?"
+                        )
+            else:
+                # If called on the driver or outside of Ray Train, return the
+                # 0th device.
+                device_ids.append(0)
+
+            devices = [torch.device(f"npu:{device_id}") for device_id in device_ids]
+        else:
+            raise RuntimeError(
+                "Using NPUTorchDeviceManager but torch npu is not available."
+            )
+
+        return devices
+
+    def set_device(self, device: Union[torch.device, int]):
+        torch.npu.set_device(device)
+
+    def supports_stream(self) -> bool:
+        """Validate if the device type support to create a stream"""
+        return True
+
+    def create_stream(self, device):
+        """Create a stream on NPU device"""
+        return torch.npu.Stream(device)
+
+    def get_stream_context(self, stream):
+        """Get a torch.stream context on NPU device"""
+        return torch.npu.stream(stream)
+
+    def get_current_stream(self):
+        """Get current stream for NPU device"""
+        return torch.npu.current_stream()

lib/python3.12/site-packages/ray/air/_internal/device_manager/nvidia_gpu.py

@@ -0,0 +1,79 @@
+import os
+from typing import List, Union
+
+import torch
+
+import ray
+from ray.air._internal.device_manager.torch_device_manager import TorchDeviceManager
+
+
+class CUDATorchDeviceManager(TorchDeviceManager):
+    """CUDA device manager"""
+
+    def is_available(self) -> bool():
+        return torch.cuda.is_available()
+
+    def get_devices(self) -> List[torch.device]:
+        """Gets the correct torch device list configured for this process.
+
+        Returns a list of torch CUDA devices allocated for the current worker.
+        If no GPUs are assigned, then it returns a list with a single CPU device.
+
+        Assumes that `CUDA_VISIBLE_DEVICES` is set and is a
+        superset of the `ray.get_gpu_ids()`.
+        """
+
+        # GPU IDs are assigned by Ray after you specify "use_gpu"
+        # GPU `ray.get_gpu_ids()` may return ints or may return strings.
+        # We should always convert to strings.
+        gpu_ids = [str(id) for id in ray.get_gpu_ids()]
+
+        device_ids = []
+
+        if len(gpu_ids) > 0:
+            cuda_visible_str = os.environ.get("CUDA_VISIBLE_DEVICES", "")
+            if cuda_visible_str and cuda_visible_str != "NoDevFiles":
+                cuda_visible_list = cuda_visible_str.split(",")
+            else:
+                cuda_visible_list = []
+
+            # By default, there should only be one GPU ID if `use_gpu=True`.
+            # If there are multiple GPUs, return a list of devices.
+            # If using fractional GPUs, these IDs are not guaranteed
+            # to be unique across different processes.
+            for gpu_id in gpu_ids:
+                try:
+                    device_ids.append(cuda_visible_list.index(gpu_id))
+                except IndexError:
+                    raise RuntimeError(
+                        "CUDA_VISIBLE_DEVICES set incorrectly. "
+                        f"Got {cuda_visible_str}, expected to include {gpu_id}. "
+                        "Did you override the `CUDA_VISIBLE_DEVICES` environment"
+                        " variable? If not, please help file an issue on Github."
+                    )
+
+        else:
+            # If called on the driver or outside of Ray Train, return the
+            # 0th device.
+            device_ids.append(0)
+
+        return [torch.device(f"cuda:{device_id}") for device_id in device_ids]
+
+    def set_device(self, device: Union[torch.device, int, str, None]):
+        torch.cuda.set_device(device)
+
+    def supports_stream(self) -> bool:
+        """Validate if the device type support create a stream"""
+        return True
+
+    def create_stream(self, device: torch.device) -> torch.cuda.Stream:
+        """Create a stream on cuda device"""
+        return torch.cuda.Stream(device)
+
+    def get_stream_context(self, stream):
+        """Get a stream context for cuda device"""
+        return torch.cuda.stream(stream)
+
+    def get_current_stream(self) -> torch.cuda.Stream:
+        """Get current stream for cuda device"""
+        return torch.cuda.current_stream()

lib/python3.12/site-packages/ray/air/_internal/device_manager/torch_device_manager.py

@@ -0,0 +1,40 @@
+from abc import ABC
+from typing import List, Union
+
+import torch
+
+
+class TorchDeviceManager(ABC):
+    """This class contains the function needed for supporting
+    an acclerator family in Ray AI Library.
+    """
+
+    def is_available(self) -> bool:
+        """Validate if device is available."""
+        ...
+
+    def get_devices(self) -> List[torch.device]:
+        """Gets the correct torch device configured for this process"""
+        ...
+
+    def set_device(self, device: Union[torch.device, int, str, None]):
+        """Set the correct device for this process"""
+        ...
+
+    def supports_stream(self) -> bool:
+        """Validate if the device type support create a stream"""
+        ...
+
+    def create_stream(self, device: torch.device):
+        """Create a device stream"""
+        ...
+
+    def get_stream_context(self, stream):
+        """Get a stream context of device. If device didn't support stream,
+        this should return a empty context manager instead of  None.
+        """
+        ...
+
+    def get_current_stream(self):
+        """Get current stream on accelerators like torch.cuda.current_stream"""
+        ...

lib/python3.12/site-packages/ray/air/_internal/filelock.py

@@ -0,0 +1,46 @@
+import hashlib
+import os
+from pathlib import Path
+
+from filelock import FileLock
+
+import ray
+
+RAY_LOCKFILE_DIR = "_ray_lockfiles"
+
+
+class TempFileLock:
+    """FileLock wrapper that uses temporary file locks.
+
+    The temporary directory that these locks are saved to can be configured via
+    the `RAY_TMPDIR` environment variable.
+
+    Args:
+        path: The file path that this temporary file lock is used for.
+            This will be used to generate the lockfile filename.
+            Ex: For concurrent writes to a file, this is the common filepath
+            that multiple processes are writing to.
+        **kwargs: Additional keyword arguments to pass to the underlying `FileLock`.
+    """
+
+    def __init__(self, path: str, **kwargs):
+        self.path = path
+        temp_dir = Path(ray._common.utils.get_user_temp_dir()).resolve()
+        self._lock_dir = temp_dir / RAY_LOCKFILE_DIR
+        self._path_hash = hashlib.sha1(
+            str(Path(self.path).resolve()).encode("utf-8")
+        ).hexdigest()
+        self._lock_path = self._lock_dir / f"{self._path_hash}.lock"
+
+        os.makedirs(str(self._lock_dir), exist_ok=True)
+        self._lock = FileLock(self._lock_path, **kwargs)
+
+    def __enter__(self):
+        self._lock.acquire()
+        return self
+
+    def __exit__(self, type, value, traceback):
+        self._lock.release()
+
+    def __getattr__(self, name):
+        return getattr(self._lock, name)

lib/python3.12/site-packages/ray/air/_internal/json.py

@@ -0,0 +1,31 @@
+import json
+import numbers
+
+import numpy as np
+
+
+class SafeFallbackEncoder(json.JSONEncoder):
+    def __init__(self, nan_str="null", **kwargs):
+        super(SafeFallbackEncoder, self).__init__(**kwargs)
+        self.nan_str = nan_str
+
+    def default(self, value):
+        try:
+            if type(value).__module__ == np.__name__ and isinstance(value, np.ndarray):
+                return value.tolist()
+
+            if isinstance(value, np.bool_):
+                return bool(value)
+
+            if np.isnan(value):
+                return self.nan_str
+
+            if issubclass(type(value), numbers.Integral):
+                return int(value)
+            if issubclass(type(value), numbers.Number):
+                return float(value)
+
+            return super(SafeFallbackEncoder, self).default(value)
+
+        except Exception:
+            return str(value)  # give up, just stringify it (ok for logs)

lib/python3.12/site-packages/ray/air/_internal/mlflow.py

@@ -0,0 +1,342 @@
+import logging
+import os
+from copy import deepcopy
+from typing import TYPE_CHECKING, Dict, Optional
+
+from packaging import version
+
+from ray._private.dict import flatten_dict
+
+if TYPE_CHECKING:
+    from mlflow.entities import Run
+    from mlflow.tracking import MlflowClient
+
+logger = logging.getLogger(__name__)
+
+
+class _MLflowLoggerUtil:
+    """Util class for setting up and logging to MLflow.
+
+    Use this util for any library that needs MLflow logging/tracking logic
+    such as Ray Tune or Ray Train.
+    """
+
+    def __init__(self):
+        import mlflow
+
+        self._mlflow = mlflow
+        self.experiment_id = None
+
+    def __deepcopy__(self, memo=None):
+        # mlflow is a module, and thus cannot be copied
+        _mlflow = self._mlflow
+        self.__dict__.pop("_mlflow")
+        dict_copy = deepcopy(self.__dict__, memo)
+        copied_object = _MLflowLoggerUtil()
+        copied_object.__dict__.update(dict_copy)
+        self._mlflow = _mlflow
+        copied_object._mlflow = _mlflow
+        return copied_object
+
+    def setup_mlflow(
+        self,
+        tracking_uri: Optional[str] = None,
+        registry_uri: Optional[str] = None,
+        experiment_id: Optional[str] = None,
+        experiment_name: Optional[str] = None,
+        tracking_token: Optional[str] = None,
+        artifact_location: Optional[str] = None,
+        create_experiment_if_not_exists: bool = True,
+    ):
+        """
+        Sets up MLflow.
+
+        Sets the Mlflow tracking uri & token, and registry URI. Also sets
+        the MLflow experiment that the logger should use, and possibly
+        creates new experiment if it does not exist.
+
+        Args:
+            tracking_uri: The tracking URI for the MLflow tracking
+                server.
+            registry_uri: The registry URI for the MLflow model registry.
+            experiment_id: The id of an already existing MLflow
+                experiment to use for logging. If None is passed in
+                here and the MFLOW_EXPERIMENT_ID is not set, or the
+                experiment with this id does not exist,
+                ``experiment_name`` will be used instead. This argument takes
+                precedence over ``experiment_name`` if both are passed in.
+            experiment_name: The experiment name to use for logging.
+                If None is passed in here, the MLFLOW_EXPERIMENT_NAME environment
+                variable is used to determine the experiment name.
+                If the experiment with the name already exists with MLflow,
+                it will be reused. If not, a new experiment will be created
+                with the provided name if
+                ``create_experiment_if_not_exists`` is set to True.
+            artifact_location: The location to store run artifacts.
+                If not provided, MLFlow picks an appropriate default.
+                Ignored if experiment already exists.
+            tracking_token: Tracking token used to authenticate with MLflow.
+            create_experiment_if_not_exists: Whether to create an
+                experiment with the provided name if it does not already
+                exist. Defaults to True.
+
+        Returns:
+            Whether setup is successful.
+        """
+        if tracking_token:
+            os.environ["MLFLOW_TRACKING_TOKEN"] = tracking_token
+
+        self._mlflow.set_tracking_uri(tracking_uri)
+        self._mlflow.set_registry_uri(registry_uri)
+
+        # First check experiment_id.
+        experiment_id = (
+            experiment_id
+            if experiment_id is not None
+            else os.environ.get("MLFLOW_EXPERIMENT_ID")
+        )
+        if experiment_id is not None:
+            from mlflow.exceptions import MlflowException
+
+            try:
+                self._mlflow.get_experiment(experiment_id=experiment_id)
+                logger.debug(
+                    f"Experiment with provided id {experiment_id} "
+                    "exists. Setting that as the experiment."
+                )
+                self.experiment_id = experiment_id
+                return
+            except MlflowException:
+                pass
+
+        # Then check experiment_name.
+        experiment_name = (
+            experiment_name
+            if experiment_name is not None
+            else os.environ.get("MLFLOW_EXPERIMENT_NAME")
+        )
+        if experiment_name is not None and self._mlflow.get_experiment_by_name(
+            name=experiment_name
+        ):
+            logger.debug(
+                f"Experiment with provided name {experiment_name} "
+                "exists. Setting that as the experiment."
+            )
+            self.experiment_id = self._mlflow.get_experiment_by_name(
+                experiment_name
+            ).experiment_id
+            return
+
+        # An experiment with the provided id or name does not exist.
+        # Create a new experiment if applicable.
+        if experiment_name and create_experiment_if_not_exists:
+            logger.debug(
+                "Existing experiment not found. Creating new "
+                f"experiment with name: {experiment_name}"
+            )
+            self.experiment_id = self._mlflow.create_experiment(
+                name=experiment_name, artifact_location=artifact_location
+            )
+            return
+
+        if create_experiment_if_not_exists:
+            raise ValueError(
+                f"Experiment with the provided experiment_id: "
+                f"{experiment_id} does not exist and no "
+                f"experiment_name provided. At least one of "
+                f"these has to be provided."
+            )
+        else:
+            raise ValueError(
+                f"Experiment with the provided experiment_id: "
+                f"{experiment_id} or experiment_name: "
+                f"{experiment_name} does not exist. Please "
+                f"create an MLflow experiment and provide "
+                f"either its id or name."
+            )
+
+    def _parse_dict(self, dict_to_log: Dict) -> Dict:
+        """Parses provided dict to convert all values to float.
+
+        MLflow can only log metrics that are floats. This does not apply to
+        logging parameters or artifacts.
+
+        Args:
+            dict_to_log: The dictionary containing the metrics to log.
+
+        Returns:
+            A dictionary containing the metrics to log with all values being
+                converted to floats, or skipped if not able to be converted.
+        """
+        new_dict = {}
+        for key, value in dict_to_log.items():
+            try:
+                value = float(value)
+                new_dict[key] = value
+            except (ValueError, TypeError):
+                logger.debug(
+                    "Cannot log key {} with value {} since the "
+                    "value cannot be converted to float.".format(key, value)
+                )
+                continue
+
+        return new_dict
+
+    def start_run(
+        self,
+        run_name: Optional[str] = None,
+        tags: Optional[Dict] = None,
+        set_active: bool = False,
+    ) -> "Run":
+        """Starts a new run and possibly sets it as the active run.
+
+        Args:
+            tags: Tags to set for the new run.
+            set_active: Whether to set the new run as the active run.
+                If an active run already exists, then that run is returned.
+
+        Returns:
+            The newly created MLflow run.
+        """
+        import mlflow
+        from mlflow.utils.mlflow_tags import MLFLOW_RUN_NAME
+
+        if tags is None:
+            tags = {}
+
+        if set_active:
+            return self._start_active_run(run_name=run_name, tags=tags)
+
+        client = self._get_client()
+        # If `mlflow==1.30.0` and we don't use `run_name`, then MLflow might error. For
+        # more information, see #29749.
+        if version.parse(mlflow.__version__) >= version.parse("1.30.0"):
+            run = client.create_run(
+                run_name=run_name, experiment_id=self.experiment_id, tags=tags
+            )
+        else:
+            tags[MLFLOW_RUN_NAME] = run_name
+            run = client.create_run(experiment_id=self.experiment_id, tags=tags)
+
+        return run
+
+    def _start_active_run(
+        self, run_name: Optional[str] = None, tags: Optional[Dict] = None
+    ) -> "Run":
+        """Starts a run and sets it as the active run if one does not exist.
+
+        If an active run already exists, then returns it.
+        """
+        active_run = self._mlflow.active_run()
+        if active_run:
+            return active_run
+
+        return self._mlflow.start_run(
+            run_name=run_name, experiment_id=self.experiment_id, tags=tags
+        )
+
+    def _run_exists(self, run_id: str) -> bool:
+        """Check if run with the provided id exists."""
+        from mlflow.exceptions import MlflowException
+
+        try:
+            self._mlflow.get_run(run_id=run_id)
+            return True
+        except MlflowException:
+            return False
+
+    def _get_client(self) -> "MlflowClient":
+        """Returns an ml.tracking.MlflowClient instance to use for logging."""
+        tracking_uri = self._mlflow.get_tracking_uri()
+        registry_uri = self._mlflow.get_registry_uri()
+
+        from mlflow.tracking import MlflowClient
+
+        return MlflowClient(tracking_uri=tracking_uri, registry_uri=registry_uri)
+
+    def log_params(self, params_to_log: Dict, run_id: Optional[str] = None):
+        """Logs the provided parameters to the run specified by run_id.
+
+        If no ``run_id`` is passed in, then logs to the current active run.
+        If there is not active run, then creates a new run and sets it as
+        the active run.
+
+        Args:
+            params_to_log: Dictionary of parameters to log.
+            run_id (Optional[str]): The ID of the run to log to.
+        """
+        params_to_log = flatten_dict(params_to_log)
+
+        if run_id and self._run_exists(run_id):
+            client = self._get_client()
+            for key, value in params_to_log.items():
+                client.log_param(run_id=run_id, key=key, value=value)
+
+        else:
+            for key, value in params_to_log.items():
+                self._mlflow.log_param(key=key, value=value)
+
+    def log_metrics(self, step, metrics_to_log: Dict, run_id: Optional[str] = None):
+        """Logs the provided metrics to the run specified by run_id.
+
+
+        If no ``run_id`` is passed in, then logs to the current active run.
+        If there is not active run, then creates a new run and sets it as
+        the active run.
+
+        Args:
+            metrics_to_log: Dictionary of metrics to log.
+            run_id (Optional[str]): The ID of the run to log to.
+        """
+        metrics_to_log = flatten_dict(metrics_to_log)
+        metrics_to_log = self._parse_dict(metrics_to_log)
+
+        if run_id and self._run_exists(run_id):
+            client = self._get_client()
+            for key, value in metrics_to_log.items():
+                client.log_metric(run_id=run_id, key=key, value=value, step=step)
+
+        else:
+            for key, value in metrics_to_log.items():
+                self._mlflow.log_metric(key=key, value=value, step=step)
+
+    def save_artifacts(self, dir: str, run_id: Optional[str] = None):
+        """Saves directory as artifact to the run specified by run_id.
+
+        If no ``run_id`` is passed in, then saves to the current active run.
+        If there is not active run, then creates a new run and sets it as
+        the active run.
+
+        Args:
+            dir: Path to directory containing the files to save.
+            run_id (Optional[str]): The ID of the run to log to.
+        """
+        if run_id and self._run_exists(run_id):
+            client = self._get_client()
+            client.log_artifacts(run_id=run_id, local_dir=dir)
+        else:
+            self._mlflow.log_artifacts(local_dir=dir)
+
+    def end_run(self, status: Optional[str] = None, run_id=None):
+        """Terminates the run specified by run_id.
+
+        If no ``run_id`` is passed in, then terminates the
+        active run if one exists.
+
+        Args:
+            status (Optional[str]): The status to set when terminating the run.
+            run_id (Optional[str]): The ID of the run to terminate.
+
+        """
+        if (
+            run_id
+            and self._run_exists(run_id)
+            and not (
+                self._mlflow.active_run()
+                and self._mlflow.active_run().info.run_id == run_id
+            )
+        ):
+            client = self._get_client()
+            client.set_terminated(run_id=run_id, status=status)
+        else:
+            self._mlflow.end_run(status=status)

lib/python3.12/site-packages/ray/air/_internal/tensorflow_utils.py

@@ -0,0 +1,137 @@
+from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union
+
+import numpy as np
+import pyarrow
+import tensorflow as tf
+
+from ray.air.util.data_batch_conversion import _unwrap_ndarray_object_type_if_needed
+from ray.air.util.tensor_extensions.arrow import get_arrow_extension_tensor_types
+
+if TYPE_CHECKING:
+    from ray.data._internal.pandas_block import PandasBlockSchema
+
+
+def convert_ndarray_to_tf_tensor(
+    ndarray: np.ndarray,
+    dtype: Optional[tf.dtypes.DType] = None,
+    type_spec: Optional[tf.TypeSpec] = None,
+) -> tf.Tensor:
+    """Convert a NumPy ndarray to a TensorFlow Tensor.
+
+    Args:
+        ndarray: A NumPy ndarray that we wish to convert to a TensorFlow Tensor.
+        dtype: A TensorFlow dtype for the created tensor; if None, the dtype will be
+            inferred from the NumPy ndarray data.
+        type_spec: A type spec that specifies the shape and dtype of the returned
+            tensor. If you specify ``dtype``, the dtype stored in the type spec is
+            ignored.
+
+    Returns: A TensorFlow Tensor.
+    """
+    if dtype is None and type_spec is not None:
+        dtype = type_spec.dtype
+
+    is_ragged = isinstance(type_spec, tf.RaggedTensorSpec)
+    ndarray = _unwrap_ndarray_object_type_if_needed(ndarray)
+    if is_ragged:
+        return tf.ragged.constant(ndarray, dtype=dtype)
+    else:
+        return tf.convert_to_tensor(ndarray, dtype=dtype)
+
+
+def convert_ndarray_batch_to_tf_tensor_batch(
+    ndarrays: Union[np.ndarray, Dict[str, np.ndarray]],
+    dtypes: Optional[Union[tf.dtypes.DType, Dict[str, tf.dtypes.DType]]] = None,
+) -> Union[tf.Tensor, Dict[str, tf.Tensor]]:
+    """Convert a NumPy ndarray batch to a TensorFlow Tensor batch.
+
+    Args:
+        ndarray: A (dict of) NumPy ndarray(s) that we wish to convert to a TensorFlow
+            Tensor.
+        dtype: A (dict of) TensorFlow dtype(s) for the created tensor; if None, the
+            dtype will be inferred from the NumPy ndarray data.
+
+    Returns: A (dict of) TensorFlow Tensor(s).
+    """
+    if isinstance(ndarrays, np.ndarray):
+        # Single-tensor case.
+        if isinstance(dtypes, dict):
+            if len(dtypes) != 1:
+                raise ValueError(
+                    "When constructing a single-tensor batch, only a single dtype "
+                    f"should be given, instead got: {dtypes}"
+                )
+            dtypes = next(iter(dtypes.values()))
+        batch = convert_ndarray_to_tf_tensor(ndarrays, dtypes)
+    else:
+        # Multi-tensor case.
+        batch = {
+            col_name: convert_ndarray_to_tf_tensor(
+                col_ndarray,
+                dtype=dtypes[col_name] if isinstance(dtypes, dict) else dtypes,
+            )
+            for col_name, col_ndarray in ndarrays.items()
+        }
+
+    return batch
+
+
+def get_type_spec(
+    schema: Union["pyarrow.lib.Schema", "PandasBlockSchema"],
+    columns: Union[str, List[str]],
+) -> Union[tf.TypeSpec, Dict[str, tf.TypeSpec]]:
+    import pyarrow as pa
+
+    from ray.data.extensions import TensorDtype
+
+    tensor_extension_types = get_arrow_extension_tensor_types()
+
+    assert not isinstance(schema, type)
+
+    dtypes: Dict[str, Union[np.dtype, pa.DataType]] = dict(
+        zip(schema.names, schema.types)
+    )
+
+    def get_dtype(dtype: Union[np.dtype, pa.DataType]) -> tf.dtypes.DType:
+        if isinstance(dtype, pa.ListType):
+            dtype = dtype.value_type
+        if isinstance(dtype, pa.DataType):
+            dtype = dtype.to_pandas_dtype()
+        if isinstance(dtype, TensorDtype):
+            dtype = dtype.element_dtype
+        res = tf.dtypes.as_dtype(dtype)
+        return res
+
+    def get_shape(dtype: Union[np.dtype, pa.DataType]) -> Tuple[int, ...]:
+        shape = (None,)
+        if isinstance(dtype, tensor_extension_types):
+            dtype = dtype.to_pandas_dtype()
+        if isinstance(dtype, pa.ListType):
+            shape += (None,)
+        elif isinstance(dtype, TensorDtype):
+            shape += dtype.element_shape
+        return shape
+
+    def get_tensor_spec(
+        dtype: Union[np.dtype, pa.DataType], *, name: str
+    ) -> tf.TypeSpec:
+
+        shape, dtype = get_shape(dtype), get_dtype(dtype)
+        # Batch dimension is always `None`. So, if there's more than one `None`-valued
+        # dimension, then the tensor is ragged.
+        is_ragged = sum(dim is None for dim in shape) > 1
+        if is_ragged:
+            type_spec = tf.RaggedTensorSpec(shape, dtype=dtype)
+        else:
+            type_spec = tf.TensorSpec(shape, dtype=dtype, name=name)
+        return type_spec
+
+    if isinstance(columns, str):
+        name, dtype = columns, dtypes[columns]
+        return get_tensor_spec(dtype, name=name)
+
+    return {
+        name: get_tensor_spec(dtype, name=name)
+        for name, dtype in dtypes.items()
+        if name in columns
+    }

lib/python3.12/site-packages/ray/air/_internal/torch_utils.py

@@ -0,0 +1,548 @@
+import warnings
+from typing import Any, Dict, List, Optional, Sequence, Union
+
+import numpy as np
+import pandas as pd
+import pyarrow
+import torch
+
+from ray._private.ray_constants import env_bool
+from ray.air._internal.device_manager import get_torch_device_manager_by_context
+from ray.air.util.data_batch_conversion import _unwrap_ndarray_object_type_if_needed
+from ray.data.collate_fn import (
+    TensorBatchReturnType,
+    TensorBatchType,
+    _is_nested_tensor_sequence,
+    _is_tensor,
+    _is_tensor_mapping,
+    _is_tensor_sequence,
+    _is_tensor_sequence_mapping,
+)
+
+# Default non-blocking transfer for tensors.
+DEFAULT_TENSOR_NON_BLOCKING_TRANSFER = env_bool(
+    "RAY_AIR_DEFAULT_TENSOR_NON_BLOCKING_TRANSFER",
+    True,
+)
+
+
+def get_devices() -> List[torch.device]:
+    """Gets the correct torch device list configured for this process.
+
+    Returns a list of torch accelerator (GPU, HPU, NPU...) devices allocated for
+    the current worker.
+    If no accelerators are assigned, then it returns a list with a single CPU device.
+    """
+    return get_torch_device_manager_by_context().get_devices()
+
+
+def convert_pandas_to_torch_tensor(
+    data_batch: pd.DataFrame,
+    columns: Optional[Union[List[str], List[List[str]]]] = None,
+    column_dtypes: Optional[Union[torch.dtype, List[torch.dtype]]] = None,
+    unsqueeze: bool = True,
+) -> Union[torch.Tensor, List[torch.Tensor]]:
+    """Converts a Pandas dataframe to a torch Tensor or list of torch Tensors.
+
+    The format of the return type will match the format of ``columns``. If a
+    list of columns is provided, the return type will be a single tensor. If
+    ``columns`` is a list of lists, then the return type will be a list of
+    tensors.
+
+    Args:
+        data_batch: The pandas dataframe to convert to a
+            torch tensor.
+        columns:
+            The names of the columns in the dataframe to include in the
+            torch tensor. If this arg is a List[List[str]], then the return
+            type will be a List of tensors. This is useful for multi-input
+            models. If None, then use all columns in the ``data_batch``.
+        column_dtype: The
+            torch dtype to use for the tensor. If set to None,
+            then automatically infer the dtype.
+        unsqueeze: If set to True, the tensors
+            will be unsqueezed (reshaped to (N, 1)) before being concatenated into
+            the final tensor. Otherwise, they will be left as is, that is
+            (N, ). Defaults to True.
+
+    Returns:
+        Either a torch tensor of size (N, len(columns)) where N is the
+        number of rows in the ``data_batch`` Dataframe, or a list of
+        tensors, where the size of item i is (N, len(columns[i])).
+
+    """
+
+    multi_input = columns and (isinstance(columns[0], (list, tuple)))
+
+    if not multi_input and column_dtypes and not isinstance(column_dtypes, torch.dtype):
+        raise TypeError(
+            "If `columns` is a list of strings, "
+            "`column_dtypes` must be None or a single `torch.dtype`."
+            f"Got {type(column_dtypes)} instead."
+        )
+
+    columns = columns if columns else []
+
+    def tensorize(vals, dtype):
+        """This recursive function allows to convert pyarrow List dtypes
+        to multi-dimensional tensors."""
+        if isinstance(vals, pd.api.extensions.ExtensionArray):
+            # torch.as_tensor() does not yet support the __array__ protocol, so we need
+            # to convert extension arrays to ndarrays manually before converting to a
+            # Torch tensor.
+            # See https://github.com/pytorch/pytorch/issues/51156.
+            vals = vals.to_numpy()
+
+        if vals.dtype.type is np.object_:
+            # Column has an object dtype which Torch can't handle, so we try to
+            # tensorize each column element and then stack the resulting tensors.
+            tensors = [tensorize(x, dtype) for x in vals]
+            try:
+                return torch.stack(tensors)
+            except RuntimeError:
+                # NOTE: RuntimeError is raised when trying to stack ragged tensors.
+                # Try to coerce the tensor to a nested tensor, if possible.
+                # If this fails, the exception will be propagated up to the caller.
+                return torch.nested_tensor(tensors)
+        else:
+            return torch.as_tensor(vals, dtype=dtype)
+
+    def get_tensor_for_columns(columns, dtype):
+        feature_tensors = []
+
+        if columns:
+            batch = data_batch[columns]
+        else:
+            batch = data_batch
+
+        for col in batch.columns:
+            col_vals = batch[col].values
+            try:
+                t = tensorize(col_vals, dtype=dtype)
+            except Exception as e:
+                raise ValueError(
+                    f"Failed to convert column {col} to a Torch Tensor of dtype "
+                    f"{dtype}. See above exception chain for the exact failure."
+                ) from e
+            if unsqueeze:
+                t = t.unsqueeze(1)
+            feature_tensors.append(t)
+
+        if len(feature_tensors) > 1:
+            feature_tensor = torch.cat(feature_tensors, dim=1)
+        else:
+            feature_tensor = feature_tensors[0]
+        return feature_tensor
+
+    if multi_input:
+        if type(column_dtypes) not in [list, tuple]:
+            column_dtypes = [column_dtypes] * len(columns)
+        return [
+            get_tensor_for_columns(columns=subcolumns, dtype=dtype)
+            for subcolumns, dtype in zip(columns, column_dtypes)
+        ]
+    else:
+        return get_tensor_for_columns(columns=columns, dtype=column_dtypes)
+
+
+def convert_ndarray_to_torch_tensor(
+    ndarray: np.ndarray,
+    dtype: Optional[torch.dtype] = None,
+    device: Optional[Union[str, "torch.device"]] = None,
+    pin_memory: bool = False,
+) -> torch.Tensor:
+    """Convert a NumPy ndarray to a Torch Tensor.
+
+    Args:
+        ndarray: A NumPy ndarray that we wish to convert to a Torch Tensor.
+        dtype: A Torch dtype for the created tensor; if None, the dtype will be
+            inferred from the NumPy ndarray data.
+        device: The device on which the tensor(s) should be placed; if None, the Torch
+            tensor(s) will be constructed on the CPU.
+        pin_memory: Whether to pin the memory of the created tensors.
+
+    Returns: A Torch Tensor.
+    """
+    ndarray = _unwrap_ndarray_object_type_if_needed(ndarray)
+
+    # Object dtype cannot be converted into PyTorch Tensor.
+    if ndarray.dtype.type is np.object_:
+        raise RuntimeError(
+            "Numpy array of object dtype cannot be converted to a Torch Tensor. This "
+            "may because the numpy array is a ragged tensor--it contains items of "
+            "different sizes. If using `iter_torch_batches()` API, you can pass in a "
+            "`collate_fn` argument to specify custom logic to convert the Numpy array "
+            "batch to a Torch tensor batch."
+        )
+
+    # The numpy array is not always writeable as it can come from the Ray object store.
+    # Numpy will throw a verbose warning here, which we suppress, as we don't write
+    # to the tensors. We also don't want to copy the array to avoid memory overhead.
+    # Original warning: https://github.com/pytorch/pytorch/blob/v1.13.0/
+    # torch/csrc/utils/tensor_numpy.cpp#L198-L206
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore")
+        result = torch.as_tensor(ndarray, dtype=dtype, device=device)
+
+    if pin_memory:
+        assert result.device.type == "cpu", (
+            "Pin memory is only supported for CPU tensors. "
+            f"Got device: {result.device} and pin_memory: {pin_memory}."
+        )
+        result = result.pin_memory()
+
+    return result
+
+
+def convert_ndarray_batch_to_torch_tensor_batch(
+    ndarrays: Union[np.ndarray, Dict[str, np.ndarray]],
+    dtypes: Optional[Union[torch.dtype, Dict[str, torch.dtype]]] = None,
+    device: Optional[Union[str, "torch.device"]] = None,
+    pin_memory: bool = False,
+) -> Union[torch.Tensor, Dict[str, torch.Tensor]]:
+    """Convert a NumPy ndarray batch to a Torch Tensor batch.
+
+    Args:
+        ndarray: A (dict of) NumPy ndarray(s) that we wish to convert to a Torch Tensor.
+        dtype: A (dict of) Torch dtype(s) for the created tensor; if None, the dtype
+            will be inferred from the NumPy ndarray data.
+        device: The device on which the tensor(s) should be placed; if None, the Torch
+            tensor(s) will be constructed on the CPU.
+        pin_memory: Whether to pin the memory of the created tensors.
+
+    Returns: A (dict of) Torch Tensor(s).
+    """
+    if isinstance(ndarrays, np.ndarray):
+        # Single-tensor case.
+        if isinstance(dtypes, dict):
+            if len(dtypes) != 1:
+                raise ValueError(
+                    "When constructing a single-tensor batch, only a single dtype "
+                    f"should be given, instead got: {dtypes}"
+                )
+            dtypes = next(iter(dtypes.values()))
+        batch = convert_ndarray_to_torch_tensor(
+            ndarrays,
+            dtype=dtypes,
+            device=device,
+            pin_memory=pin_memory,
+        )
+    else:
+        # Multi-tensor case.
+        batch = {
+            col_name: convert_ndarray_to_torch_tensor(
+                col_ndarray,
+                dtype=dtypes[col_name] if isinstance(dtypes, dict) else dtypes,
+                device=device,
+                pin_memory=pin_memory,
+            )
+            for col_name, col_ndarray in ndarrays.items()
+        }
+
+    return batch
+
+
+def load_torch_model(
+    saved_model: Union[torch.nn.Module, Dict],
+    model_definition: Optional[torch.nn.Module] = None,
+) -> torch.nn.Module:
+    """Loads a PyTorch model from the provided ``saved_model``.
+
+    ``model_definition`` is only used when ``saved_model`` is
+    a torch state dict, which will be loaded into ``model_definition``.
+    Otherwise, ``model_definition`` is discarded.
+    """
+    if isinstance(saved_model, torch.nn.Module):
+        return saved_model
+    elif isinstance(saved_model, dict):
+        if not model_definition:
+            raise ValueError(
+                "Attempting to load torch model from a "
+                "state_dict, but no `model_definition` was "
+                "provided."
+            )
+        model_definition.load_state_dict(saved_model)
+        return model_definition
+    else:
+        raise ValueError(
+            f"Saved model is of type {type(saved_model)}. "
+            f"The model saved in the checkpoint is expected "
+            f"to be of type `torch.nn.Module`, or a model "
+            f"state dict of type dict."
+        )
+
+
+def contains_tensor(obj):
+    if isinstance(obj, torch.Tensor):
+        return True
+    elif isinstance(obj, dict):
+        for k, v in obj.items():
+            if contains_tensor(k):
+                return True
+            if contains_tensor(v):
+                return True
+    elif isinstance(obj, (list, tuple)):
+        for v in obj:
+            if contains_tensor(v):
+                return True
+    return False
+
+
+# Not present in torch<=1.7.0
+# Adapted from https://github.com/pytorch/pytorch/blob/\
+# c18da597e0bb1c1aecc97c77a73fed1849057fa4/torch/nn/modules/utils.py
+def consume_prefix_in_state_dict_if_present_not_in_place(
+    state_dict: Dict[str, Any], prefix: str
+) -> Dict[str, Any]:
+    """Strip the prefix in state_dict, if any and return a new dict.
+
+    Adapted from https://github.com/pytorch/pytorch/blob/\
+c18da597e0bb1c1aecc97c77a73fed1849057fa4/torch/nn/modules/utils.py
+    The original method modified the dict in-place.
+
+    Args:
+        state_dict: a state-dict to be loaded to the model.
+        prefix: prefix.
+
+    """
+    copied = False
+
+    for key in state_dict:
+        if key.startswith(prefix):
+            newkey = key[len(prefix) :]
+            if not copied:
+                # We are doing shallow copies here, so the performance
+                # impact should be negligible anyway, but this is
+                # a simple optimization.
+                state_dict = state_dict.copy()
+                copied = True
+            state_dict[newkey] = state_dict.pop(key)
+
+    if "_metadata" in state_dict:
+        state_dict["_metadata"] = state_dict["_metadata"].copy()
+        metadata = state_dict["_metadata"]
+        for key in metadata:
+            if len(key) == 0:
+                continue
+            newkey = key[len(prefix) :]
+            metadata[newkey] = metadata.pop(key)
+
+    return state_dict
+
+
+def convert_ndarray_list_to_torch_tensor_list(
+    ndarrays: Dict[str, List[np.ndarray]],
+    dtypes: Optional[Union[torch.dtype, Dict[str, torch.dtype]]] = None,
+    device: Optional[Union[str, "torch.device"]] = None,
+    pin_memory: bool = False,
+) -> Dict[str, List[torch.Tensor]]:
+    """Convert a dict mapping column names to lists of ndarrays to Torch Tensors.
+
+    Args:
+        ndarrays: A dict mapping column names to lists of ndarrays that we wish to convert
+            to Torch Tensors.
+        dtypes: A (dict of) Torch dtype(s) for the created tensors; if None, the dtype
+            will be inferred from the NumPy ndarray data.
+        device: The device on which the tensor(s) should be placed; if None, the Torch
+            tensor(s) will be constructed on the CPU.
+        pin_memory: Whether to pin the memory of the created tensors.
+
+    Returns:
+        A dict mapping column names to lists of Tensors.
+    """
+    return {
+        col_name: [
+            convert_ndarray_batch_to_torch_tensor_batch(
+                ndarray,
+                dtypes=dtypes[col_name] if isinstance(dtypes, dict) else dtypes,
+                device=device,
+                pin_memory=pin_memory,
+            )
+            for ndarray in col_ndarrays
+        ]
+        for col_name, col_ndarrays in ndarrays.items()
+    }
+
+
+def arrow_batch_to_tensors(
+    batch: pyarrow.Table,
+    dtypes: Optional[Union[torch.dtype, Dict[str, torch.dtype]]] = None,
+    combine_chunks: bool = False,
+    pin_memory: bool = False,
+) -> Dict[str, List[torch.Tensor]]:
+    """Convert PyArrow batch to PyTorch tensors.
+
+    Args:
+        batch: PyArrow batch to convert
+        dtypes: A (dict of) Torch dtype(s) for the created tensors; if None, the dtype
+            will be inferred from the NumPy ndarray data.
+        combine_chunks: If True, combine chunks in Arrow batch before converting to
+            tensors.
+        pin_memory: Whether to pin the memory of the created tensors.
+
+    Returns:
+        A dictionary of column name to list of tensors. For non-chunked columns,
+        the list will contain a single tensor.
+    """
+    from ray.data._internal.arrow_block import ArrowBlockAccessor
+    from ray.data._internal.arrow_ops import transform_pyarrow
+
+    if combine_chunks:
+        numpy_batch = ArrowBlockAccessor(batch).to_batch_format("numpy")
+        return {
+            col_name: convert_ndarray_batch_to_torch_tensor_batch(
+                col_array,
+                dtypes=dtypes[col_name] if isinstance(dtypes, dict) else dtypes,
+                pin_memory=pin_memory,
+            )
+            for col_name, col_array in numpy_batch.items()
+        }
+    else:
+        numpy_list = transform_pyarrow.table_to_numpy_dict_chunked(
+            batch,
+        )
+        return convert_ndarray_list_to_torch_tensor_list(
+            numpy_list,
+            dtypes=dtypes,
+            pin_memory=pin_memory,
+        )
+
+
+@torch.no_grad()
+def concat_tensors_to_device(
+    tensor_sequence: Sequence[torch.Tensor],
+    device: Optional[Union[str, "torch.device"]] = None,
+    non_blocking: bool = DEFAULT_TENSOR_NON_BLOCKING_TRANSFER,
+) -> torch.Tensor:
+    """Stack sequence of tensors into a contiguous GPU tensor.
+
+    Args:
+        tensor_sequence: Sequence of tensors to stack
+        device: The device to move tensors to
+        non_blocking: If True, perform device transfer without forcing a
+            synchronization.
+
+    Returns:
+        A contiguous tensor on the target device
+    """
+    # Assumes tensors have the same shape/dtype
+    assert (
+        tensor_sequence
+    ), f"Cannot stack empty sequence of tensors. Received: {tensor_sequence}"
+
+    assert all(
+        isinstance(t, torch.Tensor) for t in tensor_sequence
+    ), "All items must be torch.Tensor. Found invalid types: " + str(
+        [type(t) for t in tensor_sequence if not isinstance(t, torch.Tensor)]
+    )
+
+    first_dtype = tensor_sequence[0].dtype
+    assert all(t.dtype == first_dtype for t in tensor_sequence), (
+        "All tensors must have the same dtype. "
+        f"Expected: {first_dtype}, got: {[t.dtype for t in tensor_sequence]}"
+    )
+
+    first_shape = tensor_sequence[0].shape[1:]
+    assert all(t.shape[1:] == first_shape for t in tensor_sequence), (
+        "All tensors must have the same shape[1:]. "
+        f"Expected: {first_shape}, got: {[t.shape[1:] for t in tensor_sequence]}"
+    )
+
+    first = tensor_sequence[0]
+    dtype = first.dtype
+    shape_tail = first.shape[1:]
+    total_rows = sum(t.shape[0] for t in tensor_sequence)
+
+    # Allocate an empty Tensor on device
+    result = torch.empty((total_rows, *shape_tail), dtype=dtype, device=device)
+
+    row_start = 0
+    for t in tensor_sequence:
+        row_end = row_start + t.shape[0]
+        result[row_start:row_end].copy_(t, non_blocking=non_blocking)
+        row_start = row_end
+
+    return result
+
+
+def _get_type_str(batch: Any) -> str:
+    """Get a string representation of the possibly nested type of the batch.
+
+    >>> import torch
+    >>> _get_type_str([1, 2, "???"])
+    'list[int | str]'
+    >>> _get_type_str({"a": [1, 2, 3], "b": 4})
+    'dict[str, int | list[int]]'
+    >>> _get_type_str({"a": torch.tensor(1), "b": [torch.tensor(2)]})
+    'dict[str, Tensor | list[Tensor]]'
+    >>> _get_type_str({"a": torch.tensor(1), "b": {"c": torch.tensor(2)}})
+    'dict[str, Tensor | dict[str, Tensor]]'
+    """
+    curr_type = type(batch).__name__
+    if isinstance(batch, (list, tuple)):
+        val_types = " | ".join(sorted({_get_type_str(v) for v in batch}))
+        invalid_type_str = f"{curr_type}[{val_types}]"
+    elif isinstance(batch, dict):
+        val_types = " | ".join(sorted({_get_type_str(v) for v in batch.values()}))
+        invalid_type_str = f"{curr_type}[str, {val_types}]"
+    else:
+        invalid_type_str = curr_type
+    return invalid_type_str
+
+
+@torch.no_grad()
+def move_tensors_to_device(
+    batch: TensorBatchType,
+    device: Optional[Union[str, "torch.device"]] = None,
+    non_blocking: bool = DEFAULT_TENSOR_NON_BLOCKING_TRANSFER,
+) -> TensorBatchReturnType:
+    """Move tensors to the specified device.
+
+    Concatenate nested lists/tuples of tensors along the first (batch) dimension.
+    For example, for the input
+    ((feature_0_chunk_0,), (feature_1_chunk_0, feature_1_chunk_1))
+    the output will be (feature_0_chunk_0, feature_1_chunk_0+1)
+    where each feature is concatenated along the batch dimension.
+
+    Args:
+        batch: A tensor or collection of tensors to move to device. Can be:
+            - A single tensor
+            - A sequence of tensors
+            - A sequence of sequences of tensors. The inner sequence of tensors is
+              combined during GPU transfer.
+            - A mapping (e.g., dict) of keys to tensors or sequences of tensors. The
+              sequence of tensors is combined during GPU transfer.
+        device: The device to move tensors to. If None, tensors are not moved.
+        non_blocking: If True, perform device transfer without forcing a
+            synchronization.
+
+    Returns:
+        The input tensors moved to the specified device
+    """
+    if device is None:
+        return batch
+
+    if _is_tensor(batch):
+        return batch.to(device, non_blocking=non_blocking)
+    elif _is_tensor_sequence(batch):
+        return type(batch)([t.to(device, non_blocking=non_blocking) for t in batch])
+    elif _is_nested_tensor_sequence(batch):
+        return type(batch)(
+            [concat_tensors_to_device(t, device, non_blocking) for t in batch]
+        )
+    elif _is_tensor_mapping(batch):
+        return {k: t.to(device, non_blocking=non_blocking) for k, t in batch.items()}
+    elif _is_tensor_sequence_mapping(batch):
+        return {
+            k: concat_tensors_to_device(v, device, non_blocking)
+            for k, v in batch.items()
+        }
+    else:
+        raise ValueError(
+            f"Invalid input type: {_get_type_str(batch)}.\n"
+            "Expected one of the following: "
+            "torch.Tensor, "
+            "List/Tuple[torch.Tensor], "
+            "Dict[str, torch.Tensor], "
+            "Mapping[str, List/Tuple[torch.Tensor]]"
+        )

lib/python3.12/site-packages/ray/air/_internal/uri_utils.py

@@ -0,0 +1,101 @@
+import os
+import urllib.parse
+from pathlib import Path
+from typing import Union
+
+
+class URI:
+    """Represents a URI, supporting path appending and retrieving parent URIs.
+
+    Example Usage:
+
+        >>> s3_uri = URI("s3://bucket/a?scheme=http&param=1")
+        >>> s3_uri
+        URI<s3://bucket/a?scheme=http&param=1>
+        >>> str(s3_uri / "b" / "c")
+        's3://bucket/a/b/c?scheme=http&param=1'
+        >>> str(s3_uri.parent)
+        's3://bucket?scheme=http&param=1'
+        >>> str(s3_uri)
+        's3://bucket/a?scheme=http&param=1'
+        >>> s3_uri.parent.name, s3_uri.name
+        ('bucket', 'a')
+        >>> local_path = URI("/tmp/local")
+        >>> str(local_path)
+        '/tmp/local'
+        >>> str(local_path.parent)
+        '/tmp'
+        >>> str(local_path / "b" / "c")
+        '/tmp/local/b/c'
+
+    Args:
+        uri: The URI to represent.
+            Ex: s3://bucket?scheme=http&endpoint_override=localhost%3A900
+            Ex: file:///a/b/c/d
+    """
+
+    def __init__(self, uri: str):
+        self._parsed = urllib.parse.urlparse(uri)
+        if not self._parsed.scheme:
+            # Just treat this as a regular path
+            self._path = Path(uri)
+        else:
+            self._path = Path(os.path.normpath(self._parsed.netloc + self._parsed.path))
+
+    def rstrip_subpath(self, subpath: Path) -> "URI":
+        """Returns a new URI that strips the given subpath from the end of this URI.
+
+        Example:
+            >>> uri = URI("s3://bucket/a/b/c/?param=1")
+            >>> str(uri.rstrip_subpath(Path("b/c")))
+            's3://bucket/a?param=1'
+
+            >>> uri = URI("/tmp/a/b/c/")
+            >>> str(uri.rstrip_subpath(Path("/b/c/.//")))
+            '/tmp/a'
+
+        """
+        assert str(self._path).endswith(str(subpath)), (self._path, subpath)
+        stripped_path = str(self._path).replace(str(subpath), "")
+        return URI(self._get_str_representation(self._parsed, stripped_path))
+
+    @property
+    def name(self) -> str:
+        return self._path.name
+
+    @property
+    def parent(self) -> "URI":
+        assert self._path.parent != ".", f"{str(self)} has no valid parent URI"
+        return URI(self._get_str_representation(self._parsed, self._path.parent))
+
+    @property
+    def scheme(self) -> str:
+        return self._parsed.scheme
+
+    @property
+    def path(self) -> str:
+        return str(self._path)
+
+    def __truediv__(self, path_to_append):
+        assert isinstance(path_to_append, str)
+        return URI(
+            self._get_str_representation(self._parsed, self._path / path_to_append)
+        )
+
+    @classmethod
+    def _get_str_representation(
+        cls, parsed_uri: urllib.parse.ParseResult, path: Union[str, Path]
+    ) -> str:
+        if not parsed_uri.scheme:
+            return str(path)
+        return parsed_uri._replace(netloc=str(path), path="").geturl()
+
+    def __repr__(self):
+        return f"URI<{str(self)}>"
+
+    def __str__(self):
+        return self._get_str_representation(self._parsed, self._path)
+
+
+def is_uri(path: str) -> bool:
+    return bool(urllib.parse.urlparse(path).scheme)

lib/python3.12/site-packages/ray/air/_internal/usage.py

@@ -0,0 +1,274 @@
+import collections
+import json
+import os
+from enum import Enum
+from typing import TYPE_CHECKING, Dict, List, Optional, Set, Union
+
+from ray._common.usage.usage_lib import TagKey, record_extra_usage_tag
+
+if TYPE_CHECKING:
+    from ray.train._internal.storage import StorageContext
+    from ray.train.trainer import BaseTrainer
+    from ray.tune import Callback
+    from ray.tune.schedulers import TrialScheduler
+    from ray.tune.search import BasicVariantGenerator, Searcher
+
+
+AIR_TRAINERS = {
+    "HorovodTrainer",
+    "LightGBMTrainer",
+    "TensorflowTrainer",
+    "TorchTrainer",
+    "XGBoostTrainer",
+}
+
+TRAIN_V2_TRAINERS = {
+    "DataParallelTrainer",
+    "JaxTrainer",
+    "LightGBMTrainer",
+    "TensorflowTrainer",
+    "TorchTrainer",
+    "XGBoostTrainer",
+}
+
+# searchers implemented by Ray Tune.
+TUNE_SEARCHERS = {
+    "AxSearch",
+    "BayesOptSearch",
+    "TuneBOHB",
+    "HEBOSearch",
+    "HyperOptSearch",
+    "NevergradSearch",
+    "OptunaSearch",
+    "ZOOptSearch",
+}
+
+# These are just wrappers around real searchers.
+# We don't want to double tag in this case, otherwise, the real tag
+# will be overwritten.
+TUNE_SEARCHER_WRAPPERS = {
+    "ConcurrencyLimiter",
+    "Repeater",
+}
+
+TUNE_SCHEDULERS = {
+    "FIFOScheduler",
+    "AsyncHyperBandScheduler",
+    "MedianStoppingRule",
+    "HyperBandScheduler",
+    "HyperBandForBOHB",
+    "PopulationBasedTraining",
+    "PopulationBasedTrainingReplay",
+    "PB2",
+    "ResourceChangingScheduler",
+}
+
+
+class AirEntrypoint(Enum):
+    TUNER = "Tuner.fit"
+    TRAINER = "Trainer.fit"
+    TUNE_RUN = "tune.run"
+    TUNE_RUN_EXPERIMENTS = "tune.run_experiments"
+
+
+def _find_class_name(obj, allowed_module_path_prefix: str, whitelist: Set[str]):
+    """Find the class name of the object. If the object is not
+    under `allowed_module_path_prefix` or if its class is not in the whitelist,
+    return "Custom".
+
+    Args:
+        obj: The object under inspection.
+        allowed_module_path_prefix: If the `obj`'s class is not under
+            the `allowed_module_path_prefix`, its class name will be anonymized.
+        whitelist: If the `obj`'s class is not in the `whitelist`,
+            it will be anonymized.
+    Returns:
+        The class name to be tagged with telemetry.
+    """
+    module_path = obj.__module__
+    cls_name = obj.__class__.__name__
+    if module_path.startswith(allowed_module_path_prefix) and cls_name in whitelist:
+        return cls_name
+    else:
+        return "Custom"
+
+
+def tag_air_trainer(trainer: "BaseTrainer"):
+    from ray.train.trainer import BaseTrainer
+
+    assert isinstance(trainer, BaseTrainer)
+    trainer_name = _find_class_name(trainer, "ray.train", AIR_TRAINERS)
+    record_extra_usage_tag(TagKey.AIR_TRAINER, trainer_name)
+
+
+def tag_train_v2_trainer(trainer):
+    from ray.train.v2.api.data_parallel_trainer import DataParallelTrainer
+
+    assert isinstance(trainer, DataParallelTrainer)
+    trainer_name = _find_class_name(trainer, "ray.train", TRAIN_V2_TRAINERS)
+    record_extra_usage_tag(TagKey.TRAIN_TRAINER, trainer_name)
+
+
+def tag_searcher(searcher: Union["BasicVariantGenerator", "Searcher"]):
+    from ray.tune.search import BasicVariantGenerator, Searcher
+
+    if isinstance(searcher, BasicVariantGenerator):
+        # Note this could be highly inflated as all train flows are treated
+        # as using BasicVariantGenerator.
+        record_extra_usage_tag(TagKey.TUNE_SEARCHER, "BasicVariantGenerator")
+    elif isinstance(searcher, Searcher):
+        searcher_name = _find_class_name(
+            searcher, "ray.tune.search", TUNE_SEARCHERS.union(TUNE_SEARCHER_WRAPPERS)
+        )
+        if searcher_name in TUNE_SEARCHER_WRAPPERS:
+            # ignore to avoid double tagging with wrapper name.
+            return
+        record_extra_usage_tag(TagKey.TUNE_SEARCHER, searcher_name)
+    else:
+        assert False, (
+            "Not expecting a non-BasicVariantGenerator, "
+            "non-Searcher type passed in for `tag_searcher`."
+        )
+
+
+def tag_scheduler(scheduler: "TrialScheduler"):
+    from ray.tune.schedulers import TrialScheduler
+
+    assert isinstance(scheduler, TrialScheduler)
+    scheduler_name = _find_class_name(scheduler, "ray.tune.schedulers", TUNE_SCHEDULERS)
+    record_extra_usage_tag(TagKey.TUNE_SCHEDULER, scheduler_name)
+
+
+def tag_setup_wandb():
+    record_extra_usage_tag(TagKey.AIR_SETUP_WANDB_INTEGRATION_USED, "1")
+
+
+def tag_setup_mlflow():
+    record_extra_usage_tag(TagKey.AIR_SETUP_MLFLOW_INTEGRATION_USED, "1")
+
+
+def _count_callbacks(callbacks: Optional[List["Callback"]]) -> Dict[str, int]:
+    """Creates a map of callback class name -> count given a list of callbacks."""
+    from ray.air.integrations.comet import CometLoggerCallback
+    from ray.air.integrations.mlflow import MLflowLoggerCallback
+    from ray.air.integrations.wandb import WandbLoggerCallback
+    from ray.tune import Callback
+    from ray.tune.logger import LoggerCallback
+    from ray.tune.logger.aim import AimLoggerCallback
+    from ray.tune.utils.callback import DEFAULT_CALLBACK_CLASSES
+
+    built_in_callbacks = (
+        WandbLoggerCallback,
+        MLflowLoggerCallback,
+        CometLoggerCallback,
+        AimLoggerCallback,
+    ) + DEFAULT_CALLBACK_CLASSES
+
+    callback_names = [callback_cls.__name__ for callback_cls in built_in_callbacks]
+    callback_counts = collections.defaultdict(int)
+
+    callbacks = callbacks or []
+    for callback in callbacks:
+        if not isinstance(callback, Callback):
+            # This will error later, but don't include this as custom usage.
+            continue
+
+        callback_name = callback.__class__.__name__
+
+        if callback_name in callback_names:
+            callback_counts[callback_name] += 1
+        elif isinstance(callback, LoggerCallback):
+            callback_counts["CustomLoggerCallback"] += 1
+        else:
+            callback_counts["CustomCallback"] += 1
+
+    return callback_counts
+
+
+def tag_callbacks(callbacks: Optional[List["Callback"]]) -> bool:
+    """Records built-in callback usage via a JSON str representing a
+    dictionary mapping callback class name -> counts.
+
+    User-defined callbacks will increment the count under the `CustomLoggerCallback`
+    or `CustomCallback` key depending on which of the provided interfaces they subclass.
+    NOTE: This will NOT track the name of the user-defined callback,
+    nor its implementation.
+
+    This will NOT report telemetry if no callbacks are provided by the user.
+
+    Returns:
+        bool: True if usage was recorded, False otherwise.
+    """
+    if not callbacks:
+        # User didn't pass in any callbacks -> no usage recorded.
+        return False
+
+    callback_counts = _count_callbacks(callbacks)
+
+    if callback_counts:
+        callback_counts_str = json.dumps(callback_counts)
+        record_extra_usage_tag(TagKey.AIR_CALLBACKS, callback_counts_str)
+
+
+def tag_storage_type(storage: "StorageContext"):
+    """Records the storage configuration of an experiment.
+
+    The storage configuration is set by `RunConfig(storage_path, storage_filesystem)`.
+
+    The possible storage types (defined by `pyarrow.fs.FileSystem.type_name`) are:
+    - 'local' = pyarrow.fs.LocalFileSystem. This includes NFS usage.
+    - 'mock' = pyarrow.fs._MockFileSystem. This is used for testing.
+    - ('s3', 'gcs', 'abfs', 'hdfs'): Various remote storage schemes
+        with default implementations in pyarrow.
+    - 'custom' = All other storage schemes, which includes ALL cases where a
+        custom `storage_filesystem` is provided.
+    - 'other' = catches any other cases not explicitly handled above.
+    """
+    whitelist = {"local", "mock", "s3", "gcs", "abfs", "hdfs"}
+
+    if storage.custom_fs_provided:
+        storage_config_tag = "custom"
+    elif storage.storage_filesystem.type_name in whitelist:
+        storage_config_tag = storage.storage_filesystem.type_name
+    else:
+        storage_config_tag = "other"
+
+    record_extra_usage_tag(TagKey.AIR_STORAGE_CONFIGURATION, storage_config_tag)
+
+
+def tag_ray_air_env_vars() -> bool:
+    """Records usage of environment variables exposed by the Ray AIR libraries.
+
+    NOTE: This does not track the values of the environment variables, nor
+    does this track environment variables not explicitly included in the
+    `all_ray_air_env_vars` allow-list.
+
+    Returns:
+        bool: True if at least one environment var is supplied by the user.
+    """
+    from ray.air.constants import AIR_ENV_VARS
+    from ray.train.constants import TRAIN_ENV_VARS
+    from ray.tune.constants import TUNE_ENV_VARS
+
+    all_ray_air_env_vars = sorted(
+        set().union(AIR_ENV_VARS, TUNE_ENV_VARS, TRAIN_ENV_VARS)
+    )
+
+    user_supplied_env_vars = []
+
+    for env_var in all_ray_air_env_vars:
+        if env_var in os.environ:
+            user_supplied_env_vars.append(env_var)
+
+    if user_supplied_env_vars:
+        env_vars_str = json.dumps(user_supplied_env_vars)
+        record_extra_usage_tag(TagKey.AIR_ENV_VARS, env_vars_str)
+        return True
+
+    return False
+
+
+def tag_air_entrypoint(entrypoint: AirEntrypoint) -> None:
+    """Records the entrypoint to an AIR training run."""
+    assert entrypoint in AirEntrypoint
+    record_extra_usage_tag(TagKey.AIR_ENTRYPOINT, entrypoint.value)

lib/python3.12/site-packages/ray/air/_internal/util.py

@@ -0,0 +1,125 @@
+import copy
+import logging
+import os
+import queue
+import threading
+from typing import Optional
+
+import numpy as np
+
+from ray.air.constants import _ERROR_REPORT_TIMEOUT
+
+logger = logging.getLogger(__name__)
+
+
+def is_nan(value):
+    return np.isnan(value)
+
+
+def is_nan_or_inf(value):
+    return is_nan(value) or np.isinf(value)
+
+
+class StartTraceback(Exception):
+    """These exceptions (and their tracebacks) can be skipped with `skip_exceptions`"""
+
+    pass
+
+
+class StartTracebackWithWorkerRank(StartTraceback):
+    def __init__(self, worker_rank: int) -> None:
+        super().__init__()
+        self.worker_rank = worker_rank
+
+    def __reduce__(self):
+        return (self.__class__, (self.worker_rank,))
+
+
+def skip_exceptions(exc: Optional[Exception]) -> Exception:
+    """Skip all contained `StartTracebacks` to reduce traceback output.
+
+    Returns a shallow copy of the exception with all `StartTracebacks` removed.
+
+    If the RAY_AIR_FULL_TRACEBACKS environment variable is set,
+    the original exception (not a copy) is returned.
+    """
+    should_not_shorten = bool(int(os.environ.get("RAY_AIR_FULL_TRACEBACKS", "0")))
+
+    if should_not_shorten:
+        return exc
+
+    if isinstance(exc, StartTraceback):
+        # If this is a StartTraceback, skip
+        return skip_exceptions(exc.__cause__)
+
+    # Perform a shallow copy to prevent recursive __cause__/__context__.
+    new_exc = copy.copy(exc).with_traceback(exc.__traceback__)
+
+    # Make sure nested exceptions are properly skipped.
+    cause = getattr(exc, "__cause__", None)
+    if cause:
+        new_exc.__cause__ = skip_exceptions(cause)
+
+    return new_exc
+
+
+def exception_cause(exc: Optional[Exception]) -> Optional[Exception]:
+    if not exc:
+        return None
+
+    return getattr(exc, "__cause__", None)
+
+
+class RunnerThread(threading.Thread):
+    """Supervisor thread that runs your script."""
+
+    def __init__(self, *args, error_queue, **kwargs):
+        threading.Thread.__init__(self, *args, **kwargs)
+        self._error_queue = error_queue
+        self._ret = None
+
+    def _propagate_exception(self, e: BaseException):
+        try:
+            # report the error but avoid indefinite blocking which would
+            # prevent the exception from being propagated in the unlikely
+            # case that something went terribly wrong
+            self._error_queue.put(e, block=True, timeout=_ERROR_REPORT_TIMEOUT)
+        except queue.Full:
+            logger.critical(
+                (
+                    "Runner Thread was unable to report error to main "
+                    "function runner thread. This means a previous error "
+                    "was not processed. This should never happen."
+                )
+            )
+
+    def run(self):
+        try:
+            self._ret = self._target(*self._args, **self._kwargs)
+        except StopIteration:
+            logger.debug(
+                (
+                    "Thread runner raised StopIteration. Interpreting it as a "
+                    "signal to terminate the thread without error."
+                )
+            )
+        except SystemExit as e:
+            # Do not propagate up for graceful termination.
+            if e.code == 0:
+                logger.debug(
+                    (
+                        "Thread runner raised SystemExit with error code 0. "
+                        "Interpreting it as a signal to terminate the thread "
+                        "without error."
+                    )
+                )
+            else:
+                # If non-zero exit code, then raise exception to main thread.
+                self._propagate_exception(e)
+        except BaseException as e:
+            # Propagate all other exceptions to the main thread.
+            self._propagate_exception(e)
+
+    def join(self, timeout=None):
+        super(RunnerThread, self).join(timeout)
+        return self._ret

lib/python3.12/site-packages/ray/air/config.py

@@ -0,0 +1,695 @@
+import logging
+import os
+import warnings
+from collections import Counter, defaultdict
+from dataclasses import _MISSING_TYPE, dataclass, fields
+from pathlib import Path
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Dict,
+    List,
+    Mapping,
+    Optional,
+    Tuple,
+    Union,
+)
+
+import pyarrow.fs
+
+import ray
+from ray._common.utils import RESOURCE_CONSTRAINT_PREFIX
+from ray._private.thirdparty.tabulate.tabulate import tabulate
+from ray.util.annotations import PublicAPI, RayDeprecationWarning
+from ray.widgets import Template, make_table_html_repr
+
+if TYPE_CHECKING:
+    import ray.tune.progress_reporter
+    from ray.tune.callback import Callback
+    from ray.tune.execution.placement_groups import PlacementGroupFactory
+    from ray.tune.experimental.output import AirVerbosity
+    from ray.tune.search.sample import Domain
+    from ray.tune.stopper import Stopper
+    from ray.tune.utils.log import Verbosity
+
+
+# Dict[str, List] is to support `tune.grid_search`:
+# TODO(sumanthratna/matt): Upstream this to Tune.
+SampleRange = Union["Domain", Dict[str, List]]
+
+
+MAX = "max"
+MIN = "min"
+_DEPRECATED_VALUE = "DEPRECATED"
+
+
+logger = logging.getLogger(__name__)
+
+
+def _repr_dataclass(obj, *, default_values: Optional[Dict[str, Any]] = None) -> str:
+    """A utility function to elegantly represent dataclasses.
+
+    In contrast to the default dataclass `__repr__`, which shows all parameters, this
+    function only shows parameters with non-default values.
+
+    Args:
+        obj: The dataclass to represent.
+        default_values: An optional dictionary that maps field names to default values.
+            Use this parameter to specify default values that are generated dynamically
+            (e.g., in `__post_init__` or by a `default_factory`). If a default value
+            isn't specified in `default_values`, then the default value is inferred from
+            the `dataclass`.
+
+    Returns:
+        A representation of the dataclass.
+    """
+    if default_values is None:
+        default_values = {}
+
+    non_default_values = {}  # Maps field name to value.
+
+    def equals(value, default_value):
+        # We need to special case None because of a bug in pyarrow:
+        # https://github.com/apache/arrow/issues/38535
+        if value is None and default_value is None:
+            return True
+        if value is None or default_value is None:
+            return False
+        return value == default_value
+
+    for field in fields(obj):
+        value = getattr(obj, field.name)
+        default_value = default_values.get(field.name, field.default)
+        is_required = isinstance(field.default, _MISSING_TYPE)
+        if is_required or not equals(value, default_value):
+            non_default_values[field.name] = value
+
+    string = f"{obj.__class__.__name__}("
+    string += ", ".join(
+        f"{name}={value!r}" for name, value in non_default_values.items()
+    )
+    string += ")"
+
+    return string
+
+
+@dataclass
+@PublicAPI(stability="stable")
+class ScalingConfig:
+    """Configuration for scaling training.
+
+    For more details, see :ref:`train_scaling_config`.
+
+    Args:
+        trainer_resources: Resources to allocate for the training coordinator.
+            The training coordinator launches the worker group and executes
+            the training function per worker, and this process does NOT require
+            GPUs. The coordinator is always scheduled on the same node as the
+            rank 0 worker, so one example use case is to set a minimum amount
+            of resources (e.g. CPU memory) required by the rank 0 node.
+            By default, this assigns 1 CPU to the training coordinator.
+        num_workers: The number of workers (Ray actors) to launch.
+            Each worker will reserve 1 CPU by default. The number of CPUs
+            reserved by each worker can be overridden with the
+            ``resources_per_worker`` argument.
+        use_gpu: If True, training will be done on GPUs (1 per worker).
+            Defaults to False. The number of GPUs reserved by each
+            worker can be overridden with the ``resources_per_worker``
+            argument.
+        resources_per_worker: If specified, the resources
+            defined in this Dict is reserved for each worker.
+            Define the ``"CPU"`` key (case-sensitive) to
+            override the number of CPUs used by each worker.
+            This can also be used to request :ref:`custom resources <custom-resources>`.
+        placement_strategy: The placement strategy to use for the
+            placement group of the Ray actors. See :ref:`Placement Group
+            Strategies <pgroup-strategy>` for the possible options.
+        accelerator_type: [Experimental] If specified, Ray Train will launch the
+            training coordinator and workers on the nodes with the specified type
+            of accelerators.
+            See :ref:`the available accelerator types <accelerator_types>`.
+            Ensure that your cluster has instances with the specified accelerator type
+            or is able to autoscale to fulfill the request.
+
+    Example:
+
+        .. code-block:: python
+
+            from ray.train import ScalingConfig
+            scaling_config = ScalingConfig(
+                # Number of distributed workers.
+                num_workers=2,
+                # Turn on/off GPU.
+                use_gpu=True,
+                # Assign extra CPU/GPU/custom resources per worker.
+                resources_per_worker={"GPU": 1, "CPU": 1, "memory": 1e9, "custom": 1.0},
+                # Try to schedule workers on different nodes.
+                placement_strategy="SPREAD",
+            )
+
+    """
+
+    trainer_resources: Optional[Union[Dict, SampleRange]] = None
+    num_workers: Union[int, SampleRange] = 1
+    use_gpu: Union[bool, SampleRange] = False
+    resources_per_worker: Optional[Union[Dict, SampleRange]] = None
+    placement_strategy: Union[str, SampleRange] = "PACK"
+    accelerator_type: Optional[str] = None
+
+    def __post_init__(self):
+        if self.resources_per_worker:
+            if not self.use_gpu and self.num_gpus_per_worker > 0:
+                raise ValueError(
+                    "`use_gpu` is False but `GPU` was found in "
+                    "`resources_per_worker`. Either set `use_gpu` to True or "
+                    "remove `GPU` from `resources_per_worker."
+                )
+
+            if self.use_gpu and self.num_gpus_per_worker == 0:
+                raise ValueError(
+                    "`use_gpu` is True but `GPU` is set to 0 in "
+                    "`resources_per_worker`. Either set `use_gpu` to False or "
+                    "request a positive number of `GPU` in "
+                    "`resources_per_worker."
+                )
+
+    def __repr__(self):
+        return _repr_dataclass(self)
+
+    def _repr_html_(self) -> str:
+        return make_table_html_repr(obj=self, title=type(self).__name__)
+
+    def __eq__(self, o: "ScalingConfig") -> bool:
+        if not isinstance(o, type(self)):
+            return False
+        return self.as_placement_group_factory() == o.as_placement_group_factory()
+
+    @property
+    def _resources_per_worker_not_none(self):
+        if self.resources_per_worker is None:
+            if self.use_gpu:
+                # Note that we don't request any CPUs, which avoids possible
+                # scheduling contention. Generally nodes have many more CPUs than
+                # GPUs, so not requesting a CPU does not lead to oversubscription.
+                resources_per_worker = {"GPU": 1}
+            else:
+                resources_per_worker = {"CPU": 1}
+        else:
+            resources_per_worker = {
+                k: v for k, v in self.resources_per_worker.items() if v != 0
+            }
+
+        if self.use_gpu:
+            resources_per_worker.setdefault("GPU", 1)
+
+        if self.accelerator_type:
+            accelerator = f"{RESOURCE_CONSTRAINT_PREFIX}{self.accelerator_type}"
+            resources_per_worker.setdefault(accelerator, 0.001)
+        return resources_per_worker
+
+    @property
+    def _trainer_resources_not_none(self):
+        if self.trainer_resources is None:
+            if self.num_workers:
+                # For Google Colab, don't allocate resources to the base Trainer.
+                # Colab only has 2 CPUs, and because of this resource scarcity,
+                # we have to be careful on where we allocate resources. Since Colab
+                # is not distributed, the concern about many parallel Ray Tune trials
+                # leading to all Trainers being scheduled on the head node if we set
+                # `trainer_resources` to 0 is no longer applicable.
+                try:
+                    import google.colab  # noqa: F401
+
+                    trainer_num_cpus = 0
+                except ImportError:
+                    trainer_num_cpus = 1
+            else:
+                # If there are no additional workers, then always reserve 1 CPU for
+                # the Trainer.
+                trainer_num_cpus = 1
+
+            trainer_resources = {"CPU": trainer_num_cpus}
+        else:
+            trainer_resources = {
+                k: v for k, v in self.trainer_resources.items() if v != 0
+            }
+
+        return trainer_resources
+
+    @property
+    def total_resources(self):
+        """Map of total resources required for the trainer."""
+        total_resource_map = defaultdict(float, self._trainer_resources_not_none)
+        for k, value in self._resources_per_worker_not_none.items():
+            total_resource_map[k] += value * self.num_workers
+        return dict(total_resource_map)
+
+    @property
+    def num_cpus_per_worker(self):
+        """The number of CPUs to set per worker."""
+        return self._resources_per_worker_not_none.get("CPU", 0)
+
+    @property
+    def num_gpus_per_worker(self):
+        """The number of GPUs to set per worker."""
+        return self._resources_per_worker_not_none.get("GPU", 0)
+
+    @property
+    def additional_resources_per_worker(self):
+        """Resources per worker, not including CPU or GPU resources."""
+        return {
+            k: v
+            for k, v in self._resources_per_worker_not_none.items()
+            if k not in ["CPU", "GPU"]
+        }
+
+    def as_placement_group_factory(self) -> "PlacementGroupFactory":
+        """Returns a PlacementGroupFactory to specify resources for Tune."""
+        from ray.tune.execution.placement_groups import PlacementGroupFactory
+
+        trainer_bundle = self._trainer_resources_not_none
+        worker_bundle = self._resources_per_worker_not_none
+
+        # Colocate Trainer and rank0 worker by merging their bundles
+        # Note: This empty bundle is required so that the Tune actor manager schedules
+        # the Trainable onto the combined bundle while taking none of its resources,
+        # rather than a non-empty head bundle.
+        combined_bundle = dict(Counter(trainer_bundle) + Counter(worker_bundle))
+        bundles = [{}, combined_bundle] + [worker_bundle] * (self.num_workers - 1)
+        return PlacementGroupFactory(bundles, strategy=self.placement_strategy)
+
+    @classmethod
+    def from_placement_group_factory(
+        cls, pgf: "PlacementGroupFactory"
+    ) -> "ScalingConfig":
+        """Create a ScalingConfig from a Tune's PlacementGroupFactory
+
+        Note that this is only needed for ResourceChangingScheduler, which
+        modifies a trial's PlacementGroupFactory but doesn't propagate
+        the changes to ScalingConfig. TrainTrainable needs to reconstruct
+        a ScalingConfig from on the trial's PlacementGroupFactory.
+        """
+
+        # pgf.bundles = [{trainer + worker}, {worker}, ..., {worker}]
+        num_workers = len(pgf.bundles)
+        combined_resources = pgf.bundles[0]
+        resources_per_worker = pgf.bundles[-1]
+        use_gpu = bool(resources_per_worker.get("GPU", False))
+        placement_strategy = pgf.strategy
+
+        # In `as_placement_group_factory`, we merged the trainer resource into the
+        # first worker resources bundle. We need to calculate the resources diff to
+        # get the trainer resources.
+        # Note: If there's only one worker, we won't be able to calculate the diff.
+        # We'll have empty trainer bundle and assign all resources to the worker.
+        trainer_resources = dict(
+            Counter(combined_resources) - Counter(resources_per_worker)
+        )
+
+        return ScalingConfig(
+            trainer_resources=trainer_resources,
+            num_workers=num_workers,
+            use_gpu=use_gpu,
+            resources_per_worker=resources_per_worker,
+            placement_strategy=placement_strategy,
+        )
+
+
+@dataclass
+@PublicAPI(stability="stable")
+class FailureConfig:
+    """Configuration related to failure handling of each training/tuning run.
+
+    Args:
+        max_failures: Tries to recover a run at least this many times.
+            Will recover from the latest checkpoint if present.
+            Setting to -1 will lead to infinite recovery retries.
+            Setting to 0 will disable retries. Defaults to 0.
+        fail_fast: Whether to fail upon the first error.
+            If fail_fast='raise' provided, the original error during training will be
+            immediately raised. fail_fast='raise' can easily leak resources and
+            should be used with caution.
+    """
+
+    max_failures: int = 0
+    fail_fast: Union[bool, str] = False
+
+    def __post_init__(self):
+        # Same check as in TuneController
+        if not (isinstance(self.fail_fast, bool) or self.fail_fast.upper() == "RAISE"):
+            raise ValueError(
+                "fail_fast must be one of {bool, 'raise'}. " f"Got {self.fail_fast}."
+            )
+
+        # Same check as in tune.run
+        if self.fail_fast and self.max_failures != 0:
+            raise ValueError(
+                f"max_failures must be 0 if fail_fast={repr(self.fail_fast)}."
+            )
+
+    def __repr__(self):
+        return _repr_dataclass(self)
+
+    def _repr_html_(self):
+        return Template("scrollableTable.html.j2").render(
+            table=tabulate(
+                {
+                    "Setting": ["Max failures", "Fail fast"],
+                    "Value": [self.max_failures, self.fail_fast],
+                },
+                tablefmt="html",
+                showindex=False,
+                headers="keys",
+            ),
+            max_height="none",
+        )
+
+
+@dataclass
+@PublicAPI(stability="stable")
+class CheckpointConfig:
+    """Configurable parameters for defining the checkpointing strategy.
+
+    Default behavior is to persist all checkpoints to disk. If
+    ``num_to_keep`` is set, the default retention policy is to keep the
+    checkpoints with maximum timestamp, i.e. the most recent checkpoints.
+
+    Args:
+        num_to_keep: The number of checkpoints to keep
+            on disk for this run. If a checkpoint is persisted to disk after
+            there are already this many checkpoints, then an existing
+            checkpoint will be deleted. If this is ``None`` then checkpoints
+            will not be deleted. Must be >= 1.
+        checkpoint_score_attribute: The attribute that will be used to
+            score checkpoints to determine which checkpoints should be kept
+            on disk when there are greater than ``num_to_keep`` checkpoints.
+            This attribute must be a key from the checkpoint
+            dictionary which has a numerical value. Per default, the last
+            checkpoints will be kept.
+        checkpoint_score_order: Either "max" or "min".
+            If "max", then checkpoints with highest values of
+            ``checkpoint_score_attribute`` will be kept.
+            If "min", then checkpoints with lowest values of
+            ``checkpoint_score_attribute`` will be kept.
+        checkpoint_frequency: Number of iterations between checkpoints. If 0
+            this will disable checkpointing.
+            Please note that most trainers will still save one checkpoint at
+            the end of training.
+            This attribute is only supported
+            by trainers that don't take in custom training loops.
+        checkpoint_at_end: If True, will save a checkpoint at the end of training.
+            This attribute is only supported by trainers that don't take in
+            custom training loops. Defaults to True for trainers that support it
+            and False for generic function trainables.
+        _checkpoint_keep_all_ranks: This experimental config is deprecated.
+            This behavior is now controlled by reporting `checkpoint=None`
+            in the workers that shouldn't persist a checkpoint.
+            For example, if you only want the rank 0 worker to persist a checkpoint
+            (e.g., in standard data parallel training), then you should save and
+            report a checkpoint if `ray.train.get_context().get_world_rank() == 0`
+            and `None` otherwise.
+        _checkpoint_upload_from_workers: This experimental config is deprecated.
+            Uploading checkpoint directly from the worker is now the default behavior.
+    """
+
+    num_to_keep: Optional[int] = None
+    checkpoint_score_attribute: Optional[str] = None
+    checkpoint_score_order: Optional[str] = MAX
+    checkpoint_frequency: Optional[int] = 0
+    checkpoint_at_end: Optional[bool] = None
+    _checkpoint_keep_all_ranks: Optional[bool] = _DEPRECATED_VALUE
+    _checkpoint_upload_from_workers: Optional[bool] = _DEPRECATED_VALUE
+
+    def __post_init__(self):
+        if self._checkpoint_keep_all_ranks != _DEPRECATED_VALUE:
+            raise DeprecationWarning(
+                "The experimental `_checkpoint_keep_all_ranks` config is deprecated. "
+                "This behavior is now controlled by reporting `checkpoint=None` "
+                "in the workers that shouldn't persist a checkpoint. "
+                "For example, if you only want the rank 0 worker to persist a "
+                "checkpoint (e.g., in standard data parallel training), "
+                "then you should save and report a checkpoint if "
+                "`ray.train.get_context().get_world_rank() == 0` "
+                "and `None` otherwise."
+            )
+
+        if self._checkpoint_upload_from_workers != _DEPRECATED_VALUE:
+            raise DeprecationWarning(
+                "The experimental `_checkpoint_upload_from_workers` config is "
+                "deprecated. Uploading checkpoint directly from the worker is "
+                "now the default behavior."
+            )
+
+        if self.num_to_keep is not None and self.num_to_keep <= 0:
+            raise ValueError(
+                f"Received invalid num_to_keep: "
+                f"{self.num_to_keep}. "
+                f"Must be None or an integer >= 1."
+            )
+        if self.checkpoint_score_order not in (MAX, MIN):
+            raise ValueError(
+                f"checkpoint_score_order must be either " f'"{MAX}" or "{MIN}".'
+            )
+
+        if self.checkpoint_frequency < 0:
+            raise ValueError(
+                f"checkpoint_frequency must be >=0, got {self.checkpoint_frequency}"
+            )
+
+    def __repr__(self):
+        return _repr_dataclass(self)
+
+    def _repr_html_(self) -> str:
+        if self.num_to_keep is None:
+            num_to_keep_repr = "All"
+        else:
+            num_to_keep_repr = self.num_to_keep
+
+        if self.checkpoint_score_attribute is None:
+            checkpoint_score_attribute_repr = "Most recent"
+        else:
+            checkpoint_score_attribute_repr = self.checkpoint_score_attribute
+
+        if self.checkpoint_at_end is None:
+            checkpoint_at_end_repr = ""
+        else:
+            checkpoint_at_end_repr = self.checkpoint_at_end
+
+        return Template("scrollableTable.html.j2").render(
+            table=tabulate(
+                {
+                    "Setting": [
+                        "Number of checkpoints to keep",
+                        "Checkpoint score attribute",
+                        "Checkpoint score order",
+                        "Checkpoint frequency",
+                        "Checkpoint at end",
+                    ],
+                    "Value": [
+                        num_to_keep_repr,
+                        checkpoint_score_attribute_repr,
+                        self.checkpoint_score_order,
+                        self.checkpoint_frequency,
+                        checkpoint_at_end_repr,
+                    ],
+                },
+                tablefmt="html",
+                showindex=False,
+                headers="keys",
+            ),
+            max_height="none",
+        )
+
+    @property
+    def _tune_legacy_checkpoint_score_attr(self) -> Optional[str]:
+        """Same as ``checkpoint_score_attr`` in ``tune.run``.
+
+        Only used for Legacy API compatibility.
+        """
+        if self.checkpoint_score_attribute is None:
+            return self.checkpoint_score_attribute
+        prefix = ""
+        if self.checkpoint_score_order == MIN:
+            prefix = "min-"
+        return f"{prefix}{self.checkpoint_score_attribute}"
+
+
+@dataclass
+@PublicAPI(stability="stable")
+class RunConfig:
+    """Runtime configuration for training and tuning runs.
+
+    Upon resuming from a training or tuning run checkpoint,
+    Ray Train/Tune will automatically apply the RunConfig from
+    the previously checkpointed run.
+
+    Args:
+        name: Name of the trial or experiment. If not provided, will be deduced
+            from the Trainable.
+        storage_path: [Beta] Path where all results and checkpoints are persisted.
+            Can be a local directory or a destination on cloud storage.
+            For multi-node training/tuning runs, this must be set to a
+            shared storage location (e.g., S3, NFS).
+            This defaults to the local ``~/ray_results`` directory.
+        storage_filesystem: [Beta] A custom filesystem to use for storage.
+            If this is provided, `storage_path` should be a path with its
+            prefix stripped (e.g., `s3://bucket/path` -> `bucket/path`).
+        failure_config: Failure mode configuration.
+        checkpoint_config: Checkpointing configuration.
+        sync_config: Configuration object for syncing. See train.SyncConfig.
+        verbose: 0, 1, or 2. Verbosity mode.
+            0 = silent, 1 = default, 2 = verbose. Defaults to 1.
+            If the ``RAY_AIR_NEW_OUTPUT=1`` environment variable is set,
+            uses the old verbosity settings:
+            0 = silent, 1 = only status updates, 2 = status and brief
+            results, 3 = status and detailed results.
+        stop: Stop conditions to consider. Refer to ray.tune.stopper.Stopper
+            for more info. Stoppers should be serializable.
+        callbacks: [DeveloperAPI] Callbacks to invoke.
+            Refer to ray.tune.callback.Callback for more info.
+            Callbacks should be serializable.
+            Currently only stateless callbacks are supported for resumed runs.
+            (any state of the callback will not be checkpointed by Tune
+            and thus will not take effect in resumed runs).
+        progress_reporter: [DeveloperAPI] Progress reporter for reporting
+            intermediate experiment progress. Defaults to CLIReporter if
+            running in command-line, or JupyterNotebookReporter if running in
+            a Jupyter notebook.
+        log_to_file: [DeveloperAPI] Log stdout and stderr to files in
+            trial directories. If this is `False` (default), no files
+            are written. If `true`, outputs are written to `trialdir/stdout`
+            and `trialdir/stderr`, respectively. If this is a single string,
+            this is interpreted as a file relative to the trialdir, to which
+            both streams are written. If this is a Sequence (e.g. a Tuple),
+            it has to have length 2 and the elements indicate the files to
+            which stdout and stderr are written, respectively.
+
+    """
+
+    name: Optional[str] = None
+    storage_path: Optional[str] = None
+    storage_filesystem: Optional[pyarrow.fs.FileSystem] = None
+    failure_config: Optional[FailureConfig] = None
+    checkpoint_config: Optional[CheckpointConfig] = None
+    sync_config: Optional["ray.train.SyncConfig"] = None
+    verbose: Optional[Union[int, "AirVerbosity", "Verbosity"]] = None
+    stop: Optional[Union[Mapping, "Stopper", Callable[[str, Mapping], bool]]] = None
+    callbacks: Optional[List["Callback"]] = None
+    progress_reporter: Optional["ray.tune.progress_reporter.ProgressReporter"] = None
+    log_to_file: Union[bool, str, Tuple[str, str]] = False
+
+    # Deprecated
+    local_dir: Optional[str] = None
+
+    def __post_init__(self):
+        from ray.train import SyncConfig
+        from ray.train.constants import DEFAULT_STORAGE_PATH
+        from ray.tune.experimental.output import AirVerbosity, get_air_verbosity
+
+        if self.local_dir is not None:
+            raise DeprecationWarning(
+                "The `RunConfig(local_dir)` argument is deprecated. "
+                "You should set the `RunConfig(storage_path)` instead."
+                "See the docs: https://docs.ray.io/en/latest/train/user-guides/"
+                "persistent-storage.html#setting-the-local-staging-directory"
+            )
+
+        if self.storage_path is None:
+            self.storage_path = DEFAULT_STORAGE_PATH
+
+            # TODO(justinvyu): [Deprecated]
+            ray_storage_uri: Optional[str] = os.environ.get("RAY_STORAGE")
+            if ray_storage_uri is not None:
+                logger.info(
+                    "Using configured Ray Storage URI as the `storage_path`: "
+                    f"{ray_storage_uri}"
+                )
+                warnings.warn(
+                    "The `RAY_STORAGE` environment variable is deprecated. "
+                    "Please use `RunConfig(storage_path)` instead.",
+                    RayDeprecationWarning,
+                    stacklevel=2,
+                )
+                self.storage_path = ray_storage_uri
+
+        if not self.failure_config:
+            self.failure_config = FailureConfig()
+
+        if not self.sync_config:
+            self.sync_config = SyncConfig()
+
+        if not self.checkpoint_config:
+            self.checkpoint_config = CheckpointConfig()
+
+        # Save the original verbose value to check for deprecations
+        self._verbose = self.verbose
+        if self.verbose is None:
+            # Default `verbose` value. For new output engine,
+            # this is AirVerbosity.DEFAULT.
+            # For old output engine, this is Verbosity.V3_TRIAL_DETAILS
+            # Todo (krfricke): Currently uses number to pass test_configs::test_repr
+            self.verbose = get_air_verbosity(AirVerbosity.DEFAULT) or 3
+
+        if isinstance(self.storage_path, Path):
+            self.storage_path = self.storage_path.as_posix()
+
+    def __repr__(self):
+        from ray.train import SyncConfig
+
+        return _repr_dataclass(
+            self,
+            default_values={
+                "failure_config": FailureConfig(),
+                "sync_config": SyncConfig(),
+                "checkpoint_config": CheckpointConfig(),
+            },
+        )
+
+    def _repr_html_(self) -> str:
+        reprs = []
+        if self.failure_config is not None:
+            reprs.append(
+                Template("title_data_mini.html.j2").render(
+                    title="Failure Config", data=self.failure_config._repr_html_()
+                )
+            )
+        if self.sync_config is not None:
+            reprs.append(
+                Template("title_data_mini.html.j2").render(
+                    title="Sync Config", data=self.sync_config._repr_html_()
+                )
+            )
+        if self.checkpoint_config is not None:
+            reprs.append(
+                Template("title_data_mini.html.j2").render(
+                    title="Checkpoint Config", data=self.checkpoint_config._repr_html_()
+                )
+            )
+
+        # Create a divider between each displayed repr
+        subconfigs = [Template("divider.html.j2").render()] * (2 * len(reprs) - 1)
+        subconfigs[::2] = reprs
+
+        settings = Template("scrollableTable.html.j2").render(
+            table=tabulate(
+                {
+                    "Name": self.name,
+                    "Local results directory": self.local_dir,
+                    "Verbosity": self.verbose,
+                    "Log to file": self.log_to_file,
+                }.items(),
+                tablefmt="html",
+                headers=["Setting", "Value"],
+                showindex=False,
+            ),
+            max_height="300px",
+        )
+
+        return Template("title_data.html.j2").render(
+            title="RunConfig",
+            data=Template("run_config.html.j2").render(
+                subconfigs=subconfigs,
+                settings=settings,
+            ),
+        )

lib/python3.12/site-packages/ray/air/constants.py

@@ -0,0 +1,94 @@
+# Key to denote the preprocessor in the checkpoint dict.
+PREPROCESSOR_KEY = "_preprocessor"
+
+# Key to denote the model in the checkpoint dict.
+MODEL_KEY = "model"
+
+# Key to denote which dataset is the evaluation dataset.
+# Only used in trainers which do not support multiple
+# evaluation datasets.
+EVALUATION_DATASET_KEY = "evaluation"
+
+# Key to denote which dataset is the training dataset.
+# This is the dataset that the preprocessor is fit on.
+TRAIN_DATASET_KEY = "train"
+
+# Name to use for the column when representing tensors in table format.
+TENSOR_COLUMN_NAME = "__value__"
+
+# The maximum length of strings returned by `__repr__` for AIR objects constructed with
+# default values.
+MAX_REPR_LENGTH = int(80 * 1.5)
+
+# Timeout used when putting exceptions raised by runner thread into the queue.
+_ERROR_REPORT_TIMEOUT = 10
+
+# Timeout when fetching new results after signaling the training function to continue.
+_RESULT_FETCH_TIMEOUT = 0.2
+
+# Timeout for fetching exceptions raised by the training function.
+_ERROR_FETCH_TIMEOUT = 1
+
+# The key used to identify whether we have already warned about ray.air.session
+# functions being used outside of the session
+SESSION_MISUSE_LOG_ONCE_KEY = "air_warn_session_misuse"
+
+# Name of attribute in Checkpoint storing current Tune ID for restoring
+# training with Ray Train
+CHECKPOINT_ID_ATTR = "_current_checkpoint_id"
+
+# Name of the marker dropped by the Trainable. If a worker detects
+# the presence of the marker in the trial dir, it will use lazy
+# checkpointing.
+LAZY_CHECKPOINT_MARKER_FILE = ".lazy_checkpoint_marker"
+
+
+# The timestamp of when the result is generated.
+# Default to when the result is processed by tune.
+TIMESTAMP = "timestamp"
+
+# (Auto-filled) Time in seconds this iteration took to run.
+# This may be overridden to override the system-computed time difference.
+TIME_THIS_ITER_S = "time_this_iter_s"
+
+# (Auto-filled) The index of this training iteration.
+TRAINING_ITERATION = "training_iteration"
+
+# File that stores parameters of the trial.
+EXPR_PARAM_FILE = "params.json"
+
+# Pickle File that stores parameters of the trial.
+EXPR_PARAM_PICKLE_FILE = "params.pkl"
+
+# File that stores the progress of the trial.
+EXPR_PROGRESS_FILE = "progress.csv"
+
+# File that stores results of the trial.
+EXPR_RESULT_FILE = "result.json"
+
+# File that stores the pickled error file
+EXPR_ERROR_PICKLE_FILE = "error.pkl"
+
+# File that stores the error file
+EXPR_ERROR_FILE = "error.txt"
+
+# File that stores the checkpoint metadata
+CHECKPOINT_TUNE_METADATA_FILE = ".tune_metadata"
+
+# ==================================================
+#               Environment Variables
+# ==================================================
+
+# Integer value which if set will copy files in reported AIR directory
+# checkpoints instead of moving them (if worker is on the same node as Trainable)
+COPY_DIRECTORY_CHECKPOINTS_INSTEAD_OF_MOVING_ENV = (
+    "TRAIN_COPY_DIRECTORY_CHECKPOINTS_INSTEAD_OF_MOVING"
+)
+
+# NOTE: When adding a new environment variable, please track it in this list.
+# TODO(ml-team): Most env var constants should get moved here.
+AIR_ENV_VARS = {
+    COPY_DIRECTORY_CHECKPOINTS_INSTEAD_OF_MOVING_ENV,
+    "RAY_AIR_FULL_TRACEBACKS",
+    "RAY_AIR_NEW_OUTPUT",
+}

lib/python3.12/site-packages/ray/air/data_batch_type.py

@@ -0,0 +1,11 @@
+from typing import TYPE_CHECKING, Dict, Union
+
+if TYPE_CHECKING:
+    import numpy
+    import pandas  # noqa: F401
+    import pyarrow
+
+# TODO de-dup with ray.data.block.DataBatch
+DataBatchType = Union[
+    "numpy.ndarray", "pyarrow.Table" "pandas.DataFrame", Dict[str, "numpy.ndarray"]
+]

lib/python3.12/site-packages/ray/air/execution/__init__.py

@@ -0,0 +1,12 @@
+from ray.air.execution.resources.fixed import FixedResourceManager
+from ray.air.execution.resources.placement_group import PlacementGroupResourceManager
+from ray.air.execution.resources.request import AcquiredResources, ResourceRequest
+from ray.air.execution.resources.resource_manager import ResourceManager
+
+__all__ = [
+    "ResourceRequest",
+    "AcquiredResources",
+    "ResourceManager",
+    "FixedResourceManager",
+    "PlacementGroupResourceManager",
+]

lib/python3.12/site-packages/ray/air/execution/_internal/__init__.py

@@ -0,0 +1,5 @@
+from ray.air.execution._internal.actor_manager import RayActorManager
+from ray.air.execution._internal.barrier import Barrier
+from ray.air.execution._internal.tracked_actor import TrackedActor
+
+__all__ = ["Barrier", "RayActorManager", "TrackedActor"]